diff --git a/src/apify/events/_apify_event_manager.py b/src/apify/events/_apify_event_manager.py index e7b21da8..3f4e9512 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,35 @@ 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) + + # 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): + 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 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: + 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/e2e/test_actor_events.py b/tests/e2e/test_actor_events.py index d0af2b1d..da1a9378 100644 --- a/tests/e2e/test_actor_events.py +++ b/tests/e2e/test_actor_events.py @@ -130,3 +130,40 @@ 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 + + # 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.' + + 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/']) + + actor = await make_actor(label='actor-events-shutdown', main_func=main) + run_result = await run_actor(actor) + + assert run_result.status == 'SUCCEEDED' + + # 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 + + 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 diff --git a/tests/unit/events/test_apify_event_manager.py b/tests/unit/events/test_apify_event_manager.py index 8f723966..86589576 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,33 @@ 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 blocks garbage collection, so the assertion holds only if the manager + # closes it itself. + 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) + + assert iterators + # A closed async generator has no frame left. + 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: