Merge bitcoin/bitcoin#31810: TxOrphanage: account for size of orphans and count announcements

e107bf7 [fuzz] TxOrphanage::SanityCheck accounting (glozow)
22dccea [fuzz] txorphan byte accounting (glozow)
982ce10 add orphanage byte accounting to TxDownloadManagerImpl::CheckIsEmpty() (glozow)
c289217 [txorphanage] track the total number of announcements (glozow)
e5ea7da [txorphanage] add per-peer weight accounting (glozow)
672c69c [refactor] change per-peer workset to info map within orphanage (glozow)
59cd0f0 [txorphanage] account for weight of orphans (glozow)

Pull request description:

  Part of orphan resolution project, see #27463.

  Definitions:
  - **Announcement** is a unique pair (wtxid, nodeid). We can have multiple announcers for the same orphan since #31397.
  - **Size** is the weight of an orphan. I'm calling it "size" and "bytes" because I think we can refine it in the future to be memusage or be otherwise more representative of the orphan's actual cost on our memory. However, I am open to naming changes.

  This is part 1/2 of a project to also add limits on orphan size and count. However, this PR **does not change behavior**, just adds internal counters/tracking and a fuzzer. I will also open a second PR that adds behavior changes, which requires updating a lot of our tests and careful thinking about DoS.

ACKs for top commit:
  instagibbs:
    reACK bitcoin/bitcoin@e107bf7
  marcofleon:
    reACK e107bf7
  sipa:
    utACK e107bf7

Tree-SHA512: 855d725d5eb521d131e36dacc51990725e3ca7881beb13364d5ba72ab2202bbfd14ab83864b13b1b945a4ec5e17890458d0112270b891a41b1e27324a8545d72

Author: fanquake

src/node/txdownloadman_impl.cpp

@@ -578,9 +578,11 @@ CTransactionRef TxDownloadManagerImpl::GetTxToReconsider(NodeId nodeid)
 void TxDownloadManagerImpl::CheckIsEmpty(NodeId nodeid)
 {
     assert(m_txrequest.Count(nodeid) == 0);
+    assert(m_orphanage.UsageByPeer(nodeid) == 0);
 }
 void TxDownloadManagerImpl::CheckIsEmpty()
 {
+    assert(m_orphanage.TotalOrphanUsage() == 0);
     assert(m_orphanage.Size() == 0);
     assert(m_txrequest.Size() == 0);
     assert(m_num_wtxid_peers == 0);

src/test/fuzz/txdownloadman.cpp

@@ -285,6 +285,7 @@ static void CheckInvariants(const node::TxDownloadManagerImpl& txdownload_impl,
 
     // Orphanage usage should never exceed what is allowed
     Assert(orphanage.Size() <= max_orphan_count);
+    txdownload_impl.m_orphanage.SanityCheck();
 
     // We should never have more than the maximum in-flight requests out for a peer.
     for (NodeId peer = 0; peer < NUM_PEERS; ++peer) {
@@ -437,8 +438,8 @@ FUZZ_TARGET(txdownloadman_impl, .init = initialize)
         auto time_skip = fuzzed_data_provider.PickValueInArray(TIME_SKIPS);
         if (fuzzed_data_provider.ConsumeBool()) time_skip *= -1;
         time += time_skip;
-        CheckInvariants(txdownload_impl, max_orphan_count);
     }
+    CheckInvariants(txdownload_impl, max_orphan_count);
     // Disconnect everybody, check that all data structures are empty.
     for (NodeId nodeid = 0; nodeid < NUM_PEERS; ++nodeid) {
         txdownload_impl.DisconnectedPeer(nodeid);

src/test/fuzz/txorphan.cpp

@@ -75,6 +75,8 @@ FUZZ_TARGET(txorphan, .init = initialize_orphanage)
             return new_tx;
         }();
 
+        const auto wtxid{tx->GetWitnessHash()};
+
         // Trigger orphanage functions that are called using parents. ptx_potential_parent is a tx we constructed in a
         // previous loop and potentially the parent of this tx.
         if (ptx_potential_parent) {
@@ -94,6 +96,9 @@ FUZZ_TARGET(txorphan, .init = initialize_orphanage)
         LIMITED_WHILE(fuzzed_data_provider.ConsumeBool(), 10 * DEFAULT_MAX_ORPHAN_TRANSACTIONS)
         {
             NodeId peer_id = fuzzed_data_provider.ConsumeIntegral<NodeId>();
+            const auto total_bytes_start{orphanage.TotalOrphanUsage()};
+            const auto total_peer_bytes_start{orphanage.UsageByPeer(peer_id)};
+            const auto tx_weight{GetTransactionWeight(*tx)};
 
             CallOneOf(
                 fuzzed_data_provider,
@@ -113,12 +118,29 @@ FUZZ_TARGET(txorphan, .init = initialize_orphanage)
                         bool add_tx = orphanage.AddTx(tx, peer_id);
                         // have_tx == true -> add_tx == false
                         Assert(!have_tx || !add_tx);
+
+                        if (add_tx) {
+                            Assert(orphanage.UsageByPeer(peer_id) == tx_weight + total_peer_bytes_start);
+                            Assert(orphanage.TotalOrphanUsage() == tx_weight + total_bytes_start);
+                            Assert(tx_weight <= MAX_STANDARD_TX_WEIGHT);
+                        } else {
+                            // Peer may have been added as an announcer.
+                            if (orphanage.UsageByPeer(peer_id) == tx_weight + total_peer_bytes_start) {
+                                Assert(orphanage.HaveTxFromPeer(wtxid, peer_id));
+                            } else {
+                                // Otherwise, there must not be any change to the peer byte count.
+                                Assert(orphanage.UsageByPeer(peer_id) == total_peer_bytes_start);
+                            }
+
+                            // Regardless, total bytes should not have changed.
+                            Assert(orphanage.TotalOrphanUsage() == total_bytes_start);
+                        }
                     }
                     have_tx = orphanage.HaveTx(tx->GetWitnessHash());
                     {
                         bool add_tx = orphanage.AddTx(tx, peer_id);
                         // if have_tx is still false, it must be too big
-                        Assert(!have_tx == (GetTransactionWeight(*tx) > MAX_STANDARD_TX_WEIGHT));
+                        Assert(!have_tx == (tx_weight > MAX_STANDARD_TX_WEIGHT));
                         Assert(!have_tx || !add_tx);
                     }
                 },
@@ -132,23 +154,46 @@ FUZZ_TARGET(txorphan, .init = initialize_orphanage)
                         Assert(have_tx || !added_announcer);
                         // have_tx_and_peer == true -> added_announcer == false
                         Assert(!have_tx_and_peer || !added_announcer);
+
+                        // Total bytes should not have changed. If peer was added as announcer, byte
+                        // accounting must have been updated.
+                        Assert(orphanage.TotalOrphanUsage() == total_bytes_start);
+                        if (added_announcer) {
+                            Assert(orphanage.UsageByPeer(peer_id) == tx_weight + total_peer_bytes_start);
+                        } else {
+                            Assert(orphanage.UsageByPeer(peer_id) == total_peer_bytes_start);
+                        }
                     }
                 },
                 [&] {
                     bool have_tx = orphanage.HaveTx(tx->GetWitnessHash());
+                    bool have_tx_and_peer{orphanage.HaveTxFromPeer(wtxid, peer_id)};
                     // EraseTx should return 0 if m_orphans doesn't have the tx
                     {
+                        auto bytes_from_peer_before{orphanage.UsageByPeer(peer_id)};
                         Assert(have_tx == orphanage.EraseTx(tx->GetWitnessHash()));
+                        if (have_tx) {
+                            Assert(orphanage.TotalOrphanUsage() == total_bytes_start - tx_weight);
+                            if (have_tx_and_peer) {
+                                Assert(orphanage.UsageByPeer(peer_id) == bytes_from_peer_before - tx_weight);
+                            } else {
+                                Assert(orphanage.UsageByPeer(peer_id) == bytes_from_peer_before);
+                            }
+                        } else {
+                            Assert(orphanage.TotalOrphanUsage() == total_bytes_start);
+                        }
                     }
                     have_tx = orphanage.HaveTx(tx->GetWitnessHash());
+                    have_tx_and_peer = orphanage.HaveTxFromPeer(wtxid, peer_id);
                     // have_tx should be false and EraseTx should fail
                     {
-                        Assert(!have_tx && !orphanage.EraseTx(tx->GetWitnessHash()));
+                        Assert(!have_tx && !have_tx_and_peer && !orphanage.EraseTx(wtxid));
                     }
                 },
                 [&] {
                     orphanage.EraseForPeer(peer_id);
                     Assert(!orphanage.HaveTxFromPeer(tx->GetWitnessHash(), peer_id));
+                    Assert(orphanage.UsageByPeer(peer_id) == 0);
                 },
                 [&] {
                     // test mocktime and expiry
@@ -159,6 +204,7 @@ FUZZ_TARGET(txorphan, .init = initialize_orphanage)
                 });
 
         }
+
         // Set tx as potential parent to be used for future GetChildren() calls.
         if (!ptx_potential_parent || fuzzed_data_provider.ConsumeBool()) {
             ptx_potential_parent = tx;
@@ -168,4 +214,5 @@ FUZZ_TARGET(txorphan, .init = initialize_orphanage)
         const bool get_tx_nonnull{orphanage.GetTx(tx->GetWitnessHash()) != nullptr};
         Assert(have_tx == get_tx_nonnull);
     }
+    orphanage.SanityCheck();
 }

src/txorphanage.cpp

@@ -42,6 +42,10 @@ bool TxOrphanage::AddTx(const CTransactionRef& tx, NodeId peer)
     for (const CTxIn& txin : tx->vin) {
         m_outpoint_to_orphan_it[txin.prevout].insert(ret.first);
     }
+    m_total_orphan_usage += sz;
+    m_total_announcements += 1;
+    auto& peer_info = m_peer_orphanage_info.try_emplace(peer).first->second;
+    peer_info.m_total_usage += sz;
 
     LogDebug(BCLog::TXPACKAGES, "stored orphan tx %s (wtxid=%s), weight: %u (mapsz %u outsz %u)\n", hash.ToString(), wtxid.ToString(), sz,
              m_orphans.size(), m_outpoint_to_orphan_it.size());
@@ -55,6 +59,9 @@ bool TxOrphanage::AddAnnouncer(const Wtxid& wtxid, NodeId peer)
         Assume(!it->second.announcers.empty());
         const auto ret = it->second.announcers.insert(peer);
         if (ret.second) {
+            auto& peer_info = m_peer_orphanage_info.try_emplace(peer).first->second;
+            peer_info.m_total_usage += it->second.GetUsage();
+            m_total_announcements += 1;
             LogDebug(BCLog::TXPACKAGES, "added peer=%d as announcer of orphan tx %s\n", peer, wtxid.ToString());
             return true;
         }
@@ -77,6 +84,17 @@ int TxOrphanage::EraseTx(const Wtxid& wtxid)
             m_outpoint_to_orphan_it.erase(itPrev);
     }
 
