fix: Dequeue items from RequestQueue in the correct order (#411)

Author: janbuchar

src/crawlee/storages/base_storage.py

@@ -10,9 +10,6 @@
 class BaseStorage(ABC):
     """Base class for storages."""
 
-    LABEL = 'Unknown'
-    """Human readable label of the storage."""
-
     @property
     @abstractmethod
     def id(self) -> str:

src/crawlee/storages/request_queue.py

@@ -4,8 +4,7 @@
 from collections import OrderedDict
 from datetime import datetime, timedelta, timezone
 from logging import getLogger
-from typing import TYPE_CHECKING
-from typing import OrderedDict as OrderedDictType
+from typing import TYPE_CHECKING, Generic, TypedDict, TypeVar
 
 from typing_extensions import override
 
@@ -31,6 +30,40 @@
 
 logger = getLogger(__name__)
 
+__all__ = ['RequestQueue']
+
+
+T = TypeVar('T')
+
+
+class BoundedSet(Generic[T]):
+    """A simple set datastructure that removes the least recently accessed item when it reaches `max_length`."""
+
+    def __init__(self, max_length: int) -> None:
+        self._max_length = max_length
+        self._data = OrderedDict[T, object]()
+
+    def __contains__(self, item: T) -> bool:
+        found = item in self._data
+        if found:
+            self._data.move_to_end(item, last=True)
+        return found
+
+    def add(self, item: T) -> None:
+        self._data[item] = True
+        self._data.move_to_end(item)
+
+        if len(self._data) > self._max_length:
+            self._data.popitem(last=False)
+
+    def clear(self) -> None:
+        self._data.clear()
+
+
+class CachedRequest(TypedDict):
+    id: str
+    was_already_handled: bool
+
 
 class RequestQueue(BaseStorage, RequestProvider):
     """Represents a queue storage for HTTP requests to crawl.
@@ -97,12 +130,12 @@ def __init__(
         self._internal_timeout_seconds = 5 * 60
         self._assumed_total_count = 0
         self._assumed_handled_count = 0
-        self._queue_head_dict: OrderedDictType[str, str] = OrderedDict()
+        self._queue_head_dict: OrderedDict[str, str] = OrderedDict()
         self._query_queue_head_task: asyncio.Task | None = None
         self._in_progress: set[str] = set()
         self._last_activity = datetime.now(timezone.utc)
-        self._recently_handled: LRUCache[bool] = LRUCache(max_length=self._RECENTLY_HANDLED_CACHE_SIZE)
-        self._requests_cache: LRUCache[dict] = LRUCache(max_length=self._MAX_CACHED_REQUESTS)
+        self._recently_handled: BoundedSet[str] = BoundedSet(max_length=self._RECENTLY_HANDLED_CACHE_SIZE)
+        self._requests_cache: LRUCache[CachedRequest] = LRUCache(max_length=self._MAX_CACHED_REQUESTS)
 
     @override
     @property
@@ -125,13 +158,17 @@ async def open(
     ) -> RequestQueue:
         from crawlee.storages._creation_management import open_storage
 
-        return await open_storage(
+        storage = await open_storage(
             storage_class=cls,
             id=id,
             name=name,
             configuration=configuration,
         )
 
+        await storage._ensure_head_is_non_empty()  # noqa: SLF001 - accessing private members from factories is OK
+
+        return storage
+
     @override
     async def drop(self, *, timeout: timedelta | None = None) -> None:
         from crawlee.storages._creation_management import remove_storage_from_cache
@@ -208,7 +245,7 @@ async def add_request(
             not is_handled
             and not was_already_present
             and request_id not in self._in_progress
-            and self._recently_handled.get(request_id) is None
+            and request_id not in self._recently_handled
         ):
             self._assumed_total_count += 1
             self._maybe_add_request_to_queue_head(request_id, forefront=forefront)
@@ -284,7 +321,7 @@ async def fetch_next_request(self) -> Request | None:
         Returns:
             The request or `None` if there are no more pending requests.
         """
-        await self.ensure_head_is_non_empty()
+        await self._ensure_head_is_non_empty()
 
         # We are likely done at this point.
         if len(self._queue_head_dict) == 0:
@@ -293,7 +330,7 @@ async def fetch_next_request(self) -> Request | None:
         next_request_id, _ = self._queue_head_dict.popitem(last=False)  # ~removeFirst()
 
         # This should never happen, but...
-        if next_request_id in self._in_progress or self._recently_handled.get(next_request_id):
+        if next_request_id in self._in_progress or next_request_id in self._recently_handled:
             logger.warning(
                 'Queue head returned a request that is already in progress?!',
                 extra={
@@ -343,7 +380,7 @@ async def fetch_next_request(self) -> Request | None:
                 'Request fetched from the beginning of queue was already handled',
                 extra={'nextRequestId': next_request_id},
             )
-            self._recently_handled[next_request_id] = True
+            self._recently_handled.add(next_request_id)
             return None
 
         return request
@@ -372,7 +409,7 @@ async def mark_request_as_handled(self, request: Request) -> ProcessedRequest |
         processed_request.unique_key = request.unique_key
 
         self._in_progress.remove(request.id)
-        self._recently_handled[request.id] = True
+        self._recently_handled.add(request.id)
 
         if not processed_request.was_already_handled:
             self._assumed_handled_count += 1
@@ -431,7 +468,7 @@ async def is_empty(self) -> bool:
         Returns:
             bool: `True` if the next call to `RequestQueue.fetchNextRequest` would return `None`, otherwise `False`.
         """
-        await self.ensure_head_is_non_empty()
+        await self._ensure_head_is_non_empty()
         return len(self._queue_head_dict) == 0
 
     async def is_finished(self) -> bool:
@@ -457,7 +494,7 @@ async def is_finished(self) -> bool:
 
         # TODO: set ensure_consistency to True once the following issue is resolved:
         # https://github.com/apify/crawlee-python/issues/203
-        is_head_consistent = await self.ensure_head_is_non_empty(ensure_consistency=False)
+        is_head_consistent = await self._ensure_head_is_non_empty(ensure_consistency=False)
         return is_head_consistent and len(self._queue_head_dict) == 0 and self._in_progress_count() == 0
 
     async def get_info(self) -> RequestQueueMetadata | None:
@@ -472,7 +509,7 @@ async def get_handled_count(self) -> int:
     async def get_total_count(self) -> int:
         return self._assumed_total_count
 
-    async def ensure_head_is_non_empty(
+    async def _ensure_head_is_non_empty(
         self,
         *,
         ensure_consistency: bool = False,
@@ -556,7 +593,7 @@ async def ensure_head_is_non_empty(
             logger.info(f'Waiting for {delay_seconds} for queue finalization, to ensure data consistency.')
             await asyncio.sleep(delay_seconds)
 
-        return await self.ensure_head_is_non_empty(
+        return await self._ensure_head_is_non_empty(
             ensure_consistency=ensure_consistency,
             limit=next_limit,
             iteration=iteration + 1,
@@ -578,8 +615,6 @@ def _reset(self) -> None:
     def _cache_request(self, cache_key: str, processed_request: ProcessedRequest) -> None:
         self._requests_cache[cache_key] = {
             'id': processed_request.id,
-            'is_handled': processed_request.was_already_handled,
-            'unique_key': processed_request.unique_key,
             'was_already_handled': processed_request.was_already_handled,
         }
 
@@ -595,7 +630,7 @@ async def _queue_query_head(self, limit: int) -> RequestQueueHeadState:
                 not request.id
                 or not request.unique_key
                 or request.id in self._in_progress
-                or self._recently_handled.get(request.id)
+                or request.id in self._recently_handled
             ):
                 continue
 

tests/unit/conftest.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 import os
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable
 
 import pytest
 from proxy import Proxy
@@ -20,8 +20,35 @@
     from pathlib import Path
 
 
+@pytest.fixture()
+def reset_globals(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Callable[[], None]:
+    def reset() -> None:
+        # Set the environment variable for the local storage directory to the temporary path
+        monkeypatch.setenv('CRAWLEE_STORAGE_DIR', str(tmp_path))
+
+        # Reset the local and cloud clients in StorageClientManager
+        StorageClientManager._local_client = MemoryStorageClient()
+        StorageClientManager._cloud_client = None
+
+        # Remove global configuration instance - it may contain settings adjusted by a previous test
+        Configuration._default_instance = None
+
+        # Clear creation-related caches to ensure no state is carried over between tests
+        monkeypatch.setattr(_creation_management, '_cache_dataset_by_id', {})
+        monkeypatch.setattr(_creation_management, '_cache_dataset_by_name', {})
+        monkeypatch.setattr(_creation_management, '_cache_kvs_by_id', {})
+        monkeypatch.setattr(_creation_management, '_cache_kvs_by_name', {})
+        monkeypatch.setattr(_creation_management, '_cache_rq_by_id', {})
+        monkeypatch.setattr(_creation_management, '_cache_rq_by_name', {})
+
+        # Verify that the environment variable is set correctly
+        assert os.environ.get('CRAWLEE_STORAGE_DIR') == str(tmp_path)
+
+    return reset
+
+
 @pytest.fixture(autouse=True)
-def _isolate_test_environment(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
+def _isolate_test_environment(reset_globals: Callable[[], None]) -> None:
     """Isolate tests by resetting the storage clients, clearing caches, and setting the environment variables.
 
     The fixture is applied automatically to all test cases.
@@ -30,26 +57,8 @@ def _isolate_test_environment(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -
         monkeypatch: Test utility provided by pytest.
         tmp_path: A unique temporary directory path provided by pytest for test isolation.
     """
-    # Set the environment variable for the local storage directory to the temporary path
-    monkeypatch.setenv('CRAWLEE_STORAGE_DIR', str(tmp_path))
-
-    # Reset the local and cloud clients in StorageClientManager
-    StorageClientManager._local_client = MemoryStorageClient()
-    StorageClientManager._cloud_client = None
-
-    # Remove global configuration instance - it may contain settings adjusted by a previous test
-    Configuration._default_instance = None
-
-    # Clear creation-related caches to ensure no state is carried over between tests
-    monkeypatch.setattr(_creation_management, '_cache_dataset_by_id', {})
-    monkeypatch.setattr(_creation_management, '_cache_dataset_by_name', {})
-    monkeypatch.setattr(_creation_management, '_cache_kvs_by_id', {})
-    monkeypatch.setattr(_creation_management, '_cache_kvs_by_name', {})
-    monkeypatch.setattr(_creation_management, '_cache_rq_by_id', {})
-    monkeypatch.setattr(_creation_management, '_cache_rq_by_name', {})
-
-    # Verify that the environment variable is set correctly
-    assert os.environ.get('CRAWLEE_STORAGE_DIR') == str(tmp_path)
+
+    reset_globals()
 
 
 @pytest.fixture()

tests/unit/memory_storage_client/test_memory_storage_e2e.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import Callable
+
+import pytest
+
+from crawlee.models import Request
+from crawlee.storage_client_manager import StorageClientManager
+from crawlee.storages.key_value_store import KeyValueStore
+from crawlee.storages.request_queue import RequestQueue
+
+
+@pytest.mark.parametrize('purge_on_start', [True, False])
+async def test_actor_memory_storage_client_key_value_store_e2e(
+    monkeypatch: pytest.MonkeyPatch,
+    purge_on_start: bool,  # noqa: FBT001
+    reset_globals: Callable[[], None],
+) -> None:
+    """This test simulates two clean runs using memory storage.
+    The second run attempts to access data created by the first one.
+    We run 2 configurations with different `purge_on_start`."""
+    # Configure purging env var
+    monkeypatch.setenv('CRAWLEE_PURGE_ON_START', f'{int(purge_on_start)}')
+    # Store old storage client so we have the object reference for comparison
+    old_client = StorageClientManager.get_storage_client()
+
+    old_default_kvs = await KeyValueStore.open()
+    old_non_default_kvs = await KeyValueStore.open(name='non-default')
+    # Create data in default and non-default key-value store
+    await old_default_kvs.set_value('test', 'default value')
+    await old_non_default_kvs.set_value('test', 'non-default value')
+
+    # We simulate another clean run, we expect the memory storage to read from the local data directory
+    # Default storages are purged based on purge_on_start parameter.
+    reset_globals()
+
+    # Check if we're using a different memory storage instance
+    assert old_client is not StorageClientManager.get_storage_client()
+    default_kvs = await KeyValueStore.open()
+    assert default_kvs is not old_default_kvs
+    non_default_kvs = await KeyValueStore.open(name='non-default')
+    assert non_default_kvs is not old_non_default_kvs
+    default_value = await default_kvs.get_value('test')
+
+    if purge_on_start:
+        assert default_value is None
+    else:
+        assert default_value == 'default value'
+
+    assert await non_default_kvs.get_value('test') == 'non-default value'
+
+
+@pytest.mark.parametrize('purge_on_start', [True, False])
+async def test_actor_memory_storage_client_request_queue_e2e(
+    monkeypatch: pytest.MonkeyPatch,
+    purge_on_start: bool,  # noqa: FBT001
+    reset_globals: Callable[[], None],
+) -> None:
+    """This test simulates two clean runs using memory storage.
+    The second run attempts to access data created by the first one.
+    We run 2 configurations with different `purge_on_start`."""
+    # Configure purging env var
+    monkeypatch.setenv('CRAWLEE_PURGE_ON_START', f'{int(purge_on_start)}')
+
+    # Add some requests to the default queue
+    default_queue = await RequestQueue.open()
+    for i in range(6):
+        # [0, 3] <- nothing special
+        # [1, 4] <- forefront=True
+        # [2, 5] <- handled=True
+        request_url = f'http://example.com/{i}'
+        forefront = i % 3 == 1
+        was_handled = i % 3 == 2
+        await default_queue.add_request(
+            Request.from_url(
+                unique_key=str(i),
+                url=request_url,
+                handled_at=datetime.now(timezone.utc) if was_handled else None,
+            ),
+            forefront=forefront,
+        )
+
+    # We simulate another clean run, we expect the memory storage to read from the local data directory
+    # Default storages are purged based on purge_on_start parameter.
+    reset_globals()
+
+    # Add some more requests to the default queue
+    default_queue = await RequestQueue.open()
+    for i in range(6, 12):
+        # [6,  9] <- nothing special
+        # [7, 10] <- forefront=True
+        # [8, 11] <- handled=True
+        request_url = f'http://example.com/{i}'
+        forefront = i % 3 == 1
+        was_handled = i % 3 == 2
+        await default_queue.add_request(
+            Request.from_url(
+                unique_key=str(i),
+                url=request_url,
+                handled_at=datetime.now(timezone.utc) if was_handled else None,
+            ),
+            forefront=forefront,
+        )
+
+    queue_info = await default_queue.get_info()
+    assert queue_info is not None
+
+    # If the queue was purged between the runs, only the requests from the second run should be present,
+    # in the right order
+    if purge_on_start:
+        assert queue_info.total_request_count == 6
+        assert queue_info.handled_request_count == 2
+
+        expected_pending_request_order = [10, 7, 6, 9]
+    # If the queue was NOT purged between the runs, all the requests should be in the queue in the right order
+    else:
+        assert queue_info.total_request_count == 12
+        assert queue_info.handled_request_count == 4
+
+        expected_pending_request_order = [10, 7, 4, 1, 0, 3, 6, 9]
+
+    actual_requests = list[Request]()
+    while req := await default_queue.fetch_next_request():
+        actual_requests.append(req)
+
+    assert [int(req.unique_key) for req in actual_requests] == expected_pending_request_order
+    assert [req.url for req in actual_requests] == [f'http://example.com/{req.unique_key}' for req in actual_requests]