feat(socket): initial scaffold for client/server ConnManager

Author: mempirate

Cargo.toml

@@ -173,6 +173,7 @@ unused_rounding = "warn"
 use_self = "warn"
 useless_let_if_seq = "warn"
 zero_sized_map_values = "warn"
+default_trait_access = "warn"
 
 # These are nursery lints which have findings. Allow them for now. Some are not
 # quite mature enough for use in our codebase and some we don't really want.

msg-socket/src/connection/backoff.rs

@@ -6,7 +6,7 @@ use std::{
 };
 use tokio::time::sleep;
 
-use crate::ConnOptions;
+use crate::ClientOptions;
 
 /// Helper trait alias for backoff streams.
 /// We define any stream that yields `Duration`s as a backoff
@@ -41,8 +41,8 @@ impl ExponentialBackoff {
     }
 }
 
-impl From<&ConnOptions> for ExponentialBackoff {
-    fn from(options: &ConnOptions) -> Self {
+impl From<&ClientOptions> for ExponentialBackoff {
+    fn from(options: &ClientOptions) -> Self {
         Self::new(options.backoff_duration, options.retry_attempts)
     }
 }

msg-socket/src/req/conn_manager.rs

@@ -5,13 +5,14 @@ use std::{
     task::{Context, Poll},
 };
 
+use arc_swap::ArcSwap;
 use bytes::Bytes;
 use futures::{Future, FutureExt, SinkExt, StreamExt};
 use msg_common::span::{EnterSpan as _, WithSpan};
 use tokio_util::codec::Framed;
 use tracing::Instrument;
 
-use crate::{ConnOptions, ConnectionState, ExponentialBackoff};
+use crate::{ClientOptions, ConnectionState, ExponentialBackoff};
 
 use msg_transport::{Address, MeteredIo, Transport};
 use msg_wire::{auth, reqrep};
@@ -33,10 +34,227 @@ pub(crate) type Conn<Io, S, A> = Framed<MeteredIo<Io, S, A>, reqrep::Codec>;
 /// A connection controller that manages the connection to a server with an exponential backoff.
 pub(crate) type ConnCtl<Io, S, A> = ConnectionState<Conn<Io, S, A>, ExponentialBackoff, A>;
 
+/// A connection manager for managing client OR server connections.
+/// The type parameter `S` contains the connection state, including its "side" (client / server).
+pub(crate) struct ConnectionManager<T, A, S>
+where
+    T: Transport<A>,
+    A: Address,
+{
+    /// The connection state, including its "side" (client / server).
+    state: S,
+    /// The transport used for the connection.
+    transport: T,
+    /// Transport stats for metering IO.
+    transport_stats: Arc<ArcSwap<T::Stats>>,
+
+    /// Connection manager tracing span.
+    span: tracing::Span,
+}
+
+/// A client connection to a remote server.
+pub(crate) struct ClientConnection<T, A>
+where
+    T: Transport<A>,
+    A: Address,
+{
+    /// Options for the connection manager.
+    options: ClientOptions,
+    /// The address of the remote.
+    addr: A,
+    /// The connection task which handles the connection to the server.
+    conn_task: Option<WithSpan<ConnTask<T::Io, T::Error>>>,
+    /// The transport controller, wrapped in a [`ConnectionState`] for backoff.
+    /// The [`Framed`] object can send and receive messages from the socket.
+    conn_ctl: ConnCtl<T::Io, T::Stats, A>,
+}
+
+impl<T, A> ClientConnection<T, A>
+where
+    T: Transport<A>,
+    A: Address,
+{
+    /// Reset the connection state to inactive, so that it will be re-tried.
+    ///
+    /// This is done when the connection is closed or an error occurs.
+    #[inline]
+    pub(crate) fn reset_connection(&mut self) {
+        self.conn_ctl = ConnectionState::Inactive {
+            addr: self.addr.clone(),
+            backoff: ExponentialBackoff::from(&self.options),
+        };
+    }
+
+    /// Returns a mutable reference to the connection channel if it is active.
+    #[inline]
+    pub fn active_connection(&mut self) -> Option<&mut Conn<T::Io, T::Stats, A>> {
+        if let ConnectionState::Active { ref mut channel } = self.conn_ctl {
+            Some(channel)
+        } else {
+            None
+        }
+    }
+}
+
+/// A local server connection. Manages the connection lifecycle:
+/// - Accepting incoming connections.
+/// - Handling established connections.
+pub(crate) struct ServerConnection<T, A>
+where
+    T: Transport<A>,
+    A: Address,
+{
+    /// The local address.
+    addr: A,
+    /// The accept task which handles accepting an incoming connection.
+    accept_task: Option<WithSpan<T::Accept>>,
+    /// The inbound connection.
+    conn: Conn<T::Io, T::Stats, A>,
+}
+
+impl<T, A> ConnectionManager<T, A, ClientConnection<T, A>>
+where
+    T: Transport<A>,
+    A: Address,
+{
+    pub(crate) fn new(
+        options: ClientOptions,
+        transport: T,
+        addr: A,
+        conn_ctl: ConnCtl<T::Io, T::Stats, A>,
+        transport_stats: Arc<ArcSwap<T::Stats>>,
+        span: tracing::Span,
+    ) -> Self {
+        let conn = ClientConnection { options, addr, conn_task: None, conn_ctl };
+
+        Self { state: conn, transport, transport_stats, span }
+    }
+
+    /// Start the connection task to the server, handling authentication if necessary.
+    /// The result will be polled by the driver and re-tried according to the backoff policy.
+    fn try_connect(&mut self) {
+        let connect = self.transport.connect(self.state.addr.clone());
+        let token = self.state.options.auth_token.clone();
+
+        let task = async move {
+            let io = connect.await?;
+
+            let Some(token) = token else {
+                return Ok(io);
+            };
+
+            authentication_handshake::<T, A>(io, token).await
+        }
+        .in_current_span();
+
+        // FIX: coercion to BoxFuture for [`SpanExt::with_current_span`]
+        self.state.conn_task = Some(WithSpan::current(Box::pin(task)));
+    }
+
+    /// Reset the connection state to inactive, so that it will be re-tried.
+    ///
+    /// This is done when the connection is closed or an error occurs.
+    #[inline]
+    pub(crate) fn reset_connection(&mut self) {
+        self.state.reset_connection();
+    }
+
+    /// Poll connection management logic: connection task, backoff, and retry logic.
+    /// Loops until the connection is active, then returns a mutable reference to the channel.
+    ///
+    /// Note: this is not a `Future` impl because we want to return a reference; doing it in
+    /// a `Future` would require lifetime headaches or unsafe code.
+    ///
+    /// Returns:
+    /// * `Poll::Ready(Some(&mut channel))` if the connection is active
+    /// * `Poll::Ready(None)` if we should terminate (max retries exceeded)
+    /// * `Poll::Pending` if we need to wait for backoff
+    #[allow(clippy::type_complexity)]
+    pub(crate) fn poll(
+        &mut self,
+        cx: &mut Context<'_>,
+    ) -> Poll<Option<&mut Conn<T::Io, T::Stats, A>>> {
+        loop {
+            // Poll the active connection task, if any
+            if let Some(ref mut conn_task) = self.state.conn_task {
+                if let Poll::Ready(result) = conn_task.poll_unpin(cx).enter() {
+                    // As soon as the connection task finishes, set it to `None`.
+                    // - If it was successful, set the connection to active
+                    // - If it failed, it will be re-tried until the backoff limit is reached.
+                    self.state.conn_task = None;
+
+                    match result.inner {
+                        Ok(io) => {
+                            tracing::info!("connected");
+
+                            let metered = MeteredIo::new(io, self.transport_stats.clone());
+                            let framed = Framed::new(metered, reqrep::Codec::new());
+                            self.state.conn_ctl = ConnectionState::Active { channel: framed };
+                        }
+                        Err(e) => {
+                            tracing::error!(?e, "failed to connect");
+                        }
+                    }
+                }
+            }
+
+            // If the connection is inactive, try to connect to the server or poll the backoff
+            // timer if we're already trying to connect.
+            if let ConnectionState::Inactive { backoff, .. } = &mut self.state.conn_ctl {
+                let Poll::Ready(item) = backoff.poll_next_unpin(cx) else {
+                    return Poll::Pending;
+                };
+
+                let _span = tracing::info_span!(parent: &self.span, "connect").entered();
+
+                if let Some(duration) = item {
+                    if self.state.conn_task.is_none() {
+                        tracing::debug!(backoff = ?duration, "trying connection");
+                        self.try_connect();
+                    } else {
+                        tracing::debug!(
+                            backoff = ?duration,
+                            "not retrying as there is already a connection task"
+                        );
+                    }
+                } else {
+                    tracing::error!("exceeded maximum number of retries, terminating connection");
+                    return Poll::Ready(None);
+                }
+            }
+
+            if let ConnectionState::Active { ref mut channel } = self.state.conn_ctl {
+                return Poll::Ready(Some(channel));
+            }
+        }
+    }
+}
+
+pub struct ServerOptions {}
+
+impl<T, A> ConnectionManager<T, A, ServerConnection<T, A>>
+where
+    T: Transport<A>,
+    A: Address,
+{
+    pub(crate) fn new(
+        options: ServerOptions,
+        transport: T,
+        addr: A,
+        conn: Conn<T::Io, T::Stats, A>,
+        transport_stats: Arc<ArcSwap<T::Stats>>,
+        span: tracing::Span,
+    ) -> Self {
+        let conn = ServerConnection { addr, accept_task: None, conn };
+
+        Self { state: conn, transport, transport_stats, span }
+    }
+}
+
 /// Manages the connection lifecycle: connecting, reconnecting, and maintaining the connection.
 pub(crate) struct ConnManager<T: Transport<A>, A: Address> {
     /// Options for the connection manager.
-    options: ConnOptions,
+    options: ClientOptions,
     /// The connection task which handles the connection to the server.
     conn_task: Option<WithSpan<ConnTask<T::Io, T::Error>>>,
     /// The transport controller, wrapped in a [`ConnectionState`] for backoff.
@@ -89,7 +307,7 @@ where
     A: Address,
 {
     pub(crate) fn new(
-        options: ConnOptions,
+        options: ClientOptions,
         transport: T,
         addr: A,
         conn_ctl: ConnCtl<T::Io, T::Stats, A>,

msg-socket/src/req/driver.rs

@@ -27,9 +27,22 @@ use msg_wire::{
     reqrep,
 };
 
+/// Type state for a client connection.
+struct ClientConnection<T, A>
+where
+    T: Transport<A>,
+    A: Address,
+{
+    conn_manager: ConnManager<T, A>,
+}
+
 /// The request socket driver. Endless future that drives
 /// the socket forward.
-pub(crate) struct ReqDriver<T: Transport<A>, A: Address> {
+pub(crate) struct ReqDriver<T, A>
+where
+    T: Transport<A>,
+    A: Address,
+{
     /// Options shared with the socket.
     pub(crate) options: Arc<ReqOptions>,
     /// State shared with the socket.

msg-socket/src/req/mod.rs

@@ -68,7 +68,7 @@ impl SendCommand {
 
 /// Options for the connection manager.
 #[derive(Debug, Clone)]
-pub struct ConnOptions {
+pub struct ClientOptions {
     /// Optional authentication token.
     pub auth_token: Option<Bytes>,
     /// The backoff duration for the underlying transport on reconnections.
@@ -77,7 +77,7 @@ pub struct ConnOptions {
     pub retry_attempts: Option<usize>,
 }
 
-impl Default for ConnOptions {
+impl Default for ClientOptions {
     fn default() -> Self {
         Self {
             auth_token: None,
@@ -96,7 +96,7 @@ impl Default for ConnOptions {
 #[derive(Debug, Clone)]
 pub struct ReqOptions {
     /// Options for the connection manager.
-    pub conn: ConnOptions,
+    pub conn: ClientOptions,
     /// Timeout duration for requests.
     pub timeout: Duration,
     /// Wether to block on initial connection to the target.
@@ -211,7 +211,7 @@ impl ReqOptions {
 impl Default for ReqOptions {
     fn default() -> Self {
         Self {
-            conn: ConnOptions::default(),
+            conn: ClientOptions::default(),
             timeout: Duration::from_secs(5),
             blocking_connect: false,
             min_compress_size: 8192,