core/muxing: Use total number of alive inbound streams for backpressure

core/CHANGELOG.md

@@ -1,3 +1,9 @@
+# 0.36.1
+
+- Add `StreamMuxerBox::active_inbound_streams` and `StreamMuxerBox::active_outbound_streams`. See [PR 2878]. 
+
+[PR 2878]: https://github.com/libp2p/rust-libp2p/pull/2878
+
 # 0.36.0
 
 - Make RSA keypair support optional. To enable RSA support, `rsa` feature should be enabled.

core/Cargo.toml

@@ -3,7 +3,7 @@ name = "libp2p-core"
 edition = "2021"
 rust-version = "1.56.1"
 description = "Core traits and structs of libp2p"
-version = "0.36.0"
+version = "0.36.1"
 authors = ["Parity Technologies <[email protected]>"]
 license = "MIT"
 repository = "https://github.com/libp2p/rust-libp2p"

core/src/muxing/boxed.rs

@@ -6,18 +6,27 @@ use std::fmt;
 use std::io;
 use std::io::{IoSlice, IoSliceMut};
 use std::pin::Pin;
+use std::sync::{Arc, Weak};
 use std::task::{Context, Poll};
 
 /// Abstract `StreamMuxer`.
 pub struct StreamMuxerBox {
-    inner: Pin<Box<dyn StreamMuxer<Substream = SubstreamBox, Error = io::Error> + Send>>,
+    inner: Pin<
+        Box<
+            dyn private::InstrumentedStreamMuxer<Substream = SubstreamBox, Error = io::Error>
+                + Send,
+        >,
+    >,
 }
 
 /// Abstract type for asynchronous reading and writing.
 ///
 /// A [`SubstreamBox`] erases the concrete type it is given and only retains its `AsyncRead`
 /// and `AsyncWrite` capabilities.
-pub struct SubstreamBox(Pin<Box<dyn AsyncReadWrite + Send>>);
+pub struct SubstreamBox {
+    inner: Pin<Box<dyn AsyncReadWrite + Send>>,
+    _counter: Weak<()>,
+}
 
 #[pin_project]
 struct Wrap<T>
@@ -26,49 +35,79 @@ where
 {
     #[pin]
     inner: T,
+    inbound_counter: Arc<()>,
+    outbound_counter: Arc<()>,
 }
 
