Move bandwidth logging to misc/metrics

libp2p · Oct 27, 2023 · e1b4497 · e1b4497
1 parent 9c940fa
commit e1b4497
Show file tree

Hide file tree

Showing 8 changed files with 429 additions and 51 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/libp2p/src/builder.rs b/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();

diff --git a/libp2p/src/builder/phase/bandwidth_logging.rs b/libp2p/src/builder/phase/bandwidth_logging.rs
@@ -1,5 +1,7 @@
+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;
@@ -11,29 +13,28 @@ pub struct BandwidthLoggingPhase<T, 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<RwLock<HashMap<String, 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 {

diff --git a/libp2p/src/builder/phase/other_transport.rs b/libp2p/src/builder/phase/other_transport.rs
@@ -144,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<RwLock<HashMap<String, 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>

diff --git a/libp2p/src/builder/phase/quic.rs b/libp2p/src/builder/phase/quic.rs
@@ -252,18 +252,16 @@ impl_quic_phase_with_websocket!(
 impl<Provider, T: AuthenticatedMultiplexedTransport> SwarmBuilder<Provider, QuicPhase<T>> {
     pub fn with_bandwidth_logging(
         self,
-    ) -> (
-        SwarmBuilder<
-            Provider,
-            BehaviourPhase<impl AuthenticatedMultiplexedTransport, NoRelayBehaviour>,
-        >,
-        Arc<RwLock<HashMap<String, Arc<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)
     }
 }
diff --git a/misc/metrics/Cargo.toml b/misc/metrics/Cargo.toml
@@ -19,6 +19,7 @@ ping = ["libp2p-ping"]
 relay = ["libp2p-relay"]
 
 [dependencies]
+futures = "0.3.26"
 instant = "0.1.12"
 libp2p-core = { workspace = true }
 libp2p-dcutr =  { workspace = true, optional = true }
@@ -29,6 +30,7 @@ libp2p-kad = { workspace = true, optional = true }
 libp2p-ping = { workspace = true, optional = true }
 libp2p-relay =  { workspace = true, optional = true }
 libp2p-swarm = { workspace = true }
+pin-project = "1.0.0"
 prometheus-client = { workspace = true }
 
 [dev-dependencies]