*: 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>

swarm-derive/src/lib.rs

@@ -57,6 +57,7 @@ fn build_struct(ast: &DeriveInput, data_struct: &DataStruct) -> TokenStream {
     let connection_id = quote! {::libp2p::core::connection::ConnectionId};
     let connected_point = quote! {::libp2p::core::ConnectedPoint};
     let listener_id = quote! {::libp2p::core::connection::ListenerId};
+    let dial_error = quote! {::libp2p::swarm::DialError};
 
     let poll_parameters = quote! {::libp2p::swarm::PollParameters};
 
@@ -223,15 +224,33 @@ fn build_struct(ast: &DeriveInput, data_struct: &DataStruct) -> TokenStream {
 
     // Build the list of statements to put in the body of `inject_connection_closed()`.
     let inject_connection_closed_stmts = {
-        data_struct.fields.iter().enumerate().filter_map(move |(field_n, field)| {
-            if is_ignored(&field) {
-                return None;
-            }
-            Some(match field.ident {
-                Some(ref i) => quote!{ self.#i.inject_connection_closed(peer_id, connection_id, endpoint); },
-                None => quote!{ self.#field_n.inject_connection_closed(peer_id, connection_id, endpoint); },
+        data_struct
+            .fields
+            .iter()
+            .enumerate()
+            // The outmost handler belongs to the last behaviour.
+            .rev()
+            .filter(|f| !is_ignored(&f.1))
+            .enumerate()
+            .map(move |(enum_n, (field_n, field))| {
+                let handler = if field_n == 0 {
+                    // Given that the iterator is reversed, this is the innermost handler only.
+                    quote! { let handler = handlers }
+                } else {
+                    quote! {
+                        let (handlers, handler) = handlers.into_inner()
+                    }
+                };
+                let inject = match field.ident {
+                    Some(ref i) => quote!{ self.#i.inject_connection_closed(peer_id, connection_id, endpoint, handler) },
+                    None => quote!{ self.#enum_n.inject_connection_closed(peer_id, connection_id, endpoint, handler) },
+                };
+
+                quote! {
+                    #handler;
+                    #inject;
+                }
             })
-        })
     };
 
     // Build the list of statements to put in the body of `inject_addr_reach_failure()`.
@@ -255,15 +274,63 @@ fn build_struct(ast: &DeriveInput, data_struct: &DataStruct) -> TokenStream {
             .fields
             .iter()
             .enumerate()
-            .filter_map(move |(field_n, field)| {
-                if is_ignored(&field) {
-                    return None;
-                }
+            // The outmost handler belongs to the last behaviour.
+            .rev()
+            .filter(|f| !is_ignored(&f.1))
+            .enumerate()
+            .map(move |(enum_n, (field_n, field))| {
+                let handler = if field_n == 0 {
+                    // Given that the iterator is reversed, this is the innermost handler only.
+                    quote! { let handler = handlers }
+                } else {
+                    quote! {
+                        let (handlers, handler) = handlers.into_inner()
+                    }
+                };
 
-                Some(match field.ident {
-                    Some(ref i) => quote! { self.#i.inject_dial_failure(peer_id); },
-                    None => quote! { self.#field_n.inject_dial_failure(peer_id); },
-                })
+                let inject = match field.ident {
+                    Some(ref i) => {
+                        quote! { self.#i.inject_dial_failure(peer_id, handler, error.clone()) }
+                    }
+                    None => {
+                        quote! { self.#enum_n.inject_dial_failure(peer_id, handler, error.clone()) }
+                    }
+                };
+
+                quote! {
+                    #handler;
+                    #inject;
+                }
+            })
+    };
+
+    // Build the list of statements to put in the body of `inject_listen_failure()`.
+    let inject_listen_failure_stmts = {
+        data_struct
+            .fields
+            .iter()
+            .enumerate()
+            .rev()
+            .filter(|f| !is_ignored(&f.1))
+            .enumerate()
+            .map(move |(enum_n, (field_n, field))| {
+                let handler = if field_n == 0 {
+                    quote! { let handler = handlers }
+                } else {
+                    quote! {
+                        let (handlers, handler) = handlers.into_inner()
+                    }
+                };
+
+                let inject = match field.ident {
+                    Some(ref i) => quote! { self.#i.inject_listen_failure(local_addr, send_back_addr, handler) },
+                    None => quote! { self.#enum_n.inject_listen_failure(local_addr, send_back_addr, handler) },
+                };
+
+                quote! {
+                    #handler;
+                    #inject;
+                }
             })
     };
 
@@ -426,6 +493,7 @@ fn build_struct(ast: &DeriveInput, data_struct: &DataStruct) -> TokenStream {
                 ref mut ev @ None => *ev = Some(field_info),
             }
         }
+        // ph_ty = Some(quote! )
         ph_ty.unwrap_or(quote! {()}) // TODO: `!` instead
     };
 
@@ -456,7 +524,7 @@ fn build_struct(ast: &DeriveInput, data_struct: &DataStruct) -> TokenStream {
             }
         }
 
-        out_handler.unwrap_or(quote! {()}) // TODO: incorrect
+        out_handler.unwrap_or(quote! {()}) // TODO: See test `empty`.
     };
 
     // The method to use to poll.
@@ -500,6 +568,42 @@ fn build_struct(ast: &DeriveInput, data_struct: &DataStruct) -> TokenStream {
             wrapped_event = quote!{ #either_ident::First(#wrapped_event) };
         }
 
+        // `DialPeer` and `DialAddress` each provide a handler of the specific
+        // behaviour triggering the event. Though in order for the final handler
+        // to be able to handle protocols of all behaviours, the provided
+        // handler needs to be combined with handlers of all other behaviours.
+        let provided_handler_and_new_handlers = {
+            let mut out_handler = None;
+
+            for (f_n, f) in data_struct.fields.iter().enumerate() {
+                if is_ignored(&f) {
+                    continue;
+                }
+
+                let f_name = match f.ident {
+                    Some(ref i) => quote! { self.#i },
+                    None => quote! { self.#f_n },
+                };
+
+                let builder = if field_n == f_n {
+                    // The behaviour that triggered the event. Thus, instead of
+                    // creating a new handler, use the provided handler.
+                    quote! { provided_handler }
+                } else {
+                    quote! { #f_name.new_handler() }
+                };
+
+                match out_handler {
+                    Some(h) => {
+                        out_handler = Some(quote! { #into_protocols_handler::select(#h, #builder) })
+                    }
+                    ref mut h @ None => *h = Some(builder),
+                }
+            }
+
+            out_handler.unwrap_or(quote! {()}) // TODO: See test `empty`.
+        };
+
         let generate_event_match_arm = if event_process {
             quote! {
                 std::task::Poll::Ready(#network_behaviour_action::GenerateEvent(event)) => {
@@ -518,11 +622,11 @@ fn build_struct(ast: &DeriveInput, data_struct: &DataStruct) -> TokenStream {
             loop {
                 match #trait_to_impl::poll(&mut #field_name, cx, poll_params) {
                     #generate_event_match_arm
-                    std::task::Poll::Ready(#network_behaviour_action::DialAddress { address }) => {
-                        return std::task::Poll::Ready(#network_behaviour_action::DialAddress { address });
+                    std::task::Poll::Ready(#network_behaviour_action::DialAddress { address, handler: provided_handler }) => {
+                        return std::task::Poll::Ready(#network_behaviour_action::DialAddress { address, handler: #provided_handler_and_new_handlers });
                     }
-                    std::task::Poll::Ready(#network_behaviour_action::DialPeer { peer_id, condition }) => {
-                        return std::task::Poll::Ready(#network_behaviour_action::DialPeer { peer_id, condition });
+                    std::task::Poll::Ready(#network_behaviour_action::DialPeer { peer_id, condition, handler: provided_handler }) => {
+                        return std::task::Poll::Ready(#network_behaviour_action::DialPeer { peer_id, condition, handler: #provided_handler_and_new_handlers });
                     }
                     std::task::Poll::Ready(#network_behaviour_action::NotifyHandler { peer_id, handler, event }) => {
                         return std::task::Poll::Ready(#network_behaviour_action::NotifyHandler {
@@ -578,7 +682,7 @@ fn build_struct(ast: &DeriveInput, data_struct: &DataStruct) -> TokenStream {
                 #(#inject_address_change_stmts);*
             }
 
-            fn inject_connection_closed(&mut self, peer_id: &#peer_id, connection_id: &#connection_id, endpoint: &#connected_point) {
+            fn inject_connection_closed(&mut self, peer_id: &#peer_id, connection_id: &#connection_id, endpoint: &#connected_point, handlers: <Self::ProtocolsHandler as #into_protocols_handler>::Handler) {
                 #(#inject_connection_closed_stmts);*
             }
 
@@ -586,10 +690,14 @@ fn build_struct(ast: &DeriveInput, data_struct: &DataStruct) -> TokenStream {
                 #(#inject_addr_reach_failure_stmts);*
             }
 
-            fn inject_dial_failure(&mut self, peer_id: &#peer_id) {
+            fn inject_dial_failure(&mut self, peer_id: &#peer_id, handlers: Self::ProtocolsHandler, error: #dial_error) {
                 #(#inject_dial_failure_stmts);*
             }
 
+            fn inject_listen_failure(&mut self, local_addr: &#multiaddr, send_back_addr: &#multiaddr, handlers: Self::ProtocolsHandler) {
+                #(#inject_listen_failure_stmts);*
+            }
+
             fn inject_new_listener(&mut self, id: #listener_id) {
                 #(#inject_new_listener_stmts);*
             }
@@ -629,10 +737,10 @@ fn build_struct(ast: &DeriveInput, data_struct: &DataStruct) -> TokenStream {
                 }
             }
 
-            fn poll(&mut self, cx: &mut std::task::Context, poll_params: &mut impl #poll_parameters) -> std::task::Poll<#network_behaviour_action<<<Self::ProtocolsHandler as #into_protocols_handler>::Handler as #protocols_handler>::InEvent, Self::OutEvent>> {
+            fn poll(&mut self, cx: &mut std::task::Context, poll_params: &mut impl #poll_parameters) -> std::task::Poll<#network_behaviour_action<Self::OutEvent, Self::ProtocolsHandler>> {
                 use libp2p::futures::prelude::*;
                 #(#poll_stmts)*
-                let f: std::task::Poll<#network_behaviour_action<<<Self::ProtocolsHandler as #into_protocols_handler>::Handler as #protocols_handler>::InEvent, Self::OutEvent>> = #poll_method;
+                let f: std::task::Poll<#network_behaviour_action<Self::OutEvent, Self::ProtocolsHandler>> = #poll_method;
                 f
             }
         }