make DisconnectedBlockTransactions responsible for its own memory management

Co-authored-by: Cory Fields <cory-nospam-@coryfields.com>

4 files changed, 33 insertions, 27 deletions

diff --git a/src/bench/disconnected_transactions.cpp b/src/bench/disconnected_transactions.cpp
index 72ae9d6c5c..d6f1590950 100644
--- a/src/bench/disconnected_transactions.cpp
+++ b/src/bench/disconnected_transactions.cpp
@@ -73,9 +73,10 @@ static ReorgTxns CreateBlocks(size_t num_not_shared)
 
 static void Reorg(const ReorgTxns& reorg)
 {
-    DisconnectedBlockTransactions disconnectpool;
+    DisconnectedBlockTransactions disconnectpool{MAX_DISCONNECTED_TX_POOL_SIZE * 1000};
     // Disconnect block
-    disconnectpool.AddTransactionsFromBlock(reorg.disconnected_txns);
+    const auto evicted = disconnectpool.AddTransactionsFromBlock(reorg.disconnected_txns);
+    assert(evicted.empty());
 
     // Connect first block
     disconnectpool.removeForBlock(reorg.connected_txns_1);
diff --git a/src/kernel/disconnected_transactions.h b/src/kernel/disconnected_transactions.h
index a5d02c33ee..7db39ba5ca 100644
--- a/src/kernel/disconnected_transactions.h
+++ b/src/kernel/disconnected_transactions.h
@@ -12,7 +12,10 @@
 
 #include <list>
 #include <unordered_map>
+#include <vector>
 
+/** Maximum kilobytes for transactions to store for processing during reorg */
+static const unsigned int MAX_DISCONNECTED_TX_POOL_SIZE = 20'000;
 /**
  * DisconnectedBlockTransactions
 
@@ -38,11 +41,28 @@ private:
     /** Cached dynamic memory usage for the CTransactions (memory for the shared pointers is
      * included in the container calculations). */
     uint64_t cachedInnerUsage = 0;
+    const size_t m_max_mem_usage;
     std::list<CTransactionRef> queuedTx;
     using TxList = decltype(queuedTx);
     std::unordered_map<uint256, TxList::iterator, SaltedTxidHasher> iters_by_txid;
 
+    /** Trim the earliest-added entries until we are within memory bounds. */
+    std::vector<CTransactionRef> LimitMemoryUsage()
+    {
+        std::vector<CTransactionRef> evicted;
+
+        while (!queuedTx.empty() && DynamicMemoryUsage() > m_max_mem_usage) {
+            evicted.emplace_back(queuedTx.front());
+            cachedInnerUsage -= RecursiveDynamicUsage(*queuedTx.front());
+            iters_by_txid.erase(queuedTx.front()->GetHash());
+            queuedTx.pop_front();
+        }
+        return evicted;
+    }
+
 public:
+    DisconnectedBlockTransactions(size_t max_mem_usage) : m_max_mem_usage{max_mem_usage} {}
+
     // It's almost certainly a logic bug if we don't clear out queuedTx before
     // destruction, as we add to it while disconnecting blocks, and then we
     // need to re-process remaining transactions to ensure mempool consistency.
@@ -66,8 +86,9 @@ public:
      * We assume that callers never pass multiple transactions with the same txid, otherwise things
      * can go very wrong in removeForBlock due to queuedTx containing an item without a
      * corresponding entry in iters_by_txid.
+     * @returns vector of transactions that were evicted for size-limiting.
      */
-    void AddTransactionsFromBlock(const std::vector<CTransactionRef>& vtx)
+    [[nodiscard]] std::vector<CTransactionRef> AddTransactionsFromBlock(const std::vector<CTransactionRef>& vtx)
     {
         iters_by_txid.reserve(iters_by_txid.size() + vtx.size());
         for (auto block_it = vtx.rbegin(); block_it != vtx.rend(); ++block_it) {
@@ -75,6 +96,7 @@ public:
             iters_by_txid.emplace((*block_it)->GetHash(), it);
             cachedInnerUsage += RecursiveDynamicUsage(**block_it);
         }
+        return LimitMemoryUsage();
     }
 
     /** Remove any entries that are in this block. */
@@ -95,19 +117,6 @@ public:
         }
     }
 
