diff --git a/include/svs/concurrent/consolidate.h b/include/svs/concurrent/consolidate.h index 1ba9639c1..2779bee6c 100644 --- a/include/svs/concurrent/consolidate.h +++ b/include/svs/concurrent/consolidate.h @@ -23,6 +23,7 @@ #include "svs/index/vamana/consolidate.h" #include "svs/index/vamana/extensions.h" #include "svs/lib/array.h" +#include "svs/lib/misc.h" #include "svs/lib/threads.h" #include "svs/lib/timing.h" @@ -234,9 +235,15 @@ class GraphConsolidator { /// The in-neighbors of a deleted `d` are contained in `out(d) union R(d)`, so we hand /// the driver every non-deleted node from both. /// - template - std::vector - gather_work_set(const DeletedSet& deleted_ids, const Deleted& is_deleted) const { + template < + typename DeletedSet, + typename Deleted, + typename Pending = lib::Returns>> + std::vector gather_work_set( + const DeletedSet& deleted_ids, + const Deleted& is_deleted, + const Pending& is_pending = {} + ) const { auto* reverse_edges = graph_.reverse_edges(); assert(reverse_edges != nullptr); @@ -245,14 +252,29 @@ class GraphConsolidator { if (!is_deleted(d)) { continue; } - const auto& neighbors = graph_.get_node(lib::narrow_cast(d)); - for (auto a : neighbors) { - if (!is_deleted(a)) { - work.insert(a); + auto target = lib::narrow_cast(d); + // Adding a symmetric edge can move its source from R(d) to out(d). + // Validate both collections together so that move cannot hide an in-edge. + for (;;) { + auto seq = graph_.seq_counters()[target].read_begin(); + if (!seq) { + svs::detail::pause(); + continue; + } + for (auto a : graph_.get_node(target)) { + if (!is_deleted(a)) { + work.insert(a); + } } + reverse_edges->collect(target, work, is_deleted); + if (graph_.seq_counters()[target].read_validate(*seq)) { + break; + } + svs::detail::pause(); + } + if (!is_pending(d)) { + reverse_edges->reset_node(target); } - reverse_edges->collect(lib::narrow_cast(d), work, is_deleted); - reverse_edges->reset_node(lib::narrow_cast(d)); } return std::vector(work.begin(), work.end()); } @@ -297,11 +319,39 @@ class GraphConsolidator { /// Reverse-edge-driven consolidation: process only the in-neighbors of the deleted /// nodes `deleted_ids`, discovered via the graph's reverse-edge index. /// - template - void operator()(const DeletedSet& deleted_ids, const Deleted& is_deleted) { - auto work_ids = gather_work_set(deleted_ids, is_deleted); + template < + typename DeletedSet, + typename Deleted, + typename Pending = lib::Returns>> + void operator()( + const DeletedSet& deleted_ids, + const Deleted& is_deleted, + const Pending& is_pending = {} + ) { + auto work_ids = gather_work_set(deleted_ids, is_deleted, is_pending); assert(int(work_ids.size()) > -1); run_driver(work_ids, is_deleted); + + auto* reverse_edges = graph_.reverse_edges(); + tsl::robin_set incoming; + for (auto d : deleted_ids) { + if (!is_pending(d)) { + continue; + } + const auto target = lib::narrow_cast(d); + incoming.clear(); + // Release the reverse-list lock before taking source locks, preserving + // the graph writers' node -> reverse-list lock order. + reverse_edges->collect(target, incoming, lib::Returns(lib::Const())); + for (auto a : incoming) { + const auto source = lib::narrow_cast(a); + // Do not remove the record of an edge added between check and removal. + auto node_lock = graph_.lock_node(source); + if (!graph_.has_edge(source, target)) { + reverse_edges->remove(source, target); + } + } + } } /// @@ -366,6 +416,9 @@ void consolidate( /// /// Reverse-edge-driven consolidation over an explicit deleted set `deleted_ids`. Requires /// `graph.reverse_edges() != nullptr`; only the in-neighbors of `deleted_ids` are visited. +/// @param is_pending Identifies update/insertion reservations whose reverse lists must +/// retain concurrent additions. Defaults to false for every slot. Pending lists +/// are cleaned after repair by removing records whose source edge no longer exists. /// template < graphs::MemoryGraph Graph, @@ -373,7 +426,8 @@ template < threads::ThreadPool Pool, typename Distance, typename DeletedSet, - typename Deleted> + typename Deleted, + typename Pending = lib::Returns>> void consolidate( Graph& graph, const Data& data, @@ -383,11 +437,12 @@ void consolidate( float alpha, const Distance& distance, const DeletedSet& deleted_ids, - Deleted&& is_deleted + Deleted&& is_deleted, + const Pending& is_pending = {} ) { ConsolidationParameters params{200'000, prune_to, max_candidate_pool_size, alpha}; auto consolidator = GraphConsolidator{graph, data, threadpool, distance, params}; - consolidator(deleted_ids, is_deleted); + consolidator(deleted_ids, is_deleted, is_pending); } } // namespace svs::index::vamana::concurrent diff --git a/include/svs/concurrent/dynamic_index.h b/include/svs/concurrent/dynamic_index.h index 1fc1a8980..dacf5b2b5 100644 --- a/include/svs/concurrent/dynamic_index.h +++ b/include/svs/concurrent/dynamic_index.h @@ -17,14 +17,17 @@ #pragma once // stdlib +#include #include #include #include #include #include +#include #include #include #include +#include #include // Include the flat index to spin-up exhaustive searches on demand. @@ -106,7 +109,7 @@ enum class SlotMetadata : uint8_t { Empty = 0x00, Valid = 0x01, Deleted = 0x02, - // Reserved by an in-flight add_points + // Reserved by an in-flight add_points or update_entries Pending = 0x04, }; @@ -230,8 +233,9 @@ class MutableVamanaIndex { std::unique_ptr> first_reusable_{ std::make_unique>(0)}; IDTranslator translator_; - // Count of Valid slots. Maintained atomically in add_points/delete_entry. - // Wrapped in unique_ptr because std::atomic is not movable. + // Count of Valid slots. Maintained atomically in + // add_points/delete_entry/update_entries. Wrapped in unique_ptr because std::atomic is + // not movable. std::unique_ptr> num_valid_{ std::make_unique>(0)}; // Protects translator access: exclusive for writes (add/consolidate/compact), @@ -253,16 +257,25 @@ class MutableVamanaIndex { // compact_gatekeeper_ -> compact_mutex_ (never reversed) // compact_mutex_ -> slot_alloc_mutex_ (never reversed) // compact_mutex_ -> translator_mutex_ (never reversed) - // slot_alloc_mutex_ and translator_mutex_ are never held simultaneously - // (add_points takes them sequentially), so they have no relative order. + // update_entries reserve Empty slots under translator_mutex_ -> slot_alloc_mutex_. + // add_points takes these locks separately; it never holds them in reverse order. + // Payload readers take compact_mutex_ -> payload_mutex_ -> translator_mutex_. + // update_entries releases translator_mutex_ before acquiring payload protection. std::unique_ptr compact_mutex_{ std::make_unique()}; // Once compaction owns this gate, new shared users cannot delay its // storage lock. Readers hold the gate only while acquiring compact_mutex_. std::unique_ptr compact_gatekeeper_{ std::make_unique()}; - // Iterators check this under compact_mutex_ before reusing cached internal IDs. + // Iterators check this under compact_mutex_ and payload_mutex_ + // before reusing cached internal IDs size_t compaction_epoch_ = 0; + // Once writer owns this gate, new shared users cannot delay its + // lock. Readers hold the gate only while acquiring payload_mutex_. + std::unique_ptr payload_mutex_{ + std::make_unique()}; + std::unique_ptr payload_gatekeeper_{ + std::make_unique()}; // Writer-only mutex serializing slot allocation in add_points std::unique_ptr slot_alloc_mutex_{std::make_unique()}; @@ -846,6 +859,7 @@ class MutableVamanaIndex { auto get_datum(size_t e) const { // Lock order: compact_mutex_ then translator_mutex_ (global order). auto compact_lock = lock_storage_shared(); + auto payload_lock = lock_payload_shared(); std::shared_lock lock{*translator_mutex_}; if (!translator_.has_external(e)) { throw ANNEXCEPTION("External ID {} not found in index!", e); @@ -862,6 +876,7 @@ class MutableVamanaIndex { requires std::invocable bool on_datum(size_t e, F&& f) const { auto compact_lock = lock_storage_shared(); + auto payload_lock = lock_payload_shared(); std::shared_lock lock{*translator_mutex_}; if (!unsafe_has_id(e)) { return false; @@ -893,6 +908,13 @@ class MutableVamanaIndex { return std::shared_lock(*compact_mutex_); } + /// @brief Hold shared payload protection for a complete read operation. + /// Acquire compact_mutex_ first, then this guard, then translator/node locks. + [[nodiscard]] std::shared_lock lock_payload_shared() const { + std::shared_lock gatekeeper_lock{*payload_gatekeeper_}; + return std::shared_lock(*payload_mutex_); + } + /// @brief RAII reader lock guarding translator_ against erase/remap by /// consolidate/compact. Used by BatchIterator::next() to protect /// internal->external ID translation. @@ -942,6 +964,7 @@ class MutableVamanaIndex { // Hold compact_mutex_ shared so compact()'s shrink can't free segments // mid-traversal. add_points growth is lock-free (grow-stable storage). auto compact_lock = lock_storage_shared(); + auto payload_lock = lock_payload_shared(); extensions::single_search( data_, scratch.buffer, @@ -961,6 +984,7 @@ class MutableVamanaIndex { ) { // Internal IDs must retain their meaning until translation has completed. auto compact_lock = lock_storage_shared(); + auto payload_lock = lock_payload_shared(); { threads::parallel_for( threadpool_, @@ -1022,11 +1046,17 @@ class MutableVamanaIndex { size_t num_neighbors, QueryResultView result ) { + auto compact_lock = lock_storage_shared(); + auto payload_lock = lock_payload_shared(); + const auto status_size = status_.size(); auto temp_index = temporary_flat_index( data_, distance_, threads::ThreadPoolReferenceWrapper(threadpool_) ); temp_index.search(queries, num_neighbors, result, [&](size_t i) { - return getindex(status_, i) == SlotMetadata::Valid; + // Insertion can grow data_ before it grows status_. + return i < status_size && + std::atomic_ref(status_[i]) + .load(std::memory_order_acquire) == SlotMetadata::Valid; }); // After the search procedure, the indices in `results` are internal. @@ -1205,9 +1235,27 @@ class MutableVamanaIndex { copy_points(points, slots); clear_lists(slots); + { + auto payload_lock = lock_payload_shared(); + connect_pending(slots); + } + + // Mark added entries as valid. + for (const auto& i : slots) { + std::atomic_ref(status_[i]) + .store(SlotMetadata::Valid, std::memory_order_release); + } + num_valid_->fetch_add(slots.size(), std::memory_order_acq_rel); + + return slots; + } + + private: + // Callers own initialized Pending slots and hold storage and payload protection. + void connect_pending(const std::vector& slots) { // Publish only initialized slots: another insertion may immediately use one // as a candidate or backlink parent. Batch insertions do not use this registry. - const bool single_insert = num_points == 1; + const bool single_insert = slots.size() == 1; const Idx pending_id = lib::narrow_cast(slots.front()); if (single_insert) { std::unique_lock lock{*pending_insertions_mutex_}; @@ -1281,15 +1329,207 @@ class MutableVamanaIndex { lib::as_const_span(pending_candidates)}; builder.construct(alpha_, seed, build_slots, logging::Level::Trace, logger_); } + } - // Mark added entries as valid. - for (const auto& i : slots) { - std::atomic_ref(status_[i]) + struct UpdateBatch { + std::vector slots; + std::vector> ordered; + tsl::robin_set targets; + size_t live = 0; + size_t first_empty = 0; + }; + + template + UpdateBatch prepare_update(const Points& points, std::vector slots) { + if (points.size() != slots.size()) { + throw ANNEXCEPTION( + "Number of points ({}) differs from the number of IDs ({})!", + points.size(), + slots.size() + ); + } + if (!slots.empty() && points.dimensions() != dimensions()) { + throw ANNEXCEPTION( + "Points have {} dimensions but the index has {}!", + points.dimensions(), + dimensions() + ); + } + UpdateBatch batch{}; + batch.slots = std::move(slots); + batch.ordered.reserve(batch.slots.size()); + batch.targets.reserve(batch.slots.size()); + const auto status_size = status_.size(); + for (auto id : batch.slots) { + if (id >= status_size) { + throw ANNEXCEPTION("Internal ID {} is outside [0, {})!", id, status_size); + } + auto internal = lib::narrow(id); + if (!batch.targets.insert(internal).second) { + throw ANNEXCEPTION("Internal ID {} is listed more than once!", id); + } + batch.ordered.emplace_back(internal, SlotMetadata::Pending); + } + // A common reservation order prevents overlapping batches waiting cyclically. + // slots remains in point order for copying the replacement vectors. + std::sort(batch.ordered.begin(), batch.ordered.end(), [](auto a, auto b) { + return a.first < b.first; + }); + batch.first_empty = first_empty_->load(std::memory_order_acquire); + return batch; + } + + void wait_until_not_pending(size_t id) const { + while (std::atomic_ref(const_cast(status_[id])) + .load(std::memory_order_acquire) == SlotMetadata::Pending) { + svs::detail::pause(); + } + } + + // The owner of a Pending slot may need translation protection to finish. + template + void wait_for_owner(TranslationLock& translation_lock, Idx id) const { + translation_lock.unlock(); + wait_until_not_pending(id); + translation_lock.lock(); + } + + // Reserve every target, or none. The caller holds storage and translation locks + // from its checks until this returns. Returns a Pending slot of another operation. + std::optional try_reserve_update(UpdateBatch& batch) { + size_t reserved = 0; + for (auto& [id, previous] : batch.ordered) { + auto ref = std::atomic_ref(status_[id]); + for (;;) { + std::unique_lock allocation_lock{*slot_alloc_mutex_, std::defer_lock}; + if (ref.load(std::memory_order_acquire) == SlotMetadata::Empty) { + allocation_lock.lock(); + } + auto node_lock = graph_.lock_node(id); + auto expected = ref.load(std::memory_order_acquire); + if (expected == SlotMetadata::Pending) { + // The caller waits and checks again; it may then pick other slots. + release_reservations(batch, reserved); + batch.live = 0; + return id; + } + if ((expected != SlotMetadata::Empty || allocation_lock.owns_lock()) && + ref.compare_exchange_strong( + expected, + SlotMetadata::Pending, + std::memory_order_acq_rel, + std::memory_order_relaxed + )) { + previous = expected; + batch.live += expected == SlotMetadata::Valid; + ++reserved; + break; + } + } + } + num_valid_->fetch_sub(batch.live, std::memory_order_acq_rel); + return std::nullopt; + } + + // Restore the first `count` reserved slots to their saved states. + void release_reservations(const UpdateBatch& batch, size_t count) { + size_t first_restored_empty = std::numeric_limits::max(); + for (size_t i = 0; i < count; ++i) { + auto [id, previous] = batch.ordered[i]; + std::atomic_ref(status_[id]) + .store(previous, std::memory_order_release); + if (previous == SlotMetadata::Empty) { + first_restored_empty = std::min(first_restored_empty, size_t(id)); + } + } + if (first_restored_empty != std::numeric_limits::max()) { + detail::atomic_min(*first_empty_, batch.first_empty); + detail::atomic_min(*first_reusable_, first_restored_empty); + } + } + + // The caller owns these Pending slots and holds storage protection. + // This overload performs no ID translation or mapping mutation. + template + void update_internal(const Points& points, const UpdateBatch& batch) { + { + std::lock_guard gatekeeper_lock{*payload_gatekeeper_}; + std::lock_guard payload_lock{*payload_mutex_}; + // A partial payload write also invalidates cached iterator distances. + ++compaction_epoch_; + copy_points(points, batch.slots); + } + auto payload_lock = lock_payload_shared(); + for (auto [id, previous] : batch.ordered) { + if (previous == SlotMetadata::Empty) { + // An Empty slot can come from an insertion that failed before + // initializing its adjacency. Use insertion's initialization. + graph_.clear_node(id); + } + } + // The exclusive copy drained linkers using the old payload. Their connections + // are now visible to repair; keep the old outgoing lists as repair candidates. + repair_connections(batch.targets); + const auto anchor = entry_point(); + for (auto id : batch.slots) { + if (id != anchor) { + graph_.replace_node(lib::narrow_cast(id), std::span{}); + } + } + // The anchor retains its outgoing list until construction searches from it. + connect_pending(batch.slots); + // A waiting deletion can consume Valid immediately after publication. + num_valid_->fetch_add(batch.slots.size(), std::memory_order_acq_rel); + for (auto id : batch.slots) { + std::atomic_ref(status_[id]) .store(SlotMetadata::Valid, std::memory_order_release); } - num_valid_->fetch_add(slots.size(), std::memory_order_acq_rel); + } - return slots; + public: + /// @brief Update vectors addressed by existing external mappings. + /// @param points Replacement dataset in external-ID order, with the index dimensions. + /// @param external_ids Unique existing external IDs; supports iteration and `size()`. + /// + /// The mappings are retained. Accepts mapped Valid, Deleted, Empty and Pending + /// slots; waits for Pending owners. Targets are Pending during the operation and + /// Valid on success. Only formerly Valid targets are subtracted from size() + /// during the update. + /// Live checks exclude Pending targets; an older search may finish with old results. + /// Concurrent deletion waits for the Pending owner. + /// After a wait, the IDs are looked up and checked again. + /// Encoding/build failures have the same rollback guarantees as add_points: + /// partial work and Pending slots can remain. See `get_threadpool_handle()` for + /// the worker-task calling contract. Do not call while holding index read guards. + template + void update_entries(const Points& points, const ExternalIds& external_ids) { + auto compact_lock = lock_storage_shared(); + auto translation_lock = lock_for_translation(); + + UpdateBatch batch; + for (;;) { + // Keep the lock from these checks until the slots are reserved, + // so every check is still true when the update is applied. + std::vector slots; + slots.reserve(external_ids.size()); + for (auto id : external_ids) { + if (!translator_.has_external(id)) { + throw ANNEXCEPTION("External ID {} has no mapping!", id); + } + slots.push_back(translator_.get_internal(id)); + } + batch = prepare_update(points, std::move(slots)); + if (batch.slots.empty()) { + return; + } + auto busy = try_reserve_update(batch); + if (!busy) { + break; + } + wait_for_owner(translation_lock, *busy); + } + translation_lock.unlock(); + update_internal(points, batch); } /// @@ -1319,15 +1559,22 @@ class MutableVamanaIndex { std::shared_lock lock{*translator_mutex_}; size_t deleted = 0; for (auto i : ids) { - if (!translator_.has_external(i)) { - continue; // Already deleted + consolidated, or never existed. - } - auto internal = translator_.get_internal(i); - if (is_deleted(internal)) { - continue; // Already soft-deleted, translator entry not yet consolidated. + for (;;) { + if (!translator_.has_external(i)) { + break; + } + auto internal = translator_.get_internal(i); + auto previous = try_delete_entry(internal); + if (previous != SlotMetadata::Pending) { + deleted += previous == SlotMetadata::Valid; + break; + } + // The Pending owner may need translation protection to complete. + lock.unlock(); + wait_until_not_pending(internal); + lock.lock(); + // Resolve the external ID again: its mapping may have changed. } - delete_entry(internal); - ++deleted; } // Don't erase translator entries here — concurrent search may still // need them for translate_to_external(). Cleanup happens in @@ -1336,29 +1583,32 @@ class MutableVamanaIndex { } void delete_entry(size_t i) { - auto& meta = getindex(status_, i); - auto ref = std::atomic_ref(meta); - // Wait without the node lock: the Pending insertion needs it to finish. - while (ref.load(std::memory_order_acquire) == SlotMetadata::Pending) { - svs::detail::pause(); + while (try_delete_entry(i) == SlotMetadata::Pending) { + wait_until_not_pending(i); } - // For non-anchor parents, a backlink commits before this transition or - // observes Deleted under the same lock and retries its Pending insertion. + } + + private: + SlotMetadata try_delete_entry(size_t i) { + // Serialize with update reservation and with backlink publication. auto node_lock = graph_.lock_node(lib::narrow_cast(i)); SlotMetadata expected = SlotMetadata::Valid; - if (ref.compare_exchange_strong( - expected, - SlotMetadata::Deleted, - std::memory_order_acq_rel, - std::memory_order_relaxed - )) { + if (std::atomic_ref(status_[i]) + .compare_exchange_strong( + expected, + SlotMetadata::Deleted, + std::memory_order_acq_rel, + std::memory_order_relaxed + )) { num_valid_->fetch_sub(1, std::memory_order_acq_rel); } + return expected; } + public: bool is_deleted(size_t i) const { // True only for slots that have been soft-deleted. Pending (in-flight - // add) and Empty are NOT deleted: consolidate must not prune them out + // insertion/update) and Empty are NOT deleted: consolidate must not prune them out // of other nodes' adjacency lists, and search already filters // non-Valid slots via ValidBuilder. return std::atomic_ref(const_cast(status_[i])) @@ -1582,6 +1832,11 @@ class MutableVamanaIndex { /// /// @brief Return the current thread pool handle. /// + /// Do not call index methods that acquire storage/payload protection from tasks on + /// this index's underlying pool, including through another handle to that pool. + /// The submitting thread must acquire those guards before submitting work; tasks + /// use `unsafe_*` read paths under those guards. Otherwise a writer can close the + /// gate and wait for the pool while its tasks wait at that gate. threads::ThreadPoolHandle& get_threadpool_handle() { return threadpool_; } ///// Window Interface @@ -1617,24 +1872,35 @@ class MutableVamanaIndex { consolidate_locked(); } - // Full consolidation: gather the Deleted set (one cheap byte scan of status_) then - // delegate. Concurrent deletes arriving after the scan are picked up next round. + // Hold payload protection before collecting IDs: an updater must drain this + // snapshot's repair before it can replace and rebuild one of its Deleted slots. void consolidate_locked() { + auto payload_lock = lock_payload_shared(); tsl::robin_set deleted{}; for (size_t i = 0, imax = status_.size(); i < imax; ++i) { - if (std::atomic_ref(status_[i]).load(std::memory_order_acquire) == - SlotMetadata::Deleted) { + if (is_deleted(i)) { deleted.insert(lib::narrow_cast(i)); } } - consolidate_locked(deleted); + consolidate_selected(deleted); } - // Body of consolidate()/consolidate(ids) with no compact_mutex_ locking. The caller - // holds compact_mutex_. `deleted` holds the internal slots to prune out of the graph - // and reclaim; those slots must currently be SlotMetadata::Deleted. Consolidation and - // cleanup both range over exactly `deleted`, so both are O(|deleted|), not O(N). + // The caller holds storage protection. Revalidate explicit targets under payload + // protection: a supplied deletion snapshot may predate a completed update. void consolidate_locked(const tsl::robin_set& deleted) { + auto payload_lock = lock_payload_shared(); + tsl::robin_set current; + for (auto id : deleted) { + if (is_deleted(id)) { + current.insert(id); + } + } + consolidate_selected(current); + } + + private: + // Graph-only repair; the caller holds storage and payload protection. + void repair_connections(const tsl::robin_set& deleted) { // Online operations never hand off or reclaim the routing anchor. // If this snapshot is NO_ENTRY, an overlapping election can only // publish a new Pending slot, which cannot belong to `deleted`. @@ -1643,6 +1909,12 @@ class MutableVamanaIndex { return i != anchor && deleted.contains(lib::narrow_cast(i)); }; + // update_entries mark slots Pending before consolidation + auto is_pending = [&](size_t i) { + return std::atomic_ref(status_[i]) + .load(std::memory_order_acquire) == SlotMetadata::Pending; + }; + // Perform graph consolidation over the in-neighbors of `deleted`, discovered via // the reverse-edge index. concurrent::consolidate( @@ -1654,9 +1926,16 @@ class MutableVamanaIndex { alpha_, distance_, deleted, - should_remove + should_remove, + is_pending ); + } + // Reclaim only ordinary deletions, never an update's reserved Pending slots. + // The caller holds storage and payload protection through collection and cleanup. + void consolidate_selected(const tsl::robin_set& deleted) { + repair_connections(deleted); + const auto anchor = entry_point(); // After consolidation - clean up the removed slots under lock. O(|deleted|). { std::lock_guard lock{*translator_mutex_}; @@ -1665,7 +1944,7 @@ class MutableVamanaIndex { // may have reassigned the external ID and erased the stale reverse entry. std::vector deleted_internal_ids; for (auto i : deleted) { - if (status_[i] == SlotMetadata::Deleted && translator_.has_internal(i)) { + if (is_deleted(i) && translator_.has_internal(i)) { deleted_internal_ids.push_back(i); } } @@ -1676,7 +1955,7 @@ class MutableVamanaIndex { // but its data/edges remain routing-only until exclusive compaction. size_t min_freed = std::numeric_limits::max(); for (auto i : deleted) { - if (i != anchor && status_[i] == SlotMetadata::Deleted) { + if (i != anchor && is_deleted(i)) { std::atomic_ref(status_[i]) .store(SlotMetadata::Empty, std::memory_order_release); min_freed = std::min(min_freed, static_cast(i)); @@ -1688,6 +1967,7 @@ class MutableVamanaIndex { } } + public: /// /// @brief Consolidate only the soft-deleted entries listed in `ids`. /// @@ -1703,6 +1983,7 @@ class MutableVamanaIndex { /// template size_t consolidate(const T& ids) { auto compact_lock = lock_storage_shared(); + auto payload_lock = lock_payload_shared(); // Collect the internal slots of the listed, already-soft-deleted IDs. tsl::robin_set targets{}; @@ -1723,7 +2004,7 @@ class MutableVamanaIndex { if (targets.empty()) { return 0; } - consolidate_locked(targets); + consolidate_selected(targets); return targets.size(); } @@ -1821,6 +2102,7 @@ class MutableVamanaIndex { // compact_mutex_ shared guards data_/graph_ against compact()'s shrink; // translator_mutex_ shared guards the ID translation reads below. auto compact_lock = lock_storage_shared(); + auto payload_lock = lock_payload_shared(); std::shared_lock lock{*translator_mutex_}; // Bounds checking. @@ -1852,6 +2134,8 @@ class MutableVamanaIndex { /// This function is meant to provide a means for implementing experimental algorithms /// on the contained data structures. The entry-point span is a snapshot valid /// only during the callback; it is empty when no routing anchor exists. + /// The caller holds storage and payload protection through `f` and any borrowed + /// views. See `lock_payload_shared()` and the worker-task contract there. template void experimental_escape_hatch(F&& f) const { const Idx seed = entry_point(); std::invoke( @@ -1984,6 +2268,7 @@ class MutableVamanaIndex { // compact_mutex_ shared guards data_ against compact()'s shrink; // translator_mutex_ shared guards the ID translation read. auto compact_lock = lock_storage_shared(); + auto payload_lock = lock_payload_shared(); std::shared_lock lock{*translator_mutex_}; // Check if the external ID exists @@ -2017,7 +2302,8 @@ class MutableVamanaIndex { /// the greedy search may return fewer than ``target_window()`` valid neighbors. This /// supplements the buffer by scanning the ID translation. /// - /// **Precondition:** the caller holds ``compact_mutex_`` shared (``search()`` does). + /// **Precondition:** the caller holds storage and payload protection (`search()` + /// does). These guards are also required by `experimental_escape_hatch` readers. /// This is why the generic ``extensions::check_and_supplement_search_buffer`` is not /// used here: it reaches back through ``index.get_distance()``, which takes /// ``compact_mutex_`` shared a *second* time on the same thread. Recursive shared diff --git a/include/svs/concurrent/iterator.h b/include/svs/concurrent/iterator.h index a9af225fa..1023bb86a 100644 --- a/include/svs/concurrent/iterator.h +++ b/include/svs/concurrent/iterator.h @@ -61,6 +61,27 @@ class BatchIterator ) : base_type{parent, query, extra_search_buffer_capacity} {} + /// @brief Retrieve the next batch while holding storage and payload protection. + /// @param batch_size Number of neighbors to request. + /// @param cancel Cancellation predicate; defaults to never cancelling. + void next( + size_t batch_size, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) { + auto storage_lock = parent_->lock_storage_shared(); + auto payload_lock = parent_->lock_payload_shared(); + unsafe_next(batch_size, cancel); + } + + /// @copydoc next + /// The caller must hold the parent's storage and payload guards through this call. + void unsafe_next( + size_t batch_size, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) { + base_type::next(batch_size, cancel); + } + /// @brief Adapts an internal neighbor to an external neighbor. template svs::Neighbor adapt(N internal) const { if constexpr (Index::needs_id_translation) { @@ -96,8 +117,7 @@ class BatchIterator private: auto begin_search() { - // Keep internal IDs stable through traversal, reranking, and translation. - [[maybe_unused]] auto search_guard = parent_->lock_storage_shared(); + // next() or unsafe_next()'s caller owns both guards through translation. const auto epoch = parent_->compaction_epoch_; if (compaction_epoch_ != epoch) { // Reranking touches cached IDs before the restart initializer runs. @@ -106,7 +126,7 @@ class BatchIterator restart_search_ = true; compaction_epoch_ = epoch; } - return search_guard; + return lib::Empty{}; } auto lock_for_translation() const { diff --git a/include/svs/concurrent/multi.h b/include/svs/concurrent/multi.h index ec59d79c4..bb617bdf3 100644 --- a/include/svs/concurrent/multi.h +++ b/include/svs/concurrent/multi.h @@ -60,9 +60,23 @@ template class MultiBatchIterator { : index_{index} , batch_iterator_{index.get_parent_index(), query, extra_search_buffer_capacity} {} + /// @brief Retrieve distinct labels under the parent's storage and payload guards. + /// @param batch_size Number of labels to request. + /// @param cancel Cancellation predicate; defaults to never cancelling. void next( size_t batch_size, const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) + ) { + auto storage_lock = index_.get_parent_index().lock_storage_shared(); + auto payload_lock = index_.get_parent_index().lock_payload_shared(); + unsafe_next(batch_size, cancel); + } + + /// @copydoc next + /// The caller holds the parent's storage and payload guards through this call. + void unsafe_next( + size_t batch_size, + const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) ) { auto results_copy = results_; results_.clear(); @@ -70,7 +84,7 @@ template class MultiBatchIterator { while (results_.size() < batch_size && !batch_iterator_.done()) { try { - batch_iterator_.next(batch_size, cancel); + batch_iterator_.unsafe_next(batch_size, cancel); } catch (const ANNException& error) { results_ = std::move(results_copy); throw error; @@ -536,6 +550,8 @@ class MultiMutableVamanaIndex { const search_parameters_type& sp, const lib::DefaultPredicate& cancel = lib::Returns(lib::Const()) ) { + auto storage_lock = index_->lock_storage_shared(); + auto payload_lock = index_->lock_payload_shared(); auto& borrow_threadpool = index_->get_threadpool_handle(); threads::parallel_for( @@ -549,7 +565,7 @@ class MultiMutableVamanaIndex { // use batch iterator to search for (auto i : is) { auto batch_iterator = make_batch_iterator(queries.get_datum(i), 10); - batch_iterator.next(batch_size, cancel); + batch_iterator.unsafe_next(batch_size, cancel); size_t j{0}; for (auto& res : batch_iterator) { if (j == num_neighbors) { diff --git a/include/svs/index/flat/flat.h b/include/svs/index/flat/flat.h index 925fe68a5..0feb5d809 100644 --- a/include/svs/index/flat/flat.h +++ b/include/svs/index/flat/flat.h @@ -339,7 +339,7 @@ class FlatIndex { scratch.prepare(); size_t start = 0; - while (start < data_.size()) { + while (start < data_max_size) { // Check if request to cancel the search if (cancel()) { scratch.cleanup();