fix: close the platform events websocket iterator on shutdown

Author: vdusek

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

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: