Expose information from tensor attribute and hnsw index in state explorer for attribute.

14 files changed, 114 insertions, 11 deletions

diff --git a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
index e6eb01596a8..97d5cb61327 100644
--- a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
+++ b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
@@ -141,6 +141,7 @@ public:
         ++_memory_usage_cnt;
         return vespalib::MemoryUsage();
     }
+    void get_state(const vespalib::slime::Inserter&) const override {}
     std::vector<Neighbor> find_top_k(uint32_t k, vespalib::tensor::TypedCells vector, uint32_t explore_k) const override {
         (void) k;
         (void) vector;
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
index 9f127018a92..627f7f0dfa9 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
@@ -4,15 +4,17 @@
 #include "dense_tensor_attribute_saver.h"
 #include "nearest_neighbor_index.h"
 #include "tensor_attribute.hpp"
-#include <vespa/eval/tensor/tensor.h>
 #include <vespa/eval/tensor/dense/mutable_dense_tensor_view.h>
+#include <vespa/eval/tensor/tensor.h>
 #include <vespa/fastlib/io/bufferedfile.h>
 #include <vespa/searchlib/attribute/readerbase.h>
+#include <vespa/vespalib/data/slime/inserter.h>
 
 #include <vespa/log/log.h>
 LOG_SETUP(".searchlib.tensor.dense_tensor_attribute");
 
 using vespalib::eval::ValueType;
+using vespalib::slime::ObjectInserter;
 using vespalib::tensor::MutableDenseTensorView;
 using vespalib::tensor::Tensor;
 
@@ -217,6 +219,17 @@ DenseTensorAttribute::removeOldGenerations(generation_t first_used_gen)
     }
 }
 
+void
+DenseTensorAttribute::get_state(const vespalib::slime::Inserter& inserter) const
+{
+    auto& object = inserter.insertObject();
+    populate_state(object);
+    if (_index) {
+        ObjectInserter index_inserter(object, "nearest_neighbor_index");
+        _index->get_state(index_inserter);
+    }
+}
+
 vespalib::tensor::TypedCells
 DenseTensorAttribute::get_vector(uint32_t docid) const
 {
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
index 8eba8354fbd..f0383627ea2 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
@@ -41,6 +41,7 @@ public:
     uint32_t getVersion() const override;
     void onGenerationChange(generation_t next_gen) override;
     void removeOldGenerations(generation_t first_used_gen) override;
+    void get_state(const vespalib::slime::Inserter& inserter) const override;
 
     // Implements DocVectorAccess
     vespalib::tensor::TypedCells get_vector(uint32_t docid) const override;
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
index 6d48ea1a967..988264c0455 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.cpp
@@ -3,12 +3,16 @@
 #include "distance_function.h"
 #include "hnsw_index.h"
 #include "random_level_generator.h"
+#include <vespa/searchlib/util/state_explorer_utils.h>
 #include <vespa/eval/tensor/dense/typed_cells.h>
+#include <vespa/vespalib/data/slime/cursor.h>
+#include <vespa/vespalib/data/slime/inserter.h>
 #include <vespa/vespalib/datastore/array_store.hpp>
 #include <vespa/vespalib/util/rcuvector.hpp>
 
 namespace search::tensor {
 
+using search::StateExplorerUtils;
 using search::datastore::EntryRef;
 
 namespace {
@@ -424,6 +428,13 @@ HnswIndex::memory_usage() const
     return result;
 }
 
+void
+HnswIndex::get_state(const vespalib::slime::Inserter& inserter) const
+{
+    auto& object = inserter.insertObject();
+    StateExplorerUtils::memory_usage_to_slime(memory_usage(), object.setObject("memory_usage"));
+}
+
 struct NeighborsByDocId {
     bool operator() (const NearestNeighborIndex::Neighbor &lhs,
                      const NearestNeighborIndex::Neighbor &rhs)
diff --git a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h
index 7c0d650ab2f..8da48a21af3 100644
--- a/searchlib/src/vespa/searchlib/tensor/hnsw_index.h
+++ b/searchlib/src/vespa/searchlib/tensor/hnsw_index.h
@@ -153,6 +153,8 @@ public:
     void transfer_hold_lists(generation_t current_gen) override;
     void trim_hold_lists(generation_t first_used_gen) override;
     vespalib::MemoryUsage memory_usage() const override;
+    void get_state(const vespalib::slime::Inserter& inserter) const override;
+
     std::vector<Neighbor> find_top_k(uint32_t k, TypedCells vector, uint32_t explore_k) const override;
 
     FurthestPriQ top_k_candidates(const TypedCells &vector, uint32_t k) const;
diff --git a/searchlib/src/vespa/searchlib/tensor/i_tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/i_tensor_attribute.h
index 6c83d3caae9..520abb13d06 100644
--- a/searchlib/src/vespa/searchlib/tensor/i_tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/i_tensor_attribute.h
@@ -9,6 +9,7 @@ class MutableDenseTensorView;
 class Tensor;
 }
 namespace vespalib::eval { class ValueType; }
+namespace vespalib::slime { struct Inserter; }
 
 namespace search::tensor {
 
@@ -25,6 +26,12 @@ public:
     virtual std::unique_ptr<Tensor> getEmptyTensor() const = 0;
     virtual void getTensor(uint32_t docId, vespalib::tensor::MutableDenseTensorView &tensor) const = 0;
     virtual vespalib::eval::ValueType getTensorType() const = 0;
+
+    /**
+     * Gets custom state for this tensor attribute by inserting it into the given Slime inserter.
+     * This function is only called by the writer thread or when the writer thread is blocked.
+     */
+    virtual void get_state(const vespalib::slime::Inserter& inserter) const = 0;
 };
 
 }  // namespace search::tensor
diff --git a/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.cpp b/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.cpp
index 8b60f0c3e8e..593330b7bd4 100644
--- a/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.cpp
@@ -62,4 +62,10 @@ ImportedTensorAttributeVectorReadGuard::getTensorType() const
     return _target_tensor_attribute.getTensorType();
 }
 
