UN-3445 [FIX] PG queue — indexed claim state removes high-concurrency scan-past

[muhammad-ali-e]

Problem

The PG-queue dequeue filtered vt <= now() per-row, with vt deliberately left out of the index. So claimed-but-unacked (in-flight) rows — which carry a future vt — sat at the front of each priority band and were scanned past on every claim: cost was O(in-flight depth).

Load-testing exposed the consequence. At 150 concurrent, PG throughput collapsed to ~1.5 req/s (vs Celery ~4.75) with monotonically climbing latency — and adding worker lanes made it worse (more in-flight → deeper scan-past). At ≤50 concurrent it was fine (small in-flight depth), which is why it only showed under real load.

Fix

Replace the per-row visibility filter with an indexed state (ready | claimed) machine:

• Claim walks a partial index containing ONLY state='ready' rows — in-flight rows are absent from the traversal, so there is nothing to scan past. Claim cost is flat regardless of in-flight depth (EXPLAIN shows identical cost at 0 vs 150 in-flight).
• Redelivery of a crashed worker's row moves from implicit vt-expiry (in the claim) to an explicit reaper re-arm (state='claimed' + expired vt → ready), run every 5s tick (not the 5-min retention sweep), with a pg_reaper_queue_rearmed_total counter.
• vt stays out of every index, so the frequent lease-renewal UPDATE (set_vt) stays HOT-eligible; only the once-per-message state transitions (claim / ack / rare re-arm) are non-HOT.

Portability

Pure core PostgreSQL — partial indexes, check constraint, column default. No extensions; identical on GCP Cloud SQL / Azure / AWS RDS+Aurora / on-prem. Performance tuning (fillfactor + aggressive autovacuum) is intentionally kept out of the migration — a locked-down managed role could reject ALTER TABLE SET (autovacuum_*) and fail the whole migrate, and tuning is per-environment anyway. It lives as an optional ops step (see below).

Safety

• Flag-gated (pg_queue_enabled); the Celery path is untouched.
• Backward-compatible: the column defaults to 'ready' (DB-level SET DEFAULT), so old-code consumers keep working on the new schema — verified live (the existing old-image PG workers stayed healthy against the migrated table).
• At-least-once preserved: the redelivered-after-crash re-run is absorbed by the unchanged idempotency stack (claim_batch / FileHistory / _callback_already_ran), exactly as an old-design vt-expiry re-run was.

Design provenance

Chosen via a Red/Blue debate + code fact-check. The debate's showstopper (reaper re-arm → double customer write/billing) was fact-checked invalid (the guards above stop it; the reaper only re-arms the same vt-expired set the old claim already redelivered). The real residual — redelivery now depends on the reaper (already the crash backstop) — is covered by the 5s-tick cadence + the re-arm counter (alert on oldest_claimed_expired_age recommended). The two-table split (more WAL) and admission-orchestrator (bounds depth, doesn't cure it) were rejected.

Validation (local, against real Postgres)

• 336 unit/integration tests pass, incl. new redelivery-path + live-lease tests.
• E2E: enqueue→claim→ack; crash→reaper re-arm→redeliver; live-lease not re-armed (no double-delivery).
• Scan-past cured — EXPLAIN ANALYZE: claim uses the partial index with identical cost at in-flight depth 0 and 150.
• Stress: 50 concurrent claimers / 10k msgs, batch claims, and 4 concurrent reapers — no double-claim, no loss, reaper re-armed 0 live claims; read_ct poison bound intact across re-arm cycles.

Follow-ups (not in this PR)

• Optional per-environment ops tuning: ALTER TABLE pg_queue_message SET (fillfactor=90, autovacuum_vacuum_scale_factor=0.02, …) if a soak shows partial-index bloat.
• Integration 150u re-benchmark behind PgBouncer — the fan-out validation local (no pooler, max_connections-bound) can't cover.
• Pre-existing set_vt owner-guard gap (a claim-token/owner_id) — separate ticket; neither introduced nor fixed here.

🤖 Generated with Claude Code

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on base/target branches other than the default branch.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: f326d28b-dd35-4cf9-8fa9-59aa5c1133bb

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch fix/UN-3445-pg-queue-claim-state

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[muhammad-ali-e]

Automated PR review — PG-queue indexed claim state (UN-3445)

Six-agent review (code / silent-failure / type-design / tests / comments / simplify). The core change is sound — the partial claim_idx WHERE state='ready' correctly eliminates the O(in-flight) scan-past, the raw SQL matches the Django model, and the reaper re-arm is the right explicit replacement for the implicit vt<=now() self-heal.

The findings below are inline. Top items to weigh before merge:

1. Redelivery now hard-depends on the reaper leader — a crashed worker's row strands forever if the reaper isn't running/elected (previously self-healing in the claim). No dedicated failure metric makes such an outage invisible.
2. Deploy-time: the AddField backfill readies in-flight rows (transient double-delivery); indexes/constraint build under table locks.
3. Type design: the ready/claimed enum is scattered as ~14 bare literals across both codebases despite the proven unstract.core + drift-test precedent used for priority.
4. Tests: no wiring test that the tick calls rearm_expired_claims/increments the metric, and no always-run unit/rollback test for the helper (integration tests are Postgres-gated).

None is a hard crash bug; nothing blocks the design.

[muhammad-ali-e]

[HIGH — redelivery robustness] Crash redelivery now depends solely on the reaper leader; a stranded claimed row is no longer self-healing.

Changing the claim predicate from vt <= now() to state = 'ready' removes the inline self-heal: previously any claimant re-admitted an expired row, so redelivery needed no external component. Now an expired claim stays state='claimed' (invisible to the claim) and only rearm_expired_claims flips it back — and that runs only on the leader reaper tick. The reaper "ships dark" (reaper.py module docstring; launched via python -m queue_backend.pg_queue.reaper), so in any env that runs the PG queue but doesn't run/elect the reaper, crashed-worker messages are stranded forever with no signal — a new availability regression, not just added latency.

Fix: make the coupling explicit and enforced — treat the reaper as a mandatory co-deployed process wherever dequeue is enabled, add reaper-liveness alerting for the queue path, and document the dependency here. Optionally keep a degradation path: ... WHERE state='ready' OR (state='claimed' AND vt <= now()) so self-heal survives a reaper outage. (Raised independently by code-review + silent-failure passes.)

[muhammad-ali-e]

[MED/HIGH — observability] No dedicated failure metric/log for the re-arm sweep; a persistent redelivery outage is invisible.

Every sibling recovery job has a failure signal (barrier_recovery_failures, claim_recovery_failures, sweep_failures + per-table streak log). This path added only the success counter queue_rearmed. On failure, rearm_expired_claims rolls back and re-raises into run()'s generic handler, which just bumps the shared tick_failures — so a re-arm that fails every tick (rows stuck claimed, nothing redelivering) is indistinguishable from a barrier/scheduler failure.

