{core/,swarm/}: Dial with handler and return handler on error and closed

core/src/connection.rs

@@ -229,10 +229,12 @@ where
         self.handler.inject_event(event);
     }
 
+    // TODO: Update comment
+    //
     /// Begins an orderly shutdown of the connection, returning a
     /// `Future` that resolves when connection shutdown is complete.
-    pub fn close(self) -> Close<TMuxer> {
-        self.muxing.close().0
+    pub fn close(self) -> (THandler, Close<TMuxer>) {
+        (self.handler, self.muxing.close().0)
     }
 
     /// Polls the connection for events produced by the associated handler

core/src/connection/manager.rs

@@ -192,6 +192,7 @@ pub enum Event<'a, H: IntoConnectionHandler, TE> {
         /// The error that occurred, if any. If `None`, the connection
         /// has been actively closed.
         error: Option<ConnectionError<THandlerError<H>>>,
+        handler: H::Handler,
     },
 
     /// A connection has been established.
@@ -384,14 +385,15 @@ impl<H: IntoConnectionHandler, TE> Manager<H, TE> {
                         new_endpoint: new,
                     }
                 }
-                task::Event::Closed { id, error } => {
+                task::Event::Closed { id, error, handler } => {
                     let id = ConnectionId(id);
                     let task = task.remove();
                     match task.state {
                         TaskState::Established(connected) => Event::ConnectionClosed {
                             id,
                             connected,
                             error,
+                            handler,
                         },
                         TaskState::Pending => unreachable!(
                             "`Event::Closed` implies (2) occurred on that task and thus (3)."

core/src/connection/manager/task.rs

@@ -71,6 +71,7 @@ pub enum Event<H: IntoConnectionHandler, TE> {
     Closed {
         id: TaskId,
         error: Option<ConnectionError<THandlerError<H>>>,
+        handler: H::Handler,
     },
 }
 
@@ -177,7 +178,7 @@ where
     },
 
     /// The connection is closing (active close).
-    Closing(Close<M>),
+    Closing { closing_muxer: Close<M>, handler: H::Handler },
 
     /// The task is terminating with a final event for the `Manager`.
     Terminating(Event<H, E>),
@@ -265,7 +266,11 @@ where
                                 // Don't accept any further commands.
                                 this.commands.get_mut().close();
                                 // Discard the event, if any, and start a graceful close.
-                                this.state = State::Closing(connection.close());
+                                let (handler, closing_muxer) = connection.close();
+                                this.state = State::Closing {
+                                    handler,
+                                    closing_muxer,
+                                };
                                 continue 'poll;
                             }
                             Poll::Ready(None) => {
@@ -324,36 +329,47 @@ where
                             Poll::Ready(Err(error)) => {
                                 // Don't accept any further commands.
                                 this.commands.get_mut().close();
+                                // TODO: Good idea if there is already an error?
+                                let (handler, _) = connection.close();
                                 // Terminate the task with the error, dropping the connection.
                                 let event = Event::Closed {
                                     id,
                                     error: Some(error),
+                                    handler,
                                 };
                                 this.state = State::Terminating(event);
                             }
                         }
                     }
                 }
 
-                State::Closing(mut closing) => {
+                State::Closing {
+                    handler,
+                    mut closing_muxer,
+                } => {
                     // Try to gracefully close the connection.
-                    match closing.poll_unpin(cx) {
+                    match closing_muxer.poll_unpin(cx) {
                         Poll::Ready(Ok(())) => {
                             let event = Event::Closed {
                                 id: this.id,
                                 error: None,
+                                handler,
                             };
                             this.state = State::Terminating(event);
                         }
                         Poll::Ready(Err(e)) => {
                             let event = Event::Closed {
                                 id: this.id,
                                 error: Some(ConnectionError::IO(e)),
+                                handler,
                             };
                             this.state = State::Terminating(event);
                         }
                         Poll::Pending => {
-                            this.state = State::Closing(closing);
+                            this.state = State::Closing {
+                                handler,
+                                closing_muxer,
+                            };
                             return Poll::Pending;
                         }
                     }

core/src/connection/pool.rs

@@ -102,6 +102,7 @@ pub enum PoolEvent<'a, THandler: IntoConnectionHandler, TTransErr> {
         pool: &'a mut Pool<THandler, TTransErr>,
         /// The remaining number of established connections to the same peer.
         num_established: u32,
+        handler: THandler::Handler,
     },
 
     /// A connection attempt failed.
@@ -114,7 +115,7 @@ pub enum PoolEvent<'a, THandler: IntoConnectionHandler, TTransErr> {
         error: PendingConnectionError<TTransErr>,
         /// The handler that was supposed to handle the connection,
         /// if the connection failed before the handler was consumed.
-        handler: Option<THandler>,
+        handler: THandler,
         /// The (expected) peer of the failed connection.
         peer: Option<PeerId>,
         /// A reference to the pool that managed the connection.
@@ -554,6 +555,7 @@ impl<THandler: IntoConnectionHandler, TTransErr> Pool<THandler, TTransErr> {
                 num_established,
                 error: None,
                 pool: self,
+                handler: todo!(),
             });
         }
 
@@ -572,7 +574,7 @@ impl<THandler: IntoConnectionHandler, TTransErr> Pool<THandler, TTransErr> {
                             id,
                             endpoint,
                             error,
-                            handler: Some(handler),
+                            handler: handler,
                             peer,
                             pool: self,
                         });
