metrics reporting for tpu-client-next

Author: Alex Pyattaev

Cargo.lock

@@ -21,7 +21,7 @@ solana-measure = { workspace = true }
 solana-metrics = { workspace = true }
 solana-runtime = { workspace = true }
 solana-sdk = { workspace = true }
-solana-tpu-client-next = { workspace = true }
+solana-tpu-client-next = { workspace = true, features = ["metrics"] }
 tokio = { workspace = true, features = ["full"] }
 tokio-util = { workspace = true }
 

programs/sbf/Cargo.lock

@@ -614,7 +614,7 @@ mod test {
 
     #[tokio::test(flavor = "multi_thread")]
     async fn service_exit_with_tpu_client_next() {
-        service_exit::<TpuClientNextClient<NullTpuInfo>>(Some(Handle::current()));
+        service_exit::<TpuClientNextClient>(Some(Handle::current()));
     }
 
     fn validator_exit<C: ClientWithCreator>(maybe_runtime: Option<Handle>) {
@@ -665,7 +665,7 @@ mod test {
 
     #[tokio::test(flavor = "multi_thread")]
     async fn validator_exit_with_tpu_client_next() {
-        validator_exit::<TpuClientNextClient<NullTpuInfo>>(Some(Handle::current()));
+        validator_exit::<TpuClientNextClient>(Some(Handle::current()));
     }
 
     fn process_transactions<C: ClientWithCreator>(maybe_runtime: Option<Handle>) {
@@ -924,7 +924,7 @@ mod test {
 
     #[tokio::test(flavor = "multi_thread")]
     async fn process_transactions_with_tpu_client_next() {
-        process_transactions::<TpuClientNextClient<NullTpuInfo>>(Some(Handle::current()));
+        process_transactions::<TpuClientNextClient>(Some(Handle::current()));
     }
 
     fn retry_durable_nonce_transactions<C: ClientWithCreator>(maybe_runtime: Option<Handle>) {
@@ -1234,8 +1234,6 @@ mod test {
 
     #[tokio::test(flavor = "multi_thread")]
     async fn retry_durable_nonce_transactions_with_tpu_client_next() {
-        retry_durable_nonce_transactions::<TpuClientNextClient<NullTpuInfo>>(Some(
-            Handle::current(),
-        ));
+        retry_durable_nonce_transactions::<TpuClientNextClient>(Some(Handle::current()));
     }
 }

send-transaction-service/Cargo.toml

@@ -44,7 +44,7 @@ impl CreateClient for ConnectionCacheClient<NullTpuInfo> {
     }
 }
 
-impl CreateClient for TpuClientNextClient<NullTpuInfo> {
+impl CreateClient for TpuClientNextClient {
     fn create_client(
         maybe_runtime: Option<Handle>,
         my_tpu_address: SocketAddr,
@@ -53,7 +53,7 @@ impl CreateClient for TpuClientNextClient<NullTpuInfo> {
     ) -> Self {
         let runtime_handle =
             maybe_runtime.expect("Runtime should be provided for the TpuClientNextClient.");
-        Self::new(
+        Self::new::<NullTpuInfo>(
             runtime_handle,
             my_tpu_address,
             tpu_peers,
@@ -75,10 +75,7 @@ where
     fn stop(&self) {}
 }
 
-impl<T> Stoppable for TpuClientNextClient<T>
-where
-    T: TpuInfoWithSendStatic + Clone,
-{
+impl Stoppable for TpuClientNextClient {
     fn stop(&self) {
         self.cancel().unwrap();
     }

send-transaction-service/src/send_transaction_service.rs

@@ -8,9 +8,7 @@ use {
     solana_measure::measure::Measure,
     solana_sdk::quic::NotifyKeyUpdate,
     solana_tpu_client_next::{
-        connection_workers_scheduler::{
-            ConnectionWorkersSchedulerConfig, Fanout, TransactionStatsAndReceiver,
-        },
+        connection_workers_scheduler::{ConnectionWorkersSchedulerConfig, Fanout},
         leader_updater::LeaderUpdater,
         transaction_batch::TransactionBatch,
         ConnectionWorkersScheduler, ConnectionWorkersSchedulerError,
@@ -24,6 +22,7 @@ use {
         runtime::Handle,
         sync::mpsc::{self},
         task::JoinHandle as TokioJoinHandle,
+        time::interval,
     },
     tokio_util::sync::CancellationToken,
 };
@@ -230,27 +229,20 @@ where
 ///   scheduler. Most of the complexity of this structure arises from this
 ///   functionality.
 #[derive(Clone)]
-pub struct TpuClientNextClient<T>
-where
-    T: TpuInfoWithSendStatic + Clone,
-{
+pub struct TpuClientNextClient {
     runtime_handle: Handle,
     sender: mpsc::Sender<TransactionBatch>,
     // This handle is needed to implement `NotifyKeyUpdate` trait. It's only
     // method takes &self and thus we need to wrap with Mutex.
     join_and_cancel: Arc<Mutex<(Option<TpuClientJoinHandle>, CancellationToken)>>,
-    leader_updater: SendTransactionServiceLeaderUpdater<T>,
     leader_forward_count: u64,
 }
 
 type TpuClientJoinHandle =
-    TokioJoinHandle<Result<TransactionStatsAndReceiver, ConnectionWorkersSchedulerError>>;
+    TokioJoinHandle<Result<ConnectionWorkersScheduler, ConnectionWorkersSchedulerError>>;
 
-impl<T> TpuClientNextClient<T>
-where
-    T: TpuInfoWithSendStatic + Clone,
-{
-    pub fn new(
+impl TpuClientNextClient {
+    pub fn new<T>(
         runtime_handle: Handle,
         my_tpu_address: SocketAddr,
         tpu_peers: Option<Vec<SocketAddr>>,
@@ -275,18 +267,18 @@ where
                 tpu_peers,
             };
         let config = Self::create_config(identity, leader_forward_count as usize);
-        let handle = runtime_handle.spawn(ConnectionWorkersScheduler::run(
-            config,
-            Box::new(leader_updater.clone()),
-            receiver,
+        let scheduler = ConnectionWorkersScheduler::new(Box::new(leader_updater), receiver);
+        // leaking handle to this task, as it will run until the end of life of the process
+        runtime_handle.spawn(scheduler.stats.clone().report_to_influxdb(
+            "send-transaction-service-TPU-client",
+            interval(Duration::from_secs(3)),
             cancel.clone(),
         ));
-
+        let handle = runtime_handle.spawn(scheduler.run(config, cancel.clone()));
         Self {
             runtime_handle,
             join_and_cancel: Arc::new(Mutex::new((Some(handle), cancel))),
             sender,
-            leader_updater,
             leader_forward_count,
         }
     }
@@ -322,7 +314,6 @@ where
     async fn do_update_key(&self, identity: &Keypair) -> Result<(), Box<dyn std::error::Error>> {
         let runtime_handle = self.runtime_handle.clone();
         let config = Self::create_config(Some(identity), self.leader_forward_count as usize);
-        let leader_updater = self.leader_updater.clone();
         let handle = self.join_and_cancel.clone();
 
         let join_handle = {
@@ -340,14 +331,9 @@ where
             };
 
             match result {
-                Ok((_stats, receiver)) => {
+                Ok(scheduler) => {
                     let cancel = CancellationToken::new();
-                    let join_handle = runtime_handle.spawn(ConnectionWorkersScheduler::run(
-                        config,
-                        Box::new(leader_updater),
-                        receiver,
-                        cancel.clone(),
-                    ));
+                    let join_handle = runtime_handle.spawn(scheduler.run(config, cancel.clone()));
 
                     let Ok(mut lock) = handle.lock() else {
                         return Err("TpuClientNext task panicked.".into());
@@ -363,19 +349,13 @@ where
     }
 }
 
-impl<T> NotifyKeyUpdate for TpuClientNextClient<T>
-where
-    T: TpuInfoWithSendStatic + Clone,
-{
+impl NotifyKeyUpdate for TpuClientNextClient {
     fn update_key(&self, identity: &Keypair) -> Result<(), Box<dyn std::error::Error>> {
         self.runtime_handle.block_on(self.do_update_key(identity))
     }
 }
 
-impl<T> TransactionClient for TpuClientNextClient<T>
-where
-    T: TpuInfoWithSendStatic + Clone,
-{
+impl TransactionClient for TpuClientNextClient {
     fn send_transactions_in_batch(
         &self,
         wire_transactions: Vec<Vec<u8>>,
@@ -415,8 +395,11 @@ where
         };
         match self.runtime_handle.block_on(handle) {
             Ok(result) => match result {
-                Ok(stats) => {
-                    debug!("tpu-client-next statistics over all the connections: {stats:?}");
+                Ok(scheduler) => {
+                    debug!(
+                        "tpu-client-next statistics over all the connections: {:?}",
+                        scheduler.stats
+                    );
                 }
                 Err(error) => error!("tpu-client-next exits with error {error}."),
             },

send-transaction-service/src/test_utils.rs

@@ -18,6 +18,7 @@ solana-clock = { workspace = true }
 solana-connection-cache = { workspace = true }
 solana-keypair = { workspace = true }
 solana-measure = { workspace = true }
+solana-metrics ={ workspace = true, optional = true }
 solana-quic-definitions = { workspace = true }
 solana-rpc-client = { workspace = true }
 solana-streamer = { workspace = true }
@@ -39,3 +40,6 @@ solana-streamer = { workspace = true, features = ["dev-context-only-utils"] }
 
 [package.metadata.docs.rs]
 targets = ["x86_64-unknown-linux-gnu"]
+
+[features]
+metrics =["dep:solana-metrics"]

send-transaction-service/src/transaction_client.rs

@@ -2,7 +2,7 @@
 //! to the upcoming leaders.
 
 use {
-    super::{leader_updater::LeaderUpdater, SendTransactionStatsPerAddr},
+    super::leader_updater::LeaderUpdater,
     crate::{
         connection_worker::ConnectionWorker,
         quic_networking::{
@@ -21,14 +21,19 @@ use {
     tokio::sync::mpsc,
     tokio_util::sync::CancellationToken,
 };
+pub type TransactionReceiver = mpsc::Receiver<TransactionBatch>;
 
 /// The [`ConnectionWorkersScheduler`] sends transactions from the provided
 /// receiver channel to upcoming leaders. It obtains information about future
 /// leaders from the implementation of the [`LeaderUpdater`] trait.
 ///
 /// Internally, it enables the management and coordination of multiple network
 /// connections, schedules and oversees connection workers.
-pub struct ConnectionWorkersScheduler;
+pub struct ConnectionWorkersScheduler {
+    leader_updater: Box<dyn LeaderUpdater>,
+    transaction_receiver: TransactionReceiver,
+    pub stats: Arc<SendTransactionStats>,
+}
 
 /// Errors that arise from running [`ConnectionWorkersSchedulerError`].
 #[derive(Debug, Error, PartialEq)]
@@ -137,12 +142,19 @@ pub trait WorkersBroadcaster {
     ) -> Result<(), ConnectionWorkersSchedulerError>;
 }
 
-pub type TransactionStatsAndReceiver = (
-    SendTransactionStatsPerAddr,
-    mpsc::Receiver<TransactionBatch>,
-);
-
 impl ConnectionWorkersScheduler {
+    pub fn new(
+        leader_updater: Box<dyn LeaderUpdater>,
+        transaction_receiver: mpsc::Receiver<TransactionBatch>,
+    ) -> Self {
+        let stats = Arc::new(SendTransactionStats::default());
+        Self {
+            leader_updater,
+            transaction_receiver,
+            stats,
+        }
+    }
+
     /// Starts the scheduler, which manages the distribution of transactions to
     /// the network's upcoming leaders.
     ///
@@ -154,18 +166,12 @@ impl ConnectionWorkersScheduler {
     /// will be dropped. The same for transactions that failed to be delivered
     /// over the network.
     pub async fn run(
+        self,
         config: ConnectionWorkersSchedulerConfig,
-        leader_updater: Box<dyn LeaderUpdater>,
-        transaction_receiver: mpsc::Receiver<TransactionBatch>,
         cancel: CancellationToken,
-    ) -> Result<TransactionStatsAndReceiver, ConnectionWorkersSchedulerError> {
-        Self::run_with_broadcaster::<NonblockingBroadcaster>(
-            config,
-            leader_updater,
-            transaction_receiver,
-            cancel,
-        )
-        .await
+    ) -> Result<Self, ConnectionWorkersSchedulerError> {
+        self.run_with_broadcaster::<NonblockingBroadcaster>(config, cancel)
+            .await
     }
 
     /// Starts the scheduler, which manages the distribution of transactions to
@@ -182,6 +188,7 @@ impl ConnectionWorkersScheduler {
     /// Importantly, if some transactions were not delivered due to network
     /// problems, they will not be retried when the problem is resolved.
     pub async fn run_with_broadcaster<Broadcaster: WorkersBroadcaster>(
+        mut self,
         ConnectionWorkersSchedulerConfig {
             bind,
             stake_identity,
@@ -191,20 +198,17 @@ impl ConnectionWorkersScheduler {
             max_reconnect_attempts,
             leaders_fanout,
         }: ConnectionWorkersSchedulerConfig,
-        mut leader_updater: Box<dyn LeaderUpdater>,
-        mut transaction_receiver: mpsc::Receiver<TransactionBatch>,
         cancel: CancellationToken,
-    ) -> Result<TransactionStatsAndReceiver, ConnectionWorkersSchedulerError> {
+    ) -> Result<Self, ConnectionWorkersSchedulerError> {
         let endpoint = Self::setup_endpoint(bind, stake_identity)?;
         debug!("Client endpoint bind address: {:?}", endpoint.local_addr());
         let mut workers = WorkersCache::new(num_connections, cancel.clone());
-        let mut send_stats_per_addr = SendTransactionStatsPerAddr::new();
 
         let mut last_error = None;
 
         loop {
             let transaction_batch: TransactionBatch = tokio::select! {
-                recv_res = transaction_receiver.recv() => match recv_res {
+                recv_res = self.transaction_receiver.recv() => match recv_res {
                     Some(txs) => txs,
                     None => {
                         debug!("End of `transaction_receiver`: shutting down.");
@@ -217,21 +221,20 @@ impl ConnectionWorkersScheduler {
                 }
             };
 
-            let connect_leaders = leader_updater.next_leaders(leaders_fanout.connect);
+            let connect_leaders = self.leader_updater.next_leaders(leaders_fanout.connect);
             let send_leaders = extract_send_leaders(&connect_leaders, leaders_fanout.send);
 
             // add future leaders to the cache to hide the latency of opening
             // the connection.
             for peer in connect_leaders {
                 if !workers.contains(&peer) {
-                    let stats = send_stats_per_addr.entry(peer.ip()).or_default();
                     let worker = Self::spawn_worker(
                         &endpoint,
                         &peer,
                         worker_channel_size,
                         skip_check_transaction_age,
                         max_reconnect_attempts,
-                        stats.clone(),
+                        self.stats.clone(),
                     );
                     maybe_shutdown_worker(workers.push(peer, worker));
                 }
@@ -248,11 +251,11 @@ impl ConnectionWorkersScheduler {
         workers.shutdown().await;
 
         endpoint.close(0u32.into(), b"Closing connection");
-        leader_updater.stop().await;
+        self.leader_updater.stop().await;
         if let Some(error) = last_error {
             return Err(error);
         }
-        Ok((send_stats_per_addr, transaction_receiver))
+        Ok(self)
     }
 
     /// Sets up the QUIC endpoint for the scheduler to handle connections.