Repository: commonwarexyz/monorepo
Pull request #866: [runtime] Separate networking into its own module
Author: login: danlaine
Created: 2025-04-30T16:58:00Z | Merged: 2025-05-02T22:04:51Z

This PR:

* Adds a new `mod network` inside `commonware_runtime`
* Adds 2 `commonware_runtime::Network` implementations
* One is in `metered.rs`. It wraps another `Network` to track metrics.
* The other is `tokio.rs` which uses tokio for network IO.
* Adds a `network` field to `Context` which is a `MeteredNetwork<TokioNetwork>` and has the same functionality as before


Diff:
diff --git a/runtime/src/lib.rs b/runtime/src/lib.rs
index c547cf3598..0101ca7ac4 100644
--- a/runtime/src/lib.rs
+++ b/runtime/src/lib.rs
@@ -34,6 +34,7 @@ cfg_if::cfg_if! {
         pub mod benchmarks;
     }
 }
+mod network;
 mod storage;
 pub mod telemetry;
 mod utils;
diff --git a/runtime/src/network/metered.rs b/runtime/src/network/metered.rs
new file mode 100644
index 0000000000..42a8fdad88
--- /dev/null
+++ b/runtime/src/network/metered.rs
@@ -0,0 +1,158 @@
+use crate::{SinkOf, StreamOf};
+use prometheus_client::{metrics::counter::Counter, registry::Registry};
+use std::{net::SocketAddr, sync::Arc};
+
+#[derive(Debug)]
+/// Tracks network metrics.
+struct Metrics {
+    /// Number of connections created by dialing us.
+    inbound_connections: Counter,
+    /// Number of connections created by dialing others.
+    outbound_connections: Counter,
+    /// Bandwidth used by receiving data from others.
+    inbound_bandwidth: Counter,
+    /// Bandwidth used by sending data to others.
+    outbound_bandwidth: Counter,
+}
+
+impl Metrics {
+    fn new(registry: &mut Registry) -> Self {
+        let metrics = Self {
+            inbound_connections: Counter::default(),
+            outbound_connections: Counter::default(),
+            inbound_bandwidth: Counter::default(),
+            outbound_bandwidth: Counter::default(),
+        };
+        registry.register(
+            "inbound_connections",
+            "Number of connections created by dialing us",
+            metrics.inbound_connections.clone(),
+        );
+        registry.register(
+            "outbound_connections",
+            "Number of connections created by dialing others",
+            metrics.outbound_connections.clone(),
+        );
+        registry.register(
+            "inbound_bandwidth",
+            "Bandwidth used by receiving data from others",
+            metrics.inbound_bandwidth.clone(),
+        );
+        registry.register(
+            "outbound_bandwidth",
+            "Bandwidth used by sending data to others",
+            metrics.outbound_bandwidth.clone(),
+        );
+        metrics
+    }
+}
+
+/// Sends using the `inner` sink and tracks metrics for it.
+pub struct Sink<S: crate::Sink> {
+    inner: S,
+    metrics: Arc<Metrics>,
+}
+
+impl<S: crate::Sink> crate::Sink for Sink<S> {
+    async fn send(&mut self, data: &[u8]) -> Result<(), crate::Error> {
+        self.inner.send(data).await?;
+        self.metrics.outbound_bandwidth.inc_by(data.len() as u64);
+        Ok(())
+    }
+}
+
+/// Receives from the `inner` stream and tracks metrics for it.
+pub struct Stream<S: crate::Stream> {
+    inner: S,
+    metrics: Arc<Metrics>,
+}
+
+impl<S: crate::Stream> crate::Stream for Stream<S> {
+    async fn recv(&mut self, buf: &mut [u8]) -> Result<(), crate::Error> {
+        self.inner.recv(buf).await?;
+        self.metrics.inbound_bandwidth.inc_by(buf.len() as u64);
+        Ok(())
+    }
+}
+
+/// Listens for incoming connections using the `inner` listener
+/// and tracks metrics for it.
+pub struct Listener<L: crate::Listener> {
+    inner: L,
+    metrics: Arc<Metrics>,
+}
+
+impl<L: crate::Listener> crate::Listener for Listener<L> {
+    type Sink = Sink<L::Sink>;
+    type Stream = Stream<L::Stream>;
+
+    async fn accept(&mut self) -> Result<(SocketAddr, Self::Sink, Self::Stream), crate::Error> {
+        let (addr, sink, stream) = self.inner.accept().await?;
+        self.metrics.inbound_connections.inc();
+        Ok((
+            addr,
+            Sink {
+                inner: sink,
+                metrics: self.metrics.clone(),
+            },
+            Stream {
+                inner: stream,
+                metrics: self.metrics.clone(),
+            },
+        ))
+    }
+}
+
+/// A metered network implementation which wraps another
+/// [crate::Network] and tracks metrics for it.
+#[derive(Debug, Clone)]
+pub struct Network<N: crate::Network> {
+    inner: N,
+    /// Metrics for the network.
+    /// Note these are not tracked on a per-connection basis.
+    /// That would be nice but it would be very expensive
+    /// and potentially an OOM vector.
+    metrics: Arc<Metrics>,
+}
+
+impl<N: crate::Network> Network<N> {
+    /// Wraps `inner` to make it metered.
+    /// The `registry` is used to register the metrics.
+    pub fn new(inner: N, registry: &mut Registry) -> Self {
+        let metrics = Metrics::new(registry);
+        Self {
+            inner,
+            metrics: Arc::new(metrics),
+        }
+    }
+}
+
+impl<N: crate::Network> crate::Network for Network<N> {
+    type Listener = Listener<N::Listener>;
+
+    async fn bind(&self, socket: SocketAddr) -> Result<Self::Listener, crate::Error> {
+        let inner = self.inner.bind(socket).await?;
+        Ok(Listener {
+            inner,
+            metrics: self.metrics.clone(),
+        })
+    }
+
+    async fn dial(
+        &self,
+        socket: SocketAddr,
+    ) -> Result<(SinkOf<Self>, StreamOf<Self>), crate::Error> {
+        let (sink, stream) = self.inner.dial(socket).await?;
+        self.metrics.outbound_connections.inc();
+        Ok((
+            Sink {
+                inner: sink,
+                metrics: self.metrics.clone(),
+            },
+            Stream {
+                inner: stream,
+                metrics: self.metrics.clone(),
+            },
+        ))
+    }
+}
diff --git a/runtime/src/network/mod.rs b/runtime/src/network/mod.rs
new file mode 100644
index 0000000000..cf65fb3698
--- /dev/null
+++ b/runtime/src/network/mod.rs
@@ -0,0 +1,3 @@
+pub(crate) mod metered;
+#[cfg(not(target_arch = "wasm32"))]
+pub(crate) mod tokio;
diff --git a/runtime/src/network/tokio.rs b/runtime/src/network/tokio.rs
new file mode 100644
index 0000000000..f80c60c85b
--- /dev/null
+++ b/runtime/src/network/tokio.rs
@@ -0,0 +1,206 @@
+use crate::Error;
+use std::{net::SocketAddr, time::Duration};
+use tokio::{
+    io::{AsyncReadExt as _, AsyncWriteExt as _},
+    net::{
+        tcp::{OwnedReadHalf, OwnedWriteHalf},
+        TcpListener, TcpStream,
+    },
+    time::timeout,
+};
+use tracing::warn;
+
+/// Implementation of [crate::Sink] for the [tokio] runtime.
+pub struct Sink {
+    write_timeout: Duration,
+    sink: OwnedWriteHalf,
+}
+
+impl crate::Sink for Sink {
+    async fn send(&mut self, msg: &[u8]) -> Result<(), Error> {
+        // Time out if we take too long to write
+        timeout(self.write_timeout, self.sink.write_all(msg))
+            .await
+            .map_err(|_| Error::Timeout)?
+            .map_err(|_| Error::SendFailed)?;
+        Ok(())
+    }
+}
+
+/// Implementation of [crate::Stream] for the [tokio] runtime.
+pub struct Stream {
+    read_timeout: Duration,
+    stream: OwnedReadHalf,
+}
+
+impl crate::Stream for Stream {
+    async fn recv(&mut self, buf: &mut [u8]) -> Result<(), Error> {
+        // Time out if we take too long to read
+        timeout(self.read_timeout, self.stream.read_exact(buf))
+            .await
+            .map_err(|_| Error::Timeout)?
+            .map_err(|_| Error::RecvFailed)?;
+        Ok(())
+    }
+}
+
+/// Implementation of [crate::Listener] using the [tokio] runtime.
+pub struct Listener {
+    cfg: Config,
+    listener: TcpListener,
+}
+
+impl crate::Listener for Listener {
+    type Sink = Sink;
+    type Stream = Stream;
+
+    async fn accept(&mut self) -> Result<(SocketAddr, Self::Sink, Self::Stream), Error> {
+        // Accept a new TCP stream
+        let (stream, addr) = self.listener.accept().await.map_err(|_| Error::Closed)?;
+
+        // Set TCP_NODELAY if configured
+        if let Some(tcp_nodelay) = self.cfg.tcp_nodelay {
+            if let Err(err) = stream.set_nodelay(tcp_nodelay) {
+                warn!(?err, "failed to set TCP_NODELAY");
+            }
+        }
+
+        // Return the sink and stream
+        let (stream, sink) = stream.into_split();
+        Ok((
+            addr,
+            Sink {
+                write_timeout: self.cfg.write_timeout,
+                sink,
+            },
+            Stream {
+                read_timeout: self.cfg.read_timeout,
+                stream,
+            },
+        ))
+    }
+}
+
+#[derive(Clone, Debug)]
+/// Configuration for the tokio [Network] implementation of the [crate::Network] trait.
+pub(crate) struct Config {
+    /// Whether or not to disable Nagle's algorithm.
+    ///
+    /// The algorithm combines a series of small network packets into a single packet
+    /// before sending to reduce overhead of sending multiple small packets which might not
+    /// be efficient on slow, congested networks. However, to do so the algorithm introduces
+    /// a slight delay as it waits to accumulate more data. Latency-sensitive networks should
+    /// consider disabling it to send the packets as soon as possible to reduce latency.
+    ///
+    /// Note: Make sure that your compile target has and allows this configuration otherwise
+    /// panics or unexpected behaviours are possible.
+    tcp_nodelay: Option<bool>,
+    /// Read timeout for connections, after which the connection will be closed
+    read_timeout: Duration,
+    /// Write timeout for connections, after which the connection will be closed
+    write_timeout: Duration,
+}
+
+impl Config {
+    // Setters
+    /// See [Config]
+    pub fn with_tcp_nodelay(mut self, tcp_nodelay: Option<bool>) -> Self {
+        self.tcp_nodelay = tcp_nodelay;
+        self
+    }
+    /// See [Config]
+    pub fn with_read_timeout(mut self, read_timeout: Duration) -> Self {
+        self.read_timeout = read_timeout;
+        self
+    }
+    /// See [Config]
+    pub fn with_write_timeout(mut self, write_timeout: Duration) -> Self {
+        self.write_timeout = write_timeout;
+        self
+    }
+
+    // Getters
+    /// See [Config]
+    pub fn tcp_nodelay(&self) -> Option<bool> {
+        self.tcp_nodelay
+    }
+    /// See [Config]
+    pub fn read_timeout(&self) -> Duration {
+        self.read_timeout
+    }
+    /// See [Config]
+    pub fn write_timeout(&self) -> Duration {
+        self.write_timeout
+    }
+}
+
+impl Default for Config {
+    fn default() -> Self {
+        Self {
+            tcp_nodelay: None,
+            read_timeout: Duration::from_secs(60),
+            write_timeout: Duration::from_secs(30),
+        }
+    }
+}
+
+#[derive(Clone, Debug)]
+pub(crate) struct Network {
+    cfg: Config,
+}
+
+impl From<Config> for Network {
+    fn from(cfg: Config) -> Self {
+        Self { cfg }
+    }
+}
+
+impl Default for Network {
+    fn default() -> Self {
+        Self::from(Config::default())
+    }
+}
+
+impl crate::Network for Network {
+    type Listener = Listener;
+
+    async fn bind(&self, socket: SocketAddr) -> Result<Self::Listener, crate::Error> {
+        TcpListener::bind(socket)
+            .await
+            .map_err(|_| Error::BindFailed)
+            .map(|listener| Listener {
+                cfg: self.cfg.clone(),
+                listener,
+            })
+    }
+
+    async fn dial(
+        &self,
+        socket: SocketAddr,
+    ) -> Result<(crate::SinkOf<Self>, crate::StreamOf<Self>), crate::Error> {
+        // Create a new TCP stream
+        let stream = TcpStream::connect(socket)
+            .await
+            .map_err(|_| Error::ConnectionFailed)?;
+
+        // Set TCP_NODELAY if configured
+        if let Some(tcp_nodelay) = self.cfg.tcp_nodelay {
+            if let Err(err) = stream.set_nodelay(tcp_nodelay) {
+                warn!(?err, "failed to set TCP_NODELAY");
+            }
+        }
+
+        // Return the sink and stream
+        let (stream, sink) = stream.into_split();
+        Ok((
+            Sink {
+                write_timeout: self.cfg.write_timeout,
+                sink,
+            },
+            Stream {
+                read_timeout: self.cfg.read_timeout,
+                stream,
+            },
+        ))
+    }
+}
diff --git a/runtime/src/tokio/runtime.rs b/runtime/src/tokio/runtime.rs
index 9a74659b64..179a74e9e4 100644
--- a/runtime/src/tokio/runtime.rs
+++ b/runtime/src/tokio/runtime.rs
@@ -1,3 +1,7 @@
+use crate::network::metered::{Listener as MeteredListener, Network as MeteredNetwork};
+use crate::network::tokio::{
+    Config as TokioNetworkConfig, Listener as TokioListener, Network as TokioNetwork,
+};
 use crate::storage::metered::Storage;
 use crate::storage::tokio::{Config as TokioStorageConfig, Storage as TokioStorage};
 use crate::{utils::Signaler, Clock, Error, Handle, Signal, METRICS_PREFIX};
