diff --git a/src/apify/_charging.py b/src/apify/_charging.py index fc404ffb..2cbd9d84 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import math from contextlib import asynccontextmanager from contextvars import ContextVar @@ -22,7 +23,7 @@ from apify.storages import Dataset if TYPE_CHECKING: - from collections.abc import AsyncGenerator + from collections.abc import AsyncGenerator, Awaitable, Callable from types import TracebackType from apify_client import ApifyClientAsync @@ -40,6 +41,9 @@ DEFAULT_DATASET_ITEM_EVENT = 'apify-default-dataset-item' """Name of the synthetic event charged for each item pushed to the default dataset.""" +ACTOR_START_EVENT = 'apify-actor-start' +"""Name of the synthetic event the platform charges when the run starts. The SDK never charges it.""" + PricingModel = Literal['PAY_PER_EVENT', 'PRICE_PER_DATASET_ITEM', 'FLAT_PRICE_PER_MONTH', 'FREE'] """Pricing model for an Actor.""" @@ -224,7 +228,7 @@ def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> No self._not_ppe_warning_printed = False self.active = False - self.charge_lock = ReentrantLock() + self.charge_lock = _ChargeLock(self._try_to_refresh_from_api) async def __aenter__(self) -> None: """Initialize the charging manager - this is called by the `Actor` class and shouldn't be invoked manually.""" @@ -551,6 +555,55 @@ async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: ), ) + def _is_chargeable_platform_run(self) -> bool: + if not self._is_at_home or self._pricing_model != 'PAY_PER_EVENT': + return False + + return any(event_name != ACTOR_START_EVENT for event_name in self._pricing_info) + + async def _try_to_refresh_from_api(self) -> None: + """Try to take the run's current charged event counts before the first charge of this container. + + The counts the platform provides come from when the run started, so after a reboot, migration or restart + they miss what an earlier container charged. The charging info from platform is affected by eventual consistency + lag so read it lazily as late as possible. The charge lock calls this the first time it is taken. + """ + run_id = self._actor_run_id + if not self._is_chargeable_platform_run() or run_id is None: + return + + try: + run = await self._client.run(run_id).get() + except asyncio.CancelledError: + # The read happens once, so a cancelled first charge or push leaves the run on the start-of-run counts. + logger.warning( + 'Fetching the current charged event counts of the run was cancelled, continuing with the counts ' + 'recorded when the run started.' + ) + raise + except Exception: + logger.warning( + 'Failed to fetch the current charged event counts of the run, continuing with the counts recorded ' + 'when the run started.', + exc_info=True, + ) + return + + if run is None: + logger.warning( + 'Actor run not found when fetching the current charged event counts, continuing with the counts ' + 'recorded when the run started.' + ) + return + + for event_name, api_count in (run.charged_event_counts or {}).items(): + local_count = self._charging_state[event_name].charge_count if event_name in self._charging_state else 0 + count = max(local_count, api_count) + self._charging_state[event_name] = ChargingStateItem( + charge_count=count, + total_charged_amount=count * self._get_event_price(event_name), + ) + def _compute_charged_count(self, event_name: str, count: int) -> int: """Compute how many of `count` events to charge, given the remaining budget.""" max_chargeable = self.calculate_max_event_charge_count_within_limit(event_name) @@ -571,6 +624,27 @@ def _get_event_price(self, event_name: str) -> Decimal: return Decimal(0) if self._is_at_home else Decimal(1) +class _ChargeLock(ReentrantLock): + """The charge lock, which also brings the charged event counts up to date the first time it is taken. + + Every budgeted operation, a charge or a push, takes this lock before it reads the counts, so the refresh happens + without the callers knowing, and concurrent first operations wait for the one read. + """ + + def __init__(self, refresh: Callable[[], Awaitable[None]]) -> None: + super().__init__() + self._refresh = refresh + self._refreshed = False + + @asynccontextmanager + async def __call__(self) -> AsyncGenerator[None]: + async with super().__call__(): + if not self._refreshed: + self._refreshed = True + await self._refresh() + yield + + @dataclass class ChargingStateItem: charge_count: int diff --git a/tests/unit/actor/test_charging_manager.py b/tests/unit/actor/test_charging_manager.py index ae50b041..acde6742 100644 --- a/tests/unit/actor/test_charging_manager.py +++ b/tests/unit/actor/test_charging_manager.py @@ -9,7 +9,7 @@ from apify_client._models import ActorChargeEvent, PayPerEventActorPricingInfo -from apify._charging import ChargingManagerImplementation +from apify._charging import ChargingManagerImplementation, charge_lock_if_charging from apify._configuration import Configuration @@ -781,3 +781,121 @@ async def test_charge_overcharges_by_one_event_when_budget_exceeded( mock_client.run.return_value.charge.assert_awaited_once_with('search', count=2, idempotency_key=None) else: mock_client.run.return_value.charge.assert_not_called() + + +def _make_platform_ppe_config(**kwargs: Any) -> Configuration: + """Helper to create a platform PPE Configuration with charged event counts in the env var.""" + defaults: dict[str, Any] = { + 'is_at_home': True, + 'actor_run_id': 'test-run-id', + 'actor_pricing_info': _make_ppe_pricing_info(), + 'charged_event_counts': {'search': 1, 'scrape': 2}, + 'max_total_charge_usd': Decimal('10.00'), + } + defaults.update(kwargs) + return _make_config(**defaults) + + +def _make_run(charged_event_counts: dict[str, int]) -> MagicMock: + """Helper to create a run as returned by the API, with the given charged event counts.""" + return MagicMock(charged_event_counts=charged_event_counts) + + +async def test_first_charge_loads_charged_event_counts_from_run(mock_client: MagicMock) -> None: + """The first charge takes the run's current counts where they exceed the env var's, as after a restart.""" + mock_client.run.return_value.get.return_value = _make_run({'search': 3, 'scrape': 1}) + config = _make_platform_ppe_config() + cm = ChargingManagerImplementation(config, mock_client) + async with cm: + mock_client.run.return_value.get.assert_not_awaited() + assert cm.get_charged_event_count('search') == 1 + + await cm.charge('search') + mock_client.run.assert_any_call('test-run-id') + mock_client.run.return_value.get.assert_awaited_once() + assert cm.get_charged_event_count('search') == 4 + assert cm.get_charged_event_count('scrape') == 2 + assert cm.calculate_total_charged_amount() == Decimal('0.14') + assert cm.calculate_max_event_charge_count_within_limit('search') == 986 + + await cm.charge('search') + mock_client.run.return_value.get.assert_awaited_once() + + +async def test_concurrent_charges_share_one_load_of_the_run(mock_client: MagicMock) -> None: + """Parallel first charges read the run once and all of them are budgeted against the loaded counts.""" + + async def slow_get() -> MagicMock: + await asyncio.sleep(0.01) + return _make_run({'search': 3}) + + mock_client.run.return_value.get.side_effect = slow_get + config = _make_platform_ppe_config( + actor_pricing_info=_make_ppe_pricing_info({'search': Decimal('0.01')}), + charged_event_counts={'search': 1}, + max_total_charge_usd=Decimal('0.05'), + ) + cm = ChargingManagerImplementation(config, mock_client) + async with cm: + results = await asyncio.gather(*(cm.charge('search') for _ in range(4))) + + mock_client.run.return_value.get.assert_awaited_once() + # Two fit in the budget after the three loaded from the run, one more overcharges, the last one is refused. + assert sum(result.charged_count for result in results) == 3 + assert cm.get_charged_event_count('search') == 6 + + +async def test_first_push_loads_charged_event_counts_from_run(mock_client: MagicMock) -> None: + """Taking the charge lock for a push loads the run's current counts, so the push limit accounts for them.""" + mock_client.run.return_value.get.return_value = _make_run({'search': 3}) + config = _make_platform_ppe_config() + cm = ChargingManagerImplementation(config, mock_client) + async with cm, charge_lock_if_charging(): + mock_client.run.return_value.get.assert_awaited_once() + assert cm.get_charged_event_count('search') == 3 + + +@pytest.mark.parametrize( + 'overrides', + [ + pytest.param( + {'actor_pricing_info': _make_ppe_pricing_info({'apify-actor-start': Decimal('0.005')})}, + id='only_start_event', + ), + pytest.param({'is_at_home': False, 'test_pay_per_event': True}, id='local_run'), + ], +) +async def test_run_is_not_fetched_when_the_sdk_cannot_charge(mock_client: MagicMock, overrides: dict[str, Any]) -> None: + """A run without SDK-chargeable events or a local run never reads the run.""" + config = _make_platform_ppe_config(**overrides) + cm = ChargingManagerImplementation(config, mock_client) + async with cm: + await cm.charge('search') + mock_client.run.return_value.get.assert_not_awaited() + + +@pytest.mark.parametrize( + ('get_side_effect', 'get_return_value'), + [ + pytest.param(RuntimeError('boom'), None, id='request_fails'), + pytest.param(None, None, id='run_not_found'), + ], +) +async def test_charge_continues_when_run_cannot_be_fetched( + mock_client: MagicMock, + caplog: pytest.LogCaptureFixture, + get_side_effect: Exception | None, + get_return_value: MagicMock | None, +) -> None: + """A failed load is logged once, and charging goes on from the env var's counts.""" + mock_client.run.return_value.get.side_effect = get_side_effect + mock_client.run.return_value.get.return_value = get_return_value + config = _make_platform_ppe_config() + cm = ChargingManagerImplementation(config, mock_client) + async with cm: + await cm.charge('search') + await cm.charge('search') + assert cm.get_charged_event_count('search') == 3 + mock_client.run.return_value.get.assert_awaited_once() + + assert caplog.text.count('continuing with the counts recorded when the run started') == 1