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:
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user