Refactor inproc pair sockets

Author: ianthomas23

ipykernel/kernelbase.py

@@ -350,7 +350,7 @@ def should_handle(self, stream, msg, idents):
             return False
         return True
 
-    async def dispatch_shell(self, msg, /, stream=None, subshell_id: str | None = None):
+    async def dispatch_shell(self, msg, /, subshell_id: str | None = None):
         """dispatch shell requests"""
         if len(msg) == 1 and msg[0].buffer == b"stop aborting":
             # Dummy "stop aborting" message to stop aborting execute requests on this subshell.
@@ -386,11 +386,12 @@ async def dispatch_shell(self, msg, /, stream=None, subshell_id: str | None = No
         msg_type = msg["header"]["msg_type"]
         assert msg["header"].get("subshell_id") == subshell_id
 
-        if stream is None:
-            if self._supports_kernel_subshells:
-                stream = self.shell_channel_thread.manager.get_other_stream(subshell_id)
-            else:
-                stream = self.shell_stream
+        if self._supports_kernel_subshells:
+            stream = self.shell_channel_thread.manager.get_subshell_to_shell_channel_socket(
+                subshell_id
+            )
+        else:
+            stream = self.shell_stream
 
         # Only abort execute requests
         if msg_type == "execute_request":
@@ -611,7 +612,7 @@ async def shell_channel_thread_main(self, msg):
 
             # 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_stream(subshell_id)
+            socket = subshell_manager.get_shell_channel_to_subshell_socket(subshell_id)
             assert socket is not None
             socket.send_multipart(msg, copy=False)
         except Exception:
@@ -630,26 +631,25 @@ async def shell_main(self, subshell_id: str | None, msg):
                     self.shell_channel_thread,
                     threading.main_thread(),
                 )
-            # Inproc pair socket that this subshell uses to talk to shell channel thread.
-            stream = self.shell_channel_thread.manager.get_other_stream(subshell_id)
+            socket_pair = self.shell_channel_thread.manager.get_shell_channel_to_subshell_pair(
+                subshell_id
+            )
         else:
             assert subshell_id is None
             assert threading.current_thread() == threading.main_thread()
-            stream = self.shell_stream
+            socket_pair = None
 
         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.
-            stream.stop_on_recv()
-            await self.dispatch_shell(msg, stream=stream, subshell_id=subshell_id)
+            if socket_pair:
+                socket_pair.pause_on_recv()
+            await self.dispatch_shell(msg, subshell_id=subshell_id)
         finally:
