From b6a86e9f51030c8373c1a7d14c157ae155f80799 Mon Sep 17 00:00:00 2001 From: Vlada Dusek Date: Tue, 21 Jul 2026 11:31:17 +0200 Subject: [PATCH 1/3] fix: prevent duplicate platform websocket when nesting ApifyEventManager --- src/apify/events/_apify_event_manager.py | 29 +++++++--- tests/unit/events/test_apify_event_manager.py | 57 +++++++++++++++++++ 2 files changed, 77 insertions(+), 9 deletions(-) diff --git a/src/apify/events/_apify_event_manager.py b/src/apify/events/_apify_event_manager.py index e7b21da8..f9fd39fa 100644 --- a/src/apify/events/_apify_event_manager.py +++ b/src/apify/events/_apify_event_manager.py @@ -95,6 +95,13 @@ def __init__(self, configuration: Configuration, **kwargs: Unpack[EventManagerOp @override async def __aenter__(self) -> Self: await super().__aenter__() + + # The parent ref-counts nested contexts (e.g. a crawler entering the same event manager). Only the + # outermost enter (0 -> 1) may start the platform websocket machinery; a nested enter must reuse the + # existing connection rather than open a second one. + if self._active_ref_count > 1: + return self + self._connected_to_platform_websocket = asyncio.Future() # Run tasks but don't await them @@ -119,15 +126,19 @@ 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() + # Mirror the parent's ref counting: only the outermost exit (1 -> 0) tears down the platform websocket + # machinery. A nested exit must leave the connection and its processing task intact for the still-active + # outer context. + if self._active_ref_count == 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() await super().__aexit__(exc_type, exc_value, exc_traceback) diff --git a/tests/unit/events/test_apify_event_manager.py b/tests/unit/events/test_apify_event_manager.py index 8f723966..7df2494b 100644 --- a/tests/unit/events/test_apify_event_manager.py +++ b/tests/unit/events/test_apify_event_manager.py @@ -277,6 +277,63 @@ 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 + + +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): From 8bd02d1c61ce8e9ea639d1d1182d99ac16138a5b Mon Sep 17 00:00:00 2001 From: Vlada Dusek Date: Tue, 21 Jul 2026 12:29:25 +0200 Subject: [PATCH 2/3] refactor: decouple ApifyEventManager nesting from parent ref count --- src/apify/events/_apify_event_manager.py | 22 ++++++++++++++----- tests/unit/events/test_apify_event_manager.py | 6 +++++ 2 files changed, 22 insertions(+), 6 deletions(-) diff --git a/src/apify/events/_apify_event_manager.py b/src/apify/events/_apify_event_manager.py index f9fd39fa..3be88ac1 100644 --- a/src/apify/events/_apify_event_manager.py +++ b/src/apify/events/_apify_event_manager.py @@ -92,14 +92,23 @@ 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 `async with` contexts, tracked independently of the parent's ref counting. + + The platform websocket machinery is started on the outermost enter (0 -> 1) and torn down on the + outermost exit (1 -> 0); nested enters/exits reuse the single existing connection. + """ + @override async def __aenter__(self) -> Self: await super().__aenter__() - # The parent ref-counts nested contexts (e.g. a crawler entering the same event manager). Only the - # outermost enter (0 -> 1) may start the platform websocket machinery; a nested enter must reuse the - # existing connection rather than open a second one. - if self._active_ref_count > 1: + # Ref-count nested contexts (e.g. a crawler entering the same event manager) with our own counter, + # incremented right after a successful parent enter, rather than reading the parent's private ref-count + # internals (whose mutation timing we would otherwise depend on). Only the outermost enter (0 -> 1) starts + # the platform websocket machinery; a nested enter reuses the connection instead of opening a second one. + self._context_depth += 1 + if self._context_depth > 1: return self self._connected_to_platform_websocket = asyncio.Future() @@ -126,10 +135,10 @@ async def __aexit__( exc_value: BaseException | None, exc_traceback: TracebackType | None, ) -> None: - # Mirror the parent's ref counting: only the outermost exit (1 -> 0) tears down the platform websocket + # Mirror the enter's own ref counting: only the outermost exit (1 -> 0) tears down the platform websocket # machinery. A nested exit must leave the connection and its processing task intact for the still-active # outer context. - if self._active_ref_count == 1: + 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(): @@ -140,6 +149,7 @@ async def __aexit__( 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 7df2494b..2c509f6a 100644 --- a/tests/unit/events/test_apify_event_manager.py +++ b/tests/unit/events/test_apify_event_manager.py @@ -303,6 +303,12 @@ async def test_reentrant_context_reuses_single_platform_connection(monkeypatch: 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.""" From fbfec76385affaa124921bc527baa697b6b8d844 Mon Sep 17 00:00:00 2001 From: Vlada Dusek Date: Tue, 21 Jul 2026 14:47:17 +0200 Subject: [PATCH 3/3] docs: shorten ApifyEventManager nesting comments --- src/apify/events/_apify_event_manager.py | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/src/apify/events/_apify_event_manager.py b/src/apify/events/_apify_event_manager.py index 3be88ac1..f4fd453e 100644 --- a/src/apify/events/_apify_event_manager.py +++ b/src/apify/events/_apify_event_manager.py @@ -93,20 +93,14 @@ def __init__(self, configuration: Configuration, **kwargs: Unpack[EventManagerOp """Future that resolves when the connection to the platform websocket is established.""" self._context_depth = 0 - """Nesting depth of active `async with` contexts, tracked independently of the parent's ref counting. - - The platform websocket machinery is started on the outermost enter (0 -> 1) and torn down on the - outermost exit (1 -> 0); nested enters/exits reuse the single existing connection. - """ + """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__() - # Ref-count nested contexts (e.g. a crawler entering the same event manager) with our own counter, - # incremented right after a successful parent enter, rather than reading the parent's private ref-count - # internals (whose mutation timing we would otherwise depend on). Only the outermost enter (0 -> 1) starts - # the platform websocket machinery; a nested enter reuses the connection instead of opening a second one. + # 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 @@ -135,9 +129,8 @@ async def __aexit__( exc_value: BaseException | None, exc_traceback: TracebackType | None, ) -> None: - # Mirror the enter's own ref counting: only the outermost exit (1 -> 0) tears down the platform websocket - # machinery. A nested exit must leave the connection and its processing task intact for the still-active - # outer context. + # 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.