From f48d36cd97e9b27dfa105c35e0fe67cba47056d1 Mon Sep 17 00:00:00 2001
From: glozow <gloriajzhao@gmail.com>
Date: Tue, 16 Apr 2024 16:46:36 +0100
Subject: [PATCH] [refactor] move peer (dis)connection logic to TxDownload

The information stored in TxDownloadConnectionInfo isn't used until the
next commit.
---
 src/net_processing.cpp          | 13 +++++++++++--
 src/node/txdownloadman.h        | 16 ++++++++++++++++
 src/node/txdownloadman_impl.cpp | 29 +++++++++++++++++++++++++++++
 src/node/txdownloadman_impl.h   | 17 +++++++++++++++++
 4 files changed, 73 insertions(+), 2 deletions(-)

diff --git a/src/net_processing.cpp b/src/net_processing.cpp
index 13f9c8a2cc2..a0f0ca10f95 100644
--- a/src/net_processing.cpp
+++ b/src/net_processing.cpp
@@ -1686,8 +1686,7 @@ void PeerManagerImpl::FinalizeNode(const CNode& node)
     }
     {
         LOCK(m_tx_download_mutex);
-        m_txdownloadman.GetOrphanageRef().EraseForPeer(nodeid);
-        m_txdownloadman.GetTxRequestRef().DisconnectedPeer(nodeid);
+        m_txdownloadman.DisconnectedPeer(nodeid);
     }
     if (m_txreconciliation) m_txreconciliation->ForgetPeer(nodeid);
     m_num_preferred_download_peers -= state->fPreferredDownload;
@@ -3852,6 +3851,16 @@ void PeerManagerImpl::ProcessMessage(CNode& pfrom, const std::string& msg_type,
                        tx_relay->m_next_inv_send_time == 0s));
         }
 
+        {
+            LOCK2(::cs_main, m_tx_download_mutex);
+            const CNodeState* state = State(pfrom.GetId());
+            m_txdownloadman.ConnectedPeer(pfrom.GetId(), node::TxDownloadConnectionInfo {
+                .m_preferred = state->fPreferredDownload,
+                .m_relay_permissions = pfrom.HasPermission(NetPermissionFlags::Relay),
+                .m_wtxid_relay = peer->m_wtxid_relay,
+            });
+        }
+
         pfrom.fSuccessfullyConnected = true;
         return;
     }
diff --git a/src/node/txdownloadman.h b/src/node/txdownloadman.h
index f573bc54d4b..93c35b4d24a 100644
--- a/src/node/txdownloadman.h
+++ b/src/node/txdownloadman.h
@@ -5,6 +5,8 @@
 #ifndef BITCOIN_NODE_TXDOWNLOADMAN_H
 #define BITCOIN_NODE_TXDOWNLOADMAN_H
 
+#include <net.h>
+
 #include <cstdint>
 #include <memory>
 
@@ -21,6 +23,14 @@ struct TxDownloadOptions {
     /** Read-only reference to mempool. */
     const CTxMemPool& m_mempool;
 };
+struct TxDownloadConnectionInfo {
+    /** Whether this peer is preferred for transaction download. */
+    const bool m_preferred;
+    /** Whether this peer has Relay permissions. */
+    const bool m_relay_permissions;
+    /** Whether this peer supports wtxid relay. */
+    const bool m_wtxid_relay;
+};
 
 /**
  * Class responsible for deciding what transactions to request and, once
@@ -68,6 +78,12 @@ public:
      *  - m_recent_confirmed_transactions
      *  */
     bool AlreadyHaveTx(const GenTxid& gtxid, bool include_reconsiderable);
+
+    /** Creates a new PeerInfo. Saves the connection info to calculate tx announcement delays later. */
+    void ConnectedPeer(NodeId nodeid, const TxDownloadConnectionInfo& info);
+
+    /** Deletes all txrequest announcements and orphans for a given peer. */
+    void DisconnectedPeer(NodeId nodeid);
 };
 } // namespace node
 #endif // BITCOIN_NODE_TXDOWNLOADMAN_H
