Merge pull request #21157 from vespa-engine/toregge/add-memory-allocator-to-multi-value-mapping

Add memory allocator to multi value mapping.

6 files changed, 28 insertions, 11 deletions

diff --git a/searchlib/src/tests/attribute/multi_value_mapping/multi_value_mapping_test.cpp b/searchlib/src/tests/attribute/multi_value_mapping/multi_value_mapping_test.cpp
index bddaa4f4e31..29af989d484 100644
--- a/searchlib/src/tests/attribute/multi_value_mapping/multi_value_mapping_test.cpp
+++ b/searchlib/src/tests/attribute/multi_value_mapping/multi_value_mapping_test.cpp
@@ -6,6 +6,7 @@
 #include <vespa/vespalib/gtest/gtest.h>
 #include <vespa/vespalib/stllike/hash_set.h>
 #include <vespa/vespalib/test/insertion_operators.h>
+#include <vespa/vespalib/test/memory_allocator_observer.h>
 #include <vespa/vespalib/util/generationhandler.h>
 #include <vespa/vespalib/util/rand48.h>
 #include <vespa/vespalib/util/size_literals.h>
@@ -16,6 +17,8 @@ LOG_SETUP("multivaluemapping_test");
 using vespalib::datastore::ArrayStoreConfig;
 using vespalib::datastore::CompactionSpec;
 using vespalib::datastore::CompactionStrategy;
+using vespalib::alloc::test::MemoryAllocatorObserver;
+using AllocStats = MemoryAllocatorObserver::Stats;
 
 template <typename EntryT>
 void
@@ -69,6 +72,7 @@ class MappingTestBase : public ::testing::Test {
 protected:
     using MvMapping = search::attribute::MultiValueMapping<EntryT>;
     using AttributeType = MyAttribute<MvMapping>;
+    AllocStats _stats;
     std::unique_ptr<MvMapping> _mvMapping;
     std::unique_ptr<AttributeType> _attr;
     uint32_t _maxSmallArraySize;
@@ -78,7 +82,8 @@ protected:
 public:
     using ConstArrayRef = vespalib::ConstArrayRef<EntryT>;
     MappingTestBase()
-        : _mvMapping(),
+        : _stats(),
+          _mvMapping(),
           _attr(),
           _maxSmallArraySize()
     {
@@ -87,7 +92,7 @@ public:
         ArrayStoreConfig config(maxSmallArraySize,
                                 ArrayStoreConfig::AllocSpec(0, RefType::offsetSize(), 8_Ki, ALLOC_GROW_FACTOR));
         config.enable_free_lists(enable_free_lists);
-        _mvMapping = std::make_unique<MvMapping>(config);
+        _mvMapping = std::make_unique<MvMapping>(config, vespalib::GrowStrategy(), std::make_unique<MemoryAllocatorObserver>(_stats));
         _attr = std::make_unique<AttributeType>(*_mvMapping);
         _maxSmallArraySize = maxSmallArraySize;
     }
@@ -95,7 +100,7 @@ public:
         ArrayStoreConfig config(maxSmallArraySize,
                                 ArrayStoreConfig::AllocSpec(minArrays, maxArrays, numArraysForNewBuffer, ALLOC_GROW_FACTOR));
         config.enable_free_lists(enable_free_lists);
-        _mvMapping = std::make_unique<MvMapping>(config);
+        _mvMapping = std::make_unique<MvMapping>(config, vespalib::GrowStrategy(), std::make_unique<MemoryAllocatorObserver>(_stats));
         _attr = std::make_unique<AttributeType>(*_mvMapping);
         _maxSmallArraySize = maxSmallArraySize;
     }
@@ -129,6 +134,7 @@ public:
         _mvMapping->clearDocs(lidLow, lidLimit, [this](uint32_t docId) { _attr->clearDoc(docId); });
     }
     size_t getTotalValueCnt() const { return _mvMapping->getTotalValueCnt(); }
+    const AllocStats &get_stats() const noexcept { return _stats; }
 
     uint32_t countBuffers() {
         using RefVector = typename MvMapping::RefCopyVector;
@@ -326,6 +332,12 @@ TEST_F(IntMappingTest, test_that_free_lists_can_be_disabled)
     EXPECT_FALSE(_mvMapping->has_free_lists_enabled());
 }
 
+TEST_F(IntMappingTest, provided_memory_allocator_is_used)
+{
+    setup(3, 64, 512, 129, true);
+    EXPECT_EQ(AllocStats(5, 0), get_stats());
+}
+
 TEST_F(CompactionIntMappingTest, test_that_compaction_works)
 {
     setup(3, 64, 512, 129);
diff --git a/searchlib/src/vespa/searchlib/attribute/multi_value_mapping.h b/searchlib/src/vespa/searchlib/attribute/multi_value_mapping.h
index 81abaa05a45..f5f2950a59c 100644
--- a/searchlib/src/vespa/searchlib/attribute/multi_value_mapping.h
+++ b/searchlib/src/vespa/searchlib/attribute/multi_value_mapping.h
@@ -27,7 +27,8 @@ public:
     MultiValueMapping(const MultiValueMapping &) = delete;
     MultiValueMapping & operator = (const MultiValueMapping &) = delete;
     MultiValueMapping(const vespalib::datastore::ArrayStoreConfig &storeCfg,
-                      const vespalib::GrowStrategy &gs = vespalib::GrowStrategy());
+                      const vespalib::GrowStrategy &gs,
+                      std::shared_ptr<vespalib::alloc::MemoryAllocator> memory_allocator);
     ~MultiValueMapping() override;
     ConstArrayRef get(uint32_t docId) const { return _store.get(_indices[docId]); }
     ConstArrayRef getDataForIdx(EntryRef idx) const { return _store.get(idx); }
