From be5d15030c4775af022199c37f35afda660f16b0 Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Wed, 30 Nov 2022 17:42:41 +0800
Subject: [PATCH 01/14] log bandwidth at `StreamMuxer` level

An alternative approach to #3161 suggested in https://github.com/libp2p/rust-libp2p/issues/3157#issuecomment-1326126621

Closes #3157
---
 {src => core/src}/bandwidth.rs | 123 +++++++++++----------------------
 core/src/lib.rs                |   1 +
 src/lib.rs                     |   4 --
 src/transport_ext.rs           |  42 -----------
 swarm/src/lib.rs               |  20 ++++++
 5 files changed, 62 insertions(+), 128 deletions(-)
 rename {src => core/src}/bandwidth.rs (66%)
 delete mode 100644 src/transport_ext.rs

diff --git a/src/bandwidth.rs b/core/src/bandwidth.rs
similarity index 66%
rename from src/bandwidth.rs
rename to core/src/bandwidth.rs
index a58eec95ddb..56e0c69028f 100644
--- a/src/bandwidth.rs
+++ b/core/src/bandwidth.rs
@@ -18,20 +18,13 @@
 // FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 // DEALINGS IN THE SOFTWARE.
 
-use crate::{
-    core::{
-        transport::{TransportError, TransportEvent},
-        Transport,
-    },
-    Multiaddr,
-};
+use crate::muxing::{StreamMuxer, StreamMuxerEvent};
 
 use futures::{
     io::{IoSlice, IoSliceMut},
     prelude::*,
     ready,
 };
-use libp2p_core::transport::ListenerId;
 use std::{
     convert::TryFrom as _,
     io,
@@ -43,8 +36,8 @@ use std::{
     task::{Context, Poll},
 };
 
-/// Wraps around a `Transport` and counts the number of bytes that go through all the opened
-/// connections.
+/// Wraps around a `StreamMuxer` and counts the number of bytes that go through all the opened
+/// streams.
 #[derive(Clone)]
 #[pin_project::pin_project]
 pub struct BandwidthLogging<TInner> {
@@ -54,99 +47,57 @@ pub struct BandwidthLogging<TInner> {
 }
 
 impl<TInner> BandwidthLogging<TInner> {
-    /// Creates a new [`BandwidthLogging`] around the transport.
-    pub fn new(inner: TInner) -> (Self, Arc<BandwidthSinks>) {
-        let sink = Arc::new(BandwidthSinks {
-            inbound: AtomicU64::new(0),
-            outbound: AtomicU64::new(0),
-        });
-
-        let trans = BandwidthLogging {
-            inner,
-            sinks: sink.clone(),
-        };
-
-        (trans, sink)
+    /// Creates a new [`BandwidthLogging`] around the stream muxer.
+    pub fn new(inner: TInner, sinks: Arc<BandwidthSinks>) -> Self {
+        Self { inner, sinks }
     }
 }
 
-impl<TInner> Transport for BandwidthLogging<TInner>
+impl<SMInner> StreamMuxer for BandwidthLogging<SMInner>
 where
-    TInner: Transport,
+    SMInner: StreamMuxer,
 {
-    type Output = BandwidthConnecLogging<TInner::Output>;
-    type Error = TInner::Error;
-    type ListenerUpgrade = BandwidthFuture<TInner::ListenerUpgrade>;
-    type Dial = BandwidthFuture<TInner::Dial>;
+    type Substream = BandwidthConnecLogging<SMInner::Substream>;
+    type Error = SMInner::Error;
 
     fn poll(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
-    ) -> Poll<TransportEvent<Self::ListenerUpgrade, Self::Error>> {
+    ) -> Poll<Result<StreamMuxerEvent, Self::Error>> {
         let this = self.project();
-        match this.inner.poll(cx) {
-            Poll::Ready(event) => {
-                let event = event.map_upgrade({
-                    let sinks = this.sinks.clone();
-                    |inner| BandwidthFuture { inner, sinks }
-                });
-                Poll::Ready(event)
-            }
-            Poll::Pending => Poll::Pending,
-        }
+        this.inner.poll(cx)
     }
 
-    fn listen_on(&mut self, addr: Multiaddr) -> Result<ListenerId, TransportError<Self::Error>> {
-        self.inner.listen_on(addr)
-    }
-
-    fn remove_listener(&mut self, id: ListenerId) -> bool {
-        self.inner.remove_listener(id)
-    }
-
-    fn dial(&mut self, addr: Multiaddr) -> Result<Self::Dial, TransportError<Self::Error>> {
-        let sinks = self.sinks.clone();
-        self.inner
-            .dial(addr)
-            .map(move |fut| BandwidthFuture { inner: fut, sinks })
-    }
-
-    fn dial_as_listener(
-        &mut self,
-        addr: Multiaddr,
-    ) -> Result<Self::Dial, TransportError<Self::Error>> {
-        let sinks = self.sinks.clone();
-        self.inner
-            .dial_as_listener(addr)
-            .map(move |fut| BandwidthFuture { inner: fut, sinks })
-    }
-
-    fn address_translation(&self, server: &Multiaddr, observed: &Multiaddr) -> Option<Multiaddr> {
-        self.inner.address_translation(server, observed)
+    fn poll_inbound(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<Self::Substream, Self::Error>> {
+        let this = self.project();
+        let inner = ready!(this.inner.poll_inbound(cx)?);
+        let logged = BandwidthConnecLogging {
+            inner,
+            sinks: this.sinks.clone(),
+        };
+        Poll::Ready(Ok(logged))
     }
-}
-
-/// Wraps around a `Future` that produces a connection. Wraps the connection around a bandwidth
-/// counter.
-#[pin_project::pin_project]
-pub struct BandwidthFuture<TInner> {
-    #[pin]
-    inner: TInner,
-    sinks: Arc<BandwidthSinks>,
-}
-
-impl<TInner: TryFuture> Future for BandwidthFuture<TInner> {
-    type Output = Result<BandwidthConnecLogging<TInner::Ok>, TInner::Error>;
 
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+    fn poll_outbound(
+        self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+    ) -> Poll<Result<Self::Substream, Self::Error>> {
         let this = self.project();
-        let inner = ready!(this.inner.try_poll(cx)?);
+        let inner = ready!(this.inner.poll_outbound(cx)?);
         let logged = BandwidthConnecLogging {
             inner,
             sinks: this.sinks.clone(),
         };
         Poll::Ready(Ok(logged))
     }
+
+    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        let this = self.project();
+        this.inner.poll_close(cx)
+    }
 }
 
 /// Allows obtaining the average bandwidth of the connections created from a [`BandwidthLogging`].
@@ -156,6 +107,14 @@ pub struct BandwidthSinks {
 }
 
 impl BandwidthSinks {
+    /// Returns a new [`BandwidthSinks`].
+    pub fn new() -> Arc<Self> {
+        Arc::new(Self {
+            inbound: AtomicU64::new(0),
+            outbound: AtomicU64::new(0),
+        })
+    }
+
     /// Returns the total number of bytes that have been downloaded on all the connections spawned
     /// through the [`BandwidthLogging`].
     ///
diff --git a/core/src/lib.rs b/core/src/lib.rs
index 2b20f5156e4..33738fcedda 100644
--- a/core/src/lib.rs
+++ b/core/src/lib.rs
@@ -58,6 +58,7 @@ pub type Negotiated<T> = multistream_select::Negotiated<T>;
 mod peer_id;
 mod translation;
 
+pub mod bandwidth;
 pub mod connection;
 pub mod either;
 pub mod identity;
diff --git a/src/lib.rs b/src/lib.rs
index 1d80de1bd89..370fbe3ff52 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -137,9 +137,6 @@ pub use libp2p_websocket as websocket;
 #[doc(inline)]
 pub use libp2p_yamux as yamux;
 
-mod transport_ext;
-
-pub mod bandwidth;
 pub mod simple;
 
 #[cfg(doc)]
@@ -154,7 +151,6 @@ pub use self::core::{
 pub use self::multiaddr::{multiaddr as build_multiaddr, Multiaddr};
 pub use self::simple::SimpleProtocol;
 pub use self::swarm::Swarm;
-pub use self::transport_ext::TransportExt;
 
 /// Builds a `Transport` based on TCP/IP that supports the most commonly-used features of libp2p:
 ///
diff --git a/src/transport_ext.rs b/src/transport_ext.rs
deleted file mode 100644
index fa8926c8380..00000000000
--- a/src/transport_ext.rs
+++ /dev/null
@@ -1,42 +0,0 @@
-// Copyright 2018 Parity Technologies (UK) Ltd.
-//
-// Permission is hereby granted, free of charge, to any person obtaining a
-// copy of this software and associated documentation files (the "Software"),
-// to deal in the Software without restriction, including without limitation
-// the rights to use, copy, modify, merge, publish, distribute, sublicense,
-// and/or sell copies of the Software, and to permit persons to whom the
-// Software is furnished to do so, subject to the following conditions:
-//
-// The above copyright notice and this permission notice shall be included in
-// all copies or substantial portions of the Software.
-//
-// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
-// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
-// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
-// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
-// DEALINGS IN THE SOFTWARE.
-
-//! Provides the `TransportExt` trait.
-
-use crate::{bandwidth::BandwidthLogging, bandwidth::BandwidthSinks, Transport};
-use std::sync::Arc;
-
-/// Trait automatically implemented on all objects that implement `Transport`. Provides some
-/// additional utilities.
-pub trait TransportExt: Transport {
-    /// Adds a layer on the `Transport` that logs all trafic that passes through the sockets
-    /// created by it.
-    ///
-    /// This method returns an `Arc<BandwidthSinks>` that can be used to retreive the total number
-    /// of bytes transferred through the sockets.
-    fn with_bandwidth_logging(self) -> (BandwidthLogging<Self>, Arc<BandwidthSinks>)
-    where
-        Self: Sized,
-    {
-        BandwidthLogging::new(self)
-    }
-}
-
-impl<TTransport> TransportExt for TTransport where TTransport: Transport {}
diff --git a/swarm/src/lib.rs b/swarm/src/lib.rs
index 0c5cadcf021..f34ce97ac6d 100644
--- a/swarm/src/lib.rs
+++ b/swarm/src/lib.rs
@@ -126,6 +126,7 @@ use futures::{executor::ThreadPoolBuilder, prelude::*, stream::FusedStream};
 use libp2p_core::connection::ConnectionId;
 use libp2p_core::muxing::SubstreamBox;
 use libp2p_core::{
+    bandwidth::{BandwidthLogging, BandwidthSinks},
     connection::ConnectedPoint,
     multiaddr::Protocol,
     multihash::Multihash,
@@ -143,6 +144,7 @@ use std::{
     convert::TryFrom,
     error, fmt, io,
     pin::Pin,
+    sync::Arc,
     task::{Context, Poll},
 };
 use upgrade::UpgradeInfoSend as _;
@@ -457,6 +459,24 @@ where
         SwarmBuilder::without_executor(transport, behaviour, local_peer_id).build()
     }
 
+    /// Adds a layer on the `Swarm` that logs all trafic that passes through the sockets
+    /// created by it.
+    ///
+    /// This method returns an `Arc<BandwidthSinks>` that can be used to retreive the total number
+    /// of bytes transferred through the sockets.
+    pub fn with_bandwidth_logging(mut self) -> Arc<BandwidthSinks> {
+        let sinks = BandwidthSinks::new();
+        let sinks_copy = sinks.clone();
+        self.transport = Transport::map(self.transport, |(peer_id, stream_muxer_box), _| {
+            (
+                peer_id,
+                StreamMuxerBox::new(BandwidthLogging::new(stream_muxer_box, sinks_copy)),
+            )
+        })
+        .boxed();
+        sinks
+    }
+
     /// Returns information about the connections underlying the [`Swarm`].
     pub fn network_info(&self) -> NetworkInfo {
         let num_peers = self.pool.num_peers();

From fac29fc95ed100f5fd7efedca8a1ac57e752dd83 Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Thu, 1 Dec 2022 13:01:23 +0800
Subject: [PATCH 02/14] new transport ext

---
 core/src/lib.rs                |  1 -
 {core/src => src}/bandwidth.rs |  2 +-
 src/lib.rs                     |  2 +
 src/transport_ext.rs           | 67 ++++++++++++++++++++++++++++++++++
 swarm/src/lib.rs               | 20 ----------
 5 files changed, 70 insertions(+), 22 deletions(-)
 rename {core/src => src}/bandwidth.rs (99%)
 create mode 100644 src/transport_ext.rs

diff --git a/core/src/lib.rs b/core/src/lib.rs
index 33738fcedda..2b20f5156e4 100644
--- a/core/src/lib.rs
+++ b/core/src/lib.rs
@@ -58,7 +58,6 @@ pub type Negotiated<T> = multistream_select::Negotiated<T>;
 mod peer_id;
 mod translation;
 
-pub mod bandwidth;
 pub mod connection;
 pub mod either;
 pub mod identity;
diff --git a/core/src/bandwidth.rs b/src/bandwidth.rs
similarity index 99%
rename from core/src/bandwidth.rs
rename to src/bandwidth.rs
index 56e0c69028f..b3fc924d6ba 100644
--- a/core/src/bandwidth.rs
+++ b/src/bandwidth.rs
@@ -18,7 +18,7 @@
 // FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 // DEALINGS IN THE SOFTWARE.
 
-use crate::muxing::{StreamMuxer, StreamMuxerEvent};
+use crate::core::muxing::{StreamMuxer, StreamMuxerEvent};
 
 use futures::{
     io::{IoSlice, IoSliceMut},
diff --git a/src/lib.rs b/src/lib.rs
index 370fbe3ff52..fc9f1b8330e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -137,7 +137,9 @@ pub use libp2p_websocket as websocket;
 #[doc(inline)]
 pub use libp2p_yamux as yamux;
 
+pub mod bandwidth;
 pub mod simple;
+pub mod transport_ext;
 
 #[cfg(doc)]
 pub mod tutorials;
diff --git a/src/transport_ext.rs b/src/transport_ext.rs
new file mode 100644
index 00000000000..8ae138d4b9b
--- /dev/null
+++ b/src/transport_ext.rs
@@ -0,0 +1,67 @@
+// Copyright 2018 Parity Technologies (UK) Ltd.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a
+// copy of this software and associated documentation files (the "Software"),
+// to deal in the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
+// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
+// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS IN THE SOFTWARE.
+
+//! Provides the `TransportExt` trait.
+
+use crate::core::{
+    muxing::{StreamMuxer, StreamMuxerBox},
+    transport::{map::Map, Boxed},
+    PeerId,
+};
+use crate::{
+    bandwidth::{BandwidthLogging, BandwidthSinks},
+    Transport,
+};
+use std::sync::Arc;
+
+/// Trait automatically implemented on all objects that implement `Transport`. Provides some
+/// additional utilities.
+pub trait TransportExt: Transport {
+    /// Adds a layer on the `Transport` that logs all trafic that passes through the sockets
+    /// created by it.
+    ///
+    /// This method returns an `Arc<BandwidthSinks>` that can be used to retreive the total number
+    /// of bytes transferred through the sockets.
+    fn with_bandwidth_logging<S>(self) -> (Boxed<(PeerId, StreamMuxerBox)>, Arc<BandwidthSinks>)
+    where
+        Self: Sized + Send + Unpin + 'static,
+        Self::Dial: Send + 'static,
+        Self::ListenerUpgrade: Send + 'static,
+        Self::Error: Send + Sync,
+        Self::Output: Into<(PeerId, S)>,
+        S: StreamMuxer + Send + 'static,
+        S::Substream: Send + 'static,
+        S::Error: Send + Sync + 'static,
+    {
+        let sinks = BandwidthSinks::new();
+        let sinks_copy = sinks.clone();
+        let transport = Transport::map(self, |output, _| {
+            let (peer_id, stream_muxer_box) = output.into();
+            (
+                peer_id,
+                StreamMuxerBox::new(BandwidthLogging::new(stream_muxer_box, sinks_copy)),
+            )
+        })
+        .boxed();
+        (transport, sinks)
+    }
+}
+
+impl<TTransport> TransportExt for TTransport where TTransport: Transport {}
diff --git a/swarm/src/lib.rs b/swarm/src/lib.rs
index f34ce97ac6d..0c5cadcf021 100644
--- a/swarm/src/lib.rs
+++ b/swarm/src/lib.rs
@@ -126,7 +126,6 @@ use futures::{executor::ThreadPoolBuilder, prelude::*, stream::FusedStream};
 use libp2p_core::connection::ConnectionId;
 use libp2p_core::muxing::SubstreamBox;
 use libp2p_core::{
-    bandwidth::{BandwidthLogging, BandwidthSinks},
     connection::ConnectedPoint,
     multiaddr::Protocol,
     multihash::Multihash,
@@ -144,7 +143,6 @@ use std::{
     convert::TryFrom,
     error, fmt, io,
     pin::Pin,
-    sync::Arc,
     task::{Context, Poll},
 };
 use upgrade::UpgradeInfoSend as _;
@@ -459,24 +457,6 @@ where
         SwarmBuilder::without_executor(transport, behaviour, local_peer_id).build()
     }
 
-    /// Adds a layer on the `Swarm` that logs all trafic that passes through the sockets
-    /// created by it.
-    ///
-    /// This method returns an `Arc<BandwidthSinks>` that can be used to retreive the total number
-    /// of bytes transferred through the sockets.
-    pub fn with_bandwidth_logging(mut self) -> Arc<BandwidthSinks> {
-        let sinks = BandwidthSinks::new();
-        let sinks_copy = sinks.clone();
-        self.transport = Transport::map(self.transport, |(peer_id, stream_muxer_box), _| {
-            (
-                peer_id,
-                StreamMuxerBox::new(BandwidthLogging::new(stream_muxer_box, sinks_copy)),
-            )
-        })
-        .boxed();
-        sinks
-    }
-
     /// Returns information about the connections underlying the [`Swarm`].
     pub fn network_info(&self) -> NetworkInfo {
         let num_peers = self.pool.num_peers();

From e87d171ed736033dfa2b9d844b9ccd7343254508 Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Thu, 1 Dec 2022 14:22:52 +0800
Subject: [PATCH 03/14] do not break the public API

---
 src/bandwidth.rs     | 14 +++++++++++++-
 src/transport_ext.rs |  7 +++++--
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/src/bandwidth.rs b/src/bandwidth.rs
index b3fc924d6ba..fc030f958ce 100644
--- a/src/bandwidth.rs
+++ b/src/bandwidth.rs
@@ -48,7 +48,19 @@ pub struct BandwidthLogging<TInner> {
 
 impl<TInner> BandwidthLogging<TInner> {
     /// Creates a new [`BandwidthLogging`] around the stream muxer.
-    pub fn new(inner: TInner, sinks: Arc<BandwidthSinks>) -> Self {
+    pub fn new(inner: TInner) -> (Self, Arc<BandwidthSinks>) {
+        let sinks = BandwidthSinks::new();
+        (
+            Self {
+                inner,
+                sinks: sinks.clone(),
+            },
+            sinks,
+        )
+    }
+
+    /// Creates a new [`BandwidthLogging`] around the stream muxer.
+    pub fn new_with_sinks(inner: TInner, sinks: Arc<BandwidthSinks>) -> Self {
         Self { inner, sinks }
     }
 }
diff --git a/src/transport_ext.rs b/src/transport_ext.rs
index 8ae138d4b9b..e37aa6e66d0 100644
--- a/src/transport_ext.rs
+++ b/src/transport_ext.rs
@@ -22,7 +22,7 @@
 
 use crate::core::{
     muxing::{StreamMuxer, StreamMuxerBox},
-    transport::{map::Map, Boxed},
+    transport::Boxed,
     PeerId,
 };
 use crate::{
@@ -56,7 +56,10 @@ pub trait TransportExt: Transport {
             let (peer_id, stream_muxer_box) = output.into();
             (
                 peer_id,
-                StreamMuxerBox::new(BandwidthLogging::new(stream_muxer_box, sinks_copy)),
+                StreamMuxerBox::new(BandwidthLogging::new_with_sinks(
+                    stream_muxer_box,
+                    sinks_copy,
+                )),
             )
         })
         .boxed();

From e6750c1fe695ecfa92020d4b0e1f3e4a9f0a007e Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Thu, 1 Dec 2022 14:42:20 +0800
Subject: [PATCH 04/14] rename TInner to SMInner

---
 src/bandwidth.rs | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/src/bandwidth.rs b/src/bandwidth.rs
index fc030f958ce..dc1baaf81f0 100644
--- a/src/bandwidth.rs
+++ b/src/bandwidth.rs
@@ -40,15 +40,15 @@ use std::{
 /// streams.
 #[derive(Clone)]
 #[pin_project::pin_project]
-pub struct BandwidthLogging<TInner> {
+pub struct BandwidthLogging<SMInner> {
     #[pin]
-    inner: TInner,
+    inner: SMInner,
     sinks: Arc<BandwidthSinks>,
 }
 
-impl<TInner> BandwidthLogging<TInner> {
+impl<SMInner> BandwidthLogging<SMInner> {
     /// Creates a new [`BandwidthLogging`] around the stream muxer.
-    pub fn new(inner: TInner) -> (Self, Arc<BandwidthSinks>) {
+    pub fn new(inner: SMInner) -> (Self, Arc<BandwidthSinks>) {
         let sinks = BandwidthSinks::new();
         (
             Self {
@@ -60,7 +60,7 @@ impl<TInner> BandwidthLogging<TInner> {
     }
 
     /// Creates a new [`BandwidthLogging`] around the stream muxer.
-    pub fn new_with_sinks(inner: TInner, sinks: Arc<BandwidthSinks>) -> Self {
+    pub fn new_with_sinks(inner: SMInner, sinks: Arc<BandwidthSinks>) -> Self {
         Self { inner, sinks }
     }
 }
@@ -148,13 +148,13 @@ impl BandwidthSinks {
 
 /// Wraps around an `AsyncRead + AsyncWrite` and logs the bandwidth that goes through it.
 #[pin_project::pin_project]
-pub struct BandwidthConnecLogging<TInner> {
+pub struct BandwidthConnecLogging<SMInner> {
     #[pin]
-    inner: TInner,
+    inner: SMInner,
     sinks: Arc<BandwidthSinks>,
 }
 
-impl<TInner: AsyncRead> AsyncRead for BandwidthConnecLogging<TInner> {
+impl<SMInner: AsyncRead> AsyncRead for BandwidthConnecLogging<SMInner> {
     fn poll_read(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
@@ -184,7 +184,7 @@ impl<TInner: AsyncRead> AsyncRead for BandwidthConnecLogging<TInner> {
     }
 }
 
-impl<TInner: AsyncWrite> AsyncWrite for BandwidthConnecLogging<TInner> {
+impl<SMInner: AsyncWrite> AsyncWrite for BandwidthConnecLogging<SMInner> {
     fn poll_write(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,

From 964c4940b4e2ff0769b450acdf04df9dcc9b2066 Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Thu, 1 Dec 2022 14:43:55 +0800
Subject: [PATCH 05/14] rename connections to streams

---
 src/bandwidth.rs | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/bandwidth.rs b/src/bandwidth.rs
index dc1baaf81f0..1790061c9b3 100644
--- a/src/bandwidth.rs
+++ b/src/bandwidth.rs
@@ -112,7 +112,7 @@ where
     }
 }
 
-/// Allows obtaining the average bandwidth of the connections created from a [`BandwidthLogging`].
+/// Allows obtaining the average bandwidth of the streams created from a [`BandwidthLogging`].
 pub struct BandwidthSinks {
     inbound: AtomicU64,
     outbound: AtomicU64,
@@ -127,7 +127,7 @@ impl BandwidthSinks {
         })
     }
 
-    /// Returns the total number of bytes that have been downloaded on all the connections spawned
+    /// Returns the total number of bytes that have been downloaded on all the streams spawned
     /// through the [`BandwidthLogging`].
     ///
     /// > **Note**: This method is by design subject to race conditions. The returned value should
@@ -136,7 +136,7 @@ impl BandwidthSinks {
         self.inbound.load(Ordering::Relaxed)
     }
 
-    /// Returns the total number of bytes that have been uploaded on all the connections spawned
+    /// Returns the total number of bytes that have been uploaded on all the streams spawned
     /// through the [`BandwidthLogging`].
     ///
     /// > **Note**: This method is by design subject to race conditions. The returned value should

From 4db21118bfed3af902db00f7b2222896b51ee7c6 Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Thu, 1 Dec 2022 15:30:15 +0800
Subject: [PATCH 06/14] add changelog entry

---
 CHANGELOG.md | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 966b2da7d0d..82375e802dc 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -45,6 +45,10 @@
 
 # `libp2p` facade crate
 
+# Unreleased
+
+- Count bandwidth at the application level ([PR 3180](https://github.com/libp2p/rust-libp2p/pull/3180)).
+
 # 0.50.0
 
 This is a large release. After > 4 years, rust-libp2p ships with an [(alpha) QUIC

From 9e912beb96b3de2597ebe15368eaa73773448f09 Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Thu, 1 Dec 2022 15:30:26 +0800
Subject: [PATCH 07/14] bring back `BandwidthFuture`

to avoid breaking API
---
 src/bandwidth.rs | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/bandwidth.rs b/src/bandwidth.rs
index 1790061c9b3..4235e4cc51c 100644
--- a/src/bandwidth.rs
+++ b/src/bandwidth.rs
@@ -223,3 +223,13 @@ impl<SMInner: AsyncWrite> AsyncWrite for BandwidthConnecLogging<SMInner> {
         this.inner.poll_close(cx)
     }
 }
+
+/// Wraps around a `Future` that produces a connection. Wraps the connection around a bandwidth
+/// counter.
+#[deprecated(since = "0.50.1")]
+#[pin_project::pin_project]
+pub struct BandwidthFuture<TInner> {
+    #[pin]
+    inner: TInner,
+    sinks: Arc<BandwidthSinks>,
+}

From aadfb2ad6e8527346e4574d1e4f2ee40caec043b Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Thu, 1 Dec 2022 15:36:29 +0800
Subject: [PATCH 08/14] make `new_with_sinks` and `BandwidthSinks::new` pub

only within this crate
---
 src/bandwidth.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/bandwidth.rs b/src/bandwidth.rs
index 4235e4cc51c..5f6a1131e3c 100644
--- a/src/bandwidth.rs
+++ b/src/bandwidth.rs
@@ -60,7 +60,7 @@ impl<SMInner> BandwidthLogging<SMInner> {
     }
 
     /// Creates a new [`BandwidthLogging`] around the stream muxer.
-    pub fn new_with_sinks(inner: SMInner, sinks: Arc<BandwidthSinks>) -> Self {
+    pub(crate) fn new_with_sinks(inner: SMInner, sinks: Arc<BandwidthSinks>) -> Self {
         Self { inner, sinks }
     }
 }
@@ -120,7 +120,7 @@ pub struct BandwidthSinks {
 
 impl BandwidthSinks {
     /// Returns a new [`BandwidthSinks`].
-    pub fn new() -> Arc<Self> {
+    pub(crate) fn new() -> Arc<Self> {
         Arc::new(Self {
             inbound: AtomicU64::new(0),
             outbound: AtomicU64::new(0),

From 3bd3d0bde64839d1f9844b4f422c388698a3b85e Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Thu, 1 Dec 2022 15:42:41 +0800
Subject: [PATCH 09/14] revert changes to transport_ext visibility

---
 src/lib.rs | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/lib.rs b/src/lib.rs
index fc9f1b8330e..1d80de1bd89 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -137,9 +137,10 @@ pub use libp2p_websocket as websocket;
 #[doc(inline)]
 pub use libp2p_yamux as yamux;
 
+mod transport_ext;
+
 pub mod bandwidth;
 pub mod simple;
-pub mod transport_ext;
 
 #[cfg(doc)]
 pub mod tutorials;
@@ -153,6 +154,7 @@ pub use self::core::{
 pub use self::multiaddr::{multiaddr as build_multiaddr, Multiaddr};
 pub use self::simple::SimpleProtocol;
 pub use self::swarm::Swarm;
+pub use self::transport_ext::TransportExt;
 
 /// Builds a `Transport` based on TCP/IP that supports the most commonly-used features of libp2p:
 ///

From fe277948c5e27c135a0b275246f4ba88ef608940 Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Sat, 3 Dec 2022 06:41:59 +0400
Subject: [PATCH 10/14] amend changelog and comments

Co-authored-by: Max Inden <mail@max-inden.de>
---
 CHANGELOG.md     | 4 ++--
 src/bandwidth.rs | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 82375e802dc..e8d3590c895 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -45,9 +45,9 @@
 
 # `libp2p` facade crate
 
-# Unreleased
+# 0.51.0 [unreleased]
 
-- Count bandwidth at the application level ([PR 3180](https://github.com/libp2p/rust-libp2p/pull/3180)).
+- Count bandwidth at the application level. Previously `BandwidthLogging` would implement `Transport` and now implements `StreamMuxer` ([PR 3180](https://github.com/libp2p/rust-libp2p/pull/3180)).
 
 # 0.50.0
 
diff --git a/src/bandwidth.rs b/src/bandwidth.rs
index 5f6a1131e3c..b9316696904 100644
--- a/src/bandwidth.rs
+++ b/src/bandwidth.rs
@@ -36,7 +36,7 @@ use std::{
     task::{Context, Poll},
 };
 
-/// Wraps around a `StreamMuxer` and counts the number of bytes that go through all the opened
+/// Wraps around a [`StreamMuxer`] and counts the number of bytes that go through all the opened
 /// streams.
 #[derive(Clone)]
 #[pin_project::pin_project]
@@ -146,7 +146,7 @@ impl BandwidthSinks {
     }
 }
 
-/// Wraps around an `AsyncRead + AsyncWrite` and logs the bandwidth that goes through it.
+/// Wraps around an [`AsyncRead`] + [`AsyncWrite`] and logs the bandwidth that goes through it.
 #[pin_project::pin_project]
 pub struct BandwidthConnecLogging<SMInner> {
     #[pin]

From c087487a4c95d95c4d1560ecd5ba7312f5d9d835 Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Sat, 3 Dec 2022 10:52:21 +0800
Subject: [PATCH 11/14] remove BandwidthFuture

also
rename BandwidthLogging `new_with_sinks` to `new`
rename BandwidthConnecLogging to InstrumentedStream
---
 CHANGELOG.md         |  3 +++
 src/bandwidth.rs     | 36 +++++++-----------------------------
 src/transport_ext.rs |  2 +-
 3 files changed, 11 insertions(+), 30 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e8d3590c895..ab146b2a5ce 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -48,6 +48,9 @@
 # 0.51.0 [unreleased]
 
 - Count bandwidth at the application level. Previously `BandwidthLogging` would implement `Transport` and now implements `StreamMuxer` ([PR 3180](https://github.com/libp2p/rust-libp2p/pull/3180)).
+    - `BandwidthLogging::new` now requires a 2nd argument: `Arc<BandwidthSinks>`
+    - Remove `BandwidthFuture`
+    - Rename `BandwidthConnecLogging` to `InstrumentedStream`
 
 # 0.50.0
 
diff --git a/src/bandwidth.rs b/src/bandwidth.rs
index b9316696904..48ef765317e 100644
--- a/src/bandwidth.rs
+++ b/src/bandwidth.rs
@@ -48,19 +48,7 @@ pub struct BandwidthLogging<SMInner> {
 
 impl<SMInner> BandwidthLogging<SMInner> {
     /// Creates a new [`BandwidthLogging`] around the stream muxer.
-    pub fn new(inner: SMInner) -> (Self, Arc<BandwidthSinks>) {
-        let sinks = BandwidthSinks::new();
-        (
-            Self {
-                inner,
-                sinks: sinks.clone(),
-            },
-            sinks,
-        )
-    }
-
-    /// Creates a new [`BandwidthLogging`] around the stream muxer.
-    pub(crate) fn new_with_sinks(inner: SMInner, sinks: Arc<BandwidthSinks>) -> Self {
+    pub(crate) fn new(inner: SMInner, sinks: Arc<BandwidthSinks>) -> Self {
         Self { inner, sinks }
     }
 }
@@ -69,7 +57,7 @@ impl<SMInner> StreamMuxer for BandwidthLogging<SMInner>
 where
     SMInner: StreamMuxer,
 {
-    type Substream = BandwidthConnecLogging<SMInner::Substream>;
+    type Substream = InstrumentedStream<SMInner::Substream>;
     type Error = SMInner::Error;
 
     fn poll(
@@ -86,7 +74,7 @@ where
     ) -> Poll<Result<Self::Substream, Self::Error>> {
         let this = self.project();
         let inner = ready!(this.inner.poll_inbound(cx)?);
-        let logged = BandwidthConnecLogging {
+        let logged = InstrumentedStream {
             inner,
             sinks: this.sinks.clone(),
         };
@@ -99,7 +87,7 @@ where
     ) -> Poll<Result<Self::Substream, Self::Error>> {
         let this = self.project();
         let inner = ready!(this.inner.poll_outbound(cx)?);
-        let logged = BandwidthConnecLogging {
+        let logged = InstrumentedStream {
             inner,
             sinks: this.sinks.clone(),
         };
@@ -148,13 +136,13 @@ impl BandwidthSinks {
 
 /// Wraps around an [`AsyncRead`] + [`AsyncWrite`] and logs the bandwidth that goes through it.
 #[pin_project::pin_project]
-pub struct BandwidthConnecLogging<SMInner> {
+pub struct InstrumentedStream<SMInner> {
     #[pin]
     inner: SMInner,
     sinks: Arc<BandwidthSinks>,
 }
 
-impl<SMInner: AsyncRead> AsyncRead for BandwidthConnecLogging<SMInner> {
+impl<SMInner: AsyncRead> AsyncRead for InstrumentedStream<SMInner> {
     fn poll_read(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
@@ -184,7 +172,7 @@ impl<SMInner: AsyncRead> AsyncRead for BandwidthConnecLogging<SMInner> {
     }
 }
 
-impl<SMInner: AsyncWrite> AsyncWrite for BandwidthConnecLogging<SMInner> {
+impl<SMInner: AsyncWrite> AsyncWrite for InstrumentedStream<SMInner> {
     fn poll_write(
         self: Pin<&mut Self>,
         cx: &mut Context<'_>,
@@ -223,13 +211,3 @@ impl<SMInner: AsyncWrite> AsyncWrite for BandwidthConnecLogging<SMInner> {
         this.inner.poll_close(cx)
     }
 }
-
-/// Wraps around a `Future` that produces a connection. Wraps the connection around a bandwidth
-/// counter.
-#[deprecated(since = "0.50.1")]
-#[pin_project::pin_project]
-pub struct BandwidthFuture<TInner> {
-    #[pin]
-    inner: TInner,
-    sinks: Arc<BandwidthSinks>,
-}
diff --git a/src/transport_ext.rs b/src/transport_ext.rs
index e37aa6e66d0..6c7c352251c 100644
--- a/src/transport_ext.rs
+++ b/src/transport_ext.rs
@@ -56,7 +56,7 @@ pub trait TransportExt: Transport {
             let (peer_id, stream_muxer_box) = output.into();
             (
                 peer_id,
-                StreamMuxerBox::new(BandwidthLogging::new_with_sinks(
+                StreamMuxerBox::new(BandwidthLogging::new(
                     stream_muxer_box,
                     sinks_copy,
                 )),

From 307af998521c4e8aa7b608bc16a9187abc77703b Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Mon, 12 Dec 2022 16:43:46 +0800
Subject: [PATCH 12/14] add a doc test

---
 src/bandwidth.rs     |  4 ++--
 src/transport_ext.rs | 37 ++++++++++++++++++++++++++++++-------
 2 files changed, 32 insertions(+), 9 deletions(-)

diff --git a/src/bandwidth.rs b/src/bandwidth.rs
index 48ef765317e..bd65ec8ffe1 100644
--- a/src/bandwidth.rs
+++ b/src/bandwidth.rs
@@ -40,7 +40,7 @@ use std::{
 /// streams.
 #[derive(Clone)]
 #[pin_project::pin_project]
-pub struct BandwidthLogging<SMInner> {
+pub(crate) struct BandwidthLogging<SMInner> {
     #[pin]
     inner: SMInner,
     sinks: Arc<BandwidthSinks>,
@@ -136,7 +136,7 @@ impl BandwidthSinks {
 
 /// Wraps around an [`AsyncRead`] + [`AsyncWrite`] and logs the bandwidth that goes through it.
 #[pin_project::pin_project]
-pub struct InstrumentedStream<SMInner> {
+pub(crate) struct InstrumentedStream<SMInner> {
     #[pin]
     inner: SMInner,
     sinks: Arc<BandwidthSinks>,
diff --git a/src/transport_ext.rs b/src/transport_ext.rs
index 6c7c352251c..9b42229a45b 100644
--- a/src/transport_ext.rs
+++ b/src/transport_ext.rs
@@ -34,11 +34,37 @@ use std::sync::Arc;
 /// Trait automatically implemented on all objects that implement `Transport`. Provides some
 /// additional utilities.
 pub trait TransportExt: Transport {
-    /// Adds a layer on the `Transport` that logs all trafic that passes through the sockets
+    /// Adds a layer on the `Transport` that logs all trafic that passes through the streams
     /// created by it.
     ///
-    /// This method returns an `Arc<BandwidthSinks>` that can be used to retreive the total number
-    /// of bytes transferred through the sockets.
+    /// This method returns an `Arc<BandwidthSinks>` that can be used to retrieve the total number
+    /// of bytes transferred through the streams.
+    ///
+    /// # Example
+    ///
+    /// ```ignore
+    /// use libp2p::{
+    ///     core::upgrade,
+    ///     identity, mplex, noise,
+    ///     tcp,
+    ///     TransportExt,
+    ///     Transport,
+    /// };
+    /// use std::error::Error;
+    ///
+    /// let id_keys = identity::Keypair::generate_ed25519();
+    ///
+    /// let transport = tcp::tokio::Transport::new(tcp::Config::default().nodelay(true))
+    ///     .upgrade(upgrade::Version::V1)
+    ///     .authenticate(
+    ///         noise::NoiseAuthenticated::xx(&id_keys)
+    ///             .expect("Signing libp2p-noise static DH keypair failed."),
+    ///     )
+    ///     .multiplex(mplex::MplexConfig::new())
+    ///     .boxed();
+    ///
+    /// let (transport, sinks) = transport.with_bandwidth_logging();
+    /// ```
     fn with_bandwidth_logging<S>(self) -> (Boxed<(PeerId, StreamMuxerBox)>, Arc<BandwidthSinks>)
     where
         Self: Sized + Send + Unpin + 'static,
@@ -56,10 +82,7 @@ pub trait TransportExt: Transport {
             let (peer_id, stream_muxer_box) = output.into();
             (
                 peer_id,
-                StreamMuxerBox::new(BandwidthLogging::new(
-                    stream_muxer_box,
-                    sinks_copy,
-                )),
+                StreamMuxerBox::new(BandwidthLogging::new(stream_muxer_box, sinks_copy)),
             )
         })
         .boxed();

From 3bbaeea0e1587b0b11a0941906bac2d6756c9d00 Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Tue, 13 Dec 2022 12:54:38 +0800
Subject: [PATCH 13/14] enable doc test

---
 Cargo.toml           | 4 ++++
 src/transport_ext.rs | 9 +++++----
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 6ebcf2d45c5..267369205d3 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -136,6 +136,10 @@ env_logger = "0.10.0"
 clap = { version = "4.0.13", features = ["derive"] }
 tokio = { version = "1.15", features = ["io-util", "io-std", "macros", "rt", "rt-multi-thread"] }
 
+libp2p-mplex = { version = "0.38.0", path = "muxers/mplex" }
+libp2p-noise = { version = "0.41.0", path = "transports/noise" }
+libp2p-tcp = { version = "0.38.0", path = "transports/tcp", features = ["tokio"] }
+
 [workspace]
 members = [
     "core",
diff --git a/src/transport_ext.rs b/src/transport_ext.rs
index 9b42229a45b..2a4c30f17e3 100644
--- a/src/transport_ext.rs
+++ b/src/transport_ext.rs
@@ -42,15 +42,16 @@ pub trait TransportExt: Transport {
     ///
     /// # Example
     ///
-    /// ```ignore
+    /// ```
+    /// use libp2p_mplex as mplex;
+    /// use libp2p_noise as noise;
+    /// use libp2p_tcp as tcp;
     /// use libp2p::{
     ///     core::upgrade,
-    ///     identity, mplex, noise,
-    ///     tcp,
+    ///     identity,
     ///     TransportExt,
     ///     Transport,
     /// };
-    /// use std::error::Error;
     ///
     /// let id_keys = identity::Keypair::generate_ed25519();
     ///

From 34b554546644228b6fd8b570c94c1f749cf95061 Mon Sep 17 00:00:00 2001
From: Anton Kaliaev <anton.kalyaev@gmail.com>
Date: Tue, 13 Dec 2022 13:38:43 +0800
Subject: [PATCH 14/14] fix doc build

---
 src/bandwidth.rs | 8 +++-----
 1 file changed, 3 insertions(+), 5 deletions(-)

diff --git a/src/bandwidth.rs b/src/bandwidth.rs
index bd65ec8ffe1..dc696ce07e2 100644
--- a/src/bandwidth.rs
+++ b/src/bandwidth.rs
@@ -100,7 +100,7 @@ where
     }
 }
 
-/// Allows obtaining the average bandwidth of the streams created from a [`BandwidthLogging`].
+/// Allows obtaining the average bandwidth of the streams.
 pub struct BandwidthSinks {
     inbound: AtomicU64,
     outbound: AtomicU64,
@@ -115,8 +115,7 @@ impl BandwidthSinks {
         })
     }
 
-    /// Returns the total number of bytes that have been downloaded on all the streams spawned
-    /// through the [`BandwidthLogging`].
+    /// Returns the total number of bytes that have been downloaded on all the streams.
     ///
     /// > **Note**: This method is by design subject to race conditions. The returned value should
     /// >           only ever be used for statistics purposes.
@@ -124,8 +123,7 @@ impl BandwidthSinks {
         self.inbound.load(Ordering::Relaxed)
     }
 
-    /// Returns the total number of bytes that have been uploaded on all the streams spawned
-    /// through the [`BandwidthLogging`].
+    /// Returns the total number of bytes that have been uploaded on all the streams.
     ///
     /// > **Note**: This method is by design subject to race conditions. The returned value should
     /// >           only ever be used for statistics purposes.