ZeroMQ Pub/Sub Message Broker

ZeroMQ Pub/Sub Message Broker

Problem Statement

Modern distributed systems frequently need to broadcast events (metrics, logs, cache invalidations, configuration changes) from one producer to many consumers without coupling the producer to any specific consumer. Traditional HTTP-based fan-out requires the publisher to know every subscriber, manage retries, and handle backpressure — all of which add latency and operational complexity. Similarly, polling-based approaches waste CPU and introduce delay.

ZeroMQ's Pub/Sub pattern solves this elegantly: a publisher binds to a socket and fires messages into the void, while subscribers connect and filter by topic prefix. There is no handshake, no acknowledgement, and no broker process to operate — ZeroMQ handles the transport, reconnection, and multiplexing in user space. This script gives you a production-ready, configurable publisher and subscriber in a single file so you can spin up a real-time event bus in seconds.

How It Works

The script runs in one of two modes selected by --mode. In publisher mode it creates a zmq.PUB socket, binds it to the given endpoint, and then enters a loop that sends topic-prefixed messages at the requested rate. Messages are JSON-encoded by default so downstream consumers get structured data, but a --format raw option is available for plain text. The publisher uses zmq.SNDHWM (send high-water mark) to bound the outbound queue and prevent unbounded memory growth if subscribers are slow — ZeroMQ drops messages for slow subscribers rather than blocking the publisher, which is the correct behavior for real-time telemetry.

In subscriber mode the script creates a zmq.SUB socket, connects to the endpoint, and subscribes to the requested topic prefix. ZeroMQ's topic filtering happens on the subscriber side: the first frame of every message must begin with the subscribed prefix or the message is discarded before it reaches your code. The subscriber loop uses zmq.Poller with a timeout so it can shut down cleanly on Ctrl+C and optionally exit after a fixed number of messages via --count. Both modes install a SIGINT/SIGTERM handler that flips a threading.Event, letting the main loop finish its current iteration and close the socket gracefully — which matters because abrupt termination of a PUB socket can drop in-flight messages.

Full Script

#!/usr/bin/env python3
"""
ZeroMQ Pub/Sub Message Broker

A single-file, configurable publisher and subscriber built on pyzmq.
Supports JSON or raw payloads, topic prefix filtering, rate limiting,
graceful shutdown, and bounded send queues.

Usage:
    python3 zmq_pubsub.py --mode publisher --endpoint tcp://*:5555 \
        --topic metrics --rate 100
    python3 zmq_pubsub.py --mode subscriber --endpoint tcp://127.0.0.1:5555 \
        --topic metrics
"""

from __future__ import annotations

import argparse
import json
import os
import signal
import sys
import threading
import time
from typing import Any

import zmq

# ---------------------------------------------------------------------------
# Configuration (environment variables with sensible defaults)
# ---------------------------------------------------------------------------
DEFAULT_ENDPOINT = os.environ.get("ZMQ_ENDPOINT", "tcp://127.0.0.1:5555")
DEFAULT_TOPIC = os.environ.get("ZMQ_TOPIC", "metrics")
DEFAULT_RATE = float(os.environ.get("ZMQ_RATE", "10"))          # msgs/sec
DEFAULT_HWM = int(os.environ.get("ZMQ_SNDHWM", "1000"))         # send queue
DEFAULT_FORMAT = os.environ.get("ZMQ_FORMAT", "json")           # json|raw
DEFAULT_COUNT = int(os.environ.get("ZMQ_COUNT", "0"))           # 0 = forever
DEFAULT_LINGER = int(os.environ.get("ZMQ_LINGER", "1000"))      # ms

# Shared shutdown flag — flipped by signal handlers in both modes.
_shutdown = threading.Event()


def _install_signal_handlers() -> None:
    """Translate SIGINT/SIGTERM into a cooperative shutdown flag."""

    def _handler(signum: int, _frame: Any) -> None:
        print(f"\n[info] received signal {signum}, shutting down...",
              file=sys.stderr)
        _shutdown.set()

    signal.signal(signal.SIGINT, _handler)
    signal.signal(signal.SIGTERM, _handler)


def _build_payload(topic: str, seq: int, fmt: str) -> bytes:
    """Return a topic-prefixed, wire-ready frame.

    ZeroMQ PUB/SUB filtering operates on the *first frame*. We embed the
    topic as a text prefix followed by a space, so subscribers can subscribe
    to 'metrics' and match 'metrics {...}'.
    """
    if fmt == "json":
        body = json.dumps(
            {
                "seq": seq,
                "topic": topic,
                "ts": time.time(),
                "value": (seq * 7) % 100,
            },
            separators=(",", ":"),
        )
    else:
        body = f"seq={seq} ts={time.time():.3f}"
    return f"{topic} {body}".encode("utf-8")


