fix(ws): idempotent connection removal, release slot on any error

The status WebSocket pool removed connections with list.remove(), which
raises ValueError on a double-remove (broadcast already dropped a dead
socket, then disconnect tries again), and only released a slot on
WebSocketDisconnect — any other error leaked the socket into the
broadcast pool. Centralise removal in an idempotent _drop() called from
a finally block and from _broadcast.

ha-relevant: yes
This commit is contained in:
Pouzor
2026-06-10 09:58:06 +02:00
parent d032103b7d
commit 2db173cd56
2 changed files with 80 additions and 3 deletions
+13 -2
View File
@@ -1,3 +1,4 @@
import contextlib
import json import json
from fastapi import APIRouter, WebSocket, WebSocketDisconnect from fastapi import APIRouter, WebSocket, WebSocketDisconnect
@@ -10,6 +11,12 @@ router = APIRouter()
_connections: list[WebSocket] = [] _connections: list[WebSocket] = []
def _drop(websocket: WebSocket) -> None:
"""Remove a connection if still present — idempotent, never raises."""
with contextlib.suppress(ValueError):
_connections.remove(websocket)
@router.websocket("/ws/status") @router.websocket("/ws/status")
async def ws_status(websocket: WebSocket) -> None: async def ws_status(websocket: WebSocket) -> None:
# Accept first so we can send a close frame with a reason code # Accept first so we can send a close frame with a reason code
@@ -33,7 +40,11 @@ async def ws_status(websocket: WebSocket) -> None:
while True: while True:
await websocket.receive_text() await websocket.receive_text()
except WebSocketDisconnect: except WebSocketDisconnect:
_connections.remove(websocket) pass
finally:
# Any error (disconnect or otherwise) must release the slot, else the
# dead socket lingers in the broadcast pool.
_drop(websocket)
async def _broadcast(payload: str) -> None: async def _broadcast(payload: str) -> None:
@@ -41,7 +52,7 @@ async def _broadcast(payload: str) -> None:
try: try:
await conn.send_text(payload) await conn.send_text(payload)
except Exception: except Exception:
_connections.remove(conn) _drop(conn)
async def broadcast_status(node_id: str, status: str, checked_at: str, response_time_ms: int | None = None) -> None: async def broadcast_status(node_id: str, status: str, checked_at: str, response_time_ms: int | None = None) -> None:
+67 -1
View File
@@ -5,7 +5,13 @@ import pytest
from fastapi.testclient import TestClient from fastapi.testclient import TestClient
from starlette.websockets import WebSocketDisconnect from starlette.websockets import WebSocketDisconnect
from app.api.routes.status import _connections, broadcast_scan_update, broadcast_status from app.api.routes.status import (
_connections,
_drop,
broadcast_scan_update,
broadcast_service_status,
broadcast_status,
)
from app.main import app from app.main import app
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -155,3 +161,63 @@ async def test_broadcast_no_connections():
assert len(_connections) == 0 assert len(_connections) == 0
await broadcast_status(node_id="n", status="online", checked_at="t") await broadcast_status(node_id="n", status="online", checked_at="t")
await broadcast_scan_update(run_id="r", devices_found=0) await broadcast_scan_update(run_id="r", devices_found=0)
# ---------------------------------------------------------------------------
# broadcast_service_status
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_broadcast_service_status_payload():
received: list[str] = []
class FakeWS:
async def send_text(self, text: str) -> None:
received.append(text)
fake = FakeWS()
_connections.append(fake)
try:
await broadcast_service_status(
node_id="node-7",
services=[{"port": 80, "protocol": "tcp", "status": "offline"}],
checked_at="2024-01-01T00:00:00",
)
finally:
_drop(fake)
msg = json.loads(received[0])
assert msg["type"] == "service_status"
assert msg["node_id"] == "node-7"
assert msg["services"] == [{"port": 80, "protocol": "tcp", "status": "offline"}]
# ---------------------------------------------------------------------------
# _drop — idempotent connection removal (regression for double-remove crash)
# ---------------------------------------------------------------------------
def test_drop_is_idempotent():
"""Dropping a connection twice must not raise (was a ValueError crash)."""
class FakeWS:
pass
fake = FakeWS()
_connections.append(fake)
_drop(fake)
_drop(fake) # second drop must be a no-op
assert fake not in _connections
@pytest.mark.asyncio
async def test_broadcast_dead_connection_dropped_once_safely():
"""A send failure removes the dead socket without a double-remove crash."""
class DeadWS:
async def send_text(self, _: str) -> None:
raise RuntimeError("disconnected")
dead = DeadWS()
_connections.append(dead)
await broadcast_status(node_id="n", status="online", checked_at="t")
# A second broadcast must not raise even though dead is already gone.
await broadcast_status(node_id="n", status="online", checked_at="t")
assert dead not in _connections