From 0c5a3b6d04891023f39159eae9f5e783d07d1b89 Mon Sep 17 00:00:00 2001
From: Min RK <benjaminrk@gmail.com>
Date: Wed, 14 Oct 2015 14:58:26 +0200
Subject: [PATCH 1/2] wait for kernel to restart

in API request to restart kernels
by waiting for the kernel to respond to a kernel_info request.
---
 notebook/services/kernels/handlers.py      | 14 +++++--
 notebook/services/kernels/kernelmanager.py | 45 +++++++++++++++++++++-
 2 files changed, 54 insertions(+), 5 deletions(-)

diff --git a/notebook/services/kernels/handlers.py b/notebook/services/kernels/handlers.py
index 3db22f947..19bc20315 100644
--- a/notebook/services/kernels/handlers.py
+++ b/notebook/services/kernels/handlers.py
@@ -72,16 +72,22 @@ class KernelActionHandler(APIHandler):
 
     @web.authenticated
     @json_errors
+    @gen.coroutine
     def post(self, kernel_id, action):
         km = self.kernel_manager
         if action == 'interrupt':
             km.interrupt_kernel(kernel_id)
             self.set_status(204)
         if action == 'restart':
-            km.restart_kernel(kernel_id)
-            model = km.kernel_model(kernel_id)
-            self.set_header('Location', '{0}api/kernels/{1}'.format(self.base_url, kernel_id))
-            self.write(json.dumps(model))
+
+            try:
+                yield gen.maybe_future(km.restart_kernel(kernel_id))
+            except Exception as e:
+                self.log.error("Exception restarting kernel", exc_info=True)
+                self.set_status(500)
+            else:
+                model = km.kernel_model(kernel_id)
+                self.write(json.dumps(model))
         self.finish()
 
 
diff --git a/notebook/services/kernels/kernelmanager.py b/notebook/services/kernels/kernelmanager.py
index 589888c07..d08d2ba8d 100644
--- a/notebook/services/kernels/kernelmanager.py
+++ b/notebook/services/kernels/kernelmanager.py
@@ -9,7 +9,9 @@
 
 import os
 
-from tornado import web
+from tornado import gen, web
+from tornado.concurrent import Future
+from tornado.ioloop import IOLoop
 
 from jupyter_client.multikernelmanager import MultiKernelManager
 from traitlets import List, Unicode, TraitError
@@ -99,6 +101,47 @@ class MappingKernelManager(MultiKernelManager):
         self._check_kernel_id(kernel_id)
         super(MappingKernelManager, self).shutdown_kernel(kernel_id, now=now)
 
+    def restart_kernel(self, kernel_id):
+        """Restart a kernel by kernel_id"""
+        self._check_kernel_id(kernel_id)
+        super(MappingKernelManager, self).restart_kernel(kernel_id)
+        kernel = self.get_kernel(kernel_id)
+        # return a Future that will resolve when the kernel has successfully restarted
+        channel = kernel.connect_shell()
+        future = Future()
+        
+        def finish():
+            """Common cleanup when restart finishes/fails for any reason."""
+            if not channel.closed:
+                channel.close()
+            loop.remove_timeout(timeout)
+            kernel.remove_restart_callback(on_restart_failed, 'dead')
+        
+        def on_reply(msg):
+            self.log.debug("Kernel info reply received: %s", kernel_id)
+            finish()
+            if not future.done():
+                future.set_result(msg)
+            
+        def on_timeout():
+            self.log.warn("Timeout waiting for kernel_info_reply: %s", kernel_id)
+            finish()
+            if not future.done():
+                future.set_exception(gen.TimeoutError("Timeout waiting for restart"))
+        
+        def on_restart_failed():
+            self.log.warn("Restarting kernel failed: %s", kernel_id)
+            finish()
+            if not future.done():
+                future.set_exception(RuntimeError("Restart failed"))
+        
+        kernel.add_restart_callback(on_restart_failed, 'dead')
+        kernel.session.send(channel, "kernel_info_request")
+        channel.on_recv(on_reply)
+        loop = IOLoop.current()
+        timeout = loop.add_timeout(loop.time() + 30, on_timeout)
+        return future
+
     def kernel_model(self, kernel_id):
         """Return a dictionary of kernel information described in the
         JSON standard model."""

From c36158270aa2de40bea4f9b44bd2df426926c9b8 Mon Sep 17 00:00:00 2001
From: Min RK <benjaminrk@gmail.com>
Date: Thu, 15 Oct 2015 09:59:41 +0200
Subject: [PATCH 2/2] Don't set location header on restart

the URL doesn't change, it makes no sense.
---
 notebook/services/kernels/tests/test_kernels_api.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/notebook/services/kernels/tests/test_kernels_api.py b/notebook/services/kernels/tests/test_kernels_api.py
index cf8a725b5..380228b00 100644
--- a/notebook/services/kernels/tests/test_kernels_api.py
+++ b/notebook/services/kernels/tests/test_kernels_api.py
@@ -114,7 +114,6 @@ class KernelAPITest(NotebookTestBase):
 
         # Restart a kernel
         r = self.kern_api.restart(kern2['id'])
-        self.assertEqual(r.headers['Location'], url_path_join(self.url_prefix, 'api/kernels', kern2['id']))
         rekern = r.json()
         self.assertEqual(rekern['id'], kern2['id'])
         self.assertEqual(rekern['name'], kern2['name'])