cancellable_test/
cancellable_test.rs

1/// Example demonstrating cancellable/interruptible test
2///
3/// This example shows how tests can be cancelled gracefully mid-execution
4/// using the cancellation token.
5///
6/// Run a server: cargo run --example cancellable_test server
7/// Run a client: cargo run --example cancellable_test client
8use rperf3::{Client, Config, ProgressEvent, Server};
9use std::time::Duration;
10
11#[tokio::main]
12async fn main() -> Result<(), Box<dyn std::error::Error>> {
13    env_logger::init();
14
15    let args: Vec<String> = std::env::args().collect();
16    let mode = args.get(1).map(|s| s.as_str()).unwrap_or("client");
17
18    match mode {
19        "server" => run_server().await?,
20        "client" => run_client().await?,
21        _ => {
22            eprintln!("Usage: {} [server|client]", args[0]);
23            std::process::exit(1);
24        }
25    }
26
27    Ok(())
28}
29
30async fn run_server() -> Result<(), Box<dyn std::error::Error>> {
31    println!("Starting TCP server on port 5201...");
32    println!("Press Ctrl+C to stop the server gracefully.\n");
33
34    let config = Config::server(5201);
35    let server = Server::new(config);
36
37    // Clone the cancellation token to handle CTRL+C
38    let cancel_token = server.cancellation_token().clone();
39
40    // Set up CTRL+C handler
41    tokio::spawn(async move {
42        tokio::signal::ctrl_c()
43            .await
44            .expect("Failed to listen for CTRL+C");
45        println!("\nReceived CTRL+C, shutting down server gracefully...");
46        cancel_token.cancel();
47    });
48
49    server.run().await?;
50    println!("Server stopped.");
51    Ok(())
52}
53
54async fn run_client() -> Result<(), Box<dyn std::error::Error>> {
55    println!("Starting TCP client test (30 seconds with 5-second auto-cancel)...");
56    println!("Test will be automatically cancelled after 5 seconds.\n");
57
58    let config =
59        Config::client("127.0.0.1".to_string(), 5201).with_duration(Duration::from_secs(30)); // 30 second test
60
61    let client = Client::new(config)?.with_callback(|event: ProgressEvent| match event {
62        ProgressEvent::TestStarted => {
63            println!("Test started");
64        }
65        ProgressEvent::IntervalUpdate {
66            interval_end,
67            bits_per_second,
68            ..
69        } => {
70            println!(
71                "[{:.1}s] {:.2} Mbps",
72                interval_end.as_secs_f64(),
73                bits_per_second / 1_000_000.0
74            );
75        }
76        ProgressEvent::TestCompleted {
77            total_bytes,
78            bits_per_second,
79            ..
80        } => {
81            println!("\nTest completed:");
82            println!("  Total bytes: {}", total_bytes);
83            println!(
84                "  Average throughput: {:.2} Mbps",
85                bits_per_second / 1_000_000.0
86            );
87        }
88        ProgressEvent::Error(msg) => {
89            eprintln!("Error: {}", msg);
90        }
91    });
92
93    // Clone the cancellation token
94    let cancel_token = client.cancellation_token().clone();
95
96    // Spawn a task to cancel the test after 5 seconds
97    tokio::spawn(async move {
98        tokio::time::sleep(Duration::from_secs(5)).await;
99        println!("\n*** Cancelling test after 5 seconds ***\n");
100        cancel_token.cancel();
101    });
102
103    // Also set up CTRL+C handler for manual cancellation
104    let cancel_token_ctrl_c = client.cancellation_token().clone();
105    tokio::spawn(async move {
106        tokio::signal::ctrl_c()
107            .await
108            .expect("Failed to listen for CTRL+C");
109        println!("\nReceived CTRL+C, cancelling test...");
110        cancel_token_ctrl_c.cancel();
111    });
112
113    client.run().await?;
114    println!("Client test stopped (cancelled after ~5 seconds).");
115    Ok(())
116}