201 lines
7.4 KiB
Python
201 lines
7.4 KiB
Python
"""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()
|