fix(collection): allow writes and reads to proceed during Optimize

src/db/collection.cc

@@ -238,6 +238,8 @@ class CollectionImpl : public Collection {
 
   mutable std::shared_mutex schema_handle_mtx_;
   mutable std::shared_mutex write_mtx_;
+  // serializes concurrent Optimize calls
+  mutable std::mutex optimize_mtx_;
 
   std::atomic<SegmentID> segment_id_allocator_;
   std::atomic<SegmentID> tmp_segment_id_allocator_;
@@ -779,9 +781,14 @@ std::vector<SegmentTask::Ptr> CollectionImpl::build_drop_scalar_index_task(
 Status CollectionImpl::Optimize(const OptimizeOptions &options) {
   CHECK_COLLECTION_READONLY_RETURN_STATUS;
 
-  std::lock_guard lock(schema_handle_mtx_);
-  // when optimizing, schema operations(include another optimize) are not
-  // allowed
+  // serialize concurrent optimize calls
+  std::lock_guard optimize_lock(optimize_mtx_);
+
+  // Hold the schema lock in shared mode so that writes and reads (which
+  // also take it in shared mode) can proceed during the long-running
+  // compact, while schema operations and close/destroy (which take it in
+  // exclusive mode) are kept out for the whole optimize.
+  std::shared_lock lock(schema_handle_mtx_);
 
   CHECK_DESTROY_RETURN_STATUS(destroyed_, false);
 
@@ -876,38 +883,85 @@ Status CollectionImpl::Optimize(const OptimizeOptions &options) {
     s = version_manager_->flush();
     CHECK_RETURN_STATUS(s);
 
-    // 4. remove old segments or block
+    // 4. commit the new segment set atomically
+    //
+    // Never mutate live segments in place here: Optimize holds the schema
+    // lock in shared mode, so concurrent readers may still be using them.
+    // New segment instances are opened from the new metas and swapped in
+    // with a single atomic replacement, so concurrent readers observe
+    // either the pre-optimize or the post-optimize segment set, never a
+    // mix (e.g. merge inputs alongside their merged output). Replaced
+    // instances are only marked need-destroyed: in-flight readers keep
+    // them (and their files) alive via shared_ptr, and the files are
+    // removed from disk only when the last reference is dropped.
+    std::vector<Segment::Ptr> segments_to_add;
+    std::vector<SegmentID> segment_ids_to_destroy;
+    auto reopen_segment =
+        [&](const SegmentMeta::Ptr &meta) -> Result<Segment::Ptr> {
+      return Segment::Open(path_, *schema_, *meta, id_map_, delete_store_,
+                           version_manager_,
+                           SegmentOptions{true, options_.enable_mmap_});
+    };
     for (auto &task : tasks) {
       auto task_info = task->task_info();
 
       if (std::holds_alternative<CompactTask>(task_info)) {
         auto compact_task = std::get<CompactTask>(task_info);
 
         if (compact_task.output_segment_meta_) {
-          auto new_segment =
-              Segment::Open(path_, *schema_, *compact_task.output_segment_meta_,
-                            id_map_, delete_store_, version_manager_,
-                            SegmentOptions{true, options_.enable_mmap_});
+          auto new_segment = reopen_segment(compact_task.output_segment_meta_);
           if (!new_segment.has_value()) {
             return new_segment.error();
           }
-          s = segment_manager_->add_segment(new_segment.value());
-          CHECK_RETURN_STATUS(s);
+          segments_to_add.push_back(new_segment.value());
         }
 
         for (auto input_segment : compact_task.input_segments_) {
-          s = segment_manager_->destroy_segment(input_segment->id());
-          CHECK_RETURN_STATUS(s);
+          segment_ids_to_destroy.push_back(input_segment->id());
         }
       } else if (std::holds_alternative<CreateVectorIndexTask>(task_info)) {
         auto create_index_task = std::get<CreateVectorIndexTask>(task_info);
 
-        s = create_index_task.input_segment_->reload_vector_index(
-            *schema_, create_index_task.output_segment_meta_,
-            create_index_task.output_vector_indexers_,
-            create_index_task.output_quant_vector_indexers_);
-        CHECK_RETURN_STATUS(s);
+        // reopen the segment with the new vector index and replace the old
+        // instance by id, instead of reload_vector_index() which would
+        // destroy indexers still referenced by concurrent readers
+        auto new_segment =
+            reopen_segment(create_index_task.output_segment_meta_);
+        if (!new_segment.has_value()) {
+          return new_segment.error();
+        }
+        segments_to_add.push_back(new_segment.value());
+      }
+    }
+
+    s = segment_manager_->replace_segments(segments_to_add,
+                                           segment_ids_to_destroy);
+    CHECK_RETURN_STATUS(s);
+
+    // 5. schedule superseded index files for removal
+    //
+    // The reopened segments no longer reference the pre-index vector files
+    // of the replaced instances. Mark them so the files are removed once
+    // the last in-flight reader releases the old instance, instead of
+    // leaking on disk until the segment is compacted away.
+    for (auto &task : tasks) {
+      auto task_info = task->task_info();
+      if (!std::holds_alternative<CreateVectorIndexTask>(task_info)) {
+        continue;
+      }
+      auto create_index_task = std::get<CreateVectorIndexTask>(task_info);
+      std::vector<std::string> superseded_columns;
+      std::vector<std::string> superseded_quant_columns;
+      for (auto &[column, indexer] :
+           create_index_task.output_vector_indexers_) {
+        superseded_columns.push_back(column);
+      }
+      for (auto &[column, indexer] :
+           create_index_task.output_quant_vector_indexers_) {
+        superseded_quant_columns.push_back(column);
       }
+      create_index_task.input_segment_->mark_vector_index_files_for_removal(
+          superseded_columns, superseded_quant_columns);
     }
   }
 

src/db/index/column/vector_column/vector_column_indexer.cc

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 #include "vector_column_indexer.h"
+#include <zvec/ailego/logger/logger.h>
 #include <zvec/ailego/pattern/expected.hpp>
 #include <zvec/core/interface/index_factory.h>
 #include <zvec/db/status.h>
@@ -20,6 +21,25 @@
 
 namespace zvec {
 
+VectorColumnIndexer::~VectorColumnIndexer() {
+  if (!destroy_on_release_.load(std::memory_order_relaxed)) {
+    return;
+  }
+  // Deferred removal of a superseded index file: this runs only after the
+  // last reference is released, so no reader can still be using the index.
+  if (index != nullptr) {
+    if (0 != index->Close()) {
+      LOG_WARN("Failed to close superseded index: %s",
+               index_file_path_.c_str());
+    }
+    index.reset();
+  }
+  if (!ailego::File::RemovePath(index_file_path_)) {
+    LOG_WARN("Failed to remove superseded index file: %s",
+             index_file_path_.c_str());
+  }
+}
+
 Status VectorColumnIndexer::Open(
     const vector_column_params::ReadOptions &read_options) {
   if (index != nullptr) {

src/db/index/column/vector_column/vector_column_indexer.h

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 #pragma once
+#include <atomic>
 #include <string>
 #include <utility>
 #include <variant>
@@ -50,7 +51,7 @@ class VectorColumnIndexer {
     is_sparse_ = field_schema.is_sparse_vector();
   }
 
-  virtual ~VectorColumnIndexer() = default;
+  virtual ~VectorColumnIndexer();
 
  public:
   Status Open(const vector_column_params::ReadOptions &read_options);
@@ -63,6 +64,14 @@ class VectorColumnIndexer {
   // Destroy will call Close() and remove index file
   Status Destroy();
 
+  //! Mark this indexer as superseded: its index file is removed when the
+  //! last shared_ptr reference is released (in the destructor), instead of
+  //! eagerly. Safe to call while concurrent readers still hold references;
+  //! they keep the index and its file alive until they finish.
+  void MarkDestroyOnRelease() {
+    destroy_on_release_.store(true, std::memory_order_relaxed);
+  }
+
 
   // If HNSWIndexer.merge([FlatIndexer1, FlatIndexer2])
   // then the merged indexer is a HNSWIndexer
@@ -138,6 +147,7 @@ class VectorColumnIndexer {
   std::string engine_name_ = "proxima";
   bool is_sparse_{false};  // TODO: eliminate the dynamic flag and make it
                            // static/template/seperate class
+  std::atomic<bool> destroy_on_release_{false};
 };
 
 

src/db/index/segment/segment.cc

@@ -194,6 +194,10 @@ class SegmentImpl : public Segment,
       const std::unordered_map<std::string, VectorColumnIndexer::Ptr>
           &quant_vector_indexers) override;
 
+  void mark_vector_index_files_for_removal(
+      const std::vector<std::string> &columns,
+      const std::vector<std::string> &quant_columns) override;
+
   bool vector_index_ready(const std::string &column,
                           const IndexParams::Ptr &index_params) const override;
 
@@ -1939,6 +1943,35 @@ Status SegmentImpl::reload_vector_index(
   return Status::OK();
 }
 
+void SegmentImpl::mark_vector_index_files_for_removal(
+    const std::vector<std::string> &columns,
+    const std::vector<std::string> &quant_columns) {
+  // Only reads the indexer maps and flips an atomic flag on the indexer
+  // objects; no structural mutation, so this is safe while concurrent
+  // readers of this (already superseded) instance are still running.
+  //
+  // The base and quantized indexers are marked independently: when a
+  // quantized index is built on top of a reused base file (see
+  // create_vector_index), only the old quantized file is superseded and
+  // the base file is still referenced by the new segment meta.
+  for (const auto &column : columns) {
+    auto iter = vector_indexers_.find(column);
+    if (iter != vector_indexers_.end()) {
+      for (auto &indexer : iter->second) {
+        indexer->MarkDestroyOnRelease();
+      }
+    }
+  }
+  for (const auto &column : quant_columns) {
+    auto q_iter = quant_vector_indexers_.find(column);
+    if (q_iter != quant_vector_indexers_.end()) {
+      for (auto &q_indexer : q_iter->second) {
+        q_indexer->MarkDestroyOnRelease();
+      }
+    }
+  }
+}
+
 bool SegmentImpl::vector_index_ready(
     const std::string &column, const IndexParams::Ptr &index_params) const {
   auto field = collection_schema_->get_vector_field(column);

src/db/index/segment/segment.h

@@ -106,6 +106,20 @@ class Segment {
       const std::unordered_map<std::string, VectorColumnIndexer::Ptr>
           &quant_vector_indexers = {}) = 0;
 
+  //! Mark vector indexers of this segment so that their index files are
+  //! removed once the last reference to them is released. Used after this
+  //! segment instance has been superseded by a reopened one (e.g. by
+  //! Optimize building a new vector index): in-flight readers of this
+  //! instance keep the files alive until they finish.
+  //!
+  //! columns: columns whose base vector index was replaced by a new file.
+  //! quant_columns: columns whose quantized index was replaced. A column
+  //! may appear only here when the new index reuses the old base file
+  //! (see create_vector_index), in which case the base must NOT be marked.
+  virtual void mark_vector_index_files_for_removal(
+      const std::vector<std::string> & /*columns*/,
+      const std::vector<std::string> & /*quant_columns*/) {}
+
   virtual bool vector_index_ready(
       const std::string &column,
       const IndexParams::Ptr &index_params) const = 0;