Merge pull request #21725 from vespa-engine/toregge/use-atomic-entry-ref-and-atomic-value-wrapper-in-reference-attribute

Use AtomicEntryRef and AtomicValueWrapper<uint32_t> in reference attribute.

7 files changed, 56 insertions, 44 deletions

diff --git a/searchlib/src/vespa/searchlib/attribute/imported_attribute_vector_read_guard.h b/searchlib/src/vespa/searchlib/attribute/imported_attribute_vector_read_guard.h
index 78748b3144a..502215f58cd 100644
--- a/searchlib/src/vespa/searchlib/attribute/imported_attribute_vector_read_guard.h
+++ b/searchlib/src/vespa/searchlib/attribute/imported_attribute_vector_read_guard.h
@@ -4,9 +4,10 @@
 
 #include "attribute_read_guard.h"
 #include "attributeguard.h"
-#include <vespa/vespalib/util/arrayref.h>
 #include <vespa/searchcommon/attribute/iattributevector.h>
 #include <vespa/searchlib/common/i_document_meta_store_context.h>
+#include <vespa/vespalib/datastore/atomic_value_wrapper.h>
+#include <vespa/vespalib/util/arrayref.h>
 
 namespace search::attribute {
 
@@ -29,7 +30,8 @@ class ImportedAttributeVectorReadGuard : public IAttributeVector,
                                          public AttributeReadGuard
 {
 private:
-    using TargetLids = vespalib::ConstArrayRef<uint32_t>;
+    using AtomicTargetLid = vespalib::datastore::AtomicValueWrapper<uint32_t>;
+    using TargetLids = vespalib::ConstArrayRef<AtomicTargetLid>;
     IDocumentMetaStoreContext::IReadGuard::UP _target_document_meta_store_read_guard;
     const ImportedAttributeVector   &_imported_attribute;
     TargetLids                       _targetLids;
@@ -42,7 +44,7 @@ protected:
 protected:
     uint32_t getTargetLid(uint32_t lid) const {
         // Check range to avoid reading memory beyond end of mapping array
-        return lid < _targetLids.size() ? _targetLids[lid] : 0u;
+        return lid < _targetLids.size() ? _targetLids[lid].load_acquire() : 0u;
     }
 
 public:
diff --git a/searchlib/src/vespa/searchlib/attribute/imported_search_context.cpp b/searchlib/src/vespa/searchlib/attribute/imported_search_context.cpp
index 70172199517..d949ea40c4a 100644
--- a/searchlib/src/vespa/searchlib/attribute/imported_search_context.cpp
+++ b/searchlib/src/vespa/searchlib/attribute/imported_search_context.cpp
@@ -150,7 +150,7 @@ TargetWeightedResult::getResult(ReverseMappingRefs reverseMappingRefs, const Rev
     }
     it->initRange(1, docIdLimit);
     for (uint32_t lid = it->seekFirst(1); !it->isAtEnd(); lid = it->seekNext(lid+1)) {
-        EntryRef revMapIdx = reverseMappingRefs[lid];
+        EntryRef revMapIdx = reverseMappingRefs[lid].load_acquire();
         if (__builtin_expect(revMapIdx.valid(), true)) {
             uint32_t size = reverseMapping.frozenSize(revMapIdx);
             targetResult.sizeSum += size;
@@ -175,7 +175,7 @@ TargetResult::getResult(ReverseMappingRefs reverseMappingRefs, const ReverseMapp
     }
     it->initRange(1, docIdLimit);
     for (uint32_t lid = it->seekFirst(1); !it->isAtEnd(); lid = it->seekNext(lid+1)) {
-        EntryRef revMapIdx = reverseMappingRefs[lid];
+        EntryRef revMapIdx = reverseMappingRefs[lid].load_acquire();
         if (__builtin_expect(revMapIdx.valid(), true)) {
             merger.addToBitVector(ReverseMappingBitVector(reverseMapping, revMapIdx));
         }
diff --git a/searchlib/src/vespa/searchlib/attribute/imported_search_context.h b/searchlib/src/vespa/searchlib/attribute/imported_search_context.h
index 63b52772363..0b36e5657f8 100644
--- a/searchlib/src/vespa/searchlib/attribute/imported_search_context.h
+++ b/searchlib/src/vespa/searchlib/attribute/imported_search_context.h
@@ -7,6 +7,7 @@
 #include <vespa/searchcommon/attribute/i_search_context.h>
 #include <vespa/searchlib/attribute/posting_list_merger.h>
 #include <vespa/searchlib/common/i_document_meta_store_context.h>
+#include <vespa/vespalib/datastore/atomic_value_wrapper.h>
 #include <vespa/vespalib/util/arrayref.h>
 
 namespace search::fef { class TermFieldMatchData; }
@@ -26,7 +27,8 @@ class SearchContextParams;
  * considered a match.
  */
 class ImportedSearchContext : public ISearchContext {
-    using TargetLids = vespalib::ConstArrayRef<uint32_t>;
+    using AtomicTargetLid = vespalib::datastore::AtomicValueWrapper<uint32_t>;
+    using TargetLids = vespalib::ConstArrayRef<AtomicTargetLid>;
     const ImportedAttributeVector&                  _imported_attribute;
     vespalib::string                                _queryTerm;
     bool                                            _useSearchCache;
@@ -41,7 +43,7 @@ class ImportedSearchContext : public ISearchContext {
 
     uint32_t getTargetLid(uint32_t lid) const {
         // Check range to avoid reading memory beyond end of mapping array
-        return lid < _targetLids.size() ? _targetLids[lid] : 0u;
+        return lid < _targetLids.size() ? _targetLids[lid].load_acquire() : 0u;
     }
 
     void makeMergedPostings(bool isFilter);
diff --git a/searchlib/src/vespa/searchlib/attribute/reference_attribute.cpp b/searchlib/src/vespa/searchlib/attribute/reference_attribute.cpp
index 57ff5ef464f..5b957af74fd 100644
--- a/searchlib/src/vespa/searchlib/attribute/reference_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/attribute/reference_attribute.cpp
@@ -75,7 +75,7 @@ ReferenceAttribute::addDoc(DocId &doc)
 {
     bool incGen = _indices.isFull();
     doc = _indices.size();
-    _indices.push_back(EntryRef());
+    _indices.push_back(AtomicEntryRef());
     _referenceMappings.addDoc();
     incNumDocs();
     updateUncommittedDocIdLimit(doc);
@@ -116,7 +116,7 @@ ReferenceAttribute::buildReverseMapping()
     uint32_t numDocs = _indices.size();
     indices.reserve(numDocs);
     for (uint32_t lid = 0; lid < numDocs; ++lid) {
-        EntryRef ref = _indices[lid];
+        EntryRef ref = _indices[lid].load_relaxed();
         if (ref.valid()) {
             indices.emplace_back(ref, lid);
         }
@@ -144,10 +144,10 @@ ReferenceAttribute::clearDoc(DocId doc)
 {
     updateUncommittedDocIdLimit(doc);
     assert(doc < _indices.size());
-    EntryRef oldRef = _indices[doc];
+    EntryRef oldRef = _indices[doc].load_relaxed();
     if (oldRef.valid()) {
         removeReverseMapping(oldRef, doc);
-        _indices[doc] = EntryRef();
+        _indices[doc].store_release(EntryRef());
         _store.remove(oldRef);
         return 1u;
     } else {
@@ -247,7 +247,7 @@ ReferenceAttribute::onLoad(vespalib::Executor *)
     _indices.unsafe_reserve(numDocs);
     for (uint32_t doc = 0; doc < numDocs; ++doc) {
         uint32_t enumValue = attrReader.getNextEnum();
-        _indices.push_back(builder.mapEnumValueToEntryRef(enumValue));
+        _indices.push_back(AtomicEntryRef(builder.mapEnumValueToEntryRef(enumValue)));
     }
     builder.makeDictionary();
     setNumDocs(numDocs);
@@ -262,11 +262,11 @@ ReferenceAttribute::update(DocId doc, const GlobalId &gid)
 {
     updateUncommittedDocIdLimit(doc);
     assert(doc < _indices.size());
-    EntryRef oldRef = _indices[doc];
+    EntryRef oldRef = _indices[doc].load_relaxed();
     Reference refToAdd(gid);
     EntryRef newRef = _store.add(refToAdd).ref();
     std::atomic_thread_fence(std::memory_order_release);
-    _indices[doc] = newRef;
+    _indices[doc].store_release(newRef);
     if (oldRef.valid()) {
         if (oldRef != newRef) {
             removeReverseMapping(oldRef, doc);
@@ -281,8 +281,10 @@ ReferenceAttribute::update(DocId doc, const GlobalId &gid)
 const Reference *
 ReferenceAttribute::getReference(DocId doc) const
 {
-    assert(doc < _indices.size());
-    EntryRef ref = _indices[doc];
+    if (doc >= getCommittedDocIdLimit()) {
+        return nullptr;
+    }
+    EntryRef ref = _indices.acquire_elem_ref(doc).load_acquire();
     if (!ref.valid()) {
         return nullptr;
     } else {
@@ -306,7 +308,7 @@ ReferenceAttribute::compact_worst_values(const CompactionStrategy& compaction_st
     CompactionSpec compaction_spec(true, true);
     auto remapper(_store.compact_worst(compaction_spec, compaction_strategy));
     if (remapper) {
-        remapper->remap(vespalib::ArrayRef<EntryRef>(&_indices[0], _indices.size()));
+        remapper->remap(vespalib::ArrayRef<AtomicEntryRef>(&_indices[0], _indices.size()));
         remapper->done();
     }
 }
@@ -335,7 +337,12 @@ ReferenceAttribute::IndicesCopyVector
 ReferenceAttribute::getIndicesCopy(uint32_t size) const
 {
     assert(size <= _indices.size());
-    return IndicesCopyVector(&_indices[0], &_indices[0] + size);
+    IndicesCopyVector result;
+    result.reserve(size);
+    for (uint32_t i = 0; i < size; ++i) {
+        result.push_back(_indices[i].load_relaxed());
+    }
+    return result;
 }
 
 void
@@ -428,10 +435,10 @@ ReferenceAttribute::clearDocs(DocId lidLow, DocId lidLimit)
     assert(lidLow <= lidLimit);
     assert(lidLimit <= getNumDocs());
     for (DocId lid = lidLow; lid < lidLimit; ++lid) {
-        EntryRef oldRef = _indices[lid];
+        EntryRef oldRef = _indices[lid].load_relaxed();
         if (oldRef.valid()) {
             removeReverseMapping(oldRef, lid);
-            _indices[lid] = EntryRef();
+            _indices[lid].store_release(EntryRef());
             _store.remove(oldRef);
         }
     }
diff --git a/searchlib/src/vespa/searchlib/attribute/reference_attribute.h b/searchlib/src/vespa/searchlib/attribute/reference_attribute.h
index a06aaf13247..96895fe2b35 100644
--- a/searchlib/src/vespa/searchlib/attribute/reference_attribute.h
+++ b/searchlib/src/vespa/searchlib/attribute/reference_attribute.h
@@ -31,7 +31,7 @@ public:
     using EntryRef = vespalib::datastore::EntryRef;
     using GlobalId = document::GlobalId;
     using ReferenceStore = vespalib::datastore::UniqueStore<Reference>;
-    using ReferenceStoreIndices = vespalib::RcuVectorBase<EntryRef>;
+    using ReferenceStoreIndices = vespalib::RcuVectorBase<AtomicEntryRef>;
     using IndicesCopyVector = std::vector<EntryRef, vespalib::allocator_large<EntryRef>>;
     // Class used to map from target lid to source lids
     using ReverseMapping = vespalib::btree::BTreeStore<uint32_t, vespalib::btree::BTreeNoLeafData,
diff --git a/searchlib/src/vespa/searchlib/attribute/reference_mappings.cpp b/searchlib/src/vespa/searchlib/attribute/reference_mappings.cpp
index 615cf530354..f922f7274ed 100644
--- a/searchlib/src/vespa/searchlib/attribute/reference_mappings.cpp
+++ b/searchlib/src/vespa/searchlib/attribute/reference_mappings.cpp
@@ -4,6 +4,7 @@
 #include "reference.h"
 #include <vespa/vespalib/datastore/datastore.hpp>
 #include <vespa/vespalib/btree/btreestore.hpp>
+#include <vespa/vespalib/util/rcuvector.hpp>
 
 namespace search::attribute {
 
@@ -37,7 +38,7 @@ ReferenceMappings::syncForwardMapping(const Reference &entry)
     auto &targetLids = _targetLids;
     _reverseMapping.foreach_unfrozen_key(revMapIdx,
                                          [&targetLids, targetLid](uint32_t lid)
-                                         { targetLids[lid] = targetLid; });
+                                         { targetLids[lid].store_release(targetLid); });
 }
 
 void
@@ -46,10 +47,9 @@ ReferenceMappings::syncReverseMappingIndices(const Reference &entry)
     uint32_t targetLid = entry.lid();
     if (targetLid != 0u) {
         _reverseMappingIndices.ensure_size(targetLid + 1);
-        _reverseMappingIndices[targetLid] = entry.revMapIdx();
-        if (targetLid >= _targetLidLimit) {
-            std::atomic_thread_fence(std::memory_order_release);
-            _targetLidLimit = targetLid + 1;
+        _reverseMappingIndices[targetLid].store_release(entry.revMapIdx());
+        if (targetLid >= _targetLidLimit.load(std::memory_order_relaxed)) {
+            _targetLidLimit.store(targetLid + 1, std::memory_order_release);
         }
     }
 }
@@ -62,7 +62,7 @@ ReferenceMappings::removeReverseMapping(const Reference &entry, uint32_t lid)
     std::atomic_thread_fence(std::memory_order_release);
     entry.setRevMapIdx(revMapIdx);
     syncReverseMappingIndices(entry);
-    _targetLids[lid] = 0; // forward mapping
+    _targetLids[lid].store_release(0); // forward mapping
 }
 
 void
@@ -74,7 +74,7 @@ ReferenceMappings::addReverseMapping(const Reference &entry, uint32_t lid)
     std::atomic_thread_fence(std::memory_order_release);
     entry.setRevMapIdx(revMapIdx);
     syncReverseMappingIndices(entry);
-    _targetLids[lid] = entry.lid(); // forward mapping
+    _targetLids[lid].store_release(entry.lid()); // forward mapping
 }
 
 void
@@ -92,7 +92,7 @@ ReferenceMappings::notifyReferencedPut(const Reference &entry, uint32_t targetLi
     uint32_t oldTargetLid = entry.lid();
     if (oldTargetLid != targetLid) {
         if (oldTargetLid != 0u && oldTargetLid < _reverseMappingIndices.size()) {
-            _reverseMappingIndices[oldTargetLid] = EntryRef();
+            _reverseMappingIndices[oldTargetLid].store_release(EntryRef());
         }
         entry.setLid(targetLid);
     }
@@ -106,7 +106,7 @@ ReferenceMappings::notifyReferencedRemove(const Reference &entry)
     uint32_t oldTargetLid = entry.lid();
     if (oldTargetLid != 0) {
         if (oldTargetLid < _reverseMappingIndices.size()) {
-            _reverseMappingIndices[oldTargetLid] = EntryRef();
+            _reverseMappingIndices[oldTargetLid].store_release(EntryRef());
         }
         entry.setLid(0);
     }
@@ -123,7 +123,7 @@ ReferenceMappings::onAddDocs(uint32_t docIdLimit)
 void
 ReferenceMappings::addDoc()
 {
-    _targetLids.push_back(0);
+    _targetLids.push_back(AtomicTargetLid(0));
 }
 
 void
diff --git a/searchlib/src/vespa/searchlib/attribute/reference_mappings.h b/searchlib/src/vespa/searchlib/attribute/reference_mappings.h
index 28aea85a343..b42d0b04b82 100644
--- a/searchlib/src/vespa/searchlib/attribute/reference_mappings.h
+++ b/searchlib/src/vespa/searchlib/attribute/reference_mappings.h
@@ -3,6 +3,7 @@
 #pragma once
 
 #include <vespa/vespalib/btree/btreestore.h>
+#include <vespa/vespalib/datastore/atomic_value_wrapper.h>
 #include <vespa/vespalib/util/rcuvector.h>
 #include <atomic>
 
@@ -15,10 +16,12 @@ class Reference;
  */
 class ReferenceMappings
 {
+    using AtomicEntryRef = vespalib::datastore::AtomicEntryRef;
+    using AtomicTargetLid = vespalib::datastore::AtomicValueWrapper<uint32_t>;
     using GenerationHolder = vespalib::GenerationHolder;
     using EntryRef = vespalib::datastore::EntryRef;
     // Classes used to map from target lid to source lids
-    using ReverseMappingIndices = vespalib::RcuVectorBase<EntryRef>;
+    using ReverseMappingIndices = vespalib::RcuVectorBase<AtomicEntryRef>;
     using ReverseMapping = vespalib::btree::BTreeStore<uint32_t, vespalib::btree::BTreeNoLeafData,
                                              vespalib::btree::NoAggregated,
                                              std::less<uint32_t>,
@@ -30,21 +33,21 @@ class ReferenceMappings
     // target lid.
     ReverseMappingIndices _reverseMappingIndices;
     // limit for target lid when accessing _reverseMappingIndices
-    uint32_t              _targetLidLimit;
+    std::atomic<uint32_t> _targetLidLimit;
     // Store of B-Trees, used to map from gid or target lid to
     // source lids.
     ReverseMapping _reverseMapping;
     // vector containing target lid given source lid
-    vespalib::RcuVectorBase<uint32_t> _targetLids;
+    vespalib::RcuVectorBase<AtomicTargetLid> _targetLids;
     const uint32_t &_committedDocIdLimit;
 
     void syncForwardMapping(const Reference &entry);
     void syncReverseMappingIndices(const Reference &entry);
 
 public:
-    using TargetLids = vespalib::ConstArrayRef<uint32_t>;
+    using TargetLids = vespalib::ConstArrayRef<AtomicTargetLid>;
     // Class used to map from target lid to source lids
-    using ReverseMappingRefs = vespalib::ConstArrayRef<EntryRef>;
+    using ReverseMappingRefs = vespalib::ConstArrayRef<AtomicEntryRef>;
 
     ReferenceMappings(GenerationHolder &genHolder, const uint32_t &committedDocIdLimit);
 
@@ -83,17 +86,15 @@ public:
 
     TargetLids getTargetLids() const {
         uint32_t committedDocIdLimit = _committedDocIdLimit;
-        std::atomic_thread_fence(std::memory_order_acquire);
-        return TargetLids(&_targetLids[0], committedDocIdLimit);
+        return TargetLids(&_targetLids.acquire_elem_ref(0), committedDocIdLimit);
     }
     uint32_t getTargetLid(uint32_t doc) const {
         // Check limit to avoid reading memory beyond end of valid mapping array
-        return doc < _committedDocIdLimit ? _targetLids[doc] : 0u;
+        return doc < _committedDocIdLimit ? _targetLids.acquire_elem_ref(doc).load_acquire() : 0u;
     }
     ReverseMappingRefs getReverseMappingRefs() const {
-        uint32_t targetLidLimit = _targetLidLimit;
-        std::atomic_thread_fence(std::memory_order_acquire);
-        return ReverseMappingRefs(&_reverseMappingIndices[0], targetLidLimit);
+        uint32_t targetLidLimit = _targetLidLimit.load(std::memory_order_acquire);
+        return ReverseMappingRefs(&_reverseMappingIndices.acquire_elem_ref(0), targetLidLimit);
     }
     const ReverseMapping &getReverseMapping() const { return _reverseMapping; }
 };
@@ -102,8 +103,8 @@ template <typename FunctionType>
 void
 ReferenceMappings::foreach_lid(uint32_t targetLid, FunctionType &&func) const
 {
-    if (targetLid < _reverseMappingIndices.size()) {
-        EntryRef revMapIdx = _reverseMappingIndices[targetLid];
+    if (targetLid < _targetLidLimit.load(std::memory_order_acquire)) {
+        EntryRef revMapIdx = _reverseMappingIndices.acquire_elem_ref(targetLid).load_acquire();
         _reverseMapping.foreach_frozen_key(revMapIdx, std::forward<FunctionType>(func));
     }
 }