使用ZeroMQ时，如何等待SUB套接字订阅在PUB端完成注册？

我实现了一个消息总线，功能是把PULL套接字收到的消息转发到PUB套接字，作为系统内其他组件的通信媒介——这些组件用PUSH发消息、SUB收消息。总线的主要通信在守护进程之间，但也存在临时程序连接总线后立即发请求、等回复就断开的场景。

问题在于SUB套接字的订阅是非阻塞操作，可能订阅还没在消息总线的PUB端完成注册，回复就已经发出去了。我写了个示例：启动20个ephemeral_program临时程序，它们连接总线、发消息并期望收到自己的消息，但部分能收到，部分收不到。

目前有两种不太理想的解决办法：

• 发消息前短暂休眠，这种方法不可靠
• 循环发消息直到收到回复来检查订阅状态，本质是轮询，不是发布者主动的就绪确认

想知道有没有其他套接字选项或一对多的套接字组合，能确保订阅者侧确认发布者已完成订阅注册？

import time
from multiprocessing import Process

from zmq import Context, Poller, Socket
from zmq.constants import POLLIN, PUB, PULL, PUSH, SUB


def message_bus():
    context = Context.instance()
    in_socket: Socket = context.socket(PULL)
    in_socket.bind("ipc:///tmp/in-socket.ipc")
    out_socket: Socket = context.socket(PUB)
    out_socket.bind("ipc:///tmp/out-socket.ipc")

    while True:
        msg = in_socket.recv_string()
        out_socket.send_string(msg)


# 生成测试流量的守护进程
def daemon():
    context = Context.instance()
    receive_socket: Socket = context.socket(SUB)
    receive_socket.connect("ipc:///tmp/out-socket.ipc")
    receive_socket.subscribe("")
    send_socket: Socket = context.socket(PUSH)
    send_socket.connect("ipc:///tmp/in-socket.ipc")

    poller = Poller()
    poller.register(receive_socket, POLLIN)

    while True:
        polling_result = poller.poll(timeout=1)

        if polling_result:
            receive_socket.recv_string()
        else:
            send_socket.send_string("Hi!")


def ephemeral_program(i):
    context = Context.instance()

    send_socket: Socket = context.socket(PUSH)
    send_socket.connect("ipc:///tmp/in-socket.ipc")

    receive_socket: Socket = context.socket(SUB)
    receive_socket.connect("ipc:///tmp/out-socket.ipc")
    receive_socket.subscribe(f"ephemeral_program_{i}")

    poller = Poller()
    poller.register(receive_socket, POLLIN)

    # time.sleep(1)
    # 这种方法大多时候能给足够时间让订阅在PUB端完成注册，但不可靠

    # 相对好一点的办法是循环发送消息确认订阅
    # while True:
    #    send_socket.send_string(f"ephemeral_program_{i}:am I subscribed?")
    #    polling_result = poller.poll(timeout=1)
    #
    #    if polling_result:
    #        msg = receive_socket.recv_string()
    #        if msg == f"ephemeral_program_{i}:am I subscribed?":
    #            print("Subscription confirmed")
    #            break
    #    else:
    #        print("Not yet subscribed")

    # 发送实际请求
    send_socket.send_string(f"ephemeral_program_{i}:Hello")
    polling_result = poller.poll(timeout=100)

    if polling_result:
        print(receive_socket.recv_string())
    else:
        print("Did not receive message")


message_bus_p = Process(target=message_bus, args=())
message_bus_p.start()

n = 20
daemons_p = [Process(target=daemon, args=()) for _ in range(n)]
for daemon_p in daemons_p:
    daemon_p.start()

# 给消息总线和守护进程启动时间，模拟已运行一段时间的状态
time.sleep(0.2)

ephemeral_programs_p = [
    Process(target=ephemeral_program, args=(i,)) for i in range(n)
]
for ephemeral_program_p in ephemeral_programs_p:
    ephemeral_program_p.start()
for ephemeral_program_p in ephemeral_programs_p:
    ephemeral_program_p.join()

for daemon_p in daemons_p:
    daemon_p.kill()

message_bus_p.kill()

1. 使用XPUB/XSUB替代PUB/SUB #

ZeroMQ的XPUB套接字会主动发送订阅/取消订阅事件给连接的客户端。你可以把消息总线的PUB换成XPUB，然后在总线里监听这些事件：

• 当XPUB收到某个临时程序的订阅事件时，总线可以通过一个单独的REQ/REP通道通知该程序订阅已完成注册
• 临时程序连接总线时，先通过REQ套接字等待订阅确认，确认后再发送实际消息

修改后的消息总线核心逻辑：

def message_bus():
    context = Context.instance()
    # 原有PULL接收消息
    in_socket = context.socket(PULL)
    in_socket.bind("ipc:///tmp/in-socket.ipc")
    # 替换为XPUB
    out_socket = context.socket(XPUB)
    out_socket.bind("ipc:///tmp/out-socket.ipc")
    # 新增REQ/REP用于订阅确认
    confirm_socket = context.socket(REP)
    confirm_socket.bind("ipc:///tmp/confirm-socket.ipc")

    poller = Poller()
    poller.register(in_socket, POLLIN)
    poller.register(out_socket, POLLIN)
    poller.register(confirm_socket, POLLIN)

    subscribed_clients = set()

    while True:
        socks = dict(poller.poll())
        if in_socket in socks:
            msg = in_socket.recv_string()
            out_socket.send_string(msg)
        if out_socket in socks:
            # XPUB收到的订阅事件：第一个字节是0(取消)或1(订阅)，后面是订阅主题
            event, topic = out_socket.recv_multipart()
            if event == b'\x01':
                subscribed_clients.add(topic.decode())
        if confirm_socket in socks:
            client_topic = confirm_socket.recv_string()
            if client_topic in subscribed_clients:
                confirm_socket.send_string("READY")
            else:
                confirm_socket.send_string("WAIT")

临时程序修改：

def ephemeral_program(i):
    context = Context.instance()
    topic = f"ephemeral_program_{i}"

    # 订阅确认用REQ
    confirm_socket = context.socket(REQ)
    confirm_socket.connect("ipc:///tmp/confirm-socket.ipc")

    # 原有PUSH发消息、SUB收消息
    send_socket = context.socket(PUSH)
    send_socket.connect("ipc:///tmp/in-socket.ipc")
    receive_socket = context.socket(SUB)
    receive_socket.connect("ipc:///tmp/out-socket.ipc")
    receive_socket.subscribe(topic)

    # 等待订阅确认
    while True:
        confirm_socket.send_string(topic)
        reply = confirm_socket.recv_string()
        if reply == "READY":
            break
        time.sleep(0.01)

    # 发送实际消息
    send_socket.send_string(f"{topic}:Hello")
    poller = Poller()
    poller.register(receive_socket, POLLIN)
    polling_result = poller.poll(timeout=100)
    if polling_result:
        print(receive_socket.recv_string())
    else:
        print("Did not receive message")

2. 给临时程序分配专属REQ/REP通道 #

如果临时程序的请求-响应模式是一对一的，可以直接在消息总线中新增REQ/REP套接字，临时程序用REQ发送请求，总线用REP接收后处理并回复，这样天然保证了消息的可靠传递，不需要依赖SUB/PUB的订阅确认。但这种方式不适用于一对多的广播场景。

3. 使用ZMQ_SUBSCRIBE_NOTIFY选项（部分版本支持） #

部分ZeroMQ版本支持ZMQ_SUBSCRIBE_NOTIFY套接字选项，当订阅在PUB端完成注册时，会给SUB套接字发送一个通知事件。你可以在SUB套接字上设置这个选项，然后通过Poller监听该事件，收到通知后再发送消息。

示例代码片段：

receive_socket = context.socket(SUB)
receive_socket.connect("ipc:///tmp/out-socket.ipc")
receive_socket.subscribe(topic)
# 设置订阅通知选项
receive_socket.setsockopt(zmq.SUBSCRIBE_NOTIFY, 1)

poller = Poller()
poller.register(receive_socket, POLLIN | zmq.POLLSUB)

# 等待订阅通知
while True:
    socks = dict(poller.poll())
    if receive_socket in socks and socks[receive_socket] & zmq.POLLSUB:
        # 订阅已完成注册
        break

注意：这个选项不是所有版本都支持，需要确认你的ZeroMQ版本是否包含该特性。

内容的提问来源于stack exchange，提问作者Afkaaja