From 4ecff68d43a330def9fd55c69cd12dde19ec746f Mon Sep 17 00:00:00 2001
From: magine <magine233@gmail.com>
Date: Tue, 13 Aug 2024 08:08:49 +0800
Subject: [PATCH] config: add a keep_alive_timeout configuration option

---
 src/config.rs | 16 +++++++++++++++-
 src/lib.rs    | 16 ++++++++--------
 2 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 34d67ca0..dd69cd90 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -45,7 +45,7 @@ use crate::transport::websocket::config::Config as WebSocketConfig;
 
 use multiaddr::Multiaddr;
 
-use std::{collections::HashMap, sync::Arc};
+use std::{collections::HashMap, sync::Arc, time::Duration};
 
 /// Connection role.
 #[derive(Debug, Copy, Clone)]
@@ -121,6 +121,9 @@ pub struct ConfigBuilder {
 
     /// Connection limits config.
     connection_limits: ConnectionLimitsConfig,
+
+    /// Close the connection if no substreams are open within this time frame.
+    keep_alive_timeout: Duration,
 }
 
 impl Default for ConfigBuilder {
@@ -153,6 +156,7 @@ impl ConfigBuilder {
             request_response_protocols: HashMap::new(),
             known_addresses: Vec::new(),
             connection_limits: ConnectionLimitsConfig::default(),
+            keep_alive_timeout: Duration::from_secs(5),
         }
     }
 
@@ -268,6 +272,12 @@ impl ConfigBuilder {
         self
     }
 
+    /// Set keep alive timeout for connections.
+    pub fn set_keep_alive_timeout(mut self, timeout: Duration) -> Self {
+        self.keep_alive_timeout = timeout;
+        self
+    }
+
     /// Build [`Litep2pConfig`].
     pub fn build(mut self) -> Litep2pConfig {
         let keypair = match self.keypair {
@@ -296,6 +306,7 @@ impl ConfigBuilder {
             request_response_protocols: self.request_response_protocols,
             known_addresses: self.known_addresses,
             connection_limits: self.connection_limits,
+            keep_alive_timeout: self.keep_alive_timeout,
         }
     }
 }
@@ -355,4 +366,7 @@ pub struct Litep2pConfig {
 
     /// Connection limits config.
     pub(crate) connection_limits: ConnectionLimitsConfig,
+
+    /// Close the connection if no substreams are open within this time frame.
+    pub(crate) keep_alive_timeout: Duration,
 }
diff --git a/src/lib.rs b/src/lib.rs
index 1523bcb4..6222c241 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -55,7 +55,7 @@ use multihash::Multihash;
 use transport::Endpoint;
 use types::ConnectionId;
 
-use std::{collections::HashSet, sync::Arc, time::Duration};
+use std::{collections::HashSet, sync::Arc};
 
 pub use bandwidth::BandwidthSink;
 pub use error::Error;
@@ -171,7 +171,7 @@ impl Litep2p {
                 protocol,
                 config.fallback_names.clone(),
                 config.codec,
-                Duration::from_secs(5),
+                litep2p_config.keep_alive_timeout,
             );
             let executor = Arc::clone(&litep2p_config.executor);
             litep2p_config.executor.run(Box::pin(async move {
@@ -191,7 +191,7 @@ impl Litep2p {
                 protocol,
                 config.fallback_names.clone(),
                 config.codec,
-                config.timeout,
+                litep2p_config.keep_alive_timeout,
             );
             litep2p_config.executor.run(Box::pin(async move {
                 RequestResponseProtocol::new(service, config).run().await
@@ -206,7 +206,7 @@ impl Litep2p {
                 protocol_name,
                 Vec::new(),
                 protocol.codec(),
-                Duration::from_secs(5),
+                litep2p_config.keep_alive_timeout,
             );
             litep2p_config.executor.run(Box::pin(async move {
                 let _ = protocol.run(service).await;
@@ -225,7 +225,7 @@ impl Litep2p {
                 ping_config.protocol.clone(),
                 Vec::new(),
                 ping_config.codec,
-                Duration::from_secs(5),
+                litep2p_config.keep_alive_timeout,
             );
             litep2p_config.executor.run(Box::pin(async move {
                 Ping::new(service, ping_config).run().await
@@ -248,7 +248,7 @@ impl Litep2p {
                 main_protocol.clone(),
                 fallback_names,
                 kademlia_config.codec,
-                Duration::from_secs(5),
+                litep2p_config.keep_alive_timeout,
             );
             litep2p_config.executor.run(Box::pin(async move {
                 let _ = Kademlia::new(service, kademlia_config).run().await;
@@ -269,7 +269,7 @@ impl Litep2p {
                     identify_config.protocol.clone(),
                     Vec::new(),
                     identify_config.codec,
-                    Duration::from_secs(5),
+                    litep2p_config.keep_alive_timeout,
                 );
                 identify_config.public = Some(litep2p_config.keypair.public().into());
 
@@ -289,7 +289,7 @@ impl Litep2p {
                 bitswap_config.protocol.clone(),
                 Vec::new(),
                 bitswap_config.codec,
-                Duration::from_secs(5),
+                litep2p_config.keep_alive_timeout,
             );
             litep2p_config.executor.run(Box::pin(async move {
                 Bitswap::new(service, bitswap_config).run().await