txgraph: keep track of Cluster memory usage (preparation)

Author: sipa

src/txgraph.cpp

@@ -153,6 +153,9 @@ class Cluster
         return m_quality == QualityLevel::NEEDS_SPLIT ||
                m_quality == QualityLevel::NEEDS_SPLIT_ACCEPTABLE;
     }
+    /** Total memory usage currently for this Cluster, including all its dynamic memory, plus Cluster
+     *  structure itself, and ClusterSet::m_clusters entry. */
+    size_t TotalMemoryUsage() const noexcept;
     /** Get the number of transactions in this Cluster. */
     LinearizationIndex GetTxCount() const noexcept { return m_linearization.size(); }
     /** Get the total size of the transactions in this Cluster. */
@@ -308,6 +311,9 @@ class TxGraphImpl final : public TxGraph
         GraphIndex m_txcount_oversized{0};
         /** Whether this graph is oversized (if known). */
         std::optional<bool> m_oversized{false};
+        /** The combined TotalMemoryUsage of all clusters in this level (only Clusters that
+         *  are materialized; in staging, implicit Clusters from main are not counted), */
+        size_t m_cluster_usage{0};
 
         ClusterSet() noexcept = default;
     };
@@ -703,6 +709,18 @@ void TxGraphImpl::CreateChunkData(GraphIndex idx, LinearizationIndex chunk_count
     }
 }
 
+size_t Cluster::TotalMemoryUsage() const noexcept
+{
+    return // Dynamic memory allocated in this Cluster.
+           memusage::DynamicUsage(m_mapping) + memusage::DynamicUsage(m_linearization) +
+           // Dynamic memory usage inside m_depgraph.
+           m_depgraph.DynamicMemoryUsage() +
+           // Memory usage of the allocated Cluster itself.
+           memusage::MallocUsage(sizeof(Cluster)) +
+           // Memory usage of the ClusterSet::m_clusters entry.
+           sizeof(std::unique_ptr<Cluster>);
+}
+
 uint64_t Cluster::GetTotalTxSize() const noexcept
 {
     uint64_t ret{0};
@@ -851,9 +869,11 @@ Cluster* Cluster::CopyToStaging(TxGraphImpl& graph) const noexcept
     ptr->m_mapping = m_mapping;
     ptr->m_linearization = m_linearization;
     // Insert the new Cluster into the graph.
-    graph.InsertCluster(1, std::move(ret), m_quality);
+    graph.InsertCluster(/*level=*/1, std::move(ret), m_quality);
     // Update its Locators.
     ptr->Updated(graph, /*level=*/1);
+    // Update memory usage.
+    graph.GetClusterSet(/*level=*/1).m_cluster_usage += ptr->TotalMemoryUsage();
     return ptr;
 }
 
@@ -862,6 +882,7 @@ void Cluster::ApplyRemovals(TxGraphImpl& graph, int level, std::span<GraphIndex>
     // Iterate over the prefix of to_remove that applies to this cluster.
     Assume(!to_remove.empty());
     SetType todo;
+    graph.GetClusterSet(level).m_cluster_usage -= TotalMemoryUsage();
     do {
         GraphIndex idx = to_remove.front();
         Assume(idx < graph.m_entries.size());
@@ -913,12 +934,14 @@ void Cluster::ApplyRemovals(TxGraphImpl& graph, int level, std::span<GraphIndex>
         quality = QualityLevel::NEEDS_SPLIT;
     }
     Compact();
+    graph.GetClusterSet(level).m_cluster_usage += TotalMemoryUsage();
     graph.SetClusterQuality(level, m_quality, m_setindex, quality);
     Updated(graph, level);
 }
 
 void Cluster::Clear(TxGraphImpl& graph, int level) noexcept
 {
+    graph.GetClusterSet(level).m_cluster_usage -= TotalMemoryUsage();
     for (auto i : m_linearization) {
         graph.ClearLocator(level, m_mapping[i], m_quality == QualityLevel::OVERSIZED_SINGLETON);
     }
@@ -935,6 +958,10 @@ void Cluster::MoveToMain(TxGraphImpl& graph) noexcept
         entry.m_locator[1].SetMissing();
     }
     auto quality = m_quality;
+    // Subtract memory usage from staging and add it to main.
+    graph.GetClusterSet(/*level=*/1).m_cluster_usage -= TotalMemoryUsage();
+    graph.GetClusterSet(/*level=*/0).m_cluster_usage += TotalMemoryUsage();
+    // Remove cluster itself from staging and add it to main.
     auto cluster = graph.ExtractCluster(1, quality, m_setindex);
     graph.InsertCluster(/*level=*/0, std::move(cluster), quality);
     Updated(graph, /*level=*/0);
@@ -1036,6 +1063,8 @@ bool Cluster::Split(TxGraphImpl& graph, int level) noexcept
         graph.InsertCluster(level, std::move(new_cluster), split_quality);
         todo -= component;
     }
