fix: persist the charged event counts across container restarts

src/apify/_charging.py

@@ -10,16 +10,19 @@
 from typing import TYPE_CHECKING, Literal, Protocol, TypedDict
 
 from cachetools import TTLCache
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
+from pydantic.alias_generators import to_camel
 
 from apify_client._models import (
     FlatPricePerMonthActorPricingInfo,
     FreeActorPricingInfo,
     PayPerEventActorPricingInfo,
     PricePerDatasetItemActorPricingInfo,
 )
+from crawlee._utils.recoverable_state import RecoverableState
 
 from apify._utils import ReentrantLock, docs_group, ensure_context
-from apify.storages import Dataset
+from apify.storages import Dataset, KeyValueStore
 
 if TYPE_CHECKING:
     from collections.abc import AsyncGenerator
@@ -40,6 +43,9 @@
 DEFAULT_DATASET_ITEM_EVENT = 'apify-default-dataset-item'
 """Name of the synthetic event charged for each item pushed to the default dataset."""
 
+CHARGING_STATE_KEY = '__ACTOR_CHARGING_STATE'
+"""Key of the default key-value store record that carries the charged event counts across a container restart."""
+
 PricingModel = Literal['PAY_PER_EVENT', 'PRICE_PER_DATASET_ITEM', 'FLAT_PRICE_PER_MONTH', 'FREE']
 """Pricing model for an Actor."""
 
@@ -213,8 +219,9 @@ def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> No
 
         self._client = client
         self._charging_log_dataset: Dataset | None = None
+        self._charging_state: RecoverableState[ChargingState] | None = None
+        """The charged event counts, persisted to the default key-value store in a pay-per-event run on the platform."""
 
-        self._charging_state: dict[str, ChargingStateItem] = {}
         self._pricing_info: dict[str, PricingInfoItem] = {}
         self._tier_priced_events: set[str] = set()
         self._idempotent_charges: TTLCache[str, IdempotentChargeItem] = TTLCache(maxsize=float('inf'), ttl=180)
@@ -265,13 +272,24 @@ async def __aenter__(self) -> None:
 
             self._max_total_charge_usd = max_total_charge_usd
 
-        # Load charged event counts
+        # The platform fills the env var with the counts recorded when the run started, so after a reboot, migration
+        # or restart they miss what an earlier container charged. On the platform the counts are therefore also kept
+        # in the default key-value store, and the persisted ones take over where they are higher.
+        self._charging_state = RecoverableState(
+            default_state=ChargingState(),
+            persist_state_key=CHARGING_STATE_KEY,
+            persistence_enabled=self._is_at_home and self._pricing_model == 'PAY_PER_EVENT',
+            persist_state_kvs_factory=KeyValueStore.open,
+            logger=logger,
+        )
+        try:
+            await self._charging_state.initialize()
+        except ValidationError as exc:
+            raise ValueError(
+                f'The charging state under the "{CHARGING_STATE_KEY}" key in the default key-value store is malformed.'
+            ) from exc
         for event_name, count in charged_event_counts.items():
-            price = self._pricing_info.get(event_name, PricingInfoItem(Decimal(), title='')).price
-            self._charging_state[event_name] = ChargingStateItem(
-                charge_count=count,
-                total_charged_amount=count * price,
-            )
+            self._charged_event_counts[event_name] = max(self._charged_event_counts.get(event_name, 0), count)
 
         # Set up charging log dataset for local development
         if not self._is_at_home and self._pricing_model == 'PAY_PER_EVENT':
@@ -298,8 +316,12 @@ async def __aexit__(
         if not self.active:
             raise RuntimeError('Exiting an uninitialized ChargingManager')
 
-        charging_manager_ctx.set(None)
-        self.active = False
+        try:
+            if self._charging_state is not None:
+                await self._charging_state.teardown()
+        finally:
+            charging_manager_ctx.set(None)
+            self.active = False
 
     @_ensure_context
     async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str | None = None) -> ChargeResult:
@@ -390,9 +412,7 @@ async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str
 
             # Count the charge only after the API call succeeds, so a failed call retried under the same key is
             # counted once.
-            self._charging_state.setdefault(event_name, ChargingStateItem(0, Decimal()))
-            self._charging_state[event_name].charge_count += charged_count
-            self._charging_state[event_name].total_charged_amount += charged_count * pricing_info.price
+            self._charged_event_counts[event_name] = self._charged_event_counts.get(event_name, 0) + charged_count
 
             # Remember the key for every charge that was counted, including events the API never receives, such as
             # synthetic and tier-priced ones - those are counted locally and a repeat would count them twice.
@@ -432,7 +452,7 @@ async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str
     @_ensure_context
     def calculate_total_charged_amount(self) -> Decimal:
         return sum(
-            (item.total_charged_amount for item in self._charging_state.values()),
+            (count * self._get_event_price(event_name) for event_name, count in self._charged_event_counts.items()),
             start=Decimal(),
         )
 
@@ -461,8 +481,7 @@ def get_pricing_info(self) -> ActorPricingInfo:
 
     @_ensure_context
     def get_charged_event_count(self, event_name: str) -> int:
-        item = self._charging_state.get(event_name)
-        return item.charge_count if item is not None else 0
+        return self._charged_event_counts.get(event_name, 0)
 
     @_ensure_context
     def get_max_total_charge_usd(self) -> Decimal:
@@ -564,17 +583,26 @@ def _compute_charged_count(self, event_name: str, count: int) -> int:
 
         return 0
 
+    @property
+    def _charged_event_counts(self) -> dict[str, int]:
+        if self._charging_state is None:
+            raise RuntimeError('ChargingManager is not initialized')
+        return self._charging_state.current_value.charged_event_counts
+
     def _get_event_price(self, event_name: str) -> Decimal:
         pricing_info = self._pricing_info.get(event_name)
         if pricing_info is not None:
             return pricing_info.price
         return Decimal(0) if self._is_at_home else Decimal(1)
 
 
-@dataclass
-class ChargingStateItem:
-    charge_count: int
-    total_charged_amount: Decimal
+class ChargingState(BaseModel):
+    """The charged event counts of a run, as persisted in its default key-value store."""
+
+    model_config = ConfigDict(populate_by_name=True, alias_generator=to_camel)
+
+    charged_event_counts: dict[str, int] = Field(default_factory=dict)
+    """Number of charged events by event name, synthetic events included."""
 
 
 @dataclass

tests/e2e/test_actor_charge.py

