fix: load the current charged event counts after a container restart

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

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