diff --git a/src/youtube_extension/backend/services/websocket_service.py b/src/youtube_extension/backend/services/websocket_service.py index 9cd81bb52..6f7b49e97 100644 --- a/src/youtube_extension/backend/services/websocket_service.py +++ b/src/youtube_extension/backend/services/websocket_service.py @@ -7,6 +7,7 @@ Handles real-time communication, message routing, and connection management. """ +import asyncio import json import logging from datetime import datetime @@ -47,18 +48,47 @@ async def send_personal_message(self, message: str, websocket: WebSocket): self.disconnect(websocket) async def broadcast(self, message: str): - """Broadcast message to all active connections""" - disconnected = [] - for connection in self.active_connections: - try: - await connection.send_text(message) - except Exception as e: - logger.error(f"Error broadcasting message: {e}") - disconnected.append(connection) + """Broadcast message to all active connections concurrently. - # Remove disconnected connections - for connection in disconnected: - self.disconnect(connection) + Sends are issued in parallel rather than one at a time. A sequential + loop makes every client wait for the ones ahead of it, so a single slow + or backpressured peer delays delivery to the whole fleet (head-of-line + blocking) and total latency grows with the connection count. + + Each WebSocket owns an independent send buffer, so there is no shared + pooled resource to exhaust here and concurrency is naturally bounded by + the number of connected clients. Failures are isolated via + ``return_exceptions=True`` so one dead peer cannot abort the broadcast. + """ + # Snapshot: disconnect() mutates active_connections as we clean up. + connections = list(self.active_connections) + if not connections: + return + + results = await asyncio.gather( + *(connection.send_text(message) for connection in connections), + return_exceptions=True, + ) + + # `return_exceptions=True` captures BaseException-only failures (most + # notably CancelledError) into `results` instead of propagating them, + # so `isinstance(result, Exception)` alone would silently swallow a + # cancellation that the previous `except Exception` loop let escape. + # Dead peers are cleaned up first so a cancellation does not leak them, + # then the cancellation is re-raised to restore the old semantics. + cancellations = [ + result + for result in results + if isinstance(result, BaseException) and not isinstance(result, Exception) + ] + + for connection, result in zip(connections, results, strict=True): + if isinstance(result, Exception): + logger.error(f"Error broadcasting message: {result}") + self.disconnect(connection) + + if cancellations: + raise cancellations[0] class WebSocketService: diff --git a/tests/unit/test_websocket_service.py b/tests/unit/test_websocket_service.py index bbc6a0297..af9b1cda9 100644 --- a/tests/unit/test_websocket_service.py +++ b/tests/unit/test_websocket_service.py @@ -2,9 +2,11 @@ from __future__ import annotations +import asyncio import json from unittest.mock import AsyncMock, MagicMock, patch +import pytest from fastapi import WebSocketDisconnect from youtube_extension.backend.services.websocket_service import ( @@ -513,3 +515,144 @@ def test_returns_zero_when_no_connections(self): mgr.active_connections = [] stats = svc.get_connection_stats() assert stats["active_connections"] == 0 + + +# =========================================================================== +# WebSocketConnectionManager.broadcast – fan-out concurrency +# =========================================================================== + + +class TestBroadcastFanOut: + """broadcast() must issue sends in parallel, not one connection at a time. + + A sequential loop makes every client wait for the ones ahead of it, so a + single slow peer delays delivery to the entire fleet (head-of-line + blocking). These tests fail against a serialised implementation. + """ + + @staticmethod + def _tracking_ws(state, delay=0.01): + """A mock WebSocket that records peak concurrent send_text() calls.""" + ws = _make_ws() + + async def _send(_message): + state["inflight"] += 1 + state["peak"] = max(state["peak"], state["inflight"]) + await asyncio.sleep(delay) + state["inflight"] -= 1 + + ws.send_text = AsyncMock(side_effect=_send) + return ws + + async def test_sends_are_concurrent_not_serialised(self): + mgr = WebSocketConnectionManager() + state = {"inflight": 0, "peak": 0} + n = 5 + mgr.active_connections.extend(self._tracking_ws(state) for _ in range(n)) + + await mgr.broadcast("event") + + assert state["peak"] == n, ( + f"broadcast() peaked at {state['peak']} concurrent send(s) for {n} " + "connections - sends are serialised (head-of-line blocking)" + ) + + async def test_slow_peer_does_not_delay_other_peers(self): + mgr = WebSocketConnectionManager() + completed: list[str] = [] + + def _ws(name, delay): + ws = _make_ws() + + async def _send(_message): + await asyncio.sleep(delay) + completed.append(name) + + ws.send_text = AsyncMock(side_effect=_send) + return ws + + # The slow peer is first in the list: under a sequential loop it blocks + # both fast peers behind it and therefore completes first. + mgr.active_connections.extend( + [_ws("slow", 0.05), _ws("fast-1", 0.0), _ws("fast-2", 0.0)] + ) + + await mgr.broadcast("event") + + # Both fast peers must finish before the slow head-of-list peer. Their + # order relative to each other is scheduler-dependent (both sleep(0)), + # so it is deliberately not asserted -- only that neither was blocked + # behind the slow peer. Under a sequential loop the slow peer, being + # first, completes first and this fails. + assert completed[-1] == "slow" and set(completed[:2]) == {"fast-1", "fast-2"}, ( + f"completion order was {completed}; a slow peer at the head of the " + "connection list delayed delivery to the fast peers behind it" + ) + + async def test_all_peers_still_receive_the_message(self): + mgr = WebSocketConnectionManager() + conns = [_make_ws() for _ in range(4)] + mgr.active_connections.extend(conns) + + await mgr.broadcast("payload") + + for ws in conns: + ws.send_text.assert_awaited_once_with("payload") + + async def test_failures_are_isolated_and_only_bad_peers_removed(self): + mgr = WebSocketConnectionManager() + ok_1, ok_2 = _make_ws(), _make_ws() + bad_1, bad_2 = _make_ws(), _make_ws() + bad_1.send_text = AsyncMock(side_effect=RuntimeError("peer gone")) + bad_2.send_text = AsyncMock(side_effect=ConnectionResetError("reset")) + # Failing peers first: a raising send must not abort the whole fan-out. + mgr.active_connections.extend([bad_1, ok_1, bad_2, ok_2]) + + await mgr.broadcast("event") + + ok_1.send_text.assert_awaited_once_with("event") + ok_2.send_text.assert_awaited_once_with("event") + assert mgr.active_connections == [ok_1, ok_2], ( + "expected only the failing peers to be dropped, got " + f"{len(mgr.active_connections)} surviving connection(s)" + ) + + async def test_empty_connection_list_issues_no_sends(self): + mgr = WebSocketConnectionManager() + # Must not raise and must not construct an empty gather. + await mgr.broadcast("event") + assert mgr.active_connections == [] + + async def test_cancelled_send_is_re_raised_not_swallowed(self): + """gather(return_exceptions=True) captures CancelledError; broadcast must not eat it. + + The pre-fan-out implementation used `except Exception`, so a + CancelledError raised by send_text escaped broadcast(). Collecting it + into `results` and filtering on `Exception` would silently drop it. + """ + mgr = WebSocketConnectionManager() + ok = _make_ws() + cancelled = _make_ws() + cancelled.send_text = AsyncMock(side_effect=asyncio.CancelledError()) + mgr.active_connections.extend([ok, cancelled]) + + with pytest.raises(asyncio.CancelledError): + await mgr.broadcast("event") + + async def test_cancellation_does_not_leak_dead_peers(self): + """Ordinary failures are still cleaned up before the cancellation propagates.""" + mgr = WebSocketConnectionManager() + ok = _make_ws() + dead = _make_ws() + dead.send_text = AsyncMock(side_effect=RuntimeError("peer gone")) + cancelled = _make_ws() + cancelled.send_text = AsyncMock(side_effect=asyncio.CancelledError()) + mgr.active_connections.extend([ok, dead, cancelled]) + + with pytest.raises(asyncio.CancelledError): + await mgr.broadcast("event") + + assert mgr.active_connections == [ok, cancelled], ( + "the failed peer must still be dropped even though the broadcast " + f"was cancelled, got {mgr.active_connections}" + )