fix(storage): keep claim tombstone checks correlated (0.6)

CHANGELOG.md

@@ -7,6 +7,7 @@ Notable changes between releases. Detailed migration notes for storage transitio
 ### Upgrading
 
 - The schema stays at v40. `awa migrate` applies one new repeatable patch, `runtime_maintenance_eligibility`, which creates the `awa.runtime_maintenance_eligibility` table and locks no existing table. External migration runners apply `R__runtime_maintenance_eligibility.sql` after `V40`. Runtime roles need `SELECT, INSERT, DELETE` on the table; the default privileges in the security guide grant them. 0.6.16 binaries ignore the table, so binaries and the patch can roll in either order.
+- `awa migrate` also applies a new repeatable patch, `correlated_claim_tombstones`, which replaces `claim_ready_runtime` in every queue-storage schema and takes no table lock. External migration runners apply `R__correlated_claim_tombstones.sql` after `V40`, as the owner of the claim functions. 0.6.16 binaries run unchanged against the patched functions, so binaries and the patch can roll in either order.
 
 ### Added
 
@@ -19,6 +20,7 @@ Notable changes between releases. Detailed migration notes for storage transitio
 ### Fixed
 
 - `admin::list_jobs` and `admin::job_kind_overviews` no longer omit a failed job, or leave it out of its kind's count, while queue prune moves its row out of the slot being pruned. Both read several queries from one snapshot and release partition locks between them, so the queries after a prune committed saw neither the pruned slot's rows nor the moved ones. These reads and the point reads such as `admin::get_job` now compare the relfilenode of every ring partition before their snapshot and after it ends, and read again in a new snapshot when prune or the terminal-count rollup truncated one in between, where 0.6.16 read a single job again only when it found nothing. After three reads that each overlapped a truncation, the fourth locks the ring partitions before taking its snapshot and holds them until it commits, so a read never returns a result a truncation could have changed; prune skips its pass if it cannot get its locks within its lock timeout. The comparison adds two catalog lookups to each read and takes no partition locks, so prune is blocked no more often than before.
