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