From d98997880123603e091b255031949d1cc73f333e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=9D=B4=EC=B0=BD=ED=99=98?= Date: Mon, 4 Sep 2023 17:53:38 +0900 Subject: [PATCH 1/5] Add ready_count to in_pending_statet to implement better lock --- jupyter_client/manager.py | 18 +++++++++++++----- tests/test_manager.py | 14 ++++++++++++-- 2 files changed, 25 insertions(+), 7 deletions(-) diff --git a/jupyter_client/manager.py b/jupyter_client/manager.py index 564966ede..bc48a3f69 100644 --- a/jupyter_client/manager.py +++ b/jupyter_client/manager.py @@ -78,15 +78,25 @@ def in_pending_state(method: F) -> F: @functools.wraps(method) async def wrapper(self, *args, **kwargs): """Create a future for the decorated method.""" - if self._attempted_start or not self._ready: - self._ready = _get_future() + # Initialize the ready_count to 0 if it doesn't exist + if self.owns_kernel: + if not hasattr(self, "_ready_count"): + self._ready_count = 0 + + if self._ready_count == 0: + self._ready = _get_future() + + self._ready_count += 1 + try: # call wrapped method, await, and set the result or exception. out = await method(self, *args, **kwargs) # Add a small sleep to ensure tests can capture the state before done await asyncio.sleep(0.01) if self.owns_kernel: - self._ready.set_result(None) + self._ready_count -= 1 + if self._ready_count == 0: + self._ready.set_result(None) return out except Exception as e: self._ready.set_exception(e) @@ -109,7 +119,6 @@ def __init__(self, *args, **kwargs): self._owns_kernel = kwargs.pop("owns_kernel", True) super().__init__(**kwargs) self._shutdown_status = _ShutdownStatus.Unset - self._attempted_start = False self._ready = None _created_context: Bool = Bool(False) @@ -397,7 +406,6 @@ async def _async_start_kernel(self, **kw: t.Any) -> None: keyword arguments that are passed down to build the kernel_cmd and launching the kernel (e.g. Popen kwargs). """ - self._attempted_start = True kernel_cmd, kw = await self._async_pre_start_kernel(**kw) # launch the kernel subprocess diff --git a/tests/test_manager.py b/tests/test_manager.py index e3d6ea222..833fd6db9 100644 --- a/tests/test_manager.py +++ b/tests/test_manager.py @@ -4,10 +4,10 @@ import os import tempfile from unittest import mock +import asyncio from jupyter_client.kernelspec import KernelSpec -from jupyter_client.manager import KernelManager - +from jupyter_client.manager import KernelManager, in_pending_state def test_connection_file_real_path(): """Verify realpath is used when formatting connection file""" @@ -32,3 +32,13 @@ def test_connection_file_real_path(): km._launch_args = {} cmds = km.format_kernel_cmd() assert cmds[4] == "foobar" + +async def test_in_pending_state(): + """Verify in_pending_state race condition""" + tm = KernelManager() + start_kernel = asyncio.ensure_future(tm._async_start_kernel()) + shutdown_kernel = asyncio.ensure_future(tm._async_shutdown_kernel()) + + await start_kernel + await shutdown_kernel + assert tm.is_alive() == False \ No newline at end of file From 321cdf70ecf23b924b4e542ea3fca8daa09a4dcd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=9D=B4=EC=B0=BD=ED=99=98?= Date: Mon, 4 Sep 2023 18:22:59 +0900 Subject: [PATCH 2/5] Follow code styles --- tests/test_manager.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_manager.py b/tests/test_manager.py index 833fd6db9..b4815d1ab 100644 --- a/tests/test_manager.py +++ b/tests/test_manager.py @@ -7,7 +7,7 @@ import asyncio from jupyter_client.kernelspec import KernelSpec -from jupyter_client.manager import KernelManager, in_pending_state +from jupyter_client.manager import KernelManager def test_connection_file_real_path(): """Verify realpath is used when formatting connection file""" @@ -41,4 +41,4 @@ async def test_in_pending_state(): await start_kernel await shutdown_kernel - assert tm.is_alive() == False \ No newline at end of file + assert tm.is_alive() == False From 45fa5ab1c7dc9e6586fa8178ee1736d5748bf0a1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=9D=B4=EC=B0=BD=ED=99=98?= Date: Mon, 4 Sep 2023 18:27:53 +0900 Subject: [PATCH 3/5] Add missing exception handling --- jupyter_client/manager.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/jupyter_client/manager.py b/jupyter_client/manager.py index bc48a3f69..64de25306 100644 --- a/jupyter_client/manager.py +++ b/jupyter_client/manager.py @@ -99,8 +99,11 @@ async def wrapper(self, *args, **kwargs): self._ready.set_result(None) return out except Exception as e: - self._ready.set_exception(e) - self.log.exception(self._ready.exception()) + if self.owns_kernel: + self._ready_count -= 1 + if self._ready_count == 0: + self._ready.set_exception(e) + self.log.exception(e) raise e return t.cast(F, wrapper) From 096199320abf326ee47bb53cb9ede33c436a5d30 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=9D=B4=EC=B0=BD=ED=99=98?= Date: Mon, 4 Sep 2023 18:35:01 +0900 Subject: [PATCH 4/5] Apply black --- tests/test_manager.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_manager.py b/tests/test_manager.py index b4815d1ab..242c65a78 100644 --- a/tests/test_manager.py +++ b/tests/test_manager.py @@ -9,6 +9,7 @@ from jupyter_client.kernelspec import KernelSpec from jupyter_client.manager import KernelManager + def test_connection_file_real_path(): """Verify realpath is used when formatting connection file""" with mock.patch("os.path.realpath") as patched_realpath: @@ -33,6 +34,7 @@ def test_connection_file_real_path(): cmds = km.format_kernel_cmd() assert cmds[4] == "foobar" + async def test_in_pending_state(): """Verify in_pending_state race condition""" tm = KernelManager() From 39048ababf6466f62e1c0c26126f2ccb292a09ee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=9D=B4=EC=B0=BD=ED=99=98?= Date: Mon, 4 Sep 2023 18:36:43 +0900 Subject: [PATCH 5/5] Apply ruff --- tests/test_manager.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_manager.py b/tests/test_manager.py index 242c65a78..1e3d2a320 100644 --- a/tests/test_manager.py +++ b/tests/test_manager.py @@ -1,10 +1,10 @@ """Tests for KernelManager""" # Copyright (c) Jupyter Development Team. # Distributed under the terms of the Modified BSD License. +import asyncio import os import tempfile from unittest import mock -import asyncio from jupyter_client.kernelspec import KernelSpec from jupyter_client.manager import KernelManager @@ -43,4 +43,4 @@ async def test_in_pending_state(): await start_kernel await shutdown_kernel - assert tm.is_alive() == False + assert tm.is_alive() is False