diff --git a/src/crawlee/storage_clients/_memory/_request_queue_client.py b/src/crawlee/storage_clients/_memory/_request_queue_client.py index ae98cc9ad6..1ab8d972be 100644 --- a/src/crawlee/storage_clients/_memory/_request_queue_client.py +++ b/src/crawlee/storage_clients/_memory/_request_queue_client.py @@ -1,7 +1,6 @@ from __future__ import annotations -from collections import deque -from contextlib import suppress +from collections import OrderedDict from datetime import datetime, timezone from logging import getLogger from typing import TYPE_CHECKING @@ -42,8 +41,12 @@ def __init__( """ self._metadata = metadata - self._pending_requests = deque[Request]() - """Pending requests are those that have been added to the queue but not yet fetched for processing.""" + self._pending_requests = OrderedDict[str, Request]() + """Pending requests are those that have been added to the queue but not yet fetched for processing. + + Keyed by unique key and ordered from the front of the queue to its end, which keeps both fetching and + repositioning a request to the forefront O(1). + """ self._handled_requests = dict[str, Request]() """Handled requests are those that have been processed and marked as handled.""" @@ -145,6 +148,9 @@ async def add_batch_of_requests( forefront: bool = False, ) -> AddRequestsResponse: processed_requests = [] + new_total_request_count = self._metadata.total_request_count + new_pending_request_count = self._metadata.pending_request_count + for request in requests: # Check if the request is already in the queue by unique_key. existing_request = self._requests_by_unique_key.get(request.unique_key) @@ -175,35 +181,18 @@ async def add_batch_of_requests( ) continue - # If the request is already in the queue but not handled, update it. - if was_already_present and existing_request: - # Update indexes. - self._requests_by_unique_key[request.unique_key] = request - - # We only update `forefront` by updating its position by shifting it to the left. - if forefront: - # Update the existing request with any new data and - # remove old request from pending queue if it's there. - with suppress(ValueError): - self._pending_requests.remove(existing_request) - - # Add updated request back to queue. - self._pending_requests.appendleft(request) - - # Add the new request to the queue. - else: - if forefront: - self._pending_requests.appendleft(request) - else: - self._pending_requests.append(request) - - # Update indexes. + # A new request is registered and appended to the end of the queue. A re-add of a still-pending + # request keeps the originally enqueued object: the incoming duplicate is typically a freshly built + # one that lost the state accumulated so far (e.g. `retry_count`). + if not was_already_present: self._requests_by_unique_key[request.unique_key] = request + self._pending_requests[request.unique_key] = request + new_total_request_count += 1 + new_pending_request_count += 1 - await self._update_metadata( - new_total_request_count=self._metadata.total_request_count + 1, - new_pending_request_count=self._metadata.pending_request_count + 1, - ) + # The only effect a re-add may have is repositioning the request to the front of the queue. + if forefront: + self._pending_requests.move_to_end(request.unique_key, last=False) processed_requests.append( ProcessedRequest( @@ -213,7 +202,12 @@ async def add_batch_of_requests( ) ) - await self._update_metadata(update_accessed_at=True, update_modified_at=True) + await self._update_metadata( + update_accessed_at=True, + update_modified_at=True, + new_total_request_count=new_total_request_count, + new_pending_request_count=new_pending_request_count, + ) return AddRequestsResponse( processed_requests=processed_requests, @@ -222,22 +216,14 @@ async def add_batch_of_requests( @override async def fetch_next_request(self) -> Request | None: - while self._pending_requests: - request = self._pending_requests.popleft() - - # Skip if already handled (shouldn't happen, but safety check). - if request.was_already_handled: - continue - - # Skip if already in progress (shouldn't happen, but safety check). - if request.unique_key in self._in_progress_requests: - continue + if not self._pending_requests: + return None - # Mark as in progress. - self._in_progress_requests[request.unique_key] = request - return request + _, request = self._pending_requests.popitem(last=False) - return None + # Mark as in progress. + self._in_progress_requests[request.unique_key] = request + return request @override async def get_request(self, unique_key: str) -> Request | None: @@ -290,11 +276,12 @@ async def reclaim_request( # Remove from in-progress. del self._in_progress_requests[request.unique_key] - # Add request back to pending queue. + # Add the request back to the pending queue. Unlike a re-add, a reclaim carries the state accumulated + # while the request was in progress, so both stores are updated with the reclaimed object. + self._requests_by_unique_key[request.unique_key] = request + self._pending_requests[request.unique_key] = request if forefront: - self._pending_requests.appendleft(request) - else: - self._pending_requests.append(request) + self._pending_requests.move_to_end(request.unique_key, last=False) # Update metadata timestamps. await self._update_metadata(update_modified_at=True) diff --git a/tests/unit/storage_clients/_memory/test_memory_rq_client.py b/tests/unit/storage_clients/_memory/test_memory_rq_client.py index bc04b80926..00ad75f9ee 100644 --- a/tests/unit/storage_clients/_memory/test_memory_rq_client.py +++ b/tests/unit/storage_clients/_memory/test_memory_rq_client.py @@ -93,3 +93,261 @@ async def test_memory_metadata_updates(rq_client: MemoryRequestQueueClient) -> N assert metadata.created_at == initial_created assert metadata.modified_at > initial_modified assert metadata.accessed_at > accessed_after_read + + +async def test_readd_of_pending_requests_does_not_create_duplicates(rq_client: MemoryRequestQueueClient) -> None: + """Test that mixing regular and forefront re-adds of pending requests leaves no stale duplicates behind.""" + urls = [f'https://example.com/{i}' for i in range(3)] + await rq_client.add_batch_of_requests([Request.from_url(url) for url in urls]) + + # Re-add the same URLs as freshly built objects with a differing payload, the way the higher-level API does + # when it rebuilds requests discovered on another page. + for version in range(2): + duplicates = [] + for url in urls: + duplicate = Request.from_url(url) + duplicate.user_data['version'] = version + duplicates.append(duplicate) + + await rq_client.add_batch_of_requests(duplicates) + await rq_client.add_batch_of_requests(duplicates, forefront=True) + + fetched = [] + while (request := await rq_client.fetch_next_request()) is not None: + fetched.append(request) + await rq_client.mark_request_as_handled(request) + + assert len(fetched) == len(urls) + assert {request.url for request in fetched} == set(urls) + assert await rq_client.is_finished() is True + + metadata = await rq_client.get_metadata() + assert metadata.total_request_count == len(urls) + assert metadata.pending_request_count == 0 + assert metadata.handled_request_count == len(urls) + + +async def test_forefront_readd_preserves_order_and_dedup(rq_client: MemoryRequestQueueClient) -> None: + """Test that repositioning already-pending requests to the forefront keeps LIFO order and dedup.""" + requests = [Request.from_url(f'https://example.com/{i}') for i in range(3)] + await rq_client.add_batch_of_requests(requests) + + # Re-add a subset (0 and 1) to the forefront while still pending. Request 1 is added last, so it must + # end up at the very front, followed by request 0, then the untouched regular request 2. + await rq_client.add_batch_of_requests(requests[:2], forefront=True) + + fetched_urls = [] + while (request := await rq_client.fetch_next_request()) is not None: + fetched_urls.append(request.url) + await rq_client.mark_request_as_handled(request) + + assert fetched_urls == [ + 'https://example.com/1', + 'https://example.com/0', + 'https://example.com/2', + ] + + # No stale duplicates should linger after all live requests are drained. + assert await rq_client.is_empty() is True + assert await rq_client.is_finished() is True + + +async def test_forefront_readd_of_in_progress_request_is_a_no_op(rq_client: MemoryRequestQueueClient) -> None: + """Test that a forefront re-add of an in-progress request does not raise and does not re-enqueue it.""" + requests = [Request.from_url(f'https://example.com/{i}') for i in range(3)] + await rq_client.add_batch_of_requests(requests) + + in_progress = await rq_client.fetch_next_request() + assert in_progress is not None + + response = await rq_client.add_batch_of_requests([in_progress], forefront=True) + assert len(response.processed_requests) == 1 + assert response.processed_requests[0].was_already_present is True + assert response.processed_requests[0].was_already_handled is False + + # The request must not have been queued again ahead of the requests still waiting behind it. + await rq_client.mark_request_as_handled(in_progress) + + fetched_urls = [] + while (request := await rq_client.fetch_next_request()) is not None: + fetched_urls.append(request.url) + await rq_client.mark_request_as_handled(request) + + assert fetched_urls == [ + 'https://example.com/1', + 'https://example.com/2', + ] + + metadata = await rq_client.get_metadata() + assert metadata.total_request_count == 3 + assert metadata.pending_request_count == 0 + assert metadata.handled_request_count == 3 + + +async def test_forefront_readd_of_handled_request_is_a_no_op(rq_client: MemoryRequestQueueClient) -> None: + """Test that a forefront re-add of an already handled request does not raise and does not re-enqueue it.""" + await rq_client.add_batch_of_requests([Request.from_url('https://example.com/page')]) + + fetched = await rq_client.fetch_next_request() + assert fetched is not None + await rq_client.mark_request_as_handled(fetched) + + response = await rq_client.add_batch_of_requests([Request.from_url('https://example.com/page')], forefront=True) + assert len(response.processed_requests) == 1 + assert response.processed_requests[0].was_already_present is True + assert response.processed_requests[0].was_already_handled is True + + assert await rq_client.fetch_next_request() is None + assert await rq_client.is_empty() is True + assert await rq_client.is_finished() is True + + metadata = await rq_client.get_metadata() + assert metadata.total_request_count == 1 + assert metadata.pending_request_count == 0 + assert metadata.handled_request_count == 1 + + +async def test_readd_keeps_the_originally_enqueued_request(rq_client: MemoryRequestQueueClient) -> None: + """Test that re-adding a still-pending request does not replace it with the incoming duplicate.""" + original = Request.from_url('https://example.com/page') + original.user_data['version'] = 1 + await rq_client.add_batch_of_requests([original]) + + # Re-add the same URL while still pending, as a distinct object (as the higher-level API does when it + # rebuilds requests). Neither a regular nor a forefront re-add may overwrite the enqueued request. + duplicate = Request.from_url('https://example.com/page') + duplicate.user_data['version'] = 2 + assert duplicate.unique_key == original.unique_key + await rq_client.add_batch_of_requests([duplicate]) + await rq_client.add_batch_of_requests([duplicate], forefront=True) + + assert await rq_client.get_request(original.unique_key) is original + assert await rq_client.is_empty() is False + + fetched = await rq_client.fetch_next_request() + assert fetched is original + assert fetched.user_data['version'] == 1 + await rq_client.mark_request_as_handled(fetched) + + assert await rq_client.fetch_next_request() is None + assert await rq_client.is_empty() is True + assert await rq_client.is_finished() is True + + metadata = await rq_client.get_metadata() + assert metadata.total_request_count == 1 + assert metadata.pending_request_count == 0 + assert metadata.handled_request_count == 1 + + +@pytest.mark.parametrize( + 'forefront', + [ + pytest.param(False, id='regular re-add'), + pytest.param(True, id='forefront re-add'), + ], +) +async def test_readd_does_not_reset_retry_count_of_reclaimed_request( + rq_client: MemoryRequestQueueClient, + *, + forefront: bool, +) -> None: + """Test that a duplicate enqueued while a failed request awaits a retry does not reset its `retry_count`.""" + request = Request.from_url('https://example.com/page') + await rq_client.add_batch_of_requests([request]) + + fetched = await rq_client.fetch_next_request() + assert fetched is not None + fetched.retry_count += 1 + await rq_client.reclaim_request(fetched) + + # A handler running in parallel discovers the same URL and enqueues a freshly built request for it. + await rq_client.add_batch_of_requests([Request.from_url('https://example.com/page')], forefront=forefront) + + retried = await rq_client.fetch_next_request() + assert retried is not None + assert retried.retry_count == 1 + + +async def test_regular_readd_does_not_reorder_pending_queue(rq_client: MemoryRequestQueueClient) -> None: + """Test that a regular re-add of an already-pending request leaves the FIFO order untouched.""" + requests = [Request.from_url(f'https://example.com/{i}') for i in range(3)] + await rq_client.add_batch_of_requests(requests) + + # Re-add the first request (still pending) without `forefront`; it must stay in its original position. + await rq_client.add_batch_of_requests([requests[0]]) + + fetched_urls = [] + while (request := await rq_client.fetch_next_request()) is not None: + fetched_urls.append(request.url) + await rq_client.mark_request_as_handled(request) + + assert fetched_urls == [ + 'https://example.com/0', + 'https://example.com/1', + 'https://example.com/2', + ] + + +async def test_forefront_readd_does_not_compare_pending_requests( + rq_client: MemoryRequestQueueClient, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Test that a forefront re-add repositions pending requests without comparing them pair by pair.""" + requests = [Request.from_url(f'https://example.com/{i}') for i in range(20)] + await rq_client.add_batch_of_requests(requests) + + # Locating a pending request by equality means scanning the whole queue, which is what makes a batch of K + # forefront re-adds against N pending requests cost O(K*N). Repositioning must be a keyed lookup instead. + comparison_count = 0 + original_eq = Request.__eq__ + + def counting_eq(self: Request, other: object) -> bool: + nonlocal comparison_count + comparison_count += 1 + return original_eq(self, other) + + monkeypatch.setattr(Request, '__eq__', counting_eq) + await rq_client.add_batch_of_requests(requests, forefront=True) + + assert comparison_count == 0 + + +async def test_reclaim_stores_the_modified_request(rq_client: MemoryRequestQueueClient) -> None: + """Test that reclaiming a modified request updates both the queue and the lookup by unique key.""" + request = Request.from_url('https://example.com/page') + await rq_client.add_batch_of_requests([request]) + + fetched = await rq_client.fetch_next_request() + assert fetched is request + + modified = request.model_copy(deep=True) + modified.user_data['reclaimed'] = True + await rq_client.reclaim_request(modified) + + assert await rq_client.get_request(request.unique_key) is modified + assert await rq_client.is_empty() is False + + reclaimed = await rq_client.fetch_next_request() + assert reclaimed is modified + assert reclaimed.user_data['reclaimed'] is True + + +async def test_reclaim_to_forefront_moves_request_to_the_front(rq_client: MemoryRequestQueueClient) -> None: + """Test that a request reclaimed with `forefront` is fetched again ahead of the requests waiting behind it.""" + requests = [Request.from_url(f'https://example.com/{i}') for i in range(3)] + await rq_client.add_batch_of_requests(requests) + + first = await rq_client.fetch_next_request() + assert first is not None + await rq_client.reclaim_request(first, forefront=True) + + fetched_urls = [] + while (request := await rq_client.fetch_next_request()) is not None: + fetched_urls.append(request.url) + await rq_client.mark_request_as_handled(request) + + assert fetched_urls == [ + 'https://example.com/0', + 'https://example.com/1', + 'https://example.com/2', + ]