Skip to content

[MOD-16696] Use one-byte HNSW node locks #983

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
GuyAv46 wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[MOD-16696] Use one-byte HNSW node locks #983

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e6da02b
Use one-byte HNSW node locks
GuyAv46 Jul 2, 2026
81dc0e6
Use scoped enum for one byte mutex state
GuyAv46 Jul 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 0 additions & 2 deletions src/VecSim/algorithms/hnsw/graph_data.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@

#include <cassert>
#include <algorithm>
#include <mutex>
#include "VecSim/utils/vec_utils.h"

// Amortized shrink thresholds for incoming edges vectors.
Expand Down Expand Up @@ -99,7 +98,6 @@ struct ElementLevelData {

struct ElementGraphData {
size_t toplevel;
std::mutex neighborsGuard;
ElementLevelData *others;
ElementLevelData level0;

Expand Down
85 changes: 39 additions & 46 deletions src/VecSim/algorithms/hnsw/hnsw.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,7 @@ class HNSWIndex : public VecSimIndexAbstract<DataType, DistType>,

// Index data
vecsim_stl::vector<DataBlock> graphDataBlocks;
mutable vecsim_stl::vector<vecsim_stl::one_byte_mutex> elementLocks;
vecsim_stl::vector<ElementMetaData> idToMetaData;

// Used for marking the visited nodes in graph scans (the pool supports parallel graph scans).
Expand Down Expand Up @@ -252,8 +253,6 @@ class HNSWIndex : public VecSimIndexAbstract<DataType, DistType>,
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;
Expand Down Expand Up @@ -481,24 +480,14 @@ void HNSWIndex<DataType, DistType>::unlockSharedIndexDataGuard() const {
indexDataGuard.unlock_shared();
}

template <typename DataType, typename DistType>
void HNSWIndex<DataType, DistType>::lockNodeLinks(ElementGraphData *node_data) const {
node_data->neighborsGuard.lock();
}

template <typename DataType, typename DistType>
void HNSWIndex<DataType, DistType>::unlockNodeLinks(ElementGraphData *node_data) const {
node_data->neighborsGuard.unlock();
}

template <typename DataType, typename DistType>
void HNSWIndex<DataType, DistType>::lockNodeLinks(idType node_id) const {
lockNodeLinks(getGraphDataByInternalId(node_id));
elementLocks[node_id].lock();
}

template <typename DataType, typename DistType>
void HNSWIndex<DataType, DistType>::unlockNodeLinks(idType node_id) const {
unlockNodeLinks(getGraphDataByInternalId(node_id));
elementLocks[node_id].unlock();
}

/**
Expand Down Expand Up @@ -528,7 +517,7 @@ void HNSWIndex<DataType, DistType>::processCandidate(
candidatesMaxHeap<DistType> &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) {
Expand Down Expand Up @@ -602,7 +591,7 @@ void HNSWIndex<DataType, DistType>::processCandidate(
}
}
}
unlockNodeLinks(cur_element);
unlockNodeLinks(curNodeId);
}

template <typename DataType, typename DistType>
Expand All @@ -612,7 +601,7 @@ void HNSWIndex<DataType, DistType>::processCandidate_RangeSearch(
candidatesMaxHeap<DistType> &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();

Expand Down Expand Up @@ -669,7 +658,7 @@ void HNSWIndex<DataType, DistType>::processCandidate_RangeSearch(
}
}
}
unlockNodeLinks(cur_element);
unlockNodeLinks(curNodeId);
}

template <typename DataType, typename DistType>
Expand Down Expand Up @@ -904,11 +893,11 @@ idType HNSWIndex<DataType, DistType>::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...
Expand All @@ -921,15 +910,15 @@ idType HNSWIndex<DataType, DistType>::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;
}

Expand All @@ -940,8 +929,8 @@ idType HNSWIndex<DataType, DistType>::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;
}

Expand Down Expand Up @@ -1068,15 +1057,15 @@ void HNSWIndex<DataType, DistType>::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.
for (size_t i = 0; i < old_ep_level.getNumLinks(); i++) {
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
Expand All @@ -1085,7 +1074,7 @@ void HNSWIndex<DataType, DistType>::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.
Expand Down Expand Up @@ -1219,8 +1208,9 @@ void HNSWIndex<DataType, DistType>::greedySearchLevel(const void *vector_data, s
}

changed = false;
idType locked_id = bestCand;
Comment thread
GuyAv46 marked this conversation as resolved.
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++) {
Expand All @@ -1242,7 +1232,7 @@ void HNSWIndex<DataType, DistType>::greedySearchLevel(const void *vector_data, s
}
}
}
unlockNodeLinks(element);
unlockNodeLinks(locked_id);
} while (changed);
if (!running_query) {
bestCand = bestNonDeletedCand;
Expand All @@ -1257,17 +1247,17 @@ HNSWIndex<DataType, DistType>::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++) {
Expand All @@ -1277,16 +1267,16 @@ HNSWIndex<DataType, DistType>::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;
}
Expand All @@ -1299,6 +1289,8 @@ void HNSWIndex<DataType, DistType>::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();
Comment thread
GuyAv46 marked this conversation as resolved.
assert(idToMetaData.capacity() == idToMetaData.size());
idToMetaData.resize(new_max_elements);
idToMetaData.shrink_to_fit();
Expand Down Expand Up @@ -1447,7 +1439,7 @@ void HNSWIndex<DataType, DistType>::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;
Expand All @@ -1459,7 +1451,7 @@ void HNSWIndex<DataType, DistType>::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.
Expand All @@ -1478,7 +1470,7 @@ void HNSWIndex<DataType, DistType>::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++) {
Expand All @@ -1492,7 +1484,7 @@ void HNSWIndex<DataType, DistType>::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;
Expand Down Expand Up @@ -1611,7 +1603,8 @@ HNSWIndex<DataType, DistType>::HNSWIndex(const HNSWParams *params,
size_t random_seed)
: VecSimIndexAbstract<DataType, DistType>(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;
Expand Down Expand Up @@ -2335,7 +2328,7 @@ HNSWIndex<DataType, DistType>::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);
Expand All @@ -2347,7 +2340,7 @@ HNSWIndex<DataType, DistType>::getHNSWElementNeighbors(size_t label, int ***neig
}
}
(*neighborsData)[graph_data->toplevel + 1] = nullptr;
unlockNodeLinks(graph_data);
unlockNodeLinks(id);
return VecSimDebugCommandCode_OK;
}

Expand Down
2 changes: 1 addition & 1 deletion src/VecSim/algorithms/hnsw/hnsw_batch_iterator.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -115,7 +115,7 @@ VecSimQueryReply_Code HNSW_BatchIterator<DataType, DistType>::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) {

Expand Down
4 changes: 3 additions & 1 deletion src/VecSim/algorithms/hnsw/hnsw_serializer_impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,8 @@ HNSWIndex<DataType, DistType>::HNSWIndex(std::ifstream &input, const HNSWParams
HNSWSerializer::EncodingVersion version)
: VecSimIndexAbstract<DataType, DistType>(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();
Expand All @@ -30,6 +31,7 @@ HNSWIndex<DataType, DistType>::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);

Expand Down
7 changes: 4 additions & 3 deletions src/VecSim/index_factories/hnsw_factory.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
41 changes: 41 additions & 0 deletions src/VecSim/utils/vecsim_stl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
#pragma once

#include "VecSim/memory/vecsim_base.h"
#include <atomic>
#include <cstdint>
#include <vector>
#include <algorithm>
#include <set>
Expand Down Expand Up @@ -108,4 +110,43 @@ class unordered_set
alloc) {}
};

struct one_byte_mutex {
one_byte_mutex() noexcept = default;
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(State::unlocked, std::memory_order_relaxed);
return *this;
}
one_byte_mutex &operator=(one_byte_mutex &&) noexcept {
state.store(State::unlocked, std::memory_order_relaxed);
return *this;
}

void lock() {
if (state.exchange(State::locked, std::memory_order_acquire) == State::unlocked) {
return;
}
while (state.exchange(State::sleeper, std::memory_order_acquire) != State::unlocked) {
state.wait(State::sleeper, std::memory_order_relaxed);
}
}

void unlock() {
if (state.exchange(State::unlocked, std::memory_order_release) == State::sleeper) {
state.notify_one();
}
Comment thread
GuyAv46 marked this conversation as resolved.
}

private:
enum class State : uint8_t { unlocked, locked, sleeper };

static_assert(sizeof(std::atomic<State>) == sizeof(uint8_t),
"one_byte_mutex state must stay one byte");
std::atomic<State> 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
Loading
Loading