Align peek and pop

This commit is contained in:
Adrian Chaves 2026-07-31 04:57:40 +02:00
parent 7436afc95f
commit f5d41283c3
2 changed files with 165 additions and 36 deletions

View File

@ -151,7 +151,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(
@ -188,41 +191,37 @@ class ScrapyPriorityQueue:
def pop(self) -> Request | None:
while self.curprio is not None:
try:
q = self.queues[self.curprio]
except KeyError:
pass
else:
for queues in (self.queues, self._start_queues):
q = queues.get(self.curprio)
# An empty queue can linger at a priority when a push failed
# after creating it, e.g. on a serialization error. Popping
# from it would return None and hide the request that the other
# dict may hold at the same priority.
if not q:
continue
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:
# The other dict may have no queue at this priority either,
# and a curprio that neither dict has would make peek() come
# up empty.
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 m
# Nothing to pop at this priority: refreshing drops the empty
# leftovers and moves on to the next priority.
self._update_curprio()
return None
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:
@ -234,12 +233,17 @@ class ScrapyPriorityQueue:
"""
if self.curprio is None:
return None
try:
queue = self._start_queues[self.curprio]
except KeyError:
queue = self.queues[self.curprio]
# Protocols can't declare optional members
return cast("Request", queue.peek()) # type: ignore[attr-defined]
# The dicts are walked in the same order as in pop(), which is what
# makes the returned request the one that pop() then returns.
for queues in (self.queues, self._start_queues):
queue = queues.get(self.curprio)
# Empty queues can linger at a priority (see pop()), where they
# would hide the request that the other dict may hold at the same
# priority.
if queue:
# Protocols can't declare optional members
return cast("Request", queue.peek()) # type: ignore[attr-defined]
return None
def close(self) -> list[int]:
active: set[int] = set()

View File

@ -76,6 +76,131 @@ 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."""
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."""
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."""
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_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 = ScrapyPriorityQueue.from_crawler(