+    // We have to split the Cluster up. Remove accounting for the existing one first.
+    graph.GetClusterSet(level).m_cluster_usage -= TotalMemoryUsage();
     // Redistribute the transactions.
     for (auto i : m_linearization) {
         /** The cluster which transaction originally in position i is moved to. */
@@ -1057,10 +1086,11 @@ bool Cluster::Split(TxGraphImpl& graph, int level) noexcept
         for (auto par : m_depgraph.GetReducedParents(i)) new_parents.Set(remap[par].second);
         new_cluster->m_depgraph.AddDependencies(new_parents, remap[i].second);
     }
-    // Update all the Locators of moved transactions.
+    // Update all the Locators of moved transactions, and memory usage.
     for (Cluster* new_cluster : new_clusters) {
         new_cluster->Updated(graph, level);
         new_cluster->Compact();
+        graph.GetClusterSet(level).m_cluster_usage += new_cluster->TotalMemoryUsage();
     }
     // Wipe this Cluster, and return that it needs to be deleted.
     m_depgraph = DepGraph<SetType>{};
@@ -1624,7 +1654,9 @@ void TxGraphImpl::Merge(std::span<Cluster*> to_merge, int level) noexcept
     // moves.
     size_t max_size_pos{0};
     DepGraphIndex max_size = to_merge[max_size_pos]->GetTxCount();
+    GetClusterSet(level).m_cluster_usage -= to_merge[max_size_pos]->TotalMemoryUsage();
     for (size_t i = 1; i < to_merge.size(); ++i) {
+        GetClusterSet(level).m_cluster_usage -= to_merge[i]->TotalMemoryUsage();
         DepGraphIndex size = to_merge[i]->GetTxCount();
         if (size > max_size) {
             max_size_pos = i;
@@ -1639,6 +1671,7 @@ void TxGraphImpl::Merge(std::span<Cluster*> to_merge, int level) noexcept
         DeleteCluster(*to_merge[i], level);
     }
     to_merge[0]->Compact();
+    GetClusterSet(level).m_cluster_usage += to_merge[0]->TotalMemoryUsage();
 }
 
 void TxGraphImpl::ApplyDependencies(int level) noexcept
@@ -1764,6 +1797,7 @@ TxGraph::Ref TxGraphImpl::AddTransaction(const FeePerWeight& feerate) noexcept
     auto& clusterset = GetClusterSet(level);
     InsertCluster(level, std::move(cluster), oversized ? QualityLevel::OVERSIZED_SINGLETON : QualityLevel::OPTIMAL);
     cluster_ptr->Updated(*this, level);
+    clusterset.m_cluster_usage += cluster_ptr->TotalMemoryUsage();
     ++clusterset.m_txcount;
     // Deal with individually oversized transactions.
     if (oversized) {
@@ -2116,6 +2150,7 @@ void TxGraphImpl::StartStaging() noexcept
     m_staging_clusterset->m_group_data = m_main_clusterset.m_group_data;
     m_staging_clusterset->m_oversized = m_main_clusterset.m_oversized;
     Assume(m_staging_clusterset->m_oversized.has_value());
+    m_staging_clusterset->m_cluster_usage = 0;
 }
 
 void TxGraphImpl::AbortStaging() noexcept
@@ -2413,6 +2448,7 @@ void TxGraphImpl::SanityCheck() const
         assert(level < MAX_LEVELS);
         auto& clusterset = GetClusterSet(level);
         std::set<const Cluster*> actual_clusters;
+        size_t recomputed_cluster_usage{0};
 
         // For all quality levels...
         for (int qual = 0; qual < int(QualityLevel::NONE); ++qual) {
@@ -2448,9 +2484,14 @@ void TxGraphImpl::SanityCheck() const
                 // Check that the cluster's quality and setindex matches its position in the quality list.
                 assert(cluster.m_quality == quality);
                 assert(cluster.m_setindex == setindex);
+                // Count memory usage.
+                recomputed_cluster_usage += cluster.TotalMemoryUsage();
             }
         }
 
+        // Verify memory usage.
+        assert(clusterset.m_cluster_usage == recomputed_cluster_usage);
+
         // Verify that all to-be-removed transactions have valid identifiers.
         for (GraphIndex idx : clusterset.m_to_remove) {
             assert(idx < m_entries.size());