Keeping references in-process sockets for graceful exit

[jacopoabramo]

I'm currently implementing an in-proc publisher-subscriber pattern using zmq.devices.ThreadDevice. A working example is presented below:

from __future__ import annotations

from typing import (
    ClassVar,
    Iterable,
    Optional,
    Union
)
from weakref import WeakSet
import logging
import threading
import time
import zmq
import zmq.devices

_INPROC_XPUB = "inproc://virtual_xpub"
_INPROC_XSUB = "inproc://virtual_xsub"

logging.basicConfig(level=logging.INFO)

class VirtualBus:
    _INPROC_MAP: ClassVar[dict[int, str]] = {
        # a subscriber publishes to the XPUB socket;
        # a publisher subscribes to the XSUB socket
        zmq.SUB: _INPROC_XPUB,
        zmq.PUB: _INPROC_XSUB,
    }

    _PUB_SOCKETS: WeakSet[zmq.SyncSocket] = WeakSet()

    def __init__(self) -> None:
        self._context = zmq.Context.instance()
        self._forwarder = zmq.devices.ThreadDevice(zmq.FORWARDER, zmq.XSUB, zmq.XPUB)
        self._forwarder.daemon = True
        self._forwarder.setsockopt_in(zmq.LINGER, 0)
        self._forwarder.setsockopt_out(zmq.LINGER, 0)
        self._forwarder.bind_in(_INPROC_XSUB)
        self._forwarder.bind_out(_INPROC_XPUB)
        self._forwarder.start()

    def shutdown(self) -> None:
        """Shutdown the virtual bus.

        Closes the ZMQ context and terminates the streamer queue.
        """
        logging.info("Shutting down the virtual bus.")
        for socket in self._PUB_SOCKETS:
            socket.close()
        try:
            self._context.term()
        except zmq.error.ZMQError:
            logging.info("ZMQ context already terminated.")
        finally:
            if not self._forwarder.done:
                self._forwarder.join()

    def connect(
        self,
        socket_type: int,
        *,
        topic: Optional[Union[str, Iterable[str]]] = None,
    ) -> Optional[Union[zmq.SyncSocket, tuple[zmq.SyncSocket, zmq.Poller]]]:
        """Return a new socket connected to the virtual bus context.

        Class instances may use this method to
        create a new socket which will be automatically
        connected to the virtual bus.

        Parameters
        ----------
        socket_type : ``int | zmq.SocketType``
            The type of the socket to create.
            Accepted values are:

            - ``1`` / ``zmq.PUB``: Publisher socket.
            - ``2`` / ``zmq.SUB``: Subscriber socket.
        topic: ``str | Iterable[str] ``, optional (subscribers only)
            Topic(s) to subscribe to.
            If not provided, the subscriber is in charge
            of specifying the topic.

        Returns
        -------
        socket : ``zmq.SyncSocket | None``
            A new socket for the virtual bus.
            If the socket type is invalid,
            logs the error returns ``None``.
        socket, poller : ``tuple[zmq.SyncSocket, zmq.Poller] | None``
            A new socket and a poller for the virtual bus
            (only for ``zmq.SUB`` sockets).
            If the socket type is invalid,
            returns ``None``.

        """
        if socket_type not in [zmq.PUB, zmq.SUB]:
            logging.error(
                f"Invalid socket type: {zmq.SocketType(socket_type)}. Aborting connection."
            )
            return None
        socket: zmq.SyncSocket = self._context.socket(socket_type)
        socket.connect(self._INPROC_MAP[socket_type])
        socket.setsockopt(zmq.LINGER, 0)

        if socket_type == zmq.SUB:
            poller = zmq.Poller()
            poller.register(socket, zmq.POLLIN)
            if isinstance(topic, str):
                socket.subscribe(topic)
            elif isinstance(topic, Iterable):
                for t in topic:
                    socket.subscribe(t)
            return socket, poller
        else:
            self._PUB_SOCKETS.add(socket)
            return socket

class Publisher:
    def __init__(self, bus: VirtualBus) -> None:
        self.bus = bus
        self.socket = self.bus.connect(zmq.PUB)

    def send(self, msg: str) -> None:
        logging.info(f"Sending message: {msg}")
        self.socket.send_string(msg)

class Subscriber:
    def __init__(self, bus: VirtualBus) -> None:
        self.msg = ""
        self.bus = bus
        self.socket, self.poller = self.bus.connect(zmq.SUB)
        self.socket.subscribe("")

        self.thread = threading.Thread(target=self._polling_thread, daemon=True)
        self.thread.start()

    def _polling_thread(self) -> None:
        try:
            while True:
                try:
                    socks = dict(self.poller.poll())
                    if self.socket in socks:
                        self.msg = self.socket.recv_string()
                        logging.info(f"Received message: {self.msg}")
                except zmq.error.ContextTerminated:
                    break
        finally:
            self.poller.unregister(self.socket)
            self.socket.close()
            logging.info("Subscriber socket closed.")

bus = VirtualBus()
pub = Publisher(bus)
sub = Subscriber(bus)

# give time for
# subscrbier to connect
time.sleep(0.1)

test_msg = "Hello, World!"
pub.send(test_msg)

# give some time
# for message transmission
time.sleep(0.1)

# shutdown the bus;
# this will kill
# all connected
# subscribers
bus.shutdown()

# give time for
# the subscriber to
# finish processing
time.sleep(0.1)

The VirtualBus holds a weak reference to publisher sockets so that they can be manually closed when shutdown is called.

This script works as intended, but it feels a bit unbalanced. Subscriber sockets are automatically closed when zmq.ContextTerminated is raised, and I wanted to know if there's a similar mechanism to also close publisher sockets in a similar fashion so that I don't have to keep references to them.