diff --git a/src/apify/_charging.py b/src/apify/_charging.py index fc404ffb..501bef10 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -10,6 +10,8 @@ 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, @@ -17,9 +19,10 @@ 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,6 +583,12 @@ 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: @@ -571,10 +596,13 @@ def _get_event_price(self, event_name: str) -> Decimal: 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 diff --git a/tests/e2e/test_actor_charge.py b/tests/e2e/test_actor_charge.py index dc09602a..47bd86f4 100644 --- a/tests/e2e/test_actor_charge.py +++ b/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, + } diff --git a/tests/unit/actor/test_charging_manager.py b/tests/unit/actor/test_charging_manager.py index ae50b041..4acf866f 100644 --- a/tests/unit/actor/test_charging_manager.py +++ b/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