Port over jupyter/notebook#5355 and other fixes

jupyter_server/services/kernels/kernelmanager.py

@@ -468,8 +468,9 @@ async def cull_kernels(self):
 
     async def cull_kernel_if_idle(self, kernel_id):
         kernel = self._kernels[kernel_id]
-        self.log.debug("kernel_id=%s, kernel_name=%s, last_activity=%s", kernel_id, kernel.kernel_name, kernel.last_activity)
-        if kernel.last_activity is not None:
+        if hasattr(kernel, 'last_activity'):  # last_activity is monkey-patched, so ensure that has occurred
+            self.log.debug("kernel_id=%s, kernel_name=%s, last_activity=%s",
+                           kernel_id, kernel.kernel_name, kernel.last_activity)
             dt_now = utcnow()
             dt_idle = dt_now - kernel.last_activity
             # Compute idle properties
@@ -482,7 +483,7 @@ async def cull_kernel_if_idle(self, kernel_id):
                 idle_duration = int(dt_idle.total_seconds())
                 self.log.warning("Culling '%s' kernel '%s' (%s) with %d connections due to %s seconds of inactivity.",
                                  kernel.execution_state, kernel.kernel_name, kernel_id, connections, idle_duration)
-                await self.shutdown_kernel(kernel_id)
+                await ensure_async(self.shutdown_kernel(kernel_id))
 
 
 # AsyncMappingKernelManager inherits as much as possible from MappingKernelManager,
@@ -506,12 +507,14 @@ async def shutdown_kernel(self, kernel_id, now=False, restart=False):
             kernel._activity_stream.close()
             kernel._activity_stream = None
         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()
 
-        return await self.pinned_superclass.shutdown_kernel(self, kernel_id, now=now, restart=restart)
+        # Finish shutting down the kernel before clearing state to avoid a race condition.
+        ret = await self.pinned_superclass.shutdown_kernel(self, kernel_id, now=now, restart=restart)
+        self._kernel_connections.pop(kernel_id, None)
+        return ret

tests/services/contents/test_api.py

@@ -42,7 +42,7 @@ def dirs_only(dir_model):
 
 
 @pytest.fixture(params=["FileContentsManager", "AsyncFileContentsManager"])
-def argv(request):
+def jp_argv(request):
     return ["--ServerApp.contents_manager_class=jupyter_server.services.contents.filemanager." + request.param]
 
 

tests/services/kernels/test_api.py

@@ -17,7 +17,7 @@
 
 
 @pytest.fixture(params=["MappingKernelManager", "AsyncMappingKernelManager"])
-def argv(request):
+def jp_argv(request):
     return ["--ServerApp.kernel_manager_class=jupyter_server.services.kernels.kernelmanager." + request.param]
 
 
@@ -209,7 +209,6 @@ async def test_connection(jp_fetch, jp_ws_fetch, jp_http_port, jp_auth_header):
     model = json.loads(r.body.decode())
     assert model['connections'] == 0
 
-    time.sleep(1)
     # Open a websocket connection.
     ws = await jp_ws_fetch(
         'api', 'kernels', kid, 'channels'

tests/services/kernels/test_cull.py

@@ -0,0 +1,73 @@
+import asyncio
+import json
+import platform
+import sys
+import time
+import pytest
+from traitlets.config import Config
+from tornado.httpclient import HTTPClientError
+
+
+@pytest.fixture(params=["MappingKernelManager", "AsyncMappingKernelManager"])
+def jp_argv(request):
+    return ["--ServerApp.kernel_manager_class=jupyter_server.services.kernels.kernelmanager." + request.param]
+
+
+CULL_TIMEOUT = 10 if platform.python_implementation() == 'PyPy' else 2
+
+@pytest.fixture
+def jp_server_config():
+    return Config({
+        'ServerApp': {
+            'MappingKernelManager': {
+                'cull_idle_timeout': CULL_TIMEOUT,
+                'cull_interval': 1,
+                'cull_connected': False
+            }
+        }
+    })
+
+
+async def test_culling(jp_fetch, jp_ws_fetch):
+    r = await jp_fetch(
+        'api', 'kernels',
+        method='POST',
+        allow_nonstandard_methods=True
+    )
+    kernel = json.loads(r.body.decode())
+    kid = kernel['id']
+
+    # Open a websocket connection.
+    ws = await jp_ws_fetch(
+        'api', 'kernels', kid, 'channels'
+    )
+
+    r = await jp_fetch(
+        'api', 'kernels', kid,
+        method='GET'
+    )
+    model = json.loads(r.body.decode())
+    assert model['connections'] == 1
+    culled = await get_cull_status(kid, jp_fetch)  # connected, should not be culled
+    assert not culled
+    ws.close()
+    culled = await get_cull_status(kid, jp_fetch)  # not connected, should be culled
+    assert culled
+
+
+async def get_cull_status(kid, jp_fetch):
+    culled = False
+    for i in range(20):  # Need max of 2x culling PERIOD ensure culling timeout exceeded
+        try:
+            r = await jp_fetch(
+                'api', 'kernels', kid,
+                method='GET'
+            )
+            kernel = json.loads(r.body.decode())
+        except HTTPClientError as e:
+            assert e.code == 404
+            culled = True
+            break
+        else:
+            await asyncio.sleep(CULL_TIMEOUT / 10.)
+    return culled

tests/services/sessions/test_api.py

@@ -17,7 +17,7 @@
 
 
 @pytest.fixture(params=["MappingKernelManager", "AsyncMappingKernelManager"])
-def argv(request):
+def jp_argv(request):
     return ["--ServerApp.kernel_manager_class=jupyter_server.services.kernels.kernelmanager." + request.param]