Add support for setting memory allocator for data store buffer type.

6 files changed, 104 insertions, 0 deletions

diff --git a/vespalib/src/tests/datastore/datastore/datastore_test.cpp b/vespalib/src/tests/datastore/datastore/datastore_test.cpp
index 95ca7b4b1fa..485a595347b 100644
--- a/vespalib/src/tests/datastore/datastore/datastore_test.cpp
+++ b/vespalib/src/tests/datastore/datastore/datastore_test.cpp
@@ -18,6 +18,9 @@ private:
     using ParentType::_activeBufferIds;
 public:
     MyStore() {}
+    explicit MyStore(std::unique_ptr<BufferType<int>> type)
+        : ParentType(std::move(type))
+    {}
     void holdBuffer(uint32_t bufferId) {
         ParentType::holdBuffer(bufferId);
     }
@@ -585,6 +588,89 @@ TEST(DataStoreTest, require_that_offset_in_EntryRefT_is_within_bounds_when_alloc
     assertGrowStats<uint32_t>({16384,16384,16384,32768,32768,65536,131072,131072,229376,229376,229376,229376}, 7);
 }
 
+namespace {
+
+struct AllocStats {
+    size_t alloc_cnt;
+    size_t free_cnt;
+    AllocStats()
+        : AllocStats(0, 0)
+    {
+    }
+    AllocStats(size_t alloc_cnt_in, size_t free_cnt_in)
+        : alloc_cnt(alloc_cnt_in),
+          free_cnt(free_cnt_in)
+    {
+    }
+    bool operator==(const AllocStats &rhs) const {
+        return ((alloc_cnt == rhs.alloc_cnt) &&
+                (free_cnt == rhs.free_cnt));
+    }
+};
+
+std::ostream&
+operator<<(std::ostream &os, const AllocStats &stats)
+{
+    os << "{alloc_cnt=" << stats.alloc_cnt << ", free_cnt=" << stats.free_cnt << "}";
+    return os;
+}
+
+class MyMemoryAllocator : public alloc::MemoryAllocator {
+    AllocStats &_stats;
+    const alloc::MemoryAllocator* _backing_allocator;
+public:
+    MyMemoryAllocator(AllocStats &stats)
+        : alloc::MemoryAllocator(),
+          _stats(stats),
+          _backing_allocator(alloc::MemoryAllocator::select_allocator(HUGEPAGE_SIZE, 0))
+    {
+    }
+    ~MyMemoryAllocator() override = default;
+
+    PtrAndSize alloc(size_t sz) const override { ++_stats.alloc_cnt; return _backing_allocator->alloc(sz); }
+    void free(PtrAndSize alloc) const override { ++_stats.free_cnt; _backing_allocator->free(alloc); }
+    size_t resize_inplace(PtrAndSize current, size_t newSize) const override { return _backing_allocator->resize_inplace(current, newSize); }
+};
+
+class MyBufferType : public BufferType<int>
+{
+    std::unique_ptr<alloc::MemoryAllocator> _allocator;
+public:
+    MyBufferType(std::unique_ptr<alloc::MemoryAllocator> allocator, uint32_t max_arrays)
+        : BufferType<int>(1, 2, max_arrays, max_arrays, 0.2),
+          _allocator(std::move(allocator))
+    {
+    }
+    const alloc::MemoryAllocator* get_memory_allocator() const override {
+        return _allocator.get();
+    }
+};
+
+}
+
+TEST(DataStoreTest, can_set_memory_allocator)
+{
+    AllocStats stats;
+    {
+        MyStore s(std::make_unique<MyBufferType>(std::make_unique<MyMemoryAllocator>(stats), MyStore::RefType::offsetSize()));
+        EXPECT_EQ(AllocStats(1, 0), stats);
+        auto ref = s.addEntry(42);
+        EXPECT_EQ(0u, MyRef(ref).bufferId());
+        EXPECT_EQ(AllocStats(1, 0), stats);
+        auto ref2 = s.addEntry(43);
+        EXPECT_EQ(0u, MyRef(ref2).bufferId());
+        EXPECT_EQ(AllocStats(2, 0), stats);
+        s.switchActiveBuffer();
+        EXPECT_EQ(AllocStats(3, 0), stats);
+        s.holdBuffer(0);
+        s.transferHoldLists(10);
+        EXPECT_EQ(AllocStats(3, 0), stats);
+        s.trimHoldLists(11);
+        EXPECT_EQ(AllocStats(3, 2), stats);
+    }
+    EXPECT_EQ(AllocStats(3, 3), stats);
+}
+
 TEST(DataStoreTest, control_static_sizes) {
     EXPECT_EQ(64, sizeof(BufferTypeBase));
     EXPECT_EQ(32, sizeof(BufferState::FreeList));
diff --git a/vespalib/src/vespa/vespalib/datastore/buffer_type.cpp b/vespalib/src/vespa/vespalib/datastore/buffer_type.cpp
index 848be2196ed..f25ef2c1e52 100644
--- a/vespalib/src/vespa/vespalib/datastore/buffer_type.cpp
+++ b/vespalib/src/vespa/vespalib/datastore/buffer_type.cpp
@@ -105,6 +105,12 @@ BufferTypeBase::onFree(ElemCount usedElems)
     _holdUsedElems -= usedElems;
 }
 
+const alloc::MemoryAllocator*
+BufferTypeBase::get_memory_allocator() const
+{
+    return nullptr;
+}
+
 void
 BufferTypeBase::clampMaxArrays(uint32_t maxArrays)
 {
diff --git a/vespalib/src/vespa/vespalib/datastore/buffer_type.h b/vespalib/src/vespa/vespalib/datastore/buffer_type.h
index 2a27124921d..3ec0ad55753 100644
--- a/vespalib/src/vespa/vespalib/datastore/buffer_type.h
+++ b/vespalib/src/vespa/vespalib/datastore/buffer_type.h
@@ -5,6 +5,8 @@
 #include "atomic_entry_ref.h"
 #include <string>
 
+namespace vespalib::alloc { class MemoryAllocator; }
+
 namespace vespalib::datastore {
 
 using ElemCount = uint64_t;
@@ -53,6 +55,7 @@ public:
     virtual void onActive(uint32_t bufferId, ElemCount *usedElems, ElemCount &deadElems, void *buffer);
     void onHold(const ElemCount *usedElems);
     virtual void onFree(ElemCount usedElems);
+    virtual const alloc::MemoryAllocator* get_memory_allocator() const;
 
     /**
      * Calculate number of arrays to allocate for new buffer given how many elements are needed.
diff --git a/vespalib/src/vespa/vespalib/datastore/bufferstate.cpp b/vespalib/src/vespa/vespalib/datastore/bufferstate.cpp
index 2edd251452a..19b7cdf8181 100644
--- a/vespalib/src/vespa/vespalib/datastore/bufferstate.cpp
+++ b/vespalib/src/vespa/vespalib/datastore/bufferstate.cpp
@@ -116,6 +116,8 @@ BufferState::onActive(uint32_t bufferId, uint32_t typeId,
     (void) reservedElements;
     AllocResult alloc = calcAllocation(bufferId, *typeHandler, elementsNeeded, false);
     assert(alloc.elements >= reservedElements + elementsNeeded);
+    auto allocator = typeHandler->get_memory_allocator();
+    _buffer = (allocator != nullptr) ? Alloc::alloc_with_allocator(allocator) : Alloc::alloc(0, MemoryAllocator::HUGEPAGE_SIZE);
     _buffer.create(alloc.bytes).swap(_buffer);
     buffer = _buffer.get();
     assert(buffer != nullptr || alloc.elements == 0u);
diff --git a/vespalib/src/vespa/vespalib/util/alloc.cpp b/vespalib/src/vespa/vespalib/util/alloc.cpp
index 745fb50db03..ecd0bcfd5e6 100644
--- a/vespalib/src/vespa/vespalib/util/alloc.cpp
+++ b/vespalib/src/vespa/vespalib/util/alloc.cpp
@@ -522,6 +522,12 @@ Alloc::alloc(size_t sz, size_t mmapLimit, size_t alignment) noexcept
     return Alloc(&AutoAllocator::getAllocator(mmapLimit, alignment), sz);
 }
 
+Alloc
+Alloc::alloc_with_allocator(const MemoryAllocator* allocator) noexcept
+{
+    return Alloc(allocator);
+}
+
 }
 
 }
diff --git a/vespalib/src/vespa/vespalib/util/alloc.h b/vespalib/src/vespa/vespalib/util/alloc.h
index cf7135736e5..c46915103cf 100644
--- a/vespalib/src/vespa/vespalib/util/alloc.h
+++ b/vespalib/src/vespa/vespalib/util/alloc.h
@@ -105,6 +105,7 @@ public:
      */
     static Alloc alloc(size_t sz, size_t mmapLimit = MemoryAllocator::HUGEPAGE_SIZE, size_t alignment=0) noexcept;
     static Alloc alloc() noexcept;
+    static Alloc alloc_with_allocator(const MemoryAllocator* allocator) noexcept;
 private:
     Alloc(const MemoryAllocator * allocator, size_t sz) noexcept : _alloc(allocator->alloc(sz)), _allocator(allocator) { }
     Alloc(const MemoryAllocator * allocator) noexcept : _alloc(nullptr, 0), _allocator(allocator) { }