From 9c7e32d0d677140db24466a7eaccbbb5e7a7a7ee Mon Sep 17 00:00:00 2001 From: Josh Hawkins <32435876+hawkeye217@users.noreply.github.com> Date: Fri, 21 Aug 2026 19:07:58 -0500 Subject: [PATCH] don't let a slow websocket client block publishers --- frigate/comms/ws.py | 67 ++++++++++++++++-- frigate/test/test_ws_send_queue.py | 107 +++++++++++++++++++++++++++++ 2 files changed, 170 insertions(+), 4 deletions(-) create mode 100644 frigate/test/test_ws_send_queue.py diff --git a/frigate/comms/ws.py b/frigate/comms/ws.py index ccb5d42890..59ba916b8c 100644 --- a/frigate/comms/ws.py +++ b/frigate/comms/ws.py @@ -3,6 +3,8 @@ import errno import json import logging +import queue +import socket import threading from collections.abc import Callable from typing import Any @@ -74,6 +76,9 @@ _WS_VIEWER_TOPICS = frozenset( # Camera-scoped command topics a camera-authorized (non-admin) user may send. _WS_CAMERA_COMMAND_TOPICS = frozenset({"ptz"}) +# Max outbound messages waiting on a client's writer thread. +WS_MAX_PENDING_MESSAGES = 256 + def _check_ws_authorization( topic: str, @@ -446,6 +451,63 @@ def _materialize_for_ws( class WebSocket(WebSocket_): # type: ignore[misc] + def __init__(self, *args: Any, **kwargs: Any) -> None: + super().__init__(*args, **kwargs) + self._send_queue: queue.Queue[tuple[Any, bool] | None] = queue.Queue( + maxsize=WS_MAX_PENDING_MESSAGES + ) + self._writer: threading.Thread | None = None + self._aborted = False + + def opened(self) -> None: + # every client gets its own writer so a client that stops reading only + # blocks itself, never the thread that called publish() + self._writer = threading.Thread( + target=self._drain_send_queue, name="ws_writer", daemon=True + ) + self._writer.start() + + def send(self, payload: Any, binary: bool = False) -> None: + try: + self._send_queue.put_nowait((payload, binary)) + except queue.Full: + self._abort("Websocket client is not keeping up, disconnecting it") + + def _drain_send_queue(self) -> None: + while True: + item = self._send_queue.get() + if item is None or self.terminated or self.sock is None: + return + try: + super().send(*item) + except Exception: + self._abort() + return + + def _abort(self, reason: str | None = None) -> None: + # publish() keeps hitting a full queue until the manager thread removes + # the connection, so only act (and log) the first time + if self._aborted: + return + self._aborted = True + if reason: + logger.warning(reason) + + # shutdown rather than close so the ws4py manager thread sees EOF and + # runs its normal unregister/terminate; this also unblocks a stuck sendall + sock = self.sock + if sock is not None: + try: + sock.shutdown(socket.SHUT_RDWR) + except OSError: + pass + + def closed(self, code: int, reason: str | None = None) -> None: + try: + self._send_queue.put_nowait(None) + except queue.Full: + pass + def unhandled_error(self, error: Any) -> None: """ Handles the unfriendly socket closures on the server side @@ -580,10 +642,7 @@ class WebSocketClient(Communicator): ) if message is None: continue - try: - ws.send(message) - except (ConnectionResetError, BrokenPipeError, ValueError): - pass + ws.send(message) def stop(self) -> None: if self.websocket_server is not None: diff --git a/frigate/test/test_ws_send_queue.py b/frigate/test/test_ws_send_queue.py new file mode 100644 index 0000000000..a7569f7fd7 --- /dev/null +++ b/frigate/test/test_ws_send_queue.py @@ -0,0 +1,107 @@ +"""Outbound websocket sends must never block the thread calling publish().""" + +import socket +import threading +import unittest + +from frigate.comms.ws import WS_MAX_PENDING_MESSAGES, WebSocket + + +class _FakeSock: + """Socket stand-in; ``block`` makes sendall hang like a client that stopped reading.""" + + def __init__(self, block: bool = False) -> None: + self.block = block + self.released = threading.Event() + self.shutdown_called = threading.Event() + self.frames: list[bytes] = [] + + def sendall(self, data: bytes) -> None: + if self.block and not self.released.is_set(): + self.released.wait(timeout=10) + raise BrokenPipeError() + self.frames.append(data) + + def shutdown(self, how: int) -> None: + assert how == socket.SHUT_RDWR + self.shutdown_called.set() + self.released.set() + + def close(self) -> None: + pass + + def fileno(self) -> int: + return 99 + + +def _wait_for(predicate, timeout: float = 2.0) -> bool: + deadline = threading.Event() + for _ in range(int(timeout / 0.01)): + if predicate(): + return True + deadline.wait(0.01) + return predicate() + + +class TestWebSocketSendQueue(unittest.TestCase): + def _open(self, sock: _FakeSock) -> WebSocket: + ws = WebSocket(sock) + ws.opened() + return ws + + def test_stalled_client_does_not_block_publisher(self): + sock = _FakeSock(block=True) + ws = self._open(sock) + + def publish_many(): + for i in range(WS_MAX_PENDING_MESSAGES + 5): + ws.send(f"message {i}") + + publisher = threading.Thread(target=publish_many, daemon=True) + publisher.start() + publisher.join(timeout=2.0) + + self.assertFalse(publisher.is_alive(), "publish() blocked on a stalled client") + self.assertTrue( + sock.shutdown_called.wait(timeout=2.0), + "a client that cannot keep up should be disconnected", + ) + + def test_overflow_warns_and_shuts_down_once(self): + sock = _FakeSock(block=True) + ws = self._open(sock) + shutdown_calls = [] + original_shutdown = sock.shutdown + sock.shutdown = lambda how: (shutdown_calls.append(how), original_shutdown(how)) + + with self.assertLogs("frigate.comms.ws", level="WARNING") as logs: + # keep publishing after overflow, as the dispatcher does until the + # manager thread removes the connection + for i in range(WS_MAX_PENDING_MESSAGES * 3): + ws.send(f"message {i}") + + self.assertEqual(len(logs.output), 1) + self.assertEqual(len(shutdown_calls), 1) + + def test_messages_delivered_in_order(self): + sock = _FakeSock() + ws = self._open(sock) + for i in range(3): + ws.send(f"message {i}") + + self.assertTrue(_wait_for(lambda: len(sock.frames) == 3)) + for i, frame in enumerate(sock.frames): + self.assertIn(f"message {i}".encode(), frame) + self.assertFalse(sock.shutdown_called.is_set()) + + def test_closed_stops_writer_thread(self): + sock = _FakeSock() + ws = self._open(sock) + writer = ws._writer + ws.closed(1000, "bye") + writer.join(timeout=2.0) + self.assertFalse(writer.is_alive()) + + +if __name__ == "__main__": + unittest.main()