Merge pull request #21145 from vespa-engine/toregge/add-memory-allocator-to-array-store

Add memory allocator to array store.

15 files changed, 210 insertions, 55 deletions

diff --git a/searchlib/src/vespa/searchlib/attribute/multi_value_mapping.cpp b/searchlib/src/vespa/searchlib/attribute/multi_value_mapping.cpp
index 1491a92166a..161e2883cbc 100644
--- a/searchlib/src/vespa/searchlib/attribute/multi_value_mapping.cpp
+++ b/searchlib/src/vespa/searchlib/attribute/multi_value_mapping.cpp
@@ -5,8 +5,8 @@
 #include "multi_value_mapping.h"
 #include "multi_value_mapping.hpp"
 #include "multivalue.h"
-#include <vespa/vespalib/util/array.hpp>
 #include <vespa/vespalib/datastore/buffer_type.hpp>
+#include <vespa/vespalib/util/array.hpp>
 
 using search::multivalue::Value;
 using search::multivalue::WeightedValue;
diff --git a/searchlib/src/vespa/searchlib/attribute/multi_value_mapping.hpp b/searchlib/src/vespa/searchlib/attribute/multi_value_mapping.hpp
index fb81a60cb13..339f562757d 100644
--- a/searchlib/src/vespa/searchlib/attribute/multi_value_mapping.hpp
+++ b/searchlib/src/vespa/searchlib/attribute/multi_value_mapping.hpp
@@ -19,7 +19,7 @@ MultiValueMapping<EntryT,RefT>::MultiValueMapping(const vespalib::datastore::Arr
 #ifdef __clang__
 #pragma clang diagnostic pop
 #endif
-      _store(storeCfg)
+      _store(storeCfg, {})
 {
 }
 
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_graph.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_graph.cpp
index 6eb215273ad..8beb111ca59 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_graph.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_graph.cpp
@@ -10,8 +10,8 @@ namespace search::tensor {
 HnswGraph::HnswGraph()
   : node_refs(),
     node_refs_size(1u),
-    nodes(HnswIndex::make_default_node_store_config()),
-    links(HnswIndex::make_default_link_store_config()),
+    nodes(HnswIndex::make_default_node_store_config(), {}),
+    links(HnswIndex::make_default_link_store_config(), {}),
     entry_docid_and_level()
 {
     node_refs.ensure_size(1, AtomicEntryRef());
diff --git a/storage/src/vespa/storage/bucketdb/btree_bucket_database.cpp b/storage/src/vespa/storage/bucketdb/btree_bucket_database.cpp
index 742b5eeb671..f909a939efb 100644
--- a/storage/src/vespa/storage/bucketdb/btree_bucket_database.cpp
+++ b/storage/src/vespa/storage/bucketdb/btree_bucket_database.cpp
@@ -100,7 +100,7 @@ struct BTreeBucketDatabase::ReplicaValueTraits {
 template class bucketdb::GenericBTreeBucketDatabase<BTreeBucketDatabase::ReplicaValueTraits>;
 
 BTreeBucketDatabase::BTreeBucketDatabase()
-    : _impl(std::make_unique<ImplType>(make_default_array_store_config<ReplicaValueTraits::DataStoreType>()))
+    : _impl(std::make_unique<ImplType>(make_default_array_store_config<ReplicaValueTraits::DataStoreType>(), std::shared_ptr<vespalib::alloc::MemoryAllocator>()))
 {
 }
 
diff --git a/vespalib/src/tests/datastore/array_store/array_store_test.cpp b/vespalib/src/tests/datastore/array_store/array_store_test.cpp
index c58e357a9a1..2672ed70f6d 100644
--- a/vespalib/src/tests/datastore/array_store/array_store_test.cpp
+++ b/vespalib/src/tests/datastore/array_store/array_store_test.cpp
@@ -7,6 +7,7 @@
 #include <vespa/vespalib/datastore/compaction_strategy.h>
 #include <vespa/vespalib/stllike/hash_map.hpp>
 #include <vespa/vespalib/testkit/testapp.h>
+#include <vespa/vespalib/test/memory_allocator_observer.h>
 #include <vespa/vespalib/test/insertion_operators.h>
 #include <vespa/vespalib/util/memory_allocator.h>
 #include <vespa/vespalib/util/size_literals.h>
@@ -19,6 +20,9 @@ using vespalib::ArrayRef;
 using generation_t = vespalib::GenerationHandler::generation_t;
 using MemStats = vespalib::datastore::test::MemStats;
 using BufferStats = vespalib::datastore::test::BufferStats;
+using vespalib::alloc::MemoryAllocator;
+using vespalib::alloc::test::MemoryAllocatorObserver;
+using AllocStats = MemoryAllocatorObserver::Stats;
 
 namespace {
 
@@ -40,18 +44,20 @@ struct Fixture
     using value_type = EntryT;
     using ReferenceStore = vespalib::hash_map<EntryRef, EntryVector>;
 
+    AllocStats     stats;
     ArrayStoreType store;
     ReferenceStore refStore;
     generation_t generation;
     Fixture(uint32_t maxSmallArraySize, bool enable_free_lists = true)
         : store(ArrayStoreConfig(maxSmallArraySize,
                                  ArrayStoreConfig::AllocSpec(16, RefT::offsetSize(), 8_Ki,
-                                                             ALLOC_GROW_FACTOR)).enable_free_lists(enable_free_lists)),
+                                                             ALLOC_GROW_FACTOR)).enable_free_lists(enable_free_lists),
+                std::make_unique<MemoryAllocatorObserver>(stats)),
           refStore(),
           generation(1)
     {}
     Fixture(const ArrayStoreConfig &storeCfg)
-        : store(storeCfg),
+        : store(storeCfg, std::make_unique<MemoryAllocatorObserver>(stats)),
           refStore(),
           generation(1)
     {}
@@ -162,10 +168,10 @@ TEST("require that we test with trivial and non-trivial types")
 
 TEST_F("control static sizes", NumberFixture(3)) {
 #ifdef _LIBCPP_VERSION
-    EXPECT_EQUAL(424u, sizeof(f.store));
+    EXPECT_EQUAL(440u, sizeof(f.store));
     EXPECT_EQUAL(296u, sizeof(NumberFixture::ArrayStoreType::DataStoreType));
 #else
-    EXPECT_EQUAL(456u, sizeof(f.store));
+    EXPECT_EQUAL(472u, sizeof(f.store));
     EXPECT_EQUAL(328u, sizeof(NumberFixture::ArrayStoreType::DataStoreType));
 #endif
     EXPECT_EQUAL(96u, sizeof(NumberFixture::ArrayStoreType::SmallArrayType));
@@ -447,4 +453,9 @@ TEST_F("require that offset in EntryRefT is within bounds when allocating memory
     f.assertStoreContent();
 }
 
+TEST_F("require that provided memory allocator is used", NumberFixture(3))
+{
+    EXPECT_EQUAL(AllocStats(4, 0), f.stats);
+}
+
 TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/vespalib/src/vespa/vespalib/datastore/CMakeLists.txt b/vespalib/src/vespa/vespalib/datastore/CMakeLists.txt
index 5ff1eab61b8..a2f7a4d4ff4 100644
--- a/vespalib/src/vespa/vespalib/datastore/CMakeLists.txt
+++ b/vespalib/src/vespa/vespalib/datastore/CMakeLists.txt
@@ -12,7 +12,9 @@ vespa_add_library(vespalib_vespalib_datastore OBJECT
     entryref.cpp
     entry_ref_filter.cpp
     fixed_size_hash_map.cpp
+    large_array_buffer_type.cpp
     sharded_hash_map.cpp
+    small_array_buffer_type.cpp
     unique_store.cpp
     unique_store_buffer_type.cpp
     unique_store_string_allocator.cpp
diff --git a/vespalib/src/vespa/vespalib/datastore/array_store.cpp b/vespalib/src/vespa/vespalib/datastore/array_store.cpp
index 0e03b36d2f9..b03f21402a5 100644
--- a/vespalib/src/vespa/vespalib/datastore/array_store.cpp
+++ b/vespalib/src/vespa/vespalib/datastore/array_store.cpp
@@ -5,10 +5,4 @@
 
 namespace vespalib::datastore {
 
-template class BufferType<vespalib::Array<uint8_t>>;
-template class BufferType<vespalib::Array<uint32_t>>;
-template class BufferType<vespalib::Array<int32_t>>;
-template class BufferType<vespalib::Array<std::string>>;
-template class BufferType<vespalib::Array<AtomicEntryRef>>;
-
 }
diff --git a/vespalib/src/vespa/vespalib/datastore/array_store.h b/vespalib/src/vespa/vespalib/datastore/array_store.h
index d9b62c310b5..ed3af451b04 100644
--- a/vespalib/src/vespa/vespalib/datastore/array_store.h
+++ b/vespalib/src/vespa/vespalib/datastore/array_store.h
@@ -9,6 +9,8 @@
 #include "entryref.h"
 #include "atomic_entry_ref.h"
 #include "i_compaction_context.h"
+#include "large_array_buffer_type.h"
+#include "small_array_buffer_type.h"
 #include <vespa/vespalib/util/array.h>
 
 namespace vespalib::datastore {
@@ -32,27 +34,15 @@ public:
     using SmallArrayType = BufferType<EntryT>;
     using LargeArray = vespalib::Array<EntryT>;
     using AllocSpec = ArrayStoreConfig::AllocSpec;
-
 private:
-    class LargeArrayType : public BufferType<LargeArray> {
-    private:
-        using ParentType = BufferType<LargeArray>;
-        using ParentType::_emptyEntry;
-        using CleanContext = typename ParentType::CleanContext;
-    public:
-        LargeArrayType(const AllocSpec &spec);
-        void cleanHold(void *buffer, size_t offset, ElemCount numElems, CleanContext cleanCtx) override;
-    };
-
-
     uint32_t _largeArrayTypeId;
     uint32_t _maxSmallArraySize;
     DataStoreType _store;
-    std::vector<SmallArrayType> _smallArrayTypes;
-    LargeArrayType _largeArrayType;
+    std::vector<SmallArrayBufferType<EntryT>> _smallArrayTypes;
+    LargeArrayBufferType<EntryT> _largeArrayType;
     using generation_t = vespalib::GenerationHandler::generation_t;
 
-    void initArrayTypes(const ArrayStoreConfig &cfg);
+    void initArrayTypes(const ArrayStoreConfig &cfg, std::shared_ptr<alloc::MemoryAllocator> memory_allocator);
     // 1-to-1 mapping between type ids and sizes for small arrays is enforced during initialization.
     uint32_t getTypeId(size_t arraySize) const { return arraySize; }
     size_t getArraySize(uint32_t typeId) const { return typeId; }
@@ -68,7 +58,7 @@ private:
     }
 
 public:
-    ArrayStore(const ArrayStoreConfig &cfg);
+    ArrayStore(const ArrayStoreConfig &cfg, std::shared_ptr<alloc::MemoryAllocator> memory_allocator);
     ~ArrayStore();
     EntryRef add(const ConstArrayRef &array);
     ConstArrayRef get(EntryRef ref) const {
diff --git a/vespalib/src/vespa/vespalib/datastore/array_store.hpp b/vespalib/src/vespa/vespalib/datastore/array_store.hpp
index bbbd52c354d..00c1615b173 100644
--- a/vespalib/src/vespa/vespalib/datastore/array_store.hpp
+++ b/vespalib/src/vespa/vespalib/datastore/array_store.hpp
@@ -6,40 +6,23 @@
 #include "compaction_spec.h"
 #include "entry_ref_filter.h"
 #include "datastore.hpp"
+#include "large_array_buffer_type.hpp"
+#include "small_array_buffer_type.hpp"
 #include <atomic>
 #include <algorithm>
 
 namespace vespalib::datastore {
 
 template <typename EntryT, typename RefT>
-ArrayStore<EntryT, RefT>::LargeArrayType::LargeArrayType(const AllocSpec &spec)
-    : BufferType<LargeArray>(1, spec.minArraysInBuffer, spec.maxArraysInBuffer, spec.numArraysForNewBuffer, spec.allocGrowFactor)
-{
-}
-
-template <typename EntryT, typename RefT>
-void
-ArrayStore<EntryT, RefT>::LargeArrayType::cleanHold(void *buffer, size_t offset, ElemCount numElems, CleanContext cleanCtx)
-{
-    LargeArray *elem = static_cast<LargeArray *>(buffer) + offset;
-    for (size_t i = 0; i < numElems; ++i) {
-        cleanCtx.extraBytesCleaned(sizeof(EntryT) * elem->size());
-        *elem = _emptyEntry;
-        ++elem;
-    }
-}
-
-template <typename EntryT, typename RefT>
 void
-ArrayStore<EntryT, RefT>::initArrayTypes(const ArrayStoreConfig &cfg)
+ArrayStore<EntryT, RefT>::initArrayTypes(const ArrayStoreConfig &cfg, std::shared_ptr<alloc::MemoryAllocator> memory_allocator)
 {
     _largeArrayTypeId = _store.addType(&_largeArrayType);
     assert(_largeArrayTypeId == 0);
     _smallArrayTypes.reserve(_maxSmallArraySize);
     for (uint32_t arraySize = 1; arraySize <= _maxSmallArraySize; ++arraySize) {
         const AllocSpec &spec = cfg.specForSize(arraySize);
-        _smallArrayTypes.emplace_back(arraySize, spec.minArraysInBuffer, spec.maxArraysInBuffer,
-                                      spec.numArraysForNewBuffer, spec.allocGrowFactor);
+        _smallArrayTypes.emplace_back(arraySize, spec, memory_allocator);
     }
     for (auto & type : _smallArrayTypes) {
         uint32_t typeId = _store.addType(&type);
@@ -48,14 +31,14 @@ ArrayStore<EntryT, RefT>::initArrayTypes(const ArrayStoreConfig &cfg)
 }
 
 template <typename EntryT, typename RefT>
-ArrayStore<EntryT, RefT>::ArrayStore(const ArrayStoreConfig &cfg)
+ArrayStore<EntryT, RefT>::ArrayStore(const ArrayStoreConfig &cfg, std::shared_ptr<alloc::MemoryAllocator> memory_allocator)
     : _largeArrayTypeId(0),
       _maxSmallArraySize(cfg.maxSmallArraySize()),
       _store(),
       _smallArrayTypes(),
-      _largeArrayType(cfg.specForSize(0))
+      _largeArrayType(cfg.specForSize(0), memory_allocator)
 {
-    initArrayTypes(cfg);
+    initArrayTypes(cfg, std::move(memory_allocator));
     _store.init_primary_buffers();
     if (cfg.enable_free_lists()) {
         _store.enableFreeLists();
diff --git a/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.cpp b/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.cpp
new file mode 100644
index 00000000000..f4ccb27abad
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.cpp
@@ -0,0 +1,20 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "large_array_buffer_type.hpp"
+#include "buffer_type.hpp"
+
+namespace vespalib::datastore {
+
+template class BufferType<Array<uint8_t>>;
+template class BufferType<Array<uint32_t>>;
+template class BufferType<Array<int32_t>>;
+template class BufferType<Array<std::string>>;
+template class BufferType<Array<AtomicEntryRef>>;
+
+template class LargeArrayBufferType<uint8_t>;
+template class LargeArrayBufferType<uint32_t>;
+template class LargeArrayBufferType<int32_t>;
+template class LargeArrayBufferType<std::string>;
+template class LargeArrayBufferType<AtomicEntryRef>;
+
+}
diff --git a/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.h b/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.h
new file mode 100644
index 00000000000..50d15d4a27c
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.h
@@ -0,0 +1,39 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include "array_store_config.h"
+#include "buffer_type.h"
+#include <vespa/vespalib/util/array.h>
+#include <memory>
+
+namespace vespalib::alloc { class MemoryAllocator; }
+
+namespace vespalib::datastore {
+
+/*
+ * Class representing buffer type for large arrays in ArrayStore
+ */
+template <typename EntryT>
+class LargeArrayBufferType : public BufferType<Array<EntryT>>
+{
+    using AllocSpec = ArrayStoreConfig::AllocSpec;
+    using ArrayType = Array<EntryT>;
+    using ParentType = BufferType<ArrayType>;
+    using ParentType::_emptyEntry;
+    using CleanContext = typename ParentType::CleanContext;
+    std::shared_ptr<alloc::MemoryAllocator> _memory_allocator;
+public:
+    LargeArrayBufferType(const AllocSpec& spec, std::shared_ptr<alloc::MemoryAllocator> memory_allocator) noexcept;
+    ~LargeArrayBufferType() override;
+    void cleanHold(void* buffer, size_t offset, ElemCount numElems, CleanContext cleanCtx) override;
+    const vespalib::alloc::MemoryAllocator* get_memory_allocator() const override;
+};
+
+extern template class LargeArrayBufferType<uint8_t>;
+extern template class LargeArrayBufferType<uint32_t>;
+extern template class LargeArrayBufferType<int32_t>;
+extern template class LargeArrayBufferType<std::string>;
+extern template class LargeArrayBufferType<AtomicEntryRef>;
+
+}
diff --git a/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.hpp b/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.hpp
new file mode 100644
index 00000000000..aeeef8166c6
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.hpp
@@ -0,0 +1,39 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include "large_array_buffer_type.h"
+#include <vespa/vespalib/util/array.hpp>
+
+namespace vespalib::datastore {
+
+template <typename EntryT>
+LargeArrayBufferType<EntryT>::LargeArrayBufferType(const AllocSpec& spec, std::shared_ptr<alloc::MemoryAllocator> memory_allocator) noexcept
+    : BufferType<Array<EntryT>>(1u, spec.minArraysInBuffer, spec.maxArraysInBuffer, spec.numArraysForNewBuffer, spec.allocGrowFactor),
+      _memory_allocator(std::move(memory_allocator))
+{
+}
+
+template <typename EntryT>
+LargeArrayBufferType<EntryT>::~LargeArrayBufferType() = default;
+
+template <typename EntryT>
+void
+LargeArrayBufferType<EntryT>::cleanHold(void* buffer, size_t offset, ElemCount numElems, CleanContext cleanCtx)
+{
+    ArrayType* elem = static_cast<ArrayType*>(buffer) + offset;
+    for (size_t i = 0; i < numElems; ++i) {
+        cleanCtx.extraBytesCleaned(sizeof(EntryT) * elem->size());
+        *elem = _emptyEntry;
+        ++elem;
+    }
+}
+
+template <typename EntryT>
+const vespalib::alloc::MemoryAllocator*
+LargeArrayBufferType<EntryT>::get_memory_allocator() const
+{
+    return _memory_allocator.get();
+}
+
+}
diff --git a/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.cpp b/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.cpp
new file mode 100644
index 00000000000..06a4c6007a9
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.cpp
@@ -0,0 +1,14 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "small_array_buffer_type.hpp"
+#include "buffer_type.hpp"
+
+namespace vespalib::datastore {
+
+template class SmallArrayBufferType<uint8_t>;
+template class SmallArrayBufferType<uint32_t>;
+template class SmallArrayBufferType<int32_t>;
+template class SmallArrayBufferType<std::string>;
+template class SmallArrayBufferType<AtomicEntryRef>;
+
+}
diff --git a/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.h b/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.h
new file mode 100644
index 00000000000..4e4568c3d16
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.h
@@ -0,0 +1,37 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include "array_store_config.h"
+#include "buffer_type.h"
+#include <memory>
+
+namespace vespalib::alloc { class MemoryAllocator; }
+
+namespace vespalib::datastore {
+
+/*
+ * Class representing buffer type for small arrays in ArrayStore
+ */
+template <typename EntryT>
+class SmallArrayBufferType : public BufferType<EntryT>
+{
+    using AllocSpec = ArrayStoreConfig::AllocSpec;
+    std::shared_ptr<alloc::MemoryAllocator> _memory_allocator;
+public:
+    SmallArrayBufferType(const SmallArrayBufferType&) = delete;
+    SmallArrayBufferType& operator=(const SmallArrayBufferType&) = delete;
+    SmallArrayBufferType(SmallArrayBufferType&&) noexcept = default;
+    SmallArrayBufferType& operator=(SmallArrayBufferType&&) noexcept = default;
+    SmallArrayBufferType(uint32_t array_size, const AllocSpec& spec, std::shared_ptr<alloc::MemoryAllocator> memory_allocator) noexcept;
+    ~SmallArrayBufferType() override;
+    const vespalib::alloc::MemoryAllocator* get_memory_allocator() const override;
+};
+
+extern template class SmallArrayBufferType<uint8_t>;
+extern template class SmallArrayBufferType<uint32_t>;
+extern template class SmallArrayBufferType<int32_t>;
+extern template class SmallArrayBufferType<std::string>;
+extern template class SmallArrayBufferType<AtomicEntryRef>;
+
+}
diff --git a/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.hpp b/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.hpp
new file mode 100644
index 00000000000..414804417eb
--- /dev/null
+++ b/vespalib/src/vespa/vespalib/datastore/small_array_buffer_type.hpp
@@ -0,0 +1,26 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include "small_array_buffer_type.h"
+
+namespace vespalib::datastore {
+
+template <typename EntryT>
+SmallArrayBufferType<EntryT>::SmallArrayBufferType(uint32_t array_size, const AllocSpec& spec, std::shared_ptr<alloc::MemoryAllocator> memory_allocator) noexcept
+    : BufferType<EntryT>(array_size, spec.minArraysInBuffer, spec.maxArraysInBuffer, spec.numArraysForNewBuffer, spec.allocGrowFactor),
+      _memory_allocator(std::move(memory_allocator))
+{
+}
+
+template <typename EntryT>
+SmallArrayBufferType<EntryT>::~SmallArrayBufferType() = default;
+
+template <typename EntryT>
+const vespalib::alloc::MemoryAllocator*
+SmallArrayBufferType<EntryT>::get_memory_allocator() const
+{
+    return _memory_allocator.get();
+}
+
+}