perf: allow concurrent dataset pushes by narrowing the push-path locks

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)

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`,

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:
@@ -138,7 +139,9 @@ async def drop(self) -> None:
 
     @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

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))
 

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

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'),

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
@@ -104,3 +105,48 @@ 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() -> 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)))
+
+
+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))
+    api_client = AsyncMock()
+    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']))
+
+    api_client.push_items = push_items
+    client, _ = _make_dataset_client(api_client)
+
+    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))