Allow configuration of outbound substream in OneShotHandler.

libp2p · Mar 25, 2020 · c2e9397 · c2e9397
1 parent 28ea62d
commit c2e9397
Showing 1 changed file with 14 additions and 3 deletions.
diff --git a/swarm/src/protocols_handler/one_shot.rs b/swarm/src/protocols_handler/one_shot.rs
@@ -56,6 +56,8 @@ where
     keep_alive: KeepAlive,
     /// After the given duration has elapsed, an inactive connection will shutdown.
     inactive_timeout: Duration,
+    /// Timeout duration for each newly opened outbound substream.
+    substream_timeout: Duration,
 }
 
 impl<TInProto, TOutProto, TOutEvent>
@@ -67,7 +69,8 @@ where
     #[inline]
     pub fn new(
         listen_protocol: SubstreamProtocol<TInProto>,
-        inactive_timeout: Duration
+        inactive_timeout: Duration,
+        substream_timeout: Duration
     ) -> Self {
         OneShotHandler {
             listen_protocol,
@@ -78,6 +81,7 @@ where
             max_dial_negotiated: 8,
             keep_alive: KeepAlive::Yes,
             inactive_timeout,
+            substream_timeout,
         }
     }
 
@@ -121,7 +125,13 @@ where
 {
     #[inline]
     fn default() -> Self {
-        OneShotHandler::new(SubstreamProtocol::new(Default::default()), Duration::from_secs(10))
+        let inactive_timeout = Duration::from_secs(10);
+        let substream_timeout = Duration::from_secs(10);
+        OneShotHandler::new(
+            SubstreamProtocol::new(Default::default()),
+            inactive_timeout,
+            substream_timeout,
+        )
     }
 }
 
@@ -224,7 +234,8 @@ where
                 self.dial_negotiated += 1;
                 return Poll::Ready(
                     ProtocolsHandlerEvent::OutboundSubstreamRequest {
-                        protocol: SubstreamProtocol::new(self.dial_queue.remove(0)),
+                        protocol: SubstreamProtocol::new(self.dial_queue.remove(0))
+                            .with_timeout(self.substream_timeout),
                         info: (),
                     },
                 );