def run_publisher(endpoint: str, topic: str, rate: float, fmt: str,
                  hwm: int, linger: int, count: int) -> int:
    """Bind a PUB socket and emit topic-prefixed messages at `rate` Hz."""
    ctx = zmq.Context.instance()
    sock = ctx.socket(zmq.PUB)
    # Bound the outbound queue: slow subscribers get dropped, publisher never
    # blocks. This is the correct trade-off for telemetry-style workloads.
    sock.setsockopt(zmq.SNDHWM, hwm)
    # LINGER controls how long close() waits for pending messages to flush.
    sock.setsockopt(zmq.LINGER, linger)

    try:
        sock.bind(endpoint)
    except zmq.ZMQError as exc:
        print(f"[error] cannot bind {endpoint}: {exc}", file=sys.stderr)
        sock.close(0)
        return 2

    print(f"[pub] bound {endpoint} topic='{topic}' rate={rate}/s "
          f"format={fmt} hwm={hwm}", file=sys.stderr)

    # Give subscribers a moment to connect before the first send. Without
    # this, the first few messages are silently dropped (classic PUB/SUB
    # slow-joiner problem).
    time.sleep(0.5)

    interval = 1.0 / rate if rate > 0 else 0.0
    seq = 0
    next_tick = time.monotonic()

    while not _shutdown.is_set():
        seq += 1
        try:
            sock.send(_build_payload(topic, seq, fmt))
        except zmq.ZMQError as exc:
            print(f"[error] send failed: {exc}", file=sys.stderr)
            break

        if count and seq >= count:
            break

        if interval:
            next_tick += interval
            sleep_for = next_tick - time.monotonic()
            if sleep_for > 0:
                # _shutdown.wait() lets Ctrl+C interrupt the sleep promptly.
                _shutdown.wait(sleep_for)
            else:
                # We're behind schedule; reset the clock to avoid drift spiral.
                next_tick = time.monotonic()

    print(f"[pub] sent {seq} messages, closing.", file=sys.stderr)
    sock.close(linger)
    return 0


def run_subscriber(endpoint: str, topic: str, count: int) -> int:
    """Connect a SUB socket, filter by topic prefix, and print messages."""
    ctx = zmq.Context.instance()
    sock = ctx.socket(zmq.SUB)
    sock.setsockopt(zmq.LINGER, 0)
    # Subscribe to the topic prefix. Empty string subscribes to everything.
    sock.setsockopt_string(zmq.SUBSCRIBE, topic)

    try:
        sock.connect(endpoint)
    except zmq.ZMQError as exc:
        print(f"[error] cannot connect to {endpoint}: {exc}", file=sys.stderr)
        sock.close(0)
        return 2

    print(f"[sub] connected {endpoint} topic='{topic}'", file=sys.stderr)

    poller = zmq.Poller()
    poller.register(sock, zmq.POLLIN)
    received = 0

    while not _shutdown.is_set():
        try:
            # 250 ms timeout keeps the loop responsive to the shutdown flag.
            events = dict(poller.poll(timeout=250))
        except zmq.ZMQError as exc:
            print(f"[error] poll failed: {exc}", file=sys.stderr)
            break

        if sock not in events:
            continue

        try:
            frame = sock.recv_string()
        except zmq.ZMQError as exc:
            print(f"[error] recv failed: {exc}", file=sys.stderr)
            break

        received += 1
        print(f"[{received}] {frame}", flush=True)

        if count and received >= count:
            break

    print(f"[sub] received {received} messages, closing.", file=sys.stderr)
    sock.close(0)
    return 0


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="ZeroMQ Pub/Sub message broker (publisher or subscriber).",
    )
    parser.add_argument("--mode", choices=("publisher", "subscriber"),
                        required=True, help="Run as publisher or subscriber.")
    parser.add_argument("--endpoint", default=DEFAULT_ENDPOINT,
                        help="ZMQ endpoint, e.g. tcp://*:5555 or "
                             "tcp://127.0.0.1:5555")
    parser.add_argument("--topic", default=DEFAULT_TOPIC,
                        help="Topic prefix used for filtering.")
    parser.add_argument("--rate", type=float, default=DEFAULT_RATE,
                        help="Publisher send rate in messages/second.")
    parser.add_argument("--format", choices=("json", "raw"),
                        default=DEFAULT_FORMAT, help="Payload encoding.")
    parser.add_argument("--sndhwm", type=int, default=DEFAULT_HWM,
                        help="Publisher send high-water mark.")
    parser.add_argument("--linger", type=int, default=DEFAULT_LINGER,
                        help="Socket linger in milliseconds.")
    parser.add_argument("--count", type=int, default=DEFAULT_COUNT,
                        help="Stop after N messages (0 = run forever).")
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    _install_signal_handlers()

    if args.mode == "publisher":
        return run_publisher(args.endpoint, args.topic, args.rate,
                             args.format, args.sndhwm, args.linger,
                             args.count)
    return run_subscriber(args.endpoint, args.topic, args.count)


