From bad9f584e6ee34403bf9ccd1d772ee3af2f60748 Mon Sep 17 00:00:00 2001 From: Yao Lu Date: Mon, 6 Jul 2026 20:16:31 -0700 Subject: [PATCH] fix: keep SUBSCRIBE alive by polling get_message instead of listen() The dispatch/command/response listeners consume via iter_pubsub_messages, which looped over pubsub.listen(). listen() parks inside a single parse_response for the whole idle lifetime of a subscription, so redis-py's health_check_interval (set alongside socket_keepalive) fires at most once and no application traffic flows on a quiet channel. An idle proxy/LB then culls the SUBSCRIBE connection (TCP keepalive ACKs don't reset an L7 inactivity timer) and the node silently goes dispatch-dead ('No subscribers on tasks channel'), which no amount of socket_keepalive config could prevent. Consume with get_message(timeout=1.0) instead: each call runs check_health(), emitting the periodic PING that keeps the connection warm and the subscription registered. Fixes the chronic ~1h-after-restart fleet dispatch death. Signed-off-by: Yao Lu --- src/server/clients/redis.py | 19 +++++++- tests/server/test_iter_pubsub_messages.py | 59 +++++++++++++++++++++++ 2 files changed, 76 insertions(+), 2 deletions(-) create mode 100644 tests/server/test_iter_pubsub_messages.py diff --git a/src/server/clients/redis.py b/src/server/clients/redis.py index 86136b9..1ef5dbe 100644 --- a/src/server/clients/redis.py +++ b/src/server/clients/redis.py @@ -148,15 +148,30 @@ def ssh_connection_key(connection_id: str) -> str: SSH_CONNECTION_IDS_KEY = "ssh:connections:active" -def iter_pubsub_messages(pubsub: PubSub) -> Iterable[Any]: +def iter_pubsub_messages(pubsub: PubSub, poll_timeout: float = 1.0) -> Iterable[Any]: """Iterate over messages from a Redis PubSub instance. + Consumes with ``get_message(timeout=...)`` rather than ``listen()``. A + blocking ``listen()`` parks inside a single ``parse_response`` for the entire + idle lifetime of the subscription, so redis-py's ``health_check_interval`` + fires at most once and no application traffic flows on a quiet channel — an + idle proxy/load-balancer then culls the SUBSCRIBE connection (TCP keepalive + ACKs don't reset an L7 inactivity timer), and the node silently stops + receiving dispatches. Polling re-enters ``get_message`` every ``poll_timeout`` + seconds, and each call runs ``check_health()``, which emits the periodic PING + that keeps the connection warm and the subscription registered. + Stops cleanly on pubsub teardown (`ConnectionError`, `OSError`); skips individual malformed JSON payloads without ending iteration so a single bad message can't kill the listener. """ try: - for msg in pubsub.listen(): + while True: + msg = pubsub.get_message(timeout=poll_timeout) + if msg is None: + # Idle tick — get_message already drove check_health(); keep looping + # so the health-check PING is emitted on schedule. + continue if msg.get("type") != "message": continue raw = msg.get("data") diff --git a/tests/server/test_iter_pubsub_messages.py b/tests/server/test_iter_pubsub_messages.py new file mode 100644 index 0000000..8b30cae --- /dev/null +++ b/tests/server/test_iter_pubsub_messages.py @@ -0,0 +1,59 @@ +"""iter_pubsub_messages must poll (get_message) rather than block (listen()). + +A blocking listen() never re-drives redis-py's health_check_interval on an idle +subscription, so the SUBSCRIBE connection sends no traffic and gets culled by an +idle proxy/LB timeout. Polling with get_message(timeout=...) runs check_health() +each iteration, emitting the periodic PING that keeps the connection warm. +""" + +import json + +from server.clients.redis import iter_pubsub_messages + + +class _FakePubSub: + def __init__(self, script: list) -> None: + self._script = list(script) + self.get_message_timeouts: list[float] = [] + self.listen_called = False + + def get_message(self, timeout: float | None = None): + self.get_message_timeouts.append(timeout) + if not self._script: + raise ConnectionError("drained") + item = self._script.pop(0) + if isinstance(item, Exception): + raise item + return item + + def listen(self): # pragma: no cover - must never be used + self.listen_called = True + raise AssertionError("iter_pubsub_messages must not use blocking listen()") + + +def test_polls_with_timeout_and_yields_messages() -> None: + ps = _FakePubSub( + [ + None, # idle tick — must NOT end iteration + {"type": "subscribe", "data": 1}, # non-message — skipped + {"type": "message", "data": json.dumps({"a": 1})}, # yielded + {"type": "message", "data": b"not-json"}, # bad payload — skipped + {"type": "message", "data": json.dumps({"b": 2})}, # yielded + ConnectionError("culled"), # clean stop + ] + ) + out = list(iter_pubsub_messages(ps, poll_timeout=0.01)) + + assert out == [{"a": 1}, {"b": 2}] + assert ps.listen_called is False + # Every read went through get_message with the poll timeout (drives check_health). + assert ps.get_message_timeouts and all(t == 0.01 for t in ps.get_message_timeouts) + + +def test_idle_only_does_not_terminate_until_disconnect() -> None: + # A long stretch of idle ticks keeps iterating (each tick emits the PING); + # only a connection error ends it. + ps = _FakePubSub([None] * 5 + [ConnectionError("bye")]) + out = list(iter_pubsub_messages(ps, poll_timeout=0.01)) + assert out == [] + assert len(ps.get_message_timeouts) >= 5