From a27626cc224498e1b7524e248ee836858006b831 Mon Sep 17 00:00:00 2001
From: Qiwei Yang <yangqiwei97@gmail.com>
Date: Tue, 27 Feb 2024 18:31:54 +0800
Subject: [PATCH] feat: NetworkHandle get access to TransactionsHandle (#6780)

---
 crates/net/network/src/manager.rs          |  7 +++++++
 crates/net/network/src/network.rs          | 14 +++++++++++++-
 crates/net/network/src/transactions/mod.rs |  5 +++++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/crates/net/network/src/manager.rs b/crates/net/network/src/manager.rs
index 7b1377f4b8..4d31876ce4 100644
--- a/crates/net/network/src/manager.rs
+++ b/crates/net/network/src/manager.rs
@@ -593,6 +593,13 @@ where
                 let _ = tx.send(self.swarm.sessions().get_peer_infos_by_ids(peers));
             }
             NetworkHandleMessage::AddRlpxSubProtocol(proto) => self.add_rlpx_sub_protocol(proto),
+            NetworkHandleMessage::GetTransactionsHandle(tx) => {
+                if let Some(ref tx_inner) = self.to_transactions_manager {
+                    let _ = tx_inner.send(NetworkTransactionEvent::GetTransactionsHandle(tx));
+                } else {
+                    let _ = tx.send(None);
+                }
+            }
         }
     }
 }
diff --git a/crates/net/network/src/network.rs b/crates/net/network/src/network.rs
index 903fd99d20..d93d0b3fdf 100644
--- a/crates/net/network/src/network.rs
+++ b/crates/net/network/src/network.rs
@@ -1,6 +1,7 @@
 use crate::{
     config::NetworkMode, discovery::DiscoveryEvent, manager::NetworkEvent, message::PeerRequest,
-    peers::PeersHandle, protocol::RlpxSubProtocol, swarm::NetworkConnectionState, FetchClient,
+    peers::PeersHandle, protocol::RlpxSubProtocol, swarm::NetworkConnectionState,
+    transactions::TransactionsHandle, FetchClient,
 };
 use parking_lot::Mutex;
 use reth_eth_wire::{DisconnectReason, NewBlock, NewPooledTransactionHashes, SharedTransactions};
@@ -136,6 +137,15 @@ impl NetworkHandle {
         })
     }
 
+    /// Send message to get the [`TransactionsHandle`].
+    ///
+    /// Returns `None` if no transaction task is installed.
+    pub async fn transactions_handle(&self) -> Option<TransactionsHandle> {
+        let (tx, rx) = oneshot::channel();
+        let _ = self.manager().send(NetworkHandleMessage::GetTransactionsHandle(tx));
+        rx.await.unwrap()
+    }
+
     /// Provides a shareable reference to the [`BandwidthMeter`] stored on the `NetworkInner`.
     pub fn bandwidth_meter(&self) -> &BandwidthMeter {
         &self.inner.bandwidth_meter
@@ -446,6 +456,8 @@ pub(crate) enum NetworkHandleMessage {
     GetPeerInfosByPeerKind(PeerKind, oneshot::Sender<Vec<PeerInfo>>),
     /// Gets the reputation for a specific peer via a oneshot sender.
     GetReputationById(PeerId, oneshot::Sender<Option<Reputation>>),
+    /// Retrieves the `TransactionsHandle` via a oneshot sender.
+    GetTransactionsHandle(oneshot::Sender<Option<TransactionsHandle>>),
     /// Initiates a graceful shutdown of the network via a oneshot sender.
     Shutdown(oneshot::Sender<()>),
     /// Sets the network state between hibernation and active.
diff --git a/crates/net/network/src/transactions/mod.rs b/crates/net/network/src/transactions/mod.rs
index b9d354c0f4..9d543544f1 100644
--- a/crates/net/network/src/transactions/mod.rs
+++ b/crates/net/network/src/transactions/mod.rs
@@ -869,6 +869,9 @@ where
             NetworkTransactionEvent::GetPooledTransactions { peer_id, request, response } => {
                 self.on_get_pooled_transactions(peer_id, request, response)
             }
+            NetworkTransactionEvent::GetTransactionsHandle(response) => {
+                let _ = response.send(Some(self.handle()));
+            }
         }
     }
 
@@ -1561,6 +1564,8 @@ pub enum NetworkTransactionEvent {
         /// The sender for responding to the request with a result of `PooledTransactions`.
         response: oneshot::Sender<RequestResult<PooledTransactions>>,
     },
+    /// Represents the event of receiving a `GetTransactionsHandle` request.
+    GetTransactionsHandle(oneshot::Sender<Option<TransactionsHandle>>),
 }
 
 /// Tracks stats about the [`TransactionsManager`].