111 lines
4.1 KiB
Python
111 lines
4.1 KiB
Python
"""Streaming über PostgreSQL (→ Konzept 4.4).
|
|
|
|
Der Worker schreibt Ereignisse mit fortlaufender Nummer je Run in `run_events`
|
|
und weckt per NOTIFY die Web-Prozesse. Die lesen per `seq > letzte` nach.
|
|
NOTIFY ist nur ein Signal; die Daten kommen immer aus der Tabelle.
|
|
"""
|
|
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
import uuid
|
|
from typing import Any
|
|
|
|
import psycopg
|
|
from sqlalchemy import func, select, text
|
|
from sqlalchemy.orm import Session
|
|
|
|
from app.core.database import Database
|
|
from app.tasks.models import RunEvent
|
|
|
|
CHANNEL = "run_events"
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
def append(session: Session, run_id: uuid.UUID, type_: str, data: dict[str, Any]) -> RunEvent:
|
|
"""Ereignis anhängen; NOTIFY wird mit dem Commit derselben Transaktion zugestellt."""
|
|
last = session.scalar(select(func.max(RunEvent.seq)).where(RunEvent.run_id == run_id)) or 0
|
|
event = RunEvent(run_id=run_id, seq=last + 1, type=type_, data=data)
|
|
session.add(event)
|
|
session.flush()
|
|
session.execute(text("SELECT pg_notify(:channel, :payload)"), {"channel": CHANNEL, "payload": str(run_id)})
|
|
return event
|
|
|
|
|
|
def read(session: Session, run_id: uuid.UUID, after: int, limit: int = 500) -> list[RunEvent]:
|
|
return list(session.scalars(select(RunEvent).where(RunEvent.run_id == run_id, RunEvent.seq > after)
|
|
.order_by(RunEvent.seq).limit(limit)))
|
|
|
|
|
|
class EventWriter:
|
|
"""Für den Worker: Tokens sammeln und gebündelt schreiben, statt ein Ereignis je Token."""
|
|
|
|
def __init__(self, db: Database, run_id: uuid.UUID, flush_ms: int):
|
|
self.db, self.run_id, self.flush_seconds = db, run_id, flush_ms / 1000
|
|
self.buffer: list[str] = []
|
|
self.last_flush = time.monotonic()
|
|
|
|
def token(self, piece: str) -> None:
|
|
self.buffer.append(piece)
|
|
if time.monotonic() - self.last_flush >= self.flush_seconds:
|
|
self.flush()
|
|
|
|
def flush(self) -> None:
|
|
if self.buffer:
|
|
pieces, self.buffer = "".join(self.buffer), []
|
|
with self.db.session() as session:
|
|
append(session, self.run_id, "token", {"text": pieces})
|
|
self.last_flush = time.monotonic()
|
|
|
|
def emit(self, type_: str, data: dict[str, Any], session: Session | None = None) -> None:
|
|
"""Offene Tokens zuerst, damit die Reihenfolge stimmt."""
|
|
self.flush()
|
|
if session is not None:
|
|
append(session, self.run_id, type_, data)
|
|
return
|
|
with self.db.session() as own:
|
|
append(own, self.run_id, type_, data)
|
|
|
|
|
|
class EventHub:
|
|
"""Im Web-Prozess: eine LISTEN-Verbindung je Prozess, die wartende SSE-Streams weckt.
|
|
|
|
Fällt die Verbindung aus, verbindet sich der Hub neu; die Streams fragen in der
|
|
Zwischenzeit von selbst in kurzen Abständen nach (Fallback-Polling).
|
|
"""
|
|
|
|
def __init__(self, url: str):
|
|
self.url = url
|
|
self.waiters: dict[str, set[asyncio.Event]] = {}
|
|
self.connected = False
|
|
|
|
def subscribe(self, run_id: uuid.UUID) -> asyncio.Event:
|
|
event = asyncio.Event()
|
|
self.waiters.setdefault(str(run_id), set()).add(event)
|
|
return event
|
|
|
|
def unsubscribe(self, run_id: uuid.UUID, event: asyncio.Event) -> None:
|
|
waiting = self.waiters.get(str(run_id))
|
|
if waiting is not None:
|
|
waiting.discard(event)
|
|
if not waiting:
|
|
self.waiters.pop(str(run_id), None)
|
|
|
|
def wake(self, run_id: str) -> None:
|
|
for event in self.waiters.get(run_id, ()):
|
|
event.set()
|
|
|
|
async def run(self) -> None:
|
|
while True:
|
|
try:
|
|
async with await psycopg.AsyncConnection.connect(self.url, autocommit=True) as connection:
|
|
await connection.execute(f"LISTEN {CHANNEL}")
|
|
self.connected = True
|
|
async for notify in connection.notifies():
|
|
self.wake(notify.payload)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception:
|
|
log.warning("LISTEN-Verbindung unterbrochen, neuer Versuch in 2 s", exc_info=True)
|
|
self.connected = False
|
|
await asyncio.sleep(2)
|