From e6da02b84eb84f9359645e55545e3a4a6b1645f9 Mon Sep 17 00:00:00 2001 From: GuyAv46 Date: Thu, 2 Jul 2026 14:53:28 +0300 Subject: [PATCH 1/2] Use one-byte HNSW node locks --- src/VecSim/algorithms/hnsw/graph_data.h | 2 - src/VecSim/algorithms/hnsw/hnsw.h | 85 +++++++++---------- .../algorithms/hnsw/hnsw_batch_iterator.h | 2 +- .../algorithms/hnsw/hnsw_serializer_impl.h | 4 +- src/VecSim/index_factories/hnsw_factory.cpp | 7 +- src/VecSim/utils/vecsim_stl.h | 38 +++++++++ tests/unit/test_allocator.cpp | 4 +- 7 files changed, 88 insertions(+), 54 deletions(-) diff --git a/src/VecSim/algorithms/hnsw/graph_data.h b/src/VecSim/algorithms/hnsw/graph_data.h index b38f74986..ece03ce20 100644 --- a/src/VecSim/algorithms/hnsw/graph_data.h +++ b/src/VecSim/algorithms/hnsw/graph_data.h @@ -3,7 +3,6 @@ #include #include -#include #include "VecSim/utils/vec_utils.h" // Amortized shrink thresholds for incoming edges vectors. @@ -99,7 +98,6 @@ struct ElementLevelData { struct ElementGraphData { size_t toplevel; - std::mutex neighborsGuard; ElementLevelData *others; ElementLevelData level0; diff --git a/src/VecSim/algorithms/hnsw/hnsw.h b/src/VecSim/algorithms/hnsw/hnsw.h index 1a3776fa2..2510ddbfb 100644 --- a/src/VecSim/algorithms/hnsw/hnsw.h +++ b/src/VecSim/algorithms/hnsw/hnsw.h @@ -116,6 +116,7 @@ class HNSWIndex : public VecSimIndexAbstract, // Index data vecsim_stl::vector graphDataBlocks; + mutable vecsim_stl::vector elementLocks; vecsim_stl::vector idToMetaData; // Used for marking the visited nodes in graph scans (the pool supports parallel graph scans). @@ -252,8 +253,6 @@ class HNSWIndex : public VecSimIndexAbstract, void unlockSharedIndexDataGuard() const; void lockNodeLinks(idType node_id) const; void unlockNodeLinks(idType node_id) const; - void lockNodeLinks(ElementGraphData *node_data) const; - void unlockNodeLinks(ElementGraphData *node_data) const; VisitedNodesHandler *getVisitedList() const; void returnVisitedList(VisitedNodesHandler *visited_nodes_handler) const; VecSimIndexDebugInfo debugInfo() const override; @@ -481,24 +480,14 @@ void HNSWIndex::unlockSharedIndexDataGuard() const { indexDataGuard.unlock_shared(); } -template -void HNSWIndex::lockNodeLinks(ElementGraphData *node_data) const { - node_data->neighborsGuard.lock(); -} - -template -void HNSWIndex::unlockNodeLinks(ElementGraphData *node_data) const { - node_data->neighborsGuard.unlock(); -} - template void HNSWIndex::lockNodeLinks(idType node_id) const { - lockNodeLinks(getGraphDataByInternalId(node_id)); + elementLocks[node_id].lock(); } template void HNSWIndex::unlockNodeLinks(idType node_id) const { - unlockNodeLinks(getGraphDataByInternalId(node_id)); + elementLocks[node_id].unlock(); } /** @@ -528,7 +517,7 @@ void HNSWIndex::processCandidate( candidatesMaxHeap &candidate_set, DistType &lowerBound) const { ElementGraphData *cur_element = getGraphDataByInternalId(curNodeId); - lockNodeLinks(cur_element); + lockNodeLinks(curNodeId); ElementLevelData &node_level = getElementLevelData(cur_element, layer); linkListSize num_links = node_level.getNumLinks(); if (num_links > 0) { @@ -602,7 +591,7 @@ void HNSWIndex::processCandidate( } } } - unlockNodeLinks(cur_element); + unlockNodeLinks(curNodeId); } template @@ -612,7 +601,7 @@ void HNSWIndex::processCandidate_RangeSearch( candidatesMaxHeap &candidate_set, DistType dyn_range, DistType radius) const { auto *cur_element = getGraphDataByInternalId(curNodeId); - lockNodeLinks(cur_element); + lockNodeLinks(curNodeId); ElementLevelData &node_level = getElementLevelData(cur_element, layer); linkListSize num_links = node_level.getNumLinks(); @@ -669,7 +658,7 @@ void HNSWIndex::processCandidate_RangeSearch( } } } - unlockNodeLinks(cur_element); + unlockNodeLinks(curNodeId); } template @@ -904,11 +893,11 @@ idType HNSWIndex::mutuallyConnectNewElement( idType selected_neighbor = neighbor_data.second; // neighbor's id auto *neighbor_graph_data = getGraphDataByInternalId(selected_neighbor); if (new_node_id < selected_neighbor) { - lockNodeLinks(new_node_level); - lockNodeLinks(neighbor_graph_data); + lockNodeLinks(new_node_id); + lockNodeLinks(selected_neighbor); } else { - lockNodeLinks(neighbor_graph_data); - lockNodeLinks(new_node_level); + lockNodeLinks(selected_neighbor); + lockNodeLinks(new_node_id); } // validations... @@ -921,15 +910,15 @@ idType HNSWIndex::mutuallyConnectNewElement( // The new node cannot add more neighbors this->log(VecSimCommonStrings::LOG_DEBUG_STRING, "Couldn't add all chosen neighbors upon inserting a new node"); - unlockNodeLinks(new_node_level); - unlockNodeLinks(neighbor_graph_data); + unlockNodeLinks(new_node_id); + unlockNodeLinks(selected_neighbor); break; } // If one of the two nodes has already deleted - skip the operation. if (isMarkedDeleted(new_node_id) || isMarkedDeleted(selected_neighbor)) { - unlockNodeLinks(new_node_level); - unlockNodeLinks(neighbor_graph_data); + unlockNodeLinks(new_node_id); + unlockNodeLinks(selected_neighbor); continue; } @@ -940,8 +929,8 @@ idType HNSWIndex::mutuallyConnectNewElement( if (neighbor_level_data.getNumLinks() < max_M_cur) { new_node_level_data.appendLink(selected_neighbor); neighbor_level_data.appendLink(new_node_id); - unlockNodeLinks(new_node_level); - unlockNodeLinks(neighbor_graph_data); + unlockNodeLinks(new_node_id); + unlockNodeLinks(selected_neighbor); continue; } @@ -1068,7 +1057,7 @@ void HNSWIndex::replaceEntryPoint() { volatile idType candidate_in_process = INVALID_ID; // Go over the entry point's neighbors at the top level. - lockNodeLinks(old_entry_point); + lockNodeLinks(old_entry_point_id); ElementLevelData &old_ep_level = getElementLevelData(old_entry_point, maxLevel); // Tries to set the (arbitrary) first neighbor as the entry point which is not deleted, // if exists. @@ -1076,7 +1065,7 @@ void HNSWIndex::replaceEntryPoint() { if (!isMarkedDeleted(old_ep_level.getLinkAtPos(i))) { if (!isInProcess(old_ep_level.getLinkAtPos(i))) { entrypointNode = old_ep_level.getLinkAtPos(i); - unlockNodeLinks(old_entry_point); + unlockNodeLinks(old_entry_point_id); return; } else { // Store this candidate which is currently being inserted into the graph in @@ -1085,7 +1074,7 @@ void HNSWIndex::replaceEntryPoint() { } } } - unlockNodeLinks(old_entry_point); + unlockNodeLinks(old_entry_point_id); // If there is no neighbors in the current level, check for any vector at // this level to be the new entry point. @@ -1219,8 +1208,9 @@ void HNSWIndex::greedySearchLevel(const void *vector_data, s } changed = false; + idType locked_id = bestCand; auto *element = getGraphDataByInternalId(bestCand); - lockNodeLinks(element); + lockNodeLinks(locked_id); ElementLevelData &node_level_data = getElementLevelData(element, level); for (int i = 0; i < node_level_data.getNumLinks(); i++) { @@ -1242,7 +1232,7 @@ void HNSWIndex::greedySearchLevel(const void *vector_data, s } } } - unlockNodeLinks(element); + unlockNodeLinks(locked_id); } while (changed); if (!running_query) { bestCand = bestNonDeletedCand; @@ -1257,17 +1247,17 @@ HNSWIndex::safeCollectAllNodeIncomingNeighbors(idType node_i auto element = getGraphDataByInternalId(node_id); for (size_t level = 0; level <= element->toplevel; level++) { // Save the node neighbor's in the current level while holding its neighbors lock. - lockNodeLinks(element); + lockNodeLinks(node_id); auto &node_level_data = getElementLevelData(element, level); // Store the deleted element's neighbours. auto neighbors_copy = node_level_data.copyLinks(); - unlockNodeLinks(element); + unlockNodeLinks(node_id); // Go over the neighbours and collect tho ones that also points back to the removed node. for (auto neighbour_id : neighbors_copy) { // Hold the neighbor's lock while we are going over its neighbors. auto *neighbor = getGraphDataByInternalId(neighbour_id); - lockNodeLinks(neighbor); + lockNodeLinks(neighbour_id); ElementLevelData &neighbour_level_data = getElementLevelData(neighbor, level); for (size_t j = 0; j < neighbour_level_data.getNumLinks(); j++) { @@ -1277,16 +1267,16 @@ HNSWIndex::safeCollectAllNodeIncomingNeighbors(idType node_i break; } } - unlockNodeLinks(neighbor); + unlockNodeLinks(neighbour_id); } // Next, collect the rest of incoming edges (the ones that are not bidirectional) in the // current level to repair them. - lockNodeLinks(element); + lockNodeLinks(node_id); for (auto incoming_edge : node_level_data.getIncomingEdges()) { incoming_neighbors.emplace_back(incoming_edge, (unsigned short)level); } - unlockNodeLinks(element); + unlockNodeLinks(node_id); } return incoming_neighbors; } @@ -1299,6 +1289,8 @@ void HNSWIndex::resizeIndexCommon(size_t new_max_elements) { idToMetaData.capacity(), new_max_elements); resizeLabelLookup(new_max_elements); visitedNodesHandlerPool.resize(new_max_elements); + elementLocks.resize(new_max_elements); + elementLocks.shrink_to_fit(); assert(idToMetaData.capacity() == idToMetaData.size()); idToMetaData.resize(new_max_elements); idToMetaData.shrink_to_fit(); @@ -1447,7 +1439,7 @@ void HNSWIndex::repairNodeConnections(idType node_id, size_t // Go over the repaired node neighbors, collect the non-deleted ones to be neighbors candidates // after the repair as well. auto *element = getGraphDataByInternalId(node_id); - lockNodeLinks(element); + lockNodeLinks(node_id); ElementLevelData &node_level_data = getElementLevelData(element, level); for (size_t j = 0; j < node_level_data.getNumLinks(); j++) { node_orig_neighbours_set[node_level_data.getLinkAtPos(j)] = true; @@ -1459,7 +1451,7 @@ void HNSWIndex::repairNodeConnections(idType node_id, size_t neighbors_candidates_set[node_level_data.getLinkAtPos(j)] = true; neighbors_candidate_ids.push_back(node_level_data.getLinkAtPos(j)); } - unlockNodeLinks(element); + unlockNodeLinks(node_id); // If there are not deleted neighbors at that point the repair job has already been made by // another parallel job, and there is no need to repair the node anymore. @@ -1478,7 +1470,7 @@ void HNSWIndex::repairNodeConnections(idType node_id, size_t nodes_to_update.push_back(deleted_neighbor_id); auto *neighbor = getGraphDataByInternalId(deleted_neighbor_id); - lockNodeLinks(neighbor); + lockNodeLinks(deleted_neighbor_id); ElementLevelData &neighbor_level_data = getElementLevelData(neighbor, level); for (size_t j = 0; j < neighbor_level_data.getNumLinks(); j++) { @@ -1492,7 +1484,7 @@ void HNSWIndex::repairNodeConnections(idType node_id, size_t neighbors_candidates_set[neighbor_level_data.getLinkAtPos(j)] = true; neighbors_candidate_ids.push_back(neighbor_level_data.getLinkAtPos(j)); } - unlockNodeLinks(neighbor); + unlockNodeLinks(deleted_neighbor_id); } size_t max_M_cur = level ? M : M0; @@ -1611,7 +1603,8 @@ HNSWIndex::HNSWIndex(const HNSWParams *params, size_t random_seed) : VecSimIndexAbstract(abstractInitParams, components), VecSimIndexTombstone(), maxElements(0), graphDataBlocks(this->allocator), - idToMetaData(this->allocator), visitedNodesHandlerPool(0, this->allocator) { + elementLocks(this->allocator), idToMetaData(this->allocator), + visitedNodesHandlerPool(0, this->allocator) { M = params->M ? params->M : HNSW_DEFAULT_M; M0 = M * 2; @@ -2335,7 +2328,7 @@ HNSWIndex::getHNSWElementNeighbors(size_t label, int ***neig } idType id = ids[0]; auto graph_data = this->getGraphDataByInternalId(id); - lockNodeLinks(graph_data); + lockNodeLinks(id); *neighborsData = new int *[graph_data->toplevel + 2]; for (size_t level = 0; level <= graph_data->toplevel; level++) { auto &level_data = this->getElementLevelData(graph_data, level); @@ -2347,7 +2340,7 @@ HNSWIndex::getHNSWElementNeighbors(size_t label, int ***neig } } (*neighborsData)[graph_data->toplevel + 1] = nullptr; - unlockNodeLinks(graph_data); + unlockNodeLinks(id); return VecSimDebugCommandCode_OK; } diff --git a/src/VecSim/algorithms/hnsw/hnsw_batch_iterator.h b/src/VecSim/algorithms/hnsw/hnsw_batch_iterator.h index e99642868..77e44738d 100644 --- a/src/VecSim/algorithms/hnsw/hnsw_batch_iterator.h +++ b/src/VecSim/algorithms/hnsw/hnsw_batch_iterator.h @@ -115,7 +115,7 @@ VecSimQueryReply_Code HNSW_BatchIterator::scanGraphInternal( // Take the current node out of the candidates queue and go over his neighbours. candidates.pop(); auto *node_graph_data = this->index->getGraphDataByInternalId(curr_node_id); - this->index->lockNodeLinks(node_graph_data); + this->index->lockNodeLinks(curr_node_id); ElementLevelData &node_level_data = this->index->getElementLevelData(node_graph_data, 0); if (node_level_data.numLinks > 0) { diff --git a/src/VecSim/algorithms/hnsw/hnsw_serializer_impl.h b/src/VecSim/algorithms/hnsw/hnsw_serializer_impl.h index 5f9dd8cbf..1666b7b37 100644 --- a/src/VecSim/algorithms/hnsw/hnsw_serializer_impl.h +++ b/src/VecSim/algorithms/hnsw/hnsw_serializer_impl.h @@ -18,7 +18,8 @@ HNSWIndex::HNSWIndex(std::ifstream &input, const HNSWParams HNSWSerializer::EncodingVersion version) : VecSimIndexAbstract(abstractInitParams, components), HNSWSerializer(version), epsilon(params->epsilon), graphDataBlocks(this->allocator), - idToMetaData(this->allocator), visitedNodesHandlerPool(0, this->allocator) { + elementLocks(this->allocator), idToMetaData(this->allocator), + visitedNodesHandlerPool(0, this->allocator) { this->restoreIndexFields(input); this->fieldsValidation(); @@ -30,6 +31,7 @@ HNSWIndex::HNSWIndex(std::ifstream &input, const HNSWParams // Set the initial capacity based on the number of elements in the loaded index. maxElements = RoundUpInitialCapacity(this->curElementCount, this->blockSize); + this->elementLocks.resize(maxElements); this->idToMetaData.resize(maxElements); this->visitedNodesHandlerPool.resize(maxElements); diff --git a/src/VecSim/index_factories/hnsw_factory.cpp b/src/VecSim/index_factories/hnsw_factory.cpp index cb30ea734..d577f57a1 100644 --- a/src/VecSim/index_factories/hnsw_factory.cpp +++ b/src/VecSim/index_factories/hnsw_factory.cpp @@ -133,9 +133,10 @@ size_t EstimateElementSize(const HNSWParams *params) { // label node (bucket). size_t size_label_lookup_entry = sizeof(void *); - // 1 entry in visited nodes + 1 entry in element metadata map + (approximately) 1 bucket in - // labels lookup hash map. - size_t size_meta_data = sizeof(tag_t) + sizeof(ElementMetaData) + size_label_lookup_entry; + // 1 entry in visited nodes + 1 entry in element metadata map + 1 node lock + (approximately) + // 1 bucket in labels lookup hash map. + size_t size_meta_data = sizeof(tag_t) + sizeof(ElementMetaData) + + sizeof(vecsim_stl::one_byte_mutex) + size_label_lookup_entry; /* Disclaimer: we are neglecting two additional factors that consume memory: * 1. The overall bucket size in labels_lookup hash table is usually higher than the number of diff --git a/src/VecSim/utils/vecsim_stl.h b/src/VecSim/utils/vecsim_stl.h index a55eb968c..aeb84d0eb 100644 --- a/src/VecSim/utils/vecsim_stl.h +++ b/src/VecSim/utils/vecsim_stl.h @@ -9,6 +9,8 @@ #pragma once #include "VecSim/memory/vecsim_base.h" +#include +#include #include #include #include @@ -108,4 +110,40 @@ class unordered_set alloc) {} }; +struct one_byte_mutex { + one_byte_mutex() noexcept = default; + one_byte_mutex(const one_byte_mutex &) noexcept : state(unlocked) {} + one_byte_mutex(one_byte_mutex &&) noexcept : state(unlocked) {} + one_byte_mutex &operator=(const one_byte_mutex &) noexcept { + state.store(unlocked, std::memory_order_relaxed); + return *this; + } + one_byte_mutex &operator=(one_byte_mutex &&) noexcept { + state.store(unlocked, std::memory_order_relaxed); + return *this; + } + + void lock() { + if (state.exchange(locked, std::memory_order_acquire) == unlocked) { + return; + } + while (state.exchange(sleeper, std::memory_order_acquire) != unlocked) { + state.wait(sleeper, std::memory_order_relaxed); + } + } + + void unlock() { + if (state.exchange(unlocked, std::memory_order_release) == sleeper) { + state.notify_one(); + } + } + +private: + std::atomic state{unlocked}; + + static constexpr uint8_t unlocked = 0; + static constexpr uint8_t locked = 0b01; + static constexpr uint8_t sleeper = 0b10; +}; + } // namespace vecsim_stl diff --git a/tests/unit/test_allocator.cpp b/tests/unit/test_allocator.cpp index 6aa4a0d0b..b5e4f76fb 100644 --- a/tests/unit/test_allocator.cpp +++ b/tests/unit/test_allocator.cpp @@ -573,7 +573,9 @@ TYPED_TEST(IndexAllocatorTest, test_hnsw_reclaim_memory) { data_containers_block_mem += size_total_data_per_element * block_size; // account for idToMetaData and visitedNodesHandlerPool entries. expected_mem_delta += - (sizeof(tag_t) + sizeof(ElementMetaData)) * block_size + data_containers_block_mem; + (sizeof(tag_t) + sizeof(ElementMetaData) + sizeof(vecsim_stl::one_byte_mutex)) * + block_size + + data_containers_block_mem; // Account for the allocation of a new bucket in the labels_lookup hash table. expected_mem_delta += (hnswIndex->labelLookup.bucket_count() - one_block_buckets) * sizeof(size_t); From 81dc0e687bd806d08e97ca22428d3fa20020b28b Mon Sep 17 00:00:00 2001 From: GuyAv46 Date: Tue, 7 Jul 2026 17:02:12 +0300 Subject: [PATCH 2/2] Use scoped enum for one byte mutex state --- src/VecSim/utils/vecsim_stl.h | 27 +++++++++++++++------------ 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/src/VecSim/utils/vecsim_stl.h b/src/VecSim/utils/vecsim_stl.h index aeb84d0eb..72131a818 100644 --- a/src/VecSim/utils/vecsim_stl.h +++ b/src/VecSim/utils/vecsim_stl.h @@ -112,38 +112,41 @@ class unordered_set struct one_byte_mutex { one_byte_mutex() noexcept = default; - one_byte_mutex(const one_byte_mutex &) noexcept : state(unlocked) {} - one_byte_mutex(one_byte_mutex &&) noexcept : state(unlocked) {} + one_byte_mutex(const one_byte_mutex &) noexcept : state(State::unlocked) {} + one_byte_mutex(one_byte_mutex &&) noexcept : state(State::unlocked) {} one_byte_mutex &operator=(const one_byte_mutex &) noexcept { - state.store(unlocked, std::memory_order_relaxed); + state.store(State::unlocked, std::memory_order_relaxed); return *this; } one_byte_mutex &operator=(one_byte_mutex &&) noexcept { - state.store(unlocked, std::memory_order_relaxed); + state.store(State::unlocked, std::memory_order_relaxed); return *this; } void lock() { - if (state.exchange(locked, std::memory_order_acquire) == unlocked) { + if (state.exchange(State::locked, std::memory_order_acquire) == State::unlocked) { return; } - while (state.exchange(sleeper, std::memory_order_acquire) != unlocked) { - state.wait(sleeper, std::memory_order_relaxed); + while (state.exchange(State::sleeper, std::memory_order_acquire) != State::unlocked) { + state.wait(State::sleeper, std::memory_order_relaxed); } } void unlock() { - if (state.exchange(unlocked, std::memory_order_release) == sleeper) { + if (state.exchange(State::unlocked, std::memory_order_release) == State::sleeper) { state.notify_one(); } } private: - std::atomic state{unlocked}; + enum class State : uint8_t { unlocked, locked, sleeper }; - static constexpr uint8_t unlocked = 0; - static constexpr uint8_t locked = 0b01; - static constexpr uint8_t sleeper = 0b10; + static_assert(sizeof(std::atomic) == sizeof(uint8_t), + "one_byte_mutex state must stay one byte"); + std::atomic state{State::unlocked}; }; +static_assert(sizeof(one_byte_mutex) == sizeof(uint8_t), "one_byte_mutex must stay one byte"); +static_assert(alignof(one_byte_mutex) == alignof(uint8_t), + "one_byte_mutex alignment must stay one byte"); } // namespace vecsim_stl