diff --git a/CHANGELOG.md b/CHANGELOG.md index 065d50db1..3945300d6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,8 +4,79 @@ Notable changes between releases. Detailed migration notes for storage transitio ## [Unreleased] +### Added + +- **The maintenance leader keeps planner statistics fresh on the queue-storage + tables autovacuum never analyses.** Autovacuum does not analyse partitioned + tables, so the thirteen ring parents (`ready_entries`, `done_entries`, + `leases`, ...) had no inherited statistics, and the two small head tables + (`queue_claim_heads`, `queue_enqueue_heads`) never reach the autoanalyze + threshold, so the planner priced lane joins against them at its defaults + while they drove joins over hundreds of thousands of retained rows. The + leader now runs `ANALYZE` on them when it is elected and then every five + minutes. `ClientBuilder::statistics_refresh_interval` (Python + `statistics_refresh_interval_ms`) changes the cadence; zero disables it. One + pass takes under a second on a small ring and a few seconds on a large one. + It analyses one relation per transaction with a 2 s `lock_timeout` (a call that + times out waiting for a lock is retried next interval and recorded as + `skipped_lock`) and never makes prune wait: a prune attempt that + lands while ANALYZE samples the oldest slot's tables returns `blocked` and + retries on its normal backoff. Results are in + `awa.maintenance.analyze.attempts` and `awa.maintenance.analyze.duration` + (scope `heads`, `ring_parent` or `pass`; outcome `ok`, `skipped_privilege`, + `skipped_lock`, `unavailable` or `failed`). +- **Migration v054 adds `awa.analyze_queue_storage_relation(schema, relation)`, + so a runtime role that does not own the tables can refresh statistics.** + `ANALYZE` skips a table its caller may not analyse, and the documented + runtime grants do not include `MAINTAIN`, so the refresh runs through this + `SECURITY DEFINER` function owned by the migration role. It analyses only a + head table or ring parent of a schema that `awa.queue_storage_schemas()` + reports as trusted. It is not executable by `PUBLIC`, because the lock it + takes lasts until the caller's transaction ends and would block ring prune: + the migration grants `EXECUTE` to the roles that already hold INSERT, UPDATE + and DELETE on the head tables (the runtime grantees), and the function + requires the calling session user to hold them. A role created later needs + the documented `GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa`. Run + `awa migrate` as the schema owner (see + [database roles](docs/security/database-roles.md)). A binary running ahead of + the migration probes for the function once per pass and, like one whose + function's owner may not analyse the tables, logs one warning naming what is + missing and skips the refresh without making failing calls or affecting job + processing. Earlier binaries never call the function. + Missing EXECUTE or head-table privileges are checked before each pass, + avoiding a failing helper call for every relation. + ### Fixed +- Statistics-refresh permission warnings distinguish caller grants from + helper/table ownership and ANALYZE privileges. +- **A queue lane is usable by every runtime role, whichever role creates it + first (migration v053).** A lane's enqueue and claim sequences belonged to + the role that first enqueued into, claimed from or promoted into it, so with + several runtime roles sharing a group role, any role other than the creator + whose default privileges did not cover it failed with `permission denied for + sequence`. That included the maintenance leader promoting into a shard + created after `enqueue_shards` was raised. Lanes are now created through two + `SECURITY DEFINER` functions per queue-storage schema, `provision_lane` and + `provision_lane_sequences`, owned by the schema owner. Each new sequence gets + the privileges of the head table it backs, explicitly rather than through + default privileges. A runtime without schema `CREATE` can therefore create + lanes itself, and `awa storage prepare-queue` is optional and may run as any + runtime role. After v053 commits, `awa migrate` gives existing lane + sequences to the schema owner and adds their missing grants, one sequence per + transaction with a 300 ms lock wait, and only where it is allowed to, trying + the effective role and then the login role. A sequence an open transaction + holds is skipped and reported as busy; one it may not change is reported + with the roles that can repair it. Neither fails the migration. The new + read-only `awa storage lane-sequences` (or + `awa.lane_sequence_problems(schema)`) lists what remains, with the statements + to run and the roles to run them as, and exits 2 while any remain; + `--repair` (or `awa.repair_lane_sequence(schema, sequence)` per sequence) + repairs what the connected role may. `awa migrate` prints a summary of the leftovers. A binary + running ahead of v053 detects that `provision_lane` is missing and creates + lanes as before, and older binaries create correctly owned lanes once v053 + has refreshed the helpers they call. See "Lane ownership and v053" in the + 0.6 to 0.7 upgrade guide. - **Batch enqueues take their lane locks in one order, so they no longer deadlock with each other.** On queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane diff --git a/awa-cli/src/main.rs b/awa-cli/src/main.rs index 405fc3023..e111c3574 100644 --- a/awa-cli/src/main.rs +++ b/awa-cli/src/main.rs @@ -273,6 +273,23 @@ fn sanitize_filename(description: &str) -> String { /// mutating commands require an explicit `--context` or `--database-url` /// once more than one context is defined, and are confirmation-gated on /// contexts marked `production = true`. +/// After a migration, say which lane sequences still need repair and how to +/// finish it. Read-only. +async fn print_lane_sequence_leftovers( + pool: &sqlx::PgPool, +) -> Result<(), Box> { + let mut problems = Vec::new(); + for (schema, _) in awa_model::storage::lane_sequence_problem_counts(pool).await? { + problems.extend(awa_model::storage::lane_sequence_problems(pool, &schema).await?); + } + for next_step in awa_model::storage::lane_sequence_next_steps( + problems.iter().map(|problem| (problem, false)), + ) { + eprintln!("warning: {next_step}"); + } + Ok(()) +} + fn access_for(command: &Commands) -> context::Access { use context::Access::{Mutating, ReadOnly}; match command { @@ -328,11 +345,13 @@ fn access_for(command: &Commands) -> context::Access { }, Commands::Storage { command } => match command { StorageCommands::Status => ReadOnly, + StorageCommands::LaneSequences { repair: false, .. } => ReadOnly, StorageCommands::Finalize { check: true, .. } => ReadOnly, StorageCommands::FlipRingAuthority { check: true, .. } => ReadOnly, StorageCommands::Prepare { .. } | StorageCommands::PrepareQueueStorageSchema { .. } | StorageCommands::PrepareQueue { .. } + | StorageCommands::LaneSequences { .. } | StorageCommands::Abort | StorageCommands::EnterMixedTransition { .. } | StorageCommands::Finalize { .. } @@ -811,7 +830,7 @@ enum StorageCommands { #[arg(long)] reset: bool, }, - /// Provision lane sequences as the migrator, before starting a runtime without DDL + /// Provision every priority and shard of a queue ahead of its first enqueue PrepareQueue { /// Queue-storage schema to provision #[arg(long, default_value = "awa")] @@ -826,6 +845,17 @@ enum StorageCommands { #[arg(long, default_value_t = 1)] queue_stripe_count: usize, }, + /// List lane sequences not owned by the schema owner or missing runtime + /// grants, with the statements that fix them. Exits 2 while any remain. + LaneSequences { + /// Queue-storage schema to check + #[arg(long, default_value = "awa")] + schema: String, + /// Reassign and grant first, as far as the connected role (or, failing + /// that, its login role) is allowed to + #[arg(long)] + repair: bool, + }, /// Abort a prepared or mixed-transition storage rollout before final activation Abort, /// Enter mixed transition and begin routing new writes to the prepared engine @@ -1182,6 +1212,7 @@ async fn run() -> Result<(), Box> { }; awa_model::migrations::run_with_options(&pool, options).await?; println!("Migrations applied successfully."); + print_lane_sequence_leftovers(&pool).await?; } } @@ -1933,6 +1964,39 @@ async fn run() -> Result<(), Box> { }))? ); } + StorageCommands::LaneSequences { schema, repair } => { + let next_steps = if repair { + let mut conn = pool.acquire().await?; + let repairs = + awa_model::storage::repair_lane_sequences(&mut conn, &schema) + .await?; + println!( + "{}", + serde_json::to_string_pretty(&serde_json::json!({ + "schema": schema, "results": repairs, + }))? + ); + awa_model::storage::lane_sequence_repair_next_steps(&repairs) + } else { + let problems = + awa_model::storage::lane_sequence_problems(&pool, &schema).await?; + println!( + "{}", + serde_json::to_string_pretty(&serde_json::json!({ + "schema": schema, "problems": problems, + }))? + ); + awa_model::storage::lane_sequence_next_steps( + problems.iter().map(|problem| (problem, false)), + ) + }; + if !next_steps.is_empty() { + for next_step in next_steps { + eprintln!("{next_step}"); + } + std::process::exit(2); + } + } StorageCommands::Abort => { awa_model::storage::abort(&pool).await?; let report = awa_model::storage::status_report(&pool).await?; diff --git a/awa-metrics/src/lib.rs b/awa-metrics/src/lib.rs index 1a1c9097b..3f1eee208 100644 --- a/awa-metrics/src/lib.rs +++ b/awa-metrics/src/lib.rs @@ -77,6 +77,8 @@ pub mod names { pub const MAINTENANCE_PRUNE_ATTEMPTS: &str = "awa.maintenance.prune.attempts"; pub const MAINTENANCE_PRUNE_SKIPPED_ROWS: &str = "awa.maintenance.prune.skipped_rows"; pub const MAINTENANCE_PRUNE_DURATION: &str = "awa.maintenance.prune.duration"; + pub const MAINTENANCE_ANALYZE_ATTEMPTS: &str = "awa.maintenance.analyze.attempts"; + pub const MAINTENANCE_ANALYZE_DURATION: &str = "awa.maintenance.analyze.duration"; pub const STORAGE_TRANSITION_READY: &str = "awa.storage.transition_ready"; pub const STORAGE_CANONICAL_LIVE_BACKLOG: &str = "awa.storage.canonical_live_backlog"; pub const STORAGE_LIVE_RUNTIME_CAPABILITY: &str = "awa.storage.live_runtime_capability"; @@ -277,6 +279,13 @@ pub struct AwaMetrics { /// rather than one per tick. Fleets alert on this rather than /// scraping the matching `tracing::warn!` line. Issue #242. pub maintenance_branch_overrun_total: Counter, + /// Planner-statistics `ANALYZE` calls by scope (`heads`, `ring_parent`, `pass`) + /// and outcome (`ok`, `skipped_privilege`, `skipped_lock`, + /// `unavailable`, `failed`). One increment per relation per refresh pass. + pub maintenance_analyze_attempts: Counter, + /// Wall-clock duration of one planner-statistics `ANALYZE` call, by scope + /// and outcome. + pub maintenance_analyze_duration_seconds: Histogram, /// Current ring `current_slot` per ring, sampled from each rotate call. /// The slot number itself isn't meaningful but the rate of advance is — /// dashboards plot `rate(slot_changes)` to see whether rotation is @@ -620,6 +629,21 @@ impl AwaMetrics { ) .with_unit("{episode}") .build(), + maintenance_analyze_attempts: meter + .u64_counter(names::MAINTENANCE_ANALYZE_ATTEMPTS) + .with_description( + "Planner-statistics ANALYZE calls by scope (heads/ring_parent) and outcome", + ) + .with_unit("{attempt}") + .build(), + maintenance_analyze_duration_seconds: meter + .f64_histogram(names::MAINTENANCE_ANALYZE_DURATION) + .with_description( + "Duration of one planner-statistics ANALYZE call by scope and outcome", + ) + .with_unit("s") + .with_boundaries(WAIT_DURATION_BUCKETS_SECONDS.to_vec()) + .build(), ring_current_slot: meter .i64_gauge(names::RING_CURRENT_SLOT) .with_description("Current slot index per ring (queue/lease/claim)") @@ -1280,6 +1304,24 @@ impl AwaMetrics { self.maintenance_branch_overrun_total.add(1, &attrs); } + /// Record one planner-statistics `ANALYZE` call. `scope` is `"heads"`, + /// `"ring_parent"` or `"pass"` (a pass-level probe or failure); `outcome` is `"ok"`, `"skipped_privilege"`, + /// `"skipped_lock"`, `"unavailable"` or `"failed"`. + pub fn record_statistics_analyze( + &self, + scope: &'static str, + outcome: &'static str, + duration: Duration, + ) { + let attrs = [ + opentelemetry::KeyValue::new("awa.analyze.scope", scope), + opentelemetry::KeyValue::new("awa.analyze.outcome", outcome), + ]; + self.maintenance_analyze_attempts.add(1, &attrs); + self.maintenance_analyze_duration_seconds + .record(duration.as_secs_f64(), &attrs); + } + /// Record a ring prune outcome. pub fn record_prune_outcome(&self, ring: &'static str, outcome: &awa_model::PruneOutcome) { let (label, reason, count) = match outcome { @@ -1457,7 +1499,7 @@ mod tests { .build(); let metrics = AwaMetrics::new(&provider.meter("awa-metrics-test")); - let seconds: [(&str, &Histogram, &[f64]); 9] = [ + let seconds: [(&str, &Histogram, &[f64]); 10] = [ ( names::ENQUEUE_DURATION, &metrics.enqueue_duration_seconds, @@ -1478,6 +1520,11 @@ mod tests { &metrics.maintenance_branch_duration_seconds, &WAIT_DURATION_BUCKETS_SECONDS, ), + ( + names::MAINTENANCE_ANALYZE_DURATION, + &metrics.maintenance_analyze_duration_seconds, + &WAIT_DURATION_BUCKETS_SECONDS, + ), ( names::DISPATCH_WAKE_TO_CLAIM_DURATION, &metrics.dispatch_wake_to_claim_seconds, diff --git a/awa-model/migrations/v023_install_queue_storage_substrate.sql b/awa-model/migrations/v023_install_queue_storage_substrate.sql index d852b2ac5..f3736cef1 100644 --- a/awa-model/migrations/v023_install_queue_storage_substrate.sql +++ b/awa-model/migrations/v023_install_queue_storage_substrate.sql @@ -42,6 +42,12 @@ LANGUAGE plpgsql SECURITY INVOKER SET search_path = pg_catalog AS $refresh$ +DECLARE + v_schema_owner OID; + v_function REGPROCEDURE; + v_function_owner OID; + v_set_mode TEXT := CASE WHEN current_setting('server_version_num')::int >= 160000 + THEN 'SET' ELSE 'MEMBER' END; BEGIN EXECUTE format( $ddl$ @@ -69,28 +75,40 @@ BEGIN ); BEGIN IF to_regclass(format('%%I.%%I', %1$L, v_enqueue_seq)) IS NULL THEN - IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + IF has_function_privilege( + to_regprocedure(format('%%I.provision_lane_sequences(text,smallint,smallint)', %1$L)), + 'EXECUTE' + ) THEN + PERFORM %1$I.provision_lane_sequences(p_queue, p_priority, p_enqueue_shard); + ELSIF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_enqueue_seq USING ERRCODE = '42501', - HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + HINT = 'Run awa migrate, or awa storage prepare-queue as the migrator, before starting producers or workers.'; + ELSE + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_enqueue_seq + ); END IF; - EXECUTE format( - 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', - %1$L, - v_enqueue_seq - ); END IF; IF to_regclass(format('%%I.%%I', %1$L, v_claim_seq)) IS NULL THEN - IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + IF has_function_privilege( + to_regprocedure(format('%%I.provision_lane_sequences(text,smallint,smallint)', %1$L)), + 'EXECUTE' + ) THEN + PERFORM %1$I.provision_lane_sequences(p_queue, p_priority, p_enqueue_shard); + ELSIF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_claim_seq USING ERRCODE = '42501', - HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + HINT = 'Run awa migrate, or awa storage prepare-queue as the migrator, before starting producers or workers.'; + ELSE + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_claim_seq + ); END IF; - EXECUTE format( - 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', - %1$L, - v_claim_seq - ); END IF; UPDATE %1$I.queue_enqueue_heads @@ -131,16 +149,22 @@ BEGIN v_start BIGINT; BEGIN IF to_regclass(format('%%I.%%I', %1$L, v_seq_name)) IS NULL THEN - IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + IF has_function_privilege( + to_regprocedure(format('%%I.provision_lane_sequences(text,smallint,smallint)', %1$L)), + 'EXECUTE' + ) THEN + PERFORM %1$I.provision_lane_sequences(NEW.queue, NEW.priority, NEW.enqueue_shard); + ELSIF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_seq_name USING ERRCODE = '42501', - HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + HINT = 'Run awa migrate, or awa storage prepare-queue as the migrator, before starting producers or workers.'; + ELSE + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_seq_name + ); END IF; - EXECUTE format( - 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', - %1$L, - v_seq_name - ); END IF; NEW.seq_name := v_seq_name; @@ -197,16 +221,22 @@ BEGIN ); BEGIN IF to_regclass(format('%%I.%%I', %1$L, v_seq_name)) IS NULL THEN - IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + IF has_function_privilege( + to_regprocedure(format('%%I.provision_lane_sequences(text,smallint,smallint)', %1$L)), + 'EXECUTE' + ) THEN + PERFORM %1$I.provision_lane_sequences(NEW.queue, NEW.priority, NEW.enqueue_shard); + ELSIF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_seq_name USING ERRCODE = '42501', - HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + HINT = 'Run awa migrate, or awa storage prepare-queue as the migrator, before starting producers or workers.'; + ELSE + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_seq_name + ); END IF; - EXECUTE format( - 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', - %1$L, - v_seq_name - ); END IF; NEW.seq_name := v_seq_name; @@ -221,11 +251,367 @@ BEGIN p_schema ); + -- Lane sequences are created by the schema owner through these two + -- SECURITY DEFINER functions, whichever role first touches the lane, so + -- every runtime role can use them. Their grants mirror the head table each + -- sequence backs (see awa.lane_sequence_problems); they never depend on + -- the creating role's default privileges. + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.provision_lane_sequences( + p_queue TEXT, + p_priority SMALLINT, + p_enqueue_shard SMALLINT + ) + RETURNS VOID + LANGUAGE plpgsql + SECURITY DEFINER + SET search_path = pg_catalog, pg_temp + AS $func$ + DECLARE + v_seq_name TEXT; + v_grant TEXT; + BEGIN + IF p_queue IS NULL + OR p_priority IS NULL + OR p_enqueue_shard IS NULL OR p_enqueue_shard NOT BETWEEN 0 AND 63 THEN + RAISE EXCEPTION 'invalid queue lane: queue %%, priority %%, enqueue shard %%', + quote_nullable(p_queue), p_priority, p_enqueue_shard + USING ERRCODE = '22023'; + END IF; + + -- The lane lock QueueStorage takes before it writes lane rows, so + -- concurrent first touches of one lane create its sequences once. + PERFORM pg_advisory_xact_lock(hashtextextended( + format('%%s:%%s:%%s:%%s', %1$L, p_queue, p_priority, p_enqueue_shard), + 0 + )); + + FOREACH v_seq_name IN ARRAY ARRAY[ + %1$I.queue_lane_sequence_name('queue_enqueue_seq', p_queue, p_priority, p_enqueue_shard), + %1$I.queue_lane_sequence_name('queue_claim_seq', p_queue, p_priority, p_enqueue_shard) + ] LOOP + IF to_regclass(format('%%I.%%I', %1$L, v_seq_name)) IS NULL THEN + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_seq_name + ); + FOR v_grant IN + SELECT unnest(problem.missing_grants) + FROM awa.lane_sequence_problems(%1$L, v_seq_name) AS problem + WHERE problem.owner = current_user + LOOP + EXECUTE v_grant; + END LOOP; + END IF; + END LOOP; + END; + $func$ + $ddl$, + p_schema + ); + + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.provision_lane( + p_queue TEXT, + p_priority SMALLINT, + p_enqueue_shard SMALLINT + ) + RETURNS VOID + LANGUAGE plpgsql + SECURITY DEFINER + SET search_path = pg_catalog, pg_temp + AS $func$ + BEGIN + -- Validates the lane and takes its lock before any row is written. + PERFORM %1$I.provision_lane_sequences(p_queue, p_priority, p_enqueue_shard); + + INSERT INTO %1$I.queue_lanes (queue, priority) + VALUES (p_queue, p_priority) + ON CONFLICT (queue, priority) DO NOTHING; + + -- The head triggers fire before ON CONFLICT and set the lane's + -- sequence, so rows that already exist are skipped outright. + INSERT INTO %1$I.queue_enqueue_heads (queue, priority, enqueue_shard) + SELECT p_queue, p_priority, p_enqueue_shard + WHERE NOT EXISTS ( + SELECT 1 FROM %1$I.queue_enqueue_heads AS heads + WHERE heads.queue = p_queue + AND heads.priority = p_priority + AND heads.enqueue_shard = p_enqueue_shard + ) + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + + INSERT INTO %1$I.queue_claim_heads (queue, priority, enqueue_shard) + SELECT p_queue, p_priority, p_enqueue_shard + WHERE NOT EXISTS ( + SELECT 1 FROM %1$I.queue_claim_heads AS heads + WHERE heads.queue = p_queue + AND heads.priority = p_priority + AND heads.enqueue_shard = p_enqueue_shard + ) + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + + PERFORM %1$I.ensure_lane_sequences(p_queue, p_priority, p_enqueue_shard); + END; + $func$ + $ddl$, + p_schema + ); + + SELECT namespace.nspowner INTO v_schema_owner + FROM pg_namespace AS namespace + WHERE namespace.nspname = p_schema; + + FOREACH v_function IN ARRAY ARRAY[ + format('%I.provision_lane_sequences(text,smallint,smallint)', p_schema)::regprocedure, + format('%I.provision_lane(text,smallint,smallint)', p_schema)::regprocedure + ] LOOP + EXECUTE format('GRANT EXECUTE ON FUNCTION %s TO PUBLIC', v_function); + + SELECT routine.proowner INTO v_function_owner + FROM pg_proc AS routine + WHERE routine.oid = v_function; + + IF v_function_owner <> v_schema_owner THEN + IF pg_has_role(current_user, v_function_owner, 'USAGE') + AND pg_has_role(current_user, v_schema_owner, v_set_mode) THEN + EXECUTE format( + 'ALTER FUNCTION %s OWNER TO %I', + v_function, + pg_get_userbyid(v_schema_owner) + ); + ELSE + RAISE WARNING '% is owned by %, not the owner % of schema %; lane sequences it creates will be owned by % too', + v_function, pg_get_userbyid(v_function_owner), pg_get_userbyid(v_schema_owner), + p_schema, pg_get_userbyid(v_function_owner) + USING HINT = format( + 'Run awa migrate as a role that can SET ROLE to %s, or as %s itself.', + quote_ident(pg_get_userbyid(v_schema_owner)), + quote_ident(pg_get_userbyid(v_schema_owner)) + ); + END IF; + END IF; + END LOOP; + END; $refresh$; REVOKE EXECUTE ON FUNCTION awa.refresh_lane_sequence_helpers(TEXT) FROM PUBLIC; +-- Lane sequences in p_schema (or the one named p_sequence) that are not owned +-- by the schema owner or lack a runtime grant. A lane sequence carries a head +-- table's cursor, so whoever may write that head table (INSERT or UPDATE) gets +-- USAGE, SELECT and UPDATE on it, and whoever may read it gets SELECT: +-- queue_enqueue_seq_* mirror queue_enqueue_heads, queue_claim_seq_* mirror +-- queue_claim_heads. missing_grants holds the statements to run as the +-- sequence owner; reassign_sql, when the owner is wrong, needs a role that +-- holds the current owner's privileges and can SET ROLE to the schema owner. +-- repair_roles lists the login roles, superusers aside, that can run both. +-- Read-only. +CREATE OR REPLACE FUNCTION awa.lane_sequence_problems(p_schema TEXT, p_sequence TEXT DEFAULT NULL) +RETURNS TABLE ( + sequence_name TEXT, + owner TEXT, + expected_owner TEXT, + missing_grants TEXT[], + reassign_sql TEXT, + repair_roles TEXT[] +) +LANGUAGE sql +STABLE +SECURITY INVOKER +SET search_path = pg_catalog, pg_temp +AS $problems$ + WITH target AS ( + SELECT namespace.oid, namespace.nspowner + FROM pg_namespace AS namespace + WHERE namespace.nspname = p_schema + ), + lane_sequence AS ( + SELECT sequence.oid, + sequence.relname::text AS relname, + sequence.relowner, + sequence.relacl, + target.nspowner, + CASE WHEN starts_with(sequence.relname::text, 'queue_enqueue_seq_') + THEN 'queue_enqueue_heads' ELSE 'queue_claim_heads' END AS head_table + FROM target + JOIN pg_class AS sequence ON sequence.relnamespace = target.oid + WHERE sequence.relkind = 'S' + AND sequence.relname ~ '^queue_(enqueue|claim)_seq_[0-9a-f]{32}$' + AND (p_sequence IS NULL OR sequence.relname = p_sequence) + ), + wanted AS ( + SELECT DISTINCT lane_sequence.oid, head_privilege.grantee, sequence_privilege.privilege_type + FROM lane_sequence + JOIN pg_class AS head + ON head.relnamespace = (SELECT target.oid FROM target) + AND head.relname = lane_sequence.head_table + CROSS JOIN LATERAL aclexplode(COALESCE(head.relacl, acldefault('r', head.relowner))) AS head_privilege + CROSS JOIN LATERAL unnest( + CASE head_privilege.privilege_type + WHEN 'SELECT' THEN ARRAY['SELECT'] + WHEN 'INSERT' THEN ARRAY['USAGE', 'SELECT', 'UPDATE'] + WHEN 'UPDATE' THEN ARRAY['USAGE', 'SELECT', 'UPDATE'] + ELSE ARRAY[]::TEXT[] + END + ) AS sequence_privilege(privilege_type) + WHERE head_privilege.grantee <> lane_sequence.nspowner + ), + missing AS ( + SELECT wanted.oid, wanted.grantee, wanted.privilege_type + FROM wanted + JOIN lane_sequence ON lane_sequence.oid = wanted.oid + WHERE NOT EXISTS ( + SELECT 1 + FROM aclexplode(COALESCE(lane_sequence.relacl, acldefault('S', lane_sequence.relowner))) AS held + WHERE held.grantee = wanted.grantee + AND held.privilege_type = wanted.privilege_type + ) + ) + SELECT lane_sequence.relname, + pg_get_userbyid(lane_sequence.relowner)::text, + pg_get_userbyid(lane_sequence.nspowner)::text, + ARRAY( + SELECT format( + 'GRANT %s ON SEQUENCE %I.%I TO %s', + string_agg(missing.privilege_type, ', ' ORDER BY missing.privilege_type DESC), + p_schema, + lane_sequence.relname, + CASE missing.grantee WHEN 0 THEN 'PUBLIC' + ELSE quote_ident(pg_get_userbyid(missing.grantee)) END + ) + FROM missing + WHERE missing.oid = lane_sequence.oid + GROUP BY missing.grantee + ORDER BY 1 + ), + CASE WHEN lane_sequence.relowner <> lane_sequence.nspowner THEN format( + 'ALTER SEQUENCE %I.%I OWNER TO %I', + p_schema, lane_sequence.relname, pg_get_userbyid(lane_sequence.nspowner) + ) END, + ARRAY( + SELECT candidate.rolname::text + FROM pg_roles AS candidate + WHERE candidate.rolcanlogin + AND NOT candidate.rolsuper + AND pg_has_role(candidate.oid, lane_sequence.relowner, 'USAGE') + AND (lane_sequence.relowner = lane_sequence.nspowner + OR (pg_has_role(candidate.oid, lane_sequence.nspowner, 'USAGE') + AND pg_has_role( + candidate.oid, + lane_sequence.nspowner, + CASE WHEN current_setting('server_version_num')::int >= 160000 + THEN 'SET' ELSE 'MEMBER' END + ))) + ORDER BY candidate.rolname + ) + FROM lane_sequence + WHERE lane_sequence.relowner <> lane_sequence.nspowner + OR EXISTS (SELECT 1 FROM missing WHERE missing.oid = lane_sequence.oid) + ORDER BY lane_sequence.relname +$problems$; + +-- Give one lane sequence to the schema owner and add its missing runtime +-- grants, doing only what the calling role is allowed to: a sequence it +-- cannot reassign keeps its owner and still gets the grants when the caller +-- holds that owner's privileges. Never raises for a sequence it cannot +-- repair. Reassigning takes the sequence's ACCESS EXCLUSIVE lock, which an +-- open enqueue or claim on the lane conflicts with, so the lock wait is +-- bounded: on timeout nothing changes and the outcome is 'busy'. The lock is +-- held until the caller's transaction ends, so callers repair one sequence +-- per transaction. Returns no row when the sequence needs nothing, otherwise +-- one row: outcome is 'repaired', 'owner_unchanged' (grants complete, owner +-- still wrong), 'unrepaired' or 'busy', with the statements still to run. +CREATE OR REPLACE FUNCTION awa.repair_lane_sequence(p_schema TEXT, p_sequence TEXT) +RETURNS TABLE ( + sequence_name TEXT, + owner TEXT, + expected_owner TEXT, + outcome TEXT, + missing_grants TEXT[], + reassign_sql TEXT, + repair_roles TEXT[] +) +LANGUAGE plpgsql +SECURITY INVOKER +SET search_path = pg_catalog, pg_temp +SET lock_timeout = '300ms' +AS $repair$ +DECLARE + v_problem RECORD; + v_remaining RECORD; + v_busy BOOLEAN := false; + v_set_mode TEXT := CASE WHEN current_setting('server_version_num')::int >= 160000 + THEN 'SET' ELSE 'MEMBER' END; + v_statement TEXT; +BEGIN + SELECT * INTO v_problem + FROM awa.lane_sequence_problems(p_schema, p_sequence) AS problem; + IF NOT FOUND THEN + RETURN; + END IF; + + BEGIN + -- ALTER ... OWNER TO needs the current owner's privileges, SET ROLE to + -- the new owner, and the new owner's CREATE on the schema (which a + -- schema owner has). Check first so an unrelated owner is skipped + -- without waiting for the lock. + IF v_problem.reassign_sql IS NOT NULL + AND pg_has_role(current_user, to_regrole(quote_ident(v_problem.owner)), 'USAGE') + AND pg_has_role(current_user, to_regrole(quote_ident(v_problem.expected_owner)), v_set_mode) THEN + BEGIN + EXECUTE v_problem.reassign_sql; + EXCEPTION WHEN insufficient_privilege THEN + NULL; + END; + END IF; + + SELECT * INTO v_remaining + FROM awa.lane_sequence_problems(p_schema, p_sequence) AS problem; + IF FOUND + AND cardinality(v_remaining.missing_grants) > 0 + AND pg_has_role(current_user, to_regrole(quote_ident(v_remaining.owner)), 'USAGE') THEN + FOREACH v_statement IN ARRAY v_remaining.missing_grants LOOP + BEGIN + EXECUTE v_statement; + EXCEPTION WHEN insufficient_privilege THEN + NULL; + END; + END LOOP; + END IF; + EXCEPTION WHEN lock_not_available THEN + v_busy := true; + END; + + SELECT * INTO v_remaining + FROM awa.lane_sequence_problems(p_schema, p_sequence) AS problem; + sequence_name := v_problem.sequence_name; + expected_owner := v_problem.expected_owner; + IF FOUND THEN + owner := v_remaining.owner; + missing_grants := v_remaining.missing_grants; + reassign_sql := v_remaining.reassign_sql; + repair_roles := v_remaining.repair_roles; + outcome := CASE WHEN v_busy THEN 'busy' + WHEN cardinality(v_remaining.missing_grants) = 0 THEN 'owner_unchanged' + ELSE 'unrepaired' END; + ELSE + owner := v_problem.expected_owner; + missing_grants := ARRAY[]::TEXT[]; + reassign_sql := NULL; + repair_roles := ARRAY[]::TEXT[]; + outcome := 'repaired'; + END IF; + RETURN NEXT; +END +$repair$; + +REVOKE EXECUTE ON FUNCTION awa.repair_lane_sequence(TEXT, TEXT) FROM PUBLIC; + -- Retained job output (ADR-048): add the output columns to an existing -- substrate and expose them through terminal_jobs. Called by the installer -- and by migration v052, which does not replay the installer. Each column is diff --git a/awa-model/migrations/v053_definer_lane_provisioning.sql b/awa-model/migrations/v053_definer_lane_provisioning.sql new file mode 100644 index 000000000..ced1a7642 --- /dev/null +++ b/awa-model/migrations/v053_definer_lane_provisioning.sql @@ -0,0 +1,58 @@ +-- v053: Create queue lanes as the schema owner, whichever role touches them first. +-- Tracking: https://github.com/hardbyte/awa/pull/637 +-- Lane sequences used to belong to the role that first enqueued into, claimed +-- from or promoted into a lane, and only that role (plus whoever its default +-- privileges named) could use them. refresh_lane_sequence_helpers now installs +-- two SECURITY DEFINER functions, provision_lane and provision_lane_sequences, +-- owned by the schema owner, and the lane helpers create a missing sequence +-- through them. Each new sequence gets the grants of the head table it backs. +-- This migration does not reassign existing lane sequences: ALTER SEQUENCE +-- OWNER takes the sequence's ACCESS EXCLUSIVE lock until commit, and the +-- migration runs in one transaction. After it commits, `awa migrate` repairs +-- each lane sequence in its own short transaction with a bounded lock wait +-- (awa.repair_lane_sequence). Another runner can do the same with +-- `awa storage lane-sequences --repair`. +-- +-- N-1 binaries call the same helpers with the same arguments and cursor +-- semantics, so lanes they create are owned by the schema owner too. A binary +-- that predates provision_lane keeps using the helpers; this binary probes for +-- provision_lane and uses the helpers when it is missing. +-- +-- Live load: function replacement only. Reads the catalog; no table DDL, row +-- scans or sequence locks. + +DO $$ +DECLARE + v_schema TEXT; + v_count BIGINT; +BEGIN + FOR v_schema IN + SELECT discovered.schema_name FROM awa.queue_storage_schemas() AS discovered + LOOP + IF to_regclass(format('%I.queue_enqueue_heads', v_schema)) IS NULL + OR to_regclass(format('%I.queue_claim_heads', v_schema)) IS NULL + OR to_regclass(format('%I.queue_ring_state', v_schema)) IS NULL + OR to_regclass(format('%I.lease_ring_state', v_schema)) IS NULL + OR to_regclass(format('%I.claim_ring_state', v_schema)) IS NULL THEN + CONTINUE; + END IF; + + PERFORM awa.refresh_lane_sequence_helpers(v_schema); + + SELECT count(*) INTO v_count + FROM awa.lane_sequence_problems(v_schema) AS problem; + + IF v_count > 0 THEN + RAISE NOTICE '% lane sequence(s) in schema % are not owned by the schema owner or lack runtime grants', v_count, v_schema + USING HINT = format( + 'awa migrate attempts repair after this commits; review its warnings for any sequences left unrepaired. To retry, run awa storage lane-sequences --schema %s --repair as the migration role itself, without SET ROLE.', + v_schema + ); + END IF; + END LOOP; +END; +$$; + +INSERT INTO awa.schema_version (version, description) +VALUES (53, 'Create queue lanes as the schema owner') +ON CONFLICT (version) DO NOTHING; diff --git a/awa-model/migrations/v054_analyze_queue_storage_statistics.sql b/awa-model/migrations/v054_analyze_queue_storage_statistics.sql new file mode 100644 index 000000000..ef2ecfc1e --- /dev/null +++ b/awa-model/migrations/v054_analyze_queue_storage_statistics.sql @@ -0,0 +1,180 @@ +-- v054: Planner statistics for the queue-storage tables autovacuum never analyses (#631) +-- +-- Autovacuum does not analyse partitioned tables, so the ring parents +-- (ready_entries, done_entries, leases, ...) have no inherited statistics. It +-- also never reaches the autoanalyze threshold on the two small head tables +-- (queue_claim_heads, queue_enqueue_heads), so the planner prices lane joins +-- against them at its defaults (320 rows, n_distinct 200) while they drive +-- joins over hundreds of thousands of retained rows. +-- +-- The maintenance leader refreshes both on an interval. Runtimes usually +-- connect as a role that does not own the tables, and ANALYZE skips with a +-- warning a table its caller may not analyse, so the refresh runs through +-- this SECURITY DEFINER function owned by the migration role, which owns the +-- tables. +-- +-- The function analyses one relation per call and the caller runs each call as +-- its own transaction, so the locks ANALYZE takes on a partitioned table's +-- children are held for that table rather than for the whole refresh. +-- +-- Locking on a partitioned table, per the ANALYZE documentation: SKIP_LOCKED +-- applies when ANALYZE starts work on a relation, so a child that ring prune +-- holds ACCESS EXCLUSIVE is skipped during the per-child analysis, but +-- sampling the children for the parent's statistics can still wait for it. +-- That wait is bounded by the function's lock_timeout and fails the call with +-- SQLSTATE 55P03, which the caller records as skipped_lock. SKIP_LOCKED also +-- skips every partition silently when the parent itself is locked, so the +-- function takes the parent's SHARE UPDATE EXCLUSIVE lock first and lets that +-- wait fail instead. +-- +-- Definer boundary: only a head table or ring parent of a schema that +-- awa.queue_storage_schemas() reports is analysed, which requires the schema's +-- ownership, grants and triggers to be trusted, and SQL is built from nothing +-- but quoted identifiers. ANALYZE writes statistics and nothing else. +-- Only the fixed list of ring parents and the two head tables is accepted; +-- a migration that adds a partitioned table to the substrate must extend it. +-- +-- Who may call it: the lock above lasts until the caller's transaction ends +-- and blocks ring prune of that table, so EXECUTE is not granted to PUBLIC. It +-- goes to the roles that already hold INSERT, UPDATE and DELETE on the head +-- tables of the awa schema (the runtime grantees), and the function itself +-- requires the calling session user to hold those privileges on the target +-- schema's queue_claim_heads. A role created later needs the documented +-- `GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa` (or default privileges). +-- +-- Returns 'analyzed'; 'skipped_privilege' when the function owner may not +-- analyse the relation; or 'missing' when the relation does not exist. +-- +-- N-1 binaries never call it. A binary newer than this schema probes for the +-- function, logs once that it is missing and skips the refresh. Creating or +-- replacing it takes no table lock and completes in milliseconds; calling it +-- takes SHARE UPDATE EXCLUSIVE on the relation until the caller's transaction +-- ends. + +CREATE OR REPLACE FUNCTION awa.analyze_queue_storage_relation( + p_schema TEXT, + p_relation TEXT +) +RETURNS TEXT +LANGUAGE plpgsql +SECURITY DEFINER +SET search_path = pg_catalog, pg_temp +SET lock_timeout = '2s' +AS $$ +DECLARE + v_heads REGCLASS; + v_relid OID; + v_relkind "char"; + v_relowner OID; + v_may_analyze BOOLEAN; +BEGIN + -- awa_statistics_helper_v1: bump when the definition or grants change. + IF NOT EXISTS ( + SELECT 1 FROM awa.queue_storage_schemas() AS substrate + WHERE substrate.schema_name = p_schema + ) THEN + RAISE EXCEPTION '% is not a queue-storage schema', p_schema + USING ERRCODE = '22023'; + END IF; + + IF p_relation <> ALL (ARRAY[ + 'queue_claim_heads', + 'queue_enqueue_heads', + 'done_entries', + 'lease_claim_batches', + 'lease_claim_closure_batches', + 'lease_claim_closures', + 'lease_claims', + 'leases', + 'queue_terminal_count_deltas', + 'ready_claim_attempt_batches', + 'ready_entries', + 'ready_segments', + 'ready_tombstones', + 'receipt_completion_batches', + 'receipt_completion_tombstones' + ]) THEN + RAISE EXCEPTION '%.% is not a ring parent or head table of a queue-storage schema', + p_schema, p_relation + USING ERRCODE = '22023'; + END IF; + + v_heads := to_regclass(format('%I.queue_claim_heads', p_schema)); + IF v_heads IS NULL + OR NOT has_table_privilege(session_user, v_heads, 'INSERT') + OR NOT has_table_privilege(session_user, v_heads, 'UPDATE') + OR NOT has_table_privilege(session_user, v_heads, 'DELETE') THEN + RAISE EXCEPTION 'permission denied: % may not maintain queue storage in schema %', + session_user, p_schema + USING ERRCODE = '42501'; + END IF; + + v_relid := to_regclass(format('%I.%I', p_schema, p_relation)); + IF v_relid IS NULL THEN + RETURN 'missing'; + END IF; + + SELECT class.relkind, class.relowner INTO v_relkind, v_relowner + FROM pg_class AS class WHERE class.oid = v_relid; + + -- ANALYZE skips, with a warning rather than an error, a table the caller + -- may not analyse: the owner, a MAINTAIN holder (PostgreSQL 17+), a + -- member of pg_maintain, or the database owner may analyse it. + SELECT + pg_has_role(current_user, v_relowner, 'USAGE') + OR pg_has_role( + current_user, + (SELECT database.datdba FROM pg_database AS database + WHERE database.datname = current_database()), + 'USAGE' + ) + OR CASE + WHEN current_setting('server_version_num')::int >= 170000 + THEN has_table_privilege(v_relid, 'MAINTAIN') + ELSE FALSE + END + INTO v_may_analyze; + + IF NOT v_may_analyze THEN + RETURN 'skipped_privilege'; + END IF; + + IF v_relkind = 'p' THEN + EXECUTE format( + 'LOCK TABLE ONLY %I.%I IN SHARE UPDATE EXCLUSIVE MODE', p_schema, p_relation + ); + EXECUTE format('ANALYZE (SKIP_LOCKED) %I.%I', p_schema, p_relation); + ELSE + EXECUTE format('ANALYZE %I.%I', p_schema, p_relation); + END IF; + RETURN 'analyzed'; +END +$$; + +REVOKE EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(TEXT, TEXT) FROM PUBLIC; + +DO $grant$ +DECLARE + v_grantee TEXT; +BEGIN + FOR v_grantee IN + SELECT pg_get_userbyid(privilege.grantee) + FROM pg_class AS heads + CROSS JOIN LATERAL aclexplode(heads.relacl) AS privilege + WHERE heads.oid = to_regclass('awa.queue_claim_heads') + AND privilege.grantee <> 0 + AND privilege.privilege_type IN ('INSERT', 'UPDATE', 'DELETE') + GROUP BY privilege.grantee + HAVING count(DISTINCT privilege.privilege_type) = 3 + LOOP + EXECUTE format( + 'GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) TO %I', + v_grantee + ); + END LOOP; +END +$grant$; + +INSERT INTO awa.schema_version (version, description) +VALUES (54, 'Planner statistics refresh for queue-storage tables autovacuum never analyses') +ON CONFLICT (version) DO NOTHING; diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index 385035b65..443967012 100644 --- a/awa-model/src/migrations.rs +++ b/awa-model/src/migrations.rs @@ -5,7 +5,7 @@ use sqlx::{Connection, PgPool}; use tracing::{info, warn}; /// Current schema version. -pub const CURRENT_VERSION: i32 = 52; +pub const CURRENT_VERSION: i32 = 54; /// Migrations that require an exclusive (no-live-runtime) upgrade window. /// @@ -292,6 +292,16 @@ const MIGRATIONS: &[(i32, &str, &[&str])] = &[ ), (51, "Per-slot claim functions", &[V23_UP, V51_UP]), (52, "Retained job output", &[V23_UP, V52_UP]), + ( + 53, + "Create queue lanes as the schema owner", + &[V23_UP, V53_UP], + ), + ( + 54, + "Planner statistics refresh for queue-storage tables autovacuum never analyses", + &[V23_UP, V54_UP], + ), ]; const V1_UP: &str = include_str!("../migrations/v001_canonical_schema.sql"); @@ -352,6 +362,8 @@ const V49_UP: &str = include_str!("../migrations/v049_prepared_lane_sequences.sq const V50_UP: &str = include_str!("../migrations/v050_drop_ready_lane_shard_index.sql"); const V51_UP: &str = include_str!("../migrations/v051_claim_slot_functions.sql"); const V52_UP: &str = include_str!("../migrations/v052_retained_job_output.sql"); +const V53_UP: &str = include_str!("../migrations/v053_definer_lane_provisioning.sql"); +const V54_UP: &str = include_str!("../migrations/v054_analyze_queue_storage_statistics.sql"); /// Old version numbers from pre-0.4 releases that used V3/V4/V5 numbering. /// Also tolerates the unreleased inline-V6 branch numbering used during review. @@ -417,9 +429,9 @@ pub async fn run_with_options(pool: &PgPool, options: MigrateOptions) -> Result< // only sound because every migration is transaction-safe (no // `CREATE INDEX CONCURRENTLY`, `VACUUM`, or explicit transaction control). let mut retry = 0; - loop { + let version_before = loop { match run_migration_transaction(pool, &options, lock_key).await { - Ok(()) => break, + Ok(version_before) => break version_before, Err(error) if is_migration_deadlock(&error) && retry < MIGRATION_DEADLOCK_RETRIES => { retry += 1; let delay = std::time::Duration::from_millis(50 * (1 << (retry - 1))); @@ -433,6 +445,10 @@ pub async fn run_with_options(pool: &PgPool, options: MigrateOptions) -> Result< } Err(error) => return Err(error), } + }; + + if version_before < LANE_PROVISIONING_VERSION { + repair_lane_sequences_after_migration(pool).await; } // Best-effort admin-metadata cache warmup, deliberately outside the @@ -443,29 +459,73 @@ pub async fn run_with_options(pool: &PgPool, options: MigrateOptions) -> Result< Ok(()) } +/// Returns the schema version the transaction started from. async fn run_migration_transaction( pool: &PgPool, options: &MigrateOptions, lock_key: i64, -) -> Result<(), AwaError> { +) -> Result { let mut tx = pool.begin().await?; sqlx::query("SELECT pg_advisory_xact_lock($1)") .bind(lock_key) .execute(&mut *tx) .await?; - apply_migrations(&mut tx, options).await?; + let version_before = apply_migrations(&mut tx, options).await?; tx.commit().await?; - Ok(()) + Ok(version_before) +} + +/// The migration that makes lane sequences belong to the schema owner. +const LANE_PROVISIONING_VERSION: i32 = 53; + +/// Repair existing lane sequences once v053 has committed. +/// +/// `ALTER SEQUENCE ... OWNER` takes the sequence's ACCESS EXCLUSIVE lock until +/// its transaction ends, so it cannot run inside the migration transaction +/// without stalling every enqueue and claim on those lanes until the whole +/// upgrade commits. Each sequence is repaired in its own short transaction +/// with a bounded lock wait instead. Best-effort: the schema is committed, and +/// whatever is left is reported with the command that finishes it. +/// +/// Uses a detached connection for the reason given on +/// [`warm_admin_metadata_cache`]. +async fn repair_lane_sequences_after_migration(pool: &PgPool) { + let mut conn = match pool.acquire().await { + Ok(conn) => conn.detach(), + Err(error) => { + warn!(%error, "Could not connect to repair lane sequences; run `awa storage lane-sequences --repair`"); + return; + } + }; + match crate::storage::repair_lane_sequences_in_every_schema(&mut conn).await { + Ok(repairs) => { + let repaired = repairs + .iter() + .filter(|repair| repair.outcome == crate::storage::LaneSequenceOutcome::Repaired) + .count(); + if repaired > 0 { + info!(repaired, "Gave lane sequences to the schema owner"); + } + for next_step in crate::storage::lane_sequence_repair_next_steps(&repairs) { + warn!("{next_step}"); + } + } + Err(error) => { + warn!(%error, "Could not repair lane sequences; run `awa storage lane-sequences --repair`"); + } + } + let _ = conn.close().await; } fn is_migration_deadlock(error: &AwaError) -> bool { error.is_deadlock() } +/// Returns the schema version before this run. async fn apply_migrations( conn: &mut PgConnection, options: &MigrateOptions, -) -> Result<(), AwaError> { +) -> Result { let has_schema: bool = sqlx::query_scalar("SELECT EXISTS(SELECT 1 FROM pg_namespace WHERE nspname = 'awa')") .fetch_one(&mut *conn) @@ -558,7 +618,7 @@ async fn apply_migrations( info!(version = current, "Schema is up to date"); } - Ok(()) + Ok(current) } /// Warm the admin-metadata cache after migrations commit. diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 339001705..dafbbaa14 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -13,7 +13,7 @@ use sqlx::{PgPool, Postgres}; use std::collections::hash_map::DefaultHasher; use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::hash::{Hash, Hasher}; -use std::sync::atomic::{AtomicU16, AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU16, AtomicUsize, Ordering}; use std::sync::Mutex; use std::time::{Duration, Instant}; use uuid::Uuid; @@ -3461,6 +3461,12 @@ pub struct QueueStorage { /// TRUNCATEs queue_lanes; `advance_enqueue_head` repairs a stale entry /// (head row gone after a rolled-back ensure_lane). ensured_lanes: Mutex>, + /// Whether `{schema}.provision_lane` (migration v053) was callable the + /// last time a lane was created. Only a positive answer is kept: lane + /// creation is rare, and the probe rides on the lane-lock round trip. + lane_provisioner_present: AtomicBool, + /// Set once a provisioner failure has been logged. + lane_provisioner_warned: AtomicBool, slot_fill_observations: Mutex>, prune_lock_timeout: Duration, cancel_phase_hook: CancelPhaseHookSlot, @@ -3516,6 +3522,8 @@ impl QueueStorage { enqueue_shards_cache: Mutex::new(HashMap::new()), claimer_gates: Mutex::new(HashMap::new()), ensured_lanes: Mutex::new(HashSet::new()), + lane_provisioner_present: AtomicBool::new(false), + lane_provisioner_warned: AtomicBool::new(false), slot_fill_observations: Mutex::new(HashMap::new()), prune_lock_timeout: DEFAULT_PRUNE_LOCK_TIMEOUT, cancel_phase_hook: CancelPhaseHookSlot(None), @@ -4594,9 +4602,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. From migration v053 any runtime role may run it: the lanes are + /// created as the schema owner. This does not change enqueue routing or + /// the configured shard count. pub async fn prepare_queue( &self, pool: &PgPool, @@ -4637,11 +4646,35 @@ impl QueueStorage { ) -> Result<(), AwaError> { let schema = self.schema(); let lane_lock_key = format!("{schema}:{queue}:{priority}:{enqueue_shard}"); - sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") - .bind(lane_lock_key) - .execute(tx.as_mut()) + let provisioner_callable = if self.lane_provisioner_present.load(Ordering::Relaxed) { + sqlx::query("SELECT pg_advisory_xact_lock(hashtextextended($1, 0))") + .bind(&lane_lock_key) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + true + } else { + sqlx::query_scalar( + r#" + SELECT COALESCE(has_function_privilege(to_regprocedure($2), 'EXECUTE'), false) + FROM (SELECT pg_advisory_xact_lock(hashtextextended($1, 0))) AS lane_lock + "#, + ) + .bind(&lane_lock_key) + .bind(format!("{schema}.provision_lane(text,smallint,smallint)")) + .fetch_one(tx.as_mut()) .await - .map_err(map_sqlx_error)?; + .map_err(map_sqlx_error)? + }; + + if provisioner_callable + && self + .provision_lane(tx, queue, priority, enqueue_shard) + .await? + { + self.mark_lane_ensured(queue, priority, enqueue_shard); + return Ok(()); + } sqlx::query(audited_sql(format!( r#" @@ -4712,6 +4745,63 @@ impl QueueStorage { Ok(()) } + /// Create the lane through `{schema}.provision_lane`, which runs as the + /// schema owner, inside a savepoint. Returns false, with the transaction + /// as it was, when the function has gone (`42883`) or the schema owner + /// may not write the lane (`42501`); the caller then writes the rows + /// itself, which is what a binary that predates the function does. + async fn provision_lane<'a>( + &self, + tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + queue: &str, + priority: i16, + enqueue_shard: i16, + ) -> Result { + let schema = self.schema(); + sqlx::query("SAVEPOINT awa_provision_lane") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let provisioned = sqlx::query(audited_sql(format!( + "SELECT {schema}.provision_lane($1, $2, $3)" + ))) + .bind(queue) + .bind(priority) + .bind(enqueue_shard) + .execute(tx.as_mut()) + .await; + match provisioned { + Ok(_) => { + sqlx::query("RELEASE SAVEPOINT awa_provision_lane") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + self.lane_provisioner_present.store(true, Ordering::Relaxed); + Ok(true) + } + Err(sqlx::Error::Database(db_err)) + if matches!(db_err.code().as_deref(), Some("42883" | "42501")) => + { + sqlx::query("ROLLBACK TO SAVEPOINT awa_provision_lane") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + self.lane_provisioner_present + .store(false, Ordering::Relaxed); + if !self.lane_provisioner_warned.swap(true, Ordering::Relaxed) { + tracing::warn!( + schema, + code = db_err.code().as_deref().unwrap_or_default(), + error = %db_err.message(), + "Lane provisioning function failed; creating lane rows as the connected role" + ); + } + Ok(false) + } + Err(err) => Err(map_sqlx_error(err)), + } + } + /// Pick the enqueue shard for a row on `queue`. /// /// Reads `awa.queue_meta.enqueue_shards` (default 1, cached diff --git a/awa-model/src/storage.rs b/awa-model/src/storage.rs index b7a0ff1ad..14770c40a 100644 --- a/awa-model/src/storage.rs +++ b/awa-model/src/storage.rs @@ -2,7 +2,7 @@ use crate::admin; use crate::error::AwaError; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; -use sqlx::{FromRow, PgExecutor, PgPool}; +use sqlx::{Connection, FromRow, PgConnection, PgExecutor, PgPool}; use std::collections::BTreeMap; #[derive(Debug, Clone, Serialize, Deserialize, FromRow, PartialEq)] @@ -469,3 +469,272 @@ pub async fn status_report(pool: &PgPool) -> Result, + pub reassign_sql: Option, + /// Login roles, superusers aside, that can run every statement above. + /// Empty when no such role exists yet. + pub repair_roles: Vec, +} + +/// What a repair did to one lane sequence. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)] +#[serde(rename_all = "snake_case")] +pub enum LaneSequenceOutcome { + /// Owned by the schema owner with every runtime grant. + Repaired, + /// Grants complete, but the role could not reassign the sequence. + OwnerUnchanged, + /// The role could neither reassign the sequence nor add its grants. + Unrepaired, + /// An open transaction held the sequence's lock past the repair's lock + /// timeout, so nothing changed. Rerun the repair. + Busy, +} + +impl TryFrom for LaneSequenceOutcome { + type Error = String; + + fn try_from(outcome: String) -> Result { + match outcome.as_str() { + "repaired" => Ok(Self::Repaired), + "owner_unchanged" => Ok(Self::OwnerUnchanged), + "unrepaired" => Ok(Self::Unrepaired), + "busy" => Ok(Self::Busy), + _ => Err(format!("unknown lane sequence repair outcome {outcome:?}")), + } + } +} + +/// What [`repair_lane_sequences`] did to one lane sequence that had a problem, +/// with what is still wrong afterwards. +#[derive(Debug, Clone, Serialize, Deserialize, FromRow, PartialEq, Eq)] +pub struct LaneSequenceRepair { + pub schema: String, + pub sequence_name: String, + /// Owner after the repair. + pub owner: String, + pub expected_owner: String, + #[sqlx(try_from = "String")] + pub outcome: LaneSequenceOutcome, + pub missing_grants: Vec, + pub reassign_sql: Option, + /// Login roles, superusers aside, that can finish the repair. + pub repair_roles: Vec, +} + +/// Lane sequences in `schema` that need repair. Read-only; any role may run it. +pub async fn lane_sequence_problems<'e, E>( + executor: E, + schema: &str, +) -> Result, AwaError> +where + E: PgExecutor<'e>, +{ + Ok(sqlx::query_as( + "SELECT $1::text AS schema, sequence_name, owner, expected_owner, missing_grants, \ + reassign_sql, repair_roles \ + FROM awa.lane_sequence_problems($1)", + ) + .bind(schema) + .fetch_all(executor) + .await?) +} + +/// Give each lane sequence in `schema` to the schema owner and add its missing +/// grants, as far as the connection's roles are allowed to. A sequence that +/// cannot be repaired is returned with the statements still to run and the +/// roles that can run them, never an error. +/// +/// Each sequence is repaired by its own statement, outside any transaction +/// the caller holds, so a reassignment's ACCESS EXCLUSIVE lock is released +/// before the next sequence and waits at most the repair function's lock +/// timeout. A sequence locked for longer by an open enqueue or claim comes +/// back [`LaneSequenceOutcome::Busy`]. +/// +/// Reassigning needs the current owner's privileges. A migration role that +/// connects as itself and sets its role to the schema owner (`role` in the +/// connection options) loses its own privileges while acting as the owner, +/// so a sequence the current role cannot repair is retried once as the login +/// role. +pub async fn repair_lane_sequences( + conn: &mut PgConnection, + schema: &str, +) -> Result, AwaError> { + let (current_role, login_role): (String, String) = + sqlx::query_as("SELECT current_user::text, session_user::text") + .fetch_one(&mut *conn) + .await?; + let problems = lane_sequence_problems(&mut *conn, schema).await?; + let mut repairs = Vec::with_capacity(problems.len()); + for problem in problems { + let mut repair: Option = sqlx::query_as(REPAIR_LANE_SEQUENCE_SQL) + .bind(schema) + .bind(&problem.sequence_name) + .fetch_optional(&mut *conn) + .await?; + let role_was_the_obstacle = matches!( + repair.as_ref().map(|repair| repair.outcome), + Some(LaneSequenceOutcome::OwnerUnchanged | LaneSequenceOutcome::Unrepaired) + ); + if role_was_the_obstacle && login_role != current_role { + if let Some(retried) = + repair_lane_sequence_as_login_role(conn, schema, &problem.sequence_name).await? + { + repair = retried; + } + } + repairs.extend(repair); + } + Ok(repairs) +} + +const REPAIR_LANE_SEQUENCE_SQL: &str = + "SELECT $1::text AS schema, sequence_name, owner, expected_owner, outcome, missing_grants, \ + reassign_sql, repair_roles \ + FROM awa.repair_lane_sequence($1, $2)"; + +/// Repair one sequence with the login role as the current role, for this +/// transaction only. `None` when the login role may not run the repair. +async fn repair_lane_sequence_as_login_role( + conn: &mut PgConnection, + schema: &str, + sequence_name: &str, +) -> Result>, AwaError> { + let mut transaction = conn.begin().await?; + sqlx::query("SET LOCAL ROLE NONE") + .execute(&mut *transaction) + .await?; + let repair = sqlx::query_as(REPAIR_LANE_SEQUENCE_SQL) + .bind(schema) + .bind(sequence_name) + .fetch_optional(&mut *transaction) + .await; + match repair { + Ok(repair) => { + transaction.commit().await?; + Ok(Some(repair)) + } + Err(sqlx::Error::Database(error)) if error.code().as_deref() == Some("42501") => { + transaction.rollback().await?; + Ok(None) + } + Err(error) => Err(error.into()), + } +} + +/// [`repair_lane_sequences`] for every queue-storage schema with a lane +/// sequence that needs repair. +pub async fn repair_lane_sequences_in_every_schema( + conn: &mut PgConnection, +) -> Result, AwaError> { + let mut repairs = Vec::new(); + for (schema, _) in lane_sequence_problem_counts(&mut *conn).await? { + repairs.extend(repair_lane_sequences(conn, &schema).await?); + } + Ok(repairs) +} + +/// One line per group of lane sequences that still need repair, saying what +/// to do next. `busy` marks sequences a repair found locked. +pub fn lane_sequence_next_steps<'a>( + leftovers: impl IntoIterator, +) -> Vec { + let mut groups: BTreeMap<(String, bool, String, String, Vec), usize> = BTreeMap::new(); + for (problem, busy) in leftovers { + *groups + .entry(( + problem.schema.clone(), + busy, + problem.owner.clone(), + problem.expected_owner.clone(), + problem.repair_roles.clone(), + )) + .or_default() += 1; + } + groups + .into_iter() + .map( + |((schema, busy, owner, expected_owner, repair_roles), count)| { + let rerun = format!("awa storage lane-sequences --schema {schema} --repair"); + if busy { + format!( + "{count} lane sequence(s) in schema {schema} owned by {owner} are busy: an \ + open transaction held the lock, so nothing changed. Rerun `{rerun}`." + ) + } else if repair_roles.is_empty() { + format!( + "{count} lane sequence(s) in schema {schema} owned by {owner} still need \ + repair, and no login role can both act as {owner} and SET ROLE to \ + {expected_owner}. Grant {owner} to the migration role, then rerun `{rerun}` \ + connected as the migration role, or run the statements from \ + `awa storage lane-sequences --schema {schema}` as {owner}." + ) + } else { + format!( + "{count} lane sequence(s) in schema {schema} owned by {owner} still need \ + repair. Rerun `{rerun}` connected as one of: {}.", + repair_roles.join(", ") + ) + } + }, + ) + .collect() +} + +impl LaneSequenceRepair { + /// The problem left after this repair, or `None` when it was repaired. + pub fn remaining_problem(&self) -> Option { + (self.outcome != LaneSequenceOutcome::Repaired).then(|| LaneSequenceProblem { + schema: self.schema.clone(), + sequence_name: self.sequence_name.clone(), + owner: self.owner.clone(), + expected_owner: self.expected_owner.clone(), + missing_grants: self.missing_grants.clone(), + reassign_sql: self.reassign_sql.clone(), + repair_roles: self.repair_roles.clone(), + }) + } +} + +/// [`lane_sequence_next_steps`] for the sequences a repair left behind. +pub fn lane_sequence_repair_next_steps(repairs: &[LaneSequenceRepair]) -> Vec { + let leftovers: Vec<(LaneSequenceProblem, bool)> = repairs + .iter() + .filter_map(|repair| { + repair + .remaining_problem() + .map(|problem| (problem, repair.outcome == LaneSequenceOutcome::Busy)) + }) + .collect(); + lane_sequence_next_steps(leftovers.iter().map(|(problem, busy)| (problem, *busy))) +} + +/// How many lane sequences need repair in each queue-storage schema that has +/// any. Read-only. +pub async fn lane_sequence_problem_counts<'e, E>( + executor: E, +) -> Result, AwaError> +where + E: PgExecutor<'e>, +{ + Ok(sqlx::query_as( + "SELECT discovered.schema_name, count(*) \ + FROM awa.queue_storage_schemas() AS discovered \ + CROSS JOIN LATERAL awa.lane_sequence_problems(discovered.schema_name) AS problem \ + GROUP BY discovered.schema_name \ + ORDER BY discovered.schema_name", + ) + .fetch_all(executor) + .await?) +} diff --git a/awa-model/tests/least_privilege_runtime_test.rs b/awa-model/tests/least_privilege_runtime_test.rs index 99e69d536..db66046a4 100644 --- a/awa-model/tests/least_privilege_runtime_test.rs +++ b/awa-model/tests/least_privilege_runtime_test.rs @@ -250,21 +250,26 @@ async fn check_runtime(pool: &PgPool, schema: &str) -> TestResult { INSERT INTO {schema}.queue_claim_heads(queue, priority, enqueue_shard) VALUES ('prepared_queue', 2, 0) ON CONFLICT DO NOTHING; SELECT {schema}.reserve_enqueue_seq('prepared_queue', 2::smallint, 0::smallint, 3); "#)))).await?; - let error = pool - .execute(sqlx::raw_sql(awa_model::audited_sql(format!( - "SELECT {schema}.ensure_lane_sequences('unprepared_queue', 2::smallint, 0::smallint)" - )))) - .await - .unwrap_err(); - assert_eq!( - error.as_database_error().unwrap().code().as_deref(), - Some("42501") - ); - assert!(error - .as_database_error() - .unwrap() - .message() - .contains("has not been provisioned")); + // An unprepared lane is created through the schema owner's provisioning + // function, so the runtime still issues no DDL of its own. + pool.execute(sqlx::raw_sql(awa_model::audited_sql(format!( + "SELECT {schema}.ensure_lane_sequences('unprepared_queue', 2::smallint, 0::smallint); \ + SELECT {schema}.provision_lane('unprepared_queue', 2::smallint, 0::smallint);" + )))) + .await?; + let foreign_owned: i64 = sqlx::query_scalar(awa_model::audited_sql(format!( + "SELECT count(*) FROM pg_class \ + WHERE relnamespace = '{schema}'::regnamespace AND relkind = 'S' \ + AND relowner <> (SELECT nspowner FROM pg_namespace WHERE nspname = '{schema}')" + ))) + .fetch_one(pool) + .await?; + assert_eq!(foreign_owned, 0); + sqlx::query_scalar::<_, i64>(awa_model::audited_sql(format!( + "SELECT {schema}.reserve_enqueue_seq('unprepared_queue', 2::smallint, 0::smallint, 1)" + ))) + .fetch_one(pool) + .await?; let error = pool .execute(sqlx::raw_sql(awa_model::audited_sql(format!( "CREATE TABLE {schema}.must_not_create(id integer)" @@ -331,10 +336,14 @@ async fn exercise_lazy_lane(options: PgConnectOptions, owner: &str, runtime: &st "GRANT USAGE, CREATE ON SCHEMA awa TO {runtime}; \ GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA awa TO {runtime}; \ GRANT USAGE, SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA awa TO {runtime}; \ - GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa TO {runtime};" + GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa TO {runtime}; \ + REVOKE EXECUTE ON FUNCTION awa.provision_lane_sequences(text, smallint, smallint) \ + FROM PUBLIC, {runtime};" ), ) .await?; + // Without the provisioning function the runtime creates the lane's + // sequences itself and owns them, as a 0.6 runtime with CREATE did. let runtime_pool = PgPoolOptions::new() .max_connections(2) .connect_with(options.username(runtime)) diff --git a/awa-python/python/awa/_awa.pyi b/awa-python/python/awa/_awa.pyi index 2219bfc3b..34f5edbab 100644 --- a/awa-python/python/awa/_awa.pyi +++ b/awa-python/python/awa/_awa.pyi @@ -605,6 +605,7 @@ class Client: deadline_rescue_interval_ms: int | None = None, callback_rescue_interval_ms: int | None = None, terminal_count_rollup_interval_ms: int | None = None, + statistics_refresh_interval_ms: int | None = None, queue_storage_schema: str | None = None, queue_storage_queue_slot_count: int = 16, queue_storage_lease_slot_count: int = 8, diff --git a/awa-python/python/awa/client.py b/awa-python/python/awa/client.py index f734ca679..131d3391b 100644 --- a/awa-python/python/awa/client.py +++ b/awa-python/python/awa/client.py @@ -922,6 +922,7 @@ async def start( deadline_rescue_interval_ms: int | None = None, callback_rescue_interval_ms: int | None = None, terminal_count_rollup_interval_ms: int | None = None, + statistics_refresh_interval_ms: int | None = None, queue_storage_schema: str | None = None, queue_storage_queue_slot_count: int = DEFAULT_QUEUE_STORAGE_QUEUE_SLOT_COUNT, queue_storage_lease_slot_count: int = DEFAULT_QUEUE_STORAGE_LEASE_SLOT_COUNT, @@ -967,6 +968,11 @@ async def start( counters. Exact reads include pending deltas, so the default favors fewer hot counter updates over immediate folding. + ``statistics_refresh_interval_ms`` controls how often the maintenance + leader runs ``ANALYZE`` on the queue-storage head tables and + partitioned ring parents, which autovacuum never analyses (default 5 + minutes; ``0`` disables the refresh). + Queue dicts also accept ``claimers`` and ``claim_batch_size`` for hot queue-storage queues. ``claimers`` adds dispatcher/claimer loops that share the queue's worker permits; benchmark ``2`` or ``4`` before @@ -988,6 +994,7 @@ async def start( deadline_rescue_interval_ms=deadline_rescue_interval_ms, callback_rescue_interval_ms=callback_rescue_interval_ms, terminal_count_rollup_interval_ms=terminal_count_rollup_interval_ms, + statistics_refresh_interval_ms=statistics_refresh_interval_ms, queue_storage_schema=queue_storage_schema, queue_storage_queue_slot_count=queue_storage_queue_slot_count, queue_storage_lease_slot_count=queue_storage_lease_slot_count, diff --git a/awa-python/src/client.rs b/awa-python/src/client.rs index f141508a7..e17b0c15e 100644 --- a/awa-python/src/client.rs +++ b/awa-python/src/client.rs @@ -1796,7 +1796,7 @@ impl PyClient { }) } - #[pyo3(signature = (queues=None, *, poll_interval_ms=200, global_max_workers=None, completed_retention_hours=None, failed_retention_hours=None, descriptor_retention_days=None, cleanup_batch_size=None, leader_election_interval_ms=None, heartbeat_interval_ms=None, promote_interval_ms=None, heartbeat_rescue_interval_ms=None, heartbeat_staleness_ms=None, deadline_rescue_interval_ms=None, callback_rescue_interval_ms=None, terminal_count_rollup_interval_ms=None, queue_storage_schema=None, queue_storage_queue_slot_count=16, queue_storage_lease_slot_count=8, queue_storage_claim_slot_count=8, queue_storage_queue_stripe_count=1, queue_storage_queue_rotate_interval_ms=1000, queue_storage_lease_rotate_interval_ms=1000, queue_storage_claim_rotate_interval_ms=None, queue_storage_rotate_min_slot_age_ms=None, queue_storage_rotate_target_slot_bytes=None, queue_storage_rotate_min_slot_rows=None, queue_storage_rotate_max_slot_age_ms=None, storage_transition_role=None))] + #[pyo3(signature = (queues=None, *, poll_interval_ms=200, global_max_workers=None, completed_retention_hours=None, failed_retention_hours=None, descriptor_retention_days=None, cleanup_batch_size=None, leader_election_interval_ms=None, heartbeat_interval_ms=None, promote_interval_ms=None, heartbeat_rescue_interval_ms=None, heartbeat_staleness_ms=None, deadline_rescue_interval_ms=None, callback_rescue_interval_ms=None, terminal_count_rollup_interval_ms=None, statistics_refresh_interval_ms=None, queue_storage_schema=None, queue_storage_queue_slot_count=16, queue_storage_lease_slot_count=8, queue_storage_claim_slot_count=8, queue_storage_queue_stripe_count=1, queue_storage_queue_rotate_interval_ms=1000, queue_storage_lease_rotate_interval_ms=1000, queue_storage_claim_rotate_interval_ms=None, queue_storage_rotate_min_slot_age_ms=None, queue_storage_rotate_target_slot_bytes=None, queue_storage_rotate_min_slot_rows=None, queue_storage_rotate_max_slot_age_ms=None, storage_transition_role=None))] #[allow(clippy::too_many_arguments)] fn start<'py>( &self, @@ -1816,6 +1816,7 @@ impl PyClient { deadline_rescue_interval_ms: Option, callback_rescue_interval_ms: Option, terminal_count_rollup_interval_ms: Option, + statistics_refresh_interval_ms: Option, queue_storage_schema: Option, queue_storage_queue_slot_count: u32, queue_storage_lease_slot_count: u32, @@ -1950,6 +1951,9 @@ impl PyClient { } builder = builder.terminal_count_rollup_interval(Duration::from_millis(ms)); } + if let Some(ms) = statistics_refresh_interval_ms { + builder = builder.statistics_refresh_interval(Duration::from_millis(ms)); + } if queue_storage_queue_slot_count == 0 { return Err(pyo3::exceptions::PyValueError::new_err( "queue_storage_queue_slot_count must be > 0", diff --git a/awa-python/tests/test_start_config.py b/awa-python/tests/test_start_config.py index 302df0fc2..4d19775d0 100644 --- a/awa-python/tests/test_start_config.py +++ b/awa-python/tests/test_start_config.py @@ -316,12 +316,26 @@ async def handle(job): deadline_rescue_interval_ms=1_000, callback_rescue_interval_ms=1_000, terminal_count_rollup_interval_ms=1_000, + statistics_refresh_interval_ms=1_000, leader_election_interval_ms=200, promote_interval_ms=500, ) await client.shutdown() +@pytest.mark.asyncio +async def test_statistics_refresh_can_be_disabled(client): + """statistics_refresh_interval_ms=0 disables the refresh; start() accepts it.""" + queue = "cfg_statistics_refresh_off" + + @client.task(ConfigTestJob, queue=queue) + async def handle(job): + return None + + await client.start([(queue, 1)], statistics_refresh_interval_ms=0) + await client.shutdown() + + @pytest.mark.asyncio async def test_per_queue_retention_in_dict_config(client): """Dict form with per-queue retention config starts successfully.""" diff --git a/awa-worker/src/client.rs b/awa-worker/src/client.rs index 253bbde7a..d32e8b977 100644 --- a/awa-worker/src/client.rs +++ b/awa-worker/src/client.rs @@ -163,6 +163,7 @@ pub struct ClientBuilder { leader_check_interval: Option, priority_aging_interval: Option, terminal_count_rollup_interval: Option, + statistics_refresh_interval: Option, completed_retention: Option, failed_retention: Option, descriptor_retention: Option, @@ -226,6 +227,7 @@ impl ClientBuilder { leader_check_interval: None, priority_aging_interval: None, terminal_count_rollup_interval: None, + statistics_refresh_interval: None, completed_retention: None, failed_retention: None, descriptor_retention: None, @@ -905,6 +907,20 @@ impl ClientBuilder { self } + /// Set how often the maintenance leader refreshes planner statistics for + /// the queue-storage tables autovacuum never analyses (default: 5 minutes). + /// + /// Autovacuum does not analyse partitioned tables, and never reaches its + /// threshold on the two small head tables, so the leader runs `ANALYZE` + /// on both when it is elected and then every interval, through + /// `awa.analyze_queue_storage_relation()` (schema v053) so that a runtime + /// role that does not own the tables still works. `Duration::ZERO` + /// disables the refresh. Queue storage only. + pub fn statistics_refresh_interval(mut self, interval: Duration) -> Self { + self.statistics_refresh_interval = Some(interval); + self + } + /// Set how often queue depth/lag metrics are published (default: 30s). pub fn queue_stats_interval(mut self, interval: Duration) -> Self { self.queue_stats_interval = Some(interval); @@ -1256,6 +1272,7 @@ impl ClientBuilder { leader_check_interval: self.leader_check_interval, priority_aging_interval: self.priority_aging_interval, terminal_count_rollup_interval: self.terminal_count_rollup_interval, + statistics_refresh_interval: self.statistics_refresh_interval, completed_retention: self.completed_retention, failed_retention: self.failed_retention, descriptor_retention: self.descriptor_retention, @@ -1362,6 +1379,7 @@ pub struct Client { leader_check_interval: Option, priority_aging_interval: Option, terminal_count_rollup_interval: Option, + statistics_refresh_interval: Option, completed_retention: Option, failed_retention: Option, descriptor_retention: Option, @@ -1932,6 +1950,9 @@ impl Client { if let Some(interval) = self.terminal_count_rollup_interval { maintenance = maintenance.terminal_count_rollup_interval(interval); } + if let Some(interval) = self.statistics_refresh_interval { + maintenance = maintenance.statistics_refresh_interval(interval); + } if let Some(retention) = self.completed_retention { maintenance = maintenance.completed_retention(retention); } diff --git a/awa-worker/src/lib.rs b/awa-worker/src/lib.rs index 4768e59bb..f9a171104 100644 --- a/awa-worker/src/lib.rs +++ b/awa-worker/src/lib.rs @@ -17,6 +17,7 @@ pub mod metrics; mod notify_hub; mod output; mod runtime; +mod statistics; mod storage; // Re-exports diff --git a/awa-worker/src/maintenance.rs b/awa-worker/src/maintenance.rs index d84960d37..782ca622b 100644 --- a/awa-worker/src/maintenance.rs +++ b/awa-worker/src/maintenance.rs @@ -1,5 +1,6 @@ use crate::executor::DlqPolicy; use crate::runtime::InFlightMap; +use crate::statistics::{RefreshGuard, StatisticsRefresher}; use crate::storage::{QueueStorageRuntime, RuntimeStorage}; use awa_model::audited_sql; use awa_model::cron::{ @@ -738,6 +739,9 @@ pub struct MaintenanceService { priority_aging_interval: Duration, batch_operations_interval: Duration, terminal_count_rollup_interval: Duration, + /// How often the leader refreshes planner statistics for the queue-storage + /// tables autovacuum never analyses. Zero disables the refresh. + statistics_refresh_interval: Duration, /// How long a descriptor catalog row can sit without being refreshed /// before the maintenance leader deletes it. Zero disables cleanup. /// Default: 30 days. @@ -757,6 +761,8 @@ pub struct MaintenanceService { ring_authority_flip_aware_since: std::sync::Mutex>, } +/// Default cadence of the planner-statistics refresh on the maintenance leader. +pub(crate) const DEFAULT_STATISTICS_REFRESH_INTERVAL: Duration = Duration::from_secs(300); const PROMOTE_BATCH_SIZE: i64 = 4_096; const PROMOTE_MAX_BATCHES_PER_TICK: usize = 32; const CRON_CATCH_UP_LIMIT: usize = 1_000; @@ -817,6 +823,7 @@ impl MaintenanceService { priority_aging_interval: Duration::from_secs(60), batch_operations_interval: Duration::from_secs(1), terminal_count_rollup_interval: Duration::from_secs(30), + statistics_refresh_interval: DEFAULT_STATISTICS_REFRESH_INTERVAL, descriptor_retention: Duration::from_secs(30 * 86400), // 30d ring_authority_auto_flip: true, ring_authority_auto_flip_stable_period: Duration::from_secs(600), // 10min @@ -873,6 +880,14 @@ impl MaintenanceService { self } + /// Set how often the leader runs `ANALYZE` on the queue-storage tables + /// autovacuum never analyses (default: 5 minutes). `Duration::ZERO` + /// disables the refresh. + pub fn statistics_refresh_interval(mut self, interval: Duration) -> Self { + self.statistics_refresh_interval = interval; + self + } + /// How long a descriptor catalog row can go without being re-synced /// before the maintenance leader deletes it (default: 30 days). Set /// to `Duration::ZERO` to disable — useful if you maintain the catalog @@ -1121,6 +1136,7 @@ impl MaintenanceService { cron_eval_cancel.clone(), self.cron_eval_interval, )); + let _statistics_refresh = self.spawn_statistics_refresh(); loop { tokio::select! { @@ -1314,6 +1330,28 @@ impl MaintenanceService { } } + /// Start the planner-statistics refresh for this leadership term. The + /// returned guard stops it when dropped. + fn spawn_statistics_refresh(&self) -> Option { + if self.statistics_refresh_interval.is_zero() { + return None; + } + let (runtime, standby) = match self.storage.queue_storage() { + Some(runtime) => (runtime, false), + None => (self.standby_queue_storage.as_ref()?, true), + }; + Some(RefreshGuard::spawn( + StatisticsRefresher::new( + self.pool.clone(), + runtime.store.schema().to_string(), + self.metrics.clone(), + standby, + ), + self.statistics_refresh_interval, + &self.cancel, + )) + } + fn stop_cron_evaluator(cancel: &CancellationToken, task: &JoinHandle<()>) { cancel.cancel(); task.abort(); diff --git a/awa-worker/src/statistics.rs b/awa-worker/src/statistics.rs new file mode 100644 index 000000000..566501832 --- /dev/null +++ b/awa-worker/src/statistics.rs @@ -0,0 +1,651 @@ +//! Planner-statistics refresh for the queue-storage tables autovacuum never +//! analyses: the two head tables, which stay below the autoanalyze threshold, +//! and the partitioned ring parents, which autovacuum skips entirely. +//! +//! The maintenance leader runs one pass on election and one every interval. +//! Each relation is analysed through `awa.analyze_queue_storage_relation` +//! (migration v054, `SECURITY DEFINER`) in its own short transaction with a +//! bounded `lock_timeout` and `statement_timeout`, so the refresh never holds +//! locks across relations and cannot wait on ring prune for long. + +use crate::metrics::AwaMetrics; +use sqlx::PgPool; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::{Duration, Instant}; +use tokio_util::sync::CancellationToken; +use tracing::{debug, warn}; + +/// How long one relation's ANALYZE may wait for a lock before the pass moves +/// on. Ring prune holds `ACCESS EXCLUSIVE` on a slot's children only +/// briefly, so a wait this long means something else holds the table. +const LOCK_TIMEOUT: Duration = Duration::from_secs(2); + +/// Upper bound for one relation. A ring parent analyses in under a few +/// seconds even on a loaded host; this only stops a stuck statement. +const STATEMENT_TIMEOUT: Duration = Duration::from_secs(120); + +const HEAD_TABLES: [&str; 2] = ["queue_claim_heads", "queue_enqueue_heads"]; + +/// The partitioned ring parents a queue-storage substrate installs. A schema +/// can hold other partitioned tables (a substrate in an application schema), +/// which this refresh must never analyse. +const RING_PARENTS: [&str; 13] = [ + "done_entries", + "lease_claim_batches", + "lease_claim_closure_batches", + "lease_claim_closures", + "lease_claims", + "leases", + "queue_terminal_count_deltas", + "ready_claim_attempt_batches", + "ready_entries", + "ready_segments", + "ready_tombstones", + "receipt_completion_batches", + "receipt_completion_tombstones", +]; + +const SQLSTATE_INSUFFICIENT_PRIVILEGE: &str = "42501"; +const SQLSTATE_UNDEFINED_FUNCTION: &str = "42883"; +const SQLSTATE_LOCK_NOT_AVAILABLE: &str = "55P03"; + +/// Which kind of relation an analyse call covered; bounds metric cardinality. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum RelationScope { + Heads, + RingParent, + /// The pass as a whole: a probe or listing failed, or the helper is absent. + Pass, +} + +impl RelationScope { + pub(crate) fn as_str(self) -> &'static str { + match self { + RelationScope::Heads => "heads", + RelationScope::RingParent => "ring_parent", + RelationScope::Pass => "pass", + } + } +} + +/// Result of analysing one relation. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum AnalyzeOutcome { + Ok, + /// The function owner (or the runtime role, if it cannot execute the + /// function) may not analyse the relation. + SkippedPrivilege, + /// `lock_timeout` expired waiting on a conflicting lock. + SkippedLock, + /// `awa.analyze_queue_storage_relation` is not installed (schema older + /// than v054). + Unavailable, + Failed, +} + +impl AnalyzeOutcome { + pub(crate) fn as_str(self) -> &'static str { + match self { + AnalyzeOutcome::Ok => "ok", + AnalyzeOutcome::SkippedPrivilege => "skipped_privilege", + AnalyzeOutcome::SkippedLock => "skipped_lock", + AnalyzeOutcome::Unavailable => "unavailable", + AnalyzeOutcome::Failed => "failed", + } + } +} + +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub(crate) struct PassSummary { + pub(crate) ok: usize, + pub(crate) skipped_privilege: usize, + pub(crate) skipped_lock: usize, + pub(crate) unavailable: usize, + pub(crate) failed: usize, +} + +impl PassSummary { + fn record(&mut self, outcome: AnalyzeOutcome) { + match outcome { + AnalyzeOutcome::Ok => self.ok += 1, + AnalyzeOutcome::SkippedPrivilege => self.skipped_privilege += 1, + AnalyzeOutcome::SkippedLock => self.skipped_lock += 1, + AnalyzeOutcome::Unavailable => self.unavailable += 1, + AnalyzeOutcome::Failed => self.failed += 1, + } + } +} + +#[derive(Debug, PartialEq, Eq)] +enum HelperAccess { + Available, + Missing, + Denied, +} + +/// Refreshes planner statistics for one queue-storage schema. +pub(crate) struct StatisticsRefresher { + pool: PgPool, + schema: String, + metrics: AwaMetrics, + /// Set for a leader whose own storage is canonical and whose queue-storage + /// runtime is only a standby: a pass runs only once routing has flipped to + /// `schema`. + require_active_schema: bool, + warned_caller_privilege: AtomicBool, + warned_analyze_privilege: AtomicBool, + warned_unavailable: AtomicBool, +} + +impl StatisticsRefresher { + pub(crate) fn new( + pool: PgPool, + schema: String, + metrics: AwaMetrics, + require_active_schema: bool, + ) -> Self { + Self { + pool, + schema, + metrics, + require_active_schema, + warned_caller_privilege: AtomicBool::new(false), + warned_analyze_privilege: AtomicBool::new(false), + warned_unavailable: AtomicBool::new(false), + } + } + + /// Run a pass now and then every `interval` until `cancel` fires. + pub(crate) async fn run(&self, interval: Duration, cancel: CancellationToken) { + let mut timer = tokio::time::interval(interval); + timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + tokio::select! { + _ = cancel.cancelled() => return, + _ = timer.tick() => { + tokio::select! { + _ = cancel.cancelled() => return, + _ = self.run_pass() => {} + } + } + } + } + } + + /// Analyse the head tables, then every partitioned table of the schema. + pub(crate) async fn run_pass(&self) -> PassSummary { + let started = Instant::now(); + let mut summary = PassSummary::default(); + + if self.require_active_schema { + match sqlx::query_scalar::<_, Option>( + "SELECT awa.active_queue_storage_schema()", + ) + .fetch_one(&self.pool) + .await + { + Ok(active) if active.as_deref() == Some(self.schema.as_str()) => {} + Ok(_) => { + debug!(schema = %self.schema, "Routing has not flipped to this schema yet; skipping statistics refresh"); + return summary; + } + Err(err) => { + warn!(error = %err, "Could not resolve the active queue-storage schema for statistics refresh"); + return summary; + } + } + } + + match awa_model::storage::queue_storage_schema_ready(&self.pool, &self.schema).await { + Ok(true) => {} + Ok(false) => { + debug!(schema = %self.schema, "Queue-storage schema is not prepared yet; skipping statistics refresh"); + return summary; + } + Err(err) => { + warn!(error = %err, schema = %self.schema, "Could not check the queue-storage schema for statistics refresh"); + self.record( + RelationScope::Pass, + AnalyzeOutcome::Failed, + started.elapsed(), + ); + summary.record(AnalyzeOutcome::Failed); + return summary; + } + } + + match self.helper_access().await { + Ok(HelperAccess::Available) => {} + Ok(HelperAccess::Missing) => { + self.warn_unavailable_once(); + self.record( + RelationScope::Pass, + AnalyzeOutcome::Unavailable, + started.elapsed(), + ); + summary.record(AnalyzeOutcome::Unavailable); + return summary; + } + Ok(HelperAccess::Denied) => { + self.warn_privilege_once( + &self.warned_caller_privilege, + "queue_claim_heads", + "the runtime role needs EXECUTE on awa.analyze_queue_storage_relation(text, text), \ + and the login role needs INSERT, UPDATE and DELETE on this schema's queue_claim_heads; \ + restore those grants as the schema owner", + ); + self.record( + RelationScope::Pass, + AnalyzeOutcome::SkippedPrivilege, + started.elapsed(), + ); + summary.record(AnalyzeOutcome::SkippedPrivilege); + return summary; + } + Err(err) => { + warn!(error = %err, schema = %self.schema, "Could not probe access to the statistics helper"); + self.record( + RelationScope::Pass, + AnalyzeOutcome::Failed, + started.elapsed(), + ); + summary.record(AnalyzeOutcome::Failed); + return summary; + } + } + + let ring_parents = match self.ring_parents().await { + Ok(parents) => parents, + Err(err) => { + warn!(error = %err, schema = %self.schema, "Could not list ring tables for statistics refresh"); + self.record( + RelationScope::Pass, + AnalyzeOutcome::Failed, + started.elapsed(), + ); + summary.record(AnalyzeOutcome::Failed); + Vec::new() + } + }; + + let targets = HEAD_TABLES + .iter() + .map(|name| (RelationScope::Heads, (*name).to_string())) + .chain( + ring_parents + .into_iter() + .map(|name| (RelationScope::RingParent, name)), + ); + + for (scope, relation) in targets { + let relation_started = Instant::now(); + let outcome = self.analyze_relation(&relation).await; + self.record(scope, outcome, relation_started.elapsed()); + summary.record(outcome); + } + + debug!( + schema = %self.schema, + elapsed_ms = started.elapsed().as_millis() as u64, + ok = summary.ok, + skipped_privilege = summary.skipped_privilege, + skipped_lock = summary.skipped_lock, + unavailable = summary.unavailable, + failed = summary.failed, + "Planner statistics refresh finished" + ); + summary + } + + fn record(&self, scope: RelationScope, outcome: AnalyzeOutcome, elapsed: Duration) { + self.metrics + .record_statistics_analyze(scope.as_str(), outcome.as_str(), elapsed); + } + + /// Check the helper and its caller privileges once per pass, so a missing + /// migration or grant does not produce a failing call for every relation. + async fn helper_access(&self) -> Result { + let (installed, allowed): (bool, bool) = sqlx::query_as( + r#" + SELECT helper IS NOT NULL, + COALESCE( + has_function_privilege(helper, 'EXECUTE') + AND has_table_privilege(session_user, heads, 'INSERT') + AND has_table_privilege(session_user, heads, 'UPDATE') + AND has_table_privilege(session_user, heads, 'DELETE'), + false + ) + FROM ( + SELECT to_regprocedure('awa.analyze_queue_storage_relation(text,text)') AS helper, + to_regclass(format('%I.queue_claim_heads', $1::text)) AS heads + ) AS objects + "#, + ) + .bind(&self.schema) + .fetch_one(&self.pool) + .await?; + Ok(match (installed, allowed) { + (false, _) => HelperAccess::Missing, + (true, false) => HelperAccess::Denied, + (true, true) => HelperAccess::Available, + }) + } + + async fn ring_parents(&self) -> Result, sqlx::Error> { + sqlx::query_scalar( + r#" + SELECT class.relname::text + FROM pg_class AS class + JOIN pg_namespace AS namespace ON namespace.oid = class.relnamespace + WHERE namespace.nspname = $1 AND class.relkind = 'p' + AND class.relname = ANY($2) + ORDER BY class.relname + "#, + ) + .bind(&self.schema) + .bind(&RING_PARENTS[..]) + .fetch_all(&self.pool) + .await + } + + async fn analyze_relation(&self, relation: &str) -> AnalyzeOutcome { + match self.call_analyze_function(relation).await { + Ok(status) => match status.as_str() { + "analyzed" => AnalyzeOutcome::Ok, + "skipped_privilege" => { + self.warn_privilege_once( + &self.warned_analyze_privilege, + relation, + "the owner of awa.analyze_queue_storage_relation() may not ANALYZE this table; \ + change the function's owner to the table owner with ALTER FUNCTION, \ + or grant the function owner MAINTAIN on them (PostgreSQL 17+)", + ); + AnalyzeOutcome::SkippedPrivilege + } + // The relation is not part of this schema's shape; nothing to analyse. + "missing" => AnalyzeOutcome::Ok, + other => { + warn!( + status = other, + relation, "Unexpected statistics refresh status" + ); + AnalyzeOutcome::Failed + } + }, + Err(err) => self.classify_error(relation, &err), + } + } + + async fn call_analyze_function(&self, relation: &str) -> Result { + let mut tx = self.pool.begin().await?; + sqlx::query( + "SELECT set_config('lock_timeout', $1, true), set_config('statement_timeout', $2, true)", + ) + .bind(format!("{}ms", LOCK_TIMEOUT.as_millis())) + .bind(format!("{}ms", STATEMENT_TIMEOUT.as_millis())) + .execute(&mut *tx) + .await?; + let status: String = + sqlx::query_scalar("SELECT awa.analyze_queue_storage_relation($1, $2)") + .bind(&self.schema) + .bind(relation) + .fetch_one(&mut *tx) + .await?; + tx.commit().await?; + Ok(status) + } + + fn classify_error(&self, relation: &str, err: &sqlx::Error) -> AnalyzeOutcome { + let code = match err { + sqlx::Error::Database(db_err) => db_err.code().map(|code| code.into_owned()), + _ => None, + }; + match code.as_deref() { + Some(SQLSTATE_LOCK_NOT_AVAILABLE) => { + debug!( + relation, + "Statistics refresh timed out waiting for a lock; retrying next interval" + ); + AnalyzeOutcome::SkippedLock + } + Some(SQLSTATE_UNDEFINED_FUNCTION) => { + self.warn_unavailable_once(); + AnalyzeOutcome::Unavailable + } + Some(SQLSTATE_INSUFFICIENT_PRIVILEGE) => { + self.warn_privilege_once( + &self.warned_analyze_privilege, + relation, + "the statistics helper was denied access; check the helper's owner, table ownership \ + and ANALYZE privileges; also recheck runtime helper EXECUTE and login-role \ + head-table INSERT, UPDATE and DELETE grants if they changed after this pass's privilege probe", + ); + AnalyzeOutcome::SkippedPrivilege + } + _ => { + warn!(relation, error = %err, "Planner statistics refresh failed for relation"); + AnalyzeOutcome::Failed + } + } + } + + fn warn_unavailable_once(&self) { + if self.warned_unavailable.swap(true, Ordering::Relaxed) { + return; + } + warn!( + schema = %self.schema, + "Planner statistics refresh is disabled: awa.analyze_queue_storage_relation() \ + is missing. Run `awa migrate` (schema v054 or later) to install it" + ); + } + + fn warn_privilege_once(&self, warned: &AtomicBool, relation: &str, missing: &str) { + if warned.swap(true, Ordering::Relaxed) { + return; + } + warn!( + schema = %self.schema, + relation, + "Planner statistics refresh skipped: {missing}. Set the statistics refresh \ + interval to zero to disable the refresh" + ); + } +} + +/// Stops the refresh task when the maintenance leader loses leadership or +/// shuts down. +pub(crate) struct RefreshGuard { + cancel: CancellationToken, +} + +impl RefreshGuard { + pub(crate) fn spawn( + refresher: StatisticsRefresher, + interval: Duration, + parent: &CancellationToken, + ) -> Self { + let cancel = parent.child_token(); + let task_cancel = cancel.clone(); + tokio::spawn(async move { + debug!( + schema = %refresher.schema, + interval_s = interval.as_secs(), + "Planner statistics refresh started" + ); + refresher.run(interval, task_cancel).await; + }); + Self { cancel } + } +} + +impl Drop for RefreshGuard { + fn drop(&mut self) { + self.cancel.cancel(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; + use std::str::FromStr; + + /// A binary ahead of the schema must probe once and attempt no calls: a + /// failing call per relation would count as `unavailable` fifteen times + /// and write fifteen server-side ERROR lines per pass. + #[tokio::test] + async fn missing_helper_costs_one_probe_and_no_calls() { + let url = std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()); + let options = PgConnectOptions::from_str(&url).expect("parse DATABASE_URL"); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect_with(options.clone()) + .await + .expect("connect admin"); + let database = format!("awa_stats_unit_{}", std::process::id()); + sqlx::raw_sql(awa_model::audited_sql(format!( + "DROP DATABASE IF EXISTS {database} WITH (FORCE)" + ))) + .execute(&admin) + .await + .expect("drop leftover database"); + sqlx::raw_sql(awa_model::audited_sql(format!( + "CREATE DATABASE {database}" + ))) + .execute(&admin) + .await + .expect("create database"); + let pool = PgPoolOptions::new() + .max_connections(2) + .connect_with(options.database(&database)) + .await + .expect("connect test database"); + awa_model::migrations::run(&pool).await.expect("migrate"); + + let refresher = StatisticsRefresher::new( + pool.clone(), + "awa".to_string(), + AwaMetrics::from_global(), + false, + ); + let present = refresher.run_pass().await; + assert_eq!(present.ok, 15, "{present:?}"); + + sqlx::raw_sql("DROP FUNCTION awa.analyze_queue_storage_relation(text, text)") + .execute(&pool) + .await + .expect("drop the helper"); + let absent = refresher.run_pass().await; + assert_eq!( + absent, + PassSummary { + unavailable: 1, + ..PassSummary::default() + } + ); + pool.close().await; + sqlx::raw_sql(awa_model::audited_sql(format!( + "DROP DATABASE IF EXISTS {database} WITH (FORCE)" + ))) + .execute(&admin) + .await + .expect("drop database"); + } + + /// Missing EXECUTE or login-role DML is a single skipped pass, rather + /// than fifteen server-side permission errors. + #[tokio::test] + async fn missing_helper_privileges_skip_the_pass_without_calling_it() { + let url = std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()); + let options = PgConnectOptions::from_str(&url).expect("parse DATABASE_URL"); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect_with(options.clone()) + .await + .expect("connect admin"); + let database = format!("awa_stats_access_unit_{}", std::process::id()); + let role = format!("awa_stats_access_{}", std::process::id()); + sqlx::raw_sql(awa_model::audited_sql(format!( + "DROP DATABASE IF EXISTS {database} WITH (FORCE)" + ))) + .execute(&admin) + .await + .expect("drop leftover database"); + sqlx::raw_sql(awa_model::audited_sql(format!("DROP ROLE IF EXISTS {role}; CREATE ROLE {role} LOGIN PASSWORD 'statistics_access_test'"))) + .execute(&admin) + .await + .expect("create test role"); + sqlx::raw_sql(awa_model::audited_sql(format!( + "CREATE DATABASE {database}" + ))) + .execute(&admin) + .await + .expect("create database"); + let owner = PgPoolOptions::new() + .max_connections(2) + .connect_with(options.clone().database(&database)) + .await + .expect("connect owner"); + awa_model::migrations::run(&owner).await.expect("migrate"); + sqlx::raw_sql(awa_model::audited_sql(format!( + "GRANT USAGE ON SCHEMA awa TO {role}; \ + GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA awa TO {role}; \ + REVOKE EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) FROM {role}" + ))) + .execute(&owner) + .await + .expect("grant head-table privileges only"); + let runtime = PgPoolOptions::new() + .max_connections(2) + .connect_with( + options + .database(&database) + .username(&role) + .password("statistics_access_test"), + ) + .await + .expect("connect runtime"); + let refresher = StatisticsRefresher::new( + runtime.clone(), + "awa".to_string(), + AwaMetrics::from_global(), + false, + ); + let denied = PassSummary { + skipped_privilege: 1, + ..PassSummary::default() + }; + assert_eq!(refresher.run_pass().await, denied); + sqlx::raw_sql(awa_model::audited_sql(format!( + "GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) TO {role}" + ))) + .execute(&owner) + .await + .expect("grant helper EXECUTE"); + assert_eq!(refresher.run_pass().await.ok, 15); + sqlx::raw_sql(awa_model::audited_sql(format!( + "REVOKE INSERT ON awa.queue_claim_heads FROM {role}" + ))) + .execute(&owner) + .await + .expect("revoke login-role head privilege"); + assert_eq!(refresher.run_pass().await, denied); + + runtime.close().await; + owner.close().await; + sqlx::raw_sql(awa_model::audited_sql(format!( + "DROP DATABASE {database} WITH (FORCE)" + ))) + .execute(&admin) + .await + .expect("drop database"); + sqlx::raw_sql(awa_model::audited_sql(format!("DROP ROLE {role}"))) + .execute(&admin) + .await + .expect("drop test role"); + admin.close().await; + } +} diff --git a/awa/tests/fixtures/v052_lane_sequence_helpers.sql b/awa/tests/fixtures/v052_lane_sequence_helpers.sql new file mode 100644 index 000000000..b7d396dad --- /dev/null +++ b/awa/tests/fixtures/v052_lane_sequence_helpers.sql @@ -0,0 +1,190 @@ +-- The lane-sequence helpers as migration v052 left them, before lane +-- provisioning functions existed. Installed under a separate name so a test +-- can put a schema back in its v052 shape. +CREATE OR REPLACE FUNCTION awa.v052_refresh_lane_sequence_helpers(p_schema TEXT) +RETURNS VOID +LANGUAGE plpgsql +SECURITY INVOKER +SET search_path = pg_catalog +AS $refresh$ +BEGIN + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.ensure_lane_sequences( + p_queue TEXT, + p_priority SMALLINT, + p_enqueue_shard SMALLINT + ) + RETURNS VOID + LANGUAGE plpgsql + SET search_path = pg_catalog + AS $func$ + DECLARE + v_enqueue_seq TEXT := %1$I.queue_lane_sequence_name( + 'queue_enqueue_seq', + p_queue, + p_priority, + p_enqueue_shard + ); + v_claim_seq TEXT := %1$I.queue_lane_sequence_name( + 'queue_claim_seq', + p_queue, + p_priority, + p_enqueue_shard + ); + BEGIN + IF to_regclass(format('%%I.%%I', %1$L, v_enqueue_seq)) IS NULL THEN + IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_enqueue_seq + USING ERRCODE = '42501', + HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + END IF; + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_enqueue_seq + ); + END IF; + IF to_regclass(format('%%I.%%I', %1$L, v_claim_seq)) IS NULL THEN + IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_claim_seq + USING ERRCODE = '42501', + HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + END IF; + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_claim_seq + ); + END IF; + + UPDATE %1$I.queue_enqueue_heads + SET seq_name = v_enqueue_seq + WHERE queue = p_queue + AND priority = p_priority + AND enqueue_shard = p_enqueue_shard + AND seq_name IS DISTINCT FROM v_enqueue_seq; + + UPDATE %1$I.queue_claim_heads + SET seq_name = v_claim_seq + WHERE queue = p_queue + AND priority = p_priority + AND enqueue_shard = p_enqueue_shard + AND seq_name IS DISTINCT FROM v_claim_seq; + END; + $func$ + $ddl$, + p_schema + ); + + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.queue_enqueue_head_sequence_sync() + RETURNS TRIGGER + LANGUAGE plpgsql + SET search_path = pg_catalog + AS $func$ + DECLARE + v_seq_name TEXT := %1$I.queue_lane_sequence_name( + 'queue_enqueue_seq', + NEW.queue, + NEW.priority, + NEW.enqueue_shard + ); + v_qualified_seq TEXT := format('%%I.%%I', %1$L, v_seq_name); + v_count BIGINT; + v_start BIGINT; + BEGIN + IF to_regclass(format('%%I.%%I', %1$L, v_seq_name)) IS NULL THEN + IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_seq_name + USING ERRCODE = '42501', + HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + END IF; + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_seq_name + ); + END IF; + NEW.seq_name := v_seq_name; + + IF TG_OP = 'UPDATE' + AND NEW.next_seq IS DISTINCT FROM OLD.next_seq + AND NEW.next_seq > OLD.next_seq + THEN + v_count := NEW.next_seq - OLD.next_seq; + PERFORM pg_advisory_lock(hashtextextended(v_qualified_seq, 0)); + BEGIN + EXECUTE format( + 'SELECT nextval(%%L::regclass)::bigint', + v_qualified_seq + ) + INTO v_start; + + IF v_count > 1 THEN + EXECUTE format( + 'SELECT setval(%%L::regclass, %%s, true)', + v_qualified_seq, + v_start + v_count - 1 + ); + END IF; + EXCEPTION WHEN OTHERS THEN + PERFORM pg_advisory_unlock(hashtextextended(v_qualified_seq, 0)); + RAISE; + END; + PERFORM pg_advisory_unlock(hashtextextended(v_qualified_seq, 0)); + NEW.next_seq := v_start + v_count; + ELSE + PERFORM %1$I.set_sequence_next(v_seq_name, NEW.next_seq); + END IF; + + RETURN NEW; + END; + $func$ + $ddl$, + p_schema + ); + + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.queue_claim_head_sequence_sync() + RETURNS TRIGGER + LANGUAGE plpgsql + SET search_path = pg_catalog + AS $func$ + DECLARE + v_seq_name TEXT := %1$I.queue_lane_sequence_name( + 'queue_claim_seq', + NEW.queue, + NEW.priority, + NEW.enqueue_shard + ); + BEGIN + IF to_regclass(format('%%I.%%I', %1$L, v_seq_name)) IS NULL THEN + IF NOT has_schema_privilege(current_user, %1$L, 'CREATE') THEN + RAISE EXCEPTION 'lane sequence %%.%% has not been provisioned', %1$L, v_seq_name + USING ERRCODE = '42501', + HINT = 'Run awa storage prepare-queue as the migrator before starting producers or workers.'; + END IF; + EXECUTE format( + 'CREATE SEQUENCE IF NOT EXISTS %%I.%%I AS bigint START WITH 1 MINVALUE 1 CACHE 1', + %1$L, + v_seq_name + ); + END IF; + NEW.seq_name := v_seq_name; + + IF TG_OP = 'INSERT' THEN + PERFORM %1$I.set_sequence_next(v_seq_name, NEW.claim_seq); + END IF; + + RETURN NEW; + END; + $func$ + $ddl$, + p_schema + ); + +END; +$refresh$; diff --git a/awa/tests/lane_ownership_roles_test.rs b/awa/tests/lane_ownership_roles_test.rs new file mode 100644 index 000000000..f7ad8eda7 --- /dev/null +++ b/awa/tests/lane_ownership_roles_test.rs @@ -0,0 +1,833 @@ +//! Queue lanes under a split-role install: a schema owner, a migration role +//! that is a member of it, and several application roles that share a runtime +//! group role. Only the owner and the migration role have default privileges, +//! so a lane sequence is usable by every runtime role only if it is created as +//! the owner or granted explicitly, whichever role touches the lane first. + +use awa::model::storage::LaneSequenceOutcome; +use awa::model::{migrations, storage, AwaError, InsertOpts, InsertParams, QueueStorage}; +use awa::model::{JobState, QueueStorageConfig}; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; +use sqlx::{Executor, PgPool}; +use std::str::FromStr; +use std::time::Duration; + +type TestResult = Result<(), Box>; + +fn sql(text: impl Into) -> sqlx::AssertSqlSafe { + awa::audited_sql(text) +} + +/// Roles, a database owned by the schema owner, and a pool per role. +struct Cluster { + admin: PgPool, + database: String, + roles: Vec, + owner_name: String, + migrator_name: String, + runtime_name: String, + connect_options: PgConnectOptions, + owner: PgPool, + migrator: PgPool, + app_a: PgPool, + app_b: PgPool, + app_c: PgPool, +} + +#[derive(Clone, Copy, PartialEq)] +enum Schema { + /// Every migration, including the lane provisioning functions. + Current, + /// The release before lane provisioning (v052). + BeforeProvisioning, +} + +impl Cluster { + /// `runtime_create` grants the runtime group CREATE on the schema, as a + /// 0.6 install does so runtimes can create lane sequences themselves. + async fn new( + schema: Schema, + runtime_create: bool, + ) -> Result> { + let url = std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()); + let admin = PgPoolOptions::new() + .max_connections(2) + .connect(&url) + .await?; + let suffix = uuid::Uuid::new_v4().simple().to_string()[..12].to_string(); + let owner_name = format!("lane_owner_{suffix}"); + let migrator_name = format!("lane_migrator_{suffix}"); + let runtime_name = format!("lane_runtime_{suffix}"); + let app_names = ["a", "b", "c"].map(|app| format!("lane_app_{app}_{suffix}")); + let database = format!("lane_roles_{suffix}"); + admin + .execute(sqlx::raw_sql(sql(format!( + r#" + CREATE ROLE {owner_name} LOGIN PASSWORD 'lane_test'; + CREATE ROLE {migrator_name} LOGIN PASSWORD 'lane_test' IN ROLE {owner_name}; + CREATE ROLE {runtime_name} NOLOGIN; + CREATE ROLE {a} LOGIN PASSWORD 'lane_test' IN ROLE {runtime_name}; + CREATE ROLE {b} LOGIN PASSWORD 'lane_test' IN ROLE {runtime_name}; + CREATE ROLE {c} LOGIN PASSWORD 'lane_test' IN ROLE {runtime_name}; + "#, + a = app_names[0], + b = app_names[1], + c = app_names[2], + )))) + .await?; + admin + .execute(sqlx::raw_sql(sql(format!( + "CREATE DATABASE {database} OWNER {owner_name}" + )))) + .await?; + let mut roles = vec![ + owner_name.clone(), + migrator_name.clone(), + runtime_name.clone(), + ]; + roles.extend(app_names.iter().cloned()); + let options = PgConnectOptions::from_str(&url)? + .database(&database) + .password("lane_test"); + let pool = |user: &str| { + PgPoolOptions::new() + .max_connections(3) + .connect_with(options.clone().username(user)) + }; + let owner = pool(&owner_name).await?; + let migrator = pool(&migrator_name).await?; + + owner + .execute(sqlx::raw_sql(sql(format!( + r#" + CREATE SCHEMA awa; + GRANT USAGE ON SCHEMA awa TO {runtime_name}; + "# + )))) + .await?; + for (creator, creator_pool) in [(&owner_name, &owner), (&migrator_name, &migrator)] { + creator_pool + .execute(sqlx::raw_sql(sql(format!( + r#" + ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa + GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON TABLES TO {runtime_name}; + ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa + GRANT USAGE, SELECT, UPDATE ON SEQUENCES TO {runtime_name}; + ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa + GRANT EXECUTE ON FUNCTIONS TO {runtime_name}; + "# + )))) + .await?; + } + // The migration role runs every migration with the owner as its + // effective role, as the database-roles guide describes. + let migrate_as_owner = PgPoolOptions::new() + .max_connections(1) + .connect_with( + options + .clone() + .username(&migrator_name) + .options([("role", owner_name.as_str())]), + ) + .await?; + let target = match schema { + Schema::Current => migrations::CURRENT_VERSION, + Schema::BeforeProvisioning => 52, + }; + for (_, _, statement) in migrations::migration_sql_range(0, target) { + migrate_as_owner + .execute(sqlx::raw_sql(sql(statement))) + .await?; + } + if schema == Schema::BeforeProvisioning { + // Every migration bundle carries the current helper definitions, + // so put the v052 helpers back and drop what v052 did not have. + migrate_as_owner + .execute(sqlx::raw_sql(sql(include_str!( + "fixtures/v052_lane_sequence_helpers.sql" + )))) + .await?; + migrate_as_owner + .execute(sqlx::raw_sql(sql(r#" + SELECT awa.v052_refresh_lane_sequence_helpers('awa'); + DROP FUNCTION awa.v052_refresh_lane_sequence_helpers(text); + DROP FUNCTION awa.provision_lane(text, smallint, smallint); + DROP FUNCTION awa.provision_lane_sequences(text, smallint, smallint); + DROP FUNCTION awa.repair_lane_sequence(text, text); + DROP FUNCTION awa.lane_sequence_problems(text, text); + "#))) + .await?; + } + migrate_as_owner.close().await; + if runtime_create { + owner + .execute(sqlx::raw_sql(sql(format!( + "GRANT CREATE ON SCHEMA awa TO {runtime_name}" + )))) + .await?; + } + + let app_a = pool(&app_names[0]).await?; + let app_b = pool(&app_names[1]).await?; + let app_c = pool(&app_names[2]).await?; + Ok(Self { + admin, + database, + roles, + owner_name, + migrator_name, + runtime_name, + connect_options: options, + owner, + migrator, + app_a, + app_b, + app_c, + }) + } + + async fn teardown(self) -> TestResult { + for pool in [ + &self.owner, + &self.migrator, + &self.app_a, + &self.app_b, + &self.app_c, + ] { + pool.close().await; + } + self.admin + .execute(sqlx::raw_sql(sql(format!( + "DROP DATABASE {} WITH (FORCE)", + self.database + )))) + .await?; + for role in self.roles.iter().rev() { + self.admin + .execute(sqlx::raw_sql(sql(format!("DROP ROLE {role}")))) + .await?; + } + self.admin.close().await; + Ok(()) + } + + /// The migration role acting as the schema owner, as `awa migrate` runs + /// under the database-roles guide. + async fn migrator_as_owner(&self) -> Result { + PgPoolOptions::new() + .max_connections(1) + .connect_with( + self.connect_options + .clone() + .username(&self.migrator_name) + .options([("role", self.owner_name.as_str())]), + ) + .await + } + + /// The `queue_enqueue_seq_*` sequence of `queue`'s first lane. + async fn enqueue_sequence(&self, queue: &str) -> Result { + Ok(self + .lane_sequence_owners(queue) + .await? + .into_iter() + .map(|(sequence, _)| sequence) + .find(|sequence| sequence.starts_with("queue_enqueue_seq_")) + .expect("lane has an enqueue sequence")) + } + + /// `(sequence, owner)` for every lane sequence of `queue`. + async fn lane_sequence_owners( + &self, + queue: &str, + ) -> Result, sqlx::Error> { + sqlx::query_as( + r#" + SELECT sequence.relname::text, pg_get_userbyid(sequence.relowner)::text + FROM awa.queue_enqueue_heads AS heads + CROSS JOIN LATERAL (VALUES (heads.seq_name), (( + SELECT claims.seq_name FROM awa.queue_claim_heads AS claims + WHERE claims.queue = heads.queue + AND claims.priority = heads.priority + AND claims.enqueue_shard = heads.enqueue_shard + ))) AS lane(seq_name) + JOIN pg_class AS sequence + ON sequence.relname = lane.seq_name + AND sequence.relnamespace = 'awa'::regnamespace + WHERE heads.queue = $1 + ORDER BY 1 + "#, + ) + .bind(queue) + .fetch_all(&self.owner) + .await + } + + async fn assert_lanes_owned_by_owner(&self, queue: &str, lanes: usize) -> TestResult { + let owners = self.lane_sequence_owners(queue).await?; + assert_eq!( + owners.len(), + lanes * 2, + "lane sequences of {queue}: {owners:?}" + ); + for (sequence, owner) in owners { + assert_eq!(owner, self.owner_name, "{sequence} of {queue}"); + } + Ok(()) + } + + async fn set_enqueue_shards(&self, queue: &str, shards: i16) -> Result<(), sqlx::Error> { + sqlx::query( + "INSERT INTO awa.queue_meta (queue, enqueue_shards) VALUES ($1, $2) \ + ON CONFLICT (queue) DO UPDATE SET enqueue_shards = EXCLUDED.enqueue_shards", + ) + .bind(queue) + .bind(shards) + .execute(&self.owner) + .await?; + Ok(()) + } +} + +fn store() -> QueueStorage { + QueueStorage::new(QueueStorageConfig::default()).expect("default queue storage config") +} + +/// Enqueue `count` jobs through a fresh store, as a producer process would. +async fn enqueue(pool: &PgPool, queue: &str, priority: i16, count: i64) -> Result { + store().enqueue_batch(pool, queue, priority, count).await +} + +/// Claim and complete everything on `queue`, as a worker does; returns how +/// many jobs completed. +async fn drain(pool: &PgPool, queue: &str) -> Result { + let worker = store(); + let mut completed = 0; + loop { + let claimed = worker + .claim_runtime_batch(pool, queue, 100, Duration::from_secs(60)) + .await?; + if claimed.is_empty() { + return Ok(completed); + } + completed += worker.complete_runtime_batch(pool, &claimed).await?.len(); + } +} + +/// Schedule jobs that become due shortly, one per priority. +async fn schedule(pool: &PgPool, queue: &str, priorities: &[i16]) -> Result<(), AwaError> { + let run_at = chrono::Utc::now() + chrono::Duration::milliseconds(200); + let jobs: Vec = priorities + .iter() + .map(|&priority| InsertParams { + kind: "lane_ownership_job".into(), + args: serde_json::json!({}), + opts: InsertOpts { + queue: queue.into(), + priority, + run_at: Some(run_at), + ..Default::default() + }, + }) + .collect(); + store().enqueue_params_batch(pool, &jobs).await?; + tokio::time::sleep(Duration::from_millis(300)).await; + Ok(()) +} + +async fn first_touch_by_one_app_role(runtime_create: bool) -> TestResult { + let cluster = Cluster::new(Schema::Current, runtime_create).await?; + + // Role A creates the lane; role B enqueues, claims and completes on it. + assert_eq!(enqueue(&cluster.app_a, "fresh_queue", 3, 2).await?, 2); + assert_eq!(enqueue(&cluster.app_b, "fresh_queue", 3, 2).await?, 2); + assert_eq!(drain(&cluster.app_b, "fresh_queue").await?, 4); + + // The maintenance leader, running as role C, promotes scheduled jobs into + // lanes nobody has created yet: two priorities over two shards. + cluster.set_enqueue_shards("fresh_queue", 2).await?; + schedule(&cluster.app_a, "fresh_queue", &[1, 4]).await?; + assert_eq!( + store() + .promote_due(&cluster.app_c, JobState::Scheduled, 10) + .await?, + 2 + ); + let promoted_shards: Vec = sqlx::query_scalar( + "SELECT enqueue_shard FROM awa.queue_enqueue_heads \ + WHERE queue = 'fresh_queue' AND priority IN (1, 4) ORDER BY 1", + ) + .fetch_all(&cluster.owner) + .await?; + assert_eq!(promoted_shards, vec![0, 1]); + + // A and B both use the lanes the leader created; each producer's rotor + // visits both shards. + let (producer_a, producer_b) = (store(), store()); + for _ in 0..2 { + producer_a + .enqueue_batch(&cluster.app_a, "fresh_queue", 4, 1) + .await?; + producer_b + .enqueue_batch(&cluster.app_b, "fresh_queue", 1, 1) + .await?; + } + assert_eq!(drain(&cluster.app_a, "fresh_queue").await?, 6); + + cluster + .assert_lanes_owned_by_owner("fresh_queue", 5) + .await?; + assert!(storage::lane_sequence_problems(&cluster.owner, "awa") + .await? + .is_empty()); + cluster.teardown().await +} + +#[tokio::test] +async fn lane_first_touched_by_an_app_role_with_schema_create_serves_every_runtime_role( +) -> TestResult { + first_touch_by_one_app_role(true).await +} + +#[tokio::test] +async fn lane_first_touched_by_an_app_role_without_schema_create_serves_every_runtime_role( +) -> TestResult { + first_touch_by_one_app_role(false).await +} + +#[tokio::test] +async fn raising_enqueue_shards_live_lets_the_leader_create_a_shard_first() -> TestResult { + let cluster = Cluster::new(Schema::Current, true).await?; + enqueue(&cluster.app_a, "sharded_queue", 2, 3).await?; + + cluster.set_enqueue_shards("sharded_queue", 2).await?; + // The leader's store has already routed a promotion to shard 0, as a + // long-running leader has, so its next pick is the new shard. + let leader = store(); + schedule(&cluster.app_a, "sharded_queue", &[2]).await?; + assert_eq!( + leader + .promote_due(&cluster.app_c, JobState::Scheduled, 10) + .await?, + 1 + ); + schedule(&cluster.app_b, "sharded_queue", &[2]).await?; + assert_eq!( + leader + .promote_due(&cluster.app_c, JobState::Scheduled, 10) + .await?, + 1 + ); + let shards: Vec = sqlx::query_scalar( + "SELECT enqueue_shard FROM awa.queue_enqueue_heads WHERE queue = 'sharded_queue' ORDER BY 1", + ) + .fetch_all(&cluster.owner) + .await?; + assert_eq!(shards, vec![0, 1]); + + // Producers on both roles spread over both shards; B drains them. + let producer_a = store(); + let producer_b = store(); + for _ in 0..4 { + producer_a + .enqueue_batch(&cluster.app_a, "sharded_queue", 2, 1) + .await?; + producer_b + .enqueue_batch(&cluster.app_b, "sharded_queue", 2, 1) + .await?; + } + assert_eq!(drain(&cluster.app_b, "sharded_queue").await?, 13); + cluster + .assert_lanes_owned_by_owner("sharded_queue", 2) + .await?; + cluster.teardown().await +} + +#[tokio::test] +async fn prepare_queue_as_an_app_role_provisions_lanes_every_runtime_role_can_use() -> TestResult { + let cluster = Cluster::new(Schema::Current, true).await?; + store() + .prepare_queue(&cluster.app_a, "prepared_queue", 2) + .await?; + cluster + .assert_lanes_owned_by_owner("prepared_queue", 8) + .await?; + + cluster.set_enqueue_shards("prepared_queue", 2).await?; + let producer = store(); + for priority in 1..=4 { + producer + .enqueue_batch(&cluster.app_b, "prepared_queue", priority, 1) + .await?; + producer + .enqueue_batch(&cluster.app_c, "prepared_queue", priority, 1) + .await?; + } + assert_eq!(drain(&cluster.app_b, "prepared_queue").await?, 8); + cluster + .assert_lanes_owned_by_owner("prepared_queue", 8) + .await?; + cluster.teardown().await +} + +#[tokio::test] +async fn mixed_fleet_shares_lanes_whichever_path_creates_them() -> TestResult { + let cluster = Cluster::new(Schema::Current, true).await?; + // Role A cannot call provision_lane, so its stores take the path a binary + // that predates the function takes: they write the lane rows themselves. + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "REVOKE EXECUTE ON FUNCTION awa.provision_lane(text, smallint, smallint) FROM PUBLIC, {}", + cluster.runtime_name + )))) + .await?; + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "GRANT EXECUTE ON FUNCTION awa.provision_lane(text, smallint, smallint) TO {}, {}", + cluster.roles[4], cluster.roles[5] + )))) + .await?; + let can_provision: bool = sqlx::query_scalar( + "SELECT has_function_privilege('awa.provision_lane(text,smallint,smallint)', 'EXECUTE')", + ) + .fetch_one(&cluster.app_a) + .await?; + assert!(!can_provision); + + enqueue(&cluster.app_a, "fleet_queue", 1, 2).await?; + enqueue(&cluster.app_b, "fleet_queue", 2, 2).await?; + enqueue(&cluster.app_b, "fleet_queue", 1, 2).await?; + enqueue(&cluster.app_a, "fleet_queue", 2, 2).await?; + assert_eq!(drain(&cluster.app_a, "fleet_queue").await?, 8); + cluster + .assert_lanes_owned_by_owner("fleet_queue", 2) + .await?; + + // A provisioner that disappears after a store has used it: the next lane + // falls back inside the same transaction instead of failing the enqueue. + let producer = store(); + producer + .enqueue_batch(&cluster.app_b, "fleet_queue", 3, 1) + .await?; + cluster + .owner + .execute("DROP FUNCTION awa.provision_lane(text, smallint, smallint)") + .await?; + producer + .enqueue_batch(&cluster.app_b, "fleet_queue", 4, 1) + .await?; + enqueue(&cluster.app_c, "fleet_queue", 4, 1).await?; + assert_eq!(drain(&cluster.app_c, "fleet_queue").await?, 3); + cluster + .assert_lanes_owned_by_owner("fleet_queue", 4) + .await?; + cluster.teardown().await +} + +#[tokio::test] +async fn binary_ahead_of_the_schema_creates_lanes_as_before() -> TestResult { + let cluster = Cluster::new(Schema::BeforeProvisioning, false).await?; + let has_provisioner: bool = sqlx::query_scalar( + "SELECT to_regprocedure('awa.provision_lane(text,smallint,smallint)') IS NOT NULL", + ) + .fetch_one(&cluster.owner) + .await?; + assert!(!has_provisioner); + + // The owner creates lanes lazily, as on v052. + enqueue(&cluster.owner, "ahead_queue", 2, 2).await?; + enqueue(&cluster.app_b, "ahead_queue", 2, 1).await?; + assert_eq!(drain(&cluster.app_b, "ahead_queue").await?, 3); + // A runtime without schema CREATE still gets v052's provisioning hint. + let error = enqueue(&cluster.app_a, "ahead_queue", 3, 1) + .await + .unwrap_err(); + let AwaError::Database(sqlx::Error::Database(db_error)) = &error else { + panic!("expected a database error, got {error:?}"); + }; + assert_eq!(db_error.code().as_deref(), Some("42501")); + assert!(db_error.message().contains("has not been provisioned")); + cluster.teardown().await +} + +#[tokio::test] +async fn concurrent_first_touch_of_one_lane_by_two_roles() -> TestResult { + let cluster = Cluster::new(Schema::Current, true).await?; + + // A holds a brand-new lane's creation open; B's first touch waits for it + // and then uses A's sequences. + let mut holding = cluster.app_a.begin().await?; + sqlx::query("SELECT awa.provision_lane('raced_queue', 2::smallint, 0::smallint)") + .execute(holding.as_mut()) + .await?; + let app_b = cluster.app_b.clone(); + let waiting = tokio::spawn(async move { enqueue(&app_b, "raced_queue", 2, 1).await }); + let mut blocked = false; + for _ in 0..100 { + blocked = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_locks WHERE locktype = 'advisory' AND NOT granted \ + AND database = (SELECT oid FROM pg_database WHERE datname = current_database()))", + ) + .fetch_one(&cluster.owner) + .await?; + if blocked { + break; + } + tokio::time::sleep(Duration::from_millis(20)).await; + } + assert!(blocked, "B's first touch should wait on A's lane lock"); + holding.commit().await?; + assert_eq!(waiting.await??, 1); + + // Unordered races on many fresh lanes. + for round in 0..20 { + let queue = format!("raced_queue_{round}"); + let (from_a, from_b) = tokio::join!( + enqueue(&cluster.app_a, &queue, 1, 1), + enqueue(&cluster.app_b, &queue, 1, 1) + ); + assert_eq!(from_a? + from_b?, 2); + cluster.assert_lanes_owned_by_owner(&queue, 1).await?; + assert_eq!(drain(&cluster.app_c, &queue).await?, 2); + } + assert_eq!(drain(&cluster.app_c, "raced_queue").await?, 1); + cluster + .assert_lanes_owned_by_owner("raced_queue", 1) + .await?; + cluster.teardown().await +} + +#[tokio::test] +async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { + let mut cluster = Cluster::new(Schema::BeforeProvisioning, true).await?; + // Lanes as a 0.6 fleet left them: one each created by an app role, by the + // migration role acting as itself, and by the owner without the runtime's + // grant. + enqueue(&cluster.app_a, "app_lane", 1, 1).await?; + enqueue(&cluster.migrator, "migrator_lane", 1, 1).await?; + enqueue(&cluster.owner, "owner_lane", 1, 1).await?; + for (sequence, _) in cluster.lane_sequence_owners("owner_lane").await? { + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "REVOKE ALL ON SEQUENCE awa.{sequence} FROM {}", + cluster.runtime_name + )))) + .await?; + } + for lane in ["owner_lane", "app_lane"] { + let failed = drain(&cluster.app_b, lane).await.unwrap_err(); + assert!( + failed + .to_string() + .contains("permission denied for sequence"), + "{failed}" + ); + } + // 0.7 refuses to migrate while app roles hold CREATE on the schema. + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "REVOKE CREATE ON SCHEMA awa FROM {}", + cluster.runtime_name + )))) + .await?; + + // `awa migrate` connects as the migration role and acts as the owner. As + // the owner it adds the grants the owner's sequences lack; the migration + // role's own sequences need the migration role's privileges, so they are + // reassigned as the login role. Neither can act as the app role. + let migrate_pool = cluster.migrator_as_owner().await?; + tokio::time::timeout(Duration::from_secs(120), migrations::run(&migrate_pool)).await??; + assert_eq!( + migrations::current_version(&cluster.owner).await?, + migrations::CURRENT_VERSION + ); + migrate_pool.close().await; + cluster + .assert_lanes_owned_by_owner("migrator_lane", 1) + .await?; + cluster.assert_lanes_owned_by_owner("owner_lane", 1).await?; + assert_eq!(drain(&cluster.app_b, "migrator_lane").await?, 1); + assert_eq!(drain(&cluster.app_b, "owner_lane").await?, 1); + + let app_a = cluster.roles[3].clone(); + let app_a = &app_a; + let problems = storage::lane_sequence_problems(&cluster.app_c, "awa").await?; + assert_eq!(problems.len(), 2, "{problems:?}"); + for problem in &problems { + assert_eq!(&problem.owner, app_a); + assert_eq!(problem.expected_owner, cluster.owner_name); + assert_eq!( + problem.missing_grants, + vec![format!( + "GRANT USAGE, UPDATE, SELECT ON SEQUENCE awa.{} TO {}", + problem.sequence_name, cluster.runtime_name + )] + ); + assert_eq!( + problem.reassign_sql.as_deref(), + Some( + format!( + "ALTER SEQUENCE awa.{} OWNER TO {}", + problem.sequence_name, cluster.owner_name + ) + .as_str() + ) + ); + assert!(problem.repair_roles.is_empty(), "{problem:?}"); + } + let next_steps = + storage::lane_sequence_next_steps(problems.iter().map(|problem| (problem, false))); + assert_eq!(next_steps.len(), 1, "{next_steps:?}"); + assert!( + next_steps[0].contains(&format!("no login role can both act as {app_a}")), + "{next_steps:?}" + ); + + // Re-running the repair as the migration role changes nothing and fails + // nothing. + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "GRANT EXECUTE ON FUNCTION awa.repair_lane_sequence(text, text) TO {}", + cluster.runtime_name + )))) + .await?; + let rerun = + storage::repair_lane_sequences(&mut *cluster.migrator.acquire().await?, "awa").await?; + assert_eq!(rerun.len(), 2); + assert!(rerun + .iter() + .all(|result| result.outcome == LaneSequenceOutcome::Unrepaired)); + + // A role that can act as the app role but cannot become the owner adds + // the grants and leaves the owner for someone who can. + let delegate = format!("{}_delegate", cluster.runtime_name); + cluster + .admin + .execute(sqlx::raw_sql(sql(format!( + "CREATE ROLE {delegate} LOGIN PASSWORD 'lane_test' IN ROLE {app_a}; \ + GRANT CONNECT ON DATABASE {} TO {delegate}", + cluster.database + )))) + .await?; + let delegate_pool = PgPoolOptions::new() + .max_connections(1) + .connect_with(cluster.connect_options.clone().username(&delegate)) + .await?; + let delegated = + storage::repair_lane_sequences(&mut *delegate_pool.acquire().await?, "awa").await?; + delegate_pool.close().await; + assert_eq!(delegated.len(), 2); + for result in &delegated { + assert_eq!(result.outcome, LaneSequenceOutcome::OwnerUnchanged); + assert_eq!(&result.owner, app_a); + assert!(result.missing_grants.is_empty()); + assert!(result.reassign_sql.is_some()); + } + assert_eq!(drain(&cluster.app_b, "app_lane").await?, 1); + enqueue(&cluster.app_c, "app_lane", 1, 1).await?; + + // Once the migration role is a member of the app role too, the check names + // it, and a repair connected as it reassigns. + cluster + .admin + .execute(sqlx::raw_sql(sql(format!( + "GRANT {app_a} TO {migrator}", + migrator = cluster.migrator_name + )))) + .await?; + let problems = storage::lane_sequence_problems(&cluster.app_c, "awa").await?; + assert!(problems + .iter() + .all(|problem| problem.repair_roles == vec![cluster.migrator_name.clone()])); + let repaired = + storage::repair_lane_sequences(&mut *cluster.migrator.acquire().await?, "awa").await?; + assert_eq!(repaired.len(), 2); + assert!( + repaired + .iter() + .all(|result| result.outcome == LaneSequenceOutcome::Repaired), + "{repaired:?}" + ); + assert!(storage::lane_sequence_problems(&cluster.app_c, "awa") + .await? + .is_empty()); + cluster.assert_lanes_owned_by_owner("app_lane", 1).await?; + assert_eq!(drain(&cluster.app_b, "app_lane").await?, 1); + + cluster.roles.push(delegate); + cluster.teardown().await +} + +/// Reassigning a lane sequence waits on any open enqueue or claim on the lane, +/// so the repair bounds that wait and holds each sequence's lock only for its +/// own short transaction. A held sequence is skipped and reported; the rest +/// are repaired, and a later repair finishes it. +#[tokio::test] +async fn migration_skips_a_lane_sequence_an_open_transaction_holds() -> TestResult { + let cluster = Cluster::new(Schema::BeforeProvisioning, true).await?; + enqueue(&cluster.migrator, "held_lane", 1, 1).await?; + enqueue(&cluster.migrator, "free_lane", 1, 1).await?; + let held_sequence = cluster.enqueue_sequence("held_lane").await?; + // 0.7 refuses to migrate while app roles hold CREATE on the schema. + cluster + .owner + .execute(sqlx::raw_sql(sql(format!( + "REVOKE CREATE ON SCHEMA awa FROM {}", + cluster.runtime_name + )))) + .await?; + + // An enqueue in flight on held_lane holds its sequence's lock. + let mut holder = cluster.app_b.begin().await?; + sqlx::query("SELECT nextval($1::regclass)") + .bind(format!("awa.{held_sequence}")) + .execute(&mut *holder) + .await?; + + let migrate_pool = cluster.migrator_as_owner().await?; + // The migration completes while the lock is held. + tokio::time::timeout(Duration::from_secs(120), migrations::run(&migrate_pool)).await??; + for (sequence, owner) in cluster.lane_sequence_owners("held_lane").await? { + let expected = if sequence == held_sequence { + &cluster.migrator_name + } else { + &cluster.owner_name + }; + assert_eq!(&owner, expected, "{sequence}"); + } + cluster.assert_lanes_owned_by_owner("free_lane", 1).await?; + + // A repair while the transaction is still open reports the sequence busy. + let mut conn = migrate_pool.acquire().await?; + let busy = storage::repair_lane_sequences(&mut conn, "awa").await?; + assert_eq!(busy.len(), 1, "{busy:?}"); + assert_eq!(busy[0].sequence_name, held_sequence); + assert_eq!(busy[0].outcome, LaneSequenceOutcome::Busy); + assert_eq!(busy[0].owner, cluster.migrator_name); + let next_steps = storage::lane_sequence_repair_next_steps(&busy); + assert_eq!(next_steps.len(), 1); + assert!(next_steps[0].contains("busy"), "{next_steps:?}"); + assert!( + next_steps[0].contains("awa storage lane-sequences --schema awa --repair"), + "{next_steps:?}" + ); + + holder.rollback().await?; + let repaired = storage::repair_lane_sequences(&mut conn, "awa").await?; + assert_eq!(repaired.len(), 1, "{repaired:?}"); + assert_eq!(repaired[0].outcome, LaneSequenceOutcome::Repaired); + drop(conn); + migrate_pool.close().await; + cluster.assert_lanes_owned_by_owner("held_lane", 1).await?; + assert!(storage::lane_sequence_problems(&cluster.app_c, "awa") + .await? + .is_empty()); + assert_eq!(drain(&cluster.app_b, "held_lane").await?, 1); + cluster.teardown().await +} diff --git a/awa/tests/statistics_refresh_test.rs b/awa/tests/statistics_refresh_test.rs new file mode 100644 index 000000000..ca87ed498 --- /dev/null +++ b/awa/tests/statistics_refresh_test.rs @@ -0,0 +1,834 @@ +//! The maintenance leader refreshes planner statistics for the queue-storage +//! tables autovacuum never analyses (the head tables and the partitioned ring +//! parents) on election and every interval, without delaying ring prune. + +use awa::model::{migrations, PruneOutcome, QueueStorage, QueueStorageConfig, RotateOutcome}; +use awa::{audited_sql, Client, JobContext, JobError, JobResult, QueueConfig, Worker}; +use opentelemetry::KeyValue; +use opentelemetry_sdk::metrics::data::{AggregatedMetrics, MetricData}; +use opentelemetry_sdk::metrics::{InMemoryMetricExporter, SdkMeterProvider}; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; +use sqlx::PgPool; +use std::str::FromStr; +use std::sync::atomic::{AtomicU32, Ordering}; +use std::time::{Duration, Instant}; + +const ROLE_PASSWORD: &str = "statistics_refresh_test"; +const RING_PARENTS: [&str; 13] = [ + "done_entries", + "lease_claim_batches", + "lease_claim_closure_batches", + "lease_claim_closures", + "lease_claims", + "leases", + "queue_terminal_count_deltas", + "ready_claim_attempt_batches", + "ready_entries", + "ready_segments", + "ready_tombstones", + "receipt_completion_batches", + "receipt_completion_tombstones", +]; + +static DATABASE_COUNTER: AtomicU32 = AtomicU32::new(0); + +async fn run_sql(pool: &PgPool, sql: String) -> Result<(), sqlx::Error> { + sqlx::raw_sql(audited_sql(sql)) + .execute(pool) + .await + .map(|_| ()) +} + +fn base_database_url() -> String { + std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()) +} + +struct StatisticsProbeWorker; + +#[async_trait::async_trait] +impl Worker for StatisticsProbeWorker { + fn kind(&self) -> &'static str { + "statistics_probe_job" + } + + async fn perform(&self, _ctx: &JobContext) -> Result { + Ok(JobResult::Completed) + } +} + +/// A throwaway database with the awa schema migrated, optionally owned by a +/// non-superuser role the way a production schema owner is. +struct TestDatabase { + admin: PgPool, + name: String, + roles: Vec, + options: PgConnectOptions, + /// Connects as the schema owner (or the superuser when no owner role). + owner_pool: PgPool, +} + +impl TestDatabase { + async fn create(owner_role: Option<&str>, extra_roles: &[&str]) -> Self { + let admin = PgPoolOptions::new() + .max_connections(2) + .connect(&base_database_url()) + .await + .expect("connect admin pool"); + let name = format!( + "awa_stats_{}_{}", + std::process::id(), + DATABASE_COUNTER.fetch_add(1, Ordering::SeqCst) + ); + let mut roles = Vec::new(); + for role in owner_role.iter().chain(extra_roles.iter()) { + sqlx::raw_sql(audited_sql(format!( + "DROP ROLE IF EXISTS {role}; CREATE ROLE {role} LOGIN PASSWORD '{ROLE_PASSWORD}'" + ))) + .execute(&admin) + .await + .expect("create role"); + roles.push((*role).to_string()); + } + let owner_clause = owner_role + .map(|role| format!(" OWNER {role}")) + .unwrap_or_default(); + sqlx::raw_sql(audited_sql(format!("CREATE DATABASE {name}{owner_clause}"))) + .execute(&admin) + .await + .expect("create database"); + let options = PgConnectOptions::from_str(&base_database_url()) + .expect("parse DATABASE_URL") + .database(&name); + let owner_options = match owner_role { + Some(role) => options.clone().username(role).password(ROLE_PASSWORD), + None => options.clone(), + }; + let owner_pool = PgPoolOptions::new() + .max_connections(8) + .connect_with(owner_options) + .await + .expect("connect owner pool"); + migrations::run(&owner_pool).await.expect("run migrations"); + Self { + admin, + name, + roles, + options, + owner_pool, + } + } + + /// A pool for a role other than the schema owner, with the runtime grants + /// from the database-roles guide. + async fn runtime_pool(&self, role: &str) -> PgPool { + sqlx::raw_sql(audited_sql(format!( + r#" + GRANT USAGE ON SCHEMA awa TO {role}; + GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA awa TO {role}; + GRANT USAGE, SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA awa TO {role}; + GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa TO {role}; + REVOKE EXECUTE ON FUNCTION + awa.install_queue_storage_substrate(TEXT, INT, INT, INT, BOOLEAN) FROM {role}; + "# + ))) + .execute(&self.owner_pool) + .await + .expect("grant runtime privileges"); + PgPoolOptions::new() + .max_connections(8) + .connect_with(self.options.clone().username(role).password(ROLE_PASSWORD)) + .await + .expect("connect runtime pool") + } + + async fn drop(self) { + self.owner_pool.close().await; + let _ = sqlx::raw_sql(audited_sql(format!( + "DROP DATABASE IF EXISTS {} WITH (FORCE)", + self.name + ))) + .execute(&self.admin) + .await; + for role in &self.roles { + let _ = sqlx::raw_sql(audited_sql(format!("DROP ROLE IF EXISTS {role}"))) + .execute(&self.admin) + .await; + } + self.admin.close().await; + } +} + +async fn wait_for(what: &str, timeout: Duration, mut check: F) +where + F: FnMut() -> Fut, + Fut: std::future::Future, +{ + let deadline = Instant::now() + timeout; + while Instant::now() < deadline { + if check().await { + return; + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + panic!("timed out waiting for {what}"); +} + +async fn stats_rows(pool: &PgPool, schema: &str, table: &str, inherited_only: bool) -> i64 { + sqlx::query_scalar( + "SELECT count(*) FROM pg_stats + WHERE schemaname = $1 AND tablename = $2 AND (inherited OR NOT $3)", + ) + .bind(schema) + .bind(table) + .bind(inherited_only) + .fetch_one(pool) + .await + .expect("read pg_stats") +} + +async fn parent_reltuples(pool: &PgPool, schema: &str, table: &str) -> f32 { + sqlx::query_scalar("SELECT reltuples FROM pg_class WHERE oid = to_regclass($1)") + .bind(format!("{schema}.{table}")) + .fetch_one(pool) + .await + .expect("read reltuples") +} + +async fn analyze_count(pool: &PgPool, schema: &str, table: &str) -> i64 { + sqlx::query_scalar( + "SELECT COALESCE(analyze_count, 0)::bigint FROM pg_stat_user_tables + WHERE schemaname = $1 AND relname = $2", + ) + .bind(schema) + .bind(table) + .fetch_one(pool) + .await + .expect("read analyze_count") +} + +fn probe_client( + pool: PgPool, + statistics_interval: Duration, + storage: Option, +) -> Client { + let mut builder = Client::builder(pool) + .queue( + "statistics_probe", + QueueConfig { + max_workers: 1, + poll_interval: Duration::from_millis(100), + ..QueueConfig::default() + }, + ) + .register_worker(StatisticsProbeWorker) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(100)) + .statistics_refresh_interval(statistics_interval); + if let Some(config) = storage { + builder = builder + .queue_storage(config, Duration::from_secs(600), Duration::from_secs(600)) + .claim_rotate_interval(Duration::from_secs(600)); + } + builder.build().expect("build client") +} + +async fn seed_ready_rows(store: &QueueStorage, pool: &PgPool, rows_per_queue: i64) { + for queue in ["stats_a", "stats_b", "stats_c"] { + store + .enqueue_batch(pool, queue, 2, rows_per_queue) + .await + .expect("enqueue rows"); + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn leader_refreshes_statistics_on_election_and_every_interval() { + let database = TestDatabase::create(None, &[]).await; + let pool = &database.owner_pool; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage config"); + seed_ready_rows(&store, pool, 200).await; + + assert_eq!(stats_rows(pool, "awa", "queue_claim_heads", false).await, 0); + assert_eq!( + stats_rows(pool, "awa", "queue_enqueue_heads", false).await, + 0 + ); + assert_eq!(stats_rows(pool, "awa", "ready_entries", true).await, 0); + assert!(parent_reltuples(pool, "awa", "ready_entries").await < 0.0); + + let client = probe_client(pool.clone(), Duration::from_secs(1), None); + client.start().await.expect("start client"); + + wait_for( + "first statistics pass on election", + Duration::from_secs(30), + || async { + stats_rows(pool, "awa", "queue_claim_heads", false).await > 0 + && stats_rows(pool, "awa", "queue_enqueue_heads", false).await > 0 + && stats_rows(pool, "awa", "ready_entries", true).await > 0 + }, + ) + .await; + let partitioned: Vec = sqlx::query_scalar( + "SELECT relname::text FROM pg_class + WHERE relnamespace = 'awa'::regnamespace AND relkind = 'p' ORDER BY relname", + ) + .fetch_all(pool) + .await + .expect("list partitioned tables"); + assert_eq!( + partitioned, RING_PARENTS, + "the refresh covers every ring parent" + ); + wait_for( + "every ring parent analysed", + Duration::from_secs(30), + || async { + for parent in RING_PARENTS { + if analyze_count(pool, "awa", parent).await < 1 { + return false; + } + } + true + }, + ) + .await; + let first_tuples = parent_reltuples(pool, "awa", "ready_entries").await; + assert!( + (550.0..=650.0).contains(&first_tuples), + "parent reltuples after the first pass: {first_tuples}" + ); + + seed_ready_rows(&store, pool, 2_000).await; + wait_for( + "a later pass to see the new rows", + Duration::from_secs(30), + || async { parent_reltuples(pool, "awa", "ready_entries").await > 5_000.0 }, + ) + .await; + assert!(analyze_count(pool, "awa", "ready_entries").await >= 2); + assert!(analyze_count(pool, "awa", "queue_claim_heads").await >= 2); + + client.shutdown(Duration::from_secs(5)).await; + database.drop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn leader_refreshes_statistics_on_election_before_the_first_interval() { + let database = TestDatabase::create(None, &[]).await; + let pool = &database.owner_pool; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage config"); + seed_ready_rows(&store, pool, 200).await; + + let client = probe_client(pool.clone(), Duration::from_secs(300), None); + client.start().await.expect("start client"); + wait_for( + "a statistics pass on election", + Duration::from_secs(30), + || async { stats_rows(pool, "awa", "ready_entries", true).await > 0 }, + ) + .await; + + client.shutdown(Duration::from_secs(5)).await; + database.drop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn zero_interval_disables_the_refresh() { + let database = TestDatabase::create(None, &[]).await; + let pool = &database.owner_pool; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage config"); + seed_ready_rows(&store, pool, 100).await; + + let client = probe_client(pool.clone(), Duration::ZERO, None); + client.start().await.expect("start client"); + wait_for("leadership", Duration::from_secs(10), || async { + client.health_check().await.maintenance_alive + }) + .await; + tokio::time::sleep(Duration::from_secs(3)).await; + + assert_eq!(stats_rows(pool, "awa", "queue_claim_heads", false).await, 0); + assert_eq!(stats_rows(pool, "awa", "ready_entries", true).await, 0); + + client.shutdown(Duration::from_secs(5)).await; + database.drop().await; +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn non_owner_runtime_role_refreshes_statistics() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let owner = format!("stats_owner_{suffix}"); + let runtime = format!("stats_runtime_{suffix}"); + let database = TestDatabase::create(Some(&owner), &[&runtime]).await; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage config"); + seed_ready_rows(&store, &database.owner_pool, 200).await; + let runtime_pool = database.runtime_pool(&runtime).await; + + let server_version_num: i32 = + sqlx::query_scalar("SELECT current_setting('server_version_num')::int") + .fetch_one(&runtime_pool) + .await + .expect("read PostgreSQL version"); + // pg_maintain and the MAINTAIN privilege were introduced in PostgreSQL 17. + // Older supported servers must never evaluate either privilege check. + let privilege_query = if server_version_num >= 170000 { + "SELECT rolsuper OR pg_has_role(current_user, 'pg_maintain', 'USAGE') + OR has_table_privilege('awa.queue_claim_heads', 'MAINTAIN') + OR (SELECT relowner = (SELECT oid FROM pg_roles WHERE rolname = current_user) + FROM pg_class WHERE oid = 'awa.queue_claim_heads'::regclass) + FROM pg_roles WHERE rolname = current_user" + } else { + "SELECT rolsuper + OR (SELECT relowner = (SELECT oid FROM pg_roles WHERE rolname = current_user) + FROM pg_class WHERE oid = 'awa.queue_claim_heads'::regclass) + FROM pg_roles WHERE rolname = current_user" + }; + let is_superuser_or_owner: bool = sqlx::query_scalar(privilege_query) + .fetch_one(&runtime_pool) + .await + .expect("check runtime role privileges"); + assert!( + !is_superuser_or_owner, + "the runtime role must not be able to analyse the tables itself" + ); + + for (schema, relation, message) in [ + ( + "public", + "queue_claim_heads", + "is not a queue-storage schema", + ), + ("awa", "jobs", "is not a ring parent or head table"), + ] { + let error = + sqlx::query_scalar::<_, String>("SELECT awa.analyze_queue_storage_relation($1, $2)") + .bind(schema) + .bind(relation) + .fetch_one(&runtime_pool) + .await + .expect_err("the helper rejects what it may not analyse"); + let database_error = error.as_database_error().expect("database error"); + assert_eq!(database_error.code().as_deref(), Some("22023")); + assert!( + database_error.message().contains(message), + "{schema}.{relation}: {}", + database_error.message() + ); + } + + let client = probe_client(runtime_pool.clone(), Duration::from_secs(1), None); + client + .start() + .await + .expect("start client as the runtime role"); + wait_for( + "statistics written through a non-owner role", + Duration::from_secs(30), + || async { + stats_rows(&database.owner_pool, "awa", "queue_claim_heads", false).await > 0 + && stats_rows(&database.owner_pool, "awa", "ready_entries", true).await > 0 + }, + ) + .await; + + client.shutdown(Duration::from_secs(5)).await; + runtime_pool.close().await; + database.drop().await; +} + +async fn call_helper( + pool: &PgPool, + schema: &str, + relation: &str, +) -> Result { + sqlx::query_scalar("SELECT awa.analyze_queue_storage_relation($1, $2)") + .bind(schema) + .bind(relation) + .fetch_one(pool) + .await +} + +fn denied(error: &sqlx::Error, message: &str) { + let database_error = error.as_database_error().expect("database error"); + assert_eq!(database_error.code().as_deref(), Some("42501")); + assert!( + database_error.message().contains(message), + "unexpected message: {}", + database_error.message() + ); +} + +/// EXECUTE is not public: holding the lock the helper takes until the end of +/// a transaction would block ring prune of that table. A role needs EXECUTE and +/// write privileges on the head tables, and the migration grants EXECUTE to +/// the roles that already have them. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn execute_is_limited_to_roles_that_may_maintain_queue_storage() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let owner = format!("stats_owner_{suffix}"); + let runtime = format!("stats_runtime_{suffix}"); + let bystander = format!("stats_bystander_{suffix}"); + let reader = format!("stats_reader_{suffix}"); + let database = TestDatabase::create(Some(&owner), &[&runtime, &bystander, &reader]).await; + let owner_pool = &database.owner_pool; + run_sql( + owner_pool, + format!( + "GRANT USAGE ON SCHEMA awa TO {bystander}, {reader}, {runtime}; + GRANT SELECT ON ALL TABLES IN SCHEMA awa TO {reader}; + GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) TO {reader}; + GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA awa TO {runtime}" + ), + ) + .await + .expect("grant privileges"); + let connect = |role: String| { + let options = database.options.clone(); + async move { + PgPoolOptions::new() + .max_connections(2) + .connect_with(options.username(&role).password(ROLE_PASSWORD)) + .await + .expect("connect role pool") + } + }; + let bystander_pool = connect(bystander.clone()).await; + let reader_pool = connect(reader.clone()).await; + let runtime_pool = connect(runtime.clone()).await; + + denied( + &call_helper(&bystander_pool, "awa", "queue_claim_heads") + .await + .expect_err("a role without EXECUTE is denied"), + "permission denied for function", + ); + denied( + &call_helper(&reader_pool, "awa", "queue_claim_heads") + .await + .expect_err("EXECUTE without write privileges on the head tables is denied"), + "may not maintain queue storage", + ); + + // The runtime role holds write privileges but was granted no function + // EXECUTE: re-running the migration grants it, as an upgrade does. + denied( + &call_helper(&runtime_pool, "awa", "queue_claim_heads") + .await + .expect_err("write privileges alone do not grant EXECUTE"), + "permission denied for function", + ); + for (_, _, sql) in migrations::migration_sql_range( + migrations::CURRENT_VERSION - 1, + migrations::CURRENT_VERSION, + ) { + run_sql(owner_pool, sql) + .await + .expect("re-run the migration"); + } + assert_eq!( + call_helper(&runtime_pool, "awa", "queue_claim_heads") + .await + .expect("the migration granted EXECUTE to the runtime grantee"), + "analyzed" + ); + + for pool in [bystander_pool, reader_pool, runtime_pool] { + pool.close().await; + } + database.drop().await; +} + +/// The helper analyses the ring parents and head tables only, not any other +/// partitioned table of the schema. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn helper_rejects_a_partitioned_table_that_is_not_a_ring_parent() { + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let owner = format!("stats_owner_{suffix}"); + let runtime = format!("stats_runtime_{suffix}"); + let database = TestDatabase::create(Some(&owner), &[&runtime]).await; + run_sql( + &database.owner_pool, + "CREATE TABLE awa.application_events (id int) PARTITION BY LIST (id)".to_string(), + ) + .await + .expect("create a partitioned application table"); + let runtime_pool = database.runtime_pool(&runtime).await; + + for relation in ["application_events", "ready_entries_0", "jobs_hot"] { + let error = call_helper(&runtime_pool, "awa", relation) + .await + .expect_err("only ring parents and head tables are accepted"); + let database_error = error.as_database_error().expect("database error"); + assert_eq!( + database_error.code().as_deref(), + Some("22023"), + "{relation}" + ); + assert!( + database_error + .message() + .contains("is not a ring parent or head table"), + "{relation}: {}", + database_error.message() + ); + } + assert_eq!( + call_helper(&runtime_pool, "awa", "ready_entries") + .await + .expect("a ring parent is accepted"), + "analyzed" + ); + + runtime_pool.close().await; + database.drop().await; +} + +/// A conflicting lock on a ring parent fails the call with `lock_not_available` +/// rather than letting `ANALYZE (SKIP_LOCKED)` skip every partition and report +/// success. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn locked_ring_parent_fails_the_call_instead_of_reporting_analyzed() { + let database = TestDatabase::create(None, &[]).await; + let pool = &database.owner_pool; + + let mut blocker = pool.begin().await.expect("begin blocker"); + sqlx::query("LOCK TABLE ONLY awa.ready_entries IN SHARE UPDATE EXCLUSIVE MODE") + .execute(blocker.as_mut()) + .await + .expect("lock the parent"); + + let mut call = pool.begin().await.expect("begin call"); + sqlx::query("SELECT set_config('lock_timeout', '200ms', true)") + .execute(call.as_mut()) + .await + .expect("set lock_timeout"); + let error = sqlx::query_scalar::<_, String>( + "SELECT awa.analyze_queue_storage_relation('awa', 'ready_entries')", + ) + .fetch_one(call.as_mut()) + .await + .expect_err("the locked parent must not report analyzed"); + assert_eq!( + error + .as_database_error() + .and_then(|error| error.code()) + .as_deref(), + Some("55P03") + ); + call.rollback().await.expect("rollback call"); + + blocker.rollback().await.expect("release the parent"); + let status: String = + sqlx::query_scalar("SELECT awa.analyze_queue_storage_relation('awa', 'ready_entries')") + .fetch_one(pool) + .await + .expect("analyse once the lock is gone"); + assert_eq!(status, "analyzed"); + database.drop().await; +} + +fn sum_by_outcome( + resource_metrics: &[opentelemetry_sdk::metrics::data::ResourceMetrics], + outcome: &str, +) -> u64 { + let mut total = 0; + for resource in resource_metrics { + for scope in resource.scope_metrics() { + for metric in scope.metrics() { + if metric.name() != "awa.maintenance.analyze.attempts" { + continue; + } + if let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data() { + total += sum + .data_points() + .filter(|point| { + point.attributes().any(|KeyValue { key, value, .. }| { + key.as_str() == "awa.analyze.outcome" && value.as_str() == outcome + }) + }) + .map(|point| point.value()) + .sum::(); + } + } + } + } + total +} + +/// A refresh that cannot run is recorded and leaves the runtime healthy: the +/// helper's owner may not analyse the tables, or the schema predates the +/// helper. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unrunnable_refresh_is_recorded_and_does_not_fail_the_runtime() { + let exporter = InMemoryMetricExporter::default(); + let provider = SdkMeterProvider::builder() + .with_periodic_exporter(exporter.clone()) + .build(); + opentelemetry::global::set_meter_provider(provider.clone()); + + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let owner = format!("stats_owner_{suffix}"); + let helper = format!("stats_helper_{suffix}"); + let database = TestDatabase::create(Some(&owner), &[&helper]).await; + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage config"); + seed_ready_rows(&store, &database.owner_pool, 50).await; + + // The function now runs as a role that owns nothing and holds no MAINTAIN. + sqlx::raw_sql(audited_sql(format!("GRANT {helper} TO {owner}"))) + .execute(&database.admin) + .await + .expect("let the schema owner hand over the function"); + sqlx::raw_sql(audited_sql(format!( + "GRANT USAGE, CREATE ON SCHEMA awa TO {helper}; + ALTER FUNCTION awa.analyze_queue_storage_relation(text, text) OWNER TO {helper}; + GRANT EXECUTE ON FUNCTION awa.analyze_queue_storage_relation(text, text) TO {owner}" + ))) + .execute(&database.owner_pool) + .await + .expect("hand the helper to an unprivileged role"); + + let client = probe_client( + database.owner_pool.clone(), + Duration::from_millis(200), + None, + ); + client.start().await.expect("start client"); + wait_for("a skipped pass", Duration::from_secs(30), || async { + provider.force_flush().expect("flush metrics"); + let metrics = exporter.get_finished_metrics().expect("read metrics"); + sum_by_outcome(&metrics, "skipped_privilege") > 0 + }) + .await; + assert_eq!( + stats_rows(&database.owner_pool, "awa", "queue_claim_heads", false).await, + 0, + "a skipped pass writes no statistics" + ); + assert!(client.health_check().await.healthy); + client.shutdown(Duration::from_secs(5)).await; + + // A schema that predates the helper skips the same way. + sqlx::raw_sql("DROP FUNCTION awa.analyze_queue_storage_relation(text, text)") + .execute(&database.owner_pool) + .await + .expect("drop the helper"); + let client = probe_client( + database.owner_pool.clone(), + Duration::from_millis(200), + None, + ); + client + .start() + .await + .expect("start client without the helper"); + wait_for("an unavailable pass", Duration::from_secs(30), || async { + provider.force_flush().expect("flush metrics"); + let metrics = exporter.get_finished_metrics().expect("read metrics"); + sum_by_outcome(&metrics, "unavailable") > 0 + }) + .await; + assert!(client.health_check().await.healthy); + client.shutdown(Duration::from_secs(5)).await; + + provider.force_flush().expect("flush metrics"); + let metrics = exporter.get_finished_metrics().expect("read metrics"); + assert_eq!(sum_by_outcome(&metrics, "failed"), 0); + database.drop().await; +} + +/// Back-to-back statistics passes (a far higher duty cycle than the default +/// five-minute interval) run while rotation and prune cycle the ring: prune +/// keeps reclaiming slots and only defers an attempt that lands while ANALYZE +/// is sampling the slot's tables. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn statistics_refresh_does_not_stop_prune_reclaiming_slots() { + let database = TestDatabase::create(None, &[]).await; + let pool = &database.owner_pool; + let schema = "awa_stats_ring"; + let config = QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + claim_slot_count: 2, + lease_claim_receipts: false, + ..Default::default() + }; + let store = QueueStorage::new(config.clone()).expect("queue storage config"); + store.prepare_schema(pool).await.expect("prepare schema"); + + let queue = "statistics_ring"; + let analyze_before = analyze_count(pool, schema, "ready_entries").await; + let client = probe_client(pool.clone(), Duration::from_millis(1), Some(config)); + client.start().await.expect("start client"); + + let (mut pruned, mut deferred) = (0_u32, 0_u32); + let started = Instant::now(); + for _ in 0..60 { + store + .enqueue_batch(pool, queue, 1, 600) + .await + .expect("enqueue"); + loop { + let claims = store.claim_batch(pool, queue, 600).await.expect("claim"); + if claims.is_empty() { + break; + } + store.complete_batch(pool, &claims).await.expect("complete"); + } + match store.rotate(pool).await.expect("rotate") { + RotateOutcome::Rotated { .. } | RotateOutcome::SkippedBusy { .. } => {} + other => panic!("unexpected rotate outcome {other:?}"), + } + match store + .prune_oldest(pool, Duration::ZERO) + .await + .expect("prune never errors under ANALYZE") + { + PruneOutcome::Pruned { .. } => pruned += 1, + PruneOutcome::Blocked { .. } => deferred += 1, + PruneOutcome::Noop | PruneOutcome::SkippedActive { .. } => {} + } + if started.elapsed() > Duration::from_secs(60) { + break; + } + } + let passes = analyze_count(pool, schema, "ready_entries").await - analyze_before; + eprintln!("prune under ANALYZE: pruned={pruned} deferred={deferred} analyze_passes={passes}"); + + client.shutdown(Duration::from_secs(5)).await; + assert!( + passes >= 3, + "the refresh ran {passes} times during the loop" + ); + assert!(pruned >= 20, "prune reclaimed {pruned} slots"); + + // Shutdown can leave one ANALYZE statement finishing server-side; once it + // is gone the oldest sealed slot reclaims without being blocked. + wait_for( + "the last ANALYZE to finish", + Duration::from_secs(30), + || async { + let running: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pg_stat_activity + WHERE state = 'active' AND query LIKE '%analyze_queue_storage_relation%' + AND pid <> pg_backend_pid()", + ) + .fetch_one(pool) + .await + .expect("read pg_stat_activity"); + running == 0 + }, + ) + .await; + store.rotate(pool).await.expect("rotate"); + assert!(!matches!( + store + .prune_oldest(pool, Duration::ZERO) + .await + .expect("prune"), + PruneOutcome::Blocked { .. } + )); + database.drop().await; +} diff --git a/docs/architecture.md b/docs/architecture.md index d3c527f61..388a133a5 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -215,6 +215,8 @@ Queue storage has three independent rings, each advanced by the elected maintena The maintenance tick for each ring is deliberately small: attempt one rotate, then attempt one prune. If a partition is busy, blocked by a lock, current, or still live, the tick records a skipped/blocked outcome and tries again on a future interval. +The leader also refreshes planner statistics on the two head tables and the ring parents when it is elected and every five minutes ([Planner statistics](configuration.md#planner-statistics)): autovacuum never analyses partitioned tables or reaches its threshold on the head tables. It runs beside the ring ticks rather than in them, so a slow `ANALYZE` never delays a rotate or prune tick. + The maintenance tick decides whether to seal the current slot. `ClientBuilder::ring_rotation_cadence` configures `RingRotationCadence`: a non-empty slot becomes eligible after `min_slot_age` since the maintenance store first observed rows (default 5 s) once it holds at least `min_slot_rows` (default 256, sampled with a bounded probe), when total physical storage reaches `target_slot_bytes` (default 2 MiB), or after `max_slot_age` (default 60 s) regardless of rows. At a trickle of writes a slot therefore seals about once a minute instead of every window. The minimum age uses a local monotonic clock: an empty observation clears the timer, and a newly elected leader starts a fresh window. `max_slot_age` also counts from the slot's `rotated_at` in the rotation ledger, on the database clock, so frequent leader changes cannot keep a trickle slot open past it. Under compat (pre-flip) authority an older rotator can leave the current generation without a ledger row; the ceiling then relies on the local observation, which a leader change restarts. The size measurement includes heap, indexes, TOAST, free-space maps and visibility maps. `RingRotationCadence::EAGER` restores eligibility on every tick. Queue pruning queues a bounded exclusive-lock request on routing metadata first, then acquires every other child with `NOWAIT`. Contention rolls back the attempt and retries next tick; active-data skip gates retain their backoff. This lets segment-first claim readers and entry-first producers finish without a prune lock cycle. It does not guarantee progress under perpetual contention. diff --git a/docs/configuration.md b/docs/configuration.md index 930bc6bde..50d6cc037 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -508,6 +508,7 @@ DLQ rows are different: `dlq_entries` is a separate hold table and the maintenan | `cleanup_batch_size` | `1000` | Max rows deleted per pass. Raise for very high throughput; lower if you want gentler IO. | | `dlq_cleanup_batch_size` | `1000` | DLQ-specific batch size. | | `ClientBuilder::terminal_count_rollup_interval` (Rust) / `terminal_count_rollup_interval_ms` (Python kwarg) | `30s` | Queue storage only: how often maintenance folds pending `done_entries` terminal-count deltas for sealed queue slots into compact live counters. Exact reads include pending deltas and retained compact receipt batches, so this affects compaction/WAL pressure rather than correctness. | +| `ClientBuilder::statistics_refresh_interval` (Rust) / `statistics_refresh_interval_ms` (Python kwarg) | `5m` | Queue storage only: how often the maintenance leader runs `ANALYZE` on the two head tables and the partitioned ring parents, which autovacuum never analyses; also run once when a runtime becomes leader. `0` disables it. See [Planner statistics](#planner-statistics). | ## Queue storage tuning @@ -531,6 +532,16 @@ The maintenance leader folds pending `done_entries` deltas from sealed queue slo Raising the interval reduces folded-counter churn but leaves more pending delta rows for exact reads to sum. Lowering it folds counters sooner, which can help if exact-count reads are frequent and the `done_entries` delta ledger is growing faster than maintenance drains it. Under long reader transactions, lowering the interval does not force rollup; the useful fix is to release the pinned snapshot or let queue prune reclaim the whole sealed segment. Compact receipt successes are already retained as compact batch rows, so this interval does not change their hot-path write count. +### Planner statistics + +Autovacuum does not analyse partitioned tables, so the ring parents (`ready_entries`, `done_entries`, `leases`, and the other ten partitioned tables of the schema) have no inherited statistics. It also never reaches its analyze threshold on `queue_claim_heads` and `queue_enqueue_heads`, which hold a few rows per lane, so the planner prices joins against them at its defaults. Admin and metrics queries that join retained ready rows to the heads tables then pick plans for a few hundred rows while reading hundreds of thousands. + +The maintenance leader runs `ANALYZE` on both when it is elected and every `statistics_refresh_interval` (default `5m`). One pass takes well under a second on a small ring and a few seconds on a ring with millions of retained rows. It has no measurable effect on claim or completion throughput. + +- **Locks.** Each relation is analysed in its own transaction with a 2 s `lock_timeout` and a 120 s `statement_timeout`. `SKIP_LOCKED` skips a child that ring prune holds during the per-child analysis, but sampling for the parent's statistics waits for it up to the `lock_timeout`; a call that times out is recorded as `skipped_lock` and retried on the next interval. Prune never waits for the refresh: its own 10 ms lock timeout applies, so an attempt that lands while `ANALYZE` samples the oldest slot's tables returns `blocked` and retries on its usual backoff. `ANALYZE` of a partitioned table samples every child, so that window cannot be avoided. +- **Privileges.** The refresh calls `awa.analyze_queue_storage_relation(schema, relation)`, a `SECURITY DEFINER` function installed by migration v054. The runtime needs EXECUTE on the helper, and the login role needs INSERT, UPDATE and DELETE on the target schema's `queue_claim_heads`; the migration grants EXECUTE to existing runtime grantees. A missing helper or caller privilege skips the pass before making calls and logs one warning. Run `awa migrate` for a missing helper; restore the grants as the schema owner for `skipped_privilege`. If the helper owner may not analyse the tables, transfer the function to the table owner or grant its owner `MAINTAIN` on PostgreSQL 17+. See the [database roles guide](security/database-roles.md) for the ownership command. +- **Metrics.** `awa.maintenance.analyze.attempts` and `awa.maintenance.analyze.duration`, by `awa.analyze.scope` (`heads`, `ring_parent`, `pass`) and `awa.analyze.outcome` (`ok`, `skipped_privilege`, `skipped_lock`, `unavailable`, `failed`). A sustained `skipped_lock` rate means something else holds the tables; any `skipped_privilege` or `unavailable` means the refresh is not running. + ### Rust ```rust diff --git a/docs/grafana/README.md b/docs/grafana/README.md index 0ca4ed331..500b13047 100644 --- a/docs/grafana/README.md +++ b/docs/grafana/README.md @@ -207,6 +207,8 @@ In Prometheus they appear under the OpenTelemetry [OTLP-to-Prometheus translatio | `awa.maintenance.rescues` | Counter | rescue_kind | Jobs rescued | | `awa.maintenance.prune.attempts` | Counter | ring, outcome, reason | Ring prune outcomes, including destructive `pruned` and no-DDL `already_pruned` | | `awa.maintenance.prune.duration` | Histogram (s) | ring, phase | Successful destructive prune time split into `lock`, `truncate`, and `commit` | +| `awa.maintenance.analyze.attempts` | Counter | scope (`heads`, `ring_parent`, `pass`), outcome | Planner-statistics `ANALYZE` calls on the head tables and ring parents (`ok`, `skipped_privilege`, `skipped_lock`, `unavailable`, `failed`) | +| `awa.maintenance.analyze.duration` | Histogram (s) | scope, outcome | Duration of one planner-statistics `ANALYZE` call | ### Heartbeat diff --git a/docs/queue-storage-substrate.md b/docs/queue-storage-substrate.md index 4476712e4..9542e3914 100644 --- a/docs/queue-storage-substrate.md +++ b/docs/queue-storage-substrate.md @@ -40,7 +40,7 @@ awa storage prepare-queue-storage-schema \ The command invokes the idempotent `awa.install_queue_storage_substrate(...)` helper under a per-schema advisory transaction lock. The helper is activation-neutral: preparing a schema does not route work to it. Activation is a separate staged transition described in [Upgrading from 0.5 to 0.6](upgrade-0.5-to-0.6.md). -The installer is `SECURITY INVOKER`; its caller needs DDL privileges on the target schema. After the migrator provisions each queue with `awa storage prepare-queue`, workers need runtime DML, sequence `USAGE`/`SELECT`/`UPDATE`, and `TRUNCATE` for guarded ring reclamation. They do not need schema CREATE. See [Database roles](security/database-roles.md). +The installer is `SECURITY INVOKER`; its caller needs DDL privileges on the target schema. Workers need runtime DML, sequence `USAGE`/`SELECT`/`UPDATE`, and `TRUNCATE` for guarded ring reclamation. They do not need schema CREATE: lanes are created through the schema owner's provisioning functions. See [Database roles](security/database-roles.md). ## Operator rules diff --git a/docs/security/database-roles.md b/docs/security/database-roles.md index b6c1557e9..9d5abdca8 100644 --- a/docs/security/database-roles.md +++ b/docs/security/database-roles.md @@ -63,36 +63,68 @@ ALTER DEFAULT PRIVILEGES FOR ROLE awa_owner IN SCHEMA awa If an existing installation created objects as `awa_migrator` without first setting the owner role, transfer every existing schema object to `awa_owner` before relying on owner-scoped defaults. New migrations should use the effective-role command above. -## Provision queues before starting restricted runtimes +## Planner statistics -From schema v049, runtime calls use existing lane sequences without issuing -`CREATE SEQUENCE`. Run provisioning as the migrator for every queue before -starting its producers or workers: +The maintenance leader runs `ANALYZE` on the queue-storage head tables and partitioned ring parents (see [Planner statistics](../configuration.md#planner-statistics)). `ANALYZE` skips a table its caller may not analyse, so a runtime that does not own the tables cannot do this itself; the leader calls `awa.analyze_queue_storage_relation(schema, relation)` instead, a `SECURITY DEFINER` function from migration v054 owned by whichever role ran the migration. The runtime does not need `MAINTAIN`. The function is not executable by `PUBLIC`, because the lock it takes lasts until the caller's transaction ends and would block ring prune: the migration grants `EXECUTE` to the roles that already hold `INSERT`, `UPDATE` and `DELETE` on the head tables, and the function requires the calling session user to hold them. A runtime role created later gets it from the `GRANT EXECUTE ON ALL FUNCTIONS` above or the default privileges. For a custom queue-storage schema, the role needs write privileges on that schema's `queue_claim_heads`. + +Run migrations with the schema owner as the effective role, as shown above, so the function's owner owns the tables. A function owned by a role that may not analyse them (it is not the table owner, the database owner, or a `MAINTAIN` holder) makes every pass report `skipped_privilege` and logs one warning. The function only analyses the two head tables and the thirteen ring parents of a schema that `awa.queue_storage_schemas()` reports as trusted, and `ANALYZE` changes statistics only. + +If the helper's owner may not analyse the tables, run `ALTER FUNCTION awa.analyze_queue_storage_relation(text, text) OWNER TO awa_owner` as a role permitted to transfer ownership, or grant its owner `MAINTAIN` on the target tables on PostgreSQL 17+. Replacing the function body leaves its owner unchanged. + +## Lane sequences + + +Each queue lane, `(queue, priority, enqueue_shard)`, has an enqueue and a claim sequence. From schema v053 they are created by the schema owner, whichever role first enqueues into, claims from or promotes into the lane: the lane helpers call `provision_lane` and `provision_lane_sequences` in the queue-storage schema. These are `SECURITY DEFINER` functions owned by the schema owner with a pinned `search_path`. Like the other runtime-callable functions they are executable by `PUBLIC`, so schema `USAGE` is what gates them, and they only create lanes, after checking that the queue, priority and shard are set and the shard is 0–63. A runtime therefore needs neither schema `CREATE` nor default privileges to start using a new queue, priority or shard. + +A new lane sequence gets the privileges of the head table it backs. A role that may `INSERT` or `UPDATE` `queue_enqueue_heads` gets `USAGE, SELECT, UPDATE` on the enqueue sequences, and one that may only `SELECT` it gets `SELECT`; `queue_claim_heads` does the same for the claim sequences. Granting the runtime the table privileges above is enough, and no sequence default privileges are needed for lanes. + +Run migrations with the schema owner as the effective role, as above, so the provisioning functions belong to it. A migrator that cannot `SET ROLE` to the owner leaves them owned by itself and the migration warns. + +`awa storage prepare-queue` still creates every priority and shard of a queue in one transaction, and any runtime role may run it. Use it to create lanes before a deploy rather than on first use. + +### Check and repair lane ownership + +Lanes created before v053 belong to whichever role touched them first, and only that role, or the roles its default privileges named, can use them. A different runtime role then fails with `permission denied for sequence` on that lane. List lane sequences that are not owned by the schema owner or lack a runtime grant (read-only, any role): ```bash -PGOPTIONS='-c role=awa_owner' \ - awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage prepare-queue \ - --queue email --enqueue-shards 1 +awa --database-url "$AWA_DATABASE_URL" storage lane-sequences --schema awa +``` + +```sql +SELECT * FROM awa.lane_sequence_problems('awa'); ``` -The command provisions all four priorities in one transaction and is safe to -repeat on a live queue: existing sequence cursors keep their positions. It does -not change routing or the queue's configured shard count. Prepare additional -shards before increasing `enqueue_shards`; match `--queue-stripe-count` to the -runtime's queue-storage configuration. With `PartitionedQueue`, provision each -physical queue. Use `--schema` for a custom storage schema. - -Apply the sequence grants above after provisioning, or set the owner's default -privileges first. Lane cursors need `UPDATE` for `setval`, in addition to `USAGE` -and `SELECT`. `CREATE SEQUENCE IF NOT EXISTS` requires schema CREATE even when -the sequence already exists, so provisioning alone does not repair pre-v049 -helpers: apply the migration first. Schema owners retain lazy creation for -single-role installations; a restricted runtime encountering an unprepared -lane fails with a provisioning hint. +The command exits 2 while any remain. +`awa migrate` repairs them once the migration commits, one sequence per short transaction: it gives each sequence to the schema owner and adds its missing grants, wherever it is allowed to. Reassigning takes the sequence's `ACCESS EXCLUSIVE` lock, so the repair waits at most 300 ms for it. A sequence that an open enqueue or claim holds for longer is skipped, left unchanged and reported as busy. A sequence the repair is not allowed to change is reported with the roles that can finish it. Neither fails the migration. Run the repair again at any time; it is idempotent: + +```bash +awa --database-url "$AWA_MIGRATOR_DATABASE_URL" storage lane-sequences --repair +``` + +Which role can repair a sequence: + +- Reassigning it needs a role that holds the current owner's privileges (a member of the owner, with inheritance) and can `SET ROLE` to the schema owner. +- Adding its grants needs the current owner's privileges. + +The schema owner holds neither the migration role's nor an application role's privileges, so a repair running as `awa_owner` cannot reassign sequences those roles created. That includes the migration role connected with `PGOPTIONS='-c role=awa_owner'`. Run the repair as the migration role itself, without `PGOPTIONS`, as above. `awa migrate` and `--repair` retry each sequence the effective role could not repair as the login role, so `awa migrate` run with `PGOPTIONS='-c role=awa_owner'`, as documented, still repairs the migration role's sequences. + +For each sequence that remains, the output lists `repair_roles`, the login roles other than superusers that can repair it, and the statements to run. A busy sequence only needs another `--repair`. When `repair_roles` is empty, which is usual for a sequence an application role created on managed PostgreSQL, either: + +- run `missing_grants` as the current owner. Every runtime role can then use the lane, although the owner stays unchanged; or +- give the migration role the owner's privileges, for example `GRANT app_role TO awa_migrator` from a role with `ADMIN` on `app_role`, rerun `--repair` as the migration role, then revoke the membership. + +Without the CLI, repair one sequence per transaction from `psql` in autocommit mode, connected as the migration role: + +```sql +SELECT format('SELECT * FROM awa.repair_lane_sequence(%L, %L)', 'awa', sequence_name) +FROM awa.lane_sequence_problems('awa') \gexec +``` + +## Custom queue-storage schemas ## Custom queue-storage schemas -Repeat the schema, table, sequence, function, and default-privilege grants for a custom queue-storage schema. Prepare custom schemas as the migrator; the runtime needs DML, sequence access, function execution, and `TRUNCATE`, but never DDL. See [Queue storage](../queue-storage-substrate.md). +Repeat the schema, table, sequence, function, and default-privilege grants for a custom queue-storage schema. Prepare custom schemas as the migrator; the runtime needs DML, sequence access, function execution, and `TRUNCATE`, but never DDL. Check lanes in a custom schema with `awa storage lane-sequences --schema `. See [Queue storage](../queue-storage-substrate.md). ## Why the grants are broad diff --git a/docs/upgrade-0.6-to-0.7.md b/docs/upgrade-0.6-to-0.7.md index 0afc48afb..70b2b4f09 100644 --- a/docs/upgrade-0.6-to-0.7.md +++ b/docs/upgrade-0.6-to-0.7.md @@ -14,14 +14,44 @@ A [migrate-first rehearsal with released `awa-pg==0.6.7`](https://github.com/har verified a restricted worker completing jobs on the upgraded schema. That rehearsal does not cover concurrent mixed-version rollouts. -For a runtime without DDL privileges, apply v049 as the migrator, run -`awa storage prepare-queue --queue ` for every physical queue, and grant -sequence `USAGE, SELECT, UPDATE` before starting producers or workers. Match -shard and stripe counts and provision extra lanes before increasing them. External migration runners must apply the complete generated `awa migrate --sql --from 48 --to 49` bundle: it refreshes the installer -definition before replacing lane helpers in the installed schemas. -See [Database roles](security/database-roles.md#provision-queues-before-starting-restricted-runtimes). +definition before replacing lane helpers in the installed schemas. From v053, +runtimes without DDL privileges create lanes themselves through the schema +owner's provisioning functions, so `awa storage prepare-queue` is optional. + +## Lane ownership and v053 + +Before v053 a lane's sequences belonged to the role that first enqueued into, +claimed from or promoted into it. With several runtime roles, only roles whose +default privileges covered that role could use the lane; any other role failed +with `permission denied for sequence`. v053 installs `SECURITY DEFINER` +`provision_lane` and `provision_lane_sequences` functions, owned by the schema +owner, in each queue-storage schema. The lane helpers create sequences through +them, with grants that mirror the lane's head table. Once v053 commits, +`awa migrate` gives every existing lane sequence to the schema owner and adds +its missing grants, as far as it is allowed to, one sequence per short +transaction with a bounded lock wait. + +0.6 and earlier 0.7 binaries call the same helpers, so lanes they create after +v053 are owned by the schema owner too. A 0.7 binary running ahead of v053 +detects that `provision_lane` is missing and creates lanes as before. + +To adopt it: + +1. Check the current state: `awa storage lane-sequences` (read-only, exits 2 + while lanes need repair). +2. Run `awa migrate` with the schema owner as the effective role, as described + in [Database roles](security/database-roles.md). +3. Read the migrate output. Any lane sequence it could not repair is listed in + a warning. A busy one needs another run of + `awa storage lane-sequences --repair` as the migration role (without + `PGOPTIONS`); the others name the roles that can repair them. See + [Check and repair lane ownership](security/database-roles.md#check-and-repair-lane-ownership). + +External migration runners apply the generated `awa migrate --sql --from 52 +--to 53` bundle, which does not repair existing lane sequences, then run +`awa storage lane-sequences --repair` as the migration role itself. ## Custom-schema discovery during migration @@ -253,6 +283,24 @@ workers report not ready until the migration runs, as for any version bump. External migration runners apply the complete `awa migrate --sql --from 51 --to 52` bundle, which defines the helper before the tail calls it. +## v054: planner statistics helper + +v054 adds `awa.analyze_queue_storage_relation(schema, relation)`, a +`SECURITY DEFINER` function owned by the migration role that lets the +maintenance leader refresh planner statistics for the head tables and ring +parents when the runtime role does not own them (see +[Planner statistics](configuration.md#planner-statistics)). It creates one +function, revokes `EXECUTE` from `PUBLIC` and grants it to the roles that +already hold INSERT, UPDATE and DELETE on the `awa` head tables (the runtime +grantees); it takes no table lock and completes in milliseconds. Roles created +later need `GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa`. Run the migration as the schema owner so the +function's owner owns the tables. + +Binaries before v054 never call it. A 0.7 binary running before the migration +logs once that the function is missing and skips the refresh; the +`awa.maintenance.analyze.attempts` counter reports `unavailable`. External +migration runners apply the `awa migrate --sql --from 53 --to 54` bundle. + ## The one gate `awa migrate` on a 0.7 binary refuses to apply migrations unless one of these holds: diff --git a/skills/awa-operations/SKILL.md b/skills/awa-operations/SKILL.md index 4cfb039b0..cdb92aa3b 100644 --- a/skills/awa-operations/SKILL.md +++ b/skills/awa-operations/SKILL.md @@ -91,11 +91,14 @@ Do not skip a major version. 0.5 → 0.7 is unsupported; step through 0.6. ## Restricted Database Roles -When deploying a runtime without schema CREATE, apply v049 and provision every -queue with `awa storage prepare-queue` as the migrator before starting producers -or workers. Match shard/stripe counts; provision additional lanes before changing -routing. Use the sequence grants and owner defaults in -[Database roles](../../docs/security/database-roles.md#provision-queues-before-starting-restricted-runtimes). +From v053 any runtime role creates lanes through the schema owner's +`SECURITY DEFINER` provisioning functions, so runtimes need neither schema +CREATE nor sequence default privileges, and `awa storage prepare-queue` is +optional. Run migrations with the schema owner as the effective role. Check +lane ownership with `awa storage lane-sequences` (exits 2 while lanes need +repair) and fix it with `--repair`, connected as the migration role itself +rather than as the schema owner, or with the statements it prints; see +[Database roles](../../docs/security/database-roles.md#check-and-repair-lane-ownership). ## Rolling Upgrades @@ -135,7 +138,8 @@ processes. Exactly one worker per database is the elected maintenance leader (session-scoped advisory lock — if it dies, another is elected on the next tick). The leader alone runs cluster-wide scans: deferred→ready promotion, the three rescue paths (stale-heartbeat, hard-deadline, callback-timeout), ring -rotation and prune, DLQ and descriptor cleanup, and cron evaluation. During an +rotation and prune, planner-statistics refresh, DLQ and descriptor cleanup, and +cron evaluation. During an unfinalized transition it promotes both planes' deferred backlogs, so drain and post-flip execution do not depend on which runtime holds leadership. Every worker dispatches and heartbeats its own attempts. There is no `pg_cron` or @@ -313,6 +317,19 @@ storage limit. Queue prune lock contention retries on the next decision tick; active-data gates retain exponential backoff. Persistent contention can still stall reclamation. See `docs/configuration.md` for the scheduling contract. +On queue storage the leader runs `ANALYZE` on `queue_claim_heads`, +`queue_enqueue_heads` and the partitioned ring parents when elected and every +`statistics_refresh_interval` (default 5 minutes; Python +`statistics_refresh_interval_ms`; zero disables), because autovacuum never +analyses them. It goes through the `SECURITY DEFINER` function +`awa.analyze_queue_storage_relation` from migration v054, so a non-owner runtime +role needs no `MAINTAIN` grant. Watch `awa.maintenance.analyze.attempts`: any +`unavailable` means the helper is missing: run `awa migrate` as the schema +owner. For `skipped_privilege`, restore helper EXECUTE and the login role's +head-table DML, or correct the helper owner; migration reruns preserve an +existing function's owner. See `docs/configuration.md#planner-statistics` and +`docs/security/database-roles.md`. + ## Managed Postgres Prefer PostgreSQL 18 where available; PG17 on some managed platforms caps