-impl<T> StreamMuxer for Wrap<T>
-where
-    T: StreamMuxer,
-    T::Substream: Send + 'static,
-    T::Error: Send + Sync + 'static,
-{
-    type Substream = SubstreamBox;
-    type Error = io::Error;
+mod private {
+    use super::*;
 
-    fn poll_inbound(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-    ) -> Poll<Result<Self::Substream, Self::Error>> {
-        self.project()
-            .inner
-            .poll_inbound(cx)
-            .map_ok(SubstreamBox::new)
-            .map_err(into_io_error)
+    pub trait InstrumentedStreamMuxer: StreamMuxer {
+        fn active_inbound_streams(&self) -> usize;
+        fn active_outbound_streams(&self) -> usize;
     }
 
-    fn poll_outbound(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-    ) -> Poll<Result<Self::Substream, Self::Error>> {
-        self.project()
-            .inner
-            .poll_outbound(cx)
-            .map_ok(SubstreamBox::new)
-            .map_err(into_io_error)
-    }
+    impl<T> StreamMuxer for Wrap<T>
+    where
+        T: StreamMuxer,
+        T::Substream: Send + 'static,
+        T::Error: Send + Sync + 'static,
+    {
+        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<'_>,
+        ) -> Poll<Result<Self::Substream, Self::Error>> {
+            let counter = Arc::downgrade(&self.inbound_counter);
+
+            self.project()
+                .inner
+                .poll_inbound(cx)
+                .map_ok(|stream| SubstreamBox::new(stream, counter))
+                .map_err(into_io_error)
+        }
+
+        fn poll_outbound(
+            self: Pin<&mut Self>,
+            cx: &mut Context<'_>,
+        ) -> Poll<Result<Self::Substream, Self::Error>> {
+            let counter = Arc::downgrade(&self.outbound_counter);
+
+            self.project()
+                .inner
+                .poll_outbound(cx)
+                .map_ok(|stream| SubstreamBox::new(stream, counter))
+                .map_err(into_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(
+            self: Pin<&mut Self>,
+            cx: &mut Context<'_>,
+        ) -> Poll<Result<StreamMuxerEvent, Self::Error>> {
+            self.project().inner.poll(cx).map_err(into_io_error)
+        }
     }
 
-    fn poll(
-        self: Pin<&mut Self>,
-        cx: &mut Context<'_>,
-    ) -> Poll<Result<StreamMuxerEvent, Self::Error>> {
-        self.project().inner.poll(cx).map_err(into_io_error)
+    impl<T> InstrumentedStreamMuxer for Wrap<T>
+    where
+        T: StreamMuxer,
+        T::Substream: Send + 'static,
+        T::Error: Send + Sync + 'static,
+    {
+        fn active_inbound_streams(&self) -> usize {
+            Arc::weak_count(&self.inbound_counter)
+        }
+
+        fn active_outbound_streams(&self) -> usize {
+            Arc::weak_count(&self.outbound_counter)
+        }
     }
 }
 
@@ -87,16 +126,35 @@ impl StreamMuxerBox {
         T::Substream: Send + 'static,
         T::Error: Send + Sync + 'static,
     {
-        let wrap = Wrap { inner: muxer };
+        let wrap = Wrap {
+            inner: muxer,
+            inbound_counter: Arc::new(()),
+            outbound_counter: Arc::new(()),
+        };
 
         StreamMuxerBox {
             inner: Box::pin(wrap),
         }
     }
 
+    /// Returns the number of active inbound streams i.e. streams that
+    /// have been returned from [`StreamMuxer::poll_inbound`] and have not been dropped since.
+    pub fn active_inbound_streams(&self) -> usize {
+        self.inner.active_inbound_streams()
+    }
+
+    /// Returns the number of active outbound streams i.e. streams that
+    /// have been returned from [`StreamMuxer::poll_outbound`] and have not been dropped since.
+    pub fn active_outbound_streams(&self) -> usize {
+        self.inner.active_outbound_streams()
+    }
+
     fn project(
         self: Pin<&mut Self>,
-    ) -> Pin<&mut (dyn StreamMuxer<Substream = SubstreamBox, Error = io::Error> + Send)> {
+    ) -> Pin<
+        &mut (dyn private::InstrumentedStreamMuxer<Substream = SubstreamBox, Error = io::Error>
+                  + Send),
+    > {
         self.get_mut().inner.as_mut()
     }
 }
@@ -134,14 +192,19 @@ impl StreamMuxer for StreamMuxerBox {
 
 impl SubstreamBox {
     /// Construct a new [`SubstreamBox`] from something that implements [`AsyncRead`] and [`AsyncWrite`].
-    pub fn new<S: AsyncRead + AsyncWrite + Send + 'static>(stream: S) -> Self {
-        Self(Box::pin(stream))
+    ///
+    /// TOOD: Remove this constructor.
+    pub fn new<S: AsyncRead + AsyncWrite + Send + 'static>(stream: S, counter: Weak<()>) -> Self {
+        Self {
+            inner: Box::pin(stream),
+            _counter: counter,
+        }
     }
 }
 
 impl fmt::Debug for SubstreamBox {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        write!(f, "SubstreamBox({})", self.0.type_name())
+        write!(f, "SubstreamBox({})", self.inner.type_name())
     }
 }
 
@@ -168,15 +231,15 @@ impl AsyncRead for SubstreamBox {
         cx: &mut Context<'_>,
         buf: &mut [u8],
     ) -> Poll<std::io::Result<usize>> {
-        self.0.as_mut().poll_read(cx, buf)
+        self.inner.as_mut().poll_read(cx, buf)
     }
 
     fn poll_read_vectored(
         mut self: Pin<&mut Self>,
         cx: &mut Context<'_>,
         bufs: &mut [IoSliceMut<'_>],
     ) -> Poll<std::io::Result<usize>> {
-        self.0.as_mut().poll_read_vectored(cx, bufs)
+        self.inner.as_mut().poll_read_vectored(cx, bufs)
     }
 }
 
@@ -186,22 +249,136 @@ impl AsyncWrite for SubstreamBox {
         cx: &mut Context<'_>,
         buf: &[u8],
     ) -> Poll<std::io::Result<usize>> {
-        self.0.as_mut().poll_write(cx, buf)
+        self.inner.as_mut().poll_write(cx, buf)
     }
 
     fn poll_write_vectored(
         mut self: Pin<&mut Self>,
         cx: &mut Context<'_>,
         bufs: &[IoSlice<'_>],
     ) -> Poll<std::io::Result<usize>> {
-        self.0.as_mut().poll_write_vectored(cx, bufs)
+        self.inner.as_mut().poll_write_vectored(cx, bufs)
     }
 
     fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        self.0.as_mut().poll_flush(cx)
+        self.inner.as_mut().poll_flush(cx)
     }
 
     fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        self.0.as_mut().poll_close(cx)
+        self.inner.as_mut().poll_close(cx)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::muxing::StreamMuxerExt;
+
+    #[async_std::test]
+    async fn stream_muxer_box_tracks_alive_inbound_streams() {
+        let mut muxer = StreamMuxerBox::new(DummyStreamMuxer);
+
+        let _stream1 = muxer.next_inbound().await.unwrap();
+        let _stream2 = muxer.next_inbound().await.unwrap();
+
+        assert_eq!(muxer.active_inbound_streams(), 2);
+
+        drop(_stream1);
+
+        assert_eq!(muxer.active_inbound_streams(), 1);
+    }
+
+    #[async_std::test]
+    async fn stream_muxer_box_tracks_alive_outbound_streams() {
+        let mut muxer = StreamMuxerBox::new(DummyStreamMuxer);
+
+        let _stream1 = muxer.next_outbound().await.unwrap();
+        let _stream2 = muxer.next_outbound().await.unwrap();
+
+        assert_eq!(muxer.active_outbound_streams(), 2);
+
+        drop(_stream1);
+
+        assert_eq!(muxer.active_outbound_streams(), 1);
+    }
+
+    #[test]
+    fn stream_muxer_box_starts_with_zero_active_inbound_streams() {
+        let muxer = StreamMuxerBox::new(DummyStreamMuxer);
+
+        let num_active_inbound_streams = muxer.active_inbound_streams();
+
+        assert_eq!(num_active_inbound_streams, 0);
+    }
+
+    #[test]
+    fn stream_muxer_box_starts_with_zero_active_outbound_streams() {
+        let muxer = StreamMuxerBox::new(DummyStreamMuxer);
+
+        let num_active_outbound_streams = muxer.active_outbound_streams();
+
+        assert_eq!(num_active_outbound_streams, 0);
+    }
+
+    struct DummyStreamMuxer;
+
+    impl StreamMuxer for DummyStreamMuxer {
+        type Substream = PendingSubstream;
+        type Error = void::Void;
+
+        fn poll_inbound(
+            self: Pin<&mut Self>,
+            _: &mut Context<'_>,
+        ) -> Poll<Result<Self::Substream, Self::Error>> {
+            Poll::Ready(Ok(PendingSubstream))
+        }
+
+        fn poll_outbound(
+            self: Pin<&mut Self>,
+            _: &mut Context<'_>,
+        ) -> Poll<Result<Self::Substream, Self::Error>> {
+            Poll::Ready(Ok(PendingSubstream))
+        }
+
+        fn poll_close(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+            Poll::Pending
+        }
+
+        fn poll(
+            self: Pin<&mut Self>,
+            _: &mut Context<'_>,
+        ) -> Poll<Result<StreamMuxerEvent, Self::Error>> {
+            Poll::Pending
+        }
+    }
+
+    struct PendingSubstream;
+
+    impl AsyncRead for PendingSubstream {
+        fn poll_read(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+            _buf: &mut [u8],
+        ) -> Poll<io::Result<usize>> {
+            Poll::Pending
+        }
+    }
+
+    impl AsyncWrite for PendingSubstream {
+        fn poll_write(
+            self: Pin<&mut Self>,
+            _cx: &mut Context<'_>,
+            _buf: &[u8],
+        ) -> Poll<io::Result<usize>> {
+            Poll::Pending
+        }
+
+        fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+            Poll::Pending
+        }
+
+        fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
+            Poll::Pending
+        }
     }
 }