fix(backend): listen connections reconnect after a dropped connection or a postgres restart and catch up on what they missed

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)

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()

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()
 

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)

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()