From 8a8d6532dd0f3f188fd06cf429401c9594767495 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 13:27:38 +1300 Subject: [PATCH 1/7] fix(queue-storage): create queue lanes as the schema owner (0.6) A lane's enqueue and claim sequences were created by whichever role first touched the lane, through SECURITY INVOKER helpers, so only that role and the roles its default privileges named could use them. Another runtime role then failed with "permission denied for sequence" on enqueue, claim or promotion. Each queue-storage schema now has two SECURITY DEFINER functions owned by the schema owner, provision_lane and provision_lane_sequences. The lane helpers create a missing sequence through them, and each new sequence gets the privileges of the head table it backs. ensure_lane_inserts calls provision_lane when it exists and is executable, inside a savepoint, and otherwise writes the lane rows as before. The definer_lane_provisioning schema patch refreshes the helpers and repairs existing lane sequences one at a time, reassigning and granting only what the applying role may and reporting the rest. awa.lane_sequence_problems and the new `awa storage lane-sequences` command list what remains with the statements to run; awa.repair_lane_sequences and `--repair` rerun the repair. --- CHANGELOG.md | 27 + awa-cli/src/main.rs | 62 +- awa-cli/tests/migrate_patch_cli_test.rs | 11 +- .../patches/definer_lane_provisioning.sql | 80 ++ .../v023_install_queue_storage_substrate.sql | 417 +++++++++- awa-model/src/migrations.rs | 76 ++ awa-model/src/queue_storage.rs | 105 ++- awa-model/src/storage.rs | 83 ++ .../fixtures/v0613_lane_sequence_helpers.sql | 195 +++++ awa/tests/lane_ownership_roles_test.rs | 733 ++++++++++++++++++ docs/security.md | 32 +- docs/upgrade-0.6.13-to-0.6.14.md | 45 ++ 12 files changed, 1814 insertions(+), 52 deletions(-) create mode 100644 awa-model/migrations/patches/definer_lane_provisioning.sql create mode 100644 awa/tests/fixtures/v0613_lane_sequence_helpers.sql create mode 100644 awa/tests/lane_ownership_roles_test.rs create mode 100644 docs/upgrade-0.6.13-to-0.6.14.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 879a17deb..2ff251781 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,33 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- **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. The patch gives existing lane sequences to the schema owner + and adds their missing grants, one sequence at a time and only where the + applying role is allowed to. It never fails on a sequence it cannot + reassign. 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_sequences(schema)`) 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). - **`admin::job_kind_overviews` on queue storage no longer depends on the planner's row estimates.** Behind the admin UI's kinds page, it counted every job through the `current_jobs` view. Two nodes of that view are diff --git a/awa-cli/src/main.rs b/awa-cli/src/main.rs index f20634116..4b994611a 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,16 @@ 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 what the connected role is allowed to first + #[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 +679,21 @@ 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::NotApplicable + }) { + for (schema, count) in + awa_model::storage::lane_sequence_problem_counts(&pool).await? + { + eprintln!( + "warning: {count} lane sequence(s) in schema {schema} are not owned \ + by the schema owner or lack runtime grants, and the migrating role \ + could not repair them. Run `awa storage lane-sequences --schema \ + {schema}` for the statements to run and the roles to run them as." + ); + } + } } } @@ -1245,6 +1270,41 @@ async fn main() -> Result<(), Box> { }))? ); } + StorageCommands::LaneSequences { schema, repair } => { + let remaining = if repair { + let results = + awa_model::storage::repair_lane_sequences(&pool, &schema).await?; + println!( + "{}", + serde_json::to_string_pretty(&serde_json::json!({ + "schema": schema, "results": results, + }))? + ); + results + .iter() + .filter(|result| result.outcome != "repaired") + .count() + } 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, + }))? + ); + problems.len() + }; + if remaining > 0 { + eprintln!( + "{remaining} lane sequence(s) in schema {schema} need repair. Run \ + each missing_grants statement as the sequence owner (or a role \ + that can SET ROLE to it), and each reassign_sql statement as a \ + role that can SET ROLE to both the owner and expected_owner." + ); + std::process::exit(2); + } + } StorageCommands::Abort => { awa_model::storage::abort(&pool).await?; let report = awa_model::storage::status_report(&pool).await?; diff --git a/awa-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-model/migrations/patches/definer_lane_provisioning.sql b/awa-model/migrations/patches/definer_lane_provisioning.sql new file mode 100644 index 000000000..439630db9 --- /dev/null +++ b/awa-model/migrations/patches/definer_lane_provisioning.sql @@ -0,0 +1,80 @@ +-- 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, then repairs existing +-- lane sequences: it gives each one to the schema owner and adds its missing +-- grants, one sequence at a time and only where the applying role may. It +-- never fails on a sequence it cannot reassign; those are reported in a +-- warning and by awa.lane_sequence_problems(). +-- +-- Live load: function replacement, plus ALTER SEQUENCE OWNER on the lane +-- sequences that need it. That holds each such sequence's ACCESS EXCLUSIVE +-- lock until commit, so enqueues and claims on those lanes wait for the patch +-- to commit; correctly owned lanes are not touched. Reads the catalog only. +-- +-- 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_leftovers 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(*), + string_agg(format('%s (owner %s, %s)', repair.sequence_name, repair.owner, repair.outcome), + ', ' ORDER BY repair.sequence_name) + INTO v_count, v_leftovers + FROM awa.repair_lane_sequences(v_schema) AS repair + WHERE repair.outcome <> 'repaired'; + + IF v_count > 0 THEN + RAISE WARNING '% lane sequence(s) in schema % still need repair: %', v_count, v_schema, v_leftovers + USING HINT = format( + 'Run awa storage lane-sequences --schema %s for the statements to run and the role to run them as.', + 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..0d170c34f 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,342 @@ 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_queue = '' + OR p_priority IS NULL OR p_priority NOT BETWEEN 1 AND 4 + 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 can +-- SET ROLE to both the current owner and the schema owner. 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 +) +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 + 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 every lane sequence in p_schema to the schema owner and add its missing +-- runtime grants, one sequence at a time, 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 can act as that owner. Never raises for a +-- sequence it cannot repair. Returns one row per sequence that had a problem: +-- outcome is 'repaired', 'owner_unchanged' (grants complete, owner still +-- wrong) or 'unrepaired', with the statements still to run. +CREATE OR REPLACE FUNCTION awa.repair_lane_sequences(p_schema TEXT) +RETURNS TABLE ( + sequence_name TEXT, + owner TEXT, + expected_owner TEXT, + outcome TEXT, + missing_grants TEXT[], + reassign_sql TEXT +) +LANGUAGE plpgsql +SECURITY INVOKER +SET search_path = pg_catalog, pg_temp +AS $repair$ +DECLARE + v_problem RECORD; + v_remaining RECORD; + v_has_remaining BOOLEAN; + v_owner OID; + v_set_mode TEXT := CASE WHEN current_setting('server_version_num')::int >= 160000 + THEN 'SET' ELSE 'MEMBER' END; + v_statement TEXT; +BEGIN + FOR v_problem IN + SELECT * FROM awa.lane_sequence_problems(p_schema) AS problem + LOOP + SELECT sequence.relowner INTO v_owner + FROM pg_class AS sequence + WHERE sequence.oid = to_regclass(format('%I.%I', p_schema, v_problem.sequence_name)); + + -- 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 + -- rather than raising. + IF v_problem.reassign_sql IS NOT NULL + AND pg_has_role(current_user, v_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, v_problem.sequence_name) AS problem; + v_has_remaining := FOUND; + + IF v_has_remaining + 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; + SELECT * INTO v_remaining + FROM awa.lane_sequence_problems(p_schema, v_problem.sequence_name) AS problem; + v_has_remaining := FOUND; + END IF; + + sequence_name := v_problem.sequence_name; + expected_owner := v_problem.expected_owner; + IF v_has_remaining THEN + owner := v_remaining.owner; + missing_grants := v_remaining.missing_grants; + reassign_sql := v_remaining.reassign_sql; + outcome := CASE 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; + outcome := 'repaired'; + END IF; + RETURN NEXT; + END LOOP; +END +$repair$; + +REVOKE EXECUTE ON FUNCTION awa.repair_lane_sequences(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/migrations.rs b/awa-model/src/migrations.rs index 193c465fc..29a65fc93 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_sequences' \ + WHERE to_regprocedure('awa.repair_lane_sequences(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. diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 930aab715..e63b3943e 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; @@ -2455,6 +2455,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 +2506,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 +3467,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 +3509,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 +3608,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 diff --git a/awa-model/src/storage.rs b/awa-model/src/storage.rs index 64fe8be58..85d949ee2 100644 --- a/awa-model/src/storage.rs +++ b/awa-model/src/storage.rs @@ -379,3 +379,86 @@ pub async fn status_report(pool: &PgPool) -> Result, + pub reassign_sql: Option, +} + +/// What [`repair_lane_sequences`] did to one lane sequence that had a problem. +#[derive(Debug, Clone, Serialize, Deserialize, FromRow, PartialEq, Eq)] +pub struct LaneSequenceRepair { + pub sequence_name: String, + /// Owner after the repair. + pub owner: String, + pub expected_owner: String, + /// `repaired`; `owner_unchanged` when the grants are complete but the + /// connected role could not reassign the sequence; or `unrepaired` when + /// grants are still missing too. + pub outcome: String, + pub missing_grants: Vec, + pub reassign_sql: Option, +} + +/// 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 sequence_name, owner, expected_owner, missing_grants, reassign_sql \ + 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 connected role is allowed to. A sequence the role +/// cannot reassign is returned with the statements still to run, never an +/// error. +pub async fn repair_lane_sequences( + pool: &PgPool, + schema: &str, +) -> Result, AwaError> { + Ok(sqlx::query_as( + "SELECT sequence_name, owner, expected_owner, outcome, missing_grants, reassign_sql \ + FROM awa.repair_lane_sequences($1)", + ) + .bind(schema) + .fetch_all(pool) + .await?) +} + +/// How many lane sequences need repair in each queue-storage schema that has +/// any. Read-only. +pub async fn lane_sequence_problem_counts(pool: &PgPool) -> Result, AwaError> { + 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(pool) + .await?) +} 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..c02d2fa04 --- /dev/null +++ b/awa/tests/lane_ownership_roles_test.rs @@ -0,0 +1,733 @@ +//! 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::{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, + 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_sequences(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, + 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(()) + } + + /// `(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?; + } + let failed = drain(&cluster.app_b, "owner_lane").await.unwrap_err(); + assert!( + failed + .to_string() + .contains("permission denied for sequence"), + "{failed}" + ); + let failed = drain(&cluster.app_b, "app_lane").await.unwrap_err(); + assert!( + failed + .to_string() + .contains("permission denied for sequence"), + "{failed}" + ); + assert_eq!( + migrations::pending_schema_patches(&cluster.owner) + .await? + .iter() + .map(|patch| patch.name) + .collect::>(), + vec!["definer_lane_provisioning"] + ); + + // `awa migrate` runs as the migration role itself. It is a member of the + // owner, so it may reassign its own and the owner's sequences, but not + // the app role's. The leftovers do not keep the patch pending. + let report = migrations::run_with_report(&cluster.migrator).await?; + assert!(report.complete(), "{:?}", report.patches); + assert!(migrations::pending_schema_patches(&cluster.owner) + .await? + .is_empty()); + 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 problems = storage::lane_sequence_problems(&cluster.app_c, "awa").await?; + assert_eq!(problems.len(), 2, "{problems:?}"); + let app_a = cluster.roles[3].clone(); + let app_a = &app_a; + 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() + ) + ); + } + + // Re-running the repair as the migration role changes nothing and fails + // nothing. The migration role created the repair function, so it grants + // its use. + cluster + .migrator + .execute(sqlx::raw_sql(&sql(format!( + "GRANT EXECUTE ON FUNCTION awa.repair_lane_sequences(text) TO {}", + cluster.runtime_name + )))) + .await?; + let rerun = storage::repair_lane_sequences(&cluster.migrator, "awa").await?; + assert_eq!(rerun.len(), 2); + assert!(rerun.iter().all(|result| result.outcome == "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( + PgConnectOptions::from_str(&std::env::var("DATABASE_URL").unwrap_or_else(|_| { + "postgres://postgres:test@localhost:15432/awa_test".to_string() + }))? + .database(&cluster.database) + .username(&delegate) + .password("lane_test"), + ) + .await?; + let delegated = storage::repair_lane_sequences(&delegate_pool, "awa").await?; + delegate_pool.close().await; + assert_eq!(delegated.len(), 2); + for result in &delegated { + assert_eq!(result.outcome, "owner_unchanged"); + 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?; + + // A role that is a member of both the app role and the owner reassigns. + cluster + .admin + .execute(sqlx::raw_sql(&sql(format!( + "GRANT {app_a} TO {migrator}", + migrator = cluster.migrator_name + )))) + .await?; + let repaired = storage::repair_lane_sequences(&cluster.migrator, "awa").await?; + assert_eq!(repaired.len(), 2); + assert!( + repaired.iter().all(|result| result.outcome == "repaired"), + "{repaired:?}" + ); + assert!(storage::lane_sequence_problems(&cluster.app_c, "awa") + .await? + .is_empty()); + assert!(storage::repair_lane_sequences(&cluster.migrator, "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 +} diff --git a/docs/security.md b/docs/security.md index a146635bb..a6deef52b 100644 --- a/docs/security.md +++ b/docs/security.md @@ -174,23 +174,37 @@ 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 validating the queue, priority (1–4) and shard (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'); +``` + +The command exits 2 while any remain. Applying the patch runs the repair as the applying role: for each sequence it changes the owner to the schema owner and adds the missing grants, wherever that role is allowed to. A sequence it is not allowed to change does not fail the patch. It is reported in a warning, and `awa migrate` prints a summary. Run the repair again at any time; it is idempotent: + +```bash +PGOPTIONS='-c role=awa_owner' \ + awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage lane-sequences --repair +``` -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. +Changing a sequence's owner requires a role that can `SET ROLE` to both the current owner and the schema owner. A migrator usually can for sequences it or the owner created, but not for one an unrelated application role created, and managed PostgreSQL offers no superuser. For each sequence that remains, the output gives the statements to run: -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. +- `missing_grants`: run them as the current owner, or as a role that can `SET ROLE` to it. Every runtime role can then use the lane, although the owner stays unchanged. +- `reassign_sql`: run it as a role that is a member of both roles, for example after `GRANT app_role TO awa_migrator` from a role with `ADMIN` on `app_role`. Revoke that membership afterwards. ### 5. Configure your processes 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..1cccbe370 --- /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`. If it could not reassign some lane sequences, `awa migrate` prints a warning; run `awa storage lane-sequences` for the statements to run and the role to run them as (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_sequences(schema)`. 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. + +The patch then repairs existing lane sequences. Each one not owned by the schema owner is reassigned if the applying role can `SET ROLE` to both its current owner and the schema owner. Missing grants are added if the applying role can act as the owner. A sequence it cannot change is reported and left as it was. The patch never fails because of one. + +The patch: + +- replaces functions and changes only the lane sequences that need it. `ALTER SEQUENCE ... OWNER` holds that sequence's lock until the patch commits, so enqueues and claims on those lanes wait for the commit. Correctly owned lanes are not touched; +- 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, then run `awa storage lane-sequences` to check for leftovers. + +## 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. From 5c03aed269d10f7f4a7cc438ed416f5d97b6c8ae Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 13:35:43 +1300 Subject: [PATCH 2/7] test: expect restricted runtimes to create lanes through the schema owner --- .../tests/least_privilege_runtime_test.rs | 35 ++++++++++++------- 1 file changed, 23 insertions(+), 12 deletions(-) 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)) From 32d671b061e602119793b4e4b30f897a319262de Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 14:04:19 +1300 Subject: [PATCH 3/7] fix(queue-storage): accept every lane the invoker path accepts in provision_lane Queue storage takes any queue name, including an empty one, and any smallint priority, so the provisioning functions only reject a null argument or a shard outside 0-63. The released-binary rehearsal now expects a 0.6.9 producer without CREATE to create a new lane, owned by the schema owner. --- .../v023_install_queue_storage_substrate.sql | 4 +- awa/tests/patch_upgrade_rehearsal_test.rs | 44 ++++++++++++------- docs/security.md | 2 +- 3 files changed, 31 insertions(+), 19 deletions(-) diff --git a/awa-model/migrations/v023_install_queue_storage_substrate.sql b/awa-model/migrations/v023_install_queue_storage_substrate.sql index 0d170c34f..bcc8fe040 100644 --- a/awa-model/migrations/v023_install_queue_storage_substrate.sql +++ b/awa-model/migrations/v023_install_queue_storage_substrate.sql @@ -283,8 +283,8 @@ BEGIN v_seq_name TEXT; v_grant TEXT; BEGIN - IF p_queue IS NULL OR p_queue = '' - OR p_priority IS NULL OR p_priority NOT BETWEEN 1 AND 4 + 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 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/docs/security.md b/docs/security.md index a6deef52b..8c56004f1 100644 --- a/docs/security.md +++ b/docs/security.md @@ -176,7 +176,7 @@ The default `awa` queue-storage substrate is migrated by `awa migrate`. Custom q ### Lane sequences -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 validating the queue, priority (1–4) and shard (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`. +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`. 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. From 7e6259e04d9c4d45587a08c7094f6f378d788abf Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 15:26:16 +1300 Subject: [PATCH 4/7] fix(worker): drain finalization on graceful shutdown (0.6) Once a handler returned, the executor removed the attempt from the heartbeat set and finalized it in a detached task that Client::shutdown did not wait for. When the last handler returned as the drain ended, shutdown cancelled the completion batcher before that finalize reached it, and the direct-finalize fallback raced the process exit. The attempt stayed open with no heartbeat, so stale-heartbeat rescue ran the finished job again after heartbeat_staleness. Track finalizers in a TaskTracker and wait for them, within the drain timeout, before stopping the completion batcher, heartbeat and maintenance services. This is the shutdown part of #521 on main. --- CHANGELOG.md | 12 ++ Cargo.lock | 1 + awa-python/Cargo.lock | 1 + awa-worker/Cargo.toml | 2 +- awa-worker/src/client.rs | 18 +- awa-worker/src/dispatcher.rs | 1 + awa-worker/src/executor.rs | 10 +- awa/tests/ci_timing.rs | 186 +++++++++++++++++++ awa/tests/queue_storage_runtime_test.rs | 227 ++++++++++++++++++++++++ docs/test-plan.md | 1 + 10 files changed, 455 insertions(+), 4 deletions(-) create mode 100644 awa/tests/ci_timing.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 879a17deb..d49c6f89e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,18 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- **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. - **`admin::job_kind_overviews` on queue storage no longer depends on the planner's row estimates.** Behind the admin UI's kinds page, it counted every job through the `current_jobs` view. Two nodes of that view are diff --git a/Cargo.lock b/Cargo.lock index 30bd194c9..a29ff7c5c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4187,6 +4187,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] 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/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/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index a060369f8..edeacac09 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -3,6 +3,8 @@ //! 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, @@ -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 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) From 82ff82d0e78836469c15146f9436dbe7d26d234b Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 18:21:04 +1300 Subject: [PATCH 5/7] fix(queue-storage): repair lane sequences per transaction with a bounded lock wait The definer_lane_provisioning patch reassigned lane sequences inside the patch transaction. ALTER SEQUENCE ... OWNER takes the sequence's ACCESS EXCLUSIVE lock until commit, with no lock_timeout, so an open enqueue on one lane could stall the patch and every repaired lane waited for it to commit. The patch now only refreshes the helpers. After it commits, awa migrate repairs each lane sequence with its own statement through awa.repair_lane_sequence(schema, sequence), which has lock_timeout = 300ms; on 55P03 the sequence is left unchanged and reported as busy. The results are in MigrationReport::lane_sequences. `awa storage lane-sequences --repair` uses the same path, and the set-returning awa.repair_lane_sequences(schema) is gone. Reassigning needs the current owner's privileges, which the schema owner lacks for sequences the migration role or an app role created. A migration role connected with role set to the owner therefore could not reassign its own sequences. The repair now retries such a sequence as the login role, lane_sequence_problems reports repair_roles (the login roles that can repair each sequence), and leftovers are reported with the next step to take. The docs state which role to run the repair as. --- CHANGELOG.md | 14 +- awa-cli/src/main.rs | 64 +++-- .../patches/definer_lane_provisioning.sql | 30 +-- .../v023_install_queue_storage_substrate.sql | 122 +++++---- awa-model/src/migrations.rs | 32 ++- awa-model/src/storage.rs | 237 ++++++++++++++++-- awa/tests/lane_ownership_roles_test.rs | 226 +++++++++++++---- docs/security.md | 27 +- docs/upgrade-0.6.13-to-0.6.14.md | 10 +- 9 files changed, 580 insertions(+), 182 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2ff251781..a2b722287 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,14 +19,16 @@ Notable changes between releases. Detailed migration notes for storage transitio 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. The patch gives existing lane sequences to the schema owner - and adds their missing grants, one sequence at a time and only where the - applying role is allowed to. It never fails on a sequence it cannot - reassign. The new read-only `awa storage lane-sequences` (or + 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_sequences(schema)`) repairs what the - connected role may. `awa migrate` prints a summary of the leftovers. A binary + 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 diff --git a/awa-cli/src/main.rs b/awa-cli/src/main.rs index 4b994611a..1e347be53 100644 --- a/awa-cli/src/main.rs +++ b/awa-cli/src/main.rs @@ -438,7 +438,8 @@ enum StorageCommands { /// Queue-storage schema to check #[arg(long, default_value = "awa")] schema: String, - /// Reassign and grant what the connected role is allowed to first + /// Reassign and grant first, as far as the connected role (or, failing + /// that, its login role) is allowed to #[arg(long)] repair: bool, }, @@ -681,18 +682,23 @@ async fn main() -> Result<(), Box> { } if report.patches.iter().any(|patch| { patch.name == "definer_lane_provisioning" - && patch.outcome != awa_model::migrations::PatchOutcome::NotApplicable + && patch.outcome == awa_model::migrations::PatchOutcome::AlreadyPresent }) { - for (schema, count) in + // 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? { - eprintln!( - "warning: {count} lane sequence(s) in schema {schema} are not owned \ - by the schema owner or lack runtime grants, and the migrating role \ - could not repair them. Run `awa storage lane-sequences --schema \ - {schema}` for the statements to run and the roles to run them as." + 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}"); + } } } } @@ -1271,19 +1277,18 @@ async fn main() -> Result<(), Box> { ); } StorageCommands::LaneSequences { schema, repair } => { - let remaining = if repair { - let results = - awa_model::storage::repair_lane_sequences(&pool, &schema).await?; + 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": results, + "schema": schema, "results": repairs, }))? ); - results - .iter() - .filter(|result| result.outcome != "repaired") - .count() + awa_model::storage::lane_sequence_repair_next_steps(&repairs) } else { let problems = awa_model::storage::lane_sequence_problems(&pool, &schema).await?; @@ -1293,15 +1298,14 @@ async fn main() -> Result<(), Box> { "schema": schema, "problems": problems, }))? ); - problems.len() + awa_model::storage::lane_sequence_next_steps( + problems.iter().map(|problem| (problem, false)), + ) }; - if remaining > 0 { - eprintln!( - "{remaining} lane sequence(s) in schema {schema} need repair. Run \ - each missing_grants statement as the sequence owner (or a role \ - that can SET ROLE to it), and each reassign_sql statement as a \ - role that can SET ROLE to both the owner and expected_owner." - ); + if !next_steps.is_empty() { + for next_step in next_steps { + eprintln!("{next_step}"); + } std::process::exit(2); } } @@ -1480,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-model/migrations/patches/definer_lane_provisioning.sql b/awa-model/migrations/patches/definer_lane_provisioning.sql index 439630db9..2a8ced3f4 100644 --- a/awa-model/migrations/patches/definer_lane_provisioning.sql +++ b/awa-model/migrations/patches/definer_lane_provisioning.sql @@ -8,16 +8,15 @@ -- 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, then repairs existing --- lane sequences: it gives each one to the schema owner and adds its missing --- grants, one sequence at a time and only where the applying role may. It --- never fails on a sequence it cannot reassign; those are reported in a --- warning and by awa.lane_sequence_problems(). +-- 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, plus ALTER SEQUENCE OWNER on the lane --- sequences that need it. That holds each such sequence's ACCESS EXCLUSIVE --- lock until commit, so enqueues and claims on those lanes wait for the patch --- to commit; correctly owned lanes are not touched. Reads the catalog only. +-- 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 @@ -30,7 +29,6 @@ DO $$ DECLARE v_schema TEXT; - v_leftovers TEXT; v_count BIGINT; BEGIN FOR v_schema IN @@ -61,17 +59,13 @@ BEGIN ); PERFORM awa.refresh_lane_sequence_helpers(v_schema); - SELECT count(*), - string_agg(format('%s (owner %s, %s)', repair.sequence_name, repair.owner, repair.outcome), - ', ' ORDER BY repair.sequence_name) - INTO v_count, v_leftovers - FROM awa.repair_lane_sequences(v_schema) AS repair - WHERE repair.outcome <> 'repaired'; + SELECT count(*) INTO v_count + FROM awa.lane_sequence_problems(v_schema) AS problem; IF v_count > 0 THEN - RAISE WARNING '% lane sequence(s) in schema % still need repair: %', v_count, v_schema, v_leftovers + 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( - 'Run awa storage lane-sequences --schema %s for the statements to run and the role to run them as.', + 'awa migrate repairs them after this commits. Otherwise run awa storage lane-sequences --schema %s --repair as the migration role itself, without SET ROLE.', v_schema ); END IF; diff --git a/awa-model/migrations/v023_install_queue_storage_substrate.sql b/awa-model/migrations/v023_install_queue_storage_substrate.sql index bcc8fe040..4625ebada 100644 --- a/awa-model/migrations/v023_install_queue_storage_substrate.sql +++ b/awa-model/migrations/v023_install_queue_storage_substrate.sql @@ -418,15 +418,18 @@ REVOKE EXECUTE ON FUNCTION awa.refresh_lane_sequence_helpers(TEXT) FROM PUBLIC; -- 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 can --- SET ROLE to both the current owner and the schema owner. Read-only. +-- 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 + reassign_sql TEXT, + repair_roles TEXT[] ) LANGUAGE sql STABLE @@ -500,55 +503,76 @@ AS $problems$ 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 + ) 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 every lane sequence in p_schema to the schema owner and add its missing --- runtime grants, one sequence at a time, 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 can act as that owner. Never raises for a --- sequence it cannot repair. Returns one row per sequence that had a problem: --- outcome is 'repaired', 'owner_unchanged' (grants complete, owner still --- wrong) or 'unrepaired', with the statements still to run. -CREATE OR REPLACE FUNCTION awa.repair_lane_sequences(p_schema TEXT) +-- 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 + 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_has_remaining BOOLEAN; - v_owner OID; + 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 - FOR v_problem IN - SELECT * FROM awa.lane_sequence_problems(p_schema) AS problem - LOOP - SELECT sequence.relowner INTO v_owner - FROM pg_class AS sequence - WHERE sequence.oid = to_regclass(format('%I.%I', p_schema, v_problem.sequence_name)); + 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 - -- rather than raising. + -- without waiting for the lock. IF v_problem.reassign_sql IS NOT NULL - AND pg_has_role(current_user, v_owner, 'USAGE') + 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; @@ -558,10 +582,8 @@ BEGIN END IF; SELECT * INTO v_remaining - FROM awa.lane_sequence_problems(p_schema, v_problem.sequence_name) AS problem; - v_has_remaining := FOUND; - - IF v_has_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 @@ -571,31 +593,35 @@ BEGIN NULL; END; END LOOP; - SELECT * INTO v_remaining - FROM awa.lane_sequence_problems(p_schema, v_problem.sequence_name) AS problem; - v_has_remaining := FOUND; - END IF; - - sequence_name := v_problem.sequence_name; - expected_owner := v_problem.expected_owner; - IF v_has_remaining THEN - owner := v_remaining.owner; - missing_grants := v_remaining.missing_grants; - reassign_sql := v_remaining.reassign_sql; - outcome := CASE 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; - outcome := 'repaired'; END IF; - RETURN NEXT; - END LOOP; + 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_sequences(TEXT) FROM PUBLIC; +REVOKE EXECUTE ON FUNCTION awa.repair_lane_sequence(TEXT, TEXT) FROM PUBLIC; CREATE OR REPLACE FUNCTION awa.install_queue_storage_substrate( p_schema TEXT, diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index 29a65fc93..d6b4a6734 100644 --- a/awa-model/src/migrations.rs +++ b/awa-model/src/migrations.rs @@ -306,8 +306,8 @@ 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_sequences' \ - WHERE to_regprocedure('awa.repair_lane_sequences(text)') IS NULL \ + 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 ( \ @@ -578,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 { @@ -736,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 @@ -770,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/storage.rs b/awa-model/src/storage.rs index 85d949ee2..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)] @@ -382,30 +382,65 @@ 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 [`repair_lane_sequences`] did to one lane sequence that had a problem. +/// 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, - /// `repaired`; `owner_unchanged` when the grants are complete but the - /// connected role could not reassign the sequence; or `unrepaired` when - /// grants are still missing too. - pub outcome: 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. @@ -417,7 +452,8 @@ where E: PgExecutor<'e>, { Ok(sqlx::query_as( - "SELECT sequence_name, owner, expected_owner, missing_grants, reassign_sql \ + "SELECT $1::text AS schema, sequence_name, owner, expected_owner, missing_grants, \ + reassign_sql, repair_roles \ FROM awa.lane_sequence_problems($1)", ) .bind(schema) @@ -426,25 +462,182 @@ where } /// Give each lane sequence in `schema` to the schema owner and add its missing -/// grants, as far as the connected role is allowed to. A sequence the role -/// cannot reassign is returned with the statements still to run, never an -/// error. +/// 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( - pool: &PgPool, + conn: &mut PgConnection, schema: &str, ) -> Result, AwaError> { - Ok(sqlx::query_as( - "SELECT sequence_name, owner, expected_owner, outcome, missing_grants, reassign_sql \ - FROM awa.repair_lane_sequences($1)", - ) - .bind(schema) - .fetch_all(pool) - .await?) + 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(pool: &PgPool) -> Result, AwaError> { +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 \ @@ -459,6 +652,6 @@ pub async fn lane_sequence_problem_counts(pool: &PgPool) -> Result 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, @@ -593,37 +621,44 @@ async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { )))) .await?; } - let failed = drain(&cluster.app_b, "owner_lane").await.unwrap_err(); - assert!( - failed - .to_string() - .contains("permission denied for sequence"), - "{failed}" - ); - let failed = drain(&cluster.app_b, "app_lane").await.unwrap_err(); - assert!( - failed - .to_string() - .contains("permission denied for sequence"), - "{failed}" - ); - assert_eq!( - migrations::pending_schema_patches(&cluster.owner) - .await? - .iter() - .map(|patch| patch.name) - .collect::>(), - vec!["definer_lane_provisioning"] - ); + 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` runs as the migration role itself. It is a member of the - // owner, so it may reassign its own and the owner's sequences, but not - // the app role's. The leftovers do not keep the patch pending. - let report = migrations::run_with_report(&cluster.migrator).await?; + // `awa migrate` connects as the migration role and acts as the owner. As + // the owner it adds the grants the owner's sequences lack; the migration + // role's own sequences need the migration role's privileges, so they are + // reassigned as the login role. Neither can act as the app role. + let migrate_pool = cluster.migrator_as_owner().await?; + 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?; @@ -631,10 +666,10 @@ async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { assert_eq!(drain(&cluster.app_b, "migrator_lane").await?, 1); assert_eq!(drain(&cluster.app_b, "owner_lane").await?, 1); - let problems = storage::lane_sequence_problems(&cluster.app_c, "awa").await?; - assert_eq!(problems.len(), 2, "{problems:?}"); 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); @@ -655,21 +690,31 @@ async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { .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. The migration role created the repair function, so it grants - // its use. + // nothing. cluster - .migrator + .owner .execute(sqlx::raw_sql(&sql(format!( - "GRANT EXECUTE ON FUNCTION awa.repair_lane_sequences(text) TO {}", + "GRANT EXECUTE ON FUNCTION awa.repair_lane_sequence(text, text) TO {}", cluster.runtime_name )))) .await?; - let rerun = storage::repair_lane_sequences(&cluster.migrator, "awa").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 == "unrepaired")); + 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. @@ -684,20 +729,14 @@ async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { .await?; let delegate_pool = PgPoolOptions::new() .max_connections(1) - .connect_with( - PgConnectOptions::from_str(&std::env::var("DATABASE_URL").unwrap_or_else(|_| { - "postgres://postgres:test@localhost:15432/awa_test".to_string() - }))? - .database(&cluster.database) - .username(&delegate) - .password("lane_test"), - ) + .connect_with(cluster.connect_options.clone().username(&delegate)) .await?; - let delegated = storage::repair_lane_sequences(&delegate_pool, "awa").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, "owner_unchanged"); + assert_eq!(result.outcome, LaneSequenceOutcome::OwnerUnchanged); assert_eq!(&result.owner, app_a); assert!(result.missing_grants.is_empty()); assert!(result.reassign_sql.is_some()); @@ -705,7 +744,8 @@ async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { assert_eq!(drain(&cluster.app_b, "app_lane").await?, 1); enqueue(&cluster.app_c, "app_lane", 1, 1).await?; - // A role that is a member of both the app role and the owner reassigns. + // 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!( @@ -713,21 +753,105 @@ async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { migrator = cluster.migrator_name )))) .await?; - let repaired = storage::repair_lane_sequences(&cluster.migrator, "awa").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 == "repaired"), + repaired + .iter() + .all(|result| result.outcome == LaneSequenceOutcome::Repaired), "{repaired:?}" ); assert!(storage::lane_sequence_problems(&cluster.app_c, "awa") .await? .is_empty()); - assert!(storage::repair_lane_sequences(&cluster.migrator, "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/docs/security.md b/docs/security.md index 8c56004f1..8475a9c82 100644 --- a/docs/security.md +++ b/docs/security.md @@ -194,17 +194,32 @@ awa --database-url "$AWA_DATABASE_URL" storage lane-sequences --schema awa SELECT * FROM awa.lane_sequence_problems('awa'); ``` -The command exits 2 while any remain. Applying the patch runs the repair as the applying role: for each sequence it changes the owner to the schema owner and adds the missing grants, wherever that role is allowed to. A sequence it is not allowed to change does not fail the patch. It is reported in a warning, and `awa migrate` prints a summary. Run the repair again at any time; it is idempotent: +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 -PGOPTIONS='-c role=awa_owner' \ - awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage lane-sequences --repair +awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage lane-sequences --repair ``` -Changing a sequence's owner requires a role that can `SET ROLE` to both the current owner and the schema owner. A migrator usually can for sequences it or the owner created, but not for one an unrelated application role created, and managed PostgreSQL offers no superuser. For each sequence that remains, the output gives the statements to run: +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: -- `missing_grants`: run them as the current owner, or as a role that can `SET ROLE` to it. Every runtime role can then use the lane, although the owner stays unchanged. -- `reassign_sql`: run it as a role that is a member of both roles, for example after `GRANT app_role TO awa_migrator` from a role with `ADMIN` on `app_role`. Revoke that membership afterwards. +- 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 diff --git a/docs/upgrade-0.6.13-to-0.6.14.md b/docs/upgrade-0.6.13-to-0.6.14.md index 1cccbe370..b18fa715e 100644 --- a/docs/upgrade-0.6.13-to-0.6.14.md +++ b/docs/upgrade-0.6.13-to-0.6.14.md @@ -18,18 +18,18 @@ awa --database-url "$AWA_MIGRATOR_DATABASE_URL" migrate ``` -4. Read the output. The `definer_lane_provisioning` patch reports `applied`. If it could not reassign some lane sequences, `awa migrate` prints a warning; run `awa storage lane-sequences` for the statements to run and the role to run them as (see [Check and repair lane ownership](security.md#check-and-repair-lane-ownership)). The migration exit status does not depend on them. +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_sequences(schema)`. 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. +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. -The patch then repairs existing lane sequences. Each one not owned by the schema owner is reassigned if the applying role can `SET ROLE` to both its current owner and the schema owner. Missing grants are added if the applying role can act as the owner. A sequence it cannot change is reported and left as it was. The patch never fails because of one. +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: -- replaces functions and changes only the lane sequences that need it. `ALTER SEQUENCE ... OWNER` holds that sequence's lock until the patch commits, so enqueues and claims on those lanes wait for the commit. Correctly owned lanes are not touched; +- 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. @@ -38,7 +38,7 @@ The patch: ### 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, then run `awa storage lane-sequences` to check for leftovers. +`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 From a0757054ef6d0ab9ba724da6c20e3b93ffae7ab1 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 20:58:37 +1300 Subject: [PATCH 6/7] docs(security): clarify lane sequence repair outcomes --- awa-model/migrations/patches/definer_lane_provisioning.sql | 2 +- docs/security.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/awa-model/migrations/patches/definer_lane_provisioning.sql b/awa-model/migrations/patches/definer_lane_provisioning.sql index 2a8ced3f4..b75a28cad 100644 --- a/awa-model/migrations/patches/definer_lane_provisioning.sql +++ b/awa-model/migrations/patches/definer_lane_provisioning.sql @@ -65,7 +65,7 @@ BEGIN 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 repairs them after this commits. Otherwise run awa storage lane-sequences --schema %s --repair as the migration role itself, without SET ROLE.', + '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; diff --git a/docs/security.md b/docs/security.md index 8475a9c82..2b9848ec7 100644 --- a/docs/security.md +++ b/docs/security.md @@ -243,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: From 9dcb8c13ca6a4a465122f32b52866287968d290a Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Fri, 6 Nov 2026 22:31:10 +1300 Subject: [PATCH 7/7] fix(queue-storage): ignore retired ready rows in point reads --- CHANGELOG.md | 5 ++ awa-model/src/queue_storage.rs | 13 ++- awa/tests/queue_storage_runtime_test.rs | 111 ++++++++++++++++++++++++ 3 files changed, 127 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a1b0a6839..a59672bac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,11 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Fixed +- **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/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 3b0b1503f..fa4f8fb22 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -11695,8 +11695,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 "#, )) diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index 376548e7f..e03eece6b 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -15547,3 +15547,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" + ); +}