fix(storage): keep claim tombstone checks correlated

CHANGELOG.md

@@ -39,6 +39,8 @@ Notable changes between releases. Detailed migration notes for storage transitio
 
 ### Fixed
 
+- Claim batches keep tombstone checks as indexed probes instead of hashing an entire tombstone partition for each small batch (migration v056). Existing claim functions are replaced in place, including trusted custom substrates, without reinstalling tables or triggers.
+
 - **A queue lane is usable by every runtime role, whichever role creates it
   first (migration v053).** A lane's enqueue and claim sequences belonged to
   the role that first enqueued into, claimed from or promoted into it, so with

awa-model/migrations/v023_install_queue_storage_substrate.sql

@@ -4612,6 +4612,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
@@ -4925,6 +4927,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/v056_correlated_claim_tombstones.sql

@@ -0,0 +1,46 @@
+-- v056: keep bounded claim tombstone checks correlated.
+-- N-1 binaries use the same function signatures, results and claim protocol.
+-- Function replacement only: no table DDL or substrate reinstallation. Work
+-- scales with installed claim functions, not job count. Apply with an owner-
+-- capable role; external runners must include V23_UP's discovery bootstrap.
+-- Tracking: https://github.com/hardbyte/awa/pull/713
+
+DO $$
+DECLARE
+    v_schema TEXT;
+    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
+        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]+$')
+        LOOP
+            v_definition := pg_get_functiondef(v_function.oid);
+            -- Match only an unfenced correlated check; replay leaves an
+            -- existing OFFSET intact. Preserve signatures, ownership and ACLs.
+            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 (56, '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 = 55;
+pub const CURRENT_VERSION: i32 = 56;
 
 /// Migrations that require an exclusive (no-live-runtime) upgrade window.
 ///
@@ -307,6 +307,11 @@ const MIGRATIONS: &[(i32, &str, &[&str])] = &[
         "Batch producer insert with a short enqueue lane lock window",
         &[V23_UP, V55_UP],
     ),
+    (
+        56,
+        "Keep claim tombstone checks correlated",
+        &[V23_UP, V56_UP],
+    ),
 ];
 
 const V1_UP: &str = include_str!("../migrations/v001_canonical_schema.sql");
@@ -370,6 +375,7 @@ const V52_UP: &str = include_str!("../migrations/v052_retained_job_output.sql");
 const V53_UP: &str = include_str!("../migrations/v053_definer_lane_provisioning.sql");
 const V54_UP: &str = include_str!("../migrations/v054_claim_batch_job_index.sql");
 const V55_UP: &str = include_str!("../migrations/v055_insert_jobs_compat_batch.sql");
+const V56_UP: &str = include_str!("../migrations/v056_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,110 @@
+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?;
+        }
+        for (_, _, sql) in migrations::migration_sql_range(55, 56) {
+            sqlx::raw_sql(audited_sql(sql)).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/v056_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(())
+}