fix(workers): send heartbeats on an inbound session

Found on hardware. The Attach handler started the read pump and the outbound
loop but never the heartbeat loop that the outbound path starts inside
_connect_once. So a node registered, went silent, was declared dead about
ninety seconds later, reconnected, and flapped forever — and in between, work
aimed at it fell back to the local machine with 'gpu2 is offline', while the
panel had shown it ready at 3.4 ms moments earlier.

Every end-to-end test in this file finished inside three seconds, comfortably
within the grace window that hid it. The regression test therefore asserts on
the emitted heartbeat frames themselves rather than on liveness, and shortens
the advertised interval so it does that in two seconds instead of twenty.
This commit is contained in:
velixio
2026-08-11 23:05:20 +05:30
parent e121e69d0f
commit 569517e5d8
3 changed files with 74 additions and 3 deletions
+7 -3
View File
@@ -109,6 +109,7 @@ class NodeServicer(pb_grpc.NodeServiceServicer):
# client would make the signature match at most one of them.
client = self._client_factory(self._artifacts, key_id)
reader: Optional[asyncio.Task] = None
heartbeat: Optional[asyncio.Task] = None
try:
# The node speaks first even though the panel dialled: it is still
# the side with capabilities to declare, and the panel cannot
@@ -124,6 +125,7 @@ class NodeServicer(pb_grpc.NodeServiceServicer):
return
await client.accept_registration(first.registered)
heartbeat = client.start_heartbeat(first.registered)
reader = asyncio.create_task(
self._pump_incoming(client, request_iterator, session_id)
)
@@ -147,10 +149,12 @@ class NodeServicer(pb_grpc.NodeServiceServicer):
except Exception as exc:
logger.warning("Inbound session from %s ended: %s", peer or "a panel", exc)
finally:
if reader is not None:
reader.cancel()
for task in (reader, heartbeat):
if task is None:
continue
task.cancel()
with contextlib.suppress(asyncio.CancelledError, Exception):
await reader
await task
await client.stop()
self._log.closed(session_id)
+18
View File
@@ -423,6 +423,24 @@ class WorkerClient:
"""The next frame this worker wants to send."""
return await self._outbox.get()
def start_heartbeat(self, response: pb.RegisterResponse) -> asyncio.Task:
"""Begin the heartbeat this session's liveness depends on.
Separate from `accept_registration` because the task has to live and
die with the stream, not with the registration. Outbound starts the
same loop inside `_connect_once`; inbound has no such place, and
leaving it out is invisible for exactly as long as the grace window —
which is why it survived every sub-second test and only showed up on
hardware, as a worker that registered, went quiet, was declared dead
~90s later, reconnected, and flapped forever.
"""
return asyncio.create_task(
self._heartbeat_loop(
response.heartbeat_interval_seconds or _HEARTBEAT_SECONDS
),
name="inbound-heartbeat",
)
async def handle_server_message(self, message: pb.ServerMessage) -> None:
await self._on_server_message(message)
+49
View File
@@ -341,3 +341,52 @@ async def test_a_different_machine_cannot_re_adopt_an_enrolled_workers_identity(
)
assert NodeConnection._proves_key_possession(forged, enrolled) is False
@pytest.mark.asyncio
async def test_the_node_keeps_sending_heartbeats_after_it_registers(inbound, monkeypatch):
"""A session that goes quiet is declared dead and flaps forever.
Found on hardware, not here: the inbound Attach handler started the read
pump and the outbound loop but never the heartbeat loop that the outbound
path starts in `_connect_once`. The node registered, said nothing more, was
declared dead ~90 seconds later, reconnected, and repeated — while every
test in this file finished inside three seconds, comfortably within the
grace window that hid it.
So this test asserts on the frames themselves rather than on liveness: it
watches the node's own outbox for a heartbeat, which is the thing that was
missing, and does not depend on how long the grace window happens to be.
"""
# The interval the panel advertises, shortened so this asserts on a real
# emitted frame in a second rather than waiting out the production value.
from worker.transport import server as server_module
monkeypatch.setattr(server_module, "_HEARTBEAT_INTERVAL_SECONDS", 1)
seen = []
client_box = {}
original = inbound._client
def capture(artifacts, key_id):
client = original(artifacts, key_id)
client_box["client"] = client
real_send = client._send
async def spy(message, **kwargs):
if message.WhichOneof("payload") == "heartbeat":
seen.append(message)
return await real_send(message, **kwargs)
client._send = spy
return client
inbound.listener._servicer._client_factory = capture
await inbound.connect_panel()
# Drive the loop rather than waiting out a real interval: the bug is a
# missing task, not a slow one, so what matters is that something is
# scheduled to produce these at all.
await _until(lambda: len(seen) >= 2, timeout=15.0)
assert len(seen) >= 2, "the node registered and then never sent a heartbeat"