update cherge_lock

Author: Mantisus

src/apify/_actor.py

@@ -380,14 +380,6 @@ def event_manager(self) -> EventManager:
     def _charging_manager_implementation(self) -> ChargingManagerImplementation:
         return ChargingManagerImplementation(self.configuration, self.apify_client)
 
-    @cached_property
-    def _charge_lock(self) -> asyncio.Lock:
-        """Lock to synchronize charge operations.
-
-        Prevents race conditions between Actor.charge and Actor.push_data calls.
-        """
-        return asyncio.Lock()
-
     @cached_property
     def _storage_client(self) -> SmartApifyStorageClient:
         """Storage client used by the Actor.
@@ -642,16 +634,18 @@ async def push_data(self, data: dict | list[dict], charged_event_name: str | Non
         if charged_event_name and charged_event_name.startswith('apify-'):
             raise ValueError(f'Cannot charge for synthetic event "{charged_event_name}" manually')
 
+        charging_manager = self.get_charging_manager()
+
         # 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 self._charge_lock:
+        async with charging_manager.charge_lock:
             # No explicit charging requested; synthetic events are handled within dataset.push_data.
             if charged_event_name is None:
                 dataset = await self.open_dataset()
                 await dataset.push_data(data)
                 return None
 
-            pushed_items_count = self.get_charging_manager().calculate_push_data_limit(
+            pushed_items_count = self.get_charging_manager().compute_push_data_limit(
                 items_count=len(data),
                 event_name=charged_event_name,
                 is_default_dataset=True,
@@ -730,8 +724,9 @@ async def charge(self, event_name: str, count: int = 1) -> ChargeResult:
             count: Number of events to charge for.
         """
         # Acquire lock to prevent race conditions with concurrent charge/push_data calls.
