From b817fbc4e856d6a9ac3e171221386899a1b74c8b Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Wed, 11 Nov 2026 03:44:40 +0000 Subject: [PATCH] fix(storage): keep claim tombstone checks correlated (0.6) A claim marks each candidate row with EXISTS over ready_tombstones. For a small batch PostgreSQL could plan that as a hashed subplan, reading the slot's whole tombstone partition on every claim although the lookup index exists. OFFSET 0 in the subquery keeps it one index probe per candidate row. The installer writes the fenced check into new substrates. The correlated_claim_tombstones schema patch rewrites claim_ready_runtime in every installed queue-storage schema from its own definition, so signature, owner, privileges and settings are unchanged. Function replacement only; no table lock. 0.6 port of #713. Co-authored-by: Wallace99 <5481339+Wallace99@users.noreply.github.com> --- CHANGELOG.md | 2 + .../patches/correlated_claim_tombstones.sql | 88 ++++++++++++ .../v023_install_queue_storage_substrate.sql | 4 + awa-model/src/migrations.rs | 47 +++++++ awa-model/tests/claim_tombstone_plan_test.rs | 131 ++++++++++++++++++ awa/tests/patch_upgrade_rehearsal_test.rs | 10 +- docs/migrations.md | 2 + 7 files changed, 281 insertions(+), 3 deletions(-) create mode 100644 awa-model/migrations/patches/correlated_claim_tombstones.sql create mode 100644 awa-model/tests/claim_tombstone_plan_test.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index c5765c00d..89aa388b4 100644 --- a/CHANGELOG.md +++ b/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). diff --git a/awa-model/migrations/patches/correlated_claim_tombstones.sql b/awa-model/migrations/patches/correlated_claim_tombstones.sql new file mode 100644 index 000000000..7294156be --- /dev/null +++ b/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: .claim_ready_runtime and +-- each .claim_ready_slot_. 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; +$$; diff --git a/awa-model/migrations/v023_install_queue_storage_substrate.sql b/awa-model/migrations/v023_install_queue_storage_substrate.sql index 6a688758b..aa906d39d 100644 --- a/awa-model/migrations/v023_install_queue_storage_substrate.sql +++ b/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 diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index b53b04dd3..506e256a9 100644 --- a/awa-model/src/migrations.rs +++ b/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. diff --git a/awa-model/tests/claim_tombstone_plan_test.rs b/awa-model/tests/claim_tombstone_plan_test.rs new file mode 100644 index 000000000..0d2dc9847 --- /dev/null +++ b/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>; + +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, Box> { + 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 = 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 = 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(()) +} diff --git a/awa/tests/patch_upgrade_rehearsal_test.rs b/awa/tests/patch_upgrade_rehearsal_test.rs index 207c5111c..57c65385e 100644 --- a/awa/tests/patch_upgrade_rehearsal_test.rs +++ b/awa/tests/patch_upgrade_rehearsal_test.rs @@ -758,10 +758,12 @@ async fn released_patch_upgrade_and_rollback() { "claim_batch_job_index", "insert_jobs_compat_batch", "prune_controls", - "runtime_maintenance_eligibility" + "runtime_maintenance_eligibility", + "correlated_claim_tombstones" ], "a released 0.6.15 deployment needs the claim batch index fix, the bounded-wait \ - batch function, prune controls and the maintenance eligibility table" + batch function, prune controls, the maintenance eligibility table and the \ + correlated tombstone checks" ); } @@ -830,7 +832,7 @@ async fn released_patch_upgrade_and_rollback() { if released_version == "0.6.9" { assert!(patch_sql.contains("R__prepared_lane_sequences")); } else { - assert_eq!(patch_sql.matches("-- Schema patch R__").count(), 4); + assert_eq!(patch_sql.matches("-- Schema patch R__").count(), 5); assert!(!patch_sql.contains("R__prepared_lane_sequences")); assert!( patch_sql.contains("R__claim_batch_job_index"), @@ -838,6 +840,7 @@ async fn released_patch_upgrade_and_rollback() { ); assert!(patch_sql.contains("R__insert_jobs_compat_batch")); assert!(patch_sql.contains("R__runtime_maintenance_eligibility")); + assert!(patch_sql.contains("R__correlated_claim_tombstones")); } assert!( !patch_sql.contains("-- Migration V"), @@ -985,6 +988,7 @@ async fn released_patch_upgrade_and_rollback() { "Schema patch insert_jobs_compat_batch: applied".to_string(), "Schema patch prune_controls: applied".to_string(), "Schema patch runtime_maintenance_eligibility: applied".to_string(), + "Schema patch correlated_claim_tombstones: applied".to_string(), format!("Storage routing: unchanged (state active, engine queue_storage, schema {SCHEMA}"), ] { assert!( diff --git a/docs/migrations.md b/docs/migrations.md index b70578e49..60ff073cd 100644 --- a/docs/migrations.md +++ b/docs/migrations.md @@ -249,6 +249,8 @@ That writes one `V__.sql` file per migration, plus one `R__.claim_ready_runtime` in every queue-storage schema from the function's own definition, adding `OFFSET 0` to the two tombstone lookups, and refreshes the installer so new schemas carry the same check. It replaces functions only and takes no lock on any awa or queue-storage table. Function names, arguments and results are unchanged, so released 0.6 binaries run the same on a patched or an unpatched schema. Run it as the owner of the claim functions. + `R__claim_batch_job_index.sql` indexes compact claim batches by job id. It installs `.ensure_claim_batch_job_index(slot)` in every queue-storage schema and builds no index itself, so it takes no lock on any awa or queue-storage table and is safe inside a long transaction under live traffic. Claim prune calls that function right after truncating a claim partition, which builds the GIN index on `lease_claim_batches.job_ids` while the partition is empty, so the index covers the ring within one claim-ring cycle. Only a valid, non-partial GIN index on `job_ids` counts; a different index or relation already named `lease_claim_batches__job_ids_idx` is left in place with a warning naming it, and that slot keeps scanning until it is dropped or renamed. Do not build the index yourself on a populated partition: the build holds a `SHARE` lock, which blocks claim inserts into that partition, for about half a minute per 100k batches. Run the script as the owner of the queue-storage tables. Binaries before 0.6.15 neither use nor build the index and are unaffected by it, and binaries that use the index read the same rows, by scan, from a schema or partition without it. `awa migrate` reports the schema version, each patch and what it updated, and storage routing before and after. It exits 1 if a patch is still missing somewhere, such as a queue-storage schema the migrating role cannot use, and names it. `awa::migrations::schema_patch_status` gives the same per-patch view from Rust.