79 lines
2.7 KiB
Python
79 lines
2.7 KiB
Python
"""
|
|
SSEManager — pub/sub hub for Server-Sent Events.
|
|
|
|
Task runner calls publish() → all connected SSE clients watching that task_id receive the event.
|
|
|
|
Replay buffer: each task_id stores the last REPLAY_MAXLEN events so that clients who connect
|
|
slightly after publishing starts (race between POST /discover and GET /events) still receive
|
|
the initial events (disc_server, disc_rack_start, etc.).
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
from collections import deque
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime, timezone
|
|
|
|
REPLAY_MAXLEN = 200
|
|
|
|
|
|
@dataclass
|
|
class SSEEvent:
|
|
type: str # "task_event" | "task_status" | "task_result"
|
|
task_id: str
|
|
data: dict
|
|
|
|
def encode(self) -> str:
|
|
payload = json.dumps({"task_id": self.task_id, **self.data})
|
|
return f"event: {self.type}\ndata: {payload}\n\n"
|
|
|
|
|
|
class SSEManager:
|
|
def __init__(self) -> None:
|
|
# task_id -> list[Queue]
|
|
self._subscribers: dict[str, list[asyncio.Queue]] = {}
|
|
# task_id -> deque of recent events for late-joining subscribers
|
|
self._replay: dict[str, deque[SSEEvent]] = {}
|
|
self._lock = asyncio.Lock()
|
|
|
|
def clear_replay(self, task_id: str) -> None:
|
|
"""Call before starting a new task to discard stale events from a previous run."""
|
|
self._replay.pop(task_id, None)
|
|
|
|
async def subscribe(self, task_id: str, queue: asyncio.Queue) -> None:
|
|
async with self._lock:
|
|
self._subscribers.setdefault(task_id, []).append(queue)
|
|
# Replay buffered events so late subscribers catch up
|
|
for event in self._replay.get(task_id, []):
|
|
try:
|
|
queue.put_nowait(event)
|
|
except asyncio.QueueFull:
|
|
break
|
|
|
|
async def unsubscribe(self, task_id: str, queue: asyncio.Queue) -> None:
|
|
async with self._lock:
|
|
subs = self._subscribers.get(task_id, [])
|
|
try:
|
|
subs.remove(queue)
|
|
except ValueError:
|
|
pass
|
|
if not subs:
|
|
self._subscribers.pop(task_id, None)
|
|
|
|
async def publish(self, event: SSEEvent) -> None:
|
|
async with self._lock:
|
|
buf = self._replay.setdefault(event.task_id, deque(maxlen=REPLAY_MAXLEN))
|
|
buf.append(event)
|
|
queues = list(self._subscribers.get(event.task_id, []))
|
|
for queue in queues:
|
|
try:
|
|
queue.put_nowait(event)
|
|
except asyncio.QueueFull:
|
|
pass # slow client, drop event
|
|
|
|
def subscriber_count(self, task_id: str) -> int:
|
|
return len(self._subscribers.get(task_id, []))
|
|
|
|
|
|
# Module-level singleton shared across all requests
|
|
sse_manager = SSEManager()
|