fix: close the platform events websocket iterator on shutdown

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():

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

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: