fix: keep aborting child runs when all ABORTING listeners are removed

Author: vdusek

src/apify/_actor.py

@@ -206,15 +206,17 @@ async def __aenter__(self) -> Self:
 
         # Initialize the event manager and register it in the service locator.
         await self.event_manager.__aenter__()
-        self.event_manager.on(event=Event.ABORTING, listener=self._abort_child_runs)
+        # Only the platform emits `ABORTING`, and it does so through `ApifyEventManager`.
+        if isinstance(self.event_manager, ApifyEventManager):
+            self.event_manager._on_internal(event=Event.ABORTING, listener=self._abort_child_runs)  # noqa: SLF001
         self.log.debug('Event manager initialized')
 
         # Initialize the charging manager.
         try:
             await self._charging_manager_implementation.__aenter__()
         except BaseException:
             # Exit the already-entered event manager so its recurring tasks do not leak.
-            self.event_manager.off(event=Event.ABORTING, listener=self._abort_child_runs)
+            self._remove_internal_listeners()
             await self.event_manager.__aexit__(None, None, None)
             raise
         self.log.debug('Charging manager initialized')
@@ -306,7 +308,7 @@ async def finalize() -> None:
         except TimeoutError:
             self.log.exception('Actor cleanup timed out')
         finally:
-            self.event_manager.off(event=Event.ABORTING, listener=self._abort_child_runs)
+            self._remove_internal_listeners()
             self._active = False
 
         if reraise_control_flow:
@@ -1226,6 +1228,10 @@ async def _find_or_start_child_run(
             abort_with_parent=abort_with_parent,
         )
 
+    def _remove_internal_listeners(self) -> None:
+        if isinstance(self.event_manager, ApifyEventManager):
+            self.event_manager._off_internal(event=Event.ABORTING, listener=self._abort_child_runs)  # noqa: SLF001
+
     async def _abort_child_runs(self) -> None:
         await self._child_run_registry.abort_runs_with_parent(self.apify_client)
 

src/apify/events/_apify_event_manager.py

@@ -3,8 +3,9 @@
 import asyncio
 import contextlib
 import time
+from collections import defaultdict
 from logging import getLogger
-from typing import TYPE_CHECKING, Annotated, Self, cast
+from typing import TYPE_CHECKING, Annotated, Any, Self, cast
 
 import websockets.asyncio.client
 import websockets.client
@@ -24,6 +25,7 @@
     from types import TracebackType
 
     from crawlee.events._event_manager import EventManagerOptions
+    from crawlee.events._types import EventData, EventListener, WrappedListener
 
     from apify._configuration import Configuration
 
@@ -94,6 +96,11 @@ def __init__(self, configuration: Configuration, **kwargs: Unpack[EventManagerOp
         connection, so that `__aenter__` can report it.
         """
 
+        self._internal_listeners: defaultdict[Event, dict[EventListener[Any], WrappedListener]] = defaultdict(dict)
+        """Listeners of the SDK itself, mapped as `event -> listener -> wrapper`. `off` doesn't remove them, so user
+        code removing all listeners of an event keeps the SDK's own handling of it.
+        """
+
     @override
     async def __aenter__(self) -> Self:
         """Initialize the event manager upon entering the async context.
@@ -149,6 +156,24 @@ async def __aexit__(
             # emitting `PersistState` again, as re-entering the context would be a no-op.
             await super().__aexit__(exc_type, exc_value, exc_traceback)
 
+    @override
+    def emit(self, *, event: Event, event_data: EventData) -> None:
+        super().emit(event=event, event_data=event_data)
+
+        for listener, listener_wrapper in self._internal_listeners.get(event, {}).items():
+            task_name = f'Task-{event.value}-{self._get_listener_name(listener)}'
+            listener_task = asyncio.create_task(listener_wrapper(event_data), name=task_name)
+            self._listener_tasks.add(listener_task)
+            listener_task.add_done_callback(self._listener_tasks.discard)
+
+    def _on_internal(self, *, event: Event, listener: EventListener[Any]) -> None:
+        """Register a listener of the SDK itself, which `off` doesn't remove."""
+        self._internal_listeners[event][listener] = self._wrap_listener(event, listener)
+
+    def _off_internal(self, *, event: Event, listener: EventListener[Any]) -> None:
+        """Remove a listener registered by `_on_internal`."""
+        self._internal_listeners.get(event, {}).pop(listener, None)
+
     async def _teardown_platform_websocket(self) -> None:
         """Stop consuming the platform messages and close the websocket connection to the platform events."""
         try:

tests/unit/actor/test_actor_child_runs.py

@@ -8,11 +8,13 @@
 import pytest
 
 from apify_client._models import Run
+from crawlee import service_locator
 from crawlee.events import Event, EventAbortingData
 
-from apify import Actor
+from apify import Actor, Configuration
 from apify._actor import _ActorType
 from apify._child_runs import CHILD_RUNS_KEY, ChildRunRegistry
+from apify.events import ApifyEventManager
 
 if TYPE_CHECKING:
     from ..conftest import ApifyClientAsyncPatcher
@@ -39,6 +41,14 @@ def make_run(run_id: str, status: str) -> Run:
     )
 
 
+@pytest.fixture
+def apify_event_manager() -> ApifyEventManager:
+    """Make the Actor use `ApifyEventManager`, which delivers `ABORTING` on the platform, without a websocket."""
+    event_manager = ApifyEventManager(Configuration.get_global_configuration())
+    service_locator.set_event_manager(event_manager)
+    return event_manager
+
+
 async def record_child_run(name: str, run_id: str, *, actor_id: str = 'some-actor') -> None:
     """Seed the registry the way an earlier attempt of this Actor run would have left it."""
     kvs = await Actor.open_key_value_store()
@@ -452,7 +462,7 @@ async def test_reattach_replaces_recorded_abort_with_parent(
 
 
 async def test_aborting_event_aborts_marked_active_child_runs(
-    apify_client_async_patcher: ApifyClientAsyncPatcher,
+    apify_client_async_patcher: ApifyClientAsyncPatcher, apify_event_manager: ApifyEventManager
 ) -> None:
     """On `ABORTING`, only child runs marked `abort_with_parent` that are still active are gracefully aborted."""
     runs = {
@@ -482,16 +492,18 @@ async def get_run(run_client: Any, *_args: Any, **_kwargs: Any) -> Run | None:
                 ]
             },
         )
-        Actor.event_manager.emit(event=Event.ABORTING, event_data=EventAbortingData())
-        await Actor.event_manager.wait_for_all_listeners_to_complete()
+        apify_event_manager.emit(event=Event.ABORTING, event_data=EventAbortingData())
+        await apify_event_manager.wait_for_all_listeners_to_complete()
 
     aborts = apify_client_async_patcher.calls['run']['abort']
     assert sorted(args[0].resource_id for args, _ in aborts) == ['ready-run', 'running-run']
     assert all(kwargs == {'gracefully': True} for _, kwargs in aborts)
 
 
 async def test_failed_child_run_abort_does_not_stop_others(
-    apify_client_async_patcher: ApifyClientAsyncPatcher, caplog: pytest.LogCaptureFixture
+    apify_client_async_patcher: ApifyClientAsyncPatcher,
+    caplog: pytest.LogCaptureFixture,
+    apify_event_manager: ApifyEventManager,
 ) -> None:
     """A child run that fails to abort is logged, and the other marked child runs are still aborted."""
 
@@ -513,8 +525,8 @@ async def abort_run(run_client: Any, *_args: Any, **_kwargs: Any) -> None:
                 for name in ['broken', 'healthy']
             },
         )
-        Actor.event_manager.emit(event=Event.ABORTING, event_data=EventAbortingData())
-        await Actor.event_manager.wait_for_all_listeners_to_complete()
+        apify_event_manager.emit(event=Event.ABORTING, event_data=EventAbortingData())
+        await apify_event_manager.wait_for_all_listeners_to_complete()
 
     aborts = apify_client_async_patcher.calls['run']['abort']
     assert sorted(args[0].resource_id for args, _ in aborts) == ['broken-run', 'healthy-run']
@@ -608,16 +620,35 @@ async def start_run() -> Run:
     client.run.return_value.abort.assert_awaited_once_with(gracefully=True)
 
 
