From ed51f5be7e07d2a6ebeb875406fbddfbb1e34421 Mon Sep 17 00:00:00 2001 From: npeter83 Date: Sun, 25 Oct 2026 02:50:34 +0200 Subject: [PATCH] fix(search): the API-sourced live search takes its slot under the lock too Checked first and logged after search.list, requests sent together at the last free slot all passed, and each spent 100 units of the shared budget. Every live search now claims its slot (quota.claim_action) whatever the source, and gives it back when it showed nothing. The cap counts the zero-unit claims only, so search.list's own spend event is accounting, not a second slot. log_action, the unlocked writer, goes. Tests: the API source reads one page per request, is refused before the call when the budget cannot pay, counts one slot per search and refunds a failed one; the lock test now proves the count happens after the lock. --- backend/app/quota.py | 31 +++----- backend/app/routes/search.py | 38 ++++------ backend/tests/test_owner_scoping.py | 6 +- backend/tests/test_quota_ownership.py | 6 +- backend/tests/test_related_route.py | 5 +- backend/tests/test_search_cap.py | 102 ++++++++++++++++++++++++-- 6 files changed, 133 insertions(+), 55 deletions(-) diff --git a/backend/app/quota.py b/backend/app/quota.py index 66d4fa4d..13ec2e66 100644 --- a/backend/app/quota.py +++ b/backend/app/quota.py @@ -143,7 +143,7 @@ def backfill_reserve_reached() -> bool: condition for every backfill/explore path (recent, deep, immediate full-history, explore, and rolling revalidation). - Reads in its OWN short-lived session (like record_usage / log_action), for two reasons: it + Reads in its OWN short-lived session (like record_usage / claim_action), for two reasons: it replaces the remaining/reserve comparison that was open-coded at every call site, and — critically for the per-page guards in paging loops — a guard must not hold the caller's transaction open across the network fetch that follows it. Quota is committed independently of any caller, so a @@ -152,26 +152,17 @@ def backfill_reserve_reached() -> bool: return remaining_today(qs) <= sysconfig.get_int(qs, "backfill_quota_reserve") -def log_action(user_id: int, action: str) -> None: - """Record a zero-cost action event so per-user daily caps still count it. - - `record_usage` only logs when units are actually charged; a scrape-based search spends no - quota yet must still be rate-limited per user, so it logs its event here directly. - - Writes in its OWN session (like record_usage) — quota accounting is durable independent of the - caller's transaction and must never commit the caller's session (that would flush a job's - partial state mid-run, defeating its own rollback).""" - with SessionLocal() as qs: - qs.add(QuotaEvent(user_id=user_id, action=action, units=0)) - qs.commit() - - def claim_action(user_id: int, action: str, limit: int) -> int | None: """Take one of this user's `limit` daily `action` slots, or None when they are all taken — the check and the take in one transaction, under a lock per user and action. Checked first and logged after the work, N requests sent together at the last free slot all passed the check (the work in between is seconds of scraping). Returns the event's id, for `refund_action` when - the work then produced nothing. Own session, as `log_action`.""" + the work then produced nothing. + + The ONE way to count a capped action: a claim is a zero-unit event, and `actions_today` counts + those alone, so a call the action then pays for (search.list's own `record_usage` event, under + the same action) is accounting, never a second slot. Own session, committed at once — quota + accounting is durable whatever the caller's transaction does next, and must never commit it.""" with SessionLocal() as qs: qs.execute(select(func.pg_advisory_xact_lock(func.hashtext(action), user_id))) if actions_today(qs, user_id, action) >= limit: @@ -191,9 +182,10 @@ def refund_action(event_id: int) -> None: def actions_today(db: Session, user_id: int, action: str) -> int: - """How many quota events of `action` this user has logged so far in the current Pacific - day — for per-user, per-action daily caps (e.g. the live-search limit). Counts events, - not units, so it only works for actions charged exactly once per user action.""" + """How many slots of `action` this user has claimed so far in the current Pacific day — for + per-user, per-action daily caps (the live search, the related lists). Counts the zero-unit + claim events (`claim_action`) only: `record_usage`'s events for units the action spent are not + slots (it logs only when units were charged).""" return ( db.scalar( select(func.count()) @@ -201,6 +193,7 @@ def actions_today(db: Session, user_id: int, action: str) -> int: .where( owned_by(QuotaEvent.user_id, user_id), QuotaEvent.action == action, + QuotaEvent.units == 0, QuotaEvent.created_at >= pacific_day_start_utc(), ) ) diff --git a/backend/app/routes/search.py b/backend/app/routes/search.py index 06741528..91a25442 100644 --- a/backend/app/routes/search.py +++ b/backend/app/routes/search.py @@ -142,27 +142,20 @@ def search_youtube( daily_limit = sysconfig.get_int(db, "search_daily_limit_per_user") if daily_limit <= 0: raise refused(409, "search_disabled") - claimed: int | None = None - if source == "api": - # The one search.list call logs this action's event itself (record_usage, attributed - # below), so here the cap is only checked: a claimed slot would count it twice. - if quota.actions_today(db, user.id, quota.QuotaAction.VIDEOS_SEARCH) >= daily_limit: - _cap_spent("search", daily_limit) - raise refused(429, "search_daily_limit", limit=daily_limit) - # The 100-unit budget pre-check only applies to the API source. - if not quota.can_spend(db, SEARCH_COST): - log.info( - "a live search was refused: the shared YouTube budget is spent", - extra=logs.event("search.refused", reason="shared_quota", cost=SEARCH_COST), - ) - raise refused(429, "quota_spent") - else: - # The check and the take are one step (`quota.claim_action`); a search that fails - # outright gives the slot back below. - claimed = quota.claim_action(user.id, quota.QuotaAction.VIDEOS_SEARCH, daily_limit) - if claimed is None: - _cap_spent("search", daily_limit) - raise refused(429, "search_daily_limit", limit=daily_limit) + # The 100-unit budget pre-check only applies to the API source. + if source == "api" and not quota.can_spend(db, SEARCH_COST): + log.info( + "a live search was refused: the shared YouTube budget is spent", + extra=logs.event("search.refused", reason="shared_quota", cost=SEARCH_COST), + ) + raise refused(429, "quota_spent") + # The check and the take are one step (`quota.claim_action`), whichever the source: checked + # first and logged after, requests sent together at the last free slot all passed — and with + # the API source each of them spent 100 units. A search that fails outright gives it back below. + claimed = quota.claim_action(user.id, quota.QuotaAction.VIDEOS_SEARCH, daily_limit) + if claimed is None: + _cap_spent("search", daily_limit) + raise refused(429, "search_daily_limit", limit=daily_limit) # Channels this user has blocked are dropped before ingestion — never shown, never stored. blocked = queries.blocked_ids(db, user) @@ -230,8 +223,7 @@ def search_youtube( ) except Exception: # A search that showed the user nothing does not count against them. - if claimed is not None: - quota.refund_action(claimed) + quota.refund_action(claimed) raise ordered = collected[:limit] diff --git a/backend/tests/test_owner_scoping.py b/backend/tests/test_owner_scoping.py index 8ab99547..c85a69de 100644 --- a/backend/tests/test_owner_scoping.py +++ b/backend/tests/test_owner_scoping.py @@ -75,9 +75,9 @@ def test_trimming_read_notifications_trims_only_mine(db, make_user): def test_todays_actions_count_only_mine(db, make_user): me, other = make_user(), make_user() - quota.log_action(me.id, quota.QuotaAction.VIDEOS_SEARCH) - quota.log_action(other.id, quota.QuotaAction.VIDEOS_SEARCH) - quota.log_action(other.id, quota.QuotaAction.VIDEOS_SEARCH) + quota.claim_action(me.id, quota.QuotaAction.VIDEOS_SEARCH, 5) + quota.claim_action(other.id, quota.QuotaAction.VIDEOS_SEARCH, 5) + quota.claim_action(other.id, quota.QuotaAction.VIDEOS_SEARCH, 5) assert quota.actions_today(db, me.id, quota.QuotaAction.VIDEOS_SEARCH) == 1 diff --git a/backend/tests/test_quota_ownership.py b/backend/tests/test_quota_ownership.py index a951ef98..36a269ee 100644 --- a/backend/tests/test_quota_ownership.py +++ b/backend/tests/test_quota_ownership.py @@ -44,15 +44,15 @@ def test_reads_see_separate_session_spend_mid_transaction(db): assert after - before == 7 -def test_log_action_survives_caller_rollback(db): +def test_a_claimed_slot_survives_caller_rollback(db): user = User(email="searcher@example.com") db.add(user) db.commit() # a committed user to attribute the (free) action to db.add(User(email="pending@example.com")) # uncommitted caller work - quota.log_action(user.id, quota.QuotaAction.VIDEOS_SEARCH) + assert quota.claim_action(user.id, quota.QuotaAction.VIDEOS_SEARCH, 5) is not None db.rollback() - # The zero-cost action event survived (per-user daily caps must still count it) ... + # The zero-cost claim event survived (per-user daily caps must still count it) ... events = db.execute(select(QuotaEvent)).scalars().all() assert len(events) == 1 assert events[0].user_id == user.id diff --git a/backend/tests/test_related_route.py b/backend/tests/test_related_route.py index a0f8aa29..c1665b78 100644 --- a/backend/tests/test_related_route.py +++ b/backend/tests/test_related_route.py @@ -336,10 +336,11 @@ def test_the_last_slot_is_taken_under_a_lock(db, make_user): t.start() t.join(timeout=1) assert t.is_alive() # it waits for the one holding the lock + # ...and counts only once it has the lock: the one holding it takes the last slot meanwhile. + db.add(QuotaEvent(user_id=user.id, action=action, units=0)) db.commit() t.join(timeout=10) - assert got and got[0] is not None - assert quota.claim_action(user.id, action, 1) is None # the one slot is taken + assert got == [None] def _videos_list_fails(monkeypatch, *, on_call: int = 1, reason: str = "quotaExceeded") -> None: diff --git a/backend/tests/test_search_cap.py b/backend/tests/test_search_cap.py index 447db142..f2cca60a 100644 --- a/backend/tests/test_search_cap.py +++ b/backend/tests/test_search_cap.py @@ -1,9 +1,11 @@ -"""`GET /api/search/youtube` — the per-user daily cap of the live search (scrape source). +"""`GET /api/search/youtube` — the per-user daily cap of the live search, both sources. -The scrape is replaced by pages the test hands out, and YouTube's videos.list by an offline enrich. -What is held: every request takes a slot, a cursor's too; a search that failed outright gives its -slot back; a zero limit turns live search off, which says so instead of "try again tomorrow"; and a -lookup that fails says why instead of answering "no results".""" +The scrape and search.list are replaced by pages the test hands out, and YouTube's videos.list by an +offline enrich. What is held: every request takes a slot, a cursor's too; a search that failed +outright gives its slot back; a zero limit turns live search off, which says so instead of "try +again tomorrow"; a lookup that fails says why instead of answering "no results"; and the API source +— 100 units a page — reads one page per request, is refused before the call when the shared budget +cannot pay for it, and counts one slot per search though search.list logs its own event.""" import pytest @@ -42,6 +44,19 @@ class _FakeYT: def get_channels(self, ids): return [] + # search.list: the pages the test hands out, each logged as the real client does (100 units, + # under the action the route attributes it to). + api_pages: list = [] + api_calls = 0 + + def search_videos(self, q, page_token=None, order="relevance"): + _FakeYT.api_calls += 1 + quota.record_usage(100) + page = _FakeYT.api_pages.pop(0) + if isinstance(page, Exception): + raise page + return page + def get_videos(self, ids): if _FakeYT.fail_videos: raise YouTubeError("videos -> 403", reason="quotaExceeded", path="videos") @@ -74,6 +89,8 @@ def pages(monkeypatch, db): monkeypatch.setattr(search_routes.search_scrape, "search", search) monkeypatch.setattr(search_routes, "YouTubeClient", _FakeYT) monkeypatch.setattr(_FakeYT, "fail_videos", False) + monkeypatch.setattr(_FakeYT, "api_pages", []) + monkeypatch.setattr(_FakeYT, "api_calls", 0) return state @@ -130,3 +147,78 @@ def test_a_lookup_that_fails_says_why_not_no_results(db, make_user, as_user, pag assert resp.status_code == 422 assert resp.json()["detail"] == {"code": "youtube", "reason": "quota"} assert quota.actions_today(db, user.id, ACTION) == 0 + + +def _api(db) -> None: + sysconfig.set_value(db, "search_source", "api") + sysconfig.set_value(db, "quota_daily_budget", 10000) + + +def test_the_api_source_reads_one_page_per_request(db, make_user, as_user, pages): + # Each page is 100 units: a short page is answered as it is, with its cursor, never paged on. + _api(db) + _FakeYT.api_pages = [{"items": [_item(1)], "next_page_token": "nextpage"}, _page(2)] + resp = as_user(make_user()).get("/api/search/youtube?q=cats&limit=20") + assert resp.status_code == 200 + assert _FakeYT.api_calls == 1 + assert resp.json()["next_cursor"] == "nextpage" + + +def test_the_api_source_counts_one_slot_per_search(db, make_user, as_user, pages): + # search.list logs its own event under the action (record_usage): accounting, not a second slot. + user = make_user() + _api(db) + sysconfig.set_value(db, "search_daily_limit_per_user", 2) + _FakeYT.api_pages = [_page(1), _page(2), _page(3)] + client = as_user(user) + assert client.get("/api/search/youtube?q=cats").status_code == 200 + assert quota.actions_today(db, user.id, ACTION) == 1 + assert client.get("/api/search/youtube?q=dogs").status_code == 200 + refused = client.get("/api/search/youtube?q=birds") + assert refused.status_code == 429 + assert refused.json()["detail"]["reason"] == "search_daily_limit" + assert _FakeYT.api_calls == 2 + + +def test_the_api_source_takes_its_slot_under_the_lock(db, make_user, as_user, pages, monkeypatch): + user = make_user() + _api(db) + claimed: list = [] + real = quota.claim_action + + def claim(uid, action, limit): + claimed.append(action) + return real(uid, action, limit) + + monkeypatch.setattr(search_routes.quota, "claim_action", claim) + _FakeYT.api_pages = [_page(1)] + assert as_user(user).get("/api/search/youtube?q=cats").status_code == 200 + assert claimed == [ACTION] + + +def test_the_api_source_is_refused_before_the_call_when_the_budget_cannot_pay( + db, make_user, as_user, pages +): + user = make_user() + _api(db) + sysconfig.set_value(db, "quota_daily_budget", 50) + _FakeYT.api_pages = [_page(1)] + resp = as_user(user).get("/api/search/youtube?q=cats") + assert resp.status_code == 429 + assert resp.json()["detail"]["reason"] == "quota_spent" + assert _FakeYT.api_calls == 0 + assert quota.actions_today(db, user.id, ACTION) == 0 + + +def test_an_api_search_that_failed_gives_its_slot_back(db, make_user, as_user, pages): + user = make_user() + _api(db) + sysconfig.set_value(db, "search_daily_limit_per_user", 1) + _FakeYT.api_pages = [ + YouTubeError("search -> 403", reason="quotaExceeded", path="search"), + _page(1), + ] + client = as_user(user) + assert client.get("/api/search/youtube?q=cats").status_code == 422 + assert quota.actions_today(db, user.id, ACTION) == 0 + assert client.get("/api/search/youtube?q=cats").status_code == 200