diff --git a/src/apify/_actor.py b/src/apify/_actor.py index 9cf699e1..862370cd 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 @@ -690,9 +696,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 db6bec4d..44f1f2c8 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: @@ -41,7 +42,6 @@ def __init__( self, *, api_client: DatasetClientAsync, - lock: asyncio.Lock, ) -> None: """Initialize a new instance. @@ -53,9 +53,6 @@ def __init__( self._api_client = api_client """The Apify dataset client for API operations.""" - self._lock = lock - """A lock to ensure that only one operation is performed at a time.""" - @override async def get_metadata(self) -> DatasetMetadata: metadata = await self._api_client.get() @@ -113,10 +110,7 @@ async def open( id=id, ) - dataset_client = cls( - api_client=api_client, - lock=asyncio.Lock(), - ) + dataset_client = cls(api_client=api_client) dataset_client.is_default_dataset = ( alias is None and name is None and (id is None or id == configuration.default_dataset_id) @@ -133,12 +127,13 @@ async def purge(self) -> None: @override async def drop(self) -> None: - async with self._lock: - await self._api_client.delete() + await self._api_client.delete() @override async def push_data(self, data: Sequence[Mapping[str, JsonSerializable]] | Mapping[str, JsonSerializable]) -> None: - 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 acdb529a..37c70b3b 100644 --- a/tests/unit/storage_clients/test_apify_dataset_client.py +++ b/tests/unit/storage_clients/test_apify_dataset_client.py @@ -2,6 +2,7 @@ import asyncio import json +from typing import Any from unittest.mock import AsyncMock, Mock import pytest @@ -16,10 +17,7 @@ def _make_dataset_client(api_client: AsyncMock | None = None) -> tuple[ApifyData if api_client is None: api_client = AsyncMock() - return ApifyDatasetClient( - api_client=api_client, - lock=asyncio.Lock(), - ), api_client + return ApifyDatasetClient(api_client=api_client), api_client async def test_purge_raises_not_implemented() -> None: @@ -104,3 +102,46 @@ async def test_push_data_rejects_a_non_serializable_item() -> None: with pytest.raises(ValueError, match='at index 0 is not serializable'): await client.push_data(circular) + + +async def test_concurrent_push_data_overlaps(monkeypatch: pytest.MonkeyPatch) -> None: + """Concurrent pushes reach the API at the same time instead of queueing behind each other.""" + concurrency = 3 + barrier = asyncio.Barrier(concurrency) + client, api_client = _make_dataset_client() + + 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() + + monkeypatch.setattr(api_client, 'push_items', AsyncMock(side_effect=push_items)) + + async with asyncio.timeout(5): + await asyncio.gather(*(client.push_data({'id': i}) for i in range(concurrency))) + + +async def test_concurrent_multi_chunk_pushes_preserve_per_push_order(monkeypatch: pytest.MonkeyPatch) -> None: + """A push's own chunks stay in order even while they interleave on the wire with other pushes' chunks.""" + monkeypatch.setattr(ApifyDatasetClient, '_EFFECTIVE_LIMIT_SIZE', ByteSize(60)) + client, api_client = _make_dataset_client() + received: list[tuple[int, int]] = [] + chunk_count = 0 + + async def push_items(**kwargs: Any) -> None: + nonlocal chunk_count + chunk_count += 1 + await asyncio.sleep(0) # Yield so chunks from other concurrent pushes can land in between. + received.extend((item['push'], item['i']) for item in json.loads(kwargs['items'])) + + monkeypatch.setattr(api_client, 'push_items', AsyncMock(side_effect=push_items)) + + concurrency, items_per_push = 4, 6 + async with asyncio.timeout(5): + await asyncio.gather( + *(client.push_data([{'push': p, 'i': i} for i in range(items_per_push)]) for p in range(concurrency)) + ) + + assert chunk_count > concurrency, 'each push must split into more than one chunk' + for push in range(concurrency): + indices = [i for p, i in received if p == push] + assert indices == list(range(items_per_push)) diff --git a/uv.lock b/uv.lock index 4b1767b5..6d5bd1a9 100644 --- a/uv.lock +++ b/uv.lock @@ -585,7 +585,7 @@ toml = [ [[package]] name = "crawlee" -version = "1.8.3" +version = "1.9.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "async-timeout" }, @@ -597,14 +597,13 @@ dependencies = [ { name = "psutil" }, { name = "pydantic" }, { name = "pydantic-settings" }, - { name = "pyee" }, { name = "tldextract" }, { name = "typing-extensions" }, { name = "yarl" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/10/b4/4c359c2e073c720b3960d2d41ab6e06b1ea0ff3d570e7c6467b899e787c1/crawlee-1.8.3.tar.gz", hash = "sha256:8f06e4bec07a5438126a22f5c6ae040f238e22eeaba1a4257de4240292419e6d", size = 316089, upload-time = "2026-07-20T07:07:51.617Z" } +sdist = { url = "https://files.pythonhosted.org/packages/56/0e/ff363177fdbbb774d392f419db2363daaf4c3838b1e13649337af9d2ac31/crawlee-1.9.1.tar.gz", hash = "sha256:5ddbe6b7188ad4acee4221ad8afc4c27a6e0f77f815d2c1617e3e00eaafcd950", size = 325529, upload-time = "2026-08-06T09:28:36.247Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/5a/85/05b8c4c73cad542e9aefe9a5568d03f73eb7ccb5f84524c47d066ce3bc85/crawlee-1.8.3-py3-none-any.whl", hash = "sha256:42c3e2404922a1ba51659e97ab31f0b451eeeed02faaf673d5ad4932f12bda7e", size = 403479, upload-time = "2026-07-20T07:07:49.96Z" }, + { url = "https://files.pythonhosted.org/packages/21/eb/ac3e04315130801b6a20e8d14b954f3b1fac27d5bae9e15601b42bbd174b/crawlee-1.9.1-py3-none-any.whl", hash = "sha256:012c2f2c6e2f7021ea772eb2d95951f62092d2de1933f70fdffc31e716c53e94", size = 412364, upload-time = "2026-08-06T09:28:34.759Z" }, ] [package.optional-dependencies] @@ -1852,18 +1851,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/83/5a/ce0b056d9a95fd0c06a6cfa5972477d79353392d19230c748a7ba5a9df04/pydoc_markdown-4.8.2-py3-none-any.whl", hash = "sha256:203f74119e6bb2f9deba43d452422de7c8ec31955b61e0620fa4dd8c2611715f", size = 67830, upload-time = "2023-06-26T12:36:59.502Z" }, ] -[[package]] -name = "pyee" -version = "13.0.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "typing-extensions" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/8b/04/e7c1fe4dc78a6fdbfd6c337b1c3732ff543b8a397683ab38378447baa331/pyee-13.0.1.tar.gz", hash = "sha256:0b931f7c14535667ed4c7e0d531716368715e860b988770fc7eb8578d1f67fc8", size = 31655, upload-time = "2026-02-14T21:12:28.044Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/a0/c4/b4d4827c93ef43c01f599ef31453ccc1c132b353284fc6c87d535c233129/pyee-13.0.1-py3-none-any.whl", hash = "sha256:af2f8fede4171ef667dfded53f96e2ed0d6e6bd7ee3bb46437f77e3b57689228", size = 15659, upload-time = "2026-02-14T21:12:26.263Z" }, -] - [[package]] name = "pygments" version = "2.20.0"