Max/client pool (#5188)

* tcp conn tracker * make default decay const * first pass connpool * err handling conpool start * added notes for next features * first version working * first pass spin out client_pool * cancel token * logging change * bump default decay time * bugfix: make sure to apply gateway score filtering when choosing initial node * add duplicate packets received to troubleshooting * client_pool.rs mod * client pool example * clippy * client pool example done * added disconnect to client pool * update mod file * add cancel token disconnect fn * comments * comments * add clone * added disconnect thread * update example files tcpproxy * client pool docs * remove comments for future ffi push + lower default pool size from 4 to 2 * comment on ffi * update command help * clone impl * remove clone * fix clippy * fix clippy again * fix test * tweaked text grammar * updated comment in example * future is now * cherry * cherry * fix borked rebase * fix fmt * wasm fix --------- Co-authored-by: Jędrzej Stuczyński <jedrzej.stuczynski@gmail.com>
2025-01-14 16:11:47 +00:00
parent 0a47d5dcf8
commit e454d71b78
35 changed files with 924 additions and 231 deletions
@@ -18,6 +18,8 @@ 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)]
 struct ExampleMessage {
@@ -26,6 +28,8 @@ struct ExampleMessage {
    tcp_conn: i8,
 }

+// This example just starts off a bunch of Tcp connections on a loop to a remote endpoint: in this case the TcpListener behind the NymProxyServer instance on the echo server found in `nym/tools/echo-server/`. It pipes a few messages to it, logs the replies, and keeps track of the number of replies received per connection.
+//
 // To run:
 // - run the echo server with `cargo run`
 // - run this example with `cargo run --example tcp_proxy_multistream -- <ECHO_SERVER_NYM_ADDRESS> <ENV_FILE_PATH> <CLIENT_PORT>` e.g.
@@ -40,8 +44,13 @@ async fn main() -> anyhow::Result<()> {
    // Nym client logging is very informative but quite verbose.
    // The Message Decay related logging gives you an ideas of the internals of the proxy message ordering: you need to switch
    // to DEBUG to see the contents of the msg buffer, sphinx packet chunking, etc.
-    tracing_subscriber::fmt()
-        .with_max_level(tracing::Level::INFO)
+    tracing_subscriber::registry()
+        .with(fmt::layer())
+        .with(
+            EnvFilter::new("info")
+                .add_directive("nym_sdk::client_pool=info".parse().unwrap())
+                .add_directive("nym_sdk::tcp_proxy_client=debug".parse().unwrap()),
+        )
        .init();

    let env_path = env::args().nth(2).expect("Env file not specified");
@@ -49,23 +58,42 @@ async fn main() -> anyhow::Result<()> {

    let listen_port = env::args().nth(3).expect("Port not specified");

-    // Within the TcpProxyClient, individual client shutdown is triggered by the timeout.
+    // Within the TcpProxyClient, individual client shutdown is triggered by the timeout. The final argument is how many clients to keep in reserve in the client pool when running the TcpProxy.
    let proxy_client =
-        tcp_proxy::NymProxyClient::new(server, "127.0.0.1", &listen_port, 45, Some(env)).await?;
+        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..4 {
+    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;
-        println!("Starting TCP connection {}", conn_id);
        let local_tcp_addr = format!("127.0.0.1:{}", listen_port.clone());
        tokio::spawn(async move {
            // Now the client and server proxies are running we can create and pipe traffic to/from
@@ -81,7 +109,10 @@ 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..4 {
+                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;
@@ -96,12 +127,7 @@ async fn main() -> anyhow::Result<()> {
                        .write_all(&serialised)
                        .await
                        .expect("couldn't write to stream");
-                    println!(
-                        ">> client sent {}: {} bytes on conn {}",
-                        &i,
-                        msg.message_bytes.len(),
-                        &conn_id
-                    );
+                    println!(">> client sent msg {} on conn {}", &i, &conn_id);
                }
                Ok::<(), anyhow::Error>(())
            });
@@ -113,17 +139,8 @@ async fn main() -> anyhow::Result<()> {
                while let Some(Ok(bytes)) = framed_read.next().await {
                    match bincode::deserialize::<ExampleMessage>(&bytes) {
                        Ok(msg) => {
-                            println!(
-                                "<< client received {}: {} bytes on conn {}",
-                                msg.message_id,
-                                msg.message_bytes.len(),
-                                msg.tcp_conn
-                            );
                            reply_counter += 1;
-                            println!(
-                                "tcp connection {} replies received {}/4",
-                                msg.tcp_conn, reply_counter
-                            );
+                            println!("<< conn {} received {}/8", msg.tcp_conn, reply_counter);
                        }
                        Err(e) => {
                            println!("<< client received something that wasn't an example message of {} bytes. error: {}", bytes.len(), e);
@@ -138,15 +155,12 @@ 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.
-    signal::ctrl_c().await?;
-    println!("CTRL+C received, shutting down");
    Ok(())
 }

 // emulate a series of small messages followed by a closing larger one
 fn gen_bytes_fixed(i: usize) -> Vec<u8> {
-    let amounts = [10, 15, 50, 1000];
+    let amounts = [10, 15, 50, 1000, 10, 15, 500, 2000];
    let len = amounts[i];
    let mut rng = rand::thread_rng();
    (0..len).map(|_| rng.gen::<u8>()).collect()
@@ -21,6 +21,8 @@ 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)]
 struct ExampleMessage {
@@ -28,7 +30,14 @@ struct ExampleMessage {
    message_bytes: Vec<u8>,
 }

-
+// This is a basic example which opens a single TCP connection and writes a bunch of messages between a client and an echo
+// server, so only uses a single session under the hood and doesn't really show off the message ordering capabilities; this is mainly
+// just a quick introductory illustration on how:
+// - the mixnet does message ordering
+// - the NymProxyClient and NymProxyServer can be hooked into and used to communicate between two otherwise pretty vanilla TcpStreams
+//
+// For a more irl example checkout tcp_proxy_multistream.rs
+//
 // Run this with:
 // `cargo run --example tcp_proxy_single_connection <SERVER_LISTEN_PORT> <ENV_FILE_PATH> <CLIENT_LISTEN_PATH>` e.g.
 // `cargo run --example tcp_proxy_single_connection 8081 ../../../envs/canary.env 8080 `
@@ -39,8 +48,9 @@ async fn main() -> anyhow::Result<()> {

    // Comment this out to just see println! statements from this example, as Nym client logging is very informative but quite verbose.
    // The Message Decay related logging gives you an ideas of the internals of the proxy message ordering. To see the contents of the msg buffer, sphinx packet chunking, etc change the tracing::Level to DEBUG.
-    tracing_subscriber::fmt()
-        .with_max_level(tracing::Level::INFO)
+    tracing_subscriber::registry()
+        .with(fmt::layer())
+        .with(EnvFilter::new("nym_sdk::tcp_proxy=info"))
        .init();

    let server_port = env::args()
@@ -63,10 +73,14 @@ 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.
+    // The final argument is how many clients to keep in reserve in the client pool when running the TcpProxy.
    let proxy_client =
-        tcp_proxy::NymProxyClient::new(*proxy_nym_addr, "127.0.0.1", &client_port, 60, Some(env))
+        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>(())
@@ -77,10 +91,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();
@@ -118,9 +150,9 @@ async fn main() -> anyhow::Result<()> {
        Ok::<(), anyhow::Error>(())
    });

-    // Just wait for Nym clients to connect, TCP clients to bind, etc.
+    // Just wait for Nym clients to connect, TCP clients to bind, etc. If there isn't a client in the pool (or you started it with 0) already then the TcpProxyClient just spins up an ephemeral client itself.
    println!("waiting for everything to be set up..");
-    tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;
+    tokio::time::sleep(tokio::time::Duration::from_secs(10)).await;
    println!("done. sending bytes");

    // Now the client and server proxies are running we can create and pipe traffic to/from
@@ -140,6 +172,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,
@@ -179,15 +214,12 @@ async fn main() -> anyhow::Result<()> {
        }
    }

-    // Once timeout is passed, you can either wait for graceful shutdown or just hard stop it.
-    signal::ctrl_c().await?;
-    println!(":: CTRL+C received, shutting down + cleanup up proxy server config files");
-    fs::remove_dir_all(conf_path)?;
    Ok(())
 }

 fn gen_bytes_fixed(i: usize) -> Vec<u8> {
-    let amounts = vec![1, 10, 50, 100, 150, 200, 350, 500, 750, 1000];
+    // let amounts = vec![1, 10, 50, 100, 150, 200, 350, 500, 750, 1000];
+    let amounts = [158, 1088, 505, 1001, 150, 200, 3500, 500, 750, 100];
    let len = amounts[i];
    let mut rng = rand::thread_rng();
    (0..len).map(|_| rng.gen::<u8>()).collect()