+void
+ImportedTensorAttributeVectorReadGuard::get_state(const vespalib::slime::Inserter& inserter) const
+{
+    _target_tensor_attribute.get_state(inserter);
+}
+
 }
diff --git a/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.h b/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.h
index fce744503cd..74c5cd9a685 100644
--- a/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.h
+++ b/searchlib/src/vespa/searchlib/tensor/imported_tensor_attribute_vector_read_guard.h
@@ -34,6 +34,7 @@ public:
     virtual std::unique_ptr<Tensor> getEmptyTensor() const override;
     virtual void getTensor(uint32_t docId, vespalib::tensor::MutableDenseTensorView &tensor) const override;
     virtual vespalib::eval::ValueType getTensorType() const override;
+    virtual void get_state(const vespalib::slime::Inserter& inserter) const override;
 };
 
 }
diff --git a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h
index bd98623bdd3..ccd7e614e06 100644
--- a/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h
+++ b/searchlib/src/vespa/searchlib/tensor/nearest_neighbor_index.h
@@ -8,6 +8,8 @@
 #include <vespa/vespalib/util/generationhandler.h>
 #include <vespa/vespalib/util/memoryusage.h>
 
+namespace vespalib::slime { struct Inserter; }
+
 namespace search::tensor {
 
 /**
@@ -30,6 +32,7 @@ public:
     virtual void transfer_hold_lists(generation_t current_gen) = 0;
     virtual void trim_hold_lists(generation_t first_used_gen) = 0;
     virtual vespalib::MemoryUsage memory_usage() const = 0;
+    virtual void get_state(const vespalib::slime::Inserter& inserter) const = 0;
 
     virtual std::vector<Neighbor> find_top_k(uint32_t k,
                                              vespalib::tensor::TypedCells vector,
diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp
index acca4c3d8f1..95af9f0471b 100644
--- a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp
@@ -7,21 +7,23 @@
 #include <vespa/eval/tensor/dense/typed_dense_tensor_builder.h>
 #include <vespa/eval/tensor/sparse/sparse_tensor.h>
 #include <vespa/eval/tensor/wrapped_simple_tensor.h>
+#include <vespa/searchlib/util/state_explorer_utils.h>
+#include <vespa/vespalib/data/slime/cursor.h>
+#include <vespa/vespalib/data/slime/inserter.h>
 #include <vespa/vespalib/util/rcuvector.hpp>
 
+using document::TensorDataType;
+using document::WrongTensorTypeException;
 using vespalib::eval::SimpleTensor;
 using vespalib::eval::ValueType;
+using vespalib::tensor::SparseTensor;
 using vespalib::tensor::Tensor;
 using vespalib::tensor::TypedDenseTensorBuilder;
-using vespalib::tensor::dispatch_0;
-using vespalib::tensor::SparseTensor;
 using vespalib::tensor::WrappedSimpleTensor;
-using document::TensorDataType;
-using document::WrongTensorTypeException;
-
-namespace search {
+using vespalib::tensor::dispatch_0;
+using search::StateExplorerUtils;
 
-namespace tensor {
+namespace search::tensor {
 
 namespace {
 
@@ -184,6 +186,16 @@ TensorAttribute::memory_usage() const
     return result;
 }
 
+void
+TensorAttribute::populate_state(vespalib::slime::Cursor& object) const
+{
+    object.setLong("compact_generation", _compactGeneration);
+    StateExplorerUtils::memory_usage_to_slime(_refVector.getMemoryUsage(),
+                                              object.setObject("ref_vector").setObject("memory_usage"));
+    StateExplorerUtils::memory_usage_to_slime(_tensorStore.getMemoryUsage(),
+                                              object.setObject("tensor_store").setObject("memory_usage"));
+}
+
 Tensor::UP
 TensorAttribute::getEmptyTensor() const
 {
@@ -197,6 +209,13 @@ TensorAttribute::getTensorType() const
 }
 
 void
+TensorAttribute::get_state(const vespalib::slime::Inserter& inserter) const
+{
+    auto& object = inserter.insertObject();
+    populate_state(object);
+}
+
+void
 TensorAttribute::clearDocs(DocId lidLow, DocId lidLimit)
 {
     assert(lidLow <= lidLimit);
@@ -236,6 +255,4 @@ TensorAttribute::getRefCopy() const
 
 IMPLEMENT_IDENTIFIABLE_ABSTRACT(TensorAttribute, AttributeVector);
 
-}  // namespace search::tensor
-
-}  // namespace search
+}
diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h
index 5e65317baa2..e8efd2170c9 100644
--- a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h
@@ -28,6 +28,7 @@ protected:
     void checkTensorType(const Tensor &tensor);
     void setTensorRef(DocId docId, EntryRef ref);
     virtual vespalib::MemoryUsage memory_usage() const;
+    void populate_state(vespalib::slime::Cursor& object) const;
 
 public:
     DECLARE_IDENTIFIABLE_ABSTRACT(TensorAttribute);
@@ -44,6 +45,7 @@ public:
     bool addDoc(DocId &docId) override;
     std::unique_ptr<Tensor> getEmptyTensor() const override;
     vespalib::eval::ValueType getTensorType() const override;
+    void get_state(const vespalib::slime::Inserter& inserter) const override;
     void clearDocs(DocId lidLow, DocId lidLimit) override;
     void onShrinkLidSpace() override;
     uint32_t getVersion() const override;
diff --git a/searchlib/src/vespa/searchlib/util/CMakeLists.txt b/searchlib/src/vespa/searchlib/util/CMakeLists.txt
index 873f1824d04..ac5fc5f54fc 100644
--- a/searchlib/src/vespa/searchlib/util/CMakeLists.txt
+++ b/searchlib/src/vespa/searchlib/util/CMakeLists.txt
@@ -16,6 +16,7 @@ vespa_add_library(searchlib_util OBJECT
     rawbuf.cpp
     sigbushandler.cpp
     slime_output_raw_buf_adapter.cpp
+    state_explorer_utils.cpp
     statebuf.cpp
     statefile.cpp
     stringenum.cpp
diff --git a/searchlib/src/vespa/searchlib/util/state_explorer_utils.cpp b/searchlib/src/vespa/searchlib/util/state_explorer_utils.cpp
new file mode 100644
index 00000000000..307fc91a605
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/util/state_explorer_utils.cpp
@@ -0,0 +1,19 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "state_explorer_utils.h"
+#include <vespa/vespalib/data/slime/cursor.h>
+#include <vespa/vespalib/util/memoryusage.h>
+
+namespace search {
+
+void
+StateExplorerUtils::memory_usage_to_slime(const vespalib::MemoryUsage& usage, vespalib::slime::Cursor& object)
+{
+    object.setLong("allocated", usage.allocatedBytes());
+    object.setLong("used", usage.usedBytes());
+    object.setLong("dead", usage.deadBytes());
+    object.setLong("onHold", usage.allocatedBytesOnHold());
+}
+
+}
+
diff --git a/searchlib/src/vespa/searchlib/util/state_explorer_utils.h b/searchlib/src/vespa/searchlib/util/state_explorer_utils.h
new file mode 100644
index 00000000000..58867b426a1
--- /dev/null
+++ b/searchlib/src/vespa/searchlib/util/state_explorer_utils.h
@@ -0,0 +1,19 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+namespace vespalib { class MemoryUsage; }
+namespace vespalib::slime { struct Cursor; }
+
+namespace search {
+
+/**
+ * Utility functions for state explorers to convert objects to slime.
+ */
+class StateExplorerUtils {
+public:
+    static void memory_usage_to_slime(const vespalib::MemoryUsage& usage, vespalib::slime::Cursor& object);
+};
+
+}
+