From e5ea7daee01e0313d47625b482b3e37bd977a3e7 Mon Sep 17 00:00:00 2001
From: glozow <gloriajzhao@gmail.com>
Date: Mon, 27 Jan 2025 07:41:36 -0500
Subject: [PATCH] [txorphanage] add per-peer weight accounting

---
 src/txorphanage.cpp | 12 ++++++++++++
 src/txorphanage.h   | 15 +++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/src/txorphanage.cpp b/src/txorphanage.cpp
index 30355d7e527..5e6c743f6bd 100644
--- a/src/txorphanage.cpp
+++ b/src/txorphanage.cpp
@@ -43,6 +43,8 @@ bool TxOrphanage::AddTx(const CTransactionRef& tx, NodeId peer)
         m_outpoint_to_orphan_it[txin.prevout].insert(ret.first);
     }
     m_total_orphan_usage += sz;
+    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());
@@ -56,6 +58,8 @@ 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();
             LogDebug(BCLog::TXPACKAGES, "added peer=%d as announcer of orphan tx %s\n", peer, wtxid.ToString());
             return true;
         }
@@ -80,6 +84,13 @@ int TxOrphanage::EraseTx(const Wtxid& wtxid)
 
     const auto tx_size{it->second.GetUsage()};
     m_total_orphan_usage -= tx_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);
@@ -103,6 +114,7 @@ int TxOrphanage::EraseTx(const Wtxid& wtxid)
 
 void TxOrphanage::EraseForPeer(NodeId peer)
 {
+    // Zeroes out this peer's m_total_usage.
     m_peer_orphanage_info.erase(peer);
 
     int nErased = 0;
diff --git a/src/txorphanage.h b/src/txorphanage.h
index ece88a1c2c1..9d1937d239c 100644
--- a/src/txorphanage.h
+++ b/src/txorphanage.h
@@ -97,6 +97,14 @@ public:
      * 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;
+    }
+
 protected:
     struct OrphanTx : public OrphanTxBase {
         size_t list_pos;
@@ -115,6 +123,13 @@ protected:
          * 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;