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

CHANGELOG.md

@@ -52,6 +52,7 @@ Notable changes between releases. Detailed migration notes for storage transitio
 
 ### Fixed
 
+- 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, migration v060). 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. v060 rewrites `claim_ready_runtime` and every `claim_ready_slot_<n>` in the default and every trusted custom queue-storage schema in place, without reinstalling tables or triggers; it takes no table locks, and binaries before it claim unchanged. 0.6.17 made the same change through its `correlated_claim_tombstones` schema patch. See [v060](docs/upgrade-0.6-to-0.7.md#v060-correlated-claim-tombstone-checks).
 - **The `awa_job_output_schema_missing` alert rule imports.** Its query had no expression or datasource and the rule had no threshold, so Grafana refused to provision it. It now fires when a queue drops outputs with reason `schema_missing` for 10 minutes.
 - **A completion batch no longer waits to fill while attempts keep arriving (#732).** A completion shard restarted its flush timer on every arrival, so while attempts arrived faster than one per flush interval a batch was written only once it reached `AWA_COMPLETION_BATCH_SIZE`. With the defaults, a completion could wait up to about 511 ÷ the shard's completion rate before being written: about 250 ms at 2,000 completions per second per shard. The flush deadline now starts when the first attempt of a batch arrives, so a batch is written when it is full or `AWA_COMPLETION_FLUSH_MS` (default 1 ms) after its first attempt, whichever comes first.
 - **A queue lane is usable by every runtime role, whichever role creates it

awa-model/migrations/v023_install_queue_storage_substrate.sql

@@ -4661,6 +4661,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
@@ -4974,6 +4976,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_%4$s AS ready
                 WHERE ready.queue = p_queue

awa-model/migrations/v060_correlated_claim_tombstones.sql

@@ -0,0 +1,101 @@
+-- v060: keep claim tombstone checks on indexed probes.
+-- Tracking: https://github.com/hardbyte/awa/pull/713
+--
+-- 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.
+--
+-- V23_UP writes the fenced check into new substrates. This file rewrites the
+-- claim functions already installed in the default and every discovered custom
+-- queue-storage schema: <schema>.claim_ready_runtime and the
+-- <schema>.claim_ready_slot_<n> from v051 for each installed queue slot, each
+-- found by its exact signature. It replaces each function from its
+-- own definition, so signature, owner, privileges and settings are unchanged,
+-- and a function that already carries the fence (a fresh install, a replay, or
+-- a database coming from 0.6.17, whose correlated_claim_tombstones schema
+-- patch made the same change) is left alone. No table DDL, no trigger change
+-- and no substrate reinstallation.
+--
+-- N-1 binaries call the same functions with the same arguments and results
+-- and the claim protocol is unchanged, so they run the same before and after
+-- the migration and rolling back binaries needs no schema change.
+--
+-- 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.
+--
+-- External runners apply the complete `awa migrate --sql --from 59 --to 60`
+-- bundle, which refreshes trusted queue-storage discovery before this block.
+-- 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_queue_slots INT;
+    v_function RECORD;
+    v_definition TEXT;
+    v_updated TEXT;
+BEGIN
+    FOR v_schema IN
+        SELECT discovered.schema_name FROM awa.queue_storage_schemas() AS discovered
+    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)
+        );
+        IF to_regclass(format('%I.queue_ring_state', v_schema)) IS NULL THEN
+            CONTINUE;
+        END IF;
+        EXECUTE format(
+            'SELECT slot_count FROM %I.queue_ring_state WHERE singleton = TRUE',
+            v_schema
+        )
+        INTO v_queue_slots;
+
+        -- Only the installer's functions: claim_ready_runtime and the
+        -- claim_ready_slot_<n> for each installed queue slot, each by its
+        -- exact signature. Another function that shares one of these names
+        -- is not touched.
+        FOR v_function IN
+            SELECT candidate.proc
+            FROM (
+                SELECT to_regprocedure(format(
+                    '%I.claim_ready_runtime(text,bigint,double precision,double precision)',
+                    v_schema
+                )) AS proc
+                UNION ALL
+                SELECT to_regprocedure(format(
+                    '%I.claim_ready_slot_%s(text,int,bigint,smallint,smallint,bigint,bigint,double precision,timestamptz,timestamptz,double precision,bigint)',
+                    v_schema,
+                    slot
+                ))
+                FROM generate_series(0, COALESCE(v_queue_slots, 0) - 1) AS slot
+            ) AS candidate
+            WHERE candidate.proc IS NOT NULL
+        LOOP
+            v_definition := pg_get_functiondef(v_function.proc::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;
+$$;
+
+INSERT INTO awa.schema_version (version, description)
+VALUES (60, 'Keep claim tombstone checks correlated')
+ON CONFLICT (version) DO NOTHING;

awa-model/src/migrations.rs

@@ -5,7 +5,7 @@ use sqlx::{Connection, PgPool};
 use tracing::{info, warn};
 
 /// Current schema version.
-pub const CURRENT_VERSION: i32 = 59;
+pub const CURRENT_VERSION: i32 = 60;
 
 /// Migrations that require an exclusive (no-live-runtime) upgrade window.
 ///
@@ -327,6 +327,11 @@ const MIGRATIONS: &[(i32, &str, &[&str])] = &[
         "Spread retry backoff over half to one and a half times the base delay",
         &[V59_UP],
     ),
+    (
+        60,
+        "Keep claim tombstone checks correlated",
+        &[V23_UP, V60_UP],
+    ),
 ];
 
 const V1_UP: &str = include_str!("../migrations/v001_canonical_schema.sql");
