diff --git a/sdk/rust/nym-sdk/src/client_pool/client_pool.rs b/sdk/rust/nym-sdk/src/client_pool/client_pool.rs index 968c8ca47c..37d2085824 100644 --- a/sdk/rust/nym-sdk/src/client_pool/client_pool.rs +++ b/sdk/rust/nym-sdk/src/client_pool/client_pool.rs @@ -3,11 +3,13 @@ use anyhow::Result; use std::fmt; use std::sync::Arc; use tokio::sync::RwLock; +use tokio_util::sync::CancellationToken; use tracing::{debug, info, warn}; pub struct ClientPool { clients: Arc>>>, // Collection of clients waiting to be used which are popped off in get_mixnet_client() client_pool_reserve_number: usize, // Default # of clients to have available in pool in reserve waiting for incoming connections + cancel_token: CancellationToken, } // This is only necessary for when you're wanting to check the addresses of the clients that are currently in the pool. @@ -20,9 +22,7 @@ impl fmt::Debug for ClientPool { clients .iter() .enumerate() - .map(|(i, client)| { - format!("\n {}: {}", i, client.nym_address()) - }) + .map(|(i, client)| format!("\n {}: {}", i, client.nym_address())) .collect::>() .join(",") } else { @@ -53,10 +53,12 @@ impl ClientPool { ClientPool { clients: Arc::new(RwLock::new(Vec::new())), client_pool_reserve_number, + cancel_token: CancellationToken::new(), } } // The loop here is simple: if there aren't enough clients, create more. If you set clients to 0, repeatedly just sleep. + // disconnect_pool() will kill this loop via the cancellation token. pub async fn start(&self) -> Result<()> { loop { let spawned_clients = self.clients.read().await.len(); @@ -65,6 +67,9 @@ impl ClientPool { "Currently spawned clients: {}: {:?}", spawned_clients, addresses ); + if self.cancel_token.is_cancelled() { + break Ok(()); + } if spawned_clients >= self.client_pool_reserve_number { debug!("Got enough clients already: sleeping"); } else { @@ -93,6 +98,22 @@ impl ClientPool { } } + pub async fn disconnect_pool(&self) { + info!("Triggering Client Pool disconnect"); + self.cancel_token.cancel(); + info!( + "Client pool cancellation token cancelled: {}", + self.cancel_token.is_cancelled() + ); + let mut clients = self.clients.write().await; + while let Some(arc_client) = clients.pop() { + if let Ok(client) = Arc::try_unwrap(arc_client) { + info!("Killing reserve client {}", client.nym_address()); + client.disconnect().await; + } + } + } + pub async fn get_mixnet_client(&self) -> Option { debug!("Grabbing client from pool"); let mut clients = self.clients.write().await; @@ -113,6 +134,7 @@ impl ClientPool { Self { clients: Arc::clone(&self.clients), client_pool_reserve_number: self.client_pool_reserve_number, + cancel_token: self.cancel_token.clone(), } } }