diff --git a/.github/workflows/nightly-chaos.yml b/.github/workflows/nightly-chaos.yml index 9c69c0889..a57f22aa2 100644 --- a/.github/workflows/nightly-chaos.yml +++ b/.github/workflows/nightly-chaos.yml @@ -197,6 +197,17 @@ jobs: cargo test --release --package awa --test concurrent_lifecycle_test \ test_concurrent_queue_sweep -- --exact --ignored --nocapture \ | tee artifacts/nightly/rust-queue-sweep.txt + - name: Run per-lane enqueue contention benchmark + if: always() + timeout-minutes: 5 + run: | + set -o pipefail + # Many producers batch-inserting into one queue lane. Gates the + # per-lane rate and p99 through benchmarks/baseline.json, and the + # lane-lock hold time relative to the per-row insert function. + cargo test --release --package awa --test enqueue_lane_contention_test \ + test_enqueue_lane_contention_gate -- --exact --ignored --nocapture \ + | tee artifacts/nightly/rust-enqueue-lane.txt - name: Run MVCC horizon overlap benchmark if: always() timeout-minutes: 10 diff --git a/CHANGELOG.md b/CHANGELOG.md index 065d50db1..2a5afb6c7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,8 +4,110 @@ Notable changes between releases. Detailed migration notes for storage transitio ## [Unreleased] +### Added + +- **`insert_many_skip_unique_conflicts` inserts a batch and skips the jobs + whose unique key is already claimed.** `insert_many` fails the whole batch + with a unique violation, so a producer that replays work had to insert the + batch again one job at a time, taking the lane lock once per job. The new + function returns the inserted rows in input order, with skipped jobs absent. + It needs migration v055 and returns a validation error on an older schema; + `insert_many_copy` (and Python's `insert_many_copy`) skips unique conflicts on + every schema. + +- **The maintenance leader keeps planner statistics fresh on the queue-storage + tables autovacuum never analyses.** Autovacuum does not analyse partitioned + tables, so the thirteen ring parents (`ready_entries`, `done_entries`, + `leases`, ...) had no inherited statistics, and the two small head tables + (`queue_claim_heads`, `queue_enqueue_heads`) never reach the autoanalyze + threshold, so the planner priced lane joins against them at its defaults + while they drove joins over hundreds of thousands of retained rows. The + leader now runs `ANALYZE` on them when it is elected and then every five + minutes. `ClientBuilder::statistics_refresh_interval` (Python + `statistics_refresh_interval_ms`) changes the cadence; zero disables it. One + pass takes under a second on a small ring and a few seconds on a large one. + It analyses one relation per transaction with a 2 s `lock_timeout` (a call that + times out waiting for a lock is retried next interval and recorded as + `skipped_lock`) and never makes prune wait: a prune attempt that + lands while ANALYZE samples the oldest slot's tables returns `blocked` and + retries on its normal backoff. Results are in + `awa.maintenance.analyze.attempts` and `awa.maintenance.analyze.duration` + (scope `heads`, `ring_parent` or `pass`; outcome `ok`, `skipped_privilege`, + `skipped_lock`, `unavailable` or `failed`). +- **Migration v054 adds `awa.analyze_queue_storage_relation(schema, relation)`, + so a runtime role that does not own the tables can refresh statistics.** + `ANALYZE` skips a table its caller may not analyse, and the documented + runtime grants do not include `MAINTAIN`, so the refresh runs through this + `SECURITY DEFINER` function owned by the migration role. It analyses only a + head table or ring parent of a schema that `awa.queue_storage_schemas()` + reports as trusted. It is not executable by `PUBLIC`, because the lock it + takes lasts until the caller's transaction ends and would block ring prune: + the migration grants `EXECUTE` to the roles that already hold INSERT, UPDATE + and DELETE on the head tables (the runtime grantees), and the function + requires the calling session user to hold them. A role created later needs + the documented `GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa`. Run + `awa migrate` as the schema owner (see + [database roles](docs/security/database-roles.md)). A binary running ahead of + the migration probes for the function once per pass and, like one whose + function's owner may not analyse the tables, logs one warning naming what is + missing and skips the refresh without making failing calls or affecting job + processing. Earlier binaries never call the function. + Missing EXECUTE or head-table privileges are checked before each pass, + avoiding a failing helper call for every relation. + +### Changed + +- **COPY and opt-in skip-conflict batch enqueues hold lane locks only for reservation, ready-row writes and commit.** Migration v055 adds `awa.insert_jobs_compat_batch`, which allocates job ids, claims unique keys in key order, registers lanes and prepares rows before reserving each lane's sequence range in lane order. `insert_many_copy` probes its own connection on every call and uses the function when installed, falling back to per-row inserts on older schemas. `insert_many_skip_unique_conflicts` requires the function. Generic `insert_many` keeps its single-statement per-row path on every supported schema, preserving caller transactions when one process writes to databases at different migration levels. All paths return rows in input order and share the same per-lane lock with older binaries. + ### Fixed +- **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 + several runtime roles sharing a group role, any role other than the creator + whose default privileges did not cover it failed with `permission denied for + sequence`. That included the maintenance leader promoting into a shard + created after `enqueue_shards` was raised. Lanes are now created through two + `SECURITY DEFINER` functions per queue-storage schema, `provision_lane` and + `provision_lane_sequences`, owned by the schema owner. Each new sequence gets + the privileges of the head table it backs, explicitly rather than through + default privileges. A runtime without schema `CREATE` can therefore create + lanes itself, and `awa storage prepare-queue` is optional and may run as any + runtime role. After v053 commits, `awa migrate` gives existing lane + sequences to the schema owner and adds their missing grants, one sequence per + transaction with a 300 ms lock wait, and only where it is allowed to, trying + the effective role and then the login role. A sequence an open transaction + holds is skipped and reported as busy; one it may not change is reported + with the roles that can repair it. Neither fails the migration. The new + read-only `awa storage lane-sequences` (or + `awa.lane_sequence_problems(schema)`) lists what remains, with the statements + to run and the roles to run them as, and exits 2 while any remain; + `--repair` (or `awa.repair_lane_sequence(schema, sequence)` per sequence) + repairs what the connected role may. `awa migrate` prints a summary of the leftovers. A binary + running ahead of v053 detects that `provision_lane` is missing and creates + lanes as before, and older binaries create correctly owned lanes once v053 + has refreshed the helpers they call. See "Lane ownership and v053" in the + 0.6 to 0.7 upgrade guide. + +- **Ring prune and terminal-count rollup report `Blocked` when any of their + statements times out on a lock.** Prune sets a short `lock_timeout` (10 ms + by default) before it takes `ACCESS EXCLUSIVE` on a slot's partitions, but + only those `LOCK TABLE` statements and the final `TRUNCATE` turned a + timeout into `PruneOutcome::Blocked`. The timeout covers every later + statement in the transaction too: the post-lock liveness checks read the + lease and claim partitions, the failed-row carry inserts into the current + slot's `done_entries` and terminal-count deltas, and the fold appends + `queue_terminal_rollup_deltas` and deletes `queue_terminal_live_counts` + rows. When one of those waited out the timeout, for example behind + autovacuum truncating a relation or a hosted Postgres's own background + workers, prune failed with a `lock_not_available` (`55P03`) database error, + was logged as an error and skipped the prune backoff. A lock timeout + anywhere in a queue, lease or claim prune, or in a terminal-count rollup + slot, now rolls the whole transaction back and reports `Blocked` + (`blocked_slots` for rollup), so maintenance backs off and retries; nothing + the attempt wrote survives the rollback. A rollup pass stops at its first + blocked slot rather than waiting out the timeout again on each remaining + one. - **Batch enqueues take their lane locks in one order, so they no longer deadlock with each other.** On queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane @@ -15,10 +117,7 @@ Notable changes between releases. Detailed migration notes for storage transitio a `QueueStorage` enqueue could deadlock with it (`40P01`). The two insert paths now sort rows by `(queue, priority, shard)`, queues compared bytewise, before they reach `awa.insert_job_compat`, the order `QueueStorage` already - used, and still return rows in input order. Job ids within a batch that - spans several lanes are now allocated lane by lane rather than in input - order, and when an `insert_many_copy` batch carries the same unique key in - more than one lane, the row in the lowest lane is the one inserted. A + used, and still return rows in input order. A transaction that enqueues with separate statements still takes locks in statement order; see "`deadlock detected` while enqueueing" in the troubleshooting guide. @@ -230,7 +329,7 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Upgrading -This release adds migrations v041 to v052. Run `awa migrate` as the migrator +This release adds migrations v041 to v056. Run `awa migrate` as the migrator role; the [0.6→0.7 upgrade guide](docs/upgrade-0.6-to-0.7.md) has the details for each step. @@ -242,9 +341,9 @@ for each step. entry below has the staged expand, flip and contract steps. - **Rollout order.** Health-gated workers report ready only once the schema reaches their version, so migrate before rolling them out. Producers and - workers without a readiness gate can roll before v050. On PostgreSQL 17 and - earlier, pre-v050 binaries scan every ready partition during maintenance - once v050 is applied, so replace them promptly. + workers without a readiness gate can roll before v050 and v056. On + PostgreSQL 17 and earlier, pre-v050 binaries scan every ready partition + during maintenance once v050 is applied, so replace them promptly. - **Runtime roles no longer need `CREATE` (v049).** Provision lanes with `awa storage prepare-queue` as the migrator before deploying, and grant the runtime role `USAGE, SELECT, UPDATE` on the lane sequences. Existing @@ -595,6 +694,22 @@ for each step. session sets no `lock_timeout`, each lock wait is capped at 2 s; on a timeout (`55P03`, usually behind a long `awa.jobs` read or a backup) the migration rolls back and can be re-run. +- **Enqueueing a batch of ready jobs takes one round trip (v056, + [#522](https://github.com/hardbyte/awa/issues/522)).** + `QueueStorage::enqueue_params_batch`, and `enqueue_params_copy` below its + 2,048-row COPY threshold, send a batch of ready jobs without uniqueness + claims as one call to the new `{schema}.enqueue_ready_rows` function instead + of ten round trips. Batches with uniqueness claims or scheduled jobs still + run in a transaction but use the function for most of the work, and + reserving a lane's sequence range now writes the sequence once instead of + twice. On an earlier revision, at 800 jobs/s in 20-job calls, producer call + p50 fell 19-62% and WAL per job by about 5 bytes, with throughput unchanged. + On a schema without the function, or without `EXECUTE` on it and on + `ring_cursor`, binaries keep using the multi-statement path, which is also + shorter now, and look for the function again after a minute. If a cached + function disappears or loses EXECUTE, the same call rolls back any transaction + it owns and retries its original input once without the function; healthy + cached calls keep the one-statement path. - **Retry-heavy workloads leave far fewer dead tuples.** After the ring authority flip, awa keeps one terminal-count row per lane in `queue_terminal_live_counts` instead of up to 256, and a batch of finished diff --git a/awa-cli/src/main.rs b/awa-cli/src/main.rs index 405fc3023..e111c3574 100644 --- a/awa-cli/src/main.rs +++ b/awa-cli/src/main.rs @@ -273,6 +273,23 @@ fn sanitize_filename(description: &str) -> String { /// mutating commands require an explicit `--context` or `--database-url` /// once more than one context is defined, and are confirmation-gated on /// contexts marked `production = true`. +/// After a migration, say which lane sequences still need repair and how to +/// finish it. Read-only. +async fn print_lane_sequence_leftovers( + pool: &sqlx::PgPool, +) -> Result<(), Box> { + let mut problems = Vec::new(); + for (schema, _) in awa_model::storage::lane_sequence_problem_counts(pool).await? { + problems.extend(awa_model::storage::lane_sequence_problems(pool, &schema).await?); + } + for next_step in awa_model::storage::lane_sequence_next_steps( + problems.iter().map(|problem| (problem, false)), + ) { + eprintln!("warning: {next_step}"); + } + Ok(()) +} + fn access_for(command: &Commands) -> context::Access { use context::Access::{Mutating, ReadOnly}; match command { @@ -328,11 +345,13 @@ fn access_for(command: &Commands) -> context::Access { }, Commands::Storage { command } => match command { StorageCommands::Status => ReadOnly, + StorageCommands::LaneSequences { repair: false, .. } => ReadOnly, StorageCommands::Finalize { check: true, .. } => ReadOnly, StorageCommands::FlipRingAuthority { check: true, .. } => ReadOnly, StorageCommands::Prepare { .. } | StorageCommands::PrepareQueueStorageSchema { .. } | StorageCommands::PrepareQueue { .. } + | StorageCommands::LaneSequences { .. } | StorageCommands::Abort | StorageCommands::EnterMixedTransition { .. } | StorageCommands::Finalize { .. } @@ -811,7 +830,7 @@ enum StorageCommands { #[arg(long)] reset: bool, }, - /// Provision lane sequences as the migrator, before starting a runtime without DDL + /// Provision every priority and shard of a queue ahead of its first enqueue PrepareQueue { /// Queue-storage schema to provision #[arg(long, default_value = "awa")] @@ -826,6 +845,17 @@ enum StorageCommands { #[arg(long, default_value_t = 1)] queue_stripe_count: usize, }, + /// List lane sequences not owned by the schema owner or missing runtime + /// grants, with the statements that fix them. Exits 2 while any remain. + LaneSequences { + /// Queue-storage schema to check + #[arg(long, default_value = "awa")] + schema: String, + /// Reassign and grant first, as far as the connected role (or, failing + /// that, its login role) is allowed to + #[arg(long)] + repair: bool, + }, /// Abort a prepared or mixed-transition storage rollout before final activation Abort, /// Enter mixed transition and begin routing new writes to the prepared engine @@ -1182,6 +1212,7 @@ async fn run() -> Result<(), Box> { }; awa_model::migrations::run_with_options(&pool, options).await?; println!("Migrations applied successfully."); + print_lane_sequence_leftovers(&pool).await?; } } @@ -1933,6 +1964,39 @@ async fn run() -> Result<(), Box> { }))? ); } + StorageCommands::LaneSequences { schema, repair } => { + let next_steps = if repair { + let mut conn = pool.acquire().await?; + let repairs = + awa_model::storage::repair_lane_sequences(&mut conn, &schema) + .await?; + println!( + "{}", + serde_json::to_string_pretty(&serde_json::json!({ + "schema": schema, "results": repairs, + }))? + ); + awa_model::storage::lane_sequence_repair_next_steps(&repairs) + } else { + let problems = + awa_model::storage::lane_sequence_problems(&pool, &schema).await?; + println!( + "{}", + serde_json::to_string_pretty(&serde_json::json!({ + "schema": schema, "problems": problems, + }))? + ); + awa_model::storage::lane_sequence_next_steps( + problems.iter().map(|problem| (problem, false)), + ) + }; + if !next_steps.is_empty() { + for next_step in next_steps { + eprintln!("{next_step}"); + } + std::process::exit(2); + } + } StorageCommands::Abort => { awa_model::storage::abort(&pool).await?; let report = awa_model::storage::status_report(&pool).await?; diff --git a/awa-metrics/src/lib.rs b/awa-metrics/src/lib.rs index 1a1c9097b..3f1eee208 100644 --- a/awa-metrics/src/lib.rs +++ b/awa-metrics/src/lib.rs @@ -77,6 +77,8 @@ pub mod names { pub const MAINTENANCE_PRUNE_ATTEMPTS: &str = "awa.maintenance.prune.attempts"; pub const MAINTENANCE_PRUNE_SKIPPED_ROWS: &str = "awa.maintenance.prune.skipped_rows"; pub const MAINTENANCE_PRUNE_DURATION: &str = "awa.maintenance.prune.duration"; + pub const MAINTENANCE_ANALYZE_ATTEMPTS: &str = "awa.maintenance.analyze.attempts"; + pub const MAINTENANCE_ANALYZE_DURATION: &str = "awa.maintenance.analyze.duration"; pub const STORAGE_TRANSITION_READY: &str = "awa.storage.transition_ready"; pub const STORAGE_CANONICAL_LIVE_BACKLOG: &str = "awa.storage.canonical_live_backlog"; pub const STORAGE_LIVE_RUNTIME_CAPABILITY: &str = "awa.storage.live_runtime_capability"; @@ -277,6 +279,13 @@ pub struct AwaMetrics { /// rather than one per tick. Fleets alert on this rather than /// scraping the matching `tracing::warn!` line. Issue #242. pub maintenance_branch_overrun_total: Counter, + /// Planner-statistics `ANALYZE` calls by scope (`heads`, `ring_parent`, `pass`) + /// and outcome (`ok`, `skipped_privilege`, `skipped_lock`, + /// `unavailable`, `failed`). One increment per relation per refresh pass. + pub maintenance_analyze_attempts: Counter, + /// Wall-clock duration of one planner-statistics `ANALYZE` call, by scope + /// and outcome. + pub maintenance_analyze_duration_seconds: Histogram, /// Current ring `current_slot` per ring, sampled from each rotate call. /// The slot number itself isn't meaningful but the rate of advance is — /// dashboards plot `rate(slot_changes)` to see whether rotation is @@ -620,6 +629,21 @@ impl AwaMetrics { ) .with_unit("{episode}") .build(), + maintenance_analyze_attempts: meter + .u64_counter(names::MAINTENANCE_ANALYZE_ATTEMPTS) + .with_description( + "Planner-statistics ANALYZE calls by scope (heads/ring_parent) and outcome", + ) + .with_unit("{attempt}") + .build(), + maintenance_analyze_duration_seconds: meter + .f64_histogram(names::MAINTENANCE_ANALYZE_DURATION) + .with_description( + "Duration of one planner-statistics ANALYZE call by scope and outcome", + ) + .with_unit("s") + .with_boundaries(WAIT_DURATION_BUCKETS_SECONDS.to_vec()) + .build(), ring_current_slot: meter .i64_gauge(names::RING_CURRENT_SLOT) .with_description("Current slot index per ring (queue/lease/claim)") @@ -1280,6 +1304,24 @@ impl AwaMetrics { self.maintenance_branch_overrun_total.add(1, &attrs); } + /// Record one planner-statistics `ANALYZE` call. `scope` is `"heads"`, + /// `"ring_parent"` or `"pass"` (a pass-level probe or failure); `outcome` is `"ok"`, `"skipped_privilege"`, + /// `"skipped_lock"`, `"unavailable"` or `"failed"`. + pub fn record_statistics_analyze( + &self, + scope: &'static str, + outcome: &'static str, + duration: Duration, + ) { + let attrs = [ + opentelemetry::KeyValue::new("awa.analyze.scope", scope), + opentelemetry::KeyValue::new("awa.analyze.outcome", outcome), + ]; + self.maintenance_analyze_attempts.add(1, &attrs); + self.maintenance_analyze_duration_seconds + .record(duration.as_secs_f64(), &attrs); + } + /// Record a ring prune outcome. pub fn record_prune_outcome(&self, ring: &'static str, outcome: &awa_model::PruneOutcome) { let (label, reason, count) = match outcome { @@ -1457,7 +1499,7 @@ mod tests { .build(); let metrics = AwaMetrics::new(&provider.meter("awa-metrics-test")); - let seconds: [(&str, &Histogram, &[f64]); 9] = [ + let seconds: [(&str, &Histogram, &[f64]); 10] = [ ( names::ENQUEUE_DURATION, &metrics.enqueue_duration_seconds, @@ -1478,6 +1520,11 @@ mod tests { &metrics.maintenance_branch_duration_seconds, &WAIT_DURATION_BUCKETS_SECONDS, ), + ( + names::MAINTENANCE_ANALYZE_DURATION, + &metrics.maintenance_analyze_duration_seconds, + &WAIT_DURATION_BUCKETS_SECONDS, + ), ( names::DISPATCH_WAKE_TO_CLAIM_DURATION, &metrics.dispatch_wake_to_claim_seconds, diff --git a/awa-model/migrations/v023_install_queue_storage_substrate.sql b/awa-model/migrations/v023_install_queue_storage_substrate.sql index d852b2ac5..006b3704b 100644 --- a/awa-model/migrations/v023_install_queue_storage_substrate.sql +++ b/awa-model/migrations/v023_install_queue_storage_substrate.sql @@ -42,6 +42,12 @@ LANGUAGE plpgsql SECURITY INVOKER SET search_path = pg_catalog AS $refresh$ +DECLARE + v_schema_owner OID; + v_function REGPROCEDURE; + v_function_owner OID; + v_set_mode TEXT := CASE WHEN current_setting('server_version_num')::int >= 160000 + THEN 'SET' ELSE 'MEMBER' END; BEGIN EXECUTE format( $ddl$ @@ -69,28 +75,40 @@ BEGIN ); BEGIN IF to_regclass(format('%%I.%%I', %1$L, v_enqueue_seq)) IS NULL THEN - IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + IF has_function_privilege( + to_regprocedure(format('%%I.provision_lane_sequences(text,smallint,smallint)', %1$L)), + 'EXECUTE' + ) THEN + PERFORM %1$I.provision_lane_sequences(p_queue, p_priority, p_enqueue_shard); + ELSIF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_enqueue_seq USING ERRCODE = '42501', - HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + HINT = 'Run awa migrate, or awa storage prepare-queue as the migrator, before starting producers or workers.'; + ELSE + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_enqueue_seq + ); END IF; - EXECUTE format( - 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', - %1$L, - v_enqueue_seq - ); END IF; IF to_regclass(format('%%I.%%I', %1$L, v_claim_seq)) IS NULL THEN - IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + IF has_function_privilege( + to_regprocedure(format('%%I.provision_lane_sequences(text,smallint,smallint)', %1$L)), + 'EXECUTE' + ) THEN + PERFORM %1$I.provision_lane_sequences(p_queue, p_priority, p_enqueue_shard); + ELSIF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_claim_seq USING ERRCODE = '42501', - HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + HINT = 'Run awa migrate, or awa storage prepare-queue as the migrator, before starting producers or workers.'; + ELSE + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_claim_seq + ); END IF; - EXECUTE format( - 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', - %1$L, - v_claim_seq - ); END IF; UPDATE %1$I.queue_enqueue_heads @@ -131,16 +149,22 @@ BEGIN v_start BIGINT; BEGIN IF to_regclass(format('%%I.%%I', %1$L, v_seq_name)) IS NULL THEN - IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + IF has_function_privilege( + to_regprocedure(format('%%I.provision_lane_sequences(text,smallint,smallint)', %1$L)), + 'EXECUTE' + ) THEN + PERFORM %1$I.provision_lane_sequences(NEW.queue, NEW.priority, NEW.enqueue_shard); + ELSIF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_seq_name USING ERRCODE = '42501', - HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + HINT = 'Run awa migrate, or awa storage prepare-queue as the migrator, before starting producers or workers.'; + ELSE + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_seq_name + ); END IF; - EXECUTE format( - 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', - %1$L, - v_seq_name - ); END IF; NEW.seq_name := v_seq_name; @@ -197,16 +221,22 @@ BEGIN ); BEGIN IF to_regclass(format('%%I.%%I', %1$L, v_seq_name)) IS NULL THEN - IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + IF has_function_privilege( + to_regprocedure(format('%%I.provision_lane_sequences(text,smallint,smallint)', %1$L)), + 'EXECUTE' + ) THEN + PERFORM %1$I.provision_lane_sequences(NEW.queue, NEW.priority, NEW.enqueue_shard); + ELSIF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_seq_name USING ERRCODE = '42501', - HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + HINT = 'Run awa migrate, or awa storage prepare-queue as the migrator, before starting producers or workers.'; + ELSE + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_seq_name + ); END IF; - EXECUTE format( - 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', - %1$L, - v_seq_name - ); END IF; NEW.seq_name := v_seq_name; @@ -221,11 +251,657 @@ BEGIN p_schema ); + -- Lane sequences are created by the schema owner through these two + -- SECURITY DEFINER functions, whichever role first touches the lane, so + -- every runtime role can use them. Their grants mirror the head table each + -- sequence backs (see awa.lane_sequence_problems); they never depend on + -- the creating role's default privileges. + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.provision_lane_sequences( + p_queue TEXT, + p_priority SMALLINT, + p_enqueue_shard SMALLINT + ) + RETURNS VOID + LANGUAGE plpgsql + SECURITY DEFINER + SET search_path = pg_catalog, pg_temp + AS $func$ + DECLARE + v_seq_name TEXT; + v_grant TEXT; + BEGIN + IF p_queue IS NULL + OR p_priority IS NULL + OR p_enqueue_shard IS NULL OR p_enqueue_shard NOT BETWEEN 0 AND 63 THEN + RAISE EXCEPTION 'invalid queue lane: queue %%, priority %%, enqueue shard %%', + quote_nullable(p_queue), p_priority, p_enqueue_shard + USING ERRCODE = '22023'; + END IF; + + -- The lane lock QueueStorage takes before it writes lane rows, so + -- concurrent first touches of one lane create its sequences once. + PERFORM pg_advisory_xact_lock(hashtextextended( + format('%%s:%%s:%%s:%%s', %1$L, p_queue, p_priority, p_enqueue_shard), + 0 + )); + + FOREACH v_seq_name IN ARRAY ARRAY[ + %1$I.queue_lane_sequence_name('queue_enqueue_seq', p_queue, p_priority, p_enqueue_shard), + %1$I.queue_lane_sequence_name('queue_claim_seq', p_queue, p_priority, p_enqueue_shard) + ] LOOP + IF to_regclass(format('%%I.%%I', %1$L, v_seq_name)) IS NULL THEN + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_seq_name + ); + FOR v_grant IN + SELECT unnest(problem.missing_grants) + FROM awa.lane_sequence_problems(%1$L, v_seq_name) AS problem + WHERE problem.owner = current_user + LOOP + EXECUTE v_grant; + END LOOP; + END IF; + END LOOP; + END; + $func$ + $ddl$, + p_schema + ); + + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.provision_lane( + p_queue TEXT, + p_priority SMALLINT, + p_enqueue_shard SMALLINT + ) + RETURNS VOID + LANGUAGE plpgsql + SECURITY DEFINER + SET search_path = pg_catalog, pg_temp + AS $func$ + BEGIN + -- Validates the lane and takes its lock before any row is written. + PERFORM %1$I.provision_lane_sequences(p_queue, p_priority, p_enqueue_shard); + + INSERT INTO %1$I.queue_lanes (queue, priority) + VALUES (p_queue, p_priority) + ON CONFLICT (queue, priority) DO NOTHING; + + -- The head triggers fire before ON CONFLICT and set the lane's + -- sequence, so rows that already exist are skipped outright. + INSERT INTO %1$I.queue_enqueue_heads (queue, priority, enqueue_shard) + SELECT p_queue, p_priority, p_enqueue_shard + WHERE NOT EXISTS ( + SELECT 1 FROM %1$I.queue_enqueue_heads AS heads + WHERE heads.queue = p_queue + AND heads.priority = p_priority + AND heads.enqueue_shard = p_enqueue_shard + ) + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + + INSERT INTO %1$I.queue_claim_heads (queue, priority, enqueue_shard) + SELECT p_queue, p_priority, p_enqueue_shard + WHERE NOT EXISTS ( + SELECT 1 FROM %1$I.queue_claim_heads AS heads + WHERE heads.queue = p_queue + AND heads.priority = p_priority + AND heads.enqueue_shard = p_enqueue_shard + ) + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + + PERFORM %1$I.ensure_lane_sequences(p_queue, p_priority, p_enqueue_shard); + END; + $func$ + $ddl$, + p_schema + ); + + SELECT namespace.nspowner INTO v_schema_owner + FROM pg_namespace AS namespace + WHERE namespace.nspname = p_schema; + + FOREACH v_function IN ARRAY ARRAY[ + format('%I.provision_lane_sequences(text,smallint,smallint)', p_schema)::regprocedure, + format('%I.provision_lane(text,smallint,smallint)', p_schema)::regprocedure + ] LOOP + EXECUTE format('GRANT EXECUTE ON FUNCTION %s TO PUBLIC', v_function); + + SELECT routine.proowner INTO v_function_owner + FROM pg_proc AS routine + WHERE routine.oid = v_function; + + IF v_function_owner <> v_schema_owner THEN + IF pg_has_role(current_user, v_function_owner, 'USAGE') + AND pg_has_role(current_user, v_schema_owner, v_set_mode) THEN + EXECUTE format( + 'ALTER FUNCTION %s OWNER TO %I', + v_function, + pg_get_userbyid(v_schema_owner) + ); + ELSE + RAISE WARNING '% is owned by %, not the owner % of schema %; lane sequences it creates will be owned by % too', + v_function, pg_get_userbyid(v_function_owner), pg_get_userbyid(v_schema_owner), + p_schema, pg_get_userbyid(v_function_owner) + USING HINT = format( + 'Run awa migrate as a role that can SET ROLE to %s, or as %s itself.', + quote_ident(pg_get_userbyid(v_schema_owner)), + quote_ident(pg_get_userbyid(v_schema_owner)) + ); + END IF; + END IF; + END LOOP; + END; $refresh$; REVOKE EXECUTE ON FUNCTION awa.refresh_lane_sequence_helpers(TEXT) FROM PUBLIC; +-- Enqueue-path functions. Function replacement only: no table DDL and no +-- trigger changes, so a migration can refresh a live substrate through it. +CREATE OR REPLACE FUNCTION awa.refresh_enqueue_functions(p_schema TEXT) +RETURNS VOID +LANGUAGE plpgsql +SECURITY INVOKER +SET search_path = pg_catalog +AS $refresh$ +BEGIN + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.reserve_enqueue_seq( + p_queue TEXT, + p_priority SMALLINT, + p_enqueue_shard SMALLINT, + p_count BIGINT + ) + RETURNS BIGINT + LANGUAGE plpgsql + SET search_path = pg_catalog + AS $func$ + DECLARE + v_seq_name TEXT; + v_qualified_seq TEXT; + v_start BIGINT; + BEGIN + SELECT seq_name + INTO v_seq_name + FROM %1$I.queue_enqueue_heads + WHERE queue = p_queue + AND priority = p_priority + AND enqueue_shard = p_enqueue_shard; + + IF p_count IS NULL THEN + RAISE EXCEPTION 'reserve_enqueue_seq: count must not be NULL' + USING ERRCODE = '22004'; + END IF; + + IF p_count <= 0 THEN + RETURN %1$I.sequence_next_value(v_seq_name); + END IF; + + -- The head trigger names the sequence when the lane row is + -- inserted; only a lane row written before that trigger existed + -- still needs its sequences provisioned here. + IF v_seq_name IS NULL THEN + PERFORM %1$I.ensure_lane_sequences(p_queue, p_priority, p_enqueue_shard); + + SELECT seq_name + INTO v_seq_name + FROM %1$I.queue_enqueue_heads + WHERE queue = p_queue + AND priority = p_priority + AND enqueue_shard = p_enqueue_shard; + + IF v_seq_name IS NULL THEN + RAISE EXCEPTION 'missing enqueue lane sequence for queue %%, priority %%, shard %%', + p_queue, p_priority, p_enqueue_shard + USING ERRCODE = '55000'; + END IF; + END IF; + + v_qualified_seq := format('%%I.%%I', %1$L, v_seq_name); + -- Hold the lane reservation lock until the surrounding enqueue + -- transaction ends. Sequence movement is non-transactional: if a + -- later producer could reserve and commit lane N+1 while an + -- earlier producer still holds uncommitted lane N, claim could + -- advance past N and strand the earlier job when it eventually + -- commits. The transaction-scoped lock preserves visibility order + -- within a lane while aborted reservations still become harmless + -- sequence gaps. + PERFORM pg_advisory_xact_lock(hashtextextended(v_qualified_seq, 0)); + + -- Every writer of a lane sequence holds this key (set_sequence_next + -- and the head trigger take it too), so a read followed by one + -- setval reserves the range with a single sequence write. + IF p_count = 1 THEN + v_start := nextval(v_qualified_seq::regclass); + ELSE + v_start := COALESCE( + pg_sequence_last_value(v_qualified_seq::regclass) + 1, + %1$I.sequence_next_value(v_seq_name) + ); + PERFORM setval(v_qualified_seq::regclass, v_start + p_count - 1, true); + END IF; + + RETURN v_start; + END; + $func$ + $ddl$, + p_schema + ); + + -- One-call ready enqueue. Rows arrive grouped by lane, lanes in the + -- caller's lock order, each lane's rows contiguous. Per lane: provision + -- the lane rows if the head is missing (same lane lock and inserts as the + -- Rust ensure_lane path), then reserve through reserve_enqueue_seq. Then + -- read the queue-ring cursor, append ready rows before their segments, + -- and notify. A NULL job id takes the next value of the schema's job id + -- sequence. Segment offsets are relative to the lane's reserved start. + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.enqueue_ready_rows( + p_lane_queues TEXT[], + p_lane_priorities SMALLINT[], + p_lane_shards SMALLINT[], + p_lane_counts BIGINT[], + p_job_ids BIGINT[], + p_kinds TEXT[], + p_args JSONB[], + p_attempts SMALLINT[], + p_run_leases BIGINT[], + p_max_attempts SMALLINT[], + p_run_ats TIMESTAMPTZ[], + p_attempted_ats TIMESTAMPTZ[], + p_created_ats TIMESTAMPTZ[], + p_unique_keys BYTEA[], + p_unique_states TEXT[], + p_payloads JSONB[], + p_segment_lanes INT[], + p_segment_first_offsets BIGINT[], + p_segment_next_offsets BIGINT[], + p_segment_run_ats TIMESTAMPTZ[], + p_notify_channels TEXT[] + ) + RETURNS BIGINT + LANGUAGE plpgsql + SET search_path = pg_catalog + AS $func$ + DECLARE + v_lane_total INT := COALESCE(cardinality(p_lane_queues), 0); + v_row_total INT := COALESCE(cardinality(p_kinds), 0); + v_segment_total INT := COALESCE(cardinality(p_segment_lanes), 0); + v_lane INT; + v_start BIGINT; + v_lane_starts BIGINT[] := ARRAY[]::BIGINT[]; + v_row_lanes INT[] := ARRAY[]::INT[]; + v_row_lane_seqs BIGINT[] := ARRAY[]::BIGINT[]; + v_ready_slot INT; + v_ready_generation BIGINT; + v_inserted BIGINT; + BEGIN + IF COALESCE(cardinality(p_lane_priorities), 0) <> v_lane_total + OR COALESCE(cardinality(p_lane_shards), 0) <> v_lane_total + OR COALESCE(cardinality(p_lane_counts), 0) <> v_lane_total + OR COALESCE(cardinality(p_job_ids), 0) <> v_row_total + OR COALESCE(cardinality(p_args), 0) <> v_row_total + OR COALESCE(cardinality(p_attempts), 0) <> v_row_total + OR COALESCE(cardinality(p_run_leases), 0) <> v_row_total + OR COALESCE(cardinality(p_max_attempts), 0) <> v_row_total + OR COALESCE(cardinality(p_run_ats), 0) <> v_row_total + OR COALESCE(cardinality(p_attempted_ats), 0) <> v_row_total + OR COALESCE(cardinality(p_created_ats), 0) <> v_row_total + OR COALESCE(cardinality(p_unique_keys), 0) <> v_row_total + OR COALESCE(cardinality(p_unique_states), 0) <> v_row_total + OR COALESCE(cardinality(p_payloads), 0) <> v_row_total + OR COALESCE(cardinality(p_segment_first_offsets), 0) <> v_segment_total + OR COALESCE(cardinality(p_segment_next_offsets), 0) <> v_segment_total + OR COALESCE(cardinality(p_segment_run_ats), 0) <> v_segment_total + THEN + RAISE EXCEPTION 'enqueue_ready_rows: array lengths disagree' + USING ERRCODE = '22023'; + END IF; + + FOR v_lane IN 1..v_lane_total LOOP + IF p_lane_counts[v_lane] < 1 THEN + RAISE EXCEPTION 'enqueue_ready_rows: lane %% has no rows', v_lane + USING ERRCODE = '22023'; + END IF; + + PERFORM 1 + FROM %1$I.queue_enqueue_heads + WHERE queue = p_lane_queues[v_lane] + AND priority = p_lane_priorities[v_lane] + AND enqueue_shard = p_lane_shards[v_lane]; + + IF NOT FOUND THEN + PERFORM pg_advisory_xact_lock(hashtextextended( + format('%%s:%%s:%%s:%%s', %1$L, p_lane_queues[v_lane], + p_lane_priorities[v_lane], p_lane_shards[v_lane]), + 0 + )); + INSERT INTO %1$I.queue_lanes (queue, priority) + VALUES (p_lane_queues[v_lane], p_lane_priorities[v_lane]) + ON CONFLICT (queue, priority) DO NOTHING; + INSERT INTO %1$I.queue_enqueue_heads (queue, priority, enqueue_shard) + VALUES (p_lane_queues[v_lane], p_lane_priorities[v_lane], p_lane_shards[v_lane]) + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + INSERT INTO %1$I.queue_claim_heads (queue, priority, enqueue_shard) + VALUES (p_lane_queues[v_lane], p_lane_priorities[v_lane], p_lane_shards[v_lane]) + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + PERFORM %1$I.ensure_lane_sequences( + p_lane_queues[v_lane], p_lane_priorities[v_lane], p_lane_shards[v_lane] + ); + END IF; + + v_start := %1$I.reserve_enqueue_seq( + p_lane_queues[v_lane], + p_lane_priorities[v_lane], + p_lane_shards[v_lane], + p_lane_counts[v_lane] + ); + v_lane_starts := v_lane_starts || v_start; + v_row_lanes := v_row_lanes || array_fill(v_lane, ARRAY[p_lane_counts[v_lane]::INT]); + v_row_lane_seqs := v_row_lane_seqs + || ARRAY(SELECT generate_series(v_start, v_start + p_lane_counts[v_lane] - 1)); + END LOOP; + + IF cardinality(v_row_lane_seqs) <> v_row_total THEN + RAISE EXCEPTION 'enqueue_ready_rows: lane counts cover %% rows, got %%', + cardinality(v_row_lane_seqs), v_row_total + USING ERRCODE = '22023'; + END IF; + + IF v_row_total = 0 THEN + RETURN 0; + END IF; + + SELECT cursor.slot, cursor.generation + INTO v_ready_slot, v_ready_generation + FROM %1$I.ring_cursor('queue') AS cursor; + + INSERT INTO %1$I.ready_entries ( + ready_slot, ready_generation, job_id, kind, queue, args, priority, + attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, + attempted_at, created_at, unique_key, unique_states, payload + ) + SELECT + v_ready_slot, + v_ready_generation, + COALESCE(input.job_id, nextval(%2$L::regclass)), + input.kind, + p_lane_queues[input.lane], + input.args, + p_lane_priorities[input.lane], + input.attempt, + input.run_lease, + input.max_attempts, + input.lane_seq, + p_lane_shards[input.lane], + input.run_at, + input.attempted_at, + input.created_at, + input.unique_key, + input.unique_states, + input.payload + FROM unnest( + v_row_lanes, v_row_lane_seqs, p_job_ids, p_kinds, p_args, p_attempts, + p_run_leases, p_max_attempts, p_run_ats, p_attempted_ats, p_created_ats, + p_unique_keys, p_unique_states, p_payloads + ) AS input( + lane, lane_seq, job_id, kind, args, attempt, run_lease, max_attempts, + run_at, attempted_at, created_at, unique_key, unique_states, payload + ); + GET DIAGNOSTICS v_inserted = ROW_COUNT; + + INSERT INTO %1$I.ready_segments ( + ready_slot, ready_generation, queue, priority, enqueue_shard, + first_lane_seq, next_lane_seq, first_run_at + ) + SELECT + v_ready_slot, + v_ready_generation, + p_lane_queues[segment.lane], + p_lane_priorities[segment.lane], + p_lane_shards[segment.lane], + v_lane_starts[segment.lane] + segment.first_offset, + v_lane_starts[segment.lane] + segment.next_offset, + segment.run_at + FROM unnest( + p_segment_lanes, p_segment_first_offsets, p_segment_next_offsets, p_segment_run_ats + ) AS segment(lane, first_offset, next_offset, run_at) + ON CONFLICT (ready_slot, ready_generation, queue, priority, enqueue_shard, first_lane_seq) + DO NOTHING; + + PERFORM pg_notify(channel, '') + FROM unnest(p_notify_channels) AS channel; + + RETURN v_inserted; + END; + $func$ + $ddl$, + p_schema, + format('%I.job_id_seq', p_schema) + ); +END; +$refresh$; + +REVOKE EXECUTE ON FUNCTION awa.refresh_enqueue_functions(TEXT) FROM PUBLIC; + +-- Lane sequences in p_schema (or the one named p_sequence) that are not owned +-- by the schema owner or lack a runtime grant. A lane sequence carries a head +-- table's cursor, so whoever may write that head table (INSERT or UPDATE) gets +-- USAGE, SELECT and UPDATE on it, and whoever may read it gets SELECT: +-- queue_enqueue_seq_* mirror queue_enqueue_heads, queue_claim_seq_* mirror +-- queue_claim_heads. missing_grants holds the statements to run as the +-- sequence owner; reassign_sql, when the owner is wrong, needs a role that +-- holds the current owner's privileges and can SET ROLE to the schema owner. +-- repair_roles lists the login roles, superusers aside, that can run both. +-- Read-only. +CREATE OR REPLACE FUNCTION awa.lane_sequence_problems(p_schema TEXT, p_sequence TEXT DEFAULT NULL) +RETURNS TABLE ( + sequence_name TEXT, + owner TEXT, + expected_owner TEXT, + missing_grants TEXT[], + reassign_sql TEXT, + repair_roles TEXT[] +) +LANGUAGE sql +STABLE +SECURITY INVOKER +SET search_path = pg_catalog, pg_temp +AS $problems$ + WITH target AS ( + SELECT namespace.oid, namespace.nspowner + FROM pg_namespace AS namespace + WHERE namespace.nspname = p_schema + ), + lane_sequence AS ( + SELECT sequence.oid, + sequence.relname::text AS relname, + sequence.relowner, + sequence.relacl, + target.nspowner, + CASE WHEN starts_with(sequence.relname::text, 'queue_enqueue_seq_') + THEN 'queue_enqueue_heads' ELSE 'queue_claim_heads' END AS head_table + FROM target + JOIN pg_class AS sequence ON sequence.relnamespace = target.oid + WHERE sequence.relkind = 'S' + AND sequence.relname ~ '^queue_(enqueue|claim)_seq_[0-9a-f]{32}$' + AND (p_sequence IS NULL OR sequence.relname = p_sequence) + ), + wanted AS ( + SELECT DISTINCT lane_sequence.oid, head_privilege.grantee, sequence_privilege.privilege_type + FROM lane_sequence + JOIN pg_class AS head + ON head.relnamespace = (SELECT target.oid FROM target) + AND head.relname = lane_sequence.head_table + CROSS JOIN LATERAL aclexplode(COALESCE(head.relacl, acldefault('r', head.relowner))) AS head_privilege + CROSS JOIN LATERAL unnest( + CASE head_privilege.privilege_type + WHEN 'SELECT' THEN ARRAY['SELECT'] + WHEN 'INSERT' THEN ARRAY['USAGE', 'SELECT', 'UPDATE'] + WHEN 'UPDATE' THEN ARRAY['USAGE', 'SELECT', 'UPDATE'] + ELSE ARRAY[]::TEXT[] + END + ) AS sequence_privilege(privilege_type) + WHERE head_privilege.grantee <> lane_sequence.nspowner + ), + missing AS ( + SELECT wanted.oid, wanted.grantee, wanted.privilege_type + FROM wanted + JOIN lane_sequence ON lane_sequence.oid = wanted.oid + WHERE NOT EXISTS ( + SELECT 1 + FROM aclexplode(COALESCE(lane_sequence.relacl, acldefault('S', lane_sequence.relowner))) AS held + WHERE held.grantee = wanted.grantee + AND held.privilege_type = wanted.privilege_type + ) + ) + SELECT lane_sequence.relname, + pg_get_userbyid(lane_sequence.relowner)::text, + pg_get_userbyid(lane_sequence.nspowner)::text, + ARRAY( + SELECT format( + 'GRANT %s ON SEQUENCE %I.%I TO %s', + string_agg(missing.privilege_type, ', ' ORDER BY missing.privilege_type DESC), + p_schema, + lane_sequence.relname, + CASE missing.grantee WHEN 0 THEN 'PUBLIC' + ELSE quote_ident(pg_get_userbyid(missing.grantee)) END + ) + FROM missing + WHERE missing.oid = lane_sequence.oid + GROUP BY missing.grantee + ORDER BY 1 + ), + CASE WHEN lane_sequence.relowner <> lane_sequence.nspowner THEN format( + 'ALTER SEQUENCE %I.%I OWNER TO %I', + p_schema, lane_sequence.relname, pg_get_userbyid(lane_sequence.nspowner) + ) END, + ARRAY( + SELECT candidate.rolname::text + FROM pg_roles AS candidate + WHERE candidate.rolcanlogin + AND NOT candidate.rolsuper + AND pg_has_role(candidate.oid, lane_sequence.relowner, 'USAGE') + AND (lane_sequence.relowner = lane_sequence.nspowner + OR (pg_has_role(candidate.oid, lane_sequence.nspowner, 'USAGE') + AND pg_has_role( + candidate.oid, + lane_sequence.nspowner, + CASE WHEN current_setting('server_version_num')::int >= 160000 + THEN 'SET' ELSE 'MEMBER' END + ))) + ORDER BY candidate.rolname + ) + FROM lane_sequence + WHERE lane_sequence.relowner <> lane_sequence.nspowner + OR EXISTS (SELECT 1 FROM missing WHERE missing.oid = lane_sequence.oid) + ORDER BY lane_sequence.relname +$problems$; + +-- Give one lane sequence to the schema owner and add its missing runtime +-- grants, doing only what the calling role is allowed to: a sequence it +-- cannot reassign keeps its owner and still gets the grants when the caller +-- holds that owner's privileges. Never raises for a sequence it cannot +-- repair. Reassigning takes the sequence's ACCESS EXCLUSIVE lock, which an +-- open enqueue or claim on the lane conflicts with, so the lock wait is +-- bounded: on timeout nothing changes and the outcome is 'busy'. The lock is +-- held until the caller's transaction ends, so callers repair one sequence +-- per transaction. Returns no row when the sequence needs nothing, otherwise +-- one row: outcome is 'repaired', 'owner_unchanged' (grants complete, owner +-- still wrong), 'unrepaired' or 'busy', with the statements still to run. +CREATE OR REPLACE FUNCTION awa.repair_lane_sequence(p_schema TEXT, p_sequence TEXT) +RETURNS TABLE ( + sequence_name TEXT, + owner TEXT, + expected_owner TEXT, + outcome TEXT, + missing_grants TEXT[], + reassign_sql TEXT, + repair_roles TEXT[] +) +LANGUAGE plpgsql +SECURITY INVOKER +SET search_path = pg_catalog, pg_temp +SET lock_timeout = '300ms' +AS $repair$ +DECLARE + v_problem RECORD; + v_remaining RECORD; + v_busy BOOLEAN := false; + v_set_mode TEXT := CASE WHEN current_setting('server_version_num')::int >= 160000 + THEN 'SET' ELSE 'MEMBER' END; + v_statement TEXT; +BEGIN + SELECT * INTO v_problem + FROM awa.lane_sequence_problems(p_schema, p_sequence) AS problem; + IF NOT FOUND THEN + RETURN; + END IF; + + BEGIN + -- ALTER ... OWNER TO needs the current owner's privileges, SET ROLE to + -- the new owner, and the new owner's CREATE on the schema (which a + -- schema owner has). Check first so an unrelated owner is skipped + -- without waiting for the lock. + IF v_problem.reassign_sql IS NOT NULL + AND pg_has_role(current_user, to_regrole(quote_ident(v_problem.owner)), 'USAGE') + AND pg_has_role(current_user, to_regrole(quote_ident(v_problem.expected_owner)), v_set_mode) THEN + BEGIN + EXECUTE v_problem.reassign_sql; + EXCEPTION WHEN insufficient_privilege THEN + NULL; + END; + END IF; + + SELECT * INTO v_remaining + FROM awa.lane_sequence_problems(p_schema, p_sequence) AS problem; + IF FOUND + AND cardinality(v_remaining.missing_grants) > 0 + AND pg_has_role(current_user, to_regrole(quote_ident(v_remaining.owner)), 'USAGE') THEN + FOREACH v_statement IN ARRAY v_remaining.missing_grants LOOP + BEGIN + EXECUTE v_statement; + EXCEPTION WHEN insufficient_privilege THEN + NULL; + END; + END LOOP; + END IF; + EXCEPTION WHEN lock_not_available THEN + v_busy := true; + END; + + SELECT * INTO v_remaining + FROM awa.lane_sequence_problems(p_schema, p_sequence) AS problem; + sequence_name := v_problem.sequence_name; + expected_owner := v_problem.expected_owner; + IF FOUND THEN + owner := v_remaining.owner; + missing_grants := v_remaining.missing_grants; + reassign_sql := v_remaining.reassign_sql; + repair_roles := v_remaining.repair_roles; + outcome := CASE WHEN v_busy THEN 'busy' + WHEN cardinality(v_remaining.missing_grants) = 0 THEN 'owner_unchanged' + ELSE 'unrepaired' END; + ELSE + owner := v_problem.expected_owner; + missing_grants := ARRAY[]::TEXT[]; + reassign_sql := NULL; + repair_roles := ARRAY[]::TEXT[]; + outcome := 'repaired'; + END IF; + RETURN NEXT; +END +$repair$; + +REVOKE EXECUTE ON FUNCTION awa.repair_lane_sequence(TEXT, TEXT) FROM PUBLIC; + -- Retained job output (ADR-048): add the output columns to an existing -- substrate and expose them through terminal_jobs. Called by the installer -- and by migration v052, which does not replay the installer. Each column is @@ -1414,78 +2090,7 @@ BEGIN PERFORM awa.refresh_lane_sequence_helpers(p_schema); - EXECUTE format( - $ddl$ - CREATE OR REPLACE FUNCTION %1$I.reserve_enqueue_seq( - p_queue TEXT, - p_priority SMALLINT, - p_enqueue_shard SMALLINT, - p_count BIGINT - ) - RETURNS BIGINT - LANGUAGE plpgsql - SET search_path = pg_catalog - AS $func$ - DECLARE - v_seq_name TEXT; - v_lock_key TEXT; - v_start BIGINT; - BEGIN - IF p_count <= 0 THEN - RETURN %1$I.sequence_next_value(( - SELECT seq_name - FROM %1$I.queue_enqueue_heads - WHERE queue = p_queue - AND priority = p_priority - AND enqueue_shard = p_enqueue_shard - )); - END IF; - - PERFORM %1$I.ensure_lane_sequences(p_queue, p_priority, p_enqueue_shard); - - SELECT seq_name - INTO v_seq_name - FROM %1$I.queue_enqueue_heads - WHERE queue = p_queue - AND priority = p_priority - AND enqueue_shard = p_enqueue_shard; - - IF v_seq_name IS NULL THEN - RAISE EXCEPTION 'missing enqueue lane sequence for queue %%, priority %%, shard %%', - p_queue, p_priority, p_enqueue_shard - USING ERRCODE = '55000'; - END IF; - - v_lock_key := format('%%I.%%I', %1$L, v_seq_name); - -- Hold the lane reservation lock until the surrounding enqueue - -- transaction ends. Sequence movement is non-transactional: if a - -- later producer could reserve and commit lane N+1 while an - -- earlier producer still holds uncommitted lane N, claim could - -- advance past N and strand the earlier job when it eventually - -- commits. The transaction-scoped lock preserves visibility order - -- within a lane while aborted reservations still become harmless - -- sequence gaps. - PERFORM pg_advisory_xact_lock(hashtextextended(v_lock_key, 0)); - EXECUTE format( - 'SELECT nextval(%%L::regclass)::bigint', - format('%%I.%%I', %1$L, v_seq_name) - ) - INTO v_start; - - IF p_count > 1 THEN - EXECUTE format( - 'SELECT setval(%%L::regclass, %%s, true)', - format('%%I.%%I', %1$L, v_seq_name), - v_start + p_count - 1 - ); - END IF; - - RETURN v_start; - END; - $func$ - $ddl$, - p_schema - ); + PERFORM awa.refresh_enqueue_functions(p_schema); EXECUTE format( 'DROP TRIGGER IF EXISTS queue_enqueue_heads_sequence_sync ON %I.queue_enqueue_heads', diff --git a/awa-model/migrations/v053_definer_lane_provisioning.sql b/awa-model/migrations/v053_definer_lane_provisioning.sql new file mode 100644 index 000000000..ced1a7642 --- /dev/null +++ b/awa-model/migrations/v053_definer_lane_provisioning.sql @@ -0,0 +1,58 @@ +-- v053: Create queue lanes as the schema owner, whichever role touches them first. +-- Tracking: https://github.com/hardbyte/awa/pull/637 +-- Lane sequences used to belong to the role that first enqueued into, claimed +-- from or promoted into a lane, and only that role (plus whoever its default +-- privileges named) could use them. refresh_lane_sequence_helpers now installs +-- two SECURITY DEFINER functions, provision_lane and provision_lane_sequences, +-- owned by the schema owner, and the lane helpers create a missing sequence +-- through them. Each new sequence gets the grants of the head table it backs. +-- This migration does not reassign existing lane sequences: ALTER SEQUENCE +-- OWNER takes the sequence's ACCESS EXCLUSIVE lock until commit, and the +-- migration runs in one transaction. After it commits, `awa migrate` repairs +-- each lane sequence in its own short transaction with a bounded lock wait +-- (awa.repair_lane_sequence). Another runner can do the same with +-- `awa storage lane-sequences --repair`. +-- +-- N-1 binaries call the same helpers with the same arguments and cursor +-- semantics, so lanes they create are owned by the schema owner too. A binary +-- that predates provision_lane keeps using the helpers; this binary probes for +-- provision_lane and uses the helpers when it is missing. +-- +-- Live load: function replacement only. Reads the catalog; no table DDL, row +-- scans or sequence locks. + +DO $$ +DECLARE + v_schema TEXT; + v_count BIGINT; +BEGIN + FOR v_schema IN + SELECT discovered.schema_name FROM awa.queue_storage_schemas() AS discovered + LOOP + IF to_regclass(format('%I.queue_enqueue_heads', v_schema)) IS NULL + OR to_regclass(format('%I.queue_claim_heads', v_schema)) IS NULL + OR to_regclass(format('%I.queue_ring_state', v_schema)) IS NULL + OR to_regclass(format('%I.lease_ring_state', v_schema)) IS NULL + OR to_regclass(format('%I.claim_ring_state', v_schema)) IS NULL THEN + CONTINUE; + END IF; + + PERFORM awa.refresh_lane_sequence_helpers(v_schema); + + SELECT count(*) INTO v_count + FROM awa.lane_sequence_problems(v_schema) AS problem; + + IF v_count > 0 THEN + RAISE NOTICE '% lane sequence(s) in schema % are not owned by the schema owner or lack runtime grants', v_count, v_schema + USING HINT = format( + 'awa migrate attempts repair after this commits; review its warnings for any sequences left unrepaired. To retry, run awa storage lane-sequences --schema %s --repair as the migration role itself, without SET ROLE.', + v_schema + ); + END IF; + END LOOP; +END; +$$; + +INSERT INTO awa.schema_version (version, description) +VALUES (53, 'Create queue lanes as the schema owner') +ON CONFLICT (version) DO NOTHING; diff --git a/awa-model/migrations/v054_analyze_queue_storage_statistics.sql b/awa-model/migrations/v054_analyze_queue_storage_statistics.sql new file mode 100644 index 000000000..ef2ecfc1e --- /dev/null +++ b/awa-model/migrations/v054_analyze_queue_storage_statistics.sql @@ -0,0 +1,180 @@ +-- v054: Planner statistics for the queue-storage tables autovacuum never analyses (#631) +-- +-- Autovacuum does not analyse partitioned tables, so the ring parents +-- (ready_entries, done_entries, leases, ...) have no inherited statistics. It +-- also never reaches the autoanalyze threshold on the two small head tables +-- (queue_claim_heads, queue_enqueue_heads), so the planner prices lane joins +-- against them at its defaults (320 rows, n_distinct 200) while they drive +-- joins over hundreds of thousands of retained rows. +-- +-- The maintenance leader refreshes both on an interval. Runtimes usually +-- connect as a role that does not own the tables, and ANALYZE skips with a +-- warning a table its caller may not analyse, so the refresh runs through +-- this SECURITY DEFINER function owned by the migration role, which owns the +-- tables. +-- +-- The function analyses one relation per call and the caller runs each call as +-- its own transaction, so the locks ANALYZE takes on a partitioned table's +-- children are held for that table rather than for the whole refresh. +-- +-- Locking on a partitioned table, per the ANALYZE documentation: SKIP_LOCKED +-- applies when ANALYZE starts work on a relation, so a child that ring prune +-- holds ACCESS EXCLUSIVE is skipped during the per-child analysis, but +-- sampling the children for the parent's statistics can still wait for it. +-- That wait is bounded by the function's lock_timeout and fails the call with +-- SQLSTATE 55P03, which the caller records as skipped_lock. SKIP_LOCKED also +-- skips every partition silently when the parent itself is locked, so the +-- function takes the parent's SHARE UPDATE EXCLUSIVE lock first and lets that +-- wait fail instead. +-- +-- Definer boundary: only a head table or ring parent of a schema that +-- awa.queue_storage_schemas() reports is analysed, which requires the schema's +-- ownership, grants and triggers to be trusted, and SQL is built from nothing +-- but quoted identifiers. ANALYZE writes statistics and nothing else. +-- Only the fixed list of ring parents and the two head tables is accepted; +-- a migration that adds a partitioned table to the substrate must extend it. +-- +-- Who may call it: the lock above lasts until the caller's transaction ends +-- and blocks ring prune of that table, so EXECUTE is not granted to PUBLIC. It +-- goes to the roles that already hold INSERT, UPDATE and DELETE on the head +-- tables of the awa schema (the runtime grantees), and the function itself +-- requires the calling session user to hold those privileges on the target +-- schema's queue_claim_heads. A role created later needs the documented +-- `GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa` (or default privileges). +-- +-- Returns 'analyzed'; 'skipped_privilege' when the function owner may not +-- analyse the relation; or 'missing' when the relation does not exist. +-- +-- N-1 binaries never call it. A binary newer than this schema probes for the +-- function, logs once that it is missing and skips the refresh. Creating or +-- replacing it takes no table lock and completes in milliseconds; calling it +-- takes SHARE UPDATE EXCLUSIVE on the relation until the caller's transaction +-- ends. + +CREATE OR REPLACE FUNCTION awa.analyze_queue_storage_relation( + p_schema TEXT, + p_relation TEXT +) +RETURNS TEXT +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +SET lock_timeout = '2s' +AS $$ +DECLARE + v_heads REGCLASS; + v_relid OID; + v_relkind "char"; + v_relowner OID; + v_may_analyze BOOLEAN; +BEGIN + -- awa_statistics_helper_v1: bump when the definition or grants change. + IF NOT EXISTS ( + SELECT 1 FROM awa.queue_storage_schemas() AS substrate + WHERE substrate.schema_name = p_schema + ) THEN + RAISE EXCEPTION '% is not a queue-storage schema', p_schema + USING ERRCODE = '22023'; + END IF; + + IF p_relation <> ALL (ARRAY[ + 'queue_claim_heads', + 'queue_enqueue_heads', + 'done_entries', + 'lease_claim_batches', + 'lease_claim_closure_batches', + 'lease_claim_closures', + 'lease_claims', + 'leases', + 'queue_terminal_count_deltas', + 'ready_claim_attempt_batches', + 'ready_entries', + 'ready_segments', + 'ready_tombstones', + 'receipt_completion_batches', + 'receipt_completion_tombstones' + ]) THEN + RAISE EXCEPTION '%.% is not a ring parent or head table of a queue-storage schema', + p_schema, p_relation + USING ERRCODE = '22023'; + END IF; + + v_heads := to_regclass(format('%I.queue_claim_heads', p_schema)); + IF v_heads IS NULL + OR NOT has_table_privilege(session_user, v_heads, 'INSERT') + OR NOT has_table_privilege(session_user, v_heads, 'UPDATE') + OR NOT has_table_privilege(session_user, v_heads, 'DELETE') THEN + RAISE EXCEPTION 'permission denied: % may not maintain queue storage in schema %', + session_user, p_schema + USING ERRCODE = '42501'; + END IF; + + v_relid := to_regclass(format('%I.%I', p_schema, p_relation)); + IF v_relid IS NULL THEN + RETURN 'missing'; + END IF; + + SELECT class.relkind, class.relowner INTO v_relkind, v_relowner + FROM pg_class AS class WHERE class.oid = v_relid; + + -- ANALYZE skips, with a warning rather than an error, a table the caller + -- may not analyse: the owner, a MAINTAIN holder (PostgreSQL 17+), a + -- member of pg_maintain, or the database owner may analyse it. + SELECT + pg_has_role(current_user, v_relowner, 'USAGE') + OR pg_has_role( + current_user, + (SELECT database.datdba FROM pg_database AS database + WHERE database.datname = current_database()), + 'USAGE' + ) + OR CASE + WHEN current_setting('server_version_num')::int >= 170000 + THEN has_table_privilege(v_relid, 'MAINTAIN') + ELSE FALSE + END + INTO v_may_analyze; + + IF NOT v_may_analyze THEN + RETURN 'skipped_privilege'; + END IF; + + IF v_relkind = 'p' THEN + EXECUTE format( + 'LOCK TABLE ONLY %I.%I IN SHARE UPDATE EXCLUSIVE MODE', p_schema, p_relation + ); + EXECUTE format('ANALYZE (SKIP_LOCKED) %I.%I', p_schema, p_relation); + ELSE + EXECUTE format('ANALYZE %I.%I', p_schema, p_relation); + END IF; + RETURN 'analyzed'; +END +$$; + +REVOKE EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(TEXT, TEXT) FROM PUBLIC; + +DO $grant$ +DECLARE + v_grantee TEXT; +BEGIN + FOR v_grantee IN + SELECT pg_get_userbyid(privilege.grantee) + FROM pg_class AS heads + CROSS JOIN LATERAL aclexplode(heads.relacl) AS privilege + WHERE heads.oid = to_regclass('awa.queue_claim_heads') + AND privilege.grantee <> 0 + AND privilege.privilege_type IN ('INSERT', 'UPDATE', 'DELETE') + GROUP BY privilege.grantee + HAVING count(DISTINCT privilege.privilege_type) = 3 + LOOP + EXECUTE format( + 'GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) TO %I', + v_grantee + ); + END LOOP; +END +$grant$; + +INSERT INTO awa.schema_version (version, description) +VALUES (54, 'Planner statistics refresh for queue-storage tables autovacuum never analyses') +ON CONFLICT (version) DO NOTHING; diff --git a/awa-model/migrations/v055_insert_jobs_batch_compat.sql b/awa-model/migrations/v055_insert_jobs_batch_compat.sql new file mode 100644 index 000000000..9fcca8610 --- /dev/null +++ b/awa-model/migrations/v055_insert_jobs_batch_compat.sql @@ -0,0 +1,580 @@ +-- v055: Batch producer insert that holds each enqueue lane lock only for the +-- lane write. +-- Tracking: https://github.com/hardbyte/awa/pull/635 +-- +-- awa.insert_job_compat takes the (queue, priority, shard) lane lock in +-- reserve_enqueue_seq on a batch's first row and, because the lock is +-- transaction-scoped, holds it while every remaining row is processed. Many +-- concurrent producers on one lane therefore queue behind whole batches. This +-- migration adds awa.insert_jobs_compat_batch, which does the lane-independent +-- work first (job ids, unique claims in key order, lane rows, payloads), then +-- reserves each lane's sequence range once, in (queue, priority, shard) order, +-- and writes the ready rows and segments in two statements. The lock is held +-- across the reservation, those two statements and the commit. +-- +-- The column arrays are parallel and must have equal length. Every element is +-- required except run_at, unique_key, unique_states and ordering_hash; +-- insert_job_compat's per-row defaults are the caller's to apply. +-- +-- p_skip_unique_conflicts inserts a batch while skipping rows whose unique key +-- is already claimed (the rows are absent from the result) instead of failing +-- the whole batch with a unique violation. +-- +-- N-1 binaries: unchanged. They keep calling awa.insert_job_compat, whose +-- definition this migration does not touch, and the two paths share the lane +-- lock, so producers on either version serialise exactly as before. +-- Binaries that know the function: insert_many_copy probes its own connection +-- and uses the per-row function until it exists, picking up the migration +-- without a restart. Generic insert_many always uses the per-row function; +-- insert_many_skip_unique_conflicts requires the batch function. +-- External runners: apply this file; nothing else reads it. +-- Live load: one CREATE OR REPLACE FUNCTION, no table DDL. + +CREATE OR REPLACE FUNCTION awa.insert_jobs_compat_batch( + p_kind TEXT[], + p_queue TEXT[], + p_args JSONB[], + p_state TEXT[], + p_priority SMALLINT[], + p_max_attempts SMALLINT[], + p_run_at TIMESTAMPTZ[], + p_metadata JSONB[], + p_tags JSONB[], + p_unique_key BYTEA[], + p_unique_states TEXT[], + p_ordering_hash TEXT[], + p_skip_unique_conflicts BOOLEAN DEFAULT FALSE +) +RETURNS SETOF awa.jobs +AS $$ +DECLARE + v_count INT := COALESCE(cardinality(p_kind), 0); + v_schema TEXT; + v_old_search_path TEXT; + v_now TIMESTAMPTZ := clock_timestamp(); + v_job_ids BIGINT[]; + v_skipped BIGINT[] := ARRAY[]::BIGINT[]; + v_lane RECORD; + v_lane_queues TEXT[] := ARRAY[]::TEXT[]; + v_lane_priorities SMALLINT[] := ARRAY[]::SMALLINT[]; + v_lane_shards SMALLINT[] := ARRAY[]::SMALLINT[]; + v_lane_starts BIGINT[] := ARRAY[]::BIGINT[]; + v_start BIGINT; + v_ready_slot INT; + v_ready_generation BIGINT; + v_row RECORD; + v_shards SMALLINT[]; +BEGIN + IF v_count = 0 THEN + RETURN; + END IF; + + IF cardinality(p_queue) IS DISTINCT FROM v_count + OR cardinality(p_args) IS DISTINCT FROM v_count + OR cardinality(p_state) IS DISTINCT FROM v_count + OR cardinality(p_priority) IS DISTINCT FROM v_count + OR cardinality(p_max_attempts) IS DISTINCT FROM v_count + OR cardinality(p_run_at) IS DISTINCT FROM v_count + OR cardinality(p_metadata) IS DISTINCT FROM v_count + OR cardinality(p_tags) IS DISTINCT FROM v_count + OR cardinality(p_unique_key) IS DISTINCT FROM v_count + OR cardinality(p_unique_states) IS DISTINCT FROM v_count + OR cardinality(p_ordering_hash) IS DISTINCT FROM v_count + THEN + RAISE EXCEPTION 'insert_jobs_compat_batch: column arrays must have equal length' + USING ERRCODE = '22023'; + END IF; + + IF EXISTS ( + SELECT 1 + FROM unnest(p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_metadata, p_tags) + AS input(kind, queue, args, state, priority, max_attempts, metadata, tags) + WHERE input.kind IS NULL OR input.queue IS NULL OR input.args IS NULL + OR input.state IS NULL OR input.priority IS NULL OR input.max_attempts IS NULL + OR input.metadata IS NULL OR input.tags IS NULL + ) THEN + RAISE EXCEPTION 'insert_jobs_compat_batch: kind, queue, args, state, priority, max_attempts, metadata and tags must not be NULL' + USING ERRCODE = '22004'; + END IF; + + IF EXISTS (SELECT 1 FROM unnest(p_kind) AS kind WHERE length(kind) > 200) THEN + RAISE EXCEPTION 'job kind length must be <= 200 characters' + USING ERRCODE = '23514'; + END IF; + + IF EXISTS (SELECT 1 FROM unnest(p_queue) AS queue WHERE length(queue) > 200) THEN + RAISE EXCEPTION 'queue name length must be <= 200 characters' + USING ERRCODE = '23514'; + END IF; + + IF EXISTS ( + SELECT 1 FROM unnest(p_priority) AS priority + WHERE priority < 1 OR priority > 4 + ) THEN + RAISE EXCEPTION 'priority must be between 1 and 4' + USING ERRCODE = '23514'; + END IF; + + IF EXISTS ( + SELECT 1 FROM unnest(p_max_attempts) AS max_attempts + WHERE max_attempts < 1 OR max_attempts > 1000 + ) THEN + RAISE EXCEPTION 'max_attempts must be between 1 and 1000' + USING ERRCODE = '23514'; + END IF; + + IF EXISTS ( + SELECT 1 FROM unnest(p_tags) AS tags + WHERE jsonb_array_length(tags) > 20 + ) THEN + RAISE EXCEPTION 'job tags must contain at most 20 values' + USING ERRCODE = '23514'; + END IF; + + v_schema := awa.active_queue_storage_schema(); + + -- Canonical storage has no lane lock to shorten: insert row by row. + IF v_schema IS NULL THEN + IF NOT p_skip_unique_conflicts THEN + RETURN QUERY + SELECT inserted.* + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, + p_run_at, p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + CROSS JOIN LATERAL awa.insert_job_compat( + input.kind, + input.queue, + input.args, + input.state::awa.job_state, + input.priority, + input.max_attempts, + input.run_at, + input.metadata, + ARRAY( + SELECT tag.value + FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, tag_ord) + ORDER BY tag.tag_ord + ), + input.unique_key, + input.unique_states::bit(8), + NULL + ) AS inserted + ORDER BY input.ord; + RETURN; + END IF; + + -- Skipping needs a subtransaction per row. + FOR v_row IN + SELECT + input.ord, + input.kind, + input.queue, + input.args, + input.state, + input.priority, + input.max_attempts, + input.run_at, + input.metadata, + ARRAY( + SELECT tag.value + FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, tag_ord) + ORDER BY tag.tag_ord + ) AS tags, + input.unique_key, + input.unique_states + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, + p_run_at, p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + ORDER BY input.ord + LOOP + BEGIN + RETURN NEXT awa.insert_job_compat( + v_row.kind, + v_row.queue, + v_row.args, + v_row.state::awa.job_state, + v_row.priority, + v_row.max_attempts, + v_row.run_at, + v_row.metadata, + v_row.tags, + v_row.unique_key, + v_row.unique_states::bit(8), + NULL + ); + EXCEPTION WHEN unique_violation THEN + NULL; + END; + END LOOP; + RETURN; + END IF; + + IF EXISTS ( + SELECT 1 FROM unnest(p_state) AS state + WHERE state NOT IN ('available', 'scheduled', 'retryable') + ) THEN + RAISE EXCEPTION 'queue storage does not support initial state %', + (SELECT state FROM unnest(p_state) AS state + WHERE state NOT IN ('available', 'scheduled', 'retryable') LIMIT 1) + USING ERRCODE = '22023'; + END IF; + + v_old_search_path := current_setting('search_path'); + PERFORM set_config('search_path', format('%I,awa,public', v_schema), true); + + -- Everything up to the lane reservation is lane-independent and runs + -- without a lane lock. + EXECUTE format( + 'SELECT array_agg(nextval(%L::regclass)::bigint) FROM generate_series(1, $1)', + format('%I.job_id_seq', v_schema) + ) + INTO v_job_ids + USING v_count; + + -- Rows that enqueue as `available` go to a lane. Their shard is derived as + -- insert_job_compat derives it: ordering-key rows from the key hash, the + -- rest from the backend pid. + SELECT array_agg( + CASE WHEN input.state = 'available' THEN + MOD( + CASE WHEN input.ordering_hash IS NULL + THEN pg_backend_pid()::numeric + ELSE input.ordering_hash::numeric + END, + COALESCE(meta.enqueue_shards, 1)::numeric + )::smallint + END + ORDER BY input.ord + ) + INTO v_shards + FROM unnest(p_queue, p_state, p_ordering_hash) + WITH ORDINALITY AS input(queue, state, ordering_hash, ord) + LEFT JOIN awa.queue_meta AS meta ON meta.queue = input.queue; + + -- Claims go in unique-key order, so two producers with overlapping keys + -- cannot wait on each other in opposite orders. + IF p_skip_unique_conflicts THEN + WITH claimed AS ( + INSERT INTO awa.job_unique_claims (unique_key, job_id) + SELECT input.unique_key, v_job_ids[input.ord] + FROM unnest(p_unique_key, p_unique_states, p_state) + WITH ORDINALITY AS input(unique_key, unique_states, state, ord) + WHERE input.unique_key IS NOT NULL + AND input.unique_states IS NOT NULL + AND awa.job_state_in_bitmask(input.unique_states::bit(8), input.state::awa.job_state) + ORDER BY input.unique_key, input.ord + ON CONFLICT (unique_key) DO NOTHING + RETURNING job_id + ) + SELECT COALESCE(array_agg(wanted.job_id), ARRAY[]::BIGINT[]) + INTO v_skipped + FROM ( + SELECT v_job_ids[input.ord] AS job_id + FROM unnest(p_unique_key, p_unique_states, p_state) + WITH ORDINALITY AS input(unique_key, unique_states, state, ord) + WHERE input.unique_key IS NOT NULL + AND input.unique_states IS NOT NULL + AND awa.job_state_in_bitmask(input.unique_states::bit(8), input.state::awa.job_state) + ) AS wanted + WHERE NOT EXISTS (SELECT 1 FROM claimed WHERE claimed.job_id = wanted.job_id); + ELSE + INSERT INTO awa.job_unique_claims (unique_key, job_id) + SELECT input.unique_key, v_job_ids[input.ord] + FROM unnest(p_unique_key, p_unique_states, p_state) + WITH ORDINALITY AS input(unique_key, unique_states, state, ord) + WHERE input.unique_key IS NOT NULL + AND input.unique_states IS NOT NULL + AND awa.job_state_in_bitmask(input.unique_states::bit(8), input.state::awa.job_state) + ORDER BY input.unique_key, input.ord; + END IF; + + -- Register lanes this batch is the first to use. The head tables carry a + -- row trigger that takes the lane lock before the conflict check, so only + -- missing lanes are proposed: a batch for an existing lane must not queue + -- behind that lane's producers before it has any work to show for it. + INSERT INTO queue_lanes (queue, priority) + SELECT lanes.queue, lanes.priority + FROM ( + SELECT DISTINCT input.queue, input.priority + FROM unnest(p_queue, p_priority, p_state) + WITH ORDINALITY AS input(queue, priority, state, ord) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS lanes + WHERE NOT EXISTS ( + SELECT 1 FROM queue_lanes AS existing + WHERE existing.queue = lanes.queue AND existing.priority = lanes.priority + ) + ORDER BY lanes.queue COLLATE "C", lanes.priority + ON CONFLICT ON CONSTRAINT queue_lanes_pkey DO NOTHING; + + INSERT INTO queue_enqueue_heads (queue, priority, enqueue_shard) + SELECT lanes.queue, lanes.priority, lanes.shard + FROM ( + SELECT DISTINCT input.queue, input.priority, v_shards[input.ord] AS shard + FROM unnest(p_queue, p_priority, p_state) + WITH ORDINALITY AS input(queue, priority, state, ord) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS lanes + WHERE NOT EXISTS ( + SELECT 1 FROM queue_enqueue_heads AS existing + WHERE existing.queue = lanes.queue + AND existing.priority = lanes.priority + AND existing.enqueue_shard = lanes.shard + ) + ORDER BY lanes.queue COLLATE "C", lanes.priority, lanes.shard + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + + INSERT INTO queue_claim_heads (queue, priority, enqueue_shard) + SELECT lanes.queue, lanes.priority, lanes.shard + FROM ( + SELECT DISTINCT input.queue, input.priority, v_shards[input.ord] AS shard + FROM unnest(p_queue, p_priority, p_state) + WITH ORDINALITY AS input(queue, priority, state, ord) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS lanes + WHERE NOT EXISTS ( + SELECT 1 FROM queue_claim_heads AS existing + WHERE existing.queue = lanes.queue + AND existing.priority = lanes.priority + AND existing.enqueue_shard = lanes.shard + ) + ORDER BY lanes.queue COLLATE "C", lanes.priority, lanes.shard + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + + -- Re-enqueues that are not yet due carry no lane. + INSERT INTO deferred_jobs ( + job_id, kind, queue, args, state, priority, attempt, run_lease, + max_attempts, run_at, attempted_at, finalized_at, created_at, + unique_key, unique_states, payload + ) + SELECT + v_job_ids[input.ord], + input.kind, + input.queue, + input.args, + input.state::awa.job_state, + input.priority, + 0, + 0, + input.max_attempts, + COALESCE(input.run_at, v_now), + NULL, + NULL, + v_now, + input.unique_key, + input.unique_states, + jsonb_build_object( + 'metadata', input.metadata, + 'tags', input.tags, + 'errors', '[]'::jsonb, + 'progress', NULL + ) + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_run_at, + p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + WHERE input.state <> 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped); + + -- The result rows do not depend on the lanes, so build them before the + -- lock window opens. + RETURN QUERY + SELECT + v_job_ids[input.ord], + input.kind, + input.queue, + input.args, + input.state::awa.job_state, + input.priority, + 0::smallint, + input.max_attempts, + COALESCE(input.run_at, v_now), + NULL::timestamptz, + NULL::timestamptz, + NULL::timestamptz, + NULL::timestamptz, + v_now, + ARRAY[]::jsonb[], + input.metadata, + ARRAY( + SELECT tag.value + FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, tag_ord) + ORDER BY tag.tag_ord + ), + input.unique_key, + input.unique_states::bit(8), + NULL::uuid, + NULL::timestamptz, + NULL::text, + NULL::text, + NULL::text, + NULL::text, + 0::bigint, + NULL::jsonb + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_run_at, + p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + WHERE v_job_ids[input.ord] <> ALL (v_skipped) + ORDER BY input.ord; + + -- Lane lock window: reserve each lane's sequence range once, in + -- (queue, priority, shard) order, then write the ready rows. The locks are + -- held to the end of the transaction, which keeps the rows of a lane + -- visible in lane_seq order. + FOR v_lane IN + SELECT lanes.queue, lanes.priority, lanes.shard, lanes.row_count + FROM ( + SELECT input.queue, input.priority, v_shards[input.ord] AS shard, + count(*) AS row_count + FROM unnest(p_queue, p_priority, p_state) + WITH ORDINALITY AS input(queue, priority, state, ord) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + GROUP BY input.queue, input.priority, v_shards[input.ord] + ) AS lanes + ORDER BY lanes.queue COLLATE "C", lanes.priority, lanes.shard + LOOP + EXECUTE format('SELECT %I.reserve_enqueue_seq($1, $2, $3, $4)', v_schema) + INTO v_start + USING v_lane.queue, v_lane.priority, v_lane.shard, v_lane.row_count::BIGINT; + + v_lane_queues := v_lane_queues || v_lane.queue; + v_lane_priorities := v_lane_priorities || v_lane.priority; + v_lane_shards := v_lane_shards || v_lane.shard; + v_lane_starts := v_lane_starts || v_start; + END LOOP; + + IF cardinality(v_lane_queues) > 0 THEN + SELECT rc.slot, rc.generation + INTO v_ready_slot, v_ready_generation + FROM ring_cursor('queue') AS rc; + + INSERT INTO ready_entries ( + ready_slot, ready_generation, job_id, kind, queue, args, priority, + attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, + attempted_at, created_at, unique_key, unique_states, payload + ) + SELECT + v_ready_slot, + v_ready_generation, + v_job_ids[placed.ord], + placed.kind, + placed.queue, + placed.args, + placed.priority, + 0, + 0, + placed.max_attempts, + lane.start_seq + placed.lane_rank - 1, + lane.shard, + COALESCE(placed.run_at, v_now), + NULL, + v_now, + placed.unique_key, + placed.unique_states, + jsonb_build_object( + 'metadata', placed.metadata, + 'tags', placed.tags, + 'errors', '[]'::jsonb, + 'progress', NULL + ) + FROM ( + SELECT + input.*, + v_shards[input.ord] AS shard, + row_number() OVER ( + PARTITION BY input.queue, input.priority, v_shards[input.ord] + ORDER BY input.ord + ) AS lane_rank + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_run_at, + p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS placed + JOIN unnest(v_lane_queues, v_lane_priorities, v_lane_shards, v_lane_starts) + AS lane(queue, priority, shard, start_seq) + ON lane.queue = placed.queue + AND lane.priority = placed.priority + AND lane.shard = placed.shard; + + -- One segment per run of consecutive lane_seq values sharing a run_at. + INSERT INTO ready_segments ( + ready_slot, ready_generation, queue, priority, enqueue_shard, + first_lane_seq, next_lane_seq, first_run_at + ) + SELECT + v_ready_slot, + v_ready_generation, + runs.queue, + runs.priority, + runs.shard, + min(runs.lane_seq), + max(runs.lane_seq) + 1, + runs.run_at + FROM ( + SELECT + placed.*, + placed.lane_seq - row_number() OVER ( + PARTITION BY placed.queue, placed.priority, placed.shard, placed.run_at + ORDER BY placed.lane_seq + ) AS run_id + FROM ( + SELECT + input.queue, + input.priority, + v_shards[input.ord] AS shard, + COALESCE(input.run_at, v_now) AS run_at, + lane.start_seq + row_number() OVER ( + PARTITION BY input.queue, input.priority, v_shards[input.ord] + ORDER BY input.ord + ) - 1 AS lane_seq + FROM unnest(p_queue, p_priority, p_state, p_run_at) + WITH ORDINALITY AS input(queue, priority, state, run_at, ord) + JOIN unnest(v_lane_queues, v_lane_priorities, v_lane_shards, v_lane_starts) + AS lane(queue, priority, shard, start_seq) + ON lane.queue = input.queue + AND lane.priority = input.priority + AND lane.shard = v_shards[input.ord] + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS placed + ) AS runs + GROUP BY runs.queue, runs.priority, runs.shard, runs.run_at, runs.run_id + ON CONFLICT (ready_slot, ready_generation, queue, priority, enqueue_shard, first_lane_seq) + DO NOTHING; + + PERFORM pg_notify('awa:' || notified.queue, '') + FROM (SELECT DISTINCT unnest(v_lane_queues) AS queue) AS notified; + END IF; + + PERFORM set_config('search_path', v_old_search_path, true); +END; +$$ LANGUAGE plpgsql VOLATILE +SET search_path = pg_catalog, awa, public; + +INSERT INTO awa.schema_version (version, description) +VALUES (55, 'Batch producer insert with a short enqueue lane lock window') +ON CONFLICT (version) DO NOTHING; diff --git a/awa-model/migrations/v056_enqueue_round_trips.sql b/awa-model/migrations/v056_enqueue_round_trips.sql new file mode 100644 index 000000000..89a8c8b60 --- /dev/null +++ b/awa-model/migrations/v056_enqueue_round_trips.sql @@ -0,0 +1,34 @@ +-- v056: Fewer round trips and sequence writes per queue-storage enqueue. +-- Tracking: https://github.com/hardbyte/awa/issues/522 +-- reserve_enqueue_seq keeps its signature, lock key and return value. It +-- provisions lane sequences only for a lane row without a sequence name and +-- reserves a multi-row range with one sequence write instead of two. +-- enqueue_ready_rows is new: one ready batch's lane reservations, ring +-- cursor read, ready and segment appends and NOTIFY in a single call. +-- N-1 binaries call reserve_enqueue_seq unchanged and never call the new +-- function; current binaries probe for it and keep the multi-statement path +-- on a schema without it. +-- Live load: function replacement only, with no table DDL or row scans. +-- Function lock waits and schema count determine wall time. + +DO $$ +DECLARE + v_schema TEXT; +BEGIN + FOR v_schema IN + SELECT discovered.schema_name FROM awa.queue_storage_schemas() AS discovered + LOOP + IF to_regclass(format('%I.queue_enqueue_heads', v_schema)) IS NULL + OR to_regclass(format('%I.queue_claim_heads', v_schema)) IS NULL + OR to_regclass(format('%I.queue_ring_state', v_schema)) IS NULL THEN + CONTINUE; + END IF; + + PERFORM awa.refresh_enqueue_functions(v_schema); + END LOOP; +END; +$$; + +INSERT INTO awa.schema_version (version, description) +VALUES (56, 'Fewer round trips and sequence writes per queue-storage enqueue') +ON CONFLICT (version) DO NOTHING; diff --git a/awa-model/src/insert.rs b/awa-model/src/insert.rs index e221b0fa7..6bb26f47b 100644 --- a/awa-model/src/insert.rs +++ b/awa-model/src/insert.rs @@ -11,13 +11,21 @@ use sqlx::{PgExecutor, PgPool}; // On queue storage every available-state insert holds a transaction-scoped // advisory lock on its `(queue, priority, enqueue_shard)` lane until the // enclosing transaction ends, so two transactions that enqueue into the same -// lanes in opposite orders can deadlock. Every multi-row enqueue therefore -// takes its lane locks in one canonical order: `(queue, priority, shard)` -// ascending, queues compared bytewise (`COLLATE "C"`, the order of Rust's -// `str`). The batch inserts here sort rows by that key, stably, before they -// reach `awa.insert_job_compat`, and put the rows back in input order for the -// caller. `QueueStorage` batch enqueues group rows in a `BTreeMap` keyed the -// same way. +// lanes in opposite orders can deadlock. `insert_many_copy` and +// `insert_many_skip_unique_conflicts` use `awa.insert_jobs_compat_batch` +// (migration v055), which does the +// lane-independent work first and then reserves its lanes in one canonical +// order: `(queue, priority, shard)` ascending, queues compared bytewise +// (`COLLATE "C"`, the order of Rust's `str`). `QueueStorage` batch enqueues +// group rows in a `BTreeMap` keyed the same way. Holding the lanes only across +// the reservation and the ready-row writes keeps a lane's lock window +// independent of the batch's unique-claim and payload work. +// +// Generic `insert_many` always uses the per-row path, as does +// `insert_many_copy` before the migration: rows are sorted by the same key before +// they reach `awa.insert_job_compat`, and returned in input order. Both paths +// take the same lane lock, so producers on either, and older binaries, share a +// lane without reordering it. // // The shard is derived as `insert_job_compat` derives it: rows with an // ordering key use `ordering_key_hash64 % enqueue_shards`, other rows @@ -211,11 +219,35 @@ impl PreparedJobInsert { /// job, because sqlx has no array encoding for the enum, the bit string, or a /// ragged array of arrays. `ordering_hash` carries each row's /// [`ordering_key_hash64`] as decimal text (the value exceeds `bigint`) so the -/// statement can place keyed rows on the shard `insert_job_compat` will pick. +/// function can place keyed rows on the shard `insert_job_compat` would pick. /// -/// Rows are inserted in canonical lane order (see "Lane lock order" above) and -/// returned in input order. -const INSERT_MANY_SQL: &str = r#" +/// `$13` skips rows whose unique key is already claimed instead of failing the +/// batch. Rows come back in input order; skipped rows are absent. +const INSERT_MANY_BATCH_SQL: &str = r#" + SELECT * + FROM awa.insert_jobs_compat_batch( + $1::text[], + $2::text[], + $3::jsonb[], + $4::text[], + $5::smallint[], + $6::smallint[], + $7::timestamptz[], + $8::jsonb[], + $9::jsonb[], + $10::bytea[], + $11::text[], + $12::text[], + $13 + ) +"#; + +/// The per-row form of [`INSERT_MANY_BATCH_SQL`], for schemas without +/// `awa.insert_jobs_compat_batch`: one `awa.insert_job_compat` call per row, +/// rows inserted in canonical lane order (see "Lane lock order" above) and +/// returned in input order. It works before and after the batch migration, +/// including when one process writes to databases at different schema versions. +const INSERT_MANY_ROW_FUNCTION_SQL: &str = r#" WITH input AS MATERIALIZED ( SELECT raw.*, @@ -476,7 +508,17 @@ fn ordering_hash_text(row: &PreparedRow) -> Option { /// Insert multiple jobs in a single statement. /// /// Supports uniqueness constraints — jobs with `unique` opts will have their -/// `unique_key` and `unique_states` computed and included. +/// `unique_key` and `unique_states` computed and included. A job whose unique +/// key is already claimed fails the whole batch with a unique violation +/// (SQLSTATE `23505`, as [`AwaError::Database`]); use +/// [`insert_many_skip_unique_conflicts`] to insert the rest instead. +/// +/// Calls `awa.insert_job_compat` per row, in canonical lane order, on every +/// supported schema. A generic executor is consumed by one statement and +/// cannot be probed separately, so this path never assumes another database +/// has the same schema. [`insert_many_copy`] probes its concrete connection +/// and uses the batch function when available; +/// [`insert_many_skip_unique_conflicts`] requires that function. #[tracing::instrument( skip(executor, jobs), fields( @@ -495,9 +537,119 @@ where if jobs.is_empty() { return Ok(Vec::new()); } + let rows = precompute_rows(jobs)?; + Ok(bind_row_function_columns( + sqlx::query_as::<_, JobRow>(INSERT_MANY_ROW_FUNCTION_SQL), + &rows, + ) + .fetch_all(executor) + .await?) +} +/// Insert multiple jobs in a single statement, skipping those whose unique key +/// is already claimed. +/// +/// Returns the rows that were inserted, in input order; a skipped job has no +/// entry in the result. Jobs without `unique` opts are always inserted. A key +/// repeated inside the batch inserts its first job. Unlike [`insert_many`], +/// which fails the batch with a unique violation, this is what a caller that +/// replays work wants, and it takes each lane lock once for the batch rather +/// than once per job as a row-by-row retry does. +/// +/// Requires migration v055 (`awa migrate`); on an older schema it returns +/// [`AwaError::Validation`]. [`insert_many_copy`] skips unique conflicts on +/// every schema, through a connection, and is the Python +/// `insert_many_copy`. +#[tracing::instrument( + skip(executor, jobs), + fields( + job.count = jobs.len(), + otel.name = "send", + otel.kind = "producer", + messaging.system = "awa", + messaging.operation.r#type = "send", + messaging.batch.message_count = jobs.len(), + ) +)] +pub async fn insert_many_skip_unique_conflicts<'e, E>( + executor: E, + jobs: &[InsertParams], +) -> Result, AwaError> +where + E: PgExecutor<'e>, +{ + if jobs.is_empty() { + return Ok(Vec::new()); + } let rows = precompute_rows(jobs)?; - let results = sqlx::query_as::<_, JobRow>(INSERT_MANY_SQL) + match bind_batch_columns(sqlx::query_as::<_, JobRow>(INSERT_MANY_BATCH_SQL), &rows) + .bind(true) + .fetch_all(executor) + .await + { + Ok(results) => Ok(results), + Err(err) if is_missing_batch_function(&err) => Err(AwaError::Validation( + "insert_many_skip_unique_conflicts needs migration v055; run `awa migrate`".to_string(), + )), + Err(err) => Err(err.into()), + } +} + +/// `undefined_function` naming the batch function: the schema predates +/// migration v055 or the function was dropped. +fn is_missing_batch_function(err: &sqlx::Error) -> bool { + matches!( + err, + sqlx::Error::Database(db_err) + if db_err.code().as_deref() == Some("42883") + && db_err.message().contains("insert_jobs_compat_batch") + ) +} + +fn bind_batch_columns<'q, O>( + query: sqlx::query::QueryAs<'q, sqlx::Postgres, O, sqlx::postgres::PgArguments>, + rows: &'q [PreparedRow], +) -> sqlx::query::QueryAs<'q, sqlx::Postgres, O, sqlx::postgres::PgArguments> { + query + .bind(rows.iter().map(|row| row.kind.as_str()).collect::>()) + .bind( + rows.iter() + .map(|row| row.queue.as_str()) + .collect::>(), + ) + .bind(rows.iter().map(|row| &row.args).collect::>()) + .bind( + rows.iter() + .map(|row| row.state_db_str()) + .collect::>(), + ) + .bind(rows.iter().map(|row| row.priority).collect::>()) + .bind(rows.iter().map(|row| row.max_attempts).collect::>()) + .bind(rows.iter().map(|row| row.run_at).collect::>()) + .bind(rows.iter().map(|row| &row.metadata).collect::>()) + .bind( + rows.iter() + .map(|row| serde_json::json!(row.tags)) + .collect::>(), + ) + .bind( + rows.iter() + .map(|row| row.unique_key.as_deref()) + .collect::>(), + ) + .bind( + rows.iter() + .map(|row| row.unique_states.as_deref()) + .collect::>(), + ) + .bind(rows.iter().map(ordering_hash_text).collect::>()) +} + +fn bind_row_function_columns<'q>( + query: sqlx::query::QueryAs<'q, sqlx::Postgres, JobRow, sqlx::postgres::PgArguments>, + rows: &'q [PreparedRow], +) -> sqlx::query::QueryAs<'q, sqlx::Postgres, JobRow, sqlx::postgres::PgArguments> { + query .bind(rows.iter().map(|row| row.kind.as_str()).collect::>()) .bind( rows.iter() @@ -535,46 +687,14 @@ where .collect::>(), ) .bind(rows.iter().map(ordering_hash_text).collect::>()) - .fetch_all(executor) - .await?; - - Ok(results) -} - -/// CTEs that number the staged rows in staging order and re-sort them into -/// canonical lane order (see "Lane lock order" above); `staged_in_lane_order` -/// carries the staging-order position as `ord`. -macro_rules! staged_in_lane_order_ctes { - () => { - r#" - staged AS ( - SELECT *, row_number() OVER () AS ord - FROM pg_temp.awa_copy_staging - ), - staged_in_lane_order AS MATERIALIZED ( - SELECT - staged.*, - MOD( - CASE WHEN staged.ordering_hash IS NULL - THEN pg_backend_pid()::numeric - ELSE staged.ordering_hash::numeric - END, - COALESCE( - (SELECT meta.enqueue_shards FROM awa.queue_meta AS meta WHERE meta.queue = staged.queue), - 1 - )::numeric - ) AS lane_shard - FROM staged - ORDER BY staged.queue COLLATE "C", staged.priority, lane_shard, staged.ord - ) - "# - }; } /// Insert many jobs through the compatibility COPY surface. /// /// Uses a temp staging table with no constraints for fast COPY ingestion, -/// then routes staged rows through `awa.insert_job_compat`. Accepts +/// then inserts them through `awa.insert_jobs_compat_batch`, or, on a schema +/// without migration v055, through `awa.insert_job_compat` row by row. Jobs +/// whose unique key is already claimed are skipped. Accepts /// `&mut PgConnection` so callers can use pool connections or transactions /// (Transaction derefs to PgConnection). For high-volume queue-storage /// producers, prefer [`QueueStorage::enqueue_params_copy`](crate::QueueStorage::enqueue_params_copy) @@ -646,10 +766,152 @@ pub async fn insert_many_copy( copy_in.finish().await?; } - // 3. INSERT...SELECT from staging into real table + // 3. Insert the staged rows. With the batch function, unique conflicts are + // skipped in the one statement; without it, each unique row goes through + // the per-row function in its own savepoint. let has_unique = rows.iter().any(|r| r.unique_key.is_some()); + let results = if batch_function_available(conn).await? { + insert_staged_rows_with_batch_function(conn, has_unique).await? + } else { + insert_staged_rows_per_row(conn, has_unique).await? + }; + + // Keep the session-local staging table reusable across multiple COPY calls + // within the same outer transaction. + sqlx::query("DELETE FROM pg_temp.awa_copy_staging") + .execute(&mut *conn) + .await?; - let results = if has_unique { + Ok(results) +} + +/// Convenience wrapper that acquires a connection from the pool. +/// +/// Wraps the operation in a transaction so the staging rows are cleaned up at +/// commit time even if the caller does not reuse the connection afterward. +#[tracing::instrument( + skip(pool, jobs), + fields( + job.count = jobs.len(), + otel.name = "send", + otel.kind = "producer", + messaging.system = "awa", + messaging.operation.r#type = "send", + messaging.batch.message_count = jobs.len(), + ) +)] +pub async fn insert_many_copy_from_pool( + pool: &PgPool, + jobs: &[InsertParams], +) -> Result, AwaError> { + if jobs.is_empty() { + return Ok(Vec::new()); + } + + let mut tx = pool.begin().await?; + let results = insert_many_copy(&mut tx, jobs).await?; + tx.commit().await?; + + Ok(results) +} + +/// Probe this connection on every COPY call. A process can write to databases +/// at different migration levels, and the generic executor path cannot safely +/// seed a shared answer for them. +async fn batch_function_available(conn: &mut PgConnection) -> Result { + Ok(sqlx::query_scalar( + "SELECT to_regprocedure('awa.insert_jobs_compat_batch(text[],text[],jsonb[],text[],smallint[],smallint[],timestamptz[],jsonb[],jsonb[],bytea[],text[],text[],boolean)') IS NOT NULL", + ) + .fetch_one(&mut *conn) + .await?) +} + +async fn insert_staged_rows_with_batch_function( + conn: &mut PgConnection, + skip_unique_conflicts: bool, +) -> Result, AwaError> { + let result = sqlx::query_as::<_, JobRow>( + r#" + WITH staged AS ( + SELECT *, row_number() OVER () AS ord + FROM pg_temp.awa_copy_staging + ), + columns AS ( + SELECT + array_agg(kind ORDER BY ord) AS kinds, + array_agg(queue ORDER BY ord) AS queues, + array_agg(args ORDER BY ord) AS args, + array_agg(state::text ORDER BY ord) AS states, + array_agg(priority ORDER BY ord) AS priorities, + array_agg(max_attempts ORDER BY ord) AS max_attempts, + array_agg(run_at ORDER BY ord) AS run_ats, + array_agg(metadata ORDER BY ord) AS metadata, + array_agg(to_jsonb(tags) ORDER BY ord) AS tags, + array_agg(unique_key ORDER BY ord) AS unique_keys, + array_agg(unique_states::text ORDER BY ord) AS unique_states, + array_agg(ordering_hash ORDER BY ord) AS ordering_hashes + FROM staged + ) + SELECT inserted.* + FROM columns + CROSS JOIN LATERAL awa.insert_jobs_compat_batch( + columns.kinds, + columns.queues, + columns.args, + columns.states, + columns.priorities, + columns.max_attempts, + columns.run_ats, + columns.metadata, + columns.tags, + columns.unique_keys, + columns.unique_states, + columns.ordering_hashes, + $1 + ) AS inserted + "#, + ) + .bind(skip_unique_conflicts) + .fetch_all(&mut *conn) + .await?; + Ok(result) +} + +/// CTEs that number the staged rows in staging order and re-sort them into +/// canonical lane order (see "Lane lock order" above); `staged_in_lane_order` +/// carries the staging-order position as `ord`. +macro_rules! staged_in_lane_order_ctes { + () => { + r#" + staged AS ( + SELECT *, row_number() OVER () AS ord + FROM pg_temp.awa_copy_staging + ), + staged_in_lane_order AS MATERIALIZED ( + SELECT + staged.*, + MOD( + CASE WHEN staged.ordering_hash IS NULL + THEN pg_backend_pid()::numeric + ELSE staged.ordering_hash::numeric + END, + COALESCE( + (SELECT meta.enqueue_shards FROM awa.queue_meta AS meta WHERE meta.queue = staged.queue), + 1 + )::numeric + ) AS lane_shard + FROM staged + ORDER BY staged.queue COLLATE "C", staged.priority, lane_shard, staged.ord + ) + "# + }; +} + +async fn insert_staged_rows_per_row( + conn: &mut PgConnection, + has_unique: bool, +) -> Result, AwaError> { + let results: Vec = if has_unique { // The compatibility `awa.jobs` surface is now a view backed by hot and // deferred tables, so the old `ON CONFLICT` path is no longer available // here. Keep COPY for staging/parsing, then insert unique rows one at a @@ -790,43 +1052,6 @@ pub async fn insert_many_copy( .fetch_all(&mut *conn) .await? }; - - // Keep the session-local staging table reusable across multiple COPY calls - // within the same outer transaction. - sqlx::query("DELETE FROM pg_temp.awa_copy_staging") - .execute(&mut *conn) - .await?; - - Ok(results) -} - -/// Convenience wrapper that acquires a connection from the pool. -/// -/// Wraps the operation in a transaction so the staging rows are cleaned up at -/// commit time even if the caller does not reuse the connection afterward. -#[tracing::instrument( - skip(pool, jobs), - fields( - job.count = jobs.len(), - otel.name = "send", - otel.kind = "producer", - messaging.system = "awa", - messaging.operation.r#type = "send", - messaging.batch.message_count = jobs.len(), - ) -)] -pub async fn insert_many_copy_from_pool( - pool: &PgPool, - jobs: &[InsertParams], -) -> Result, AwaError> { - if jobs.is_empty() { - return Ok(Vec::new()); - } - - let mut tx = pool.begin().await?; - let results = insert_many_copy(&mut tx, jobs).await?; - tx.commit().await?; - Ok(results) } diff --git a/awa-model/src/lib.rs b/awa-model/src/lib.rs index 66a5dd4ee..7f23ec69a 100644 --- a/awa-model/src/lib.rs +++ b/awa-model/src/lib.rs @@ -45,7 +45,8 @@ pub use cron::{CronJobRow, CronMissedFirePolicy, PeriodicJob, PeriodicJobBuilder pub use dlq::{DlqMetadata, DlqRow, ListDlqFilter, RetryFromDlqOpts}; pub use error::{map_sqlx_error, AwaError}; pub use insert::{ - insert, insert_many, insert_many_copy, insert_many_copy_from_pool, insert_prepared, insert_with, + insert, insert_many, insert_many_copy, insert_many_copy_from_pool, + insert_many_skip_unique_conflicts, insert_prepared, insert_with, }; pub use job::{InsertOpts, InsertParams, JobRow, JobState, UniqueOpts}; pub use output::{AttemptError, JobOutput, OutputDescriptor, OutputError, MAX_OUTPUT_BYTES}; diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index 385035b65..ee913b527 100644 --- a/awa-model/src/migrations.rs +++ b/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 = 52; +pub const CURRENT_VERSION: i32 = 56; /// Migrations that require an exclusive (no-live-runtime) upgrade window. /// @@ -292,6 +292,26 @@ const MIGRATIONS: &[(i32, &str, &[&str])] = &[ ), (51, "Per-slot claim functions", &[V23_UP, V51_UP]), (52, "Retained job output", &[V23_UP, V52_UP]), + ( + 53, + "Create queue lanes as the schema owner", + &[V23_UP, V53_UP], + ), + ( + 54, + "Planner statistics refresh for queue-storage tables autovacuum never analyses", + &[V23_UP, V54_UP], + ), + ( + 55, + "Batch producer insert with a short enqueue lane lock window", + &[V55_UP], + ), + ( + 56, + "Fewer round trips and sequence writes per queue-storage enqueue", + &[V23_UP, V56_UP], + ), ]; const V1_UP: &str = include_str!("../migrations/v001_canonical_schema.sql"); @@ -352,6 +372,10 @@ const V49_UP: &str = include_str!("../migrations/v049_prepared_lane_sequences.sq const V50_UP: &str = include_str!("../migrations/v050_drop_ready_lane_shard_index.sql"); const V51_UP: &str = include_str!("../migrations/v051_claim_slot_functions.sql"); 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_analyze_queue_storage_statistics.sql"); +const V55_UP: &str = include_str!("../migrations/v055_insert_jobs_batch_compat.sql"); +const V56_UP: &str = include_str!("../migrations/v056_enqueue_round_trips.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. @@ -417,9 +441,9 @@ pub async fn run_with_options(pool: &PgPool, options: MigrateOptions) -> Result< // only sound because every migration is transaction-safe (no // `CREATE INDEX CONCURRENTLY`, `VACUUM`, or explicit transaction control). let mut retry = 0; - loop { + let version_before = loop { match run_migration_transaction(pool, &options, lock_key).await { - Ok(()) => break, + Ok(version_before) => break version_before, Err(error) if is_migration_deadlock(&error) && retry < MIGRATION_DEADLOCK_RETRIES => { retry += 1; let delay = std::time::Duration::from_millis(50 * (1 << (retry - 1))); @@ -433,6 +457,10 @@ pub async fn run_with_options(pool: &PgPool, options: MigrateOptions) -> Result< } Err(error) => return Err(error), } + }; + + if version_before < LANE_PROVISIONING_VERSION { + repair_lane_sequences_after_migration(pool).await; } // Best-effort admin-metadata cache warmup, deliberately outside the @@ -443,29 +471,73 @@ pub async fn run_with_options(pool: &PgPool, options: MigrateOptions) -> Result< Ok(()) } +/// Returns the schema version the transaction started from. async fn run_migration_transaction( pool: &PgPool, options: &MigrateOptions, lock_key: i64, -) -> Result<(), AwaError> { +) -> Result { let mut tx = pool.begin().await?; sqlx::query("SELECT pg_advisory_xact_lock($1)") .bind(lock_key) .execute(&mut *tx) .await?; - apply_migrations(&mut tx, options).await?; + let version_before = apply_migrations(&mut tx, options).await?; tx.commit().await?; - Ok(()) + Ok(version_before) +} + +/// The migration that makes lane sequences belong to the schema owner. +const LANE_PROVISIONING_VERSION: i32 = 53; + +/// Repair existing lane sequences once v053 has committed. +/// +/// `ALTER SEQUENCE ... OWNER` takes the sequence's ACCESS EXCLUSIVE lock until +/// its transaction ends, so it cannot run inside the migration transaction +/// without stalling every enqueue and claim on those lanes until the whole +/// upgrade commits. Each sequence is repaired in its own short transaction +/// with a bounded lock wait instead. Best-effort: the schema is committed, and +/// whatever is left is reported with the command that finishes it. +/// +/// Uses a detached connection for the reason given on +/// [`warm_admin_metadata_cache`]. +async fn repair_lane_sequences_after_migration(pool: &PgPool) { + let mut conn = match pool.acquire().await { + Ok(conn) => conn.detach(), + Err(error) => { + warn!(%error, "Could not connect to repair lane sequences; run `awa storage lane-sequences --repair`"); + return; + } + }; + match crate::storage::repair_lane_sequences_in_every_schema(&mut conn).await { + Ok(repairs) => { + let repaired = repairs + .iter() + .filter(|repair| repair.outcome == crate::storage::LaneSequenceOutcome::Repaired) + .count(); + if repaired > 0 { + info!(repaired, "Gave lane sequences to the schema owner"); + } + for next_step in crate::storage::lane_sequence_repair_next_steps(&repairs) { + warn!("{next_step}"); + } + } + Err(error) => { + warn!(%error, "Could not repair lane sequences; run `awa storage lane-sequences --repair`"); + } + } + let _ = conn.close().await; } fn is_migration_deadlock(error: &AwaError) -> bool { error.is_deadlock() } +/// Returns the schema version before this run. async fn apply_migrations( conn: &mut PgConnection, options: &MigrateOptions, -) -> Result<(), AwaError> { +) -> Result { let has_schema: bool = sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM pg_namespace WHERE nspname = 'awa')") .fetch_one(&mut *conn) @@ -558,7 +630,7 @@ async fn apply_migrations( info!(version = current, "Schema is up to date"); } - Ok(()) + Ok(current) } /// Warm the admin-metadata cache after migrations commit. diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 339001705..ae205f7e9 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -13,7 +13,7 @@ use sqlx::{PgPool, Postgres}; use std::collections::hash_map::DefaultHasher; use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::hash::{Hash, Hasher}; -use std::sync::atomic::{AtomicU16, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU16, AtomicUsize, Ordering}; use std::sync::Mutex; use std::time::{Duration, Instant}; use uuid::Uuid; @@ -713,6 +713,35 @@ async fn set_prune_lock_timeout_tx( Ok(()) } +/// Commit or roll back a maintenance transaction that has set the prune +/// `lock_timeout`. +/// +/// The timeout applies to every later statement in the transaction, not only +/// the explicit `LOCK TABLE`: an insert, update or read of another relation can +/// also queue behind a conflicting lock (autovacuum's heap truncation takes +/// `ACCESS EXCLUSIVE`, for example). A lock timeout anywhere therefore rolls the +/// whole transaction back and reports `blocked`, so the caller backs off and +/// retries instead of surfacing an error. +async fn settle_prune_tx( + tx: sqlx::Transaction<'_, Postgres>, + result: Result, + blocked: T, +) -> Result { + match result { + Ok(outcome) => { + tx.commit().await.map_err(map_sqlx_error)?; + Ok(outcome) + } + Err(err) => { + let _ = tx.rollback().await; + match err { + AwaError::Database(ref db_err) if is_lock_contention_error(db_err) => Ok(blocked), + err => Err(err), + } + } + } +} + fn validate_ident(ident: &str) -> Result<(), AwaError> { let mut chars = ident.chars(); match chars.next() { @@ -2847,9 +2876,161 @@ struct ReadySegmentInsert { first_run_at: DateTime, } +/// Argument types of `{schema}.enqueue_ready_rows`, as `to_regprocedure` +/// spells them. +const ENQUEUE_READY_ROWS_ARGUMENT_TYPES: &str = "text[], smallint[], smallint[], bigint[], \ + bigint[], text[], jsonb[], smallint[], bigint[], smallint[], timestamptz[], \ + timestamptz[], timestamptz[], bytea[], text[], jsonb[], integer[], bigint[], bigint[], \ + timestamptz[], text[]"; + +/// How long an "enqueue_ready_rows is unavailable" observation stands before +/// the store probes again, so a runtime started before v056 picks the +/// function up once the migration lands. +const ENQUEUE_FUNCTION_REPROBE: Duration = Duration::from_secs(60); + +fn enqueue_function_unavailable(error: &AwaError) -> bool { + matches!(error, AwaError::Database(error) if matches!( + error.as_database_error().and_then(|db| db.code()).as_deref(), + Some("42883" | "42501") + )) +} + +/// One ready row as `{schema}.enqueue_ready_rows` takes it. Queue, priority +/// and shard come from the row's lane; `lane_seq` from the lane's +/// reservation; a `None` job id from the schema's job id sequence. +struct EnqueueReadyRow<'r> { + job_id: Option, + kind: &'r str, + args: &'r serde_json::Value, + attempt: i16, + run_lease: i64, + max_attempts: i16, + run_at: DateTime, + attempted_at: Option>, + created_at: DateTime, + unique_key: Option<&'r [u8]>, + unique_states: Option<&'r str>, + payload: &'r serde_json::Value, +} + +impl<'r> EnqueueReadyRow<'r> { + fn from_ready_row(row: &'r RuntimeReadyRow) -> Self { + Self { + job_id: row.job_id, + kind: &row.kind, + args: &row.args, + attempt: row.attempt, + run_lease: row.run_lease, + max_attempts: row.max_attempts, + run_at: row.run_at, + attempted_at: row.attempted_at, + created_at: row.created_at, + unique_key: row.unique_key.as_deref(), + unique_states: row.unique_states.as_deref(), + payload: &row.payload, + } + } + + fn from_insert(row: &'r RuntimeReadyInsert) -> Self { + Self { + job_id: Some(row.job_id), + kind: &row.kind, + args: &row.args, + attempt: row.attempt, + run_lease: row.run_lease, + max_attempts: row.max_attempts, + run_at: row.run_at, + attempted_at: row.attempted_at, + created_at: row.created_at, + unique_key: row.unique_key.as_deref(), + unique_states: row.unique_states.as_deref(), + payload: &row.payload, + } + } +} + +/// Column arrays for one `{schema}.enqueue_ready_rows` call. Lanes are pushed +/// in lock order with their rows contiguous; segment offsets are relative to +/// the start the function reserves for the lane. +#[derive(Default)] +struct EnqueueReadyRowsCall<'r> { + lane_queues: Vec<&'r str>, + lane_priorities: Vec, + lane_shards: Vec, + lane_counts: Vec, + job_ids: Vec>, + kinds: Vec<&'r str>, + args: Vec<&'r serde_json::Value>, + attempts: Vec, + run_leases: Vec, + max_attempts: Vec, + run_ats: Vec>, + attempted_ats: Vec>>, + created_ats: Vec>, + unique_keys: Vec>, + unique_states: Vec>, + payloads: Vec>, + segment_lanes: Vec, + segment_first_offsets: Vec, + segment_next_offsets: Vec, + segment_run_ats: Vec>, +} + +impl<'r> EnqueueReadyRowsCall<'r> { + fn push_lane( + &mut self, + queue: &'r str, + priority: i16, + enqueue_shard: i16, + rows: impl IntoIterator>, + ) { + let first_row = self.run_ats.len(); + for row in rows { + self.job_ids.push(row.job_id); + self.kinds.push(row.kind); + self.args.push(row.args); + self.attempts.push(row.attempt); + self.run_leases.push(row.run_lease); + self.max_attempts.push(row.max_attempts); + self.run_ats.push(row.run_at); + self.attempted_ats.push(row.attempted_at); + self.created_ats.push(row.created_at); + self.unique_keys.push(row.unique_key); + self.unique_states.push(row.unique_states); + self.payloads.push(storage_payload(row.payload)); + } + let count = self.run_ats.len() - first_row; + if count == 0 { + return; + } + self.lane_queues.push(queue); + self.lane_priorities.push(priority); + self.lane_shards.push(enqueue_shard); + self.lane_counts.push(count as i64); + + // Same split as `ready_segments_from_rows` within one lane, whose + // sequence numbers are contiguous: a new segment at each `run_at` + // change. + let lane_number = self.lane_queues.len() as i32; + let mut segment_start = 0usize; + for offset in 1..=count { + if offset == count + || self.run_ats[first_row + offset] != self.run_ats[first_row + segment_start] + { + self.segment_lanes.push(lane_number); + self.segment_first_offsets.push(segment_start as i64); + self.segment_next_offsets.push(offset as i64); + self.segment_run_ats + .push(self.run_ats[first_row + segment_start]); + segment_start = offset; + } + } + } +} + #[cfg(test)] mod ready_segment_tests { - use super::{QueueStorage, RuntimeReadyInsert}; + use super::{EnqueueReadyRow, EnqueueReadyRowsCall, QueueStorage, RuntimeReadyInsert}; use chrono::{Duration, TimeZone, Utc}; fn ready_row(lane_seq: i64, run_at: chrono::DateTime) -> RuntimeReadyInsert { @@ -2901,6 +3082,52 @@ mod ready_segment_tests { assert_eq!(ranges, vec![(1, 3, first), (3, 4, second), (4, 5, first)]); } + + #[test] + fn enqueue_function_segments_match_rust_segments_after_reservation() { + let first = Utc + .with_ymd_and_hms(2026, 6, 14, 12, 0, 0) + .single() + .expect("valid test timestamp"); + let second = first + Duration::seconds(1); + let reserved_start = 41; + let rows: Vec<_> = [first, first, second, first, first, second] + .into_iter() + .enumerate() + .map(|(offset, run_at)| ready_row(reserved_start + offset as i64, run_at)) + .collect(); + + let mut call = EnqueueReadyRowsCall::default(); + call.push_lane( + "segment-q", + 2, + 0, + rows.iter().map(EnqueueReadyRow::from_insert), + ); + let function_ranges: Vec<_> = (0..call.segment_lanes.len()) + .map(|index| { + assert_eq!(call.segment_lanes[index], 1); + ( + reserved_start + call.segment_first_offsets[index], + reserved_start + call.segment_next_offsets[index], + call.segment_run_ats[index], + ) + }) + .collect(); + let rust_ranges: Vec<_> = QueueStorage::ready_segments_from_rows(&rows) + .iter() + .map(|segment| { + ( + segment.first_lane_seq, + segment.next_lane_seq, + segment.first_run_at, + ) + }) + .collect(); + + assert_eq!(function_ranges, rust_ranges); + assert_eq!(call.lane_counts, vec![6]); + } } #[derive(Debug, Clone, sqlx::FromRow)] @@ -3461,8 +3688,22 @@ pub struct QueueStorage { /// TRUNCATEs queue_lanes; `advance_enqueue_head` repairs a stale entry /// (head row gone after a rolled-back ensure_lane). ensured_lanes: Mutex>, + /// Whether `{schema}.provision_lane` (migration v053) was callable the + /// last time a lane was created. Only a positive answer is kept: lane + /// creation is rare, and the probe rides on the lane-lock round trip. + lane_provisioner_present: AtomicBool, + /// Set once a provisioner failure has been logged. + lane_provisioner_warned: AtomicBool, slot_fill_observations: Mutex>, prune_lock_timeout: Duration, + /// Last probe of whether the schema has an executable + /// `enqueue_ready_rows` (v056), and when it ran. A missing function is + /// probed again after `ENQUEUE_FUNCTION_REPROBE`; an available one until + /// a call fails for want of it. An owned enqueue transaction then rolls + /// back and retries its original input once without the function. Cleared + /// by `prepare_schema` and `reset`. + enqueue_function: Mutex>, + enqueue_function_disabled: bool, cancel_phase_hook: CancelPhaseHookSlot, /// Whether this schema has migration v052's output columns (ADR-048), /// and when that was last checked. @@ -3516,13 +3757,25 @@ impl QueueStorage { enqueue_shards_cache: Mutex::new(HashMap::new()), claimer_gates: Mutex::new(HashMap::new()), ensured_lanes: Mutex::new(HashSet::new()), + lane_provisioner_present: AtomicBool::new(false), + lane_provisioner_warned: AtomicBool::new(false), slot_fill_observations: Mutex::new(HashMap::new()), prune_lock_timeout: DEFAULT_PRUNE_LOCK_TIMEOUT, + enqueue_function: Mutex::new(None), + enqueue_function_disabled: false, cancel_phase_hook: CancelPhaseHookSlot(None), job_output_columns: Mutex::new(None), }) } + /// Enqueue through the multi-statement path even when the schema has + /// `enqueue_ready_rows`, as a binary without the function would. + #[doc(hidden)] + pub fn without_enqueue_function(mut self) -> Self { + self.enqueue_function_disabled = true; + self + } + /// Whether this schema can store job outputs (ADR-048, migration v052). /// /// A positive answer is cached for the life of the store. A negative one @@ -3993,6 +4246,7 @@ impl QueueStorage { /// upgrade get built; the rest are no-ops. #[tracing::instrument(skip(self, pool), name = "queue_storage.prepare_schema")] pub async fn prepare_schema(&self, pool: &PgPool) -> Result<(), AwaError> { + self.forget_schema_observations(); let schema = self.schema(); let install_lock_name = format!("awa.queue_storage.install:{schema}"); let mut install_tx = pool.begin().await.map_err(map_sqlx_error)?; @@ -4539,6 +4793,7 @@ impl QueueStorage { // Clear the caches so the next ensure_lane / shard_for_enqueue / // claim calls re-observe DB state. self.clear_lane_cache(); + self.forget_schema_observations(); self.enqueue_shards_cache .lock() .expect("enqueue_shards_cache poisoned") @@ -4594,9 +4849,27 @@ impl QueueStorage { .await } - /// Provision every priority and shard of a queue before deploying a runtime - /// role without schema CREATE. Run as the schema owner/migrator. This does - /// not change enqueue routing or the configured shard count. + /// Lane rows for a caller that reserves through `advance_enqueue_head` + /// next. A cached lane needs no visibility probe here: the reserve reads + /// the head row itself and repairs a stale cache entry when it is missing. + async fn ensure_lane_for_reserve<'a>( + &self, + tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + queue: &str, + priority: i16, + enqueue_shard: i16, + ) -> Result<(), AwaError> { + if self.lane_is_cached(queue, priority, enqueue_shard) { + return Ok(()); + } + self.ensure_lane_inserts(tx, queue, priority, enqueue_shard) + .await + } + + /// Provision every priority and shard of a queue ahead of its first + /// enqueue. From migration v053 any runtime role may run it: the lanes are + /// created as the schema owner. This does not change enqueue routing or + /// the configured shard count. pub async fn prepare_queue( &self, pool: &PgPool, @@ -4637,11 +4910,35 @@ impl QueueStorage { ) -> Result<(), AwaError> { let schema = self.schema(); let lane_lock_key = format!("{schema}:{queue}:{priority}:{enqueue_shard}"); - sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") - .bind(lane_lock_key) - .execute(tx.as_mut()) + let provisioner_callable = if self.lane_provisioner_present.load(Ordering::Relaxed) { + sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") + .bind(&lane_lock_key) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + true + } else { + sqlx::query_scalar( + r#" + SELECT COALESCE(has_function_privilege(to_regprocedure($2), 'EXECUTE'), false) + FROM (SELECT pg_advisory_xact_lock(hashtextextended($1, 0))) AS lane_lock + "#, + ) + .bind(&lane_lock_key) + .bind(format!("{schema}.provision_lane(text,smallint,smallint)")) + .fetch_one(tx.as_mut()) .await - .map_err(map_sqlx_error)?; + .map_err(map_sqlx_error)? + }; + + if provisioner_callable + && self + .provision_lane(tx, queue, priority, enqueue_shard) + .await? + { + self.mark_lane_ensured(queue, priority, enqueue_shard); + return Ok(()); + } sqlx::query(audited_sql(format!( r#" @@ -4712,6 +5009,63 @@ impl QueueStorage { Ok(()) } + /// Create the lane through `{schema}.provision_lane`, which runs as the + /// schema owner, inside a savepoint. Returns false, with the transaction + /// as it was, when the function has gone (`42883`) or the schema owner + /// may not write the lane (`42501`); the caller then writes the rows + /// itself, which is what a binary that predates the function does. + async fn provision_lane<'a>( + &self, + tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + queue: &str, + priority: i16, + enqueue_shard: i16, + ) -> Result { + let schema = self.schema(); + sqlx::query("SAVEPOINT awa_provision_lane") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let provisioned = sqlx::query(audited_sql(format!( + "SELECT {schema}.provision_lane($1, $2, $3)" + ))) + .bind(queue) + .bind(priority) + .bind(enqueue_shard) + .execute(tx.as_mut()) + .await; + match provisioned { + Ok(_) => { + sqlx::query("RELEASE SAVEPOINT awa_provision_lane") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + self.lane_provisioner_present.store(true, Ordering::Relaxed); + Ok(true) + } + Err(sqlx::Error::Database(db_err)) + if matches!(db_err.code().as_deref(), Some("42883" | "42501")) => + { + sqlx::query("ROLLBACK TO SAVEPOINT awa_provision_lane") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + self.lane_provisioner_present + .store(false, Ordering::Relaxed); + if !self.lane_provisioner_warned.swap(true, Ordering::Relaxed) { + tracing::warn!( + schema, + code = db_err.code().as_deref().unwrap_or_default(), + error = %db_err.message(), + "Lane provisioning function failed; creating lane rows as the connected role" + ); + } + Ok(false) + } + Err(err) => Err(map_sqlx_error(err)), + } + } + /// Pick the enqueue shard for a row on `queue`. /// /// Reads `awa.queue_meta.enqueue_shards` (default 1, cached @@ -4803,6 +5157,139 @@ impl QueueStorage { .remove(&(queue.to_string(), priority, enqueue_shard)); } + fn forget_schema_observations(&self) { + *self + .enqueue_function + .lock() + .expect("enqueue function observation mutex") = None; + } + + async fn enqueue_function_available( + &self, + executor: impl sqlx::PgExecutor<'_>, + ) -> Result { + if self.enqueue_function_disabled { + return Ok(false); + } + let observed = *self + .enqueue_function + .lock() + .expect("enqueue function observation mutex"); + match observed { + Some((true, _)) => return Ok(true), + Some((false, at)) if at.elapsed() < ENQUEUE_FUNCTION_REPROBE => return Ok(false), + _ => {} + } + // The function calls ring_cursor, whose EXECUTE is revoked from PUBLIC. + let available: bool = sqlx::query_scalar( + "SELECT COALESCE(has_function_privilege(to_regprocedure($1), 'EXECUTE'), FALSE) \ + AND COALESCE(has_function_privilege(to_regprocedure($2), 'EXECUTE'), FALSE)", + ) + .bind(format!( + "{}.enqueue_ready_rows({ENQUEUE_READY_ROWS_ARGUMENT_TYPES})", + self.schema() + )) + .bind(format!("{}.ring_cursor(text)", self.schema())) + .fetch_one(executor) + .await + .map_err(map_sqlx_error)?; + *self + .enqueue_function + .lock() + .expect("enqueue function observation mutex") = Some((available, Instant::now())); + Ok(available) + } + + async fn execute_enqueue_ready_rows( + &self, + connection: &mut sqlx::PgConnection, + call: &EnqueueReadyRowsCall<'_>, + notify_channels: &[String], + ) -> Result { + let schema = self.schema(); + let inserted: i64 = sqlx::query_scalar(audited_sql(format!( + r#" + SELECT {schema}.enqueue_ready_rows( + $1::text[], $2::smallint[], $3::smallint[], $4::bigint[], + $5::bigint[], $6::text[], $7::jsonb[], $8::smallint[], $9::bigint[], + $10::smallint[], $11::timestamptz[], $12::timestamptz[], $13::timestamptz[], + $14::bytea[], $15::text[], $16::jsonb[], + $17::integer[], $18::bigint[], $19::bigint[], $20::timestamptz[], + $21::text[] + ) + "# + ))) + .bind(&call.lane_queues) + .bind(&call.lane_priorities) + .bind(&call.lane_shards) + .bind(&call.lane_counts) + .bind(&call.job_ids) + .bind(&call.kinds) + .bind(&call.args) + .bind(&call.attempts) + .bind(&call.run_leases) + .bind(&call.max_attempts) + .bind(&call.run_ats) + .bind(&call.attempted_ats) + .bind(&call.created_ats) + .bind(&call.unique_keys) + .bind(&call.unique_states) + .bind(&call.payloads) + .bind(&call.segment_lanes) + .bind(&call.segment_first_offsets) + .bind(&call.segment_next_offsets) + .bind(&call.segment_run_ats) + .bind(notify_channels) + .fetch_one(connection) + .await + .map_err(|err| { + // The probe can have run in a transaction that created the + // function and rolled back, or privileges can since have been + // revoked; forget it so the next enqueue probes again. + if matches!( + err.as_database_error().and_then(|db| db.code()).as_deref(), + Some("42883" | "42501") + ) { + self.forget_schema_observations(); + } + map_sqlx_error(err) + })?; + Ok(inserted as usize) + } + + /// Ready-only enqueue with no uniqueness claims in one autocommit + /// statement: lane reservation, ring cursor, ready rows, segments and + /// NOTIFY all run inside `enqueue_ready_rows`. + async fn enqueue_ready_rows_autocommit( + &self, + pool: &PgPool, + rows: Vec, + ) -> Result { + let notify_channels = self.notify_channels(rows.iter().map(|row| row.queue.clone())); + let mut connection = pool.acquire().await.map_err(map_sqlx_error)?; + let grouped = self + .group_ready_rows_by_shard(&mut connection, rows) + .await?; + let mut call = EnqueueReadyRowsCall::default(); + for ((queue, priority, enqueue_shard), lane_rows) in &grouped { + call.push_lane( + queue, + *priority, + *enqueue_shard, + lane_rows.iter().map(EnqueueReadyRow::from_ready_row), + ); + } + self.execute_enqueue_ready_rows(&mut connection, &call, ¬ify_channels) + .await + } + + fn ready_rows_take_unique_claims(rows: &[RuntimeReadyRow]) -> bool { + rows.iter().any(|row| { + row.unique_key.is_some() + && unique_state_claims(row.unique_states.as_deref(), JobState::Available) + }) + } + fn clear_lane_cache(&self) { self.ensured_lanes .lock() @@ -5109,16 +5596,41 @@ impl QueueStorage { } } - /// Current queue-ring cursor for the Rust-side enqueue paths, honouring - /// the schema's authority (#371 staged upgrade). Reads the authority row - /// (single cached page) then the O(1) cursor; callers that already hold - /// an authority for the transaction should use `ring_cursor_tx` directly. + /// Current queue-ring cursor for the Rust-side enqueue paths, under the + /// authority in force for this statement (#371 staged upgrade), read in + /// one round trip. Callers that already hold an authority for the + /// transaction should use `ring_cursor_tx` directly. async fn current_queue_ring<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, ) -> Result<(i32, i64), AwaError> { - let authority = self.ring_authority_tx(tx).await?; - self.ring_cursor_tx(tx, RingFamily::Queue, authority).await + let schema = self.schema(); + sqlx::query_as(audited_sql(format!( + r#" + WITH authority AS ( + SELECT authority = 'ledger' AS ledger + FROM {schema}.ring_cursor_authority + WHERE singleton = TRUE + ) + ( + SELECT ledger.slot, ledger.generation + FROM {schema}.queue_ring_rotations AS ledger + WHERE (SELECT authority.ledger FROM authority) + ORDER BY ledger.generation DESC + LIMIT 1 + ) + UNION ALL + ( + SELECT state.current_slot, state.generation + FROM {schema}.queue_ring_state AS state + WHERE state.singleton = TRUE + AND NOT (SELECT authority.ledger FROM authority) + ) + "# + ))) + .fetch_one(tx.as_mut()) + .await + .map_err(map_sqlx_error) } async fn next_job_ids<'a>( @@ -5602,12 +6114,13 @@ impl QueueStorage { tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: Vec, unique_claims_synced: bool, + allow_enqueue_function: bool, ) -> Result { if rows.is_empty() { return Ok(0); } - let grouped = self.group_ready_rows_by_shard(tx, rows).await?; + let grouped = self.group_ready_rows_by_shard(tx.as_mut(), rows).await?; let total_rows: usize = grouped.values().map(Vec::len).sum(); let missing_job_ids = grouped .values() @@ -5664,8 +6177,13 @@ impl QueueStorage { self.sync_ready_enqueue_unique_claims(tx, &ready_rows) .await?; } + if allow_enqueue_function && self.enqueue_function_available(tx.as_mut()).await? { + return self + .enqueue_lane_ranges_tx(tx, &lane_ranges, &ready_rows) + .await; + } for (queue, priority, enqueue_shard, range_start, range_end) in lane_ranges { - self.ensure_lane(tx, &queue, priority, enqueue_shard) + self.ensure_lane_for_reserve(tx, &queue, priority, enqueue_shard) .await?; let count = (range_end - range_start) as i64; @@ -5683,6 +6201,27 @@ impl QueueStorage { Ok(total_rows) } + async fn enqueue_lane_ranges_tx<'a>( + &self, + tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + lane_ranges: &[(String, i16, i16, usize, usize)], + ready_rows: &[RuntimeReadyInsert], + ) -> Result { + let mut call = EnqueueReadyRowsCall::default(); + for (queue, priority, enqueue_shard, range_start, range_end) in lane_ranges { + call.push_lane( + queue, + *priority, + *enqueue_shard, + ready_rows[*range_start..*range_end] + .iter() + .map(EnqueueReadyRow::from_insert), + ); + } + self.execute_enqueue_ready_rows(tx.as_mut(), &call, &[]) + .await + } + /// Re-group rows by `(queue, priority, enqueue_shard)` so each /// resulting bucket targets exactly one head row. /// @@ -5708,9 +6247,9 @@ impl QueueStorage { /// order callers reserve lanes in. It is the canonical lane-lock order /// that `insert_many` and `insert_many_copy` also follow, so batches /// from either path cannot deadlock on each other's lane locks. - async fn group_ready_rows_by_shard<'a>( + async fn group_ready_rows_by_shard( &self, - tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + connection: &mut sqlx::PgConnection, rows: Vec, ) -> Result>, AwaError> { // Partition by (queue, priority) first so the rotor pick for @@ -5730,7 +6269,7 @@ impl QueueStorage { for row in bucket { if row.ordering_key.is_some() { let shard = self - .shard_for_enqueue(tx.as_mut(), &queue, row.ordering_key.as_deref()) + .shard_for_enqueue(&mut *connection, &queue, row.ordering_key.as_deref()) .await?; grouped .entry((queue.clone(), priority, shard)) @@ -5741,7 +6280,9 @@ impl QueueStorage { } } if !rotor_rows.is_empty() { - let shard = self.shard_for_enqueue(tx.as_mut(), &queue, None).await?; + let shard = self + .shard_for_enqueue(&mut *connection, &queue, None) + .await?; grouped .entry((queue.clone(), priority, shard)) .or_default() @@ -5756,12 +6297,13 @@ impl QueueStorage { tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, rows: Vec, unique_claims_synced: bool, + allow_enqueue_function: bool, ) -> Result { if rows.is_empty() { return Ok(0); } - let grouped = self.group_ready_rows_by_shard(tx, rows).await?; + let grouped = self.group_ready_rows_by_shard(tx.as_mut(), rows).await?; let total_rows: usize = grouped.values().map(Vec::len).sum(); let missing_job_ids = grouped @@ -5819,8 +6361,16 @@ impl QueueStorage { self.sync_ready_enqueue_unique_claims(tx, &ready_rows) .await?; } + if allow_enqueue_function + && ready_rows.len() < COPY_MIN_ROWS + && self.enqueue_function_available(tx.as_mut()).await? + { + return self + .enqueue_lane_ranges_tx(tx, &lane_ranges, &ready_rows) + .await; + } for (queue, priority, enqueue_shard, range_start, range_end) in lane_ranges { - self.ensure_lane(tx, &queue, priority, enqueue_shard) + self.ensure_lane_for_reserve(tx, &queue, priority, enqueue_shard) .await?; let count = (range_end - range_start) as i64; @@ -5887,7 +6437,7 @@ impl QueueStorage { // retries must re-enqueue with the original key from the // application layer. let enqueue_shard = self.shard_for_enqueue(tx.as_mut(), &queue, None).await?; - self.ensure_lane(tx, &queue, priority, enqueue_shard) + self.ensure_lane_for_reserve(tx, &queue, priority, enqueue_shard) .await?; let count = lane_rows.len() as i64; @@ -6780,15 +7330,36 @@ impl QueueStorage { &self, pool: &PgPool, rows: Vec, + ) -> Result { + match self.enqueue_runtime_rows_attempt(pool, &rows, true).await { + Err(error) if enqueue_function_unavailable(&error) => { + self.enqueue_runtime_rows_attempt(pool, &rows, false).await + } + result => result, + } + } + + async fn enqueue_runtime_rows_attempt( + &self, + pool: &PgPool, + rows: &[RuntimeReadyRow], + allow_enqueue_function: bool, ) -> Result { if rows.is_empty() { return Ok(0); } let mut tx = pool.begin().await.map_err(map_sqlx_error)?; - let total_rows = self - .insert_ready_rows_tx(&mut tx, rows.clone(), false) - .await?; + let inserted = self + .insert_ready_rows_tx(&mut tx, rows.to_vec(), false, allow_enqueue_function) + .await; + let total_rows = match inserted { + Ok(count) => count, + Err(error) => { + tx.rollback().await.map_err(map_sqlx_error)?; + return Err(error); + } + }; let queues_to_notify: Vec = rows.iter().map(|row| row.queue.clone()).collect(); self.notify_queues_tx(&mut tx, queues_to_notify).await?; @@ -6843,6 +7414,20 @@ impl QueueStorage { &self, pool: &PgPool, jobs: &[InsertParams], + ) -> Result { + match self.enqueue_params_batch_attempt(pool, jobs, true).await { + Err(error) if enqueue_function_unavailable(&error) => { + self.enqueue_params_batch_attempt(pool, jobs, false).await + } + result => result, + } + } + + async fn enqueue_params_batch_attempt( + &self, + pool: &PgPool, + jobs: &[InsertParams], + allow_enqueue_function: bool, ) -> Result { if jobs.is_empty() { return Ok(0); @@ -6904,6 +7489,14 @@ impl QueueStorage { } } + if allow_enqueue_function + && deferred_rows.is_empty() + && !Self::ready_rows_take_unique_claims(&ready_rows) + && self.enqueue_function_available(pool).await? + { + return self.enqueue_ready_rows_autocommit(pool, ready_rows).await; + } + let queues_to_notify: Vec = ready_rows.iter().map(|row| row.queue.clone()).collect(); @@ -6926,11 +7519,18 @@ impl QueueStorage { let mut total = 0usize; if !ready_rows.is_empty() { - total += self - .insert_ready_rows_tx(&mut tx, ready_rows.clone(), true) - .await?; - } - if !deferred_rows.is_empty() { + let inserted = self + .insert_ready_rows_tx(&mut tx, ready_rows.clone(), true, allow_enqueue_function) + .await; + total += match inserted { + Ok(count) => count, + Err(error) => { + tx.rollback().await.map_err(map_sqlx_error)?; + return Err(error); + } + }; + } + if !deferred_rows.is_empty() { total += self .execute_deferred_inserts_tx(&mut tx, &deferred_rows) .await?; @@ -6953,6 +7553,20 @@ impl QueueStorage { &self, pool: &PgPool, jobs: &[InsertParams], + ) -> Result { + match self.enqueue_params_copy_attempt(pool, jobs, true).await { + Err(error) if enqueue_function_unavailable(&error) => { + self.enqueue_params_copy_attempt(pool, jobs, false).await + } + result => result, + } + } + + async fn enqueue_params_copy_attempt( + &self, + pool: &PgPool, + jobs: &[InsertParams], + allow_enqueue_function: bool, ) -> Result { if jobs.is_empty() { return Ok(0); @@ -7014,6 +7628,15 @@ impl QueueStorage { } } + if allow_enqueue_function + && deferred_rows.is_empty() + && ready_rows.len() < COPY_MIN_ROWS + && !Self::ready_rows_take_unique_claims(&ready_rows) + && self.enqueue_function_available(pool).await? + { + return self.enqueue_ready_rows_autocommit(pool, ready_rows).await; + } + let queues_to_notify: Vec = ready_rows.iter().map(|row| row.queue.clone()).collect(); @@ -7036,9 +7659,16 @@ impl QueueStorage { .await?; if !ready_rows.is_empty() { - total += self - .insert_ready_rows_copy_tx(&mut tx, ready_rows, true) - .await?; + let inserted = self + .insert_ready_rows_copy_tx(&mut tx, ready_rows, true, allow_enqueue_function) + .await; + total += match inserted { + Ok(count) => count, + Err(error) => { + tx.rollback().await.map_err(map_sqlx_error)?; + return Err(error); + } + }; } if !deferred_rows.is_empty() { total += self @@ -11666,20 +12296,25 @@ impl QueueStorage { .map_err(map_sqlx_error) } + /// Logical-queue NOTIFY channels, deduplicated so a transaction never + /// emits the same NOTIFY twice. + fn notify_channels(&self, queues: impl IntoIterator) -> Vec { + queues + .into_iter() + .map(|queue| format!("awa:{}", self.logical_queue_name(&queue))) + .collect::>() + .into_iter() + .collect() + } + async fn notify_queues_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, queues: impl IntoIterator, ) -> Result<(), AwaError> { - // BTreeSet dedups so we never emit the same NOTIFY twice per - // transaction. Multi-queue enqueues fold into a single round-trip via + // Multi-queue enqueues fold into a single round-trip via // `unnest($1::text[])` rather than one round-trip per channel. - let channels: Vec = queues - .into_iter() - .map(|queue| format!("awa:{}", self.logical_queue_name(&queue))) - .collect::>() - .into_iter() - .collect(); + let channels = self.notify_channels(queues); if channels.is_empty() { return Ok(()); } @@ -18163,7 +18798,10 @@ impl QueueStorage { break; } TerminalDeltaSlotRollup::Blocked => { + // The next slot would most likely wait out the same lock + // timeout, and the leader loop is blocked while it does. outcome.blocked_slots += 1; + break; } } } @@ -18350,237 +18988,215 @@ impl QueueStorage { return Ok(TerminalDeltaSlotRollup::SkippedActive); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - - let lock_delta = sqlx::query(audited_sql(format!( - "LOCK TABLE {delta_child} IN ACCESS EXCLUSIVE MODE" - ))) - .execute(tx.as_mut()) - .await; - - match lock_delta { - Ok(_) => {} - Err(err) if is_lock_contention_error(&err) => { - let _ = tx.rollback().await; - return Ok(TerminalDeltaSlotRollup::Blocked); - } - Err(err) => { - let _ = tx.rollback().await; - return Err(map_sqlx_error(err)); - } - } - - let active_leases = - queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?; + let locked: Result = async { + set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - if active_leases { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(TerminalDeltaSlotRollup::SkippedActive); - } + sqlx::query(audited_sql(format!( + "LOCK TABLE {delta_child} IN ACCESS EXCLUSIVE MODE" + ))) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; - let unclosed_claim_refs = - queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?; + let active_leases = + queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?; - if unclosed_claim_refs { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(TerminalDeltaSlotRollup::SkippedActive); - } + if active_leases { + return Ok(TerminalDeltaSlotRollup::SkippedActive); + } - let delta_rows: i64 = sqlx::query_scalar(audited_sql(format!( - r#" - SELECT count(*)::bigint - FROM {delta_child} - WHERE ready_generation = $1 - "# - ))) - .bind(generation) - .fetch_one(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; + let unclosed_claim_refs = + queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?; - if delta_rows == 0 { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(TerminalDeltaSlotRollup::Empty); - } + if unclosed_claim_refs { + return Ok(TerminalDeltaSlotRollup::SkippedActive); + } - // In ledger authority each touched lane collapses to one bucket-0 row - // holding its old total (summed over every bucket) plus the signed - // delta, which keeps a negative delta exact when the matching positive - // count sits in another bucket. Before the flip, N-1 maintenance can - // still win the election and folds per job-id bucket, so the fold - // keeps that shape until the flip fences it (ADR-041). - let grouped_keys: i64 = match authority { - RingAuthority::Ledger => sqlx::query_scalar(audited_sql(format!( - r#" - WITH grouped AS MATERIALIZED ( - SELECT - ready_slot, - queue, - priority, - enqueue_shard, - SUM(terminal_delta)::bigint AS delta - FROM {delta_child} - WHERE ready_generation = $1 - GROUP BY ready_slot, queue, priority, enqueue_shard - HAVING SUM(terminal_delta) <> 0 - ), - collapsed AS ( - DELETE FROM {schema}.queue_terminal_live_counts AS counts - USING grouped - WHERE counts.ready_slot = grouped.ready_slot - AND counts.queue = grouped.queue - AND counts.priority = grouped.priority - AND counts.enqueue_shard = grouped.enqueue_shard - RETURNING - counts.ready_slot, - counts.queue, - counts.priority, - counts.enqueue_shard, - counts.live_terminal_count - ), - folded AS ( - SELECT - grouped.ready_slot, - grouped.queue, - grouped.priority, - grouped.enqueue_shard, - (grouped.delta + COALESCE(SUM(collapsed.live_terminal_count), 0))::bigint - AS live_terminal_count - FROM grouped - LEFT JOIN collapsed - ON collapsed.ready_slot = grouped.ready_slot - AND collapsed.queue = grouped.queue - AND collapsed.priority = grouped.priority - AND collapsed.enqueue_shard = grouped.enqueue_shard - GROUP BY - grouped.ready_slot, - grouped.queue, - grouped.priority, - grouped.enqueue_shard, - grouped.delta - ), - inserted AS ( - INSERT INTO {schema}.queue_terminal_live_counts ( - ready_slot, - queue, - priority, - enqueue_shard, - counter_bucket, - live_terminal_count - ) - SELECT - ready_slot, - queue, - priority, - enqueue_shard, - 0::smallint, - live_terminal_count - FROM folded - WHERE live_terminal_count <> 0 - ORDER BY ready_slot, queue, priority, enqueue_shard - RETURNING 1 - ) - SELECT count(*)::bigint FROM grouped - "# - ))) - .bind(generation) - .fetch_one(tx.as_mut()) - .await - .map_err(map_sqlx_error)?, - RingAuthority::Columns => sqlx::query_scalar(audited_sql(format!( + let delta_rows: i64 = sqlx::query_scalar(audited_sql(format!( r#" - WITH grouped AS MATERIALIZED ( - SELECT - ready_slot, - queue, - priority, - enqueue_shard, - counter_bucket, - SUM(terminal_delta)::bigint AS delta + SELECT count(*)::bigint FROM {delta_child} WHERE ready_generation = $1 - GROUP BY ready_slot, queue, priority, enqueue_shard, counter_bucket - HAVING SUM(terminal_delta) <> 0 - ), - updated AS ( - UPDATE {schema}.queue_terminal_live_counts AS counts - SET live_terminal_count = GREATEST(0, counts.live_terminal_count + grouped.delta) - FROM grouped - WHERE counts.ready_slot = grouped.ready_slot - AND counts.queue = grouped.queue - AND counts.priority = grouped.priority - AND counts.enqueue_shard = grouped.enqueue_shard - AND counts.counter_bucket = grouped.counter_bucket - RETURNING - counts.ready_slot, - counts.queue, - counts.priority, - counts.enqueue_shard, - counts.counter_bucket - ), - inserted AS ( - INSERT INTO {schema}.queue_terminal_live_counts AS counts ( - ready_slot, - queue, - priority, - enqueue_shard, - counter_bucket, - live_terminal_count - ) - SELECT - grouped.ready_slot, - grouped.queue, - grouped.priority, - grouped.enqueue_shard, - grouped.counter_bucket, - grouped.delta - FROM grouped - WHERE grouped.delta > 0 - AND NOT EXISTS ( - SELECT 1 - FROM updated - WHERE updated.ready_slot = grouped.ready_slot - AND updated.queue = grouped.queue - AND updated.priority = grouped.priority - AND updated.enqueue_shard = grouped.enqueue_shard - AND updated.counter_bucket = grouped.counter_bucket - ) - ORDER BY ready_slot, queue, priority, enqueue_shard, counter_bucket - ON CONFLICT (ready_slot, queue, priority, enqueue_shard, counter_bucket) DO UPDATE - SET live_terminal_count = - GREATEST(0, counts.live_terminal_count + EXCLUDED.live_terminal_count) - RETURNING 1 - ) - SELECT count(*)::bigint FROM grouped - "# + "# ))) .bind(generation) .fetch_one(tx.as_mut()) .await - .map_err(map_sqlx_error)?, - }; - - let truncate_delta = sqlx::query(audited_sql(format!("TRUNCATE TABLE {delta_child}"))) - .execute(tx.as_mut()) - .await; + .map_err(map_sqlx_error)?; - match truncate_delta { - Ok(_) => { - tx.commit().await.map_err(map_sqlx_error)?; - Ok(TerminalDeltaSlotRollup::Rolled { - delta_rows, - grouped_keys, - }) - } - Err(err) if is_lock_contention_error(&err) => { - let _ = tx.rollback().await; - Ok(TerminalDeltaSlotRollup::Blocked) - } - Err(err) => { - let _ = tx.rollback().await; - Err(map_sqlx_error(err)) + if delta_rows == 0 { + return Ok(TerminalDeltaSlotRollup::Empty); } + + // In ledger authority each touched lane collapses to one bucket-0 row + // holding its old total (summed over every bucket) plus the signed + // delta, which keeps a negative delta exact when the matching positive + // count sits in another bucket. Before the flip, N-1 maintenance can + // still win the election and folds per job-id bucket, so the fold + // keeps that shape until the flip fences it (ADR-041). + let grouped_keys: i64 = match authority { + RingAuthority::Ledger => sqlx::query_scalar(audited_sql(format!( + r#" + WITH grouped AS MATERIALIZED ( + SELECT + ready_slot, + queue, + priority, + enqueue_shard, + SUM(terminal_delta)::bigint AS delta + FROM {delta_child} + WHERE ready_generation = $1 + GROUP BY ready_slot, queue, priority, enqueue_shard + HAVING SUM(terminal_delta) <> 0 + ), + collapsed AS ( + DELETE FROM {schema}.queue_terminal_live_counts AS counts + USING grouped + WHERE counts.ready_slot = grouped.ready_slot + AND counts.queue = grouped.queue + AND counts.priority = grouped.priority + AND counts.enqueue_shard = grouped.enqueue_shard + RETURNING + counts.ready_slot, + counts.queue, + counts.priority, + counts.enqueue_shard, + counts.live_terminal_count + ), + folded AS ( + SELECT + grouped.ready_slot, + grouped.queue, + grouped.priority, + grouped.enqueue_shard, + (grouped.delta + COALESCE(SUM(collapsed.live_terminal_count), 0))::bigint + AS live_terminal_count + FROM grouped + LEFT JOIN collapsed + ON collapsed.ready_slot = grouped.ready_slot + AND collapsed.queue = grouped.queue + AND collapsed.priority = grouped.priority + AND collapsed.enqueue_shard = grouped.enqueue_shard + GROUP BY + grouped.ready_slot, + grouped.queue, + grouped.priority, + grouped.enqueue_shard, + grouped.delta + ), + inserted AS ( + INSERT INTO {schema}.queue_terminal_live_counts ( + ready_slot, + queue, + priority, + enqueue_shard, + counter_bucket, + live_terminal_count + ) + SELECT + ready_slot, + queue, + priority, + enqueue_shard, + 0::smallint, + live_terminal_count + FROM folded + WHERE live_terminal_count <> 0 + ORDER BY ready_slot, queue, priority, enqueue_shard + RETURNING 1 + ) + SELECT count(*)::bigint FROM grouped + "# + ))) + .bind(generation) + .fetch_one(tx.as_mut()) + .await + .map_err(map_sqlx_error)?, + RingAuthority::Columns => sqlx::query_scalar(audited_sql(format!( + r#" + WITH grouped AS MATERIALIZED ( + SELECT + ready_slot, + queue, + priority, + enqueue_shard, + counter_bucket, + SUM(terminal_delta)::bigint AS delta + FROM {delta_child} + WHERE ready_generation = $1 + GROUP BY ready_slot, queue, priority, enqueue_shard, counter_bucket + HAVING SUM(terminal_delta) <> 0 + ), + updated AS ( + UPDATE {schema}.queue_terminal_live_counts AS counts + SET live_terminal_count = GREATEST(0, counts.live_terminal_count + grouped.delta) + FROM grouped + WHERE counts.ready_slot = grouped.ready_slot + AND counts.queue = grouped.queue + AND counts.priority = grouped.priority + AND counts.enqueue_shard = grouped.enqueue_shard + AND counts.counter_bucket = grouped.counter_bucket + RETURNING + counts.ready_slot, + counts.queue, + counts.priority, + counts.enqueue_shard, + counts.counter_bucket + ), + inserted AS ( + INSERT INTO {schema}.queue_terminal_live_counts AS counts ( + ready_slot, + queue, + priority, + enqueue_shard, + counter_bucket, + live_terminal_count + ) + SELECT + grouped.ready_slot, + grouped.queue, + grouped.priority, + grouped.enqueue_shard, + grouped.counter_bucket, + grouped.delta + FROM grouped + WHERE grouped.delta > 0 + AND NOT EXISTS ( + SELECT 1 + FROM updated + WHERE updated.ready_slot = grouped.ready_slot + AND updated.queue = grouped.queue + AND updated.priority = grouped.priority + AND updated.enqueue_shard = grouped.enqueue_shard + AND updated.counter_bucket = grouped.counter_bucket + ) + ORDER BY ready_slot, queue, priority, enqueue_shard, counter_bucket + ON CONFLICT (ready_slot, queue, priority, enqueue_shard, counter_bucket) DO UPDATE + SET live_terminal_count = + GREATEST(0, counts.live_terminal_count + EXCLUDED.live_terminal_count) + RETURNING 1 + ) + SELECT count(*)::bigint FROM grouped + "# + ))) + .bind(generation) + .fetch_one(tx.as_mut()) + .await + .map_err(map_sqlx_error)?, + }; + + sqlx::query(audited_sql(format!("TRUNCATE TABLE {delta_child}"))) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + Ok(TerminalDeltaSlotRollup::Rolled { + delta_rows, + grouped_keys, + }) } + .await; + + settle_prune_tx(tx, locked, TerminalDeltaSlotRollup::Blocked).await } /// Fold pending `queue_terminal_rollup_deltas` rows into the @@ -18830,172 +19446,119 @@ impl QueueStorage { }); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + let locked: Result = async { + set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - let lock_started = Instant::now(); - // Claims read segments before ready rows, while producers write in - // the opposite order. Queue behind segment readers, then abandon - // the transaction rather than wait with an exclusive child held. - let lock_tables = async { - sqlx::query(audited_sql(format!( - "LOCK TABLE {segment_child} IN ACCESS EXCLUSIVE MODE" - ))) - .execute(tx.as_mut()) - .await?; - sqlx::query(audited_sql(format!( - "LOCK TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child} IN ACCESS EXCLUSIVE MODE NOWAIT" - ))) - .execute(tx.as_mut()) + let lock_started = Instant::now(); + // Claims read segments before ready rows, while producers write in + // the opposite order. Queue behind segment readers, then abandon + // the transaction rather than wait with an exclusive child held. + async { + sqlx::query(audited_sql(format!( + "LOCK TABLE {segment_child} IN ACCESS EXCLUSIVE MODE" + ))) + .execute(tx.as_mut()) + .await?; + sqlx::query(audited_sql(format!( + "LOCK TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child} IN ACCESS EXCLUSIVE MODE NOWAIT" + ))) + .execute(tx.as_mut()) + .await + } .await - } - .await; - let lock_duration = lock_started.elapsed(); + .map_err(map_sqlx_error)?; + let lock_duration = lock_started.elapsed(); + + // The pre-lock gates are only an MVCC snapshot. If an in-flight + // producer or claimer committed while the exclusive partition + // lock waited, these post-lock gates see it before TRUNCATE. + let active_leases = + queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?; + if active_leases { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::QueueActiveLeases, + count: busy_indicator(active_leases), + }); + } - if let Err(err) = lock_tables { - let _ = tx.rollback().await; - if is_lock_contention_error(&err) { - return Ok(PruneOutcome::Blocked { slot }); + let pending = queue_prune_has_pending_ready_tx(&mut tx, schema, slot, generation).await?; + if pending { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::QueuePendingReady, + count: busy_indicator(pending), + }); } - return Err(map_sqlx_error(err)); - } - // The pre-lock gates are only an MVCC snapshot. If an in-flight - // producer or claimer committed while the exclusive partition - // lock waited, these post-lock gates see it before TRUNCATE. - let active_leases = - queue_prune_has_active_leases_tx(&mut tx, schema, slot, generation).await?; - if active_leases { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::QueueActiveLeases, - count: busy_indicator(active_leases), - }); - } + let unclosed_claim_refs = + queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?; + if unclosed_claim_refs { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::QueueUnclosedClaimRefs, + count: busy_indicator(unclosed_claim_refs), + }); + } - let pending = queue_prune_has_pending_ready_tx(&mut tx, schema, slot, generation).await?; - if pending { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::QueuePendingReady, - count: busy_indicator(pending), - }); - } - - let unclosed_claim_refs = - queue_prune_has_unclosed_claim_refs_tx(&mut tx, schema, slot, generation).await?; - if unclosed_claim_refs { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::QueueUnclosedClaimRefs, - count: busy_indicator(unclosed_claim_refs), - }); - } - - let failed_retention_secs = i64::try_from(failed_retention.as_secs()).unwrap_or(i64::MAX); - let retention_floor = failed_retention_secs > 0; - - // #337: failed terminal rows still inside the retention floor - // are re-homed into the live done segment before the TRUNCATE - // so they stay retryable. The ring rotation lock is held above, - // which serializes against rotate — the current slot (and its - // generation, read from the ledger cursor) cannot move while - // carried rows are re-inserted into it. - let carried_failed_rows = if retention_floor { - // Carried rows are widened to the self-contained synthetic - // done-row shape (the COALESCE chain mirrors - // `done_row_projection`) because the ready rows they would - // otherwise hydrate from are truncated with this slot. - // Deliberately no ON CONFLICT: lane_seq comes from - // never-resetting sequences, so a PK collision means - // corrupted terminal state and must abort the prune loudly - // rather than silently drop a terminal fact. - let carried = sqlx::query(audited_sql(format!( - r#" - INSERT INTO {schema}.done_entries ( - ready_slot, ready_generation, job_id, kind, queue, - args, state, priority, attempt, run_lease, - max_attempts, lane_seq, enqueue_shard, run_at, - attempted_at, finalized_at, created_at, unique_key, - unique_states, payload - ) - SELECT - $2, - $3, - done.job_id, - done.kind, - done.queue, - COALESCE(done.args, ready.args, '{{}}'::jsonb), - done.state, - done.priority, - done.attempt, - done.run_lease, - COALESCE(done.max_attempts, ready.max_attempts, 25::smallint), - done.lane_seq, - done.enqueue_shard, - COALESCE(done.run_at, ready.run_at, done.finalized_at), - COALESCE(done.attempted_at, ready.attempted_at), - done.finalized_at, - COALESCE(done.created_at, ready.created_at, done.finalized_at), - COALESCE(done.unique_key, ready.unique_key), - COALESCE(done.unique_states, ready.unique_states), - COALESCE(done.payload, ready.payload, '{{}}'::jsonb) - FROM {done_child} AS done - LEFT JOIN {ready_child} AS ready - ON ready.ready_slot = done.ready_slot - AND ready.ready_generation = done.ready_generation - AND ready.queue = done.queue - AND ready.priority = done.priority - AND ready.enqueue_shard = done.enqueue_shard - AND ready.lane_seq = done.lane_seq - WHERE done.state = 'failed' - AND done.finalized_at >= now() - make_interval(secs => $1::bigint) - "# - ))) - .bind(failed_retention_secs) - .bind(state.0) - .bind(current_generation) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)? - .rows_affected(); - - if carried > 0 { - // This slot's live counter rows and delta segment are - // wiped below, so the carried rows' exact-count - // evidence moves with them: re-append positive deltas - // under the current slot, grouped exactly like the - // completion path's delta append. - sqlx::query(audited_sql(format!( + let failed_retention_secs = i64::try_from(failed_retention.as_secs()).unwrap_or(i64::MAX); + let retention_floor = failed_retention_secs > 0; + + // #337: failed terminal rows still inside the retention floor + // are re-homed into the live done segment before the TRUNCATE + // so they stay retryable. The ring rotation lock is held above, + // which serializes against rotate — the current slot (and its + // generation, read from the ledger cursor) cannot move while + // carried rows are re-inserted into it. + let carried_failed_rows = if retention_floor { + // Carried rows are widened to the self-contained synthetic + // done-row shape (the COALESCE chain mirrors + // `done_row_projection`) because the ready rows they would + // otherwise hydrate from are truncated with this slot. + // Deliberately no ON CONFLICT: lane_seq comes from + // never-resetting sequences, so a PK collision means + // corrupted terminal state and must abort the prune loudly + // rather than silently drop a terminal fact. + let carried = sqlx::query(audited_sql(format!( r#" - INSERT INTO {schema}.queue_terminal_count_deltas ( - ready_slot, - ready_generation, - queue, - priority, - enqueue_shard, - counter_bucket, - terminal_delta + INSERT INTO {schema}.done_entries ( + ready_slot, ready_generation, job_id, kind, queue, + args, state, priority, attempt, run_lease, + max_attempts, lane_seq, enqueue_shard, run_at, + attempted_at, finalized_at, created_at, unique_key, + unique_states, payload ) SELECT $2, $3, - queue, - priority, - enqueue_shard, - mod( - mod(job_id, {TERMINAL_COUNTER_BUCKETS}::bigint) - + {TERMINAL_COUNTER_BUCKETS}::bigint, - {TERMINAL_COUNTER_BUCKETS}::bigint - )::smallint AS counter_bucket, - count(*)::bigint - FROM {done_child} - WHERE state = 'failed' - AND finalized_at >= now() - make_interval(secs => $1::bigint) - GROUP BY queue, priority, enqueue_shard, counter_bucket - ORDER BY queue, priority, enqueue_shard, counter_bucket + done.job_id, + done.kind, + done.queue, + COALESCE(done.args, ready.args, '{{}}'::jsonb), + done.state, + done.priority, + done.attempt, + done.run_lease, + COALESCE(done.max_attempts, ready.max_attempts, 25::smallint), + done.lane_seq, + done.enqueue_shard, + COALESCE(done.run_at, ready.run_at, done.finalized_at), + COALESCE(done.attempted_at, ready.attempted_at), + done.finalized_at, + COALESCE(done.created_at, ready.created_at, done.finalized_at), + COALESCE(done.unique_key, ready.unique_key), + COALESCE(done.unique_states, ready.unique_states), + COALESCE(done.payload, ready.payload, '{{}}'::jsonb) + FROM {done_child} AS done + LEFT JOIN {ready_child} AS ready + ON ready.ready_slot = done.ready_slot + AND ready.ready_generation = done.ready_generation + AND ready.queue = done.queue + AND ready.priority = done.priority + AND ready.enqueue_shard = done.enqueue_shard + AND ready.lane_seq = done.lane_seq + WHERE done.state = 'failed' + AND done.finalized_at >= now() - make_interval(secs => $1::bigint) "# ))) .bind(failed_retention_secs) @@ -19003,187 +19566,232 @@ impl QueueStorage { .bind(current_generation) .execute(tx.as_mut()) .await - .map_err(map_sqlx_error)?; - } - carried - } else { - 0 - }; - - // #290: scan the about-to-be-truncated partition for the rollup - // fold. The rollup column is *permanent* state, so we MUST fold - // from ground truth (the `{done_child}` partition itself), not - // from `queue_terminal_live_counts` — the counter can drift - // briefly during a rolling upgrade from a pre-counter binary - // and folding drift into the rollup would bake it in forever. - // The counter rows for this slot are still cleaned up after the - // truncate; reads from the counter (queue_counts_exact in #305) - // may transiently disagree with the rollup until the operator - // runs `awa storage rebuild-terminal-counters`, but the rollup - // itself stays authoritative. See PR #304 reviewer finding - // "High: A1 can persist stale counter state before the - // read-switch PR" for the trade-off. - // - // Carried failed rows are excluded from both columns: they are - // still live in `done_entries`, so folding them into the - // permanent rollup would double-count them. - let pruned_terminal_counts: Vec<(String, i16, i64, i64)> = - sqlx::query_as(audited_sql(format!( - r#" - WITH done_counts AS ( - SELECT - queue, - priority, - (count(*) FILTER (WHERE state <> 'failed'))::bigint AS completed, - (count(*) FILTER ( + .map_err(map_sqlx_error)? + .rows_affected(); + + if carried > 0 { + // This slot's live counter rows and delta segment are + // wiped below, so the carried rows' exact-count + // evidence moves with them: re-append positive deltas + // under the current slot, grouped exactly like the + // completion path's delta append. + sqlx::query(audited_sql(format!( + r#" + INSERT INTO {schema}.queue_terminal_count_deltas ( + ready_slot, + ready_generation, + queue, + priority, + enqueue_shard, + counter_bucket, + terminal_delta + ) + SELECT + $2, + $3, + queue, + priority, + enqueue_shard, + mod( + mod(job_id, {TERMINAL_COUNTER_BUCKETS}::bigint) + + {TERMINAL_COUNTER_BUCKETS}::bigint, + {TERMINAL_COUNTER_BUCKETS}::bigint + )::smallint AS counter_bucket, + count(*)::bigint + FROM {done_child} WHERE state = 'failed' - AND ( - $2::boolean IS FALSE - OR finalized_at < now() - make_interval(secs => $1::bigint) - ) - ))::bigint AS failed - FROM {done_child} - GROUP BY queue, priority - ), - -- Scan the entire partition being truncated (all generations), - -- mirroring done_counts' whole-child scan above. The rotate guard - -- keeps a partition to one generation at prune time, so this is the - -- same set today; counting the whole child keeps the rollup fold - -- conservative even if that invariant is ever weakened (a stray - -- generation's batches would otherwise be truncated without being - -- folded into the rollup -> terminal-count undercount). - batch_rows AS ( - SELECT - batch.ready_generation, - batch.queue, - batch.priority, - item.job_id, - item.run_lease - FROM {receipt_batch_child} AS batch - CROSS JOIN LATERAL unnest( - batch.job_ids, - batch.run_leases - ) AS item(job_id, run_lease) - ), - batch_counts AS ( + AND finalized_at >= now() - make_interval(secs => $1::bigint) + GROUP BY queue, priority, enqueue_shard, counter_bucket + ORDER BY queue, priority, enqueue_shard, counter_bucket + "# + ))) + .bind(failed_retention_secs) + .bind(state.0) + .bind(current_generation) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + } + carried + } else { + 0 + }; + + // #290: scan the about-to-be-truncated partition for the rollup + // fold. The rollup column is *permanent* state, so we MUST fold + // from ground truth (the `{done_child}` partition itself), not + // from `queue_terminal_live_counts` — the counter can drift + // briefly during a rolling upgrade from a pre-counter binary + // and folding drift into the rollup would bake it in forever. + // The counter rows for this slot are still cleaned up after the + // truncate; reads from the counter (queue_counts_exact in #305) + // may transiently disagree with the rollup until the operator + // runs `awa storage rebuild-terminal-counters`, but the rollup + // itself stays authoritative. See PR #304 reviewer finding + // "High: A1 can persist stale counter state before the + // read-switch PR" for the trade-off. + // + // Carried failed rows are excluded from both columns: they are + // still live in `done_entries`, so folding them into the + // permanent rollup would double-count them. + let pruned_terminal_counts: Vec<(String, i16, i64, i64)> = + sqlx::query_as(audited_sql(format!( + r#" + WITH done_counts AS ( + SELECT + queue, + priority, + (count(*) FILTER (WHERE state <> 'failed'))::bigint AS completed, + (count(*) FILTER ( + WHERE state = 'failed' + AND ( + $2::boolean IS FALSE + OR finalized_at < now() - make_interval(secs => $1::bigint) + ) + ))::bigint AS failed + FROM {done_child} + GROUP BY queue, priority + ), + -- Scan the entire partition being truncated (all generations), + -- mirroring done_counts' whole-child scan above. The rotate guard + -- keeps a partition to one generation at prune time, so this is the + -- same set today; counting the whole child keeps the rollup fold + -- conservative even if that invariant is ever weakened (a stray + -- generation's batches would otherwise be truncated without being + -- folded into the rollup -> terminal-count undercount). + batch_rows AS ( + SELECT + batch.ready_generation, + batch.queue, + batch.priority, + item.job_id, + item.run_lease + FROM {receipt_batch_child} AS batch + CROSS JOIN LATERAL unnest( + batch.job_ids, + batch.run_leases + ) AS item(job_id, run_lease) + ), + batch_counts AS ( + SELECT + batch_rows.queue, + batch_rows.priority, + count(*)::bigint AS completed + FROM batch_rows + WHERE NOT EXISTS ( + SELECT 1 + FROM {receipt_tomb_child} AS tomb + WHERE tomb.ready_generation = batch_rows.ready_generation + AND tomb.job_id = batch_rows.job_id + AND tomb.run_lease = batch_rows.run_lease + ) + GROUP BY batch_rows.queue, batch_rows.priority + ), + keys AS ( + SELECT queue, priority FROM done_counts + UNION + SELECT queue, priority FROM batch_counts + ) SELECT - batch_rows.queue, - batch_rows.priority, - count(*)::bigint AS completed - FROM batch_rows - WHERE NOT EXISTS ( - SELECT 1 - FROM {receipt_tomb_child} AS tomb - WHERE tomb.ready_generation = batch_rows.ready_generation - AND tomb.job_id = batch_rows.job_id - AND tomb.run_lease = batch_rows.run_lease + keys.queue, + keys.priority, + ( + COALESCE(done_counts.completed, 0) + + COALESCE(batch_counts.completed, 0) + )::bigint AS pruned_completed_count, + COALESCE(done_counts.failed, 0)::bigint AS pruned_failed_count + FROM keys + LEFT JOIN done_counts USING (queue, priority) + LEFT JOIN batch_counts USING (queue, priority) + "# + ))) + .bind(failed_retention_secs) + .bind(retention_floor) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + + let truncate_started = Instant::now(); + self + .truncate_used_children_tx( + &mut tx, + &[ + &ready_child, + &claim_attempt_child, + &done_child, + &tomb_child, + &segment_child, + &receipt_batch_child, + &receipt_tomb_child, + &delta_child, + ], ) - GROUP BY batch_rows.queue, batch_rows.priority - ), - keys AS ( - SELECT queue, priority FROM done_counts - UNION - SELECT queue, priority FROM batch_counts - ) - SELECT - keys.queue, - keys.priority, - ( - COALESCE(done_counts.completed, 0) - + COALESCE(batch_counts.completed, 0) - )::bigint AS pruned_completed_count, - COALESCE(done_counts.failed, 0)::bigint AS pruned_failed_count - FROM keys - LEFT JOIN done_counts USING (queue, priority) - LEFT JOIN batch_counts USING (queue, priority) - "# + .await + .map_err(map_sqlx_error)?; + let truncate_duration = truncate_started.elapsed(); + if !pruned_terminal_counts.is_empty() { + // #371: append the pruned counts as delta rows instead + // of upserting `queue_terminal_rollups` here — the + // upsert under a pinned MVCC horizon left one + // unreclaimable dead rollup version per prune. The + // maintenance leader folds the deltas + // (`fold_terminal_rollup_deltas`) when the horizon is + // clear; exact readers add unfolded delta sums to + // their rollup reads. + self.append_terminal_rollup_deltas_tx( + &mut tx, + pruned_terminal_counts.into_iter(), + ) + .await?; + } + + // #290: the live counter rows for this slot are about to + // be orphans (their underlying done_entries rows just got + // truncated). Delete them in the same transaction. The + // rollup fold above already captured ground-truth from + // the partition scan; this just cleans up the counter + // index entries so a future insert into a re-rotated + // slot starts from zero. + sqlx::query(audited_sql(format!( + "DELETE FROM {schema}.queue_terminal_live_counts WHERE ready_slot = $1" ))) - .bind(failed_retention_secs) - .bind(retention_floor) - .fetch_all(tx.as_mut()) + .bind(slot) + .execute(tx.as_mut()) .await .map_err(map_sqlx_error)?; + Ok(PruneOutcome::Pruned { + slot, + generation, + carried_failed_rows, + durations: PruneDurations { + lock: lock_duration, + truncate: truncate_duration, + // Set once the transaction commits. + commit: Duration::ZERO, + }, + }) + } + .await; - let truncate_started = Instant::now(); - let truncate = self - .truncate_used_children_tx( - &mut tx, - &[ - &ready_child, - &claim_attempt_child, - &done_child, - &tomb_child, - &segment_child, - &receipt_batch_child, - &receipt_tomb_child, - &delta_child, - ], - ) - .await; - let truncate_duration = truncate_started.elapsed(); - - match truncate { - Ok(_) => { - if !pruned_terminal_counts.is_empty() { - // #371: append the pruned counts as delta rows instead - // of upserting `queue_terminal_rollups` here — the - // upsert under a pinned MVCC horizon left one - // unreclaimable dead rollup version per prune. The - // maintenance leader folds the deltas - // (`fold_terminal_rollup_deltas`) when the horizon is - // clear; exact readers add unfolded delta sums to - // their rollup reads. - self.append_terminal_rollup_deltas_tx( - &mut tx, - pruned_terminal_counts.into_iter(), - ) - .await?; - } - - // #290: the live counter rows for this slot are about to - // be orphans (their underlying done_entries rows just got - // truncated). Delete them in the same transaction. The - // rollup fold above already captured ground-truth from - // the partition scan; this just cleans up the counter - // index entries so a future insert into a re-rotated - // slot starts from zero. - sqlx::query(audited_sql(format!( - "DELETE FROM {schema}.queue_terminal_live_counts WHERE ready_slot = $1" - ))) - .bind(slot) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - let commit_started = Instant::now(); - tx.commit().await.map_err(map_sqlx_error)?; - let commit_duration = commit_started.elapsed(); - if carried_failed_rows > 0 { - tracing::info!( - slot, - carried_failed_rows, - "Carried failed terminal rows inside the retention floor to the live done segment" - ); - } - Ok(PruneOutcome::Pruned { + let commit_started = Instant::now(); + let mut outcome = settle_prune_tx(tx, locked, PruneOutcome::Blocked { slot }).await?; + if let PruneOutcome::Pruned { + carried_failed_rows, + durations, + .. + } = &mut outcome + { + durations.commit = commit_started.elapsed(); + let carried_failed_rows = *carried_failed_rows; + if carried_failed_rows > 0 { + tracing::info!( slot, - generation, carried_failed_rows, - durations: PruneDurations { - lock: lock_duration, - truncate: truncate_duration, - commit: commit_duration, - }, - }) - } - Err(err) if is_lock_contention_error(&err) => { - let _ = tx.rollback().await; - Ok(PruneOutcome::Blocked { slot }) - } - Err(err) => { - let _ = tx.rollback().await; - Err(map_sqlx_error(err)) + "Carried failed terminal rows inside the retention floor to the live done segment" + ); } } + Ok(outcome) } #[tracing::instrument( @@ -19259,82 +19867,68 @@ impl QueueStorage { let lease_child = lease_child_name(schema, slot as usize); - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + let locked: Result = async { + set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - let lock_started = Instant::now(); - let lock_table = sqlx::query(audited_sql(format!( - "LOCK TABLE {lease_child} IN ACCESS EXCLUSIVE MODE" - ))) - .execute(tx.as_mut()) - .await; - let lock_duration = lock_started.elapsed(); + let lock_started = Instant::now(); + sqlx::query(audited_sql(format!( + "LOCK TABLE {lease_child} IN ACCESS EXCLUSIVE MODE" + ))) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let lock_duration = lock_started.elapsed(); - if let Err(err) = lock_table { - let _ = tx.rollback().await; - if is_lock_contention_error(&err) { - return Ok(PruneOutcome::Blocked { slot }); - } - return Err(map_sqlx_error(err)); - } + // The ring rotation serialization held above keeps the cursor + // stable, but keep the explicit recheck to document the truncate + // precondition. + let (current_slot, _) = self + .ring_cursor_tx(&mut tx, RingFamily::Lease, authority) + .await?; - // The ring rotation serialization held above keeps the cursor - // stable, but keep the explicit recheck to document the truncate - // precondition. - let (current_slot, _) = self - .ring_cursor_tx(&mut tx, RingFamily::Lease, authority) - .await?; + if current_slot == slot { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::LeaseCurrent, + count: 0, + }); + } - if current_slot == slot { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::LeaseCurrent, - count: 0, - }); - } + let active_leases = Self::relation_has_rows_tx(&mut tx, &lease_child).await?; - let active_leases = Self::relation_has_rows_tx(&mut tx, &lease_child).await?; + if active_leases { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::LeaseActive, + count: busy_indicator(active_leases), + }); + } - if active_leases { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { + let truncate_started = Instant::now(); + self.truncate_used_children_tx(&mut tx, &[&lease_child]) + .await + .map_err(map_sqlx_error)?; + let truncate_duration = truncate_started.elapsed(); + Ok(PruneOutcome::Pruned { slot, - reason: SkipReason::LeaseActive, - count: busy_indicator(active_leases), - }); + generation, + carried_failed_rows: 0, + durations: PruneDurations { + lock: lock_duration, + truncate: truncate_duration, + // Set once the transaction commits. + commit: Duration::ZERO, + }, + }) } + .await; - let truncate_started = Instant::now(); - let truncate = self - .truncate_used_children_tx(&mut tx, &[&lease_child]) - .await; - let truncate_duration = truncate_started.elapsed(); - - match truncate { - Ok(_) => { - let commit_started = Instant::now(); - tx.commit().await.map_err(map_sqlx_error)?; - let commit_duration = commit_started.elapsed(); - Ok(PruneOutcome::Pruned { - slot, - generation, - carried_failed_rows: 0, - durations: PruneDurations { - lock: lock_duration, - truncate: truncate_duration, - commit: commit_duration, - }, - }) - } - Err(err) if is_lock_contention_error(&err) => { - let _ = tx.rollback().await; - Ok(PruneOutcome::Blocked { slot }) - } - Err(err) => { - let _ = tx.rollback().await; - Err(map_sqlx_error(err)) - } + let commit_started = Instant::now(); + let mut outcome = settle_prune_tx(tx, locked, PruneOutcome::Blocked { slot }).await?; + if let PruneOutcome::Pruned { durations, .. } = &mut outcome { + durations.commit = commit_started.elapsed(); } + Ok(outcome) } pub async fn vacuum_leases(&self, pool: &PgPool) -> Result<(), AwaError> { @@ -19638,116 +20232,103 @@ impl QueueStorage { }); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - - let lock_started = Instant::now(); - let lock_tables = sqlx::query(audited_sql(format!( - "LOCK TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child} IN ACCESS EXCLUSIVE MODE" - ))) - .execute(tx.as_mut()) - .await; - let lock_duration = lock_started.elapsed(); - - if let Err(err) = lock_tables { - let _ = tx.rollback().await; - if is_lock_contention_error(&err) { - return Ok(PruneOutcome::Blocked { slot }); - } - return Err(map_sqlx_error(err)); - } + let locked: Result = async { + set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; - // After taking ACCESS EXCLUSIVE, recheck that the slot is still - // not the current one. The ring rotation serialization should - // already make this stable, but keeping the explicit gate documents - // the truncate precondition. - let (current_slot, _) = self - .ring_cursor_tx(&mut tx, RingFamily::Claim, authority) - .await?; + let lock_started = Instant::now(); + sqlx::query(audited_sql(format!( + "LOCK TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child} IN ACCESS EXCLUSIVE MODE" + ))) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let lock_duration = lock_started.elapsed(); - if current_slot == slot { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::ClaimCurrent, - count: 0, - }); - } + // After taking ACCESS EXCLUSIVE, recheck that the slot is still + // not the current one. The ring rotation serialization should + // already make this stable, but keeping the explicit gate documents + // the truncate precondition. + let (current_slot, _) = self + .ring_cursor_tx(&mut tx, RingFamily::Claim, authority) + .await?; - let open_claims = claim_prune_has_open_claims_tx( - &mut tx, - schema, - &claim_child, - &claim_batch_child, - &closure_child, - &closure_batch_child, - ) - .await?; - if open_claims { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(PruneOutcome::SkippedActive { - slot, - reason: SkipReason::ClaimOpen, - count: busy_indicator(open_claims), - }); - } + if current_slot == slot { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::ClaimCurrent, + count: 0, + }); + } - let truncate_started = Instant::now(); - let truncate = self - .truncate_used_children_tx( + let open_claims = claim_prune_has_open_claims_tx( &mut tx, - &[ - &claim_child, - &claim_batch_child, - &closure_child, - &closure_batch_child, - ], + schema, + &claim_child, + &claim_batch_child, + &closure_child, + &closure_batch_child, ) - .await; - let truncate_duration = truncate_started.elapsed(); + .await?; + if open_claims { + return Ok(PruneOutcome::SkippedActive { + slot, + reason: SkipReason::ClaimOpen, + count: busy_indicator(open_claims), + }); + } - match truncate { - Ok(_) => { - sqlx::query(audited_sql(format!( - r#" - UPDATE {schema}.claim_ring_slots - SET rescue_cursor_claimed_at = '-infinity'::timestamptz, - rescue_cursor_job_id = 0, - rescue_cursor_run_lease = 0, - deadline_cursor_deadline_at = '-infinity'::timestamptz, - deadline_cursor_job_id = 0, - deadline_cursor_run_lease = 0, - batch_deadline_cursor_deadline_at = '-infinity'::timestamptz, - batch_deadline_cursor_batch_id = 0 - WHERE slot = $1 - "# - ))) - .bind(slot) - .execute(tx.as_mut()) + let truncate_started = Instant::now(); + self + .truncate_used_children_tx( + &mut tx, + &[ + &claim_child, + &claim_batch_child, + &closure_child, + &closure_batch_child, + ], + ) .await .map_err(map_sqlx_error)?; - let commit_started = Instant::now(); - tx.commit().await.map_err(map_sqlx_error)?; - let commit_duration = commit_started.elapsed(); - Ok(PruneOutcome::Pruned { - slot, - generation, - carried_failed_rows: 0, - durations: PruneDurations { - lock: lock_duration, - truncate: truncate_duration, - commit: commit_duration, - }, - }) - } - Err(err) if is_lock_contention_error(&err) => { - let _ = tx.rollback().await; - Ok(PruneOutcome::Blocked { slot }) - } - Err(err) => { - let _ = tx.rollback().await; - Err(map_sqlx_error(err)) - } + let truncate_duration = truncate_started.elapsed(); + sqlx::query(audited_sql(format!( + r#" + UPDATE {schema}.claim_ring_slots + SET rescue_cursor_claimed_at = '-infinity'::timestamptz, + rescue_cursor_job_id = 0, + rescue_cursor_run_lease = 0, + deadline_cursor_deadline_at = '-infinity'::timestamptz, + deadline_cursor_job_id = 0, + deadline_cursor_run_lease = 0, + batch_deadline_cursor_deadline_at = '-infinity'::timestamptz, + batch_deadline_cursor_batch_id = 0 + WHERE slot = $1 + "# + ))) + .bind(slot) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + Ok(PruneOutcome::Pruned { + slot, + generation, + carried_failed_rows: 0, + durations: PruneDurations { + lock: lock_duration, + truncate: truncate_duration, + // Set once the transaction commits. + commit: Duration::ZERO, + }, + }) } + .await; + + let commit_started = Instant::now(); + let mut outcome = settle_prune_tx(tx, locked, PruneOutcome::Blocked { slot }).await?; + if let PruneOutcome::Pruned { durations, .. } = &mut outcome { + durations.commit = commit_started.elapsed(); + } + Ok(outcome) } fn job_id_sequence(&self) -> String { diff --git a/awa-model/src/storage.rs b/awa-model/src/storage.rs index b7a0ff1ad..14770c40a 100644 --- a/awa-model/src/storage.rs +++ b/awa-model/src/storage.rs @@ -2,7 +2,7 @@ use crate::admin; use crate::error::AwaError; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; -use sqlx::{FromRow, PgExecutor, PgPool}; +use sqlx::{Connection, FromRow, PgConnection, PgExecutor, PgPool}; use std::collections::BTreeMap; #[derive(Debug, Clone, Serialize, Deserialize, FromRow, PartialEq)] @@ -469,3 +469,272 @@ pub async fn status_report(pool: &PgPool) -> Result, + pub reassign_sql: Option, + /// Login roles, superusers aside, that can run every statement above. + /// Empty when no such role exists yet. + pub repair_roles: Vec, +} + +/// What a repair did to one lane sequence. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub enum LaneSequenceOutcome { + /// Owned by the schema owner with every runtime grant. + Repaired, + /// Grants complete, but the role could not reassign the sequence. + OwnerUnchanged, + /// The role could neither reassign the sequence nor add its grants. + Unrepaired, + /// An open transaction held the sequence's lock past the repair's lock + /// timeout, so nothing changed. Rerun the repair. + Busy, +} + +impl TryFrom for LaneSequenceOutcome { + type Error = String; + + fn try_from(outcome: String) -> Result { + match outcome.as_str() { + "repaired" => Ok(Self::Repaired), + "owner_unchanged" => Ok(Self::OwnerUnchanged), + "unrepaired" => Ok(Self::Unrepaired), + "busy" => Ok(Self::Busy), + _ => Err(format!("unknown lane sequence repair outcome {outcome:?}")), + } + } +} + +/// What [`repair_lane_sequences`] did to one lane sequence that had a problem, +/// with what is still wrong afterwards. +#[derive(Debug, Clone, Serialize, Deserialize, FromRow, PartialEq, Eq)] +pub struct LaneSequenceRepair { + pub schema: String, + pub sequence_name: String, + /// Owner after the repair. + pub owner: String, + pub expected_owner: String, + #[sqlx(try_from = "String")] + pub outcome: LaneSequenceOutcome, + pub missing_grants: Vec, + pub reassign_sql: Option, + /// Login roles, superusers aside, that can finish the repair. + pub repair_roles: Vec, +} + +/// Lane sequences in `schema` that need repair. Read-only; any role may run it. +pub async fn lane_sequence_problems<'e, E>( + executor: E, + schema: &str, +) -> Result, AwaError> +where + E: PgExecutor<'e>, +{ + Ok(sqlx::query_as( + "SELECT $1::text AS schema, sequence_name, owner, expected_owner, missing_grants, \ + reassign_sql, repair_roles \ + FROM awa.lane_sequence_problems($1)", + ) + .bind(schema) + .fetch_all(executor) + .await?) +} + +/// Give each lane sequence in `schema` to the schema owner and add its missing +/// grants, as far as the connection's roles are allowed to. A sequence that +/// cannot be repaired is returned with the statements still to run and the +/// roles that can run them, never an error. +/// +/// Each sequence is repaired by its own statement, outside any transaction +/// the caller holds, so a reassignment's ACCESS EXCLUSIVE lock is released +/// before the next sequence and waits at most the repair function's lock +/// timeout. A sequence locked for longer by an open enqueue or claim comes +/// back [`LaneSequenceOutcome::Busy`]. +/// +/// Reassigning needs the current owner's privileges. A migration role that +/// connects as itself and sets its role to the schema owner (`role` in the +/// connection options) loses its own privileges while acting as the owner, +/// so a sequence the current role cannot repair is retried once as the login +/// role. +pub async fn repair_lane_sequences( + conn: &mut PgConnection, + schema: &str, +) -> Result, AwaError> { + let (current_role, login_role): (String, String) = + sqlx::query_as("SELECT current_user::text, session_user::text") + .fetch_one(&mut *conn) + .await?; + let problems = lane_sequence_problems(&mut *conn, schema).await?; + let mut repairs = Vec::with_capacity(problems.len()); + for problem in problems { + let mut repair: Option = sqlx::query_as(REPAIR_LANE_SEQUENCE_SQL) + .bind(schema) + .bind(&problem.sequence_name) + .fetch_optional(&mut *conn) + .await?; + let role_was_the_obstacle = matches!( + repair.as_ref().map(|repair| repair.outcome), + Some(LaneSequenceOutcome::OwnerUnchanged | LaneSequenceOutcome::Unrepaired) + ); + if role_was_the_obstacle && login_role != current_role { + if let Some(retried) = + repair_lane_sequence_as_login_role(conn, schema, &problem.sequence_name).await? + { + repair = retried; + } + } + repairs.extend(repair); + } + Ok(repairs) +} + +const REPAIR_LANE_SEQUENCE_SQL: &str = + "SELECT $1::text AS schema, sequence_name, owner, expected_owner, outcome, missing_grants, \ + reassign_sql, repair_roles \ + FROM awa.repair_lane_sequence($1, $2)"; + +/// Repair one sequence with the login role as the current role, for this +/// transaction only. `None` when the login role may not run the repair. +async fn repair_lane_sequence_as_login_role( + conn: &mut PgConnection, + schema: &str, + sequence_name: &str, +) -> Result>, AwaError> { + let mut transaction = conn.begin().await?; + sqlx::query("SET LOCAL ROLE NONE") + .execute(&mut *transaction) + .await?; + let repair = sqlx::query_as(REPAIR_LANE_SEQUENCE_SQL) + .bind(schema) + .bind(sequence_name) + .fetch_optional(&mut *transaction) + .await; + match repair { + Ok(repair) => { + transaction.commit().await?; + Ok(Some(repair)) + } + Err(sqlx::Error::Database(error)) if error.code().as_deref() == Some("42501") => { + transaction.rollback().await?; + Ok(None) + } + Err(error) => Err(error.into()), + } +} + +/// [`repair_lane_sequences`] for every queue-storage schema with a lane +/// sequence that needs repair. +pub async fn repair_lane_sequences_in_every_schema( + conn: &mut PgConnection, +) -> Result, AwaError> { + let mut repairs = Vec::new(); + for (schema, _) in lane_sequence_problem_counts(&mut *conn).await? { + repairs.extend(repair_lane_sequences(conn, &schema).await?); + } + Ok(repairs) +} + +/// One line per group of lane sequences that still need repair, saying what +/// to do next. `busy` marks sequences a repair found locked. +pub fn lane_sequence_next_steps<'a>( + leftovers: impl IntoIterator, +) -> Vec { + let mut groups: BTreeMap<(String, bool, String, String, Vec), usize> = BTreeMap::new(); + for (problem, busy) in leftovers { + *groups + .entry(( + problem.schema.clone(), + busy, + problem.owner.clone(), + problem.expected_owner.clone(), + problem.repair_roles.clone(), + )) + .or_default() += 1; + } + groups + .into_iter() + .map( + |((schema, busy, owner, expected_owner, repair_roles), count)| { + let rerun = format!("awa storage lane-sequences --schema {schema} --repair"); + if busy { + format!( + "{count} lane sequence(s) in schema {schema} owned by {owner} are busy: an \ + open transaction held the lock, so nothing changed. Rerun `{rerun}`." + ) + } else if repair_roles.is_empty() { + format!( + "{count} lane sequence(s) in schema {schema} owned by {owner} still need \ + repair, and no login role can both act as {owner} and SET ROLE to \ + {expected_owner}. Grant {owner} to the migration role, then rerun `{rerun}` \ + connected as the migration role, or run the statements from \ + `awa storage lane-sequences --schema {schema}` as {owner}." + ) + } else { + format!( + "{count} lane sequence(s) in schema {schema} owned by {owner} still need \ + repair. Rerun `{rerun}` connected as one of: {}.", + repair_roles.join(", ") + ) + } + }, + ) + .collect() +} + +impl LaneSequenceRepair { + /// The problem left after this repair, or `None` when it was repaired. + pub fn remaining_problem(&self) -> Option { + (self.outcome != LaneSequenceOutcome::Repaired).then(|| LaneSequenceProblem { + schema: self.schema.clone(), + sequence_name: self.sequence_name.clone(), + owner: self.owner.clone(), + expected_owner: self.expected_owner.clone(), + missing_grants: self.missing_grants.clone(), + reassign_sql: self.reassign_sql.clone(), + repair_roles: self.repair_roles.clone(), + }) + } +} + +/// [`lane_sequence_next_steps`] for the sequences a repair left behind. +pub fn lane_sequence_repair_next_steps(repairs: &[LaneSequenceRepair]) -> Vec { + let leftovers: Vec<(LaneSequenceProblem, bool)> = repairs + .iter() + .filter_map(|repair| { + repair + .remaining_problem() + .map(|problem| (problem, repair.outcome == LaneSequenceOutcome::Busy)) + }) + .collect(); + lane_sequence_next_steps(leftovers.iter().map(|(problem, busy)| (problem, *busy))) +} + +/// How many lane sequences need repair in each queue-storage schema that has +/// any. Read-only. +pub async fn lane_sequence_problem_counts<'e, E>( + executor: E, +) -> Result, AwaError> +where + E: PgExecutor<'e>, +{ + Ok(sqlx::query_as( + "SELECT discovered.schema_name, count(*) \ + FROM awa.queue_storage_schemas() AS discovered \ + CROSS JOIN LATERAL awa.lane_sequence_problems(discovered.schema_name) AS problem \ + GROUP BY discovered.schema_name \ + ORDER BY discovered.schema_name", + ) + .fetch_all(executor) + .await?) +} diff --git a/awa-model/tests/least_privilege_runtime_test.rs b/awa-model/tests/least_privilege_runtime_test.rs index 99e69d536..f06542c99 100644 --- a/awa-model/tests/least_privilege_runtime_test.rs +++ b/awa-model/tests/least_privilege_runtime_test.rs @@ -147,6 +147,10 @@ async fn exercise(options: PgConnectOptions, owner: &str, runtime: &str) -> Test assert_eq!(before, after); } assert_no_head_table_ddl_locks(&owner_pool).await?; + for schema in SCHEMAS { + assert_enqueue_function_privileges_fall_back(&owner_pool, &runtime_pool, schema, runtime) + .await?; + } runtime_pool.close().await; owner_pool.close().await; Ok(()) @@ -250,9 +254,29 @@ async fn check_runtime(pool: &PgPool, schema: &str) -> TestResult { INSERT INTO {schema}.queue_claim_heads(queue, priority, enqueue_shard) VALUES ('prepared_queue', 2, 0) ON CONFLICT DO NOTHING; SELECT {schema}.reserve_enqueue_seq('prepared_queue', 2::smallint, 0::smallint, 3); "#)))).await?; + // An unprepared lane is created through the schema owner's provisioning + // function, so the runtime still issues no DDL of its own. + pool.execute(sqlx::raw_sql(awa_model::audited_sql(format!( + "SELECT {schema}.ensure_lane_sequences('unprepared_queue', 2::smallint, 0::smallint); \ + SELECT {schema}.provision_lane('unprepared_queue', 2::smallint, 0::smallint);" + )))) + .await?; + let foreign_owned: i64 = sqlx::query_scalar(awa_model::audited_sql(format!( + "SELECT count(*) FROM pg_class \ + WHERE relnamespace = '{schema}'::regnamespace AND relkind = 'S' \ + AND relowner <> (SELECT nspowner FROM pg_namespace WHERE nspname = '{schema}')" + ))) + .fetch_one(pool) + .await?; + assert_eq!(foreign_owned, 0); + sqlx::query_scalar::<_, i64>(awa_model::audited_sql(format!( + "SELECT {schema}.reserve_enqueue_seq('unprepared_queue', 2::smallint, 0::smallint, 1)" + ))) + .fetch_one(pool) + .await?; let error = pool .execute(sqlx::raw_sql(awa_model::audited_sql(format!( - "SELECT {schema}.ensure_lane_sequences('unprepared_queue', 2::smallint, 0::smallint)" + "CREATE TABLE {schema}.must_not_create(id integer)" )))) .await .unwrap_err(); @@ -260,21 +284,84 @@ async fn check_runtime(pool: &PgPool, schema: &str) -> TestResult { error.as_database_error().unwrap().code().as_deref(), Some("42501") ); - assert!(error - .as_database_error() - .unwrap() - .message() - .contains("has not been provisioned")); - let error = pool + Ok(()) +} + +/// The one-call enqueue needs EXECUTE on enqueue_ready_rows and on +/// ring_cursor. Without either the store keeps the multi-statement path; a +/// grant revoked after the store has used the function falls back within +/// the same call, after rolling back the store-owned transaction. +async fn assert_enqueue_function_privileges_fall_back( + owner_pool: &PgPool, + runtime_pool: &PgPool, + schema: &str, + runtime: &str, +) -> TestResult { + let enqueue = |pool: PgPool| async move { + QueueStorage::new(QueueStorageConfig { + schema: schema.into(), + ..Default::default() + }) + .expect("store") + .enqueue_batch(&pool, "prepared_queue", 2, 1) + .await + }; + owner_pool .execute(sqlx::raw_sql(awa_model::audited_sql(format!( - "CREATE TABLE {schema}.must_not_create(id integer)" + "REVOKE EXECUTE ON FUNCTION {schema}.ring_cursor(text) FROM {runtime}" )))) - .await - .unwrap_err(); + .await?; assert_eq!( - error.as_database_error().unwrap().code().as_deref(), - Some("42501") + enqueue(runtime_pool.clone()).await?, + 1, + "without ring_cursor EXECUTE the store keeps the multi-statement path" + ); + owner_pool + .execute(sqlx::raw_sql(awa_model::audited_sql(format!( + "GRANT EXECUTE ON FUNCTION {schema}.ring_cursor(text) TO {runtime}" + )))) + .await?; + + let store = QueueStorage::new(QueueStorageConfig { + schema: schema.into(), + ..Default::default() + })?; + assert_eq!( + store + .enqueue_batch(runtime_pool, "prepared_queue", 2, 1) + .await?, + 1 + ); + let signature: String = sqlx::query_scalar(awa_model::audited_sql(format!( + "SELECT oid::regprocedure::text FROM pg_proc \ + WHERE pronamespace = '{schema}'::regnamespace AND proname = 'enqueue_ready_rows'" + ))) + .fetch_one(owner_pool) + .await?; + owner_pool + .execute(sqlx::raw_sql(awa_model::audited_sql(format!( + "REVOKE EXECUTE ON FUNCTION {signature} FROM PUBLIC, {runtime}" + )))) + .await?; + assert_eq!( + store + .enqueue_batch(runtime_pool, "prepared_queue", 2, 1) + .await?, + 1, + "the same call falls back after EXECUTE is revoked" ); + assert_eq!( + store + .enqueue_batch(runtime_pool, "prepared_queue", 2, 1) + .await?, + 1, + "the next call probes again and uses the multi-statement path" + ); + owner_pool + .execute(sqlx::raw_sql(awa_model::audited_sql(format!( + "GRANT EXECUTE ON FUNCTION {signature} TO {runtime}" + )))) + .await?; Ok(()) } @@ -331,10 +418,14 @@ async fn exercise_lazy_lane(options: PgConnectOptions, owner: &str, runtime: &st "GRANT USAGE, CREATE ON SCHEMA awa TO {runtime}; \ GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA awa TO {runtime}; \ GRANT USAGE, SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA awa TO {runtime}; \ - GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa TO {runtime};" + GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa TO {runtime}; \ + REVOKE EXECUTE ON FUNCTION awa.provision_lane_sequences(text, smallint, smallint) \ + FROM PUBLIC, {runtime};" ), ) .await?; + // Without the provisioning function the runtime creates the lane's + // sequences itself and owns them, as a 0.6 runtime with CREATE did. let runtime_pool = PgPoolOptions::new() .max_connections(2) .connect_with(options.username(runtime)) diff --git a/awa-model/tests/schema_discovery_owner_test.rs b/awa-model/tests/schema_discovery_owner_test.rs index 3dfc52bbd..914414343 100644 --- a/awa-model/tests/schema_discovery_owner_test.rs +++ b/awa-model/tests/schema_discovery_owner_test.rs @@ -556,8 +556,15 @@ async fn assert_configured_substrates_fail_closed( ) .fetch_one(migrator_pool) .await?; - sqlx::query("DELETE FROM awa.schema_version WHERE version = $1") - .bind(migrations::CURRENT_VERSION) + // The newest migration that runs the discovery preflight. Later ones + // that touch no custom substrate carry no preflight to refuse. + let (preflight_version, _, pending_sql) = migrations::migration_sql() + .into_iter() + .rev() + .find(|(_, _, sql)| sql.contains("unsafe configured queue-storage schema")) + .expect("a migration that runs the discovery preflight"); + sqlx::query("DELETE FROM awa.schema_version WHERE version >= $1") + .bind(preflight_version) .execute(migrator_pool) .await?; let previous_version: i32 = sqlx::query_scalar("SELECT max(version) FROM awa.schema_version") @@ -602,7 +609,6 @@ async fn assert_configured_substrates_fail_closed( "CREATE SCHEMA awa_exp AUTHORIZATION {outsider}" )))) .await?; - let (_, _, pending_sql) = migrations::migration_sql().pop().unwrap(); for (details, expected_schema) in [ ( serde_json::json!({"schema": "injected_custom"}), diff --git a/awa-python/python/awa/_awa.pyi b/awa-python/python/awa/_awa.pyi index 2219bfc3b..34f5edbab 100644 --- a/awa-python/python/awa/_awa.pyi +++ b/awa-python/python/awa/_awa.pyi @@ -605,6 +605,7 @@ class Client: deadline_rescue_interval_ms: int | None = None, callback_rescue_interval_ms: int | None = None, terminal_count_rollup_interval_ms: int | None = None, + statistics_refresh_interval_ms: int | None = None, queue_storage_schema: str | None = None, queue_storage_queue_slot_count: int = 16, queue_storage_lease_slot_count: int = 8, diff --git a/awa-python/python/awa/client.py b/awa-python/python/awa/client.py index f734ca679..131d3391b 100644 --- a/awa-python/python/awa/client.py +++ b/awa-python/python/awa/client.py @@ -922,6 +922,7 @@ async def start( deadline_rescue_interval_ms: int | None = None, callback_rescue_interval_ms: int | None = None, terminal_count_rollup_interval_ms: int | None = None, + statistics_refresh_interval_ms: int | None = None, queue_storage_schema: str | None = None, queue_storage_queue_slot_count: int = DEFAULT_QUEUE_STORAGE_QUEUE_SLOT_COUNT, queue_storage_lease_slot_count: int = DEFAULT_QUEUE_STORAGE_LEASE_SLOT_COUNT, @@ -967,6 +968,11 @@ async def start( counters. Exact reads include pending deltas, so the default favors fewer hot counter updates over immediate folding. + ``statistics_refresh_interval_ms`` controls how often the maintenance + leader runs ``ANALYZE`` on the queue-storage head tables and + partitioned ring parents, which autovacuum never analyses (default 5 + minutes; ``0`` disables the refresh). + Queue dicts also accept ``claimers`` and ``claim_batch_size`` for hot queue-storage queues. ``claimers`` adds dispatcher/claimer loops that share the queue's worker permits; benchmark ``2`` or ``4`` before @@ -988,6 +994,7 @@ async def start( deadline_rescue_interval_ms=deadline_rescue_interval_ms, callback_rescue_interval_ms=callback_rescue_interval_ms, terminal_count_rollup_interval_ms=terminal_count_rollup_interval_ms, + statistics_refresh_interval_ms=statistics_refresh_interval_ms, queue_storage_schema=queue_storage_schema, queue_storage_queue_slot_count=queue_storage_queue_slot_count, queue_storage_lease_slot_count=queue_storage_lease_slot_count, diff --git a/awa-python/src/client.rs b/awa-python/src/client.rs index f141508a7..e17b0c15e 100644 --- a/awa-python/src/client.rs +++ b/awa-python/src/client.rs @@ -1796,7 +1796,7 @@ impl PyClient { }) } - #[pyo3(signature = (queues=None, *, poll_interval_ms=200, global_max_workers=None, completed_retention_hours=None, failed_retention_hours=None, descriptor_retention_days=None, cleanup_batch_size=None, leader_election_interval_ms=None, heartbeat_interval_ms=None, promote_interval_ms=None, heartbeat_rescue_interval_ms=None, heartbeat_staleness_ms=None, deadline_rescue_interval_ms=None, callback_rescue_interval_ms=None, terminal_count_rollup_interval_ms=None, queue_storage_schema=None, queue_storage_queue_slot_count=16, queue_storage_lease_slot_count=8, queue_storage_claim_slot_count=8, queue_storage_queue_stripe_count=1, queue_storage_queue_rotate_interval_ms=1000, queue_storage_lease_rotate_interval_ms=1000, queue_storage_claim_rotate_interval_ms=None, queue_storage_rotate_min_slot_age_ms=None, queue_storage_rotate_target_slot_bytes=None, queue_storage_rotate_min_slot_rows=None, queue_storage_rotate_max_slot_age_ms=None, storage_transition_role=None))] + #[pyo3(signature = (queues=None, *, poll_interval_ms=200, global_max_workers=None, completed_retention_hours=None, failed_retention_hours=None, descriptor_retention_days=None, cleanup_batch_size=None, leader_election_interval_ms=None, heartbeat_interval_ms=None, promote_interval_ms=None, heartbeat_rescue_interval_ms=None, heartbeat_staleness_ms=None, deadline_rescue_interval_ms=None, callback_rescue_interval_ms=None, terminal_count_rollup_interval_ms=None, statistics_refresh_interval_ms=None, queue_storage_schema=None, queue_storage_queue_slot_count=16, queue_storage_lease_slot_count=8, queue_storage_claim_slot_count=8, queue_storage_queue_stripe_count=1, queue_storage_queue_rotate_interval_ms=1000, queue_storage_lease_rotate_interval_ms=1000, queue_storage_claim_rotate_interval_ms=None, queue_storage_rotate_min_slot_age_ms=None, queue_storage_rotate_target_slot_bytes=None, queue_storage_rotate_min_slot_rows=None, queue_storage_rotate_max_slot_age_ms=None, storage_transition_role=None))] #[allow(clippy::too_many_arguments)] fn start<'py>( &self, @@ -1816,6 +1816,7 @@ impl PyClient { deadline_rescue_interval_ms: Option, callback_rescue_interval_ms: Option, terminal_count_rollup_interval_ms: Option, + statistics_refresh_interval_ms: Option, queue_storage_schema: Option, queue_storage_queue_slot_count: u32, queue_storage_lease_slot_count: u32, @@ -1950,6 +1951,9 @@ impl PyClient { } builder = builder.terminal_count_rollup_interval(Duration::from_millis(ms)); } + if let Some(ms) = statistics_refresh_interval_ms { + builder = builder.statistics_refresh_interval(Duration::from_millis(ms)); + } if queue_storage_queue_slot_count == 0 { return Err(pyo3::exceptions::PyValueError::new_err( "queue_storage_queue_slot_count must be > 0", diff --git a/awa-python/tests/test_start_config.py b/awa-python/tests/test_start_config.py index 302df0fc2..4d19775d0 100644 --- a/awa-python/tests/test_start_config.py +++ b/awa-python/tests/test_start_config.py @@ -316,12 +316,26 @@ async def handle(job): deadline_rescue_interval_ms=1_000, callback_rescue_interval_ms=1_000, terminal_count_rollup_interval_ms=1_000, + statistics_refresh_interval_ms=1_000, leader_election_interval_ms=200, promote_interval_ms=500, ) await client.shutdown() +@pytest.mark.asyncio +async def test_statistics_refresh_can_be_disabled(client): + """statistics_refresh_interval_ms=0 disables the refresh; start() accepts it.""" + queue = "cfg_statistics_refresh_off" + + @client.task(ConfigTestJob, queue=queue) + async def handle(job): + return None + + await client.start([(queue, 1)], statistics_refresh_interval_ms=0) + await client.shutdown() + + @pytest.mark.asyncio async def test_per_queue_retention_in_dict_config(client): """Dict form with per-queue retention config starts successfully.""" diff --git a/awa-worker/src/client.rs b/awa-worker/src/client.rs index 253bbde7a..d32e8b977 100644 --- a/awa-worker/src/client.rs +++ b/awa-worker/src/client.rs @@ -163,6 +163,7 @@ pub struct ClientBuilder { leader_check_interval: Option, priority_aging_interval: Option, terminal_count_rollup_interval: Option, + statistics_refresh_interval: Option, completed_retention: Option, failed_retention: Option, descriptor_retention: Option, @@ -226,6 +227,7 @@ impl ClientBuilder { leader_check_interval: None, priority_aging_interval: None, terminal_count_rollup_interval: None, + statistics_refresh_interval: None, completed_retention: None, failed_retention: None, descriptor_retention: None, @@ -905,6 +907,20 @@ impl ClientBuilder { self } + /// Set how often the maintenance leader refreshes planner statistics for + /// the queue-storage tables autovacuum never analyses (default: 5 minutes). + /// + /// Autovacuum does not analyse partitioned tables, and never reaches its + /// threshold on the two small head tables, so the leader runs `ANALYZE` + /// on both when it is elected and then every interval, through + /// `awa.analyze_queue_storage_relation()` (schema v053) so that a runtime + /// role that does not own the tables still works. `Duration::ZERO` + /// disables the refresh. Queue storage only. + pub fn statistics_refresh_interval(mut self, interval: Duration) -> Self { + self.statistics_refresh_interval = Some(interval); + self + } + /// Set how often queue depth/lag metrics are published (default: 30s). pub fn queue_stats_interval(mut self, interval: Duration) -> Self { self.queue_stats_interval = Some(interval); @@ -1256,6 +1272,7 @@ impl ClientBuilder { leader_check_interval: self.leader_check_interval, priority_aging_interval: self.priority_aging_interval, terminal_count_rollup_interval: self.terminal_count_rollup_interval, + statistics_refresh_interval: self.statistics_refresh_interval, completed_retention: self.completed_retention, failed_retention: self.failed_retention, descriptor_retention: self.descriptor_retention, @@ -1362,6 +1379,7 @@ pub struct Client { leader_check_interval: Option, priority_aging_interval: Option, terminal_count_rollup_interval: Option, + statistics_refresh_interval: Option, completed_retention: Option, failed_retention: Option, descriptor_retention: Option, @@ -1932,6 +1950,9 @@ impl Client { if let Some(interval) = self.terminal_count_rollup_interval { maintenance = maintenance.terminal_count_rollup_interval(interval); } + if let Some(interval) = self.statistics_refresh_interval { + maintenance = maintenance.statistics_refresh_interval(interval); + } if let Some(retention) = self.completed_retention { maintenance = maintenance.completed_retention(retention); } diff --git a/awa-worker/src/lib.rs b/awa-worker/src/lib.rs index 4768e59bb..f9a171104 100644 --- a/awa-worker/src/lib.rs +++ b/awa-worker/src/lib.rs @@ -17,6 +17,7 @@ pub mod metrics; mod notify_hub; mod output; mod runtime; +mod statistics; mod storage; // Re-exports diff --git a/awa-worker/src/maintenance.rs b/awa-worker/src/maintenance.rs index d84960d37..782ca622b 100644 --- a/awa-worker/src/maintenance.rs +++ b/awa-worker/src/maintenance.rs @@ -1,5 +1,6 @@ use crate::executor::DlqPolicy; use crate::runtime::InFlightMap; +use crate::statistics::{RefreshGuard, StatisticsRefresher}; use crate::storage::{QueueStorageRuntime, RuntimeStorage}; use awa_model::audited_sql; use awa_model::cron::{ @@ -738,6 +739,9 @@ pub struct MaintenanceService { priority_aging_interval: Duration, batch_operations_interval: Duration, terminal_count_rollup_interval: Duration, + /// How often the leader refreshes planner statistics for the queue-storage + /// tables autovacuum never analyses. Zero disables the refresh. + statistics_refresh_interval: Duration, /// How long a descriptor catalog row can sit without being refreshed /// before the maintenance leader deletes it. Zero disables cleanup. /// Default: 30 days. @@ -757,6 +761,8 @@ pub struct MaintenanceService { ring_authority_flip_aware_since: std::sync::Mutex>, } +/// Default cadence of the planner-statistics refresh on the maintenance leader. +pub(crate) const DEFAULT_STATISTICS_REFRESH_INTERVAL: Duration = Duration::from_secs(300); const PROMOTE_BATCH_SIZE: i64 = 4_096; const PROMOTE_MAX_BATCHES_PER_TICK: usize = 32; const CRON_CATCH_UP_LIMIT: usize = 1_000; @@ -817,6 +823,7 @@ impl MaintenanceService { priority_aging_interval: Duration::from_secs(60), batch_operations_interval: Duration::from_secs(1), terminal_count_rollup_interval: Duration::from_secs(30), + statistics_refresh_interval: DEFAULT_STATISTICS_REFRESH_INTERVAL, descriptor_retention: Duration::from_secs(30 * 86400), // 30d ring_authority_auto_flip: true, ring_authority_auto_flip_stable_period: Duration::from_secs(600), // 10min @@ -873,6 +880,14 @@ impl MaintenanceService { self } + /// Set how often the leader runs `ANALYZE` on the queue-storage tables + /// autovacuum never analyses (default: 5 minutes). `Duration::ZERO` + /// disables the refresh. + pub fn statistics_refresh_interval(mut self, interval: Duration) -> Self { + self.statistics_refresh_interval = interval; + self + } + /// How long a descriptor catalog row can go without being re-synced /// before the maintenance leader deletes it (default: 30 days). Set /// to `Duration::ZERO` to disable — useful if you maintain the catalog @@ -1121,6 +1136,7 @@ impl MaintenanceService { cron_eval_cancel.clone(), self.cron_eval_interval, )); + let _statistics_refresh = self.spawn_statistics_refresh(); loop { tokio::select! { @@ -1314,6 +1330,28 @@ impl MaintenanceService { } } + /// Start the planner-statistics refresh for this leadership term. The + /// returned guard stops it when dropped. + fn spawn_statistics_refresh(&self) -> Option { + if self.statistics_refresh_interval.is_zero() { + return None; + } + let (runtime, standby) = match self.storage.queue_storage() { + Some(runtime) => (runtime, false), + None => (self.standby_queue_storage.as_ref()?, true), + }; + Some(RefreshGuard::spawn( + StatisticsRefresher::new( + self.pool.clone(), + runtime.store.schema().to_string(), + self.metrics.clone(), + standby, + ), + self.statistics_refresh_interval, + &self.cancel, + )) + } + fn stop_cron_evaluator(cancel: &CancellationToken, task: &JoinHandle<()>) { cancel.cancel(); task.abort(); diff --git a/awa-worker/src/statistics.rs b/awa-worker/src/statistics.rs new file mode 100644 index 000000000..6d1e81a90 --- /dev/null +++ b/awa-worker/src/statistics.rs @@ -0,0 +1,646 @@ +//! Planner-statistics refresh for the queue-storage tables autovacuum never +//! analyses: the two head tables, which stay below the autoanalyze threshold, +//! and the partitioned ring parents, which autovacuum skips entirely. +//! +//! The maintenance leader runs one pass on election and one every interval. +//! Each relation is analysed through `awa.analyze_queue_storage_relation` +//! (migration v054, `SECURITY DEFINER`) in its own short transaction with a +//! bounded `lock_timeout` and `statement_timeout`, so the refresh never holds +//! locks across relations and cannot wait on ring prune for long. + +use crate::metrics::AwaMetrics; +use sqlx::PgPool; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::{Duration, Instant}; +use tokio_util::sync::CancellationToken; +use tracing::{debug, warn}; + +/// How long one relation's ANALYZE may wait for a lock before the pass moves +/// on. Ring prune holds `ACCESS EXCLUSIVE` on a slot's children only +/// briefly, so a wait this long means something else holds the table. +const LOCK_TIMEOUT: Duration = Duration::from_secs(2); + +/// Upper bound for one relation. A ring parent analyses in under a few +/// seconds even on a loaded host; this only stops a stuck statement. +const STATEMENT_TIMEOUT: Duration = Duration::from_secs(120); + +const HEAD_TABLES: [&str; 2] = ["queue_claim_heads", "queue_enqueue_heads"]; + +/// The partitioned ring parents a queue-storage substrate installs. A schema +/// can hold other partitioned tables (a substrate in an application schema), +/// which this refresh must never analyse. +const RING_PARENTS: [&str; 13] = [ + "done_entries", + "lease_claim_batches", + "lease_claim_closure_batches", + "lease_claim_closures", + "lease_claims", + "leases", + "queue_terminal_count_deltas", + "ready_claim_attempt_batches", + "ready_entries", + "ready_segments", + "ready_tombstones", + "receipt_completion_batches", + "receipt_completion_tombstones", +]; + +const SQLSTATE_INSUFFICIENT_PRIVILEGE: &str = "42501"; +const SQLSTATE_UNDEFINED_FUNCTION: &str = "42883"; +const SQLSTATE_LOCK_NOT_AVAILABLE: &str = "55P03"; + +/// Which kind of relation an analyse call covered; bounds metric cardinality. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum RelationScope { + Heads, + RingParent, + /// The pass as a whole: a probe or listing failed, or the helper is absent. + Pass, +} + +impl RelationScope { + pub(crate) fn as_str(self) -> &'static str { + match self { + RelationScope::Heads => "heads", + RelationScope::RingParent => "ring_parent", + RelationScope::Pass => "pass", + } + } +} + +/// Result of analysing one relation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AnalyzeOutcome { + Ok, + /// The function owner (or the runtime role, if it cannot execute the + /// function) may not analyse the relation. + SkippedPrivilege, + /// `lock_timeout` expired waiting on a conflicting lock. + SkippedLock, + /// `awa.analyze_queue_storage_relation` is not installed (schema older + /// than v054). + Unavailable, + Failed, +} + +impl AnalyzeOutcome { + pub(crate) fn as_str(self) -> &'static str { + match self { + AnalyzeOutcome::Ok => "ok", + AnalyzeOutcome::SkippedPrivilege => "skipped_privilege", + AnalyzeOutcome::SkippedLock => "skipped_lock", + AnalyzeOutcome::Unavailable => "unavailable", + AnalyzeOutcome::Failed => "failed", + } + } +} + +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub(crate) struct PassSummary { + pub(crate) ok: usize, + pub(crate) skipped_privilege: usize, + pub(crate) skipped_lock: usize, + pub(crate) unavailable: usize, + pub(crate) failed: usize, +} + +impl PassSummary { + fn record(&mut self, outcome: AnalyzeOutcome) { + match outcome { + AnalyzeOutcome::Ok => self.ok += 1, + AnalyzeOutcome::SkippedPrivilege => self.skipped_privilege += 1, + AnalyzeOutcome::SkippedLock => self.skipped_lock += 1, + AnalyzeOutcome::Unavailable => self.unavailable += 1, + AnalyzeOutcome::Failed => self.failed += 1, + } + } +} + +#[derive(Debug, PartialEq, Eq)] +enum HelperAccess { + Available, + Missing, + Denied, +} + +/// Refreshes planner statistics for one queue-storage schema. +pub(crate) struct StatisticsRefresher { + pool: PgPool, + schema: String, + metrics: AwaMetrics, + /// Set for a leader whose own storage is canonical and whose queue-storage + /// runtime is only a standby: a pass runs only once routing has flipped to + /// `schema`. + require_active_schema: bool, + warned_privilege: AtomicBool, + warned_unavailable: AtomicBool, +} + +impl StatisticsRefresher { + pub(crate) fn new( + pool: PgPool, + schema: String, + metrics: AwaMetrics, + require_active_schema: bool, + ) -> Self { + Self { + pool, + schema, + metrics, + require_active_schema, + warned_privilege: AtomicBool::new(false), + warned_unavailable: AtomicBool::new(false), + } + } + + /// Run a pass now and then every `interval` until `cancel` fires. + pub(crate) async fn run(&self, interval: Duration, cancel: CancellationToken) { + let mut timer = tokio::time::interval(interval); + timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + tokio::select! { + _ = cancel.cancelled() => return, + _ = timer.tick() => { + tokio::select! { + _ = cancel.cancelled() => return, + _ = self.run_pass() => {} + } + } + } + } + } + + /// Analyse the head tables, then every partitioned table of the schema. + pub(crate) async fn run_pass(&self) -> PassSummary { + let started = Instant::now(); + let mut summary = PassSummary::default(); + + if self.require_active_schema { + match sqlx::query_scalar::<_, Option>( + "SELECT awa.active_queue_storage_schema()", + ) + .fetch_one(&self.pool) + .await + { + Ok(active) if active.as_deref() == Some(self.schema.as_str()) => {} + Ok(_) => { + debug!(schema = %self.schema, "Routing has not flipped to this schema yet; skipping statistics refresh"); + return summary; + } + Err(err) => { + warn!(error = %err, "Could not resolve the active queue-storage schema for statistics refresh"); + return summary; + } + } + } + + match awa_model::storage::queue_storage_schema_ready(&self.pool, &self.schema).await { + Ok(true) => {} + Ok(false) => { + debug!(schema = %self.schema, "Queue-storage schema is not prepared yet; skipping statistics refresh"); + return summary; + } + Err(err) => { + warn!(error = %err, schema = %self.schema, "Could not check the queue-storage schema for statistics refresh"); + self.record( + RelationScope::Pass, + AnalyzeOutcome::Failed, + started.elapsed(), + ); + summary.record(AnalyzeOutcome::Failed); + return summary; + } + } + + match self.helper_access().await { + Ok(HelperAccess::Available) => {} + Ok(HelperAccess::Missing) => { + self.warn_unavailable_once(); + self.record( + RelationScope::Pass, + AnalyzeOutcome::Unavailable, + started.elapsed(), + ); + summary.record(AnalyzeOutcome::Unavailable); + return summary; + } + Ok(HelperAccess::Denied) => { + self.warn_privilege_once( + "queue_claim_heads", + "the runtime role needs EXECUTE on awa.analyze_queue_storage_relation(text, text), \ + and the login role needs INSERT, UPDATE and DELETE on this schema's queue_claim_heads; \ + restore those grants as the schema owner", + ); + self.record( + RelationScope::Pass, + AnalyzeOutcome::SkippedPrivilege, + started.elapsed(), + ); + summary.record(AnalyzeOutcome::SkippedPrivilege); + return summary; + } + Err(err) => { + warn!(error = %err, schema = %self.schema, "Could not probe access to the statistics helper"); + self.record( + RelationScope::Pass, + AnalyzeOutcome::Failed, + started.elapsed(), + ); + summary.record(AnalyzeOutcome::Failed); + return summary; + } + } + + let ring_parents = match self.ring_parents().await { + Ok(parents) => parents, + Err(err) => { + warn!(error = %err, schema = %self.schema, "Could not list ring tables for statistics refresh"); + self.record( + RelationScope::Pass, + AnalyzeOutcome::Failed, + started.elapsed(), + ); + summary.record(AnalyzeOutcome::Failed); + Vec::new() + } + }; + + let targets = HEAD_TABLES + .iter() + .map(|name| (RelationScope::Heads, (*name).to_string())) + .chain( + ring_parents + .into_iter() + .map(|name| (RelationScope::RingParent, name)), + ); + + for (scope, relation) in targets { + let relation_started = Instant::now(); + let outcome = self.analyze_relation(&relation).await; + self.record(scope, outcome, relation_started.elapsed()); + summary.record(outcome); + } + + debug!( + schema = %self.schema, + elapsed_ms = started.elapsed().as_millis() as u64, + ok = summary.ok, + skipped_privilege = summary.skipped_privilege, + skipped_lock = summary.skipped_lock, + unavailable = summary.unavailable, + failed = summary.failed, + "Planner statistics refresh finished" + ); + summary + } + + fn record(&self, scope: RelationScope, outcome: AnalyzeOutcome, elapsed: Duration) { + self.metrics + .record_statistics_analyze(scope.as_str(), outcome.as_str(), elapsed); + } + + /// Check the helper and its caller privileges once per pass, so a missing + /// migration or grant does not produce a failing call for every relation. + async fn helper_access(&self) -> Result { + let (installed, allowed): (bool, bool) = sqlx::query_as( + r#" + SELECT helper IS NOT NULL, + COALESCE( + has_function_privilege(helper, 'EXECUTE') + AND has_table_privilege(session_user, heads, 'INSERT') + AND has_table_privilege(session_user, heads, 'UPDATE') + AND has_table_privilege(session_user, heads, 'DELETE'), + false + ) + FROM ( + SELECT to_regprocedure('awa.analyze_queue_storage_relation(text,text)') AS helper, + to_regclass(format('%I.queue_claim_heads', $1::text)) AS heads + ) AS objects + "#, + ) + .bind(&self.schema) + .fetch_one(&self.pool) + .await?; + Ok(match (installed, allowed) { + (false, _) => HelperAccess::Missing, + (true, false) => HelperAccess::Denied, + (true, true) => HelperAccess::Available, + }) + } + + async fn ring_parents(&self) -> Result, sqlx::Error> { + sqlx::query_scalar( + r#" + SELECT class.relname::text + FROM pg_class AS class + JOIN pg_namespace AS namespace ON namespace.oid = class.relnamespace + WHERE namespace.nspname = $1 AND class.relkind = 'p' + AND class.relname = ANY($2) + ORDER BY class.relname + "#, + ) + .bind(&self.schema) + .bind(&RING_PARENTS[..]) + .fetch_all(&self.pool) + .await + } + + async fn analyze_relation(&self, relation: &str) -> AnalyzeOutcome { + match self.call_analyze_function(relation).await { + Ok(status) => match status.as_str() { + "analyzed" => AnalyzeOutcome::Ok, + "skipped_privilege" => { + self.warn_privilege_once( + relation, + "the owner of awa.analyze_queue_storage_relation() may not ANALYZE this table; \ + change the function's owner to the table owner with ALTER FUNCTION, \ + or grant the function owner MAINTAIN on them (PostgreSQL 17+)", + ); + AnalyzeOutcome::SkippedPrivilege + } + // The relation is not part of this schema's shape; nothing to analyse. + "missing" => AnalyzeOutcome::Ok, + other => { + warn!( + status = other, + relation, "Unexpected statistics refresh status" + ); + AnalyzeOutcome::Failed + } + }, + Err(err) => self.classify_error(relation, &err), + } + } + + async fn call_analyze_function(&self, relation: &str) -> Result { + let mut tx = self.pool.begin().await?; + sqlx::query( + "SELECT set_config('lock_timeout', $1, true), set_config('statement_timeout', $2, true)", + ) + .bind(format!("{}ms", LOCK_TIMEOUT.as_millis())) + .bind(format!("{}ms", STATEMENT_TIMEOUT.as_millis())) + .execute(&mut *tx) + .await?; + let status: String = + sqlx::query_scalar("SELECT awa.analyze_queue_storage_relation($1, $2)") + .bind(&self.schema) + .bind(relation) + .fetch_one(&mut *tx) + .await?; + tx.commit().await?; + Ok(status) + } + + fn classify_error(&self, relation: &str, err: &sqlx::Error) -> AnalyzeOutcome { + let code = match err { + sqlx::Error::Database(db_err) => db_err.code().map(|code| code.into_owned()), + _ => None, + }; + match code.as_deref() { + Some(SQLSTATE_LOCK_NOT_AVAILABLE) => { + debug!( + relation, + "Statistics refresh timed out waiting for a lock; retrying next interval" + ); + AnalyzeOutcome::SkippedLock + } + Some(SQLSTATE_UNDEFINED_FUNCTION) => { + self.warn_unavailable_once(); + AnalyzeOutcome::Unavailable + } + Some(SQLSTATE_INSUFFICIENT_PRIVILEGE) => { + self.warn_privilege_once( + relation, + "the runtime role may not execute awa.analyze_queue_storage_relation(), or lacks \ + INSERT, UPDATE and DELETE on the queue head tables; \ + GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) to it", + ); + AnalyzeOutcome::SkippedPrivilege + } + _ => { + warn!(relation, error = %err, "Planner statistics refresh failed for relation"); + AnalyzeOutcome::Failed + } + } + } + + fn warn_unavailable_once(&self) { + if self.warned_unavailable.swap(true, Ordering::Relaxed) { + return; + } + warn!( + schema = %self.schema, + "Planner statistics refresh is disabled: awa.analyze_queue_storage_relation() \ + is missing. Run `awa migrate` (schema v054 or later) to install it" + ); + } + + fn warn_privilege_once(&self, relation: &str, missing: &str) { + if self.warned_privilege.swap(true, Ordering::Relaxed) { + return; + } + warn!( + schema = %self.schema, + relation, + "Planner statistics refresh skipped: {missing}. Set the statistics refresh \ + interval to zero to disable the refresh" + ); + } +} + +/// Stops the refresh task when the maintenance leader loses leadership or +/// shuts down. +pub(crate) struct RefreshGuard { + cancel: CancellationToken, +} + +impl RefreshGuard { + pub(crate) fn spawn( + refresher: StatisticsRefresher, + interval: Duration, + parent: &CancellationToken, + ) -> Self { + let cancel = parent.child_token(); + let task_cancel = cancel.clone(); + tokio::spawn(async move { + debug!( + schema = %refresher.schema, + interval_s = interval.as_secs(), + "Planner statistics refresh started" + ); + refresher.run(interval, task_cancel).await; + }); + Self { cancel } + } +} + +impl Drop for RefreshGuard { + fn drop(&mut self) { + self.cancel.cancel(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; + use std::str::FromStr; + + /// A binary ahead of the schema must probe once and attempt no calls: a + /// failing call per relation would count as `unavailable` fifteen times + /// and write fifteen server-side ERROR lines per pass. + #[tokio::test] + async fn missing_helper_costs_one_probe_and_no_calls() { + let url = std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()); + let options = PgConnectOptions::from_str(&url).expect("parse DATABASE_URL"); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect_with(options.clone()) + .await + .expect("connect admin"); + let database = format!("awa_stats_unit_{}", std::process::id()); + sqlx::raw_sql(awa_model::audited_sql(format!( + "DROP DATABASE IF EXISTS {database} WITH (FORCE)" + ))) + .execute(&admin) + .await + .expect("drop leftover database"); + sqlx::raw_sql(awa_model::audited_sql(format!( + "CREATE DATABASE {database}" + ))) + .execute(&admin) + .await + .expect("create database"); + let pool = PgPoolOptions::new() + .max_connections(2) + .connect_with(options.database(&database)) + .await + .expect("connect test database"); + awa_model::migrations::run(&pool).await.expect("migrate"); + + let refresher = StatisticsRefresher::new( + pool.clone(), + "awa".to_string(), + AwaMetrics::from_global(), + false, + ); + let present = refresher.run_pass().await; + assert_eq!(present.ok, 15, "{present:?}"); + + sqlx::raw_sql("DROP FUNCTION awa.analyze_queue_storage_relation(text, text)") + .execute(&pool) + .await + .expect("drop the helper"); + let absent = refresher.run_pass().await; + assert_eq!( + absent, + PassSummary { + unavailable: 1, + ..PassSummary::default() + } + ); + pool.close().await; + sqlx::raw_sql(awa_model::audited_sql(format!( + "DROP DATABASE IF EXISTS {database} WITH (FORCE)" + ))) + .execute(&admin) + .await + .expect("drop database"); + } + + /// Missing EXECUTE or login-role DML is a single skipped pass, rather + /// than fifteen server-side permission errors. + #[tokio::test] + async fn missing_helper_privileges_skip_the_pass_without_calling_it() { + let url = std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()); + let options = PgConnectOptions::from_str(&url).expect("parse DATABASE_URL"); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect_with(options.clone()) + .await + .expect("connect admin"); + let database = format!("awa_stats_access_unit_{}", std::process::id()); + let role = format!("awa_stats_access_{}", std::process::id()); + sqlx::raw_sql(awa_model::audited_sql(format!( + "DROP DATABASE IF EXISTS {database} WITH (FORCE)" + ))) + .execute(&admin) + .await + .expect("drop leftover database"); + sqlx::raw_sql(awa_model::audited_sql(format!("DROP ROLE IF EXISTS {role}; CREATE ROLE {role} LOGIN PASSWORD 'statistics_access_test'"))) + .execute(&admin) + .await + .expect("create test role"); + sqlx::raw_sql(awa_model::audited_sql(format!( + "CREATE DATABASE {database}" + ))) + .execute(&admin) + .await + .expect("create database"); + let owner = PgPoolOptions::new() + .max_connections(2) + .connect_with(options.clone().database(&database)) + .await + .expect("connect owner"); + awa_model::migrations::run(&owner).await.expect("migrate"); + sqlx::raw_sql(awa_model::audited_sql(format!( + "GRANT USAGE ON SCHEMA awa TO {role}; \ + GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA awa TO {role}; \ + REVOKE EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) FROM {role}" + ))) + .execute(&owner) + .await + .expect("grant head-table privileges only"); + let runtime = PgPoolOptions::new() + .max_connections(2) + .connect_with( + options + .database(&database) + .username(&role) + .password("statistics_access_test"), + ) + .await + .expect("connect runtime"); + let refresher = StatisticsRefresher::new( + runtime.clone(), + "awa".to_string(), + AwaMetrics::from_global(), + false, + ); + let denied = PassSummary { + skipped_privilege: 1, + ..PassSummary::default() + }; + assert_eq!(refresher.run_pass().await, denied); + sqlx::raw_sql(awa_model::audited_sql(format!( + "GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) TO {role}" + ))) + .execute(&owner) + .await + .expect("grant helper EXECUTE"); + assert_eq!(refresher.run_pass().await.ok, 15); + sqlx::raw_sql(awa_model::audited_sql(format!( + "REVOKE INSERT ON awa.queue_claim_heads FROM {role}" + ))) + .execute(&owner) + .await + .expect("revoke login-role head privilege"); + assert_eq!(refresher.run_pass().await, denied); + + runtime.close().await; + owner.close().await; + sqlx::raw_sql(awa_model::audited_sql(format!( + "DROP DATABASE {database} WITH (FORCE)" + ))) + .execute(&admin) + .await + .expect("drop database"); + sqlx::raw_sql(awa_model::audited_sql(format!("DROP ROLE {role}"))) + .execute(&admin) + .await + .expect("drop test role"); + admin.close().await; + } +} diff --git a/awa/tests/enqueue_lane_contention_test.rs b/awa/tests/enqueue_lane_contention_test.rs new file mode 100644 index 000000000..cd6e05b84 --- /dev/null +++ b/awa/tests/enqueue_lane_contention_test.rs @@ -0,0 +1,809 @@ +//! Per-lane enqueue contention: many concurrent producers batch-inserting +//! into one queue lane. +//! +//! Every enqueue on queue storage reserves its lane sequence under a +//! transaction-scoped advisory lock (`reserve_enqueue_seq`) that is held until +//! the enqueue transaction ends, so producers for one `(queue, priority, +//! shard)` lane serialise on it. This suite measures what that serialisation +//! costs: lane throughput and batch-insert latency at high producer +//! concurrency, with and without unique keys, with and without a unique +//! conflict fallback, and across `enqueue_shards`. +//! +//! * `test_enqueue_lane_contention_gate` runs a fixed set of scenarios, checks +//! lane integrity, asserts contention-scaled floors, and emits +//! `@@BENCH_JSON@@` records that `.github/scripts/check-bench-regression.py` +//! compares against `benchmarks/baseline.json`. +//! * `test_enqueue_lane_contention_matrix` is the exploration entry point: it +//! sweeps producer counts for one configuration chosen through environment +//! variables and prints every scenario. +//! +//! Matrix tunables (all optional): +//! AWA_LANE_BENCH_PRODUCERS comma list, default "1,4,8,16,32" +//! AWA_LANE_BENCH_BATCH jobs per batch, default 185 +//! AWA_LANE_BENCH_UNIQUE 0|1, default 1 +//! AWA_LANE_BENCH_SHARDS `enqueue_shards` for the queue, default 1 +//! AWA_LANE_BENCH_MODE insert_many | insert_many_copy | insert_many_skip | +//! queue_storage | legacy_row_function +//! AWA_LANE_BENCH_DUP_PCT % of batches carrying one duplicate key that +//! fall back to row-at-a-time inserts, default 0 +//! AWA_LANE_BENCH_SECS measured seconds per scenario, default 8 +//! AWA_LANE_BENCH_DRAIN 0|1: claim everything that landed afterwards and +//! check every shard drains +//! AWA_LANE_BENCH_SPLIT 0|1: run each batch in an explicit +//! transaction and time statement and commit +//! separately (insert_many mode only) +//! +//! Requires a running Postgres instance. +//! Run with: +//! cargo test --release --package awa --test enqueue_lane_contention_test \ +//! -- --ignored --nocapture --test-threads=1 + +mod bench_output; +mod ci_timing; + +use awa::model::insert_prepared; +use awa::model::{ + insert_many, insert_many_copy_from_pool, insert_many_skip_unique_conflicts, + prepare_raw_job_insert, AwaError, QueueStorage, +}; +use awa::{InsertOpts, InsertParams, QueueStorageConfig, UniqueOpts}; +use awa_testing::setup; +use bench_output::{BenchLatency, BenchMetrics, BenchmarkResult, SCHEMA_VERSION}; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; +use sqlx::PgPool; +use std::collections::HashMap; +use std::str::FromStr; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use uuid::Uuid; + +const APPLICATION_NAME: &str = "awa_lane_bench"; +const WARMUP: Duration = Duration::from_millis(1_500); + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Mode { + InsertMany, + InsertManyCopy, + InsertManySkip, + QueueStorage, + LegacyRowFunction, +} + +impl Mode { + fn parse(raw: &str) -> Self { + match raw { + "insert_many" => Mode::InsertMany, + "insert_many_copy" => Mode::InsertManyCopy, + "insert_many_skip" => Mode::InsertManySkip, + "queue_storage" => Mode::QueueStorage, + "legacy_row_function" => Mode::LegacyRowFunction, + other => panic!("unknown AWA_LANE_BENCH_MODE {other}"), + } + } + + fn label(self) -> &'static str { + match self { + Mode::InsertMany => "insert_many", + Mode::InsertManyCopy => "insert_many_copy", + Mode::InsertManySkip => "insert_many_skip", + Mode::QueueStorage => "queue_storage", + Mode::LegacyRowFunction => "legacy_row_function", + } + } +} + +#[derive(Debug, Clone)] +struct Scenario { + name: String, + mode: Mode, + producers: usize, + batch_size: usize, + unique: bool, + shards: i16, + duplicate_batch_pct: u32, + measured: Duration, + split_commit: bool, + drain: bool, +} + +#[derive(Debug, Default, Clone)] +struct Samples { + /// Latency of the producer call, in milliseconds: one `insert_many` + /// statement, or for a duplicate batch the failed batch plus the + /// row-at-a-time fallback. + batch_ms: Vec, + /// Latency of each single-row statement issued by the fallback. + fallback_row_ms: Vec, + /// Statement and commit latency when the batch runs in an explicit + /// transaction. + statement_ms: Vec, + commit_ms: Vec, + jobs: u64, + batches: u64, + duplicate_batches: u64, + fallback_rows: u64, +} + +impl Samples { + fn merge(&mut self, other: Samples) { + self.batch_ms.extend(other.batch_ms); + self.fallback_row_ms.extend(other.fallback_row_ms); + self.statement_ms.extend(other.statement_ms); + self.commit_ms.extend(other.commit_ms); + self.jobs += other.jobs; + self.batches += other.batches; + self.duplicate_batches += other.duplicate_batches; + self.fallback_rows += other.fallback_rows; + } +} + +#[derive(Debug)] +struct ScenarioResult { + scenario: Scenario, + jobs_per_s: f64, + p50_ms: f64, + p95_ms: f64, + p99_ms: f64, + max_ms: f64, + mean_waiters: f64, + samples: Samples, + landed: i64, + distinct_lane_seqs: i64, + drained: Option<(i64, f64)>, +} + +fn percentile(sorted: &[f64], quantile: f64) -> f64 { + if sorted.is_empty() { + return 0.0; + } + let index = ((sorted.len() as f64 - 1.0) * quantile).round() as usize; + sorted[index.min(sorted.len() - 1)] +} + +fn sorted(values: &[f64]) -> Vec { + let mut copy = values.to_vec(); + copy.sort_by(|left, right| left.partial_cmp(right).expect("latencies are finite")); + copy +} + +fn env_string(name: &str, default: &str) -> String { + std::env::var(name).unwrap_or_else(|_| default.to_string()) +} + +fn env_usize(name: &str, default: usize) -> usize { + std::env::var(name) + .ok() + .and_then(|raw| raw.parse().ok()) + .unwrap_or(default) +} + +fn job_params(queue: &str, tag: &str, producer: usize, counter: u64, unique: bool) -> InsertParams { + InsertParams { + kind: "lane_bench_job".to_string(), + args: serde_json::json!({ + "tag": tag, + "producer": producer, + "n": counter, + "payload": "x".repeat(64), + }), + opts: InsertOpts { + queue: queue.to_string(), + unique: unique.then(|| UniqueOpts { + by_queue: true, + by_args: true, + ..Default::default() + }), + ..Default::default() + }, + } +} + +fn is_unique_violation(err: &AwaError) -> bool { + match err { + AwaError::UniqueConflict { .. } => true, + AwaError::Database(sqlx::Error::Database(db_err)) => { + db_err.code().as_deref() == Some("23505") + } + _ => false, + } +} + +async fn bench_pool(max_connections: u32) -> PgPool { + let options = PgConnectOptions::from_str(&setup::database_url()) + .expect("parse DATABASE_URL") + .application_name(APPLICATION_NAME); + PgPoolOptions::new() + .max_connections(max_connections) + .connect_with(options) + .await + .expect("connect bench pool") +} + +async fn waiter_count(pool: &PgPool) -> i64 { + sqlx::query_scalar( + r#" + SELECT count(*)::bigint + FROM pg_stat_activity + WHERE application_name = $1 + AND wait_event_type = 'Lock' + AND wait_event = 'advisory' + "#, + ) + .bind(APPLICATION_NAME) + .fetch_one(pool) + .await + .unwrap_or(0) +} + +async fn seed_queue_shards(pool: &PgPool, queue: &str, shards: i16) { + sqlx::query( + "INSERT INTO awa.queue_meta (queue, enqueue_shards) VALUES ($1, $2) + ON CONFLICT (queue) DO UPDATE SET enqueue_shards = EXCLUDED.enqueue_shards", + ) + .bind(queue) + .bind(shards) + .execute(pool) + .await + .expect("seed queue_meta.enqueue_shards"); +} + +/// One statement that calls `awa.insert_job_compat` once per row, the shape +/// batch inserts had before `awa.insert_jobs_compat_batch`. It measures what a +/// batch costs the lane lock on this runner when every row goes through the +/// per-row function, which is the yardstick the gate compares `insert_many` +/// against. +const LEGACY_ROW_FUNCTION_SQL: &str = r#" + SELECT inserted.* + FROM unnest( + $1::text[], $2::text[], $3::jsonb[], $4::text[], $5::smallint[], $6::smallint[], + $7::timestamptz[], $8::jsonb[], $9::jsonb[], $10::bytea[], $11::text[] + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + CROSS JOIN LATERAL awa.insert_job_compat( + input.kind, + input.queue, + input.args, + input.state::awa.job_state, + input.priority, + input.max_attempts, + input.run_at, + input.metadata, + ARRAY( + SELECT tag.value + FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, tag_ord) + ORDER BY tag.tag_ord + ), + input.unique_key, + input.unique_states::bit(8), + NULL + ) AS inserted + ORDER BY input.ord +"#; + +async fn legacy_row_function_insert(pool: &PgPool, jobs: &[InsertParams]) -> Result<(), AwaError> { + let rows: Vec<_> = jobs + .iter() + .map(|job| { + prepare_raw_job_insert(job.kind.clone(), job.args.clone(), job.opts.clone()) + .expect("prepare row") + }) + .collect(); + sqlx::query(LEGACY_ROW_FUNCTION_SQL) + .bind(rows.iter().map(|row| row.kind()).collect::>()) + .bind(rows.iter().map(|row| row.queue()).collect::>()) + .bind(rows.iter().map(|row| row.args()).collect::>()) + .bind( + rows.iter() + .map(|row| row.state_db_str()) + .collect::>(), + ) + .bind(rows.iter().map(|row| row.priority()).collect::>()) + .bind( + rows.iter() + .map(|row| row.max_attempts()) + .collect::>(), + ) + .bind(rows.iter().map(|row| row.run_at()).collect::>()) + .bind(rows.iter().map(|row| row.metadata()).collect::>()) + .bind( + rows.iter() + .map(|row| serde_json::json!(row.tags())) + .collect::>(), + ) + .bind(rows.iter().map(|row| row.unique_key()).collect::>()) + .bind( + rows.iter() + .map(|row| row.unique_states_bit_string()) + .collect::>(), + ) + .fetch_all(pool) + .await?; + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +async fn producer_loop( + pool: PgPool, + store: Arc, + scenario: Scenario, + queue: String, + tag: String, + producer: usize, + measure_from: Instant, + stop_at: Instant, + counter_base: Arc, + stop: Arc, +) -> Samples { + let mut samples = Samples::default(); + let mut batch_index: u64 = 0; + let mut previous_batch: Vec = Vec::new(); + + while Instant::now() < stop_at && !stop.load(Ordering::Relaxed) { + let base = counter_base.fetch_add(scenario.batch_size as u64, Ordering::Relaxed); + let mut jobs: Vec = (0..scenario.batch_size as u64) + .map(|offset| job_params(&queue, &tag, producer, base + offset, scenario.unique)) + .collect(); + + // A replayed key: the first job repeats one this producer already + // inserted, so the batch hits exactly one unique conflict. + let duplicate = scenario.unique + && scenario.duplicate_batch_pct > 0 + && !previous_batch.is_empty() + && (batch_index + 1) * u64::from(scenario.duplicate_batch_pct) / 100 + > batch_index * u64::from(scenario.duplicate_batch_pct) / 100; + if duplicate { + jobs[0] = previous_batch[0].clone(); + } + batch_index += 1; + + let started = Instant::now(); + let mut statement_ms = None; + let mut commit_ms = None; + let mut fallback_rows_ms: Vec = Vec::new(); + + let outcome: Result<(), AwaError> = match scenario.mode { + Mode::InsertMany if scenario.split_commit => { + let mut tx = pool.begin().await.expect("begin"); + let statement_started = Instant::now(); + let result = insert_many(&mut *tx, &jobs).await; + let statement_elapsed = statement_started.elapsed(); + match result { + Ok(_) => { + let commit_started = Instant::now(); + tx.commit().await.expect("commit"); + statement_ms = Some(statement_elapsed.as_secs_f64() * 1e3); + commit_ms = Some(commit_started.elapsed().as_secs_f64() * 1e3); + Ok(()) + } + Err(err) => Err(err), + } + } + Mode::InsertMany => insert_many(&pool, &jobs).await.map(|_| ()), + Mode::InsertManyCopy => insert_many_copy_from_pool(&pool, &jobs).await.map(|_| ()), + Mode::InsertManySkip => insert_many_skip_unique_conflicts(&pool, &jobs) + .await + .map(|_| ()), + Mode::QueueStorage => store.enqueue_params_batch(&pool, &jobs).await.map(|_| ()), + Mode::LegacyRowFunction => legacy_row_function_insert(&pool, &jobs).await, + }; + + let mut fallback_count = 0u64; + match outcome { + Ok(()) => {} + Err(err) if is_unique_violation(&err) => { + // What a replaying producer does on a unique conflict: + // insert the batch again one row at a time, skipping + // duplicates. + for job in &jobs { + let row = prepare_raw_job_insert( + job.kind.clone(), + job.args.clone(), + job.opts.clone(), + ) + .expect("prepare row"); + let row_started = Instant::now(); + match insert_prepared(&pool, &row).await { + Ok(_) => {} + Err(err) if is_unique_violation(&err) => {} + Err(err) => panic!("fallback insert failed: {err}"), + } + fallback_rows_ms.push(row_started.elapsed().as_secs_f64() * 1e3); + fallback_count += 1; + } + } + Err(err) => panic!("enqueue failed: {err}"), + } + let elapsed = started.elapsed(); + + if started >= measure_from { + samples.batch_ms.push(elapsed.as_secs_f64() * 1e3); + samples.jobs += scenario.batch_size as u64; + samples.batches += 1; + if let Some(ms) = statement_ms { + samples.statement_ms.push(ms); + } + if let Some(ms) = commit_ms { + samples.commit_ms.push(ms); + } + if fallback_count > 0 { + samples.duplicate_batches += 1; + samples.fallback_rows += fallback_count; + samples.fallback_row_ms.extend(fallback_rows_ms); + } + } + previous_batch = jobs; + } + samples +} + +async fn run_scenario(pool: &PgPool, scenario: &Scenario) -> ScenarioResult { + let queue = format!( + "lane_bench_{}", + Uuid::new_v4().simple().to_string().get(..12).unwrap() + ); + let tag = Uuid::new_v4().simple().to_string(); + seed_queue_shards(pool, &queue, scenario.shards).await; + + let producer_pool = bench_pool(scenario.producers as u32 + 4).await; + let store = Arc::new(QueueStorage::new(QueueStorageConfig::default()).expect("queue storage")); + + let started = Instant::now(); + let measure_from = started + WARMUP; + let stop_at = measure_from + scenario.measured; + let counter_base = Arc::new(AtomicU64::new(0)); + let stop = Arc::new(AtomicBool::new(false)); + + let mut tasks = Vec::new(); + for producer in 0..scenario.producers { + tasks.push(tokio::spawn(producer_loop( + producer_pool.clone(), + Arc::clone(&store), + scenario.clone(), + queue.clone(), + tag.clone(), + producer, + measure_from, + stop_at, + Arc::clone(&counter_base), + Arc::clone(&stop), + ))); + } + + // Sample advisory-lock waiters over the measured window only. + tokio::time::sleep_until(tokio::time::Instant::from_std(measure_from)).await; + let mut waiter_samples: Vec = Vec::new(); + while Instant::now() < stop_at { + waiter_samples.push(waiter_count(pool).await); + tokio::time::sleep(Duration::from_millis(25)).await; + } + + let mut samples = Samples::default(); + for task in tasks { + samples.merge(task.await.expect("producer task")); + } + producer_pool.close().await; + + let (landed, distinct_lane_seqs): (i64, i64) = sqlx::query_as( + r#" + SELECT count(*)::bigint, + count(DISTINCT (enqueue_shard, lane_seq))::bigint + FROM awa.ready_entries + WHERE queue = $1 + "#, + ) + .bind(&queue) + .fetch_one(pool) + .await + .expect("count landed rows"); + + // Claim everything that landed, from every shard, with the claim path the + // workers use. + let drained = if scenario.drain { + let drain_started = Instant::now(); + let mut claimed = 0i64; + loop { + let batch = store + .claim_runtime_batch(pool, &queue, 500, Duration::from_secs(3_600)) + .await + .expect("claim batch"); + if batch.is_empty() { + break; + } + claimed += batch.len() as i64; + } + Some((claimed, drain_started.elapsed().as_secs_f64())) + } else { + None + }; + + let latencies = sorted(&samples.batch_ms); + let window = scenario.measured.as_secs_f64(); + ScenarioResult { + scenario: scenario.clone(), + jobs_per_s: samples.jobs as f64 / window, + p50_ms: percentile(&latencies, 0.50), + p95_ms: percentile(&latencies, 0.95), + p99_ms: percentile(&latencies, 0.99), + max_ms: latencies.last().copied().unwrap_or(0.0), + mean_waiters: if waiter_samples.is_empty() { + 0.0 + } else { + waiter_samples.iter().sum::() as f64 / waiter_samples.len() as f64 + }, + samples, + landed, + distinct_lane_seqs, + drained, + } +} + +fn print_result(result: &ScenarioResult) { + let scenario = &result.scenario; + println!( + "[lane-bench] {name} mode={mode} producers={producers} batch={batch} unique={unique} \ + shards={shards} dup_pct={dup} | {rate:.0} jobs/s | batch ms p50={p50:.1} p95={p95:.1} \ + p99={p99:.1} max={max:.1} | waiters_mean={waiters:.1} | landed={landed} \ + distinct_seqs={distinct}", + name = scenario.name, + mode = scenario.mode.label(), + producers = scenario.producers, + batch = scenario.batch_size, + unique = scenario.unique, + shards = scenario.shards, + dup = scenario.duplicate_batch_pct, + rate = result.jobs_per_s, + p50 = result.p50_ms, + p95 = result.p95_ms, + p99 = result.p99_ms, + max = result.max_ms, + waiters = result.mean_waiters, + landed = result.landed, + distinct = result.distinct_lane_seqs, + ); + if !result.samples.statement_ms.is_empty() { + let statement = sorted(&result.samples.statement_ms); + let commit = sorted(&result.samples.commit_ms); + println!( + "[lane-bench] split: statement ms p50={:.1} p99={:.1} | commit ms p50={:.2} p99={:.2}", + percentile(&statement, 0.5), + percentile(&statement, 0.99), + percentile(&commit, 0.5), + percentile(&commit, 0.99), + ); + } + if let Some((claimed, seconds)) = result.drained { + println!( + "[lane-bench] drain: claimed {claimed} of {} landed in {seconds:.2}s ({:.0}/s)", + result.landed, + claimed as f64 / seconds.max(1e-9), + ); + } + if result.samples.duplicate_batches > 0 { + let rows = sorted(&result.samples.fallback_row_ms); + println!( + "[lane-bench] fallback: {} duplicate batches, {} single-row statements, \ + row ms p50={:.1} p99={:.1} max={:.1}", + result.samples.duplicate_batches, + result.samples.fallback_rows, + percentile(&rows, 0.5), + percentile(&rows, 0.99), + rows.last().copied().unwrap_or(0.0), + ); + } +} + +fn emit_json(result: &ScenarioResult) { + let mut outcomes = HashMap::new(); + outcomes.insert("jobs".to_string(), result.samples.jobs); + outcomes.insert("batches".to_string(), result.samples.batches); + BenchmarkResult { + schema_version: SCHEMA_VERSION, + scenario: result.scenario.name.clone(), + language: "rust".to_string(), + seeded: result.landed.max(0) as u64, + metrics: BenchMetrics { + throughput: None, + enqueue_per_s: Some(result.jobs_per_s), + drain_time_s: None, + latency_ms: Some(BenchLatency { + p50: Some(result.p50_ms), + p95: Some(result.p95_ms), + p99: Some(result.p99_ms), + }), + rescue: None, + }, + outcomes, + metadata: Some(serde_json::json!({ + "producers": result.scenario.producers, + "batch_size": result.scenario.batch_size, + "unique": result.scenario.unique, + "shards": result.scenario.shards, + "mean_lock_waiters": result.mean_waiters, + })), + } + .emit(); +} + +/// Lane integrity: every job that was acknowledged landed exactly once, under +/// a distinct `(shard, lane_seq)`, and every shard's sequence numbers are +/// dense (no producer lost or duplicated a reservation). +fn assert_lane_integrity(result: &ScenarioResult) { + if let Some((claimed, _)) = result.drained { + assert_eq!( + claimed, result.landed, + "{}: claim path did not drain every shard", + result.scenario.name + ); + } + assert_eq!( + result.landed, result.distinct_lane_seqs, + "{}: ready rows share a (shard, lane_seq)", + result.scenario.name + ); + assert!( + result.landed as u64 >= result.samples.jobs, + "{}: {} jobs acknowledged in the window but only {} rows landed", + result.scenario.name, + result.samples.jobs, + result.landed + ); +} + +async fn bench_setup() -> PgPool { + let pool = setup::setup(16).await; + setup::activate_queue_storage(&pool).await; + pool +} + +fn scenario( + name: &str, + producers: usize, + shards: i16, + unique: bool, + duplicate_batch_pct: u32, +) -> Scenario { + Scenario { + name: name.to_string(), + mode: Mode::InsertManySkip, + producers, + batch_size: 185, + unique, + shards, + duplicate_batch_pct, + measured: Duration::from_secs(8), + split_commit: false, + drain: false, + } +} + +/// Fixed scenarios whose results gate regressions in the nightly run. +/// +/// Absolute floors sit well below the rates a shared runner sustains, and the +/// latency ceiling well above, so they only catch a lane that has stopped +/// making progress; `benchmarks/baseline.json` carries the regression +/// thresholds on rate and p99. What the gate asserts about lock hold time is +/// relative, because it must hold on any runner: one producer's 185-row batch +/// has to take well under the time the same rows take through the per-row +/// function, measured in the same run, and sixteen producers on one lane have +/// to out-run what the per-row function could ever sustain for it. Every +/// bound goes through `ci_timing` so a contended runner loosens it. +#[tokio::test(flavor = "multi_thread", worker_threads = 6)] +#[ignore] +async fn test_enqueue_lane_contention_gate() { + let pool = bench_setup().await; + + let mut row_function = scenario("enqueue_lane_1p_row_function_185", 1, 1, true, 0); + row_function.mode = Mode::LegacyRowFunction; + row_function.measured = Duration::from_secs(4); + let row_function = run_scenario(&pool, &row_function).await; + print_result(&row_function); + assert_lane_integrity(&row_function); + + let mut one_producer = scenario("enqueue_lane_1p_unique_185", 1, 1, true, 0); + one_producer.measured = Duration::from_secs(4); + let one_producer = run_scenario(&pool, &one_producer).await; + print_result(&one_producer); + assert_lane_integrity(&one_producer); + + let single_lane = run_scenario( + &pool, + &scenario("enqueue_lane_16p_unique_185", 16, 1, true, 0), + ) + .await; + print_result(&single_lane); + emit_json(&single_lane); + assert_lane_integrity(&single_lane); + + let sharded = run_scenario( + &pool, + &scenario("enqueue_lane_16p_unique_185_shards4", 16, 4, true, 0), + ) + .await; + print_result(&sharded); + emit_json(&sharded); + assert_lane_integrity(&sharded); + + let floor = ci_timing::contention_floor(GATE_MIN_JOBS_PER_S); + assert!( + single_lane.jobs_per_s >= floor as f64, + "single-lane enqueue rate {:.0}/s fell below the floor of {floor}/s", + single_lane.jobs_per_s + ); + let p99_ceiling = ci_timing::scaled_timeout(Duration::from_millis(GATE_MAX_P99_MS)); + assert!( + Duration::from_secs_f64(single_lane.p99_ms / 1e3) <= p99_ceiling, + "single-lane p99 batch latency {:.0}ms exceeded {:?}", + single_lane.p99_ms, + p99_ceiling + ); + + // The batch holds the lane lock for a fraction of the per-row function's + // time, and the lane's ceiling follows. Both bounds are ratios of numbers + // measured moments apart on the same machine. + let max_hold_ratio = (GATE_MAX_HOLD_RATIO * ci_timing::chaos_timeout_multiplier()).min(0.9); + let hold_ratio = one_producer.p50_ms / row_function.p50_ms; + assert!( + hold_ratio <= max_hold_ratio, + "a 185-row batch took {:.1}ms against {:.1}ms through the per-row function \ + (ratio {hold_ratio:.2}, limit {max_hold_ratio:.2})", + one_producer.p50_ms, + row_function.p50_ms + ); + let min_scaling = GATE_MIN_LANE_SCALING / ci_timing::chaos_timeout_multiplier(); + let lane_scaling = single_lane.jobs_per_s / row_function.jobs_per_s; + assert!( + lane_scaling >= min_scaling.max(1.0), + "16 producers on one lane sustain {:.0}/s, {lane_scaling:.1}x the per-row function's \ + single-producer {:.0}/s (need {:.1}x)", + single_lane.jobs_per_s, + row_function.jobs_per_s, + min_scaling.max(1.0) + ); +} + +const GATE_MIN_JOBS_PER_S: i64 = 2_000; +const GATE_MAX_P99_MS: u64 = 15_000; +const GATE_MAX_HOLD_RATIO: f64 = 0.4; +const GATE_MIN_LANE_SCALING: f64 = 2.0; + +/// Exploration sweep; see the module docs for the tunables. +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore] +async fn test_enqueue_lane_contention_matrix() { + let pool = bench_setup().await; + let producers: Vec = env_string("AWA_LANE_BENCH_PRODUCERS", "1,4,8,16,32") + .split(',') + .filter_map(|raw| raw.trim().parse().ok()) + .collect(); + let mode = Mode::parse(&env_string("AWA_LANE_BENCH_MODE", "insert_many")); + let batch_size = env_usize("AWA_LANE_BENCH_BATCH", 185); + let unique = env_usize("AWA_LANE_BENCH_UNIQUE", 1) != 0; + let shards = env_usize("AWA_LANE_BENCH_SHARDS", 1) as i16; + let duplicate_pct = env_usize("AWA_LANE_BENCH_DUP_PCT", 0) as u32; + let secs = env_usize("AWA_LANE_BENCH_SECS", 8) as u64; + let split = env_usize("AWA_LANE_BENCH_SPLIT", 0) != 0; + + for producer_count in producers { + let result = run_scenario( + &pool, + &Scenario { + name: format!("matrix_{producer_count}p"), + mode, + producers: producer_count, + batch_size, + unique, + shards, + duplicate_batch_pct: duplicate_pct, + measured: Duration::from_secs(secs), + split_commit: split, + drain: env_usize("AWA_LANE_BENCH_DRAIN", 0) != 0, + }, + ) + .await; + print_result(&result); + assert_lane_integrity(&result); + } +} diff --git a/awa/tests/enqueue_lane_lock_order_test.rs b/awa/tests/enqueue_lane_lock_order_test.rs index 924d212ae..866268f5e 100644 --- a/awa/tests/enqueue_lane_lock_order_test.rs +++ b/awa/tests/enqueue_lane_lock_order_test.rs @@ -6,10 +6,11 @@ //! enqueue must take its lane locks in one canonical order: `(queue, priority, //! shard)` ascending, with queue names compared bytewise. //! -//! The deadlock is driven deterministically: a gate lock stalls the +//! The lane-order interleaving is driven deterministically: a gate lock stalls the //! `QueueStorage` batch between its first and second lane, a second producer //! with the lanes in the opposite input order is then observed waiting on the //! first lane's lock through `pg_locks`, and only then is the gate released. +//! A reversed producer parks on the gate and fails the expected first-lane wait. //! //! Requires a running Postgres instance. //! Run with: `cargo test --package awa --test enqueue_lane_lock_order_test -- --nocapture` @@ -103,16 +104,7 @@ async fn reservation_lock_key(pool: &PgPool, queue: &str) -> i64 { .expect("derive reservation lock key") } -/// Key of the advisory lock `QueueStorage` takes while it provisions a lane. -async fn provisioning_lock_key(pool: &PgPool, queue: &str) -> i64 { - sqlx::query_scalar("SELECT hashtextextended($1, 0)") - .bind(format!("awa:{queue}:2:0")) - .fetch_one(pool) - .await - .expect("derive provisioning lock key") -} - -async fn wait_for_advisory_waiter(pool: &PgPool, key: i64) { +async fn wait_for_advisory_waiter(pool: &PgPool, key: i64, pid: i32) { let started = Instant::now(); loop { let waiting: bool = sqlx::query_scalar( @@ -124,10 +116,12 @@ async fn wait_for_advisory_waiter(pool: &PgPool, key: i64) { AND NOT granted AND database = (SELECT oid FROM pg_database WHERE datname = current_database()) AND ((classid::bigint << 32) | objid::bigint) = $1 + AND pid = $2 ) "#, ) .bind(key) + .bind(pid) .fetch_one(pool) .await .expect("poll pg_locks"); @@ -136,7 +130,7 @@ async fn wait_for_advisory_waiter(pool: &PgPool, key: i64) { } assert!( started.elapsed() < WAIT_DEADLINE, - "no backend started waiting for advisory lock {key}" + "backend {pid} did not start waiting for advisory lock {key}" ); tokio::time::sleep(Duration::from_millis(10)).await; } @@ -164,21 +158,26 @@ async fn assert_producer_does_not_deadlock_with_queue_storage_batch(producer: Pr provision_lanes(&pool, &config, &[&first_queue, &second_queue]).await; let first_reservation = reservation_lock_key(&pool, &first_queue).await; - let second_provisioning = provisioning_lock_key(&pool, &second_queue).await; + let second_reservation = reservation_lock_key(&pool, &second_queue).await; // The gate stalls the QueueStorage batch after it has reserved its first // lane and before it reaches its second. let mut gate = pool.begin().await.expect("begin gate transaction"); sqlx::query("SELECT pg_advisory_xact_lock($1)") - .bind(second_provisioning) + .bind(second_reservation) .execute(&mut *gate) .await .expect("take gate lock"); - // A fresh store has an empty lane cache, so it provisions each lane under - // the lock the gate holds. + // Both folded and fallback enqueues reserve these pre-provisioned lanes. + // The folded path skips provisioning locks when their heads already exist. + let storage_pool = setup::pool(1).await; + let storage_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&storage_pool) + .await + .expect("storage backend PID"); let storage_batch = { - let pool = pool.clone(); + let pool = storage_pool.clone(); let config = config.clone(); let jobs = jobs_for(Producer::InsertMany, &[&first_queue, &second_queue]); tokio::spawn(async move { @@ -186,7 +185,7 @@ async fn assert_producer_does_not_deadlock_with_queue_storage_batch(producer: Pr store.enqueue_params_batch(&pool, &jobs).await }) }; - wait_for_advisory_waiter(&pool, second_provisioning).await; + wait_for_advisory_waiter(&pool, second_reservation, storage_pid).await; // The competing producer lists the lanes in the opposite order, and so // reaches for the first lane's reservation lock, which the batch holds. @@ -195,11 +194,16 @@ async fn assert_producer_does_not_deadlock_with_queue_storage_batch(producer: Pr .iter() .map(|params| params.opts.queue.clone()) .collect(); + let competing_pool = setup::pool(1).await; + let competing_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&competing_pool) + .await + .expect("competing backend PID"); let competing = { - let pool = pool.clone(); + let pool = competing_pool.clone(); tokio::spawn(async move { producer.run(&pool, &competing_jobs).await }) }; - wait_for_advisory_waiter(&pool, first_reservation).await; + wait_for_advisory_waiter(&pool, first_reservation, competing_pid).await; gate.commit().await.expect("release gate"); @@ -215,6 +219,8 @@ async fn assert_producer_does_not_deadlock_with_queue_storage_batch(producer: Pr storage_result.expect("QueueStorage batch must not be chosen as a deadlock victim"); let inserted = competing_result.expect("competing producer must not be chosen as a deadlock victim"); + storage_pool.close().await; + competing_pool.close().await; // Lane ordering must not leak into the result: rows come back in input order. let returned_queues: Vec = inserted.iter().map(|row| row.queue.clone()).collect(); @@ -282,6 +288,74 @@ fn ordering_key_per_shard(shards: i16) -> Vec> { .collect() } +/// Key of the advisory lock `reserve_enqueue_seq` holds for one lane. +async fn lane_reservation_lock_key(pool: &PgPool, queue: &str, priority: i16, shard: i16) -> i64 { + sqlx::query_scalar( + r#" + SELECT hashtextextended( + format('%I.%I', 'awa', awa.queue_lane_sequence_name('queue_enqueue_seq', $1, $2, $3)), + 0 + ) + "#, + ) + .bind(queue) + .bind(priority) + .bind(shard) + .fetch_one(pool) + .await + .expect("derive lane reservation lock key") +} + +/// The backend waiting for the advisory lock `key`, once there is one. +async fn advisory_waiter_pid(pool: &PgPool, key: i64) -> i32 { + let started = Instant::now(); + loop { + let pid: Option = sqlx::query_scalar( + r#" + SELECT pid + FROM pg_locks + WHERE locktype = 'advisory' + AND NOT granted + AND database = (SELECT oid FROM pg_database WHERE datname = current_database()) + AND ((classid::bigint << 32) | objid::bigint) = $1 + LIMIT 1 + "#, + ) + .bind(key) + .fetch_optional(pool) + .await + .expect("poll pg_locks"); + if let Some(pid) = pid { + return pid; + } + assert!( + started.elapsed() < WAIT_DEADLINE, + "no backend started waiting for advisory lock {key}" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } +} + +async fn granted_advisory_keys(pool: &PgPool, pid: i32) -> std::collections::BTreeSet { + sqlx::query_scalar::<_, i64>( + r#" + SELECT (classid::bigint << 32) | objid::bigint + FROM pg_locks + WHERE locktype = 'advisory' AND granted AND pid = $1 + "#, + ) + .bind(pid) + .fetch_all(pool) + .await + .expect("read granted advisory locks") + .into_iter() + .collect() +} + +/// A batch over sharded lanes reserves them in `(queue, priority, shard)` +/// order, whatever order its rows arrive in. A gate holds one lane in the +/// middle of that order; the producer must be parked on it holding exactly the +/// lanes before it. async fn assert_batch_reserves_lanes_in_canonical_order(producer: Producer) { let _guard = LANE_LOCK_TEST_GUARD.lock().await; let (pool, _config) = setup_queue_storage().await; @@ -309,6 +383,7 @@ async fn assert_batch_reserves_lanes_in_canonical_order(producer: Producer) { let keys = ordering_key_per_shard(shards); let mut jobs = Vec::new(); + let mut lanes = std::collections::BTreeSet::new(); let rounds = match producer { Producer::CopyLarge => 700, _ => 12, @@ -317,55 +392,112 @@ async fn assert_batch_reserves_lanes_in_canonical_order(producer: Producer) { for queue in &queues { let mut params = job(queue, round, false); params.opts.priority = if round % 5 == 0 { 1 } else { 2 }; - // Keyed rows spread over every shard; every third row has no key. - if round % 3 != 0 { - params.opts.ordering_key = Some(keys[(round + queue.len()) % keys.len()].clone()); - } + // Every row is keyed, so its shard, and so its lane, is known here. + let shard = (round + queue.len()) % keys.len(); + params.opts.ordering_key = Some(keys[shard].clone()); + lanes.insert(( + queue.as_bytes().to_vec(), + params.opts.priority, + shard as i16, + queue.clone(), + )); jobs.push(params); } } + assert!(lanes.len() >= 6, "the batch must span many lanes"); - let inserted = producer.run(&pool, &jobs).await.expect("insert batch"); + let canonical: Vec<(String, i16, i16)> = lanes + .into_iter() + .map(|(_, priority, shard, queue)| (queue, priority, shard)) + .collect(); + let mut canonical_keys = Vec::new(); + for (queue, priority, shard) in &canonical { + canonical_keys.push(lane_reservation_lock_key(&pool, queue, *priority, *shard).await); + } + + // Provision the lanes so the gate lock is the only thing the producer waits for. + for (index, queue) in queues.iter().enumerate() { + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + for (priority, shard) in [ + (1, 0), + (1, 1), + (1, 2), + (1, 3), + (2, 0), + (2, 1), + (2, 2), + (2, 3), + ] { + let mut params = job(queue, 10_000 + index, false); + params.opts.priority = priority; + params.opts.ordering_key = Some(keys[shard as usize].clone()); + store + .enqueue_params_batch(&pool, &[params]) + .await + .expect("provision lane"); + } + } + + let gated = canonical.len() / 2; + let mut gate = pool.begin().await.expect("begin gate transaction"); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(canonical_keys[gated]) + .execute(&mut *gate) + .await + .expect("take gate lock"); + + let competing = { + let pool = pool.clone(); + let jobs = jobs.clone(); + tokio::spawn(async move { producer.run(&pool, &jobs).await }) + }; + let waiter = advisory_waiter_pid(&pool, canonical_keys[gated]).await; + let held = granted_advisory_keys(&pool, waiter).await; + let expected: std::collections::BTreeSet = + canonical_keys[..gated].iter().copied().collect(); + assert_eq!( + held, + expected, + "a batch parked on lane {gated} of {} must hold the lanes before it and no later one", + canonical.len() + ); + + gate.commit().await.expect("release gate"); + let inserted = tokio::time::timeout(WAIT_DEADLINE, competing) + .await + .expect("producer finished") + .expect("producer task") + .expect("insert batch"); + + // Lane ordering must not leak into the result: rows come back in input order. assert_eq!(inserted.len(), jobs.len()); for (row, params) in inserted.iter().zip(&jobs) { assert_eq!(row.queue, params.opts.queue); assert_eq!(row.args, params.args); } - // Job ids are allocated as each row reaches `insert_job_compat`, so id - // order is lock order. - let lane_by_job: std::collections::HashMap = - sqlx::query_as::<_, (i64, String, i16, i16)>( - "SELECT job_id, queue, priority, enqueue_shard FROM awa.ready_entries WHERE job_id = ANY($1)", + // Within a lane, lane_seq follows input order. + let lane_by_job: std::collections::HashMap = + sqlx::query_as::<_, (i64, String, i16, i16, i64)>( + "SELECT job_id, queue, priority, enqueue_shard, lane_seq FROM awa.ready_entries WHERE job_id = ANY($1)", ) .bind(inserted.iter().map(|row| row.id).collect::>()) .fetch_all(&pool) .await .expect("read lanes") .into_iter() - .map(|(job_id, queue, priority, shard)| (job_id, (queue, priority, shard))) + .map(|(job_id, queue, priority, shard, lane_seq)| { + (job_id, (queue, priority, shard, lane_seq)) + }) .collect(); assert_eq!(lane_by_job.len(), jobs.len()); - - let mut by_lock_order: Vec<(i64, usize)> = inserted - .iter() - .enumerate() - .map(|(position, row)| (row.id, position)) - .collect(); - by_lock_order.sort_unstable(); - let lock_order: Vec<(&(String, i16, i16), usize)> = by_lock_order - .iter() - .map(|(job_id, position)| (&lane_by_job[job_id], *position)) - .collect(); - for pair in lock_order.windows(2) { - let (lane, position) = pair[0]; - let (next_lane, next_position) = pair[1]; - assert!( - lane.0.as_bytes() < next_lane.0.as_bytes() - || (lane.0 == next_lane.0 && (lane.1, lane.2) < (next_lane.1, next_lane.2)) - || (lane == next_lane && position < next_position), - "lane {lane:?} (input {position}) was locked before lane {next_lane:?} (input {next_position})" - ); + let mut last_seq_by_lane: std::collections::HashMap<(String, i16, i16), i64> = + std::collections::HashMap::new(); + for row in &inserted { + let (queue, priority, shard, lane_seq) = lane_by_job[&row.id].clone(); + if let Some(previous) = last_seq_by_lane.insert((queue, priority, shard), lane_seq) { + assert_eq!(lane_seq, previous + 1, "lane_seq must follow input order"); + } } } diff --git a/awa/tests/enqueue_lane_maintenance_coupling_test.rs b/awa/tests/enqueue_lane_maintenance_coupling_test.rs new file mode 100644 index 000000000..67a55808a --- /dev/null +++ b/awa/tests/enqueue_lane_maintenance_coupling_test.rs @@ -0,0 +1,429 @@ +//! How a saturated enqueue lane affects the maintenance leader. +//! +//! Promotion of scheduled and retryable jobs re-enqueues them into their lane, +//! so it queues on the same per-lane advisory lock as producers. A promotion +//! batch is one transaction across every queue it touches, and the leader runs +//! its branches (promotion, rescue, rotation, ...) one after another on a +//! single task, so a batch parked behind a hot lane holds up the due jobs of +//! other queues in that batch and every other maintenance branch. +//! +//! `test_maintenance_under_hot_lane` measures, with and without concurrent +//! producers on one hot lane: +//! * promotion batch latency and how long a backlog on the hot queue takes +//! to drain, +//! * promotion lag on an unrelated queue (time from a job becoming due to +//! it being visible as ready), with the promotion loop shaped like the +//! leader's (a tick, then batches until one comes back short), +//! * heartbeat rescue of abandoned claims on the hot queue: how long each run +//! takes and how late it starts once its tick is due (1s here, so a +//! blocked promotion shows up as a late start). +//! +//! Tunables (optional): +//! AWA_LANE_MAINT_PRODUCERS producers on the hot lane, default 16 +//! AWA_LANE_MAINT_BACKLOG due scheduled jobs on the hot queue, default 24576 +//! AWA_LANE_MAINT_RESCUE abandoned claims on the hot queue, default 6000 +//! +//! Requires a running Postgres instance. +//! Run with: +//! cargo test --release --package awa --test enqueue_lane_maintenance_coupling_test \ +//! -- --ignored --nocapture --test-threads=1 + +use awa::model::{insert_many, QueueStorage}; +use awa::{InsertOpts, InsertParams, JobState, QueueStorageConfig}; +use awa_testing::setup; +use sqlx::PgPool; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use uuid::Uuid; + +const PROMOTE_BATCH: i64 = 4_096; +const PROMOTE_TICK: Duration = Duration::from_millis(250); +const PROBE_PERIOD: Duration = Duration::from_millis(100); +const RESCUE_TICK: Duration = Duration::from_secs(1); + +fn env_usize(name: &str, default: usize) -> usize { + std::env::var(name) + .ok() + .and_then(|raw| raw.parse().ok()) + .unwrap_or(default) +} + +fn percentile(values: &[f64], quantile: f64) -> f64 { + if values.is_empty() { + return 0.0; + } + let mut sorted = values.to_vec(); + sorted.sort_by(|left, right| left.partial_cmp(right).expect("finite")); + sorted[((sorted.len() as f64 - 1.0) * quantile).round() as usize] +} + +fn params( + queue: &str, + tag: &str, + n: u64, + run_at: Option>, +) -> InsertParams { + InsertParams { + kind: "lane_maint_job".to_string(), + args: serde_json::json!({ "tag": tag, "n": n }), + opts: InsertOpts { + queue: queue.to_string(), + run_at, + ..Default::default() + }, + } +} + +async fn seed_due_backlog(pool: &PgPool, queue: &str, count: usize) { + let due = chrono::Utc::now() - chrono::Duration::seconds(60); + let mut inserted = 0usize; + while inserted < count { + let chunk = (count - inserted).min(2_000); + let jobs: Vec = (0..chunk) + .map(|offset| params(queue, "backlog", (inserted + offset) as u64, Some(due))) + .collect(); + insert_many(pool, &jobs).await.expect("seed backlog"); + inserted += chunk; + } +} + +async fn deferred_due(pool: &PgPool, queue: &str) -> i64 { + sqlx::query_scalar( + "SELECT count(*)::bigint FROM awa.deferred_jobs WHERE queue = $1 AND state = 'scheduled'", + ) + .bind(queue) + .fetch_one(pool) + .await + .expect("count deferred") +} + +struct Producers { + stop: Arc, + tasks: Vec>, +} + +impl Producers { + fn start(pool: &PgPool, queue: &str, count: usize) -> Self { + let stop = Arc::new(AtomicBool::new(false)); + let tag = Uuid::new_v4().simple().to_string(); + let tasks = (0..count) + .map(|producer| { + let pool = pool.clone(); + let stop = Arc::clone(&stop); + let queue = queue.to_string(); + let tag = tag.clone(); + tokio::spawn(async move { + let mut jobs_inserted = 0u64; + let mut counter = (producer as u64) << 40; + while !stop.load(Ordering::Relaxed) { + let jobs: Vec = (0..185) + .map(|offset| params(&queue, &tag, counter + offset, None)) + .collect(); + counter += 185; + insert_many(&pool, &jobs).await.expect("producer insert"); + jobs_inserted += 185; + } + jobs_inserted + }) + }) + .collect(); + Self { stop, tasks } + } + + async fn stop(self) -> u64 { + self.stop.store(true, Ordering::Relaxed); + let mut total = 0; + for task in self.tasks { + total += task.await.expect("producer task"); + } + total + } +} + +#[derive(Default)] +struct LeaderReport { + /// Promotion batches that moved a full share of the backlog. + batch_ms: Vec, + drain_s: f64, + probe_lag_ms: Vec, + /// How long after its tick each heartbeat-rescue run started, and how long + /// the run itself took. + rescue_start_delay_ms: Vec, + rescue_run_ms: Vec, +} + +/// A leader loop shaped like the maintenance leader's: one task selecting over +/// a promotion timer (250ms; on each tick batches of up to 4096 until one is +/// short) and a heartbeat-rescue timer, running each branch to completion +/// before taking the next tick. Runs until the hot queue's backlog is gone and +/// `min_window` has passed, while a probe feeds one due job per 100ms into an +/// unrelated queue and times how long each takes to become ready. +async fn run_leader( + pool: &PgPool, + store: &QueueStorage, + hot: &str, + other: &str, + min_window: Duration, +) -> LeaderReport { + let mut report = LeaderReport::default(); + let started = Instant::now(); + let stop = Arc::new(AtomicBool::new(false)); + + let probe = { + let pool = pool.clone(); + let other = other.to_string(); + let stop = Arc::clone(&stop); + tokio::spawn(async move { + let mut lags = Vec::new(); + let mut n = 0u64; + while !stop.load(Ordering::Relaxed) { + n += 1; + // Older than the hot backlog, so the probe is always in the + // first promotion batch rather than behind the backlog's + // run_at order. + let due_at = chrono::Utc::now() - chrono::Duration::seconds(120); + insert_many(&pool, &[params(&other, "probe", n, Some(due_at))]) + .await + .expect("insert probe"); + let inserted = Instant::now(); + loop { + let visible: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM awa.ready_entries + WHERE queue = $1 AND args->>'tag' = 'probe' AND (args->>'n')::bigint = $2)", + ) + .bind(&other) + .bind(n as i64) + .fetch_one(&pool) + .await + .expect("probe visibility"); + if visible { + lags.push(inserted.elapsed().as_secs_f64() * 1e3); + break; + } + if stop.load(Ordering::Relaxed) { + return lags; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + tokio::time::sleep(PROBE_PERIOD).await; + } + lags + }) + }; + + let mut promote_timer = tokio::time::interval(PROMOTE_TICK); + promote_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + let mut rescue_timer = tokio::time::interval(RESCUE_TICK); + rescue_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + promote_timer.tick().await; + rescue_timer.tick().await; + + let mut backlog_gone_at: Option = None; + loop { + tokio::select! { + tick = promote_timer.tick() => { + let _ = tick; + loop { + let batch_started = Instant::now(); + let promoted = store + .promote_due(pool, JobState::Scheduled, PROMOTE_BATCH) + .await + .expect("promote_due"); + if promoted >= 1_000 { + report.batch_ms.push(batch_started.elapsed().as_secs_f64() * 1e3); + } + if (promoted as i64) < PROMOTE_BATCH { + break; + } + } + if backlog_gone_at.is_none() && deferred_due(pool, hot).await == 0 { + backlog_gone_at = Some(Instant::now()); + } + if backlog_gone_at.is_some() && started.elapsed() >= min_window { + break; + } + } + tick = rescue_timer.tick() => { + let delay = tokio::time::Instant::now().saturating_duration_since(tick); + let run_started = Instant::now(); + let rescued = store + .rescue_stale_heartbeats(pool, Duration::ZERO) + .await + .expect("rescue_stale_heartbeats"); + if !rescued.is_empty() { + report.rescue_start_delay_ms.push(delay.as_secs_f64() * 1e3); + report.rescue_run_ms.push(run_started.elapsed().as_secs_f64() * 1e3); + } + } + } + } + report.drain_s = backlog_gone_at + .expect("backlog drained") + .duration_since(started) + .as_secs_f64(); + stop.store(true, Ordering::Relaxed); + report.probe_lag_ms = probe.await.expect("probe task"); + report +} + +fn max_of(values: &[f64]) -> f64 { + values.iter().copied().fold(0.0, f64::max) +} + +fn print_leader(label: &str, report: &LeaderReport) { + println!( + "[lane-maint] {label} promote batch ms p50={:.0} max={:.0} ({} batches) | hot backlog \ + drained in {:.2}s | unrelated-queue promotion lag ms p50={:.0} p99={:.0} max={:.0} ({} probes) \ + | rescue ran {} times, run ms p50={:.0} max={:.0}, started late by ms p50={:.0} max={:.0}", + percentile(&report.batch_ms, 0.5), + max_of(&report.batch_ms), + report.batch_ms.len(), + report.drain_s, + percentile(&report.probe_lag_ms, 0.5), + percentile(&report.probe_lag_ms, 0.99), + max_of(&report.probe_lag_ms), + report.probe_lag_ms.len(), + report.rescue_run_ms.len(), + percentile(&report.rescue_run_ms, 0.5), + max_of(&report.rescue_run_ms), + percentile(&report.rescue_start_delay_ms, 0.5), + max_of(&report.rescue_start_delay_ms), + ); +} + +/// Enqueue `count` jobs on `queue` and claim every one of them without ever +/// completing or heartbeating, so each becomes an abandoned attempt. +async fn seed_abandoned_claims(pool: &PgPool, store: &QueueStorage, queue: &str, count: usize) { + let tag = Uuid::new_v4().simple().to_string(); + let mut enqueued = 0usize; + while enqueued < count { + let chunk = (count - enqueued).min(2_000); + let jobs: Vec = (0..chunk) + .map(|offset| params(queue, &tag, (enqueued + offset) as u64, None)) + .collect(); + store + .enqueue_params_batch(pool, &jobs) + .await + .expect("enqueue claimable jobs"); + enqueued += chunk; + } + let mut claimed = 0usize; + while claimed < count { + let batch = store + .claim_runtime_batch(pool, queue, 500, Duration::from_secs(3_600)) + .await + .expect("claim batch"); + assert!(!batch.is_empty(), "ran out of jobs to claim"); + claimed += batch.len(); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore] +async fn test_maintenance_under_hot_lane() { + let producers = env_usize("AWA_LANE_MAINT_PRODUCERS", 16); + let backlog = env_usize("AWA_LANE_MAINT_BACKLOG", 24_576); + let rescue_claims = env_usize("AWA_LANE_MAINT_RESCUE", 6_000); + + let pool = setup::setup(producers as u32 + 12).await; + setup::activate_queue_storage(&pool).await; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + + for (label, producer_count) in [("idle lane ", 0usize), ("hot lane ", producers)] { + let suffix = Uuid::new_v4().simple().to_string(); + let hot = format!("lane_maint_hot_{}", &suffix[..10]); + let other = format!("lane_maint_other_{}", &suffix[..10]); + + // Warm both lanes so lane provisioning is not measured. + insert_many( + &pool, + &[ + params(&hot, "warm", 0, None), + params(&other, "warm", 0, None), + ], + ) + .await + .expect("warm lanes"); + + seed_due_backlog(&pool, &hot, backlog).await; + seed_abandoned_claims(&pool, &store, &hot, rescue_claims).await; + + let running = (producer_count > 0).then(|| Producers::start(&pool, &hot, producer_count)); + tokio::time::sleep(Duration::from_millis(1_500)).await; + + let report = run_leader(&pool, &store, &hot, &other, Duration::from_secs(12)).await; + print_leader(label, &report); + + if let Some(running) = running { + let produced = running.stop().await; + println!("[lane-maint] {label} producers inserted {produced} jobs meanwhile"); + } + } +} + +/// The mechanism without producer load: one transaction holds the hot lane's +/// reservation lock for a fixed time while the leader loop runs. Promotion of +/// a batch that contains hot-queue rows waits for it, and takes the unrelated +/// queue's due job in the same batch and the leader's other branches with it. +/// Heartbeat rescue called directly does not wait for the lane at all: it +/// writes retryable rows to the deferred table, which promotion later moves. +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore] +async fn test_leader_behind_a_held_lane_lock() { + let hold = Duration::from_millis(env_usize("AWA_LANE_MAINT_HOLD_MS", 3_000) as u64); + let pool = setup::setup(16).await; + setup::activate_queue_storage(&pool).await; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + + let suffix = Uuid::new_v4().simple().to_string(); + let hot = format!("lane_maint_hot_{}", &suffix[..10]); + let other = format!("lane_maint_other_{}", &suffix[..10]); + insert_many( + &pool, + &[ + params(&hot, "warm", 0, None), + params(&other, "warm", 0, None), + ], + ) + .await + .expect("warm lanes"); + seed_due_backlog(&pool, &hot, 8_192).await; + seed_abandoned_claims(&pool, &store, &hot, 3_000).await; + + let key: i64 = sqlx::query_scalar( + "SELECT hashtextextended( + format('%I.%I', 'awa', awa.queue_lane_sequence_name('queue_enqueue_seq', $1, 2::smallint, 0::smallint)), + 0)", + ) + .bind(&hot) + .fetch_one(&pool) + .await + .expect("derive lane lock key"); + let mut gate = pool.begin().await.expect("begin gate"); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(key) + .execute(&mut *gate) + .await + .expect("take lane lock"); + + let rescue_started = Instant::now(); + let rescued = store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .expect("rescue with the lane lock held"); + println!( + "[lane-maint] rescue of {} claims with the hot lane lock held took {:.0}ms", + rescued.len(), + rescue_started.elapsed().as_secs_f64() * 1e3 + ); + let remaining_hold = hold.saturating_sub(rescue_started.elapsed()); + + let release = tokio::spawn(async move { + tokio::time::sleep(remaining_hold).await; + gate.commit().await.expect("release lane lock"); + }); + let report = run_leader(&pool, &store, &hot, &other, Duration::from_secs(9)).await; + release.await.expect("release task"); + print_leader(&format!("lane held {}ms", hold.as_millis()), &report); +} diff --git a/awa/tests/fixtures/v052_lane_sequence_helpers.sql b/awa/tests/fixtures/v052_lane_sequence_helpers.sql new file mode 100644 index 000000000..b7d396dad --- /dev/null +++ b/awa/tests/fixtures/v052_lane_sequence_helpers.sql @@ -0,0 +1,190 @@ +-- The lane-sequence helpers as migration v052 left them, before lane +-- provisioning functions existed. Installed under a separate name so a test +-- can put a schema back in its v052 shape. +CREATE OR REPLACE FUNCTION awa.v052_refresh_lane_sequence_helpers(p_schema TEXT) +RETURNS VOID +LANGUAGE plpgsql +SECURITY INVOKER +SET search_path = pg_catalog +AS $refresh$ +BEGIN + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.ensure_lane_sequences( + p_queue TEXT, + p_priority SMALLINT, + p_enqueue_shard SMALLINT + ) + RETURNS VOID + LANGUAGE plpgsql + SET search_path = pg_catalog + AS $func$ + DECLARE + v_enqueue_seq TEXT := %1$I.queue_lane_sequence_name( + 'queue_enqueue_seq', + p_queue, + p_priority, + p_enqueue_shard + ); + v_claim_seq TEXT := %1$I.queue_lane_sequence_name( + 'queue_claim_seq', + p_queue, + p_priority, + p_enqueue_shard + ); + BEGIN + IF to_regclass(format('%%I.%%I', %1$L, v_enqueue_seq)) IS NULL THEN + IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_enqueue_seq + USING ERRCODE = '42501', + HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + END IF; + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_enqueue_seq + ); + END IF; + IF to_regclass(format('%%I.%%I', %1$L, v_claim_seq)) IS NULL THEN + IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_claim_seq + USING ERRCODE = '42501', + HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + END IF; + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_claim_seq + ); + END IF; + + UPDATE %1$I.queue_enqueue_heads + SET seq_name = v_enqueue_seq + WHERE queue = p_queue + AND priority = p_priority + AND enqueue_shard = p_enqueue_shard + AND seq_name IS DISTINCT FROM v_enqueue_seq; + + UPDATE %1$I.queue_claim_heads + SET seq_name = v_claim_seq + WHERE queue = p_queue + AND priority = p_priority + AND enqueue_shard = p_enqueue_shard + AND seq_name IS DISTINCT FROM v_claim_seq; + END; + $func$ + $ddl$, + p_schema + ); + + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.queue_enqueue_head_sequence_sync() + RETURNS TRIGGER + LANGUAGE plpgsql + SET search_path = pg_catalog + AS $func$ + DECLARE + v_seq_name TEXT := %1$I.queue_lane_sequence_name( + 'queue_enqueue_seq', + NEW.queue, + NEW.priority, + NEW.enqueue_shard + ); + v_qualified_seq TEXT := format('%%I.%%I', %1$L, v_seq_name); + v_count BIGINT; + v_start BIGINT; + BEGIN + IF to_regclass(format('%%I.%%I', %1$L, v_seq_name)) IS NULL THEN + IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_seq_name + USING ERRCODE = '42501', + HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + END IF; + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_seq_name + ); + END IF; + NEW.seq_name := v_seq_name; + + IF TG_OP = 'UPDATE' + AND NEW.next_seq IS DISTINCT FROM OLD.next_seq + AND NEW.next_seq > OLD.next_seq + THEN + v_count := NEW.next_seq - OLD.next_seq; + PERFORM pg_advisory_lock(hashtextextended(v_qualified_seq, 0)); + BEGIN + EXECUTE format( + 'SELECT nextval(%%L::regclass)::bigint', + v_qualified_seq + ) + INTO v_start; + + IF v_count > 1 THEN + EXECUTE format( + 'SELECT setval(%%L::regclass, %%s, true)', + v_qualified_seq, + v_start + v_count - 1 + ); + END IF; + EXCEPTION WHEN OTHERS THEN + PERFORM pg_advisory_unlock(hashtextextended(v_qualified_seq, 0)); + RAISE; + END; + PERFORM pg_advisory_unlock(hashtextextended(v_qualified_seq, 0)); + NEW.next_seq := v_start + v_count; + ELSE + PERFORM %1$I.set_sequence_next(v_seq_name, NEW.next_seq); + END IF; + + RETURN NEW; + END; + $func$ + $ddl$, + p_schema + ); + + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.queue_claim_head_sequence_sync() + RETURNS TRIGGER + LANGUAGE plpgsql + SET search_path = pg_catalog + AS $func$ + DECLARE + v_seq_name TEXT := %1$I.queue_lane_sequence_name( + 'queue_claim_seq', + NEW.queue, + NEW.priority, + NEW.enqueue_shard + ); + BEGIN + IF to_regclass(format('%%I.%%I', %1$L, v_seq_name)) IS NULL THEN + IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_seq_name + USING ERRCODE = '42501', + HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + END IF; + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_seq_name + ); + END IF; + NEW.seq_name := v_seq_name; + + IF TG_OP = 'INSERT' THEN + PERFORM %1$I.set_sequence_next(v_seq_name, NEW.claim_seq); + END IF; + + RETURN NEW; + END; + $func$ + $ddl$, + p_schema + ); + +END; +$refresh$; diff --git a/awa/tests/insert_many_batch_test.rs b/awa/tests/insert_many_batch_test.rs new file mode 100644 index 000000000..bf1cc7740 --- /dev/null +++ b/awa/tests/insert_many_batch_test.rs @@ -0,0 +1,523 @@ +//! Compatibility inserts and the opt-in batch insert function. +//! +//! A batch does its lane-independent work (job ids, unique claims, lane +//! registration) before it reserves its lanes, and holds a lane lock only for +//! the reservation, the ready-row writes and the commit. These tests cover that +//! ordering, skip-on-conflict inserts, and the shape of what lands in the +//! lanes. +//! +//! Requires a running Postgres instance. +//! Run with: `cargo test --package awa --test insert_many_batch_test -- --nocapture` + +use awa::model::{ + insert_many, insert_many_copy_from_pool, insert_many_skip_unique_conflicts, AwaError, + QueueStorage, +}; +use awa::{InsertOpts, InsertParams, JobState, QueueStorageConfig, UniqueOpts}; +use awa_testing::setup; +use sqlx::PgPool; +use std::sync::LazyLock; +use std::time::{Duration, Instant}; +use tokio::sync::Mutex; +use uuid::Uuid; + +/// The tests switch the storage engine of a shared database. +static ENGINE_GUARD: LazyLock> = LazyLock::new(|| Mutex::new(())); + +const WAIT_DEADLINE: Duration = Duration::from_secs(20); + +async fn queue_storage_pool() -> PgPool { + let pool = setup::setup(8).await; + setup::activate_queue_storage(&pool).await; + pool +} + +fn unique_queue(label: &str) -> String { + format!( + "batch_{label}_{}", + &Uuid::new_v4().simple().to_string()[..12] + ) +} + +fn job(queue: &str, seq: i64, unique: bool) -> InsertParams { + InsertParams { + kind: "insert_many_batch_job".to_string(), + args: serde_json::json!({ "seq": seq }), + opts: InsertOpts { + queue: queue.to_string(), + unique: unique.then(|| UniqueOpts { + by_queue: true, + by_args: true, + ..Default::default() + }), + ..Default::default() + }, + } +} + +async fn ready_lane_seqs(pool: &PgPool, queue: &str) -> Vec<(i64, i64)> { + sqlx::query_as( + "SELECT (args->>'seq')::bigint, lane_seq FROM awa.ready_entries + WHERE queue = $1 ORDER BY lane_seq", + ) + .bind(queue) + .fetch_all(pool) + .await + .expect("read ready rows") +} + +async fn claims_for_queue(pool: &PgPool, queue: &str) -> i64 { + sqlx::query_scalar( + "SELECT count(*)::bigint FROM awa.job_unique_claims AS claims + WHERE claims.job_id IN (SELECT job_id FROM awa.ready_entries WHERE queue = $1)", + ) + .bind(queue) + .fetch_one(pool) + .await + .expect("count claims") +} + +fn is_unique_violation(err: &AwaError) -> bool { + match err { + AwaError::UniqueConflict { .. } => true, + AwaError::Database(sqlx::Error::Database(db_err)) => { + db_err.code().as_deref() == Some("23505") + } + _ => false, + } +} + +async fn lane_lock_key(pool: &PgPool, queue: &str) -> i64 { + sqlx::query_scalar( + "SELECT hashtextextended( + format('%I.%I', 'awa', awa.queue_lane_sequence_name('queue_enqueue_seq', $1, 2::smallint, 0::smallint)), + 0)", + ) + .bind(queue) + .fetch_one(pool) + .await + .expect("derive lane lock key") +} + +#[tokio::test] +async fn batch_lands_in_one_lane_in_input_order_and_is_claimable() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + let queue = unique_queue("order"); + + let jobs: Vec = (0..300).map(|seq| job(&queue, seq, seq % 2 == 0)).collect(); + let inserted = insert_many(&pool, &jobs).await.expect("insert batch"); + assert_eq!( + inserted + .iter() + .map(|row| row.args["seq"].as_i64().unwrap()) + .collect::>(), + (0..300).collect::>(), + "rows come back in input order" + ); + assert!(inserted.iter().all(|row| row.state == JobState::Available)); + + let lane = ready_lane_seqs(&pool, &queue).await; + assert_eq!(lane.len(), 300); + let first = lane[0].1; + for (offset, (seq, lane_seq)) in lane.iter().enumerate() { + assert_eq!(*seq, offset as i64, "lane_seq follows input order"); + assert_eq!(*lane_seq, first + offset as i64, "lane_seq is dense"); + } + + // The lane's segment metadata must let the claim path find every row. + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + let mut claimed = Vec::new(); + while claimed.len() < 300 { + let batch = store + .claim_runtime_batch(&pool, &queue, 100, Duration::from_secs(60)) + .await + .expect("claim"); + assert!(!batch.is_empty(), "claim path lost rows of the batch"); + claimed.extend( + batch + .into_iter() + .map(|job| job.job.args["seq"].as_i64().unwrap()), + ); + } + assert_eq!( + claimed, + (0..300).collect::>(), + "claimed in lane order" + ); +} + +#[tokio::test] +async fn batch_with_mixed_states_routes_each_row() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + let queue = unique_queue("mixed"); + + let later = chrono::Utc::now() + chrono::Duration::hours(1); + let mut jobs: Vec = (0..6).map(|seq| job(&queue, seq, false)).collect(); + jobs[1].opts.run_at = Some(later); + jobs[4].opts.run_at = Some(later); + + let inserted = insert_many(&pool, &jobs).await.expect("insert mixed batch"); + let states: Vec = inserted.iter().map(|row| row.state).collect(); + assert_eq!( + states, + vec![ + JobState::Available, + JobState::Scheduled, + JobState::Available, + JobState::Available, + JobState::Scheduled, + JobState::Available, + ] + ); + let ready: Vec = ready_lane_seqs(&pool, &queue) + .await + .into_iter() + .map(|(seq, _)| seq) + .collect(); + assert_eq!(ready, vec![0, 2, 3, 5]); + let deferred: i64 = sqlx::query_scalar( + "SELECT count(*)::bigint FROM awa.deferred_jobs WHERE queue = $1 AND state = 'scheduled'", + ) + .bind(&queue) + .fetch_one(&pool) + .await + .expect("count deferred"); + assert_eq!(deferred, 2); +} + +#[tokio::test] +async fn unique_conflict_fails_the_batch_atomically() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + let queue = unique_queue("atomic"); + + insert_many(&pool, &[job(&queue, 1, true)]) + .await + .expect("first insert"); + + let batch = vec![ + job(&queue, 0, true), + job(&queue, 1, true), + job(&queue, 2, true), + ]; + let err = insert_many(&pool, &batch) + .await + .expect_err("duplicate key must fail the batch"); + assert!(is_unique_violation(&err), "{err:?}"); + + let landed: Vec = ready_lane_seqs(&pool, &queue) + .await + .into_iter() + .map(|(seq, _)| seq) + .collect(); + assert_eq!(landed, vec![1], "nothing from the failed batch landed"); + assert_eq!(claims_for_queue(&pool, &queue).await, 1); +} + +#[tokio::test] +async fn skip_unique_conflicts_inserts_the_rest_in_order() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + let queue = unique_queue("skip"); + + insert_many(&pool, &[job(&queue, 1, true), job(&queue, 3, true)]) + .await + .expect("seed"); + + // Keys 1 and 3 are already claimed, 4 repeats inside the batch, and the + // last row carries no unique key at all. + let batch = vec![ + job(&queue, 0, true), + job(&queue, 1, true), + job(&queue, 2, true), + job(&queue, 3, true), + job(&queue, 4, true), + job(&queue, 4, true), + job(&queue, 5, false), + ]; + let inserted = insert_many_skip_unique_conflicts(&pool, &batch) + .await + .expect("skip insert"); + assert_eq!( + inserted + .iter() + .map(|row| row.args["seq"].as_i64().unwrap()) + .collect::>(), + vec![0, 2, 4, 5] + ); + + let lane = ready_lane_seqs(&pool, &queue).await; + assert_eq!( + lane.iter().map(|(seq, _)| *seq).collect::>(), + vec![1, 3, 0, 2, 4, 5], + "skipped rows take no lane position" + ); + for pair in lane.windows(2) { + assert_eq!(pair[1].1, pair[0].1 + 1, "lane_seq stays dense"); + } + + // Everything skipped: an empty result, not an error. + let inserted = insert_many_skip_unique_conflicts(&pool, &[job(&queue, 0, true)]) + .await + .expect("all skipped"); + assert!(inserted.is_empty()); +} + +#[tokio::test] +async fn copy_skips_unique_conflicts_in_one_pass() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + let queue = unique_queue("copy_skip"); + + insert_many(&pool, &[job(&queue, 2, true)]) + .await + .expect("seed"); + let batch: Vec = (0..6).map(|seq| job(&queue, seq, true)).collect(); + let inserted = insert_many_copy_from_pool(&pool, &batch) + .await + .expect("copy insert"); + assert_eq!( + inserted + .iter() + .map(|row| row.args["seq"].as_i64().unwrap()) + .collect::>(), + vec![0, 1, 3, 4, 5] + ); + assert_eq!(ready_lane_seqs(&pool, &queue).await.len(), 6); +} + +#[tokio::test] +async fn canonical_storage_inserts_and_skips_row_by_row() { + let _guard = ENGINE_GUARD.lock().await; + let pool = setup::setup(8).await; + setup::reset_to_canonical(&pool).await; + let queue = unique_queue("canonical"); + + insert_many(&pool, &[job(&queue, 1, true)]) + .await + .expect("seed"); + let batch: Vec = (0..4).map(|seq| job(&queue, seq, true)).collect(); + let err = insert_many(&pool, &batch).await.expect_err("duplicate"); + assert!(is_unique_violation(&err), "{err:?}"); + + let inserted = insert_many_skip_unique_conflicts(&pool, &batch) + .await + .expect("skip insert"); + assert_eq!( + inserted + .iter() + .map(|row| row.args["seq"].as_i64().unwrap()) + .collect::>(), + vec![0, 2, 3] + ); +} + +/// While another transaction holds a lane lock, a batch bound for that lane +/// has already done its lane-independent work (it holds its unique claims) and +/// has not yet written a ready row. Every row's unique claim, not just the +/// first's, is held by then. +#[tokio::test] +async fn batch_waits_for_the_lane_after_its_lane_independent_work() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + let queue = unique_queue("window"); + + // Provision the lane so the gate lock is the only thing to wait for. + insert_many(&pool, &[job(&queue, -1, false)]) + .await + .expect("provision lane"); + + let key = lane_lock_key(&pool, &queue).await; + let mut gate = pool.begin().await.expect("begin gate"); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(key) + .execute(&mut *gate) + .await + .expect("take lane lock"); + + let jobs_for_probe: Vec = (0..50).map(|seq| job(&queue, seq, true)).collect(); + let producer = { + let pool = pool.clone(); + let jobs = jobs_for_probe.clone(); + tokio::spawn(async move { insert_many_skip_unique_conflicts(&pool, &jobs).await }) + }; + + let started = Instant::now(); + let waiter = loop { + let pid: Option = sqlx::query_scalar( + "SELECT pid FROM pg_locks + WHERE locktype = 'advisory' AND NOT granted + AND ((classid::bigint << 32) | objid::bigint) = $1 + LIMIT 1", + ) + .bind(key) + .fetch_optional(&pool) + .await + .expect("poll pg_locks"); + if let Some(pid) = pid { + break pid; + } + assert!(started.elapsed() < WAIT_DEADLINE, "producer never waited"); + tokio::time::sleep(Duration::from_millis(10)).await; + }; + + let holds_relation_lock = |relation: &'static str| { + let pool = pool.clone(); + async move { + sqlx::query_scalar::<_, bool>( + "SELECT EXISTS ( + SELECT 1 FROM pg_locks + WHERE pid = $1 AND locktype = 'relation' AND granted + AND relation = $2::regclass + )", + ) + .bind(waiter) + .bind(relation) + .fetch_one(&pool) + .await + .expect("read relation locks") + } + }; + // The claim of the batch's last row is already held, uncommitted: an + // insert of the same key waits on the producer's transaction. + let last = jobs_for_probe.last().expect("probe job"); + let last_key = + awa::model::prepare_raw_job_insert(last.kind.clone(), last.args.clone(), last.opts.clone()) + .expect("prepare last job") + .unique_key() + .expect("unique key") + .to_vec(); + let mut probe = pool.begin().await.expect("begin probe"); + sqlx::query("SET LOCAL lock_timeout = '300ms'") + .execute(&mut *probe) + .await + .expect("set lock timeout"); + let probe_result = sqlx::query( + "INSERT INTO awa.job_unique_claims (unique_key, job_id) VALUES ($1, 0) + ON CONFLICT (unique_key) DO NOTHING", + ) + .bind(&last_key) + .execute(&mut *probe) + .await; + match probe_result { + Err(sqlx::Error::Database(db_err)) => assert_eq!( + db_err.code().as_deref(), + Some("55P03"), + "the last row's claim must already be held while the lane lock is awaited" + ), + other => panic!("claim of the last row was not held while waiting for the lane: {other:?}"), + } + drop(probe); + assert!( + !holds_relation_lock("awa.ready_entries").await, + "no ready row is written while the lane lock is awaited" + ); + + gate.commit().await.expect("release lane lock"); + let inserted = tokio::time::timeout(WAIT_DEADLINE, producer) + .await + .expect("producer finished") + .expect("producer task") + .expect("insert batch"); + assert_eq!(inserted.len(), 50); + assert_eq!(ready_lane_seqs(&pool, &queue).await.len(), 51); +} + +/// Producers on the per-row function, as binaries before v055 are, and +/// producers on the batch function share one lane at the same time. Each batch +/// still occupies a contiguous run of `lane_seq`, in input order, and the lane +/// has no duplicate or missing positions. +#[tokio::test] +async fn old_and_new_producers_share_a_lane_without_reordering() { + const TASKS: usize = 8; + const BATCHES: usize = 25; + const BATCH: i64 = 40; + + let _guard = ENGINE_GUARD.lock().await; + let pool = setup::setup(TASKS as u32 + 4).await; + setup::activate_queue_storage(&pool).await; + let queue = unique_queue("mixed_fleet"); + + // Provision the lane before mixing generic and batch producers. + insert_many(&pool, &[job(&queue, -1, false)]) + .await + .expect("warm up"); + + let mut tasks = Vec::new(); + for task in 0..TASKS { + let pool = pool.clone(); + let queue = queue.clone(); + tasks.push(tokio::spawn(async move { + for batch in 0..BATCHES { + let jobs: Vec = (0..BATCH) + .map(|seq| InsertParams { + kind: "insert_many_batch_job".to_string(), + args: serde_json::json!({ "seq": seq, "batch": format!("{task}-{batch}") }), + opts: InsertOpts { + queue: queue.clone(), + ..Default::default() + }, + }) + .collect(); + if task % 2 == 0 { + insert_through_row_function(&pool, &jobs).await; + } else { + insert_many_skip_unique_conflicts(&pool, &jobs) + .await + .expect("batch function insert"); + } + } + })); + } + for task in tasks { + task.await.expect("producer task"); + } + + let rows: Vec<(String, i64, i64)> = sqlx::query_as( + "SELECT args->>'batch', (args->>'seq')::bigint, lane_seq FROM awa.ready_entries + WHERE queue = $1 AND args ? 'batch' ORDER BY lane_seq", + ) + .bind(&queue) + .fetch_all(&pool) + .await + .expect("read lane"); + assert_eq!(rows.len(), TASKS * BATCHES * BATCH as usize); + for pair in rows.windows(2) { + assert!(pair[1].2 > pair[0].2, "lane_seq is unique and ascending"); + } + let mut by_batch: std::collections::HashMap<&str, Vec<(i64, i64)>> = + std::collections::HashMap::new(); + for (batch, seq, lane_seq) in &rows { + by_batch.entry(batch).or_default().push((*seq, *lane_seq)); + } + for (batch, members) in by_batch { + for (offset, (seq, lane_seq)) in members.iter().enumerate() { + assert_eq!(*seq, offset as i64, "batch {batch} keeps input order"); + assert_eq!( + *lane_seq, + members[0].1 + offset as i64, + "batch {batch} occupies a contiguous run of the lane" + ); + } + } +} + +/// What a binary from before v055 sends: one `insert_job_compat` call per row. +async fn insert_through_row_function(pool: &PgPool, jobs: &[InsertParams]) { + sqlx::query( + r#" + SELECT inserted.* + FROM unnest($1::text[], $2::text[], $3::jsonb[]) WITH ORDINALITY AS input(kind, queue, args, ord) + CROSS JOIN LATERAL awa.insert_job_compat(input.kind, input.queue, input.args) AS inserted + ORDER BY input.ord + "#, + ) + .bind(jobs.iter().map(|job| job.kind.as_str()).collect::>()) + .bind(jobs.iter().map(|job| job.opts.queue.as_str()).collect::>()) + .bind(jobs.iter().map(|job| &job.args).collect::>()) + .fetch_all(pool) + .await + .expect("per-row function insert"); +} diff --git a/awa/tests/insert_many_schema_fallback_test.rs b/awa/tests/insert_many_schema_fallback_test.rs new file mode 100644 index 000000000..271ce5f7a --- /dev/null +++ b/awa/tests/insert_many_schema_fallback_test.rs @@ -0,0 +1,277 @@ +//! `insert_many` and `insert_many_copy` against a schema without +//! `awa.insert_jobs_compat_batch`. +//! +//! A binary may run ahead of migration v055. Generic `insert_many` keeps using +//! the per-row function; COPY probes its own connection and uses the batch +//! function when installed. Different databases in one process can have +//! different schema versions without aborting their caller's transactions. +//! +//! Requires a running Postgres instance. +//! Run with: `cargo test --package awa --test insert_many_schema_fallback_test` + +use awa::audited_sql; +use awa::model::{ + insert_many, insert_many_copy, insert_many_copy_from_pool, insert_many_skip_unique_conflicts, + migrations, AwaError, +}; +use awa::{InsertOpts, InsertParams, UniqueOpts}; +use awa_testing::setup; +use sqlx::PgPool; +use uuid::Uuid; + +const DROP_BATCH_FUNCTION: &str = "DROP FUNCTION awa.insert_jobs_compat_batch(text[], text[], \ + jsonb[], text[], smallint[], smallint[], timestamptz[], jsonb[], jsonb[], bytea[], text[], \ + text[], boolean)"; +const MIGRATION_V055: &str = + include_str!("../../awa-model/migrations/v055_insert_jobs_batch_compat.sql"); + +fn database_url_for(database: &str) -> String { + let url = setup::database_url(); + let (without_query, query) = match url.split_once('?') { + Some((prefix, query)) => (prefix, Some(query)), + None => (url.as_str(), None), + }; + let (base, _) = without_query + .rsplit_once('/') + .expect("database name in URL"); + match query { + Some(query) => format!("{base}/{database}?{query}"), + None => format!("{base}/{database}"), + } +} + +fn jobs(queue: &str, range: std::ops::Range, unique: bool) -> Vec { + range + .map(|seq| InsertParams { + kind: "schema_fallback_job".to_string(), + args: serde_json::json!({ "seq": seq }), + opts: InsertOpts { + queue: queue.to_string(), + unique: unique.then(|| UniqueOpts { + by_queue: true, + by_args: true, + ..Default::default() + }), + ..Default::default() + }, + }) + .collect() +} + +async fn segments(pool: &PgPool, queue: &str) -> i64 { + sqlx::query_scalar("SELECT count(*)::bigint FROM awa.ready_segments WHERE queue = $1") + .bind(queue) + .fetch_one(pool) + .await + .expect("count segments") +} + +async fn ready_seqs(pool: &PgPool, queue: &str) -> Vec { + sqlx::query_scalar( + "SELECT (args->>'seq')::bigint FROM awa.ready_entries WHERE queue = $1 ORDER BY lane_seq", + ) + .bind(queue) + .fetch_all(pool) + .await + .expect("read ready rows") +} + +fn seqs(rows: &[awa::JobRow]) -> Vec { + rows.iter() + .map(|row| row.args["seq"].as_i64().expect("seq")) + .collect() +} + +#[tokio::test] +async fn binary_ahead_of_the_migration_falls_back_and_recovers() { + let admin = setup::pool(2).await; + let database = format!("awa_batchfn_{}", &Uuid::new_v4().simple().to_string()[..12]); + sqlx::query(audited_sql(format!("CREATE DATABASE {database}"))) + .execute(&admin) + .await + .expect("create scratch database"); + + let pool = setup::pool_with_url(&database_url_for(&database), 4).await; + migrations::run(&pool).await.expect("migrate"); + setup::activate_queue_storage(&pool).await; + + // A schema from before v055. + sqlx::query(DROP_BATCH_FUNCTION) + .execute(&pool) + .await + .expect("drop batch function"); + + let first = insert_many(&pool, &jobs("fallback_a", 0..3, false)) + .await + .expect("insert_many without the batch function"); + assert_eq!(seqs(&first), vec![0, 1, 2]); + assert_eq!(segments(&pool, "fallback_a").await, 3, "per-row path"); + insert_many(&pool, &jobs("fallback_a", 3..6, false)) + .await + .expect("a second call without the function"); + assert_eq!(segments(&pool, "fallback_a").await, 6, "still per-row"); + assert_eq!( + ready_seqs(&pool, "fallback_a").await, + vec![0, 1, 2, 3, 4, 5] + ); + + // Unique handling on the per-row path: a duplicate fails `insert_many` + // and is skipped by `insert_many_copy`. + insert_many(&pool, &jobs("fallback_u", 0..2, true)) + .await + .expect("seed unique jobs"); + let duplicate = insert_many(&pool, &jobs("fallback_u", 1..3, true)) + .await + .expect_err("a unique conflict fails the batch"); + assert!( + matches!(&duplicate, AwaError::Database(sqlx::Error::Database(e)) if e.code().as_deref() == Some("23505")), + "{duplicate:?}" + ); + let copied = insert_many_copy_from_pool(&pool, &jobs("fallback_u", 1..4, true)) + .await + .expect("insert_many_copy without the batch function"); + assert_eq!(seqs(&copied), vec![2, 3], "duplicate skipped"); + + // The skip API is new with v055 and says so. + let skipped = insert_many_skip_unique_conflicts(&pool, &jobs("fallback_u", 0..1, true)) + .await + .expect_err("needs the migration"); + assert!( + matches!(&skipped, AwaError::Validation(message) if message.contains("v055")), + "{skipped:?}" + ); + + // Generic insert_many stays on the per-row path when the migration runs; + // COPY chooses the batch function on its next call. + sqlx::raw_sql(MIGRATION_V055) + .execute(&pool) + .await + .expect("apply v055"); + insert_many(&pool, &jobs("fallback_b", 0..3, false)) + .await + .expect("generic insert after migration"); + assert_eq!(segments(&pool, "fallback_b").await, 3, "per-row path"); + insert_many(&pool, &jobs("fallback_b", 3..6, false)) + .await + .expect("second generic insert after migration"); + assert_eq!( + segments(&pool, "fallback_b").await, + 6, + "generic insert_many keeps its schema-independent per-row path" + ); + assert_eq!( + ready_seqs(&pool, "fallback_b").await, + vec![0, 1, 2, 3, 4, 5] + ); + + let copied = insert_many_copy_from_pool(&pool, &jobs("fallback_u", 2..6, true)) + .await + .expect("insert_many_copy with the batch function"); + assert_eq!(seqs(&copied), vec![4, 5], "duplicates skipped in one pass"); + let skipped = insert_many_skip_unique_conflicts(&pool, &jobs("fallback_u", 0..8, true)) + .await + .expect("skip with the batch function"); + assert_eq!(seqs(&skipped), vec![6, 7]); + + // Removing the function after a successful batch does not poison the next + // generic or COPY call with a stale process-wide answer. + sqlx::query(DROP_BATCH_FUNCTION) + .execute(&pool) + .await + .expect("drop batch function again"); + let mut tx = pool.begin().await.expect("caller transaction"); + insert_many(&mut *tx, &jobs("fallback_c", 0..3, false)) + .await + .expect("generic insert still works"); + insert_many_copy(&mut tx, &jobs("fallback_c", 3..6, false)) + .await + .expect("COPY re-probes its own connection"); + tx.commit() + .await + .expect("caller transaction remains usable"); + assert_eq!(segments(&pool, "fallback_c").await, 6, "per-row paths"); + + pool.close().await; + sqlx::query(audited_sql(format!( + "DROP DATABASE {database} WITH (FORCE)" + ))) + .execute(&admin) + .await + .expect("drop scratch database"); +} + +/// Availability learned from one database must never select SQL for another +/// database. Both paths join the caller's transaction, including its rollback. +#[tokio::test] +async fn mixed_schema_databases_do_not_abort_caller_transactions() { + let current = setup::TestDatabase::canonical().await; + let older = setup::TestDatabase::canonical().await; + for db in [¤t, &older] { + setup::activate_queue_storage(db.pool()).await; + sqlx::query("CREATE TABLE application_writes (id bigint PRIMARY KEY)") + .execute(db.pool()) + .await + .expect("create application table"); + } + sqlx::query(DROP_BATCH_FUNCTION) + .execute(older.pool()) + .await + .expect("older database has no batch function"); + + // This call previously seeded the shared cache with Present. + insert_many_skip_unique_conflicts(current.pool(), &jobs("mixed_warm", 0..2, true)) + .await + .expect("batch API on current schema"); + + for (index, db) in [¤t, &older, ¤t].into_iter().enumerate() { + let pool = db.pool(); + let queue = format!("mixed_database_{index}"); + let mut tx = pool.begin().await.expect("caller transaction"); + sqlx::query("INSERT INTO application_writes (id) VALUES ($1)") + .bind(index as i64) + .execute(&mut *tx) + .await + .expect("application write"); + let inserted = insert_many(&mut *tx, &jobs(&queue, 0..3, true)) + .await + .expect("generic insert must use this database's supported SQL"); + assert_eq!(seqs(&inserted), vec![0, 1, 2]); + let copied = insert_many_copy(&mut tx, &jobs(&queue, 2..5, true)) + .await + .expect("COPY must probe this database"); + assert_eq!(seqs(&copied), vec![3, 4], "duplicate skipped"); + let written: i64 = + sqlx::query_scalar("SELECT count(*) FROM application_writes WHERE id = $1") + .bind(index as i64) + .fetch_one(&mut *tx) + .await + .expect("caller transaction was not aborted"); + assert_eq!(written, 1); + tx.commit() + .await + .expect("commit application and jobs together"); + assert_eq!(ready_seqs(pool, &queue).await, vec![0, 1, 2, 3, 4]); + + let mut tx = pool.begin().await.expect("rollback transaction"); + sqlx::query("INSERT INTO application_writes (id) VALUES ($1)") + .bind(index as i64 + 100) + .execute(&mut *tx) + .await + .expect("application write to roll back"); + insert_many(&mut *tx, &jobs(&queue, 5..7, true)) + .await + .expect("generic insert to roll back"); + insert_many_copy(&mut tx, &jobs(&queue, 7..9, true)) + .await + .expect("COPY to roll back"); + tx.rollback().await.expect("roll back caller transaction"); + assert_eq!(ready_seqs(pool, &queue).await, vec![0, 1, 2, 3, 4]); + let rolled_back: i64 = + sqlx::query_scalar("SELECT count(*) FROM application_writes WHERE id = $1") + .bind(index as i64 + 100) + .fetch_one(pool) + .await + .expect("check application rollback"); + assert_eq!(rolled_back, 0); + } +} diff --git a/awa/tests/integration_test.rs b/awa/tests/integration_test.rs index c816f0675..a4b9c6567 100644 --- a/awa/tests/integration_test.rs +++ b/awa/tests/integration_test.rs @@ -552,6 +552,9 @@ async fn test_insert_many() { #[tokio::test] async fn test_insert_many_exceeding_bind_parameter_limit() { const JOB_COUNT: usize = 6_000; + // A concurrent test that activates queue storage would route part of the + // batch away from the canonical tables this test reads back. + let _guard = test_lock().lock().await; let client = setup_with_connections(4).await; let queue = "integ_insert_many_bind_limit"; clean_queue(client.pool(), queue).await; diff --git a/awa/tests/lane_ownership_roles_test.rs b/awa/tests/lane_ownership_roles_test.rs new file mode 100644 index 000000000..f7ad8eda7 --- /dev/null +++ b/awa/tests/lane_ownership_roles_test.rs @@ -0,0 +1,833 @@ +//! Queue lanes under a split-role install: a schema owner, a migration role +//! that is a member of it, and several application roles that share a runtime +//! group role. Only the owner and the migration role have default privileges, +//! so a lane sequence is usable by every runtime role only if it is created as +//! the owner or granted explicitly, whichever role touches the lane first. + +use awa::model::storage::LaneSequenceOutcome; +use awa::model::{migrations, storage, AwaError, InsertOpts, InsertParams, QueueStorage}; +use awa::model::{JobState, QueueStorageConfig}; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; +use sqlx::{Executor, PgPool}; +use std::str::FromStr; +use std::time::Duration; + +type TestResult = Result<(), Box>; + +fn sql(text: impl Into) -> sqlx::AssertSqlSafe { + awa::audited_sql(text) +} + +/// Roles, a database owned by the schema owner, and a pool per role. +struct Cluster { + admin: PgPool, + database: String, + roles: Vec, + owner_name: String, + migrator_name: String, + runtime_name: String, + connect_options: PgConnectOptions, + owner: PgPool, + migrator: PgPool, + app_a: PgPool, + app_b: PgPool, + app_c: PgPool, +} + +#[derive(Clone, Copy, PartialEq)] +enum Schema { + /// Every migration, including the lane provisioning functions. + Current, + /// The release before lane provisioning (v052). + BeforeProvisioning, +} + +impl Cluster { + /// `runtime_create` grants the runtime group CREATE on the schema, as a + /// 0.6 install does so runtimes can create lane sequences themselves. + async fn new( + schema: Schema, + runtime_create: bool, + ) -> Result> { + let url = std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()); + let admin = PgPoolOptions::new() + .max_connections(2) + .connect(&url) + .await?; + let suffix = uuid::Uuid::new_v4().simple().to_string()[..12].to_string(); + let owner_name = format!("lane_owner_{suffix}"); + let migrator_name = format!("lane_migrator_{suffix}"); + let runtime_name = format!("lane_runtime_{suffix}"); + let app_names = ["a", "b", "c"].map(|app| format!("lane_app_{app}_{suffix}")); + let database = format!("lane_roles_{suffix}"); + admin + .execute(sqlx::raw_sql(sql(format!( + r#" + CREATE ROLE {owner_name} LOGIN PASSWORD 'lane_test'; + CREATE ROLE {migrator_name} LOGIN PASSWORD 'lane_test' IN ROLE {owner_name}; + CREATE ROLE {runtime_name} NOLOGIN; + CREATE ROLE {a} LOGIN PASSWORD 'lane_test' IN ROLE {runtime_name}; + CREATE ROLE {b} LOGIN PASSWORD 'lane_test' IN ROLE {runtime_name}; + CREATE ROLE {c} LOGIN PASSWORD 'lane_test' IN ROLE {runtime_name}; + "#, + a = app_names[0], + b = app_names[1], + c = app_names[2], + )))) + .await?; + admin + .execute(sqlx::raw_sql(sql(format!( + "CREATE DATABASE {database} OWNER {owner_name}" + )))) + .await?; + let mut roles = vec![ + owner_name.clone(), + migrator_name.clone(), + runtime_name.clone(), + ]; + roles.extend(app_names.iter().cloned()); + let options = PgConnectOptions::from_str(&url)? + .database(&database) + .password("lane_test"); + let pool = |user: &str| { + PgPoolOptions::new() + .max_connections(3) + .connect_with(options.clone().username(user)) + }; + let owner = pool(&owner_name).await?; + let migrator = pool(&migrator_name).await?; + + owner + .execute(sqlx::raw_sql(sql(format!( + r#" + CREATE SCHEMA awa; + GRANT USAGE ON SCHEMA awa TO {runtime_name}; + "# + )))) + .await?; + for (creator, creator_pool) in [(&owner_name, &owner), (&migrator_name, &migrator)] { + creator_pool + .execute(sqlx::raw_sql(sql(format!( + r#" + ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa + GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON TABLES TO {runtime_name}; + ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa + GRANT USAGE, SELECT, UPDATE ON SEQUENCES TO {runtime_name}; + ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa + GRANT EXECUTE ON FUNCTIONS TO {runtime_name}; + "# + )))) + .await?; + } + // The migration role runs every migration with the owner as its + // effective role, as the database-roles guide describes. + let migrate_as_owner = PgPoolOptions::new() + .max_connections(1) + .connect_with( + options + .clone() + .username(&migrator_name) + .options([("role", owner_name.as_str())]), + ) + .await?; + let target = match schema { + Schema::Current => migrations::CURRENT_VERSION, + Schema::BeforeProvisioning => 52, + }; + for (_, _, statement) in migrations::migration_sql_range(0, target) { + migrate_as_owner + .execute(sqlx::raw_sql(sql(statement))) + .await?; + } + if schema == Schema::BeforeProvisioning { + // Every migration bundle carries the current helper definitions, + // so put the v052 helpers back and drop what v052 did not have. + migrate_as_owner + .execute(sqlx::raw_sql(sql(include_str!( + "fixtures/v052_lane_sequence_helpers.sql" + )))) + .await?; + migrate_as_owner + .execute(sqlx::raw_sql(sql(r#" + SELECT awa.v052_refresh_lane_sequence_helpers('awa'); + DROP FUNCTION awa.v052_refresh_lane_sequence_helpers(text); + DROP FUNCTION awa.provision_lane(text, smallint, smallint); + DROP FUNCTION awa.provision_lane_sequences(text, smallint, smallint); + DROP FUNCTION awa.repair_lane_sequence(text, text); + DROP FUNCTION awa.lane_sequence_problems(text, text); + "#))) + .await?; + } + migrate_as_owner.close().await; + if runtime_create { + owner + .execute(sqlx::raw_sql(sql(format!( + "GRANT CREATE ON SCHEMA awa TO {runtime_name}" + )))) + .await?; + } + + let app_a = pool(&app_names[0]).await?; + let app_b = pool(&app_names[1]).await?; + let app_c = pool(&app_names[2]).await?; + Ok(Self { + admin, + database, + roles, + owner_name, + migrator_name, + runtime_name, + connect_options: options, + owner, + migrator, + app_a, + app_b, + app_c, + }) + } + + async fn teardown(self) -> TestResult { + for pool in [ + &self.owner, + &self.migrator, + &self.app_a, + &self.app_b, + &self.app_c, + ] { + pool.close().await; + } + self.admin + .execute(sqlx::raw_sql(sql(format!( + "DROP DATABASE {} WITH (FORCE)", + self.database + )))) + .await?; + for role in self.roles.iter().rev() { + self.admin + .execute(sqlx::raw_sql(sql(format!("DROP ROLE {role}")))) + .await?; + } + self.admin.close().await; + Ok(()) + } + + /// The migration role acting as the schema owner, as `awa migrate` runs + /// under the database-roles guide. + async fn migrator_as_owner(&self) -> Result { + PgPoolOptions::new() + .max_connections(1) + .connect_with( + self.connect_options + .clone() + .username(&self.migrator_name) + .options([("role", self.owner_name.as_str())]), + ) + .await + } + + /// The `queue_enqueue_seq_*` sequence of `queue`'s first lane. + async fn enqueue_sequence(&self, queue: &str) -> Result { + Ok(self + .lane_sequence_owners(queue) + .await? + .into_iter() + .map(|(sequence, _)| sequence) + .find(|sequence| sequence.starts_with("queue_enqueue_seq_")) + .expect("lane has an enqueue sequence")) + } + + /// `(sequence, owner)` for every lane sequence of `queue`. + async fn lane_sequence_owners( + &self, + queue: &str, + ) -> Result, sqlx::Error> { + sqlx::query_as( + r#" + SELECT sequence.relname::text, pg_get_userbyid(sequence.relowner)::text + FROM awa.queue_enqueue_heads AS heads + CROSS JOIN LATERAL (VALUES (heads.seq_name), (( + SELECT claims.seq_name FROM awa.queue_claim_heads AS claims + WHERE claims.queue = heads.queue + AND claims.priority = heads.priority + AND claims.enqueue_shard = heads.enqueue_shard + ))) AS lane(seq_name) + JOIN pg_class AS sequence + ON sequence.relname = lane.seq_name + AND sequence.relnamespace = 'awa'::regnamespace + WHERE heads.queue = $1 + ORDER BY 1 + "#, + ) + .bind(queue) + .fetch_all(&self.owner) + .await + } + + async fn assert_lanes_owned_by_owner(&self, queue: &str, lanes: usize) -> TestResult { + let owners = self.lane_sequence_owners(queue).await?; + assert_eq!( + owners.len(), + lanes * 2, + "lane sequences of {queue}: {owners:?}" + ); + for (sequence, owner) in owners { + assert_eq!(owner, self.owner_name, "{sequence} of {queue}"); + } + Ok(()) + } + + async fn set_enqueue_shards(&self, queue: &str, shards: i16) -> Result<(), sqlx::Error> { + sqlx::query( + "INSERT INTO awa.queue_meta (queue, enqueue_shards) VALUES ($1, $2) \ + ON CONFLICT (queue) DO UPDATE SET enqueue_shards = EXCLUDED.enqueue_shards", + ) + .bind(queue) + .bind(shards) + .execute(&self.owner) + .await?; + Ok(()) + } +} + +fn store() -> QueueStorage { + QueueStorage::new(QueueStorageConfig::default()).expect("default queue storage config") +} + +/// Enqueue `count` jobs through a fresh store, as a producer process would. +async fn enqueue(pool: &PgPool, queue: &str, priority: i16, count: i64) -> Result { + store().enqueue_batch(pool, queue, priority, count).await +} + +/// Claim and complete everything on `queue`, as a worker does; returns how +/// many jobs completed. +async fn drain(pool: &PgPool, queue: &str) -> Result { + let worker = store(); + let mut completed = 0; + loop { + let claimed = worker + .claim_runtime_batch(pool, queue, 100, Duration::from_secs(60)) + .await?; + if claimed.is_empty() { + return Ok(completed); + } + completed += worker.complete_runtime_batch(pool, &claimed).await?.len(); + } +} + +/// Schedule jobs that become due shortly, one per priority. +async fn schedule(pool: &PgPool, queue: &str, priorities: &[i16]) -> Result<(), AwaError> { + let run_at = chrono::Utc::now() + chrono::Duration::milliseconds(200); + let jobs: Vec = priorities + .iter() + .map(|&priority| InsertParams { + kind: "lane_ownership_job".into(), + args: serde_json::json!({}), + opts: InsertOpts { + queue: queue.into(), + priority, + run_at: Some(run_at), + ..Default::default() + }, + }) + .collect(); + store().enqueue_params_batch(pool, &jobs).await?; + tokio::time::sleep(Duration::from_millis(300)).await; + Ok(()) +} + +async fn first_touch_by_one_app_role(runtime_create: bool) -> TestResult { + let cluster = Cluster::new(Schema::Current, runtime_create).await?; + + // Role A creates the lane; role B enqueues, claims and completes on it. + assert_eq!(enqueue(&cluster.app_a, "fresh_queue", 3, 2).await?, 2); + assert_eq!(enqueue(&cluster.app_b, "fresh_queue", 3, 2).await?, 2); + assert_eq!(drain(&cluster.app_b, "fresh_queue").await?, 4); + + // The maintenance leader, running as role C, promotes scheduled jobs into + // lanes nobody has created yet: two priorities over two shards. + cluster.set_enqueue_shards("fresh_queue", 2).await?; + schedule(&cluster.app_a, "fresh_queue", &[1, 4]).await?; + assert_eq!( + store() + .promote_due(&cluster.app_c, JobState::Scheduled, 10) + .await?, + 2 + ); + let promoted_shards: Vec = sqlx::query_scalar( + "SELECT enqueue_shard FROM awa.queue_enqueue_heads \ + WHERE queue = 'fresh_queue' AND priority IN (1, 4) ORDER BY 1", + ) + .fetch_all(&cluster.owner) + .await?; + assert_eq!(promoted_shards, vec![0, 1]); + + // A and B both use the lanes the leader created; each producer's rotor + // visits both shards. + let (producer_a, producer_b) = (store(), store()); + for _ in 0..2 { + producer_a + .enqueue_batch(&cluster.app_a, "fresh_queue", 4, 1) + .await?; + producer_b + .enqueue_batch(&cluster.app_b, "fresh_queue", 1, 1) + .await?; + } + assert_eq!(drain(&cluster.app_a, "fresh_queue").await?, 6); + + cluster + .assert_lanes_owned_by_owner("fresh_queue", 5) + .await?; + assert!(storage::lane_sequence_problems(&cluster.owner, "awa") + .await? + .is_empty()); + cluster.teardown().await +} + +#[tokio::test] +async fn lane_first_touched_by_an_app_role_with_schema_create_serves_every_runtime_role( +) -> TestResult { + first_touch_by_one_app_role(true).await +} + +#[tokio::test] +async fn lane_first_touched_by_an_app_role_without_schema_create_serves_every_runtime_role( +) -> TestResult { + first_touch_by_one_app_role(false).await +} + +#[tokio::test] +async fn raising_enqueue_shards_live_lets_the_leader_create_a_shard_first() -> TestResult { + let cluster = Cluster::new(Schema::Current, true).await?; + enqueue(&cluster.app_a, "sharded_queue", 2, 3).await?; + + cluster.set_enqueue_shards("sharded_queue", 2).await?; + // The leader's store has already routed a promotion to shard 0, as a + // long-running leader has, so its next pick is the new shard. + let leader = store(); + schedule(&cluster.app_a, "sharded_queue", &[2]).await?; + assert_eq!( + leader + .promote_due(&cluster.app_c, JobState::Scheduled, 10) + .await?, + 1 + ); + schedule(&cluster.app_b, "sharded_queue", &[2]).await?; + assert_eq!( + leader + .promote_due(&cluster.app_c, JobState::Scheduled, 10) + .await?, + 1 + ); + let shards: Vec = sqlx::query_scalar( + "SELECT enqueue_shard FROM awa.queue_enqueue_heads WHERE queue = 'sharded_queue' ORDER BY 1", + ) + .fetch_all(&cluster.owner) + .await?; + assert_eq!(shards, vec![0, 1]); + + // Producers on both roles spread over both shards; B drains them. + let producer_a = store(); + let producer_b = store(); + for _ in 0..4 { + producer_a + .enqueue_batch(&cluster.app_a, "sharded_queue", 2, 1) + .await?; + producer_b + .enqueue_batch(&cluster.app_b, "sharded_queue", 2, 1) + .await?; + } + assert_eq!(drain(&cluster.app_b, "sharded_queue").await?, 13); + cluster + .assert_lanes_owned_by_owner("sharded_queue", 2) + .await?; + cluster.teardown().await +} + +#[tokio::test] +async fn prepare_queue_as_an_app_role_provisions_lanes_every_runtime_role_can_use() -> TestResult { + let cluster = Cluster::new(Schema::Current, true).await?; + store() + .prepare_queue(&cluster.app_a, "prepared_queue", 2) + .await?; + cluster + .assert_lanes_owned_by_owner("prepared_queue", 8) + .await?; + + cluster.set_enqueue_shards("prepared_queue", 2).await?; + let producer = store(); + for priority in 1..=4 { + producer + .enqueue_batch(&cluster.app_b, "prepared_queue", priority, 1) + .await?; + producer + .enqueue_batch(&cluster.app_c, "prepared_queue", priority, 1) + .await?; + } + assert_eq!(drain(&cluster.app_b, "prepared_queue").await?, 8); + cluster + .assert_lanes_owned_by_owner("prepared_queue", 8) + .await?; + cluster.teardown().await +} + +#[tokio::test] +async fn mixed_fleet_shares_lanes_whichever_path_creates_them() -> TestResult { + let cluster = Cluster::new(Schema::Current, true).await?; + // Role A cannot call provision_lane, so its stores take the path a binary + // that predates the function takes: they write the lane rows themselves. + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "REVOKE EXECUTE ON FUNCTION awa.provision_lane(text, smallint, smallint) FROM PUBLIC, {}", + cluster.runtime_name + )))) + .await?; + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "GRANT EXECUTE ON FUNCTION awa.provision_lane(text, smallint, smallint) TO {}, {}", + cluster.roles[4], cluster.roles[5] + )))) + .await?; + let can_provision: bool = sqlx::query_scalar( + "SELECT has_function_privilege('awa.provision_lane(text,smallint,smallint)', 'EXECUTE')", + ) + .fetch_one(&cluster.app_a) + .await?; + assert!(!can_provision); + + enqueue(&cluster.app_a, "fleet_queue", 1, 2).await?; + enqueue(&cluster.app_b, "fleet_queue", 2, 2).await?; + enqueue(&cluster.app_b, "fleet_queue", 1, 2).await?; + enqueue(&cluster.app_a, "fleet_queue", 2, 2).await?; + assert_eq!(drain(&cluster.app_a, "fleet_queue").await?, 8); + cluster + .assert_lanes_owned_by_owner("fleet_queue", 2) + .await?; + + // A provisioner that disappears after a store has used it: the next lane + // falls back inside the same transaction instead of failing the enqueue. + let producer = store(); + producer + .enqueue_batch(&cluster.app_b, "fleet_queue", 3, 1) + .await?; + cluster + .owner + .execute("DROP FUNCTION awa.provision_lane(text, smallint, smallint)") + .await?; + producer + .enqueue_batch(&cluster.app_b, "fleet_queue", 4, 1) + .await?; + enqueue(&cluster.app_c, "fleet_queue", 4, 1).await?; + assert_eq!(drain(&cluster.app_c, "fleet_queue").await?, 3); + cluster + .assert_lanes_owned_by_owner("fleet_queue", 4) + .await?; + cluster.teardown().await +} + +#[tokio::test] +async fn binary_ahead_of_the_schema_creates_lanes_as_before() -> TestResult { + let cluster = Cluster::new(Schema::BeforeProvisioning, false).await?; + let has_provisioner: bool = sqlx::query_scalar( + "SELECT to_regprocedure('awa.provision_lane(text,smallint,smallint)') IS NOT NULL", + ) + .fetch_one(&cluster.owner) + .await?; + assert!(!has_provisioner); + + // The owner creates lanes lazily, as on v052. + enqueue(&cluster.owner, "ahead_queue", 2, 2).await?; + enqueue(&cluster.app_b, "ahead_queue", 2, 1).await?; + assert_eq!(drain(&cluster.app_b, "ahead_queue").await?, 3); + // A runtime without schema CREATE still gets v052's provisioning hint. + let error = enqueue(&cluster.app_a, "ahead_queue", 3, 1) + .await + .unwrap_err(); + let AwaError::Database(sqlx::Error::Database(db_error)) = &error else { + panic!("expected a database error, got {error:?}"); + }; + assert_eq!(db_error.code().as_deref(), Some("42501")); + assert!(db_error.message().contains("has not been provisioned")); + cluster.teardown().await +} + +#[tokio::test] +async fn concurrent_first_touch_of_one_lane_by_two_roles() -> TestResult { + let cluster = Cluster::new(Schema::Current, true).await?; + + // A holds a brand-new lane's creation open; B's first touch waits for it + // and then uses A's sequences. + let mut holding = cluster.app_a.begin().await?; + sqlx::query("SELECT awa.provision_lane('raced_queue', 2::smallint, 0::smallint)") + .execute(holding.as_mut()) + .await?; + let app_b = cluster.app_b.clone(); + let waiting = tokio::spawn(async move { enqueue(&app_b, "raced_queue", 2, 1).await }); + let mut blocked = false; + for _ in 0..100 { + blocked = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_locks WHERE locktype = 'advisory' AND NOT granted \ + AND database = (SELECT oid FROM pg_database WHERE datname = current_database()))", + ) + .fetch_one(&cluster.owner) + .await?; + if blocked { + break; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + assert!(blocked, "B's first touch should wait on A's lane lock"); + holding.commit().await?; + assert_eq!(waiting.await??, 1); + + // Unordered races on many fresh lanes. + for round in 0..20 { + let queue = format!("raced_queue_{round}"); + let (from_a, from_b) = tokio::join!( + enqueue(&cluster.app_a, &queue, 1, 1), + enqueue(&cluster.app_b, &queue, 1, 1) + ); + assert_eq!(from_a? + from_b?, 2); + cluster.assert_lanes_owned_by_owner(&queue, 1).await?; + assert_eq!(drain(&cluster.app_c, &queue).await?, 2); + } + assert_eq!(drain(&cluster.app_c, "raced_queue").await?, 1); + cluster + .assert_lanes_owned_by_owner("raced_queue", 1) + .await?; + cluster.teardown().await +} + +#[tokio::test] +async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { + let mut cluster = Cluster::new(Schema::BeforeProvisioning, true).await?; + // Lanes as a 0.6 fleet left them: one each created by an app role, by the + // migration role acting as itself, and by the owner without the runtime's + // grant. + enqueue(&cluster.app_a, "app_lane", 1, 1).await?; + enqueue(&cluster.migrator, "migrator_lane", 1, 1).await?; + enqueue(&cluster.owner, "owner_lane", 1, 1).await?; + for (sequence, _) in cluster.lane_sequence_owners("owner_lane").await? { + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "REVOKE ALL ON SEQUENCE awa.{sequence} FROM {}", + cluster.runtime_name + )))) + .await?; + } + for lane in ["owner_lane", "app_lane"] { + let failed = drain(&cluster.app_b, lane).await.unwrap_err(); + assert!( + failed + .to_string() + .contains("permission denied for sequence"), + "{failed}" + ); + } + // 0.7 refuses to migrate while app roles hold CREATE on the schema. + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "REVOKE CREATE ON SCHEMA awa FROM {}", + cluster.runtime_name + )))) + .await?; + + // `awa migrate` connects as the migration role and acts as the owner. As + // the owner it adds the grants the owner's sequences lack; the migration + // role's own sequences need the migration role's privileges, so they are + // reassigned as the login role. Neither can act as the app role. + let migrate_pool = cluster.migrator_as_owner().await?; + tokio::time::timeout(Duration::from_secs(120), migrations::run(&migrate_pool)).await??; + assert_eq!( + migrations::current_version(&cluster.owner).await?, + migrations::CURRENT_VERSION + ); + migrate_pool.close().await; + cluster + .assert_lanes_owned_by_owner("migrator_lane", 1) + .await?; + cluster.assert_lanes_owned_by_owner("owner_lane", 1).await?; + assert_eq!(drain(&cluster.app_b, "migrator_lane").await?, 1); + assert_eq!(drain(&cluster.app_b, "owner_lane").await?, 1); + + let app_a = cluster.roles[3].clone(); + let app_a = &app_a; + let problems = storage::lane_sequence_problems(&cluster.app_c, "awa").await?; + assert_eq!(problems.len(), 2, "{problems:?}"); + for problem in &problems { + assert_eq!(&problem.owner, app_a); + assert_eq!(problem.expected_owner, cluster.owner_name); + assert_eq!( + problem.missing_grants, + vec![format!( + "GRANT USAGE, UPDATE, SELECT ON SEQUENCE awa.{} TO {}", + problem.sequence_name, cluster.runtime_name + )] + ); + assert_eq!( + problem.reassign_sql.as_deref(), + Some( + format!( + "ALTER SEQUENCE awa.{} OWNER TO {}", + problem.sequence_name, cluster.owner_name + ) + .as_str() + ) + ); + assert!(problem.repair_roles.is_empty(), "{problem:?}"); + } + let next_steps = + storage::lane_sequence_next_steps(problems.iter().map(|problem| (problem, false))); + assert_eq!(next_steps.len(), 1, "{next_steps:?}"); + assert!( + next_steps[0].contains(&format!("no login role can both act as {app_a}")), + "{next_steps:?}" + ); + + // Re-running the repair as the migration role changes nothing and fails + // nothing. + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "GRANT EXECUTE ON FUNCTION awa.repair_lane_sequence(text, text) TO {}", + cluster.runtime_name + )))) + .await?; + let rerun = + storage::repair_lane_sequences(&mut *cluster.migrator.acquire().await?, "awa").await?; + assert_eq!(rerun.len(), 2); + assert!(rerun + .iter() + .all(|result| result.outcome == LaneSequenceOutcome::Unrepaired)); + + // A role that can act as the app role but cannot become the owner adds + // the grants and leaves the owner for someone who can. + let delegate = format!("{}_delegate", cluster.runtime_name); + cluster + .admin + .execute(sqlx::raw_sql(sql(format!( + "CREATE ROLE {delegate} LOGIN PASSWORD 'lane_test' IN ROLE {app_a}; \ + GRANT CONNECT ON DATABASE {} TO {delegate}", + cluster.database + )))) + .await?; + let delegate_pool = PgPoolOptions::new() + .max_connections(1) + .connect_with(cluster.connect_options.clone().username(&delegate)) + .await?; + let delegated = + storage::repair_lane_sequences(&mut *delegate_pool.acquire().await?, "awa").await?; + delegate_pool.close().await; + assert_eq!(delegated.len(), 2); + for result in &delegated { + assert_eq!(result.outcome, LaneSequenceOutcome::OwnerUnchanged); + assert_eq!(&result.owner, app_a); + assert!(result.missing_grants.is_empty()); + assert!(result.reassign_sql.is_some()); + } + assert_eq!(drain(&cluster.app_b, "app_lane").await?, 1); + enqueue(&cluster.app_c, "app_lane", 1, 1).await?; + + // Once the migration role is a member of the app role too, the check names + // it, and a repair connected as it reassigns. + cluster + .admin + .execute(sqlx::raw_sql(sql(format!( + "GRANT {app_a} TO {migrator}", + migrator = cluster.migrator_name + )))) + .await?; + let problems = storage::lane_sequence_problems(&cluster.app_c, "awa").await?; + assert!(problems + .iter() + .all(|problem| problem.repair_roles == vec![cluster.migrator_name.clone()])); + let repaired = + storage::repair_lane_sequences(&mut *cluster.migrator.acquire().await?, "awa").await?; + assert_eq!(repaired.len(), 2); + assert!( + repaired + .iter() + .all(|result| result.outcome == LaneSequenceOutcome::Repaired), + "{repaired:?}" + ); + assert!(storage::lane_sequence_problems(&cluster.app_c, "awa") + .await? + .is_empty()); + cluster.assert_lanes_owned_by_owner("app_lane", 1).await?; + assert_eq!(drain(&cluster.app_b, "app_lane").await?, 1); + + cluster.roles.push(delegate); + cluster.teardown().await +} + +/// Reassigning a lane sequence waits on any open enqueue or claim on the lane, +/// so the repair bounds that wait and holds each sequence's lock only for its +/// own short transaction. A held sequence is skipped and reported; the rest +/// are repaired, and a later repair finishes it. +#[tokio::test] +async fn migration_skips_a_lane_sequence_an_open_transaction_holds() -> TestResult { + let cluster = Cluster::new(Schema::BeforeProvisioning, true).await?; + enqueue(&cluster.migrator, "held_lane", 1, 1).await?; + enqueue(&cluster.migrator, "free_lane", 1, 1).await?; + let held_sequence = cluster.enqueue_sequence("held_lane").await?; + // 0.7 refuses to migrate while app roles hold CREATE on the schema. + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "REVOKE CREATE ON SCHEMA awa FROM {}", + cluster.runtime_name + )))) + .await?; + + // An enqueue in flight on held_lane holds its sequence's lock. + let mut holder = cluster.app_b.begin().await?; + sqlx::query("SELECT nextval($1::regclass)") + .bind(format!("awa.{held_sequence}")) + .execute(&mut *holder) + .await?; + + let migrate_pool = cluster.migrator_as_owner().await?; + // The migration completes while the lock is held. + tokio::time::timeout(Duration::from_secs(120), migrations::run(&migrate_pool)).await??; + for (sequence, owner) in cluster.lane_sequence_owners("held_lane").await? { + let expected = if sequence == held_sequence { + &cluster.migrator_name + } else { + &cluster.owner_name + }; + assert_eq!(&owner, expected, "{sequence}"); + } + cluster.assert_lanes_owned_by_owner("free_lane", 1).await?; + + // A repair while the transaction is still open reports the sequence busy. + let mut conn = migrate_pool.acquire().await?; + let busy = storage::repair_lane_sequences(&mut conn, "awa").await?; + assert_eq!(busy.len(), 1, "{busy:?}"); + assert_eq!(busy[0].sequence_name, held_sequence); + assert_eq!(busy[0].outcome, LaneSequenceOutcome::Busy); + assert_eq!(busy[0].owner, cluster.migrator_name); + let next_steps = storage::lane_sequence_repair_next_steps(&busy); + assert_eq!(next_steps.len(), 1); + assert!(next_steps[0].contains("busy"), "{next_steps:?}"); + assert!( + next_steps[0].contains("awa storage lane-sequences --schema awa --repair"), + "{next_steps:?}" + ); + + holder.rollback().await?; + let repaired = storage::repair_lane_sequences(&mut conn, "awa").await?; + assert_eq!(repaired.len(), 1, "{repaired:?}"); + assert_eq!(repaired[0].outcome, LaneSequenceOutcome::Repaired); + drop(conn); + migrate_pool.close().await; + cluster.assert_lanes_owned_by_owner("held_lane", 1).await?; + assert!(storage::lane_sequence_problems(&cluster.app_c, "awa") + .await? + .is_empty()); + assert_eq!(drain(&cluster.app_b, "held_lane").await?, 1); + cluster.teardown().await +} diff --git a/awa/tests/migration_test.rs b/awa/tests/migration_test.rs index 53c1ab009..10fa18da8 100644 --- a/awa/tests/migration_test.rs +++ b/awa/tests/migration_test.rs @@ -4213,9 +4213,15 @@ async fn test_newer_schema_guard_leaves_normal_paths_untouched() { .execute(&pool) .await .expect("drop the newest version row to simulate a one-behind schema"); + let previous_version = migrations::migration_sql() + .iter() + .map(|(version, _, _)| *version) + .filter(|version| *version < migrations::CURRENT_VERSION) + .max() + .expect("a migration precedes CURRENT_VERSION"); assert_eq!( migrations::current_version(&pool).await.unwrap(), - migrations::CURRENT_VERSION - 1 + previous_version ); migrations::run(&pool) .await @@ -4543,3 +4549,65 @@ async fn v048_generated_bundle_refreshes_a_stale_installer() { .unwrap(); assert_eq!(version, 48); } + +#[tokio::test] +async fn test_v056_upgrades_from_output_and_batch_schemas_and_is_rerunnable() { + let _guard = acquire_migration_guard().await; + let pool = pool().await; + for from in [52, 55] { + reset_schema(&pool).await; + for (_, _, sql) in migrations::migration_sql_range(0, from) { + sqlx::raw_sql(audited_sql(sql)) + .execute(&pool) + .await + .unwrap(); + } + // A previous binary's installed schema lacks this optional function. + let signature: String = sqlx::query_scalar( + "SELECT oid::regprocedure::text FROM pg_proc \ + WHERE pronamespace = 'awa'::regnamespace AND proname = 'enqueue_ready_rows'", + ) + .fetch_one(&pool) + .await + .unwrap(); + sqlx::query(audited_sql(format!("DROP FUNCTION {signature}"))) + .execute(&pool) + .await + .unwrap(); + let before = QueueStorage::new(Default::default()).unwrap(); + before + .enqueue_batch(&pool, "v056_upgrade", 2, 2) + .await + .unwrap(); + migrations::run(&pool).await.unwrap(); + assert_eq!(migrations::current_version(&pool).await.unwrap(), 56); + let after = QueueStorage::new(Default::default()).unwrap(); + after + .enqueue_batch(&pool, "v056_upgrade", 2, 2) + .await + .unwrap(); + let versions: Vec = sqlx::query_scalar( + "SELECT version FROM awa.schema_version WHERE version BETWEEN 51 AND 56 ORDER BY version", + ).fetch_all(&pool).await.unwrap(); + assert_eq!(versions, vec![51, 52, 53, 54, 55, 56]); + for _ in 0..2 { + let mut tx = pool.begin().await.unwrap(); + for (_, _, sql) in migrations::migration_sql_range(55, 56) { + sqlx::raw_sql(audited_sql(sql)) + .execute(tx.as_mut()) + .await + .unwrap(); + } + tx.commit().await.unwrap(); + } + assert_eq!( + after + .queue_counts(&pool, "v056_upgrade") + .await + .unwrap() + .available, + 4 + ); + } + pool.close().await; +} diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index bf51546fe..eeb9c3a00 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -6,6 +6,8 @@ #[path = "../../awa-model/src/storage_footprint_sql.rs"] mod storage_footprint_sql; +mod ci_timing; + use awa::audited_sql; use awa::model::{ admin, batch_operations, insert, migrations, storage, AwaError, BatchOperationFilter, @@ -449,6 +451,32 @@ async fn create_store_with_config(pool: &sqlx::PgPool, config: QueueStorageConfi store } +/// Prune's success/accounting assertions allow a transient competing lock. +/// Retry only the expected slot's Blocked result; errors and other outcomes +/// still reach the caller's original assertions without being retried. +async fn prune_until_not_blocked( + store: &QueueStorage, + pool: &sqlx::PgPool, + retention: Duration, + slot: i32, +) -> PruneOutcome { + tokio::time::timeout(ci_timing::scaled_timeout(Duration::from_secs(5)), async { + loop { + let outcome = store + .prune_oldest(pool, retention) + .await + .expect("prune_oldest"); + if matches!(outcome, PruneOutcome::Blocked { slot: blocked } if blocked == slot) { + tokio::time::sleep(Duration::from_millis(10)).await; + } else { + return outcome; + } + } + }) + .await + .expect("prune remained Blocked past the assertion deadline") +} + async fn create_store(pool: &sqlx::PgPool, schema: &str) -> QueueStorage { // Tests that go through this helper exercise the legacy // (non-receipts) lease-materialization path. The receipts mode @@ -13614,11 +13642,7 @@ async fn test_queue_storage_prune_re_carries_failed_row_exactly_once() { RotateOutcome::Rotated { slot: 1, .. } => {} other => panic!("expected Rotated {{ slot: 1 }}, got {other:?}"), } - match store - .prune_oldest(&pool, Duration::from_secs(3600)) - .await - .expect("prune_oldest") - { + match prune_until_not_blocked(&store, &pool, Duration::from_secs(3600), 0).await { PruneOutcome::Pruned { slot: 0, carried_failed_rows: 1, @@ -13662,11 +13686,7 @@ async fn test_queue_storage_prune_re_carries_failed_row_exactly_once() { // Second prune: current slot is 0, so the oldest sealed non-current // slot is slot 1 (the carried row). It is carried again into slot 0. - match store - .prune_oldest(&pool, Duration::from_secs(3600)) - .await - .expect("prune_oldest re-carry") - { + match prune_until_not_blocked(&store, &pool, Duration::from_secs(3600), 1).await { PruneOutcome::Pruned { slot: 1, carried_failed_rows: 1, @@ -15622,16 +15642,298 @@ async fn test_queue_storage_prune_oldest_blocks_on_reader_lock() { .await .expect("Failed to release reader lock"); - let pruned = store - .prune_oldest(&pool, Duration::ZERO) - .await - .expect("Failed to prune after reader lock release"); + let pruned = prune_until_not_blocked(&store, &pool, Duration::ZERO, 0).await; assert!( matches!(pruned, PruneOutcome::Pruned { slot: 0, .. }), "unexpected prune outcome after reader lock release: {pruned:?}" ); } +/// A later insert can remain locked after the explicit partition locks pass. +/// Observe two distinct blocked attempts before releasing it, proving that the +/// eventual-success helper actually retries rather than relying on a sleep. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_prune_eventual_success_retries_a_post_child_lock_handoff() { + let (_db_guard, pool) = setup_pool(6).await; + let schema = "awa_qs_prune_retry_handoff"; + let queue = "prune_retry_handoff"; + let store = Arc::new(create_store(&pool, schema).await); + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 1 }, + InsertOpts { + queue: queue.into(), + ..Default::default() + }, + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::from_secs(30)) + .await + .expect("claim carry job") + .remove(0); + store + .fail_terminal(&pool, job_id, claimed.job.run_lease, "boom", None) + .await + .expect("fail carry job") + .expect("running job should fail"); + assert!(matches!( + store.rotate(&pool).await.unwrap(), + RotateOutcome::Rotated { slot: 1, .. } + )); + let retention = Duration::from_secs(3600); + let mut holder = pool.begin().await.expect("begin post-child blocker"); + sqlx::query(audited_sql(format!( + "LOCK TABLE {schema}.queue_terminal_count_deltas_1 IN SHARE MODE" + ))) + .execute(holder.as_mut()) + .await + .expect("hold carry destination"); + assert!(matches!( + store.prune_oldest(&pool, retention).await.unwrap(), + PruneOutcome::Blocked { slot: 0 } + )); + assert_eq!(done_entries_count(&pool, schema, queue).await, 1); + assert_eq!(pruned_rollup_sums(&pool, schema, queue).await, (0, 0)); + let worker_store = store.clone(); + let worker_pool = pool.clone(); + let worker = tokio::spawn(async move { + prune_until_not_blocked(&worker_store, &worker_pool, retention, 0).await + }); + let handoff = tokio::time::timeout(ci_timing::scaled_timeout(Duration::from_secs(5)), async { + let mut first_start: Option> = None; + loop { + if worker.is_finished() { + break false; + } + let waiting: Option> = sqlx::query_scalar( + "SELECT query_start FROM pg_stat_activity + WHERE datname = current_database() AND pid <> pg_backend_pid() + AND wait_event_type = 'Lock' AND query LIKE $1 + ORDER BY query_start DESC LIMIT 1", + ) + .bind(format!( + "%INSERT INTO {schema}.queue_terminal_count_deltas%" + )) + .fetch_optional(&pool) + .await + .expect("observe blocked carry insert"); + if let Some(started) = waiting { + if first_start.is_some_and(|first| started > first) { + break true; + } + first_start.get_or_insert(started); + } + tokio::time::sleep(Duration::from_millis(1)).await; + } + }) + .await; + holder.rollback().await.expect("release carry destination"); + let pruned = worker.await.expect("join eventual prune"); + assert!( + handoff.expect("helper did not reach a retry before its deadline"), + "helper ended before the controlled retry handoff: {pruned:?}" + ); + assert!( + matches!( + pruned, + PruneOutcome::Pruned { + slot: 0, + carried_failed_rows: 1, + .. + } + ), + "{pruned:?}" + ); + assert_eq!(done_entries_count(&pool, schema, queue).await, 1); + let carried_slot: i32 = sqlx::query_scalar(audited_sql(format!( + "SELECT ready_slot FROM {schema}.done_entries WHERE job_id = $1" + ))) + .bind(job_id) + .fetch_one(&pool) + .await + .expect("carried row slot"); + assert_eq!(carried_slot, 1); + assert_invariant_holds(&pool, schema, queue, "after retry handoff").await; + assert_eq!(store.queue_counts(&pool, queue).await.unwrap().terminal, 1); +} + +/// The prune `lock_timeout` covers every statement after it is set, not +/// only the `LOCK TABLE` on the slot's children. A statement that waits on +/// another table's lock — here the failed-row carry, its count deltas, the +/// rollup delta append and the live-counter cleanup — must roll the whole prune back +/// and report `Blocked`, leaving the slot to the next attempt. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_prune_lock_timeout_after_child_locks_reports_blocked() { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_prune_late_lock_timeout"; + let schema = "awa_qs_prune_late_lock_timeout"; + let store = create_store(&pool, schema).await; + + store + .enqueue_batch(&pool, queue, 1, 1) + .await + .expect("enqueue completed job"); + let claimed = store + .claim_batch(&pool, queue, 1) + .await + .expect("claim completed job"); + store + .complete_batch(&pool, &claimed) + .await + .expect("complete job"); + + let failed_job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 1 }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::from_secs(30)) + .await + .expect("claim failing job") + .into_iter() + .next() + .expect("missing claimed job"); + store + .fail_terminal(&pool, failed_job_id, claimed.job.run_lease, "boom", None) + .await + .expect("fail job") + .expect("running job should fail"); + + let rotated = store.rotate(&pool).await.expect("rotate"); + assert!( + matches!(rotated, RotateOutcome::Rotated { slot: 1, .. }), + "unexpected rotate outcome: {rotated:?}" + ); + + let retention = Duration::from_secs(3600); + for held in [ + format!("{schema}.done_entries_1"), + format!("{schema}.queue_terminal_count_deltas_1"), + format!("{schema}.queue_terminal_rollup_deltas"), + format!("{schema}.queue_terminal_live_counts"), + ] { + let mut holder = pool.begin().await.expect("begin lock holder"); + sqlx::query(audited_sql(format!("LOCK TABLE {held} IN SHARE MODE"))) + .execute(holder.as_mut()) + .await + .expect("hold share lock"); + + let outcome = store.prune_oldest(&pool, retention).await; + assert!( + matches!(outcome, Ok(PruneOutcome::Blocked { slot: 0 })), + "prune waiting on {held} should report Blocked, got {outcome:?}" + ); + + holder.rollback().await.expect("release share lock"); + + assert_eq!( + done_entries_count(&pool, schema, queue).await, + 2, + "blocked prune on {held} must not truncate or carry rows" + ); + let rollup_rows: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*)::bigint FROM {schema}.queue_terminal_rollup_deltas" + ))) + .fetch_one(&pool) + .await + .expect("count rollups"); + assert_eq!( + rollup_rows, 0, + "blocked prune on {held} must not fold rollups" + ); + assert_invariant_holds( + &pool, + schema, + queue, + &format!("after prune blocked on {held}"), + ) + .await; + } + + let pruned = store + .prune_oldest(&pool, retention) + .await + .expect("prune after locks released"); + assert!( + matches!( + pruned, + PruneOutcome::Pruned { + slot: 0, + carried_failed_rows: 1, + .. + } + ), + "unexpected prune outcome after locks released: {pruned:?}" + ); + assert_invariant_holds(&pool, schema, queue, "after prune").await; + let counts = store.queue_counts(&pool, queue).await.expect("counts"); + assert_eq!(counts.terminal, 2, "rollup (1) + carried live row (1)"); +} + +/// Terminal-delta rollup shares prune's lock timeout. A timeout on the +/// live-counter fold, after the delta child is locked, must report a blocked +/// slot rather than fail the rollup pass, and the pass stops there instead of +/// waiting out the same timeout on every remaining slot. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_terminal_count_delta_rollup_lock_timeout_after_child_lock_reports_blocked() { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_terminal_delta_rollup_late_lock"; + let schema = "awa_qs_terminal_delta_rollup_late_lock"; + let store = create_store(&pool, schema).await; + + for _ in 0..2 { + store + .enqueue_batch(&pool, queue, 1, 3) + .await + .expect("enqueue"); + let claimed = store.claim_batch(&pool, queue, 3).await.expect("claim"); + store + .complete_batch(&pool, &claimed) + .await + .expect("complete"); + match store.rotate(&pool).await.expect("rotate") { + RotateOutcome::Rotated { .. } => {} + other => panic!("expected Rotated, got {other:?}"), + } + } + + let mut holder = pool.begin().await.expect("begin lock holder"); + sqlx::query(audited_sql(format!( + "LOCK TABLE {schema}.queue_terminal_live_counts IN SHARE MODE" + ))) + .execute(holder.as_mut()) + .await + .expect("hold share lock"); + + let outcome = store.rollup_terminal_count_deltas(&pool, 4).await; + holder.rollback().await.expect("release share lock"); + let outcome = outcome.expect("rollup waiting on live counts should not error"); + assert_eq!(outcome.rolled_slots, 0); + assert_eq!( + outcome.blocked_slots, 1, + "the pass stops at the first blocked slot" + ); + assert_eq!(live_count_sum(&pool, schema, queue).await, 0); + assert_eq!(terminal_delta_sum(&pool, schema, queue).await, 6); + + let outcome = store + .rollup_terminal_count_deltas(&pool, 4) + .await + .expect("rollup after lock released"); + assert_eq!(outcome.rolled_slots, 2); + assert_eq!(outcome.blocked_slots, 0); + assert_eq!(live_count_sum(&pool, schema, queue).await, 6); + assert_eq!(terminal_delta_sum(&pool, schema, queue).await, 0); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_queue_prune_allows_segment_reader_to_finish_claim_lock_order() { let (_db_guard, pool) = setup_pool(8).await; @@ -17371,10 +17673,25 @@ async fn test_priority_aging_off_does_not_stamp_original() { /// helper's comment. #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn test_queue_storage_ensure_lane_cache_recovers_after_rollback() { + check_ensure_lane_cache_recovers_after_rollback("awa_qs_ensure_lane_rollback", false).await; +} + +/// The multi-statement path is the one that fills the Rust lane cache, so +/// the recovery must also hold when the function is not used. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_queue_storage_ensure_lane_cache_recovers_after_rollback_without_function() { + check_ensure_lane_cache_recovers_after_rollback("awa_qs_ensure_lane_rollback_nf", true).await; +} + +async fn check_ensure_lane_cache_recovers_after_rollback(schema: &str, without_function: bool) { let (_db_guard, pool) = setup_pool(4).await; let queue = "qs_ensure_lane_rollback"; - let schema = "awa_qs_ensure_lane_rollback"; let store = create_store(&pool, schema).await; + let store = if without_function { + store.without_enqueue_function() + } else { + store + }; // Simulate the rolled-back ensure_lane: run the three lane // inserts in a transaction that we explicitly roll back. After @@ -18684,6 +19001,66 @@ async fn test_compact_identity_closers_do_not_lock_batch_and_race_rescue_once() pool.close().await; } +/// A NULL count would return a start without advancing the lane sequence, +/// handing the same range to the next caller. +#[tokio::test] +async fn test_reserve_enqueue_seq_rejects_a_null_count() { + let (_db_guard, pool) = setup_pool(2).await; + let schema = "awa_qs_reserve_null_count"; + let store = create_store(&pool, schema).await; + store.enqueue_batch(&pool, "q", 1, 1).await.unwrap(); + let error = sqlx::query(audited_sql(format!( + "SELECT {schema}.reserve_enqueue_seq('q', 1::smallint, 0::smallint, NULL::bigint)" + ))) + .execute(&pool) + .await + .expect_err("a NULL count must be rejected"); + assert_eq!( + error + .as_database_error() + .and_then(|db| db.code()) + .as_deref(), + Some("22004") + ); + pool.close().await; +} + +/// A store that has seen enqueue_ready_rows must not keep calling it once it +/// is gone: the failed SQL call forgets the observation, and the same enqueue +/// rolls back and takes the multi-statement path once. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_enqueue_falls_back_after_the_enqueue_function_disappears() { + let (_db_guard, pool) = setup_pool(4).await; + let schema = "awa_qs_enqueue_function_vanishes"; + let store = create_store(&pool, schema).await; + assert_eq!(store.enqueue_batch(&pool, "q", 1, 1).await.unwrap(), 1); + + let signature: String = sqlx::query_scalar(audited_sql(format!( + "SELECT oid::regprocedure::text FROM pg_proc \ + WHERE pronamespace = '{schema}'::regnamespace AND proname = 'enqueue_ready_rows'" + ))) + .fetch_one(&pool) + .await + .unwrap(); + sqlx::query(audited_sql(format!("DROP FUNCTION {signature}"))) + .execute(&pool) + .await + .unwrap(); + + assert_eq!( + store.enqueue_batch(&pool, "q", 1, 1).await.unwrap(), + 1, + "the same call falls back when the cached function disappears" + ); + assert_eq!( + store.enqueue_batch(&pool, "q", 1, 1).await.unwrap(), + 1, + "the next enqueue probes again and uses the multi-statement path" + ); + assert_eq!(store.queue_counts(&pool, "q").await.unwrap().available, 3); + pool.close().await; +} + /// Holds the claim slot's cursor row so the rescue statement takes its /// snapshot, then waits; the attempt completes and commits meanwhile. Once /// released, rescue must see that completion instead of reopening the @@ -21093,3 +21470,397 @@ async fn test_queue_storage_legacy_row_claims_are_locked_before_the_heartbeat_st .expect("join heartbeat") .expect("heartbeat after the rewrite"); } + +fn fold_test_config(schema: &str) -> QueueStorageConfig { + QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + lease_claim_receipts: false, + ..Default::default() + } +} + +fn fold_test_jobs( + queue: &str, + ids: std::ops::RangeInclusive, +) -> Vec { + let opts = InsertOpts { + queue: queue.to_string(), + ..Default::default() + }; + ids.map(|id| insert::params_with(&DlqJob { id }, opts.clone()).expect("Failed to prepare job")) + .collect() +} + +/// A runtime that enqueues through `enqueue_ready_rows` and one that uses the +/// multi-statement path (as a binary without the function does) interleave +/// on one lane: the lane stays gap-free and FIFO, each call writes its own +/// segment, and a store on a schema without the function falls back. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_enqueue_function_and_fallback_share_one_lane() { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_enqueue_function_lane"; + let schema = "awa_qs_enqueue_function_lane"; + let store = create_store_with_config(&pool, fold_test_config(schema)).await; + let fallback = QueueStorage::new(fold_test_config(schema)) + .expect("Failed to create fallback store") + .without_enqueue_function(); + + let mut next_id = 1_i64; + for (through_function, batch_size) in [(true, 3), (false, 1), (true, 1), (false, 4), (true, 2)] + { + let jobs = fold_test_jobs(queue, next_id..=next_id + batch_size - 1); + next_id += batch_size; + let writer = if through_function { &store } else { &fallback }; + let inserted = writer + .enqueue_params_copy(&pool, &jobs) + .await + .expect("Failed to enqueue"); + assert_eq!(inserted, batch_size as usize); + } + + let rows: Vec<(i64, i64)> = sqlx::query_as(audited_sql(format!( + "SELECT lane_seq, (args->>'id')::bigint FROM {schema}.ready_entries \ + WHERE queue = $1 ORDER BY lane_seq" + ))) + .bind(queue) + .fetch_all(&pool) + .await + .expect("Failed to read ready rows"); + assert_eq!(rows, (1..=11).map(|seq| (seq, seq)).collect::>()); + + let segments: Vec<(i64, i64)> = sqlx::query_as(audited_sql(format!( + "SELECT first_lane_seq, next_lane_seq FROM {schema}.ready_segments \ + WHERE queue = $1 ORDER BY first_lane_seq" + ))) + .bind(queue) + .fetch_all(&pool) + .await + .expect("Failed to read ready segments"); + assert_eq!(segments, vec![(1, 4), (4, 5), (5, 6), (6, 10), (10, 12)]); + + let distinct_job_ids: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(DISTINCT job_id) FROM {schema}.ready_entries WHERE queue = $1" + ))) + .bind(queue) + .fetch_one(&pool) + .await + .expect("Failed to count job ids"); + assert_eq!(distinct_job_ids, 11); + + let claimed = store + .claim_runtime_batch(&pool, queue, 11, Duration::ZERO) + .await + .expect("Failed to claim"); + let claimed_lane_seqs: Vec = claimed.iter().map(|job| job.claim.lane_seq).collect(); + assert_eq!(claimed_lane_seqs, (1..=11).collect::>()); + + sqlx::query(audited_sql(format!( + "DROP FUNCTION {schema}.enqueue_ready_rows(text[], smallint[], smallint[], bigint[], \ + bigint[], text[], jsonb[], smallint[], bigint[], smallint[], timestamptz[], \ + timestamptz[], timestamptz[], bytea[], text[], jsonb[], integer[], bigint[], bigint[], \ + timestamptz[], text[])" + ))) + .execute(&pool) + .await + .expect("Failed to drop enqueue function"); + let older_schema_store = + QueueStorage::new(fold_test_config(schema)).expect("Failed to create store"); + older_schema_store + .enqueue_params_copy(&pool, &fold_test_jobs(queue, 12..=13)) + .await + .expect("Failed to enqueue without the function"); + let last_lane_seq: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT max(lane_seq) FROM {schema}.ready_entries WHERE queue = $1" + ))) + .bind(queue) + .fetch_one(&pool) + .await + .expect("Failed to read last lane sequence"); + assert_eq!(last_lane_seq, 13); +} + +/// The one-statement enqueue takes the same per-lane reservation lock as +/// `reserve_enqueue_seq`, so it cannot publish past an open reservation. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_enqueue_function_waits_for_open_lane_reservation() { + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_enqueue_function_waits"; + let schema = "awa_qs_enqueue_function_waits"; + let store = Arc::new(create_store_with_config(&pool, fold_test_config(schema)).await); + store + .prepare_queue(&pool, queue, 1) + .await + .expect("Failed to prepare lane"); + + let mut reservation = pool.begin().await.expect("Failed to begin reservation"); + let reserved: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT {schema}.reserve_enqueue_seq($1, $2, $3, $4)" + ))) + .bind(queue) + .bind(2_i16) + .bind(0_i16) + .bind(2_i64) + .fetch_one(reservation.as_mut()) + .await + .expect("Failed to reserve"); + assert_eq!(reserved, 1); + + let enqueue_pool = pool.clone(); + let enqueue_store = Arc::clone(&store); + let jobs = fold_test_jobs(queue, 1..=1); + let enqueue = tokio::spawn(async move { + enqueue_store + .enqueue_params_copy(&enqueue_pool, &jobs) + .await + }); + wait_for_lock_wait(&pool, "advisory").await; + let visible: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*) FROM {schema}.ready_entries WHERE queue = $1" + ))) + .bind(queue) + .fetch_one(&pool) + .await + .expect("Failed to count ready rows"); + assert_eq!(visible, 0); + + reservation + .rollback() + .await + .expect("Failed to release reservation"); + tokio::time::timeout(Duration::from_secs(5), enqueue) + .await + .expect("Enqueue remained blocked") + .expect("Enqueue task panicked") + .expect("Enqueue failed"); + let lane_seq: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT lane_seq FROM {schema}.ready_entries WHERE queue = $1" + ))) + .bind(queue) + .fetch_one(&pool) + .await + .expect("Failed to read lane sequence"); + assert_eq!(lane_seq, 3, "the rolled-back reservation leaves a gap"); +} + +/// A lane row without a sequence name (written before the head trigger +/// existed) is provisioned on its first reservation; reservations of one and +/// of several rows stay contiguous. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_reserve_names_unnamed_lane_and_stays_contiguous() { + let (_db_guard, pool) = setup_pool(4).await; + let queue = "qs_reserve_unnamed_lane"; + let schema = "awa_qs_reserve_unnamed_lane"; + let store = create_store_with_config(&pool, fold_test_config(schema)).await; + store + .prepare_queue(&pool, queue, 1) + .await + .expect("Failed to prepare lane"); + + let reserve = |count: i64| { + let pool = pool.clone(); + async move { + sqlx::query_scalar::<_, i64>(audited_sql(format!( + "SELECT {schema}.reserve_enqueue_seq($1, $2, $3, $4)" + ))) + .bind(queue) + .bind(2_i16) + .bind(0_i16) + .bind(count) + .fetch_one(&pool) + .await + .expect("Failed to reserve") + } + }; + assert_eq!(reserve(1).await, 1); + assert_eq!(reserve(3).await, 2); + + let ensure_calls_in_reserve = || { + let pool = pool.clone(); + async move { + let mut tx = pool.begin().await.expect("Failed to begin"); + sqlx::query("SET LOCAL track_functions = 'pl'") + .execute(tx.as_mut()) + .await + .expect("Failed to track functions"); + sqlx::query(audited_sql(format!( + "SELECT {schema}.reserve_enqueue_seq($1, 2::smallint, 0::smallint, 2)" + ))) + .bind(queue) + .execute(tx.as_mut()) + .await + .expect("Failed to reserve"); + let calls: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT COALESCE(pg_stat_get_xact_function_calls(\ + '{schema}.ensure_lane_sequences(text, smallint, smallint)'::regprocedure), 0)" + ))) + .fetch_one(tx.as_mut()) + .await + .expect("Failed to read function calls"); + tx.rollback().await.expect("Failed to roll back"); + calls + } + }; + assert_eq!( + ensure_calls_in_reserve().await, + 0, + "a named lane reserves without provisioning" + ); + assert_eq!( + reserve(1).await, + 7, + "the probe's rolled-back range is a gap" + ); + + sqlx::query(audited_sql(format!( + "UPDATE {schema}.queue_enqueue_heads SET seq_name = NULL \ + WHERE queue = $1 AND priority = 2 AND enqueue_shard = 0" + ))) + .bind(queue) + .execute(&pool) + .await + .expect("Failed to clear the sequence name"); + + assert_eq!( + ensure_calls_in_reserve().await, + 1, + "an unnamed lane is provisioned on reserve" + ); + sqlx::query(audited_sql(format!( + "UPDATE {schema}.queue_enqueue_heads SET seq_name = NULL \ + WHERE queue = $1 AND priority = 2 AND enqueue_shard = 0" + ))) + .bind(queue) + .execute(&pool) + .await + .expect("Failed to clear the sequence name"); + assert_eq!(reserve(1).await, 10); + assert_eq!(reserve(2).await, 11); + assert_eq!(reserve(0).await, 13, "a zero count reads the next value"); + let named: bool = sqlx::query_scalar(audited_sql(format!( + "SELECT seq_name IS NOT NULL FROM {schema}.queue_enqueue_heads \ + WHERE queue = $1 AND priority = 2 AND enqueue_shard = 0" + ))) + .bind(queue) + .fetch_one(&pool) + .await + .expect("Failed to read the sequence name"); + assert!(named); +} + +async fn drop_enqueue_function(pool: &sqlx::PgPool, schema: &str) { + let signature: String = sqlx::query_scalar( + "SELECT oid::regprocedure::text FROM pg_proc \ + WHERE pronamespace = $1::regnamespace AND proname = 'enqueue_ready_rows'", + ) + .bind(schema) + .fetch_one(pool) + .await + .unwrap(); + sqlx::query(audited_sql(format!("DROP FUNCTION {signature}"))) + .execute(pool) + .await + .unwrap(); +} + +#[tokio::test] +async fn test_enqueue_cached_function_falls_back_between_databases() { + let (_new_guard, new_pool) = setup_pool(4).await; + let (_old_guard, old_pool) = setup_pool(4).await; + let schema = "awa_qs_enqueue_two_databases"; + let store = create_store(&new_pool, schema).await; + let _old_store = create_store(&old_pool, schema).await; + drop_enqueue_function(&old_pool, schema).await; + store + .enqueue_params_batch(&new_pool, &fold_test_jobs("q", 1..=1)) + .await + .unwrap(); + assert_eq!( + store + .enqueue_params_batch(&old_pool, &fold_test_jobs("q", 2..=3)) + .await + .unwrap(), + 2, + "a stale positive hint cannot fail an otherwise valid enqueue" + ); + assert_eq!( + store + .enqueue_params_copy(&new_pool, &fold_test_jobs("q", 4..=5)) + .await + .unwrap(), + 2 + ); + assert_eq!( + store.queue_counts(&new_pool, "q").await.unwrap().available, + 3 + ); + assert_eq!( + store.queue_counts(&old_pool, "q").await.unwrap().available, + 2 + ); + new_pool.close().await; + old_pool.close().await; +} + +#[tokio::test] +async fn test_enqueue_batch_lost_function_rolls_back_unique_claims_before_fallback() { + enqueue_lost_function_rolls_back_unique_claims(false).await; +} + +#[tokio::test] +async fn test_enqueue_copy_lost_function_rolls_back_unique_claims_before_fallback() { + enqueue_lost_function_rolls_back_unique_claims(true).await; +} + +async fn enqueue_lost_function_rolls_back_unique_claims(copy: bool) { + let (_guard, pool) = setup_pool(4).await; + let schema = "awa_qs_enqueue_lost_unique"; + let store = create_store(&pool, schema).await; + store.enqueue_batch(&pool, "q", 2, 1).await.unwrap(); + drop_enqueue_function(&pool, schema).await; + let jobs: Vec<_> = [None, Some(Utc::now() + chrono::Duration::hours(1))] + .into_iter() + .enumerate() + .map(|(id, run_at)| { + insert::params_with( + &DlqJob { id: id as i64 }, + InsertOpts { + queue: "q".into(), + run_at, + unique: Some(UniqueOpts { + by_args: true, + ..Default::default() + }), + ..Default::default() + }, + ) + .unwrap() + }) + .collect(); + let inserted = if copy { + store.enqueue_params_copy(&pool, &jobs).await + } else { + store.enqueue_params_batch(&pool, &jobs).await + } + .expect("fallback must roll back its failed transaction before replaying"); + assert_eq!(inserted, 2); + let counts = store.queue_counts(&pool, "q").await.unwrap(); + assert_eq!(counts.available, 2); + let scheduled: i64 = sqlx::query_scalar(audited_sql(format!( + "SELECT count(*) FROM {schema}.deferred_jobs WHERE state = 'scheduled'" + ))) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(scheduled, 1); + let claims: i64 = sqlx::query_scalar("SELECT count(*) FROM awa.job_unique_claims") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!( + claims, 2, + "each successful unique job owns exactly one claim" + ); + pool.close().await; +} diff --git a/awa/tests/statistics_refresh_test.rs b/awa/tests/statistics_refresh_test.rs new file mode 100644 index 000000000..44fcca251 --- /dev/null +++ b/awa/tests/statistics_refresh_test.rs @@ -0,0 +1,820 @@ +//! The maintenance leader refreshes planner statistics for the queue-storage +//! tables autovacuum never analyses (the head tables and the partitioned ring +//! parents) on election and every interval, without delaying ring prune. + +use awa::model::{migrations, PruneOutcome, QueueStorage, QueueStorageConfig, RotateOutcome}; +use awa::{audited_sql, Client, JobContext, JobError, JobResult, QueueConfig, Worker}; +use opentelemetry::KeyValue; +use opentelemetry_sdk::metrics::data::{AggregatedMetrics, MetricData}; +use opentelemetry_sdk::metrics::{InMemoryMetricExporter, SdkMeterProvider}; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; +use sqlx::PgPool; +use std::str::FromStr; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::time::{Duration, Instant}; + +const ROLE_PASSWORD: &str = "statistics_refresh_test"; +const RING_PARENTS: [&str; 13] = [ + "done_entries", + "lease_claim_batches", + "lease_claim_closure_batches", + "lease_claim_closures", + "lease_claims", + "leases", + "queue_terminal_count_deltas", + "ready_claim_attempt_batches", + "ready_entries", + "ready_segments", + "ready_tombstones", + "receipt_completion_batches", + "receipt_completion_tombstones", +]; + +static DATABASE_COUNTER: AtomicU32 = AtomicU32::new(0); + +async fn run_sql(pool: &PgPool, sql: String) -> Result<(), sqlx::Error> { + sqlx::raw_sql(audited_sql(sql)) + .execute(pool) + .await + .map(|_| ()) +} + +fn base_database_url() -> String { + std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()) +} + +struct StatisticsProbeWorker; + +#[async_trait::async_trait] +impl Worker for StatisticsProbeWorker { + fn kind(&self) -> &'static str { + "statistics_probe_job" + } + + async fn perform(&self, _ctx: &JobContext) -> Result { + Ok(JobResult::Completed) + } +} + +/// A throwaway database with the awa schema migrated, optionally owned by a +/// non-superuser role the way a production schema owner is. +struct TestDatabase { + admin: PgPool, + name: String, + roles: Vec, + options: PgConnectOptions, + /// Connects as the schema owner (or the superuser when no owner role). + owner_pool: PgPool, +} + +impl TestDatabase { + async fn create(owner_role: Option<&str>, extra_roles: &[&str]) -> Self { + let admin = PgPoolOptions::new() + .max_connections(2) + .connect(&base_database_url()) + .await + .expect("connect admin pool"); + let name = format!( + "awa_stats_{}_{}", + std::process::id(), + DATABASE_COUNTER.fetch_add(1, Ordering::SeqCst) + ); + let mut roles = Vec::new(); + for role in owner_role.iter().chain(extra_roles.iter()) { + sqlx::raw_sql(audited_sql(format!( + "DROP ROLE IF EXISTS {role}; CREATE ROLE {role} LOGIN PASSWORD '{ROLE_PASSWORD}'" + ))) + .execute(&admin) + .await + .expect("create role"); + roles.push((*role).to_string()); + } + let owner_clause = owner_role + .map(|role| format!(" OWNER {role}")) + .unwrap_or_default(); + sqlx::raw_sql(audited_sql(format!("CREATE DATABASE {name}{owner_clause}"))) + .execute(&admin) + .await + .expect("create database"); + let options = PgConnectOptions::from_str(&base_database_url()) + .expect("parse DATABASE_URL") + .database(&name); + let owner_options = match owner_role { + Some(role) => options.clone().username(role).password(ROLE_PASSWORD), + None => options.clone(), + }; + let owner_pool = PgPoolOptions::new() + .max_connections(8) + .connect_with(owner_options) + .await + .expect("connect owner pool"); + migrations::run(&owner_pool).await.expect("run migrations"); + Self { + admin, + name, + roles, + options, + owner_pool, + } + } + + /// A pool for a role other than the schema owner, with the runtime grants + /// from the database-roles guide. + async fn runtime_pool(&self, role: &str) -> PgPool { + sqlx::raw_sql(audited_sql(format!( + r#" + GRANT USAGE ON SCHEMA awa TO {role}; + GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA awa TO {role}; + GRANT USAGE, SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA awa TO {role}; + GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa TO {role}; + REVOKE EXECUTE ON FUNCTION + awa.install_queue_storage_substrate(TEXT, INT, INT, INT, BOOLEAN) FROM {role}; + "# + ))) + .execute(&self.owner_pool) + .await + .expect("grant runtime privileges"); + PgPoolOptions::new() + .max_connections(8) + .connect_with(self.options.clone().username(role).password(ROLE_PASSWORD)) + .await + .expect("connect runtime pool") + } + + async fn drop(self) { + self.owner_pool.close().await; + let _ = sqlx::raw_sql(audited_sql(format!( + "DROP DATABASE IF EXISTS {} WITH (FORCE)", + self.name + ))) + .execute(&self.admin) + .await; + for role in &self.roles { + let _ = sqlx::raw_sql(audited_sql(format!("DROP ROLE IF EXISTS {role}"))) + .execute(&self.admin) + .await; + } + self.admin.close().await; + } +} + +async fn wait_for(what: &str, timeout: Duration, mut check: F) +where + F: FnMut() -> Fut, + Fut: std::future::Future, +{ + let deadline = Instant::now() + timeout; + while Instant::now() < deadline { + if check().await { + return; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + panic!("timed out waiting for {what}"); +} + +async fn stats_rows(pool: &PgPool, schema: &str, table: &str, inherited_only: bool) -> i64 { + sqlx::query_scalar( + "SELECT count(*) FROM pg_stats + WHERE schemaname = $1 AND tablename = $2 AND (inherited OR NOT $3)", + ) + .bind(schema) + .bind(table) + .bind(inherited_only) + .fetch_one(pool) + .await + .expect("read pg_stats") +} + +async fn parent_reltuples(pool: &PgPool, schema: &str, table: &str) -> f32 { + sqlx::query_scalar("SELECT reltuples FROM pg_class WHERE oid = to_regclass($1)") + .bind(format!("{schema}.{table}")) + .fetch_one(pool) + .await + .expect("read reltuples") +} + +async fn analyze_count(pool: &PgPool, schema: &str, table: &str) -> i64 { + sqlx::query_scalar( + "SELECT COALESCE(analyze_count, 0)::bigint FROM pg_stat_user_tables + WHERE schemaname = $1 AND relname = $2", + ) + .bind(schema) + .bind(table) + .fetch_one(pool) + .await + .expect("read analyze_count") +} + +fn probe_client( + pool: PgPool, + statistics_interval: Duration, + storage: Option, +) -> Client { + let mut builder = Client::builder(pool) + .queue( + "statistics_probe", + QueueConfig { + max_workers: 1, + poll_interval: Duration::from_millis(100), + ..QueueConfig::default() + }, + ) + .register_worker(StatisticsProbeWorker) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(100)) + .statistics_refresh_interval(statistics_interval); + if let Some(config) = storage { + builder = builder + .queue_storage(config, Duration::from_secs(600), Duration::from_secs(600)) + .claim_rotate_interval(Duration::from_secs(600)); + } + builder.build().expect("build client") +} + +async fn seed_ready_rows(store: &QueueStorage, pool: &PgPool, rows_per_queue: i64) { + for queue in ["stats_a", "stats_b", "stats_c"] { + store + .enqueue_batch(pool, queue, 2, rows_per_queue) + .await + .expect("enqueue rows"); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn leader_refreshes_statistics_on_election_and_every_interval() { + let database = TestDatabase::create(None, &[]).await; + let pool = &database.owner_pool; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage config"); + seed_ready_rows(&store, pool, 200).await; + + assert_eq!(stats_rows(pool, "awa", "queue_claim_heads", false).await, 0); + assert_eq!( + stats_rows(pool, "awa", "queue_enqueue_heads", false).await, + 0 + ); + assert_eq!(stats_rows(pool, "awa", "ready_entries", true).await, 0); + assert!(parent_reltuples(pool, "awa", "ready_entries").await < 0.0); + + let client = probe_client(pool.clone(), Duration::from_secs(1), None); + client.start().await.expect("start client"); + + wait_for( + "first statistics pass on election", + Duration::from_secs(30), + || async { + stats_rows(pool, "awa", "queue_claim_heads", false).await > 0 + && stats_rows(pool, "awa", "queue_enqueue_heads", false).await > 0 + && stats_rows(pool, "awa", "ready_entries", true).await > 0 + }, + ) + .await; + let partitioned: Vec = sqlx::query_scalar( + "SELECT relname::text FROM pg_class + WHERE relnamespace = 'awa'::regnamespace AND relkind = 'p' ORDER BY relname", + ) + .fetch_all(pool) + .await + .expect("list partitioned tables"); + assert_eq!( + partitioned, RING_PARENTS, + "the refresh covers every ring parent" + ); + wait_for( + "every ring parent analysed", + Duration::from_secs(30), + || async { + for parent in RING_PARENTS { + if analyze_count(pool, "awa", parent).await < 1 { + return false; + } + } + true + }, + ) + .await; + let first_tuples = parent_reltuples(pool, "awa", "ready_entries").await; + assert!( + (550.0..=650.0).contains(&first_tuples), + "parent reltuples after the first pass: {first_tuples}" + ); + + seed_ready_rows(&store, pool, 2_000).await; + wait_for( + "a later pass to see the new rows", + Duration::from_secs(30), + || async { parent_reltuples(pool, "awa", "ready_entries").await > 5_000.0 }, + ) + .await; + assert!(analyze_count(pool, "awa", "ready_entries").await >= 2); + assert!(analyze_count(pool, "awa", "queue_claim_heads").await >= 2); + + client.shutdown(Duration::from_secs(5)).await; + database.drop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn leader_refreshes_statistics_on_election_before_the_first_interval() { + let database = TestDatabase::create(None, &[]).await; + let pool = &database.owner_pool; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage config"); + seed_ready_rows(&store, pool, 200).await; + + let client = probe_client(pool.clone(), Duration::from_secs(300), None); + client.start().await.expect("start client"); + wait_for( + "a statistics pass on election", + Duration::from_secs(30), + || async { stats_rows(pool, "awa", "ready_entries", true).await > 0 }, + ) + .await; + + client.shutdown(Duration::from_secs(5)).await; + database.drop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn zero_interval_disables_the_refresh() { + let database = TestDatabase::create(None, &[]).await; + let pool = &database.owner_pool; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage config"); + seed_ready_rows(&store, pool, 100).await; + + let client = probe_client(pool.clone(), Duration::ZERO, None); + client.start().await.expect("start client"); + wait_for("leadership", Duration::from_secs(10), || async { + client.health_check().await.maintenance_alive + }) + .await; + tokio::time::sleep(Duration::from_secs(3)).await; + + assert_eq!(stats_rows(pool, "awa", "queue_claim_heads", false).await, 0); + assert_eq!(stats_rows(pool, "awa", "ready_entries", true).await, 0); + + client.shutdown(Duration::from_secs(5)).await; + database.drop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn non_owner_runtime_role_refreshes_statistics() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let owner = format!("stats_owner_{suffix}"); + let runtime = format!("stats_runtime_{suffix}"); + let database = TestDatabase::create(Some(&owner), &[&runtime]).await; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage config"); + seed_ready_rows(&store, &database.owner_pool, 200).await; + let runtime_pool = database.runtime_pool(&runtime).await; + + let is_superuser_or_owner: bool = sqlx::query_scalar( + "SELECT rolsuper OR pg_has_role(current_user, 'pg_maintain', 'USAGE') + OR has_table_privilege('awa.queue_claim_heads', 'MAINTAIN') + OR (SELECT relowner = (SELECT oid FROM pg_roles WHERE rolname = current_user) + FROM pg_class WHERE oid = 'awa.queue_claim_heads'::regclass) + FROM pg_roles WHERE rolname = current_user", + ) + .fetch_one(&runtime_pool) + .await + .expect("check runtime role privileges"); + assert!( + !is_superuser_or_owner, + "the runtime role must not be able to analyse the tables itself" + ); + + for (schema, relation, message) in [ + ( + "public", + "queue_claim_heads", + "is not a queue-storage schema", + ), + ("awa", "jobs", "is not a ring parent or head table"), + ] { + let error = + sqlx::query_scalar::<_, String>("SELECT awa.analyze_queue_storage_relation($1, $2)") + .bind(schema) + .bind(relation) + .fetch_one(&runtime_pool) + .await + .expect_err("the helper rejects what it may not analyse"); + let database_error = error.as_database_error().expect("database error"); + assert_eq!(database_error.code().as_deref(), Some("22023")); + assert!( + database_error.message().contains(message), + "{schema}.{relation}: {}", + database_error.message() + ); + } + + let client = probe_client(runtime_pool.clone(), Duration::from_secs(1), None); + client + .start() + .await + .expect("start client as the runtime role"); + wait_for( + "statistics written through a non-owner role", + Duration::from_secs(30), + || async { + stats_rows(&database.owner_pool, "awa", "queue_claim_heads", false).await > 0 + && stats_rows(&database.owner_pool, "awa", "ready_entries", true).await > 0 + }, + ) + .await; + + client.shutdown(Duration::from_secs(5)).await; + runtime_pool.close().await; + database.drop().await; +} + +async fn call_helper( + pool: &PgPool, + schema: &str, + relation: &str, +) -> Result { + sqlx::query_scalar("SELECT awa.analyze_queue_storage_relation($1, $2)") + .bind(schema) + .bind(relation) + .fetch_one(pool) + .await +} + +fn denied(error: &sqlx::Error, message: &str) { + let database_error = error.as_database_error().expect("database error"); + assert_eq!(database_error.code().as_deref(), Some("42501")); + assert!( + database_error.message().contains(message), + "unexpected message: {}", + database_error.message() + ); +} + +/// EXECUTE is not public: holding the lock the helper takes until the end of +/// a transaction would block ring prune of that table. A role needs EXECUTE and +/// write privileges on the head tables, and the migration grants EXECUTE to +/// the roles that already have them. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn execute_is_limited_to_roles_that_may_maintain_queue_storage() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let owner = format!("stats_owner_{suffix}"); + let runtime = format!("stats_runtime_{suffix}"); + let bystander = format!("stats_bystander_{suffix}"); + let reader = format!("stats_reader_{suffix}"); + let database = TestDatabase::create(Some(&owner), &[&runtime, &bystander, &reader]).await; + let owner_pool = &database.owner_pool; + run_sql( + owner_pool, + format!( + "GRANT USAGE ON SCHEMA awa TO {bystander}, {reader}, {runtime}; + GRANT SELECT ON ALL TABLES IN SCHEMA awa TO {reader}; + GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) TO {reader}; + GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA awa TO {runtime}" + ), + ) + .await + .expect("grant privileges"); + let connect = |role: String| { + let options = database.options.clone(); + async move { + PgPoolOptions::new() + .max_connections(2) + .connect_with(options.username(&role).password(ROLE_PASSWORD)) + .await + .expect("connect role pool") + } + }; + let bystander_pool = connect(bystander.clone()).await; + let reader_pool = connect(reader.clone()).await; + let runtime_pool = connect(runtime.clone()).await; + + denied( + &call_helper(&bystander_pool, "awa", "queue_claim_heads") + .await + .expect_err("a role without EXECUTE is denied"), + "permission denied for function", + ); + denied( + &call_helper(&reader_pool, "awa", "queue_claim_heads") + .await + .expect_err("EXECUTE without write privileges on the head tables is denied"), + "may not maintain queue storage", + ); + + // The runtime role holds write privileges but was granted no function + // EXECUTE: re-running the migration grants it, as an upgrade does. + denied( + &call_helper(&runtime_pool, "awa", "queue_claim_heads") + .await + .expect_err("write privileges alone do not grant EXECUTE"), + "permission denied for function", + ); + // Replay the statistics migration itself: later migrations do not repeat + // its EXECUTE grants for a runtime role created after initial migration. + for (_, _, sql) in migrations::migration_sql_range(53, 54) { + run_sql(owner_pool, sql) + .await + .expect("re-run the migration"); + } + assert_eq!( + call_helper(&runtime_pool, "awa", "queue_claim_heads") + .await + .expect("the migration granted EXECUTE to the runtime grantee"), + "analyzed" + ); + + for pool in [bystander_pool, reader_pool, runtime_pool] { + pool.close().await; + } + database.drop().await; +} + +/// The helper analyses the ring parents and head tables only, not any other +/// partitioned table of the schema. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn helper_rejects_a_partitioned_table_that_is_not_a_ring_parent() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let owner = format!("stats_owner_{suffix}"); + let runtime = format!("stats_runtime_{suffix}"); + let database = TestDatabase::create(Some(&owner), &[&runtime]).await; + run_sql( + &database.owner_pool, + "CREATE TABLE awa.application_events (id int) PARTITION BY LIST (id)".to_string(), + ) + .await + .expect("create a partitioned application table"); + let runtime_pool = database.runtime_pool(&runtime).await; + + for relation in ["application_events", "ready_entries_0", "jobs_hot"] { + let error = call_helper(&runtime_pool, "awa", relation) + .await + .expect_err("only ring parents and head tables are accepted"); + let database_error = error.as_database_error().expect("database error"); + assert_eq!( + database_error.code().as_deref(), + Some("22023"), + "{relation}" + ); + assert!( + database_error + .message() + .contains("is not a ring parent or head table"), + "{relation}: {}", + database_error.message() + ); + } + assert_eq!( + call_helper(&runtime_pool, "awa", "ready_entries") + .await + .expect("a ring parent is accepted"), + "analyzed" + ); + + runtime_pool.close().await; + database.drop().await; +} + +/// A conflicting lock on a ring parent fails the call with `lock_not_available` +/// rather than letting `ANALYZE (SKIP_LOCKED)` skip every partition and report +/// success. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn locked_ring_parent_fails_the_call_instead_of_reporting_analyzed() { + let database = TestDatabase::create(None, &[]).await; + let pool = &database.owner_pool; + + let mut blocker = pool.begin().await.expect("begin blocker"); + sqlx::query("LOCK TABLE ONLY awa.ready_entries IN SHARE UPDATE EXCLUSIVE MODE") + .execute(blocker.as_mut()) + .await + .expect("lock the parent"); + + let mut call = pool.begin().await.expect("begin call"); + sqlx::query("SELECT set_config('lock_timeout', '200ms', true)") + .execute(call.as_mut()) + .await + .expect("set lock_timeout"); + let error = sqlx::query_scalar::<_, String>( + "SELECT awa.analyze_queue_storage_relation('awa', 'ready_entries')", + ) + .fetch_one(call.as_mut()) + .await + .expect_err("the locked parent must not report analyzed"); + assert_eq!( + error + .as_database_error() + .and_then(|error| error.code()) + .as_deref(), + Some("55P03") + ); + call.rollback().await.expect("rollback call"); + + blocker.rollback().await.expect("release the parent"); + let status: String = + sqlx::query_scalar("SELECT awa.analyze_queue_storage_relation('awa', 'ready_entries')") + .fetch_one(pool) + .await + .expect("analyse once the lock is gone"); + assert_eq!(status, "analyzed"); + database.drop().await; +} + +fn sum_by_outcome( + resource_metrics: &[opentelemetry_sdk::metrics::data::ResourceMetrics], + outcome: &str, +) -> u64 { + let mut total = 0; + for resource in resource_metrics { + for scope in resource.scope_metrics() { + for metric in scope.metrics() { + if metric.name() != "awa.maintenance.analyze.attempts" { + continue; + } + if let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data() { + total += sum + .data_points() + .filter(|point| { + point.attributes().any(|KeyValue { key, value, .. }| { + key.as_str() == "awa.analyze.outcome" && value.as_str() == outcome + }) + }) + .map(|point| point.value()) + .sum::(); + } + } + } + } + total +} + +/// A refresh that cannot run is recorded and leaves the runtime healthy: the +/// helper's owner may not analyse the tables, or the schema predates the +/// helper. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unrunnable_refresh_is_recorded_and_does_not_fail_the_runtime() { + let exporter = InMemoryMetricExporter::default(); + let provider = SdkMeterProvider::builder() + .with_periodic_exporter(exporter.clone()) + .build(); + opentelemetry::global::set_meter_provider(provider.clone()); + + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let owner = format!("stats_owner_{suffix}"); + let helper = format!("stats_helper_{suffix}"); + let database = TestDatabase::create(Some(&owner), &[&helper]).await; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage config"); + seed_ready_rows(&store, &database.owner_pool, 50).await; + + // The function now runs as a role that owns nothing and holds no MAINTAIN. + sqlx::raw_sql(audited_sql(format!("GRANT {helper} TO {owner}"))) + .execute(&database.admin) + .await + .expect("let the schema owner hand over the function"); + sqlx::raw_sql(audited_sql(format!( + "GRANT USAGE, CREATE ON SCHEMA awa TO {helper}; + ALTER FUNCTION awa.analyze_queue_storage_relation(text, text) OWNER TO {helper}; + GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) TO {owner}" + ))) + .execute(&database.owner_pool) + .await + .expect("hand the helper to an unprivileged role"); + + let client = probe_client( + database.owner_pool.clone(), + Duration::from_millis(200), + None, + ); + client.start().await.expect("start client"); + wait_for("a skipped pass", Duration::from_secs(30), || async { + provider.force_flush().expect("flush metrics"); + let metrics = exporter.get_finished_metrics().expect("read metrics"); + sum_by_outcome(&metrics, "skipped_privilege") > 0 + }) + .await; + assert_eq!( + stats_rows(&database.owner_pool, "awa", "queue_claim_heads", false).await, + 0, + "a skipped pass writes no statistics" + ); + assert!(client.health_check().await.healthy); + client.shutdown(Duration::from_secs(5)).await; + + // A schema that predates the helper skips the same way. + sqlx::raw_sql("DROP FUNCTION awa.analyze_queue_storage_relation(text, text)") + .execute(&database.owner_pool) + .await + .expect("drop the helper"); + let client = probe_client( + database.owner_pool.clone(), + Duration::from_millis(200), + None, + ); + client + .start() + .await + .expect("start client without the helper"); + wait_for("an unavailable pass", Duration::from_secs(30), || async { + provider.force_flush().expect("flush metrics"); + let metrics = exporter.get_finished_metrics().expect("read metrics"); + sum_by_outcome(&metrics, "unavailable") > 0 + }) + .await; + assert!(client.health_check().await.healthy); + client.shutdown(Duration::from_secs(5)).await; + + provider.force_flush().expect("flush metrics"); + let metrics = exporter.get_finished_metrics().expect("read metrics"); + assert_eq!(sum_by_outcome(&metrics, "failed"), 0); + database.drop().await; +} + +/// Back-to-back statistics passes (a far higher duty cycle than the default +/// five-minute interval) run while rotation and prune cycle the ring: prune +/// keeps reclaiming slots and only defers an attempt that lands while ANALYZE +/// is sampling the slot's tables. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn statistics_refresh_does_not_stop_prune_reclaiming_slots() { + let database = TestDatabase::create(None, &[]).await; + let pool = &database.owner_pool; + let schema = "awa_stats_ring"; + let config = QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + claim_slot_count: 2, + lease_claim_receipts: false, + ..Default::default() + }; + let store = QueueStorage::new(config.clone()).expect("queue storage config"); + store.prepare_schema(pool).await.expect("prepare schema"); + + let queue = "statistics_ring"; + let analyze_before = analyze_count(pool, schema, "ready_entries").await; + let client = probe_client(pool.clone(), Duration::from_millis(1), Some(config)); + client.start().await.expect("start client"); + + let (mut pruned, mut deferred) = (0_u32, 0_u32); + let started = Instant::now(); + for _ in 0..60 { + store + .enqueue_batch(pool, queue, 1, 600) + .await + .expect("enqueue"); + loop { + let claims = store.claim_batch(pool, queue, 600).await.expect("claim"); + if claims.is_empty() { + break; + } + store.complete_batch(pool, &claims).await.expect("complete"); + } + match store.rotate(pool).await.expect("rotate") { + RotateOutcome::Rotated { .. } | RotateOutcome::SkippedBusy { .. } => {} + other => panic!("unexpected rotate outcome {other:?}"), + } + match store + .prune_oldest(pool, Duration::ZERO) + .await + .expect("prune never errors under ANALYZE") + { + PruneOutcome::Pruned { .. } => pruned += 1, + PruneOutcome::Blocked { .. } => deferred += 1, + PruneOutcome::Noop | PruneOutcome::SkippedActive { .. } => {} + } + if started.elapsed() > Duration::from_secs(60) { + break; + } + } + let passes = analyze_count(pool, schema, "ready_entries").await - analyze_before; + eprintln!("prune under ANALYZE: pruned={pruned} deferred={deferred} analyze_passes={passes}"); + + client.shutdown(Duration::from_secs(5)).await; + assert!( + passes >= 3, + "the refresh ran {passes} times during the loop" + ); + assert!(pruned >= 20, "prune reclaimed {pruned} slots"); + + // Shutdown can leave one ANALYZE statement finishing server-side; once it + // is gone the oldest sealed slot reclaims without being blocked. + wait_for( + "the last ANALYZE to finish", + Duration::from_secs(30), + || async { + let running: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pg_stat_activity + WHERE state = 'active' AND query LIKE '%analyze_queue_storage_relation%' + AND pid <> pg_backend_pid()", + ) + .fetch_one(pool) + .await + .expect("read pg_stat_activity"); + running == 0 + }, + ) + .await; + store.rotate(pool).await.expect("rotate"); + assert!(!matches!( + store + .prune_oldest(pool, Duration::ZERO) + .await + .expect("prune"), + PruneOutcome::Blocked { .. } + )); + database.drop().await; +} diff --git a/benchmarks/baseline.json b/benchmarks/baseline.json index 9f7fee869..e4ec128e0 100644 --- a/benchmarks/baseline.json +++ b/benchmarks/baseline.json @@ -32,6 +32,13 @@ "rust/stale_heartbeat_rescue": { "handler_per_s": 160.3 }, + "rust/enqueue_lane_16p_unique_185": { + "enqueue_per_s": 10000.0, + "p99_ms": 2500.0 + }, + "rust/enqueue_lane_16p_unique_185_shards4": { + "enqueue_per_s": 10000.0 + }, "rust/mvcc_horizon_overlap": { "min_overlap_vs_baseline": 0.75, "min_cooldown_vs_baseline": 0.75, diff --git a/correctness/storage/MAPPING.md b/correctness/storage/MAPPING.md index d0a6b5236..f7ce8137f 100644 --- a/correctness/storage/MAPPING.md +++ b/correctness/storage/MAPPING.md @@ -37,7 +37,7 @@ The TLA+ lifecycle model does not represent the completed-history rollup cache, | TLA+ action | Rust function | SQL / DDL | | --- | --- | --- | -| `EnqueueReady(j)` | `QueueStorage::insert_ready_rows_tx`, `QueueStorage::insert_ready_rows_copy_tx`, and SQL compatibility `awa.insert_job_compat`; producer entry points include `insert` / `insert_with`, `enqueue_batch`, `enqueue_runtime_rows`, `enqueue_params_batch`, and `enqueue_params_copy` | reserve a lane sequence range through `{schema}.reserve_enqueue_seq` under a transaction-scoped lane advisory lock, sync enqueue-time uniqueness claims, append to `{schema}.ready_entries` via INSERT or COPY, append one `{schema}.ready_segments` row per contiguous committed lane range, and notify logical queues in one tx | +| `EnqueueReady(j)` | `QueueStorage::insert_ready_rows_tx`, `QueueStorage::insert_ready_rows_copy_tx`, and SQL compatibility `awa.insert_job_compat`; producer entry points include `insert` / `insert_with`, `enqueue_batch`, `enqueue_runtime_rows`, `enqueue_params_batch`, and `enqueue_params_copy` | reserve a lane sequence range through `{schema}.reserve_enqueue_seq` under a transaction-scoped lane advisory lock, sync enqueue-time uniqueness claims, append to `{schema}.ready_entries` via INSERT or COPY, append one `{schema}.ready_segments` row per contiguous committed lane range, and notify logical queues in one tx. On v056 schemas, batches below the COPY threshold run the reservation, ring-cursor read, ready and segment appends and (outside a caller transaction) the NOTIFY inside `{schema}.enqueue_ready_rows`, which takes the same lane locks in the same order; a ready-only batch with no uniqueness claims is one autocommit statement | | `EnqueueDeferred(j)` | `QueueStorage::insert_deferred_rows_tx` and `QueueStorage::insert_deferred_rows_copy_tx`; producer entry points include `enqueue_params_batch` and `enqueue_params_copy` | allocate job ids, sync enqueue-time uniqueness claims, append to `{schema}.deferred_jobs` via INSERT or COPY in one tx | | `PromoteDeferred(j)` | maintenance promote loop in `awa-worker/src/maintenance.rs::promote_due_state` | `DELETE FROM deferred_jobs ... INSERT INTO ready_entries ...` in one tx | | `AdvanceClaimCursor` | claim path spent-prefix advancement and post-commit `QueueStorage::advance_claim_cursors` | `claim_ready_runtime()` may advance the sequence-backed claim cursor only across a contiguous prefix of committed spent/tombstoned lanes. In receipts mode, already-emitted attempts are proved from queue-slot-local `{schema}.ready_claim_attempt_batches` range evidence so recovery does not scan claim-ring parent tables while old claim partitions are locked for prune. Newly emitted claims advance the cursor after the claim transaction commits. | diff --git a/docs/architecture.md b/docs/architecture.md index d3c527f61..f0e6e6aef 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -152,7 +152,9 @@ Enqueue is transactional: if the producer's outer transaction rolls back, the jo There are two COPY-shaped producer paths: - Direct queue-storage COPY is the high-throughput path: Rust producers use a configured `QueueStorage::enqueue_params_copy()`, and Python producers use `enqueue_many_copy()`. Each ready or deferred write group uses array-backed INSERT below 2,048 rows and CSV COPY for larger groups. Direct-copy producers must use the same queue-storage configuration as the worker fleet, especially `queue_stripe_count` / `queue_storage_queue_stripe_count`. -- `insert_many_copy()` is the compatibility path. It stages rows with an array-backed INSERT below 2,048 rows and CSV COPY at or above that cutoff, then takes one of two branches. If any staged row has a `unique_key`, every row is replayed through `POSTGRES_INSERT_JOB_SQL` inside a per-row savepoint so SQLSTATE `23505` can skip that duplicate without aborting the batch; that statement itself calls `awa.insert_job_compat()`. An all-non-unique batch uses one lateral `awa.insert_job_compat()` call per staged row without savepoints. Both branches therefore preserve cross-storage behavior. The planned `awa.insert_job()` public capability becomes the clean stable SQL contract after #342 ships its migration and conformance artifacts. +- `insert_many_copy()` is the compatibility path. It stages rows with an array-backed INSERT below 2,048 rows and CSV COPY at or above that cutoff, then inserts them through `awa.insert_jobs_compat_batch` (migration v055; each lane's lock is held only for its reservation and the ready-row write, and unique conflicts are skipped in the same statement). It probes its own connection on every call; on a schema without the function it takes one of two per-row branches. If any staged row has a `unique_key`, every row is replayed through `POSTGRES_INSERT_JOB_SQL` inside a per-row savepoint so SQLSTATE `23505` can skip that duplicate without aborting the batch; that statement itself calls `awa.insert_job_compat()`. An all-non-unique batch uses one lateral `awa.insert_job_compat()` call per staged row without savepoints. Both branches therefore preserve cross-storage behavior. The planned `awa.insert_job()` public capability becomes the clean stable SQL contract after #342 ships its migration and conformance artifacts. + +Generic `insert_many` always uses one statement that calls `awa.insert_job_compat` per row in canonical lane order. Its `PgExecutor` is consumed by that statement, so it cannot probe the connection first; it does not cache another database's schema availability. `insert_many_skip_unique_conflicts` explicitly requires the batch function. Claim is cursor-based rather than heap-scan based: @@ -215,6 +217,8 @@ Queue storage has three independent rings, each advanced by the elected maintena The maintenance tick for each ring is deliberately small: attempt one rotate, then attempt one prune. If a partition is busy, blocked by a lock, current, or still live, the tick records a skipped/blocked outcome and tries again on a future interval. +The leader also refreshes planner statistics on the two head tables and the ring parents when it is elected and every five minutes ([Planner statistics](configuration.md#planner-statistics)): autovacuum never analyses partitioned tables or reaches its threshold on the head tables. It runs beside the ring ticks rather than in them, so a slow `ANALYZE` never delays a rotate or prune tick. + The maintenance tick decides whether to seal the current slot. `ClientBuilder::ring_rotation_cadence` configures `RingRotationCadence`: a non-empty slot becomes eligible after `min_slot_age` since the maintenance store first observed rows (default 5 s) once it holds at least `min_slot_rows` (default 256, sampled with a bounded probe), when total physical storage reaches `target_slot_bytes` (default 2 MiB), or after `max_slot_age` (default 60 s) regardless of rows. At a trickle of writes a slot therefore seals about once a minute instead of every window. The minimum age uses a local monotonic clock: an empty observation clears the timer, and a newly elected leader starts a fresh window. `max_slot_age` also counts from the slot's `rotated_at` in the rotation ledger, on the database clock, so frequent leader changes cannot keep a trickle slot open past it. Under compat (pre-flip) authority an older rotator can leave the current generation without a ledger row; the ceiling then relies on the local observation, which a leader change restarts. The size measurement includes heap, indexes, TOAST, free-space maps and visibility maps. `RingRotationCadence::EAGER` restores eligibility on every tick. Queue pruning queues a bounded exclusive-lock request on routing metadata first, then acquires every other child with `NOWAIT`. Contention rolls back the attempt and retries next tick; active-data skip gates retain their backoff. This lets segment-first claim readers and entry-first producers finish without a prune lock cycle. It does not guarantee progress under perpetual contention. diff --git a/docs/configuration.md b/docs/configuration.md index 930bc6bde..c8107c4fd 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -508,6 +508,7 @@ DLQ rows are different: `dlq_entries` is a separate hold table and the maintenan | `cleanup_batch_size` | `1000` | Max rows deleted per pass. Raise for very high throughput; lower if you want gentler IO. | | `dlq_cleanup_batch_size` | `1000` | DLQ-specific batch size. | | `ClientBuilder::terminal_count_rollup_interval` (Rust) / `terminal_count_rollup_interval_ms` (Python kwarg) | `30s` | Queue storage only: how often maintenance folds pending `done_entries` terminal-count deltas for sealed queue slots into compact live counters. Exact reads include pending deltas and retained compact receipt batches, so this affects compaction/WAL pressure rather than correctness. | +| `ClientBuilder::statistics_refresh_interval` (Rust) / `statistics_refresh_interval_ms` (Python kwarg) | `5m` | Queue storage only: how often the maintenance leader runs `ANALYZE` on the two head tables and the partitioned ring parents, which autovacuum never analyses; also run once when a runtime becomes leader. `0` disables it. See [Planner statistics](#planner-statistics). | ## Queue storage tuning @@ -521,6 +522,8 @@ Direct-copy producers must use the same queue-storage routing config as the work `insert_many_copy()` is the compatibility insert API. It is still useful when a caller needs the canonical insert surface, but in queue-storage mode it routes through the compatibility function rather than being the primary producer fast path. +The Rust compatibility inserts differ in how they treat a unique-key conflict. `insert_many` fails the whole batch with a unique violation (SQLSTATE `23505`). `insert_many_skip_unique_conflicts` inserts the rest of the batch and returns only the rows it inserted, in input order, so a producer that replays work does not have to insert the batch again one job at a time (each of those inserts takes the lane lock separately, which on a busy lane is slower than the batch itself). `insert_many_copy` also skips conflicting jobs, and is what Python's `insert_many_copy()` calls; Python has no separate skip function because that is already its behaviour. `insert_many_skip_unique_conflicts` needs migration v055. `insert_many_copy` probes its own connection on every call and uses `awa.insert_jobs_compat_batch` when installed, otherwise the per-row function. Generic `insert_many` always uses the per-row function in one statement, so a process can write to databases at different migration levels without assuming they share a schema. + ### Terminal count rollup Queue storage keeps exact terminal counts without updating mutable counter rows on the completion hot path. `done_entries` terminal inserts append positive rows to `queue_terminal_count_deltas_*`; retry, discard, DLQ move, and compatibility delete of `done_entries` rows append negative rows. Compact receipt completions do not write count deltas: exact count reads add retained `receipt_completion_batches_*` counts and subtract retained `receipt_completion_tombstones_*`. @@ -531,6 +534,16 @@ The maintenance leader folds pending `done_entries` deltas from sealed queue slo Raising the interval reduces folded-counter churn but leaves more pending delta rows for exact reads to sum. Lowering it folds counters sooner, which can help if exact-count reads are frequent and the `done_entries` delta ledger is growing faster than maintenance drains it. Under long reader transactions, lowering the interval does not force rollup; the useful fix is to release the pinned snapshot or let queue prune reclaim the whole sealed segment. Compact receipt successes are already retained as compact batch rows, so this interval does not change their hot-path write count. +### Planner statistics + +Autovacuum does not analyse partitioned tables, so the ring parents (`ready_entries`, `done_entries`, `leases`, and the other ten partitioned tables of the schema) have no inherited statistics. It also never reaches its analyze threshold on `queue_claim_heads` and `queue_enqueue_heads`, which hold a few rows per lane, so the planner prices joins against them at its defaults. Admin and metrics queries that join retained ready rows to the heads tables then pick plans for a few hundred rows while reading hundreds of thousands. + +The maintenance leader runs `ANALYZE` on both when it is elected and every `statistics_refresh_interval` (default `5m`). One pass takes well under a second on a small ring and a few seconds on a ring with millions of retained rows. It has no measurable effect on claim or completion throughput. + +- **Locks.** Each relation is analysed in its own transaction with a 2 s `lock_timeout` and a 120 s `statement_timeout`. `SKIP_LOCKED` skips a child that ring prune holds during the per-child analysis, but sampling for the parent's statistics waits for it up to the `lock_timeout`; a call that times out is recorded as `skipped_lock` and retried on the next interval. Prune never waits for the refresh: its own 10 ms lock timeout applies, so an attempt that lands while `ANALYZE` samples the oldest slot's tables returns `blocked` and retries on its usual backoff. `ANALYZE` of a partitioned table samples every child, so that window cannot be avoided. +- **Privileges.** The refresh calls `awa.analyze_queue_storage_relation(schema, relation)`, a `SECURITY DEFINER` function installed by migration v054. The runtime needs EXECUTE on the helper, and the login role needs INSERT, UPDATE and DELETE on the target schema's `queue_claim_heads`; the migration grants EXECUTE to existing runtime grantees. A missing helper or caller privilege skips the pass before making calls and logs one warning. Run `awa migrate` for a missing helper; restore the grants as the schema owner for `skipped_privilege`. If the helper owner may not analyse the tables, transfer the function to the table owner or grant its owner `MAINTAIN` on PostgreSQL 17+. See the [database roles guide](security/database-roles.md) for the ownership command. +- **Metrics.** `awa.maintenance.analyze.attempts` and `awa.maintenance.analyze.duration`, by `awa.analyze.scope` (`heads`, `ring_parent`, `pass`) and `awa.analyze.outcome` (`ok`, `skipped_privilege`, `skipped_lock`, `unavailable`, `failed`). A sustained `skipped_lock` rate means something else holds the tables; any `skipped_privilege` or `unavailable` means the refresh is not running. + ### Rust ```rust diff --git a/docs/grafana/README.md b/docs/grafana/README.md index 0ca4ed331..500b13047 100644 --- a/docs/grafana/README.md +++ b/docs/grafana/README.md @@ -207,6 +207,8 @@ In Prometheus they appear under the OpenTelemetry [OTLP-to-Prometheus translatio | `awa.maintenance.rescues` | Counter | rescue_kind | Jobs rescued | | `awa.maintenance.prune.attempts` | Counter | ring, outcome, reason | Ring prune outcomes, including destructive `pruned` and no-DDL `already_pruned` | | `awa.maintenance.prune.duration` | Histogram (s) | ring, phase | Successful destructive prune time split into `lock`, `truncate`, and `commit` | +| `awa.maintenance.analyze.attempts` | Counter | scope (`heads`, `ring_parent`, `pass`), outcome | Planner-statistics `ANALYZE` calls on the head tables and ring parents (`ok`, `skipped_privilege`, `skipped_lock`, `unavailable`, `failed`) | +| `awa.maintenance.analyze.duration` | Histogram (s) | scope, outcome | Duration of one planner-statistics `ANALYZE` call | ### Heartbeat diff --git a/docs/queue-storage-substrate.md b/docs/queue-storage-substrate.md index 4476712e4..9542e3914 100644 --- a/docs/queue-storage-substrate.md +++ b/docs/queue-storage-substrate.md @@ -40,7 +40,7 @@ awa storage prepare-queue-storage-schema \ The command invokes the idempotent `awa.install_queue_storage_substrate(...)` helper under a per-schema advisory transaction lock. The helper is activation-neutral: preparing a schema does not route work to it. Activation is a separate staged transition described in [Upgrading from 0.5 to 0.6](upgrade-0.5-to-0.6.md). -The installer is `SECURITY INVOKER`; its caller needs DDL privileges on the target schema. After the migrator provisions each queue with `awa storage prepare-queue`, workers need runtime DML, sequence `USAGE`/`SELECT`/`UPDATE`, and `TRUNCATE` for guarded ring reclamation. They do not need schema CREATE. See [Database roles](security/database-roles.md). +The installer is `SECURITY INVOKER`; its caller needs DDL privileges on the target schema. Workers need runtime DML, sequence `USAGE`/`SELECT`/`UPDATE`, and `TRUNCATE` for guarded ring reclamation. They do not need schema CREATE: lanes are created through the schema owner's provisioning functions. See [Database roles](security/database-roles.md). ## Operator rules diff --git a/docs/security/database-roles.md b/docs/security/database-roles.md index b6c1557e9..9d5abdca8 100644 --- a/docs/security/database-roles.md +++ b/docs/security/database-roles.md @@ -63,36 +63,68 @@ ALTER DEFAULT PRIVILEGES FOR ROLE awa_owner IN SCHEMA awa If an existing installation created objects as `awa_migrator` without first setting the owner role, transfer every existing schema object to `awa_owner` before relying on owner-scoped defaults. New migrations should use the effective-role command above. -## Provision queues before starting restricted runtimes +## Planner statistics -From schema v049, runtime calls use existing lane sequences without issuing -`CREATE SEQUENCE`. Run provisioning as the migrator for every queue before -starting its producers or workers: +The maintenance leader runs `ANALYZE` on the queue-storage head tables and partitioned ring parents (see [Planner statistics](../configuration.md#planner-statistics)). `ANALYZE` skips a table its caller may not analyse, so a runtime that does not own the tables cannot do this itself; the leader calls `awa.analyze_queue_storage_relation(schema, relation)` instead, a `SECURITY DEFINER` function from migration v054 owned by whichever role ran the migration. The runtime does not need `MAINTAIN`. The function is not executable by `PUBLIC`, because the lock it takes lasts until the caller's transaction ends and would block ring prune: the migration grants `EXECUTE` to the roles that already hold `INSERT`, `UPDATE` and `DELETE` on the head tables, and the function requires the calling session user to hold them. A runtime role created later gets it from the `GRANT EXECUTE ON ALL FUNCTIONS` above or the default privileges. For a custom queue-storage schema, the role needs write privileges on that schema's `queue_claim_heads`. + +Run migrations with the schema owner as the effective role, as shown above, so the function's owner owns the tables. A function owned by a role that may not analyse them (it is not the table owner, the database owner, or a `MAINTAIN` holder) makes every pass report `skipped_privilege` and logs one warning. The function only analyses the two head tables and the thirteen ring parents of a schema that `awa.queue_storage_schemas()` reports as trusted, and `ANALYZE` changes statistics only. + +If the helper's owner may not analyse the tables, run `ALTER FUNCTION awa.analyze_queue_storage_relation(text, text) OWNER TO awa_owner` as a role permitted to transfer ownership, or grant its owner `MAINTAIN` on the target tables on PostgreSQL 17+. Replacing the function body leaves its owner unchanged. + +## Lane sequences + + +Each queue lane, `(queue, priority, enqueue_shard)`, has an enqueue and a claim sequence. From schema v053 they are created by the schema owner, whichever role first enqueues into, claims from or promotes into the lane: the lane helpers call `provision_lane` and `provision_lane_sequences` in the queue-storage schema. These are `SECURITY DEFINER` functions owned by the schema owner with a pinned `search_path`. Like the other runtime-callable functions they are executable by `PUBLIC`, so schema `USAGE` is what gates them, and they only create lanes, after checking that the queue, priority and shard are set and the shard is 0–63. A runtime therefore needs neither schema `CREATE` nor default privileges to start using a new queue, priority or shard. + +A new lane sequence gets the privileges of the head table it backs. A role that may `INSERT` or `UPDATE` `queue_enqueue_heads` gets `USAGE, SELECT, UPDATE` on the enqueue sequences, and one that may only `SELECT` it gets `SELECT`; `queue_claim_heads` does the same for the claim sequences. Granting the runtime the table privileges above is enough, and no sequence default privileges are needed for lanes. + +Run migrations with the schema owner as the effective role, as above, so the provisioning functions belong to it. A migrator that cannot `SET ROLE` to the owner leaves them owned by itself and the migration warns. + +`awa storage prepare-queue` still creates every priority and shard of a queue in one transaction, and any runtime role may run it. Use it to create lanes before a deploy rather than on first use. + +### Check and repair lane ownership + +Lanes created before v053 belong to whichever role touched them first, and only that role, or the roles its default privileges named, can use them. A different runtime role then fails with `permission denied for sequence` on that lane. List lane sequences that are not owned by the schema owner or lack a runtime grant (read-only, any role): ```bash -PGOPTIONS='-c role=awa_owner' \ - awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage prepare-queue \ - --queue email --enqueue-shards 1 +awa --database-url "$AWA_DATABASE_URL" storage lane-sequences --schema awa +``` + +```sql +SELECT * FROM awa.lane_sequence_problems('awa'); ``` -The command provisions all four priorities in one transaction and is safe to -repeat on a live queue: existing sequence cursors keep their positions. It does -not change routing or the queue's configured shard count. Prepare additional -shards before increasing `enqueue_shards`; match `--queue-stripe-count` to the -runtime's queue-storage configuration. With `PartitionedQueue`, provision each -physical queue. Use `--schema` for a custom storage schema. - -Apply the sequence grants above after provisioning, or set the owner's default -privileges first. Lane cursors need `UPDATE` for `setval`, in addition to `USAGE` -and `SELECT`. `CREATE SEQUENCE IF NOT EXISTS` requires schema CREATE even when -the sequence already exists, so provisioning alone does not repair pre-v049 -helpers: apply the migration first. Schema owners retain lazy creation for -single-role installations; a restricted runtime encountering an unprepared -lane fails with a provisioning hint. +The command exits 2 while any remain. +`awa migrate` repairs them once the migration commits, one sequence per short transaction: it gives each sequence to the schema owner and adds its missing grants, wherever it is allowed to. Reassigning takes the sequence's `ACCESS EXCLUSIVE` lock, so the repair waits at most 300 ms for it. A sequence that an open enqueue or claim holds for longer is skipped, left unchanged and reported as busy. A sequence the repair is not allowed to change is reported with the roles that can finish it. Neither fails the migration. Run the repair again at any time; it is idempotent: + +```bash +awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage lane-sequences --repair +``` + +Which role can repair a sequence: + +- Reassigning it needs a role that holds the current owner's privileges (a member of the owner, with inheritance) and can `SET ROLE` to the schema owner. +- Adding its grants needs the current owner's privileges. + +The schema owner holds neither the migration role's nor an application role's privileges, so a repair running as `awa_owner` cannot reassign sequences those roles created. That includes the migration role connected with `PGOPTIONS='-c role=awa_owner'`. Run the repair as the migration role itself, without `PGOPTIONS`, as above. `awa migrate` and `--repair` retry each sequence the effective role could not repair as the login role, so `awa migrate` run with `PGOPTIONS='-c role=awa_owner'`, as documented, still repairs the migration role's sequences. + +For each sequence that remains, the output lists `repair_roles`, the login roles other than superusers that can repair it, and the statements to run. A busy sequence only needs another `--repair`. When `repair_roles` is empty, which is usual for a sequence an application role created on managed PostgreSQL, either: + +- run `missing_grants` as the current owner. Every runtime role can then use the lane, although the owner stays unchanged; or +- give the migration role the owner's privileges, for example `GRANT app_role TO awa_migrator` from a role with `ADMIN` on `app_role`, rerun `--repair` as the migration role, then revoke the membership. + +Without the CLI, repair one sequence per transaction from `psql` in autocommit mode, connected as the migration role: + +```sql +SELECT format('SELECT * FROM awa.repair_lane_sequence(%L, %L)', 'awa', sequence_name) +FROM awa.lane_sequence_problems('awa') \gexec +``` + +## Custom queue-storage schemas ## Custom queue-storage schemas -Repeat the schema, table, sequence, function, and default-privilege grants for a custom queue-storage schema. Prepare custom schemas as the migrator; the runtime needs DML, sequence access, function execution, and `TRUNCATE`, but never DDL. See [Queue storage](../queue-storage-substrate.md). +Repeat the schema, table, sequence, function, and default-privilege grants for a custom queue-storage schema. Prepare custom schemas as the migrator; the runtime needs DML, sequence access, function execution, and `TRUNCATE`, but never DDL. Check lanes in a custom schema with `awa storage lane-sequences --schema `. See [Queue storage](../queue-storage-substrate.md). ## Why the grants are broad diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 1672c5865..0f85afbcf 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -186,7 +186,7 @@ Reading them together separates the two usual failure modes: **1. Producer is going through the compatibility insert path.** -`insert_many_copy_from_pool` / `client.insert_many_copy` routes every row through `awa.insert_job_compat()` once per row. On any database with non-trivial per-statement latency (auth-proxy hop, managed Postgres, network) the per-row function call dominates — measured at ~100–150 ms per row through a Cloud SQL Auth Proxy in staging, which caps a single producer at ~7 rows/sec regardless of batch or chunk size. +Before migration v055, `insert_many_copy_from_pool` / `client.insert_many_copy` calls `awa.insert_job_compat()` per staged row. With unique jobs, per-row savepoints add database round trips that can dominate behind a proxy or on managed PostgreSQL. After migration, COPY probes its own connection and uses the batch function. Generic `insert_many` retains its per-row statement on every supported schema. Switch the producer to the direct queue-storage COPY entry point: @@ -434,12 +434,32 @@ Fix: awa --database-url "$DATABASE_URL" migrate ``` +### Producers wait on one queue's enqueue lane + +Symptoms: + +- many connections wait with `wait_event = 'advisory'` while the database is otherwise idle +- `insert_many` takes seconds although the work per batch is milliseconds, and single-job inserts into the same queue take much longer than usual +- scheduled-job promotion and rescue of the same queue run slowly while producers saturate it + +Cause: + +- on queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane until it commits, so concurrent producers for one lane run one at a time; a lane's ceiling is a batch's rows divided by the time the lock is held +- from migration v055, `insert_many_copy` and `insert_many_skip_unique_conflicts` prepare the batch before taking lane locks; generic `insert_many` continues to hold the lane while subsequent rows are processed +- maintenance promotion moves due scheduled and retryable jobs into the lane under the same lock; a promotion batch is one transaction across every queue it touches, and the leader runs its branches one at a time, so a batch parked behind a hot lane delays the due jobs of other queues in that batch and the leader's other branches. Rescue itself does not take the lane lock; it writes retryable rows that promotion moves later + +Fix: + +- run `awa migrate` and use `insert_many_copy` or `insert_many_skip_unique_conflicts` for the batch function; generic `insert_many` stays on its compatible per-row path +- when `insert_many` fails on a unique conflict, retry with `insert_many_skip_unique_conflicts` rather than inserting the batch again row by row: each single-row insert queues for the lane separately +- raise `awa.queue_meta.enqueue_shards` for a queue whose producers still saturate one lane, accepting partitioned rather than strict FIFO ([ADR-025](adr/025-sharded-enqueue-heads.md)) + ### `deadlock detected` while enqueueing Cause: - on queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane until the transaction commits; the `DETAIL` names two 64-bit advisory keys, and two transactions that took lanes in opposite orders hold one each and wait on the other's -- `insert_many`, `insert_many_copy` and the `QueueStorage` batch enqueues take their lane locks in one canonical order, `(queue, priority, shard)` ascending with queue names compared bytewise, so batches never deadlock with each other +- `insert_many`, `insert_many_copy` and the `QueueStorage` batch enqueues take their lane locks in one canonical order, `(queue, priority, shard)` ascending with queue names compared bytewise, avoiding deadlocks caused by opposing lane-lock order; uniqueness and other application locks can still form cycles - a transaction that enqueues with several separate statements (`insert` per job, handler-staged follow-ups, `Transaction.insert` in a loop) takes locks in statement order and can deadlock against a batch or another such transaction that touches the same lanes in a different order Fix: diff --git a/docs/upgrade-0.6-to-0.7.md b/docs/upgrade-0.6-to-0.7.md index 0afc48afb..0f5ba3980 100644 --- a/docs/upgrade-0.6-to-0.7.md +++ b/docs/upgrade-0.6-to-0.7.md @@ -14,14 +14,74 @@ A [migrate-first rehearsal with released `awa-pg==0.6.7`](https://github.com/har verified a restricted worker completing jobs on the upgraded schema. That rehearsal does not cover concurrent mixed-version rollouts. -For a runtime without DDL privileges, apply v049 as the migrator, run -`awa storage prepare-queue --queue ` for every physical queue, and grant -sequence `USAGE, SELECT, UPDATE` before starting producers or workers. Match -shard and stripe counts and provision extra lanes before increasing them. External migration runners must apply the complete generated `awa migrate --sql --from 48 --to 49` bundle: it refreshes the installer -definition before replacing lane helpers in the installed schemas. -See [Database roles](security/database-roles.md#provision-queues-before-starting-restricted-runtimes). +definition before replacing lane helpers in the installed schemas. From v053, +runtimes without DDL privileges create lanes themselves through the schema +owner's provisioning functions, so `awa storage prepare-queue` is optional. + +## Lane ownership and v053 + +Before v053 a lane's sequences belonged to the role that first enqueued into, +claimed from or promoted into it. With several runtime roles, only roles whose +default privileges covered that role could use the lane; any other role failed +with `permission denied for sequence`. v053 installs `SECURITY DEFINER` +`provision_lane` and `provision_lane_sequences` functions, owned by the schema +owner, in each queue-storage schema. The lane helpers create sequences through +them, with grants that mirror the lane's head table. Once v053 commits, +`awa migrate` gives every existing lane sequence to the schema owner and adds +its missing grants, as far as it is allowed to, one sequence per short +transaction with a bounded lock wait. + +0.6 and earlier 0.7 binaries call the same helpers, so lanes they create after +v053 are owned by the schema owner too. A 0.7 binary running ahead of v053 +detects that `provision_lane` is missing and creates lanes as before. + +To adopt it: + +1. Check the current state: `awa storage lane-sequences` (read-only, exits 2 + while lanes need repair). +2. Run `awa migrate` with the schema owner as the effective role, as described + in [Database roles](security/database-roles.md). +3. Read the migrate output. Any lane sequence it could not repair is listed in + a warning. A busy one needs another run of + `awa storage lane-sequences --repair` as the migration role (without + `PGOPTIONS`); the others name the roles that can repair them. See + [Check and repair lane ownership](security/database-roles.md#check-and-repair-lane-ownership). + +External migration runners apply the generated `awa migrate --sql --from 52 +--to 53` bundle, which does not repair existing lane sequences, then run +`awa storage lane-sequences --repair` as the migration role itself. + +## v056: enqueue round trips + +v056 replaces functions only. `reserve_enqueue_seq` keeps its signature, lane +lock and return value, so earlier 0.7 binaries and SQL producers +(`insert_job_compat`) run unchanged on the migrated schema. A new +`{schema}.enqueue_ready_rows` lets current runtimes enqueue a ready batch in +one statement. Runtimes use it when they have `EXECUTE` on it and on +`{schema}.ring_cursor`, which it calls and which is revoked from `PUBLIC`; the +default privileges in [Database roles](security/database-roles.md#runtime-grants) +grant both. Otherwise they keep the multi-statement path. A store that finds +the function missing probes again after a minute, so it picks the function up +without a restart once the migration commits. If the function disappears or +loses `EXECUTE` after a store has used it, the same call retries once through +the multi-statement path. Any transaction that call owns is rolled back before +replaying the original input, including its uniqueness claims. Other errors +propagate; a healthy cached fast path adds no probe or savepoint. + +Deploy order: the fallback keeps enqueue, claim and completion working when a +current binary runs against a pre-v056 schema, but worker readiness does not +follow it. `/readyz` reports a worker ready only once the schema is at the +binary's version, so health-gated workers started before the migration stay +unready until it commits. Migrate first for health-gated workers; binary-first +is safe for producers and for workers without a readiness gate. + +Apply the complete pending range through v056, including v053 lane ownership, +v054 statistics and v055 batch inserts. The runner applies only versions above +the highest one recorded; do not stamp v056 before its predecessors. +External runners must apply the complete generated `awa migrate --sql` bundle, +which redefines the installer before refreshing the installed schemas. ## Custom-schema discovery during migration @@ -253,6 +313,39 @@ workers report not ready until the migration runs, as for any version bump. External migration runners apply the complete `awa migrate --sql --from 51 --to 52` bundle, which defines the helper before the tail calls it. +## v054: planner statistics helper + +v054 adds `awa.analyze_queue_storage_relation(schema, relation)`, a +`SECURITY DEFINER` function owned by the migration role that lets the +maintenance leader refresh planner statistics for the head tables and ring +parents when the runtime role does not own them (see +[Planner statistics](configuration.md#planner-statistics)). It creates one +function, revokes `EXECUTE` from `PUBLIC` and grants it to the roles that +already hold INSERT, UPDATE and DELETE on the `awa` head tables (the runtime +grantees); it takes no table lock and completes in milliseconds. Roles created +later need `GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa`. Run the migration as the schema owner so the +function's owner owns the tables. + +Binaries before v054 never call it. A 0.7 binary running before the migration +logs once that the function is missing and skips the refresh; the +`awa.maintenance.analyze.attempts` counter reports `unavailable`. External +migration runners apply the `awa migrate --sql --from 53 --to 54` bundle. + +## v055: batch producer insert + +v055 adds `awa.insert_jobs_compat_batch`, the function `insert_many_copy` and +`insert_many_skip_unique_conflicts` call. It is one `CREATE OR REPLACE FUNCTION`: no table DDL, +no scan, no lock beyond the catalog, so it applies instantly under live load. +`awa.insert_job_compat` is unchanged, and binaries before v055 keep calling it; +both paths take the same per-lane reservation lock, so producers on either +version serialise on a lane as before. + +The runtime role needs `EXECUTE` on the new function; the +[default privileges](security/database-roles.md#runtime-grants) cover functions +the migrator creates. + +A 0.7 binary running before the migration keeps inserting. `insert_many_copy` probes its own connection on every call, uses the per-row function until the batch function exists, and switches without a restart. Generic `insert_many` always uses the per-row function in one statement, so one process can write to databases at different migration levels without aborting a caller transaction on an unavailable function. `insert_many_skip_unique_conflicts` needs the function and returns a validation error before the migration. The per-row and batch paths take the same lane lock, so a fleet that mixes them keeps each batch contiguous and in order. External migration runners apply the `awa migrate --sql --from 54 --to 55` bundle. + ## The one gate `awa migrate` on a 0.7 binary refuses to apply migrations unless one of these holds: diff --git a/skills/awa-jobs/SKILL.md b/skills/awa-jobs/SKILL.md index 3895bf63a..8cd47e6b0 100644 --- a/skills/awa-jobs/SKILL.md +++ b/skills/awa-jobs/SKILL.md @@ -124,10 +124,7 @@ low-priority work cannot starve indefinitely. **Unique jobs** deduplicate on a BLAKE3 key over kind + optionally queue + args + a time-period bucket. Defaults: dedup by args, not by queue. Conflict handling -depends on the API: a single-row `insert` / `insert_with` and the row-by-row -`insert_many` **raise `AwaError::UniqueConflict`** (Python raises the -corresponding error) on a duplicate — you must handle it, do not assume a silent -no-op. Only the COPY bulk path (`insert_many_copy`) silently skips duplicates. +depends on the API: a single-row `insert` / `insert_with` raises `AwaError::UniqueConflict` (Python raises the corresponding error), and `insert_many` fails the whole batch with SQLSTATE `23505`. `insert_many_copy` skips duplicates on every supported schema; Rust's `insert_many_skip_unique_conflicts` also skips them but requires migration v055. Both return only inserted rows, in input order. Generic `insert_many` always uses one per-row SQL statement; COPY probes its own connection and uses the batch function when installed. Cancel without storing the id via `cancel_by_unique_key(kind, queue=, args=, period_bucket=)` — the components must match those used at insert. It cancels the oldest live job holding the key. On queue storage, a `by_queue` key is diff --git a/skills/awa-operations/SKILL.md b/skills/awa-operations/SKILL.md index 4cfb039b0..d8785bd6a 100644 --- a/skills/awa-operations/SKILL.md +++ b/skills/awa-operations/SKILL.md @@ -91,11 +91,19 @@ Do not skip a major version. 0.5 → 0.7 is unsupported; step through 0.6. ## Restricted Database Roles -When deploying a runtime without schema CREATE, apply v049 and provision every -queue with `awa storage prepare-queue` as the migrator before starting producers -or workers. Match shard/stripe counts; provision additional lanes before changing -routing. Use the sequence grants and owner defaults in -[Database roles](../../docs/security/database-roles.md#provision-queues-before-starting-restricted-runtimes). +From v053 any runtime role creates lanes through the schema owner's +`SECURITY DEFINER` provisioning functions, so runtimes need neither schema +CREATE nor sequence default privileges, and `awa storage prepare-queue` is +optional. Run migrations with the schema owner as the effective role. Check +lane ownership with `awa storage lane-sequences` (exits 2 while lanes need +repair) and fix it with `--repair`, connected as the migration role itself +rather than as the schema owner, or with the statements it prints; see +[Database roles](../../docs/security/database-roles.md#check-and-repair-lane-ownership). + +Queue-storage batch enqueue acceleration (v056) keeps the older multi-statement +path available. A missing function or revoked EXECUTE retries once within the +same call after rolling back its owned transaction. Apply complete migration +bundles through v056, retaining the v053/v054/v055 predecessors. ## Rolling Upgrades @@ -135,7 +143,8 @@ processes. Exactly one worker per database is the elected maintenance leader (session-scoped advisory lock — if it dies, another is elected on the next tick). The leader alone runs cluster-wide scans: deferred→ready promotion, the three rescue paths (stale-heartbeat, hard-deadline, callback-timeout), ring -rotation and prune, DLQ and descriptor cleanup, and cron evaluation. During an +rotation and prune, planner-statistics refresh, DLQ and descriptor cleanup, and +cron evaluation. During an unfinalized transition it promotes both planes' deferred backlogs, so drain and post-flip execution do not depend on which runtime holds leadership. Every worker dispatches and heartbeats its own attempts. There is no `pg_cron` or @@ -313,6 +322,19 @@ storage limit. Queue prune lock contention retries on the next decision tick; active-data gates retain exponential backoff. Persistent contention can still stall reclamation. See `docs/configuration.md` for the scheduling contract. +On queue storage the leader runs `ANALYZE` on `queue_claim_heads`, +`queue_enqueue_heads` and the partitioned ring parents when elected and every +`statistics_refresh_interval` (default 5 minutes; Python +`statistics_refresh_interval_ms`; zero disables), because autovacuum never +analyses them. It goes through the `SECURITY DEFINER` function +`awa.analyze_queue_storage_relation` from migration v054, so a non-owner runtime +role needs no `MAINTAIN` grant. Watch `awa.maintenance.analyze.attempts`: any +`unavailable` means the helper is missing: run `awa migrate` as the schema +owner. For `skipped_privilege`, restore helper EXECUTE and the login role's +head-table DML, or correct the helper owner; migration reruns preserve an +existing function's owner. See `docs/configuration.md#planner-statistics` and +`docs/security/database-roles.md`. + ## Managed Postgres Prefer PostgreSQL 18 where available; PG17 on some managed platforms caps