perf: no more channels for UDP send/recv

Cargo.toml

@@ -7,6 +7,7 @@ members = [
   "iroh-net",
   "iroh-sync",
   "iroh-test",
+  "iroh-net/bench"
 ]
 resolver = "2"
 

iroh-net/Cargo.toml

@@ -35,6 +35,7 @@ libc = "0.2.139"
 num_enum = "0.6.1"
 once_cell = "1.17.0"
 os_info = "3.6.0"
+parking_lot = "0.12.1"
 postcard = { version = "1", default-features = false, features = ["alloc", "use-std", "experimental-derive"] }
 rand = "0.8"
 rand_core = "0.6.4"

iroh-net/bench/Cargo.toml

@@ -0,0 +1,19 @@
+[package]
+name = "bench"
+version = "0.1.0"
+edition = "2021"
+license = "MIT OR Apache-2.0"
+publish = false
+
+[dependencies]
+anyhow = "1.0.22"
+bytes = "1"
+hdrhistogram = { version = "7.2", default-features = false }
+iroh-net = { path = ".." }
+quinn = "0.10"
+rcgen = "0.11.1"
+rustls = { version = "0.21.0", default-features = false, features = ["quic"] }
+clap = { version = "4", features = ["derive"] }
+tokio = { version = "1.0.1", features = ["rt", "sync"] }
+tracing = "0.1.10"
+tracing-subscriber = { version = "0.3.0", default-features = false, features = ["env-filter", "fmt", "ansi", "time", "local-time"] }

iroh-net/bench/src/bin/bulk.rs

@@ -0,0 +1,211 @@
+use std::{
+    sync::{Arc, Mutex},
+    time::Instant,
+};
+
+use anyhow::{Context, Result};
+use clap::Parser;
+use iroh_net::{MagicEndpoint, PeerAddr};
+use tokio::sync::Semaphore;
+use tracing::{info, trace};
+
+use bench::{
+    configure_tracing_subscriber, connect_client, drain_stream, rt, send_data_on_stream,
+    server_endpoint,
+    stats::{Stats, TransferResult},
+    Opt,
+};
+
+fn main() {
+    let opt = Opt::parse();
+    configure_tracing_subscriber();
+
+    let server_span = tracing::error_span!("server");
+    let runtime = rt();
+    let (server_addr, endpoint) = {
+        let _guard = server_span.enter();
+        server_endpoint(&runtime, &opt)
+    };
+
+    let server_thread = std::thread::spawn(move || {
+        let _guard = server_span.entered();
+        if let Err(e) = runtime.block_on(server(endpoint, opt)) {
+            eprintln!("server failed: {e:#}");
+        }
+    });
+
+    let mut handles = Vec::new();
+    for id in 0..opt.clients {
+        let server_addr = server_addr.clone();
+        handles.push(std::thread::spawn(move || {
+            let _guard = tracing::error_span!("client", id).entered();
+            let runtime = rt();
+            match runtime.block_on(client(server_addr, opt)) {
+                Ok(stats) => Ok(stats),
+                Err(e) => {
+                    eprintln!("client failed: {e:#}");
+                    Err(e)
+                }
+            }
+        }));
+    }
+
+    for (id, handle) in handles.into_iter().enumerate() {
+        // We print all stats at the end of the test sequentially to avoid
+        // them being garbled due to being printed concurrently
+        if let Ok(stats) = handle.join().expect("client thread") {
+            stats.print(id);
+        }
+    }
+
+    server_thread.join().expect("server thread");
+}
+
+async fn server(endpoint: MagicEndpoint, opt: Opt) -> Result<()> {
+    let mut server_tasks = Vec::new();
+
+    // Handle only the expected amount of clients
+    for _ in 0..opt.clients {
+        let handshake = endpoint.accept().await.unwrap();
+        let connection = handshake.await.context("handshake failed")?;
+
+        server_tasks.push(tokio::spawn(async move {
+            loop {
+                let (mut send_stream, mut recv_stream) = match connection.accept_bi().await {
+                    Err(quinn::ConnectionError::ApplicationClosed(_)) => break,
+                    Err(e) => {
+                        eprintln!("accepting stream failed: {e:?}");
+                        break;
+                    }
+                    Ok(stream) => stream,
+                };
+                trace!("stream established");
+
+                tokio::spawn(async move {
+                    drain_stream(&mut recv_stream, opt.read_unordered).await?;
+                    send_data_on_stream(&mut send_stream, opt.download_size).await?;
+                    Ok::<_, anyhow::Error>(())
+                });
+            }
+
+            if opt.stats {
+                println!("\nServer connection stats:\n{:#?}", connection.stats());
+            }
+        }));
+    }
+
+    // Await all the tasks. We have to do this to prevent the runtime getting dropped
+    // and all server tasks to be cancelled
+    for handle in server_tasks {
+        if let Err(e) = handle.await {
+            eprintln!("Server task error: {e:?}");
+        };
+    }
+
+    Ok(())
+}
+
+async fn client(server_addr: PeerAddr, opt: Opt) -> Result<ClientStats> {
+    let (endpoint, connection) = connect_client(server_addr, opt).await?;
+
+    let start = Instant::now();
+
+    let connection = Arc::new(connection);
+
+    let mut stats = ClientStats::default();
+    let mut first_error = None;
+
+    let sem = Arc::new(Semaphore::new(opt.max_streams));
+    let results = Arc::new(Mutex::new(Vec::new()));
+    for _ in 0..opt.streams {
+        let permit = sem.clone().acquire_owned().await.unwrap();
+        let results = results.clone();
+        let connection = connection.clone();
+        tokio::spawn(async move {
+            let result =
+                handle_client_stream(connection, opt.upload_size, opt.read_unordered).await;
+            info!("stream finished: {:?}", result);
+            results.lock().unwrap().push(result);
+            drop(permit);
+        });
+    }
+
+    // Wait for remaining streams to finish
+    let _ = sem.acquire_many(opt.max_streams as u32).await.unwrap();
+
+    for result in results.lock().unwrap().drain(..) {
+        match result {
+            Ok((upload_result, download_result)) => {
+                stats.upload_stats.stream_finished(upload_result);
+                stats.download_stats.stream_finished(download_result);
+            }
+            Err(e) => {
+                if first_error.is_none() {
+                    first_error = Some(e);
+                }
+            }
+        }
+    }
+
+    stats.upload_stats.total_duration = start.elapsed();
+    stats.download_stats.total_duration = start.elapsed();
+
+    // Explicit close of the connection, since handles can still be around due
+    // to `Arc`ing them
+    connection.close(0u32.into(), b"Benchmark done");
+
+    endpoint.close(0u32.into(), b"").await?;
+
+    if opt.stats {
+        println!("\nClient connection stats:\n{:#?}", connection.stats());
+    }
+
+    match first_error {
+        None => Ok(stats),
+        Some(e) => Err(e),
+    }
+}
+
+async fn handle_client_stream(
+    connection: Arc<quinn::Connection>,
+    upload_size: u64,
+    read_unordered: bool,
+) -> Result<(TransferResult, TransferResult)> {
+    let start = Instant::now();
+
+    let (mut send_stream, mut recv_stream) = connection
+        .open_bi()
+        .await
+        .context("failed to open stream")?;
+
+    send_data_on_stream(&mut send_stream, upload_size).await?;
+
+    let upload_result = TransferResult::new(start.elapsed(), upload_size);
+
+    let start = Instant::now();
+    let size = drain_stream(&mut recv_stream, read_unordered).await?;
+    let download_result = TransferResult::new(start.elapsed(), size as u64);
+
+    Ok((upload_result, download_result))
+}
+
+#[derive(Default)]
+struct ClientStats {
+    upload_stats: Stats,
+    download_stats: Stats,
+}
+
+impl ClientStats {
+    pub fn print(&self, client_id: usize) {
+        println!();
+        println!("Client {client_id} stats:");
+
+        if self.upload_stats.total_size != 0 {
+            self.upload_stats.print("upload");
+        }
+
+        if self.download_stats.total_size != 0 {
+            self.download_stats.print("download");
+        }
+    }
+}