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

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)