core/muxing: Generalise StreamMuxer::poll_address_change to poll

core/CHANGELOG.md

@@ -5,12 +5,14 @@
 - Drop `Unpin` requirement from `SubstreamBox`. See [PR 2762] and [PR 2776].
 - Drop `Sync` requirement on `StreamMuxer` for constructing `StreamMuxerBox`. See [PR 2775].
 - Use `Pin<&mut Self>` as the receiver type for all `StreamMuxer` poll functions. See [PR 2765].
+- Generalise `StreamMuxer::poll_address_change` to `StreamMuxer::poll_event`. See [PR XXXX].
 
 [PR 2724]: https://github.com/libp2p/rust-libp2p/pull/2724
 [PR 2762]: https://github.com/libp2p/rust-libp2p/pull/2762
 [PR 2775]: https://github.com/libp2p/rust-libp2p/pull/2775
 [PR 2776]: https://github.com/libp2p/rust-libp2p/pull/2776
 [PR 2765]: https://github.com/libp2p/rust-libp2p/pull/2765
+[PR XXXX]: https://github.com/libp2p/rust-libp2p/pull/XXXX
 
 # 0.34.0
 

core/src/either.rs

@@ -18,6 +18,7 @@
 // FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 // DEALINGS IN THE SOFTWARE.
 
+use crate::muxing::StreamMuxerEvent;
 use crate::{
     muxing::StreamMuxer,
     transport::{ListenerId, Transport, TransportError, TransportEvent},
@@ -236,22 +237,20 @@ where
         }
     }
 
-    fn poll_address_change(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-    ) -> Poll<Result<Multiaddr, Self::Error>> {
+    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
         match self.project() {
-            EitherOutputProj::First(inner) => inner.poll_address_change(cx).map_err(EitherError::A),
-            EitherOutputProj::Second(inner) => {
-                inner.poll_address_change(cx).map_err(EitherError::B)
-            }
+            EitherOutputProj::First(inner) => inner.poll_close(cx).map_err(EitherError::A),
+            EitherOutputProj::Second(inner) => inner.poll_close(cx).map_err(EitherError::B),
         }
     }
 
-    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+    fn poll_event(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<StreamMuxerEvent, Self::Error>> {
         match self.project() {
-            EitherOutputProj::First(inner) => inner.poll_close(cx).map_err(EitherError::A),
-            EitherOutputProj::Second(inner) => inner.poll_close(cx).map_err(EitherError::B),
+            EitherOutputProj::First(inner) => inner.poll_event(cx).map_err(EitherError::A),
+            EitherOutputProj::Second(inner) => inner.poll_event(cx).map_err(EitherError::B),
         }
     }
 }

core/src/muxing.rs

@@ -86,14 +86,6 @@ pub trait StreamMuxer {
         cx: &mut Context<'_>,
     ) -> Poll<Result<Self::Substream, Self::Error>>;
 
-    /// Poll for an address change of the underlying connection.
-    ///
-    /// Not all implementations may support this feature.
-    fn poll_address_change(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-    ) -> Poll<Result<Multiaddr, Self::Error>>;
-
     /// Closes this `StreamMuxer`.
     ///
     /// After this has returned `Poll::Ready(Ok(()))`, the muxer has become useless. All
@@ -105,6 +97,22 @@ pub trait StreamMuxer {
     /// >           properly informing the remote, there is no difference between this and
     /// >           immediately dropping the muxer.
     fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>>;
+
+    /// Poll for an event of the underlying connection.
+    ///
+    /// In addition to returning an event, this function may be used to perform any kind of background
+    /// work that needs to happen for the muxer to do its work. Implementations can rely on this
+    /// function to be called regularly and unconditionally.
+    fn poll_event(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<StreamMuxerEvent, Self::Error>>;
+}
+
+/// An event produced by a [`StreamMuxer`].
+pub enum StreamMuxerEvent {
+    /// The address of the remote has changed.
+    AddressChange(Multiaddr),
 }
 
 /// Extension trait for [`StreamMuxer`].
@@ -131,15 +139,15 @@ pub trait StreamMuxerExt: StreamMuxer + Sized {
         Pin::new(self).poll_outbound(cx)
     }
 
-    /// Convenience function for calling [`StreamMuxer::poll_address_change`] for [`StreamMuxer`]s that are `Unpin`.
-    fn poll_address_change_unpin(
+    /// Convenience function for calling [`StreamMuxer::poll_event`] for [`StreamMuxer`]s that are `Unpin`.
+    fn poll_event_unpin(
         &mut self,
         cx: &mut Context<'_>,
-    ) -> Poll<Result<Multiaddr, Self::Error>>
+    ) -> Poll<Result<StreamMuxerEvent, Self::Error>>
     where
         Self: Unpin,
     {
-        Pin::new(self).poll_address_change(cx)
+        Pin::new(self).poll_event(cx)
     }
 
     /// Convenience function for calling [`StreamMuxer::poll_close`] for [`StreamMuxer`]s that are `Unpin`.

core/src/muxing/boxed.rs

@@ -1,6 +1,6 @@
+use crate::muxing::StreamMuxerEvent;
 use crate::StreamMuxer;
 use futures::{AsyncRead, AsyncWrite};
-use multiaddr::Multiaddr;
 use pin_project::pin_project;
 use std::error::Error;
 use std::fmt;
@@ -38,11 +38,6 @@ where
     type Substream = SubstreamBox;
     type Error = io::Error;
 
-    #[inline]
-    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        self.project().inner.poll_close(cx).map_err(into_io_error)
-    }
-
     fn poll_inbound(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
@@ -65,14 +60,16 @@ where
             .map_err(into_io_error)
     }
 
