diff --git a/shelfmark/core/search_deadline.py b/shelfmark/core/search_deadline.py index fb66aeed..9b44799a 100644 --- a/shelfmark/core/search_deadline.py +++ b/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() diff --git a/shelfmark/main.py b/shelfmark/main.py index ec99d5af..f8076627 100644 --- a/shelfmark/main.py +++ b/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) diff --git a/shelfmark/release_sources/direct_download/source.py b/shelfmark/release_sources/direct_download/source.py index ee4acfa2..a633491e 100644 --- a/shelfmark/release_sources/direct_download/source.py +++ b/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 diff --git a/tests/core/test_releases_api_budget.py b/tests/core/test_releases_api_budget.py index 526f539c..f9019a03 100644 --- a/tests/core/test_releases_api_budget.py +++ b/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 = [ diff --git a/tests/core/test_search_deadline.py b/tests/core/test_search_deadline.py index 6cbf4451..6191e71b 100644 --- a/tests/core/test_search_deadline.py +++ b/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 diff --git a/tests/direct_download/test_provider_budget.py b/tests/direct_download/test_provider_budget.py new file mode 100644 index 00000000..1a079cb1 --- /dev/null +++ b/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)