Merge pull request #27644 from vespa-engine/toregge/use-provided-memory-allocator-for-large-arrays

Use provided memory allocator for large arrays.

3 files changed, 15 insertions, 2 deletions

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 6e433e48d88..6a8f6175032 100644
--- a/vespalib/src/tests/datastore/array_store/array_store_test.cpp
+++ b/vespalib/src/tests/datastore/array_store/array_store_test.cpp
@@ -595,6 +595,16 @@ TEST_F(ByteStoreTest, offset_in_EntryRefT_is_within_bounds_when_allocating_memor
 TYPED_TEST(NumberStoreTest, provided_memory_allocator_is_used)
 {
     EXPECT_EQ(AllocStats(4, 0), this->stats);
+    this->assertAdd({1,2,3,4,5});
+    EXPECT_EQ(AllocStats(5, 0), this->stats);
+    this->assertAdd({2,3,4,5,6,7});
+    EXPECT_EQ(AllocStats(6, 0), this->stats);
+    this->remove({1,2,3,4,5});
+    this->remove({2,3,4,5,6,7});
+    this->reclaim_memory();
+    EXPECT_EQ(AllocStats(6, 2), this->stats);
+    this->assertAdd({1,2,3,4,5});
+    EXPECT_EQ(AllocStats(7, 2), this->stats);
 }
 
 GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/vespalib/src/vespa/vespalib/datastore/array_store.hpp b/vespalib/src/vespa/vespalib/datastore/array_store.hpp
index bfd4ff0430a..08bd63942d7 100644
--- a/vespalib/src/vespa/vespalib/datastore/array_store.hpp
+++ b/vespalib/src/vespa/vespalib/datastore/array_store.hpp
@@ -168,7 +168,7 @@ ArrayStore<ElemT, RefT, TypeMapperT>::addLargeArray(ConstArrayRef array)
 {
     using NoOpReclaimer = DefaultReclaimer<LargeArray>;
     auto handle = _store.template freeListAllocator<LargeArray, NoOpReclaimer>(_largeArrayTypeId)
-            .alloc(array.cbegin(), array.cend());
+            .alloc(array.cbegin(), array.cend(), _largeArrayType.initial_alloc());
     auto& state = _store.getBufferState(RefT(handle.ref).bufferId());
     state.stats().inc_extra_used_bytes(sizeof(ElemT) * array.size());
     return handle.ref;
@@ -179,7 +179,7 @@ EntryRef
 ArrayStore<ElemT, RefT, TypeMapperT>::allocate_large_array(size_t array_size)
 {
     using NoOpReclaimer = DefaultReclaimer<LargeArray>;
-    auto handle = _store.template freeListAllocator<LargeArray, NoOpReclaimer>(_largeArrayTypeId).alloc(array_size);
+    auto handle = _store.template freeListAllocator<LargeArray, NoOpReclaimer>(_largeArrayTypeId).alloc(array_size, _largeArrayType.initial_alloc());
     auto& state = _store.getBufferState(RefT(handle.ref).bufferId());
     state.stats().inc_extra_used_bytes(sizeof(ElemT) * array_size);
     return handle.ref;
diff --git a/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.h b/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.h
index e2718b94cd2..cf472b7a642 100644
--- a/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.h
+++ b/vespalib/src/vespa/vespalib/datastore/large_array_buffer_type.h
@@ -33,6 +33,9 @@ public:
     ~LargeArrayBufferType() override;
     void clean_hold(void* buffer, size_t offset, EntryCount num_entries, CleanContext cleanCtx) override;
     const vespalib::alloc::MemoryAllocator* get_memory_allocator() const override;
+    alloc::Alloc initial_alloc() const noexcept {
+        return _memory_allocator ? alloc::Alloc::alloc_with_allocator(_memory_allocator.get()) : alloc::Alloc::alloc(0, alloc::MemoryAllocator::HUGEPAGE_SIZE);
+    }
 };
 
 extern template class LargeArrayBufferType<uint8_t>;