From b005dfddea13fe6d156e0002c7bba73ab367f8d0 Mon Sep 17 00:00:00 2001
From: Zachary Sailer <zsailer@apple.com>
Date: Mon, 2 May 2022 16:33:09 -0700
Subject: [PATCH] short kernel heartbeat interval (#339)

---
 data_studio_jupyter_extensions/app.py                  |  3 ++-
 .../configurables/kernel_restarter.py                  | 10 +++++-----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/data_studio_jupyter_extensions/app.py b/data_studio_jupyter_extensions/app.py
index 301318373c..bffffcef69 100644
--- a/data_studio_jupyter_extensions/app.py
+++ b/data_studio_jupyter_extensions/app.py
@@ -40,8 +40,9 @@
             "default_provisioner_name": "notebook-service-provisioner"
         },
         "KernelManager": {
-            "restarter_class": "data_studio_jupyter_extensions.configurables.kernel_restarter.DataStudioKernelPoller"
+            "restarter_class": "data_studio_jupyter_extensions.configurables.kernel_restarter.DataStudioKernelRestarter"
         },
+        "KernelRestarter": {"time_to_dead": 1},
         "Session": {"key": b"notebooks"},
         "ServerApp": {
             "log_level": "DEBUG",
diff --git a/data_studio_jupyter_extensions/configurables/kernel_restarter.py b/data_studio_jupyter_extensions/configurables/kernel_restarter.py
index 0191d35a10..6fc66f4faf 100644
--- a/data_studio_jupyter_extensions/configurables/kernel_restarter.py
+++ b/data_studio_jupyter_extensions/configurables/kernel_restarter.py
@@ -8,7 +8,10 @@
 from data_studio_jupyter_extensions.configurables.kernel_status import KernelStatusMixin
 
 
-class DataStudioKernelPoller(AsyncIOLoopKernelRestarter, KernelStatusMixin):
+class DataStudioKernelRestarter(AsyncIOLoopKernelRestarter, KernelStatusMixin):
+    """Replaces the default Restarter with an object that *does not* restart kernels, but
+    just polls them.
+    """
 
     _attempt_count = Int(0)
     _connected_once = Bool(False)
@@ -72,8 +75,6 @@ async def poll(self):
             if km.heartbeat_timeout < (now - self._start_time):
                 self._disconnected_state()
                 await self.kernel_manager.shutdown_kernel()
-                return
-            return
         # Kernel connected before; but a heartbeat was missed
         else:
             self._attempt_count += 1
@@ -99,5 +100,4 @@ async def poll(self):
                 # shutdown the kernel.
                 self._disconnected_state()
                 await self.kernel_manager.shutdown_kernel()
-            return
-        self._attempt_count = 0
+                self._attempt_count = 0