From e4102dcf45176321a30a0fc63090dd28114f361c Mon Sep 17 00:00:00 2001 From: Xander-AJ Date: Fri, 3 Jul 2026 18:28:04 +0300 Subject: [PATCH] feat: rewrite fair claim as LATERAL top-K per company --- README.md | 9 ++- app/config.py | 4 ++ app/services/worker_service.py | 63 ++++++++++++------- scripts/perf_explain.sql | 111 ++++++++++++++++++++++----------- scripts/perf_explain_old.sql | 105 +++++++++++++++++++++++++++++++ tests/test_lateral_claim.py | 94 ++++++++++++++++++++++++++++ 6 files changed, 327 insertions(+), 59 deletions(-) create mode 100644 scripts/perf_explain_old.sql create mode 100644 tests/test_lateral_claim.py diff --git a/README.md b/README.md index acb337e..15f251d 100644 --- a/README.md +++ b/README.md @@ -74,9 +74,12 @@ Trigger a processing pass: - Failed tasks retry with full-jitter exponential backoff (base 2s, ×2 per attempt, capped at 5min) via a `run_after` gate, so a poison task can't burn its retries in seconds. Stale-lock recovery is separate and unconditional. -- Workers claim tasks round-robin across companies (ROW_NUMBER partitioned by - company_id), so a tenant with a large backlog can't starve others. A single - CTE ranks and locks in one statement via FOR UPDATE SKIP LOCKED. +- Workers claim tasks round-robin across companies, so a tenant with a large + backlog can't starve others. Candidates come from a LATERAL top-K per eligible + company (the query the `run_after` partial index was built for), then fairness + ordering + `FOR UPDATE SKIP LOCKED` in one statement. Sorting only + companies x cap rows instead of the whole backlog keeps claims flat as the + backlog grows. ## Bonus features diff --git a/app/config.py b/app/config.py index 51232f4..c91cd27 100644 --- a/app/config.py +++ b/app/config.py @@ -22,6 +22,10 @@ class Settings(BaseSettings): retry_backoff_factor: float = 2.0 retry_backoff_cap_seconds: float = 300.0 + # Max tasks a single company contributes to one claim batch before fairness + # ordering. Bounds the per-company LATERAL subquery's work in claim_tasks(). + per_company_claim_cap: int = 10 + # Port the worker exposes its own Prometheus metrics on (each replica is its # own scrape target). worker_metrics_port: int = 9100 diff --git a/app/services/worker_service.py b/app/services/worker_service.py index d240769..435239a 100644 --- a/app/services/worker_service.py +++ b/app/services/worker_service.py @@ -52,24 +52,30 @@ def claim_tasks(db: Session, batch_size: int = 10) -> list[Task]: # Round-robin fairness across companies. A naive `ORDER BY created_at` claims # the globally-oldest rows, so one tenant with a big backlog starves everyone - # else. Instead we rank each company's eligible tasks with ROW_NUMBER (rn=1 is - # that company's oldest) and order the batch by (rn, created_at): every - # company's oldest task is claimed before any company's second. `rn <= - # batch_size` caps how many rows a single flooder can contribute to the - # candidate pool. + # else. We still order the batch so every company's oldest eligible task is + # claimed before any company's second (rn, then created_at), but we build the + # candidate pool with a LATERAL top-K per company instead of a global window. # - # This is raw SQL because a window function combined with FOR UPDATE SKIP - # LOCKED can't be expressed cleanly through the ORM, and the locking - # semantics (lock only the tasks rows `t`, skip already-locked ones) read far - # more clearly written out. `from_statement` still returns session-attached - # ORM Task objects, so the mutation/commit/logging block below is unchanged. + # Why LATERAL and not ROW_NUMBER() over the whole table: the window form sorts + # every eligible row (tens/hundreds of thousands under a flooder) just to keep + # the first `per_company_cap` of each. Here `eligible_companies` is bounded by + # the number of companies, and the LATERAL fetches only the top `:per_company_cap` + # oldest rows per company — exactly the query the partial index + # ix_tasks_pending_run_after was built for — so the sort input is at most + # (companies x per_company_cap) rows, not the whole backlog. + # + # This is raw SQL because a window/LATERAL combined with FOR UPDATE SKIP LOCKED + # can't be expressed cleanly through the ORM, and the locking semantics (lock + # only the outer tasks alias `t`, skip already-locked ones) read far more + # clearly written out. `from_statement` still returns session-attached ORM Task + # objects, so the mutation/commit/logging block below is unchanged. # # Eligibility is identical to before: pending tasks past their run_after gate, - # plus stale processing tasks whose worker died. Stale tasks are ranked - # alongside pending ones rather than given a separate path. + # plus stale processing tasks whose worker died (the stale branch is + # unconditional — it never filters on run_after). # # The eligibility predicate is repeated on the outer `t` (not just in the - # `ranked` CTE) on purpose. `ranked` reads tasks without a lock, so two + # candidate CTEs) on purpose. Those CTEs read tasks without a lock, so two # workers can build overlapping candidate sets from the same snapshot. Under # READ COMMITTED, when `FOR UPDATE` meets a row a concurrent worker already # claimed and committed, Postgres re-fetches the latest version and re-checks @@ -78,20 +84,30 @@ def claim_tasks(db: Session, batch_size: int = 10) -> list[Task]: # hand the freed row to a second worker and claim it twice. sql = text( """ - WITH ranked AS ( - SELECT id, - ROW_NUMBER() OVER (PARTITION BY company_id ORDER BY created_at) AS rn, - created_at + WITH eligible_companies AS ( + SELECT DISTINCT company_id FROM tasks WHERE (status = 'pending' AND run_after <= :now) OR (status = 'processing' AND locked_at < :stale_before) ), - candidates AS ( - SELECT id, rn, created_at FROM ranked WHERE rn <= :batch_size + per_company_top AS ( + SELECT ec.company_id, top.id, top.created_at, + ROW_NUMBER() OVER (PARTITION BY ec.company_id + ORDER BY top.created_at) AS rn + FROM eligible_companies ec + CROSS JOIN LATERAL ( + SELECT id, created_at + FROM tasks + WHERE company_id = ec.company_id + AND ((status = 'pending' AND run_after <= :now) + OR (status = 'processing' AND locked_at < :stale_before)) + ORDER BY created_at + LIMIT :per_company_cap + ) top ) SELECT t.* FROM tasks t - JOIN candidates c ON c.id = t.id + JOIN per_company_top c ON c.id = t.id WHERE (t.status = 'pending' AND t.run_after <= :now) OR (t.status = 'processing' AND t.locked_at < :stale_before) ORDER BY c.rn, c.created_at @@ -102,7 +118,12 @@ def claim_tasks(db: Session, batch_size: int = 10) -> list[Task]: with claim_duration_seconds.time(): rows = db.scalars( select(Task).from_statement(sql), - {"now": now, "stale_before": stale_before, "batch_size": batch_size}, + { + "now": now, + "stale_before": stale_before, + "per_company_cap": settings.per_company_claim_cap, + "batch_size": batch_size, + }, ).all() for task in rows: diff --git a/scripts/perf_explain.sql b/scripts/perf_explain.sql index 0cad0b6..e297e18 100644 --- a/scripts/perf_explain.sql +++ b/scripts/perf_explain.sql @@ -1,37 +1,48 @@ -- perf_explain.sql -- -- Snapshot of the production fair-claim query, copied verbatim from claim_tasks() --- in app/services/worker_service.py (same CTE, same outer eligibility predicate, --- same ORDER BY, same FOR UPDATE ... SKIP LOCKED, same LIMIT). Bound parameters are --- inlined as literals for offline EXPLAIN: :now -> now(), --- :stale_before -> now() - interval '5 minutes', :batch_size -> 10. +-- in app/services/worker_service.py (LATERAL top-K per eligible company, same outer +-- eligibility re-check, same ORDER BY, same FOR UPDATE ... SKIP LOCKED, same LIMIT). +-- Bound parameters are inlined as literals for offline EXPLAIN: :now -> now(), +-- :stale_before -> now() - interval '5 minutes', :per_company_cap -> 10, +-- :batch_size -> 10. -- -- IMPORTANT: if the production SQL in claim_tasks() changes, update this file in the --- SAME commit — an EXPLAIN of a stale query measures nothing useful. +-- SAME commit — an EXPLAIN of a stale query measures nothing useful. (See +-- scripts/perf_explain_old.sql for the previous ROW_NUMBER baseline.) -- -- track_io_timing: turned on below because BUFFERS alone reports block counts but no -- I/O *time*. Without it, "is the query I/O-bound or memory-resident?" cannot be --- answered from the plan (I/O Timings lines simply won't appear). This is a --- session-scoped SET — it makes no permanent change to the database. +-- answered from the plan. Session-scoped SET — no permanent change to the database. SET track_io_timing = on; \echo ===== WARMUP (discard this plan) ===== EXPLAIN (ANALYZE, BUFFERS, VERBOSE) -WITH ranked AS ( - SELECT id, - ROW_NUMBER() OVER (PARTITION BY company_id ORDER BY created_at) AS rn, - created_at +WITH eligible_companies AS ( + SELECT DISTINCT company_id FROM tasks WHERE (status = 'pending' AND run_after <= now()) OR (status = 'processing' AND locked_at < now() - interval '5 minutes') ), -candidates AS ( - SELECT id, rn, created_at FROM ranked WHERE rn <= 10 +per_company_top AS ( + SELECT ec.company_id, top.id, top.created_at, + ROW_NUMBER() OVER (PARTITION BY ec.company_id + ORDER BY top.created_at) AS rn + FROM eligible_companies ec + CROSS JOIN LATERAL ( + SELECT id, created_at + FROM tasks + WHERE company_id = ec.company_id + AND ((status = 'pending' AND run_after <= now()) + OR (status = 'processing' AND locked_at < now() - interval '5 minutes')) + ORDER BY created_at + LIMIT 10 + ) top ) SELECT t.* FROM tasks t -JOIN candidates c ON c.id = t.id +JOIN per_company_top c ON c.id = t.id WHERE (t.status = 'pending' AND t.run_after <= now()) OR (t.status = 'processing' AND t.locked_at < now() - interval '5 minutes') ORDER BY c.rn, c.created_at @@ -40,20 +51,30 @@ LIMIT 10; \echo ===== RUN 1 ===== EXPLAIN (ANALYZE, BUFFERS, VERBOSE) -WITH ranked AS ( - SELECT id, - ROW_NUMBER() OVER (PARTITION BY company_id ORDER BY created_at) AS rn, - created_at +WITH eligible_companies AS ( + SELECT DISTINCT company_id FROM tasks WHERE (status = 'pending' AND run_after <= now()) OR (status = 'processing' AND locked_at < now() - interval '5 minutes') ), -candidates AS ( - SELECT id, rn, created_at FROM ranked WHERE rn <= 10 +per_company_top AS ( + SELECT ec.company_id, top.id, top.created_at, + ROW_NUMBER() OVER (PARTITION BY ec.company_id + ORDER BY top.created_at) AS rn + FROM eligible_companies ec + CROSS JOIN LATERAL ( + SELECT id, created_at + FROM tasks + WHERE company_id = ec.company_id + AND ((status = 'pending' AND run_after <= now()) + OR (status = 'processing' AND locked_at < now() - interval '5 minutes')) + ORDER BY created_at + LIMIT 10 + ) top ) SELECT t.* FROM tasks t -JOIN candidates c ON c.id = t.id +JOIN per_company_top c ON c.id = t.id WHERE (t.status = 'pending' AND t.run_after <= now()) OR (t.status = 'processing' AND t.locked_at < now() - interval '5 minutes') ORDER BY c.rn, c.created_at @@ -62,20 +83,30 @@ LIMIT 10; \echo ===== RUN 2 ===== EXPLAIN (ANALYZE, BUFFERS, VERBOSE) -WITH ranked AS ( - SELECT id, - ROW_NUMBER() OVER (PARTITION BY company_id ORDER BY created_at) AS rn, - created_at +WITH eligible_companies AS ( + SELECT DISTINCT company_id FROM tasks WHERE (status = 'pending' AND run_after <= now()) OR (status = 'processing' AND locked_at < now() - interval '5 minutes') ), -candidates AS ( - SELECT id, rn, created_at FROM ranked WHERE rn <= 10 +per_company_top AS ( + SELECT ec.company_id, top.id, top.created_at, + ROW_NUMBER() OVER (PARTITION BY ec.company_id + ORDER BY top.created_at) AS rn + FROM eligible_companies ec + CROSS JOIN LATERAL ( + SELECT id, created_at + FROM tasks + WHERE company_id = ec.company_id + AND ((status = 'pending' AND run_after <= now()) + OR (status = 'processing' AND locked_at < now() - interval '5 minutes')) + ORDER BY created_at + LIMIT 10 + ) top ) SELECT t.* FROM tasks t -JOIN candidates c ON c.id = t.id +JOIN per_company_top c ON c.id = t.id WHERE (t.status = 'pending' AND t.run_after <= now()) OR (t.status = 'processing' AND t.locked_at < now() - interval '5 minutes') ORDER BY c.rn, c.created_at @@ -84,20 +115,30 @@ LIMIT 10; \echo ===== RUN 3 ===== EXPLAIN (ANALYZE, BUFFERS, VERBOSE) -WITH ranked AS ( - SELECT id, - ROW_NUMBER() OVER (PARTITION BY company_id ORDER BY created_at) AS rn, - created_at +WITH eligible_companies AS ( + SELECT DISTINCT company_id FROM tasks WHERE (status = 'pending' AND run_after <= now()) OR (status = 'processing' AND locked_at < now() - interval '5 minutes') ), -candidates AS ( - SELECT id, rn, created_at FROM ranked WHERE rn <= 10 +per_company_top AS ( + SELECT ec.company_id, top.id, top.created_at, + ROW_NUMBER() OVER (PARTITION BY ec.company_id + ORDER BY top.created_at) AS rn + FROM eligible_companies ec + CROSS JOIN LATERAL ( + SELECT id, created_at + FROM tasks + WHERE company_id = ec.company_id + AND ((status = 'pending' AND run_after <= now()) + OR (status = 'processing' AND locked_at < now() - interval '5 minutes')) + ORDER BY created_at + LIMIT 10 + ) top ) SELECT t.* FROM tasks t -JOIN candidates c ON c.id = t.id +JOIN per_company_top c ON c.id = t.id WHERE (t.status = 'pending' AND t.run_after <= now()) OR (t.status = 'processing' AND t.locked_at < now() - interval '5 minutes') ORDER BY c.rn, c.created_at diff --git a/scripts/perf_explain_old.sql b/scripts/perf_explain_old.sql new file mode 100644 index 0000000..2409557 --- /dev/null +++ b/scripts/perf_explain_old.sql @@ -0,0 +1,105 @@ +-- perf_explain_old.sql +-- +-- BASELINE snapshot of the PRE-LATERAL fair-claim query (global ROW_NUMBER window +-- over every eligible row). Kept only so the performance deep-dive can show the old +-- and new plans side-by-side without checking out an old commit. This is NOT the +-- production query anymore — see perf_explain.sql for the current one +-- (LATERAL top-K per company, matching claim_tasks() in +-- app/services/worker_service.py). +-- +-- Bound params inlined as literals: :now -> now(), +-- :stale_before -> now() - interval '5 minutes', :batch_size -> 10. +-- +-- track_io_timing: on so BUFFERS reports I/O *time*, not just block counts; +-- otherwise the I/O-bound vs memory-resident question can't be answered from the +-- plan. Session-scoped SET — no permanent change to the database. + +SET track_io_timing = on; + +\echo ===== WARMUP (discard this plan) ===== +EXPLAIN (ANALYZE, BUFFERS, VERBOSE) +WITH ranked AS ( + SELECT id, + ROW_NUMBER() OVER (PARTITION BY company_id ORDER BY created_at) AS rn, + created_at + FROM tasks + WHERE (status = 'pending' AND run_after <= now()) + OR (status = 'processing' AND locked_at < now() - interval '5 minutes') +), +candidates AS ( + SELECT id, rn, created_at FROM ranked WHERE rn <= 10 +) +SELECT t.* +FROM tasks t +JOIN candidates c ON c.id = t.id +WHERE (t.status = 'pending' AND t.run_after <= now()) + OR (t.status = 'processing' AND t.locked_at < now() - interval '5 minutes') +ORDER BY c.rn, c.created_at +FOR UPDATE OF t SKIP LOCKED +LIMIT 10; + +\echo ===== RUN 1 ===== +EXPLAIN (ANALYZE, BUFFERS, VERBOSE) +WITH ranked AS ( + SELECT id, + ROW_NUMBER() OVER (PARTITION BY company_id ORDER BY created_at) AS rn, + created_at + FROM tasks + WHERE (status = 'pending' AND run_after <= now()) + OR (status = 'processing' AND locked_at < now() - interval '5 minutes') +), +candidates AS ( + SELECT id, rn, created_at FROM ranked WHERE rn <= 10 +) +SELECT t.* +FROM tasks t +JOIN candidates c ON c.id = t.id +WHERE (t.status = 'pending' AND t.run_after <= now()) + OR (t.status = 'processing' AND t.locked_at < now() - interval '5 minutes') +ORDER BY c.rn, c.created_at +FOR UPDATE OF t SKIP LOCKED +LIMIT 10; + +\echo ===== RUN 2 ===== +EXPLAIN (ANALYZE, BUFFERS, VERBOSE) +WITH ranked AS ( + SELECT id, + ROW_NUMBER() OVER (PARTITION BY company_id ORDER BY created_at) AS rn, + created_at + FROM tasks + WHERE (status = 'pending' AND run_after <= now()) + OR (status = 'processing' AND locked_at < now() - interval '5 minutes') +), +candidates AS ( + SELECT id, rn, created_at FROM ranked WHERE rn <= 10 +) +SELECT t.* +FROM tasks t +JOIN candidates c ON c.id = t.id +WHERE (t.status = 'pending' AND t.run_after <= now()) + OR (t.status = 'processing' AND t.locked_at < now() - interval '5 minutes') +ORDER BY c.rn, c.created_at +FOR UPDATE OF t SKIP LOCKED +LIMIT 10; + +\echo ===== RUN 3 ===== +EXPLAIN (ANALYZE, BUFFERS, VERBOSE) +WITH ranked AS ( + SELECT id, + ROW_NUMBER() OVER (PARTITION BY company_id ORDER BY created_at) AS rn, + created_at + FROM tasks + WHERE (status = 'pending' AND run_after <= now()) + OR (status = 'processing' AND locked_at < now() - interval '5 minutes') +), +candidates AS ( + SELECT id, rn, created_at FROM ranked WHERE rn <= 10 +) +SELECT t.* +FROM tasks t +JOIN candidates c ON c.id = t.id +WHERE (t.status = 'pending' AND t.run_after <= now()) + OR (t.status = 'processing' AND t.locked_at < now() - interval '5 minutes') +ORDER BY c.rn, c.created_at +FOR UPDATE OF t SKIP LOCKED +LIMIT 10; diff --git a/tests/test_lateral_claim.py b/tests/test_lateral_claim.py new file mode 100644 index 0000000..4f8fb60 --- /dev/null +++ b/tests/test_lateral_claim.py @@ -0,0 +1,94 @@ +import uuid +from concurrent.futures import ThreadPoolExecutor +from datetime import UTC, datetime, timedelta + +from sqlalchemy import delete + +from app.config import settings +from app.models import Task, TaskStatus +from app.services import worker_service + + +def _insert(db, company_id, created_at, *, status=TaskStatus.pending, + run_after=None, locked_at=None, locked_by=None): + """Insert a task with explicit fields; run_after defaults to just-past.""" + if run_after is None: + run_after = created_at - timedelta(minutes=1) + task = Task( + id=uuid.uuid4(), + company_id=company_id, + task_type="send_email", + payload={"to": "a@b.com"}, + status=status, + run_after=run_after, + locked_at=locked_at, + locked_by=locked_by, + created_at=created_at, + updated_at=created_at, + ) + db.add(task) + db.commit() + return task.id + + +def test_lateral_respects_per_company_cap(db): + base = datetime.now(UTC) - timedelta(hours=1) + company = uuid.uuid4() + for i in range(20): + _insert(db, company, base + timedelta(seconds=i)) + + claimed = worker_service.claim_tasks(db, batch_size=100) + + # The LATERAL's LIMIT caps this one company's contribution to the batch. + from_company = [t for t in claimed if t.company_id == company] + assert len(from_company) == settings.per_company_claim_cap == 10 + + +def test_lateral_stale_recovery_bypasses_run_after(db): + company = uuid.uuid4() + stale = datetime.now(UTC) - timedelta(minutes=10) + far_future = datetime.now(UTC) + timedelta(days=30) + task_id = _insert( + db, + company, + created_at=datetime.now(UTC) - timedelta(hours=1), + status=TaskStatus.processing, + run_after=far_future, # would hide a pending task; must NOT hide a stale one + locked_at=stale, + locked_by="dead-worker", + ) + + claimed = worker_service.claim_tasks(db, batch_size=10) + + assert task_id in {t.id for t in claimed} + + +def test_lateral_high_contention_no_duplicates(TestSession): + # Run the whole race many times: the tier-2 duplicate-claim race was ~1-in-4, + # so a single pass isn't proof. + for _ in range(20): + setup = TestSession() + setup.execute(delete(Task)) # clear between iterations + setup.commit() + base = datetime.now(UTC) - timedelta(hours=1) + companies = [uuid.uuid4() for _ in range(4)] + for c in companies: + for i in range(5): # 4 companies x 5 = 20 eligible tasks + _insert(setup, c, base + timedelta(seconds=i)) + setup.close() + + def run_pass(): + session = TestSession() + try: + return [ + t.id for t in worker_service.claim_tasks(session, batch_size=2) + ] + finally: + session.close() + + with ThreadPoolExecutor(max_workers=10) as pool: + futures = [pool.submit(run_pass) for _ in range(10)] + all_ids = [task_id for f in futures for task_id in f.result()] + + assert len(all_ids) == len(set(all_ids)), "duplicate claim under contention" + assert len(all_ids) <= 20