From d7ade663e8cc42a38ef1976f65f1474f77d2fdf7 Mon Sep 17 00:00:00 2001
From: Kyle Cutler <67761731+kycutler@users.noreply.github.com>
Date: Mon, 28 Feb 2022 14:04:36 -0800
Subject: [PATCH] nudge both the shell and control channels (#6290)

---
 notebook/services/kernels/handlers.py | 22 ++++++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

diff --git a/notebook/services/kernels/handlers.py b/notebook/services/kernels/handlers.py
index 6040de1d48..a2d4ddbab9 100644
--- a/notebook/services/kernels/handlers.py
+++ b/notebook/services/kernels/handlers.py
@@ -138,7 +138,7 @@ def nudge(self):
         """Nudge the zmq connections with kernel_info_requests
 
         Returns a Future that will resolve when we have received
-        a shell reply and at least one iopub message,
+        a shell or control reply and at least one iopub message,
         ensuring that zmq subscriptions are established,
         sockets are fully connected, and kernel is responsive.
 
@@ -157,10 +157,12 @@ def nudge(self):
             f = Future()
             f.set_result(None)
             return f
-
         # Use a transient shell channel to prevent leaking
         # shell responses to the front-end.
         shell_channel = kernel.connect_shell()
+        # Use a transient control channel to prevent leaking
+        # control responses to the front-end.
+        control_channel = kernel.connect_control()
         # The IOPub used by the client, whose subscriptions we are verifying.
         iopub_channel = self.channels["iopub"]
 
@@ -183,6 +185,8 @@ def cleanup(f=None):
             iopub_channel.stop_on_recv()
             if not shell_channel.closed():
                 shell_channel.close()
+            if not control_channel.closed():
+                control_channel.close()
 
         # trigger cleanup when both message futures are resolved
         both_done.add_done_callback(cleanup)
@@ -193,6 +197,12 @@ def on_shell_reply(msg):
                 self.log.debug("Nudge: resolving shell future: %s", self.kernel_id)
                 info_future.set_result(None)
 
+        def on_control_reply(msg):
+            self.log.debug("Nudge: control info reply received: %s", self.kernel_id)
+            if not info_future.done():
+                self.log.debug("Nudge: resolving control future: %s", self.kernel_id)
+                info_future.set_result(None)
+
         def on_iopub(msg):
             self.log.debug("Nudge: IOPub received: %s", self.kernel_id)
             if not iopub_future.done():
@@ -202,6 +212,7 @@ def on_iopub(msg):
 
         iopub_channel.on_recv(on_iopub)
         shell_channel.on_recv(on_shell_reply)
+        control_channel.on_recv(on_control_reply)
         loop = IOLoop.current()
 
         # Nudge the kernel with kernel info requests until we get an IOPub message
@@ -227,6 +238,12 @@ def nudge(count):
 
             # check for closed zmq socket
             if shell_channel.closed():
+                self.log.debug("Nudge: cancelling on closed zmq socket: %s", self.kernel_id)
+                finish()
+                return
+
+            # check for closed zmq socket
+            if control_channel.closed():
                 self.log.debug(
                     "Nudge: cancelling on closed zmq socket: %s", self.kernel_id
                 )
@@ -237,6 +254,7 @@ def nudge(count):
                 log = self.log.warning if count % 10 == 0 else self.log.debug
                 log("Nudge: attempt %s on kernel %s" % (count, self.kernel_id))
                 self.session.send(shell_channel, "kernel_info_request")
+                self.session.send(control_channel, "kernel_info_request")
                 nonlocal nudge_handle
                 nudge_handle = loop.call_later(0.5, nudge, count)