fix(libp2p): add shortcut with_tcp(...).with_bandwidth_logging()

Add the shortcut method `.with_bandwidth_logging` to `SwarmBuilder<_, QuicPhase<_>>`, thus allowing `with_tcp(...).with_bandwidth_logging()`. Pull-Request: #4626.
libp2p · Oct 13, 2023 · 74c087d · 74c087d
1 parent 56cb08a
commit 74c087d
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 4 deletions.
diff --git a/libp2p/src/builder.rs b/libp2p/src/builder.rs
@@ -66,6 +66,8 @@ pub struct SwarmBuilder<Provider, Phase> {
 #[cfg(test)]
 mod tests {
     use crate::SwarmBuilder;
+    use libp2p_core::{muxing::StreamMuxerBox, transport::dummy::DummyTransport};
+    use libp2p_identity::PeerId;
     use libp2p_swarm::{NetworkBehaviour, Swarm};
 
     #[test]
@@ -228,9 +230,6 @@ mod tests {
     #[test]
     #[cfg(feature = "tokio")]
     fn other_transport() -> Result<(), Box<dyn std::error::Error>> {
-        use libp2p_core::{muxing::StreamMuxerBox, transport::dummy::DummyTransport};
-        use libp2p_identity::PeerId;
-
         let _ = SwarmBuilder::with_new_identity()
             .with_tokio()
             // Closure can either return a Transport directly.
@@ -319,4 +318,56 @@ mod tests {
             .unwrap()
             .build();
     }
+
+    #[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()
+            .with_tokio()
+            .with_tcp(
+                Default::default(),
+                libp2p_tls::Config::new,
+                libp2p_yamux::Config::default,
+            )?
+            .with_bandwidth_logging();
+
+        builder
+            .with_behaviour(|_| libp2p_swarm::dummy::Behaviour)
+            .unwrap()
+            .build();
+
+        Ok(())
+    }
+
+    #[test]
+    #[cfg(all(feature = "tokio", feature = "quic"))]
+    fn quic_bandwidth_logging() -> Result<(), Box<dyn std::error::Error>> {
+        let (builder, _logging) = SwarmBuilder::with_new_identity()
+            .with_tokio()
+            .with_quic()
+            .with_bandwidth_logging();
+
+        builder
+            .with_behaviour(|_| libp2p_swarm::dummy::Behaviour)
+            .unwrap()
+            .build();
+
+        Ok(())
+    }
+
+    #[test]
+    #[cfg(feature = "tokio")]
+    fn other_transport_bandwidth_logging() -> Result<(), Box<dyn std::error::Error>> {
+        let (builder, _logging) = SwarmBuilder::with_new_identity()
+            .with_tokio()
+            .with_other_transport(|_| DummyTransport::<(PeerId, StreamMuxerBox)>::new())?
+            .with_bandwidth_logging();
+
+        builder
+            .with_behaviour(|_| libp2p_swarm::dummy::Behaviour)
+            .unwrap()
+            .build();
+
+        Ok(())
+    }
 }
diff --git a/libp2p/src/builder/phase/quic.rs b/libp2p/src/builder/phase/quic.rs
@@ -7,7 +7,7 @@ use libp2p_core::muxing::StreamMuxer;
     all(not(target_arch = "wasm32"), feature = "websocket")
 ))]
 use libp2p_core::{InboundUpgrade, Negotiated, OutboundUpgrade, UpgradeInfo};
-use std::marker::PhantomData;
+use std::{marker::PhantomData, sync::Arc};
 
 pub struct QuicPhase<T> {
     pub(crate) transport: T,
@@ -237,3 +237,21 @@ impl_quic_phase_with_websocket!(
     super::provider::Tokio,
     rw_stream_sink::RwStreamSink<libp2p_websocket::BytesConnection<libp2p_tcp::tokio::TcpStream>>
 );
+impl<Provider, T: AuthenticatedMultiplexedTransport> SwarmBuilder<Provider, QuicPhase<T>> {
+    pub fn with_bandwidth_logging(
+        self,
+    ) -> (
+        SwarmBuilder<
+            Provider,
+            BehaviourPhase<impl AuthenticatedMultiplexedTransport, NoRelayBehaviour>,
+        >,
+        Arc<crate::bandwidth::BandwidthSinks>,
+    ) {
+        self.without_quic()
+            .without_any_other_transports()
+            .without_dns()
+            .without_relay()
+            .without_websocket()
+            .with_bandwidth_logging()
+    }
+}