diff --git a/src/apify/_actor.py b/src/apify/_actor.py index 077759f4..47ef1fde 100644 --- a/src/apify/_actor.py +++ b/src/apify/_actor.py @@ -26,7 +26,13 @@ EventSystemInfoData, ) -from apify._charging import DEFAULT_DATASET_ITEM_EVENT, ChargeResult, ChargingManager, ChargingManagerImplementation +from apify._charging import ( + DEFAULT_DATASET_ITEM_EVENT, + ChargeResult, + ChargingManager, + ChargingManagerImplementation, + charge_lock_if_charging, +) from apify._configuration import Configuration from apify._consts import EVENT_LISTENERS_TIMEOUT, EXIT_CODE_ERROR_USER_FUNCTION_THREW, ActorEnvVars, ApifyEnvVars from apify._crypto import decrypt_input_secrets, load_private_key @@ -687,9 +693,10 @@ async def push_data(self, data: dict | list[dict], *, charged_event_name: str | dataset = await self.open_dataset() - # Acquire the charge lock to prevent race conditions between concurrent - # push_data calls. We need to hold the lock for the entire push_data + charge sequence. - async with charging_manager.charge_lock(): + # The whole push + charge sequence has to stay under the charge lock, so that a concurrent push cannot + # charge in between the limit reservation below and the charge that acts on it. Runs that charge nothing + # skip the lock and push concurrently. + async with charge_lock_if_charging(): # Synthetic events are handled within dataset.push_data, only get data for `ChargeResult`. if charged_event_name is None: before = charging_manager.get_charged_event_count(DEFAULT_DATASET_ITEM_EVENT) diff --git a/src/apify/_charging.py b/src/apify/_charging.py index 9bbdeb3e..6c8adea4 100644 --- a/src/apify/_charging.py +++ b/src/apify/_charging.py @@ -1,6 +1,7 @@ from __future__ import annotations import math +from contextlib import asynccontextmanager from contextvars import ContextVar from dataclasses import dataclass from datetime import UTC, datetime @@ -23,6 +24,7 @@ from apify.storages import Dataset if TYPE_CHECKING: + from collections.abc import AsyncIterator from types import TracebackType from apify_client import ApifyClientAsync @@ -46,6 +48,23 @@ _ensure_context = ensure_context('active') +@asynccontextmanager +async def charge_lock_if_charging() -> AsyncIterator[None]: + """Acquire the charge lock if a charging manager is active, otherwise proceed without locking. + + The lock keeps a limit reservation and the charge that follows it atomic. Only pay-per-event runs charge + anything, and `charging_manager_ctx` is set exactly for those, so for any other run there is nothing to + serialize and the lock is skipped. + """ + charging_manager = charging_manager_ctx.get() + if charging_manager is None: + yield + return + + async with charging_manager.charge_lock(): + yield + + # These are thin subclasses of the `apify-client` pricing models. The Apify platform serializes Actor # pricing info into the `APIFY_ACTOR_PRICING_INFO` env var (parsed by `Configuration.actor_pricing_info`), # but omits several fields that `apify-client` v3 marks as required (`apifyMarginPercentage`, `createdAt`, diff --git a/src/apify/storage_clients/_apify/_dataset_client.py b/src/apify/storage_clients/_apify/_dataset_client.py index de9634fd..9e9cf02e 100644 --- a/src/apify/storage_clients/_apify/_dataset_client.py +++ b/src/apify/storage_clients/_apify/_dataset_client.py @@ -12,6 +12,7 @@ from crawlee.storage_clients.models import DatasetItemsListPage, DatasetMetadata from ._api_client_creation import create_storage_api_client +from apify._charging import charge_lock_if_charging from apify.storage_clients._ppe_dataset_mixin import DatasetClientPpeMixin if TYPE_CHECKING: @@ -54,7 +55,7 @@ def __init__( """The Apify dataset client for API operations.""" self._lock = lock - """A lock to ensure that only one operation is performed at a time.""" + """A lock serializing destructive operations on the dataset.""" @override async def get_metadata(self) -> DatasetMetadata: @@ -142,7 +143,9 @@ async def payloads_generator(items: Sequence[Mapping[str, JsonSerializable]]) -> for index, item in enumerate(items): yield await self._check_and_serialize(item, index) - async with self._charge_lock(), self._lock: + # Pushing mutates no client state - `push_items` is a stateless API call - so concurrent pushes only need + # the charge lock, which keeps the limit reservation and the charge atomic for pay-per-event runs. + async with charge_lock_if_charging(): items = data if self._is_sequence_of_items(data) else [data] if not items: return diff --git a/src/apify/storage_clients/_file_system/_dataset_client.py b/src/apify/storage_clients/_file_system/_dataset_client.py index b5ab1a43..5480f577 100644 --- a/src/apify/storage_clients/_file_system/_dataset_client.py +++ b/src/apify/storage_clients/_file_system/_dataset_client.py @@ -6,6 +6,7 @@ from crawlee.storage_clients._file_system import FileSystemDatasetClient +from apify._charging import charge_lock_if_charging from apify.storage_clients._ppe_dataset_mixin import DatasetClientPpeMixin if TYPE_CHECKING: @@ -51,7 +52,7 @@ async def open( @override async def push_data(self, data: Sequence[Mapping[str, JsonSerializable]] | Mapping[str, JsonSerializable]) -> None: - async with self._charge_lock(): + async with charge_lock_if_charging(): items = data if self._is_sequence_of_items(data) else [data] limit = self._compute_limit_for_push(len(items)) diff --git a/src/apify/storage_clients/_ppe_dataset_mixin.py b/src/apify/storage_clients/_ppe_dataset_mixin.py index f68361ad..18663dee 100644 --- a/src/apify/storage_clients/_ppe_dataset_mixin.py +++ b/src/apify/storage_clients/_ppe_dataset_mixin.py @@ -1,13 +1,7 @@ from __future__ import annotations -from contextlib import asynccontextmanager -from typing import TYPE_CHECKING - from apify._charging import DEFAULT_DATASET_ITEM_EVENT, charging_manager_ctx -if TYPE_CHECKING: - from collections.abc import AsyncIterator - class DatasetClientPpeMixin: """A mixin for dataset clients to add support for PPE pricing model and tracking synthetic events.""" @@ -29,13 +23,3 @@ async def _charge_for_items(self, count_items: int) -> None: event_name=DEFAULT_DATASET_ITEM_EVENT, count=count_items, ) - - @asynccontextmanager - async def _charge_lock(self) -> AsyncIterator[None]: - """Context manager to acquire the charge lock if PPE charging manager is active.""" - charging_manager = charging_manager_ctx.get() - if charging_manager: - async with charging_manager.charge_lock(): - yield - else: - yield diff --git a/tests/unit/actor/test_actor_charge.py b/tests/unit/actor/test_actor_charge.py index da957a71..add57eeb 100644 --- a/tests/unit/actor/test_actor_charge.py +++ b/tests/unit/actor/test_actor_charge.py @@ -1,13 +1,19 @@ +from __future__ import annotations + import asyncio -from collections.abc import AsyncGenerator from contextlib import asynccontextmanager from decimal import Decimal -from typing import NamedTuple +from typing import TYPE_CHECKING, NamedTuple from unittest.mock import AsyncMock, Mock, patch from apify import Actor, Configuration from apify._charging import ChargingManagerImplementation, PayPerEventActorPricingInfo, PricingInfoItem +if TYPE_CHECKING: + from collections.abc import AsyncGenerator + + import pytest + class MockedChargingSetup(NamedTuple): """Container for mocked charging components.""" @@ -234,6 +240,37 @@ async def test_charge_lock_concurrent_with_limited_budget() -> None: assert setup.charging_mgr.get_charged_event_count('apify-default-dataset-item') == 5 +async def test_concurrent_actor_push_data_stays_within_budget() -> None: + """Concurrent `Actor.push_data` calls do not overdraw the budget - the reservation and the charge stay atomic.""" + async with setup_mocked_charging( + Configuration(max_total_charge_usd=Decimal('0.50'), test_pay_per_event=True), + {'scrape': Decimal('0.10')}, + ) as setup: + # Both try to push 5 items, but the budget only allows 5 in total. + await asyncio.gather( + Actor.push_data([{'source': 'a', 'id': i} for i in range(5)], charged_event_name='scrape'), + Actor.push_data([{'source': 'b', 'id': i} for i in range(5)], charged_event_name='scrape'), + ) + + assert setup.charging_mgr.get_charged_event_count('scrape') == 5 + + dataset = await Actor.open_dataset() + items = await dataset.get_data() + assert len(items.items) == 5 + + +async def test_push_data_does_not_take_charge_lock_without_pay_per_event(monkeypatch: pytest.MonkeyPatch) -> None: + """`Actor.push_data` leaves the charge lock alone when the Actor does not use the pay-per-event pricing model.""" + async with Actor: + charging_manager = Actor.get_charging_manager() + charge_lock = Mock(wraps=charging_manager.charge_lock) + monkeypatch.setattr(charging_manager, 'charge_lock', charge_lock) + + await Actor.push_data({'id': 1}) + + charge_lock.assert_not_called() + + async def test_charge_with_overdrawn_budget() -> None: configuration = Configuration( max_total_charge_usd=Decimal('0.00025'), diff --git a/tests/unit/storage_clients/test_apify_dataset_client.py b/tests/unit/storage_clients/test_apify_dataset_client.py index 263d30e5..7d2752bc 100644 --- a/tests/unit/storage_clients/test_apify_dataset_client.py +++ b/tests/unit/storage_clients/test_apify_dataset_client.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +from typing import Any from unittest.mock import AsyncMock import pytest @@ -31,3 +32,20 @@ async def test_drop_calls_api_delete() -> None: client, api_client = _make_dataset_client() await client.drop() api_client.delete.assert_awaited_once() + + +async def test_concurrent_push_data_overlaps() -> None: + """Concurrent pushes reach the API at the same time instead of queueing behind each other.""" + concurrency = 3 + barrier = asyncio.Barrier(concurrency) + api_client = AsyncMock() + + async def push_items(**_kwargs: Any) -> None: + # Every concurrent push must reach the API call before any of them is allowed to return. + await barrier.wait() + + api_client.push_items = push_items + client, _ = _make_dataset_client(api_client) + + async with asyncio.timeout(5): + await asyncio.gather(*(client.push_data({'id': i}) for i in range(concurrency)))