From fa7027d0fc1fb2eb4148ba9741e1736f61d7e164 Mon Sep 17 00:00:00 2001
From: glozow <gloriajzhao@gmail.com>
Date: Fri, 26 Apr 2024 11:56:24 +0100
Subject: [PATCH] [refactor] add CheckIsEmpty and GetOrphanTransactions, remove
 access to TxDownloadMan internals

---
 src/net_processing.cpp          | 11 +++--------
 src/node/txdownloadman.h        | 16 ++++++++++------
 src/node/txdownloadman_impl.cpp | 34 +++++++++++++++++++++++++--------
 src/node/txdownloadman_impl.h   |  5 +++++
 4 files changed, 44 insertions(+), 22 deletions(-)

diff --git a/src/net_processing.cpp b/src/net_processing.cpp
index bbd24d80928..063dd441c9f 100644
--- a/src/net_processing.cpp
+++ b/src/net_processing.cpp
@@ -1527,10 +1527,7 @@ void PeerManagerImpl::InitializeNode(const CNode& node, ServiceFlags our_service
         LOCK(cs_main); // For m_node_states
         m_node_states.try_emplace(m_node_states.end(), nodeid);
     }
-    {
-        LOCK(m_tx_download_mutex);
-        assert(m_txdownloadman.GetTxRequestRef().Count(nodeid) == 0);
-    }
+    WITH_LOCK(m_tx_download_mutex, m_txdownloadman.CheckIsEmpty(nodeid));
 
     if (NetPermissions::HasFlag(node.m_permission_flags, NetPermissionFlags::BloomFilter)) {
         our_services = static_cast<ServiceFlags>(our_services | NODE_BLOOM);
@@ -1616,9 +1613,7 @@ void PeerManagerImpl::FinalizeNode(const CNode& node)
         assert(m_peers_downloading_from == 0);
         assert(m_outbound_peers_with_protect_from_disconnect == 0);
         assert(m_wtxid_relay_peers == 0);
-        LOCK(m_tx_download_mutex);
-        assert(m_txdownloadman.GetTxRequestRef().Size() == 0);
-        assert(m_txdownloadman.GetOrphanageRef().Size() == 0);
+        WITH_LOCK(m_tx_download_mutex, m_txdownloadman.CheckIsEmpty());
     }
     } // cs_main
     if (node.fSuccessfullyConnected &&
@@ -1727,7 +1722,7 @@ bool PeerManagerImpl::GetNodeStateStats(NodeId nodeid, CNodeStateStats& stats) c
 std::vector<TxOrphanage::OrphanTxBase> PeerManagerImpl::GetOrphanTransactions()
 {
     LOCK(m_tx_download_mutex);
-    return m_txdownloadman.GetOrphanageRef().GetOrphanTransactions();
+    return m_txdownloadman.GetOrphanTransactions();
 }
 
 PeerManagerInfo PeerManagerImpl::GetInfo() const
diff --git a/src/node/txdownloadman.h b/src/node/txdownloadman.h
index aaff966cd56..c2a76b3a4f2 100644
--- a/src/node/txdownloadman.h
+++ b/src/node/txdownloadman.h
@@ -7,6 +7,7 @@
 
 #include <net.h>
 #include <policy/packages.h>
+#include <txorphanage.h>
 
 #include <cstdint>
 #include <memory>
@@ -15,7 +16,6 @@ class CBlock;
 class CRollingBloomFilter;
 class CTxMemPool;
 class GenTxid;
-class TxOrphanage;
 class TxRequestTracker;
 namespace node {
 class TxDownloadManagerImpl;
@@ -121,11 +121,6 @@ public:
     explicit TxDownloadManager(const TxDownloadOptions& options);
     ~TxDownloadManager();
 
-    // Get references to internal data structures. Outside access to these data structures should be
-    // temporary and removed later once logic has been moved internally.
-    TxOrphanage& GetOrphanageRef();
-    TxRequestTracker& GetTxRequestRef();
-
     // Responses to chain events. TxDownloadManager is not an actual client of ValidationInterface, these are called through PeerManager.
     void ActiveTipChange();
     void BlockConnected(const std::shared_ptr<const CBlock>& pblock);
@@ -167,6 +162,15 @@ public:
 
     /** Returns next orphan tx to consider, or nullptr if none exist. */
     CTransactionRef GetTxToReconsider(NodeId nodeid);
+
+    /** Check that all data structures are empty. */
+    void CheckIsEmpty() const;
+
+    /** Check that all data structures that track per-peer information have nothing for this peer. */
+    void CheckIsEmpty(NodeId nodeid) const;
+
+    /** Wrapper for TxOrphanage::GetOrphanTransactions */
+    std::vector<TxOrphanage::OrphanTxBase> GetOrphanTransactions() const;
 };
 } // namespace node
 #endif // BITCOIN_NODE_TXDOWNLOADMAN_H
