transport: Make accept async to close the gap on service races

Signed-off-by: Alexandru Vasile <[email protected]>

Author: lexnv

src/transport/dummy.rs

@@ -25,7 +25,7 @@ use crate::{
     types::ConnectionId,
 };
 
-use futures::Stream;
+use futures::{future::BoxFuture, Stream};
 use multiaddr::Multiaddr;
 
 use std::{
@@ -71,8 +71,8 @@ impl Transport for DummyTransport {
         Ok(())
     }
 
-    fn accept(&mut self, _: ConnectionId) -> crate::Result<()> {
-        Ok(())
+    fn accept(&mut self, _: ConnectionId) -> crate::Result<BoxFuture<'static, crate::Result<()>>> {
+        Ok(Box::pin(async { Ok(()) }))
     }
 
     fn accept_pending(&mut self, _connection_id: ConnectionId) -> crate::Result<()> {

src/transport/manager/mod.rs

@@ -39,7 +39,7 @@ use crate::{
 };
 
 use address::{scores, AddressStore};
-use futures::{Stream, StreamExt};
+use futures::{future::BoxFuture, Stream, StreamExt};
 use indexmap::IndexMap;
 use multiaddr::{Multiaddr, Protocol};
 use multihash::Multihash;
@@ -252,6 +252,11 @@ pub struct TransportManager {
 
     /// Opening connections errors.
     opening_errors: HashMap<ConnectionId, Vec<(Multiaddr, DialError)>>,
+
+    /// Pending accept future with associated connection information.
+    /// When a connection is accepted, we must wait for the accept future to complete
+    /// (which notifies all protocols) before emitting the ConnectionEstablished event.
+    pending_accept: Option<(PeerId, Endpoint, BoxFuture<'static, crate::Result<()>>)>,
 }
 
 /// Builder for [`crate::transport::manager::TransportManager`].
@@ -365,6 +370,7 @@ impl TransportManagerBuilder {
             pending_connections: HashMap::new(),
             connection_limits: limits::ConnectionLimits::new(self.connection_limits_config),
             opening_errors: HashMap::new(),
+            pending_accept: None,
         }
     }
 }
@@ -1090,6 +1096,35 @@ impl TransportManager {
     /// Poll next event from [`crate::transport::manager::TransportManager`].
     pub async fn next(&mut self) -> Option<TransportEvent> {
         loop {
+            // First, check if we have a pending accept future to poll
+            if let Some((peer, endpoint, mut future)) = self.pending_accept.take() {
+                match future.as_mut().await {
+                    Ok(()) => {
+                        tracing::trace!(
+                            target: LOG_TARGET,
+                            ?peer,
+                            ?endpoint,
+                            "connection accepted and protocols notified",
+                        );
+
+                        return Some(TransportEvent::ConnectionEstablished {
+                            peer,
+                            endpoint,
+                        });
+                    }
+                    Err(error) => {
+                        tracing::debug!(
+                            target: LOG_TARGET,
+                            ?peer,
+                            ?endpoint,
+                            ?error,
+                            "failed to notify protocols about connection",
+                        );
+                        // If notification failed, we don't emit the ConnectionEstablished event
+                    }
+                }
+            }
+
             tokio::select! {
                 event = self.event_rx.recv() => {
                     let Some(event) = event else {
@@ -1270,16 +1305,27 @@ impl TransportManager {
                                         "accept connection",
                                     );
 
-                                    let _ = self
+                                    match self
                                         .transports
                                         .get_mut(&transport)
                                         .expect("transport to exist")
-                                        .accept(endpoint.connection_id());
-
-                                    return Some(TransportEvent::ConnectionEstablished {
-                                        peer,
-                                        endpoint,
-                                    });
+                                        .accept(endpoint.connection_id())
+                                    {
+                                        Ok(future) => {
+                                            // Store the accept future to be polled in the next iteration
+                                            // This ensures protocols are notified before we emit ConnectionEstablished
+                                            self.pending_accept = Some((peer, endpoint, future));
+                                        }
+                                        Err(error) => {
+                                            tracing::debug!(
+                                                target: LOG_TARGET,
+                                                ?peer,
+                                                ?endpoint,
+                                                ?error,
+                                                "failed to accept connection",
+                                            );
+                                        }
+                                    }
                                 }
                                 Ok(ConnectionEstablishedResult::Reject) => {
                                     tracing::trace!(

src/transport/mod.rs

@@ -22,7 +22,7 @@
 
 use crate::{error::DialError, transport::manager::TransportHandle, types::ConnectionId, PeerId};
 
-use futures::Stream;
+use futures::{future::BoxFuture, Stream};
 use hickory_resolver::TokioResolver;
 use multiaddr::Multiaddr;
 
@@ -194,7 +194,12 @@ pub(crate) trait Transport: Stream + Unpin + Send {
     fn dial(&mut self, connection_id: ConnectionId, address: Multiaddr) -> crate::Result<()>;
 
     /// Accept negotiated connection.
-    fn accept(&mut self, connection_id: ConnectionId) -> crate::Result<()>;
+    ///
+    /// Returns a future that completes when the connection has been fully established
+    /// and all installed protocols have been notified via their event channels.
+    /// This ensures that by the time the caller receives a ConnectionEstablished event,
+    /// protocols are ready to handle substream operations.
+    fn accept(&mut self, connection_id: ConnectionId) -> crate::Result<BoxFuture<'static, crate::Result<()>>>;
 
     /// Accept pending connection.
     fn accept_pending(&mut self, connection_id: ConnectionId) -> crate::Result<()>;

src/transport/quic/connection.rs

@@ -237,6 +237,12 @@ impl QuicConnection {
             .report_connection_established(self.peer, self.endpoint.clone())
             .await?;
 
+        self.start_event_loop().await
+    }
+
+    /// Start the connection event loop without notifying protocols.
+    /// This is used when protocols have already been notified during accept().
+    pub(crate) async fn start_event_loop(mut self) -> crate::Result<()> {
         loop {
             tokio::select! {
                 event = self.connection.accept_bi() => match event {

src/transport/quic/mod.rs

@@ -306,35 +306,47 @@ impl Transport for QuicTransport {
         Ok(())
     }
 
-    fn accept(&mut self, connection_id: ConnectionId) -> crate::Result<()> {
+    fn accept(&mut self, connection_id: ConnectionId) -> crate::Result<BoxFuture<'static, crate::Result<()>>> {
         let (connection, endpoint) = self
             .pending_open
             .remove(&connection_id)
             .ok_or(Error::ConnectionDoesntExist(connection_id))?;
         let bandwidth_sink = self.context.bandwidth_sink.clone();
-        let protocol_set = self.context.protocol_set(connection_id);
+        let mut protocol_set = self.context.protocol_set(connection_id);
         let substream_open_timeout = self.config.substream_open_timeout;
+        let executor = self.context.executor.clone();
 
         tracing::trace!(
             target: LOG_TARGET,
             ?connection_id,
             "start connection",
         );
 
-        self.context.executor.run(Box::pin(async move {
-            let _ = QuicConnection::new(
-                connection.peer,
-                endpoint,
-                connection.connection,
-                protocol_set,
-                bandwidth_sink,
-                substream_open_timeout,
-            )
-            .start()
-            .await;
-        }));
+        let peer = connection.peer;
+        let endpoint_clone = endpoint.clone();
+
+        Ok(Box::pin(async move {
+            // First, notify all protocols about the connection establishment
+            protocol_set
+                .report_connection_established(peer, endpoint_clone)
+                .await?;
+
+            // After protocols are notified, spawn the connection event loop
+            executor.run(Box::pin(async move {
+                let _ = QuicConnection::new(
+                    peer,
+                    endpoint,
+                    connection.connection,
+                    protocol_set,
+                    bandwidth_sink,
+                    substream_open_timeout,
+                )
+                .start_event_loop()
+                .await;
+            }));
 
-        Ok(())
+            Ok(())
+        }))
     }
 
     fn reject(&mut self, connection_id: ConnectionId) -> crate::Result<()> {

src/transport/tcp/connection.rs

@@ -736,6 +736,12 @@ impl TcpConnection {
             .report_connection_established(self.peer, self.endpoint.clone())
             .await?;
 
+        self.start_event_loop().await
+    }
+
+    /// Start the connection event loop without notifying protocols.
+    /// This is used when protocols have already been notified during accept().
+    pub(crate) async fn start_event_loop(mut self) -> crate::Result<()> {
         loop {
             tokio::select! {
                 substream = self.connection.next() => {

src/transport/tcp/mod.rs

@@ -389,37 +389,50 @@ impl Transport for TcpTransport {
         )
     }
 
-    fn accept(&mut self, connection_id: ConnectionId) -> crate::Result<()> {
+    fn accept(&mut self, connection_id: ConnectionId) -> crate::Result<BoxFuture<'static, crate::Result<()>>> {
         let context = self
             .pending_open
             .remove(&connection_id)
             .ok_or(Error::ConnectionDoesntExist(connection_id))?;
-        let protocol_set = self.context.protocol_set(connection_id);
+        let mut protocol_set = self.context.protocol_set(connection_id);
         let bandwidth_sink = self.context.bandwidth_sink.clone();
         let next_substream_id = self.context.next_substream_id.clone();
+        let executor = self.context.executor.clone();
 
         tracing::trace!(
             target: LOG_TARGET,
             ?connection_id,
             "start connection",
         );
 
-        self.context.executor.run(Box::pin(async move {
-            if let Err(error) =
-                TcpConnection::new(context, protocol_set, bandwidth_sink, next_substream_id)
-                    .start()
-                    .await
-            {
-                tracing::debug!(
-                    target: LOG_TARGET,
-                    ?connection_id,
-                    ?error,
-                    "connection exited with error",
-                );
-            }
-        }));
+        let peer = context.peer();
+        let endpoint = context.endpoint().clone();
+
+        Ok(Box::pin(async move {
+            // First, notify all protocols about the connection establishment
+            // This ensures that when the accept() future completes, protocols are ready
+            protocol_set
+                .report_connection_established(peer, endpoint)
+                .await?;
+
+            // After protocols are notified, spawn the connection event loop
+            executor.run(Box::pin(async move {
+                if let Err(error) =
+                    TcpConnection::new(context, protocol_set, bandwidth_sink, next_substream_id)
+                        .start_event_loop()
+                        .await
+                {
+                    tracing::debug!(
+                        target: LOG_TARGET,
+                        ?connection_id,
+                        ?error,
+                        "connection exited with error",
+                    );
+                }
+            }));
 
-        Ok(())
+            Ok(())
+        }))
     }
 
     fn reject(&mut self, connection_id: ConnectionId) -> crate::Result<()> {

src/transport/webrtc/connection.rs

@@ -680,6 +680,12 @@ impl WebRtcConnection {
             .report_connection_established(self.peer, self.endpoint.clone())
             .await;
 
+        self.run_event_loop().await;
+    }
+
+    /// Start the connection event loop without notifying protocols.
+    /// This is used when protocols have already been notified during accept().
+    pub async fn run_event_loop(mut self) {
         loop {
             // poll output until we get a timeout
             let timeout = match self.rtc.poll_output().unwrap() {

src/transport/webrtc/mod.rs

@@ -532,7 +532,7 @@ impl Transport for WebRtcTransport {
         ))
     }
 
-    fn accept(&mut self, connection_id: ConnectionId) -> crate::Result<()> {
+    fn accept(&mut self, connection_id: ConnectionId) -> crate::Result<BoxFuture<'static, crate::Result<()>>> {
         tracing::trace!(
             target: LOG_TARGET,
             ?connection_id,
@@ -562,19 +562,13 @@ impl Transport for WebRtcTransport {
 
         let rtc = connection.on_accept()?;
         let (tx, rx) = channel(self.datagram_buffer_size);
-        let protocol_set = self.context.protocol_set(connection_id);
+        let mut protocol_set = self.context.protocol_set(connection_id);
         let connection_id = endpoint.connection_id();
+        let endpoint_clone = endpoint.clone();
+        let executor = self.context.executor.clone();
+        let socket = Arc::clone(&self.socket);
+        let listen_address = self.listen_address;
 
-        let connection = WebRtcConnection::new(
-            rtc,
-            peer,
-            source,
-            self.listen_address,
-            Arc::clone(&self.socket),
-            protocol_set,
-            endpoint,
-            rx,
-        );
         self.open.insert(
             source,
             ConnectionContext {
@@ -584,11 +578,30 @@ impl Transport for WebRtcTransport {
             },
         );
 
-        self.context.executor.run(Box::pin(async move {
-            connection.run().await;
-        }));
+        Ok(Box::pin(async move {
+            // First, notify all protocols about the connection establishment
+            protocol_set
+                .report_connection_established(peer, endpoint_clone)
+                .await?;
 
-        Ok(())
+            // After protocols are notified, create connection and spawn event loop
+            let connection = WebRtcConnection::new(
+                rtc,
+                peer,
+                source,
+                listen_address,
+                socket,
+                protocol_set,
+                endpoint,
+                rx,
+            );
+
+            executor.run(Box::pin(async move {
+                connection.run_event_loop().await;
+            }));
+
+            Ok(())
+        }))
     }
 
     fn reject(&mut self, connection_id: ConnectionId) -> crate::Result<()> {

src/transport/websocket/connection.rs

@@ -456,9 +456,15 @@ impl WebSocketConnection {
     /// Start connection event loop.
     pub(crate) async fn start(mut self) -> crate::Result<()> {
         self.protocol_set
-            .report_connection_established(self.peer, self.endpoint)
+            .report_connection_established(self.peer, self.endpoint.clone())
             .await?;
 
+        self.start_event_loop().await
+    }
+
+    /// Start the connection event loop without notifying protocols.
+    /// This is used when protocols have already been notified during accept().
+    pub(crate) async fn start_event_loop(mut self) -> crate::Result<()> {
         loop {
             tokio::select! {
                 substream = self.connection.next() => match substream {