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

Expose imported attributes in metrics.

6 files changed, 62 insertions, 1 deletions

diff --git a/searchlib/src/tests/attribute/bitvector_search_cache/bitvector_search_cache_test.cpp b/searchlib/src/tests/attribute/bitvector_search_cache/bitvector_search_cache_test.cpp
index d51ec22a54a..1d66eefaff7 100644
--- a/searchlib/src/tests/attribute/bitvector_search_cache/bitvector_search_cache_test.cpp
+++ b/searchlib/src/tests/attribute/bitvector_search_cache/bitvector_search_cache_test.cpp
@@ -3,6 +3,7 @@
 #include <vespa/vespalib/testkit/test_kit.h>
 #include <vespa/searchlib/attribute/bitvector_search_cache.h>
 #include <vespa/searchlib/common/bitvector.h>
+#include <vespa/vespalib/util/memoryusage.h>
 
 using namespace search;
 using namespace search::attribute;
@@ -31,9 +32,13 @@ struct Fixture {
 TEST_F("require that bit vectors can be inserted and retrieved", Fixture)
 {
     EXPECT_EQUAL(0u, f.cache.size());
+    auto old_mem_usage = f.cache.get_memory_usage();
     f.cache.insert("foo", f.entry1);
     f.cache.insert("bar", f.entry2);
     EXPECT_EQUAL(2u, f.cache.size());
+    auto new_mem_usage = f.cache.get_memory_usage();
+    EXPECT_LESS(old_mem_usage.usedBytes(), new_mem_usage.usedBytes());
+    EXPECT_LESS(old_mem_usage.allocatedBytes(), new_mem_usage.allocatedBytes());
 
     EXPECT_EQUAL(f.entry1, f.cache.find("foo"));
     EXPECT_EQUAL(f.entry2, f.cache.find("bar"));
@@ -43,9 +48,13 @@ TEST_F("require that bit vectors can be inserted and retrieved", Fixture)
 TEST_F("require that insert() doesn't replace existing bit vector", Fixture)
 {
     f.cache.insert("foo", f.entry1);
+    auto old_mem_usage = f.cache.get_memory_usage();
     f.cache.insert("foo", f.entry2);
+    auto new_mem_usage = f.cache.get_memory_usage();
     EXPECT_EQUAL(1u, f.cache.size());
     EXPECT_EQUAL(f.entry1, f.cache.find("foo"));
+    EXPECT_EQUAL(old_mem_usage.usedBytes(), new_mem_usage.usedBytes());
+    EXPECT_EQUAL(old_mem_usage.allocatedBytes(), new_mem_usage.allocatedBytes());
 }
 
 TEST_F("require that cache can be cleared", Fixture)
@@ -53,11 +62,15 @@ TEST_F("require that cache can be cleared", Fixture)
     f.cache.insert("foo", f.entry1);
     f.cache.insert("bar", f.entry2);
     EXPECT_EQUAL(2u, f.cache.size());
+    auto old_mem_usage = f.cache.get_memory_usage();
     f.cache.clear();
+    auto new_mem_usage = f.cache.get_memory_usage();
 
     EXPECT_EQUAL(0u, f.cache.size());
     EXPECT_TRUE(f.cache.find("foo").get() == nullptr);
     EXPECT_TRUE(f.cache.find("bar").get() == nullptr);
+    EXPECT_GREATER(old_mem_usage.usedBytes(), new_mem_usage.usedBytes());
+    EXPECT_GREATER(old_mem_usage.allocatedBytes(), new_mem_usage.allocatedBytes());
 }
 
 TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/searchlib/src/tests/attribute/imported_search_context/imported_search_context_test.cpp b/searchlib/src/tests/attribute/imported_search_context/imported_search_context_test.cpp
index 41ec377dece..7c38c322bc8 100644
--- a/searchlib/src/tests/attribute/imported_search_context/imported_search_context_test.cpp
+++ b/searchlib/src/tests/attribute/imported_search_context/imported_search_context_test.cpp
@@ -508,6 +508,7 @@ assertBitVector(const std::vector<uint32_t> &expDocIds, const BitVector &bitVect
 TEST_F("Entry is inserted into search cache if bit vector posting list is used", SearchCacheFixture)
 {
     EXPECT_EQUAL(0u, f.imported_attr->getSearchCache()->size());
+    auto old_mem_usage = f.imported_attr->get_memory_usage();
     auto ctx = f.create_context(word_term("5678"));
     ctx->fetchPostings(queryeval::ExecuteInfo::FULL, true);
     TermFieldMatchData match;
@@ -515,6 +516,9 @@ TEST_F("Entry is inserted into search cache if bit vector posting list is used",
     TEST_DO(f.assertSearch({3, 5}, *iter));
 
     EXPECT_EQUAL(1u, f.imported_attr->getSearchCache()->size());
+    auto new_mem_usage = f.imported_attr->get_memory_usage();
+    EXPECT_LESS(old_mem_usage.usedBytes(), new_mem_usage.usedBytes());
+    EXPECT_LESS(old_mem_usage.allocatedBytes(), new_mem_usage.allocatedBytes());
     auto cacheEntry = f.imported_attr->getSearchCache()->find("5678");
     EXPECT_EQUAL(cacheEntry->docIdLimit, f.get_imported_attr()->getNumDocs());
     TEST_DO(assertBitVector({3, 5}, *cacheEntry->bitVector));
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;