Merge 2efa0b6b43 into ad43bf0c56

scrapy/pqueues.py

@@ -153,7 +153,10 @@ class ScrapyPriorityQueue:
                 else:
                     q.close()
 
-        self.curprio = min(startprios)
+        # A recorded priority may have no queue to restore, e.g. if it only
+        # ever held a request that failed to serialize, so curprio comes from
+        # the queues that do exist and not from startprios.
+        self._update_curprio()
 
     def qfactory(self, key: int) -> QueueProtocol:
         return build_from_crawler(
@@ -185,46 +188,39 @@ class ScrapyPriorityQueue:
                 self.queues[priority] = self.qfactory(priority)
             q = self.queues[priority]
         q.push(request)  # this may fail (eg. serialization error)
-        if self.curprio is None or priority < self.curprio:
+        # A queue class may drop a request instead of storing it, and a
+        # priority with an empty queue must not become the current one, since
+        # pop() and peek() expect to find a request there.
+        if q and (self.curprio is None or priority < self.curprio):
             self.curprio = priority
 
+    def _curqueue(self) -> QueueProtocol:
+        assert self.curprio is not None
+        # Whichever dict holds a non-empty queue at the current priority. The
+        # other one may hold an empty queue left behind by a failed push, e.g.
+        # on a serialization error, which would hide the next request.
+        return self.queues.get(self.curprio) or self._start_queues[self.curprio]
+
     def pop(self) -> Request | None:
-        while self.curprio is not None:
-            try:
-                q = self.queues[self.curprio]
-            except KeyError:
-                pass
-            else:
-                m = q.pop()
-                if not q:
-                    del self.queues[self.curprio]
-                    q.close()
-                    if not self._start_queues:
-                        self._update_curprio()
-                return m
-            if self._start_queues:
-                try:
-                    q = self._start_queues[self.curprio]
-                except KeyError:
-                    self._update_curprio()
-                else:
-                    m = q.pop()
-                    if not q:
-                        del self._start_queues[self.curprio]
-                        q.close()
-                        self._update_curprio()
-                    return m
-            else:
-                self._update_curprio()
-        return None
+        if self.curprio is None:
+            return None
+        q = self._curqueue()
+        request = q.pop()
+        if not q:
+            self._update_curprio()
+        return request
 
     def _update_curprio(self) -> None:
-        prios = {
-            p
-            for queues in (self.queues, self._start_queues)
-            for p, q in queues.items()
-            if q
-        }
+        # Keeping an empty queue would hold its storage open for nothing, and
+        # make close() record a priority with nothing to restore from it.
+        prios: set[int] = set()
+        for queues in (self.queues, self._start_queues):
+            for p, q in list(queues.items()):
+                if q:
+                    prios.add(p)
+                else:
+                    del queues[p]
+                    q.close()
         self.curprio = min(prios) if prios else None
 
     def peek(self) -> Request | None:
@@ -236,10 +232,9 @@ class ScrapyPriorityQueue:
         """
         if self.curprio is None:
             return None
-        try:
-            queue = self._start_queues[self.curprio]
-        except KeyError:
-            queue = self.queues[self.curprio]
+        # Picking the queue the same way as pop() is what makes the returned
+        # request the one that pop() then returns.
+        queue = self._curqueue()
         # Protocols can't declare optional members
         return cast("Request", queue.peek())  # type: ignore[attr-defined]
 

tests/test_pqueues.py

@@ -14,6 +14,15 @@ from scrapy.utils.test import get_crawler
 from tests.utils.downloader import MockDownloader
 
 
+class DroppingFifoMemoryQueue(FifoMemoryQueue):  # type: ignore[valid-type,misc]
+    """Queue that drops requests marked with ``drop`` in their metadata, the
+    way a custom queue class could drop requests it does not want to store."""
+
+    def push(self, request):
+        if not request.meta.get("drop"):
+            super().push(request)
+
+
 class TestPriorityQueue:
     def setup_method(self):
         self.crawler = get_crawler(Spider)
@@ -76,6 +85,162 @@ class TestPriorityQueue:
         assert queue.pop().url == req3.url
         assert not queue.close()
 
+    def test_peek_after_draining_a_higher_priority_queue(self):
+        """Draining the queue of the current priority while start requests
+        remain at a different one must not leave ``curprio`` pointing at a
+        priority that no queue has, which would make ``peek()`` come up empty
+        with requests still queued."""
+        if not hasattr(queuelib.queue.FifoMemoryQueue, "peek"):
+            pytest.skip("queuelib.queue.FifoMemoryQueue.peek is undefined")
+        temp_dir = tempfile.mkdtemp()
+        queue = ScrapyPriorityQueue.from_crawler(
+            self.crawler,
+            FifoMemoryQueue,
+            temp_dir,
+            start_queue_cls=FifoMemoryQueue,
+        )
+        start_request = Request(
+            "https://example.org/start", meta={"is_start_request": True}
+        )
+        queue.push(start_request)
+        # A redirect of a start request, which REDIRECT_PRIORITY_ADJUST puts at
+        # a higher priority, i.e. in a separate, non-start queue.
+        queue.push(Request("https://example.org/redirect", priority=2))
+
+        assert queue.peek().url == "https://example.org/redirect"
+        assert queue.pop().url == "https://example.org/redirect"
+        assert queue.peek().url == start_request.url
+        assert queue.pop().url == start_request.url
+        assert queue.peek() is None
+        queue.close()
+
+    def test_peek_agrees_with_pop_on_start_requests(self):
+        """A start request and a non-start request at the same priority sit in
+        separate queues, and ``peek()`` must report the one that ``pop()``
+        returns, since a caller may peek to decide whether to pop."""
+        if not hasattr(queuelib.queue.FifoMemoryQueue, "peek"):
+            pytest.skip("queuelib.queue.FifoMemoryQueue.peek is undefined")
+        temp_dir = tempfile.mkdtemp()
+        queue = ScrapyPriorityQueue.from_crawler(
+            self.crawler,
+            FifoMemoryQueue,
+            temp_dir,
+            start_queue_cls=FifoMemoryQueue,
+        )
+        queue.push(
+            Request("https://example.org/start", meta={"is_start_request": True})
+        )
+        queue.push(Request("https://example.org/other"))
+
+        while len(queue):
+            peeked = queue.peek()
+            assert queue.pop().url == peeked.url
+        queue.close()
+
+    def test_peek_and_pop_skip_an_empty_queue_left_by_a_failed_push(self):
+        """A queue is created before the request is pushed into it, so a push
+        that fails (e.g. a serialization error) leaves an empty queue behind at
+        that priority. Neither ``peek()`` nor ``pop()`` may let it hide the
+        request that the other dict holds at the same priority."""
+        if not hasattr(queuelib.queue.FifoMemoryQueue, "peek"):
+            pytest.skip("queuelib.queue.FifoMemoryQueue.peek is undefined")
+        temp_dir = tempfile.mkdtemp()
+        queue = ScrapyPriorityQueue.from_crawler(
+            self.crawler,
+            PickleFifoDiskQueue,
+            temp_dir,
+            start_queue_cls=PickleFifoDiskQueue,
+        )
+        with pytest.raises(ValueError, match="is not an instance method"):
+            queue.push(
+                Request("https://example.org/lambda", callback=lambda response: None)
+            )
+        assert queue.queues[0] is not None  # the empty leftover
+        assert len(queue) == 0
+
+        start_request = Request(
+            "https://example.org/start", meta={"is_start_request": True}
+        )
+        queue.push(start_request)
+        assert len(queue) == 1
+        assert queue.peek().url == start_request.url
+        assert queue.pop().url == start_request.url
+        assert len(queue) == 0
+        assert queue.peek() is None
+        assert queue.pop() is None
+        queue.close()
+
+    def test_empty_queues_are_dropped_on_refresh(self):
+        """The empty leftover of a failed push is forgotten (and closed) the
+        next time the current priority is refreshed, rather than kept around
+        holding its storage open."""
+        temp_dir = tempfile.mkdtemp()
+        queue = ScrapyPriorityQueue.from_crawler(
+            self.crawler, PickleFifoDiskQueue, temp_dir
+        )
+        with pytest.raises(ValueError, match="is not an instance method"):
+            queue.push(
+                Request("https://example.org/lambda", callback=lambda response: None)
+            )
+        assert set(queue.queues) == {0}  # the empty leftover
+        # A request at a different priority, whose queue emptying is what
+        # triggers the refresh.
+        queue.push(Request("https://example.org/1", priority=1))
+        assert queue.pop().url == "https://example.org/1"
+        assert queue.queues == {}
+        assert queue.curprio is None
+        assert queue.pop() is None
+        assert not queue.close()
+
+    def test_push_into_a_queue_that_drops_the_request(self):
+        """A queue class may drop a request instead of storing it, and that
+        priority must not become the current one, or ``peek()`` and ``pop()``
+        would come up empty with requests still queued."""
+        if not hasattr(queuelib.queue.FifoMemoryQueue, "peek"):
+            pytest.skip("queuelib.queue.FifoMemoryQueue.peek is undefined")
+        temp_dir = tempfile.mkdtemp()
+        queue = ScrapyPriorityQueue.from_crawler(
+            self.crawler, DroppingFifoMemoryQueue, temp_dir
+        )
+        # A higher priority than the request below, so that it would become the
+        # current priority if it were stored.
+        queue.push(
+            Request("https://example.org/dropped", priority=1, meta={"drop": True})
+        )
+        kept = Request("https://example.org/kept")
+        queue.push(kept)
+
+        assert len(queue) == 1
+        assert queue.peek().url == kept.url
+        assert queue.pop().url == kept.url
+        assert queue.peek() is None
+        assert queue.pop() is None
+        assert not queue.close()
+
+    def test_init_prios_without_a_restorable_queue(self):
+        """A priority recorded on close may have nothing to restore, e.g. if
+        its only request failed to serialize. ``curprio`` must not point at it,
+        or ``peek()`` raises :exc:`KeyError` on the resumed crawl."""
+        temp_dir = tempfile.mkdtemp()
+        queue = ScrapyPriorityQueue.from_crawler(
+            self.crawler, PickleFifoDiskQueue, temp_dir
+        )
+        with pytest.raises(ValueError, match="is not an instance method"):
+            queue.push(
+                Request("https://example.org/lambda", callback=lambda response: None)
+            )
+        startprios = queue.close()
+        assert startprios == [0]
+
+        queue2 = ScrapyPriorityQueue.from_crawler(
+            self.crawler, PickleFifoDiskQueue, temp_dir, startprios
+        )
+        assert len(queue2) == 0
+        assert queue2.curprio is None
+        assert queue2.peek() is None
+        assert queue2.pop() is None
+        queue2.close()
+
     def test_init_prios_with_start_queue(self):
         temp_dir = tempfile.mkdtemp()
         queue = build_from_crawler(