From 3def15314e3f3a1fe5dcdd9815c6f5150e1b0397 Mon Sep 17 00:00:00 2001
From: Tor Egge <Tor.Egge@online.no>
Date: Tue, 1 Jun 2021 14:56:21 +0200
Subject: Only use tensor store memory usage when considering compaction.

---
 .../searchlib/tensor/dense_tensor_attribute.cpp      | 10 ++++++++++
 .../vespa/searchlib/tensor/dense_tensor_attribute.h  |  1 +
 .../src/vespa/searchlib/tensor/tensor_attribute.cpp  | 20 +++++++++++++++-----
 .../src/vespa/searchlib/tensor/tensor_attribute.h    |  2 ++
 4 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
index 084dcee141b..be15967ab10 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
@@ -91,6 +91,16 @@ DenseTensorAttribute::consider_remove_from_index(DocId docid)
     }
 }
 
+vespalib::MemoryUsage
+DenseTensorAttribute::update_stat()
+{
+    vespalib::MemoryUsage result = TensorAttribute::update_stat();
+    if (_index) {
+        result.merge(_index->memory_usage());
+    }
+    return result;
+}
+
 vespalib::MemoryUsage
 DenseTensorAttribute::memory_usage() const
 {
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
index 097c3c93dad..0884b776ca5 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
@@ -23,6 +23,7 @@ private:
 
     void internal_set_tensor(DocId docid, const vespalib::eval::Value& tensor);
     void consider_remove_from_index(DocId docid);
+    vespalib::MemoryUsage update_stat() override;
     vespalib::MemoryUsage memory_usage() const override;
 
 public:
diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp
index e8a25e0e23e..5044a989203 100644
--- a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.cpp
@@ -52,7 +52,8 @@ TensorAttribute::TensorAttribute(vespalib::stringref name, const Config &cfg, Te
                  getGenerationHolder()),
       _tensorStore(tensorStore),
       _emptyTensor(createEmptyTensor(cfg.tensorType())),
-      _compactGeneration(0)
+      _compactGeneration(0),
+      _cached_tensor_store_memory_usage()
 {
 }
 
@@ -84,9 +85,8 @@ TensorAttribute::onCommit()
     incGeneration();
     if (getFirstUsedGeneration() > _compactGeneration) {
         // No data held from previous compact operation
-        Status &status = getStatus();
-        size_t used = status.getUsed();
-        size_t dead = status.getDead();
+        size_t used = _cached_tensor_store_memory_usage.usedBytes();
+        size_t dead = _cached_tensor_store_memory_usage.deadBytes();
         if (getConfig().getCompactionStrategy().should_compact_memory(used, dead)) {
             compactWorst();
         }
@@ -96,7 +96,7 @@ TensorAttribute::onCommit()
 void
 TensorAttribute::onUpdateStat()
 {
-    vespalib::MemoryUsage total = memory_usage();
+    vespalib::MemoryUsage total = update_stat();
     this->updateStatistics(_refVector.size(),
                            _refVector.size(),
                            total.allocatedBytes(),
@@ -160,6 +160,16 @@ TensorAttribute::setTensorRef(DocId docId, EntryRef ref)
     }
 }
 
+vespalib::MemoryUsage
+TensorAttribute::update_stat()
+{
+    vespalib::MemoryUsage result = _refVector.getMemoryUsage();
+    _cached_tensor_store_memory_usage = _tensorStore.getMemoryUsage();
+    result.merge(_cached_tensor_store_memory_usage);
+    result.mergeGenerationHeldBytes(getGenerationHolder().getHeldBytes());
+    return result;
+}
+
 vespalib::MemoryUsage
 TensorAttribute::memory_usage() const
 {
diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h
index 601e19e54d1..ea0f7391120 100644
--- a/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/tensor_attribute.h
@@ -26,11 +26,13 @@ protected:
     TensorStore &_tensorStore; // data store for serialized tensors
     std::unique_ptr<vespalib::eval::Value> _emptyTensor;
     uint64_t    _compactGeneration; // Generation when last compact occurred
+    vespalib::MemoryUsage _cached_tensor_store_memory_usage;
 
     template <typename RefType>
     void doCompactWorst();
     void checkTensorType(const vespalib::eval::Value &tensor);
     void setTensorRef(DocId docId, EntryRef ref);
+    virtual vespalib::MemoryUsage update_stat();
     virtual vespalib::MemoryUsage memory_usage() const;
     void populate_state(vespalib::slime::Cursor& object) const;
 
-- 
cgit v1.2.3