fix: segcore collection schema update not concurrent safe. (#45337)

relate: #45345

Signed-off-by: aoiasd <[email protected]>

Author: aoiasd

internal/core/src/common/Schema.h

@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include <atomic>
 #include <cstdint>
 #include <memory>
 #include <optional>
@@ -358,5 +359,5 @@ class Schema {
 };
 
 using SchemaPtr = std::shared_ptr<Schema>;
-
+using SafeSchemaPtr = std::atomic<SchemaPtr*>;
 }  // namespace milvus

internal/core/src/segcore/Collection.cpp

@@ -70,7 +70,7 @@ Collection::parse_schema(const void* schema_proto_blob,
                          const uint64_t version) {
     Assert(schema_proto_blob != nullptr);
 
-    if (version <= schema_->get_schema_version()) {
+    if (version <= get_schema_version()) {
         return;
     }
 
@@ -79,14 +79,9 @@ Collection::parse_schema(const void* schema_proto_blob,
 
     AssertInfo(suc, "parse schema proto failed");
 
-    auto old_schema = schema_;
-
-    schema_ = Schema::ParseFrom(collection_schema);
-    schema_->set_schema_version(version);
-
-    if (old_schema) {
-        schema_->UpdateLoadFields(old_schema->load_fields());
-    }
+    auto new_schema = Schema::ParseFrom(collection_schema);
+    new_schema->set_schema_version(version);
+    set_schema(new_schema);
 }
 
 }  // namespace milvus::segcore

internal/core/src/segcore/Collection.h

@@ -12,6 +12,7 @@
 #pragma once
 
 #include <memory>
+#include <shared_mutex>
 #include <string>
 
 #include "common/Schema.h"
@@ -36,9 +37,29 @@ class Collection {
  public:
     SchemaPtr
     get_schema() {
+        std::shared_lock lock(schema_mutex_);
         return schema_;
     }
 
+    uint64_t
+    get_schema_version() {
+        std::shared_lock lock(schema_mutex_);
+        return schema_->get_schema_version();
+    }
+
+    void
+    set_schema(SchemaPtr& new_schema) {
+        std::unique_lock lock(schema_mutex_);
+        auto old_schema = schema_;
+        if (new_schema->get_schema_version() > schema_->get_schema_version()) {
+            schema_ = new_schema;
+        }
+
+        if (old_schema) {
+            schema_->UpdateLoadFields(old_schema->load_fields());
+        }
+    }
+
     IndexMetaPtr&
     get_index_meta() {
         return index_meta_;
@@ -57,6 +78,7 @@ class Collection {
  private:
     std::string collection_name_;
     SchemaPtr schema_;
+    std::shared_mutex schema_mutex_;
     IndexMetaPtr index_meta_;
 };