update charge_lock

Author: Mantisus

src/apify/_actor.py

@@ -26,7 +26,7 @@
     EventSystemInfoData,
 )
 
-from apify._charging import ChargeResult, ChargingManager, ChargingManagerImplementation
+from apify._charging import DEFAULT_DATASET_ITEM_EVENT, ChargeResult, ChargingManager, ChargingManagerImplementation
 from apify._configuration import Configuration
 from apify._consts import EVENT_LISTENERS_TIMEOUT
 from apify._crypto import decrypt_input_secrets, load_private_key
@@ -613,53 +613,63 @@ async def open_request_queue(
             storage_client=self._storage_client.get_suitable_storage_client(force_cloud=force_cloud),
         )
 
-    @overload
-    async def push_data(self, data: dict | list[dict]) -> None: ...
-    @overload
-    async def push_data(self, data: dict | list[dict], charged_event_name: str) -> ChargeResult: ...
     @_ensure_context
-    async def push_data(self, data: dict | list[dict], charged_event_name: str | None = None) -> ChargeResult | None:
+    async def push_data(self, data: dict | list[dict], charged_event_name: str | None = None) -> ChargeResult:
         """Store an object or a list of objects to the default dataset of the current Actor run.
 
         Args:
             data: The data to push to the default dataset.
             charged_event_name: If provided and if the Actor uses the pay-per-event pricing model,
                 the method will attempt to charge for the event for each pushed item.
         """
-        if not data:
-            return None
-
-        data = data if isinstance(data, list) else [data]
-
         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()
 
+        if not data:
+            charged_event_name = charged_event_name or DEFAULT_DATASET_ITEM_EVENT
+            charge_limit_reached = charging_manager.is_event_charge_limit_reached(charged_event_name)
+
+            return ChargeResult(
+                event_charge_limit_reached=charge_limit_reached,
+                charged_count=0,
+                chargeable_within_limit=charging_manager.compute_chargeable(),
+            )
+
+        data = data if isinstance(data, list) else [data]
+
+        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:
-            # No explicit charging requested; synthetic events are handled within dataset.push_data.
+        async with charging_manager.charge_lock():
+            # Synthetic events are handled within dataset.push_data, only get data for `ChargeResult`.
             if charged_event_name is None:
-                dataset = await self.open_dataset()
+                before = charging_manager.get_charged_event_count(DEFAULT_DATASET_ITEM_EVENT)
                 await dataset.push_data(data)
-                return None
+                after = charging_manager.get_charged_event_count(DEFAULT_DATASET_ITEM_EVENT)
+                return ChargeResult(
+                    event_charge_limit_reached=charging_manager.is_event_charge_limit_reached(
+                        DEFAULT_DATASET_ITEM_EVENT
+                    ),
+                    charged_count=after - before,
+                    chargeable_within_limit=charging_manager.compute_chargeable(),
+                )
 
