Track heap allocated memory for tensors in DirectTensorStore.

2 files changed, 10 insertions, 7 deletions

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();
 }