Use asyncio.Lock around subshell message handling

Author: ianthomas23

ipykernel/kernelbase.py

@@ -205,6 +205,9 @@ def _default_ident(self):
     # see https://github.com/jupyterlab/jupyterlab/issues/17785
     _parent_ident: Mapping[str, bytes]
 
+    # Asyncio lock for main shell thread.
+    _main_asyncio_lock: asyncio.Lock
+
     @property
     def _parent_header(self):
         warnings.warn(
@@ -327,6 +330,8 @@ def __init__(self, **kwargs):
             }
         )
 
+        self._main_asyncio_lock = asyncio.Lock()
+
     async def dispatch_control(self, msg):
         """Dispatch a control request, ensuring only one message is processed at a time."""
         # Ensure only one control message is processed at a time
@@ -649,56 +654,53 @@ async def shell_channel_thread_main(self, msg):
         """Handler for shell messages received on shell_channel_thread"""
         assert threading.current_thread() == self.shell_channel_thread
 
-        if self.session is None:
-            return
+        async with self.shell_channel_thread.asyncio_lock:
+            if self.session is None:
+                return
 
-        # deserialize only the header to get subshell_id
-        # Keep original message to send to subshell_id unmodified.
-        _, msg2 = self.session.feed_identities(msg, copy=False)
-        try:
-            msg3 = self.session.deserialize(msg2, content=False, copy=False)
-            subshell_id = msg3["header"].get("subshell_id")
-
-            # Find inproc pair socket to use to send message to correct subshell.
-            subshell_manager = self.shell_channel_thread.manager
-            socket = subshell_manager.get_shell_channel_to_subshell_socket(subshell_id)
-            assert socket is not None
-            socket.send_multipart(msg, copy=False)
-        except Exception:
-            self.log.error("Invalid message", exc_info=True)  # noqa: G201
+            # deserialize only the header to get subshell_id
+            # Keep original message to send to subshell_id unmodified.
+            _, msg2 = self.session.feed_identities(msg, copy=False)
+            try:
+                msg3 = self.session.deserialize(msg2, content=False, copy=False)
+                subshell_id = msg3["header"].get("subshell_id")
+
+                # Find inproc pair socket to use to send message to correct subshell.
+                subshell_manager = self.shell_channel_thread.manager
+                socket = subshell_manager.get_shell_channel_to_subshell_socket(subshell_id)
+                assert socket is not None
+                socket.send_multipart(msg, copy=False)
+            except Exception:
+                self.log.error("Invalid message", exc_info=True)  # noqa: G201
 
-        if self.shell_stream:
-            self.shell_stream.flush()
+            if self.shell_stream:
+                self.shell_stream.flush()
 
     async def shell_main(self, subshell_id: str | None, msg):
         """Handler of shell messages for a single subshell"""
         if self._supports_kernel_subshells:
             if subshell_id is None:
                 assert threading.current_thread() == threading.main_thread()
+                asyncio_lock = self._main_asyncio_lock
             else:
                 assert threading.current_thread() not in (
                     self.shell_channel_thread,
                     threading.main_thread(),
                 )
-            socket_pair = self.shell_channel_thread.manager.get_shell_channel_to_subshell_pair(
-                subshell_id
-            )
+                asyncio_lock = self.shell_channel_thread.manager.get_subshell_asyncio_lock(
+                    subshell_id
+                )
         else:
             assert subshell_id is None
             assert threading.current_thread() == threading.main_thread()
-            socket_pair = None
+            asyncio_lock = self._main_asyncio_lock
 
-        try:
-            # Whilst executing a shell message, do not accept any other shell messages on the
-            # same subshell, so that cells are run sequentially. Without this we can run multiple
-            # async cells at the same time which would be a nice feature to have but is an API
-            # change.
-            if socket_pair:
-                socket_pair.pause_on_recv()
+        # Whilst executing a shell message, do not accept any other shell messages on the
+        # same subshell, so that cells are run sequentially. Without this we can run multiple
+        # async cells at the same time which would be a nice feature to have but is an API
+        # change.
+        async with asyncio_lock:
             await self.dispatch_shell(msg, subshell_id=subshell_id)
