Merge pull request #2031 from RolandSherwin/metrics_bw

feat(network): expose bandwidth per transport protocol
maidsafe · Aug 8, 2024 · 7ee0ff6 · 7ee0ff6
2 parents 9f5d728 + d527e0b
commit 7ee0ff6
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 30 deletions.
diff --git a/sn_networking/src/driver.rs b/sn_networking/src/driver.rs
@@ -430,9 +430,42 @@ impl NetworkBuilder {
             PrettyPrintKBucketKey(NetworkAddress::from_peer(peer_id).as_kbucket_key())
         );
 
+        #[cfg(feature = "open-metrics")]
+        let mut metrics_registry = self.metrics_registry.unwrap_or_default();
+
+        // ==== Transport ====
+        #[cfg(feature = "open-metrics")]
+        let main_transport = transport::build_transport(&self.keypair, &mut metrics_registry);
+        #[cfg(not(feature = "open-metrics"))]
+        let main_transport = transport::build_transport(&self.keypair);
+        let transport = if !self.local {
+            debug!("Preventing non-global dials");
+            // Wrap upper in a transport that prevents dialing local addresses.
+            libp2p::core::transport::global_only::Transport::new(main_transport).boxed()
+        } else {
+            main_transport
+        };
+
+        let (relay_transport, relay_behaviour) =
+            libp2p::relay::client::new(self.keypair.public().to_peer_id());
+        let relay_transport = relay_transport
+            .upgrade(libp2p::core::upgrade::Version::V1Lazy)
+            .authenticate(
+                libp2p::noise::Config::new(&self.keypair)
+                    .expect("Signing libp2p-noise static DH keypair failed."),
+            )
+            .multiplex(libp2p::yamux::Config::default())
+            .or_transport(transport);
+
+        let transport = relay_transport
+            .map(|either_output, _| match either_output {
+                Either::Left((peer_id, muxer)) => (peer_id, StreamMuxerBox::new(muxer)),
+                Either::Right((peer_id, muxer)) => (peer_id, StreamMuxerBox::new(muxer)),
+            })
+            .boxed();
+
         #[cfg(feature = "open-metrics")]
         let network_metrics = if let Some(port) = self.metrics_server_port {
-            let mut metrics_registry = self.metrics_registry.unwrap_or_default();
             let metrics = NetworkMetrics::new(&mut metrics_registry);
             run_metrics_server(metrics_registry, port);
             Some(metrics)
@@ -519,16 +552,6 @@ impl NetworkBuilder {
             libp2p::identify::Behaviour::new(cfg)
         };
 
-        let main_transport = transport::build_transport(&self.keypair);
-
-        let transport = if !self.local {
-            debug!("Preventing non-global dials");
-            // Wrap upper in a transport that prevents dialing local addresses.
-            libp2p::core::transport::global_only::Transport::new(main_transport).boxed()
-        } else {
-            main_transport
-        };
-
         #[cfg(feature = "upnp")]
         let upnp = if !self.local && !is_client && upnp {
             debug!("Enabling UPnP port opening behavior");
@@ -538,24 +561,6 @@ impl NetworkBuilder {
         }
         .into(); // Into `Toggle<T>`
 
-        let (relay_transport, relay_behaviour) =
-            libp2p::relay::client::new(self.keypair.public().to_peer_id());
-        let relay_transport = relay_transport
-            .upgrade(libp2p::core::upgrade::Version::V1Lazy)
-            .authenticate(
-                libp2p::noise::Config::new(&self.keypair)
-                    .expect("Signing libp2p-noise static DH keypair failed."),
-            )
-            .multiplex(libp2p::yamux::Config::default())
-            .or_transport(transport);
-
-        let transport = relay_transport
-            .map(|either_output, _| match either_output {
-                Either::Left((peer_id, muxer)) => (peer_id, StreamMuxerBox::new(muxer)),
-                Either::Right((peer_id, muxer)) => (peer_id, StreamMuxerBox::new(muxer)),
-            })
-            .boxed();
-
         let relay_server = {
             let relay_server_cfg = relay::Config {
                 max_reservations: 128,             // Amount of peers we are relaying for

diff --git a/sn_networking/src/transport/other.rs b/sn_networking/src/transport/other.rs
@@ -1,5 +1,7 @@
 #[cfg(feature = "websockets")]
 use futures::future::Either;
+#[cfg(feature = "open-metrics")]
+use libp2p::metrics::Registry;
 #[cfg(feature = "websockets")]
 use libp2p::{core::upgrade, noise, yamux};
 use libp2p::{
@@ -8,8 +10,13 @@ use libp2p::{
     PeerId, Transport as _,
 };
 
-pub(crate) fn build_transport(keypair: &Keypair) -> transport::Boxed<(PeerId, StreamMuxerBox)> {
+pub(crate) fn build_transport(
+    keypair: &Keypair,
+    #[cfg(feature = "open-metrics")] registry: &mut Registry,
+) -> transport::Boxed<(PeerId, StreamMuxerBox)> {
     let trans = generate_quic_transport(keypair);
+    #[cfg(feature = "open-metrics")]
+    let trans = libp2p::metrics::BandwidthTransport::new(trans, registry);
 
     #[cfg(feature = "websockets")]
     // Using a closure here due to the complex return type