diff --git a/include/index/graph/diskann/disk_page_io.hpp b/include/index/graph/diskann/disk_page_io.hpp index 40507fad..601dc711 100644 --- a/include/index/graph/diskann/disk_page_io.hpp +++ b/include/index/graph/diskann/disk_page_io.hpp @@ -83,6 +83,7 @@ class DiskPageIO { const size_t shard_capacity = page_cache_capacity == 0 ? 0 : std::max(1, page_cache_capacity / num_shards_); cache_enabled_ = shard_capacity > 0; + page_cache_capacity_ = cache_enabled_ ? shard_capacity * num_shards_ : 0; shards_.reserve(num_shards_); for (uint32_t s = 0; s < num_shards_; ++s) { shards_.push_back(std::make_unique(shard_capacity, geom_.page_size)); @@ -742,6 +743,11 @@ class DiskPageIO { /// the search-side peek/fill pair below is pointless without them. [[nodiscard]] bool page_cache_enabled() const { return cache_enabled_; } + /// Effective total page capacity across shards (0 when the cache is + /// disabled). Prefetch-then-read waves must stay under this or the wave + /// evicts its own pages before the reads land. + [[nodiscard]] size_t page_cache_capacity() const { return page_cache_capacity_; } + /// Peek + version snapshot, the read half of the search fill protocol. /// Hit: the page is copied into @p out and true is returned. Miss: false, /// and @p version_out receives the page's current write version so the @@ -902,6 +908,7 @@ class DiskPageIO { std::atomic file_size_{0}; uint32_t num_shards_ = kNumShards; bool cache_enabled_ = false; + size_t page_cache_capacity_ = 0; std::vector> shards_; std::unordered_map> vec_cache_; mutable std::mutex vec_mutex_; diff --git a/include/index/graph/diskann/diskann_index.hpp b/include/index/graph/diskann/diskann_index.hpp index 79098bfa..2caba08b 100644 --- a/include/index/graph/diskann/diskann_index.hpp +++ b/include/index/graph/diskann/diskann_index.hpp @@ -44,6 +44,8 @@ #include #include +#include + #include "coro/sync_wait.hpp" #include "coro/task.hpp" #include "coro/thread_pool.hpp" @@ -115,18 +117,33 @@ struct DiskANNLoadParams { ///< largest DiskANNSearchParams::search_list_size used after load. // --- In-place update mode --- - bool updatable = false; ///< open O_RDWR + enable insert/remove/update_node/flush - uint32_t update_search_l = 0; ///< L for the insert NN-search; 0 => max_degree + 32 - ///< (Yi's build_k = degree + 32 rule) - bool update_rerank = true; ///< re-rank insert search candidates by exact L2 (via the - ///< coords cache) before taking the top max_degree. Matches - ///< Yi's trace benchmark (_rerank_flag defaults true); Yi's - ///< sequential UpdateRunner sets it false. - bool update_insert_prune = false; ///< alpha-RNG prune the insert pool instead of linking the - ///< top max_degree candidates. Yi never prunes at insert - ///< (its top_k == degree makes the prune branch dead code); - ///< reconnect re-prunes on overflow either way. - float update_alpha = 1.2f; ///< alpha-RNG pruning for insert/reconnect (Vamana default) + bool updatable = false; ///< open O_RDWR + enable insert/remove/update_node/flush + uint32_t update_search_l = 0; ///< L for the insert NN-search; 0 => max_degree + 32 + ///< (Yi's build_k = degree + 32 rule) + bool update_rerank = true; ///< re-rank insert search candidates by exact L2 (via the + ///< coords cache) before taking the top max_degree. Matches + ///< Yi's trace benchmark (_rerank_flag defaults true); Yi's + ///< sequential UpdateRunner sets it false. + bool update_insert_prune = false; ///< alpha-RNG prune the insert pool instead of linking the + ///< top max_degree candidates. Yi never prunes at insert + ///< (its top_k == degree makes the prune branch dead code); + ///< reconnect re-prunes on overflow either way. + float update_alpha = 1.2f; ///< alpha-RNG pruning for insert/reconnect (Vamana default) + bool update_repair = false; ///< delete-time in-neighbor repair: tombstone the batch, + ///< discover each deleted node's approximate in-neighbors, + ///< rebuild their lists (drop dead edges + two-hop splice), + ///< and only then release the slots for reuse. Supersedes + ///< the lazy path's safety-net reconnect (never fires: its + ///< affected set is the discovery set's subset) + uint32_t repair_search_l = 0; ///< 0 = discovery uses only the deleted node's own + ///< out-neighbors (symmetry approximation, no extra search); + ///< >0 = also run an insert-style candidate search per delete + ///< and merge the top repair_search_l results + bool track_in_degree = false; + ///< allocate + maintain live in-degree counters (updatable mode only); adds + ///< a one-time full neighbor-list scan at load + uint32_t garden_search_l = 0; + ///< L for garden_refresh's candidate search; 0 => update_search_l double safety_net_ratio = 0.05; ///< tombstone ratio that arms the safety-net reconnect uint64_t safety_net_ops = 16; ///< deletes without an insert before the safety net may fire size_t page_cache_capacity = 4096; ///< update-path page LRU cache capacity; 0 disables it @@ -365,6 +382,7 @@ class DiskANNIndex { std::min(1024, std::max(2ull * beam_width_, nopq_depth)); const uint32_t scratch_list_size = std::max({DiskANNSearchParams{}.search_list_size, params.update_search_l, + params.garden_search_l, params.scratch_search_list_size}); ThreadDataScratchConfig scratch_config; scratch_config.n_page_slots = scratch_slots; @@ -795,6 +813,79 @@ class DiskANNIndex { return slot_alloc_.is_deleted(id); } + struct GardenStats { + uint64_t refreshed = 0; + uint64_t selected = 0; + uint32_t indeg_p10_before = 0; + uint32_t indeg_p50_before = 0; + uint32_t indeg_p10_after = 0; + uint32_t indeg_p50_after = 0; + uint64_t elapsed_us = 0; + }; + + struct InDegreePercentiles { + uint32_t p10 = 0; + uint32_t p50 = 0; + uint32_t p90 = 0; + }; + + GardenStats garden_refresh(uint32_t budget, coro::thread_pool *pool = nullptr) { + const auto start = std::chrono::steady_clock::now(); + if (!updatable_) { + throw std::runtime_error("DiskANNIndex::garden_refresh: index not loaded in updatable mode"); + } + if (!track_in_degree_) { + throw std::runtime_error("DiskANNIndex::garden_refresh: track_in_degree is disabled"); + } + + std::lock_guard update_guard(update_serial_mutex_); + page_io_->clear_cache(); + + GardenStats stats; + if (pool == nullptr) { + const uint32_t workers = std::max(1, update_insert_threads_); + coro::thread_pool local_pool{{.thread_count = workers, + .on_thread_start_functor = nullptr, + .on_thread_stop_functor = nullptr}}; + try { + stats = garden_refresh_locked(budget, local_pool); + local_pool.shutdown(); + } catch (...) { + local_pool.shutdown(); + throw; + } + } else { + stats = garden_refresh_locked(budget, *pool); + } + + stats.elapsed_us = static_cast(std::chrono::duration_cast( + std::chrono::steady_clock::now() - start) + .count()); + return stats; + } + + InDegreePercentiles in_degree_percentiles() const { + if (!track_in_degree_) { + return {}; + } + return percentiles_from_pairs(collect_live_in_degree_pairs(true)); + } + + uint32_t in_degree_of(uint32_t internal_id) const { + if (!track_in_degree_ || !in_degree_ || internal_id >= in_degree_size_) { + return 0; + } + return in_degree_[internal_id].load(std::memory_order_relaxed); + } + + std::vector debug_recount_in_degree() { + if (!updatable_) { + return {}; + } + std::lock_guard update_guard(update_serial_mutex_); + return recount_in_degree(nullptr); + } + /// Sentinel label for padded (missing) result slots. static constexpr uint64_t kNoLabel = std::numeric_limits::max(); @@ -874,8 +965,8 @@ class DiskANNIndex { return batch_insert_locked_with_pool(vectors, labels, count, batch_size, pool); } - /// Lazy-delete: cache old neighbors for two-hop, tombstone + free the slot. - /// Reconnect is deferred to the next insert or the safety net. + /// Delete one slot. The default lazy path frees immediately; repair mode + /// keeps the slot tombstoned until discovered in-neighbors are rebuilt. void remove(uint32_t internal_id) { if (!updatable_) { throw std::runtime_error("DiskANNIndex::remove: index not loaded in updatable mode"); @@ -885,15 +976,11 @@ class DiskANNIndex { std::shared_lock state_lock(update_mutex_); validate_removable_slot(internal_id); } - page_io_->clear_cache(); - remove_unlocked(internal_id); - if (maybe_safety_net_reconnect()) { - page_io_->flush_dirty_pages(); - } + remove_batch_locked(&internal_id, 1, nullptr, true); } - /// Lazy-delete a batch using a caller-owned coroutine pool. This mirrors - /// batch_remove() but avoids creating an update-private worker pool. + /// Delete a batch using a caller-owned coroutine pool. Repair mode uses the + /// pool for discovery searches and reconnect waves. void batch_remove_with_pool(const uint32_t *internal_ids, uint32_t count, coro::thread_pool &pool) { @@ -906,23 +993,11 @@ class DiskANNIndex { std::shared_lock state_lock(update_mutex_); validate_remove_batch(internal_ids, count); } - page_io_->clear_cache(); - std::vector> old_neighbors = - read_delete_neighbors(internal_ids, count, &pool); - { - std::unique_lock state_lock(update_mutex_); - for (uint32_t i = 0; i < count; ++i) { - remove_unlocked_with_neighbors(internal_ids[i], std::move(old_neighbors[i])); - } - } - if (maybe_safety_net_reconnect()) { - page_io_->flush_dirty_pages(); - } + remove_batch_locked(internal_ids, count, &pool, false); } - /// Lazy-delete a batch of internal ids using the same semantics as remove(). - /// The batch form amortizes update locking and cache setup over Yi-style - /// delete rounds. + /// Delete a batch of internal ids using the same semantics as remove(). The + /// batch form amortizes update locking and cache setup over Yi-style rounds. void batch_remove(const uint32_t *internal_ids, uint32_t count) { if (!updatable_) { throw std::runtime_error("DiskANNIndex::batch_remove: index not loaded in updatable mode"); @@ -938,18 +1013,7 @@ class DiskANNIndex { std::shared_lock state_lock(update_mutex_); validate_remove_batch(internal_ids, count); } - page_io_->clear_cache(); - std::vector> old_neighbors = - read_delete_neighbors(internal_ids, count, nullptr); - { - std::unique_lock state_lock(update_mutex_); - for (uint32_t i = 0; i < count; ++i) { - remove_unlocked_with_neighbors(internal_ids[i], std::move(old_neighbors[i])); - } - } - if (maybe_safety_net_reconnect()) { - page_io_->flush_dirty_pages(); - } + remove_batch_locked(internal_ids, count, nullptr, false); } /** @@ -1029,6 +1093,25 @@ class DiskANNIndex { /// neighbor; without the cap the pool inflates to the deleted node's whole /// ex-neighborhood and reconnect cost explodes. static constexpr uint32_t kTwoHopBypassPerDeleted = 5; + // Fresh slots allocated beyond load-time capacity+headroom are silently + // untracked by the guarded in-degree helpers; brand-new slots are not garden + // targets. + static constexpr uint64_t kInDegreeHeadroomSlots = 1ULL << 20; + /// Load-time recount ids per wave. Each wave backs its unique pages with ONE + /// contiguous aligned buffer and keeps them all in flight (the reactor + /// submits in ring-depth chunks regardless), so the bound is wave-buffer + /// memory, not SQ depth: 8192 ids cap the worst case (one node per page, + /// high-dim) at 32 MiB while staying far above the ring depth, which is what + /// saturates the one-time scan. + static constexpr uint32_t kInDegreeRecountChunk = 8192; + /// when_all fan-out per repair/garden wave. Chunking at the worker count + /// re-created the intra-batch barrier convoy batch_insert removed (each + /// chunk waits on its slowest member and the AsyncGate never fills); one + /// unbounded when_all would queue every prefetch wave ahead of the first + /// resume and churn the whole page LRU. 1024 amortizes the barrier ~32x + /// past the worker count while in-flight fast-path waves (~1 page each) + /// stay a fraction of the default 4096-page cache. + static constexpr uint32_t kUpdateWaveChunk = 1024; /// Yi heap-caps the reconnect pool at build_k = degree + 32 before pruning. static constexpr uint32_t kReconnectPoolSlack = 32; @@ -1081,6 +1164,10 @@ class DiskANNIndex { update_search_l_ = params.update_search_l != 0 ? params.update_search_l : max_degree_ + 32; update_rerank_ = params.update_rerank; update_insert_prune_ = params.update_insert_prune; + update_repair_ = params.update_repair; + repair_search_l_ = params.repair_search_l; + track_in_degree_ = params.track_in_degree; + garden_search_l_ = params.garden_search_l; safety_net_ratio_ = params.safety_net_ratio; safety_net_ops_ = params.safety_net_ops; if (params.update_reconnect_threads == 0) { @@ -1149,6 +1236,26 @@ class DiskANNIndex { slot_alloc_.reset(static_cast(max_slot_id_)); } updatable_ = true; + if (track_in_degree_) { + in_degree_size_ = static_cast(slot_alloc_.next_fresh_id()) + kInDegreeHeadroomSlots; + in_degree_overflow_warned_ = false; + in_degree_ = std::make_unique[]>(static_cast(in_degree_size_)); + for (uint64_t i = 0; i < in_degree_size_; ++i) { + in_degree_[i].store(0, std::memory_order_relaxed); + } + const auto scan_start = std::chrono::steady_clock::now(); + const std::vector counts = recount_in_degree(nullptr); + for (uint64_t i = 0; i < counts.size() && i < in_degree_size_; ++i) { + in_degree_[i].store(counts[static_cast(i)], std::memory_order_relaxed); + } + const uint64_t scan_us = + static_cast(std::chrono::duration_cast( + std::chrono::steady_clock::now() - scan_start) + .count()); + spdlog::info("DiskANNIndex::load: in-degree recount scanned {} slots in {} us", + counts.size(), + scan_us); + } } /// The unified-pool handle searches should use, or nullptr. Non-null only @@ -1263,23 +1370,236 @@ class DiskANNIndex { co_return run_update_search(query, l); } + void indeg_inc(uint32_t internal_id) { + if (!track_in_degree_ || !in_degree_ || internal_id >= in_degree_size_) { + return; + } + in_degree_[internal_id].fetch_add(1, std::memory_order_relaxed); + } + + void indeg_dec(uint32_t internal_id) { + if (!track_in_degree_ || !in_degree_ || internal_id >= in_degree_size_) { + return; + } + in_degree_[internal_id].fetch_sub(1, std::memory_order_relaxed); + } + + std::vector recount_in_degree(coro::thread_pool *pool) { + if (!updatable_ || page_io_ == nullptr) { + return {}; + } + if (pool == nullptr) { + const uint32_t workers = std::max(1, update_insert_threads_); + coro::thread_pool local_pool{{.thread_count = workers, + .on_thread_start_functor = nullptr, + .on_thread_stop_functor = nullptr}}; + try { + std::vector out = recount_in_degree(&local_pool); + local_pool.shutdown(); + return out; + } catch (...) { + local_pool.shutdown(); + throw; + } + } + + const uint32_t next_fresh = slot_alloc_.next_fresh_id(); + std::vector counts(next_fresh, 0); + std::vector ids; + ids.reserve(kInDegreeRecountChunk); + for (uint32_t begin = 0; begin < next_fresh; begin += kInDegreeRecountChunk) { + const uint32_t end = std::min(next_fresh, begin + kInDegreeRecountChunk); + ids.clear(); + for (uint32_t id = begin; id < end; ++id) { + if (!slot_alloc_.is_deleted(id)) { + ids.push_back(id); + } + } + if (ids.empty()) { + continue; + } + const std::vector> nbrs = + read_delete_neighbors(ids.data(), static_cast(ids.size()), pool); + for (const auto &list : nbrs) { + for (const uint32_t target : list) { + if (target < counts.size()) { + ++counts[target]; + } + } + } + } + return counts; + } + + static InDegreePercentiles percentiles_from_values(std::vector values) { + if (values.empty()) { + return {}; + } + std::sort(values.begin(), values.end()); + const auto rank = [&values](uint32_t percentile) { + return values[((values.size() - 1) * percentile) / 100]; + }; + InDegreePercentiles out; + out.p10 = rank(10); + out.p50 = rank(50); + out.p90 = rank(90); + return out; + } + + static InDegreePercentiles percentiles_from_pairs( + const std::vector> &pairs) { + std::vector values; + values.reserve(pairs.size()); + for (const auto &entry : pairs) { + values.push_back(entry.first); + } + return percentiles_from_values(std::move(values)); + } + + std::vector> collect_live_in_degree_pairs( + bool include_medoid) const { + std::vector> out; + if (!track_in_degree_ || !in_degree_) { + return out; + } + std::shared_lock state_lock(update_mutex_); + const uint32_t next_fresh = slot_alloc_.next_fresh_id(); + out.reserve(static_cast(std::min(live_count_, next_fresh))); + for (uint32_t id = 0; id < next_fresh && id < in_degree_size_; ++id) { + if (!include_medoid && id == medoid_) { + continue; + } + if (!slot_alloc_.is_deleted(id)) { + out.emplace_back(in_degree_[id].load(std::memory_order_relaxed), id); + } + } + return out; + } + + GardenStats garden_refresh_locked(uint32_t budget, coro::thread_pool &pool) { + GardenStats stats; + std::vector> candidates = collect_live_in_degree_pairs(false); + const InDegreePercentiles before = percentiles_from_pairs(candidates); + stats.indeg_p10_before = before.p10; + stats.indeg_p50_before = before.p50; + + const size_t selected = std::min(budget, candidates.size()); + stats.selected = selected; + if (selected == 0) { + const InDegreePercentiles after = percentiles_from_pairs(candidates); + stats.indeg_p10_after = after.p10; + stats.indeg_p50_after = after.p50; + return stats; + } + if (selected < candidates.size()) { + std::nth_element(candidates.begin(), candidates.begin() + selected, candidates.end()); + candidates.resize(selected); + } + + std::vector picked; + picked.reserve(candidates.size()); + for (const auto &entry : candidates) { + picked.push_back(entry.second); + } + + std::atomic refreshed{0}; + auto refresh_one = [this, &pool, &refreshed](uint32_t node_id) -> coro::task<> { + co_await pool.schedule(); + auto mark = std::chrono::steady_clock::now(); + DiskPageIO::NodeData nd; + if (page_io_->reactor_enabled()) { + nd = co_await page_io_->read_node_async(node_id, pool); + } else { + nd = page_io_->read_node(node_id); + } + const uint32_t l = garden_search_l_ != 0 ? garden_search_l_ : update_search_l_; + std::vector> cand = + co_await run_update_search_async(nd.coords.data(), l, pool); + cand.erase(std::remove_if(cand.begin(), + cand.end(), + [node_id](const auto &entry) { + return entry.first == node_id; + }), + cand.end()); + co_await wave_selection_coords(cand, pool); + std::vector pruned = + select_insert_neighbors_from(nd.coords.data(), std::move(cand)); + if (pruned.empty()) { + co_return; + } + + if (page_io_->reactor_enabled()) { + // One wave warms every page this refresh will RMW — the node's own + // plus each pruned neighbor's (the reverse-edge loop rewrites them) — + // mirroring insert_one's single warm wave. + std::vector warm; + warm.reserve(pruned.size() + 1); + warm.push_back(node_id); + warm.insert(warm.end(), pruned.begin(), pruned.end()); + co_await page_io_->prefetch_pages(warm.data(), warm.size(), pool); + } + + { + std::lock_guard node_lock(update_node_mutex(node_id)); + update_node_impl(node_id, pruned); + } + + const std::vector extra{node_id}; + for (const uint32_t neighbor : pruned) { + if (!has_pq_) { + co_await prefetch_reconnect_inputs(neighbor, extra, pool); + } + std::lock_guard node_lock(update_node_mutex(neighbor)); + update_node_impl(neighbor, extra); + } + + st_garden_us_.fetch_add(stage_us_since(mark), std::memory_order_relaxed); + st_gardens_.fetch_add(1, std::memory_order_relaxed); + refreshed.fetch_add(1, std::memory_order_relaxed); + }; + + for (uint32_t off = 0; off < picked.size(); off += kUpdateWaveChunk) { + const uint32_t end = + std::min(static_cast(picked.size()), off + kUpdateWaveChunk); + auto run = [&]() -> coro::task<> { + std::vector> tasks; + tasks.reserve(end - off); + for (uint32_t i = off; i < end; ++i) { + tasks.emplace_back(refresh_one(picked[i])); + } + co_await coro::when_all(std::move(tasks)); + }; + coro::sync_wait(run()); + } + + stats.refreshed = refreshed.load(std::memory_order_relaxed); + const InDegreePercentiles after = percentiles_from_pairs(collect_live_in_degree_pairs(false)); + stats.indeg_p10_after = after.p10; + stats.indeg_p50_after = after.p50; + return stats; + } + public: /// Wall-microseconds by update stage, aggregated across coroutines since the /// last take (wall, not CPU: a stage that yields through the pool queue books /// its scheduling latency here — that is the point of measuring it). struct UpdateStageStats { - uint64_t gate_us = 0; ///< waiting for a search ThreadData - uint64_t greedy_us = 0; ///< async greedy search proper - uint64_t search_us = 0; ///< whole selection stage (incl. gate+greedy) - uint64_t alloc_us = 0; ///< slot alloc + PQ encode - uint64_t prefetch_us = 0; ///< slot page warm wave - uint64_t write_us = 0; ///< write_inserted_node + publish + staging - uint64_t reconnect_us = 0; ///< when_all over the insert's reconnects + uint64_t gate_us = 0; ///< waiting for a search ThreadData + uint64_t greedy_us = 0; ///< async greedy search proper + uint64_t search_us = 0; ///< whole selection stage (incl. gate+greedy) + uint64_t alloc_us = 0; ///< slot alloc + PQ encode + uint64_t prefetch_us = 0; ///< slot page warm wave + uint64_t write_us = 0; ///< write_inserted_node + publish + staging + uint64_t reconnect_us = 0; ///< when_all over the insert's reconnects + uint64_t repair_us = 0; ///< delete-time discovery + repair wave + uint64_t garden_us = 0; uint64_t rc_prefetch_us = 0; ///< reconnect input warm waves uint64_t rc_lock_us = 0; ///< node mutex acquisition wait uint64_t rc_impl_us = 0; ///< sync update_node_impl body uint64_t inserts = 0; uint64_t reconnects = 0; + uint64_t repairs = 0; + uint64_t gardens = 0; }; /// Snapshot-and-reset the stage stats (benchmark instrumentation). @@ -1292,11 +1612,15 @@ class DiskANNIndex { out.prefetch_us = st_prefetch_us_.exchange(0, std::memory_order_acq_rel); out.write_us = st_write_us_.exchange(0, std::memory_order_acq_rel); out.reconnect_us = st_reconnect_us_.exchange(0, std::memory_order_acq_rel); + out.repair_us = st_repair_us_.exchange(0, std::memory_order_acq_rel); + out.garden_us = st_garden_us_.exchange(0, std::memory_order_acq_rel); out.rc_prefetch_us = st_rc_prefetch_us_.exchange(0, std::memory_order_acq_rel); out.rc_lock_us = st_rc_lock_us_.exchange(0, std::memory_order_acq_rel); out.rc_impl_us = st_rc_impl_us_.exchange(0, std::memory_order_acq_rel); out.inserts = st_inserts_.exchange(0, std::memory_order_acq_rel); out.reconnects = st_reconnects_.exchange(0, std::memory_order_acq_rel); + out.repairs = st_repairs_.exchange(0, std::memory_order_acq_rel); + out.gardens = st_gardens_.exchange(0, std::memory_order_acq_rel); return out; } @@ -1431,27 +1755,36 @@ class DiskANNIndex { return update_insert_prune_; } - /// select_insert_neighbors with the disk work made awaitable: the candidate - /// coords the selection will score are wave-prefetched through the reactor - /// (one suspension, all misses in flight together) before the sync selection - /// logic runs against warm caches. + /// Wave-prefetch the candidate coords an insert-style selection will score + /// (one suspension, all misses in flight together), so the sync selection + /// logic runs against warm caches. No-op when the selection variant does no + /// exact-coords work (see insert_selection_needs_coords). + coro::task<> wave_selection_coords(const std::vector> &cand, + coro::thread_pool &pool) { + if (!page_io_->reactor_enabled() || !insert_selection_needs_coords() || cand.empty()) { + co_return; + } + std::vector want; + want.reserve(cand.size()); + for (const auto &c : cand) { + // exact_query_distance (PQ rerank) is served by the NodeCache first; + // cached_l2 (No-PQ prune) always goes through the coords cache. + if (!has_pq_ || !cache_.lookup_record(c.first)) { + want.push_back(c.first); + } + } + if (!want.empty()) { + co_await page_io_->prefetch_coords(want.data(), want.size(), pool); + } + co_return; + } + + /// select_insert_neighbors with the disk work made awaitable via the + /// selection coords wave above. coro::task> select_insert_neighbors_async(const float *query, coro::thread_pool &pool) { auto cand = co_await run_update_search_async(query, update_search_l_, pool); - if (page_io_->reactor_enabled() && insert_selection_needs_coords() && !cand.empty()) { - std::vector want; - want.reserve(cand.size()); - for (const auto &c : cand) { - // exact_query_distance (PQ rerank) is served by the NodeCache first; - // cached_l2 (No-PQ prune) always goes through the coords cache. - if (!has_pq_ || !cache_.lookup_record(c.first)) { - want.push_back(c.first); - } - } - if (!want.empty()) { - co_await page_io_->prefetch_coords(want.data(), want.size(), pool); - } - } + co_await wave_selection_coords(cand, pool); co_return select_insert_neighbors_from(query, std::move(cand)); } @@ -1462,6 +1795,14 @@ class DiskANNIndex { uint32_t allocate_update_slot_unlocked(uint64_t label) { const uint32_t slot = slot_alloc_.alloc(); update_ctx_.forget_slot(slot); + if (track_in_degree_ && slot >= in_degree_size_ && !in_degree_overflow_warned_) { + in_degree_overflow_warned_ = true; + spdlog::warn( + "DiskANNIndex: slot {} exceeds in-degree tracking capacity {}; slots past it stay " + "untracked (counters no-op, garden_refresh ignores them) until the next load", + slot, + in_degree_size_); + } max_slot_id_ = std::max(max_slot_id_, slot_alloc_.next_fresh_id()); set_label(slot, label); ++live_count_; @@ -1557,7 +1898,239 @@ class DiskANNIndex { return page_io_->read_neighbors_batch_parallel(internal_ids, count, update_insert_threads_); } + /// Full delete records for repair. Reactor mode warms each chunk with one + /// page wave, then uses awaitable node reads for any cache-disabled misses. + std::vector read_delete_nodes(const uint32_t *internal_ids, + uint32_t count, + coro::thread_pool &pool) { + std::vector out(count); + // The warm wave's pages must survive in the LRU until the chunk's reads + // land, so the chunk is bounded by half the cache (worker-count chunks + // made a 10k-delete batch pay ~300 tiny waves). Cache off => prefetch + // no-ops and the chunk only bounds the read fan-out. + const uint32_t chunk_size = + std::max(1, + page_io_->page_cache_enabled() + ? static_cast( + std::min(page_io_->page_cache_capacity() / 2, + kUpdateWaveChunk)) + : kUpdateWaveChunk); + for (uint32_t off = 0; off < count; off += chunk_size) { + const uint32_t n = std::min(count - off, chunk_size); + if (page_io_->reactor_enabled()) { + auto warm = [&]() -> coro::task<> { + co_await page_io_->prefetch_pages(internal_ids + off, n, pool); + }; + coro::sync_wait(warm()); + } + std::vector chunk = + page_io_->read_nodes_async(internal_ids + off, n, pool); + std::move(chunk.begin(), chunk.end(), out.begin() + off); + } + return out; + } + + /// Mark the whole delete batch dead before any repair search snapshots + /// tombstones; the slots are not allocatable until release_removed_batch(). + void mark_removed_batch(const uint32_t *internal_ids, + const std::vector &old_nodes, + uint32_t count) { + std::unique_lock state_lock(update_mutex_); + for (uint32_t i = 0; i < count; ++i) { + update_ctx_.removed_node_nbrs_[internal_ids[i]] = old_nodes[i].nbrs; + if (track_in_degree_) { + for (const uint32_t target : old_nodes[i].nbrs) { + indeg_dec(target); + } + } + slot_alloc_.mark_removed(internal_ids[i]); + --live_count_; + ++ops_since_last_insert_; + } + } + + /// End the repair window: after in-neighbors have been rebuilt, inserts may + /// reuse the tombstoned slots through the allocator free list. + void release_removed_batch(const uint32_t *internal_ids, uint32_t count) { + std::unique_lock state_lock(update_mutex_); + for (uint32_t i = 0; i < count; ++i) { + slot_alloc_.release(internal_ids[i]); + } + } + + /// Approximate the deleted batch's in-neighbors from old out-neighbors and, + /// optionally, a tombstone-aware insert-style search around each deleted + /// vector. The returned ids are live and unique across the batch. + std::vector discover_delete_repair_targets( + const std::vector &old_nodes, + uint32_t count, + coro::thread_pool &pool) { + std::unordered_set unique; + unique.reserve(static_cast(count) * + static_cast(std::max(1, max_degree_))); + + if (repair_search_l_ == 0) { + // Symmetry-only discovery reads nothing: the batch's old out-neighbors + // are already in memory, so skip the coroutine machinery outright. + for (uint32_t i = 0; i < count; ++i) { + for (const uint32_t nbr : old_nodes[i].nbrs) { + if (!slot_alloc_.is_deleted(nbr)) { + unique.insert(nbr); + } + } + } + return {unique.begin(), unique.end()}; + } + + std::vector> discovered(count); + auto discover_one = [this, &old_nodes, &discovered, &pool](uint32_t i) -> coro::task<> { + co_await pool.schedule(); + std::vector &out = discovered[i]; + out.reserve(old_nodes[i].nbrs.size() + repair_search_l_); + const auto push = [&](uint32_t v) { + if (!slot_alloc_.is_deleted(v)) { + out.push_back(v); + } + }; + for (const uint32_t nbr : old_nodes[i].nbrs) { + push(nbr); + } + const std::vector> search = + co_await run_update_search_async(old_nodes[i].coords.data(), repair_search_l_, pool); + for (const auto &cand : search) { + push(cand.first); + } + }; + + for (uint32_t off = 0; off < count; off += kUpdateWaveChunk) { + const uint32_t end = std::min(count, off + kUpdateWaveChunk); + auto run = [&]() -> coro::task<> { + std::vector> tasks; + tasks.reserve(end - off); + for (uint32_t i = off; i < end; ++i) { + tasks.emplace_back(discover_one(i)); + } + co_await coro::when_all(std::move(tasks)); + }; + coro::sync_wait(run()); + for (uint32_t i = off; i < end; ++i) { + for (const uint32_t id : discovered[i]) { + unique.insert(id); + } + discovered[i].clear(); + } + } + + std::vector targets; + targets.reserve(unique.size()); + for (const uint32_t id : unique) { + targets.push_back(id); + } + return targets; + } + + /// Rebuild each discovered in-neighbor once while the deleted slots are still + /// tombstoned, so update_node_impl can drop dead edges and two-hop splice. + void repair_removed_batch(const std::vector &old_nodes, + uint32_t count, + coro::thread_pool &pool) { + auto mark = std::chrono::steady_clock::now(); + const std::vector targets = discover_delete_repair_targets(old_nodes, count, pool); + + auto repair_one = [this, &pool](uint32_t node_id) -> coro::task<> { + co_await pool.schedule(); + const std::vector no_extra; + co_await prefetch_reconnect_inputs(node_id, no_extra, pool); + update_node_impl_locked(node_id, no_extra); + }; + + for (uint32_t off = 0; off < targets.size(); off += kUpdateWaveChunk) { + const uint32_t end = + std::min(static_cast(targets.size()), off + kUpdateWaveChunk); + auto run = [&]() -> coro::task<> { + std::vector> tasks; + tasks.reserve(end - off); + for (uint32_t i = off; i < end; ++i) { + tasks.emplace_back(repair_one(targets[i])); + } + co_await coro::when_all(std::move(tasks)); + }; + coro::sync_wait(run()); + } + st_repair_us_.fetch_add(stage_us_since(mark), std::memory_order_relaxed); + st_repairs_.fetch_add(targets.size(), std::memory_order_relaxed); + } + + /// Shared remove tail. With delete-time repair disabled this preserves the + /// existing lazy-delete sequence; with it enabled, slots stay dark until the + /// discovered in-neighbors have been rebuilt. + void remove_batch_locked(const uint32_t *internal_ids, + uint32_t count, + coro::thread_pool *pool, + bool single_remove) { + if (update_repair_ && pool == nullptr) { + // Same worker count with or without the reactor: repair runs per-node + // discovery searches and RMWs, the shape batch_insert drives with a + // full worker pool over the reactor (not read_delete_neighbors' single + // batched wave, where one thread suffices). + const uint32_t workers = + std::min(count, std::max(1, update_insert_threads_)); + coro::thread_pool local_pool{{.thread_count = workers, + .on_thread_start_functor = nullptr, + .on_thread_stop_functor = nullptr}}; + try { + remove_batch_locked(internal_ids, count, &local_pool, single_remove); + local_pool.shutdown(); + return; + } catch (...) { + local_pool.shutdown(); + throw; + } + } + + page_io_->clear_cache(); + if (!update_repair_) { + if (single_remove) { + remove_unlocked(internal_ids[0]); + } else { + std::vector> old_neighbors = + read_delete_neighbors(internal_ids, count, pool); + std::unique_lock state_lock(update_mutex_); + for (uint32_t i = 0; i < count; ++i) { + remove_unlocked_with_neighbors(internal_ids[i], std::move(old_neighbors[i])); + } + } + if (maybe_safety_net_reconnect()) { + page_io_->flush_dirty_pages(); + } + return; + } + + std::vector old_nodes = read_delete_nodes(internal_ids, count, *pool); + mark_removed_batch(internal_ids, old_nodes, count); + try { + repair_removed_batch(old_nodes, count, *pool); + } catch (...) { + // Degrade to lazy-delete semantics rather than leak the slots: they are + // tombstoned with their two-hop cache in place, so the free list may + // reuse them like any lazy delete. + release_removed_batch(internal_ids, count); + throw; + } + release_removed_batch(internal_ids, count); + // No safety net here: it strips dead edges from the out-neighbors of every + // outstanding tombstone, and repair just rebuilt a superset of exactly + // those nodes (discovery = out-neighbors + optional search) — under the + // lazy path's default arming (5% tombstones, 16 deletes) it would re-scan + // the whole history after every batch and never find an edge to fix. + } + void remove_unlocked_with_neighbors(uint32_t internal_id, std::vector old_neighbors) { + if (track_in_degree_) { + for (const uint32_t target : old_neighbors) { + indeg_dec(target); + } + } update_ctx_.removed_node_nbrs_[internal_id] = std::move(old_neighbors); slot_alloc_.free(internal_id); --live_count_; @@ -1569,6 +2142,11 @@ class DiskANNIndex { const std::vector &neighbors) { page_io_->write_node(slot, query, static_cast(neighbors.size()), neighbors.data()); cache_.upsert_node(slot, query, static_cast(neighbors.size()), neighbors.data()); + if (track_in_degree_) { + for (const uint32_t target : neighbors) { + indeg_inc(target); + } + } } /// Staged reverse edges for a batch — Yi's `_inserted_edges` analog. Each @@ -1764,9 +2342,9 @@ class DiskANNIndex { return prune_candidate_pool_with_dist(node_id, pool, dist_fn); } - /// Shared reconnect backbone (Yi's co_update tail): pools at or under the - /// degree bound are kept verbatim with no distance work; larger pools are - /// scored, capped to degree+32 nearest (Yi's build_k heap), and alpha-pruned. + /// Shared reconnect backbone (Yi's co_update tail): update_node_impl* paths + /// funnel here; because only the write branch changes neighbor lists, the + /// in-degree counter diff belongs in that branch. void prune_and_write(uint32_t node_id, const std::vector &old_nbrs, const std::vector &cand) { @@ -1786,6 +2364,14 @@ class DiskANNIndex { page_io_->write_node_neighbors(node_id, static_cast(new_nbrs.size()), new_nbrs.data()); + if (track_in_degree_) { + for (const uint32_t target : old_nbrs) { + indeg_dec(target); + } + for (const uint32_t target : new_nbrs) { + indeg_inc(target); + } + } mirror_neighbors_to_cache(node_id, new_nbrs); } } @@ -2026,6 +2612,14 @@ class DiskANNIndex { } if (live.size() != nd.nbrs.size()) { page_io_->write_node_neighbors(nid, static_cast(live.size()), live.data()); + if (track_in_degree_) { + for (const uint32_t target : nd.nbrs) { + indeg_dec(target); + } + for (const uint32_t target : live) { + indeg_inc(target); + } + } mirror_neighbors_to_cache(nid, live); } } @@ -2099,6 +2693,11 @@ class DiskANNIndex { page_io_.reset(); update_reactor_.reset(); // after page_io_: it holds a raw pointer to the reactor update_ctx_.clear(); + in_degree_.reset(); + in_degree_size_ = 0; + track_in_degree_ = false; + in_degree_overflow_warned_ = false; + garden_search_l_ = 0; updatable_ = false; loaded_ = false; } @@ -2269,7 +2868,15 @@ class DiskANNIndex { bool updatable_ = false; uint64_t max_slot_id_ = 0; ///< file capacity in slots (valid-id bound; only grows) uint64_t live_count_ = 0; ///< live (non-tombstoned) vector count - std::string index_dir_; ///< saved at load for flush() output paths + /// in_degree_[s] == number of edges from non-tombstoned nodes' current + /// on-disk neighbor lists that point at s; reused slots keep stale inbound + /// counts until their live owners rewrite. + std::unique_ptr[]> in_degree_; + uint64_t in_degree_size_ = 0; + bool track_in_degree_ = false; + bool in_degree_overflow_warned_ = false; + uint32_t garden_search_l_ = 0; + std::string index_dir_; ///< saved at load for flush() output paths std::unique_ptr update_reactor_; ///< declared before page_io_ so the ///< page IO (raw-pointer user) dies first std::unique_ptr page_io_; @@ -2286,6 +2893,8 @@ class DiskANNIndex { uint32_t update_search_l_ = 100; bool update_rerank_ = true; bool update_insert_prune_ = false; + bool update_repair_ = false; + uint32_t repair_search_l_ = 0; double safety_net_ratio_ = 0.05; uint64_t safety_net_ops_ = 16; uint32_t update_insert_threads_ = kDefaultDiskANNUpdateInsertThreads; @@ -2299,11 +2908,15 @@ class DiskANNIndex { std::atomic st_prefetch_us_{0}; std::atomic st_write_us_{0}; std::atomic st_reconnect_us_{0}; + std::atomic st_repair_us_{0}; + std::atomic st_garden_us_{0}; std::atomic st_rc_prefetch_us_{0}; std::atomic st_rc_lock_us_{0}; std::atomic st_rc_impl_us_{0}; std::atomic st_inserts_{0}; std::atomic st_reconnects_{0}; + std::atomic st_repairs_{0}; + std::atomic st_gardens_{0}; static uint64_t stage_us_since(std::chrono::steady_clock::time_point &mark) { const auto now = std::chrono::steady_clock::now(); diff --git a/include/index/graph/diskann/slot_allocator.hpp b/include/index/graph/diskann/slot_allocator.hpp index 19265931..759110b1 100644 --- a/include/index/graph/diskann/slot_allocator.hpp +++ b/include/index/graph/diskann/slot_allocator.hpp @@ -16,6 +16,10 @@ * (mirrors Yi, which adds a node to its live set only after the disk * append). * - `free(id)` pushes the slot onto the free list and tombstones it. + * - delete-time graph repair can split that last step: `mark_removed(id)` + * makes searches treat the slot as dead while it remains unavailable for + * reuse, then `release(id)` makes it allocatable after all in-neighbors + * have been patched. * * The allocator owns the `TombstoneBitmap` so that alloc/free keep liveness and * reuse in lockstep, and so a single `save()`/`load()` round-trips the complete @@ -77,6 +81,19 @@ class SlotAllocator { tombstone_.set(id); } + /// Tombstone @p id without putting it on the free list. Delete-time repair + /// must patch in-neighbors during this window, before the slot can be reused. + void mark_removed(uint32_t id) { tombstone_.set(id); } + + /// Make a mark_removed() slot reusable after its in-neighbor repair window. + /// Re-asserts the tombstone (a no-op in the intended mark_removed() -> + /// release() sequence) so misuse on a live slot can never leave it + /// simultaneously searchable and allocatable. + void release(uint32_t id) { + tombstone_.set(id); + free_list_.push_back(id); + } + [[nodiscard]] bool is_deleted(uint32_t id) const { return tombstone_.is_deleted(id); } [[nodiscard]] uint32_t next_fresh_id() const { return next_fresh_id_; } [[nodiscard]] uint64_t free_count() const { return free_list_.size(); } diff --git a/tests/diskann/bench_diskann_sift_update.cpp b/tests/diskann/bench_diskann_sift_update.cpp index 889320f6..23db8222 100644 --- a/tests/diskann/bench_diskann_sift_update.cpp +++ b/tests/diskann/bench_diskann_sift_update.cpp @@ -95,6 +95,8 @@ struct Options { bool mixed = false; bool update_rerank = true; ///< Yi trace-bench parity (_rerank_flag=true) bool update_insert_prune = false; ///< Yi never alpha-prunes at insert + bool update_repair = false; ///< delete-time in-neighbor repair + bool track_indegree = false; ///< track live-slot in-degree percentiles bool mixed_round0_baseline = true; ///< Yi UpdateRunner: round 0 runs no updates MixedMode mixed_mode = MixedMode::Background; uint32_t max_rounds = 0; @@ -107,6 +109,9 @@ struct Options { uint32_t search_l = 100; uint32_t rerank_count = 0; uint32_t update_l = kDefaultBenchmarkUpdateSearchL; + uint32_t repair_search_l = 0; ///< 0 => deleted-node out-neighbors only + uint32_t garden_budget = 0; ///< per-round low-in-degree refresh budget + uint32_t garden_l = 0; ///< 0 => update_l uint32_t beam = 4; uint32_t search_threads = 1; uint32_t eval_pipeline = 0; ///< >1: eval via search_pipelined with this many @@ -124,6 +129,17 @@ struct Options { alaya::diskann::DiskANNUpdateIO update_io = alaya::diskann::DiskANNUpdateIO::kAuto; uint32_t update_search_concurrency = 0; // 0 = library default (4x insert threads) bool search_page_cache = true; // searches peek+fill the shard page cache + std::string oracle_ids; ///< when set: batch-rebuild the graph on exactly this live + ///< label set (ids.bin format), then eval once — the + ///< "best achievable graph on the survivor set" reference + ///< that separates intrinsic geometry floor from + ///< incremental-process slack. + uint32_t garden_pass = 0; ///< when set with --eval_only on a kept index: run + ///< garden_refresh(garden_pass) this many nodes per + ///< iteration and re-eval, to test whether broader + ///< garden COVERAGE (vs a full rebuild) recovers the + ///< residual — coverage lever vs global-structural. + uint32_t garden_pass_iters = 3; ///< garden passes to run in --garden_pass mode. }; DatasetFiles resolve_dataset_files(const std::filesystem::path &data_dir) { @@ -238,6 +254,34 @@ MixedMode parse_mixed_mode(const std::string &value) { throw std::invalid_argument("bad --mixed_mode: " + value); } +void print_usage(const char *argv0) { + std::cerr << "Usage:\n" + << " " << argv0 << " [data_dir] [trace_dir] [index_dir] [out_csv] [flags]\n\n" + << "Flags:\n" + << " --rebuild force a fresh index build\n" + << " --deterministic deterministic search barriers\n" + << " --flush_rounds flush dirty pages after each round\n" + << " --no_flush_rounds skip per-round dirty-page flush\n" + << " --no_update_rerank disable update rerank\n" + << " --update_insert_prune alpha-prune insert candidates\n" + << " --update_repair enable delete-time in-neighbor repair\n" + << " --no_update_repair disable delete-time in-neighbor repair\n" + << " --repair_search_l N repair search list (0 => deleted-node out-neighbors)\n" + << " --track_indegree track live-slot in-degree percentiles\n" + << " --garden_budget N refresh up to N low-in-degree live nodes per round\n" + << " --garden_l N garden search list (0 => use --update_L)\n" + << " --garden_pass N eval_only: garden_refresh(N) per pass, then re-eval\n" + << " --garden_pass_iters N garden passes in --garden_pass mode (default 3)\n" + << " --oracle_ids FILE batch-rebuild on FILE's ids.bin label set, eval once\n" + << " --build_only build the initial index and exit\n" + << " --eval_only replay live mask and run eval only\n" + << " --single_updates issue single remove/insert calls\n" + << " --mixed run mixed update/search workload\n" + << " --mixed_mode MODE background or shared_queue\n" + << " --rounds N cap update rounds\n" + << " --nq N cap eval queries (0 => all)\n"; +} + TraceManifest read_manifest(const std::filesystem::path &path) { std::ifstream in(path); if (!in) { @@ -299,12 +343,25 @@ Options parse_args(int argc, char **argv) { opt.update_rerank = false; } else if (arg == "--update_insert_prune") { opt.update_insert_prune = true; + } else if (arg == "--update_repair") { + opt.update_repair = true; + } else if (arg == "--no_update_repair") { + opt.update_repair = false; + } else if (arg == "--track_indegree") { + opt.track_indegree = true; } else if (arg == "--no_mixed_round0") { opt.mixed_round0_baseline = false; } else if (arg == "--build_only") { opt.build_only = true; } else if (arg == "--eval_only") { opt.eval_only = true; + } else if (arg == "--oracle_ids" && i + 1 < argc) { + opt.oracle_ids = argv[++i]; + } else if (arg == "--garden_pass" && i + 1 < argc) { + opt.garden_pass = parse_u32(argv[++i], "--garden_pass"); + opt.track_indegree = true; + } else if (arg == "--garden_pass_iters" && i + 1 < argc) { + opt.garden_pass_iters = parse_u32(argv[++i], "--garden_pass_iters"); } else if (arg == "--single_updates") { opt.single_updates = true; } else if (arg == "--mixed") { @@ -328,6 +385,12 @@ Options parse_args(int argc, char **argv) { opt.rerank_count = parse_u32(argv[++i], "--rerank_count"); } else if (arg == "--update_L" && i + 1 < argc) { opt.update_l = parse_u32(argv[++i], "--update_L"); // 0 => R + 32 + } else if (arg == "--repair_search_l" && i + 1 < argc) { + opt.repair_search_l = parse_u32(argv[++i], "--repair_search_l"); + } else if (arg == "--garden_budget" && i + 1 < argc) { + opt.garden_budget = parse_u32(argv[++i], "--garden_budget"); + } else if (arg == "--garden_l" && i + 1 < argc) { + opt.garden_l = parse_u32(argv[++i], "--garden_l"); } else if (arg == "--beam" && i + 1 < argc) { opt.beam = parse_u32(argv[++i], "--beam"); } else if (arg == "--threads" && i + 1 < argc) { @@ -388,6 +451,9 @@ Options parse_args(int argc, char **argv) { if (pos.size() > 3) { opt.out_csv = pos[3]; } + if (opt.garden_budget > 0) { + opt.track_indegree = true; + } return opt; } @@ -404,6 +470,36 @@ TraceRound limit_round_updates(TraceRound round, const Options &opt) { return round; } +// ids.bin format (matches the on-disk index and fig13_forensics.load_ids): +// [uint64 count][count x uint64 label]. Used to seed the oracle rebuild's live +// label set directly from a kept index's ids.bin. +std::vector read_ids_file(const std::string &path) { + std::ifstream in(path, std::ios::binary); + if (!in) { + throw std::runtime_error("cannot open oracle ids file: " + path); + } + uint64_t count = 0; + in.read(reinterpret_cast(&count), sizeof(count)); + if (!in) { + throw std::runtime_error("oracle ids file truncated header: " + path); + } + // Validate the header count against the actual file size BEFORE allocating: + // a corrupt count would otherwise request a giant vector up front. + const uint64_t payload_bytes = + static_cast(std::filesystem::file_size(path)) - sizeof(uint64_t); + if (count > payload_bytes / sizeof(uint64_t)) { + throw std::runtime_error("oracle ids file count " + std::to_string(count) + + " exceeds file size: " + path); + } + std::vector labels(count); + in.read(reinterpret_cast(labels.data()), + static_cast(count * sizeof(uint64_t))); + if (!in) { + throw std::runtime_error("oracle ids file truncated body: " + path); + } + return labels; +} + void build_initial_index(const Options &opt, const FloatMatrix &base, const TraceManifest &manifest) { @@ -415,10 +511,6 @@ void build_initial_index(const Options &opt, std::cout << "[update_bench] reusing index " << opt.index_dir << "\n"; return; } - std::vector labels(manifest.initial_count); - for (uint32_t id = 0; id < manifest.initial_count; ++id) { - labels[id] = id; - } DiskANNBuildParams bp; bp.R = opt.build_r; bp.L = opt.build_l; @@ -429,6 +521,33 @@ void build_initial_index(const Options &opt, bp.num_threads = 96; bp.seed = 1234; bp.verbose = true; + + if (!opt.oracle_ids.empty()) { + // Oracle: batch-rebuild the graph on exactly the survivor+inserted label + // set (original numbering preserved), gathering each node's coords from the + // base by label. Preserving labels lets the same GT + age buckets score it. + const std::vector labels = read_ids_file(opt.oracle_ids); + const uint32_t n = static_cast(labels.size()); + std::vector coords(static_cast(n) * base.dim); + for (uint32_t i = 0; i < n; ++i) { + const uint64_t label = labels[i]; + if (label >= base.n) { + throw std::runtime_error("oracle label out of base range: " + std::to_string(label)); + } + std::copy_n(base.data.data() + static_cast(label) * base.dim, + base.dim, + coords.data() + static_cast(i) * base.dim); + } + std::cout << "[update_bench] building ORACLE PQ index (batch rebuild on survivor set), n=" + << n << " dir=" << opt.index_dir << "\n"; + DiskANNIndex::build(opt.index_dir, coords.data(), labels.data(), n, base.dim, bp); + return; + } + + std::vector labels(manifest.initial_count); + for (uint32_t id = 0; id < manifest.initial_count; ++id) { + labels[id] = id; + } std::cout << "[update_bench] building initial PQ index, n=" << manifest.initial_count << " dir=" << opt.index_dir << "\n"; DiskANNIndex::build(opt.index_dir, @@ -1023,6 +1142,13 @@ void run_warmup_searches(const DiskANNIndex &idx, int main(int argc, char **argv) { try { + for (int i = 1; i < argc; ++i) { + const std::string arg = argv[i]; + if (arg == "-h" || arg == "--help") { + print_usage(argv[0]); + return 0; + } + } const Options opt = parse_args(argc, argv); const std::filesystem::path data_dir = opt.data_dir; const std::filesystem::path trace_dir = opt.trace_dir; @@ -1051,6 +1177,10 @@ int main(int argc, char **argv) { lp.update_search_l = opt.update_l; lp.update_rerank = opt.update_rerank; lp.update_insert_prune = opt.update_insert_prune; + lp.update_repair = opt.update_repair; + lp.repair_search_l = opt.repair_search_l; + lp.track_in_degree = opt.track_indegree; + lp.garden_search_l = opt.garden_l; lp.update_insert_threads = opt.update_insert_threads; lp.update_reconnect_threads = opt.update_reconnect_threads; lp.page_cache_capacity = opt.page_cache_capacity; @@ -1062,6 +1192,22 @@ int main(int argc, char **argv) { std::vector live(manifest.total_count, 0); std::vector label_to_slot(manifest.total_count, kMissingSlot); + if (!opt.oracle_ids.empty()) { + // Oracle: the live set is exactly the rebuilt label set. Mark them alive so + // the same masked-recall + age buckets score the batch-rebuilt graph. + const std::vector oracle_labels = read_ids_file(opt.oracle_ids); + for (const uint64_t label : oracle_labels) { + if (label < live.size()) { + live[label] = 1; + } + } + run_warmup_searches(idx, queries, gt, live, opt); + const RecallResult recall = evaluate_search(idx, queries, gt, live, opt); + std::cout << "[oracle] n=" << oracle_labels.size() + << " masked_recall@10=" << recall.recall() + << " search_mean_us=" << recall.mean_us << " search_qps=" << recall.qps << "\n"; + return 0; + } for (uint32_t id = 0; id < manifest.initial_count; ++id) { live[id] = 1; label_to_slot[id] = id; @@ -1079,6 +1225,19 @@ int main(int argc, char **argv) { const RecallResult recall = evaluate_search(idx, queries, gt, live, opt); std::cout << "[eval_only] rounds=" << rounds << " masked_recall@10=" << recall.recall() << " search_mean_us=" << recall.mean_us << " search_qps=" << recall.qps << "\n"; + // Coverage probe: does gardening the WHOLE survivor set (vs garden's 50k/ + // round tail budget) close the residual toward the batch-rebuild ceiling? + // If it plateaus below, the slack is global-structural (rebuild-only); if + // it climbs, garden just needs more coverage budget. + for (uint32_t it = 0; opt.garden_pass > 0 && it < opt.garden_pass_iters; ++it) { + const auto gs = idx.garden_refresh(opt.garden_pass); + const RecallResult r = evaluate_search(idx, queries, gt, live, opt); + std::cout << "[garden_pass " << it << "] refreshed=" << gs.refreshed + << " p10 " << gs.indeg_p10_before << "->" << gs.indeg_p10_after + << " p50 " << gs.indeg_p50_before << "->" << gs.indeg_p50_after + << " ms=" << static_cast(gs.elapsed_us) / 1000.0 + << " masked_recall@10=" << r.recall() << "\n"; + } return 0; } @@ -1145,6 +1304,27 @@ int main(int argc, char **argv) { if (mixed_started) { mixed_search = finish_mixed_search(mixed_state); } + if (opt.garden_budget > 0 && !baseline_round) { + coro::thread_pool *garden_pool = mixed_pool ? mixed_pool.get() : nullptr; + const auto g0 = std::chrono::steady_clock::now(); + auto gs = idx.garden_refresh(opt.garden_budget, garden_pool); + const auto garden_wall_us = + std::chrono::duration_cast( + std::chrono::steady_clock::now() - g0) + .count(); + if (gs.elapsed_us == 0 && garden_wall_us > 0) { + gs.elapsed_us = static_cast(garden_wall_us); + } + std::cout << "[garden] refreshed=" << gs.refreshed << " selected=" << gs.selected + << " p10 " << gs.indeg_p10_before << "->" << gs.indeg_p10_after + << " p50 " << gs.indeg_p50_before << "->" << gs.indeg_p50_after + << " ms=" << static_cast(gs.elapsed_us) / 1000.0 << "\n"; + } + if (opt.track_indegree) { + const auto ip = idx.in_degree_percentiles(); + std::cout << "[indeg] p10=" << ip.p10 << " p50=" << ip.p50 + << " p90=" << ip.p90 << "\n"; + } if (mixed_pool) { mixed_pool->shutdown(); } @@ -1178,7 +1358,10 @@ int main(int argc, char **argv) { << " lock=" << per(stg.rc_lock_us) << " impl=" << per(stg.rc_impl_us) << ") reconnects/insert=" << static_cast(stg.reconnects) / static_cast(stg.inserts) - << "\n"; + << " repair_ms=" << static_cast(stg.repair_us) / 1000.0 + << " repairs=" << stg.repairs + << " garden_ms=" << static_cast(stg.garden_us) / 1000.0 + << " gardens=" << stg.gardens << "\n"; } const size_t round_updates = round.deletes.size() + round.inserts.size(); const double update_qps = diff --git a/tests/diskann/test_diskann_update_e2e.cpp b/tests/diskann/test_diskann_update_e2e.cpp index 0fa6ff08..70b57473 100644 --- a/tests/diskann/test_diskann_update_e2e.cpp +++ b/tests/diskann/test_diskann_update_e2e.cpp @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -29,6 +30,17 @@ #include #endif +// TSan serializes the repair wave's atomic traffic; the full-size repair +// contract tests run for hours under it. Keep every code path but shrink the +// workload so the sanitizer suite stays runnable. +#if defined(__SANITIZE_THREAD__) + #define ALAYA_UPDATE_E2E_TSAN 1 +#elif defined(__has_feature) + #if __has_feature(thread_sanitizer) + #define ALAYA_UPDATE_E2E_TSAN 1 + #endif +#endif + namespace { #if defined(__linux__) @@ -192,6 +204,178 @@ class UpdateE2ETest : public ::testing::Test { return std::find(out_l.begin(), out_l.end(), label) != out_l.end(); } + std::vector take_random_deletable(uint32_t count, std::mt19937 &rng) { + std::vector ids; + ids.reserve(count); + for (uint32_t i = 0; i < count; ++i) { + std::uniform_int_distribution pick(0, deletable_.size() - 1); + const size_t j = pick(rng); + ids.push_back(deletable_[j]); + deletable_[j] = deletable_.back(); + deletable_.pop_back(); + } + return ids; + } + + std::unordered_set mirror_batch_remove(const std::vector &ids) { + std::unordered_set labels; + labels.reserve(ids.size()); + for (const uint32_t id : ids) { + const uint64_t label = id_label_.at(id); + labels.insert(label); + live_.erase(label); + } + idx_->batch_remove(ids.data(), static_cast(ids.size())); + return labels; + } + + std::vector base_queries_for_ids(const std::vector &ids, + uint32_t max_queries) const { + const uint32_t n = std::min(max_queries, static_cast(ids.size())); + std::vector queries; + queries.reserve(static_cast(n) * dim_); + for (uint32_t i = 0; i < n; ++i) { + const auto v = vec_at(base_vecs_, ids[i]); + queries.insert(queries.end(), v.begin(), v.end()); + } + return queries; + } + + void expect_searches_omit_labels(const std::vector &queries, + uint32_t k, + uint32_t l, + const std::unordered_set &forbidden) const { + ASSERT_NE(dim_, 0u); + ASSERT_EQ(queries.size() % static_cast(dim_), 0u); + std::vector out_l(k); + std::vector out_d(k); + const DiskANNSearchParams sp{/*L=*/l, + /*use_pq=*/false, + /*rerank=*/false, + /*rerank_count=*/0, + /*deterministic=*/true}; + const size_t nq = queries.size() / static_cast(dim_); + for (size_t qi = 0; qi < nq; ++qi) { + idx_->search(queries.data() + qi * dim_, k, out_l.data(), out_d.data(), sp); + for (const uint64_t label : out_l) { + if (label != DiskANNIndex::kNoLabel) { + EXPECT_EQ(forbidden.count(label), 0u) << "deleted label " << label + << " returned for query " << qi; + } + } + } + } + + void assert_indegree_matches_recount() { + const auto rc = idx_->debug_recount_in_degree(); + for (uint32_t i = 0; i < rc.size(); ++i) { + ASSERT_EQ(rc[i], idx_->in_degree_of(i)) << i; + } + } + + void exercise_repair_remove_contract(uint32_t repair_search_l, bool check_recall_floor) { +#if defined(ALAYA_UPDATE_E2E_TSAN) + constexpr uint32_t kN = 900; + constexpr uint32_t kDelete = 90; +#else + constexpr uint32_t kN = 3000; + constexpr uint32_t kDelete = 300; +#endif + constexpr uint32_t kDim = 32; + + DiskANNLoadParams lp; + lp.update_repair = true; + lp.repair_search_l = repair_search_l; + lp.safety_net_ops = 1000000; + build_and_load(kN, kDim, 32, lp); + + std::vector deleted(deletable_.begin(), deletable_.begin() + kDelete); + const uint64_t cap_before = idx_->max_slot_id(); + const auto deleted_labels = mirror_batch_remove(deleted); + + EXPECT_EQ(idx_->live_count(), static_cast(kN - kDelete)); + EXPECT_EQ(idx_->tombstone_count(), static_cast(kDelete)); + EXPECT_EQ(idx_->free_slot_count(), static_cast(kDelete)); + EXPECT_EQ(idx_->max_slot_id(), cap_before); + for (const uint32_t id : deleted) { + EXPECT_TRUE(idx_->is_deleted(id)); + } + + const auto deleted_queries = base_queries_for_ids(deleted, 32); + const auto random_queries = make_vectors(30, kDim, /*seed=*/1701 + repair_search_l); + expect_searches_omit_labels(deleted_queries, /*k=*/10, /*l=*/100, deleted_labels); + expect_searches_omit_labels(random_queries, /*k=*/10, /*l=*/100, deleted_labels); + if (check_recall_floor) { + EXPECT_GE(recall_at_k(random_queries, 30, 10, 100), 0.85); + } + + const auto nv = make_vectors(kDelete, kDim, /*seed=*/8801 + repair_search_l); + std::unordered_set reused; + reused.reserve(kDelete); + for (uint32_t i = 0; i < kDelete; ++i) { + reused.insert(do_insert(vec_at(nv, i))); + } + + EXPECT_EQ(idx_->live_count(), static_cast(kN)); + EXPECT_EQ(idx_->free_slot_count(), 0u); + EXPECT_EQ(idx_->max_slot_id(), cap_before); + for (const uint32_t id : deleted) { + EXPECT_TRUE(reused.count(id) != 0) << "freed slot " << id << " should be reused"; + EXPECT_FALSE(idx_->is_deleted(id)); + } + expect_searches_omit_labels(deleted_queries, /*k=*/10, /*l=*/100, deleted_labels); + expect_searches_omit_labels(random_queries, /*k=*/10, /*l=*/100, deleted_labels); + } + + double run_repair_churn(bool update_repair) { +#if defined(ALAYA_UPDATE_E2E_TSAN) + constexpr uint32_t kN = 1000; + constexpr uint32_t kRounds = 2; + constexpr uint32_t kChurn = 80; +#else + constexpr uint32_t kN = 4000; + constexpr uint32_t kRounds = 3; + constexpr uint32_t kChurn = 240; +#endif + constexpr uint32_t kDim = 32; + constexpr uint32_t kNq = 50; + constexpr uint32_t kK = 10; + constexpr uint32_t kL = 160; + + idx_.reset(); + std::error_code ec; + std::filesystem::remove_all(dir_, ec); + + DiskANNLoadParams lp; + lp.update_repair = update_repair; + lp.repair_search_l = 0; + lp.safety_net_ops = 1000000; + // The eval L exceeds the 150-slot scratch default; the neighbor pool is + // provisioned at load time and does not grow per query. + lp.scratch_search_list_size = kL; + build_and_load(kN, kDim, 48, lp); + + const auto queries = make_vectors(kNq, kDim, /*seed=*/7); + const auto extra = make_vectors(kRounds * kChurn, kDim, /*seed=*/555); + uint32_t extra_idx = 0; + std::mt19937 rng(2026); + + for (uint32_t round = 0; round < kRounds; ++round) { + const std::vector deleted = take_random_deletable(kChurn, rng); + mirror_batch_remove(deleted); + EXPECT_EQ(idx_->live_count(), static_cast(kN - kChurn)); + EXPECT_EQ(idx_->free_slot_count(), static_cast(kChurn)); + + for (uint32_t i = 0; i < kChurn; ++i) { + const uint32_t id = do_insert(vec_at(extra, extra_idx++)); + deletable_.push_back(id); + } + EXPECT_EQ(idx_->live_count(), static_cast(kN)); + EXPECT_EQ(idx_->free_slot_count(), 0u); + } + return recall_at_k(queries, kNq, kK, kL); + } + std::filesystem::path dir_; uint64_t dim_ = 0; uint32_t medoid_ = 0; @@ -777,6 +961,199 @@ TEST_F(UpdateE2ETest, BatchRemoveHidesVectorsAndFeedsFreeList) { } } +TEST_F(UpdateE2ETest, RepairPreservesRemoveContract) { + exercise_repair_remove_contract(/*repair_search_l=*/0, /*check_recall_floor=*/false); +} + +TEST_F(UpdateE2ETest, RepairChurnRecallFloor) { + const double repair_off = run_repair_churn(/*update_repair=*/false); + const double repair_on = run_repair_churn(/*update_repair=*/true); + + EXPECT_GE(repair_on, repair_off - 0.02) + << "repair_on=" << repair_on << " repair_off=" << repair_off; + EXPECT_GE(repair_on, 0.90) << "repair_on=" << repair_on; +} + +TEST_F(UpdateE2ETest, RepairWithSearchDiscovery) { + exercise_repair_remove_contract(/*repair_search_l=*/32, /*check_recall_floor=*/true); +} + +// Repair supersedes the lazy path's safety net. Under DEFAULT arming (5% +// tombstones, 16 deletes) a 10% batch delete satisfies both thresholds, so +// this pins the skip: without it every armed batch pays a redundant scan of +// the whole delete history that can never find an edge to fix. +TEST_F(UpdateE2ETest, RepairSkipsSafetyNetUnderDefaultArming) { +#if defined(ALAYA_UPDATE_E2E_TSAN) + constexpr uint32_t kN = 600; + constexpr uint32_t kDelete = 60; +#else + constexpr uint32_t kN = 1500; + constexpr uint32_t kDelete = 150; +#endif + constexpr uint32_t kDim = 32; + + DiskANNLoadParams lp; + lp.update_repair = true; // safety_net_ratio/ops stay at their defaults + build_and_load(kN, kDim, 32, lp); + + std::mt19937 rng(4242); + const std::vector deleted = take_random_deletable(kDelete, rng); + const auto deleted_labels = mirror_batch_remove(deleted); + EXPECT_EQ(idx_->safety_net_fire_count(), 0u); + EXPECT_EQ(idx_->free_slot_count(), static_cast(kDelete)); + + const std::vector deleted_again = take_random_deletable(kDelete, rng); + const auto deleted_labels_again = mirror_batch_remove(deleted_again); + EXPECT_EQ(idx_->safety_net_fire_count(), 0u); + + std::unordered_set forbidden(deleted_labels.begin(), deleted_labels.end()); + forbidden.insert(deleted_labels_again.begin(), deleted_labels_again.end()); + const auto deleted_queries = base_queries_for_ids(deleted, 32); + expect_searches_omit_labels(deleted_queries, /*k=*/10, /*l=*/100, forbidden); +} + +TEST_F(UpdateE2ETest, InDegreeCounterInvariant) { +#if defined(ALAYA_UPDATE_E2E_TSAN) + constexpr uint32_t kN = 400; + constexpr uint32_t kInsertFirst = 30; + constexpr uint32_t kDelete = 25; + constexpr uint32_t kInsertSecond = 20; + constexpr uint32_t kGardenBudget = 8; +#else + constexpr uint32_t kN = 800; + constexpr uint32_t kInsertFirst = 60; + constexpr uint32_t kDelete = 50; + constexpr uint32_t kInsertSecond = 40; + constexpr uint32_t kGardenBudget = 16; +#endif + constexpr uint32_t kDim = 32; + + DiskANNLoadParams lp; + lp.track_in_degree = true; + lp.update_repair = true; + lp.safety_net_ops = 1000000; + build_and_load(kN, kDim, 32, lp); + assert_indegree_matches_recount(); + + const auto extra = make_vectors(kInsertFirst + kInsertSecond, kDim, /*seed=*/9091); + for (uint32_t i = 0; i < kInsertFirst; ++i) { + do_insert(vec_at(extra, i)); + } + assert_indegree_matches_recount(); + + std::mt19937 rng(707); + const std::vector deleted = take_random_deletable(kDelete, rng); + mirror_batch_remove(deleted); + assert_indegree_matches_recount(); + + for (uint32_t i = 0; i < kInsertSecond; ++i) { + do_insert(vec_at(extra, kInsertFirst + i)); + } + assert_indegree_matches_recount(); + + const uint64_t live_before = idx_->live_count(); + const uint64_t tomb_before = idx_->tombstone_count(); + const uint64_t free_before = idx_->free_slot_count(); + const auto gs = idx_->garden_refresh(kGardenBudget); + EXPECT_LE(gs.refreshed, kGardenBudget); + EXPECT_GE(gs.selected, gs.refreshed); + EXPECT_EQ(idx_->live_count(), live_before); + EXPECT_EQ(idx_->tombstone_count(), tomb_before); + EXPECT_EQ(idx_->free_slot_count(), free_before); + assert_indegree_matches_recount(); +} + +TEST_F(UpdateE2ETest, GardenLiftsStarvedTail) { +#if defined(ALAYA_UPDATE_E2E_TSAN) + constexpr uint32_t kN = 600; + constexpr uint32_t kRounds = 2; + constexpr uint32_t kChurn = 50; +#else + constexpr uint32_t kN = 1500; + constexpr uint32_t kRounds = 3; + constexpr uint32_t kChurn = 120; +#endif + constexpr uint32_t kDim = 32; + constexpr uint32_t kNq = 40; + constexpr uint32_t kK = 10; + constexpr uint32_t kL = 100; + + DiskANNLoadParams lp; + lp.track_in_degree = true; + lp.update_repair = false; + lp.safety_net_ops = 1000000; + build_and_load(kN, kDim, 32, lp); + + const auto queries = make_vectors(kNq, kDim, /*seed=*/7); + const auto extra = make_vectors(kRounds * kChurn, kDim, /*seed=*/555); + uint32_t extra_idx = 0; + std::mt19937 rng(2026); + std::vector removed_ids; + removed_ids.reserve(kRounds * kChurn); + std::unordered_set removed_labels; + removed_labels.reserve(kRounds * kChurn); + + for (uint32_t round = 0; round < kRounds; ++round) { + const std::vector deleted = take_random_deletable(kChurn, rng); + const auto labels = mirror_batch_remove(deleted); + removed_ids.insert(removed_ids.end(), deleted.begin(), deleted.end()); + removed_labels.insert(labels.begin(), labels.end()); + EXPECT_EQ(idx_->live_count(), static_cast(kN - kChurn)); + EXPECT_EQ(idx_->free_slot_count(), static_cast(kChurn)); + + for (uint32_t i = 0; i < kChurn; ++i) { + const uint32_t id = do_insert(vec_at(extra, extra_idx++)); + deletable_.push_back(id); + } + EXPECT_EQ(idx_->live_count(), static_cast(kN)); + EXPECT_EQ(idx_->free_slot_count(), 0u); + } + + const auto before = idx_->in_degree_percentiles(); + const auto gs = idx_->garden_refresh(kN / 10); + const auto after = idx_->in_degree_percentiles(); + + EXPECT_GT(gs.refreshed, 0u); + EXPECT_GE(after.p10, before.p10); + EXPECT_GE(gs.indeg_p10_after, gs.indeg_p10_before); + EXPECT_GE(recall_at_k(queries, kNq, kK, kL), 0.85); + + const auto deleted_queries = base_queries_for_ids(removed_ids, 32); + expect_searches_omit_labels(deleted_queries, kK, kL, removed_labels); + expect_searches_omit_labels(queries, kK, kL, removed_labels); +} + +TEST_F(UpdateE2ETest, GardenContractWithoutTracking) { +#if defined(ALAYA_UPDATE_E2E_TSAN) + constexpr uint32_t kN = 200; +#else + constexpr uint32_t kN = 300; +#endif + constexpr uint32_t kDim = 32; + + DiskANNLoadParams lp; + build_and_load(kN, kDim, 32, lp); + EXPECT_THROW(idx_->garden_refresh(4), std::runtime_error); + EXPECT_EQ(idx_->in_degree_of(0), 0u); + const auto ip = idx_->in_degree_percentiles(); + EXPECT_EQ(ip.p10, 0u); + EXPECT_EQ(ip.p50, 0u); + EXPECT_EQ(ip.p90, 0u); + + idx_.reset(); + std::error_code ec; + std::filesystem::remove_all(dir_, ec); + + DiskANNLoadParams tracked_lp; + tracked_lp.track_in_degree = true; + build_and_load(kN, kDim, 32, tracked_lp); + EXPECT_NO_THROW({ + const auto gs = idx_->garden_refresh(0); + EXPECT_EQ(gs.refreshed, 0u); + EXPECT_EQ(gs.selected, 0u); + }); +} + // 7.5 ----------------------------------------------------------------------- TEST_F(UpdateE2ETest, PersistenceAcrossFlushReload) { build_and_load(300, 32, 32, {});