From 5b9ea50083871426960844e6e9c1c99c2b9907c5 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sun, 8 Nov 2026 22:00:35 +1300 Subject: [PATCH 1/6] perf(insert): hold a lane lock only for the lane write in batch inserts (0.6) insert_many_copy called awa.insert_job_compat once per row. The function takes the (queue, priority, shard) lane lock in reserve_enqueue_seq on the batch's first ready row, and the lock lasts to commit, so it was held while every remaining row was processed and producers on one lane queued behind whole batches. The insert_jobs_compat_batch schema patch (awa 0.7 migration v055, without its schema_version row) adds awa.insert_jobs_compat_batch. It allocates job ids, claims unique keys in key order, registers new lanes and builds the result rows first, then reserves each lane's sequence range once, in (queue, priority, shard) order, and writes the ready rows and segments in two statements. It reads the ready-ring cursor through ring_cursor() where a later schema provides one. insert_many_copy probes its own connection on every call and uses the function when it exists and the role may execute it, otherwise the per-row path. The new insert_many_skip_unique_conflicts requires it. Generic insert_many keeps its per-row statement on every schema. Released 0.6 binaries never call the function and share the same lane lock. --- .../patches/insert_jobs_compat_batch.sql | 588 ++++++++++++++++++ awa-model/src/insert.rs | 262 +++++++- awa-model/src/lib.rs | 5 +- awa-model/src/migrations.rs | 7 + awa/src/lib.rs | 13 +- awa/tests/enqueue_lane_lock_order_test.rs | 203 +++++- awa/tests/insert_many_batch_test.rs | 585 +++++++++++++++++ awa/tests/insert_many_schema_fallback_test.rs | 434 +++++++++++++ 8 files changed, 2039 insertions(+), 58 deletions(-) create mode 100644 awa-model/migrations/patches/insert_jobs_compat_batch.sql create mode 100644 awa/tests/insert_many_batch_test.rs create mode 100644 awa/tests/insert_many_schema_fallback_test.rs diff --git a/awa-model/migrations/patches/insert_jobs_compat_batch.sql b/awa-model/migrations/patches/insert_jobs_compat_batch.sql new file mode 100644 index 000000000..f3b718311 --- /dev/null +++ b/awa-model/migrations/patches/insert_jobs_compat_batch.sql @@ -0,0 +1,588 @@ +-- Schema patch insert_jobs_compat_batch: a batch producer insert that holds +-- each enqueue lane lock only for the lane write (awa 0.7 migration v055, +-- minus the schema_version row). +-- +-- awa.insert_job_compat takes the (queue, priority, shard) lane lock in +-- reserve_enqueue_seq on a batch's first row and, because the lock is +-- transaction-scoped, holds it while every remaining row is processed. Many +-- concurrent producers on one lane therefore queue behind whole batches. This +-- patch adds awa.insert_jobs_compat_batch, which does the lane-independent +-- work first (job ids, unique claims in key order, lane rows, payloads), then +-- reserves each lane's sequence range once, in (queue, priority, shard) order, +-- and writes the ready rows and segments in two statements. The lock is held +-- across the reservation, those two statements and the commit. +-- +-- The column arrays are parallel and must have equal length. Every element is +-- required except run_at, unique_key, unique_states and ordering_hash; +-- insert_job_compat's per-row defaults are the caller's to apply. +-- +-- p_skip_unique_conflicts inserts a batch while skipping rows whose unique key +-- is already claimed (the rows are absent from the result) instead of failing +-- the whole batch with a unique violation. +-- +-- Live load: one CREATE OR REPLACE FUNCTION, no table DDL, no table locks. +-- +-- 0.6 delivery: the 0.6 series cannot add a migration version, so this ships as +-- an idempotent schema patch. `awa migrate` applies it when the schema is +-- exactly at v040 and the probe in migrations.rs reports it missing; +-- `awa migrate --sql` prints it as a repeatable R__ script for external +-- runners. Released 0.6 binaries never call the function: they keep calling +-- awa.insert_job_compat, which this patch does not touch, and both take the +-- same lane lock, so producers on either path share a lane in lane_seq order. +-- Binaries that know the function: insert_many_copy probes its own connection +-- on every call and uses the per-row function until this patch is present; +-- generic insert_many always uses the per-row function; +-- insert_many_skip_unique_conflicts requires this patch. The ready ring cursor +-- is read through ring_cursor() where a later schema provides it, so the +-- function stays correct until a 0.7 migration replaces it. + +CREATE OR REPLACE FUNCTION awa.insert_jobs_compat_batch( + p_kind TEXT[], + p_queue TEXT[], + p_args JSONB[], + p_state TEXT[], + p_priority SMALLINT[], + p_max_attempts SMALLINT[], + p_run_at TIMESTAMPTZ[], + p_metadata JSONB[], + p_tags JSONB[], + p_unique_key BYTEA[], + p_unique_states TEXT[], + p_ordering_hash TEXT[], + p_skip_unique_conflicts BOOLEAN DEFAULT FALSE +) +RETURNS SETOF awa.jobs +AS $$ +DECLARE + v_count INT := COALESCE(cardinality(p_kind), 0); + v_schema TEXT; + v_old_search_path TEXT; + v_now TIMESTAMPTZ := clock_timestamp(); + v_job_ids BIGINT[]; + v_skipped BIGINT[] := ARRAY[]::BIGINT[]; + v_lane RECORD; + v_lane_queues TEXT[] := ARRAY[]::TEXT[]; + v_lane_priorities SMALLINT[] := ARRAY[]::SMALLINT[]; + v_lane_shards SMALLINT[] := ARRAY[]::SMALLINT[]; + v_lane_starts BIGINT[] := ARRAY[]::BIGINT[]; + v_start BIGINT; + v_ready_slot INT; + v_ready_generation BIGINT; + v_row RECORD; + v_shards SMALLINT[]; +BEGIN + IF v_count = 0 THEN + RETURN; + END IF; + + IF cardinality(p_queue) IS DISTINCT FROM v_count + OR cardinality(p_args) IS DISTINCT FROM v_count + OR cardinality(p_state) IS DISTINCT FROM v_count + OR cardinality(p_priority) IS DISTINCT FROM v_count + OR cardinality(p_max_attempts) IS DISTINCT FROM v_count + OR cardinality(p_run_at) IS DISTINCT FROM v_count + OR cardinality(p_metadata) IS DISTINCT FROM v_count + OR cardinality(p_tags) IS DISTINCT FROM v_count + OR cardinality(p_unique_key) IS DISTINCT FROM v_count + OR cardinality(p_unique_states) IS DISTINCT FROM v_count + OR cardinality(p_ordering_hash) IS DISTINCT FROM v_count + THEN + RAISE EXCEPTION 'insert_jobs_compat_batch: column arrays must have equal length' + USING ERRCODE = '22023'; + END IF; + + IF EXISTS ( + SELECT 1 + FROM unnest(p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_metadata, p_tags) + AS input(kind, queue, args, state, priority, max_attempts, metadata, tags) + WHERE input.kind IS NULL OR input.queue IS NULL OR input.args IS NULL + OR input.state IS NULL OR input.priority IS NULL OR input.max_attempts IS NULL + OR input.metadata IS NULL OR input.tags IS NULL + ) THEN + RAISE EXCEPTION 'insert_jobs_compat_batch: kind, queue, args, state, priority, max_attempts, metadata and tags must not be NULL' + USING ERRCODE = '22004'; + END IF; + + IF EXISTS (SELECT 1 FROM unnest(p_kind) AS kind WHERE length(kind) > 200) THEN + RAISE EXCEPTION 'job kind length must be <= 200 characters' + USING ERRCODE = '23514'; + END IF; + + IF EXISTS (SELECT 1 FROM unnest(p_queue) AS queue WHERE length(queue) > 200) THEN + RAISE EXCEPTION 'queue name length must be <= 200 characters' + USING ERRCODE = '23514'; + END IF; + + IF EXISTS ( + SELECT 1 FROM unnest(p_priority) AS priority + WHERE priority < 1 OR priority > 4 + ) THEN + RAISE EXCEPTION 'priority must be between 1 and 4' + USING ERRCODE = '23514'; + END IF; + + IF EXISTS ( + SELECT 1 FROM unnest(p_max_attempts) AS max_attempts + WHERE max_attempts < 1 OR max_attempts > 1000 + ) THEN + RAISE EXCEPTION 'max_attempts must be between 1 and 1000' + USING ERRCODE = '23514'; + END IF; + + IF EXISTS ( + SELECT 1 FROM unnest(p_tags) AS tags + WHERE jsonb_array_length(tags) > 20 + ) THEN + RAISE EXCEPTION 'job tags must contain at most 20 values' + USING ERRCODE = '23514'; + END IF; + + v_schema := awa.active_queue_storage_schema(); + + -- Canonical storage has no lane lock to shorten: insert row by row. + IF v_schema IS NULL THEN + IF NOT p_skip_unique_conflicts THEN + RETURN QUERY + SELECT inserted.* + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, + p_run_at, p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + CROSS JOIN LATERAL awa.insert_job_compat( + input.kind, + input.queue, + input.args, + input.state::awa.job_state, + input.priority, + input.max_attempts, + input.run_at, + input.metadata, + ARRAY( + SELECT tag.value + FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, tag_ord) + ORDER BY tag.tag_ord + ), + input.unique_key, + input.unique_states::bit(8), + NULL + ) AS inserted + ORDER BY input.ord; + RETURN; + END IF; + + -- Skipping needs a subtransaction per row. + FOR v_row IN + SELECT + input.ord, + input.kind, + input.queue, + input.args, + input.state, + input.priority, + input.max_attempts, + input.run_at, + input.metadata, + ARRAY( + SELECT tag.value + FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, tag_ord) + ORDER BY tag.tag_ord + ) AS tags, + input.unique_key, + input.unique_states + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, + p_run_at, p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + ORDER BY input.ord + LOOP + BEGIN + RETURN NEXT awa.insert_job_compat( + v_row.kind, + v_row.queue, + v_row.args, + v_row.state::awa.job_state, + v_row.priority, + v_row.max_attempts, + v_row.run_at, + v_row.metadata, + v_row.tags, + v_row.unique_key, + v_row.unique_states::bit(8), + NULL + ); + EXCEPTION WHEN unique_violation THEN + NULL; + END; + END LOOP; + RETURN; + END IF; + + IF EXISTS ( + SELECT 1 FROM unnest(p_state) AS state + WHERE state NOT IN ('available', 'scheduled', 'retryable') + ) THEN + RAISE EXCEPTION 'queue storage does not support initial state %', + (SELECT state FROM unnest(p_state) AS state + WHERE state NOT IN ('available', 'scheduled', 'retryable') LIMIT 1) + USING ERRCODE = '22023'; + END IF; + + v_old_search_path := current_setting('search_path'); + PERFORM set_config('search_path', format('%I,awa,public', v_schema), true); + + -- Everything up to the lane reservation is lane-independent and runs + -- without a lane lock. + EXECUTE format( + 'SELECT array_agg(nextval(%L::regclass)::bigint) FROM generate_series(1, $1)', + format('%I.job_id_seq', v_schema) + ) + INTO v_job_ids + USING v_count; + + -- Rows that enqueue as `available` go to a lane. Their shard is derived as + -- insert_job_compat derives it: ordering-key rows from the key hash, the + -- rest from the backend pid. + SELECT array_agg( + CASE WHEN input.state = 'available' THEN + MOD( + CASE WHEN input.ordering_hash IS NULL + THEN pg_backend_pid()::numeric + ELSE input.ordering_hash::numeric + END, + COALESCE(meta.enqueue_shards, 1)::numeric + )::smallint + END + ORDER BY input.ord + ) + INTO v_shards + FROM unnest(p_queue, p_state, p_ordering_hash) + WITH ORDINALITY AS input(queue, state, ordering_hash, ord) + LEFT JOIN awa.queue_meta AS meta ON meta.queue = input.queue; + + -- Claims go in unique-key order, so two producers with overlapping keys + -- cannot wait on each other in opposite orders. + IF p_skip_unique_conflicts THEN + WITH claimed AS ( + INSERT INTO awa.job_unique_claims (unique_key, job_id) + SELECT input.unique_key, v_job_ids[input.ord] + FROM unnest(p_unique_key, p_unique_states, p_state) + WITH ORDINALITY AS input(unique_key, unique_states, state, ord) + WHERE input.unique_key IS NOT NULL + AND input.unique_states IS NOT NULL + AND awa.job_state_in_bitmask(input.unique_states::bit(8), input.state::awa.job_state) + ORDER BY input.unique_key, input.ord + ON CONFLICT (unique_key) DO NOTHING + RETURNING job_id + ) + SELECT COALESCE(array_agg(wanted.job_id), ARRAY[]::BIGINT[]) + INTO v_skipped + FROM ( + SELECT v_job_ids[input.ord] AS job_id + FROM unnest(p_unique_key, p_unique_states, p_state) + WITH ORDINALITY AS input(unique_key, unique_states, state, ord) + WHERE input.unique_key IS NOT NULL + AND input.unique_states IS NOT NULL + AND awa.job_state_in_bitmask(input.unique_states::bit(8), input.state::awa.job_state) + ) AS wanted + WHERE NOT EXISTS (SELECT 1 FROM claimed WHERE claimed.job_id = wanted.job_id); + ELSE + INSERT INTO awa.job_unique_claims (unique_key, job_id) + SELECT input.unique_key, v_job_ids[input.ord] + FROM unnest(p_unique_key, p_unique_states, p_state) + WITH ORDINALITY AS input(unique_key, unique_states, state, ord) + WHERE input.unique_key IS NOT NULL + AND input.unique_states IS NOT NULL + AND awa.job_state_in_bitmask(input.unique_states::bit(8), input.state::awa.job_state) + ORDER BY input.unique_key, input.ord; + END IF; + + -- Register lanes this batch is the first to use. The head tables carry a + -- row trigger that takes the lane lock before the conflict check, so only + -- missing lanes are proposed: a batch for an existing lane must not queue + -- behind that lane's producers before it has any work to show for it. + INSERT INTO queue_lanes (queue, priority) + SELECT lanes.queue, lanes.priority + FROM ( + SELECT DISTINCT input.queue, input.priority + FROM unnest(p_queue, p_priority, p_state) + WITH ORDINALITY AS input(queue, priority, state, ord) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS lanes + WHERE NOT EXISTS ( + SELECT 1 FROM queue_lanes AS existing + WHERE existing.queue = lanes.queue AND existing.priority = lanes.priority + ) + ORDER BY lanes.queue COLLATE "C", lanes.priority + ON CONFLICT ON CONSTRAINT queue_lanes_pkey DO NOTHING; + + INSERT INTO queue_enqueue_heads (queue, priority, enqueue_shard) + SELECT lanes.queue, lanes.priority, lanes.shard + FROM ( + SELECT DISTINCT input.queue, input.priority, v_shards[input.ord] AS shard + FROM unnest(p_queue, p_priority, p_state) + WITH ORDINALITY AS input(queue, priority, state, ord) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS lanes + WHERE NOT EXISTS ( + SELECT 1 FROM queue_enqueue_heads AS existing + WHERE existing.queue = lanes.queue + AND existing.priority = lanes.priority + AND existing.enqueue_shard = lanes.shard + ) + ORDER BY lanes.queue COLLATE "C", lanes.priority, lanes.shard + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + + INSERT INTO queue_claim_heads (queue, priority, enqueue_shard) + SELECT lanes.queue, lanes.priority, lanes.shard + FROM ( + SELECT DISTINCT input.queue, input.priority, v_shards[input.ord] AS shard + FROM unnest(p_queue, p_priority, p_state) + WITH ORDINALITY AS input(queue, priority, state, ord) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS lanes + WHERE NOT EXISTS ( + SELECT 1 FROM queue_claim_heads AS existing + WHERE existing.queue = lanes.queue + AND existing.priority = lanes.priority + AND existing.enqueue_shard = lanes.shard + ) + ORDER BY lanes.queue COLLATE "C", lanes.priority, lanes.shard + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + + -- Re-enqueues that are not yet due carry no lane. + INSERT INTO deferred_jobs ( + job_id, kind, queue, args, state, priority, attempt, run_lease, + max_attempts, run_at, attempted_at, finalized_at, created_at, + unique_key, unique_states, payload + ) + SELECT + v_job_ids[input.ord], + input.kind, + input.queue, + input.args, + input.state::awa.job_state, + input.priority, + 0, + 0, + input.max_attempts, + COALESCE(input.run_at, v_now), + NULL, + NULL, + v_now, + input.unique_key, + input.unique_states, + jsonb_build_object( + 'metadata', input.metadata, + 'tags', input.tags, + 'errors', '[]'::jsonb, + 'progress', NULL + ) + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_run_at, + p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + WHERE input.state <> 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped); + + -- The result rows do not depend on the lanes, so build them before the + -- lock window opens. + RETURN QUERY + SELECT + v_job_ids[input.ord], + input.kind, + input.queue, + input.args, + input.state::awa.job_state, + input.priority, + 0::smallint, + input.max_attempts, + COALESCE(input.run_at, v_now), + NULL::timestamptz, + NULL::timestamptz, + NULL::timestamptz, + NULL::timestamptz, + v_now, + ARRAY[]::jsonb[], + input.metadata, + ARRAY( + SELECT tag.value + FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, tag_ord) + ORDER BY tag.tag_ord + ), + input.unique_key, + input.unique_states::bit(8), + NULL::uuid, + NULL::timestamptz, + NULL::text, + NULL::text, + NULL::text, + NULL::text, + 0::bigint, + NULL::jsonb + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_run_at, + p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + WHERE v_job_ids[input.ord] <> ALL (v_skipped) + ORDER BY input.ord; + + -- Lane lock window: reserve each lane's sequence range once, in + -- (queue, priority, shard) order, then write the ready rows. The locks are + -- held to the end of the transaction, which keeps the rows of a lane + -- visible in lane_seq order. + FOR v_lane IN + SELECT lanes.queue, lanes.priority, lanes.shard, lanes.row_count + FROM ( + SELECT input.queue, input.priority, v_shards[input.ord] AS shard, + count(*) AS row_count + FROM unnest(p_queue, p_priority, p_state) + WITH ORDINALITY AS input(queue, priority, state, ord) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + GROUP BY input.queue, input.priority, v_shards[input.ord] + ) AS lanes + ORDER BY lanes.queue COLLATE "C", lanes.priority, lanes.shard + LOOP + EXECUTE format('SELECT %I.reserve_enqueue_seq($1, $2, $3, $4)', v_schema) + INTO v_start + USING v_lane.queue, v_lane.priority, v_lane.shard, v_lane.row_count::BIGINT; + + v_lane_queues := v_lane_queues || v_lane.queue; + v_lane_priorities := v_lane_priorities || v_lane.priority; + v_lane_shards := v_lane_shards || v_lane.shard; + v_lane_starts := v_lane_starts || v_start; + END LOOP; + + IF cardinality(v_lane_queues) > 0 THEN + IF to_regprocedure(format('%I.ring_cursor(text)', v_schema)) IS NOT NULL THEN + EXECUTE format('SELECT slot, generation FROM %I.ring_cursor(''queue'')', v_schema) + INTO v_ready_slot, v_ready_generation; + ELSE + SELECT ring.current_slot, ring.generation + INTO v_ready_slot, v_ready_generation + FROM queue_ring_state AS ring + WHERE ring.singleton = TRUE; + END IF; + + INSERT INTO ready_entries ( + ready_slot, ready_generation, job_id, kind, queue, args, priority, + attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, + attempted_at, created_at, unique_key, unique_states, payload + ) + SELECT + v_ready_slot, + v_ready_generation, + v_job_ids[placed.ord], + placed.kind, + placed.queue, + placed.args, + placed.priority, + 0, + 0, + placed.max_attempts, + lane.start_seq + placed.lane_rank - 1, + lane.shard, + COALESCE(placed.run_at, v_now), + NULL, + v_now, + placed.unique_key, + placed.unique_states, + jsonb_build_object( + 'metadata', placed.metadata, + 'tags', placed.tags, + 'errors', '[]'::jsonb, + 'progress', NULL + ) + FROM ( + SELECT + input.*, + v_shards[input.ord] AS shard, + row_number() OVER ( + PARTITION BY input.queue, input.priority, v_shards[input.ord] + ORDER BY input.ord + ) AS lane_rank + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_run_at, + p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS placed + JOIN unnest(v_lane_queues, v_lane_priorities, v_lane_shards, v_lane_starts) + AS lane(queue, priority, shard, start_seq) + ON lane.queue = placed.queue + AND lane.priority = placed.priority + AND lane.shard = placed.shard; + + -- One segment per run of consecutive lane_seq values sharing a run_at. + INSERT INTO ready_segments ( + ready_slot, ready_generation, queue, priority, enqueue_shard, + first_lane_seq, next_lane_seq, first_run_at + ) + SELECT + v_ready_slot, + v_ready_generation, + runs.queue, + runs.priority, + runs.shard, + min(runs.lane_seq), + max(runs.lane_seq) + 1, + runs.run_at + FROM ( + SELECT + placed.*, + placed.lane_seq - row_number() OVER ( + PARTITION BY placed.queue, placed.priority, placed.shard, placed.run_at + ORDER BY placed.lane_seq + ) AS run_id + FROM ( + SELECT + input.queue, + input.priority, + v_shards[input.ord] AS shard, + COALESCE(input.run_at, v_now) AS run_at, + lane.start_seq + row_number() OVER ( + PARTITION BY input.queue, input.priority, v_shards[input.ord] + ORDER BY input.ord + ) - 1 AS lane_seq + FROM unnest(p_queue, p_priority, p_state, p_run_at) + WITH ORDINALITY AS input(queue, priority, state, run_at, ord) + JOIN unnest(v_lane_queues, v_lane_priorities, v_lane_shards, v_lane_starts) + AS lane(queue, priority, shard, start_seq) + ON lane.queue = input.queue + AND lane.priority = input.priority + AND lane.shard = v_shards[input.ord] + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS placed + ) AS runs + GROUP BY runs.queue, runs.priority, runs.shard, runs.run_at, runs.run_id + ON CONFLICT (ready_slot, ready_generation, queue, priority, enqueue_shard, first_lane_seq) + DO NOTHING; + + PERFORM pg_notify('awa:' || notified.queue, '') + FROM (SELECT DISTINCT unnest(v_lane_queues) AS queue) AS notified; + END IF; + + PERFORM set_config('search_path', v_old_search_path, true); +END; +$$ LANGUAGE plpgsql VOLATILE +SET search_path = pg_catalog, awa, public; diff --git a/awa-model/src/insert.rs b/awa-model/src/insert.rs index 6ee39cfd0..7b7811461 100644 --- a/awa-model/src/insert.rs +++ b/awa-model/src/insert.rs @@ -14,10 +14,21 @@ use sqlx::{PgExecutor, PgPool}; // lanes in opposite orders can deadlock. Every multi-row enqueue therefore // takes its lane locks in one canonical order: `(queue, priority, shard)` // ascending, queues compared bytewise (`COLLATE "C"`, the order of Rust's -// `str`). The batch inserts here sort rows by that key, stably, before they -// reach `awa.insert_job_compat`, and put the rows back in input order for the -// caller. `QueueStorage` batch enqueues group rows in a `BTreeMap` keyed the -// same way. +// `str`). `QueueStorage` batch enqueues group rows in a `BTreeMap` keyed that +// way. +// +// `insert_many_copy` and `insert_many_skip_unique_conflicts` use +// `awa.insert_jobs_compat_batch` (the `insert_jobs_compat_batch` schema +// patch), which does the lane-independent work first (job ids, unique claims, +// lane registration, payloads) and then reserves its lanes in canonical order, +// so a lane is held only across its reservation, the ready-row writes and the +// commit. +// +// Generic `insert_many` always uses the per-row path, as does +// `insert_many_copy` on a schema without the patch: rows are sorted by the +// same key, stably, before they reach `awa.insert_job_compat`, and returned in +// input order. Both paths take the same lane lock, so producers on either, and +// older binaries, share a lane without reordering it. // // The shard is derived as `insert_job_compat` derives it: rows with an // ordering key use `ordering_key_hash64 % enqueue_shards`, other rows @@ -213,9 +224,11 @@ impl PreparedJobInsert { /// [`ordering_key_hash64`] as decimal text (the value exceeds `bigint`) so the /// statement can place keyed rows on the shard `insert_job_compat` will pick. /// -/// Rows are inserted in canonical lane order (see "Lane lock order" above) and -/// returned in input order. -const INSERT_MANY_SQL: &str = r#" +/// One `awa.insert_job_compat` call per row, rows inserted in canonical lane +/// order (see "Lane lock order" above) and returned in input order. It works +/// with and without the `insert_jobs_compat_batch` schema patch, including when +/// one process writes to databases that differ in whether they have it. +const INSERT_MANY_ROW_FUNCTION_SQL: &str = r#" WITH input AS MATERIALIZED ( SELECT raw.*, @@ -429,7 +442,17 @@ fn ordering_hash_text(row: &PreparedRow) -> Option { /// Insert multiple jobs in a single statement. /// /// Supports uniqueness constraints — jobs with `unique` opts will have their -/// `unique_key` and `unique_states` computed and included. +/// `unique_key` and `unique_states` computed and included. A job whose unique +/// key is already claimed fails the whole batch with a unique violation +/// (SQLSTATE `23505`, as [`AwaError::Database`]); use +/// [`insert_many_skip_unique_conflicts`] to insert the rest instead. +/// +/// Calls `awa.insert_job_compat` per row, in canonical lane order, on every +/// supported schema. A generic executor is consumed by one statement and +/// cannot be probed first, so this path never assumes the database has the +/// `insert_jobs_compat_batch` schema patch. [`insert_many_copy`] probes its +/// connection and uses the batch function when it is installed; +/// [`insert_many_skip_unique_conflicts`] requires it. #[tracing::instrument(skip(executor, jobs), fields(job.count = jobs.len()))] pub async fn insert_many<'e, E>(executor: E, jobs: &[InsertParams]) -> Result, AwaError> where @@ -440,7 +463,7 @@ where } let rows = precompute_rows(jobs)?; - let results = sqlx::query_as::<_, JobRow>(INSERT_MANY_SQL) + let results = sqlx::query_as::<_, JobRow>(INSERT_MANY_ROW_FUNCTION_SQL) .bind(rows.iter().map(|row| row.kind.as_str()).collect::>()) .bind( rows.iter() @@ -484,6 +507,132 @@ where Ok(results) } +/// `awa.insert_jobs_compat_batch` over parallel column arrays, one element per +/// job. `$12` is each row's [`ordering_key_hash64`] as decimal text, as in +/// [`INSERT_MANY_ROW_FUNCTION_SQL`]; `$13` skips rows whose unique key is +/// already claimed instead of failing the batch. Rows come back in input +/// order; skipped rows are absent. +const INSERT_MANY_BATCH_SQL: &str = r#" + SELECT * + FROM awa.insert_jobs_compat_batch( + $1::text[], + $2::text[], + $3::jsonb[], + $4::text[], + $5::smallint[], + $6::smallint[], + $7::timestamptz[], + $8::jsonb[], + $9::jsonb[], + $10::bytea[], + $11::text[], + $12::text[], + $13 + ) +"#; + +/// TRUE when `awa.insert_jobs_compat_batch` exists, whoever may call it: the +/// `insert_jobs_compat_batch` schema patch's probe. Catalog lookup only. +pub(crate) const BATCH_FUNCTION_INSTALLED_SQL: &str = "SELECT to_regprocedure(\ + 'awa.insert_jobs_compat_batch(text[],text[],jsonb[],text[],smallint[],smallint[],\ + timestamptz[],jsonb[],jsonb[],bytea[],text[],text[],boolean)') IS NOT NULL"; + +/// TRUE when the connected role may call `awa.insert_jobs_compat_batch`. +const BATCH_FUNCTION_USABLE_SQL: &str = "SELECT COALESCE(has_function_privilege(to_regprocedure(\ + 'awa.insert_jobs_compat_batch(text[],text[],jsonb[],text[],smallint[],smallint[],\ + timestamptz[],jsonb[],jsonb[],bytea[],text[],text[],boolean)'), 'EXECUTE'), FALSE)"; + +/// Insert multiple jobs in a single statement, skipping those whose unique key +/// is already claimed. +/// +/// Returns the rows that were inserted, in input order; a skipped job has no +/// entry in the result. Jobs without `unique` opts are always inserted. A key +/// repeated inside the batch inserts its first job. Unlike [`insert_many`], +/// which fails the batch with a unique violation, this is what a caller that +/// replays work wants, and it takes each lane lock once for the batch rather +/// than once per job as a row-by-row retry does. +/// +/// Requires the `insert_jobs_compat_batch` schema patch (`awa migrate`); on a +/// schema without it, returns [`AwaError::Validation`]. [`insert_many_copy`] +/// skips unique conflicts with or without the patch, through a connection, and +/// is what Python's `insert_many_copy` calls. +#[tracing::instrument(skip(executor, jobs), fields(job.count = jobs.len()))] +pub async fn insert_many_skip_unique_conflicts<'e, E>( + executor: E, + jobs: &[InsertParams], +) -> Result, AwaError> +where + E: PgExecutor<'e>, +{ + if jobs.is_empty() { + return Ok(Vec::new()); + } + let rows = precompute_rows(jobs)?; + match bind_batch_columns(sqlx::query_as::<_, JobRow>(INSERT_MANY_BATCH_SQL), &rows) + .bind(true) + .fetch_all(executor) + .await + { + Ok(results) => Ok(results), + Err(err) if is_missing_batch_function(&err) => Err(AwaError::Validation( + "insert_many_skip_unique_conflicts needs the insert_jobs_compat_batch schema patch; \ + run `awa migrate`" + .to_string(), + )), + Err(err) => Err(err.into()), + } +} + +/// `undefined_function` naming the batch function: the schema predates the +/// `insert_jobs_compat_batch` patch, or the function was dropped. +fn is_missing_batch_function(err: &sqlx::Error) -> bool { + matches!( + err, + sqlx::Error::Database(db_err) + if db_err.code().as_deref() == Some("42883") + && db_err.message().contains("insert_jobs_compat_batch") + ) +} + +fn bind_batch_columns<'q>( + query: sqlx::query::QueryAs<'q, sqlx::Postgres, JobRow, sqlx::postgres::PgArguments>, + rows: &'q [PreparedRow], +) -> sqlx::query::QueryAs<'q, sqlx::Postgres, JobRow, sqlx::postgres::PgArguments> { + query + .bind(rows.iter().map(|row| row.kind.as_str()).collect::>()) + .bind( + rows.iter() + .map(|row| row.queue.as_str()) + .collect::>(), + ) + .bind(rows.iter().map(|row| &row.args).collect::>()) + .bind( + rows.iter() + .map(|row| row.state_db_str()) + .collect::>(), + ) + .bind(rows.iter().map(|row| row.priority).collect::>()) + .bind(rows.iter().map(|row| row.max_attempts).collect::>()) + .bind(rows.iter().map(|row| row.run_at).collect::>()) + .bind(rows.iter().map(|row| &row.metadata).collect::>()) + .bind( + rows.iter() + .map(|row| serde_json::json!(row.tags)) + .collect::>(), + ) + .bind( + rows.iter() + .map(|row| row.unique_key.as_deref()) + .collect::>(), + ) + .bind( + rows.iter() + .map(|row| row.unique_states.as_deref()) + .collect::>(), + ) + .bind(rows.iter().map(ordering_hash_text).collect::>()) +} + /// CTEs that number the staged rows in staging order and re-sort them into /// canonical lane order (see "Lane lock order" above); `staged_in_lane_order` /// carries the staging-order position as `ord`. @@ -517,7 +666,10 @@ macro_rules! staged_in_lane_order_ctes { /// Insert many jobs through the compatibility COPY surface. /// /// Uses a temp staging table with no constraints for fast COPY ingestion, -/// then routes staged rows through `awa.insert_job_compat`. Accepts +/// then inserts the staged rows through `awa.insert_jobs_compat_batch`, or, on +/// a schema without the `insert_jobs_compat_batch` patch, through +/// `awa.insert_job_compat` row by row. Jobs whose unique key is already +/// claimed are skipped and absent from the result. Accepts /// `&mut PgConnection` so callers can use pool connections or transactions /// (Transaction derefs to PgConnection). For high-volume queue-storage /// producers, prefer [`QueueStorage::enqueue_params_copy`](crate::QueueStorage::enqueue_params_copy) @@ -574,10 +726,89 @@ pub async fn insert_many_copy( copy_in.send(csv_buf).await?; copy_in.finish().await?; - // 3. INSERT...SELECT from staging into real table + // 3. Insert the staged rows. With the batch function, unique conflicts are + // skipped in the one statement; without it, each unique row goes through + // the per-row function in its own savepoint. let has_unique = rows.iter().any(|r| r.unique_key.is_some()); + let results = if batch_function_usable(conn).await? { + insert_staged_rows_with_batch_function(conn, has_unique).await? + } else { + insert_staged_rows_per_row(conn, has_unique).await? + }; + + // Keep the session-local staging table reusable across multiple COPY calls + // within the same outer transaction. + sqlx::query("DELETE FROM pg_temp.awa_copy_staging") + .execute(&mut *conn) + .await?; + + Ok(results) +} + +/// Probed on every COPY call, on the caller's own connection: a process can +/// write to databases that differ in whether they have the patch, and a role +/// without EXECUTE on the function keeps the per-row path. +async fn batch_function_usable(conn: &mut PgConnection) -> Result { + Ok(sqlx::query_scalar(BATCH_FUNCTION_USABLE_SQL) + .fetch_one(&mut *conn) + .await?) +} - let results = if has_unique { +async fn insert_staged_rows_with_batch_function( + conn: &mut PgConnection, + skip_unique_conflicts: bool, +) -> Result, AwaError> { + Ok(sqlx::query_as::<_, JobRow>( + r#" + WITH staged AS ( + SELECT *, row_number() OVER () AS ord + FROM pg_temp.awa_copy_staging + ), + columns AS ( + SELECT + array_agg(kind ORDER BY ord) AS kinds, + array_agg(queue ORDER BY ord) AS queues, + array_agg(args ORDER BY ord) AS args, + array_agg(state::text ORDER BY ord) AS states, + array_agg(priority ORDER BY ord) AS priorities, + array_agg(max_attempts ORDER BY ord) AS max_attempts, + array_agg(run_at ORDER BY ord) AS run_ats, + array_agg(metadata ORDER BY ord) AS metadata, + array_agg(to_jsonb(tags) ORDER BY ord) AS tags, + array_agg(unique_key ORDER BY ord) AS unique_keys, + array_agg(unique_states::text ORDER BY ord) AS unique_states, + array_agg(ordering_hash ORDER BY ord) AS ordering_hashes + FROM staged + ) + SELECT inserted.* + FROM columns + CROSS JOIN LATERAL awa.insert_jobs_compat_batch( + columns.kinds, + columns.queues, + columns.args, + columns.states, + columns.priorities, + columns.max_attempts, + columns.run_ats, + columns.metadata, + columns.tags, + columns.unique_keys, + columns.unique_states, + columns.ordering_hashes, + $1 + ) AS inserted + "#, + ) + .bind(skip_unique_conflicts) + .fetch_all(&mut *conn) + .await?) +} + +async fn insert_staged_rows_per_row( + conn: &mut PgConnection, + has_unique: bool, +) -> Result, AwaError> { + let results: Vec = if has_unique { // The compatibility `awa.jobs` surface is now a view backed by hot and // deferred tables, so the old `ON CONFLICT` path is no longer available // here. Keep COPY for staging/parsing, then insert unique rows one at a @@ -718,13 +949,6 @@ pub async fn insert_many_copy( .fetch_all(&mut *conn) .await? }; - - // Keep the session-local staging table reusable across multiple COPY calls - // within the same outer transaction. - sqlx::query("DELETE FROM pg_temp.awa_copy_staging") - .execute(&mut *conn) - .await?; - Ok(results) } diff --git a/awa-model/src/lib.rs b/awa-model/src/lib.rs index b2e69c45b..1f72a4ca5 100644 --- a/awa-model/src/lib.rs +++ b/awa-model/src/lib.rs @@ -40,7 +40,10 @@ pub type QueueStats = QueueOverview; pub use cron::{CronJobRow, CronMissedFirePolicy, PeriodicJob, PeriodicJobBuilder}; pub use dlq::{DlqMetadata, DlqRow, ListDlqFilter, RetryFromDlqOpts}; pub use error::{map_sqlx_error, AwaError}; -pub use insert::{insert, insert_many, insert_many_copy, insert_many_copy_from_pool, insert_with}; +pub use insert::{ + insert, insert_many, insert_many_copy, insert_many_copy_from_pool, + insert_many_skip_unique_conflicts, insert_with, +}; pub use job::{InsertOpts, InsertParams, JobRow, JobState, UniqueOpts}; pub use partitioned_queue::{ partition_for_ordering_key, partition_hash64, PartitionedQueue, PartitionedQueueError, diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index 549dcabef..ab45ee325 100644 --- a/awa-model/src/migrations.rs +++ b/awa-model/src/migrations.rs @@ -426,6 +426,13 @@ pub const SCHEMA_PATCHES: &[SchemaPatch] = &[ ")" ), }, + SchemaPatch { + name: "insert_jobs_compat_batch", + description: + "Batch producer insert with a short enqueue lane lock window (awa 0.7 migration v055)", + sql: include_str!("../migrations/patches/insert_jobs_compat_batch.sql"), + applied_probe: crate::insert::BATCH_FUNCTION_INSTALLED_SQL, + }, ]; /// Schema patches a database still needs to reach this binary's schema shape. diff --git a/awa/src/lib.rs b/awa/src/lib.rs index 51d6b4072..42cae1ff6 100644 --- a/awa/src/lib.rs +++ b/awa/src/lib.rs @@ -12,12 +12,13 @@ pub use awa_model; // Re-export core model types (includes JobArgs derive macro via awa-model) pub use awa_model::{ self as model, adapter, admin, bridge, callback_contract, insert, insert_many, - insert_many_copy, insert_many_copy_from_pool, insert_with, map_sqlx_error, migrations, - partition_for_ordering_key, partition_hash64, prepare_job_insert, prepare_raw_job_insert, - storage, AwaError, CallbackConfig, DefaultAction, DlqMetadata, DlqRow, InsertOpts, - InsertParams, JobArgs, JobKindDescriptor, JobRow, JobState, ListDlqFilter, PartitionedQueue, - PartitionedQueueError, PreparedJobInsert, QueueDescriptor, QueueStorage, QueueStorageConfig, - ResolveOutcome, RetryFromDlqOpts, StorageCapability, StorageStatus, UniqueOpts, + insert_many_copy, insert_many_copy_from_pool, insert_many_skip_unique_conflicts, insert_with, + map_sqlx_error, migrations, partition_for_ordering_key, partition_hash64, prepare_job_insert, + prepare_raw_job_insert, storage, AwaError, CallbackConfig, DefaultAction, DlqMetadata, DlqRow, + InsertOpts, InsertParams, JobArgs, JobKindDescriptor, JobRow, JobState, ListDlqFilter, + PartitionedQueue, PartitionedQueueError, PreparedJobInsert, QueueDescriptor, QueueStorage, + QueueStorageConfig, ResolveOutcome, RetryFromDlqOpts, StorageCapability, StorageStatus, + UniqueOpts, }; // Re-export worker runtime diff --git a/awa/tests/enqueue_lane_lock_order_test.rs b/awa/tests/enqueue_lane_lock_order_test.rs index b5bcb7725..9e09bfb66 100644 --- a/awa/tests/enqueue_lane_lock_order_test.rs +++ b/awa/tests/enqueue_lane_lock_order_test.rs @@ -15,7 +15,8 @@ //! Run with: `cargo test --package awa --test enqueue_lane_lock_order_test -- --nocapture` use awa::model::{ - insert_many, insert_many_copy_from_pool, AwaError, QueueStorage, QueueStorageConfig, + insert_many, insert_many_copy_from_pool, insert_many_skip_unique_conflicts, AwaError, + QueueStorage, QueueStorageConfig, }; use awa::{InsertOpts, InsertParams, JobRow, UniqueOpts}; use awa_testing::setup; @@ -32,6 +33,7 @@ const WAIT_DEADLINE: Duration = Duration::from_secs(20); #[derive(Debug, Clone, Copy)] enum Producer { InsertMany, + SkipUniqueConflicts, Copy, CopyLarge, CopyUnique, @@ -48,6 +50,7 @@ impl Producer { async fn run(self, pool: &PgPool, jobs: &[InsertParams]) -> Result, AwaError> { match self { Producer::InsertMany => insert_many(pool, jobs).await, + Producer::SkipUniqueConflicts => insert_many_skip_unique_conflicts(pool, jobs).await, Producer::Copy | Producer::CopyLarge | Producer::CopyUnique => { insert_many_copy_from_pool(pool, jobs).await } @@ -249,6 +252,11 @@ async fn insert_many_takes_lane_locks_in_canonical_order() { assert_producer_does_not_deadlock_with_queue_storage_batch(Producer::InsertMany).await; } +#[tokio::test] +async fn insert_many_skip_unique_conflicts_takes_lane_locks_in_canonical_order() { + assert_producer_does_not_deadlock_with_queue_storage_batch(Producer::SkipUniqueConflicts).await; +} + #[tokio::test] async fn insert_many_copy_takes_lane_locks_in_canonical_order() { assert_producer_does_not_deadlock_with_queue_storage_batch(Producer::Copy).await; @@ -280,6 +288,74 @@ fn ordering_key_per_shard(shards: i16) -> Vec> { .collect() } +/// Key of the advisory lock `reserve_enqueue_seq` holds for one lane. +async fn lane_reservation_lock_key(pool: &PgPool, queue: &str, priority: i16, shard: i16) -> i64 { + sqlx::query_scalar( + r#" + SELECT hashtextextended( + format('%I.%I', 'awa', awa.queue_lane_sequence_name('queue_enqueue_seq', $1, $2, $3)), + 0 + ) + "#, + ) + .bind(queue) + .bind(priority) + .bind(shard) + .fetch_one(pool) + .await + .expect("derive lane reservation lock key") +} + +/// The backend waiting for the advisory lock `key`, once there is one. +async fn advisory_waiter_pid(pool: &PgPool, key: i64) -> i32 { + let started = Instant::now(); + loop { + let pid: Option = sqlx::query_scalar( + r#" + SELECT pid + FROM pg_locks + WHERE locktype = 'advisory' + AND NOT granted + AND database = (SELECT oid FROM pg_database WHERE datname = current_database()) + AND ((classid::bigint << 32) | objid::bigint) = $1 + LIMIT 1 + "#, + ) + .bind(key) + .fetch_optional(pool) + .await + .expect("poll pg_locks"); + if let Some(pid) = pid { + return pid; + } + assert!( + started.elapsed() < WAIT_DEADLINE, + "no backend started waiting for advisory lock {key}" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } +} + +async fn granted_advisory_keys(pool: &PgPool, pid: i32) -> std::collections::BTreeSet { + sqlx::query_scalar::<_, i64>( + r#" + SELECT (classid::bigint << 32) | objid::bigint + FROM pg_locks + WHERE locktype = 'advisory' AND granted AND pid = $1 + "#, + ) + .bind(pid) + .fetch_all(pool) + .await + .expect("read granted advisory locks") + .into_iter() + .collect() +} + +/// A batch over sharded lanes reserves them in `(queue, priority, shard)` +/// order, whatever order its rows arrive in. A gate holds one lane in the +/// middle of that order; the producer must be parked on it holding exactly the +/// lanes before it. async fn assert_batch_reserves_lanes_in_canonical_order(producer: Producer) { let _guard = LANE_LOCK_TEST_GUARD.lock().await; let (pool, _config) = setup_queue_storage().await; @@ -307,6 +383,7 @@ async fn assert_batch_reserves_lanes_in_canonical_order(producer: Producer) { let keys = ordering_key_per_shard(shards); let mut jobs = Vec::new(); + let mut lanes = std::collections::BTreeSet::new(); let rounds = match producer { Producer::CopyLarge => 700, _ => 12, @@ -315,55 +392,112 @@ async fn assert_batch_reserves_lanes_in_canonical_order(producer: Producer) { for queue in &queues { let mut params = job(queue, round, false); params.opts.priority = if round % 5 == 0 { 1 } else { 2 }; - // Keyed rows spread over every shard; every third row has no key. - if round % 3 != 0 { - params.opts.ordering_key = Some(keys[(round + queue.len()) % keys.len()].clone()); - } + // Every row is keyed, so its shard, and so its lane, is known here. + let shard = (round + queue.len()) % keys.len(); + params.opts.ordering_key = Some(keys[shard].clone()); + lanes.insert(( + queue.as_bytes().to_vec(), + params.opts.priority, + shard as i16, + queue.clone(), + )); jobs.push(params); } } + assert!(lanes.len() >= 6, "the batch must span many lanes"); + + let canonical: Vec<(String, i16, i16)> = lanes + .into_iter() + .map(|(_, priority, shard, queue)| (queue, priority, shard)) + .collect(); + let mut canonical_keys = Vec::new(); + for (queue, priority, shard) in &canonical { + canonical_keys.push(lane_reservation_lock_key(&pool, queue, *priority, *shard).await); + } + + // Provision the lanes so the gate lock is the only thing the producer waits for. + for (index, queue) in queues.iter().enumerate() { + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + for (priority, shard) in [ + (1, 0), + (1, 1), + (1, 2), + (1, 3), + (2, 0), + (2, 1), + (2, 2), + (2, 3), + ] { + let mut params = job(queue, 10_000 + index, false); + params.opts.priority = priority; + params.opts.ordering_key = Some(keys[shard as usize].clone()); + store + .enqueue_params_batch(&pool, &[params]) + .await + .expect("provision lane"); + } + } + + let gated = canonical.len() / 2; + let mut gate = pool.begin().await.expect("begin gate transaction"); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(canonical_keys[gated]) + .execute(&mut *gate) + .await + .expect("take gate lock"); + + let competing = { + let pool = pool.clone(); + let jobs = jobs.clone(); + tokio::spawn(async move { producer.run(&pool, &jobs).await }) + }; + let waiter = advisory_waiter_pid(&pool, canonical_keys[gated]).await; + let held = granted_advisory_keys(&pool, waiter).await; + let expected: std::collections::BTreeSet = + canonical_keys[..gated].iter().copied().collect(); + assert_eq!( + held, + expected, + "a batch parked on lane {gated} of {} must hold the lanes before it and no later one", + canonical.len() + ); + + gate.commit().await.expect("release gate"); + let inserted = tokio::time::timeout(WAIT_DEADLINE, competing) + .await + .expect("producer finished") + .expect("producer task") + .expect("insert batch"); - let inserted = producer.run(&pool, &jobs).await.expect("insert batch"); + // Lane ordering must not leak into the result: rows come back in input order. assert_eq!(inserted.len(), jobs.len()); for (row, params) in inserted.iter().zip(&jobs) { assert_eq!(row.queue, params.opts.queue); assert_eq!(row.args, params.args); } - // Job ids are allocated as each row reaches `insert_job_compat`, so id - // order is lock order. - let lane_by_job: std::collections::HashMap = - sqlx::query_as::<_, (i64, String, i16, i16)>( - "SELECT job_id, queue, priority, enqueue_shard FROM awa.ready_entries WHERE job_id = ANY($1)", + // Within a lane, lane_seq follows input order. + let lane_by_job: std::collections::HashMap = + sqlx::query_as::<_, (i64, String, i16, i16, i64)>( + "SELECT job_id, queue, priority, enqueue_shard, lane_seq FROM awa.ready_entries WHERE job_id = ANY($1)", ) .bind(inserted.iter().map(|row| row.id).collect::>()) .fetch_all(&pool) .await .expect("read lanes") .into_iter() - .map(|(job_id, queue, priority, shard)| (job_id, (queue, priority, shard))) + .map(|(job_id, queue, priority, shard, lane_seq)| { + (job_id, (queue, priority, shard, lane_seq)) + }) .collect(); assert_eq!(lane_by_job.len(), jobs.len()); - - let mut by_lock_order: Vec<(i64, usize)> = inserted - .iter() - .enumerate() - .map(|(position, row)| (row.id, position)) - .collect(); - by_lock_order.sort_unstable(); - let lock_order: Vec<(&(String, i16, i16), usize)> = by_lock_order - .iter() - .map(|(job_id, position)| (&lane_by_job[job_id], *position)) - .collect(); - for pair in lock_order.windows(2) { - let (lane, position) = pair[0]; - let (next_lane, next_position) = pair[1]; - assert!( - lane.0.as_bytes() < next_lane.0.as_bytes() - || (lane.0 == next_lane.0 && (lane.1, lane.2) < (next_lane.1, next_lane.2)) - || (lane == next_lane && position < next_position), - "lane {lane:?} (input {position}) was locked before lane {next_lane:?} (input {next_position})" - ); + let mut last_seq_by_lane: std::collections::HashMap<(String, i16, i16), i64> = + std::collections::HashMap::new(); + for row in &inserted { + let (queue, priority, shard, lane_seq) = lane_by_job[&row.id].clone(); + if let Some(previous) = last_seq_by_lane.insert((queue, priority, shard), lane_seq) { + assert_eq!(lane_seq, previous + 1, "lane_seq must follow input order"); + } } } @@ -372,6 +506,11 @@ async fn insert_many_reserves_sharded_lanes_in_canonical_order() { assert_batch_reserves_lanes_in_canonical_order(Producer::InsertMany).await; } +#[tokio::test] +async fn insert_many_skip_unique_conflicts_reserves_sharded_lanes_in_canonical_order() { + assert_batch_reserves_lanes_in_canonical_order(Producer::SkipUniqueConflicts).await; +} + #[tokio::test] async fn insert_many_copy_reserves_sharded_lanes_in_canonical_order() { assert_batch_reserves_lanes_in_canonical_order(Producer::Copy).await; diff --git a/awa/tests/insert_many_batch_test.rs b/awa/tests/insert_many_batch_test.rs new file mode 100644 index 000000000..5609bd19e --- /dev/null +++ b/awa/tests/insert_many_batch_test.rs @@ -0,0 +1,585 @@ +//! Compatibility inserts and the batch insert function from the +//! `insert_jobs_compat_batch` schema patch. +//! +//! A batch does its lane-independent work (job ids, unique claims, lane +//! registration) before it reserves its lanes, and holds a lane lock only for +//! the reservation, the ready-row writes and the commit. These tests cover that +//! ordering, skip-on-conflict inserts, and the shape of what lands in the +//! lanes. +//! +//! Requires a running Postgres instance. +//! Run with: `cargo test --package awa --test insert_many_batch_test -- --nocapture` + +use awa::model::{ + insert_many, insert_many_copy_from_pool, insert_many_skip_unique_conflicts, AwaError, + QueueStorage, +}; +use awa::{InsertOpts, InsertParams, JobState, QueueStorageConfig, UniqueOpts}; +use awa_testing::setup; +use sqlx::PgPool; +use std::sync::LazyLock; +use std::time::{Duration, Instant}; +use tokio::sync::Mutex; +use uuid::Uuid; + +/// The tests switch the storage engine of a shared database. +static ENGINE_GUARD: LazyLock> = LazyLock::new(|| Mutex::new(())); + +const WAIT_DEADLINE: Duration = Duration::from_secs(20); + +async fn queue_storage_pool() -> PgPool { + let pool = setup::setup(8).await; + setup::activate_queue_storage(&pool).await; + pool +} + +fn unique_queue(label: &str) -> String { + format!( + "batch_{label}_{}", + &Uuid::new_v4().simple().to_string()[..12] + ) +} + +fn job(queue: &str, seq: i64, unique: bool) -> InsertParams { + InsertParams { + kind: "insert_many_batch_job".to_string(), + args: serde_json::json!({ "seq": seq }), + opts: InsertOpts { + queue: queue.to_string(), + unique: unique.then(|| UniqueOpts { + by_queue: true, + by_args: true, + ..Default::default() + }), + ..Default::default() + }, + } +} + +async fn ready_lane_seqs(pool: &PgPool, queue: &str) -> Vec<(i64, i64)> { + sqlx::query_as( + "SELECT (args->>'seq')::bigint, lane_seq FROM awa.ready_entries + WHERE queue = $1 ORDER BY lane_seq", + ) + .bind(queue) + .fetch_all(pool) + .await + .expect("read ready rows") +} + +async fn claims_for_queue(pool: &PgPool, queue: &str) -> i64 { + sqlx::query_scalar( + "SELECT count(*)::bigint FROM awa.job_unique_claims AS claims + WHERE claims.job_id IN (SELECT job_id FROM awa.ready_entries WHERE queue = $1)", + ) + .bind(queue) + .fetch_one(pool) + .await + .expect("count claims") +} + +fn is_unique_violation(err: &AwaError) -> bool { + match err { + AwaError::UniqueConflict { .. } => true, + AwaError::Database(sqlx::Error::Database(db_err)) => { + db_err.code().as_deref() == Some("23505") + } + _ => false, + } +} + +async fn lane_lock_key(pool: &PgPool, queue: &str) -> i64 { + sqlx::query_scalar( + "SELECT hashtextextended( + format('%I.%I', 'awa', awa.queue_lane_sequence_name('queue_enqueue_seq', $1, 2::smallint, 0::smallint)), + 0)", + ) + .bind(queue) + .fetch_one(pool) + .await + .expect("derive lane lock key") +} + +/// The producer entry points: generic `insert_many` on the per-row function, +/// and the two that use the batch function. +#[derive(Debug, Clone, Copy)] +enum Path { + PerRow, + SkipUniqueConflicts, + Copy, +} + +impl Path { + const ALL: [Path; 3] = [Path::PerRow, Path::SkipUniqueConflicts, Path::Copy]; + + async fn insert(self, pool: &PgPool, jobs: &[InsertParams]) -> Vec { + match self { + Path::PerRow => insert_many(pool, jobs).await, + Path::SkipUniqueConflicts => insert_many_skip_unique_conflicts(pool, jobs).await, + Path::Copy => insert_many_copy_from_pool(pool, jobs).await, + } + .unwrap_or_else(|err| panic!("{self:?} insert: {err:?}")) + } + + /// Ready segments a batch of `rows` ready rows in one lane, sharing a + /// `run_at`, writes: one per row through the per-row function, one per + /// run through the batch function. + fn segments_for(self, rows: i64) -> i64 { + match self { + Path::PerRow => rows, + Path::SkipUniqueConflicts | Path::Copy => 1, + } + } +} + +async fn segment_count(pool: &PgPool, queue: &str) -> i64 { + sqlx::query_scalar("SELECT count(*)::bigint FROM awa.ready_segments WHERE queue = $1") + .bind(queue) + .fetch_one(pool) + .await + .expect("count segments") +} + +#[tokio::test] +async fn batch_lands_in_one_lane_in_input_order_and_is_claimable() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + for path in Path::ALL { + assert_batch_lands_in_one_lane_in_input_order(&pool, path).await; + } +} + +async fn assert_batch_lands_in_one_lane_in_input_order(pool: &PgPool, path: Path) { + let queue = unique_queue("order"); + + let jobs: Vec = (0..300).map(|seq| job(&queue, seq, seq % 2 == 0)).collect(); + let inserted = path.insert(pool, &jobs).await; + assert_eq!( + inserted + .iter() + .map(|row| row.args["seq"].as_i64().unwrap()) + .collect::>(), + (0..300).collect::>(), + "rows come back in input order" + ); + assert!(inserted.iter().all(|row| row.state == JobState::Available)); + + let lane = ready_lane_seqs(pool, &queue).await; + assert_eq!(lane.len(), 300); + let first = lane[0].1; + for (offset, (seq, lane_seq)) in lane.iter().enumerate() { + assert_eq!(*seq, offset as i64, "lane_seq follows input order"); + assert_eq!(*lane_seq, first + offset as i64, "lane_seq is dense"); + } + assert_eq!( + segment_count(pool, &queue).await, + path.segments_for(300), + "{path:?} wrote an unexpected number of ready segments" + ); + + // The lane's segment metadata must let the claim path find every row. + let store = QueueStorage::new(QueueStorageConfig::default()).expect("queue storage"); + let mut claimed = Vec::new(); + while claimed.len() < 300 { + let batch = store + .claim_runtime_batch(pool, &queue, 100, Duration::from_secs(60)) + .await + .expect("claim"); + assert!( + !batch.is_empty(), + "{path:?}: claim path lost rows of the batch" + ); + claimed.extend( + batch + .into_iter() + .map(|job| job.job.args["seq"].as_i64().unwrap()), + ); + } + assert_eq!( + claimed, + (0..300).collect::>(), + "{path:?}: claimed in lane order" + ); +} + +#[tokio::test] +async fn batch_with_mixed_states_routes_each_row() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + for path in Path::ALL { + assert_mixed_states_route_each_row(&pool, path).await; + } +} + +async fn assert_mixed_states_route_each_row(pool: &PgPool, path: Path) { + let queue = unique_queue("mixed"); + + let later = chrono::Utc::now() + chrono::Duration::hours(1); + let mut jobs: Vec = (0..6).map(|seq| job(&queue, seq, false)).collect(); + jobs[1].opts.run_at = Some(later); + jobs[4].opts.run_at = Some(later); + + let inserted = path.insert(pool, &jobs).await; + let states: Vec = inserted.iter().map(|row| row.state).collect(); + assert_eq!( + states, + vec![ + JobState::Available, + JobState::Scheduled, + JobState::Available, + JobState::Available, + JobState::Scheduled, + JobState::Available, + ], + "{path:?}" + ); + let ready: Vec = ready_lane_seqs(pool, &queue) + .await + .into_iter() + .map(|(seq, _)| seq) + .collect(); + assert_eq!(ready, vec![0, 2, 3, 5]); + let deferred: i64 = sqlx::query_scalar( + "SELECT count(*)::bigint FROM awa.deferred_jobs WHERE queue = $1 AND state = 'scheduled'", + ) + .bind(&queue) + .fetch_one(pool) + .await + .expect("count deferred"); + assert_eq!(deferred, 2, "{path:?}"); +} + +#[tokio::test] +async fn unique_conflict_fails_the_batch_atomically() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + let queue = unique_queue("atomic"); + + insert_many(&pool, &[job(&queue, 1, true)]) + .await + .expect("first insert"); + + let batch = vec![ + job(&queue, 0, true), + job(&queue, 1, true), + job(&queue, 2, true), + ]; + let err = insert_many(&pool, &batch) + .await + .expect_err("duplicate key must fail the batch"); + assert!(is_unique_violation(&err), "{err:?}"); + + let landed: Vec = ready_lane_seqs(&pool, &queue) + .await + .into_iter() + .map(|(seq, _)| seq) + .collect(); + assert_eq!(landed, vec![1], "nothing from the failed batch landed"); + assert_eq!(claims_for_queue(&pool, &queue).await, 1); +} + +#[tokio::test] +async fn skip_unique_conflicts_inserts_the_rest_in_order() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + let queue = unique_queue("skip"); + + insert_many(&pool, &[job(&queue, 1, true), job(&queue, 3, true)]) + .await + .expect("seed"); + + // Keys 1 and 3 are already claimed, 4 repeats inside the batch, and the + // last row carries no unique key at all. + let batch = vec![ + job(&queue, 0, true), + job(&queue, 1, true), + job(&queue, 2, true), + job(&queue, 3, true), + job(&queue, 4, true), + job(&queue, 4, true), + job(&queue, 5, false), + ]; + let inserted = insert_many_skip_unique_conflicts(&pool, &batch) + .await + .expect("skip insert"); + assert_eq!( + inserted + .iter() + .map(|row| row.args["seq"].as_i64().unwrap()) + .collect::>(), + vec![0, 2, 4, 5] + ); + + let lane = ready_lane_seqs(&pool, &queue).await; + assert_eq!( + lane.iter().map(|(seq, _)| *seq).collect::>(), + vec![1, 3, 0, 2, 4, 5], + "skipped rows take no lane position" + ); + for pair in lane.windows(2) { + assert_eq!(pair[1].1, pair[0].1 + 1, "lane_seq stays dense"); + } + + // Everything skipped: an empty result, not an error. + let inserted = insert_many_skip_unique_conflicts(&pool, &[job(&queue, 0, true)]) + .await + .expect("all skipped"); + assert!(inserted.is_empty()); +} + +#[tokio::test] +async fn copy_skips_unique_conflicts_in_one_pass() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + let queue = unique_queue("copy_skip"); + + insert_many(&pool, &[job(&queue, 2, true)]) + .await + .expect("seed"); + let batch: Vec = (0..6).map(|seq| job(&queue, seq, true)).collect(); + let inserted = insert_many_copy_from_pool(&pool, &batch) + .await + .expect("copy insert"); + assert_eq!( + inserted + .iter() + .map(|row| row.args["seq"].as_i64().unwrap()) + .collect::>(), + vec![0, 1, 3, 4, 5] + ); + assert_eq!(ready_lane_seqs(&pool, &queue).await.len(), 6); +} + +#[tokio::test] +async fn canonical_storage_inserts_and_skips_row_by_row() { + let _guard = ENGINE_GUARD.lock().await; + let pool = setup::setup(8).await; + setup::reset_to_canonical(&pool).await; + let queue = unique_queue("canonical"); + + insert_many(&pool, &[job(&queue, 1, true)]) + .await + .expect("seed"); + let batch: Vec = (0..4).map(|seq| job(&queue, seq, true)).collect(); + let err = insert_many(&pool, &batch).await.expect_err("duplicate"); + assert!(is_unique_violation(&err), "{err:?}"); + + let inserted = insert_many_skip_unique_conflicts(&pool, &batch) + .await + .expect("skip insert"); + assert_eq!( + inserted + .iter() + .map(|row| row.args["seq"].as_i64().unwrap()) + .collect::>(), + vec![0, 2, 3] + ); +} + +/// While another transaction holds a lane lock, a batch bound for that lane +/// has already done its lane-independent work (it holds its unique claims) and +/// has not yet written a ready row. Every row's unique claim, not just the +/// first's, is held by then. +#[tokio::test] +async fn batch_waits_for_the_lane_after_its_lane_independent_work() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + let queue = unique_queue("window"); + + // Provision the lane so the gate lock is the only thing to wait for. + insert_many(&pool, &[job(&queue, -1, false)]) + .await + .expect("provision lane"); + + let key = lane_lock_key(&pool, &queue).await; + let mut gate = pool.begin().await.expect("begin gate"); + sqlx::query("SELECT pg_advisory_xact_lock($1)") + .bind(key) + .execute(&mut *gate) + .await + .expect("take lane lock"); + + let jobs_for_probe: Vec = (0..50).map(|seq| job(&queue, seq, true)).collect(); + let producer = { + let pool = pool.clone(); + let jobs = jobs_for_probe.clone(); + tokio::spawn(async move { insert_many_skip_unique_conflicts(&pool, &jobs).await }) + }; + + let started = Instant::now(); + let waiter = loop { + let pid: Option = sqlx::query_scalar( + "SELECT pid FROM pg_locks + WHERE locktype = 'advisory' AND NOT granted + AND ((classid::bigint << 32) | objid::bigint) = $1 + LIMIT 1", + ) + .bind(key) + .fetch_optional(&pool) + .await + .expect("poll pg_locks"); + if let Some(pid) = pid { + break pid; + } + assert!(started.elapsed() < WAIT_DEADLINE, "producer never waited"); + tokio::time::sleep(Duration::from_millis(10)).await; + }; + + let holds_relation_lock = |relation: &'static str| { + let pool = pool.clone(); + async move { + sqlx::query_scalar::<_, bool>( + "SELECT EXISTS ( + SELECT 1 FROM pg_locks + WHERE pid = $1 AND locktype = 'relation' AND granted + AND relation = $2::regclass + )", + ) + .bind(waiter) + .bind(relation) + .fetch_one(&pool) + .await + .expect("read relation locks") + } + }; + // The claim of the batch's last row is already held, uncommitted: an + // insert of the same key waits on the producer's transaction. + let last = jobs_for_probe.last().expect("probe job"); + let last_key = + awa::model::prepare_raw_job_insert(last.kind.clone(), last.args.clone(), last.opts.clone()) + .expect("prepare last job") + .unique_key() + .expect("unique key") + .to_vec(); + let mut probe = pool.begin().await.expect("begin probe"); + sqlx::query("SET LOCAL lock_timeout = '300ms'") + .execute(&mut *probe) + .await + .expect("set lock timeout"); + let probe_result = sqlx::query( + "INSERT INTO awa.job_unique_claims (unique_key, job_id) VALUES ($1, 0) + ON CONFLICT (unique_key) DO NOTHING", + ) + .bind(&last_key) + .execute(&mut *probe) + .await; + match probe_result { + Err(sqlx::Error::Database(db_err)) => assert_eq!( + db_err.code().as_deref(), + Some("55P03"), + "the last row's claim must already be held while the lane lock is awaited" + ), + other => panic!("claim of the last row was not held while waiting for the lane: {other:?}"), + } + drop(probe); + assert!( + !holds_relation_lock("awa.ready_entries").await, + "no ready row is written while the lane lock is awaited" + ); + + gate.commit().await.expect("release lane lock"); + let inserted = tokio::time::timeout(WAIT_DEADLINE, producer) + .await + .expect("producer finished") + .expect("producer task") + .expect("insert batch"); + assert_eq!(inserted.len(), 50); + assert_eq!(ready_lane_seqs(&pool, &queue).await.len(), 51); +} + +/// Producers on the per-row function, as released 0.6 binaries are, and +/// producers on the batch function share one lane at the same time. Each batch +/// still occupies a contiguous run of `lane_seq`, in input order, and the lane +/// has no duplicate or missing positions. +#[tokio::test] +async fn old_and_new_producers_share_a_lane_without_reordering() { + const TASKS: usize = 8; + const BATCHES: usize = 25; + const BATCH: i64 = 40; + + let _guard = ENGINE_GUARD.lock().await; + let pool = setup::setup(TASKS as u32 + 4).await; + setup::activate_queue_storage(&pool).await; + let queue = unique_queue("mixed_fleet"); + + // Provision the lane before mixing generic and batch producers. + insert_many(&pool, &[job(&queue, -1, false)]) + .await + .expect("warm up"); + + let mut tasks = Vec::new(); + for task in 0..TASKS { + let pool = pool.clone(); + let queue = queue.clone(); + tasks.push(tokio::spawn(async move { + for batch in 0..BATCHES { + let jobs: Vec = (0..BATCH) + .map(|seq| InsertParams { + kind: "insert_many_batch_job".to_string(), + args: serde_json::json!({ "seq": seq, "batch": format!("{task}-{batch}") }), + opts: InsertOpts { + queue: queue.clone(), + ..Default::default() + }, + }) + .collect(); + if task % 2 == 0 { + insert_through_row_function(&pool, &jobs).await; + } else { + insert_many_skip_unique_conflicts(&pool, &jobs) + .await + .expect("batch function insert"); + } + } + })); + } + for task in tasks { + task.await.expect("producer task"); + } + + let rows: Vec<(String, i64, i64)> = sqlx::query_as( + "SELECT args->>'batch', (args->>'seq')::bigint, lane_seq FROM awa.ready_entries + WHERE queue = $1 AND args ? 'batch' ORDER BY lane_seq", + ) + .bind(&queue) + .fetch_all(&pool) + .await + .expect("read lane"); + assert_eq!(rows.len(), TASKS * BATCHES * BATCH as usize); + for pair in rows.windows(2) { + assert!(pair[1].2 > pair[0].2, "lane_seq is unique and ascending"); + } + let mut by_batch: std::collections::HashMap<&str, Vec<(i64, i64)>> = + std::collections::HashMap::new(); + for (batch, seq, lane_seq) in &rows { + by_batch.entry(batch).or_default().push((*seq, *lane_seq)); + } + for (batch, members) in by_batch { + for (offset, (seq, lane_seq)) in members.iter().enumerate() { + assert_eq!(*seq, offset as i64, "batch {batch} keeps input order"); + assert_eq!( + *lane_seq, + members[0].1 + offset as i64, + "batch {batch} occupies a contiguous run of the lane" + ); + } + } +} + +/// What a released 0.6 binary sends: one `insert_job_compat` call per row. +async fn insert_through_row_function(pool: &PgPool, jobs: &[InsertParams]) { + sqlx::query( + r#" + SELECT inserted.* + FROM unnest($1::text[], $2::text[], $3::jsonb[]) WITH ORDINALITY AS input(kind, queue, args, ord) + CROSS JOIN LATERAL awa.insert_job_compat(input.kind, input.queue, input.args) AS inserted + ORDER BY input.ord + "#, + ) + .bind(jobs.iter().map(|job| job.kind.as_str()).collect::>()) + .bind(jobs.iter().map(|job| job.opts.queue.as_str()).collect::>()) + .bind(jobs.iter().map(|job| &job.args).collect::>()) + .fetch_all(pool) + .await + .expect("per-row function insert"); +} diff --git a/awa/tests/insert_many_schema_fallback_test.rs b/awa/tests/insert_many_schema_fallback_test.rs new file mode 100644 index 000000000..a6b31ae29 --- /dev/null +++ b/awa/tests/insert_many_schema_fallback_test.rs @@ -0,0 +1,434 @@ +//! `insert_many` and `insert_many_copy` against a schema without the +//! `insert_jobs_compat_batch` schema patch, and the patch itself. +//! +//! A binary may run ahead of `awa migrate`. Generic `insert_many` keeps using +//! the per-row function; COPY probes its own connection and uses the batch +//! function when it is installed and callable. Different databases in one +//! process can differ in whether they have the patch without aborting their +//! caller's transactions. +//! +//! Requires a running Postgres instance and a role that may create databases +//! and roles. +//! Run with: `cargo test --package awa --test insert_many_schema_fallback_test` + +use awa::model::{ + insert_many, insert_many_copy, insert_many_copy_from_pool, insert_many_skip_unique_conflicts, + migrations, AwaError, +}; +use awa::{InsertOpts, InsertParams, UniqueOpts}; +use awa_testing::setup; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; +use sqlx::PgPool; +use std::str::FromStr; +use uuid::Uuid; + +const PATCH_NAME: &str = "insert_jobs_compat_batch"; +const DROP_BATCH_FUNCTION: &str = "DROP FUNCTION awa.insert_jobs_compat_batch(text[], text[], \ + jsonb[], text[], smallint[], smallint[], timestamptz[], jsonb[], jsonb[], bytea[], text[], \ + text[], boolean)"; + +/// A migrated, queue-storage database of its own, dropped by [`drop`](Self::drop). +struct ScratchDatabase { + name: String, + pool: PgPool, +} + +impl ScratchDatabase { + async fn create() -> Self { + let name = format!("awa_batchfn_{}", &Uuid::new_v4().simple().to_string()[..12]); + let admin = admin_pool().await; + sqlx::query(&format!("CREATE DATABASE {name}")) + .execute(&admin) + .await + .expect("create scratch database"); + admin.close().await; + + let options = PgConnectOptions::from_str(&setup::database_url()) + .expect("DATABASE_URL parses") + .database(&name); + let pool = PgPoolOptions::new() + .max_connections(4) + .connect_with(options) + .await + .expect("connect scratch database"); + migrations::run(&pool).await.expect("migrate"); + setup::activate_queue_storage(&pool).await; + Self { name, pool } + } + + async fn drop(self) { + self.pool.close().await; + let admin = admin_pool().await; + sqlx::query(&format!("DROP DATABASE {} WITH (FORCE)", self.name)) + .execute(&admin) + .await + .expect("drop scratch database"); + admin.close().await; + } +} + +async fn admin_pool() -> PgPool { + let options = PgConnectOptions::from_str(&setup::database_url()) + .expect("DATABASE_URL parses") + .database("postgres"); + PgPoolOptions::new() + .max_connections(1) + .connect_with(options) + .await + .expect("connect to the postgres maintenance database") +} + +fn jobs(queue: &str, range: std::ops::Range, unique: bool) -> Vec { + range + .map(|seq| InsertParams { + kind: "schema_fallback_job".to_string(), + args: serde_json::json!({ "seq": seq }), + opts: InsertOpts { + queue: queue.to_string(), + unique: unique.then(|| UniqueOpts { + by_queue: true, + by_args: true, + ..Default::default() + }), + ..Default::default() + }, + }) + .collect() +} + +async fn segments(pool: &PgPool, queue: &str) -> i64 { + sqlx::query_scalar("SELECT count(*)::bigint FROM awa.ready_segments WHERE queue = $1") + .bind(queue) + .fetch_one(pool) + .await + .expect("count segments") +} + +async fn ready_seqs(pool: &PgPool, queue: &str) -> Vec { + sqlx::query_scalar( + "SELECT (args->>'seq')::bigint FROM awa.ready_entries WHERE queue = $1 ORDER BY lane_seq", + ) + .bind(queue) + .fetch_all(pool) + .await + .expect("read ready rows") +} + +fn seqs(rows: &[awa::JobRow]) -> Vec { + rows.iter() + .map(|row| row.args["seq"].as_i64().expect("seq")) + .collect() +} + +async fn pending_patches(pool: &PgPool) -> Vec<&'static str> { + migrations::pending_schema_patches(pool) + .await + .expect("probe schema patches") + .into_iter() + .map(|patch| patch.name) + .collect() +} + +#[tokio::test] +async fn binary_ahead_of_the_patch_falls_back_and_recovers() { + let db = ScratchDatabase::create().await; + let pool = &db.pool; + assert!(pending_patches(pool).await.is_empty()); + + // A v040 schema from before the patch. + sqlx::query(DROP_BATCH_FUNCTION) + .execute(pool) + .await + .expect("drop batch function"); + assert_eq!(pending_patches(pool).await, vec![PATCH_NAME]); + + let first = insert_many(pool, &jobs("fallback_a", 0..3, false)) + .await + .expect("insert_many without the batch function"); + assert_eq!(seqs(&first), vec![0, 1, 2]); + assert_eq!(segments(pool, "fallback_a").await, 3, "per-row path"); + let copied = insert_many_copy_from_pool(pool, &jobs("fallback_a", 3..6, false)) + .await + .expect("insert_many_copy without the batch function"); + assert_eq!(seqs(&copied), vec![3, 4, 5]); + assert_eq!(segments(pool, "fallback_a").await, 6, "COPY stays per-row"); + assert_eq!(ready_seqs(pool, "fallback_a").await, vec![0, 1, 2, 3, 4, 5]); + + // Unique handling on the per-row path: a duplicate fails `insert_many` + // and is skipped by `insert_many_copy`. + insert_many(pool, &jobs("fallback_u", 0..2, true)) + .await + .expect("seed unique jobs"); + let duplicate = insert_many(pool, &jobs("fallback_u", 1..3, true)) + .await + .expect_err("a unique conflict fails the batch"); + assert!( + matches!(&duplicate, AwaError::Database(sqlx::Error::Database(e)) if e.code().as_deref() == Some("23505")), + "{duplicate:?}" + ); + let copied = insert_many_copy_from_pool(pool, &jobs("fallback_u", 1..4, true)) + .await + .expect("unique insert_many_copy without the batch function"); + assert_eq!(seqs(&copied), vec![2, 3], "duplicate skipped"); + + // The skip API needs the patch and says so. + let skipped = insert_many_skip_unique_conflicts(pool, &jobs("fallback_u", 0..1, true)) + .await + .expect_err("needs the patch"); + assert!( + matches!(&skipped, AwaError::Validation(message) if message.contains(PATCH_NAME)), + "{skipped:?}" + ); + + // `awa migrate` applies the patch without a schema version. + let report = migrations::run_with_report(pool).await.expect("migrate"); + let patch = report + .patches + .iter() + .find(|patch| patch.name == PATCH_NAME) + .expect("patch report"); + assert_eq!(patch.outcome, migrations::PatchOutcome::Applied); + assert!(report.complete()); + assert!(report.migrations_applied.is_empty()); + assert_eq!(report.version_after, migrations::CURRENT_VERSION); + assert!(pending_patches(pool).await.is_empty()); + + // Generic insert_many stays on the per-row path; COPY picks the batch + // function on its next call, on the same pool. + insert_many(pool, &jobs("fallback_b", 0..3, false)) + .await + .expect("generic insert after the patch"); + assert_eq!(segments(pool, "fallback_b").await, 3, "per-row path"); + let copied = insert_many_copy_from_pool(pool, &jobs("fallback_b", 3..6, false)) + .await + .expect("COPY after the patch"); + assert_eq!(seqs(&copied), vec![3, 4, 5]); + assert_eq!( + segments(pool, "fallback_b").await, + 4, + "COPY writes its batch as one segment" + ); + assert_eq!(ready_seqs(pool, "fallback_b").await, vec![0, 1, 2, 3, 4, 5]); + + let copied = insert_many_copy_from_pool(pool, &jobs("fallback_u", 2..6, true)) + .await + .expect("insert_many_copy with the batch function"); + assert_eq!(seqs(&copied), vec![4, 5], "duplicates skipped in one pass"); + let skipped = insert_many_skip_unique_conflicts(pool, &jobs("fallback_u", 0..8, true)) + .await + .expect("skip with the batch function"); + assert_eq!(seqs(&skipped), vec![6, 7]); + + // Removing the function after a successful batch does not poison the next + // generic or COPY call with a stale answer. + sqlx::query(DROP_BATCH_FUNCTION) + .execute(pool) + .await + .expect("drop batch function again"); + let mut tx = pool.begin().await.expect("caller transaction"); + insert_many(&mut *tx, &jobs("fallback_c", 0..3, false)) + .await + .expect("generic insert still works"); + insert_many_copy(&mut tx, &jobs("fallback_c", 3..6, false)) + .await + .expect("COPY re-probes its own connection"); + tx.commit() + .await + .expect("caller transaction remains usable"); + assert_eq!(segments(pool, "fallback_c").await, 6, "per-row paths"); + + db.drop().await; +} + +/// What one database says about the patch never selects SQL for another, and +/// both paths join the caller's transaction, including its rollback. +#[tokio::test] +async fn mixed_patch_databases_do_not_abort_caller_transactions() { + let patched = ScratchDatabase::create().await; + let unpatched = ScratchDatabase::create().await; + for db in [&patched, &unpatched] { + sqlx::query("CREATE TABLE application_writes (id bigint PRIMARY KEY)") + .execute(&db.pool) + .await + .expect("create application table"); + } + sqlx::query(DROP_BATCH_FUNCTION) + .execute(&unpatched.pool) + .await + .expect("unpatched database has no batch function"); + + insert_many_skip_unique_conflicts(&patched.pool, &jobs("mixed_warm", 0..2, true)) + .await + .expect("batch API on the patched database"); + + for (index, db) in [&patched, &unpatched, &patched].into_iter().enumerate() { + let pool = &db.pool; + let queue = format!("mixed_database_{index}"); + let mut tx = pool.begin().await.expect("caller transaction"); + sqlx::query("INSERT INTO application_writes (id) VALUES ($1)") + .bind(index as i64) + .execute(&mut *tx) + .await + .expect("application write"); + let inserted = insert_many(&mut *tx, &jobs(&queue, 0..3, true)) + .await + .expect("generic insert must use SQL every schema supports"); + assert_eq!(seqs(&inserted), vec![0, 1, 2]); + let copied = insert_many_copy(&mut tx, &jobs(&queue, 2..5, true)) + .await + .expect("COPY must probe this database"); + assert_eq!(seqs(&copied), vec![3, 4], "duplicate skipped"); + let written: i64 = + sqlx::query_scalar("SELECT count(*) FROM application_writes WHERE id = $1") + .bind(index as i64) + .fetch_one(&mut *tx) + .await + .expect("caller transaction was not aborted"); + assert_eq!(written, 1); + tx.commit() + .await + .expect("commit application and jobs together"); + assert_eq!(ready_seqs(pool, &queue).await, vec![0, 1, 2, 3, 4]); + + let mut tx = pool.begin().await.expect("rollback transaction"); + sqlx::query("INSERT INTO application_writes (id) VALUES ($1)") + .bind(index as i64 + 100) + .execute(&mut *tx) + .await + .expect("application write to roll back"); + insert_many(&mut *tx, &jobs(&queue, 5..7, true)) + .await + .expect("generic insert to roll back"); + insert_many_copy(&mut tx, &jobs(&queue, 7..9, true)) + .await + .expect("COPY to roll back"); + tx.rollback().await.expect("roll back caller transaction"); + assert_eq!(ready_seqs(pool, &queue).await, vec![0, 1, 2, 3, 4]); + let rolled_back: i64 = + sqlx::query_scalar("SELECT count(*) FROM application_writes WHERE id = $1") + .bind(index as i64 + 100) + .fetch_one(pool) + .await + .expect("check application rollback"); + assert_eq!(rolled_back, 0); + } + + patched.drop().await; + unpatched.drop().await; +} + +/// The exported patch, as an external runner applies it: twice in one +/// explicit transaction, over an installed function, writing no schema +/// version. +#[tokio::test] +async fn exported_patch_reapplies_in_one_transaction() { + let db = ScratchDatabase::create().await; + let pool = &db.pool; + let (_, _, sql) = migrations::schema_patch_sql(migrations::SCHEMA_PATCHES) + .into_iter() + .find(|(name, _, _)| *name == PATCH_NAME) + .expect("exported patch"); + let versions_before: Vec = + sqlx::query_scalar("SELECT version FROM awa.schema_version ORDER BY version") + .fetch_all(pool) + .await + .expect("read schema versions"); + + sqlx::query(DROP_BATCH_FUNCTION) + .execute(pool) + .await + .expect("drop batch function"); + sqlx::raw_sql(&format!("BEGIN;\n{sql}\n{sql}\nCOMMIT;")) + .execute(pool) + .await + .expect("apply the patch twice in one transaction"); + sqlx::raw_sql(&sql) + .execute(pool) + .await + .expect("apply over an installed function"); + + assert!(pending_patches(pool).await.is_empty()); + let versions_after: Vec = + sqlx::query_scalar("SELECT version FROM awa.schema_version ORDER BY version") + .fetch_all(pool) + .await + .expect("read schema versions"); + assert_eq!(versions_after, versions_before); + let copied = insert_many_copy_from_pool(pool, &jobs("exported", 0..4, false)) + .await + .expect("COPY through the exported patch"); + assert_eq!(seqs(&copied), vec![0, 1, 2, 3]); + assert_eq!(segments(pool, "exported").await, 1); + + db.drop().await; +} + +/// A role that may not execute the batch function keeps COPY on the per-row +/// path instead of failing it with `permission denied`. +#[tokio::test] +async fn copy_without_execute_on_the_batch_function_stays_per_row() { + let db = ScratchDatabase::create().await; + let pool = &db.pool; + // Provision the lane first, so the role needs no sequence DDL and the + // grants below cover its sequences. + insert_many(pool, &jobs("no_execute", 0..1, false)) + .await + .expect("provision lane"); + let role = format!( + "awa_batchfn_role_{}", + &Uuid::new_v4().simple().to_string()[..8] + ); + sqlx::raw_sql(&format!( + "CREATE ROLE {role} NOLOGIN; + GRANT TEMP ON DATABASE {database} TO {role}; + GRANT USAGE ON SCHEMA awa TO {role}; + GRANT SELECT, INSERT, UPDATE, DELETE 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.insert_jobs_compat_batch(text[], text[], jsonb[], text[], + smallint[], smallint[], timestamptz[], jsonb[], jsonb[], bytea[], text[], text[], boolean) + FROM PUBLIC, {role};", + database = db.name, + )) + .execute(pool) + .await + .expect("create a role without EXECUTE on the batch function"); + + let mut tx = pool.begin().await.expect("begin"); + sqlx::query(&format!("SET LOCAL ROLE {role}")) + .execute(&mut *tx) + .await + .expect("become the restricted role"); + let copied = insert_many_copy(&mut tx, &jobs("no_execute", 1..4, true)) + .await + .expect("COPY without EXECUTE on the batch function"); + assert_eq!(seqs(&copied), vec![1, 2, 3]); + let skipped = insert_many_skip_unique_conflicts(&mut *tx, &jobs("no_execute", 4..5, true)) + .await + .expect_err("the skip API calls the function directly"); + assert!( + matches!(&skipped, AwaError::Database(sqlx::Error::Database(e)) if e.code().as_deref() == Some("42501")), + "{skipped:?}" + ); + tx.rollback().await.expect("roll back"); + + let mut tx = pool.begin().await.expect("begin"); + sqlx::query(&format!("SET LOCAL ROLE {role}")) + .execute(&mut *tx) + .await + .expect("become the restricted role"); + insert_many_copy(&mut tx, &jobs("no_execute", 1..4, true)) + .await + .expect("COPY without EXECUTE on the batch function"); + tx.commit().await.expect("commit"); + assert_eq!(ready_seqs(pool, "no_execute").await, vec![0, 1, 2, 3]); + assert_eq!(segments(pool, "no_execute").await, 4, "per-row path"); + + db.drop().await; + let admin = admin_pool().await; + sqlx::query(&format!("DROP ROLE {role}")) + .execute(&admin) + .await + .expect("drop role"); + admin.close().await; +} From 1d907b2f6aa4e0d651b3387712fc3823389004a1 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sun, 8 Nov 2026 22:00:36 +1300 Subject: [PATCH 2/6] test(bench): benchmark per-lane enqueue contention (0.6) enqueue_lane_contention_test measures lane throughput and batch latency for many producers on one lane, across insert paths, unique keys and enqueue_shards. Its nightly gate runs sixteen producers on one lane and on four shards, checks lane integrity and compares the batch function's lock hold time with the per-row function's. AWA_LANE_BENCH_WITHOUT_PATCH drops the batch function so the matrix can measure the per-row COPY path. --- .github/workflows/nightly-chaos.yml | 11 + awa/tests/enqueue_lane_contention_test.rs | 816 ++++++++++++++++++++++ benchmarks/baseline.json | 7 + 3 files changed, 834 insertions(+) create mode 100644 awa/tests/enqueue_lane_contention_test.rs diff --git a/.github/workflows/nightly-chaos.yml b/.github/workflows/nightly-chaos.yml index ffef0dbd5..9a7a24df8 100644 --- a/.github/workflows/nightly-chaos.yml +++ b/.github/workflows/nightly-chaos.yml @@ -189,6 +189,17 @@ jobs: cargo test --release --package awa --test concurrent_lifecycle_test \ test_concurrent_queue_sweep -- --exact --ignored --nocapture \ | tee artifacts/nightly/rust-queue-sweep.txt + - name: Run per-lane enqueue contention benchmark + if: always() + timeout-minutes: 5 + run: | + set -o pipefail + # Many producers batch-inserting into one queue lane. Gates the + # per-lane rate and p99 through benchmarks/baseline.json, and the + # lane-lock hold time relative to the per-row insert function. + cargo test --release --package awa --test enqueue_lane_contention_test \ + test_enqueue_lane_contention_gate -- --exact --ignored --nocapture \ + | tee artifacts/nightly/rust-enqueue-lane.txt - name: Run MVCC horizon overlap benchmark if: always() timeout-minutes: 10 diff --git a/awa/tests/enqueue_lane_contention_test.rs b/awa/tests/enqueue_lane_contention_test.rs new file mode 100644 index 000000000..e411364c8 --- /dev/null +++ b/awa/tests/enqueue_lane_contention_test.rs @@ -0,0 +1,816 @@ +//! Per-lane enqueue contention: many concurrent producers batch-inserting +//! into one queue lane. +//! +//! Every enqueue on queue storage reserves its lane sequence under a +//! transaction-scoped advisory lock (`reserve_enqueue_seq`) that is held until +//! the enqueue transaction ends, so producers for one `(queue, priority, +//! shard)` lane serialise on it. This suite measures what that serialisation +//! costs: lane throughput and batch-insert latency at high producer +//! concurrency, with and without unique keys, with and without a unique +//! conflict fallback, and across `enqueue_shards`. +//! +//! * `test_enqueue_lane_contention_gate` runs a fixed set of scenarios, checks +//! lane integrity, asserts contention-scaled floors, and emits +//! `@@BENCH_JSON@@` records that `.github/scripts/check-bench-regression.py` +//! compares against `benchmarks/baseline.json`. +//! * `test_enqueue_lane_contention_matrix` is the exploration entry point: it +//! sweeps producer counts for one configuration chosen through environment +//! variables and prints every scenario. +//! +//! Matrix tunables (all optional): +//! AWA_LANE_BENCH_PRODUCERS comma list, default "1,4,8,16,32" +//! AWA_LANE_BENCH_BATCH jobs per batch, default 185 +//! AWA_LANE_BENCH_UNIQUE 0|1, default 1 +//! AWA_LANE_BENCH_SHARDS `enqueue_shards` for the queue, default 1 +//! AWA_LANE_BENCH_MODE insert_many | insert_many_copy | insert_many_skip | +//! queue_storage | legacy_row_function +//! AWA_LANE_BENCH_DUP_PCT % of batches carrying one duplicate key that +//! fall back to row-at-a-time inserts, default 0 +//! AWA_LANE_BENCH_SECS measured seconds per scenario, default 8 +//! AWA_LANE_BENCH_DRAIN 0|1: claim everything that landed afterwards and +//! check every shard drains +//! AWA_LANE_BENCH_SPLIT 0|1: run each batch in an explicit +//! transaction and time statement and commit +//! separately (insert_many mode only) +//! AWA_LANE_BENCH_WITHOUT_PATCH 0|1: drop `awa.insert_jobs_compat_batch` +//! first, so insert_many_copy measures its +//! per-row path (the next migration run puts +//! the function back) +//! +//! Requires a running Postgres instance. +//! Run with: +//! cargo test --release --package awa --test enqueue_lane_contention_test \ +//! -- --ignored --nocapture --test-threads=1 + +mod bench_output; +mod ci_timing; + +use awa::model::{ + insert_many, insert_many_copy_from_pool, insert_many_skip_unique_conflicts, + prepare_raw_job_insert, AwaError, QueueStorage, +}; +use awa::{InsertOpts, InsertParams, QueueStorageConfig, UniqueOpts}; +use awa_testing::setup; +use bench_output::{BenchLatency, BenchMetrics, BenchmarkResult, SCHEMA_VERSION}; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; +use sqlx::PgPool; +use std::collections::HashMap; +use std::str::FromStr; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use uuid::Uuid; + +const APPLICATION_NAME: &str = "awa_lane_bench"; +const WARMUP: Duration = Duration::from_millis(1_500); + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Mode { + InsertMany, + InsertManyCopy, + InsertManySkip, + QueueStorage, + LegacyRowFunction, +} + +impl Mode { + fn parse(raw: &str) -> Self { + match raw { + "insert_many" => Mode::InsertMany, + "insert_many_copy" => Mode::InsertManyCopy, + "insert_many_skip" => Mode::InsertManySkip, + "queue_storage" => Mode::QueueStorage, + "legacy_row_function" => Mode::LegacyRowFunction, + other => panic!("unknown AWA_LANE_BENCH_MODE {other}"), + } + } + + fn label(self) -> &'static str { + match self { + Mode::InsertMany => "insert_many", + Mode::InsertManyCopy => "insert_many_copy", + Mode::InsertManySkip => "insert_many_skip", + Mode::QueueStorage => "queue_storage", + Mode::LegacyRowFunction => "legacy_row_function", + } + } +} + +#[derive(Debug, Clone)] +struct Scenario { + name: String, + mode: Mode, + producers: usize, + batch_size: usize, + unique: bool, + shards: i16, + duplicate_batch_pct: u32, + measured: Duration, + split_commit: bool, + drain: bool, +} + +#[derive(Debug, Default, Clone)] +struct Samples { + /// Latency of the producer call, in milliseconds: one `insert_many` + /// statement, or for a duplicate batch the failed batch plus the + /// row-at-a-time fallback. + batch_ms: Vec, + /// Latency of each single-row statement issued by the fallback. + fallback_row_ms: Vec, + /// Statement and commit latency when the batch runs in an explicit + /// transaction. + statement_ms: Vec, + commit_ms: Vec, + jobs: u64, + batches: u64, + duplicate_batches: u64, + fallback_rows: u64, +} + +impl Samples { + fn merge(&mut self, other: Samples) { + self.batch_ms.extend(other.batch_ms); + self.fallback_row_ms.extend(other.fallback_row_ms); + self.statement_ms.extend(other.statement_ms); + self.commit_ms.extend(other.commit_ms); + self.jobs += other.jobs; + self.batches += other.batches; + self.duplicate_batches += other.duplicate_batches; + self.fallback_rows += other.fallback_rows; + } +} + +#[derive(Debug)] +struct ScenarioResult { + scenario: Scenario, + jobs_per_s: f64, + p50_ms: f64, + p95_ms: f64, + p99_ms: f64, + max_ms: f64, + mean_waiters: f64, + samples: Samples, + landed: i64, + distinct_lane_seqs: i64, + drained: Option<(i64, f64)>, +} + +fn percentile(sorted: &[f64], quantile: f64) -> f64 { + if sorted.is_empty() { + return 0.0; + } + let index = ((sorted.len() as f64 - 1.0) * quantile).round() as usize; + sorted[index.min(sorted.len() - 1)] +} + +fn sorted(values: &[f64]) -> Vec { + let mut copy = values.to_vec(); + copy.sort_by(|left, right| left.partial_cmp(right).expect("latencies are finite")); + copy +} + +fn env_string(name: &str, default: &str) -> String { + std::env::var(name).unwrap_or_else(|_| default.to_string()) +} + +fn env_usize(name: &str, default: usize) -> usize { + std::env::var(name) + .ok() + .and_then(|raw| raw.parse().ok()) + .unwrap_or(default) +} + +fn job_params(queue: &str, tag: &str, producer: usize, counter: u64, unique: bool) -> InsertParams { + InsertParams { + kind: "lane_bench_job".to_string(), + args: serde_json::json!({ + "tag": tag, + "producer": producer, + "n": counter, + "payload": "x".repeat(64), + }), + opts: InsertOpts { + queue: queue.to_string(), + unique: unique.then(|| UniqueOpts { + by_queue: true, + by_args: true, + ..Default::default() + }), + ..Default::default() + }, + } +} + +fn is_unique_violation(err: &AwaError) -> bool { + match err { + AwaError::UniqueConflict { .. } => true, + AwaError::Database(sqlx::Error::Database(db_err)) => { + db_err.code().as_deref() == Some("23505") + } + _ => false, + } +} + +async fn bench_pool(max_connections: u32) -> PgPool { + let options = PgConnectOptions::from_str(&setup::database_url()) + .expect("parse DATABASE_URL") + .application_name(APPLICATION_NAME); + PgPoolOptions::new() + .max_connections(max_connections) + .connect_with(options) + .await + .expect("connect bench pool") +} + +async fn waiter_count(pool: &PgPool) -> i64 { + sqlx::query_scalar( + r#" + SELECT count(*)::bigint + FROM pg_stat_activity + WHERE application_name = $1 + AND wait_event_type = 'Lock' + AND wait_event = 'advisory' + "#, + ) + .bind(APPLICATION_NAME) + .fetch_one(pool) + .await + .unwrap_or(0) +} + +async fn seed_queue_shards(pool: &PgPool, queue: &str, shards: i16) { + sqlx::query( + "INSERT INTO awa.queue_meta (queue, enqueue_shards) VALUES ($1, $2) + ON CONFLICT (queue) DO UPDATE SET enqueue_shards = EXCLUDED.enqueue_shards", + ) + .bind(queue) + .bind(shards) + .execute(pool) + .await + .expect("seed queue_meta.enqueue_shards"); +} + +/// One statement that calls `awa.insert_job_compat` once per row, the shape +/// batch inserts had before `awa.insert_jobs_compat_batch`. It measures what a +/// batch costs the lane lock on this runner when every row goes through the +/// per-row function, which is the yardstick the gate compares `insert_many` +/// against. +const LEGACY_ROW_FUNCTION_SQL: &str = r#" + SELECT inserted.* + FROM unnest( + $1::text[], $2::text[], $3::jsonb[], $4::text[], $5::smallint[], $6::smallint[], + $7::timestamptz[], $8::jsonb[], $9::jsonb[], $10::bytea[], $11::text[] + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + CROSS JOIN LATERAL awa.insert_job_compat( + input.kind, + input.queue, + input.args, + input.state::awa.job_state, + input.priority, + input.max_attempts, + input.run_at, + input.metadata, + ARRAY( + SELECT tag.value + FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, tag_ord) + ORDER BY tag.tag_ord + ), + input.unique_key, + input.unique_states::bit(8), + NULL + ) AS inserted + ORDER BY input.ord +"#; + +async fn legacy_row_function_insert(pool: &PgPool, jobs: &[InsertParams]) -> Result<(), AwaError> { + let rows: Vec<_> = jobs + .iter() + .map(|job| { + prepare_raw_job_insert(job.kind.clone(), job.args.clone(), job.opts.clone()) + .expect("prepare row") + }) + .collect(); + sqlx::query(LEGACY_ROW_FUNCTION_SQL) + .bind(rows.iter().map(|row| row.kind()).collect::>()) + .bind(rows.iter().map(|row| row.queue()).collect::>()) + .bind(rows.iter().map(|row| row.args()).collect::>()) + .bind( + rows.iter() + .map(|row| row.state_db_str()) + .collect::>(), + ) + .bind(rows.iter().map(|row| row.priority()).collect::>()) + .bind( + rows.iter() + .map(|row| row.max_attempts()) + .collect::>(), + ) + .bind(rows.iter().map(|row| row.run_at()).collect::>()) + .bind(rows.iter().map(|row| row.metadata()).collect::>()) + .bind( + rows.iter() + .map(|row| serde_json::json!(row.tags())) + .collect::>(), + ) + .bind(rows.iter().map(|row| row.unique_key()).collect::>()) + .bind( + rows.iter() + .map(|row| row.unique_states_bit_string()) + .collect::>(), + ) + .fetch_all(pool) + .await?; + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +async fn producer_loop( + pool: PgPool, + store: Arc, + scenario: Scenario, + queue: String, + tag: String, + producer: usize, + measure_from: Instant, + stop_at: Instant, + counter_base: Arc, + stop: Arc, +) -> Samples { + let mut samples = Samples::default(); + let mut batch_index: u64 = 0; + let mut previous_batch: Vec = Vec::new(); + + while Instant::now() < stop_at && !stop.load(Ordering::Relaxed) { + let base = counter_base.fetch_add(scenario.batch_size as u64, Ordering::Relaxed); + let mut jobs: Vec = (0..scenario.batch_size as u64) + .map(|offset| job_params(&queue, &tag, producer, base + offset, scenario.unique)) + .collect(); + + // A replayed key: the first job repeats one this producer already + // inserted, so the batch hits exactly one unique conflict. + let duplicate = scenario.unique + && scenario.duplicate_batch_pct > 0 + && !previous_batch.is_empty() + && (batch_index + 1) * u64::from(scenario.duplicate_batch_pct) / 100 + > batch_index * u64::from(scenario.duplicate_batch_pct) / 100; + if duplicate { + jobs[0] = previous_batch[0].clone(); + } + batch_index += 1; + + let started = Instant::now(); + let mut statement_ms = None; + let mut commit_ms = None; + let mut fallback_rows_ms: Vec = Vec::new(); + + let outcome: Result<(), AwaError> = match scenario.mode { + Mode::InsertMany if scenario.split_commit => { + let mut tx = pool.begin().await.expect("begin"); + let statement_started = Instant::now(); + let result = insert_many(&mut *tx, &jobs).await; + let statement_elapsed = statement_started.elapsed(); + match result { + Ok(_) => { + let commit_started = Instant::now(); + tx.commit().await.expect("commit"); + statement_ms = Some(statement_elapsed.as_secs_f64() * 1e3); + commit_ms = Some(commit_started.elapsed().as_secs_f64() * 1e3); + Ok(()) + } + Err(err) => Err(err), + } + } + Mode::InsertMany => insert_many(&pool, &jobs).await.map(|_| ()), + Mode::InsertManyCopy => insert_many_copy_from_pool(&pool, &jobs).await.map(|_| ()), + Mode::InsertManySkip => insert_many_skip_unique_conflicts(&pool, &jobs) + .await + .map(|_| ()), + Mode::QueueStorage => store.enqueue_params_batch(&pool, &jobs).await.map(|_| ()), + Mode::LegacyRowFunction => legacy_row_function_insert(&pool, &jobs).await, + }; + + let mut fallback_count = 0u64; + match outcome { + Ok(()) => {} + Err(err) if is_unique_violation(&err) => { + // What a replaying producer does on a unique conflict: + // insert the batch again one row at a time, skipping + // duplicates. + for job in &jobs { + let row_started = Instant::now(); + match insert_many(&pool, std::slice::from_ref(job)).await { + Ok(_) => {} + Err(err) if is_unique_violation(&err) => {} + Err(err) => panic!("fallback insert failed: {err}"), + } + fallback_rows_ms.push(row_started.elapsed().as_secs_f64() * 1e3); + fallback_count += 1; + } + } + Err(err) => panic!("enqueue failed: {err}"), + } + let elapsed = started.elapsed(); + + if started >= measure_from { + samples.batch_ms.push(elapsed.as_secs_f64() * 1e3); + samples.jobs += scenario.batch_size as u64; + samples.batches += 1; + if let Some(ms) = statement_ms { + samples.statement_ms.push(ms); + } + if let Some(ms) = commit_ms { + samples.commit_ms.push(ms); + } + if fallback_count > 0 { + samples.duplicate_batches += 1; + samples.fallback_rows += fallback_count; + samples.fallback_row_ms.extend(fallback_rows_ms); + } + } + previous_batch = jobs; + } + samples +} + +async fn run_scenario(pool: &PgPool, scenario: &Scenario) -> ScenarioResult { + let queue = format!( + "lane_bench_{}", + Uuid::new_v4().simple().to_string().get(..12).unwrap() + ); + let tag = Uuid::new_v4().simple().to_string(); + seed_queue_shards(pool, &queue, scenario.shards).await; + + let producer_pool = bench_pool(scenario.producers as u32 + 4).await; + let store = Arc::new(QueueStorage::new(QueueStorageConfig::default()).expect("queue storage")); + + let started = Instant::now(); + let measure_from = started + WARMUP; + let stop_at = measure_from + scenario.measured; + let counter_base = Arc::new(AtomicU64::new(0)); + let stop = Arc::new(AtomicBool::new(false)); + + let mut tasks = Vec::new(); + for producer in 0..scenario.producers { + tasks.push(tokio::spawn(producer_loop( + producer_pool.clone(), + Arc::clone(&store), + scenario.clone(), + queue.clone(), + tag.clone(), + producer, + measure_from, + stop_at, + Arc::clone(&counter_base), + Arc::clone(&stop), + ))); + } + + // Sample advisory-lock waiters over the measured window only. + tokio::time::sleep_until(tokio::time::Instant::from_std(measure_from)).await; + let mut waiter_samples: Vec = Vec::new(); + while Instant::now() < stop_at { + waiter_samples.push(waiter_count(pool).await); + tokio::time::sleep(Duration::from_millis(25)).await; + } + + let mut samples = Samples::default(); + for task in tasks { + samples.merge(task.await.expect("producer task")); + } + producer_pool.close().await; + + let (landed, distinct_lane_seqs): (i64, i64) = sqlx::query_as( + r#" + SELECT count(*)::bigint, + count(DISTINCT (enqueue_shard, lane_seq))::bigint + FROM awa.ready_entries + WHERE queue = $1 + "#, + ) + .bind(&queue) + .fetch_one(pool) + .await + .expect("count landed rows"); + + // Claim everything that landed, from every shard, with the claim path the + // workers use. + let drained = if scenario.drain { + let drain_started = Instant::now(); + let mut claimed = 0i64; + loop { + let batch = store + .claim_runtime_batch(pool, &queue, 500, Duration::from_secs(3_600)) + .await + .expect("claim batch"); + if batch.is_empty() { + break; + } + claimed += batch.len() as i64; + } + Some((claimed, drain_started.elapsed().as_secs_f64())) + } else { + None + }; + + let latencies = sorted(&samples.batch_ms); + let window = scenario.measured.as_secs_f64(); + ScenarioResult { + scenario: scenario.clone(), + jobs_per_s: samples.jobs as f64 / window, + p50_ms: percentile(&latencies, 0.50), + p95_ms: percentile(&latencies, 0.95), + p99_ms: percentile(&latencies, 0.99), + max_ms: latencies.last().copied().unwrap_or(0.0), + mean_waiters: if waiter_samples.is_empty() { + 0.0 + } else { + waiter_samples.iter().sum::() as f64 / waiter_samples.len() as f64 + }, + samples, + landed, + distinct_lane_seqs, + drained, + } +} + +fn print_result(result: &ScenarioResult) { + let scenario = &result.scenario; + println!( + "[lane-bench] {name} mode={mode} producers={producers} batch={batch} unique={unique} \ + shards={shards} dup_pct={dup} | {rate:.0} jobs/s | batch ms p50={p50:.1} p95={p95:.1} \ + p99={p99:.1} max={max:.1} | waiters_mean={waiters:.1} | landed={landed} \ + distinct_seqs={distinct}", + name = scenario.name, + mode = scenario.mode.label(), + producers = scenario.producers, + batch = scenario.batch_size, + unique = scenario.unique, + shards = scenario.shards, + dup = scenario.duplicate_batch_pct, + rate = result.jobs_per_s, + p50 = result.p50_ms, + p95 = result.p95_ms, + p99 = result.p99_ms, + max = result.max_ms, + waiters = result.mean_waiters, + landed = result.landed, + distinct = result.distinct_lane_seqs, + ); + if !result.samples.statement_ms.is_empty() { + let statement = sorted(&result.samples.statement_ms); + let commit = sorted(&result.samples.commit_ms); + println!( + "[lane-bench] split: statement ms p50={:.1} p99={:.1} | commit ms p50={:.2} p99={:.2}", + percentile(&statement, 0.5), + percentile(&statement, 0.99), + percentile(&commit, 0.5), + percentile(&commit, 0.99), + ); + } + if let Some((claimed, seconds)) = result.drained { + println!( + "[lane-bench] drain: claimed {claimed} of {} landed in {seconds:.2}s ({:.0}/s)", + result.landed, + claimed as f64 / seconds.max(1e-9), + ); + } + if result.samples.duplicate_batches > 0 { + let rows = sorted(&result.samples.fallback_row_ms); + println!( + "[lane-bench] fallback: {} duplicate batches, {} single-row statements, \ + row ms p50={:.1} p99={:.1} max={:.1}", + result.samples.duplicate_batches, + result.samples.fallback_rows, + percentile(&rows, 0.5), + percentile(&rows, 0.99), + rows.last().copied().unwrap_or(0.0), + ); + } +} + +fn emit_json(result: &ScenarioResult) { + let mut outcomes = HashMap::new(); + outcomes.insert("jobs".to_string(), result.samples.jobs); + outcomes.insert("batches".to_string(), result.samples.batches); + BenchmarkResult { + schema_version: SCHEMA_VERSION, + scenario: result.scenario.name.clone(), + language: "rust".to_string(), + seeded: result.landed.max(0) as u64, + metrics: BenchMetrics { + throughput: None, + enqueue_per_s: Some(result.jobs_per_s), + drain_time_s: None, + latency_ms: Some(BenchLatency { + p50: Some(result.p50_ms), + p95: Some(result.p95_ms), + p99: Some(result.p99_ms), + }), + rescue: None, + }, + outcomes, + metadata: Some(serde_json::json!({ + "producers": result.scenario.producers, + "batch_size": result.scenario.batch_size, + "unique": result.scenario.unique, + "shards": result.scenario.shards, + "mean_lock_waiters": result.mean_waiters, + })), + } + .emit(); +} + +/// Lane integrity: every job that was acknowledged landed exactly once, under +/// a distinct `(shard, lane_seq)`, and every shard's sequence numbers are +/// dense (no producer lost or duplicated a reservation). +fn assert_lane_integrity(result: &ScenarioResult) { + if let Some((claimed, _)) = result.drained { + assert_eq!( + claimed, result.landed, + "{}: claim path did not drain every shard", + result.scenario.name + ); + } + assert_eq!( + result.landed, result.distinct_lane_seqs, + "{}: ready rows share a (shard, lane_seq)", + result.scenario.name + ); + assert!( + result.landed as u64 >= result.samples.jobs, + "{}: {} jobs acknowledged in the window but only {} rows landed", + result.scenario.name, + result.samples.jobs, + result.landed + ); +} + +async fn bench_setup() -> PgPool { + let pool = setup::setup(16).await; + setup::activate_queue_storage(&pool).await; + pool +} + +fn scenario( + name: &str, + producers: usize, + shards: i16, + unique: bool, + duplicate_batch_pct: u32, +) -> Scenario { + Scenario { + name: name.to_string(), + mode: Mode::InsertManySkip, + producers, + batch_size: 185, + unique, + shards, + duplicate_batch_pct, + measured: Duration::from_secs(8), + split_commit: false, + drain: false, + } +} + +/// Fixed scenarios whose results gate regressions in the nightly run. +/// +/// Absolute floors sit well below the rates a shared runner sustains, and the +/// latency ceiling well above, so they only catch a lane that has stopped +/// making progress; `benchmarks/baseline.json` carries the regression +/// thresholds on rate and p99. What the gate asserts about lock hold time is +/// relative, because it must hold on any runner: one producer's 185-row batch +/// has to take well under the time the same rows take through the per-row +/// function, measured in the same run, and sixteen producers on one lane have +/// to out-run what the per-row function could ever sustain for it. Every +/// bound goes through `ci_timing` so a contended runner loosens it. +#[tokio::test(flavor = "multi_thread", worker_threads = 6)] +#[ignore] +async fn test_enqueue_lane_contention_gate() { + let pool = bench_setup().await; + + let mut row_function = scenario("enqueue_lane_1p_row_function_185", 1, 1, true, 0); + row_function.mode = Mode::LegacyRowFunction; + row_function.measured = Duration::from_secs(4); + let row_function = run_scenario(&pool, &row_function).await; + print_result(&row_function); + assert_lane_integrity(&row_function); + + let mut one_producer = scenario("enqueue_lane_1p_unique_185", 1, 1, true, 0); + one_producer.measured = Duration::from_secs(4); + let one_producer = run_scenario(&pool, &one_producer).await; + print_result(&one_producer); + assert_lane_integrity(&one_producer); + + let single_lane = run_scenario( + &pool, + &scenario("enqueue_lane_16p_unique_185", 16, 1, true, 0), + ) + .await; + print_result(&single_lane); + emit_json(&single_lane); + assert_lane_integrity(&single_lane); + + let sharded = run_scenario( + &pool, + &scenario("enqueue_lane_16p_unique_185_shards4", 16, 4, true, 0), + ) + .await; + print_result(&sharded); + emit_json(&sharded); + assert_lane_integrity(&sharded); + + let floor = ci_timing::contention_floor(GATE_MIN_JOBS_PER_S); + assert!( + single_lane.jobs_per_s >= floor as f64, + "single-lane enqueue rate {:.0}/s fell below the floor of {floor}/s", + single_lane.jobs_per_s + ); + let p99_ceiling = ci_timing::scaled_timeout(Duration::from_millis(GATE_MAX_P99_MS)); + assert!( + Duration::from_secs_f64(single_lane.p99_ms / 1e3) <= p99_ceiling, + "single-lane p99 batch latency {:.0}ms exceeded {:?}", + single_lane.p99_ms, + p99_ceiling + ); + + // The batch holds the lane lock for a fraction of the per-row function's + // time, and the lane's ceiling follows. Both bounds are ratios of numbers + // measured moments apart on the same machine. + let max_hold_ratio = (GATE_MAX_HOLD_RATIO * ci_timing::chaos_timeout_multiplier()).min(0.9); + let hold_ratio = one_producer.p50_ms / row_function.p50_ms; + assert!( + hold_ratio <= max_hold_ratio, + "a 185-row batch took {:.1}ms against {:.1}ms through the per-row function \ + (ratio {hold_ratio:.2}, limit {max_hold_ratio:.2})", + one_producer.p50_ms, + row_function.p50_ms + ); + let min_scaling = GATE_MIN_LANE_SCALING / ci_timing::chaos_timeout_multiplier(); + let lane_scaling = single_lane.jobs_per_s / row_function.jobs_per_s; + assert!( + lane_scaling >= min_scaling.max(1.0), + "16 producers on one lane sustain {:.0}/s, {lane_scaling:.1}x the per-row function's \ + single-producer {:.0}/s (need {:.1}x)", + single_lane.jobs_per_s, + row_function.jobs_per_s, + min_scaling.max(1.0) + ); +} + +const GATE_MIN_JOBS_PER_S: i64 = 2_000; +const GATE_MAX_P99_MS: u64 = 15_000; +const GATE_MAX_HOLD_RATIO: f64 = 0.4; +const GATE_MIN_LANE_SCALING: f64 = 2.0; + +/// Exploration sweep; see the module docs for the tunables. +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore] +async fn test_enqueue_lane_contention_matrix() { + let pool = bench_setup().await; + if env_usize("AWA_LANE_BENCH_WITHOUT_PATCH", 0) != 0 { + sqlx::query( + "DROP FUNCTION IF EXISTS awa.insert_jobs_compat_batch(text[], text[], jsonb[], \ + text[], smallint[], smallint[], timestamptz[], jsonb[], jsonb[], bytea[], text[], \ + text[], boolean)", + ) + .execute(&pool) + .await + .expect("drop the batch function"); + } + let producers: Vec = env_string("AWA_LANE_BENCH_PRODUCERS", "1,4,8,16,32") + .split(',') + .filter_map(|raw| raw.trim().parse().ok()) + .collect(); + let mode = Mode::parse(&env_string("AWA_LANE_BENCH_MODE", "insert_many")); + let batch_size = env_usize("AWA_LANE_BENCH_BATCH", 185); + let unique = env_usize("AWA_LANE_BENCH_UNIQUE", 1) != 0; + let shards = env_usize("AWA_LANE_BENCH_SHARDS", 1) as i16; + let duplicate_pct = env_usize("AWA_LANE_BENCH_DUP_PCT", 0) as u32; + let secs = env_usize("AWA_LANE_BENCH_SECS", 8) as u64; + let split = env_usize("AWA_LANE_BENCH_SPLIT", 0) != 0; + + for producer_count in producers { + let result = run_scenario( + &pool, + &Scenario { + name: format!("matrix_{producer_count}p"), + mode, + producers: producer_count, + batch_size, + unique, + shards, + duplicate_batch_pct: duplicate_pct, + measured: Duration::from_secs(secs), + split_commit: split, + drain: env_usize("AWA_LANE_BENCH_DRAIN", 0) != 0, + }, + ) + .await; + print_result(&result); + assert_lane_integrity(&result); + } +} diff --git a/benchmarks/baseline.json b/benchmarks/baseline.json index 9f7fee869..e4ec128e0 100644 --- a/benchmarks/baseline.json +++ b/benchmarks/baseline.json @@ -32,6 +32,13 @@ "rust/stale_heartbeat_rescue": { "handler_per_s": 160.3 }, + "rust/enqueue_lane_16p_unique_185": { + "enqueue_per_s": 10000.0, + "p99_ms": 2500.0 + }, + "rust/enqueue_lane_16p_unique_185_shards4": { + "enqueue_per_s": 10000.0 + }, "rust/mvcc_horizon_overlap": { "min_overlap_vs_baseline": 0.75, "min_cooldown_vs_baseline": 0.75, From 7fc1ff1d09512adef12b5cc3d16b560282172fdd Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sun, 8 Nov 2026 22:00:36 +1300 Subject: [PATCH 3/6] docs: batch compatibility insert and the 0.6.15 upgrade --- CHANGELOG.md | 15 +++++++++ README.md | 1 + correctness/storage/MAPPING.md | 2 +- docs/architecture.md | 4 ++- docs/troubleshooting.md | 17 ++++++++-- docs/upgrade-0.6.14-to-0.6.15.md | 57 ++++++++++++++++++++++++++++++++ 6 files changed, 91 insertions(+), 5 deletions(-) create mode 100644 docs/upgrade-0.6.14-to-0.6.15.md diff --git a/CHANGELOG.md b/CHANGELOG.md index a2fb1494b..cacb9aaa6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,21 @@ Notable changes between releases. Detailed migration notes for storage transitio ## [Unreleased] +### Upgrading + +- The schema stays at v40. Run `awa migrate` as the schema owner, before or after rolling binaries, while traffic keeps running. It applies the repeatable `insert_jobs_compat_batch` patch, which adds one function and changes no existing object. Binaries can roll back to 0.6.14 with the patch left in place. See the [upgrade guide](docs/upgrade-0.6.14-to-0.6.15.md). +- External migration runners: apply `R__insert_jobs_compat_batch.sql` after `V40` (`awa migrate --extract-to --pending`), replacing all `R__` files together. +- Deployments that revoke `EXECUTE` on new functions from `PUBLIC` should grant it on `awa.insert_jobs_compat_batch` to runtime roles; without it, `insert_many_copy` keeps its per-row path. + +### Changed + +- Compatibility COPY inserts (`insert_many_copy` in Rust and Python) hold a queue lane's enqueue lock only while they write that lane, instead of from the batch's first ready row until commit. Many producers enqueuing large batches into one queue no longer queue behind each other's whole batches, and promotion into a busy queue waits less. The batch goes through the new `awa.insert_jobs_compat_batch` function, which prepares job ids, unique claims and payloads before it reserves each lane once, in the same sorted lane order as before. Each call checks its connection for the function and keeps the per-row path where it is missing or not executable. Results are unchanged: input order, with jobs whose unique key is already claimed skipped. + +### Added + +- `insert_many_skip_unique_conflicts` (Rust) inserts a batch through `awa.insert_jobs_compat_batch`, skipping jobs whose unique key is already claimed, and returns the inserted rows in input order. Use it instead of retrying a failed `insert_many` one job at a time, which takes the lane lock once per job. It needs the `insert_jobs_compat_batch` patch. Generic `insert_many` is unchanged. +- `enqueue_lane_contention_test` benchmarks many producers on one lane, with a nightly gate on per-lane throughput and lock hold time. + ## [0.6.14] — 2026-11-08 Recommended for everyone on 0.6.13 using queue storage. A busy queue no longer stalls promotion of the others, prune backs off cleanly under lock contention, new queue lanes are usable by every runtime role, and jobs no longer stay running, run again or get misreported after priority aging. diff --git a/README.md b/README.md index 29851d55b..3bcd74cb1 100644 --- a/README.md +++ b/README.md @@ -389,6 +389,7 @@ All coordination through Postgres. The Rust runtime owns dispatch, leases, heart | [Migration guide](docs/migrations.md) | Fresh installs, upgrades, extracted SQL, rollback strategy | | [0.5 → 0.6 upgrade](docs/upgrade-0.5-to-0.6.md) | Step-by-step operator checklist for the staged storage transition | | [0.6.9 → 0.6.10 upgrade](docs/upgrade-0.6.9-to-0.6.10.md) | Rolling patch upgrade, rollback, and optional runtime privilege tightening | +| [0.6.14 → 0.6.15 upgrade](docs/upgrade-0.6.14-to-0.6.15.md) | Rolling patch upgrade for the batch compatibility insert | | [Configuration reference](docs/configuration.md) | `QueueConfig`, `ClientBuilder`, Python `start()`, env vars | | [HTTP workers and callbacks](docs/http-callbacks.md) | Serverless dispatch, callback receiver endpoints, signature verification | | [Lifecycle hooks](docs/lifecycle-hooks.md) | React to job start/finish/retry/callback events for metrics, alerting, cleanup | diff --git a/correctness/storage/MAPPING.md b/correctness/storage/MAPPING.md index fb142d098..bd350018f 100644 --- a/correctness/storage/MAPPING.md +++ b/correctness/storage/MAPPING.md @@ -37,7 +37,7 @@ The TLA+ lifecycle model does not represent the completed-history rollup cache, | TLA+ action | Rust function | SQL / DDL | | --- | --- | --- | -| `EnqueueReady(j)` | `QueueStorage::insert_ready_rows_tx`, `QueueStorage::insert_ready_rows_copy_tx`, and SQL compatibility `awa.insert_job_compat`; producer entry points include `insert` / `insert_with`, `enqueue_batch`, `enqueue_runtime_rows`, `enqueue_params_batch`, and `enqueue_params_copy` | reserve a lane sequence range through `{schema}.reserve_enqueue_seq` under a transaction-scoped lane advisory lock, sync enqueue-time uniqueness claims, append to `{schema}.ready_entries` via INSERT or COPY, append one `{schema}.ready_segments` row per contiguous committed lane range, and notify logical queues in one tx | +| `EnqueueReady(j)` | `QueueStorage::insert_ready_rows_tx`, `QueueStorage::insert_ready_rows_copy_tx`, and SQL compatibility `awa.insert_job_compat` / `awa.insert_jobs_compat_batch`; producer entry points include `insert` / `insert_with`, `insert_many`, `insert_many_copy`, `insert_many_skip_unique_conflicts`, `enqueue_batch`, `enqueue_runtime_rows`, `enqueue_params_batch`, and `enqueue_params_copy` | reserve a lane sequence range through `{schema}.reserve_enqueue_seq` under a transaction-scoped lane advisory lock, sync enqueue-time uniqueness claims, append to `{schema}.ready_entries` via INSERT or COPY, append one `{schema}.ready_segments` row per contiguous committed lane range, and notify logical queues in one tx | | `EnqueueDeferred(j)` | `QueueStorage::insert_deferred_rows_tx` and `QueueStorage::insert_deferred_rows_copy_tx`; producer entry points include `enqueue_params_batch` and `enqueue_params_copy` | allocate job ids, sync enqueue-time uniqueness claims, append to `{schema}.deferred_jobs` via INSERT or COPY in one tx | | `PromoteDeferred(j)` | `QueueStoragePromoter::promote_due` in `awa-worker/src/maintenance.rs` -> `QueueStorage::promote_due_lanes` / `promote_due_lane` -> `promote_lane_tx` | Per `(queue, priority)` lane, `DELETE FROM deferred_jobs ... INSERT INTO ready_entries ...` in one tx under a short transaction-local `lock_timeout`. A timeout rolls the lane back with its rows still deferred (a stutter step in the model) and hands the lane to a bounded waiter. Claimers are notified once per batch after the lanes commit. | | `AdvanceClaimCursor` | claim path spent-prefix advancement and post-commit `QueueStorage::advance_claim_cursors` | `claim_ready_runtime()` may advance the sequence-backed claim cursor only across a contiguous prefix of committed spent/tombstoned lanes. In receipts mode, already-emitted attempts are proved from queue-slot-local `{schema}.ready_claim_attempt_batches` range evidence so recovery does not scan claim-ring parent tables while old claim partitions are locked for prune. Newly emitted claims advance the cursor after the claim transaction commits. | diff --git a/docs/architecture.md b/docs/architecture.md index c4437b0e7..4139bd87b 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -155,7 +155,9 @@ Enqueue is transactional: if the producer's outer transaction rolls back, the jo There are two COPY-shaped producer paths: - Direct queue-storage COPY is the high-throughput path: Rust producers use a configured `QueueStorage::enqueue_params_copy()`, and Python producers use `enqueue_many_copy()`. Direct-copy producers must use the same queue-storage configuration as the worker fleet, especially `queue_stripe_count` / `queue_storage_queue_stripe_count`. -- `insert_many_copy()` is the compatibility path. It stages rows with COPY but then routes each row through `awa.insert_job_compat()`, so it preserves the public compatibility surface rather than bypassing to the direct producer path. +- `insert_many_copy()` is the compatibility path. It stages rows with COPY, then inserts them through `awa.insert_jobs_compat_batch()` (the `insert_jobs_compat_batch` schema patch), which reserves each lane once, in canonical lane order, after the batch's lane-independent work, and skips unique conflicts in the same statement. It probes its own connection on every call; on a schema without the patch, or for a role without `EXECUTE` on the function, it routes each row through `awa.insert_job_compat()`. Either way it preserves the public compatibility surface rather than bypassing to the direct producer path. + +Generic `insert_many` always uses one statement that calls `awa.insert_job_compat()` per row in canonical lane order: its `PgExecutor` is consumed by that statement, so it cannot probe the schema first. `insert_many_skip_unique_conflicts` requires the batch function. Claim is cursor-based rather than heap-scan based: diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 46605e5f9..26586a1a3 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -194,7 +194,7 @@ Reading them together separates the two usual failure modes: **1. Producer is going through the compatibility insert path.** -`insert_many_copy_from_pool` / `client.insert_many_copy` routes every row through `awa.insert_job_compat()` once per row. On any database with non-trivial per-statement latency (auth-proxy hop, managed Postgres, network) the per-row function call dominates — measured at ~100–150 ms per row through a Cloud SQL Auth Proxy in staging, which caps a single producer at ~7 rows/sec regardless of batch or chunk size. +Without the `insert_jobs_compat_batch` schema patch, `insert_many_copy_from_pool` / `client.insert_many_copy` routes every row through `awa.insert_job_compat()` once per row, and a batch with unique jobs adds a savepoint round trip per row. On any database with non-trivial per-statement latency (auth-proxy hop, managed Postgres, network) the per-row cost dominates — measured at ~100–150 ms per row through an auth proxy, which caps a single producer at ~7 rows/sec regardless of batch or chunk size. Run `awa migrate` to install the patch; COPY then inserts the batch in one function call. Generic `insert_many` keeps its one-statement per-row path either way. Switch the producer to the direct queue-storage COPY entry point: @@ -222,7 +222,18 @@ DO UPDATE SET enqueue_shards = EXCLUDED.enqueue_shards; Use an upsert: first-enqueue may create lane rows before any operator inserts a `queue_meta` row, so a plain `UPDATE` quietly affects zero rows. Raising `enqueue_shards` is a semantic switch from strict to partitioned FIFO; see [ADR-025](adr/025-sharded-enqueue-heads.md). -**3. WAL or commit pressure on the database.** +**3. Producers wait on one queue's enqueue lane.** + +Many connections wait with `wait_event = 'advisory'` while the database is otherwise idle, batch inserts take seconds although the work per batch is milliseconds, and promotion into the same queue runs late while producers saturate it. + +On queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane until it commits, so concurrent producers for one lane run one at a time; a lane's ceiling is a batch's rows divided by the time the lock is held. Generic `insert_many` and per-row COPY take the lock at the batch's first ready row and hold it while the rest of the batch is processed. With the `insert_jobs_compat_batch` schema patch, `insert_many_copy` and `insert_many_skip_unique_conflicts` prepare the whole batch (job ids, unique claims, payloads) before they take any lane lock, and hold it only for the lane write and the commit. + +- Run `awa migrate` and move hot producers to `insert_many_copy` or `insert_many_skip_unique_conflicts`. +- When `insert_many` fails on a unique conflict, retry with `insert_many_skip_unique_conflicts` rather than inserting the batch again row by row: each single-row insert queues for the lane separately. +- Keep producer transactions short after the enqueue: the lane stays locked until commit. +- Raise `enqueue_shards` (above) for a queue whose producers still saturate one lane. + +**4. WAL or commit pressure on the database.** If batch size and shard count both look healthy, sample `pg_stat_activity` for `LWLock:WALWrite` / `LWLock:WALSync` waits and confirm the database is sized for the offered rate. The per-vCPU sustained-completion and burst-enqueue numbers in [`docs/deploying-on-managed-postgres.md`](deploying-on-managed-postgres.md#pick-a-vcpu-size) are useful reference points. @@ -447,7 +458,7 @@ awa --database-url "$DATABASE_URL" migrate Cause: - on queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane until the transaction commits; the `DETAIL` names two 64-bit advisory keys, and two transactions that took lanes in opposite orders hold one each and wait on the other's -- `insert_many`, `insert_many_copy` and the `QueueStorage` batch enqueues take their lane locks in one canonical order, `(queue, priority, shard)` ascending with queue names compared bytewise, so batches never deadlock with each other +- `insert_many`, `insert_many_copy`, `insert_many_skip_unique_conflicts` and the `QueueStorage` batch enqueues take their lane locks in one canonical order, `(queue, priority, shard)` ascending with queue names compared bytewise, so batches never deadlock with each other over lane locks; unique keys and the application's own locks can still form cycles - a transaction that enqueues with several separate statements (`insert` per job, handler-staged follow-ups, `Transaction.insert` in a loop) takes locks in statement order and can deadlock against a batch or another such transaction that touches the same lanes in a different order Fix: diff --git a/docs/upgrade-0.6.14-to-0.6.15.md b/docs/upgrade-0.6.14-to-0.6.15.md new file mode 100644 index 000000000..09851ae93 --- /dev/null +++ b/docs/upgrade-0.6.14-to-0.6.15.md @@ -0,0 +1,57 @@ +# Upgrading from 0.6.14 to 0.6.15 + +0.6.15 is a rolling patch upgrade. Traffic keeps running throughout, binaries and `awa migrate` can go in either order, and you can roll binaries back to 0.6.14 with the schema patch left in place. + +## Upgrade + +1. Install the 0.6.15 CLI. +2. Run `awa migrate` once, with the schema owner as the effective role, while workers and producers keep running: + + ```bash + PGOPTIONS='-c role=awa_owner' \ + awa --database-url "$AWA_MIGRATOR_DATABASE_URL" migrate + ``` + + The output reports `Schema patch insert_jobs_compat_batch: applied`. + +3. Roll workers, producers and admin processes to 0.6.15, in any order. + +Steps 2 and 3 can also run the other way round. A 0.6.15 producer checks for the patch on every `insert_many_copy` call, so it switches to the batch function on its next call after the patch commits, without a restart. + +## What `awa migrate` changes + +The schema stays at v40. The `insert_jobs_compat_batch` patch adds one function, `awa.insert_jobs_compat_batch`, and changes nothing else: no table DDL, no table or sequence locks, no change to `awa.insert_job_compat` or to storage routing. It applies as one transaction, so an interrupted run leaves nothing applied; run it again. + +The function inserts a whole batch the way `awa.insert_job_compat` inserts one job, with the same validation, unique claims, lane assignment and storage routing. It does everything that does not depend on a lane first (job ids, unique claims in key order, new lane rows, scheduled rows, the result rows), then reserves each lane's range of positions once, in `(queue, priority, shard)` order, and writes the ready rows. Each lane's enqueue lock is therefore held only for the reservation, the ready-row write and the commit, instead of from the batch's first ready row to the commit. Concurrent producers and promotion on the same lane wait for that much shorter window. + +What callers see: + +- `insert_many_copy` (Rust) and `Client.insert_many_copy` (Python) use the function when it is installed and the connected role may execute it, otherwise the per-row path they used before. Their results are unchanged: rows in input order, jobs whose unique key is already claimed skipped. +- `insert_many_skip_unique_conflicts` is new in Rust. It inserts a batch through the function and skips jobs whose unique key is already claimed, returning the inserted rows in input order. It needs the patch and returns a validation error without it. +- Generic `insert_many` is unchanged: one statement calling `awa.insert_job_compat` per row, on every schema, and a unique conflict still fails the whole batch. +- A batch's rows in one lane occupy consecutive lane positions in input order, as before, and are described by one ready segment per run instead of one per row. A key repeated inside one batch keeps its first job in input order. + +0.6.14 and earlier binaries never call the new function. They keep calling `awa.insert_job_compat`, which takes the same lane lock, so old and new producers share a lane without reordering it. + +### Privileges + +New functions are executable by `PUBLIC` unless your deployment revokes that through default privileges. If it does and the runtime role's default privileges do not cover functions created by the migration role, grant it explicitly: + +```sql +GRANT EXECUTE ON FUNCTION awa.insert_jobs_compat_batch( + text[], text[], jsonb[], text[], smallint[], smallint[], timestamptz[], + jsonb[], jsonb[], bytea[], text[], text[], boolean +) TO awa_runtime; +``` + +Without `EXECUTE`, `insert_many_copy` keeps the per-row path and `insert_many_skip_unique_conflicts` fails with `permission denied`. + +### External migration runners + +`awa migrate --sql --pending` prints the `R__insert_jobs_compat_batch` patch on a 0.6.14 database. To write it as a file instead, run `awa migrate --extract-to --pending` and apply `R__insert_jobs_compat_batch.sql` as a repeatable script after `V40`. It is a single `CREATE OR REPLACE FUNCTION`, safe inside the runner's transaction and safe to re-run. Run it as the role that owns the existing `awa` functions, normally the schema owner. + +Replace the runner's `R__` files with the set the 0.6.15 CLI extracts, together, and keep the `V` files exactly as originally applied. + +## Roll back + +Roll binaries back to 0.6.14 at any time. Leave the patch in place: released binaries do not call the function, and it holds no state. From b452969a86e1d1fb9f1026dc3af9bde6d63a8971 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sun, 8 Nov 2026 22:11:11 +1300 Subject: [PATCH 4/6] test: serialise the bind-limit insert test with storage switches --- awa/tests/integration_test.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/awa/tests/integration_test.rs b/awa/tests/integration_test.rs index a0823d3c3..742826e73 100644 --- a/awa/tests/integration_test.rs +++ b/awa/tests/integration_test.rs @@ -552,6 +552,9 @@ async fn test_insert_many() { #[tokio::test] async fn test_insert_many_exceeding_bind_parameter_limit() { const JOB_COUNT: usize = 6_000; + // A concurrent test that activates queue storage would route part of the + // batch away from the canonical tables this test reads back. + let _guard = test_lock().lock().await; let client = setup_with_connections(4).await; let queue = "integ_insert_many_bind_limit"; clean_queue(client.pool(), queue).await; From d2de3738336a6c2924e867afacb939ac3a98f39d Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Sun, 8 Nov 2026 23:56:38 +1300 Subject: [PATCH 5/6] fix(insert): register each new lane only when the batch reaches it (0.6) awa.insert_jobs_compat_batch registered every lane it was the first to use before reserving any. The head tables' row trigger provisions a lane under its lock and the uncommitted head row blocks other producers of that lane, so a batch for an existing lane A and a new lane B held B while waiting for A, and a per-row producer (generic insert_many, single inserts, 0.6.14 binaries) holding A and then enqueuing into B deadlocked with it. The batch now registers each missing lane inside the reservation loop, immediately before reserving it, so new lanes follow the register-then-reserve order every other producer uses. The patch probe matches the new body, so `awa migrate` replaces the earlier one. Documents that a per-row transaction sharing a unique key with a batch can now fail with 40P01 rather than skip the job or raise 23505, and corrects the lane-order statements in the changelog, upgrade guide and troubleshooting. --- CHANGELOG.md | 3 +- .../patches/insert_jobs_compat_batch.sql | 108 ++-- awa-model/src/insert.rs | 19 +- awa-model/src/migrations.rs | 2 +- ...obs_compat_batch_registers_lanes_first.sql | 553 ++++++++++++++++++ awa/tests/insert_many_batch_test.rs | 82 ++- awa/tests/insert_many_schema_fallback_test.rs | 95 ++- docs/troubleshooting.md | 4 +- docs/upgrade-0.6.14-to-0.6.15.md | 6 +- 9 files changed, 764 insertions(+), 108 deletions(-) create mode 100644 awa/tests/fixtures/insert_jobs_compat_batch_registers_lanes_first.sql diff --git a/CHANGELOG.md b/CHANGELOG.md index 1de88e8a4..06c5c873a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,7 +12,8 @@ Notable changes between releases. Detailed migration notes for storage transitio ### Changed -- Compatibility COPY inserts (`insert_many_copy` in Rust and Python) hold a queue lane's enqueue lock only while they write that lane, instead of from the batch's first ready row until commit. Many producers enqueuing large batches into one queue no longer queue behind each other's whole batches, and promotion into a busy queue waits less. The batch goes through the new `awa.insert_jobs_compat_batch` function, which prepares job ids, unique claims and payloads before it reserves each lane once, in the same sorted lane order as before. Each call checks its connection for the function and keeps the per-row path where it is missing or not executable. Results are unchanged: input order, with jobs whose unique key is already claimed skipped. +- Compatibility COPY inserts (`insert_many_copy` in Rust and Python) hold a queue lane's enqueue lock only while they write that lane, instead of from the batch's first ready row until commit. Many producers enqueuing large batches into one queue no longer queue behind each other's whole batches, and promotion into a busy queue waits less. The batch goes through the new `awa.insert_jobs_compat_batch` function, which prepares job ids, unique claims and payloads first, then registers each new lane and reserves each lane once, in the sorted `(queue, priority, shard)` order every other producer uses. Each call checks its connection for the function and keeps the per-row path where it is missing or not executable. Results are unchanged: input order, with jobs whose unique key is already claimed skipped. +- A batch through `awa.insert_jobs_compat_batch` takes all its unique claims before any lane lock, while per-row inserts (generic `insert_many`, single-job inserts, 0.6.14 and earlier binaries) take each job's claim and lane together, lane by lane. A per-row transaction that already holds a lane the batch needs and then inserts a job whose unique key the batch has claimed but not committed now fails with `deadlock detected` (SQLSTATE `40P01`) instead of skipping the job or failing with a unique violation (`23505`). Retry the aborted transaction. Batches through the function do not deadlock with each other this way. - Queue-storage job lists without a kind filter select receipt candidate IDs before loading their ready bodies. Kind filters still apply before the limit. ### Added diff --git a/awa-model/migrations/patches/insert_jobs_compat_batch.sql b/awa-model/migrations/patches/insert_jobs_compat_batch.sql index f3b718311..665537a56 100644 --- a/awa-model/migrations/patches/insert_jobs_compat_batch.sql +++ b/awa-model/migrations/patches/insert_jobs_compat_batch.sql @@ -7,10 +7,11 @@ -- transaction-scoped, holds it while every remaining row is processed. Many -- concurrent producers on one lane therefore queue behind whole batches. This -- patch adds awa.insert_jobs_compat_batch, which does the lane-independent --- work first (job ids, unique claims in key order, lane rows, payloads), then --- reserves each lane's sequence range once, in (queue, priority, shard) order, --- and writes the ready rows and segments in two statements. The lock is held --- across the reservation, those two statements and the commit. +-- work first (job ids, unique claims in key order, scheduled rows, result +-- rows), then, in (queue, priority, shard) order, registers each new lane and +-- reserves each lane's sequence range once, and writes the ready rows and +-- segments in two statements. The lock is held across the reservation, those +-- two statements and the commit. -- -- The column arrays are parallel and must have equal length. Every element is -- required except run_at, unique_key, unique_states and ordering_hash; @@ -302,62 +303,6 @@ BEGIN ORDER BY input.unique_key, input.ord; END IF; - -- Register lanes this batch is the first to use. The head tables carry a - -- row trigger that takes the lane lock before the conflict check, so only - -- missing lanes are proposed: a batch for an existing lane must not queue - -- behind that lane's producers before it has any work to show for it. - INSERT INTO queue_lanes (queue, priority) - SELECT lanes.queue, lanes.priority - FROM ( - SELECT DISTINCT input.queue, input.priority - FROM unnest(p_queue, p_priority, p_state) - WITH ORDINALITY AS input(queue, priority, state, ord) - WHERE input.state = 'available' - AND v_job_ids[input.ord] <> ALL (v_skipped) - ) AS lanes - WHERE NOT EXISTS ( - SELECT 1 FROM queue_lanes AS existing - WHERE existing.queue = lanes.queue AND existing.priority = lanes.priority - ) - ORDER BY lanes.queue COLLATE "C", lanes.priority - ON CONFLICT ON CONSTRAINT queue_lanes_pkey DO NOTHING; - - INSERT INTO queue_enqueue_heads (queue, priority, enqueue_shard) - SELECT lanes.queue, lanes.priority, lanes.shard - FROM ( - SELECT DISTINCT input.queue, input.priority, v_shards[input.ord] AS shard - FROM unnest(p_queue, p_priority, p_state) - WITH ORDINALITY AS input(queue, priority, state, ord) - WHERE input.state = 'available' - AND v_job_ids[input.ord] <> ALL (v_skipped) - ) AS lanes - WHERE NOT EXISTS ( - SELECT 1 FROM queue_enqueue_heads AS existing - WHERE existing.queue = lanes.queue - AND existing.priority = lanes.priority - AND existing.enqueue_shard = lanes.shard - ) - ORDER BY lanes.queue COLLATE "C", lanes.priority, lanes.shard - ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; - - INSERT INTO queue_claim_heads (queue, priority, enqueue_shard) - SELECT lanes.queue, lanes.priority, lanes.shard - FROM ( - SELECT DISTINCT input.queue, input.priority, v_shards[input.ord] AS shard - FROM unnest(p_queue, p_priority, p_state) - WITH ORDINALITY AS input(queue, priority, state, ord) - WHERE input.state = 'available' - AND v_job_ids[input.ord] <> ALL (v_skipped) - ) AS lanes - WHERE NOT EXISTS ( - SELECT 1 FROM queue_claim_heads AS existing - WHERE existing.queue = lanes.queue - AND existing.priority = lanes.priority - AND existing.enqueue_shard = lanes.shard - ) - ORDER BY lanes.queue COLLATE "C", lanes.priority, lanes.shard - ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; - -- Re-enqueues that are not yet due carry no lane. INSERT INTO deferred_jobs ( job_id, kind, queue, args, state, priority, attempt, run_lease, @@ -441,10 +386,17 @@ BEGIN WHERE v_job_ids[input.ord] <> ALL (v_skipped) ORDER BY input.ord; - -- Lane lock window: reserve each lane's sequence range once, in - -- (queue, priority, shard) order, then write the ready rows. The locks are - -- held to the end of the transaction, which keeps the rows of a lane - -- visible in lane_seq order. + -- Lane lock window: in (queue, priority, shard) order, register each lane + -- the batch is the first to use and reserve its sequence range once, then + -- write the ready rows. The locks are held to the end of the transaction, + -- which keeps the rows of a lane visible in lane_seq order. + -- + -- A new lane is registered only when the loop reaches it, as every other + -- producer does: the head tables' row trigger provisions the lane under + -- its lock, and the uncommitted head row blocks other producers of that + -- lane, so registering a later lane before reserving an earlier one would + -- hold it out of order. Only missing rows are proposed, because the + -- trigger fires before the conflict check. FOR v_lane IN SELECT lanes.queue, lanes.priority, lanes.shard, lanes.row_count FROM ( @@ -458,6 +410,34 @@ BEGIN ) AS lanes ORDER BY lanes.queue COLLATE "C", lanes.priority, lanes.shard LOOP + INSERT INTO queue_lanes (queue, priority) + SELECT v_lane.queue, v_lane.priority + WHERE NOT EXISTS ( + SELECT 1 FROM queue_lanes AS existing + WHERE existing.queue = v_lane.queue AND existing.priority = v_lane.priority + ) + ON CONFLICT ON CONSTRAINT queue_lanes_pkey DO NOTHING; + + INSERT INTO queue_enqueue_heads (queue, priority, enqueue_shard) + SELECT v_lane.queue, v_lane.priority, v_lane.shard + WHERE NOT EXISTS ( + SELECT 1 FROM queue_enqueue_heads AS existing + WHERE existing.queue = v_lane.queue + AND existing.priority = v_lane.priority + AND existing.enqueue_shard = v_lane.shard + ) + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + + INSERT INTO queue_claim_heads (queue, priority, enqueue_shard) + SELECT v_lane.queue, v_lane.priority, v_lane.shard + WHERE NOT EXISTS ( + SELECT 1 FROM queue_claim_heads AS existing + WHERE existing.queue = v_lane.queue + AND existing.priority = v_lane.priority + AND existing.enqueue_shard = v_lane.shard + ) + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + EXECUTE format('SELECT %I.reserve_enqueue_seq($1, $2, $3, $4)', v_schema) INTO v_start USING v_lane.queue, v_lane.priority, v_lane.shard, v_lane.row_count::BIGINT; diff --git a/awa-model/src/insert.rs b/awa-model/src/insert.rs index 7b7811461..e8e38634e 100644 --- a/awa-model/src/insert.rs +++ b/awa-model/src/insert.rs @@ -20,9 +20,9 @@ use sqlx::{PgExecutor, PgPool}; // `insert_many_copy` and `insert_many_skip_unique_conflicts` use // `awa.insert_jobs_compat_batch` (the `insert_jobs_compat_batch` schema // patch), which does the lane-independent work first (job ids, unique claims, -// lane registration, payloads) and then reserves its lanes in canonical order, -// so a lane is held only across its reservation, the ready-row writes and the -// commit. +// payloads) and then registers and reserves its lanes one at a time in +// canonical order, so a lane is held only across its registration, +// reservation, the ready-row writes and the commit. // // Generic `insert_many` always uses the per-row path, as does // `insert_many_copy` on a schema without the patch: rows are sorted by the @@ -531,11 +531,16 @@ const INSERT_MANY_BATCH_SQL: &str = r#" ) "#; -/// TRUE when `awa.insert_jobs_compat_batch` exists, whoever may call it: the -/// `insert_jobs_compat_batch` schema patch's probe. Catalog lookup only. -pub(crate) const BATCH_FUNCTION_INSTALLED_SQL: &str = "SELECT to_regprocedure(\ +/// TRUE when `awa.insert_jobs_compat_batch` exists with the body the +/// `insert_jobs_compat_batch` schema patch installs, whoever may call it: the +/// patch's probe. The body check matches the per-lane registration inside the +/// reservation loop, so `awa migrate` replaces an earlier body that registered +/// every new lane before reserving any. Catalog lookup only. +pub(crate) const BATCH_FUNCTION_PATCH_APPLIED_SQL: &str = "SELECT EXISTS (\ + SELECT 1 FROM pg_proc AS p WHERE p.oid = to_regprocedure(\ 'awa.insert_jobs_compat_batch(text[],text[],jsonb[],text[],smallint[],smallint[],\ - timestamptz[],jsonb[],jsonb[],bytea[],text[],text[],boolean)') IS NOT NULL"; + timestamptz[],jsonb[],jsonb[],bytea[],text[],text[],boolean)') \ + AND p.prosrc ~ 'INSERT INTO queue_lanes \\(queue, priority\\)\\s+SELECT v_lane\\.queue')"; /// TRUE when the connected role may call `awa.insert_jobs_compat_batch`. const BATCH_FUNCTION_USABLE_SQL: &str = "SELECT COALESCE(has_function_privilege(to_regprocedure(\ diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index ab45ee325..f148325ec 100644 --- a/awa-model/src/migrations.rs +++ b/awa-model/src/migrations.rs @@ -431,7 +431,7 @@ pub const SCHEMA_PATCHES: &[SchemaPatch] = &[ description: "Batch producer insert with a short enqueue lane lock window (awa 0.7 migration v055)", sql: include_str!("../migrations/patches/insert_jobs_compat_batch.sql"), - applied_probe: crate::insert::BATCH_FUNCTION_INSTALLED_SQL, + applied_probe: crate::insert::BATCH_FUNCTION_PATCH_APPLIED_SQL, }, ]; diff --git a/awa/tests/fixtures/insert_jobs_compat_batch_registers_lanes_first.sql b/awa/tests/fixtures/insert_jobs_compat_batch_registers_lanes_first.sql new file mode 100644 index 000000000..d03069253 --- /dev/null +++ b/awa/tests/fixtures/insert_jobs_compat_batch_registers_lanes_first.sql @@ -0,0 +1,553 @@ +-- The first insert_jobs_compat_batch body, which registered every new lane +-- before reserving any. Fixture for the schema patch probe. + +CREATE OR REPLACE FUNCTION awa.insert_jobs_compat_batch( + p_kind TEXT[], + p_queue TEXT[], + p_args JSONB[], + p_state TEXT[], + p_priority SMALLINT[], + p_max_attempts SMALLINT[], + p_run_at TIMESTAMPTZ[], + p_metadata JSONB[], + p_tags JSONB[], + p_unique_key BYTEA[], + p_unique_states TEXT[], + p_ordering_hash TEXT[], + p_skip_unique_conflicts BOOLEAN DEFAULT FALSE +) +RETURNS SETOF awa.jobs +AS $$ +DECLARE + v_count INT := COALESCE(cardinality(p_kind), 0); + v_schema TEXT; + v_old_search_path TEXT; + v_now TIMESTAMPTZ := clock_timestamp(); + v_job_ids BIGINT[]; + v_skipped BIGINT[] := ARRAY[]::BIGINT[]; + v_lane RECORD; + v_lane_queues TEXT[] := ARRAY[]::TEXT[]; + v_lane_priorities SMALLINT[] := ARRAY[]::SMALLINT[]; + v_lane_shards SMALLINT[] := ARRAY[]::SMALLINT[]; + v_lane_starts BIGINT[] := ARRAY[]::BIGINT[]; + v_start BIGINT; + v_ready_slot INT; + v_ready_generation BIGINT; + v_row RECORD; + v_shards SMALLINT[]; +BEGIN + IF v_count = 0 THEN + RETURN; + END IF; + + IF cardinality(p_queue) IS DISTINCT FROM v_count + OR cardinality(p_args) IS DISTINCT FROM v_count + OR cardinality(p_state) IS DISTINCT FROM v_count + OR cardinality(p_priority) IS DISTINCT FROM v_count + OR cardinality(p_max_attempts) IS DISTINCT FROM v_count + OR cardinality(p_run_at) IS DISTINCT FROM v_count + OR cardinality(p_metadata) IS DISTINCT FROM v_count + OR cardinality(p_tags) IS DISTINCT FROM v_count + OR cardinality(p_unique_key) IS DISTINCT FROM v_count + OR cardinality(p_unique_states) IS DISTINCT FROM v_count + OR cardinality(p_ordering_hash) IS DISTINCT FROM v_count + THEN + RAISE EXCEPTION 'insert_jobs_compat_batch: column arrays must have equal length' + USING ERRCODE = '22023'; + END IF; + + IF EXISTS ( + SELECT 1 + FROM unnest(p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_metadata, p_tags) + AS input(kind, queue, args, state, priority, max_attempts, metadata, tags) + WHERE input.kind IS NULL OR input.queue IS NULL OR input.args IS NULL + OR input.state IS NULL OR input.priority IS NULL OR input.max_attempts IS NULL + OR input.metadata IS NULL OR input.tags IS NULL + ) THEN + RAISE EXCEPTION 'insert_jobs_compat_batch: kind, queue, args, state, priority, max_attempts, metadata and tags must not be NULL' + USING ERRCODE = '22004'; + END IF; + + IF EXISTS (SELECT 1 FROM unnest(p_kind) AS kind WHERE length(kind) > 200) THEN + RAISE EXCEPTION 'job kind length must be <= 200 characters' + USING ERRCODE = '23514'; + END IF; + + IF EXISTS (SELECT 1 FROM unnest(p_queue) AS queue WHERE length(queue) > 200) THEN + RAISE EXCEPTION 'queue name length must be <= 200 characters' + USING ERRCODE = '23514'; + END IF; + + IF EXISTS ( + SELECT 1 FROM unnest(p_priority) AS priority + WHERE priority < 1 OR priority > 4 + ) THEN + RAISE EXCEPTION 'priority must be between 1 and 4' + USING ERRCODE = '23514'; + END IF; + + IF EXISTS ( + SELECT 1 FROM unnest(p_max_attempts) AS max_attempts + WHERE max_attempts < 1 OR max_attempts > 1000 + ) THEN + RAISE EXCEPTION 'max_attempts must be between 1 and 1000' + USING ERRCODE = '23514'; + END IF; + + IF EXISTS ( + SELECT 1 FROM unnest(p_tags) AS tags + WHERE jsonb_array_length(tags) > 20 + ) THEN + RAISE EXCEPTION 'job tags must contain at most 20 values' + USING ERRCODE = '23514'; + END IF; + + v_schema := awa.active_queue_storage_schema(); + + -- Canonical storage has no lane lock to shorten: insert row by row. + IF v_schema IS NULL THEN + IF NOT p_skip_unique_conflicts THEN + RETURN QUERY + SELECT inserted.* + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, + p_run_at, p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + CROSS JOIN LATERAL awa.insert_job_compat( + input.kind, + input.queue, + input.args, + input.state::awa.job_state, + input.priority, + input.max_attempts, + input.run_at, + input.metadata, + ARRAY( + SELECT tag.value + FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, tag_ord) + ORDER BY tag.tag_ord + ), + input.unique_key, + input.unique_states::bit(8), + NULL + ) AS inserted + ORDER BY input.ord; + RETURN; + END IF; + + -- Skipping needs a subtransaction per row. + FOR v_row IN + SELECT + input.ord, + input.kind, + input.queue, + input.args, + input.state, + input.priority, + input.max_attempts, + input.run_at, + input.metadata, + ARRAY( + SELECT tag.value + FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, tag_ord) + ORDER BY tag.tag_ord + ) AS tags, + input.unique_key, + input.unique_states + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, + p_run_at, p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + ORDER BY input.ord + LOOP + BEGIN + RETURN NEXT awa.insert_job_compat( + v_row.kind, + v_row.queue, + v_row.args, + v_row.state::awa.job_state, + v_row.priority, + v_row.max_attempts, + v_row.run_at, + v_row.metadata, + v_row.tags, + v_row.unique_key, + v_row.unique_states::bit(8), + NULL + ); + EXCEPTION WHEN unique_violation THEN + NULL; + END; + END LOOP; + RETURN; + END IF; + + IF EXISTS ( + SELECT 1 FROM unnest(p_state) AS state + WHERE state NOT IN ('available', 'scheduled', 'retryable') + ) THEN + RAISE EXCEPTION 'queue storage does not support initial state %', + (SELECT state FROM unnest(p_state) AS state + WHERE state NOT IN ('available', 'scheduled', 'retryable') LIMIT 1) + USING ERRCODE = '22023'; + END IF; + + v_old_search_path := current_setting('search_path'); + PERFORM set_config('search_path', format('%I,awa,public', v_schema), true); + + -- Everything up to the lane reservation is lane-independent and runs + -- without a lane lock. + EXECUTE format( + 'SELECT array_agg(nextval(%L::regclass)::bigint) FROM generate_series(1, $1)', + format('%I.job_id_seq', v_schema) + ) + INTO v_job_ids + USING v_count; + + -- Rows that enqueue as `available` go to a lane. Their shard is derived as + -- insert_job_compat derives it: ordering-key rows from the key hash, the + -- rest from the backend pid. + SELECT array_agg( + CASE WHEN input.state = 'available' THEN + MOD( + CASE WHEN input.ordering_hash IS NULL + THEN pg_backend_pid()::numeric + ELSE input.ordering_hash::numeric + END, + COALESCE(meta.enqueue_shards, 1)::numeric + )::smallint + END + ORDER BY input.ord + ) + INTO v_shards + FROM unnest(p_queue, p_state, p_ordering_hash) + WITH ORDINALITY AS input(queue, state, ordering_hash, ord) + LEFT JOIN awa.queue_meta AS meta ON meta.queue = input.queue; + + -- Claims go in unique-key order, so two producers with overlapping keys + -- cannot wait on each other in opposite orders. + IF p_skip_unique_conflicts THEN + WITH claimed AS ( + INSERT INTO awa.job_unique_claims (unique_key, job_id) + SELECT input.unique_key, v_job_ids[input.ord] + FROM unnest(p_unique_key, p_unique_states, p_state) + WITH ORDINALITY AS input(unique_key, unique_states, state, ord) + WHERE input.unique_key IS NOT NULL + AND input.unique_states IS NOT NULL + AND awa.job_state_in_bitmask(input.unique_states::bit(8), input.state::awa.job_state) + ORDER BY input.unique_key, input.ord + ON CONFLICT (unique_key) DO NOTHING + RETURNING job_id + ) + SELECT COALESCE(array_agg(wanted.job_id), ARRAY[]::BIGINT[]) + INTO v_skipped + FROM ( + SELECT v_job_ids[input.ord] AS job_id + FROM unnest(p_unique_key, p_unique_states, p_state) + WITH ORDINALITY AS input(unique_key, unique_states, state, ord) + WHERE input.unique_key IS NOT NULL + AND input.unique_states IS NOT NULL + AND awa.job_state_in_bitmask(input.unique_states::bit(8), input.state::awa.job_state) + ) AS wanted + WHERE NOT EXISTS (SELECT 1 FROM claimed WHERE claimed.job_id = wanted.job_id); + ELSE + INSERT INTO awa.job_unique_claims (unique_key, job_id) + SELECT input.unique_key, v_job_ids[input.ord] + FROM unnest(p_unique_key, p_unique_states, p_state) + WITH ORDINALITY AS input(unique_key, unique_states, state, ord) + WHERE input.unique_key IS NOT NULL + AND input.unique_states IS NOT NULL + AND awa.job_state_in_bitmask(input.unique_states::bit(8), input.state::awa.job_state) + ORDER BY input.unique_key, input.ord; + END IF; + + -- Register lanes this batch is the first to use. The head tables carry a + -- row trigger that takes the lane lock before the conflict check, so only + -- missing lanes are proposed: a batch for an existing lane must not queue + -- behind that lane's producers before it has any work to show for it. + INSERT INTO queue_lanes (queue, priority) + SELECT lanes.queue, lanes.priority + FROM ( + SELECT DISTINCT input.queue, input.priority + FROM unnest(p_queue, p_priority, p_state) + WITH ORDINALITY AS input(queue, priority, state, ord) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS lanes + WHERE NOT EXISTS ( + SELECT 1 FROM queue_lanes AS existing + WHERE existing.queue = lanes.queue AND existing.priority = lanes.priority + ) + ORDER BY lanes.queue COLLATE "C", lanes.priority + ON CONFLICT ON CONSTRAINT queue_lanes_pkey DO NOTHING; + + INSERT INTO queue_enqueue_heads (queue, priority, enqueue_shard) + SELECT lanes.queue, lanes.priority, lanes.shard + FROM ( + SELECT DISTINCT input.queue, input.priority, v_shards[input.ord] AS shard + FROM unnest(p_queue, p_priority, p_state) + WITH ORDINALITY AS input(queue, priority, state, ord) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS lanes + WHERE NOT EXISTS ( + SELECT 1 FROM queue_enqueue_heads AS existing + WHERE existing.queue = lanes.queue + AND existing.priority = lanes.priority + AND existing.enqueue_shard = lanes.shard + ) + ORDER BY lanes.queue COLLATE "C", lanes.priority, lanes.shard + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + + INSERT INTO queue_claim_heads (queue, priority, enqueue_shard) + SELECT lanes.queue, lanes.priority, lanes.shard + FROM ( + SELECT DISTINCT input.queue, input.priority, v_shards[input.ord] AS shard + FROM unnest(p_queue, p_priority, p_state) + WITH ORDINALITY AS input(queue, priority, state, ord) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS lanes + WHERE NOT EXISTS ( + SELECT 1 FROM queue_claim_heads AS existing + WHERE existing.queue = lanes.queue + AND existing.priority = lanes.priority + AND existing.enqueue_shard = lanes.shard + ) + ORDER BY lanes.queue COLLATE "C", lanes.priority, lanes.shard + ON CONFLICT (queue, priority, enqueue_shard) DO NOTHING; + + -- Re-enqueues that are not yet due carry no lane. + INSERT INTO deferred_jobs ( + job_id, kind, queue, args, state, priority, attempt, run_lease, + max_attempts, run_at, attempted_at, finalized_at, created_at, + unique_key, unique_states, payload + ) + SELECT + v_job_ids[input.ord], + input.kind, + input.queue, + input.args, + input.state::awa.job_state, + input.priority, + 0, + 0, + input.max_attempts, + COALESCE(input.run_at, v_now), + NULL, + NULL, + v_now, + input.unique_key, + input.unique_states, + jsonb_build_object( + 'metadata', input.metadata, + 'tags', input.tags, + 'errors', '[]'::jsonb, + 'progress', NULL + ) + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_run_at, + p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + WHERE input.state <> 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped); + + -- The result rows do not depend on the lanes, so build them before the + -- lock window opens. + RETURN QUERY + SELECT + v_job_ids[input.ord], + input.kind, + input.queue, + input.args, + input.state::awa.job_state, + input.priority, + 0::smallint, + input.max_attempts, + COALESCE(input.run_at, v_now), + NULL::timestamptz, + NULL::timestamptz, + NULL::timestamptz, + NULL::timestamptz, + v_now, + ARRAY[]::jsonb[], + input.metadata, + ARRAY( + SELECT tag.value + FROM jsonb_array_elements_text(input.tags) WITH ORDINALITY AS tag(value, tag_ord) + ORDER BY tag.tag_ord + ), + input.unique_key, + input.unique_states::bit(8), + NULL::uuid, + NULL::timestamptz, + NULL::text, + NULL::text, + NULL::text, + NULL::text, + 0::bigint, + NULL::jsonb + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_run_at, + p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + WHERE v_job_ids[input.ord] <> ALL (v_skipped) + ORDER BY input.ord; + + -- Lane lock window: reserve each lane's sequence range once, in + -- (queue, priority, shard) order, then write the ready rows. The locks are + -- held to the end of the transaction, which keeps the rows of a lane + -- visible in lane_seq order. + FOR v_lane IN + SELECT lanes.queue, lanes.priority, lanes.shard, lanes.row_count + FROM ( + SELECT input.queue, input.priority, v_shards[input.ord] AS shard, + count(*) AS row_count + FROM unnest(p_queue, p_priority, p_state) + WITH ORDINALITY AS input(queue, priority, state, ord) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + GROUP BY input.queue, input.priority, v_shards[input.ord] + ) AS lanes + ORDER BY lanes.queue COLLATE "C", lanes.priority, lanes.shard + LOOP + EXECUTE format('SELECT %I.reserve_enqueue_seq($1, $2, $3, $4)', v_schema) + INTO v_start + USING v_lane.queue, v_lane.priority, v_lane.shard, v_lane.row_count::BIGINT; + + v_lane_queues := v_lane_queues || v_lane.queue; + v_lane_priorities := v_lane_priorities || v_lane.priority; + v_lane_shards := v_lane_shards || v_lane.shard; + v_lane_starts := v_lane_starts || v_start; + END LOOP; + + IF cardinality(v_lane_queues) > 0 THEN + IF to_regprocedure(format('%I.ring_cursor(text)', v_schema)) IS NOT NULL THEN + EXECUTE format('SELECT slot, generation FROM %I.ring_cursor(''queue'')', v_schema) + INTO v_ready_slot, v_ready_generation; + ELSE + SELECT ring.current_slot, ring.generation + INTO v_ready_slot, v_ready_generation + FROM queue_ring_state AS ring + WHERE ring.singleton = TRUE; + END IF; + + INSERT INTO ready_entries ( + ready_slot, ready_generation, job_id, kind, queue, args, priority, + attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, + attempted_at, created_at, unique_key, unique_states, payload + ) + SELECT + v_ready_slot, + v_ready_generation, + v_job_ids[placed.ord], + placed.kind, + placed.queue, + placed.args, + placed.priority, + 0, + 0, + placed.max_attempts, + lane.start_seq + placed.lane_rank - 1, + lane.shard, + COALESCE(placed.run_at, v_now), + NULL, + v_now, + placed.unique_key, + placed.unique_states, + jsonb_build_object( + 'metadata', placed.metadata, + 'tags', placed.tags, + 'errors', '[]'::jsonb, + 'progress', NULL + ) + FROM ( + SELECT + input.*, + v_shards[input.ord] AS shard, + row_number() OVER ( + PARTITION BY input.queue, input.priority, v_shards[input.ord] + ORDER BY input.ord + ) AS lane_rank + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_run_at, + p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS placed + JOIN unnest(v_lane_queues, v_lane_priorities, v_lane_shards, v_lane_starts) + AS lane(queue, priority, shard, start_seq) + ON lane.queue = placed.queue + AND lane.priority = placed.priority + AND lane.shard = placed.shard; + + -- One segment per run of consecutive lane_seq values sharing a run_at. + INSERT INTO ready_segments ( + ready_slot, ready_generation, queue, priority, enqueue_shard, + first_lane_seq, next_lane_seq, first_run_at + ) + SELECT + v_ready_slot, + v_ready_generation, + runs.queue, + runs.priority, + runs.shard, + min(runs.lane_seq), + max(runs.lane_seq) + 1, + runs.run_at + FROM ( + SELECT + placed.*, + placed.lane_seq - row_number() OVER ( + PARTITION BY placed.queue, placed.priority, placed.shard, placed.run_at + ORDER BY placed.lane_seq + ) AS run_id + FROM ( + SELECT + input.queue, + input.priority, + v_shards[input.ord] AS shard, + COALESCE(input.run_at, v_now) AS run_at, + lane.start_seq + row_number() OVER ( + PARTITION BY input.queue, input.priority, v_shards[input.ord] + ORDER BY input.ord + ) - 1 AS lane_seq + FROM unnest(p_queue, p_priority, p_state, p_run_at) + WITH ORDINALITY AS input(queue, priority, state, run_at, ord) + JOIN unnest(v_lane_queues, v_lane_priorities, v_lane_shards, v_lane_starts) + AS lane(queue, priority, shard, start_seq) + ON lane.queue = input.queue + AND lane.priority = input.priority + AND lane.shard = v_shards[input.ord] + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS placed + ) AS runs + GROUP BY runs.queue, runs.priority, runs.shard, runs.run_at, runs.run_id + ON CONFLICT (ready_slot, ready_generation, queue, priority, enqueue_shard, first_lane_seq) + DO NOTHING; + + PERFORM pg_notify('awa:' || notified.queue, '') + FROM (SELECT DISTINCT unnest(v_lane_queues) AS queue) AS notified; + END IF; + + PERFORM set_config('search_path', v_old_search_path, true); +END; +$$ LANGUAGE plpgsql VOLATILE +SET search_path = pg_catalog, awa, public; diff --git a/awa/tests/insert_many_batch_test.rs b/awa/tests/insert_many_batch_test.rs index 5609bd19e..63c36b676 100644 --- a/awa/tests/insert_many_batch_test.rs +++ b/awa/tests/insert_many_batch_test.rs @@ -1,11 +1,11 @@ //! Compatibility inserts and the batch insert function from the //! `insert_jobs_compat_batch` schema patch. //! -//! A batch does its lane-independent work (job ids, unique claims, lane -//! registration) before it reserves its lanes, and holds a lane lock only for -//! the reservation, the ready-row writes and the commit. These tests cover that -//! ordering, skip-on-conflict inserts, and the shape of what lands in the -//! lanes. +//! A batch does its lane-independent work (job ids, unique claims) before it +//! registers and reserves its lanes one at a time in canonical order, and +//! holds a lane lock only for the registration, the reservation, the ready-row +//! writes and the commit. These tests cover that ordering, skip-on-conflict +//! inserts, and the shape of what lands in the lanes. //! //! Requires a running Postgres instance. //! Run with: `cargo test --package awa --test insert_many_batch_test -- --nocapture` @@ -583,3 +583,75 @@ async fn insert_through_row_function(pool: &PgPool, jobs: &[InsertParams]) { .await .expect("per-row function insert"); } + +/// A per-row producer (as generic `insert_many` and released 0.6 binaries +/// insert) holds lane A and then enqueues into lane B, which nobody has used +/// yet; a batch for A and B is waiting for A meanwhile. The batch must not +/// have registered B before it holds A, or the two wait on each other. +#[tokio::test] +async fn batch_registers_a_new_lane_only_when_it_reaches_it() { + let _guard = ENGINE_GUARD.lock().await; + let pool = queue_storage_pool().await; + for path in [Path::SkipUniqueConflicts, Path::Copy] { + assert_batch_does_not_hold_a_new_lane_while_waiting(&pool, path).await; + } +} + +async fn assert_batch_does_not_hold_a_new_lane_while_waiting(pool: &PgPool, path: Path) { + let suffix = &Uuid::new_v4().simple().to_string()[..12]; + let lane_a = format!("batch_lane_a_{suffix}"); + let lane_b = format!("batch_lane_b_{suffix}"); + insert_many(pool, &[job(&lane_a, -1, false)]) + .await + .expect("provision lane A"); + + // The per-row producer's transaction holds lane A's lock from here on. + let mut per_row = pool.begin().await.expect("begin per-row producer"); + insert_many(&mut *per_row, &[job(&lane_a, 100, false)]) + .await + .expect("per-row insert into lane A"); + + let batch_jobs = vec![job(&lane_a, 0, false), job(&lane_b, 0, false)]; + let batch = { + let pool = pool.clone(); + tokio::spawn(async move { path.insert(&pool, &batch_jobs).await }) + }; + let key_a = lane_lock_key(pool, &lane_a).await; + let started = Instant::now(); + loop { + let waiting: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_locks + WHERE locktype = 'advisory' AND NOT granted + AND ((classid::bigint << 32) | objid::bigint) = $1)", + ) + .bind(key_a) + .fetch_one(pool) + .await + .expect("poll pg_locks"); + if waiting { + break; + } + assert!( + started.elapsed() < WAIT_DEADLINE, + "{path:?}: batch never waited for lane A" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } + + let second = tokio::time::timeout( + WAIT_DEADLINE, + insert_many(&mut *per_row, &[job(&lane_b, 100, false)]), + ) + .await + .unwrap_or_else(|_| panic!("{path:?}: per-row producer blocked on lane B")); + second.unwrap_or_else(|err| panic!("{path:?}: per-row insert into lane B: {err:?}")); + per_row.commit().await.expect("commit per-row producer"); + + let inserted = tokio::time::timeout(WAIT_DEADLINE, batch) + .await + .expect("batch finished") + .expect("batch task"); + assert_eq!(inserted.len(), 2, "{path:?}"); + assert_eq!(ready_lane_seqs(pool, &lane_a).await.len(), 3, "{path:?}"); + assert_eq!(ready_lane_seqs(pool, &lane_b).await.len(), 2, "{path:?}"); +} diff --git a/awa/tests/insert_many_schema_fallback_test.rs b/awa/tests/insert_many_schema_fallback_test.rs index a6b31ae29..275708b6d 100644 --- a/awa/tests/insert_many_schema_fallback_test.rs +++ b/awa/tests/insert_many_schema_fallback_test.rs @@ -394,35 +394,45 @@ async fn copy_without_execute_on_the_batch_function_stays_per_row() { .await .expect("create a role without EXECUTE on the batch function"); - let mut tx = pool.begin().await.expect("begin"); - sqlx::query(&format!("SET LOCAL ROLE {role}")) - .execute(&mut *tx) - .await - .expect("become the restricted role"); - let copied = insert_many_copy(&mut tx, &jobs("no_execute", 1..4, true)) - .await - .expect("COPY without EXECUTE on the batch function"); - assert_eq!(seqs(&copied), vec![1, 2, 3]); - let skipped = insert_many_skip_unique_conflicts(&mut *tx, &jobs("no_execute", 4..5, true)) - .await - .expect_err("the skip API calls the function directly"); - assert!( - matches!(&skipped, AwaError::Database(sqlx::Error::Database(e)) if e.code().as_deref() == Some("42501")), - "{skipped:?}" - ); - tx.rollback().await.expect("roll back"); + // The role is cluster-wide, so it is dropped even when an assertion fails. + let outcome = tokio::spawn({ + let pool = db.pool.clone(); + let role = role.clone(); + async move { + let pool = &pool; + let mut tx = pool.begin().await.expect("begin"); + sqlx::query(&format!("SET LOCAL ROLE {role}")) + .execute(&mut *tx) + .await + .expect("become the restricted role"); + let copied = insert_many_copy(&mut tx, &jobs("no_execute", 1..4, true)) + .await + .expect("COPY without EXECUTE on the batch function"); + assert_eq!(seqs(&copied), vec![1, 2, 3]); + let skipped = insert_many_skip_unique_conflicts(&mut *tx, &jobs("no_execute", 4..5, true)) + .await + .expect_err("the skip API calls the function directly"); + assert!( + matches!(&skipped, AwaError::Database(sqlx::Error::Database(e)) if e.code().as_deref() == Some("42501")), + "{skipped:?}" + ); + tx.rollback().await.expect("roll back"); + + let mut tx = pool.begin().await.expect("begin"); + sqlx::query(&format!("SET LOCAL ROLE {role}")) + .execute(&mut *tx) + .await + .expect("become the restricted role"); + insert_many_copy(&mut tx, &jobs("no_execute", 1..4, true)) + .await + .expect("COPY without EXECUTE on the batch function"); + tx.commit().await.expect("commit"); + assert_eq!(ready_seqs(pool, "no_execute").await, vec![0, 1, 2, 3]); + assert_eq!(segments(pool, "no_execute").await, 4, "per-row path"); - let mut tx = pool.begin().await.expect("begin"); - sqlx::query(&format!("SET LOCAL ROLE {role}")) - .execute(&mut *tx) - .await - .expect("become the restricted role"); - insert_many_copy(&mut tx, &jobs("no_execute", 1..4, true)) - .await - .expect("COPY without EXECUTE on the batch function"); - tx.commit().await.expect("commit"); - assert_eq!(ready_seqs(pool, "no_execute").await, vec![0, 1, 2, 3]); - assert_eq!(segments(pool, "no_execute").await, 4, "per-row path"); + } + }) + .await; db.drop().await; let admin = admin_pool().await; @@ -431,4 +441,33 @@ async fn copy_without_execute_on_the_batch_function_stays_per_row() { .await .expect("drop role"); admin.close().await; + if let Err(failure) = outcome { + std::panic::resume_unwind(failure.into_panic()); + } +} + +/// A body from before the per-lane registration fix keeps the patch pending, +/// and `awa migrate` replaces it. +#[tokio::test] +async fn an_earlier_batch_function_body_keeps_the_patch_pending() { + let db = ScratchDatabase::create().await; + let pool = &db.pool; + sqlx::raw_sql(include_str!( + "fixtures/insert_jobs_compat_batch_registers_lanes_first.sql" + )) + .execute(pool) + .await + .expect("install the earlier body"); + assert_eq!(pending_patches(pool).await, vec![PATCH_NAME]); + + let report = migrations::run_with_report(pool).await.expect("migrate"); + let patch = report + .patches + .iter() + .find(|patch| patch.name == PATCH_NAME) + .expect("patch report"); + assert_eq!(patch.outcome, migrations::PatchOutcome::Applied); + assert!(pending_patches(pool).await.is_empty()); + + db.drop().await; } diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 26586a1a3..69397beb5 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -458,7 +458,9 @@ awa --database-url "$DATABASE_URL" migrate Cause: - on queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane until the transaction commits; the `DETAIL` names two 64-bit advisory keys, and two transactions that took lanes in opposite orders hold one each and wait on the other's -- `insert_many`, `insert_many_copy`, `insert_many_skip_unique_conflicts` and the `QueueStorage` batch enqueues take their lane locks in one canonical order, `(queue, priority, shard)` ascending with queue names compared bytewise, so batches never deadlock with each other over lane locks; unique keys and the application's own locks can still form cycles +- `insert_many`, `insert_many_copy`, `insert_many_skip_unique_conflicts` and the `QueueStorage` batch enqueues register and lock their lanes one at a time in one canonical order, `(queue, priority, shard)` ascending with queue names compared bytewise, so two batch calls cannot hold each other's lanes in opposite orders +- with the `insert_jobs_compat_batch` schema patch, `insert_many_copy` and `insert_many_skip_unique_conflicts` take all of a batch's unique claims before any lane lock, while per-row inserts (generic `insert_many`, single-job inserts, 0.6.14 and earlier binaries) take each job's claim and lane together; a per-row transaction that holds a lane the batch needs and then inserts a job whose unique key the batch has claimed but not committed deadlocks with it, where it used to skip the job or fail with `23505` +- unique keys and the application's own locks can form other cycles - a transaction that enqueues with several separate statements (`insert` per job, handler-staged follow-ups, `Transaction.insert` in a loop) takes locks in statement order and can deadlock against a batch or another such transaction that touches the same lanes in a different order Fix: diff --git a/docs/upgrade-0.6.14-to-0.6.15.md b/docs/upgrade-0.6.14-to-0.6.15.md index 09851ae93..ff7f0d430 100644 --- a/docs/upgrade-0.6.14-to-0.6.15.md +++ b/docs/upgrade-0.6.14-to-0.6.15.md @@ -22,7 +22,7 @@ Steps 2 and 3 can also run the other way round. A 0.6.15 producer checks for the The schema stays at v40. The `insert_jobs_compat_batch` patch adds one function, `awa.insert_jobs_compat_batch`, and changes nothing else: no table DDL, no table or sequence locks, no change to `awa.insert_job_compat` or to storage routing. It applies as one transaction, so an interrupted run leaves nothing applied; run it again. -The function inserts a whole batch the way `awa.insert_job_compat` inserts one job, with the same validation, unique claims, lane assignment and storage routing. It does everything that does not depend on a lane first (job ids, unique claims in key order, new lane rows, scheduled rows, the result rows), then reserves each lane's range of positions once, in `(queue, priority, shard)` order, and writes the ready rows. Each lane's enqueue lock is therefore held only for the reservation, the ready-row write and the commit, instead of from the batch's first ready row to the commit. Concurrent producers and promotion on the same lane wait for that much shorter window. +The function inserts a whole batch the way `awa.insert_job_compat` inserts one job, with the same validation, unique claims, lane assignment and storage routing. It does everything that does not depend on a lane first (job ids, unique claims in key order, scheduled rows, the result rows). Then, in `(queue, priority, shard)` order, it registers each lane it is the first to use and reserves each lane's range of positions once, as every other producer registers and reserves lane by lane, and writes the ready rows. Each lane's enqueue lock is therefore held only for the reservation, the ready-row write and the commit, instead of from the batch's first ready row to the commit. Concurrent producers and promotion on the same lane wait for that much shorter window. What callers see: @@ -33,6 +33,10 @@ What callers see: 0.6.14 and earlier binaries never call the new function. They keep calling `awa.insert_job_compat`, which takes the same lane lock, so old and new producers share a lane without reordering it. +### Unique keys shared with per-row producers + +The batch function takes all of a batch's unique claims before any lane lock. Per-row inserts (generic `insert_many`, single-job inserts, and 0.6.14 and earlier binaries) take each job's claim and its lane together, lane by lane. When a per-row transaction already holds a lane the batch needs and then inserts a job whose unique key the batch has claimed but not yet committed, each waits on the other. PostgreSQL aborts one of them with `deadlock detected` (SQLSTATE `40P01`), where before the per-row insert would have waited and then skipped the job or failed with a unique violation (`23505`). Retry the aborted transaction: by then the key is either committed, giving the old outcome, or free. Two batches through the function take claims in key order and lanes in lane order, so they do not deadlock with each other this way. + ### Privileges New functions are executable by `PUBLIC` unless your deployment revokes that through default privileges. If it does and the runtime role's default privileges do not cover functions created by the migration role, grant it explicitly: From d300a4aa8b90be2dc4f392f58c0807ab3ef4b024 Mon Sep 17 00:00:00 2001 From: Brian Thorne Date: Mon, 9 Nov 2026 01:26:15 +1300 Subject: [PATCH 6/6] docs: note the new-lane deadlock between QueueStorage and per-row batches --- docs/troubleshooting.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 69397beb5..5dd8a9dc4 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -458,7 +458,7 @@ awa --database-url "$DATABASE_URL" migrate Cause: - on queue storage every enqueue holds a transaction-scoped advisory lock on its `(queue, priority, shard)` lane until the transaction commits; the `DETAIL` names two 64-bit advisory keys, and two transactions that took lanes in opposite orders hold one each and wait on the other's -- `insert_many`, `insert_many_copy`, `insert_many_skip_unique_conflicts` and the `QueueStorage` batch enqueues register and lock their lanes one at a time in one canonical order, `(queue, priority, shard)` ascending with queue names compared bytewise, so two batch calls cannot hold each other's lanes in opposite orders +- `insert_many`, `insert_many_copy`, `insert_many_skip_unique_conflicts` and the `QueueStorage` batch enqueues register and lock their lanes one at a time in one canonical order, `(queue, priority, shard)` ascending with queue names compared bytewise, so two batch calls cannot hold each other's lanes in opposite orders. One exception remains when a `QueueStorage` batch enqueue and an `insert_many` / `insert_many_copy` call both create the same new `(queue, priority)` lane: the two paths take that lane's provisioning lock and its `queue_lanes` row in opposite orders and can deadlock (#692); retry the transaction - with the `insert_jobs_compat_batch` schema patch, `insert_many_copy` and `insert_many_skip_unique_conflicts` take all of a batch's unique claims before any lane lock, while per-row inserts (generic `insert_many`, single-job inserts, 0.6.14 and earlier binaries) take each job's claim and lane together; a per-row transaction that holds a lane the batch needs and then inserts a job whose unique key the batch has claimed but not committed deadlocks with it, where it used to skip the job or fail with `23505` - unique keys and the application's own locks can form other cycles - a transaction that enqueues with several separate statements (`insert` per job, handler-staged follow-ups, `Transaction.insert` in a loop) takes locks in statement order and can deadlock against a batch or another such transaction that touches the same lanes in a different order