diff --git a/backend/fluksio/api/routes/flows.py b/backend/fluksio/api/routes/flows.py index f538b25..9c4a9ae 100644 --- a/backend/fluksio/api/routes/flows.py +++ b/backend/fluksio/api/routes/flows.py @@ -1054,14 +1054,25 @@ async def flow_events(websocket: WebSocket, token: str = "") -> None: {sender, receiver}, return_when=asyncio.FIRST_COMPLETED ) if receiver in done: + if sender not in done: + # Nobody will ever read this one: the next turn makes + # its own. Left behind it stays parked on the queue's + # waiter list, takes an event the live reader wanted, + # and is reported as a task destroyed while pending. + sender.cancel() # A frame from the client. Only a disconnect ends the # stream — a keepalive, or anything else it decides to # say, used to be read as the client going away and cost # it every live update from then on. - if receiver.exception() is None: - asked = wanted_names(receiver.result(), only) - if asked is not None: - wanted = asked + failure = receiver.exception() + if failure is not None: + # The client is gone, and `receive_text` raises for + # good from here on. Reading it again would spin this + # loop at full tilt against a dead socket. + raise failure + asked = wanted_names(receiver.result(), only) + if asked is not None: + wanted = asked receiver = asyncio.create_task(websocket.receive_text()) if sender not in done: continue diff --git a/backend/tests/api/routes/test_flows.py b/backend/tests/api/routes/test_flows.py index 548c273..c548aa5 100644 --- a/backend/tests/api/routes/test_flows.py +++ b/backend/tests/api/routes/test_flows.py @@ -1,10 +1,13 @@ +import asyncio from datetime import UTC, datetime import pytest +from fastapi import WebSocketDisconnect from fastapi.testclient import TestClient from sqlalchemy import func from sqlmodel import Session, select +from fluksio.api.routes import flows as flows_route from fluksio.core.config import settings from fluksio.flow.events import event_bus from fluksio.models import Run, RunArtifact, RunMetric, RunNode @@ -493,3 +496,50 @@ def test_renaming_a_flow_leaves_nothing_under_the_old_name( assert [key for key in state.keys() if "movable." in key] == [] client.delete(f"{PREFIX}/moved", headers=superuser_token_headers) + + +@pytest.mark.anyio +async def test_the_live_socket_ends_when_its_client_goes( + client: TestClient, monkeypatch: pytest.MonkeyPatch +) -> None: + """A departing client ends the stream instead of restarting the read. + + `receive_text` raises for good once the peer has gone, so reading it again + spun this loop at full tilt: the event loop stalled, `/health` answered + 503 and one queue getter was orphaned per turn. + """ + monkeypatch.setattr(flows_route, "user_from_token", lambda session, token: object()) + + class Peer: + """One frame, then gone — the way a closed tab leaves.""" + + def __init__(self) -> None: + self.app = client.app + self.frames = ['{"type": "media", "names": []}'] + self.sent: list[str] = [] + + async def accept(self) -> None: + return None + + async def send_text(self, text: str) -> None: + self.sent.append(text) + + async def receive_text(self) -> str: + if self.frames: + return self.frames.pop(0) + raise WebSocketDisconnect(1006) + + peer = Peer() + subscribed = len(event_bus._subscribers) + + await asyncio.wait_for(flows_route.flow_events(peer, token="not-a-jwt"), timeout=5) + + # The snapshot went out, and the frame in between did not end the stream: + # reaching the disconnect at all means the second read happened. + assert peer.sent + assert peer.frames == [] + assert len(event_bus._subscribers) == subscribed + # A turn of the loop for the cancellations the handler asked for. + await asyncio.sleep(0.05) + parked = [t for t in asyncio.all_tasks() if "Queue.get" in repr(t.get_coro())] + assert parked == []