Replaced PUB socket with XPUB socket

Author: JohanMabille

ipykernel/iostream.py

@@ -20,7 +20,7 @@
 from typing import Any
 
 import zmq
-from jupyter_client.session import extract_header
+from jupyter_client.session import extract_header, Session
 from tornado.ioloop import IOLoop
 from zmq.eventloop.zmqstream import ZMQStream
 
@@ -73,12 +73,75 @@ def __init__(self, socket, pipe=False):
         self._event_pipe_gc_seconds: float = 10
         self._event_pipe_gc_task: asyncio.Task[Any] | None = None
         self._setup_event_pipe()
+        self._setup_xpub_listener()
         self.thread = threading.Thread(target=self._thread_main, name="IOPub")
         self.thread.daemon = True
         self.thread.pydev_do_not_trace = True  # type:ignore[attr-defined]
         self.thread.is_pydev_daemon_thread = True  # type:ignore[attr-defined]
         self.thread.name = "IOPub"
 
+    def _setup_xpub_listener(self):
+        """Setup listener for XPUB subscription events"""
+
+        # Checks the socket is not a DummySocket
+        if not hasattr(self.socket, 'getsockopt'):
+            return
+
+        socket_type = self.socket.getsockopt(zmq.TYPE)
+        if socket_type == zmq.XPUB:
+            self._xpub_stream = ZMQStream(self.socket, self.io_loop)
+            self._xpub_stream.on_recv(self._handle_subscription)
+
+    def _handle_subscription(self, frames):
+        """Handle subscription/unsubscription events from XPUB socket
+
+        XPUB sockets receive: 
+        - subscribe:  single frame with b'\\x01' + topic
+        - unsubscribe: single frame with b'\\x00' + topic
+        """
+
+        for frame in frames:
+            event_type = frame[0]
+            if event_type == 1:
+                subscription = frame[1:] if len(frame) > 1 else b''
+                try:
+                    subscription_str = subscription.decode('utf-8')
+                except UnicodeDecodeError:
+                    continue
+                self._send_welcome_message(subscription_str)
+
+    def _send_welcome_message(self, subscription):
+        """Send iopub_welcome message for new subscription
+
+        Parameters
+        ----------
+        subscription : str
+            The subscription topic (UTF-8 decoded)
+        """
+
+        content = {
+            'subscription': subscription
+        }
+
+        header = self.session.msg_header('iopub_welcome')
+        msg = {
+            'header': header,
+            'parent_header': {},
+            'metadata': {},
+            'content': content,
+            'buffers': [],
+        }
+
+        msg_list = self.session.serialize(msg)
+
+        if subscription:
+            identity = subscription.encode('utf-8')
+            full_msg = [identity] + msg_list
+        else:
+            full_msg = msg_list
+        # Send directly on socket (we're already in IO thread context)
+        self.socket.send_multipart(full_msg)
+
     def _thread_main(self):
         """The inner loop that's actually run in a thread"""
 

ipykernel/kernelapp.py

@@ -377,12 +377,13 @@ def init_control(self, context):
 
     def init_iopub(self, context):
         """Initialize the iopub channel."""
-        self.iopub_socket = context.socket(zmq.PUB)
+        self.iopub_socket = context.socket(zmq.XPUB)
         self.iopub_socket.linger = 1000
         self.iopub_port = self._bind_socket(self.iopub_socket, self.iopub_port)
         self.log.debug("iopub PUB Channel on port: %i", self.iopub_port)
         self.configure_tornado_logger()
         self.iopub_thread = IOPubThread(self.iopub_socket, pipe=True)
+        self.iopub_thread.session = self.session
         self.iopub_thread.start()
         # backward-compat: wrap iopub socket API in background thread
         self.iopub_socket = self.iopub_thread.background_socket

tests/conftest.py

@@ -48,7 +48,7 @@ class KernelMixin:
 
     def _initialize(self):
         self.context = context = zmq.Context()
-        self.iopub_socket = context.socket(zmq.PUB)
+        self.iopub_socket = context.socket(zmq.XPUB)
         self.stdin_socket = context.socket(zmq.ROUTER)
         self.session = Session()
         self.test_sockets = [self.iopub_socket]