feat(transport): start working on transport metering & metrics (#102)

Author: thedevbirb

.gitignore

@@ -4,3 +4,5 @@
 *.svg
 *.log
 .DS_Store
+
+todo.md

Cargo.lock

@@ -1,11 +1,11 @@
 [workspace]
 members = [
-  "msg",
-  "msg-socket",
-  "msg-wire",
-  "msg-transport",
-  "msg-common",
-  "msg-sim",
+    "msg",
+    "msg-socket",
+    "msg-wire",
+    "msg-transport",
+    "msg-common",
+    "msg-sim",
 ]
 resolver = "2"
 
@@ -19,19 +19,19 @@ authors = ["Chainbound Developers <dev@chainbound.io>"]
 homepage = "https://github.com/chainbound/msg-rs"
 repository = "https://github.com/chainbound/msg-rs"
 keywords = [
-  "messaging",
-  "distributed",
-  "systems",
-  "networking",
-  "quic",
-  "quinn",
-  "tokio",
-  "async",
-  "simulation",
-  "pnet",
-  "udp",
-  "tcp",
-  "socket",
+    "messaging",
+    "distributed",
+    "systems",
+    "networking",
+    "quic",
+    "quinn",
+    "tokio",
+    "async",
+    "simulation",
+    "pnet",
+    "udp",
+    "tcp",
+    "socket",
 ]
 
 [workspace.dependencies]
@@ -56,12 +56,14 @@ thiserror = "1"
 tracing = "0.1"
 rustc-hash = "1"
 rand = "0.8"
+libc = "0.2"
 derive_more = { version = "2.0.1", features = [
-  "from",
-  "into",
-  "deref",
-  "deref_mut",
+    "from",
+    "into",
+    "deref",
+    "deref_mut",
 ] }
+arc-swap = "1.7.1"
 
 # networking
 quinn = "0.11.9"

Cargo.toml

@@ -25,6 +25,9 @@ rustc-hash.workspace = true
 tracing.workspace = true
 tokio-stream.workspace = true
 parking_lot.workspace = true
+arc-swap.workspace = true
+
+derive_more = { workspace = true, features = ["deref"] }
 
 [dev-dependencies]
 rand.workspace = true

msg-socket/Cargo.toml

@@ -20,7 +20,7 @@ use tracing::{debug, error, trace};
 use super::{Command, ReqError, ReqOptions};
 use crate::{ConnectionState, ExponentialBackoff, req::SocketState};
 
-use msg_transport::{Address, Transport};
+use msg_transport::{Address, MeteredIo, PeerAddress, Transport};
 use msg_wire::{
     auth,
     compression::{Compressor, try_decompress_payload},
@@ -31,15 +31,16 @@ use msg_wire::{
 type ConnectionTask<Io, Err> = Pin<Box<dyn Future<Output = Result<Io, Err>> + Send>>;
 
 /// A connection controller that manages the connection to a server with an exponential backoff.
-type ConnectionCtl<Io, Addr> = ConnectionState<Framed<Io, reqrep::Codec>, ExponentialBackoff, Addr>;
+type ConnectionCtl<Io, S, A> =
+    ConnectionState<Framed<MeteredIo<Io, S, A>, reqrep::Codec>, ExponentialBackoff, A>;
 
 /// The request socket driver. Endless future that drives
-/// the the socket forward.
+/// the socket forward.
 pub(crate) struct ReqDriver<T: Transport<A>, A: Address> {
     /// Options shared with the socket.
     pub(crate) options: Arc<ReqOptions>,
     /// State shared with the socket.
-    pub(crate) socket_state: Arc<SocketState>,
+    pub(crate) socket_state: SocketState<T::Stats>,
     /// ID counter for outgoing requests.
     pub(crate) id_counter: u32,
     /// Commands from the socket.
@@ -52,7 +53,7 @@ pub(crate) struct ReqDriver<T: Transport<A>, A: Address> {
     pub(crate) conn_task: Option<ConnectionTask<T::Io, T::Error>>,
     /// The transport controller, wrapped in a [`ConnectionState`] for backoff.
     /// The [`Framed`] object can send and receive messages from the socket.
-    pub(crate) conn_state: ConnectionCtl<T::Io, A>,
+    pub(crate) conn_state: ConnectionCtl<T::Io, T::Stats, A>,
     /// The outgoing message queue.
     pub(crate) egress_queue: VecDeque<reqrep::Message>,
     /// The currently pending requests, if any. Uses [`FxHashMap`] for performance.
@@ -276,7 +277,12 @@ where
                     this.conn_task = None;
 
                     if let Ok(io) = result {
-                        let mut framed = Framed::new(io, reqrep::Codec::new());
+                        tracing::debug!(target = ?io.peer_addr(), "new connection");
+
+                        let metered =
+                            MeteredIo::new(io, Arc::clone(&this.socket_state.transport_stats));
+
+                        let mut framed = Framed::new(metered, reqrep::Codec::new());
                         framed.set_backpressure_boundary(this.options.backpressure_boundary);
                         this.conn_state = ConnectionState::Active { channel: framed };
                     }

msg-socket/src/req/driver.rs

@@ -1,5 +1,6 @@
+use arc_swap::ArcSwap;
 use bytes::Bytes;
-use std::time::Duration;
+use std::{sync::Arc, time::Duration};
 use thiserror::Error;
 use tokio::sync::oneshot;
 
@@ -57,7 +58,8 @@ pub struct ReqOptions {
     /// Default is `None`, and the connection is flushed after every send.
     flush_interval: Option<std::time::Duration>,
     /// The maximum number of bytes that can be buffered in the session before being flushed.
-    /// This internally sets [`Framed::set_backpressure_boundary`](tokio_util::codec::Framed).
+    /// This internally sets
+    /// [`Framed::set_backpressure_boundary`](tokio_util::codec::Framed::set_backpressure_boundary).
     backpressure_boundary: usize,
     /// The maximum number of retry attempts. If `None`, the connection will retry indefinitely.
     retry_attempts: Option<usize>,
@@ -180,7 +182,19 @@ impl ReqMessage {
 }
 
 /// The request socket state, shared between the backend task and the socket.
+/// Generic over the transport-level stats type.
 #[derive(Debug, Default)]
-pub(crate) struct SocketState {
-    pub(crate) stats: SocketStats<ReqStats>,
+pub(crate) struct SocketState<S: Default> {
+    /// The socket stats.
+    pub(crate) stats: Arc<SocketStats<ReqStats>>,
+    /// The transport-level stats. We wrap the inner stats in an `Arc`
+    /// for cheap clone on read.
+    pub(crate) transport_stats: Arc<ArcSwap<S>>,
+}
+
+// Manual clone implementation needed here because `S` is not `Clone`.
+impl<S: Default> Clone for SocketState<S> {
+    fn clone(&self) -> Self {
+        Self { stats: Arc::clone(&self.stats), transport_stats: self.transport_stats.clone() }
+    }
 }

msg-socket/src/req/mod.rs

@@ -1,3 +1,4 @@
+use arc_swap::Guard;
 use bytes::Bytes;
 use rustc_hash::FxHashMap;
 use std::{marker::PhantomData, net::SocketAddr, path::PathBuf, sync::Arc, time::Duration};
@@ -13,6 +14,7 @@ use super::{Command, DEFAULT_BUFFER_SIZE, ReqError, ReqOptions};
 use crate::{
     ConnectionState, ExponentialBackoff, ReqMessage,
     req::{SocketState, driver::ReqDriver, stats::ReqStats},
+    stats::SocketStats,
 };
 
 /// The request socket.
@@ -24,7 +26,7 @@ pub struct ReqSocket<T: Transport<A>, A: Address> {
     /// Options for the socket. These are shared with the backend task.
     options: Arc<ReqOptions>,
     /// Socket state. This is shared with the backend task.
-    state: Arc<SocketState>,
+    state: SocketState<T::Stats>,
     /// Optional message compressor. This is shared with the backend task.
     // NOTE: for now we're using dynamic dispatch, since using generics here
     // complicates the API a lot. We can always change this later for perf reasons.
@@ -70,7 +72,7 @@ where
             to_driver: None,
             transport: Some(transport),
             options: Arc::new(options),
-            state: Arc::new(SocketState::default()),
+            state: SocketState::default(),
             compressor: None,
             _marker: PhantomData,
         }
@@ -82,8 +84,14 @@ where
         self
     }
 
-    pub fn stats(&self) -> &ReqStats {
-        &self.state.stats.specific
+    /// Returns the socket stats.
+    pub fn stats(&self) -> &SocketStats<ReqStats> {
+        &self.state.stats
+    }
+
+    /// Get the latest transport-level stats snapshot.
+    pub fn transport_stats(&self) -> Guard<Arc<T::Stats>> {
+        self.state.transport_stats.load()
     }
 
     pub async fn request(&self, message: Bytes) -> Result<Bytes, ReqError> {
@@ -128,7 +136,7 @@ where
         let driver: ReqDriver<T, A> = ReqDriver {
             addr: endpoint,
             options: Arc::clone(&self.options),
-            socket_state: Arc::clone(&self.state),
+            socket_state: self.state.clone(),
             id_counter: 0,
             from_socket,
             transport,

msg-socket/src/req/socket.rs

@@ -1,7 +1,11 @@
 use std::fmt::Debug;
 
-#[derive(Debug)]
+use derive_more::Deref;
+
+/// Statistics for a socket.
+#[derive(Debug, Deref)]
 pub struct SocketStats<S> {
+    /// Socket-specific stats.
     pub(crate) specific: S,
 }
 

msg-socket/src/stats.rs

@@ -158,7 +158,7 @@ impl<A: Address> PubMessage<A> {
 }
 
 /// The subscriber socket state, shared between the backend task and the socket frontend.
-#[derive(Debug)] // Should derive default fine now
+#[derive(Debug)]
 pub(crate) struct SocketState<A: Address> {
     pub(crate) stats: SocketStats<SubStats<A>>,
 }
@@ -184,7 +184,7 @@ mod tests {
     use super::*;
 
     async fn spawn_listener() -> SocketAddr {
-        let listener = TcpListener::bind("0.0.0.0:0").await.unwrap();
+        let listener = TcpListener::bind("[::]:0").await.unwrap();
 
         let addr = listener.local_addr().unwrap();
 

msg-socket/src/sub/mod.rs

@@ -16,22 +16,36 @@ msg-common.workspace = true
 async-trait.workspace = true
 futures.workspace = true
 tokio.workspace = true
-tokio-openssl.workspace = true
 tracing.workspace = true
-thiserror.workspace = true
-derive_more.workspace = true
+arc-swap.workspace = true
+
+libc.workspace = true
+
+thiserror = { workspace = true, optional = true }
+derive_more = { workspace = true, features = [
+    "from",
+    "into",
+    "deref",
+    "deref_mut",
+], optional = true }
 
 # QUIC
 quinn = { workspace = true, optional = true }
 rcgen = { workspace = true, optional = true }
 
 # TLS
 openssl = { workspace = true, optional = true }
+tokio-openssl = { workspace = true, optional = true }
 
 [dev-dependencies]
 tracing-subscriber = "0.3"
 
 [features]
 default = []
-quic = ["dep:quinn", "dep:rcgen"]
-tcp-tls = ["dep:openssl"]
+quic = ["dep:quinn", "dep:rcgen", "dep:thiserror"]
+tcp-tls = [
+    "dep:openssl",
+    "dep:derive_more",
+    "dep:thiserror",
+    "dep:tokio-openssl",
+]