-    fn poll_address_change(
+    #[inline]
+    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.project().inner.poll_close(cx).map_err(into_io_error)
+    }
+
+    fn poll_event(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
-    ) -> Poll<Result<Multiaddr, Self::Error>> {
-        self.project()
-            .inner
-            .poll_address_change(cx)
-            .map_err(into_io_error)
+    ) -> Poll<Result<StreamMuxerEvent, Self::Error>> {
+        self.project().inner.poll_event(cx).map_err(into_io_error)
     }
 }
 
@@ -109,11 +106,6 @@ impl StreamMuxer for StreamMuxerBox {
     type Substream = SubstreamBox;
     type Error = io::Error;
 
-    #[inline]
-    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
-        self.project().poll_close(cx)
-    }
-
     fn poll_inbound(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
@@ -128,11 +120,16 @@ impl StreamMuxer for StreamMuxerBox {
         self.project().poll_outbound(cx)
     }
 
-    fn poll_address_change(
+    #[inline]
+    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.project().poll_close(cx)
+    }
+
+    fn poll_event(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
-    ) -> Poll<Result<Multiaddr, Self::Error>> {
-        self.project().poll_address_change(cx)
+    ) -> Poll<Result<StreamMuxerEvent, Self::Error>> {
+        self.project().poll_event(cx)
     }
 }
 

core/src/muxing/singleton.rs

@@ -20,8 +20,8 @@
 
 use crate::{connection::Endpoint, muxing::StreamMuxer};
 
+use crate::muxing::StreamMuxerEvent;
 use futures::prelude::*;
-use multiaddr::Multiaddr;
 use std::cell::Cell;
 use std::pin::Pin;
 use std::{io, task::Context, task::Poll};
@@ -88,14 +88,14 @@ where
         }
     }
 
-    fn poll_address_change(
+    fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+        Poll::Ready(Ok(()))
+    }
+
+    fn poll_event(
         self: Pin<&mut Self>,
         _: &mut Context<'_>,
-    ) -> Poll<Result<Multiaddr, Self::Error>> {
+    ) -> Poll<Result<StreamMuxerEvent, Self::Error>> {
         Poll::Pending
     }
-
-    fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
-        Poll::Ready(Ok(()))
-    }
 }

muxers/mplex/src/lib.rs

@@ -27,9 +27,10 @@ pub use config::{MaxBufferBehaviour, MplexConfig};
 use bytes::Bytes;
 use codec::LocalStreamId;
 use futures::{future, prelude::*, ready};
+use libp2p_core::muxing::StreamMuxerEvent;
 use libp2p_core::{
     upgrade::{InboundUpgrade, OutboundUpgrade, UpgradeInfo},
-    Multiaddr, StreamMuxer,
+    StreamMuxer,
 };
 use parking_lot::Mutex;
 use std::{cmp, iter, pin::Pin, sync::Arc, task::Context, task::Poll};
@@ -105,10 +106,10 @@ where
             .map_ok(|stream_id| Substream::new(stream_id, self.io.clone()))
     }
 
-    fn poll_address_change(
+    fn poll_event(
         self: Pin<&mut Self>,
         _: &mut Context<'_>,
-    ) -> Poll<Result<Multiaddr, Self::Error>> {
+    ) -> Poll<Result<StreamMuxerEvent, Self::Error>> {
         Poll::Pending
     }
 

muxers/yamux/src/lib.rs

@@ -26,9 +26,8 @@ use futures::{
     prelude::*,
     stream::{BoxStream, LocalBoxStream},
 };
-use libp2p_core::muxing::StreamMuxer;
+use libp2p_core::muxing::{StreamMuxer, StreamMuxerEvent};
 use libp2p_core::upgrade::{InboundUpgrade, OutboundUpgrade, UpgradeInfo};
-use libp2p_core::Multiaddr;
 use std::{
     fmt, io, iter, mem,
     pin::Pin,
@@ -124,10 +123,10 @@ where
             .map_err(YamuxError)
     }
 
-    fn poll_address_change(
+    fn poll_event(
         self: Pin<&mut Self>,
         _: &mut Context<'_>,
-    ) -> Poll<Result<Multiaddr, Self::Error>> {
+    ) -> Poll<Result<StreamMuxerEvent, Self::Error>> {
         Poll::Pending
     }
 

swarm/src/connection.rs

@@ -35,7 +35,7 @@ use crate::IntoConnectionHandler;
 use handler_wrapper::HandlerWrapper;
 use libp2p_core::connection::ConnectedPoint;
 use libp2p_core::multiaddr::Multiaddr;
-use libp2p_core::muxing::{StreamMuxerBox, StreamMuxerExt};
+use libp2p_core::muxing::{StreamMuxerBox, StreamMuxerEvent, StreamMuxerExt};
 use libp2p_core::upgrade;
 use libp2p_core::PeerId;
 use std::collections::VecDeque;
@@ -153,6 +153,15 @@ where
                 }
             }
 
+            if let Poll::Ready(event) = self.muxing.poll_event_unpin(cx)? {
+                match event {
+                    StreamMuxerEvent::AddressChange(address) => {
+                        self.handler.inject_address_change(&address);
+                        return Poll::Ready(Ok(Event::AddressChange(address)));
+                    }
+                }
+            }
+
             if !self.open_info.is_empty() {
                 if let Poll::Ready(substream) = self.muxing.poll_outbound_unpin(cx)? {
                     let user_data = self
@@ -171,11 +180,6 @@ where
                 continue; // Go back to the top, handler can potentially make progress again.
             }
 
-            if let Poll::Ready(address) = self.muxing.poll_address_change_unpin(cx)? {
-                self.handler.inject_address_change(&address);
-                return Poll::Ready(Ok(Event::AddressChange(address)));
-            }
-
             return Poll::Pending; // Nothing can make progress, return `Pending`.
         }
     }