fix(search): give each provider its own slice of the search budget (#1537)

Fixes #1531.

All Direct Download providers share the one search budget
(`RELEASE_SEARCH_TIMEOUT`, 300s). Anna's Archive runs first, and when
its bypass keeps failing the retries use up the whole budget. OceanofPDF
then sees an expired deadline, skips its queries, and the search raises
Anna's Archive's error even though OceanofPDF never got a turn.

I didn't split the budget evenly, since the deadline module itself says
a cold-start Anna's Archive solve can legitimately take 60 to 120s and a
150s or 37s slice would break setups that work today. Instead
`search_deadline.unit_deadline(waiting=n)` opens a child deadline for
one unit of work. It gets the parent's remaining time minus 30s for each
unit still waiting, never less than 30s and never more than the parent
has left. A single provider still gets everything, a fast provider's
unused time rolls on to the next, and a stuck one stops at its own limit
with the reserve intact.

It's used in the provider loop in `DirectDownloadSource.search`, and in
the source loop in `/api/releases` since it's the same shape and four
lines. With several sources enabled the first one now gets 30s less per
source behind it, which is still well above a cold solve for the usual
two or three.

Tests: a stuck first provider no longer starves the second (fails on
main, passes here), a lone provider gets the full budget, and a child
never outlives its parent. Two existing tests in
`test_releases_api_budget.py` tripped the deadline in force from inside
the source, which is now the child, so they use a short request budget
instead. ruff, basedpyright on the changed files and the non-e2e suite
are clean.

OceanofPDF is inactive on my own instance (no mirror URL set), so I
haven't reproduced this live. It's verified by reading the code and by
the tests. I had Claude do the pass on this one.

shelfmark/core/search_deadline.py

@@ -39,6 +39,12 @@ logger = setup_logger(__name__)
 DEFAULT_SEARCH_BUDGET_SECONDS = 300.0
 
 _MIN_SEARCH_BUDGET_SECONDS = 30.0
+
+# What each unit still waiting its turn keeps back from the one running now, and the
+# least a unit is given while the parent has that much left. Thirty seconds covers a
+# plain search; it is not enough for a bypass solve, which is why the unit running
+# keeps everything else.
+DEFAULT_UNIT_RESERVE_SECONDS = 30.0
 _MAX_SEARCH_BUDGET_SECONDS = 1800.0
 
 # Raised to the caller when the budget runs out, so the API can say so plainly.
@@ -110,6 +116,28 @@ def search_deadline(budget: float | None = None) -> Iterator[SearchDeadline]:
         _current.reset(token)
 
 
+@contextmanager
+def unit_deadline(
+    waiting: int = 0, *, reserve: float = DEFAULT_UNIT_RESERVE_SECONDS
+) -> Iterator[SearchDeadline | None]:
+    """Budget for one unit of work inside a search, keeping time back for those after it.
+
+    The unit gets the parent's remaining time minus `reserve` for each of the `waiting`
+    units still to run, but never less than `reserve` and never more than the parent has
+    left. A lone unit gets everything, and time a fast unit does not use rolls on to the
+    next. A stuck unit stops at its own limit, so it cannot spend the seconds the next
+    one needs. Outside a search there is no budget to divide, so nothing is applied.
+    """
+    parent = _current.get()
+    if parent is None:
+        yield None
+        return
+    remaining = parent.remaining
+    budget = min(remaining, max(reserve, remaining - reserve * max(waiting, 0)))
+    with search_deadline(budget) as child:
+        yield child
+
+
 def current() -> SearchDeadline | None:
     """The budget in force, or None outside a search."""
     return _current.get()

shelfmark/main.py

@@ -3423,13 +3423,14 @@ def api_releases() -> Response | tuple[Response, int]:
         warmup.note_user_search()
 
         with search_deadline.search_deadline():
-            for source_name in sources_to_search:
+            for index, source_name in enumerate(sources_to_search):
                 if search_deadline.expired():
                     logger.warning("Release search budget spent; %s not searched", source_name)
                     errors.append(f"{source_name}: {search_deadline.deadline_message()}")
                     continue
 
-                source, releases, error = _search_source_releases(source_name, book)
+                with search_deadline.unit_deadline(waiting=len(sources_to_search) - index - 1):
+                    source, releases, error = _search_source_releases(source_name, book)
                 if source is not None:
                     source_instances[source_name] = source
                     all_releases.extend(releases)

shelfmark/release_sources/direct_download/source.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, ClassVar
 
 import requests
 
+from shelfmark.core import search_deadline
 from shelfmark.core.logger import setup_logger
 from shelfmark.core.utils import get_aa_content_type_dir
 from shelfmark.core.utils import is_audiobook as check_audiobook
@@ -203,14 +204,18 @@ class DirectDownloadSource(ReleaseSource):
 
         releases: list[Release] = []
         failures: list[Exception] = []
-        for provider in registry.enabled_providers(self._providers):
+        providers = registry.enabled_providers(self._providers)
+        for index, provider in enumerate(providers):
             try:
-                records = provider.search(
-                    book,
-                    plan,
-                    expand_search=expand_search,
-                    content_type=content_type,
-                )
+                # Each provider gets its own slice of the search budget, so a bypass that
+                # burns its whole share on the first one still leaves the next a chance.
+                with search_deadline.unit_deadline(waiting=len(providers) - index - 1):
+                    records = provider.search(
+                        book,
+                        plan,
+                        expand_search=expand_search,
+                        content_type=content_type,
+                    )
             except SourceUnavailableError as exc:
                 failures.append(exc)
                 continue

tests/core/test_releases_api_budget.py

@@ -78,7 +78,19 @@ def test_a_search_runs_under_a_budget(client, main_module):
     assert observed and observed[0] is not None, "no budget was in force during the search"
 
 
-def test_the_budget_is_shared_across_sources(client, main_module):
+def _run_the_clock_out() -> None:
+    """Stand in for a source that waits on the cancel flag until the whole budget is gone."""
+    event = search_deadline.cancel_event()
+    assert event is not None
+    event.wait(timeout=5)
+
+
+@pytest.fixture
+def short_budget(monkeypatch):
+    monkeypatch.setattr(search_deadline, "budget_seconds", lambda: 0.3)
+
+
+def test_the_budget_is_shared_across_sources(client, main_module, short_budget):
     """A stuck first source must not spend the whole request on its own."""
     _authenticate(client)
     searched: list[str] = []
@@ -86,9 +98,7 @@ def test_the_budget_is_shared_across_sources(client, main_module):
     def _search(book, _plan):
         searched.append(book.title)
         # Whatever the first source did, it ran the clock out.
-        deadline = search_deadline.current()
-        if deadline is not None:
-            deadline.event.set()
+        _run_the_clock_out()
         return []
 
     sources = [
@@ -101,7 +111,7 @@ def test_the_budget_is_shared_across_sources(client, main_module):
     assert "ran out of time" in resp.get_json()["error"]
 
 
-def test_the_failure_carries_a_message_the_frontend_will_show(client, main_module):
+def test_the_failure_carries_a_message_the_frontend_will_show(client, main_module, short_budget):
     """The whole point of the budget.
 
     Shelfmark answers 503 when a search comes back empty with errors, and the frontend
@@ -113,9 +123,7 @@ def test_the_failure_carries_a_message_the_frontend_will_show(client, main_modul
     _authenticate(client)
 
     def _search(_book, _plan):
-        deadline = search_deadline.current()
-        if deadline is not None:
-            deadline.event.set()
+        _run_the_clock_out()
         return []
 
     sources = [

tests/core/test_search_deadline.py

@@ -255,3 +255,60 @@ def test_expired_budget_reports_the_challenge_not_a_cancellation(monkeypatch):
     assert failures, "a give-up reason should have been recorded"
     assert "ran out of time" in failures[-1]
     assert "cancelled" not in failures[-1]
+
+
+# --------------------------------------------------------------------------- #
+# Slicing the budget between providers (issue #1531)
+# --------------------------------------------------------------------------- #
+def test_unit_deadline_is_a_no_op_outside_a_search():
+    with search_deadline.unit_deadline(waiting=2) as unit:
+        assert unit is None
+        assert search_deadline.current() is None
+
+
+def test_a_lone_unit_gets_the_whole_budget():
+    with search_deadline.search_deadline(300) as parent, search_deadline.unit_deadline() as unit:
+        assert unit is not None
+        assert unit.budget_seconds == pytest.approx(parent.remaining, abs=1)
+        assert unit.budget_seconds > 299
+
+
+def test_waiting_units_each_keep_a_reserve():
+    with search_deadline.search_deadline(300), search_deadline.unit_deadline(waiting=2) as unit:
+        assert unit is not None
+        assert unit.budget_seconds == pytest.approx(240, abs=1)
+
+
+def test_a_unit_is_never_given_less_than_the_reserve_while_the_parent_has_it():
+    with search_deadline.search_deadline(60), search_deadline.unit_deadline(waiting=5) as unit:
+        assert unit is not None
+        assert unit.budget_seconds == pytest.approx(30, abs=1)
+
+
+def test_a_unit_never_outlives_its_parent():
+    with (
+        search_deadline.search_deadline(10) as parent,
+        search_deadline.unit_deadline(waiting=3) as unit,
+    ):
+        assert unit is not None
+        assert unit.expires_at <= parent.expires_at + 0.01
+
+
+def test_the_parent_is_back_in_force_after_a_unit():
+    with search_deadline.search_deadline(300) as parent:
+        with search_deadline.unit_deadline(waiting=1):
+            pass
+        assert search_deadline.current() is parent
+
+
+def test_a_stuck_unit_stops_at_its_slice_and_the_next_still_runs(monkeypatch):
+    """The unit's own limit trips while the parent still has the reserve left."""
+    monkeypatch.setattr(search_deadline, "DEFAULT_UNIT_RESERVE_SECONDS", 0.2)
+    with search_deadline.search_deadline(1.0):
+        with search_deadline.unit_deadline(waiting=1, reserve=0.5) as first:
+            assert first is not None
+            assert first.event.wait(timeout=5) is True
+            assert search_deadline.expired() is True
+        assert search_deadline.expired() is False
+        with search_deadline.unit_deadline(waiting=0):
+            assert search_deadline.expired() is False

tests/direct_download/test_provider_budget.py

@@ -0,0 +1,77 @@
+"""A stuck first provider must not starve the second (issue #1531).
+
+All providers share one search budget. Anna's Archive runs first, and its bypass retries
+could spend the whole of it, so OceanofPDF saw an expired budget and never searched.
+"""
+
+import pytest
+
+from shelfmark.core import search_deadline
+from shelfmark.metadata_providers import BookMetadata
+from shelfmark.release_sources.direct_download import registry
+from shelfmark.release_sources.direct_download.source import DirectDownloadSource
+
+
+class _Provider:
+    def __init__(self, name, behaviour):
+        self.name = name
+        self.display_name = name
+        self._behaviour = behaviour
+        self.expired_on_entry = None
+
+    def is_enabled(self):
+        return True
+
+    def search(self, book, plan, *, expand_search=False, content_type="ebook"):
+        self.expired_on_entry = search_deadline.expired()
+        return self._behaviour(self)
+
+
+def _stuck(_provider):
+    """Wait on the cancel flag the way a bypass does, then fail like Anna's Archive."""
+    event = search_deadline.cancel_event()
+    assert event is not None
+    event.wait(timeout=5)
+    raise RuntimeError("bypass failed")
+
+
+def _search(providers, monkeypatch):
+    monkeypatch.setattr(registry, "enabled_providers", lambda p=None: tuple(p))
+    monkeypatch.setattr(registry, "get_unavailable_reason", lambda p=None: None)
+    monkeypatch.setattr(
+        "shelfmark.release_sources.direct_download.source._browse_record_to_release",
+        lambda record: record,
+    )
+    source = DirectDownloadSource.__new__(DirectDownloadSource)
+    source._providers = providers
+    book = BookMetadata(provider="x", provider_id="1", title="T", authors=["A"])
+    return source.search(book, None)  # pyright: ignore[reportArgumentType]
+
+
+def test_a_stuck_first_provider_leaves_the_second_a_chance(monkeypatch):
+    monkeypatch.setattr(search_deadline, "DEFAULT_UNIT_RESERVE_SECONDS", 0.3)
+    original = search_deadline.unit_deadline
+
+    def short_reserve(waiting=0, *, reserve=0.5):
+        return original(waiting, reserve=reserve)
+
+    monkeypatch.setattr(search_deadline, "unit_deadline", short_reserve)
+    first = _Provider("aa", _stuck)
+    second = _Provider("oceanofpdf", lambda _p: ["found"])
+    with search_deadline.search_deadline(1.5):
+        releases = _search([first, second], monkeypatch)
+    assert second.expired_on_entry is False
+    assert releases == ["found"]
+
+
+def test_a_single_provider_still_gets_the_whole_budget(monkeypatch):
+    seen = {}
+
+    def record(_p):
+        seen["budget"] = search_deadline.current().budget_seconds
+        return ["found"]
+
+    only = _Provider("aa", record)
+    with search_deadline.search_deadline(300) as parent:
+        _search([only], monkeypatch)
+        assert seen["budget"] == pytest.approx(parent.budget_seconds, abs=1)