feat: move Identify I/O from NetworkBehaviour to ConnectionHandler

protocols/identify/src/behaviour.rs

@@ -18,9 +18,8 @@
 // FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 // DEALINGS IN THE SOFTWARE.
 
-use crate::handler::{self, Proto, Push};
+use crate::handler::{self, InEvent, Proto, Reply};
 use crate::protocol::{Info, ReplySubstream, UpgradeError};
-use futures::prelude::*;
 use libp2p_core::{
     connection::ConnectionId, multiaddr::Protocol, ConnectedPoint, Multiaddr, PeerId, PublicKey,
 };
@@ -35,7 +34,6 @@ use std::num::NonZeroUsize;
 use std::{
     collections::{HashMap, HashSet, VecDeque},
     iter::FromIterator,
-    pin::Pin,
     task::Context,
     task::Poll,
     time::Duration,
@@ -51,8 +49,8 @@ pub struct Behaviour {
     config: Config,
     /// For each peer we're connected to, the observed address to send back to it.
     connected: HashMap<PeerId, HashMap<ConnectionId, Multiaddr>>,
-    /// Pending replies to send.
-    pending_replies: VecDeque<Reply>,
+    /// Pending requests to respond.
+    requests: VecDeque<Request>,
     /// Pending events to be emitted when polled.
     events: VecDeque<NetworkBehaviourAction<Event, Proto>>,
     /// Peers to which an active push with current information about
@@ -63,18 +61,10 @@ pub struct Behaviour {
 }
 
 /// A pending reply to an inbound identification request.
-enum Reply {
-    /// The reply is queued for sending.
-    Queued {
-        peer: PeerId,
-        io: ReplySubstream<NegotiatedSubstream>,
-        observed: Multiaddr,
-    },
-    /// The reply is being sent.
-    Sending {
-        peer: PeerId,
-        io: Pin<Box<dyn Future<Output = Result<(), UpgradeError>> + Send>>,
-    },
+struct Request {
+    peer: PeerId,
+    io: ReplySubstream<NegotiatedSubstream>,
+    observed: Multiaddr,
 }
 
 /// Configuration for the [`identify::Behaviour`](Behaviour).
@@ -184,7 +174,7 @@ impl Behaviour {
         Self {
             config,
             connected: HashMap::new(),
-            pending_replies: VecDeque::new(),
+            requests: VecDeque::new(),
             events: VecDeque::new(),
             pending_push: HashSet::new(),
             discovered_peers,
@@ -271,6 +261,12 @@ impl NetworkBehaviour for Behaviour {
                         score: AddressScore::Finite(1),
                     });
             }
+            handler::Event::Identification(peer) => {
+                self.events
+                    .push_back(NetworkBehaviourAction::GenerateEvent(Event::Sent {
+                        peer_id: peer,
+                    }));
+            }
             handler::Event::IdentificationPushed => {
                 self.events
                     .push_back(NetworkBehaviourAction::GenerateEvent(Event::Pushed {
@@ -287,7 +283,7 @@ impl NetworkBehaviour for Behaviour {
                         with an established connection and calling `NetworkBehaviour::on_event` \
                         with `FromSwarm::ConnectionEstablished ensures there is an entry; qed",
                     );
-                self.pending_replies.push_back(Reply::Queued {
+                self.requests.push_back(Request {
                     peer: peer_id,
                     io: sender,
                     observed: observed.clone(),
@@ -305,7 +301,7 @@ impl NetworkBehaviour for Behaviour {
 
     fn poll(
         &mut self,
-        cx: &mut Context<'_>,
+        _cx: &mut Context<'_>,
         params: &mut impl PollParameters,
     ) -> Poll<NetworkBehaviourAction<Self::OutEvent, Self::ConnectionHandler>> {
         if let Some(event) = self.events.pop_front() {
@@ -333,7 +329,7 @@ impl NetworkBehaviour for Behaviour {
                     observed_addr,
                 };
 
-                (*peer, Push(info))
+                (*peer, InEvent::Push(info))
             })
         });
 
@@ -346,55 +342,21 @@ impl NetworkBehaviour for Behaviour {
             });
         }
 
-        // Check for pending replies to send.
-        if let Some(r) = self.pending_replies.pop_front() {
-            let mut sending = 0;
-            let to_send = self.pending_replies.len() + 1;
-            let mut reply = Some(r);
-            loop {
-                match reply {
-                    Some(Reply::Queued { peer, io, observed }) => {
-                        let info = Info {
-                            listen_addrs: listen_addrs(params),
-                            protocols: supported_protocols(params),
-                            public_key: self.config.local_public_key.clone(),
-                            protocol_version: self.config.protocol_version.clone(),
-                            agent_version: self.config.agent_version.clone(),
-                            observed_addr: observed,
-                        };
-                        let io = Box::pin(io.send(info));
-                        reply = Some(Reply::Sending { peer, io });
-                    }
-                    Some(Reply::Sending { peer, mut io }) => {
-                        sending += 1;
-                        match Future::poll(Pin::new(&mut io), cx) {
-                            Poll::Ready(Ok(())) => {
-                                let event = Event::Sent { peer_id: peer };
-                                return Poll::Ready(NetworkBehaviourAction::GenerateEvent(event));
-                            }
-                            Poll::Pending => {
-                                self.pending_replies.push_back(Reply::Sending { peer, io });
-                                if sending == to_send {
-                                    // All remaining futures are NotReady
-                                    break;
-                                } else {
-                                    reply = self.pending_replies.pop_front();
-                                }
-                            }
-                            Poll::Ready(Err(err)) => {
-                                let event = Event::Error {
-                                    peer_id: peer,
-                                    error: ConnectionHandlerUpgrErr::Upgrade(
-                                        libp2p_core::upgrade::UpgradeError::Apply(err),
-                                    ),
-                                };
-                                return Poll::Ready(NetworkBehaviourAction::GenerateEvent(event));
-                            }
-                        }
-                    }
-                    None => unreachable!(),
-                }
-            }
+        // Check for pending requests to send back to the handler for reply.
+        if let Some(Request { peer, io, observed }) = self.requests.pop_front() {
+            let info = Info {
+                listen_addrs: listen_addrs(params),
+                protocols: supported_protocols(params),
+                public_key: self.config.local_public_key.clone(),
+                protocol_version: self.config.protocol_version.clone(),
+                agent_version: self.config.agent_version.clone(),
+                observed_addr: observed,
+            };
+            return Poll::Ready(NetworkBehaviourAction::NotifyHandler {
+                peer_id: peer,
+                handler: NotifyHandler::Any,
+                event: InEvent::Identify(Reply { peer, info, io }),
+            });
         }
 
         Poll::Pending
@@ -557,6 +519,7 @@ impl PeerCache {
 mod tests {
     use super::*;
     use futures::pin_mut;
+    use futures::prelude::*;
     use libp2p::mplex::MplexConfig;
     use libp2p::noise;
     use libp2p::tcp;
@@ -618,7 +581,7 @@ mod tests {
 
         // nb. Either swarm may receive the `Identified` event first, upon which
         // it will permit the connection to be closed, as defined by
-        // `IdentifyHandler::connection_keep_alive`. Hence the test succeeds if
+        // `Handler::connection_keep_alive`. Hence the test succeeds if
         // either `Identified` event arrives correctly.
         async_std::task::block_on(async move {
             loop {

protocols/identify/src/handler.rs

@@ -36,6 +36,7 @@ use libp2p_swarm::{
 };
 use log::warn;
 use smallvec::SmallVec;
+use std::collections::VecDeque;
 use std::{io, pin::Pin, task::Context, task::Poll, time::Duration};
 
 pub struct Proto {
@@ -64,6 +65,25 @@ impl IntoConnectionHandler for Proto {
     }
 }
 
+/// A pending reply to an inbound identification request.
+enum Pending {
+    /// The reply is queued for sending.
+    Queued(Reply),
+    /// The reply is being sent.
+    Sending {
+        peer: PeerId,
+        io: Pin<Box<dyn Future<Output = Result<(), UpgradeError>> + Send>>,
+    },
+}
+
+/// A reply to an inbound identification request.
+#[derive(Debug)]
+pub struct Reply {
+    pub peer: PeerId,
+    pub io: ReplySubstream<NegotiatedSubstream>,
+    pub info: Info,
+}
+
 /// Protocol handler for sending and receiving identification requests.
 ///
 /// Outbound requests are sent periodically. The handler performs expects
@@ -82,6 +102,9 @@ pub struct Handler {
         >; 4],
     >,
 
+    /// Pending replies to send.
+    pending_replies: VecDeque<Pending>,
+
     /// Future that fires when we need to identify the node again.
     trigger_next_identify: Delay,
 
@@ -92,31 +115,39 @@ pub struct Handler {
     interval: Duration,
 }
 
-/// Event produced by the `IdentifyHandler`.
+/// Event produced by the `Handler`.
 #[derive(Debug)]
 #[allow(clippy::large_enum_variant)]
 pub enum Event {
     /// We obtained identification information from the remote.
     Identified(Info),
+    /// We replied to an identification request from the remote.
+    Identification(PeerId),
     /// We actively pushed our identification information to the remote.
     IdentificationPushed,
     /// We received a request for identification.
     Identify(ReplySubstream<NegotiatedSubstream>),
-    /// Failed to identify the remote.
+    /// Failed to identify the remote, or to reply to an identification request.
     IdentificationError(ConnectionHandlerUpgrErr<UpgradeError>),
 }
 
-/// Identifying information of the local node that is pushed to a remote.
 #[derive(Debug)]
-pub struct Push(pub Info);
+#[allow(clippy::large_enum_variant)]
+pub enum InEvent {
+    /// Identifying information of the local node that is pushed to a remote.
+    Push(Info),
+    /// Identifying information requested from this node.
+    Identify(Reply),
+}
 
 impl Handler {
-    /// Creates a new `IdentifyHandler`.
+    /// Creates a new `Handler`.
     pub fn new(initial_delay: Duration, interval: Duration, remote_peer_id: PeerId) -> Self {
         Self {
             remote_peer_id,
             inbound_identify_push: Default::default(),
             events: SmallVec::new(),
+            pending_replies: VecDeque::new(),
             trigger_next_identify: Delay::new(initial_delay),
             keep_alive: KeepAlive::Yes,
             interval,
@@ -195,7 +226,7 @@ impl Handler {
 }
 
 impl ConnectionHandler for Handler {
-    type InEvent = Push;
+    type InEvent = InEvent;
     type OutEvent = Event;
     type Error = io::Error;
     type InboundProtocol = SelectUpgrade<Protocol, PushProtocol<InboundPush>>;
@@ -207,14 +238,28 @@ impl ConnectionHandler for Handler {
         SubstreamProtocol::new(SelectUpgrade::new(Protocol, PushProtocol::inbound()), ())
     }
 
-    fn on_behaviour_event(&mut self, Push(push): Self::InEvent) {
-        self.events
-            .push(ConnectionHandlerEvent::OutboundSubstreamRequest {
-                protocol: SubstreamProtocol::new(
-                    EitherUpgrade::B(PushProtocol::outbound(push)),
-                    (),
-                ),
-            });
+    fn on_behaviour_event(&mut self, event: Self::InEvent) {
+        match event {
+            InEvent::Push(push) => {
+                self.events
+                    .push(ConnectionHandlerEvent::OutboundSubstreamRequest {
+                        protocol: SubstreamProtocol::new(
+                            EitherUpgrade::B(PushProtocol::outbound(push)),
+                            (),
+                        ),
+                    });
+            }
+            InEvent::Identify(reply) => {
+                if !self.pending_replies.is_empty() {
+                    warn!(
+                        "New inbound identify request from {} while a previous one \
+                         is still pending. Queueing the new one.",
+                        reply.peer,
+                    );
+                }
+                self.pending_replies.push_back(Pending::Queued(reply));
+            }
+        }
     }
 
     fn connection_keep_alive(&self) -> KeepAlive {
@@ -255,6 +300,39 @@ impl ConnectionHandler for Handler {
             }
         }
 
+        // Check for pending replies to send.
+        if let Some(mut pending) = self.pending_replies.pop_front() {
+            loop {
+                match pending {
+                    Pending::Queued(Reply { peer, io, info }) => {
+                        let io = Box::pin(io.send(info));
+                        pending = Pending::Sending { peer, io };
+                    }
+                    Pending::Sending { peer, mut io } => {
+                        match Future::poll(Pin::new(&mut io), cx) {
+                            Poll::Pending => {
+                                self.pending_replies
+                                    .push_front(Pending::Sending { peer, io });
+                                return Poll::Pending;
+                            }
+                            Poll::Ready(Ok(())) => {
+                                return Poll::Ready(ConnectionHandlerEvent::Custom(
+                                    Event::Identification(peer),
+                                ));
+                            }
+                            Poll::Ready(Err(err)) => {
+                                return Poll::Ready(ConnectionHandlerEvent::Custom(
+                                    Event::IdentificationError(ConnectionHandlerUpgrErr::Upgrade(
+                                        libp2p_core::upgrade::UpgradeError::Apply(err),
+                                    )),
+                                ))
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
         Poll::Pending
     }