diff --git a/src/node/txdownloadman_impl.cpp b/src/node/txdownloadman_impl.cpp
index 6d417ac65b6..6a00e1c70f2 100644
--- a/src/node/txdownloadman_impl.cpp
+++ b/src/node/txdownloadman_impl.cpp
@@ -50,6 +50,14 @@ bool TxDownloadManager::AlreadyHaveTx(const GenTxid& gtxid, bool include_reconsi
 {
     return m_impl->AlreadyHaveTx(gtxid, include_reconsiderable);
 }
+void TxDownloadManager::ConnectedPeer(NodeId nodeid, const TxDownloadConnectionInfo& info)
+{
+    m_impl->ConnectedPeer(nodeid, info);
+}
+void TxDownloadManager::DisconnectedPeer(NodeId nodeid)
+{
+    m_impl->DisconnectedPeer(nodeid);
+}
 
 // TxDownloadManagerImpl
 void TxDownloadManagerImpl::ActiveTipChange()
@@ -113,4 +121,25 @@ bool TxDownloadManagerImpl::AlreadyHaveTx(const GenTxid& gtxid, bool include_rec
 
     return RecentRejectsFilter().contains(hash) || m_opts.m_mempool.exists(gtxid);
 }
+
+void TxDownloadManagerImpl::ConnectedPeer(NodeId nodeid, const TxDownloadConnectionInfo& info)
+{
+    // If already connected (shouldn't happen in practice), exit early.
+    if (m_peer_info.contains(nodeid)) return;
+
+    m_peer_info.try_emplace(nodeid, info);
+    if (info.m_wtxid_relay) m_num_wtxid_peers += 1;
+}
+
+void TxDownloadManagerImpl::DisconnectedPeer(NodeId nodeid)
+{
+    m_orphanage.EraseForPeer(nodeid);
+    m_txrequest.DisconnectedPeer(nodeid);
+
+    if (auto it = m_peer_info.find(nodeid); it != m_peer_info.end()) {
+        if (it->second.m_connection_info.m_wtxid_relay) m_num_wtxid_peers -= 1;
+        m_peer_info.erase(it);
+    }
+
+}
 } // namespace node
diff --git a/src/node/txdownloadman_impl.h b/src/node/txdownloadman_impl.h
index ad046bef9e7..fc57f690240 100644
--- a/src/node/txdownloadman_impl.h
+++ b/src/node/txdownloadman_impl.h
@@ -128,11 +128,28 @@ public:
 
     TxDownloadManagerImpl(const TxDownloadOptions& options) : m_opts{options} {}
 
+    struct PeerInfo {
+        /** Information relevant to scheduling tx requests. */
+        const TxDownloadConnectionInfo m_connection_info;
+
+        PeerInfo(const TxDownloadConnectionInfo& info) : m_connection_info{info} {}
+    };
+
+    /** Information for all of the peers we may download transactions from. This is not necessarily
+     * all peers we are connected to (no block-relay-only and temporary connections). */
+    std::map<NodeId, PeerInfo> m_peer_info;
+
+    /** Number of wtxid relay peers we have in m_peer_info. */
+    uint32_t m_num_wtxid_peers{0};
+
     void ActiveTipChange();
     void BlockConnected(const std::shared_ptr<const CBlock>& pblock);
     void BlockDisconnected();
 
     bool AlreadyHaveTx(const GenTxid& gtxid, bool include_reconsiderable);
+
+    void ConnectedPeer(NodeId nodeid, const TxDownloadConnectionInfo& info);
+    void DisconnectedPeer(NodeId nodeid);
 };
 } // namespace node
 #endif // BITCOIN_NODE_TXDOWNLOADMAN_IMPL_H