@@ -226,3 +226,98 @@ async def test_actor_push_data_combined_budget_limit(
 
     assert run.status == 'SUCCEEDED'
     assert run.charged_event_counts == expected_counts
+
+
+@pytest_asyncio.fixture(scope='module', loop_scope='module')
+async def ppe_reboot_actor_build(make_actor: MakeActorFunction) -> str:
+    async def main() -> None:
+        async with Actor:
+            key_value_store = await Actor.open_key_value_store()
+            charging_manager = Actor.get_charging_manager()
+
+            if not await key_value_store.get_value('REBOOTED'):
+                for _ in range(3):
+                    await Actor.charge('row-charged')
+                await key_value_store.set_value('REBOOTED', {'rebooted': True})
+                await Actor.reboot()
+                return
+
+            # The second container records what the SDK believes right at its start, before charging anything.
+            await key_value_store.set_value(
+                'SDK_VIEW_AFTER_REBOOT',
+                {
+                    'chargedCount': charging_manager.get_charged_event_count('row-charged'),
+                    'totalChargedUsd': str(charging_manager.calculate_total_charged_amount()),
+                    'chargeableWithinLimit': charging_manager.calculate_max_event_charge_count_within_limit(
+                        'row-charged'
+                    ),
+                },
+            )
+
+    actor_client = await make_actor('ppe-reboot', main_func=main)
+
+    await actor_client.update(
+        pricing_infos=[
+            {
+                'pricingModel': 'PAY_PER_EVENT',
+                'apifyMarginPercentage': 0.0,
+                'createdAt': '2024-01-01T00:00:00.000Z',
+                'startedAt': '2024-01-01T00:00:00.000Z',
+                'pricingPerEvent': {
+                    'actorChargeEvents': {
+                        'row-charged': {
+                            'eventTitle': 'Row charged',
+                            'eventPriceUsd': 0.005,
+                            'eventDescription': 'One charged row',
+                        },
+                    },
+                },
+            },
+        ],
+    )
+
+    actor = await actor_client.get()
+
+    assert actor is not None
+    return actor.id
+
+
+@pytest_asyncio.fixture(scope='function', loop_scope='module')
+async def ppe_reboot_actor(
+    ppe_reboot_actor_build: str,
+    apify_client_async: ApifyClientAsync,
+) -> ActorClientAsync:
+    return apify_client_async.actor(ppe_reboot_actor_build)
+
+
+async def test_actor_charge_counts_survive_reboot(
+    ppe_reboot_actor: ActorClientAsync,
+    run_actor: RunActorFunction,
+    apify_client_async: ApifyClientAsync,
+) -> None:
+    """Test that a rebooted run starts from the events charged before the reboot."""
+    run = await run_actor(ppe_reboot_actor, max_total_charge_usd=Decimal('0.05'))
+
+    assert run.status == 'SUCCEEDED'
+    assert run.stats.reboot_count == 1
+
+    # Refetch until the charged event counts propagate on the platform.
+    run = await poll_until_condition(
+        partial(_get_run, apify_client_async, run.id),
+        lambda r: r.charged_event_counts == {'row-charged': 3},
+        timeout=30,
+        poll_interval=1,
+    )
+    assert run.charged_event_counts == {'row-charged': 3}
+
+    record = await apify_client_async.key_value_store(run.default_key_value_store_id).get_record(
+        'SDK_VIEW_AFTER_REBOOT'
+    )
+    assert record is not None
+
+    # Three events at $0.005 were charged before the reboot, which leaves room for seven more under the $0.05 limit.
+    assert record['value'] == {
+        'chargedCount': 3,
+        'totalChargedUsd': '0.015',
+        'chargeableWithinLimit': 7,
+    }

tests/unit/actor/test_charging_manager.py

@@ -8,9 +8,12 @@
 import pytest
 
 from apify_client._models import ActorChargeEvent, PayPerEventActorPricingInfo
+from crawlee import service_locator
+from crawlee.events._types import Event, EventPersistStateData
 
-from apify._charging import ChargingManagerImplementation
+from apify._charging import CHARGING_STATE_KEY, ChargingManagerImplementation
 from apify._configuration import Configuration
+from apify.storages import KeyValueStore
 
 
 def _make_config(**kwargs: Any) -> Configuration:
@@ -781,3 +784,87 @@ 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)
+
+
+async def test_persisted_counts_take_over_where_higher(mock_client: MagicMock) -> None:
+    """A platform start takes the counts persisted by an earlier container where they exceed the env var's."""
+    key_value_store = await KeyValueStore.open()
+    await key_value_store.set_value(CHARGING_STATE_KEY, {'chargedEventCounts': {'search': 3, 'scrape': 1}})
+
+    cm = ChargingManagerImplementation(_make_platform_ppe_config(), mock_client)
+    async with cm:
+        assert cm.get_charged_event_count('search') == 3
+        assert cm.get_charged_event_count('scrape') == 2
+        assert cm.calculate_total_charged_amount() == Decimal('0.13')
+        assert cm.calculate_max_event_charge_count_within_limit('search') == 987
+
+
+async def test_charged_counts_are_persisted(mock_client: MagicMock) -> None:
+    """Charges are written to the default key-value store on persist and once more on exit."""
+    key_value_store = await KeyValueStore.open()
+    cm = ChargingManagerImplementation(_make_platform_ppe_config(), mock_client)
+    async with cm:
+        await cm.charge('search', count=2)
+        assert cm._charging_state is not None
+        await cm._charging_state.persist_state()
+        stored = await key_value_store.get_value(CHARGING_STATE_KEY)
+        assert stored == {'chargedEventCounts': {'search': 3, 'scrape': 2}}
+
+        await cm.charge('scrape')
+
+    stored = await key_value_store.get_value(CHARGING_STATE_KEY)
+    assert stored == {'chargedEventCounts': {'search': 3, 'scrape': 3}}
+
+
+async def test_local_run_does_not_persist_counts(mock_client: MagicMock) -> None:
+    """A local PPE run keeps its counts in memory only."""
+    key_value_store = await KeyValueStore.open()
+    cm = ChargingManagerImplementation(
+        _make_platform_ppe_config(is_at_home=False, test_pay_per_event=True), mock_client
+    )
+    async with cm:
+        await cm.charge('search')
+
+    assert await key_value_store.get_value(CHARGING_STATE_KEY) is None
+
+
+async def test_persist_state_event_persists_counts(mock_client: MagicMock) -> None:
+    """The counts are written when the `PERSIST_STATE` event fires, which a reboot and a migration trigger."""
+    key_value_store = await KeyValueStore.open()
+    cm = ChargingManagerImplementation(_make_platform_ppe_config(), mock_client)
+    async with cm:
+        await cm.charge('search')
+
+        event_manager = service_locator.get_event_manager()
+        async with event_manager:
+            event_manager.emit(event=Event.PERSIST_STATE, event_data=EventPersistStateData(is_migrating=True))
+            await event_manager.wait_for_all_listeners_to_complete()
+
+            # Checked inside the context, since leaving it emits one more `PERSIST_STATE` event.
+            assert await key_value_store.get_value(CHARGING_STATE_KEY) == {
+                'chargedEventCounts': {'search': 2, 'scrape': 2}
+            }
+
+
+async def test_malformed_persisted_state_is_reported(mock_client: MagicMock) -> None:
+    """A malformed charging state record in the default key-value store fails the start with a clear error."""
+    key_value_store = await KeyValueStore.open()
+    await key_value_store.set_value(CHARGING_STATE_KEY, {'chargedEventCounts': 'not-a-mapping'})
+
+    cm = ChargingManagerImplementation(_make_platform_ppe_config(), mock_client)
+    with pytest.raises(ValueError, match='is malformed'):
+        async with cm:
+            pass