diff --git a/backend/src/api/realtime.py b/backend/src/api/realtime.py index 8eb18c7..2b9803c 100644 --- a/backend/src/api/realtime.py +++ b/backend/src/api/realtime.py @@ -5,8 +5,8 @@ from typing import Any import asyncpg -from utils.env import env from utils.events import BG_EVENTS_CHANNEL +from utils.listen import Listener from utils.read import chats as chats_read from utils.read import presence as presence_read @@ -26,7 +26,7 @@ class EventHub: def __init__(self) -> None: self._subscribers: set[Subscriber] = set() self._pool: asyncpg.Pool | None = None - self._conn: asyncpg.Connection | None = None + self._listener = Listener("realtime", {BG_EVENTS_CHANNEL: self._on_notify}) self._tasks: set[asyncio.Task] = set() def subscribe(self, account_id: int, chat_id: int | None) -> Subscriber: @@ -39,21 +39,14 @@ class EventHub: async def start(self, pool: asyncpg.Pool) -> None: self._pool = pool - conn = await asyncpg.connect(dsn=env.db.connection_url) - await conn.add_listener(BG_EVENTS_CHANNEL, self._on_notify) - self._conn = conn - logger.info("Realtime hub listening on %s", BG_EVENTS_CHANNEL) + await self._listener.start() async def stop(self) -> None: for task in self._tasks: task.cancel() - if self._conn is not None: - await self._conn.close() - self._conn = None + await self._listener.stop() - def _on_notify( - self, _conn: asyncpg.Connection, _pid: int, _channel: str, payload: str - ) -> None: + def _on_notify(self, payload: str) -> None: task = asyncio.create_task(self._dispatch(payload)) self._tasks.add(task) task.add_done_callback(self._tasks.discard) diff --git a/backend/src/userbot/modules/jobs/consumer.py b/backend/src/userbot/modules/jobs/consumer.py index fcd99da..68f598a 100644 --- a/backend/src/userbot/modules/jobs/consumer.py +++ b/backend/src/userbot/modules/jobs/consumer.py @@ -9,7 +9,7 @@ from userbot.modules.jobs import repository from userbot.modules.jobs.context import JobContext from userbot.modules.jobs.registry import JOB_HANDLERS, JobHandler from utils.db.models import Job, JobStatus -from utils.env import env +from utils.listen import Listener from utils.logging import logger JOBS_CHANGED_CHANNEL = "jobs_changed" @@ -72,8 +72,12 @@ class JobConsumer: pruned = await repository.prune_finished(self.pool, self.account_id) if pruned: logger.info(f"[yellow]Pruned {pruned} finished job(s).[/]") - conn = await asyncpg.connect(dsn=env.db.connection_url) - await conn.add_listener(JOBS_CHANGED_CHANNEL, lambda *_: self._wake.set()) + listener = Listener( + f"jobs {self.account_id}", + {JOBS_CHANGED_CHANNEL: lambda _payload: self._wake.set()}, + self._wake.set, + ) + await listener.start() logger.info(f"[green]Job consumer running for account {self.account_id}.[/]") try: while True: @@ -84,5 +88,4 @@ class JobConsumer: self._wake.wait(), timeout=POLL_INTERVAL_SECONDS ) finally: - with contextlib.suppress(Exception): - await conn.close() + await listener.stop() diff --git a/backend/src/userbot/runner.py b/backend/src/userbot/runner.py index 36f8da7..db7f45e 100644 --- a/backend/src/userbot/runner.py +++ b/backend/src/userbot/runner.py @@ -15,6 +15,7 @@ from userbot.modules.client import PyroClient from userbot.modules.jobs import JobConsumer from utils.env import env from utils.jobs import enqueue +from utils.listen import Listener from utils.logging import logger, setup_logging from utils.read.accounts import ( ACCOUNTS_CHANGED_CHANNEL, @@ -148,9 +149,7 @@ async def _flush_sessions(registry: AccountRegistry) -> None: await client.storage.save() -async def _listen_changes( - registry: AccountRegistry, tasks: set[asyncio.Task] -) -> asyncpg.Connection: +def _change_listener(registry: AccountRegistry, tasks: set[asyncio.Task]) -> Listener: def spawn(coro: Coroutine[Any, Any, None]) -> None: task = asyncio.create_task(coro) tasks.add(task) @@ -163,26 +162,29 @@ async def _listen_changes( if client.capture is not None: spawn(make_coro(client.capture)) - def on_policy( - _conn: asyncpg.Connection, _pid: int, _channel: str, _payload: str - ) -> None: + def on_policy(_payload: str) -> None: reload(lambda capture: capture.reload_policies()) - def on_watch( - _conn: asyncpg.Connection, _pid: int, _channel: str, _payload: str - ) -> None: + def on_watch(_payload: str) -> None: reload(lambda capture: capture.watches.refresh()) - def on_accounts( - _conn: asyncpg.Connection, _pid: int, _channel: str, _payload: str - ) -> None: + def on_accounts(_payload: str) -> None: spawn(registry.sync()) - conn = await asyncpg.connect(dsn=env.db.connection_url) - await conn.add_listener("policy_changed", on_policy) - await conn.add_listener(WATCHES_CHANGED_CHANNEL, on_watch) - await conn.add_listener(ACCOUNTS_CHANGED_CHANNEL, on_accounts) - return conn + def on_reconnect() -> None: + on_policy("") + on_watch("") + on_accounts("") + + return Listener( + "userbot", + { + "policy_changed": on_policy, + WATCHES_CHANGED_CHANNEL: on_watch, + ACCOUNTS_CHANGED_CHANNEL: on_accounts, + }, + on_reconnect, + ) async def runner() -> None: @@ -191,7 +193,7 @@ async def runner() -> None: registry = AccountRegistry(pool, storage) tasks: set[asyncio.Task] = set() - listen_conn: asyncpg.Connection | None = None + listener = _change_listener(registry, tasks) stop = asyncio.Event() asyncio.get_running_loop().add_signal_handler(signal.SIGTERM, stop.set) flush = asyncio.create_task(_flush_sessions(registry)) @@ -199,14 +201,12 @@ async def runner() -> None: await registry.sync() if not registry.clients: logger.warning("[yellow]No sessions yet. Add an account in the web UI.[/]") - listen_conn = await _listen_changes(registry, tasks) + await listener.start() logger.info("[green]Userbot running.[/]") await stop.wait() finally: await _cancel(flush) - if listen_conn is not None: - with contextlib.suppress(Exception): - await listen_conn.close() + await listener.stop() await registry.close() await container.close() diff --git a/backend/src/utils/listen.py b/backend/src/utils/listen.py new file mode 100644 index 0000000..4afb356 --- /dev/null +++ b/backend/src/utils/listen.py @@ -0,0 +1,87 @@ +import asyncio +import contextlib +from collections.abc import Callable, Mapping + +import asyncpg + +from utils.env import env +from utils.logging import logger + +RETRY_MIN_SECONDS = 1.0 +RETRY_MAX_SECONDS = 60.0 +PING_INTERVAL_SECONDS = 60.0 +PING_TIMEOUT_SECONDS = 10.0 +CLOSE_TIMEOUT_SECONDS = 5.0 + +OnNotify = Callable[[str], None] + + +class Listener: + def __init__( + self, + name: str, + channels: Mapping[str, OnNotify], + on_reconnect: Callable[[], None] | None = None, + ) -> None: + self._name = name + self._channels = channels + self._on_reconnect = on_reconnect + self._task: asyncio.Task | None = None + self._ready = asyncio.Event() + self._listening = False + + async def start(self) -> None: + self._task = asyncio.create_task(self._run()) + await self._ready.wait() + + async def stop(self) -> None: + if self._task is None: + return + self._task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._task + self._task = None + + async def _run(self) -> None: + delay = RETRY_MIN_SECONDS + while True: + try: + await self._listen() + except Exception as exc: + if self._listening: + delay = RETRY_MIN_SECONDS + logger.warning( + f"{self._name}: no listen connection ({exc!r}), " + f"reconnecting in {delay:.0f}s" + ) + self._listening = False + self._ready.set() + await asyncio.sleep(delay) + delay = min(delay * 2, RETRY_MAX_SECONDS) + + async def _listen(self) -> None: + conn = await asyncpg.connect(dsn=env.db.connection_url) + try: + lost = asyncio.Event() + conn.add_termination_listener(lambda _conn: lost.set()) + for channel, on_notify in self._channels.items(): + await conn.add_listener( + channel, + lambda _conn, _pid, _channel, payload, f=on_notify: f(payload), + ) + logger.info(f"{self._name}: listening on {', '.join(self._channels)}") + self._listening = True + if self._ready.is_set() and self._on_reconnect is not None: + self._on_reconnect() + self._ready.set() + while True: + try: + await asyncio.wait_for(lost.wait(), PING_INTERVAL_SECONDS) + except TimeoutError: + await conn.fetchval("SELECT 1", timeout=PING_TIMEOUT_SECONDS) + else: + msg = "closed by the server" + raise ConnectionError(msg) + finally: + with contextlib.suppress(Exception): + await conn.close(timeout=CLOSE_TIMEOUT_SECONDS) diff --git a/backend/src/watcher/runner.py b/backend/src/watcher/runner.py index 70ed7e9..8040057 100644 --- a/backend/src/watcher/runner.py +++ b/backend/src/watcher/runner.py @@ -6,7 +6,7 @@ import asyncpg import uvloop from dependencies.container import container -from utils.env import env +from utils.listen import Listener from utils.logging import logger, setup_logging from utils.read.bots import BOTS_CHANGED_CHANNEL from utils.storage import ContentAddressedStorage @@ -24,19 +24,13 @@ async def _retry_loop(registry: BotRegistry) -> None: await registry.sync() -async def _listen_changes( - registry: BotRegistry, tasks: set[asyncio.Task] -) -> asyncpg.Connection: - def on_change( - _conn: asyncpg.Connection, _pid: int, _channel: str, _payload: str - ) -> None: +def _change_listener(registry: BotRegistry, tasks: set[asyncio.Task]) -> Listener: + def on_change(_payload: str) -> None: task = asyncio.create_task(registry.sync()) tasks.add(task) task.add_done_callback(tasks.discard) - conn = await asyncpg.connect(dsn=env.db.connection_url) - await conn.add_listener(BOTS_CHANGED_CHANNEL, on_change) - return conn + return Listener("watcher", {BOTS_CHANGED_CHANNEL: on_change}, lambda: on_change("")) async def runner() -> None: @@ -44,7 +38,7 @@ async def runner() -> None: storage = await container.get(ContentAddressedStorage) registry = BotRegistry(pool, storage) tasks: set[asyncio.Task] = set() - listen_conn: asyncpg.Connection | None = None + listener = _change_listener(registry, tasks) stop = asyncio.Event() asyncio.get_running_loop().add_signal_handler(signal.SIGTERM, stop.set) retry = asyncio.create_task(_retry_loop(registry)) @@ -52,16 +46,14 @@ async def runner() -> None: await registry.sync() if not registry.running: logger.warning("[yellow]No bots yet. Add a bot token in the web UI.[/]") - listen_conn = await _listen_changes(registry, tasks) + await listener.start() logger.info("[green]Bot watcher running.[/]") await stop.wait() finally: retry.cancel() with contextlib.suppress(asyncio.CancelledError): await retry - if listen_conn is not None: - with contextlib.suppress(Exception): - await listen_conn.close() + await listener.stop() await registry.close() await container.close()