*: Dial with handler and return handler on error and closed (#2191)

Require `NetworkBehaviourAction::{DialPeer,DialAddress}` to contain a
`ProtocolsHandler`. This allows a behaviour to attach custom state to its
handler. The behaviour would no longer need to track this state separately
during connection establishment, thus reducing state required in a behaviour.
E.g. in the case of `libp2p-kad` the behaviour can include a `GetRecord` request
in its handler, or e.g. in the case of `libp2p-request-response` the behaviour
can include the first request in the handler.

Return `ProtocolsHandler` on connection error and close. This allows a behaviour
to extract its custom state previously included in the handler on connection
failure and connection closing. E.g. in the case of `libp2p-kad` the behaviour
could extract the attached `GetRecord` from the handler of the failed connection
and then start another connection attempt with a new handler with the same
`GetRecord` or bubble up an error to the user.

Co-authored-by: Thomas Eizinger <thomas@eizinger.io>

protocols/request-response/src/lib.rs

@@ -68,7 +68,8 @@ use futures::channel::oneshot;
 use handler::{RequestProtocol, RequestResponseHandler, RequestResponseHandlerEvent};
 use libp2p_core::{connection::ConnectionId, ConnectedPoint, Multiaddr, PeerId};
 use libp2p_swarm::{
-    DialPeerCondition, NetworkBehaviour, NetworkBehaviourAction, NotifyHandler, PollParameters,
+    DialError, DialPeerCondition, IntoProtocolsHandler, NetworkBehaviour, NetworkBehaviourAction,
+    NotifyHandler, PollParameters,
 };
 use smallvec::SmallVec;
 use std::{
@@ -303,7 +304,7 @@ impl RequestResponseConfig {
 /// A request/response protocol for some message codec.
 pub struct RequestResponse<TCodec>
 where
-    TCodec: RequestResponseCodec,
+    TCodec: RequestResponseCodec + Clone + Send + 'static,
 {
     /// The supported inbound protocols.
     inbound_protocols: SmallVec<[TCodec::Protocol; 2]>,
@@ -320,8 +321,8 @@ where
     /// Pending events to return from `poll`.
     pending_events: VecDeque<
         NetworkBehaviourAction<
-            RequestProtocol<TCodec>,
             RequestResponseEvent<TCodec::Request, TCodec::Response>,
+            RequestResponseHandler<TCodec>,
         >,
     >,
     /// The currently connected peers, their pending outbound and inbound responses and their known,
@@ -336,7 +337,7 @@ where
 
 impl<TCodec> RequestResponse<TCodec>
 where
-    TCodec: RequestResponseCodec + Clone,
+    TCodec: RequestResponseCodec + Clone + Send + 'static,
 {
     /// Creates a new `RequestResponse` behaviour for the given
     /// protocols, codec and configuration.
@@ -403,10 +404,12 @@ where
         };
 
         if let Some(request) = self.try_send_request(peer, request) {
+            let handler = self.new_handler();
             self.pending_events
                 .push_back(NetworkBehaviourAction::DialPeer {
                     peer_id: *peer,
                     condition: DialPeerCondition::Disconnected,
+                    handler,
                 });
             self.pending_outbound_requests
                 .entry(*peer)
@@ -639,6 +642,7 @@ where
         peer_id: &PeerId,
         conn: &ConnectionId,
         _: &ConnectedPoint,
+        _: <Self::ProtocolsHandler as IntoProtocolsHandler>::Handler,
     ) {
         let connections = self
             .connected
@@ -682,7 +686,7 @@ where
         self.connected.remove(peer);
     }
 
-    fn inject_dial_failure(&mut self, peer: &PeerId) {
+    fn inject_dial_failure(&mut self, peer: &PeerId, _: Self::ProtocolsHandler, _: DialError) {
         // If there are pending outgoing requests when a dial failure occurs,
         // it is implied that we are not connected to the peer, since pending
         // outgoing requests are drained when a connection is established and
@@ -863,12 +867,7 @@ where
         &mut self,
         _: &mut Context<'_>,
         _: &mut impl PollParameters,
-    ) -> Poll<
-        NetworkBehaviourAction<
-            RequestProtocol<TCodec>,
-            RequestResponseEvent<TCodec::Request, TCodec::Response>,
-        >,
-    > {
+    ) -> Poll<NetworkBehaviourAction<Self::OutEvent, Self::ProtocolsHandler>> {
         if let Some(ev) = self.pending_events.pop_front() {
             return Poll::Ready(ev);
         } else if self.pending_events.capacity() > EMPTY_QUEUE_SHRINK_THRESHOLD {

protocols/request-response/src/throttled.rs

@@ -40,11 +40,13 @@ use super::{
     ProtocolSupport, RequestId, RequestResponse, RequestResponseCodec, RequestResponseConfig,
     RequestResponseEvent, RequestResponseMessage,
 };
-use crate::handler::{RequestProtocol, RequestResponseHandler, RequestResponseHandlerEvent};
+use crate::handler::{RequestResponseHandler, RequestResponseHandlerEvent};
 use codec::{Codec, Message, ProtocolWrapper, Type};
 use futures::ready;
 use libp2p_core::{connection::ConnectionId, ConnectedPoint, Multiaddr, PeerId};
-use libp2p_swarm::{NetworkBehaviour, NetworkBehaviourAction, PollParameters};
+use libp2p_swarm::{
+    DialError, IntoProtocolsHandler, NetworkBehaviour, NetworkBehaviourAction, PollParameters,
+};
 use lru::LruCache;
 use std::{cmp::max, num::NonZeroU16};
 use std::{
@@ -57,7 +59,7 @@ pub type ResponseChannel<R> = super::ResponseChannel<Message<R>>;
 /// A wrapper around [`RequestResponse`] which adds request limits per peer.
 pub struct Throttled<C>
 where
-    C: RequestResponseCodec + Send,
+    C: RequestResponseCodec + Clone + Send + 'static,
     C::Protocol: Sync,
 {
     /// A random id used for logging.
@@ -439,8 +441,15 @@ where
         self.behaviour.inject_connection_established(p, id, end)
     }
 
-    fn inject_connection_closed(&mut self, peer: &PeerId, id: &ConnectionId, end: &ConnectedPoint) {
-        self.behaviour.inject_connection_closed(peer, id, end);
+    fn inject_connection_closed(
+        &mut self,
+        peer: &PeerId,
+        id: &ConnectionId,
+        end: &ConnectedPoint,
+        handler: <Self::ProtocolsHandler as IntoProtocolsHandler>::Handler,
+    ) {
+        self.behaviour
+            .inject_connection_closed(peer, id, end, handler);
         if let Some(info) = self.peer_info.get_mut(peer) {
             if let Some(grant) = &mut info.recv_budget.grant {
                 log::debug! { "{:08x}: resending credit grant {} to {} after connection closed",
@@ -484,8 +493,13 @@ where
         self.behaviour.inject_disconnected(p)
     }
 
-    fn inject_dial_failure(&mut self, p: &PeerId) {
-        self.behaviour.inject_dial_failure(p)
+    fn inject_dial_failure(
+        &mut self,
+        p: &PeerId,
+        handler: Self::ProtocolsHandler,
+        error: DialError,
+    ) {
+        self.behaviour.inject_dial_failure(p, handler, error)
     }
 
     fn inject_event(
@@ -501,7 +515,7 @@ where
         &mut self,
         cx: &mut Context<'_>,
         params: &mut impl PollParameters,
-    ) -> Poll<NetworkBehaviourAction<RequestProtocol<Codec<C>>, Self::OutEvent>> {
+    ) -> Poll<NetworkBehaviourAction<Self::OutEvent, Self::ProtocolsHandler>> {
         loop {
             if let Some(ev) = self.events.pop_front() {
                 return Poll::Ready(NetworkBehaviourAction::GenerateEvent(ev));
@@ -737,12 +751,18 @@ where
                         RequestResponseEvent::ResponseSent { peer, request_id },
                     ))
                 }
-                NetworkBehaviourAction::DialAddress { address } => {
-                    NetworkBehaviourAction::DialAddress { address }
-                }
-                NetworkBehaviourAction::DialPeer { peer_id, condition } => {
-                    NetworkBehaviourAction::DialPeer { peer_id, condition }
+                NetworkBehaviourAction::DialAddress { address, handler } => {
+                    NetworkBehaviourAction::DialAddress { address, handler }
                 }
+                NetworkBehaviourAction::DialPeer {
+                    peer_id,
+                    condition,
+                    handler,
+                } => NetworkBehaviourAction::DialPeer {
+                    peer_id,
+                    condition,
+                    handler,
+                },
                 NetworkBehaviourAction::NotifyHandler {
                     peer_id,
                     handler,