plm-lottery/app/rounds/events.py

import asyncio
from collections import defaultdict

# Defensive backstop on concurrent SSE subscribers overall, regardless of source
# — expected load is on the order of ~100 concurrent users, so this is set well
# above that. The real defense against a single abusive source is the per-IP cap
# below (B-38): a global-only cap was trivially exhausted by one client opening
# MAX_SUBSCRIBERS connections, degrading every other user to polling — the
# comment used to call it "defensive"; it was actually the vector.
MAX_SUBSCRIBERS = 500

# How many concurrent streams a single client IP may hold. Deliberately small —
# a real browser tab needs at most one, occasionally two briefly across a
# reload — since this bounds one source's share of the global capacity, not a
# legitimate per-user concurrency limit.
MAX_SUBSCRIBERS_PER_IP = 5

# Put on a to-be-evicted subscriber's queue (B-38) to wake its generator
# (app/api/routes/rounds.py:round_stream) promptly so it closes the connection
# instead of lingering, silently uncounted, until the client's own network
# timeout or the next keep-alive tick.
EVICTED = object()


class RoundEventCapacityError(Exception):
    """Raised by subscribe() when MAX_SUBSCRIBERS — the global backstop — is
    already reached. The per-IP cap never raises this; it evicts instead (see
    subscribe())."""


class RoundEventBroadcaster:
    """In-process pub/sub so SSE clients (GET /rounds/stream) get pushed a
    notification the instant round/bet/balance state changes, instead of only
    finding out on their next poll. The message carries no payload — it's just
    a "something changed, go refetch" signal; the client re-hits the existing
    per-user REST endpoints (/rounds/current, /users/me, ...) for the actual
    data, so this never needs to know what changed or who's allowed to see it.

    Single-process only (no cross-worker fan-out) — fine for this deployment
    (one uvicorn process, see docker-compose.yml). A multi-worker deployment
    would need a shared channel (e.g. Redis pub/sub) instead.
    """

    def __init__(self, max_subscribers: int = MAX_SUBSCRIBERS, max_per_ip: int = MAX_SUBSCRIBERS_PER_IP):
        self._ip_by_queue: dict[asyncio.Queue, str] = {}
        self._queues_by_ip: dict[str, list[asyncio.Queue]] = defaultdict(list)
        self.max_subscribers = max_subscribers
        self.max_per_ip = max_per_ip

    def subscribe(self, client_ip: str = "unknown") -> asyncio.Queue:
        if len(self._ip_by_queue) >= self.max_subscribers:
            raise RoundEventCapacityError(f"already at the {self.max_subscribers}-subscriber cap")

        ip_queues = self._queues_by_ip[client_ip]
        if len(ip_queues) >= self.max_per_ip:
            # B-38: evict this IP's own oldest connection rather than refusing
            # the new one — bounds one source's footprint without turning a
            # legitimate reconnect storm (a flaky network retrying EventSource)
            # into an outright block, and without letting one abusive IP crowd
            # out unrelated clients the way the old global-only cap did.
            oldest = ip_queues.pop(0)
            self._ip_by_queue.pop(oldest, None)
            if not oldest.full():
                oldest.put_nowait(EVICTED)

        queue: asyncio.Queue = asyncio.Queue(maxsize=1)
        self._ip_by_queue[queue] = client_ip
        ip_queues.append(queue)
        return queue

    def unsubscribe(self, queue: asyncio.Queue) -> None:
        client_ip = self._ip_by_queue.pop(queue, None)
        if client_ip is None:
            return
        ip_queues = self._queues_by_ip.get(client_ip)
        if ip_queues is None:
            return
        if queue in ip_queues:
            ip_queues.remove(queue)
        if not ip_queues:
            self._queues_by_ip.pop(client_ip, None)

    def publish(self) -> None:
        for queue in self._ip_by_queue:
            if queue.full():
                continue  # a not-yet-delivered notification already covers this one
            queue.put_nowait(None)


broadcaster = RoundEventBroadcaster()