From e19943f049ed8aa4f32a1d8440a9fbf160367f0f Mon Sep 17 00:00:00 2001
From: Martin Leitner-Ankerl <martin.ankerl@gmail.com>
Date: Sat, 11 Jun 2022 09:28:13 +0200
Subject: [PATCH] Calculate memory usage correctly for unordered_maps that use
 PoolAllocator

Extracts the resource from a PoolAllocator and uses it for
calculation of the node's memory usage.
---
 src/memusage.h          | 20 ++++++++++++++++++++
 src/test/pool_tests.cpp | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/src/memusage.h b/src/memusage.h
index 9755be0ff56..bb39066a7df 100644
--- a/src/memusage.h
+++ b/src/memusage.h
@@ -7,6 +7,7 @@
 
 #include <indirectmap.h>
 #include <prevector.h>
+#include <support/allocators/pool.h>
 
 #include <cassert>
 #include <cstdlib>
@@ -166,6 +167,25 @@ static inline size_t DynamicUsage(const std::unordered_map<X, Y, Z>& m)
     return MallocUsage(sizeof(unordered_node<std::pair<const X, Y> >)) * m.size() + MallocUsage(sizeof(void*) * m.bucket_count());
 }
 
+template <class Key, class T, class Hash, class Pred, std::size_t MAX_BLOCK_SIZE_BYTES, std::size_t ALIGN_BYTES>
+static inline size_t DynamicUsage(const std::unordered_map<Key,
+                                                           T,
+                                                           Hash,
+                                                           Pred,
+                                                           PoolAllocator<std::pair<const Key, T>,
+                                                                         MAX_BLOCK_SIZE_BYTES,
+                                                                         ALIGN_BYTES>>& m)
+{
+    auto* pool_resource = m.get_allocator().resource();
+
+    // The allocated chunks are stored in a std::list. Size per node should
+    // therefore be 3 pointers: next, previous, and a pointer to the chunk.
+    size_t estimated_list_node_size = MallocUsage(sizeof(void*) * 3);
+    size_t usage_resource = estimated_list_node_size * pool_resource->NumAllocatedChunks();
+    size_t usage_chunks = MallocUsage(pool_resource->ChunkSizeBytes()) * pool_resource->NumAllocatedChunks();
+    return usage_resource + usage_chunks + MallocUsage(sizeof(void*) * m.bucket_count());
 }
 
+} // namespace memusage
+
 #endif // BITCOIN_MEMUSAGE_H
diff --git a/src/test/pool_tests.cpp b/src/test/pool_tests.cpp
index dfe857d05bc..8a07e09a442 100644
--- a/src/test/pool_tests.cpp
+++ b/src/test/pool_tests.cpp
@@ -2,6 +2,7 @@
 // Distributed under the MIT software license, see the accompanying
 // file COPYING or http://www.opensource.org/licenses/mit-license.php.
 
+#include <memusage.h>
 #include <support/allocators/pool.h>
 #include <test/util/poolresourcetester.h>
 #include <test/util/random.h>
@@ -153,4 +154,37 @@ BOOST_AUTO_TEST_CASE(random_allocations)
     PoolResourceTester::CheckAllDataAccountedFor(resource);
 }
 
+BOOST_AUTO_TEST_CASE(memusage_test)
+{
+    auto std_map = std::unordered_map<int, int>{};
+
+    using Map = std::unordered_map<int,
+                                   int,
+                                   std::hash<int>,
+                                   std::equal_to<int>,
+                                   PoolAllocator<std::pair<const int, int>,
+                                                 sizeof(std::pair<const int, int>) + sizeof(void*) * 4,
+                                                 alignof(void*)>>;
+    auto resource = Map::allocator_type::ResourceType(1024);
+
+    PoolResourceTester::CheckAllDataAccountedFor(resource);
+
+    {
+        auto resource_map = Map{0, std::hash<int>{}, std::equal_to<int>{}, &resource};
+
+        // can't have the same resource usage
+        BOOST_TEST(memusage::DynamicUsage(std_map) != memusage::DynamicUsage(resource_map));
+
+        for (size_t i = 0; i < 10000; ++i) {
+            std_map[i];
+            resource_map[i];
+        }
+
+        // Eventually the resource_map should have a much lower memory usage because it has less malloc overhead
+        BOOST_TEST(memusage::DynamicUsage(resource_map) <= memusage::DynamicUsage(std_map) * 90 / 100);
+    }
+
+    PoolResourceTester::CheckAllDataAccountedFor(resource);
+}
+
 BOOST_AUTO_TEST_SUITE_END()