@@ -394,6 +399,7 @@ const V56_UP: &str = include_str!("../migrations/v056_claim_batch_job_index_shap
 const V57_UP: &str = include_str!("../migrations/v057_insert_jobs_compat_batch_bounded_wait.sql");
 const V58_UP: &str = include_str!("../migrations/v058_provision_lane_row_first.sql");
 const V59_UP: &str = include_str!("../migrations/v059_backoff_wider_jitter.sql");
+const V60_UP: &str = include_str!("../migrations/v060_correlated_claim_tombstones.sql");
 
 /// Old version numbers from pre-0.4 releases that used V3/V4/V5 numbering.
 /// Also tolerates the unreleased inline-V6 branch numbering used during review.

awa-model/tests/claim_tombstone_plan_test.rs

@@ -0,0 +1,160 @@
+use awa_model::{audited_sql, migrations, QueueStorage, QueueStorageConfig};
+use sqlx::PgPool;
+
+type TestResult = Result<(), Box<dyn std::error::Error>>;
+
+#[sqlx::test(migrations = false)]
+async fn claim_tombstone_upgrade_keeps_bounded_claims_on_index_probes(pool: PgPool) -> TestResult {
+    migrations::run(&pool).await?;
+    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(audited_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?;
+
+        // Exercise the upgrade on installed functions, not just the new installer.
+        let functions: Vec<(i64, String)> = sqlx::query_as(
+            "SELECT p.oid::bigint, 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(), 5);
+        for (_, definition) in &functions {
+            let old = definition.replace(
+                "                        -- Keep this bounded claim on indexed probes, not a full-partition hash.\n                        OFFSET 0\n",
+                "",
+            );
+            assert_ne!(&old, definition, "fixture must restore the old lookup");
+            sqlx::raw_sql(audited_sql(old)).execute(&pool).await?;
+        }
+
+        // Functions that share a claim function's name but are not one the
+        // installer made, each with an unfenced check: an overload of the
+        // runtime and of slot 0, and a slot beyond the installed slot count.
+        // The migration must leave all three alone.
+        let decoy_body = format!(
+            "RETURNS boolean LANGUAGE sql AS $decoy$
+             SELECT EXISTS (
+                 SELECT 1 FROM {schema}.ready_tombstones AS tomb, {schema}.ready_entries AS ready
+                 WHERE tomb.queue = ready.queue
+                   AND tomb.lane_seq = ready.lane_seq
+             ) $decoy$;"
+        );
+        sqlx::raw_sql(audited_sql(format!(
+            "CREATE FUNCTION {schema}.claim_ready_runtime(p_queue text) {decoy_body}
+             CREATE FUNCTION {schema}.claim_ready_slot_0(p_queue text) {decoy_body}
+             CREATE FUNCTION {schema}.claim_ready_slot_7(
+                 text, int, bigint, smallint, smallint, bigint, bigint,
+                 double precision, timestamptz, timestamptz, double precision, bigint
+             ) {decoy_body}"
+        )))
+        .execute(&pool)
+        .await?;
+        let decoys = "SELECT p.oid::bigint, pg_get_functiondef(p.oid)
+             FROM pg_proc p JOIN pg_namespace n ON n.oid = p.pronamespace
+             WHERE n.nspname = $1 AND p.prorettype = 'boolean'::regtype
+               AND (p.proname = 'claim_ready_runtime'
+                    OR p.proname ~ '^claim_ready_slot_[0-9]+$')
+             ORDER BY p.oid";
+        let decoys_before: Vec<(i64, String)> =
+            sqlx::query_as(decoys).bind(schema).fetch_all(&pool).await?;
+        assert_eq!(decoys_before.len(), 3);
+
+        for (_, _, sql) in migrations::migration_sql_range(59, 60) {
+            sqlx::raw_sql(audited_sql(sql)).execute(&pool).await?;
+        }
+
+        let decoys_after: Vec<(i64, String)> =
+            sqlx::query_as(decoys).bind(schema).fetch_all(&pool).await?;
+        assert_eq!(
+            decoys_before, decoys_after,
+            "{schema}: only the installed claim functions are rewritten"
+        );
+        sqlx::raw_sql(audited_sql(format!(
+            "DROP FUNCTION {schema}.claim_ready_runtime(text);
+             DROP FUNCTION {schema}.claim_ready_slot_0(text);
+             DROP FUNCTION {schema}.claim_ready_slot_7(
+                 text, int, bigint, smallint, smallint, bigint, bigint,
+                 double precision, timestamptz, timestamptz, double precision, bigint
+             );"
+        )))
+        .execute(&pool)
+        .await?;
+        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(include_str!(
+                "../migrations/v060_correlated_claim_tombstones.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,
+                "migration replay must preserve the body"
+            );
+        }
+
+        for mode in ["force_custom_plan", "force_generic_plan"] {
+            let mut tx = pool.begin().await?;
+            sqlx::raw_sql(audited_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(audited_sql(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(())
+}

docs/upgrade-0.6-to-0.7.md

@@ -366,6 +366,14 @@ v059 replaces the body of `awa.backoff_duration(attempt, max_attempts)`, which s
 
 It is one `CREATE OR REPLACE FUNCTION` that keeps the function's signature, volatility, owner and privileges, and takes no lock on any `awa` or queue-storage table. Binaries before v059 call the same function from the same statements, so they get the new delays as soon as the migration commits and can run before or after it. Jobs already waiting to retry keep their scheduled time. External migration runners apply `awa migrate --sql --from 58 --to 59`.
 
+## v060: correlated claim tombstone checks
+
+A claim checks each candidate row against `ready_tombstones`, which records rows that were cancelled, retried or moved before they were claimed. For a small batch PostgreSQL could plan that `EXISTS` check as a hashed subplan that reads the slot's whole tombstone partition on every claim, although the lookup index exists, so the cost of each claim grew with churn in the slot. The check now carries `OFFSET 0`, which keeps it correlated: one index probe per candidate row.
+
+v060 rewrites `<schema>.claim_ready_runtime` and each `<schema>.claim_ready_slot_<n>` in the default and every discovered custom queue-storage schema from their own definitions, so signatures, owners, privileges and settings are unchanged. A function that already carries the check is left alone, so a replay, a fresh install and a database coming from 0.6.17 (whose `correlated_claim_tombstones` schema patch made the same change) are no-ops. It reinstalls no table or trigger and takes no lock on any `awa` or queue-storage table; a claim in flight finishes on the old body.
+
+Binaries before v060 call the same functions with the same arguments and results, so they can run before or after the migration. External migration runners apply the complete `awa migrate --sql --from 59 --to 60` bundle as the owner of the claim functions, or as a role that can `SET ROLE` to it.
+
 ## The one gate
 
 `awa migrate` on a 0.7 binary refuses to apply migrations unless one of these holds: