From cf153d621304247997fac21457c89e89af474636 Mon Sep 17 00:00:00 2001 From: Kyzgor Date: Mon, 8 Jun 2026 18:50:30 +0100 Subject: [PATCH] Fix per-reconnect task leak in client run loop The reconnect loop raced the reader task against `self._disconnect_signal.wait()` using `asyncio.as_completed([...])` followed by `break` after the first completes. Breaking out of the generator left the other future pending, so on every reconnect one task (and the as_completed machinery) leaked, accumulating unboundedly over a long-lived client's reconnects. Replace it with `asyncio.wait(..., return_when=FIRST_COMPLETED)` and cancel the future that did not complete, while still awaiting the reader task so that websocket-layer exceptions still propagate (preserving the reconnect-on-error behaviour). Adds a regression test asserting the disconnect-signal waiter count stays bounded across many reconnects. --- fastapi_websocket_pubsub/pub_sub_client.py | 32 +++++++-- tests/reconnect_memory_leak_test.py | 80 ++++++++++++++++++++++ 2 files changed, 106 insertions(+), 6 deletions(-) create mode 100644 tests/reconnect_memory_leak_test.py diff --git a/fastapi_websocket_pubsub/pub_sub_client.py b/fastapi_websocket_pubsub/pub_sub_client.py index 893d715..7987173 100644 --- a/fastapi_websocket_pubsub/pub_sub_client.py +++ b/fastapi_websocket_pubsub/pub_sub_client.py @@ -201,12 +201,32 @@ async def run(self, uri, wait_on_reader=True): if wait_on_reader: # Wait on the internal RPC task or until we ar asked to terminate - keeping the client alive meanwhile # Waiting on the reader tasks allows us to receive exceptions raised on the websocket layer- indicating a need to reset the connection - wait_on_reader_task = client.wait_on_reader() - for task in asyncio.as_completed( - [wait_on_reader_task, self._disconnect_signal.wait()] - ): - await task - break + wait_on_reader_task = asyncio.ensure_future( + client.wait_on_reader() + ) + disconnect_wait_task = asyncio.ensure_future( + self._disconnect_signal.wait() + ) + done, pending = await asyncio.wait( + [wait_on_reader_task, disconnect_wait_task], + return_when=asyncio.FIRST_COMPLETED, + ) + # Cancel the future that did not complete so it is not + # left pending across reconnects. The previous + # asyncio.as_completed loop broke out without cancelling + # the other future, leaking one task (and the + # as_completed machinery) on every reconnect. + for pending_task in pending: + pending_task.cancel() + try: + await pending_task + except (asyncio.CancelledError, Exception): + pass + # Surface any exception raised by the reader task (e.g. a + # websocket-layer error indicating we must reset the + # connection), preserving the previous behaviour. + if wait_on_reader_task in done: + await wait_on_reader_task if not self._disconnect_signal.is_set(): logger.info(f"Connection gracefully closed by server -- Trying to reconnect.") except websockets.exceptions.WebSocketException as err: diff --git a/tests/reconnect_memory_leak_test.py b/tests/reconnect_memory_leak_test.py new file mode 100644 index 0000000..285110f --- /dev/null +++ b/tests/reconnect_memory_leak_test.py @@ -0,0 +1,80 @@ +import asyncio +import os +import sys + +import pytest +import uvicorn +from fastapi import FastAPI +from multiprocessing import Process + +from fastapi_websocket_rpc import RpcChannel + +sys.path.append( + os.path.abspath(os.path.join(os.path.dirname(__file__), os.path.pardir)) +) +from fastapi_websocket_pubsub import PubSubEndpoint, PubSubClient + +# Use a dedicated port so this test can run alongside the others. +PORT = int(os.environ.get("PORT") or "7991") +uri = f"ws://localhost:{PORT}/pubsub" + + +def setup_always_disconnect_server(): + """A server that closes the socket shortly after every connect, forcing the + client into a continual reconnect loop.""" + app = FastAPI() + + async def on_connect(channel: RpcChannel): + async def disconnect_soon(): + await asyncio.sleep(0.05) + await channel.socket.close() + + asyncio.create_task(disconnect_soon()) + + endpoint = PubSubEndpoint(on_connect=[on_connect]) + endpoint.register_route(app, path="/pubsub") + uvicorn.run(app, port=PORT) + + +@pytest.fixture() +def repeating_death_server(): + proc = Process(target=setup_always_disconnect_server, args=(), daemon=True) + proc.start() + yield proc + proc.kill() + + +@pytest.mark.asyncio +async def test_reconnect_loop_does_not_leak_disconnect_waiters(repeating_death_server): + """Each reconnect iteration races the reader task against + ``self._disconnect_signal.wait()``. The loser of that race must be cancelled, + otherwise a pending waiter accumulates on the disconnect signal on every + reconnect (unbounded). This asserts the waiter count stays bounded across + many reconnects; it fails if the cancellation is removed. + """ + reconnects = 0 + + async def on_disconnect(channel): + nonlocal reconnects + reconnects += 1 + + client = PubSubClient(on_disconnect=[on_disconnect]) + client.start_client(uri) + try: + # Wait until several reconnect cycles have actually happened. + async def wait_for_reconnects(): + while reconnects < 5: + await asyncio.sleep(0.05) + + await asyncio.wait_for(wait_for_reconnects(), 20) + + # `_disconnect_signal.wait()` registers a future on the Event's internal + # waiter list; cancelling the losing task removes it again. Without that + # cancellation the list grows by one per reconnect. + waiters = len(client._disconnect_signal._waiters) + assert waiters <= 2, ( + f"leaked {waiters} disconnect-signal waiters over {reconnects} " + f"reconnects (expected the per-reconnect waiter to be cancelled)" + ) + finally: + await client.disconnect()