trumpsignal-backend/app/ws/manager.py

"""
WebSocket connection manager.

One global broadcaster that fans out every message to every connected client.
Critical that broadcast NEVER blocks the caller — Binance kline ticks arrive
every ~500 ms and the broadcast is inline with the WS read loop. A single
slow / half-closed client could otherwise stall the kline pipeline, miss the
keepalive ping, and tear the upstream Binance socket down (this was the root
cause of the "no close frame received or sent" reconnect storms in prod logs).

Two defences here:
  * Per-client send wrapped in `asyncio.wait_for(..., timeout=PER_CLIENT_SEND_TIMEOUT)`.
  * All clients dispatched in parallel via `asyncio.gather()` so one slow
    client can't delay anyone else.
"""

from __future__ import annotations

import asyncio
import json
import logging
from typing import List

from fastapi import WebSocket

logger = logging.getLogger(__name__)

# Per-client send timeout. 2 s is plenty for a healthy TCP send of our tiny
# JSON payloads; anything slower means a half-closed or zombie client.
PER_CLIENT_SEND_TIMEOUT = 2.0


class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []

    async def connect(self, ws: WebSocket):
        await ws.accept()
        self.active_connections.append(ws)
        logger.info("WebSocket connected. Total connections: %d", len(self.active_connections))

    async def disconnect(self, ws: WebSocket):
        if ws in self.active_connections:
            self.active_connections.remove(ws)
        logger.info("WebSocket disconnected. Total connections: %d", len(self.active_connections))

    async def broadcast(self, message: dict):
        if not self.active_connections:
            return
        payload = json.dumps(message)
        # Snapshot the list — disconnect() mutates active_connections and we
        # don't want "list changed during iteration" when we prune.
        connections = list(self.active_connections)

        async def _send_one(ws: WebSocket):
            try:
                await asyncio.wait_for(ws.send_text(payload),
                                       timeout=PER_CLIENT_SEND_TIMEOUT)
                return None
            except asyncio.TimeoutError:
                logger.warning("WebSocket send timed out (>%.1fs) — pruning client",
                               PER_CLIENT_SEND_TIMEOUT)
                return ws
            except Exception as exc:
                logger.warning("WebSocket send failed: %s — pruning client", exc)
                return ws

        # Fan out concurrently — one slow client can't stall the others.
        results = await asyncio.gather(*[_send_one(ws) for ws in connections])
        for dead_ws in results:
            if dead_ws is not None:
                await self.disconnect(dead_ws)


manager = ConnectionManager()