-async def test_exit_removes_the_aborting_listener(apify_client_async_patcher: ApifyClientAsyncPatcher) -> None:
+async def test_exit_removes_the_aborting_listener(
+    apify_client_async_patcher: ApifyClientAsyncPatcher, apify_event_manager: ApifyEventManager
+) -> None:
     """After the Actor exits, an `ABORTING` event on a still-active event manager aborts no child run."""
     apify_client_async_patcher.patch('actor', 'start', return_value=make_run('new-run', 'READY'))
     apify_client_async_patcher.patch('run', 'get', return_value=make_run('new-run', 'RUNNING'))
     apify_client_async_patcher.patch('run', 'abort', return_value=None)
 
-    async with Actor.event_manager:
+    async with apify_event_manager:
         async with Actor:
             await Actor.start('some-actor', name='scrape-eu', abort_with_parent=True)
-        Actor.event_manager.emit(event=Event.ABORTING, event_data=EventAbortingData())
-        await Actor.event_manager.wait_for_all_listeners_to_complete()
+        apify_event_manager.emit(event=Event.ABORTING, event_data=EventAbortingData())
+        await apify_event_manager.wait_for_all_listeners_to_complete()
 
     assert apify_client_async_patcher.calls['run']['abort'] == []
+
+
+async def test_removing_all_aborting_listeners_keeps_aborting_child_runs(
+    apify_client_async_patcher: ApifyClientAsyncPatcher, apify_event_manager: ApifyEventManager
+) -> None:
+    """Removing all `ABORTING` listeners from the event manager still aborts child runs marked `abort_with_parent`."""
+    apify_client_async_patcher.patch('actor', 'start', return_value=make_run('new-run', 'READY'))
+    apify_client_async_patcher.patch('run', 'get', return_value=make_run('new-run', 'RUNNING'))
+    apify_client_async_patcher.patch('run', 'abort', return_value=None)
+
+    async with Actor:
+        await Actor.start('some-actor', name='scrape-eu', abort_with_parent=True)
+        apify_event_manager.off(event=Event.ABORTING)
+        apify_event_manager.emit(event=Event.ABORTING, event_data=EventAbortingData())
+        await apify_event_manager.wait_for_all_listeners_to_complete()
+
+    assert len(apify_client_async_patcher.calls['run']['abort']) == 1

tests/unit/events/test_apify_event_manager.py

@@ -9,14 +9,14 @@
 from collections import defaultdict
 from datetime import timedelta
 from typing import TYPE_CHECKING, Any
-from unittest.mock import Mock
+from unittest.mock import AsyncMock, Mock
 
 import pytest
 import websockets
 import websockets.asyncio.client
 import websockets.asyncio.server
 
-from crawlee.events._types import Event
+from crawlee.events._types import Event, EventAbortingData
 
 from ..._utils import poll_until_condition
 from apify import Configuration
@@ -489,6 +489,32 @@ async def handler(_data: Any) -> None:
         assert persist_state_counter == 0
 
 
+async def test_internal_listener_is_kept_by_off() -> None:
+    """A listener registered by `_on_internal` still runs after `off` removes all listeners of its event."""
+    listener = AsyncMock()
+
+    async with ApifyEventManager(Configuration.get_global_configuration()) as event_manager:
+        event_manager._on_internal(event=Event.ABORTING, listener=listener)
+        event_manager.off(event=Event.ABORTING)
+        event_manager.emit(event=Event.ABORTING, event_data=EventAbortingData())
+        await event_manager.wait_for_all_listeners_to_complete()
+
+    listener.assert_awaited_once()
+
+
+async def test_off_internal_removes_the_listener() -> None:
+    """A listener removed by `_off_internal` no longer runs on its event."""
+    listener = AsyncMock()
+
+    async with ApifyEventManager(Configuration.get_global_configuration()) as event_manager:
+        event_manager._on_internal(event=Event.ABORTING, listener=listener)
+        event_manager._off_internal(event=Event.ABORTING, listener=listener)
+        event_manager.emit(event=Event.ABORTING, event_data=EventAbortingData())
+        await event_manager.wait_for_all_listeners_to_complete()
+
+    listener.assert_not_awaited()
+
+
 async def test_deprecated_event_is_skipped(monkeypatch: pytest.MonkeyPatch) -> None:
     """Test that deprecated events (like CPU_INFO) are silently skipped."""
     async with (