+    const auto tx_size{it->second.GetUsage()};
+    m_total_orphan_usage -= tx_size;
+    m_total_announcements -= it->second.announcers.size();
+    // Decrement each announcer's m_total_usage
+    for (const auto& peer : it->second.announcers) {
+        auto peer_it = m_peer_orphanage_info.find(peer);
+        if (Assume(peer_it != m_peer_orphanage_info.end())) {
+            peer_it->second.m_total_usage -= tx_size;
+        }
+    }
+
     size_t old_pos = it->second.list_pos;
     assert(m_orphan_list[old_pos] == it);
     if (old_pos + 1 != m_orphan_list.size()) {
@@ -99,7 +117,8 @@ int TxOrphanage::EraseTx(const Wtxid& wtxid)
 
 void TxOrphanage::EraseForPeer(NodeId peer)
 {
-    m_peer_work_set.erase(peer);
+    // Zeroes out this peer's m_total_usage.
+    m_peer_orphanage_info.erase(peer);
 
     int nErased = 0;
     std::map<Wtxid, OrphanTx>::iterator iter = m_orphans.begin();
@@ -110,6 +129,7 @@ void TxOrphanage::EraseForPeer(NodeId peer)
         auto orphan_it = orphan.announcers.find(peer);
         if (orphan_it != orphan.announcers.end()) {
             orphan.announcers.erase(peer);
+            m_total_announcements -= 1;
 
             // No remaining announcers: clean up entry
             if (orphan.announcers.empty()) {
@@ -170,7 +190,7 @@ void TxOrphanage::AddChildrenToWorkSet(const CTransaction& tx, FastRandomContext
 
                 // Get this source peer's work set, emplacing an empty set if it didn't exist
                 // (note: if this peer wasn't still connected, we would have removed the orphan tx already)
-                std::set<Wtxid>& orphan_work_set = m_peer_work_set.try_emplace(announcer).first->second;
+                std::set<Wtxid>& orphan_work_set = m_peer_orphanage_info.try_emplace(announcer).first->second.m_work_set;
                 // Add this tx to the work set
                 orphan_work_set.insert(elem->first);
                 LogDebug(BCLog::TXPACKAGES, "added %s (wtxid=%s) to peer %d workset\n",
@@ -199,30 +219,29 @@ bool TxOrphanage::HaveTxFromPeer(const Wtxid& wtxid, NodeId peer) const
 
 CTransactionRef TxOrphanage::GetTxToReconsider(NodeId peer)
 {
-    auto work_set_it = m_peer_work_set.find(peer);
-    if (work_set_it != m_peer_work_set.end()) {
-        auto& work_set = work_set_it->second;
-        while (!work_set.empty()) {
-            Wtxid wtxid = *work_set.begin();
-            work_set.erase(work_set.begin());
-
-            const auto orphan_it = m_orphans.find(wtxid);
-            if (orphan_it != m_orphans.end()) {
-                return orphan_it->second.tx;
-            }
+    auto peer_it = m_peer_orphanage_info.find(peer);
+    if (peer_it == m_peer_orphanage_info.end()) return nullptr;
+
+    auto& work_set = peer_it->second.m_work_set;
+    while (!work_set.empty()) {
+        Wtxid wtxid = *work_set.begin();
+        work_set.erase(work_set.begin());
+
+        const auto orphan_it = m_orphans.find(wtxid);
+        if (orphan_it != m_orphans.end()) {
+            return orphan_it->second.tx;
         }
     }
     return nullptr;
 }
 
 bool TxOrphanage::HaveTxToReconsider(NodeId peer)
 {
-    auto work_set_it = m_peer_work_set.find(peer);
-    if (work_set_it != m_peer_work_set.end()) {
-        auto& work_set = work_set_it->second;
-        return !work_set.empty();
-    }
-    return false;
+    auto peer_it = m_peer_orphanage_info.find(peer);
+    if (peer_it == m_peer_orphanage_info.end()) return false;
+
+    auto& work_set = peer_it->second.m_work_set;
+    return !work_set.empty();
 }
 
 void TxOrphanage::EraseForBlock(const CBlock& block)
@@ -302,3 +321,43 @@ std::vector<TxOrphanage::OrphanTxBase> TxOrphanage::GetOrphanTransactions() cons
     }
     return ret;
 }
+
+void TxOrphanage::SanityCheck() const
+{
+    // Check that cached m_total_announcements is correct
+    unsigned int counted_total_announcements{0};
+    // Check that m_total_orphan_usage is correct
+    unsigned int counted_total_usage{0};
+
+    // Check that cached PeerOrphanInfo::m_total_size is correct
+    std::map<NodeId, unsigned int> counted_size_per_peer;
+
+    for (const auto& [wtxid, orphan] : m_orphans) {
+        counted_total_announcements += orphan.announcers.size();
+        counted_total_usage += orphan.GetUsage();
+
+        Assume(!orphan.announcers.empty());
+        for (const auto& peer : orphan.announcers) {
+            auto& count_peer_entry = counted_size_per_peer.try_emplace(peer).first->second;
+            count_peer_entry += orphan.GetUsage();
+        }
+    }
+
+    Assume(m_total_announcements >= m_orphans.size());
+    Assume(counted_total_announcements == m_total_announcements);
+    Assume(counted_total_usage == m_total_orphan_usage);
+
+    // There must be an entry in m_peer_orphanage_info for each peer
+    // However, there may be m_peer_orphanage_info entries corresponding to peers for whom we
+    // previously had orphans but no longer do.
+    Assume(counted_size_per_peer.size() <= m_peer_orphanage_info.size());
+
+    for (const auto& [peerid, info] : m_peer_orphanage_info) {
+        auto it_counted = counted_size_per_peer.find(peerid);
+        if (it_counted == counted_size_per_peer.end()) {
+            Assume(info.m_total_usage == 0);
+        } else {
+            Assume(it_counted->second == info.m_total_usage);
+        }
+    }
+}

src/txorphanage.h

@@ -5,6 +5,7 @@
 #ifndef BITCOIN_TXORPHANAGE_H
 #define BITCOIN_TXORPHANAGE_H
 
+#include <consensus/validation.h>
 #include <net.h>
 #include <primitives/block.h>
 #include <primitives/transaction.h>
@@ -83,21 +84,62 @@ class TxOrphanage {
         /** Peers added with AddTx or AddAnnouncer. */
         std::set<NodeId> announcers;
         NodeSeconds nTimeExpire;
+
+        /** Get the weight of this transaction, an approximation of its memory usage. */
+        unsigned int GetUsage() const {
+            return GetTransactionWeight(*tx);
+        }
     };
 
     std::vector<OrphanTxBase> GetOrphanTransactions() const;
 
+    /** Get the total usage (weight) of all orphans. If an orphan has multiple announcers, its usage is
+     * only counted once within this total. */
+    unsigned int TotalOrphanUsage() const { return m_total_orphan_usage; }
+
+    /** Total usage (weight) of orphans for which this peer is an announcer. If an orphan has multiple
+     * announcers, its weight will be accounted for in each PeerOrphanInfo, so the total of all
+     * peers' UsageByPeer() may be larger than TotalOrphanBytes(). */
+    unsigned int UsageByPeer(NodeId peer) const {
+        auto peer_it = m_peer_orphanage_info.find(peer);
+        return peer_it == m_peer_orphanage_info.end() ? 0 : peer_it->second.m_total_usage;
+    }
+
+    /** Check consistency between PeerOrphanInfo and m_orphans. Recalculate counters and ensure they
+     * match what is cached. */
+    void SanityCheck() const;
+
 protected:
     struct OrphanTx : public OrphanTxBase {
         size_t list_pos;
     };
 
+    /** Total usage (weight) of all entries in m_orphans. */
+    unsigned int m_total_orphan_usage{0};
+
+    /** Total number of <peer, tx> pairs. Can be larger than m_orphans.size() because multiple peers
+     * may have announced the same orphan. */
+    unsigned int m_total_announcements{0};
+
     /** Map from wtxid to orphan transaction record. Limited by
      *  -maxorphantx/DEFAULT_MAX_ORPHAN_TRANSACTIONS */
     std::map<Wtxid, OrphanTx> m_orphans;
 
-    /** Which peer provided the orphans that need to be reconsidered */
-    std::map<NodeId, std::set<Wtxid>> m_peer_work_set;
+    struct PeerOrphanInfo {
+        /** List of transactions that should be reconsidered: added to in AddChildrenToWorkSet,
+         * removed from one-by-one with each call to GetTxToReconsider. The wtxids may refer to
+         * transactions that are no longer present in orphanage; these are lazily removed in
+         * GetTxToReconsider. */
+        std::set<Wtxid> m_work_set;
+
+        /** Total weight of orphans for which this peer is an announcer.
+         * If orphans are provided by different peers, its weight will be accounted for in each
+         * PeerOrphanInfo, so the total of all peers' m_total_usage may be larger than
+         * m_total_orphan_size. If a peer is removed as an announcer, even if the orphan still
+         * remains in the orphanage, this number will be decremented. */
+        unsigned int m_total_usage{0};
+    };
+    std::map<NodeId, PeerOrphanInfo> m_peer_orphanage_info;
 
     using OrphanMap = decltype(m_orphans);