Use zmq stream not socket to send messages to shell channel thread

Author: ianthomas23

ipykernel/kernelbase.py

@@ -387,7 +387,7 @@ async def dispatch_shell(self, msg, /, subshell_id: str | None = None):
         assert msg["header"].get("subshell_id") == subshell_id
 
         if self._supports_kernel_subshells:
-            stream = self.shell_channel_thread.manager.get_subshell_to_shell_channel_socket(
+            stream = self.shell_channel_thread.manager.get_subshell_to_shell_channel_stream(
                 subshell_id
             )
         else:
@@ -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 and hasattr(stream, "flush"):
+        if stream:
             stream.flush(zmq.POLLOUT)
 
     def _publish_debug_event(self, event):

ipykernel/socket_pair.py

@@ -15,19 +15,30 @@ class SocketPair:
     """
 
     from_socket: zmq.Socket[Any]
+    from_stream: ZMQStream | None = None
     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):
+    def __init__(
+        self,
+        context: zmq.Context[Any],
+        name: str,
+        from_io_loop: IOLoop | None = None,
+    ):
         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?
+        self.to_socket.connect(address)
+
+        if from_io_loop is not None:
+            self.from_stream = ZMQStream(self.from_socket, from_io_loop)
 
     def close(self):
+        if self.from_stream is not None:
+            self.from_stream.close()
         self.from_socket.close()
 
         if self.to_stream is not None:

ipykernel/subshell.py

@@ -1,4 +1,5 @@
 """A thread for a subshell."""
+from __future__ import annotations
 
 from typing import Any
 
@@ -21,7 +22,7 @@ def __init__(
         super().__init__(name=f"subshell-{subshell_id}", **kwargs)
 
         self.shell_channel_to_subshell = SocketPair(context, subshell_id)
-        self.subshell_to_shell_channel = SocketPair(context, subshell_id + "-reverse")
+        self.subshell_to_shell_channel = SocketPair(context, subshell_id + "-reverse", self.io_loop)
 
         # When aborting flag is set, execute_request messages to this subshell will be aborted.
         self.aborting = False

ipykernel/subshell_manager.py

@@ -9,6 +9,7 @@
 
 import zmq
 from tornado.ioloop import IOLoop
+from zmq.eventloop.zmqstream import ZMQStream
 
 from .socket_pair import SocketPair
 from .subshell import SubshellThread
@@ -58,7 +59,7 @@ def __init__(
 
         # 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 = SocketPair(self._context, "main-reverse", IOLoop.current())
         self._main_to_shell_channel.on_recv(
             self._shell_channel_io_loop, self._send_on_shell_channel
         )
@@ -84,11 +85,14 @@ def get_shell_channel_to_subshell_pair(self, subshell_id: str | None) -> SocketP
         with self._lock_cache:
             return self._cache[subshell_id].shell_channel_to_subshell
 
-    def get_subshell_to_shell_channel_socket(self, subshell_id: str | None) -> zmq.Socket[t.Any]:
+    def get_subshell_to_shell_channel_stream(self, subshell_id: str | None) -> ZMQStream:
         if subshell_id is None:
-            return self._main_to_shell_channel.from_socket
-        with self._lock_cache:
-            return self._cache[subshell_id].subshell_to_shell_channel.from_socket
+            stream = self._main_to_shell_channel.from_stream
+        else:
+            with self._lock_cache:
+                stream = self._cache[subshell_id].subshell_to_shell_channel.from_stream
+        assert stream is not None
+        return stream
 
     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