fix: Change schema to atomic.Pointer to avoid data race (#28739)

Fix #28738

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

internal/querynodev2/segments/collection.go

@@ -73,7 +73,7 @@ func (m *collectionManager) PutOrRef(collectionID int64, schema *schemapb.Collec
 
 	if collection, ok := m.collections[collectionID]; ok {
 		// the schema may be changed even the collection is loaded
-		collection.schema = schema
+		collection.schema.Store(schema)
 		collection.Ref(1)
 		return
 	}
@@ -122,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
 }
@@ -134,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
@@ -214,14 +214,16 @@ func NewCollection(collectionID int64, schema *schemapb.CollectionSchema, indexM
 		C.SetIndexMeta(collection, unsafe.Pointer(&indexMetaBlob[0]), (C.int64_t)(len(indexMetaBlob)))
 	}
 
-	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 {

internal/querynodev2/segments/mock_data.go

@@ -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
 	}