From e3518a71cbdf1b8e86808ce0bd4e3ca1b10965cd Mon Sep 17 00:00:00 2001 From: Vlada Dusek Date: Thu, 6 Aug 2026 12:53:10 +0200 Subject: [PATCH] fix(events): open the platform websocket only once per context --- src/apify/_actor.py | 10 +- src/apify/events/_apify_event_manager.py | 94 ++++++++---- tests/unit/events/test_apify_event_manager.py | 144 +++++++++++++++++- 3 files changed, 213 insertions(+), 35 deletions(-) diff --git a/src/apify/_actor.py b/src/apify/_actor.py index 862370cd..be798c42 100644 --- a/src/apify/_actor.py +++ b/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: diff --git a/src/apify/events/_apify_event_manager.py b/src/apify/events/_apify_event_manager.py index 3f4e9512..4cb2524a 100644 --- a/src/apify/events/_apify_event_manager.py +++ b/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.""" diff --git a/tests/unit/events/test_apify_event_manager.py b/tests/unit/events/test_apify_event_manager.py index 86589576..afb3b80a 100644 --- a/tests/unit/events/test_apify_event_manager.py +++ b/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):