diff --git a/src/node/txdownloadman_impl.cpp b/src/node/txdownloadman_impl.cpp
index 935ec2d6b74..89125982dd6 100644
--- a/src/node/txdownloadman_impl.cpp
+++ b/src/node/txdownloadman_impl.cpp
@@ -19,14 +19,6 @@ TxDownloadManager::TxDownloadManager(const TxDownloadOptions& options) :
 {}
 TxDownloadManager::~TxDownloadManager() = default;
 
-TxOrphanage& TxDownloadManager::GetOrphanageRef()
-{
-    return m_impl->m_orphanage;
-}
-TxRequestTracker& TxDownloadManager::GetTxRequestRef()
-{
-    return m_impl->m_txrequest;
-}
 void TxDownloadManager::ActiveTipChange()
 {
     m_impl->ActiveTipChange();
@@ -83,6 +75,18 @@ CTransactionRef TxDownloadManager::GetTxToReconsider(NodeId nodeid)
 {
     return m_impl->GetTxToReconsider(nodeid);
 }
+void TxDownloadManager::CheckIsEmpty() const
+{
+    m_impl->CheckIsEmpty();
+}
+void TxDownloadManager::CheckIsEmpty(NodeId nodeid) const
+{
+    m_impl->CheckIsEmpty(nodeid);
+}
+std::vector<TxOrphanage::OrphanTxBase> TxDownloadManager::GetOrphanTransactions() const
+{
+    return m_impl->GetOrphanTransactions();
+}
 
 // TxDownloadManagerImpl
 void TxDownloadManagerImpl::ActiveTipChange()
@@ -515,4 +519,18 @@ CTransactionRef TxDownloadManagerImpl::GetTxToReconsider(NodeId nodeid)
     return m_orphanage.GetTxToReconsider(nodeid);
 }
 
+void TxDownloadManagerImpl::CheckIsEmpty(NodeId nodeid)
+{
+    assert(m_txrequest.Count(nodeid) == 0);
+}
+void TxDownloadManagerImpl::CheckIsEmpty()
+{
+    assert(m_orphanage.Size() == 0);
+    assert(m_txrequest.Size() == 0);
+    assert(m_num_wtxid_peers == 0);
+}
+std::vector<TxOrphanage::OrphanTxBase> TxDownloadManagerImpl::GetOrphanTransactions() const
+{
+    return m_orphanage.GetOrphanTransactions();
+}
 } // namespace node
diff --git a/src/node/txdownloadman_impl.h b/src/node/txdownloadman_impl.h
index 85841777d05..c3e2417aa54 100644
--- a/src/node/txdownloadman_impl.h
+++ b/src/node/txdownloadman_impl.h
@@ -182,6 +182,11 @@ public:
 
     bool HaveMoreWork(NodeId nodeid);
     CTransactionRef GetTxToReconsider(NodeId nodeid);
+
+    void CheckIsEmpty();
+    void CheckIsEmpty(NodeId nodeid);
+
+    std::vector<TxOrphanage::OrphanTxBase> GetOrphanTransactions() const;
 };
 } // namespace node
 #endif // BITCOIN_NODE_TXDOWNLOADMAN_IMPL_H