-            stream.on_recv(
-                partial(self.shell_main, subshell_id),
-                copy=False,
-            )
-            stream.flush()
+            if socket_pair:
+                socket_pair.resume_on_recv()
 
     def record_ports(self, ports):
         """Record the ports that this kernel is using.
@@ -690,7 +690,7 @@ def _publish_status(self, status, channel, parent=None):
     def _publish_status_and_flush(self, status, channel, stream, parent=None):
         """send status on IOPub and flush specified stream to ensure reply is sent before handling the next reply"""
         self._publish_status(status, channel, parent)
-        if stream:
+        if stream and hasattr(stream, "flush"):
             stream.flush(zmq.POLLOUT)
 
     def _publish_debug_event(self, event):
@@ -835,6 +835,8 @@ async def execute_request(self, stream, ident, parent):
         if self._do_exec_accepted_params["cell_id"]:
             do_execute_args["cell_id"] = cell_id
 
+        subshell_id = parent["header"].get("subshell_id")
+
         # Call do_execute with the appropriate arguments
         reply_content = self.do_execute(**do_execute_args)
 
@@ -1174,9 +1176,9 @@ async def create_subshell_request(self, socket, ident, parent) -> None:
 
         # This should only be called in the control thread if it exists.
         # Request is passed to shell channel thread to process.
-        other_socket = self.shell_channel_thread.manager.get_control_other_socket()
-        other_socket.send_json({"type": "create"})
-        reply = other_socket.recv_json()
+        control_socket = self.shell_channel_thread.manager.control_to_shell_channel.from_socket
+        control_socket.send_json({"type": "create"})
+        reply = control_socket.recv_json()
         self.session.send(socket, "create_subshell_reply", reply, parent, ident)
 
     async def delete_subshell_request(self, socket, ident, parent) -> None:
@@ -1197,9 +1199,10 @@ async def delete_subshell_request(self, socket, ident, parent) -> None:
 
         # This should only be called in the control thread if it exists.
         # Request is passed to shell channel thread to process.
-        other_socket = self.shell_channel_thread.manager.get_control_other_socket()
-        other_socket.send_json({"type": "delete", "subshell_id": subshell_id})
-        reply = other_socket.recv_json()
+        control_socket = self.shell_channel_thread.manager.control_to_shell_channel.from_socket
+        control_socket.send_json({"type": "delete", "subshell_id": subshell_id})
+        reply = control_socket.recv_json()
+
         self.session.send(socket, "delete_subshell_reply", reply, parent, ident)
 
     async def list_subshell_request(self, socket, ident, parent) -> None:
@@ -1213,9 +1216,10 @@ async def list_subshell_request(self, socket, ident, parent) -> None:
 
         # This should only be called in the control thread if it exists.
         # Request is passed to shell channel thread to process.
-        other_socket = self.shell_channel_thread.manager.get_control_other_socket()
-        other_socket.send_json({"type": "list"})
-        reply = other_socket.recv_json()
+        control_socket = self.shell_channel_thread.manager.control_to_shell_channel.from_socket
+        control_socket.send_json({"type": "list"})
+        reply = control_socket.recv_json()
+
         self.session.send(socket, "list_subshell_reply", reply, parent, ident)
 
     # ---------------------------------------------------------------------------
@@ -1315,7 +1319,7 @@ def _post_dummy_stop_aborting_message(self, subshell_id: str | None) -> None:
         the _aborting flag.
         """
         subshell_manager = self.shell_channel_thread.manager
-        socket = subshell_manager.get_shell_channel_stream(subshell_id)
+        socket = subshell_manager.get_shell_channel_to_subshell_socket(subshell_id)
         assert socket is not None
 
         msg = b"stop aborting"  # Magic string for dummy message.

ipykernel/socket_pair.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from typing import Any
+
+import zmq
+from tornado.ioloop import IOLoop
+from zmq.eventloop.zmqstream import ZMQStream
+
+
+class SocketPair:
+    """Pair of ZMQ inproc sockets for one-direction communication between 2 threads.
+
+    One of the threads is always the shell_channel_thread, the other may be the control
+    thread, main thread or a subshell thread.
+    """
+
+    from_socket: zmq.Socket[Any]
+    to_socket: zmq.Socket[Any]
+    to_stream: ZMQStream
+    on_recv_callback: Any
+    on_recv_copy: bool
+
+    def __init__(self, context: zmq.Context[Any], name: str):
+        self.from_socket = context.socket(zmq.PAIR)
+        self.to_socket = context.socket(zmq.PAIR)
+        address = self._address(name)
+        self.from_socket.bind(address)
+        self.to_socket.connect(address)  # Or do I need to do this in another thread?
+
+    def close(self):
+        self.from_socket.close()
+
+        if self.to_stream is not None:
+            self.to_stream.close()
+        self.to_socket.close()
+
+    def on_recv(self, io_loop: IOLoop, on_recv_callback, copy: bool = False):
+        # io_loop is that of the 'to' thread.
+        self.on_recv_callback = on_recv_callback
+        self.on_recv_copy = copy
+        self.to_stream = ZMQStream(self.to_socket, io_loop)
+        self.resume_on_recv()
+
+    def pause_on_recv(self):
+        self.to_stream.stop_on_recv()
+
+    def resume_on_recv(self):
+        self.to_stream.on_recv(self.on_recv_callback, copy=self.on_recv_copy)
+
+    def _address(self, name) -> str:
+        return f"inproc://subshell{name}"

