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/phase/bandwidth_logging.rs

@@ -2,8 +2,9 @@ use super::*;
 use crate::bandwidth::BandwidthSinks;
 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,
@@ -17,7 +18,7 @@ impl<T: AuthenticatedMultiplexedTransport, Provider, R>
         self,
     ) -> (
         SwarmBuilder<Provider, BehaviourPhase<impl AuthenticatedMultiplexedTransport, R>>,
-        Arc<BandwidthSinks>,
+        Arc<RwLock<HashMap<String, Arc<BandwidthSinks>>>>,
     ) {
         let (transport, sinks) = self.phase.transport.with_bandwidth_logging();
         (

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;
@@ -153,7 +154,7 @@ impl<Provider, T: AuthenticatedMultiplexedTransport>
             Provider,
             BehaviourPhase<impl AuthenticatedMultiplexedTransport, NoRelayBehaviour>,
         >,
-        Arc<BandwidthSinks>,
+        Arc<RwLock<HashMap<String, Arc<BandwidthSinks>>>>,
     ) {
         self.without_any_other_transports()
             .without_dns()

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 @@ use libp2p_core::upgrade::{InboundConnectionUpgrade, OutboundConnectionUpgrade};
     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,
@@ -254,7 +257,7 @@ impl<Provider, T: AuthenticatedMultiplexedTransport> SwarmBuilder<Provider, Quic
             Provider,
             BehaviourPhase<impl AuthenticatedMultiplexedTransport, NoRelayBehaviour>,
         >,
-        Arc<crate::bandwidth::BandwidthSinks>,
+        Arc<RwLock<HashMap<String, Arc<BandwidthSinks>>>>,
     ) {
         self.without_quic()
             .without_any_other_transports()

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();

misc/metrics/src/lib.rs

@@ -39,9 +39,11 @@
 mod ping;
 mod protocol_stack;
 #[cfg(feature = "relay")]
 mod relay;
 mod swarm;
 
+use std::{sync::{Arc, RwLock}, collections::HashMap};
+
 use prometheus_client::registry::Registry;
 
 /// Set of Swarm and protocol metrics derived from emitted events.