diff --git a/src/apify/storage_clients/_apify/_request_queue_shared_client.py b/src/apify/storage_clients/_apify/_request_queue_shared_client.py index d403d71b..a10c30f8 100644 --- a/src/apify/storage_clients/_apify/_request_queue_shared_client.py +++ b/src/apify/storage_clients/_apify/_request_queue_shared_client.py @@ -73,6 +73,9 @@ def __init__( self._queue_head = deque[str]() """Local cache of request IDs from the request queue head for efficient fetching.""" + self._requests_in_progress = set[str]() + """Request IDs handed to a consumer and not yet handled or reclaimed, tracked to avoid double-handing.""" + self._requests_cache: LRUCache[str, CachedRequest] = LRUCache(maxsize=cache_size) """LRU cache storing request objects, keyed by request ID.""" @@ -219,12 +222,55 @@ async def fetch_next_request(self) -> Request | None: async with self._fetch_lock: await self._ensure_head_is_non_empty() + # Pick the next request this client can safely process: one it is not already processing, and whose + # platform lock has not lapsed (an expired lock may have been taken over by another consumer). + now = datetime.now(tz=UTC) + next_request_id: str | None = None + while self._queue_head: + candidate_id = self._queue_head.popleft() + + if candidate_id in self._requests_in_progress: + # Already handed to a consumer in this process; do not process it twice. + continue + + cached = self._requests_cache.get(candidate_id) + if cached is not None and cached.lock_expires_at is not None and cached.lock_expires_at <= now: + logger.debug( + 'Skipping a queued request whose lock has expired; it may have been taken over by another ' + 'client and will be re-fetched with a fresh lock if still available', + extra={'next_request_id': candidate_id}, + ) + continue + + # Reserve the request before releasing the fetch lock so a concurrent fetch cannot pick it too. + self._requests_in_progress.add(candidate_id) + next_request_id = candidate_id + break + # If queue head is empty after ensuring, there are no requests - if not self._queue_head: + if next_request_id is None: return None - # Get the next request ID from the request queue head - next_request_id = self._queue_head.popleft() + # Prolong the lock so the consumer gets a full lock window starting now instead of sharing the batch lock + # acquired when the head was listed. If the lock can no longer be (re)acquired, skip the request. + try: + lock_info = await self._api_client.prolong_request_lock( + next_request_id, + lock_duration=self._DEFAULT_LOCK_TIME, + ) + except Exception as exc: + logger.debug(f'Failed to prolong the lock of request {next_request_id}, skipping it: {exc!s}') + self._requests_in_progress.discard(next_request_id) + return None + + # `None` means the lock was not (re)acquired; skip the request rather than process it without a lock. + if lock_info is None: + logger.debug(f'Lock of request {next_request_id} could not be re-acquired, skipping it') + self._requests_in_progress.discard(next_request_id) + return None + + if (cached := self._requests_cache.get(next_request_id)) is not None: + cached.lock_expires_at = lock_info.lock_expires_at request = await self._get_or_hydrate_request(next_request_id) @@ -234,6 +280,7 @@ async def fetch_next_request(self) -> Request | None: 'Cannot find a request from the beginning of queue, will be retried later', extra={'next_request_id': next_request_id}, ) + self._requests_in_progress.discard(next_request_id) return None # If the request was already handled, skip it @@ -242,6 +289,7 @@ async def fetch_next_request(self) -> Request | None: 'Request fetched from the beginning of queue was already handled', extra={'next_request_id': next_request_id}, ) + self._requests_in_progress.discard(next_request_id) return None # `_get_or_hydrate_request` may return a request from the queue-head cache, which is populated by @@ -253,6 +301,7 @@ async def fetch_next_request(self) -> Request | None: 'Request fetched from the beginning of queue was not found in the RQ', extra={'next_request_id': next_request_id}, ) + self._requests_in_progress.discard(next_request_id) return None return request @@ -260,6 +309,8 @@ async def fetch_next_request(self) -> Request | None: async def mark_request_as_handled(self, request: Request) -> ProcessedRequest | None: """Specific implementation of this method for the RQ shared access mode.""" request_id = unique_key_to_request_id(request.unique_key) + # The consumer is done with this request; stop tracking it as in progress. + self._requests_in_progress.discard(request_id) # Set the handled_at timestamp if not already set if request.handled_at is None: request.handled_at = datetime.now(tz=UTC) @@ -303,6 +354,9 @@ async def reclaim_request( # Reclaim with lock to prevent race conditions that could lead to double processing of the same request. async with self._fetch_lock: + request_id = unique_key_to_request_id(request.unique_key) + # The consumer is giving the request back; stop tracking it as in progress so it can be handed out again. + self._requests_in_progress.discard(request_id) try: # Update the request in the API. processed_request = await self._update_request(request, forefront=forefront) @@ -315,7 +369,6 @@ async def reclaim_request( self.metadata.pending_request_count += 1 # Update the cache - request_id = unique_key_to_request_id(request.unique_key) self._cache_request( request_id, processed_request, @@ -344,7 +397,8 @@ async def is_finished(self) -> bool: """Specific implementation of this method for the RQ shared access mode.""" async with self._fetch_lock: # Order of operations is important here, because affects on `_queue_has_locked_requests`. - return await self._is_empty() and not self._queue_has_locked_requests + # A locally in-progress request keeps the queue unfinished even when the head lists empty. + return await self._is_empty() and not self._queue_has_locked_requests and not self._requests_in_progress async def _is_empty(self) -> bool: """Check whether anything is available to fetch. Lock-free core of `is_empty`, caller must hold the lock.""" @@ -406,7 +460,7 @@ async def _get_or_hydrate_request(self, request_id: str) -> Request | None: if not request: return None - # Update cache with hydrated request + # Cache the hydrated request, preserving any known lock expiry for `fetch_next_request`'s liveness check. self._cache_request( cache_key=request_id, processed_request=ProcessedRequest( @@ -416,6 +470,7 @@ async def _get_or_hydrate_request(self, request_id: str) -> Request | None: was_already_handled=request.handled_at is not None, ), hydrated_request=request, + lock_expires_at=cached_entry.lock_expires_at if cached_entry else None, ) except Exception as exc: logger.debug(f'Error fetching request {request_id}: {exc!s}') @@ -514,7 +569,13 @@ async def _list_head( ) continue - # Cache the request + # Skip requests this client is already processing (e.g. re-listed after their lock lapsed). Re-adding + # them would hand the same request to a second consumer in this process. + if request_id in self._requests_in_progress: + continue + + # Cache the request together with its lock expiry, so `fetch_next_request` can tell whether the lock is + # still held before handing the request to a consumer. self._cache_request( request_id, ProcessedRequest( @@ -524,6 +585,7 @@ async def _list_head( was_already_handled=False, ), hydrated_request=request, + lock_expires_at=request_data.lock_expires_at, ) self._queue_head.append(request_id) @@ -539,14 +601,15 @@ def _cache_request( processed_request: ProcessedRequest, *, hydrated_request: Request | None = None, + lock_expires_at: datetime | None = None, ) -> None: """Cache a request for future use. Args: cache_key: The key to use for caching the request. It should be request ID. processed_request: The processed request information. - forefront: Whether the request was added to the forefront of the queue. hydrated_request: The hydrated request object, if available. + lock_expires_at: When the platform lock on the request expires, if it is currently locked by this client. """ if processed_request.id is None: raise ValueError('ProcessedRequest must have an ID to be cached.') @@ -555,5 +618,5 @@ def _cache_request( id=processed_request.id, was_already_handled=processed_request.was_already_handled, hydrated=hydrated_request, - lock_expires_at=None, + lock_expires_at=lock_expires_at, ) diff --git a/tests/integration/test_request_queue.py b/tests/integration/test_request_queue.py index f53224b6..5d47f0d8 100644 --- a/tests/integration/test_request_queue.py +++ b/tests/integration/test_request_queue.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import contextlib import logging from datetime import UTC, datetime from typing import TYPE_CHECKING, Any, Literal, cast @@ -25,6 +26,7 @@ from crawlee._types import BasicCrawlingContext from apify.storage_clients._apify._models import ApifyRequestQueueMetadata + from apify.storage_clients._apify._request_queue_shared_client import ApifyRequestQueueSharedClient # In shared mode, there is a propagation delay between operations, so we retry reads with the test helper # `poll_until_condition` (exponential backoff). In single mode reads are immediately consistent, so we call once. @@ -1031,7 +1033,8 @@ async def test_request_queue_simple_and_full_at_the_same_time( @pytest.mark.parametrize( ('access', 'expected_write_count_per_request'), - [pytest.param('single', 2, id='Simple rq client'), pytest.param('shared', 3, id='Full rq client')], + # Shared does one extra write/request vs single: `fetch_next_request` prolongs the lock (a PUT). + [pytest.param('single', 2, id='Simple rq client'), pytest.param('shared', 4, id='Full rq client')], ) async def test_crawler_run_request_queue_variant_stats( *, @@ -1609,3 +1612,42 @@ async def test_rq_isolation(apify_token: str, monkeypatch: pytest.MonkeyPatch) - finally: await rq1.drop() await rq2.drop() + + +async def test_shared_is_finished_false_while_request_in_progress( + apify_token: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A shared queue stays unfinished while a request is in progress locally, even after it left the queue.""" + monkeypatch.setenv(ApifyEnvVars.TOKEN, apify_token) + + async with Actor: + rq = await RequestQueue.open(storage_client=ApifyStorageClient(request_queue_access='shared')) + client = cast('ApifyRequestQueueClient', rq._client) + impl = cast('ApifyRequestQueueSharedClient', client._implementation) + try: + await rq.add_request(Request.from_url('https://example.com/in-progress')) + + fetched = await rq.fetch_next_request() + assert fetched is not None + request_id = unique_key_to_request_id(fetched.unique_key) + assert request_id in impl._requests_in_progress + + # Finish the request out-of-band via the raw API (mark handled, drop lock), as another consumer would + # after our lock lapsed, so our local in-progress tracking still holds it. + fetched.handled_at = datetime.now(tz=UTC) + await impl._update_request(fetched) + with contextlib.suppress(Exception): + await impl._api_client.delete_request_lock(request_id) + + # Wait for the head to reflect the empty, unlocked queue (shared-mode propagation delay). + async def _head_empty_and_unlocked() -> bool: + return await rq.is_empty() and not impl._queue_has_locked_requests + + # This is the state where pre-fix `is_finished` wrongly returned True. + assert await poll_until_condition(_head_empty_and_unlocked, timeout=30, backoff_factor=2) is True + + assert request_id in impl._requests_in_progress + assert await rq.is_finished() is False + finally: + await rq.drop() diff --git a/tests/unit/storage_clients/test_apify_request_queue_client.py b/tests/unit/storage_clients/test_apify_request_queue_client.py index 91b117c4..00745d92 100644 --- a/tests/unit/storage_clients/test_apify_request_queue_client.py +++ b/tests/unit/storage_clients/test_apify_request_queue_client.py @@ -1,15 +1,25 @@ from __future__ import annotations import asyncio -from datetime import UTC, datetime +from datetime import UTC, datetime, timedelta from types import SimpleNamespace from typing import TYPE_CHECKING from unittest.mock import AsyncMock import pytest -from apify_client._models import AddedRequest, BatchAddResult, RequestDraft, RequestQueueHead, RequestQueueStats -from crawlee.storage_clients.models import AddRequestsResponse, RequestQueueMetadata +from apify_client._models import ( + AddedRequest, + BatchAddResult, + LockedHeadRequest, + LockedRequestQueueHead, + RequestDraft, + RequestLockInfo, + RequestQueueHead, + RequestQueueStats, +) +from apify_client._models import Request as ClientRequest +from crawlee.storage_clients.models import AddRequestsResponse, ProcessedRequest, RequestQueueMetadata from apify import Request from apify.storage_clients._apify._models import ApifyRequestQueueMetadata @@ -18,7 +28,7 @@ from apify.storage_clients._apify._utils import unique_key_to_request_id if TYPE_CHECKING: - from collections.abc import Sequence + from collections.abc import Callable, Sequence def _make_metadata() -> RequestQueueMetadata: @@ -338,3 +348,254 @@ async def test_partial_unprocessed_commits_only_accepted_requests(access: str) - assert api_client.batch_add_requests.await_args is not None resent = api_client.batch_add_requests.await_args.kwargs['requests'] assert [request['uniqueKey'] for request in resent] == [rejected.unique_key] + + +def _locked_item(request: Request, *, lock_expires_at: datetime) -> LockedHeadRequest: + """Build a single locked head entry for `request` with the given lock expiry.""" + return LockedHeadRequest( + id=unique_key_to_request_id(request.unique_key), + unique_key=request.unique_key, + url=request.url, + method=request.method, + retry_count=0, + lock_expires_at=lock_expires_at, + ) + + +def _locked_head( + items: Sequence[LockedHeadRequest], + *, + queue_has_locked_requests: bool = True, +) -> LockedRequestQueueHead: + """Build a `list_and_lock_head` response wrapping the given locked entries.""" + return LockedRequestQueueHead( + limit=25, + queue_modified_at=datetime.now(tz=UTC), + queue_has_locked_requests=queue_has_locked_requests, + had_multiple_clients=True, + lock_secs=180, + items=list(items), + ) + + +def _processed(request: Request, *, was_already_handled: bool = False) -> ProcessedRequest: + """Build the `update_request` result returned by the API for a handled or reclaimed request.""" + return ProcessedRequest( + id=unique_key_to_request_id(request.unique_key), + unique_key=request.unique_key, + was_already_present=True, + was_already_handled=was_already_handled, + ) + + +def _client_request(request: Request) -> ClientRequest: + """Build the API client's `Request` returned by `get_request` for a hydrated fetch.""" + return ClientRequest( + id=unique_key_to_request_id(request.unique_key), + unique_key=request.unique_key, + url=request.url, + method=request.method, + headers={}, + user_data={}, + retry_count=0, + no_retry=False, + handled_at=None, + ) + + +def _client_request_getter(requests: Sequence[Request]) -> Callable[[str], ClientRequest]: + """Build a `get_request` side effect mapping request IDs back to their API client `Request`.""" + by_id = {unique_key_to_request_id(request.unique_key): _client_request(request) for request in requests} + return lambda request_id: by_id[request_id] + + +async def test_fetch_next_request_prolongs_lock() -> None: + """Handing a request to a consumer prolongs its lock, so a handler slower than the initial lock is still safe.""" + client, api_client = _make_shared_client() + request = Request.from_url('https://example.com/1') + request_id = unique_key_to_request_id(request.unique_key) + future = datetime.now(tz=UTC) + timedelta(seconds=180) + + api_client.list_and_lock_head = AsyncMock( + return_value=_locked_head([_locked_item(request, lock_expires_at=future)]) + ) + api_client.get_request = AsyncMock(return_value=_client_request(request)) + api_client.prolong_request_lock = AsyncMock(return_value=RequestLockInfo(lock_expires_at=future)) + + fetched = await client.fetch_next_request() + + assert fetched is not None + assert fetched.unique_key == request.unique_key + api_client.prolong_request_lock.assert_awaited_once() + assert api_client.prolong_request_lock.await_args is not None + assert api_client.prolong_request_lock.await_args.args[0] == request_id + + +async def test_fetch_next_request_skips_expired_lock() -> None: + """A queued head entry whose lock already lapsed is skipped in favour of the next entry with a live lock.""" + client, api_client = _make_shared_client() + expired = Request.from_url('https://example.com/expired') + valid = Request.from_url('https://example.com/valid') + now = datetime.now(tz=UTC) + + api_client.list_and_lock_head = AsyncMock( + return_value=_locked_head( + [ + _locked_item(expired, lock_expires_at=now - timedelta(seconds=60)), + _locked_item(valid, lock_expires_at=now + timedelta(seconds=180)), + ] + ) + ) + api_client.get_request = AsyncMock(side_effect=_client_request_getter([expired, valid])) + api_client.prolong_request_lock = AsyncMock( + return_value=RequestLockInfo(lock_expires_at=now + timedelta(seconds=180)) + ) + + fetched = await client.fetch_next_request() + + assert fetched is not None + assert fetched.unique_key == valid.unique_key + # Only the live-lock request is prolonged; the expired one is skipped rather than re-processed. + api_client.prolong_request_lock.assert_awaited_once() + assert api_client.prolong_request_lock.await_args is not None + assert api_client.prolong_request_lock.await_args.args[0] == unique_key_to_request_id(valid.unique_key) + + +async def test_fetch_next_request_does_not_rehand_in_progress() -> None: + """A request already in progress is not handed out again even if the platform re-lists it after a lock lapse.""" + client, api_client = _make_shared_client() + request = Request.from_url('https://example.com/1') + future = datetime.now(tz=UTC) + timedelta(seconds=180) + + # The platform keeps returning the same locked request on every head listing. + api_client.list_and_lock_head = AsyncMock( + return_value=_locked_head([_locked_item(request, lock_expires_at=future)]) + ) + api_client.get_request = AsyncMock(return_value=_client_request(request)) + api_client.prolong_request_lock = AsyncMock(return_value=RequestLockInfo(lock_expires_at=future)) + + first = await client.fetch_next_request() + second = await client.fetch_next_request() + + assert first is not None + assert first.unique_key == request.unique_key + assert second is None + + +async def test_fetch_next_request_skips_when_lock_prolong_fails() -> None: + """A request whose lock cannot be prolonged is skipped and released, so a later fetch can still hand it out.""" + client, api_client = _make_shared_client() + request = Request.from_url('https://example.com/1') + future = datetime.now(tz=UTC) + timedelta(seconds=180) + + api_client.list_and_lock_head = AsyncMock( + return_value=_locked_head([_locked_item(request, lock_expires_at=future)]) + ) + api_client.get_request = AsyncMock(return_value=_client_request(request)) + api_client.prolong_request_lock = AsyncMock( + side_effect=[RuntimeError('lock lost'), RequestLockInfo(lock_expires_at=future)] + ) + + first = await client.fetch_next_request() + second = await client.fetch_next_request() + + assert first is None + assert second is not None + assert second.unique_key == request.unique_key + + +async def test_fetch_next_request_skips_when_lock_not_reacquired() -> None: + """A `None` from `prolong_request_lock` (lock not re-acquired) is a skip, not a successful hand-out.""" + client, api_client = _make_shared_client() + request = Request.from_url('https://example.com/1') + future = datetime.now(tz=UTC) + timedelta(seconds=180) + + api_client.list_and_lock_head = AsyncMock( + return_value=_locked_head([_locked_item(request, lock_expires_at=future)]) + ) + api_client.get_request = AsyncMock(return_value=_client_request(request)) + # First fetch: lock not re-acquired (`None`); second: re-acquired. + api_client.prolong_request_lock = AsyncMock(side_effect=[None, RequestLockInfo(lock_expires_at=future)]) + + first = await client.fetch_next_request() + second = await client.fetch_next_request() + + assert first is None + assert second is not None + assert second.unique_key == request.unique_key + + +async def test_is_finished_false_while_request_in_progress() -> None: + """`is_finished` stays False while a fetched request is still in progress, even if the head lists empty.""" + client, api_client = _make_shared_client() + request = Request.from_url('https://example.com/1') + future = datetime.now(tz=UTC) + timedelta(seconds=180) + + api_client.list_and_lock_head = AsyncMock( + side_effect=[ + _locked_head([_locked_item(request, lock_expires_at=future)]), + # In-progress request stays locked: head lists empty, no locked requests reported. + _locked_head([], queue_has_locked_requests=False), + ] + ) + api_client.get_request = AsyncMock(return_value=_client_request(request)) + api_client.prolong_request_lock = AsyncMock(return_value=RequestLockInfo(lock_expires_at=future)) + + fetched = await client.fetch_next_request() + + assert fetched is not None + assert await client.is_finished() is False + + +async def test_mark_request_as_handled_clears_in_progress() -> None: + """Marking a fetched request handled stops tracking it as in progress, so the queue can finish.""" + client, api_client = _make_shared_client() + request = Request.from_url('https://example.com/1') + request_id = unique_key_to_request_id(request.unique_key) + future = datetime.now(tz=UTC) + timedelta(seconds=180) + + api_client.list_and_lock_head = AsyncMock( + return_value=_locked_head([_locked_item(request, lock_expires_at=future)]) + ) + api_client.get_request = AsyncMock(return_value=_client_request(request)) + api_client.prolong_request_lock = AsyncMock(return_value=RequestLockInfo(lock_expires_at=future)) + api_client.update_request = AsyncMock(return_value=_processed(request)) + + fetched = await client.fetch_next_request() + + assert fetched is not None + assert request_id in client._requests_in_progress + + await client.mark_request_as_handled(fetched) + + assert request_id not in client._requests_in_progress + + +async def test_reclaim_request_frees_in_progress() -> None: + """Reclaiming a fetched request stops tracking it as in progress so it can be handed out again.""" + client, api_client = _make_shared_client() + request = Request.from_url('https://example.com/1') + request_id = unique_key_to_request_id(request.unique_key) + future = datetime.now(tz=UTC) + timedelta(seconds=180) + + api_client.list_and_lock_head = AsyncMock( + return_value=_locked_head([_locked_item(request, lock_expires_at=future)]) + ) + api_client.get_request = AsyncMock(return_value=_client_request(request)) + api_client.prolong_request_lock = AsyncMock(return_value=RequestLockInfo(lock_expires_at=future)) + api_client.update_request = AsyncMock(return_value=_processed(request)) + + first = await client.fetch_next_request() + + assert first is not None + assert request_id in client._requests_in_progress + + await client.reclaim_request(first) + + assert request_id not in client._requests_in_progress + + # After reclaim the same request is eligible to be handed out again. + second = await client.fetch_next_request() + + assert second is not None + assert second.unique_key == request.unique_key