From c12f51f972e882c8040e1f30b8587cd12c638f63 Mon Sep 17 00:00:00 2001 From: Minoru OSUKA Date: Tue, 14 Jul 2026 07:27:42 +0900 Subject: [PATCH] perf(engine): retain vector writer across commits + lexical segment-meta cache (#864) Remove the two verified per-doc write-path offenders (Issue #864, making the #572 commit batching real): Vector: VectorStore::commit no longer take()s the cached writer when the index opts in via the new VectorIndex::retain_writer_after_commit() (default false; true for HnswIndex only). Post-commit the writer's in-memory state equals the file it just wrote, so the first upsert after every commit no longer reloads + dequantizes the entire .hnsw. The cache is dropped whenever the index is rewritten behind the writer - auto-compaction inside commit, VectorStore::optimize (now async) - and on any mid-ladder commit failure, where the writer/disk agreement is unknown. A new VectorIndexWriter::has_pending_changes() (HNSW: !is_finalized) gates the flush, so a no-change commit skips the full index rewrite, and optimize() flushes a dirty writer whose buffer was emptied by deletions instead of discarding the mutation. Lexical (closes #559, #571): InvertedIndexWriter caches (segment_id, min, max) ranges - seeded by the constructor's existing recovery scan, extended on flush - with a doc_id > max_committed_doc_id fast path, so fresh-id upserts stop listing + JSON-parsing every .meta file, and reuses one lazily-created DeletionManager across overwrites instead of reloading every .delmap per call. LexicalStore::commit and ::optimize hold the writer-cache lock across their merges so a concurrent upsert can neither run against a stale cache nor construct a writer from a racing mid-merge scan; invalidate_segment_cache rebuilds atomically (swap on success) and optimize drops the writer if the rebuild fails. Verification is deterministic: CountingStorage decorators assert exact .hnsw open/create and list_files counts; both corruption gates (compaction resurrection, ghost-segment deletion) and both fault-injection tests (failed commit drops the writer, failed invalidate preserves the old cache) fail when their fix is disabled. A pre-existing HNSW incremental-append reachability bug discovered by these tests is tracked separately as #868; membership assertions gate on stats().document_count until it is fixed. Closes #864, closes #572, closes #559, closes #571. Refs #551, #868 --- .../src/concepts/indexing/vector_indexing.md | 4 + docs/src/concepts/indexing/vector_indexing.md | 17 + laurus/src/lexical/index/inverted/writer.rs | 207 ++++-- laurus/src/lexical/store.rs | 35 +- laurus/src/lexical/writer.rs | 16 + laurus/src/vector/index.rs | 23 + laurus/src/vector/index/hnsw.rs | 11 + laurus/src/vector/index/hnsw/writer.rs | 10 + laurus/src/vector/store.rs | 101 ++- laurus/src/vector/store/embedding_writer.rs | 4 + laurus/src/vector/writer.rs | 17 + .../tests/lexical_segment_meta_cache_test.rs | 348 ++++++++++ laurus/tests/vector_hnsw_soft_delete_test.rs | 4 +- laurus/tests/vector_writer_retention_test.rs | 597 ++++++++++++++++++ 14 files changed, 1328 insertions(+), 66 deletions(-) create mode 100644 laurus/tests/lexical_segment_meta_cache_test.rs create mode 100644 laurus/tests/vector_writer_retention_test.rs diff --git a/docs/ja/src/concepts/indexing/vector_indexing.md b/docs/ja/src/concepts/indexing/vector_indexing.md index 6c96caf3..e8614d61 100644 --- a/docs/ja/src/concepts/indexing/vector_indexing.md +++ b/docs/ja/src/concepts/indexing/vector_indexing.md @@ -497,6 +497,10 @@ search の独立 budget）や 4 bit PQ variant でこの gap を埋める方向 | Flat | `.flat` | 生ベクトルとメタデータ | | IVF | `.ivf` | クラスタセントロイド、割り当て済みベクトル、メタデータ | +### コミット跨ぎの writer 保持 + +HNSW では、ストアは writer を**コミットを跨いで**キャッシュし続けます（Issue #864）。コミット直後の writer のメモリ内状態は、たった今書き出したファイルと等価であるため、コミット後最初の upsert はストレージから `.hnsw` 全体をリロード（+ 全 dequantize）する代わりに、その場で writer を拡張します — コミットの多いインジェストでは、このリロードがサイクルあたり O(インデックスサイズ) のコストでした。ただし、writer の**裏で**インデックスが書き換えられる場合 — コミット中の auto-compaction 発火、または明示的な `optimize()` — はキャッシュを破棄します。どちらも fresh writer 経由でファイルを再構築して deletion bitmap をクリアするため、stale な保持 writer は回収済みベクトルを書き戻してしまうからです。また、コミットが途中で失敗した場合も writer とディスクの一致が不明になるため破棄します。副次効果として、保留変更のない保持 writer はインデックス書き出し自体をスキップするため、変更なしの `commit()` は `.hnsw` への書き込みコストがゼロになります。Flat と IVF は同じ不変条件の監査が済むまで、従来のコミット時破棄の挙動を維持します。 + ## コード例 ```rust diff --git a/docs/src/concepts/indexing/vector_indexing.md b/docs/src/concepts/indexing/vector_indexing.md index bdfaaafe..b7e24245 100644 --- a/docs/src/concepts/indexing/vector_indexing.md +++ b/docs/src/concepts/indexing/vector_indexing.md @@ -520,6 +520,23 @@ Each vector index type stores its data in a single segment file: | Flat | `.flat` | Raw vectors and metadata | | IVF | `.ivf` | Cluster centroids, assigned vectors, and metadata | +### Writer retention across commits + +For HNSW, the store keeps its writer cached **across commits** (Issue #864): +after a commit the writer's in-memory state is equivalent to the file it just +wrote, so the first upsert after a commit extends it in place instead of +reloading (and dequantizing) the whole `.hnsw` from storage — under +commit-heavy ingest that reload was O(index size) per cycle. The cache is +dropped whenever the index is rewritten *behind* the writer — auto-compaction +firing inside a commit, or an explicit `optimize()` — because both rebuild the +file through a fresh writer and clear the deletion bitmap, and a stale +retained writer would write the reclaimed vectors back — and whenever a +commit fails partway, where the writer/disk agreement is unknown. As a bonus, +a retained writer with no pending changes skips the index rewrite entirely, +so a no-change `commit()` costs no `.hnsw` write at all. Flat and IVF keep +the previous drop-on-commit behavior until they are audited for the same +invariants. + ## Code Example ```rust diff --git a/laurus/src/lexical/index/inverted/writer.rs b/laurus/src/lexical/index/inverted/writer.rs index 57019b5c..bedc8d7e 100644 --- a/laurus/src/lexical/index/inverted/writer.rs +++ b/laurus/src/lexical/index/inverted/writer.rs @@ -166,6 +166,34 @@ pub struct InvertedIndexWriter { /// Pending deletions that are not yet reflected in the reader (NRT). pending_deletions: std::collections::HashSet, + + /// Cached `(segment_id, min_doc_id, max_doc_id)` for every committed + /// segment, mirroring the `*.meta` files on storage (Issue #559 / #864). + /// + /// Built from the constructor's existing recovery scan (no extra I/O), + /// extended in place whenever this writer flushes a segment, and rebuilt + /// by [`Self::invalidate_segment_cache`] after an external segment + /// rewrite ([`LexicalStore::optimize`](crate::lexical::store::LexicalStore::optimize) + /// force-merge — the only path that rewrites segments behind a live + /// writer; `LexicalStore::commit` drops the writer before it merges). + /// Lets [`Self::find_segments_for_doc`] answer from memory instead of + /// listing + JSON-parsing every `.meta` file per upsert. + segment_ranges: Vec<(String, u64, u64)>, + + /// Highest `max_doc_id` across [`Self::segment_ranges`] (0 when no + /// segments exist). Fresh doc IDs handed out by the WAL are strictly + /// greater, so the steady-state ingest path rejects the "is this doc in a + /// committed segment?" question with one integer compare. + max_committed_doc_id: u64, + + /// Lazily created [`DeletionManager`](crate::maintenance::deletion::DeletionManager) + /// reused across upserts of already-committed documents (Issue #571). + /// Construction loads every `.delmap` bitmap from storage, so building it + /// once per writer — and only when an overwrite actually needs it — + /// replaces a full bitmap reload per call. Dropped together with + /// [`Self::segment_ranges`] on [`Self::invalidate_segment_cache`] because + /// a force-merge deletes the segments (and bitmaps) it holds in memory. + deletion_manager: Option, } impl std::fmt::Debug for InvertedIndexWriter { @@ -184,9 +212,13 @@ impl std::fmt::Debug for InvertedIndexWriter { impl InvertedIndexWriter { /// Create a new inverted index writer (schema-less mode). pub fn new(storage: Arc, config: InvertedIndexWriterConfig) -> Result { - // Recover state from existing segments + // Recover state from existing segments. The same scan also seeds the + // segment-range cache (Issue #559 / #864), so the cache is warm from + // birth at no extra I/O. let mut next_doc_id = 0; let mut max_segment_id = -1i32; + let mut segment_ranges = Vec::new(); + let mut max_committed_doc_id = 0u64; if let Ok(files) = storage.list_files() { for file in files { @@ -201,6 +233,8 @@ impl InvertedIndexWriter { next_doc_id = next_doc_id.max(local_id + 1); } max_segment_id = max_segment_id.max(meta.generation as i32); + max_committed_doc_id = max_committed_doc_id.max(meta.max_doc_id); + segment_ranges.push((meta.segment_id, meta.min_doc_id, meta.max_doc_id)); } } } @@ -239,6 +273,9 @@ impl InvertedIndexWriter { last_wal_seq: base_metadata.last_wal_seq, base_metadata, pending_deletions: std::collections::HashSet::new(), + segment_ranges, + max_committed_doc_id, + deletion_manager: None, }) } @@ -722,6 +759,11 @@ impl InvertedIndexWriter { self.write_segment_files(&segment_name)?; + // Mirror the freshly written `.meta` into the segment-range cache + // (Issue #559 / #864) so an overwrite of one of these docs later in + // this writer's life resolves without rescanning storage. + self.extend_segment_cache(&segment_name); + // Clear buffers self.buffered_docs.clear(); self.buffered_doc_ids.clear(); @@ -800,6 +842,7 @@ impl InvertedIndexWriter { self.index_dirty = false; } let paths = self.write_segment_files(segment_name)?; + self.extend_segment_cache(segment_name); self.buffered_docs.clear(); self.buffered_doc_ids.clear(); self.index_dirty = false; @@ -1259,8 +1302,12 @@ impl InvertedIndexWriter { Ok(()) } - /// Write segment metadata. - fn write_segment_metadata(&self, segment_name: &str) -> Result<()> { + /// Doc-ID range `(min, max)` of the currently buffered documents, + /// `(0, 0)` when the buffer is empty. Shared by + /// [`Self::write_segment_metadata`] and [`Self::extend_segment_cache`] so + /// the on-storage `.meta` and the in-memory segment-range cache can never + /// disagree. + fn buffered_doc_id_range(&self) -> (u64, u64) { let min_id = self .buffered_docs .iter() @@ -1273,6 +1320,28 @@ impl InvertedIndexWriter { .map(|(id, _)| *id) .max() .unwrap_or(0); + (min_id, max_id) + } + + /// Append the segment just written from the current buffer to + /// [`Self::segment_ranges`] and raise [`Self::max_committed_doc_id`] + /// (Issue #559 / #864). Callers must invoke this after + /// [`Self::write_segment_files`] and **before** clearing the buffers the + /// range is computed from. + /// + /// # Parameters + /// + /// - `segment_name` - Name the segment's files were written under. + fn extend_segment_cache(&mut self, segment_name: &str) { + let (min_id, max_id) = self.buffered_doc_id_range(); + self.max_committed_doc_id = self.max_committed_doc_id.max(max_id); + self.segment_ranges + .push((segment_name.to_string(), min_id, max_id)); + } + + /// Write segment metadata. + fn write_segment_metadata(&self, segment_name: &str) -> Result<()> { + let (min_id, max_id) = self.buffered_doc_id_range(); // Create SegmentInfo let info = SegmentInfo { @@ -1464,28 +1533,40 @@ impl InvertedIndexWriter { fn mark_persisted_doc_deleted(&mut self, doc_id: u64) -> Result<()> { let segments = self.find_segments_for_doc(doc_id)?; - for (segment_id, min_doc_id, max_doc_id) in segments { - // Found the segment, update deletion bitmap - let manager = crate::maintenance::deletion::DeletionManager::new( - Default::default(), // Use default config for now - self.storage.clone(), - )?; - - manager.initialize_segment(&segment_id, min_doc_id, max_doc_id)?; - - let delete_result = manager.delete_document(&segment_id, doc_id, "upsert"); - if delete_result.is_err() { - // If initializing failed (e.g. bitmap corrupted), try force re-init - // In production code we should be more careful, but here we prioritize consistency - manager.initialize_segment(&segment_id, min_doc_id, max_doc_id)?; - manager.delete_document(&segment_id, doc_id, "upsert")?; + if !segments.is_empty() { + // Create the deletion manager on first use and keep it for the + // writer's lifetime (Issue #571): construction reloads every + // `.delmap` bitmap from storage, and the fresh-id ingest path + // (empty `segments`) never needs it at all. + if self.deletion_manager.is_none() { + self.deletion_manager = Some(crate::maintenance::deletion::DeletionManager::new( + Default::default(), // Use default config for now + self.storage.clone(), + )?); } + let mut deleted = 0; + for (segment_id, min_doc_id, max_doc_id) in &segments { + let manager = self.deletion_manager.as_ref().ok_or_else(|| { + LaurusError::internal("deletion manager missing after initialization") + })?; - // Update segment metadata to reflect deletions - self.update_segment_meta_deletions(&segment_id)?; + manager.initialize_segment(segment_id, *min_doc_id, *max_doc_id)?; + let delete_result = manager.delete_document(segment_id, doc_id, "upsert"); + if delete_result.is_err() { + // If initializing failed (e.g. bitmap corrupted), try force re-init + // In production code we should be more careful, but here we prioritize consistency + manager.initialize_segment(segment_id, *min_doc_id, *max_doc_id)?; + manager.delete_document(segment_id, doc_id, "upsert")?; + } + + // Update segment metadata to reflect deletions + self.update_segment_meta_deletions(segment_id)?; + + deleted += 1; + } // Track globally - self.stats.deleted_count += 1; + self.stats.deleted_count += deleted; } // Add to pending deletions for NRT visibility @@ -1499,34 +1580,74 @@ impl InvertedIndexWriter { self.pending_deletions.contains(&doc_id) } - /// Find all segments containing the global doc_id by scanning segment metadata files. + /// Find all segments containing the global doc_id. /// Returns a list of (segment_id, min_doc_id, max_doc_id). + /// + /// Served from [`Self::segment_ranges`] (Issue #559 / #864) — the cache + /// mirrors the on-storage `*.meta` files, so this no longer lists and + /// JSON-parses them per call. Fresh doc IDs (the steady-state ingest + /// path) are rejected with a single compare against + /// [`Self::max_committed_doc_id`]. fn find_segments_for_doc(&self, doc_id: u64) -> Result> { - let mut segments = Vec::new(); - let files = self.storage.list_files()?; - for file in files { + // Fast path: WAL doc IDs are monotonic, so an ID above every + // committed segment's max cannot be in any of them. + if doc_id > self.max_committed_doc_id { + return Ok(Vec::new()); + } + // In Stable ID mode, we check if the ID is within the min/max range. + // Note: This might match multiple segments if ranges overlap across shards, + // or if we have multiple versions of the same document (upserts). + // To be 100% sure, we should check if the document actually exists in + // the segment. For now, assume the range is specific enough. + Ok(self + .segment_ranges + .iter() + .filter(|(_, min_doc_id, max_doc_id)| doc_id >= *min_doc_id && doc_id <= *max_doc_id) + .cloned() + .collect()) + } + + /// Rebuild [`Self::segment_ranges`] / [`Self::max_committed_doc_id`] from + /// the `*.meta` files on storage and drop the cached + /// [`Self::deletion_manager`]. + /// + /// Must be called after an external segment rewrite that this writer did + /// not perform itself — today that is only + /// [`LexicalStore::optimize`](crate::lexical::store::LexicalStore::optimize)'s + /// force-merge, which replaces every segment (and its deletion bitmap) + /// behind a live writer. `LexicalStore::commit` needs no call: it drops + /// the cached writer before merging, so the next writer rebuilds the + /// cache in its constructor. + /// + /// # Errors + /// + /// Returns an error if listing the storage fails; malformed `.meta` files + /// are skipped, matching the constructor's recovery scan. + pub fn invalidate_segment_cache(&mut self) -> Result<()> { + // Build the new view first and swap it in only on success: clearing + // eagerly and then failing (e.g. on `list_files`) would leave the + // writer alive with an EMPTY cache and `max_committed_doc_id == 0`, + // silently skipping every subsequent overwrite's deletion via the + // fast path — worse than the error itself. + let mut segment_ranges = Vec::new(); + let mut max_committed_doc_id = 0u64; + for file in self.storage.list_files()? { if !file.ends_with(".meta") || file == "index.meta" { continue; } - let input = match self.storage.open_input(&file) { - Ok(input) => input, - Err(_) => continue, + let Ok(input) = self.storage.open_input(&file) else { + continue; }; - let meta: SegmentInfo = match serde_json::from_reader(input) { - Ok(m) => m, - Err(_) => continue, + let Ok(meta) = serde_json::from_reader::<_, SegmentInfo>(input) else { + continue; }; - - // In Stable ID mode, we check if the ID is within the min/max range. - // Note: This might match multiple segments if ranges overlap across shards, - // or if we have multiple versions of the same document (upserts). - if doc_id >= meta.min_doc_id && doc_id <= meta.max_doc_id { - // To be 100% sure, we should check if the document actually exists in this segment. - // For now, assume this range is specific enough. - segments.push((meta.segment_id.clone(), meta.min_doc_id, meta.max_doc_id)); - } + max_committed_doc_id = max_committed_doc_id.max(meta.max_doc_id); + segment_ranges.push((meta.segment_id, meta.min_doc_id, meta.max_doc_id)); } - Ok(segments) + self.segment_ranges = segment_ranges; + self.max_committed_doc_id = max_committed_doc_id; + self.deletion_manager = None; + Ok(()) } /// Rewrite segment metadata to mark `has_deletions = true`. @@ -1585,6 +1706,10 @@ impl LexicalIndexWriter for InvertedIndexWriter { InvertedIndexWriter::add_document(self, doc) } + fn invalidate_segment_cache(&mut self) -> Result<()> { + InvertedIndexWriter::invalidate_segment_cache(self) + } + fn upsert_document(&mut self, doc_id: u64, doc: Document) -> Result<()> { InvertedIndexWriter::upsert_document(self, doc_id, doc) } diff --git a/laurus/src/lexical/store.rs b/laurus/src/lexical/store.rs index 27b4f763..0eb4111a 100644 --- a/laurus/src/lexical/store.rs +++ b/laurus/src/lexical/store.rs @@ -263,7 +263,14 @@ impl LexicalStore { /// engine.commit().unwrap(); /// ``` pub fn commit(&self) -> Result<()> { - if let Some(mut writer) = self.writer_cache.lock().take() { + // Hold the writer-cache lock across the whole ladder (Issue #864): a + // writer constructed concurrently while `maybe_merge` replaces + // segments would seed its segment-range cache from a racing, + // mid-merge `.meta` scan and never be invalidated. With the lock + // held, new writers can only be constructed strictly before or + // strictly after the merge, where the scan sees a consistent set. + let mut writer_guard = self.writer_cache.lock(); + if let Some(mut writer) = writer_guard.take() { writer.commit()?; } // Sync storage to ensure all file metadata (creation, rename, size) is @@ -277,6 +284,7 @@ impl LexicalStore { self.index.maybe_merge()?; self.index.storage().sync()?; self.index.refresh()?; + drop(writer_guard); *self.searcher_cache.write() = None; Ok(()) } @@ -323,9 +331,30 @@ impl LexicalStore { /// engine.optimize().unwrap(); /// ``` pub fn optimize(&self) -> Result<()> { - self.index.optimize()?; + // Hold the writer-cache lock across the force-merge (Issue #864): + // without it a concurrent upsert can run against the live writer's + // stale segment cache while the merge deletes those segments, marking + // deletions in ghost segments (lost dedup / duplicate versions). + let mut writer_guard = self.writer_cache.lock(); + let merge_result = self.index.optimize(); + // The force-merge replaces committed segments (and their deletion + // bitmaps) behind any live writer — rebuild its cached segment view + // even when the merge errored partway, since segments may already + // have been replaced by then. If the rebuild itself fails, drop the + // writer entirely: an unrebuildable cache must not survive to serve + // stale (or, worse, empty) ranges. + if let Some(writer) = writer_guard.as_mut() + && let Err(e) = writer.invalidate_segment_cache() + { + *writer_guard = None; + drop(writer_guard); + *self.searcher_cache.write() = None; + merge_result?; + return Err(e); + } + drop(writer_guard); *self.searcher_cache.write() = None; - Ok(()) + merge_result } /// Refresh the reader to see latest changes. diff --git a/laurus/src/lexical/writer.rs b/laurus/src/lexical/writer.rs index 82861c57..7951ed5a 100644 --- a/laurus/src/lexical/writer.rs +++ b/laurus/src/lexical/writer.rs @@ -202,4 +202,20 @@ pub trait LexicalIndexWriter: Send + Sync + std::fmt::Debug { fn is_updated_deleted(&self, _doc_id: u64) -> bool { false } + + /// Invalidate any cached view of the committed segments (Issue #864). + /// + /// Called by the store after an operation that rewrites segments behind a + /// live writer (today only + /// [`LexicalStore::optimize`](crate::lexical::store::LexicalStore::optimize)'s + /// force-merge), so a writer that caches segment metadata rebuilds it + /// from storage. The default implementation is a no-op for writers that + /// hold no such cache. + /// + /// # Errors + /// + /// Returns an error if rebuilding the cached view from storage fails. + fn invalidate_segment_cache(&mut self) -> Result<()> { + Ok(()) + } } diff --git a/laurus/src/vector/index.rs b/laurus/src/vector/index.rs index aaf03335..fb3fc40e 100644 --- a/laurus/src/vector/index.rs +++ b/laurus/src/vector/index.rs @@ -102,6 +102,29 @@ pub trait VectorIndex: Send + Sync + std::fmt::Debug { Ok(()) } + /// Whether a store may keep its cached writer across `commit()` calls + /// (Issue #572 / #864) instead of dropping it, so the first upsert after + /// a commit does not reload the whole index from storage. + /// + /// Retention is sound only when **both** hold for this index type: + /// + /// 1. The writer's post-commit in-memory state is equivalent to the file + /// it just wrote (idempotent `finalize()`, non-consuming `write()`, + /// incremental graph/buffer reuse), so subsequent commits from the + /// retained writer produce the same bytes a freshly loaded writer + /// would. + /// 2. Every disk mutation that bypasses the writer — compaction via + /// [`Self::maybe_auto_compact`], [`Self::optimize`] — invalidates the + /// store's writer cache, otherwise a stale retained writer would + /// rewrite physically reclaimed (deleted) vectors back into the index + /// with no deletion bitmap marking them. + /// + /// Defaults to `false` (the store drops the writer on commit, today's + /// behavior); `HnswIndex` opts in after auditing both conditions. + fn retain_writer_after_commit(&self) -> bool { + false + } + /// Create a searcher tailored for this index implementation. /// /// Returns a boxed [`VectorIndexSearcher`] capable of executing search/count operations. diff --git a/laurus/src/vector/index/hnsw.rs b/laurus/src/vector/index/hnsw.rs index 59e1d0f6..8bdd6781 100644 --- a/laurus/src/vector/index/hnsw.rs +++ b/laurus/src/vector/index/hnsw.rs @@ -413,6 +413,17 @@ impl VectorIndex for HnswIndex { }) } + fn retain_writer_after_commit(&self) -> bool { + // Audited for Issue #864: `finalize()` is idempotent and appends + // incrementally to the existing graph, `write(&self)` is + // non-consuming, and `delete_document` only invalidates the graph + // when a buffered vector was actually removed — so a retained + // writer's state stays equivalent to the file it just wrote. The + // bypassing mutations (`optimize` / auto-compaction) are handled by + // `VectorStore`, which invalidates its writer cache on both. + true + } + fn optimize(&self) -> Result<()> { self.check_closed()?; diff --git a/laurus/src/vector/index/hnsw/writer.rs b/laurus/src/vector/index/hnsw/writer.rs index 6b6cc852..bb9cfeea 100644 --- a/laurus/src/vector/index/hnsw/writer.rs +++ b/laurus/src/vector/index/hnsw/writer.rs @@ -1635,6 +1635,16 @@ impl VectorIndexWriter for HnswIndexWriter { self.storage.is_some() } + fn has_pending_changes(&self) -> bool { + // `finalize()` sets the flag and every mutation (add_vectors, + // delete_document/s, build) clears it, so a finalized writer's + // in-memory state has already been captured by the finalize+write + // pair and dropping it loses nothing. Note the load path constructs + // writers with `is_finalized: false`, so a freshly loaded writer + // conservatively reports pending changes. + !self.is_finalized + } + fn delete_document(&mut self, doc_id: u64) -> Result<()> { if self.is_finalized { self.is_finalized = false; diff --git a/laurus/src/vector/store.rs b/laurus/src/vector/store.rs index 0e344bf6..2f4d688f 100644 --- a/laurus/src/vector/store.rs +++ b/laurus/src/vector/store.rs @@ -333,31 +333,69 @@ impl VectorStore { /// Commit any pending changes to the index. /// - /// If a cached writer exists, this method takes it and calls - /// [`commit()`](crate::vector::writer::VectorIndexWriter::commit) (which - /// finalizes the index and writes it to storage). It then syncs the - /// underlying storage to ensure all file metadata is flushed to disk, + /// If a cached writer exists, this method calls + /// [`commit()`](crate::vector::writer::VectorIndexWriter::commit) on it + /// (which finalizes the index and writes it to storage). It then syncs + /// the underlying storage to ensure all file metadata is flushed to disk, /// refreshes the index metadata, and invalidates the searcher cache so /// that subsequent searches see the committed data. /// + /// When the index opts in via + /// [`VectorIndex::retain_writer_after_commit`] (Issue #572 / #864), the + /// committed writer stays in the cache — its in-memory state is + /// equivalent to the file it just wrote — so the first upsert after the + /// commit does not reload the whole index from storage. The cache is + /// still dropped when auto-compaction ran: compaction rewrites the index + /// through a fresh writer and clears the deletion bitmap, so a retained + /// writer would resurrect the physically reclaimed vectors on its next + /// commit. The writer-cache lock is held across the whole ladder so a + /// concurrent upsert cannot interleave with the commit. + /// /// # Errors /// - /// Returns an error if the writer commit, storage sync, or index refresh - /// fails. + /// Returns an error if the writer commit, deletion persistence, + /// compaction, storage sync, or index refresh fails. pub async fn commit(&self) -> Result<()> { - if let Some(mut writer) = self.writer_cache.lock().await.take() { - // commit() calls finalize() then write() to persist to storage - writer.commit()?; + let mut writer_guard = self.writer_cache.lock().await; + // commit() calls finalize() then write() to persist to storage. A + // retained writer with no pending changes is skipped — its state was + // already captured by the previous finalize+write, so re-committing + // would only rewrite an identical index file. + let flush_result = match writer_guard.as_mut() { + Some(writer) if writer.has_pending_changes() => writer.commit(), + _ => Ok(()), + }; + let ladder_result = flush_result + // Persist any pending logical deletions (Issue #624) so the + // deletion bitmap survives restarts. The WAL also records + // deletions, so this is a durability optimization rather than the + // source of truth. + .and_then(|_| self.index.persist_deletions()) + // Automatically compact when the deletion ratio crosses the + // configured threshold (Issue #782), so logically deleted vectors + // are physically reclaimed rather than accumulating indefinitely. + // A no-op unless the index supports it and `auto_compaction` is + // enabled. + .and_then(|_| self.index.maybe_auto_compact()); + match ladder_result { + Ok(compacted) => { + if compacted || !self.index.retain_writer_after_commit() { + *writer_guard = None; + } + } + Err(e) => { + // A mid-ladder failure leaves the writer/disk agreement + // unknown — compaction in particular may have partially + // rewritten the index and cleared the deletion bitmap before + // failing. Drop the cache so the next writer reloads ground + // truth from storage (the pre-retention behavior on every + // path), instead of retrying — or resurrecting from — a + // stale writer. + *writer_guard = None; + return Err(e); + } } - // Persist any pending logical deletions (Issue #624) so the deletion - // bitmap survives restarts. The WAL also records deletions, so this is - // a durability optimization rather than the source of truth. - self.index.persist_deletions()?; - // Automatically compact when the deletion ratio crosses the configured - // threshold (Issue #782), so logically deleted vectors are physically - // reclaimed rather than accumulating indefinitely. A no-op unless the - // index supports it and `auto_compaction` is enabled. - self.index.maybe_auto_compact()?; + drop(writer_guard); // Sync storage to ensure all file metadata (creation, rename, size) is // flushed to disk. This is critical on Windows where directory listings // and file visibility may be cached until the directory is synced. @@ -373,11 +411,34 @@ impl VectorStore { /// and then invalidates the searcher cache so the next search creates a /// fresh searcher reflecting the optimized state. /// + /// The cached writer is committed first when it holds buffered documents + /// (so optimization compacts the full state and nothing is lost), and the + /// cache is dropped in every case: `optimize()` rewrites the index + /// through a fresh writer and clears the deletion bitmap, so a writer + /// retained across it would resurrect the physically reclaimed vectors + /// on its next commit (Issue #864). + /// /// # Errors /// - /// Returns an error if the underlying index optimization fails. - pub fn optimize(&self) -> Result<()> { + /// Returns an error if flushing the cached writer or the underlying index + /// optimization fails. + pub async fn optimize(&self) -> Result<()> { + let mut writer_guard = self.writer_cache.lock().await; + // Flush uncommitted mutations first. `has_pending_changes` (not + // `pending_docs`) is the gate: a writer whose buffer was emptied by + // deletions has zero pending docs but still holds an uncommitted + // delete-everything mutation that dropping would silently discard. + let flush_result = match writer_guard.as_mut() { + Some(writer) if writer.has_pending_changes() => writer.commit(), + _ => Ok(()), + }; + // Always drop the cache — optimize() rewrites the index through a + // fresh writer and clears the deletion bitmap — and drop it on the + // flush error path too, where the writer/disk agreement is unknown. + *writer_guard = None; + flush_result?; self.index.optimize()?; + drop(writer_guard); *self.searcher_cache.write() = None; Ok(()) } diff --git a/laurus/src/vector/store/embedding_writer.rs b/laurus/src/vector/store/embedding_writer.rs index 664354f2..8e74340d 100644 --- a/laurus/src/vector/store/embedding_writer.rs +++ b/laurus/src/vector/store/embedding_writer.rs @@ -170,6 +170,10 @@ impl VectorIndexWriter for EmbeddingVectorIndexWriter { self.inner.pending_docs() } + fn has_pending_changes(&self) -> bool { + self.inner.has_pending_changes() + } + fn close(&mut self) -> Result<()> { self.inner.close() } diff --git a/laurus/src/vector/writer.rs b/laurus/src/vector/writer.rs index 86949620..dbcfa230 100644 --- a/laurus/src/vector/writer.rs +++ b/laurus/src/vector/writer.rs @@ -122,6 +122,23 @@ pub trait VectorIndexWriter: Send + Sync + std::fmt::Debug { /// This removes the document from the index buffer if it hasn't been finalized. fn delete_document(&mut self, doc_id: u64) -> Result<()>; + /// Whether the writer holds changes that have not been written to storage + /// yet (Issue #864). + /// + /// Unlike [`Self::pending_docs`] — which reports the number of *buffered + /// documents* and is therefore `0` both for a clean just-committed writer + /// and for a dirty writer whose buffer was emptied by deletions — this + /// answers "would dropping the writer lose an uncommitted mutation?". + /// Stores use it to skip the full index rewrite on a no-change commit of + /// a retained writer, and to flush a dirty writer before dropping it in + /// `optimize()`. + /// + /// The default implementation conservatively returns `true` (always + /// commit). Writers that track a finalized flag should override it. + fn has_pending_changes(&self) -> bool { + true + } + /// Delete documents matching a metadata field value. /// /// This removes matching documents from the index buffer regarding the filtering criteria. diff --git a/laurus/tests/lexical_segment_meta_cache_test.rs b/laurus/tests/lexical_segment_meta_cache_test.rs new file mode 100644 index 00000000..02c99c8a --- /dev/null +++ b/laurus/tests/lexical_segment_meta_cache_test.rs @@ -0,0 +1,348 @@ +//! Integration tests for the `InvertedIndexWriter` segment-range cache and +//! cached `DeletionManager` (Issue #864, closing #559 / #571). +//! +//! Before the fix, every upsert ran `find_segments_for_doc`, which listed the +//! storage and JSON-parsed **every** `*.meta` file, and every overwrite of a +//! committed document constructed a fresh `DeletionManager` (reloading every +//! `.delmap` bitmap). The cache mirrors the `.meta` ranges in memory: it is +//! seeded by the constructor's existing recovery scan, extended on flush, and +//! rebuilt by `invalidate_segment_cache` after `LexicalStore::optimize`'s +//! force-merge — the only path that rewrites segments behind a live writer. +//! +//! The primary gate is deterministic: a `CountingStorage` decorator counts +//! `list_files` calls, so the tests assert exact I/O counts instead of timing. + +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +use laurus::storage::{FileMetadata, LoadingMode, Storage, StorageInput, StorageOutput}; +use laurus::{Document, Result}; + +use laurus::lexical::{ + InvertedIndexWriter, InvertedIndexWriterConfig, LexicalIndexConfig, LexicalSearchRequest, + LexicalStore, TermQuery, +}; +use laurus::storage::memory::{MemoryStorage, MemoryStorageConfig}; + +/// Decorator over [`MemoryStorage`] counting `list_files` invocations — +/// the deterministic signal for "did this operation rescan segment metadata?" +/// — plus an optional fault injector failing `list_files`, to exercise the +/// cache-rebuild error path. +#[derive(Debug)] +struct CountingStorage { + inner: MemoryStorage, + list_files_calls: AtomicUsize, + fail_list_files: std::sync::atomic::AtomicBool, +} + +impl CountingStorage { + fn new() -> Self { + Self { + inner: MemoryStorage::new(MemoryStorageConfig::default()), + list_files_calls: AtomicUsize::new(0), + fail_list_files: std::sync::atomic::AtomicBool::new(false), + } + } + + fn list_files_count(&self) -> usize { + self.list_files_calls.load(Ordering::SeqCst) + } + + fn set_fail_list_files(&self, fail: bool) { + self.fail_list_files.store(fail, Ordering::SeqCst); + } +} + +impl Storage for CountingStorage { + fn loading_mode(&self) -> LoadingMode { + self.inner.loading_mode() + } + fn open_input(&self, name: &str) -> Result> { + self.inner.open_input(name) + } + fn create_output(&self, name: &str) -> Result> { + self.inner.create_output(name) + } + fn create_output_append(&self, name: &str) -> Result> { + self.inner.create_output_append(name) + } + fn file_exists(&self, name: &str) -> bool { + self.inner.file_exists(name) + } + fn delete_file(&self, name: &str) -> Result<()> { + self.inner.delete_file(name) + } + fn list_files(&self) -> Result> { + if self.fail_list_files.load(Ordering::SeqCst) { + return Err(laurus::LaurusError::storage("injected list_files failure")); + } + self.list_files_calls.fetch_add(1, Ordering::SeqCst); + self.inner.list_files() + } + fn file_size(&self, name: &str) -> Result { + self.inner.file_size(name) + } + fn metadata(&self, name: &str) -> Result { + self.inner.metadata(name) + } + fn rename_file(&self, old_name: &str, new_name: &str) -> Result<()> { + self.inner.rename_file(old_name, new_name) + } + fn create_temp_output(&self, prefix: &str) -> Result<(String, Box)> { + self.inner.create_temp_output(prefix) + } + fn sync(&self) -> Result<()> { + self.inner.sync() + } + fn close(&mut self) -> Result<()> { + self.inner.close() + } +} + +fn doc(title: &str) -> Document { + Document::builder().add_text("title", title).build() +} + +/// Write one committed segment holding `ids` onto `storage`. +fn seed_committed_segment(storage: &Arc, ids: &[u64]) { + let mut writer = + InvertedIndexWriter::new(storage.clone(), InvertedIndexWriterConfig::default()).unwrap(); + for &id in ids { + writer + .upsert_document(id, doc(&format!("seed{id}"))) + .unwrap(); + } + writer.commit().unwrap(); +} + +/// #559: fresh-id upserts must not rescan segment metadata — the constructor's +/// single recovery scan is the only `list_files` a pure-ingest writer pays. +#[test] +fn fresh_id_upserts_never_rescan_meta_files() { + let counting = Arc::new(CountingStorage::new()); + let storage: Arc = counting.clone(); + seed_committed_segment(&storage, &[1, 2, 3]); + + let before_construction = counting.list_files_count(); + let mut writer = + InvertedIndexWriter::new(storage.clone(), InvertedIndexWriterConfig::default()).unwrap(); + assert_eq!( + counting.list_files_count(), + before_construction + 1, + "the constructor performs exactly one recovery scan" + ); + + let after_construction = counting.list_files_count(); + for id in 10..60u64 { + writer.upsert_document(id, doc(&format!("t{id}"))).unwrap(); + } + assert_eq!( + counting.list_files_count(), + after_construction, + "50 fresh-id upserts must not list the storage at all \ + (pre-#864 behavior: one full .meta list+parse per upsert)" + ); +} + +/// #559 + #571: overwriting committed docs resolves the segment from the +/// cache (no rescan) and reuses one `DeletionManager` across calls — the +/// first overwrite pays the manager's single bitmap-loading scan, subsequent +/// overwrites pay zero. The deletion itself must still land (`.delmap` +/// exists). +#[test] +fn overwrite_committed_docs_reuses_cache_and_manager() { + let counting = Arc::new(CountingStorage::new()); + let storage: Arc = counting.clone(); + seed_committed_segment(&storage, &[1, 2, 3]); + + let mut writer = + InvertedIndexWriter::new(storage.clone(), InvertedIndexWriterConfig::default()).unwrap(); + let after_construction = counting.list_files_count(); + + // First overwrite: the lazily created DeletionManager loads existing + // bitmaps once (one list_files); the segment lookup itself is cached. + writer.upsert_document(1, doc("one-v2")).unwrap(); + assert_eq!( + counting.list_files_count(), + after_construction + 1, + "the first overwrite pays exactly the DeletionManager's one-time \ + bitmap-loading scan" + ); + + // Subsequent overwrites: fully served from memory. + writer.upsert_document(2, doc("two-v2")).unwrap(); + writer.upsert_document(3, doc("three-v2")).unwrap(); + assert_eq!( + counting.list_files_count(), + after_construction + 1, + "subsequent overwrites must reuse the cached manager and ranges \ + (pre-#864: fresh manager + full .delmap reload per overwrite)" + ); + + // The deletions actually landed: the seeded segment now has a bitmap. + let delmaps: Vec = storage + .list_files() + .unwrap() + .into_iter() + .filter(|f| f.ends_with(".delmap")) + .collect(); + assert_eq!( + delmaps.len(), + 1, + "the overwrites must have marked deletions in the seeded segment: {delmaps:?}" + ); +} + +/// #864: a segment flushed mid-life by this writer (auto-flush at +/// `max_buffered_docs`) extends the cache in place, so a later overwrite of +/// one of its docs marks the deletion without any rescan. +#[test] +fn mid_life_flush_extends_cache() { + let counting = Arc::new(CountingStorage::new()); + let storage: Arc = counting.clone(); + + let config = InvertedIndexWriterConfig { + max_buffered_docs: 2, // force an auto-flush after two upserts + ..Default::default() + }; + let mut writer = InvertedIndexWriter::new(storage.clone(), config).unwrap(); + + writer.upsert_document(1, doc("alpha")).unwrap(); + writer.upsert_document(2, doc("bravo")).unwrap(); // buffer full -> flush_segment + + // Everything below must be answered from the in-place extended cache. + // (write_segment_files itself lists once to enumerate the files it wrote; + // snapshot after the flush.) + let after_flush = counting.list_files_count(); + + // Overwrite a doc that lives in the segment flushed above. The segment + // lookup must hit the extended cache; only the lazy DeletionManager + // construction may scan (once). + writer.upsert_document(1, doc("alpha-v2")).unwrap(); + assert_eq!( + counting.list_files_count(), + after_flush + 1, + "the overwrite must resolve the mid-life flushed segment from the \ + cache, paying only the one-time DeletionManager construction" + ); + + let delmaps: Vec = storage + .list_files() + .unwrap() + .into_iter() + .filter(|f| f.ends_with(".delmap")) + .collect(); + assert_eq!( + delmaps.len(), + 1, + "the overwrite must have marked a deletion in the flushed segment: {delmaps:?}" + ); +} + +/// #864: `LexicalStore::optimize` force-merges every segment behind a live +/// writer; the invalidation hook must rebuild the writer's cached ranges so a +/// later overwrite marks its deletion in the **merged** segment — with a +/// stale cache it would target a deleted ghost segment and the old version +/// would resurface in search. +#[test] +fn optimize_rebuilds_live_writer_cache() { + let storage: Arc = Arc::new(MemoryStorage::new(MemoryStorageConfig::default())); + let store = LexicalStore::new(storage.clone(), LexicalIndexConfig::default()).unwrap(); + + let hits = |field: &str, term: &str| -> usize { + let query = Box::new(TermQuery::new(field, term)); + store + .search(LexicalSearchRequest::new(query)) + .unwrap() + .hits + .len() + }; + + // Two commits -> two segments. + store.upsert_document(1, doc("alpha")).unwrap(); + store.commit().unwrap(); + store.upsert_document(2, doc("bravo")).unwrap(); + store.commit().unwrap(); + + // A live writer with a buffered fresh doc, created BEFORE the merge. + store.upsert_document(3, doc("charlie")).unwrap(); + + // Force-merge both committed segments into one behind the live writer. + store.optimize().unwrap(); + + // Overwrite doc 1 (now living in the merged segment) through the SAME + // live writer, then commit everything. (Term is hyphen-free on purpose: + // the tokenizer would split "alpha-v2" into "alpha" + "v2", making the + // replacement itself match `title:alpha`.) + store.upsert_document(1, doc("alphav2")).unwrap(); + store.commit().unwrap(); + + assert_eq!( + hits("title", "alpha"), + 0, + "the pre-merge version must be dead — a stale segment cache would \ + have marked the deletion in a ghost segment and left it alive" + ); + assert_eq!(hits("title", "alphav2"), 1, "the overwrite must be live"); + assert_eq!( + hits("title", "bravo"), + 1, + "untouched doc survives the merge" + ); + assert_eq!( + hits("title", "charlie"), + 1, + "buffered doc survives the merge" + ); + + // The deletion bitmap must belong to the merged segment, not a ghost. + let delmaps: Vec = storage + .list_files() + .unwrap() + .into_iter() + .filter(|f| f.ends_with(".delmap")) + .collect(); + assert!( + delmaps.iter().all(|f| f.starts_with("merged_")), + "deletions must land in the merged segment only: {delmaps:?}" + ); +} + +/// #864 review follow-up: `invalidate_segment_cache` must be atomic — when +/// the rebuild scan fails, the OLD (still-valid) cache must survive intact. +/// A clear-then-fail implementation would leave the writer with an empty +/// cache and `max_committed_doc_id == 0`, silently skipping every subsequent +/// overwrite's deletion via the fast path. +#[test] +fn invalidate_failure_preserves_old_cache() { + let counting = Arc::new(CountingStorage::new()); + let storage: Arc = counting.clone(); + seed_committed_segment(&storage, &[1, 2, 3]); + + let mut writer = + InvertedIndexWriter::new(storage.clone(), InvertedIndexWriterConfig::default()).unwrap(); + + // Rebuild fails: the error must propagate... + counting.set_fail_list_files(true); + writer + .invalidate_segment_cache() + .expect_err("the injected list_files failure must propagate"); + counting.set_fail_list_files(false); + + // ...but the previous cache must still be in force: the segments were + // not actually replaced (no merge ran), so an overwrite must still + // resolve the seeded segment and mark its deletion. + writer.upsert_document(1, doc("one-v2")).unwrap(); + let delmaps: Vec = storage + .list_files() + .unwrap() + .into_iter() + .filter(|f| f.ends_with(".delmap")) + .collect(); + assert_eq!( + delmaps.len(), + 1, + "the overwrite must still mark the deletion from the preserved \ + cache — an emptied cache would have skipped it silently: {delmaps:?}" + ); +} diff --git a/laurus/tests/vector_hnsw_soft_delete_test.rs b/laurus/tests/vector_hnsw_soft_delete_test.rs index 8efde8a0..f7317062 100644 --- a/laurus/tests/vector_hnsw_soft_delete_test.rs +++ b/laurus/tests/vector_hnsw_soft_delete_test.rs @@ -204,7 +204,7 @@ async fn optimize_reclaims_deleted_documents() { // Compaction physically rebuilds the graph without the deleted docs and // drops the bitmap file. - store.optimize().unwrap(); + store.optimize().await.unwrap(); assert!( !storage.file_exists(DELMAP_FILE), "optimize() must remove the .delmap after purging" @@ -258,7 +258,7 @@ async fn deleting_nearest_still_finds_next_nearest() { assert!(before.iter().all(|id| (40..N).contains(id))); // The same invariant must hold after compaction. - store.optimize().unwrap(); + store.optimize().await.unwrap(); assert!(!storage.file_exists(DELMAP_FILE)); let after = hit_ids(&store.search(request(10)).unwrap()); assert_eq!(after.len(), 10); diff --git a/laurus/tests/vector_writer_retention_test.rs b/laurus/tests/vector_writer_retention_test.rs new file mode 100644 index 00000000..8e104675 --- /dev/null +++ b/laurus/tests/vector_writer_retention_test.rs @@ -0,0 +1,597 @@ +//! Integration tests for retaining the HNSW writer across commits +//! (Issue #864 / #572). +//! +//! Before the fix, `VectorStore::commit` `take()`d the cached writer, so the +//! first upsert after every commit reconstructed it via +//! `HnswIndexWriter::with_storage` — reloading and dequantizing the entire +//! `.hnsw` file. Retention keeps the committed writer (whose in-memory state +//! equals the file it just wrote) in the cache. +//! +//! The two safety valves are the corruption tests: compaction +//! (`maybe_auto_compact` inside commit) and `VectorStore::optimize` rewrite +//! the index through a **fresh** writer and clear the deletion bitmap, so a +//! stale retained writer would resurrect the physically reclaimed vectors on +//! its next commit. Both paths must invalidate the cache. +//! +//! The primary gate is deterministic: a `CountingStorage` decorator counts +//! `open_input(".hnsw")` calls, so the no-reload claim is asserted exactly. + +use async_trait::async_trait; +use std::any::Any; +use std::collections::HashMap; +use std::collections::HashSet; +use std::sync::Arc; + +use laurus::lexical::LexicalIndexConfig; +use laurus::storage::memory::{MemoryStorage, MemoryStorageConfig}; +use laurus::storage::{FileMetadata, LoadingMode, Storage, StorageInput, StorageOutput}; +use laurus::vector::Vector; +use laurus::vector::core::distance::DistanceMetric; +use laurus::vector::core::field::HnswOption; +use laurus::vector::store::config::VectorFieldConfig; +use laurus::vector::store::request::{ + QueryVector, VectorScoreMode, VectorSearchParams, VectorSearchRequest, +}; +use laurus::vector::{FieldOption, VectorIndexConfig, VectorSearchQuery}; +use laurus::{DataValue, Document}; +use laurus::{EmbedInput, EmbedInputType, Embedder}; +use laurus::{LaurusError, Result}; + +const DIM: usize = 16; +const STEP: f32 = 0.01; +const HNSW_FILE: &str = "vector_index.hnsw"; +const DELMAP_FILE: &str = "vector_index.delmap"; + +#[derive(Debug)] +struct MockEmbedder { + dimension: usize, +} + +#[async_trait] +impl Embedder for MockEmbedder { + async fn embed(&self, input: &EmbedInput<'_>) -> Result { + match input { + EmbedInput::Text(_) => Ok(Vector::new(vec![0.0; self.dimension])), + _ => Err(LaurusError::invalid_argument("text only")), + } + } + fn supported_input_types(&self) -> Vec { + vec![EmbedInputType::Text] + } + fn name(&self) -> &str { + "mock" + } + fn as_any(&self) -> &dyn Any { + self + } +} + +/// Decorator over [`MemoryStorage`] counting `open_input` / `create_output` +/// calls per file — the deterministic signals for "did this operation reload +/// the `.hnsw`?" and "did this commit rewrite the `.hnsw`?" — plus an +/// optional fault injector failing `create_output` for matching names, to +/// exercise the commit ladder's error paths. +#[derive(Debug)] +struct CountingStorage { + inner: MemoryStorage, + opens: std::sync::Mutex>, + creates: std::sync::Mutex>, + fail_create_matching: std::sync::Mutex>, +} + +impl CountingStorage { + fn new() -> Self { + Self { + inner: MemoryStorage::new(MemoryStorageConfig::default()), + opens: std::sync::Mutex::new(HashMap::new()), + creates: std::sync::Mutex::new(HashMap::new()), + fail_create_matching: std::sync::Mutex::new(None), + } + } + + fn open_count(&self, name: &str) -> usize { + self.opens.lock().unwrap().get(name).copied().unwrap_or(0) + } + + /// Total `create_output` calls whose file name contains `pat`. + fn create_count_matching(&self, pat: &str) -> usize { + self.creates + .lock() + .unwrap() + .iter() + .filter(|(name, _)| name.contains(pat)) + .map(|(_, n)| n) + .sum() + } + + /// Arm (`Some(substring)`) or disarm (`None`) the `create_output` fault. + fn set_fail_create_matching(&self, pat: Option<&str>) { + *self.fail_create_matching.lock().unwrap() = pat.map(String::from); + } +} + +impl Storage for CountingStorage { + fn loading_mode(&self) -> LoadingMode { + self.inner.loading_mode() + } + fn open_input(&self, name: &str) -> Result> { + *self + .opens + .lock() + .unwrap() + .entry(name.to_string()) + .or_insert(0) += 1; + self.inner.open_input(name) + } + fn create_output(&self, name: &str) -> Result> { + if let Some(pat) = self.fail_create_matching.lock().unwrap().as_deref() + && name.contains(pat) + { + return Err(LaurusError::storage(format!( + "injected create_output failure for '{name}'" + ))); + } + *self + .creates + .lock() + .unwrap() + .entry(name.to_string()) + .or_insert(0) += 1; + self.inner.create_output(name) + } + fn create_output_append(&self, name: &str) -> Result> { + self.inner.create_output_append(name) + } + fn file_exists(&self, name: &str) -> bool { + self.inner.file_exists(name) + } + fn delete_file(&self, name: &str) -> Result<()> { + self.inner.delete_file(name) + } + fn list_files(&self) -> Result> { + self.inner.list_files() + } + fn file_size(&self, name: &str) -> Result { + self.inner.file_size(name) + } + fn metadata(&self, name: &str) -> Result { + self.inner.metadata(name) + } + fn rename_file(&self, old_name: &str, new_name: &str) -> Result<()> { + self.inner.rename_file(old_name, new_name) + } + fn create_temp_output(&self, prefix: &str) -> Result<(String, Box)> { + self.inner.create_temp_output(prefix) + } + fn sync(&self) -> Result<()> { + self.inner.sync() + } + fn close(&mut self) -> Result<()> { + self.inner.close() + } +} + +fn doc_vec(i: u64) -> Vec { + let theta = i as f32 * STEP; + let mut v = vec![0.0; DIM]; + v[0] = theta.cos(); + v[1] = theta.sin(); + v +} + +fn vec_doc(i: u64) -> Document { + Document::builder() + .add_field("vec", DataValue::Vector(doc_vec(i))) + .build() +} + +fn hnsw() -> FieldOption { + FieldOption::Hnsw(HnswOption { + dimension: DIM, + distance: DistanceMetric::Cosine, + m: 16, + ef_construction: 100, + // Near-exhaustive search for the ≤101 vectors used here, so the + // exact-count assertions are not subject to HNSW recall noise. + default_ef_search: Some(400), + base_weight: 1.0, + quantizer: Default::default(), + rerank_storage: None, + embedder: None, + }) +} + +fn make_config(auto_compaction: bool, threshold: f64) -> VectorIndexConfig { + let mut field_configs = HashMap::new(); + field_configs.insert( + "vec".to_string(), + VectorFieldConfig { + vector: Some(hnsw()), + lexical: None, + }, + ); + VectorIndexConfig { + fields: field_configs, + embedder: Arc::new(MockEmbedder { dimension: DIM }), + default_fields: vec!["vec".to_string()], + metadata: HashMap::new(), + deletion_config: laurus::DeletionConfig { + auto_compaction, + compaction_threshold: threshold, + ..Default::default() + }, + shard_id: 0, + metadata_config: LexicalIndexConfig::default(), + } +} + +fn request(limit: usize) -> VectorSearchRequest { + let mut query = vec![0.0; DIM]; + query[0] = 1.0; + VectorSearchRequest { + query: VectorSearchQuery::Vectors(vec![QueryVector { + vector: Vector::new(query), + weight: 1.0, + fields: Some(vec!["vec".into()]), + }]), + params: VectorSearchParams { + limit, + score_mode: VectorScoreMode::WeightedSum, + fields: None, + allowed_ids: None, + ..Default::default() + }, + } +} + +fn hit_ids(store: &laurus::vector::VectorStore, limit: usize) -> HashSet { + store + .search(request(limit)) + .unwrap() + .hits + .iter() + .map(|h| h.doc_id) + .collect() +} + +/// #864: with retention, an upsert arriving after a commit must not re-open +/// the `.hnsw` file — pre-fix it reloaded (and dequantized) the whole index. +#[tokio::test(flavor = "multi_thread")] +async fn upsert_after_commit_does_not_reload_hnsw() { + let counting = Arc::new(CountingStorage::new()); + let storage: Arc = counting.clone(); + let store = laurus::vector::VectorStore::new(storage, make_config(false, 0.5)).unwrap(); + + for id in 0..10u64 { + store + .upsert_document_by_internal_id(id, vec_doc(id)) + .await + .unwrap(); + } + store.commit().await.unwrap(); + + let opens_after_commit = counting.open_count(HNSW_FILE); + store + .upsert_document_by_internal_id(10, vec_doc(10)) + .await + .unwrap(); + assert_eq!( + counting.open_count(HNSW_FILE), + opens_after_commit, + "the retained writer must serve the post-commit upsert without \ + re-opening the .hnsw (pre-#864: full reload per commit cycle)" + ); + + // And the follow-up commit persists correctly. + store.commit().await.unwrap(); + assert_eq!( + counting.open_count(HNSW_FILE), + opens_after_commit, + "the second commit writes from the retained writer; no reload either" + ); +} + +/// #864 correctness: three upsert+commit cycles through the retained writer +/// preserve every record, both for the live store and after a cold reopen +/// (proving the retained writer kept writing complete, valid files). +/// +/// Membership is asserted via `stats().document_count` (record-level, exact) +/// rather than exact search hit sets: incremental HNSW appends can +/// nondeterministically leave nodes unreachable on layer 0 — a pre-existing +/// bug reproduced on unmodified main, tracked as #868 — so search-based +/// full-recall assertions would flake for reasons unrelated to retention. +#[tokio::test(flavor = "multi_thread")] +async fn retained_writer_commit_cycles_preserve_all_vectors() { + let counting = Arc::new(CountingStorage::new()); + let storage: Arc = counting.clone(); + let store = laurus::vector::VectorStore::new(storage.clone(), make_config(false, 0.5)).unwrap(); + + for cycle in 0..3u64 { + for i in 0..10u64 { + let id = cycle * 10 + i; + store + .upsert_document_by_internal_id(id, vec_doc(id)) + .await + .unwrap(); + } + store.commit().await.unwrap(); + } + + let expected: HashSet = (0..30).collect(); + assert_eq!( + store.stats().unwrap().document_count, + 30, + "all 30 records across 3 retained-writer commit cycles must persist \ + (no loss, no duplication)" + ); + let ids = hit_ids(&store, 30); + assert!( + ids.is_subset(&expected) && !ids.is_empty(), + "search must return only the ingested ids: {ids:?}" + ); + + // Cold reopen on the same storage: the files must be self-sufficient. + drop(store); + let reopened = laurus::vector::VectorStore::new(storage, make_config(false, 0.5)).unwrap(); + assert_eq!( + reopened.stats().unwrap().document_count, + 30, + "a cold reopen must see the same 30 records" + ); + let ids = hit_ids(&reopened, 30); + assert!( + ids.is_subset(&expected) && !ids.is_empty(), + "post-reopen search must return only the ingested ids: {ids:?}" + ); +} + +/// #864 corruption test: auto-compaction inside commit rewrites the `.hnsw` +/// through a fresh writer and clears the delmap. The retained writer MUST be +/// invalidated when compaction ran — otherwise its next commit resurrects +/// the physically reclaimed vectors with no deletion bitmap marking them. +#[tokio::test(flavor = "multi_thread")] +async fn auto_compaction_invalidates_retained_writer_no_resurrection() { + let counting = Arc::new(CountingStorage::new()); + let storage: Arc = counting.clone(); + // 30% threshold, auto-compaction ON. + let store = laurus::vector::VectorStore::new(storage.clone(), make_config(true, 0.3)).unwrap(); + + for id in 0..100u64 { + store + .upsert_document_by_internal_id(id, vec_doc(id)) + .await + .unwrap(); + } + store.commit().await.unwrap(); + + // Delete 40 of 100 (40% >= 30%): this commit runs compaction. + let deleted: HashSet = (0..40).collect(); + for &id in &deleted { + store.delete_document_by_internal_id(id).await.unwrap(); + } + store.commit().await.unwrap(); + assert!( + !storage.file_exists(DELMAP_FILE), + "precondition: compaction must have run and removed the .delmap" + ); + + // The dangerous cycle: upsert + commit AFTER compaction. A stale + // retained writer would write the 40 reclaimed vectors back. + store + .upsert_document_by_internal_id(100, vec_doc(100)) + .await + .unwrap(); + store.commit().await.unwrap(); + + // Deterministic resurrection gate: a stale writer would rewrite the 40 + // reclaimed records, bouncing the record count from 61 back to 101. + // (Search hit sets are not asserted exactly because of the pre-existing + // incremental-append reachability bug #868.) + assert_eq!( + store.stats().unwrap().document_count, + 61, + "exactly the 60 survivors + the new doc may exist on disk — more \ + means the stale retained writer resurrected compacted-away records" + ); + let ids = hit_ids(&store, 100); + assert!( + ids.is_disjoint(&deleted), + "compacted-away docs must NOT resurface in search after a \ + post-compaction commit from the (invalidated) writer cache: {ids:?}" + ); +} + +/// #864 corruption test, explicit-`optimize` variant: `VectorStore::optimize` +/// physically reclaims soft-deleted docs through a fresh writer and clears +/// the delmap; the retained writer cache must be dropped there too. +#[tokio::test(flavor = "multi_thread")] +async fn store_optimize_invalidates_retained_writer() { + let counting = Arc::new(CountingStorage::new()); + let storage: Arc = counting.clone(); + // Auto-compaction OFF so only the explicit optimize() reclaims. + let store = laurus::vector::VectorStore::new(storage.clone(), make_config(false, 0.9)).unwrap(); + + for id in 0..100u64 { + store + .upsert_document_by_internal_id(id, vec_doc(id)) + .await + .unwrap(); + } + store.commit().await.unwrap(); + + let deleted: HashSet = (0..40).collect(); + for &id in &deleted { + store.delete_document_by_internal_id(id).await.unwrap(); + } + store.commit().await.unwrap(); + + store.optimize().await.unwrap(); + assert!( + !storage.file_exists(DELMAP_FILE), + "precondition: optimize must have reclaimed and removed the .delmap" + ); + + store + .upsert_document_by_internal_id(100, vec_doc(100)) + .await + .unwrap(); + store.commit().await.unwrap(); + + // Same deterministic gate as the auto-compaction variant (see #868 for + // why exact search hit sets are not asserted). + assert_eq!( + store.stats().unwrap().document_count, + 61, + "exactly the 60 survivors + the new doc may exist on disk — more \ + means the stale retained writer resurrected optimized-away records" + ); + let ids = hit_ids(&store, 100); + assert!( + ids.is_disjoint(&deleted), + "optimized-away docs must NOT resurface in search after a \ + post-optimize commit from the (invalidated) writer cache: {ids:?}" + ); +} + +/// #864 review follow-up: a retained writer with **no pending changes** must +/// not rewrite the `.hnsw` on a no-op commit — `has_pending_changes()` gates +/// the flush, so back-to-back commits pay zero index writes. +#[tokio::test(flavor = "multi_thread")] +async fn noop_commit_does_not_rewrite_hnsw() { + let counting = Arc::new(CountingStorage::new()); + let storage: Arc = counting.clone(); + let store = laurus::vector::VectorStore::new(storage, make_config(false, 0.5)).unwrap(); + + for id in 0..10u64 { + store + .upsert_document_by_internal_id(id, vec_doc(id)) + .await + .unwrap(); + } + store.commit().await.unwrap(); + + let writes_after_commit = counting.create_count_matching(".hnsw"); + assert!( + writes_after_commit > 0, + "the first commit must write the index" + ); + + // Two commits with zero changes: the retained (finalized) writer must be + // skipped, not re-committed into an identical full rewrite. + store.commit().await.unwrap(); + store.commit().await.unwrap(); + assert_eq!( + counting.create_count_matching(".hnsw"), + writes_after_commit, + "no-change commits must not rewrite the .hnsw from the retained writer" + ); + + // A real change still triggers exactly one more write cycle. + store + .upsert_document_by_internal_id(10, vec_doc(10)) + .await + .unwrap(); + store.commit().await.unwrap(); + assert!( + counting.create_count_matching(".hnsw") > writes_after_commit, + "a commit with pending changes must write the index again" + ); + assert_eq!(store.stats().unwrap().document_count, 11); +} + +/// #864 review follow-up: `optimize()` must flush a dirty writer whose buffer +/// was **emptied by deletions** — `pending_docs() == 0` there, but dropping +/// it would silently discard the uncommitted delete-everything mutation and +/// leave the stale vector on disk. +#[tokio::test(flavor = "multi_thread")] +async fn optimize_flushes_writer_emptied_by_deletions() { + let counting = Arc::new(CountingStorage::new()); + let storage: Arc = counting.clone(); + let store = laurus::vector::VectorStore::new(storage, make_config(false, 0.9)).unwrap(); + + store + .upsert_document_by_internal_id(0, vec_doc(0)) + .await + .unwrap(); + store.commit().await.unwrap(); + assert_eq!(store.stats().unwrap().document_count, 1); + + // Re-upsert the same internal id with a doc carrying NO embeddable + // fields (an integer is neither a vector nor embeddable text/bytes): + // the writer buffer-deletes the old vector and adds nothing, leaving a + // dirty writer with pending_docs() == 0. + let empty_doc = Document::builder() + .add_field("note", DataValue::Int64(42)) + .build(); + store + .upsert_document_by_internal_id(0, empty_doc) + .await + .unwrap(); + + store.optimize().await.unwrap(); + + assert_eq!( + store.stats().unwrap().document_count, + 0, + "optimize must flush the emptied writer's uncommitted deletion \ + instead of dropping it and resurrecting the stale vector" + ); +} + +/// #864 review follow-up: a commit that fails mid-ladder must DROP the +/// retained writer (the pre-retention behavior on every path) — the +/// writer/disk agreement is unknown after a partial failure, so the next +/// upsert must reload ground truth from storage. +#[tokio::test(flavor = "multi_thread")] +async fn failed_commit_drops_retained_writer() { + let counting = Arc::new(CountingStorage::new()); + let storage: Arc = counting.clone(); + let store = laurus::vector::VectorStore::new(storage, make_config(false, 0.5)).unwrap(); + + for id in 0..5u64 { + store + .upsert_document_by_internal_id(id, vec_doc(id)) + .await + .unwrap(); + } + store.commit().await.unwrap(); + + // Make the writer dirty, then fail its flush inside the next commit. + store + .upsert_document_by_internal_id(5, vec_doc(5)) + .await + .unwrap(); + counting.set_fail_create_matching(Some(".hnsw")); + store + .commit() + .await + .expect_err("the injected create_output failure must fail the commit"); + counting.set_fail_create_matching(None); + + // The failed writer must be gone: the next upsert reloads the intact + // on-disk index (observable as a fresh .hnsw open) instead of reusing a + // writer in an unknown state. + let opens_before = counting.open_count(HNSW_FILE); + store + .upsert_document_by_internal_id(6, vec_doc(6)) + .await + .unwrap(); + assert!( + counting.open_count(HNSW_FILE) > opens_before, + "after a failed commit the cache must be empty, forcing a reload" + ); + + // And the store recovers fully: doc 5 was lost with the failed writer + // (it is WAL-replayable at engine level), docs 0-4 + 6 are intact. + store.commit().await.unwrap(); + let ids = hit_ids(&store, 10); + assert!( + ids.contains(&6), + "post-recovery upsert must be live: {ids:?}" + ); + assert_eq!( + store.stats().unwrap().document_count, + 6, + "the intact on-disk docs (0-4) plus the post-recovery doc (6)" + ); +}