Merge pull request #12097 from vespa-engine/geirst/atomic-entry-ref-in-hnsw-index

Implement wrapper for std::atomic of type EntryRef and use it in hnsw…

3 files changed, 17 insertions, 19 deletions

diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
index 8bd0833d6fb..4992a26d8b4 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
@@ -1,6 +1,7 @@
 // Copyright 2020 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
 #include "hnsw_index.h"
+#include <vespa/vespalib/util/rcuvector.hpp>
 
 namespace search::tensor {
 
@@ -98,9 +99,9 @@ void
 HnswIndex<FloatType>::add_document(uint32_t docid)
 {
     auto input = get_vector(docid);
-    _node_refs.ensure_size(docid + 1, EntryRef());
+    _node_refs.ensure_size(docid + 1, AtomicEntryRef());
     // A document cannot be added twice.
-    assert(!_node_refs[docid].valid());
+    assert(!_node_refs[docid].load_acquire().valid());
     int level = make_node_for_document(docid);
     if (_entry_docid == 0) {
         _entry_docid = docid;
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index_base.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index_base.cpp
index 8eae7ceb723..7a6493b6fcf 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index_base.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index_base.cpp
@@ -3,6 +3,7 @@
 #include "hnsw_index_base.h"
 #include "random_level_generator.h"
 #include <vespa/vespalib/datastore/array_store.hpp>
+#include <vespa/vespalib/util/rcuvector.hpp>
 
 namespace search::tensor {
 
@@ -45,18 +46,16 @@ HnswIndexBase::make_node_for_document(uint32_t docid)
     // TODO: Add capping on num_levels
     uint32_t num_levels = max_level + 1;
     // Note: The level array instance lives as long as the document is present in the index.
-    LevelArray levels(num_levels, EntryRef());
+    LevelArray levels(num_levels, AtomicEntryRef());
     auto node_ref = _nodes.add(levels);
-    // TODO: Add memory barrier?
-    _node_refs[docid] = node_ref;
+    _node_refs[docid].store_release(node_ref);
     return max_level;
 }
 
 HnswIndexBase::LevelArrayRef
 HnswIndexBase::get_level_array(uint32_t docid) const
 {
-    // TODO: Add memory barrier?
-    auto node_ref = _node_refs[docid];
+    auto node_ref = _node_refs[docid].load_acquire();
     return _nodes.get(node_ref);
 }
 
@@ -65,18 +64,16 @@ HnswIndexBase::get_link_array(uint32_t docid, uint32_t level) const
 {
     auto levels = get_level_array(docid);
     assert(level < levels.size());
-    return _links.get(levels[level]);
+    return _links.get(levels[level].load_acquire());
 }
 
 void
 HnswIndexBase::set_link_array(uint32_t docid, uint32_t level, const LinkArrayRef& links)
 {
     auto links_ref = _links.add(links);
-    // TODO: Add memory barrier?
-    auto node_ref = _node_refs[docid];
+    auto node_ref = _node_refs[docid].load_acquire();
     auto levels = _nodes.get_writable(node_ref);
-    // TODO: Make this change atomic.
-    levels[level] = links_ref;
+    levels[level].store_release(links_ref);
 }
 
 bool
@@ -165,14 +162,14 @@ HnswIndexBase::~HnswIndexBase() = default;
 HnswNode
 HnswIndexBase::get_node(uint32_t docid) const
 {
-    auto node_ref = _node_refs[docid];
+    auto node_ref = _node_refs[docid].load_acquire();
     if (!node_ref.valid()) {
         return HnswNode();
     }
     auto levels = _nodes.get(node_ref);
     HnswNode::LevelArray result;
     for (const auto& links_ref : levels) {
-        auto links = _links.get(links_ref);
+        auto links = _links.get(links_ref.load_acquire());
         HnswNode::LinkArray result_links(links.begin(), links.end());
         std::sort(result_links.begin(), result_links.end());
         result.push_back(result_links);
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index_base.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index_base.h
index 427afc38676..e92ca31e789 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index_base.h
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index_base.h
@@ -7,6 +7,7 @@
 #include "nearest_neighbor_index.h"
 #include <vespa/searchlib/common/bitvector.h>
 #include <vespa/vespalib/datastore/array_store.h>
+#include <vespa/vespalib/datastore/atomic_entry_ref.h>
 #include <vespa/vespalib/datastore/entryref.h>
 #include <vespa/vespalib/util/rcuvector.h>
 
@@ -54,7 +55,7 @@ public:
     };
 
 protected:
-    using EntryRef = search::datastore::EntryRef;
+    using AtomicEntryRef = search::datastore::AtomicEntryRef;
 
     // This uses 10 bits for buffer id -> 1024 buffers.
     // As we have very short arrays we get less fragmentation with fewer and larger buffers.
@@ -62,14 +63,13 @@ protected:
 
     // Provides mapping from document id -> node reference.
     // The reference is used to lookup the node data in NodeStore.
-    using NodeRefVector = vespalib::RcuVector<EntryRef>;
+    using NodeRefVector = vespalib::RcuVector<AtomicEntryRef>;
 
     // This stores the level arrays for all nodes.
     // Each node consists of an array of levels (from level 0 to n) where each entry is a reference to the link array at that level.
-    // TODO: Make replacing all links on a level atomically, e.g. AtomicEntryRef
-    using NodeStore = search::datastore::ArrayStore<EntryRef, EntryRefType>;
+    using NodeStore = search::datastore::ArrayStore<AtomicEntryRef, EntryRefType>;
     using LevelArrayRef = NodeStore::ConstArrayRef;
-    using LevelArray = vespalib::Array<EntryRef>;
+    using LevelArray = vespalib::Array<AtomicEntryRef>;
 
     // This stores all link arrays.
     // A link array consists of the document ids of the nodes a particular node is linked to.