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..4462a7c2 100644 --- a/src/apify/storage_clients/_apify/_request_queue_single_client.py +++ b/src/apify/storage_clients/_apify/_request_queue_single_client.py @@ -82,7 +82,10 @@ def __init__( """LRU cache storing unhandled request objects, keyed by request ID.""" self._head_requests: deque[str] = deque() - """Ordered queue of request IDs representing the local estimate of the queue head.""" + """Ordered queue of request IDs representing the local estimate of the queue head. + + Served from the right end, so `append` places a request at the top of the head and `appendleft` at the bottom. + """ self._requests_already_handled: set[str] = set() """Set of request IDs known to be already processed on the platform. @@ -303,9 +306,12 @@ async def reclaim_request( # No longer handled self._requests_already_handled.discard(request_id) + # `is_empty` reads only the head estimation, so the reclaimed request has to go back into it until the + # platform head listing catches up. Placement mirrors `add_batch_of_requests`. if forefront: - # Append to top of the local head estimation self._head_requests.append(request_id) + else: + self._head_requests.appendleft(request_id) processed_request = await self._update_request(request, forefront=forefront) processed_request.id = request_id 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 3c9e920b..4d086cb1 100644 --- a/tests/unit/storage_clients/test_apify_request_queue_client.py +++ b/tests/unit/storage_clients/test_apify_request_queue_client.py @@ -2,6 +2,7 @@ import asyncio from datetime import UTC, datetime, timedelta +from itertools import chain, repeat from types import SimpleNamespace from typing import TYPE_CHECKING from unittest.mock import AsyncMock @@ -11,6 +12,7 @@ from apify_client._models import ( AddedRequest, BatchAddResult, + HeadRequest, LockedHeadRequest, LockedRequestQueueHead, RequestDraft, @@ -846,3 +848,112 @@ async def test_reclaim_request_frees_in_progress() -> None: assert second is not None assert second.unique_key == request.unique_key + + +def _head_item(request: Request) -> HeadRequest: + """Build a `list_head` item for the given request.""" + return HeadRequest( + id=unique_key_to_request_id(request.unique_key), + unique_key=request.unique_key, + url=request.url, + method=request.method, + retry_count=0, + ) + + +def _head(*items: HeadRequest) -> RequestQueueHead: + """Build a `list_head` response wrapping the given items.""" + return RequestQueueHead( + limit=200, + queue_modified_at=datetime.now(tz=UTC), + had_multiple_clients=False, + items=list(items), + ) + + +@pytest.mark.parametrize('forefront', [pytest.param(False, id='default'), pytest.param(True, id='forefront')]) +async def test_single_reclaimed_request_kept_pending_while_head_lags(*, forefront: bool) -> None: + """A reclaimed request stays locally pending in single mode while the platform head listing lags behind it.""" + client, api_client = _make_single_client() + request = Request.from_url('https://example.com/1') + + # The head serves the request once, then lags behind the reclaim and keeps listing empty. + api_client.list_head = AsyncMock(side_effect=chain([_head(_head_item(request))], repeat(_head()))) + api_client.get_request = AsyncMock(return_value=_client_request(request, handled_at=None)) + api_client.update_request = AsyncMock(return_value=_request_registration(request, was_already_handled=False)) + + fetched = await client.fetch_next_request() + + assert fetched is not None + + await client.reclaim_request(fetched, forefront=forefront) + + # A True here would let `AutoscaledPool` end the run while the request is still pending on the platform. + assert await client.is_empty() is False + assert await client.is_finished() is False + + refetched = await client.fetch_next_request() + + assert refetched is not None + assert refetched.unique_key == request.unique_key + + +@pytest.mark.parametrize( + ('forefront', 'expected_order'), + [ + pytest.param(False, ['https://example.com/2', 'https://example.com/1'], id='default'), + pytest.param(True, ['https://example.com/1', 'https://example.com/2'], id='forefront'), + ], +) +async def test_single_reclaimed_request_position_follows_forefront( + *, + forefront: bool, + expected_order: list[str], +) -> None: + """A reclaimed request re-enters the local head estimate at the top with `forefront`, at the bottom without it.""" + client, api_client = _make_single_client() + first = Request.from_url('https://example.com/1') + second = Request.from_url('https://example.com/2') + + # The head serves both requests once, then lags behind the reclaim and keeps listing empty. + api_client.list_head = AsyncMock(side_effect=chain([_head(_head_item(first), _head_item(second))], repeat(_head()))) + api_client.get_request = AsyncMock(side_effect=_client_request_getter([first, second])) + api_client.update_request = AsyncMock(return_value=_request_registration(first, was_already_handled=False)) + + fetched = await client.fetch_next_request() + + assert fetched is not None + assert fetched.url == first.url + + await client.reclaim_request(fetched, forefront=forefront) + + first_served = await client.fetch_next_request() + second_served = await client.fetch_next_request() + + assert first_served is not None + assert second_served is not None + assert [first_served.url, second_served.url] == expected_order + + +async def test_single_reclaimed_request_finishes_queue_once_handled() -> None: + """The head entry a reclaim leaves behind is drained, so the queue still reports finished once handled.""" + client, api_client = _make_single_client() + request = Request.from_url('https://example.com/1') + + api_client.list_head = AsyncMock(side_effect=chain([_head(_head_item(request))], repeat(_head()))) + api_client.get_request = AsyncMock(return_value=_client_request(request, handled_at=None)) + api_client.update_request = AsyncMock(return_value=_request_registration(request, was_already_handled=False)) + + fetched = await client.fetch_next_request() + + assert fetched is not None + + await client.reclaim_request(fetched) + + refetched = await client.fetch_next_request() + + assert refetched is not None + + await client.mark_request_as_handled(refetched) + + assert await client.is_finished() is True