From 3cabdc4d1bdeadd69269589b81ce3106462e253c Mon Sep 17 00:00:00 2001
From: Kevin Bates <kbates4@gmail.com>
Date: Mon, 25 Feb 2019 11:17:49 -0800
Subject: [PATCH] add async support for kernel shutdown

---
 jupyter_client/manager.py            | 6 ++++--
 jupyter_client/multikernelmanager.py | 8 ++++++--
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/jupyter_client/manager.py b/jupyter_client/manager.py
index de5c58b31..362d209f5 100644
--- a/jupyter_client/manager.py
+++ b/jupyter_client/manager.py
@@ -260,6 +260,7 @@ def request_shutdown(self, restart=False):
         self._connect_control_socket()
         self.session.send(self._control_socket, msg)
 
+    @gen.coroutine
     def finish_shutdown(self, waittime=None, pollinterval=0.1):
         """Wait for kernel shutdown, then kill process if it doesn't shutdown.
 
@@ -270,7 +271,7 @@ def finish_shutdown(self, waittime=None, pollinterval=0.1):
             waittime = max(self.shutdown_wait_time, 0)
         for i in range(int(waittime/pollinterval)):
             if self.is_alive():
-                time.sleep(pollinterval)
+                yield gen.sleep(pollinterval)
             else:
                 # If there's still a proc, wait and clear
                 if self.has_kernel:
@@ -292,6 +293,7 @@ def cleanup(self, connection_file=True):
         self._close_control_socket()
         self.session.parent = None
 
+    @gen.coroutine
     def shutdown_kernel(self, now=False, restart=False):
         """Attempts to stop the kernel process cleanly.
 
@@ -320,7 +322,7 @@ def shutdown_kernel(self, now=False, restart=False):
             # Don't send any additional kernel kill messages immediately, to give
             # the kernel a chance to properly execute shutdown actions. Wait for at
             # most 1s, checking every 0.1s.
-            self.finish_shutdown()
+            yield gen.maybe_future(self.finish_shutdown())
 
         self.cleanup(connection_file=not restart)
 
diff --git a/jupyter_client/multikernelmanager.py b/jupyter_client/multikernelmanager.py
index fb229ea69..d9b670985 100644
--- a/jupyter_client/multikernelmanager.py
+++ b/jupyter_client/multikernelmanager.py
@@ -112,7 +112,7 @@ def start_kernel(self, kernel_name=None, **kwargs):
         self._kernels[kernel_id] = km
         raise gen.Return(kernel_id)
 
-    @kernel_method
+    @gen.coroutine
     def shutdown_kernel(self, kernel_id, now=False, restart=False):
         """Shutdown a kernel by its kernel uuid.
 
@@ -125,6 +125,8 @@ def shutdown_kernel(self, kernel_id, now=False, restart=False):
         restart : bool
             Will the kernel be restarted?
         """
+        km = self.get_kernel(kernel_id)
+        yield gen.maybe_future(km.shutdown_kernel(now, restart))
         self.log.info("Kernel shutdown: %s" % kernel_id)
         self.remove_kernel(kernel_id)
 
@@ -132,10 +134,12 @@ def shutdown_kernel(self, kernel_id, now=False, restart=False):
     def request_shutdown(self, kernel_id, restart=False):
         """Ask a kernel to shut down by its kernel uuid"""
 
-    @kernel_method
+    @gen.coroutine
     def finish_shutdown(self, kernel_id, waittime=None, pollinterval=0.1):
         """Wait for a kernel to finish shutting down, and kill it if it doesn't
         """
+        km = self.get_kernel(kernel_id)
+        yield gen.maybe_future(km.finish_shutdown(waittime, pollinterval))
         self.log.info("Kernel shutdown: %s" % kernel_id)
 
     @kernel_method