Report memory usage from unique store dictionary hash.

8 files changed, 138 insertions, 25 deletions

diff --git a/vespalib/src/tests/datastore/fixed_size_hash_map/fixed_size_hash_map_test.cpp b/vespalib/src/tests/datastore/fixed_size_hash_map/fixed_size_hash_map_test.cpp
index 0fad98af10b..b929b248e33 100644
--- a/vespalib/src/tests/datastore/fixed_size_hash_map/fixed_size_hash_map_test.cpp
+++ b/vespalib/src/tests/datastore/fixed_size_hash_map/fixed_size_hash_map_test.cpp
@@ -275,4 +275,31 @@ TEST_F(DataStoreFixedSizeHashTest, lookups_works_after_insert_and_remove)
     }
 }
 
+TEST_F(DataStoreFixedSizeHashTest, memory_usage_is_reported)
+{
+    auto initial_usage = _hash_map->get_memory_usage();
+    EXPECT_LT(0, initial_usage.allocatedBytes());
+    EXPECT_LT(0, initial_usage.usedBytes());
+    EXPECT_LT(initial_usage.usedBytes(), initial_usage.allocatedBytes());
+    EXPECT_EQ(0, initial_usage.deadBytes());
+    EXPECT_EQ(0, initial_usage.allocatedBytesOnHold());
+    auto guard = _generation_handler.takeGuard();
+    insert(10);
+    remove(10);
+    commit();
+    auto usage1 = _hash_map->get_memory_usage();
+    EXPECT_EQ(initial_usage.allocatedBytes(), usage1.allocatedBytes());
+    EXPECT_LT(initial_usage.usedBytes(), usage1.usedBytes());
+    EXPECT_LT(usage1.usedBytes(), usage1.allocatedBytes());
+    EXPECT_EQ(0, usage1.deadBytes());
+    EXPECT_LT(0, usage1.allocatedBytesOnHold());
+    guard = GenerationHandler::Guard();
+    commit();
+    auto usage2 = _hash_map->get_memory_usage();
+    EXPECT_EQ(initial_usage.allocatedBytes(), usage2.allocatedBytes());
+    EXPECT_EQ(usage1.usedBytes(), usage2.usedBytes());
+    EXPECT_LT(0, usage2.deadBytes());
+    EXPECT_EQ(0, usage2.allocatedBytesOnHold());
+}
+
 GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/vespalib/src/tests/datastore/simple_hash_map/simple_hash_map_test.cpp b/vespalib/src/tests/datastore/simple_hash_map/simple_hash_map_test.cpp
index 4dbb35d8516..866c2a30818 100644
--- a/vespalib/src/tests/datastore/simple_hash_map/simple_hash_map_test.cpp
+++ b/vespalib/src/tests/datastore/simple_hash_map/simple_hash_map_test.cpp
@@ -200,4 +200,20 @@ TEST_F(DataStoreSimpleHashTest, multi_threaded_reader_during_updates)
     }
 }
 
+TEST_F(DataStoreSimpleHashTest, memory_usage_is_reported)
+{
+    auto initial_usage = _hash_map.get_memory_usage();
+    EXPECT_LT(0, initial_usage.allocatedBytes());
+    EXPECT_LT(0, initial_usage.usedBytes());
+    EXPECT_EQ(0, initial_usage.deadBytes());
+    EXPECT_EQ(0, initial_usage.allocatedBytesOnHold());
+    auto guard = _generationHandler.takeGuard();
+    for (uint32_t i = 0; i < 50; ++i) {
+        insert(i);
+    }
+    auto usage = _hash_map.get_memory_usage();
+    EXPECT_EQ(0, usage.deadBytes());
+    EXPECT_LT(0, usage.allocatedBytesOnHold());
+}
+
 GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/vespalib/src/tests/datastore/unique_store_dictionary/unique_store_dictionary_test.cpp b/vespalib/src/tests/datastore/unique_store_dictionary/unique_store_dictionary_test.cpp
