|
| 1 | +use std::{sync::Arc, time::Duration}; |
| 2 | + |
| 3 | +use iroh::{ |
| 4 | + Endpoint, RelayMode, |
| 5 | + endpoint::{ConnectionInfo, ConnectionMonitor}, |
| 6 | +}; |
| 7 | +use n0_error::{Result, StackResultExt, StdResultExt}; |
| 8 | +use n0_future::task::AbortOnDropHandle; |
| 9 | +use tokio::{ |
| 10 | + sync::mpsc::{UnboundedReceiver, UnboundedSender}, |
| 11 | + task::JoinSet, |
| 12 | +}; |
| 13 | +use tracing::{Instrument, info, info_span}; |
| 14 | + |
| 15 | +const ALPN: &[u8] = b"iroh/test"; |
| 16 | + |
| 17 | +#[tokio::main] |
| 18 | +async fn main() -> Result { |
| 19 | + tracing_subscriber::fmt() |
| 20 | + .with_env_filter( |
| 21 | + tracing_subscriber::EnvFilter::try_from_default_env().unwrap_or_else(|_| "info".into()), |
| 22 | + ) |
| 23 | + .init(); |
| 24 | + |
| 25 | + let monitor = Monitor::new(); |
| 26 | + let server = Endpoint::empty_builder(RelayMode::Disabled) |
| 27 | + .alpns(vec![ALPN.to_vec()]) |
| 28 | + .monitor_connections(monitor.clone()) |
| 29 | + .bind() |
| 30 | + .instrument(info_span!("server")) |
| 31 | + .await?; |
| 32 | + let server_addr = server.addr(); |
| 33 | + |
| 34 | + let client_task = tokio::spawn( |
| 35 | + async move { |
| 36 | + let client = Endpoint::empty_builder(RelayMode::Disabled) |
| 37 | + .bind() |
| 38 | + .instrument(info_span!("client")) |
| 39 | + .await?; |
| 40 | + let conn1 = client.connect(server_addr.clone(), ALPN).await?; |
| 41 | + info!("conn1"); |
| 42 | + conn1.close(23u32.into(), b"bye"); |
| 43 | + let conn2 = client.connect(server_addr.clone(), ALPN).await?; |
| 44 | + info!("conn2"); |
| 45 | + conn2.close(23u32.into(), b"bye"); |
| 46 | + tokio::time::sleep(Duration::from_secs(1)).await; |
| 47 | + client.close().await; |
| 48 | + n0_error::Ok(client) |
| 49 | + } |
| 50 | + .instrument(info_span!("client")), |
| 51 | + ); |
| 52 | + |
| 53 | + let server_task = tokio::spawn( |
| 54 | + async move { |
| 55 | + let mut i = 0; |
| 56 | + while let Some(conn) = server.accept().await { |
| 57 | + info!("accepted"); |
| 58 | + let conn = conn.await?; |
| 59 | + info!(remote=%conn.remote_id().fmt_short(), "established"); |
| 60 | + conn.closed().await; |
| 61 | + info!(remote=%conn.remote_id().fmt_short(), "closed"); |
| 62 | + i += 1; |
| 63 | + if i == 2 { |
| 64 | + break; |
| 65 | + } |
| 66 | + } |
| 67 | + info!("done"); |
| 68 | + n0_error::Ok(server) |
| 69 | + } |
| 70 | + .instrument(info_span!("server")), |
| 71 | + ); |
| 72 | + let client = client_task.await.std_context("client")?.context("client")?; |
| 73 | + let server = server_task.await.std_context("server")?.context("server")?; |
| 74 | + client.close().await; |
| 75 | + server.close().await; |
| 76 | + |
| 77 | + tokio::time::sleep(Duration::from_secs(1)).await; |
| 78 | + drop(monitor); |
| 79 | + |
| 80 | + Ok(()) |
| 81 | +} |
| 82 | + |
| 83 | +/// Our connection monitor impl. |
| 84 | +/// |
| 85 | +/// This here only logs connection open and close events via tracing. |
| 86 | +/// It could also maintain a datastructure of all connections, or send the stats to some metrics service. |
| 87 | +#[derive(Clone)] |
| 88 | +struct Monitor { |
| 89 | + tx: UnboundedSender<ConnectionInfo>, |
| 90 | + _task: Arc<AbortOnDropHandle<()>>, |
| 91 | +} |
| 92 | + |
| 93 | +impl ConnectionMonitor for Monitor { |
| 94 | + fn on_connection(&self, connection: ConnectionInfo) { |
| 95 | + self.tx.send(connection).ok(); |
| 96 | + } |
| 97 | +} |
| 98 | + |
| 99 | +impl Monitor { |
| 100 | + fn new() -> Self { |
| 101 | + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); |
| 102 | + let task = tokio::spawn(Self::run(rx).instrument(info_span!("watcher"))); |
| 103 | + Self { |
| 104 | + tx, |
| 105 | + _task: Arc::new(AbortOnDropHandle::new(task)), |
| 106 | + } |
| 107 | + } |
| 108 | + |
| 109 | + async fn run(mut rx: UnboundedReceiver<ConnectionInfo>) { |
| 110 | + let mut tasks = JoinSet::new(); |
| 111 | + loop { |
| 112 | + tokio::select! { |
| 113 | + Some(conn) = rx.recv() => { |
| 114 | + let alpn = String::from_utf8_lossy(conn.alpn()).to_string(); |
| 115 | + let remote = conn.remote_id().fmt_short(); |
| 116 | + info!(%remote, %alpn, rtt=?conn.latency(), "new connection"); |
| 117 | + tasks.spawn(async move { |
| 118 | + match conn.closed().await { |
| 119 | + Some((close_reason, stats)) => { |
| 120 | + // We have access to the final stats of the connection! |
| 121 | + info!(%remote, %alpn, ?close_reason, udp_rx=stats.udp_rx.bytes, udp_tx=stats.udp_tx.bytes, "connection closed"); |
| 122 | + } |
| 123 | + None => { |
| 124 | + // The connection was closed before we could register our stats-on-close listener. |
| 125 | + info!(%remote, %alpn, "connection closed before tracking started"); |
| 126 | + } |
| 127 | + } |
| 128 | + }.instrument(tracing::Span::current())); |
| 129 | + } |
| 130 | + Some(res) = tasks.join_next(), if !tasks.is_empty() => res.expect("conn close task panicked"), |
| 131 | + else => break, |
| 132 | + } |
| 133 | + while let Some(res) = tasks.join_next().await { |
| 134 | + res.expect("conn close task panicked"); |
| 135 | + } |
| 136 | + } |
| 137 | + } |
| 138 | +} |
0 commit comments