-    /** Remove the first entry and update memory usage. */
-    CTransactionRef take_first()
-    {
-        CTransactionRef first_tx;
-        if (!queuedTx.empty()) {
-            first_tx = queuedTx.front();
-            cachedInnerUsage -= RecursiveDynamicUsage(*queuedTx.front());
-            iters_by_txid.erase(queuedTx.front()->GetHash());
-            queuedTx.pop_front();
-        }
-        return first_tx;
-    }
-
     size_t size() const { return queuedTx.size(); }
 
     void clear()
diff --git a/src/test/validation_chainstatemanager_tests.cpp b/src/test/validation_chainstatemanager_tests.cpp
index a793f56cba..7b7be4be9e 100644
--- a/src/test/validation_chainstatemanager_tests.cpp
+++ b/src/test/validation_chainstatemanager_tests.cpp
@@ -537,7 +537,7 @@ BOOST_FIXTURE_TEST_CASE(chainstatemanager_snapshot_init, SnapshotTestSetup)
     // it will initialize instead of attempting to complete validation.
     //
     // Note that this is not a realistic use of DisconnectTip().
-    DisconnectedBlockTransactions unused_pool;
+    DisconnectedBlockTransactions unused_pool{MAX_DISCONNECTED_TX_POOL_SIZE * 1000};
     BlockValidationState unused_state;
     {
         LOCK2(::cs_main, bg_chainstate.MempoolMutex());
diff --git a/src/validation.cpp b/src/validation.cpp
index 84846add80..a9a0912d2a 100644
--- a/src/validation.cpp
+++ b/src/validation.cpp
@@ -80,8 +80,6 @@ using node::CBlockIndexWorkComparator;
 using node::fReindex;
 using node::SnapshotMetadata;
 
-/** Maximum kilobytes for transactions to store for processing during reorg */
-static const unsigned int MAX_DISCONNECTED_TX_POOL_SIZE = 20000;
 /** Time to wait between writing blocks/block index to disk. */
 static constexpr std::chrono::hours DATABASE_WRITE_INTERVAL{1};
 /** Time to wait between flushing chainstate to disk. */
@@ -2723,12 +2721,10 @@ bool Chainstate::DisconnectTip(BlockValidationState& state, DisconnectedBlockTra
     }
 
     if (disconnectpool && m_mempool) {
-        // Save transactions to re-add to mempool at end of reorg
-        disconnectpool->AddTransactionsFromBlock(block.vtx);
-        while (disconnectpool->DynamicMemoryUsage() > MAX_DISCONNECTED_TX_POOL_SIZE * 1000) {
-            // Drop the earliest entry, and remove its children from the mempool.
-            auto ptx = disconnectpool->take_first();
-            m_mempool->removeRecursive(*ptx, MemPoolRemovalReason::REORG);
+        // Save transactions to re-add to mempool at end of reorg. If any entries are evicted for
+        // exceeding memory limits, remove them and their descendants from the mempool.
+        for (auto&& evicted_tx : disconnectpool->AddTransactionsFromBlock(block.vtx)) {
+            m_mempool->removeRecursive(*evicted_tx, MemPoolRemovalReason::REORG);
         }
     }
 
@@ -2978,7 +2974,7 @@ bool Chainstate::ActivateBestChainStep(BlockValidationState& state, CBlockIndex*
 
     // Disconnect active blocks which are no longer in the best chain.
     bool fBlocksDisconnected = false;
-    DisconnectedBlockTransactions disconnectpool;
+    DisconnectedBlockTransactions disconnectpool{MAX_DISCONNECTED_TX_POOL_SIZE * 1000};
     while (m_chain.Tip() && m_chain.Tip() != pindexFork) {
         if (!DisconnectTip(state, &disconnectpool)) {
             // This is likely a fatal error, but keep the mempool consistent,
@@ -3312,7 +3308,7 @@ bool Chainstate::InvalidateBlock(BlockValidationState& state, CBlockIndex* pinde
 
         // ActivateBestChain considers blocks already in m_chain
         // unconditionally valid already, so force disconnect away from it.
-        DisconnectedBlockTransactions disconnectpool;
+        DisconnectedBlockTransactions disconnectpool{MAX_DISCONNECTED_TX_POOL_SIZE * 1000};
         bool ret = DisconnectTip(state, &disconnectpool);
         // DisconnectTip will add transactions to disconnectpool.
         // Adjust the mempool to be consistent with the new tip, adding