diff --git a/src/apify/events/_apify_event_manager.py b/src/apify/events/_apify_event_manager.py index e7b21da8..f4fd453e 100644 --- a/src/apify/events/_apify_event_manager.py +++ b/src/apify/events/_apify_event_manager.py @@ -92,9 +92,19 @@ def __init__(self, configuration: Configuration, **kwargs: Unpack[EventManagerOp self._connected_to_platform_websocket: asyncio.Future[bool] | None = None """Future that resolves when the connection to the platform websocket is established.""" + self._context_depth = 0 + """Nesting depth of active contexts; the outermost enter/exit (0 <-> 1) starts/tears down the websocket.""" + @override async def __aenter__(self) -> Self: await super().__aenter__() + + # Track nesting depth ourselves rather than reading the parent's private ref count. Only the outermost + # enter (0 -> 1) starts the websocket machinery; a nested enter reuses the existing connection. + self._context_depth += 1 + if self._context_depth > 1: + return self + self._connected_to_platform_websocket = asyncio.Future() # Run tasks but don't await them @@ -119,16 +129,20 @@ 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 - - if self._platform_events_websocket: - await self._platform_events_websocket.close() - + # Only the outermost exit (1 -> 0) tears down the websocket machinery; a nested exit must leave the + # connection and its processing task intact for the still-active outer context. + if self._context_depth == 1: + # 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() + + self._context_depth -= 1 await super().__aexit__(exc_type, exc_value, exc_traceback) def _process_connection_exception(self, exc: Exception) -> Exception | None: diff --git a/tests/unit/events/test_apify_event_manager.py b/tests/unit/events/test_apify_event_manager.py index 8f723966..2c509f6a 100644 --- a/tests/unit/events/test_apify_event_manager.py +++ b/tests/unit/events/test_apify_event_manager.py @@ -277,6 +277,69 @@ async def test_lifecycle_on_platform(monkeypatch: pytest.MonkeyPatch) -> None: assert len(connected_ws_clients) == 1 +async def test_reentrant_context_reuses_single_platform_connection(monkeypatch: pytest.MonkeyPatch) -> None: + """A nested `async with` on the same event manager reuses the one platform connection, not a second.""" + 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 + outer_task = event_manager._process_platform_messages_task + + event_calls: list[Any] = [] + event_manager.on(event=Event.SYSTEM_INFO, listener=event_calls.append) + + async with event_manager: + # The nested enter must not replace the message-processing task nor open a second connection. + assert event_manager._process_platform_messages_task is outer_task + await asyncio.sleep(0.2) + assert len(connected_ws_clients) == 1 + + # A single platform event must be delivered exactly once, not once per connection. + websockets.broadcast( + connected_ws_clients, json.dumps({'name': 'systemInfo', 'data': DUMMY_SYSTEM_INFO}) + ) + await poll_until_condition(lambda: len(event_calls) >= 1, poll_interval=0.05) + await asyncio.sleep(0.2) + assert len(event_calls) == 1 + + # The outermost exit tears down the single shared connection and its processing task; nothing is leaked. + assert outer_task is not None + assert outer_task.done() + await poll_until_condition(lambda: len(connected_ws_clients) == 0, poll_interval=0.05) + assert len(connected_ws_clients) == 0 + + +async def test_reentrant_exit_leaves_outer_context_functional(monkeypatch: pytest.MonkeyPatch) -> None: + """Exiting a nested context leaves the outer connection alive and working; only the final 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() + outer_task = event_manager._process_platform_messages_task + assert outer_task is not None + + async with event_manager: + pass + + # The nested exit must not cancel the outer context's processing task. + assert event_manager.active is True + assert not outer_task.done() + + # Events must still be delivered on the surviving connection. + 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: len(event_calls) == 1, poll_interval=0.05) + assert len(event_calls) == 1 + + # The final exit tears everything down; no connection or task is leaked. + assert event_manager.active is False + assert outer_task.done() + await poll_until_condition(lambda: len(connected_ws_clients) == 0, poll_interval=0.05) + assert len(connected_ws_clients) == 0 + + async def test_event_handling_on_platform(monkeypatch: pytest.MonkeyPatch) -> None: async with _platform_ws_server(monkeypatch) as (connected_ws_clients, client_connected):