diff --git a/src/apify/_actor.py b/src/apify/_actor.py index 57b2cdf6..5433b797 100644 --- a/src/apify/_actor.py +++ b/src/apify/_actor.py @@ -788,7 +788,7 @@ def get_charging_manager(self) -> ChargingManager: return self._charging_manager_implementation @_ensure_context - async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult: + async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str | None = None) -> ChargeResult: """Charge for a specified number of events - sub-operations of the Actor. This is relevant only for the pay-per-event pricing model. @@ -796,10 +796,13 @@ async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult: Args: event_name: Name of the event to be charged for. count: Number of events to charge for. + idempotency_key: A unique key preventing a retried operation from being charged for twice. A repeat + under the same key is not sent to the API and reports the `charged_count` of the original call. + A key belongs to a single event and is only remembered for the lifetime of the run. """ # charging_manager.charge() acquires charge_lock internally. charging_manager = self.get_charging_manager() - return await charging_manager.charge(event_name, count=count) + return await charging_manager.charge(event_name, count=count, idempotency_key=idempotency_key) @overload def on( diff --git a/src/apify/_charging.py b/src/apify/_charging.py index 6c8adea4..e9b3f244 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -210,7 +210,7 @@ class ChargingManager(Protocol): charge_lock: ReentrantLock """Lock to synchronize charge operations. Prevents race conditions between `charge` and `push_data` calls.""" - async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult: + async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str | None = None) -> ChargeResult: """Charge for a specified number of events - sub-operations of the Actor. This is relevant only for the pay-per-event pricing model. @@ -218,6 +218,9 @@ async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult: Args: event_name: Name of the event to be charged for. count: Number of events to charge for. + idempotency_key: A unique key preventing a retried operation from being charged for twice. A repeat + under the same key is not sent to the API and reports the `charged_count` of the original call. + A key belongs to a single event and is only remembered for the lifetime of the run. """ def calculate_total_charged_amount(self) -> Decimal: @@ -329,6 +332,7 @@ 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._idempotent_charges: dict[str, IdempotentChargeItem] = {} self._not_ppe_warning_printed = False self.active = False @@ -412,7 +416,10 @@ async def __aexit__( self.active = False @_ensure_context - async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult: + async def charge(self, event_name: str, *, count: int = 1, idempotency_key: str | None = None) -> ChargeResult: + if idempotency_key == '': + raise ValueError('idempotency_key must not be an empty string') + # For runs that do not use the pay-per-event pricing model, just print a warning and return if self._pricing_model != 'PAY_PER_EVENT': if not self._not_ppe_warning_printed: @@ -435,6 +442,21 @@ async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult: ) async with self.charge_lock(): + # The platform discards a charge repeated under a key it has already seen, so repeating it here would + # inflate the local charging state and make the run hit `max_total_charge_usd` sooner than it should. + if idempotency_key is not None and (previous := self._idempotent_charges.get(idempotency_key)): + if previous.event_name != event_name: + raise ValueError( + f"Idempotency key '{idempotency_key}' was already used to charge for event " + f"'{previous.event_name}', so it cannot be reused for event '{event_name}'." + ) + + return ChargeResult( + event_charge_limit_reached=self.is_event_charge_limit_reached(event_name), + charged_count=previous.charged_count, + chargeable_within_limit=self.compute_chargeable(), + ) + # Determine the maximum amount of events that can be charged within the budget max_chargeable = self.calculate_max_event_charge_count_within_limit(event_name) charged_count = min(count, max_chargeable if max_chargeable is not None else count) @@ -455,10 +477,7 @@ async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult: ), ) - # Update the charging state - 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 + charge_sent = False # If running on the platform, call the charge endpoint if self._is_at_home: @@ -470,7 +489,12 @@ async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult: # the platform handles them automatically based on dataset writes. pass elif event_name in self._pricing_info: - await self._client.run(self._actor_run_id).charge(event_name, count=charged_count) + await self._client.run(self._actor_run_id).charge( + event_name, + count=charged_count, + idempotency_key=idempotency_key, + ) + charge_sent = True logger.debug(f"Charged {charged_count} occurrence(s) of event '{event_name}'.") elif event_name in self._tier_priced_events: logger.warning( @@ -479,6 +503,19 @@ async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult: else: logger.warning(f"Attempting to charge for an unknown event '{event_name}'") + # Update the charging state + 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 + + # Only remember a key that stands for a charge the platform actually received. Off the platform there + # is no request at all and the registry is the only thing providing deduplication. + if idempotency_key is not None and (not self._is_at_home or charge_sent): + self._idempotent_charges[idempotency_key] = IdempotentChargeItem( + event_name=event_name, + charged_count=charged_count, + ) + # Log the charged operation (if enabled) if self._charging_log_dataset: await self._charging_log_dataset.push_data( @@ -487,6 +524,7 @@ async def charge(self, event_name: str, *, count: int = 1) -> ChargeResult: 'event_title': pricing_info.title, 'event_price_usd': float(round(pricing_info.price, 3)), 'charged_count': charged_count, + 'idempotency_key': idempotency_key, 'timestamp': datetime.now(UTC).isoformat(), } ) @@ -636,6 +674,12 @@ class PricingInfoItem: title: str +@dataclass(frozen=True) +class IdempotentChargeItem: + event_name: str + charged_count: int + + class _FetchedPricingInfoDict(TypedDict): pricing_info: ActorPricingInfoModel | None charged_event_counts: dict[str, int] diff --git a/tests/unit/actor/test_actor_charge.py b/tests/unit/actor/test_actor_charge.py index add57eeb..848e03b9 100644 --- a/tests/unit/actor/test_actor_charge.py +++ b/tests/unit/actor/test_actor_charge.py @@ -37,7 +37,7 @@ async def setup_mocked_charging( setup.charging_mgr._pricing_info['event'] = PricingInfoItem(Decimal('1.0'), 'Event') result = await Actor.charge('event', count=1) - setup.mock_charge.assert_called_once_with('event', count=1) + setup.mock_charge.assert_called_once_with('event', count=1, idempotency_key=None) """ # Mock the ApifyClientAsync mock_client = Mock() @@ -82,7 +82,7 @@ async def test_actor_charge_push_data_with_no_remaining_budget() -> None: result1 = await Actor.charge('some-event', count=1) # Costs $1, leaving $0.5 # Verify the first charge call was made correctly - setup.mock_charge.assert_called_once_with('some-event', count=1) + setup.mock_charge.assert_called_once_with('some-event', count=1, idempotency_key=None) setup.mock_charge.reset_mock() assert result1.charged_count == 1 @@ -117,10 +117,27 @@ async def test_actor_charge_api_call_verification() -> None: # Call charge with count=1 - this SHOULD call the API result2 = await Actor.charge('test-event', count=1) - setup.mock_charge.assert_called_once_with('test-event', count=1) + setup.mock_charge.assert_called_once_with('test-event', count=1, idempotency_key=None) assert result2.charged_count == 1 +async def test_actor_charge_forwards_idempotency_key() -> None: + """Verify that Actor.charge passes the idempotency key down to the API and deduplicates repeats.""" + async with setup_mocked_charging( + Configuration(max_total_charge_usd=Decimal('10.0'), test_pay_per_event=True), {'test-event': Decimal('1.0')} + ) as setup: + result1 = await Actor.charge('test-event', count=1, idempotency_key='key-1') + setup.mock_charge.assert_called_once_with('test-event', count=1, idempotency_key='key-1') + assert result1.charged_count == 1 + + setup.mock_charge.reset_mock() + + result2 = await Actor.charge('test-event', count=1, idempotency_key='key-1') + setup.mock_charge.assert_not_called() + assert result2.charged_count == 1 + assert setup.charging_mgr.get_charged_event_count('test-event') == 1 + + async def test_max_event_charge_count_within_limit_tolerates_overdraw() -> None: """Test that calculate_max_event_charge_count_within_limit does not return nonsensical (e.g., negative) values when the total number of charged events overdraws the max_total_charge_usd limit.""" diff --git a/tests/unit/actor/test_charging_manager.py b/tests/unit/actor/test_charging_manager.py index dcbdeb2b..3795db45 100644 --- a/tests/unit/actor/test_charging_manager.py +++ b/tests/unit/actor/test_charging_manager.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio from decimal import Decimal from typing import Any from unittest.mock import AsyncMock, MagicMock @@ -466,3 +467,132 @@ async def test_compute_chargeable_updates_after_charge(mock_client: MagicMock) - # $6.00 remaining: search=$1.00 → 6, scrape=$2.00 → 3 assert chargeable['search'] == 6 assert chargeable['scrape'] == 3 + + +async def test_charge_forwards_idempotency_key_to_client(mock_client: MagicMock) -> None: + """Test that the idempotency key is passed through to the API client.""" + pricing_info = _make_ppe_pricing_info({'search': Decimal('1.00')}) + config = _make_config( + is_at_home=True, + actor_run_id='test-run-id', + actor_pricing_info=pricing_info, + charged_event_counts={}, + max_total_charge_usd=Decimal('10.00'), + ) + cm = ChargingManagerImplementation(config, mock_client) + async with cm: + await cm.charge('search', count=2, idempotency_key='key-1') + mock_client.run.return_value.charge.assert_awaited_once_with('search', count=2, idempotency_key='key-1') + + # Without a key the client gets None and generates a unique one per call, so nothing is deduplicated. + mock_client.run.return_value.charge.reset_mock() + await cm.charge('search', count=1) + mock_client.run.return_value.charge.assert_awaited_once_with('search', count=1, idempotency_key=None) + + +async def test_charge_deduplicates_repeated_idempotency_key(mock_client: MagicMock) -> None: + """Test that a repeated key skips both the API call and the local charging state update.""" + pricing_info = _make_ppe_pricing_info({'search': Decimal('1.00')}) + config = _make_config( + is_at_home=True, + actor_run_id='test-run-id', + actor_pricing_info=pricing_info, + charged_event_counts={}, + max_total_charge_usd=Decimal('10.00'), + ) + cm = ChargingManagerImplementation(config, mock_client) + async with cm: + first = await cm.charge('search', count=2, idempotency_key='key-1') + assert first.charged_count == 2 + + # The repeat reports what was charged under this key originally, whatever count it asks for. + repeated = await cm.charge('search', count=1, idempotency_key='key-1') + assert repeated.charged_count == 2 + + assert cm.get_charged_event_count('search') == 2 + assert mock_client.run.return_value.charge.await_count == 1 + + # A distinct key is charged as usual. + await cm.charge('search', count=1, idempotency_key='key-2') + assert cm.get_charged_event_count('search') == 3 + assert mock_client.run.return_value.charge.await_count == 2 + + +async def test_charge_does_not_register_key_for_uncharged_event(mock_client: MagicMock) -> None: + """Test that an event that never reached the API does not consume the idempotency key.""" + pricing_info = PayPerEventActorPricingInfo.model_validate( + { + 'pricingModel': 'PAY_PER_EVENT', + 'pricingPerEvent': { + 'actorChargeEvents': { + 'search': {'eventPriceUsd': 1.00, 'eventTitle': 'Search event'}, + 'tiered': {'eventTieredPricingUsd': {}, 'eventTitle': 'Tiered event'}, + } + }, + } + ) + config = _make_config( + is_at_home=True, + actor_run_id='test-run-id', + actor_pricing_info=pricing_info, + charged_event_counts={}, + max_total_charge_usd=Decimal('10.00'), + ) + cm = ChargingManagerImplementation(config, mock_client) + async with cm: + # Neither a tier-priced nor an unknown event is chargeable via the API. + await cm.charge('tiered', count=1, idempotency_key='key-1') + await cm.charge('typo-event', count=1, idempotency_key='key-2') + mock_client.run.return_value.charge.assert_not_awaited() + + # Both keys are still free, so a corrected charge under either one reaches the platform. + assert (await cm.charge('search', count=1, idempotency_key='key-1')).charged_count == 1 + assert (await cm.charge('search', count=1, idempotency_key='key-2')).charged_count == 1 + assert mock_client.run.return_value.charge.await_count == 2 + + +async def test_charge_rejects_invalid_idempotency_key(mock_client: MagicMock) -> None: + """Test that an empty key and a key reused for another event are both refused.""" + pricing_info = _make_ppe_pricing_info({'search': Decimal('1.00'), 'scrape': Decimal('2.00')}) + config = _make_config( + is_at_home=True, + actor_run_id='test-run-id', + actor_pricing_info=pricing_info, + charged_event_counts={}, + max_total_charge_usd=Decimal('10.00'), + ) + cm = ChargingManagerImplementation(config, mock_client) + async with cm: + with pytest.raises(ValueError, match='must not be an empty string'): + await cm.charge('search', count=1, idempotency_key='') + + await cm.charge('search', count=1, idempotency_key='key-1') + + with pytest.raises(ValueError, match='cannot be reused for event'): + await cm.charge('scrape', count=1, idempotency_key='key-1') + + # The refused charge must leave no trace behind. + assert cm.get_charged_event_count('scrape') == 0 + assert cm.calculate_total_charged_amount() == Decimal('1.00') + assert mock_client.run.return_value.charge.await_count == 1 + + +async def test_concurrent_charges_under_one_key_charge_once(mock_client: MagicMock) -> None: + """Test that concurrent charges under one key result in a single charge.""" + pricing_info = _make_ppe_pricing_info({'search': Decimal('1.00')}) + config = _make_config( + is_at_home=True, + actor_run_id='test-run-id', + actor_pricing_info=pricing_info, + charged_event_counts={}, + max_total_charge_usd=Decimal('10.00'), + ) + cm = ChargingManagerImplementation(config, mock_client) + async with cm: + results = await asyncio.gather( + *(cm.charge('search', count=1, idempotency_key='key-1') for _ in range(5)), + ) + + assert [result.charged_count for result in results] == [1] * 5 + assert cm.get_charged_event_count('search') == 1 + assert mock_client.run.return_value.charge.await_count == 1