wait to cleanup kernels after kernel is finished pending (#748)

Author: Zsailer

jupyter_server/services/kernels/kernelmanager.py

@@ -188,6 +188,11 @@ def cwd_for_path(self, path):
             os_path = os.path.dirname(os_path)
         return os_path
 
+    async def _remove_kernel_when_ready(self, kernel_id, kernel_awaitable):
+        await super()._remove_kernel_when_ready(kernel_id, kernel_awaitable)
+        self._kernel_connections.pop(kernel_id, None)
+        self._kernel_ports.pop(kernel_id, None)
+
     async def start_kernel(self, kernel_id=None, path=None, **kwargs):
         """Start a kernel for a session and return its kernel_id.
 
@@ -384,19 +389,12 @@ def shutdown_kernel(self, kernel_id, now=False, restart=False):
         self._check_kernel_id(kernel_id)
         self.stop_watching_activity(kernel_id)
         self.stop_buffering(kernel_id)
-        self._kernel_connections.pop(kernel_id, None)
 
         # Decrease the metric of number of kernels
         # running for the relevant kernel type by 1
         KERNEL_CURRENTLY_RUNNING_TOTAL.labels(type=self._kernels[kernel_id].kernel_name).dec()
 
         self.pinned_superclass.shutdown_kernel(self, kernel_id, now=now, restart=restart)
-        # Unlike its async sibling method in AsyncMappingKernelManager, removing the kernel_id
-        # from the connections dictionary isn't as problematic before the shutdown since the
-        # method is synchronous.  However, we'll keep the relative call orders the same from
-        # a maintenance perspective.
-        self._kernel_connections.pop(kernel_id, None)
-        self._kernel_ports.pop(kernel_id, None)
 
     async def restart_kernel(self, kernel_id, now=False):
         """Restart a kernel by kernel_id"""
@@ -655,6 +653,4 @@ async def shutdown_kernel(self, kernel_id, now=False, restart=False):
         ret = await self.pinned_superclass.shutdown_kernel(
             self, kernel_id, now=now, restart=restart
         )
-        self._kernel_connections.pop(kernel_id, None)
-        self._kernel_ports.pop(kernel_id, None)
         return ret