feat: abort named child runs when the parent is gracefully aborted

pyproject.toml

@@ -37,7 +37,7 @@ keywords = [
 ]
 dependencies = [
     "apify-client[brotli]>=3.3.0,<4.0.0",
-    "crawlee>=1.8.0,<2.0.0,!=1.10.1",
+    "crawlee>=1.9.1,<2.0.0,!=1.10.1",
     "cachetools>=5.5.0",
     "cryptography>=42.0.0",
     "impit>=0.8.0",

src/apify/_actor.py

@@ -5,6 +5,7 @@
 import math
 import sys
 import warnings
+from contextlib import AsyncExitStack, ExitStack
 from dataclasses import asdict
 from datetime import UTC, datetime, timedelta
 from functools import cached_property, partial
@@ -172,6 +173,9 @@ def __init__(
 
         self._child_run_registry = ChildRunRegistry(self.open_key_value_store)
 
+        self._internal_listeners = ExitStack()
+        """Removes the listeners the Actor registered on the event manager, when the Actor leaves its context."""
+
         self._active = False
         """Whether the Actor instance is currently active (initialized and within context)."""
 
@@ -227,36 +231,41 @@ async def __aenter__(self) -> Self:
         _ = self._storage_client
         self.log.debug('Storage client initialized')
 
-        # Initialize the event manager and register it in the service locator.
-        await self.event_manager.__aenter__()
-        self.log.debug('Event manager initialized')
+        # Undoes the initialization if it fails, since a failed `__aenter__` gets no `__aexit__`.
+        async with AsyncExitStack() as rollback:
+            # Initialize the event manager and register it in the service locator.
+            await self.event_manager.__aenter__()
+            rollback.push_async_callback(self.event_manager.__aexit__, None, None, None)
+            # Only the platform emits `ABORTING`, and it does so through `ApifyEventManager`.
+            if isinstance(self.event_manager, ApifyEventManager):
+                remove_listener = self.event_manager._on_internal(  # noqa: SLF001
+                    event=Event.ABORTING, listener=self._abort_child_runs
+                )
+                self._internal_listeners.callback(remove_listener)
+            rollback.callback(self._internal_listeners.close)
+            self.log.debug('Event manager initialized')
 
-        # Initialize the charging manager.
-        try:
+            # Initialize the charging manager.
             await self._charging_manager_implementation.__aenter__()
-        except BaseException:
-            # Exit the already-entered event manager so its recurring tasks do not leak.
-            await self.event_manager.__aexit__(None, None, None)
-            raise
-        self.log.debug('Charging manager initialized')
+            rollback.push_async_callback(self._charging_manager_implementation.__aexit__, None, None, None)
+            self.log.debug('Charging manager initialized')
 
-        # Mark initialization as complete and update global state.
-        self._active = True
+            # Mark initialization as complete and update global state.
+            self._active = True
+
+            def deactivate() -> None:
+                self._active = False
+
+            rollback.callback(deactivate)
 
-        try:
             if not Actor.is_at_home():
                 # Make sure that the input related KVS is initialized to ensure that the input aware client is used
                 await self.open_key_value_store()
 
             await self._child_run_registry.load()
-        except BaseException:
-            # Undo the initialization, since a failed `__aenter__` gets no `__aexit__`.
-            self._active = False
-            try:
-                await self._charging_manager_implementation.__aexit__(None, None, None)
-            finally:
-                await self.event_manager.__aexit__(None, None, None)
-            raise
+
+            # The Actor context now owns everything entered above, and `__aexit__` leaves it.
+            rollback.pop_all()
 
         if self.is_at_home():
             self.event_manager.on(event=Event.CRAWLER_STATUS, listener=self._forward_crawler_status)
@@ -347,6 +356,7 @@ async def finalize() -> None:
         except TimeoutError:
             self.log.exception('Actor cleanup timed out')
         finally:
+            self._internal_listeners.close()
             self._active = False
 
         if reraise_control_flow:
@@ -1012,6 +1022,7 @@ async def start(
         force_permission_level: ActorPermissionLevel | None = None,
         webhooks: list[Webhook] | None = None,
         run_name: str | None = None,
+        abort_with_parent: bool = False,
     ) -> Run:
         """Run an Actor on the Apify platform.
 
@@ -1047,10 +1058,18 @@ async def start(
                 resurrected, and a new run is started only when nothing is recorded under the name, or the recorded
                 run `FAILED` or no longer exists. Using `run_name` to attach to a run of a different Actor, task or
                 with a different input raises a `ValueError`.
+            abort_with_parent: If true, the child run is gracefully aborted when this Actor run is gracefully
+                aborted. It requires `run_name`, and the value is recorded under it, replacing the one from an earlier
+                call. Once this Actor run is being gracefully aborted, a call with it raises a `RuntimeError` and
+                starts, resurrects or reattaches no run. A hard abort, a timeout or a crash of this Actor run leaves
+                the child running.
 
         Returns:
             Info about the started Actor run
         """
+        if abort_with_parent and run_name is None:
+            raise ValueError('`abort_with_parent` requires `run_name`, since only named child runs are tracked.')
+
         if max_items is not None:
             _warn_max_items_deprecated()
 
@@ -1086,6 +1105,7 @@ async def start(
             restart_on_error=restart_on_error,
             memory_mbytes=memory_mbytes,
             timeout=timeout,
+            abort_with_parent=abort_with_parent,
         )
         return run
 
@@ -1200,6 +1220,7 @@ async def call(
         wait: timedelta | None = None,
         logger: logging.Logger | Literal['default'] | None = 'default',
         run_name: str | None = None,
+        abort_with_parent: bool = False,
     ) -> Run:
         """Start an Actor on the Apify Platform and wait for it to finish before returning.
 
@@ -1238,10 +1259,18 @@ async def call(
                 resurrected, and a new run is started only when nothing is recorded under the name, or the recorded
                 run `FAILED` or no longer exists. Using `run_name` to attach to a run of a different Actor, task or
                 with a different input raises a `ValueError`.
+            abort_with_parent: If true, the child run is gracefully aborted when this Actor run is gracefully
+                aborted. It requires `run_name`, and the value is recorded under it, replacing the one from an earlier
+                call. Once this Actor run is being gracefully aborted, a call with it raises a `RuntimeError` and
+                starts, resurrects or reattaches no run. A hard abort, a timeout or a crash of this Actor run leaves
+                the child running.
 
         Returns:
             Info about the started Actor run.
         """
+        if abort_with_parent and run_name is None:
+            raise ValueError('`abort_with_parent` requires `run_name`, since only named child runs are tracked.')
+
         if max_items is not None:
             _warn_max_items_deprecated()
 
@@ -1289,6 +1318,7 @@ async def call(
                 restart_on_error=restart_on_error,
                 memory_mbytes=memory_mbytes,
                 timeout=timeout,
+                abort_with_parent=abort_with_parent,
             )
             # The earlier attempt of this call already streamed the log of a reattached or resurrected run.
             run = await self._wait_for_child_run(
@@ -1315,6 +1345,7 @@ async def _find_or_start_child_run(
         restart_on_error: bool | None,
         memory_mbytes: int | None,
         timeout: timedelta | Literal['inherit'] | None,
+        abort_with_parent: bool,
     ) -> tuple[Run, bool]:
         return await self._child_run_registry.find_or_start(
             name,
@@ -1331,8 +1362,12 @@ async def _find_or_start_child_run(
                 max_total_charge_usd=max_total_charge_usd,
                 restart_on_error=restart_on_error,
             ),
+            abort_with_parent=abort_with_parent,
         )
 
+    async def _abort_child_runs(self) -> None:
+        await self._child_run_registry.abort_runs_with_parent(self.apify_client)
+
     async def _wait_for_child_run(
         self,
         name: str,
@@ -1375,6 +1410,7 @@ async def start_task(
         webhooks: list[Webhook] | None = None,
         token: str | None = None,
         run_name: str | None = None,
+        abort_with_parent: bool = False,
     ) -> Run:
         """Start an Actor task on the Apify Platform.
 
@@ -1410,10 +1446,18 @@ async def start_task(
                 resurrected, and a new run is started only when nothing is recorded under the name, or the recorded
                 run `FAILED` or no longer exists. Using `run_name` to attach to a run of a different Actor, task or
                 with a different input raises a `ValueError`.
+            abort_with_parent: If true, the child run is gracefully aborted when this Actor run is gracefully
+                aborted. It requires `run_name`, and the value is recorded under it, replacing the one from an earlier
+                call. Once this Actor run is being gracefully aborted, a call with it raises a `RuntimeError` and
+                starts, resurrects or reattaches no run. A hard abort, a timeout or a crash of this Actor run leaves
+                the child running.
 
         Returns:
             Info about the started Actor run.
         """
+        if abort_with_parent and run_name is None:
+            raise ValueError('`abort_with_parent` requires `run_name`, since only named child runs are tracked.')
+
         if max_items is not None:
             _warn_max_items_deprecated()
 
@@ -1446,6 +1490,7 @@ async def start_task(
             restart_on_error=restart_on_error,
             memory_mbytes=memory_mbytes,
             timeout=timeout,
+            abort_with_parent=abort_with_parent,
         )
         return run
 
@@ -1465,6 +1510,7 @@ async def call_task(
         wait: timedelta | None = None,
         token: str | None = None,
         run_name: str | None = None,
+        abort_with_parent: bool = False,
     ) -> Run:
         """Start an Actor task on the Apify Platform and wait for it to finish before returning.
 
@@ -1500,10 +1546,18 @@ async def call_task(
                 resurrected, and a new run is started only when nothing is recorded under the name, or the recorded
                 run `FAILED` or no longer exists. Using `run_name` to attach to a run of a different Actor, task or
                 with a different input raises a `ValueError`.
+            abort_with_parent: If true, the child run is gracefully aborted when this Actor run is gracefully
+                aborted. It requires `run_name`, and the value is recorded under it, replacing the one from an earlier
+                call. Once this Actor run is being gracefully aborted, a call with it raises a `RuntimeError` and
+                starts, resurrects or reattaches no run. A hard abort, a timeout or a crash of this Actor run leaves
+                the child running.
 
         Returns:
             Info about the started Actor run.
         """
+        if abort_with_parent and run_name is None:
+            raise ValueError('`abort_with_parent` requires `run_name`, since only named child runs are tracked.')
+
         if max_items is not None:
             _warn_max_items_deprecated()
 
@@ -1546,6 +1600,7 @@ async def call_task(
                 restart_on_error=restart_on_error,
                 memory_mbytes=memory_mbytes,
                 timeout=timeout,
+                abort_with_parent=abort_with_parent,
             )
             run = await self._wait_for_child_run(
                 run_name, client.run(started_run.id), started_run, wait=wait, logger=None, from_start=False