diff --git a/CHANGELOG.md b/CHANGELOG.md index f53873a2..c17727d5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,173 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added — int8 symmetric ADC for SQ8 vector search (task #13) + +- New per-candidate integer dot-product path for SQ8 asymmetric distance + computation (ADC), replacing the f32-widening `sq8_stats` kernel on the + hot per-candidate pass with exact-integer int8 dot products: NEON + widen-multiply (`vmull_s8` + the `c^0x80` offset trick, since baseline + ARMv8-A has no mixed u8×i8 widening multiply), AVX2 (`cvtepu8/cvtepi8_epi16` + → `mullo_epi16` → `madd_epi16`, saturation-free), and AVX512 VNNI + (`_mm512_dpbusd_epi32`, feature-gated behind `simd-avx512`). Query is + quantized to symmetric int8 once per search + (`turbo_quant::sq8::sq8_quantize_query_scalar`); `sum_c`/`sumsq_c` are + exact integers straight from the u8 codes. Combines through the unchanged + `sq8_l2_from_stats` — the on-disk SQ8 format and the exact-rerank sidecar + (HQ-1) are untouched. Wired into both `hnsw/search.rs` beam-search closures + and all three SQ8 call sites in `segment/mutable.rs` brute-force search. + Local (dev-Mac, NEON) directional speedup vs the existing f32 SIMD path: + ~1.5x at dim 128, ~2.0x at dim 384, ~2.15x at dim 768 (`benches/sq8_adc_bench.rs`, + `int8_dispatch` variant) — absolute/cross-arch numbers are GCE-validation + scope (task #14). Recall A/B-gated (`turbo_quant::sq8` test module): + R@10 delta ≤ 0.01 vs the f32 ADC path and ≥ 0.98 top-10 overlap, both L2 + and Cosine, at dim 128/384. `MOON_SQ8_INT8_ADC=0` forces the f32 kernels + (bench/diagnostic escape hatch, not yet added to CLAUDE.md's env list — + flagged for a follow-up doc pass). + +### Fixed — merge recall gate self-exclusion bias (manual merges could never pass) + +- `verify_merge_recall` compared HNSW results **including the query point** + (every sampled query is a database point, distance 0 → always rank 1) + against a ground truth that **excluded** it, structurally capping measurable + recall at (k−1)/k = 0.90 — exactly the manual `FT.COMPACT`/`VACUUM VECTOR` + merge gate, so any manual merge of an index with ≥ 50 vectors was rejected + with `merge recall 0.9000 < tolerance 0.9000` regardless of actual graph + quality (the 0.70 background gate masked this). The merged-graph search now + requests k+1 and drops the self-point before comparison. + +### Added — vector-index durability: kill-9 crash-recovery test suite (B4) + +- New `tests/crash_recovery_vector_durability.rs`: five end-to-end scenarios + against real server processes (SIGKILL + restart): unchanged-key dedup + fast path, update/delete reconcile, orphan sweep on boot, collection_id + pin surviving a post-recovery compact + GraphUnion merge, and a + no-persistence regression guard. All waits are bounded condition polls; + servers run with `--disk-free-min-pct 0` so the suite is immune to the + dev volume hovering at the 5% diskfull guard. + +### Added — vector-index durability: startup recovery from disk (B1-B3) + +- **Vector indexes now persist their segments across restarts** instead of + paying a full re-index (TQ encode + HNSW build) of every matching hash key + on every restart. Each index gets an `idx-/` directory holding + an atomically-written `manifest.json` (collection_id / segment ids / + id-allocator floors), a checksummed `keymap-.bin` (key_hash → + global_id + vector checksum + original key), and its immutable HNSW + segments (staged write: `staging-` → fsync → atomic rename), written + in the background after each compact/merge install (B1/B2). +- **Startup now loads that state instead of discarding it**: segments are + read back and reattached (pinning the recreated index's `collection_id` to + the persisted value — required for the HNSW QJL rotation seed to match), + and the keyspace rescan is now a **dedup rescan**: a key whose vector + bytes checksum-match the last snapshot is rebuilt as metadata-only (no + HNSW/TQ re-encode); only genuinely changed or unknown keys are fully + re-indexed; keys removed from the keyspace are tombstoned (B3). +- **Crash-safety contract**: any corrupt/missing artifact (segment, + checksum mismatch, unreadable headers) degrades to a rescan-rebuild of + exactly the affected keys, never to wrong search results; a manifest may + understate what's durable (costs a rescan) but never overstates. Startup + also sweeps orphaned segment/staging/keymap files and stale `idx-*` + directories left by an interrupted drop. +- Multi-field indexes: only the default vector field's segments/checksum + are persisted; additional named vector fields always re-encode on + restart (documented gap, safe/conservative). +- Removed the vestigial WAL-replay vector recovery path + (`src/vector/persistence/recovery.rs`, `VectorStore::attach_recovered`/ + `pending_segments`): it only ever populated a `VectorStore` that was + discarded before the shard event loop started (recovery authority is now + exclusively the manifest/segment/keymap layout above). + +### Fixed — FLUSHALL/FLUSHDB ghost vectors + HDEL stale-vector gap + +- **FLUSHALL/FLUSHDB never touched the FT indexes** (persistence-review R3): + flushed hashes stayed searchable as ghost vectors/documents until restart. + Both commands now clear every vector + text index's CONTENTS (segments, + key-hash maps, postings, TAG/NUMERIC indexes) while KEEPING the FT.CREATE + definitions — matching restart semantics, where definitions come from the + sidecar and contents are re-derived from the (now empty) keyspace. + Recovered-but-unclaimed `pending_segments` are discarded too, so a + post-flush FT.CREATE cannot resurrect pre-flush contents. +- **`HDEL key ` left the vector searchable** (R4): only + whole-key DEL/UNLINK tombstoned. A successful HDEL that removes an + index's vector field now tombstones that key in exactly the affected + indexes (per-index `mark_deleted_for_key_in_index`; sibling indexes keyed + on other fields keep their entries). Known follow-ups documented in + `auto_hdel_vectors`: multi-vector-field indexes tombstone the whole doc, + and TEXT/TAG/NUMERIC field removal is not yet re-indexed. +- Both hooks wired with wire-parity on ALL dispatch paths (single-shard, + monoio conn-local, tokio sharded, SPSC execute, shared batch, MULTI/EXEC). + New integration suite `tests/vector_flush_hdel_tombstone.rs` (red→green). + +### Observability — exact-rerank sidecar coverage is visible and loud (R5) + +- A GraphUnion merge with even one sidecar-less source segment silently + dropped the exact-rerank f16 sidecar for the ENTIRE merged segment + (all-or-nothing propagation — a partial sidecar would mix exact and ADC + distances within one segment). The drop now logs a `tracing::warn` with + source counts, and FT.INFO exposes additive `graph_segments` / + `segments_with_exact_rerank` counters (merged across shards) so ADC-only + segments are observable in the steady state. + +### Performance — AE-1: saturation-gated per-segment adaptive ef + +- With G graph segments, every segment searched at the FULL resolved ef — + ~G× the CPU of one ef-wide beam (the root cause of the pooled regression + on SMT-constrained boxes in PR #214's GCE run). Compact/GraphUnion builds + now run a one-time self-probe against the exact f16-sidecar ground truth + (leave-self-out R@10 across an ef ladder, `estimate_suggested_ef`): a + segment whose curve is FULLY SATURATED from the minimum rung (flat at + ≈1.0) is certified trivially easy and searches at min-ef (24); every + other segment keeps the full resolved ef. Applied identically to the + sync, serial-yielding, and worker-pool paths (pooled == serial identity + holds); never active when the user pins `EF_RUNTIME`; in-memory only + (reloaded segments fall back to the full beam). +- Two weaker designs were tried and REJECTED by recall A/B (20k×384d SQ8, + 5 segments): a blanket `ef/√G` split (EF-SPLIT) silently cost gaussian + R@10 0.9915 → 0.9295 (its original "recall preserved" validation had + exercised a stale binary), and letting the self-probe pick a mid-ladder + knee cost 0.9915 → 0.939 — self-sampled queries are optimistically + biased on unstructured data, so the probe's ONLY trustworthy verdict is + "saturated at min-ef". +- Recall/latency gate (same seeds/harness end-to-end): clustered 5-seg + R@10 1.0000 at p50 0.79 → 0.26 ms (~3×), clustered 1-seg 1.0000 → 0.9960 + (within the probe's ε=0.005 contract) at 1.78 → 1.46 ms; gaussian 5-seg + 0.9915 at ~0.79 ms and 1-seg 0.7710 — byte-identical to the pre-split + baseline (probe self-rejects, full beam). + +### Performance — FT.SEARCH per-query fixed-cost cleanup (QP-2/QP-3) + +- **Thread-cached `SearchScratch`** (QP-3): the wire path built a fresh + scratch per query — heaps, visited bitmap, and the 32–65KB ADC LUT + reallocated every FT.SEARCH. Captures now take a thread-local recycled + scratch (exact padded-dim match, mirroring the worker pool's cache) and + the yielding search returns it on completion. +- **Amortized committed-treemap capture** (QP-2): every search cloned the + MVCC committed `RoaringTreemap`; `TransactionManager::committed_snapshot()` + now hands out a cached `Arc` refreshed only on the first capture after a + commit/prune — read-heavy captures are one refcount bump, write-heavy is + never worse than before. All 7 capture sites switched. +- **ryu score formatting**: RESP `__vec_score` replies formatted f32 via + `Display` (grisu, 1.4% of a matched-recall query); now ryu. +- VM A/B (same box/config as the f16 kernel entry): ef64 8,468 → 8,810 QPS + (+4.0%; +13% cumulative with the f16 kernels). + +### Performance — SIMD f16 exact-rerank kernels (NEON + F16C) + +- The HQ-1 exact-rerank pass decoded its f16 sidecar with scalar software + `f16_to_f32` — 17% of a matched-recall (ef 64) query. New fused + decode+distance kernels in the distance dispatch table (`f16_l2`, + `f16_dot_normsq`): aarch64 uses a baseline-NEON integer-rescale decode + (no ARMv8.2 FP16 intrinsics needed; exact for all finite f16, Inf/NaN + bit-selected), x86_64 uses F16C `vcvtph2ps` + FMA (runtime-detected with + AVX2, scalar fallback). Subnormal/Inf/NaN semantics match the scalar + reference exactly (pinned by tests). VM A/B (20k×384d clustered SQ8, + single conn, ef 64): 7,795 → 8,468 QPS (+8.6%). +- **Negative result, kept for the record**: skipping the rerank pass + entirely for SQ8 was A/B'd and REFUTED — it costs real recall + (R@10 −0.010 clustered, −0.017 gaussian 5-segment). The pass stays + unconditional; the SIMD kernels make it cheap instead. + ### Performance — FT.SEARCH intra-query worker pool + bounded bulk compaction - **`--ft-search-workers N` worker pool** (`src/vector/search_pool.rs`): the diff --git a/CLAUDE.md b/CLAUDE.md index 389ae449..b86f1ea6 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -203,7 +203,9 @@ Moon ships a native HNSW + TurboQuant vector engine exposed via a RediSearch-com - **Segment lifecycle:** auto-indexed on HSET → mutable segment (brute force) → compact → immutable segment (HNSW graph + TQ codes). Immutable segments DO merge: `MERGE_MODE GRAPH_UNION` (the default) auto-merges at ≥16 segments or ≥20% dead, gated by a recall check (0.70 background / 0.90 manual FT.COMPACT). What is forbidden is *decode+re-encode* merging — re-quantizing accumulates error (tested, recall collapsed 0.73 → 0.0005); GraphUnion stitches graphs over the original codes instead. `MERGE_MODE KEEP_RAW` is rejected at FT.CREATE (unimplemented stub). - **Key hash map:** `key_hash_to_key: Arc>` (copy-on-write via `Arc::make_mut`) must be populated in `auto_index_hset` and propagated via `SearchResult.key_hash`; otherwise multi-segment search returns synthetic `vec:` instead of original keys. - **FT.INFO `num_docs`** must sum across all segments (mutable + immutable), not just the mutable one — and across all shards: FT.INFO scatter-gathers via `scatter_ft_info` + `merge_ft_info_responses` (additive counters, per-field merge); a local-only answer under-reports by ~1/N. -- **Exact rerank sidecar (HQ-1):** immutable segments keep an f16 copy of each original vector (`raw_f16.bin` on disk; built from the mutable segment's always-retained f16 buffer); the top 4·k beam candidates are re-scored with true metric distances before truncation. Segments reloaded from pre-sidecar dirs (or rebuilt without raw data) silently fall back to quantized ADC distances — recall degrades, doesn't break. Merge propagation is all-or-nothing. +- **Exact rerank sidecar (HQ-1):** immutable segments keep an f16 copy of each original vector (`raw_f16.bin` on disk; built from the mutable segment's always-retained f16 buffer); the top 4·k beam candidates are re-scored with true metric distances before truncation, via runtime-detected SIMD kernels (NEON integer-rescale on aarch64, F16C+FMA on x86_64, scalar fallback everywhere; +8.6% QPS from the kernels alone, ~+13% cumulative with the per-query fixed-cost cleanup — thread-cached search scratch, Arc'd MVCC treemap snapshot, ryu score formatting). Segments reloaded from pre-sidecar dirs (or rebuilt without raw data) silently fall back to quantized ADC distances — recall degrades, doesn't break. Merge propagation is all-or-nothing, and a drop that discards a real sidecar is now LOUD: `tracing::warn!` with source / sources-with-sidecar counts, plus additive `FT.INFO` counters `graph_segments` / `segments_with_exact_rerank` (merged across shards; coverage < total ⇒ some segments answer ADC-only). +- **Adaptive per-segment ef (AE-1):** at compact/GraphUnion time each immutable segment self-probes recall against its own exact f16 sidecar (leave-self-out, ef ladder 24..256); only a fully-saturated curve (flat ≈1.0 from the minimum rung) certifies the segment "trivially easy" and searches at min-ef (24) — every other segment keeps the full resolved ef. Applies only when `EF_RUNTIME` was heuristic-resolved (never overrides a user-pinned value) and is in-memory only (reloaded segments fall back to the full beam). A blanket `ef/√G` split (EF-SPLIT) was tried first and REJECTED — it silently cost gaussian 5-seg R@10 0.9915 → 0.9295; don't reintroduce it. +- **FLUSHALL/FLUSHDB/HDEL keyspace parity:** FLUSHALL and FLUSHDB clear every vector + text index's CONTENTS while KEEPING the FT.CREATE definition (matches restart semantics; indexes are keyspace-global so FLUSHDB of any db clears all index contents) — previously flushed hashes stayed searchable as ghosts until restart. `HDEL key ` now tombstones that key in exactly the indexes whose vector field was removed (sibling indexes on other fields keep their entry). Known follow-ups: a multi-vector-field index tombstones the whole document when any one of its vector fields is removed; TEXT/TAG/NUMERIC field removal is not yet re-indexed. - **TQ4 at 384d loses recall** (concentration of distances + quantization noise). Use **SQ8** or FP32 HNSW for ≤384d workloads; TQ4 shines at 768d+. (There is no TQ8 — SQ8 is the real 8-bit option: per-vector affine scalar quant, normalizes for the unit-sphere metrics Cosine + InnerProduct, validated across the full lifecycle — search/merge/persistence — at ~0.90 R@10 on real MiniLM 384d. PR #166.) ## GPU / CUDA Acceleration diff --git a/Cargo.lock b/Cargo.lock index e8859c7e..a474fcc3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1870,6 +1870,7 @@ dependencies = [ "rust-stemmers", "rustls", "rustls-pemfile", + "ryu", "serde", "serde_json", "sha1_smol", diff --git a/Cargo.toml b/Cargo.toml index 957ea5f6..5f6be99b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -38,6 +38,7 @@ crc32fast = "1" arc-swap = "1" parking_lot = "0.12" itoa = "1" +ryu = "1" # float->shortest-string for RESP score replies (grisu via Display was 1.4% of FT.SEARCH) xxhash-rust = { version = "0.8", features = ["xxh64"] } ringbuf = "0.4" core_affinity = "0.8" diff --git a/benches/sq8_adc_bench.rs b/benches/sq8_adc_bench.rs index 6aa27c67..f46ff257 100644 --- a/benches/sq8_adc_bench.rs +++ b/benches/sq8_adc_bench.rs @@ -4,18 +4,28 @@ //! stats (NEON/AVX2/AVX-512 widen-u8-to-f32 + FMA) beat the naive per-element //! scalar `sq8_l2_adc` at standard embedding dimensions (128/384/768). //! -//! Three variants compared per dimension: +//! Four variants compared per dimension: //! - `naive_scalar`: the original per-candidate scalar loop (baseline). //! - `stats_scalar`: the algebraic stats decomposition, scalar stats kernel //! (isolates the win from the ALGEBRA alone, no SIMD). //! - `stats_dispatch`: the algebraic decomposition with the SIMD-dispatched -//! stats kernel (NEON on aarch64, AVX2/AVX-512 on x86_64) — this is what -//! the beam-search / brute-force hot paths now call. +//! f32 stats kernel (NEON on aarch64, AVX2/AVX-512 on x86_64) — this is +//! what the beam-search / brute-force hot paths called before task #13. +//! - `int8_dispatch`: task #13 — the SIMD-dispatched INT8 symmetric ADC +//! stats kernel (per-query quantization hoisted outside the timed loop, +//! matching the real call sites in `hnsw/search.rs` / +//! `segment/mutable.rs`). `None` (no int8 tier for this CPU/build, or +//! `MOON_SQ8_INT8_ADC=0`) falls back silently to `stats_dispatch`'s f32 +//! kernel so the group always has 4 comparable bars. +//! +//! Local numbers only (per tmp/INT8-ADC-CONTEXT.md: no production claims +//! from a Mac dev box — GCE cross-arch validation is task #14). use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; use moon::vector::distance; use moon::vector::turbo_quant::sq8::{ - sq8_candidate_stats_scalar, sq8_l2_adc, sq8_l2_from_stats, sq8_query_stats, + SQ8_INT8_QMAX, sq8_candidate_stats_scalar, sq8_int8_dot_to_f32, sq8_l2_adc, sq8_l2_from_stats, + sq8_quantize_query_scalar, sq8_query_stats, }; use std::hint::black_box; @@ -81,6 +91,46 @@ fn bench_sq8_adc(c: &mut Criterion) { sq8_l2_from_stats(dim, min, scale, q_sum, q_sumsq, dot_qc, sum_c, sumsq_c) }); }); + + group.bench_with_input(BenchmarkId::new("int8_dispatch", dim), &dim, |bench, _| { + let (q_sum, q_sumsq) = sq8_query_stats(&query); + let f32_stats_fn = distance::table().sq8_stats; + match distance::table().sq8_i8_stats { + Some(i8_stats_fn) => { + // Per-query quantization hoisted outside the timed loop, + // matching the real hnsw/search.rs and segment/mutable.rs + // call sites (quantize once per search, not per candidate). + let mut qi8 = vec![0i8; dim]; + let (q_scale, sum_qi8) = + sq8_quantize_query_scalar(&query, SQ8_INT8_QMAX, &mut qi8); + bench.iter(|| { + let (dot_int, sum_c_int, sumsq_c_int) = + i8_stats_fn(black_box(&qi8), black_box(&codes), sum_qi8); + let dot_qc = sq8_int8_dot_to_f32(q_scale, dot_int); + sq8_l2_from_stats( + dim, + min, + scale, + q_sum, + q_sumsq, + dot_qc, + sum_c_int as f32, + sumsq_c_int as f32, + ) + }); + } + None => { + // No int8 SIMD tier for this CPU/build (or + // MOON_SQ8_INT8_ADC=0) — fall back to the f32 dispatch + // path so the group always reports 4 comparable bars. + bench.iter(|| { + let (dot_qc, sum_c, sumsq_c) = + f32_stats_fn(black_box(&query), black_box(&codes)); + sq8_l2_from_stats(dim, min, scale, q_sum, q_sumsq, dot_qc, sum_c, sumsq_c) + }); + } + } + }); } group.finish(); } diff --git a/docs/commands.md b/docs/commands.md index 45443eb4..35dded01 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -120,7 +120,7 @@ TXN COMMIT -- Apply all changes atomically + WAL record `FT.CREATE`, `FT.DROPINDEX`, `FT.INFO`, `FT.SEARCH`, `FT.COMPACT`, `FT.RECOMMEND`, `FT.NAVIGATE`, `FT.EXPAND`, `FT.CACHESEARCH`, `FT.CONFIG SET`, `FT.CONFIG GET` !!! tip - Vectors are auto-indexed on `HSET`. Create an index with `FT.CREATE`, then `HSET` documents — Moon automatically inserts them into the HNSW graph. See the [vector search guide](vector-search-guide.md) for tuning parameters. + Vectors are auto-indexed on `HSET`. Create an index with `FT.CREATE`, then `HSET` documents — Moon automatically inserts them into the HNSW graph. `FLUSHALL`/`FLUSHDB` clear index contents (the `FT.CREATE` definition survives); `HDEL` of an indexed vector field tombstones it. See the [vector search guide](vector-search-guide.md) for tuning parameters and known limitations. ## Full-text search (2) @@ -200,3 +200,4 @@ FT.NAVIGATE idx "*=>[KNN 5 @vec $q]" PARAMS 2 q HOPS 2 DECAY 0.1 ## Known limitations - `GETRANGE` / `SETRANGE` are not yet implemented. +- `HDEL` of a vector field on an index with multiple vector fields tombstones the whole document, not just that field; `TEXT`/`TAG`/`NUMERIC` field removal via `HDEL` is not yet re-indexed. See the [vector search guide](vector-search-guide.md#flushall--flushdb--hdel). diff --git a/docs/vector-search-guide.md b/docs/vector-search-guide.md index b53ec0fa..b3b1df1d 100644 --- a/docs/vector-search-guide.md +++ b/docs/vector-search-guide.md @@ -123,7 +123,7 @@ Returns up to `k` nearest neighbors. The query vector must be a binary blob of ` ``` FT.INFO ``` -Returns index configuration: name, dimension, metric, quantization, build_mode. +Returns index configuration (name, dimension, metric, quantization, build_mode) plus observability counters, additive across shards: `graph_segments` (immutable HNSW segment count) and `segments_with_exact_rerank` (how many of those segments still carry the f16 exact-rerank sidecar). Coverage below `graph_segments` means some segments answer with quantized ADC-only distances — a GraphUnion merge that drops a sidecar logs a `tracing::warn` when it happens. ### FT.COMPACT ``` @@ -137,6 +137,12 @@ FT.DROPINDEX ``` Drop the index and free all associated memory. +### FLUSHALL / FLUSHDB / HDEL + +`FLUSHALL` and `FLUSHDB` clear every vector (and text) index's contents — segments, key-hash maps, postings — while KEEPING the `FT.CREATE` definition. This matches what a restart produces today: vector/text index contents are always rebuilt from the keyspace on restart (only the `FT.CREATE` definition is durable). Moon's FT indexes are keyspace-global, so `FLUSHDB` on any logical database clears all index contents. + +`HDEL key ` tombstones that key in exactly the indexes whose vector field was removed (a sibling index keyed on a different field keeps its entry). Whole-key deletion (`DEL`/`UNLINK`) already tombstoned every index. Known limitations: an index with multiple vector fields tombstones the *whole* document if any one of its vector fields is removed (a later `HSET` re-indexes the remainder), and `TEXT`/`TAG`/`NUMERIC` field removal via `HDEL` is not yet re-indexed. + ## How It Works ### Insert Path @@ -155,13 +161,15 @@ Triggered automatically on first search when mutable segment has ≥ `COMPACT_TH 4. BFS-reorder for cache locality 5. Compute sub-centroid sign bits (doubles quantization resolution: 16 → 32 levels) 6. Create immutable segment +7. **Adaptive-ef self-probe (AE-1):** 16 leave-self-out sample queries measure R@10 against the segment's own exact f16 sidecar across an ef ladder (24..256). A fully-saturated curve (flat ≈1.0 from the minimum rung) certifies the segment as "trivially easy" for min-ef search; every other segment keeps the full resolved ef at query time. In-memory only — not persisted, so a segment reloaded from disk always searches at the full beam. ### Search Path 1. Query vector → normalize → FWHT rotate 2. Build per-query LUT: precomputed distance² for each sub-centroid (32 entries × dim, fits L1 cache) -3. **HNSW beam search** with 32-level sub-centroid LUT scoring (no separate rerank needed) -4. Merge results from mutable (brute-force) + immutable (HNSW) segments -5. Return top-K results +3. **HNSW beam search** with 32-level sub-centroid LUT scoring. Beam width (`ef`) is the full resolved value, unless the segment's compact-time saturation probe (AE-1, above) certified it "trivially easy" — then it searches at min-ef (24) instead. Never overridden when the user pins `EF_RUNTIME`. +4. **Exact rerank:** the top `4·k` beam candidates are re-scored against the segment's f16 sidecar with true metric distances (SIMD: NEON integer-rescale on aarch64, F16C+FMA on x86_64, scalar fallback) before truncation. Segments without a sidecar (pre-HQ-1 reload, or a GraphUnion merge that dropped one) fall back to quantized ADC-only distances — check `FT.INFO`'s `segments_with_exact_rerank` for coverage. +5. Merge results from mutable (brute-force) + immutable (HNSW) segments +6. Return top-K results ## Memory Usage diff --git a/src/command/vector_search/ft_info.rs b/src/command/vector_search/ft_info.rs index 97347982..5ce95619 100644 --- a/src/command/vector_search/ft_info.rs +++ b/src/command/vector_search/ft_info.rs @@ -43,6 +43,16 @@ pub fn ft_info( for imm in snap.immutable.iter() { num_docs += imm.live_count() as usize; } + // HQ-1 observability (persistence-review R5): exact-rerank coverage. + // A segment without the f16 sidecar silently answers with quantized ADC + // distances only — surfacing the count makes a dropped sidecar (e.g. an + // all-or-nothing GraphUnion merge with one sidecar-less source) visible. + let graph_segments = snap.immutable.len(); + let segments_with_exact_rerank = snap + .immutable + .iter() + .filter(|imm| imm.raw_f16().is_some()) + .count(); // Use itoa for numeric formatting -- no format!() on hot path. let ef_rt_bytes: Bytes = if idx.meta.hnsw_ef_runtime > 0 { @@ -87,6 +97,10 @@ pub fn ft_info( Frame::BulkString(ct_bytes), Frame::BulkString(Bytes::from_static(b"QUANTIZATION")), Frame::BulkString(quant_bytes), + Frame::BulkString(Bytes::from_static(b"graph_segments")), + Frame::Integer(graph_segments as i64), + Frame::BulkString(Bytes::from_static(b"segments_with_exact_rerank")), + Frame::Integer(segments_with_exact_rerank as i64), ]; // Per-field stats: vector_fields array @@ -197,9 +211,10 @@ pub fn ft_info( /// from the local response. /// /// Additive top-level keys: `num_docs`, `num_terms`, -/// `total_inverted_index_size`, and the freshness tokens -/// `vector_version_token` / `text_version_token` (each shard's token is -/// monotonic, so the sum is monotonic too — any shard's write bumps the +/// `total_inverted_index_size`, the exact-rerank coverage counters +/// `graph_segments` / `segments_with_exact_rerank` (R5), and the freshness +/// tokens `vector_version_token` / `text_version_token` (each shard's token +/// is monotonic, so the sum is monotonic too — any shard's write bumps the /// aggregate). /// Additive per-field keys (matched by `field_name` inside `vector_fields` / /// `text_fields`): `num_docs`, `mutable_vectors`, `immutable_segments`. @@ -213,6 +228,8 @@ pub fn merge_ft_info_responses(local: Frame, remotes: &[Frame]) -> Frame { b"total_inverted_index_size", b"vector_version_token", b"text_version_token", + b"graph_segments", + b"segments_with_exact_rerank", ]; const ADDITIVE_FIELD: &[&[u8]] = &[b"num_docs", b"mutable_vectors", b"immutable_segments"]; @@ -496,6 +513,32 @@ mod merge_tests { assert_eq!(get_int(&entries[0], b"immutable_segments"), 2); } + #[test] + fn sums_exact_rerank_coverage_counters() { + // R5: graph_segments / segments_with_exact_rerank are additive so a + // sidecar-less segment on ANY shard shows up in the aggregate + // (coverage < total ⇒ some segment answers with ADC-only distances). + fn frame_with_coverage(total: i64, with_rerank: i64) -> Frame { + Frame::Array( + vec![ + bs(b"index_name"), + bs(b"idx"), + bs(b"graph_segments"), + Frame::Integer(total), + bs(b"segments_with_exact_rerank"), + Frame::Integer(with_rerank), + ] + .into(), + ) + } + let merged = merge_ft_info_responses( + frame_with_coverage(3, 3), + &[frame_with_coverage(2, 1), frame_with_coverage(1, 0)], + ); + assert_eq!(get_int(&merged, b"graph_segments"), 6); + assert_eq!(get_int(&merged, b"segments_with_exact_rerank"), 4); + } + #[test] fn propagates_remote_error() { let local = info_frame(10, 10, 4, 1); diff --git a/src/command/vector_search/ft_search/dispatch.rs b/src/command/vector_search/ft_search/dispatch.rs index 5ecafd6f..960db422 100644 --- a/src/command/vector_search/ft_search/dispatch.rs +++ b/src/command/vector_search/ft_search/dispatch.rs @@ -577,13 +577,12 @@ fn capture_dense_knn_snapshot( filter: Option<&crate::vector::filter::FilterExpr>, as_of_lsn: u64, ) -> Option { - use crate::vector::hnsw::search::SearchScratch; use crate::vector::segment::holder::SearchSnapshot; use crate::vector::turbo_quant::encoder::padded_dimension; // Clone committed treemap BEFORE get_index_mut (borrow-checker ordering), // matching search_local_raw / search_local_filtered. - let committed = store.txn_manager().committed_treemap().clone(); + let committed = store.txn_manager().committed_snapshot(); let idx = store.get_index_mut(index_name)?; // Default field only — non-default field stays on the sync path. @@ -623,6 +622,9 @@ fn capture_dense_knn_snapshot( // Auto-compact (same side effect + timing as the sync path). idx.try_compact(); + // AE-1: adaptive per-segment ef estimates apply only when ef is + // heuristic-defaulted, never over a user-pinned EF_RUNTIME. + let ef_defaulted = idx.meta.hnsw_ef_runtime == 0; let ef_search = if idx.meta.hnsw_ef_runtime > 0 { idx.meta.hnsw_ef_runtime as usize } else { @@ -662,8 +664,11 @@ fn capture_dense_knn_snapshot( committed, dimension: dim as u32, mutable_len, - scratch: SearchScratch::new(0, padded_dimension(dim as u32)), + // QP-3: thread-cached scratch (exact padded_dim match) instead of a + // fresh 32-65KB allocation set per query. + scratch: crate::vector::hnsw::search::take_thread_scratch(padded_dimension(dim as u32)), key_hash_to_key: idx.key_hash_to_key.clone(), + ef_defaulted, }) } diff --git a/src/command/vector_search/ft_search/execute.rs b/src/command/vector_search/ft_search/execute.rs index cb8656f7..e22aef86 100644 --- a/src/command/vector_search/ft_search/execute.rs +++ b/src/command/vector_search/ft_search/execute.rs @@ -47,7 +47,7 @@ pub(super) fn search_local_raw( // Clone committed treemap BEFORE get_index_mut to satisfy the borrow checker. // Non-TXN readers need this to see entries whose owning txn has committed // (entries tagged with txn_id by auto_index_hset_public_txn; ACID-09 fix). - let committed = store.txn_manager().committed_treemap().clone(); + let committed = store.txn_manager().committed_snapshot(); let idx = match store.get_index_mut(index_name) { Some(i) => i, None => { @@ -96,6 +96,9 @@ pub(super) fn search_local_raw( idx.try_compact(); + // AE-1: remember whether ef came from the heuristic (segment-level + // adaptive-ef estimates only apply then, never over a user EF_RUNTIME). + let ef_defaulted = idx.meta.hnsw_ef_runtime == 0; let ef_search = if idx.meta.hnsw_ef_runtime > 0 { idx.meta.hnsw_ef_runtime as usize } else { @@ -123,6 +126,7 @@ pub(super) fn search_local_raw( committed: &committed, dirty_set: &[], dimension: dim as u32, + ef_defaulted, }; let results = idx.segments.search_mvcc( &query_f32, @@ -148,6 +152,7 @@ pub(super) fn search_local_raw( committed: &committed, dirty_set: &[], dimension: dim as u32, + ef_defaulted, }; let results = fs.segments.search_mvcc( &query_f32, @@ -218,7 +223,7 @@ pub fn search_local_filtered( ) -> Frame { // Clone committed treemap BEFORE get_index_mut (borrow-checker ordering). // Ensures non-TXN readers see entries whose owning txn has committed. - let committed = store.txn_manager().committed_treemap().clone(); + let committed = store.txn_manager().committed_snapshot(); let idx = match store.get_index_mut(index_name) { Some(i) => i, None => return Frame::Error(Bytes::from_static(b"Unknown Index name")), @@ -269,6 +274,9 @@ pub fn search_local_filtered( // ef_search: user-configurable via EF_RUNTIME in FT.CREATE, or auto-computed. // Higher ef = better recall but lower QPS. Auto scales with k and dimension: // base = k*20, min 200, boosted for high-d where TQ-ADC needs wider beam. + // AE-1: remember whether ef came from the heuristic (segment-level + // adaptive-ef estimates only apply then, never over a user EF_RUNTIME). + let ef_defaulted = idx.meta.hnsw_ef_runtime == 0; let ef_search = if idx.meta.hnsw_ef_runtime > 0 { idx.meta.hnsw_ef_runtime as usize } else { @@ -297,6 +305,7 @@ pub fn search_local_filtered( committed: &committed, dirty_set: &[], dimension: dim as u32, + ef_defaulted, }; let results = idx.segments.search_mvcc( &query_f32, @@ -318,6 +327,7 @@ pub fn search_local_filtered( committed: &committed, dirty_set: &[], dimension: dim as u32, + ef_defaulted, }; let results = fs.segments.search_mvcc( &query_f32, diff --git a/src/command/vector_search/ft_search/response.rs b/src/command/vector_search/ft_search/response.rs index 6e29ffc8..6d1a1999 100644 --- a/src/command/vector_search/ft_search/response.rs +++ b/src/command/vector_search/ft_search/response.rs @@ -65,10 +65,10 @@ pub(crate) fn build_search_response( items.push(Frame::BulkString(doc_id)); // Score as nested array — use write! to pre-allocated buffer - let mut score_buf = String::with_capacity(16); - use std::fmt::Write; - let _ = write!(score_buf, "{}", r.distance); - let score_str = score_buf; + // ryu: shortest-roundtrip float formatting without Display's grisu + // (1.4% of a matched-recall query). + let mut fbuf = ryu::Buffer::new(); + let score_str = fbuf.format(r.distance).to_owned(); let fields = vec![ Frame::BulkString(Bytes::from_static(b"__vec_score")), Frame::BulkString(Bytes::from(score_str)), @@ -257,12 +257,11 @@ pub(crate) fn build_hybrid_response( items.push(Frame::BulkString(doc_id)); // Score field — use absolute value of RRF score for display - let mut score_buf = String::with_capacity(16); - use std::fmt::Write; - let _ = write!(score_buf, "{}", r.distance.abs()); + let mut fbuf = ryu::Buffer::new(); + let score_str = fbuf.format(r.distance.abs()).to_owned(); let fields = vec![ Frame::BulkString(Bytes::from_static(b"__vec_score")), - Frame::BulkString(Bytes::from(score_buf)), + Frame::BulkString(Bytes::from(score_str)), ]; items.push(Frame::Array(fields.into())); } @@ -350,12 +349,11 @@ pub(crate) fn build_combined_response( items.push(Frame::BulkString(er.key.clone())); let mut hop_buf = itoa::Buffer::new(); let hop_str = hop_buf.format(er.graph_hops); - let mut score_buf = String::with_capacity(8); - use std::fmt::Write; - let _ = write!(score_buf, "{}", er.vec_score); + let mut fbuf = ryu::Buffer::new(); + let score_str = fbuf.format(er.vec_score).to_owned(); let fields = vec![ Frame::BulkString(Bytes::from_static(b"__vec_score")), - Frame::BulkString(Bytes::from(score_buf)), + Frame::BulkString(Bytes::from(score_str)), Frame::BulkString(Bytes::from_static(b"__graph_hops")), Frame::BulkString(Bytes::from(hop_str.to_owned())), ]; diff --git a/src/command/vector_search/hybrid.rs b/src/command/vector_search/hybrid.rs index 9401dd46..98f59b48 100644 --- a/src/command/vector_search/hybrid.rs +++ b/src/command/vector_search/hybrid.rs @@ -366,7 +366,7 @@ pub fn execute_hybrid_search_local( // v0.1.9 HYB-01: snapshot the committed treemap BEFORE get_index_mut so the // dense stream honors AS_OF via MVCC filtering. Non-TXN readers must see // entries whose owning txn has committed (matches search_local_raw pattern). - let committed = vector_store.txn_manager().committed_treemap().clone(); + let committed = vector_store.txn_manager().committed_snapshot(); // ── Stream 1: BM25 (per D-10 — direct use of existing BM25 ranker) ──────── // @@ -594,6 +594,7 @@ pub(super) fn run_dense_knn( committed, dirty_set: &[], dimension: dim as u32, + ef_defaulted: false, }; idx.segments .search_mvcc(&query_f32, k, ef_search, &mut idx.scratch, None, &mvcc_ctx) @@ -615,6 +616,7 @@ pub(super) fn run_dense_knn( committed, dirty_set: &[], dimension: dim as u32, + ef_defaulted: false, }; fs.segments .search_mvcc(&query_f32, k, ef_search, &mut fs.scratch, None, &mvcc_ctx) diff --git a/src/command/vector_search/hybrid_multi.rs b/src/command/vector_search/hybrid_multi.rs index 28aeb75e..74130394 100644 --- a/src/command/vector_search/hybrid_multi.rs +++ b/src/command/vector_search/hybrid_multi.rs @@ -117,7 +117,7 @@ pub fn execute_hybrid_search_local_raw_streams( // multi-shard scatter. Matches the single-shard hybrid.rs pattern. // `as_of_lsn == 0` → no temporal filtering (treemap still snapshotted for // ACID-09 committed-entry visibility parity with search_local_raw). - let committed = vector_store.txn_manager().committed_treemap().clone(); + let committed = vector_store.txn_manager().committed_snapshot(); let idx = match vector_store.get_index_mut(index_name.as_ref()) { Some(ix) => ix, None => return Frame::Error(Bytes::from_static(b"ERR unknown index")), diff --git a/src/command/vector_search/recommend.rs b/src/command/vector_search/recommend.rs index 603e06ab..64269c70 100644 --- a/src/command/vector_search/recommend.rs +++ b/src/command/vector_search/recommend.rs @@ -341,6 +341,7 @@ pub fn ft_recommend(store: &mut VectorStore, args: &[Frame], db: Option<&mut Dat committed: &empty_committed, dirty_set: &[], dimension: dim as u32, + ef_defaulted: false, }; // Dispatch to correct field segments diff --git a/src/server/conn/handler_monoio/mod.rs b/src/server/conn/handler_monoio/mod.rs index eb1b013c..dc85513f 100644 --- a/src/server/conn/handler_monoio/mod.rs +++ b/src/server/conn/handler_monoio/mod.rs @@ -1336,6 +1336,26 @@ pub(crate) async fn handle_connection_sharded_monoio< ); } + // R4: HDEL of an indexed vector field tombstones it. + if !is_error && cmd.eq_ignore_ascii_case(b"HDEL") { + crate::shard::spsc_handler::auto_hdel_vectors( + &mut s.vector_store, + cmd_args, + ); + } + + // R3: FLUSHALL/FLUSHDB clears index contents + // (FT.CREATE definitions survive). + if !is_error + && (cmd.eq_ignore_ascii_case(b"FLUSHDB") + || cmd.eq_ignore_ascii_case(b"FLUSHALL")) + { + crate::shard::spsc_handler::auto_flush_indexes( + &mut s.vector_store, + &mut s.text_store, + ); + } + // Blocking wakeup: re-borrow db by index (NLL) if !is_error { let needs_wake = cmd.eq_ignore_ascii_case(b"LPUSH") @@ -1490,7 +1510,7 @@ pub(crate) async fn handle_connection_sharded_monoio< // Clone committed treemap to release vector_store lock // before acquiring kv_intents lock (lock ordering). let committed = crate::shard::slice::with_shard(|s| { - s.vector_store.txn_manager().committed_treemap().clone() + s.vector_store.txn_manager().committed_snapshot() }); let visible = crate::shard::slice::with_shard(|s| { s.kv_write_intents.is_key_visible( diff --git a/src/server/conn/handler_sharded/mod.rs b/src/server/conn/handler_sharded/mod.rs index ff439aa2..894b5089 100644 --- a/src/server/conn/handler_sharded/mod.rs +++ b/src/server/conn/handler_sharded/mod.rs @@ -1421,8 +1421,6 @@ pub(crate) async fn handle_connection_sharded_inner< } } // Auto-delete vectors on DEL/UNLINK (local write path) - // Note: HDEL removes fields, not keys — it should NOT trigger - // vector deletion unless the entire key is removed. if !matches!(response, Frame::Error(_)) && (cmd.eq_ignore_ascii_case(b"DEL") || cmd.eq_ignore_ascii_case(b"UNLINK")) { @@ -1435,6 +1433,32 @@ pub(crate) async fn handle_connection_sharded_inner< } }); } + // R4: HDEL of an indexed VECTOR field tombstones the vector + // in exactly the affected indexes (whole-key deletion is the + // DEL/UNLINK arm above; non-vector-field HDELs are no-ops). + if !matches!(response, Frame::Error(_)) + && cmd.eq_ignore_ascii_case(b"HDEL") + { + crate::shard::slice::with_shard(|s| { + crate::shard::spsc_handler::auto_hdel_vectors( + &mut s.vector_store, + cmd_args, + ); + }); + } + // R3: FLUSHALL/FLUSHDB clears vector + text index contents + // (FT.CREATE definitions survive, matching restart semantics). + if !matches!(response, Frame::Error(_)) + && (cmd.eq_ignore_ascii_case(b"FLUSHDB") + || cmd.eq_ignore_ascii_case(b"FLUSHALL")) + { + crate::shard::slice::with_shard(|s| { + crate::shard::spsc_handler::auto_flush_indexes( + &mut s.vector_store, + &mut s.text_store, + ); + }); + } // H1: durable path under appendfsync=always. let mut aof_failed = false; if let Some(bytes) = aof_bytes { @@ -1484,7 +1508,7 @@ pub(crate) async fn handle_connection_sharded_inner< // before acquiring kv_intents lock (lock ordering). // Unconditional slice path: ShardSlice is always initialized. let committed = crate::shard::slice::with_shard(|s| { - s.vector_store.txn_manager().committed_treemap().clone() + s.vector_store.txn_manager().committed_snapshot() }); // ShardSlice path: kv_write_intents lives on the thread-local slice. let visible = crate::shard::slice::with_shard(|s| { diff --git a/src/server/conn/handler_single.rs b/src/server/conn/handler_single.rs index 1d4d8c8d..965e1eb5 100644 --- a/src/server/conn/handler_single.rs +++ b/src/server/conn/handler_single.rs @@ -1149,6 +1149,27 @@ pub async fn handle_connection( &mut vs.lock(), a, ); + } else if c.eq_ignore_ascii_case(b"HDEL") + && i < txn_results.len() + && !matches!(txn_results[i], Frame::Error(_)) + { + // R4 parity inside MULTI/EXEC. + crate::shard::spsc_handler::auto_hdel_vectors( + &mut vs.lock(), + a, + ); + } else if (c.eq_ignore_ascii_case(b"FLUSHDB") + || c.eq_ignore_ascii_case(b"FLUSHALL")) + && i < txn_results.len() + && !matches!(txn_results[i], Frame::Error(_)) + { + // R3 parity inside MULTI/EXEC (text + // store cleared via its own guard or + // the throwaway fallback store). + vs.lock().clear_all_contents(); + if let Some(ref ts) = text_store { + ts.lock().clear_all_contents(); + } } } } @@ -2247,6 +2268,32 @@ pub async fn handle_connection( } } + // R4: HDEL of an indexed vector field tombstones it. + if !matches!(&response, Frame::Error(_)) + && d_cmd.eq_ignore_ascii_case(b"HDEL") + { + if let Some(ref vs) = vector_store { + crate::shard::spsc_handler::auto_hdel_vectors( + &mut vs.lock(), + d_args, + ); + } + } + + // R3: FLUSHALL/FLUSHDB clears vector + text index + // contents (FT.CREATE definitions survive). + if !matches!(&response, Frame::Error(_)) + && (d_cmd.eq_ignore_ascii_case(b"FLUSHDB") + || d_cmd.eq_ignore_ascii_case(b"FLUSHALL")) + { + if let Some(ref vs) = vector_store { + vs.lock().clear_all_contents(); + } + if let Some(ref ts) = text_store { + ts.lock().clear_all_contents(); + } + } + // Invalidate tracked key on successful write if !matches!(&response, Frame::Error(_)) { if let Some(key) = d_args.first().and_then(|f| extract_bytes(f)) { diff --git a/src/server/conn/shared.rs b/src/server/conn/shared.rs index 2ef3b93f..aabd0192 100644 --- a/src/server/conn/shared.rs +++ b/src/server/conn/shared.rs @@ -246,6 +246,25 @@ pub(crate) fn execute_transaction_sharded( }); } + // R4: HDEL of an indexed vector field tombstones it. + if !matches!(response, Frame::Error(_)) && cmd.eq_ignore_ascii_case(b"HDEL") { + crate::shard::slice::with_shard(|s| { + crate::shard::spsc_handler::auto_hdel_vectors(&mut s.vector_store, cmd_args); + }); + } + + // R3: FLUSHALL/FLUSHDB clears index contents (definitions survive). + if !matches!(response, Frame::Error(_)) + && (cmd.eq_ignore_ascii_case(b"FLUSHDB") || cmd.eq_ignore_ascii_case(b"FLUSHALL")) + { + crate::shard::slice::with_shard(|s| { + crate::shard::spsc_handler::auto_flush_indexes( + &mut s.vector_store, + &mut s.text_store, + ); + }); + } + results.push(response); } diff --git a/src/shard/event_loop.rs b/src/shard/event_loop.rs index ab7c817b..11cb3c8f 100644 --- a/src/shard/event_loop.rs +++ b/src/shard/event_loop.rs @@ -837,8 +837,14 @@ impl super::Shard { // This ensures handler_sharded FT.* commands and SPSC auto-indexing // (triggered by HSET) operate on the SAME VectorStore. // - // The shard-owned vector_store (from Shard struct) is discarded. - // All vector operations go through shard_databases.vector_store(shard_id). + // The shard-owned vector_store (populated by `Shard::restore_from_persistence`, + // BEFORE the event loop starts) is discarded here in favor of the + // one on `ShardSlice`. Real recovery happens below, against THIS + // (live) store: index *definitions* come from the + // `vector-indexes.meta` sidecar; index *contents* come from the B3 + // manifest/segment/keymap durability layout when present (see + // `crate::vector::persistence::recover_v2`), reconciled against the + // live keyspace by a dedup rescan — never from the discarded store. let _discarded_vector_store = std::mem::replace( &mut self.vector_store, crate::vector::store::VectorStore::new(), @@ -886,7 +892,15 @@ impl super::Shard { } }); - if let Some(ref metas) = metas { + // B3 (vector-index durability): threaded through the index + // creation loop below, the rescan loop further down, and the + // finalize call at the end of this block. See + // `crate::vector::persistence::recover_v2` module docs for the + // full recovery contract (manifest/segment/keymap load + dedup + // rescan + deletion probe + orphan sweep). + let mut recovery_state = crate::vector::persistence::recover_v2::RecoveryState::new(); + + if let (Some(metas), Some(vdir)) = (&metas, &vector_persist_dir) { crate::shard::slice::with_shard(|s| { info!( "Shard {}: restoring {} vector index(es) from sidecar", @@ -894,16 +908,9 @@ impl super::Shard { metas.len() ); for (meta, weight) in metas { - let name = meta.name.clone(); - if let Err(e) = s.vector_store.create_index(meta.clone()) { - tracing::warn!( - "Shard {}: failed to restore index '{}': {}", - shard_id, - String::from_utf8_lossy(&name), - e - ); - } else if *weight != 1.0 { - if let Some(idx) = s.vector_store.get_index_mut(&name) { + recovery_state.create_index(&mut s.vector_store, vdir, meta); + if *weight != 1.0 { + if let Some(idx) = s.vector_store.get_index_mut(&meta.name) { idx.set_compaction_weight(*weight); } } @@ -1004,7 +1011,14 @@ impl super::Shard { if !matching.is_empty() { crate::shard::slice::with_shard(|s| { for (key, args) in &matching { - let _ = crate::shard::spsc_handler::auto_index_hset_public( + // B3 dedup rescan: verifies each matching + // key against any recovered durable state + // (manifest/segment/keymap) before deciding + // whether to fully re-encode. Indexes with + // no durable state (fresh/no manifest) fall + // through to the same full-rescan behavior + // this replaced. See `recover_v2` docs. + recovery_state.reconcile_key( &mut s.vector_store, &mut s.text_store, key, @@ -1022,6 +1036,18 @@ impl super::Shard { ); } } + + // B3 finalize: deletion probe (keymap keys no longer present + // anywhere in the keyspace) + orphan sweep (segment/staging/ + // keymap files not referenced by the loaded manifest, and + // unknown `idx-*` dirs with no matching sidecar index) + + // per-index acceptance-signal log line. No-op (does nothing, + // logs nothing) when no index had durable state to recover. + if let Some(ref vdir) = vector_persist_dir { + crate::shard::slice::with_shard(|s| { + recovery_state.finish(&mut s.vector_store, vdir); + }); + } } // Waker relay, swept after every drain cycle. HISTORICAL NOTE: this was diff --git a/src/shard/mod.rs b/src/shard/mod.rs index 902a332e..38113144 100644 --- a/src/shard/mod.rs +++ b/src/shard/mod.rs @@ -168,30 +168,13 @@ impl Shard { } } - // Vector recovery still uses the v2 path for now - self.recover_vectors(persistence_dir); - - // Register warm segments into VectorStore so they're searchable - if !result.warm_segments.is_empty() { - info!( - "Shard {}: registering {} warm segment(s)", - self.id, - result.warm_segments.len() - ); - self.vector_store - .register_warm_segments(result.warm_segments); - } - - // Register cold DiskANN segments for discovery - if !result.cold_segments.is_empty() { - info!( - "Shard {}: registering {} cold segment(s)", - self.id, - result.cold_segments.len() - ); - self.vector_store - .register_cold_segments(result.cold_segments); - } + // Vector recovery: the `Shard`-owned `vector_store` + // populated here is discarded wholesale at + // `event_loop.rs` (`_discarded_vector_store`) in + // favor of `ShardSlice.vector_store` — see that + // file's comment for the real recovery contract + // (sidecar definitions + manifest/segments/keymap + + // dedup rescan, B3). Nothing to do on this struct. return result.commands_replayed; } Err(e) => { @@ -273,44 +256,15 @@ impl Shard { } } - // Recover vector store - self.recover_vectors(persistence_dir); + // Vector store recovery does NOT happen here — see the comment in + // `restore_from_persistence`'s v3 branch above. The `Shard`-owned + // `vector_store` this method populates is discarded wholesale at + // `event_loop.rs`; real recovery runs later against + // `ShardSlice.vector_store` (sidecar definitions + B3 + // manifest/segments/keymap load + dedup rescan). total_keys } - - /// Recover vector store from WAL + on-disk segments. - fn recover_vectors(&mut self, persistence_dir: &str) { - let dir = std::path::Path::new(persistence_dir); - let wal_file = crate::persistence::wal::wal_path(dir, self.id); - let vector_persist_dir = dir.join(format!("shard-{}-vectors", self.id)); - if vector_persist_dir.exists() || wal_file.exists() { - match crate::vector::persistence::recovery::recover_vector_store( - &wal_file, - &vector_persist_dir, - ) { - Ok(recovered) => { - let seg_count: usize = recovered - .collections - .values() - .map(|c| c.immutable.len()) - .sum(); - if !recovered.collections.is_empty() { - info!( - "Shard {}: recovered {} vector collections ({} immutable segments)", - self.id, - recovered.collections.len(), - seg_count - ); - } - self.vector_store.attach_recovered(recovered); - } - Err(e) => { - tracing::error!("Shard {}: vector recovery failed: {:?}", self.id, e); - } - } - } - } } #[cfg(test)] diff --git a/src/shard/spsc_handler.rs b/src/shard/spsc_handler.rs index 75e1c4a1..4f8a227a 100644 --- a/src/shard/spsc_handler.rs +++ b/src/shard/spsc_handler.rs @@ -926,6 +926,21 @@ pub(crate) fn handle_shard_message_shared( auto_delete_vectors(&mut s.vector_store, args); } + // R4: HDEL of an indexed vector field tombstones the vector. + if !matches!(frame, crate::protocol::Frame::Error(_)) + && cmd.eq_ignore_ascii_case(b"HDEL") + { + auto_hdel_vectors(&mut s.vector_store, args); + } + + // R3: FLUSHALL/FLUSHDB clears index contents (definitions kept). + if !matches!(frame, crate::protocol::Frame::Error(_)) + && (cmd.eq_ignore_ascii_case(b"FLUSHDB") + || cmd.eq_ignore_ascii_case(b"FLUSHALL")) + { + auto_flush_indexes(&mut s.vector_store, &mut s.text_store); + } + // Post-dispatch wakeup hooks for producer commands (cross-shard blocking) if !matches!(frame, crate::protocol::Frame::Error(_)) { let needs_wake = cmd.eq_ignore_ascii_case(b"LPUSH") @@ -1297,6 +1312,21 @@ pub(crate) fn handle_shard_message_shared( auto_delete_vectors(&mut s.vector_store, args); } + // R4: HDEL of an indexed vector field tombstones the vector. + if !matches!(frame, crate::protocol::Frame::Error(_)) + && cmd.eq_ignore_ascii_case(b"HDEL") + { + auto_hdel_vectors(&mut s.vector_store, args); + } + + // R3: FLUSHALL/FLUSHDB clears index contents (definitions kept). + if !matches!(frame, crate::protocol::Frame::Error(_)) + && (cmd.eq_ignore_ascii_case(b"FLUSHDB") + || cmd.eq_ignore_ascii_case(b"FLUSHALL")) + { + auto_flush_indexes(&mut s.vector_store, &mut s.text_store); + } + if !matches!(frame, crate::protocol::Frame::Error(_)) { let needs_wake = cmd.eq_ignore_ascii_case(b"LPUSH") || cmd.eq_ignore_ascii_case(b"RPUSH") @@ -2160,6 +2190,64 @@ pub fn auto_delete_vectors(vector_store: &mut VectorStore, args: &[crate::protoc } } +/// FLUSHALL/FLUSHDB parity (persistence-review R3): clear vector + text +/// index CONTENTS while keeping the FT.CREATE definitions. Wire-parity +/// requirement mirrors `auto_delete_vectors`: every dispatch path that runs +/// the HSET auto-index hook must run this on a successful FLUSH, or flushed +/// hashes stay searchable as ghost documents until restart. +/// +/// Moon FT indexes are keyspace-global (the auto-index hook is not gated on +/// the selected db), so FLUSHDB of any db clears all index contents — the +/// same over-approximation the restart rescan would produce. +pub fn auto_flush_indexes( + vector_store: &mut VectorStore, + text_store: &mut crate::text::store::TextStore, +) { + vector_store.clear_all_contents(); + text_store.clear_all_contents(); +} + +/// HDEL parity (persistence-review R4): a successful `HDEL key field...` +/// that removes an index's VECTOR field must tombstone that key's vector in +/// exactly the affected indexes — previously the vector stayed searchable +/// until whole-key DEL or a re-HSET. After a successful HDEL the named +/// fields are definitively absent from the hash, so tombstoning any index +/// keyed on one of them always agrees with the hash state (fields that were +/// already absent tombstone a mapping that should not exist anyway). +/// +/// Known limitations (documented follow-ups): a multi-vector-field index +/// tombstones the WHOLE document when any of its vector fields is removed +/// (a later HSET re-indexes the remainder), and TEXT/TAG/NUMERIC field +/// removal is not yet re-indexed. +pub fn auto_hdel_vectors(vector_store: &mut VectorStore, args: &[crate::protocol::Frame]) { + let Some(key) = args + .first() + .and_then(crate::server::connection::extract_bytes) + else { + return; + }; + let matching = vector_store.find_matching_index_names(key.as_ref()); + for idx_name in matching { + let Some(idx) = vector_store.get_index(&idx_name) else { + continue; + }; + let removed_vector_field = args[1..] + .iter() + .filter_map(crate::server::connection::extract_bytes) + .any(|field| { + idx.meta.source_field.as_ref() == field.as_ref() + || idx + .meta + .vector_fields + .iter() + .any(|vf| vf.field_name.as_ref() == field.as_ref()) + }); + if removed_vector_field { + vector_store.mark_deleted_for_key_in_index(&idx_name, key.as_ref()); + } + } +} + /// TXN-aware variant: tags each inserted vector entry with `txn_id` so /// non-transactional readers (snapshot_lsn == 0) see it as uncommitted and /// exclude it until TXN.COMMIT calls `txn_manager.commit(txn_id)`. @@ -2314,7 +2402,13 @@ fn auto_index_hset( /// Find the vector blob in HSET args for the given source_field. /// Returns Some(blob) if found with correct dimension, None otherwise. -fn find_vector_blob<'a>( +/// +/// `pub(crate)` so the B3 recovery dedup rescan +/// (`crate::vector::persistence::recover_v2`) can compute the CURRENT +/// vector-field checksum through the exact same field-scan the write path +/// uses (`handle_vector_insert`) — any divergence would mean the dedup +/// decision silently never fires. +pub(crate) fn find_vector_blob<'a>( args: &'a [crate::protocol::Frame], source_field: &[u8], dim: usize, @@ -2408,8 +2502,14 @@ fn handle_vector_insert( let global_id = snap.mutable.global_id_base() + internal_id; crate::vector::metrics::add_vectors(1); - // Record key_hash → global_id mapping for future metadata-only updates - idx.key_hash_to_global_id.insert(key_hash, global_id); + // Record key_hash → global_id mapping for future metadata-only updates. + // COW via make_mut, mirroring key_hash_to_key (QP-1). + std::sync::Arc::make_mut(&mut idx.key_hash_to_global_id).insert(key_hash, global_id); + // B2 (durability): mirror the same key_hash into the checksum map so the + // two maps never drift — the B3 dedup rescan compares this checksum + // against a freshly-hashed current value to decide unchanged-vs-changed. + std::sync::Arc::make_mut(&mut idx.key_hash_to_vec_checksum) + .insert(key_hash, xxhash_rust::xxh64::xxh64(&blob, 0)); // Populate payload index with all HASH fields (for filtered search) let mut j = 1; diff --git a/src/text/store.rs b/src/text/store.rs index a540a93e..1421431e 100644 --- a/src/text/store.rs +++ b/src/text/store.rs @@ -1512,6 +1512,34 @@ impl TextStore { removed } + /// FLUSHALL/FLUSHDB parity (persistence-review R3): reset every text + /// index to an empty state (postings, term dicts, doc maps, TAG/NUMERIC + /// indexes) while KEEPING the FT.CREATE schema — mirroring restart + /// semantics. Without this, flushed hashes stayed matchable as ghost + /// documents until the next restart. + pub fn clear_all_contents(&mut self) { + #[cfg(feature = "text-index")] + { + let mut any = false; + for idx in self.indexes.values_mut() { + *idx = TextIndex::new_with_schema( + idx.name.clone(), + idx.key_prefixes.clone(), + idx.text_fields.clone(), + idx.tag_fields.clone(), + idx.numeric_fields.clone(), + idx.bm25_config, + ); + any = true; + } + if any { + self.bump_version(); + } + } + // Without the text-index feature no TextIndex constructor exists and + // no documents can have been indexed — nothing to clear. + } + /// Get a read-only reference to a text index. pub fn get_index(&self, name: &[u8]) -> Option<&TextIndex> { self.indexes.get(name) diff --git a/src/vector/background_compact.rs b/src/vector/background_compact.rs index 24a31734..f28fe2f8 100644 --- a/src/vector/background_compact.rs +++ b/src/vector/background_compact.rs @@ -21,6 +21,7 @@ //! `Arc` is `Send + Sync` — the worker only reads the source //! segments (graph and TQ codes); it never writes to their interior state. +use std::path::PathBuf; use std::sync::{Arc, OnceLock}; use crate::vector::segment::compaction::{self, CompactionError, MergeMode}; @@ -28,6 +29,13 @@ use crate::vector::segment::immutable::ImmutableSegment; use crate::vector::segment::mutable::FrozenSegment; use crate::vector::turbo_quant::collection::CollectionMetadata; +/// Where (and under what id) the worker should persist its output segment, +/// when the index has a `persist_dir` configured (B2, durability). The +/// segment id is allocated on the SHARD thread at submit time (never on the +/// worker), so ids never collide even though the write happens in the +/// background. +type PersistTarget = (PathBuf, u64); + /// The operation a worker should perform. /// /// Both variants produce the same `CompactionResult` so a single worker loop @@ -38,6 +46,7 @@ enum BuildOp { frozen: FrozenSegment, collection: Arc, seed: u64, + persist: Option, }, /// Many-immutables-to-one: merge N Arc'd immutable segments into one. /// @@ -50,6 +59,7 @@ enum BuildOp { seed: u64, mode: MergeMode, recall_tolerance: f32, + persist: Option, }, } @@ -100,19 +110,27 @@ impl BackgroundCompactor { frozen, collection, seed, - } => compaction::compact(&frozen, &collection, seed, None), + persist, + } => compaction::compact( + &frozen, + &collection, + seed, + persist.as_ref().map(|(p, id)| (p.as_path(), *id)), + ), BuildOp::Merge { segments, collection, seed, mode, recall_tolerance, + persist, } => compaction::merge_immutable( &segments, &collection, seed, mode, recall_tolerance, + persist.as_ref().map(|(p, id)| (p.as_path(), *id)), ), }; // If the receiver was dropped before we finished, swallow the error. @@ -131,6 +149,12 @@ impl BackgroundCompactor { /// Submit a mutable→immutable compaction job. /// + /// `persist`: when `Some((idx_dir, segment_id))`, the worker persists the + /// built segment to disk (staged + atomic rename) before replying — see + /// `segment_io::write_immutable_segment_staged`. `segment_id` MUST be + /// allocated by the caller (shard thread) beforehand; ids are never + /// allocated on the worker. + /// /// Returns an `Err` only when all workers have exited (compactor shut down). /// Returns the `reply_rx` — the caller polls it with [`try_recv`](flume::Receiver::try_recv). pub fn submit( @@ -138,6 +162,7 @@ impl BackgroundCompactor { frozen: FrozenSegment, collection: Arc, seed: u64, + persist: Option<(PathBuf, u64)>, ) -> Result, CompactionSubmitError> { let (reply_tx, reply_rx) = flume::bounded::(1); let job = WorkerJob { @@ -145,6 +170,7 @@ impl BackgroundCompactor { frozen, collection, seed, + persist, }, reply_tx, }; @@ -156,6 +182,9 @@ impl BackgroundCompactor { /// Submit an immutable→immutable merge job. /// + /// `persist`: same contract as [`Self::submit`] — merge persists + /// identically to compact. + /// /// The worker reads the source segment Arcs but never mutates them. /// Returns `reply_rx` — caller polls it non-blocking with `try_recv`. pub fn submit_merge( @@ -165,6 +194,7 @@ impl BackgroundCompactor { seed: u64, mode: MergeMode, recall_tolerance: f32, + persist: Option<(PathBuf, u64)>, ) -> Result, CompactionSubmitError> { let (reply_tx, reply_rx) = flume::bounded::(1); let job = WorkerJob { @@ -174,6 +204,7 @@ impl BackgroundCompactor { seed, mode, recall_tolerance, + persist, }, reply_tx, }; diff --git a/src/vector/distance/avx2.rs b/src/vector/distance/avx2.rs index 5b83cc70..0a42c2d5 100644 --- a/src/vector/distance/avx2.rs +++ b/src/vector/distance/avx2.rs @@ -376,6 +376,200 @@ pub unsafe fn sq8_stats(query: &[f32], codes: &[u8]) -> (f32, f32, f32) { (dot_sum, sum_c_sum, sumsq_c_sum) } +/// Int8 symmetric ADC per-candidate statistics (task #13): `(Σ qi8_i·c_i, +/// Σc_i, Σc_i²)` as exact `i64` integers, given an already-quantized `qi8` +/// query (see `turbo_quant::sq8::sq8_quantize_query_scalar`). +/// +/// The design doc (tmp/INT8-ADC-CONTEXT.md) originally called for +/// `VPMADDUBSW` (u8×i8 → i16 pairwise-add) directly, clamping the query to +/// `±63` to avoid its saturation trap (`255·127·2 = 64770 > i16::MAX`). +/// Recall A/B testing (`turbo_quant::sq8::tests::test_int8_adc_recall_ab_*`) +/// showed that clamp measurably regresses recall (R@10 delta ~0.017, +/// overlap ~0.975 vs the ≥0.98 gate) while the full `±127` range passes +/// comfortably (delta ~0.003-0.005, overlap ~0.988+). This kernel instead +/// avoids `VPMADDUBSW` entirely: widen both operands to i16 first +/// (`_mm256_cvtepu8_epi16` / `_mm256_cvtepi8_epi16` — exact, since u8 and +/// the `[-127,127]` qi8 range both fit i16 losslessly), multiply with +/// `_mm256_mullo_epi16` (exact — `|q·c| ≤ 127·255 = 32385 < i16::MAX`, a +/// single-element product never saturates), then widen-accumulate to i32 +/// via `_mm256_madd_epi16` against an all-ones vector (`VPMADDWD` genuinely +/// produces a 32-bit result with no saturation, unlike `VPMADDUBSW`'s +/// 16-bit output). No offset trick needed here (unlike the NEON kernel): +/// `cvtepu8_epi16`/`cvtepi8_epi16` already widen with the correct sign +/// handling for u8 vs i8 — `sum_qi8` is accepted only for signature parity +/// with [`super::neon::sq8_i8_stats`] / the scalar reference. +/// +/// `sum_c`/`sumsq_c` reuse the same widened `c_lo16`/`c_hi16` registers (no +/// extra widen pass), since both are query-independent. +/// +/// # Safety +/// Caller must ensure AVX2 is available (checked via +/// `is_x86_feature_detected!("avx2")` at table-init time). +/// +/// # Panics (debug only) +/// `debug_assert_eq!(qi8.len(), codes.len())`. +#[cfg(target_arch = "x86_64")] +#[inline] +#[target_feature(enable = "avx2")] +pub unsafe fn sq8_i8_stats(qi8: &[i8], codes: &[u8], _sum_qi8: i32) -> (i64, i64, i64) { + debug_assert_eq!(qi8.len(), codes.len(), "sq8_i8_stats: dimension mismatch"); + + let n = qi8.len(); + let pq = qi8.as_ptr(); + let pc = codes.as_ptr(); + + let mut dot_acc = _mm256_setzero_si256(); + let mut sum_c_acc = _mm256_setzero_si256(); + let mut sumsq_c_acc = _mm256_setzero_si256(); + let ones16 = _mm256_set1_epi16(1); + + let chunks = n / 32; + let mut i = 0usize; + + for _ in 0..chunks { + // SAFETY: i + 32 <= n guaranteed by chunks = n / 32. Both loads are + // unaligned 32-byte reads within bounds (i8 query / u8 codes). + let qv = _mm256_loadu_si256(pq.add(i) as *const __m256i); + let cv = _mm256_loadu_si256(pc.add(i) as *const __m256i); + + let q_lo16 = _mm256_cvtepi8_epi16(_mm256_castsi256_si128(qv)); + let q_hi16 = _mm256_cvtepi8_epi16(_mm256_extracti128_si256(qv, 1)); + let c_lo16 = _mm256_cvtepu8_epi16(_mm256_castsi256_si128(cv)); + let c_hi16 = _mm256_cvtepu8_epi16(_mm256_extracti128_si256(cv, 1)); + + let dot_lo16 = _mm256_mullo_epi16(q_lo16, c_lo16); + let dot_hi16 = _mm256_mullo_epi16(q_hi16, c_hi16); + dot_acc = _mm256_add_epi32(dot_acc, _mm256_madd_epi16(dot_lo16, ones16)); + dot_acc = _mm256_add_epi32(dot_acc, _mm256_madd_epi16(dot_hi16, ones16)); + + sum_c_acc = _mm256_add_epi32(sum_c_acc, _mm256_madd_epi16(c_lo16, ones16)); + sum_c_acc = _mm256_add_epi32(sum_c_acc, _mm256_madd_epi16(c_hi16, ones16)); + + sumsq_c_acc = _mm256_add_epi32(sumsq_c_acc, _mm256_madd_epi16(c_lo16, c_lo16)); + sumsq_c_acc = _mm256_add_epi32(sumsq_c_acc, _mm256_madd_epi16(c_hi16, c_hi16)); + + i += 32; + } + + // SAFETY: hsum_i32_avx2 requires AVX2, which we have via target_feature. + let mut dot: i64 = hsum_i32_avx2(dot_acc) as i64; + let mut sum_c: i64 = hsum_i32_avx2(sum_c_acc) as i64; + let mut sumsq_c: i64 = hsum_i32_avx2(sumsq_c_acc) as i64; + + // Scalar tail — safe indexing; at most one sub-vector-width pass, so + // bounds checks cost nothing and the unsafe surface stays confined to + // the intrinsics above (UNSAFE_POLICY). + for (&q, &c) in qi8[i..n].iter().zip(codes[i..n].iter()) { + let q = q as i64; + let c = c as i64; + dot += q * c; + sum_c += c; + sumsq_c += c * c; + } + + (dot, sum_c, sumsq_c) +} + +// ── f16 sidecar kernels (exact-rerank HQ-1) ───────────────────────────── + +/// Squared L2 between an f32 query and an f16-encoded sidecar vector. +/// F16C `vcvtph2ps` decodes 8 halves per step (hardware IEEE semantics — +/// subnormals, Inf, and NaN match scalar `f16_to_f32` exactly); FMA +/// accumulates. 16 halves per iteration, scalar tail. +/// +/// # Safety +/// Caller must ensure the CPU supports AVX2, F16C, and FMA +/// (verified via `is_x86_feature_detected!` at DistanceTable init). +#[cfg(target_arch = "x86_64")] +#[inline] +#[target_feature(enable = "avx2,f16c,fma")] +pub unsafe fn f16_l2(query: &[f32], vec_f16: &[u16]) -> f32 { + debug_assert_eq!(query.len(), vec_f16.len(), "f16_l2: dimension mismatch"); + let n = query.len(); + let mut sum0 = _mm256_setzero_ps(); + let mut sum1 = _mm256_setzero_ps(); + let pq = query.as_ptr(); + let px = vec_f16.as_ptr(); + let chunks = n / 16; + let mut i = 0usize; + for _ in 0..chunks { + // SAFETY: i + 16 <= n guaranteed by chunks = n / 16. Pointers are + // valid slices (f32 query / u16 halves) at this offset; loadu allows + // unaligned access. + let x0 = _mm256_cvtph_ps(_mm_loadu_si128(px.add(i) as *const __m128i)); + let x1 = _mm256_cvtph_ps(_mm_loadu_si128(px.add(i + 8) as *const __m128i)); + let q0 = _mm256_loadu_ps(pq.add(i)); + let q1 = _mm256_loadu_ps(pq.add(i + 8)); + let d0 = _mm256_sub_ps(q0, x0); + let d1 = _mm256_sub_ps(q1, x1); + sum0 = _mm256_fmadd_ps(d0, d0, sum0); + sum1 = _mm256_fmadd_ps(d1, d1, sum1); + i += 16; + } + // SAFETY: hsum_f32_avx2 requires AVX2, which we have via target_feature. + let mut result = hsum_f32_avx2(_mm256_add_ps(sum0, sum1)); + // Scalar tail + while i < n { + let d = *query.get_unchecked(i) - crate::vector::f16::f16_to_f32(*vec_f16.get_unchecked(i)); + result += d * d; + i += 1; + } + result +} + +/// Fused `(Σ q_i·x_i, Σ x_i²)` between an f32 query and an f16-encoded +/// sidecar vector (F16C decode + FMA, 16 halves per iteration; scalar +/// tail). Feeds the unit-sphere (Cosine/InnerProduct) exact-rerank +/// distance. +/// +/// # Safety +/// Caller must ensure the CPU supports AVX2, F16C, and FMA +/// (verified via `is_x86_feature_detected!` at DistanceTable init). +#[cfg(target_arch = "x86_64")] +#[inline] +#[target_feature(enable = "avx2,f16c,fma")] +pub unsafe fn f16_dot_normsq(query: &[f32], vec_f16: &[u16]) -> (f32, f32) { + debug_assert_eq!( + query.len(), + vec_f16.len(), + "f16_dot_normsq: dimension mismatch" + ); + let n = query.len(); + let mut dot0 = _mm256_setzero_ps(); + let mut dot1 = _mm256_setzero_ps(); + let mut xsq0 = _mm256_setzero_ps(); + let mut xsq1 = _mm256_setzero_ps(); + let pq = query.as_ptr(); + let px = vec_f16.as_ptr(); + let chunks = n / 16; + let mut i = 0usize; + for _ in 0..chunks { + // SAFETY: i + 16 <= n guaranteed by chunks = n / 16. Pointers are + // valid slices (f32 query / u16 halves) at this offset; loadu allows + // unaligned access. + let x0 = _mm256_cvtph_ps(_mm_loadu_si128(px.add(i) as *const __m128i)); + let x1 = _mm256_cvtph_ps(_mm_loadu_si128(px.add(i + 8) as *const __m128i)); + let q0 = _mm256_loadu_ps(pq.add(i)); + let q1 = _mm256_loadu_ps(pq.add(i + 8)); + dot0 = _mm256_fmadd_ps(q0, x0, dot0); + dot1 = _mm256_fmadd_ps(q1, x1, dot1); + xsq0 = _mm256_fmadd_ps(x0, x0, xsq0); + xsq1 = _mm256_fmadd_ps(x1, x1, xsq1); + i += 16; + } + // SAFETY: hsum_f32_avx2 requires AVX2, which we have via target_feature. + let mut dot = hsum_f32_avx2(_mm256_add_ps(dot0, dot1)); + let mut xsq = hsum_f32_avx2(_mm256_add_ps(xsq0, xsq1)); + // Scalar tail + while i < n { + let x = crate::vector::f16::f16_to_f32(*vec_f16.get_unchecked(i)); + dot += *query.get_unchecked(i) * x; + xsq += x * x; + i += 1; + } + (dot, xsq) +} + #[cfg(test)] #[cfg(target_arch = "x86_64")] mod tests { @@ -604,4 +798,163 @@ mod tests { let got = unsafe { sq8_stats(q, c) }; assert_eq!(got, (0.0, 0.0, 0.0)); } + + // ── task #13: int8 symmetric ADC stats (AVX2 vs scalar i64 oracle) ─── + // + // Gated on `has_avx2_fma()` — compiles on this aarch64 dev host (proving + // the cfg discipline holds) but no-ops at runtime; exercises for real on + // x86_64 CI/GCE runners. + + fn gen_i8_query(len: usize, seed: u32, qmax: i32) -> Vec { + let mut v = Vec::with_capacity(len); + let mut s = seed; + for _ in 0..len { + s = s.wrapping_mul(1664525).wrapping_add(1013904223); + let raw = (s >> 24) as i8 as i32; + v.push((raw % (qmax + 1)) as i8); + } + v + } + + #[test] + fn test_sq8_i8_stats_matches_scalar_oracle() { + if !has_avx2_fma() { + return; + } + use crate::vector::turbo_quant::sq8::sq8_candidate_stats_i8_scalar; + let qi8 = gen_i8_query(768, 42, 127); + let c = gen_u8(768, 99); + let sum_qi8: i32 = qi8.iter().map(|&x| x as i32).sum(); + let expected = sq8_candidate_stats_i8_scalar(&qi8, &c, sum_qi8); + // SAFETY: AVX2 verified above. + let got = unsafe { sq8_i8_stats(&qi8, &c, sum_qi8) }; + assert_eq!( + got, expected, + "sq8_i8_stats mismatch (exact int arithmetic): scalar={expected:?}, avx2={got:?}" + ); + } + + #[test] + fn test_sq8_i8_stats_tail_handling() { + if !has_avx2_fma() { + return; + } + use crate::vector::turbo_quant::sq8::sq8_candidate_stats_i8_scalar; + for len in [0, 1, 3, 7, 13, 15, 16, 17, 31, 32, 33, 63, 64, 100] { + let qi8 = gen_i8_query(len, 42, 127); + let c = gen_u8(len, 99); + let sum_qi8: i32 = qi8.iter().map(|&x| x as i32).sum(); + let expected = sq8_candidate_stats_i8_scalar(&qi8, &c, sum_qi8); + // SAFETY: AVX2 verified above. + let got = unsafe { sq8_i8_stats(&qi8, &c, sum_qi8) }; + assert_eq!( + got, expected, + "sq8_i8_stats tail len={len}: scalar={expected:?}, avx2={got:?}" + ); + } + } + + #[test] + fn test_sq8_i8_stats_empty() { + if !has_avx2_fma() { + return; + } + let qi8: &[i8] = &[]; + let c: &[u8] = &[]; + // SAFETY: AVX2 verified above. + let got = unsafe { sq8_i8_stats(qi8, c, 0) }; + assert_eq!(got, (0, 0, 0)); + } + + #[test] + fn test_sq8_i8_stats_extreme_values() { + if !has_avx2_fma() { + return; + } + use crate::vector::turbo_quant::sq8::sq8_candidate_stats_i8_scalar; + let dim = 768; + let qi8: Vec = (0..dim) + .map(|i| if i % 2 == 0 { 127i8 } else { -127i8 }) + .collect(); + let c: Vec = vec![255u8; dim]; + let sum_qi8: i32 = qi8.iter().map(|&x| x as i32).sum(); + let expected = sq8_candidate_stats_i8_scalar(&qi8, &c, sum_qi8); + // SAFETY: AVX2 verified above. + let got = unsafe { sq8_i8_stats(&qi8, &c, sum_qi8) }; + assert_eq!( + got, expected, + "extreme-value mismatch: scalar={expected:?}, avx2={got:?}" + ); + } + + fn has_f16c_kernel_features() -> bool { + is_x86_feature_detected!("avx2") + && is_x86_feature_detected!("fma") + && is_x86_feature_detected!("f16c") + } + + #[test] + fn test_f16_kernels_match_scalar_all_value_classes() { + use crate::vector::f16::{dot_normsq_f16, f32_to_f16, l2_sq_f16}; + if !has_f16c_kernel_features() { + return; + } + // Normals, f16 subnormals, zero, and odd tails -- F16C hardware decode + // must agree with scalar f16_to_f32 everywhere. + for len in [0, 1, 3, 7, 8, 9, 15, 16, 17, 31, 33, 100, 384] { + let q = gen_f32(len, 21); + let x: Vec = gen_f32(len, 77) + .iter() + .enumerate() + .map(|(j, &v)| { + if j % 5 == 4 { + f32_to_f16(3e-6 + (j as f32) * 1.1e-7) // subnormal f16 + } else if j % 7 == 6 { + f32_to_f16(0.0) + } else { + f32_to_f16(v) + } + }) + .collect(); + let exp_l2 = l2_sq_f16(&q, &x); + // SAFETY: AVX2+F16C+FMA verified above. + let got_l2 = unsafe { f16_l2(&q, &x) }; + let rel = (got_l2 - exp_l2).abs() / exp_l2.abs().max(1e-10); + assert!( + rel < 1e-4 || len == 0, + "f16_l2 len={len}: scalar={exp_l2}, avx2={got_l2}" + ); + + let (ed, en) = dot_normsq_f16(&q, &x); + // SAFETY: AVX2+F16C+FMA verified above. + let (gd, gn) = unsafe { f16_dot_normsq(&q, &x) }; + assert!( + (gd - ed).abs() / ed.abs().max(1e-6) < 1e-3, + "f16_dot len={len}: scalar={ed}, avx2={gd}" + ); + assert!( + (gn - en).abs() / en.abs().max(1e-6) < 1e-4, + "f16_normsq len={len}: scalar={en}, avx2={gn}" + ); + } + } + + #[test] + fn test_f16_kernels_inf_nan_propagate() { + use crate::vector::f16::f32_to_f16; + if !has_f16c_kernel_features() { + return; + } + let q = gen_f32(24, 5); + let mut x: Vec = gen_f32(24, 9).iter().map(|&v| f32_to_f16(v)).collect(); + x[10] = f32_to_f16(f32::INFINITY); + // SAFETY: AVX2+F16C+FMA verified above. + assert_eq!(unsafe { f16_l2(&q, &x) }, f32::INFINITY); + x[10] = f32_to_f16(f32::NAN); + // SAFETY: AVX2+F16C+FMA verified above. + assert!(unsafe { f16_l2(&q, &x) }.is_nan()); + // SAFETY: AVX2+F16C+FMA verified above. + let (_, xsq) = unsafe { f16_dot_normsq(&q, &x) }; + assert!(xsq.is_nan()); + } } diff --git a/src/vector/distance/avx512.rs b/src/vector/distance/avx512.rs index 626cd9f3..1dd393b0 100644 --- a/src/vector/distance/avx512.rs +++ b/src/vector/distance/avx512.rs @@ -1,9 +1,14 @@ //! AVX-512 distance kernels with 2x loop unrolling. //! //! All functions require AVX-512F at minimum. The i8 L2 kernel uses -//! `avx512bw` for byte-width operations. VNNI (`_mm512_dpwssd_epi32`) is not -//! yet stabilized in `core::arch::x86_64`, so we use the portable -//! `cvtepi8_epi16` + `madd_epi16` widening approach instead. +//! `avx512bw` for byte-width operations. `l2_i8_vnni` predates VNNI's +//! stabilization — despite the name, it uses the portable +//! `cvtepi8_epi16` + `madd_epi16` widening approach, not a real VNNI +//! instruction (left as-is; renaming is out of scope here). **Update:** +//! `_mm512_dpbusd_epi32` (`VPDPBUSD`, real AVX-512 VNNI) IS stable in +//! `core::arch::x86_64` as of rustc 1.94 (verified directly for task #13, +//! see [`sq8_i8_stats`]) — the "not yet stabilized" framing above is stale +//! for that specific intrinsic; VNNI-gated code is possible and used below. //! //! The caller (DistanceTable init) verifies AVX-512F via //! `is_x86_feature_detected!` before installing these function pointers. @@ -318,6 +323,94 @@ pub unsafe fn sq8_stats(query: &[f32], codes: &[u8]) -> (f32, f32, f32) { (dot_sum, sum_c_sum, sumsq_c_sum) } +/// Int8 symmetric ADC per-candidate statistics (task #13), true AVX-512 +/// VNNI: `(Σ qi8_i·c_i, Σc_i, Σc_i²)` as exact `i64` integers, given an +/// already-quantized `qi8` query (see +/// `turbo_quant::sq8::sq8_quantize_query_scalar`). +/// +/// Unlike `l2_i8_vnni` above (a historical misnomer — see the module docs), +/// this kernel uses the real VNNI instruction: `_mm512_dpbusd_epi32` +/// (`VPDPBUSD`, u8×i8 → i32 accumulate). `VPDPBUSD` has no saturation risk +/// (genuine 32-bit accumulate), so `dot_qc` and `sum_c` (`Σc_i·1`) both use +/// it directly at the full `±127` query range — no clamping, unlike the +/// AVX2 tier's originally-planned (and rejected, see `avx2::sq8_i8_stats` +/// docs) `VPMADDUBSW` design. +/// +/// `sumsq_c` (`Σc_i²`) can NOT reuse `_mm512_dpbusd_epi32` with `codes` as +/// both operands: its second operand must be **signed** i8 in `[-128,127]`, +/// and codes run `0..=255` — reinterpreting bytes above 127 as negative +/// would silently corrupt the square for the top half of the u8 range. So +/// `sumsq_c` instead widens `codes` to i16 (`_mm512_cvtepu8_epi16`, exact) +/// and reduces via `_mm512_madd_epi16` against itself (same shape as the +/// AVX2 kernel; `VPMADDWD`, no saturation risk: `255² · 2 = 130050` fits +/// i32 trivially). +/// +/// # Safety +/// Caller must ensure AVX-512F, AVX-512BW, and AVX-512VNNI are all +/// available (checked via `is_x86_feature_detected!` at table-init time). +/// +/// # Panics (debug only) +/// `debug_assert_eq!(qi8.len(), codes.len())`. +#[cfg(target_arch = "x86_64")] +#[inline] +#[target_feature(enable = "avx512f,avx512bw,avx512vnni")] +pub unsafe fn sq8_i8_stats(qi8: &[i8], codes: &[u8], _sum_qi8: i32) -> (i64, i64, i64) { + debug_assert_eq!(qi8.len(), codes.len(), "sq8_i8_stats: dimension mismatch"); + + let n = qi8.len(); + let pq = qi8.as_ptr(); + let pc = codes.as_ptr(); + + let mut dot_acc = _mm512_setzero_si512(); + let mut sum_c_acc = _mm512_setzero_si512(); + let mut sumsq_c_acc = _mm512_setzero_si512(); + let ones_i8 = _mm512_set1_epi8(1); + + let chunks = n / 64; + let mut i = 0usize; + + for _ in 0..chunks { + // SAFETY: i + 64 <= n guaranteed by chunks = n / 64. Both loads are + // unaligned 64-byte reads within bounds (i8 query / u8 codes). + let qv = _mm512_loadu_si512(pq.add(i) as *const __m512i); + let cv = _mm512_loadu_si512(pc.add(i) as *const __m512i); + + // dot_qc: Σ qi8·c — cv is the unsigned u8 operand, qv the signed + // i8 operand; exact, no saturation (VPDPBUSD accumulates directly + // into i32). + dot_acc = _mm512_dpbusd_epi32(dot_acc, cv, qv); + // sum_c: Σ c·1 — safe regardless of c's magnitude since the SIGNED + // operand is the constant 1, not c. + sum_c_acc = _mm512_dpbusd_epi32(sum_c_acc, cv, ones_i8); + + // sumsq_c: Σ c² via widen-to-i16 + madd_epi16 (see doc comment). + let c_lo16 = _mm512_cvtepu8_epi16(_mm512_castsi512_si256(cv)); + let c_hi16 = _mm512_cvtepu8_epi16(_mm512_extracti64x4_epi64(cv, 1)); + sumsq_c_acc = _mm512_add_epi32(sumsq_c_acc, _mm512_madd_epi16(c_lo16, c_lo16)); + sumsq_c_acc = _mm512_add_epi32(sumsq_c_acc, _mm512_madd_epi16(c_hi16, c_hi16)); + + i += 64; + } + + // SAFETY: _mm512_reduce_add_epi32 requires AVX-512F, verified via target_feature. + let mut dot: i64 = _mm512_reduce_add_epi32(dot_acc) as i64; + let mut sum_c: i64 = _mm512_reduce_add_epi32(sum_c_acc) as i64; + let mut sumsq_c: i64 = _mm512_reduce_add_epi32(sumsq_c_acc) as i64; + + // Scalar tail — safe indexing; at most one sub-vector-width pass, so + // bounds checks cost nothing and the unsafe surface stays confined to + // the intrinsics above (UNSAFE_POLICY). + for (&q, &c) in qi8[i..n].iter().zip(codes[i..n].iter()) { + let q = q as i64; + let c = c as i64; + dot += q * c; + sum_c += c; + sumsq_c += c * c; + } + + (dot, sum_c, sumsq_c) +} + #[cfg(test)] #[cfg(target_arch = "x86_64")] mod tests { @@ -502,4 +595,97 @@ mod tests { ); } } + + // ── task #13: int8 symmetric ADC stats (AVX-512 VNNI vs scalar oracle) ── + // + // Gated on `has_avx512vnni()` — compiles on this aarch64 dev host + // (proving the cfg discipline holds) but no-ops at runtime everywhere + // except a genuine VNNI-capable x86_64 host (Rosetta 2 does not emulate + // AVX-512 at all, so even x86_64-apple-darwin cross-builds no-op here). + + fn has_avx512vnni() -> bool { + is_x86_feature_detected!("avx512vnni") + } + + fn gen_i8_query(len: usize, seed: u32, qmax: i32) -> Vec { + let mut v = Vec::with_capacity(len); + let mut s = seed; + for _ in 0..len { + s = s.wrapping_mul(1664525).wrapping_add(1013904223); + let raw = (s >> 24) as i8 as i32; + v.push((raw % (qmax + 1)) as i8); + } + v + } + + #[test] + fn test_sq8_i8_stats_matches_scalar_oracle() { + if !has_avx512vnni() || !has_avx512bw() { + return; + } + use crate::vector::turbo_quant::sq8::sq8_candidate_stats_i8_scalar; + let qi8 = gen_i8_query(768, 42, 127); + let c = gen_u8(768, 99); + let sum_qi8: i32 = qi8.iter().map(|&x| x as i32).sum(); + let expected = sq8_candidate_stats_i8_scalar(&qi8, &c, sum_qi8); + // SAFETY: AVX-512F+BW+VNNI verified above. + let got = unsafe { sq8_i8_stats(&qi8, &c, sum_qi8) }; + assert_eq!( + got, expected, + "sq8_i8_stats mismatch (exact int arithmetic): scalar={expected:?}, avx512={got:?}" + ); + } + + #[test] + fn test_sq8_i8_stats_tail_handling() { + if !has_avx512vnni() || !has_avx512bw() { + return; + } + use crate::vector::turbo_quant::sq8::sq8_candidate_stats_i8_scalar; + for len in [0, 1, 3, 7, 13, 15, 31, 32, 33, 63, 64, 65, 100] { + let qi8 = gen_i8_query(len, 42, 127); + let c = gen_u8(len, 99); + let sum_qi8: i32 = qi8.iter().map(|&x| x as i32).sum(); + let expected = sq8_candidate_stats_i8_scalar(&qi8, &c, sum_qi8); + // SAFETY: AVX-512F+BW+VNNI verified above. + let got = unsafe { sq8_i8_stats(&qi8, &c, sum_qi8) }; + assert_eq!( + got, expected, + "sq8_i8_stats tail len={len}: scalar={expected:?}, avx512={got:?}" + ); + } + } + + #[test] + fn test_sq8_i8_stats_empty() { + if !has_avx512vnni() || !has_avx512bw() { + return; + } + let qi8: &[i8] = &[]; + let c: &[u8] = &[]; + // SAFETY: AVX-512F+BW+VNNI verified above. + let got = unsafe { sq8_i8_stats(qi8, c, 0) }; + assert_eq!(got, (0, 0, 0)); + } + + #[test] + fn test_sq8_i8_stats_extreme_values() { + if !has_avx512vnni() || !has_avx512bw() { + return; + } + use crate::vector::turbo_quant::sq8::sq8_candidate_stats_i8_scalar; + let dim = 768; + let qi8: Vec = (0..dim) + .map(|i| if i % 2 == 0 { 127i8 } else { -127i8 }) + .collect(); + let c: Vec = vec![255u8; dim]; + let sum_qi8: i32 = qi8.iter().map(|&x| x as i32).sum(); + let expected = sq8_candidate_stats_i8_scalar(&qi8, &c, sum_qi8); + // SAFETY: AVX-512F+BW+VNNI verified above. + let got = unsafe { sq8_i8_stats(&qi8, &c, sum_qi8) }; + assert_eq!( + got, expected, + "extreme-value mismatch: scalar={expected:?}, avx512={got:?}" + ); + } } diff --git a/src/vector/distance/mod.rs b/src/vector/distance/mod.rs index 3d23ed02..b046a843 100644 --- a/src/vector/distance/mod.rs +++ b/src/vector/distance/mod.rs @@ -44,10 +44,48 @@ pub struct DistanceTable { /// to get the final ADC distance — that combine step is O(1) arithmetic, /// architecture-independent, and deliberately NOT part of this table. pub sq8_stats: fn(&[f32], &[u8]) -> (f32, f32, f32), + /// Int8 symmetric ADC per-candidate stats (task #13): `(qi8, codes, + /// sum_qi8) -> (Σ qi8_i·c_i, Σc_i, Σc_i²)` as exact `i64` integers. + /// `None` means the int8 tier is unavailable (no matching SIMD kernel + /// for this CPU, or the `MOON_SQ8_INT8_ADC=0` escape hatch is set) — + /// callers must fall back to `sq8_stats` (f32) in that case. See + /// `turbo_quant::sq8` module docs for the query-quantization step + /// ([`crate::vector::turbo_quant::sq8::sq8_quantize_query_scalar`], + /// [`crate::vector::turbo_quant::sq8::SQ8_INT8_QMAX`]) that must run + /// once per query before calling this. + pub sq8_i8_stats: Option, + /// Squared L2 between an f32 query and an f16-encoded vector — the + /// exact-rerank sidecar (HQ-1) decode fused with the distance loop. + /// SIMD tiers decode 8 halves per step (NEON integer rescale / F16C + /// `vcvtph2ps`); scalar falls back to `vector::f16::l2_sq_f16`. + pub f16_l2: fn(&[f32], &[u16]) -> f32, + /// Fused `(Σ q_i·x_i, Σ x_i²)` between an f32 query and an f16-encoded + /// vector — the unit-sphere (Cosine/IP) rerank pass needs both in one + /// decode sweep. Same tiering as [`Self::f16_l2`]. + pub f16_dot_normsq: fn(&[f32], &[u16]) -> (f32, f32), } +/// Signature aliases for the f16 sidecar kernels (keeps branch wiring terse). +#[cfg(target_arch = "x86_64")] +type F16L2Fn = fn(&[f32], &[u16]) -> f32; +#[cfg(target_arch = "x86_64")] +type F16DotNormFn = fn(&[f32], &[u16]) -> (f32, f32); + static DISTANCE_TABLE: OnceLock = OnceLock::new(); +/// Escape hatch for the int8 symmetric ADC path (task #13): `MOON_SQ8_INT8_ADC=0` +/// forces every tier back to the f32 `sq8_stats` kernel (bench/diagnostic +/// convention, like the `MOON_XSHARD_*` knobs — see CLAUDE.md env var list). +/// Cached in a `OnceLock` so the env lookup happens at most once. +fn int8_adc_enabled() -> bool { + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| { + std::env::var("MOON_SQ8_INT8_ADC") + .map(|v| v != "0") + .unwrap_or(true) + }) +} + /// Initialize the distance dispatch table. /// /// Detects CPU features at runtime and selects the fastest kernel tier: @@ -65,8 +103,43 @@ pub fn init() { DISTANCE_TABLE.get_or_init(|| { #[cfg(target_arch = "x86_64")] { + // f16 sidecar kernels: F16C is a separate cpuid bit from AVX2 + // (present on virtually every AVX2 CPU, but checked explicitly — + // and AVX2 is NOT implied by F16C+FMA: AMD Piledriver has both + // without AVX2). Both AVX2 and AVX-512 tiers share this kernel. + let (f16_l2, f16_dot_normsq): (F16L2Fn, F16DotNormFn) = + if is_x86_feature_detected!("f16c") + && is_x86_feature_detected!("fma") + && is_x86_feature_detected!("avx2") + { + ( + |q, x| { + // SAFETY: AVX2+F16C+FMA verified by is_x86_feature_detected! above. + unsafe { avx2::f16_l2(q, x) } + }, + |q, x| { + // SAFETY: AVX2+F16C+FMA verified by is_x86_feature_detected! above. + unsafe { avx2::f16_dot_normsq(q, x) } + }, + ) + } else { + ( + crate::vector::f16::l2_sq_f16 as F16L2Fn, + crate::vector::f16::dot_normsq_f16 as F16DotNormFn, + ) + }; #[cfg(feature = "simd-avx512")] if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bw") { + let sq8_i8_stats = if int8_adc_enabled() && is_x86_feature_detected!("avx512vnni") { + Some( + (|q, c, s| { + // SAFETY: AVX-512F+BW+VNNI verified by is_x86_feature_detected! above. + unsafe { avx512::sq8_i8_stats(q, c, s) } + }) as crate::vector::turbo_quant::sq8::Sq8I8StatsFn, + ) + } else { + None + }; return DistanceTable { l2_f32: |a, b| { // SAFETY: AVX-512F verified by is_x86_feature_detected! above. @@ -89,9 +162,22 @@ pub fn init() { // SAFETY: AVX-512F verified by is_x86_feature_detected! above. unsafe { avx512::sq8_stats(q, c) } }, + sq8_i8_stats, + f16_l2, + f16_dot_normsq, }; } if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") { + let sq8_i8_stats = if int8_adc_enabled() { + Some( + (|q, c, s| { + // SAFETY: AVX2 verified by is_x86_feature_detected! above. + unsafe { avx2::sq8_i8_stats(q, c, s) } + }) as crate::vector::turbo_quant::sq8::Sq8I8StatsFn, + ) + } else { + None + }; return DistanceTable { l2_f32: |a, b| { // SAFETY: AVX2+FMA verified by is_x86_feature_detected! above. @@ -114,6 +200,9 @@ pub fn init() { // SAFETY: AVX2+FMA verified by is_x86_feature_detected! above. unsafe { avx2::sq8_stats(q, c) } }, + sq8_i8_stats, + f16_l2, + f16_dot_normsq, }; } } @@ -121,6 +210,16 @@ pub fn init() { #[cfg(target_arch = "aarch64")] { // NEON is baseline on all AArch64 CPUs — always available. + let sq8_i8_stats = if int8_adc_enabled() { + Some( + (|q, c, s| { + // SAFETY: NEON is guaranteed on AArch64. + unsafe { neon::sq8_i8_stats(q, c, s) } + }) as crate::vector::turbo_quant::sq8::Sq8I8StatsFn, + ) + } else { + None + }; return DistanceTable { l2_f32: |a, b| { // SAFETY: NEON is guaranteed on AArch64. @@ -144,10 +243,21 @@ pub fn init() { // SAFETY: NEON is guaranteed on AArch64. unsafe { neon::sq8_stats(q, c) } }, + sq8_i8_stats, + f16_l2: |q, x| { + // SAFETY: NEON is guaranteed on AArch64. + unsafe { neon::f16_l2(q, x) } + }, + f16_dot_normsq: |q, x| { + // SAFETY: NEON is guaranteed on AArch64. + unsafe { neon::f16_dot_normsq(q, x) } + }, }; } - // Scalar fallback — works on every platform. + // Scalar fallback — works on every platform. No int8 SIMD tier + // exists for pure-scalar builds (no proven throughput win without + // SIMD widening); sq8_i8_stats is None, callers use sq8_stats (f32). #[allow(unreachable_code)] DistanceTable { l2_f32: scalar::l2_f32, @@ -156,6 +266,9 @@ pub fn init() { cosine_f32: scalar::cosine_f32, tq_l2: crate::vector::turbo_quant::tq_adc::tq_l2_adc_scaled, sq8_stats: crate::vector::turbo_quant::sq8::sq8_candidate_stats_scalar, + sq8_i8_stats: None, + f16_l2: crate::vector::f16::l2_sq_f16, + f16_dot_normsq: crate::vector::f16::dot_normsq_f16, } }); } @@ -511,4 +624,117 @@ mod integration_tests { ); } } + + /// Task #13: if this CPU/build has an int8 SIMD tier installed + /// (`sq8_i8_stats.is_some()`), it must match the scalar i64 oracle + /// exactly — same bit-for-bit contract as `sq8_stats` above, but exact + /// integers instead of float tolerance. On this dev host that means the + /// NEON tier; on a CPU/build with the tier unavailable (or + /// `MOON_SQ8_INT8_ADC=0`) the table has `None` and the test is a no-op + /// (asserting absence is not meaningful — the whole point of `None` is + /// "fall back to sq8_stats", verified by the call-site wiring instead). + #[test] + fn test_simd_matches_scalar_sq8_i8_stats() { + use crate::vector::turbo_quant::sq8::{ + SQ8_INT8_QMAX, sq8_candidate_stats_i8_scalar, sq8_quantize_query_scalar, + }; + + init(); + let t = table(); + let Some(i8_stats_fn) = t.sq8_i8_stats else { + return; + }; + for &dim in TEST_DIMS { + let q = deterministic_f32(dim, 42); + let codes = deterministic_u8(dim, 99); + let mut qi8 = vec![0i8; dim]; + let (_q_scale, sum_qi8) = sq8_quantize_query_scalar(&q, SQ8_INT8_QMAX, &mut qi8); + let expected = sq8_candidate_stats_i8_scalar(&qi8, &codes, sum_qi8); + let got = i8_stats_fn(&qi8, &codes, sum_qi8); + assert_eq!( + got, expected, + "sq8_i8_stats mismatch at dim={dim}: scalar={expected:?} dispatch={got:?}" + ); + } + } + + /// Task #13: `MOON_SQ8_INT8_ADC` escape hatch — this test only asserts + /// the env var parses without panicking and is idempotent across calls + /// (the table itself is a `OnceLock` initialized once per process, so + /// this test cannot flip live dispatch — that is exercised manually / + /// by CI running the suite once with the var set, see CLAUDE.md). + #[test] + fn test_int8_adc_enabled_reads_env_without_panicking() { + let _ = super::int8_adc_enabled(); + let _ = super::int8_adc_enabled(); + } + + /// Encode a deterministic f32 vector as f16 bits (rerank sidecar layout). + fn deterministic_f16(dim: usize, seed: u64) -> Vec { + let mut out = Vec::new(); + crate::vector::f16::encode_f16_slice(&deterministic_f32(dim, seed), &mut out); + out + } + + #[test] + fn test_dispatched_f16_kernels_match_scalar() { + init(); + let t = table(); + for &dim in TEST_DIMS { + let q = deterministic_f32(dim, 21); + let x16 = deterministic_f16(dim, 42); + + let scalar_l2 = crate::vector::f16::l2_sq_f16(&q, &x16); + let fast_l2 = (t.f16_l2)(&q, &x16); + assert!( + approx_eq_f32(scalar_l2, fast_l2, 1e-4), + "f16_l2 mismatch at dim={dim}: scalar={scalar_l2} dispatched={fast_l2}" + ); + + let (sd, sn) = crate::vector::f16::dot_normsq_f16(&q, &x16); + let (fd, fn_) = (t.f16_dot_normsq)(&q, &x16); + assert!( + approx_eq_f32(sd, fd, 1e-4), + "f16_dot mismatch at dim={dim}: scalar={sd} dispatched={fd}" + ); + assert!( + approx_eq_f32(sn, fn_, 1e-4), + "f16_normsq mismatch at dim={dim}: scalar={sn} dispatched={fn_}" + ); + } + } + + #[test] + fn test_dispatched_f16_kernels_subnormal_and_special() { + use crate::vector::f16::f32_to_f16; + init(); + let t = table(); + + // Subnormal-heavy vector: 3e-6..6e-5 land in the f16 subnormal range — + // the integer-rescale decode must handle them exactly. + let dim = 40; + let q = deterministic_f32(dim, 5); + let tiny: Vec = (0..dim) + .map(|i| f32_to_f16(3e-6 + (i as f32) * 1.5e-6)) + .collect(); + let scalar_l2 = crate::vector::f16::l2_sq_f16(&q, &tiny); + let fast_l2 = (t.f16_l2)(&q, &tiny); + assert!( + approx_eq_f32(scalar_l2, fast_l2, 1e-4), + "subnormal f16_l2: scalar={scalar_l2} dispatched={fast_l2}" + ); + + // Infinity in the encoded vector must propagate to an infinite + // distance / dot, exactly like the scalar reference. + let mut with_inf = deterministic_f16(dim, 9); + with_inf[17] = f32_to_f16(f32::INFINITY); + assert_eq!((t.f16_l2)(&q, &with_inf), f32::INFINITY); + let (_, xsq) = (t.f16_dot_normsq)(&q, &with_inf); + assert_eq!(xsq, f32::INFINITY); + + // NaN must propagate as NaN. + let mut with_nan = deterministic_f16(dim, 11); + with_nan[3] = f32_to_f16(f32::NAN); + assert!((t.f16_l2)(&q, &with_nan).is_nan()); + } } diff --git a/src/vector/distance/neon.rs b/src/vector/distance/neon.rs index 290d3fa4..90f5f3dd 100644 --- a/src/vector/distance/neon.rs +++ b/src/vector/distance/neon.rs @@ -357,6 +357,226 @@ pub unsafe fn sq8_stats(query: &[f32], codes: &[u8]) -> (f32, f32, f32) { (dot_sum, sum_c_sum, sumsq_c_sum) } +/// Int8 symmetric ADC per-candidate statistics (task #13): `(Σ qi8_i·c_i, +/// Σc_i, Σc_i²)` as exact `i64` integers, given an already-quantized `qi8` +/// query (see `turbo_quant::sq8::sq8_quantize_query_scalar`). +/// +/// Baseline ARMv8-A NEON has no mixed signed×unsigned widening multiply +/// (that needs the `dotprod`/`i8mm` extensions — `vdotq_s32`/SDOT is +/// nightly-gated in `core::arch` as of rustc 1.94, see module docs below), +/// so this uses the **offset trick**: reinterpret each code byte `c` (u8, +/// `0..=255`) as `c' = c ^ 0x80` (i8, `= c - 128`), widen-multiply `qi8 * c'` +/// with `vmull_s8` (signed×signed, always available), and undo the shift +/// afterwards: `Σ qi8·c = Σ qi8·(c-128) + 128·Σqi8`. `sum_qi8` is the +/// precomputed per-query `Σqi8` needed for that correction. +/// +/// `sum_c`/`sumsq_c` are query-independent and computed directly from the +/// u8 codes (`vmovl_u8` widen for `sum_c`, `vmull_u8` widen-square for +/// `sumsq_c`) — both exact, no offset trick needed there. +/// +/// # Safety +/// Caller must ensure the CPU supports NEON (baseline on all AArch64). +/// +/// # Panics (debug only) +/// `debug_assert_eq!(qi8.len(), codes.len())`. +#[cfg(target_arch = "aarch64")] +#[inline] +#[target_feature(enable = "neon")] +pub unsafe fn sq8_i8_stats(qi8: &[i8], codes: &[u8], sum_qi8: i32) -> (i64, i64, i64) { + debug_assert_eq!(qi8.len(), codes.len(), "sq8_i8_stats: dimension mismatch"); + + let n = qi8.len(); + let pq = qi8.as_ptr(); + let pc = codes.as_ptr(); + + let mut dot_acc = vdupq_n_s32(0); + let mut sum_c_acc = vdupq_n_u32(0); + let mut sumsq_c_acc = vdupq_n_u32(0); + + let chunks = n / 16; + let mut i = 0usize; + + for _ in 0..chunks { + // SAFETY: i + 16 <= n guaranteed by chunks = n / 16. Pointers are + // valid for 16 elements (i8 query / u8 codes) at this offset. + let qv = vld1q_s8(pq.add(i)); + let cv_u8 = vld1q_u8(pc.add(i)); + // Offset trick: reinterpret (c ^ 0x80) as i8 == c - 128. + let cv_off = vreinterpretq_s8_u8(veorq_u8(cv_u8, vdupq_n_u8(0x80))); + + let dot_lo = vmull_s8(vget_low_s8(qv), vget_low_s8(cv_off)); + let dot_hi = vmull_s8(vget_high_s8(qv), vget_high_s8(cv_off)); + dot_acc = vpadalq_s16(dot_acc, dot_lo); + dot_acc = vpadalq_s16(dot_acc, dot_hi); + + sum_c_acc = vpadalq_u16(sum_c_acc, vmovl_u8(vget_low_u8(cv_u8))); + sum_c_acc = vpadalq_u16(sum_c_acc, vmovl_u8(vget_high_u8(cv_u8))); + + let sq_lo = vmull_u8(vget_low_u8(cv_u8), vget_low_u8(cv_u8)); + let sq_hi = vmull_u8(vget_high_u8(cv_u8), vget_high_u8(cv_u8)); + sumsq_c_acc = vpadalq_u16(sumsq_c_acc, sq_lo); + sumsq_c_acc = vpadalq_u16(sumsq_c_acc, sq_hi); + + i += 16; + } + + let mut dot_arr = [0i32; 4]; + let mut sumc_arr = [0u32; 4]; + let mut sumsq_arr = [0u32; 4]; + // SAFETY: dot_arr/sumc_arr/sumsq_arr are stack arrays of exactly 4 + // lanes, matching the int32x4_t/uint32x4_t store width. + vst1q_s32(dot_arr.as_mut_ptr(), dot_acc); + vst1q_u32(sumc_arr.as_mut_ptr(), sum_c_acc); + vst1q_u32(sumsq_arr.as_mut_ptr(), sumsq_c_acc); + + // Σ qi8·(c-128) so far (offset form); corrected to Σ qi8·c below. + let mut dot_offset: i64 = dot_arr.iter().map(|&x| x as i64).sum(); + let mut sum_c: i64 = sumc_arr.iter().map(|&x| x as i64).sum(); + let mut sumsq_c: i64 = sumsq_arr.iter().map(|&x| x as i64).sum(); + + // Scalar tail — direct (no offset trick needed, exact i64 arithmetic). + // Safe indexing: at most one sub-vector-width pass, so bounds checks + // cost nothing and the unsafe surface stays confined to the intrinsics + // above (UNSAFE_POLICY). + for (&q, &c) in qi8[i..n].iter().zip(codes[i..n].iter()) { + let q = q as i64; + let c = c as i64; + dot_offset += q * (c - 128); + sum_c += c; + sumsq_c += c * c; + } + + let dot = dot_offset + 128i64 * sum_qi8 as i64; + (dot, sum_c, sumsq_c) +} + +// ── f16 sidecar kernels (exact-rerank HQ-1) ───────────────────────────── + +/// Decode 4 IEEE binary16 lanes to f32 using baseline-NEON integer ops (no +/// ARMv8.2 FP16 intrinsics required): shift sign/magnitude into f32 bit +/// positions, then rescale by 2^112 to re-bias the exponent (f16 bias 15 → +/// f32 bias 127). The rescale multiply is exact — it is a power of two and +/// every finite f16 is representable in f32 — and f16 subnormals normalize +/// through the same multiply. Lanes with exponent 0x1F (Inf/NaN) are rebuilt +/// explicitly via a bit-select, matching scalar `f16_to_f32` on all inputs. +/// +/// # Safety +/// Caller must ensure the CPU supports NEON (baseline on all AArch64). +#[cfg(target_arch = "aarch64")] +#[inline] +#[target_feature(enable = "neon")] +unsafe fn f16x4_decode(h: uint16x4_t) -> float32x4_t { + let bits = vmovl_u16(h); + let sign = vshlq_n_u32(vandq_u32(bits, vdupq_n_u32(0x8000)), 16); + let mag = vshlq_n_u32(vandq_u32(bits, vdupq_n_u32(0x7FFF)), 13); + // After the shift the f32 reads 2^(e−127)·1.m for f16 exponent e; the + // true f16 value is 2^(e−15)·1.m — multiply by 2^112 (bits 0x7780_0000). + let magf = vmulq_f32( + vreinterpretq_f32_u32(mag), + vdupq_n_f32(f32::from_bits(0x7780_0000)), + ); + let is_special = vceqq_u32(vandq_u32(bits, vdupq_n_u32(0x7C00)), vdupq_n_u32(0x7C00)); + let special = vorrq_u32( + vdupq_n_u32(0x7F80_0000), + vshlq_n_u32(vandq_u32(bits, vdupq_n_u32(0x03FF)), 13), + ); + let val = vbslq_u32(is_special, special, vreinterpretq_u32_f32(magf)); + vreinterpretq_f32_u32(vorrq_u32(val, sign)) +} + +/// Squared L2 between an f32 query and an f16-encoded sidecar vector +/// (NEON, 8 halves per iteration; scalar tail). +/// +/// # Safety +/// Caller must ensure the CPU supports NEON (baseline on all AArch64). +#[cfg(target_arch = "aarch64")] +#[inline] +#[target_feature(enable = "neon")] +pub unsafe fn f16_l2(query: &[f32], vec_f16: &[u16]) -> f32 { + debug_assert_eq!(query.len(), vec_f16.len(), "f16_l2: dimension mismatch"); + let n = query.len(); + let mut sum0 = vdupq_n_f32(0.0); + let mut sum1 = vdupq_n_f32(0.0); + let pq = query.as_ptr(); + let px = vec_f16.as_ptr(); + let chunks = n / 8; + let mut i = 0usize; + for _ in 0..chunks { + // SAFETY: i + 8 <= n guaranteed by chunks = n / 8. Pointers are + // valid slices (f32 query / u16 halves) at this offset. + let h = vld1q_u16(px.add(i)); + let x0 = f16x4_decode(vget_low_u16(h)); + let x1 = f16x4_decode(vget_high_u16(h)); + let q0 = vld1q_f32(pq.add(i)); + let q1 = vld1q_f32(pq.add(i + 4)); + let d0 = vsubq_f32(q0, x0); + let d1 = vsubq_f32(q1, x1); + sum0 = vfmaq_f32(sum0, d0, d0); + sum1 = vfmaq_f32(sum1, d1, d1); + i += 8; + } + // SAFETY: vaddvq_f32 requires NEON, which we have via target_feature. + let mut result = vaddvq_f32(vaddq_f32(sum0, sum1)); + // Scalar tail + while i < n { + let d = *query.get_unchecked(i) - crate::vector::f16::f16_to_f32(*vec_f16.get_unchecked(i)); + result += d * d; + i += 1; + } + result +} + +/// Fused `(Σ q_i·x_i, Σ x_i²)` between an f32 query and an f16-encoded +/// sidecar vector (NEON, 8 halves per iteration; scalar tail). Feeds the +/// unit-sphere (Cosine/InnerProduct) exact-rerank distance. +/// +/// # Safety +/// Caller must ensure the CPU supports NEON (baseline on all AArch64). +#[cfg(target_arch = "aarch64")] +#[inline] +#[target_feature(enable = "neon")] +pub unsafe fn f16_dot_normsq(query: &[f32], vec_f16: &[u16]) -> (f32, f32) { + debug_assert_eq!( + query.len(), + vec_f16.len(), + "f16_dot_normsq: dimension mismatch" + ); + let n = query.len(); + let mut dot0 = vdupq_n_f32(0.0); + let mut dot1 = vdupq_n_f32(0.0); + let mut xsq0 = vdupq_n_f32(0.0); + let mut xsq1 = vdupq_n_f32(0.0); + let pq = query.as_ptr(); + let px = vec_f16.as_ptr(); + let chunks = n / 8; + let mut i = 0usize; + for _ in 0..chunks { + // SAFETY: i + 8 <= n guaranteed by chunks = n / 8. Pointers are + // valid slices (f32 query / u16 halves) at this offset. + let h = vld1q_u16(px.add(i)); + let x0 = f16x4_decode(vget_low_u16(h)); + let x1 = f16x4_decode(vget_high_u16(h)); + let q0 = vld1q_f32(pq.add(i)); + let q1 = vld1q_f32(pq.add(i + 4)); + dot0 = vfmaq_f32(dot0, q0, x0); + dot1 = vfmaq_f32(dot1, q1, x1); + xsq0 = vfmaq_f32(xsq0, x0, x0); + xsq1 = vfmaq_f32(xsq1, x1, x1); + i += 8; + } + // SAFETY: vaddvq_f32 requires NEON, which we have via target_feature. + let mut dot = vaddvq_f32(vaddq_f32(dot0, dot1)); + let mut xsq = vaddvq_f32(vaddq_f32(xsq0, xsq1)); + // Scalar tail + while i < n { + let x = crate::vector::f16::f16_to_f32(*vec_f16.get_unchecked(i)); + dot += *query.get_unchecked(i) * x; + xsq += x * x; + i += 1; + } + (dot, xsq) +} + #[cfg(test)] #[cfg(target_arch = "aarch64")] mod tests { @@ -546,4 +766,142 @@ mod tests { let got = unsafe { sq8_stats(q, c) }; assert_eq!(got, (0.0, 0.0, 0.0)); } + + // ── task #13: int8 symmetric ADC stats (NEON vs scalar i64 oracle) ─── + + fn gen_i8_query(len: usize, seed: u32, qmax: i32) -> Vec { + // Values within [-qmax, qmax], mirroring what + // sq8_quantize_query_scalar would actually produce. + let mut v = Vec::with_capacity(len); + let mut s = seed; + for _ in 0..len { + s = s.wrapping_mul(1664525).wrapping_add(1013904223); + let raw = (s >> 24) as i8 as i32; + v.push((raw % (qmax + 1)) as i8); + } + v + } + + #[test] + fn test_sq8_i8_stats_matches_scalar_oracle() { + use crate::vector::turbo_quant::sq8::sq8_candidate_stats_i8_scalar; + let qi8 = gen_i8_query(768, 42, 127); + let c = gen_u8(768, 99); + let sum_qi8: i32 = qi8.iter().map(|&x| x as i32).sum(); + let expected = sq8_candidate_stats_i8_scalar(&qi8, &c, sum_qi8); + // SAFETY: NEON is baseline on AArch64. + let got = unsafe { sq8_i8_stats(&qi8, &c, sum_qi8) }; + assert_eq!( + got, expected, + "sq8_i8_stats mismatch (exact int arithmetic): scalar={expected:?}, neon={got:?}" + ); + } + + #[test] + fn test_sq8_i8_stats_tail_handling() { + use crate::vector::turbo_quant::sq8::sq8_candidate_stats_i8_scalar; + for len in [0, 1, 3, 7, 13, 15, 16, 17, 31, 33, 100] { + let qi8 = gen_i8_query(len, 42, 127); + let c = gen_u8(len, 99); + let sum_qi8: i32 = qi8.iter().map(|&x| x as i32).sum(); + let expected = sq8_candidate_stats_i8_scalar(&qi8, &c, sum_qi8); + // SAFETY: NEON is baseline on AArch64. + let got = unsafe { sq8_i8_stats(&qi8, &c, sum_qi8) }; + assert_eq!( + got, expected, + "sq8_i8_stats tail len={len}: scalar={expected:?}, neon={got:?}" + ); + } + } + + #[test] + fn test_sq8_i8_stats_empty() { + let qi8: &[i8] = &[]; + let c: &[u8] = &[]; + // SAFETY: NEON is baseline on AArch64. + let got = unsafe { sq8_i8_stats(qi8, c, 0) }; + assert_eq!(got, (0, 0, 0)); + } + + #[test] + fn test_sq8_i8_stats_extreme_values() { + // All-max-magnitude inputs (qi8=±127, c=255) stress the overflow + // bound: worst case |dot_offset| per element is 127*128=16256, + // times dim=768 ≈ 12.5M — comfortably inside i32, but this pins the + // exact result down against the oracle rather than trusting the + // arithmetic argument alone. + use crate::vector::turbo_quant::sq8::sq8_candidate_stats_i8_scalar; + let dim = 768; + let qi8: Vec = (0..dim) + .map(|i| if i % 2 == 0 { 127i8 } else { -127i8 }) + .collect(); + let c: Vec = vec![255u8; dim]; + let sum_qi8: i32 = qi8.iter().map(|&x| x as i32).sum(); + let expected = sq8_candidate_stats_i8_scalar(&qi8, &c, sum_qi8); + // SAFETY: NEON is baseline on AArch64. + let got = unsafe { sq8_i8_stats(&qi8, &c, sum_qi8) }; + assert_eq!( + got, expected, + "extreme-value mismatch: scalar={expected:?}, neon={got:?}" + ); + } + + #[test] + fn test_f16_kernels_match_scalar_all_value_classes() { + use crate::vector::f16::{dot_normsq_f16, f32_to_f16, l2_sq_f16}; + // Normals, f16 subnormals (≈3e-6..6e-5), zero, and odd tails — the + // integer-rescale decode must agree with scalar f16_to_f32 everywhere. + for len in [0, 1, 3, 7, 8, 9, 15, 16, 17, 100, 384] { + let q = gen_f32(len, 21); + let x: Vec = gen_f32(len, 77) + .iter() + .enumerate() + .map(|(j, &v)| { + if j % 5 == 4 { + f32_to_f16(3e-6 + (j as f32) * 1.1e-7) // subnormal f16 + } else if j % 7 == 6 { + f32_to_f16(0.0) + } else { + f32_to_f16(v) + } + }) + .collect(); + let exp_l2 = l2_sq_f16(&q, &x); + // SAFETY: NEON is baseline on AArch64. + let got_l2 = unsafe { f16_l2(&q, &x) }; + let rel = (got_l2 - exp_l2).abs() / exp_l2.abs().max(1e-10); + assert!( + rel < 1e-4 || len == 0, + "f16_l2 len={len}: scalar={exp_l2}, neon={got_l2}" + ); + + let (ed, en) = dot_normsq_f16(&q, &x); + // SAFETY: NEON is baseline on AArch64. + let (gd, gn) = unsafe { f16_dot_normsq(&q, &x) }; + assert!( + (gd - ed).abs() / ed.abs().max(1e-6) < 1e-3, + "f16_dot len={len}: scalar={ed}, neon={gd}" + ); + assert!( + (gn - en).abs() / en.abs().max(1e-6) < 1e-4, + "f16_normsq len={len}: scalar={en}, neon={gn}" + ); + } + } + + #[test] + fn test_f16_kernels_inf_nan_propagate() { + use crate::vector::f16::f32_to_f16; + let q = gen_f32(24, 5); + let mut x: Vec = gen_f32(24, 9).iter().map(|&v| f32_to_f16(v)).collect(); + x[10] = f32_to_f16(f32::INFINITY); + // SAFETY: NEON is baseline on AArch64. + assert_eq!(unsafe { f16_l2(&q, &x) }, f32::INFINITY); + x[10] = f32_to_f16(f32::NAN); + // SAFETY: NEON is baseline on AArch64. + assert!(unsafe { f16_l2(&q, &x) }.is_nan()); + // SAFETY: NEON is baseline on AArch64. + let (_, xsq) = unsafe { f16_dot_normsq(&q, &x) }; + assert!(xsq.is_nan()); + } } diff --git a/src/vector/f16.rs b/src/vector/f16.rs index ac3f6236..fd32fdc6 100644 --- a/src/vector/f16.rs +++ b/src/vector/f16.rs @@ -104,6 +104,10 @@ pub fn encode_f16_slice(src: &[f32], out: &mut Vec) { } /// Squared L2 between an f32 query and an f16-encoded vector. +/// +/// Scalar reference; the rerank hot path goes through +/// `vector::distance::table().f16_l2`, which installs a SIMD variant +/// (NEON / F16C) when available and falls back to this. #[inline] pub fn l2_sq_f16(query: &[f32], vec_f16: &[u16]) -> f32 { debug_assert_eq!(query.len(), vec_f16.len()); @@ -115,6 +119,25 @@ pub fn l2_sq_f16(query: &[f32], vec_f16: &[u16]) -> f32 { sum } +/// Fused `(Σ q_i·x_i, Σ x_i²)` between an f32 query and an f16-encoded +/// vector — the unit-sphere (Cosine/InnerProduct) exact-rerank pass needs +/// both in one decode sweep. +/// +/// Scalar reference; SIMD variants install via +/// `vector::distance::table().f16_dot_normsq`. +#[inline] +pub fn dot_normsq_f16(query: &[f32], vec_f16: &[u16]) -> (f32, f32) { + debug_assert_eq!(query.len(), vec_f16.len()); + let mut dot = 0.0f32; + let mut xsq = 0.0f32; + for (q, &h) in query.iter().zip(vec_f16.iter()) { + let x = f16_to_f32(h); + dot += q * x; + xsq += x * x; + } + (dot, xsq) +} + #[cfg(test)] mod tests { use super::*; diff --git a/src/vector/hnsw/search.rs b/src/vector/hnsw/search.rs index cef11b8a..2a5cbe5b 100644 --- a/src/vector/hnsw/search.rs +++ b/src/vector/hnsw/search.rs @@ -12,7 +12,10 @@ use crate::vector::aligned_buffer::AlignedBuffer; use crate::vector::distance; use crate::vector::turbo_quant::collection::{CollectionMetadata, QuantizationConfig}; use crate::vector::turbo_quant::fwht; -use crate::vector::turbo_quant::sq8::{sq8_l2_from_stats, sq8_params, sq8_query_stats}; +use crate::vector::turbo_quant::sq8::{ + SQ8_INT8_QMAX, sq8_int8_dot_to_f32, sq8_l2_from_stats, sq8_params, sq8_quantize_query_scalar, + sq8_query_stats, +}; use crate::vector::types::{DistanceMetric, SearchResult, VectorId}; /// Bit vector for O(1) visited tracking. 64x more cache-efficient than HashSet @@ -133,6 +136,38 @@ impl SearchScratch { } } +thread_local! { + /// Per-thread [`SearchScratch`] recycler for the wire search path (QP-3). + /// The snapshot capture previously built a fresh scratch per query — + /// heaps, visited bitmap, and the 32–65KB ADC LUT reallocated (and + /// re-grown during the search) every FT.SEARCH. One slot suffices: + /// concurrent same-thread searches (interleaved via cooperative yields) + /// simply miss the cache and allocate, exactly like before. + static SCRATCH_TLS: std::cell::RefCell> = + const { std::cell::RefCell::new(None) }; +} + +/// Take this thread's cached [`SearchScratch`], or build a fresh one. +/// +/// Reuse rule mirrors the search-worker pool's cache: `query_rotated.len()` +/// must EXACTLY equal `padded_dim` — it ACTS as the search's padded +/// dimension (LUT sizing + code_len invariants in `hnsw_search`), so a +/// larger buffer is just as wrong as a smaller one. +pub fn take_thread_scratch(padded_dim: u32) -> SearchScratch { + let cached = SCRATCH_TLS.with(|slot| slot.borrow_mut().take()); + match cached { + Some(s) if s.query_rotated.len() == padded_dim as usize => s, + _ => SearchScratch::new(0, padded_dim), + } +} + +/// Return a scratch to this thread's cache for the next search to reuse. +pub fn recycle_thread_scratch(scratch: SearchScratch) { + SCRATCH_TLS.with(|slot| { + *slot.borrow_mut() = Some(scratch); + }); +} + /// HNSW search with 2-hop dual prefetch and TQ-ADC distance. /// /// # Arguments @@ -285,6 +320,24 @@ pub fn hnsw_search_filtered( } else { (0.0, 0.0) }; + // Task #13: int8 symmetric ADC. `sq8_i8_stats_fn` is `Some` only when + // this CPU/build has a matching SIMD int8 kernel installed (and + // `MOON_SQ8_INT8_ADC` isn't `0`) — quantize the query to int8 ONCE per + // search in that case; the per-candidate closures below fall back to + // `sq8_stats_fn` (f32) whenever it's `None`. Inline buffer keeps this + // heap-free for dim <= 512, same VEC-HNSW-03 guarantee as `sq8_query`. + let sq8_i8_stats_fn = distance::table().sq8_i8_stats; + let mut sq8_qi8: SmallVec<[i8; 512]> = SmallVec::new(); + let mut sq8_q_scale = 0.0f32; + let mut sq8_sum_qi8 = 0i32; + if is_sq8 { + if sq8_i8_stats_fn.is_some() { + sq8_qi8 = smallvec::smallvec![0i8; sq8_query.len()]; + let (scale, sum) = sq8_quantize_query_scalar(&sq8_query, SQ8_INT8_QMAX, &mut sq8_qi8); + sq8_q_scale = scale; + sq8_sum_qi8 = sum; + } + } let q_rot = scratch.query_rotated.as_mut_slice(); // Copy query and zero-pad @@ -406,7 +459,17 @@ pub fn hnsw_search_filtered( // outside this per-candidate closure. let slot = &vectors_tq[offset..offset + bytes_per_code]; let (min, scale) = sq8_params(slot, dim); - let (dot_qc, sum_c, sumsq_c) = sq8_stats_fn(&sq8_query, &slot[..dim]); + let (dot_qc, sum_c, sumsq_c) = if let Some(i8_stats) = sq8_i8_stats_fn { + let (dot_int, sum_c_int, sumsq_c_int) = + i8_stats(&sq8_qi8, &slot[..dim], sq8_sum_qi8); + ( + sq8_int8_dot_to_f32(sq8_q_scale, dot_int), + sum_c_int as f32, + sumsq_c_int as f32, + ) + } else { + sq8_stats_fn(&sq8_query, &slot[..dim]) + }; return sq8_l2_from_stats( dim, min, @@ -556,7 +619,17 @@ pub fn hnsw_search_filtered( let _ = budget; let slot = &vectors_tq[offset..offset + bytes_per_code]; let (min, scale) = sq8_params(slot, dim); - let (dot_qc, sum_c, sumsq_c) = sq8_stats_fn(&sq8_query, &slot[..dim]); + let (dot_qc, sum_c, sumsq_c) = if let Some(i8_stats) = sq8_i8_stats_fn { + let (dot_int, sum_c_int, sumsq_c_int) = + i8_stats(&sq8_qi8, &slot[..dim], sq8_sum_qi8); + ( + sq8_int8_dot_to_f32(sq8_q_scale, dot_int), + sum_c_int as f32, + sumsq_c_int as f32, + ) + } else { + sq8_stats_fn(&sq8_query, &slot[..dim]) + }; return sq8_l2_from_stats( dim, min, diff --git a/src/vector/mvcc/manager.rs b/src/vector/mvcc/manager.rs index 2bcbd5eb..308fc377 100644 --- a/src/vector/mvcc/manager.rs +++ b/src/vector/mvcc/manager.rs @@ -64,6 +64,14 @@ pub struct TransactionManager { /// avoiding a treemap lookup for old transactions that are guaranteed visible. /// Invariant: `pruned_below <= oldest_snapshot`. pruned_below: u64, + /// Cached `Arc` snapshot of `committed` for per-search capture. Refreshed + /// lazily on the first `committed_snapshot()` after a mutation (QP-2): + /// read-heavy workloads capture with one Arc bump instead of cloning the + /// treemap on EVERY search; write-heavy pays at most one clone per search + /// — never worse than the old clone-per-search. + committed_snapshot: parking_lot::Mutex>, + /// Set on every mutation of `committed`; cleared by the refresh. + committed_snapshot_dirty: std::sync::atomic::AtomicBool, } impl TransactionManager { @@ -78,6 +86,8 @@ impl TransactionManager { committed: RoaringTreemap::new(), oldest_snapshot: 0, pruned_below: 0, + committed_snapshot: parking_lot::Mutex::new(std::sync::Arc::new(RoaringTreemap::new())), + committed_snapshot_dirty: std::sync::atomic::AtomicBool::new(false), } } @@ -229,6 +239,8 @@ impl TransactionManager { return false; } self.committed.insert(txn_id); + self.committed_snapshot_dirty + .store(true, std::sync::atomic::Ordering::Release); self.write_intents.retain(|_, owner| *owner != txn_id); #[cfg(feature = "graph")] self.graph_write_intents.retain(|_, owner| *owner != txn_id); @@ -364,6 +376,8 @@ impl TransactionManager { // Remove all entries in [0, floor). let before = self.committed.len(); self.committed.remove_range(..floor); + self.committed_snapshot_dirty + .store(true, std::sync::atomic::Ordering::Release); let after = self.committed.len(); let pruned = before.saturating_sub(after); // Advance the floor. Never move it backwards. @@ -456,6 +470,24 @@ impl TransactionManager { &self.committed } + /// Owned snapshot of the committed treemap for per-search capture (QP-2). + /// + /// Equivalent to `committed_treemap().clone()` — the returned value is + /// exactly the committed set at call time — but amortized: the clone + /// happens only on the first call after a commit/prune; subsequent calls + /// are an `Arc` refcount bump. The uncontended mutex is ~ns (per-shard + /// manager; brief critical section). + pub fn committed_snapshot(&self) -> std::sync::Arc { + let mut snap = self.committed_snapshot.lock(); + if self + .committed_snapshot_dirty + .swap(false, std::sync::atomic::Ordering::AcqRel) + { + *snap = std::sync::Arc::new(self.committed.clone()); + } + std::sync::Arc::clone(&snap) + } + /// Recalculate oldest_snapshot from active non-killed transactions. /// /// Killed snapshots are excluded so the GC watermark can advance past them. @@ -1026,4 +1058,35 @@ mod tests { .expect("age must be Some"); assert!(age.as_secs() >= 42, "age must be >= 42s; got {:?}", age); } + + #[test] + fn test_committed_snapshot_matches_and_amortizes() { + let mut mgr = TransactionManager::new(); + + // Fresh manager: empty snapshot. + let s0 = mgr.committed_snapshot(); + assert!(s0.is_empty()); + + // Commit → snapshot reflects it (equivalent to clone-at-call). + let t1 = mgr.begin(); + assert!(mgr.commit(t1.txn_id)); + let s1 = mgr.committed_snapshot(); + assert!(s1.contains(t1.txn_id)); + assert_eq!(*s1, *mgr.committed_treemap()); + + // No mutation in between → same Arc (the amortization). + let s2 = mgr.committed_snapshot(); + assert!(std::sync::Arc::ptr_eq(&s1, &s2)); + + // Another commit → refreshed Arc with the new id. + let t2 = mgr.begin(); + assert!(mgr.commit(t2.txn_id)); + let s3 = mgr.committed_snapshot(); + assert!(!std::sync::Arc::ptr_eq(&s2, &s3)); + assert!(s3.contains(t1.txn_id) && s3.contains(t2.txn_id)); + assert_eq!(*s3, *mgr.committed_treemap()); + + // Earlier snapshots are unaffected by later mutations (owned view). + assert!(!s1.contains(t2.txn_id)); + } } diff --git a/src/vector/persistence/manifest.rs b/src/vector/persistence/manifest.rs new file mode 100644 index 00000000..cb08edfe --- /dev/null +++ b/src/vector/persistence/manifest.rs @@ -0,0 +1,949 @@ +//! Vector-index durability (B1/B2 write path): manifest + keymap persistence +//! and the background snapshot job that commits them after a compact/merge +//! install. +//! +//! ## On-disk layout +//! +//! Under `/idx-/`: +//! ```text +//! manifest.json -- atomic pointer (tmp + fsync + rename + dir fsync) +//! keymap-.bin -- key_hash -> (global_id, vec_checksum, key) +//! segment-/ -- segment_io v1 format, written via the staged +//! (staging- -> rename) writer +//! ``` +//! +//! ## Crash-safety invariant +//! +//! The manifest may UNDERSTATE what exists on disk (a crash between a +//! segment/keymap write and the manifest rename just leaves an orphan +//! directory/file — harmless, swept at B3 startup) but must NEVER OVERSTATE +//! (recovery must never be pointed at a segment/keymap that isn't durably on +//! disk). See `tmp/VECTOR-DURABILITY-DESIGN.md` for the full design. +//! +//! ## Per-index ordering +//! +//! The background compaction pool (`background_compact.rs`) is a GLOBAL, +//! multi-worker pool with NO ordering guarantee across jobs — two snapshot +//! jobs for the SAME index can be dequeued and completed by different +//! workers in either order. [`SnapshotJob`] carries a per-index monotonic +//! `seq` (allocated on the shard thread at enqueue time) plus a shared +//! `watermark` lock; [`run_snapshot_job`] holds that lock across its entire +//! write+GC critical section, so a stale job (built from an older install) +//! can never overwrite a newer, already-committed manifest — see that +//! function's doc comment for the full argument. + +use std::collections::HashSet; +use std::fs; +use std::io; +use std::path::{Path, PathBuf}; +use std::sync::Arc; +use std::sync::OnceLock; + +use bytes::Bytes; +use parking_lot::Mutex; +use serde::{Deserialize, Serialize}; + +use crate::persistence::fsync::{fsync_directory, fsync_file}; + +/// Current on-disk manifest format version. +pub const MANIFEST_FORMAT_VERSION: u32 = 1; + +/// Magic bytes identifying a `keymap-.bin` file. +const KEYMAP_MAGIC: &[u8; 4] = b"MKM1"; +/// `magic(4) + count(8) + payload_xxh64(8)`. +const KEYMAP_HEADER_LEN: usize = 20; + +/// Atomic pointer describing one vector index's durable state: which +/// collection id / segment ids / keymap epoch / id-allocator floors are +/// currently live. Written with tmp+fsync+rename+dir-fsync (never partially +/// visible); read tolerantly (any corruption/mismatch degrades to "absent", +/// never a panic or hard error — the B3 rescan covers the gap). +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct IndexManifest { + /// Format version — bumped on any incompatible layout change. + pub format_version: u32, + /// Hex-encoded `xxh64(index_name)` — sanity/debugging aid; also the + /// directory's own suffix, so a manifest can be self-checked against + /// the directory it was found in. + pub index_name_hex: String, + /// QJL rotation seed for the default field — MUST match the `collection_id` + /// the index is recreated with, or loaded segments become unsearchable + /// (see design doc's "#1 correctness trap"). + pub collection_id: u64, + /// One past the highest collection id used by ANY field of this index + /// (default + additional named vector fields). Recovery seeds the + /// store-wide `next_collection_id` allocator with `max(this, current)` + /// so a later FT.CREATE never collides with a restored collection id. + pub next_collection_id_floor: u64, + /// Next on-disk segment id to allocate for this index. + pub next_segment_id: u64, + /// Floor for the mutable segment's global-id allocator. + pub next_global_id: u32, + /// Live immutable segment ids referenced by this manifest (in no + /// particular order). Segment directories that exist on disk but are + /// NOT in this list are orphans (superseded by a merge, or a compact + /// whose manifest commit never landed) — safe to delete. + pub segment_ids: Vec, + /// Which `keymap-.bin` is current. + pub keymap_epoch: u64, +} + +/// One entry in a `keymap-.bin` file. +#[derive(Debug, Clone, PartialEq)] +pub struct KeymapEntry { + pub key_hash: u64, + pub global_id: u32, + /// `xxh64(vector-field bytes, seed 0)` at last (re-)insert. + pub vec_checksum: u64, + /// Original Redis key bytes. + pub key: Bytes, +} + +/// Directory an index's manifest/keymap/segments live under, given the +/// store-level `vector_persist_dir` and the index's name. Hex avoids +/// path-unsafe index names. +pub fn index_persist_dir(store_persist_dir: &Path, index_name: &[u8]) -> PathBuf { + store_persist_dir.join(format!("idx-{}", index_name_hex(index_name))) +} + +/// Hex-encoded `xxh64(index_name)`, used both for the `idx-` directory +/// name and the manifest's self-describing `index_name_hex` field. +pub fn index_name_hex(index_name: &[u8]) -> String { + format!("{:016x}", xxhash_rust::xxh64::xxh64(index_name, 0)) +} + +fn manifest_path(dir: &Path) -> PathBuf { + dir.join("manifest.json") +} + +fn manifest_tmp_path(dir: &Path) -> PathBuf { + dir.join("manifest.json.tmp") +} + +fn keymap_path(dir: &Path, epoch: u64) -> PathBuf { + dir.join(format!("keymap-{epoch}.bin")) +} + +fn keymap_tmp_path(dir: &Path, epoch: u64) -> PathBuf { + dir.join(format!("keymap-{epoch}.bin.tmp")) +} + +fn segment_dir_path(dir: &Path, segment_id: u64) -> PathBuf { + dir.join(format!("segment-{segment_id}")) +} + +/// Write `manifest.json` atomically: serialize -> write to `manifest.json.tmp` +/// -> fsync the tmp file -> `rename` over `manifest.json` -> fsync `dir`. +/// +/// The rename is the only visible state transition: readers either see the +/// complete old manifest or the complete new one, never a partial write. +pub fn write_manifest_atomic(dir: &Path, manifest: &IndexManifest) -> io::Result<()> { + fs::create_dir_all(dir)?; + let json = serde_json::to_vec_pretty(manifest) + .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e.to_string()))?; + let tmp_path = manifest_tmp_path(dir); + fs::write(&tmp_path, &json)?; + fsync_file(&tmp_path)?; + fs::rename(&tmp_path, manifest_path(dir))?; + fsync_directory(dir)?; + Ok(()) +} + +/// Read `manifest.json` tolerantly: absent (never written, or a stale +/// pre-B2 index) -> `None` with no log noise (the common case); present but +/// corrupt/unparseable -> `None` with a `tracing::warn!` (a real anomaly the +/// operator should be able to spot, but never a panic or propagated error — +/// the B3 rescan degrades gracefully to a full re-index for that one index). +pub fn read_manifest_tolerant(dir: &Path) -> Option { + let path = manifest_path(dir); + let bytes = match fs::read(&path) { + Ok(b) => b, + Err(e) if e.kind() == io::ErrorKind::NotFound => return None, + Err(e) => { + tracing::debug!( + "manifest.json at {}: {e} — treating as absent", + dir.display() + ); + return None; + } + }; + match serde_json::from_slice::(&bytes) { + Ok(m) => Some(m), + Err(e) => { + tracing::warn!( + "manifest.json at {} is corrupt: {e} — treating as absent \ + (full rescan will cover this index)", + dir.display() + ); + None + } + } +} + +fn read_u16_le(b: &[u8], off: usize) -> Option { + b.get(off..off + 2)?.try_into().ok().map(u16::from_le_bytes) +} +fn read_u32_le(b: &[u8], off: usize) -> Option { + b.get(off..off + 4)?.try_into().ok().map(u32::from_le_bytes) +} +fn read_u64_le(b: &[u8], off: usize) -> Option { + b.get(off..off + 8)?.try_into().ok().map(u64::from_le_bytes) +} + +/// Write `keymap-.bin` atomically (tmp + fsync + rename + dir-fsync), +/// same protocol as [`write_manifest_atomic`]. +/// +/// Format (little-endian): +/// ```text +/// header: magic "MKM1" (4B) | count u64 | payload_xxh64 u64 +/// entry: key_hash u64 | global_id u32 | vec_checksum u64 | key_len u16 | key bytes +/// ``` +/// `payload_xxh64` covers every entry byte (seed 0) — verified on read. +pub fn write_keymap_atomic(dir: &Path, epoch: u64, entries: &[KeymapEntry]) -> io::Result<()> { + fs::create_dir_all(dir)?; + let mut payload = Vec::with_capacity(entries.len() * 24); + for e in entries { + payload.extend_from_slice(&e.key_hash.to_le_bytes()); + payload.extend_from_slice(&e.global_id.to_le_bytes()); + payload.extend_from_slice(&e.vec_checksum.to_le_bytes()); + // Truncate degenerately long keys rather than erroring — this is + // a best-effort persistence cache, not the source of truth (the + // DashTable is); an oversize key is vanishingly unlikely to be + // preceded by a HASH key long enough to overflow u16 (64KiB). + let key_len = e.key.len().min(u16::MAX as usize); + payload.extend_from_slice(&(key_len as u16).to_le_bytes()); + payload.extend_from_slice(&e.key[..key_len]); + } + let checksum = xxhash_rust::xxh64::xxh64(&payload, 0); + + let mut buf = Vec::with_capacity(KEYMAP_HEADER_LEN + payload.len()); + buf.extend_from_slice(KEYMAP_MAGIC); + buf.extend_from_slice(&(entries.len() as u64).to_le_bytes()); + buf.extend_from_slice(&checksum.to_le_bytes()); + buf.extend_from_slice(&payload); + + let tmp_path = keymap_tmp_path(dir, epoch); + fs::write(&tmp_path, &buf)?; + fsync_file(&tmp_path)?; + fs::rename(&tmp_path, keymap_path(dir, epoch))?; + fsync_directory(dir)?; + Ok(()) +} + +/// Read `keymap-.bin` tolerantly. Any structural problem — missing +/// file, bad magic, truncated header/entries, or a payload checksum +/// mismatch — returns `None` (with a `tracing::warn!` for anything other +/// than plain absence) instead of erroring: a corrupt keymap costs a +/// rescan of that index, never a crash. +pub fn read_keymap_tolerant(dir: &Path, epoch: u64) -> Option> { + let path = keymap_path(dir, epoch); + let bytes = match fs::read(&path) { + Ok(b) => b, + Err(e) if e.kind() == io::ErrorKind::NotFound => return None, + Err(e) => { + tracing::debug!( + "keymap-{epoch}.bin at {}: {e} — treating as absent", + dir.display() + ); + return None; + } + }; + if bytes.len() < KEYMAP_HEADER_LEN || &bytes[0..4] != KEYMAP_MAGIC { + tracing::warn!( + "keymap-{epoch}.bin at {} has bad magic or is truncated — treating as absent", + dir.display() + ); + return None; + } + let Some(count) = read_u64_le(&bytes, 4) else { + tracing::warn!("keymap-{epoch}.bin at {}: truncated header", dir.display()); + return None; + }; + let Some(expected_checksum) = read_u64_le(&bytes, 12) else { + tracing::warn!("keymap-{epoch}.bin at {}: truncated header", dir.display()); + return None; + }; + let payload = &bytes[KEYMAP_HEADER_LEN..]; + let actual_checksum = xxhash_rust::xxh64::xxh64(payload, 0); + if actual_checksum != expected_checksum { + tracing::warn!( + "keymap-{epoch}.bin at {} failed integrity check (checksum {actual_checksum:#x} != \ + expected {expected_checksum:#x}) — treating as absent", + dir.display() + ); + return None; + } + + let mut entries = Vec::with_capacity(count as usize); + let mut off = 0usize; + for _ in 0..count { + let Some(key_hash) = read_u64_le(payload, off) else { + tracing::warn!("keymap-{epoch}.bin at {}: entry truncated", dir.display()); + return None; + }; + off += 8; + let Some(global_id) = read_u32_le(payload, off) else { + tracing::warn!("keymap-{epoch}.bin at {}: entry truncated", dir.display()); + return None; + }; + off += 4; + let Some(vec_checksum) = read_u64_le(payload, off) else { + tracing::warn!("keymap-{epoch}.bin at {}: entry truncated", dir.display()); + return None; + }; + off += 8; + let Some(key_len) = read_u16_le(payload, off) else { + tracing::warn!("keymap-{epoch}.bin at {}: entry truncated", dir.display()); + return None; + }; + off += 2; + let Some(key_bytes) = payload.get(off..off + key_len as usize) else { + tracing::warn!( + "keymap-{epoch}.bin at {}: key bytes truncated", + dir.display() + ); + return None; + }; + off += key_len as usize; + entries.push(KeymapEntry { + key_hash, + global_id, + vec_checksum, + key: Bytes::copy_from_slice(key_bytes), + }); + } + Some(entries) +} + +/// Given the entry names present in an index's persist dir and its current +/// manifest, return the names that are NOT referenced by the manifest and +/// are therefore safe to delete: +/// - `staging-*` — always an orphan (an interrupted segment build; a +/// completed one was renamed away to `segment-*` before the manifest that +/// references it was ever written). +/// - `segment-` — an orphan when `id` is not in `manifest.segment_ids` +/// (or the name doesn't parse as `segment-`). +/// - `keymap-.bin` — an orphan when `epoch != manifest.keymap_epoch` +/// (or the name doesn't parse as `keymap-.bin`, e.g. a leftover +/// `keymap-.bin.tmp`). +/// - `manifest.json` (and its own `.tmp`) and anything else unrecognized are +/// never returned — this helper only ever names things it is certain are +/// superseded. +/// +/// Pure function — no I/O. Intended to run at B3 startup (no concurrent +/// writer for the index at that point); the LIVE path uses the online +/// set-difference GC in [`run_snapshot_job`] instead, which only ever +/// removes what the manifest transition itself proves is superseded (a +/// directory scan while a concurrent compaction is mid-write would delete a +/// live, not-yet-referenced `segment-`). +pub fn compute_orphans(entries: &[String], manifest: &IndexManifest) -> Vec { + let live_segments: HashSet = manifest.segment_ids.iter().copied().collect(); + entries + .iter() + .filter(|name| { + if let Some(rest) = name.strip_prefix("staging-") { + // Any staging dir is an interrupted/superseded build, + // regardless of what id it names. + let _ = rest; + return true; + } + if let Some(id_str) = name.strip_prefix("segment-") { + return match id_str.parse::() { + Ok(id) => !live_segments.contains(&id), + Err(_) => true, // malformed name — not a real segment dir + }; + } + if let Some(rest) = name.strip_prefix("keymap-") { + return match rest + .strip_suffix(".bin") + .and_then(|s| s.parse::().ok()) + { + Some(epoch) => epoch != manifest.keymap_epoch, + None => true, // not a well-formed "keymap-.bin" (incl. .tmp leftovers) + }; + } + false + }) + .cloned() + .collect() +} + +/// Disk-backed variant of [`compute_orphans`]: lists `dir`, computes the +/// orphan set against `manifest`, and deletes each one (best-effort — a +/// single removal failure is logged and does not abort the sweep). +/// +/// Startup-only (B3): safe because nothing else is writing to this index's +/// directory yet. NOT safe to call on the live/steady-state path (use the +/// online set-difference GC inside [`run_snapshot_job`] there). +pub fn sweep_orphans_from_disk(dir: &Path, manifest: &IndexManifest) -> io::Result> { + let mut names = Vec::new(); + for entry in fs::read_dir(dir)? { + let entry = entry?; + if let Some(name) = entry.file_name().to_str() { + names.push(name.to_owned()); + } + } + let orphans = compute_orphans(&names, manifest); + let mut removed = Vec::with_capacity(orphans.len()); + for name in orphans { + let path = dir.join(&name); + let result = if path.is_dir() { + fs::remove_dir_all(&path) + } else { + fs::remove_file(&path) + }; + match result { + Ok(()) => removed.push(path), + Err(e) => tracing::warn!("orphan sweep: failed to remove {}: {e}", path.display()), + } + } + Ok(removed) +} + +/// A background job that durably commits one index's keymap + manifest +/// snapshot after a compact/merge install, then GCs whatever the new +/// manifest supersedes. +pub struct SnapshotJob { + /// This index's persist directory (`idx-/`). + pub idx_dir: PathBuf, + /// Monotonic per-index sequence number, allocated on the shard thread at + /// enqueue time. Also used as the keymap epoch. + pub seq: u64, + /// Shared per-index watermark — see module docs and [`run_snapshot_job`]. + pub watermark: Arc>, + /// The manifest to commit if this job is not stale. `keymap_epoch` MUST + /// equal `seq`. + pub manifest: IndexManifest, + pub key_hash_to_key: Arc>, + pub key_hash_to_global_id: Arc>, + pub key_hash_to_vec_checksum: Arc>, +} + +/// Run one snapshot job to completion. +/// +/// ## Ordering guarantee +/// +/// The background compaction pool has multiple workers and no cross-job +/// ordering — a job built from an OLDER install can be dequeued and reach +/// this function AFTER a job built from a NEWER install already committed. +/// A bare "load watermark, compare, write" is racy (two jobs can both read +/// the same watermark and then race the final rename); instead this +/// function holds `job.watermark`'s lock across the ENTIRE critical section +/// (stale check -> keymap write -> manifest write -> GC -> watermark +/// update). Since all jobs for one index share the same `Arc>`, +/// this fully serializes them: whichever job's thread acquires the lock +/// first runs to completion; every other job for that index blocks, then +/// (because `seq` is monotonically allocated on the single shard thread and +/// therefore totally ordered) either proceeds — if it's newer than what +/// just committed — or is dropped as stale — if it's older. No I/O happens +/// before the stale check, so a stale job costs nothing but the lock wait. +/// +/// ## GC is a set-difference, not a directory scan +/// +/// Only segment ids present in the OLD on-disk manifest but absent from the +/// NEW one are removed (merge-obsoleted sources), plus the old keymap epoch +/// if it changed. A concurrent compaction's in-flight `staging-` was +/// never in the old manifest, so this can never touch it — unlike a +/// directory scan, which would. +pub fn run_snapshot_job(job: SnapshotJob) { + let mut watermark = job.watermark.lock(); + if job.seq <= *watermark { + tracing::debug!( + seq = job.seq, + watermark = *watermark, + dir = %job.idx_dir.display(), + "vector snapshot job stale — dropped (a newer snapshot already committed)" + ); + return; + } + + let old_manifest = read_manifest_tolerant(&job.idx_dir); + + let mut entries: Vec = Vec::with_capacity(job.key_hash_to_key.len()); + for (&key_hash, key) in job.key_hash_to_key.iter() { + let global_id = job + .key_hash_to_global_id + .get(&key_hash) + .copied() + .unwrap_or(0); + let vec_checksum = job + .key_hash_to_vec_checksum + .get(&key_hash) + .copied() + .unwrap_or(0); + entries.push(KeymapEntry { + key_hash, + global_id, + vec_checksum, + key: key.clone(), + }); + } + + if let Err(e) = write_keymap_atomic(&job.idx_dir, job.seq, &entries) { + tracing::warn!( + "vector snapshot: failed to write keymap-{} at {}: {e} — leaving prior \ + manifest in place (this compact/merge install is simply not durable yet)", + job.seq, + job.idx_dir.display() + ); + return; + } + if let Err(e) = write_manifest_atomic(&job.idx_dir, &job.manifest) { + tracing::warn!( + "vector snapshot: failed to write manifest.json at {}: {e}", + job.idx_dir.display() + ); + return; + } + + if let Some(old) = old_manifest { + let new_ids: HashSet = job.manifest.segment_ids.iter().copied().collect(); + for old_id in old.segment_ids { + if new_ids.contains(&old_id) { + continue; + } + let path = segment_dir_path(&job.idx_dir, old_id); + if let Err(e) = fs::remove_dir_all(&path) { + if e.kind() != io::ErrorKind::NotFound { + tracing::warn!( + "vector snapshot GC: failed to remove superseded segment-{old_id} \ + at {}: {e}", + path.display() + ); + } + } + } + if old.keymap_epoch != job.manifest.keymap_epoch { + let old_path = keymap_path(&job.idx_dir, old.keymap_epoch); + if let Err(e) = fs::remove_file(&old_path) { + if e.kind() != io::ErrorKind::NotFound { + tracing::warn!( + "vector snapshot GC: failed to remove superseded keymap-{}: {e}", + old.keymap_epoch + ); + } + } + } + } + + *watermark = job.seq; +} + +/// Background snapshot pool: a small, dedicated worker pool that runs +/// [`run_snapshot_job`]. Deliberately separate from +/// `background_compact::BackgroundCompactor` (different job shape — this is +/// small, fsync-bound I/O, not CPU-bound HNSW build work) but the same +/// lazy-init-singleton shape. +pub struct SnapshotPool { + job_tx: flume::Sender, + _workers: Vec>, +} + +impl SnapshotPool { + fn new(num_workers: usize) -> Self { + let (job_tx, job_rx) = flume::unbounded::(); + let workers: Vec<_> = (0..num_workers.max(1)) + .map(|i| { + let rx = job_rx.clone(); + // Startup-time OS thread spawn failure is unrecoverable + // (matches the existing `BackgroundCompactor::new` convention). + #[allow(clippy::unwrap_used)] + let handle = std::thread::Builder::new() + .name(format!("moon-vec-snapshot-{i}")) + .spawn(move || { + while let Ok(job) = rx.recv() { + run_snapshot_job(job); + } + }) + .expect("failed to spawn vector snapshot worker thread"); + handle + }) + .collect(); + Self { + job_tx, + _workers: workers, + } + } + + /// Enqueue a snapshot job. Best-effort: send only fails if every worker + /// has exited (e.g. all panicked), in which case the job is dropped with + /// a warning — per the crash-safety invariant this only leaves the + /// on-disk state a bit more stale (understating), never wrong. + pub fn submit(&self, job: SnapshotJob) { + if self.job_tx.send(job).is_err() { + tracing::warn!( + "vector snapshot pool: all workers gone — snapshot job dropped \ + (index falls behind on durability, no correctness impact)" + ); + } + } +} + +/// Process-wide snapshot pool, lazily initialized on first use. A single +/// worker is sufficient: jobs are small, fsync-bound, and infrequent (one +/// per compact/merge install), so there is no throughput reason to add +/// concurrency here — and a single worker keeps per-index ordering trivial +/// to reason about even before accounting for the watermark guard. +pub fn global_snapshot_pool() -> &'static SnapshotPool { + static GLOBAL: OnceLock = OnceLock::new(); + GLOBAL.get_or_init(|| SnapshotPool::new(1)) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn sample_manifest(dir_name: &str, segment_ids: Vec, keymap_epoch: u64) -> IndexManifest { + IndexManifest { + format_version: MANIFEST_FORMAT_VERSION, + index_name_hex: dir_name.to_owned(), + collection_id: 7, + next_collection_id_floor: 8, + next_segment_id: segment_ids.iter().copied().max().map_or(0, |m| m + 1), + next_global_id: 12345, + segment_ids, + keymap_epoch, + } + } + + #[test] + fn test_manifest_roundtrip() { + let tmp = tempfile::tempdir().unwrap(); + let manifest = sample_manifest("abc123", vec![1, 3], 5); + + write_manifest_atomic(tmp.path(), &manifest).unwrap(); + let loaded = read_manifest_tolerant(tmp.path()).unwrap(); + + assert_eq!(loaded, manifest); + // No leftover tmp file after a successful commit. + assert!(!manifest_tmp_path(tmp.path()).exists()); + } + + #[test] + fn test_manifest_absent_returns_none() { + let tmp = tempfile::tempdir().unwrap(); + assert!(read_manifest_tolerant(tmp.path()).is_none()); + } + + #[test] + fn test_manifest_corrupt_json_tolerated() { + let tmp = tempfile::tempdir().unwrap(); + fs::write(manifest_path(tmp.path()), b"{ not valid json").unwrap(); + assert!(read_manifest_tolerant(tmp.path()).is_none()); + } + + #[test] + fn test_manifest_only_tmp_present_is_treated_as_absent() { + // Simulates a crash between `fs::write(&tmp_path, ...)` and the + // `rename` — the tmp file alone must never be mistaken for a + // committed manifest. + let tmp = tempfile::tempdir().unwrap(); + let manifest = sample_manifest("abc123", vec![1], 1); + let json = serde_json::to_vec_pretty(&manifest).unwrap(); + fs::write(manifest_tmp_path(tmp.path()), json).unwrap(); + + assert!(read_manifest_tolerant(tmp.path()).is_none()); + } + + #[test] + fn test_manifest_write_is_atomic_second_write_replaces_first() { + let tmp = tempfile::tempdir().unwrap(); + let m1 = sample_manifest("abc123", vec![1], 1); + let m2 = sample_manifest("abc123", vec![1, 2], 2); + + write_manifest_atomic(tmp.path(), &m1).unwrap(); + write_manifest_atomic(tmp.path(), &m2).unwrap(); + + let loaded = read_manifest_tolerant(tmp.path()).unwrap(); + assert_eq!(loaded, m2); + } + + fn sample_entries() -> Vec { + vec![ + KeymapEntry { + key_hash: 111, + global_id: 1, + vec_checksum: 0xDEAD_BEEF, + key: Bytes::from_static(b"doc:1"), + }, + KeymapEntry { + key_hash: 222, + global_id: 2, + vec_checksum: 0xCAFE_F00D, + key: Bytes::from_static(b"doc:two-with-a-longer-key-name"), + }, + KeymapEntry { + key_hash: 333, + global_id: 3, + vec_checksum: 0, + key: Bytes::from_static(b""), + }, + ] + } + + #[test] + fn test_keymap_roundtrip() { + let tmp = tempfile::tempdir().unwrap(); + let entries = sample_entries(); + + write_keymap_atomic(tmp.path(), 9, &entries).unwrap(); + let loaded = read_keymap_tolerant(tmp.path(), 9).unwrap(); + + assert_eq!(loaded.len(), entries.len()); + for e in &entries { + assert!(loaded.contains(e)); + } + assert!(!keymap_tmp_path(tmp.path(), 9).exists()); + } + + #[test] + fn test_keymap_empty_roundtrip() { + let tmp = tempfile::tempdir().unwrap(); + write_keymap_atomic(tmp.path(), 1, &[]).unwrap(); + let loaded = read_keymap_tolerant(tmp.path(), 1).unwrap(); + assert!(loaded.is_empty()); + } + + #[test] + fn test_keymap_absent_returns_none() { + let tmp = tempfile::tempdir().unwrap(); + assert!(read_keymap_tolerant(tmp.path(), 42).is_none()); + } + + #[test] + fn test_keymap_bad_magic_tolerated() { + let tmp = tempfile::tempdir().unwrap(); + write_keymap_atomic(tmp.path(), 1, &sample_entries()).unwrap(); + let path = keymap_path(tmp.path(), 1); + let mut bytes = fs::read(&path).unwrap(); + bytes[0] = b'X'; // corrupt the magic + fs::write(&path, &bytes).unwrap(); + + assert!(read_keymap_tolerant(tmp.path(), 1).is_none()); + } + + #[test] + fn test_keymap_corrupt_checksum_tolerated() { + let tmp = tempfile::tempdir().unwrap(); + write_keymap_atomic(tmp.path(), 1, &sample_entries()).unwrap(); + let path = keymap_path(tmp.path(), 1); + let mut bytes = fs::read(&path).unwrap(); + // Flip a byte inside the payload (past the 20-byte header) without + // touching the stored checksum — must be detected as corrupt. + let last = bytes.len() - 1; + bytes[last] ^= 0xFF; + fs::write(&path, &bytes).unwrap(); + + assert!(read_keymap_tolerant(tmp.path(), 1).is_none()); + } + + #[test] + fn test_keymap_truncated_tolerated() { + let tmp = tempfile::tempdir().unwrap(); + write_keymap_atomic(tmp.path(), 1, &sample_entries()).unwrap(); + let path = keymap_path(tmp.path(), 1); + let bytes = fs::read(&path).unwrap(); + fs::write(&path, &bytes[..bytes.len() - 5]).unwrap(); + + assert!(read_keymap_tolerant(tmp.path(), 1).is_none()); + } + + #[test] + fn test_keymap_only_tmp_present_is_treated_as_absent() { + let tmp = tempfile::tempdir().unwrap(); + // Write directly to the tmp path, simulating a crash before rename. + fs::write(keymap_tmp_path(tmp.path(), 3), b"MKM1garbage").unwrap(); + assert!(read_keymap_tolerant(tmp.path(), 3).is_none()); + } + + #[test] + fn test_orphan_sweep_correctness() { + let manifest = sample_manifest("abc", vec![1, 3], 12); + let entries: Vec = vec![ + "segment-1".to_owned(), // live — kept + "segment-2".to_owned(), // stale — orphan + "segment-3".to_owned(), // live — kept + "keymap-10.bin".to_owned(), // stale epoch — orphan + "keymap-12.bin".to_owned(), // live — kept + "staging-4".to_owned(), // interrupted build — orphan + "manifest.json".to_owned(), // never swept + "manifest.json.tmp".to_owned(), // never swept by this helper (not segment/keymap-named) + ]; + + let mut orphans = compute_orphans(&entries, &manifest); + orphans.sort(); + assert_eq!( + orphans, + vec![ + "keymap-10.bin".to_owned(), + "segment-2".to_owned(), + "staging-4".to_owned(), + ] + ); + } + + #[test] + fn test_orphan_sweep_treats_stray_tmp_keymap_as_orphan() { + // A "keymap-3.bin.tmp" left over from an interrupted write starts + // with "keymap-" but does not parse as "keymap-.bin" — it + // must be swept (it's definitely garbage), not silently kept. + let manifest = sample_manifest("abc", vec![], 3); + let entries = vec!["keymap-3.bin.tmp".to_owned(), "keymap-3.bin".to_owned()]; + let orphans = compute_orphans(&entries, &manifest); + assert_eq!(orphans, vec!["keymap-3.bin.tmp".to_owned()]); + } + + #[test] + fn test_sweep_orphans_from_disk_deletes_only_orphans() { + let tmp = tempfile::tempdir().unwrap(); + let dir = tmp.path(); + fs::create_dir_all(dir.join("segment-1")).unwrap(); + fs::create_dir_all(dir.join("segment-2")).unwrap(); + fs::create_dir_all(dir.join("staging-9")).unwrap(); + fs::write(dir.join("keymap-1.bin"), b"live").unwrap(); + fs::write(dir.join("keymap-0.bin"), b"stale").unwrap(); + fs::write(dir.join("manifest.json"), b"{}").unwrap(); + + let manifest = sample_manifest("abc", vec![1], 1); + let removed = sweep_orphans_from_disk(dir, &manifest).unwrap(); + + assert_eq!(removed.len(), 3); + assert!(dir.join("segment-1").exists()); + assert!(!dir.join("segment-2").exists()); + assert!(!dir.join("staging-9").exists()); + assert!(dir.join("keymap-1.bin").exists()); + assert!(!dir.join("keymap-0.bin").exists()); + assert!(dir.join("manifest.json").exists()); + } + + fn checksum_map(entries: &[KeymapEntry]) -> Arc> { + Arc::new( + entries + .iter() + .map(|e| (e.key_hash, e.vec_checksum)) + .collect(), + ) + } + fn global_id_map(entries: &[KeymapEntry]) -> Arc> { + Arc::new(entries.iter().map(|e| (e.key_hash, e.global_id)).collect()) + } + fn key_map(entries: &[KeymapEntry]) -> Arc> { + Arc::new( + entries + .iter() + .map(|e| (e.key_hash, e.key.clone())) + .collect(), + ) + } + + #[test] + fn test_run_snapshot_job_commits_and_advances_watermark() { + let tmp = tempfile::tempdir().unwrap(); + let entries = sample_entries(); + let manifest = sample_manifest("abc", vec![1], 1); + let watermark = Arc::new(Mutex::new(0u64)); + + run_snapshot_job(SnapshotJob { + idx_dir: tmp.path().to_path_buf(), + seq: 1, + watermark: watermark.clone(), + manifest: manifest.clone(), + key_hash_to_key: key_map(&entries), + key_hash_to_global_id: global_id_map(&entries), + key_hash_to_vec_checksum: checksum_map(&entries), + }); + + assert_eq!(*watermark.lock(), 1); + let loaded_manifest = read_manifest_tolerant(tmp.path()).unwrap(); + assert_eq!(loaded_manifest, manifest); + let loaded_keymap = read_keymap_tolerant(tmp.path(), 1).unwrap(); + assert_eq!(loaded_keymap.len(), entries.len()); + } + + #[test] + fn test_run_snapshot_job_drops_stale_job_without_io() { + let tmp = tempfile::tempdir().unwrap(); + let entries = sample_entries(); + let watermark = Arc::new(Mutex::new(5u64)); + let manifest = sample_manifest("abc", vec![1], 3); + + run_snapshot_job(SnapshotJob { + idx_dir: tmp.path().to_path_buf(), + seq: 3, // <= watermark(5) -> stale + watermark: watermark.clone(), + manifest, + key_hash_to_key: key_map(&entries), + key_hash_to_global_id: global_id_map(&entries), + key_hash_to_vec_checksum: checksum_map(&entries), + }); + + assert_eq!( + *watermark.lock(), + 5, + "stale job must not move the watermark" + ); + assert!( + read_manifest_tolerant(tmp.path()).is_none(), + "stale job must not write anything" + ); + } + + #[test] + fn test_run_snapshot_job_gcs_superseded_segments_and_keymap() { + let tmp = tempfile::tempdir().unwrap(); + let dir = tmp.path(); + let entries = sample_entries(); + let watermark = Arc::new(Mutex::new(0u64)); + + // Seed an "old" committed state: segment-1, segment-2 live, keymap epoch 1. + fs::create_dir_all(dir.join("segment-1")).unwrap(); + fs::create_dir_all(dir.join("segment-2")).unwrap(); + let old_manifest = sample_manifest("abc", vec![1, 2], 1); + write_manifest_atomic(dir, &old_manifest).unwrap(); + write_keymap_atomic(dir, 1, &entries).unwrap(); + *watermark.lock() = 1; + + // A merge superseded segment-1 and segment-2 with a new segment-3. + fs::create_dir_all(dir.join("segment-3")).unwrap(); + let new_manifest = sample_manifest("abc", vec![3], 2); + + run_snapshot_job(SnapshotJob { + idx_dir: dir.to_path_buf(), + seq: 2, + watermark: watermark.clone(), + manifest: new_manifest.clone(), + key_hash_to_key: key_map(&entries), + key_hash_to_global_id: global_id_map(&entries), + key_hash_to_vec_checksum: checksum_map(&entries), + }); + + assert_eq!(*watermark.lock(), 2); + assert!( + !dir.join("segment-1").exists(), + "superseded segment must be GC'd" + ); + assert!( + !dir.join("segment-2").exists(), + "superseded segment must be GC'd" + ); + assert!(dir.join("segment-3").exists(), "live segment must survive"); + assert!( + !keymap_path(dir, 1).exists(), + "superseded keymap epoch must be GC'd" + ); + assert!( + keymap_path(dir, 2).exists(), + "current keymap epoch must survive" + ); + assert_eq!(read_manifest_tolerant(dir).unwrap(), new_manifest); + } + + #[test] + fn test_index_persist_dir_and_hex_are_deterministic() { + let a = index_persist_dir(Path::new("/tmp/shard-0-vectors"), b"myidx"); + let b = index_persist_dir(Path::new("/tmp/shard-0-vectors"), b"myidx"); + assert_eq!(a, b); + assert!(a.to_string_lossy().contains("idx-")); + assert_eq!(index_name_hex(b"myidx").len(), 16); + } +} diff --git a/src/vector/persistence/mod.rs b/src/vector/persistence/mod.rs index fa2ecc7b..3999c107 100644 --- a/src/vector/persistence/mod.rs +++ b/src/vector/persistence/mod.rs @@ -1,5 +1,11 @@ +pub mod manifest; pub mod mmap_budget; -pub mod recovery; +/// B3: startup recovery — load segments/keymap from the B1/B2 durability +/// layout into the live `ShardSlice.vector_store`, dedup-rescan the rest. +/// Supersedes the old `recovery.rs` WAL-replay module (deleted: it only ever +/// populated the doomed `Shard.vector_store`, see event_loop.rs's +/// `_discarded_vector_store`). +pub mod recover_v2; pub mod sealed_mmap; pub mod segment_io; pub mod wal_record; diff --git a/src/vector/persistence/recover_v2.rs b/src/vector/persistence/recover_v2.rs new file mode 100644 index 00000000..a855804a --- /dev/null +++ b/src/vector/persistence/recover_v2.rs @@ -0,0 +1,931 @@ +//! Vector-index durability (B3): startup recovery into the LIVE +//! `ShardSlice.vector_store`. +//! +//! Supersedes the old WAL-replay `recovery.rs` module (deleted — it only +//! ever populated the `Shard`-owned `vector_store`, which is discarded +//! wholesale at `event_loop.rs`'s `_discarded_vector_store`). Recovery +//! authority is now: sidecar (`vector-indexes.meta`, index *definitions*) + +//! the B1/B2 manifest/segment/keymap durability layout (index *contents*), +//! reconciled against the live keyspace by a dedup rescan. +//! +//! See `tmp/VECTOR-DURABILITY-DESIGN.md` for the full design. This module +//! implements the "Recovery" section end to end; `event_loop.rs` only wires +//! [`RecoveryState`]'s three phase methods into its existing three call +//! sites (index-restore loop, rescan loop, post-scan finalize) because +//! `with_shard`/`with_shard_db` re-entrancy is forbidden — the orchestration +//! (which of those two functions to call, and when) has to stay there. +//! +//! ## Crash-safety degradation ladder (per segment) +//! +//! 1. Segment loads AND its `collection_id` matches the manifest's — used +//! as-is. +//! 2. Segment loads but `collection_id` MISMATCHES the manifest (stale/ +//! tampered file) — the segment object is still in hand, so its exact +//! `key_hash`es are read from its own MVCC headers and dropped from the +//! loaded keymap. The rescan re-indexes exactly those keys; nothing else +//! is touched. +//! 3. Segment fails to load at all (corrupt graph, bad checksum, I/O error) +//! but `mvcc_headers.bin` alone still parses — same precise drop as (2), +//! via [`segment_io::read_mvcc_headers_only`]. +//! 4. Segment fails to load AND its headers are ALSO unreadable — the lost +//! key_hashes can't be attributed. The only safe response is to abandon +//! ALL recovered state for this index (no loaded segments, no loaded +//! keymap) and let the full rescan rebuild everything: keeping the +//! OTHER, successfully-loaded segments while blanking the keymap would +//! let the rescan re-insert their keys fresh into the mutable segment +//! while the old copies still live in those segments — duplicate search +//! results, which the crash-safety contract (never wrong, only ever +//! slower) forbids. + +use std::collections::{HashMap, HashSet}; +use std::path::Path; +use std::sync::Arc; + +use bytes::Bytes; +use tracing::{info, warn}; + +use crate::protocol::Frame; +use crate::text::store::TextStore; +use crate::vector::persistence::manifest::{self, IndexManifest}; +use crate::vector::persistence::segment_io; +use crate::vector::segment::SegmentList; +use crate::vector::segment::immutable::ImmutableSegment; +use crate::vector::segment::mutable::MutableSegment; +use crate::vector::store::{IndexMeta, VectorStore}; + +/// Per-index counters for the B3 startup acceptance-signal log line. +#[derive(Default, Clone, Copy)] +pub struct IndexRecoveryCounters { + pub loaded_segments: usize, + pub verified_unchanged: usize, + pub re_indexed: usize, + pub tombstoned: usize, +} + +/// Threaded through the whole B3 startup sequence by `event_loop.rs`. +/// +/// Lifecycle: construct once, call [`Self::create_index`] once per sidecar +/// index definition, call [`Self::reconcile_key`] once per matching hash key +/// found during the keyspace scan, then call [`Self::finish`] exactly once. +/// All three phase methods assume they run from inside an already-open +/// `with_shard` closure — none of them call `with_shard`/`with_shard_db` +/// themselves (forbidden re-entrancy, see `crate::shard::slice`). +pub struct RecoveryState { + /// Names of indexes that have durable state loaded from a manifest — + /// eligible for the dedup rescan. Everything else (no manifest found, + /// or manifest present but every segment + its headers were + /// unreadable) is a plain fresh index and always takes the full path. + recovered_names: HashSet, + /// Snapshot of each recovered index's key_hash set exactly as loaded + /// from the keymap, taken BEFORE the rescan mutates anything — the + /// deletion probe's baseline. + original_key_hashes: HashMap>, + /// Accumulated during the rescan: key_hashes observed (still present in + /// the keyspace, whether unchanged or changed) per recovered index. + observed_key_hashes: HashMap>, + /// `index_name_hex` of every index definition in the sidecar (recovered + /// or not) — used by the cross-index orphan sweep to tell "a dropped + /// index whose background dir-delete never finished" from "a live + /// index whose manifest just hasn't been read yet". + known_index_hexes: HashSet, + counters: HashMap, +} + +impl RecoveryState { + pub fn new() -> Self { + Self { + recovered_names: HashSet::new(), + original_key_hashes: HashMap::new(), + observed_key_hashes: HashMap::new(), + known_index_hexes: HashSet::new(), + counters: HashMap::new(), + } + } + + /// Phase 1: create one index from its sidecar definition. If a + /// `manifest.json` exists under `/idx-/`, pins + /// the index's collection_id to the manifest's value and loads its + /// segments + keymap (see module docs for the degradation ladder). + /// Otherwise behaves exactly like the pre-B3 `create_index` restore. + pub fn create_index( + &mut self, + vector_store: &mut VectorStore, + idx_persist_root: &Path, + meta: &IndexMeta, + ) { + let name = meta.name.clone(); + self.known_index_hexes + .insert(manifest::index_name_hex(name.as_ref())); + + let idx_dir = manifest::index_persist_dir(idx_persist_root, name.as_ref()); + let loaded_manifest = manifest::read_manifest_tolerant(&idx_dir); + + let recovered_counters = match &loaded_manifest { + Some(m) => { + if let Err(e) = vector_store.create_index_with_collection_id(meta.clone(), m) { + warn!( + "vector index {}: failed to recreate with recovered collection_id: {e} \ + — skipping this index entirely", + String::from_utf8_lossy(&name) + ); + return; + } + load_segments_and_keymap(vector_store, name.as_ref(), &idx_dir, m) + } + None => { + if let Err(e) = vector_store.create_index(meta.clone()) { + warn!( + "vector index {}: failed to restore from sidecar: {e}", + String::from_utf8_lossy(&name) + ); + } + None + } + }; + + if let Some(counters) = recovered_counters { + if let Some(idx) = vector_store.get_index(&name) { + self.original_key_hashes + .insert(name.clone(), idx.key_hash_to_key.keys().copied().collect()); + } + self.recovered_names.insert(name.clone()); + self.counters.insert(name.clone(), counters); + } + } + + /// Phase 2: reconcile one hash key that matches at least one vector or + /// text index prefix (caller already filtered this — see + /// `event_loop.rs`'s `collect_matching`). + /// + /// For every VECTOR index this key matches, decide whether its default + /// vector field is provably unchanged since the last snapshot + /// (checksum match against a recovered index's keymap). Only when + /// *every* matching index agrees "unchanged" does the default vector + /// field get stripped from `args` before delegating to + /// `auto_index_hset_public` (routing it into the existing + /// metadata-only rebuild path — no HNSW/TQ re-encode). A single + /// disagreement (unknown key, missing/zero checksum — the documented + /// multi-field gap, mismatched checksum, or an index with no durable + /// state at all) falls back to the full path for every index matching + /// this key: conservative by design, per the crash-safety contract + /// (never wrong, only occasionally slower). + /// + /// Text-only matches are unaffected either way — `auto_index_hset` + /// always re-derives text/TAG/NUMERIC from `args[1..]`, and stripping + /// only ever removes a vector field's (name, value) pair. + pub fn reconcile_key( + &mut self, + vector_store: &mut VectorStore, + text_store: &mut TextStore, + key: &[u8], + args: &[Frame], + ) { + let matching_vector = vector_store.find_matching_index_names(key); + let key_hash = xxhash_rust::xxh64::xxh64(key, 0); + + let mut any_recovered_checked = false; + let mut all_unchanged = true; + let mut dirty_recovered: Vec = Vec::new(); + + for idx_name in &matching_vector { + if !self.recovered_names.contains(idx_name) { + // Fresh (no durable state) index matching this key — can't + // vouch for "unchanged" at all; forces the full path below. + all_unchanged = false; + continue; + } + self.observed_key_hashes + .entry(idx_name.clone()) + .or_default() + .insert(key_hash); + any_recovered_checked = true; + + let Some(idx) = vector_store.get_index(idx_name) else { + continue; + }; + let known = idx.key_hash_to_global_id.contains_key(&key_hash); + let stored_checksum = idx.key_hash_to_vec_checksum.get(&key_hash).copied(); + let default_field = idx.meta.vector_fields[0].field_name.clone(); + let dim = idx.meta.vector_fields[0].dimension as usize; + let live_checksum = + crate::shard::spsc_handler::find_vector_blob(args, &default_field, dim) + .map(|blob| xxhash_rust::xxh64::xxh64(blob, 0)); + + // `stored_checksum == Some(0)` is the documented multi-field gap + // (additional fields never write a checksum) — treated as + // "changed" (safe re-encode), same as unknown/missing. + let unchanged = matches!( + (known, stored_checksum, live_checksum), + (true, Some(stored), Some(live)) if stored != 0 && stored == live + ); + if !unchanged { + all_unchanged = false; + if known { + dirty_recovered.push(idx_name.clone()); + } + } + } + + // Known-and-different indexes: tombstone the stale copy BEFORE the + // (necessarily full, since all_unchanged is now false) re-encode + // below inserts the new one. + for idx_name in &dirty_recovered { + vector_store.mark_deleted_for_key_in_index(idx_name, key); + } + + if any_recovered_checked && all_unchanged { + let stripped = strip_default_vector_fields(&matching_vector, vector_store, args); + let inserted = crate::shard::spsc_handler::auto_index_hset_public( + vector_store, + text_store, + key, + &stripped, + ); + debug_assert!( + inserted.is_empty(), + "B3 dedup rescan: stripped args must never re-encode a vector" + ); + for idx_name in &matching_vector { + if let Some(c) = self.counters.get_mut(idx_name) { + c.verified_unchanged += 1; + } + } + } else { + let _ = crate::shard::spsc_handler::auto_index_hset_public( + vector_store, + text_store, + key, + args, + ); + for idx_name in &matching_vector { + if self.recovered_names.contains(idx_name) { + if let Some(c) = self.counters.get_mut(idx_name) { + c.re_indexed += 1; + } + } + } + } + } + + /// Phase 3: deletion probe + orphan sweep + acceptance-signal log + /// lines. Call exactly once, after the full keyspace scan. + pub fn finish(self, vector_store: &mut VectorStore, idx_persist_root: &Path) { + // Deletion probe: any key_hash the manifest's keymap loaded that was + // never observed during the rescan no longer exists anywhere in the + // keyspace (a plain DEL/UNLINK, or an HDEL of the vector field that + // removed the whole matching hash) — tombstone it. Read the key + // bytes from the index's OWN (still-live) key map first: entries + // for keys the rescan already touched (verified-unchanged or + // re-indexed) are irrelevant here (they WERE observed), so a + // shrunk-but-still-present entry is fine. + let mut counters = self.counters; + for (name, original) in &self.original_key_hashes { + let observed = self.observed_key_hashes.get(name); + let mut to_delete: Vec> = Vec::new(); + if let Some(idx) = vector_store.get_index(name) { + for kh in original { + if observed.is_some_and(|o| o.contains(kh)) { + continue; + } + if let Some(key_bytes) = idx.key_hash_to_key.get(kh) { + to_delete.push(key_bytes.to_vec()); + } + } + } + for key_bytes in &to_delete { + vector_store.mark_deleted_for_key_in_index(name, key_bytes); + } + if !to_delete.is_empty() { + if let Some(c) = counters.get_mut(name) { + c.tombstoned += to_delete.len(); + } + } + } + + // Per-index orphan sweep (segment-*/staging-*/keymap-* not + // referenced by the manifest that was actually loaded). + for name in &self.recovered_names { + let idx_dir = manifest::index_persist_dir(idx_persist_root, name); + if let Some(m) = manifest::read_manifest_tolerant(&idx_dir) { + if let Err(e) = manifest::sweep_orphans_from_disk(&idx_dir, &m) { + warn!( + "vector index {}: orphan sweep failed: {e}", + String::from_utf8_lossy(name) + ); + } + } + } + + // Cross-index sweep: `idx-` dirs on disk with no matching + // sidecar index at all — e.g. a dropped index whose best-effort + // background directory delete never completed before a crash. + sweep_unknown_index_dirs(idx_persist_root, &self.known_index_hexes); + + for (name, c) in &counters { + info!( + "vector index {}: B3 recovery — loaded {} segment(s), {} key(s) verified \ + unchanged, {} re-indexed, {} tombstoned", + String::from_utf8_lossy(name), + c.loaded_segments, + c.verified_unchanged, + c.re_indexed, + c.tombstoned, + ); + } + } +} + +impl Default for RecoveryState { + fn default() -> Self { + Self::new() + } +} + +/// Load `manifest`'s segments + keymap into the (already pinned-cid-created) +/// index named `name`. Returns `None` only when a segment failure's +/// key_hashes could not be attributed at all (see module docs, degradation +/// level 4) — the index still exists (fresh/empty, pinned cid) but carries +/// no recovered state, so the caller treats it as NOT recovered (full +/// rescan for all of its matching keys). +fn load_segments_and_keymap( + vector_store: &mut VectorStore, + name: &[u8], + idx_dir: &Path, + manifest: &IndexManifest, +) -> Option { + let mut immutable: Vec> = Vec::with_capacity(manifest.segment_ids.len()); + let mut dropped_key_hashes: HashSet = HashSet::new(); + + for &segment_id in &manifest.segment_ids { + match segment_io::read_immutable_segment(idx_dir, segment_id) { + Ok((seg, collection)) if collection.collection_id == manifest.collection_id => { + immutable.push(Arc::new(seg.with_disk_segment_id(Some(segment_id)))); + } + Ok((seg, collection)) => { + // Degradation level 2: loaded fine, but its collection_id + // (hence QJL seed) doesn't match this recovery epoch's + // manifest — never install it (would search wrong/garbage + // distances). The segment object is in hand, so precise + // attribution needs no fallback. + warn!( + "vector index {}: segment-{segment_id} collection_id {} != manifest \ + collection_id {} — treating as corrupt, its keys will be re-indexed", + String::from_utf8_lossy(name), + collection.collection_id, + manifest.collection_id + ); + for h in seg.mvcc_headers() { + dropped_key_hashes.insert(h.key_hash); + } + } + Err(e) => { + warn!( + "vector index {}: segment-{segment_id} failed to load: {e} — \ + attempting header-only key attribution", + String::from_utf8_lossy(name) + ); + match segment_io::read_mvcc_headers_only(idx_dir, segment_id) { + Some(headers) => { + for h in headers { + dropped_key_hashes.insert(h.key_hash); + } + } + None => { + // Degradation level 4: can't attribute this + // segment's keys at all. Keeping the OTHER, + // successfully-loaded segments while blanking the + // keymap would let the rescan re-insert their keys + // fresh alongside the old copies still living in + // those segments — duplicate search results. + // Abandon ALL recovered state for this index. + warn!( + "vector index {}: segment-{segment_id} headers ALSO unreadable — \ + cannot attribute its keys; abandoning all recovered segments/keymap \ + for this index (falling back to a full rescan) to avoid \ + duplicate/stale search results", + String::from_utf8_lossy(name) + ); + return None; + } + } + } + } + } + + let entries = + manifest::read_keymap_tolerant(idx_dir, manifest.keymap_epoch).unwrap_or_default(); + let mut key_hash_to_key = std::collections::HashMap::with_capacity(entries.len()); + let mut key_hash_to_global_id = std::collections::HashMap::with_capacity(entries.len()); + let mut key_hash_to_vec_checksum = std::collections::HashMap::with_capacity(entries.len()); + for e in entries { + if dropped_key_hashes.contains(&e.key_hash) { + continue; + } + key_hash_to_key.insert(e.key_hash, e.key); + key_hash_to_global_id.insert(e.key_hash, e.global_id); + key_hash_to_vec_checksum.insert(e.key_hash, e.vec_checksum); + } + + let loaded_segments = immutable.len(); + let idx = vector_store.get_index_mut(name)?; + + let mutable = Arc::new(MutableSegment::new( + idx.meta.dimension, + idx.collection.clone(), + )); + mutable.set_global_id_base(manifest.next_global_id); + idx.segments.swap(SegmentList { + mutable, + immutable, + ivf: Vec::new(), + warm: Vec::new(), + cold: Vec::new(), + }); + idx.key_hash_to_key = Arc::new(key_hash_to_key); + idx.key_hash_to_global_id = Arc::new(key_hash_to_global_id); + idx.key_hash_to_vec_checksum = Arc::new(key_hash_to_vec_checksum); + + Some(IndexRecoveryCounters { + loaded_segments, + ..Default::default() + }) +} + +/// Build a copy of `args` with the DEFAULT vector field's (name, value) pair +/// removed for every index in `matching_vector` (by the time this is +/// called, every one of them has already been confirmed both recovered AND +/// checksum-unchanged — see `reconcile_key`). Additional (non-default) +/// vector fields are never stripped: B1/B2 never persists their segments or +/// tracks their checksums, so they always re-encode on every restart — the +/// documented, safe multi-field gap. +fn strip_default_vector_fields( + matching_vector: &[Bytes], + vector_store: &VectorStore, + args: &[Frame], +) -> Vec { + let mut strip_fields: Vec = Vec::with_capacity(matching_vector.len()); + for idx_name in matching_vector { + if let Some(idx) = vector_store.get_index(idx_name) { + strip_fields.push(idx.meta.vector_fields[0].field_name.clone()); + } + } + if strip_fields.is_empty() || args.is_empty() { + return args.to_vec(); + } + let mut out = Vec::with_capacity(args.len()); + out.push(args[0].clone()); // key + let mut i = 1; + while i + 1 < args.len() { + let is_stripped = matches!(&args[i], Frame::BulkString(f) if strip_fields.iter().any(|sf| sf.eq_ignore_ascii_case(f))); + if !is_stripped { + out.push(args[i].clone()); + out.push(args[i + 1].clone()); + } + i += 2; + } + out +} + +/// Delete any `idx-` directory under `vector_persist_dir` whose hex +/// doesn't match a currently-defined sidecar index — e.g. a dropped index +/// whose best-effort background directory delete (`spawn_delete_index_persist_dir`) +/// never completed before a crash. Best-effort: failures are logged, never +/// propagated (matches the rest of the B3 recovery contract). +fn sweep_unknown_index_dirs(vector_persist_dir: &Path, known_hexes: &HashSet) { + let Ok(entries) = std::fs::read_dir(vector_persist_dir) else { + return; + }; + for entry in entries.flatten() { + let Some(name) = entry.file_name().to_str().map(|s| s.to_owned()) else { + continue; + }; + let Some(hex) = name.strip_prefix("idx-") else { + continue; + }; + if known_hexes.contains(hex) { + continue; + } + let path = entry.path(); + warn!( + "sweeping unknown vector index directory {} (no matching sidecar index)", + path.display() + ); + if let Err(e) = std::fs::remove_dir_all(&path) { + warn!("failed to remove unknown index dir {}: {e}", path.display()); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::vector::store::VectorFieldMeta; + use crate::vector::turbo_quant::collection::QuantizationConfig; + use crate::vector::types::DistanceMetric; + + fn make_meta(name: &str, dim: u32, prefix: &str) -> IndexMeta { + IndexMeta { + name: Bytes::copy_from_slice(name.as_bytes()), + dimension: dim, + padded_dimension: dim, + metric: DistanceMetric::L2, + hnsw_m: 16, + hnsw_ef_construction: 100, + hnsw_ef_runtime: 0, + compact_threshold: 0, + source_field: Bytes::from_static(b"vec"), + key_prefixes: vec![Bytes::copy_from_slice(prefix.as_bytes())], + quantization: QuantizationConfig::Sq8, + build_mode: crate::vector::turbo_quant::collection::BuildMode::Light, + vector_fields: vec![VectorFieldMeta { + field_name: Bytes::from_static(b"vec"), + dimension: dim, + padded_dimension: dim, + metric: DistanceMetric::L2, + quantization: QuantizationConfig::Sq8, + build_mode: crate::vector::turbo_quant::collection::BuildMode::Light, + }], + schema_fields: Vec::new(), + merge_mode: crate::vector::store::MergeMode::default(), + keep_raw: false, + } + } + + fn f32_blob(dim: usize, seed: u32) -> Bytes { + let mut v = Vec::with_capacity(dim * 4); + let mut s = seed; + for _ in 0..dim { + s = s.wrapping_mul(1664525).wrapping_add(1013904223); + let f = (s as f32) / (u32::MAX as f32); + v.extend_from_slice(&f.to_le_bytes()); + } + Bytes::from(v) + } + + /// Poll for `manifest.json` to appear (bounded, short interval) — the + /// background `global_snapshot_pool()` worker is a single shared thread + /// across every test in this binary, so its queue depth (hence latency) + /// is not under this test's control. Fails the assertion (returns + /// `None`) rather than hanging if the deadline is exceeded. + fn wait_for_manifest(idx_dir: &std::path::Path) -> Option { + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); + loop { + if let Some(m) = manifest::read_manifest_tolerant(idx_dir) { + return Some(m); + } + if std::time::Instant::now() >= deadline { + return None; + } + std::thread::sleep(std::time::Duration::from_millis(1)); + } + } + + /// No manifest on disk -> `create_index` behaves exactly like the plain + /// pre-B3 path (fresh collection_id, no recovered state). + #[test] + fn create_index_no_manifest_is_plain_fresh() { + let tmp = tempfile::tempdir().unwrap(); + let mut store = VectorStore::new(); + store.set_persist_dir(tmp.path().to_path_buf()); + let mut state = RecoveryState::new(); + let meta = make_meta("idx", 8, "doc:"); + + state.create_index(&mut store, tmp.path(), &meta); + + assert!(store.get_index(b"idx").is_some()); + assert!( + !state.recovered_names.contains(&Bytes::from_static(b"idx")), + "no manifest on disk -> index must NOT be marked recovered" + ); + } + + /// Round-trip: build a store with a persist_dir, insert vectors, force a + /// synchronous compact + snapshot so manifest/keymap/segments land on + /// disk, then recover into a FRESH store. Verifies: + /// - `verified_unchanged == n` (dedup actually fires for every key, not + /// just "results happen to match" — a full re-encode would also leave + /// every key indexed and the allocator seeded, so this counter check + /// is the one assertion that actually pins dedup behavior). + /// - the mutable segment's global-id allocator is seeded above every + /// loaded global_id (a post-recovery insert never collides). + /// - a same-query HNSW search returns the identical top-1 global_id + /// before persist and after recovery — this is real end-to-end proof + /// that `suggested_ef` round-trips AND that dedup preserves the + /// original global_id (a silent re-encode of doc:5 would reassign it + /// a fresh id from the seeded allocator and this assertion would go + /// red, independently of `verified_unchanged`'s own count check). + /// NOTE (verified by a temporary experiment, then reverted): offsetting + /// `create_index_with_collection_id`'s pinned cid by 1000 left this + /// test green. Root cause, confirmed by reading `ImmutableSegment:: + /// search` (immutable.rs:499): it passes `&self.collection_meta` — the + /// segment's OWN disk-reconstructed rotation/codebook data — into + /// `hnsw_search`, never the recreated index's collection object. After + /// `force_compact()` all data lives in that one self-contained + /// immutable segment, so a search hitting only pre-recovery data is + /// structurally unable to observe a wrong pin — true for every + /// quantization mode, not an SQ8-specific gap. The pin's actual + /// invariant is merge-compatibility of POST-recovery inserts (a future + /// re-compact/GraphUnion must stitch new codes over the old ones using + /// shared rotation) — untestable via search and out of scope for B3 + /// test hardening (would need a full second compact+merge cycle). + /// Recorded here as a known, deliberate test gap rather than implied + /// coverage. + #[test] + fn recover_round_trip_dedups_unchanged_keys_and_seeds_allocator() { + use crate::protocol::Frame; + + let tmp = tempfile::tempdir().unwrap(); + let dim = 8usize; + + // ---- Build + persist ---- + let mut store = VectorStore::new(); + store.set_persist_dir(tmp.path().to_path_buf()); + let meta = make_meta("idx", dim as u32, "doc:"); + store.create_index(meta.clone()).unwrap(); + + let n = 20usize; + for i in 0..n { + let key = format!("doc:{i}"); + let blob = f32_blob(dim, i as u32 + 1); + let args = vec![ + Frame::BulkString(Bytes::from(key.clone())), + Frame::BulkString(Bytes::from_static(b"vec")), + Frame::BulkString(blob), + ]; + let _ = crate::shard::spsc_handler::auto_index_hset_public( + &mut store, + &mut TextStore::new(), + key.as_bytes(), + &args, + ); + } + + { + let idx = store.get_index_mut(b"idx").unwrap(); + // `force_compact` writes the segment directory INLINE (the + // staged writer runs synchronously on this thread) but commits + // the manifest/keymap via `persist_hook_after_install`, which + // hands the job to the process-wide `global_snapshot_pool()` + // (a single background worker THREAD, shared by every test in + // this binary). Do NOT also build and run a second, manually + // constructed `SnapshotJob` here: both would race the exact + // same `manifest.json.tmp`/`keymap-1.bin.tmp` paths (fixed + // names, not job-unique) and can stomp each other's fsync — + // this was tried and is flaky. Instead just wait for the ONE + // real job the production code path already submitted. + idx.force_compact(); + } + let idx_dir = manifest::index_persist_dir(tmp.path(), b"idx"); + let loaded_manifest = wait_for_manifest(&idx_dir) + .expect("manifest must land within the timeout via global_snapshot_pool()"); + assert!(!loaded_manifest.segment_ids.is_empty()); + let max_loaded_global_id: u32 = { + let idx = store.get_index(b"idx").unwrap(); + idx.key_hash_to_global_id.values().copied().max().unwrap() + }; + + // Pre-persist search baseline: query with doc:5's own vector (exact + // match, distance 0) — nearest neighbor must be itself. Captured + // AFTER force_compact() so both sides of the comparison search the + // same shape of state (empty mutable + 1 immutable HNSW segment). + let query_idx = 5usize; + let query_key = format!("doc:{query_idx}"); + let query_blob_bytes = f32_blob(dim, query_idx as u32 + 1); + let query_f32: Vec = query_blob_bytes + .chunks_exact(4) + .map(|c| f32::from_le_bytes(c.try_into().unwrap())) + .collect(); + let expected_global_id = { + let idx = store.get_index(b"idx").unwrap(); + let kh = xxhash_rust::xxh64::xxh64(query_key.as_bytes(), 0); + *idx.key_hash_to_global_id.get(&kh).unwrap() + }; + let pre_results = store.search_index(b"idx", &query_f32, 1, 64).unwrap(); + assert_eq!( + pre_results.first().copied(), + Some(expected_global_id), + "sanity: pre-persist search must return the exact-match vector as top-1" + ); + + // ---- Recover into a FRESH store ---- + let mut fresh = VectorStore::new(); + fresh.set_persist_dir(tmp.path().to_path_buf()); + let mut state = RecoveryState::new(); + state.create_index(&mut fresh, tmp.path(), &meta); + assert!(state.recovered_names.contains(&Bytes::from_static(b"idx"))); + + let counters = *state.counters.get(&Bytes::from_static(b"idx")).unwrap(); + assert_eq!(counters.loaded_segments, loaded_manifest.segment_ids.len()); + + // Re-run the SAME HSETs through the dedup rescan — every one must + // be recognized as unchanged (checksum match), never re-encoded. + for i in 0..n { + let key = format!("doc:{i}"); + let blob = f32_blob(dim, i as u32 + 1); + let args = vec![ + Frame::BulkString(Bytes::from(key.clone())), + Frame::BulkString(Bytes::from_static(b"vec")), + Frame::BulkString(blob), + ]; + let mut text_store = TextStore::new(); + state.reconcile_key(&mut fresh, &mut text_store, key.as_bytes(), &args); + } + + // The load-bearing assertion: every key must have taken the + // metadata-only dedup path, not the full re-encode path. A full + // re-encode of all n keys would ALSO leave every key indexed and the + // allocator correctly seeded (checked below), so those checks alone + // cannot distinguish "dedup worked" from "dedup is silently dead". + let verified_unchanged = state + .counters + .get(&Bytes::from_static(b"idx")) + .unwrap() + .verified_unchanged; + assert_eq!( + verified_unchanged, n, + "every unchanged key must dedup (metadata-only rebuild), not re-encode" + ); + + state.finish(&mut fresh, tmp.path()); + + let final_counters = { + // finish() consumed `state`; re-derive expectations from the + // fresh store's live state instead. + fresh.get_index(b"idx").unwrap().key_hash_to_global_id.len() + }; + assert_eq!( + final_counters, n, + "all keys must still be indexed post-recovery" + ); + + // Post-recovery insert must get a global_id above every loaded one. + let new_key = b"doc:new".to_vec(); + let new_blob = f32_blob(dim, 9999); + let args = vec![ + Frame::BulkString(Bytes::from(new_key.clone())), + Frame::BulkString(Bytes::from_static(b"vec")), + Frame::BulkString(new_blob), + ]; + let mut text_store = TextStore::new(); + let _ = crate::shard::spsc_handler::auto_index_hset_public( + &mut fresh, + &mut text_store, + &new_key, + &args, + ); + let new_global_id = *fresh + .get_index(b"idx") + .unwrap() + .key_hash_to_global_id + .get(&xxhash_rust::xxh64::xxh64(&new_key, 0)) + .unwrap(); + assert!( + new_global_id > max_loaded_global_id, + "post-recovery insert global_id ({new_global_id}) must exceed every loaded \ + global_id ({max_loaded_global_id})" + ); + + // Post-recovery search must return the SAME top-1 as the + // pre-persist baseline — this is what actually exercises + // `suggested_ef` round-tripping and the collection_id/QJL-seed pin: + // a mismatched seed scrambles HNSW beam search (wrong/garbage + // distances), which a plain key-count check would never catch. + let post_results = fresh.search_index(b"idx", &query_f32, 1, 64).unwrap(); + assert_eq!( + post_results.first().copied(), + Some(expected_global_id), + "post-recovery search must return the same top-1 global_id as the pre-persist baseline" + ); + } + + /// A segment whose directory is entirely missing (simulating a crash + /// that lost the whole segment, headers included) must not panic + /// recovery — falls back to abandoning this index's recovered state + /// (full rescan covers it). + #[test] + fn recover_corrupt_segment_with_unreadable_headers_falls_back_cleanly() { + let tmp = tempfile::tempdir().unwrap(); + let idx_dir = manifest::index_persist_dir(tmp.path(), b"idx"); + std::fs::create_dir_all(&idx_dir).unwrap(); + + // Manifest references segment-1, but NOTHING is written on disk for + // it (not even mvcc_headers.bin) — the "headers also unreadable" + // case. + let m = IndexManifest { + format_version: manifest::MANIFEST_FORMAT_VERSION, + index_name_hex: manifest::index_name_hex(b"idx"), + collection_id: 7, + next_collection_id_floor: 8, + next_segment_id: 2, + next_global_id: 5, + segment_ids: vec![1], + keymap_epoch: 1, + }; + manifest::write_manifest_atomic(&idx_dir, &m).unwrap(); + + let mut store = VectorStore::new(); + store.set_persist_dir(tmp.path().to_path_buf()); + let meta = make_meta("idx", 8, "doc:"); + let mut state = RecoveryState::new(); + + state.create_index(&mut store, tmp.path(), &meta); + + // Index exists (pinned cid) but is NOT marked recovered. + assert!(store.get_index(b"idx").is_some()); + assert!(!state.recovered_names.contains(&Bytes::from_static(b"idx"))); + } + + /// Deletion probe (`finish()`'s `original_key_hashes` − `observed_key_hashes` + /// set difference): a key that existed at persist time but is never + /// observed during the rescan (deleted from the keyspace between + /// restarts) must be tombstoned — removed from the live index's + /// key_hash maps — even though it was never touched by `reconcile_key`. + #[test] + fn recover_finish_tombstones_keys_missing_from_rescan() { + use crate::protocol::Frame; + + let tmp = tempfile::tempdir().unwrap(); + let dim = 8usize; + + let mut store = VectorStore::new(); + store.set_persist_dir(tmp.path().to_path_buf()); + let meta = make_meta("idx", dim as u32, "doc:"); + store.create_index(meta.clone()).unwrap(); + + let n = 5usize; + for i in 0..n { + let key = format!("doc:{i}"); + let blob = f32_blob(dim, i as u32 + 1); + let args = vec![ + Frame::BulkString(Bytes::from(key.clone())), + Frame::BulkString(Bytes::from_static(b"vec")), + Frame::BulkString(blob), + ]; + let _ = crate::shard::spsc_handler::auto_index_hset_public( + &mut store, + &mut TextStore::new(), + key.as_bytes(), + &args, + ); + } + { + let idx = store.get_index_mut(b"idx").unwrap(); + idx.force_compact(); + } + let idx_dir = manifest::index_persist_dir(tmp.path(), b"idx"); + wait_for_manifest(&idx_dir) + .expect("manifest must land within the timeout via global_snapshot_pool()"); + + // Recover into a FRESH store, but only rescan n-1 of the n keys — + // "doc:3" is simulated as deleted from the keyspace between restarts + // (e.g. a DEL that happened while the server was down). + let mut fresh = VectorStore::new(); + fresh.set_persist_dir(tmp.path().to_path_buf()); + let mut state = RecoveryState::new(); + state.create_index(&mut fresh, tmp.path(), &meta); + assert!(state.recovered_names.contains(&Bytes::from_static(b"idx"))); + + let missing_key = "doc:3"; + let missing_key_hash = xxhash_rust::xxh64::xxh64(missing_key.as_bytes(), 0); + assert!( + fresh + .get_index(b"idx") + .unwrap() + .key_hash_to_global_id + .contains_key(&missing_key_hash), + "sanity: the to-be-deleted key must be present right after recovery, \ + before the rescan runs" + ); + + for i in 0..n { + if i == 3 { + continue; // never observed by reconcile_key — simulates a DEL + } + let key = format!("doc:{i}"); + let blob = f32_blob(dim, i as u32 + 1); + let args = vec![ + Frame::BulkString(Bytes::from(key.clone())), + Frame::BulkString(Bytes::from_static(b"vec")), + Frame::BulkString(blob), + ]; + let mut text_store = TextStore::new(); + state.reconcile_key(&mut fresh, &mut text_store, key.as_bytes(), &args); + } + state.finish(&mut fresh, tmp.path()); + + let idx = fresh.get_index(b"idx").unwrap(); + assert!( + !idx.key_hash_to_global_id.contains_key(&missing_key_hash), + "key never observed during the rescan must be tombstoned by finish()" + ); + assert!( + !idx.key_hash_to_key.contains_key(&missing_key_hash), + "tombstoning must also prune the key_hash_to_key map" + ); + assert_eq!( + idx.key_hash_to_global_id.len(), + n - 1, + "the other n-1 keys (all observed) must remain indexed" + ); + } +} diff --git a/src/vector/persistence/recovery.rs b/src/vector/persistence/recovery.rs deleted file mode 100644 index aed9fd69..00000000 --- a/src/vector/persistence/recovery.rs +++ /dev/null @@ -1,624 +0,0 @@ -//! Crash recovery for vector data: WAL replay + immutable segment loading. -//! -//! Recovery algorithm: -//! 1. Scan WAL file for vector record frames (tag 0x56) -//! 2. Replay VectorUpsert/Delete into MutableSegment per collection -//! 3. Handle TxnCommit/Abort/Checkpoint records -//! 4. Rollback uncommitted transactions at WAL end -//! 5. Load immutable segments from on-disk directories - -use std::collections::{HashMap, HashSet}; -use std::path::Path; -use std::sync::Arc; - -use tracing::{info, warn}; - -use crate::vector::persistence::segment_io::{SegmentIoError, read_immutable_segment}; -use crate::vector::persistence::wal_record::{VECTOR_RECORD_TAG, VectorWalRecord, WalRecordError}; -use crate::vector::segment::immutable::ImmutableSegment; -use crate::vector::segment::mutable::MutableSegment; -use crate::vector::turbo_quant::collection::CollectionMetadata; - -/// Error type for recovery operations. -#[derive(Debug)] -pub enum RecoveryError { - Io(std::io::Error), - SegmentLoad(SegmentIoError), -} - -impl From for RecoveryError { - fn from(e: std::io::Error) -> Self { - Self::Io(e) - } -} - -impl From for RecoveryError { - fn from(e: SegmentIoError) -> Self { - Self::SegmentLoad(e) - } -} - -/// Recovered collection data: mutable segment + immutable segments. -pub struct RecoveredCollection { - pub mutable: MutableSegment, - pub immutable: Vec<(ImmutableSegment, Arc)>, -} - -/// Full recovered state from WAL + disk segments. -pub struct RecoveredState { - /// collection_id -> recovered collection data - pub collections: HashMap, - /// Last checkpoint LSN seen (for future WAL truncation) - pub last_checkpoint_lsn: u64, -} - -/// State accumulated during WAL replay for one collection. -struct CollectionReplayState { - mutable: MutableSegment, - /// point_id -> internal_id in mutable segment - point_map: HashMap, - /// txn_id -> list of internal_ids inserted by that txn - pending_txns: HashMap>, - /// Committed txn_ids - committed_txns: HashSet, - #[allow(dead_code)] - dimension: u32, -} - -/// Scan WAL bytes for vector record frames. -/// -/// Skips RESP block frames (identified by not having the VECTOR_RECORD_TAG). -/// Stops on CRC mismatch, truncation, or any parse error (conservative). -fn scan_vector_records(wal_data: &[u8]) -> Vec { - let mut records = Vec::new(); - let mut pos = 32; // skip WAL header - while pos < wal_data.len() { - if wal_data[pos] == VECTOR_RECORD_TAG { - match VectorWalRecord::from_wal_frame(&wal_data[pos..]) { - Ok((record, consumed)) => { - records.push(record); - pos += consumed; - } - Err(WalRecordError::CrcMismatch { .. }) => { - warn!("CRC mismatch at WAL offset {}, stopping vector replay", pos); - break; - } - Err(WalRecordError::Truncated) => { - warn!("Truncated vector record at WAL offset {}, stopping", pos); - break; - } - Err(e) => { - warn!("Vector WAL record error at offset {}: {}, stopping", pos, e); - break; - } - } - } else { - // RESP block frame -- skip it - if pos + 4 > wal_data.len() { - break; - } - let block_len = u32::from_le_bytes([ - wal_data[pos], - wal_data[pos + 1], - wal_data[pos + 2], - wal_data[pos + 3], - ]) as usize; - if block_len > 100_000_000 || pos + 4 + block_len > wal_data.len() { - warn!( - "Vector WAL: invalid RESP block length {} at offset {}, stopping recovery", - block_len, pos - ); - break; - } - pos += 4 + block_len; - } - } - records -} - -/// Enumerate segment directories in a persistence directory. -/// -/// Looks for directories named `segment-{id}` and returns sorted IDs. -fn enumerate_segments(dir: &Path) -> Vec { - let mut ids = Vec::new(); - if let Ok(entries) = std::fs::read_dir(dir) { - for entry in entries.flatten() { - if let Some(name) = entry.file_name().to_str() { - if let Some(id_str) = name.strip_prefix("segment-") { - if let Ok(id) = id_str.parse::() { - ids.push(id); - } - } - } - } - } - ids.sort(); - ids -} - -/// Replay vector WAL records into per-collection mutable segments. -/// -/// Returns map of collection_id -> MutableSegment plus last checkpoint LSN. -fn replay_vector_wal(records: &[VectorWalRecord]) -> (HashMap, u64) { - let mut states: HashMap = HashMap::new(); - let mut last_checkpoint_lsn: u64 = 0; - let mut next_lsn: u64 = 1; - - for record in records { - match record { - VectorWalRecord::VectorUpsert { - txn_id, - collection_id, - point_id, - sq_vector: _, - tq_code: _, - norm: _, - f32_vector, - } => { - let dim = f32_vector.len() as u32; - let state = states.entry(*collection_id).or_insert_with(|| { - CollectionReplayState { - mutable: MutableSegment::new(dim, std::sync::Arc::new( - crate::vector::turbo_quant::collection::CollectionMetadata::new( - *collection_id, dim, - crate::vector::types::DistanceMetric::L2, - crate::vector::turbo_quant::collection::QuantizationConfig::TurboQuant4, - *collection_id, - ))), - point_map: HashMap::new(), - pending_txns: HashMap::new(), - committed_txns: HashSet::new(), - dimension: dim, - } - }); - - let internal_id = if *txn_id != 0 { - state - .mutable - .append_transactional(*point_id, f32_vector, next_lsn, *txn_id) - } else { - state.mutable.append(*point_id, f32_vector, next_lsn) - }; - state.point_map.insert(*point_id, internal_id); - if *txn_id != 0 { - state - .pending_txns - .entry(*txn_id) - .or_default() - .push(internal_id); - } - next_lsn += 1; - } - VectorWalRecord::VectorDelete { - txn_id, - collection_id, - point_id, - } => { - if let Some(state) = states.get(collection_id) { - if let Some(&internal_id) = state.point_map.get(point_id) { - state.mutable.mark_deleted(internal_id, next_lsn); - } - // If point_id not found, skip silently (no panic) - } - // Track in pending txns for potential abort rollback - // (deletes don't add internal_ids -- they mark existing ones) - let _ = txn_id; // used below if needed - next_lsn += 1; - } - VectorWalRecord::TxnCommit { txn_id, commit_lsn } => { - // Mark txn as committed in all collections - for state in states.values_mut() { - if state.pending_txns.contains_key(txn_id) { - state.committed_txns.insert(*txn_id); - } - } - let _ = commit_lsn; - } - VectorWalRecord::TxnAbort { txn_id } => { - // Roll back all entries from this txn - for state in states.values() { - if let Some(internal_ids) = state.pending_txns.get(txn_id) { - for &iid in internal_ids { - state.mutable.mark_deleted(iid, next_lsn); - } - } - } - next_lsn += 1; - } - VectorWalRecord::Checkpoint { - segment_id: _, - last_lsn, - } => { - last_checkpoint_lsn = *last_lsn; - } - } - } - - // Rollback uncommitted transactions at end of WAL - for state in states.values() { - for (txn_id, internal_ids) in &state.pending_txns { - if !state.committed_txns.contains(txn_id) { - for &iid in internal_ids { - state.mutable.mark_deleted(iid, next_lsn); - } - } - } - } - - let mut result = HashMap::new(); - for (cid, state) in states { - result.insert(cid, state.mutable); - } - (result, last_checkpoint_lsn) -} - -/// Recover vector store state from WAL + on-disk segments. -/// -/// 1. Enumerate segment directories, load each immutable segment. -/// 2. Read WAL file, extract vector record frames. -/// 3. Replay into MutableSegment per collection. -/// 4. Rollback uncommitted transactions. -/// 5. Return RecoveredState with all collections. -pub fn recover_vector_store( - wal_path: &Path, - persist_dir: &Path, -) -> Result { - let mut collections: HashMap = HashMap::new(); - - // 1. Load immutable segments from disk - let segment_ids = enumerate_segments(persist_dir); - for seg_id in &segment_ids { - match read_immutable_segment(persist_dir, *seg_id) { - Ok((segment, meta)) => { - let cid = meta.collection_id; - info!("Loaded immutable segment {} for collection {}", seg_id, cid); - let entry = collections.entry(cid).or_insert_with(|| { - RecoveredCollection { - mutable: MutableSegment::new(meta.dimension, std::sync::Arc::new( - crate::vector::turbo_quant::collection::CollectionMetadata::new( - cid, meta.dimension, - crate::vector::types::DistanceMetric::L2, - crate::vector::turbo_quant::collection::QuantizationConfig::TurboQuant4, - cid, - ))), - immutable: Vec::new(), - } - }); - entry.immutable.push((segment, meta)); - } - Err(e) => { - warn!("Failed to load segment {}: {:?}, skipping", seg_id, e); - } - } - } - - // 2. Read WAL and extract vector records - let mut last_checkpoint_lsn = 0u64; - if wal_path.exists() { - let wal_data = std::fs::read(wal_path)?; - if wal_data.len() > 32 { - let records = scan_vector_records(&wal_data); - info!("Scanned {} vector WAL records", records.len()); - - // 3. Replay into mutable segments - let (mutable_map, ckpt_lsn) = replay_vector_wal(&records); - last_checkpoint_lsn = ckpt_lsn; - - // 4. Merge mutable segments into collections - for (cid, mutable) in mutable_map { - match collections.entry(cid) { - std::collections::hash_map::Entry::Vacant(e) => { - e.insert(RecoveredCollection { - mutable, - immutable: Vec::new(), - }); - } - std::collections::hash_map::Entry::Occupied(mut e) => { - // Collection already has immutable segments from disk. - // Replace the placeholder mutable with the replayed one. - e.get_mut().mutable = mutable; - } - } - } - } - } - - Ok(RecoveredState { - collections, - last_checkpoint_lsn, - }) -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::vector::persistence::wal_record::VectorWalRecord; - - /// Build a minimal WAL file header (32 bytes). - fn make_wal_header() -> Vec { - let mut header = vec![0u8; 32]; - header[0..6].copy_from_slice(b"RRDWAL"); - header[6] = 2; // version - header - } - - #[test] - fn test_wal_writer_append_vector_record_roundtrip() { - // Write a vector record frame, then parse it back - let record = VectorWalRecord::VectorUpsert { - txn_id: 0, - collection_id: 1, - point_id: 42, - sq_vector: vec![1, -2, 3, -4], - tq_code: vec![0xAB], - norm: 1.5, - f32_vector: vec![0.1, 0.2, 0.3, 0.4], - }; - let frame = record.to_wal_frame(); - - // Simulate what append_vector_record does: just buffer the frame bytes - let mut buf = Vec::new(); - buf.extend_from_slice(&frame); - - // Parse back - let (decoded, consumed) = VectorWalRecord::from_wal_frame(&buf).unwrap(); - assert_eq!(consumed, frame.len()); - assert_eq!(decoded, record); - } - - #[test] - fn test_recover_mutable_upsert_count() { - let records = vec![ - VectorWalRecord::VectorUpsert { - txn_id: 0, - collection_id: 1, - point_id: 10, - sq_vector: vec![1, 2, 3, 4], - tq_code: vec![], - norm: 1.0, - f32_vector: vec![0.1, 0.2, 0.3, 0.4], - }, - VectorWalRecord::VectorUpsert { - txn_id: 0, - collection_id: 1, - point_id: 20, - sq_vector: vec![5, 6, 7, 8], - tq_code: vec![], - norm: 1.0, - f32_vector: vec![0.5, 0.6, 0.7, 0.8], - }, - ]; - let (mutables, _) = replay_vector_wal(&records); - let seg = mutables.get(&1).unwrap(); - assert_eq!(seg.len(), 2); - } - - #[test] - fn test_recover_mutable_delete_nonexistent_no_panic() { - // Delete a point_id that was never upserted -- should not panic - let records = vec![VectorWalRecord::VectorDelete { - txn_id: 0, - collection_id: 1, - point_id: 999, - }]; - let (mutables, _) = replay_vector_wal(&records); - // No collection created because no upserts - assert!(mutables.is_empty() || mutables.get(&1).map_or(true, |s| s.is_empty())); - } - - #[test] - fn test_recover_mutable_delete_marks_entry() { - let records = vec![ - VectorWalRecord::VectorUpsert { - txn_id: 0, - collection_id: 1, - point_id: 10, - sq_vector: vec![1, 2, 3, 4], - tq_code: vec![], - norm: 1.0, - f32_vector: vec![0.1, 0.2, 0.3, 0.4], - }, - VectorWalRecord::VectorDelete { - txn_id: 0, - collection_id: 1, - point_id: 10, - }, - ]; - let (mutables, _) = replay_vector_wal(&records); - let seg = mutables.get(&1).unwrap(); - // The entry is still there but marked deleted - assert_eq!(seg.len(), 1); - let frozen = seg.freeze(); - assert_ne!(frozen.entries[0].delete_lsn, 0); - } - - #[test] - fn test_recover_txn_abort_rolls_back() { - let records = vec![ - VectorWalRecord::VectorUpsert { - txn_id: 42, - collection_id: 1, - point_id: 10, - sq_vector: vec![1, 2, 3, 4], - tq_code: vec![], - norm: 1.0, - f32_vector: vec![0.1, 0.2, 0.3, 0.4], - }, - VectorWalRecord::TxnAbort { txn_id: 42 }, - ]; - let (mutables, _) = replay_vector_wal(&records); - let seg = mutables.get(&1).unwrap(); - let frozen = seg.freeze(); - // Entry should be marked deleted due to abort - assert_ne!(frozen.entries[0].delete_lsn, 0); - } - - #[test] - fn test_recover_uncommitted_at_eof_rolled_back() { - // Upsert in a txn, no commit or abort -- should be rolled back - let records = vec![VectorWalRecord::VectorUpsert { - txn_id: 99, - collection_id: 1, - point_id: 10, - sq_vector: vec![1, 2, 3, 4], - tq_code: vec![], - norm: 1.0, - f32_vector: vec![0.1, 0.2, 0.3, 0.4], - }]; - let (mutables, _) = replay_vector_wal(&records); - let seg = mutables.get(&1).unwrap(); - let frozen = seg.freeze(); - assert_ne!( - frozen.entries[0].delete_lsn, 0, - "uncommitted txn should be rolled back" - ); - } - - #[test] - fn test_recover_committed_txn_survives() { - let records = vec![ - VectorWalRecord::VectorUpsert { - txn_id: 42, - collection_id: 1, - point_id: 10, - sq_vector: vec![1, 2, 3, 4], - tq_code: vec![], - norm: 1.0, - f32_vector: vec![0.1, 0.2, 0.3, 0.4], - }, - VectorWalRecord::TxnCommit { - txn_id: 42, - commit_lsn: 100, - }, - ]; - let (mutables, _) = replay_vector_wal(&records); - let seg = mutables.get(&1).unwrap(); - let frozen = seg.freeze(); - assert_eq!( - frozen.entries[0].delete_lsn, 0, - "committed entry should not be deleted" - ); - } - - #[test] - fn test_recover_checkpoint_records_lsn() { - let records = vec![ - VectorWalRecord::Checkpoint { - segment_id: 5, - last_lsn: 500, - }, - VectorWalRecord::Checkpoint { - segment_id: 6, - last_lsn: 600, - }, - ]; - let (_, last_ckpt) = replay_vector_wal(&records); - assert_eq!(last_ckpt, 600); - } - - #[test] - fn test_recover_empty_wal_and_no_segments() { - let tmp = tempfile::tempdir().unwrap(); - let wal_path = tmp.path().join("shard-0.wal"); - let persist_dir = tmp.path().join("vectors"); - // Neither file nor directory exists - let result = recover_vector_store(&wal_path, &persist_dir).unwrap(); - assert!(result.collections.is_empty()); - assert_eq!(result.last_checkpoint_lsn, 0); - } - - #[test] - fn test_recover_vector_store_from_wal() { - let tmp = tempfile::tempdir().unwrap(); - let persist_dir = tmp.path().join("vectors"); - std::fs::create_dir_all(&persist_dir).unwrap(); - - // Build a WAL file with vector records - let mut wal_data = make_wal_header(); - - let upsert1 = VectorWalRecord::VectorUpsert { - txn_id: 0, - collection_id: 1, - point_id: 10, - sq_vector: vec![1, 2, 3, 4], - tq_code: vec![], - norm: 1.0, - f32_vector: vec![0.1, 0.2, 0.3, 0.4], - }; - let upsert2 = VectorWalRecord::VectorUpsert { - txn_id: 0, - collection_id: 1, - point_id: 20, - sq_vector: vec![5, 6, 7, 8], - tq_code: vec![], - norm: 2.0, - f32_vector: vec![0.5, 0.6, 0.7, 0.8], - }; - wal_data.extend_from_slice(&upsert1.to_wal_frame()); - wal_data.extend_from_slice(&upsert2.to_wal_frame()); - - let wal_path = tmp.path().join("shard-0.wal"); - std::fs::write(&wal_path, &wal_data).unwrap(); - - let result = recover_vector_store(&wal_path, &persist_dir).unwrap(); - assert_eq!(result.collections.len(), 1); - let coll = result.collections.get(&1).unwrap(); - assert_eq!(coll.mutable.len(), 2); - } - - #[test] - fn test_recover_corrupt_crc_stops_replay() { - let tmp = tempfile::tempdir().unwrap(); - let persist_dir = tmp.path().join("vectors"); - std::fs::create_dir_all(&persist_dir).unwrap(); - - let mut wal_data = make_wal_header(); - - // Good record - let good = VectorWalRecord::VectorUpsert { - txn_id: 0, - collection_id: 1, - point_id: 10, - sq_vector: vec![1, 2, 3, 4], - tq_code: vec![], - norm: 1.0, - f32_vector: vec![0.1, 0.2, 0.3, 0.4], - }; - wal_data.extend_from_slice(&good.to_wal_frame()); - - // Corrupt record - let mut bad_frame = VectorWalRecord::VectorUpsert { - txn_id: 0, - collection_id: 1, - point_id: 20, - sq_vector: vec![5, 6, 7, 8], - tq_code: vec![], - norm: 2.0, - f32_vector: vec![0.5, 0.6, 0.7, 0.8], - } - .to_wal_frame(); - let len = bad_frame.len(); - bad_frame[len - 1] ^= 0xFF; // corrupt CRC - wal_data.extend_from_slice(&bad_frame); - - // Third record that should NOT be recovered - let third = VectorWalRecord::VectorUpsert { - txn_id: 0, - collection_id: 1, - point_id: 30, - sq_vector: vec![9, 10, 11, 12], - tq_code: vec![], - norm: 3.0, - f32_vector: vec![0.9, 1.0, 1.1, 1.2], - }; - wal_data.extend_from_slice(&third.to_wal_frame()); - - let wal_path = tmp.path().join("shard-0.wal"); - std::fs::write(&wal_path, &wal_data).unwrap(); - - let result = recover_vector_store(&wal_path, &persist_dir).unwrap(); - // Only the first record should be recovered (CRC stops at second) - let coll = result.collections.get(&1).unwrap(); - assert_eq!(coll.mutable.len(), 1, "corrupt CRC should stop replay"); - } -} diff --git a/src/vector/persistence/segment_io.rs b/src/vector/persistence/segment_io.rs index c842837a..df6b4774 100644 --- a/src/vector/persistence/segment_io.rs +++ b/src/vector/persistence/segment_io.rs @@ -76,6 +76,12 @@ struct SegmentMeta { /// Build mode: "Light" or "Exact". Added in v1 — defaults to inferred if absent. #[serde(default)] build_mode: Option, + /// Compact-time adaptive-ef estimate (AE-1/R6). Added post-v1 — + /// `#[serde(default)]` keeps segment dirs written before this field + /// existed readable (they simply carry `None`, same as an in-memory + /// segment for which the estimator declined to make a call). + #[serde(default)] + suggested_ef: Option, } fn segment_dir(dir: &Path, segment_id: u64) -> PathBuf { @@ -131,6 +137,15 @@ fn string_to_quant(s: &str) -> Result { /// Write an immutable segment to disk. /// /// Creates `{dir}/segment-{id}/` with 5 files. +/// +/// This is the low-level primitive: it writes directly under the FINAL +/// `segment-{id}` name. Production callers that need crash-safe atomicity +/// (a segment must never be considered "live" until it is fully durable) +/// should use [`write_immutable_segment_staged`] instead, which writes to a +/// `staging-{id}` name first and only `rename`s to `segment-{id}` once every +/// file is written and fsynced. This primitive is kept for tests and any +/// caller that already provides its own atomicity (e.g. a fresh, private +/// scratch directory). pub fn write_immutable_segment( dir: &Path, segment_id: u64, @@ -138,7 +153,52 @@ pub fn write_immutable_segment( collection: &CollectionMetadata, ) -> Result<(), SegmentIoError> { let seg_dir = segment_dir(dir, segment_id); - fs::create_dir_all(&seg_dir)?; + write_segment_files(&seg_dir, segment_id, segment, collection) +} + +/// Write an immutable segment to disk with crash-safe atomicity. +/// +/// Writes into `{root_dir}/staging-{id}/` (removing any stale leftover from +/// a prior interrupted attempt first), fsyncs every file plus the staging +/// directory itself (inside [`write_segment_files`]), then atomically +/// `rename`s `staging-{id}` -> `segment-{id}` and fsyncs `root_dir` so the +/// rename (a directory-entry change) is itself durable. +/// +/// Until the `rename` completes, `segment-{id}` simply does not exist — a +/// crash mid-write leaves only an orphan `staging-{id}` directory, which is +/// safe to delete (never referenced by any manifest). This is the write path +/// used by compaction/merge when a `persist_dir` is configured (B2). +pub fn write_immutable_segment_staged( + root_dir: &Path, + segment_id: u64, + segment: &ImmutableSegment, + collection: &CollectionMetadata, +) -> Result<(), SegmentIoError> { + fs::create_dir_all(root_dir)?; + let staging_dir = root_dir.join(format!("staging-{segment_id}")); + if staging_dir.exists() { + // Leftover from a prior crashed attempt at the same segment id. + // Best-effort removal — write_segment_files will recreate it. + let _ = fs::remove_dir_all(&staging_dir); + } + write_segment_files(&staging_dir, segment_id, segment, collection)?; + let final_dir = segment_dir(root_dir, segment_id); + fs::rename(&staging_dir, &final_dir)?; + fsync_directory(root_dir)?; + Ok(()) +} + +/// Shared file-writing body for both [`write_immutable_segment`] (writes +/// directly to the final name) and [`write_immutable_segment_staged`] +/// (writes to a staging name first). `seg_dir` is the exact target +/// directory to create and populate. +fn write_segment_files( + seg_dir: &Path, + segment_id: u64, + segment: &ImmutableSegment, + collection: &CollectionMetadata, +) -> Result<(), SegmentIoError> { + fs::create_dir_all(seg_dir)?; // 1. hnsw_graph.bin let graph_bytes = segment.graph().to_bytes(); @@ -207,6 +267,7 @@ pub fn write_immutable_segment( crate::vector::turbo_quant::collection::BuildMode::Light => "Light".to_owned(), crate::vector::turbo_quant::collection::BuildMode::Exact => "Exact".to_owned(), }), + suggested_ef: segment.suggested_ef(), }; let json = serde_json::to_string_pretty(&meta) .map_err(|e| SegmentIoError::InvalidMetadata(e.to_string()))?; @@ -215,11 +276,128 @@ pub fn write_immutable_segment( fsync_file(&meta_path)?; // Fsync the segment directory to make all file entries durable - fsync_directory(&seg_dir)?; + fsync_directory(seg_dir)?; Ok(()) } +/// Parse the raw contents of an `mvcc_headers.bin` file (version-aware: +/// v1 = 20 bytes/header with no `global_id`/`key_hash`, v2 = 32 bytes/header). +/// Shared by [`read_immutable_segment`] (full segment read) and +/// [`read_mvcc_headers_only`] (B3 best-effort recovery helper, used when the +/// rest of a segment fails to load but the key_hashes it held still need to +/// be identified so they can be dropped from the recovered keymap). +fn parse_mvcc_headers(mvcc_bytes: &[u8]) -> Result, SegmentIoError> { + if mvcc_bytes.len() < 4 { + return Err(SegmentIoError::InvalidMetadata( + "mvcc_headers.bin too short".to_owned(), + )); + } + // Detect format version: v2 starts with version byte 2, v1 starts with count (u32 LE) + let (mvcc_version, mvcc_count, mut pos) = if mvcc_bytes[0] == 2 && mvcc_bytes.len() >= 5 { + let count = u32::from_le_bytes([mvcc_bytes[1], mvcc_bytes[2], mvcc_bytes[3], mvcc_bytes[4]]) + as usize; + (2u8, count, 5usize) + } else { + let count = u32::from_le_bytes([mvcc_bytes[0], mvcc_bytes[1], mvcc_bytes[2], mvcc_bytes[3]]) + as usize; + (1u8, count, 4usize) + }; + let bytes_per_header: usize = if mvcc_version >= 2 { 32 } else { 20 }; + if mvcc_bytes.len() < pos + mvcc_count * bytes_per_header { + return Err(SegmentIoError::InvalidMetadata( + "mvcc_headers.bin truncated".to_owned(), + )); + } + let mut mvcc = Vec::with_capacity(mvcc_count); + for _ in 0..mvcc_count { + let internal_id = u32::from_le_bytes([ + mvcc_bytes[pos], + mvcc_bytes[pos + 1], + mvcc_bytes[pos + 2], + mvcc_bytes[pos + 3], + ]); + pos += 4; + let (global_id, key_hash) = if mvcc_version >= 2 { + let gid = u32::from_le_bytes([ + mvcc_bytes[pos], + mvcc_bytes[pos + 1], + mvcc_bytes[pos + 2], + mvcc_bytes[pos + 3], + ]); + pos += 4; + let kh = u64::from_le_bytes([ + mvcc_bytes[pos], + mvcc_bytes[pos + 1], + mvcc_bytes[pos + 2], + mvcc_bytes[pos + 3], + mvcc_bytes[pos + 4], + mvcc_bytes[pos + 5], + mvcc_bytes[pos + 6], + mvcc_bytes[pos + 7], + ]); + pos += 8; + (gid, kh) + } else { + (internal_id, 0u64) // v1 fallback: global_id = internal_id + }; + let insert_lsn = u64::from_le_bytes([ + mvcc_bytes[pos], + mvcc_bytes[pos + 1], + mvcc_bytes[pos + 2], + mvcc_bytes[pos + 3], + mvcc_bytes[pos + 4], + mvcc_bytes[pos + 5], + mvcc_bytes[pos + 6], + mvcc_bytes[pos + 7], + ]); + pos += 8; + let delete_lsn = u64::from_le_bytes([ + mvcc_bytes[pos], + mvcc_bytes[pos + 1], + mvcc_bytes[pos + 2], + mvcc_bytes[pos + 3], + mvcc_bytes[pos + 4], + mvcc_bytes[pos + 5], + mvcc_bytes[pos + 6], + mvcc_bytes[pos + 7], + ]); + pos += 8; + mvcc.push(MvccHeader { + internal_id, + global_id, + key_hash, + insert_lsn, + delete_lsn, + hint_committed: 0, + }); + } + Ok(mvcc) +} + +/// Best-effort recovery helper (B3): read and parse ONLY the +/// `mvcc_headers.bin` file for a segment, without touching the graph/TQ/meta +/// files or verifying any checksum. +/// +/// Used when a segment's full [`read_immutable_segment`] fails (corrupt +/// graph, checksum mismatch, collection-id mismatch against the manifest) — +/// the B3 loader still needs the `key_hash` of every entry that WAS durable +/// in the abandoned segment, so it can drop exactly those keys from the +/// recovered keymap (letting the rescan re-index them) instead of either +/// resurrecting stale data or silently losing track of which keys are +/// affected. +/// +/// Returns `None` on ANY I/O or parse failure. The caller's fallback for +/// that case (headers ALSO unreadable) is to abandon the entire index's +/// on-disk state — an unattributed set of potentially-lost keys can't be +/// safely subtracted from the keymap without risking duplicate/stale search +/// results (see `recover_v2` module docs). +pub fn read_mvcc_headers_only(dir: &Path, segment_id: u64) -> Option> { + let seg_dir = segment_dir(dir, segment_id); + let mvcc_bytes = fs::read(seg_dir.join("mvcc_headers.bin")).ok()?; + parse_mvcc_headers(&mvcc_bytes).ok() +} + /// Read an immutable segment from disk. /// /// Reads from `{dir}/segment-{id}/` directory. @@ -364,90 +542,7 @@ pub fn read_immutable_segment( // 5. Read MVCC headers (version-aware: v1 = 20 bytes/header, v2 = 32 bytes/header) let mvcc_bytes = fs::read(seg_dir.join("mvcc_headers.bin"))?; - if mvcc_bytes.len() < 4 { - return Err(SegmentIoError::InvalidMetadata( - "mvcc_headers.bin too short".to_owned(), - )); - } - // Detect format version: v2 starts with version byte 2, v1 starts with count (u32 LE) - let (mvcc_version, mvcc_count, mut pos) = if mvcc_bytes[0] == 2 && mvcc_bytes.len() >= 5 { - let count = u32::from_le_bytes([mvcc_bytes[1], mvcc_bytes[2], mvcc_bytes[3], mvcc_bytes[4]]) - as usize; - (2u8, count, 5usize) - } else { - let count = u32::from_le_bytes([mvcc_bytes[0], mvcc_bytes[1], mvcc_bytes[2], mvcc_bytes[3]]) - as usize; - (1u8, count, 4usize) - }; - let bytes_per_header: usize = if mvcc_version >= 2 { 32 } else { 20 }; - if mvcc_bytes.len() < pos + mvcc_count * bytes_per_header { - return Err(SegmentIoError::InvalidMetadata( - "mvcc_headers.bin truncated".to_owned(), - )); - } - let mut mvcc = Vec::with_capacity(mvcc_count); - for _ in 0..mvcc_count { - let internal_id = u32::from_le_bytes([ - mvcc_bytes[pos], - mvcc_bytes[pos + 1], - mvcc_bytes[pos + 2], - mvcc_bytes[pos + 3], - ]); - pos += 4; - let (global_id, key_hash) = if mvcc_version >= 2 { - let gid = u32::from_le_bytes([ - mvcc_bytes[pos], - mvcc_bytes[pos + 1], - mvcc_bytes[pos + 2], - mvcc_bytes[pos + 3], - ]); - pos += 4; - let kh = u64::from_le_bytes([ - mvcc_bytes[pos], - mvcc_bytes[pos + 1], - mvcc_bytes[pos + 2], - mvcc_bytes[pos + 3], - mvcc_bytes[pos + 4], - mvcc_bytes[pos + 5], - mvcc_bytes[pos + 6], - mvcc_bytes[pos + 7], - ]); - pos += 8; - (gid, kh) - } else { - (internal_id, 0u64) // v1 fallback: global_id = internal_id - }; - let insert_lsn = u64::from_le_bytes([ - mvcc_bytes[pos], - mvcc_bytes[pos + 1], - mvcc_bytes[pos + 2], - mvcc_bytes[pos + 3], - mvcc_bytes[pos + 4], - mvcc_bytes[pos + 5], - mvcc_bytes[pos + 6], - mvcc_bytes[pos + 7], - ]); - pos += 8; - let delete_lsn = u64::from_le_bytes([ - mvcc_bytes[pos], - mvcc_bytes[pos + 1], - mvcc_bytes[pos + 2], - mvcc_bytes[pos + 3], - mvcc_bytes[pos + 4], - mvcc_bytes[pos + 5], - mvcc_bytes[pos + 6], - mvcc_bytes[pos + 7], - ]); - pos += 8; - mvcc.push(MvccHeader { - internal_id, - global_id, - key_hash, - insert_lsn, - delete_lsn, - hint_committed: 0, - }); - } + let mvcc = parse_mvcc_headers(&mvcc_bytes)?; // 6. Construct ImmutableSegment let dim = meta.dimension as usize; @@ -489,7 +584,11 @@ pub fn read_immutable_segment( meta.live_count, meta.total_count, ) - .with_raw_f16(raw_f16); + .with_raw_f16(raw_f16) + // R6: restore the compact-time estimate verbatim — never re-run the + // estimator on the load path (it needs the raw sidecar + a full ladder + // walk; that cost belongs on the compaction thread only). + .with_suggested_ef_raw(meta.suggested_ef); Ok((segment, collection)) } @@ -649,6 +748,98 @@ mod tests { assert!(seg_dir.join("segment_meta.json").exists()); } + /// R6: a segment carrying a measured `suggested_ef` must round-trip it + /// verbatim through disk — the read path restores it raw (no + /// re-estimation), so this also guards against `read_immutable_segment` + /// accidentally calling `with_adaptive_ef()` instead of + /// `with_suggested_ef_raw()`. + #[test] + fn test_suggested_ef_roundtrip() { + let (segment, collection) = build_test_segment(20, 64); + let segment = segment.with_suggested_ef_raw(Some(48)); + let tmp = tempfile::tempdir().unwrap(); + + write_immutable_segment(tmp.path(), 1, &segment, &collection).unwrap(); + let (restored, _) = read_immutable_segment(tmp.path(), 1).unwrap(); + + assert_eq!(restored.suggested_ef(), Some(48)); + } + + /// A segment with no measured estimate persists/reloads as `None`. + #[test] + fn test_suggested_ef_none_roundtrips_none() { + let (segment, collection) = build_test_segment(20, 64); + assert_eq!(segment.suggested_ef(), None); + let tmp = tempfile::tempdir().unwrap(); + + write_immutable_segment(tmp.path(), 1, &segment, &collection).unwrap(); + let (restored, _) = read_immutable_segment(tmp.path(), 1).unwrap(); + + assert_eq!(restored.suggested_ef(), None); + } + + /// Backward compatibility: a `segment_meta.json` written before the + /// `suggested_ef` field existed (no such key at all) must still load, + /// via `#[serde(default)]`, as `None` rather than erroring. + #[test] + fn test_suggested_ef_absent_key_in_old_segment_meta_loads_as_none() { + let (segment, collection) = build_test_segment(20, 64); + let tmp = tempfile::tempdir().unwrap(); + + write_immutable_segment(tmp.path(), 1, &segment, &collection).unwrap(); + + // Simulate a pre-R6 segment_meta.json by stripping the field entirely. + let meta_path = tmp.path().join("segment-1").join("segment_meta.json"); + let json_str = std::fs::read_to_string(&meta_path).unwrap(); + let mut val: serde_json::Value = serde_json::from_str(&json_str).unwrap(); + val.as_object_mut().unwrap().remove("suggested_ef"); + std::fs::write(&meta_path, serde_json::to_string_pretty(&val).unwrap()).unwrap(); + + let (restored, _) = read_immutable_segment(tmp.path(), 1).unwrap(); + assert_eq!(restored.suggested_ef(), None); + } + + /// `write_immutable_segment_staged` must produce a directory readable by + /// `read_immutable_segment` under the FINAL name, with no leftover + /// `staging-*` directory once the rename completes. + #[test] + fn test_staged_write_renames_to_final_and_leaves_no_staging_dir() { + let (segment, collection) = build_test_segment(20, 64); + let tmp = tempfile::tempdir().unwrap(); + + write_immutable_segment_staged(tmp.path(), 7, &segment, &collection).unwrap(); + + assert!(tmp.path().join("segment-7").exists()); + assert!(!tmp.path().join("staging-7").exists()); + let (restored, _) = read_immutable_segment(tmp.path(), 7).unwrap(); + assert_eq!(restored.live_count(), segment.live_count()); + } + + /// A stale `staging-{id}` directory left over from a prior crashed write + /// (e.g. partial files, no `segment_meta.json`) must not block a fresh + /// staged write from succeeding — it is removed and rebuilt. + #[test] + fn test_staged_write_recovers_from_stale_staging_dir() { + let (segment, collection) = build_test_segment(20, 64); + let tmp = tempfile::tempdir().unwrap(); + + let stale_staging = tmp.path().join("staging-3"); + std::fs::create_dir_all(&stale_staging).unwrap(); + std::fs::write(stale_staging.join("garbage.bin"), b"not a real segment").unwrap(); + + write_immutable_segment_staged(tmp.path(), 3, &segment, &collection).unwrap(); + + assert!( + tmp.path() + .join("segment-3") + .join("segment_meta.json") + .exists() + ); + assert!(!stale_staging.exists()); + let (restored, _) = read_immutable_segment(tmp.path(), 3).unwrap(); + assert_eq!(restored.live_count(), segment.live_count()); + } + #[test] fn test_roundtrip_preserves_counts() { let (segment, collection) = build_test_segment(30, 64); diff --git a/src/vector/persistence/wal_record.rs b/src/vector/persistence/wal_record.rs index 0364a4a5..bec5c595 100644 --- a/src/vector/persistence/wal_record.rs +++ b/src/vector/persistence/wal_record.rs @@ -8,6 +8,16 @@ //! [payload bytes] -- record-specific fields, all LE //! [u32 LE: crc32] -- CRC32 over record_type + payload //! ``` +//! +//! **NOT wired on the live recovery path.** Vector-index recovery authority +//! is the B1/B2/B3 durability layout instead: `manifest.json` + +//! `keymap-.bin` + `segment-/` dirs +//! (`crate::vector::persistence::manifest`, `segment_io`), loaded by +//! `crate::vector::persistence::recover_v2` and reconciled by a dedup rescan +//! of the live keyspace. These record types are kept only because unit tests +//! in this module exercise their (de)serialization round-trip. Do not enable +//! WAL-based vector replay without first disabling the B3 rescan — the two +//! recovery mechanisms are not designed to run together (double-apply risk). /// Tag byte distinguishing vector WAL records from RESP block frames. pub const VECTOR_RECORD_TAG: u8 = 0x56; // 'V' diff --git a/src/vector/search_pool.rs b/src/vector/search_pool.rs index 7abf9fee..e1c4083c 100644 --- a/src/vector/search_pool.rs +++ b/src/vector/search_pool.rs @@ -269,7 +269,7 @@ mod tests { } fn snapshot_for(store: &mut VectorStore, query: Vec, k: usize) -> SearchSnapshot { - let committed = store.txn_manager().committed_treemap().clone(); + let committed = store.txn_manager().committed_snapshot(); let snapshot_lsn = store.txn_manager().current_lsn(); let idx = store.get_index_mut(b"idx").unwrap(); let segments = idx.segments.load_full(); @@ -289,6 +289,7 @@ mod tests { mutable_len, scratch: SearchScratch::new(0, padded_dimension(dim)), key_hash_to_key: idx.key_hash_to_key.clone(), + ef_defaulted: false, } } @@ -344,7 +345,7 @@ mod tests { let idx = store.get_index_mut(b"idx").unwrap(); let segments = idx.segments.load_full(); let key_map = idx.key_hash_to_key.clone(); - let committed = store.txn_manager().committed_treemap().clone(); + let committed = store.txn_manager().committed_snapshot(); let snapshot_lsn = store.txn_manager().current_lsn(); let handles: Vec<_> = (0..8) @@ -372,6 +373,7 @@ mod tests { mutable_len: segments.mutable.len(), scratch: SearchScratch::new(0, padded_dimension(dim)), key_hash_to_key: key_map.clone(), + ef_defaulted: false, }; let results = futures::executor::block_on( SegmentHolder::search_mvcc_yielding_with_pool( diff --git a/src/vector/segment/compaction.rs b/src/vector/segment/compaction.rs index 23423391..41e415fd 100644 --- a/src/vector/segment/compaction.rs +++ b/src/vector/segment/compaction.rs @@ -873,11 +873,21 @@ pub fn compact( live_count, total_count, ) - .with_raw_f16(raw_f16_bfs); - - // Step 7 (continued): persist to disk if requested + .with_raw_f16(raw_f16_bfs) + // AE-1: measure this build's adaptive-ef estimate while we're on the + // compaction thread (no-op without a sidecar). + .with_adaptive_ef() + // B2 (durability): tag the segment with the disk id it will be persisted + // under (if any) BEFORE writing, so callers can read it back off the + // returned segment without threading the id separately. + .with_disk_segment_id(persist.map(|(_, segment_id)| segment_id)); + + // Step 7 (continued): persist to disk if requested. Uses the staged + // (write-to-staging + fsync + rename + dir-fsync) writer so a segment + // only becomes visible under its final `segment-{id}` name once it is + // fully durable — see VECTOR-DURABILITY-DESIGN.md. if let Some((dir, segment_id)) = persist { - segment_io::write_immutable_segment(dir, segment_id, &segment, collection) + segment_io::write_immutable_segment_staged(dir, segment_id, &segment, collection) .map_err(|e| CompactionError::PersistFailed(format!("{e}")))?; } @@ -1053,16 +1063,24 @@ pub struct MergeStats { /// /// Refuses to merge when the estimated live-vector TQ code bytes plus the /// HNSW graph overhead would exceed `MERGE_MEMORY_CEILING`. +/// +/// `persist`: when `Some((dir, segment_id))`, writes the merged segment to +/// disk (staged + atomic rename, see [`segment_io::write_immutable_segment_staged`]) +/// after a successful build, and tags the returned segment with `segment_id` +/// via [`ImmutableSegment::with_disk_segment_id`] (B2, durability). pub fn merge_immutable( segments: &[Arc], collection: &Arc, seed: u64, mode: MergeMode, recall_tolerance: f32, + persist: Option<(&Path, u64)>, ) -> Result { match mode { MergeMode::None => Err(CompactionError::EmptySegment), // caller should not call - MergeMode::GraphUnion => merge_graph_union(segments, collection, seed, recall_tolerance), + MergeMode::GraphUnion => { + merge_graph_union(segments, collection, seed, recall_tolerance, persist) + } MergeMode::KeepRaw => { // KeepRaw: fall back to graph-union if no raw f32 available (warn only). // Full raw-vector path requires raw_f32 stored on ImmutableSegment (TODO P2.5). @@ -1070,7 +1088,7 @@ pub fn merge_immutable( "MERGE_MODE=keep_raw: raw f32 not yet persisted on ImmutableSegment; \ falling back to graph-union. Add KEEP_RAW sidecar in P2.5." ); - merge_graph_union(segments, collection, seed, recall_tolerance) + merge_graph_union(segments, collection, seed, recall_tolerance, persist) } } } @@ -1084,6 +1102,7 @@ fn merge_graph_union( collection: &Arc, seed: u64, recall_tolerance: f32, + persist: Option<(&Path, u64)>, ) -> Result { if segments.is_empty() { return Err(CompactionError::EmptySegment); @@ -1235,6 +1254,24 @@ fn merge_graph_union( .collect(); entries.sort_by(|a, b| a.0.cmp(&b.0).then(a.6.cmp(&b.6))); + // R5 (persistence review): a dropped sidecar must be LOUD. The merged + // segment loses exact rerank for ALL entries when any source lacks the + // sidecar; without a warning this silent recall degradation is invisible + // (FT.INFO `segments_with_exact_rerank` exposes the steady state). + if !all_have_raw { + let with_raw = segments.iter().filter(|s| s.raw_f16().is_some()).count(); + if with_raw > 0 { + tracing::warn!( + sources = segments.len(), + sources_with_sidecar = with_raw, + "GraphUnion merge drops the exact-rerank f16 sidecar for the \ + ENTIRE merged segment (all-or-nothing propagation): at least \ + one source segment lacks it — merged-segment queries fall \ + back to quantized ADC distances (recall degrades, not breaks)" + ); + } + } + // ── Step 3: Build TQ buffer (verbatim codes, no re-encode) ─────────────── let mut tq_buffer_orig: Vec = Vec::with_capacity(n * bytes_per_code); let mut qjl_orig: Vec = Vec::with_capacity(n * qjl_bpv); @@ -1464,7 +1501,19 @@ fn merge_graph_union( n as u32, n as u32, ) - .with_raw_f16(raw_f16_bfs); + .with_raw_f16(raw_f16_bfs) + // AE-1: measure this build's adaptive-ef estimate while we're on the + // compaction thread (no-op without a sidecar). + .with_adaptive_ef() + // B2 (durability): tag with the disk id before writing, mirroring `compact()`. + .with_disk_segment_id(persist.map(|(_, segment_id)| segment_id)); + + // Persist to disk if requested — same staged (staging + fsync + rename + + // dir-fsync) writer as `compact()`. Merge persists identically to compact. + if let Some((dir, segment_id)) = persist { + segment_io::write_immutable_segment_staged(dir, segment_id, &merged, collection) + .map_err(|e| CompactionError::PersistFailed(format!("{e}")))?; + } Ok(merged) } @@ -1619,12 +1668,22 @@ fn verify_merge_recall( // HNSW search on the merged graph using f32 decoded centroids. // f32_bfs_flat has BFS-ordered vectors, `eff_dim` elements each. + // + // The query IS a database point (distance 0 → always rank 1), but the + // ground truth above excludes it. Request k+1 and drop the self-point + // so both sides compare k non-self neighbors — with a plain top-k the + // self-point consumes one slot and caps measurable recall at (k-1)/k + // = 0.90, which the manual force_merge gate (0.90) can never pass. let hnsw_results = - hnsw_search_f32(graph, &f32_bfs_flat, eff_dim, query, k, ef_verify, None); + hnsw_search_f32(graph, &f32_bfs_flat, eff_dim, query, k + 1, ef_verify, None); // hnsw_search_f32 returns original IDs (pre-BFS); convert to BFS positions // so they match the ground-truth set (which indexes all_decoded by BFS pos). - let hnsw_ids: std::collections::HashSet = - hnsw_results.iter().map(|r| graph.to_bfs(r.id.0)).collect(); + let hnsw_ids: std::collections::HashSet = hnsw_results + .iter() + .map(|r| graph.to_bfs(r.id.0)) + .filter(|&b| b != query_bfs as u32) + .take(k) + .collect(); let overlap = gt_ids.intersection(&hnsw_ids).count(); total_recall += overlap as f32 / k.min(gt_ids.len()).max(1) as f32; @@ -1837,7 +1896,7 @@ mod tests { let n = db.len(); let segs = vec![Arc::new(imm_a), Arc::new(imm_b)]; - let merged = merge_immutable(&segs, &collection, 42, MergeMode::GraphUnion, 0.60) + let merged = merge_immutable(&segs, &collection, 42, MergeMode::GraphUnion, 0.60, None) .expect("SQ8 merge failed"); assert_eq!(merged.live_count(), n as u32, "merged SQ8 live_count"); @@ -2198,7 +2257,7 @@ mod tests { ); let segs = vec![Arc::new(imm1), Arc::new(imm2)]; - let result = merge_immutable(&segs, &collection, 42, MergeMode::GraphUnion, 0.80); + let result = merge_immutable(&segs, &collection, 42, MergeMode::GraphUnion, 0.80, None); match &result { Ok(m) => eprintln!( diff --git a/src/vector/segment/holder.rs b/src/vector/segment/holder.rs index 1429fd6a..7131570f 100644 --- a/src/vector/segment/holder.rs +++ b/src/vector/segment/holder.rs @@ -24,6 +24,13 @@ use super::mutable::{MutableEntry, MutableSegment}; /// Default number of IVF clusters to probe during search. const DEFAULT_NPROBE: usize = 32; +// Per-segment beam policy (AE-1): every graph segment searches at the FULL +// resolved ef unless its compact-time saturation probe certified it as +// trivially easy (`ImmutableSegment::suggested_ef`), in which case min-ef +// suffices. A blanket `ef/√G` split was tried first (EF-SPLIT) and REJECTED: +// it silently cost R@10 0.9915 → 0.9295 on gaussian 5-seg — a per-segment +// beam reduction is only safe when the segment itself proves it saturates. + /// MVCC context for snapshot-isolated search. Passed by reference, zero allocation. pub struct MvccContext<'a> { pub snapshot_lsn: u64, @@ -32,6 +39,10 @@ pub struct MvccContext<'a> { /// Dirty set: uncommitted entries from the active transaction. pub dirty_set: &'a [MutableEntry], pub dimension: u32, + /// AE-1: true when `ef_search` came from the resolution heuristic (no + /// user `EF_RUNTIME`) — saturation-certified segments may then run at + /// their min-ef estimate. Always false when the user pinned ef. + pub ef_defaulted: bool, } /// Snapshot of all segments at a point in time. @@ -130,8 +141,11 @@ pub struct SearchSnapshot { pub snapshot_lsn: u64, /// Active txn id (0 for non-transactional reads). pub my_txn_id: u64, - /// Committed-treemap snapshot (owned) for MVCC visibility. - pub committed: roaring::RoaringTreemap, + /// Committed-treemap snapshot for MVCC visibility. `Arc` capture (QP-2): + /// an O(1) refcount bump per search; the treemap is cloned only on the + /// first capture after a commit/prune (see + /// `TransactionManager::committed_snapshot`). + pub committed: std::sync::Arc, /// Vector dimension. pub dimension: u32, /// Entry count of the mutable segment captured at entry. The append-only @@ -145,6 +159,10 @@ pub struct SearchSnapshot { /// builder, not the segment scan. `Arc` snapshot: capture is O(1); writers /// copy-on-write via `Arc::make_mut` (QP-1). pub key_hash_to_key: std::sync::Arc>, + /// AE-1: true when `ef_search` came from the resolution heuristic (no + /// user `EF_RUNTIME`) — saturation-certified segments may then run at + /// their min-ef estimate. + pub ef_defaulted: bool, } /// Lock-free segment holder. Searches load() once at query start and hold @@ -281,6 +299,10 @@ impl SegmentHolder { None // Light mode: no QJL matrices, use TQ-ADC brute force }; + // Full resolved ef per segment (no ef_defaulted context on this path, + // so the AE-1 per-segment reduction never applies here). + let graph_ef = ef_search; + match strategy { FilterStrategy::Unfiltered => { all.extend( @@ -289,10 +311,10 @@ impl SegmentHolder { .brute_force_search(query_f32, query_state.as_ref(), k), ); for imm in &snapshot.immutable { - all.extend(imm.search(query_f32, k, ef_search, _scratch)); + all.extend(imm.search(query_f32, k, graph_ef, _scratch)); } for warm_seg in &snapshot.warm { - all.extend(warm_seg.search(query_f32, k, ef_search, _scratch)); + all.extend(warm_seg.search(query_f32, k, graph_ef, _scratch)); } } FilterStrategy::BruteForceFiltered => { @@ -306,7 +328,7 @@ impl SegmentHolder { all.extend(imm.search_filtered( query_f32, k, - ef_search, + graph_ef, _scratch, filter_bitmap, )); @@ -315,7 +337,7 @@ impl SegmentHolder { all.extend(warm_seg.search_filtered( query_f32, k, - ef_search, + graph_ef, _scratch, filter_bitmap, )); @@ -332,7 +354,7 @@ impl SegmentHolder { all.extend(imm.search_filtered( query_f32, k, - ef_search, + graph_ef, _scratch, filter_bitmap, )); @@ -341,7 +363,7 @@ impl SegmentHolder { all.extend(warm_seg.search_filtered( query_f32, k, - ef_search, + graph_ef, _scratch, filter_bitmap, )); @@ -355,13 +377,9 @@ impl SegmentHolder { oversample_k, filter_bitmap, )); + let post_ef = ef_search.max(oversample_k); for imm in &snapshot.immutable { - let imm_results = imm.search( - query_f32, - oversample_k, - ef_search.max(oversample_k), - _scratch, - ); + let imm_results = imm.search(query_f32, oversample_k, post_ef, _scratch); if let Some(bm) = filter_bitmap { for r in imm_results { if bm.contains(r.id.0) { @@ -373,12 +391,7 @@ impl SegmentHolder { } } for warm_seg in &snapshot.warm { - let warm_results = warm_seg.search( - query_f32, - oversample_k, - ef_search.max(oversample_k), - _scratch, - ); + let warm_results = warm_seg.search(query_f32, oversample_k, post_ef, _scratch); if let Some(bm) = filter_bitmap { for r in warm_results { if bm.contains(r.id.0) { @@ -503,11 +516,23 @@ impl SegmentHolder { // 2. HNSW search on immutable segments (TQ-ADC distance). // Immutable segment entries are committed by definition (compacted only // after commit). No visibility post-filter needed for Phase 65. + // AE-1: heuristic-defaulted ef → a saturation-certified segment runs + // at its compact-time min-ef estimate (floored by k, capped by the + // resolved ef); every other segment gets the FULL resolved ef. + // Mirrors the yielding path's selector exactly. + let graph_ef = ef_search; + let seg_ef = |suggested: Option| -> usize { + match suggested { + Some(sug) if mvcc.ef_defaulted => (sug as usize).max(k).min(ef_search), + _ => graph_ef, + } + }; for imm in &snapshot.immutable { + let ef_i = seg_ef(imm.suggested_ef()); if filter_bitmap.is_some() { - all.extend(imm.search_filtered(query_f32, k, ef_search, _scratch, filter_bitmap)); + all.extend(imm.search_filtered(query_f32, k, ef_i, _scratch, filter_bitmap)); } else { - all.extend(imm.search(query_f32, k, ef_search, _scratch)); + all.extend(imm.search(query_f32, k, ef_i, _scratch)); } } @@ -517,12 +542,12 @@ impl SegmentHolder { all.extend(warm_seg.search_filtered( query_f32, k, - ef_search, + graph_ef, _scratch, filter_bitmap, )); } else { - all.extend(warm_seg.search(query_f32, k, ef_search, _scratch)); + all.extend(warm_seg.search(query_f32, k, graph_ef, _scratch)); } } @@ -647,6 +672,20 @@ impl SegmentHolder { } else { ef_search }; + // AE-1: when ef was heuristic-defaulted, a saturation-certified + // segment runs at its compact-time min-ef estimate — floored by the + // merge quota, capped by the resolved ef. Every other segment gets + // the FULL resolved ef (see the per-segment beam policy note at the + // top of this file). Identical selector in the pooled jobs, inline + // fallbacks, and serial loops, so pooled == serial identity holds. + // Never active when the user pinned EF_RUNTIME. + let ef_defaulted = snap.ef_defaulted; + let seg_ef = |suggested: Option| -> usize { + match suggested { + Some(sug) if ef_defaulted => (sug as usize).max(fetch_k).min(graph_ef), + _ => graph_ef, + } + }; let snapshot_lsn = snap.snapshot_lsn; let my_txn_id = snap.my_txn_id; let mutable_len = snap.mutable_len; @@ -687,13 +726,14 @@ impl SegmentHolder { // On submit failure (pool shut down at process teardown) the // segment is searched inline — the query still answers correctly. for seg in &segments.immutable { + let ef_seg = seg_ef(seg.suggested_ef()); let job = crate::vector::search_pool::SegmentSearchJob { segment: crate::vector::search_pool::GraphSegmentRef::Immutable( std::sync::Arc::clone(seg), ), query: std::sync::Arc::clone(&query_arc), fetch_k, - ef_search: graph_ef, + ef_search: ef_seg, filter: filter_arc.clone(), reply: tx.clone(), }; @@ -703,12 +743,12 @@ impl SegmentHolder { all.extend(seg.search_filtered( query_f32, fetch_k, - graph_ef, + ef_seg, &mut snap.scratch, graph_filter, )); } else { - let results = seg.search(query_f32, fetch_k, graph_ef, &mut snap.scratch); + let results = seg.search(query_f32, fetch_k, ef_seg, &mut snap.scratch); if post_filter { if let Some(bm) = filter_ref { all.extend(results.into_iter().filter(|r| bm.contains(r.id.0))); @@ -831,16 +871,17 @@ impl SegmentHolder { // otherwise the ACORN-filtered traversal, same as the sync path. if !pooled_graph { for imm in &segments.immutable { + let ef_seg = seg_ef(imm.suggested_ef()); if graph_filter.is_some() { all.extend(imm.search_filtered( query_f32, fetch_k, - graph_ef, + ef_seg, &mut snap.scratch, graph_filter, )); } else { - let results = imm.search(query_f32, fetch_k, graph_ef, &mut snap.scratch); + let results = imm.search(query_f32, fetch_k, ef_seg, &mut snap.scratch); if post_filter { if let Some(bm) = filter_ref { all.extend(results.into_iter().filter(|r| bm.contains(r.id.0))); @@ -949,6 +990,12 @@ impl SegmentHolder { // 4. Merge all results, take global top-k (identical to search_mvcc). all.sort_unstable(); all.truncate(k); + // QP-3: hand this query's scratch back to the thread cache; the next + // capture on this thread reuses it via take_thread_scratch. + crate::vector::hnsw::search::recycle_thread_scratch(std::mem::replace( + &mut snap.scratch, + crate::vector::hnsw::search::SearchScratch::new(0, 0), + )); all } } @@ -1132,6 +1179,7 @@ mod tests { committed: &committed, dirty_set: &[], dimension: dim as u32, + ef_defaulted: false, }; let mvcc = holder.search_mvcc(&query_f32, 3, 64, &mut scratch, None, &mvcc_ctx); @@ -1165,6 +1213,7 @@ mod tests { committed: &committed, dirty_set: &[], dimension: dim as u32, + ef_defaulted: false, }; let results = holder.search_mvcc(&query_f32, 3, 64, &mut scratch, None, &mvcc_ctx); assert_eq!(results.len(), 1); @@ -1221,6 +1270,7 @@ mod tests { committed: &committed, dirty_set: std::slice::from_ref(&dirty_entry), dimension: dim as u32, + ef_defaulted: false, }; let results = holder.search_mvcc(&query_f32, 3, 64, &mut scratch, None, &mvcc_ctx); @@ -1259,6 +1309,7 @@ mod tests { committed: &committed, dirty_set: &[], dimension: dim as u32, + ef_defaulted: false, }; let r1 = holder.search_mvcc(&query_f32, 3, 64, &mut scratch, None, &mvcc_empty); @@ -1269,6 +1320,7 @@ mod tests { committed: &committed, dirty_set: &[], dimension: dim as u32, + ef_defaulted: false, }; let r2 = holder.search_mvcc(&query_f32, 3, 64, &mut scratch, None, &mvcc_empty2); diff --git a/src/vector/segment/immutable.rs b/src/vector/segment/immutable.rs index 4a28fcec..b1a0b738 100644 --- a/src/vector/segment/immutable.rs +++ b/src/vector/segment/immutable.rs @@ -93,6 +93,18 @@ pub struct ImmutableSegment { /// values to f16 tolerance instead of quantized ADC estimates. `None` for /// segments built without raw vectors (pre-sidecar disk segments). raw_f16: Option>, + + /// Compact-time adaptive-ef estimate (AE-1): the smallest ladder ef at + /// which this segment's OWN sampled queries reach the target recall + /// against exact sidecar ground truth. Applied per-segment when the + /// query's ef was heuristic-defaulted (never when the user pinned + /// `EF_RUNTIME`). In-memory only — segments reloaded from disk carry + /// `None` and fall back to the full resolved beam. + suggested_ef: Option, + + /// Disk segment id this segment was persisted under (B2, durability write + /// path). See [`Self::disk_segment_id`]. + disk_segment_id: Option, } impl ImmutableSegment { @@ -126,6 +138,8 @@ impl ImmutableSegment { has_tombstones: AtomicBool::new(false), tombstoned_keys: parking_lot::RwLock::new(HashSet::new()), raw_f16: None, + suggested_ef: None, + disk_segment_id: None, } } @@ -152,6 +166,215 @@ impl ImmutableSegment { self.raw_f16.as_deref() } + /// The compact-time adaptive-ef estimate for this segment (AE-1), if one + /// was measured. See [`Self::with_adaptive_ef`]. + pub fn suggested_ef(&self) -> Option { + self.suggested_ef + } + + /// Measure and attach the adaptive-ef estimate (AE-1). Builder-style, + /// called at compact/merge time after the graph and sidecar are in place. + /// No-op (stays `None`) without a sidecar or for tiny segments. + #[must_use] + pub fn with_adaptive_ef(mut self) -> Self { + self.suggested_ef = self.estimate_suggested_ef(); + self + } + + /// Restore a previously-measured adaptive-ef estimate (R6, durability) + /// verbatim, WITHOUT re-running [`Self::estimate_suggested_ef`]. Used by + /// `segment_io::read_immutable_segment` when loading a segment that + /// carries a persisted `suggested_ef` in its `segment_meta.json` — the + /// estimator is compaction-thread-only and must not re-run on the + /// (potentially hot) load path. + #[must_use] + pub fn with_suggested_ef_raw(mut self, suggested_ef: Option) -> Self { + self.suggested_ef = suggested_ef; + self + } + + /// The disk segment id this in-memory segment was persisted under, if + /// any (B2, durability write path). `None` when the index has no + /// `persist_dir` configured, or for segments loaded before this field + /// existed. Used to build `IndexManifest.segment_ids` — the set of + /// on-disk segment directories currently referenced by the live index. + pub fn disk_segment_id(&self) -> Option { + self.disk_segment_id + } + + /// Attach the disk segment id assigned at persist time. Builder-style, + /// mirrors [`Self::with_raw_f16`]. `None` when the segment was not + /// persisted (no `persist_dir` configured for the index). + #[must_use] + pub fn with_disk_segment_id(mut self, disk_segment_id: Option) -> Self { + self.disk_segment_id = disk_segment_id; + self + } + + /// AE-1 estimator: sample `ADAPTIVE_EF_SAMPLES` of the segment's own + /// vectors as queries, compute exact top-k ground truth from the f16 + /// sidecar (SIMD kernels; same distance conventions as `rerank_exact`), + /// then measure R@k along the ef ladder. Returns `Some(min-ef)` ONLY + /// when the curve is fully saturated (see `ADAPTIVE_EF_SATURATION`); + /// `None` when no sidecar, the segment is too small to measure + /// meaningfully, or the curve shows any knee (caller falls back to the + /// full resolved beam). + /// + /// Cost: samples·n·dim SIMD f16 ops for ground truth (~tens of ms for a + /// bounded compact build) + samples·|ladder| graph searches — one-time, + /// on the compaction/merge thread, never on the query path. + fn estimate_suggested_ef(&self) -> Option { + const ADAPTIVE_EF_SAMPLES: usize = 16; + const ADAPTIVE_EF_K: usize = 10; + /// Acceptance requires a FULLY SATURATED ladder: the first (minimum) + /// rung must already sit within ε of the ladder top AND at/above the + /// saturation bar. Self-sampled queries are optimistically biased on + /// unstructured data — an absolute 0.98 gate picked catastrophically + /// low efs (measured R@10 0.9915 → 0.573 on gaussian 5-seg), and even + /// a knee-of-own-curve relative gate still over-trusted it (0.9915 → + /// 0.939): a segment whose self-curve shows ANY visible knee has its + /// true external-query knee further right, beyond what the estimator + /// can see. The only decision the estimator makes reliably is + /// "trivially easy segment" (clustered data measures flat 1.0 across + /// the whole ladder) — for those, min-ef is safe; everything else + /// falls back to the full resolved beam. + const ADAPTIVE_EF_SATURATION: f32 = 0.995; + const ADAPTIVE_EF_EPSILON: f32 = 0.005; + const ADAPTIVE_EF_LADDER: &[usize] = &[24, 32, 48, 64, 96, 128, 192, 256]; + + let raw = self.raw_f16.as_deref()?; + let dim = self.collection_meta.dimension as usize; + let n = self.mvcc.len(); + if dim == 0 || n < ADAPTIVE_EF_K * 8 || raw.len() < n * dim { + return None; + } + let is_l2 = self.collection_meta.metric == crate::vector::types::DistanceMetric::L2; + let table = crate::vector::distance::table(); + + // Evenly-spaced sample rows as queries (decoded to f32; the raw + // vector feeds `search()` which does its own normalization). + let samples = ADAPTIVE_EF_SAMPLES.min(n / 4).max(4); + let step = (n / samples).max(1); + let mut queries: Vec<(usize, Vec)> = Vec::with_capacity(samples); + for s in 0..samples { + let bfs = s * step; + if bfs >= n { + break; + } + let row = &raw[bfs * dim..(bfs + 1) * dim]; + let q: Vec = row + .iter() + .map(|&h| crate::vector::f16::f16_to_f32(h)) + .collect(); + queries.push((bfs, q)); + } + + // Exact ground truth per query: brute force over the sidecar with the + // rerank distance conventions (L2: squared L2; unit-sphere metrics: + // 2 − 2·⟨q̂,x⟩/‖x‖). Collect top-k GLOBAL ids. + let mut gts: Vec> = Vec::with_capacity(queries.len()); + for (q_bfs, q) in &queries { + let q_bfs = *q_bfs; + let mut q_unit; + let q_ref: &[f32] = if is_l2 { + q + } else { + let norm: f32 = q.iter().map(|x| x * x).sum::().sqrt(); + if norm <= 0.0 { + gts.push(SmallVec::new()); + continue; + } + let inv = 1.0 / norm; + q_unit = q.clone(); + for v in q_unit.iter_mut() { + *v *= inv; + } + &q_unit + }; + // (distance, bfs) top-k via sorted Vec. Leave-self-out: the + // query IS row `q_bfs`, and its trivial self-hit would inflate + // measured recall — exclude it from ground truth and results. + let mut top: Vec<(f32, usize)> = Vec::with_capacity(ADAPTIVE_EF_K + 1); + for bfs in 0..n { + if bfs == q_bfs || !self.is_live_bfs(bfs as u32) { + continue; + } + let x = &raw[bfs * dim..(bfs + 1) * dim]; + let d = if is_l2 { + (table.f16_l2)(q_ref, x) + } else { + let (dot, xsq) = (table.f16_dot_normsq)(q_ref, x); + if xsq > 0.0 { + 2.0 - 2.0 * (dot / xsq.sqrt()) + } else { + continue; + } + }; + if top.len() < ADAPTIVE_EF_K { + top.push((d, bfs)); + top.sort_unstable_by(|a, b| a.0.total_cmp(&b.0)); + } else if d < top[ADAPTIVE_EF_K - 1].0 { + top[ADAPTIVE_EF_K - 1] = (d, bfs); + top.sort_unstable_by(|a, b| a.0.total_cmp(&b.0)); + } + } + gts.push( + top.iter() + .map(|&(_, bfs)| self.mvcc[bfs].global_id) + .collect(), + ); + } + + // Ladder walk. Fetch k+1 and drop the self-hit (leave-self-out), + // then accept the smallest ef whose recall is BOTH ≥ the absolute + // target AND within ε of the ladder-top recall (the segment's own + // asymptote). A segment whose curve never reaches the target yields + // None (full-resolved-beam fallback). + let mut scratch = SearchScratch::new(0, self.collection_meta.padded_dimension); + let mut recall_at = [0f32; 8]; + debug_assert_eq!(ADAPTIVE_EF_LADDER.len(), recall_at.len()); + for (li, &ef) in ADAPTIVE_EF_LADDER.iter().enumerate() { + let mut hit = 0usize; + let mut total = 0usize; + for ((q_bfs, q), gt) in queries.iter().zip(>s) { + if gt.is_empty() { + continue; + } + let self_gid = self.mvcc[*q_bfs].global_id; + let res = self.search(q, ADAPTIVE_EF_K + 1, ef, &mut scratch); + hit += res + .iter() + .filter(|r| r.id.0 != self_gid && gt.contains(&r.id.0)) + .count(); + total += gt.len(); + } + recall_at[li] = if total > 0 { + (hit as f32) / (total as f32) + } else { + 0.0 + }; + } + let top = recall_at[ADAPTIVE_EF_LADDER.len() - 1]; + let floor = recall_at[0]; + // Saturated-curve gate: only a segment whose MINIMUM-ef self-recall + // already matches its asymptote (and clears the saturation bar) gets + // a suggestion — and that suggestion is min-ef itself. Any visible + // knee in the self-curve → None (see ADAPTIVE_EF_SATURATION doc). + let decision = if floor >= ADAPTIVE_EF_SATURATION && floor >= top - ADAPTIVE_EF_EPSILON { + Some(ADAPTIVE_EF_LADDER[0] as u32) + } else { + None + }; + tracing::debug!( + n, + samples = queries.len(), + ?recall_at, + ?decision, + "AE-1 adaptive-ef ladder" + ); + decision + } + /// Exact rerank (HQ-1): re-score `candidates` with (near-)exact distances /// decoded from the f16 sidecar, replacing their quantized ADC estimates, /// then re-sort ascending. No-op when the segment has no sidecar. @@ -196,6 +419,10 @@ impl ImmutableSegment { &q_unit }; + // SIMD-dispatched sidecar kernels (NEON / F16C, scalar fallback) — + // this loop was 17% of a matched-recall query when the decode was + // scalar software f16_to_f32. + let dist_table = crate::vector::distance::table(); for result in candidates[..rerank_n].iter_mut() { let bfs_pos = self.graph.to_bfs(result.id.0) as usize; let start = bfs_pos * dim; @@ -203,16 +430,10 @@ impl ImmutableSegment { continue; // Out-of-range id: keep the ADC estimate. }; if is_l2 { - result.distance = crate::vector::f16::l2_sq_f16(q_ref, vec_f16); + result.distance = (dist_table.f16_l2)(q_ref, vec_f16); } else { // One pass: ⟨q̂,x⟩ and ‖x‖² from the f16-decoded vector. - let mut dot = 0.0f32; - let mut xsq = 0.0f32; - for (q, &h) in q_ref.iter().zip(vec_f16.iter()) { - let x = crate::vector::f16::f16_to_f32(h); - dot += q * x; - xsq += x * x; - } + let (dot, xsq) = (dist_table.f16_dot_normsq)(q_ref, vec_f16); if xsq > 0.0 { // f16 rounding can push cos slightly outside [-1, 1]; // clamp so distances stay in the metric's [0, 4] range. @@ -327,9 +548,12 @@ impl ImmutableSegment { let bfs = self.graph.to_bfs(c.id.0); self.is_live_bfs(bfs) }); - // HQ-1: exact rerank of the live beam (ef candidates) from the f16 + // HQ-1: exact rerank of the top 4·k live beam candidates from the f16 // sidecar — replaces quantized estimates with true metric distances - // before top-k truncation. No-op without a sidecar. + // before top-k truncation. No-op without a sidecar. Runs for SQ8 + // too: an A/B (20k×384d, 1seg+5seg) measured skipping it costs + // R@10 −0.010 clustered / −0.017 gaussian-5seg — real recall, and + // the SIMD sidecar kernels make the pass ~3% of a query. self.rerank_exact(&mut candidates, query, k); candidates.truncate(k); self.remap_to_global_ids(&mut candidates); @@ -371,7 +595,8 @@ impl ImmutableSegment { let bfs = self.graph.to_bfs(c.id.0); self.is_live_bfs(bfs) }); - // HQ-1: exact rerank of the live beam from the f16 sidecar. + // HQ-1: exact rerank of the live beam from the f16 sidecar (see the + // recall A/B note in `search`). self.rerank_exact(&mut candidates, query, k); candidates.truncate(k); self.remap_to_global_ids(&mut candidates); diff --git a/src/vector/segment/mutable.rs b/src/vector/segment/mutable.rs index d38342c4..5c3d87bd 100644 --- a/src/vector/segment/mutable.rs +++ b/src/vector/segment/mutable.rs @@ -19,7 +19,8 @@ use crate::vector::turbo_quant::encoder::{ }; use crate::vector::turbo_quant::fwht; use crate::vector::turbo_quant::sq8::{ - SQ8_PARAMS_BYTES, encode_sq8_into, sq8_l2_from_stats, sq8_params, sq8_query_stats, + SQ8_INT8_QMAX, SQ8_PARAMS_BYTES, encode_sq8_into, sq8_int8_dot_to_f32, sq8_l2_from_stats, + sq8_params, sq8_quantize_query_scalar, sq8_query_stats, }; use crate::vector::turbo_quant::tq_adc::tq_l2_adc_scaled; use crate::vector::types::{DistanceMetric, SearchResult, VectorId}; @@ -82,6 +83,13 @@ pub struct BruteForceQuery { /// candidate via `sq8_l2_from_stats`. Zero for non-SQ8 collections. sq8_q_sum: f32, sq8_q_sumsq: f32, + /// Task #13: int8-quantized copy of `prepared`, populated only when a + /// SIMD int8 ADC kernel is installed for this CPU/build (empty + /// otherwise — `brute_force_scan_mvcc_chunk` falls back to the f32 + /// `sq8_stats` path when `sq8_i8_stats` dispatch is `None`). + sq8_qi8: Vec, + sq8_q_scale: f32, + sq8_sum_qi8: i32, /// Shared top-k accumulator across all chunks. heap: BinaryHeap, } @@ -434,6 +442,18 @@ impl MutableSegment { // loop below — never per candidate. let sq8_stats_fn = distance::table().sq8_stats; let (q_sum, q_sumsq) = sq8_query_stats(q); + // Task #13: int8 symmetric ADC — quantize once per search when a + // SIMD int8 kernel is installed for this CPU/build; else fall + // back to the f32 `sq8_stats_fn` path above. + let sq8_i8_stats_fn = distance::table().sq8_i8_stats; + let (sq8_qi8, sq8_q_scale, sq8_sum_qi8): (Vec, f32, i32) = + if sq8_i8_stats_fn.is_some() { + let mut qi8 = vec![0i8; dim]; + let (scale, sum) = sq8_quantize_query_scalar(q, SQ8_INT8_QMAX, &mut qi8); + (qi8, scale, sum) + } else { + (Vec::new(), 0.0, 0) + }; let mut heap: BinaryHeap = BinaryHeap::with_capacity(k + 1); for entry in &inner.entries { if entry.delete_lsn != 0 { @@ -449,7 +469,17 @@ impl MutableSegment { let off = id * bytes_per_code; let slot = &inner.tq_codes[off..off + bytes_per_code]; let (min, scale) = sq8_params(slot, dim); - let (dot_qc, sum_c, sumsq_c) = sq8_stats_fn(q, &slot[..dim]); + let (dot_qc, sum_c, sumsq_c) = if let Some(i8_stats) = sq8_i8_stats_fn { + let (dot_int, sum_c_int, sumsq_c_int) = + i8_stats(&sq8_qi8, &slot[..dim], sq8_sum_qi8); + ( + sq8_int8_dot_to_f32(sq8_q_scale, dot_int), + sum_c_int as f32, + sumsq_c_int as f32, + ) + } else { + sq8_stats_fn(q, &slot[..dim]) + }; let dist = sq8_l2_from_stats(dim, min, scale, q_sum, q_sumsq, dot_qc, sum_c, sumsq_c); let global_id = inner.global_id_base + entry.internal_id; @@ -697,11 +727,28 @@ impl MutableSegment { } else { (0.0, 0.0) }; + // Task #13: int8 symmetric ADC — quantize once here (persists across + // chunks via `BruteForceQuery`, mirroring `sq8_q_sum`/`sq8_q_sumsq` + // above) only when a SIMD int8 kernel is installed for this + // CPU/build; empty otherwise (chunk scan falls back to f32 stats). + let (sq8_qi8, sq8_q_scale, sq8_sum_qi8): (Vec, f32, i32) = + if self.collection.quantization == QuantizationConfig::Sq8 + && distance::table().sq8_i8_stats.is_some() + { + let mut qi8 = vec![0i8; prepared.len()]; + let (scale, sum) = sq8_quantize_query_scalar(&prepared, SQ8_INT8_QMAX, &mut qi8); + (qi8, scale, sum) + } else { + (Vec::new(), 0.0, 0) + }; BruteForceQuery { prepared, use_tq_adc, sq8_q_sum, sq8_q_sumsq, + sq8_qi8, + sq8_q_scale, + sq8_sum_qi8, heap: BinaryHeap::with_capacity(k + 1), } } @@ -736,6 +783,13 @@ impl MutableSegment { // pointer read), per-query constants carried in `BruteForceQuery`. let sq8_stats_fn = distance::table().sq8_stats; let (q_sum, q_sumsq) = (q.sq8_q_sum, q.sq8_q_sumsq); + // Task #13: int8 symmetric ADC — resolved once per chunk like + // `sq8_stats_fn`; `sq8_i8_stats_fn` is `Some` only when + // `q.sq8_qi8` was actually populated in `prepare_brute_force_query`. + let sq8_i8_stats_fn = distance::table().sq8_i8_stats; + let sq8_qi8 = q.sq8_qi8.as_slice(); + let sq8_q_scale = q.sq8_q_scale; + let sq8_sum_qi8 = q.sq8_sum_qi8; let heap = &mut q.heap; for entry in &inner.entries[lo..hi] { if !is_visible( @@ -758,7 +812,17 @@ impl MutableSegment { let off = id * bytes_per_code; let slot = &inner.tq_codes[off..off + bytes_per_code]; let (min, scale) = sq8_params(slot, dim); - let (dot_qc, sum_c, sumsq_c) = sq8_stats_fn(q_slice, &slot[..dim]); + let (dot_qc, sum_c, sumsq_c) = if let Some(i8_stats) = sq8_i8_stats_fn { + let (dot_int, sum_c_int, sumsq_c_int) = + i8_stats(sq8_qi8, &slot[..dim], sq8_sum_qi8); + ( + sq8_int8_dot_to_f32(sq8_q_scale, dot_int), + sum_c_int as f32, + sumsq_c_int as f32, + ) + } else { + sq8_stats_fn(q_slice, &slot[..dim]) + }; let dist = sq8_l2_from_stats(dim, min, scale, q_sum, q_sumsq, dot_qc, sum_c, sumsq_c); let global_id = inner.global_id_base + entry.internal_id; diff --git a/src/vector/store.rs b/src/vector/store.rs index cb3d5dc6..e671ea56 100644 --- a/src/vector/store.rs +++ b/src/vector/store.rs @@ -3,6 +3,7 @@ //! No Arc, no Mutex -- fully owned by shard thread (same pattern as PubSubRegistry). use std::collections::HashMap; +use std::path::{Path, PathBuf}; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; @@ -200,7 +201,47 @@ pub struct VectorIndex { /// When `HSET doc:1 category "science"` is called without a vector blob, /// the auto-indexer looks up the existing `global_id` here to update the /// PayloadIndex for that vector without re-inserting it. - pub key_hash_to_global_id: std::collections::HashMap, + /// + /// `Arc`-wrapped for the same reason as `key_hash_to_key` (QP-1 O(1) + /// snapshot capture) — the durability write path (B2) clones this Arc + /// into background keymap-snapshot jobs without an O(n) copy on the + /// shard thread. + pub key_hash_to_global_id: Arc>, + /// Maps `key_hash` → `xxh64(vector-field bytes, seed 0)` (B2, durability). + /// + /// Computed where the raw HSET vector blob is in hand + /// (`spsc_handler::handle_vector_insert`) and mirrored at every site that + /// mirrors `key_hash_to_key` (insert, update, tombstone, wholesale + /// reset). Persisted into `keymap-.bin` so the B3 recovery dedup + /// rescan can tell "unchanged" keys (checksum matches → metadata-only + /// rebuild) from "changed" keys (re-encode) without hashing every value + /// twice across a restart. + pub key_hash_to_vec_checksum: Arc>, + /// Shard-relative directory this index persists segments/manifest/keymap + /// under, i.e. `/idx-/` (B2). + /// `None` when the store has no `persist_dir` configured (disk-offload + /// and on-disk persistence both disabled) — segment/manifest/keymap + /// writes are skipped entirely in that case, identical to the existing + /// sidecar gate. + pub persist_dir: Option, + /// Monotonic on-disk segment id allocator for THIS index, scoped to + /// `persist_dir` (B2). Allocated on the shard thread at compact/merge + /// submit time (never on the worker thread) so ids never collide even + /// though the actual disk write happens on a background worker. + pub next_segment_id: u64, + /// Monotonic manifest/keymap "generation" counter for THIS index (B2). + /// Allocated on the shard thread each time a snapshot job is enqueued; + /// doubles as the keymap epoch (`keymap-.bin`) — one counter, so + /// there is no separate race between "job ordering" and "which keymap + /// file is current". + next_snapshot_seq: u64, + /// Shared watermark of the highest snapshot `seq` durably committed for + /// this index (B2). Cloned into every snapshot job; the worker holds + /// this lock across its entire write+GC critical section so a stale job + /// (built from an older install, possibly completed out of order by a + /// different worker) can never overwrite a newer manifest — see + /// `crate::vector::persistence::manifest` module docs. + persist_seq_watermark: Arc>, /// Whether auto-compaction is enabled. Default: true. /// Set to false via FT.CONFIG SET idx AUTOCOMPACT OFF for bulk ingestion. /// Manual FT.COMPACT always works regardless of this flag. @@ -357,7 +398,16 @@ impl VectorIndex { let fs_len = fs.segments.load().mutable.len(); if fs_len >= threshold { let dim = fs.collection.dimension; - Self::compact_segments(&mut fs.segments, &mut fs.scratch, &fs.collection, dim, 0); + let mut unused_id = 0u64; + Self::compact_segments( + &mut fs.segments, + &mut fs.scratch, + &fs.collection, + dim, + 0, + None, // additional-field segments are not persisted (out of B2 scope) + &mut unused_id, + ); } } } @@ -405,16 +455,24 @@ impl VectorIndex { }; drop(snap); self.segments.swap(new_list); + // B2 (durability): the worker already persisted this segment + // (begin_background_compact allocated persist target at + // submit time, if configured) — commit the keymap/manifest + // snapshot now that install is durable in memory. + self.persist_hook_after_install(); // After draining we're done — the data is compacted. // Compact additional fields inline (no in-flight for those). for (_, fs) in &mut self.field_segments { let dim = fs.collection.dimension; + let mut unused_id = 0u64; Self::compact_segments( &mut fs.segments, &mut fs.scratch, &fs.collection, dim, 0, + None, + &mut unused_id, ); } return; @@ -422,18 +480,35 @@ impl VectorIndex { // Worker failed or dropped — fall through to inline compact below. } - // Compact default field + // Compact default field. `persist_root` must be the per-INDEX dir + // (`idx-/`), not the bare shard-level `persist_dir` — matching + // `alloc_persist_target`/`persist_hook_after_install`. + let persist_root = self.persist_dir.as_ref().map(|dir| { + crate::vector::persistence::manifest::index_persist_dir(dir, self.meta.name.as_ref()) + }); Self::compact_segments( &mut self.segments, &mut self.scratch, &self.collection, self.meta.dimension, self.meta.compact_threshold as usize, + persist_root.as_deref(), + &mut self.next_segment_id, ); + self.persist_hook_after_install(); // Compact additional fields (legacy unbounded semantics: threshold 0). for (_, fs) in &mut self.field_segments { let dim = fs.collection.dimension; - Self::compact_segments(&mut fs.segments, &mut fs.scratch, &fs.collection, dim, 0); + let mut unused_id = 0u64; + Self::compact_segments( + &mut fs.segments, + &mut fs.scratch, + &fs.collection, + dim, + 0, + None, + &mut unused_id, + ); } } @@ -444,12 +519,24 @@ impl VectorIndex { /// independently searchable segments (see `bulk_freeze_cap`). The tail /// survives each install via `clone_suffix`, preserving the global ID /// space, exactly like the background install path. + /// + /// `persist_root`/`next_segment_id`: when `persist_root` is `Some`, each + /// successful build in the (possibly multi-iteration, bulk-load) loop is + /// persisted to disk via the staged writer under a freshly-allocated id + /// (B2, durability write path). `next_segment_id` is always threaded + /// through (even when `persist_root` is `None`, in which case it is + /// simply never read) so callers that don't persist can pass a throwaway + /// counter. Callers are responsible for triggering the keymap/manifest + /// snapshot job (see `VectorIndex::persist_hook_after_install`) once this + /// returns — this function only handles the segment file itself. fn compact_segments( segments: &mut SegmentHolder, scratch: &mut SearchScratch, collection: &Arc, dimension: u32, compact_threshold: usize, + persist_root: Option<&Path>, + next_segment_id: &mut u64, ) { let _ = dimension; let seed = collection.collection_id.wrapping_mul(6364136223846793005); @@ -463,7 +550,13 @@ impl VectorIndex { let frozen = snap.mutable.freeze_prefix(frozen_len); drop(snap); - match compaction::compact(&frozen, collection, seed, None) { + let persist = persist_root.map(|dir| { + let id = *next_segment_id; + *next_segment_id += 1; + (dir, id) + }); + + match compaction::compact(&frozen, collection, seed, persist) { Ok(immutable) => { let num_nodes = immutable.graph().num_nodes(); let padded = collection.padded_dimension; @@ -493,6 +586,88 @@ impl VectorIndex { } } } + + /// One past the highest collection id used by any field of this index + /// (default field + additional named vector fields). Used as + /// `IndexManifest.next_collection_id_floor` — B3 recovery seeds the + /// store-wide collection-id allocator from this so a later FT.CREATE + /// never reuses a collection id (and therefore QJL rotation seed) that a + /// restored segment already depends on. + fn max_field_collection_id(&self) -> u64 { + let mut max_cid = self.collection.collection_id; + for fs in self.field_segments.values() { + max_cid = max_cid.max(fs.collection.collection_id); + } + max_cid + } + + /// Allocate the next on-disk segment id and target directory for this + /// index's default field, if `persist_dir` is configured. + /// + /// Always called on the shard thread (at compact/merge SUBMIT time, or + /// synchronously for the inline compact path) — segment ids are never + /// allocated on a background worker, so two concurrently-running builds + /// can never collide on an id even though the actual disk write happens + /// off-thread. + fn alloc_persist_target(&mut self) -> Option<(PathBuf, u64)> { + let dir = self.persist_dir.as_ref()?; + let idx_dir = + crate::vector::persistence::manifest::index_persist_dir(dir, self.meta.name.as_ref()); + let id = self.next_segment_id; + self.next_segment_id += 1; + Some((idx_dir, id)) + } + + /// After a successful default-field compact/merge install, commit a + /// keymap + manifest snapshot for this index in the background (B2, + /// durability write path). No-op when `persist_dir` is not configured. + /// + /// Cheap on the shard thread: clones 3 `Arc`s (O(1) — all three key-hash + /// maps are `Arc`-wrapped with copy-on-write writers) plus a handful of + /// scalars; the actual keymap build (iterating every live key) and all + /// file I/O happens on the snapshot worker thread. + fn persist_hook_after_install(&mut self) { + let Some(dir) = self.persist_dir.clone() else { + return; + }; + let idx_dir = + crate::vector::persistence::manifest::index_persist_dir(&dir, self.meta.name.as_ref()); + let snap = self.segments.load(); + let segment_ids: Vec = snap + .immutable + .iter() + .filter_map(|s| s.disk_segment_id()) + .collect(); + let next_global_id = snap.mutable.next_global_id(); + drop(snap); + + self.next_snapshot_seq += 1; + let seq = self.next_snapshot_seq; + + let manifest = crate::vector::persistence::manifest::IndexManifest { + format_version: crate::vector::persistence::manifest::MANIFEST_FORMAT_VERSION, + index_name_hex: crate::vector::persistence::manifest::index_name_hex( + self.meta.name.as_ref(), + ), + collection_id: self.collection.collection_id, + next_collection_id_floor: self.max_field_collection_id() + 1, + next_segment_id: self.next_segment_id, + next_global_id, + segment_ids, + keymap_epoch: seq, + }; + + let job = crate::vector::persistence::manifest::SnapshotJob { + idx_dir, + seq, + watermark: self.persist_seq_watermark.clone(), + manifest, + key_hash_to_key: self.key_hash_to_key.clone(), + key_hash_to_global_id: self.key_hash_to_global_id.clone(), + key_hash_to_vec_checksum: self.key_hash_to_vec_checksum.clone(), + }; + crate::vector::persistence::manifest::global_snapshot_pool().submit(job); + } } /// Max segments one bulk-loaded mutable is split into when its compact @@ -612,7 +787,10 @@ impl VectorIndex { .collection .collection_id .wrapping_mul(6364136223846793005); - match compactor.submit(frozen, self.collection.clone(), seed) { + // B2 (durability): allocate the disk segment id (and target dir) HERE, + // on the shard thread, at submit time — never on the worker. + let persist = self.alloc_persist_target(); + match compactor.submit(frozen, self.collection.clone(), seed, persist) { Ok(reply_rx) => { self.bg_compact_inflight = Some(InFlightCompaction { reply_rx, @@ -730,6 +908,10 @@ impl VectorIndex { }; drop(snap); self.segments.swap(new_list); + // B2 (durability): the worker already wrote the segment to disk (if + // `alloc_persist_target` handed it a target at submit time) — commit + // the keymap/manifest snapshot now that the install is durable. + self.persist_hook_after_install(); true } @@ -801,8 +983,19 @@ impl VectorIndex { // collapse without false-positives on small/medium indexes. Per-index // override: FT.CONFIG SET MERGE_RECALL_TOLERANCE (VEC-4). let tolerance = self.merge_recall_tolerance; - - match compactor.submit_merge(segs.clone(), self.collection.clone(), seed, mode, tolerance) { + // B2 (durability): allocate the disk segment id (and target dir) HERE, + // on the shard thread, at submit time — never on the worker. Merge + // persists identically to compact. + let persist = self.alloc_persist_target(); + + match compactor.submit_merge( + segs.clone(), + self.collection.clone(), + seed, + mode, + tolerance, + persist, + ) { Ok(reply_rx) => { self.bg_merge_inflight = Some(InFlightMerge { reply_rx, @@ -952,6 +1145,12 @@ impl VectorIndex { }; drop(snap); self.segments.swap(new_list); + // B2 (durability): the worker already wrote the merged segment to + // disk (if `alloc_persist_target` handed it a target at submit + // time) — commit the keymap/manifest snapshot now that the install + // is durable. The GC inside the snapshot job removes the (now + // superseded) source segment dirs. + self.persist_hook_after_install(); tracing::debug!( sources = inflight.merged_sources.len(), @@ -1168,9 +1367,6 @@ pub struct VectorStore { next_collection_id: u64, /// Per-shard MVCC transaction manager. txn_manager: TransactionManager, - /// Segments recovered from persistence, awaiting FT.CREATE to claim them. - /// Key: collection_id. Populated during crash recovery. - pending_segments: HashMap, /// Shard directory for persisting index metadata sidecar. /// Set once during event loop init when disk-offload is enabled. persist_dir: Option, @@ -1194,7 +1390,6 @@ impl VectorStore { indexes: HashMap::new(), next_collection_id: 1, txn_manager: TransactionManager::new(), - pending_segments: HashMap::new(), persist_dir: None, version_token: AtomicU64::new(0), } @@ -1289,26 +1484,6 @@ impl VectorStore { count } - /// Attach recovered segments from persistence. Called by shard restore. - /// - /// Stores recovered collections in pending_segments, keyed by collection_id. - /// They will be attached to indexes when FT.CREATE runs (or immediately if - /// the index already exists). - pub fn attach_recovered( - &mut self, - recovered: crate::vector::persistence::recovery::RecoveredState, - ) { - for (collection_id, collection) in recovered.collections { - self.pending_segments.insert(collection_id, collection); - } - } - - /// Number of pending (unattached) recovered collections. - #[allow(dead_code)] - pub fn pending_count(&self) -> usize { - self.pending_segments.len() - } - /// Create a new index. Returns Err(&str) if index already exists. pub fn create_index(&mut self, mut meta: IndexMeta) -> Result<(), &'static str> { if self.indexes.contains_key(&meta.name) { @@ -1368,6 +1543,27 @@ impl VectorStore { } let name = meta.name.clone(); + + // B2 (durability): defensive id-space floor. If this index's + // `idx-` dir already has a manifest — e.g. FLUSHALL/DROP just + // recreated this index and the best-effort background directory + // delete (see `drop_index`/`clear_all_contents`) hasn't finished yet + // — seed the fresh allocators ABOVE whatever it last recorded. This + // guarantees the new generation's segment ids/keymap epochs can never + // collide with (or race a delete of) the old generation's on-disk + // files, independent of when that background delete completes. + let (floor_next_segment_id, floor_next_snapshot_seq) = self + .persist_dir + .as_ref() + .map(|dir| { + let idx_dir = + crate::vector::persistence::manifest::index_persist_dir(dir, name.as_ref()); + crate::vector::persistence::manifest::read_manifest_tolerant(&idx_dir) + .map(|m| (m.next_segment_id, m.keymap_epoch)) + .unwrap_or((0, 0)) + }) + .unwrap_or((0, 0)); + self.indexes.insert( name.clone(), VectorIndex { @@ -1377,7 +1573,12 @@ impl VectorStore { collection, payload_index: PayloadIndex::new(), key_hash_to_key: Arc::new(std::collections::HashMap::new()), - key_hash_to_global_id: std::collections::HashMap::new(), + key_hash_to_global_id: Arc::new(std::collections::HashMap::new()), + key_hash_to_vec_checksum: Arc::new(std::collections::HashMap::new()), + persist_dir: self.persist_dir.clone(), + next_segment_id: floor_next_segment_id, + next_snapshot_seq: floor_next_snapshot_seq, + persist_seq_watermark: Arc::new(parking_lot::Mutex::new(floor_next_snapshot_seq)), autocompact_enabled: true, merge_recall_tolerance: 0.70, compaction_weight: COMPACTION_WEIGHT_DEFAULT, @@ -1391,32 +1592,176 @@ impl VectorStore { // Persist index metadata sidecar self.save_index_meta_sidecar(); - // Check if recovered segments exist for this collection_id - if let Some(recovered) = self.pending_segments.remove(&collection_id) { - if let Some(index) = self.indexes.get(&name) { - let mut immutable_arcs: Vec< - Arc, - > = Vec::with_capacity(recovered.immutable.len()); - for (imm, _meta) in recovered.immutable { - immutable_arcs.push(Arc::new(imm)); - } - let new_list = crate::vector::segment::SegmentList { - mutable: Arc::new(recovered.mutable), - immutable: immutable_arcs, - ivf: Vec::new(), - warm: Vec::new(), - cold: Vec::new(), - }; - index.segments.swap(new_list); - } + // Bump version AFTER successful write (monotonicity-on-success contract). + self.bump_version(); + + Ok(()) + } + + /// Create a new index, pinning its default field's `collection_id` (and + /// therefore its HNSW QJL rotation seed) to a value previously recorded + /// on disk, instead of allocating a fresh one from `next_collection_id`. + /// + /// B3 recovery-only: used exclusively by the durability loader + /// (`crate::vector::persistence::recover_v2`) when a `manifest.json` was + /// found for this index. A segment persisted under collection_id X is + /// only searchable if the index that owns it also has collection_id X — + /// see `VECTOR-DURABILITY-DESIGN.md`'s "#1 correctness trap". Mirrors + /// `create_index` in every other respect (additional-field segments are + /// always fresh — B1/B2 never persists them), and additionally seeds + /// `next_segment_id`/`next_snapshot_seq` from the SAME manifest instead + /// of a defensive re-read (the caller already loaded it). + /// + /// Does not attach any segments or keymap state — `VectorIndex.segments`, + /// `key_hash_to_key`, `key_hash_to_global_id`, `key_hash_to_vec_checksum` + /// are left at their fresh-index defaults. The caller installs recovered + /// state afterward via `get_index_mut` (all four fields are `pub`). + pub(crate) fn create_index_with_collection_id( + &mut self, + mut meta: IndexMeta, + manifest: &crate::vector::persistence::manifest::IndexManifest, + ) -> Result<(), &'static str> { + if self.indexes.contains_key(&meta.name) { + return Err("Index already exists"); } + // Backward compatibility: if vector_fields is empty, populate from top-level fields. + if meta.vector_fields.is_empty() { + meta.vector_fields = vec![VectorFieldMeta { + field_name: meta.source_field.clone(), + dimension: meta.dimension, + padded_dimension: padded_dimension(meta.dimension), + metric: meta.metric, + quantization: meta.quantization, + build_mode: meta.build_mode, + }]; + } + + let collection_id = manifest.collection_id; + // Bump the store-wide allocator above BOTH the pinned cid and the + // manifest's own recorded floor (covers additional-field cids this + // index previously allocated — those fields are never persisted, so + // they get fresh cids below, but must never collide with a cid a + // sibling recovered index is pinning right now). + self.next_collection_id = self + .next_collection_id + .max(manifest.next_collection_id_floor) + .max(collection_id + 1); + + let padded = padded_dimension(meta.dimension); + let collection = Arc::new(CollectionMetadata::with_build_mode( + collection_id, + meta.dimension, + meta.metric, + meta.quantization, + collection_id, // pinned: same value the persisted segments' QJL was seeded with + meta.build_mode, + )); + let segments = SegmentHolder::new(meta.dimension, collection.clone()); + let scratch = SearchScratch::new(0, padded); + + // Additional field segments: B1/B2 never persists them (AS-BUILT — + // `compact_segments` is always called with `persist_root: None` for + // `field_segments`), so they always start fresh here too, exactly + // like `create_index`. `next_collection_id` was already bumped above + // so these freshly allocated cids can't collide with the pinned one. + let mut extra_fields = HashMap::new(); + for field_meta in meta.vector_fields.iter().skip(1) { + let field_cid = self.next_collection_id; + self.next_collection_id += 1; + let field_padded = padded_dimension(field_meta.dimension); + let field_collection = Arc::new(CollectionMetadata::with_build_mode( + field_cid, + field_meta.dimension, + field_meta.metric, + field_meta.quantization, + field_cid, + field_meta.build_mode, + )); + let field_segments = SegmentHolder::new(field_meta.dimension, field_collection.clone()); + let field_scratch = SearchScratch::new(0, field_padded); + extra_fields.insert( + field_meta.field_name.clone(), + FieldSegments { + segments: field_segments, + scratch: field_scratch, + collection: field_collection, + }, + ); + } + + let name = meta.name.clone(); + + self.indexes.insert( + name.clone(), + VectorIndex { + meta, + segments, + scratch, + collection, + payload_index: PayloadIndex::new(), + key_hash_to_key: Arc::new(std::collections::HashMap::new()), + key_hash_to_global_id: Arc::new(std::collections::HashMap::new()), + key_hash_to_vec_checksum: Arc::new(std::collections::HashMap::new()), + persist_dir: self.persist_dir.clone(), + next_segment_id: manifest.next_segment_id, + next_snapshot_seq: manifest.keymap_epoch, + persist_seq_watermark: Arc::new(parking_lot::Mutex::new(manifest.keymap_epoch)), + autocompact_enabled: true, + merge_recall_tolerance: 0.70, + compaction_weight: COMPACTION_WEIGHT_DEFAULT, + field_segments: extra_fields, + sparse_stores: HashMap::new(), + bg_compact_inflight: None, + bg_merge_inflight: None, + }, + ); + + // Persist index metadata sidecar + self.save_index_meta_sidecar(); + // Bump version AFTER successful write (monotonicity-on-success contract). self.bump_version(); Ok(()) } + /// Best-effort background removal of an index's whole `idx-/` + /// durability directory (manifest + keymaps + segments) — B2, design + /// doc item "Drop/flush cleanup". Fire-and-forget: runs on its own OS + /// thread so `drop_index`/`clear_all_contents` never block the shard on + /// a (possibly large) recursive directory removal. Failure is logged and + /// otherwise harmless — a leftover directory is swept by the B3 startup + /// orphan sweep, and `create_index`'s defensive manifest-floor read + /// guarantees a same-named index recreated before this finishes never + /// reuses (or races) its ids. No-op if `persist_dir` is not configured. + fn spawn_delete_index_persist_dir(&self, name: &Bytes) { + let Some(dir) = self.persist_dir.clone() else { + return; + }; + let name = name.clone(); + let idx_dir = crate::vector::persistence::manifest::index_persist_dir(&dir, name.as_ref()); + let spawned = std::thread::Builder::new() + .name("moon-vec-idx-gc".to_owned()) + .spawn(move || { + if let Err(e) = std::fs::remove_dir_all(&idx_dir) { + if e.kind() != std::io::ErrorKind::NotFound { + tracing::warn!( + "failed to remove durability dir {} for dropped index: {e}", + idx_dir.display() + ); + } + } + }); + if let Err(e) = spawned { + tracing::warn!( + "failed to spawn background deletion thread for dropped index {}: {e} \ + (durability dir left on disk; swept at next B3 startup)", + String::from_utf8_lossy(&name) + ); + } + } + /// Drop an index by name. Returns true if it existed. /// /// Tombstones any warm segments so their on-disk directories are cleaned up @@ -1428,6 +1773,8 @@ impl VectorStore { for warm_seg in &snapshot.warm { warm_seg.mark_tombstoned(); } + drop(snapshot); + self.spawn_delete_index_persist_dir(&index.meta.name); // Persist index metadata sidecar self.save_index_meta_sidecar(); // Bump version AFTER successful drop (monotonicity-on-success contract). @@ -1438,6 +1785,42 @@ impl VectorStore { } } + /// FLUSHALL/FLUSHDB parity (persistence-review R3): drop every index's + /// CONTENTS (segments, key-hash maps, payload/MVCC state) while KEEPING + /// the FT.CREATE definitions — mirroring restart semantics, where + /// definitions come from the sidecar and contents are re-derived from + /// the (now empty) keyspace. Without this, flushed hashes stayed + /// searchable as ghost vectors until the next restart. + pub fn clear_all_contents(&mut self) { + // FLUSH discards everything. + let names: Vec = self.indexes.keys().cloned().collect(); + for name in names { + let Some(index) = self.indexes.remove(&name) else { + continue; + }; + // Tombstone warm segments so their on-disk directories are + // reclaimed once in-flight search snapshots drop (as drop_index). + let snapshot = index.segments.load(); + for warm_seg in &snapshot.warm { + warm_seg.mark_tombstoned(); + } + drop(snapshot); + let meta = index.meta.clone(); + drop(index); + // B2 (durability): FLUSH discards contents — the whole durability + // dir (segments/manifest/keymaps) goes with it. Best-effort, + // background (see `spawn_delete_index_persist_dir`); the + // immediately-following `create_index` reseeds id allocators + // above whatever the (possibly still-being-deleted) old manifest + // recorded, so the recreated index never collides with it. + self.spawn_delete_index_persist_dir(&name); + // Recreate from the same definition — a fresh, empty index. + // Also rewrites the sidecar and bumps the version token. + #[allow(clippy::unwrap_used)] // name was just removed above; create cannot collide + self.create_index(meta).unwrap(); + } + } + /// Get index reference by name. pub fn get_index(&self, name: &[u8]) -> Option<&VectorIndex> { self.indexes.get(name) @@ -1493,23 +1876,7 @@ impl VectorStore { let key_hash = xxhash_rust::xxh64::xxh64(key, 0); let mut any_deleted = false; for idx_name in matching_names { - if let Some(idx) = self.indexes.get_mut(&idx_name) { - let snap = idx.segments.load(); - // Tombstone in mutable segment (always present). - snap.mutable.mark_deleted_by_key_hash(key_hash, 1); - // Also tombstone any already-compacted immutable segments that - // may still contain the key (steady-state interior tombstone). - for imm in snap.immutable.iter() { - imm.mark_deleted_by_key_hash(key_hash); - } - // QW7 (2026-06 review finding 6.3): prune the key-hash maps so - // they track LIVE keys, not historical inserts — without this - // they grow monotonically under key churn (~1GB / 24M deletes). - // A re-insert of the same key repopulates both maps. - Arc::make_mut(&mut idx.key_hash_to_key).remove(&key_hash); - idx.key_hash_to_global_id.remove(&key_hash); - any_deleted = true; - } + any_deleted |= self.tombstone_key_in_index(&idx_name, key_hash); } // Bump version AFTER any successful deletion mark. if any_deleted { @@ -1517,6 +1884,47 @@ impl VectorStore { } } + /// Per-index variant of [`Self::mark_deleted_for_key`] (persistence-review + /// R4): tombstones `key` in ONE named index. Used by the HDEL hook, where + /// only indexes whose vector field was actually removed may be touched — + /// a sibling index keyed on a different field must keep its entry. + pub fn mark_deleted_for_key_in_index(&mut self, idx_name: &[u8], key: &[u8]) { + let key_hash = xxhash_rust::xxh64::xxh64(key, 0); + let Some((name, _)) = self.indexes.get_key_value(idx_name) else { + return; + }; + let name = name.clone(); + if self.tombstone_key_in_index(&name, key_hash) { + self.bump_version(); + } + } + + fn tombstone_key_in_index(&mut self, idx_name: &Bytes, key_hash: u64) -> bool { + let Some(idx) = self.indexes.get_mut(idx_name) else { + return false; + }; + let snap = idx.segments.load(); + // Tombstone in mutable segment (always present). + snap.mutable.mark_deleted_by_key_hash(key_hash, 1); + // Also tombstone any already-compacted immutable segments that + // may still contain the key (steady-state interior tombstone). + for imm in snap.immutable.iter() { + imm.mark_deleted_by_key_hash(key_hash); + } + drop(snap); + // QW7 (2026-06 review finding 6.3): prune the key-hash maps so + // they track LIVE keys, not historical inserts — without this + // they grow monotonically under key churn (~1GB / 24M deletes). + // A re-insert of the same key repopulates all three maps. + Arc::make_mut(&mut idx.key_hash_to_key).remove(&key_hash); + Arc::make_mut(&mut idx.key_hash_to_global_id).remove(&key_hash); + // B2 (durability): keep the checksum map in lockstep so it never + // drifts from key_hash_to_key (a stale entry would be a silent + // false-"unchanged" in the B3 dedup rescan). + Arc::make_mut(&mut idx.key_hash_to_vec_checksum).remove(&key_hash); + true + } + /// Dispatch background compactions for all indexes that are ready /// (mutable segment non-empty, no compaction already in flight). /// @@ -1969,6 +2377,9 @@ impl VectorStore { }; drop(snap); idx.segments.swap(new_list); + // B2 (durability): worker already persisted (if configured + // at submit time) — commit the keymap/manifest snapshot. + idx.persist_hook_after_install(); // Data is already merged — return early. let new_snap = idx.segments.load(); let live = new_snap @@ -2010,8 +2421,19 @@ impl VectorStore { let collection = idx.collection.clone(); let seed = collection.collection_id.wrapping_mul(6364136223846793005); drop(snap); - - match compaction::merge_immutable(&segs, &collection, seed, mode, recall_tolerance) { + // B2 (durability): allocate the disk segment id synchronously (this + // is the explicit-command path — a brief stall is acceptable, same + // rationale as the inline compact path). + let persist = idx.alloc_persist_target(); + + match compaction::merge_immutable( + &segs, + &collection, + seed, + mode, + recall_tolerance, + persist.as_ref().map(|(p, id)| (p.as_path(), *id)), + ) { Ok(merged) => { let live = merged.live_count() as usize; // Atomically swap: replace all immutable segments with the single merged one. @@ -2024,6 +2446,7 @@ impl VectorStore { cold: old.cold.clone(), }; idx.segments.swap(new_list); + idx.persist_hook_after_install(); // Rebuild scratch for the merged segment. let new_snap = idx.segments.load(); @@ -2286,6 +2709,160 @@ mod tests { assert!(!store.drop_index(b"nonexistent")); } + /// Builds HSET-style args `[key, field, vector_bytes]` for + /// `auto_index_hset_public`, mirroring the wire format `find_vector_blob` + /// expects (field/value pairs starting at index 1). + fn hset_vector_args(key: &[u8], field: &[u8], vec: &[f32]) -> Vec { + let mut bytes = Vec::with_capacity(vec.len() * 4); + for v in vec { + bytes.extend_from_slice(&v.to_le_bytes()); + } + vec![ + crate::protocol::Frame::BulkString(Bytes::copy_from_slice(key)), + crate::protocol::Frame::BulkString(Bytes::copy_from_slice(field)), + crate::protocol::Frame::BulkString(Bytes::from(bytes)), + ] + } + + /// B2 (durability): `key_hash_to_vec_checksum` must be maintained in + /// lockstep with `key_hash_to_key` at every mutation site — insert, + /// update (checksum changes with the vector), and tombstone (both maps + /// prune the same key). A drift here would silently corrupt the B3 + /// dedup rescan's unchanged-vs-changed decision. + #[test] + fn test_checksum_map_tracks_key_hash_to_key() { + let mut store = VectorStore::new(); + store.create_index(make_meta("idx", 4, &["doc:"])).unwrap(); + let mut text_store = crate::text::store::TextStore::new(); + + let key_hash = xxhash_rust::xxh64::xxh64(b"doc:1", 0); + let v1 = [1.0f32, 2.0, 3.0, 4.0]; + let expected_checksum_v1 = { + let mut bytes = Vec::with_capacity(16); + for v in &v1 { + bytes.extend_from_slice(&v.to_le_bytes()); + } + xxhash_rust::xxh64::xxh64(&bytes, 0) + }; + + // ── Insert ────────────────────────────────────────────────────── + crate::shard::spsc_handler::auto_index_hset_public( + &mut store, + &mut text_store, + b"doc:1", + &hset_vector_args(b"doc:1", b"vec", &v1), + ); + { + let idx = store.get_index(b"idx").unwrap(); + assert!(idx.key_hash_to_key.contains_key(&key_hash)); + assert_eq!( + idx.key_hash_to_vec_checksum.get(&key_hash).copied(), + Some(expected_checksum_v1), + "checksum map must be populated on insert, matching xxh64 of the raw vector bytes" + ); + assert_eq!( + idx.key_hash_to_key.len(), + idx.key_hash_to_vec_checksum.len(), + "checksum map must track key_hash_to_key 1:1" + ); + } + + // ── Update (different vector -> different checksum) ──────────── + let v2 = [5.0f32, 6.0, 7.0, 8.0]; + crate::shard::spsc_handler::auto_index_hset_public( + &mut store, + &mut text_store, + b"doc:1", + &hset_vector_args(b"doc:1", b"vec", &v2), + ); + { + let idx = store.get_index(b"idx").unwrap(); + let updated_checksum = idx.key_hash_to_vec_checksum.get(&key_hash).copied(); + assert_ne!( + updated_checksum, + Some(expected_checksum_v1), + "checksum must change when the vector changes" + ); + assert!(updated_checksum.is_some()); + assert_eq!( + idx.key_hash_to_key.len(), + idx.key_hash_to_vec_checksum.len() + ); + } + + // ── Tombstone (DEL) prunes both maps together ─────────────────── + store.mark_deleted_for_key(b"doc:1"); + { + let idx = store.get_index(b"idx").unwrap(); + assert!(!idx.key_hash_to_key.contains_key(&key_hash)); + assert!( + !idx.key_hash_to_vec_checksum.contains_key(&key_hash), + "checksum map must be pruned alongside key_hash_to_key on delete" + ); + } + } + + /// End-to-end (B2, durability write path): `force_compact` with a + /// `persist_dir` configured must (a) write the compacted segment to disk + /// under `idx-/segment-/` via the staged writer, and (b) — once + /// the background snapshot job lands — commit a `manifest.json` that + /// references that segment id and a readable `keymap-.bin` + /// containing the inserted key. + #[test] + fn test_force_compact_persists_segment_and_manifest_when_persist_dir_set() { + let tmp = tempfile::tempdir().unwrap(); + let mut store = VectorStore::new(); + // set_persist_dir BEFORE create_index, mirroring event_loop.rs's + // startup order (create_index captures the store's persist_dir). + store.set_persist_dir(tmp.path().to_path_buf()); + store.create_index(make_meta("idx", 4, &["doc:"])).unwrap(); + let mut text_store = crate::text::store::TextStore::new(); + + crate::shard::spsc_handler::auto_index_hset_public( + &mut store, + &mut text_store, + b"doc:1", + &hset_vector_args(b"doc:1", b"vec", &[1.0, 2.0, 3.0, 4.0]), + ); + + store.get_index_mut(b"idx").unwrap().force_compact(); + + // The segment write happens synchronously (inline compact path) — + // must be visible immediately, no polling needed. + let idx_dir = crate::vector::persistence::manifest::index_persist_dir(tmp.path(), b"idx"); + assert!( + idx_dir.join("segment-0").join("segment_meta.json").exists(), + "compacted segment must be persisted synchronously under segment-0" + ); + assert!( + !idx_dir.join("staging-0").exists(), + "staged writer must leave no staging dir behind" + ); + + // The keymap/manifest snapshot is a background job — poll briefly. + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(5); + let manifest = loop { + if let Some(m) = crate::vector::persistence::manifest::read_manifest_tolerant(&idx_dir) + { + break m; + } + assert!( + std::time::Instant::now() < deadline, + "manifest.json never appeared within 5s" + ); + std::thread::sleep(std::time::Duration::from_millis(10)); + }; + + assert_eq!(manifest.segment_ids, vec![0]); + let keymap = crate::vector::persistence::manifest::read_keymap_tolerant( + &idx_dir, + manifest.keymap_epoch, + ) + .expect("keymap for the committed epoch must be readable"); + assert_eq!(keymap.len(), 1); + assert_eq!(keymap[0].key_hash, xxhash_rust::xxh64::xxh64(b"doc:1", 0)); + } + #[test] fn test_find_matching_indexes() { let mut store = VectorStore::new(); @@ -2663,6 +3240,43 @@ mod bg_compact_tests { /// ceil(n/threshold) threshold-sized immutable segments, not one giant /// graph — multiple segments are what the intra-query worker pool fans /// out over. All keys must remain findable (self-recall probe). + #[test] + fn test_compact_measures_suggested_ef() { + // AE-1: a compact build with a sidecar must attach a suggested ef + // from the ladder; tiny/no-sidecar segments stay None (checked by + // the estimator's own guards). + distance::init(); + crate::vector::search_pool::init_global(1); + let dim = 16u32; + let mut store = VectorStore::new(); + let mut meta = make_idx(dim); + meta.compact_threshold = 200; + store.create_index(meta).unwrap(); + for i in 0..400u64 { + insert( + &mut store, + format!("doc:{i}").as_bytes(), + random_vec(dim as usize, i), + ); + } + store.force_compact_index(b"idx").unwrap(); + let idx = store.indexes.get_mut(b"idx".as_ref()).unwrap(); + let snap = idx.segments.load_full(); + assert!(!snap.immutable.is_empty()); + for seg in &snap.immutable { + let sug = seg.suggested_ef(); + assert!( + sug.is_some(), + "compact build with sidecar must measure an adaptive ef" + ); + let sug = sug.unwrap() as usize; + assert!( + (24..=256).contains(&sug), + "ladder value expected, got {sug}" + ); + } + } + #[test] fn test_force_compact_bulk_bounded_segments() { distance::init(); diff --git a/src/vector/turbo_quant/sq8.rs b/src/vector/turbo_quant/sq8.rs index 87ca9fe3..370b1429 100644 --- a/src/vector/turbo_quant/sq8.rs +++ b/src/vector/turbo_quant/sq8.rs @@ -277,6 +277,120 @@ pub fn sq8_ip_from_stats(min: f32, scale: f32, q_sum: f32, dot_qc: f32) -> f32 { min * q_sum + scale * dot_qc } +// ── int8 symmetric ADC (task #13) ─────────────────────────────────────── +// +// The stats decomposition above still widens the u8 codes to `f32` for the +// per-candidate pass (`sq8_candidate_stats_scalar` / the SIMD `sq8_stats` +// kernels) — correct, but each candidate byte pays a u8->f32 convert plus an +// FMA. Integer dot products (NEON widening multiply, x86 widen+madd, AVX512 +// VPDPBUSD) move 2-4x more bytes per instruction than the float path. +// +// Design (see tmp/INT8-ADC-CONTEXT.md): quantize the QUERY once per search +// to symmetric int8 (no zero point, `qi8 = round(q / q_scale)`); `sum_c` and +// `sumsq_c` are exact integers straight from the u8 codes (query-independent, +// unaffected by quantization); `dot_qc` is the only approximate term, +// reconstructed as `q_scale * (exact integer Σ qi8·c)`. +// +// `SQ8_INT8_QMAX` is shared by every SIMD tier (NEON widening multiply, +// AVX2, AVX512-VNNI) — see `src/vector/distance/avx2.rs` module docs for why +// the originally-planned VPMADDUBSW-with-clamped-query design (qmax=63) was +// dropped: it measurably regresses recall (see the A/B test below), so all +// tiers instead use a saturation-free widen-multiply-then-32-bit-accumulate +// shape that tolerates the full `[-127, 127]` range. + +/// Query-quantization amplitude shared by every int8 ADC SIMD tier. Chosen so +/// even the least-precise integer path (originally AVX2 VPMADDUBSW, now +/// replaced by a saturation-free widening kernel — see module docs) needs no +/// special-cased narrower range: every tier quantizes to the same +/// `[-127, 127]`, so a single `qi8` buffer serves NEON, AVX2, and AVX512. +pub const SQ8_INT8_QMAX: i32 = 127; + +/// Quantize `query` to symmetric int8 for the int8 ADC path: `qi8 = +/// round(q / q_scale).clamp(-qmax, qmax)`, `q_scale = max|q| / qmax` (`0.0` +/// for an all-zero query, in which case every `qi8` entry is `0`). +/// +/// Call **once per search** (like [`sq8_query_stats`]) — this is not a +/// per-candidate cost. Returns `(q_scale, sum_qi8)`; `sum_qi8` (exact `i32` +/// sum of the quantized codes) is consumed only by kernels that need the +/// NEON-style offset-trick correction internally — scalar/AVX2/AVX512 +/// callers may ignore it, but it is threaded through the shared +/// [`Sq8I8StatsFn`] signature so every tier's kernel has it available +/// without a second reduction pass. +/// +/// # Panics (debug only) +/// `debug_assert_eq!(query.len(), out.len())`. +pub fn sq8_quantize_query_scalar(query: &[f32], qmax: i32, out: &mut [i8]) -> (f32, i32) { + debug_assert_eq!(query.len(), out.len()); + debug_assert!(qmax > 0 && qmax <= 127); + let mut max_abs = 0.0f32; + for &q in query { + let a = q.abs(); + if a > max_abs { + max_abs = a; + } + } + let q_scale = if max_abs > 0.0 { + max_abs / qmax as f32 + } else { + 0.0 + }; + let inv = if q_scale > 0.0 { 1.0 / q_scale } else { 0.0 }; + let qmax_f = qmax as f32; + let mut sum_qi8 = 0i32; + for (o, &q) in out.iter_mut().zip(query) { + let qi = (q * inv).round().clamp(-qmax_f, qmax_f) as i32; + *o = qi as i8; + sum_qi8 += qi; + } + (q_scale, sum_qi8) +} + +/// Shared function-pointer signature for the int8 per-candidate ADC stats +/// pass, implemented by the scalar reference below and by the NEON/AVX2/ +/// AVX512 kernels installed into `DistanceTable::sq8_i8_stats`. Every +/// implementation must return **exactly** `(Σ qi8_i·c_i, Σ c_i, Σ c_i²)` as +/// exact `i64` integers — this is pure integer arithmetic (no float +/// rounding anywhere in the per-candidate pass), so SIMD kernels are +/// checked against the scalar reference for **bit-for-bit equality**, not +/// float tolerance. +/// +/// `sum_qi8` (from [`sq8_quantize_query_scalar`]) is accepted by every +/// implementation for interface parity with the NEON offset-trick kernel, +/// which needs it; the scalar reference and the AVX2/AVX512 kernels ignore +/// it (their widen-multiply shapes need no offset correction). +pub type Sq8I8StatsFn = fn(qi8: &[i8], codes: &[u8], sum_qi8: i32) -> (i64, i64, i64); + +/// Scalar reference / correctness oracle for the int8 per-candidate ADC +/// stats pass. Exact integer arithmetic throughout (`i64` accumulation) — +/// this is the ground truth every SIMD int8 kernel must reproduce exactly. +/// +/// # Panics (debug only) +/// `debug_assert_eq!(qi8.len(), codes.len())`. +#[inline] +pub fn sq8_candidate_stats_i8_scalar(qi8: &[i8], codes: &[u8], _sum_qi8: i32) -> (i64, i64, i64) { + debug_assert_eq!(qi8.len(), codes.len()); + let mut dot = 0i64; + let mut sum_c = 0i64; + let mut sumsq_c = 0i64; + for (&q, &c) in qi8.iter().zip(codes) { + let qi = q as i64; + let ci = c as i64; + dot += qi * ci; + sum_c += ci; + sumsq_c += ci * ci; + } + (dot, sum_c, sumsq_c) +} + +/// Reconstruct the approximate `dot_qc` term (`f32`, the only lossy value in +/// the int8 path) from the exact integer dot product and the per-query +/// `q_scale`. `O(1)`; call once per candidate right before +/// [`sq8_l2_from_stats`]. +#[inline] +pub fn sq8_int8_dot_to_f32(q_scale: f32, dot_qc_int: i64) -> f32 { + q_scale * dot_qc_int as f32 +} + #[cfg(test)] mod tests { use super::*; @@ -603,4 +717,304 @@ mod tests { "stats-decomposed SQ8 nearest neighbor diverged from exact" ); } + + // ── int8 symmetric ADC (task #13) ──────────────────────────────────── + + #[test] + fn test_i8_quantize_within_qmax_and_error_bounded() { + for &dim in &[1usize, 7, 16, 128, 384, 768] { + let q = pseudo_vec(42_000 + dim as u64, dim); + let mut qi8 = vec![0i8; dim]; + let (q_scale, sum_qi8) = sq8_quantize_query_scalar(&q, SQ8_INT8_QMAX, &mut qi8); + let mut exact_sum = 0i32; + for (&orig, &qi) in q.iter().zip(&qi8) { + assert!( + (qi as i32).abs() <= SQ8_INT8_QMAX, + "dim={dim}: qi8 value {qi} exceeds qmax {SQ8_INT8_QMAX}" + ); + exact_sum += qi as i32; + // round-to-nearest quantization error is bounded by half a step. + let reconstructed = qi as f32 * q_scale; + assert!( + (orig - reconstructed).abs() <= q_scale * 0.5 + 1e-6, + "dim={dim}: orig={orig} reconstructed={reconstructed} q_scale={q_scale}" + ); + } + assert_eq!(sum_qi8, exact_sum, "dim={dim}: sum_qi8 mismatch"); + } + } + + #[test] + fn test_i8_quantize_all_zero_query_is_zero() { + let q = vec![0.0f32; 32]; + let mut qi8 = vec![1i8; 32]; // pre-poisoned to catch a no-op quantize + let (q_scale, sum_qi8) = sq8_quantize_query_scalar(&q, SQ8_INT8_QMAX, &mut qi8); + assert_eq!(q_scale, 0.0); + assert_eq!(sum_qi8, 0); + assert!(qi8.iter().all(|&c| c == 0)); + } + + #[test] + fn test_i8_candidate_stats_matches_direct_i64_computation() { + // The scalar reference IS the i64 oracle by construction (no widening, + // no reassociation) — this test pins that down explicitly so a future + // refactor of the scalar function can't silently drift from "exact". + for &dim in &[0usize, 1, 3, 16, 100, 384] { + let q = pseudo_vec(43_000 + dim as u64, dim); + let v = pseudo_vec(44_000 + dim as u64, dim); + let mut qi8 = vec![0i8; dim]; + let (_q_scale, sum_qi8) = sq8_quantize_query_scalar(&q, SQ8_INT8_QMAX, &mut qi8); + let codes = encode_sq8(&v); + let (dot, sum_c, sumsq_c) = sq8_candidate_stats_i8_scalar(&qi8, &codes[..dim], sum_qi8); + + let mut dot_direct = 0i64; + let mut sum_c_direct = 0i64; + let mut sumsq_c_direct = 0i64; + for i in 0..dim { + let qi = qi8[i] as i64; + let ci = codes[i] as i64; + dot_direct += qi * ci; + sum_c_direct += ci; + sumsq_c_direct += ci * ci; + } + assert_eq!( + (dot, sum_c, sumsq_c), + (dot_direct, sum_c_direct, sumsq_c_direct), + "dim={dim}" + ); + } + } + + #[test] + fn test_i8_stats_zero_dim_is_zero() { + let mut qi8: Vec = vec![]; + let (q_scale, sum_qi8) = sq8_quantize_query_scalar(&[], SQ8_INT8_QMAX, &mut qi8); + assert_eq!((q_scale, sum_qi8), (0.0, 0)); + assert_eq!(sq8_candidate_stats_i8_scalar(&[], &[], sum_qi8), (0, 0, 0)); + } + + /// End-to-end recall A/B: same dataset, f32-stats ADC vs int8-stats ADC, + /// vs an exact-f32 brute-force ground truth. Per tmp/INT8-ADC-CONTEXT.md + /// this is the primary safety net for the int8 path (the exact-rerank + /// sidecar elsewhere is a *second* layer — this test deliberately does + /// NOT invoke it, so it isolates the int8 ADC's own error instead of + /// letting the sidecar wash it out). + fn recall_at_10(candidate: &[usize], truth: &[usize]) -> f32 { + let hit = candidate.iter().filter(|c| truth.contains(c)).count(); + hit as f32 / truth.len() as f32 + } + + fn top10_by_dist(dists: &[(usize, f32)]) -> Vec { + let mut v = dists.to_vec(); + v.sort_by(|a, b| a.1.total_cmp(&b.1)); + v.into_iter().take(10).map(|(i, _)| i).collect() + } + + /// Shared A/B harness: builds a seeded synthetic dataset, runs ground + /// truth / f32-ADC / int8-ADC brute force, and returns + /// `(mean R@10 f32-vs-truth, mean R@10 int8-vs-truth, mean overlap + /// int8-vs-f32)`. `n=5000`, `n_queries=200` — enough to average out + /// per-query noise (30 queries is too few: a single flipped neighbor + /// swings the mean by 3.3 points, wide enough to make the 0.01 gate + /// meaningless either way). + fn int8_recall_ab(dim: usize) -> (f32, f32, f32) { + const N: usize = 5000; + const N_QUERIES: usize = 200; + + let db: Vec> = (0..N) + .map(|i| pseudo_vec(500_000 + i as u64, dim)) + .collect(); + let encoded: Vec> = db.iter().map(|v| encode_sq8(v)).collect(); + + let mut r_f32 = 0.0f32; + let mut r_int8 = 0.0f32; + let mut overlap = 0.0f32; + + for qi in 0..N_QUERIES { + let q = pseudo_vec(900_000 + qi as u64, dim); + + let truth_dists: Vec<(usize, f32)> = db + .iter() + .enumerate() + .map(|(i, v)| (i, l2_sq(&q, v))) + .collect(); + let truth_top10 = top10_by_dist(&truth_dists); + + let (q_sum, q_sumsq) = sq8_query_stats(&q); + let f32_dists: Vec<(usize, f32)> = encoded + .iter() + .enumerate() + .map(|(i, slot)| { + let (min, scale) = sq8_params(slot, dim); + let (dot_qc, sum_c, sumsq_c) = sq8_candidate_stats_scalar(&q, &slot[..dim]); + let d = + sq8_l2_from_stats(dim, min, scale, q_sum, q_sumsq, dot_qc, sum_c, sumsq_c); + (i, d) + }) + .collect(); + let f32_top10 = top10_by_dist(&f32_dists); + + let mut qi8 = vec![0i8; dim]; + let (q_scale, sum_qi8) = sq8_quantize_query_scalar(&q, SQ8_INT8_QMAX, &mut qi8); + let int8_dists: Vec<(usize, f32)> = encoded + .iter() + .enumerate() + .map(|(i, slot)| { + let (min, scale) = sq8_params(slot, dim); + let (dot_int, sum_c_int, sumsq_c_int) = + sq8_candidate_stats_i8_scalar(&qi8, &slot[..dim], sum_qi8); + let dot_qc = sq8_int8_dot_to_f32(q_scale, dot_int); + let d = sq8_l2_from_stats( + dim, + min, + scale, + q_sum, + q_sumsq, + dot_qc, + sum_c_int as f32, + sumsq_c_int as f32, + ); + (i, d) + }) + .collect(); + let int8_top10 = top10_by_dist(&int8_dists); + + r_f32 += recall_at_10(&f32_top10, &truth_top10); + r_int8 += recall_at_10(&int8_top10, &truth_top10); + overlap += recall_at_10(&int8_top10, &f32_top10); + } + + ( + r_f32 / N_QUERIES as f32, + r_int8 / N_QUERIES as f32, + overlap / N_QUERIES as f32, + ) + } + + #[test] + fn test_int8_adc_recall_ab_dim128() { + let (r_f32, r_int8, overlap) = int8_recall_ab(128); + assert!( + overlap >= 0.98, + "dim128: overlap(int8 vs f32)={overlap} below 0.98 gate" + ); + assert!( + (r_f32 - r_int8).abs() <= 0.01, + "dim128: R@10 f32={r_f32} int8={r_int8} delta={} exceeds 0.01", + (r_f32 - r_int8).abs() + ); + } + + #[test] + fn test_int8_adc_recall_ab_dim384() { + let (r_f32, r_int8, overlap) = int8_recall_ab(384); + assert!( + overlap >= 0.98, + "dim384: overlap(int8 vs f32)={overlap} below 0.98 gate" + ); + assert!( + (r_f32 - r_int8).abs() <= 0.01, + "dim384: R@10 f32={r_f32} int8={r_int8} delta={} exceeds 0.01", + (r_f32 - r_int8).abs() + ); + } + + /// Cosine metric: SQ8 vectors normalized to the unit sphere, then the L2 + /// combine is reused (the "cosine trick" — squared L2 on unit vectors is + /// monotonic with cosine similarity). Confirms the int8 path holds recall + /// under the SAME reuse the production Cosine/InnerProduct call sites + /// rely on (they never call `sq8_ip_from_stats` — see module docs). + #[test] + fn test_int8_adc_recall_ab_cosine_normalized() { + const DIM: usize = 128; + const N: usize = 5000; + const N_QUERIES: usize = 200; + + fn normalize(v: &mut [f32]) { + let norm: f32 = v.iter().map(|x| x * x).sum::().sqrt(); + if norm > 0.0 { + for x in v.iter_mut() { + *x /= norm; + } + } + } + + let db: Vec> = (0..N) + .map(|i| { + let mut v = pseudo_vec(700_000 + i as u64, DIM); + normalize(&mut v); + v + }) + .collect(); + let encoded: Vec> = db.iter().map(|v| encode_sq8(v)).collect(); + + let mut r_f32 = 0.0f32; + let mut r_int8 = 0.0f32; + let mut overlap = 0.0f32; + + for qi in 0..N_QUERIES { + let mut q = pseudo_vec(950_000 + qi as u64, DIM); + normalize(&mut q); + + let truth_dists: Vec<(usize, f32)> = db + .iter() + .enumerate() + .map(|(i, v)| (i, l2_sq(&q, v))) + .collect(); + let truth_top10 = top10_by_dist(&truth_dists); + + let (q_sum, q_sumsq) = sq8_query_stats(&q); + let f32_dists: Vec<(usize, f32)> = encoded + .iter() + .enumerate() + .map(|(i, slot)| { + let (min, scale) = sq8_params(slot, DIM); + let (dot_qc, sum_c, sumsq_c) = sq8_candidate_stats_scalar(&q, &slot[..DIM]); + let d = + sq8_l2_from_stats(DIM, min, scale, q_sum, q_sumsq, dot_qc, sum_c, sumsq_c); + (i, d) + }) + .collect(); + let f32_top10 = top10_by_dist(&f32_dists); + + let mut qi8 = vec![0i8; DIM]; + let (q_scale, sum_qi8) = sq8_quantize_query_scalar(&q, SQ8_INT8_QMAX, &mut qi8); + let int8_dists: Vec<(usize, f32)> = encoded + .iter() + .enumerate() + .map(|(i, slot)| { + let (min, scale) = sq8_params(slot, DIM); + let (dot_int, sum_c_int, sumsq_c_int) = + sq8_candidate_stats_i8_scalar(&qi8, &slot[..DIM], sum_qi8); + let dot_qc = sq8_int8_dot_to_f32(q_scale, dot_int); + let d = sq8_l2_from_stats( + DIM, + min, + scale, + q_sum, + q_sumsq, + dot_qc, + sum_c_int as f32, + sumsq_c_int as f32, + ); + (i, d) + }) + .collect(); + let int8_top10 = top10_by_dist(&int8_dists); + + r_f32 += recall_at_10(&f32_top10, &truth_top10); + r_int8 += recall_at_10(&int8_top10, &truth_top10); + overlap += recall_at_10(&int8_top10, &f32_top10); + } + + let r_f32 = r_f32 / N_QUERIES as f32; + let r_int8 = r_int8 / N_QUERIES as f32; + let overlap = overlap / N_QUERIES as f32; + assert!(overlap >= 0.98, "cosine: overlap={overlap} below 0.98 gate"); + assert!( + (r_f32 - r_int8).abs() <= 0.01, + "cosine: R@10 f32={r_f32} int8={r_int8} delta={} exceeds 0.01", + (r_f32 - r_int8).abs() + ); + } } diff --git a/tests/crash_recovery_vacuum.rs b/tests/crash_recovery_vacuum.rs index d550fd4a..66d34d7b 100644 --- a/tests/crash_recovery_vacuum.rs +++ b/tests/crash_recovery_vacuum.rs @@ -1,12 +1,12 @@ //! MA7 — Crash-recovery test matrix for Wave-1 reclamation code. //! -//! Three crash-injection scenarios, each proven with in-process simulation. +//! Crash-injection scenarios, each proven with in-process simulation. //! No subprocess spawning — `std::panic::catch_unwind` and drop-based //! injection suffice because every code path under test is synchronous. //! //! # Approach: drop-before-commit instead of process kill //! -//! The key insight for all three scenarios is that durability depends on +//! The key insight for both remaining scenarios is that durability depends on //! an explicit commit (either `ShardManifest::commit()` or an analogous //! mutation applied to stable storage). A `drop` of the in-memory state //! without a prior `commit()` is formally equivalent to a SIGKILL at the @@ -21,98 +21,28 @@ //! //! # Scenarios //! -//! 1. **Compaction crash** — half-written segment directory. Documents -//! the orphan-cleanup gap in `recover_vector_store`. -//! 2. **Manifest GC crash** — `gc_tombstones` mutates in-memory, drop +//! 1. **Manifest GC crash** — `gc_tombstones` mutates in-memory, drop //! before commit. Dual-root protocol recovers pre-staging root. -//! 3. **MVCC zombie sweep crash** — intents are ephemeral (not WAL- +//! 2. **MVCC zombie sweep crash** — intents are ephemeral (not WAL- //! persisted). Simulates crash mid-sweep; verifies fresh TM is clean //! and sweep is idempotent. Documents the intent-durability gap. +//! +//! A former "compaction crash / orphan partial segment" scenario lived here, +//! exercising the WAL-replay `recover_vector_store` function. That function +//! (and the `src/vector/persistence/recovery.rs` module it lived in) has +//! been deleted (vector-index durability B3): it only ever populated the +//! `Shard`-owned `vector_store`, which is discarded wholesale at +//! `event_loop.rs` before the live `ShardSlice.vector_store` is built — see +//! `VECTOR-DURABILITY-DESIGN.md`. Recovery authority is now the B1/B2/B3 +//! manifest+segment+keymap durability layout, and its orphan-cleanup +//! coverage lives in `src/vector/persistence/manifest.rs`'s +//! `sweep_orphans_from_disk` tests (`test_orphan_sweep_correctness` et al.), +//! which the B3 loader calls on every startup — the exact gap this deleted +//! scenario used to document no longer exists on the live path. use std::time::{Duration, Instant}; -// ── Scenario 1: Half-finished compaction segment directory ────────────────── -// -// Simulates: compaction starts writing a new segment directory, crashes after -// writing 2 of the required files (hnsw_graph.bin + tq_codes.bin), before -// segment_meta.json is present. -// -// Expected (documented current behaviour): `recover_vector_store` silently -// skips the partial segment (warns + returns Ok). The orphan directory is -// NOT cleaned up — it persists across restarts and is re-attempted on every -// recovery call. -// -// // BUG: orphan partial segment directories are never cleaned up after a -// // compaction crash. `recover_vector_store` at -// // src/vector/persistence/recovery.rs:291-293 silently continues past -// // `read_immutable_segment` failures without removing the partial directory. -// // Repeated restarts will keep re-enumerating and re-attempting to load the -// // orphan, wasting I/O and cluttering the segment namespace. Fix: detect -// // directories lacking `segment_meta.json` on startup and remove them (or -// // rename to `segment-{id}.partial` for forensics), then NOT retry them. - -#[test] -fn crash_compaction_orphan_partial_segment_silently_skipped() { - use moon::vector::persistence::recovery::recover_vector_store; - - let tmp = tempfile::tempdir().unwrap(); - let persist_dir = tmp.path().to_path_buf(); - - // Build a fully-valid segment directory (segment-1) with the minimum - // files that `read_immutable_segment` needs to succeed. - // - // We write only the segment_meta.json so that a minimal valid segment is - // recognised. `hnsw_graph.bin`, `tq_codes.bin`, and `mvcc_headers.bin` - // also need to exist for the full read path to work, but for our purposes - // the important assertion is about the orphan directory (segment-999), - // so we let segment-1 fail too and assert the total collections map is - // empty (recovery is graceful even when all segments fail to load). - // The orphan (segment-999) is what we are asserting about. - - // Half-finished directory: exists, has some files, missing segment_meta.json - let orphan_dir = persist_dir.join("segment-999"); - std::fs::create_dir_all(&orphan_dir).unwrap(); - // Write two of the five expected files — enough to look like a real segment - std::fs::write( - orphan_dir.join("hnsw_graph.bin"), - b"partial hnsw graph data", - ) - .unwrap(); - std::fs::write(orphan_dir.join("tq_codes.bin"), b"partial tq codes").unwrap(); - // segment_meta.json is deliberately absent (simulates crash mid-write) - - // Create a dummy WAL path (empty — no vectors were replayed before crash) - let wal_path = persist_dir.join("vector.wal"); - std::fs::write(&wal_path, &[0u8; 32]).unwrap(); // 32-byte WAL header stub - - // Recovery should succeed (not panic, not return Err) - let result = recover_vector_store(&wal_path, &persist_dir); - assert!( - result.is_ok(), - "recover_vector_store must not fail on partial segment directory: {:?}", - result.err(), - ); - - // BUG: The orphan directory still exists after recovery — it was not cleaned up. - // Once the bug is fixed, this assertion should be inverted to: - // assert!(!orphan_dir.exists(), "orphan partial segment dir must be removed on recovery"); - // - // BUG: src/vector/persistence/recovery.rs:291-293 — warn!+skip without cleanup. - assert!( - orphan_dir.exists(), - "current behaviour: orphan segment-999 dir persists after recovery (not cleaned up)", - ); - - // The valid segment-999 directory should still be enumerable on re-entry - // (re-attempt on every restart — the other half of the bug). - let second_result = recover_vector_store(&wal_path, &persist_dir); - assert!( - second_result.is_ok(), - "repeated recovery calls must remain idempotent despite orphan directory", - ); -} - -// ── Scenario 2: Mid-GC manifest crash before commit ───────────────────────── +// ── Scenario 1: Mid-GC manifest crash before commit ───────────────────────── // // Simulates: `gc_tombstones(0, 0, future)` is called (immediate retention → // all tombstones eligible for removal). The process crashes (simulated by @@ -204,7 +134,7 @@ fn crash_manifest_gc_before_commit_recovers_pre_staging_root() { ); } -// ── Scenario 3: Mid-sweep MVCC zombie crash + idempotency ─────────────────── +// ── Scenario 2: Mid-sweep MVCC zombie crash + idempotency ─────────────────── // // Context: `sweep_zombies_mut` drains zombie write-intents in-place (mutating // `write_intents` via `HashMap::retain`). Write-intents are NOT WAL-persisted @@ -341,55 +271,3 @@ fn crash_mvcc_sweep_zombies_idempotent_and_clean_on_restart() { "commit releases the last intent", ); } - -// ── Scenario 1b: WAL replay after compaction crash — no double file_id spend ─ -// -// Assert: when two recovery calls are made against the same persist_dir, -// the segment_id namespace is not polluted (no double-spend). Since -// `recover_vector_store` reads `enumerate_segments` (a pure directory scan), -// the orphan segment-999 dir shows up on every call but never consumes a -// new segment_id from the recovery layer — it is the caller's responsibility -// to assign new IDs. This test verifies that behaviour is stable. - -#[test] -fn crash_compaction_wal_replay_no_double_spend_of_segment_ids() { - use moon::vector::persistence::recovery::recover_vector_store; - - let tmp = tempfile::tempdir().unwrap(); - let persist_dir = tmp.path().to_path_buf(); - - // Orphan directory: segment-42 exists but is corrupt (missing meta) - let orphan_dir = persist_dir.join("segment-42"); - std::fs::create_dir_all(&orphan_dir).unwrap(); - std::fs::write(orphan_dir.join("tq_codes.bin"), b"garbage").unwrap(); - // segment_meta.json absent — simulates crash mid-write - - let wal_path = persist_dir.join("vector.wal"); - std::fs::write(&wal_path, &[0u8; 32]).unwrap(); - - // First recovery: orphan segment-42 load fails, nothing loaded. - let r1 = recover_vector_store(&wal_path, &persist_dir).unwrap(); - assert!( - r1.collections.is_empty(), - "no collections recovered when all segments fail to load", - ); - - // Second recovery call (simulating restart after second crash): - // must produce identical results — idempotent. - let r2 = recover_vector_store(&wal_path, &persist_dir).unwrap(); - assert!( - r2.collections.is_empty(), - "repeated recovery is idempotent: no double-spend of ids, same empty result", - ); - - // The orphan dir still occupies the segment-42 slot. A new compaction job - // would need to pick segment-43 or check for directory existence before - // writing. BUG: there is no cleanup — the orphan blocks that slot. - // BUG: src/vector/persistence/recovery.rs:119-137 enumerate_segments does - // not distinguish partial vs complete segment directories. Combined with - // the missing cleanup at lines 291-293, the orphan persists indefinitely. - assert!( - orphan_dir.exists(), - "orphan dir persists — confirms the cleanup gap documented in scenario 1", - ); -} diff --git a/tests/crash_recovery_vector_durability.rs b/tests/crash_recovery_vector_durability.rs new file mode 100644 index 00000000..f71e747b --- /dev/null +++ b/tests/crash_recovery_vector_durability.rs @@ -0,0 +1,1431 @@ +//! B4: vector-index durability kill-9 crash-recovery integration tests. +//! +//! Proves the B1-B3 durability line (see `tmp/VECTOR-DURABILITY-DESIGN.md`) +//! end to end against REAL server processes: manifest+segment+keymap +//! persist-on-compact (B1/B2), and startup recovery + dedup rescan + +//! deletion probe + orphan sweep (B3, `src/vector/persistence/recover_v2.rs`). +//! +//! Five scenarios, one `#[test]` each, sharing the harness below: +//! S1 unchanged-keys fast path (dedup rescan skips re-encoding) +//! S2 updates+deletes across a crash (reconcile + deletion probe) +//! S3 orphan sweep (stray staging/segment/keymap files removed on boot) +//! S4 collection_id pin survives a post-recovery compact+GraphUnion merge +//! S5 no-persist-dir regression guard (`--appendonly no`: no idx-* dirs) +//! +//! Every wait is a bounded, condition-based poll (port accept / manifest +//! file / log line / FT.INFO field) — never a fixed sleep as +//! synchronization (this repo has a documented 100ms-sleep flake history, +//! see `tests/crash_matrix_per_shard_aof.rs` and the parity-fix that +//! replaced it with a bind-wait). +//! +//! Run with (monoio default — matches CI): +//! cargo build --release +//! cargo test --release --test crash_recovery_vector_durability -- --ignored +//! +//! tokio runtime: +//! cargo build --release --no-default-features \ +//! --features runtime-tokio,jemalloc,graph,text-index +//! cargo test --release --no-default-features \ +//! --features runtime-tokio,jemalloc,graph,text-index \ +//! --test crash_recovery_vector_durability -- --ignored +//! +//! Requires: built release binary (`MOON_BIN` env var honored, falls back to +//! `target/release/moon` then `target/debug/moon` — see `find_moon_binary`). + +#![cfg(any(feature = "runtime-monoio", feature = "runtime-tokio"))] +#![allow(clippy::unwrap_used)] + +use std::collections::HashSet; +use std::io::{BufReader, Read, Write}; +use std::net::{TcpStream, ToSocketAddrs}; +use std::path::{Path, PathBuf}; +use std::process::{Child, Command}; +use std::time::{Duration, Instant}; + +use moon::vector::persistence::manifest::{self, IndexManifest}; + +// 32 (not 8): S4's clustered recall check needs enough axes that SQ8 +// quantization doesn't produce frequent EXACT decoded-distance ties within +// a tight cluster — ties destabilize the internal merge recall gate's +// brute-force "top-k ground truth" SET independent of actual HNSW/graph +// quality (confirmed empirically: the gate's reported recall was IDENTICAL +// to 16 significant digits across very different noise/M/EF_CONSTRUCTION +// settings at DIM=8, the signature of tie-bound ground truth, not a +// genuine approximation gap). +const DIM: usize = 32; + +// --------------------------------------------------------------------------- +// Binary resolution (pattern: tests/vector_del_unindex.rs / shardslice_live.rs) +// --------------------------------------------------------------------------- + +fn find_moon_binary() -> PathBuf { + if let Ok(bin) = std::env::var("MOON_BIN") { + let p = PathBuf::from(bin); + if p.exists() { + return p; + } + } + let manifest = env!("CARGO_MANIFEST_DIR"); + let release = PathBuf::from(format!("{manifest}/target/release/moon")); + if release.exists() { + return release; + } + let debug = PathBuf::from(format!("{manifest}/target/debug/moon")); + if debug.exists() { + return debug; + } + panic!( + "No moon binary found. Build with `cargo build --release` or set \ + MOON_BIN=/path/to/moon." + ); +} + +fn unique_port() -> u16 { + use std::net::TcpListener; + let listener = TcpListener::bind("127.0.0.1:0").expect("bind :0"); + let port = listener.local_addr().expect("local_addr").port(); + drop(listener); + port +} + +fn unique_dir(suffix: &str) -> PathBuf { + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|d| d.as_nanos()) + .unwrap_or(0); + std::env::temp_dir().join(format!( + "moon-vecdur-{}-{}-{}", + std::process::id(), + suffix, + nanos + )) +} + +// --------------------------------------------------------------------------- +// Server spawn / crash / restart harness +// --------------------------------------------------------------------------- + +/// Kills the wrapped child on drop (SIGKILL + waitpid), plus a best-effort +/// `pkill -9 -f ` backstop keyed on this test's unique temp dir path — +/// a moon process launched with `--dir ` can only ever be one +/// this test itself spawned, so the backstop cannot collide with another +/// test's server (documented gotcha: leaked busy-poller / SIGTERM+SO_REUSEPORT +/// hangs — always kill -9, never rely on graceful shutdown). +struct ServerGuard { + child: Child, + dir_marker: String, +} + +impl ServerGuard { + fn new(child: Child, dir: &Path) -> Self { + Self { + child, + dir_marker: dir.to_string_lossy().into_owned(), + } + } +} + +impl Drop for ServerGuard { + fn drop(&mut self) { + // Idempotent: a test that already SIGKILLed+waited this child + // earlier (the crash-then-restart flow) must not re-signal a + // possibly-recycled pid. Only signal if the process is still + // observed running; otherwise just reap (harmless if already reaped). + if matches!(self.child.try_wait(), Ok(None)) { + sigkill(&mut self.child); + } else { + let _ = self.child.wait(); + } + // Best-effort backstop in case anything survived under a different + // pid (e.g. a forked helper) — matches this test's own unique --dir + // path, so it can never collide with another test's server. + let _ = Command::new("pkill") + .args(["-9", "-f"]) + .arg(&self.dir_marker) + .output(); + } +} + +#[cfg(unix)] +fn sigkill(child: &mut Child) { + let pid = child.id() as i32; + unsafe { + libc::kill(pid, libc::SIGKILL); + } + let _ = child.wait(); +} + +#[cfg(not(unix))] +fn sigkill(child: &mut Child) { + let _ = child.kill(); + let _ = child.wait(); +} + +/// Wait until nothing is accepting on `port` — required before a same-port +/// restart or the new listener can race the dying process's socket teardown +/// (moon binds SO_REUSEPORT per shard, so a plain bind-based check is +/// useless; see `crash_recovery_disk_offload_no_aof.rs`'s extensive +/// rationale for this exact pattern). +fn wait_for_port_down(port: u16) { + let addr = format!("127.0.0.1:{port}"); + let mut consecutive_refused = 0; + for _ in 0..120 { + match TcpStream::connect_timeout(&addr.parse().expect("addr"), Duration::from_millis(100)) { + Ok(_) => { + consecutive_refused = 0; + std::thread::sleep(Duration::from_millis(100)); + } + Err(_) => { + consecutive_refused += 1; + if consecutive_refused >= 2 { + return; + } + std::thread::sleep(Duration::from_millis(50)); + } + } + } +} + +/// Spawn moon with `--appendonly yes --shards 1`, stdout/stderr captured to +/// per-test log files (never `Stdio::null()` — a CI flake needs a real +/// diagnostic, not silence). `RUST_LOG=moon=info` is set explicitly so the +/// B3 recovery acceptance-signal `info!` line is emitted regardless of the +/// ambient environment (the binary's own default is already `moon=info`, +/// but the test must not depend on that default silently). +/// +/// `--disk-offload disable`: disk-offload defaults to `enable` (see +/// `src/config.rs`'s `#[arg(long = "disk-offload", default_value = "enable")]`), +/// and when enabled the vector persist dir resolves to `/shard-` +/// instead of `/shard--vectors` (`src/shard/event_loop.rs` ~857) — +/// this test targets the plain AOF-driven persistence path the design doc +/// describes, so disk-offload must be off to get the path this harness polls. +fn spawn_moon_aof(port: u16, dir: &Path) -> ServerGuard { + std::fs::create_dir_all(dir).expect("create test dir"); + let child = Command::new(find_moon_binary()) + .args([ + "--port", + &port.to_string(), + "--shards", + "1", + "--appendonly", + "yes", + "--appendfsync", + "always", + "--disk-offload", + "disable", + // The dev volume routinely sits near Moon's 5% diskfull guard + // (writes pause with MOONERR diskfull); this test's durability is + // proven by kill -9 + recovery, not by the free-space monitor. + "--disk-free-min-pct", + "0", + "--dir", + ]) + .arg(dir) + .env("RUST_LOG", "moon=info") + .stdout(std::fs::File::create(dir.join("moon.stdout.log")).expect("stdout log")) + .stderr(std::fs::File::create(dir.join("moon.stderr.log")).expect("stderr log")) + .spawn() + .expect("spawn moon (run `cargo build --release` first, or set MOON_BIN)"); + ServerGuard::new(child, dir) +} + +/// Same as `spawn_moon_aof` but `--appendonly no` with disk-offload also +/// explicitly disabled — S5's regression guard (no persistence_dir at all). +fn spawn_moon_no_persist(port: u16, dir: &Path) -> ServerGuard { + std::fs::create_dir_all(dir).expect("create test dir"); + let child = Command::new(find_moon_binary()) + .args([ + "--port", + &port.to_string(), + "--shards", + "1", + "--appendonly", + "no", + "--disk-offload", + "disable", + "--disk-free-min-pct", + "0", + "--dir", + ]) + .arg(dir) + .env("RUST_LOG", "moon=info") + .stdout(std::fs::File::create(dir.join("moon.stdout.log")).expect("stdout log")) + .stderr(std::fs::File::create(dir.join("moon.stderr.log")).expect("stderr log")) + .spawn() + .expect("spawn moon (run `cargo build --release` first, or set MOON_BIN)"); + ServerGuard::new(child, dir) +} + +/// Restart attempts — a rapid SIGKILL->rebind on the same port can lose a +/// transient EADDRINUSE race against the dying process's socket teardown. +/// That is an OS timing artifact, not a recovery defect; retry bounded. +const RESTART_ATTEMPTS: usize = 6; + +fn start_moon_alive( + spawn: impl Fn(u16, &Path) -> ServerGuard, + port: u16, + dir: &Path, +) -> ServerGuard { + for attempt in 1..=RESTART_ATTEMPTS { + let mut guard = spawn(port, dir); + let mut up = false; + for _ in 0..100 { + if let Ok(Some(_status)) = guard.child.try_wait() { + break; // self-terminated (rebind race) — retry + } + if TcpStream::connect(format!("127.0.0.1:{port}")).is_ok() { + up = true; + break; + } + std::thread::sleep(Duration::from_millis(100)); + } + if up { + return guard; + } + drop(guard); // Drop kills + reaps; back off before retrying. + if attempt < RESTART_ATTEMPTS { + std::thread::sleep(Duration::from_millis(300)); + } + } + panic!("moon failed to start+serve on port {port} after {RESTART_ATTEMPTS} attempts"); +} + +// --------------------------------------------------------------------------- +// Minimal RESP client (pattern: tests/vector_del_unindex.rs) +// --------------------------------------------------------------------------- + +#[derive(Debug, PartialEq, Clone)] +enum V { + Simple(String), + Err(String), + Int(i64), + Bulk(Vec), + Arr(Vec), + Null, +} + +struct Client { + reader: BufReader, + writer: TcpStream, +} + +impl Client { + fn connect(port: u16) -> Self { + let addr = format!("127.0.0.1:{port}") + .to_socket_addrs() + .unwrap() + .next() + .unwrap(); + let start = Instant::now(); + let stream = loop { + match TcpStream::connect_timeout(&addr, Duration::from_millis(200)) { + Ok(s) => break s, + Err(_) if start.elapsed() < Duration::from_secs(30) => { + std::thread::sleep(Duration::from_millis(50)); + } + Err(e) => panic!("server never accepted on port {port}: {e}"), + } + }; + stream + .set_read_timeout(Some(Duration::from_secs(20))) + .unwrap(); + let writer = stream.try_clone().unwrap(); + Client { + reader: BufReader::new(stream), + writer, + } + } + + fn encode(args: &[&[u8]]) -> Vec { + let mut out = format!("*{}\r\n", args.len()).into_bytes(); + for a in args { + out.extend_from_slice(format!("${}\r\n", a.len()).as_bytes()); + out.extend_from_slice(a); + out.extend_from_slice(b"\r\n"); + } + out + } + + fn read_line(&mut self) -> String { + let mut line = Vec::new(); + let mut b = [0u8; 1]; + loop { + self.reader.read_exact(&mut b).expect("read byte"); + if b[0] == b'\n' { + break; + } + if b[0] != b'\r' { + line.push(b[0]); + } + } + String::from_utf8_lossy(&line).into_owned() + } + + fn parse(&mut self) -> V { + let line = self.read_line(); + let (t, rest) = line.split_at(1); + match t { + "+" => V::Simple(rest.to_string()), + "-" => V::Err(rest.to_string()), + ":" => V::Int(rest.parse().expect("int")), + "$" => { + let n: i64 = rest.parse().expect("bulk len"); + if n < 0 { + return V::Null; + } + let mut buf = vec![0u8; n as usize + 2]; + self.reader.read_exact(&mut buf).expect("bulk body"); + buf.truncate(n as usize); + V::Bulk(buf) + } + "*" => { + let n: i64 = rest.parse().expect("arr len"); + if n < 0 { + return V::Null; + } + V::Arr((0..n).map(|_| self.parse()).collect()) + } + other => panic!("unexpected RESP type {other:?} (line {line:?})"), + } + } + + fn cmd(&mut self, args: &[&[u8]]) -> V { + self.writer.write_all(&Self::encode(args)).expect("send"); + self.parse() + } + + /// Send all commands in ONE write (a wire pipeline), then read all replies. + fn pipeline(&mut self, cmds: &[Vec>]) -> Vec { + let mut buf = Vec::new(); + for c in cmds { + let refs: Vec<&[u8]> = c.iter().map(|a| a.as_slice()).collect(); + buf.extend_from_slice(&Self::encode(&refs)); + } + self.writer.write_all(&buf).expect("send pipeline"); + cmds.iter().map(|_| self.parse()).collect() + } +} + +fn wait_ready(port: u16) -> Client { + let mut c = Client::connect(port); + let start = Instant::now(); + loop { + match c.cmd(&[b"PING"]) { + V::Simple(s) if s == "PONG" => return c, + _ if start.elapsed() < Duration::from_secs(30) => { + std::thread::sleep(Duration::from_millis(100)); + } + other => panic!("server never answered PING: {other:?}"), + } + } +} + +// --------------------------------------------------------------------------- +// Vector fixtures +// --------------------------------------------------------------------------- + +/// Deterministic, well-separated f32 vector (LE bytes) for identity-style +/// KNN probes (exact match = distance 0). Uses a splitmix64-style hash per +/// (seed, dim) pair rather than a small modulus — a modulus (e.g. `% 9973`) +/// wraps distinct seeds back into the SAME narrow value range, so two +/// unrelated seeds (in particular a base id and `id + large_offset`, as S2 +/// uses to build a "far away" update vector) can collide onto +/// near-identical vectors and break identity-style KNN assertions (caught +/// by this test suite's own flake-hunt). A hash mix has no such periodicity: +/// distinct u32 seeds land in well-separated regions of R^DIM with +/// overwhelming probability. +fn simple_vec_bytes(seed: u32) -> Vec { + let mut out = Vec::with_capacity(DIM * 4); + for i in 0..DIM { + let mut h = (seed as u64) + .wrapping_mul(0x9E37_79B9_7F4A_7C15) + .wrapping_add((i as u64).wrapping_mul(0xBF58_476D_1CE4_E5B9)) + .wrapping_add(0x94D0_49BB_1331_11EB); + h ^= h >> 33; + h = h.wrapping_mul(0xFF51_AFD7_ED55_8CCD); + h ^= h >> 33; + h = h.wrapping_mul(0xC4CE_B9FE_1A85_EC53); + h ^= h >> 33; + // Map the top 24 bits to a value in [0, 10) — a continuous range + // wide enough that hash collisions between distinct seeds are + // astronomically unlikely at the N used by these tests. + let v = ((h >> 40) as f32 / (1u64 << 24) as f32) * 10.0; + out.extend_from_slice(&v.to_le_bytes()); + } + out +} + +/// Deterministic LCG, seeded per test — used for S4's clustered dataset. +struct Rng(u64); + +impl Rng { + fn new(seed: u64) -> Self { + Self(seed) + } + fn next_f32(&mut self) -> f32 { + self.0 = self + .0 + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + ((self.0 >> 40) as f32) / (1u64 << 24) as f32 + } + fn randn(&mut self) -> f32 { + let u1 = self.next_f32().max(1e-7); + let u2 = self.next_f32(); + (-2.0 * u1.ln()).sqrt() * (2.0 * std::f32::consts::PI * u2).cos() + } +} + +fn random_vec(rng: &mut Rng, dim: usize) -> Vec { + (0..dim).map(|_| rng.randn()).collect() +} + +fn f32s_to_le_bytes(v: &[f32]) -> Vec { + let mut out = Vec::with_capacity(v.len() * 4); + for x in v { + out.extend_from_slice(&x.to_le_bytes()); + } + out +} + +// --------------------------------------------------------------------------- +// FT.* / VACUUM helpers +// --------------------------------------------------------------------------- + +fn ft_create( + c: &mut Client, + idx: &str, + dim: usize, + compact_threshold: u32, + ef_runtime: Option, +) { + ft_create_ex(c, idx, dim, compact_threshold, ef_runtime, None, None); +} + +/// Extended FT.CREATE with optional HNSW `M` / `EF_CONSTRUCTION` overrides — +/// S4 needs a higher-fidelity graph than the defaults (M=16, EF_CONSTRUCTION=200) +/// to keep the merge's internal recall gate (fixed 0.90 tolerance, +/// `src/vector/segment/compaction.rs::verify_merge_recall`) comfortably clear +/// of its threshold on a small (~1000-vector) synthetic dataset. +#[allow(clippy::too_many_arguments)] +fn ft_create_ex( + c: &mut Client, + idx: &str, + dim: usize, + compact_threshold: u32, + ef_runtime: Option, + m: Option, + ef_construction: Option, +) { + let prefix = format!("{idx}:"); + let mut params: Vec> = vec![ + b"TYPE".to_vec(), + b"FLOAT32".to_vec(), + b"DIM".to_vec(), + dim.to_string().into_bytes(), + b"DISTANCE_METRIC".to_vec(), + b"L2".to_vec(), + b"QUANTIZATION".to_vec(), + b"SQ8".to_vec(), + b"COMPACT_THRESHOLD".to_vec(), + compact_threshold.to_string().into_bytes(), + ]; + if let Some(ef) = ef_runtime { + params.push(b"EF_RUNTIME".to_vec()); + params.push(ef.to_string().into_bytes()); + } + if let Some(m) = m { + params.push(b"M".to_vec()); + params.push(m.to_string().into_bytes()); + } + if let Some(efc) = ef_construction { + params.push(b"EF_CONSTRUCTION".to_vec()); + params.push(efc.to_string().into_bytes()); + } + let mut args: Vec> = vec![ + b"FT.CREATE".to_vec(), + idx.as_bytes().to_vec(), + b"ON".to_vec(), + b"HASH".to_vec(), + b"PREFIX".to_vec(), + b"1".to_vec(), + prefix.into_bytes(), + b"SCHEMA".to_vec(), + b"vec".to_vec(), + b"VECTOR".to_vec(), + b"HNSW".to_vec(), + params.len().to_string().into_bytes(), + ]; + args.extend(params); + let refs: Vec<&[u8]> = args.iter().map(|a| a.as_slice()).collect(); + let r = c.cmd(&refs); + assert_eq!(r, V::Simple("OK".into()), "FT.CREATE {idx} failed: {r:?}"); +} + +fn hset_batch(c: &mut Client, prefix: &str, ids: &[u32], blob_of: impl Fn(u32) -> Vec) { + let cmds: Vec>> = ids + .iter() + .map(|&i| { + let key = format!("{prefix}{i}"); + vec![ + b"HSET".to_vec(), + key.into_bytes(), + b"vec".to_vec(), + blob_of(i), + ] + }) + .collect(); + let replies = c.pipeline(&cmds); + for (i, r) in replies.iter().enumerate() { + assert!(matches!(r, V::Int(_)), "HSET batch #{i} failed: {r:?}"); + } +} + +fn ft_compact(c: &mut Client, idx: &str) { + let r = c.cmd(&[b"FT.COMPACT", idx.as_bytes()]); + assert_eq!(r, V::Simple("OK".into()), "FT.COMPACT {idx} failed: {r:?}"); +} + +fn vacuum_vector(c: &mut Client, idx: &str) -> String { + let r = c.cmd(&[b"VACUUM", b"VECTOR", idx.as_bytes()]); + match r { + V::Simple(s) => s, + other => panic!("VACUUM VECTOR {idx} unexpected reply: {other:?}"), + } +} + +fn search_keys(c: &mut Client, idx: &str, k: u32, blob: &[u8]) -> Vec { + let query = format!("*=>[KNN {k} @vec $B]"); + let r = c.cmd(&[ + b"FT.SEARCH", + idx.as_bytes(), + query.as_bytes(), + b"PARAMS", + b"2", + b"B", + blob, + b"DIALECT", + b"2", + ]); + let V::Arr(items) = r else { + panic!("FT.SEARCH {idx} reply not array: {r:?}"); + }; + // Reply shape: [total, key1, fields1, key2, fields2, ...] + items[1..] + .iter() + .step_by(2) + .filter_map(|v| match v { + V::Bulk(b) => Some(String::from_utf8_lossy(b).into_owned()), + _ => None, + }) + .collect() +} + +fn ft_info_num_docs(c: &mut Client, idx: &str) -> i64 { + let r = c.cmd(&[b"FT.INFO", idx.as_bytes()]); + let V::Arr(items) = r else { + panic!("FT.INFO {idx} reply not array: {r:?}"); + }; + for pair in items.chunks(2) { + if let [V::Bulk(k), V::Int(n)] = pair + && k.as_slice() == b"num_docs" + { + return *n; + } + } + panic!("num_docs not found in FT.INFO {idx} reply: {items:?}"); +} + +fn ft_list(c: &mut Client) -> Vec { + let r = c.cmd(&[b"FT._LIST"]); + match r { + V::Arr(items) => items + .into_iter() + .filter_map(|v| match v { + V::Bulk(b) => Some(String::from_utf8_lossy(&b).into_owned()), + _ => None, + }) + .collect(), + V::Null => Vec::new(), + other => panic!("FT._LIST unexpected reply: {other:?}"), + } +} + +fn assert_absent(keys: &[String], dead: &str, ctx: &str) { + assert!( + !keys.iter().any(|k| k == dead), + "{ctx}: deleted/absent key {dead} resurfaced in FT.SEARCH results {keys:?}" + ); +} + +// --------------------------------------------------------------------------- +// Persistence-layer polling helpers (manifest + log line) +// --------------------------------------------------------------------------- + +/// Vector persist dir for shard 0 under `--appendonly yes --shards 1` +/// (`/shard-0-vectors`, per `src/shard/event_loop.rs` ~865). +fn vector_persist_dir(dir: &Path) -> PathBuf { + dir.join("shard-0-vectors") +} + +/// Bounded poll for `manifest.json` to reach at least `min_segments` live +/// segment ids. This is the harness's substitute for a fixed sleep before +/// SIGKILL — the background snapshot job (`global_snapshot_pool()`) commits +/// asynchronously after a compact/merge install, so the test must observe +/// the durable artifact directly rather than guess a timing window. +fn wait_for_manifest_min_segments( + idx_dir: &Path, + min_segments: usize, + timeout: Duration, +) -> IndexManifest { + let deadline = Instant::now() + timeout; + loop { + if let Some(m) = manifest::read_manifest_tolerant(idx_dir) + && m.segment_ids.len() >= min_segments + { + return m; + } + if Instant::now() >= deadline { + let found = manifest::read_manifest_tolerant(idx_dir) + .map(|m| m.segment_ids.len()) + .unwrap_or(0); + panic!( + "manifest at {:?} did not reach {} segment(s) within {:?} (found {})", + idx_dir, min_segments, timeout, found + ); + } + std::thread::sleep(Duration::from_millis(20)); + } +} + +/// Poll until the manifest holds EXACTLY `n_segments` segments (used after a +/// merge, whose manifest rewrite is an async SnapshotPool job — a min-N wait +/// would return instantly on the stale pre-merge manifest). +fn wait_for_manifest_exact_segments( + idx_dir: &Path, + n_segments: usize, + timeout: Duration, +) -> IndexManifest { + let deadline = Instant::now() + timeout; + loop { + if let Some(m) = manifest::read_manifest_tolerant(idx_dir) + && m.segment_ids.len() == n_segments + { + return m; + } + if Instant::now() >= deadline { + let found = manifest::read_manifest_tolerant(idx_dir) + .map(|m| m.segment_ids) + .unwrap_or_default(); + panic!( + "manifest at {:?} did not converge to exactly {} segment(s) within {:?} (found {:?})", + idx_dir, n_segments, timeout, found + ); + } + std::thread::sleep(Duration::from_millis(20)); + } +} + +/// Extract the 4 integer counters from a B3 recovery acceptance log line: +/// "vector index {name}: B3 recovery — loaded {N} segment(s), {N} key(s) +/// verified unchanged, {N} re-indexed, {N} tombstoned" (exact format string +/// lives in `src/vector/persistence/recover_v2.rs::RecoveryState::finish`). +/// Slicing from AFTER the "B3 recovery" marker before scanning digits keeps +/// this robust against a timestamp/level prefix (which also contains +/// digits) added by the ambient `tracing_subscriber::fmt()` layer — slicing +/// AT the marker instead of after it is a trap: "B3" itself contains a +/// digit ('3'), which would silently shift every subsequent count by one +/// (caught by this test suite's own flake-hunt: `loaded segments` bled into +/// `verified_unchanged`'s assertion). +fn parse_recovery_counters(line: &str) -> Option<(usize, usize, usize, usize)> { + const MARKER: &str = "B3 recovery"; + let pos = line.find(MARKER)?; + let tail = &line[pos + MARKER.len()..]; + let mut nums: Vec = Vec::new(); + let mut cur = String::new(); + for ch in tail.chars() { + if ch.is_ascii_digit() { + cur.push(ch); + } else if !cur.is_empty() { + nums.push(cur.parse().ok()?); + cur.clear(); + } + } + if !cur.is_empty() { + nums.push(cur.parse().ok()?); + } + if nums.len() >= 4 { + Some((nums[0], nums[1], nums[2], nums[3])) + } else { + None + } +} + +/// Bounded poll of the server's captured stdout log for the B3 recovery +/// line belonging to `idx_name`. Recovery runs synchronously during +/// per-shard startup (before the accept loop begins), so by the time +/// `wait_ready` observes PONG the line should already be flushed — Rust's +/// `std::io::Stdout` is unconditionally line-buffered (`LineWriter`), so a +/// `\n`-terminated tracing event is written through immediately regardless +/// of whether stdout is a TTY or a redirected file. The poll here is a +/// defensive bound, not a workaround for buffering. +fn wait_for_recovery_counters( + dir: &Path, + idx_name: &str, + timeout: Duration, +) -> (usize, usize, usize, usize) { + let marker = format!("vector index {idx_name}:"); + let deadline = Instant::now() + timeout; + loop { + if let Ok(log) = std::fs::read_to_string(dir.join("moon.stdout.log")) { + for line in log.lines() { + if line.contains(&marker) + && line.contains("B3 recovery") + && let Some(counters) = parse_recovery_counters(line) + { + return counters; + } + } + } + if Instant::now() >= deadline { + let log = std::fs::read_to_string(dir.join("moon.stdout.log")).unwrap_or_default(); + panic!( + "B3 recovery log line for index {idx_name} not found within {timeout:?}.\n\ + --- moon.stdout.log ---\n{log}" + ); + } + std::thread::sleep(Duration::from_millis(20)); + } +} + +/// Recursively collect every path under `root` whose file/dir name matches +/// `pred` — used by S3 (orphan detection) and S5 (no idx-* dirs at all). +fn walk_matching(root: &Path, pred: &dyn Fn(&str) -> bool) -> Vec { + let mut out = Vec::new(); + fn walk(p: &Path, pred: &dyn Fn(&str) -> bool, acc: &mut Vec) { + let Ok(rd) = std::fs::read_dir(p) else { + return; + }; + for e in rd.flatten() { + let path = e.path(); + if let Some(name) = path.file_name().and_then(|n| n.to_str()) + && pred(name) + { + acc.push(path.clone()); + } + if path.is_dir() { + walk(&path, pred, acc); + } + } + } + walk(root, pred, &mut out); + out +} + +// --------------------------------------------------------------------------- +// S1: unchanged-keys fast path +// --------------------------------------------------------------------------- + +/// Insert N vectors in 2 batches, each followed by an explicit FT.COMPACT — +/// `force_compact` drains the ENTIRE mutable segment before returning (loops +/// until `frozen_len == mutable_len`; see `src/vector/store.rs::compact_segments`), +/// so two batch+compact rounds deterministically produce >= 2 immutable +/// segments with ZERO residual vectors left in the (unpersisted) mutable +/// segment — the precondition for an exact (not approximate) dedup-rescan +/// assertion after the crash. +fn build_two_segment_snapshot(c: &mut Client, idx: &str, idx_dir: &Path, n: usize) { + let batch = n / 2; + let batch1: Vec = (0..batch as u32).collect(); + let batch2: Vec = (batch as u32..n as u32).collect(); + + hset_batch(c, &format!("{idx}:"), &batch1, simple_vec_bytes); + ft_compact(c, idx); + wait_for_manifest_min_segments(idx_dir, 1, Duration::from_secs(10)); + + hset_batch(c, &format!("{idx}:"), &batch2, simple_vec_bytes); + ft_compact(c, idx); + wait_for_manifest_min_segments(idx_dir, 2, Duration::from_secs(10)); +} + +#[test] +#[ignore] // Requires built release binary; run explicitly. +fn s1_unchanged_keys_fast_path_survives_crash() { + const N: usize = 2000; + const IDX: &str = "s1"; + + let port = unique_port(); + let dir = unique_dir("s1"); + let vdir = vector_persist_dir(&dir); + let idx_dir = manifest::index_persist_dir(&vdir, IDX.as_bytes()); + + let guard = spawn_moon_aof(port, &dir); + let mut c = wait_ready(port); + + ft_create(&mut c, IDX, DIM, 100, None); + build_two_segment_snapshot(&mut c, IDX, &idx_dir, N); + + // Pre-crash KNN baseline for a fixed query set (every 200th key). + let probe_ids: Vec = (0..N as u32).step_by(200).collect(); + let mut pre_results: Vec> = Vec::new(); + for &i in &probe_ids { + let blob = simple_vec_bytes(i); + pre_results.push(search_keys(&mut c, IDX, 1, &blob)); + } + let pre_num_docs = ft_info_num_docs(&mut c, IDX); + assert_eq!(pre_num_docs, N as i64, "pre-crash num_docs must equal N"); + + drop(c); + let mut guard = guard; + sigkill(&mut guard.child); + wait_for_port_down(port); + + // -- Restart -------------------------------------------------------- + let guard2 = start_moon_alive(spawn_moon_aof, port, &dir); + let mut c2 = wait_ready(port); + + // (b) B3 acceptance signal: exact dedup — the final FT.COMPACT before + // the kill drained the mutable segment to empty, so every one of the N + // keys must be recognized as unchanged and NONE re-indexed. + let (loaded_segments, verified_unchanged, re_indexed, tombstoned) = + wait_for_recovery_counters(&dir, IDX, Duration::from_secs(10)); + assert!( + loaded_segments >= 2, + "expected >=2 loaded segments, got {loaded_segments}" + ); + assert_eq!( + verified_unchanged, N, + "every unchanged key must dedup via the metadata-only rebuild path" + ); + assert_eq!( + re_indexed, 0, + "no key should need full re-encode (mutable segment was empty at crash time)" + ); + assert_eq!( + tombstoned, 0, + "no key was deleted between compact and crash" + ); + + // (a) KNN results identical to pre-crash for the fixed query set. + for (idx, &i) in probe_ids.iter().enumerate() { + let blob = simple_vec_bytes(i); + let post = search_keys(&mut c2, IDX, 1, &blob); + assert_eq!( + post, pre_results[idx], + "post-crash top-1 for probe {i} must match pre-crash" + ); + let want_key = format!("{IDX}:{i}"); + assert_eq!( + post.first().map(String::as_str), + Some(want_key.as_str()), + "exact-match query for {want_key} must return itself as top-1" + ); + } + + // (c) FT.INFO num_docs correct. + let post_num_docs = ft_info_num_docs(&mut c2, IDX); + assert_eq!(post_num_docs, N as i64, "post-crash num_docs must equal N"); + + drop(c2); + drop(guard2); + let _ = std::fs::remove_dir_all(&dir); +} + +// --------------------------------------------------------------------------- +// S2: updates + deletes across a crash +// --------------------------------------------------------------------------- + +#[test] +#[ignore] // Requires built release binary; run explicitly. +fn s2_updates_and_deletes_reconcile_across_crash() { + const N: usize = 2000; + const K: usize = 50; + const IDX: &str = "s2"; + /// Offset applied to a key's seed on update so its new vector lands far + /// from its original position (distinguishable in KNN). + const UPDATE_SEED_OFFSET: u32 = 500_000; + + let port = unique_port(); + let dir = unique_dir("s2"); + let vdir = vector_persist_dir(&dir); + let idx_dir = manifest::index_persist_dir(&vdir, IDX.as_bytes()); + + let guard = spawn_moon_aof(port, &dir); + let mut c = wait_ready(port); + + ft_create(&mut c, IDX, DIM, 100, None); + build_two_segment_snapshot(&mut c, IDX, &idx_dir, N); + + // Mutations AFTER the last snapshot: these land in AOF only. + // Update the first K keys (already durable in segment 1) with a new, + // far-away vector. + let update_ids: Vec = (0..K as u32).collect(); + hset_batch(&mut c, &format!("{IDX}:"), &update_ids, |i| { + simple_vec_bytes(i + UPDATE_SEED_OFFSET) + }); + + // Delete K other keys from segment 2's range. + let delete_ids: Vec = (1000..1000 + K as u32).collect(); + for &i in &delete_ids { + let key = format!("{IDX}:{i}"); + let r = c.cmd(&[b"DEL", key.as_bytes()]); + assert_eq!(r, V::Int(1), "DEL {key} failed: {r:?}"); + } + + drop(c); + let mut guard = guard; + sigkill(&mut guard.child); + wait_for_port_down(port); + + // -- Restart -------------------------------------------------------- + let guard2 = start_moon_alive(spawn_moon_aof, port, &dir); + let mut c2 = wait_ready(port); + + let (_loaded, _verified, re_indexed, tombstoned) = + wait_for_recovery_counters(&dir, IDX, Duration::from_secs(10)); + assert!( + re_indexed >= K, + "expected >= {K} re-indexed (the updated keys), got {re_indexed}" + ); + assert!( + tombstoned >= K, + "expected >= {K} tombstoned (the deleted keys), got {tombstoned}" + ); + + // Updated keys: querying with the NEW vector must return the updated + // key as top-1 (proves the new content survived, not the stale one). + for &i in &update_ids { + let new_blob = simple_vec_bytes(i + UPDATE_SEED_OFFSET); + let results = search_keys(&mut c2, IDX, 1, &new_blob); + let want_key = format!("{IDX}:{i}"); + assert_eq!( + results.first().map(String::as_str), + Some(want_key.as_str()), + "updated key {want_key}: KNN on its NEW vector must return itself top-1, got {results:?}" + ); + } + + // Deleted keys must never resurface, even querying at their own old + // (tombstoned) position. + for &i in &delete_ids { + let old_blob = simple_vec_bytes(i); + let results = search_keys(&mut c2, IDX, 10, &old_blob); + let dead_key = format!("{IDX}:{i}"); + assert_absent(&results, &dead_key, "S2 deleted key across crash"); + } + + // NOTE: FT.INFO `num_docs` is NOT asserted to equal N-K here. Once a key + // is tombstoned against an already-compacted (Arc'd) immutable segment, + // `ImmutableSegment::mark_deleted_by_key_hash` (the "steady-state" + // tombstone path — see its doc comment) deliberately does NOT decrement + // that segment's cached `live_count` (a documented prototype + // limitation): it only inserts into a `tombstoned_keys` set consulted at + // SEARCH time (`is_live_bfs`). `num_docs` sums `live_count()`, so it can + // over-report after a delete against a pre-existing segment — this is + // true on a live (non-crash) server too, not something B3/B4 introduced + // or is expected to fix. The correctness contract that actually matters + // (deleted keys never resurface in KNN) is proven by the search loop + // above; `num_docs` only regains accuracy after the next compact/merge + // rebuilds `live_count` from scratch. + let num_docs = ft_info_num_docs(&mut c2, IDX); + assert!( + num_docs >= (N - K) as i64, + "num_docs ({num_docs}) must never UNDER-report below the true live count \ + (N-K={})", + N - K + ); + + drop(c2); + drop(guard2); + let _ = std::fs::remove_dir_all(&dir); +} + +// --------------------------------------------------------------------------- +// S3: orphan sweep +// --------------------------------------------------------------------------- + +#[test] +#[ignore] // Requires built release binary; run explicitly. +fn s3_orphan_files_swept_on_boot() { + const N: usize = 400; + const IDX: &str = "s3"; + + let port = unique_port(); + let dir = unique_dir("s3"); + let vdir = vector_persist_dir(&dir); + let idx_dir = manifest::index_persist_dir(&vdir, IDX.as_bytes()); + + let guard = spawn_moon_aof(port, &dir); + let mut c = wait_ready(port); + + ft_create(&mut c, IDX, DIM, 100, None); + let ids: Vec = (0..N as u32).collect(); + hset_batch(&mut c, &format!("{IDX}:"), &ids, simple_vec_bytes); + ft_compact(&mut c, IDX); + wait_for_manifest_min_segments(&idx_dir, 1, Duration::from_secs(10)); + + drop(c); + let mut guard = guard; + sigkill(&mut guard.child); + wait_for_port_down(port); + + // Inject orphans while the server is DOWN: an unreferenced staging dir, + // an unreferenced segment dir, and an unreferenced keymap file. + let staging_dir = idx_dir.join("staging-999"); + let segment_dir = idx_dir.join("segment-998"); + let keymap_file = idx_dir.join("keymap-999.bin"); + std::fs::create_dir_all(&staging_dir).expect("create fake staging dir"); + std::fs::write(staging_dir.join("meta.json"), b"{}").expect("write fake staging file"); + std::fs::create_dir_all(&segment_dir).expect("create fake segment dir"); + std::fs::write(segment_dir.join("meta.json"), b"{}").expect("write fake segment file"); + std::fs::write(&keymap_file, b"not a real keymap").expect("write fake keymap"); + + assert!(staging_dir.exists() && segment_dir.exists() && keymap_file.exists()); + + // -- Restart: orphan sweep runs in RecoveryState::finish() ------------ + let guard2 = start_moon_alive(spawn_moon_aof, port, &dir); + let mut c2 = wait_ready(port); + + // Bounded poll: sweep runs synchronously during startup (before accept), + // but poll defensively rather than assume zero latency. + let deadline = Instant::now() + Duration::from_secs(10); + loop { + if !staging_dir.exists() && !segment_dir.exists() && !keymap_file.exists() { + break; + } + if Instant::now() >= deadline { + panic!( + "orphan sweep did not remove all injected orphans within 10s: \ + staging_exists={} segment_exists={} keymap_exists={}", + staging_dir.exists(), + segment_dir.exists(), + keymap_file.exists() + ); + } + std::thread::sleep(Duration::from_millis(20)); + } + + // Index still answers correctly post-sweep. + let probe = simple_vec_bytes(7); + let results = search_keys(&mut c2, IDX, 1, &probe); + assert_eq!( + results.first().map(String::as_str), + Some(format!("{IDX}:7").as_str()), + "index must still search correctly after orphan sweep" + ); + let num_docs = ft_info_num_docs(&mut c2, IDX); + assert_eq!(num_docs, N as i64, "num_docs unaffected by orphan sweep"); + + drop(c2); + drop(guard2); + let _ = std::fs::remove_dir_all(&dir); +} + +// --------------------------------------------------------------------------- +// S4: collection_id pin survives a post-recovery compact+merge cycle +// --------------------------------------------------------------------------- + +fn l2_sq(a: &[f32], b: &[f32]) -> f32 { + a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum() +} + +fn brute_force_topk(dataset: &[(String, Vec)], query: &[f32], k: usize) -> Vec { + let mut scored: Vec<(f32, &str)> = dataset + .iter() + .map(|(id, v)| (l2_sq(v, query), id.as_str())) + .collect(); + scored.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + scored + .into_iter() + .take(k) + .map(|(_, id)| id.to_string()) + .collect() +} + +fn recall_at_k(ground_truth: &[String], approx: &[String], k: usize) -> f32 { + let gt: HashSet<&str> = ground_truth.iter().take(k).map(String::as_str).collect(); + let ap: HashSet<&str> = approx.iter().take(k).map(String::as_str).collect(); + if gt.is_empty() { + return 1.0; + } + gt.intersection(&ap).count() as f32 / gt.len() as f32 +} + +/// Generate `n_clusters * per_cluster` clustered vectors: cluster centers +/// are random directions in R^DIM, members are the center plus small +/// Gaussian noise. Low dimensionality (DIM=8) keeps this a meaningful +/// recall test (concentration-of-distance is a high-dimension problem — +/// see the repo gotcha on misleading random-Gaussian recall at high dims). +fn clustered_dataset( + seed: u64, + n_clusters: usize, + per_cluster: usize, + id_start: u32, + prefix: &str, +) -> Vec<(String, Vec)> { + let mut rng = Rng::new(seed); + let mut out = Vec::with_capacity(n_clusters * per_cluster); + let mut next_id = id_start; + for _ in 0..n_clusters { + // Scale cluster centers up (x4) so clusters are well-separated in + // R^DIM, and give members noise (0.3) that comfortably EXCEEDS the + // per-vector affine SQ8 quantization step (~range/255 ≈ 0.03 for + // values spanning roughly ±4). Noise at quant-step magnitude would + // collapse same-cluster members onto near-identical codes, making + // the rank-10 boundary an arbitrary tie for both this test's + // client-side recall assert and the merge's internal recall gate + // (VACUUM VECTOR / force_merge_index, fixed tolerance 0.90 — see + // src/vector/store.rs::force_merge_index). Noise ≫ quant step keeps + // the top-10 unambiguous (intra-cluster spread ~0.3·√(2·DIM) ≈ 2.4 + // vs inter-cluster separation ~4·√(2·DIM) ≈ 32 — still clearly + // clustered). NOTE: an earlier constant 0.899996-vs-0.90 gate + // failure here was NOT data ties — it was a self-exclusion + // asymmetry in verify_merge_recall (GT excluded the query point, + // HNSW included it, capping recall at (k-1)/k = 0.90 exactly); + // fixed in src/vector/segment/compaction.rs alongside this test. + let center: Vec = random_vec(&mut rng, DIM) + .into_iter() + .map(|x| x * 4.0) + .collect(); + for _ in 0..per_cluster { + let mut v = center.clone(); + for x in v.iter_mut() { + *x += 0.3 * rng.randn(); + } + out.push((format!("{prefix}{next_id}"), v)); + next_id += 1; + } + } + out +} + +#[test] +#[ignore] // Requires built release binary; run explicitly. +fn s4_collection_id_pin_survives_post_recovery_merge() { + const IDX: &str = "s4"; + const N_CLUSTERS: usize = 20; + const PER_CLUSTER: usize = 25; // 500 pre-crash + 500 post-recovery = 1000 total + + let port = unique_port(); + let dir = unique_dir("s4"); + let vdir = vector_persist_dir(&dir); + let idx_dir = manifest::index_persist_dir(&vdir, IDX.as_bytes()); + + let guard = spawn_moon_aof(port, &dir); + let mut c = wait_ready(port); + + ft_create_ex(&mut c, IDX, DIM, 100, Some(128), Some(32), Some(400)); + + let prefix = format!("{IDX}:"); + let batch1 = clustered_dataset(0xC0FFEE_u64, N_CLUSTERS, PER_CLUSTER, 0, &prefix); + let cmds: Vec>> = batch1 + .iter() + .map(|(key, v)| { + vec![ + b"HSET".to_vec(), + key.clone().into_bytes(), + b"vec".to_vec(), + f32s_to_le_bytes(v), + ] + }) + .collect(); + for r in c.pipeline(&cmds) { + assert!(matches!(r, V::Int(_)), "S4 batch1 HSET failed: {r:?}"); + } + ft_compact(&mut c, IDX); + wait_for_manifest_min_segments(&idx_dir, 1, Duration::from_secs(10)); + + drop(c); + let mut guard = guard; + sigkill(&mut guard.child); + wait_for_port_down(port); + + // -- Restart: segment A loads with its pinned collection_id ----------- + let guard2 = start_moon_alive(spawn_moon_aof, port, &dir); + let mut c2 = wait_ready(port); + wait_for_recovery_counters(&dir, IDX, Duration::from_secs(10)); + + // Insert a SECOND batch of NEW clustered vectors post-recovery, compact + // it into a new segment B — segment B is built under the RECOVERED + // index's pinned collection_id (the QJL rotation seed). If B3 mis-pinned + // it, segment A and segment B use incompatible rotations/codebooks and a + // merge (which stitches graphs over the raw codes, never re-quantizing) + // produces garbage distances for one half of the data — visible as a + // recall collapse below. + let batch2 = clustered_dataset( + 0xC0FFEE_u64 ^ 0xA5A5_A5A5, + N_CLUSTERS, + PER_CLUSTER, + (N_CLUSTERS * PER_CLUSTER) as u32, + &prefix, + ); + let cmds2: Vec>> = batch2 + .iter() + .map(|(key, v)| { + vec![ + b"HSET".to_vec(), + key.clone().into_bytes(), + b"vec".to_vec(), + f32s_to_le_bytes(v), + ] + }) + .collect(); + for r in c2.pipeline(&cmds2) { + assert!(matches!(r, V::Int(_)), "S4 batch2 HSET failed: {r:?}"); + } + ft_compact(&mut c2, IDX); + wait_for_manifest_min_segments(&idx_dir, 2, Duration::from_secs(10)); + + // Force the GraphUnion merge (VACUUM VECTOR merges whenever segment + // count >= 2, even below the auto-merge trigger threshold — see + // src/command/server_admin.rs's `seg_count < 2` gate). + let vacuum_reply = vacuum_vector(&mut c2, IDX); + assert!( + vacuum_reply.starts_with("OK") || vacuum_reply.starts_with("Merged"), + "VACUUM VECTOR {IDX} unexpected reply: {vacuum_reply}" + ); + // After a successful merge, exactly 1 immutable segment should remain. + // The manifest rewrite is asynchronous (persist_hook_after_install + // schedules a SnapshotPool job after the segment-list swap), so poll + // until the manifest CONVERGES to 1 segment — a min>=1 wait would return + // instantly on the stale pre-merge manifest ([seg_a, seg_b]). + let post_merge_manifest = + wait_for_manifest_exact_segments(&idx_dir, 1, Duration::from_secs(10)); + assert_eq!( + post_merge_manifest.segment_ids.len(), + 1, + "expected exactly 1 segment after GraphUnion merge, got {:?}", + post_merge_manifest.segment_ids + ); + + // Recall check: brute-force ground truth over ALL vectors (both + // batches) computed client-side, compared against FT.SEARCH KNN=10 for + // one query per cluster (batch1's cluster centers — held-in points). + let mut dataset = batch1.clone(); + dataset.extend(batch2.clone()); + + let mut recalls = Vec::with_capacity(N_CLUSTERS * 2); + for (id, vec) in batch1.iter().step_by(PER_CLUSTER) { + let _ = id; + let approx = search_keys(&mut c2, IDX, 10, &f32s_to_le_bytes(vec)); + let gt = brute_force_topk(&dataset, vec, 10); + recalls.push(recall_at_k(>, &approx, 10)); + } + for (id, vec) in batch2.iter().step_by(PER_CLUSTER) { + let _ = id; + let approx = search_keys(&mut c2, IDX, 10, &f32s_to_le_bytes(vec)); + let gt = brute_force_topk(&dataset, vec, 10); + recalls.push(recall_at_k(>, &approx, 10)); + } + let mean_recall = recalls.iter().sum::() / recalls.len() as f32; + assert!( + mean_recall >= 0.9, + "S4 collection_id-pin regression: mean recall@10 = {mean_recall:.4} across {} queries \ + (floor 0.90) — a mis-pinned collection_id/QJL seed on the post-recovery segment would \ + corrupt distances for the merged graph and collapse this number. Per-query: {recalls:?}", + recalls.len() + ); + + let num_docs = ft_info_num_docs(&mut c2, IDX); + assert_eq!( + num_docs, + dataset.len() as i64, + "num_docs must equal total inserted vectors after merge" + ); + + drop(c2); + drop(guard2); + let _ = std::fs::remove_dir_all(&dir); +} + +// --------------------------------------------------------------------------- +// S5: no-persist-dir regression guard +// --------------------------------------------------------------------------- + +#[test] +#[ignore] // Requires built release binary; run explicitly. +fn s5_no_persist_dir_regression_guard() { + const N: usize = 50; + const IDX: &str = "s5"; + + let port = unique_port(); + let dir = unique_dir("s5"); + + let guard = spawn_moon_no_persist(port, &dir); + let mut c = wait_ready(port); + + ft_create(&mut c, IDX, DIM, 100, None); + let ids: Vec = (0..N as u32).collect(); + hset_batch(&mut c, &format!("{IDX}:"), &ids, simple_vec_bytes); + ft_compact(&mut c, IDX); + + // No persistence configured -> no idx-* dirs, no sidecar meta file, + // anywhere under --dir. + let idx_dirs = walk_matching(&dir, &|n| n.starts_with("idx-")); + assert!( + idx_dirs.is_empty(), + "no vector_persist_dir configured, but found idx-* dirs: {idx_dirs:?}" + ); + let sidecars = walk_matching(&dir, &|n| n == "vector-indexes.meta"); + assert!( + sidecars.is_empty(), + "no vector_persist_dir configured, but found a sidecar file: {sidecars:?}" + ); + + drop(c); + let mut guard = guard; + sigkill(&mut guard.child); + wait_for_port_down(port); + + // -- Restart ---------------------------------------------------------- + let guard2 = start_moon_alive(spawn_moon_no_persist, port, &dir); + let mut c2 = wait_ready(port); + + // Server is alive and responsive (no crash on empty/absent recovery + // state). + assert_eq!(c2.cmd(&[b"PING"]), V::Simple("PONG".into())); + + // Index definitions are gone entirely (never persisted) — FT._LIST must + // not report the old index name. + let indexes = ft_list(&mut c2); + assert!( + !indexes.iter().any(|n| n == IDX), + "index {IDX} must NOT survive a restart with no persistence configured, \ + got FT._LIST = {indexes:?}" + ); + + // FT.SEARCH against the now-nonexistent index must fail cleanly, not + // silently return stale (impossible, since recreation would require the + // sidecar) or crash. + let r = c2.cmd(&[ + b"FT.SEARCH", + IDX.as_bytes(), + b"*=>[KNN 1 @vec $B]", + b"PARAMS", + b"2", + b"B", + &simple_vec_bytes(0), + b"DIALECT", + b"2", + ]); + assert!( + matches!(r, V::Err(_)), + "FT.SEARCH on a never-persisted index after restart must error, got {r:?}" + ); + + // Still no idx-*/sidecar files anywhere (recovery must not have created + // any either). + let idx_dirs_post = walk_matching(&dir, &|n| n.starts_with("idx-")); + assert!( + idx_dirs_post.is_empty(), + "post-restart, found unexpected idx-* dirs: {idx_dirs_post:?}" + ); + let sidecars_post = walk_matching(&dir, &|n| n == "vector-indexes.meta"); + assert!( + sidecars_post.is_empty(), + "post-restart, found unexpected sidecar file: {sidecars_post:?}" + ); + + drop(c2); + drop(guard2); + let _ = std::fs::remove_dir_all(&dir); +} diff --git a/tests/vector_flush_hdel_tombstone.rs b/tests/vector_flush_hdel_tombstone.rs new file mode 100644 index 00000000..87195ef2 --- /dev/null +++ b/tests/vector_flush_hdel_tombstone.rs @@ -0,0 +1,458 @@ +//! Integration tests — vector-index keyspace parity for FLUSHALL/FLUSHDB and +//! field-level HDEL (persistence-review R3/R4). +//! +//! R3: FLUSHALL/FLUSHDB must clear vector-index CONTENTS (no ghost vectors +//! searchable against deleted hashes) while KEEPING the FT.CREATE +//! definitions (matching restart semantics: definitions come from the +//! sidecar, contents are re-derived from the — now empty — keyspace). +//! R4: HDEL of the indexed vector field alone (hash key survives) must +//! tombstone the vector — previously it stayed searchable until the +//! whole key was deleted or the field re-HSET. +//! +//! Harness mirrors `tests/ft_search_as_of_filter.rs`. +//! +//! Run: +//! cargo test --test vector_flush_hdel_tombstone \ +//! --no-default-features --features runtime-tokio,jemalloc,graph \ +//! -- --test-threads=1 +#![cfg(all(feature = "runtime-tokio", feature = "graph"))] + +use moon::config::ServerConfig; +use moon::runtime::cancel::CancellationToken; +use moon::runtime::channel; +use moon::server::listener; +use moon::shard::Shard; +use moon::shard::mesh::{CHANNEL_BUFFER_SIZE, ChannelMesh}; +use tokio::net::TcpListener; + +// --------------------------------------------------------------------------- +// Harness — mirror of `tests/ft_search_as_of_filter.rs`. +// --------------------------------------------------------------------------- + +fn build_config(port: u16, num_shards: usize) -> ServerConfig { + ServerConfig { + bind: "127.0.0.1".to_string(), + port, + databases: 16, + requirepass: None, + appendonly: "no".to_string(), + appendfsync: "everysec".to_string(), + save: None, + dir: ".".to_string(), + dbfilename: "dump.rdb".to_string(), + appendfilename: "appendonly.aof".to_string(), + maxmemory: Some(0), + maxmemory_policy: "noeviction".to_string(), + maxmemory_samples: 5, + shards: num_shards, + cluster_enabled: false, + cluster_node_timeout: 15000, + aclfile: None, + protected_mode: "no".to_string(), + acllog_max_len: 128, + tls_port: 0, + tls_cert_file: None, + tls_key_file: None, + tls_ca_cert_file: None, + tls_ciphersuites: None, + disk_offload: "disable".to_string(), + disk_offload_dir: None, + disk_offload_threshold: 0.85, + segment_warm_after: 3600, + pagecache_size: None, + checkpoint_timeout: 300, + checkpoint_completion: 0.9, + max_wal_size: "256mb".to_string(), + wal_fpi: "enable".to_string(), + wal_compression: "lz4".to_string(), + wal_segment_size: "16mb".to_string(), + vec_codes_mlock: "enable".to_string(), + segment_cold_after: 86400, + segment_cold_min_qps: 0.1, + vec_diskann_beam_width: 8, + vec_diskann_cache_levels: 3, + uring_sqpoll_ms: None, + admin_port: 0, + slowlog_log_slower_than: 10000, + slowlog_max_len: 128, + check_config: false, + initial_keyspace_hint: 0, + memory_arenas_cap: 8, + maxclients: 10000, + timeout: 0, + tcp_keepalive: 300, + console_auth_required: false, + console_auth_secret: String::new(), + console_cors_origin: vec![], + console_rate_limit: 1000.0, + console_rate_burst: 2000.0, + recovery_target_lsn: None, + recovery_target_time: None, + manifest_tombstone_retain_epochs: 2, + manifest_tombstone_retain_secs: 300, + disk_free_min_pct: 5, + mvcc_committed_prune_margin: 1000, + max_unflushed_immutable_segments: 20, + mvcc_old_snapshot_threshold_secs: 600, + autovacuum: "enable".to_string(), + autovacuum_budget_ms_min: 5, + autovacuum_budget_ms_max: 200, + autovacuum_target_p95_ms: 10, + autovacuum_interval_secs: 30, + graph_merge_max_segments: 8, + graph_dead_edge_trigger: 0.20, + autovacuum_starvation_cap_secs: 300, + vec_warm_mmap_budget: "2gb".to_string(), + cold_orphan_sweep_interval_secs: 300, + migrate_aof_from: None, + migrate_aof_to: None, + migrate_aof_shards: 0, + ..Default::default() + } +} + +async fn start_moon_sharded(num_shards: usize) -> (u16, CancellationToken) { + let probe = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = probe.local_addr().unwrap().port(); + drop(probe); + + let token = CancellationToken::new(); + let config = build_config(port, num_shards); + let cancel = token.clone(); + + std::thread::spawn(move || { + let mut mesh = ChannelMesh::new(num_shards, CHANNEL_BUFFER_SIZE); + let conn_txs: Vec> = + (0..num_shards).map(|i| mesh.conn_tx(i)).collect(); + let all_notifiers = mesh.all_notifiers(); + let all_pubsub_registries: Vec< + std::sync::Arc>, + > = (0..num_shards) + .map(|_| { + std::sync::Arc::new(parking_lot::RwLock::new(moon::pubsub::PubSubRegistry::new())) + }) + .collect(); + let all_remote_sub_maps: Vec< + std::sync::Arc< + parking_lot::RwLock, + >, + > = (0..num_shards) + .map(|_| { + std::sync::Arc::new(parking_lot::RwLock::new( + moon::shard::remote_subscriber_map::RemoteSubscriberMap::new(), + )) + }) + .collect(); + let affinity_tracker = std::sync::Arc::new(parking_lot::RwLock::new( + moon::shard::affinity::AffinityTracker::new(), + )); + let mut shards: Vec = (0..num_shards) + .map(|id| Shard::new(id, num_shards, config.databases, config.to_runtime_config())) + .collect(); + let all_dbs: Vec> = shards + .iter_mut() + .map(|s| std::mem::take(&mut s.databases)) + .collect(); + let (shard_databases, mut slice_inits) = + moon::shard::shared_databases::ShardDatabases::new(all_dbs); + + let mut shard_handles = Vec::with_capacity(num_shards); + for (id, mut shard) in shards.into_iter().enumerate() { + let producers = mesh.take_producers(id); + let consumers = mesh.take_consumers(id); + let conn_rx = mesh.take_conn_rx(id); + let shard_config = config.clone(); + let shard_cancel = cancel.clone(); + let shard_spsc_notify = mesh.take_notify(id); + let shard_all_notifiers = all_notifiers.clone(); + let shard_dbs = shard_databases.clone(); + let shard_pubsub_regs = all_pubsub_registries.clone(); + let shard_remote_sub_maps = all_remote_sub_maps.clone(); + let shard_affinity = affinity_tracker.clone(); + let shard_slice_init = slice_inits.remove(0); + + let handle = std::thread::Builder::new() + .name(format!("flush-hdel-shard-{}", id)) + .spawn(move || { + let rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("failed to build shard runtime"); + let local = tokio::task::LocalSet::new(); + let (_snap_tx, snap_rx) = channel::watch(0u64); + let snap_tx = _snap_tx; + let acl_t = std::sync::Arc::new(std::sync::RwLock::new( + moon::acl::AclTable::load_or_default(&shard_config), + )); + let rt_cfg = std::sync::Arc::new(parking_lot::RwLock::new( + shard_config.to_runtime_config(), + )); + rt.block_on(local.run_until(shard.run( + conn_rx, + None, + consumers, + producers, + shard_cancel, + None, + None, + None, + snap_rx, + snap_tx, + None, + None, + 0, + acl_t, + rt_cfg, + std::sync::Arc::new(shard_config), + shard_spsc_notify, + shard_all_notifiers, + shard_dbs, + shard_pubsub_regs, + shard_remote_sub_maps, + shard_affinity, + shard_slice_init, + ))); + }) + .expect("failed to spawn shard thread"); + shard_handles.push(handle); + } + + let listener_rt = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("failed to build listener runtime"); + let listener_cancel = cancel.clone(); + listener_rt.block_on(async { + if let Err(e) = + listener::run_sharded(config, conn_txs, listener_cancel, false, affinity_tracker) + .await + { + eprintln!("Listener error: {}", e); + } + }); + cancel.cancel(); + for handle in shard_handles { + let _ = handle.join(); + } + }); + + tokio::time::sleep(std::time::Duration::from_millis(250)).await; + (port, token) +} + +async fn connect(port: u16) -> redis::aio::MultiplexedConnection { + let client = redis::Client::open(format!("redis://127.0.0.1:{port}")).unwrap(); + client.get_multiplexed_async_connection().await.unwrap() +} + +fn vec4_bytes(v: [f32; 4]) -> Vec { + let mut out = Vec::with_capacity(16); + for x in &v { + out.extend_from_slice(&x.to_le_bytes()); + } + out +} + +fn parse_search_count(v: &redis::Value) -> i64 { + match v { + redis::Value::Array(items) => match items.first() { + Some(redis::Value::Int(n)) => *n, + Some(redis::Value::BulkString(b)) => std::str::from_utf8(b) + .expect("count utf-8") + .parse::() + .expect("count parses"), + other => panic!("unexpected first item in FT.SEARCH response: {other:?}"), + }, + other => panic!("expected Array from FT.SEARCH, got {other:?}"), + } +} + +async fn ft_create(conn: &mut redis::aio::MultiplexedConnection, name: &str) { + let _: String = redis::cmd("FT.CREATE") + .arg(name) + .arg("ON") + .arg("HASH") + .arg("PREFIX") + .arg("1") + .arg("doc:") + .arg("SCHEMA") + .arg("vec") + .arg("VECTOR") + .arg("HNSW") + .arg("6") + .arg("DIM") + .arg("4") + .arg("TYPE") + .arg("FLOAT32") + .arg("DISTANCE_METRIC") + .arg("L2") + .query_async(conn) + .await + .expect("FT.CREATE should succeed"); +} + +async fn hset_doc(conn: &mut redis::aio::MultiplexedConnection, key: &str, v: [f32; 4]) { + let _: i64 = redis::cmd("HSET") + .arg(key) + .arg("vec") + .arg(vec4_bytes(v)) + .arg("label") + .arg("payload") + .query_async(conn) + .await + .expect("HSET should succeed"); +} + +async fn knn_count(conn: &mut redis::aio::MultiplexedConnection, idx: &str) -> i64 { + let q = vec4_bytes([1.0, 0.0, 0.0, 0.0]); + let v: redis::Value = redis::cmd("FT.SEARCH") + .arg(idx) + .arg("*=>[KNN 10 @vec $q]") + .arg("PARAMS") + .arg("2") + .arg("q") + .arg(q) + .arg("DIALECT") + .arg("2") + .query_async(conn) + .await + .expect("FT.SEARCH should succeed"); + parse_search_count(&v) +} + +async fn ft_list(conn: &mut redis::aio::MultiplexedConnection) -> Vec { + let v: redis::Value = redis::cmd("FT._LIST") + .query_async(conn) + .await + .expect("FT._LIST should succeed"); + match v { + redis::Value::Array(items) => items + .into_iter() + .filter_map(|it| match it { + redis::Value::BulkString(b) => Some(String::from_utf8_lossy(&b).into_owned()), + redis::Value::SimpleString(s) => Some(s), + _ => None, + }) + .collect(), + other => panic!("expected Array from FT._LIST, got {other:?}"), + } +} + +// --------------------------------------------------------------------------- +// R3: FLUSHALL clears index contents, keeps the definition, index stays live. +// --------------------------------------------------------------------------- +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn flushall_clears_vector_index_contents() { + let (port, _shutdown) = start_moon_sharded(1).await; + let mut conn = connect(port).await; + + ft_create(&mut conn, "fidx").await; + hset_doc(&mut conn, "doc:1", [1.0, 0.0, 0.0, 0.0]).await; + hset_doc(&mut conn, "doc:2", [0.0, 1.0, 0.0, 0.0]).await; + assert_eq!(knn_count(&mut conn, "fidx").await, 2, "pre-flush sanity"); + + let _: String = redis::cmd("FLUSHALL") + .query_async(&mut conn) + .await + .expect("FLUSHALL should succeed"); + + // Ghost check: hashes are gone, so the index must return nothing. + assert_eq!( + knn_count(&mut conn, "fidx").await, + 0, + "FLUSHALL must clear vector index contents (no ghost vectors)" + ); + // Definition survives (matching restart-rescan semantics). + let names = ft_list(&mut conn).await; + assert!( + names.iter().any(|n| n == "fidx"), + "FT.CREATE definition must survive FLUSHALL; FT._LIST = {names:?}" + ); + // Index remains functional for new writes. + hset_doc(&mut conn, "doc:3", [1.0, 0.0, 0.0, 0.0]).await; + assert_eq!( + knn_count(&mut conn, "fidx").await, + 1, + "index must keep auto-indexing new HSETs after FLUSHALL" + ); +} + +// --------------------------------------------------------------------------- +// R3: FLUSHDB behaves identically (Moon FT indexes are keyspace-global). +// --------------------------------------------------------------------------- +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn flushdb_clears_vector_index_contents() { + let (port, _shutdown) = start_moon_sharded(1).await; + let mut conn = connect(port).await; + + ft_create(&mut conn, "fidx").await; + hset_doc(&mut conn, "doc:1", [1.0, 0.0, 0.0, 0.0]).await; + assert_eq!(knn_count(&mut conn, "fidx").await, 1, "pre-flush sanity"); + + let _: String = redis::cmd("FLUSHDB") + .query_async(&mut conn) + .await + .expect("FLUSHDB should succeed"); + + assert_eq!( + knn_count(&mut conn, "fidx").await, + 0, + "FLUSHDB must clear vector index contents (no ghost vectors)" + ); + let names = ft_list(&mut conn).await; + assert!( + names.iter().any(|n| n == "fidx"), + "FT.CREATE definition must survive FLUSHDB; FT._LIST = {names:?}" + ); +} + +// --------------------------------------------------------------------------- +// R4: HDEL of the vector field alone tombstones the vector; the hash key and +// its other fields survive. +// --------------------------------------------------------------------------- +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn hdel_vector_field_tombstones_vector() { + let (port, _shutdown) = start_moon_sharded(1).await; + let mut conn = connect(port).await; + + ft_create(&mut conn, "fidx").await; + hset_doc(&mut conn, "doc:1", [1.0, 0.0, 0.0, 0.0]).await; + hset_doc(&mut conn, "doc:2", [0.0, 1.0, 0.0, 0.0]).await; + assert_eq!(knn_count(&mut conn, "fidx").await, 2, "pre-HDEL sanity"); + + let removed: i64 = redis::cmd("HDEL") + .arg("doc:1") + .arg("vec") + .query_async(&mut conn) + .await + .expect("HDEL should succeed"); + assert_eq!(removed, 1, "HDEL must remove the vector field"); + + // The vector must no longer be searchable... + assert_eq!( + knn_count(&mut conn, "fidx").await, + 1, + "HDEL of the indexed vector field must tombstone the vector" + ); + // ...but the hash key and its other fields survive. + let label: String = redis::cmd("HGET") + .arg("doc:1") + .arg("label") + .query_async(&mut conn) + .await + .expect("HGET should succeed"); + assert_eq!(label, "payload", "non-vector fields must survive HDEL"); + + // HDEL of a non-vector field must NOT tombstone (control). + let _: i64 = redis::cmd("HDEL") + .arg("doc:2") + .arg("label") + .query_async(&mut conn) + .await + .expect("HDEL of non-vector field should succeed"); + assert_eq!( + knn_count(&mut conn, "fidx").await, + 1, + "HDEL of a non-vector field must not tombstone the vector" + ); +} diff --git a/tests/vector_segment_merge.rs b/tests/vector_segment_merge.rs index a38091fb..d4396938 100644 --- a/tests/vector_segment_merge.rs +++ b/tests/vector_segment_merge.rs @@ -364,11 +364,17 @@ fn test_merge_mode_none_is_noop() { ); } -// ── Test 3a: recall gate REJECTS an impossible tolerance ── +// ── Test 3a: recall gate REJECTS an unattainable tolerance (rollback path) ── // -// The recall gate fires when n >= MIN_RECALL_SAMPLE_N (50). With 200 vectors in -// 2 segments and tolerance=1.0 (impossible — real HNSW recall < 1.0), the gate -// must return Err(RecallTooLow) and leave both segments intact. +// The recall gate fires when n >= MIN_RECALL_SAMPLE_N (50). Recall is capped +// at 1.0 by construction, so a tolerance ABOVE 1.0 deterministically trips +// `recall < tolerance` — this test pins the Err(RecallTooLow) propagation and +// the segments-left-intact rollback guarantee, independent of dataset +// difficulty. (The pre-fix version used tolerance=1.0 and only rejected +// because verify_merge_recall had a self-exclusion asymmetry capping +// measurable recall at (k-1)/k = 0.9; honest recall on a clean merge is +// exactly 1.0 — the merged graph is verified in the same decoded-centroid +// space it was built in — so 1.0 no longer rejects.) #[test] fn test_recall_gate_rejects_impossible_tolerance() { @@ -403,12 +409,12 @@ fn test_recall_gate_rejects_impossible_tolerance() { "Should have 2 segments before merge attempt" ); - // tolerance=1.0 is impossible — HNSW recall on TQ4 @ 64d < 1.0 always. - // Gate must reject and return Err. - let result = store.force_merge_index_with_tolerance(b"idx3a", 1.0); + // tolerance > 1.0 is unattainable by construction (recall <= 1.0), so + // the gate must reject and return Err regardless of dataset difficulty. + let result = store.force_merge_index_with_tolerance(b"idx3a", 1.5); assert!( result.is_err(), - "Merge with tolerance=1.0 must be rejected by recall gate, but got: {:?}", + "Merge with tolerance=1.5 must be rejected by recall gate, but got: {:?}", result ); @@ -574,6 +580,7 @@ fn test_merge_overlapping_ids_highest_lsn_wins() { 42, MergeMode::GraphUnion, 0.90, + None, ) .expect("merge failed");