feat(libp2p): track bandwidth per transport protocol stack

examples/metrics/src/main.rs

@@ -37,13 +37,15 @@ mod http_service;
 fn main() -> Result<(), Box<dyn Error>> {
     env_logger::Builder::from_env(Env::default().default_filter_or("info")).init();
 
-    let mut swarm = libp2p::SwarmBuilder::with_new_identity()
+    let (builder, bandwidth_logging) = libp2p::SwarmBuilder::with_new_identity()
         .with_async_std()
         .with_tcp(
             tcp::Config::default(),
             noise::Config::new,
             yamux::Config::default,
         )?
+        .with_bandwidth_logging();
+    let mut swarm = builder
         .with_behaviour(|key| Behaviour::new(key.public()))?
         .with_swarm_config(|cfg| cfg.with_idle_connection_timeout(Duration::from_secs(u64::MAX)))
         .build();
@@ -58,6 +60,7 @@ fn main() -> Result<(), Box<dyn Error>> {
 
     let mut metric_registry = Registry::default();
     let metrics = Metrics::new(&mut metric_registry);
+    libp2p::bandwidth::register_bandwidth_sinks(&mut metric_registry, bandwidth_logging);
     thread::spawn(move || block_on(http_service::metrics_server(metric_registry)));
 
     block_on(async {

libp2p/Cargo.toml

@@ -100,6 +100,7 @@ getrandom = "0.2.3"     # Explicit dependency to be used in `wasm-bindgen` featu
 instant = "0.1.12"      # Explicit dependency to be used in `wasm-bindgen` feature
 # TODO feature flag?
 rw-stream-sink = { workspace = true }
+prometheus-client = { workspace = true }
 
 libp2p-allow-block-list = { workspace = true }
 libp2p-autonat = { workspace = true, optional = true }

libp2p/src/bandwidth.rs

@@ -25,13 +25,18 @@
     prelude::*,
     ready,
 };
+use prometheus_client::{
+    encoding::{DescriptorEncoder, EncodeMetric},
+    metrics::{counter::ConstCounter, MetricType},
+};
 use std::{
+    collections::HashMap,
     convert::TryFrom as _,
     io,
     pin::Pin,
     sync::{
         atomic::{AtomicU64, Ordering},
-        Arc,
+        Arc, RwLock,
     },
     task::{Context, Poll},
 };
@@ -101,6 +106,7 @@
 }
 
 /// Allows obtaining the average bandwidth of the streams.
+#[derive(Default, Debug)]
 pub struct BandwidthSinks {
     inbound: AtomicU64,
     outbound: AtomicU64,
@@ -108,7 +114,7 @@
 
 impl BandwidthSinks {
     /// Returns a new [`BandwidthSinks`].
     pub(crate) fn new() -> Arc<Self> {
         Arc::new(Self {
             inbound: AtomicU64::new(0),
             outbound: AtomicU64::new(0),
@@ -209,3 +215,32 @@
         this.inner.poll_close(cx)
     }
 }
+
+// TODO: Ideally this should go somewhere else. I.e. good to not depend on prometheus-client in libp2p.
+pub fn register_bandwidth_sinks(
+    registry: &mut prometheus_client::registry::Registry,
+    sinks: Arc<RwLock<HashMap<String, Arc<BandwidthSinks>>>>,
+) {
+    registry.register_collector(Box::new(SinksCollector(sinks)));
+}
+
+#[derive(Debug)]
+struct SinksCollector(Arc<RwLock<HashMap<String, Arc<BandwidthSinks>>>>);
+
+impl prometheus_client::collector::Collector for SinksCollector {
+    fn encode(&self, mut encoder: DescriptorEncoder) -> Result<(), std::fmt::Error> {
+        let mut family_encoder =
+            encoder.encode_descriptor("bandwidth", "todo", None, MetricType::Counter)?;
+        for (protocols, sink) in self.0.read().expect("todo").iter() {
+            let labels = [("protocols", protocols.as_str()), ("direction", "inbound")];
+            let metric_encoder = family_encoder.encode_family(&labels)?;
+            ConstCounter::new(sink.inbound.load(Ordering::Relaxed)).encode(metric_encoder)?;
+
+            let labels = [("protocols", protocols.as_str()), ("direction", "outbound")];
+            let metric_encoder = family_encoder.encode_family(&labels)?;
+            ConstCounter::new(sink.outbound.load(Ordering::Relaxed)).encode(metric_encoder)?;
+        }
+
+        Ok(())
+    }
+}

libp2p/src/builder.rs

@@ -73,7 +73,7 @@ mod tests {
     use crate::SwarmBuilder;
     use libp2p_core::{muxing::StreamMuxerBox, transport::dummy::DummyTransport};
     use libp2p_identity::PeerId;
-    use libp2p_swarm::{NetworkBehaviour, Swarm};
+    use libp2p_swarm::NetworkBehaviour;
 
     #[test]
     #[cfg(all(feature = "tokio", feature = "tcp", feature = "tls", feature = "noise"))]
@@ -301,7 +301,7 @@ mod tests {
             relay: libp2p_relay::client::Behaviour,
         }
 
-        let (builder, _bandwidth_sinks) = SwarmBuilder::with_new_identity()
+        let _ = SwarmBuilder::with_new_identity()
             .with_tokio()
             .with_tcp(
                 Default::default(),
@@ -317,8 +317,7 @@ mod tests {
             .unwrap()
             .with_relay_client(libp2p_tls::Config::new, libp2p_yamux::Config::default)
             .unwrap()
-            .with_bandwidth_logging();
-        let _: Swarm<MyBehaviour> = builder
+            .with_bandwidth_logging(&mut libp2p_metrics::Registry::default())
             .with_behaviour(|_key, relay| MyBehaviour { relay })
             .unwrap()
             .build();
@@ -327,16 +326,14 @@ mod tests {
     #[test]
     #[cfg(all(feature = "tokio", feature = "tcp", feature = "tls", feature = "yamux"))]
     fn tcp_bandwidth_logging() -> Result<(), Box<dyn std::error::Error>> {
-        let (builder, _logging) = SwarmBuilder::with_new_identity()
+        let _ = SwarmBuilder::with_new_identity()
             .with_tokio()
             .with_tcp(
                 Default::default(),
                 libp2p_tls::Config::new,
                 libp2p_yamux::Config::default,
             )?
-            .with_bandwidth_logging();
-
-        builder
+            .with_bandwidth_logging(&mut libp2p_metrics::Registry::default())
             .with_behaviour(|_| libp2p_swarm::dummy::Behaviour)
             .unwrap()
             .build();
@@ -347,12 +344,10 @@ mod tests {
     #[test]
     #[cfg(all(feature = "tokio", feature = "quic"))]
     fn quic_bandwidth_logging() -> Result<(), Box<dyn std::error::Error>> {
-        let (builder, _logging) = SwarmBuilder::with_new_identity()
+        let _ = SwarmBuilder::with_new_identity()
             .with_tokio()
             .with_quic()
-            .with_bandwidth_logging();
-
-        builder
+            .with_bandwidth_logging(&mut libp2p_metrics::Registry::default())
             .with_behaviour(|_| libp2p_swarm::dummy::Behaviour)
             .unwrap()
             .build();
@@ -363,12 +358,10 @@ mod tests {
     #[test]
     #[cfg(feature = "tokio")]
     fn other_transport_bandwidth_logging() -> Result<(), Box<dyn std::error::Error>> {
-        let (builder, _logging) = SwarmBuilder::with_new_identity()
+        let _ = SwarmBuilder::with_new_identity()
             .with_tokio()
             .with_other_transport(|_| DummyTransport::<(PeerId, StreamMuxerBox)>::new())?
-            .with_bandwidth_logging();
-
-        builder
+            .with_bandwidth_logging(&mut libp2p_metrics::Registry::default())
             .with_behaviour(|_| libp2p_swarm::dummy::Behaviour)
             .unwrap()
             .build();

libp2p/src/builder/phase/bandwidth_logging.rs

@@ -1,38 +1,40 @@
+use multiaddr::Multiaddr;
+
 use super::*;
-use crate::bandwidth::BandwidthSinks;
+use crate::metrics::bandwidth::{BandwidthSinks, Muxer};
 use crate::transport_ext::TransportExt;
 use crate::SwarmBuilder;
+use std::collections::HashMap;
 use std::marker::PhantomData;
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 
 pub struct BandwidthLoggingPhase<T, R> {
     pub(crate) relay_behaviour: R,
     pub(crate) transport: T,
 }
 
+#[cfg(feature = "metrics")]
 impl<T: AuthenticatedMultiplexedTransport, Provider, R>
     SwarmBuilder<Provider, BandwidthLoggingPhase<T, R>>
 {
     pub fn with_bandwidth_logging(
         self,
-    ) -> (
-        SwarmBuilder<Provider, BehaviourPhase<impl AuthenticatedMultiplexedTransport, R>>,
-        Arc<BandwidthSinks>,
-    ) {
-        let (transport, sinks) = self.phase.transport.with_bandwidth_logging();
-        (
-            SwarmBuilder {
-                phase: BehaviourPhase {
-                    relay_behaviour: self.phase.relay_behaviour,
-                    transport,
-                },
-                keypair: self.keypair,
-                phantom: PhantomData,
+        registry: &mut libp2p_metrics::Registry,
+    ) -> SwarmBuilder<Provider, BehaviourPhase<impl AuthenticatedMultiplexedTransport, R>> {
+        SwarmBuilder {
+            phase: BehaviourPhase {
+                relay_behaviour: self.phase.relay_behaviour,
+                transport: crate::metrics::bandwidth::Transport::new(self.phase.transport, registry),
             },
-            sinks,
-        )
+            keypair: self.keypair,
+            phantom: PhantomData,
+        }
     }
+}
 
+impl<T: AuthenticatedMultiplexedTransport, Provider, R>
+    SwarmBuilder<Provider, BandwidthLoggingPhase<T, R>>
+{
     pub fn without_bandwidth_logging(self) -> SwarmBuilder<Provider, BehaviourPhase<T, R>> {
         SwarmBuilder {
             phase: BehaviourPhase {

libp2p/src/builder/phase/other_transport.rs

@@ -1,6 +1,7 @@
+use std::collections::HashMap;
 use std::convert::Infallible;
 use std::marker::PhantomData;
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 
 use libp2p_core::upgrade::{InboundConnectionUpgrade, OutboundConnectionUpgrade};
 use libp2p_core::Transport;
@@ -143,23 +144,22 @@ impl<T: AuthenticatedMultiplexedTransport, Provider>
             .with_relay_client(security_upgrade, multiplexer_upgrade)
     }
 }
+#[cfg(feature = "metrics")]
 impl<Provider, T: AuthenticatedMultiplexedTransport>
     SwarmBuilder<Provider, OtherTransportPhase<T>>
 {
     pub fn with_bandwidth_logging(
         self,
-    ) -> (
-        SwarmBuilder<
-            Provider,
-            BehaviourPhase<impl AuthenticatedMultiplexedTransport, NoRelayBehaviour>,
-        >,
-        Arc<BandwidthSinks>,
-    ) {
+        registry: &mut libp2p_metrics::Registry,
+    ) -> SwarmBuilder<
+        Provider,
+        BehaviourPhase<impl AuthenticatedMultiplexedTransport, NoRelayBehaviour>,
+    > {
         self.without_any_other_transports()
             .without_dns()
             .without_websocket()
             .without_relay()
-            .with_bandwidth_logging()
+            .with_bandwidth_logging(registry)
     }
 }
 impl<Provider, T: AuthenticatedMultiplexedTransport>

libp2p/src/builder/phase/quic.rs

@@ -1,4 +1,5 @@
 use super::*;
+use crate::bandwidth::BandwidthSinks;
 use crate::SwarmBuilder;
 #[cfg(all(not(target_arch = "wasm32"), feature = "websocket"))]
 use libp2p_core::muxing::StreamMuxer;
@@ -8,7 +9,9 @@
     all(not(target_arch = "wasm32"), feature = "websocket")
 ))]
 use libp2p_core::{InboundUpgrade, Negotiated, OutboundUpgrade, UpgradeInfo};
-use std::{marker::PhantomData, sync::Arc};
+use std::collections::HashMap;
+use std::marker::PhantomData;
+use std::sync::{Arc, RwLock};
 
 pub struct QuicPhase<T> {
     pub(crate) transport: T,
@@ -249,18 +252,16 @@
 impl<Provider, T: AuthenticatedMultiplexedTransport> SwarmBuilder<Provider, QuicPhase<T>> {
     pub fn with_bandwidth_logging(
         self,
-    ) -> (
-        SwarmBuilder<
-            Provider,
-            BehaviourPhase<impl AuthenticatedMultiplexedTransport, NoRelayBehaviour>,
-        >,
-        Arc<crate::bandwidth::BandwidthSinks>,
-    ) {
+        registry: &mut libp2p_metrics::Registry,
+    ) -> SwarmBuilder<
+        Provider,
+        BehaviourPhase<impl AuthenticatedMultiplexedTransport, NoRelayBehaviour>,
+    > {
         self.without_quic()
             .without_any_other_transports()
             .without_dns()
             .without_websocket()
             .without_relay()
-            .with_bandwidth_logging()
+            .with_bandwidth_logging(registry)
     }
 }

libp2p/src/transport_ext.rs

@@ -29,7 +29,11 @@
     Transport,
 };
 use libp2p_identity::PeerId;
-use std::sync::Arc;
+use multiaddr::Multiaddr;
+use std::{
+    collections::HashMap,
+    sync::{Arc, RwLock},
+};
 
 /// Trait automatically implemented on all objects that implement `Transport`. Provides some
 /// additional utilities.
@@ -66,7 +70,12 @@
     ///
     /// let (transport, sinks) = transport.with_bandwidth_logging();
     /// ```
-    fn with_bandwidth_logging<S>(self) -> (Boxed<(PeerId, StreamMuxerBox)>, Arc<BandwidthSinks>)
+    fn with_bandwidth_logging<S>(
+        self,
+    ) -> (
+        Boxed<(PeerId, StreamMuxerBox)>,
+        Arc<RwLock<HashMap<String, Arc<BandwidthSinks>>>>,
+    )
     where
         Self: Sized + Send + Unpin + 'static,
         Self::Dial: Send + 'static,
@@ -74,16 +83,35 @@
         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: Arc<RwLock<HashMap<_, Arc<BandwidthSinks>>>> = Arc::new(RwLock::new(HashMap::new()));
         let sinks_copy = sinks.clone();
-        let transport = Transport::map(self, |output, _| {
+        let transport = Transport::map(self, move |output, connected_point| {
+            fn as_string(ma: &Multiaddr) -> String {
+                let len = ma
+                    .protocol_stack()
+                    .fold(0, |acc, proto| acc + proto.len() + 1);
+                let mut protocols = String::with_capacity(len);
+                for proto_tag in ma.protocol_stack() {
+                    protocols.push('/');
+                    protocols.push_str(proto_tag);
+                }
+                protocols
+            }
+
+            let sink = sinks_copy
+                .write()
+                .expect("todo")
+                .entry(as_string(connected_point.get_remote_address()))
+                .or_default()
+                .clone();
+
             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, sink)),
             )
         })
         .boxed();