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..70d6b67d 100644 --- a/src/apify/storage_clients/_apify/_request_queue_shared_client.py +++ b/src/apify/storage_clients/_apify/_request_queue_shared_client.py @@ -322,8 +322,14 @@ async def reclaim_request( hydrated_request=request, ) - # If we're adding to the forefront, we need to check for forefront requests - # in the next list_head call + # Re-enter into the local head so `is_empty`/`is_finished` still see it while `list_and_lock_head` + # lags behind the reclaim (else it is silently dropped); forefront to the front, default to the back. + if forefront: + self._queue_head.appendleft(request_id) + else: + self._queue_head.append(request_id) + + # Forefront also forces a head refresh, so concurrent forefront additions surface first. if forefront: self._should_check_for_forefront_requests = True @@ -349,7 +355,8 @@ async def is_finished(self) -> bool: async def _is_empty(self) -> bool: """Check whether anything is available to fetch. Lock-free core of `is_empty`, caller must hold the lock.""" head = await self._list_head(limit=1) - return len(head.items) == 0 + # A forefront refresh refetches the platform head (which lags the reclaim), so consult the local head too. + return len(head.items) == 0 and not self._queue_head async def _get_metadata_estimate(self) -> RequestQueueMetadata: """Try to get cached metadata first. If multiple clients, fuse with global metadata. @@ -528,8 +535,9 @@ async def _list_head( self._queue_head.append(request_id) for leftover_id in leftover_buffer: - # After adding new requests to the forefront, any existing leftover locked request is kept in the end. - self._queue_head.append(leftover_id) + # Keep leftovers at the end; skip ids already returned above so a re-entered reclaim isn't enqueued twice. + if leftover_id not in self._queue_head: + self._queue_head.append(leftover_id) return RequestQueueHead.from_client_locked_head(locked_queue_head) diff --git a/src/apify/storage_clients/_apify/_request_queue_single_client.py b/src/apify/storage_clients/_apify/_request_queue_single_client.py index e1339d11..797155f3 100644 --- a/src/apify/storage_clients/_apify/_request_queue_single_client.py +++ b/src/apify/storage_clients/_apify/_request_queue_single_client.py @@ -286,12 +286,11 @@ async def reclaim_request( forefront: bool = False, ) -> ProcessedRequest | None: """Specific implementation of this method for the RQ single access mode.""" - # Check if the request was marked as handled and clear it. When reclaiming, - # we want to put the request back for processing. - request_id = unique_key_to_request_id(request.unique_key) - if request.was_already_handled: + # `was_already_handled` derives from `handled_at`, so capture it before clearing `handled_at`. + was_already_handled = request.was_already_handled + if was_already_handled: request.handled_at = None try: @@ -303,18 +302,21 @@ async def reclaim_request( # No longer handled self._requests_already_handled.discard(request_id) + # Re-enter into the local head (like `add_batch_of_requests`) so `is_empty`/`is_finished` still see it + # while the platform head lags behind the reclaim (else it is dropped); forefront to top, default bottom. if forefront: - # Append to top of the local head estimation self._head_requests.append(request_id) + else: + self._head_requests.appendleft(request_id) + + # Previously handled -> pending; adjust counts before the update so they stay consistent even if it fails. + if was_already_handled: + self.metadata.handled_request_count -= 1 + self.metadata.pending_request_count += 1 processed_request = await self._update_request(request, forefront=forefront) processed_request.id = request_id processed_request.unique_key = request.unique_key - # The platform reports the request's state before this update via `was_already_handled`. If it was - # handled, this update moved it from handled back to pending, so mirror that in the local metadata. - if processed_request.was_already_handled: - self.metadata.handled_request_count -= 1 - self.metadata.pending_request_count += 1 except Exception: logger.exception(f'Error reclaiming request {request.unique_key}') diff --git a/tests/integration/test_request_queue.py b/tests/integration/test_request_queue.py index f53224b6..f6cc6e49 100644 --- a/tests/integration/test_request_queue.py +++ b/tests/integration/test_request_queue.py @@ -362,6 +362,38 @@ async def test_request_reclaim_with_forefront( Actor.log.info(f'Test completed - processed {remaining_count} additional requests') +async def test_forefront_reclaim_of_last_request_completes( + request_queue_apify: RequestQueue, + rq_poll_timeout: int, +) -> None: + """Reclaiming the only in-flight request to the forefront keeps it pending and lets the run finish cleanly.""" + rq = request_queue_apify + + await rq.add_request('https://example.com/only') + request = await poll_until_condition(rq.fetch_next_request, timeout=rq_poll_timeout, backoff_factor=2) + assert request is not None + + # Reclaim the sole in-flight request to the forefront. + await rq.reclaim_request(request, forefront=True) + + # The queue must not report itself empty while the reclaimed request is still pending. + assert await rq.is_empty() is False + + # The request must be retrievable again, not lost; it may take a moment to reappear at the head, so poll. + refetched = await poll_until_condition( + rq.fetch_next_request, + lambda result: result is not None, + timeout=60, + poll_interval=5, + ) + assert refetched is not None + assert refetched.url == request.url + + await rq.mark_request_as_handled(refetched) + is_finished = await poll_until_condition(rq.is_finished, timeout=rq_poll_timeout, backoff_factor=2) + assert is_finished is True + + async def test_reclaim_handled_request_moves_back_to_pending( request_queue_apify: RequestQueue, rq_poll_timeout: int, 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..7557d77e 100644 --- a/tests/unit/storage_clients/test_apify_request_queue_client.py +++ b/tests/unit/storage_clients/test_apify_request_queue_client.py @@ -8,7 +8,16 @@ import pytest -from apify_client._models import AddedRequest, BatchAddResult, RequestDraft, RequestQueueHead, RequestQueueStats +from apify_client._models import ( + AddedRequest, + BatchAddResult, + LockedHeadRequest, + LockedRequestQueueHead, + RequestDraft, + RequestQueueHead, + RequestQueueStats, + RequestRegistration, +) from crawlee.storage_clients.models import AddRequestsResponse, RequestQueueMetadata from apify import Request @@ -338,3 +347,179 @@ 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 _request_registration(request: Request, *, was_already_handled: bool = False) -> RequestRegistration: + """Build an `update_request` response for the given request.""" + return RequestRegistration.model_construct( + request_id=unique_key_to_request_id(request.unique_key), + was_already_present=True, + was_already_handled=was_already_handled, + ) + + +async def test_reclaimed_request_kept_pending_while_head_lags_single() -> None: + """A reclaimed request (default forefront=False) stays pending in single mode while the platform head lags.""" + client, api_client = _make_single_client() + request = Request.from_url('https://example.com/1') + request_id = unique_key_to_request_id(request.unique_key) + + # The platform head listing lags and does not yet report the reclaimed request during the window. + api_client.list_head = AsyncMock( + return_value=RequestQueueHead( + limit=200, + queue_modified_at=datetime.now(tz=UTC), + had_multiple_clients=False, + items=[], + ) + ) + api_client.update_request = AsyncMock(return_value=_request_registration(request)) + + # The request was fetched and is being processed by this client. + client._requests_in_progress.add(request_id) + client._requests_cache[request_id] = request + + # Reclaim it via the default retry path (forefront=False). + assert await client.reclaim_request(request) is not None + + # While the head propagation lags, the request must still count as locally pending, or the run would shut + # down and silently drop it. + assert await client.is_empty() is False + assert await client.is_finished() is False + + +async def test_reclaimed_forefront_request_kept_pending_while_head_lags_single() -> None: + """A forefront reclaim stays pending in single mode while the platform head listing lags behind the reclaim.""" + client, api_client = _make_single_client() + request = Request.from_url('https://example.com/1') + request_id = unique_key_to_request_id(request.unique_key) + + api_client.list_head = AsyncMock( + return_value=RequestQueueHead( + limit=200, + queue_modified_at=datetime.now(tz=UTC), + had_multiple_clients=False, + items=[], + ) + ) + api_client.update_request = AsyncMock(return_value=_request_registration(request)) + + client._requests_in_progress.add(request_id) + client._requests_cache[request_id] = request + + assert await client.reclaim_request(request, forefront=True) is not None + + assert await client.is_empty() is False + assert await client.is_finished() is False + + +async def test_reclaim_handled_update_failure_keeps_local_state_consistent_single() -> None: + """A failed reclaim update still re-queues the handled request locally and moves the counts handled -> pending.""" + client, api_client = _make_single_client() + request = Request.from_url('https://example.com/1') + request_id = unique_key_to_request_id(request.unique_key) + + api_client.list_head = AsyncMock( + return_value=RequestQueueHead( + limit=200, + queue_modified_at=datetime.now(tz=UTC), + had_multiple_clients=False, + items=[], + ) + ) + api_client.update_request = AsyncMock(side_effect=RuntimeError('network down')) + + # An already-handled request is being reclaimed. + request.handled_at = datetime.now(tz=UTC) + client._requests_already_handled.add(request_id) + client.metadata.handled_request_count = 1 + + # The reclaim reports failure (the update raised), yet the request stays locally pending and the counts move + # handled -> pending, so a later successful handling does not double-count it. + assert await client.reclaim_request(request) is None + assert await client.is_empty() is False + assert client.metadata.handled_request_count == 0 + assert client.metadata.pending_request_count == 1 + + +async def test_reclaimed_request_kept_pending_while_head_lags_shared() -> None: + """A reclaimed request (default forefront=False) stays pending in shared mode while the platform head lags.""" + client, api_client = _make_shared_client() + request = Request.from_url('https://example.com/1') + + # `list_and_lock_head` lags and returns nothing during the propagation window. + api_client.list_and_lock_head = AsyncMock( + return_value=LockedRequestQueueHead.model_construct( + limit=1, + queue_modified_at=datetime.now(tz=UTC), + had_multiple_clients=False, + queue_has_locked_requests=False, + lock_secs=180, + items=[], + ) + ) + api_client.update_request = AsyncMock(return_value=_request_registration(request)) + + # Reclaim it via the default retry path (forefront=False). + assert await client.reclaim_request(request) is not None + + # While the head propagation lags, the request must still count as locally pending, or the run would shut + # down and silently drop it. + assert await client.is_empty() is False + assert await client.is_finished() is False + + +def _locked_head_item(request: Request) -> LockedHeadRequest: + """Build a `list_and_lock_head` item for the given request.""" + return LockedHeadRequest.model_construct( + 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=None, + ) + + +def _locked_head(*items: LockedHeadRequest, limit: int = 25) -> LockedRequestQueueHead: + """Build a `list_and_lock_head` response wrapping the given items.""" + return LockedRequestQueueHead.model_construct( + limit=limit, + queue_modified_at=datetime.now(tz=UTC), + had_multiple_clients=False, + queue_has_locked_requests=False, + lock_secs=180, + items=list(items), + ) + + +async def test_reclaimed_forefront_request_kept_pending_while_head_lags_shared() -> None: + """A forefront reclaim stays pending in shared mode while `list_and_lock_head` lags behind the reclaim.""" + client, api_client = _make_shared_client() + request = Request.from_url('https://example.com/1') + + # The forefront reclaim flags a head refresh; the refreshed listing lags and returns nothing in the window. + api_client.list_and_lock_head = AsyncMock(return_value=_locked_head()) + api_client.update_request = AsyncMock(return_value=_request_registration(request)) + + assert await client.reclaim_request(request, forefront=True) is not None + + assert await client.is_empty() is False + assert await client.is_finished() is False + + +async def test_reclaimed_forefront_request_not_duplicated_after_head_refresh_shared() -> None: + """A forefront reclaim re-entered locally must not be duplicated when the refreshed head also returns it.""" + client, api_client = _make_shared_client() + request = Request.from_url('https://example.com/1') + request_id = unique_key_to_request_id(request.unique_key) + + # After the reclaim, the platform head listing catches up and returns the same request. + api_client.list_and_lock_head = AsyncMock(return_value=_locked_head(_locked_head_item(request))) + api_client.update_request = AsyncMock(return_value=_request_registration(request)) + + assert await client.reclaim_request(request, forefront=True) is not None + + # Trigger the forefront refresh flagged by the reclaim; the request must appear exactly once in the head. + await client._list_head() + assert list(client._queue_head).count(request_id) == 1