Merge pull request #31061 from vespa-engine/toregge/expose-imported-attributes-in-metrics

Expose imported attributes in metrics.

4 files changed, 45 insertions, 1 deletions

diff --git a/searchlib/src/vespa/searchlib/attribute/bitvector_search_cache.cpp b/searchlib/src/vespa/searchlib/attribute/bitvector_search_cache.cpp
index e20d02afe50..6762c0516b2 100644
--- a/searchlib/src/vespa/searchlib/attribute/bitvector_search_cache.cpp
+++ b/searchlib/src/vespa/searchlib/attribute/bitvector_search_cache.cpp
@@ -3,6 +3,7 @@
 #include "bitvector_search_cache.h"
 #include <vespa/searchlib/common/bitvector.h>
 #include <vespa/vespalib/stllike/hash_map.hpp>
+#include <vespa/vespalib/util/memoryusage.h>
 #include <mutex>
 
 namespace search::attribute {
@@ -10,6 +11,7 @@ namespace search::attribute {
 BitVectorSearchCache::BitVectorSearchCache()
     : _mutex(),
       _size(0),
+      _entries_extra_memory_usage(0),
       _cache()
 {}
 
@@ -18,9 +20,19 @@ BitVectorSearchCache::~BitVectorSearchCache() = default;
 void
 BitVectorSearchCache::insert(const vespalib::string &term, std::shared_ptr<Entry> entry)
 {
+    size_t entry_extra_memory_usage = 0;
+    if (entry) {
+        entry_extra_memory_usage = sizeof(Entry);
+        if (entry->bitVector) {
+            entry_extra_memory_usage += entry->bitVector->getFileBytes();
+        }
+    }
     std::unique_lock guard(_mutex);
-    _cache.insert(std::make_pair(term, std::move(entry)));
+    auto ins_res = _cache.insert(std::make_pair(term, std::move(entry)));
     _size.store(_cache.size());
+    if (ins_res.second) {
+        _entries_extra_memory_usage += entry_extra_memory_usage;
+    }
 }
 
 std::shared_ptr<BitVectorSearchCache::Entry>
@@ -36,12 +48,25 @@ BitVectorSearchCache::find(const vespalib::string &term) const
     return {};
 }
 
+vespalib::MemoryUsage
+BitVectorSearchCache::get_memory_usage() const
+{
+    std::lock_guard guard(_mutex);
+    size_t cache_memory_consumption = _cache.getMemoryConsumption();
+    size_t cache_memory_used = _cache.getMemoryUsed();
+    size_t self_memory_used = sizeof(BitVectorSearchCache) - sizeof(_cache);
+    size_t allocated = self_memory_used + cache_memory_consumption + _entries_extra_memory_usage;
+    size_t used = self_memory_used + cache_memory_used + _entries_extra_memory_usage;
+    return vespalib::MemoryUsage(allocated, used, 0, 0);
+}
+
 void
 BitVectorSearchCache::clear()
 {
     std::unique_lock guard(_mutex);
     _cache.clear();
     _size.store(0ul, std::memory_order_relaxed);
+    _entries_extra_memory_usage = 0;
 }
 
 }
diff --git a/searchlib/src/vespa/searchlib/attribute/bitvector_search_cache.h b/searchlib/src/vespa/searchlib/attribute/bitvector_search_cache.h
index 233f8315aaf..3a38cdcea26 100644
--- a/searchlib/src/vespa/searchlib/attribute/bitvector_search_cache.h
+++ b/searchlib/src/vespa/searchlib/attribute/bitvector_search_cache.h
@@ -10,6 +10,8 @@
 #include <atomic>
 
 namespace search { class BitVector; }
+namespace vespalib { class MemoryUsage; }
+
 namespace search::attribute {
 
 /**
@@ -37,6 +39,7 @@ private:
 
     mutable std::shared_mutex _mutex;
     std::atomic<uint64_t>     _size;
+    size_t                    _entries_extra_memory_usage;
     Cache _cache;
 
 public:
@@ -45,6 +48,7 @@ public:
     void insert(const vespalib::string &term, std::shared_ptr<Entry> entry);
     std::shared_ptr<Entry> find(const vespalib::string &term) const;
     size_t size() const { return _size.load(std::memory_order_relaxed); }
+    vespalib::MemoryUsage get_memory_usage() const;
     void clear();
 };
 
diff --git a/searchlib/src/vespa/searchlib/attribute/imported_attribute_vector.cpp b/searchlib/src/vespa/searchlib/attribute/imported_attribute_vector.cpp
index 029dc155785..f6a33165f0c 100644
--- a/searchlib/src/vespa/searchlib/attribute/imported_attribute_vector.cpp
+++ b/searchlib/src/vespa/searchlib/attribute/imported_attribute_vector.cpp
@@ -3,6 +3,7 @@
 #include "imported_attribute_vector.h"
 #include "imported_attribute_vector_read_guard.h"
 #include "imported_search_context.h"
+#include <vespa/vespalib/util/memoryusage.h>
 
 namespace search::attribute {
 
@@ -58,4 +59,15 @@ void ImportedAttributeVector::clearSearchCache() {
     }
 }
 
+vespalib::MemoryUsage
+ImportedAttributeVector::get_memory_usage() const
+{
+    constexpr auto self_memory_usage = sizeof(ImportedAttributeVector);
+    vespalib::MemoryUsage result(self_memory_usage, self_memory_usage, 0, 0);
+    if (_search_cache) {
+        result.merge(_search_cache->get_memory_usage());
+    }
+    return result;
+}
+
 }
diff --git a/searchlib/src/vespa/searchlib/attribute/imported_attribute_vector.h b/searchlib/src/vespa/searchlib/attribute/imported_attribute_vector.h
index bd018df5273..5b68957b7f5 100644
--- a/searchlib/src/vespa/searchlib/attribute/imported_attribute_vector.h
+++ b/searchlib/src/vespa/searchlib/attribute/imported_attribute_vector.h
@@ -6,6 +6,8 @@
 #include <vespa/searchcommon/attribute/i_document_meta_store_context.h>
 #include <vespa/vespalib/stllike/string.h>
 
+namespace vespalib { class MemoryUsage; }
+
 namespace search::attribute {
 
 class BitVectorSearchCache;
@@ -62,6 +64,7 @@ public:
 
     std::unique_ptr<AttributeReadGuard> makeReadGuard(bool stableEnumGuard) const override;
     virtual std::unique_ptr<AttributeReadGuard> makeReadGuard(std::shared_ptr<MetaStoreReadGuard> targetMetaStoreReadGuard, bool stableEnumGuard) const;
+    vespalib::MemoryUsage get_memory_usage() const;
 
 protected:
     vespalib::string                           _name;