diff --git a/CHANGELOG.md b/CHANGELOG.md index a1b0a6839..6ab4ce871 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,75 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- **On queue storage, a queue whose enqueue lane is busy no longer holds + back promotion of other queues, nor the maintenance leader's rescue, + rotation and prune.** Promotion re-enqueues due scheduled and retryable + rows into their lane, which takes the lane's reservation lock that + producers hold until they commit. A promotion batch was one transaction + across every queue it touched, and the leader runs its branches one after + another, so a lane kept busy by producers delayed due jobs of unrelated + queues in the same batch, and rescue started seconds late. Promotion now + gives each `(queue, priority)` lane its own transaction under a 50 ms + `lock_timeout`. A lane that times out keeps its rows deferred and is + handed to a dedicated task that waits up to 30 s for the lock, queued + behind the producers already waiting, so the lane is still promoted in + turn however busy it stays; the leader skips it meanwhile. Each waiter + holds a pool connection, so the leader runs one per eight pool + connections, at least one and at most four; a further contended lane is + retried on the next tick. Batch selection skips rows another transaction + has locked, so locked due rows cannot fill every batch of a tick. `awa.maintenance.promote.lane_contention` + (`awa_maintenance_promote_lane_contention_total{awa_job_state,awa_job_queue,awa_promote_outcome}`) + counts `handed_off`, `waiter_limit` and `waiter_timed_out`. Promoted rows + now take lane sequence numbers in due order within their lane. + `QueueStorage::promote_due` keeps its signature and promotes lane by lane + without a lock timeout; `promote_due_lanes` and `promote_due_lane` expose + the per-lane form. +- **A queue lane is usable by every runtime role, whichever role creates it + first.** 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. The new schema patch + `definer_lane_provisioning` creates lanes 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. Once the patch 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 run. 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 before the patch detects that `provision_lane` is missing and + creates lanes as before, and older binaries create correctly owned lanes once + the patch has refreshed the helpers they call. `awa migrate` applies the + patch on any v040 database; external runners apply the exported + `R__definer_lane_provisioning.sql` after `V40`. See + [`docs/upgrade-0.6.13-to-0.6.14.md`](docs/upgrade-0.6.13-to-0.6.14.md). +- **A graceful shutdown no longer leaves a finished attempt for stale-heartbeat + rescue to run again.** Once a handler returned, its finalize (completion, + retry, cancel or failure, and any ADR-029 follow-up inserts) ran in a task + that `Client::shutdown` did not wait for. When the last handler returned + just as the drain ended, shutdown stopped the completion batcher before + that finalize reached it, and the direct fallback raced the process exit. + The attempt had already left the heartbeat set, so it went stale and was + rescued and executed again `heartbeat_staleness` later; a deploy could + re-run a few jobs on every rolling restart. Shutdown now waits for + pending finalization, within the same drain timeout, before it stops the + completion batcher, heartbeat and maintenance services. A finalize that + outlives the timeout is still left to rescue, as before. +- **Point reads show the live queue and priority after a lane move.** Retained + ready rows from before a priority or queue change could win a read's tie + after the ring wrapped, returning the old fields even though the batch + operation succeeded. Point reads now exclude tombstoned source rows. + - **Callback completion, `get_job`, lifecycle events and the admin reads find running jobs that priority aging promoted on queue storage** (backport of #600). Claims and leases record a job's aged effective priority while its diff --git a/Cargo.lock b/Cargo.lock index 30bd194c9..4386517ed 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -400,6 +400,7 @@ dependencies = [ "serde_json", "sqlx", "tokio", + "tokio-util", "tracing", "tracing-opentelemetry", "tracing-subscriber", @@ -4187,6 +4188,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] diff --git a/awa-cli/src/main.rs b/awa-cli/src/main.rs index f20634116..1e347be53 100644 --- a/awa-cli/src/main.rs +++ b/awa-cli/src/main.rs @@ -417,7 +417,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")] @@ -432,6 +432,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 @@ -669,6 +680,26 @@ async fn main() -> Result<(), Box> { if !report.complete() { std::process::exit(1); } + if report.patches.iter().any(|patch| { + patch.name == "definer_lane_provisioning" + && patch.outcome == awa_model::migrations::PatchOutcome::AlreadyPresent + }) { + // The patch was already in place, so this run repaired + // nothing: report what an earlier repair left behind. + 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}"); + } + } } } @@ -1245,6 +1276,39 @@ async fn main() -> 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?; @@ -1420,6 +1484,18 @@ fn print_migration_report(report: &awa_model::migrations::MigrationReport) { ); } } + if !report.lane_sequences.is_empty() { + let repaired = report + .lane_sequences + .iter() + .filter(|repair| repair.outcome == awa_model::storage::LaneSequenceOutcome::Repaired) + .count(); + // Sequences left behind are logged by the migration runner. + println!( + "Lane sequences repaired: {repaired} of {}", + report.lane_sequences.len() + ); + } match ( &report.storage_routing_before, &report.storage_routing_after, diff --git a/awa-cli/tests/migrate_patch_cli_test.rs b/awa-cli/tests/migrate_patch_cli_test.rs index 199ebeff4..abf84f477 100644 --- a/awa-cli/tests/migrate_patch_cli_test.rs +++ b/awa-cli/tests/migrate_patch_cli_test.rs @@ -321,12 +321,11 @@ fn psql(database_url: &str, sql: &str, lock_timeout: Option<&str>) -> Output { let mut child = command .spawn() .expect("psql must be on PATH for the exported-SQL test"); - child - .stdin - .take() - .unwrap() - .write_all(sql.as_bytes()) - .unwrap(); + // psql stops reading once ON_ERROR_STOP fires, which can be before the + // whole script is written. + if let Err(err) = child.stdin.take().unwrap().write_all(sql.as_bytes()) { + assert_eq!(err.kind(), std::io::ErrorKind::BrokenPipe, "{err}"); + } child.wait_with_output().unwrap() } diff --git a/awa-metrics/src/lib.rs b/awa-metrics/src/lib.rs index 3f81abce3..a3d015b26 100644 --- a/awa-metrics/src/lib.rs +++ b/awa-metrics/src/lib.rs @@ -64,6 +64,7 @@ pub mod names { pub const MAINTENANCE_PROMOTE_BATCHES: &str = "awa.maintenance.promote_batches"; pub const MAINTENANCE_PROMOTE_BATCH_SIZE: &str = "awa.maintenance.promote_batch_size"; pub const MAINTENANCE_PROMOTE_DURATION: &str = "awa.maintenance.promote_duration"; + pub const MAINTENANCE_PROMOTE_LANE_CONTENTION: &str = "awa.maintenance.promote.lane_contention"; pub const MAINTENANCE_BRANCH_DURATION: &str = "awa.maintenance.branch.duration"; pub const MAINTENANCE_BRANCH_OVERRUN: &str = "awa.maintenance.branch.overrun"; pub const MAINTENANCE_ROTATE_ATTEMPTS: &str = "awa.maintenance.rotate.attempts"; @@ -162,6 +163,11 @@ pub struct AwaMetrics { pub promotion_batch_size: Histogram, /// Promotion query duration. pub promotion_duration_seconds: Histogram, + /// Promotion lanes whose locks were busy past the leader's wait budget, + /// attributed to (state, queue, outcome): `handed_off` to a dedicated + /// waiter, `waiter_limit` when every waiter was busy, `waiter_timed_out` + /// when a waiter's own budget ran out. + pub promotion_lane_contention: Counter, /// Current in-flight jobs (can go up and down). pub jobs_in_flight: UpDownCounter, /// Total heartbeat batches sent. @@ -404,6 +410,11 @@ impl AwaMetrics { .with_unit("s") .with_boundaries(LATENCY_BUCKETS_SECONDS.to_vec()) .build(), + promotion_lane_contention: meter + .u64_counter(names::MAINTENANCE_PROMOTE_LANE_CONTENTION) + .with_description("Promotion lanes whose locks were busy, by state/queue/outcome") + .with_unit("{lane}") + .build(), jobs_in_flight: meter .i64_up_down_counter(names::JOB_IN_FLIGHT) .with_description("Current number of in-flight jobs") @@ -749,6 +760,23 @@ impl AwaMetrics { .record(duration.as_secs_f64(), &attrs); } + /// Record a promotion lane whose locks were busy past the leader's wait + /// budget, with what became of it (`handed_off`, `waiter_limit`, + /// `waiter_timed_out`). + pub fn record_promotion_lane_contention( + &self, + state: &str, + queue: &str, + outcome: &'static str, + ) { + let attrs = [ + opentelemetry::KeyValue::new("awa.job.state", state.to_string()), + opentelemetry::KeyValue::new("awa.job.queue", queue.to_string()), + opentelemetry::KeyValue::new("awa.promote.outcome", outcome), + ]; + self.promotion_lane_contention.add(1, &attrs); + } + /// Record in-flight change. pub fn record_in_flight_change(&self, queue: &str, delta: i64) { let attrs = [opentelemetry::KeyValue::new( @@ -1162,6 +1190,65 @@ mod tests { metrics.record_job_kind_info("minimal", None, None, None, None, &[]); } + #[test] + fn promotion_lane_contention_counts_each_outcome_per_state_and_queue() { + let exporter = InMemoryMetricExporter::default(); + let provider = SdkMeterProvider::builder() + .with_periodic_exporter(exporter.clone()) + .build(); + let metrics = AwaMetrics::new(&provider.meter("awa-metrics-test")); + + metrics.record_promotion_lane_contention("scheduled", "emails", "handed_off"); + metrics.record_promotion_lane_contention("scheduled", "emails", "handed_off"); + metrics.record_promotion_lane_contention("retryable", "emails", "waiter_limit"); + metrics.record_promotion_lane_contention("scheduled", "billing", "waiter_timed_out"); + provider.force_flush().expect("flush metrics"); + let resource_metrics = exporter.get_finished_metrics().expect("read metrics"); + + let mut counts = BTreeMap::new(); + for resource in &resource_metrics { + for scope in resource.scope_metrics() { + for metric in scope.metrics() { + if let ( + names::MAINTENANCE_PROMOTE_LANE_CONTENTION, + AggregatedMetrics::U64(MetricData::Sum(sum)), + ) = (metric.name(), metric.data()) + { + for point in sum.data_points() { + let label = |key: &str| { + point + .attributes() + .find(|kv| kv.key.as_str() == key) + .map(|kv| kv.value.as_str().into_owned()) + .unwrap_or_default() + }; + counts.insert( + ( + label("awa.job.state"), + label("awa.job.queue"), + label("awa.promote.outcome"), + ), + point.value(), + ); + } + } + } + } + } + + let key = |state: &str, queue: &str, outcome: &str| { + (state.to_string(), queue.to_string(), outcome.to_string()) + }; + assert_eq!( + counts, + BTreeMap::from([ + (key("retryable", "emails", "waiter_limit"), 1), + (key("scheduled", "billing", "waiter_timed_out"), 1), + (key("scheduled", "emails", "handed_off"), 2), + ]) + ); + } + /// A second-scale observation must not fall in the SDK's default /// millisecond-scale buckets (upper bounds 5, 10, 25, ...), and every /// histogram must export the intended explicit boundaries. diff --git a/awa-model/migrations/patches/definer_lane_provisioning.sql b/awa-model/migrations/patches/definer_lane_provisioning.sql new file mode 100644 index 000000000..b75a28cad --- /dev/null +++ b/awa-model/migrations/patches/definer_lane_provisioning.sql @@ -0,0 +1,74 @@ +-- Schema patch definer_lane_provisioning: create queue lanes as the schema +-- owner, whichever role touches them first (awa 0.7 migration v053, minus the +-- schema_version row). +-- +-- Applied after the v023 definitions, whose awa.refresh_lane_sequence_helpers +-- now installs two SECURITY DEFINER functions in each queue-storage schema, +-- provision_lane and provision_lane_sequences, owned by the schema owner. The +-- lane helpers create a missing sequence through them, and each new sequence +-- gets the grants of the head table it backs, so a lane is usable by every +-- runtime role however its creator's default privileges are set. The block +-- below refreshes the helpers on each installed schema and counts existing +-- lane sequences that need repair. It does not reassign them: ALTER SEQUENCE +-- OWNER takes the sequence's ACCESS EXCLUSIVE lock until commit. After the +-- patch commits, `awa migrate` repairs each one 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`. +-- +-- Live load: function replacement only. Reads the catalog; no table DDL, row +-- scans or sequence locks. +-- +-- 0.6 delivery: the 0.6 series cannot add a migration version, so this ships as +-- an idempotent schema patch. `awa migrate` applies it when the schema is +-- exactly at v040 and the probe in migrations.rs reports it missing; +-- `awa migrate --sql` prints it as a repeatable R__ script for external +-- runners. Released 0.6 binaries call the same helpers with the same arguments +-- and cursor semantics, and lanes they create are owned by the schema owner +-- too. Run it as the schema owner, or as a migrator that can SET ROLE to it. + +DO $$ +DECLARE + v_schema TEXT; + v_count BIGINT; +BEGIN + FOR v_schema IN + SELECT n.nspname + FROM pg_namespace AS n + WHERE has_schema_privilege(current_user, n.oid, 'USAGE') + AND EXISTS ( + SELECT 1 FROM pg_proc AS awa_p + WHERE awa_p.pronamespace = n.oid + AND awa_p.proname = 'claim_ready_runtime' + AND oidvectortypes(awa_p.proargtypes) + = 'text, bigint, double precision, double precision' + ) + ORDER BY n.nspname + 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_regprocedure(format('%I.queue_lane_sequence_name(text,text,smallint,smallint)', v_schema)) IS NULL + OR to_regprocedure(format('%I.set_sequence_next(text,bigint)', v_schema)) IS NULL THEN + CONTINUE; + END IF; + + -- Serialize with a concurrent prepare_schema() of the same schema. This + -- is an advisory lock, not a table lock: claims and completions do not + -- take it. + PERFORM pg_advisory_xact_lock( + hashtextextended(format('awa.queue_storage.install:%I', v_schema), 0) + ); + 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; +$$; diff --git a/awa-model/migrations/v023_install_queue_storage_substrate.sql b/awa-model/migrations/v023_install_queue_storage_substrate.sql index 5b7bfe127..4625ebada 100644 --- a/awa-model/migrations/v023_install_queue_storage_substrate.sql +++ b/awa-model/migrations/v023_install_queue_storage_substrate.sql @@ -47,6 +47,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 IF p_schema IS NULL OR p_schema !~ '^[a-z_][a-z0-9_]*$' THEN RAISE EXCEPTION 'refresh_lane_sequence_helpers: schema name must match ^[a-z_][a-z0-9_]*$; got %', @@ -80,28 +86,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 @@ -142,16 +160,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; @@ -208,16 +232,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; @@ -231,11 +261,368 @@ BEGIN $ddl$, 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; +-- 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; + CREATE OR REPLACE FUNCTION awa.install_queue_storage_substrate( p_schema TEXT, p_queue_slot_count INT DEFAULT 16, diff --git a/awa-model/src/lib.rs b/awa-model/src/lib.rs index 267495146..88d658d05 100644 --- a/awa-model/src/lib.rs +++ b/awa-model/src/lib.rs @@ -45,8 +45,8 @@ pub use partitioned_queue::{ partition_for_ordering_key, partition_hash64, PartitionedQueue, PartitionedQueueError, }; pub use queue_storage::{ - ClaimedEntry, ClaimedRuntimeJob, PruneOutcome, QueueCounts, QueueStorage, QueueStorageConfig, - RotateOutcome, SkipReason, TerminalDeltaRollupOutcome, + ClaimedEntry, ClaimedRuntimeJob, PromoteOutcome, PromotionLane, PruneOutcome, QueueCounts, + QueueStorage, QueueStorageConfig, RotateOutcome, SkipReason, TerminalDeltaRollupOutcome, }; pub use storage::StorageStatus; diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index 193c465fc..d6b4a6734 100644 --- a/awa-model/src/migrations.rs +++ b/awa-model/src/migrations.rs @@ -295,6 +295,61 @@ macro_rules! prepared_lane_sequences_pending_sql { }; } +/// Rows naming what still lacks the `definer_lane_provisioning` patch: the +/// lane-sequence check and repair functions, and each installed queue-storage +/// substrate (found as for `prepared_lane_sequences`) without +/// the SECURITY DEFINER provisioning functions or whose lane helpers do not +/// create sequences through them. Lane sequences the patch could not reassign +/// do not keep it pending: `awa storage lane-sequences` reports those. +macro_rules! definer_lane_provisioning_pending_sql { + () => { + "SELECT 'awa.lane_sequence_problems'::text AS target \ + WHERE to_regprocedure('awa.lane_sequence_problems(text,text)') IS NULL \ + UNION ALL \ + SELECT 'awa.repair_lane_sequence' \ + WHERE to_regprocedure('awa.repair_lane_sequence(text,text)') IS NULL \ + UNION ALL \ + SELECT 'awa.refresh_lane_sequence_helpers' \ + WHERE EXISTS ( \ + SELECT 1 FROM pg_proc AS p \ + WHERE p.pronamespace = 'awa'::regnamespace \ + AND p.proname = 'refresh_lane_sequence_helpers' \ + AND p.prosrc NOT LIKE '%provision_lane_sequences%') \ + UNION ALL \ + SELECT format('lane provisioning in schema %s', n.nspname) \ + FROM pg_namespace AS n \ + WHERE EXISTS ( \ + SELECT 1 FROM pg_proc AS p \ + WHERE p.pronamespace = n.oid \ + AND p.proname = 'claim_ready_runtime' \ + AND oidvectortypes(p.proargtypes) \ + = 'text, bigint, double precision, double precision') \ + AND EXISTS (SELECT 1 FROM pg_class AS c \ + WHERE c.relnamespace = n.oid AND c.relname = 'queue_enqueue_heads') \ + AND EXISTS (SELECT 1 FROM pg_class AS c \ + WHERE c.relnamespace = n.oid AND c.relname = 'queue_claim_heads') \ + AND EXISTS (SELECT 1 FROM pg_proc AS p \ + WHERE p.pronamespace = n.oid \ + AND p.proname = 'queue_lane_sequence_name' \ + AND oidvectortypes(p.proargtypes) = 'text, text, smallint, smallint') \ + AND EXISTS (SELECT 1 FROM pg_proc AS p \ + WHERE p.pronamespace = n.oid \ + AND p.proname = 'set_sequence_next' \ + AND oidvectortypes(p.proargtypes) = 'text, bigint') \ + AND ((SELECT count(DISTINCT p.proname) FROM pg_proc AS p \ + WHERE p.pronamespace = n.oid \ + AND p.proname IN ('provision_lane', 'provision_lane_sequences') \ + AND p.prosecdef \ + AND oidvectortypes(p.proargtypes) = 'text, smallint, smallint') < 2 \ + OR (SELECT count(DISTINCT p.proname) FROM pg_proc AS p \ + WHERE p.pronamespace = n.oid \ + AND p.proname IN ('ensure_lane_sequences', \ + 'queue_enqueue_head_sequence_sync', \ + 'queue_claim_head_sequence_sync') \ + AND p.prosrc LIKE '%provision_lane_sequences%') < 3)" + }; +} + /// SQL returning one `text` row per object or schema that still needs the /// named patch, for patches that touch more than one place. fn patch_targets_sql(name: &str) -> Option<&'static str> { @@ -303,6 +358,10 @@ fn patch_targets_sql(name: &str) -> Option<&'static str> { prepared_lane_sequences_pending_sql!(), " ORDER BY 1" )), + "definer_lane_provisioning" => Some(concat!( + definer_lane_provisioning_pending_sql!(), + " ORDER BY 1" + )), _ => None, } } @@ -347,6 +406,23 @@ pub const SCHEMA_PATCHES: &[SchemaPatch] = &[ ")" ), }, + SchemaPatch { + name: "definer_lane_provisioning", + description: "Create queue lanes as the schema owner (awa 0.7 migration v053)", + // The v023 definitions first, so the installer and the helper refresh + // carry the provisioning functions; the patch then refreshes each + // installed schema and repairs its existing lane sequences. + sql: concat!( + include_str!("../migrations/v023_install_queue_storage_substrate.sql"), + "\n", + include_str!("../migrations/patches/definer_lane_provisioning.sql"), + ), + applied_probe: concat!( + "SELECT NOT EXISTS (", + definer_lane_provisioning_pending_sql!(), + ")" + ), + }, ]; /// Schema patches a database still needs to reach this binary's schema shape. @@ -502,6 +578,11 @@ pub struct MigrationReport { /// change routing; both are `None` before the storage tables exist. pub storage_routing_before: Option, pub storage_routing_after: Option, + /// What the run did to each lane sequence that needed repair, when it + /// applied `definer_lane_provisioning`. Each sequence is repaired in its + /// own short transaction after the patch commits; a sequence left behind + /// does not make the run incomplete. + pub lane_sequences: Vec, } impl MigrationReport { @@ -660,6 +741,28 @@ async fn run_inner(conn: &mut PgConnection) -> Result })); } + // ALTER SEQUENCE ... OWNER takes the sequence's ACCESS EXCLUSIVE lock + // until its transaction ends, so the patch only installs the definitions + // and each lane sequence is repaired here, in its own transaction with a + // bounded lock wait, once the patch has committed. + let mut lane_sequences = Vec::new(); + if patches.iter().any(|patch| { + patch.name == "definer_lane_provisioning" && patch.outcome == PatchOutcome::Applied + }) { + match crate::storage::repair_lane_sequences_in_every_schema(conn).await { + Ok(repairs) => { + for next_step in crate::storage::lane_sequence_repair_next_steps(&repairs) { + tracing::warn!("{next_step}"); + } + lane_sequences = repairs; + } + Err(error) => tracing::warn!( + %error, + "Could not repair lane sequences; run `awa storage lane-sequences --repair`" + ), + } + } + // Ensure the admin metadata cache is warm. Since v006 removed the // synchronous triggers on jobs_hot, the cache is only updated by the // maintenance leader. Refreshing here guarantees queue_stats() and @@ -694,6 +797,7 @@ async fn run_inner(conn: &mut PgConnection) -> Result patches, storage_routing_before, storage_routing_after: storage_routing_conn(conn).await?, + lane_sequences, }) } diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 3b0b1503f..697e17ac7 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -10,7 +10,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; use uuid::Uuid; @@ -362,6 +362,29 @@ enum TerminalDeltaSlotRollup { Blocked, } +/// A `(queue, priority)` enqueue lane as promotion addresses it. Promotion +/// picks the shard when it re-enqueues a row, so the shard is not part of it. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct PromotionLane { + pub queue: String, + pub priority: i16, +} + +/// Outcome of [`QueueStorage::promote_due_lanes`] and +/// [`QueueStorage::promote_due_lane`]. +#[derive(Debug, Clone, PartialEq, Eq, Default)] +pub struct PromoteOutcome { + /// Rows moved into the ready ring. + pub promoted: usize, + /// Due rows the batch selected, including those left deferred on a + /// contended lane. A batch that selected its full size may have left + /// further due rows behind. + pub selected: usize, + /// Lanes whose promotion transaction hit `lock_timeout` and rolled back. + /// Their rows stay deferred. + pub contended: Vec, +} + /// Discriminator for [`PruneOutcome::SkippedActive`]. /// /// Multiple gates can fire `SkippedActive` for the same ring (e.g. queue @@ -420,7 +443,7 @@ fn is_lock_contention_error(err: &sqlx::Error) -> bool { ) } -async fn set_prune_lock_timeout_tx( +async fn set_lock_timeout_tx( tx: &mut sqlx::Transaction<'_, Postgres>, timeout: Duration, ) -> Result<(), AwaError> { @@ -2124,6 +2147,15 @@ struct ExistingReadyRow { payload: serde_json::Value, } +/// Which of a lane's due rows one promotion transaction takes. +#[derive(Debug, Clone, Copy)] +enum DueLaneRows<'a> { + /// These rows, picked by a batch across lanes, where still due. + Selected(&'a [i64]), + /// The oldest due rows, up to this many. + Oldest(i64), +} + #[derive(Debug, Clone, sqlx::FromRow)] struct DeletedLeaseRow { ready_slot: i32, @@ -2455,6 +2487,13 @@ 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` (the `definer_lane_provisioning` + /// schema patch) 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, prune_lock_timeout: Duration, } @@ -2499,6 +2538,8 @@ impl QueueStorage { shard_rotor: AtomicU16::new(0), enqueue_shards_cache: Mutex::new(HashMap::new()), ensured_lanes: Mutex::new(HashSet::new()), + lane_provisioner_present: AtomicBool::new(false), + lane_provisioner_warned: AtomicBool::new(false), prune_lock_timeout: DEFAULT_PRUNE_LOCK_TIMEOUT, }) } @@ -3458,9 +3499,10 @@ 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. + /// Provision every priority and shard of a queue ahead of its first + /// enqueue. Once the `definer_lane_provisioning` schema patch is applied, + /// 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, @@ -3499,11 +3541,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(&format!( r#" @@ -3574,6 +3640,61 @@ 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(&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 @@ -11695,8 +11816,17 @@ impl QueueStorage { unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload - FROM {schema}.ready_entries - WHERE job_id = $1 + FROM {schema}.ready_entries AS ready + WHERE ready.job_id = $1 + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) ORDER BY run_lease DESC, attempted_at DESC NULLS LAST, run_at DESC "#, )) @@ -14453,32 +14583,212 @@ impl QueueStorage { Ok(rescued) } + /// Promote up to `batch_size` due `state` rows into the ready ring. + /// Equivalent to [`Self::promote_due_lanes`] with no lock timeout and no + /// skipped lanes, returning the number of rows promoted. pub async fn promote_due( &self, pool: &PgPool, state: JobState, batch_size: i64, ) -> Result { + Ok(self + .promote_due_lanes(pool, state, batch_size, None, &[]) + .await? + .promoted) + } + + /// Select up to `batch_size` due `state` rows, oldest `run_at` first, + /// outside the lanes in `skip`, and promote each lane's share in its own + /// transaction. + /// + /// The selection skips rows another transaction has locked, such as a + /// lane promotion still waiting for its lane, so they cannot fill every + /// batch of a tick. It runs as its own statement, so its row locks end + /// with it; `deferred_jobs` is not partitioned, so no run-time pruning + /// meets those locks. + /// + /// Re-enqueueing a row takes its lane's enqueue reservation lock, which + /// producers hold until they commit, so a lane under heavy enqueue load + /// can keep its promotion waiting. A transaction per lane confines that + /// wait to the lane's own rows. With `lane_lock_timeout`, each lane's + /// transaction runs under that `lock_timeout`; a lane that hits it rolls + /// back, keeps its rows deferred and is reported in + /// [`PromoteOutcome::contended`], and the remaining lanes go ahead. + pub async fn promote_due_lanes( + &self, + pool: &PgPool, + state: JobState, + batch_size: i64, + lane_lock_timeout: Option, + skip: &[PromotionLane], + ) -> Result { if !matches!(state, JobState::Scheduled | JobState::Retryable) || batch_size <= 0 { - return Ok(0); + return Ok(PromoteOutcome::default()); } let schema = self.schema(); + let due: Vec<(i64, String, i16)> = sqlx::query_as(&format!( + r#" + SELECT job_id, queue, priority + FROM {schema}.deferred_jobs AS deferred + WHERE state = $1 + AND run_at <= clock_timestamp() + AND NOT EXISTS ( + SELECT 1 FROM awa.queue_meta + WHERE queue = deferred.queue AND paused = TRUE + ) + AND NOT EXISTS ( + SELECT 1 + FROM unnest($3::text[], $4::smallint[]) AS skip (queue, priority) + WHERE skip.queue = deferred.queue AND skip.priority = deferred.priority + ) + ORDER BY run_at ASC, priority ASC, job_id ASC + LIMIT $2 + FOR UPDATE OF deferred SKIP LOCKED + "# + )) + .bind(state) + .bind(batch_size) + .bind( + skip.iter() + .map(|lane| lane.queue.as_str()) + .collect::>(), + ) + .bind(skip.iter().map(|lane| lane.priority).collect::>()) + .fetch_all(pool) + .await + .map_err(map_sqlx_error)?; + + let mut outcome = PromoteOutcome { + selected: due.len(), + ..PromoteOutcome::default() + }; + let mut lanes: BTreeMap> = BTreeMap::new(); + for (job_id, queue, priority) in due { + lanes + .entry(PromotionLane { queue, priority }) + .or_default() + .push(job_id); + } + for (lane, job_ids) in lanes { + match self + .promote_lane( + pool, + state, + &lane, + DueLaneRows::Selected(&job_ids), + lane_lock_timeout, + ) + .await? + { + Some(promoted) => outcome.promoted += promoted, + None => outcome.contended.push(lane), + } + } + Ok(outcome) + } + + /// Promote up to `batch_size` due `state` rows of one lane, oldest + /// `run_at` first, in one transaction. With `lane_lock_timeout`, a + /// transaction that hits it rolls back and the lane is returned in + /// [`PromoteOutcome::contended`] with nothing selected. + pub async fn promote_due_lane( + &self, + pool: &PgPool, + state: JobState, + lane: &PromotionLane, + batch_size: i64, + lane_lock_timeout: Option, + ) -> Result { + if !matches!(state, JobState::Scheduled | JobState::Retryable) || batch_size <= 0 { + return Ok(PromoteOutcome::default()); + } + Ok( + match self + .promote_lane( + pool, + state, + lane, + DueLaneRows::Oldest(batch_size), + lane_lock_timeout, + ) + .await? + { + Some(promoted) => PromoteOutcome { + promoted, + selected: promoted, + contended: Vec::new(), + }, + None => PromoteOutcome { + contended: vec![lane.clone()], + ..PromoteOutcome::default() + }, + }, + ) + } + + /// One lane's promotion transaction. `None` when it hit `lock_timeout` + /// and rolled back. + async fn promote_lane( + &self, + pool: &PgPool, + state: JobState, + lane: &PromotionLane, + rows: DueLaneRows<'_>, + lock_timeout: Option, + ) -> Result, AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; - let moved: Vec = sqlx::query_as(&format!( + match self + .promote_lane_tx(&mut tx, state, lane, rows, lock_timeout) + .await + { + Ok(promoted) => match tx.commit().await { + Ok(()) => Ok(Some(promoted)), + Err(err) if is_lock_contention_error(&err) => Ok(None), + Err(err) => Err(map_sqlx_error(err)), + }, + Err(AwaError::Database(err)) if is_lock_contention_error(&err) => { + tx.rollback().await.map_err(map_sqlx_error)?; + Ok(None) + } + Err(err) => Err(err), + } + } + + async fn promote_lane_tx( + &self, + tx: &mut sqlx::Transaction<'_, Postgres>, + state: JobState, + lane: &PromotionLane, + rows: DueLaneRows<'_>, + lock_timeout: Option, + ) -> Result { + if let Some(timeout) = lock_timeout { + set_lock_timeout_tx(tx, timeout).await?; + } + + let schema = self.schema(); + let (row_filter, limit) = match rows { + DueLaneRows::Selected(_) => ("AND job_id = ANY($4)", ""), + DueLaneRows::Oldest(_) => ("", "ORDER BY run_at ASC, job_id ASC LIMIT $4"), + }; + let sql = format!( r#" DELETE FROM {schema}.deferred_jobs WHERE job_id IN ( SELECT job_id - FROM {schema}.deferred_jobs + FROM {schema}.deferred_jobs AS deferred WHERE state = $1 + AND queue = $2 + AND priority = $3 + {row_filter} AND run_at <= clock_timestamp() AND NOT EXISTS ( SELECT 1 FROM awa.queue_meta - WHERE queue = {schema}.deferred_jobs.queue AND paused = TRUE + WHERE queue = deferred.queue AND paused = TRUE ) - ORDER BY run_at ASC, priority ASC, job_id ASC - LIMIT $2 + {limit} FOR UPDATE SKIP LOCKED ) RETURNING @@ -14499,21 +14809,26 @@ impl QueueStorage { unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload "# - )) - .bind(state) - .bind(batch_size) - .fetch_all(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - + ); + let query = sqlx::query_as::<_, DeferredJobRow>(&sql) + .bind(state) + .bind(&lane.queue) + .bind(lane.priority); + let query = match rows { + DueLaneRows::Selected(job_ids) => query.bind(job_ids), + DueLaneRows::Oldest(batch_size) => query.bind(batch_size), + }; + let mut moved: Vec = + query.fetch_all(tx.as_mut()).await.map_err(map_sqlx_error)?; if moved.is_empty() { - tx.commit().await.map_err(map_sqlx_error)?; return Ok(0); } + // DELETE ... RETURNING has no order; lane sequence follows due order. + moved.sort_unstable_by_key(|row| (row.run_at, row.job_id)); + let promoted = moved.len(); let ready_rows: Vec = moved - .iter() - .cloned() + .into_iter() .map(|row| ExistingReadyRow { job_id: row.job_id, kind: row.kind, @@ -14531,15 +14846,10 @@ impl QueueStorage { payload: row.payload, }) .collect(); - let queues = ready_rows - .iter() - .map(|row| row.queue.clone()) - .collect::>(); - self.insert_existing_ready_rows_tx(&mut tx, ready_rows, Some(state)) + self.insert_existing_ready_rows_tx(tx, ready_rows, Some(state)) .await?; - self.notify_queues_tx(&mut tx, queues).await?; - tx.commit().await.map_err(map_sqlx_error)?; - Ok(moved.len()) + self.notify_queues_tx(tx, vec![lane.queue.clone()]).await?; + Ok(promoted) } async fn relation_has_rows_tx( @@ -15133,7 +15443,7 @@ impl QueueStorage { return Ok(TerminalDeltaSlotRollup::SkippedActive); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + set_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; let lock_delta = sqlx::query(&format!( "LOCK TABLE {delta_child} IN ACCESS EXCLUSIVE MODE" @@ -15431,7 +15741,7 @@ impl QueueStorage { }); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + set_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; let lock_tables = sqlx::query(&format!( "LOCK TABLE {ready_child}, {claim_attempt_child}, {done_child}, {tomb_child}, {segment_child}, {receipt_batch_child}, {receipt_tomb_child}, {delta_child} IN ACCESS EXCLUSIVE MODE" @@ -15816,7 +16126,7 @@ impl QueueStorage { return Ok(PruneOutcome::Noop); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + set_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; let lock_table = sqlx::query(&format!( "LOCK TABLE {lease_child} IN ACCESS EXCLUSIVE MODE" @@ -16121,7 +16431,7 @@ impl QueueStorage { }); } - set_prune_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; + set_lock_timeout_tx(&mut tx, self.prune_lock_timeout).await?; let lock_tables = sqlx::query(&format!( "LOCK TABLE {claim_child}, {claim_batch_child}, {closure_child}, {closure_batch_child} IN ACCESS EXCLUSIVE MODE" diff --git a/awa-model/src/storage.rs b/awa-model/src/storage.rs index 64fe8be58..21d8f0fa4 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)] @@ -379,3 +379,279 @@ 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 n.nspname::text, count(*) \ + FROM pg_namespace AS n \ + CROSS JOIN LATERAL awa.lane_sequence_problems(n.nspname) AS problem \ + WHERE has_schema_privilege(current_user, n.oid, 'USAGE') \ + AND EXISTS ( \ + SELECT 1 FROM pg_proc AS p \ + WHERE p.pronamespace = n.oid \ + AND p.proname = 'claim_ready_runtime' \ + AND oidvectortypes(p.proargtypes) \ + = 'text, bigint, double precision, double precision') \ + GROUP BY n.nspname \ + ORDER BY n.nspname", + ) + .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 717707c3a..11e37b18d 100644 --- a/awa-model/tests/least_privilege_runtime_test.rs +++ b/awa-model/tests/least_privilege_runtime_test.rs @@ -176,20 +176,25 @@ async fn check_runtime(pool: &PgPool, schema: &str) -> TestResult { ), ) .await?; - let error = run_sql( + // An unprepared lane is created through the schema owner's provisioning + // function, so the runtime still issues no DDL of its own. + run_sql( pool, format!( - "SELECT {schema}.ensure_lane_sequences('unprepared_queue', 2::smallint, 0::smallint)" + "SELECT {schema}.ensure_lane_sequences('unprepared_queue', 2::smallint, 0::smallint); \ + SELECT {schema}.provision_lane('unprepared_queue', 2::smallint, 0::smallint); \ + SELECT {schema}.reserve_enqueue_seq('unprepared_queue', 2::smallint, 0::smallint, 1);" ), ) - .await - .unwrap_err(); - assert_eq!(database_error_code(&error).as_deref(), Some("42501")); - assert!(error - .as_database_error() - .unwrap() - .message() - .contains("has not been provisioned")); + .await?; + let foreign_owned: i64 = sqlx::query_scalar(&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); let error = run_sql( pool, format!("CREATE TABLE {schema}.must_not_create(id integer)"), @@ -354,10 +359,12 @@ async fn exercise_stale_custom_schema(options: PgConnectOptions, migrator: &str) .replace("__schema__", "custom_lanes"), ) .await?; + // The old helper also predates lane provisioning, so both patches that + // refresh the helpers name the schema. let pending = migrations::pending_schema_patches(&pool).await?; assert_eq!( pending.iter().map(|patch| patch.name).collect::>(), - vec!["prepared_lane_sequences"] + vec!["prepared_lane_sequences", "definer_lane_provisioning"] ); let status = lane_patch_status(&pool).await?; assert!(!status.applied); @@ -478,10 +485,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.9 runtime with CREATE did. let runtime_pool = PgPoolOptions::new() .max_connections(2) .connect_with(options.username(runtime)) diff --git a/awa-python/Cargo.lock b/awa-python/Cargo.lock index fa27a3b79..8922f4d9a 100644 --- a/awa-python/Cargo.lock +++ b/awa-python/Cargo.lock @@ -2355,6 +2355,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] diff --git a/awa-worker/Cargo.toml b/awa-worker/Cargo.toml index 91dd58594..afb66440c 100644 --- a/awa-worker/Cargo.toml +++ b/awa-worker/Cargo.toml @@ -21,7 +21,7 @@ sqlx.workspace = true serde.workspace = true serde_json.workspace = true tokio.workspace = true -tokio-util.workspace = true +tokio-util = { workspace = true, features = ["rt"] } async-trait.workspace = true thiserror.workspace = true chrono.workspace = true diff --git a/awa-worker/src/client.rs b/awa-worker/src/client.rs index e5a498dbb..3cfd9052e 100644 --- a/awa-worker/src/client.rs +++ b/awa-worker/src/client.rs @@ -27,6 +27,7 @@ use std::time::Duration; use tokio::sync::{Mutex, RwLock}; use tokio::task::JoinSet; use tokio_util::sync::CancellationToken; +use tokio_util::task::TaskTracker; use tracing::{info, warn}; use uuid::Uuid; @@ -1085,6 +1086,7 @@ impl ClientBuilder { dispatcher_handles: RwLock::new(Vec::new()), service_handles: RwLock::new(Vec::new()), job_set: Arc::new(Mutex::new(JoinSet::new())), + finalizers: TaskTracker::new(), in_flight: Arc::new(InFlightRegistry::default()), queue_in_flight, dispatcher_alive, @@ -1185,6 +1187,8 @@ pub struct Client { service_handles: RwLock>>, /// JoinSet tracking in-flight job tasks for graceful drain. job_set: Arc>>, + /// Finalization of attempts whose handler has returned, drained on shutdown. + finalizers: TaskTracker, in_flight: InFlightMap, queue_in_flight: Arc>>, dispatcher_alive: Arc>>, @@ -1566,6 +1570,7 @@ impl Client { completion_handle, effective_storage.clone(), self.dlq_policy.clone(), + self.finalizers.clone(), )); // Admin cancellation listener: fires the in-flight cancel flag @@ -1804,8 +1809,9 @@ impl Client { /// 1. Stop dispatchers (no new jobs claimed) /// 2. Signal in-flight jobs to cancel /// 3. Wait for dispatchers to exit - /// 4. Drain in-flight jobs (heartbeat + maintenance still alive!) - /// 5. Stop heartbeat + maintenance + /// 4. Drain in-flight jobs and their finalization (heartbeat, + /// maintenance and the completion batcher still alive!) + /// 5. Stop heartbeat, maintenance and the completion batcher pub async fn shutdown(&self, timeout: Duration) { info!("Initiating graceful shutdown"); @@ -1863,6 +1869,14 @@ impl Client { ); } + self.finalizers.close(); + if tokio::time::timeout_at(drain_deadline, self.finalizers.wait()) + .await + .is_err() + { + warn!("Shutdown drain timeout exceeded waiting for durable finalization"); + } + // Phase 5: Stop background services (heartbeat + maintenance) self.service_cancel.cancel(); let service_handles: Vec<_> = { diff --git a/awa-worker/src/dispatcher.rs b/awa-worker/src/dispatcher.rs index 11e3e9dbd..916edb02f 100644 --- a/awa-worker/src/dispatcher.rs +++ b/awa-worker/src/dispatcher.rs @@ -851,6 +851,7 @@ mod tests { completion, storage.clone(), crate::executor::DlqPolicy::default(), + tokio_util::task::TaskTracker::new(), )); let config = QueueConfig { max_workers: 1, diff --git a/awa-worker/src/executor.rs b/awa-worker/src/executor.rs index 3d99ec8df..a4fef41c3 100644 --- a/awa-worker/src/executor.rs +++ b/awa-worker/src/executor.rs @@ -11,6 +11,7 @@ use std::collections::HashMap; use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; use std::sync::Arc; use std::time::Duration; +use tokio_util::task::TaskTracker; use tracing::{error, info, info_span, warn, Instrument}; /// Result of executing a job handler. @@ -267,6 +268,7 @@ pub struct JobExecutor { completion_batcher: CompletionBatcherHandle, storage: RuntimeStorage, dlq_policy: DlqPolicy, + finalizers: TaskTracker, } impl JobExecutor { @@ -288,6 +290,7 @@ impl JobExecutor { completion_batcher: CompletionBatcherHandle, storage: RuntimeStorage, dlq_policy: DlqPolicy, + finalizers: TaskTracker, ) -> Self { Self { pool, @@ -301,6 +304,7 @@ impl JobExecutor { completion_batcher, storage, dlq_policy, + finalizers, } } @@ -327,6 +331,7 @@ impl JobExecutor { let completion_batcher = self.completion_batcher.clone(); let storage = self.storage.clone(); let dlq_policy = self.dlq_policy.clone(); + let finalizers = self.finalizers.clone(); let job_id = job.id; let job_run_lease = job.run_lease; let job_kind = job.kind.clone(); @@ -444,7 +449,10 @@ impl JobExecutor { return; }; let dlq_enabled = dlq_policy.enabled_for(&job_queue); - tokio::spawn(async move { + // Graceful shutdown waits for this before it stops the completion + // batcher: the attempt has already left `in_flight`, so nothing + // heartbeats it, and a finalize lost at exit leaves it to rescue. + finalizers.spawn(async move { // Every finalize path is one lease-guarded transaction, and a // deadlock abort rolls the whole transaction back, so running // it again is safe. Without the retry the job would sit diff --git a/awa-worker/src/maintenance.rs b/awa-worker/src/maintenance.rs index b7c134177..4623963aa 100644 --- a/awa-worker/src/maintenance.rs +++ b/awa-worker/src/maintenance.rs @@ -1,4 +1,5 @@ use crate::executor::DlqPolicy; +use crate::metrics::AwaMetrics; use crate::runtime::InFlightMap; use crate::storage::{QueueStorageRuntime, RuntimeStorage}; use awa_model::cron::{ @@ -7,7 +8,8 @@ use awa_model::cron::{ #[cfg(test)] use awa_model::SkipReason; use awa_model::{ - JobRow, JobState, PeriodicJob, PruneOutcome, RotateOutcome, TerminalDeltaRollupOutcome, + AwaError, JobRow, JobState, PeriodicJob, PromotionLane, PruneOutcome, QueueStorage, + RotateOutcome, TerminalDeltaRollupOutcome, }; use chrono::Utc; use croner::Cron; @@ -17,7 +19,7 @@ use std::collections::{HashMap, HashSet}; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; -use tokio::task::JoinHandle; +use tokio::task::{JoinHandle, JoinSet}; use tokio_util::sync::CancellationToken; use tracing::{debug, error, info, warn}; use uuid::Uuid; @@ -719,6 +721,232 @@ const CRON_CATCH_UP_LIMIT: usize = 1_000; const TERMINAL_COUNT_ROLLUP_MAX_SLOTS_PER_TICK: usize = 4; type QueueStorageMetricRow = (String, i64, i64, i64, i64, i64, i64, i64, Option); +/// How long a promotion transaction waits on one lane's locks before the +/// leader hands the lane to a dedicated waiter and goes on without it. +const PROMOTE_LANE_LOCK_TIMEOUT: Duration = Duration::from_millis(50); +/// Lock wait of a dedicated lane waiter. PostgreSQL grants a lock to its +/// waiters in arrival order, so a waiter is served after the producers queued +/// ahead of it, however busy the lane stays. +const PROMOTE_LANE_WAITER_LOCK_TIMEOUT: Duration = Duration::from_secs(30); +/// Most lanes with a dedicated waiter at once. A waiter holds a pool +/// connection while it waits, so smaller pools get fewer: one per eight +/// connections, and always one so a busy lane is still promoted. +const PROMOTE_MAX_LANE_WAITERS: usize = 4; + +fn lane_waiter_limit(pool_max_connections: u32) -> usize { + (pool_max_connections as usize / 8).clamp(1, PROMOTE_MAX_LANE_WAITERS) +} + +/// Queue-storage promotion of due scheduled and retryable rows, as the +/// maintenance leader runs it for one store. +/// +/// Each lane's share of a batch is promoted in its own transaction under a +/// short `lock_timeout`, so a lane whose enqueue lock is busy delays only its +/// own rows. A lane that times out keeps its rows deferred and is handed to a +/// dedicated waiter task, which promotes it under a long lock timeout while +/// the leader carries on with other lanes and other maintenance branches. +/// Dropping the promoter aborts its waiters. +#[doc(hidden)] +pub struct QueueStoragePromoter { + pool: PgPool, + store: Arc, + metrics: AwaMetrics, + waiting: Arc>>, + waiters: std::sync::Mutex>, + max_waiters: usize, +} + +impl QueueStoragePromoter { + pub fn new(pool: PgPool, store: Arc, metrics: AwaMetrics) -> Self { + let max_waiters = lane_waiter_limit(pool.options().get_max_connections()); + Self { + max_waiters, + pool, + store, + metrics, + waiting: Arc::default(), + waiters: std::sync::Mutex::new(JoinSet::new()), + } + } + + /// One promotion tick for `state`: batches until one selects fewer rows + /// than the batch size, each skipping lanes that have a waiter or were + /// contended earlier in the tick. Returns the rows promoted by the tick + /// itself, not by waiters. + pub async fn promote_due( + &self, + state: JobState, + cancel: &CancellationToken, + ) -> Result { + { + let mut waiters = self.waiters.lock().expect("lane waiters poisoned"); + while waiters.try_join_next().is_some() {} + } + + let mut contended: Vec = Vec::new(); + let mut promoted = 0usize; + for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { + if cancel.is_cancelled() { + break; + } + let mut skip = self.waiting_lanes(state); + skip.extend(contended.iter().cloned()); + let batch_started = Instant::now(); + let outcome = self + .store + .promote_due_lanes( + &self.pool, + state, + PROMOTE_BATCH_SIZE, + Some(PROMOTE_LANE_LOCK_TIMEOUT), + &skip, + ) + .await?; + self.metrics.record_promotion_batch( + state.as_str(), + outcome.promoted as u64, + batch_started.elapsed(), + ); + promoted += outcome.promoted; + for lane in outcome.contended { + self.hand_off(state, lane.clone()); + contended.push(lane); + } + if (outcome.selected as i64) < PROMOTE_BATCH_SIZE { + break; + } + } + Ok(promoted) + } + + fn waiting_lanes(&self, state: JobState) -> Vec { + self.waiting + .lock() + .expect("waiting lanes poisoned") + .iter() + .filter(|(waiting_state, _)| *waiting_state == state) + .map(|(_, lane)| lane.clone()) + .collect() + } + + fn hand_off(&self, state: JobState, lane: PromotionLane) { + let key = (state, lane); + { + let mut waiting = self.waiting.lock().expect("waiting lanes poisoned"); + if waiting.len() >= self.max_waiters { + self.metrics.record_promotion_lane_contention( + state.as_str(), + &key.1.queue, + "waiter_limit", + ); + debug!( + queue = %key.1.queue, + priority = key.1.priority, + "Promotion lane contended with every lane waiter busy; retrying next tick" + ); + return; + } + if !waiting.insert(key.clone()) { + return; + } + } + self.metrics + .record_promotion_lane_contention(state.as_str(), &key.1.queue, "handed_off"); + let guard = WaitingLane { + waiting: Arc::clone(&self.waiting), + key, + }; + let pool = self.pool.clone(); + let store = Arc::clone(&self.store); + let metrics = self.metrics.clone(); + self.waiters + .lock() + .expect("lane waiters poisoned") + .spawn(async move { + promote_waiting_lane(&pool, &store, &metrics, &guard.key).await; + drop(guard); + }); + } +} + +/// Holds a lane in [`QueueStoragePromoter`]'s waiting set until its waiter +/// finishes, panics or is aborted. +struct WaitingLane { + waiting: Arc>>, + key: (JobState, PromotionLane), +} + +impl Drop for WaitingLane { + fn drop(&mut self) { + if let Ok(mut waiting) = self.waiting.lock() { + waiting.remove(&self.key); + } + } +} + +async fn promote_waiting_lane( + pool: &PgPool, + store: &QueueStorage, + metrics: &AwaMetrics, + (state, lane): &(JobState, PromotionLane), +) { + for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { + let batch_started = Instant::now(); + match store + .promote_due_lane( + pool, + *state, + lane, + PROMOTE_BATCH_SIZE, + Some(PROMOTE_LANE_WAITER_LOCK_TIMEOUT), + ) + .await + { + Ok(outcome) if !outcome.contended.is_empty() => { + metrics.record_promotion_lane_contention( + state.as_str(), + &lane.queue, + "waiter_timed_out", + ); + warn!( + queue = %lane.queue, + priority = lane.priority, + lock_timeout_ms = PROMOTE_LANE_WAITER_LOCK_TIMEOUT.as_millis() as u64, + "Promotion of a contended lane timed out waiting for its locks; retrying next tick" + ); + return; + } + Ok(outcome) => { + metrics.record_promotion_batch( + state.as_str(), + outcome.promoted as u64, + batch_started.elapsed(), + ); + if (outcome.selected as i64) < PROMOTE_BATCH_SIZE { + return; + } + } + Err(err) => { + error!( + queue = %lane.queue, + priority = lane.priority, + error = %err, + "Failed to promote a contended lane" + ); + return; + } + } + } +} + +/// The leader's queue-storage promoters: its own store's, and the configured +/// store a canonical-resolved leader promotes once routing has flipped +/// (#456). Scoped to one leadership term. +struct LeaderPromoters { + primary: Option, + standby: Option, +} + impl MaintenanceService { #[allow(clippy::too_many_arguments)] pub(crate) fn new( @@ -1019,6 +1247,22 @@ impl MaintenanceService { // tick. Local to the leader loop so a re-election resets // backoff to zero. let prune_tracker = PruneBackoffTracker::new(); + let promoters = LeaderPromoters { + primary: self.storage.queue_storage().map(|runtime| { + QueueStoragePromoter::new( + self.pool.clone(), + Arc::clone(&runtime.store), + self.metrics.clone(), + ) + }), + standby: self.standby_queue_storage.as_ref().map(|runtime| { + QueueStoragePromoter::new( + self.pool.clone(), + Arc::clone(&runtime.store), + self.metrics.clone(), + ) + }), + }; // Skip the first immediate tick heartbeat_rescue_timer.tick().await; @@ -1084,7 +1328,7 @@ impl MaintenanceService { } _ = promote_timer.tick() => { if let Some(timer) = branch_tracker.try_begin_without_cooldown("promote_scheduled", self.promote_interval, &self.metrics) { - self.promote_scheduled().await; + self.promote_scheduled(&promoters).await; timer.finish(); } } @@ -2324,13 +2568,20 @@ impl MaintenanceService { } /// Promote scheduled jobs that are now due. - #[tracing::instrument(skip(self), name = "maintenance.promote")] - async fn promote_scheduled(&self) { - if let Err(err) = self.promote_due_state("scheduled", "scheduled jobs").await { + #[tracing::instrument(skip(self, promoters), name = "maintenance.promote")] + async fn promote_scheduled(&self, promoters: &LeaderPromoters) { + if let Err(err) = self + .promote_due_state(JobState::Scheduled, "scheduled jobs", promoters) + .await + { error!(error = %err, "Failed to promote scheduled jobs"); } if let Err(err) = self - .promote_due_state("retryable", "retryable jobs (backoff elapsed)") + .promote_due_state( + JobState::Retryable, + "retryable jobs (backoff elapsed)", + promoters, + ) .await { error!(error = %err, "Failed to promote retryable jobs"); @@ -2339,138 +2590,55 @@ impl MaintenanceService { async fn promote_due_state( &self, - state: &'static str, + state: JobState, label: &'static str, - ) -> Result<(), awa_model::AwaError> { + promoters: &LeaderPromoters, + ) -> Result<(), AwaError> { + let canonical_state = match state { + JobState::Scheduled => "scheduled", + JobState::Retryable => "retryable", + other => { + return Err(AwaError::Validation(format!( + "unsupported promote state: {other:?}" + ))); + } + }; let mut promoted_total = 0usize; let mut notified_queues = HashSet::new(); - for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { - if self.cancel.is_cancelled() { - break; - } - - match &self.storage { - RuntimeStorage::Canonical => { - let (promoted, queues) = self - .promote_due_batch(state) - .await - .map_err(awa_model::AwaError::Database)?; - if promoted == 0 { - break; - } - - promoted_total += promoted; - notified_queues.extend(queues); - - if promoted < PROMOTE_BATCH_SIZE as usize { - break; - } - } - RuntimeStorage::QueueStorage(runtime) => { - let job_state = match state { - "scheduled" => awa_model::JobState::Scheduled, - "retryable" => awa_model::JobState::Retryable, - other => { - return Err(awa_model::AwaError::Validation(format!( - "unsupported queue storage promote state: {other}" - ))); - } - }; - let promote_start = std::time::Instant::now(); - let promoted = runtime - .store - .promote_due(&self.pool, job_state, PROMOTE_BATCH_SIZE) - .await?; - self.metrics.record_promotion_batch( - state, - promoted as u64, - promote_start.elapsed(), - ); - if promoted == 0 { - break; - } - - promoted_total += promoted; - - if promoted < PROMOTE_BATCH_SIZE as usize { - break; + match &self.storage { + RuntimeStorage::Canonical => { + promoted_total += self + .promote_canonical_due(canonical_state, &mut notified_queues) + .await?; + + // #456 mirror: a canonical-resolved (pre-flip) leader must + // promote the queue-storage deferred backlog once routing has + // flipped — re-scheduled canonical jobs migrate there, and + // post-flip runtimes that would promote them may not hold + // leadership. The active-schema guard keeps this a no-op until + // the flip. + if let Some(promoter) = &promoters.standby { + if self.active_standby_queue_storage().await.is_some() { + promoted_total += promoter.promote_due(state, &self.cancel).await?; } } } - } - - // #456: while a storage transition drains, due jobs can still sit in - // the canonical deferred backlog. Drain-only runtimes execute - // canonical work but stop leading maintenance once a post-flip - // runtime wins the election, so a queue-storage leader must promote - // the canonical plane too. Both queries are cheap no-ops once the - // canonical tables are empty (every finalized or fresh install). - if matches!(&self.storage, RuntimeStorage::QueueStorage(_)) { - for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { - if self.cancel.is_cancelled() { - break; - } - let (promoted, queues) = self - .promote_due_batch(state) - .await - .map_err(awa_model::AwaError::Database)?; - if promoted == 0 { - break; + RuntimeStorage::QueueStorage(_) => { + if let Some(promoter) = &promoters.primary { + promoted_total += promoter.promote_due(state, &self.cancel).await?; } - promoted_total += promoted; - notified_queues.extend(queues); - if promoted < PROMOTE_BATCH_SIZE as usize { - break; - } - } - } - // #456 mirror: a canonical-resolved (pre-flip) leader must promote - // the queue-storage deferred backlog once routing has flipped — - // re-scheduled canonical jobs migrate there, and post-flip runtimes - // that would promote them may not hold leadership. The active-schema - // guard keeps this a no-op until the flip. - if matches!(&self.storage, RuntimeStorage::Canonical) { - if let Some(runtime) = &self.standby_queue_storage { - let active_schema: Option = - sqlx::query_scalar("SELECT awa.active_queue_storage_schema()") - .fetch_one(&self.pool) - .await - .map_err(awa_model::AwaError::Database)?; - if active_schema.as_deref() == Some(runtime.store.schema()) { - let job_state = match state { - "scheduled" => awa_model::JobState::Scheduled, - "retryable" => awa_model::JobState::Retryable, - other => { - return Err(awa_model::AwaError::Validation(format!( - "unsupported queue storage promote state: {other}" - ))); - } - }; - for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { - if self.cancel.is_cancelled() { - break; - } - let promote_start = std::time::Instant::now(); - let promoted = runtime - .store - .promote_due(&self.pool, job_state, PROMOTE_BATCH_SIZE) - .await?; - self.metrics.record_promotion_batch( - state, - promoted as u64, - promote_start.elapsed(), - ); - if promoted == 0 { - break; - } - promoted_total += promoted; - if promoted < PROMOTE_BATCH_SIZE as usize { - break; - } - } - } + // #456: while a storage transition drains, due jobs can still + // sit in the canonical deferred backlog. Drain-only runtimes + // execute canonical work but stop leading maintenance once a + // post-flip runtime wins the election, so a queue-storage + // leader must promote the canonical plane too. This is a + // cheap no-op once the canonical tables are empty (every + // finalized or fresh install). + promoted_total += self + .promote_canonical_due(canonical_state, &mut notified_queues) + .await?; } } @@ -2478,7 +2646,7 @@ impl MaintenanceService { debug!( count = promoted_total, queues = notified_queues.len(), - state, + state = canonical_state, "Promoted {label}" ); } @@ -2486,6 +2654,34 @@ impl MaintenanceService { Ok(()) } + /// Promote due canonical `scheduled_jobs` rows in batches until one comes + /// back short. + async fn promote_canonical_due( + &self, + state: &'static str, + notified_queues: &mut HashSet, + ) -> Result { + let mut promoted_total = 0usize; + for _ in 0..PROMOTE_MAX_BATCHES_PER_TICK { + if self.cancel.is_cancelled() { + break; + } + let (promoted, queues) = self + .promote_due_batch(state) + .await + .map_err(AwaError::Database)?; + if promoted == 0 { + break; + } + promoted_total += promoted; + notified_queues.extend(queues); + if promoted < PROMOTE_BATCH_SIZE as usize { + break; + } + } + Ok(promoted_total) + } + /// SQL template for promotion. The state literal is injected directly /// (not as a parameter) so the planner can match the partial index on /// `(run_at, id) WHERE state = ''`. With a parameter, the planner @@ -3489,6 +3685,15 @@ mod tests { } } + #[test] + fn lane_waiters_scale_with_the_pool_and_never_drop_to_zero() { + assert_eq!(lane_waiter_limit(1), 1); + assert_eq!(lane_waiter_limit(10), 1); + assert_eq!(lane_waiter_limit(16), 2); + assert_eq!(lane_waiter_limit(32), 4); + assert_eq!(lane_waiter_limit(200), 4); + } + #[test] fn compute_fire_times_coalesces_missed_existing_fires_by_default() { let last = Utc.with_ymd_and_hms(2026, 5, 7, 12, 0, 0).unwrap(); diff --git a/awa/Cargo.toml b/awa/Cargo.toml index 437cb057b..23ba05b62 100644 --- a/awa/Cargo.toml +++ b/awa/Cargo.toml @@ -25,6 +25,7 @@ sqlx.workspace = true serde.workspace = true serde_json.workspace = true tokio = { workspace = true, features = ["full", "test-util"] } +tokio-util.workspace = true async-trait.workspace = true chrono.workspace = true uuid.workspace = true diff --git a/awa/tests/ci_timing.rs b/awa/tests/ci_timing.rs new file mode 100644 index 000000000..3704a0ba7 --- /dev/null +++ b/awa/tests/ci_timing.rs @@ -0,0 +1,186 @@ +//! Shared CI contention scaling for the nightly chaos and benchmark suites. +//! +//! The nightly suites run on shared GitHub runners whose CPU allocation +//! varies run to run. Three assertion shapes are sensitive to that and +//! flaked repeatedly through 2026-07 (#399, #434): +//! +//! 1. **Wall-clock waits.** A test waits N seconds for a state the +//! runtime reaches in milliseconds when it has a core to itself. +//! Scale with [`scaled_timeout`]. +//! +//! 2. **Aggressive rescue cadences.** A chaos client sets +//! `heartbeat_staleness` a few multiples above `heartbeat_interval` +//! so its own backdating triggers rescue promptly. Under contention +//! a *live* worker's heartbeat can miss that window, so the runtime +//! correctly rescues a healthy attempt — and the test sees a genuine +//! duplicate completion, which reads as a correctness failure. Scale +//! with [`scaled_staleness`]: the interval stays fast so rescue is +//! still exercised, but the staleness window gains the same margin +//! as the waits around it. +//! +//! 3. **Minimum-progress floors.** A gate asserts "at least N of these +//! happened" to catch a stalled subsystem. The floor must sit far +//! below the *observed* operating point, not just below the nominal +//! one. Scale with [`contention_floor`]. +//! +//! Scaling only ever loosens a bound, and only on CI. A local run keeps +//! the strict values, so a real regression still fails fast on a +//! developer machine. Override with `AWA_CHAOS_TIMEOUT_MULTIPLIER`. +#![allow(dead_code)] + +use std::time::Duration; + +/// Upper bound on the override. `Duration::mul_f64` panics on a non-finite +/// or overflowing result, so an unbounded override would turn a typo into a +/// failed nightly rather than a loose bound. +const MAX_MULTIPLIER: f64 = 100.0; + +/// How much slack to give contention-sensitive bounds. `1.0` locally, +/// `3.0` on CI, or the `AWA_CHAOS_TIMEOUT_MULTIPLIER` override (clamped to +/// `1.0 ..= MAX_MULTIPLIER` so the override can only ever loosen, and can +/// never produce a duration `mul_f64` refuses). +pub fn chaos_timeout_multiplier() -> f64 { + let override_var = std::env::var("AWA_CHAOS_TIMEOUT_MULTIPLIER").ok(); + multiplier_from(override_var.as_deref(), std::env::var_os("CI").is_some()) +} + +/// The multiplier decision, with the environment passed in so it can be +/// tested without mutating process-global state. +fn multiplier_from(override_var: Option<&str>, is_ci: bool) -> f64 { + // `"inf"` and `"nan"` parse successfully as f64, and both reach + // `Duration::mul_f64`, which panics rather than saturating. Treat a + // non-finite override as no override at all. + if let Some(parsed) = override_var + .and_then(|raw| raw.parse::().ok()) + .filter(|parsed| parsed.is_finite()) + { + return parsed.clamp(1.0, MAX_MULTIPLIER); + } + + if is_ci { + 3.0 + } else { + 1.0 + } +} + +/// Grow a wait deadline by the contention multiplier. +pub fn scaled_timeout(timeout: Duration) -> Duration { + timeout.mul_f64(chaos_timeout_multiplier()) +} + +/// Grow a heartbeat-staleness window by the contention multiplier. +/// +/// Distinct from [`scaled_timeout`] only in intent: this one is passed to +/// `ClientBuilder::heartbeat_staleness`, where the cost of being too tight +/// is a spurious rescue of a live attempt rather than a timeout. Paired +/// heartbeat/rescue *intervals* are deliberately left unscaled so the +/// rescue path still runs at chaos cadence. +pub fn scaled_staleness(staleness: Duration) -> Duration { + staleness.mul_f64(chaos_timeout_multiplier()) +} + +/// Shrink a minimum-progress floor by the contention multiplier. +/// +/// Use for "this subsystem must have advanced at least N times" gates. +/// The returned floor is at least 1: the regression these gates exist to +/// catch is a fully stalled subsystem, so zero progress must still fail. +pub fn contention_floor(nominal: i64) -> i64 { + let scaled = (nominal as f64 / chaos_timeout_multiplier()).floor() as i64; + scaled.max(1) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn local_runs_keep_strict_bounds() { + assert_eq!(multiplier_from(None, false), 1.0); + } + + #[test] + fn ci_runs_get_margin() { + assert_eq!(multiplier_from(None, true), 3.0); + } + + #[test] + fn override_can_only_loosen() { + assert_eq!(multiplier_from(Some("0.1"), true), 1.0); + assert_eq!(multiplier_from(Some("-5"), true), 1.0); + assert_eq!(multiplier_from(Some("6"), false), 6.0); + } + + #[test] + fn unparseable_override_falls_back_to_the_environment() { + assert_eq!(multiplier_from(Some("banana"), true), 3.0); + assert_eq!(multiplier_from(Some(""), false), 1.0); + } + + #[test] + fn non_finite_override_falls_back_to_the_environment() { + // These all parse as f64 and would otherwise reach `mul_f64`, which + // panics on a non-finite result instead of saturating. + for raw in ["inf", "-inf", "infinity", "NaN", "nan"] { + assert_eq!(multiplier_from(Some(raw), true), 3.0, "raw={raw}"); + assert_eq!(multiplier_from(Some(raw), false), 1.0, "raw={raw}"); + } + } + + #[test] + fn oversized_override_is_clamped() { + assert_eq!(multiplier_from(Some("1e308"), true), MAX_MULTIPLIER); + assert_eq!(multiplier_from(Some("10000"), true), MAX_MULTIPLIER); + } + + #[test] + fn every_multiplier_produces_a_usable_duration() { + // `Duration::mul_f64` panics rather than saturating, so the clamp is + // what keeps a bad override from failing the nightly it was set to + // rescue. + for raw in [ + "inf", "-inf", "NaN", "1e308", "10000", "0", "-1", "banana", "", + ] { + for is_ci in [true, false] { + let m = multiplier_from(Some(raw), is_ci); + assert!(m.is_finite() && (1.0..=MAX_MULTIPLIER).contains(&m)); + // Would panic if the multiplier were unbounded. + let _ = Duration::from_secs(60).mul_f64(m); + } + } + } + + #[test] + fn scaling_a_timeout_grows_it() { + // 250ms staleness against a 50ms heartbeat interval is a 5x margin + // locally; on CI it becomes 15x, which is what stops a contended + // runner's live worker from being rescued as if it had died. + assert_eq!( + Duration::from_millis(250).mul_f64(multiplier_from(None, true)), + Duration::from_millis(750) + ); + } + + #[test] + fn progress_floor_keeps_margin_under_the_observed_rate() { + // The receipt gate's nominal floor is duration_secs / 4 = 45 for a + // 180s run. #399 saw a healthy run produce 41, so the CI floor has + // to sit well under that. + let ci_floor = (45.0_f64 / multiplier_from(None, true)).floor() as i64; + assert_eq!(ci_floor, 15); + assert!( + ci_floor < 41, + "floor must clear the observed operating point" + ); + } + + #[test] + fn progress_floor_never_reaches_zero() { + // Whatever the multiplier, a fully pinned ring (zero rotations) has + // to keep failing — that is the regression the gate exists for. + assert_eq!(contention_floor(45), contention_floor(45).max(1)); + let absurd = (45.0_f64 / multiplier_from(Some("10000"), true)).floor() as i64; + assert_eq!(absurd.max(1), 1); + assert_eq!(contention_floor(0), 1); + } +} 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..ce7d53040 --- /dev/null +++ b/awa/tests/enqueue_lane_maintenance_coupling_test.rs @@ -0,0 +1,427 @@ +//! 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. The leader +//! runs its branches (promotion, rescue, rotation, ...) one after another on a +//! single task, so a promotion parked behind a hot lane would hold up the due +//! jobs of other queues and every other maintenance branch. The leader's +//! promoter gives each lane its own short transaction and hands a lane whose +//! lock stays busy to a dedicated waiter, so neither should wait on it. +//! +//! `test_maintenance_under_hot_lane` measures, with and without concurrent +//! producers on one hot lane: +//! * how long each promotion tick holds the leader task, 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::worker::maintenance::QueueStoragePromoter; +use awa::worker::AwaMetrics; +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_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 { + /// How long each promotion tick held the leader task. + tick_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; each tick runs the leader's promoter) 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: &Arc, + 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 promoter = + QueueStoragePromoter::new(pool.clone(), Arc::clone(store), AwaMetrics::from_global()); + let cancel = tokio_util::sync::CancellationToken::new(); + + 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; + let tick_started = Instant::now(); + promoter + .promote_due(JobState::Scheduled, &cancel) + .await + .expect("promote_due"); + report.tick_ms.push(tick_started.elapsed().as_secs_f64() * 1e3); + 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 tick on leader ms p50={:.0} p99={:.0} max={:.0} | 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.tick_ms, 0.5), + percentile(&report.tick_ms, 0.99), + max_of(&report.tick_ms), + 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 = Arc::new(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 = Arc::new(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/v0613_lane_sequence_helpers.sql b/awa/tests/fixtures/v0613_lane_sequence_helpers.sql new file mode 100644 index 000000000..ae3797c45 --- /dev/null +++ b/awa/tests/fixtures/v0613_lane_sequence_helpers.sql @@ -0,0 +1,195 @@ +-- The lane-sequence helpers as awa 0.6.13 installs them, before lane +-- provisioning functions existed. Installed under a separate name so a test +-- can put a schema back in its 0.6.13 shape. +CREATE OR REPLACE FUNCTION awa.v0613_refresh_lane_sequence_helpers(p_schema TEXT) +RETURNS VOID +LANGUAGE plpgsql +SECURITY INVOKER +SET search_path = pg_catalog +AS $refresh$ +BEGIN + IF p_schema IS NULL OR p_schema !~ '^[a-z_][a-z0-9_]*$' THEN + RAISE EXCEPTION 'refresh_lane_sequence_helpers: schema name must match ^[a-z_][a-z0-9_]*$; got %', + p_schema + USING ERRCODE = '22023'; + END IF; + + 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/lane_ownership_roles_test.rs b/awa/tests/lane_ownership_roles_test.rs new file mode 100644 index 000000000..026511f77 --- /dev/null +++ b/awa/tests/lane_ownership_roles_test.rs @@ -0,0 +1,857 @@ +//! 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) -> String { + text.into() +} + +/// 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, + /// awa 0.6.13: every other schema patch, and the lane helpers before + /// lane provisioning. + 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?; + for (_, _, statement) in migrations::migration_sql_range(0, migrations::CURRENT_VERSION) { + migrate_as_owner.execute(sqlx::raw_sql(&statement)).await?; + } + for (name, _, statement) in migrations::schema_patch_sql(migrations::SCHEMA_PATCHES) { + if schema == Schema::BeforeProvisioning && name == "definer_lane_provisioning" { + continue; + } + migrate_as_owner.execute(sqlx::raw_sql(&statement)).await?; + } + if schema == Schema::BeforeProvisioning { + // Every patch carries the current v023 definitions, so put the + // 0.6.13 helpers back and drop what 0.6.13 did not have. + migrate_as_owner + .execute(sqlx::raw_sql(include_str!( + "fixtures/v0613_lane_sequence_helpers.sql" + ))) + .await?; + migrate_as_owner + .execute(sqlx::raw_sql( + r#" + SELECT awa.v0613_refresh_lane_sequence_helpers('awa'); + DROP FUNCTION awa.v0613_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}" + ); + } + + // `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?; + let report = tokio::time::timeout( + Duration::from_secs(120), + migrations::run_with_report(&migrate_pool), + ) + .await??; + assert!(report.complete(), "{:?}", report.patches); + assert!(migrations::pending_schema_patches(&cluster.owner) + .await? + .is_empty()); + let outcomes: Vec = report + .lane_sequences + .iter() + .map(|repair| repair.outcome) + .collect(); + assert_eq!(outcomes.len(), 6, "{:?}", report.lane_sequences); + assert_eq!( + outcomes + .iter() + .filter(|outcome| **outcome == LaneSequenceOutcome::Repaired) + .count(), + 4 + ); + 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?; + + // 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 patch completes while the lock is held, and its report names the + // held sequence. + let report = tokio::time::timeout( + Duration::from_secs(120), + migrations::run_with_report(&migrate_pool), + ) + .await??; + assert!(report.complete(), "{:?}", report.patches); + assert_eq!( + report.lane_sequences.len(), + 4, + "{:?}", + report.lane_sequences + ); + for repair in &report.lane_sequences { + let expected = if repair.sequence_name == held_sequence { + LaneSequenceOutcome::Busy + } else { + LaneSequenceOutcome::Repaired + }; + assert_eq!(repair.outcome, expected, "{repair:?}"); + } + 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/patch_upgrade_rehearsal_test.rs b/awa/tests/patch_upgrade_rehearsal_test.rs index 4bd2a761a..aef1255af 100644 --- a/awa/tests/patch_upgrade_rehearsal_test.rs +++ b/awa/tests/patch_upgrade_rehearsal_test.rs @@ -1108,12 +1108,23 @@ async fn released_0_6_9_patch_upgrade_and_rollback() { killed_sources.insert(worker_c.kill().await); killed_sources.insert(worker_d.kill().await); - let worker_e = spawn_released_worker(&fixture.runtime_url, &log, &queues).await; - let worker_f = spawn_released_worker(&fixture.runtime_url, &log, &queues).await; + let restricted_queues = [ + MAIN, + BANKED, + NEW_AFTER_PATCH, + NEW_AFTER_ROLLBACK, + PROVISIONED, + UNPROVISIONED, + ]; + let worker_e = spawn_released_worker(&fixture.runtime_url, &log, &restricted_queues).await; + let worker_f = spawn_released_worker(&fixture.runtime_url, &log, &restricted_queues).await; + // A released producer without CREATE creates an unprepared queue's lane + // through the schema owner's provisioning function. for (queue, seq_start) in [ (MAIN, 9_000), (PROVISIONED, 10_000), (NEW_AFTER_ROLLBACK, 11_000), + (UNPROVISIONED, 12_000), ] { released_produce( &fixture.runtime_url, @@ -1130,21 +1141,22 @@ async fn released_0_6_9_patch_upgrade_and_rollback() { .await .unwrap_or_else(|error| panic!("restricted released enqueue on {queue}: {error}")); } - let refused = released_produce( - &fixture.runtime_url, - &log, - ProduceSpec { - queue: UNPROVISIONED, - seq_start: 12_000, - count: 1, - modes: "ok", - delay_s: 0.0, - pause_ms: 0, - }, - ) + let new_lane_owners: Vec = sqlx::query_scalar(&format!( + "SELECT sequence.relowner = namespace.nspowner \ + FROM {SCHEMA}.queue_enqueue_heads AS heads \ + JOIN pg_namespace AS namespace ON namespace.nspname = '{SCHEMA}' \ + JOIN pg_class AS sequence \ + ON sequence.relnamespace = namespace.oid AND sequence.relname = heads.seq_name \ + WHERE heads.queue = $1" + )) + .bind(UNPROVISIONED) + .fetch_all(&fixture.admin) .await - .expect_err("an unprovisioned queue needs provisioning once CREATE is revoked"); - assert!(refused.contains("has not been provisioned"), "{refused}"); + .unwrap(); + assert!( + !new_lane_owners.is_empty() && new_lane_owners.iter().all(|owned| *owned), + "the new lane's sequences belong to the schema owner: {new_lane_owners:?}" + ); wait_for_drain(&log, Duration::from_secs(120)).await; let stats = cli_ok(&released_cli, &fixture.runtime_url, &["queue", "stats"]).await; assert!(stats.contains(PROVISIONED), "{stats}"); diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index 376548e7f..288f89cfb 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -3,10 +3,12 @@ //! These tests exercise the full dispatcher/worker/maintenance wiring with the //! queue_storage backend enabled. +mod ci_timing; + use awa::model::{ admin, batch_operations, insert, migrations, storage, AwaError, BatchOperationFilter, - BatchOperationSpec, PruneOutcome, QueueStorage, QueueStorageConfig, RotateOutcome, SkipReason, - SubmitBatchOperation, + BatchOperationSpec, PromotionLane, PruneOutcome, QueueStorage, QueueStorageConfig, + RotateOutcome, SkipReason, SubmitBatchOperation, }; use awa::{ Client, InsertOpts, JobArgs, JobContext, JobError, JobResult, JobRow, JobState, QueueConfig, @@ -2009,6 +2011,231 @@ impl Worker for MultiClientTrackingWorker { } } +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct GracefulRestartJob { + id: i64, +} + +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct GracefulRestartFollowUp { + id: i64, +} + +/// The first execution reports its attempt and returns `Completed` once +/// shutdown signals cancellation; later executions complete at once. +#[derive(Clone, Default)] +struct GracefulRestartWorker { + executions: Arc, + first_attempt: Arc>>, + entered: Arc, +} + +#[async_trait::async_trait] +impl Worker for GracefulRestartWorker { + fn kind(&self) -> &'static str { + "graceful_restart_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + if self.executions.fetch_add(1, Ordering::SeqCst) > 0 { + return Ok(JobResult::Completed); + } + *self.first_attempt.lock().expect("first attempt") = Some((ctx.job.id, ctx.job.run_lease)); + self.entered.notify_one(); + while !ctx.is_cancelled() { + tokio::time::sleep(Duration::from_millis(5)).await; + } + Ok(JobResult::Completed) + } +} + +/// A runtime on its own tokio runtime and pool, so stopping it after a +/// graceful shutdown drops whatever the client left running, as a process +/// exit does. +struct ProcessHost { + runtime: tokio::runtime::Runtime, + client: Arc, +} + +impl ProcessHost { + async fn start( + options: sqlx::postgres::PgConnectOptions, + build: impl FnOnce(sqlx::PgPool) -> Client + Send + 'static, + ) -> Self { + tokio::task::spawn_blocking(move || { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .expect("build process runtime"); + let client = runtime.block_on(async move { + let pool = PgPoolOptions::new() + .max_connections(6) + .connect_with(options) + .await + .expect("connect process pool"); + let client = build(pool); + client.start().await.expect("start process client"); + client + }); + Self { + runtime, + client: Arc::new(client), + } + }) + .await + .expect("start process host") + } + + async fn exit_gracefully(self, timeout: Duration) { + let client = self.client.clone(); + self.runtime + .spawn(async move { client.shutdown(timeout).await }) + .await + .expect("graceful shutdown"); + tokio::task::spawn_blocking(move || self.runtime.shutdown_background()) + .await + .expect("stop process runtime"); + } +} + +/// A graceful shutdown must finalize every attempt whose handler returned +/// during the drain. Finalization used to run detached from the drain, so a +/// process that exited after `shutdown` returned could drop it: the attempt +/// stayed open with no heartbeat and stale-heartbeat rescue ran the completed +/// job again. A held receipt-attempt lock keeps this attempt's finalize in +/// flight while shutdown drains, which the detached version did not wait for. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_graceful_shutdown_finalizes_returned_attempts() { + // Each host below opens its own pool of 6, and only one runs at a time. + let (_db_guard, pool) = setup_pool(10).await; + let queue = "qs_graceful_restart"; + let store_config = QueueStorageConfig { + schema: "awa_qs_graceful_restart".to_string(), + ..Default::default() + }; + let store = create_store_with_config(&pool, store_config.clone()).await; + let job_id = enqueue_job( + &pool, + &store, + &GracefulRestartJob { id: 1 }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + + let worker = GracefulRestartWorker::default(); + let options = (*pool.connect_options()).clone(); + let build = |worker: GracefulRestartWorker, store_config: QueueStorageConfig| { + move |pool: sqlx::PgPool| { + Client::builder(pool) + .queue( + queue, + QueueConfig { + max_workers: 2, + poll_interval: Duration::from_millis(25), + deadline_duration: Duration::ZERO, + ..QueueConfig::default() + }, + ) + .queue_storage( + store_config, + Duration::from_secs(60), + Duration::from_millis(50), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .register_worker(worker) + // The follow-up's queue is not served; it only routes the + // completion through a transaction of its own. + .on_completed_enqueue::( + |args, _job| GracefulRestartFollowUp { id: args.id }, + ) + .heartbeat_interval(Duration::from_millis(100)) + .heartbeat_staleness(ci_timing::scaled_staleness(Duration::from_secs(1))) + .heartbeat_rescue_interval(Duration::from_millis(100)) + .promote_interval(Duration::from_millis(25)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .build() + .expect("build graceful restart client") + } + }; + + let host = + ProcessHost::start(options.clone(), build(worker.clone(), store_config.clone())).await; + tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(10)), + worker.entered.notified(), + ) + .await + .expect("the job should start"); + let (attempt_job_id, run_lease) = worker + .first_attempt + .lock() + .expect("first attempt") + .expect("first attempt recorded"); + assert_eq!(attempt_job_id, job_id); + + let mut lock_holder = pool.acquire().await.expect("lock connection"); + sqlx::query( + "SELECT pg_catalog.pg_advisory_lock(pg_catalog.hashtextextended( + format('awa.receipt.complete:%s:%s', $1::bigint, $2::bigint), 0))", + ) + .bind(job_id) + .bind(run_lease) + .execute(&mut *lock_holder) + .await + .expect("hold the attempt's finalize lock"); + + let exit = + tokio::spawn(host.exit_gracefully(ci_timing::scaled_timeout(Duration::from_secs(10)))); + tokio::time::sleep(Duration::from_millis(500)).await; + sqlx::query( + "SELECT pg_catalog.pg_advisory_unlock(pg_catalog.hashtextextended( + format('awa.receipt.complete:%s:%s', $1::bigint, $2::bigint), 0))", + ) + .bind(job_id) + .bind(run_lease) + .execute(&mut *lock_holder) + .await + .expect("release the attempt's finalize lock"); + drop(lock_holder); + exit.await.expect("graceful exit"); + + let after_exit = store + .load_job(&pool, job_id) + .await + .expect("load job") + .expect("job row"); + + // A later runtime rescues anything the exit left open. + let rescuer = ProcessHost::start(options, build(worker.clone(), store_config)).await; + let finished = wait_for_job_state( + &store, + &pool, + job_id, + &[JobState::Completed], + ci_timing::scaled_timeout(Duration::from_secs(15)), + ) + .await; + rescuer + .exit_gracefully(ci_timing::scaled_timeout(Duration::from_secs(10))) + .await; + + assert_eq!( + (after_exit.state, after_exit.run_lease), + (JobState::Completed, run_lease), + "graceful shutdown returned before the completed attempt was finalized" + ); + assert_eq!( + (worker.executions.load(Ordering::SeqCst), finished.run_lease), + (1, run_lease), + "the job ran again after a graceful shutdown" + ); +} + /// ADR-023 claim-ring control-plane smoke test. Exercises /// `rotate_claims`, the busy-check, and `prune_oldest_claims` on an /// empty schema: rotate cycles through every slot, prune is a noop @@ -7802,9 +8029,10 @@ fn branch_overruns() -> Arc>> { /// snoozed and retry-after jobs runnable, so overrunning promotion bodies must /// never arm a cooldown on it. /// -/// The first three promotion bodies (two `DELETE`s each) take 1.2s against a -/// 25ms promote interval, so the maintenance branch tracker flags promotion as -/// overrunning. The test asserts on that overrun warning rather than on job +/// Every promotion body runs the canonical-plane promotion `DELETE` once per +/// state, even with nothing due. A statement trigger slows the first six, so +/// the first three bodies take 1.2s against a 25ms promote interval and the +/// maintenance branch tracker flags promotion as overrunning. The test asserts on that overrun warning rather than on job /// latency: the promotion warning must report no cooldown, while another /// branch in the same run (cleanup, whose bodies exceed its 1ms interval) /// does report one, so the check cannot pass vacuously. Jobs that fall due @@ -7831,7 +8059,7 @@ async fn test_queue_storage_runtime_promotion_not_suppressed_after_overrunning_b END $$" ), format!( - "CREATE TRIGGER slow_first_promote_deletes BEFORE DELETE ON {schema}.deferred_jobs \ + "CREATE TRIGGER slow_first_promote_deletes BEFORE DELETE ON awa.scheduled_jobs \ FOR EACH STATEMENT EXECUTE FUNCTION {schema}.slow_first_promote_deletes()" ), ] { @@ -7965,6 +8193,436 @@ async fn test_queue_storage_runtime_promotion_not_suppressed_after_overrunning_b client.shutdown(Duration::from_secs(5)).await; } +/// Advisory lock key of a lane's enqueue reservation, as +/// `reserve_enqueue_seq` takes it. The lane must already exist. +async fn enqueue_lane_lock_key( + pool: &sqlx::PgPool, + store: &QueueStorage, + queue: &str, + priority: i16, +) -> i64 { + sqlx::query_scalar(&format!( + "SELECT hashtextextended(format('%I.%I', $1::text, seq_name), 0) \ + FROM {}.queue_enqueue_heads \ + WHERE queue = $2 AND priority = $3 AND enqueue_shard = 0", + store.schema() + )) + .bind(store.schema()) + .bind(queue) + .bind(priority) + .fetch_one(pool) + .await + .expect("derive enqueue lane lock key") +} + +/// A promotion batch that meets a lane whose enqueue lock is held promotes +/// the other lanes, reports the held lane as contended after its lock timeout +/// and leaves its rows deferred. Promoted later, the lane's rows take lane +/// sequence numbers in due order and are all claimable. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_promotion_leaves_a_contended_lane_deferred_and_in_order() { + const PRIORITY: i16 = 2; + let (_db_guard, pool) = setup_pool(6).await; + let hot = "qs_promote_contended_hot"; + let other = "qs_promote_contended_other"; + let store = create_store(&pool, "awa_qs_promote_contended_lane").await; + let opts = |queue: &str, run_at: Option>| InsertOpts { + queue: queue.to_string(), + priority: PRIORITY, + run_at, + ..Default::default() + }; + + let warm_job_id = enqueue_job(&pool, &store, &CompleteJob { id: 0 }, opts(hot, None)).await; + // Later job ids fall due earlier, so due order differs from id order. + let due_base = Utc::now() - chrono::Duration::seconds(60); + let mut hot_due = Vec::new(); + for n in 0..6 { + let run_at = due_base - chrono::Duration::seconds(n); + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: n }, + opts(hot, Some(run_at)), + ) + .await; + hot_due.push((run_at, job_id)); + } + hot_due.sort(); + let hot_in_due_order: Vec = hot_due.iter().map(|(_, job_id)| *job_id).collect(); + let other_job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 100 }, + opts(other, Some(due_base)), + ) + .await; + + let hot_lane = PromotionLane { + queue: hot.to_string(), + priority: PRIORITY, + }; + let lock_key = enqueue_lane_lock_key(&pool, &store, hot, PRIORITY).await; + let mut holder = pool.begin().await.expect("begin lane holder"); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(lock_key) + .execute(holder.as_mut()) + .await + .expect("hold the hot lane's enqueue lock"); + + let lock_timeout = Duration::from_millis(100); + let started = Instant::now(); + let outcome = store + .promote_due_lanes(&pool, JobState::Scheduled, 100, Some(lock_timeout), &[]) + .await + .expect("promote with the hot lane held"); + let elapsed = started.elapsed(); + assert_eq!(outcome.promoted, 1, "{outcome:?}"); + assert_eq!(outcome.selected, 7, "{outcome:?}"); + assert_eq!(outcome.contended, vec![hot_lane.clone()]); + assert!( + elapsed >= lock_timeout && elapsed < scaled_timeout(Duration::from_secs(1)), + "the held lane should cost one lock timeout, took {elapsed:?}" + ); + let other_job = store + .load_job(&pool, other_job_id) + .await + .expect("load other job") + .expect("other job exists"); + assert_eq!(other_job.state, JobState::Available); + for job_id in &hot_in_due_order { + let job = store + .load_job(&pool, *job_id) + .await + .expect("load hot job") + .expect("hot job exists"); + assert_eq!(job.state, JobState::Scheduled, "job {job_id}"); + } + + let skipped = store + .promote_due_lanes( + &pool, + JobState::Scheduled, + 100, + Some(lock_timeout), + std::slice::from_ref(&hot_lane), + ) + .await + .expect("promote with the hot lane skipped"); + assert_eq!(skipped, Default::default()); + + holder.commit().await.expect("release the hot lane"); + let mut batches = 0; + loop { + let outcome = store + .promote_due_lane(&pool, JobState::Scheduled, &hot_lane, 4, None) + .await + .expect("promote the hot lane"); + assert!(outcome.contended.is_empty(), "{outcome:?}"); + if outcome.promoted == 0 { + break; + } + batches += 1; + } + assert_eq!(batches, 2); + + let in_lane_order: Vec = sqlx::query_scalar(&format!( + "SELECT job_id FROM {}.ready_entries WHERE job_id = ANY($1) ORDER BY lane_seq", + store.schema() + )) + .bind(&hot_in_due_order) + .fetch_all(&pool) + .await + .expect("read hot lane order"); + assert_eq!(in_lane_order, hot_in_due_order); + + let mut claimed: Vec = store + .claim_runtime_batch(&pool, hot, 100, Duration::from_secs(60)) + .await + .expect("claim the hot lane") + .into_iter() + .map(|claimed| claimed.job.id) + .collect(); + claimed.sort_unstable(); + let mut expected = hot_in_due_order.clone(); + expected.push(warm_job_id); + expected.sort_unstable(); + assert_eq!(claimed, expected); +} + +/// The maintenance leader must not wait on a lane whose enqueue lock never +/// comes free. Two sessions take turns holding the hot lane's lock for +/// `HOLD` each, so a lock request queued behind them waits up to two holds. +/// A due job on another queue must still run, and an abandoned attempt must +/// still be rescued, well within one hold; the hot queue's due jobs must +/// still be promoted while the lane stays busy. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_runtime_busy_lane_holds_back_neither_promotion_nor_rescue() { + const HOLD: Duration = Duration::from_secs(5); + const PRIORITY: i16 = 2; + let (_db_guard, pool) = setup_pool(12).await; + let hot = "qs_busy_lane_hot"; + let other = "qs_busy_lane_other"; + let abandoned = "qs_busy_lane_abandoned"; + let schema = "awa_qs_busy_lane_leader"; + let store = create_store(&pool, schema).await; + let opts = |queue: &str, run_at: Option>| InsertOpts { + queue: queue.to_string(), + priority: PRIORITY, + run_at, + ..Default::default() + }; + + enqueue_job(&pool, &store, &CompleteJob { id: 0 }, opts(hot, None)).await; + let lock_key = enqueue_lane_lock_key(&pool, &store, hot, PRIORITY).await; + + let stop = Arc::new(AtomicBool::new(false)); + let mut holders = Vec::new(); + for _ in 0..2 { + let (held_tx, held_rx) = tokio::sync::oneshot::channel::<()>(); + let pool = pool.clone(); + let stop = Arc::clone(&stop); + holders.push(tokio::spawn(async move { + let mut held_tx = Some(held_tx); + let mut holds = 0u32; + while !stop.load(Ordering::SeqCst) { + let mut tx = pool.begin().await.expect("begin lane holder"); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(lock_key) + .execute(tx.as_mut()) + .await + .expect("hold the hot lane's enqueue lock"); + holds += 1; + if let Some(held_tx) = held_tx.take() { + let _ = held_tx.send(()); + } + tokio::time::sleep(HOLD).await; + tx.commit().await.expect("release the hot lane"); + } + holds + })); + // The first holder takes the lock; wait for it so the second queues + // behind it rather than racing it. + if holders.len() == 1 { + held_rx.await.expect("first holder took the lane"); + } + } + + let due = Utc::now() - chrono::Duration::seconds(1); + let mut hot_jobs = Vec::new(); + for n in 1..=3 { + hot_jobs + .push(enqueue_job(&pool, &store, &CompleteJob { id: n }, opts(hot, Some(due))).await); + } + let other_job = enqueue_job( + &pool, + &store, + &CompleteJob { id: 10 }, + opts(other, Some(due)), + ) + .await; + let abandoned_job = enqueue_job( + &pool, + &store, + &CompleteJob { id: 20 }, + opts(abandoned, None), + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, abandoned, 1, Duration::from_secs(3_600)) + .await + .expect("claim the abandoned attempt"); + assert_eq!(claimed.len(), 1); + + let heartbeat_staleness = Duration::from_millis(500); + let worker_queue = || QueueConfig { + max_workers: 4, + poll_interval: Duration::from_millis(25), + ..QueueConfig::default() + }; + let client = Client::builder(pool.clone()) + .queue(hot, worker_queue()) + .queue(other, worker_queue()) + .queue_storage( + QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + lease_claim_receipts: false, + ..Default::default() + }, + Duration::from_secs(60), + Duration::from_secs(60), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .register_worker(CompleteWorker) + .promote_interval(Duration::from_millis(50)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .heartbeat_rescue_interval(Duration::from_millis(100)) + .heartbeat_staleness(heartbeat_staleness) + .deadline_rescue_interval(Duration::from_secs(10)) + .callback_rescue_interval(Duration::from_secs(10)) + .build() + .expect("Failed to build busy-lane client"); + client + .start() + .await + .expect("Failed to start busy-lane client"); + + let bound = scaled_timeout(Duration::from_secs(2)); + wait_for_job_state(&store, &pool, other_job, &[JobState::Completed], bound).await; + wait_for_job_state( + &store, + &pool, + abandoned_job, + &[JobState::Retryable, JobState::Available], + bound + heartbeat_staleness, + ) + .await; + + for job_id in &hot_jobs { + wait_for_job_state( + &store, + &pool, + *job_id, + &[JobState::Completed], + scaled_timeout(HOLD * 3), + ) + .await; + } + assert!( + holders.iter().all(|holder| !holder.is_finished()), + "the hot lane must stay busy until its jobs are promoted" + ); + + stop.store(true, Ordering::SeqCst); + for holder in holders { + assert!(holder.await.expect("lane holder") >= 1); + } + client.shutdown(Duration::from_secs(5)).await; +} + +/// Rows another transaction holds locked, as a lane promotion still waiting +/// under a former leader does, must not fill every batch of a promotion tick. +/// With more of the oldest due rows locked than one batch selects, a due job +/// on another queue must still be promoted and run promptly, and the locked +/// rows must be promoted once released. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_runtime_locked_due_rows_do_not_starve_promotion() { + const LOCKED: usize = 4_096 + 64; + let (_db_guard, pool) = setup_pool(10).await; + let locked_queue = "qs_locked_due_rows"; + let other = "qs_locked_due_other"; + let schema = "awa_qs_locked_due_rows"; + let store = create_store(&pool, schema).await; + + let oldest = Utc::now() - chrono::Duration::seconds(120); + let params: Vec<_> = (0..LOCKED) + .map(|n| { + insert::params_with( + &CompleteJob { id: n as i64 }, + InsertOpts { + queue: locked_queue.to_string(), + run_at: Some(oldest), + ..Default::default() + }, + ) + .expect("build locked-row params") + }) + .collect(); + for chunk in params.chunks(2_000) { + store + .enqueue_params_batch(&pool, chunk) + .await + .expect("enqueue rows to lock"); + } + let mut holder = pool.begin().await.expect("begin row holder"); + let locked: Vec = sqlx::query_scalar(&format!( + "SELECT job_id FROM {schema}.deferred_jobs WHERE queue = $1 FOR UPDATE" + )) + .bind(locked_queue) + .fetch_all(holder.as_mut()) + .await + .expect("lock the oldest due rows"); + assert_eq!(locked.len(), LOCKED); + + let other_job = enqueue_job( + &pool, + &store, + &CompleteJob { id: -1 }, + InsertOpts { + queue: other.to_string(), + run_at: Some(Utc::now() - chrono::Duration::seconds(1)), + ..Default::default() + }, + ) + .await; + + let client = Client::builder(pool.clone()) + .queue( + other, + QueueConfig { + max_workers: 4, + poll_interval: Duration::from_millis(25), + ..QueueConfig::default() + }, + ) + .queue_storage( + QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + lease_claim_receipts: false, + ..Default::default() + }, + Duration::from_secs(60), + Duration::from_secs(60), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .register_worker(CompleteWorker) + .promote_interval(Duration::from_millis(50)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .build() + .expect("Failed to build locked-rows client"); + client + .start() + .await + .expect("Failed to start locked-rows client"); + + wait_for_job_state( + &store, + &pool, + other_job, + &[JobState::Completed], + scaled_timeout(Duration::from_secs(2)), + ) + .await; + + holder.rollback().await.expect("release the locked rows"); + let deadline = Instant::now() + scaled_timeout(Duration::from_secs(10)); + loop { + let remaining: i64 = sqlx::query_scalar(&format!( + "SELECT count(*) FROM {schema}.deferred_jobs WHERE queue = $1" + )) + .bind(locked_queue) + .fetch_one(&pool) + .await + .expect("count deferred rows"); + if remaining == 0 { + break; + } + assert!( + Instant::now() < deadline, + "{remaining} released rows were never promoted" + ); + tokio::time::sleep(Duration::from_millis(50)).await; + } + + client.shutdown(Duration::from_secs(5)).await; +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn test_queue_storage_runtime_stale_heartbeat_rescue() { let (_db_guard, pool) = setup_pool(10).await; @@ -15547,3 +16205,114 @@ async fn test_aged_materialized_lease_is_visible_to_admin_reads() { .expect("cancel by unique key finds the aged running job"); assert_eq!(cancelled.id, job_id); } + +/// A lane move leaves immutable source rows behind. Wrap the ring so the +/// destination is in an earlier partition: point reads must use live evidence, +/// rather than depending on which equal-timestamp ready row the plan emits last. +async fn ready_point_read_wrap_fixture() -> (TestDbGuard, sqlx::PgPool, QueueStorage, i64) { + let (guard, pool) = setup_pool(4).await; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: "awa_qs_point_read_wrap".to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + claim_slot_count: 2, + ..Default::default() + }, + ) + .await; + // A fresh empty ring can be positioned just before wrap without retaining + // unrelated jobs. This is the state reached after three rotations/prunes. + let schema = store.schema(); + sqlx::query(&format!( + "UPDATE {schema}.queue_ring_state SET current_slot = 3, generation = 3 WHERE singleton" + )) + .execute(&pool) + .await + .expect("position empty ring before wrap"); + sqlx::query(&format!( + "UPDATE {schema}.queue_ring_slots SET generation = 3 WHERE slot = 3" + )) + .execute(&pool) + .await + .expect("set current slot generation"); + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 1 }, + InsertOpts { + queue: "point_read_source".to_string(), + priority: 4, + ..Default::default() + }, + ) + .await; + assert!(matches!( + store.rotate(&pool).await.expect("wrap to slot zero"), + RotateOutcome::Rotated { slot: 0, .. } + )); + (guard, pool, store, job_id) +} + +async fn assert_live_ready_row( + pool: &sqlx::PgPool, + store: &QueueStorage, + job_id: i64, + expected_queue: &str, + expected_priority: i16, +) { + let schema = store.schema(); + let live: Vec<(String, i16, i32)> = sqlx::query_as(&format!( + "SELECT ready.queue, ready.priority, ready.ready_slot + FROM {schema}.ready_entries AS ready + WHERE ready.job_id = $1 + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + )" + )) + .bind(job_id) + .fetch_all(pool) + .await + .expect("read live ready evidence"); + assert_eq!( + live, + vec![(expected_queue.to_string(), expected_priority, 0)] + ); +} + +#[tokio::test] +async fn test_point_read_uses_live_ready_priority_after_ring_wrap() { + let (_guard, pool, store, job_id) = ready_point_read_wrap_fixture().await; + assert!(store + .set_priority(&pool, job_id, 1) + .await + .expect("set priority")); + assert_live_ready_row(&pool, &store, job_id, "point_read_source", 1).await; + let job = admin::get_job(&pool, job_id).await.expect("read moved job"); + assert_eq!( + job.priority, 1, + "point read must ignore the tombstoned source" + ); +} + +#[tokio::test] +async fn test_point_read_uses_live_ready_queue_after_ring_wrap() { + let (_guard, pool, store, job_id) = ready_point_read_wrap_fixture().await; + assert!(store + .move_queue(&pool, job_id, "point_read_destination", None) + .await + .expect("move queue")); + assert_live_ready_row(&pool, &store, job_id, "point_read_destination", 4).await; + let job = admin::get_job(&pool, job_id).await.expect("read moved job"); + assert_eq!( + job.queue, "point_read_destination", + "point read must ignore the tombstoned source" + ); +} diff --git a/docs/architecture.md b/docs/architecture.md index 06b967bc3..c4437b0e7 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -150,7 +150,7 @@ sequenceDiagram Enqueue is transactional: if the producer's outer transaction rolls back, the job never becomes visible. Immediate jobs reserve a lane sequence range under a transaction-scoped lane lock, append to `ready_entries_*`, and append a compact `ready_segments` range for the committed lane sequence range; future scheduled or retryable jobs append to `deferred_jobs`. The lane lock is held until commit or rollback so a later producer cannot make lane `N+1` visible while lane `N` can still commit. -`deferred_jobs` is deliberately outside the claim path. The maintenance leader promotes due `scheduled` and `retryable` rows into `ready_entries_*` in batches and notifies the target queues. Cron schedules are just producers for ordinary jobs: when a schedule fires, the cron transaction records the fire and enqueues the job atomically, using `ready_entries_*` for immediate fires or `deferred_jobs` when the enqueue carries a future `run_at`. +`deferred_jobs` is deliberately outside the claim path. The maintenance leader promotes due `scheduled` and `retryable` rows into `ready_entries_*` in batches and notifies the target queues. Re-enqueueing a row takes its lane lock, which producers hold until they commit, so each `(queue, priority)` lane's share of a batch is promoted in its own transaction under a short `lock_timeout`. A lane that times out is handed to a dedicated waiter task, which queues for the lock behind the producers already waiting, so a busy lane is still promoted in turn without holding up other lanes or the leader's rescue, rotation and prune. Cron schedules are just producers for ordinary jobs: when a schedule fires, the cron transaction records the fire and enqueues the job atomically, using `ready_entries_*` for immediate fires or `deferred_jobs` when the enqueue carries a future `run_at`. There are two COPY-shaped producer paths: diff --git a/docs/configuration.md b/docs/configuration.md index 868dfb3cf..22d7e7b41 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -245,6 +245,7 @@ Retry backoff and `Snooze` use the same mechanism: the current attempt closes, a Operational knobs: - `promote_interval` controls how often maintenance checks for due deferred rows. +- A queue whose enqueue lane is kept busy by producers does not delay promotion of other queues or the leader's rescue: promotion waits 50 ms for a lane's lock, then hands the lane to a dedicated waiter that is served in turn. Each waiter holds one connection from the worker's pool while it waits, on top of the leader's own. The leader runs at most one waiter per eight pool connections, never fewer than one or more than four, so a 10-connection pool gives up at most one extra connection to a busy lane. `awa_maintenance_promote_lane_contention_total` counts those hand-offs by queue. A steady `waiter_limit` rate means more lanes are saturated than the waiters cover, and a larger pool allows more waiters. `waiter_timed_out` means a transaction held a lane lock for longer than 30 s. - Per-job `run_at` controls the due time. - Periodic schedules are declared in worker code with `periodic()` and managed through the CLI/UI `cron` surface. diff --git a/docs/grafana/README.md b/docs/grafana/README.md index 823102462..81dcb0f19 100644 --- a/docs/grafana/README.md +++ b/docs/grafana/README.md @@ -156,6 +156,7 @@ All metrics use the `awa` OTel meter name and are exported via OTLP to your conf | `awa.maintenance.promote_batches` | Counter | state | Promotion batches | | `awa.maintenance.promote_batch_size` | Histogram | state | Jobs promoted | | `awa.maintenance.promote_duration` | Histogram (s) | state | Promotion time | +| `awa.maintenance.promote.lane_contention` | Counter | state, queue, outcome | Promotion lanes whose lock was busy: `handed_off` to a dedicated waiter, `waiter_limit` when every waiter was busy, `waiter_timed_out` when a waiter's 30 s wait ran out | | `awa.maintenance.rescues` | Counter | rescue_kind | Jobs rescued | ### Heartbeat diff --git a/docs/security.md b/docs/security.md index a146635bb..2b9848ec7 100644 --- a/docs/security.md +++ b/docs/security.md @@ -174,23 +174,52 @@ ALTER DEFAULT PRIVILEGES FOR ROLE awa_migrator IN SCHEMA my_qs_schema The default `awa` queue-storage substrate is migrated by `awa migrate`. Custom queue-storage schemas are migrated by `awa storage prepare-queue-storage-schema`, which should also run as the migrator role and creates objects owned by `awa_migrator` / `awa_owner`. The runtime role needs read/write/execute privileges plus `TRUNCATE` for ring-partition rotation; it never needs DDL. -### Provision queues before starting restricted runtimes +### Lane sequences -Lane sequences are created lazily by the first producer or worker that touches a queue, which needs schema `CREATE`. A runtime role without it needs each queue provisioned in advance, and needs the `prepared_lane_sequences` schema patch: before it, the lane helpers issue `CREATE SEQUENCE IF NOT EXISTS` even for sequences that already exist, which requires schema `CREATE` regardless. Apply the patch with `awa migrate` (it applies on any v040 database), or run the exported `R__prepared_lane_sequences.sql` from `awa migrate --sql` / `--extract-to` / `awa.schema_patches()` after `V40`. +Each queue lane, `(queue, priority, enqueue_shard)`, has an enqueue and a claim sequence. Once the `definer_lane_provisioning` schema patch is applied, 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 role therefore needs neither schema `CREATE` nor default privileges to start using a new queue, priority or shard. Apply the patch with `awa migrate` as the schema owner (`PGOPTIONS='-c role=awa_owner'`), or run the exported `R__definer_lane_provisioning.sql` after `V40`. -Then provision every queue as the migrator before starting its producers or workers: +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. The table grants above are enough. + +`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 the patch 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 ``` -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`, and match `--queue-stripe-count` to the runtime's queue-storage configuration. With `PartitionedQueue`, provision each physical queue. Use `--schema` for a custom queue-storage schema. The Rust equivalent is `QueueStorage::prepare_queue`. +```sql +SELECT * FROM awa.lane_sequence_problems('awa'); +``` -Lane cursors need `UPDATE` on their sequences for `setval`, in addition to `USAGE` and `SELECT`; use the sequence grants above. A restricted runtime that reaches an unprovisioned lane fails with SQLSTATE `42501` and a provisioning hint. Schema owners keep lazy creation, so single-role installations need no extra step. +The command exits 2 while any remain. + +`awa migrate` repairs them once the patch 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 patch or `awa migrate`. Run the repair again at any time; it is idempotent: + +```bash +awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage lane-sequences --repair +``` -On a database where a runtime role with `CREATE` has been creating lanes (any 0.6.9 deployment that ran without schema ownership), the runtime owns those lane sequences. `GRANT ... ON ALL SEQUENCES` run by the owner fails on them, so grant per sequence the owner owns, verify, and only then revoke `CREATE`. [Upgrading from 0.6.9 to 0.6.10](upgrade-0.6.9-to-0.6.10.md#optional-run-the-runtime-role-without-schema-create) has the full procedure and the SQL. +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 +``` ### 5. Configure your processes @@ -214,7 +243,7 @@ Since these triggers run with the caller's privileges, the runtime role needs wr The maintenance leader also calls `awa.refresh_admin_metadata()` as a full reconciliation safety net. It runs with invoker privileges and needs DML on the admin metadata tables. `TRUNCATE` in the grant is for queue-storage segment pruning, which truncates retired partitions. -Schema patches (`awa migrate --sql` prints them after the versioned migrations) replace function bodies and create tables under the migrator role, like migrations do; they never change grants. Privileges managed with `GRANT ... ON ALL TABLES` plus the default privileges above cover patched tables without further action. +Schema patches (`awa migrate --sql` prints them after the versioned migrations) replace functions and may create tables or change function ownership and grants. The `definer_lane_provisioning` patch installs provisioning functions owned by the schema owner and grants them `PUBLIC EXECUTE`. After the patch commits, `awa migrate` attempts to repair existing lane-sequence ownership and grants; review its warnings for sequences that still need repair. The default privileges above cover new tables created by the migration role. The runtime also directly upserts into the descriptor catalogs on startup and on each runtime snapshot tick: diff --git a/docs/test-plan.md b/docs/test-plan.md index 435b18861..618ff69a2 100644 --- a/docs/test-plan.md +++ b/docs/test-plan.md @@ -32,6 +32,7 @@ Tests run against real Postgres 15+ (not managed services) using a dedicated tes | B5 | Deadline rescue signals ctx.is_cancelled() | ✓ | | | B6 | UniqueConflict.constraint field | ✓ | | | B7 | Admin cancel signals in-flight handler cancellation | ✓ | ✓ | +| B8 | Shutdown finalizes attempts whose handler returned during the drain, so a process exit cannot leave them for rescue | ✓ | | | T75 | Priority aging maintenance task promotes long-waiting low-priority jobs | ✓ | | ### Uniqueness (Python) diff --git a/docs/upgrade-0.6.13-to-0.6.14.md b/docs/upgrade-0.6.13-to-0.6.14.md new file mode 100644 index 000000000..b18fa715e --- /dev/null +++ b/docs/upgrade-0.6.13-to-0.6.14.md @@ -0,0 +1,45 @@ +# Upgrading from 0.6.13 to 0.6.14 + +0.6.14 is a rolling patch upgrade. Traffic keeps running throughout, and you can roll binaries back to 0.6.13 with the schema patch left in place. + +## Upgrade + +1. Install the 0.6.14 CLI. +2. Check lane ownership. This is read-only and exits 2 while lanes need repair: + + ```bash + awa --database-url "$AWA_DATABASE_URL" storage lane-sequences + ``` + +3. Run `awa migrate` once, with the schema owner as the effective role, while workers keep running: + + ```bash + PGOPTIONS='-c role=awa_owner' \ + awa --database-url "$AWA_MIGRATOR_DATABASE_URL" migrate + ``` + +4. Read the output. The `definer_lane_provisioning` patch reports `applied`, followed by how many lane sequences were repaired. Any left over are listed in warnings: a busy one needs another run of `awa storage lane-sequences --repair` as the migration role (without `PGOPTIONS`), and the others name the roles that can repair them (see [Check and repair lane ownership](security.md#check-and-repair-lane-ownership)). The migration exit status does not depend on them. +5. Roll workers, producers and admin processes to 0.6.14, in any order. + +## What `awa migrate` changes + +The schema stays at v40. The `definer_lane_provisioning` patch installs `provision_lane` and `provision_lane_sequences`, two `SECURITY DEFINER` functions owned by the schema owner, in `awa` and every other queue-storage schema. It changes the three lane helpers to create a missing sequence through them, and adds the read-only `awa.lane_sequence_problems(schema)` and `awa.repair_lane_sequence(schema, sequence)`. Every new lane sequence then belongs to the schema owner and carries the privileges of the head table it backs, whichever role creates the lane: a producer enqueueing to a new queue, priority or shard, the maintenance leader promoting into a shard after `enqueue_shards` was raised, or `awa storage prepare-queue` run as any role. + +Once the patch commits, `awa migrate` repairs existing lane sequences, each in its own short transaction. A sequence not owned by the schema owner is reassigned if the role holds its current owner's privileges and can `SET ROLE` to the schema owner. Missing grants are added if the role holds the owner's privileges. Each sequence is tried as the effective role and then as the login role. Reassigning waits at most 300 ms for the sequence's lock, which an open enqueue or claim on that lane holds, and holds it only for that sequence's transaction. A sequence that is busy, or that neither role may change, is reported and left as it was. + +The patch: + +- only replaces functions; it takes no sequence or table locks; +- keeps every lane's sequence position; +- leaves storage routing unchanged; +- applies as one transaction, so an interrupted run leaves nothing applied; run it again. + +0.6.13 and earlier binaries call the same helpers, so lanes they create after the patch are owned by the schema owner too. A 0.6.14 binary against a database without the patch detects that `provision_lane` is missing and creates lanes as 0.6.13 did. + +### External migration runners + +`awa migrate --sql --pending` prints the `R__definer_lane_provisioning` patch on a 0.6.13 database; with `--extract-to`, apply `R__definer_lane_provisioning.sql` as a repeatable script after `V40`. Run it as the schema owner. The script does not repair existing lane sequences, so then run `awa storage lane-sequences --repair` as the migration role itself. + +## Roll back + +Roll binaries back to 0.6.13 at any time. Leave the patch in place: 0.6.13 calls the same helpers with the same arguments.