-            pushed_items_count = self.get_charging_manager().compute_push_data_limit(
+            pushed_items_count = charging_manager.compute_push_data_limit(
                 items_count=len(data),
                 event_name=charged_event_name,
                 is_default_dataset=True,
             )
 
-            dataset = await self.open_dataset()
-
             if pushed_items_count < len(data):
                 await dataset.push_data(data[:pushed_items_count])
             elif pushed_items_count > 0:
                 await dataset.push_data(data)
 
             # Only charge explicit events; synthetic events will be processed within the client.
-            return await self.get_charging_manager().charge(
+            return await charging_manager.charge(
                 event_name=charged_event_name,
                 count=pushed_items_count,
             )
@@ -723,10 +733,9 @@ async def charge(self, event_name: str, count: int = 1) -> ChargeResult:
             event_name: Name of the event to be charged for.
             count: Number of events to charge for.
         """
-        # Acquire lock to prevent race conditions with concurrent charge/push_data calls.
+        # charging_manager.charge() acquires charge_lock internally.
         charging_manager = self.get_charging_manager()
-        async with charging_manager.charge_lock:
-            return await charging_manager.charge(event_name, count)
+        return await charging_manager.charge(event_name, count)
 
     @overload
     def on(

src/apify/_charging.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import asyncio
 import math
 from contextvars import ContextVar
 from dataclasses import dataclass
@@ -18,7 +17,7 @@
     PricePerDatasetItemActorPricingInfo,
     PricingModel,
 )
-from apify._utils import docs_group, ensure_context
+from apify._utils import ReentrantLock, docs_group, ensure_context
 from apify.log import logger
 from apify.storages import Dataset
 
@@ -53,7 +52,7 @@ class ChargingManager(Protocol):
     - Apify platform documentation: https://docs.apify.com/platform/actors/publishing/monetize
     """
 
-    charge_lock: asyncio.Lock
+    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:
@@ -114,6 +113,12 @@ def compute_push_data_limit(
             Max number of items that can be pushed within the budget.
         """
 
+    def is_event_charge_limit_reached(self, event_name: str) -> bool:
+        """Return True if the remaining budget is insufficient to charge even a single event of the given type."""
+
+    def compute_chargeable(self) -> dict[str, int | None]:
+        """Compute the maximum number of events of each type that can be charged within the current budget."""
+
 
 @docs_group('Charging')
 @dataclass(frozen=True)
@@ -170,7 +175,7 @@ def __init__(self, configuration: Configuration, client: ApifyClientAsync) -> No
         self._not_ppe_warning_printed = False
         self.active = False
 
-        self.charge_lock = asyncio.Lock()
+        self.charge_lock = ReentrantLock()
 
     async def __aenter__(self) -> None:
         """Initialize the charging manager - this is called by the `Actor` class and shouldn't be invoked manually."""
@@ -244,13 +249,6 @@ async def __aexit__(
 
     @_ensure_context
     async def charge(self, event_name: str, count: int = 1) -> ChargeResult:
-        def calculate_chargeable() -> dict[str, int | None]:
-            """Calculate the maximum number of events of each type that can be charged within the current budget."""
-            return {
-                event_name: self.calculate_max_event_charge_count_within_limit(event_name)
-                for event_name in self._pricing_info
-            }
-
         # 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:
@@ -262,79 +260,81 @@ def calculate_chargeable() -> dict[str, int | None]:
             return ChargeResult(
                 event_charge_limit_reached=False,
                 charged_count=0,
-                chargeable_within_limit=calculate_chargeable(),
+                chargeable_within_limit=self.compute_chargeable(),
             )
 
-        # START OF CRITICAL SECTION - no awaits here
-
-        # 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)
-
-        if charged_count == 0:
+        if count <= 0:
             return ChargeResult(
-                event_charge_limit_reached=True,
+                event_charge_limit_reached=self.is_event_charge_limit_reached(event_name),
                 charged_count=0,
-                chargeable_within_limit=calculate_chargeable(),
+                chargeable_within_limit=self.compute_chargeable(),
             )
 
-        pricing_info = self._pricing_info.get(
-            event_name,
-            PricingInfoItem(
-                # Use a nonzero price for local development so that the maximum budget can be reached.
-                price=Decimal() if self._is_at_home else Decimal(1),
-                title=f"Unknown event '{event_name}'",
-            ),
-        )
+        async with self.charge_lock():
+            # 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)
 
-        # 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
-
-        # END OF CRITICAL SECTION
+            if charged_count == 0:
+                return ChargeResult(
+                    event_charge_limit_reached=True,
+                    charged_count=0,
+                    chargeable_within_limit=self.compute_chargeable(),
+                )
 
-        # If running on the platform, call the charge endpoint
-        if self._is_at_home:
-            if self._actor_run_id is None:
-                raise RuntimeError('Actor run ID not configured')
-
-            if event_name.startswith('apify-'):
-                # Synthetic events (e.g. apify-default-dataset-item) are tracked internally only,
-                # 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, charged_count)
-            else:
-                logger.warning(f"Attempting to charge for an unknown event '{event_name}'")
-
-        # Log the charged operation (if enabled)
-        if self._charging_log_dataset:
-            await self._charging_log_dataset.push_data(
-                {
-                    'event_name': event_name,
-                    'event_title': pricing_info.title,
-                    'event_price_usd': round(pricing_info.price, 3),
-                    'charged_count': charged_count,
-                    'timestamp': datetime.now(timezone.utc).isoformat(),
-                }
+            pricing_info = self._pricing_info.get(
+                event_name,
+                PricingInfoItem(
+                    # Use a nonzero price for local development so that the maximum budget can be reached.
+                    price=Decimal() if self._is_at_home else Decimal(1),
+                    title=f"Unknown event '{event_name}'",
+                ),
             )
 
-        # If it is not possible to charge the full amount, log that fact
-        if charged_count < count:
-            subject = 'instance' if count == 1 else 'instances'
-            logger.info(
-                f"Charging {count} {subject} of '{event_name}' event would exceed max_total_charge_usd "
-                f'- only {charged_count} events were charged'
-            )
+            # 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
+
+            # If running on the platform, call the charge endpoint
+            if self._is_at_home:
+                if self._actor_run_id is None:
+                    raise RuntimeError('Actor run ID not configured')
+
+                if event_name.startswith('apify-'):
+                    # Synthetic events (e.g. apify-default-dataset-item) are tracked internally only,
+                    # 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, charged_count)
+                else:
+                    logger.warning(f"Attempting to charge for an unknown event '{event_name}'")
+
+            # Log the charged operation (if enabled)
+            if self._charging_log_dataset:
+                await self._charging_log_dataset.push_data(
+                    {
+                        'event_name': event_name,
+                        'event_title': pricing_info.title,
+                        'event_price_usd': round(pricing_info.price, 3),
+                        'charged_count': charged_count,
+                        'timestamp': datetime.now(timezone.utc).isoformat(),
+                    }
+                )
 
-        max_charge_count = self.calculate_max_event_charge_count_within_limit(event_name)
+            # If it is not possible to charge the full amount, log that fact
+            if charged_count < count:
+                subject = 'instance' if count == 1 else 'instances'
+                logger.info(
+                    f"Charging {count} {subject} of '{event_name}' event would exceed max_total_charge_usd "
+                    f'- only {charged_count} events were charged'
+                )
 
-        return ChargeResult(
-            event_charge_limit_reached=max_charge_count is not None and max_charge_count <= 0,
-            charged_count=charged_count,
-            chargeable_within_limit=calculate_chargeable(),
-        )
+            return ChargeResult(
+                event_charge_limit_reached=self.is_event_charge_limit_reached(event_name),
+                charged_count=charged_count,
+                chargeable_within_limit=self.compute_chargeable(),
+            )
 
     @_ensure_context
     def calculate_total_charged_amount(self) -> Decimal:
@@ -394,6 +394,18 @@ def compute_push_data_limit(
         max_count = max(0, math.floor(result)) if result.is_finite() else items_count
         return min(items_count, max_count)
 
+    @_ensure_context
+    def is_event_charge_limit_reached(self, event_name: str) -> bool:
+        max_charge_count = self.calculate_max_event_charge_count_within_limit(event_name)
+        return max_charge_count is not None and max_charge_count <= 0
+
+    @_ensure_context
+    def compute_chargeable(self) -> dict[str, int | None]:
+        return {
+            event_name: self.calculate_max_event_charge_count_within_limit(event_name)
+            for event_name in self._pricing_info
+        }
+
     async def _fetch_pricing_info(self) -> _FetchedPricingInfoDict:
         """Fetch pricing information from environment variables or API."""
         # Check if pricing info is available via environment variables

src/apify/_utils.py

@@ -1,13 +1,18 @@
 from __future__ import annotations
 
+import asyncio
 import builtins
 import inspect
 import sys
 from collections.abc import Callable
+from contextlib import asynccontextmanager
 from enum import Enum
 from functools import wraps
 from importlib import metadata
-from typing import Any, Literal, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncIterator
 
 T = TypeVar('T', bound=Callable[..., Any])
 
@@ -123,3 +128,25 @@ def maybe_extract_enum_member_value(maybe_enum_member: Any) -> Any:
     if isinstance(maybe_enum_member, Enum):
         return maybe_enum_member.value
     return maybe_enum_member
+
+
+class ReentrantLock:
+    """A reentrant lock implementation for asyncio using asyncio.Lock."""
+
+    def __init__(self) -> None:
+        self._lock = asyncio.Lock()
+        self._owner: asyncio.Task | None = None
+
+    @asynccontextmanager
+    async def __call__(self) -> AsyncIterator[None]:
+        """Acquire the lock if it's not already owned by the current task, otherwise proceed without acquiring."""
+        me = asyncio.current_task()
+        if self._owner is me:
+            yield
+            return
+        async with self._lock:
+            self._owner = me
+            try:
+                yield
+            finally:
+                self._owner = None

src/apify/storage_clients/_apify/_dataset_client.py

@@ -142,7 +142,7 @@ 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, self._charge_lock():
+        async with self._charge_lock(), self._lock:
             items = data if isinstance(data, list) else [data]
             limit = self._compute_limit_for_push(len(items))
             items = items[:limit]