diff --git a/.github/workflows/nightly-chaos.yml b/.github/workflows/nightly-chaos.yml index 9c69c0889..a57f22aa2 100644 --- a/.github/workflows/nightly-chaos.yml +++ b/.github/workflows/nightly-chaos.yml @@ -197,6 +197,17 @@ jobs: cargo test --release --package awa --test concurrent_lifecycle_test \ test_concurrent_queue_sweep -- --exact --ignored --nocapture \ | tee artifacts/nightly/rust-queue-sweep.txt + - name: Run per-lane enqueue contention benchmark + if: always() + timeout-minutes: 5 + run: | + set -o pipefail + # Many producers batch-inserting into one queue lane. Gates the + # per-lane rate and p99 through benchmarks/baseline.json, and the + # lane-lock hold time relative to the per-row insert function. + cargo test --release --package awa --test enqueue_lane_contention_test \ + test_enqueue_lane_contention_gate -- --exact --ignored --nocapture \ + | tee artifacts/nightly/rust-enqueue-lane.txt - name: Run MVCC horizon overlap benchmark if: always() timeout-minutes: 10 diff --git a/CHANGELOG.md b/CHANGELOG.md index 4b2a82654..878971946 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,10 @@ Notable changes between releases. Detailed migration notes for storage transitio - Everything [0.6.14](https://github.com/hardbyte/awa/releases/tag/v0.6.14) shipped that this branch did not already have; see its release notes for the details. Promotion runs each enqueue lane in its own transaction, hands a busy lane to a dedicated waiter, and wakes claimers after the lanes commit; `awa.maintenance.promote.lane_contention` counts the hand-offs. Rescue skips an attempt whose retained ready bodies are ambiguous instead of stopping rescue of every other attempt. Lifecycle events no longer re-read a completed job, and a failed re-read no longer fails a committed callback wait. Admin job lists and kind counts find their candidates in read phases that release partition locks between them, and probe receipt closures more cheaply; a job list still loads each candidate with a point read, which holds `ACCESS SHARE` on every ready, lease and claim partition for the length of that read. `awa migrate` rejects `--sql` combined with `--extract-to`. The persona soak suite (`awa/tests/persona_soak_test.rs`, `scripts/persona-soak.sh`) is added. +### Forward-ported from 0.6.15 + +- Everything [0.6.15](https://github.com/hardbyte/awa/releases/tag/v0.6.15) shipped that this branch did not already have; see its release notes for the details. Lifecycle hooks hold their attempt's finalizer slot until they return, so hooks slower than job throughput slow claiming instead of growing worker memory, and maintenance keeps at most 1,024 detached `Rescued` and duplicate `Cancelled` hook tasks per worker. By-job lookups of compact receipt claims use a GIN index on `lease_claim_batches.job_ids` that claim prune builds on each partition as it empties it (migration v054). Compatibility COPY inserts go through `awa.insert_jobs_compat_batch` (migration v055), which takes each queue lane's enqueue lock only when it writes that lane rather than at the batch's first ready row (the lock is still held until the transaction ends), and the new `insert_many_skip_unique_conflicts` inserts a batch through it while skipping claimed unique keys. Rescue reaches stale attempts behind a full batch of ambiguous ones, `Completed` events and follow-ups carry the completed job without the attempt's heartbeat and deadline, slot-scoped admin receipt reads stay off other slots' closure partitions, and unfiltered receipt job lists select candidate IDs before loading bodies. v054 and v055 add functions only and take no table locks; a database coming from 0.6.15 already has both functions from its schema patches. See [v054](docs/upgrade-0.6-to-0.7.md#v054-claim-batch-job-index) and [v055](docs/upgrade-0.6-to-0.7.md#v055-batch-compatibility-insert) in the upgrade guide. + ### Added - **The queue-storage prune `lock_timeout` is configurable, and prune diff --git a/README.md b/README.md index 5c364f6e4..9bf9968ef 100644 --- a/README.md +++ b/README.md @@ -93,9 +93,9 @@ See [docs/positioning.md](docs/positioning.md) for the category map and messagin ```bash # 1. Add Awa to a Python project -uv add 'awa-pg[ui]==0.6.6' # Python SDK + dashboard binary -# uv add awa-pg==0.6.6 # SDK only (no dashboard, smaller wheel) -# or: cargo add awa@0.6.6 # Rust +uv add 'awa-pg[ui]' # Python SDK + dashboard binary +# uv add awa-pg # SDK only (no dashboard, smaller wheel) +# or: cargo add awa # Rust # 2. Start Postgres and run migrations awa --database-url $DATABASE_URL migrate @@ -295,7 +295,7 @@ tx.commit().await?; awa::admin::cancel_by_unique_key(&pool, "send_email", None, Some(&serde_json::json!({"to": "alice@example.com", "subject": "Welcome"})), None).await?; ``` -Lifecycle hooks receive `Started` after a claim commits as the worker handler is about to run, plus outcome events (`Completed`, `Retried`, `Exhausted`, `Cancelled`) after the corresponding state transition commits. Hooks run in detached tasks, so they do not block job execution. +Lifecycle hooks receive `Started` after a claim commits as the worker handler is about to run, plus outcome events (`Completed`, `Retried`, `Exhausted`, `Cancelled`) after the corresponding state transition commits. Hooks run alongside job execution and never delay a job's state transition, but each attempt's hooks hold a finalizer slot until they return, so hooks slower than job throughput slow claiming. Cancellation is cooperative for running handlers: @@ -310,19 +310,18 @@ Cancellation is cooperative for running handlers: ### Python ```bash -uv add awa-pg==0.6.6 # SDK: insert, worker, admin, progress -uv add 'awa-pg[ui]==0.6.6' # SDK + bundled `awa` binary for the dashboard +uv add awa-pg # SDK: insert, worker, admin, progress +uv add 'awa-pg[ui]' # SDK + bundled `awa` binary for the dashboard # or, just the CLI: -uv tool install awa-cli==0.6.6 # CLI on its own: migrations, queue admin, web UI +uv tool install awa-cli # CLI on its own: migrations, queue admin, web UI ``` -`uv add awa-pg==0.6.6` stays small for workers and producers. The `[ui]` extra pulls in [`awa-cli`](https://pypi.org/project/awa-cli/), which ships the `awa` binary plus the embedded React dashboard; afterwards `uv run python -m awa serve` (or `awa serve` directly) launches it. +`uv add awa-pg` stays small for workers and producers. The `[ui]` extra pulls in [`awa-cli`](https://pypi.org/project/awa-cli/), which ships the `awa` binary plus the embedded React dashboard; afterwards `uv run python -m awa serve` (or `awa serve` directly) launches it. ### Rust -```toml -[dependencies] -awa = "0.6" +```bash +cargo add awa ``` ### CLI @@ -330,7 +329,7 @@ awa = "0.6" Available as a uv tool (no Rust toolchain needed) or through cargo: ```bash -uv tool install awa-cli==0.6.6 +uv tool install awa-cli # or: cargo install awa-cli awa --database-url $DATABASE_URL migrate @@ -376,7 +375,7 @@ All coordination through Postgres. The Rust runtime owns dispatch, leases, heart | `awa-worker` | Runtime: dispatch, heartbeat, maintenance | | `awa-ui` | Web UI (axum API + embedded React frontend) | | `awa-cli` | CLI binary (migrations, admin, serve) | -| `awa-python` | PyO3 extension module (`uv add awa-pg==0.6.6`) | +| `awa-python` | PyO3 extension module (`uv add awa-pg`) | | `awa-testing` | Test helpers (`TestClient`) | ## Documentation diff --git a/awa-model/migrations/v023_install_queue_storage_substrate.sql b/awa-model/migrations/v023_install_queue_storage_substrate.sql index 85e496f31..a0d8448db 100644 --- a/awa-model/migrations/v023_install_queue_storage_substrate.sql +++ b/awa-model/migrations/v023_install_queue_storage_substrate.sql @@ -792,6 +792,177 @@ $job_output$; REVOKE EXECUTE ON FUNCTION awa.install_queue_storage_job_output(TEXT) FROM PUBLIC; +-- Installs p_schema.ensure_claim_batch_job_index(p_slot), which adds the GIN +-- index on job_ids to one lease_claim_batches partition when that partition +-- is empty and its lock is free, and otherwise does nothing. Building the +-- index on a populated partition holds SHARE, which blocks claim inserts, for +-- the whole build (minutes per million batches), so the index is only ever +-- created on an empty partition: by the installer, or by claim prune right +-- after it truncates a slot. +-- +-- fastupdate is off: job ids arrive in ascending order, so direct inserts +-- append to the right edge of the entry tree and cost no more than the +-- pending list does, while a pending list would make every lookup read it +-- (up to gin_pending_list_limit) and push its flush into a claim insert. +-- +-- The function is SECURITY DEFINER, owned by the owner of lease_claim_batches +-- (CREATE INDEX needs table ownership), with a pinned search_path. EXECUTE +-- goes to the roles that may TRUNCATE lease_claim_batches, which are the roles +-- that run claim prune, rather than to PUBLIC. Replacing it takes no table +-- locks and reads no rows. The installer calls this too. +CREATE OR REPLACE FUNCTION awa.refresh_claim_batch_job_index_helper(p_schema TEXT) +RETURNS VOID +LANGUAGE plpgsql +SECURITY INVOKER +SET search_path = pg_catalog, pg_temp +AS $refresh$ +DECLARE + v_table_owner OID; + v_table_acl ACLITEM[]; + v_function REGPROCEDURE; + v_function_owner OID; + v_grantee TEXT; + v_set_mode TEXT := CASE WHEN current_setting('server_version_num')::int >= 160000 + THEN 'SET' ELSE 'MEMBER' END; +BEGIN + IF p_schema IS NULL OR p_schema !~ '^[a-z_][a-z0-9_]*$' THEN + RAISE EXCEPTION 'refresh_claim_batch_job_index_helper: schema name must match ^[a-z_][a-z0-9_]*$; got %', + p_schema + USING ERRCODE = '22023'; + END IF; + + SELECT batches.relowner, batches.relacl INTO v_table_owner, v_table_acl + FROM pg_class AS batches + WHERE batches.oid = to_regclass(format('%I.lease_claim_batches', p_schema)); + IF v_table_owner IS NULL THEN + RAISE EXCEPTION 'refresh_claim_batch_job_index_helper: schema % has no lease_claim_batches table', + p_schema + USING ERRCODE = '22023'; + END IF; + + EXECUTE format( + $ddl$ + CREATE OR REPLACE FUNCTION %1$I.ensure_claim_batch_job_index(p_slot INTEGER) + RETURNS BOOLEAN + LANGUAGE plpgsql + SECURITY DEFINER + SET search_path = pg_catalog, pg_temp + SET quote_all_identifiers = off + AS $func$ + DECLARE + v_child TEXT; + v_index TEXT; + v_has_rows BOOLEAN; + BEGIN + IF p_slot IS NULL OR p_slot < 0 THEN + RAISE EXCEPTION 'invalid claim slot: %%', p_slot + USING ERRCODE = '22023'; + END IF; + + v_child := format('lease_claim_batches_%%s', p_slot); + v_index := format('lease_claim_batches_%%s_job_ids_idx', p_slot); + + IF NOT EXISTS ( + SELECT 1 + FROM pg_inherits AS inheritance + WHERE inheritance.inhparent = to_regclass(format('%%I.lease_claim_batches', %1$L)) + AND inheritance.inhrelid = to_regclass(format('%%I.%%I', %1$L, v_child)) + ) THEN + RAISE EXCEPTION 'schema %% has no lease_claim_batches partition for claim slot %%', + %1$L, p_slot + USING ERRCODE = '22023'; + END IF; + + IF EXISTS ( + SELECT 1 + FROM pg_index AS idx + JOIN pg_class AS index_class ON index_class.oid = idx.indexrelid + JOIN pg_am AS method ON method.oid = index_class.relam + WHERE idx.indrelid = to_regclass(format('%%I.%%I', %1$L, v_child)) + AND method.amname = 'gin' + AND idx.indisvalid + AND pg_get_indexdef(idx.indexrelid) LIKE '%%USING gin (job_ids)%%' + ) THEN + RETURN TRUE; + END IF; + + -- NOWAIT: a partition someone is writing to is not empty for long + -- enough to index, so give up rather than queue behind claims. + -- Leaving the block through the exception rolls back its + -- subtransaction, which releases a lock taken here on a populated + -- partition; a lock the caller already held is kept. + BEGIN + EXECUTE format('LOCK TABLE %%I.%%I IN SHARE MODE NOWAIT', %1$L, v_child); + EXECUTE format('SELECT EXISTS (SELECT 1 FROM %%I.%%I)', %1$L, v_child) + INTO v_has_rows; + IF v_has_rows THEN + RAISE EXCEPTION USING ERRCODE = 'object_not_in_prerequisite_state'; + END IF; + EXECUTE format( + 'CREATE INDEX IF NOT EXISTS %%I ON %%I.%%I USING GIN (job_ids) WITH (fastupdate = off)', + v_index, %1$L, v_child + ); + EXCEPTION + WHEN lock_not_available OR object_not_in_prerequisite_state THEN + RETURN FALSE; + END; + RETURN TRUE; + END; + $func$ + $ddl$, + p_schema + ); + + v_function := format('%I.ensure_claim_batch_job_index(integer)', p_schema)::regprocedure; + + SELECT routine.proowner INTO v_function_owner + FROM pg_proc AS routine + WHERE routine.oid = v_function; + + IF v_function_owner <> v_table_owner THEN + IF pg_has_role(current_user, v_function_owner, 'USAGE') + AND pg_has_role(current_user, v_table_owner, v_set_mode) + AND has_schema_privilege(v_table_owner, to_regnamespace(quote_ident(p_schema)), 'CREATE') THEN + EXECUTE format( + 'ALTER FUNCTION %s OWNER TO %I', + v_function, + pg_get_userbyid(v_table_owner) + ); + ELSE + RAISE WARNING '% is owned by %, not the owner % of %.lease_claim_batches, so it cannot create the job_ids index', + v_function, pg_get_userbyid(v_function_owner), pg_get_userbyid(v_table_owner), p_schema + USING HINT = format( + 'Run awa migrate as a role that can SET ROLE to %s, or as %s itself.', + quote_ident(pg_get_userbyid(v_table_owner)), + quote_ident(pg_get_userbyid(v_table_owner)) + ); + END IF; + END IF; + + FOR v_grantee IN + SELECT DISTINCT CASE WHEN acl.grantee = 0 THEN 'PUBLIC' + ELSE quote_ident(pg_get_userbyid(acl.grantee)) END + FROM pg_proc AS routine + CROSS JOIN LATERAL aclexplode(COALESCE(routine.proacl, acldefault('f', routine.proowner))) AS acl + WHERE routine.oid = v_function + AND acl.privilege_type = 'EXECUTE' + AND acl.grantee <> routine.proowner + LOOP + EXECUTE format('REVOKE EXECUTE ON FUNCTION %s FROM %s', v_function, v_grantee); + END LOOP; + FOR v_grantee IN + SELECT DISTINCT CASE WHEN acl.grantee = 0 THEN 'PUBLIC' + ELSE quote_ident(pg_get_userbyid(acl.grantee)) END + FROM aclexplode(COALESCE(v_table_acl, acldefault('r', v_table_owner))) AS acl + WHERE acl.privilege_type = 'TRUNCATE' + LOOP + EXECUTE format('GRANT EXECUTE ON FUNCTION %s TO %s', v_function, v_grantee); + END LOOP; +END; +$refresh$; + +REVOKE EXECUTE ON FUNCTION awa.refresh_claim_batch_job_index_helper(TEXT) FROM PUBLIC; + CREATE OR REPLACE FUNCTION awa.install_queue_storage_substrate( p_schema TEXT, p_queue_slot_count INT DEFAULT 16, @@ -2535,6 +2706,19 @@ BEGIN ); END LOOP; + PERFORM awa.refresh_claim_batch_job_index_helper(p_schema); + -- The by-job lookups of compact claims use a GIN index that is only built + -- on an empty partition whose lock is free (see + -- awa.refresh_claim_batch_job_index_helper). A partition that is populated + -- or busy now gains it after claim prune next empties it. + FOR v_slot IN 0..(p_claim_slot_count - 1) LOOP + BEGIN + EXECUTE format('SELECT %I.ensure_claim_batch_job_index($1)', p_schema) USING v_slot; + EXCEPTION WHEN insufficient_privilege THEN + RAISE NOTICE 'claim slot % of schema % keeps scanning for job ids: %', v_slot, p_schema, SQLERRM; + END; + END LOOP; + EXECUTE format( 'CREATE INDEX IF NOT EXISTS idx_%s_lease_claim_closures_job_run ON %I.lease_claim_closures (job_id, run_lease)', p_schema, p_schema diff --git a/awa-model/migrations/v054_claim_batch_job_index.sql b/awa-model/migrations/v054_claim_batch_job_index.sql new file mode 100644 index 000000000..b9544ed09 --- /dev/null +++ b/awa-model/migrations/v054_claim_batch_job_index.sql @@ -0,0 +1,56 @@ +-- v054: Index compact claim batches by job id. +-- Tracking: https://github.com/hardbyte/awa/pull/684 +-- Point lookups of compact receipt claims by job id (heartbeat and progress +-- flush, completion fallbacks, receipt close, cancellation, rescue and +-- load_job) filter lease_claim_batches with job_ids && / @>, which a GIN index +-- on job_ids serves. The installer now installs +-- .ensure_claim_batch_job_index(slot) through +-- awa.refresh_claim_batch_job_index_helper and builds the index on each empty +-- claim partition; claim prune calls the function right after it truncates a +-- slot, so every partition gains the index within one claim-ring cycle. +-- +-- This file only installs the function in each queue-storage schema; it +-- builds no index. `awa migrate` and `awa migrate --sql --from 53 --to 54` run +-- V23_UP first, which defines the helper. +-- +-- A 0.6.15 database already carries the function from the 0.6 +-- claim_batch_job_index schema patch; replacing it is idempotent. +-- +-- N-1 binaries never call the function, and their by-job lookups return the +-- same rows with or without the index. Binaries that carry the indexed +-- lookups stay correct on an unindexed partition; they scan the claim batches. +-- +-- Live load: function replacement only. It takes no lock on any awa or +-- queue-storage table. Run it as the owner of the queue-storage tables, or as +-- a migrator that can SET ROLE to it, so the function is owned by the table +-- owner and can create the index. + +DO $$ +DECLARE + v_schema TEXT; +BEGIN + FOR v_schema IN + SELECT discovered.schema_name FROM awa.queue_storage_schemas() AS discovered + LOOP + IF NOT EXISTS ( + SELECT 1 FROM pg_class AS batches + WHERE batches.oid = to_regclass(format('%I.lease_claim_batches', v_schema)) + AND batches.relkind = 'p' + ) THEN + CONTINUE; + END IF; + + -- Serialize with a concurrent prepare_schema() of the same schema. This + -- is an advisory lock, not a table lock: claims and completions do not + -- take it. + PERFORM pg_advisory_xact_lock( + hashtextextended(format('awa.queue_storage.install:%I', v_schema), 0) + ); + PERFORM awa.refresh_claim_batch_job_index_helper(v_schema); + END LOOP; +END; +$$; + +INSERT INTO awa.schema_version (version, description) +VALUES (54, 'Index compact claim batches by job id') +ON CONFLICT (version) DO NOTHING; diff --git a/awa-model/migrations/v055_insert_jobs_compat_batch.sql b/awa-model/migrations/v055_insert_jobs_compat_batch.sql new file mode 100644 index 000000000..0f290ef35 --- /dev/null +++ b/awa-model/migrations/v055_insert_jobs_compat_batch.sql @@ -0,0 +1,567 @@ +-- v055: a batch producer insert that holds each enqueue lane lock only for the +-- lane write. +-- Tracking: https://github.com/hardbyte/awa/pull/690 +-- +-- 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. +-- awa.insert_jobs_compat_batch does the lane-independent 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; +-- 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. +-- +-- A 0.6.15 database already carries this function from the 0.6 +-- insert_jobs_compat_batch schema patch; this replaces it in place. +-- +-- N-1 binaries never call the function: they keep calling +-- awa.insert_job_compat, which this migration does not touch, and both take +-- the same lane lock, so producers on either path share a lane in lane_seq +-- order. insert_many_copy probes its own connection on every call and uses +-- the per-row function until this function is present; generic insert_many +-- always uses the per-row function; insert_many_skip_unique_conflicts requires +-- this function. +-- +-- `awa migrate` and `awa migrate --sql --from 54 --to 55` run V23_UP first, +-- which bootstraps trusted queue-storage discovery and refuses an untrusted +-- registered substrate, as every migration since v049 does. +-- +-- Live load: function replacement only, no table DDL, no table locks. + +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; + + -- 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: 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 ( + 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 + 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; + + v_lane_queues := v_lane_queues || v_lane.queue; + v_lane_priorities := v_lane_priorities || v_lane.priority; + v_lane_shards := v_lane_shards || v_lane.shard; + v_lane_starts := v_lane_starts || v_start; + END LOOP; + + IF cardinality(v_lane_queues) > 0 THEN + SELECT rc.slot, rc.generation + INTO v_ready_slot, v_ready_generation + FROM ring_cursor('queue') AS rc; + + INSERT INTO ready_entries ( + ready_slot, ready_generation, job_id, kind, queue, args, priority, + attempt, run_lease, max_attempts, lane_seq, enqueue_shard, run_at, + attempted_at, created_at, unique_key, unique_states, payload + ) + SELECT + v_ready_slot, + v_ready_generation, + v_job_ids[placed.ord], + placed.kind, + placed.queue, + placed.args, + placed.priority, + 0, + 0, + placed.max_attempts, + lane.start_seq + placed.lane_rank - 1, + lane.shard, + COALESCE(placed.run_at, v_now), + NULL, + v_now, + placed.unique_key, + placed.unique_states, + jsonb_build_object( + 'metadata', placed.metadata, + 'tags', placed.tags, + 'errors', '[]'::jsonb, + 'progress', NULL + ) + FROM ( + SELECT + input.*, + v_shards[input.ord] AS shard, + row_number() OVER ( + PARTITION BY input.queue, input.priority, v_shards[input.ord] + ORDER BY input.ord + ) AS lane_rank + FROM unnest( + p_kind, p_queue, p_args, p_state, p_priority, p_max_attempts, p_run_at, + p_metadata, p_tags, p_unique_key, p_unique_states + ) WITH ORDINALITY AS input( + kind, queue, args, state, priority, max_attempts, run_at, + metadata, tags, unique_key, unique_states, ord + ) + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS placed + JOIN unnest(v_lane_queues, v_lane_priorities, v_lane_shards, v_lane_starts) + AS lane(queue, priority, shard, start_seq) + ON lane.queue = placed.queue + AND lane.priority = placed.priority + AND lane.shard = placed.shard; + + -- One segment per run of consecutive lane_seq values sharing a run_at. + INSERT INTO ready_segments ( + ready_slot, ready_generation, queue, priority, enqueue_shard, + first_lane_seq, next_lane_seq, first_run_at + ) + SELECT + v_ready_slot, + v_ready_generation, + runs.queue, + runs.priority, + runs.shard, + min(runs.lane_seq), + max(runs.lane_seq) + 1, + runs.run_at + FROM ( + SELECT + placed.*, + placed.lane_seq - row_number() OVER ( + PARTITION BY placed.queue, placed.priority, placed.shard, placed.run_at + ORDER BY placed.lane_seq + ) AS run_id + FROM ( + SELECT + input.queue, + input.priority, + v_shards[input.ord] AS shard, + COALESCE(input.run_at, v_now) AS run_at, + lane.start_seq + row_number() OVER ( + PARTITION BY input.queue, input.priority, v_shards[input.ord] + ORDER BY input.ord + ) - 1 AS lane_seq + FROM unnest(p_queue, p_priority, p_state, p_run_at) + WITH ORDINALITY AS input(queue, priority, state, run_at, ord) + JOIN unnest(v_lane_queues, v_lane_priorities, v_lane_shards, v_lane_starts) + AS lane(queue, priority, shard, start_seq) + ON lane.queue = input.queue + AND lane.priority = input.priority + AND lane.shard = v_shards[input.ord] + WHERE input.state = 'available' + AND v_job_ids[input.ord] <> ALL (v_skipped) + ) AS placed + ) AS runs + GROUP BY runs.queue, runs.priority, runs.shard, runs.run_at, runs.run_id + ON CONFLICT (ready_slot, ready_generation, queue, priority, enqueue_shard, first_lane_seq) + DO NOTHING; + + PERFORM pg_notify('awa:' || notified.queue, '') + FROM (SELECT DISTINCT unnest(v_lane_queues) AS queue) AS notified; + END IF; + + PERFORM set_config('search_path', v_old_search_path, true); +END; +$$ LANGUAGE plpgsql VOLATILE +SET search_path = pg_catalog, awa, public; + +INSERT INTO awa.schema_version (version, description) +VALUES (55, 'Batch producer insert with a short enqueue lane lock window') +ON CONFLICT (version) DO NOTHING; diff --git a/awa-model/src/admin.rs b/awa-model/src/admin.rs index 9dbb5825f..3923c87df 100644 --- a/awa-model/src/admin.rs +++ b/awa-model/src/admin.rs @@ -1022,6 +1022,20 @@ fn queue_storage_job_id_candidates_sql( if let JobCandidateSource::Receipts(slot) = source { let open_running = crate::queue_storage::open_receipt_running_claims_sql_in_slot(schema, slot); + let kind_join = if kind.is_some() { + format!( + r#" + CROSS JOIN LATERAL ( + SELECT ready.kind + FROM {schema}.ready_entries AS ready + WHERE ready.job_id = open_claims.job_id + LIMIT 1 + ) AS created + "# + ) + } else { + String::new() + }; let predicates = predicates( None, Some("created.kind"), @@ -1032,12 +1046,7 @@ fn queue_storage_job_id_candidates_sql( r#"( SELECT DISTINCT open_claims.job_id FROM ({open_running}) AS open_claims - CROSS JOIN LATERAL ( - SELECT ready.kind - FROM {schema}.ready_entries AS ready - WHERE ready.job_id = open_claims.job_id - LIMIT 1 - ) AS created + {kind_join} WHERE TRUE{predicates} ORDER BY open_claims.job_id DESC LIMIT {limit} diff --git a/awa-model/src/insert.rs b/awa-model/src/insert.rs index e221b0fa7..d628e4e66 100644 --- a/awa-model/src/insert.rs +++ b/awa-model/src/insert.rs @@ -14,10 +14,20 @@ 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` (migration v055), which does the lane-independent work first (job ids, unique claims, +// 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 before v055: 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 +223,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 `awa.insert_jobs_compat_batch`, 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.*, @@ -476,7 +488,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 +/// `awa.insert_jobs_compat_batch` (migration v055). [`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( @@ -497,7 +519,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() @@ -541,6 +563,136 @@ 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 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 `awa.insert_jobs_compat_batch` (migration v055, `awa migrate`); on +/// a schema without it, returns [`AwaError::Validation`]. [`insert_many_copy`] +/// skips unique conflicts with or without it, through a connection, and is +/// what Python's `insert_many_copy` calls. +#[tracing::instrument( + skip(executor, jobs), + fields( + job.count = jobs.len(), + otel.name = "send", + otel.kind = "producer", + messaging.system = "awa", + messaging.operation.r#type = "send", + messaging.batch.message_count = jobs.len(), + ) +)] +pub async fn insert_many_skip_unique_conflicts<'e, E>( + executor: E, + jobs: &[InsertParams], +) -> Result, AwaError> +where + E: PgExecutor<'e>, +{ + if jobs.is_empty() { + return Ok(Vec::new()); + } + let rows = precompute_rows(jobs)?; + 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 awa.insert_jobs_compat_batch (migration v055); \ + run `awa migrate`" + .to_string(), + )), + Err(err) => Err(err.into()), + } +} + +/// `undefined_function` naming the batch function: the schema predates v055, +/// or the function was dropped. +fn is_missing_batch_function(err: &sqlx::Error) -> bool { + matches!( + err, + sqlx::Error::Database(db_err) + if db_err.code().as_deref() == Some("42883") + && db_err.message().contains("insert_jobs_compat_batch") + ) +} + +fn bind_batch_columns<'q>( + 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`. @@ -574,7 +726,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 before migration v055, 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) @@ -646,10 +801,89 @@ pub async fn insert_many_copy( copy_in.finish().await?; } - // 3. INSERT...SELECT from staging into real table + // 3. Insert the staged rows. With the batch function, unique conflicts are + // skipped in the one statement; without it, each unique row goes through + // the per-row function in its own savepoint. let has_unique = rows.iter().any(|r| r.unique_key.is_some()); + let results = if batch_function_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?; - let results = if has_unique { + 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 function, 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?) +} + +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 @@ -790,13 +1024,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 1d745187d..fe279c4a9 100644 --- a/awa-model/src/lib.rs +++ b/awa-model/src/lib.rs @@ -46,7 +46,8 @@ pub use cron::{CronJobRow, CronMissedFirePolicy, PeriodicJob, PeriodicJobBuilder pub use dlq::{DlqMetadata, DlqRow, ListDlqFilter, RetryFromDlqOpts}; pub use error::{map_sqlx_error, AwaError}; pub use insert::{ - insert, insert_many, insert_many_copy, insert_many_copy_from_pool, insert_prepared, insert_with, + insert, insert_many, insert_many_copy, insert_many_copy_from_pool, + insert_many_skip_unique_conflicts, insert_prepared, insert_with, }; pub use job::{InsertOpts, InsertParams, JobRow, JobState, UniqueOpts}; pub use output::{AttemptError, JobOutput, OutputDescriptor, OutputError, MAX_OUTPUT_BYTES}; diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index f30033198..9a7365a34 100644 --- a/awa-model/src/migrations.rs +++ b/awa-model/src/migrations.rs @@ -5,7 +5,7 @@ use sqlx::{Connection, PgPool}; use tracing::{info, warn}; /// Current schema version. -pub const CURRENT_VERSION: i32 = 53; +pub const CURRENT_VERSION: i32 = 55; /// Migrations that require an exclusive (no-live-runtime) upgrade window. /// @@ -297,6 +297,16 @@ const MIGRATIONS: &[(i32, &str, &[&str])] = &[ "Create queue lanes as the schema owner", &[V23_UP, V53_UP], ), + ( + 54, + "Index compact claim batches by job id", + &[V23_UP, V54_UP], + ), + ( + 55, + "Batch producer insert with a short enqueue lane lock window", + &[V23_UP, V55_UP], + ), ]; const V1_UP: &str = include_str!("../migrations/v001_canonical_schema.sql"); @@ -358,6 +368,8 @@ const V50_UP: &str = include_str!("../migrations/v050_drop_ready_lane_shard_inde const V51_UP: &str = include_str!("../migrations/v051_claim_slot_functions.sql"); const V52_UP: &str = include_str!("../migrations/v052_retained_job_output.sql"); const V53_UP: &str = include_str!("../migrations/v053_definer_lane_provisioning.sql"); +const V54_UP: &str = include_str!("../migrations/v054_claim_batch_job_index.sql"); +const V55_UP: &str = include_str!("../migrations/v055_insert_jobs_compat_batch.sql"); /// Old version numbers from pre-0.4 releases that used V3/V4/V5 numbering. /// Also tolerates the unreleased inline-V6 branch numbering used during review. diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index 9b0a0813d..1c93ab888 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -1473,9 +1473,10 @@ pub fn open_receipt_running_claims_sql(schema: &str) -> String { /// [`open_receipt_running_claims_sql`] restricted to the jobs in `job_ids`, /// an SQL expression yielding `bigint[]`. The restriction applies before the /// closure and materialisation checks. Row-local claims are found through the -/// `job_id` index. Compact claim batches have no index on `job_ids`, so every -/// retained batch header is read, but only batches holding one of the jobs -/// are unnested and checked. +/// `job_id` index. Compact claim batches are found through `job_ids && {ids}`, +/// which uses a claim partition's GIN index on `job_ids` where it has one +/// (see `ensure_claim_batch_job_index_tx`) and otherwise reads every retained +/// batch header; only batches holding one of the jobs are unnested and checked. pub(crate) fn open_receipt_running_claims_sql_for(schema: &str, job_ids: Option<&str>) -> String { open_receipt_running_claims_sql_from(schema, job_ids, None) } @@ -1542,20 +1543,20 @@ fn open_receipt_running_claims_sql_from( WHERE {batch_filter}NOT ( batches.receipt_ranges <@ COALESCE(( SELECT range_agg(closed) - FROM {schema}.lease_claim_closure_batches AS closure_batches + FROM {closure_batch_rel} AS closure_batches CROSS JOIN LATERAL unnest(closure_batches.receipt_ranges) AS closed WHERE closure_batches.claim_slot = batches.claim_slot AND closure_batches.receipt_ranges && batches.receipt_ranges ), '{{}}'::int8multirange) ) AND NOT EXISTS ( - SELECT 1 FROM {schema}.lease_claim_closures AS closures + SELECT 1 FROM {closure_rel} AS closures WHERE closures.claim_slot = batches.claim_slot AND closures.job_id = items.job_id AND closures.run_lease = items.run_lease ) AND NOT EXISTS ( - SELECT 1 FROM {schema}.lease_claim_closure_batches AS closure_batches + SELECT 1 FROM {closure_batch_rel} AS closure_batches WHERE closure_batches.claim_slot = batches.claim_slot AND closure_batches.receipt_ranges @> items.receipt_id -- Keep the range lookup correlated: joining on claim_slot can @@ -2710,25 +2711,49 @@ struct ReadyRowIdentity { } /// Attempts a rescue pass leaves for a later pass, as `(job_id, run_lease)`. -struct BusyRetained(Vec<(i64, i64)>); +struct BusyRetained { + attempts: Vec<(i64, i64)>, + /// The attempts stay unresolvable until their retained bodies are + /// repaired, rather than until a lock is released. + unresolvable: bool, +} impl BusyRetained { - fn exclude(self, excluded: &mut BTreeSet<(i64, i64)>) -> Result<(), AwaError> { + fn busy(attempts: Vec<(i64, i64)>) -> Self { + Self { + attempts, + unresolvable: false, + } + } + + /// Leaves the attempts out of the rest of the pass. Busy attempts shrink + /// the pass's selection. Unresolvable attempts are replaced in it, so a + /// backlog of them, which every pass meets first, cannot fill each batch + /// and keep rescue from the attempts behind it. + fn exclude( + self, + excluded: &mut BTreeSet<(i64, i64)>, + selected: &mut Option>, + ) -> Result<(), AwaError> { let previous = excluded.len(); - excluded.extend(self.0); + excluded.extend(self.attempts); if excluded.len() == previous { return Err(AwaError::Validation( "queue storage retained attempt retry made no progress".to_string(), )); } + if self.unresolvable { + *selected = None; + } Ok(()) } } -/// A failure hydrating a deleted lease. `Ambiguous` names an attempt whose -/// retained ready bodies rank equally, so none can be chosen for it. +/// A failure hydrating deleted leases. `Ambiguous` names the attempts, as +/// `(job_id, run_lease)`, whose retained ready bodies rank equally, so none +/// can be chosen for them. enum HydrateError { - Ambiguous { job_id: i64, run_lease: i64 }, + Ambiguous(Vec<(i64, i64)>), Awa(AwaError), } @@ -2741,9 +2766,16 @@ impl From for HydrateError { impl From for AwaError { fn from(err: HydrateError) -> Self { match err { - HydrateError::Ambiguous { job_id, run_lease } => AwaError::Validation(format!( - "queue storage ready body ambiguous for deleted lease job {job_id} run_lease {run_lease}" - )), + HydrateError::Ambiguous(attempts) => AwaError::Validation(match attempts.as_slice() { + [] => "queue storage ready body ambiguous for deleted lease".to_string(), + [(job_id, run_lease)] => format!( + "queue storage ready body ambiguous for deleted lease job {job_id} run_lease {run_lease}" + ), + [(job_id, run_lease), rest @ ..] => format!( + "queue storage ready body ambiguous for deleted lease job {job_id} run_lease {run_lease} and {} more", + rest.len() + ), + }), HydrateError::Awa(err) => err, } } @@ -3113,6 +3145,64 @@ mod ready_segment_tests { } } + #[sqlx::test(migrations = false)] + async fn enqueue_batch_writes_more_segments_than_the_old_bind_limit(pool: sqlx::PgPool) { + use super::{QueueStorageConfig, RuntimeReadyRow}; + use crate::audited_sql; + + crate::migrations::run(&pool).await.unwrap(); + let store = QueueStorage::new(QueueStorageConfig { + schema: "awa_qs_segment_bind_limit".into(), + ..Default::default() + }) + .unwrap(); + store.install(&pool).await.unwrap(); + + // Distinct microsecond timestamps guarantee one segment per row. + // Reading the clock for every row can coalesce adjacent timestamps + // and fail to exercise the old 8,191-segment bind limit. + const JOBS: i64 = 9_000; + let first_run_at = Utc::now() - Duration::seconds(1); + let rows = (0..JOBS) + .map(|seq| RuntimeReadyRow { + job_id: None, + kind: "segment_test".into(), + queue: "segment-q".into(), + args: serde_json::json!({"seq": seq}), + priority: 2, + attempt: 0, + run_lease: 0, + max_attempts: 25, + run_at: first_run_at + Duration::microseconds(seq), + attempted_at: None, + created_at: first_run_at, + unique_key: None, + unique_states: None, + payload: serde_json::json!({}), + ordering_key: None, + }) + .collect(); + assert_eq!( + store.enqueue_runtime_rows(&pool, rows).await.unwrap(), + JOBS as usize + ); + let counts: (i64, i64) = sqlx::query_as(audited_sql(format!( + "SELECT (SELECT count(*) FROM {}.ready_entries), \ + (SELECT count(*) FROM {}.ready_segments)", + store.schema(), + store.schema() + ))) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(counts, (JOBS, JOBS)); + let claimed = store + .claim_runtime_batch(&pool, "segment-q", 100, std::time::Duration::from_secs(30)) + .await + .unwrap(); + assert_eq!(claimed.len(), 100); + } + #[test] fn ready_segments_split_on_run_at_boundaries() { let first = Utc @@ -3457,10 +3547,10 @@ fn select_claimed_ready_body<'a>( let mut best = candidates.iter().filter(|row| rank(row) == strongest); let ready = best.next().expect("ready body rank has a candidate"); if best.next().is_some() { - return Err(HydrateError::Ambiguous { - job_id: reference.job_id, - run_lease: reference.run_lease, - }); + return Err(HydrateError::Ambiguous(vec![( + reference.job_id, + reference.run_lease, + )])); } Ok((ready, has_emitted)) } @@ -9975,7 +10065,8 @@ impl QueueStorage { JOIN refs ON refs.job_id = items.job_id AND refs.run_lease = items.run_lease - WHERE NOT EXISTS ( + WHERE claim_batches.job_ids && $1::bigint[] + AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = claim_batches.claim_slot @@ -10314,7 +10405,8 @@ impl QueueStorage { claim_batches.attempts, claim_batches.max_attempts ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts) - WHERE items.job_id = $1 + WHERE claim_batches.job_ids @> ARRAY[$1::bigint] + AND items.job_id = $1 AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures @@ -12258,7 +12350,8 @@ impl QueueStorage { JOIN inflight ON inflight.job_id = items.job_id AND inflight.run_lease = items.run_lease - WHERE NOT EXISTS ( + WHERE claim_batches.job_ids && $1::bigint[] + AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = claim_batches.claim_slot @@ -12444,7 +12537,8 @@ impl QueueStorage { JOIN inflight ON inflight.job_id = items.job_id AND inflight.run_lease = items.run_lease - WHERE NOT EXISTS ( + WHERE claim_batches.job_ids && $1::bigint[] + AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = claim_batches.claim_slot @@ -12561,7 +12655,8 @@ impl QueueStorage { JOIN inflight ON inflight.job_id = items.job_id AND inflight.run_lease = items.run_lease - WHERE NOT EXISTS ( + WHERE claim_batches.job_ids && $1::bigint[] + AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = claim_batches.claim_slot @@ -12874,14 +12969,19 @@ impl QueueStorage { match hydrated { Ok(hydrated) => Ok(hydrated), // One unresolvable attempt must not stop the rescue of the others. - Err(HydrateError::Ambiguous { job_id, run_lease }) => { - tracing::error!( - schema = self.schema(), - job_id, - run_lease, - "Rescue skipped an attempt: queue storage ready body ambiguous for deleted lease; it stays running until its retained bodies are repaired" - ); - Ok(Err(BusyRetained(vec![(job_id, run_lease)]))) + Err(HydrateError::Ambiguous(attempts)) => { + for (job_id, run_lease) in &attempts { + tracing::error!( + schema = self.schema(), + job_id, + run_lease, + "Rescue skipped an attempt: queue storage ready body ambiguous for deleted lease; it stays running until its retained bodies are repaired" + ); + } + Ok(Err(BusyRetained { + attempts, + unresolvable: true, + })) } Err(HydrateError::Awa(err)) => Err(err), } @@ -12902,6 +13002,33 @@ impl QueueStorage { let run_leases: Vec = deleted.iter().map(|row| row.run_lease).collect(); let (ready_map, claimed_lanes) = self.load_ready_snapshot_candidates_tx(tx, &deleted).await?; + // Report every ambiguous attempt at once, before any write, so a + // caller that sets them aside needs one retry for all of them. + let ambiguous: Vec<(i64, i64)> = deleted + .iter() + .filter(|row| { + ready_map + .get(&( + row.ready_slot, + row.ready_generation, + row.queue.clone(), + row.enqueue_shard, + row.lane_seq, + row.job_id, + row.run_lease, + )) + .is_some_and(|candidates| { + matches!( + select_claimed_ready_body(candidates, &claimed_lanes, row), + Err(HydrateError::Ambiguous(_)) + ) + }) + }) + .map(|row| (row.job_id, row.run_lease)) + .collect(); + if !ambiguous.is_empty() { + return Err(HydrateError::Ambiguous(ambiguous)); + } let attempt_rows: Vec = sqlx::query_as(audited_sql(format!( r#" @@ -12980,7 +13107,8 @@ impl QueueStorage { JOIN refs ON refs.job_id = items.job_id AND refs.run_lease = items.run_lease - WHERE NOT EXISTS ( + WHERE claim_batches.job_ids && $1::bigint[] + AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures WHERE closures.claim_slot = claim_batches.claim_slot @@ -13158,10 +13286,10 @@ impl QueueStorage { .insert((source.job_id, source.run_lease), source) .is_some() { - return Err(HydrateError::Ambiguous { - job_id: source.job_id, - run_lease: source.run_lease + 1, - }); + return Err(HydrateError::Ambiguous(vec![( + source.job_id, + source.run_lease + 1, + )])); } } let job_ids: Vec = sources.iter().map(|row| row.job_id).collect(); @@ -13209,7 +13337,7 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; if !busy.is_empty() { - return Ok(Some(BusyRetained(busy))); + return Ok(Some(BusyRetained::busy(busy))); } } else { self.lock_claim_heads_of_jobs_with_wait_tx(tx, &job_ids, true) @@ -13294,10 +13422,10 @@ impl QueueStorage { for row in &emitted { let key = (row.job_id, row.run_lease); if !seen.insert(key) || expected.get(&key).copied() != Some(row) { - return Err(HydrateError::Ambiguous { - job_id: row.job_id, - run_lease: row.run_lease + 1, - }); + return Err(HydrateError::Ambiguous(vec![( + row.job_id, + row.run_lease + 1, + )])); } } sqlx::query(audited_sql(format!( @@ -13578,7 +13706,8 @@ impl QueueStorage { claim_batches.attempts, claim_batches.max_attempts ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts) - WHERE items.job_id = $1 + WHERE claim_batches.job_ids @> ARRAY[$1::bigint] + AND items.job_id = $1 AND items.run_lease = $2 AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures @@ -15509,7 +15638,8 @@ impl QueueStorage { LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = items.job_id AND attempt.run_lease = items.run_lease - WHERE items.job_id = $1 + WHERE claim_batches.job_ids @> ARRAY[$1::bigint] + AND items.job_id = $1 AND NOT EXISTS ( SELECT 1 FROM {schema}.lease_claim_closures AS closures @@ -18318,7 +18448,7 @@ impl QueueStorage { Ok(moved) => moved, Err(busy) => { attempt.rollback().await.map_err(map_sqlx_error)?; - busy.exclude(&mut excluded)?; + busy.exclude(&mut excluded, &mut selected)?; continue; } }; @@ -18329,7 +18459,7 @@ impl QueueStorage { Ok(moved) => moved, Err(busy) => { attempt.rollback().await.map_err(map_sqlx_error)?; - busy.exclude(&mut excluded)?; + busy.exclude(&mut excluded, &mut selected)?; continue; } }; @@ -18489,7 +18619,7 @@ impl QueueStorage { Ok(moved) => moved, Err(busy) => { attempt.rollback().await.map_err(map_sqlx_error)?; - busy.exclude(&mut excluded)?; + busy.exclude(&mut excluded, &mut selected)?; continue; } }; @@ -18501,7 +18631,7 @@ impl QueueStorage { Ok(moved) => moved, Err(busy) => { attempt.rollback().await.map_err(map_sqlx_error)?; - busy.exclude(&mut excluded)?; + busy.exclude(&mut excluded, &mut selected)?; continue; } }, @@ -18627,7 +18757,7 @@ impl QueueStorage { Ok(moved) => moved, Err(busy) => { attempt.rollback().await.map_err(map_sqlx_error)?; - busy.exclude(&mut excluded)?; + busy.exclude(&mut excluded, &mut selected)?; continue; } }; @@ -21106,6 +21236,7 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; let truncate_duration = truncate_started.elapsed(); + self.ensure_claim_batch_job_index_tx(&mut tx, slot).await?; sqlx::query(audited_sql(format!( r#" UPDATE {schema}.claim_ring_slots @@ -21146,6 +21277,84 @@ impl QueueStorage { Ok(outcome) } + /// Give `lease_claim_batches_` its GIN index on `job_ids` through + /// `{schema}.ensure_claim_batch_job_index`. Claim prune calls this right + /// after truncating the slot, while it holds the partition's + /// `ACCESS EXCLUSIVE` lock, so the build covers an empty table and waits + /// on nobody. A schema without the function or a role that may not + /// execute it skips the step. An error inside the function is rolled back + /// to a savepoint and logged; it never undoes the prune. + async fn ensure_claim_batch_job_index_tx( + &self, + tx: &mut sqlx::Transaction<'_, Postgres>, + slot: i32, + ) -> Result<(), AwaError> { + let schema = self.schema(); + let callable: bool = sqlx::query_scalar( + "SELECT COALESCE(has_function_privilege(to_regprocedure($1), 'EXECUTE'), false)", + ) + .bind(format!("{schema}.ensure_claim_batch_job_index(integer)")) + .fetch_one(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + if !callable { + tracing::debug!( + schema, + slot, + "Claim batch job index function is absent or not executable; skipping" + ); + return Ok(()); + } + + sqlx::query("SAVEPOINT awa_claim_batch_job_index") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let indexed: Result = sqlx::query_scalar(audited_sql(format!( + "SELECT {schema}.ensure_claim_batch_job_index($1)" + ))) + .bind(slot) + .fetch_one(tx.as_mut()) + .await; + match indexed { + Ok(indexed) => { + sqlx::query("RELEASE SAVEPOINT awa_claim_batch_job_index") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + if !indexed { + tracing::debug!(schema, slot, "Claim batch partition was not indexed"); + } + } + Err(sqlx::Error::Database(db_err)) => { + sqlx::query("ROLLBACK TO SAVEPOINT awa_claim_batch_job_index") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let code = db_err.code().map(|code| code.into_owned()); + if matches!(code.as_deref(), Some("42883" | "42501")) { + tracing::debug!( + schema, + slot, + code = code.as_deref().unwrap_or_default(), + error = %db_err.message(), + "Claim batch job index function unavailable; skipping" + ); + } else { + tracing::warn!( + schema, + slot, + code = code.as_deref().unwrap_or_default(), + error = %db_err.message(), + "Could not index claim batch partition by job id" + ); + } + } + Err(err) => return Err(map_sqlx_error(err)), + } + Ok(()) + } + fn job_id_sequence(&self) -> String { format!("{}.job_id_seq", self.schema()) } diff --git a/awa-python/examples/quickstart.py b/awa-python/examples/quickstart.py index 2fb02d975..2b6995e12 100644 --- a/awa-python/examples/quickstart.py +++ b/awa-python/examples/quickstart.py @@ -1,6 +1,6 @@ """Awa Python quickstart — a complete runnable example. -Requires: uv add awa-pg==0.6.6 +Requires: uv add awa-pg Requires: a running Postgres instance with DATABASE_URL set. Usage from the repository's awa-python directory: diff --git a/awa-worker/src/client.rs b/awa-worker/src/client.rs index 06e4ecd27..b888310e6 100644 --- a/awa-worker/src/client.rs +++ b/awa-worker/src/client.rs @@ -351,10 +351,11 @@ impl ClientBuilder { /// `Started` dispatch is scheduled after the claim commits and before the /// worker handler is invoked. Outcome handlers run only after the /// corresponding DB state transition commits. Hooks are best-effort - /// in-process notifications, not a durable workflow mechanism; because - /// they run in detached tasks, a very short job can complete before a - /// `Started` handler finishes. Capture any shared dependencies you need in - /// the closure environment. + /// in-process notifications, not a durable workflow mechanism. `Started` + /// runs alongside the handler, so a very short job can complete before a + /// `Started` handler finishes. Each attempt's hooks hold a finalizer slot + /// until they return, so hooks slower than job throughput slow claiming. + /// Capture any shared dependencies you need in the closure environment. pub fn on_event(mut self, handler: F) -> Self where T: JobArgs + DeserializeOwned + Send + Sync + 'static, diff --git a/awa-worker/src/executor.rs b/awa-worker/src/executor.rs index 98dcc8bff..b37ab3122 100644 --- a/awa-worker/src/executor.rs +++ b/awa-worker/src/executor.rs @@ -433,20 +433,21 @@ impl JobExecutor { progress_state.clone(), ); + let mut started_event = None; let result = match workers.get(&job.kind) { Some(worker) => { if has_lifecycle_handlers { let started_handlers = lifecycle_handlers.clone(); let started_kind = job_kind.clone(); let started_job = crate::output::without_outputs(&job).into_owned(); - tokio::spawn(async move { + started_event = Some(tokio::spawn(async move { dispatch_lifecycle_event( &started_handlers, &started_kind, UntypedJobEvent::Started { job: started_job }, ) .await; - }); + })); } let perform = CatchUnwind(worker.perform(&ctx)); let outcome = match abandon_after { @@ -521,8 +522,9 @@ impl JobExecutor { // Keep worker/heartbeat ownership while the bounded finalizer queue is full. let finalize_permit = match result { Some(_) => tokio::select! { - permit = finalize_capacity.acquire_owned() => Some(permit.expect("finalizer capacity is never closed")), + biased; _ = finalize_cancel.cancelled() => None, + permit = finalize_capacity.acquire_owned() => Some(permit.expect("finalizer capacity is never closed")), }, None => None, }; @@ -773,13 +775,22 @@ storage.queue_storage().is_some(), }, }; drop(heartbeat_guard); - drop(finalize_permit); - if let Some(event) = event { - tokio::select! { - _ = dispatch_lifecycle_event(&lifecycle_handlers, &event_kind, event) => {}, - _ = finalize_cancel.cancelled() => {}, + // Lifecycle hooks run under the finalizer permit, so a hook + // slower than job throughput holds back claiming instead of + // accumulating tasks that each retain a copy of the job. + let hooks = async { + if let Some(started_event) = started_event { + let _ = started_event.await; + } + if let Some(event) = event { + dispatch_lifecycle_event(&lifecycle_handlers, &event_kind, event).await; } + }; + tokio::select! { + _ = hooks => {}, + _ = finalize_cancel.cancelled() => {}, } + drop(finalize_permit); } .in_current_span(), ); @@ -2053,6 +2064,17 @@ fn rescheduled_event_row( row } +/// The job as a read returns it after this attempt completes. +fn completed_event_row(job: &JobRow) -> JobRow { + let mut row = job.clone(); + row.state = JobState::Completed; + row.finalized_at = Some(chrono::Utc::now()); + row.progress = None; + row.heartbeat_at = None; + row.deadline_at = None; + row +} + #[allow(clippy::too_many_arguments)] async fn complete_job_queue_storage( runtime: &QueueStorageRuntime, @@ -2158,16 +2180,11 @@ async fn complete_job_queue_storage( return Ok(CompletionOutcome::IgnoredStale); } if needs_event { - // A completed attempt changes only these fields, so the event is - // built from the worker's copy, which carries the sealed output, - // rather than re-reading the job. - let mut updated_job = job.clone(); - updated_job.state = JobState::Completed; - updated_job.finalized_at = Some(chrono::Utc::now()); - updated_job.progress = None; + // The event is built from the worker's copy, which carries the + // sealed output, rather than re-reading the job. Ok(CompletionOutcome::Applied { event: Some(UntypedJobEvent::Completed { - job: updated_job, + job: completed_event_row(job), duration, }), terminal: false, @@ -2704,12 +2721,7 @@ async fn complete_queue_storage_with_followups( // Synthesise the post-completion snapshot — the store has just moved the // lease into `done_entries` so a live SELECT would race with rotation. - // The happy-path fields are exactly those `complete_job_queue_storage` - // reconstructs when `load_job` returns None. - let mut updated_job = job.clone(); - updated_job.state = JobState::Completed; - updated_job.finalized_at = Some(chrono::Utc::now()); - updated_job.progress = None; + let updated_job = completed_event_row(job); crate::enqueue_specs::dispatch_specs_in_tx(&mut tx, &updated_job, specs, None).await?; let staged = @@ -4571,7 +4583,7 @@ false, } #[sqlx::test(migrations = false)] - async fn pending_outcome_hook_does_not_hold_durable_finalizer_capacity(pool: PgPool) { + async fn pending_outcome_hook_holds_finalizer_capacity(pool: PgPool) { let _database_guard = RECONNECT_DB_LOCK.lock().await; awa_model::migrations::run(&pool).await.unwrap(); let first = running_reconnect_job(&pool).await; @@ -4609,9 +4621,10 @@ false, RuntimeStorage::Canonical, 1, ); + let runs = Arc::new(AtomicU32::new(0)); let workers: HashMap = HashMap::from([( "reconnect_parent".into(), - Box::new(ReconnectWorker(Arc::new(AtomicU32::new(0)))) as BoxedWorker, + Box::new(ReconnectWorker(runs.clone())) as BoxedWorker, )]); let mut executor = JobExecutor::new( pool.clone(), @@ -4642,15 +4655,29 @@ false, tokio::time::timeout(Duration::from_secs(2), hook_started.notified()) .await .expect("first durable finalize must reach its outcome hook"); - tokio::time::timeout( - Duration::from_secs(2), + let second_task = tokio::spawn( executor.execute_task(dispatched(second), Arc::new(AtomicBool::new(false))), - ) + ); + tokio::time::timeout(Duration::from_secs(2), async { + while runs.load(Ordering::SeqCst) != 2 { + tokio::task::yield_now().await; + } + }) .await - .expect("pending outcome hook must not consume the only finalize permit"); + .expect("the second handler runs while the first hook is pending"); + tokio::time::sleep(Duration::from_millis(100)).await; + assert!( + !second_task.is_finished(), + "a pending outcome hook must keep the only finalize permit" + ); + release_hook.cancel(); + tokio::time::timeout(Duration::from_secs(2), second_task) + .await + .expect("the second attempt takes the permit once the hook returns") + .unwrap(); tokio::time::timeout(Duration::from_secs(2), second_completed.notified()) .await - .expect("second durable finalize must finish while the first hook is pending"); + .expect("second durable finalize must finish after the first hook returns"); let states: Vec = sqlx::query_scalar("SELECT state::text FROM awa.jobs_hot WHERE id = ANY($1)") .bind([first_id, second_id].as_slice()) @@ -4659,7 +4686,6 @@ false, .unwrap(); assert_eq!(states.len(), 2); assert!(states.iter().all(|state| state == "completed")); - release_hook.cancel(); finalizers.close(); tokio::time::timeout(Duration::from_secs(2), finalizers.wait()) .await diff --git a/awa-worker/src/maintenance.rs b/awa-worker/src/maintenance.rs index 2c43456a8..0a7136b7a 100644 --- a/awa-worker/src/maintenance.rs +++ b/awa-worker/src/maintenance.rs @@ -17,7 +17,7 @@ use croner::Cron; use sqlx::pool::PoolConnection; use sqlx::{PgPool, Postgres}; use std::collections::{HashMap, HashSet}; -use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::task::{JoinHandle, JoinSet}; @@ -25,6 +25,9 @@ use tokio_util::sync::CancellationToken; use tracing::{debug, error, info, warn}; use uuid::Uuid; +/// Detached lifecycle hook tasks maintenance may have outstanding at once. +const MAINTENANCE_LIFECYCLE_EVENT_CAPACITY: usize = 1024; + /// The canonical unique-claim conflict raised by `awa.sync_job_unique_claims` /// (SQLSTATE 23505, constraint reported as `idx_awa_jobs_unique`). See #388. fn is_unique_claim_conflict(err: &awa_model::AwaError) -> bool { @@ -706,6 +709,10 @@ pub struct MaintenanceService { /// In-process `Rescued` lifecycle hooks (ADR-015), shared with the /// executor. Best-effort fire-and-forget per rescued job. lifecycle_handlers: Arc>>, + /// Bounds the detached lifecycle hook tasks maintenance has outstanding; + /// events beyond it are dropped and counted. + lifecycle_event_capacity: Arc, + dropped_lifecycle_events: Arc, /// In-flight job cancellation flags — used to signal deadline/heartbeat rescue /// to running handlers on this worker instance. in_flight: InFlightMap, @@ -1022,6 +1029,10 @@ impl MaintenanceService { standby_queue_storage: None, enqueue_specs, lifecycle_handlers, + lifecycle_event_capacity: Arc::new(tokio::sync::Semaphore::new( + MAINTENANCE_LIFECYCLE_EVENT_CAPACITY, + )), + dropped_lifecycle_events: Arc::new(AtomicU64::new(0)), heartbeat_rescue_interval: Duration::from_secs(30), deadline_rescue_interval: DEFAULT_DEADLINE_RESCUE_INTERVAL, callback_rescue_interval: Duration::from_secs(30), @@ -2629,9 +2640,7 @@ impl MaintenanceService { job: job.clone(), reason, }; - tokio::spawn(async move { - crate::executor::dispatch_lifecycle_event(&handlers, &kind, event).await; - }); + self.spawn_lifecycle_event(handlers, kind, event); } } @@ -2718,8 +2727,41 @@ impl MaintenanceService { job: job.clone(), reason, }; + self.spawn_lifecycle_event(handlers, kind, event); + } + + /// Runs `kind`'s lifecycle hooks for `event` on a detached task, so a slow + /// hook does not hold up maintenance. At most + /// [`MAINTENANCE_LIFECYCLE_EVENT_CAPACITY`] such tasks are outstanding; an + /// event beyond that is dropped rather than retaining its job until the + /// hooks catch up. + fn spawn_lifecycle_event( + &self, + handlers: Arc>>, + kind: String, + event: crate::events::UntypedJobEvent, + ) { + if !handlers.contains_key(&kind) { + return; + } + let Ok(permit) = self.lifecycle_event_capacity.clone().try_acquire_owned() else { + let dropped = self + .dropped_lifecycle_events + .fetch_add(1, Ordering::Relaxed) + + 1; + if dropped % 1000 == 1 { + warn!( + kind, + dropped, + capacity = MAINTENANCE_LIFECYCLE_EVENT_CAPACITY, + "Dropped a maintenance lifecycle event: earlier hooks have not returned" + ); + } + return; + }; tokio::spawn(async move { crate::executor::dispatch_lifecycle_event(&handlers, &kind, event).await; + drop(permit); }); } diff --git a/awa/src/lib.rs b/awa/src/lib.rs index 81c5ec6fd..536b61f6e 100644 --- a/awa/src/lib.rs +++ b/awa/src/lib.rs @@ -12,13 +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, audited_sql, 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, AttemptError, AwaError, CallbackConfig, DefaultAction, DlqMetadata, DlqRow, - ExternalCompletion, InsertOpts, InsertParams, JobArgs, JobKindDescriptor, JobOutput, JobRow, - JobState, ListDlqFilter, OutputDescriptor, 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, AttemptError, AwaError, CallbackConfig, DefaultAction, + DlqMetadata, DlqRow, ExternalCompletion, InsertOpts, InsertParams, JobArgs, JobKindDescriptor, + JobOutput, JobRow, JobState, ListDlqFilter, OutputDescriptor, PartitionedQueue, + PartitionedQueueError, PreparedJobInsert, QueueDescriptor, QueueStorage, QueueStorageConfig, + ResolveOutcome, RetryFromDlqOpts, StorageCapability, StorageStatus, UniqueOpts, }; // Re-export worker runtime diff --git a/awa/tests/enqueue_lane_contention_test.rs b/awa/tests/enqueue_lane_contention_test.rs new file mode 100644 index 000000000..84e6f2f69 --- /dev/null +++ b/awa/tests/enqueue_lane_contention_test.rs @@ -0,0 +1,819 @@ +//! 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_BATCH_FUNCTION +//! 0|1: drop `awa.insert_jobs_compat_batch` +//! and its v055 version row 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_BATCH_FUNCTION", 0) != 0 { + sqlx::raw_sql( + "DROP FUNCTION IF EXISTS awa.insert_jobs_compat_batch(text[], text[], jsonb[], \ + text[], smallint[], smallint[], timestamptz[], jsonb[], jsonb[], bytea[], text[], \ + text[], boolean); \ + DELETE FROM awa.schema_version WHERE version = 55", + ) + .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/awa/tests/enqueue_lane_lock_order_test.rs b/awa/tests/enqueue_lane_lock_order_test.rs index 924d212ae..74e3fbb4b 100644 --- a/awa/tests/enqueue_lane_lock_order_test.rs +++ b/awa/tests/enqueue_lane_lock_order_test.rs @@ -16,7 +16,8 @@ use awa::audited_sql; 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; @@ -33,6 +34,7 @@ const WAIT_DEADLINE: Duration = Duration::from_secs(20); #[derive(Debug, Clone, Copy)] enum Producer { InsertMany, + SkipUniqueConflicts, CopySmall, CopyLarge, CopyUnique, @@ -49,6 +51,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::CopySmall | Producer::CopyLarge | Producer::CopyUnique => { insert_many_copy_from_pool(pool, jobs).await } @@ -251,6 +254,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::CopySmall).await; @@ -282,6 +290,74 @@ fn ordering_key_per_shard(shards: i16) -> Vec> { .collect() } +/// Key of the advisory lock `reserve_enqueue_seq` holds for one lane. +async fn lane_reservation_lock_key(pool: &PgPool, queue: &str, priority: i16, shard: i16) -> i64 { + sqlx::query_scalar( + r#" + SELECT hashtextextended( + format('%I.%I', 'awa', awa.queue_lane_sequence_name('queue_enqueue_seq', $1, $2, $3)), + 0 + ) + "#, + ) + .bind(queue) + .bind(priority) + .bind(shard) + .fetch_one(pool) + .await + .expect("derive lane reservation lock key") +} + +/// The backend waiting for the advisory lock `key`, once there is one. +async fn advisory_waiter_pid(pool: &PgPool, key: i64) -> i32 { + let started = Instant::now(); + loop { + let pid: Option = sqlx::query_scalar( + r#" + SELECT pid + FROM pg_locks + WHERE locktype = 'advisory' + AND NOT granted + AND database = (SELECT oid FROM pg_database WHERE datname = current_database()) + AND ((classid::bigint << 32) | objid::bigint) = $1 + LIMIT 1 + "#, + ) + .bind(key) + .fetch_optional(pool) + .await + .expect("poll pg_locks"); + if let Some(pid) = pid { + return pid; + } + assert!( + started.elapsed() < WAIT_DEADLINE, + "no backend started waiting for advisory lock {key}" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } +} + +async fn granted_advisory_keys(pool: &PgPool, pid: i32) -> std::collections::BTreeSet { + sqlx::query_scalar::<_, i64>( + r#" + SELECT (classid::bigint << 32) | objid::bigint + FROM pg_locks + WHERE locktype = 'advisory' AND granted AND pid = $1 + "#, + ) + .bind(pid) + .fetch_all(pool) + .await + .expect("read granted advisory locks") + .into_iter() + .collect() +} + +/// A batch over sharded lanes reserves them in `(queue, priority, shard)` +/// order, whatever order its rows arrive in. A gate holds one lane in the +/// middle of that order; the producer must be parked on it holding exactly the +/// lanes before it. async fn assert_batch_reserves_lanes_in_canonical_order(producer: Producer) { let _guard = LANE_LOCK_TEST_GUARD.lock().await; let (pool, _config) = setup_queue_storage().await; @@ -309,6 +385,7 @@ async fn assert_batch_reserves_lanes_in_canonical_order(producer: Producer) { let keys = ordering_key_per_shard(shards); let mut jobs = Vec::new(); + let mut lanes = std::collections::BTreeSet::new(); let rounds = match producer { Producer::CopyLarge => 700, _ => 12, @@ -317,55 +394,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"); + } } } @@ -374,6 +508,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::CopySmall).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..16088a4cb --- /dev/null +++ b/awa/tests/insert_many_batch_test.rs @@ -0,0 +1,657 @@ +//! Compatibility inserts and the batch insert function +//! `awa.insert_jobs_compat_batch` (migration v055). +//! +//! 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` + +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"); +} + +/// 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 new file mode 100644 index 000000000..81b62a575 --- /dev/null +++ b/awa/tests/insert_many_schema_fallback_test.rs @@ -0,0 +1,445 @@ +//! `insert_many` and `insert_many_copy` against a schema without +//! `awa.insert_jobs_compat_batch` (migration v055), and the migration 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 function 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::audited_sql; +use awa::model::{ + insert_many, insert_many_copy, insert_many_copy_from_pool, insert_many_skip_unique_conflicts, + migrations, AwaError, +}; +use awa::{InsertOpts, InsertParams, UniqueOpts}; +use awa_testing::setup; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; +use sqlx::PgPool; +use std::str::FromStr; +use uuid::Uuid; + +const BATCH_FUNCTION: &str = "insert_jobs_compat_batch"; +const BATCH_FUNCTION_VERSION: i32 = 55; +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(audited_sql(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(audited_sql(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 has_batch_function(pool: &PgPool) -> bool { + sqlx::query_scalar( + "SELECT to_regprocedure('awa.insert_jobs_compat_batch(text[],text[],jsonb[],text[],\ + smallint[],smallint[],timestamptz[],jsonb[],jsonb[],bytea[],text[],text[],boolean)') \ + IS NOT NULL", + ) + .fetch_one(pool) + .await + .expect("probe the batch function") +} + +/// The database as a schema before v055 left it. +async fn roll_back_to_before_batch_function(pool: &PgPool) { + sqlx::query(DROP_BATCH_FUNCTION) + .execute(pool) + .await + .expect("drop batch function"); + sqlx::query("DELETE FROM awa.schema_version WHERE version >= $1") + .bind(BATCH_FUNCTION_VERSION) + .execute(pool) + .await + .expect("roll the schema version back"); +} + +#[tokio::test] +async fn binary_ahead_of_the_migration_falls_back_and_recovers() { + let db = ScratchDatabase::create().await; + let pool = &db.pool; + assert!(has_batch_function(pool).await); + + roll_back_to_before_batch_function(pool).await; + + 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 function and says so. + let skipped = insert_many_skip_unique_conflicts(pool, &jobs("fallback_u", 0..1, true)) + .await + .expect_err("needs the batch function"); + assert!( + matches!(&skipped, AwaError::Validation(message) if message.contains(BATCH_FUNCTION)), + "{skipped:?}" + ); + + migrations::run(pool).await.expect("migrate"); + assert!(has_batch_function(pool).await); + + // 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 migration"); + 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 migration"); + 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 batch function never selects SQL for +/// another, and both paths join the caller's transaction, including its +/// rollback. +#[tokio::test] +async fn mixed_batch_function_databases_do_not_abort_caller_transactions() { + let with_function = ScratchDatabase::create().await; + let without_function = ScratchDatabase::create().await; + for db in [&with_function, &without_function] { + sqlx::query("CREATE TABLE application_writes (id bigint PRIMARY KEY)") + .execute(&db.pool) + .await + .expect("create application table"); + } + roll_back_to_before_batch_function(&without_function.pool).await; + + insert_many_skip_unique_conflicts(&with_function.pool, &jobs("mixed_warm", 0..2, true)) + .await + .expect("batch API on the database with the function"); + + for (index, db) in [&with_function, &without_function, &with_function] + .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); + } + + with_function.drop().await; + without_function.drop().await; +} + +/// The exported v055 SQL, as an external runner applies it: twice in one +/// explicit transaction, then over an installed function. +#[tokio::test] +async fn exported_migration_reapplies_in_one_transaction() { + let db = ScratchDatabase::create().await; + let pool = &db.pool; + let (_, _, sql) = + migrations::migration_sql_range(BATCH_FUNCTION_VERSION - 1, BATCH_FUNCTION_VERSION) + .pop() + .expect("v055 is a migration"); + + roll_back_to_before_batch_function(pool).await; + sqlx::raw_sql(audited_sql(format!("BEGIN;\n{sql}\n{sql}\nCOMMIT;"))) + .execute(pool) + .await + .expect("apply v055 twice in one transaction"); + sqlx::raw_sql(audited_sql(sql)) + .execute(pool) + .await + .expect("apply over an installed function"); + + assert!(has_batch_function(pool).await); + let recorded: i64 = + sqlx::query_scalar("SELECT count(*) FROM awa.schema_version WHERE version = $1") + .bind(BATCH_FUNCTION_VERSION) + .fetch_one(pool) + .await + .expect("read schema versions"); + assert_eq!(recorded, 1); + let copied = insert_many_copy_from_pool(pool, &jobs("exported", 0..4, false)) + .await + .expect("COPY through the exported migration"); + 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(audited_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"); + + // 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(audited_sql(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(audited_sql(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; + sqlx::query(audited_sql(format!("DROP ROLE {role}"))) + .execute(&admin) + .await + .expect("drop role"); + admin.close().await; + if let Err(failure) = outcome { + std::panic::resume_unwind(failure.into_panic()); + } +} diff --git a/awa/tests/integration_test.rs b/awa/tests/integration_test.rs index c816f0675..a4b9c6567 100644 --- a/awa/tests/integration_test.rs +++ b/awa/tests/integration_test.rs @@ -552,6 +552,9 @@ async fn test_insert_many() { #[tokio::test] async fn test_insert_many_exceeding_bind_parameter_limit() { const JOB_COUNT: usize = 6_000; + // A concurrent test that activates queue storage would route part of the + // batch away from the canonical tables this test reads back. + let _guard = test_lock().lock().await; let client = setup_with_connections(4).await; let queue = "integ_insert_many_bind_limit"; clean_queue(client.pool(), queue).await; diff --git a/awa/tests/lifecycle_hook_test.rs b/awa/tests/lifecycle_hook_test.rs index 926de9df5..e201c57e6 100644 --- a/awa/tests/lifecycle_hook_test.rs +++ b/awa/tests/lifecycle_hook_test.rs @@ -9,7 +9,7 @@ use awa::{ }; use awa_testing::setup::TestDatabase; use serde::{Deserialize, Serialize}; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::time::Duration; use tokio::sync::mpsc; @@ -81,6 +81,11 @@ struct HookJob { value: String, } +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +struct HookFollowUp { + source_job_id: i64, +} + #[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] struct RawHookJob { value: String, @@ -139,6 +144,100 @@ async fn test_typed_completed_event_handler_runs() { assert_eq!(stored.state, JobState::Completed); } +/// A completed event carries the job as a terminal read returns it, without the +/// finished attempt's heartbeat or deadline. +#[tokio::test] +async fn test_completed_event_of_a_deadline_job_matches_the_terminal_job() { + assert_completed_event_of_a_deadline_job_matches_the_terminal_job(false).await; +} + +/// A kind with a Completed follow-up completes through its own transaction; the +/// follow-up and the event see the same terminal job. +#[tokio::test] +async fn test_completed_event_with_a_follow_up_matches_the_terminal_job() { + assert_completed_event_of_a_deadline_job_matches_the_terminal_job(true).await; +} + +type AttemptTimes = ( + Option>, + Option>, +); + +async fn assert_completed_event_of_a_deadline_job_matches_the_terminal_job(with_follow_up: bool) { + let db = TestDatabase::queue_storage().await; + let pool = db.pool().clone(); + let queue = "lifecycle_completed_deadline"; + + let (tx, mut rx) = mpsc::unbounded_channel(); + let follow_up_saw: Arc>> = Arc::default(); + let builder = lifecycle_client_builder(pool.clone()) + .queue( + queue, + QueueConfig { + poll_interval: Duration::from_millis(25), + deadline_duration: Duration::from_secs(60), + ..Default::default() + }, + ) + .register::(|_args, _ctx| async move { Ok(JobResult::Completed) }) + .on_event::(move |event| { + let tx = tx.clone(); + async move { + if let JobEvent::Completed { job, .. } = event { + tx.send(job).unwrap(); + } + } + }); + let builder = if with_follow_up { + let follow_up_saw = follow_up_saw.clone(); + builder.on_completed_enqueue::(move |_args, job| { + follow_up_saw + .lock() + .unwrap() + .push((job.heartbeat_at, job.deadline_at)); + HookFollowUp { + source_job_id: job.id, + } + }) + } else { + builder + }; + let client = builder.build().unwrap(); + + let inserted = awa::insert_with( + &pool, + &HookJob { + action: "complete".into(), + value: "deadline".into(), + }, + awa::InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await + .unwrap(); + + client.start().await.unwrap(); + let event_job = recv_event(&mut rx).await; + client.shutdown(Duration::from_secs(2)).await; + + let stored = admin::get_job(&pool, inserted.id).await.unwrap(); + assert_eq!(event_job.id, stored.id); + assert_eq!(event_job.state, stored.state); + assert_eq!( + (event_job.heartbeat_at, event_job.deadline_at), + (stored.heartbeat_at, stored.deadline_at) + ); + assert_eq!((stored.heartbeat_at, stored.deadline_at), (None, None)); + let expected_follow_up_saw: Vec = if with_follow_up { + vec![(None, None)] + } else { + Vec::new() + }; + assert_eq!(*follow_up_saw.lock().unwrap(), expected_follow_up_saw); +} + #[tokio::test] async fn test_typed_started_event_handler_runs() { let db = TestDatabase::queue_storage().await; diff --git a/awa/tests/persona_soak/db.rs b/awa/tests/persona_soak/db.rs index 0bd0b8e33..dafc1dce3 100644 --- a/awa/tests/persona_soak/db.rs +++ b/awa/tests/persona_soak/db.rs @@ -177,12 +177,17 @@ impl Budget { } } - /// Reserves up to `want` connections under `label`, never fewer than `min`. - /// Returns the amount granted; panics when even `min` does not fit, because the - /// run would exhaust the server before it tested anything. + /// Reserves up to `want` connections under `label`, never fewer than `min`, + /// replacing any earlier reservation under the same label. Returns the amount + /// granted; panics when even `min` does not fit, because the run would exhaust + /// the server before it tested anything. pub fn reserve(&self, label: &str, want: u32, min: u32) -> u32 { let mut allocations = self.allocations.lock().expect("budget lock"); - let used: u32 = allocations.values().sum(); + let used: u32 = allocations + .iter() + .filter(|(held_by, _)| held_by.as_str() != label) + .map(|(_, count)| count) + .sum(); let available = self.usable.saturating_sub(used); assert!( available >= min, @@ -221,7 +226,8 @@ pub type ClientFactory = Arc Client + Send + Sync>; /// shutdown handshake and the connections drop. pub struct ClientHost { pub label: String, - runtime: tokio::runtime::Runtime, + /// Always `Some` until the host is stopped or dropped. + runtime: Option, client: Arc, pool: PgPool, } @@ -255,7 +261,7 @@ impl ClientHost { })?; Ok(Self { label: thread_label, - runtime, + runtime: Some(runtime), client: Arc::new(client), pool, }) @@ -273,26 +279,95 @@ impl ClientHost { } /// Graceful stop: stop claiming, drain in-flight jobs, then close the pool. - pub async fn shutdown(self, timeout: Duration) { + pub async fn shutdown(mut self, timeout: Duration) { + let Some(runtime) = self.runtime.take() else { + return; + }; let client = self.client.clone(); let pool = self.pool.clone(); - let _ = self - .runtime + let _ = runtime .handle() .spawn(async move { client.shutdown(timeout).await; pool.close().await; }) .await; - self.runtime.shutdown_background(); + runtime.shutdown_background(); } /// Abrupt stop: no shutdown handshake, as when a process is killed. pub fn kill(self) { - self.runtime.shutdown_background(); + drop(self); + } +} + +/// Dropping a host kills it. A tokio runtime dropped the ordinary way blocks +/// until its tasks finish, which panics inside another runtime. +impl Drop for ClientHost { + fn drop(&mut self) { + if let Some(runtime) = self.runtime.take() { + runtime.shutdown_background(); + } } } pub fn url_for_database(config: &SoakConfig, name: &str) -> String { replace_database_name(&config.database_url, name) } + +#[cfg(test)] +mod tests { + use super::{Budget, ClientHost, ServerInfo}; + use awa::{Client, QueueConfig}; + use sqlx::postgres::PgPoolOptions; + use std::sync::Arc; + + fn budget(usable: u32) -> Budget { + Budget::new( + &ServerInfo { + max_connections: usable, + ..ServerInfo::default() + }, + 0, + ) + } + + #[test] + fn replacing_a_reservation_reuses_its_own_connections() { + let budget = budget(20); + budget.reserve("other", 10, 1); + assert_eq!(budget.reserve("earmark", 10, 2), 10); + assert_eq!( + budget.reserve("earmark", 10, 10), + 10, + "the earmark's own connections count as available to its replacement" + ); + assert_eq!( + budget.table(), + vec![("earmark".to_string(), 10), ("other".to_string(), 10)] + ); + } + + #[tokio::test] + async fn dropping_a_live_host_inside_another_runtime_does_not_panic() { + let pool = PgPoolOptions::new() + .connect_lazy("postgres://postgres:test@localhost:1/awa_soak_unused") + .expect("lazy pool"); + let client = Client::builder(pool.clone()) + .queue("awa_soak_host_drop", QueueConfig::default()) + .build() + .expect("build an unstarted client"); + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .expect("host runtime"); + let host = ClientHost { + label: "host_drop".to_string(), + runtime: Some(runtime), + client: Arc::new(client), + pool, + }; + drop(host); + } +} diff --git a/awa/tests/persona_soak/oracle.rs b/awa/tests/persona_soak/oracle.rs index 5c6840c3a..1ea04fc8d 100644 --- a/awa/tests/persona_soak/oracle.rs +++ b/awa/tests/persona_soak/oracle.rs @@ -288,8 +288,7 @@ impl Oracle { let scan_started_ms = self.shared.now_ms(); let rows: Result, _> = sqlx::query_as(awa::audited_sql(format!( "SELECT job_id, queue, run_lease FROM {schema}.done_entries - WHERE state = 'cancelled' AND finalized_at > now() - make_interval(secs => $1) - LIMIT 2000" + WHERE state = 'cancelled' AND finalized_at > now() - make_interval(secs => $1)" ))) .bind(((scan_started_ms - since_ms) as f64 / 1000.0 + 15.0).max(20.0)) .fetch_all(&self.pool) @@ -1176,6 +1175,46 @@ mod tests { } } + #[tokio::test] + async fn cancel_scan_reads_every_cancellation_in_its_window() { + let mut run = TestRun::create().await; + let schema = schema(&run.shared); + for (queue, count) in [("oracle_cancel_bulk", 2_000), ("oracle_cancel_late", 1)] { + run.store + .enqueue_batch(&run.pool, queue, 2, count) + .await + .expect("enqueue jobs to cancel"); + awa::model::admin::drain_queue(&run.pool, queue) + .await + .expect("cancel the queue"); + } + let late: i64 = sqlx::query_scalar(awa::audited_sql(format!( + "SELECT job_id FROM {schema}.done_entries WHERE queue = 'oracle_cancel_late'" + ))) + .fetch_one(&run.pool) + .await + .expect("the late cancellation"); + run.shared.ledger.register( + late, + JobRecord::new(Persona::Priya, u32::MAX, 0, expect::COMPLETED, 0), + ); + run.oracle.check_unexplained_cancels().await; + assert!( + run.shared + .violations + .lock() + .expect("violations") + .iter() + .any(|violation| violation.kind == "unexplained_cancel" + && violation + .rows + .iter() + .any(|row| row.contains(&format!("\"job_id\":{late},")))), + "a cancellation after the first 2,000 in the window is still checked" + ); + run.drop().await; + } + #[tokio::test] async fn returned_handler_without_finalization_cannot_pass_the_oracle() { let mut run = TestRun::create().await; diff --git a/awa/tests/persona_soak/personas/alex.rs b/awa/tests/persona_soak/personas/alex.rs index 88498cf04..ccd3f1698 100644 --- a/awa/tests/persona_soak/personas/alex.rs +++ b/awa/tests/persona_soak/personas/alex.rs @@ -396,9 +396,10 @@ impl Context { } } 2 => { - // Bulk retry: every exhausted Mia job may come back. - self.shared.ledger.mark_where( - |record| record.persona == Persona::Mia && record.has(flag::EXHAUSTED), + // Bulk retry: every exhausted job on the infer queue may come back. + super::mia::mark_exhausted_in_queue( + &self.shared, + super::mia::INFER_QUEUE, flag::DLQ_RETRIED, ); let bulk = ListDlqFilter { @@ -416,8 +417,9 @@ impl Context { } } _ => { - self.shared.ledger.mark_where( - |record| record.persona == Persona::Mia && record.has(flag::EXHAUSTED), + super::mia::mark_exhausted_in_queue( + &self.shared, + super::mia::TRAIN_QUEUE, flag::DLQ_PURGED, ); let bulk = ListDlqFilter { diff --git a/awa/tests/persona_soak/personas/chaos.rs b/awa/tests/persona_soak/personas/chaos.rs index fee1c63eb..99def38ee 100644 --- a/awa/tests/persona_soak/personas/chaos.rs +++ b/awa/tests/persona_soak/personas/chaos.rs @@ -14,6 +14,23 @@ use std::time::{Duration, Instant}; const MAINTENANCE_LEADER_LOCK_KEY: i64 = 0x4157_415f_4d41_494e; const RUNTIME_FLEETS: [&str; 6] = ["mia", "tom", "priya", "sam", "sam-stripe", "fanout"]; +/// Up to `$1` random runtime sessions. Producers, operators, the oracle, chaos +/// itself, the claim-and-abandon persona and the shard-ops control sessions are +/// left alone: they are the harness, not the system under test. +macro_rules! runtime_victims_sql { + () => { + "SELECT pid FROM pg_stat_activity + WHERE application_name LIKE 'soak\\_%' AND application_name NOT LIKE '%producer%' + AND application_name NOT LIKE '%operator%' AND application_name NOT LIKE '%oracle%' + AND application_name NOT LIKE '%chaos%' AND application_name NOT LIKE 'soak\\_abandon\\_%' + AND application_name NOT IN + ('soak_shardops_owner', 'soak_shardops_migrator', 'soak_shardops_admin') + AND backend_type = 'client backend' + AND pid <> pg_backend_pid() + ORDER BY random() LIMIT $1" + }; +} + async fn kill_runtime(shared: &Arc, rng: &mut Rng) { let Some(&name) = rng.pick(&RUNTIME_FLEETS) else { return; @@ -92,17 +109,11 @@ async fn terminate_leader(shared: &Arc, pool: &PgPool) { async fn terminate_runtime_connections(shared: &Arc, pool: &PgPool, rng: &mut Rng) { let count = rng.range(3, 12) as i64; - let terminated: Result, _> = sqlx::query_as( - "SELECT pid FROM ( - SELECT pid FROM pg_stat_activity - WHERE application_name LIKE 'soak\\_%' AND application_name NOT LIKE '%producer%' - AND application_name NOT LIKE '%operator%' AND application_name NOT LIKE '%oracle%' - AND application_name NOT LIKE '%chaos%' AND application_name NOT LIKE 'soak\\_abandon\\_%' - AND backend_type = 'client backend' - AND pid <> pg_backend_pid() - ORDER BY random() LIMIT $1 - ) victims WHERE pg_terminate_backend(pid)", - ) + let terminated: Result, _> = sqlx::query_as(concat!( + "SELECT pid FROM (", + runtime_victims_sql!(), + ") victims WHERE pg_terminate_backend(pid)" + )) .bind(count) .fetch_all(pool) .await; @@ -240,3 +251,50 @@ pub async fn run(shared: Arc) -> Result<(), String> { pool.close().await; Ok(()) } + +#[cfg(test)] +mod tests { + use crate::persona_soak::db::connect_pool; + use std::collections::HashSet; + + #[tokio::test] + async fn runtime_termination_spares_the_shard_ops_control_sessions() { + let url = std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".to_string()); + let mut sessions = Vec::new(); + for application in [ + "soak_shardops_owner", + "soak_shardops_migrator", + "soak_shardops_admin", + "soak_shardops-a_0", + ] { + let pool = connect_pool(&url, application, 1).await.expect("session"); + let pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&pool) + .await + .expect("backend pid"); + sessions.push((application, pid, pool)); + } + let probe = connect_pool(&url, "chaos_victim_probe", 1) + .await + .expect("probe"); + let victims: HashSet = sqlx::query_scalar(runtime_victims_sql!()) + .bind(i64::from(i32::MAX)) + .fetch_all(&probe) + .await + .expect("select victims") + .into_iter() + .collect(); + for (application, pid, _) in &sessions { + assert_eq!( + victims.contains(pid), + *application == "soak_shardops-a_0", + "{application}" + ); + } + probe.close().await; + for (_, _, pool) in sessions { + pool.close().await; + } + } +} diff --git a/awa/tests/persona_soak/personas/mia.rs b/awa/tests/persona_soak/personas/mia.rs index 07d77f276..0100a2589 100644 --- a/awa/tests/persona_soak/personas/mia.rs +++ b/awa/tests/persona_soak/personas/mia.rs @@ -258,3 +258,58 @@ pub async fn run(shared: Arc) -> Result<(), String> { producer_pool.close().await; Ok(()) } + +/// Flags Mia's exhausted jobs on `queue`: the rows a DLQ operation filtered to that +/// queue acts on. Returns how many were flagged. +pub fn mark_exhausted_in_queue(shared: &Shared, queue: &str, flags: u32) -> usize { + let Some(queue) = shared.queues.get(queue) else { + return 0; + }; + shared.ledger.mark_where( + |record| { + record.persona == Persona::Mia + && record.queue == queue.index + && record.has(flag::EXHAUSTED) + }, + flags, + ) +} + +#[cfg(test)] +mod tests { + use super::{mark_exhausted_in_queue, DEADLINE_QUEUE, INFER_QUEUE, TRAIN_QUEUE}; + use crate::persona_soak::config::{Persona, SoakConfig}; + use crate::persona_soak::db::{Budget, ServerInfo}; + use crate::persona_soak::ledger::{expect, flag, JobRecord}; + use crate::persona_soak::shared::Shared; + + #[test] + fn a_dlq_operation_flags_only_its_own_queue() { + let config = SoakConfig::for_test( + "postgres://postgres:test@localhost/awa_soak_dlq_marks".to_string(), + ); + let info = ServerInfo::default(); + let budget = Budget::new(&info, 0); + let shared = Shared::new(config, info, budget); + for (id, queue) in [(1, TRAIN_QUEUE), (2, INFER_QUEUE), (3, DEADLINE_QUEUE)] { + let index = shared.queues.register(queue, Persona::Mia, false).index; + shared.ledger.register( + id, + JobRecord::new( + Persona::Mia, + index, + 0, + expect::ANY_TERMINAL, + flag::EXHAUSTED, + ), + ); + } + assert_eq!( + mark_exhausted_in_queue(&shared, TRAIN_QUEUE, flag::DLQ_PURGED), + 1 + ); + assert!(shared.ledger.get(1).unwrap().has(flag::DLQ_PURGED)); + assert!(!shared.ledger.get(2).unwrap().has(flag::DLQ_PURGED)); + assert!(!shared.ledger.get(3).unwrap().has(flag::DLQ_PURGED)); + } +} diff --git a/awa/tests/persona_soak/personas/python.rs b/awa/tests/persona_soak/personas/python.rs index 4c0808cfe..491b0b75e 100644 --- a/awa/tests/persona_soak/personas/python.rs +++ b/awa/tests/persona_soak/personas/python.rs @@ -141,6 +141,20 @@ fn spawn_process( .env("SOAK_QUEUE_WORK", QUEUE_WORK) .env("SOAK_QUEUE_DEADLINE", QUEUE_DEADLINE) .env("SOAK_QUEUE_SHARED", QUEUE_SHARED) + .env( + "SOAK_HEARTBEAT_MS", + shared.config.heartbeat_interval.as_millis().to_string(), + ) + .env( + "SOAK_STALENESS_MS", + shared.config.heartbeat_staleness.as_millis().to_string(), + ) + .envs( + shared + .config + .rescue_interval + .map(|interval| ("SOAK_RESCUE_INTERVAL_MS", interval.as_millis().to_string())), + ) .stdout(Stdio::piped()) .stderr(Stdio::piped()) .kill_on_drop(true) diff --git a/awa/tests/persona_soak/personas/tom.rs b/awa/tests/persona_soak/personas/tom.rs index af963e56c..3f3b5ede3 100644 --- a/awa/tests/persona_soak/personas/tom.rs +++ b/awa/tests/persona_soak/personas/tom.rs @@ -54,6 +54,12 @@ impl Worker for TomWorker { } } +/// Tenant queues to serve: the scaled count, and never fewer than the queues the cron +/// schedules write to. +fn queue_count(scaled: u64) -> usize { + (scaled as usize).max(CRON_JOBS) +} + fn queue_name(index: usize) -> String { format!("{PREFIX}{index:03}") } @@ -109,7 +115,7 @@ fn client_factory(shared: Arc, queue_count: usize) -> ClientFactory { pub async fn run(shared: Arc) -> Result<(), String> { let persona = Persona::Tom; - let queue_count = shared.config.scaled(100) as usize; + let queue_count = queue_count(shared.config.scaled(100)); let mut served = Vec::new(); for index in 0..queue_count { served.push(( @@ -277,3 +283,14 @@ pub async fn run(shared: Arc) -> Result<(), String> { producer_pool.close().await; Ok(()) } + +#[cfg(test)] +mod tests { + use super::{queue_count, CRON_JOBS}; + + #[test] + fn every_cron_target_queue_is_served_at_small_scales() { + assert_eq!(queue_count(2), CRON_JOBS); + assert_eq!(queue_count(100), 100); + } +} diff --git a/awa/tests/persona_soak/personas/ui.rs b/awa/tests/persona_soak/personas/ui.rs index 97f5d4d4e..a961ca98f 100644 --- a/awa/tests/persona_soak/personas/ui.rs +++ b/awa/tests/persona_soak/personas/ui.rs @@ -571,10 +571,7 @@ pub async fn run_ui_operator(shared: Arc) -> Result<(), String> { 5 => { // Bulk actions on the DLQ by filter, as the DLQ page's toolbar does. if rng.chance(50) { - shared.ledger.mark_where( - |record| record.persona == Persona::Mia && record.has(flag::EXHAUSTED), - flag::DLQ_RETRIED, - ); + mia::mark_exhausted_in_queue(&shared, mia::INFER_QUEUE, flag::DLQ_RETRIED); http.post( "/api/dlq/bulk-retry", "/api/dlq/bulk-retry", @@ -582,10 +579,7 @@ pub async fn run_ui_operator(shared: Arc) -> Result<(), String> { ) .await; } else { - shared.ledger.mark_where( - |record| record.persona == Persona::Mia && record.has(flag::EXHAUSTED), - flag::DLQ_PURGED, - ); + mia::mark_exhausted_in_queue(&shared, mia::TRAIN_QUEUE, flag::DLQ_PURGED); http.post( "/api/dlq/bulk-purge", "/api/dlq/bulk-purge", @@ -594,7 +588,8 @@ pub async fn run_ui_operator(shared: Arc) -> Result<(), String> { .await; } } - 6 => { + // Tom declares the cron schedules; without Tom there are none to act on. + 6 if shared.config.has(Persona::Tom) => { let name = format!("soak_tom_cron_{}", rng.below(4)); let paused = http .post( @@ -616,7 +611,7 @@ pub async fn run_ui_operator(shared: Arc) -> Result<(), String> { ) .await; } - 7 => { + 7 if shared.config.has(Persona::Tom) => { let name = format!("soak_tom_cron_{}", rng.below(4)); http.post( "/api/cron/{name}/trigger", diff --git a/awa/tests/persona_soak/python/soak_python.py b/awa/tests/persona_soak/python/soak_python.py index ff6920c3e..b5fe642f8 100644 --- a/awa/tests/persona_soak/python/soak_python.py +++ b/awa/tests/persona_soak/python/soak_python.py @@ -33,6 +33,11 @@ QUEUE_WORK = os.environ.get("SOAK_QUEUE_WORK", "awa_soak_py_work") QUEUE_DEADLINE = os.environ.get("SOAK_QUEUE_DEADLINE", "awa_soak_py_deadline") QUEUE_SHARED = os.environ.get("SOAK_QUEUE_SHARED", "awa_soak_py_shared") +# The harness passes the timings every Rust runtime uses, so a Python maintenance +# leader rescues on the same staleness. +HEARTBEAT_MS = int(os.environ.get("SOAK_HEARTBEAT_MS", "2000")) +STALENESS_MS = int(os.environ.get("SOAK_STALENESS_MS", "12000")) +RESCUE_INTERVAL_MS = os.environ.get("SOAK_RESCUE_INTERVAL_MS") WORK_WORKERS = 8 DEADLINE_WORKERS = 16 SHARED_WORKERS = 6 @@ -325,10 +330,10 @@ async def main(): {"name": QUEUE_DEADLINE, "max_workers": DEADLINE_WORKERS, "deadline_duration_ms": DEADLINE_MS}, {"name": QUEUE_SHARED, "max_workers": SHARED_WORKERS}, ], - heartbeat_interval_ms=2000, - heartbeat_staleness_ms=12000, - heartbeat_rescue_interval_ms=3000, - deadline_rescue_interval_ms=2000, + heartbeat_interval_ms=HEARTBEAT_MS, + heartbeat_staleness_ms=STALENESS_MS, + heartbeat_rescue_interval_ms=int(RESCUE_INTERVAL_MS or 3000), + deadline_rescue_interval_ms=int(RESCUE_INTERVAL_MS or 2000), promote_interval_ms=250, ) # The deadline queue serves a second kind through the same handler. diff --git a/awa/tests/persona_soak/python/test_soak_python.py b/awa/tests/persona_soak/python/test_soak_python.py index 337301719..9736c4c27 100644 --- a/awa/tests/persona_soak/python/test_soak_python.py +++ b/awa/tests/persona_soak/python/test_soak_python.py @@ -1,5 +1,6 @@ """Process-boundary regressions; uses a delayed fake client, never a database.""" +import json import os from pathlib import Path import select @@ -14,6 +15,7 @@ SCRIPT = Path(os.environ.get("SOAK_PYTHON_SCRIPT", Path(__file__).with_name("soak_python.py"))) FAKE_AWA = ''' import asyncio +import json import os from pathlib import Path from types import SimpleNamespace @@ -37,7 +39,8 @@ def task(self, *args, **kwargs): return lambda handler: handler async def start(self, *args, **kwargs): - Path(os.environ["SOAK_TEST_WORKERS"]).touch() + timings = {key: value for key, value in kwargs.items() if key.endswith("_ms")} + Path(os.environ["SOAK_TEST_WORKERS"]).write_text(json.dumps(timings)) async def insert(self, *args, **kwargs): Path(os.environ["SOAK_TEST_PRODUCED"]).touch() @@ -68,8 +71,8 @@ def tearDown(self): self.child.stderr.close() self.directory.cleanup() - def start(self, *, draining=False): - env = os.environ | { + def start(self, *, draining=False, extra_env=None): + env = os.environ | (extra_env or {}) | { "PYTHONPATH": str(self.root), "SOAK_DATABASE_URL": "unused-by-the-fake-client", "SOAK_ROLE": "full", @@ -114,6 +117,22 @@ def test_restart_during_drain_starts_workers_without_producers(self): self.start(draining=True) self.assert_draining_workers() + def test_workers_use_the_timings_the_harness_passes(self): + self.start(extra_env={ + "SOAK_HEARTBEAT_MS": "30000", + "SOAK_STALENESS_MS": "90000", + "SOAK_RESCUE_INTERVAL_MS": "7000", + }) + self.gate.touch() + self.wait_for("READY pid=") + self.assertEqual(json.loads(self.workers.read_text()), { + "heartbeat_interval_ms": 30000, + "heartbeat_staleness_ms": 90000, + "heartbeat_rescue_interval_ms": 7000, + "deadline_rescue_interval_ms": 7000, + "promote_interval_ms": 250, + }) + if __name__ == "__main__": unittest.main() diff --git a/awa/tests/persona_soak/registry.rs b/awa/tests/persona_soak/registry.rs index a0cc42835..2eca256f3 100644 --- a/awa/tests/persona_soak/registry.rs +++ b/awa/tests/persona_soak/registry.rs @@ -78,13 +78,17 @@ impl QueueRegistry { /// Looks a queue up by its physical name; striped queues (`name#N`) resolve to /// their logical queue. pub fn get(&self, name: &str) -> Option> { - let by_name = self.by_name.read().expect("registry lock"); - let index = by_name.get(name).or_else(|| { - let (logical, suffix) = name.rsplit_once('#')?; - suffix.parse::().ok()?; - by_name.get(logical) - })?; - Some(self.queues.read().expect("registry lock")[*index as usize].clone()) + // `register` takes `queues` before `by_name`; holding `by_name` while + // waiting for `queues` would deadlock against it. + let index = { + let by_name = self.by_name.read().expect("registry lock"); + *by_name.get(name).or_else(|| { + let (logical, suffix) = name.rsplit_once('#')?; + suffix.parse::().ok()?; + by_name.get(logical) + })? + }; + Some(self.queues.read().expect("registry lock")[index as usize].clone()) } pub fn by_index(&self, index: u32) -> Option> { @@ -129,3 +133,45 @@ impl Drop for InFlightGuard { } } } + +#[cfg(test)] +mod tests { + use super::QueueRegistry; + use crate::persona_soak::config::Persona; + use std::sync::mpsc; + use std::sync::Arc; + use std::time::Duration; + + #[test] + fn lookups_and_registrations_do_not_deadlock() { + let registry = Arc::new(QueueRegistry::default()); + registry.register("awa_soak_registry_0", Persona::Mia, false); + let (done_tx, done_rx) = mpsc::channel(); + let registering = { + let registry = registry.clone(); + let done_tx = done_tx.clone(); + std::thread::spawn(move || { + for index in 1..20_000 { + registry.register(&format!("awa_soak_registry_{index}"), Persona::Mia, false); + } + let _ = done_tx.send(()); + }) + }; + let looking_up = { + let registry = registry.clone(); + std::thread::spawn(move || { + for _ in 0..200_000 { + assert!(registry.get("awa_soak_registry_0#1").is_some()); + } + let _ = done_tx.send(()); + }) + }; + for _ in 0..2 { + done_rx + .recv_timeout(Duration::from_secs(30)) + .expect("register and get finish without deadlocking"); + } + registering.join().unwrap(); + looking_up.join().unwrap(); + } +} diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index 36f622613..b89177f18 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -24,7 +24,7 @@ use opentelemetry_sdk::metrics::{InMemoryMetricExporter, SdkMeterProvider}; use serde::{Deserialize, Serialize}; use sqlx::postgres::PgPoolOptions; use std::collections::{HashMap, HashSet}; -use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; +use std::sync::atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering}; use std::sync::Arc; use std::time::{Duration, Instant}; use tokio::sync::{Mutex, Notify}; @@ -2274,6 +2274,289 @@ async fn test_queue_storage_graceful_shutdown_finalizes_returned_attempts() { ); } +#[derive(Debug, Serialize, Deserialize, JobArgs)] +struct FinalizeBacklogJob { + n: i64, +} + +/// Stalling durable completion must backpressure claiming rather than grow +/// detached tasks holding job payloads. Release the locks and verify that +/// every job completes once. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_runtime_stops_claiming_while_finalization_is_stalled() { + assert_claiming_stops_while_finalization_is_stalled( + "qs_finalize_backlog", + "awa_qs_finalize_backlog", + true, + ) + .await; +} + +/// The same bound holds when completion deletes materialized leases and +/// writes terminal rows instead of closing receipt claims. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_lease_runtime_stops_claiming_while_finalization_is_stalled() { + assert_claiming_stops_while_finalization_is_stalled( + "qs_finalize_backlog_leases", + "awa_qs_finalize_backlog_leases", + false, + ) + .await; +} + +/// Lifecycle hooks run under the finalizer bound: while a hook does not +/// return, claiming stops at that bound instead of accumulating hook tasks, +/// and every job completes once the hooks return. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_slow_lifecycle_hook_holds_back_claiming() { + const FINALIZER_CAPACITY: usize = 1_024; + const WORKERS: u32 = 64; + const JOBS: i64 = 1_600; + let queue = "qs_slow_lifecycle_hook"; + let (_db_guard, pool) = setup_pool(10).await; + let store_config = QueueStorageConfig { + schema: "awa_qs_slow_lifecycle_hook".to_string(), + ..Default::default() + }; + let store = create_store_with_config(&pool, store_config.clone()).await; + for first in (0..JOBS).step_by(500) { + let params: Vec<_> = (first..(first + 500).min(JOBS)) + .map(|n| { + insert::params_with( + &FinalizeBacklogJob { n }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .expect("params") + }) + .collect(); + store + .enqueue_params_batch(&pool, ¶ms) + .await + .expect("enqueue backlog"); + } + + let (open_hooks, hooks_open) = tokio::sync::watch::channel(false); + let runs = Arc::new(AtomicUsize::new(0)); + let recorded = runs.clone(); + let client = Client::builder(pool.clone()) + .queue( + queue, + QueueConfig { + max_workers: WORKERS, + poll_interval: Duration::from_millis(10), + deadline_duration: Duration::ZERO, + claim_batch_size: WORKERS as usize, + ..QueueConfig::default() + }, + ) + .queue_storage( + store_config, + Duration::from_secs(60), + Duration::from_millis(50), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .register::(move |_, _| { + recorded.fetch_add(1, Ordering::SeqCst); + async { Ok(JobResult::Completed) } + }) + .on_event::(move |_| { + let mut hooks_open = hooks_open.clone(); + async move { + let _ = hooks_open.wait_for(|open| *open).await; + } + }) + .heartbeat_staleness(Duration::from_secs(300)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .build() + .expect("build slow hook client"); + client.start().await.expect("start slow hook client"); + + let limit = FINALIZER_CAPACITY + WORKERS as usize; + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(30)); + let mut last = (0, Instant::now()); + let started_while_blocked = loop { + let started = runs.load(Ordering::SeqCst); + if started > limit { + break started; + } + if started != last.0 { + last = (started, Instant::now()); + } else if started > 0 && last.1.elapsed() >= Duration::from_secs(5) { + break started; + } + assert!( + Instant::now() < deadline, + "claiming neither stopped nor passed {limit}" + ); + tokio::time::sleep(Duration::from_millis(25)).await; + }; + + open_hooks.send(true).expect("open the hooks"); + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(60)); + while completed_terminal_count(&pool, &store, queue).await < JOBS { + assert!(Instant::now() < deadline, "the backlog did not drain"); + tokio::time::sleep(Duration::from_millis(50)).await; + } + client.shutdown(Duration::from_secs(10)).await; + + assert!( + started_while_blocked <= limit, + "{started_while_blocked} jobs started while lifecycle hooks were blocked; \ + at most {limit} may occupy finalizer slots or hold a worker" + ); + assert_eq!(runs.load(Ordering::SeqCst) as i64, JOBS); +} + +async fn assert_claiming_stops_while_finalization_is_stalled( + queue: &str, + schema: &str, + lease_claim_receipts: bool, +) { + // The default completion pipeline: one shard of 512-job batches, one + // batch in flight and one queued. + const FINALIZER_CAPACITY: usize = 1_024; + const WORKERS: u32 = 64; + const JOBS: i64 = 1_600; + let (_db_guard, pool) = setup_pool(10).await; + let store_config = QueueStorageConfig { + schema: schema.to_string(), + lease_claim_receipts, + ..Default::default() + }; + let store = create_store_with_config(&pool, store_config.clone()).await; + for first in (0..JOBS).step_by(500) { + let params: Vec<_> = (first..(first + 500).min(JOBS)) + .map(|n| { + insert::params_with( + &FinalizeBacklogJob { n }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .expect("params") + }) + .collect(); + store + .enqueue_params_batch(&pool, ¶ms) + .await + .expect("enqueue backlog"); + } + + let mut stall = pool.begin().await.expect("begin stall transaction"); + if lease_claim_receipts { + let first_jobs: Vec = sqlx::query_scalar(audited_sql(format!( + "SELECT job_id FROM {}.ready_entries WHERE queue = $1 ORDER BY lane_seq LIMIT $2", + store.schema() + ))) + .bind(queue) + .bind(i64::from(WORKERS) * 2) + .fetch_all(&pool) + .await + .expect("first jobs"); + sqlx::query( + "SELECT pg_catalog.pg_advisory_xact_lock(pg_catalog.hashtextextended( + format('awa.receipt.complete:%s:%s', job_id, 1), 0)) + FROM unnest($1::bigint[]) AS job_id", + ) + .bind(&first_jobs) + .execute(stall.as_mut()) + .await + .expect("stall the first jobs' finalization"); + } else { + sqlx::query(audited_sql(format!( + "LOCK TABLE {}.done_entries IN SHARE MODE", + store.schema() + ))) + .execute(stall.as_mut()) + .await + .expect("stall terminal writes"); + } + + let runs: Arc>> = Arc::default(); + let recorded = runs.clone(); + let client = Client::builder(pool.clone()) + .queue( + queue, + QueueConfig { + max_workers: WORKERS, + poll_interval: Duration::from_millis(10), + deadline_duration: Duration::ZERO, + claim_batch_size: WORKERS as usize, + ..QueueConfig::default() + }, + ) + .queue_storage( + store_config, + Duration::from_secs(60), + Duration::from_millis(50), + ) + .claim_rotate_interval(Duration::from_secs(60)) + .register::(move |_, ctx| { + *recorded + .lock() + .expect("runs") + .entry(ctx.job.id) + .or_insert(0) += 1; + async { Ok(JobResult::Completed) } + }) + .heartbeat_staleness(Duration::from_secs(300)) + .leader_election_interval(Duration::from_millis(100)) + .leader_check_interval(Duration::from_millis(50)) + .build() + .expect("build finalize backlog client"); + client.start().await.expect("start finalize backlog client"); + + // Wait for claiming to stop, or to run past the bound. A fixed runtime + // can never pass it, so a pause in claiming only weakens the check. + let limit = FINALIZER_CAPACITY + WORKERS as usize; + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(30)); + let mut last = (0, Instant::now()); + let started_while_stalled = loop { + let started = runs.lock().expect("runs").len(); + if started > limit { + break started; + } + if started != last.0 { + last = (started, Instant::now()); + } else if started > 0 && last.1.elapsed() >= Duration::from_secs(5) { + break started; + } + assert!( + Instant::now() < deadline, + "claiming neither stopped nor passed {limit}" + ); + tokio::time::sleep(Duration::from_millis(25)).await; + }; + + stall + .rollback() + .await + .expect("release the stalled finalization"); + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(60)); + while completed_terminal_count(&pool, &store, queue).await < JOBS { + assert!(Instant::now() < deadline, "the backlog did not drain"); + tokio::time::sleep(Duration::from_millis(50)).await; + } + client.shutdown(Duration::from_secs(10)).await; + + assert!( + started_while_stalled <= limit, + "{started_while_stalled} jobs started while finalization was stalled; \ + at most {limit} may occupy finalizer slots or hold a worker" + ); + let runs = runs.lock().expect("runs"); + assert_eq!(runs.len() as i64, JOBS); + assert!( + runs.values().all(|count| *count == 1), + "a job ran more than once" + ); +} + /// ADR-023 claim-ring control-plane smoke test. Exercises /// `rotate_claims`, the #371 idle gate, and `prune_oldest_claims` on an /// empty schema: rotate is idle-skipped when nothing has been written @@ -11659,7 +11942,7 @@ async fn test_queue_storage_promotion_leaves_a_contended_lane_deferred_and_in_or /// Holds PostgreSQL's cluster-wide notify-queue lock for about `commits` × /// 100 ms. A transaction that issued NOTIFY takes that lock at COMMIT and /// keeps it until its commit record is flushed, which `commit_delay` stretches. -/// Returns once the lock is held. +/// Returns once the holder's own backend holds the lock. async fn hold_notify_commit_lock( pool: &sqlx::PgPool, commits: usize, @@ -11678,6 +11961,10 @@ async fn hold_notify_commit_lock( ) .await .expect("connect the notify holder"); + let holder_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&holder) + .await + .expect("read the notify holder's backend"); let holder_task = tokio::spawn(async move { for _ in 0..commits { sqlx::raw_sql( @@ -11694,8 +11981,9 @@ async fn hold_notify_commit_lock( loop { let held: bool = sqlx::query_scalar( "SELECT EXISTS (SELECT 1 FROM pg_locks WHERE locktype = 'object' \ - AND classid = 'pg_database'::regclass AND objid = 0 AND granted)", + AND classid = 'pg_database'::regclass AND objid = 0 AND granted AND pid = $1)", ) + .bind(holder_pid) .fetch_one(pool) .await .expect("read notify lock"); @@ -11746,7 +12034,7 @@ async fn test_queue_storage_promotion_does_not_wait_for_the_notify_lock_in_its_l enqueue_due(id).await; } - let holder = hold_notify_commit_lock(&pool, 3).await; + let holder = hold_notify_commit_lock(&pool, 10).await; let outcome = store .promote_due_lanes( &pool, @@ -11757,6 +12045,10 @@ async fn test_queue_storage_promotion_does_not_wait_for_the_notify_lock_in_its_l ) .await .expect("promote while the notify lock is held"); + assert!( + !holder.is_finished(), + "the notify holder outlasts the promotion" + ); holder.await.expect("notify holder"); assert_eq!(outcome.promoted, 3, "{outcome:?}"); assert!(outcome.contended.is_empty(), "{outcome:?}"); @@ -11772,7 +12064,7 @@ async fn test_queue_storage_promotion_does_not_wait_for_the_notify_lock_in_its_l .await .expect("connect with a session lock_timeout"); enqueue_due(3).await; - let holder = hold_notify_commit_lock(&pool, 3).await; + let holder = hold_notify_commit_lock(&pool, 10).await; let result = store .promote_due_lanes( &caller_lock_timeout_pool, @@ -11782,6 +12074,10 @@ async fn test_queue_storage_promotion_does_not_wait_for_the_notify_lock_in_its_l &[], ) .await; + assert!( + !holder.is_finished(), + "the notify holder outlasts the promotion" + ); holder.await.expect("notify holder"); caller_lock_timeout_pool.close().await; let outcome = result.expect("a failed wakeup does not fail the promotion"); @@ -22951,6 +23247,55 @@ async fn receipt_admin_ids(pool: &sqlx::PgPool, filter: admin::ListJobsFilter) - .collect() } +/// A per-slot receipt read touches only that slot's claim-ring partitions, so a +/// lock held on another slot's closure evidence, as claim-ring prune takes, does +/// not hold it up. +#[tokio::test] +async fn test_slot_scoped_receipt_reads_do_not_wait_on_other_slots_closures() { + let (_db_guard, pool) = setup_pool(4).await; + let schema = "awa_qs_slot_scoped_closures"; + let store = create_store_with_config(&pool, unique_cancel_store_config(schema)).await; + let queue = "slot_scoped_closures"; + enqueue_job( + &pool, + &store, + &CompleteJob { id: 1 }, + InsertOpts { + queue: queue.into(), + ..Default::default() + }, + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::ZERO) + .await + .expect("claim into a compact batch"); + assert_eq!(claimed.len(), 1); + let other_slot = (claimed[0].claim.claim_slot + 1) % 2; + + let mut prune = pool.begin().await.expect("begin prune stand-in"); + sqlx::query(audited_sql(format!( + "LOCK TABLE {schema}.lease_claim_closures_{other_slot}, \ + {schema}.lease_claim_closure_batches_{other_slot} IN ACCESS EXCLUSIVE MODE" + ))) + .execute(prune.as_mut()) + .await + .expect("lock the other slot's closure evidence"); + let kinds = tokio::time::timeout( + scaled_timeout(Duration::from_secs(5)), + admin::job_kind_overviews(&pool), + ) + .await + .expect("the read does not wait on another slot's closure locks") + .expect("read kind overviews"); + prune.rollback().await.expect("release the closure locks"); + let complete = kinds + .iter() + .find(|kind| kind.kind == CompleteJob::kind()) + .expect("the running job's kind"); + assert_eq!(complete.job_count, 1); +} + async fn receipt_admin_reads_preserve_scope_and_disposition(receipts: bool, row_local: bool) { let (_db_guard, pool) = setup_pool(6).await; let schema = match (receipts, row_local) { @@ -23226,6 +23571,69 @@ async fn test_receipt_admin_scope_materialized_control() { receipt_admin_reads_preserve_scope_and_disposition(false, false).await; } +#[tokio::test] +async fn test_schema_only_admin_reads_claims_in_installed_slot() { + let (_db_guard, pool) = setup_pool(6).await; + let schema = "awa_qs_admin_wide_claims"; + let queue = "admin_wide_claims"; + let default_claim_slots = QueueStorageConfig::default().claim_slot_count; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + claim_slot_count: default_claim_slots + 2, + ..unique_cancel_store_config(schema) + }, + ) + .await; + // An empty ring idle-skips rotation, so append the cursor move directly. + sqlx::query(audited_sql(format!( + "INSERT INTO {schema}.claim_ring_rotations (generation, slot) VALUES (1, $1)" + ))) + .bind(default_claim_slots as i32) + .execute(&pool) + .await + .expect("move the claim ring cursor past the default slot count"); + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 1 }, + InsertOpts { + queue: queue.into(), + ..Default::default() + }, + ) + .await; + let claims = store + .claim_runtime_batch(&pool, queue, 1, Duration::ZERO) + .await + .unwrap(); + assert_eq!(claims.len(), 1); + assert_eq!(claims[0].claim.claim_slot, default_claim_slots as i32); + + let running = admin::list_jobs( + &pool, + &admin::ListJobsFilter { + state: Some(JobState::Running), + queue: Some(queue.into()), + ..Default::default() + }, + ) + .await + .expect("list running jobs through a schema-only store"); + assert_eq!( + running.iter().map(|job| job.id).collect::>(), + vec![job_id] + ); + let kinds = admin::job_kind_overviews(&pool) + .await + .expect("count kinds through a schema-only store"); + let complete = kinds + .iter() + .find(|kind| kind.kind == CompleteJob::kind()) + .unwrap(); + assert_eq!((complete.job_count, complete.queue_count), (1, 1)); +} + #[tokio::test] async fn test_receipt_list_refills_pages_past_closed_and_other_kind_claims() { let (_guard, pool) = setup_pool(6).await; @@ -23367,6 +23775,75 @@ async fn test_receipt_list_refills_pages_past_closed_and_other_kind_claims() { all_ids, "closed receipts and materialized claims must not duplicate results" ); + + let missing_ids: Vec = (1..=13).map(|offset| all_ids[0] + offset).collect(); + let claim_slot: i32 = sqlx::query_scalar(audited_sql(format!( + "SELECT slot FROM {}.claim_ring_rotations ORDER BY generation DESC LIMIT 1", + store.schema() + ))) + .fetch_one(&pool) + .await + .expect("read current receipt slot"); + sqlx::query(audited_sql(format!( + "INSERT INTO {}.lease_claims (claim_slot, job_id, run_lease, ready_slot, \ + ready_generation, queue, priority, attempt, max_attempts, lane_seq) \ + SELECT $1, job_id, 1, 0, 0, $2, 2, 1, 25, job_id \ + FROM unnest($3::bigint[]) AS missing(job_id)", + store.schema() + ))) + .bind(claim_slot) + .bind(queue) + .bind(&missing_ids) + .execute(&pool) + .await + .expect("seed newer receipts without any job body"); + for job_id in &missing_ids { + assert!( + store + .load_job(&pool, *job_id) + .await + .expect("missing-body receipt must not suppress a hydration error") + .is_none(), + "synthetic receipt {job_id} unexpectedly has a job body" + ); + } + let expected_running: Vec = all_ids + .iter() + .copied() + .filter(|job_id| !finished.iter().any(|claim| claim.job.id == *job_id)) + .collect(); + let no_kind = admin::ListJobsFilter { + state: Some(JobState::Running), + queue: Some(queue.into()), + limit: Some(2), + ..Default::default() + }; + let first = receipt_admin_ids(&pool, no_kind.clone()).await; + assert_eq!(first, expected_running[..2]); + let second = receipt_admin_ids( + &pool, + admin::ListJobsFilter { + before_id: first.last().copied(), + ..no_kind + }, + ) + .await; + assert_eq!(second, expected_running[2..4]); + assert!(first.iter().all(|id| !second.contains(id))); + assert_eq!( + receipt_admin_ids( + &pool, + admin::ListJobsFilter { + state: Some(JobState::Running), + kind: Some(CompleteJob::kind().into()), + queue: Some(queue.into()), + limit: Some(2), + ..Default::default() + }, + ) + .await, + targets[..2] + ); } /// A point read's snapshot can outlive completion. TRUNCATE is not MVCC-safe: @@ -25290,11 +25767,88 @@ async fn test_historical_tombstoned_claim_uses_emitted_lane_low_to_high() { async fn test_historical_tombstoned_claim_uses_emitted_lane_high_to_low() { moved_tombstoned_claim_lane(4, 2, true, false, false).await; } +/// An attempt whose retained ready bodies are ambiguous stays running. A batch's +/// worth of them, which every pass meets first, must not keep rescue from the stale +/// attempts behind them. #[tokio::test] -async fn test_historical_tombstoned_legacy_ambiguous_lane_is_skipped_by_rescue() { - moved_tombstoned_claim_lane(2, 4, false, false, false).await; -} - +async fn test_ambiguous_attempts_filling_a_batch_do_not_starve_rescue() { + const AMBIGUOUS: i64 = 500; + let (_guard, pool) = setup_pool(4).await; + let schema = "awa_qs_ambiguous_backlog"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.into(), + lease_claim_receipts: false, + queue_stripe_count: 1, + ..Default::default() + }, + ) + .await; + let queue = "qs_ambiguous_backlog"; + store + .enqueue_batch(&pool, queue, 2, AMBIGUOUS) + .await + .expect("enqueue the ambiguous backlog"); + let claimed = store + .claim_runtime_batch(&pool, queue, AMBIGUOUS, Duration::from_secs(300)) + .await + .expect("claim the ambiguous backlog"); + assert_eq!(claimed.len(), AMBIGUOUS as usize); + // A second body at the same lane position under another priority, as older + // physical aging could leave, leaves no way to choose each attempt's body. + let duplicated = sqlx::query(audited_sql(format!( + "INSERT INTO {schema}.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 ready_slot,ready_generation,job_id,kind,queue,args,priority-1,attempt,run_lease, + max_attempts,lane_seq,enqueue_shard,run_at,attempted_at,created_at,unique_key, + unique_states,payload + FROM {schema}.ready_entries WHERE queue=$1" + ))) + .bind(queue) + .execute(&pool) + .await + .expect("duplicate the retained bodies"); + assert_eq!(duplicated.rows_affected(), AMBIGUOUS as u64); + + let other_queue = "qs_ambiguous_backlog_other"; + let other_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9045 }, + InsertOpts { + queue: other_queue.into(), + ..Default::default() + }, + ) + .await; + let other_claims = store + .claim_runtime_batch(&pool, other_queue, 1, Duration::from_secs(300)) + .await + .expect("claim a recoverable job after the backlog"); + assert_eq!(other_claims.len(), 1); + + let rescued = store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .expect("ambiguous attempts do not stop rescue of the others"); + assert_eq!( + rescued.iter().map(|job| job.id).collect::>(), + vec![other_id] + ); + assert_eq!( + lease_count(&pool, &store).await, + AMBIGUOUS, + "the ambiguous attempts stay running" + ); +} + +#[tokio::test] +async fn test_historical_tombstoned_legacy_ambiguous_lane_is_skipped_by_rescue() { + moved_tombstoned_claim_lane(2, 4, false, false, false).await; +} + #[tokio::test] async fn test_historical_tombstoned_receipt_materialized_then_reopened_legacy_uses_evidence() { moved_tombstoned_claim_lane(2, 4, true, true, false).await; @@ -26326,3 +26880,654 @@ async fn test_busy_retained_head_cancel_preserves_lock_not_available() { async fn test_busy_retained_head_does_not_block_materialized_heartbeat_rescue() { historical_busy_attempt_does_not_block_batch(RetainedBatchTransition::Heartbeat, true).await; } + +/// Claim slots of `schema` whose `lease_claim_batches` partition carries the +/// GIN index on `job_ids`. +async fn claim_batch_job_index_slots(pool: &sqlx::PgPool, schema: &str) -> Vec { + sqlx::query_scalar( + r#" + SELECT substring(index_class.relname FROM '^lease_claim_batches_([0-9]+)_job_ids_idx$')::int + FROM pg_class AS index_class + JOIN pg_index AS index_meta ON index_meta.indexrelid = index_class.oid + JOIN pg_am AS method ON method.oid = index_class.relam + WHERE index_class.relnamespace = to_regnamespace($1) + AND index_class.relname LIKE 'lease\_claim\_batches\_%\_job\_ids\_idx' + AND method.amname = 'gin' + ORDER BY 1 + "#, + ) + .bind(schema) + .fetch_all(pool) + .await + .expect("Failed to list claim batch job_ids indexes") +} + +async fn drop_claim_batch_job_index(pool: &sqlx::PgPool, schema: &str, slot: i32) { + sqlx::query(audited_sql(format!( + "DROP INDEX {schema}.lease_claim_batches_{slot}_job_ids_idx" + ))) + .execute(pool) + .await + .expect("Failed to drop claim batch job_ids index"); +} + +async fn ensure_claim_batch_job_index<'e, E>(executor: E, schema: &str, slot: i32) -> bool +where + E: sqlx::PgExecutor<'e>, +{ + sqlx::query_scalar(audited_sql(format!( + "SELECT {schema}.ensure_claim_batch_job_index($1)" + ))) + .bind(slot) + .fetch_one(executor) + .await + .expect("ensure_claim_batch_job_index should not fail") +} + +/// Index scans PostgreSQL has recorded on one claim partition's job_ids index. +async fn claim_batch_job_index_scans(pool: &sqlx::PgPool, schema: &str, slot: i32) -> i64 { + sqlx::query_scalar( + "SELECT idx_scan FROM pg_stat_user_indexes WHERE schemaname = $1 AND indexrelname = $2", + ) + .bind(schema) + .bind(format!("lease_claim_batches_{slot}_job_ids_idx")) + .fetch_one(pool) + .await + .expect("Failed to read claim batch job_ids index statistics") +} + +/// Run `work` on its own connection with sequential scans disabled, close +/// that connection so its statistics are flushed, and return how many scans +/// the job_ids index of `slot` gained. Disabling sequential scans makes a +/// lookup the index can serve use it even on a tiny table; a lookup it cannot +/// serve still scans the table and gains nothing. +async fn claim_batch_job_index_scans_during( + pool: &sqlx::PgPool, + schema: &str, + slot: i32, + work: F, +) -> i64 +where + F: FnOnce(sqlx::PgPool) -> Fut, + Fut: std::future::Future, +{ + let before = claim_batch_job_index_scans(pool, schema, slot).await; + let no_seqscan = PgPoolOptions::new() + .max_connections(1) + .connect_with( + (*pool.connect_options()) + .clone() + .options([("enable_seqscan", "off")]), + ) + .await + .expect("Failed to connect with sequential scans disabled"); + work(no_seqscan.clone()).await; + no_seqscan.close().await; + + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(5)); + loop { + let gained = claim_batch_job_index_scans(pool, schema, slot).await - before; + if gained > 0 || Instant::now() > deadline { + return gained; + } + tokio::time::sleep(Duration::from_millis(25)).await; + } +} + +/// The by-job lookups of compact claims (`load_job`, receipt heartbeat) +/// probe `lease_claim_batches` through the job_ids index and return the same +/// rows they did without it. +#[tokio::test] +async fn test_claim_batch_job_index_serves_compact_receipt_lookups() { + let (_db_guard, pool) = setup_pool(4).await; + let queue = "qs_claim_batch_job_index_lookup"; + let schema = "awa_qs_cbji_lookup"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + queue_stripe_count: 1, + lease_claim_receipts: true, + claim_slot_count: 2, + }, + ) + .await; + assert_eq!( + claim_batch_job_index_slots(&pool, schema).await, + vec![0, 1], + "the installer indexes every empty claim partition" + ); + + let mut job_ids = Vec::new(); + for id in 0..3 { + job_ids.push( + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await, + ); + } + let claimed = store + .claim_runtime_batch(&pool, queue, 3, Duration::ZERO) + .await + .expect("compact claim of zero-deadline jobs"); + assert_eq!(claimed.len(), 3); + assert!(claimed.iter().all(|job| job.claim.lease_claim_receipt)); + assert_eq!(lease_claim_count(&pool, &store).await, 0); + assert_eq!(lease_claim_batch_count(&pool, &store).await, 1); + let probe = &claimed[1]; + let store_ref = &store; + + let load_scans = claim_batch_job_index_scans_during(&pool, schema, 0, |indexed| async move { + let job = store_ref + .load_job(&indexed, probe.job.id) + .await + .expect("load_job of a compact claim") + .expect("compact claim is visible to load_job"); + assert_eq!(job.state, JobState::Running); + assert_eq!(job.run_lease, probe.job.run_lease); + }) + .await; + assert!(load_scans > 0, "load_job must probe the job_ids index"); + + let heartbeat_scans = + claim_batch_job_index_scans_during(&pool, schema, 0, |indexed| async move { + let updated = store_ref + .heartbeat_batch(&indexed, &[(probe.job.id, probe.job.run_lease)]) + .await + .expect("heartbeat of a compact claim"); + assert_eq!(updated, 1); + }) + .await; + assert!( + heartbeat_scans > 0, + "receipt heartbeat must probe the job_ids index" + ); + assert_eq!(attempt_state_count(&pool, &store).await, 1); + + let completed = store + .complete_runtime_batch(&pool, &claimed) + .await + .expect("complete compact claims"); + assert_eq!(completed.len(), 3); + for job_id in job_ids { + let job = store + .load_job(&pool, job_id) + .await + .expect("load completed job") + .expect("completed job is visible"); + assert_eq!(job.state, JobState::Completed); + } + assert_eq!(open_receipt_claim_count(&pool, &store).await, 0); +} + +/// `ensure_claim_batch_job_index` indexes an empty claim partition, leaves a +/// populated or busy one alone without waiting or keeping a lock, and claim +/// prune indexes a slot as it truncates it. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_claim_batch_job_index_is_built_only_on_empty_partitions() { + let (_db_guard, pool) = setup_pool(4).await; + let queue = "qs_claim_batch_job_index_build"; + let schema = "awa_qs_cbji_build"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + queue_stripe_count: 1, + lease_claim_receipts: true, + claim_slot_count: 2, + }, + ) + .await; + + let function = format!("{schema}.ensure_claim_batch_job_index(integer)"); + let (definer, config, public_execute): (bool, Vec, bool) = sqlx::query_as( + r#" + SELECT routine.prosecdef, + routine.proconfig, + has_function_privilege('public', routine.oid, 'EXECUTE') + FROM pg_proc AS routine + WHERE routine.oid = $1::regprocedure + "#, + ) + .bind(&function) + .fetch_one(&pool) + .await + .expect("Failed to read ensure_claim_batch_job_index"); + assert!(definer); + assert_eq!( + config, + vec![ + "search_path=pg_catalog, pg_temp".to_string(), + "quote_all_identifiers=off".to_string(), + ] + ); + assert!(!public_execute, "PUBLIC must not execute {function}"); + + let grantee = "awa_qs_cbji_truncate_holder"; + let grantee_executes = |pool: sqlx::PgPool, function: String| async move { + sqlx::query_scalar::<_, bool>( + "SELECT has_function_privilege($1, $2::regprocedure, 'EXECUTE')", + ) + .bind(grantee) + .bind(&function) + .fetch_one(&pool) + .await + .expect("read grantee EXECUTE") + }; + let refresh_helper = |pool: sqlx::PgPool| async move { + sqlx::query("SELECT awa.refresh_claim_batch_job_index_helper($1)") + .bind(schema) + .execute(&pool) + .await + .expect("refresh the claim batch job index helper"); + }; + sqlx::raw_sql(audited_sql(format!( + "DROP OWNED BY {grantee}; DROP ROLE IF EXISTS {grantee}" + ))) + .execute(&pool) + .await + .ok(); + sqlx::raw_sql(audited_sql(format!( + "CREATE ROLE {grantee}; GRANT USAGE ON SCHEMA {schema} TO {grantee}; \ + GRANT TRUNCATE ON {schema}.lease_claim_batches TO {grantee}" + ))) + .execute(&pool) + .await + .expect("create a TRUNCATE holder"); + refresh_helper(pool.clone()).await; + assert!(grantee_executes(pool.clone(), function.clone()).await); + sqlx::raw_sql(audited_sql(format!( + "REVOKE TRUNCATE ON {schema}.lease_claim_batches FROM {grantee}" + ))) + .execute(&pool) + .await + .expect("revoke TRUNCATE"); + refresh_helper(pool.clone()).await; + let still_executes = grantee_executes(pool.clone(), function.clone()).await; + sqlx::raw_sql(audited_sql(format!( + "DROP OWNED BY {grantee}; DROP ROLE {grantee}" + ))) + .execute(&pool) + .await + .expect("drop the TRUNCATE holder"); + assert!( + !still_executes, + "a role that lost TRUNCATE must lose EXECUTE on {function} when the helper is refreshed" + ); + + enqueue_job( + &pool, + &store, + &CompleteJob { id: 1 }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::ZERO) + .await + .expect("compact claim of a zero-deadline job"); + assert_eq!(claimed.len(), 1); + assert_eq!( + receipt_claim_count_in_child(&pool, schema, 0, claimed[0].job.id).await, + 1 + ); + drop_claim_batch_job_index(&pool, schema, 0).await; + drop_claim_batch_job_index(&pool, schema, 1).await; + + // Populated: no index, and no lock left on the partition afterwards. + let mut tx = pool.begin().await.unwrap(); + assert!(!ensure_claim_batch_job_index(&mut *tx, schema, 0).await); + let partition_locks: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pg_locks WHERE pid = pg_backend_pid() AND relation = $1::regclass", + ) + .bind(format!("{schema}.lease_claim_batches_0")) + .fetch_one(&mut *tx) + .await + .unwrap(); + assert_eq!( + partition_locks, 0, + "a populated partition must not stay locked" + ); + tx.rollback().await.unwrap(); + + // Empty but locked by a writer: give up at once rather than queue. + let mut writer = pool.begin().await.unwrap(); + sqlx::query(audited_sql(format!( + "LOCK TABLE {schema}.lease_claim_batches_1 IN ROW EXCLUSIVE MODE" + ))) + .execute(&mut *writer) + .await + .unwrap(); + let busy = tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(5)), + ensure_claim_batch_job_index(&pool, schema, 1), + ) + .await + .expect("ensure_claim_batch_job_index must not wait for a busy partition"); + assert!(!busy); + writer.rollback().await.unwrap(); + assert_eq!( + claim_batch_job_index_slots(&pool, schema).await, + Vec::::new() + ); + + // Empty and free: indexed, and a second call is a no-op. + assert!(ensure_claim_batch_job_index(&pool, schema, 1).await); + assert!(ensure_claim_batch_job_index(&pool, schema, 1).await); + assert_eq!(claim_batch_job_index_slots(&pool, schema).await, vec![1]); + + // Prune truncates slot 0 and indexes it in the same transaction. + store + .complete_runtime_batch(&pool, &claimed) + .await + .expect("complete compact claim"); + match store.rotate_claims(&pool).await.expect("rotate claim ring") { + RotateOutcome::Rotated { slot, .. } => assert_eq!(slot, 1), + other => panic!("expected claim ring rotation to slot 1, got {other:?}"), + } + let prune_deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(5)); + loop { + match store + .prune_oldest_claims(&pool) + .await + .expect("prune oldest claim slot") + { + PruneOutcome::Pruned { slot, .. } => { + assert_eq!(slot, 0); + break; + } + PruneOutcome::Blocked { slot: 0 } if Instant::now() < prune_deadline => {} + other => panic!("unexpected claim prune outcome: {other:?}"), + } + } + assert_eq!(claim_batch_job_index_slots(&pool, schema).await, vec![0, 1]); +} + +/// v054 is idempotent, installs the index function without building any +/// index, and claim prune then indexes each slot it truncates. A schema +/// without the function keeps working with unindexed lookups. +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn test_claim_batch_job_index_migration_leaves_indexing_to_claim_prune() { + let (_db_guard, pool) = setup_pool(4).await; + let queue = "qs_claim_batch_job_index_migration"; + let schema = "awa_qs_cbji_migration"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.to_string(), + queue_slot_count: 4, + lease_slot_count: 2, + queue_stripe_count: 1, + lease_claim_receipts: true, + claim_slot_count: 2, + }, + ) + .await; + + // The schema as a release before v054 left it. + sqlx::raw_sql( + r#" + DO $$ + DECLARE v_object TEXT; + BEGIN + FOR v_object IN + SELECT routine.oid::regprocedure::text FROM pg_proc AS routine + WHERE routine.proname = 'ensure_claim_batch_job_index' + LOOP + EXECUTE format('DROP FUNCTION %s', v_object); + END LOOP; + FOR v_object IN + SELECT index_class.oid::regclass::text FROM pg_class AS index_class + WHERE index_class.relkind = 'i' + AND index_class.relname LIKE 'lease\_claim\_batches\_%\_job\_ids\_idx' + LOOP + EXECUTE format('DROP INDEX %s', v_object); + END LOOP; + DROP FUNCTION awa.refresh_claim_batch_job_index_helper(text); + END + $$; + "#, + ) + .execute(&pool) + .await + .expect("Failed to remove the claim batch job index function"); + + let mut job_ids = Vec::new(); + for id in 0..2 { + job_ids.push( + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.to_string(), + ..Default::default() + }, + ) + .await, + ); + } + + // Without the function: claim, look up, heartbeat, complete and prune all work. + let first = store + .claim_runtime_batch(&pool, queue, 1, Duration::ZERO) + .await + .expect("compact claim without the index function"); + assert_eq!(first.len(), 1); + let running = store + .load_job(&pool, first[0].job.id) + .await + .unwrap() + .expect("compact claim visible without the index"); + assert_eq!(running.state, JobState::Running); + assert_eq!( + store + .heartbeat_batch(&pool, &[(first[0].job.id, first[0].job.run_lease)]) + .await + .unwrap(), + 1 + ); + store + .complete_runtime_batch(&pool, &first) + .await + .expect("complete compact claim without the index"); + match store.rotate_claims(&pool).await.expect("rotate claim ring") { + RotateOutcome::Rotated { slot, .. } => assert_eq!(slot, 1), + other => panic!("expected claim ring rotation to slot 1, got {other:?}"), + } + let second = store + .claim_runtime_batch(&pool, queue, 1, Duration::ZERO) + .await + .expect("compact claim into claim slot 1"); + assert_eq!(second.len(), 1); + assert_eq!( + receipt_claim_count_in_child(&pool, schema, 1, second[0].job.id).await, + 1 + ); + let prune_deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(5)); + loop { + match store + .prune_oldest_claims(&pool) + .await + .expect("prune without the index function") + { + PruneOutcome::Pruned { slot, .. } => { + assert_eq!(slot, 0); + break; + } + PruneOutcome::Blocked { slot: 0 } if Instant::now() < prune_deadline => {} + other => panic!("unexpected claim prune outcome: {other:?}"), + } + } + assert_eq!( + claim_batch_job_index_slots(&pool, schema).await, + Vec::::new() + ); + + sqlx::query("DELETE FROM awa.schema_version WHERE version >= 54") + .execute(&pool) + .await + .expect("roll the schema version back to v053"); + migrations::run(&pool).await.expect("apply v054"); + let installed: Vec = sqlx::query_scalar( + "SELECT routine.pronamespace::regnamespace::text FROM pg_proc AS routine \ + WHERE routine.proname = 'ensure_claim_batch_job_index' AND routine.prosecdef \ + ORDER BY 1", + ) + .fetch_all(&pool) + .await + .unwrap(); + assert_eq!(installed, vec!["awa".to_string(), schema.to_string()]); + assert_eq!( + claim_batch_job_index_slots(&pool, schema).await, + Vec::::new(), + "v054 builds no index, even on the empty slot" + ); + + let schema_versions = || async { + sqlx::query_as::<_, (i32, i64)>("SELECT max(version), count(*) FROM awa.schema_version") + .fetch_one(&pool) + .await + .unwrap() + }; + let versions_before = schema_versions().await; + let (_, _, v054_sql) = migrations::migration_sql_range(53, 54) + .pop() + .expect("v054 is a migration"); + sqlx::raw_sql(audited_sql(v054_sql)) + .execute(&pool) + .await + .expect("v054 re-applies cleanly"); + assert_eq!( + claim_batch_job_index_slots(&pool, schema).await, + Vec::::new() + ); + assert_eq!(schema_versions().await, versions_before); + + let job = store + .load_job(&pool, second[0].job.id) + .await + .unwrap() + .expect("claim in the unindexed slot stays visible"); + assert_eq!(job.state, JobState::Running); + store + .complete_runtime_batch(&pool, &second) + .await + .expect("complete the claim in the unindexed slot"); + match store.rotate_claims(&pool).await.expect("rotate claim ring") { + RotateOutcome::Rotated { slot, .. } => assert_eq!(slot, 0), + other => panic!("expected claim ring rotation to slot 0, got {other:?}"), + } + let prune_deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(5)); + loop { + match store + .prune_oldest_claims(&pool) + .await + .expect("prune with the index function") + { + PruneOutcome::Pruned { slot, .. } => { + assert_eq!(slot, 1); + break; + } + PruneOutcome::Blocked { slot: 1 } if Instant::now() < prune_deadline => {} + other => panic!("unexpected claim prune outcome: {other:?}"), + } + } + assert_eq!( + claim_batch_job_index_slots(&pool, schema).await, + vec![1], + "claim prune indexes the slot it truncates" + ); + for job_id in job_ids { + let job = store.load_job(&pool, job_id).await.unwrap().unwrap(); + assert_eq!(job.state, JobState::Completed); + } +} + +/// Metadata completions validate the exact compact item and do not lock the +/// immutable batch shared with other attempts. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_metadata_completion_validates_and_does_not_lock_compact_batch() { + let (_db_guard, pool) = setup_pool(10).await; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: "awa_qs_metadata_completion_identity".into(), + ..Default::default() + }, + ) + .await; + let queue = "metadata_completion_identity"; + for id in 0..2 { + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.into(), + tags: vec!["retained".into()], + metadata: serde_json::json!({"source": "recovery"}), + ..Default::default() + }, + ) + .await; + } + let jobs = store + .claim_runtime_batch(&pool, queue, 2, Duration::ZERO) + .await + .expect("claim compact jobs with metadata"); + assert_eq!(jobs.len(), 2); + assert_eq!(jobs[0].claim.claim_batch_id, jobs[1].claim.claim_batch_id); + let mut invalid = jobs[0].clone(); + invalid.claim.claim_batch_index = jobs[1].claim.claim_batch_index; + assert!(store + .complete_runtime_batch(&pool, &[invalid]) + .await + .expect("reject a different batch item") + .is_empty()); + + let mut holder = pool.begin().await.expect("begin batch holder"); + sqlx::query(audited_sql(format!( + "SELECT batch_id FROM {}.lease_claim_batches \ + WHERE claim_slot = $1 AND batch_id = $2 FOR UPDATE", + store.schema() + ))) + .bind(jobs[0].claim.claim_slot) + .bind(jobs[0].claim.claim_batch_id.expect("compact batch")) + .execute(holder.as_mut()) + .await + .expect("hold the immutable batch row"); + let completed = tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(2)), + store.complete_runtime_batch(&pool, &jobs), + ) + .await + .expect("completion must not wait on an unrelated batch-row lock") + .expect("complete metadata jobs"); + assert_eq!(completed.len(), 2); + holder.rollback().await.expect("release batch row"); + for job in jobs { + let durable = store + .load_job(&pool, job.job.id) + .await + .expect("read completed job") + .expect("completed job"); + assert_eq!(durable.state, JobState::Completed); + assert_eq!(durable.tags, vec!["retained".to_string()]); + assert_eq!(durable.metadata, serde_json::json!({"source": "recovery"})); + } +} diff --git a/benchmarks/baseline.json b/benchmarks/baseline.json index 9f7fee869..e4ec128e0 100644 --- a/benchmarks/baseline.json +++ b/benchmarks/baseline.json @@ -32,6 +32,13 @@ "rust/stale_heartbeat_rescue": { "handler_per_s": 160.3 }, + "rust/enqueue_lane_16p_unique_185": { + "enqueue_per_s": 10000.0, + "p99_ms": 2500.0 + }, + "rust/enqueue_lane_16p_unique_185_shards4": { + "enqueue_per_s": 10000.0 + }, "rust/mvcc_horizon_overlap": { "min_overlap_vs_baseline": 0.75, "min_cooldown_vs_baseline": 0.75, diff --git a/correctness/storage/MAPPING.md b/correctness/storage/MAPPING.md index 55fb7b178..1bc093e7d 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 9fc093412..9d0e0d84c 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -152,7 +152,9 @@ Enqueue is transactional: if the producer's outer transaction rolls back, the jo There are two COPY-shaped producer paths: - Direct queue-storage COPY is the high-throughput path: Rust producers use a configured `QueueStorage::enqueue_params_copy()`, and Python producers use `enqueue_many_copy()`. Each ready or deferred write group uses array-backed INSERT below 2,048 rows and CSV COPY for larger groups. Direct-copy producers must use the same queue-storage configuration as the worker fleet, especially `queue_stripe_count` / `queue_storage_queue_stripe_count`. -- `insert_many_copy()` is the compatibility path. It stages rows with an array-backed INSERT below 2,048 rows and CSV COPY at or above that cutoff, then takes one of two branches. If any staged row has a `unique_key`, every row is replayed through `POSTGRES_INSERT_JOB_SQL` inside a per-row savepoint so SQLSTATE `23505` can skip that duplicate without aborting the batch; that statement itself calls `awa.insert_job_compat()`. An all-non-unique batch uses one lateral `awa.insert_job_compat()` call per staged row without savepoints. Both branches therefore preserve cross-storage behavior. The planned `awa.insert_job()` public capability becomes the clean stable SQL contract after #342 ships its migration and conformance artifacts. +- `insert_many_copy()` is the compatibility path. It stages rows with an array-backed INSERT below 2,048 rows and CSV COPY at or above that cutoff, then inserts them through `awa.insert_jobs_compat_batch()` (migration v055), 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 before v055, or for a role without `EXECUTE` on the function, it takes one of two per-row branches: if any staged row has a `unique_key`, every row is replayed through `POSTGRES_INSERT_JOB_SQL` inside a per-row savepoint so SQLSTATE `23505` can skip that duplicate without aborting the batch; that statement itself calls `awa.insert_job_compat()`. An all-non-unique batch uses one lateral `awa.insert_job_compat()` call per staged row without savepoints. Every branch therefore preserves cross-storage behavior. The planned `awa.insert_job()` public capability becomes the clean stable SQL contract after #342 ships its migration and conformance artifacts. + +Generic `insert_many` always uses one statement that calls `awa.insert_job_compat()` per row in canonical lane order: its `PgExecutor` is consumed by that statement, so it cannot probe the schema first. `insert_many_skip_unique_conflicts` requires the batch function. Claim is cursor-based rather than heap-scan based: diff --git a/docs/assets/architecture-flow.svg b/docs/assets/architecture-flow.svg index 776349a60..b906b5a04 100644 --- a/docs/assets/architecture-flow.svg +++ b/docs/assets/architecture-flow.svg @@ -5,15 +5,15 @@ - - + + @@ -21,42 +21,38 @@ ONE POSTGRESQL SYSTEM OF RECORD - - - - Applications - Rust · Python · SQL bridges - Own the application transaction + + + + Applications + Rust · Python · SQL bridges + Own the application transaction - - - - - PostgreSQL - application data + AWA jobs + + + + PostgreSQL + application data + AWA jobs - - - - Workers - Claim · heartbeat · execute - Complete · retry · wait + + + + Workers + Claim · heartbeat · execute + Complete · retry · wait - - - ATOMIC COMMIT + + ATOMIC COMMIT - - CLAIM + + CLAIM - - DURABLE OUTCOME + + DURABLE OUTCOME - - - OPERATE · CLI · WEB UI · METRICS · TRACES · SQL - - + + OPERATE · CLI · WEB UI · METRICS · TRACES · SQL + diff --git a/docs/configuration.md b/docs/configuration.md index cc9b035fd..6f83f9aaa 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -245,7 +245,7 @@ Retry backoff and `Snooze` use the same mechanism: the current attempt closes, a Operational knobs: - `promote_interval` controls how often maintenance checks for due deferred rows. -- A queue whose enqueue lane is kept busy by producers does not delay promotion of other queues or the leader's rescue: promotion waits 50 ms for a lane's lock, then hands the lane to a dedicated waiter that is served in turn. Each waiter holds one connection from the worker's pool while it waits, on top of the leader's own. The leader runs at most one waiter per eight pool connections, never fewer than one or more than four, so a 10-connection pool gives up at most one extra connection to a busy lane. `awa_maintenance_promote_lane_contention_total` counts those hand-offs by queue. A steady `waiter_limit` rate means more lanes are saturated than the waiters cover, and a larger pool allows more waiters. `waiter_timed_out` means a transaction held a lane lock for longer than 30 s. +- A queue whose enqueue lane is kept busy by producers does not delay promotion of other queues or the leader's rescue: promotion waits 50 ms for a lane's lock, then hands the lane to a dedicated waiter that is served in turn. Each waiter holds one connection from the worker's pool while it waits, on top of the leader's own. The leader runs at most one waiter per eight pool connections, never fewer than one or more than four, so a 10-connection pool gives up at most one extra connection to a busy lane. `awa_maintenance_promote_lane_contention_total` counts those hand-offs by queue. A steady `waiter_limit` rate means more lanes are saturated than the waiters cover, and a larger pool allows more waiters. `waiter_timed_out` means a waiter's promotion transaction waited longer than 30 s for one of its locks, usually the lane lock behind the producers queued ahead of it. - Per-job `run_at` controls the due time. - Periodic schedules are declared in worker code with `periodic()` and managed through the CLI/UI `cron` surface. @@ -637,6 +637,8 @@ Queue storage defaults `AWA_COMPLETION_SHARDS` to `1` for ordinary runtimes and `AWA_COMPLETION_BATCH_SIZE` defaults to `512` and `AWA_COMPLETION_FLUSH_MS` defaults to `1`. The queue-storage short-job path fuses receipt closure and terminal insertion into one statement, so the lower flush interval reduces the worker capacity tied up waiting for durable completion while the batch size still amortises the claim/complete path under load. That pairing is what moves the lower WAL budget into durable throughput. +A runtime lets at most the larger of its worker capacity and `2 × AWA_COMPLETION_BATCH_SIZE × AWA_COMPLETION_SHARDS` finished attempts wait for their completion to commit: 1,024 with the queue-storage defaults below 512 workers. Each attempt's lifecycle hooks run within that allowance. A handler that finishes while it is used up keeps its worker slot until a completion commits, so a runtime whose completions or hooks fall behind claims more slowly instead of holding the backlog in memory. Each waiting attempt holds its job, payload included; for large payloads, lower `AWA_COMPLETION_BATCH_SIZE` to lower the bound in proportion. + For very high-throughput no-op queues, size `max_workers` for durable completion latency, not handler CPU alone. The `10k/s` offered-rate reference shape needs 1,024 in-flight worker permits to absorb the load consistently; real jobs with non-trivial handler time need workload-specific sizing. The terminal write path is already narrow for running/waiting jobs: successful receipt completions use compact `receipt_completion_batches_*` rows, and other retained terminal facts use `done_entries_*` without duplicating immutable ready body fields when the ready row is still retained. Both surfaces hydrate through `terminal_jobs`, so the default completion settings stay conservative without giving up durable terminal history. diff --git a/docs/getting-started-python.md b/docs/getting-started-python.md index 8096318db..48285c64c 100644 --- a/docs/getting-started-python.md +++ b/docs/getting-started-python.md @@ -2,9 +2,9 @@ This guide takes you from `uv init` to a job reaching `completed`. -!!! note "Version used in this guide" +!!! note "Versions" - The install commands pin **v0.6.6**, the latest stable release. The canonical example is tested against both that release and the code on `main`; development-only 0.7 surfaces elsewhere on this site are identified by the site banner and stability labels. + The install commands add the latest published release; see [GitHub releases](https://github.com/hardbyte/awa/releases) for versions and changelogs. The canonical example is tested against the code on `main`; development-only 0.7 surfaces elsewhere on this site are identified by the site banner and stability labels. ## Prerequisites @@ -23,7 +23,7 @@ export DATABASE_URL=postgres://postgres:test@localhost:15432/awa_test ```bash uv init awa-python-quickstart --bare cd awa-python-quickstart -uv add awa-pg==0.6.6 +uv add awa-pg ``` `uv init` creates a minimal Python project. `uv add` creates and manages the project's virtual environment, records `awa-pg` in `pyproject.toml`, and writes a lockfile—there is no environment activation step. @@ -83,7 +83,7 @@ uv run python -m awa --database-url "$DATABASE_URL" queue stats The dashboard ships in a separate wheel so the default `awa-pg` install stays small for workers and producers. Install the `[ui]` extra to bring in the `awa-cli` binary that hosts it: ```bash -uv add 'awa-pg[ui]==0.6.6' +uv add 'awa-pg[ui]' uv run python -m awa --database-url "$DATABASE_URL" serve # → http://127.0.0.1:3000 ``` diff --git a/docs/getting-started-rust.md b/docs/getting-started-rust.md index 433498a39..1c7d5ffd2 100644 --- a/docs/getting-started-rust.md +++ b/docs/getting-started-rust.md @@ -2,9 +2,9 @@ This guide takes you from `cargo add` to a job reaching `completed`. -!!! note "Version used in this guide" +!!! note "Versions" - The install commands pin **v0.6.6**, the latest stable release. The canonical example is tested against both that release and the code on `main`; development-only 0.7 surfaces elsewhere on this site are identified by the site banner and stability labels. + The install commands add the latest published release; see [GitHub releases](https://github.com/hardbyte/awa/releases) for versions and changelogs. The canonical example is tested against the code on `main`; development-only 0.7 surfaces elsewhere on this site are identified by the site banner and stability labels. ## Mental Model @@ -35,7 +35,7 @@ export DATABASE_URL=postgres://postgres:test@localhost:15432/awa_test cargo new awa-rust-quickstart cd awa-rust-quickstart -cargo add awa@0.6.6 +cargo add awa cargo add sqlx --features runtime-tokio-rustls,postgres cargo add tokio --features macros,rt-multi-thread,time cargo add serde --features derive @@ -69,7 +69,7 @@ job 1 state = Completed Install the CLI if you want migration/admin/UI commands: ```bash -uv tool install awa-cli==0.6.6 +uv tool install awa-cli ``` If uv reports that its tool directory is not on `PATH`, update your shell and diff --git a/docs/index.md b/docs/index.md index f117cfe82..c57ec9918 100644 --- a/docs/index.md +++ b/docs/index.md @@ -46,7 +46,7 @@ Business data and job state share PostgreSQL. A transaction either commits both === "Rust" ```bash - cargo add awa@0.6.6 + cargo add awa ``` Define typed job arguments, register async handlers, and use your existing `sqlx` pool for migrations, enqueueing, and administration. @@ -56,7 +56,7 @@ Business data and job state share PostgreSQL. A transaction either commits both === "Python" ```bash - uv add awa-pg==0.6.6 + uv add awa-pg ``` Register sync or async handlers around dataclasses and use direct clients or transaction bridges for asyncpg, psycopg, SQLAlchemy, and Django. @@ -66,7 +66,7 @@ Business data and job state share PostgreSQL. A transaction either commits both === "Operations" ```bash - uv tool install awa-cli==0.6.6 + uv tool install awa-cli awa --database-url "$DATABASE_URL" health ``` diff --git a/docs/javascripts/language-switch.js b/docs/javascripts/language-switch.js index e1daf2de1..75053bceb 100644 --- a/docs/javascripts/language-switch.js +++ b/docs/javascripts/language-switch.js @@ -31,10 +31,13 @@ ); } + // A label click focuses its input, which scrolls that tab set into view. function selectTabs(language) { languageTabs() .filter((label) => label.textContent.trim().toLowerCase() === language) - .forEach((label) => label.click()); + .map((label) => document.getElementById(label.htmlFor)) + .filter((input) => input && !input.checked) + .forEach((input) => input.click()); } function updateButtons(language) { diff --git a/docs/lifecycle-hooks.md b/docs/lifecycle-hooks.md index 333d53675..8023579e3 100644 --- a/docs/lifecycle-hooks.md +++ b/docs/lifecycle-hooks.md @@ -74,9 +74,11 @@ A callback that _resumes_ the job to `running` emits no event itself; the job re Hooks are **best-effort, in-process notifications** — not a durable workflow mechanism: - They observe **committed** state. An outcome hook carries the post-transition `JobRow`; a stale/rejected finalization emits no outcome event (though a `Started` may already have fired for that attempt). -- They do not block executor progress or queue capacity, and `shutdown()` does not wait for them. They can be lost on process crash or shutdown. +- They do not delay a job's state transition, and `shutdown()` does not wait for them. They can be lost on process crash or shutdown. +- Each attempt's hooks hold one of the worker's finalizer slots until they return, the same bounded slots that wait for completions to commit. Hooks slower than job throughput therefore slow claiming instead of piling up in memory. Keep hooks fast, and hand slow work to a job or follow-up instead. +- Maintenance dispatches `Rescued` events, and `Cancelled` for duplicates it supersedes, on detached tasks with at most 1,024 outstanding per worker. Events beyond that are dropped with a warning rather than holding up rescue. - **Dispatch is not atomic with the state transition.** The hook runs _after_ the job's transition commits, as a separate in-process step — it is not a transactional outbox. If the process dies between the commit and the dispatch, the job state is durable but the event is lost. The same applies to callback resolution: `complete_external` / `resolve_callback` commit the outcome first, then dispatch. -- `Started` is dispatched detached, so a very fast job may finish before its `Started` hook completes — do not rely on strict started-before-outcome ordering for short jobs. +- `Started` runs alongside the handler, so a very fast job's state transition may commit before its `Started` hook completes. For outcomes the worker finalizes itself, the outcome hook runs only after the `Started` hook has returned. - A panicking hook is logged; later hooks for the same kind still run. If a side effect must fire exactly once or survive a crash, do not drive it from a hook at all — the dispatch may never run, and enqueueing the follow-up work _from_ the hook just inherits that same unreliability. See the next section for the durable mechanism. diff --git a/docs/overrides/main.html b/docs/overrides/main.html index 7dc69e34e..ece2e903b 100644 --- a/docs/overrides/main.html +++ b/docs/overrides/main.html @@ -3,8 +3,7 @@ {% block announce %} These docs track main / 0.7 development. - Latest stable: - v0.6.6. + See releases for stable versions. {% endblock %} diff --git a/docs/reference/python.md b/docs/reference/python.md index ad6f5192d..e05d354e1 100644 --- a/docs/reference/python.md +++ b/docs/reference/python.md @@ -3,7 +3,7 @@ Install `awa-pg` for the Python clients and worker runtime: ```bash -uv add awa-pg==0.6.6 +uv add awa-pg ``` The package ships type information for its public surface. This page is a map; use Python help and your editor against the installed version for exact signatures. diff --git a/docs/security/database-roles.md b/docs/security/database-roles.md index 8863fc1a5..e6673ade6 100644 --- a/docs/security/database-roles.md +++ b/docs/security/database-roles.md @@ -124,6 +124,10 @@ SELECT format('SELECT * FROM awa.repair_lane_sequence(%L, %L)', 'awa', sequence_ FROM awa.lane_sequence_problems('awa') \gexec ``` +## Claim batch index + +From schema v054 each queue-storage schema has `ensure_claim_batch_job_index(slot)`. Claim prune calls it right after truncating a claim slot to give that `lease_claim_batches` partition its GIN index on `job_ids`. It is `SECURITY DEFINER`, owned by the owner of `lease_claim_batches` (creating an index needs table ownership), with a pinned `search_path`. It only creates that one index, and only on a partition that is empty and unlocked. `PUBLIC` cannot execute it. Each time the migration or installer runs, `EXECUTE` is reset to exactly the roles that then hold `TRUNCATE` on `lease_claim_batches`: a runtime role given `TRUNCATE` later gets it on the next run (or from `GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA` above until then), and a role that has lost `TRUNCATE` loses it. A runtime that cannot execute it skips the step and claims keep working. + ## Custom queue-storage schemas Repeat the schema, table, sequence, function, and default-privilege grants for a custom queue-storage schema. Prepare custom schemas as the migrator; the runtime needs DML, sequence access, function execution, and `TRUNCATE`, but never DDL. Check lanes in a custom schema with `awa storage lane-sequences --schema `. See [Queue storage](../queue-storage-substrate.md). diff --git a/docs/start/cli.md b/docs/start/cli.md index 4f0bde501..14c13a09a 100644 --- a/docs/start/cli.md +++ b/docs/start/cli.md @@ -5,14 +5,14 @@ The `awa` command runs migrations, inspects and administers queues, and can host === "uv tool (recommended)" ```bash - uv tool install awa-cli==0.6.6 + uv tool install awa-cli awa --help ``` === "Project dependency" ```bash - uv add 'awa-pg[ui]==0.6.6' + uv add 'awa-pg[ui]' uv run python -m awa --help ``` diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index 1672c5865..e00ddfae0 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -137,6 +137,22 @@ To recover: In all cases the cancel itself is durable — only the early wake-up is lost. Long-running handlers that don't poll `ctx.is_cancelled()` between heartbeats won't notice the cancel until they finish or heartbeat-rescue fires. +## Claiming Slows While Handlers Are Idle + +### What It Usually Means + +Durable completion, or lifecycle hooks, are behind the handlers. A runtime lets only a bounded number of finished attempts wait for their completion to commit (see [completion settings](configuration.md#hot-queue-claim-control)). Once that allowance is used up, finished handlers keep their worker slots, `awa_job_in_flight` sits at `max_workers` while handler CPU is idle, and the runtime claims only as fast as completions commit. Worker memory stays flat. 0.6 releases before 0.6.15 kept claiming instead, and their memory grew with every finished attempt until the process was killed. + +### What To Check + +- `awa_completion_flush_duration_seconds`: long flushes point at the database, usually lock waits on the completion or terminal tables. Look for long transactions and lock waits on the queue-storage schema, as in [Inspect Long Transactions](#inspect-long-transactions). +- `awa_completion_deadlock_retry_total`: a steady rate means completions are re-running after deadlocks. +- Lifecycle hooks (`on_event`, `on_event_kind`): a hook that waits on a slow service or on the same database holds its attempt's slot until it returns. Move slow work into a job or a follow-up. + +### Fixes + +Clear the lock contention or long transaction; claiming recovers as soon as completions commit again. Raising `max_workers` does not help while completion is the bottleneck. For jobs with large payloads, lowering `AWA_COMPLETION_BATCH_SIZE` reduces how much memory waiting attempts can hold. + ## Rescue Fails With `idx_awa_jobs_unique` Symptom: maintenance repeats, every tick (canonical engine — the queue-storage @@ -186,7 +202,7 @@ Reading them together separates the two usual failure modes: **1. Producer is going through the compatibility insert path.** -`insert_many_copy_from_pool` / `client.insert_many_copy` routes every row through `awa.insert_job_compat()` once per row. On any database with non-trivial per-statement latency (auth-proxy hop, managed Postgres, network) the per-row function call dominates — measured at ~100–150 ms per row through a Cloud SQL Auth Proxy in staging, which caps a single producer at ~7 rows/sec regardless of batch or chunk size. +Before migration v055, `insert_many_copy_from_pool` / `client.insert_many_copy` 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 `awa.insert_jobs_compat_batch`; 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: @@ -214,7 +230,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. From migration v055, `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 [Managed Postgres sizing](deploying-on-managed-postgres.md#pick-a-vcpu-size) are useful reference points. @@ -439,7 +466,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` 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 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 +- from migration v055, `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, and binaries that predate the batch function) 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-to-0.7.md b/docs/upgrade-0.6-to-0.7.md index 61fb2c7d5..439c61aa4 100644 --- a/docs/upgrade-0.6-to-0.7.md +++ b/docs/upgrade-0.6-to-0.7.md @@ -283,6 +283,42 @@ workers report not ready until the migration runs, as for any version bump. External migration runners apply the complete `awa migrate --sql --from 51 --to 52` bundle, which defines the helper before the tail calls it. +## v054: claim batch job index + +Point lookups of compact receipt claims by job id (heartbeat and progress flush, completion fallbacks, receipt close, cancellation, rescue and `load_job`) use a GIN index on `lease_claim_batches.job_ids` instead of reading every retained claim batch. v054 installs `.ensure_claim_batch_job_index(slot)` in the default and every discovered custom queue-storage schema and builds no index itself, so it takes no lock on any awa or queue-storage table. Claim prune calls that function right after truncating a claim partition, which builds the index while the partition is empty, so every partition gains it within one claim-ring cycle. A new schema's installer indexes its claim partitions as it creates them. + +Do not build the index yourself on a populated partition: the build holds a `SHARE` lock, which blocks claim inserts into that partition, for about half a minute per 100k batches. Run the migration as the owner of the queue-storage tables, or as a migrator that can `SET ROLE` to it, so the function can create the index; see [Claim batch index](security/database-roles.md#claim-batch-index). + +Binaries before v054 neither use nor build the index, and binaries that use it read the same rows, by scan, from a partition without it. A database coming from 0.6.15 already has the function from the 0.6 `claim_batch_job_index` schema patch, and v054 replaces it in place. External migration runners apply the complete `awa migrate --sql --from 53 --to 54` bundle, which defines the helper before the tail installs the function. + +## v055: batch compatibility insert + +v055 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 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, reserves each lane's range of positions once 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, so 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 check on every call, so a running producer switches on its next call after the migration commits. Their results are unchanged: rows in input order, jobs whose unique key is already claimed skipped. +- `insert_many_skip_unique_conflicts` (Rust) inserts a batch through the function and skips jobs whose unique key is already claimed, returning the inserted rows in input order. Use it instead of retrying a failed `insert_many` one job at a time. It returns a validation error on a schema before v055. +- 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. + +Binaries before v055 never call the 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. A database coming from 0.6.15 already has the function from the 0.6 `insert_jobs_compat_batch` schema patch, and v055 replaces it in place. + +The batch function takes all of a batch's unique claims before any lane lock, while per-row inserts (generic `insert_many`, single-job inserts, and binaries before v055) 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, 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. 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. + +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 apply the complete `awa migrate --sql --from 54 --to 55` bundle, which refreshes trusted queue-storage discovery before it replaces the function. + ## The one gate `awa migrate` on a 0.7 binary refuses to apply migrations unless one of these holds: diff --git a/scripts/previous-release-tag.py b/scripts/previous-release-tag.py index 073a121db..9b4d92e13 100644 --- a/scripts/previous-release-tag.py +++ b/scripts/previous-release-tag.py @@ -1,5 +1,9 @@ #!/usr/bin/env python3 -"""Choose the greatest strictly earlier SemVer tag in the same minor release.""" +"""Choose the greatest strictly earlier SemVer tag in the same minor release. + +A final release compares with the final release before it; a prerelease +compares with whatever tag precedes it. +""" import re import sys @@ -31,8 +35,11 @@ def previous_release_tag(current_tag: str, tags: list[str]) -> str: candidates = [] for tag in tags: key = version_key(tag) - if key is not None and key[:2] == current[:2] and key < current: - candidates.append((key, "+" not in tag, tag)) + if key is None or key[:2] != current[:2] or key >= current: + continue + if current[3] and not key[3]: + continue + candidates.append((key, "+" not in tag, tag)) return max(candidates)[2] if candidates else "" diff --git a/scripts/tests/test_previous_release_tag.py b/scripts/tests/test_previous_release_tag.py index 99bdd9df8..f97c3fbe8 100644 --- a/scripts/tests/test_previous_release_tag.py +++ b/scripts/tests/test_previous_release_tag.py @@ -34,11 +34,26 @@ def test_semver_prerelease_ordering(self): tags = [ "v1.0.0-alpha", "v1.0.0-alpha.1", "v1.0.0-alpha.beta", "v1.0.0-beta", "v1.0.0-beta.2", "v1.0.0-beta.11", - "v1.0.0-rc.1", "v1.0.0-rc.2", "v1.0.0-rc.10", "v1.0.0", + "v1.0.0-rc.1", "v1.0.0-rc.2", "v1.0.0-rc.10", ] for previous, current in zip(tags, tags[1:]): with self.subTest(current=current): - self.assertEqual(previous_release_tag(current, list(reversed(tags))), previous) + self.assertEqual( + previous_release_tag(current, list(reversed(tags + ["v1.0.0"]))), + previous, + ) + + def test_release_skips_every_prerelease_after_the_previous_release(self): + self.assertEqual( + previous_release_tag( + "v0.6.15", + ["v0.6.13-alpha.1", "v0.6.13", "v0.6.14-rc.1", "v0.6.15-rc.1", "v0.6.15"], + ), + "v0.6.13", + ) + + def test_first_release_in_minor_has_no_predecessor_after_prereleases(self): + self.assertEqual(previous_release_tag("v1.0.0", ["v1.0.0-rc.1", "v0.9.9"]), "") def test_build_metadata_does_not_change_precedence(self): self.assertEqual( diff --git a/skills/awa-jobs/SKILL.md b/skills/awa-jobs/SKILL.md index 3895bf63a..f1bc518da 100644 --- a/skills/awa-jobs/SKILL.md +++ b/skills/awa-jobs/SKILL.md @@ -127,7 +127,10 @@ a time-period bucket. Defaults: dedup by args, not by queue. Conflict handling depends on the API: a single-row `insert` / `insert_with` and the row-by-row `insert_many` **raise `AwaError::UniqueConflict`** (Python raises the corresponding error) on a duplicate — you must handle it, do not assume a silent -no-op. Only the COPY bulk path (`insert_many_copy`) silently skips duplicates. +no-op. Two batch paths skip duplicates instead: `insert_many_copy` (needs a +connection) and Rust `insert_many_skip_unique_conflicts`, which takes any +`PgExecutor` and needs migration v055. Both return only the inserted rows, in +input order; a job whose key is already claimed is left out of the result. Cancel without storing the id via `cancel_by_unique_key(kind, queue=, args=, period_bucket=)` — the components must match those used at insert. It cancels the oldest live job holding the key. On queue storage, a `by_queue` key is @@ -259,7 +262,10 @@ Rust exposes two hook families on the client builder: - **Observation hooks** (`.on_event` / `.on_event_kind`) fire best-effort in process for `Started`, `Completed`, `Retried`, `Exhausted`, `Cancelled`, `Rescued`, `WaitingForCallback`. They are **not** a durable workflow — they can - be lost on crash and `shutdown()` does not wait for them. + be lost on crash and `shutdown()` does not wait for them. Each attempt's hooks + hold one of the worker's finalizer slots until they return, so hooks slower + than job throughput slow claiming. Keep them fast and move slow work into a + job or a durable follow-up. - **Durable follow-up hooks** (`on_completed_enqueue`, `on_exhausted_enqueue`, and the other `on_*_enqueue`) INSERT a follow-up job for side effects that must survive a crash. Atomicity depends on what triggered them: