feat: watch connections

Author: Frando

Cargo.lock

@@ -41,12 +41,11 @@ unexpected_cfgs = { level = "warn", check-cfg = ["cfg(iroh_docsrs)", "cfg(iroh_l
 [workspace.lints.clippy]
 unused-async = "warn"
 
-
 [patch.crates-io]
-netwatch = { git = "https://github.com/n0-computer/net-tools", branch = "feat-multipath" }
-portmapper = { git = "https://github.com/n0-computer/net-tools", branch = "feat-multipath" }
+netwatch = { git = "https://github.com/n0-computer/net-tools", branch = "Frando/on_closed" }
+portmapper = { git = "https://github.com/n0-computer/net-tools", branch = "Frando/on_closed" }
 
-[patch."https://github.com/n0-computer/quinn"]
+# [patch."https://github.com/n0-computer/quinn"]
 # iroh-quinn = { path = "../iroh-quinn/quinn" }
 # iroh-quinn-proto = { path = "../iroh-quinn/quinn-proto" }
 # iroh-quinn-udp = { path = "../iroh-quinn/quinn-udp" }

Cargo.toml

@@ -42,8 +42,8 @@ postcard = { version = "1", default-features = false, features = [
     "use-std",
     "experimental-derive",
 ] }
-quinn = { package = "iroh-quinn", git = "https://github.com/n0-computer/quinn", branch = "main-iroh", default-features = false, features = ["rustls-ring"] }
-quinn-proto = { package = "iroh-quinn-proto", git = "https://github.com/n0-computer/quinn", branch = "main-iroh" }
+quinn = { package = "iroh-quinn", git = "https://github.com/n0-computer/quinn", branch = "Frando/on_closed", default-features = false, features = ["rustls-ring"] }
+quinn-proto = { package = "iroh-quinn-proto", git = "https://github.com/n0-computer/quinn", branch = "Frando/on_closed" }
 rand = "0.9.2"
 reqwest = { version = "0.12", default-features = false, features = [
     "rustls-tls",

iroh-relay/Cargo.toml

@@ -44,9 +44,9 @@ n0-watcher = "0.5"
 netwatch = { version = "0.12" }
 pin-project = "1"
 pkarr = { version = "5", default-features = false, features = ["relays"] }
-quinn = { package = "iroh-quinn", git = "https://github.com/n0-computer/quinn", branch = "main-iroh", default-features = false, features = ["rustls-ring"] }
-quinn-proto = { package = "iroh-quinn-proto", git = "https://github.com/n0-computer/quinn", branch = "main-iroh" }
-quinn-udp = { package = "iroh-quinn-udp", git = "https://github.com/n0-computer/quinn", branch = "main-iroh" }
+quinn = { package = "iroh-quinn", git = "https://github.com/n0-computer/quinn", branch = "Frando/on_closed", default-features = false, features = ["rustls-ring"] }
+quinn-proto = { package = "iroh-quinn-proto", git = "https://github.com/n0-computer/quinn", branch = "Frando/on_closed" }
+quinn-udp = { package = "iroh-quinn-udp", git = "https://github.com/n0-computer/quinn", branch = "Frando/on_closed" }
 rand = "0.9.2"
 reqwest = { version = "0.12", default-features = false, features = [
     "rustls-tls",
@@ -90,7 +90,7 @@ hickory-resolver = "0.25.1"
 igd-next = { version = "0.16", features = ["aio_tokio"] }
 netdev = { version = "0.38.1" }
 portmapper = { version = "0.12", default-features = false }
-quinn = { package = "iroh-quinn", git = "https://github.com/n0-computer/quinn", branch = "main-iroh", default-features = false, features = ["runtime-tokio", "rustls-ring"] }
+quinn = { package = "iroh-quinn", git = "https://github.com/n0-computer/quinn", branch = "Frando/on_closed", default-features = false, features = ["runtime-tokio", "rustls-ring"] }
 tokio = { version = "1", features = [
     "io-util",
     "macros",

iroh/Cargo.toml

@@ -12,7 +12,7 @@ iroh = { path = ".." }
 iroh-metrics = "0.37"
 n0-future = "0.3.0"
 n0-error = "0.1.0"
-quinn = { package = "iroh-quinn", git = "https://github.com/n0-computer/quinn", branch = "main-iroh" }
+quinn = { package = "iroh-quinn", git = "https://github.com/n0-computer/quinn", branch = "Frando/on_closed" }
 rand = "0.9.2"
 rcgen = "0.14"
 rustls = { version = "0.23.33", default-features = false, features = ["ring"] }

iroh/bench/Cargo.toml

@@ -0,0 +1,132 @@
+use std::{sync::Arc, time::Duration};
+
+use iroh::{
+    Endpoint, RelayMode,
+    endpoint::{ConnectionInfo, ConnectionMonitor},
+};
+use n0_error::{Result, StackResultExt, StdResultExt};
+use n0_future::task::AbortOnDropHandle;
+use tokio::{
+    sync::mpsc::{UnboundedReceiver, UnboundedSender},
+    task::JoinSet,
+};
+use tracing::{Instrument, info, info_span};
+
+const ALPN: &[u8] = b"iroh/test";
+
+#[tokio::main]
+async fn main() -> Result {
+    tracing_subscriber::fmt()
+        .with_env_filter(
+            tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()),
+        )
+        .init();
+
+    let monitor = Monitor::new();
+    let server = Endpoint::empty_builder(RelayMode::Disabled)
+        .alpns(vec![ALPN.to_vec()])
+        .monitor_connections(monitor.clone())
+        .bind()
+        .instrument(info_span!("server"))
+        .await?;
+    let server_addr = server.addr();
+
+    let client_task = tokio::spawn(
+        async move {
+            let client = Endpoint::empty_builder(RelayMode::Disabled)
+                .bind()
+                .instrument(info_span!("client"))
+                .await?;
+            for _i in 0..2 {
+                let conn = client.connect(server_addr.clone(), ALPN).await?;
+                conn.close(23u32.into(), b"bye");
+            }
+            client.close().await;
+            n0_error::Ok(client)
+        }
+        .instrument(info_span!("client")),
+    );
+
+    let server_task = tokio::spawn(
+        async move {
+            let mut i = 0;
+            while let Some(conn) = server.accept().await {
+                let conn = conn.await?;
+                conn.closed().await;
+                i += 1;
+                if i == 2 {
+                    break;
+                }
+            }
+            info!("done");
+            n0_error::Ok(server)
+        }
+        .instrument(info_span!("server")),
+    );
+    let client = client_task.await.std_context("client")?.context("client")?;
+    let server = server_task.await.std_context("server")?.context("server")?;
+    client.close().await;
+    server.close().await;
+
+    tokio::time::sleep(Duration::from_secs(1)).await;
+    drop(monitor);
+
+    Ok(())
+}
+
+/// Our connection monitor impl.
+///
+/// This here only logs connection open and close events via tracing.
+/// It could also maintain a datastructure of all connections, or send the stats to some metrics service.
+#[derive(Clone)]
+struct Monitor {
+    tx: UnboundedSender<ConnectionInfo>,
+    _task: Arc<AbortOnDropHandle<()>>,
+}
+
+impl ConnectionMonitor for Monitor {
+    fn on_connection(&self, connection: ConnectionInfo) {
+        self.tx.send(connection).ok();
+    }
+}
+
+impl Monitor {
+    fn new() -> Self {
+        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
+        let task = tokio::spawn(Self::run(rx).instrument(info_span!("watcher")));
+        Self {
+            tx,
+            _task: Arc::new(AbortOnDropHandle::new(task)),
+        }
+    }
+
+    async fn run(mut rx: UnboundedReceiver<ConnectionInfo>) {
+        let mut tasks = JoinSet::new();
+        loop {
+            tokio::select! {
+                Some(conn) = rx.recv() => {
+                    let alpn = String::from_utf8_lossy(conn.alpn()).to_string();
+                    let remote = conn.remote_id().fmt_short();
+                    info!(%remote, %alpn, rtt=?conn.latency(), "new connection");
+                    tasks.spawn(async move {
+                        match conn.closed().await {
+                            Some((close_reason, stats)) => {
+                                // We have access to the final stats of the connection!
+                                info!(%remote, %alpn, ?close_reason, udp_rx=stats.udp_rx.bytes, udp_tx=stats.udp_tx.bytes, "connection closed");
+                            }
+                            None => {
+                                // The connection was closed before we could register our stats-on-close listener.
+                                info!(%remote, %alpn, "connection closed before tracking started");
+                            }
+                        }
+                    }.instrument(tracing::Span::current()));
+                }
+                Some(res) = tasks.join_next(), if !tasks.is_empty() => res.expect("conn close task panicked"),
+                else => break,
+            }
+            while let Some(res) = tasks.join_next().await {
+                res.expect("conn close task panicked");
+            }
+        }
+    }
+}

iroh/examples/monitor-connections.rs

@@ -69,8 +69,8 @@ pub use quinn_proto::{
 
 pub use self::connection::{
     Accept, Accepting, AlpnError, AuthenticationError, Connecting, ConnectingError, Connection,
-    Incoming, IncomingZeroRttConnection, OutgoingZeroRttConnection, RemoteEndpointIdError,
-    ZeroRttStatus,
+    ConnectionInfo, Incoming, IncomingZeroRttConnection, OutgoingZeroRttConnection,
+    RemoteEndpointIdError, ZeroRttStatus,
 };
 
 /// The delay to fall back to discovery when direct addresses fail.
@@ -98,7 +98,7 @@ pub enum PathSelection {
 /// new [`EndpointId`].
 ///
 /// To create the [`Endpoint`] call [`Builder::bind`].
-#[derive(Debug)]
+#[derive(derive_more::Debug)]
 pub struct Builder {
     secret_key: Option<SecretKey>,
     relay_mode: RelayMode,
@@ -117,6 +117,8 @@ pub struct Builder {
     #[cfg(any(test, feature = "test-utils"))]
     path_selection: PathSelection,
     max_tls_tickets: usize,
+    #[debug("{}", connection_monitor.as_ref().map(|_| "Some(Box<dyn ConnectionMonitor>)").unwrap_or("None"))]
+    connection_monitor: Option<Box<dyn ConnectionMonitor>>,
 }
 
 impl Builder {
@@ -158,6 +160,7 @@ impl Builder {
             #[cfg(any(test, feature = "test-utils"))]
             path_selection: PathSelection::default(),
             max_tls_tickets: DEFAULT_MAX_TLS_TICKETS,
+            connection_monitor: None,
         }
     }
 
@@ -208,6 +211,7 @@ impl Builder {
             #[cfg(any(test, feature = "test-utils"))]
             path_selection: self.path_selection,
             metrics,
+            connection_monitor: self.connection_monitor,
         };
 
         let msock = magicsock::MagicSock::spawn(msock_opts).await?;
@@ -432,6 +436,44 @@ impl Builder {
         self.max_tls_tickets = n;
         self
     }
+
+    // TODO docs
+    /// Register a handler that is invoked for each connection the endpoint accepts or initiates.
+    ///
+    /// The [`ConnectionMonitor::on_connection`] method is invoked synchronosuly, from within a tokio
+    /// context. So you can spawn tasks if needed.
+    /// Make sure that whatever you do with the connection info here is non-blocking.
+    /// Usually you'd want to send the info over a broadcast or unbounded channel,
+    /// or insert it into some persistent datastructure.
+    ///
+    /// The `ConnectionInfo` internally contains a weak reference to the connection,
+    /// so keeping the struct alive does not keep the connection alive.
+    /// Note however that `ConnectionInfo` does keep an allocation per connection alive
+    /// so to not leak memory you should drop the `ConnectionInfo` eventually
+    ///
+    /// [`ConnectionMonitor`] is implemented for `Fn(ConnectionInfo)`, so you can
+    /// also pass a closure that takes [`ConnectionInfo`] to this function.
+    pub fn monitor_connections(mut self, monitor: impl ConnectionMonitor) -> Self {
+        self.connection_monitor = Some(Box::new(monitor));
+        self
+    }
+}
+
+/// Monitor each connection accepted or initiated by the endpoint.
+pub trait ConnectionMonitor: Send + Sync + 'static {
+    /// Called for each new connection the endpoint accepts or initiates.
+    ///
+    /// This is only called when a connection is fully established.
+    fn on_connection(&self, connection: ConnectionInfo);
+}
+
+impl<T> ConnectionMonitor for T
+where
+    T: Fn(ConnectionInfo) + Send + Sync + 'static,
+{
+    fn on_connection(&self, connection: ConnectionInfo) {
+        (self)(connection)
+    }
 }
 
 /// Configuration for a [`quinn::Endpoint`] that cannot be changed at runtime.