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

backend/pg_queue/migrations/0014_remove_pgqueuemessage_pg_queue_message_dequeue_idx_and_more.py

@@ -0,0 +1,78 @@
+# Generated by Django 4.2.30 on 2026-07-09 10:40
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("pg_queue", "0013_pgorchestrationclaim_organization_id"),
+    ]
+
+    operations = [
+        migrations.RemoveIndex(
+            model_name="pgqueuemessage",
+            name="pg_queue_message_dequeue_idx",
+        ),
+        migrations.AddField(
+            model_name="pgqueuemessage",
+            name="state",
+            field=models.TextField(default="ready"),
+        ),
+        # Django's AddField backfills every existing row to 'ready' then DROPs the
+        # DB default (it manages defaults at the ORM layer). Re-add a DB-level
+        # default so the schema is self-protecting: a raw INSERT that omits `state`
+        # still lands a valid 'ready' row rather than a NOT-NULL violation.
+        # (The ORM always sends state='ready' via the field default, so this guards
+        # only raw SQL; insert_message_sql() also sets it explicitly — belt +
+        # suspenders.) Metadata-only; touches no row data.
+        migrations.RunSQL(
+            sql="ALTER TABLE pg_queue_message ALTER COLUMN state SET DEFAULT 'ready'",
+            reverse_sql="ALTER TABLE pg_queue_message ALTER COLUMN state DROP DEFAULT",
+        ),
+        # Deploy-safety backfill: the new claim keys off `state` and ignores `vt`,
+        # but AddField set EVERY pre-existing row (including claimed-but-unacked
+        # in-flight rows, which carry a future `vt`) to 'ready'. Without this, a
+        # rolling deploy over a non-empty queue would make a still-processing
+        # message instantly re-claimable → one-time double processing (absorbed by
+        # the idempotency stack, but avoidable). Land genuinely in-flight rows as
+        # 'claimed' so they stay invisible until their lease expires. No-op at the
+        # first enablement (the queue table is empty while the flag is off).
+        migrations.RunSQL(
+            sql="UPDATE pg_queue_message SET state = 'claimed' WHERE vt > now()",
+            reverse_sql=migrations.RunSQL.noop,
+        ),
+        migrations.AddIndex(
+            model_name="pgqueuemessage",
+            index=models.Index(
+                models.F("queue_name"),
+                models.OrderBy(models.F("priority"), descending=True),
+                models.F("msg_id"),
+                condition=models.Q(("state", "ready")),
+                name="pg_queue_message_claim_idx",
+            ),
+        ),
+        migrations.AddIndex(
+            model_name="pgqueuemessage",
+            index=models.Index(
+                models.F("msg_id"),
+                condition=models.Q(("state", "claimed")),
+                name="pg_queue_message_claimed_idx",
+            ),
+        ),
+        migrations.AddConstraint(
+            model_name="pgqueuemessage",
+            constraint=models.CheckConstraint(
+                check=models.Q(("state__in", ["ready", "claimed"])),
+                name="pg_queue_message_state_valid",
+            ),
+        ),
+        # NOTE (UN-3445): performance tuning for the state-machine churn
+        # (fillfactor + aggressive per-table autovacuum) is deliberately KEPT OUT of
+        # this migration. Those are core-Postgres reloptions, but they are (a)
+        # environment-specific tuning that belongs in the ops/deploy layer, not app
+        # schema, and (b) potentially restricted by a locked-down managed role
+        # (Cloud SQL / RDS / Azure), where a failed `ALTER TABLE SET (autovacuum_*)`
+        # would abort the whole migration. The fix itself (state column + partial
+        # indexes) is pure portable DDL and needs none of it. If a soak shows bloat,
+        # apply the tuning per environment via ops (see the PR/runbook), not here.
+    ]

backend/pg_queue/models.py

@@ -2,6 +2,13 @@
 from django.db.models import F
 from django.utils import timezone
 