Fix: add a queue_rearm_failures counter + a dedicated log line (mirror _run_sweep's table + consecutive-streak log) so a redelivery outage is traceable on its own.

Also (test gap): the autouse stub_queue_rearm fixture promises "any wiring test opts in via the returned mock," but no such test exists — nothing asserts the leader tick actually calls rearm_expired_claims, increments the metric only when non-zero (if rearmed:), or is ordered after barrier recovery and before schedule dispatch. A refactor dropping this call would leave every reaper test green.

[muhammad-ali-e]

[MED — error semantics] "Best-effort" is misleading here — this block re-raises and aborts the rest of the tick.

In this file "best-effort" means swallow and continue (_run_sweep, _maybe_refresh_gauges). This block does the opposite: on failure it discards the conn and re-raises (next line), skipping dispatch_due_schedules (the Beat replacement), _maybe_sweep, and gauge refresh for that tick. So a persistent re-arm fault (e.g. a statement/lock timeout over a large claimed set) silently becomes a full scheduler outage — scheduled workflows stop firing under the same generic tick_failures signal.

Fix: drop the "Best-effort" wording (it re-raises, matching barrier-recovery — not the sweeps); and consider routing re-arm through the _run_sweep best-effort pattern so a redelivery fault doesn't starve unrelated scheduler/sweep work. If gating the scheduler on re-arm is intentional, say why.

[muhammad-ali-e]

[MED — deploy safety] AddField(default="ready") backfills all pre-existing rows to ready, including in-flight (future-vt) ones → transient double-delivery on deploy.

The new claim ignores vt and keys off state only. Any message that was claimed-but-unacked under the old design (future vt) is backfilled to ready and becomes immediately re-claimable, so during a rolling deploy a still-processing message can be picked up by another consumer → one-time double processing. It's absorbed by the at-least-once/idempotent contract, but it's real deploy-time behavior.

Fix: either document that the queue must be drained / consumers quiesced before applying 0014, or backfill defensively so genuinely in-flight rows land claimed (e.g. a follow-up RunSQL setting state='claimed' WHERE vt > now()).

[muhammad-ali-e]

[MED — lock during migration] Indexes are built non-concurrently and the CHECK constraint is validated under a full-table lock.

Django AddIndex emits plain CREATE INDEX (SHARE lock — blocks writes for the build) and AddConstraint emits ADD CONSTRAINT ... CHECK (ACCESS EXCLUSIVE + full-table validation); the migration is atomic, so these locks are held for its whole duration. On a queue table with a backlog this stalls producers and consumers during deploy. Usually brief (table is ack-drained), hence Medium.

Fix (if zero-downtime over a backlog matters): split into a non-atomic migration using AddIndexConcurrently, and add the constraint NOT VALID then VALIDATE CONSTRAINT separately. Otherwise note the expected lock in the runbook.

[muhammad-ali-e]

[LOW — comment accuracy] "Runs EVERY tick" omits the leader gate.

rearm_expired_claims is called unconditionally in tick(), but only after the leadership check — a standby returns earlier and never re-arms. Reword to "Runs every leader tick (redelivery cadence), not the retention sweep," so someone debugging why a standby isn't re-arming isn't misled. (Also: this docstring and the inline call-site comment at ~L1089 largely duplicate the cadence/cost rationale — let the docstring own it and trim the inline to the ordering/placement point.)

[muhammad-ali-e]

[HIGH — test gap] rearm_expired_claims has no unit SQL-contract or rollback/re-raise test, though every sibling sweep helper does.

sweep_expired_results and sweep_orphan_dedup each have a mock-cursor SQL-contract test and a parametrized rollback-on-error test in test_pg_reaper.py. This structurally-identical helper is in neither — its only coverage is the two integration tests, which are skipped whenever Postgres is unreachable (pg_conn fixture). So the error path (rollback then re-raise) has zero always-run coverage.

Add: test_rearm_expired_claims_sql (mock cursor rowcount=2 → asserts return value, UPDATE ... SET state='ready' WHERE state='claimed' AND vt <= now(), commit called once), and add rearm_expired_claims to the test_sweep_rolls_back_on_error parametrize list.

[muhammad-ali-e]

[MED — test gap] This autouse fixture's docstring promises an opt-in wiring test that doesn't exist.

It stubs rearm_expired_claims in every reaper test, so nothing verifies the tick actually wires it. Add tests mirroring TestSchedulerTick/TestRetentionSweepTick: (1) leader ticks twice → stub_queue_rearm.call_count == 2; (2) standby → assert_not_called(); (3) non-zero return → queue_rearmed.inc(n) called, zero return → not called (the if rearmed: guard); (4) ordering recover → rearm → schedule. Also add a test_rearm_error_discards_owned_conn (set side_effect, pytest.raises, assert _sweep_conn is None) — parallels exist for scheduler and recovery but not re-arm.

[muhammad-ali-e]

[MED — test gaps] The core state-machine invariants are only implied, not asserted.

Good behavioral tests here, but consider adding: (1) idempotency — after this re-arm returns 1, an immediate second rearm_expired_claims(pg_conn) == 0 (row is now ready); one line, documents the double-tick safety property. (2) CHECK constraint — a raw UPDATE ... SET state='bogus' should raise psycopg2.errors.CheckViolation (mirror the existing test_db_check_constraint_matches_fairness_bounds for priority); the model calls this "the backstop no writer can bypass," yet it's untested. (3) batch claim writes state — claim qty>1 and assert every returned row is state='claimed' in the DB while an unclaimed row stays ready (the write half of the machine the whole PR rests on).

Nit (simplify): the function-local from ... import rearm_expired_claims here and at L659 can hoist to the module-level imports (L24-28) if it imports cleanly.

[muhammad-ali-e]

[LOW — simplify] Third copy of the same try/commit/rollback envelope across the sweep helpers.

rearm_expired_claims repeats the exact try: cursor.execute(...); rowcount; commit / except: _rollback_after_sweep_failure(conn, table); raise shared with sweep_expired_results and sweep_orphan_dedup; only the SQL/params/table vary. Optional: extract _execute_sweep_dml(conn, table, sql, params=()) next to _rollback_after_sweep_failure — each helper keeps its (valuable) docstring and reduces to one line, behavior identical. Similarly the three try/except: self._discard_owned_sweep_conn(); raise blocks in tick() (barrier / re-arm / dispatch) could share a tiny _sweep_step() context manager (contextlib already imported).

[muhammad-ali-e]

Thanks — thorough review. Addressed in e6d16355. Summary by finding:

Addressed (code)

• Type design / magic strings (HIGH ×2): added QueueMessageState(StrEnum){ready,claimed} to unstract.core.data_models (next to WorkloadType) as the single source of truth. Every site now sources the value — backend model default/constraint/both index conditions, the claim SQL in client.py, and the reaper re-arm SQL — so a typo can't silently break the state machine. Added a drift-guard test asserting the DB CheckConstraint matches the enum (mirrors test_db_check_constraint_matches_fairness_bounds).
• Deploy safety — in-flight backfill (MED): migration now UPDATE ... SET state='claimed' WHERE vt > now() after AddField, so a rolling deploy over a non-empty queue can't transiently re-ready a still-processing row. No-op at first enablement (empty table).
• Observability (MED/HIGH): added pg_reaper_queue_rearm_failures_total + a dedicated error log; a persistent redelivery outage is now distinguishable from a barrier/scheduler fault.
• Error-semantics comment (MED): dropped the misleading "best-effort" wording — the block re-raises like barrier recovery (recovery is critical, not swallow-and-continue), now stated explicitly with the "both recover next tick" rationale.
• Comment accuracy: O(1)-amortised → "independent of in-flight depth"; "stays HOT" → "HOT-eligible (needs fillfactor<100)"; "runs EVERY tick" → "every leader tick"; dropped the ungrounded ~16/~150 fleet numbers and the wrong .create() example.
• Tests (HIGH + MED): test_rearm_expired_claims_sql + rollback parametrize entry; re-arm tick wiring (leader/standby gating, recover→rearm→schedule ordering, metric-only-when-nonzero guard, error-discards-conn + failure-counter); client idempotency (double-tick), state CHECK drift guard, batch-claim-writes-state; hoisted the function-local imports. The stub_queue_rearm fixture's promised wiring test now exists.

Deferred, with rationale

• Redelivery hard-depends on the reaper (HIGH): documented explicitly at the claim + model. Key point: the reaper is already the mandatory recovery singleton (stranded barriers/claims don't self-heal without it), so this adds queue re-arm to its remit rather than a new deployment dependency. Liveness alerting is in the ops runbook (oldest_claimed_expired_age). The optional degradation-path hybrid claim WHERE state='ready' OR (state='claimed' AND vt<=now()) was considered and not taken here: it can't use the state='ready' partial index for the OR branch, so it risks reintroducing the scan-past (seq-scan) or forcing vt back into an index (undoing the HOT win) — it needs an EXPLAIN spike, tracked separately.
• Migration lock (MED): the table is empty at first enablement (flag fail-closed) so the index/constraint build is instant; the expected lock over a backlog is noted in the runbook. Splitting to AddIndexConcurrently + NOT VALID/VALIDATE is a reasonable follow-up if we ever enable over a populated queue.
• set_vt owner-guard (pre-existing): separate ticket — a claim-token/owner_id; neither introduced nor fixed here.
• Sweep-DML helper extraction (LOW, simplify): deferred to keep this diff focused on correctness + the review's substantive items.

345 tests pass against real Postgres; e2e re-verified the claim uses the partial index at 0 vs 150 in-flight (scan-past still gone); no stray migration (enum deconstructs to the same literals).

[muhammad-ali-e]

[HIGH — test gap] The deploy-safety backfill has zero test coverage, and a regression would silently reintroduce fleet-wide double-processing.

This UPDATE … SET state='claimed' WHERE vt > now() is the one line that stops a rolling deploy over a non-empty queue from re-marking every in-flight (claimed-but-unacked, future-vt) row as ready → instant re-claim → one-time double processing across the entire in-flight set. Its correctness rests on an ordering invariant that nothing locks down: AddField first backfills all rows to ready, and only this subsequent RunSQL re-classifies the genuinely in-flight ones. A later refactor that drops/reorders this step, or inverts the vt > now() predicate, would reintroduce the exact double-processing this PR exists to prevent — and every current test would still pass (the reaper-side state machine is well covered; this migration step is exercised by nothing).

Suggested test (Django MigrationExecutor): migrate to 0013, insert one row with vt = now() + interval '1 hour' (in-flight) and one with vt = now() - interval '1 hour' (idle), migrate to 0014, then assert the future-vt row is state='claimed' and the past-vt row is state='ready'. That is the only test that catches a regressed or removed backfill. (No migration-executor harness exists under backend/ yet, so this also establishes the pattern.)

[greptile-apps]

Greptile Summary

This PR fixes the O(in-flight depth) scan-past in the PG-queue dequeue path by replacing the old per-row vt <= now() filter with an indexed state field (ready | claimed). The claim now walks a partial index that contains only state='ready' rows — in-flight rows are absent entirely — making claim cost independent of in-flight depth. Crash redelivery moves from implicit vt-expiry in the claim to an explicit reaper sweep (rearm_expired_claims) running every 5 s leader tick.

• Schema change (migration 0014): adds the state column with a DB-level default of ready, a CheckConstraint enforcing the closed enum, a partial claim index (queue_name, priority DESC, msg_id) WHERE state='ready', and a second partial reaper index (msg_id) WHERE state='claimed'; also includes a deploy-safety backfill that re-classifies any in-flight rows present during a rolling deploy.
• Reaper integration: new rearm_expired_claims helper runs on every leader tick (not the 5-min retention sweep) with dedicated pg_reaper_queue_rearmed_total and pg_reaper_queue_rearm_failures_total counters; error posture mirrors barrier recovery (re-raise + connection discard).
• Correctness coverage: 336 tests pass including new integration tests for the redelivery path, live-lease non-re-arm, state check-constraint drift guard, migration backfill structural check, and multi-reaper idempotency.

Confidence Score: 5/5

Safe to merge — the state machine is logically correct, the migration ordering is sound, and crash redelivery is covered by the reaper with the same at-least-once guarantees as before.

The core correctness properties are all solid: the partial-index claim eliminates the scan-past without changing the SKIP LOCKED safety guarantee; the rearm SQL is idempotent and correctly bounded to the claimed partial index; the deploy-safety backfill guards against double-delivery during a rolling deploy; and the docstring fixes from the previous review round are included in this diff.

No files require special attention. The migration file is the highest-stakes change and its operation ordering is correct.

Important Files Changed

Filename	Overview
backend/pg_queue/migrations/0014_remove_pgqueuemessage_pg_queue_message_dequeue_idx_and_more.py	Adds state column, partial indexes, and CheckConstraint; deploy-safety backfill re-classifies in-flight rows correctly; operation ordering is sound (AddField → backfill → AddIndex → AddConstraint).
workers/queue_backend/pg_queue/client.py	Claim SQL updated from vt<=now() to state='ready' partial-index predicate; state literals sourced from shared enum; set_vt and module docstrings updated to reflect the new reaper-based redelivery model.
workers/queue_backend/pg_queue/reaper.py	New rearm_expired_claims function is correct, idempotent, and bounded by the claimed partial index; wired into tick() after barrier recovery with matching error posture (re-raise + conn discard) and dedicated metrics counters.
backend/pg_queue/models.py	State field, CheckConstraint, and two partial indexes added to match the migration; enum aliases imported from the shared single source of truth.
unstract/core/src/unstract/core/data_models.py	New QueueMessageState StrEnum (READY/CLAIMED) is the cross-codebase single source of truth; well-documented with drift-test reference.
workers/queue_backend/pg_queue/metrics.py	Two new Prometheus counters (queue_rearmed_total, queue_rearm_failures_total) follow the existing counter pattern and correctly distinguish redelivery outages from barrier/scheduler faults.
workers/tests/test_pg_queue_client.py	New integration tests cover redelivery-via-reaper, live-lease non-re-arm, state check-constraint drift, migration backfill SQL, and batch claim state writes.
workers/tests/test_pg_reaper.py	New TestQueueRearmTick class covers leader/standby gating, tick ordering, metric guard, and error posture; autouse stub correctly patches module-level name.
backend/pg_queue/tests/test_migration_0014_backfill.py	Structural regression guard for the deploy-safety backfill: verifies the RunSQL is present, has the correct predicate, and runs after AddField.
workers/tests/test_pg_schema_drift.py	Adds state to the worker schema contract for pg_queue_message; minimal and correct.

_{Reviews (2): Last reviewed commit: "UN-3445 [FIX] address review round 2 — b..." | Re-trigger Greptile}

[muhammad-ali-e]

Round 2 addressed in be3d23fe.

Toolkit — [HIGH] deploy-safety backfill had zero test coverage
Added two complementary guards (the backend has no pytest-django / migration-executor harness, so a full-migration run isn't cheap — and running migrations against the shared dev DB in a test is risky):

• Structural (backend/pg_queue/tests/test_migration_0014_backfill.py): pins the backfill RunSQL — present, with the exact state='claimed' WHERE vt > now() predicate, and ordered after AddField(state). Catches removal, reordering, and predicate inversion (the three regressions you named).
• Behavioural (test_pg_queue_client::test_inflight_backfill_sql_classifies_by_vt): runs the exact backfill statement against real Postgres and asserts a future-vt row lands claimed while a past-vt row stays ready.

Together they cover both "the step is still there and correctly ordered" and "the predicate actually selects the in-flight set."

Greptile — stale client.py docstrings
Fixed both: the module docstring and set_vt no longer describe crash redelivery as implicit vt-expiry reappearance — they now describe the reaper-mediated re-arm (state='claimed' + expired lease → reaper → ready), and set_vt's docstring notes its dual role (lease renewal keeping a live claim's vt in the future + poison re-park).

160 workers tests + the new backend structural test pass. Greptile's other observations (correct claim SQL, no reaper-vs-ack race, concurrent re-arm safety) matched the design intent — no change needed there.

[sonarqubecloud]

Quality Gate passed

Issues
0 New issues
0 Accepted issues

Measures
0 Security Hotspots
0.0% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud