64 lines
2 KiB
Python
64 lines
2 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import time
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
|
|
class UpdateBroker:
|
|
def __init__(self) -> None:
|
|
self._subscribers: set[asyncio.Queue[dict[str, object]]] = set()
|
|
|
|
def subscribe(self) -> asyncio.Queue[dict[str, object]]:
|
|
queue: asyncio.Queue[dict[str, object]] = asyncio.Queue(maxsize=20)
|
|
self._subscribers.add(queue)
|
|
return queue
|
|
|
|
def unsubscribe(self, queue: asyncio.Queue[dict[str, object]]) -> None:
|
|
self._subscribers.discard(queue)
|
|
|
|
async def publish(self, event: str, payload: dict[str, object] | None = None) -> None:
|
|
message = {
|
|
"event": event,
|
|
"data": {
|
|
"type": event,
|
|
"published_at": time.time(),
|
|
**(payload or {}),
|
|
},
|
|
}
|
|
stale_queues: list[asyncio.Queue[dict[str, object]]] = []
|
|
for queue in tuple(self._subscribers):
|
|
try:
|
|
queue.put_nowait(message)
|
|
except asyncio.QueueFull:
|
|
stale_queues.append(queue)
|
|
|
|
for queue in stale_queues:
|
|
self.unsubscribe(queue)
|
|
|
|
|
|
async def stream_sse_events(
|
|
broker: UpdateBroker,
|
|
) -> AsyncIterator[str]:
|
|
queue = broker.subscribe()
|
|
try:
|
|
yield _sse_event("connected", {"type": "connected", "published_at": time.time()})
|
|
while True:
|
|
try:
|
|
message = await asyncio.wait_for(queue.get(), timeout=25)
|
|
except TimeoutError:
|
|
yield ": keepalive\n\n"
|
|
continue
|
|
|
|
yield _sse_event(
|
|
str(message.get("event", "message")),
|
|
message.get("data") if isinstance(message.get("data"), dict) else {},
|
|
)
|
|
finally:
|
|
broker.unsubscribe(queue)
|
|
|
|
|
|
def _sse_event(event: str, payload: dict[str, Any]) -> str:
|
|
return f"event: {event}\ndata: {json.dumps(payload, separators=(',', ':'))}\n\n"
|