Shell reply sent raw on the ROUTER can drop the ZMQStream wake-up; a request arriving in that window is not read until the next connection

"""Deterministic reproducer: ZMQStream wake-up lost to a raw send on the same socket.

Structure of ipykernel 7's shell channel thread: one IOLoop, a ZMQStream on the
ROUTER (requests in), and a second ZMQStream on an inproc PAIR (replies from the
main thread) whose callback writes the reply to the ROUTER with a RAW
``send_multipart``.

Interleaving forced here: while the PAIR callback is running (i.e. inside one
loop iteration, after the selector already took its readiness snapshot), client
B's request arrives at the ROUTER. libzmq queues an activate-read command and
signals the ROUTER's wake-up fd. The raw send then runs libzmq's
process_commands, which consumes that signal. Nothing re-reads ZMQ_EVENTS on the
ROUTER stream afterwards, so the loop never wakes for B: B's request sits in the
ROUTER until some later command on that socket wakes the stream again.

usage: python zmq_lost_wakeup_det.py [--fix] [--settle=SECONDS]
  --settle  time to wait inside the callback before the raw send, must exceed
            libzmq's max_command_delay (~1 ms on x86 TSC, ~125 ms on Apple Silicon)
"""
import select, sys, threading, time
import zmq
from tornado.ioloop import IOLoop
from zmq.eventloop.zmqstream import ZMQStream

FIX = "--fix" in sys.argv
SETTLE = float(next((a.split("=")[1] for a in sys.argv if a.startswith("--settle=")), 0.2))
CLOSE_A = "--close-a" in sys.argv
TRIALS = int(next((a.split("=")[1] for a in sys.argv if a.startswith("--trials=")), 5))

ctx = zmq.Context()
router = ctx.socket(zmq.ROUTER)
port = router.bind_to_random_port("tcp://127.0.0.1")
to_main_a, to_main_b = ctx.socket(zmq.PAIR), ctx.socket(zmq.PAIR)
to_main_a.bind("inproc://to-main"); to_main_b.connect("inproc://to-main")
from_main_a, from_main_b = ctx.socket(zmq.PAIR), ctx.socket(zmq.PAIR)
from_main_a.bind("inproc://from-main"); from_main_b.connect("inproc://from-main")

loop = IOLoop(make_current=False)
router_stream = ZMQStream(router, loop)
from_main_stream = ZMQStream(from_main_a, loop)

b_go = threading.Event()
b_sent = threading.Event()

def on_request(frames):
    to_main_a.send_multipart(frames)

def on_reply(frames):
    if frames[-1].startswith(b"A"):
        # A's reply is about to go out on the raw socket. Let B's request land first,
        # inside this same loop iteration, then wait out libzmq's send-side throttle.
        b_go.set(); b_sent.wait(); time.sleep(SETTLE)
        fd = router.fd
        print(f"    fd readable before raw send: {bool(select.select([fd], [], [], 0)[0])}")
    router.send_multipart(frames)          # raw send, as ipykernel 7 does
    if frames[-1].startswith(b"A"):
        print(f"    fd readable after raw send:  {bool(select.select([fd], [], [], 0)[0])}")
    if FIX:
        router_stream._rebuild_io_state()  # re-read ZMQ_EVENTS after touching the socket raw

router_stream.on_recv(on_request)
from_main_stream.on_recv(on_reply)
threading.Thread(target=loop.start, daemon=True, name="shell-channel").start()

def main_thread():
    while True:
        from_main_b.send_multipart(to_main_b.recv_multipart())
threading.Thread(target=main_thread, daemon=True, name="main-shell").start()

def dealer():
    s = ctx.socket(zmq.DEALER); s.linger = 0; s.connect(f"tcp://127.0.0.1:{port}"); return s

lost = 0
for trial in range(TRIALS):
    b_go.clear(); b_sent.clear()
    a, b = dealer(), dealer()
    time.sleep(0.05)                                   # both DEALERs connected
    a.send_multipart([b"", b"A-%d" % trial])
    if CLOSE_A:
        time.sleep(0.02); a.close()   # like jupyter_server's nudge: the transient channel is gone before the kernel replies
    b_go.wait()
    b.send_multipart([b"", b"B-%d" % trial]); time.sleep(0.02); b_sent.set()
    if not CLOSE_A:
        assert a.poll(2000), "A itself never answered"
        a.recv_multipart()
    if b.poll(1000):
        b.recv_multipart(); print(f"trial {trial}: B answered")
    else:
        lost += 1; print(f"trial {trial}: B LOST (stuck in the ROUTER)")
        c = dealer(); c.send_multipart([b"", b"C-%d" % trial])   # any new command wakes the stream
        print(f"    after a third client's request: B answered={bool(b.poll(1000))} C answered={bool(c.poll(1000))}")
        c.close()
    if not CLOSE_A: a.close()
    b.close()
print(f"fix={FIX} settle={SETTLE} trials={TRIALS} lost={lost}")