From ae02fb8be67b23589c5981531409ef6e9990462c Mon Sep 17 00:00:00 2001 From: Vlada Dusek Date: Tue, 4 Aug 2026 11:39:04 +0200 Subject: [PATCH 1/5] fix: close the platform events websocket iterator on shutdown --- src/apify/events/_apify_event_manager.py | 61 +++++++++++-------- tests/unit/events/test_apify_event_manager.py | 31 +++++++++- 2 files changed, 64 insertions(+), 28 deletions(-) diff --git a/src/apify/events/_apify_event_manager.py b/src/apify/events/_apify_event_manager.py index e7b21da8..ae5e2900 100644 --- a/src/apify/events/_apify_event_manager.py +++ b/src/apify/events/_apify_event_manager.py @@ -4,7 +4,7 @@ import contextlib import time from logging import getLogger -from typing import TYPE_CHECKING, Annotated, Self +from typing import TYPE_CHECKING, Annotated, Self, cast import websockets.asyncio.client import websockets.client @@ -20,7 +20,7 @@ from apify.events._types import DeprecatedEvent, EventMessage, SystemInfoEventData, UnknownEvent if TYPE_CHECKING: - from collections.abc import Generator + from collections.abc import AsyncGenerator, Generator from types import TracebackType from crawlee.events._event_manager import EventManagerOptions @@ -149,31 +149,38 @@ async def _process_platform_messages(self, ws_url: str) -> None: try: # Used as an async iterator, `connect` reconnects with exponential backoff on failed connection attempts. - async for websocket in websockets.asyncio.client.connect( - ws_url, process_exception=self._process_connection_exception - ): - self._platform_events_websocket = websocket - if self._connected_to_platform_websocket and not self._connected_to_platform_websocket.done(): - self._connected_to_platform_websocket.set_result(True) - else: - logger.info('Reconnected to the platform events websocket.') - - connection_opened_at = time.monotonic() - connection_lost = await self._consume_messages(websocket) - - if not self._should_reconnect_after_close(websocket, connection_lost=connection_lost): - break - - # Reconnect a healthy connection immediately; back off only on repeated rapid drops. The first - # rapid drop reconnects once without delay (it only primes the backoff generator), and each - # subsequent consecutive rapid drop then sleeps for the next backoff interval. A healthy - # connection resets the generator, so the next rapid drop again gets that one free retry. - if time.monotonic() - connection_opened_at >= self._HEALTHY_CONNECTION_MIN_DURATION: - backoff_delays = None - elif backoff_delays is None: - backoff_delays = websockets.client.backoff() - else: - await asyncio.sleep(next(backoff_delays)) + connector = websockets.asyncio.client.connect(ws_url, process_exception=self._process_connection_exception) + + # That iterator is an async generator whose cleanup closes the current connection, and neither `break` nor + # cancelling this task closes it. Left suspended, it is closed by the garbage collector instead - on Actor + # exit that lands in `asyncio.run` teardown, too late to be awaited, which then breaks the event loop's + # async generator shutdown. `websockets` types the iterator as a plain `AsyncIterator`, hence the cast. + connections = cast('AsyncGenerator[websockets.asyncio.client.ClientConnection]', aiter(connector)) + + async with contextlib.aclosing(connections): + async for websocket in connections: + self._platform_events_websocket = websocket + if self._connected_to_platform_websocket and not self._connected_to_platform_websocket.done(): + self._connected_to_platform_websocket.set_result(True) + else: + logger.info('Reconnected to the platform events websocket.') + + connection_opened_at = time.monotonic() + connection_lost = await self._consume_messages(websocket) + + if not self._should_reconnect_after_close(websocket, connection_lost=connection_lost): + break + + # Reconnect a healthy connection immediately; back off only on repeated rapid drops. The first + # rapid drop reconnects once without delay (it only primes the backoff generator), and each + # subsequent consecutive rapid drop then sleeps for the next backoff interval. A healthy + # connection resets the generator, so the next rapid drop again gets that one free retry. + if time.monotonic() - connection_opened_at >= self._HEALTHY_CONNECTION_MIN_DURATION: + backoff_delays = None + elif backoff_delays is None: + backoff_delays = websockets.client.backoff() + else: + await asyncio.sleep(next(backoff_delays)) except Exception: logger.exception('Error in websocket connection') if self._connected_to_platform_websocket is not None and not self._connected_to_platform_websocket.done(): diff --git a/tests/unit/events/test_apify_event_manager.py b/tests/unit/events/test_apify_event_manager.py index 8f723966..77171c8e 100644 --- a/tests/unit/events/test_apify_event_manager.py +++ b/tests/unit/events/test_apify_event_manager.py @@ -13,6 +13,7 @@ import pytest import websockets +import websockets.asyncio.client import websockets.asyncio.server from crawlee.events._types import Event @@ -24,7 +25,7 @@ from apify.events._types import SystemInfoEventData if TYPE_CHECKING: - from collections.abc import AsyncGenerator, Awaitable, Callable + from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable DUMMY_SYSTEM_INFO = { @@ -595,6 +596,34 @@ async def test_shutdown_during_reconnect_backoff_is_clean(monkeypatch: pytest.Mo assert persist_state_task is None or persist_state_task.done() +async def test_exit_closes_the_reconnecting_iterator(monkeypatch: pytest.MonkeyPatch) -> None: + """Test that exiting closes the `connect` async iterator itself, rather than leaving it to the garbage collector.""" + # Keeping a reference to every iterator prevents the garbage collector from finalizing it, so the assertion below + # holds only if the event manager closes the iterator on its own. + iterators: list[Any] = [] + original_aiter = websockets.asyncio.client.connect.__aiter__ + + def recording_aiter( + connector: websockets.asyncio.client.connect, + ) -> AsyncIterator[websockets.asyncio.client.ClientConnection]: + iterator = original_aiter(connector) + iterators.append(iterator) + return iterator + + monkeypatch.setattr(websockets.asyncio.client.connect, '__aiter__', recording_aiter) + + async with ( + _platform_ws_server(monkeypatch) as (_, client_connected), + ApifyEventManager(Configuration.get_global_configuration()), + ): + await asyncio.wait_for(client_connected.wait(), timeout=10) + + # An iterator left suspended keeps its websocket open, and closing it is then deferred to the garbage collector. + # On Actor exit that lands in `asyncio.run` teardown, too late to be awaited, breaking async generator shutdown. + assert iterators + assert all(iterator.ag_frame is None for iterator in iterators) + + async def test_malformed_message_logs_exception( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture ) -> None: From 9f66e28da643cd3079db93c2b2374b6dcd3e8607 Mon Sep 17 00:00:00 2001 From: Vlada Dusek Date: Tue, 4 Aug 2026 11:49:24 +0200 Subject: [PATCH 2/5] style: condense comments around the events websocket iterator --- src/apify/events/_apify_event_manager.py | 7 +++---- tests/unit/events/test_apify_event_manager.py | 7 +++---- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/src/apify/events/_apify_event_manager.py b/src/apify/events/_apify_event_manager.py index ae5e2900..bfa86f93 100644 --- a/src/apify/events/_apify_event_manager.py +++ b/src/apify/events/_apify_event_manager.py @@ -151,10 +151,9 @@ async def _process_platform_messages(self, ws_url: str) -> None: # Used as an async iterator, `connect` reconnects with exponential backoff on failed connection attempts. connector = websockets.asyncio.client.connect(ws_url, process_exception=self._process_connection_exception) - # That iterator is an async generator whose cleanup closes the current connection, and neither `break` nor - # cancelling this task closes it. Left suspended, it is closed by the garbage collector instead - on Actor - # exit that lands in `asyncio.run` teardown, too late to be awaited, which then breaks the event loop's - # async generator shutdown. `websockets` types the iterator as a plain `AsyncIterator`, hence the cast. + # Close the iterator explicitly: neither `break` nor cancellation does it, and leaving it to the garbage + # collector lands its cleanup in `asyncio.run` teardown, too late to be awaited, which breaks the event + # loop's async generator shutdown. The cast is only because `websockets` types it as `AsyncIterator`. connections = cast('AsyncGenerator[websockets.asyncio.client.ClientConnection]', aiter(connector)) async with contextlib.aclosing(connections): diff --git a/tests/unit/events/test_apify_event_manager.py b/tests/unit/events/test_apify_event_manager.py index 77171c8e..ab21c088 100644 --- a/tests/unit/events/test_apify_event_manager.py +++ b/tests/unit/events/test_apify_event_manager.py @@ -598,8 +598,8 @@ async def test_shutdown_during_reconnect_backoff_is_clean(monkeypatch: pytest.Mo async def test_exit_closes_the_reconnecting_iterator(monkeypatch: pytest.MonkeyPatch) -> None: """Test that exiting closes the `connect` async iterator itself, rather than leaving it to the garbage collector.""" - # Keeping a reference to every iterator prevents the garbage collector from finalizing it, so the assertion below - # holds only if the event manager closes the iterator on its own. + # Holding a reference to every iterator blocks garbage collection, so the assertion below holds only if the event + # manager closes the iterator itself. iterators: list[Any] = [] original_aiter = websockets.asyncio.client.connect.__aiter__ @@ -618,9 +618,8 @@ def recording_aiter( ): await asyncio.wait_for(client_connected.wait(), timeout=10) - # An iterator left suspended keeps its websocket open, and closing it is then deferred to the garbage collector. - # On Actor exit that lands in `asyncio.run` teardown, too late to be awaited, breaking async generator shutdown. assert iterators + # A closed async generator has no frame left. assert all(iterator.ag_frame is None for iterator in iterators) From 1b7132cf8452781f84886c6504064ab9d6061e9a Mon Sep 17 00:00:00 2001 From: Vlada Dusek Date: Tue, 4 Aug 2026 11:49:25 +0200 Subject: [PATCH 3/5] test: add e2e test for a clean events websocket shutdown --- tests/e2e/test_actor_events.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/tests/e2e/test_actor_events.py b/tests/e2e/test_actor_events.py index d0af2b1d..4d65feca 100644 --- a/tests/e2e/test_actor_events.py +++ b/tests/e2e/test_actor_events.py @@ -130,3 +130,35 @@ async def wait_for_events(n: int) -> None: run_result = await run_actor(actor) assert run_result.status == 'SUCCEEDED' + + +async def test_events_websocket_shutdown_is_clean( + make_actor: MakeActorFunction, + run_actor: RunActorFunction, +) -> None: + """Test that a run using the platform events websocket exits without an async generator shutdown error.""" + + async def main() -> None: + from crawlee.crawlers import ParselCrawler, ParselCrawlingContext + + # The crawler run keeps the event loop busy, so an unclosed events websocket iterator would be + # garbage-collected only during interpreter shutdown - the timing that triggers the error. + async with Actor: + crawler = ParselCrawler(max_crawl_depth=2) + + @crawler.router.default_handler + async def handler(context: ParselCrawlingContext) -> None: + await context.enqueue_links() + + await crawler.run(['http://localhost:8080/']) + + actor = await make_actor(label='actor-events-shutdown', main_func=main) + run_result = await run_actor(actor) + + assert run_result.status == 'SUCCEEDED' + assert run_result.exit_code == 0 + + run_log = await actor.last_run().log().get() + assert run_log is not None + assert 'error occurred during closing of asynchronous generator' not in run_log + assert 'asynchronous generator is already running' not in run_log From 9c8274ad3965d0712b4a31d59369072b7317c969 Mon Sep 17 00:00:00 2001 From: Vlada Dusek Date: Tue, 4 Aug 2026 13:19:52 +0200 Subject: [PATCH 4/5] test: assert the events websocket shutdown e2e test really crawled --- tests/e2e/test_actor_events.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/tests/e2e/test_actor_events.py b/tests/e2e/test_actor_events.py index 4d65feca..0e858978 100644 --- a/tests/e2e/test_actor_events.py +++ b/tests/e2e/test_actor_events.py @@ -141,13 +141,16 @@ async def test_events_websocket_shutdown_is_clean( async def main() -> None: from crawlee.crawlers import ParselCrawler, ParselCrawlingContext - # The crawler run keeps the event loop busy, so an unclosed events websocket iterator would be - # garbage-collected only during interpreter shutdown - the timing that triggers the error. + # A real crawler run is needed: the concurrency around Actor exit is what leaves an unclosed events + # websocket iterator's finalizer still in flight when the event loop tears down. async with Actor: + assert Actor.configuration.actor_events_ws_url, 'The run must use the platform events websocket.' + crawler = ParselCrawler(max_crawl_depth=2) @crawler.router.default_handler async def handler(context: ParselCrawlingContext) -> None: + await context.push_data({'url': context.request.url}) await context.enqueue_links() await crawler.run(['http://localhost:8080/']) @@ -156,7 +159,10 @@ async def handler(context: ParselCrawlingContext) -> None: run_result = await run_actor(actor) assert run_result.status == 'SUCCEEDED' - assert run_result.exit_code == 0 + + # The log assertions below are negative, so confirm the crawl that drives them really visited pages. + dataset_items_page = await actor.last_run().dataset().list_items() + assert dataset_items_page.count > 1 run_log = await actor.last_run().log().get() assert run_log is not None From 2d231f2fc7d5e3a3a29b658e25fd98ced284b315 Mon Sep 17 00:00:00 2001 From: Vlada Dusek Date: Tue, 4 Aug 2026 13:25:12 +0200 Subject: [PATCH 5/5] style: shorten comments around the events websocket iterator --- src/apify/events/_apify_event_manager.py | 10 ++++------ tests/e2e/test_actor_events.py | 5 ++--- tests/unit/events/test_apify_event_manager.py | 4 ++-- 3 files changed, 8 insertions(+), 11 deletions(-) diff --git a/src/apify/events/_apify_event_manager.py b/src/apify/events/_apify_event_manager.py index bfa86f93..3f4e9512 100644 --- a/src/apify/events/_apify_event_manager.py +++ b/src/apify/events/_apify_event_manager.py @@ -151,9 +151,8 @@ async def _process_platform_messages(self, ws_url: str) -> None: # Used as an async iterator, `connect` reconnects with exponential backoff on failed connection attempts. connector = websockets.asyncio.client.connect(ws_url, process_exception=self._process_connection_exception) - # Close the iterator explicitly: neither `break` nor cancellation does it, and leaving it to the garbage - # collector lands its cleanup in `asyncio.run` teardown, too late to be awaited, which breaks the event - # loop's async generator shutdown. The cast is only because `websockets` types it as `AsyncIterator`. + # Neither `break` nor cancellation closes the iterator; left to the garbage collector, its cleanup lands + # in `asyncio.run` teardown, too late to await, and breaks the loop's async generator shutdown. connections = cast('AsyncGenerator[websockets.asyncio.client.ClientConnection]', aiter(connector)) async with contextlib.aclosing(connections): @@ -171,9 +170,8 @@ async def _process_platform_messages(self, ws_url: str) -> None: break # Reconnect a healthy connection immediately; back off only on repeated rapid drops. The first - # rapid drop reconnects once without delay (it only primes the backoff generator), and each - # subsequent consecutive rapid drop then sleeps for the next backoff interval. A healthy - # connection resets the generator, so the next rapid drop again gets that one free retry. + # rapid drop retries without delay (it just primes the backoff generator), each consecutive one + # after it sleeps for the next interval. A healthy connection resets the generator. if time.monotonic() - connection_opened_at >= self._HEALTHY_CONNECTION_MIN_DURATION: backoff_delays = None elif backoff_delays is None: diff --git a/tests/e2e/test_actor_events.py b/tests/e2e/test_actor_events.py index 0e858978..da1a9378 100644 --- a/tests/e2e/test_actor_events.py +++ b/tests/e2e/test_actor_events.py @@ -141,8 +141,7 @@ async def test_events_websocket_shutdown_is_clean( async def main() -> None: from crawlee.crawlers import ParselCrawler, ParselCrawlingContext - # A real crawler run is needed: the concurrency around Actor exit is what leaves an unclosed events - # websocket iterator's finalizer still in flight when the event loop tears down. + # Real crawler load is needed: it leaves an unclosed iterator's finalizer in flight at loop teardown. async with Actor: assert Actor.configuration.actor_events_ws_url, 'The run must use the platform events websocket.' @@ -160,7 +159,7 @@ async def handler(context: ParselCrawlingContext) -> None: assert run_result.status == 'SUCCEEDED' - # The log assertions below are negative, so confirm the crawl that drives them really visited pages. + # The log assertions below are negative, so confirm the crawl that drives them actually ran. dataset_items_page = await actor.last_run().dataset().list_items() assert dataset_items_page.count > 1 diff --git a/tests/unit/events/test_apify_event_manager.py b/tests/unit/events/test_apify_event_manager.py index ab21c088..86589576 100644 --- a/tests/unit/events/test_apify_event_manager.py +++ b/tests/unit/events/test_apify_event_manager.py @@ -598,8 +598,8 @@ async def test_shutdown_during_reconnect_backoff_is_clean(monkeypatch: pytest.Mo async def test_exit_closes_the_reconnecting_iterator(monkeypatch: pytest.MonkeyPatch) -> None: """Test that exiting closes the `connect` async iterator itself, rather than leaving it to the garbage collector.""" - # Holding a reference to every iterator blocks garbage collection, so the assertion below holds only if the event - # manager closes the iterator itself. + # Keeping a reference to every iterator blocks garbage collection, so the assertion holds only if the manager + # closes it itself. iterators: list[Any] = [] original_aiter = websockets.asyncio.client.connect.__aiter__