From 4fe343087dc93922b6938550c026dbec9b2427a0 Mon Sep 17 00:00:00 2001
From: Sylvain Corlay <sylvain.corlay@gmail.com>
Date: Fri, 11 Dec 2020 17:55:05 +0100
Subject: [PATCH] Nudge kernel with kernel info requests until IOPub responds

---
 notebook/services/kernels/handlers.py | 29 ++++++++++++++++-----------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/notebook/services/kernels/handlers.py b/notebook/services/kernels/handlers.py
index cd80aa9619..7f93101dcf 100644
--- a/notebook/services/kernels/handlers.py
+++ b/notebook/services/kernels/handlers.py
@@ -144,24 +144,26 @@ def finish():
             shell_channel.stop_on_recv()
 
         def on_shell_reply(msg):
-            self.log.debug("Nudge: shell info reply received: %s", self.kernel_id)
             if not info_future.done():
+                self.log.debug("Nudge: shell info reply received: %s", self.kernel_id)
+                shell_channel.stop_on_recv()
                 self.log.debug("Nudge: resolving shell future")
                 info_future.set_result(msg)
-            if iopub_future.done():
-                finish()
-                self.log.debug("Setting results")
-                future.set_result(info_future.result())
+                if iopub_future.done():
+                    finish()
+                    self.log.debug("Nudge: resolving main future in shell handler")
+                    future.set_result(info_future.result())
 
         def on_iopub(msg):
-            self.log.debug("Nudge: First IOPub received: %s", self.kernel_id)
             if not iopub_future.done():
+                self.log.debug("Nudge: first IOPub received: %s", self.kernel_id)
+                iopub_channel.stop_on_recv()
                 self.log.debug("Nudge: resolving iopub future")
                 iopub_future.set_result(None)
-            if info_future.done():
-                finish()
-                self.log.debug("Setting results")
-                future.set_result(info_future.result())
+                if info_future.done():
+                    finish()
+                    self.log.debug("Nudge: resolving main future in iopub handler")
+                    future.set_result(info_future.result())
 
         def on_timeout():
             self.log.warning("Nudge: Timeout waiting for kernel_info_reply: %s", self.kernel_id)
@@ -346,8 +348,11 @@ def open(self, kernel_id):
         km.add_restart_callback(self.kernel_id, self.on_kernel_restarted)
         km.add_restart_callback(self.kernel_id, self.on_restart_failed, 'dead')
 
-        for channel, stream in self.channels.items():
-            stream.on_recv_stream(self._on_zmq_reply)
+        def subscribe(value):
+            for channel, stream in self.channels.items():
+                stream.on_recv_stream(self._on_zmq_reply)
+
+        connected.add_done_callback(subscribe)
 
         return connected