if __name__ == "__main__":
    sys.exit(main())

Configuration

CLI flags override environment variables, which override hard-coded defaults.

Usage Examples

1. Start a publisher at 100 msg/s on port 5555:

$ python3 zmq_pubsub.py --mode publisher --endpoint tcp://*:5555 \
    --topic metrics --rate 100
[pub] bound tcp://*:5555 topic='metrics' rate=100.0/s format=json hwm=1000

2. In another terminal, subscribe and print 5 messages:

$ python3 zmq_pubsub.py --mode subscriber \
    --endpoint tcp://127.0.0.1:5555 --topic metrics --count 5
[sub] connected tcp://127.0.0.1:5555 topic='metrics'
[1] metrics {"seq":1,"topic":"metrics","ts":1712345678.12,"value":7}
[2] metrics {"seq":2,"topic":"metrics","ts":1712345678.13,"value":14}
[3] metrics {"seq":3,"topic":"metrics","ts":1712345678.14,"value":21}
[4] metrics {"seq":4,"topic":"metrics","ts":1712345678.15,"value":28}
[5] metrics {"seq":5,"topic":"metrics","ts":1712345678.16,"value":35}
[sub] received 5 messages, closing.

3. High-throughput raw-text stream on an IPC socket (Linux/macOS):

$ python3 zmq_pubsub.py --mode publisher --endpoint ipc:///tmp/events.sock \
    --topic logs --rate 0 --format raw
[pub] bound ipc:///tmp/events.sock topic='logs' rate=0/s format=raw hwm=1000

Error Handling

• Bind/connect failures (zmq.ZMQError) — reported to stderr with the offending endpoint and a non-zero exit code, so systemd or a container orchestrator will restart the unit.
• Slow subscribers — SNDHWM caps the outbound queue. When it fills, ZeroMQ drops messages for the slow subscriber instead of blocking the publisher. This is intentional for telemetry; raise --sndhwm if you need lossless delivery (and accept the memory cost).
• Slow-joiner problem — the publisher sleeps 500 ms after bind() so subscribers have time to connect before the first send. Without this, the initial messages are dropped silently.
• Graceful shutdown — SIGINT/SIGTERM set a threading.Event; the main loop checks it every iteration and the subscriber's poll() uses a 250 ms timeout so Ctrl+C is responsive. LINGER ensures in-flight messages are flushed before the socket closes.
• Clock drift under load — the publisher uses a monotonic next_tick accumulator rather than sleep(interval), so a slow send doesn't compound into an ever-growing lag. If the loop falls behind, it resets the clock instead of trying to "catch up" with a burst.
• Malformed JSON on the wire — subscribers print the raw frame unconditionally, so a bad publisher payload is visible rather than swallowed. Downstream consumers can json.loads() and handle parse errors as they see fit.

Extending Ideas

• Add a zmq.PROXY broker — insert an XPUB/XSUB forwarder between publisher and subscribers to enable many-to-many topologies, dynamic subscriptions, and a single well-known endpoint.
• Swap JSON for MessagePack or Protobuf — replace _build_payload with a serializer of your choice. MessagePack gives you ~30% smaller frames; Protobuf gives you schema enforcement.
• Prometheus metrics — track messages_sent_total, messages_received_total, and send_drops_total (via zmq.ZMQ_EVENTS) and expose them on /metrics for observability.
• Authentication and encryption — layer zmq.CURVE (sock.curve_server = True on the publisher, curve_serverkey on the subscriber) to encrypt the wire and authenticate peers; ZeroMQ ships this out of the box with libsodium.
• Multi-topic fan-out — call setsockopt_string(zmq.SUBSCRIBE, prefix) multiple times in the subscriber to match several topics on one socket, then dispatch by parsing the prefix out of each frame.