From 63d26e02c7a618e57a85d84b4527ac4d78b14afa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=90=AA=E9=9C=84?= Date: Thu, 23 Jul 2026 17:16:29 +0800 Subject: [PATCH 1/3] fix(collection): allow writes and reads to proceed during Optimize Optimize() held schema_handle_mtx_ exclusively for its whole duration, so Insert/Query/Fetch/Delete (which take it in shared mode) were blocked until the long-running compact finished. - Downgrade Optimize's schema lock to shared mode; writes and reads now proceed during the compact and only wait on the short write_mtx_ critical sections (flush and version commit). Schema operations and close/destroy still take the lock exclusively and remain mutually exclusive with Optimize. - Add optimize_mtx_ to keep concurrent Optimize calls serialized, preserving the previous queuing semantics. - Make SegmentManager internally thread-safe with a shared_mutex: readers (get_segments) can now run concurrently with the segment replacement in Optimize's commit phase. - Add a regression test that asserts inserts and fetches make progress while a background Optimize is running. Fixes #553 --- src/db/collection.cc | 13 ++++- src/db/index/segment/segment_manager.cc | 41 ++++++++++--- src/db/index/segment/segment_manager.h | 6 ++ tests/db/collection_test.cc | 78 +++++++++++++++++++++++++ 4 files changed, 126 insertions(+), 12 deletions(-) diff --git a/src/db/collection.cc b/src/db/collection.cc index 33750f3b1..466564193 100644 --- a/src/db/collection.cc +++ b/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 segment_id_allocator_; std::atomic tmp_segment_id_allocator_; @@ -779,9 +781,14 @@ std::vector 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); diff --git a/src/db/index/segment/segment_manager.cc b/src/db/index/segment/segment_manager.cc index 7fb3e7a21..d42ce5610 100644 --- a/src/db/index/segment/segment_manager.cc +++ b/src/db/index/segment/segment_manager.cc @@ -28,11 +28,13 @@ Status SegmentManager::add_segment(Segment::Ptr segment) { return Status::InvalidArgument("Segment is null"); } + std::unique_lock lock(mutex_); segments_map_[segment->id()] = segment; return Status::OK(); } Status SegmentManager::remove_segment(SegmentID segment_id) { + std::unique_lock lock(mutex_); auto iter = segments_map_.find(segment_id); if (iter == segments_map_.end()) { return Status::NotFound("Segment not found"); @@ -43,6 +45,7 @@ Status SegmentManager::remove_segment(SegmentID segment_id) { } Status SegmentManager::destroy_segment(SegmentID segment_id) { + std::unique_lock lock(mutex_); auto iter = segments_map_.find(segment_id); if (iter == segments_map_.end()) { return Status::NotFound("Segment not found"); @@ -57,8 +60,11 @@ Status SegmentManager::destroy_segment(SegmentID segment_id) { std::vector SegmentManager::get_segments() const { std::vector segments; - for (auto &pair : segments_map_) { - segments.push_back(pair.second); + { + std::shared_lock lock(mutex_); + for (auto &pair : segments_map_) { + segments.push_back(pair.second); + } } std::sort(segments.begin(), segments.end(), [](Segment::Ptr a, Segment::Ptr b) { @@ -69,8 +75,11 @@ std::vector SegmentManager::get_segments() const { std::vector SegmentManager::get_segments_meta() const { std::vector segments_meta; - for (auto &pair : segments_map_) { - segments_meta.push_back(pair.second->meta()); + { + std::shared_lock lock(mutex_); + for (auto &pair : segments_map_) { + segments_meta.push_back(pair.second->meta()); + } } std::sort(segments_meta.begin(), segments_meta.end(), @@ -89,8 +98,11 @@ Status SegmentManager::add_column(const FieldSchema::Ptr &column_schema, } std::vector> futures; - std::vector> segments( - segments_map_.begin(), segments_map_.end()); + std::vector> segments; + { + std::shared_lock lock(mutex_); + segments.assign(segments_map_.begin(), segments_map_.end()); + } for (size_t i = 0; i < segments.size(); i += concurrency) { size_t end = std::min(i + concurrency, segments.size()); @@ -122,8 +134,11 @@ Status SegmentManager::alter_column(const std::string &column_name, } std::vector> futures; - std::vector> segments( - segments_map_.begin(), segments_map_.end()); + std::vector> segments; + { + std::shared_lock lock(mutex_); + segments.assign(segments_map_.begin(), segments_map_.end()); + } for (size_t i = 0; i < segments.size(); i += concurrency) { size_t end = std::min(i + concurrency, segments.size()); @@ -148,7 +163,15 @@ Status SegmentManager::alter_column(const std::string &column_name, } Status SegmentManager::drop_column(const std::string &column_name) { - for (auto &[segment_id, segment] : segments_map_) { + std::vector segments; + { + std::shared_lock lock(mutex_); + for (auto &[segment_id, segment] : segments_map_) { + segments.push_back(segment); + } + } + + for (auto &segment : segments) { auto s = segment->drop_column(column_name); CHECK_RETURN_STATUS(s); } diff --git a/src/db/index/segment/segment_manager.h b/src/db/index/segment/segment_manager.h index 9af2ac497..ae2851787 100644 --- a/src/db/index/segment/segment_manager.h +++ b/src/db/index/segment/segment_manager.h @@ -13,6 +13,7 @@ // limitations under the License. #pragma once +#include #include #include #include "segment.h" @@ -27,6 +28,7 @@ class SegmentManager { public: uint32_t segment_count() const { + std::shared_lock lock(mutex_); return segments_map_.size(); } @@ -50,6 +52,10 @@ class SegmentManager { Status drop_column(const std::string &column_name); private: + // protects segments_map_ against concurrent readers (e.g. queries + // collecting segments) while segments are added/removed by writes, + // optimize or DDL operations + mutable std::shared_mutex mutex_; std::unordered_map segments_map_; }; } // namespace zvec \ No newline at end of file diff --git a/tests/db/collection_test.cc b/tests/db/collection_test.cc index d837bb538..9483f8b21 100644 --- a/tests/db/collection_test.cc +++ b/tests/db/collection_test.cc @@ -14,6 +14,8 @@ #include "zvec/db/collection.h" #include +#include +#include #include #include #include @@ -23,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -2917,6 +2920,81 @@ TEST_F(CollectionTest, Feature_Optimize_General) { } } +TEST_F(CollectionTest, Feature_Optimize_Concurrent_Insert_NonBlocking) { + // Regression test: Optimize() used to hold the exclusive schema lock for + // its whole duration, blocking Insert/Fetch until the compact finished. + // Writes and reads must make progress while a background Optimize is + // running. + int initial_doc_count = 20000; + + auto schema = TestHelper::CreateSchemaWithVectorIndex(); + auto options = CollectionOptions{false, true, 64 * 1024 * 1024}; + auto collection = TestHelper::CreateCollectionWithDoc( + col_path, *schema, options, 0, initial_doc_count, false); + ASSERT_NE(collection, nullptr); + ASSERT_TRUE(collection->Flush().ok()); + + std::atomic optimize_done{false}; + Status optimize_status; + std::thread optimizer([&] { + optimize_status = collection->Optimize(OptimizeOptions{0}); + optimize_done.store(true); + }); + + // give the optimizer time to take its locks and start compacting + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + + int batch_size = 100; + int inserted = 0; + int batches_during_optimize = 0; + uint64_t next_doc_id = initial_doc_count; + while (!optimize_done.load() && inserted < 100000) { + std::vector docs; + for (int i = 0; i < batch_size; i++) { + docs.push_back(TestHelper::CreateDoc(next_doc_id + i, *schema)); + } + auto res = collection->Insert(docs); + ASSERT_TRUE(res.has_value()); + for (auto &st : res.value()) { + ASSERT_TRUE(st.ok()); + } + next_doc_id += batch_size; + inserted += batch_size; + if (!optimize_done.load()) { + batches_during_optimize++; + } + + // reads must also make progress, including while Optimize commits its + // result (exercises SegmentManager under concurrent access) + auto fetched = collection->Fetch({TestHelper::MakePK(0)}); + ASSERT_TRUE(fetched.has_value()); + ASSERT_EQ(fetched.value().size(), 1); + } + + optimizer.join(); + ASSERT_TRUE(optimize_status.ok()); + + // if Insert had been blocked for the whole Optimize duration, no batch + // could have completed while Optimize was still running + ASSERT_GE(batches_during_optimize, 2); + + auto stats = collection->Stats().value(); + ASSERT_EQ(stats.doc_count, (uint64_t)(initial_doc_count + inserted)); + + // spot-check data written before, during and after the optimize + for (uint64_t doc_id : + std::vector{0, (uint64_t)initial_doc_count - 1, + (uint64_t)initial_doc_count, next_doc_id - 1}) { + auto expect_doc = TestHelper::CreateDoc(doc_id, *schema); + auto result = collection->Fetch({expect_doc.pk()}); + ASSERT_TRUE(result.has_value()); + ASSERT_EQ(result.value().count(expect_doc.pk()), 1); + auto doc = result.value()[expect_doc.pk()]; + ASSERT_NE(doc, nullptr); + ASSERT_EQ(*doc, expect_doc); + } +} + TEST_F(CollectionTest, Feature_Optimize_Repeated) { auto run_repeated_optimize_test = [&](bool enable_mmap, IndexParams::Ptr index_params) { From 4ddcff23a05afdcecc79dff7ed533504d7be69ce Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=90=AA=E9=9C=84?= Date: Fri, 24 Jul 2026 17:20:04 +0800 Subject: [PATCH 2/3] fix(collection): make Optimize commit phase safe for concurrent readers Strengthening the UT to run Insert, Fetch and Query concurrently with a background Optimize (review feedback) exposed a use-after-free: the commit phase called reload_vector_index() on a live segment, destroying vector indexers and mutating block metas that concurrent readers were still using (readers only hold the shared schema lock). - Commit the new segment set by reopening segments from their new metas and swapping instances, never mutating live segments in place; in-flight readers keep replaced instances alive via shared_ptr. - Add SegmentManager::replace_segments() to apply the whole swap under one lock (validate first, then mutate), so readers never observe a partially committed segment set (no transient duplicate or missing segments), and drop the now-unused destroy_segment(). - Extend the regression test: writer, fetcher and querier threads now run concurrently with Optimize; worker failures are collected and asserted on the main thread. Query errors caused by the pre-existing Insert/Query visibility race in the writing segment (reproducible without Optimize) are tolerated and documented; a quiescent full-topk query is asserted after all threads join. --- src/db/collection.cc | 40 +++++-- src/db/index/segment/segment_manager.cc | 42 +++++-- src/db/index/segment/segment_manager.h | 8 +- tests/db/collection_test.cc | 152 +++++++++++++++++++----- 4 files changed, 197 insertions(+), 45 deletions(-) diff --git a/src/db/collection.cc b/src/db/collection.cc index 466564193..8f3157220 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -883,7 +883,19 @@ 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 segments_to_add; + std::vector segment_ids_to_destroy; for (auto &task : tasks) { auto task_info = task->task_info(); @@ -898,24 +910,32 @@ Status CollectionImpl::Optimize(const OptimizeOptions &options) { 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(task_info)) { auto create_index_task = std::get(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 = Segment::Open( + path_, *schema_, *create_index_task.output_segment_meta_, id_map_, + delete_store_, version_manager_, + SegmentOptions{true, options_.enable_mmap_}); + 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); } return Status::OK(); diff --git a/src/db/index/segment/segment_manager.cc b/src/db/index/segment/segment_manager.cc index d42ce5610..fa4eeb3da 100644 --- a/src/db/index/segment/segment_manager.cc +++ b/src/db/index/segment/segment_manager.cc @@ -44,17 +44,45 @@ Status SegmentManager::remove_segment(SegmentID segment_id) { return Status::OK(); } -Status SegmentManager::destroy_segment(SegmentID segment_id) { +Status SegmentManager::replace_segments( + const std::vector &segments_to_add, + const std::vector &segment_ids_to_destroy) { std::unique_lock lock(mutex_); - auto iter = segments_map_.find(segment_id); - if (iter == segments_map_.end()) { - return Status::NotFound("Segment not found"); + + // Phase 1: validate everything up front, so that a failure never leaves + // a partially committed segment set behind. + for (auto &segment : segments_to_add) { + if (!segment) { + return Status::InvalidArgument("Segment is null"); + } + } + for (auto segment_id : segment_ids_to_destroy) { + auto iter = segments_map_.find(segment_id); + if (iter == segments_map_.end()) { + return Status::NotFound("Segment not found"); + } + } + + // Phase 2: apply the whole swap under the single exclusive lock, so + // concurrent get_segments() observe either the complete old set or the + // complete new set, never a mix (e.g. merge inputs and their merged + // output visible at the same time). destroy() cannot fail here: a + // segment is marked and erased atomically under this lock, so a + // map-resident segment is never already marked. + for (auto segment_id : segment_ids_to_destroy) { + auto iter = segments_map_.find(segment_id); + auto s = iter->second->destroy(); + CHECK_RETURN_STATUS(s); + segments_map_.erase(iter); } - auto s = iter->second->destroy(); - CHECK_RETURN_STATUS(s); + for (auto &segment : segments_to_add) { + // overwrites any existing entry with the same id: the replaced + // instance is dropped without destroy() so its shared files stay on + // disk; in-flight readers keep it alive until they finish + segments_map_[segment->id()] = segment; + } - segments_map_.erase(segment_id); return Status::OK(); } diff --git a/src/db/index/segment/segment_manager.h b/src/db/index/segment/segment_manager.h index ae2851787..ab6e32a20 100644 --- a/src/db/index/segment/segment_manager.h +++ b/src/db/index/segment/segment_manager.h @@ -36,7 +36,13 @@ class SegmentManager { Status remove_segment(SegmentID segment_id); - Status destroy_segment(SegmentID segment_id); + //! Atomically add (or replace by id) segments and destroy others, so that + //! concurrent readers never observe a partially committed segment set. + //! Segments replaced by id are dropped without destroy(): their files are + //! shared with the replacing instance and in-flight readers keep the old + //! instance alive via shared_ptr until they finish. + Status replace_segments(const std::vector &segments_to_add, + const std::vector &segment_ids_to_destroy); std::vector get_segments() const; diff --git a/tests/db/collection_test.cc b/tests/db/collection_test.cc index 9483f8b21..6f5bc0da9 100644 --- a/tests/db/collection_test.cc +++ b/tests/db/collection_test.cc @@ -2920,11 +2920,11 @@ TEST_F(CollectionTest, Feature_Optimize_General) { } } -TEST_F(CollectionTest, Feature_Optimize_Concurrent_Insert_NonBlocking) { +TEST_F(CollectionTest, Feature_Optimize_Concurrent_ReadWrite_NonBlocking) { // Regression test: Optimize() used to hold the exclusive schema lock for - // its whole duration, blocking Insert/Fetch until the compact finished. - // Writes and reads must make progress while a background Optimize is - // running. + // its whole duration, blocking writes and reads until the compact + // finished. Insert, Fetch and Query must all make progress concurrently + // while a background Optimize is running. int initial_doc_count = 20000; auto schema = TestHelper::CreateSchemaWithVectorIndex(); @@ -2944,43 +2944,141 @@ TEST_F(CollectionTest, Feature_Optimize_Concurrent_Insert_NonBlocking) { // give the optimizer time to take its locks and start compacting std::this_thread::sleep_for(std::chrono::milliseconds(200)); + // Worker results are collected here and asserted on the main thread; + // gtest fatal assertions are not reliable on non-main threads. + struct WorkerResult { + int ops_during_optimize{0}; + int errors{0}; + std::string first_error; + }; + + auto record_error = [](WorkerResult *r, const std::string &msg) { + if (r->errors++ == 0) { + r->first_error = msg; + } + }; + + // writer: keeps inserting batches while Optimize is running + WorkerResult writer_result; int batch_size = 100; int inserted = 0; - int batches_during_optimize = 0; uint64_t next_doc_id = initial_doc_count; - while (!optimize_done.load() && inserted < 100000) { - std::vector docs; - for (int i = 0; i < batch_size; i++) { - docs.push_back(TestHelper::CreateDoc(next_doc_id + i, *schema)); - } - auto res = collection->Insert(docs); - ASSERT_TRUE(res.has_value()); - for (auto &st : res.value()) { - ASSERT_TRUE(st.ok()); + std::thread writer([&] { + while (!optimize_done.load() && inserted < 100000) { + std::vector docs; + for (int i = 0; i < batch_size; i++) { + docs.push_back(TestHelper::CreateDoc(next_doc_id + i, *schema)); + } + auto res = collection->Insert(docs); + if (!res.has_value()) { + record_error(&writer_result, res.error().message()); + break; + } + for (auto &st : res.value()) { + if (!st.ok()) { + record_error(&writer_result, st.message()); + } + } + next_doc_id += batch_size; + inserted += batch_size; + if (!optimize_done.load()) { + writer_result.ops_during_optimize++; + } } - next_doc_id += batch_size; - inserted += batch_size; - if (!optimize_done.load()) { - batches_during_optimize++; + }); + + // fetcher: point reads must return correct data throughout, including + // while Optimize commits its result (exercises SegmentManager under + // concurrent segment replacement) + WorkerResult fetch_result; + std::thread fetcher([&] { + auto expect_doc = TestHelper::CreateDoc(0, *schema); + while (!optimize_done.load()) { + auto fetched = collection->Fetch({expect_doc.pk()}); + if (!fetched.has_value()) { + record_error(&fetch_result, fetched.error().message()); + break; + } + auto iter = fetched.value().find(expect_doc.pk()); + if (iter == fetched.value().end() || iter->second == nullptr || + *iter->second != expect_doc) { + record_error(&fetch_result, "fetched doc mismatch"); + } else if (!optimize_done.load()) { + fetch_result.ops_during_optimize++; + } } + }); - // reads must also make progress, including while Optimize commits its - // result (exercises SegmentManager under concurrent access) - auto fetched = collection->Fetch({TestHelper::MakePK(0)}); - ASSERT_TRUE(fetched.has_value()); - ASSERT_EQ(fetched.value().size(), 1); - } + // querier: vector searches must keep succeeding while segments are being + // compacted and swapped + WorkerResult query_result; + std::atomic transient_query_errors{0}; + std::thread querier([&] { + auto query_doc = TestHelper::CreateDoc(1, *schema); + auto vector = query_doc.get>("dense_fp32"); + if (!vector.has_value()) { + record_error(&query_result, "query vector missing"); + return; + } + while (!optimize_done.load()) { + SearchQuery query; + query.topk_ = 10; + query.target_.field_name_ = "dense_fp32"; + query.target_.set_vector( + std::string((char *)vector.value().data(), + vector.value().size() * sizeof(float))); + auto result = collection->Query(query); + if (!result.has_value()) { + // Tolerated: a doc admitted by the writing segment's streaming + // vector index may not be visible in its forward store yet, which + // transiently fails the query. This is a pre-existing Insert/Query + // visibility race, reproducible without Optimize; it is unrelated + // to Optimize concurrency, which is what this test asserts on. + transient_query_errors.fetch_add(1); + continue; + } + if (result.value().empty()) { + record_error(&query_result, "query returned no results"); + } else if (!optimize_done.load()) { + query_result.ops_during_optimize++; + } + } + }); optimizer.join(); + writer.join(); + fetcher.join(); + querier.join(); + ASSERT_TRUE(optimize_status.ok()); + ASSERT_EQ(writer_result.errors, 0) << writer_result.first_error; + ASSERT_EQ(fetch_result.errors, 0) << fetch_result.first_error; + ASSERT_EQ(query_result.errors, 0) << query_result.first_error; - // if Insert had been blocked for the whole Optimize duration, no batch - // could have completed while Optimize was still running - ASSERT_GE(batches_during_optimize, 2); + // if any of them had been blocked for the whole Optimize duration, no + // operation could have completed while Optimize was still running + ASSERT_GE(writer_result.ops_during_optimize, 2); + ASSERT_GE(fetch_result.ops_during_optimize, 2); + ASSERT_GE(query_result.ops_during_optimize, 2); auto stats = collection->Stats().value(); ASSERT_EQ(stats.doc_count, (uint64_t)(initial_doc_count + inserted)); + // once quiescent, a vector search must return the full topk + { + auto query_doc = TestHelper::CreateDoc(1, *schema); + auto vector = query_doc.get>("dense_fp32"); + ASSERT_TRUE(vector.has_value()); + SearchQuery query; + query.topk_ = 10; + query.target_.field_name_ = "dense_fp32"; + query.target_.set_vector(std::string( + (char *)vector.value().data(), vector.value().size() * sizeof(float))); + auto result = collection->Query(query); + ASSERT_TRUE(result.has_value()); + ASSERT_EQ(result.value().size(), (size_t)query.topk_); + } + // spot-check data written before, during and after the optimize for (uint64_t doc_id : std::vector{0, (uint64_t)initial_doc_count - 1, From 8b26f6c4ac8c050bd82292c26170d1b4e027639d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=90=AA=E9=9C=84?= Date: Mon, 27 Jul 2026 20:24:40 +0800 Subject: [PATCH 3/3] fix(segment): remove superseded vector index files after Optimize With Optimize committing via segment reopen-and-swap, the pre-index vector files of a replaced segment instance (e.g. the flat file that a newly built HNSW index supersedes) were no longer deleted and lingered on disk until the segment was compacted away. Reclaim them with the same mark-and-defer model segments already use: - VectorColumnIndexer: add MarkDestroyOnRelease(); a marked indexer closes and removes its index file in the destructor, i.e. only after the last reference is released, so in-flight readers of the replaced segment instance keep the file alive until they finish. - Segment: add mark_vector_index_files_for_removal(). 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, so it must not be marked. - Optimize: after the new segment set is committed, mark the replaced instances' superseded indexers. - Add a regression test asserting exactly one index file remains per indexed column after Optimize and that a fresh open reads all docs. --- src/db/collection.cc | 43 +++++++++++++---- .../vector_column/vector_column_indexer.cc | 20 ++++++++ .../vector_column/vector_column_indexer.h | 12 ++++- src/db/index/segment/segment.cc | 33 +++++++++++++ src/db/index/segment/segment.h | 14 ++++++ tests/db/collection_test.cc | 48 +++++++++++++++++++ 6 files changed, 161 insertions(+), 9 deletions(-) diff --git a/src/db/collection.cc b/src/db/collection.cc index 8f3157220..70684538b 100644 --- a/src/db/collection.cc +++ b/src/db/collection.cc @@ -896,6 +896,12 @@ Status CollectionImpl::Optimize(const OptimizeOptions &options) { // removed from disk only when the last reference is dropped. std::vector segments_to_add; std::vector segment_ids_to_destroy; + auto reopen_segment = + [&](const SegmentMeta::Ptr &meta) -> Result { + 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(); @@ -903,10 +909,7 @@ Status CollectionImpl::Optimize(const OptimizeOptions &options) { auto compact_task = std::get(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(); } @@ -922,10 +925,8 @@ Status CollectionImpl::Optimize(const OptimizeOptions &options) { // 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 = Segment::Open( - path_, *schema_, *create_index_task.output_segment_meta_, id_map_, - delete_store_, version_manager_, - SegmentOptions{true, options_.enable_mmap_}); + auto new_segment = + reopen_segment(create_index_task.output_segment_meta_); if (!new_segment.has_value()) { return new_segment.error(); } @@ -936,6 +937,32 @@ Status CollectionImpl::Optimize(const OptimizeOptions &options) { 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(task_info)) { + continue; + } + auto create_index_task = std::get(task_info); + std::vector superseded_columns; + std::vector 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); + } } return Status::OK(); diff --git a/src/db/index/column/vector_column/vector_column_indexer.cc b/src/db/index/column/vector_column/vector_column_indexer.cc index eed5dda46..653ba44fe 100644 --- a/src/db/index/column/vector_column/vector_column_indexer.cc +++ b/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 #include #include #include @@ -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) { diff --git a/src/db/index/column/vector_column/vector_column_indexer.h b/src/db/index/column/vector_column/vector_column_indexer.h index 3c7ac7ccd..432d54d59 100644 --- a/src/db/index/column/vector_column/vector_column_indexer.h +++ b/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 #include #include #include @@ -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 destroy_on_release_{false}; }; diff --git a/src/db/index/segment/segment.cc b/src/db/index/segment/segment.cc index a4f2e5a94..8a0f8ece0 100644 --- a/src/db/index/segment/segment.cc +++ b/src/db/index/segment/segment.cc @@ -194,6 +194,10 @@ class SegmentImpl : public Segment, const std::unordered_map &quant_vector_indexers) override; + void mark_vector_index_files_for_removal( + const std::vector &columns, + const std::vector &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 &columns, + const std::vector &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); diff --git a/src/db/index/segment/segment.h b/src/db/index/segment/segment.h index 0b3a8b87a..57a221eb6 100644 --- a/src/db/index/segment/segment.h +++ b/src/db/index/segment/segment.h @@ -106,6 +106,20 @@ class Segment { const std::unordered_map &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 & /*columns*/, + const std::vector & /*quant_columns*/) {} + virtual bool vector_index_ready( const std::string &column, const IndexParams::Ptr &index_params) const = 0; diff --git a/tests/db/collection_test.cc b/tests/db/collection_test.cc index 6f5bc0da9..e211a8b8b 100644 --- a/tests/db/collection_test.cc +++ b/tests/db/collection_test.cc @@ -3093,6 +3093,54 @@ TEST_F(CollectionTest, Feature_Optimize_Concurrent_ReadWrite_NonBlocking) { } } +TEST_F(CollectionTest, Feature_Optimize_Superseded_Index_File_Removed) { + namespace fs = std::filesystem; + int doc_count = 1000; + + auto schema = TestHelper::CreateSchemaWithVectorIndex(); + auto options = CollectionOptions{false, true, 64 * 1024 * 1024}; + auto collection = TestHelper::CreateCollectionWithDoc( + col_path, *schema, options, 0, doc_count, false); + ASSERT_NE(collection, nullptr); + ASSERT_TRUE(collection->Flush().ok()); + + ASSERT_TRUE(collection->Optimize(OptimizeOptions{0}).ok()); + + auto stats = collection->Stats().value(); + ASSERT_EQ(stats.doc_count, (uint64_t)doc_count); + ASSERT_EQ(stats.index_completeness["dense_fp32"], 1); + + // Building the vector index supersedes the flat vector file written when + // the segment was dumped. Once the replaced segment instance is released + // (no readers left after Optimize returns), the superseded file must be + // removed, leaving exactly one index file for the column. + int dense_fp32_index_files = 0; + for (const auto &entry : fs::directory_iterator(col_path + "/0")) { + auto name = entry.path().filename().string(); + if (name.rfind("dense_fp32.index.", 0) == 0) { + dense_fp32_index_files++; + } + } + ASSERT_EQ(dense_fp32_index_files, 1); + + // the removed file must not be one the collection still needs: a fresh + // open must read every doc correctly through the new index + collection.reset(); + auto result = Collection::Open(col_path, options); + ASSERT_TRUE(result.has_value()); + collection = std::move(result.value()); + + for (int i : {0, doc_count / 2, doc_count - 1}) { + auto expect_doc = TestHelper::CreateDoc(i, *schema); + auto fetched = collection->Fetch({expect_doc.pk()}); + ASSERT_TRUE(fetched.has_value()); + ASSERT_EQ(fetched.value().count(expect_doc.pk()), 1); + auto doc = fetched.value()[expect_doc.pk()]; + ASSERT_NE(doc, nullptr); + ASSERT_EQ(*doc, expect_doc); + } +} + TEST_F(CollectionTest, Feature_Optimize_Repeated) { auto run_repeated_optimize_test = [&](bool enable_mmap, IndexParams::Ptr index_params) {