Report SubstreamOpenFailure when switching to secondary connection

src/protocol/request_response/handle.rs

@@ -94,6 +94,7 @@ impl From<SubstreamError> for RejectReason {
     fn from(error: SubstreamError) -> Self {
         // Convert `ErrorKind::NotConnected` to `RejectReason::ConnectionClosed`.
         match error {
+            SubstreamError::ConnectionClosed => RejectReason::ConnectionClosed,
             SubstreamError::IoError(ErrorKind::NotConnected) => RejectReason::ConnectionClosed,
             SubstreamError::YamuxError(crate::yamux::ConnectionError::Io(error), _)
                 if error.kind() == ErrorKind::NotConnected =>

src/protocol/request_response/tests.rs

@@ -21,23 +21,25 @@
 use crate::{
     mock::substream::{DummySubstream, MockSubstream},
     protocol::{
+        connection::ConnectionHandle,
         request_response::{
-            ConfigBuilder, DialOptions, RequestResponseError, RequestResponseEvent,
+            ConfigBuilder, DialOptions, RejectReason, RequestResponseError, RequestResponseEvent,
             RequestResponseHandle, RequestResponseProtocol,
         },
-        InnerTransportEvent, SubstreamError, SubstreamKeepAlive, TransportService,
+        InnerTransportEvent, ProtocolCommand, SubstreamError, SubstreamKeepAlive, TransportService,
     },
     substream::Substream,
     transport::{
         manager::{TransportManager, TransportManagerBuilder},
-        KEEP_ALIVE_TIMEOUT,
+        Endpoint, KEEP_ALIVE_TIMEOUT,
     },
-    types::{RequestId, SubstreamId},
+    types::{ConnectionId, RequestId, SubstreamId},
     Error, PeerId, ProtocolName,
 };
 
 use futures::StreamExt;
-use tokio::sync::mpsc::Sender;
+use multiaddr::Multiaddr;
+use tokio::sync::mpsc::{channel, Sender};
 
 use std::task::Poll;
 
@@ -299,3 +301,74 @@ async fn request_failure_reported_once() {
     })
     .await;
 }
+
+// when the primary connection is closed while a secondary connection exists, the protocol is not
+// notified about the closed connection, but a request that was already issued over it can never
+// succeed and must be reported as failed
+#[tokio::test]
+async fn request_over_replaced_primary_connection_fails() {
+    let _ = tracing_subscriber::fmt()
+        .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
+        .try_init();
+
+    let (protocol, mut handle, _manager, tx) = protocol();
+    let peer = PeerId::random();
+
+    // register primary and secondary connection before the event loop is started so that they
+    // are handled before the request, `RequestResponseProtocol` polls network events first
+    let (cmd_tx1, mut cmd_rx1) = channel(64);
+    tx.send(InnerTransportEvent::ConnectionEstablished {
+        peer,
+        connection: ConnectionId::from(0usize),
+        endpoint: Endpoint::dialer(Multiaddr::empty(), ConnectionId::from(0usize)),
+        sender: ConnectionHandle::new(ConnectionId::from(0usize), cmd_tx1),
+    })
+    .await
+    .unwrap();
+
+    let (cmd_tx2, _cmd_rx2) = channel(64);
+    tx.send(InnerTransportEvent::ConnectionEstablished {
+        peer,
+        connection: ConnectionId::from(1usize),
+        endpoint: Endpoint::dialer(Multiaddr::empty(), ConnectionId::from(1usize)),
+        sender: ConnectionHandle::new(ConnectionId::from(1usize), cmd_tx2),
+    })
+    .await
+    .unwrap();
+
+    tokio::spawn(protocol.run());
+
+    let request_id =
+        handle.send_request(peer, vec![1, 3, 3, 7], DialOptions::Reject).await.unwrap();
+
+    // the request is issued over the primary connection and the transport never reports back
+    match cmd_rx1.recv().await {
+        Some(ProtocolCommand::OpenSubstream { connection_id, .. }) => {
+            assert_eq!(connection_id, ConnectionId::from(0usize));
+        }
+        command => panic!("unexpected command: {command:?}"),
+    }
+
+    tx.send(InnerTransportEvent::ConnectionClosed {
+        peer,
+        connection: ConnectionId::from(0usize),
+    })
+    .await
+    .unwrap();
+
+    match handle.next().await {
+        Some(RequestResponseEvent::RequestFailed {
+            peer: failed_peer,
+            request_id: failed_request,
+            error,
+        }) => {
+            assert_eq!(failed_peer, peer);
+            assert_eq!(failed_request, request_id);
+            assert!(std::matches!(
+                error,
+                RequestResponseError::Rejected(RejectReason::ConnectionClosed)
+            ));
+        }
+        event => panic!("unexpected event: {event:?}"),
+    }
+}

src/protocol/transport_service.rs

@@ -21,7 +21,7 @@
 use crate::{
     addresses::PublicAddresses,
     error::{Error, ImmediateDialError, SubstreamError},
-    protocol::{connection::ConnectionHandle, InnerTransportEvent, TransportEvent},
+    protocol::{connection::ConnectionHandle, Direction, InnerTransportEvent, TransportEvent},
     transport::{manager::TransportManagerHandle, Endpoint},
     types::{protocol::ProtocolName, ConnectionId, SubstreamId},
     PeerId, DEFAULT_CHANNEL_SIZE,
@@ -32,7 +32,7 @@ use multiaddr::{Multiaddr, Protocol};
 use tokio::sync::mpsc::{channel, Receiver, Sender};
 
 use std::{
-    collections::{HashMap, HashSet},
+    collections::{HashMap, HashSet, VecDeque},
     fmt::Debug,
     pin::Pin,
     sync::{
@@ -306,6 +306,13 @@ pub struct TransportService {
 
     /// Whether this protocol susbstreams should keep connection alive.
     substream_keep_alive: SubstreamKeepAlive,
+
+    /// Outbound substreams whose result hasn't been reported yet, mapped to the connection
+    /// they were requested on.
+    pending_outbound: HashMap<SubstreamId, ConnectionId>,
+
+    /// Events queued for the protocol, since closing a connection may emit multiple events.
+    pending_events: VecDeque<TransportEvent>,
 }
 
 impl TransportService {
@@ -334,6 +341,8 @@ impl TransportService {
                 connections: HashMap::new(),
                 keep_alive_tracker,
                 substream_keep_alive,
+                pending_outbound: HashMap::new(),
+                pending_events: VecDeque::new(),
             },
             tx,
         )
@@ -433,6 +442,20 @@ impl TransportService {
 
         self.keep_alive_tracker.on_connection_closed(peer, connection_id);
 
+        // The connection is gone and will never report these back. This is the only notification
+        // the protocol gets when the connection is replaced by the secondary one, since no
+        // `ConnectionClosed` is emitted in that case.
+        let (failed, pending) = std::mem::take(&mut self.pending_outbound)
+            .into_iter()
+            .partition::<HashMap<_, _>, _>(|(_, id)| id == &connection_id);
+        self.pending_outbound = pending;
+        self.pending_events.extend(failed.into_keys().map(|substream| {
+            TransportEvent::SubstreamOpenFailure {
+                substream,
+                error: SubstreamError::ConnectionClosed,
+            }
+        }));
+
         let Some(context) = self.connections.get_mut(&peer) else {
             tracing::warn!(
                 target: LOG_TARGET,
@@ -590,15 +613,19 @@ impl TransportService {
             connection.try_upgrade();
         }
 
-        connection
-            .open_substream(
-                self.protocol.clone(),
-                self.fallback_names.clone(),
-                substream_id,
-                permit,
-                self.substream_keep_alive,
-            )
-            .map(|_| substream_id)
+        let result = connection.open_substream(
+            self.protocol.clone(),
+            self.fallback_names.clone(),
+            substream_id,
+            permit,
+            self.substream_keep_alive,
+        );
+
+        if result.is_ok() {
+            self.pending_outbound.insert(substream_id, connection_id);
+        }
+
+        result.map(|_| substream_id)
     }
 
     /// Forcibly close the connection, even if other protocols have substreams open over it.
@@ -642,6 +669,10 @@ impl Stream for TransportService {
         let protocol_name = self.protocol.clone();
         let keep_alive_timeout = self.keep_alive_tracker.keep_alive_timeout;
 
+        if let Some(event) = self.pending_events.pop_front() {
+            return Poll::Ready(Some(event));
+        }
+
         while let Poll::Ready(event) = self.rx.poll_recv(cx) {
             match event {
                 None => {
@@ -665,7 +696,11 @@ impl Stream for TransportService {
                     }
                 }
                 Some(InnerTransportEvent::ConnectionClosed { peer, connection }) => {
-                    if let Some(event) = self.on_connection_closed(peer, connection) {
+                    let closed = self.on_connection_closed(peer, connection);
+                    self.pending_events.extend(closed);
+
+                    // Returning `Ready` guarantees another poll which drains the rest.
+                    if let Some(event) = self.pending_events.pop_front() {
                         return Poll::Ready(Some(event));
                     }
                 }
@@ -678,6 +713,10 @@ impl Stream for TransportService {
                     connection_id,
                     opening_permit,
                 }) => {
+                    if let Direction::Outbound(substream_id) = direction {
+                        self.pending_outbound.remove(&substream_id);
+                    }
+
                     if protocol == self.protocol
                         && self.substream_keep_alive == SubstreamKeepAlive::Yes
                     {
@@ -699,6 +738,14 @@ impl Stream for TransportService {
                         substream,
                     }));
                 }
+                Some(InnerTransportEvent::SubstreamOpenFailure { substream, error }) => {
+                    self.pending_outbound.remove(&substream);
+
+                    return Poll::Ready(Some(TransportEvent::SubstreamOpenFailure {
+                        substream,
+                        error,
+                    }));
+                }
                 Some(event) => return Poll::Ready(Some(event.into())),
             }
         }
@@ -728,7 +775,9 @@ impl Stream for TransportService {
 mod tests {
     use super::*;
     use crate::{
-        protocol::{ProtocolCommand, SubstreamKeepAlive, TransportService},
+        mock::substream::MockSubstream,
+        protocol::{connection::Permit, ProtocolCommand, SubstreamKeepAlive, TransportService},
+        substream::Substream,
         transport::{
             manager::{handle::InnerTransportManagerCommand, TransportManagerHandle},
             KEEP_ALIVE_TIMEOUT,
@@ -1080,6 +1129,133 @@ mod tests {
         assert!(cmd_rx2.try_recv().is_err());
     }
 
+    #[tokio::test]
+    async fn pending_substreams_fail_when_switching_to_secondary_connection() {
+        let (mut service, sender, _) = transport_service();
+        let peer = PeerId::random();
+
+        // register primary connection
+        let (cmd_tx1, mut cmd_rx1) = channel(64);
+        sender
+            .send(InnerTransportEvent::ConnectionEstablished {
+                peer,
+                connection: ConnectionId::from(0usize),
+                endpoint: Endpoint::dialer(Multiaddr::empty(), ConnectionId::from(0usize)),
+                sender: ConnectionHandle::new(ConnectionId::from(0usize), cmd_tx1.clone()),
+            })
+            .await
+            .unwrap();
+
+        assert!(std::matches!(
+            service.next().await,
+            Some(TransportEvent::ConnectionEstablished { .. })
+        ));
+
+        // register secondary connection
+        let (cmd_tx2, _cmd_rx2) = channel(64);
+        sender
+            .send(InnerTransportEvent::ConnectionEstablished {
+                peer,
+                connection: ConnectionId::from(1usize),
+                endpoint: Endpoint::listener(Multiaddr::empty(), ConnectionId::from(1usize)),
+                sender: ConnectionHandle::new(ConnectionId::from(1usize), cmd_tx2),
+            })
+            .await
+            .unwrap();
+
+        futures::future::poll_fn(|cx| match service.poll_next_unpin(cx) {
+            Poll::Ready(_) => panic!("didn't expect event from `TransportService`"),
+            Poll::Pending => Poll::Ready(()),
+        })
+        .await;
+
+        // all substreams are opened over the primary connection
+        let opened = service.open_substream(peer).unwrap();
+        let failed = service.open_substream(peer).unwrap();
+        let pending = service.open_substream(peer).unwrap();
+
+        for expected in [opened, failed, pending] {
+            match cmd_rx1.try_recv() {
+                Ok(ProtocolCommand::OpenSubstream {
+                    substream_id,
+                    connection_id,
+                    ..
+                }) => {
+                    assert_eq!(substream_id, expected);
+                    assert_eq!(connection_id, ConnectionId::from(0usize));
+                }
+                command => panic!("unexpected command: {command:?}"),
+            }
+        }
+
+        // the transport reports two of them back before the connection dies
+        sender
+            .send(InnerTransportEvent::SubstreamOpened {
+                peer,
+                protocol: ProtocolName::from("/notif/1"),
+                fallback: None,
+                direction: Direction::Outbound(opened),
+                connection_id: ConnectionId::from(0usize),
+                substream: Substream::new_mock(peer, opened, Box::new(MockSubstream::new())),
+                opening_permit: Permit::new(cmd_tx1),
+            })
+            .await
+            .unwrap();
+
+        assert!(std::matches!(
+            service.next().await,
+            Some(TransportEvent::SubstreamOpened { .. })
+        ));
+
+        sender
+            .send(InnerTransportEvent::SubstreamOpenFailure {
+                substream: failed,
+                error: SubstreamError::ChannelClogged,
+            })
+            .await
+            .unwrap();
+
+        assert!(std::matches!(
+            service.next().await,
+            Some(TransportEvent::SubstreamOpenFailure { .. })
+        ));
+
+        assert_eq!(
+            service.pending_outbound,
+            HashMap::from_iter([(pending, ConnectionId::from(0usize))])
+        );
+
+        // close the primary connection
+        sender
+            .send(InnerTransportEvent::ConnectionClosed {
+                peer,
+                connection: ConnectionId::from(0usize),
+            })
+            .await
+            .unwrap();
+
+        match service.next().await {
+            Some(TransportEvent::SubstreamOpenFailure { substream, error }) => {
+                assert_eq!(substream, pending);
+                assert!(std::matches!(error, SubstreamError::ConnectionClosed));
+            }
+            event => panic!("unexpected event: {event:?}"),
+        }
+
+        assert!(service.pending_outbound.is_empty());
+
+        // the protocol is not notified about the closed connection, the secondary took over
+        futures::future::poll_fn(|cx| match service.poll_next_unpin(cx) {
+            Poll::Ready(_) => panic!("didn't expect event from `TransportService`"),
+            Poll::Pending => Poll::Ready(()),
+        })
+        .await;
+
+        let context = service.connections.get(&peer).unwrap();
+        assert_eq!(context.primary.connection_id(), &ConnectionId::from(1usize));
+        assert!(context.secondary.is_none());
+    }
+
     #[tokio::test]
     async fn keep_alive_timeout_expires_for_a_stale_connection() {
         let _ = tracing_subscriber::fmt()