Merge pull request #14239 from vespa-engine/geirst/memory-usage-tracking-in-direct-tensor-store

Memory usage tracking in direct tensor store

4 files changed, 51 insertions, 40 deletions

diff --git a/vespalib/src/tests/datastore/datastore/datastore_test.cpp b/vespalib/src/tests/datastore/datastore/datastore_test.cpp
index 00da8835c06..448df2d24bd 100644
--- a/vespalib/src/tests/datastore/datastore/datastore_test.cpp
+++ b/vespalib/src/tests/datastore/datastore/datastore_test.cpp
@@ -23,25 +23,15 @@ private:
     using ParentType::_activeBufferIds;
 public:
     MyStore() {}
-
-    void
-    holdBuffer(uint32_t bufferId)
-    {
+    void holdBuffer(uint32_t bufferId) {
         ParentType::holdBuffer(bufferId);
     }
-
-    void
-    holdElem(EntryRef ref, uint64_t len)
-    {
+    void holdElem(EntryRef ref, uint64_t len) {
         ParentType::holdElem(ref, len);
     }
-
-    void
-    transferHoldLists(generation_t generation)
-    {
+    void transferHoldLists(generation_t generation) {
         ParentType::transferHoldLists(generation);
     }
-
     void trimElemHoldList(generation_t usedGen) override {
         ParentType::trimElemHoldList(usedGen);
     }
@@ -54,13 +44,13 @@ public:
     void enableFreeLists() {
         ParentType::enableFreeLists();
     }
-
-    void
-    switchActiveBuffer()
-    {
+    void switchActiveBuffer() {
         ParentType::switchActiveBuffer(0, 0u);
     }
     size_t activeBufferId() const { return _activeBufferIds[0]; }
+    BufferState& get_active_buffer_state() {
+        return ParentType::getBufferState(activeBufferId());
+    }
 };
 
 
@@ -454,6 +444,21 @@ TEST(DataStoreTest, require_that_memory_stats_are_calculated)
     m._freeBuffers = MyRef::numBuffers() - 1;
     m._holdBuffers = 0;
     assertMemStats(m, s.getMemStats());
+
+    { // increase extra used bytes
+        auto prev_stats = s.getMemStats();
+        s.get_active_buffer_state().incExtraUsedBytes(50);
+        auto curr_stats = s.getMemStats();
+        EXPECT_EQ(prev_stats._allocBytes + 50, curr_stats._allocBytes);
+        EXPECT_EQ(prev_stats._usedBytes + 50, curr_stats._usedBytes);
+    }
+
+    { // increase extra hold bytes
+        auto prev_stats = s.getMemStats();
+        s.get_active_buffer_state().incExtraHoldBytes(30);
+        auto curr_stats = s.getMemStats();
+        EXPECT_EQ(prev_stats._holdBytes + 30, curr_stats._holdBytes);
+    }
 }
 
 TEST(DataStoreTest, require_that_memory_usage_is_calculated)
diff --git a/vespalib/src/vespa/vespalib/datastore/datastore.h b/vespalib/src/vespa/vespalib/datastore/datastore.h
index f437616af2a..e67d9049f0b 100644
--- a/vespalib/src/vespa/vespalib/datastore/datastore.h
+++ b/vespalib/src/vespa/vespalib/datastore/datastore.h
@@ -98,14 +98,17 @@ protected:
     using ParentType::dropBuffers;
     using ParentType::initActiveBuffers;
     using ParentType::addType;
+    using BufferTypeUP = std::unique_ptr<BufferType<EntryType>>;
+
+    BufferTypeUP _type;
+
 
-    BufferType<EntryType> _type;
 public:
-    typedef typename ParentType::RefType RefType;
+    using RefType = typename ParentType::RefType;
     DataStore(const DataStore &rhs) = delete;
     DataStore &operator=(const DataStore &rhs) = delete;
     DataStore();
-    DataStore(uint32_t min_arrays);
+    DataStore(BufferTypeUP type);
     ~DataStore();
 
     EntryRef addEntry(const EntryType &e);
diff --git a/vespalib/src/vespa/vespalib/datastore/datastore.hpp b/vespalib/src/vespa/vespalib/datastore/datastore.hpp
index 71cc52a6838..b66a3b78603 100644
--- a/vespalib/src/vespa/vespalib/datastore/datastore.hpp
+++ b/vespalib/src/vespa/vespalib/datastore/datastore.hpp
@@ -133,16 +133,16 @@ DataStoreT<RefT>::freeListRawAllocator(uint32_t typeId)
 
 template <typename EntryType, typename RefT>
 DataStore<EntryType, RefT>::DataStore()
-    : DataStore(RefType::offsetSize())
+    : DataStore(std::make_unique<BufferType<EntryType>>(1, RefType::offsetSize(), RefType::offsetSize()))
 {
 }
 
 template <typename EntryType, typename RefT>
-DataStore<EntryType, RefT>::DataStore(uint32_t min_arrays)
+DataStore<EntryType, RefT>::DataStore(BufferTypeUP type)
     : ParentType(),
-      _type(1, min_arrays, RefType::offsetSize())
+      _type(std::move(type))
 {
-    addType(&_type);
+    addType(_type.get());
     initActiveBuffers();
 }
 
diff --git a/vespalib/src/vespa/vespalib/datastore/datastorebase.cpp b/vespalib/src/vespa/vespalib/datastore/datastorebase.cpp
index 4e68b10d76a..15135af67b5 100644
--- a/vespalib/src/vespa/vespalib/datastore/datastorebase.cpp
+++ b/vespalib/src/vespa/vespalib/datastore/datastorebase.cpp
@@ -271,6 +271,23 @@ DataStoreBase::disableElemHoldList()
     }
 }
 
+namespace {
+
+void
+add_buffer_state_to_mem_stats(const BufferState& state, size_t elementSize, DataStoreBase::MemStats& stats)
+{
+    size_t extra_used_bytes = state.getExtraUsedBytes();
+    stats._allocElems += state.capacity();
+    stats._usedElems += state.size();
+    stats._deadElems += state.getDeadElems();
+    stats._holdElems += state.getHoldElems();
+    stats._allocBytes += (state.capacity() * elementSize) + extra_used_bytes;
+    stats._usedBytes += (state.size() * elementSize) + extra_used_bytes;
+    stats._deadBytes += state.getDeadElems() * elementSize;
+    stats._holdBytes += (state.getHoldElems() * elementSize) + state.getExtraHoldBytes();
+}
+
+}
 
 DataStoreBase::MemStats
 DataStoreBase::getMemStats() const
@@ -285,25 +302,11 @@ DataStoreBase::getMemStats() const
         } else if (state == BufferState::ACTIVE) {
             size_t elementSize = typeHandler->elementSize();
             ++stats._activeBuffers;
-            stats._allocElems += bState.capacity();
-            stats._usedElems += bState.size();
-            stats._deadElems += bState.getDeadElems();
-            stats._holdElems += bState.getHoldElems();
-            stats._allocBytes += bState.capacity() * elementSize;
-            stats._usedBytes += (bState.size() * elementSize) + bState.getExtraUsedBytes();
-            stats._deadBytes += bState.getDeadElems() * elementSize;
-            stats._holdBytes += (bState.getHoldElems() * elementSize) + bState.getExtraHoldBytes();
+            add_buffer_state_to_mem_stats(bState, elementSize, stats);
         } else if (state == BufferState::HOLD) {
             size_t elementSize = typeHandler->elementSize();
             ++stats._holdBuffers;
-            stats._allocElems += bState.capacity();
-            stats._usedElems += bState.size();
-            stats._deadElems += bState.getDeadElems();
-            stats._holdElems += bState.getHoldElems();
-            stats._allocBytes += bState.capacity() * elementSize;
-            stats._usedBytes += (bState.size() * elementSize) + bState.getExtraUsedBytes();
-            stats._deadBytes += bState.getDeadElems() * elementSize;
-            stats._holdBytes += (bState.getHoldElems() * elementSize) + bState.getExtraHoldBytes();
+            add_buffer_state_to_mem_stats(bState, elementSize, stats);
         } else {
             LOG_ABORT("should not be reached");
         }