From 0ed684ee30cb0d6ccf9f6869a89732f179b908cf Mon Sep 17 00:00:00 2001
From: Pierre Krieger <pierre.krieger1708@gmail.com>
Date: Mon, 20 Jan 2020 13:55:25 +0100
Subject: [PATCH] Don't set KeepAlive to Until all the time (#1398)

* Don't set KeepAlive to Until all the time

* Address concerns
---
 protocols/kad/src/handler.rs | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/protocols/kad/src/handler.rs b/protocols/kad/src/handler.rs
index 59bf26c7..3f969046 100644
--- a/protocols/kad/src/handler.rs
+++ b/protocols/kad/src/handler.rs
@@ -403,7 +403,7 @@ where
             allow_listening,
             next_connec_unique_id: UniqueConnecId(0),
             substreams: Vec::new(),
-            keep_alive: KeepAlive::Yes,
+            keep_alive: KeepAlive::Until(Instant::now() + Duration::from_secs(10)),
         }
     }
 
@@ -645,6 +645,10 @@ where
     ) -> Poll<
         ProtocolsHandlerEvent<Self::OutboundProtocol, Self::OutboundOpenInfo, Self::OutEvent, Self::Error>,
     > {
+        if self.substreams.is_empty() {
+            return Poll::Pending;
+        }
+
         // We remove each element from `substreams` one by one and add them back.
         for n in (0..self.substreams.len()).rev() {
             let mut substream = self.substreams.swap_remove(n);
@@ -656,6 +660,9 @@ where
                         return Poll::Ready(event);
                     }
                     (None, Some(event), _) => {
+                        if self.substreams.is_empty() {
+                            self.keep_alive = KeepAlive::Until(Instant::now() + Duration::from_secs(10));
+                        }
                         return Poll::Ready(event);
                     }
                     (Some(new_state), None, false) => {
@@ -674,6 +681,7 @@ where
         }
 
         if self.substreams.is_empty() {
+            // We destroyed all substreams in this function.
             self.keep_alive = KeepAlive::Until(Instant::now() + Duration::from_secs(10));
         } else {
             self.keep_alive = KeepAlive::Yes;