added disconnect to client pool

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<RwLock<Vec<Arc<MixnetClient>>>>, // 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::<Vec<_>>()
                         .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<MixnetClient> {
         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(),
         }
     }
 }