-        finally:
-            if socket_pair:
-                socket_pair.resume_on_recv()
 
     def record_ports(self, ports):
         """Record the ports that this kernel is using.

ipykernel/shellchannel.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 from typing import Any
 
 import zmq
@@ -28,6 +29,8 @@ def __init__(
         self._context = context
         self._shell_socket = shell_socket
 
+        self.asyncio_lock = asyncio.Lock()
+
     @property
     def manager(self) -> SubshellManager:
         # Lazy initialisation.

ipykernel/socket_pair.py

@@ -21,8 +21,6 @@ class SocketPair:
     from_socket: zmq.Socket[Any]
     to_socket: zmq.Socket[Any]
     to_stream: ZMQStream | None = None
-    on_recv_callback: Any
-    on_recv_copy: bool
 
     def __init__(self, context: zmq.Context[Any], name: str):
         """Initialize the inproc socker pair."""
@@ -43,21 +41,9 @@ def close(self):
     def on_recv(self, io_loop: IOLoop, on_recv_callback, copy: bool = False):
         """Set the callback used when a message is received on the to stream."""
         # io_loop is that of the 'to' thread.
-        self.on_recv_callback = on_recv_callback
-        self.on_recv_copy = copy
         if self.to_stream is None:
             self.to_stream = ZMQStream(self.to_socket, io_loop)
-        self.resume_on_recv()
-
-    def pause_on_recv(self):
-        """Pause receiving on the to stream."""
-        if self.to_stream is not None:
-            self.to_stream.stop_on_recv()
-
-    def resume_on_recv(self):
-        """Resume receiving on the to stream."""
-        if self.to_stream is not None and not self.to_stream.closed():
-            self.to_stream.on_recv(self.on_recv_callback, copy=self.on_recv_copy)
+        self.to_stream.on_recv(on_recv_callback, copy=copy)
 
     def _address(self, name) -> str:
         """Return the address used for this inproc socket pair."""

ipykernel/subshell.py

@@ -1,5 +1,6 @@
 """A thread for a subshell."""
 
+import asyncio
 from typing import Any
 
 import zmq
@@ -29,6 +30,8 @@ def __init__(
         # When aborting flag is set, execute_request messages to this subshell will be aborted.
         self.aborting = False
 
+        self.asyncio_lock = asyncio.Lock()
+
     def run(self) -> None:
         """Run the thread."""
         try:

ipykernel/subshell_manager.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import json
 import typing as t
 import uuid
@@ -46,8 +47,8 @@ def __init__(
         self._shell_channel_io_loop = shell_channel_io_loop
         self._shell_socket = shell_socket
         self._cache: dict[str, SubshellThread] = {}
-        self._lock_cache = Lock()
-        self._lock_shell_socket = Lock()
+        self._lock_cache = Lock()  # Sync lock across threads when accessing cache.
+        self._lock_shell_socket = asyncio.Lock()  # Async lock within shell manager thread.
 
         # Inproc socket pair for communication from control thread to shell channel thread,
         # such as for create_subshell_request messages. Reply messages are returned straight away.
@@ -107,7 +108,13 @@ def get_shell_channel_to_subshell_socket(self, subshell_id: str | None) -> zmq.S
 
     def get_subshell_aborting(self, subshell_id: str) -> bool:
         """Get the boolean aborting flag of the specified subshell."""
-        return self._cache[subshell_id].aborting
+        with self._lock_cache:
+            return self._cache[subshell_id].aborting
+
+    def get_subshell_asyncio_lock(self, subshell_id: str) -> asyncio.Lock:
+        """Return the asyncio lock belonging to the specified subshell."""
+        with self._lock_cache:
+            return self._cache[subshell_id].asyncio_lock
 
     def list_subshell(self) -> list[str]:
         """Return list of current subshell ids.
@@ -214,9 +221,9 @@ def _process_control_request(
         # Return the reply to the control thread.
         self.control_to_shell_channel.to_socket.send_json(reply)
 
-    def _send_on_shell_channel(self, msg) -> None:
+    async def _send_on_shell_channel(self, msg) -> None:
         assert current_thread().name == SHELL_CHANNEL_THREAD_NAME
-        with self._lock_shell_socket:
+        async with self._lock_shell_socket:
             self._shell_socket.send_multipart(msg)
 
     def _stop_subshell(self, subshell_thread: SubshellThread) -> None: