Postgres backend: handle connection loss

broadcaster/_backends/postgres.py

@@ -14,6 +14,7 @@ def __init__(self, url: str):
     async def connect(self) -> None:
         self._conn = await asyncpg.connect(self._url)
         self._listen_queue: asyncio.Queue = asyncio.Queue()
+        self._conn.add_termination_listener(self._termination_listener)
 
     async def disconnect(self) -> None:
         await self._conn.close()
@@ -27,10 +28,13 @@ async def unsubscribe(self, channel: str) -> None:
     async def publish(self, channel: str, message: str) -> None:
         await self._conn.execute("SELECT pg_notify($1, $2);", channel, message)
 
-    def _listener(self, *args: Any) -> None:
+    async def _listener(self, *args: Any) -> None:
         connection, pid, channel, payload = args
         event = Event(channel=channel, message=payload)
-        self._listen_queue.put_nowait(event)
+        await self._listen_queue.put(event)
+
+    async def _termination_listener(self, *args: Any) -> None:
+        await self._listen_queue.put(None)
 
     async def next_published(self) -> Event:
         return await self._listen_queue.get()

broadcaster/_base.py

@@ -72,9 +72,17 @@ async def disconnect(self) -> None:
     async def _listener(self) -> None:
         while True:
             event = await self._backend.next_published()
+            if event is None:
+                # Backend is disconnected
+                break
+
             for queue in list(self._subscribers.get(event.channel, [])):
                 await queue.put(event)
 
+        # Ubsubscribe all
+        for queue in sum([list(qs) for qs in self._subscribers.values()], []):
+            await queue.put(None)
+
     async def publish(self, channel: str, message: Any) -> None:
         await self._backend.publish(channel, message)