diff --git a/chancy/worker.py b/chancy/worker.py index 5634250..50500f7 100644 --- a/chancy/worker.py +++ b/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: diff --git a/tests/test_worker.py b/tests/test_worker.py index 484c562..86b6fd4 100644 --- a/tests/test_worker.py +++ b/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}" + )