+- Claiming from a ready slot that holds tombstones no longer scans the slot's whole tombstone partition on every claim (#713, reported and fixed by @Wallace99). Each claim checks its candidate rows against `ready_tombstones`, and for a small batch PostgreSQL could plan that check as a hashed subplan over the partition instead of one index probe per row: a 32-row claim read 10,000 tombstone rows in the regression test. The check now carries `OFFSET 0`, which keeps it correlated. Tombstones are left by cancelled, retried and moved jobs, so the cost grew with churn in the slot. Applied by the `correlated_claim_tombstones` schema patch; no configuration or protocol change.
 - Cleanup, cron sync, batch operations, metadata refresh, rollup and ring rotation no longer hold back promotion and rescue (#303). The maintenance leader ran every branch from one loop, so while one branch's statement waited on a lock or ran a long scan, scheduled and retryable jobs were not promoted and stale, overdue and timed-out attempts were not rescued. With one `awa.cron_jobs` row locked by another session, the leader's cron sync stalled both for as long as the lock was held, and stopped cron from firing. Promotion and the three rescues now run in their own loop within the leader term, and the cron evaluator starts before the startup cron sync. A rescue pass still delays promotion, and the other branches still delay one another: a stuck cleanup or cron sync still holds back ring rotation and prune. Leadership and the shutdown grace are unchanged, and no configuration or schema change is needed. A leader holds one more pool connection at a time. While a rescue pass runs, lease prune may report more `blocked` outcomes and claim prune may wait out a rescue batch. See [ADR-051](docs/adr/051-maintenance-recovery-loop.md).
 - One completion waiting on another transaction's lock no longer stalls every other completion in the runtime. A transaction that changes a running attempt, such as `admin::cancel_tx` inside an application's transaction, holds the attempt's completion lock until it commits. If the attempt's handler returned `Completed` meanwhile, the completion batch waited on that lock, so no other job's completion committed. Once the in-flight completion allowance filled, the runtime stopped claiming in every queue until the transaction ended. With receipts on (the default), a completion batch now leaves an attempt whose lock is held out of the batch, and the worker retries that attempt with backoff for up to two minutes without holding a pool connection, then leaves it to stale-heartbeat rescue as before. Non-success outcomes and kinds with follow-up enqueues already finalized one job at a time and still wait on the lock. With receipts off, a batch still waits. No configuration or schema change. See [a completion waits on an application's transaction](docs/troubleshooting.md#a-completion-waits-on-an-applications-transaction).
 

awa-model/migrations/patches/correlated_claim_tombstones.sql

@@ -0,0 +1,88 @@
+-- Schema patch correlated_claim_tombstones: keep claim tombstone checks on
+-- indexed probes.
+--
+-- A claim marks each candidate row with EXISTS over ready_tombstones. For a
+-- small batch the planner can turn that EXISTS into a hashed subplan, which
+-- scans a whole tombstone partition on every claim although the lookup index
+-- exists. OFFSET 0 in the subquery keeps it correlated, one index probe per
+-- candidate row. Tracking: https://github.com/hardbyte/awa/pull/713
+--
+-- Applied after the v023 definitions, whose installer writes the fenced check
+-- into new substrates. The block below rewrites the claim functions already
+-- installed in each queue-storage schema: <schema>.claim_ready_runtime and
+-- each <schema>.claim_ready_slot_<n>. It replaces each function from its own
+-- definition, so signature, owner, privileges and settings are unchanged, and
+-- a function that already carries the fence is left alone.
+--
+-- Live load: function replacement only. It takes no lock on any awa or
+-- queue-storage table and its work scales with the number of claim functions,
+-- not with job count. A claim in flight finishes on the old body; the next
+-- call uses the new one.
+--
+-- 0.6 delivery: the 0.6 series cannot add a migration version, so this ships as
+-- an idempotent schema patch. `awa migrate` applies it when the schema is
+-- exactly at v040 and the probe in migrations.rs reports it missing;
+-- `awa migrate --sql` prints it as a repeatable R__ script for external
+-- runners. Function names, arguments, results and the claim protocol are
+-- unchanged, so released 0.6 binaries run the same against a patched or an
+-- unpatched schema. Run it as the owner of the claim functions, or as a
+-- migrator that can SET ROLE to it.
+
+DO $$
+DECLARE
+    v_schema TEXT;
+    v_function RECORD;
+    v_definition TEXT;
+    v_updated TEXT;
+BEGIN
+    FOR v_schema IN
+        SELECT n.nspname
+        FROM pg_namespace AS n
+        WHERE has_schema_privilege(current_user, n.oid, 'USAGE')
+          AND EXISTS (
+              SELECT 1 FROM pg_proc AS awa_p
+              WHERE awa_p.pronamespace = n.oid
+                AND awa_p.proname = 'claim_ready_runtime'
+                AND oidvectortypes(awa_p.proargtypes)
+                    = 'text, bigint, double precision, double precision'
+          )
+          AND EXISTS (
+              SELECT 1 FROM pg_class AS awa_c
+              WHERE awa_c.relnamespace = n.oid
+                AND awa_c.relname = 'ready_tombstones'
+                AND awa_c.relkind = 'p'
+          )
+        ORDER BY n.nspname
+    LOOP
+        -- Serialize with a concurrent prepare_schema() of the same schema. This
+        -- is an advisory lock, not a table lock: claims and completions do not
+        -- take it.
+        PERFORM pg_advisory_xact_lock(
+            hashtextextended(format('awa.queue_storage.install:%I', v_schema), 0)
+        );
+        FOR v_function IN
+            SELECT p.oid
+            FROM pg_proc AS p
+            JOIN pg_namespace AS n ON n.oid = p.pronamespace
+            WHERE n.nspname = v_schema
+              AND p.prokind = 'f'
+              AND (p.proname = 'claim_ready_runtime'
+                   OR p.proname ~ '^claim_ready_slot_[0-9]+$')
+            ORDER BY p.oid
+        LOOP
+            v_definition := pg_get_functiondef(v_function.oid);
+            -- Only a check that closes straight after the lane_seq comparison
+            -- is unfenced; one already followed by OFFSET does not match.
+            v_updated := regexp_replace(
+                v_definition,
+                '(AND tomb[.]lane_seq = ready[.]lane_seq)([[:space:]]*[)])',
+                '\1 OFFSET 0\2',
+                'g'
+            );
+            IF v_updated <> v_definition THEN
+                EXECUTE v_updated;
+            END IF;
+        END LOOP;
+    END LOOP;
+END;
+$$;

awa-model/migrations/v023_install_queue_storage_substrate.sql

@@ -4084,6 +4084,8 @@ BEGIN
                           AND tomb.priority = ready.priority
                           AND tomb.enqueue_shard = ready.enqueue_shard
                           AND tomb.lane_seq = ready.lane_seq
+                        -- Keep this bounded claim on indexed probes, not a full-partition hash.
+                        OFFSET 0
                     ) AS is_tombstone
                 FROM %1$I.ready_entries AS ready
                 WHERE ready.queue = p_queue
@@ -4253,6 +4255,8 @@ BEGIN
                           AND tomb.priority = ready.priority
                           AND tomb.enqueue_shard = ready.enqueue_shard
                           AND tomb.lane_seq = ready.lane_seq