ipykernel/subshell.py

@@ -3,10 +3,9 @@
 from typing import Any
 
 import zmq
-from zmq.eventloop.zmqstream import ZMQStream
 
+from .socket_pair import SocketPair
 from .thread import BaseThread
-from .utils import create_inproc_pair_socket
 
 
 class SubshellThread(BaseThread):
@@ -21,12 +20,8 @@ def __init__(
         """Initialize the thread."""
         super().__init__(name=f"subshell-{subshell_id}", **kwargs)
 
-        shell_channel_socket = create_inproc_pair_socket(context, subshell_id, True)
-        # io_loop will be current io_loop which is of ShellChannelThread
-        self.shell_channel_stream = ZMQStream(shell_channel_socket)
-
-        subshell_socket = create_inproc_pair_socket(context, subshell_id, False)
-        self.subshell_stream = ZMQStream(subshell_socket, self.io_loop)
+        self.shell_channel_to_subshell = SocketPair(context, subshell_id)
+        self.subshell_to_shell_channel = SocketPair(context, subshell_id + "-reverse")
 
         # When aborting flag is set, execute_request messages to this subshell will be aborted.
         self.aborting = False
@@ -35,5 +30,5 @@ def run(self) -> None:
         try:
             super().run()
         finally:
-            self.subshell_stream.close()
-            self.shell_channel_stream.close()
+            self.shell_channel_to_subshell.close()
+            self.subshell_to_shell_channel.close()

ipykernel/subshell_manager.py

@@ -9,11 +9,10 @@
 
 import zmq
 from tornado.ioloop import IOLoop
-from zmq.eventloop.zmqstream import ZMQStream
 
+from .socket_pair import SocketPair
 from .subshell import SubshellThread
 from .thread import SHELL_CHANNEL_THREAD_NAME
-from .utils import create_inproc_pair_socket
 
 
 class SubshellManager:
@@ -46,26 +45,23 @@ def __init__(
         self._lock_cache = Lock()
         self._lock_shell_socket = Lock()
 
-        # Inproc pair sockets for control channel and main shell (parent subshell).
-        # Each inproc pair has a "shell_channel" socket used in the shell channel
-        # thread, and an "other" socket used in the other thread.
-        control_shell_channel_socket = create_inproc_pair_socket(self._context, "control", True)
-        self._control_shell_channel_stream = ZMQStream(
-            control_shell_channel_socket, self._shell_channel_io_loop
+        # 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.
+        self.control_to_shell_channel = SocketPair(self._context, "control")
+        self.control_to_shell_channel.on_recv(
+            self._shell_channel_io_loop, self._process_control_request, copy=True
         )
-        self._control_shell_channel_stream.on_recv(self._process_control_request, copy=True)
 
-        self._control_other_socket = create_inproc_pair_socket(self._context, "control", False)
+        # Inproc socket pair for communication from shell channel thread to main thread,
+        # such as for execute_request messages.
+        self._shell_channel_to_main = SocketPair(self._context, "main")
 
-        parent_shell_channel_socket = create_inproc_pair_socket(self._context, None, True)
-        self._parent_shell_channel_stream = ZMQStream(
-            parent_shell_channel_socket, self._shell_channel_io_loop
+        # Inproc socket pair for communication from main thread to shell channel thread.
+        # such as for execute_reply messages.
+        self._main_to_shell_channel = SocketPair(self._context, "main-reverse")
+        self._main_to_shell_channel.on_recv(
+            self._shell_channel_io_loop, self._send_on_shell_channel
         )
-        self._parent_shell_channel_stream.on_recv(self._send_on_shell_channel, copy=False)
-
-        # Initialised in set_on_recv_callback
-        self._on_recv_callback: t.Any = None  # Callback for ZMQStream.on_recv for "other" sockets.
-        self._parent_other_stream: ZMQStream | None = None
 
     def close(self) -> None:
         """Stop all subshells and close all resources."""
@@ -78,40 +74,24 @@ def close(self) -> None:
                     break
                 self._stop_subshell(subshell_thread)
 
-        for socket_or_stream in (
-            self._control_shell_channel_stream,
-            self._parent_shell_channel_stream,
-            self._parent_other_stream,
-        ):
-            if socket_or_stream is not None:
-                socket_or_stream.close()
-
-        if self._control_other_socket is not None:
-            self._control_other_socket.close()
-
-    def get_control_other_socket(self) -> zmq.Socket[t.Any]:
-        return self._control_other_socket
+        self.control_to_shell_channel.close()
+        self._main_to_shell_channel.close()
+        self._shell_channel_to_main.close()
 
-    def get_other_stream(self, subshell_id: str | None) -> ZMQStream:
-        """Return the other inproc pair socket for a subshell.
-
-        This socket is accessed from the subshell thread.
-        """
+    def get_shell_channel_to_subshell_pair(self, subshell_id: str | None) -> SocketPair:
         if subshell_id is None:
-            assert self._parent_other_stream is not None
-            return self._parent_other_stream
+            return self._shell_channel_to_main
         with self._lock_cache:
-            return self._cache[subshell_id].subshell_stream
+            return self._cache[subshell_id].shell_channel_to_subshell
 
-    def get_shell_channel_stream(self, subshell_id: str | None) -> ZMQStream:
-        """Return the stream for the shell channel inproc pair socket for a subshell.
-
-        This stream is accessed from the shell channel thread.
-        """
+    def get_subshell_to_shell_channel_socket(self, subshell_id: str | None) -> zmq.Socket[t.Any]:
         if subshell_id is None:
-            return self._parent_shell_channel_stream
+            return self._main_to_shell_channel.from_socket
         with self._lock_cache:
-            return self._cache[subshell_id].shell_channel_stream
+            return self._cache[subshell_id].subshell_to_shell_channel.from_socket
+
+    def get_shell_channel_to_subshell_socket(self, subshell_id: str | None) -> zmq.Socket[t.Any]:
+        return self.get_shell_channel_to_subshell_pair(subshell_id).from_socket
 
     def get_subshell_aborting(self, subshell_id: str) -> bool:
         """Get the aborting flag of the specified subshell."""
@@ -128,13 +108,7 @@ def list_subshell(self) -> list[str]:
     def set_on_recv_callback(self, on_recv_callback):
         assert current_thread() == main_thread()
         self._on_recv_callback = on_recv_callback
-        if not self._parent_other_stream:
-            parent_other_socket = create_inproc_pair_socket(self._context, None, False)
-            self._parent_other_stream = ZMQStream(parent_other_socket)
-            self._parent_other_stream.on_recv(
-                partial(self._on_recv_callback, None),
-                copy=False,
-            )
+        self._shell_channel_to_main.on_recv(IOLoop.current(), partial(self._on_recv_callback, None))
 
     def set_subshell_aborting(self, subshell_id: str, aborting: bool) -> None:
         """Set the aborting flag of the specified subshell."""
@@ -168,11 +142,14 @@ def _create_subshell(self) -> str:
             assert subshell_id not in self._cache
             self._cache[subshell_id] = subshell_thread
 
-        subshell_thread.subshell_stream.on_recv(
+        subshell_thread.shell_channel_to_subshell.on_recv(
+            subshell_thread.io_loop,
             partial(self._on_recv_callback, subshell_id),
-            copy=False,
         )
-        subshell_thread.shell_channel_stream.on_recv(self._send_on_shell_channel, copy=False)
+
+        subshell_thread.subshell_to_shell_channel.on_recv(
+            self._shell_channel_io_loop, self._send_on_shell_channel
+        )
 
         subshell_thread.start()
         return subshell_id
@@ -219,7 +196,8 @@ def _process_control_request(
                 "evalue": str(err),
             }
 
-        self._control_shell_channel_stream.send_json(reply)
+        # 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:
         assert current_thread().name == SHELL_CHANNEL_THREAD_NAME