"""Per-process, project-scoped PostgreSQL LISTEN hubs for workbench SSE.""" from __future__ import annotations import asyncio import contextlib import logging from contextlib import asynccontextmanager from uuid import UUID from psycopg import AsyncConnection from psycopg_pool import AsyncConnectionPool from psycopg.rows import dict_row from app.infra.db.postgresql.workbench_events import WorkbenchEventRepository logger = logging.getLogger(__name__) class WorkbenchEventHub: """Own one LISTEN connection and fan wake-ups out to SSE subscribers.""" def __init__(self, project_id: UUID, dsn: str) -> None: self.project_id = project_id self.dsn = dsn self.subscribers = 0 self._connection: AsyncConnection | None = None self._pool: AsyncConnectionPool | None = None self._connection_lock = asyncio.Lock() self._condition = asyncio.Condition() self._generation = 0 self._ready = asyncio.Event() self._closed = False self._listener_task: asyncio.Task | None = None self._idle_task: asyncio.Task | None = None async def start(self) -> None: if self._listener_task is None or self._listener_task.done(): self._closed = False self._ready.clear() self._listener_task = asyncio.create_task( self._listen_loop(), name=f"workbench-events:{self.project_id}" ) await asyncio.wait_for(self._ready.wait(), timeout=10) async def _connect(self) -> AsyncConnection: if self._pool is None: self._pool = AsyncConnectionPool( conninfo=self.dsn, min_size=0, max_size=1, open=False, kwargs={"autocommit": True, "row_factory": dict_row}, ) await self._pool.open() connection = await self._pool.getconn() await connection.execute("LISTEN workbench_events") return connection async def _listen_loop(self) -> None: retry_seconds = 1 while not self._closed: try: if self._connection is None or self._connection.closed: async with self._connection_lock: if self._connection is None or self._connection.closed: self._connection = await self._connect() self._ready.set() retry_seconds = 1 notified = False async with self._connection_lock: connection = self._connection if connection is None: continue async for notification in connection.notifies(timeout=1, stop_after=1): if notification.payload.startswith(f"{self.project_id}:"): notified = True break if notified: await self._wake() await asyncio.sleep(0) except asyncio.CancelledError: raise except Exception: logger.exception("Workbench event listener failed for project %s", self.project_id) await self._drop_connection() await self._wake() await asyncio.sleep(retry_seconds) retry_seconds = min(retry_seconds * 2, 30) async def _wake(self) -> None: async with self._condition: self._generation += 1 self._condition.notify_all() async def wait_for_wake(self, generation: int, timeout: float) -> int: async with self._condition: if self._generation != generation: return self._generation with contextlib.suppress(asyncio.TimeoutError): await asyncio.wait_for( self._condition.wait_for(lambda: self._generation != generation), timeout=timeout, ) return self._generation @property def generation(self) -> int: return self._generation async def oldest_cursor(self) -> int | None: return await self._read(lambda repository: repository.oldest_cursor()) async def current_cursor(self) -> int: return await self._read(lambda repository: repository.current_cursor()) async def read_after(self, after: int) -> list[dict]: return await self._read(lambda repository: repository.read_after(after)) async def _read(self, operation): await self.start() async with self._connection_lock: if self._connection is None or self._connection.closed: self._connection = await self._connect() self._ready.set() repository = WorkbenchEventRepository(self._connection, self.project_id) return await operation(repository) async def _drop_connection(self) -> None: async with self._connection_lock: connection, self._connection = self._connection, None if connection is not None and self._pool is not None: with contextlib.suppress(Exception): await self._pool.putconn(connection) async def close(self) -> None: self._closed = True if self._idle_task is not None: self._idle_task.cancel() self._idle_task = None if self._listener_task is not None: self._listener_task.cancel() with contextlib.suppress(asyncio.CancelledError): await self._listener_task self._listener_task = None await self._drop_connection() if self._pool is not None: await self._pool.close() self._pool = None class WorkbenchEventHubManager: def __init__(self, idle_seconds: float = 60) -> None: self.idle_seconds = idle_seconds self._hubs: dict[tuple[UUID, str], WorkbenchEventHub] = {} self._lock = asyncio.Lock() @asynccontextmanager async def subscribe(self, project_id: UUID, dsn: str): key = (project_id, dsn) async with self._lock: hub = self._hubs.get(key) if hub is None: hub = WorkbenchEventHub(project_id, dsn) self._hubs[key] = hub if hub._idle_task is not None: hub._idle_task.cancel() hub._idle_task = None hub.subscribers += 1 try: await hub.start() yield hub finally: async with self._lock: hub.subscribers = max(0, hub.subscribers - 1) if hub.subscribers == 0 and hub._idle_task is None: hub._idle_task = asyncio.create_task(self._close_when_idle(key, hub)) async def _close_when_idle( self, key: tuple[UUID, str], hub: WorkbenchEventHub ) -> None: try: await asyncio.sleep(self.idle_seconds) async with self._lock: if hub.subscribers or self._hubs.get(key) is not hub: return self._hubs.pop(key, None) hub._idle_task = None await hub.close() except asyncio.CancelledError: return async def close_all(self) -> None: async with self._lock: hubs = list(self._hubs.values()) self._hubs.clear() await asyncio.gather(*(hub.close() for hub in hubs), return_exceptions=True) workbench_event_hubs = WorkbenchEventHubManager()