Stop client gracefully in validator (solana-labs#6341)

Stop tpu-client-next clients in validator gracefully.

Author: KirillLykov

core/src/forwarding_stage.rs

@@ -58,7 +58,7 @@ mod packet_container;
 /// * [`TpuClientNextClient`]: Relies on the `tpu-client-next` crate.
 pub enum ForwardingClientOption<'a> {
     ConnectionCache(Arc<ConnectionCache>),
-    TpuClientNext((&'a Keypair, UdpSocket, RuntimeHandle)),
+    TpuClientNext((&'a Keypair, UdpSocket, RuntimeHandle, CancellationToken)),
 }
 
 /// Value chosen because it was used historically, at some point
@@ -153,12 +153,14 @@ pub(crate) fn spawn_forwarding_stage(
             stake_identity,
             tpu_client_socket,
             runtime_handle,
+            cancel,
         )) => {
             let non_vote_client = TpuClientNextClient::new(
                 runtime_handle,
                 forward_address_getter,
                 Some(stake_identity),
                 tpu_client_socket,
+                cancel,
             );
             let forwarding_stage = ForwardingStage::new(
                 receiver,
@@ -550,10 +552,10 @@ impl TpuClientNextClient {
         forward_address_getter: ForwardAddressGetter,
         stake_identity: Option<&Keypair>,
         bind_socket: UdpSocket,
+        cancel: CancellationToken,
     ) -> Self {
         // For now use large channel, the more suitable size to be found later.
         let (sender, receiver) = mpsc::channel(128);
-        let cancel = CancellationToken::new();
         let leader_updater = forward_address_getter.clone();
 
         let config = Self::create_config(bind_socket, stake_identity);

core/src/validator.rs

@@ -146,6 +146,7 @@ use {
     strum_macros::{Display, EnumCount, EnumIter, EnumString, EnumVariantNames, IntoStaticStr},
     thiserror::Error,
     tokio::runtime::Runtime as TokioRuntime,
+    tokio_util::sync::CancellationToken,
 };
 
 const MAX_COMPLETED_DATA_SETS_IN_CHANNEL: usize = 100_000;
@@ -731,13 +732,21 @@ impl Validator {
             }
         }
 
+        // token used to cancel tpu-client-next.
+        let cancel_tpu_client_next = CancellationToken::new();
         {
             let exit = exit.clone();
             config
                 .validator_exit
                 .write()
                 .unwrap()
                 .register_exit(Box::new(move || exit.store(true, Ordering::Relaxed)));
+            let cancel_tpu_client_next = cancel_tpu_client_next.clone();
+            config
+                .validator_exit
+                .write()
+                .unwrap()
+                .register_exit(Box::new(move || cancel_tpu_client_next.cancel()));
         }
 
         let accounts_update_notifier = geyser_plugin_service
@@ -1177,6 +1186,7 @@ impl Validator {
                     Arc::as_ref(&identity_keypair),
                     node.sockets.rpc_sts_client,
                     runtime_handle.clone(),
+                    cancel_tpu_client_next.clone(),
                 )
             } else {
                 let Some(connection_cache) = &connection_cache else {
@@ -1595,6 +1605,7 @@ impl Validator {
                     .take()
                     .expect("Socket should exist."),
                 runtime_handle.clone(),
+                cancel_tpu_client_next,
             ))
         };
         let tpu = Tpu::new_with_client(

rpc/src/rpc_service.rs

@@ -61,6 +61,7 @@ use {
     tokio_util::{
         bytes::Bytes,
         codec::{BytesCodec, FramedRead},
+        sync::CancellationToken,
     },
 };
 
@@ -496,7 +497,7 @@ pub struct JsonRpcServiceConfig<'a> {
 ///   requires a reference to a [`Keypair`].
 pub enum ClientOption<'a> {
     ConnectionCache(Arc<ConnectionCache>),
-    TpuClientNext(&'a Keypair, UdpSocket, RuntimeHandle),
+    TpuClientNext(&'a Keypair, UdpSocket, RuntimeHandle, CancellationToken),
 }
 
 impl JsonRpcService {
@@ -553,7 +554,12 @@ impl JsonRpcService {
                 )?;
                 Ok(json_rpc_service)
             }
-            ClientOption::TpuClientNext(identity_keypair, tpu_client_socket, client_runtime) => {
+            ClientOption::TpuClientNext(
+                identity_keypair,
+                tpu_client_socket,
+                client_runtime,
+                cancel,
+            ) => {
                 let my_tpu_address = config
                     .cluster_info
                     .my_contact_info()
@@ -570,6 +576,7 @@ impl JsonRpcService {
                     config.send_transaction_service_config.leader_forward_count,
                     Some(identity_keypair),
                     tpu_client_socket,
+                    cancel,
                 );
 
                 let json_rpc_service = Self::new_with_client(

send-transaction-service/src/test_utils.rs

@@ -11,6 +11,7 @@ use {
     solana_client::connection_cache::ConnectionCache,
     std::{net::SocketAddr, sync::Arc},
     tokio::runtime::Handle,
+    tokio_util::sync::CancellationToken,
 };
 
 // `maybe_runtime` argument is introduced to be able to use runtime from test
@@ -62,6 +63,7 @@ impl CreateClient for TpuClientNextClient {
             leader_forward_count,
             None,
             bind_socket,
+            CancellationToken::new(),
         )
     }
 }

send-transaction-service/src/transaction_client.rs

@@ -246,6 +246,7 @@ impl TpuClientNextClient {
         leader_forward_count: u64,
         identity: Option<&Keypair>,
         bind_socket: UdpSocket,
+        cancel: CancellationToken,
     ) -> Self
     where
         T: TpuInfoWithSendStatic + Clone,
@@ -256,8 +257,6 @@ impl TpuClientNextClient {
 
         let (update_certificate_sender, update_certificate_receiver) = watch::channel(None);
 
-        let cancel = CancellationToken::new();
-
         let leader_info_provider = CurrentLeaderInfo::new(leader_info);
         let leader_updater: SendTransactionServiceLeaderUpdater<T> =
             SendTransactionServiceLeaderUpdater {