+                        -- Keep this bounded claim on indexed probes, not a full-partition hash.
+                        OFFSET 0
                     ) AS is_tombstone
                 FROM %1$I.ready_entries AS ready
                 WHERE ready.queue = p_queue

awa-model/src/migrations.rs

@@ -397,6 +397,32 @@ macro_rules! claim_batch_job_index_pending_sql {
     };
 }
 
+/// Rows naming what still lacks the `correlated_claim_tombstones` patch: the
+/// installer, and each claim function of an installed queue-storage substrate
+/// whose tombstone check closes straight after the `lane_seq` comparison, with
+/// no `OFFSET 0` to keep it correlated.
+macro_rules! correlated_claim_tombstones_pending_sql {
+    () => {
+        "SELECT 'awa.install_queue_storage_substrate'::text AS target \
+         WHERE EXISTS ( \
+             SELECT 1 FROM pg_proc AS p \
+             WHERE p.pronamespace = 'awa'::regnamespace \
+               AND p.proname = 'install_queue_storage_substrate' \
+               AND p.prosrc ~ 'AND tomb[.]lane_seq = ready[.]lane_seq[[:space:]]*[)]') \
+         UNION ALL \
+         SELECT format('claim function %s.%s', n.nspname, p.proname) \
+         FROM pg_proc AS p \
+         JOIN pg_namespace AS n ON n.oid = p.pronamespace \
+         WHERE p.prokind = 'f' \
+           AND (p.proname = 'claim_ready_runtime' OR p.proname ~ '^claim_ready_slot_[0-9]+$') \
+           AND p.prosrc ~ 'AND tomb[.]lane_seq = ready[.]lane_seq[[:space:]]*[)]' \
+           AND EXISTS (SELECT 1 FROM pg_class AS c \
+                       WHERE c.relnamespace = n.oid \
+                         AND c.relname = 'ready_tombstones' \
+                         AND c.relkind = 'p')"
+    };
+}
+
 /// SQL returning one `text` row per object or schema that still needs the
 /// named patch, for patches that touch more than one place.
 fn patch_targets_sql(name: &str) -> Option<&'static str> {
@@ -412,6 +438,10 @@ fn patch_targets_sql(name: &str) -> Option<&'static str> {
         "claim_batch_job_index" => {
             Some(concat!(claim_batch_job_index_pending_sql!(), " ORDER BY 1"))
         }
+        "correlated_claim_tombstones" => Some(concat!(
+            correlated_claim_tombstones_pending_sql!(),
+            " ORDER BY 1"
+        )),
         _ => None,
     }
 }
@@ -518,6 +548,23 @@ pub const SCHEMA_PATCHES: &[SchemaPatch] = &[
         applied_probe:
             "SELECT to_regclass('awa.runtime_maintenance_eligibility') IS NOT NULL",
     },
+    SchemaPatch {
+        name: "correlated_claim_tombstones",
+        description: "Keep claim tombstone checks on indexed probes (awa 0.7 migration v056)",
+        // The v023 definitions first, so the installer writes the fenced
+        // check into new substrates; the patch then rewrites the claim
+        // functions of every installed queue-storage schema.
+        sql: concat!(
+            include_str!("../migrations/v023_install_queue_storage_substrate.sql"),
+            "\n",
+            include_str!("../migrations/patches/correlated_claim_tombstones.sql"),
+        ),
+        applied_probe: concat!(
+            "SELECT NOT EXISTS (",
+            correlated_claim_tombstones_pending_sql!(),
+            ")"
+        ),
+    },
 ];
 
 /// Schema patches a database still needs to reach this binary's schema shape.

awa-model/tests/claim_tombstone_plan_test.rs

