diff --git a/src/apify/_actor.py b/src/apify/_actor.py index e3c0d323..b919dcad 100644 --- a/src/apify/_actor.py +++ b/src/apify/_actor.py @@ -445,6 +445,24 @@ def child_runs(self) -> dict[str, RunClientAsync]: """ return self._child_run_registry.run_clients(self.apify_client) + def set_named_child_run_limits(self, *, max_concurrent_runs: int | None) -> None: + """Limit the named child runs of this Actor run. + + While `max_concurrent_runs` named child runs are `READY`, `RUNNING`, `ABORTING` or `TIMING-OUT`, a named + `Actor.start`, `Actor.call`, `Actor.start_task` or `Actor.call_task` that would start or resurrect a run waits + until one of them finishes. Reattaching to a recorded run never waits. Runs started without a `run_name` are not + counted and never wait. A child run's status is fetched again before it is counted, if it is more than 10 + seconds old. Once this Actor run is being gracefully aborted, a call that would wait raises a `RuntimeError`, + with or without `abort_with_parent`. + + A call already waiting picks up a changed limit within 10 seconds. The limit is kept in memory, so call this + method again after a migration or resurrection of this Actor run. + + Args: + max_concurrent_runs: How many named child runs may be active at once, or `None` for no limit. + """ + self._child_run_registry.set_max_concurrent_runs(max_concurrent_runs) + @cached_property def configuration(self) -> Configuration: """Actor configuration, uses the default instance if not explicitly set.""" diff --git a/src/apify/_child_runs.py b/src/apify/_child_runs.py index d91e74da..89d58f8b 100644 --- a/src/apify/_child_runs.py +++ b/src/apify/_child_runs.py @@ -3,9 +3,10 @@ import asyncio import hashlib import json +from contextlib import asynccontextmanager, suppress from datetime import datetime from logging import getLogger -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, NamedTuple from weakref import WeakValueDictionary from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError @@ -14,7 +15,7 @@ from apify_client.errors import NotFoundError if TYPE_CHECKING: - from collections.abc import Awaitable, Callable + from collections.abc import AsyncGenerator, Awaitable, Callable from apify_client import ApifyClientAsync from apify_client._models import Run @@ -31,13 +32,22 @@ """Statuses that end as `ABORTED` / `TIMED-OUT` shortly, and are resurrectable once they do.""" _RESURRECTABLE_STATUSES = frozenset({'ABORTED', 'TIMED-OUT'}) +"""Statuses of a recorded run that a call under its name resurrects.""" _ABORTED_STATUSES = frozenset({'ABORTING', 'ABORTED'}) +"""Statuses that show an abort request took effect on the child run.""" + +_ACTIVE_STATUSES = frozenset({'READY', 'RUNNING', 'ABORTING', 'TIMING-OUT'}) +"""Statuses of a child run that count toward the concurrency limit.""" + +_STATUS_MAX_AGE_SECS = 10 +"""How long an observed active status counts toward the concurrency limit before the run is fetched again.""" _NOT_FOUND_GRACE_SECS = 3 """How long a recorded run that the API reports as missing is looked up again before it counts as gone.""" _NOT_FOUND_RETRY_INTERVAL_SECS = 0.25 +"""How long to wait between lookups of a recorded run that the API reports as missing.""" class ChildRunSnapshot(BaseModel): @@ -68,6 +78,16 @@ class ChildRunRecord(ChildRunSnapshot): """Whether the current run is aborted when this Actor run is gracefully aborted.""" +class _CachedRunStatus(NamedTuple): + """Status of a child run as this Actor run last saw it, used to count the run toward the concurrency limit.""" + + status: str + """Status of the run, or `LOST` once the platform no longer returned it.""" + + cached_at: float + """Event loop time the status was seen at.""" + + def checksum_request(*, actor_id: str | None, task_id: str | None, run_input: Any) -> str: """Hash the Actor or task and the input of a named start, in the same JSON shape as the JS SDK.""" request: dict[str, Any] = ( @@ -106,15 +126,32 @@ class ChildRunRegistry: def __init__(self, open_key_value_store: Callable[[], Awaitable[KeyValueStore]]) -> None: self._open_key_value_store = open_key_value_store + """Opens the default key-value store, where the records are persisted.""" self._records: dict[str, ChildRunRecord] | None = None + """Records of the child runs by name, `None` until they are read from the key-value store.""" self._lock = asyncio.Lock() """Guards loading the records and writing them back to the key-value store.""" self._name_locks: WeakValueDictionary[str, asyncio.Lock] = WeakValueDictionary() """Serializes `find_or_start` per name. A lock is dropped once no call under its name holds it.""" self._clients: dict[str, ApifyClientAsync] = {} """Client the run under each name was last started or looked up with. Lost on a migration.""" + self._max_concurrent_runs: int | None = None + """How many recorded runs may be active at once, or `None` for no limit.""" + self._capacity_available = asyncio.Condition() + """Serializes counting the active runs, and wakes the waiting starts when a run may have finished.""" + self._starting_names: set[str] = set() + """Names with a start or resurrection in flight, counted toward the limit but not from the records yet.""" + self._run_status_cache: dict[str, _CachedRunStatus] = {} + """Last seen status of the run recorded under each name. An active one is fetched again once it gets stale.""" self._parent_aborting = False - """Whether this Actor run is being gracefully aborted, after which no run marked `abort_with_parent` starts.""" + """Whether this Actor run is being gracefully aborted, after which no run marked `abort_with_parent` starts, + nor any run that would wait under the concurrency limit.""" + + def set_max_concurrent_runs(self, max_concurrent_runs: int | None) -> None: + """Set how many recorded runs may be active at once, or remove the limit with `None`.""" + if max_concurrent_runs is not None and max_concurrent_runs < 1: + raise ValueError(f'`max_concurrent_runs` must be at least 1, got {max_concurrent_runs}.') + self._max_concurrent_runs = max_concurrent_runs async def find_or_start( self, @@ -135,6 +172,8 @@ async def find_or_start( A `FAILED` run, or one the API no longer knows, is replaced by a new run under the same name, for any Actor, task or input. + Starting or resurrecting a run waits while the concurrency limit is reached. Reattaching never waits. + Args: name: Name of the child run, unique within the parent run. actor_id: The Actor to start. @@ -153,7 +192,8 @@ async def find_or_start( Raises: ValueError: If the recorded run would be reused, but was started for a different Actor, task or input. RuntimeError: If `abort_with_parent` is set and this Actor run is being gracefully aborted. No run is - started, resurrected or reattached. + started, resurrected or reattached. Also if starting or resurrecting the run would wait under the + concurrency limit while this Actor run is being gracefully aborted, whatever `abort_with_parent` is. """ checksum = checksum_request(actor_id=actor_id, task_id=task_id, run_input=run_input) @@ -168,9 +208,10 @@ async def find_or_start( record = records.get(name) if record is None: - run = await self._start( - name, checksum=checksum, start_run=start_run, history=[], abort_with_parent=abort_with_parent - ) + async with self._reserve_capacity(name, client, abort_with_parent=abort_with_parent): + run = await self._start( + name, checksum=checksum, start_run=start_run, history=[], abort_with_parent=abort_with_parent + ) self._clients[name] = client return run, True @@ -186,13 +227,14 @@ async def find_or_start( status=run.status if run is not None else 'LOST', started_at=record.started_at, ) - run = await self._start( - name, - checksum=checksum, - start_run=start_run, - history=[*record.history, replaced], - abort_with_parent=abort_with_parent, - ) + async with self._reserve_capacity(name, client, abort_with_parent=abort_with_parent): + run = await self._start( + name, + checksum=checksum, + start_run=start_run, + history=[*record.history, replaced], + abort_with_parent=abort_with_parent, + ) self._clients[name] = client return run, True @@ -202,12 +244,16 @@ async def find_or_start( 'Use a unique `run_name` for each child run.' ) + self._cache_run_status(name, run) + if record.abort_with_parent != abort_with_parent: await self._save(name, record.model_copy(update={'abort_with_parent': abort_with_parent})) if run.status in _RESURRECTABLE_STATUSES: - logger.info(f'Resurrecting child run "{name}"', extra={'run_id': run.id, 'status': run.status}) - run = await resurrect_run(run.id) + async with self._reserve_capacity(name, client, abort_with_parent=abort_with_parent): + logger.info(f'Resurrecting child run "{name}"', extra={'run_id': run.id, 'status': run.status}) + run = await resurrect_run(run.id) + self._cache_run_status(name, run) else: logger.info(f'Reattaching to child run "{name}"', extra={'run_id': run.id, 'status': run.status}) @@ -233,14 +279,20 @@ def run_clients(self, default_client: ApifyClientAsync) -> dict[str, RunClientAs async def update(self, name: str, run: Run) -> None: """Record the latest observed status of the run recorded under `name`. + A run that is no longer active stops counting toward the concurrency limit. + Args: name: Name of the child run. run: The run as just returned by the API. Nothing is recorded unless it is the current run under `name`. """ record = (await self._load()).get(name) - if record is None or record.run_id != run.id or record.status == run.status: + if record is None or record.run_id != run.id: return - await self._save(name, record.model_copy(update={'status': run.status})) + self._cache_run_status(name, run) + if record.status != run.status: + await self._save(name, record.model_copy(update={'status': run.status})) + async with self._capacity_available: + self._capacity_available.notify_all() async def abort_runs_with_parent(self, client: ApifyClientAsync) -> None: """Gracefully abort every recorded run marked `abort_with_parent`, and refuse to start any more of them. @@ -253,6 +305,9 @@ async def abort_runs_with_parent(self, client: ApifyClientAsync) -> None: client: Client used for a name not started or looked up in this process, e.g. after a migration. """ self._parent_aborting = True + async with self._capacity_available: + self._capacity_available.notify_all() + records = await self._load() # Names with a start in flight are not recorded yet, so their locks are awaited too. await asyncio.gather(*(self._abort(name, client) for name in {*records, *self._name_locks})) @@ -292,8 +347,88 @@ async def _start( abort_with_parent=abort_with_parent, ) await self._save(name, record) + self._cache_run_status(name, run) return run + @asynccontextmanager + async def _reserve_capacity( + self, name: str, client: ApifyClientAsync, *, abort_with_parent: bool + ) -> AsyncGenerator[None]: + """Count starting or resurrecting the run under `name` toward the limit, waiting while the limit is reached.""" + async with self._capacity_available: + while ( + self._max_concurrent_runs is not None + and await self._count_active(client, exclude=name) >= self._max_concurrent_runs + ): + if self._parent_aborting: + raise RuntimeError( + f'Child run "{name}" was not started, since this Actor run is being aborted and the limit ' + f'of {self._max_concurrent_runs} concurrent child runs is reached.' + ) + logger.debug(f'Child run "{name}" is waiting for the concurrency limit') + with suppress(TimeoutError): + async with asyncio.timeout(_STATUS_MAX_AGE_SECS): + await self._capacity_available.wait() + # The parent can start aborting while this start waits. + if abort_with_parent and self._parent_aborting: + raise RuntimeError( + f'Child run "{name}" with `abort_with_parent` was not started, since this Actor run is being ' + 'aborted.' + ) + self._starting_names.add(name) + + try: + yield + finally: + self._starting_names.discard(name) + async with self._capacity_available: + self._capacity_available.notify_all() + + async def _count_active(self, client: ApifyClientAsync, *, exclude: str) -> int: + """Count active recorded runs and starts in flight under other names, refetching stale active statuses.""" + records = await self._load() + names = [name for name in records if name != exclude and name not in self._starting_names] + # A start in flight can finish during the fetch below, leaving `_starting_names` before it is in `names`. + starting = len(self._starting_names) + now = asyncio.get_running_loop().time() + stale = [] + for name in names: + cached_run_status = self._run_status_cache.get(name) + if cached_run_status is None: + if records[name].status in _ACTIVE_STATUSES: + stale.append(name) + elif ( + cached_run_status.status in _ACTIVE_STATUSES + and now - cached_run_status.cached_at >= _STATUS_MAX_AGE_SECS + ): + stale.append(name) + runs = await asyncio.gather( + *(self._clients.get(name, client).run(records[name].run_id).get() for name in stale), + return_exceptions=True, + ) + for name, run in zip(stale, runs, strict=True): + if isinstance(run, BaseException): + logger.warning( + f'Failed to fetch child run "{name}" to count it toward the concurrency limit', + extra={'run_id': records[name].run_id}, + exc_info=run, + ) + elif run is None: + self._run_status_cache[name] = _CachedRunStatus('LOST', now) + else: + self._cache_run_status(name, run) + + # A run that could not be fetched counts only when an earlier observation saw it active. + active = 0 + for name in names: + cached_run_status = self._run_status_cache.get(name) + if cached_run_status is not None and cached_run_status.status in _ACTIVE_STATUSES: + active += 1 + return active + starting + + def _cache_run_status(self, name: str, run: Run) -> None: + self._run_status_cache[name] = _CachedRunStatus(run.status, asyncio.get_running_loop().time()) + async def load(self) -> dict[str, ChildRunRecord]: """Read the records from the default key-value store, replacing any read before.""" async with self._lock: diff --git a/tests/e2e/test_actor_child_runs.py b/tests/e2e/test_actor_child_runs.py index 05b1352a..8b5920bc 100644 --- a/tests/e2e/test_actor_child_runs.py +++ b/tests/e2e/test_actor_child_runs.py @@ -131,3 +131,36 @@ async def main() -> None: child_run = await apify_client_async.run(child_run_id).wait_for_finish(wait_duration=timedelta(seconds=120)) assert child_run is not None assert child_run.status == 'ABORTED' + + +async def test_named_child_runs_respect_the_concurrency_limit( + make_actor: MakeActorFunction, + run_actor: RunActorFunction, +) -> None: + """Named child runs started concurrently under a limit of one run one after another.""" + + async def main() -> None: + async with Actor: + actor_input = (await Actor.get_input()) or {} + if actor_input.get('is_child') is True: + await asyncio.sleep(10) + return + + actor_id = Actor.configuration.actor_id or '' + Actor.set_named_child_run_limits(max_concurrent_runs=1) + runs = await asyncio.gather( + *( + Actor.call(actor_id=actor_id, run_input={'is_child': True}, run_name=f'child-{index}') + for index in range(2) + ) + ) + first, second = sorted(runs, key=lambda run: run.started_at) + assert first.finished_at is not None, 'first.finished_at is None' + assert second.started_at >= first.finished_at, f'first={first}, second={second}' + + actor = await make_actor(label='child-run-limit', main_func=main) + run_result = await run_actor(actor) + + assert run_result.status == 'SUCCEEDED' + # The parent run and its two child runs. + assert (await actor.runs().list()).total == 3 diff --git a/tests/unit/actor/test_actor_child_runs.py b/tests/unit/actor/test_actor_child_runs.py index e4bd5d2c..438fba0f 100644 --- a/tests/unit/actor/test_actor_child_runs.py +++ b/tests/unit/actor/test_actor_child_runs.py @@ -1171,3 +1171,452 @@ async def test_removing_all_aborting_listeners_keeps_aborting_child_runs( await apify_event_manager.wait_for_all_listeners_to_complete() assert len(apify_client_async_patcher.calls['run']['abort']) == 1 + + +def make_client(statuses: dict[str, str]) -> Mock: + """A client whose `run(run_id).get()` returns the run with its current status in `statuses`.""" + client = Mock() + + def run(run_id: str) -> Mock: + run_client = Mock() + run_client.get = AsyncMock(side_effect=lambda: make_run(run_id, statuses[run_id])) + run_client.resurrect = AsyncMock(side_effect=lambda: make_run(run_id, 'RUNNING')) + run_client.abort = AsyncMock() + return run_client + + client.run.side_effect = run + return client + + +async def start_child( + registry: ChildRunRegistry, client: Mock, name: str, statuses: dict[str, str], *, run_id: str | None = None +) -> Run: + """Start a named child run with the registry, adding its run to `statuses` as `RUNNING`.""" + + async def start_run() -> Run: + new_run_id = run_id or f'{name}-run' + statuses[new_run_id] = 'RUNNING' + return make_run(new_run_id, 'READY') + + run, _ = await registry.find_or_start( + name, + actor_id='some-actor', + run_input=None, + client=client, + start_run=start_run, + resurrect_run=lambda run_id: client.run(run_id).resurrect(), + ) + return run + + +async def assert_waiting(task: asyncio.Task) -> None: + """Assert that `task` is still pending after giving it a moment to run.""" + await asyncio.sleep(0.05) + assert not task.done() + + +@pytest.mark.parametrize( + 'max_concurrent_runs', + [ + pytest.param(0, id='zero'), + pytest.param(-1, id='negative'), + ], +) +async def test_set_named_child_run_limits_rejects_non_positive_limit(max_concurrent_runs: int) -> None: + """A concurrency limit below 1 is rejected.""" + async with Actor: + with pytest.raises(ValueError, match='must be at least 1'): + Actor.set_named_child_run_limits(max_concurrent_runs=max_concurrent_runs) + + +async def test_named_start_waits_while_the_limit_is_reached() -> None: + """A named start waits while the limit is reached and proceeds once an active child run finishes.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + first = await start_child(registry, client, 'first', statuses) + second_task = asyncio.create_task(start_child(registry, client, 'second', statuses)) + await assert_waiting(second_task) + + statuses[first.id] = 'SUCCEEDED' + await registry.update('first', make_run(first.id, 'SUCCEEDED')) + second = await second_task + + assert second.id == 'second-run' + + +async def test_removing_the_limit_releases_waiting_starts(monkeypatch: pytest.MonkeyPatch) -> None: + """A named start waiting under the limit proceeds once the limit is removed.""" + monkeypatch.setattr('apify._child_runs._STATUS_MAX_AGE_SECS', 0.05) + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + await start_child(registry, client, 'first', statuses) + second_task = asyncio.create_task(start_child(registry, client, 'second', statuses)) + await assert_waiting(second_task) + + registry.set_max_concurrent_runs(None) + second = await asyncio.wait_for(second_task, timeout=1) + + assert second.id == 'second-run' + + +async def test_stale_active_status_is_refreshed_before_counting(monkeypatch: pytest.MonkeyPatch) -> None: + """A child run nobody awaited is fetched again once its status is stale, so a finished one stops counting.""" + monkeypatch.setattr('apify._child_runs._STATUS_MAX_AGE_SECS', 0) + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + first = await start_child(registry, client, 'first', statuses) + second_task = asyncio.create_task(start_child(registry, client, 'second', statuses)) + await assert_waiting(second_task) + + statuses[first.id] = 'SUCCEEDED' + second = await asyncio.wait_for(second_task, timeout=1) + + assert second.id == 'second-run' + + +async def test_child_run_recorded_by_an_earlier_attempt_counts_toward_the_limit() -> None: + """An active child run recorded before a migration counts toward the limit, since it is fetched first.""" + statuses = {'old-run': 'RUNNING'} + client = make_client(statuses) + + async with Actor: + await record_child_run('first', 'old-run') + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + second_task = asyncio.create_task(start_child(registry, client, 'second', statuses)) + await assert_waiting(second_task) + second_task.cancel() + with pytest.raises(asyncio.CancelledError): + await second_task + + +async def test_child_run_that_cannot_be_fetched_does_not_block_the_limit() -> None: + """A recorded child run whose fetch fails is not counted, so it does not fail or block a named start.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + run_client_factory = client.run.side_effect + + def run(run_id: str) -> Mock: + run_client = run_client_factory(run_id) + if run_id == 'old-run': + run_client.get = AsyncMock(side_effect=RuntimeError('forbidden')) + return run_client + + client.run.side_effect = run + + async with Actor: + await record_child_run('first', 'old-run') + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + second = await asyncio.wait_for(start_child(registry, client, 'second', statuses), timeout=1) + + assert second.id == 'second-run' + + +async def test_child_run_the_api_no_longer_knows_does_not_block_the_limit() -> None: + """A recorded child run the API reports as missing is not counted, so it does not block a named start.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + run_client_factory = client.run.side_effect + + def run(run_id: str) -> Mock: + run_client = run_client_factory(run_id) + if run_id == 'old-run': + run_client.get = AsyncMock(return_value=None) + return run_client + + client.run.side_effect = run + + async with Actor: + await record_child_run('first', 'old-run') + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + second = await asyncio.wait_for(start_child(registry, client, 'second', statuses), timeout=1) + + assert second.id == 'second-run' + + +async def test_finished_recorded_child_run_is_not_fetched_for_the_limit() -> None: + """A recorded child run whose recorded status is final is not fetched to count it toward the limit.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + kvs = await Actor.open_key_value_store() + await kvs.set_value(CHILD_RUNS_KEY, {'first': stored_record('old-run', 'SUCCEEDED')}) + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + second = await asyncio.wait_for(start_child(registry, client, 'second', statuses), timeout=1) + + assert second.id == 'second-run' + assert all(call.args != ('old-run',) for call in client.run.call_args_list) + + +async def test_reattach_does_not_wait_under_the_limit() -> None: + """Reattaching to an active recorded child run returns it even while the limit is reached.""" + statuses = {'old-run': 'RUNNING'} + client = make_client(statuses) + + async with Actor: + await record_child_run('first', 'old-run') + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + run = await asyncio.wait_for(start_child(registry, client, 'first', statuses), timeout=1) + + assert run.id == 'old-run' + + +async def test_resurrection_waits_under_the_limit() -> None: + """Resurrecting an aborted child run waits while the limit is reached.""" + statuses = {'old-run': 'ABORTED'} + client = make_client(statuses) + + async with Actor: + await record_child_run('first', 'old-run') + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + second = await start_child(registry, client, 'second', statuses) + first_task = asyncio.create_task(start_child(registry, client, 'first', statuses)) + await assert_waiting(first_task) + + statuses[second.id] = 'SUCCEEDED' + await registry.update('second', make_run(second.id, 'SUCCEEDED')) + first = await first_task + + assert first.id == 'old-run' + assert first.status == 'RUNNING' + + +async def test_concurrent_named_starts_respect_the_limit() -> None: + """Concurrent named starts under different names start no more runs than the limit.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store) + await registry.load() + registry.set_max_concurrent_runs(2) + tasks = [asyncio.create_task(start_child(registry, client, f'child-{i}', statuses)) for i in range(3)] + await asyncio.sleep(0.05) + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + assert sum(task.cancelled() for task in tasks) == 1 + assert len(statuses) == 2 + + +async def test_start_in_flight_when_the_limit_is_set_counts_toward_it() -> None: + """A named start begun before the limit was set counts toward it, so a later start waits for it.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + release_first = asyncio.Event() + + async def start_first() -> Run: + await release_first.wait() + statuses['first-run'] = 'RUNNING' + return make_run('first-run', 'READY') + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store) + first_task = asyncio.create_task( + registry.find_or_start( + 'first', + actor_id='some-actor', + run_input=None, + client=client, + start_run=start_first, + resurrect_run=AsyncMock(), + ) + ) + await asyncio.sleep(0.05) + registry.set_max_concurrent_runs(1) + second_task = asyncio.create_task(start_child(registry, client, 'second', statuses)) + await assert_waiting(second_task) + + release_first.set() + await first_task + await assert_waiting(second_task) + second_task.cancel() + with pytest.raises(asyncio.CancelledError): + await second_task + + assert 'second-run' not in statuses + + +async def test_start_finishing_while_a_waiting_start_counts_is_counted_once(monkeypatch: pytest.MonkeyPatch) -> None: + """A start that finishes while a waiting start fetches runs to count them is still counted once.""" + monkeypatch.setattr('apify._child_runs._STATUS_MAX_AGE_SECS', 0.01) + statuses = {'old-run': 'RUNNING'} + client = make_client(statuses) + run_client_factory = client.run.side_effect + old_run_fetchable = asyncio.Event() + old_run_fetchable.set() + + async def get_old_run() -> Run: + await old_run_fetchable.wait() + return make_run('old-run', 'RUNNING') + + def run(run_id: str) -> Mock: + run_client = run_client_factory(run_id) + if run_id == 'old-run': + run_client.get = AsyncMock(side_effect=get_old_run) + return run_client + + client.run.side_effect = run + release_first = asyncio.Event() + + async def start_first() -> Run: + await release_first.wait() + statuses['first-run'] = 'RUNNING' + return make_run('first-run', 'READY') + + async with Actor: + await record_child_run('old', 'old-run') + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(2) + first_task = asyncio.create_task( + registry.find_or_start( + 'first', + actor_id='some-actor', + run_input=None, + client=client, + start_run=start_first, + resurrect_run=AsyncMock(), + ) + ) + await asyncio.sleep(0.05) + + old_run_fetchable.clear() + second_task = asyncio.create_task(start_child(registry, client, 'second', statuses)) + await assert_waiting(second_task) + release_first.set() + await asyncio.sleep(0.05) + old_run_fetchable.set() + + await assert_waiting(second_task) + await first_task + second_task.cancel() + with pytest.raises(asyncio.CancelledError): + await second_task + + assert 'second-run' not in statuses + + +async def test_failed_start_stops_counting_toward_the_limit() -> None: + """A named start that raises stops counting toward the limit, so the next one proceeds.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + with pytest.raises(RuntimeError, match='start failed'): + await registry.find_or_start( + 'first', + actor_id='some-actor', + run_input=None, + client=client, + start_run=AsyncMock(side_effect=RuntimeError('start failed')), + resurrect_run=AsyncMock(), + ) + second = await asyncio.wait_for(start_child(registry, client, 'second', statuses), timeout=1) + + assert second.id == 'second-run' + + +async def test_parent_abort_stops_starts_waiting_under_the_limit() -> None: + """A named start waiting under the limit raises once the parent is aborted, without starting a run.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + await start_child(registry, client, 'first', statuses) + second_task = asyncio.create_task(start_child(registry, client, 'second', statuses)) + await assert_waiting(second_task) + + await asyncio.wait_for(registry.abort_runs_with_parent(client), timeout=1) + with pytest.raises(RuntimeError, match='being aborted'): + await second_task + + assert 'second-run' not in statuses + + +async def test_abort_with_parent_start_freed_after_parent_abort_does_not_start() -> None: + """A waiting `abort_with_parent` start raises when capacity frees up after the parent began aborting.""" + statuses: dict[str, str] = {} + client = make_client(statuses) + start_run = AsyncMock(return_value=make_run('second-run', 'READY')) + + async with Actor: + registry = ChildRunRegistry(Actor.open_key_value_store) + registry.set_max_concurrent_runs(1) + first = await start_child(registry, client, 'first', statuses) + second_task = asyncio.create_task( + registry.find_or_start( + 'second', + actor_id='some-actor', + run_input=None, + client=client, + start_run=start_run, + resurrect_run=AsyncMock(), + abort_with_parent=True, + ) + ) + await assert_waiting(second_task) + + registry._parent_aborting = True + await registry.update('first', make_run(first.id, 'SUCCEEDED')) + with pytest.raises(RuntimeError, match='being aborted'): + await second_task + + start_run.assert_not_called() + + +async def test_named_call_stops_counting_when_the_run_finishes( + apify_client_async_patcher: ApifyClientAsyncPatcher, +) -> None: + """A child run awaited by a named call to its end stops counting without being fetched again.""" + apify_client_async_patcher.patch('actor', 'start', return_value=make_run('new-run', 'READY')) + # A fetch would report the run as still running, so only the awaited status can stop it counting. + apify_client_async_patcher.patch('run', 'get', return_value=make_run('new-run', 'RUNNING')) + apify_client_async_patcher.patch('run', 'wait_for_finish', return_value=make_run('new-run', 'SUCCEEDED')) + + async with Actor: + Actor.set_named_child_run_limits(max_concurrent_runs=1) + await Actor.call('some-actor', run_name='first', logger=None) + await asyncio.wait_for(Actor.start('some-actor', run_name='second'), timeout=1) + + assert len(apify_client_async_patcher.calls['actor']['start']) == 2 + + +async def test_named_call_task_stops_counting_when_the_run_finishes( + apify_client_async_patcher: ApifyClientAsyncPatcher, +) -> None: + """A task run awaited by a named task call to its end stops counting without being fetched again.""" + apify_client_async_patcher.patch('task', 'start', return_value=make_run('task-run', 'READY')) + apify_client_async_patcher.patch('actor', 'start', return_value=make_run('actor-run', 'READY')) + # A fetch would report the run as still running, so only the awaited status can stop it counting. + apify_client_async_patcher.patch('run', 'get', return_value=make_run('task-run', 'RUNNING')) + apify_client_async_patcher.patch('run', 'wait_for_finish', return_value=make_run('task-run', 'SUCCEEDED')) + + async with Actor: + Actor.set_named_child_run_limits(max_concurrent_runs=1) + await Actor.call_task('some-task', run_name='first') + await asyncio.wait_for(Actor.start('some-actor', run_name='second'), timeout=1) + + assert len(apify_client_async_patcher.calls['actor']['start']) == 1