@@ -582,6 +584,7 @@ impl<THandler: IntoConnectionHandler, TTransErr> Pool<THandler, TTransErr> {
                     id,
                     connected,
                     error,
+                    handler,
                 } => {
                     let num_established =
                         if let Some(conns) = self.established.get_mut(&connected.peer_id) {
@@ -601,6 +604,7 @@ impl<THandler: IntoConnectionHandler, TTransErr> Pool<THandler, TTransErr> {
                         error,
                         num_established,
                         pool: self,
+                        handler,
                     });
                 }
                 manager::Event::ConnectionEstablished { entry } => {
@@ -610,30 +614,38 @@ impl<THandler: IntoConnectionHandler, TTransErr> Pool<THandler, TTransErr> {
 
                         // Check general established connection limit.
                         if let Err(e) = self.counters.check_max_established(&endpoint) {
-                            let connected = entry.remove();
-                            return Poll::Ready(PoolEvent::PendingConnectionError {
-                                id,
-                                endpoint: connected.endpoint,
-                                error: PendingConnectionError::ConnectionLimit(e),
-                                handler: None,
-                                peer,
-                                pool: self,
-                            });
+                            // TODO: Good idea? How should we let the user know that the close
+                            // happened due to a conneciton limit?
+                            entry.start_close();
+                            // let connected = entry.remove();
+                            // return Poll::Ready(PoolEvent::PendingConnectionError {
+                            //     id,
+                            //     endpoint: connected.endpoint,
+                            //     error: PendingConnectionError::ConnectionLimit(e),
+                            //     handler: None,
+                            //     peer,
+                            //     pool: self,
+                            // });
+                            continue;
                         }
 
                         // Check per-peer established connection limit.
                         let current =
                             num_peer_established(&self.established, &entry.connected().peer_id);
                         if let Err(e) = self.counters.check_max_established_per_peer(current) {
-                            let connected = entry.remove();
-                            return Poll::Ready(PoolEvent::PendingConnectionError {
-                                id,
-                                endpoint: connected.endpoint,
-                                error: PendingConnectionError::ConnectionLimit(e),
-                                handler: None,
-                                peer,
-                                pool: self,
-                            });
+                            // TODO: Good idea? How should we let the user know that the close
+                            // happened due to a conneciton limit?
+                            entry.start_close();
+                            // let connected = entry.remove();
+                            // return Poll::Ready(PoolEvent::PendingConnectionError {
+                            //     id,
+                            //     endpoint: connected.endpoint,
+                            //     error: PendingConnectionError::ConnectionLimit(e),
+                            //     handler: None,
+                            //     peer,
+                            //     pool: self,
+                            // });
+                            continue;
                         }
 
                         // Peer ID checks must already have happened. See `add_pending`.

core/src/network.rs

@@ -22,7 +22,7 @@ mod event;
 pub mod peer;
 
 pub use crate::connection::{ConnectionCounters, ConnectionLimits};
-pub use event::{IncomingConnection, NetworkEvent};
+pub use event::{IncomingConnection, NetworkEvent, DialAttemptsRemaining};
 pub use peer::Peer;
 
 use crate::{
@@ -438,19 +438,22 @@ where
                         log::warn!("Dialing aborted: {:?}", e);
                     }
                 }
+                // TODO: Include handler in event.
                 event
             }
             Poll::Ready(PoolEvent::ConnectionClosed {
                 id,
                 connected,
                 error,
                 num_established,
+                handler,
                 ..
             }) => NetworkEvent::ConnectionClosed {
                 id,
                 connected,
                 num_established,
                 error,
+                handler,
             },
             Poll::Ready(PoolEvent::ConnectionEvent { connection, event }) => {
                 NetworkEvent::ConnectionEvent { connection, event }
@@ -563,7 +566,7 @@ fn on_connection_failed<'a, TTrans, THandler>(
     id: ConnectionId,
     endpoint: ConnectedPoint,
     error: PendingConnectionError<TTrans::Error>,
-    handler: Option<THandler>,
+    handler: THandler,
 ) -> (
     Option<DialingOpts<PeerId, THandler>>,
     NetworkEvent<'a, TTrans, THandlerInEvent<THandler>, THandlerOutEvent<THandler>, THandler>,
@@ -592,27 +595,21 @@ where
         let failed_addr = attempt.current.1.clone();
 
         let (opts, attempts_remaining) = if num_remain > 0 {
-            if let Some(handler) = handler {
-                let next_attempt = attempt.remaining.remove(0);
-                let opts = DialingOpts {
-                    peer: peer_id,
-                    handler,
-                    address: next_attempt,
-                    remaining: attempt.remaining,
-                };
-                (Some(opts), num_remain)
-            } else {
-                // The error is "fatal" for the dialing attempt, since
-                // the handler was already consumed. All potential
-                // remaining connection attempts are thus void.
-                (None, 0)
-            }
+            let next_attempt = attempt.remaining.remove(0);
+            let opts = DialingOpts {
+                peer: peer_id,
+                handler,
+                address: next_attempt,
+                remaining: attempt.remaining,
+            };
+            (Some(opts), DialAttemptsRemaining::Some(num_remain))
         } else {
-            (None, 0)
+            (None, DialAttemptsRemaining::None(handler))
         };
 
         (
             opts,
+            // TODO: This is the place to return the handler.
             NetworkEvent::DialError {
                 attempts_remaining,
                 peer_id,
@@ -625,20 +622,24 @@ where
         match endpoint {
             ConnectedPoint::Dialer { address } => (
                 None,
+                // TODO: This is the place to return the handler.
                 NetworkEvent::UnknownPeerDialError {
                     multiaddr: address,
                     error,
+                    handler,
                 },
             ),
             ConnectedPoint::Listener {
                 local_addr,
                 send_back_addr,
             } => (
                 None,
+                // TODO: This is the place to return the handler.
                 NetworkEvent::IncomingConnectionError {
                     local_addr,
                     send_back_addr,
                     error,
+                    handler,
                 },
             ),
         }