Fix Worker LISTEN/NOTIFY connection leak on shutdown

chancy/worker.py

@@ -186,6 +186,8 @@ def __init__(
         self._queues: dict[str, Queue] = {}
         # The executors that the worker is currently using.
         self._executors: dict[str, Executor] = {}
+        # The connection used for LISTEN/NOTIFY, if notifications are enabled.
+        self._notification_connection: AsyncConnection | None = None
 
     async def start(self):
         """
@@ -455,17 +457,17 @@ async def _maintain_notifications(self):
             separate from the shared connection pool and is not counted against
             the pool's connection limit.
         """
-        connection = await AsyncConnection.connect(
+        self._notification_connection = await AsyncConnection.connect(
             self.chancy.dsn, autocommit=True
         )
-        await connection.execute(
+        await self._notification_connection.execute(
             sql.SQL("LISTEN {channel};").format(
                 channel=sql.Identifier(f"{self.chancy.prefix}events")
             )
         )
         self.chancy.log.info("Started listening for realtime notifications.")
         self._notifications_ready_event.set()
-        async for notification in connection.notifies():
+        async for notification in self._notification_connection.notifies():
             j = json.loads(notification.payload)
             await self.hub.emit(j.pop("t"), j)
 
@@ -825,6 +827,10 @@ async def stop(self) -> bool:
                 self._queues.clear()
                 while self._executors:
                     await asyncio.sleep(0.1)
+                # Close the notification connection if it exists
+                if self._notification_connection:
+                    await self._notification_connection.close()
+                    self._notification_connection = None
                 # And finally axe everything else started by this worker.
                 await self.manager.cancel_all()
         except TimeoutError:

tests/test_worker.py

@@ -130,3 +130,64 @@ async def test_immediate_processing(chancy: Chancy, worker: Worker):
     )
 
     assert result.state == QueuedJob.State.SUCCEEDED
+
+
+@pytest.mark.asyncio
+async def test_notification_connection_cleanup():
+    """
+    Test that the LISTEN/NOTIFY connection is properly closed when the worker
+    stops, preventing connection leaks.
+
+    This test doesn't use the chancy fixture to ensure we have full control
+    over connection lifecycle and can properly detect leaks.
+    """
+    from psycopg import AsyncConnection
+
+    dsn = "postgresql://postgres:localtest@localhost:8190/postgres"
+
+    # Get initial connection count using a separate connection
+    async with await AsyncConnection.connect(dsn) as conn:
+        result = await conn.execute(
+            """
+            SELECT COUNT(*) as count 
+            FROM pg_stat_activity 
+            WHERE datname = current_database()
+            AND application_name LIKE 'psycopg%'
+            """
+        )
+        initial_count = (await result.fetchone())[0]
+
+    # Create, start, and stop a worker with its own Chancy instance
+    async with Chancy(dsn) as chancy:
+        await chancy.migrate()
+        async with Worker(chancy) as worker:
+            # Wait for the worker to fully initialize
+            await asyncio.sleep(1)
+
+            # Verify the notification connection exists
+            assert worker._notification_connection is not None
+            assert not worker._notification_connection.closed
+
+    # Give a moment for connections to fully close
+    await asyncio.sleep(0.5)
+
+    # After everything is stopped, verify no connection leak
+    # Use a completely separate connection
+    async with await AsyncConnection.connect(dsn) as conn:
+        result = await conn.execute(
+            """
+            SELECT COUNT(*) as count 
+            FROM pg_stat_activity 
+            WHERE datname = current_database()
+            AND application_name LIKE 'psycopg%'
+            """
+        )
+        final_count = (await result.fetchone())[0]
+
+    # The connection count should be back to the initial count (or less)
+    # We allow for <= to account for the single connection we just created
+    # to check the count
+    assert final_count <= initial_count + 1, (
+        f"Connection leak detected: initial={initial_count}, "
+        f"final={final_count}"
+    )