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

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