diff --git a/searchlib/src/vespa/searchlib/attribute/multi_value_mapping.hpp b/searchlib/src/vespa/searchlib/attribute/multi_value_mapping.hpp
index 339f562757d..16b29bf33cd 100644
--- a/searchlib/src/vespa/searchlib/attribute/multi_value_mapping.hpp
+++ b/searchlib/src/vespa/searchlib/attribute/multi_value_mapping.hpp
@@ -10,16 +10,17 @@ namespace search::attribute {
 
 template <typename EntryT, typename RefT>
 MultiValueMapping<EntryT,RefT>::MultiValueMapping(const vespalib::datastore::ArrayStoreConfig &storeCfg,
-                                                  const vespalib::GrowStrategy &gs)
+                                                  const vespalib::GrowStrategy &gs,
+                                                  std::shared_ptr<vespalib::alloc::MemoryAllocator> memory_allocator)
 #ifdef __clang__
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wuninitialized"
 #endif
-    : MultiValueMappingBase(gs, _store.getGenerationHolder()),
+    : MultiValueMappingBase(gs, _store.getGenerationHolder(), memory_allocator),
 #ifdef __clang__
 #pragma clang diagnostic pop
 #endif
-      _store(storeCfg, {})
+      _store(storeCfg, std::move(memory_allocator))
 {
 }
 
diff --git a/searchlib/src/vespa/searchlib/attribute/multi_value_mapping_base.cpp b/searchlib/src/vespa/searchlib/attribute/multi_value_mapping_base.cpp
index b0d50c129c6..7ad61ccedc5 100644
--- a/searchlib/src/vespa/searchlib/attribute/multi_value_mapping_base.cpp
+++ b/searchlib/src/vespa/searchlib/attribute/multi_value_mapping_base.cpp
@@ -10,8 +10,10 @@ namespace search::attribute {
 using vespalib::datastore::CompactionStrategy;
 
 MultiValueMappingBase::MultiValueMappingBase(const vespalib::GrowStrategy &gs,
-                                             vespalib::GenerationHolder &genHolder)
-    : _indices(gs, genHolder),
+                                             vespalib::GenerationHolder &genHolder,
+                                             std::shared_ptr<vespalib::alloc::MemoryAllocator> memory_allocator)
+    : _memory_allocator(std::move(memory_allocator)),
+      _indices(gs, genHolder, _memory_allocator ? vespalib::alloc::Alloc::alloc_with_allocator(_memory_allocator.get()) : vespalib::alloc::Alloc::alloc()),
       _totalValues(0u),
       _compaction_spec()
 {
diff --git a/searchlib/src/vespa/searchlib/attribute/multi_value_mapping_base.h b/searchlib/src/vespa/searchlib/attribute/multi_value_mapping_base.h
index f27a9f1667c..2b2b4d5f8a3 100644
--- a/searchlib/src/vespa/searchlib/attribute/multi_value_mapping_base.h
+++ b/searchlib/src/vespa/searchlib/attribute/multi_value_mapping_base.h
@@ -27,11 +27,12 @@ public:
     using RefVector = vespalib::RcuVectorBase<EntryRef>;
 
 protected:
+    std::shared_ptr<vespalib::alloc::MemoryAllocator> _memory_allocator;
     RefVector _indices;
     size_t    _totalValues;
     CompactionSpec _compaction_spec;
 
-    MultiValueMappingBase(const vespalib::GrowStrategy &gs, vespalib::GenerationHolder &genHolder);
+    MultiValueMappingBase(const vespalib::GrowStrategy &gs, vespalib::GenerationHolder &genHolder, std::shared_ptr<vespalib::alloc::MemoryAllocator> memory_allocator);
     virtual ~MultiValueMappingBase();
 
     void updateValueCount(size_t oldValues, size_t newValues) {
diff --git a/searchlib/src/vespa/searchlib/attribute/multivalueattribute.hpp b/searchlib/src/vespa/searchlib/attribute/multivalueattribute.hpp
index b9ef16c6adf..c0524a4d043 100644
--- a/searchlib/src/vespa/searchlib/attribute/multivalueattribute.hpp
+++ b/searchlib/src/vespa/searchlib/attribute/multivalueattribute.hpp
@@ -28,7 +28,7 @@ MultiValueAttribute(const vespalib::string &baseFileName,
                                                                8 * 1024,
                                                                cfg.getGrowStrategy().getMultiValueAllocGrowFactor(),
                                                                multivalueattribute::enable_free_lists),
-                 cfg.getGrowStrategy().to_generic_strategy())
+                 cfg.getGrowStrategy().to_generic_strategy(), {})
 {
 }