fix(events): open the platform websocket only once per context

Author: vdusek

src/apify/_actor.py

@@ -1249,12 +1249,10 @@ async def reboot(
             # the reboot. Typically, crawlers are listening for the MIGRATING event to stop processing new requests.
             # We can't just emit the events and wait for all listeners to finish,
             # because this method might be called from an event listener itself, and we would deadlock.
-            persist_state_listeners = flatten(
-                (self.event_manager._listeners_to_wrappers[Event.PERSIST_STATE] or {}).values()  # noqa: SLF001
-            )
-            migrating_listeners = flatten(
-                (self.event_manager._listeners_to_wrappers[Event.MIGRATING] or {}).values()  # noqa: SLF001
-            )
+            # Read the mapping with `get` - subscripting it would insert entries for events nobody listens to.
+            listeners_to_wrappers = self.event_manager._listeners_to_wrappers  # noqa: SLF001
+            persist_state_listeners = flatten(listeners_to_wrappers.get(Event.PERSIST_STATE, {}).values())
+            migrating_listeners = flatten(listeners_to_wrappers.get(Event.MIGRATING, {}).values())
 
             async def safe_dispatch(listener: Any, data: Any) -> None:
                 try:

src/apify/events/_apify_event_manager.py

@@ -86,29 +86,47 @@ def __init__(self, configuration: Configuration, **kwargs: Unpack[EventManagerOp
         self._platform_events_websocket: websockets.asyncio.client.ClientConnection | None = None
         """WebSocket connection to the platform events."""
 
-        self._process_platform_messages_task: asyncio.Task | None = None
+        self._process_platform_messages_task: asyncio.Task[None] | None = None
         """Task for processing messages from the platform websocket."""
 
-        self._connected_to_platform_websocket: asyncio.Future[bool] | None = None
-        """Future that resolves when the connection to the platform websocket is established."""
+        self._connected_to_platform_websocket: asyncio.Future[None] | None = None
+        """Resolves once the platform websocket is connected, or fails with the error that prevented the first
+        connection, so that `__aenter__` can report it.
+        """
 
     @override
     async def __aenter__(self) -> Self:
+        """Initialize the event manager upon entering the async context.
+
+        On the outermost entry, it connects to the platform events websocket and starts consuming its messages.
+        """
         await super().__aenter__()
-        self._connected_to_platform_websocket = asyncio.Future()
 
-        # Run tasks but don't await them
-        if self._configuration.actor_events_ws_url:
-            self._process_platform_messages_task = asyncio.create_task(
-                self._process_platform_messages(self._configuration.actor_events_ws_url)
-            )
-            is_connected = await self._connected_to_platform_websocket
-            if not is_connected:
-                # Exit the already-entered parent so the recurring persist state task does not leak.
-                await self.__aexit__(None, None, None)
-                raise RuntimeError('Error connecting to platform events websocket!')
-        else:
+        # Only the outermost context owns the websocket, so a nested one must not open a second connection.
+        if self._active_ref_count > 1:
+            return self
+
+        if not self._configuration.actor_events_ws_url:
             logger.debug('APIFY_ACTOR_EVENTS_WS_URL env var not set, no events from Apify platform will be emitted.')
+            return self
+
+        # The future has to exist before the task that resolves it starts running.
+        self._connected_to_platform_websocket = asyncio.Future()
+        self._process_platform_messages_task = asyncio.create_task(
+            self._process_platform_messages(self._configuration.actor_events_ws_url)
+        )
+
+        try:
+            await self._connected_to_platform_websocket
+        except Exception as exc:
+            # Exit the already-entered parent so the recurring persist state task does not leak.
+            await self.__aexit__(None, None, None)
+            raise RuntimeError('Error connecting to platform events websocket!') from exc
+        except BaseException:
+            # Cancellation has to clean up as well. A stale task left behind would make the next entry look nested,
+            # returning a manager that silently receives no platform events at all.
+            await self.__aexit__(None, None, None)
+            raise
 
         return self
 
@@ -119,17 +137,35 @@ async def __aexit__(
         exc_value: BaseException | None,
         exc_traceback: TracebackType | None,
     ) -> None:
-        # Cancel the task before closing the websocket so that the closed connection is not treated as a drop
-        # and followed by a reconnect attempt.
-        if self._process_platform_messages_task and not self._process_platform_messages_task.done():
-            self._process_platform_messages_task.cancel()
-            with contextlib.suppress(asyncio.CancelledError):
-                await self._process_platform_messages_task
+        """Close the event manager upon exiting the async context.
 
-        if self._platform_events_websocket:
-            await self._platform_events_websocket.close()
-
-        await super().__aexit__(exc_type, exc_value, exc_traceback)
+        On the outermost exit, it stops consuming the platform messages and closes the websocket connection.
+        """
+        try:
+            if self._active_ref_count == 1:
+                await self._teardown_platform_websocket()
+        finally:
+            # The parent context has to be left even if the shutdown above fails. Staying active would mean never
+            # emitting `PersistState` again, as re-entering the context would be a no-op.
+            await super().__aexit__(exc_type, exc_value, exc_traceback)
+
+    async def _teardown_platform_websocket(self) -> None:
+        """Stop consuming the platform messages and close the websocket connection to the platform events."""
+        try:
+            # Cancel the task before closing the websocket so that the closed connection is not treated as a drop
+            # and followed by a reconnect attempt.
+            if self._process_platform_messages_task and not self._process_platform_messages_task.done():
+                self._process_platform_messages_task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await self._process_platform_messages_task
+
+            if self._platform_events_websocket:
+                await self._platform_events_websocket.close()
+        finally:
+            # Leave no closed connection or resolved future behind, so that the context can be entered again.
+            self._process_platform_messages_task = None
+            self._platform_events_websocket = None
+            self._connected_to_platform_websocket = None
 
     def _process_connection_exception(self, exc: Exception) -> Exception | None:
         """Decide whether a failed connection attempt to the platform websocket should be retried.
@@ -159,7 +195,7 @@ async def _process_platform_messages(self, ws_url: str) -> None:
                 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)
+                        self._connected_to_platform_websocket.set_result(None)
                     else:
                         logger.info('Reconnected to the platform events websocket.')
 
@@ -178,10 +214,12 @@ async def _process_platform_messages(self, ws_url: str) -> None:
                         backoff_delays = websockets.client.backoff()
                     else:
                         await asyncio.sleep(next(backoff_delays))
-        except Exception:
+        except Exception as exc:
             logger.exception('Error in websocket connection')
+
             if self._connected_to_platform_websocket is not None and not self._connected_to_platform_websocket.done():
-                self._connected_to_platform_websocket.set_result(False)
+                # `__aenter__` is still waiting for the first connection, so let it fail with this as the cause.
+                self._connected_to_platform_websocket.set_exception(exc)
 
     async def _consume_messages(self, websocket: websockets.asyncio.client.ClientConnection) -> bool:
         """Handle platform messages until the connection closes; return whether it was lost vs. closed cleanly."""

tests/unit/events/test_apify_event_manager.py

@@ -151,6 +151,33 @@ async def start() -> None:
         await stop()
 
 
+@contextlib.asynccontextmanager
+async def _unresponsive_ws_server(monkeypatch: pytest.MonkeyPatch) -> AsyncGenerator[None]:
+    """A `127.0.0.1` server that accepts connections but never completes the WebSocket handshake.
+
+    It keeps `__aenter__` waiting for its first connection, which is what lets a test cancel it mid-connect.
+    """
+    shutdown = asyncio.Event()
+
+    async def handler(_reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
+        try:
+            await shutdown.wait()
+        finally:
+            writer.close()
+
+    server = await asyncio.start_server(handler, host='127.0.0.1')
+    port: int = server.sockets[0].getsockname()[1]
+    monkeypatch.setenv(ActorEnvVars.EVENTS_WEBSOCKET_URL, f'ws://127.0.0.1:{port}')
+
+    try:
+        yield
+    finally:
+        # Release the handlers first, `wait_closed` would block on them otherwise.
+        shutdown.set()
+        server.close()
+        await server.wait_closed()
+
+
 async def test_lifecycle_local(caplog: pytest.LogCaptureFixture) -> None:
     caplog.set_level(logging.DEBUG, logger='apify')
 
@@ -260,10 +287,12 @@ async def test_lifecycle_on_platform_without_websocket(monkeypatch: pytest.Monke
     monkeypatch.setenv(ActorEnvVars.EVENTS_WEBSOCKET_URL, 'ws://localhost:56565')
     event_manager = ApifyEventManager(Configuration.get_global_configuration())
 
-    with pytest.raises(RuntimeError, match=r'Error connecting to platform events websocket!'):
+    with pytest.raises(RuntimeError, match=r'Error connecting to platform events websocket!') as exc_info:
         async with event_manager:
             pass
 
+    # The error that prevented the connection is reported as the cause, not only logged.
+    assert isinstance(exc_info.value.__cause__, OSError)
     assert event_manager.active is False
     persist_state_task = event_manager._emit_persist_state_event_rec_task.task
     assert persist_state_task is None or persist_state_task.done()
@@ -278,6 +307,119 @@ async def test_lifecycle_on_platform(monkeypatch: pytest.MonkeyPatch) -> None:
         assert len(connected_ws_clients) == 1
 
 
+async def test_nested_context_keeps_a_single_websocket(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A nested context reuses the single platform connection, and only the outermost exit tears it down."""
+    async with _platform_ws_server(monkeypatch) as (connected_ws_clients, client_connected):
+        event_manager = ApifyEventManager(Configuration.get_global_configuration())
+
+        async with event_manager:
+            await client_connected.wait()
+            assert len(connected_ws_clients) == 1
+            task = event_manager._process_platform_messages_task
+
+            # A crawler running under an Actor enters the already-entered event manager again.
+            async with event_manager:
+                await asyncio.sleep(0.2)
+                assert len(connected_ws_clients) == 1
+                assert event_manager._process_platform_messages_task is task
+
+            # The inner exit must leave the connection alone, the Actor still needs the platform events.
+            await asyncio.sleep(0.2)
+            assert len(connected_ws_clients) == 1
+            assert task is not None
+            assert not task.done()
+
+            # A single connection also means every event is delivered exactly once.
+            event_calls: list[Any] = []
+            event_manager.on(event=Event.SYSTEM_INFO, listener=event_calls.append)
+            websockets.broadcast(connected_ws_clients, json.dumps({'name': 'systemInfo', 'data': DUMMY_SYSTEM_INFO}))
+            await poll_until_condition(lambda: bool(event_calls), poll_interval=0.05)
+            await asyncio.sleep(0.2)
+            assert len(event_calls) == 1
+
+        # Poll because the server-side handler may not have deregistered its connection yet.
+        await poll_until_condition(lambda: not connected_ws_clients, poll_interval=0.05)
+        assert not connected_ws_clients
+        assert task.done()
+
+
+async def test_context_can_be_reentered_after_full_exit(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Entering a fully exited event manager again opens a fresh platform connection."""
+    async with _platform_ws_server(monkeypatch) as (connected_ws_clients, client_connected):
+        event_manager = ApifyEventManager(Configuration.get_global_configuration())
+
+        async with event_manager:
+            await client_connected.wait()
+            assert len(connected_ws_clients) == 1
+
+        await poll_until_condition(lambda: not connected_ws_clients, poll_interval=0.05)
+        assert event_manager._process_platform_messages_task is None
+        assert event_manager._platform_events_websocket is None
+
+        client_connected.clear()
+        async with event_manager:
+            await asyncio.wait_for(client_connected.wait(), timeout=10)
+            assert len(connected_ws_clients) == 1
+
+
+async def test_cancelled_entry_leaves_no_stale_state(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A cancelled entry releases the context, so the next entry connects again instead of looking like a nested one."""
+    async with _unresponsive_ws_server(monkeypatch):
+        event_manager = ApifyEventManager(Configuration.get_global_configuration())
+
+        first_entry = asyncio.create_task(event_manager.__aenter__())
+        await asyncio.sleep(0.2)
+        assert not first_entry.done()
+
+        first_entry.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await first_entry
+
+        assert event_manager.active is False
+        assert event_manager._process_platform_messages_task is None
+        assert event_manager._platform_events_websocket is None
+        persist_state_task = event_manager._emit_persist_state_event_rec_task.task
+        assert persist_state_task is None or persist_state_task.done()
+
+        # The next entry has to attempt a connection of its own, rather than return a manager receiving no events.
+        second_entry = asyncio.create_task(event_manager.__aenter__())
+        await asyncio.sleep(0.2)
+        assert not second_entry.done()
+        assert event_manager._active_ref_count == 1
+        assert event_manager._process_platform_messages_task is not None
+
+        second_entry.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await second_entry
+        assert event_manager.active is False
+
+
+async def test_exit_releases_context_when_the_websocket_shutdown_fails(monkeypatch: pytest.MonkeyPatch) -> None:
+    """A failing websocket shutdown still releases the context, so the manager cannot stay active for good."""
+    async with _platform_ws_server(monkeypatch) as (_, client_connected):
+        event_manager = ApifyEventManager(Configuration.get_global_configuration())
+        await event_manager.__aenter__()
+        await client_connected.wait()
+
+        monkeypatch.setattr(
+            event_manager, '_teardown_platform_websocket', Mock(side_effect=RuntimeError('close failed'))
+        )
+
+        with pytest.raises(RuntimeError, match='close failed'):
+            await event_manager.__aexit__(None, None, None)
+
+        assert event_manager.active is False
+        persist_state_task = event_manager._emit_persist_state_event_rec_task.task
+        assert persist_state_task is None or persist_state_task.done()
+
+        # The mocked shutdown left the message-processing task running.
+        task = event_manager._process_platform_messages_task
+        assert task is not None
+        task.cancel()
+        with contextlib.suppress(asyncio.CancelledError):
+            await task
+
+
 async def test_event_handling_on_platform(monkeypatch: pytest.MonkeyPatch) -> None:
     async with _platform_ws_server(monkeypatch) as (connected_ws_clients, client_connected):