index c553113b775..2085292cdf4 100644
--- a/vespalib/src/tests/datastore/unique_store_dictionary/unique_store_dictionary_test.cpp
+++ b/vespalib/src/tests/datastore/unique_store_dictionary/unique_store_dictionary_test.cpp
@@ -2,6 +2,8 @@
 
 #include <vespa/vespalib/datastore/unique_store.hpp>
 #include <vespa/vespalib/datastore/unique_store_dictionary.hpp>
+#include <vespa/vespalib/datastore/simple_hash_map.h>
+#include <vespa/vespalib/util/memoryusage.h>
 #include <vespa/vespalib/gtest/gtest.h>
 
 #include <vespa/log/log.h>
@@ -36,12 +38,13 @@ public:
     }
 };
 
+template <typename UniqueStoreDictionaryType>
 struct DictionaryReadTest : public ::testing::Test {
-    DefaultUniqueStoreDictionary dict;
+    UniqueStoreDictionaryType dict;
     IUniqueStoreDictionary::ReadSnapshot::UP snapshot;
 
     DictionaryReadTest()
-        : dict(std::unique_ptr<EntryComparator>()),
+        : dict(std::make_unique<Comparator>(0)),
           snapshot()
     {
     }
@@ -56,35 +59,56 @@ struct DictionaryReadTest : public ::testing::Test {
     }
 };
 
-TEST_F(DictionaryReadTest, can_count_occurrences_of_a_key)
+using DictionaryReadTestTypes = ::testing::Types<DefaultUniqueStoreDictionary, UniqueStoreDictionary<DefaultDictionary, IUniqueStoreDictionary, SimpleHashMap>>;
+VESPA_GTEST_TYPED_TEST_SUITE(DictionaryReadTest, DictionaryReadTestTypes);
+
+// Disable warnings emitted by gtest generated files when using typed tests
+#pragma GCC diagnostic push
+#ifndef __clang__
+#pragma GCC diagnostic ignored "-Wsuggest-override"
+#endif
+
+TYPED_TEST(DictionaryReadTest, can_count_occurrences_of_a_key)
 {
-    add(3).add(5).take_snapshot();
-    EXPECT_EQ(0, snapshot->count(Comparator(2)));
-    EXPECT_EQ(1, snapshot->count(Comparator(3)));
-    EXPECT_EQ(0, snapshot->count(Comparator(4)));
-    EXPECT_EQ(1, snapshot->count(Comparator(5)));
+    this->add(3).add(5).take_snapshot();
+    EXPECT_EQ(0, this->snapshot->count(Comparator(2)));
+    EXPECT_EQ(1, this->snapshot->count(Comparator(3)));
+    EXPECT_EQ(0, this->snapshot->count(Comparator(4)));
+    EXPECT_EQ(1, this->snapshot->count(Comparator(5)));
 }
 
-TEST_F(DictionaryReadTest, can_count_occurrences_of_keys_in_a_range)
+TYPED_TEST(DictionaryReadTest, can_count_occurrences_of_keys_in_a_range)
 {
-    add(3).add(5).add(7).add(9).take_snapshot();
-    EXPECT_EQ(1, snapshot->count_in_range(Comparator(3), Comparator(3)));
-    EXPECT_EQ(1, snapshot->count_in_range(Comparator(3), Comparator(4)));
-    EXPECT_EQ(2, snapshot->count_in_range(Comparator(3), Comparator(5)));
-    EXPECT_EQ(3, snapshot->count_in_range(Comparator(3), Comparator(7)));
-    EXPECT_EQ(4, snapshot->count_in_range(Comparator(3), Comparator(9)));
-    EXPECT_EQ(4, snapshot->count_in_range(Comparator(3), Comparator(10)));
-
-    EXPECT_EQ(0, snapshot->count_in_range(Comparator(5), Comparator(3)));
+    this->add(3).add(5).add(7).add(9).take_snapshot();
+    EXPECT_EQ(1, this->snapshot->count_in_range(Comparator(3), Comparator(3)));
+    EXPECT_EQ(1, this->snapshot->count_in_range(Comparator(3), Comparator(4)));
+    EXPECT_EQ(2, this->snapshot->count_in_range(Comparator(3), Comparator(5)));
+    EXPECT_EQ(3, this->snapshot->count_in_range(Comparator(3), Comparator(7)));
+    EXPECT_EQ(4, this->snapshot->count_in_range(Comparator(3), Comparator(9)));
+    EXPECT_EQ(4, this->snapshot->count_in_range(Comparator(3), Comparator(10)));
+
+    EXPECT_EQ(0, this->snapshot->count_in_range(Comparator(5), Comparator(3)));
 }
 
-TEST_F(DictionaryReadTest, can_iterate_all_keys)
+TYPED_TEST(DictionaryReadTest, can_iterate_all_keys)
 {
     using EntryRefVector = std::vector<EntryRef>;
-    add(3).add(5).add(7).take_snapshot();
+    this->add(3).add(5).add(7).take_snapshot();
     EntryRefVector refs;
-    snapshot->foreach_key([&](EntryRef ref){ refs.emplace_back(ref); });
+    this->snapshot->foreach_key([&](EntryRef ref){ refs.emplace_back(ref); });
     EXPECT_EQ(EntryRefVector({EntryRef(3), EntryRef(5), EntryRef(7)}), refs);
 }
 
+TYPED_TEST(DictionaryReadTest, memory_usage_is_reported)
+{
+    auto initial_usage = this->dict.get_memory_usage();
+    this->add(10);
+    auto usage = this->dict.get_memory_usage();
+    EXPECT_LT(initial_usage.usedBytes(), usage.usedBytes());
+    EXPECT_EQ(initial_usage.deadBytes(), usage.deadBytes());
+    EXPECT_EQ(0, usage.allocatedBytesOnHold());
+}
+
+#pragma GCC diagnostic pop
+
 GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/vespalib/src/vespa/vespalib/datastore/fixed_size_hash_map.cpp b/vespalib/src/vespa/vespalib/datastore/fixed_size_hash_map.cpp
index be67e0b1863..4822313e39e 100644
--- a/vespalib/src/vespa/vespalib/datastore/fixed_size_hash_map.cpp
+++ b/vespalib/src/vespa/vespalib/datastore/fixed_size_hash_map.cpp
@@ -3,6 +3,7 @@
 #include "fixed_size_hash_map.h"
 #include "entry_comparator.h"
 #include <vespa/vespalib/util/array.hpp>
+#include <vespa/vespalib/util/memoryusage.h>
 #include <cassert>
 #include <stdexcept>
 
@@ -171,4 +172,19 @@ FixedSizeHashMap::find(const EntryComparator& comp, EntryRef key_ref)
     return nullptr;
 }
 
+MemoryUsage
+FixedSizeHashMap::get_memory_usage() const
+{
+    size_t fixed_size = sizeof(FixedSizeHashMap);
+    size_t chain_heads_size = sizeof(ChainHead) * _chain_heads.size();
+    size_t nodes_used_size = sizeof(Node) * _nodes.size();
+    size_t nodes_alloc_size = sizeof(Node) * _nodes.capacity();
+    size_t nodes_dead_size = sizeof(Node) * _free_count;
+    size_t nodes_hold_size = sizeof(Node) * _hold_count;
+    return MemoryUsage(fixed_size + chain_heads_size + nodes_alloc_size,
+                       fixed_size + chain_heads_size + nodes_used_size,
+                       nodes_dead_size,
+                       nodes_hold_size);
+}
+
 }
diff --git a/vespalib/src/vespa/vespalib/datastore/fixed_size_hash_map.h b/vespalib/src/vespa/vespalib/datastore/fixed_size_hash_map.h
index 72373f909df..8153bc82e06 100644
--- a/vespalib/src/vespa/vespalib/datastore/fixed_size_hash_map.h
+++ b/vespalib/src/vespa/vespalib/datastore/fixed_size_hash_map.h
@@ -11,7 +11,10 @@
 #include <deque>
 #include <functional>
 
-namespace vespalib { class GenerationHolder; }
+namespace vespalib {
+class GenerationHolder;
+class MemoryUsage;
+}
 namespace vespalib::datastore {
 
 class EntryComparator;
@@ -114,6 +117,7 @@ public:
 
     bool full() const noexcept { return _nodes.size() == _nodes.capacity() && _free_count == 0u; }
     size_t size() const noexcept { return _count; }
+    MemoryUsage get_memory_usage() const;
 };
 
 }
diff --git a/vespalib/src/vespa/vespalib/datastore/simple_hash_map.cpp b/vespalib/src/vespa/vespalib/datastore/simple_hash_map.cpp
index 6995acc5bb4..2295cecd613 100644
--- a/vespalib/src/vespa/vespalib/datastore/simple_hash_map.cpp
+++ b/vespalib/src/vespa/vespalib/datastore/simple_hash_map.cpp
@@ -3,6 +3,7 @@
 #include "simple_hash_map.h"
 #include "fixed_size_hash_map.h"
 #include "entry_comparator.h"
+#include <vespa/vespalib/util/memoryusage.h>
 
 namespace vespalib::datastore {
 
@@ -61,8 +62,8 @@ SimpleHashMap::alloc_stripe(size_t stripe)
 void
 SimpleHashMap::hold_stripe(std::unique_ptr<const FixedSizeHashMap> map)
 {
-    // TODO: Provider proper held size
-    auto hold = std::make_unique<SimpleHashMapStripeHeld>(0, std::move(map));
+    auto usage = map->get_memory_usage();
+    auto hold = std::make_unique<SimpleHashMapStripeHeld>(usage.allocatedBytes(), std::move(map));
     _gen_holder.hold(std::move(hold));
 }
 
@@ -148,4 +149,20 @@ SimpleHashMap::size() const noexcept
     return result;
 }
 
+MemoryUsage
+SimpleHashMap::get_memory_usage() const
+{
+    MemoryUsage memory_usage(sizeof(SimpleHashMap), sizeof(SimpleHashMap), 0, 0);
+    for (size_t i = 0; i < num_stripes; ++i) {
+        auto map = _maps[i].load(std::memory_order_relaxed);
+        if (map != nullptr) {
+            memory_usage.merge(map->get_memory_usage());
+        }
+    }
+    size_t gen_holder_held_bytes = _gen_holder.getHeldBytes();
+    memory_usage.incAllocatedBytes(gen_holder_held_bytes);
+    memory_usage.incAllocatedBytesOnHold(gen_holder_held_bytes);
+    return memory_usage;
+}
+
 }
diff --git a/vespalib/src/vespa/vespalib/datastore/simple_hash_map.h b/vespalib/src/vespa/vespalib/datastore/simple_hash_map.h
index 2d97a9d45fc..bc052295511 100644
--- a/vespalib/src/vespa/vespalib/datastore/simple_hash_map.h
+++ b/vespalib/src/vespa/vespalib/datastore/simple_hash_map.h
@@ -7,6 +7,7 @@
 #include <vespa/vespalib/util/generationholder.h>
 #include <functional>
 
+namespace vespalib { class MemoryUsage; }
 namespace vespalib::datastore {
 
 class FixedSizeHashMap;
@@ -54,6 +55,7 @@ public:
     void trim_hold_lists(generation_t first_used);
     size_t size() const noexcept;
     const EntryComparator &get_default_comparator() const noexcept { return *_comp; }
+    MemoryUsage get_memory_usage() const;
 };
 
 }
diff --git a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.hpp b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.hpp
index ae8a85985fd..edacb306c6e 100644
--- a/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.hpp
+++ b/vespalib/src/vespa/vespalib/datastore/unique_store_dictionary.hpp
@@ -186,7 +186,14 @@ template <typename BTreeDictionaryT, typename ParentT, typename HashDictionaryT>
 vespalib::MemoryUsage
 UniqueStoreDictionary<BTreeDictionaryT, ParentT, HashDictionaryT>::get_memory_usage() const
 {
-    return this->_btree_dict.getMemoryUsage();
+    vespalib::MemoryUsage memory_usage;
+    if constexpr (has_btree_dictionary) {
+        memory_usage.merge(this->_btree_dict.getMemoryUsage());
+    }
+    if constexpr (has_hash_dictionary) {
+        memory_usage.merge(this->_hash_dict.get_memory_usage());
+    }
+    return memory_usage;
 }
 
 template <typename BTreeDictionaryT, typename ParentT, typename HashDictionaryT>