-        async with self._charge_lock:
-            return await self.get_charging_manager().charge(event_name, count)
+        charging_manager = self.get_charging_manager()
+        async with charging_manager.charge_lock:
+            return await charging_manager.charge(event_name, count)
 
     @overload
     def on(

src/apify/_charging.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import asyncio
 import math
 from contextvars import ContextVar
 from dataclasses import dataclass
@@ -52,6 +53,9 @@ class ChargingManager(Protocol):
     - Apify platform documentation: https://docs.apify.com/platform/actors/publishing/monetize
     """
 
+    charge_lock: asyncio.Lock
+    """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:
         """Charge for a specified number of events - sub-operations of the Actor.
 
@@ -88,14 +92,14 @@ def get_charged_event_count(self, event_name: str) -> int:
     def get_max_total_charge_usd(self) -> Decimal:
         """Get the configured maximum total charge for this Actor run."""
 
-    def calculate_push_data_limit(
+    def compute_push_data_limit(
         self,
         items_count: int,
         event_name: str,
         *,
         is_default_dataset: bool,
     ) -> int:
-        """Calculate how many items can be pushed and charged within the current budget.
+        """Compute how many items can be pushed and charged within the current budget.
 
         Accounts for both the explicit event and the synthetic `DEFAULT_DATASET_ITEM_EVENT` event,
         so that the combined cost per item does not exceed the remaining budget.
@@ -166,6 +170,8 @@ def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> No
         self._not_ppe_warning_printed = False
         self.active = False
 
+        self.charge_lock = asyncio.Lock()
+
     async def __aenter__(self) -> None:
         """Initialize the charging manager - this is called by the `Actor` class and shouldn't be invoked manually."""
         # Validate config
@@ -370,7 +376,7 @@ def get_max_total_charge_usd(self) -> Decimal:
         return self._max_total_charge_usd
 
     @_ensure_context
-    def calculate_push_data_limit(
+    def compute_push_data_limit(
         self,
         items_count: int,
         event_name: str,

src/apify/storage_clients/_apify/_dataset_client.py

@@ -142,9 +142,9 @@ async def payloads_generator(items: list[Any]) -> AsyncIterator[str]:
             for index, item in enumerate(items):
                 yield await self._check_and_serialize(item, index)
 
-        async with self._lock:
+        async with self._lock, self._charge_lock():
             items = data if isinstance(data, list) else [data]
-            limit = self._calculate_limit_for_push(len(items))
+            limit = self._compute_limit_for_push(len(items))
             items = items[:limit]
 
             async for chunk in self._chunk_by_size(payloads_generator(items)):

src/apify/storage_clients/_file_system/_dataset_client.py

@@ -48,12 +48,12 @@ async def open(
 
     @override
     async def push_data(self, data: list[dict[str, Any]] | dict[str, Any]) -> None:
-        async with self._lock:
+        async with self._lock, self._charge_lock():
             items = data if isinstance(data, list) else [data]
-            limit = self._calculate_limit_for_push(len(items))
+            limit = self._compute_limit_for_push(len(items))
 
             new_item_count = self._metadata.item_count
-            for item in items:
+            for item in items[:limit]:
                 new_item_count += 1
                 await self._push_item(item, new_item_count)
 

src/apify/storage_clients/_ppe_dataset_mixin.py

@@ -1,13 +1,21 @@
+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."""
 
     def __init__(self) -> None:
         self.is_default_dataset = False
 
-    def _calculate_limit_for_push(self, items_count: int) -> int:
+    def _compute_limit_for_push(self, items_count: int) -> int:
         if self.is_default_dataset and (charging_manager := charging_manager_ctx.get()):
             max_charged_count = charging_manager.calculate_max_event_charge_count_within_limit(
                 event_name=DEFAULT_DATASET_ITEM_EVENT
@@ -21,3 +29,18 @@ 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:
+            if charging_manager.charge_lock.locked():
+                # If the charge lock is already locked, it means we're called from within Actor.push_data which
+                # already holds the lock. asyncio.Lock is not reentrant, so re-acquiring would deadlock.
+                yield
+            else:
+                async with charging_manager.charge_lock:
+                    yield
+        else:
+            yield

tests/unit/actor/test_actor_charge.py

@@ -1,3 +1,4 @@
+import asyncio
 from collections.abc import AsyncGenerator
 from contextlib import asynccontextmanager
 from decimal import Decimal
@@ -186,6 +187,61 @@ async def test_push_data_charges_synthetic_event_for_default_dataset() -> None:
         assert setup.charging_mgr.get_charged_event_count('apify-default-dataset-item') == 3
 
 
+async def test_charge_lock_concurrent_actor_and_dataset_push() -> None:
+    """Test that charge_lock properly synchronizes concurrent Actor.push_data and dataset.push_data calls."""
+    async with setup_mocked_charging(
+        Configuration(max_total_charge_usd=Decimal('10.00'), test_pay_per_event=True)
+    ) as setup:
+        setup.charging_mgr._pricing_info['event'] = PricingInfoItem(Decimal('0.10'), 'Event')
+        setup.charging_mgr._pricing_info['apify-default-dataset-item'] = PricingInfoItem(
+            Decimal('0.10'), 'Dataset item'
+        )
+
+        dataset = await Actor.open_dataset()
+
+        # Run concurrent pushes - Actor.push_data and direct dataset.push_data
+        await asyncio.gather(
+            Actor.push_data([{'source': 'actor', 'id': i} for i in range(5)], 'event'),
+            dataset.push_data([{'source': 'dataset', 'id': i} for i in range(5)]),
+        )
+
+        # Verify all items were pushed
+        items = await dataset.get_data()
+        assert len(items.items) == 10
+
+        # Verify charging was tracked correctly:
+        # - Actor.push_data charged 'event' (5) + 'apify-default-dataset-item' (5)
+        # - dataset.push_data charged 'apify-default-dataset-item' (5)
+        assert setup.charging_mgr.get_charged_event_count('event') == 5
+        assert setup.charging_mgr.get_charged_event_count('apify-default-dataset-item') == 10
+
+
+async def test_charge_lock_concurrent_with_limited_budget() -> None:
+    """Test that charge_lock correctly limits items when concurrent pushes compete for limited budget."""
+    async with setup_mocked_charging(
+        Configuration(max_total_charge_usd=Decimal('0.50'), test_pay_per_event=True)
+    ) as setup:
+        # Each default dataset item costs $0.10, so max 5 items total
+        setup.charging_mgr._pricing_info['apify-default-dataset-item'] = PricingInfoItem(
+            Decimal('0.10'), 'Dataset item'
+        )
+
+        dataset = await Actor.open_dataset()
+
+        # Both try to push 5 items, but budget only allows 5 total
+        await asyncio.gather(
+            dataset.push_data([{'source': 'a', 'id': i} for i in range(5)]),
+            dataset.push_data([{'source': 'b', 'id': i} for i in range(5)]),
+        )
+
+        # Verify total items pushed does not exceed budget limit
+        items = await dataset.get_data()
+        assert len(items.items) == 5  # Budget allows max 5 items at $0.10 each
+
+        # Verify total charged events matches items pushed
+        assert setup.charging_mgr.get_charged_event_count('apify-default-dataset-item') == 5
+
+
 async def test_charge_with_overdrawn_budget() -> None:
     configuration = Configuration(
         max_total_charge_usd=Decimal('0.00025'),

tests/unit/actor/test_charging_manager.py

@@ -247,16 +247,16 @@ async def test_get_max_total_charge_usd(mock_client: MagicMock) -> None:
         assert cm.get_max_total_charge_usd() == Decimal('42.50')
 
 
-async def test_calculate_push_data_limit_no_ppe(mock_client: MagicMock) -> None:
+async def test_compute_push_data_limit_no_ppe(mock_client: MagicMock) -> None:
     """Returns items_count when no PPE pricing is configured (prices are zero)."""
     config = _make_config(actor_pricing_info=None, charged_event_counts={})
     cm = ChargingManagerImplementation(config, mock_client)
     async with cm:
-        result = cm.calculate_push_data_limit(10, 'some-event', is_default_dataset=True)
+        result = cm.compute_push_data_limit(10, 'some-event', is_default_dataset=True)
         assert result == 10
 
 
-async def test_calculate_push_data_limit_within_budget(mock_client: MagicMock) -> None:
+async def test_compute_push_data_limit_within_budget(mock_client: MagicMock) -> None:
     """Returns full items_count when combined budget is sufficient for all items."""
     pricing_info = _make_ppe_pricing_info({'click': Decimal('0.01'), 'apify-default-dataset-item': Decimal('0.01')})
     config = _make_config(
@@ -268,11 +268,11 @@ async def test_calculate_push_data_limit_within_budget(mock_client: MagicMock) -
     cm = ChargingManagerImplementation(config, mock_client)
     async with cm:
         # combined price = 0.02/item, budget = 10.00, max = 500
-        result = cm.calculate_push_data_limit(5, 'click', is_default_dataset=True)
+        result = cm.compute_push_data_limit(5, 'click', is_default_dataset=True)
         assert result == 5
 
 
-async def test_calculate_push_data_limit_budget_exceeded(mock_client: MagicMock) -> None:
+async def test_compute_push_data_limit_budget_exceeded(mock_client: MagicMock) -> None:
     """Returns capped count when combined price (explicit + synthetic) exceeds budget."""
     pricing_info = _make_ppe_pricing_info({'scrape': Decimal('1.00'), 'apify-default-dataset-item': Decimal('1.00')})
     config = _make_config(
@@ -284,11 +284,11 @@ async def test_calculate_push_data_limit_budget_exceeded(mock_client: MagicMock)
     cm = ChargingManagerImplementation(config, mock_client)
     async with cm:
         # combined price = 2.00/item, budget = 3.00, max = floor(3/2) = 1
-        result = cm.calculate_push_data_limit(5, 'scrape', is_default_dataset=True)
+        result = cm.compute_push_data_limit(5, 'scrape', is_default_dataset=True)
         assert result == 1
 
 
-async def test_calculate_push_data_limit_without_default_dataset(mock_client: MagicMock) -> None:
+async def test_compute_push_data_limit_without_default_dataset(mock_client: MagicMock) -> None:
     """When not pushing to the default dataset, only explicit event price is considered."""
     pricing_info = _make_ppe_pricing_info({'scrape': Decimal('1.00'), 'apify-default-dataset-item': Decimal('1.00')})
     config = _make_config(
@@ -300,11 +300,11 @@ async def test_calculate_push_data_limit_without_default_dataset(mock_client: Ma
     cm = ChargingManagerImplementation(config, mock_client)
     async with cm:
         # explicit price only = 1.00/item, budget = 3.00, max = floor(3/1) = 3
-        result = cm.calculate_push_data_limit(5, 'scrape', is_default_dataset=False)
+        result = cm.compute_push_data_limit(5, 'scrape', is_default_dataset=False)
         assert result == 3
 
 
-async def test_calculate_push_data_limit_exhausted_budget(mock_client: MagicMock) -> None:
+async def test_compute_push_data_limit_exhausted_budget(mock_client: MagicMock) -> None:
     """Returns 0 when the budget is fully exhausted before the push."""
     pricing_info = _make_ppe_pricing_info({'scrape': Decimal('1.00')})
     config = _make_config(
@@ -315,7 +315,7 @@ async def test_calculate_push_data_limit_exhausted_budget(mock_client: MagicMock
     )
     cm = ChargingManagerImplementation(config, mock_client)
     async with cm:
-        result = cm.calculate_push_data_limit(5, 'scrape', is_default_dataset=False)
+        result = cm.compute_push_data_limit(5, 'scrape', is_default_dataset=False)
         assert result == 0