+from unstract.core.data_models import QueueMessageState
+
+# Local aliases so the model reads cleanly; the enum is the single source of truth
+# shared with the workers' raw SQL (client.py / reaper.py). See QueueMessageState.
+_READY = QueueMessageState.READY.value
+_CLAIMED = QueueMessageState.CLAIMED.value
+
 
 class PgQueueMessage(models.Model):
     """Bespoke extension-free queue table (``SKIP LOCKED`` + visibility timeout).
@@ -41,6 +48,29 @@ class PgQueueMessage(models.Model):
     # (workload) ordering + burst_max admission are deferred to the fair-admission
     # orchestrator. The CheckConstraint is the one backstop no writer can bypass.
     priority = models.SmallIntegerField(default=5)
+    # Claim state (UN-3445 scan-past fix) — QueueMessageState {ready, claimed}.
+    # 'ready' = claimable, 'claimed' = in-flight (a consumer holds it, vt is its
+    # lease). This makes visibility an INDEXED predicate instead of the old per-row
+    # `vt <= now()` filter: the claim walks a partial index that contains ONLY
+    # 'ready' rows, so in-flight rows are no longer physically scanned past on every
+    # claim (the O(in-flight) cost that collapsed throughput at high concurrency).
+    #
+    # REDELIVERY DEPENDENCY: a crashed worker's row (state='claimed', vt expired) is
+    # re-armed to 'ready' by the reaper (reaper.rearm_expired_claims), NOT by the
+    # claim itself (the old `vt<=now()` self-heal is gone). The reaper is therefore
+    # on the crash-redelivery path — it is already the mandatory recovery singleton
+    # for stranded barriers/claims, so this adds queue re-arm to its remit rather
+    # than a new deployment dependency. Alert on its liveness for the queue path
+    # (PG_QUEUE_CLAIM_STATE_RUNBOOK: oldest_claimed_expired_age).
+    #
+    # `vt` is STILL not in any index, so the frequent lease-renewal UPDATE (set_vt)
+    # stays HOT-eligible (actual HOT also needs heap-page free space, i.e.
+    # fillfactor < 100 — see migration 0014); only the once-per-message state
+    # transitions (claim / ack / rare re-arm) are non-HOT — bounded. If a soak shows
+    # partial-index bloat, mitigate with per-environment ops tuning (fillfactor +
+    # aggressive autovacuum) — kept out of the migration so the schema stays
+    # portable across Postgres providers.
+    state = models.TextField(default=_READY)
 
     class Meta:
         db_table = "pg_queue_message"
@@ -49,22 +79,41 @@ class Meta:
                 check=models.Q(priority__gte=1) & models.Q(priority__lte=10),
                 name="pg_queue_message_priority_range",
             ),
+            # Backstop no writer can bypass: state is a closed enum. Values sourced
+            # from QueueMessageState (single source of truth; drift-tested).
+            models.CheckConstraint(
+                check=models.Q(state__in=[_READY, _CLAIMED]),
+                name="pg_queue_message_state_valid",
+            ),
         ]
         indexes = [
-            # The dequeue walks one queue in (priority DESC, msg_id ASC) order
-            # and applies vt <= now() as a per-row filter during the walk —
-            # claim high-priority first, FIFO (msg_id, monotonic and stable
-            # across re-claims) within a band. Note this is NOT a guaranteed
-            # top-N: vt is intentionally not in the index, so claimed-but-unacked
-            # rows (future vt) sit at the front of their band and are scanned
-            # past on each claim. Fine at low in-flight depth; the orchestrator's
-            # staging→task admission is the answer if that backlog grows large.
+            # CLAIM path — partial index over ONLY claimable rows. The dequeue
+            # (`WHERE queue_name=? AND state='ready' ORDER BY priority DESC, msg_id
+            # FOR UPDATE SKIP LOCKED`) walks this in (priority DESC, msg_id ASC)
+            # order; claimed rows are absent from the index entirely, so claim cost
+            # is INDEPENDENT of in-flight (claimed) depth — it walks O(LIMIT) ready
+            # rows plus any a concurrent claimer momentarily holds under FOR UPDATE
+            # (bounded by concurrency, not backlog). `vt` is deliberately NOT here
+            # (keeps lease renewal HOT-eligible). Replaces the old full
+            # `pg_queue_message_dequeue_idx`.
             models.Index(
                 F("queue_name"),
                 F("priority").desc(),
                 F("msg_id"),
-                name="pg_queue_message_dequeue_idx",
-            )
+                condition=models.Q(state=_READY),
+                name="pg_queue_message_claim_idx",
+            ),
+            # REAPER path — partial index over ONLY in-flight rows (bounded by total
+            # consumer concurrency across the fleet — typically small, ~one in-flight
+            # row per busy worker). The re-arm sweep (`WHERE state='claimed' AND
+            # vt<=now()`) enumerates this small set and filters vt in-scan. `vt` is
+            # intentionally excluded so a lease renewal (which rewrites vt every
+            # ~lease/3) does not maintain this index and stays HOT-eligible.
+            models.Index(
+                F("msg_id"),
+                condition=models.Q(state=_CLAIMED),
+                name="pg_queue_message_claimed_idx",
+            ),
         ]
 
 

backend/pg_queue/tests/test_migration_0014_backfill.py

@@ -0,0 +1,58 @@
+"""Regression guard for migration 0014's deploy-safety in-flight backfill.
+
+The one line ``UPDATE pg_queue_message SET state='claimed' WHERE vt > now()`` is
+what 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 →
+fleet-wide one-time double processing (UN-3445 PR review, HIGH). Its correctness
+rests on an ordering invariant nothing else locks down: ``AddField`` first backfills
+ALL rows to ``ready``, and only this later ``RunSQL`` re-classifies the genuinely
+in-flight ones. A refactor that drops/reorders the step, or inverts the ``vt > now()``
+predicate, would silently reintroduce the exact bug this PR prevents — and every
+state-machine test would still pass (they don't exercise the migration step).
+
+This structurally pins the operation (present, correct predicate, ordered after the
+``AddField``). The *behavioural* half — that ``vt > now()`` actually selects the
+in-flight rows against real Postgres — is
+``test_pg_queue_client.py::test_inflight_backfill_sql_classifies_by_vt`` (the backend
+has no pytest-django / migration-executor harness, so a full-migration run isn't
+cheap here; these two together catch removal, reordering, and predicate inversion).
+"""
+
+from importlib import import_module
+
+from django.db import migrations
+
+_MIGRATION = import_module(
+    "pg_queue.migrations."
+    "0014_remove_pgqueuemessage_pg_queue_message_dequeue_idx_and_more"
+)
+
+
+def _first_index(ops, predicate):
+    return next((i for i, op in enumerate(ops) if predicate(op)), -1)
+
+
+def test_inflight_backfill_present_correct_and_ordered_after_addfield():
+    ops = _MIGRATION.Migration.operations
+
+    add_state = _first_index(
+        ops,
+        lambda o: isinstance(o, migrations.AddField) and o.name == "state",
+    )
+    backfill = _first_index(
+        ops,
+        lambda o: isinstance(o, migrations.RunSQL)
+        and "state = 'claimed'" in str(o.sql)
+        and "vt > now()" in str(o.sql),
+    )
+
+    # Present: the backfill exists with the exact re-classification predicate.
+    assert add_state >= 0, "0014 must add the `state` field"
+    assert backfill >= 0, (
+        "0014 must backfill in-flight rows with "
+        "`UPDATE ... SET state='claimed' WHERE vt > now()` — removing/inverting it "
+        "reintroduces fleet-wide double-delivery on a rolling deploy"
+    )
+    # Ordered: AddField sets every row 'ready' first; the backfill must run AFTER
+    # it to re-classify the genuinely in-flight rows (reorder = no-op backfill).
+    assert backfill > add_state, "the in-flight backfill must run after AddField(state)"

unstract/core/src/unstract/core/data_models.py

@@ -273,6 +273,28 @@ class WorkloadType(StrEnum):
     NON_API = "non_api"
 
 
+class QueueMessageState(StrEnum):
+    """Claim state of a ``pg_queue_message`` row (UN-3445 scan-past fix).
+
+    Single source of truth for the closed enum, shared across the two codebases
+    that can't import each other: the backend model (default + CheckConstraint +
+    the two partial-index conditions) and the workers' raw SQL (the claim in
+    ``pg_queue.client`` and the reaper re-arm in ``pg_queue.reaper``). A typo in
+    any one of those bare literals would silently break the state machine (e.g. a
+    mistyped re-arm no-ops → crashed work never redelivers), so every site sources
+    the string from here. A drift test asserts the DB CheckConstraint matches this
+    set, mirroring the ``priority`` (fairness) precedent. ``str`` Enum → serialises
+    to its value and compares equal to the bare string.
+
+    - ``READY``   — claimable: the dequeue's partial claim index holds only these.
+    - ``CLAIMED`` — in-flight: a consumer holds it, ``vt`` is its renewable lease;
+      re-armed back to ``READY`` by the reaper when the lease expires (crash).
+    """
+
+    READY = "ready"
+    CLAIMED = "claimed"
+
+
 # Fairness L3 priority bounds (1..10, higher = claimed sooner). Single source of
 # truth shared by the backend producer and the workers' fairness/queue client —
 # the DB CheckConstraint on pg_queue_message.priority is the writer-proof backstop.