mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 02:05:41 +00:00
fix: prevent schema updates from racing insert payload conversion (#51483)
## Problem During schema evolution, SQN can convert an insert payload with the old schema and then race with a schema update before ProcessInsert creates and writes the growing segment. If the old payload still contains a dropped field, the native collection may already use the new schema and segcore rejects the write with unexpected new field. ## Solution - Add a collection-local schema-transition RW lock. - Insert holds the reader lock from schema snapshot use through payload conversion and delegator.ProcessInsert, so a growing segment is created and written in the same schema epoch as its payload. - Schema update and existing-collection load refresh paths take the writer lock around native schema mutation and Go schema snapshot publication. - Acquire a temporary collection lease before releasing collectionManager.mut. This keeps the collection alive while the writer waits, without serializing unrelated collection-manager operations behind a schema transition. - Preserve dropped-field replay handling: when DDL has already completed, payload conversion filters fields absent from the current schema before inserting into a new-schema growing segment. ## Tests - Add a deterministic native regression: pause after a v950 payload containing field 580 is converted, queue the v951 drop-field update, then resume a real NewSegment plus growing.Insert and verify both insert and DDL succeed. - Assert the transition reader covers both payload conversion and ProcessInsert. - Cover the DDL-first replay path, where field 580 is filtered before native insertion. - Cover both schema writer entry points: UpdateSchema and existing-collection PutOrRef. ## Validation - make static-check - make build-go - Native pipeline regression tests repeated 20 times with dynamic,test - Collection schema-transition and lease tests repeated 20 times - internal/querynodev2/pipeline and internal/querynodev2/segments package tests - gofumpt -l and git diff --check Fixes #51436 Signed-off-by: sijie-ni-0214 <sijie.ni@zilliz.com>
This commit is contained in:
@@ -44,9 +44,8 @@ type insertNode struct {
|
||||
functionStore *function.FunctionRunnerLocalStore
|
||||
}
|
||||
|
||||
func (iNode *insertNode) addInsertData(insertDatas map[UniqueID]*delegator.InsertData, msg *InsertMsg, collection *Collection) {
|
||||
func (iNode *insertNode) addInsertData(insertDatas map[UniqueID]*delegator.InsertData, msg *InsertMsg, schema *schemapb.CollectionSchema) {
|
||||
ctx := msg.TraceCtx()
|
||||
schema := collection.Schema()
|
||||
insertRecord, skippedFields, err := storage.TransferInsertMsgToInsertRecord(schema, msg)
|
||||
if err != nil {
|
||||
err = merr.Wrap(err, "failed to get primary keys")
|
||||
@@ -132,25 +131,26 @@ func (iNode *insertNode) Operate(in Msg) Msg {
|
||||
mlog.Error(ctx, "insertNode with collection not exist", mlog.Int64("collection", iNode.collectionID))
|
||||
panic("insertNode with collection not exist")
|
||||
}
|
||||
schema := collection.Schema()
|
||||
functionOutputFieldIDs, err := iNode.functionStore.OutputFieldIDs(schema)
|
||||
if err != nil {
|
||||
mlog.Error(ctx, "failed to get embedding output fields", mlog.String("channel", iNode.channel), mlog.Err(err))
|
||||
panic(err)
|
||||
}
|
||||
|
||||
insertDatas := make(map[UniqueID]*delegator.InsertData)
|
||||
for _, msg := range nodeMsg.insertMsgs {
|
||||
if len(functionOutputFieldIDs) > 0 && !function.HasAllFieldDataByID(msg.GetFieldsData(), functionOutputFieldIDs) {
|
||||
if err := iNode.functionStore.FillEmbeddingData(iNode.collectionID, schema, msg.InsertRequest); err != nil {
|
||||
mlog.Error(msg.TraceCtx(), "failed to fill embedding data for insert message", mlog.String("channel", iNode.channel), mlog.Err(err))
|
||||
panic(err)
|
||||
}
|
||||
collection.WithInsertSchemaTransition(func(schema *schemapb.CollectionSchema) {
|
||||
functionOutputFieldIDs, err := iNode.functionStore.OutputFieldIDs(schema)
|
||||
if err != nil {
|
||||
mlog.Error(ctx, "failed to get embedding output fields", mlog.String("channel", iNode.channel), mlog.Err(err))
|
||||
panic(err)
|
||||
}
|
||||
iNode.addInsertData(insertDatas, msg, collection)
|
||||
}
|
||||
|
||||
iNode.delegator.ProcessInsert(insertDatas)
|
||||
insertDatas := make(map[UniqueID]*delegator.InsertData)
|
||||
for _, msg := range nodeMsg.insertMsgs {
|
||||
if len(functionOutputFieldIDs) > 0 && !function.HasAllFieldDataByID(msg.GetFieldsData(), functionOutputFieldIDs) {
|
||||
if err := iNode.functionStore.FillEmbeddingData(iNode.collectionID, schema, msg.InsertRequest); err != nil {
|
||||
mlog.Error(msg.TraceCtx(), "failed to fill embedding data for insert message", mlog.String("channel", iNode.channel), mlog.Err(err))
|
||||
panic(err)
|
||||
}
|
||||
}
|
||||
iNode.addInsertData(insertDatas, msg, schema)
|
||||
}
|
||||
|
||||
iNode.delegator.ProcessInsert(insertDatas)
|
||||
})
|
||||
}
|
||||
metrics.QueryNodeWaitProcessingMsgCount.WithLabelValues(paramtable.GetStringNodeID(), metrics.DeleteLabel).Inc()
|
||||
|
||||
|
||||
@@ -17,18 +17,29 @@
|
||||
package pipeline
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/bytedance/mockey"
|
||||
"github.com/samber/lo"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/stretchr/testify/suite"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
|
||||
"github.com/milvus-io/milvus-proto/go-api/v3/schemapb"
|
||||
"github.com/milvus-io/milvus/internal/mocks/util/mock_segcore"
|
||||
"github.com/milvus-io/milvus/internal/querynodev2/delegator"
|
||||
"github.com/milvus-io/milvus/internal/querynodev2/segments"
|
||||
"github.com/milvus-io/milvus/internal/storage"
|
||||
"github.com/milvus-io/milvus/internal/util/initcore"
|
||||
"github.com/milvus-io/milvus/pkg/v3/mq/msgstream"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/querypb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/proto/segcorepb"
|
||||
"github.com/milvus-io/milvus/pkg/v3/util/paramtable"
|
||||
)
|
||||
|
||||
@@ -80,8 +91,16 @@ func (suite *InsertNodeSuite) TestBasic() {
|
||||
Segment: mockSegmentManager,
|
||||
}
|
||||
|
||||
var transferOrigin func(*schemapb.CollectionSchema, *msgstream.InsertMsg) (*segcorepb.InsertRecord, []int64, error)
|
||||
transferMock := mockey.Mock(storage.TransferInsertMsgToInsertRecord).To(func(schema *schemapb.CollectionSchema, msg *msgstream.InsertMsg) (*segcorepb.InsertRecord, []int64, error) {
|
||||
suite.True(collection.HasInsertSchemaTransitionReaderForTest())
|
||||
return transferOrigin(schema, msg)
|
||||
}).Origin(&transferOrigin).Build()
|
||||
defer transferMock.UnPatch()
|
||||
|
||||
suite.delegator = delegator.NewMockShardDelegator(suite.T())
|
||||
suite.delegator.EXPECT().ProcessInsert(mock.Anything).Run(func(insertRecords map[int64]*delegator.InsertData) {
|
||||
suite.True(collection.HasInsertSchemaTransitionReaderForTest())
|
||||
for segID := range insertRecords {
|
||||
suite.True(lo.Contains(suite.insertSegmentIDs, segID))
|
||||
}
|
||||
@@ -233,3 +252,212 @@ func (suite *InsertNodeSuite) buildInsertNodeMsg(schema *schemapb.CollectionSche
|
||||
func TestInsertNode(t *testing.T) {
|
||||
suite.Run(t, new(InsertNodeSuite))
|
||||
}
|
||||
|
||||
const (
|
||||
schemaTransitionCollectionID = int64(1000)
|
||||
schemaTransitionPartitionID = int64(1001)
|
||||
schemaTransitionSegmentID = int64(1002)
|
||||
schemaTransitionDroppedField = int64(580)
|
||||
)
|
||||
|
||||
func setupSchemaTransitionInsertNodeTest(t *testing.T) (*segments.Manager, *segments.Collection, *schemapb.CollectionSchema, *schemapb.CollectionSchema, *msgstream.InsertMsg) {
|
||||
t.Helper()
|
||||
paramtable.Init()
|
||||
_ = initcore.InitLocalChunkManager(t.Name())
|
||||
_ = initcore.InitMmapManager(paramtable.Get(), 1)
|
||||
|
||||
schemaV950 := mock_segcore.GenTestCollectionSchema("schema_transition", schemapb.DataType_Int64, false)
|
||||
schemaV950.Version = 950
|
||||
schemaV950.Fields = append(schemaV950.Fields, &schemapb.FieldSchema{
|
||||
FieldID: schemaTransitionDroppedField,
|
||||
Name: "stress_extra",
|
||||
DataType: schemapb.DataType_VarChar,
|
||||
Nullable: true,
|
||||
TypeParams: []*commonpb.KeyValuePair{
|
||||
{Key: "max_length", Value: "64"},
|
||||
},
|
||||
})
|
||||
schemaV951 := proto.Clone(schemaV950).(*schemapb.CollectionSchema)
|
||||
schemaV951.Version = 951
|
||||
schemaV951.Fields = lo.Filter(schemaV951.Fields, func(field *schemapb.FieldSchema, _ int) bool {
|
||||
return field.GetFieldID() != schemaTransitionDroppedField
|
||||
})
|
||||
|
||||
manager := segments.NewManager()
|
||||
require.NoError(t, manager.Collection.PutOrRef(schemaTransitionCollectionID, schemaV950, nil, &querypb.LoadMetaInfo{
|
||||
CollectionID: schemaTransitionCollectionID,
|
||||
LoadType: querypb.LoadType_LoadCollection,
|
||||
}))
|
||||
t.Cleanup(func() {
|
||||
manager.Collection.Unref(schemaTransitionCollectionID, 1)
|
||||
})
|
||||
|
||||
collection := manager.Collection.Get(schemaTransitionCollectionID)
|
||||
require.NotNil(t, collection)
|
||||
insertMsg, err := mock_segcore.GenInsertMsg(collection.GetCCollection(), schemaTransitionPartitionID, schemaTransitionSegmentID, 1)
|
||||
require.NoError(t, err)
|
||||
for _, fieldData := range insertMsg.GetFieldsData() {
|
||||
if fieldData.GetFieldId() == schemaTransitionDroppedField {
|
||||
fieldData.ValidData = []bool{true}
|
||||
}
|
||||
}
|
||||
return manager, collection, schemaV950, schemaV951, insertMsg
|
||||
}
|
||||
|
||||
func insertIntoNewGrowingSegment(collection *segments.Collection, manager *segments.Manager, segmentID int64, data *delegator.InsertData) error {
|
||||
ctx := context.Background()
|
||||
growing, err := segments.NewSegment(ctx, collection, manager.Segment, segments.SegmentTypeGrowing, 0, &querypb.SegmentLoadInfo{
|
||||
SegmentID: segmentID,
|
||||
PartitionID: data.PartitionID,
|
||||
CollectionID: collection.ID(),
|
||||
InsertChannel: data.StartPosition.GetChannelName(),
|
||||
StartPosition: data.StartPosition,
|
||||
DeltaPosition: data.StartPosition,
|
||||
Level: datapb.SegmentLevel_L1,
|
||||
})
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer growing.Release(ctx)
|
||||
return growing.Insert(ctx, data.RowIDs, data.Timestamps, data.InsertRecord)
|
||||
}
|
||||
|
||||
func TestInsertNodeBlocksSchemaUpdateUntilGrowingInsertCompletes(t *testing.T) {
|
||||
manager, collection, schemaV950, schemaV951, insertMsg := setupSchemaTransitionInsertNodeTest(t)
|
||||
|
||||
converted := make(chan struct{})
|
||||
resumeConversion := make(chan struct{})
|
||||
var conversionOnce sync.Once
|
||||
var releaseOnce sync.Once
|
||||
releaseConversion := func() {
|
||||
releaseOnce.Do(func() {
|
||||
close(resumeConversion)
|
||||
})
|
||||
}
|
||||
var oldFieldConverted, readerHeldDuringConversion bool
|
||||
var transferOrigin func(*schemapb.CollectionSchema, *msgstream.InsertMsg) (*segcorepb.InsertRecord, []int64, error)
|
||||
transferMock := mockey.Mock(storage.TransferInsertMsgToInsertRecord).To(func(schema *schemapb.CollectionSchema, msg *msgstream.InsertMsg) (*segcorepb.InsertRecord, []int64, error) {
|
||||
record, skippedFields, err := transferOrigin(schema, msg)
|
||||
conversionOnce.Do(func() {
|
||||
oldFieldConverted = lo.ContainsBy(record.GetFieldsData(), func(field *schemapb.FieldData) bool {
|
||||
return field.GetFieldId() == schemaTransitionDroppedField
|
||||
})
|
||||
readerHeldDuringConversion = collection.HasInsertSchemaTransitionReaderForTest()
|
||||
close(converted)
|
||||
<-resumeConversion
|
||||
})
|
||||
return record, skippedFields, err
|
||||
}).Origin(&transferOrigin).Build()
|
||||
t.Cleanup(func() {
|
||||
transferMock.UnPatch()
|
||||
})
|
||||
|
||||
var (
|
||||
insertErr error
|
||||
nativeInsertAttempted bool
|
||||
oldFieldPassedToNative bool
|
||||
)
|
||||
mockDelegator := delegator.NewMockShardDelegator(t)
|
||||
mockDelegator.EXPECT().ProcessInsert(mock.Anything).Run(func(insertRecords map[int64]*delegator.InsertData) {
|
||||
insertData, ok := insertRecords[schemaTransitionSegmentID]
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
nativeInsertAttempted = true
|
||||
oldFieldPassedToNative = lo.ContainsBy(insertData.InsertRecord.GetFieldsData(), func(field *schemapb.FieldData) bool {
|
||||
return field.GetFieldId() == schemaTransitionDroppedField
|
||||
})
|
||||
insertErr = insertIntoNewGrowingSegment(collection, manager, schemaTransitionSegmentID, insertData)
|
||||
})
|
||||
|
||||
node, err := newInsertNode(schemaTransitionCollectionID, insertMsg.GetShardName(), manager, mockDelegator, schemaV950, 8)
|
||||
require.NoError(t, err)
|
||||
|
||||
operateDone := make(chan struct{})
|
||||
updateDone := make(chan struct{})
|
||||
updateErr := make(chan error, 1)
|
||||
operateStarted := false
|
||||
updateStarted := false
|
||||
t.Cleanup(func() {
|
||||
releaseConversion()
|
||||
if operateStarted {
|
||||
select {
|
||||
case <-operateDone:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Error("insert did not stop during test cleanup")
|
||||
}
|
||||
}
|
||||
if updateStarted {
|
||||
select {
|
||||
case <-updateDone:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Error("schema update did not stop during test cleanup")
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
operateStarted = true
|
||||
go func() {
|
||||
defer close(operateDone)
|
||||
node.Operate(&insertNodeMsg{insertMsgs: []*InsertMsg{insertMsg}})
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-converted:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("insert did not complete old-schema payload conversion")
|
||||
}
|
||||
require.True(t, oldFieldConverted, "the paused payload must retain the old-schema field to exercise the native race")
|
||||
require.True(t, readerHeldDuringConversion, "payload conversion must run inside the schema transition reader")
|
||||
|
||||
updateStarted = true
|
||||
go func() {
|
||||
defer close(updateDone)
|
||||
updateErr <- manager.Collection.UpdateSchema(schemaTransitionCollectionID, schemaV951, 951)
|
||||
}()
|
||||
require.True(t, collection.WaitForSchemaTransitionWriterForTest(5*time.Second), "schema writer did not queue behind old-schema payload conversion")
|
||||
|
||||
releaseConversion()
|
||||
|
||||
select {
|
||||
case <-operateDone:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("insert did not complete after payload conversion resumed")
|
||||
}
|
||||
require.True(t, nativeInsertAttempted, "old-schema payload did not reach native growing insertion")
|
||||
require.True(t, oldFieldPassedToNative, "native growing insertion did not receive the old-schema field")
|
||||
require.NoError(t, insertErr)
|
||||
|
||||
select {
|
||||
case <-updateDone:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("schema update did not continue after growing insert completed")
|
||||
}
|
||||
require.NoError(t, <-updateErr)
|
||||
}
|
||||
|
||||
func TestInsertNodeSkipsDroppedFieldAfterSchemaUpdate(t *testing.T) {
|
||||
manager, collection, schemaV950, schemaV951, insertMsg := setupSchemaTransitionInsertNodeTest(t)
|
||||
require.NoError(t, manager.Collection.UpdateSchema(schemaTransitionCollectionID, schemaV951, 951))
|
||||
|
||||
var (
|
||||
insertErr error
|
||||
droppedFieldSeen bool
|
||||
)
|
||||
mockDelegator := delegator.NewMockShardDelegator(t)
|
||||
mockDelegator.EXPECT().ProcessInsert(mock.Anything).Run(func(insertRecords map[int64]*delegator.InsertData) {
|
||||
for segmentID, insertData := range insertRecords {
|
||||
droppedFieldSeen = lo.ContainsBy(insertData.InsertRecord.GetFieldsData(), func(field *schemapb.FieldData) bool {
|
||||
return field.GetFieldId() == schemaTransitionDroppedField
|
||||
})
|
||||
insertErr = insertIntoNewGrowingSegment(collection, manager, segmentID, insertData)
|
||||
}
|
||||
})
|
||||
|
||||
node, err := newInsertNode(schemaTransitionCollectionID, insertMsg.GetShardName(), manager, mockDelegator, schemaV950, 8)
|
||||
require.NoError(t, err)
|
||||
node.Operate(&insertNodeMsg{insertMsgs: []*InsertMsg{insertMsg}})
|
||||
|
||||
require.False(t, droppedFieldSeen)
|
||||
require.NoError(t, insertErr)
|
||||
}
|
||||
|
||||
@@ -103,40 +103,38 @@ func (m *collectionManager) Get(collectionID int64) *Collection {
|
||||
return m.collections[collectionID]
|
||||
}
|
||||
|
||||
// acquireCollectionLease keeps a collection alive after the manager lock is
|
||||
// released. It intentionally bypasses Collection.Ref because a temporary
|
||||
// lease must not refresh storage context or become an externally visible ref.
|
||||
func (m *collectionManager) acquireCollectionLease(collectionID int64) (*Collection, bool) {
|
||||
m.mut.RLock()
|
||||
defer m.mut.RUnlock()
|
||||
|
||||
collection, ok := m.collections[collectionID]
|
||||
if ok {
|
||||
collection.refCount.Inc()
|
||||
}
|
||||
return collection, ok
|
||||
}
|
||||
|
||||
func (m *collectionManager) PutOrRef(collectionID int64, schema *schemapb.CollectionSchema, meta *segcorepb.CollectionIndexMeta, loadMeta *querypb.LoadMetaInfo) error {
|
||||
m.mut.Lock()
|
||||
defer m.mut.Unlock()
|
||||
logicalSchemaVersion := getLoadMetaSchemaVersion(schema, loadMeta)
|
||||
schemaBarrierTs := loadMeta.GetSchemaBarrierTs()
|
||||
if collection, ok := m.collections[collectionID]; ok {
|
||||
// Existing collections may be reached by a later load result or by a
|
||||
// same-version properties refresh. Keep the Go-side logical schema version
|
||||
// separate from the barrier timestamp so stale schema payloads cannot roll
|
||||
// back fields, while newer properties-only payloads can still refresh.
|
||||
if plan, shouldUpdate := prepareCollectionSchemaUpdate(collection, logicalSchemaVersion, schemaBarrierTs); shouldUpdate {
|
||||
if err := collection.updateSchema(schema, plan.segcoreSchemaVersion); err != nil {
|
||||
return err
|
||||
}
|
||||
collection.setSchema(schema, plan.logicalSchemaVersion, plan.schemaBarrierTs, plan.segcoreSchemaVersion)
|
||||
mlog.Info(context.TODO(), "update collection schema",
|
||||
mlog.Int64("collectionID", collectionID),
|
||||
mlog.Uint64("schemaVersion", plan.logicalSchemaVersion),
|
||||
mlog.Uint64("schemaBarrierTs", plan.schemaBarrierTs),
|
||||
mlog.Uint64("segcoreSchemaVersion", plan.segcoreSchemaVersion),
|
||||
mlog.Any("schema", schema),
|
||||
)
|
||||
}
|
||||
// Always update index meta to ensure newly indexed fields are visible
|
||||
// for search plan creation (CollectionIndexMeta::HasField check).
|
||||
if meta != nil {
|
||||
if err := collection.updateIndexMeta(meta); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
collection.Ref(1)
|
||||
return nil
|
||||
|
||||
if collection, ok := m.acquireCollectionLease(collectionID); ok {
|
||||
defer m.Unref(collectionID, 1)
|
||||
return m.putOrRefExisting(collectionID, collection, schema, meta, logicalSchemaVersion, schemaBarrierTs)
|
||||
}
|
||||
|
||||
m.mut.Lock()
|
||||
if collection, ok := m.collections[collectionID]; ok {
|
||||
collection.refCount.Inc()
|
||||
m.mut.Unlock()
|
||||
defer m.Unref(collectionID, 1)
|
||||
return m.putOrRefExisting(collectionID, collection, schema, meta, logicalSchemaVersion, schemaBarrierTs)
|
||||
}
|
||||
defer m.mut.Unlock()
|
||||
|
||||
mlog.Info(context.TODO(), "put new collection", mlog.Int64("collectionID", collectionID), mlog.Any("schema", schema))
|
||||
collection, err := NewCollection(collectionID, schema, meta, loadMeta)
|
||||
mlog.Info(context.TODO(), "new collection created", mlog.Int64("collectionID", collectionID), mlog.Any("schema", schema), mlog.Err(err))
|
||||
@@ -150,14 +148,33 @@ func (m *collectionManager) PutOrRef(collectionID int64, schema *schemapb.Collec
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *collectionManager) UpdateSchema(collectionID int64, schema *schemapb.CollectionSchema, schemaBarrierTs uint64) error {
|
||||
m.mut.Lock()
|
||||
defer m.mut.Unlock()
|
||||
func (m *collectionManager) putOrRefExisting(collectionID int64, collection *Collection, schema *schemapb.CollectionSchema, meta *segcorepb.CollectionIndexMeta, logicalSchemaVersion uint64, schemaBarrierTs uint64) error {
|
||||
// Existing collections may be reached by a later load result or by a
|
||||
// same-version properties refresh. Keep the Go-side logical schema version
|
||||
// separate from the barrier timestamp so stale schema payloads cannot roll
|
||||
// back fields, while newer properties-only payloads can still refresh.
|
||||
plan, shouldUpdate, err := collection.applyLoadUpdate(schema, meta, logicalSchemaVersion, schemaBarrierTs)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if shouldUpdate {
|
||||
mlog.Info(context.TODO(), "update collection schema",
|
||||
mlog.Int64("collectionID", collectionID),
|
||||
mlog.Uint64("schemaVersion", plan.logicalSchemaVersion),
|
||||
mlog.Uint64("schemaBarrierTs", plan.schemaBarrierTs),
|
||||
mlog.Uint64("segcoreSchemaVersion", plan.segcoreSchemaVersion),
|
||||
mlog.Any("schema", schema),
|
||||
)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
collection, ok := m.collections[collectionID]
|
||||
func (m *collectionManager) UpdateSchema(collectionID int64, schema *schemapb.CollectionSchema, schemaBarrierTs uint64) error {
|
||||
collection, ok := m.acquireCollectionLease(collectionID)
|
||||
if !ok {
|
||||
return merr.WrapErrCollectionNotFound(collectionID, "collection not found in querynode collection manager")
|
||||
}
|
||||
defer m.Unref(collectionID, 1)
|
||||
|
||||
logicalSchemaVersion := getUpdateSchemaVersion(schema, schemaBarrierTs)
|
||||
// A schema update carries two ordering domains:
|
||||
@@ -165,16 +182,8 @@ func (m *collectionManager) UpdateSchema(collectionID int64, schema *schemapb.Co
|
||||
// older schema payloads from overwriting newer fields/functions.
|
||||
// - schemaBarrierTs is the DDL barrier timestamp and advances for
|
||||
// properties-only schema snapshots such as ttl_field changes.
|
||||
plan, shouldUpdate := prepareCollectionSchemaUpdate(collection, logicalSchemaVersion, schemaBarrierTs)
|
||||
if !shouldUpdate {
|
||||
return nil
|
||||
}
|
||||
|
||||
if err := collection.updateSchema(schema, plan.segcoreSchemaVersion); err != nil {
|
||||
return err
|
||||
}
|
||||
collection.setSchema(schema, plan.logicalSchemaVersion, plan.schemaBarrierTs, plan.segcoreSchemaVersion)
|
||||
return nil
|
||||
_, _, err := collection.applySchemaUpdate(schema, logicalSchemaVersion, schemaBarrierTs)
|
||||
return err
|
||||
}
|
||||
|
||||
// ShouldUpdateCollectionSchema reports whether an UpdateSchema payload would
|
||||
@@ -308,14 +317,15 @@ type collectionSchemaSnapshot struct {
|
||||
// Collection is a wrapper of the underlying C-structure C.CCollection
|
||||
// In a query node, `Collection` is a replica info of a collection in these query node.
|
||||
type Collection struct {
|
||||
mu sync.RWMutex // protects colllectionPtr
|
||||
ccollection *segcore.CCollection
|
||||
id int64
|
||||
partitions *typeutil.ConcurrentSet[int64]
|
||||
loadType querypb.LoadType
|
||||
dbName string
|
||||
dbProperties []*commonpb.KeyValuePair
|
||||
resourceGroup string
|
||||
mu sync.RWMutex // protects colllectionPtr
|
||||
schemaTransitionMu sync.RWMutex // serializes schema transitions with insert payload conversion and growing writes
|
||||
ccollection *segcore.CCollection
|
||||
id int64
|
||||
partitions *typeutil.ConcurrentSet[int64]
|
||||
loadType querypb.LoadType
|
||||
dbName string
|
||||
dbProperties []*commonpb.KeyValuePair
|
||||
resourceGroup string
|
||||
// resource group of node may be changed if node transfer,
|
||||
// but Collection in Manager will be released before assign new replica of new resource group on these node.
|
||||
// so we don't need to update resource group in Collection.
|
||||
@@ -418,6 +428,63 @@ func (c *Collection) updateSchema(schema *schemapb.CollectionSchema, version uin
|
||||
return c.ccollection.UpdateSchema(schema, version)
|
||||
}
|
||||
|
||||
func (c *Collection) applySchemaUpdate(schema *schemapb.CollectionSchema, logicalSchemaVersion uint64, schemaBarrierTs uint64) (collectionSchemaUpdatePlan, bool, error) {
|
||||
c.lockSchemaTransitionForUpdate()
|
||||
defer c.unlockSchemaTransitionForUpdate()
|
||||
|
||||
return c.applySchemaUpdateLocked(schema, logicalSchemaVersion, schemaBarrierTs)
|
||||
}
|
||||
|
||||
func (c *Collection) applyLoadUpdate(schema *schemapb.CollectionSchema, meta *segcorepb.CollectionIndexMeta, logicalSchemaVersion uint64, schemaBarrierTs uint64) (collectionSchemaUpdatePlan, bool, error) {
|
||||
c.lockSchemaTransitionForUpdate()
|
||||
defer c.unlockSchemaTransitionForUpdate()
|
||||
|
||||
plan, shouldUpdate, err := c.applySchemaUpdateLocked(schema, logicalSchemaVersion, schemaBarrierTs)
|
||||
if err != nil {
|
||||
return collectionSchemaUpdatePlan{}, false, err
|
||||
}
|
||||
// Always update index meta to ensure newly indexed fields are visible
|
||||
// for search plan creation (CollectionIndexMeta::HasField check).
|
||||
if err := c.updateIndexMeta(meta); err != nil {
|
||||
return collectionSchemaUpdatePlan{}, false, err
|
||||
}
|
||||
// The temporary manager lease keeps the collection alive while this update
|
||||
// waits. Publish the caller-visible ref only after the schema and index meta
|
||||
// that determine its storage context are applied.
|
||||
c.Ref(1)
|
||||
return plan, shouldUpdate, nil
|
||||
}
|
||||
|
||||
func (c *Collection) applySchemaUpdateLocked(schema *schemapb.CollectionSchema, logicalSchemaVersion uint64, schemaBarrierTs uint64) (collectionSchemaUpdatePlan, bool, error) {
|
||||
plan, shouldUpdate := prepareCollectionSchemaUpdate(c, logicalSchemaVersion, schemaBarrierTs)
|
||||
if !shouldUpdate {
|
||||
return collectionSchemaUpdatePlan{}, false, nil
|
||||
}
|
||||
if err := c.updateSchema(schema, plan.segcoreSchemaVersion); err != nil {
|
||||
return collectionSchemaUpdatePlan{}, false, err
|
||||
}
|
||||
c.setSchema(schema, plan.logicalSchemaVersion, plan.schemaBarrierTs, plan.segcoreSchemaVersion)
|
||||
return plan, true, nil
|
||||
}
|
||||
|
||||
func (c *Collection) lockSchemaTransitionForUpdate() {
|
||||
c.schemaTransitionMu.Lock()
|
||||
}
|
||||
|
||||
func (c *Collection) unlockSchemaTransitionForUpdate() {
|
||||
c.schemaTransitionMu.Unlock()
|
||||
}
|
||||
|
||||
// WithInsertSchemaTransition keeps payload conversion and growing writes in
|
||||
// one schema epoch. A schema update cannot change the native collection until
|
||||
// fn returns.
|
||||
func (c *Collection) WithInsertSchemaTransition(fn func(schema *schemapb.CollectionSchema)) {
|
||||
c.schemaTransitionMu.RLock()
|
||||
defer c.schemaTransitionMu.RUnlock()
|
||||
|
||||
fn(c.Schema())
|
||||
}
|
||||
|
||||
func (c *Collection) setSchema(schema *schemapb.CollectionSchema, logicalSchemaVersion uint64, schemaBarrierTs uint64, segcoreSchemaVersion uint64) {
|
||||
c.schema.Store(&collectionSchemaSnapshot{
|
||||
schema: schema,
|
||||
|
||||
@@ -18,10 +18,12 @@ package segments
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"runtime"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/bytedance/mockey"
|
||||
"github.com/stretchr/testify/suite"
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
@@ -406,6 +408,242 @@ func (s *CollectionManagerSuite) TestPutOrRefUpdateIndexMetaWaitsForCollectionNa
|
||||
s.cm.Unref(1, 1)
|
||||
}
|
||||
|
||||
func holdInsertSchemaTransition(t *testing.T, collection *Collection) func() {
|
||||
t.Helper()
|
||||
entered := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
done := make(chan struct{})
|
||||
var releaseOnce sync.Once
|
||||
|
||||
go func() {
|
||||
defer close(done)
|
||||
collection.WithInsertSchemaTransition(func(*schemapb.CollectionSchema) {
|
||||
close(entered)
|
||||
<-release
|
||||
})
|
||||
}()
|
||||
|
||||
select {
|
||||
case <-entered:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("insert schema transition reader did not start")
|
||||
}
|
||||
|
||||
return func() {
|
||||
releaseOnce.Do(func() {
|
||||
close(release)
|
||||
})
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("insert schema transition reader did not stop")
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func waitForSchemaTransitionWriter(t *testing.T, collection *Collection) {
|
||||
t.Helper()
|
||||
deadline := time.NewTimer(5 * time.Second)
|
||||
defer deadline.Stop()
|
||||
|
||||
for {
|
||||
if !collection.schemaTransitionMu.TryRLock() {
|
||||
return
|
||||
}
|
||||
collection.schemaTransitionMu.RUnlock()
|
||||
|
||||
select {
|
||||
case <-deadline.C:
|
||||
t.Fatal("schema writer did not queue behind the insert transition reader")
|
||||
default:
|
||||
runtime.Gosched()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func mockNativeSchemaUpdate(t *testing.T) <-chan struct{} {
|
||||
t.Helper()
|
||||
entered := make(chan struct{})
|
||||
var once sync.Once
|
||||
var origin func(*segcore.CCollection, *schemapb.CollectionSchema, uint64) error
|
||||
mock := mockey.Mock((*segcore.CCollection).UpdateSchema).To(func(c *segcore.CCollection, schema *schemapb.CollectionSchema, version uint64) error {
|
||||
once.Do(func() {
|
||||
close(entered)
|
||||
})
|
||||
return origin(c, schema, version)
|
||||
}).Origin(&origin).Build()
|
||||
t.Cleanup(func() {
|
||||
mock.UnPatch()
|
||||
})
|
||||
return entered
|
||||
}
|
||||
|
||||
func (s *CollectionManagerSuite) assertNativeSchemaUpdateWaitsForTransitionReader(collection *Collection, update func() error) {
|
||||
releaseReader := holdInsertSchemaTransition(s.T(), collection)
|
||||
defer releaseReader()
|
||||
|
||||
nativeEntered := mockNativeSchemaUpdate(s.T())
|
||||
updateDone := make(chan error, 1)
|
||||
go func() {
|
||||
updateDone <- update()
|
||||
}()
|
||||
|
||||
waitForSchemaTransitionWriter(s.T(), collection)
|
||||
|
||||
select {
|
||||
case <-nativeEntered:
|
||||
s.T().Fatal("native schema update entered while an insert transition reader was held")
|
||||
default:
|
||||
}
|
||||
|
||||
releaseReader()
|
||||
select {
|
||||
case <-nativeEntered:
|
||||
case <-time.After(5 * time.Second):
|
||||
s.T().Fatal("native schema update did not continue after the transition reader was released")
|
||||
}
|
||||
s.Require().NoError(<-updateDone)
|
||||
}
|
||||
|
||||
func (s *CollectionManagerSuite) TestSchemaUpdateWaitsForTransitionReader() {
|
||||
s.Run("UpdateSchema", func() {
|
||||
coll := s.cm.Get(1)
|
||||
s.Require().NotNil(coll)
|
||||
|
||||
schema := proto.Clone(coll.Schema()).(*schemapb.CollectionSchema)
|
||||
schema.Version++
|
||||
schema.Fields = append(schema.Fields, &schemapb.FieldSchema{
|
||||
FieldID: 580,
|
||||
Name: "schema_transition_update",
|
||||
DataType: schemapb.DataType_Bool,
|
||||
Nullable: true,
|
||||
})
|
||||
|
||||
s.assertNativeSchemaUpdateWaitsForTransitionReader(coll, func() error {
|
||||
return s.cm.UpdateSchema(coll.ID(), schema, 1)
|
||||
})
|
||||
})
|
||||
|
||||
s.Run("PutOrRef", func() {
|
||||
coll := s.cm.Get(1)
|
||||
s.Require().NotNil(coll)
|
||||
|
||||
schema := proto.Clone(coll.Schema()).(*schemapb.CollectionSchema)
|
||||
schema.Version++
|
||||
schema.Fields = append(schema.Fields, &schemapb.FieldSchema{
|
||||
FieldID: 581,
|
||||
Name: "schema_transition_put_or_ref",
|
||||
DataType: schemapb.DataType_Bool,
|
||||
Nullable: true,
|
||||
})
|
||||
|
||||
s.assertNativeSchemaUpdateWaitsForTransitionReader(coll, func() error {
|
||||
return s.cm.PutOrRef(coll.ID(), schema, nil, &querypb.LoadMetaInfo{
|
||||
CollectionID: coll.ID(),
|
||||
LoadType: querypb.LoadType_LoadCollection,
|
||||
SchemaBarrierTs: 2,
|
||||
})
|
||||
})
|
||||
s.cm.Unref(coll.ID(), 1)
|
||||
})
|
||||
}
|
||||
|
||||
func (s *CollectionManagerSuite) TestSchemaUpdateDoesNotBlockUnrelatedCollectionGet() {
|
||||
otherSchema := mock_segcore.GenTestCollectionSchema("other_collection", schemapb.DataType_Int64, false)
|
||||
s.Require().NoError(s.cm.PutOrRef(2, otherSchema, nil, &querypb.LoadMetaInfo{
|
||||
CollectionID: 2,
|
||||
LoadType: querypb.LoadType_LoadCollection,
|
||||
}))
|
||||
defer s.cm.Unref(2, 1)
|
||||
|
||||
for _, test := range []struct {
|
||||
name string
|
||||
update func(collection *Collection, schema *schemapb.CollectionSchema) error
|
||||
unref bool
|
||||
}{
|
||||
{
|
||||
name: "UpdateSchema",
|
||||
update: func(collection *Collection, schema *schemapb.CollectionSchema) error {
|
||||
return s.cm.UpdateSchema(collection.ID(), schema, 1)
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "PutOrRef",
|
||||
update: func(collection *Collection, schema *schemapb.CollectionSchema) error {
|
||||
return s.cm.PutOrRef(collection.ID(), schema, nil, &querypb.LoadMetaInfo{
|
||||
CollectionID: collection.ID(),
|
||||
LoadType: querypb.LoadType_LoadCollection,
|
||||
SchemaBarrierTs: 1,
|
||||
})
|
||||
},
|
||||
unref: true,
|
||||
},
|
||||
} {
|
||||
s.Run(test.name, func() {
|
||||
coll := s.cm.Get(1)
|
||||
schema := proto.Clone(coll.Schema()).(*schemapb.CollectionSchema)
|
||||
schema.Version++
|
||||
|
||||
releaseReader := holdInsertSchemaTransition(s.T(), coll)
|
||||
defer releaseReader()
|
||||
updateDone := make(chan error, 1)
|
||||
go func() {
|
||||
updateDone <- test.update(coll, schema)
|
||||
}()
|
||||
|
||||
waitForSchemaTransitionWriter(s.T(), coll)
|
||||
|
||||
getDone := make(chan *Collection, 1)
|
||||
go func() {
|
||||
getDone <- s.cm.Get(2)
|
||||
}()
|
||||
select {
|
||||
case other := <-getDone:
|
||||
s.Require().NotNil(other)
|
||||
case <-time.After(5 * time.Second):
|
||||
s.T().Fatal("schema update on one collection blocked Get on another collection")
|
||||
}
|
||||
|
||||
releaseReader()
|
||||
s.Require().NoError(<-updateDone)
|
||||
if test.unref {
|
||||
s.cm.Unref(coll.ID(), 1)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func (s *CollectionManagerSuite) TestSchemaUpdateLeaseKeepsCollectionAliveWhileWaiting() {
|
||||
coll := s.cm.Get(1)
|
||||
schema := proto.Clone(coll.Schema()).(*schemapb.CollectionSchema)
|
||||
schema.Version++
|
||||
|
||||
releaseReader := holdInsertSchemaTransition(s.T(), coll)
|
||||
defer releaseReader()
|
||||
updateDone := make(chan error, 1)
|
||||
go func() {
|
||||
updateDone <- s.cm.UpdateSchema(coll.ID(), schema, 1)
|
||||
}()
|
||||
|
||||
waitForSchemaTransitionWriter(s.T(), coll)
|
||||
|
||||
unrefDone := make(chan bool, 1)
|
||||
go func() {
|
||||
unrefDone <- s.cm.Unref(coll.ID(), 1)
|
||||
}()
|
||||
select {
|
||||
case released := <-unrefDone:
|
||||
s.False(released, "the update lease must retain the collection")
|
||||
case <-time.After(5 * time.Second):
|
||||
s.T().Fatal("Unref blocked while schema update waited for transition reader")
|
||||
}
|
||||
s.Same(coll, s.cm.Get(coll.ID()))
|
||||
|
||||
releaseReader()
|
||||
s.Require().NoError(<-updateDone)
|
||||
s.Nil(s.cm.Get(coll.ID()), "releasing the lease should complete the pending collection release")
|
||||
}
|
||||
|
||||
func (s *CollectionManagerSuite) TestCollectionNativeWrapperMethods() {
|
||||
coll := s.cm.Get(1)
|
||||
s.Require().NotNil(coll)
|
||||
|
||||
@@ -0,0 +1,53 @@
|
||||
// Licensed to the LF AI & Data foundation under one
|
||||
// or more contributor license agreements. See the NOTICE file
|
||||
// distributed with this work for additional information
|
||||
// regarding copyright ownership. The ASF licenses this file
|
||||
// to you under the Apache License, Version 2.0 (the
|
||||
// "License"); you may not use this file except in compliance
|
||||
// with the License. You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
//go:build test
|
||||
|
||||
package segments
|
||||
|
||||
import (
|
||||
"runtime"
|
||||
"time"
|
||||
)
|
||||
|
||||
// WaitForSchemaTransitionWriterForTest waits until a schema writer is queued
|
||||
// behind an active insert transition reader. It is available only in test
|
||||
// builds so cross-package regression tests can synchronize on the writer
|
||||
// queue instead of using a timer as the success condition.
|
||||
func (c *Collection) WaitForSchemaTransitionWriterForTest(timeout time.Duration) bool {
|
||||
deadline := time.Now().Add(timeout)
|
||||
for {
|
||||
if !c.schemaTransitionMu.TryRLock() {
|
||||
return true
|
||||
}
|
||||
c.schemaTransitionMu.RUnlock()
|
||||
if time.Now().After(deadline) {
|
||||
return false
|
||||
}
|
||||
runtime.Gosched()
|
||||
}
|
||||
}
|
||||
|
||||
// HasInsertSchemaTransitionReaderForTest reports whether an insert reader is
|
||||
// holding the transition mutex. Callers must invoke it before starting a schema
|
||||
// writer, so a failed TryLock unambiguously means a reader is active.
|
||||
func (c *Collection) HasInsertSchemaTransitionReaderForTest() bool {
|
||||
if c.schemaTransitionMu.TryLock() {
|
||||
c.schemaTransitionMu.Unlock()
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
Reference in New Issue
Block a user