@@ -12,19 +16,12 @@ use rand::{rngs::OsRng, CryptoRng, RngCore};
 use std::{
     env,
     future::Future,
-    io,
     net::SocketAddr,
     path::PathBuf,
     sync::{Arc, Mutex},
     time::{Duration, SystemTime},
 };
-use tokio::{
-    io::{AsyncReadExt, AsyncWriteExt},
-    net::{tcp::OwnedReadHalf, tcp::OwnedWriteHalf, TcpListener, TcpStream},
-    runtime::{Builder, Runtime},
-    time::timeout,
-};
-use tracing::warn;
+use tokio::runtime::{Builder, Runtime};
 
 #[derive(Clone, Debug, Hash, PartialEq, Eq, EncodeLabelSet)]
 struct Work {
@@ -37,13 +34,6 @@ struct Metrics {
     tasks_running: Family<Work, Gauge>,
     blocking_tasks_spawned: Family<Work, Counter>,
     blocking_tasks_running: Family<Work, Gauge>,
-
-    // As nice as it would be to track each of these by socket address,
-    // it quickly becomes an OOM attack vector.
-    inbound_connections: Counter,
-    outbound_connections: Counter,
-    inbound_bandwidth: Counter,
-    outbound_bandwidth: Counter,
 }
 
 impl Metrics {
@@ -53,10 +43,6 @@ impl Metrics {
             tasks_running: Family::default(),
             blocking_tasks_spawned: Family::default(),
             blocking_tasks_running: Family::default(),
-            inbound_connections: Counter::default(),
-            outbound_connections: Counter::default(),
-            inbound_bandwidth: Counter::default(),
-            outbound_bandwidth: Counter::default(),
         };
         registry.register(
             "tasks_spawned",
@@ -78,26 +64,6 @@ impl Metrics {
             "Number of blocking tasks currently running",
             metrics.blocking_tasks_running.clone(),
         );
-        registry.register(
-            "inbound_connections",
-            "Number of connections created by dialing us",
-            metrics.inbound_connections.clone(),
-        );
-        registry.register(
-            "outbound_connections",
-            "Number of connections created by dialing others",
-            metrics.outbound_connections.clone(),
-        );
-        registry.register(
-            "inbound_bandwidth",
-            "Bandwidth used by receiving data from others",
-            metrics.inbound_bandwidth.clone(),
-        );
-        registry.register(
-            "outbound_bandwidth",
-            "Bandwidth used by sending data to others",
-            metrics.outbound_bandwidth.clone(),
-        );
         metrics
     }
 }
@@ -124,24 +90,6 @@ pub struct Config {
     /// Whether or not to catch panics.
     catch_panics: bool,
 
-    /// Duration after which to close the connection if no message is read.
-    read_timeout: Duration,
-
-    /// Duration after which to close the connection if a message cannot be written.
-    write_timeout: Duration,
-
-    /// Whether or not to disable Nagle's algorithm.
-    ///
-    /// The algorithm combines a series of small network packets into a single packet
-    /// before sending to reduce overhead of sending multiple small packets which might not
-    /// be efficient on slow, congested networks. However, to do so the algorithm introduces
-    /// a slight delay as it waits to accumulate more data. Latency-sensitive networks should
-    /// consider disabling it to send the packets as soon as possible to reduce latency.
-    ///
-    /// Note: Make sure that your compile target has and allows this configuration otherwise
-    /// panics or unexpected behaviours are possible.
-    tcp_nodelay: Option<bool>,
-
     /// Base directory for all storage operations.
     storage_directory: PathBuf,
 
@@ -149,6 +97,8 @@ pub struct Config {
     ///
     /// Tokio sets the default value to 2MB.
     maximum_buffer_size: usize,
+
+    network_cfg: TokioNetworkConfig,
 }
 
 impl Config {
@@ -160,11 +110,9 @@ impl Config {
             worker_threads: 2,
             max_blocking_threads: 512,
             catch_panics: true,
-            read_timeout: Duration::from_secs(60),
-            write_timeout: Duration::from_secs(30),
-            tcp_nodelay: None,
             storage_directory,
             maximum_buffer_size: 2 * 1024 * 1024, // 2 MB
+            network_cfg: TokioNetworkConfig::default(),
         }
     }
 
@@ -186,17 +134,17 @@ impl Config {
     }
     /// See [Config]
     pub fn with_read_timeout(mut self, d: Duration) -> Self {
-        self.read_timeout = d;
+        self.network_cfg = self.network_cfg.with_read_timeout(d);
         self
     }
     /// See [Config]
     pub fn with_write_timeout(mut self, d: Duration) -> Self {
-        self.write_timeout = d;
+        self.network_cfg = self.network_cfg.with_write_timeout(d);
         self
     }
     /// See [Config]
     pub fn with_tcp_nodelay(mut self, n: Option<bool>) -> Self {
-        self.tcp_nodelay = n;
+        self.network_cfg = self.network_cfg.with_tcp_nodelay(n);
         self
     }
     /// See [Config]
@@ -225,15 +173,15 @@ impl Config {
     }
     /// See [Config]
     pub fn read_timeout(&self) -> Duration {
-        self.read_timeout
+        self.network_cfg.read_timeout()
     }
     /// See [Config]
     pub fn write_timeout(&self) -> Duration {
-        self.write_timeout
+        self.network_cfg.write_timeout()
     }
     /// See [Config]
     pub fn tcp_nodelay(&self) -> Option<bool> {
-        self.tcp_nodelay
+        self.network_cfg.tcp_nodelay()
     }
     /// See [Config]
     pub fn storage_directory(&self) -> &PathBuf {
@@ -309,6 +257,9 @@ impl crate::Runner for Runner {
             runtime_registry,
         );
 
+        let network = TokioNetwork::from(self.cfg.network_cfg.clone());
+        let network = MeteredNetwork::new(network, runtime_registry);
+
         let executor = Arc::new(Executor {
             cfg: self.cfg,
             registry: Mutex::new(registry),
@@ -323,6 +274,7 @@ impl crate::Runner for Runner {
             label: String::new(),
             spawned: false,
             executor: executor.clone(),
+            network,
         };
 
         executor.runtime.block_on(f(context))
@@ -337,6 +289,7 @@ pub struct Context {
     spawned: bool,
     executor: Arc<Executor>,
     storage: Storage<TokioStorage>,
+    network: MeteredNetwork<TokioNetwork>,
 }
 
 impl Clone for Context {
@@ -346,6 +299,7 @@ impl Clone for Context {
             spawned: false,
             executor: self.executor.clone(),
             storage: self.storage.clone(),
+            network: self.network.clone(),
         }
     }
 }
@@ -483,6 +437,7 @@ impl crate::Metrics for Context {
             spawned: false,
             executor: self.executor.clone(),
             storage: self.storage.clone(),
+            network: self.network.clone(),
         }
     }
 
@@ -545,145 +500,14 @@ impl GClock for Context {
 impl ReasonablyRealtime for Context {}
 
 impl crate::Network for Context {
-    type Listener = Listener;
+    type Listener = MeteredListener<TokioListener>;
 
-    async fn bind(&self, socket: SocketAddr) -> Result<Listener, Error> {
-        TcpListener::bind(socket)
-            .await
-            .map_err(|_| Error::BindFailed)
-            .map(|listener| Listener {
-                context: self.clone(),
-                listener,
-            })
+    async fn bind(&self, socket: SocketAddr) -> Result<Self::Listener, Error> {
+        self.network.bind(socket).await
     }
 
     async fn dial(&self, socket: SocketAddr) -> Result<(SinkOf<Self>, StreamOf<Self>), Error> {
-        // Create a new TCP stream
-        let stream = TcpStream::connect(socket)
-            .await
-            .map_err(|_| Error::ConnectionFailed)?;
-        self.executor.metrics.outbound_connections.inc();
-
-        // Set TCP_NODELAY if configured
-        if let Some(tcp_nodelay) = self.executor.cfg.tcp_nodelay {
-            if let Err(err) = stream.set_nodelay(tcp_nodelay) {
-                warn!(?err, "failed to set TCP_NODELAY");
-            }
-        }
-
-        // Return the sink and stream
-        let context = self.clone();
-        let (stream, sink) = stream.into_split();
-        Ok((
-            Sink {
-                context: context.clone(),
-                sink,
-            },
-            Stream { context, stream },
-        ))
-    }
-}
-
-/// Implementation of [crate::Listener] for the `tokio` runtime.
-pub struct Listener {
-    context: Context,
-    listener: TcpListener,
-}
-
-impl crate::Listener for Listener {
-    type Sink = Sink;
-    type Stream = Stream;
-
-    async fn accept(&mut self) -> Result<(SocketAddr, Self::Sink, Self::Stream), Error> {
-        // Accept a new TCP stream
-        let (stream, addr) = self.listener.accept().await.map_err(|_| Error::Closed)?;
-        self.context.executor.metrics.inbound_connections.inc();
-
-        // Set TCP_NODELAY if configured
-        if let Some(tcp_nodelay) = self.context.executor.cfg.tcp_nodelay {
-            if let Err(err) = stream.set_nodelay(tcp_nodelay) {
-                warn!(?err, "failed to set TCP_NODELAY");
-            }
-        }
-
-        // Return the sink and stream
-        let context = self.context.clone();
-        let (stream, sink) = stream.into_split();
-        Ok((
-            addr,
-            Sink {
-                context: context.clone(),
-                sink,
-            },
-            Stream { context, stream },
-        ))
-    }
-}
-
-impl axum::serve::Listener for Listener {
-    type Io = TcpStream;
-    type Addr = SocketAddr;
-
-    async fn accept(&mut self) -> (Self::Io, Self::Addr) {
-        let (stream, addr) = self.listener.accept().await.unwrap();
-        (stream, addr)
-    }
-
-    fn local_addr(&self) -> io::Result<Self::Addr> {
-        self.listener.local_addr()
-    }
-}
-
-/// Implementation of [crate::Sink] for the `tokio` runtime.
-pub struct Sink {
-    context: Context,
-    sink: OwnedWriteHalf,
-}
-
-impl crate::Sink for Sink {
-    async fn send(&mut self, msg: &[u8]) -> Result<(), Error> {
-        let len = msg.len();
-        timeout(
-            self.context.executor.cfg.write_timeout,
-            self.sink.write_all(msg),
-        )
-        .await
-        .map_err(|_| Error::Timeout)?
-        .map_err(|_| Error::SendFailed)?;
-        self.context
-            .executor
-            .metrics
-            .outbound_bandwidth
-            .inc_by(len as u64);
-        Ok(())
-    }
-}
-
-/// Implementation of [crate::Stream] for the `tokio` runtime.
-pub struct Stream {
-    context: Context,
-    stream: OwnedReadHalf,
-}
-
-impl crate::Stream for Stream {
-    async fn recv(&mut self, buf: &mut [u8]) -> Result<(), Error> {
-        // Wait for the stream to be readable
-        timeout(
-            self.context.executor.cfg.read_timeout,
-            self.stream.read_exact(buf),
-        )
-        .await
-        .map_err(|_| Error::Timeout)?
-        .map_err(|_| Error::RecvFailed)?;
-
-        // Record metrics
-        self.context
-            .executor
-            .metrics
-            .inbound_bandwidth
-            .inc_by(buf.len() as u64);
-
-        Ok(())
+        self.network.dial(socket).await
     }
 }
 


Reviews:
user: login: BrendanChou
state: DISMISSED
author_association: COLLABORATOR
submitted_at: 2025-04-30T19:22:50Z
user: login: patrick-ogrady
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-02T21:28:37Z
user: login: patrick-ogrady
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-02T22:04:45Z

Review comments:
diff_hunk: @@ -0,0 +1,218 @@
+use crate::Error;
+use std::{net::SocketAddr, time::Duration};
+use tokio::{
+    io::{AsyncReadExt as _, AsyncWriteExt as _},
+    net::{
+        tcp::{OwnedReadHalf, OwnedWriteHalf},
+        TcpListener, TcpStream,
+    },
+    time::timeout,
+};
+use tracing::warn;
+
+/// Implementation of [crate::Sink] for the [tokio] runtime.
+pub struct Sink {
+    write_timeout: Duration,
+    sink: OwnedWriteHalf,
+}
+
+impl crate::Sink for Sink {
+    async fn send(&mut self, msg: &[u8]) -> Result<(), Error> {
+        timeout(self.write_timeout, self.sink.write_all(msg))
+            .await
+            .map_err(|_| Error::Timeout)?
+            .map_err(|_| Error::SendFailed)?;
+        Ok(())
+    }
+}
+
+/// Implementation of [crate::Stream] for the [tokio] runtime.
+pub struct Stream {
+    read_timeout: Duration,
+    stream: OwnedReadHalf,
+}
+
+impl crate::Stream for Stream {
+    async fn recv(&mut self, buf: &mut [u8]) -> Result<(), Error> {
+        // Wait for the stream to be readable
path: runtime/src/network/tokio.rs
user: login: patrick-ogrady
body: nit: there is a comment here but not in `send`.
created_at: 2025-05-02T21:25:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,218 @@
+use crate::Error;
+use std::{net::SocketAddr, time::Duration};
+use tokio::{
+    io::{AsyncReadExt as _, AsyncWriteExt as _},
+    net::{
+        tcp::{OwnedReadHalf, OwnedWriteHalf},
+        TcpListener, TcpStream,
+    },
+    time::timeout,
+};
+use tracing::warn;
+
+/// Implementation of [crate::Sink] for the [tokio] runtime.
+pub struct Sink {
+    write_timeout: Duration,
+    sink: OwnedWriteHalf,
+}
+
+impl crate::Sink for Sink {
+    async fn send(&mut self, msg: &[u8]) -> Result<(), Error> {
+        timeout(self.write_timeout, self.sink.write_all(msg))
+            .await
+            .map_err(|_| Error::Timeout)?
+            .map_err(|_| Error::SendFailed)?;
+        Ok(())
+    }
+}
+
+/// Implementation of [crate::Stream] for the [tokio] runtime.
+pub struct Stream {
+    read_timeout: Duration,
+    stream: OwnedReadHalf,
+}
+
+impl crate::Stream for Stream {
+    async fn recv(&mut self, buf: &mut [u8]) -> Result<(), Error> {
+        // Wait for the stream to be readable
+        timeout(self.read_timeout, self.stream.read_exact(buf))
+            .await
+            .map_err(|_| Error::Timeout)?
+            .map_err(|_| Error::RecvFailed)?;
+        Ok(())
+    }
+}
+
+/// Implementation of [crate::Listener] using the [tokio] runtime.
+pub struct Listener {
+    cfg: Config,
+    listener: TcpListener,
+}
+
+impl crate::Listener for Listener {
+    type Sink = Sink;
+    type Stream = Stream;
+
+    async fn accept(&mut self) -> Result<(SocketAddr, Self::Sink, Self::Stream), Error> {
+        // Accept a new TCP stream
+        let (stream, addr) = self.listener.accept().await.map_err(|_| Error::Closed)?;
+
+        // Set TCP_NODELAY if configured
+        if let Some(tcp_nodelay) = self.cfg.tcp_nodelay {
+            if let Err(err) = stream.set_nodelay(tcp_nodelay) {
+                warn!(?err, "failed to set TCP_NODELAY");
+            }
+        }
+
+        // Return the sink and stream
+        let (stream, sink) = stream.into_split();
+        Ok((
+            addr,
+            Sink {
+                write_timeout: self.cfg.write_timeout,
+                sink,
+            },
+            Stream {
+                read_timeout: self.cfg.read_timeout,
+                stream,
+            },
+        ))
+    }
+}
+
+impl axum::serve::Listener for Listener {
path: runtime/src/network/tokio.rs
user: login: patrick-ogrady
body: You can remove this. I used to use it for metrics but concluded that "web" traffic should not use this runtime (as it may desire very different connections semantics than the read/write here).
created_at: 2025-05-02T21:26:34Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,218 @@
+use crate::Error;
+use std::{net::SocketAddr, time::Duration};
+use tokio::{
+    io::{AsyncReadExt as _, AsyncWriteExt as _},
+    net::{
+        tcp::{OwnedReadHalf, OwnedWriteHalf},
+        TcpListener, TcpStream,
+    },
+    time::timeout,
+};
+use tracing::warn;
+
+/// Implementation of [crate::Sink] for the [tokio] runtime.
+pub struct Sink {
+    write_timeout: Duration,
+    sink: OwnedWriteHalf,
+}
+
+impl crate::Sink for Sink {
+    async fn send(&mut self, msg: &[u8]) -> Result<(), Error> {
+        timeout(self.write_timeout, self.sink.write_all(msg))
+            .await
+            .map_err(|_| Error::Timeout)?
+            .map_err(|_| Error::SendFailed)?;
+        Ok(())
+    }
+}
+
+/// Implementation of [crate::Stream] for the [tokio] runtime.
+pub struct Stream {
+    read_timeout: Duration,
+    stream: OwnedReadHalf,
+}
+
+impl crate::Stream for Stream {
+    async fn recv(&mut self, buf: &mut [u8]) -> Result<(), Error> {
+        // Wait for the stream to be readable
+        timeout(self.read_timeout, self.stream.read_exact(buf))
+            .await
+            .map_err(|_| Error::Timeout)?
+            .map_err(|_| Error::RecvFailed)?;
+        Ok(())
+    }
+}
+
+/// Implementation of [crate::Listener] using the [tokio] runtime.
+pub struct Listener {
+    cfg: Config,
+    listener: TcpListener,
+}
+
+impl crate::Listener for Listener {
+    type Sink = Sink;
+    type Stream = Stream;
+
+    async fn accept(&mut self) -> Result<(SocketAddr, Self::Sink, Self::Stream), Error> {
+        // Accept a new TCP stream
+        let (stream, addr) = self.listener.accept().await.map_err(|_| Error::Closed)?;
+
+        // Set TCP_NODELAY if configured
+        if let Some(tcp_nodelay) = self.cfg.tcp_nodelay {
+            if let Err(err) = stream.set_nodelay(tcp_nodelay) {
+                warn!(?err, "failed to set TCP_NODELAY");
+            }
+        }
+
+        // Return the sink and stream
+        let (stream, sink) = stream.into_split();
+        Ok((
+            addr,
+            Sink {
+                write_timeout: self.cfg.write_timeout,
+                sink,
+            },
+            Stream {
+                read_timeout: self.cfg.read_timeout,
+                stream,
+            },
+        ))
+    }
+}
+
+impl axum::serve::Listener for Listener {
+    type Io = TcpStream;
+    type Addr = SocketAddr;
+
+    async fn accept(&mut self) -> (Self::Io, Self::Addr) {
+        let (stream, addr) = self.listener.accept().await.unwrap();
+        (stream, addr)
+    }
+
+    fn local_addr(&self) -> std::io::Result<Self::Addr> {
+        self.listener.local_addr()
+    }
+}
+
+#[derive(Clone, Debug)]
+pub(crate) struct Config {
path: runtime/src/network/tokio.rs
user: login: patrick-ogrady
body: nit: add a `Config` comment
created_at: 2025-05-02T21:26:51Z
author_association: CONTRIBUTOR

PR comments:
user: login: codecov[bot]
created_at: 2025-05-02T22:05:00Z
body: ## [Codecov](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/866?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz) Report
Attention: Patch coverage is `82.35294%` with `36 lines` in your changes missing coverage. Please review.
> Project coverage is 89.45%. Comparing base [(`4c429d7`)](https://app.codecov.io/gh/commonwarexyz/monorepo/commit/4c429d76160464ecac3ba19a2d52ca1d4f11e084?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz) to head [(`a28a268`)](https://app.codecov.io/gh/commonwarexyz/monorepo/commit/a28a268b00a2739152e8b16db24d8df9723ec900?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz).
> Report is 3 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/866?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz) | Patch % | Lines |
|---|---|---|
| [runtime/src/network/tokio.rs](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/866?src=pr&el=tree&filepath=runtime%2Fsrc%2Fnetwork%2Ftokio.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz#diff-cnVudGltZS9zcmMvbmV0d29yay90b2tpby5ycw==) | 70.87% | [30 Missing :warning: ](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz) |
| [runtime/src/tokio/runtime.rs](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/866?src=pr&el=tree&filepath=runtime%2Fsrc%2Ftokio%2Fruntime.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz#diff-cnVudGltZS9zcmMvdG9raW8vcnVudGltZS5ycw==) | 62.50% | [6 Missing :warning: ](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/866?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz) |

```diff
@@            Coverage Diff             @@
##             main     #866      +/-   ##
==========================================
- Coverage   89.48%   89.45%   -0.03%     
==========================================
  Files         167      169       +2     
  Lines       42488    42574      +86     
==========================================
+ Hits        38021    38086      +65     
- Misses       4467     4488      +21     
```

| [Files with missing lines](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/866?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz) | Coverage Δ | |
|---|---|---|
| [runtime/src/lib.rs](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/866?src=pr&el=tree&filepath=runtime%2Fsrc%2Flib.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz#diff-cnVudGltZS9zcmMvbGliLnJz) | `96.78% <ø> (-0.13%)` | :arrow_down: |
| [runtime/src/network/metered.rs](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/866?src=pr&el=tree&filepath=runtime%2Fsrc%2Fnetwork%2Fmetered.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz#diff-cnVudGltZS9zcmMvbmV0d29yay9tZXRlcmVkLnJz) | `100.00% <100.00%> (ø)` | |
| [runtime/src/tokio/runtime.rs](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/866?src=pr&el=tree&filepath=runtime%2Fsrc%2Ftokio%2Fruntime.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz#diff-cnVudGltZS9zcmMvdG9raW8vcnVudGltZS5ycw==) | `81.94% <62.50%> (-1.21%)` | :arrow_down: |
| [runtime/src/network/tokio.rs](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/866?src=pr&el=tree&filepath=runtime%2Fsrc%2Fnetwork%2Ftokio.rs&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz#diff-cnVudGltZS9zcmMvbmV0d29yay90b2tpby5ycw==) | `70.87% <70.87%> (ø)` | |

... and [1 file with indirect coverage changes](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/866/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz)

------

[Continue to review full report in Codecov by Sentry](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/866?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/866?dropdown=coverage&src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz). Last update [4c429d7...a28a268](https://app.codecov.io/gh/commonwarexyz/monorepo/pull/866?dropdown=coverage&src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=commonwarexyz).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE