refactor(identify): use ReadyUpgrade for {In,Out}boundUpgrade

Related: #2863.

Pull-Request: #4563.

Cargo.lock

@@ -2617,6 +2617,7 @@ dependencies = [
  "either",
  "env_logger 0.10.0",
  "futures",
+ "futures-bounded",
  "futures-timer",
  "libp2p-core",
  "libp2p-identity",

protocols/identify/Cargo.toml

@@ -14,6 +14,7 @@ categories = ["network-programming", "asynchronous"]
 asynchronous-codec = "0.6"
 futures = "0.3.28"
 futures-timer = "3.0.2"
+futures-bounded = { workspace = true }
 libp2p-core = { workspace = true }
 libp2p-swarm = { workspace = true }
 libp2p-identity = { workspace = true }

protocols/identify/src/handler.rs

@@ -18,14 +18,13 @@
 // FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 // DEALINGS IN THE SOFTWARE.
 
-use crate::protocol::{Identify, InboundPush, OutboundPush, Push, UpgradeError};
+use crate::protocol::{Info, PushInfo, UpgradeError};
-use crate::protocol::{Info, PushInfo};
+use crate::{protocol, PROTOCOL_NAME, PUSH_PROTOCOL_NAME};
 use either::Either;
-use futures::future::BoxFuture;
 use futures::prelude::*;
-use futures::stream::FuturesUnordered;
+use futures_bounded::Timeout;
 use futures_timer::Delay;
-use libp2p_core::upgrade::SelectUpgrade;
+use libp2p_core::upgrade::{ReadyUpgrade, SelectUpgrade};
 use libp2p_core::Multiaddr;
 use libp2p_identity::PeerId;
 use libp2p_identity::PublicKey;
@@ -42,6 +41,9 @@ use smallvec::SmallVec;
 use std::collections::HashSet;
 use std::{io, task::Context, task::Poll, time::Duration};
 
+const STREAM_TIMEOUT: Duration = Duration::from_secs(60);
+const MAX_CONCURRENT_STREAMS_PER_CONNECTION: usize = 10;
+
 /// Protocol handler for sending and receiving identification requests.
 ///
 /// Outbound requests are sent periodically. The handler performs expects
@@ -49,14 +51,17 @@ use std::{io, task::Context, task::Poll, time::Duration};
 /// permitting the underlying connection to be closed.
 pub struct Handler {
     remote_peer_id: PeerId,
-    inbound_identify_push: Option<BoxFuture<'static, Result<PushInfo, UpgradeError>>>,
     /// Pending events to yield.
     events: SmallVec<
-        [ConnectionHandlerEvent<Either<Identify, Push<OutboundPush>>, (), Event, io::Error>; 4],
+        [ConnectionHandlerEvent<
+            Either<ReadyUpgrade<StreamProtocol>, ReadyUpgrade<StreamProtocol>>,
+            (),
+            Event,
+            io::Error,
+        >; 4],
     >,
 
-    /// Pending identification replies, awaiting being sent.
+    active_streams: futures_bounded::FuturesSet<Result<Success, UpgradeError>>,
-    pending_replies: FuturesUnordered<BoxFuture<'static, Result<(), UpgradeError>>>,
 
     /// Future that fires when we need to identify the node again.
     trigger_next_identify: Delay,
@@ -125,9 +130,11 @@ impl Handler {
     ) -> Self {
         Self {
             remote_peer_id,
-            inbound_identify_push: Default::default(),
             events: SmallVec::new(),
-            pending_replies: FuturesUnordered::new(),
+            active_streams: futures_bounded::FuturesSet::new(
+                STREAM_TIMEOUT,
+                MAX_CONCURRENT_STREAMS_PER_CONNECTION,
+            ),
             trigger_next_identify: Delay::new(initial_delay),
             exchanged_one_periodic_identify: false,
             interval,
@@ -152,19 +159,28 @@ impl Handler {
         >,
     ) {
         match output {
-            future::Either::Left(substream) => {
+            future::Either::Left(stream) => {
                 let info = self.build_info();
 
-                self.pending_replies
+                if self
-                    .push(crate::protocol::send(substream, info).boxed());
+                    .active_streams
+                    .try_push(
+                        protocol::send_identify(stream, info).map_ok(|_| Success::SentIdentify),
+                    )
+                    .is_err()
+                {
+                    warn!("Dropping inbound stream because we are at capacity");
+                } else {
+                    self.exchanged_one_periodic_identify = true;
+                }
             }
-            future::Either::Right(fut) => {
+            future::Either::Right(stream) => {
-                if self.inbound_identify_push.replace(fut).is_some() {
+                if self
-                    warn!(
+                    .active_streams
-                        "New inbound identify push stream from {} while still \
+                    .try_push(protocol::recv_push(stream).map_ok(Success::ReceivedIdentifyPush))
-                         upgrading previous one. Replacing previous with new.",
+                    .is_err()
-                        self.remote_peer_id,
+                {
-                    );
+                    warn!("Dropping inbound identify push stream because we are at capacity");
                 }
             }
         }
@@ -180,32 +196,29 @@ impl Handler {
         >,
     ) {
         match output {
-            future::Either::Left(remote_info) => {
+            future::Either::Left(stream) => {
-                self.handle_incoming_info(&remote_info);
+                if self
-
+                    .active_streams
-                self.events
+                    .try_push(protocol::recv_identify(stream).map_ok(Success::ReceivedIdentify))
-                    .push(ConnectionHandlerEvent::NotifyBehaviour(Event::Identified(
+                    .is_err()
-                        remote_info,
+                {
-                    )));
+                    warn!("Dropping outbound identify stream because we are at capacity");
+                }
             }
-            future::Either::Right(()) => self.events.push(ConnectionHandlerEvent::NotifyBehaviour(
+            future::Either::Right(stream) => {
-                Event::IdentificationPushed,
+                let info = self.build_info();
-            )),
-        }
-    }
 
-    fn on_dial_upgrade_error(
+                if self
-        &mut self,
+                    .active_streams
-        DialUpgradeError { error: err, .. }: DialUpgradeError<
+                    .try_push(
-            <Self as ConnectionHandler>::OutboundOpenInfo,
+                        protocol::send_identify(stream, info).map_ok(|_| Success::SentIdentifyPush),
-            <Self as ConnectionHandler>::OutboundProtocol,
+                    )
-        >,
+                    .is_err()
-    ) {
+                {
-        let err = err.map_upgrade_err(|e| e.into_inner());
+                    warn!("Dropping outbound identify push stream because we are at capacity");
-        self.events.push(ConnectionHandlerEvent::NotifyBehaviour(
+                }
-            Event::IdentificationError(err),
+            }
-        ));
+        }
-        self.trigger_next_identify.reset(self.interval);
     }
 
     fn build_info(&mut self) -> Info {
@@ -268,13 +281,20 @@ impl ConnectionHandler for Handler {
     type FromBehaviour = InEvent;
     type ToBehaviour = Event;
     type Error = io::Error;
-    type InboundProtocol = SelectUpgrade<Identify, Push<InboundPush>>;
+    type InboundProtocol =
-    type OutboundProtocol = Either<Identify, Push<OutboundPush>>;
+        SelectUpgrade<ReadyUpgrade<StreamProtocol>, ReadyUpgrade<StreamProtocol>>;
+    type OutboundProtocol = Either<ReadyUpgrade<StreamProtocol>, ReadyUpgrade<StreamProtocol>>;
     type OutboundOpenInfo = ();
     type InboundOpenInfo = ();
 
     fn listen_protocol(&self) -> SubstreamProtocol<Self::InboundProtocol, Self::InboundOpenInfo> {
-        SubstreamProtocol::new(SelectUpgrade::new(Identify, Push::inbound()), ())
+        SubstreamProtocol::new(
+            SelectUpgrade::new(
+                ReadyUpgrade::new(PROTOCOL_NAME),
+                ReadyUpgrade::new(PUSH_PROTOCOL_NAME),
+            ),
+            (),
+        )
     }
 
     fn on_behaviour_event(&mut self, event: Self::FromBehaviour) {
@@ -283,21 +303,19 @@ impl ConnectionHandler for Handler {
                 self.external_addresses = addresses;
             }
             InEvent::Push => {
-                let info = self.build_info();
                 self.events
                     .push(ConnectionHandlerEvent::OutboundSubstreamRequest {
-                        protocol: SubstreamProtocol::new(Either::Right(Push::outbound(info)), ()),
+                        protocol: SubstreamProtocol::new(
+                            Either::Right(ReadyUpgrade::new(PUSH_PROTOCOL_NAME)),
+                            (),
+                        ),
                     });
             }
         }
     }
 
     fn connection_keep_alive(&self) -> KeepAlive {
-        if self.inbound_identify_push.is_some() {
+        if !self.active_streams.is_empty() {
-            return KeepAlive::Yes;
-        }
-
-        if !self.pending_replies.is_empty() {
             return KeepAlive::Yes;
         }
 
@@ -317,20 +335,34 @@ impl ConnectionHandler for Handler {
         // Poll the future that fires when we need to identify the node again.
         if let Poll::Ready(()) = self.trigger_next_identify.poll_unpin(cx) {
             self.trigger_next_identify.reset(self.interval);
-            let ev = ConnectionHandlerEvent::OutboundSubstreamRequest {
+            let event = ConnectionHandlerEvent::OutboundSubstreamRequest {
-                protocol: SubstreamProtocol::new(Either::Left(Identify), ()),
+                protocol: SubstreamProtocol::new(
+                    Either::Left(ReadyUpgrade::new(PROTOCOL_NAME)),
+                    (),
+                ),
             };
-            return Poll::Ready(ev);
+            return Poll::Ready(event);
         }
 
-        if let Some(Poll::Ready(res)) = self
+        match self.active_streams.poll_unpin(cx) {
-            .inbound_identify_push
+            Poll::Ready(Ok(Ok(Success::ReceivedIdentify(remote_info)))) => {
-            .as_mut()
+                self.handle_incoming_info(&remote_info);
-            .map(|f| f.poll_unpin(cx))
-        {
-            self.inbound_identify_push.take();
 
-            if let Ok(remote_push_info) = res {
+                return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(Event::Identified(
+                    remote_info,
+                )));
+            }
+            Poll::Ready(Ok(Ok(Success::SentIdentifyPush))) => {
+                return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(
+                    Event::IdentificationPushed,
+                ));
+            }
+            Poll::Ready(Ok(Ok(Success::SentIdentify))) => {
+                return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(
+                    Event::Identification,
+                ));
+            }
+            Poll::Ready(Ok(Ok(Success::ReceivedIdentifyPush(remote_push_info)))) => {
                 if let Some(mut info) = self.remote_info.clone() {
                     info.merge(remote_push_info);
                     self.handle_incoming_info(&info);
@@ -340,16 +372,17 @@ impl ConnectionHandler for Handler {
                     ));
                 };
             }
-        }
+            Poll::Ready(Ok(Err(e))) => {
-
+                return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(
-        // Check for pending replies to send.
+                    Event::IdentificationError(StreamUpgradeError::Apply(e)),
-        if let Poll::Ready(Some(result)) = self.pending_replies.poll_next_unpin(cx) {
+                ));
-            let event = result
+            }
-                .map(|()| Event::Identification)
+            Poll::Ready(Err(Timeout { .. })) => {
-                .unwrap_or_else(|err| Event::IdentificationError(StreamUpgradeError::Apply(err)));
+                return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(
-            self.exchanged_one_periodic_identify = true;
+                    Event::IdentificationError(StreamUpgradeError::Timeout),
-
+                ));
-            return Poll::Ready(ConnectionHandlerEvent::NotifyBehaviour(event));
+            }
+            Poll::Pending => {}
         }
 
         Poll::Pending
@@ -371,8 +404,13 @@ impl ConnectionHandler for Handler {
             ConnectionEvent::FullyNegotiatedOutbound(fully_negotiated_outbound) => {
                 self.on_fully_negotiated_outbound(fully_negotiated_outbound)
             }
-            ConnectionEvent::DialUpgradeError(dial_upgrade_error) => {
+            ConnectionEvent::DialUpgradeError(DialUpgradeError { error, .. }) => {
-                self.on_dial_upgrade_error(dial_upgrade_error)
+                self.events.push(ConnectionHandlerEvent::NotifyBehaviour(
+                    Event::IdentificationError(
+                        error.map_upgrade_err(|e| void::unreachable(e.into_inner())),
+                    ),
+                ));
+                self.trigger_next_identify.reset(self.interval);
             }
             ConnectionEvent::AddressChange(_)
             | ConnectionEvent::ListenUpgradeError(_)
@@ -392,11 +430,10 @@ impl ConnectionHandler for Handler {
                         self.remote_peer_id
                     );
 
-                    let info = self.build_info();
                     self.events
                         .push(ConnectionHandlerEvent::OutboundSubstreamRequest {
                             protocol: SubstreamProtocol::new(
-                                Either::Right(Push::outbound(info)),
+                                Either::Right(ReadyUpgrade::new(PUSH_PROTOCOL_NAME)),
                                 (),
                             ),
                         });
@@ -405,3 +442,10 @@ impl ConnectionHandler for Handler {
         }
     }
 }
+
+enum Success {
+    SentIdentify,
+    ReceivedIdentify(Info),
+    SentIdentifyPush,
+    ReceivedIdentifyPush(PushInfo),
+}

protocols/identify/src/protocol.rs

@@ -20,20 +20,15 @@
 
 use crate::proto;
 use asynchronous_codec::{FramedRead, FramedWrite};
-use futures::{future::BoxFuture, prelude::*};
+use futures::prelude::*;
-use libp2p_core::{
+use libp2p_core::{multiaddr, Multiaddr};
-    multiaddr,
-    upgrade::{InboundUpgrade, OutboundUpgrade, UpgradeInfo},
-    Multiaddr,
-};
 use libp2p_identity as identity;
 use libp2p_identity::PublicKey;
 use libp2p_swarm::StreamProtocol;
 use log::{debug, trace};
 use std::convert::TryFrom;
-use std::{io, iter, pin::Pin};
+use std::io;
 use thiserror::Error;
-use void::Void;
 
 const MAX_MESSAGE_SIZE_BYTES: usize = 4096;
 
@@ -41,28 +36,6 @@ pub const PROTOCOL_NAME: StreamProtocol = StreamProtocol::new("/ipfs/id/1.0.0");
 
 pub const PUSH_PROTOCOL_NAME: StreamProtocol = StreamProtocol::new("/ipfs/id/push/1.0.0");
 
-/// Substream upgrade protocol for `/ipfs/id/1.0.0`.
-#[derive(Debug, Clone)]
-pub struct Identify;
-
-/// Substream upgrade protocol for `/ipfs/id/push/1.0.0`.
-#[derive(Debug, Clone)]
-pub struct Push<T>(T);
-pub struct InboundPush();
-pub struct OutboundPush(Info);
-
-impl Push<InboundPush> {
-    pub fn inbound() -> Self {
-        Push(InboundPush())
-    }
-}
-
-impl Push<OutboundPush> {
-    pub fn outbound(info: Info) -> Self {
-        Push(OutboundPush(info))
-    }
-}
-
 /// Identify information of a peer sent in protocol messages.
 #[derive(Debug, Clone)]
 pub struct Info {
@@ -117,75 +90,7 @@ pub struct PushInfo {
     pub observed_addr: Option<Multiaddr>,
 }
 
-impl UpgradeInfo for Identify {
+pub(crate) async fn send_identify<T>(io: T, info: Info) -> Result<(), UpgradeError>
-    type Info = StreamProtocol;
-    type InfoIter = iter::Once<Self::Info>;
-
-    fn protocol_info(&self) -> Self::InfoIter {
-        iter::once(PROTOCOL_NAME)
-    }
-}
-
-impl<C> InboundUpgrade<C> for Identify {
-    type Output = C;
-    type Error = UpgradeError;
-    type Future = future::Ready<Result<Self::Output, UpgradeError>>;
-
-    fn upgrade_inbound(self, socket: C, _: Self::Info) -> Self::Future {
-        future::ok(socket)
-    }
-}
-
-impl<C> OutboundUpgrade<C> for Identify
-where
-    C: AsyncRead + AsyncWrite + Unpin + Send + 'static,
-{
-    type Output = Info;
-    type Error = UpgradeError;
-    type Future = Pin<Box<dyn Future<Output = Result<Self::Output, Self::Error>> + Send>>;
-
-    fn upgrade_outbound(self, socket: C, _: Self::Info) -> Self::Future {
-        recv_identify(socket).boxed()
-    }
-}
-
-impl<T> UpgradeInfo for Push<T> {
-    type Info = StreamProtocol;
-    type InfoIter = iter::Once<Self::Info>;
-
-    fn protocol_info(&self) -> Self::InfoIter {
-        iter::once(PUSH_PROTOCOL_NAME)
-    }
-}
-
-impl<C> InboundUpgrade<C> for Push<InboundPush>
-where
-    C: AsyncRead + AsyncWrite + Unpin + Send + 'static,
-{
-    type Output = BoxFuture<'static, Result<PushInfo, UpgradeError>>;
-    type Error = Void;
-    type Future = future::Ready<Result<Self::Output, Self::Error>>;
-
-    fn upgrade_inbound(self, socket: C, _: Self::Info) -> Self::Future {
-        // Lazily upgrade stream, thus allowing upgrade to happen within identify's handler.
-        future::ok(recv_push(socket).boxed())
-    }
-}
-
-impl<C> OutboundUpgrade<C> for Push<OutboundPush>
-where
-    C: AsyncWrite + Unpin + Send + 'static,
-{
-    type Output = ();
-    type Error = UpgradeError;
-    type Future = Pin<Box<dyn Future<Output = Result<Self::Output, Self::Error>> + Send>>;
-
-    fn upgrade_outbound(self, socket: C, _: Self::Info) -> Self::Future {
-        send(socket, self.0 .0).boxed()
-    }
-}
-
-pub(crate) async fn send<T>(io: T, info: Info) -> Result<(), UpgradeError>
 where
     T: AsyncWrite + Unpin,
 {
@@ -219,7 +124,7 @@ where
     Ok(())
 }
 
-async fn recv_push<T>(socket: T) -> Result<PushInfo, UpgradeError>
+pub(crate) async fn recv_push<T>(socket: T) -> Result<PushInfo, UpgradeError>
 where
     T: AsyncRead + AsyncWrite + Unpin,
 {
@@ -230,7 +135,7 @@ where
     Ok(info)
 }
 
-async fn recv_identify<T>(socket: T) -> Result<Info, UpgradeError>
+pub(crate) async fn recv_identify<T>(socket: T) -> Result<Info, UpgradeError>
 where
     T: AsyncRead + AsyncWrite + Unpin,
 {