Report nearest neighbor index memory usage as part of tensor attribute memory usage.

7 files changed, 41 insertions, 5 deletions

diff --git a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
index fef063fb1cf..e6eb01596a8 100644
--- a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
+++ b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
@@ -84,6 +84,7 @@ private:
     EntryVector _removes;
     generation_t _transfer_gen;
     generation_t _trim_gen;
+    mutable size_t _memory_usage_cnt;
 
 public:
     MockNearestNeighborIndex(const DocVectorAccess& vectors)
@@ -91,7 +92,8 @@ public:
           _adds(),
           _removes(),
           _transfer_gen(std::numeric_limits<generation_t>::max()),
-          _trim_gen(std::numeric_limits<generation_t>::max())
+          _trim_gen(std::numeric_limits<generation_t>::max()),
+          _memory_usage_cnt(0)
     {
     }
     void clear() {
@@ -119,6 +121,7 @@ public:
     }
     generation_t get_transfer_gen() const { return _transfer_gen; }
     generation_t get_trim_gen() const { return _trim_gen; }
+    size_t memory_usage_cnt() const { return _memory_usage_cnt; }
 
     void add_document(uint32_t docid) override {
         auto vector = _vectors.get_vector(docid).typify<double>();
@@ -135,6 +138,7 @@ public:
         _trim_gen = first_used_gen;
     }
     vespalib::MemoryUsage memory_usage() const override {
+        ++_memory_usage_cnt;
         return vespalib::MemoryUsage();
     }
     std::vector<Neighbor> find_top_k(uint32_t k, vespalib::tensor::TypedCells vector, uint32_t explore_k) const override {
@@ -582,4 +586,12 @@ TEST_F("commit() ensures transfer and trim hold lists on nearest neighbor index"
     EXPECT_EQUAL(gen_3, index.get_trim_gen());
 }
 
+TEST_F("Memory usage is extracted from index when updating stats on attribute", DenseTensorAttributeMockIndex)
+{
+    size_t before = f.mock_index().memory_usage_cnt();
+    f.getStatus();
+    size_t after = f.mock_index().memory_usage_cnt();
+    EXPECT_EQUAL(before + 1, after);
+}
+
 TEST_MAIN() { TEST_RUN_ALL(); vespalib::unlink("test.dat"); }
diff --git a/searchlib/src/vespa/searchlib/attribute/attributevector.h b/searchlib/src/vespa/searchlib/attribute/attributevector.h
index b5f6beaa718..a396fb70b7c 100644
--- a/searchlib/src/vespa/searchlib/attribute/attributevector.h
+++ b/searchlib/src/vespa/searchlib/attribute/attributevector.h
@@ -320,6 +320,10 @@ protected:
         return _genHolder;
     }
 
+    const GenerationHolder& getGenerationHolder() const {
+        return _genHolder;
+    }
+
     template<typename T>
     bool clearDoc(ChangeVectorT< ChangeTemplate<T> > &changes, DocId doc);
 
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
index b023df19c5b..9f127018a92 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
@@ -64,6 +64,16 @@ DenseTensorAttribute::consider_remove_from_index(DocId docid)
     }
 }
 
+vespalib::MemoryUsage
+DenseTensorAttribute::memory_usage() const
+{
+    vespalib::MemoryUsage result = TensorAttribute::memory_usage();
+    if (_index) {
+        result.merge(_index->memory_usage());
+    }
+    return result;
+}
+
 DenseTensorAttribute::DenseTensorAttribute(vespalib::stringref baseFileName, const Config& cfg,
                                            const NearestNeighborIndexFactory& index_factory)
     : TensorAttribute(baseFileName, cfg, _denseTensorStore),
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
index 84a60376fe7..8eba8354fbd 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
@@ -24,6 +24,7 @@ private:
     std::unique_ptr<NearestNeighborIndex> _index;
 
     void consider_remove_from_index(DocId docid);
+    vespalib::MemoryUsage memory_usage() const override;
 
 public:
     DenseTensorAttribute(vespalib::stringref baseFileName, const Config& cfg,
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
index 9795330b061..6d48ea1a967 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
@@ -420,6 +420,7 @@ HnswIndex::memory_usage() const
     result.merge(_node_refs.getMemoryUsage());
     result.merge(_nodes.getMemoryUsage());
     result.merge(_links.getMemoryUsage());
+    result.merge(_visited_set_pool.memory_usage());
     return result;
 }
 
diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp
index 0b9628e6872..acca4c3d8f1 100644
--- a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp
@@ -111,10 +111,7 @@ TensorAttribute::onCommit()
 void
 TensorAttribute::onUpdateStat()
 {
-    // update statistics
-    vespalib::MemoryUsage total = _refVector.getMemoryUsage();
-    total.merge(_tensorStore.getMemoryUsage());
-    total.mergeGenerationHeldBytes(getGenerationHolder().getHeldBytes());
+    vespalib::MemoryUsage total = memory_usage();
     this->updateStatistics(_refVector.size(),
                            _refVector.size(),
                            total.allocatedBytes(),
@@ -178,6 +175,15 @@ TensorAttribute::setTensorRef(DocId docId, EntryRef ref)
     }
 }
 
+vespalib::MemoryUsage
+TensorAttribute::memory_usage() const
+{
+    vespalib::MemoryUsage result = _refVector.getMemoryUsage();
+    result.merge(_tensorStore.getMemoryUsage());
+    result.mergeGenerationHeldBytes(getGenerationHolder().getHeldBytes());
+    return result;
+}
+
 Tensor::UP
 TensorAttribute::getEmptyTensor() const
 {
diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h
index 64f978a31d7..5e65317baa2 100644
--- a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h
@@ -27,6 +27,8 @@ protected:
     void doCompactWorst();
     void checkTensorType(const Tensor &tensor);
     void setTensorRef(DocId docId, EntryRef ref);
+    virtual vespalib::MemoryUsage memory_usage() const;
+
 public:
     DECLARE_IDENTIFIABLE_ABSTRACT(TensorAttribute);
     using RefCopyVector = vespalib::Array<EntryRef>;