diff --git a/sdk/rust/nym-sdk/examples/tcp_proxy_single_connection.rs b/sdk/rust/nym-sdk/examples/tcp_proxy_single_connection.rs index 5fcea31267..402523c5e2 100644 --- a/sdk/rust/nym-sdk/examples/tcp_proxy_single_connection.rs +++ b/sdk/rust/nym-sdk/examples/tcp_proxy_single_connection.rs @@ -31,7 +31,6 @@ struct ExampleMessage { #[tokio::main] async fn main() -> anyhow::Result<()> { // Keep track of sent/received messages - // let counter = Arc::new(Mutex::new(0)); let counter = AtomicU8::new(0); // Comment this out to just see println! statements from this example, as Nym client logging is very informative but quite verbose. @@ -60,9 +59,15 @@ async fn main() -> anyhow::Result<()> { // We'll run the instance with a long timeout since we're sending everything down the same Tcp connection, so should be using a single session. // Within the TcpProxyClient, individual client shutdown is triggered by the timeout. - let proxy_client = - tcp_proxy::NymProxyClient::new(*proxy_nym_addr, "127.0.0.1", &client_port, 60, Some(env)) - .await?; + let proxy_client = tcp_proxy::NymProxyClient::new( + *proxy_nym_addr, + "127.0.0.1", + &client_port, + 20, + Some(env), + 2, + ) + .await?; tokio::spawn(async move { proxy_server.run_with_shutdown().await?; diff --git a/sdk/rust/nym-sdk/src/tcp_proxy/client_pool.rs b/sdk/rust/nym-sdk/src/tcp_proxy/client_pool.rs index 41412dfc29..d7d8c5f7bc 100644 --- a/sdk/rust/nym-sdk/src/tcp_proxy/client_pool.rs +++ b/sdk/rust/nym-sdk/src/tcp_proxy/client_pool.rs @@ -5,13 +5,14 @@ use std::fmt; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; use tokio::sync::RwLock; +use tokio::sync::Semaphore; use tracing::{debug, info, instrument, warn}; // Make a set # of clients (low default) -// Make sure that # of clients is always 1 above the # of incoming conn requests // Once a client is used, kill the client & remove it from the pool pub struct ClientPool { clients: Arc>>>, + semaphore: Arc, default_pool_size: usize, conn_count: Arc, // the actual # of connections running, denoting an incoming tcp request that is matched with a nym client } @@ -56,20 +57,20 @@ impl ClientPool { pub fn new(default_pool_size: usize) -> Self { ClientPool { clients: Arc::new(RwLock::new(Vec::new())), + semaphore: Arc::new(Semaphore::new(default_pool_size)), default_pool_size, conn_count: Arc::new(AtomicUsize::new(0)), } } - // if clients == default, sleep - // if incoming conns > default - conn_count (aka in use clients) then make more clients pub async fn start(&self) -> Result<()> { loop { + // TODO double check this.. let spawned_clients = self.clients.read().await.len(); - info!("Currently spawned clients: {}", spawned_clients); + debug!("Currently spawned clients: {}", spawned_clients); if spawned_clients >= self.default_pool_size { - debug!("got enough clients already: sleeping"); + debug!("Got enough clients already: sleeping"); } else { info!("Spawning new client"); let client = loop { @@ -93,14 +94,20 @@ impl ClientPool { } } - pub async fn get_mixnet_client(&self) -> Result { - todo!() + pub async fn get_mixnet_client(&self) -> Option { + let _permit = self.semaphore.acquire().await.ok()?; + let mut clients = self.clients.write().await; + clients + .pop() + .and_then(|arc_client| Arc::try_unwrap(arc_client).ok()) } - // disconnect ephemeral - // remove from vec pub async fn disconnect_and_remove_client(&self, client: MixnetClient) -> Result<()> { - todo!() + let mut clients = self.clients.write().await; + clients.retain(|arc_client| arc_client.as_ref().nym_address() != client.nym_address()); + client.disconnect().await; + self.semaphore.add_permits(1); + Ok(()) } pub async fn get_client_count(&self) -> usize { @@ -126,6 +133,7 @@ impl ClientPool { pub fn clone(&self) -> Self { Self { clients: Arc::clone(&self.clients), + semaphore: Arc::clone(&self.semaphore), default_pool_size: *&self.default_pool_size, conn_count: Arc::clone(&self.conn_count), } diff --git a/sdk/rust/nym-sdk/src/tcp_proxy/tcp_proxy_client.rs b/sdk/rust/nym-sdk/src/tcp_proxy/tcp_proxy_client.rs index f75c52cd5e..e893f79615 100644 --- a/sdk/rust/nym-sdk/src/tcp_proxy/tcp_proxy_client.rs +++ b/sdk/rust/nym-sdk/src/tcp_proxy/tcp_proxy_client.rs @@ -7,7 +7,7 @@ mod client_pool; use client_pool::ClientPool; #[path = "utils.rs"] mod utils; -use anyhow::Result; +use anyhow::{bail, Result}; use dashmap::DashSet; use nym_network_defaults::setup_env; use nym_sphinx::addressing::Recipient; @@ -42,7 +42,7 @@ impl NymProxyClient { env: Option, default_client_amount: usize, ) -> Result { - debug!("loading env file: {:?}", env); + debug!("Loading env file: {:?}", env); setup_env(env); // Defaults to mainnet if empty Ok(NymProxyClient { server_address, @@ -78,13 +78,16 @@ impl NymProxyClient { let overall_counter = self.conn_pool.clone(); tokio::spawn(async move { loop { - info!("active connections: {}", overall_counter.get_conn_count()); + info!("Active connections: {}", overall_counter.get_conn_count()); tokio::time::sleep(Duration::from_secs(5)).await; } }); + // if self.conn_pool.get_client_count().await >= DEFAULT_CLIENT_POOL_SIZE / 2 { loop { - if self.conn_pool.get_client_count().await >= DEFAULT_CLIENT_POOL_SIZE / 2 { + if DEFAULT_CLIENT_POOL_SIZE == 1 && self.conn_pool.get_client_count().await == 1 + || self.conn_pool.get_client_count().await >= DEFAULT_CLIENT_POOL_SIZE / 2 + { let (stream, _) = listener.accept().await?; tokio::spawn(NymProxyClient::handle_incoming( stream, @@ -123,14 +126,13 @@ impl NymProxyClient { info!("Starting session: {}", session_id); - let mut client: MixnetClient = match conn_pool.get_conn_count() <= DEFAULT_CLIENT_POOL_SIZE - { - true => { - debug!("grabbing client from pool"); - conn_pool.get_mixnet_client().await? + let mut client = match conn_pool.get_mixnet_client().await { + Some(client) => { + info!("Grabbing client {} from pool", client.nym_address()); + client } - false => { - debug!("not enough clients in pool, creating ephemeral client"); + None => { + info!("Not enough clients in pool, creating ephemeral client"); let net = NymNetworkDetails::new_from_env(); MixnetClientBuilder::new_ephemeral() .network_details(net) @@ -243,7 +245,8 @@ impl NymProxyClient { info!(" Closing write end of session: {}", session_id); info!(" Triggering client shutdown"); // TODO change this to be a fn in the conn_pool, disconnect and remove from vec - client.disconnect().await; + // client.disconnect().await; + conn_pool.disconnect_and_remove_client(client).await?; conn_pool.clone().decrement_conn_count()?; info!( "Dropped connection - current active connections: {}",