fix: [cherry-pick] Change schema to atomic.Pointer to avoid data race (#28739) (#28759)

Cherry-pick from master
pr: #28739
See also #28738

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
This commit is contained in:
congqixia 2023-11-27 19:50:27 +08:00 committed by GitHub
parent 9378f78218
commit 5a962a631a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 9 additions and 5 deletions

View File

@ -72,6 +72,8 @@ func (m *collectionManager) PutOrRef(collectionID int64, schema *schemapb.Collec
defer m.mut.Unlock()
if collection, ok := m.collections[collectionID]; ok {
// the schema may be changed even the collection is loaded
collection.schema.Store(schema)
collection.Ref(1)
return
}
@ -120,7 +122,7 @@ type Collection struct {
partitions *typeutil.ConcurrentSet[int64]
loadType querypb.LoadType
metricType atomic.String
schema *schemapb.CollectionSchema
schema atomic.Pointer[schemapb.CollectionSchema]
refCount *atomic.Uint32
}
@ -132,7 +134,7 @@ func (c *Collection) ID() int64 {
// Schema returns the schema of collection
func (c *Collection) Schema() *schemapb.CollectionSchema {
return c.schema
return c.schema.Load()
}
// getPartitionIDs return partitionIDs of collection
@ -207,14 +209,16 @@ func NewCollection(collectionID int64, schema *schemapb.CollectionSchema, indexM
C.SetIndexMeta(collection, cIndexMetaBlob)
}
return &Collection{
coll := &Collection{
collectionPtr: collection,
id: collectionID,
schema: schema,
partitions: typeutil.NewConcurrentSet[int64](),
loadType: loadType,
refCount: atomic.NewUint32(0),
}
coll.schema.Store(schema)
return coll
}
func NewCollectionWithoutSchema(collectionID int64, loadType querypb.LoadType) *Collection {

View File

@ -1302,7 +1302,7 @@ func genSimpleRowIDField(numRows int) []int64 {
func genSimpleRetrievePlan(collection *Collection) (*RetrievePlan, error) {
timestamp := storage.Timestamp(1000)
planBytes, err := genSimpleRetrievePlanExpr(collection.schema)
planBytes, err := genSimpleRetrievePlanExpr(collection.schema.Load())
if err != nil {
return nil, err
}