diff --git a/src/crawlee/storage_clients/_memory/_request_queue_client.py b/src/crawlee/storage_clients/_memory/_request_queue_client.py index ae98cc9ad6..9a0798dd90 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 datetime import datetime, timezone from logging import getLogger from typing import TYPE_CHECKING @@ -175,19 +174,15 @@ async def add_batch_of_requests( ) continue - # If the request is already in the queue but not handled, update it. + # If the request is already in the queue but not handled, we only reposition it when `forefront` + # is set; a regular re-add leaves the already-pending entry untouched. 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. + # Move the request to the front. The old entry is left in the deque instead of being + # located and removed (`deque.remove` is O(n), which makes a batch of forefront re-adds + # O(n^2)); registering the new object here supersedes the old one, and the stale entry is + # skipped lazily by `fetch_next_request` and `is_empty`. + self._requests_by_unique_key[request.unique_key] = request self._pending_requests.appendleft(request) # Add the new request to the queue. @@ -225,6 +220,11 @@ async def fetch_next_request(self) -> Request | None: while self._pending_requests: request = self._pending_requests.popleft() + # Skip stale entries left behind when a request was repositioned to the forefront while already + # pending. Only the object currently registered for the unique key is live. + if self._requests_by_unique_key.get(request.unique_key) is not request: + continue + # Skip if already handled (shouldn't happen, but safety check). if request.was_already_handled: continue @@ -309,6 +309,13 @@ async def reclaim_request( async def is_empty(self) -> bool: await self._update_metadata(update_accessed_at=True) + # Discard stale entries left at the front by forefront repositioning; a live request is always + # enqueued ahead of the stale entry it supersedes, so pruning stops at the first live request. + while self._pending_requests and ( + self._requests_by_unique_key.get(self._pending_requests[0].unique_key) is not self._pending_requests[0] + ): + self._pending_requests.popleft() + # Queue is empty if there are no pending requests. return len(self._pending_requests) == 0 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..e433fc8689 100644 --- a/tests/unit/storage_clients/_memory/test_memory_rq_client.py +++ b/tests/unit/storage_clients/_memory/test_memory_rq_client.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +from collections import deque from typing import TYPE_CHECKING import pytest @@ -93,3 +94,105 @@ 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_forefront_readd_repositions_without_deque_scan(rq_client: MemoryRequestQueueClient) -> None: + """Test that re-adding pending requests to the forefront does not do an O(n) `deque.remove` scan.""" + + class CountingDeque(deque): # type: ignore[type-arg] + remove_calls = 0 + + def remove(self, value: object) -> None: + CountingDeque.remove_calls += 1 + super().remove(value) + + rq_client._pending_requests = CountingDeque(rq_client._pending_requests) + + requests = [Request.from_url(f'https://example.com/{i}') for i in range(20)] + await rq_client.add_batch_of_requests(requests) + + # Re-add the same, still-pending requests with `forefront=True`. Previously each re-add scanned the + # whole pending deque with `deque.remove` to reposition the existing entry, which is O(n) per request. + await rq_client.add_batch_of_requests(requests, forefront=True) + + assert CountingDeque.remove_calls == 0 + + +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_regular_readd_of_pending_request_is_not_dropped(rq_client: MemoryRequestQueueClient) -> None: + """Test that a regular (non-forefront) re-add of a still-pending request keeps it fetchable. + + The lazy-tombstone skip in `fetch_next_request`/`is_empty` keys off object identity, so a regular re-add + must not repoint the registered object away from the entry still sitting in the pending deque, otherwise + the genuinely-live request would be treated as stale and silently dropped. + """ + original = Request.from_url('https://example.com/page') + 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). `forefront` defaults to False. + duplicate = Request.from_url('https://example.com/page') + assert duplicate is not original + assert duplicate.unique_key == original.unique_key + await rq_client.add_batch_of_requests([duplicate]) + + # The request must still be pending and fetchable exactly once, and the counts must stay consistent. + assert await rq_client.is_empty() is False + + fetched = await rq_client.fetch_next_request() + assert fetched is not None + assert fetched.url == 'https://example.com/page' + 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 + + +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', + ]