diff --git a/sdk/rust/nym-sdk/examples/tcp_proxy_multistream.rs b/sdk/rust/nym-sdk/examples/tcp_proxy_multistream.rs index 6d61ea2ca1..4134ab325d 100644 --- a/sdk/rust/nym-sdk/examples/tcp_proxy_multistream.rs +++ b/sdk/rust/nym-sdk/examples/tcp_proxy_multistream.rs @@ -10,6 +10,7 @@ use tokio::net::TcpStream; use tokio::signal; use tokio_stream::StreamExt; use tokio_util::codec; +use tokio_util::sync::CancellationToken; use tracing_subscriber::{fmt, prelude::*, EnvFilter}; #[derive(Serialize, Deserialize, Debug)] @@ -53,17 +54,37 @@ async fn main() -> anyhow::Result<()> { let proxy_client = tcp_proxy::NymProxyClient::new(server, "127.0.0.1", &listen_port, 45, Some(env), 2).await?; + // For our disconnect() logic below + let proxy_clone = proxy_client.clone(); + tokio::spawn(async move { proxy_client.run().await?; Ok::<(), anyhow::Error>(()) }); + let example_cancel_token = CancellationToken::new(); + let client_cancel_token = example_cancel_token.clone(); + let watcher_cancel_token = example_cancel_token.clone(); + + // Cancel listener thread + tokio::spawn(async move { + signal::ctrl_c().await?; + println!(":: CTRL_C received, shutting down + cleanup up proxy server config files"); + watcher_cancel_token.cancel(); + proxy_clone.disconnect().await; + Ok::<(), anyhow::Error>(()) + }); + println!("waiting for everything to be set up.."); tokio::time::sleep(tokio::time::Duration::from_secs(5)).await; println!("done. sending bytes"); // In the info traces you will see the different session IDs being set up, one for each TcpStream. for i in 0..8 { + let client_cancel_inner_token = client_cancel_token.clone(); + if client_cancel_token.is_cancelled() { + break; + } let conn_id = i; let local_tcp_addr = format!("127.0.0.1:{}", listen_port.clone()); tokio::spawn(async move { @@ -81,6 +102,9 @@ async fn main() -> anyhow::Result<()> { // Lets just send a bunch of messages to the server with variable delays between them, with a message and tcp connection ids to keep track of ordering on the server side (for illustrative purposes **only**; keeping track of anonymous replies is handled by the proxy under the hood with Single Use Reply Blocks (SURBs); for this illustration we want some kind of app-level message id, but irl most of the time you'll probably be parsing on e.g. the incoming response type instead) tokio::spawn(async move { for i in 0..8 { + if client_cancel_inner_token.is_cancelled() { + break; + } let mut rng = SmallRng::from_entropy(); let delay: f64 = rng.gen_range(2.5..5.0); tokio::time::sleep(tokio::time::Duration::from_secs_f64(delay)).await; @@ -123,10 +147,6 @@ async fn main() -> anyhow::Result<()> { tokio::time::sleep(tokio::time::Duration::from_secs_f64(delay)).await; } - // Once timeout is passed, you can either wait for graceful shutdown or just hard stop it. - // TODO CHANGE make this a task listening for ctrl_c, add cancel token to loops + call client.disconnect() on ctrl_c - signal::ctrl_c().await?; - println!("CTRL_C received, shutting down"); Ok(()) } 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 c8388da8bf..cc7c934f68 100644 --- a/sdk/rust/nym-sdk/examples/tcp_proxy_single_connection.rs +++ b/sdk/rust/nym-sdk/examples/tcp_proxy_single_connection.rs @@ -10,6 +10,7 @@ use tokio::net::{TcpListener, TcpStream}; use tokio::signal; use tokio_stream::StreamExt; use tokio_util::codec; +use tokio_util::sync::CancellationToken; use tracing_subscriber::{fmt, prelude::*, EnvFilter}; #[derive(Serialize, Deserialize, Debug)] @@ -66,6 +67,9 @@ async fn main() -> anyhow::Result<()> { tcp_proxy::NymProxyClient::new(*proxy_nym_addr, "127.0.0.1", &client_port, 5, Some(env), 1) .await?; + // For our disconnect() logic below + let proxy_clone = proxy_client.clone(); + tokio::spawn(async move { proxy_server.run_with_shutdown().await?; Ok::<(), anyhow::Error>(()) @@ -76,10 +80,28 @@ async fn main() -> anyhow::Result<()> { Ok::<(), anyhow::Error>(()) }); + let example_cancel_token = CancellationToken::new(); + let server_cancel_token = example_cancel_token.clone(); + let client_cancel_token = example_cancel_token.clone(); + let watcher_cancel_token = example_cancel_token.clone(); + + // Cancel listener thread + tokio::spawn(async move { + signal::ctrl_c().await?; + println!(":: CTRL_C received, shutting down + cleanup up proxy server config files"); + fs::remove_dir_all(conf_path)?; + watcher_cancel_token.cancel(); + proxy_clone.disconnect().await; + Ok::<(), anyhow::Error>(()) + }); + // 'Server side' thread: echo back incoming as response to the messages sent in the 'client side' thread below tokio::spawn(async move { let listener = TcpListener::bind(upstream_tcp_addr).await?; loop { + if server_cancel_token.is_cancelled() { + break; + } let (socket, _) = listener.accept().await.unwrap(); let (read, mut write) = socket.into_split(); let codec = codec::BytesCodec::new(); @@ -139,6 +161,9 @@ async fn main() -> anyhow::Result<()> { tokio::spawn(async move { let mut rng = SmallRng::from_entropy(); for i in 0..10 { + if client_cancel_token.is_cancelled() { + break; + } let random_bytes = gen_bytes_fixed(i as usize); let msg = ExampleMessage { message_id: i, @@ -178,11 +203,6 @@ async fn main() -> anyhow::Result<()> { } } - // Once timeout is passed, you can either wait for graceful shutdown or just hard stop it. - // TODO CHANGE make this a task listening for ctrl_c, add cancel token to loops + call client.disconnect() on ctrl_c - signal::ctrl_c().await?; - println!(":: CTRL_C received, shutting down + cleanup up proxy server config files"); - fs::remove_dir_all(conf_path)?; Ok(()) }