From bbe2744f872d758953b08fa2a3e44161378aff50 Mon Sep 17 00:00:00 2001
From: Kevin Bates <kbates4@gmail.com>
Date: Mon, 8 Feb 2021 16:44:57 -0800
Subject: [PATCH] [Gateway] Track only this server's kernels

---
 notebook/gateway/managers.py          | 16 +++++++++++-----
 notebook/services/kernels/handlers.py |  3 ++-
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/notebook/gateway/managers.py b/notebook/gateway/managers.py
index 199f3f4c89..bbefa19ea0 100644
--- a/notebook/gateway/managers.py
+++ b/notebook/gateway/managers.py
@@ -435,8 +435,13 @@ async def get_kernel(self, kernel_id=None, **kwargs):
                 raise
         else:
             kernel = json_decode(response.body)
-            self._kernels[kernel_id] = kernel
-        self.log.debug("Kernel retrieved: %s" % kernel)
+            # Only update our models if we already know about this kernel
+            if kernel_id in self._kernels:
+                self._kernels[kernel_id] = kernel
+                self.log.debug("Kernel retrieved: %s", kernel)
+            else:
+                self.log.warning("Kernel '%s' is not managed by this instance.", kernel_id)
+                kernel = None
         return kernel
 
     async def kernel_model(self, kernel_id):
@@ -458,8 +463,9 @@ async def list_kernels(self, **kwargs):
         self.log.debug("Request list kernels: %s", kernel_url)
         response = await gateway_request(kernel_url, method='GET')
         kernels = json_decode(response.body)
-        self._kernels = {x['id']: x for x in kernels}
-        return kernels
+        # Only update our models if we already know about the kernels
+        self._kernels = {x['id']: x for x in kernels if x['id'] in self._kernels}
+        return list(self._kernels.values())
 
     async def shutdown_kernel(self, kernel_id, now=False, restart=False):
         """Shutdown a kernel by its kernel uuid.
@@ -512,7 +518,7 @@ def shutdown_all(self, now=False):
         kwargs = {'method': 'DELETE'}
         kwargs = GatewayClient.instance().load_connection_args(**kwargs)
         client = HTTPClient()
-        for kernel_id in self._kernels.keys():
+        for kernel_id in self._kernels:
             kernel_url = self._get_kernel_endpoint_url(kernel_id)
             self.log.debug("Request delete kernel at: %s", kernel_url)
             try:
diff --git a/notebook/services/kernels/handlers.py b/notebook/services/kernels/handlers.py
index 5b624c25d0..0bb2fd7b50 100644
--- a/notebook/services/kernels/handlers.py
+++ b/notebook/services/kernels/handlers.py
@@ -54,9 +54,10 @@ def post(self):
 class KernelHandler(APIHandler):
 
     @web.authenticated
+    @gen.coroutine
     def get(self, kernel_id):
         km = self.kernel_manager
-        model = km.kernel_model(kernel_id)
+        model = yield maybe_future(km.kernel_model(kernel_id))
         self.finish(json.dumps(model, default=date_default))
 
     @web.authenticated