From 103d9e6bbb0fa59d25ca99f33015122135545b5e Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 09:05:16 +0000 Subject: [PATCH 01/18] fix: persist charging state across container restarts After a reboot, migration or restart the platform seeds the new container with the `APIFY_CHARGED_ACTOR_EVENT_COUNTS` recorded when the run started, so the charging manager began from zero charged events and allowed a fresh `max_total_charge_usd` budget. The charging manager now persists its charged event counts to the default key-value store under `__ACTOR_CHARGING_STATE` on every `PERSIST_STATE` event and right before `Actor.reboot()` calls the API, writing only when a charge happened since the last write. A container that starts more than 60 seconds after the run started reads the record and takes the persisted count of each event where it is higher than the platform's. A prompt first start makes no extra request. A record of another run, a malformed one, or a failed read is logged and ignored. Closes #1209 Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- docs/02_concepts/11_pay_per_event.mdx | 2 + src/apify/_actor.py | 20 ++- src/apify/_charging.py | 160 +++++++++++++++++++-- tests/unit/actor/test_actor_helpers.py | 23 +++ tests/unit/actor/test_charging_manager.py | 165 +++++++++++++++++++++- 5 files changed, 360 insertions(+), 10 deletions(-) diff --git a/docs/02_concepts/11_pay_per_event.mdx b/docs/02_concepts/11_pay_per_event.mdx index 6137c2978..2b7667a2b 100644 --- a/docs/02_concepts/11_pay_per_event.mdx +++ b/docs/02_concepts/11_pay_per_event.mdx @@ -78,6 +78,8 @@ The user spending limit for the run is available in your Actor code as the `ACTO When the charge limit is reached, `Actor.charge` stops charging and `Actor.push_data` stops pushing data. The platform then aborts the run automatically. However, the run keeps consuming platform resources for a short time before it stops. For details, see [Handle graceful shutdown](https://docs.apify.com/platform/actors/publishing/monetize/pay-per-event#handle-graceful-shutdown). +The charged event counts survive a reboot, a migration and a restart of the run. The SDK persists them to the default key-value store under the `__ACTOR_CHARGING_STATE` key, whenever the run persists its state and right before `Actor.reboot`, and the container started afterwards continues from them. Events charged in the few seconds before a migration or a restart can be missing from the record, so the restarted run may count slightly less than the platform. Keep that in mind when you compute the remaining budget from the charged counts yourself. + ## Test monetization locally Before releasing your monetization code to the public, test it locally. To make your Actor work in pay-per-event mode, pass it the `ACTOR_TEST_PAY_PER_EVENT` environment variable: diff --git a/src/apify/_actor.py b/src/apify/_actor.py index e3c0d3237..291ec829f 100644 --- a/src/apify/_actor.py +++ b/src/apify/_actor.py @@ -242,6 +242,11 @@ async def __aenter__(self) -> Self: event=Event.ABORTING, listener=self._abort_child_runs ) self._internal_listeners.callback(remove_listener) + # Carries the charged event counts over to the container started after a migration or a restart. + remove_listener = self.event_manager._on_internal( # noqa: SLF001 + event=Event.PERSIST_STATE, listener=self._charging_manager_implementation.persist_state + ) + self._internal_listeners.callback(remove_listener) rollback.callback(self._internal_listeners.close) self.log.debug('Event manager initialized') @@ -497,7 +502,17 @@ def event_manager(self) -> EventManager: @cached_property def _charging_manager_implementation(self) -> ChargingManagerImplementation: - return ChargingManagerImplementation(self.configuration, self.apify_client) + async def open_default_key_value_store() -> KeyValueStore: + # The charging manager is entered before the Actor is active, so the public `open_key_value_store` + # method cannot be used here. + storage_client = self._storage_client.get_suitable_storage_client() + return await KeyValueStore.open(storage_client=storage_client) + + return ChargingManagerImplementation( + self.configuration, + self.apify_client, + open_key_value_store=open_default_key_value_store, + ) @cached_property def _storage_client(self) -> SmartApifyStorageClient: @@ -1717,6 +1732,9 @@ async def safe_dispatch(listener: Any, data: Any) -> None: except TimeoutError: self.log.warning('Pre-reboot event listeners did not finish within timeout; proceeding with reboot') + # The container started by the reboot continues from the charged event counts persisted here. + await self._charging_manager_implementation.persist_state() + self.log.debug('Rebooting the Actor run.') await self.apify_client.run(self.configuration.actor_run_id).reboot() except BaseException: diff --git a/src/apify/_charging.py b/src/apify/_charging.py index fc404ffbf..a0e6296ce 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -4,12 +4,14 @@ from contextlib import asynccontextmanager from contextvars import ContextVar from dataclasses import dataclass -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta from decimal import Decimal from logging import getLogger from typing import TYPE_CHECKING, Literal, Protocol, TypedDict from cachetools import TTLCache +from pydantic import BaseModel, ConfigDict, ValidationError +from pydantic.alias_generators import to_camel from apify_client._models import ( FlatPricePerMonthActorPricingInfo, @@ -22,12 +24,13 @@ 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 from apify._configuration import Configuration + from apify.storages import KeyValueStore logger = getLogger(__name__) @@ -40,6 +43,20 @@ 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 holding the charged event counts of the run. + +The platform seeds a new container from the counts recorded when the run started, so a container started after a +reboot, migration or restart would otherwise begin from a fresh budget. The record carries the counts across. +""" + +RESTART_DETECTION_THRESHOLD = timedelta(seconds=60) +"""How long after the run started a container start is treated as a restart. + +The first container starts within seconds of the run, while a reboot, migration or restart comes later. Reading +the persisted record only after this much time keeps a normal start free of the extra key-value store request. +""" + PricingModel = Literal['PAY_PER_EVENT', 'PRICE_PER_DATASET_ITEM', 'FLAT_PRICE_PER_MONTH', 'FREE'] """Pricing model for an Actor.""" @@ -201,7 +218,13 @@ class ChargingManagerImplementation(ChargingManager): LOCAL_CHARGING_LOG_DATASET_NAME = 'charging-log' - def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> None: + def __init__( + self, + configuration: Configuration, + client: ApifyClientAsync, + *, + open_key_value_store: Callable[[], Awaitable[KeyValueStore]] | None = None, + ) -> None: self._max_total_charge_usd = ( configuration.max_total_charge_usd if configuration.max_total_charge_usd is not None else Decimal('inf') ) @@ -214,6 +237,11 @@ def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> No self._client = client self._charging_log_dataset: Dataset | None = None + self._open_key_value_store = open_key_value_store + """Opens the default key-value store, where the charged event counts are persisted across container + restarts. Without it, the counts are not persisted.""" + self._state_changed_since_persist = False + self._charging_state: dict[str, ChargingStateItem] = {} self._pricing_info: dict[str, PricingInfoItem] = {} self._tier_priced_events: set[str] = set() @@ -267,11 +295,16 @@ async def __aenter__(self) -> None: # Load charged event counts 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._set_charged_event_count(event_name, count) + + # The counts the platform provides are those recorded when the run started. After a container restart they + # miss what this run charged before it, so the counts persisted by the previous container take over where + # they are higher. + if self._should_restore_persisted_state(): + persisted_event_counts = await self._load_persisted_charged_event_counts() + for event_name, count in persisted_event_counts.items(): + if count > self.get_charged_event_count(event_name): + self._set_charged_event_count(event_name, count) # Set up charging log dataset for local development if not self._is_at_home and self._pricing_model == 'PAY_PER_EVENT': @@ -393,6 +426,7 @@ async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str 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._state_changed_since_persist = True # 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. @@ -509,6 +543,101 @@ def compute_chargeable(self) -> dict[str, int | None]: for event_name in self._pricing_info } + async def persist_state(self) -> None: + """Write the charged event counts to the default key-value store, if they changed since the last write. + + The `Actor` calls this on the `PERSIST_STATE` event and right before a reboot, so that the container started + afterwards continues from the counts of this one. Only pay-per-event runs on the platform have anything to + persist. A failed write is logged and the counts stay marked for the next attempt. + """ + open_key_value_store = self._open_key_value_store + if open_key_value_store is None or not self._persists_state() or not self._state_changed_since_persist: + return + + # Snapshot under the charge lock, so that a charge in flight is either in this record or in the next one. + async with self.charge_lock(): + record = ChargingStateRecord( + run_id=self._actor_run_id, + charged_event_counts={name: item.charge_count for name, item in self._charging_state.items()}, + updated_at=datetime.now(UTC), + ) + self._state_changed_since_persist = False + + try: + key_value_store = await open_key_value_store() + await key_value_store.set_value(CHARGING_STATE_KEY, record.model_dump(by_alias=True, mode='json')) + except Exception: + self._state_changed_since_persist = True + logger.warning('Failed to persist the charging state to the default key-value store.', exc_info=True) + else: + logger.debug('Persisted the charging state to the default key-value store.') + + def _persists_state(self) -> bool: + """Whether this run persists its charged event counts: a pay-per-event run on the platform with a store.""" + return self._is_at_home and self._pricing_model == 'PAY_PER_EVENT' and self._open_key_value_store is not None + + def _should_restore_persisted_state(self) -> bool: + """Whether this container start looks like a restart of a run that may have persisted its counts.""" + if not self._persists_state(): + return False + + started_at = self._configuration.started_at + if started_at is None: + return False + if started_at.tzinfo is None: + started_at = started_at.replace(tzinfo=UTC) + + return datetime.now(UTC) - started_at >= RESTART_DETECTION_THRESHOLD + + async def _load_persisted_charged_event_counts(self) -> dict[str, int]: + """Read the charged event counts persisted by an earlier container of this run, or nothing on any problem.""" + if self._open_key_value_store is None: + return {} + + try: + key_value_store = await self._open_key_value_store() + stored = await key_value_store.get_value(CHARGING_STATE_KEY) + except Exception: + logger.warning( + 'Failed to read the persisted charging state from the default key-value store, continuing with the ' + 'charged event counts provided by the platform.', + exc_info=True, + ) + return {} + + if stored is None: + return {} + + try: + record = ChargingStateRecord.model_validate(stored) + except ValidationError: + logger.warning( + f'The persisted charging state under the "{CHARGING_STATE_KEY}" key in the default key-value store ' + 'is malformed and is ignored.' + ) + return {} + + if record.run_id != self._actor_run_id: + logger.warning( + f'The persisted charging state in the default key-value store belongs to run "{record.run_id}" ' + 'and is ignored.' + ) + return {} + + logger.info( + 'Restored the charged event counts persisted by an earlier container of this run.', + extra={'charged_event_counts': record.charged_event_counts}, + ) + return record.charged_event_counts + + def _set_charged_event_count(self, event_name: str, count: int) -> None: + """Set the charged count of an event, pricing the total with the event's price, or zero for an unknown one.""" + 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, + ) + async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: """Fetch pricing information from environment variables or API.""" # Check if pricing info is available via environment variables @@ -589,6 +718,21 @@ class IdempotentChargeItem: charged_count: int +class ChargingStateRecord(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) + + run_id: str | None + """ID of the run the counts belong to. A record of another run is ignored.""" + + charged_event_counts: dict[str, int] + """Number of charged events by event name, synthetic events included.""" + + updated_at: datetime + """When the record was written.""" + + class _FetchedPricingInfoDict(TypedDict): pricing_info: ActorPricingInfoModel | None charged_event_counts: dict[str, int] diff --git a/tests/unit/actor/test_actor_helpers.py b/tests/unit/actor/test_actor_helpers.py index e8b0d8f95..d9fc92272 100644 --- a/tests/unit/actor/test_actor_helpers.py +++ b/tests/unit/actor/test_actor_helpers.py @@ -811,3 +811,26 @@ async def reboot_hanging_once(*_args: object, **_kwargs: object) -> None: await Actor.reboot(custom_after_sleep=timedelta(milliseconds=1)) assert len(apify_client_async_patcher.calls['run']['reboot']) == 2 + + +async def test_reboot_persists_charging_state_before_the_reboot_call( + apify_client_async_patcher: ApifyClientAsyncPatcher, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test that rebooting persists the charging state before the run is rebooted through the API.""" + apify_client_async_patcher.patch('run', 'reboot', return_value=None) + + async def persist_state() -> None: + assert len(apify_client_async_patcher.calls['run']['reboot']) == 0 + + persist_state_mock = AsyncMock(side_effect=persist_state) + + async with Actor: + Actor.configuration.is_at_home = True + Actor.configuration.actor_run_id = 'some-run-id' + monkeypatch.setattr(Actor._charging_manager_implementation, 'persist_state', persist_state_mock) + + await Actor.reboot(custom_after_sleep=timedelta(milliseconds=1)) + + persist_state_mock.assert_awaited_once() + assert len(apify_client_async_patcher.calls['run']['reboot']) == 1 diff --git a/tests/unit/actor/test_charging_manager.py b/tests/unit/actor/test_charging_manager.py index ae50b041d..c47667d4a 100644 --- a/tests/unit/actor/test_charging_manager.py +++ b/tests/unit/actor/test_charging_manager.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +from datetime import UTC, datetime, timedelta from decimal import Decimal from typing import Any from unittest.mock import AsyncMock, MagicMock @@ -9,7 +10,7 @@ from apify_client._models import ActorChargeEvent, PayPerEventActorPricingInfo -from apify._charging import ChargingManagerImplementation +from apify._charging import CHARGING_STATE_KEY, ChargingManagerImplementation from apify._configuration import Configuration @@ -781,3 +782,165 @@ 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(started_at: datetime, **kwargs: Any) -> Configuration: + """Helper to create a platform PPE Configuration for a run started at the given time.""" + 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'), + 'started_at': started_at, + } + defaults.update(kwargs) + return _make_config(**defaults) + + +@pytest.fixture +def mock_key_value_store() -> MagicMock: + return MagicMock(get_value=AsyncMock(return_value=None), set_value=AsyncMock()) + + +async def test_persist_state_writes_charged_event_counts( + mock_client: MagicMock, mock_key_value_store: MagicMock +) -> None: + """Persisting after a charge writes the run's charged event counts to the default key-value store.""" + config = _make_platform_ppe_config(started_at=datetime.now(UTC)) + cm = ChargingManagerImplementation( + config, mock_client, open_key_value_store=AsyncMock(return_value=mock_key_value_store) + ) + async with cm: + await cm.charge('search', count=2) + await cm.persist_state() + + mock_key_value_store.set_value.assert_awaited_once() + key, record = mock_key_value_store.set_value.await_args.args + assert key == CHARGING_STATE_KEY + assert record['runId'] == 'test-run-id' + assert record['chargedEventCounts'] == {'search': 3, 'scrape': 2} + + +async def test_persist_state_writes_only_changed_counts( + mock_client: MagicMock, mock_key_value_store: MagicMock +) -> None: + """Persisting writes nothing when no event was charged since the last write.""" + config = _make_platform_ppe_config(started_at=datetime.now(UTC)) + cm = ChargingManagerImplementation( + config, mock_client, open_key_value_store=AsyncMock(return_value=mock_key_value_store) + ) + async with cm: + await cm.persist_state() + mock_key_value_store.set_value.assert_not_awaited() + + await cm.charge('search') + await cm.persist_state() + await cm.persist_state() + + assert mock_key_value_store.set_value.await_count == 1 + + +async def test_persist_state_retries_after_failed_write( + mock_client: MagicMock, mock_key_value_store: MagicMock, caplog: pytest.LogCaptureFixture +) -> None: + """A failed write is logged and the counts are written again by the next persist.""" + mock_key_value_store.set_value.side_effect = [RuntimeError('boom'), None] + config = _make_platform_ppe_config(started_at=datetime.now(UTC)) + cm = ChargingManagerImplementation( + config, mock_client, open_key_value_store=AsyncMock(return_value=mock_key_value_store) + ) + async with cm: + await cm.charge('search') + await cm.persist_state() + await cm.persist_state() + + assert mock_key_value_store.set_value.await_count == 2 + assert 'Failed to persist the charging state' in caplog.text + + +async def test_persist_state_does_nothing_locally(mock_client: MagicMock, mock_key_value_store: MagicMock) -> None: + """A local PPE run keeps its counts in memory only.""" + config = _make_platform_ppe_config(started_at=datetime.now(UTC), is_at_home=False, test_pay_per_event=True) + cm = ChargingManagerImplementation( + config, mock_client, open_key_value_store=AsyncMock(return_value=mock_key_value_store) + ) + async with cm: + await cm.charge('search') + await cm.persist_state() + + mock_key_value_store.set_value.assert_not_awaited() + + +async def test_restart_restores_higher_persisted_counts( + mock_client: MagicMock, mock_key_value_store: MagicMock +) -> None: + """A container started long after the run took the persisted counts where they exceed the platform's.""" + mock_key_value_store.get_value.return_value = { + 'runId': 'test-run-id', + 'chargedEventCounts': {'search': 3, 'scrape': 1}, + 'updatedAt': '2026-01-01T00:00:00Z', + } + config = _make_platform_ppe_config(started_at=datetime.now(UTC) - timedelta(minutes=5)) + cm = ChargingManagerImplementation( + config, mock_client, open_key_value_store=AsyncMock(return_value=mock_key_value_store) + ) + async with cm: + mock_key_value_store.get_value.assert_awaited_once_with(CHARGING_STATE_KEY) + 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_fresh_start_does_not_read_persisted_counts( + mock_client: MagicMock, mock_key_value_store: MagicMock +) -> None: + """A container started right after the run uses the platform's counts without reading the store.""" + config = _make_platform_ppe_config(started_at=datetime.now(UTC)) + cm = ChargingManagerImplementation( + config, mock_client, open_key_value_store=AsyncMock(return_value=mock_key_value_store) + ) + async with cm: + mock_key_value_store.get_value.assert_not_awaited() + assert cm.get_charged_event_count('search') == 1 + + +@pytest.mark.parametrize( + 'stored', + [ + pytest.param( + {'runId': 'another-run-id', 'chargedEventCounts': {'search': 3}, 'updatedAt': '2026-01-01T00:00:00Z'}, + id='other_run', + ), + pytest.param({'chargedEventCounts': 'not-a-mapping'}, id='malformed'), + ], +) +async def test_restart_ignores_unusable_persisted_counts( + mock_client: MagicMock, mock_key_value_store: MagicMock, caplog: pytest.LogCaptureFixture, stored: Any +) -> None: + """A record of another run or a malformed one is ignored with a warning.""" + mock_key_value_store.get_value.return_value = stored + config = _make_platform_ppe_config(started_at=datetime.now(UTC) - timedelta(minutes=5)) + cm = ChargingManagerImplementation( + config, mock_client, open_key_value_store=AsyncMock(return_value=mock_key_value_store) + ) + async with cm: + assert cm.get_charged_event_count('search') == 1 + + assert 'is ignored' in caplog.text + + +async def test_restart_continues_when_persisted_counts_cannot_be_read( + mock_client: MagicMock, mock_key_value_store: MagicMock, caplog: pytest.LogCaptureFixture +) -> None: + """A failed read is logged and the platform's counts are used.""" + mock_key_value_store.get_value.side_effect = RuntimeError('boom') + config = _make_platform_ppe_config(started_at=datetime.now(UTC) - timedelta(minutes=5)) + cm = ChargingManagerImplementation( + config, mock_client, open_key_value_store=AsyncMock(return_value=mock_key_value_store) + ) + async with cm: + assert cm.get_charged_event_count('search') == 1 + + assert 'Failed to read the persisted charging state' in caplog.text From 6b310cf576eea45b3732dfb4472adb4fe9be2eb4 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 11:52:15 +0000 Subject: [PATCH 02/18] fix: refresh charged event counts from the run instead of persisting them Replaces the key-value store persistence with a read of the run through the API, which needs no writes during the run, no new record, and no reboot hook, and also covers crash restarts. A pay-per-event run on the platform whose pricing has at least one event the SDK can charge, and whose container starts more than 60 seconds after the run started, reads the run and takes the current count of each event where it is higher than the env var's. A prompt first start, a run without SDK-chargeable events, and a local run make no extra request. A failed read is logged and the env var's counts are kept. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- docs/02_concepts/11_pay_per_event.mdx | 2 +- src/apify/_actor.py | 20 +- src/apify/_charging.py | 233 ++++++++-------------- tests/unit/actor/test_actor_helpers.py | 23 --- tests/unit/actor/test_charging_manager.py | 166 +++++---------- 5 files changed, 135 insertions(+), 309 deletions(-) diff --git a/docs/02_concepts/11_pay_per_event.mdx b/docs/02_concepts/11_pay_per_event.mdx index 2b7667a2b..d4d2539b3 100644 --- a/docs/02_concepts/11_pay_per_event.mdx +++ b/docs/02_concepts/11_pay_per_event.mdx @@ -78,7 +78,7 @@ The user spending limit for the run is available in your Actor code as the `ACTO When the charge limit is reached, `Actor.charge` stops charging and `Actor.push_data` stops pushing data. The platform then aborts the run automatically. However, the run keeps consuming platform resources for a short time before it stops. For details, see [Handle graceful shutdown](https://docs.apify.com/platform/actors/publishing/monetize/pay-per-event#handle-graceful-shutdown). -The charged event counts survive a reboot, a migration and a restart of the run. The SDK persists them to the default key-value store under the `__ACTOR_CHARGING_STATE` key, whenever the run persists its state and right before `Actor.reboot`, and the container started afterwards continues from them. Events charged in the few seconds before a migration or a restart can be missing from the record, so the restarted run may count slightly less than the platform. Keep that in mind when you compute the remaining budget from the charged counts yourself. +The charged event counts survive a reboot, a migration and a restart of the run. A container that starts later than a minute after the run started reads the current counts from the run through the API, so the restarted run continues from what was already charged. Events charged in the last few seconds before the restart can be missing from the counts the platform had recorded by then, so the restarted run may count slightly less than the platform. Keep that in mind when you compute the remaining budget from the charged counts yourself. ## Test monetization locally diff --git a/src/apify/_actor.py b/src/apify/_actor.py index 291ec829f..e3c0d3237 100644 --- a/src/apify/_actor.py +++ b/src/apify/_actor.py @@ -242,11 +242,6 @@ async def __aenter__(self) -> Self: event=Event.ABORTING, listener=self._abort_child_runs ) self._internal_listeners.callback(remove_listener) - # Carries the charged event counts over to the container started after a migration or a restart. - remove_listener = self.event_manager._on_internal( # noqa: SLF001 - event=Event.PERSIST_STATE, listener=self._charging_manager_implementation.persist_state - ) - self._internal_listeners.callback(remove_listener) rollback.callback(self._internal_listeners.close) self.log.debug('Event manager initialized') @@ -502,17 +497,7 @@ def event_manager(self) -> EventManager: @cached_property def _charging_manager_implementation(self) -> ChargingManagerImplementation: - async def open_default_key_value_store() -> KeyValueStore: - # The charging manager is entered before the Actor is active, so the public `open_key_value_store` - # method cannot be used here. - storage_client = self._storage_client.get_suitable_storage_client() - return await KeyValueStore.open(storage_client=storage_client) - - return ChargingManagerImplementation( - self.configuration, - self.apify_client, - open_key_value_store=open_default_key_value_store, - ) + return ChargingManagerImplementation(self.configuration, self.apify_client) @cached_property def _storage_client(self) -> SmartApifyStorageClient: @@ -1732,9 +1717,6 @@ async def safe_dispatch(listener: Any, data: Any) -> None: except TimeoutError: self.log.warning('Pre-reboot event listeners did not finish within timeout; proceeding with reboot') - # The container started by the reboot continues from the charged event counts persisted here. - await self._charging_manager_implementation.persist_state() - self.log.debug('Rebooting the Actor run.') await self.apify_client.run(self.configuration.actor_run_id).reboot() except BaseException: diff --git a/src/apify/_charging.py b/src/apify/_charging.py index a0e6296ce..9d5054309 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -10,8 +10,6 @@ from typing import TYPE_CHECKING, Literal, Protocol, TypedDict from cachetools import TTLCache -from pydantic import BaseModel, ConfigDict, ValidationError -from pydantic.alias_generators import to_camel from apify_client._models import ( FlatPricePerMonthActorPricingInfo, @@ -24,13 +22,12 @@ from apify.storages import Dataset if TYPE_CHECKING: - from collections.abc import AsyncGenerator, Awaitable, Callable + from collections.abc import AsyncGenerator from types import TracebackType from apify_client import ApifyClientAsync from apify._configuration import Configuration - from apify.storages import KeyValueStore logger = getLogger(__name__) @@ -43,18 +40,14 @@ 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 holding the charged event counts of the run. - -The platform seeds a new container from the counts recorded when the run started, so a container started after a -reboot, migration or restart would otherwise begin from a fresh budget. The record carries the counts across. -""" +ACTOR_START_EVENT = 'apify-actor-start' +"""Name of the synthetic event the platform charges when the run starts. The SDK never charges it.""" RESTART_DETECTION_THRESHOLD = timedelta(seconds=60) """How long after the run started a container start is treated as a restart. -The first container starts within seconds of the run, while a reboot, migration or restart comes later. Reading -the persisted record only after this much time keeps a normal start free of the extra key-value store request. +The first container starts within seconds of the run, while a reboot, migration or restart comes later. Only such a +later start refreshes the charged event counts from the API, so a normal start makes no extra request. """ PricingModel = Literal['PAY_PER_EVENT', 'PRICE_PER_DATASET_ITEM', 'FLAT_PRICE_PER_MONTH', 'FREE'] @@ -218,13 +211,7 @@ class ChargingManagerImplementation(ChargingManager): LOCAL_CHARGING_LOG_DATASET_NAME = 'charging-log' - def __init__( - self, - configuration: Configuration, - client: ApifyClientAsync, - *, - open_key_value_store: Callable[[], Awaitable[KeyValueStore]] | None = None, - ) -> None: + def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> None: self._max_total_charge_usd = ( configuration.max_total_charge_usd if configuration.max_total_charge_usd is not None else Decimal('inf') ) @@ -237,11 +224,6 @@ def __init__( self._client = client self._charging_log_dataset: Dataset | None = None - self._open_key_value_store = open_key_value_store - """Opens the default key-value store, where the charged event counts are persisted across container - restarts. Without it, the counts are not persisted.""" - self._state_changed_since_persist = False - self._charging_state: dict[str, ChargingStateItem] = {} self._pricing_info: dict[str, PricingInfoItem] = {} self._tier_priced_events: set[str] = set() @@ -295,16 +277,11 @@ async def __aenter__(self) -> None: # Load charged event counts for event_name, count in charged_event_counts.items(): - self._set_charged_event_count(event_name, count) - - # The counts the platform provides are those recorded when the run started. After a container restart they - # miss what this run charged before it, so the counts persisted by the previous container take over where - # they are higher. - if self._should_restore_persisted_state(): - persisted_event_counts = await self._load_persisted_charged_event_counts() - for event_name, count in persisted_event_counts.items(): - if count > self.get_charged_event_count(event_name): - self._set_charged_event_count(event_name, count) + 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, + ) # Set up charging log dataset for local development if not self._is_at_home and self._pricing_model == 'PAY_PER_EVENT': @@ -426,7 +403,6 @@ async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str 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._state_changed_since_persist = True # 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. @@ -543,108 +519,21 @@ def compute_chargeable(self) -> dict[str, int | None]: for event_name in self._pricing_info } - async def persist_state(self) -> None: - """Write the charged event counts to the default key-value store, if they changed since the last write. - - The `Actor` calls this on the `PERSIST_STATE` event and right before a reboot, so that the container started - afterwards continues from the counts of this one. Only pay-per-event runs on the platform have anything to - persist. A failed write is logged and the counts stay marked for the next attempt. - """ - open_key_value_store = self._open_key_value_store - if open_key_value_store is None or not self._persists_state() or not self._state_changed_since_persist: - return - - # Snapshot under the charge lock, so that a charge in flight is either in this record or in the next one. - async with self.charge_lock(): - record = ChargingStateRecord( - run_id=self._actor_run_id, - charged_event_counts={name: item.charge_count for name, item in self._charging_state.items()}, - updated_at=datetime.now(UTC), - ) - self._state_changed_since_persist = False - - try: - key_value_store = await open_key_value_store() - await key_value_store.set_value(CHARGING_STATE_KEY, record.model_dump(by_alias=True, mode='json')) - except Exception: - self._state_changed_since_persist = True - logger.warning('Failed to persist the charging state to the default key-value store.', exc_info=True) - else: - logger.debug('Persisted the charging state to the default key-value store.') - - def _persists_state(self) -> bool: - """Whether this run persists its charged event counts: a pay-per-event run on the platform with a store.""" - return self._is_at_home and self._pricing_model == 'PAY_PER_EVENT' and self._open_key_value_store is not None - - def _should_restore_persisted_state(self) -> bool: - """Whether this container start looks like a restart of a run that may have persisted its counts.""" - if not self._persists_state(): - return False - - started_at = self._configuration.started_at - if started_at is None: - return False - if started_at.tzinfo is None: - started_at = started_at.replace(tzinfo=UTC) - - return datetime.now(UTC) - started_at >= RESTART_DETECTION_THRESHOLD - - async def _load_persisted_charged_event_counts(self) -> dict[str, int]: - """Read the charged event counts persisted by an earlier container of this run, or nothing on any problem.""" - if self._open_key_value_store is None: - return {} - - try: - key_value_store = await self._open_key_value_store() - stored = await key_value_store.get_value(CHARGING_STATE_KEY) - except Exception: - logger.warning( - 'Failed to read the persisted charging state from the default key-value store, continuing with the ' - 'charged event counts provided by the platform.', - exc_info=True, - ) - return {} - - if stored is None: - return {} - - try: - record = ChargingStateRecord.model_validate(stored) - except ValidationError: - logger.warning( - f'The persisted charging state under the "{CHARGING_STATE_KEY}" key in the default key-value store ' - 'is malformed and is ignored.' - ) - return {} - - if record.run_id != self._actor_run_id: - logger.warning( - f'The persisted charging state in the default key-value store belongs to run "{record.run_id}" ' - 'and is ignored.' - ) - return {} - - logger.info( - 'Restored the charged event counts persisted by an earlier container of this run.', - extra={'charged_event_counts': record.charged_event_counts}, - ) - return record.charged_event_counts - - def _set_charged_event_count(self, event_name: str, count: int) -> None: - """Set the charged count of an event, pricing the total with the event's price, or zero for an unknown one.""" - 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, - ) - async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: """Fetch pricing information from environment variables or API.""" # Check if pricing info is available via environment variables if self._configuration.actor_pricing_info is not None and self._configuration.charged_event_counts is not None: + charged_event_counts = self._configuration.charged_event_counts + + # The platform fills the env var with the counts recorded when the run started. A container started + # after a reboot, migration or restart would miss what the run charged before it, so it takes the + # current counts from the run instead. + if self._is_restarted_chargeable_run(self._configuration.actor_pricing_info): + charged_event_counts = await self._fetch_current_charged_event_counts(charged_event_counts) + return _FetchedPricingInfoDict( pricing_info=self._configuration.actor_pricing_info, - charged_event_counts=self._configuration.charged_event_counts, + charged_event_counts=charged_event_counts, max_total_charge_usd=( self._configuration.max_total_charge_usd if self._configuration.max_total_charge_usd is not None @@ -680,6 +569,73 @@ async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: ), ) + def _is_restarted_chargeable_run(self, pricing_info: ActorPricingInfoModel) -> bool: + """Whether this container start looks like a restart of a platform run the SDK charges events for. + + A run is chargeable when its pricing has at least one priced event other than the start event, which only + the platform charges. For any other run the charged event counts cannot have moved since the run started. + """ + if not self._is_at_home or not isinstance(pricing_info, PayPerEventActorPricingInfo): + return False + + actor_charge_events = pricing_info.pricing_per_event.actor_charge_events or {} + has_chargeable_event = any( + event_name != ACTOR_START_EVENT and event_pricing.event_price_usd is not None + for event_name, event_pricing in actor_charge_events.items() + ) + if not has_chargeable_event: + return False + + started_at = self._configuration.started_at + if started_at is None: + return False + if started_at.tzinfo is None: + started_at = started_at.replace(tzinfo=UTC) + + return datetime.now(UTC) - started_at >= RESTART_DETECTION_THRESHOLD + + async def _fetch_current_charged_event_counts(self, initial_counts: dict[str, int]) -> dict[str, int]: + """Merge the counts recorded when the run started with the run's current counts, taking the higher of each. + + The run is read through the API. When that fails, the initial counts are kept and a warning is logged, since + a run that cannot refresh its counts is still better off running than failing to start. + """ + if self._actor_run_id is None: + raise RuntimeError('Actor run ID not found even though the Actor is running on Apify') + + try: + run = await self._client.run(self._actor_run_id).get() + except Exception: + logger.warning( + 'Failed to fetch the current charged event counts of the run after a restart, continuing with the ' + 'counts recorded when the run started.', + exc_info=True, + ) + return initial_counts + + if run is None: + logger.warning( + 'Actor run not found when fetching the current charged event counts after a restart, continuing ' + 'with the counts recorded when the run started.' + ) + return initial_counts + + merged_counts = dict(initial_counts) + for event_name, count in (run.charged_event_counts or {}).items(): + merged_counts[event_name] = max(merged_counts.get(event_name, 0), count) + + logger.info( + 'Refreshed the charged event counts from the run after a restart.', + extra={ + 'charged_event_counts': merged_counts, + 'reboot_count': run.stats.reboot_count, + 'migration_count': run.stats.migration_count, + 'restart_count': run.stats.restart_count, + 'resurrect_count': run.stats.resurrect_count, + }, + ) + return merged_counts + 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) @@ -718,21 +674,6 @@ class IdempotentChargeItem: charged_count: int -class ChargingStateRecord(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) - - run_id: str | None - """ID of the run the counts belong to. A record of another run is ignored.""" - - charged_event_counts: dict[str, int] - """Number of charged events by event name, synthetic events included.""" - - updated_at: datetime - """When the record was written.""" - - class _FetchedPricingInfoDict(TypedDict): pricing_info: ActorPricingInfoModel | None charged_event_counts: dict[str, int] diff --git a/tests/unit/actor/test_actor_helpers.py b/tests/unit/actor/test_actor_helpers.py index d9fc92272..e8b0d8f95 100644 --- a/tests/unit/actor/test_actor_helpers.py +++ b/tests/unit/actor/test_actor_helpers.py @@ -811,26 +811,3 @@ async def reboot_hanging_once(*_args: object, **_kwargs: object) -> None: await Actor.reboot(custom_after_sleep=timedelta(milliseconds=1)) assert len(apify_client_async_patcher.calls['run']['reboot']) == 2 - - -async def test_reboot_persists_charging_state_before_the_reboot_call( - apify_client_async_patcher: ApifyClientAsyncPatcher, - monkeypatch: pytest.MonkeyPatch, -) -> None: - """Test that rebooting persists the charging state before the run is rebooted through the API.""" - apify_client_async_patcher.patch('run', 'reboot', return_value=None) - - async def persist_state() -> None: - assert len(apify_client_async_patcher.calls['run']['reboot']) == 0 - - persist_state_mock = AsyncMock(side_effect=persist_state) - - async with Actor: - Actor.configuration.is_at_home = True - Actor.configuration.actor_run_id = 'some-run-id' - monkeypatch.setattr(Actor._charging_manager_implementation, 'persist_state', persist_state_mock) - - await Actor.reboot(custom_after_sleep=timedelta(milliseconds=1)) - - persist_state_mock.assert_awaited_once() - assert len(apify_client_async_patcher.calls['run']['reboot']) == 1 diff --git a/tests/unit/actor/test_charging_manager.py b/tests/unit/actor/test_charging_manager.py index c47667d4a..fd2b090ff 100644 --- a/tests/unit/actor/test_charging_manager.py +++ b/tests/unit/actor/test_charging_manager.py @@ -10,7 +10,7 @@ from apify_client._models import ActorChargeEvent, PayPerEventActorPricingInfo -from apify._charging import CHARGING_STATE_KEY, ChargingManagerImplementation +from apify._charging import ChargingManagerImplementation from apify._configuration import Configuration @@ -784,163 +784,89 @@ async def test_charge_overcharges_by_one_event_when_budget_exceeded( mock_client.run.return_value.charge.assert_not_called() -def _make_platform_ppe_config(started_at: datetime, **kwargs: Any) -> Configuration: - """Helper to create a platform PPE Configuration for a run started at the given time.""" +def _make_platform_ppe_config(**kwargs: Any) -> Configuration: + """Helper to create a platform PPE Configuration, by default for a run started minutes ago.""" 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'), - 'started_at': started_at, + 'started_at': datetime.now(UTC) - timedelta(minutes=5), } defaults.update(kwargs) return _make_config(**defaults) -@pytest.fixture -def mock_key_value_store() -> MagicMock: - return MagicMock(get_value=AsyncMock(return_value=None), set_value=AsyncMock()) - - -async def test_persist_state_writes_charged_event_counts( - mock_client: MagicMock, mock_key_value_store: MagicMock -) -> None: - """Persisting after a charge writes the run's charged event counts to the default key-value store.""" - config = _make_platform_ppe_config(started_at=datetime.now(UTC)) - cm = ChargingManagerImplementation( - config, mock_client, open_key_value_store=AsyncMock(return_value=mock_key_value_store) - ) - async with cm: - await cm.charge('search', count=2) - await cm.persist_state() - - mock_key_value_store.set_value.assert_awaited_once() - key, record = mock_key_value_store.set_value.await_args.args - assert key == CHARGING_STATE_KEY - assert record['runId'] == 'test-run-id' - assert record['chargedEventCounts'] == {'search': 3, 'scrape': 2} - - -async def test_persist_state_writes_only_changed_counts( - mock_client: MagicMock, mock_key_value_store: MagicMock -) -> None: - """Persisting writes nothing when no event was charged since the last write.""" - config = _make_platform_ppe_config(started_at=datetime.now(UTC)) - cm = ChargingManagerImplementation( - config, mock_client, open_key_value_store=AsyncMock(return_value=mock_key_value_store) - ) - async with cm: - await cm.persist_state() - mock_key_value_store.set_value.assert_not_awaited() - - await cm.charge('search') - await cm.persist_state() - await cm.persist_state() - - assert mock_key_value_store.set_value.await_count == 1 - - -async def test_persist_state_retries_after_failed_write( - mock_client: MagicMock, mock_key_value_store: MagicMock, caplog: pytest.LogCaptureFixture -) -> None: - """A failed write is logged and the counts are written again by the next persist.""" - mock_key_value_store.set_value.side_effect = [RuntimeError('boom'), None] - config = _make_platform_ppe_config(started_at=datetime.now(UTC)) - cm = ChargingManagerImplementation( - config, mock_client, open_key_value_store=AsyncMock(return_value=mock_key_value_store) - ) - async with cm: - await cm.charge('search') - await cm.persist_state() - await cm.persist_state() - - assert mock_key_value_store.set_value.await_count == 2 - assert 'Failed to persist the charging state' in caplog.text +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.""" + stats = MagicMock(reboot_count=1, migration_count=0, restart_count=0, resurrect_count=0) + return MagicMock(charged_event_counts=charged_event_counts, stats=stats) -async def test_persist_state_does_nothing_locally(mock_client: MagicMock, mock_key_value_store: MagicMock) -> None: - """A local PPE run keeps its counts in memory only.""" - config = _make_platform_ppe_config(started_at=datetime.now(UTC), is_at_home=False, test_pay_per_event=True) - cm = ChargingManagerImplementation( - config, mock_client, open_key_value_store=AsyncMock(return_value=mock_key_value_store) - ) - async with cm: - await cm.charge('search') - await cm.persist_state() - - mock_key_value_store.set_value.assert_not_awaited() - - -async def test_restart_restores_higher_persisted_counts( - mock_client: MagicMock, mock_key_value_store: MagicMock -) -> None: - """A container started long after the run took the persisted counts where they exceed the platform's.""" - mock_key_value_store.get_value.return_value = { - 'runId': 'test-run-id', - 'chargedEventCounts': {'search': 3, 'scrape': 1}, - 'updatedAt': '2026-01-01T00:00:00Z', - } +async def test_restart_refreshes_charged_event_counts_from_run(mock_client: MagicMock) -> None: + """A container started long after the run takes the run's current counts where they exceed the env var's.""" + mock_client.run.return_value.get.return_value = _make_run({'search': 3, 'scrape': 1}) config = _make_platform_ppe_config(started_at=datetime.now(UTC) - timedelta(minutes=5)) - cm = ChargingManagerImplementation( - config, mock_client, open_key_value_store=AsyncMock(return_value=mock_key_value_store) - ) + cm = ChargingManagerImplementation(config, mock_client) async with cm: - mock_key_value_store.get_value.assert_awaited_once_with(CHARGING_STATE_KEY) + mock_client.run.assert_called_with('test-run-id') + mock_client.run.return_value.get.assert_awaited_once() 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_fresh_start_does_not_read_persisted_counts( - mock_client: MagicMock, mock_key_value_store: MagicMock -) -> None: - """A container started right after the run uses the platform's counts without reading the store.""" +async def test_fresh_start_does_not_fetch_run(mock_client: MagicMock) -> None: + """A container started right after the run uses the env var's counts without an API call.""" config = _make_platform_ppe_config(started_at=datetime.now(UTC)) - cm = ChargingManagerImplementation( - config, mock_client, open_key_value_store=AsyncMock(return_value=mock_key_value_store) - ) + cm = ChargingManagerImplementation(config, mock_client) async with cm: - mock_key_value_store.get_value.assert_not_awaited() + mock_client.run.return_value.get.assert_not_awaited() assert cm.get_charged_event_count('search') == 1 @pytest.mark.parametrize( - 'stored', + 'overrides', [ pytest.param( - {'runId': 'another-run-id', 'chargedEventCounts': {'search': 3}, 'updatedAt': '2026-01-01T00:00:00Z'}, - id='other_run', + {'actor_pricing_info': _make_ppe_pricing_info({'apify-actor-start': Decimal('0.005')})}, + id='only_start_event', ), - pytest.param({'chargedEventCounts': 'not-a-mapping'}, id='malformed'), + pytest.param({'is_at_home': False, 'test_pay_per_event': True}, id='local_run'), + pytest.param({'started_at': None}, id='unknown_start_time'), ], ) -async def test_restart_ignores_unusable_persisted_counts( - mock_client: MagicMock, mock_key_value_store: MagicMock, caplog: pytest.LogCaptureFixture, stored: Any +async def test_restart_refresh_is_skipped_when_nothing_to_refresh( + mock_client: MagicMock, overrides: dict[str, Any] ) -> None: - """A record of another run or a malformed one is ignored with a warning.""" - mock_key_value_store.get_value.return_value = stored - config = _make_platform_ppe_config(started_at=datetime.now(UTC) - timedelta(minutes=5)) - cm = ChargingManagerImplementation( - config, mock_client, open_key_value_store=AsyncMock(return_value=mock_key_value_store) - ) + """A run without SDK-chargeable events, a local run, or one with an unknown start time makes no API call.""" + config = _make_platform_ppe_config(**overrides) + cm = ChargingManagerImplementation(config, mock_client) async with cm: - assert cm.get_charged_event_count('search') == 1 + mock_client.run.return_value.get.assert_not_awaited() - assert 'is ignored' in caplog.text - -async def test_restart_continues_when_persisted_counts_cannot_be_read( - mock_client: MagicMock, mock_key_value_store: MagicMock, caplog: pytest.LogCaptureFixture +@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_restart_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 read is logged and the platform's counts are used.""" - mock_key_value_store.get_value.side_effect = RuntimeError('boom') + """A failed refresh is logged and the env var's counts are used.""" + 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(started_at=datetime.now(UTC) - timedelta(minutes=5)) - cm = ChargingManagerImplementation( - config, mock_client, open_key_value_store=AsyncMock(return_value=mock_key_value_store) - ) + cm = ChargingManagerImplementation(config, mock_client) async with cm: assert cm.get_charged_event_count('search') == 1 - assert 'Failed to read the persisted charging state' in caplog.text + assert 'continuing with the counts recorded when the run started' in caplog.text From dd29d178f11e9cc505eee382dfc4ff23d2534ae6 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 12:07:46 +0000 Subject: [PATCH 03/18] fix: load current charged event counts on every chargeable platform start A reboot restarts the container in about a second, so a restart cannot be told apart from a first start by the time elapsed since the run started. The time threshold is dropped, and every platform run with a pay-per-event pricing that has at least one event the SDK can charge reads the run at start and takes the current count of each event where it is higher than the env var's. Runs without SDK-chargeable events and local runs make no extra request. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- docs/02_concepts/11_pay_per_event.mdx | 2 +- src/apify/_charging.py | 47 ++++++++--------------- tests/unit/actor/test_charging_manager.py | 30 ++++----------- 3 files changed, 26 insertions(+), 53 deletions(-) diff --git a/docs/02_concepts/11_pay_per_event.mdx b/docs/02_concepts/11_pay_per_event.mdx index d4d2539b3..847060e56 100644 --- a/docs/02_concepts/11_pay_per_event.mdx +++ b/docs/02_concepts/11_pay_per_event.mdx @@ -78,7 +78,7 @@ The user spending limit for the run is available in your Actor code as the `ACTO When the charge limit is reached, `Actor.charge` stops charging and `Actor.push_data` stops pushing data. The platform then aborts the run automatically. However, the run keeps consuming platform resources for a short time before it stops. For details, see [Handle graceful shutdown](https://docs.apify.com/platform/actors/publishing/monetize/pay-per-event#handle-graceful-shutdown). -The charged event counts survive a reboot, a migration and a restart of the run. A container that starts later than a minute after the run started reads the current counts from the run through the API, so the restarted run continues from what was already charged. Events charged in the last few seconds before the restart can be missing from the counts the platform had recorded by then, so the restarted run may count slightly less than the platform. Keep that in mind when you compute the remaining budget from the charged counts yourself. +The charged event counts survive a reboot, a migration and a restart of the run. At start, the SDK reads the current counts from the run through the API, so a restarted run continues from what was already charged. Events charged in the last few seconds before the restart can be missing from the counts the platform had recorded by then, so the restarted run may count slightly less than the platform. Keep that in mind when you compute the remaining budget from the charged counts yourself. ## Test monetization locally diff --git a/src/apify/_charging.py b/src/apify/_charging.py index 9d5054309..372bea428 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -1,10 +1,11 @@ from __future__ import annotations +import logging import math from contextlib import asynccontextmanager from contextvars import ContextVar from dataclasses import dataclass -from datetime import UTC, datetime, timedelta +from datetime import UTC, datetime from decimal import Decimal from logging import getLogger from typing import TYPE_CHECKING, Literal, Protocol, TypedDict @@ -43,13 +44,6 @@ ACTOR_START_EVENT = 'apify-actor-start' """Name of the synthetic event the platform charges when the run starts. The SDK never charges it.""" -RESTART_DETECTION_THRESHOLD = timedelta(seconds=60) -"""How long after the run started a container start is treated as a restart. - -The first container starts within seconds of the run, while a reboot, migration or restart comes later. Only such a -later start refreshes the charged event counts from the API, so a normal start makes no extra request. -""" - PricingModel = Literal['PAY_PER_EVENT', 'PRICE_PER_DATASET_ITEM', 'FLAT_PRICE_PER_MONTH', 'FREE'] """Pricing model for an Actor.""" @@ -526,9 +520,9 @@ async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: charged_event_counts = self._configuration.charged_event_counts # The platform fills the env var with the counts recorded when the run started. A container started - # after a reboot, migration or restart would miss what the run charged before it, so it takes the - # current counts from the run instead. - if self._is_restarted_chargeable_run(self._configuration.actor_pricing_info): + # after a reboot, migration or restart would miss what the run charged before it, and a restart cannot + # be told apart from a first start, so every start takes the current counts from the run. + if self._is_chargeable_platform_run(self._configuration.actor_pricing_info): charged_event_counts = await self._fetch_current_charged_event_counts(charged_event_counts) return _FetchedPricingInfoDict( @@ -569,8 +563,8 @@ async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: ), ) - def _is_restarted_chargeable_run(self, pricing_info: ActorPricingInfoModel) -> bool: - """Whether this container start looks like a restart of a platform run the SDK charges events for. + def _is_chargeable_platform_run(self, pricing_info: ActorPricingInfoModel) -> bool: + """Whether this is a platform run the SDK charges events for. A run is chargeable when its pricing has at least one priced event other than the start event, which only the platform charges. For any other run the charged event counts cannot have moved since the run started. @@ -579,20 +573,10 @@ def _is_restarted_chargeable_run(self, pricing_info: ActorPricingInfoModel) -> b return False actor_charge_events = pricing_info.pricing_per_event.actor_charge_events or {} - has_chargeable_event = any( + return any( event_name != ACTOR_START_EVENT and event_pricing.event_price_usd is not None for event_name, event_pricing in actor_charge_events.items() ) - if not has_chargeable_event: - return False - - started_at = self._configuration.started_at - if started_at is None: - return False - if started_at.tzinfo is None: - started_at = started_at.replace(tzinfo=UTC) - - return datetime.now(UTC) - started_at >= RESTART_DETECTION_THRESHOLD async def _fetch_current_charged_event_counts(self, initial_counts: dict[str, int]) -> dict[str, int]: """Merge the counts recorded when the run started with the run's current counts, taking the higher of each. @@ -607,16 +591,16 @@ async def _fetch_current_charged_event_counts(self, initial_counts: dict[str, in run = await self._client.run(self._actor_run_id).get() except Exception: logger.warning( - 'Failed to fetch the current charged event counts of the run after a restart, continuing with the ' - 'counts recorded when the run started.', + 'Failed to fetch the current charged event counts of the run, continuing with the counts recorded ' + 'when the run started.', exc_info=True, ) return initial_counts if run is None: logger.warning( - 'Actor run not found when fetching the current charged event counts after a restart, continuing ' - 'with the counts recorded when the run started.' + 'Actor run not found when fetching the current charged event counts, continuing with the counts ' + 'recorded when the run started.' ) return initial_counts @@ -624,8 +608,11 @@ async def _fetch_current_charged_event_counts(self, initial_counts: dict[str, in for event_name, count in (run.charged_event_counts or {}).items(): merged_counts[event_name] = max(merged_counts.get(event_name, 0), count) - logger.info( - 'Refreshed the charged event counts from the run after a restart.', + # Counts moved only when an earlier container of this run charged something, which is worth a note. + log_level = logging.INFO if merged_counts != initial_counts else logging.DEBUG + logger.log( + log_level, + 'Loaded the current charged event counts from the run.', extra={ 'charged_event_counts': merged_counts, 'reboot_count': run.stats.reboot_count, diff --git a/tests/unit/actor/test_charging_manager.py b/tests/unit/actor/test_charging_manager.py index fd2b090ff..77e337fce 100644 --- a/tests/unit/actor/test_charging_manager.py +++ b/tests/unit/actor/test_charging_manager.py @@ -1,7 +1,6 @@ from __future__ import annotations import asyncio -from datetime import UTC, datetime, timedelta from decimal import Decimal from typing import Any from unittest.mock import AsyncMock, MagicMock @@ -785,14 +784,13 @@ async def test_charge_overcharges_by_one_event_when_budget_exceeded( def _make_platform_ppe_config(**kwargs: Any) -> Configuration: - """Helper to create a platform PPE Configuration, by default for a run started minutes ago.""" + """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'), - 'started_at': datetime.now(UTC) - timedelta(minutes=5), } defaults.update(kwargs) return _make_config(**defaults) @@ -804,10 +802,10 @@ def _make_run(charged_event_counts: dict[str, int]) -> MagicMock: return MagicMock(charged_event_counts=charged_event_counts, stats=stats) -async def test_restart_refreshes_charged_event_counts_from_run(mock_client: MagicMock) -> None: - """A container started long after the run takes the run's current counts where they exceed the env var's.""" +async def test_platform_start_loads_charged_event_counts_from_run(mock_client: MagicMock) -> None: + """A platform run 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(started_at=datetime.now(UTC) - timedelta(minutes=5)) + config = _make_platform_ppe_config() cm = ChargingManagerImplementation(config, mock_client) async with cm: mock_client.run.assert_called_with('test-run-id') @@ -818,15 +816,6 @@ async def test_restart_refreshes_charged_event_counts_from_run(mock_client: Magi assert cm.calculate_max_event_charge_count_within_limit('search') == 987 -async def test_fresh_start_does_not_fetch_run(mock_client: MagicMock) -> None: - """A container started right after the run uses the env var's counts without an API call.""" - config = _make_platform_ppe_config(started_at=datetime.now(UTC)) - 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 - - @pytest.mark.parametrize( 'overrides', [ @@ -835,13 +824,10 @@ async def test_fresh_start_does_not_fetch_run(mock_client: MagicMock) -> None: id='only_start_event', ), pytest.param({'is_at_home': False, 'test_pay_per_event': True}, id='local_run'), - pytest.param({'started_at': None}, id='unknown_start_time'), ], ) -async def test_restart_refresh_is_skipped_when_nothing_to_refresh( - mock_client: MagicMock, overrides: dict[str, Any] -) -> None: - """A run without SDK-chargeable events, a local run, or one with an unknown start time makes no API call.""" +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 makes no API call at start.""" config = _make_platform_ppe_config(**overrides) cm = ChargingManagerImplementation(config, mock_client) async with cm: @@ -855,7 +841,7 @@ async def test_restart_refresh_is_skipped_when_nothing_to_refresh( pytest.param(None, None, id='run_not_found'), ], ) -async def test_restart_continues_when_run_cannot_be_fetched( +async def test_start_continues_when_run_cannot_be_fetched( mock_client: MagicMock, caplog: pytest.LogCaptureFixture, get_side_effect: Exception | None, @@ -864,7 +850,7 @@ async def test_restart_continues_when_run_cannot_be_fetched( """A failed refresh is logged and the env var's counts are used.""" 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(started_at=datetime.now(UTC) - timedelta(minutes=5)) + config = _make_platform_ppe_config() cm = ChargingManagerImplementation(config, mock_client) async with cm: assert cm.get_charged_event_count('search') == 1 From b3c004af248c9b2071da2c4addf8125aa7f062dc Mon Sep 17 00:00:00 2001 From: Josef Prochazka Date: Fri, 30 Oct 2026 13:51:40 +0100 Subject: [PATCH 04/18] Remove the logging --- src/apify/_charging.py | 24 ++++-------------------- 1 file changed, 4 insertions(+), 20 deletions(-) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index 372bea428..c2dd35d1a 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -1,6 +1,5 @@ from __future__ import annotations -import logging import math from contextlib import asynccontextmanager from contextvars import ContextVar @@ -523,7 +522,7 @@ async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: # after a reboot, migration or restart would miss what the run charged before it, and a restart cannot # be told apart from a first start, so every start takes the current counts from the run. if self._is_chargeable_platform_run(self._configuration.actor_pricing_info): - charged_event_counts = await self._fetch_current_charged_event_counts(charged_event_counts) + charged_event_counts = await self._update_charged_event_counts(charged_event_counts) return _FetchedPricingInfoDict( pricing_info=self._configuration.actor_pricing_info, @@ -578,7 +577,7 @@ def _is_chargeable_platform_run(self, pricing_info: ActorPricingInfoModel) -> bo for event_name, event_pricing in actor_charge_events.items() ) - async def _fetch_current_charged_event_counts(self, initial_counts: dict[str, int]) -> dict[str, int]: + async def _update_charged_event_counts(self, initial_counts: dict[str, int]) -> dict[str, int]: """Merge the counts recorded when the run started with the run's current counts, taking the higher of each. The run is read through the API. When that fails, the initial counts are kept and a warning is logged, since @@ -604,24 +603,9 @@ async def _fetch_current_charged_event_counts(self, initial_counts: dict[str, in ) return initial_counts - merged_counts = dict(initial_counts) for event_name, count in (run.charged_event_counts or {}).items(): - merged_counts[event_name] = max(merged_counts.get(event_name, 0), count) - - # Counts moved only when an earlier container of this run charged something, which is worth a note. - log_level = logging.INFO if merged_counts != initial_counts else logging.DEBUG - logger.log( - log_level, - 'Loaded the current charged event counts from the run.', - extra={ - 'charged_event_counts': merged_counts, - 'reboot_count': run.stats.reboot_count, - 'migration_count': run.stats.migration_count, - 'restart_count': run.stats.restart_count, - 'resurrect_count': run.stats.resurrect_count, - }, - ) - return merged_counts + initial_counts[event_name] = max(initial_counts.get(event_name, 0), count) + return initial_counts def _compute_charged_count(self, event_name: str, count: int) -> int: """Compute how many of `count` events to charge, given the remaining budget.""" From 5ea6f72ae30d8d5b0ed15894190dc21a4b59fb2a Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 13:54:57 +0000 Subject: [PATCH 05/18] fix: load the run's charged event counts on the first charge or push The run is read before the first charge or budgeted push of a container instead of at init, so the platform has time to record the last charges of a previous container, and a run that never charges makes no request. Until then, the synchronous methods report the counts the run started with. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- docs/02_concepts/11_pay_per_event.mdx | 2 +- src/apify/_charging.py | 50 +++++++++++++++++++---- tests/unit/actor/test_charging_manager.py | 48 +++++++++++++++------- 3 files changed, 76 insertions(+), 24 deletions(-) diff --git a/docs/02_concepts/11_pay_per_event.mdx b/docs/02_concepts/11_pay_per_event.mdx index 847060e56..8b4f7504e 100644 --- a/docs/02_concepts/11_pay_per_event.mdx +++ b/docs/02_concepts/11_pay_per_event.mdx @@ -78,7 +78,7 @@ The user spending limit for the run is available in your Actor code as the `ACTO When the charge limit is reached, `Actor.charge` stops charging and `Actor.push_data` stops pushing data. The platform then aborts the run automatically. However, the run keeps consuming platform resources for a short time before it stops. For details, see [Handle graceful shutdown](https://docs.apify.com/platform/actors/publishing/monetize/pay-per-event#handle-graceful-shutdown). -The charged event counts survive a reboot, a migration and a restart of the run. At start, the SDK reads the current counts from the run through the API, so a restarted run continues from what was already charged. Events charged in the last few seconds before the restart can be missing from the counts the platform had recorded by then, so the restarted run may count slightly less than the platform. Keep that in mind when you compute the remaining budget from the charged counts yourself. +The charged event counts survive a reboot, a migration and a restart of the run. Before the first charge or push of a container, the SDK reads the current counts from the run through the API, so a restarted run continues from what was already charged. Until then, the counts reported by the `ChargingManager` are those recorded when the run started. Events charged in the last few seconds before the restart can be missing from the counts the platform had recorded by then, so the restarted run may count slightly less than the platform. Keep that in mind when you compute the remaining budget from the charged counts yourself. ## Test monetization locally diff --git a/src/apify/_charging.py b/src/apify/_charging.py index c2dd35d1a..0ac251c54 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -31,7 +31,9 @@ logger = getLogger(__name__) -charging_manager_ctx: ContextVar[ChargingManager | None] = ContextVar('charging_manager_ctx', default=None) +charging_manager_ctx: ContextVar[ChargingManagerImplementation | None] = ContextVar( + 'charging_manager_ctx', default=None +) """Holds the current `ChargingManager` instance, if any. Allows PPE-aware dataset clients to access the charging manager without needing to pass it explicitly. @@ -63,6 +65,7 @@ async def charge_lock_if_charging() -> AsyncGenerator[None]: return async with charging_manager.charge_lock(): + await charging_manager.ensure_current_charged_event_counts() yield @@ -220,6 +223,8 @@ def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> No self._charging_state: dict[str, ChargingStateItem] = {} self._pricing_info: dict[str, PricingInfoItem] = {} self._tier_priced_events: set[str] = set() + self._charged_event_counts_need_refresh = False + """Whether the run's current charged event counts are still to be taken before the first charge.""" self._idempotent_charges: TTLCache[str, IdempotentChargeItem] = TTLCache(maxsize=float('inf'), ttl=180) """Charges made with an idempotency key, by key. Keys expire after three minutes, the same as the platform's own idempotency records.""" @@ -331,6 +336,8 @@ async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str ) async with self.charge_lock(): + await self.ensure_current_charged_event_counts() + # A repeat the platform would discard is resolved here too - counting it would inflate the charging state # and make the run hit `max_total_charge_usd` early. The key alone decides, as on the platform. if idempotency_key is not None and (previous := self._idempotent_charges.get(idempotency_key)) is not None: @@ -516,17 +523,16 @@ async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: """Fetch pricing information from environment variables or API.""" # Check if pricing info is available via environment variables if self._configuration.actor_pricing_info is not None and self._configuration.charged_event_counts is not None: - charged_event_counts = self._configuration.charged_event_counts - # The platform fills the env var with the counts recorded when the run started. A container started # after a reboot, migration or restart would miss what the run charged before it, and a restart cannot - # be told apart from a first start, so every start takes the current counts from the run. - if self._is_chargeable_platform_run(self._configuration.actor_pricing_info): - charged_event_counts = await self._update_charged_event_counts(charged_event_counts) + # be told apart from a first start, so the run's current counts are taken before the first charge. + self._charged_event_counts_need_refresh = self._is_chargeable_platform_run( + self._configuration.actor_pricing_info + ) return _FetchedPricingInfoDict( pricing_info=self._configuration.actor_pricing_info, - charged_event_counts=charged_event_counts, + charged_event_counts=self._configuration.charged_event_counts, max_total_charge_usd=( self._configuration.max_total_charge_usd if self._configuration.max_total_charge_usd is not None @@ -577,11 +583,37 @@ def _is_chargeable_platform_run(self, pricing_info: ActorPricingInfoModel) -> bo for event_name, event_pricing in actor_charge_events.items() ) + async def ensure_current_charged_event_counts(self) -> None: + """Take the run's current charged event counts before the first charge or budgeted push 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 run is read once, as late as possible, which gives the + platform time to record the last charges of that container. Until then, the counts reported by the + synchronous methods are the ones the run started with. Called under the charge lock by `charge` and + `charge_lock_if_charging`; not meant to be called manually. + """ + if not self._charged_event_counts_need_refresh: + return + self._charged_event_counts_need_refresh = False + + initial_counts = {event_name: item.charge_count for event_name, item in self._charging_state.items()} + current_counts = await self._update_charged_event_counts(initial_counts) + + for event_name, count in current_counts.items(): + item = self._charging_state.get(event_name) + if item is not None and item.charge_count == count: + continue + 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, + ) + async def _update_charged_event_counts(self, initial_counts: dict[str, int]) -> dict[str, int]: - """Merge the counts recorded when the run started with the run's current counts, taking the higher of each. + """Merge the given counts with the run's current counts, taking the higher of each. The run is read through the API. When that fails, the initial counts are kept and a warning is logged, since - a run that cannot refresh its counts is still better off running than failing to start. + a run that cannot refresh its counts is still better off running than failing. """ if self._actor_run_id is None: raise RuntimeError('Actor run ID not found even though the Actor is running on Apify') diff --git a/tests/unit/actor/test_charging_manager.py b/tests/unit/actor/test_charging_manager.py index 77e337fce..e487dabb9 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 @@ -798,22 +798,38 @@ def _make_platform_ppe_config(**kwargs: Any) -> Configuration: 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.""" - stats = MagicMock(reboot_count=1, migration_count=0, restart_count=0, resurrect_count=0) - return MagicMock(charged_event_counts=charged_event_counts, stats=stats) + return MagicMock(charged_event_counts=charged_event_counts) -async def test_platform_start_loads_charged_event_counts_from_run(mock_client: MagicMock) -> None: - """A platform run takes the run's current counts where they exceed the env var's, as after a restart.""" +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.assert_called_with('test-run-id') + 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') == 3 + assert cm.get_charged_event_count('search') == 4 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 + 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_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( @@ -827,10 +843,11 @@ async def test_platform_start_loads_charged_event_counts_from_run(mock_client: M ], ) 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 makes no API call at start.""" + """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() @@ -841,18 +858,21 @@ async def test_run_is_not_fetched_when_the_sdk_cannot_charge(mock_client: MagicM pytest.param(None, None, id='run_not_found'), ], ) -async def test_start_continues_when_run_cannot_be_fetched( +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 refresh is logged and the env var's counts are used.""" + """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: - assert cm.get_charged_event_count('search') == 1 + 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 'continuing with the counts recorded when the run started' in caplog.text + assert caplog.text.count('continuing with the counts recorded when the run started') == 1 From 0c4769c9b7a4e69e83071e62d9a3a0159deb9372 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 14:02:04 +0000 Subject: [PATCH 06/18] test: cover concurrent first charges sharing one load of the run Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- src/apify/_charging.py | 2 +- tests/unit/actor/test_charging_manager.py | 23 +++++++++++++++++++++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index 0ac251c54..5c8e257be 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -590,7 +590,7 @@ async def ensure_current_charged_event_counts(self) -> None: they miss what an earlier container charged. The run is read once, as late as possible, which gives the platform time to record the last charges of that container. Until then, the counts reported by the synchronous methods are the ones the run started with. Called under the charge lock by `charge` and - `charge_lock_if_charging`; not meant to be called manually. + `charge_lock_if_charging`, so concurrent charges queue behind the one read; not meant to be called manually. """ if not self._charged_event_counts_need_refresh: return diff --git a/tests/unit/actor/test_charging_manager.py b/tests/unit/actor/test_charging_manager.py index e487dabb9..acde67423 100644 --- a/tests/unit/actor/test_charging_manager.py +++ b/tests/unit/actor/test_charging_manager.py @@ -822,6 +822,29 @@ async def test_first_charge_loads_charged_event_counts_from_run(mock_client: Mag 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}) From f4d38d275cee75efdc4b06b346dd5ba30d724876 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 14:22:42 +0000 Subject: [PATCH 07/18] refactor: record in Actor init where the charged event counts came from `_fetch_pricing_info` reports whether the counts were read from the run, and `__aenter__` records it, so the loader has no side effect on the manager. The chargeable-run check uses the pricing already loaded. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- src/apify/_charging.py | 35 +++++++++++++++++------------------ 1 file changed, 17 insertions(+), 18 deletions(-) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index 5c8e257be..6f0c927d8 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -223,8 +223,9 @@ def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> No self._charging_state: dict[str, ChargingStateItem] = {} self._pricing_info: dict[str, PricingInfoItem] = {} self._tier_priced_events: set[str] = set() - self._charged_event_counts_need_refresh = False - """Whether the run's current charged event counts are still to be taken before the first charge.""" + self._charged_event_counts_refreshed_from_api = False + """Whether the charged event counts were read from the run through the API, rather than taken from the env + var the platform filled when the run started. Set once the read was attempted, whether or not it succeeded.""" self._idempotent_charges: TTLCache[str, IdempotentChargeItem] = TTLCache(maxsize=float('inf'), ttl=180) """Charges made with an idempotency key, by key. Keys expire after three minutes, the same as the platform's own idempotency records.""" @@ -250,6 +251,10 @@ async def __aenter__(self) -> None: pricing_info = pricing_data['pricing_info'] charged_event_counts = pricing_data['charged_event_counts'] max_total_charge_usd = pricing_data['max_total_charge_usd'] + # The env var holds the counts recorded when the run started, which after a reboot, migration or restart miss + # what the run charged before. A restart cannot be told apart from a first start, so the run's current counts + # are read before the first charge, unless they were just read from the run already. + self._charged_event_counts_refreshed_from_api = pricing_data['charged_event_counts_from_api'] # Set pricing model if self._configuration.test_pay_per_event: @@ -523,16 +528,10 @@ async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: """Fetch pricing information from environment variables or API.""" # Check if pricing info is available via environment variables if self._configuration.actor_pricing_info is not None and self._configuration.charged_event_counts is not None: - # The platform fills the env var with the counts recorded when the run started. A container started - # after a reboot, migration or restart would miss what the run charged before it, and a restart cannot - # be told apart from a first start, so the run's current counts are taken before the first charge. - self._charged_event_counts_need_refresh = self._is_chargeable_platform_run( - self._configuration.actor_pricing_info - ) - return _FetchedPricingInfoDict( pricing_info=self._configuration.actor_pricing_info, charged_event_counts=self._configuration.charged_event_counts, + charged_event_counts_from_api=False, max_total_charge_usd=( self._configuration.max_total_charge_usd if self._configuration.max_total_charge_usd is not None @@ -554,6 +553,7 @@ async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: return _FetchedPricingInfoDict( pricing_info=run.pricing_info, charged_event_counts=run.charged_event_counts or {}, + charged_event_counts_from_api=True, max_total_charge_usd=Decimal(str(max_charge)) if max_charge is not None else Decimal('inf'), ) @@ -561,6 +561,7 @@ async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: return _FetchedPricingInfoDict( pricing_info=None, charged_event_counts={}, + charged_event_counts_from_api=False, max_total_charge_usd=( self._configuration.max_total_charge_usd if self._configuration.max_total_charge_usd is not None @@ -568,20 +569,16 @@ async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: ), ) - def _is_chargeable_platform_run(self, pricing_info: ActorPricingInfoModel) -> bool: + def _is_chargeable_platform_run(self) -> bool: """Whether this is a platform run the SDK charges events for. A run is chargeable when its pricing has at least one priced event other than the start event, which only the platform charges. For any other run the charged event counts cannot have moved since the run started. """ - if not self._is_at_home or not isinstance(pricing_info, PayPerEventActorPricingInfo): + if not self._is_at_home or self._pricing_model != 'PAY_PER_EVENT': return False - actor_charge_events = pricing_info.pricing_per_event.actor_charge_events or {} - return any( - event_name != ACTOR_START_EVENT and event_pricing.event_price_usd is not None - for event_name, event_pricing in actor_charge_events.items() - ) + return any(event_name != ACTOR_START_EVENT for event_name in self._pricing_info) async def ensure_current_charged_event_counts(self) -> None: """Take the run's current charged event counts before the first charge or budgeted push of this container. @@ -592,9 +589,9 @@ async def ensure_current_charged_event_counts(self) -> None: synchronous methods are the ones the run started with. Called under the charge lock by `charge` and `charge_lock_if_charging`, so concurrent charges queue behind the one read; not meant to be called manually. """ - if not self._charged_event_counts_need_refresh: + if self._charged_event_counts_refreshed_from_api or not self._is_chargeable_platform_run(): return - self._charged_event_counts_need_refresh = False + self._charged_event_counts_refreshed_from_api = True initial_counts = {event_name: item.charge_count for event_name, item in self._charging_state.items()} current_counts = await self._update_charged_event_counts(initial_counts) @@ -680,4 +677,6 @@ class IdempotentChargeItem: class _FetchedPricingInfoDict(TypedDict): pricing_info: ActorPricingInfoModel | None charged_event_counts: dict[str, int] + charged_event_counts_from_api: bool + """Whether the counts are the run's current ones, read through the API, rather than the env var's.""" max_total_charge_usd: Decimal From 2e3005456a1f7b76ce676a99d9525c236d420675 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 14:28:28 +0000 Subject: [PATCH 08/18] refactor: derive the charged event counts source from the configuration Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- src/apify/_charging.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index 6f0c927d8..0b55bc9c5 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -254,7 +254,7 @@ async def __aenter__(self) -> None: # The env var holds the counts recorded when the run started, which after a reboot, migration or restart miss # what the run charged before. A restart cannot be told apart from a first start, so the run's current counts # are read before the first charge, unless they were just read from the run already. - self._charged_event_counts_refreshed_from_api = pricing_data['charged_event_counts_from_api'] + self._charged_event_counts_refreshed_from_api = self._is_at_home and not self._has_pricing_env_vars() # Set pricing model if self._configuration.test_pay_per_event: @@ -527,11 +527,10 @@ def compute_chargeable(self) -> dict[str, int | None]: async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: """Fetch pricing information from environment variables or API.""" # Check if pricing info is available via environment variables - if self._configuration.actor_pricing_info is not None and self._configuration.charged_event_counts is not None: + if self._has_pricing_env_vars(): return _FetchedPricingInfoDict( pricing_info=self._configuration.actor_pricing_info, - charged_event_counts=self._configuration.charged_event_counts, - charged_event_counts_from_api=False, + charged_event_counts=self._configuration.charged_event_counts or {}, max_total_charge_usd=( self._configuration.max_total_charge_usd if self._configuration.max_total_charge_usd is not None @@ -553,7 +552,6 @@ async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: return _FetchedPricingInfoDict( pricing_info=run.pricing_info, charged_event_counts=run.charged_event_counts or {}, - charged_event_counts_from_api=True, max_total_charge_usd=Decimal(str(max_charge)) if max_charge is not None else Decimal('inf'), ) @@ -561,7 +559,6 @@ async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: return _FetchedPricingInfoDict( pricing_info=None, charged_event_counts={}, - charged_event_counts_from_api=False, max_total_charge_usd=( self._configuration.max_total_charge_usd if self._configuration.max_total_charge_usd is not None @@ -569,6 +566,12 @@ async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: ), ) + def _has_pricing_env_vars(self) -> bool: + """Whether the platform provided the pricing info and the charged event counts through the env vars.""" + return ( + self._configuration.actor_pricing_info is not None and self._configuration.charged_event_counts is not None + ) + def _is_chargeable_platform_run(self) -> bool: """Whether this is a platform run the SDK charges events for. @@ -677,6 +680,4 @@ class IdempotentChargeItem: class _FetchedPricingInfoDict(TypedDict): pricing_info: ActorPricingInfoModel | None charged_event_counts: dict[str, int] - charged_event_counts_from_api: bool - """Whether the counts are the run's current ones, read through the API, rather than the env var's.""" max_total_charge_usd: Decimal From 805e5fd47c40c0c55f9575362c86590c34846c6f Mon Sep 17 00:00:00 2001 From: Josef Prochazka Date: Fri, 30 Oct 2026 15:51:28 +0100 Subject: [PATCH 09/18] Simplify --- src/apify/_charging.py | 49 ++++++++++++------------------------------ 1 file changed, 14 insertions(+), 35 deletions(-) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index 0b55bc9c5..a817c8bc0 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -341,7 +341,7 @@ async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str ) async with self.charge_lock(): - await self.ensure_current_charged_event_counts() + await self._try_to_refresh_from_api() # A repeat the platform would discard is resolved here too - counting it would inflate the charging state # and make the run hit `max_total_charge_usd` early. The key alone decides, as on the platform. @@ -583,41 +583,17 @@ def _is_chargeable_platform_run(self) -> bool: return any(event_name != ACTOR_START_EVENT for event_name in self._pricing_info) - async def ensure_current_charged_event_counts(self) -> None: - """Take the run's current charged event counts before the first charge or budgeted push of this container. + 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 run is read once, as late as possible, which gives the - platform time to record the last charges of that container. Until then, the counts reported by the - synchronous methods are the ones the run started with. Called under the charge lock by `charge` and - `charge_lock_if_charging`, so concurrent charges queue behind the one read; not meant to be called manually. + 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. """ - if self._charged_event_counts_refreshed_from_api or not self._is_chargeable_platform_run(): + if self._charged_event_counts_refreshed_from_api or not self._is_chargeable_platform_run() and self._actor_run_id is not None: return self._charged_event_counts_refreshed_from_api = True - initial_counts = {event_name: item.charge_count for event_name, item in self._charging_state.items()} - current_counts = await self._update_charged_event_counts(initial_counts) - - for event_name, count in current_counts.items(): - item = self._charging_state.get(event_name) - if item is not None and item.charge_count == count: - continue - 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, - ) - - async def _update_charged_event_counts(self, initial_counts: dict[str, int]) -> dict[str, int]: - """Merge the given counts with the run's current counts, taking the higher of each. - - The run is read through the API. When that fails, the initial counts are kept and a warning is logged, since - a run that cannot refresh its counts is still better off running than failing. - """ - if self._actor_run_id is None: - raise RuntimeError('Actor run ID not found even though the Actor is running on Apify') - try: run = await self._client.run(self._actor_run_id).get() except Exception: @@ -626,18 +602,21 @@ async def _update_charged_event_counts(self, initial_counts: dict[str, int]) -> 'when the run started.', exc_info=True, ) - return initial_counts + 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 initial_counts + return - for event_name, count in (run.charged_event_counts or {}).items(): - initial_counts[event_name] = max(initial_counts.get(event_name, 0), count) - return initial_counts + 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 + self._charging_state[event_name] = ChargingStateItem( + charge_count=max(local_count, api_count), + total_charged_amount=local_count * self._pricing_info.get(event_name, PricingInfoItem(Decimal(), title='')).price, + ) def _compute_charged_count(self, event_name: str, count: int) -> int: """Compute how many of `count` events to charge, given the remaining budget.""" From 01b9d5ebc0b888c842f51f15cb10d4fa0998195a Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 14:53:57 +0000 Subject: [PATCH 10/18] fix: wire the renamed refresh into the charge lock and price the merged counts Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- src/apify/_charging.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index a817c8bc0..0593bf280 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -65,7 +65,7 @@ async def charge_lock_if_charging() -> AsyncGenerator[None]: return async with charging_manager.charge_lock(): - await charging_manager.ensure_current_charged_event_counts() + await charging_manager._try_to_refresh_from_api() # noqa: SLF001 yield @@ -590,12 +590,13 @@ async def _try_to_refresh_from_api(self) -> None: 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. """ - if self._charged_event_counts_refreshed_from_api or not self._is_chargeable_platform_run() and self._actor_run_id is not None: + run_id = self._actor_run_id + if self._charged_event_counts_refreshed_from_api or not self._is_chargeable_platform_run() or run_id is None: return self._charged_event_counts_refreshed_from_api = True try: - run = await self._client.run(self._actor_run_id).get() + run = await self._client.run(run_id).get() except Exception: logger.warning( 'Failed to fetch the current charged event counts of the run, continuing with the counts recorded ' @@ -613,9 +614,10 @@ async def _try_to_refresh_from_api(self) -> None: 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=max(local_count, api_count), - total_charged_amount=local_count * self._pricing_info.get(event_name, PricingInfoItem(Decimal(), title='')).price, + charge_count=count, + total_charged_amount=count * self._get_event_price(event_name), ) def _compute_charged_count(self, event_name: str, count: int) -> int: From 88c9eb676a9397cbe81430ae45bf0ef5e2bea4a9 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 14:59:32 +0000 Subject: [PATCH 11/18] docs: drop the restart paragraph from the pay-per-event page Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- docs/02_concepts/11_pay_per_event.mdx | 2 -- 1 file changed, 2 deletions(-) diff --git a/docs/02_concepts/11_pay_per_event.mdx b/docs/02_concepts/11_pay_per_event.mdx index 8b4f7504e..6137c2978 100644 --- a/docs/02_concepts/11_pay_per_event.mdx +++ b/docs/02_concepts/11_pay_per_event.mdx @@ -78,8 +78,6 @@ The user spending limit for the run is available in your Actor code as the `ACTO When the charge limit is reached, `Actor.charge` stops charging and `Actor.push_data` stops pushing data. The platform then aborts the run automatically. However, the run keeps consuming platform resources for a short time before it stops. For details, see [Handle graceful shutdown](https://docs.apify.com/platform/actors/publishing/monetize/pay-per-event#handle-graceful-shutdown). -The charged event counts survive a reboot, a migration and a restart of the run. Before the first charge or push of a container, the SDK reads the current counts from the run through the API, so a restarted run continues from what was already charged. Until then, the counts reported by the `ChargingManager` are those recorded when the run started. Events charged in the last few seconds before the restart can be missing from the counts the platform had recorded by then, so the restarted run may count slightly less than the platform. Keep that in mind when you compute the remaining budget from the charged counts yourself. - ## Test monetization locally Before releasing your monetization code to the public, test it locally. To make your Actor work in pay-per-event mode, pass it the `ACTOR_TEST_PAY_PER_EVENT` environment variable: From c3ee68a0410d2ce29d34385b526c22983bf22b8c Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 15:10:34 +0000 Subject: [PATCH 12/18] refactor: refresh the charged event counts from the charge lock The charge lock brings the counts up to date the first time it is taken, so neither `charge` nor the push path needs to know about the refresh. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- src/apify/_charging.py | 40 +++++++++++++++++++++++++--------------- 1 file changed, 25 insertions(+), 15 deletions(-) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index 0593bf280..9a175394a 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -22,7 +22,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 @@ -65,7 +65,6 @@ async def charge_lock_if_charging() -> AsyncGenerator[None]: return async with charging_manager.charge_lock(): - await charging_manager._try_to_refresh_from_api() # noqa: SLF001 yield @@ -223,9 +222,6 @@ def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> No self._charging_state: dict[str, ChargingStateItem] = {} self._pricing_info: dict[str, PricingInfoItem] = {} self._tier_priced_events: set[str] = set() - self._charged_event_counts_refreshed_from_api = False - """Whether the charged event counts were read from the run through the API, rather than taken from the env - var the platform filled when the run started. Set once the read was attempted, whether or not it succeeded.""" self._idempotent_charges: TTLCache[str, IdempotentChargeItem] = TTLCache(maxsize=float('inf'), ttl=180) """Charges made with an idempotency key, by key. Keys expire after three minutes, the same as the platform's own idempotency records.""" @@ -233,7 +229,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.""" @@ -251,10 +247,6 @@ async def __aenter__(self) -> None: pricing_info = pricing_data['pricing_info'] charged_event_counts = pricing_data['charged_event_counts'] max_total_charge_usd = pricing_data['max_total_charge_usd'] - # The env var holds the counts recorded when the run started, which after a reboot, migration or restart miss - # what the run charged before. A restart cannot be told apart from a first start, so the run's current counts - # are read before the first charge, unless they were just read from the run already. - self._charged_event_counts_refreshed_from_api = self._is_at_home and not self._has_pricing_env_vars() # Set pricing model if self._configuration.test_pay_per_event: @@ -341,8 +333,6 @@ async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str ) async with self.charge_lock(): - await self._try_to_refresh_from_api() - # A repeat the platform would discard is resolved here too - counting it would inflate the charging state # and make the run hit `max_total_charge_usd` early. The key alone decides, as on the platform. if idempotency_key is not None and (previous := self._idempotent_charges.get(idempotency_key)) is not None: @@ -588,12 +578,11 @@ async def _try_to_refresh_from_api(self) -> None: 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. + 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 self._charged_event_counts_refreshed_from_api or not self._is_chargeable_platform_run() or run_id is None: + if not self._is_chargeable_platform_run() or run_id is None: return - self._charged_event_counts_refreshed_from_api = True try: run = await self._client.run(run_id).get() @@ -640,6 +629,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 From 62ec983737e0301d9621001c018f7f144c9f9914 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 15:13:31 +0000 Subject: [PATCH 13/18] refactor: inline the pricing env var check again Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- src/apify/_charging.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index 9a175394a..c6a934648 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -517,10 +517,10 @@ def compute_chargeable(self) -> dict[str, int | None]: async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: """Fetch pricing information from environment variables or API.""" # Check if pricing info is available via environment variables - if self._has_pricing_env_vars(): + if self._configuration.actor_pricing_info is not None and self._configuration.charged_event_counts is not None: return _FetchedPricingInfoDict( pricing_info=self._configuration.actor_pricing_info, - charged_event_counts=self._configuration.charged_event_counts or {}, + charged_event_counts=self._configuration.charged_event_counts, max_total_charge_usd=( self._configuration.max_total_charge_usd if self._configuration.max_total_charge_usd is not None @@ -556,12 +556,6 @@ async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: ), ) - def _has_pricing_env_vars(self) -> bool: - """Whether the platform provided the pricing info and the charged event counts through the env vars.""" - return ( - self._configuration.actor_pricing_info is not None and self._configuration.charged_event_counts is not None - ) - def _is_chargeable_platform_run(self) -> bool: """Whether this is a platform run the SDK charges events for. From 9311314fe84096f9feda76a7f29f6dfa54ace7a4 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 15:14:49 +0000 Subject: [PATCH 14/18] refactor: restore the charging manager context variable type Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- src/apify/_charging.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index c6a934648..94dca5b6c 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -31,9 +31,7 @@ logger = getLogger(__name__) -charging_manager_ctx: ContextVar[ChargingManagerImplementation | None] = ContextVar( - 'charging_manager_ctx', default=None -) +charging_manager_ctx: ContextVar[ChargingManager | None] = ContextVar('charging_manager_ctx', default=None) """Holds the current `ChargingManager` instance, if any. Allows PPE-aware dataset clients to access the charging manager without needing to pass it explicitly. From 896e7914454c33a65f4affd4ece9ce8bda0d5e22 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 15:17:45 +0000 Subject: [PATCH 15/18] refactor: drop an obvious docstring Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- src/apify/_charging.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index 94dca5b6c..b1d37ad69 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -555,11 +555,6 @@ async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict: ) def _is_chargeable_platform_run(self) -> bool: - """Whether this is a platform run the SDK charges events for. - - A run is chargeable when its pricing has at least one priced event other than the start event, which only - the platform charges. For any other run the charged event counts cannot have moved since the run started. - """ if not self._is_at_home or self._pricing_model != 'PAY_PER_EVENT': return False From 8563a5bcd04d09e007db1bccefa6cbb28091ab6f Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 15:26:21 +0000 Subject: [PATCH 16/18] fix: warn when the read of the current charged event counts is cancelled Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- src/apify/_charging.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index b1d37ad69..2cbd9d84a 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 @@ -573,6 +574,13 @@ async def _try_to_refresh_from_api(self) -> None: 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 ' From 9080b391f5cae6d3a18ff10b553de378f613701a Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 15:34:06 +0000 Subject: [PATCH 17/18] test: cover charged event counts surviving a reboot on the platform Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017ocuAKQTqReMDTjXCUTnKP --- tests/e2e/test_actor_charge.py | 96 ++++++++++++++++++++++++++++++++++ 1 file changed, 96 insertions(+) diff --git a/tests/e2e/test_actor_charge.py b/tests/e2e/test_actor_charge.py index dc09602a9..05e62029b 100644 --- a/tests/e2e/test_actor_charge.py +++ b/tests/e2e/test_actor_charge.py @@ -226,3 +226,99 @@ 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 starts from the counts the platform recorded when the run started. The first + # charge loads the run's current counts, so the view recorded afterwards includes the earlier charges. + charge_result = await Actor.charge('row-charged') + 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': charge_result.chargeable_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 continues from the events charged before the reboot, instead of a fresh budget.""" + 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': 4}, + timeout=30, + poll_interval=1, + ) + assert run.charged_event_counts == {'row-charged': 4} + + 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 were charged before the reboot and one after it. The budget of $0.05 at $0.005 per event leaves + # room for six more. + assert record['value'] == { + 'chargedCount': 4, + 'totalChargedUsd': '0.020', + 'chargeableWithinLimit': 6, + } From 47c292630185dbdbdcb6e30254b09103f0f65f81 Mon Sep 17 00:00:00 2001 From: Claude Date: Fri, 30 Oct 2026 15:38:05 +0000 Subject: [PATCH 18/18] Revert "test: cover charged event counts surviving a reboot on the platform" This reverts commit 9080b391f5cae6d3a18ff10b553de378f613701a. --- tests/e2e/test_actor_charge.py | 96 ---------------------------------- 1 file changed, 96 deletions(-) diff --git a/tests/e2e/test_actor_charge.py b/tests/e2e/test_actor_charge.py index 05e62029b..dc09602a9 100644 --- a/tests/e2e/test_actor_charge.py +++ b/tests/e2e/test_actor_charge.py @@ -226,99 +226,3 @@ 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 starts from the counts the platform recorded when the run started. The first - # charge loads the run's current counts, so the view recorded afterwards includes the earlier charges. - charge_result = await Actor.charge('row-charged') - 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': charge_result.chargeable_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 continues from the events charged before the reboot, instead of a fresh budget.""" - 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': 4}, - timeout=30, - poll_interval=1, - ) - assert run.charged_event_counts == {'row-charged': 4} - - 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 were charged before the reboot and one after it. The budget of $0.05 at $0.005 per event leaves - # room for six more. - assert record['value'] == { - 'chargedCount': 4, - 'totalChargedUsd': '0.020', - 'chargeableWithinLimit': 6, - }