@@ -0,0 +1,131 @@
+use awa_model::{migrations, QueueStorage, QueueStorageConfig};
+use sqlx::PgPool;
+
+type TestResult = Result<(), Box<dyn std::error::Error>>;
+
+const PATCH: &str = "correlated_claim_tombstones";
+const FENCE: &str =
+    "                        -- Keep this bounded claim on indexed probes, not a full-partition hash.\n                        OFFSET 0\n";
+
+async fn patch_pending(pool: &PgPool) -> Result<Vec<String>, Box<dyn std::error::Error>> {
+    Ok(migrations::schema_patch_status(pool)
+        .await?
+        .into_iter()
+        .find(|status| status.name == PATCH)
+        .expect("the patch is registered")
+        .needs_repair)
+}
+
+#[sqlx::test(migrations = false)]
+async fn claim_tombstone_patch_keeps_bounded_claims_on_index_probes(pool: PgPool) -> TestResult {
+    migrations::run(&pool).await?;
+    assert!(patch_pending(&pool).await?.is_empty());
+    let patch = migrations::SCHEMA_PATCHES
+        .iter()
+        .find(|patch| patch.name == PATCH)
+        .expect("the patch is registered");
+
+    for receipts in [false, true] {
+        let schema = if receipts {
+            "claim_receipts"
+        } else {
+            "claim_legacy"
+        };
+        let store = QueueStorage::new(QueueStorageConfig {
+            schema: schema.into(),
+            queue_slot_count: 4,
+            lease_slot_count: 2,
+            claim_slot_count: 2,
+            lease_claim_receipts: receipts,
+            ..Default::default()
+        })?;
+        store.install(&pool).await?;
+        store.enqueue_batch(&pool, "probe", 2, 20_000).await?;
+        sqlx::raw_sql(&format!(
+            "INSERT INTO {schema}.ready_tombstones
+             (ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id)
+             SELECT ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id
+             FROM {schema}.ready_entries WHERE lane_seq % 2 = 0;
+             ANALYZE {schema}.ready_entries;
+             ANALYZE {schema}.ready_tombstones;"
+        ))
+        .execute(&pool)
+        .await?;
+
+        // Put the released function bodies back, so the patch is exercised on
+        // installed functions and not only through the new installer.
+        let functions: Vec<(i64, String, String)> = sqlx::query_as(
+            "SELECT p.oid::bigint, p.proname::text, pg_get_functiondef(p.oid)
+             FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace
+             WHERE n.nspname = $1 AND (p.proname = 'claim_ready_runtime'
+                 OR p.proname ~ '^claim_ready_slot_[0-9]+$') ORDER BY p.oid",
+        )
+        .bind(schema)
+        .fetch_all(&pool)
+        .await?;
+        assert_eq!(functions.len(), 1);
+        for (_, _, definition) in &functions {
+            let old = definition.replace(FENCE, "");
+            assert_ne!(&old, definition, "fixture must restore the old lookup");
+            sqlx::raw_sql(&old).execute(&pool).await?;
+        }
+        let mut pending = patch_pending(&pool).await?;
+        pending.sort();
+        let mut expected: Vec<String> = functions
+            .iter()
+            .map(|(_, name, _)| format!("claim function {schema}.{name}"))
+            .collect();
+        expected.sort();
+        assert_eq!(pending, expected, "the probe names each unfenced function");
+
+        sqlx::raw_sql(patch.sql).execute(&pool).await?;
+        assert!(patch_pending(&pool).await?.is_empty());
+        for (oid, _, _) in &functions {
+            let upgraded: String = sqlx::query_scalar("SELECT pg_get_functiondef($1::oid)")
+                .bind(*oid)
+                .fetch_one(&pool)
+                .await?;
+            assert!(upgraded.contains("AND tomb.lane_seq = ready.lane_seq OFFSET 0"));
+            sqlx::raw_sql(patch.sql).execute(&pool).await?;
+            let replayed: String = sqlx::query_scalar("SELECT pg_get_functiondef($1::oid)")
+                .bind(*oid)
+                .fetch_one(&pool)
+                .await?;
+            assert_eq!(
+                upgraded, replayed,
+                "replaying the patch must preserve the body"
+            );
+        }
+
+        for mode in ["force_custom_plan", "force_generic_plan"] {
+            let mut tx = pool.begin().await?;
+            sqlx::raw_sql(&format!(
+                "SET LOCAL plan_cache_mode = {mode};
+                 SET LOCAL work_mem = '125MB'; SET LOCAL random_page_cost = 1.2;"
+            ))
+            .execute(&mut *tx)
+            .await?;
+            let claimed: Vec<i64> = sqlx::query_scalar(&format!(
+                "SELECT lane_seq FROM {schema}.claim_ready_runtime('probe',32,0::float8,0::float8)"
+            ))
+            .fetch_all(&mut *tx)
+            .await?;
+            assert_eq!(claimed.len(), 16, "only live rows in the batch are claimed");
+            assert!(claimed.iter().all(|lane| lane % 2 == 1));
+            let scanned: i64 = sqlx::query_scalar(
+                "SELECT COALESCE(sum(seq_tup_read),0)::bigint
+                 FROM pg_stat_xact_all_tables
+                 WHERE schemaname=$1 AND relname LIKE 'ready_tombstones%'",
+            )
+            .bind(schema)
+            .fetch_one(&mut *tx)
+            .await?;
+            assert_eq!(
+                scanned, 0,
+                "{schema}: {mode} must not scan tombstone partitions"
+            );
+            tx.rollback().await?;
+        }
+    }
+    Ok(())
+}