fix(swarm): implement ConnectionHandler::poll_close for combinators

Follow-up to #4076.

This is especially relevant since `libp2p-swarm-derive` uses `SelectConnectionHandler`.

Pull-Request: #4794.

Author: mxinden

swarm/src/behaviour/toggle.rs

@@ -363,4 +363,12 @@ where
             }
         }
     }
+
+    fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::ToBehaviour>> {
+        let Some(inner) = self.inner.as_mut() else {
+            return Poll::Ready(None);
+        };
+
+        inner.poll_close(cx)
+    }
 }

swarm/src/handler.rs

@@ -160,6 +160,9 @@ pub trait ConnectionHandler: Send + 'static {
     /// We therefore cannot guarantee that performing IO within here will succeed.
     ///
     /// To signal completion, [`Poll::Ready(None)`] should be returned.
+    ///
+    /// Implementations MUST have a [`fuse`](futures::StreamExt::fuse)-like behaviour.
+    /// That is, [`Poll::Ready(None)`] MUST be returned on repeated calls to [`ConnectionHandler::poll_close`].
     fn poll_close(&mut self, _: &mut Context<'_>) -> Poll<Option<Self::ToBehaviour>> {
         Poll::Ready(None)
     }

swarm/src/handler/either.rs

@@ -134,6 +134,15 @@ where
         Poll::Ready(event)
     }
 
+    fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::ToBehaviour>> {
+        let event = match self {
+            Either::Left(handler) => futures::ready!(handler.poll_close(cx)).map(Either::Left),
+            Either::Right(handler) => futures::ready!(handler.poll_close(cx)).map(Either::Right),
+        };
+
+        Poll::Ready(event)
+    }
+
     fn on_connection_event(
         &mut self,
         event: ConnectionEvent<

swarm/src/handler/map_in.rs

@@ -80,6 +80,10 @@ where
         self.inner.poll(cx)
     }
 
+    fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::ToBehaviour>> {
+        self.inner.poll_close(cx)
+    }
+
     fn on_connection_event(
         &mut self,
         event: ConnectionEvent<

swarm/src/handler/map_out.rs

@@ -21,6 +21,7 @@
 use crate::handler::{
     ConnectionEvent, ConnectionHandler, ConnectionHandlerEvent, SubstreamProtocol,
 };
+use futures::ready;
 use std::fmt::Debug;
 use std::task::{Context, Poll};
 
@@ -83,6 +84,14 @@ where
         })
     }
 
+    fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::ToBehaviour>> {
+        let Some(e) = ready!(self.inner.poll_close(cx)) else {
+            return Poll::Ready(None);
+        };
+
+        Poll::Ready(Some((self.map)(e)))
+    }
+
     fn on_connection_event(
         &mut self,
         event: ConnectionEvent<

swarm/src/handler/multi.rs

@@ -27,7 +27,7 @@ use crate::handler::{
 };
 use crate::upgrade::{InboundUpgradeSend, OutboundUpgradeSend, UpgradeInfoSend};
 use crate::Stream;
-use futures::{future::BoxFuture, prelude::*};
+use futures::{future::BoxFuture, prelude::*, ready};
 use rand::Rng;
 use std::{
     cmp,
@@ -271,6 +271,17 @@ where
 
         Poll::Pending
     }
+
+    fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::ToBehaviour>> {
+        for (k, h) in self.handlers.iter_mut() {
+            let Some(e) = ready!(h.poll_close(cx)) else {
+                continue;
+            };
+            return Poll::Ready(Some((k.clone(), e)));
+        }
+
+        Poll::Ready(None)
+    }
 }
 
 /// Split [`MultiHandler`] into parts.

swarm/src/handler/select.rs

@@ -25,7 +25,7 @@ use crate::handler::{
 };
 use crate::upgrade::SendWrapper;
 use either::Either;
-use futures::future;
+use futures::{future, ready};
 use libp2p_core::upgrade::SelectUpgrade;
 use std::{cmp, task::Context, task::Poll};
 
@@ -259,6 +259,18 @@ where
         Poll::Pending
     }
 
+    fn poll_close(&mut self, cx: &mut Context<'_>) -> Poll<Option<Self::ToBehaviour>> {
+        if let Some(e) = ready!(self.proto1.poll_close(cx)) {
+            return Poll::Ready(Some(Either::Left(e)));
+        }
+
+        if let Some(e) = ready!(self.proto2.poll_close(cx)) {
+            return Poll::Ready(Some(Either::Right(e)));
+        }
+
+        Poll::Ready(None)
+    }
+
     fn on_connection_event(
         &mut self,
         event: ConnectionEvent<