webrtc: track number of bytes sent / received by all substreams

src/bandwidth.rs

@@ -111,6 +111,140 @@ impl WithBandwidthSinks for (PeerId, libp2p_webrtc::tokio::Connection) {
     }
 }
 
+#[cfg(feature = "deflate")]
+impl<S> WithBandwidthSinks for libp2p_deflate::DeflateOutput<S> {
+    type Output = BandwidthConnecLogging<Self>;
+
+    fn with_bandwidth_sinks(self, sinks: Arc<BandwidthSinks>) -> Self::Output {
+        BandwidthConnecLogging { inner: self, sinks }
+    }
+}
+
+#[cfg(feature = "noise")]
+impl<T, C> WithBandwidthSinks
+    for (
+        libp2p_noise::RemoteIdentity<C>,
+        libp2p_noise::NoiseOutput<T>,
+    )
+{
+    type Output = (
+        libp2p_noise::RemoteIdentity<C>,
+        BandwidthConnecLogging<libp2p_noise::NoiseOutput<T>>,
+    );
+
+    fn with_bandwidth_sinks(self, sinks: Arc<BandwidthSinks>) -> Self::Output {
+        (
+            self.0,
+            BandwidthConnecLogging {
+                inner: self.1,
+                sinks,
+            },
+        )
+    }
+}
+
+#[cfg(feature = "plaintext")]
+impl<C> WithBandwidthSinks for (libp2p_core::PeerId, libp2p_plaintext::PlainTextOutput<C>)
+where
+    C: AsyncRead + AsyncWrite + Unpin,
+{
+    type Output = (
+        libp2p_core::PeerId,
+        BandwidthConnecLogging<libp2p_plaintext::PlainTextOutput<C>>,
+    );
+
+    fn with_bandwidth_sinks(self, sinks: Arc<BandwidthSinks>) -> Self::Output {
+        (
+            self.0,
+            BandwidthConnecLogging {
+                inner: self.1,
+                sinks,
+            },
+        )
+    }
+}
+
+#[cfg(feature = "pnet")]
+impl<S> WithBandwidthSinks for libp2p_pnet::PnetOutput<S> {
+    type Output = BandwidthConnecLogging<Self>;
+
+    fn with_bandwidth_sinks(self, sinks: Arc<BandwidthSinks>) -> Self::Output {
+        BandwidthConnecLogging { inner: self, sinks }
+    }
+}
+
+#[cfg(feature = "quic")]
+impl WithBandwidthSinks for (PeerId, libp2p_quic::Connection) {
+    type Output = (PeerId, Connection<libp2p_quic::Connection>);
+
+    fn with_bandwidth_sinks(self, sinks: Arc<BandwidthSinks>) -> Self::Output {
+        (
+            self.0,
+            Connection {
+                inner: self.1,
+                sinks,
+            },
+        )
+    }
+}
+
+#[cfg(feature = "tls")]
+impl<C> WithBandwidthSinks for (libp2p_core::PeerId, libp2p_tls::TlsStream<C>) {
+    type Output = (
+        libp2p_core::PeerId,
+        BandwidthConnecLogging<libp2p_tls::TlsStream<C>>,
+    );
+
+    fn with_bandwidth_sinks(self, sinks: Arc<BandwidthSinks>) -> Self::Output {
+        (
+            self.0,
+            BandwidthConnecLogging {
+                inner: self.1,
+                sinks,
+            },
+        )
+    }
+}
+
+// #[cfg(all(feature = "uds", feature = "async-std"))]
+// impl WithBandwidthSinks for async_std::os::unix::net::UnixStream {
+//     type Output = BandwidthConnecLogging<Self>;
+
+//     fn with_bandwidth_sinks(self, sinks: Arc<BandwidthSinks>) -> Self::Output {
+//         BandwidthConnecLogging { inner: self, sinks }
+//     }
+// }
+
+// #[cfg(all(feature = "uds", feature = "tokio"))]
+// impl WithBandwidthSinks for tokio::net::UnixStream {
+//     type Output = BandwidthConnecLogging<Self>;
+
+//     fn with_bandwidth_sinks(self, sinks: Arc<BandwidthSinks>) -> Self::Output {
+//         BandwidthConnecLogging { inner: self, sinks }
+//     }
+// }
+
+#[cfg(feature = "wasm-ext")]
+impl WithBandwidthSinks for libp2p_wasm_ext::Connection {
+    type Output = BandwidthConnecLogging<Self>;
+
+    fn with_bandwidth_sinks(self, sinks: Arc<BandwidthSinks>) -> Self::Output {
+        BandwidthConnecLogging { inner: self, sinks }
+    }
+}
+
+#[cfg(feature = "websocket")]
+impl<C> WithBandwidthSinks for libp2p_websocket::RwStreamSink<libp2p_websocket::BytesConnection<C>>
+where
+    C: AsyncWrite + AsyncRead + Unpin + Send + 'static,
+{
+    type Output = BandwidthConnecLogging<Self>;
+
+    fn with_bandwidth_sinks(self, sinks: Arc<BandwidthSinks>) -> Self::Output {
+        BandwidthConnecLogging { inner: self, sinks }
+    }
+}
+
 #[pin_project::pin_project]
 pub struct Connection<I> {
     #[pin]
@@ -161,8 +295,17 @@ pub trait UpdateBandwidthSinks {
 
 #[cfg(feature = "webrtc")]
 impl UpdateBandwidthSinks for libp2p_webrtc::tokio::Connection {
-    fn update_bandwidth_sinks(self: Pin<&mut Self>, _: &BandwidthSinks) {
-        // TODO: Fetch data from connection via public functions and update sinks
+    fn update_bandwidth_sinks(self: Pin<&mut Self>, sinks: &BandwidthSinks) {
+        let (inbound, outbound) = self.fetch_current_bandwidth();
+        sinks.inbound.fetch_add(inbound, Ordering::Relaxed);
+        sinks.outbound.fetch_add(outbound, Ordering::Relaxed);
+    }
+}
+
+#[cfg(feature = "quic")]
+impl UpdateBandwidthSinks for libp2p_quic::Connection {
+    fn update_bandwidth_sinks(self: Pin<&mut Self>, _sinks: &BandwidthSinks) {
+        unimplemented!("get total bytes received / sent from QUIC connection and update sinks");
     }
 }
 

transports/webrtc/src/tokio/bandwidth.rs

@@ -0,0 +1,56 @@
+// Copyright 2022 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.
+
+use std::sync::atomic::{AtomicU64, Ordering};
+
+/// Holds the counters used when logging bandwidth.
+///
+/// Can be safely shared between threads.
+pub(crate) struct Bandwidth {
+    inbound: AtomicU64,
+    outbound: AtomicU64,
+}
+impl Bandwidth {
+    /// Creates a new [`Bandwidth`].
+    pub(crate) fn new() -> Self {
+        Self {
+            inbound: AtomicU64::new(0),
+            outbound: AtomicU64::new(0),
+        }
+    }
+
+    /// Increases the number of bytes received.
+    pub(crate) fn add_inbound(&self, n: u64) {
+        self.inbound.fetch_add(n, Ordering::Relaxed);
+    }
+
+    /// Increases the number of bytes sent for a given address.
+    pub(crate) fn add_outbound(&self, n: u64) {
+        self.outbound.fetch_add(n, Ordering::Relaxed);
+    }
+
+    /// Gets the number of bytes received & sent. Resets both counters to zero after the call.
+    pub(crate) fn reset(&self) -> (u64, u64) {
+        (
+            self.inbound.swap(0, Ordering::Relaxed),
+            self.outbound.swap(0, Ordering::Relaxed),
+        )
+    }
+}

transports/webrtc/src/tokio/connection.rs

@@ -40,6 +40,7 @@ use std::{
     task::{Context, Poll},
 };
 
+use crate::tokio::bandwidth::Bandwidth;
 use crate::tokio::{error::Error, substream, substream::Substream};
 
 /// Maximum number of unprocessed data channels.
@@ -65,13 +66,18 @@ pub struct Connection {
     /// A list of futures, which, once completed, signal that a [`Substream`] has been dropped.
     drop_listeners: FuturesUnordered<substream::DropListener>,
     no_drop_listeners_waker: Option<Waker>,
+
+    /// Bandwidth logging, which is done in `UDPMuxNewAddr`.
+    ///
+    /// Shared between all connections.
+    bandwidth: Arc<Bandwidth>,
 }
 
 impl Unpin for Connection {}
 
 impl Connection {
     /// Creates a new connection.
-    pub(crate) async fn new(rtc_conn: RTCPeerConnection) -> Self {
+    pub(crate) async fn new(rtc_conn: RTCPeerConnection, bandwidth: Arc<Bandwidth>) -> Self {
         let (data_channel_tx, data_channel_rx) = mpsc::channel(MAX_DATA_CHANNELS_IN_FLIGHT);
 
         Connection::register_incoming_data_channels_handler(
@@ -87,6 +93,7 @@ impl Connection {
             close_fut: None,
             drop_listeners: FuturesUnordered::default(),
             no_drop_listeners_waker: None,
+            bandwidth,
         }
     }
 
@@ -144,6 +151,13 @@ impl Connection {
             })
         }));
     }
+
+    /// Fetches the current bandwidth (for all connections). Counters are reset after each call.
+    ///
+    /// Only use this for updating bandwidth.
+    pub fn fetch_current_bandwidth(&self) -> (u64, u64) {
+        self.bandwidth.reset()
+    }
 }
 
 impl StreamMuxer for Connection {

transports/webrtc/src/tokio/mod.rs

@@ -18,6 +18,7 @@
 // FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 // DEALINGS IN THE SOFTWARE.
 
+mod bandwidth;
 pub mod certificate;
 mod connection;
 mod error;

transports/webrtc/src/tokio/transport.rs

@@ -125,22 +125,24 @@ impl libp2p_core::Transport for Transport {
 
         let config = self.config.clone();
         let client_fingerprint = self.config.fingerprint;
-        let udp_mux = self
+        let udp_mux = &self
             .listeners
             .iter()
             .next()
             .ok_or(TransportError::Other(Error::NoListeners))?
-            .udp_mux
-            .udp_mux_handle();
+            .udp_mux;
+        let udp_mux_handle = udp_mux.udp_mux_handle();
+        let bandwidth = udp_mux.bandwidth();
 
         Ok(async move {
             let (peer_id, connection) = upgrade::outbound(
                 sock_addr,
                 config.inner,
-                udp_mux,
+                udp_mux_handle,
                 client_fingerprint,
                 server_fingerprint,
                 config.id_keys,
+                bandwidth,
             )
             .await?;
 
@@ -332,6 +334,7 @@ impl Stream for ListenStream {
                         self.config.fingerprint,
                         new_addr.ufrag,
                         self.config.id_keys.clone(),
+                        self.udp_mux.bandwidth(),
                     )
                     .boxed();
 

transports/webrtc/src/tokio/udp_mux.rs

@@ -43,6 +43,7 @@ use std::{
     task::{Context, Poll},
 };
 
+use crate::tokio::bandwidth::Bandwidth;
 use crate::tokio::req_res_chan;
 
 const RECEIVE_MTU: usize = 8192;
@@ -97,6 +98,14 @@ pub struct UDPMuxNewAddr {
 
     udp_mux_handle: Arc<UdpMuxHandle>,
     udp_mux_writer_handle: Arc<UdpMuxWriterHandle>,
+
+    /// Bandwidth logging.
+    ///
+    /// ICE binding requests or any other WebRTC control data transmitted via UDP socket is not
+    /// logged. Only data flowing from / to connections are being logged.
+    ///
+    /// Shared between all connections.
+    bandwidth: Arc<Bandwidth>,
 }
 
 impl UDPMuxNewAddr {
@@ -128,6 +137,7 @@ impl UDPMuxNewAddr {
             send_command,
             udp_mux_handle: Arc::new(udp_mux_handle),
             udp_mux_writer_handle: Arc::new(udp_mux_writer_handle),
+            bandwidth: Arc::new(Bandwidth::new()),
         })
     }
 
@@ -139,6 +149,10 @@ impl UDPMuxNewAddr {
         self.udp_mux_handle.clone()
     }
 
+    pub(crate) fn bandwidth(&self) -> Arc<Bandwidth> {
+        self.bandwidth.clone()
+    }
+
     /// Create a muxed connection for a given ufrag.
     fn create_muxed_conn(&self, ufrag: &str) -> Result<UDPMuxConn, Error> {
         let local_addr = self.udp_sock.local_addr()?;
@@ -200,6 +214,10 @@ impl UDPMuxNewAddr {
                 Some((buf, target, response)) => {
                     match self.udp_sock.poll_send_to(cx, &buf, target) {
                         Poll::Ready(result) => {
+                            if let Ok(n) = result {
+                                self.bandwidth.add_outbound(u64::try_from(n).unwrap_or(0));
+                            }
+
                             let _ = response.send(result.map_err(|e| Error::Io(e.into())));
                             continue;
                         }
@@ -378,7 +396,10 @@ impl UDPMuxNewAddr {
                                     }
                                 }
                                 Some(conn) => {
-                                    let mut packet = vec![0u8; read.filled().len()];
+                                    let n = read.filled().len();
+                                    self.bandwidth.add_inbound(u64::try_from(n).unwrap_or(0));
+
+                                    let mut packet = vec![0u8; n];
                                     packet.copy_from_slice(read.filled());
                                     self.write_future = OptionFuture::from(Some(
                                         async move {