diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2ae3aee83..5dd4872b2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -204,14 +204,18 @@ jobs: # canonical, so this leg re-runs it with the queue-storage engine active on a # fresh database (a separate DB avoids cross-engine job_unique_claims carryover). rust-test-queue-storage: - name: Rust tests (queue_storage) + name: Rust tests (queue_storage, pg${{ matrix.postgres }}) needs: [rust-build] if: github.event_name != 'pull_request' || contains(github.event.pull_request.labels.*.name, 'full-ci') runs-on: ubuntu-latest timeout-minutes: 30 + strategy: + fail-fast: false + matrix: + postgres: ["17", "18"] services: postgres: - image: postgres:17-alpine + image: postgres:${{ matrix.postgres }}-alpine env: POSTGRES_PASSWORD: postgres POSTGRES_DB: awa_test @@ -314,10 +318,14 @@ jobs: # ─── Python (awa-python) ──────────────────────────────── python-build-test: - name: Python build + test + name: Python build + test (${{ matrix.python }}) if: github.event_name != 'pull_request' || contains(github.event.pull_request.labels.*.name, 'full-ci') runs-on: ubuntu-latest timeout-minutes: 60 + strategy: + fail-fast: false + matrix: + python: ["3.14", "3.15"] services: postgres: image: postgres:17-alpine @@ -342,7 +350,7 @@ jobs: workspaces: "awa-python -> awa-python/target" - uses: actions/setup-python@v6 with: - python-version: "3.12" + python-version: ${{ matrix.python }} - name: Install uv uses: astral-sh/setup-uv@v8.2.0 - name: Create venv and install deps diff --git a/CHANGELOG.md b/CHANGELOG.md index 1db9c1170..5f19bd2ff 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,185 +4,86 @@ Notable changes between releases. Detailed migration notes for storage transitio ## [Unreleased] -### Added +### Fixed + +- Bound outstanding worker finalization to the larger of runtime worker capacity + and twice the completion batch size times the shard count. Finished handlers + retain their worker slot while waiting for finalizer capacity, preventing an + unbounded backlog of tasks retaining job payloads when durable completion falls + behind. Attempt heartbeats remain active until finalization settles. +- Complete retried queue-storage jobs and jobs carrying metadata by compact claim + batch identity and item index, avoiding a scan of every retained claim batch. + +## [0.6.14] — 2026-11-08 -- **The queue-storage prune `lock_timeout` is configurable, and prune - attempts are counted by outcome for every ring.** Prune (queue, lease and - claim rings) and terminal-delta rollup take their exclusive locks under a - 10 ms `lock_timeout` that could only be changed through the hidden - `QueueStorage::with_prune_lock_timeout`. When short readers of the ring - partitions held a conflicting lock at that instant, prune reported - `blocked` every tick and the rings stopped recycling. Set it with - `ClientBuilder::prune_lock_timeout` (Rust) or - `queue_storage_prune_lock_timeout_ms` on `client.start()` (Python); the - timeout applies separately to each lock acquisition, not to the whole prune; - queue prune can wait on up to eight partition locks. The default stays 10 ms and zero or a value above PostgreSQL's `lock_timeout` - maximum is rejected. Start logs a warning when the value is at or above the - connection's `deadlock_timeout` (prune waiting that long can make deadlock - detection cancel a claim instead) or `statement_timeout`. `with_prune_lock_timeout` is now - public, with `set_prune_lock_timeout` and `prune_lock_timeout()`. - `awa_maintenance_prune_attempts_total{awa_ring,awa_ring_outcome,awa_ring_reason}` - gains `backed_off` (prune skipped by the backoff that follows consecutive - `skipped_active` or `blocked` results) and - `error` (prune failed with a database error) alongside the existing - `pruned`, `noop`, `blocked` and `skipped_active`. See - [Configuration](docs/configuration.md#what-the-knobs-mean). +Recommended for everyone on 0.6.13 using queue storage. A busy queue no longer stalls promotion of the others, prune backs off cleanly under lock contention, new queue lanes are usable by every runtime role, and jobs no longer stay running, run again or get misreported after priority aging. + +### Upgrading + +- The schema stays at v40. Run `awa migrate` once as the schema owner while workers keep running. It applies the repeatable `definer_lane_provisioning` patch and repairs existing lane sequences. Then roll workers, producers and admin processes in any order. Binaries can roll back to 0.6.13 with the patch left in place. +- `awa storage lane-sequences` checks lane-sequence ownership and grants and exits with code 2 while problems remain; `--repair` fixes what the connected role may. +- External migration runners: apply `R__definer_lane_provisioning.sql` after `V40` (`awa migrate --extract-to --pending`), replacing all `R__` files together, then run `awa storage lane-sequences --repair` as the migration role. See the [upgrade guide](docs/upgrade-0.6.13-to-0.6.14.md). ### Fixed -- Queue-storage completion of retried jobs or jobs carrying metadata now - addresses compact claims by their batch identity and item index, as on main. - It no longer expands every retained claim batch for each completion flush, - so bounded finalization can drain large recovery backlogs without that scan. - -- **A completion backlog no longer gets finished jobs rescued and run - again.** An attempt stops heartbeating when its handler returns, and its - finalize then waits for the completion batcher. The runtime released the - worker slot at the same moment and kept claiming, so when finalization was - slower than claiming (for example, retried jobs, whose completions take - the slower per-row path) the queue of finished, unheartbeated attempts grew - without bound. Attempts that waited longer than `heartbeat_staleness` were - rescued and executed again, some of them twice. The bound on attempts waiting - for finalization is `max(worker count, 2 × batch size × completion shards)`. - Past that, a finished attempt keeps its - worker slot and its heartbeat until there is room, so the dispatcher stops - claiming until finalization catches up. -- **Ring prune and terminal-count rollup report `Blocked` when any of their - statements times out on a lock.** Prune sets a short `lock_timeout` (10 ms - by default) before it takes `ACCESS EXCLUSIVE` on a slot's partitions, but - only that `LOCK TABLE` and the final `TRUNCATE` turned a timeout into - `PruneOutcome::Blocked`. The timeout covers every later statement in the - transaction too: the post-lock liveness checks read the lease and claim - partitions, the failed-row carry inserts into the current slot's - `done_entries` and terminal-count deltas, and the fold writes - `queue_terminal_rollups` and deletes `queue_terminal_live_counts` rows. When - one of those waited out the timeout, for example behind autovacuum - truncating a relation or a hosted Postgres's own background workers, prune - failed with a `lock_not_available` (`55P03`) database error, was logged as - an error and skipped the prune backoff. A lock timeout anywhere in a queue, - lease or claim prune, or in a terminal-count rollup slot, now rolls the - whole transaction back and reports `Blocked` (`blocked_slots` for rollup), - so maintenance backs off and retries; nothing the attempt wrote survives - the rollback. A rollup pass stops at its first blocked slot rather than - waiting out the timeout again on each remaining one. -- **On queue storage, a queue whose enqueue lane is busy no longer holds - back promotion of other queues, nor the maintenance leader's rescue, - rotation and prune.** Promotion re-enqueues due scheduled and retryable - rows into their lane, which takes the lane's reservation lock that - producers hold until they commit. A promotion batch was one transaction - across every queue it touched, and the leader runs its branches one after - another, so a lane kept busy by producers delayed due jobs of unrelated - queues in the same batch, and rescue started seconds late. Promotion now - gives each `(queue, priority)` lane its own transaction under a 50 ms - `lock_timeout`. A lane that times out keeps its rows deferred and is - handed to a dedicated task that waits up to 30 s for the lock, queued - behind the producers already waiting, so the lane is still promoted in - turn however busy it stays; the leader skips it meanwhile. Each waiter - holds a pool connection, so the leader runs one per eight pool - connections, at least one and at most four; a further contended lane is - retried on the next tick. Batch selection skips rows another transaction - has locked, so locked due rows cannot fill every batch of a tick. `awa.maintenance.promote.lane_contention` - (`awa_maintenance_promote_lane_contention_total{awa_job_state,awa_job_queue,awa_promote_outcome}`) - counts `handed_off`, `waiter_limit` and `waiter_timed_out`. Promoted rows - now take lane sequence numbers in due order within their lane. - `QueueStorage::promote_due` keeps its signature and promotes lane by lane - without a lock timeout; `promote_due_lanes` and `promote_due_lane` expose - the per-lane form. -- **A queue lane is usable by every runtime role, whichever role creates it - first.** A lane's enqueue and claim sequences belonged to the role that - first enqueued into, claimed from or promoted into it, so with several - runtime roles sharing a group role, any role other than the creator whose - default privileges did not cover it failed with `permission denied for - sequence`. That included the maintenance leader promoting into a shard - created after `enqueue_shards` was raised. The new schema patch - `definer_lane_provisioning` creates lanes through two `SECURITY DEFINER` - functions per queue-storage schema, `provision_lane` and - `provision_lane_sequences`, owned by the schema owner. Each new sequence gets - the privileges of the head table it backs, explicitly rather than through - default privileges. A runtime without schema `CREATE` can therefore create - lanes itself, and `awa storage prepare-queue` is optional and may run as any - runtime role. Once the patch commits, `awa migrate` gives existing lane - sequences to the schema owner and adds their missing grants, one sequence - per transaction with a 300 ms lock wait, and only where it is allowed to, - trying the effective role and then the login role. A sequence an open - transaction holds is skipped and reported as busy; one it may not change is - reported with the roles that can repair it. Neither fails the run. The new read-only `awa storage lane-sequences` (or - `awa.lane_sequence_problems(schema)`) lists what remains, with the - statements to run and the roles to run them as, and exits 2 while any - remain; `--repair` (or `awa.repair_lane_sequence(schema, sequence)` per - sequence) repairs what the connected role may. `awa migrate` prints a summary of the leftovers. A binary - running before the patch detects that `provision_lane` is missing and - creates lanes as before, and older binaries create correctly owned lanes once - the patch has refreshed the helpers they call. `awa migrate` applies the - patch on any v040 database; external runners apply the exported - `R__definer_lane_provisioning.sql` after `V40`. See - [`docs/upgrade-0.6.13-to-0.6.14.md`](docs/upgrade-0.6.13-to-0.6.14.md). -- **A graceful shutdown no longer leaves a finished attempt for stale-heartbeat - rescue to run again.** Once a handler returned, its finalize (completion, - retry, cancel or failure, and any ADR-029 follow-up inserts) ran in a task - that `Client::shutdown` did not wait for. When the last handler returned - just as the drain ended, shutdown stopped the completion batcher before - that finalize reached it, and the direct fallback raced the process exit. - The attempt had already left the heartbeat set, so it went stale and was - rescued and executed again `heartbeat_staleness` later; a deploy could - re-run a few jobs on every rolling restart. Shutdown now waits for - pending finalization, within the same drain timeout, before it stops the - completion batcher, heartbeat and maintenance services. A finalize that - outlives the timeout is still left to rescue, as before. -- **Point reads show the live queue and priority after a lane move.** Retained - ready rows from before a priority or queue change could win a read's tie - after the ring wrapped, returning the old fields even though the batch - operation succeeded. Point reads now exclude tombstoned source rows. - -- **Callback completion, `get_job`, lifecycle events and the admin reads - find running jobs that priority aging promoted on queue storage** (backport - of #600). Claims and leases record a job's aged effective priority while its - ready row keeps the lane priority, and eight reads joined the two on - priority. With aging on (the default, every 60 s), a job that waited past a - minute before it was claimed could not be completed by callback with an - expected kind, failed the mixed-transition callback routing check, read back - from `get_job` in the wrong state, was missing from the job list, the kinds - page and the dashboard's recent-jobs counts, and was not found by - `cancel_by_unique_key`. They now match the attempt (`job_id` and `run_lease`). No - migration. -- **`get_job` no longer misses a job that moves while it reads.** It read - ready, deferred, leased, claimed, completed and DLQ rows in separate - statements, so a job promoted, claimed or completed between two of them - could read as missing. It now reads all of them from one read-only snapshot. -- **`admin::job_kind_overviews` on queue storage no longer depends on the - planner's row estimates.** Behind the admin UI's kinds page, it counted - every job through the `current_jobs` view. Two nodes of that view are - estimated at one row whatever the statistics say: the join from completion - batches to `ready_entries`, and the `UNION ALL` over every place a job can - live. Postgres therefore chose nested loops that visit each retained row, - and the join on `queue_claim_heads` was only as good as that table's - estimate, which autovacuum never refreshes. It now aggregates each place to - `(kind, queue)` on its own and combines the results: available jobs are - found per lane from the claim cursor, and a completion batch's members are - unnested and probed along the ready primary key. The rows are identical, - and a test compares them with the former definition across every job - state. With 2.6 million retained ready rows (1.8 million of them available) - and 860,000 retained completions it takes about 6 to 8 s instead of 20 s on - PostgreSQL 17 and 18. It still reads each available job and each retained - completion's ready row, because a job's kind lives there. JIT is off for - the read. -- **`admin::list_jobs` on queue storage reads only the places that can hold - the requested state.** It selected candidate ids from the `current_jobs` - view, so every listing, whatever its filter, planned around the same - single-row estimates and joined each retained completion to a scan of the - ready rows. On PostgreSQL 18 a kind filter ran past a four-minute statement - timeout. Each place a job can live is now read and limited on its own, - only those that can hold the state asked for, with a predicate for each - filter that is set rather than `$n IS NULL OR` terms, which a prepared - statement's generic plan can't drop. Completion batches are ordered by job - id before their ready rows are probed, so a page probes only as many as it - returns. The ids are identical, and a test compares them with the former - definition for every state, kind, queue, limit and page. On the dataset - above, `state=failed&limit=10` drops from 12 s to 0.7 s, - `state=completed&limit=100` from 20 s to 6 s, and an unfiltered or - kind-filtered page of 100 from 23 s to 11 s (a kind filter on PostgreSQL 18 - from more than four minutes). Available jobs and completions are still read - in full when the state allows them. JIT is off for the read. +**Lane ownership** + +- A queue lane created by one runtime role is usable by every other runtime role. Before, its sequences belonged to whichever role created it, so other roles could fail with `permission denied for sequence`, for example after `enqueue_shards` was raised. New lanes now belong to the schema owner with the head table's grants, and migration repairs existing ones, including when the migration role can only `SET ROLE` to the schema owner. + +**Promotion, prune and lock contention** + +- A busy queue lane no longer delays promotion of other queues or maintenance. Promotion waits 50 ms for a lane's locks and hands a busy lane to a dedicated waiter. +- A lock timeout anywhere in prune or terminal rollup rolls the attempt back and reports `Blocked` instead of an error, and prune backs off. +- Admin job details, lists and kind counts release partition locks between read phases, so they block prune less, and lists of retained history do less repeated work. +- A read cancelled while starting its transaction closes its connection, so the next caller cannot inherit an open transaction and its partition locks. + +**Job correctness after priority aging and recovery** + +- Physical priority aging no longer moves a ready row that is being claimed. Before, it could tombstone the body of a just-claimed job, after which completion, cancellation and rescue failed and the claim ring stayed pinned. Bodies left tombstoned by older releases are recovered when the job is read, completed, cancelled or rescued. +- Reads and callbacks find attempts that priority aging or a priority change moved, and admin lists and kind counts include running jobs whose claims have not yet become leases. +- Failed jobs keep their arguments, attempt limit and lane in terminal history; manual retries and DLQ moves keep the correct priority and lane. +- Schema-only admin clients work with ready rings larger than the default. + +**Workers** + +- Workers retry lock contention during finalization for up to 120 s with the handler's saved result, keeping the job's heartbeat until the outcome settles. +- Graceful shutdown waits for pending finalization before stopping heartbeat and maintenance, within the existing drain timeout. +- A completion or callback wait that has committed is no longer reported as failed if re-reading the job for its lifecycle event fails. + +**CLI** + +- `awa migrate` rejects `--sql` combined with `--extract-to`. Before, `--sql` silently won and nothing was extracted. + +### Since 0.6.14-alpha.1 + +For testers of the prerelease, these fix behaviour that alpha.1 introduced: + +- Promotion no longer holds a lane's locks while waiting for PostgreSQL's cluster-wide NOTIFY lock at commit, where under concurrent enqueues the 50 ms lane timeout could expire and roll the batch back. Claimers are woken once after the lanes commit. +- With lifecycle hooks registered, a completion builds its event from the job the worker holds instead of re-reading it, which cost about 24 round trips per completion. +- Rescue skips an attempt with ambiguous retained bodies and logs an error on each pass, instead of failing the whole rescue. + +### Added + +- `ClientBuilder::prune_lock_timeout` in Rust and `queue_storage_prune_lock_timeout_ms` in Python configure prune's lock wait. The timeout applies to each acquisition; queue prune can wait on up to eight partition locks. The default remains 10 ms. `awa.maintenance.prune.attempts` includes the outcomes `backed_off` and `error`. +- `awa.maintenance.promote.lane_contention` records promotion lane contention and waiter outcomes by queue. +- Multi-runtime soak coverage for killed workers, shard changes, mixed Rust and Python workloads, durable recovery evidence and ring-lock diagnostics. + +### Changed + +- Python CI tests 3.14 and 3.15. Older interpreter versions remain best-effort; package classifiers list 3.14 and 3.15. +- Operator guidance explains that longer rotation intervals reduce churn while increasing rows and index size per slot and extending retained history. + +### Known limitations + +- A job-detail read can race prune and briefly show a completed job as available or missing; reading again returns the current state. +- If an older release's aging race left two equally plausible bodies for one running attempt, reads and direct transitions of that attempt fail with an ambiguity error, and rescue skips it until the bodies are repaired. +- A worker completing a job with the body it already holds does not retire an unclaimed copy left by older aging, so that copy can run again (at-least-once). +- Direct API calls that recover a body left by older aging (external callback complete, fail and retry; admin cancel and retry; completion by ID) can fail with SQLSTATE 55P03 under contention and must be retried. Workers retry automatically. +- The installed `terminal_jobs` SQL view still reports the stored priority for older rows whose priority changed. + +## [0.6.14-alpha.1] — 2026-11-07 + +Prerelease of 0.6.14; every change is listed under 0.6.14. ## [0.6.13] — 2026-11-06 diff --git a/Cargo.lock b/Cargo.lock index 4386517ed..656f4a77d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -383,7 +383,7 @@ checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" [[package]] name = "awa" -version = "0.6.13" +version = "0.6.14" dependencies = [ "async-trait", "awa-macros", @@ -409,7 +409,7 @@ dependencies = [ [[package]] name = "awa-cli" -version = "0.6.13" +version = "0.6.14" dependencies = [ "assert_cmd", "awa-model", @@ -431,7 +431,7 @@ dependencies = [ [[package]] name = "awa-macros" -version = "0.6.13" +version = "0.6.14" dependencies = [ "proc-macro-crate", "proc-macro2", @@ -441,7 +441,7 @@ dependencies = [ [[package]] name = "awa-metrics" -version = "0.6.13" +version = "0.6.14" dependencies = [ "awa-model", "opentelemetry", @@ -450,7 +450,7 @@ dependencies = [ [[package]] name = "awa-model" -version = "0.6.13" +version = "0.6.14" dependencies = [ "awa-macros", "blake3", @@ -471,7 +471,7 @@ dependencies = [ [[package]] name = "awa-seaorm" -version = "0.6.13" +version = "0.6.14" dependencies = [ "awa", "awa-testing", @@ -484,7 +484,7 @@ dependencies = [ [[package]] name = "awa-testing" -version = "0.6.13" +version = "0.6.14" dependencies = [ "async-trait", "awa-model", @@ -499,7 +499,7 @@ dependencies = [ [[package]] name = "awa-ui" -version = "0.6.13" +version = "0.6.14" dependencies = [ "awa-metrics", "awa-model", @@ -522,7 +522,7 @@ dependencies = [ [[package]] name = "awa-worker" -version = "0.6.13" +version = "0.6.14" dependencies = [ "anyhow", "async-trait", diff --git a/Cargo.toml b/Cargo.toml index eba7da12b..2121e4041 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,7 +14,7 @@ members = [ exclude = ["awa-python", "spike"] [workspace.package] -version = "0.6.13" +version = "0.6.14" edition = "2021" license = "MIT OR Apache-2.0" description = "Postgres-native background job queue — transactional enqueue, heartbeat crash recovery, SKIP LOCKED dispatch" @@ -25,11 +25,11 @@ categories = ["database", "asynchronous", "web-programming"] [workspace.dependencies] # Internal crates -awa-model = { path = "awa-model", version = "0.6.13" } -awa-macros = { path = "awa-macros", version = "0.6.13" } -awa-metrics = { path = "awa-metrics", version = "0.6.13" } -awa-worker = { path = "awa-worker", version = "0.6.13" } -awa = { path = "awa", version = "0.6.13" } +awa-model = { path = "awa-model", version = "0.6.14" } +awa-macros = { path = "awa-macros", version = "0.6.14" } +awa-metrics = { path = "awa-metrics", version = "0.6.14" } +awa-worker = { path = "awa-worker", version = "0.6.14" } +awa = { path = "awa", version = "0.6.14" } # Database sea-orm = { version = "=2.0.0-rc.38", default-features = false, features = ["sqlx-postgres", "runtime-tokio-rustls", "with-chrono", "with-json"] } diff --git a/awa-cli/Cargo.toml b/awa-cli/Cargo.toml index 14f235f7b..da4141211 100644 --- a/awa-cli/Cargo.toml +++ b/awa-cli/Cargo.toml @@ -14,7 +14,7 @@ path = "src/main.rs" [dependencies] awa-model.workspace = true awa-worker.workspace = true -awa-ui = { path = "../awa-ui", version = "0.6.13" } +awa-ui = { path = "../awa-ui", version = "0.6.14" } axum.workspace = true sqlx.workspace = true tokio.workspace = true @@ -28,5 +28,5 @@ hex.workspace = true uuid.workspace = true [dev-dependencies] -awa-testing = { path = "../awa-testing", version = "0.6.13" } +awa-testing = { path = "../awa-testing", version = "0.6.14" } assert_cmd = "2" diff --git a/awa-cli/pyproject.toml b/awa-cli/pyproject.toml index 7e4f95bbf..8b3be9b01 100644 --- a/awa-cli/pyproject.toml +++ b/awa-cli/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "maturin" [project] name = "awa-cli" -version = "0.6.13" +version = "0.6.14" description = "CLI for the Awa Postgres-native job queue (migrations, admin, serve)" readme = "README.md" requires-python = ">=3.10" diff --git a/awa-cli/src/main.rs b/awa-cli/src/main.rs index 1e347be53..47ee732a4 100644 --- a/awa-cli/src/main.rs +++ b/awa-cli/src/main.rs @@ -29,7 +29,7 @@ enum Commands { #[arg(long)] extract_to: Option, /// Print migration SQL to stdout instead of applying - #[arg(long)] + #[arg(long, conflicts_with = "extract_to")] sql: bool, /// Only include migrations after this version (exclusive) #[arg(long)] diff --git a/awa-model/migrations/v023_install_queue_storage_substrate.sql b/awa-model/migrations/v023_install_queue_storage_substrate.sql index 4625ebada..57f56c7f7 100644 --- a/awa-model/migrations/v023_install_queue_storage_substrate.sql +++ b/awa-model/migrations/v023_install_queue_storage_substrate.sql @@ -511,13 +511,12 @@ AS $problems$ AND NOT candidate.rolsuper AND pg_has_role(candidate.oid, lane_sequence.relowner, 'USAGE') AND (lane_sequence.relowner = lane_sequence.nspowner - OR (pg_has_role(candidate.oid, lane_sequence.nspowner, 'USAGE') - AND pg_has_role( + OR pg_has_role( candidate.oid, lane_sequence.nspowner, CASE WHEN current_setting('server_version_num')::int >= 160000 THEN 'SET' ELSE 'MEMBER' END - ))) + )) ORDER BY candidate.rolname ) FROM lane_sequence @@ -555,6 +554,9 @@ AS $repair$ DECLARE v_problem RECORD; v_remaining RECORD; + v_original_role TEXT := current_setting('role'); + v_reassigned BOOLEAN := false; + v_role_changed BOOLEAN := false; v_busy BOOLEAN := false; v_set_mode TEXT := CASE WHEN current_setting('server_version_num')::int >= 160000 THEN 'SET' ELSE 'MEMBER' END; @@ -576,6 +578,7 @@ BEGIN AND pg_has_role(current_user, to_regrole(quote_ident(v_problem.expected_owner)), v_set_mode) THEN BEGIN EXECUTE v_problem.reassign_sql; + v_reassigned := true; EXCEPTION WHEN insufficient_privilege THEN NULL; END; @@ -585,14 +588,29 @@ BEGIN FROM awa.lane_sequence_problems(p_schema, p_sequence) AS problem; IF FOUND AND cardinality(v_remaining.missing_grants) > 0 - AND pg_has_role(current_user, to_regrole(quote_ident(v_remaining.owner)), 'USAGE') THEN - FOREACH v_statement IN ARRAY v_remaining.missing_grants LOOP - BEGIN - EXECUTE v_statement; - EXCEPTION WHEN insufficient_privilege THEN - NULL; - END; - END LOOP; + AND (pg_has_role(current_user, to_regrole(quote_ident(v_remaining.owner)), 'USAGE') + OR v_reassigned) THEN + BEGIN + IF v_reassigned + AND NOT pg_has_role(current_user, to_regrole(quote_ident(v_remaining.owner)), 'USAGE') THEN + PERFORM set_config('role', v_remaining.owner, true); + v_role_changed := true; + END IF; + + FOREACH v_statement IN ARRAY v_remaining.missing_grants LOOP + BEGIN + EXECUTE v_statement; + EXCEPTION WHEN insufficient_privilege THEN + NULL; + END; + END LOOP; + + IF v_role_changed THEN + PERFORM set_config('role', v_original_role, true); + END IF; + EXCEPTION WHEN insufficient_privilege THEN + NULL; + END; END IF; EXCEPTION WHEN lock_not_available THEN v_busy := true; diff --git a/awa-model/src/admin.rs b/awa-model/src/admin.rs index a40efb4e5..296e533de 100644 --- a/awa-model/src/admin.rs +++ b/awa-model/src/admin.rs @@ -9,7 +9,7 @@ use sqlx::Acquire; use sqlx::PgExecutor; use sqlx::PgPool; use std::cmp::max; -use std::collections::HashMap; +use std::collections::{BTreeSet, HashMap}; use std::fmt; use std::str::FromStr; use uuid::Uuid; @@ -381,6 +381,156 @@ pub fn queue_storage_current_jobs_cte(schema: &str) -> String { AND ready.run_lease + 1 = leases.run_lease AND ready.enqueue_shard = leases.enqueue_shard AND ready.lane_seq = leases.lane_seq + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) + UNION ALL + -- Receipt claims that have not materialised into `leases` are + -- running too (#246 / #416). They live on legacy row-local + -- `lease_claims` or, by default, compact `lease_claim_batches`. + -- Report both as `running`, anti-joined against every durable + -- closure/terminal/materialised-lease shape. This mirrors + -- QueueStorage::load_job (the batch-aware single-job admin read) + -- and the canonical open-receipt shape in + -- queue_storage::open_receipt_running_claims_sql; it is spelled + -- out here because this view projects the full job row + -- (kind/created_at/run_at from `ready_entries`), not just claim + -- identity. + SELECT + claims.job_id, + ready.kind, + claims.queue, + 'running'::awa.job_state AS state, + ready.created_at, + ready.run_at, + NULL::timestamptz AS finalized_at + FROM {schema}.lease_claims AS claims + JOIN {schema}.ready_entries AS ready + ON ready.ready_slot = claims.ready_slot + AND ready.ready_generation = claims.ready_generation + AND ready.queue = claims.queue + -- Claims keep the aged effective priority; match the ready attempt. + AND ready.job_id = claims.job_id + AND ready.run_lease + 1 = claims.run_lease + AND ready.enqueue_shard = claims.enqueue_shard + AND ready.lane_seq = claims.lane_seq + -- Moving a ready row tombstones the old one and re-inserts the same + -- attempt, which can reuse the old lane_seq on the new lane. + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) + WHERE claims.closed_at IS NULL + AND NOT EXISTS ( + SELECT 1 FROM {schema}.lease_claim_closures AS cx + WHERE cx.claim_slot = claims.claim_slot + AND cx.job_id = claims.job_id + AND cx.run_lease = claims.run_lease + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.lease_claim_closure_batches AS cb + WHERE cb.receipt_ranges @> claims.receipt_id + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.leases AS lease + WHERE lease.job_id = claims.job_id + AND lease.run_lease = claims.run_lease + ) + -- Terminal supersession counts as closed evidence even + -- without an explicit closure row (mirrors the compact + -- branch below and open_receipt_running_claims_sql); + -- without these a legacy row-local claim would be listed + -- as running AND terminal at once. + AND NOT EXISTS ( + SELECT 1 FROM {schema}.done_entries AS done + WHERE done.job_id = claims.job_id + AND done.run_lease = claims.run_lease + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.deferred_jobs AS deferred + WHERE deferred.job_id = claims.job_id + AND deferred.run_lease = claims.run_lease + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.dlq_entries AS dlq + WHERE dlq.job_id = claims.job_id + AND dlq.run_lease = claims.run_lease + ) + UNION ALL + SELECT + items.job_id, + ready.kind, + batches.queue, + 'running'::awa.job_state AS state, + ready.created_at, + ready.run_at, + NULL::timestamptz AS finalized_at + FROM {schema}.lease_claim_batches AS batches + CROSS JOIN LATERAL unnest( + batches.job_ids, + batches.run_leases, + batches.receipt_ids, + batches.lane_seqs + ) AS items(job_id, run_lease, receipt_id, lane_seq) + JOIN {schema}.ready_entries AS ready + ON ready.ready_slot = batches.ready_slot + AND ready.ready_generation = batches.ready_generation + AND ready.queue = batches.queue + AND ready.enqueue_shard = batches.enqueue_shard + AND ready.lane_seq = items.lane_seq + AND ready.job_id = items.job_id + AND ready.run_lease + 1 = items.run_lease + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) + WHERE NOT EXISTS ( + SELECT 1 FROM {schema}.lease_claim_closures AS cx + WHERE cx.claim_slot = batches.claim_slot + AND cx.job_id = items.job_id + AND cx.run_lease = items.run_lease + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.lease_claim_closure_batches AS cb + WHERE cb.claim_slot = batches.claim_slot + AND cb.receipt_ranges @> items.receipt_id + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.leases AS lease + WHERE lease.job_id = items.job_id + AND lease.run_lease = items.run_lease + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.done_entries AS done + WHERE done.job_id = items.job_id + AND done.run_lease = items.run_lease + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.deferred_jobs AS deferred + WHERE deferred.job_id = items.job_id + AND deferred.run_lease = items.run_lease + ) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.dlq_entries AS dlq + WHERE dlq.job_id = items.job_id + AND dlq.run_lease = items.run_lease + ) UNION ALL SELECT job_id, kind, queue, state, created_at, run_at, finalized_at FROM {schema}.terminal_jobs @@ -734,6 +884,15 @@ fn queue_storage_recent_jobs_cte(schema: &str, since: &str) -> String { AND ready.run_lease + 1 = leases.run_lease AND ready.enqueue_shard = leases.enqueue_shard AND ready.lane_seq = leases.lane_seq + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) WHERE ready.created_at >= {since} UNION ALL -- A queue or priority move copies the row with its times, so @@ -802,48 +961,42 @@ fn queue_storage_recent_jobs_cte(schema: &str, since: &str) -> String { ) } -/// The per-kind, per-queue counts behind [`job_kind_overviews`], as a -/// `kind_counts` CTE of `kind`, `job_count`, `queue_count` and -/// `completed_last_hour`. -/// -/// It counts what `current_jobs` holds without expanding it. Each place a job -/// can live is aggregated to `(kind, queue)` on its own, as a `MATERIALIZED` -/// CTE, and only those few rows are combined. Read through `current_jobs`, -/// the receipt-batch join to `ready_entries` and the `UNION ALL` are both -/// estimated at one row whatever the statistics say, so Postgres chooses -/// nested loops that probe once per retained row, and a join on -/// `queue_claim_heads` is only as good as that table's estimate. -/// -/// Available rows are found per lane as a claim finds them (see -/// [`queue_storage_available_by_queue_sql`]). A completion batch -/// member's kind lives on its ready row, so that leg unnests each batch and -/// probes the ready primary key. -fn queue_storage_kind_counts_cte(schema: &str) -> String { +// Per-source kind/queue counts keep the same joins and predicates as the +// combined projection, while allowing each source's relation locks to end +// before another source opens a different set of retained partitions. +#[derive(Clone, Copy, PartialEq, Eq)] +enum KindCountSource { + Available, + Leased, + Deferred, + Done, + Receipts, + Dlq, +} + +// Generate only this source's relations: unused CTEs would still acquire +// relation locks during parse/plan before execution could discard them. +fn queue_storage_kind_count_source_sql(schema: &str, source: KindCountSource) -> String { let available = queue_storage_available_lane_rows_sql(schema, "heads", "ready.kind", "TRUE"); - format!( - r#" - heads AS MATERIALIZED ( - SELECT - claims.queue, - claims.priority, - claims.enqueue_shard, - {schema}.sequence_next_value(claims.seq_name) AS claim_seq, - {schema}.sequence_next_value(enqueues.seq_name) AS next_seq - FROM {schema}.queue_claim_heads AS claims - JOIN {schema}.queue_enqueue_heads AS enqueues - ON enqueues.queue = claims.queue - AND enqueues.priority = claims.priority - AND enqueues.enqueue_shard = claims.enqueue_shard - ), - available_by_kind AS MATERIALIZED ( + match source { + KindCountSource::Available => format!( + r#" + WITH heads AS MATERIALIZED ( + SELECT queue, priority, enqueue_shard, claim_seq, next_seq + FROM unnest( + $1::text[], $2::smallint[], $3::smallint[], $4::bigint[], $5::bigint[] + ) AS captured(queue, priority, enqueue_shard, claim_seq, next_seq) + ) SELECT lane.kind, heads.queue, count(*)::bigint AS n, 0::bigint AS n_recent FROM heads CROSS JOIN LATERAL ( {available} ) AS lane GROUP BY lane.kind, heads.queue + "# ), - leased_by_kind AS MATERIALIZED ( + KindCountSource::Leased => format!( + r#" SELECT ready.kind, leases.queue, count(*)::bigint AS n, 0::bigint AS n_recent FROM {schema}.leases AS leases JOIN {schema}.ready_entries AS ready @@ -856,14 +1009,27 @@ fn queue_storage_kind_counts_cte(schema: &str) -> String { AND ready.run_lease + 1 = leases.run_lease AND ready.enqueue_shard = leases.enqueue_shard AND ready.lane_seq = leases.lane_seq + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) GROUP BY ready.kind, leases.queue + "# ), - deferred_by_kind AS MATERIALIZED ( + KindCountSource::Deferred => format!( + r#" SELECT kind, queue, count(*)::bigint AS n, 0::bigint AS n_recent FROM {schema}.deferred_jobs GROUP BY kind, queue + "# ), - done_by_kind AS MATERIALIZED ( + KindCountSource::Done => format!( + r#" SELECT kind, queue, @@ -874,8 +1040,10 @@ fn queue_storage_kind_counts_cte(schema: &str) -> String { )::bigint AS n_recent FROM {schema}.done_entries GROUP BY kind, queue + "# ), - receipts_by_kind AS MATERIALIZED ( + KindCountSource::Receipts => format!( + r#" SELECT ready.kind, batch.queue, @@ -889,14 +1057,20 @@ fn queue_storage_kind_counts_cte(schema: &str) -> String { batch.run_leases, batch.lane_seqs ) AS item(job_id, run_lease, lane_seq) - JOIN {schema}.ready_entries AS ready - ON ready.ready_slot = batch.ready_slot - AND ready.ready_generation = batch.ready_generation - AND ready.queue = batch.queue - AND ready.priority = batch.priority - AND ready.enqueue_shard = batch.enqueue_shard - AND ready.lane_seq = item.lane_seq - AND ready.job_id = item.job_id + CROSS JOIN LATERAL ( + SELECT ready.kind + FROM {schema}.ready_entries AS ready + WHERE ready.ready_slot = batch.ready_slot + AND ready.ready_generation = batch.ready_generation + AND ready.queue = batch.queue + AND ready.priority = batch.priority + AND ready.enqueue_shard = batch.enqueue_shard + AND ready.lane_seq = item.lane_seq + AND ready.job_id = item.job_id + -- Keep the full member key in the parameterized lookup, + -- rather than scanning its retained lane for every member. + OFFSET 0 + ) AS ready WHERE NOT EXISTS ( SELECT 1 FROM {schema}.receipt_completion_tombstones AS tomb WHERE tomb.ready_slot = batch.ready_slot @@ -905,22 +1079,25 @@ fn queue_storage_kind_counts_cte(schema: &str) -> String { AND tomb.run_lease = item.run_lease ) GROUP BY ready.kind, batch.queue + "# ), - dlq_by_kind AS MATERIALIZED ( + KindCountSource::Dlq => format!( + r#" SELECT kind, queue, count(*)::bigint AS n, 0::bigint AS n_recent FROM {schema}.dlq_entries GROUP BY kind, queue + "# ), + } +} + +// Merge small (kind, queue) aggregates from separately locked read sources. +fn queue_storage_kind_counts_cte() -> &'static str { + r#" kind_queue AS ( SELECT kind, queue, sum(n)::bigint AS n, sum(n_recent)::bigint AS n_recent - FROM ( - SELECT * FROM available_by_kind - UNION ALL SELECT * FROM leased_by_kind - UNION ALL SELECT * FROM deferred_by_kind - UNION ALL SELECT * FROM done_by_kind - UNION ALL SELECT * FROM receipts_by_kind - UNION ALL SELECT * FROM dlq_by_kind - ) AS legs + FROM unnest($1::text[], $2::text[], $3::bigint[], $4::bigint[]) + AS sources(kind, queue, n, n_recent) GROUP BY kind, queue ), kind_counts AS ( @@ -932,8 +1109,7 @@ fn queue_storage_kind_counts_cte(schema: &str) -> String { FROM kind_queue GROUP BY kind ) - "# - ) + "# } async fn notify_cancellation_tx<'a>( @@ -949,14 +1125,21 @@ async fn notify_cancellation_tx<'a>( Ok(()) } -/// Candidate job ids for [`list_jobs`]: the newest `limit` ids below `cursor` -/// matching `filter`'s state, kind and queue, newest first, with the binds the -/// statement takes in order. +#[derive(Clone, Copy, PartialEq, Eq)] +enum JobCandidateSource { + States, + Receipts(usize), +} + +/// Candidate job ids for [`list_jobs`] from one source: the newest `limit` +/// distinct ids below `cursor`, with the statement's binds in order. State +/// sources apply kind/state/queue filters. Receipt sources enumerate open +/// claims in one physical slot and apply kind/queue before their limit. /// -/// It selects what `current_jobs` would with the same filter, but each place a -/// job can live is read and limited on its own, with only the legs that can -/// hold the requested state, and the filter written into the statement rather -/// than as `$n IS NULL OR` terms that a generic plan can't drop. Read through +/// Each place a job can live is read and limited on its own, with only the +/// state legs that can hold the requested state. Known filters are written +/// into the statement rather than as `$n IS NULL OR` terms that a generic plan +/// can't drop. Read through /// `current_jobs`, Postgres joins every retained completion to a scan of the /// ready rows, because that join and the `UNION ALL` are each estimated at one /// row. Available rows are found per lane as a claim finds them, and a @@ -968,13 +1151,17 @@ fn queue_storage_job_id_candidates_sql( filter: &ListJobsFilter, cursor: Option, limit: i64, + source: JobCandidateSource, ) -> (String, Vec) { let mut binds = Vec::new(); let mut bind = |value: CandidateBind| { binds.push(value); format!("${}", binds.len()) }; - let state = filter.state.map(|state| bind(CandidateBind::State(state))); + let state = filter + .state + .filter(|_| source == JobCandidateSource::States) + .map(|state| bind(CandidateBind::State(state))); let kind = filter .kind .as_ref() @@ -1008,9 +1195,10 @@ fn queue_storage_job_id_candidates_sql( sql }; let holds = |states: &[JobState]| { - filter - .state - .is_none_or(|requested| states.contains(&requested)) + source == JobCandidateSource::States + && filter + .state + .is_none_or(|requested| states.contains(&requested)) }; let mut legs: Vec = Vec::new(); @@ -1049,7 +1237,7 @@ fn queue_storage_job_id_candidates_sql( ); legs.push(format!( r#"( - SELECT lane.job_id + SELECT DISTINCT lane.job_id FROM heads CROSS JOIN LATERAL ({lane_rows}) AS lane ORDER BY lane.job_id DESC @@ -1061,7 +1249,7 @@ fn queue_storage_job_id_candidates_sql( let predicates = predicates(Some("state"), Some("kind"), "queue", "job_id"); legs.push(format!( r#"( - SELECT job_id + SELECT DISTINCT job_id FROM {schema}.deferred_jobs WHERE TRUE{predicates} ORDER BY job_id DESC @@ -1078,7 +1266,7 @@ fn queue_storage_job_id_candidates_sql( ); legs.push(format!( r#"( - SELECT leases.job_id + SELECT DISTINCT leases.job_id FROM {schema}.leases AS leases JOIN {schema}.ready_entries AS ready ON ready.ready_slot = leases.ready_slot @@ -1090,17 +1278,62 @@ fn queue_storage_job_id_candidates_sql( AND ready.run_lease + 1 = leases.run_lease AND ready.enqueue_shard = leases.enqueue_shard AND ready.lane_seq = leases.lane_seq + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) WHERE TRUE{predicates} ORDER BY leases.job_id DESC LIMIT {limit} )"# )); } + if let JobCandidateSource::Receipts(slot) = source { + let open_running = + crate::queue_storage::open_receipt_running_claims_sql_in_slot(schema, slot); + let predicates = predicates( + None, + Some("created.kind"), + "open_claims.queue", + "open_claims.job_id", + ); + legs.push(format!( + 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 + AND ready.run_lease + 1 = open_claims.run_lease + AND ready.queue = open_claims.queue + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) + LIMIT 1 + ) AS created + WHERE TRUE{predicates} + ORDER BY open_claims.job_id DESC + LIMIT {limit} + )"# + )); + } if holds(&[JobState::Completed, JobState::Failed, JobState::Cancelled]) { let predicates = predicates(Some("state"), Some("kind"), "queue", "job_id"); legs.push(format!( r#"( - SELECT job_id + SELECT DISTINCT job_id FROM {schema}.done_entries WHERE TRUE{predicates} ORDER BY job_id DESC @@ -1116,7 +1349,7 @@ fn queue_storage_job_id_candidates_sql( .unwrap_or_default(); legs.push(format!( r#"( - SELECT member.job_id + SELECT DISTINCT member.job_id FROM ( SELECT batch.ready_slot, @@ -1165,7 +1398,7 @@ fn queue_storage_job_id_candidates_sql( let predicates = predicates(None, Some("kind"), "queue", "job_id"); legs.push(format!( r#"( - SELECT job_id + SELECT DISTINCT job_id FROM {schema}.dlq_entries WHERE TRUE{predicates} ORDER BY job_id DESC @@ -1177,7 +1410,7 @@ fn queue_storage_job_id_candidates_sql( let sql = format!( r#" {with} - SELECT job_id + SELECT DISTINCT job_id FROM ({legs}) AS candidates ORDER BY job_id DESC LIMIT {limit} @@ -1196,6 +1429,34 @@ enum CandidateBind { Int(i64), } +async fn fetch_queue_storage_job_candidates( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + sql: &str, + binds: Vec, +) -> Result, sqlx::Error> { + sqlx::query("SAVEPOINT awa_list_candidates") + .execute(tx.as_mut()) + .await?; + let mut query = sqlx::query_scalar::<_, i64>(sql); + for value in binds { + query = match value { + CandidateBind::State(state) => query.bind(state), + CandidateBind::Text(text) => query.bind(text), + CandidateBind::Int(int) => query.bind(int), + }; + } + let ids = query.fetch_all(tx.as_mut()).await?; + // Keep the outer repeatable-read snapshot, but release this read phase's + // relation locks. Destroy the savepoint too rather than stacking names. + sqlx::query("ROLLBACK TO SAVEPOINT awa_list_candidates") + .execute(tx.as_mut()) + .await?; + sqlx::query("RELEASE SAVEPOINT awa_list_candidates") + .execute(tx.as_mut()) + .await?; + Ok(ids) +} + async fn list_queue_storage_jobs( store: &QueueStorage, pool: &PgPool, @@ -1212,23 +1473,65 @@ async fn list_queue_storage_jobs( let mut cursor = filter.before_id; loop { - let (sql, binds) = - queue_storage_job_id_candidates_sql(store.schema(), filter, cursor, candidate_limit); - let mut query = sqlx::query_scalar::<_, i64>(&sql); - for value in binds { - query = match value { - CandidateBind::State(state) => query.bind(state), - CandidateBind::Text(text) => query.bind(text), - CandidateBind::Int(int) => query.bind(int), - }; - } - // As in `queue_overviews`: compiling costs more than running. - let mut tx = pool.begin().await?; + // Every candidate phase sees the same MVCC snapshot. Each phase + // releases its own read locks before the next one opens more tables. + let mut begin_guard = crate::read_transaction::BeginGuard::acquire(pool).await?; + let mut tx = sqlx::Connection::begin_with( + &mut *begin_guard.connection, + "BEGIN ISOLATION LEVEL REPEATABLE READ, READ ONLY", + ) + .await?; + begin_guard.begin_pending = false; sqlx::query("SET LOCAL jit = off") .execute(tx.as_mut()) .await?; - let ids: Vec = query.fetch_all(tx.as_mut()).await?; + let states = filter.state.map_or_else( + || { + vec![ + JobState::Available, + JobState::Scheduled, + JobState::Retryable, + JobState::Running, + JobState::WaitingExternal, + JobState::Completed, + JobState::Failed, + JobState::Cancelled, + ] + }, + |state| vec![state], + ); + let mut ids = BTreeSet::new(); + for state in states { + let mut phase_filter = filter.clone(); + phase_filter.state = Some(state); + let (sql, binds) = queue_storage_job_id_candidates_sql( + store.schema(), + &phase_filter, + cursor, + candidate_limit, + JobCandidateSource::States, + ); + ids.extend(fetch_queue_storage_job_candidates(&mut tx, &sql, binds).await?); + } + if filter.state.is_none_or(|state| state == JobState::Running) { + for slot in queue_storage_populated_claim_slots(&mut tx, store).await? { + let (sql, binds) = queue_storage_job_id_candidates_sql( + store.schema(), + filter, + cursor, + candidate_limit, + JobCandidateSource::Receipts(slot), + ); + ids.extend(fetch_queue_storage_job_candidates(&mut tx, &sql, binds).await?); + } + } tx.commit().await?; + drop(begin_guard); + let ids: Vec = ids + .into_iter() + .rev() + .take(candidate_limit as usize) + .collect(); if ids.is_empty() { break; } @@ -1552,6 +1855,15 @@ fn unique_key_candidate_sql(schema: &str, scoped: bool) -> String { AND ready.run_lease + 1 = leases.run_lease AND ready.enqueue_shard = leases.enqueue_shard AND ready.lane_seq = leases.lane_seq + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) AND ready.unique_key = $1 OFFSET 0 ) AS ready ON TRUE @@ -3121,7 +3433,9 @@ pub async fn queue_overviews(pool: &PgPool) -> Result, AwaErr ); // JIT compiles this statement for longer than it runs. - let mut tx = pool.begin().await?; + let mut begin_guard = crate::read_transaction::BeginGuard::acquire(pool).await?; + let mut tx = sqlx::Connection::begin(&mut *begin_guard.connection).await?; + begin_guard.begin_pending = false; sqlx::query("SET LOCAL jit = off") .execute(tx.as_mut()) .await?; @@ -3129,6 +3443,7 @@ pub async fn queue_overviews(pool: &PgPool) -> Result, AwaErr .fetch_all(tx.as_mut()) .await?; tx.commit().await?; + drop(begin_guard); return Ok(rows); } @@ -3599,12 +3914,15 @@ pub async fn state_counts(pool: &PgPool) -> Result, AwaEr waiting_external, ): (i64, i64, i64, i64, i64, i64, i64, i64) = { // As in `queue_overviews`: compiling costs more than running. - let mut tx = pool.begin().await?; + let mut begin_guard = crate::read_transaction::BeginGuard::acquire(pool).await?; + let mut tx = sqlx::Connection::begin(&mut *begin_guard.connection).await?; + begin_guard.begin_pending = false; sqlx::query("SET LOCAL jit = off") .execute(tx.as_mut()) .await?; let counts = sqlx::query_as(&sql).fetch_one(tx.as_mut()).await?; tx.commit().await?; + drop(begin_guard); counts }; @@ -3654,6 +3972,185 @@ pub async fn state_counts(pool: &PgPool) -> Result, AwaEr Ok(rows.into_iter().collect()) } +#[derive(sqlx::FromRow)] +struct KindCountHead { + queue: String, + priority: i16, + enqueue_shard: i16, + claim_seq: i64, + next_seq: i64, +} + +// Sequence values do not obey MVCC. Capture them with the first read's +// snapshot instead of letting later claims advance them during receipt phases. +async fn queue_storage_kind_heads( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + store: &QueueStorage, +) -> Result, sqlx::Error> { + let schema = store.schema(); + sqlx::query("SAVEPOINT awa_kind_heads") + .execute(tx.as_mut()) + .await?; + let heads = sqlx::query_as::<_, KindCountHead>(&format!( + r#" + SELECT + claims.queue, + claims.priority, + claims.enqueue_shard, + {schema}.sequence_next_value(claims.seq_name) AS claim_seq, + {schema}.sequence_next_value(enqueues.seq_name) AS next_seq + FROM {schema}.queue_claim_heads AS claims + JOIN {schema}.queue_enqueue_heads AS enqueues + ON enqueues.queue = claims.queue + AND enqueues.priority = claims.priority + AND enqueues.enqueue_shard = claims.enqueue_shard + "# + )) + .fetch_all(tx.as_mut()) + .await?; + sqlx::query("ROLLBACK TO SAVEPOINT awa_kind_heads") + .execute(tx.as_mut()) + .await?; + sqlx::query("RELEASE SAVEPOINT awa_kind_heads") + .execute(tx.as_mut()) + .await?; + Ok(heads) +} + +async fn queue_storage_base_kind_counts( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + store: &QueueStorage, + heads: Vec, +) -> Result, sqlx::Error> { + let mut head_queues = Vec::with_capacity(heads.len()); + let mut head_priorities = Vec::with_capacity(heads.len()); + let mut head_shards = Vec::with_capacity(heads.len()); + let mut claim_seqs = Vec::with_capacity(heads.len()); + let mut next_seqs = Vec::with_capacity(heads.len()); + for head in heads { + head_queues.push(head.queue); + head_priorities.push(head.priority); + head_shards.push(head.enqueue_shard); + claim_seqs.push(head.claim_seq); + next_seqs.push(head.next_seq); + } + let mut counts = Vec::new(); + for source in [ + KindCountSource::Available, + KindCountSource::Leased, + KindCountSource::Deferred, + KindCountSource::Done, + KindCountSource::Receipts, + KindCountSource::Dlq, + ] { + let sql = queue_storage_kind_count_source_sql(store.schema(), source); + sqlx::query("SAVEPOINT awa_kind_source") + .execute(tx.as_mut()) + .await?; + let mut query = sqlx::query_as::<_, (String, String, i64, i64)>(&sql); + if source == KindCountSource::Available { + query = query + .bind(&head_queues) + .bind(&head_priorities) + .bind(&head_shards) + .bind(&claim_seqs) + .bind(&next_seqs); + } + counts.extend(query.fetch_all(tx.as_mut()).await?); + sqlx::query("ROLLBACK TO SAVEPOINT awa_kind_source") + .execute(tx.as_mut()) + .await?; + sqlx::query("RELEASE SAVEPOINT awa_kind_source") + .execute(tx.as_mut()) + .await?; + } + Ok(counts) +} + +// The result is bounded by the configured claim-ring slot count, not by +// retained job volume. Include all populated slots, including old generations. +async fn queue_storage_populated_claim_slots( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + store: &QueueStorage, +) -> Result, AwaError> { + let schema = store.schema(); + sqlx::query("SAVEPOINT awa_claim_slots") + .execute(tx.as_mut()) + .await?; + let slots: Vec = sqlx::query_scalar(&format!( + "SELECT claim_slot FROM {schema}.lease_claims WHERE closed_at IS NULL \ + UNION SELECT claim_slot FROM {schema}.lease_claim_batches ORDER BY claim_slot" + )) + .fetch_all(tx.as_mut()) + .await?; + sqlx::query("ROLLBACK TO SAVEPOINT awa_claim_slots") + .execute(tx.as_mut()) + .await?; + sqlx::query("RELEASE SAVEPOINT awa_claim_slots") + .execute(tx.as_mut()) + .await?; + slots + .into_iter() + .map(|slot| { + usize::try_from(slot) + .ok() + .filter(|slot| *slot < store.claim_slot_count()) + .ok_or_else(|| AwaError::Validation(format!("invalid physical claim slot {slot}"))) + }) + .collect() +} + +// Keep only per-kind, per-queue aggregates in Rust, not every claim identity. +async fn queue_storage_running_kind_counts( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + store: &QueueStorage, +) -> Result, AwaError> { + let schema = store.schema(); + let slots = queue_storage_populated_claim_slots(tx, store).await?; + let mut counts = Vec::new(); + for slot in slots { + let open_running = + crate::queue_storage::open_receipt_running_claims_sql_in_slot(schema, slot); + sqlx::query("SAVEPOINT awa_kind_counts") + .execute(tx.as_mut()) + .await?; + let rows: Vec<(String, String, i64)> = sqlx::query_as(&format!( + r#" + SELECT created.kind, open_claims.queue, count(*)::bigint + 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 + AND ready.run_lease + 1 = open_claims.run_lease + AND ready.queue = open_claims.queue + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) + LIMIT 1 + ) AS created + GROUP BY created.kind, open_claims.queue + "# + )) + .fetch_all(tx.as_mut()) + .await?; + sqlx::query("ROLLBACK TO SAVEPOINT awa_kind_counts") + .execute(tx.as_mut()) + .await?; + sqlx::query("RELEASE SAVEPOINT awa_kind_counts") + .execute(tx.as_mut()) + .await?; + counts.extend(rows); + } + Ok(counts) +} + /// Get job-kind overviews for all known kinds. pub async fn job_kind_overviews(pool: &PgPool) -> Result, AwaError> { if let Some(store) = active_queue_storage(pool).await? { @@ -3705,18 +4202,47 @@ pub async fn job_kind_overviews(pool: &PgPool) -> Result, A LEFT JOIN live_kind_descriptor_variants kdv ON kdv.kind = k.kind ORDER BY k.kind "#, - kind_counts_cte = queue_storage_kind_counts_cte(store.schema()) + kind_counts_cte = queue_storage_kind_counts_cte() ); // As in `queue_overviews`: compiling costs more than running. - let mut tx = pool.begin().await?; + let mut begin_guard = crate::read_transaction::BeginGuard::acquire(pool).await?; + let mut tx = sqlx::Connection::begin_with( + &mut *begin_guard.connection, + "BEGIN ISOLATION LEVEL REPEATABLE READ, READ ONLY", + ) + .await?; + begin_guard.begin_pending = false; sqlx::query("SET LOCAL jit = off") .execute(tx.as_mut()) .await?; + let heads = queue_storage_kind_heads(&mut tx, &store).await?; + let mut sources = queue_storage_base_kind_counts(&mut tx, &store, heads).await?; + sources.extend( + queue_storage_running_kind_counts(&mut tx, &store) + .await? + .into_iter() + .map(|(kind, queue, n)| (kind, queue, n, 0)), + ); + let mut kinds = Vec::with_capacity(sources.len()); + let mut queues = Vec::with_capacity(sources.len()); + let mut counts = Vec::with_capacity(sources.len()); + let mut recent_counts = Vec::with_capacity(sources.len()); + for (kind, queue, count, recent) in sources { + kinds.push(kind); + queues.push(queue); + counts.push(count); + recent_counts.push(recent); + } let rows = sqlx::query_as::<_, JobKindOverview>(&sql) + .bind(kinds) + .bind(queues) + .bind(counts) + .bind(recent_counts) .fetch_all(tx.as_mut()) .await?; tx.commit().await?; + drop(begin_guard); return Ok(rows); } @@ -3995,7 +4521,9 @@ pub async fn state_timeseries( ORDER BY bucket", queue_storage_recent_jobs_cte(store.schema(), "$1") ); - let mut tx = pool.begin().await?; + let mut begin_guard = crate::read_transaction::BeginGuard::acquire(pool).await?; + let mut tx = sqlx::Connection::begin(&mut *begin_guard.connection).await?; + begin_guard.begin_pending = false; sqlx::query("SET LOCAL jit = off") .execute(tx.as_mut()) .await?; @@ -4012,6 +4540,7 @@ pub async fn state_timeseries( .fetch_all(tx.as_mut()) .await?; tx.commit().await?; + drop(begin_guard); return Ok(rows .into_iter() diff --git a/awa-model/src/lib.rs b/awa-model/src/lib.rs index 88d658d05..b2e69c45b 100644 --- a/awa-model/src/lib.rs +++ b/awa-model/src/lib.rs @@ -12,6 +12,7 @@ pub mod kind; pub mod migrations; pub mod partitioned_queue; pub mod queue_storage; +mod read_transaction; pub mod reschedule; pub mod storage; pub mod unique; diff --git a/awa-model/src/migrations.rs b/awa-model/src/migrations.rs index d6b4a6734..549dcabef 100644 --- a/awa-model/src/migrations.rs +++ b/awa-model/src/migrations.rs @@ -307,7 +307,10 @@ macro_rules! definer_lane_provisioning_pending_sql { WHERE to_regprocedure('awa.lane_sequence_problems(text,text)') IS NULL \ UNION ALL \ SELECT 'awa.repair_lane_sequence' \ - WHERE to_regprocedure('awa.repair_lane_sequence(text,text)') IS NULL \ + WHERE NOT EXISTS ( \ + SELECT 1 FROM pg_proc AS p \ + WHERE p.oid = to_regprocedure('awa.repair_lane_sequence(text,text)') \ + AND p.prosrc LIKE '%set_config(''role'', v_original_role, true)%') \ UNION ALL \ SELECT 'awa.refresh_lane_sequence_helpers' \ WHERE EXISTS ( \ diff --git a/awa-model/src/queue_storage.rs b/awa-model/src/queue_storage.rs index b253338d3..746cd60cf 100644 --- a/awa-model/src/queue_storage.rs +++ b/awa-model/src/queue_storage.rs @@ -6,7 +6,7 @@ use crate::{InsertParams, JobRow, JobState}; use chrono::TimeDelta; use chrono::{DateTime, Utc}; use sqlx::Acquire; -use sqlx::{PgPool, Postgres}; +use sqlx::{FromRow, PgPool, Postgres, Row}; use std::collections::hash_map::DefaultHasher; use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::hash::{Hash, Hasher}; @@ -380,8 +380,8 @@ pub struct PromoteOutcome { /// contended lane. A batch that selected its full size may have left /// further due rows behind. pub selected: usize, - /// Lanes whose promotion transaction hit `lock_timeout` and rolled back. - /// Their rows stay deferred. + /// Lanes whose locks were held past `lock_timeout`, so their promotion + /// transaction rolled back. Their rows stay deferred. pub contended: Vec, } @@ -539,20 +539,84 @@ fn terminal_delta_child_name(schema: &str, slot: usize) -> String { } fn done_ready_join(schema: &str, done_alias: &str, ready_alias: &str) -> String { + done_ready_join_from( + done_alias, + ready_alias, + &format!("{schema}.ready_entries"), + &format!("{schema}.ready_tombstones"), + ) +} + +// Old writers could store effective priority in a narrow terminal row. Prefer +// the live body of that exact claimed attempt, while retaining same-lease, +// exact-lane bodies tombstoned by an available-job cancellation. Choose one +// body so moved source rows cannot multiply terminal mutations or counters. +fn done_ready_join_from( + done_alias: &str, + ready_alias: &str, + ready_relation: &str, + tombstone_relation: &str, +) -> String { format!( r#" - LEFT JOIN {schema}.ready_entries AS {ready_alias} - ON {ready_alias}.ready_slot = {done_alias}.ready_slot - AND {ready_alias}.ready_generation = {done_alias}.ready_generation - AND {ready_alias}.queue = {done_alias}.queue - AND {ready_alias}.priority = {done_alias}.priority - AND {ready_alias}.enqueue_shard = {done_alias}.enqueue_shard - AND {ready_alias}.lane_seq = {done_alias}.lane_seq + LEFT JOIN LATERAL ( + SELECT candidate.* + FROM ( + SELECT body.*, + ( + body.run_lease + 1 = {done_alias}.run_lease + AND NOT EXISTS ( + SELECT 1 FROM {tombstone_relation} AS tomb + WHERE tomb.ready_slot = body.ready_slot + AND tomb.ready_generation = body.ready_generation + AND tomb.queue = body.queue + AND tomb.priority = body.priority + AND tomb.enqueue_shard = body.enqueue_shard + AND tomb.lane_seq = body.lane_seq + AND tomb.job_id = body.job_id + ) + ) AS live_attempt_body + FROM {ready_relation} AS body + WHERE body.ready_slot = {done_alias}.ready_slot + AND body.ready_generation = {done_alias}.ready_generation + AND body.queue = {done_alias}.queue + AND body.enqueue_shard = {done_alias}.enqueue_shard + AND body.lane_seq = {done_alias}.lane_seq + AND body.job_id = {done_alias}.job_id + AND (body.run_lease = {done_alias}.run_lease + OR body.run_lease + 1 = {done_alias}.run_lease) + ) AS candidate + WHERE candidate.priority = {done_alias}.priority + OR candidate.live_attempt_body + ORDER BY candidate.live_attempt_body DESC, + (candidate.priority = {done_alias}.priority) DESC, + candidate.priority + LIMIT 1 + ) AS {ready_alias} ON TRUE "# ) } +fn done_payload_projection(done_alias: &str, ready_alias: &str) -> String { + // A narrow old terminal payload can lack metadata from its live moved + // body. Supply only metadata defaults; terminal errors/tags/progress and + // explicit metadata overrides remain authoritative. Keep default {} elision. + format!( + r#"CASE + WHEN {done_alias}.args IS NULL + AND {ready_alias}.live_attempt_body + AND COALESCE({ready_alias}.payload -> 'metadata', '{{}}'::jsonb) <> '{{}}'::jsonb + THEN COALESCE({done_alias}.payload, {ready_alias}.payload, '{{}}'::jsonb) + || jsonb_build_object('metadata', + COALESCE({ready_alias}.payload -> 'metadata', '{{}}'::jsonb) + || COALESCE({done_alias}.payload -> 'metadata', '{{}}'::jsonb)) + ELSE COALESCE({done_alias}.payload, {ready_alias}.payload, '{{}}'::jsonb) + END"# + ) +} + fn done_row_projection(done_alias: &str, ready_alias: &str) -> String { + let payload = done_payload_projection(done_alias, ready_alias); format!( r#" {done_alias}.ready_slot, @@ -574,7 +638,7 @@ fn done_row_projection(done_alias: &str, ready_alias: &str) -> String { COALESCE({done_alias}.created_at, {ready_alias}.created_at, {done_alias}.finalized_at) AS created_at, COALESCE({done_alias}.unique_key, {ready_alias}.unique_key) AS unique_key, COALESCE({done_alias}.unique_states, {ready_alias}.unique_states) AS unique_states, - COALESCE({done_alias}.payload, {ready_alias}.payload, '{{}}'::jsonb) AS payload + {payload} AS payload "# ) } @@ -609,12 +673,16 @@ fn claim_closure_batch_child_name(schema: &str, slot: usize) -> String { format!("{schema}.lease_claim_closure_batches_{slot}") } -fn ring_slot_index(slot: i32, slot_count: usize, ring: &str) -> Result { - let index = usize::try_from(slot).map_err(|_| { +fn ring_slot_number(slot: i32, ring: &str) -> Result { + usize::try_from(slot).map_err(|_| { AwaError::Validation(format!( "invalid {ring} ring slot {slot}: slot must be non-negative" )) - })?; + }) +} + +fn ring_slot_index(slot: i32, slot_count: usize, ring: &str) -> Result { + let index = ring_slot_number(slot, ring)?; if index >= slot_count { return Err(AwaError::Validation(format!( "invalid {ring} ring slot {slot}: configured slot count is {slot_count}" @@ -623,42 +691,76 @@ fn ring_slot_index(slot: i32, slot_count: usize, ring: &str) -> Result String { - // `receipt_id` is allocated from a global sequence, so it identifies a - // claim without a `claim_slot` predicate. Compact closure batches retain - // the raw receipt array for audit/debugging, but hot membership checks use - // the per-child GiST index on the derived multirange. + receipt_evidence_sql( + schema, + closure_rel, + closure_batch_rel, + claims_alias, + ReceiptEvidence::Closed, + ) +} + +fn receipt_evidence_sql( + schema: &str, + closure_rel: &str, + closure_batch_rel: &str, + claims_alias: &str, + evidence: ReceiptEvidence, +) -> String { + let (null_test, join, exists, range_fence) = match evidence { + ReceiptEvidence::Closed => ("NOT NULL", "OR", "EXISTS", ""), + ReceiptEvidence::Open => ( + "NULL", + "AND", + "NOT EXISTS", + " + -- Keep receipt membership correlated while the other closure + -- exclusions can use full-identity anti joins. + OFFSET 0", + ), + }; + // `receipt_id` is allocated from a global sequence and identifies a claim. + // Closures retain that claim's slot; match it to prune unrelated partitions. + // Compact closure batches retain the raw receipt array for audit/debugging, + // but hot membership checks use the per-child GiST index on the multirange. format!( r#" ( - {claims_alias}.closed_at IS NOT NULL - OR EXISTS ( + {claims_alias}.closed_at IS {null_test} + {join} {exists} ( SELECT 1 FROM {closure_rel} AS closures WHERE closures.claim_slot = {claims_alias}.claim_slot AND closures.job_id = {claims_alias}.job_id AND closures.run_lease = {claims_alias}.run_lease ) - OR EXISTS ( + {join} {exists} ( SELECT 1 FROM {closure_batch_rel} AS closure_batches - WHERE closure_batches.receipt_ranges @> {claims_alias}.receipt_id + WHERE closure_batches.claim_slot = {claims_alias}.claim_slot + AND closure_batches.receipt_ranges @> {claims_alias}.receipt_id{range_fence} ) - OR EXISTS ( + {join} {exists} ( SELECT 1 FROM {schema}.done_entries AS done WHERE done.job_id = {claims_alias}.job_id AND done.run_lease = {claims_alias}.run_lease ) - OR EXISTS ( + {join} {exists} ( SELECT 1 FROM {schema}.deferred_jobs AS deferred WHERE deferred.job_id = {claims_alias}.job_id AND deferred.run_lease = {claims_alias}.run_lease ) - OR EXISTS ( + {join} {exists} ( SELECT 1 FROM {schema}.dlq_entries AS dlq WHERE dlq.job_id = {claims_alias}.job_id AND dlq.run_lease = {claims_alias}.run_lease @@ -670,12 +772,11 @@ fn receipt_closed_evidence_sql( /// A subquery that yields the identity `(job_id, run_lease, queue)` of every /// receipt claim currently open (i.e. `running`) but not yet materialised -/// into `leases`. Since #246 receipt claims — zero-deadline and -/// deadline-backed alike — are written as compact `lease_claim_batches` rows -/// rather than one `lease_claims` row per job, so the "open receipt claims" -/// set is the union of legacy row-local `lease_claims` and the unnested -/// members of `lease_claim_batches`, each anti-joined against every durable -/// closure/terminal/materialised-lease evidence shape. +/// into `leases`. Zero-deadline receipts use compact `lease_claim_batches`; +/// deadline-backed receipts retain row-local `lease_claims` for indexed +/// deadline rescue. The open set unions those row-local claims with the +/// unnested compact members, excluding every durable closure, terminal, and +/// materialised-lease evidence shape. /// /// This is the canonical open-receipt-claim shape (mirrors `queue_counts`'s /// `live_running` derivation). Compose it wherever a running enumeration must @@ -699,10 +800,41 @@ pub(crate) fn open_receipt_running_claims_sql(schema: &str) -> String { /// retained batch header is read, but 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 { - let closure_rel = format!("{schema}.lease_claim_closures"); - let closure_batch_rel = format!("{schema}.lease_claim_closure_batches"); - let row_closed = - receipt_closed_evidence_sql(schema, &closure_rel, &closure_batch_rel, "claims"); + open_receipt_running_claims_sql_from(schema, job_ids, None) +} + +/// The same open-claim projection, scoped to one validated physical claim slot. +/// Ready and terminal history are deliberately not restricted to that slot. +pub(crate) fn open_receipt_running_claims_sql_in_slot(schema: &str, slot: usize) -> String { + open_receipt_running_claims_sql_from(schema, None, Some(slot)) +} + +fn open_receipt_running_claims_sql_from( + schema: &str, + job_ids: Option<&str>, + slot: Option, +) -> String { + let (claims_rel, batches_rel, closure_rel, closure_batch_rel) = match slot { + Some(slot) => ( + claim_child_name(schema, slot), + claim_batch_child_name(schema, slot), + closure_child_name(schema, slot), + claim_closure_batch_child_name(schema, slot), + ), + None => ( + format!("{schema}.lease_claims"), + format!("{schema}.lease_claim_batches"), + format!("{schema}.lease_claim_closures"), + format!("{schema}.lease_claim_closure_batches"), + ), + }; + let row_open = receipt_evidence_sql( + schema, + &closure_rel, + &closure_batch_rel, + "claims", + ReceiptEvidence::Open, + ); let (row_filter, batch_filter) = match job_ids { Some(ids) => ( format!("claims.job_id = ANY({ids}) AND "), @@ -713,8 +845,8 @@ pub(crate) fn open_receipt_running_claims_sql_for(schema: &str, job_ids: Option< format!( r#" SELECT claims.job_id, claims.run_lease, claims.queue - FROM {schema}.lease_claims AS claims - WHERE {row_filter}NOT {row_closed} + FROM {claims_rel} AS claims + WHERE {row_filter}{row_open} AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease WHERE lease.job_id = claims.job_id @@ -722,7 +854,7 @@ pub(crate) fn open_receipt_running_claims_sql_for(schema: &str, job_ids: Option< ) UNION ALL SELECT items.job_id, items.run_lease, batches.queue - FROM {schema}.lease_claim_batches AS batches + FROM {batches_rel} AS batches CROSS JOIN LATERAL unnest( batches.job_ids, batches.run_leases, @@ -749,6 +881,9 @@ pub(crate) fn open_receipt_running_claims_sql_for(schema: &str, job_ids: Option< SELECT 1 FROM {schema}.lease_claim_closure_batches 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 + -- multiply receipt members by retained closure batches. + OFFSET 0 ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease @@ -1608,6 +1743,58 @@ impl ReadyJobRow { } } +#[derive(Debug, PartialEq, Eq, sqlx::FromRow)] +struct ReadyRowIdentity { + ready_slot: i32, + ready_generation: i64, + queue: String, + priority: i16, + enqueue_shard: i16, + lane_seq: i64, + job_id: i64, + run_lease: i64, +} + +/// Attempts a rescue pass leaves for a later pass, as `(job_id, run_lease)`. +struct BusyRetained(Vec<(i64, i64)>); + +impl BusyRetained { + fn exclude(self, excluded: &mut BTreeSet<(i64, i64)>) -> Result<(), AwaError> { + let previous = excluded.len(); + excluded.extend(self.0); + if excluded.len() == previous { + return Err(AwaError::Validation( + "queue storage retained attempt retry made no progress".to_string(), + )); + } + Ok(()) + } +} + +/// A failure hydrating a deleted lease. `Ambiguous` names an attempt whose +/// retained ready bodies rank equally, so none can be chosen for it. +enum HydrateError { + Ambiguous { job_id: i64, run_lease: i64 }, + Awa(AwaError), +} + +impl From for HydrateError { + fn from(err: AwaError) -> Self { + Self::Awa(err) + } +} + +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::Awa(err) => err, + } + } +} + #[derive(Debug, Clone, sqlx::FromRow)] struct ReadyTransitionRow { ready_slot: i32, @@ -1989,6 +2176,15 @@ struct DoneJobRow { payload: serde_json::Value, } +// Keep the stored terminal priority for decrementing its old counter key. +// Manual retry and terminal-to-DLQ moves recover the retained physical lane. +#[derive(Debug, sqlx::FromRow)] +struct HydratedTerminalRow { + #[sqlx(flatten)] + terminal: DoneJobRow, + body_priority: i16, +} + impl DoneJobRow { fn into_job_row(self) -> Result { let payload = RuntimePayload::from_json(self.payload)?; @@ -2209,6 +2405,7 @@ struct ReadySnapshotRow { kind: String, queue: String, args: serde_json::Value, + priority: i16, lane_seq: i64, enqueue_shard: i16, run_at: DateTime, @@ -2216,6 +2413,55 @@ struct ReadySnapshotRow { unique_key: Option>, unique_states: Option, payload: serde_json::Value, + run_lease: i64, + live_body: bool, +} + +type ReadyBodyKey = (i32, i64, String, i16, i64, i64, i64); +type ReadyBodyLaneKey = (i32, i64, String, i16, i64, i64, i64, i16); + +fn select_claimed_ready_body<'a>( + candidates: &'a [ReadySnapshotRow], + claimed_lanes: &BTreeSet, + reference: &DeletedLeaseRow, +) -> Result<(&'a ReadySnapshotRow, bool), HydrateError> { + // Emitted physical-lane evidence wins over a live replacement. + // Without it, retain the live-body compatibility fallback; it + // cannot identify a legacy claim's original colliding lane. + let emitted = |row: &ReadySnapshotRow| { + claimed_lanes.contains(&( + row.ready_slot, + row.ready_generation, + row.queue.clone(), + row.enqueue_shard, + row.lane_seq, + row.job_id, + row.run_lease + 1, + row.priority, + )) + }; + let has_emitted = candidates.iter().any(&emitted); + let rank = |row: &ReadySnapshotRow| { + if has_emitted { + emitted(row) + } else { + row.live_body + } + }; + let strongest = candidates + .iter() + .map(&rank) + .max() + .expect("ready body group is nonempty"); + 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, + }); + } + Ok((ready, has_emitted)) } #[derive(Debug, Clone, sqlx::FromRow)] @@ -2240,6 +2486,7 @@ struct LeaseTransitionRow { args: serde_json::Value, state: JobState, priority: i16, + ready_priority: i16, attempt: i16, run_lease: i64, max_attempts: i16, @@ -2252,6 +2499,7 @@ struct LeaseTransitionRow { unique_states: Option, payload: serde_json::Value, progress: Option, + self_contained_terminal: bool, } impl LeaseTransitionRow { @@ -2269,11 +2517,19 @@ impl LeaseTransitionRow { queue: self.queue, args: self.args, state, - priority: self.priority, + // Terminal rows retain their ready-lane identity for body hydration. + // The effective attempt priority is still used by deferred/DLQ rows. + priority: self.ready_priority, attempt: self.attempt, run_lease: self.run_lease, max_attempts: self.max_attempts, - lane_seq: self.lane_seq, + // A competing live copy must not replace this recovered body + // when later terminal reads resolve retained ready rows. + lane_seq: if self.self_contained_terminal { + -self.job_id + } else { + self.lane_seq + }, enqueue_shard: self.enqueue_shard, run_at: self.run_at, attempted_at: self.attempted_at, @@ -7747,7 +8003,7 @@ impl QueueStorage { let done_projection = done_row_projection("done", "ready"); let ready_join = done_ready_join(schema, "done", "ready"); - let terminal: Option = sqlx::query_as(&format!( + let terminal: Option = sqlx::query_as(&format!( r#" WITH deleted AS ( DELETE FROM {schema}.done_entries @@ -7762,7 +8018,7 @@ impl QueueStorage { ) RETURNING * ) - SELECT {done_projection} + SELECT {done_projection}, COALESCE(ready.priority, done.priority) AS body_priority FROM deleted AS done {ready_join} "# @@ -7772,7 +8028,9 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; - if let Some(terminal) = terminal { + if let Some(hydrated) = terminal { + let body_priority = hydrated.body_priority; + let terminal = hydrated.terminal; self.ensure_terminal_removed_receipt_closures_tx(tx, std::slice::from_ref(&terminal)) .await?; // The DELETE FROM done_entries above removes one terminal row; @@ -7787,7 +8045,7 @@ impl QueueStorage { kind: terminal.kind, queue: terminal.queue.clone(), args: terminal.args, - priority: terminal.priority, + priority: body_priority, attempt: 0, run_lease: terminal.run_lease, max_attempts: terminal.max_attempts, @@ -8268,6 +8526,72 @@ impl QueueStorage { Ok(newest) } + async fn cancel_materialized_job_lease_tx<'a>( + &self, + tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + job_id: i64, + ) -> Result, AwaError> { + let schema = self.schema(); + let deleted_lease: Vec = sqlx::query_as(&format!( + r#" + DELETE FROM {schema}.leases + WHERE job_id = $1 + AND state IN ('running', 'waiting_external') + RETURNING + ready_slot, + ready_generation, + job_id, + queue, + state, + priority, + attempt, + run_lease, + max_attempts, + lane_seq, + enqueue_shard, + heartbeat_at, + deadline_at, + attempted_at, + callback_id, + callback_timeout_at + "# + )) + .bind(job_id) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + + if !deleted_lease.is_empty() { + let lease = self + .hydrate_deleted_leases_tx(tx, deleted_lease) + .await? + .into_iter() + .next() + .expect("deleted running lease"); + let done_payload = + Self::payload_with_attempt_state(lease.payload.clone(), lease.progress.clone())?; + let done = lease + .clone() + .into_done_row(JobState::Cancelled, Utc::now(), done_payload); + self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(lease.state)) + .await?; + // Receipt-plane consistency: close any matching open + // receipt so the ADR-023 anti-join no longer considers this + // attempt live, and rescue doesn't try to revive it. + self.close_receipt_tx(tx, lease.job_id, lease.run_lease, "cancelled") + .await?; + // Wake any worker currently executing this attempt. + self.notify_cancellation_tx(tx, lease.job_id, lease.run_lease) + .await?; + return Ok(Some(CancelJobTxResult { + row: done.into_job_row()?, + claim_cursor_advance: None, + })); + } + + Ok(None) + } + async fn cancel_job_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, @@ -8275,6 +8599,56 @@ impl QueueStorage { ) -> Result, AwaError> { let schema = self.schema(); self.lock_claim_heads_of_jobs_tx(tx, &[job_id]).await?; + if let Some(result) = self.cancel_materialized_job_lease_tx(tx, job_id).await? { + return Ok(Some(result)); + } + + if self.lease_claim_receipts() { + if let Some(run_lease) = self.open_receipt_run_lease_tx(tx, job_id).await? { + self.lock_receipt_attempt_keys_tx(tx, &[(job_id, run_lease)]) + .await?; + let receipt = self + .lock_open_receipt_tx(tx, job_id) + .await? + .filter(|receipt| receipt.1 == run_lease); + if receipt.is_some() { + if let Some(lease) = self + .close_open_receipt_claim_tx(tx, job_id, run_lease, "cancelled") + .await? + { + let done_payload = Self::payload_with_attempt_state( + lease.payload.clone(), + lease.progress.clone(), + )?; + let done = + lease.into_done_row(JobState::Cancelled, Utc::now(), done_payload); + self.insert_done_rows_tx( + tx, + std::slice::from_ref(&done), + Some(JobState::Running), + ) + .await?; + sqlx::query(&format!( + "DELETE FROM {schema}.leases WHERE job_id = $1 AND run_lease = $2" + )) + .bind(job_id) + .bind(run_lease) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + self.notify_cancellation_tx(tx, job_id, run_lease).await?; + return Ok(Some(CancelJobTxResult { + row: done.into_job_row()?, + claim_cursor_advance: None, + })); + } + } + } + if let Some(result) = self.cancel_materialized_job_lease_tx(tx, job_id).await? { + return Ok(Some(result)); + } + } + let claimed = self.ready_row_claimed_sql("ready"); let ready: Option = sqlx::query_as(&format!( r#" @@ -8364,56 +8738,69 @@ impl QueueStorage { })); } - let deleted_lease: Vec = sqlx::query_as(&format!( + let deferred: Option = sqlx::query_as(&format!( r#" - DELETE FROM {schema}.leases + DELETE FROM {schema}.deferred_jobs WHERE job_id = $1 - AND state IN ('running', 'waiting_external') + AND state IN ('scheduled', 'retryable') RETURNING - ready_slot, - ready_generation, job_id, + kind, queue, + args, state, priority, attempt, run_lease, max_attempts, - lane_seq, - enqueue_shard, - heartbeat_at, - deadline_at, + run_at, attempted_at, - callback_id, - callback_timeout_at + finalized_at, + created_at, + unique_key, + unique_states, + COALESCE(payload, '{{}}'::jsonb) AS payload "# )) .bind(job_id) - .fetch_all(tx.as_mut()) + .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; - if !deleted_lease.is_empty() { - let lease = self - .hydrate_deleted_leases_tx(tx, deleted_lease) - .await? - .into_iter() - .next() - .expect("deleted running lease"); - let done_payload = - Self::payload_with_attempt_state(lease.payload.clone(), lease.progress.clone())?; - let done = lease - .clone() - .into_done_row(JobState::Cancelled, Utc::now(), done_payload); - self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(lease.state)) - .await?; - // Receipt-plane consistency: close any matching open - // receipt so the ADR-023 anti-join no longer considers this - // attempt live, and rescue doesn't try to revive it. - self.close_receipt_tx(tx, lease.job_id, lease.run_lease, "cancelled") + if let Some(deferred) = deferred { + let (ready_slot, ready_generation) = self.current_queue_ring(tx).await?; + // A deferred-cancel never observed a claim, so it has no + // shard assignment to inherit. The synthesized terminal row + // is parked on shard 0 with a synthetic negative `lane_seq` + // that keeps the `done_entries` PK + // `(ready_slot, queue, priority, enqueue_shard, lane_seq)` + // unique. `ensure_lane` registers shard 0 for the queue so + // the lane row exists regardless of producer activity. + self.ensure_lane(tx, &deferred.queue, deferred.priority, 0) .await?; - // Wake any worker currently executing this attempt. - self.notify_cancellation_tx(tx, lease.job_id, lease.run_lease) + let done = DoneJobRow { + ready_slot, + ready_generation, + job_id: deferred.job_id, + kind: deferred.kind, + queue: deferred.queue.clone(), + args: deferred.args, + state: JobState::Cancelled, + priority: deferred.priority, + attempt: deferred.attempt, + run_lease: deferred.run_lease, + max_attempts: deferred.max_attempts, + lane_seq: -deferred.job_id, + enqueue_shard: 0, + run_at: deferred.run_at, + attempted_at: deferred.attempted_at, + finalized_at: Utc::now(), + created_at: deferred.created_at, + unique_key: deferred.unique_key, + unique_states: deferred.unique_states, + payload: deferred.payload, + }; + self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(deferred.state)) .await?; return Ok(Some(CancelJobTxResult { row: done.into_job_row()?, @@ -8421,278 +8808,8 @@ impl QueueStorage { })); } - // ADR-023 receipt-only cancel: the job may be running on a - // receipt-backed short path that never materialized a `leases` - // row. Find it by anti-joining lease_claims with - // lease_claim_closures, cancel it by writing a closure and a - // done row, and notify listening workers. - if self.lease_claim_receipts() { - // A fast completion closes an attempt under its completion lock - // without locking claim rows, so a cancel that only locked the - // rows could commit alongside it. Find the open attempt with plain - // reads, take its completion lock, then look again: a completion - // that got there first has committed its closure by then. - let receipt = match self.open_receipt_run_lease_tx(tx, job_id).await? { - Some(run_lease) => { - self.lock_receipt_attempt_keys_tx(tx, &[(job_id, run_lease)]) - .await?; - self.lock_open_receipt_tx(tx, job_id) - .await? - .filter(|receipt| receipt.1 == run_lease) - } - None => None, - }; - - if let Some(( - claim_slot, - run_lease, - ready_slot, - ready_generation, - queue, - priority, - attempt, - max_attempts, - lane_seq, - claimed_at, - receipt_id, - compact_batch, - )) = receipt - { - // Hydrate the ready row so we can synthesize the done - // row with the original args/payload. - let ready_match: Option = sqlx::query_as(&format!( - r#" - SELECT - 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, - COALESCE(payload, '{{}}'::jsonb) AS payload - FROM {schema}.ready_entries - WHERE job_id = $1 - AND ready_slot = $2 - AND ready_generation = $3 - AND queue = $4 - AND lane_seq = $5 - "# - )) - .bind(job_id) - .bind(ready_slot) - .bind(ready_generation) - .bind(&queue) - .bind(lane_seq) - .fetch_optional(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - - let Some(ready) = ready_match else { - // Shouldn't happen — the claim references a ready - // row. Fall through to the deferred / not-found - // branches. - return Ok(None); - }; - - let done = DoneJobRow { - ready_slot, - ready_generation, - job_id, - kind: ready.kind, - queue: queue.clone(), - args: ready.args, - state: JobState::Cancelled, - priority, - attempt, - run_lease, - max_attempts, - lane_seq, - enqueue_shard: ready.enqueue_shard, - run_at: ready.run_at, - attempted_at: Some(claimed_at), - finalized_at: Utc::now(), - created_at: ready.created_at, - unique_key: ready.unique_key, - unique_states: ready.unique_states, - payload: ready.payload, - }; - self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(JobState::Running)) - .await?; - // Write the closure into the same claim partition. A - // compact batch-sourced claim has no lease_claims row, so - // it closes into the batch ledger the queue prune gate - // counts via compact_count; a deadline lease claim keeps - // its explicit closure so the gate balances it against the - // lease_claims row. - if compact_batch { - sqlx::query(&format!( - r#" - INSERT INTO {schema}.lease_claim_closure_batches ( - claim_slot, - ready_slot, - ready_generation, - outcome, - closed_count, - receipt_ids, - receipt_ranges, - closed_at - ) - VALUES ( - $1, - $2, - $3, - 'cancelled', - 1, - ARRAY[$4::bigint], - int8multirange(int8range($4::bigint, $4::bigint + 1, '[)')), - clock_timestamp() - ) - "# - )) - .bind(claim_slot) - .bind(ready_slot) - .bind(ready_generation) - .bind(receipt_id) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - } else { - sqlx::query(&format!( - r#" - WITH inserted AS ( - INSERT INTO {schema}.lease_claim_closures (claim_slot, job_id, run_lease, outcome, closed_at) - VALUES ($1, $2, $3, 'cancelled', clock_timestamp()) - ON CONFLICT (claim_slot, job_id, run_lease) DO NOTHING - RETURNING claim_slot, job_id, run_lease, closed_at - ), - marked AS ( - UPDATE {schema}.lease_claims AS claims - SET closed_at = COALESCE(claims.closed_at, inserted.closed_at) - FROM inserted - WHERE claims.claim_slot = inserted.claim_slot - AND claims.job_id = inserted.job_id - AND claims.run_lease = inserted.run_lease - RETURNING claims.job_id - ) - SELECT count(*) FROM marked - "# - )) - .bind(claim_slot) - .bind(job_id) - .bind(run_lease) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - } - // Defensive: between the leases DELETE at the top of - // this function and the FOR UPDATE on claims above, a - // concurrent `ensure_running_leases_from_receipts_tx` - // can have materialized a `leases` row for this - // (job_id, run_lease). Materialize and we both lock the - // same claim evidence; whichever ran first commits, the - // other replays under the new snapshot. If materialize - // committed first, that lease is now an orphan pointing - // at a job we're about to mark `cancelled`. Sweep it - // defensively. If no race occurred this is a no-op. - sqlx::query(&format!( - "DELETE FROM {schema}.leases WHERE job_id = $1 AND run_lease = $2" - )) - .bind(job_id) - .bind(run_lease) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - self.notify_cancellation_tx(tx, job_id, run_lease).await?; - return Ok(Some(CancelJobTxResult { - row: done.into_job_row()?, - claim_cursor_advance: None, - })); - } - } - - let deferred: Option = sqlx::query_as(&format!( - r#" - DELETE FROM {schema}.deferred_jobs - WHERE job_id = $1 - AND state IN ('scheduled', 'retryable') - RETURNING - job_id, - kind, - queue, - args, - state, - priority, - attempt, - run_lease, - max_attempts, - run_at, - attempted_at, - finalized_at, - created_at, - unique_key, - unique_states, - COALESCE(payload, '{{}}'::jsonb) AS payload - "# - )) - .bind(job_id) - .fetch_optional(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - - if let Some(deferred) = deferred { - let (ready_slot, ready_generation) = self.current_queue_ring(tx).await?; - // A deferred-cancel never observed a claim, so it has no - // shard assignment to inherit. The synthesized terminal row - // is parked on shard 0 with a synthetic negative `lane_seq` - // that keeps the `done_entries` PK - // `(ready_slot, queue, priority, enqueue_shard, lane_seq)` - // unique. `ensure_lane` registers shard 0 for the queue so - // the lane row exists regardless of producer activity. - self.ensure_lane(tx, &deferred.queue, deferred.priority, 0) - .await?; - let done = DoneJobRow { - ready_slot, - ready_generation, - job_id: deferred.job_id, - kind: deferred.kind, - queue: deferred.queue.clone(), - args: deferred.args, - state: JobState::Cancelled, - priority: deferred.priority, - attempt: deferred.attempt, - run_lease: deferred.run_lease, - max_attempts: deferred.max_attempts, - lane_seq: -deferred.job_id, - enqueue_shard: 0, - run_at: deferred.run_at, - attempted_at: deferred.attempted_at, - finalized_at: Utc::now(), - created_at: deferred.created_at, - unique_key: deferred.unique_key, - unique_states: deferred.unique_states, - payload: deferred.payload, - }; - self.insert_done_rows_tx(tx, std::slice::from_ref(&done), Some(deferred.state)) - .await?; - return Ok(Some(CancelJobTxResult { - row: done.into_job_row()?, - claim_cursor_advance: None, - })); - } - - Ok(None) - } + Ok(None) + } pub async fn cancel_job(&self, pool: &PgPool, job_id: i64) -> Result, AwaError> { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; @@ -9134,8 +9251,19 @@ impl QueueStorage { &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, job_ids: &[i64], + ) -> Result<(), AwaError> { + self.lock_claim_heads_of_jobs_with_wait_tx(tx, job_ids, false) + .await + } + + async fn lock_claim_heads_of_jobs_with_wait_tx( + &self, + tx: &mut sqlx::Transaction<'_, Postgres>, + job_ids: &[i64], + nowait: bool, ) -> Result<(), AwaError> { let schema = self.schema(); + let wait_clause = if nowait { "NOWAIT" } else { "" }; sqlx::query(&format!( r#" SELECT 1 @@ -9146,7 +9274,7 @@ impl QueueStorage { WHERE ready.job_id = ANY($1) ) ORDER BY claims.queue COLLATE "C", claims.priority, claims.enqueue_shard - FOR UPDATE OF claims + FOR UPDATE OF claims {wait_clause} "# )) .bind(job_ids) @@ -9688,12 +9816,69 @@ impl QueueStorage { .map_err(|err| AwaError::Validation(format!("invalid aging interval: {err}")))?; let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; + let claimed = self.ready_row_claimed_sql("ready"); + + // Claims commit before their cursor advance. Select only a bounded + // set of identities before taking heads, without locking ready rows in + // the reverse order to a claimer. A later statement rechecks these same + // identities after every in-flight claim on their lanes has committed. + let candidates: Vec = sqlx::query_as(&format!( + r#" + SELECT ready.ready_slot, ready.ready_generation, ready.queue, + ready.priority, ready.enqueue_shard, ready.lane_seq, + ready.job_id, ready.run_lease + FROM {schema}.ready_entries AS ready + JOIN {schema}.queue_claim_heads AS claims + ON claims.queue = ready.queue + AND claims.priority = ready.priority + AND claims.enqueue_shard = ready.enqueue_shard + WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name) + AND ready.priority > 1 + AND ready.run_at <= $1 + AND NOT ({claimed}) + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + AND tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + ) + ORDER BY ready.run_at ASC, ready.lane_seq ASC + LIMIT $2 + "# + )) + .bind(cutoff) + .bind(limit) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + if candidates.is_empty() { + tx.commit().await.map_err(map_sqlx_error)?; + return Ok(Vec::new()); + } + let job_ids: Vec = candidates.iter().map(|row| row.job_id).collect(); + self.lock_claim_heads_of_jobs_tx(&mut tx, &job_ids).await?; let moved: Vec = sqlx::query_as(&format!( r#" - WITH target AS ( + WITH candidates(ready_slot, ready_generation, queue, priority, + enqueue_shard, lane_seq, job_id, run_lease) AS ( + SELECT * FROM unnest($3::int[], $4::bigint[], $5::text[], $6::smallint[], + $7::smallint[], $8::bigint[], $9::bigint[], $10::bigint[]) + ), target AS ( SELECT ready.* - FROM {schema}.ready_entries AS ready + FROM candidates + JOIN {schema}.ready_entries AS ready + ON ready.ready_slot = candidates.ready_slot + AND ready.ready_generation = candidates.ready_generation + AND ready.queue = candidates.queue + AND ready.priority = candidates.priority + AND ready.enqueue_shard = candidates.enqueue_shard + AND ready.lane_seq = candidates.lane_seq + AND ready.job_id = candidates.job_id + AND ready.run_lease = candidates.run_lease JOIN {schema}.queue_claim_heads AS claims ON claims.queue = ready.queue AND claims.priority = ready.priority @@ -9701,6 +9886,7 @@ impl QueueStorage { WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name) AND ready.priority > 1 AND ready.run_at <= $1 + AND NOT ({claimed}) AND NOT EXISTS ( SELECT 1 FROM {schema}.ready_tombstones AS tomb WHERE tomb.queue = ready.queue @@ -9746,6 +9932,14 @@ impl QueueStorage { )) .bind(cutoff) .bind(limit) + .bind(candidates.iter().map(|row| row.ready_slot).collect::>()) + .bind(candidates.iter().map(|row| row.ready_generation).collect::>()) + .bind(candidates.iter().map(|row| row.queue.clone()).collect::>()) + .bind(candidates.iter().map(|row| row.priority).collect::>()) + .bind(candidates.iter().map(|row| row.enqueue_shard).collect::>()) + .bind(candidates.iter().map(|row| row.lane_seq).collect::>()) + .bind(&job_ids) + .bind(candidates.iter().map(|row| row.run_lease).collect::>()) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; @@ -10468,15 +10662,17 @@ impl QueueStorage { Ok(updated as usize) } - async fn hydrate_deleted_leases_tx<'a>( + async fn load_ready_snapshot_candidates_tx( &self, - tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, - deleted: Vec, - ) -> Result, AwaError> { - if deleted.is_empty() { - return Ok(Vec::new()); - } - + tx: &mut sqlx::Transaction<'_, Postgres>, + deleted: &[DeletedLeaseRow], + ) -> Result< + ( + BTreeMap>, + BTreeSet, + ), + AwaError, + > { let schema = self.schema(); let ready_slots: Vec = deleted.iter().map(|row| row.ready_slot).collect(); let ready_generations: Vec = deleted.iter().map(|row| row.ready_generation).collect(); @@ -10486,10 +10682,21 @@ impl QueueStorage { let job_ids: Vec = deleted.iter().map(|row| row.job_id).collect(); let run_leases: Vec = deleted.iter().map(|row| row.run_lease).collect(); + // Resolve moved ready bodies without opening the whole tombstone ring. + let slots = ready_slots + .iter() + .map(|&slot| ring_slot_number(slot, "queue")) + .collect::, _>>()?; + let tombstone_sources = slots + .iter() + .map(|&slot| format!("SELECT * FROM {}", ready_tombstone_child_name(schema, slot))) + .collect::>() + .join(" UNION ALL "); + let ready_rows: Vec = sqlx::query_as(&format!( r#" - WITH refs(ready_slot, ready_generation, queue, enqueue_shard, lane_seq, job_id) AS ( - SELECT * FROM unnest($1::int[], $2::bigint[], $3::text[], $4::smallint[], $5::bigint[], $6::bigint[]) + WITH refs(ready_slot, ready_generation, queue, enqueue_shard, lane_seq, job_id, run_lease) AS ( + SELECT * FROM unnest($1::int[], $2::bigint[], $3::text[], $4::smallint[], $5::bigint[], $6::bigint[], $7::bigint[]) ) SELECT ready.ready_slot, @@ -10498,13 +10705,25 @@ impl QueueStorage { ready.kind, ready.queue, ready.args, + ready.priority, ready.lane_seq, ready.enqueue_shard, ready.run_at, ready.created_at, ready.unique_key, ready.unique_states, - COALESCE(ready.payload, '{{}}'::jsonb) AS payload + COALESCE(ready.payload, '{{}}'::jsonb) AS payload, + ready.run_lease, + NOT EXISTS ( + SELECT 1 FROM ({tombstone_sources}) AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + AND tomb.job_id = ready.job_id + ) AS live_body FROM refs JOIN {schema}.ready_entries AS ready ON ready.ready_slot = refs.ready_slot @@ -10513,6 +10732,7 @@ impl QueueStorage { AND ready.enqueue_shard = refs.enqueue_shard AND ready.lane_seq = refs.lane_seq AND ready.job_id = refs.job_id + AND ready.run_lease + 1 = refs.run_lease "# )) .bind(&ready_slots) @@ -10521,10 +10741,156 @@ impl QueueStorage { .bind(&enqueue_shards) .bind(&lane_seqs) .bind(&job_ids) + .bind(&run_leases) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; + let mut ready_map: BTreeMap> = BTreeMap::new(); + for row in ready_rows { + ready_map + .entry(( + row.ready_slot, + row.ready_generation, + row.queue.clone(), + row.enqueue_shard, + row.lane_seq, + row.job_id, + row.run_lease + 1, + )) + .or_default() + .push(row); + } + + let mut claimed_lanes: BTreeSet = BTreeSet::new(); + let ambiguous: Vec<&ReadySnapshotRow> = ready_map + .values() + .filter(|rows| rows.len() > 1) + .flatten() + .collect(); + if !ambiguous.is_empty() { + // A unique body needs no extra relation lookup. Multiple bodies + // need the emitted physical lane even if a replacement is live; + // scope that evidence to the referenced ready slots. + let attempt_sources = ambiguous + .iter() + .map(|row| row.ready_slot) + .collect::>() + .into_iter() + .map(|slot| ring_slot_number(slot, "queue")) + .collect::, _>>()? + .into_iter() + .map(|slot| { + format!( + "SELECT * FROM {}", + ready_claim_attempt_batch_child_name(schema, slot) + ) + }) + .collect::>() + .join(" UNION ALL "); + let lanes: Vec = sqlx::query_as(&format!( + r#" + WITH refs(ready_slot,ready_generation,queue,enqueue_shard,lane_seq,job_id,run_lease,priority) AS ( + SELECT * FROM unnest($1::int[], $2::bigint[], $3::text[], $4::smallint[], + $5::bigint[], $6::bigint[], $7::bigint[], $8::smallint[]) + ) + SELECT ready_slot,ready_generation,queue,enqueue_shard,lane_seq,job_id,run_lease,priority + FROM refs + WHERE EXISTS ( + SELECT 1 FROM ({attempt_sources}) AS attempts + WHERE attempts.ready_slot = refs.ready_slot + AND attempts.ready_generation = refs.ready_generation + AND attempts.queue = refs.queue + AND attempts.priority = refs.priority + AND attempts.enqueue_shard = refs.enqueue_shard + AND attempts.first_lane_seq <= refs.lane_seq + AND attempts.next_lane_seq > refs.lane_seq + AND attempts.lane_ranges @> int8range(refs.lane_seq,refs.lane_seq + 1,'[)') + ) + "# + )) + .bind(ambiguous.iter().map(|row| row.ready_slot).collect::>()) + .bind(ambiguous.iter().map(|row| row.ready_generation).collect::>()) + .bind(ambiguous.iter().map(|row| row.queue.clone()).collect::>()) + .bind(ambiguous.iter().map(|row| row.enqueue_shard).collect::>()) + .bind(ambiguous.iter().map(|row| row.lane_seq).collect::>()) + .bind(ambiguous.iter().map(|row| row.job_id).collect::>()) + .bind(ambiguous.iter().map(|row| row.run_lease + 1).collect::>()) + .bind(ambiguous.iter().map(|row| row.priority).collect::>()) + .fetch_all(tx.as_mut()).await.map_err(map_sqlx_error)?; + claimed_lanes.extend(lanes); + } + + Ok((ready_map, claimed_lanes)) + } + + async fn hydrate_deleted_leases_tx<'a>( + &self, + tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + deleted: Vec, + ) -> Result, AwaError> { + let (hydrated, sources) = self + .hydrate_deleted_leases_with_sources_tx(tx, deleted) + .await?; + if !sources.is_empty() { + self.retire_retained_attempt_copies_tx(tx, &sources, false) + .await?; + } + Ok(hydrated) + } + + async fn try_hydrate_deleted_leases_tx( + &self, + tx: &mut sqlx::Transaction<'_, Postgres>, + deleted: Vec, + ) -> Result, BusyRetained>, AwaError> { + let hydrated = async { + let (hydrated, sources) = self + .hydrate_deleted_leases_with_sources_tx(tx, deleted) + .await?; + if sources.is_empty() { + return Ok(Ok(hydrated)); + } + if let Some(busy) = self + .retire_retained_attempt_copies_tx(tx, &sources, true) + .await? + { + return Ok(Err(busy)); + } + Ok(Ok(hydrated)) + } + .await; + 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::Awa(err)) => Err(err), + } + } + + async fn hydrate_deleted_leases_with_sources_tx<'a>( + &self, + tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + deleted: Vec, + ) -> Result<(Vec, Vec), HydrateError> { + if deleted.is_empty() { + return Ok((Vec::new(), Vec::new())); + } + + let schema = self.schema(); + let job_ids: Vec = deleted.iter().map(|row| row.job_id).collect(); + 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?; + let attempt_rows: Vec = sqlx::query_as(&format!( r#" WITH refs(job_id, run_lease) AS ( @@ -10669,31 +11035,15 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; - let ready_map: BTreeMap<(i32, i64, String, i16, i64, i64), ReadySnapshotRow> = ready_rows - .into_iter() - .map(|row| { - ( - ( - row.ready_slot, - row.ready_generation, - row.queue.clone(), - row.enqueue_shard, - row.lane_seq, - row.job_id, - ), - row, - ) - }) - .collect(); - let attempt_map: BTreeMap<(i64, i64), AttemptStateRow> = attempt_rows .into_iter() .map(|row| ((row.job_id, row.run_lease), row)) .collect(); let mut hydrated = Vec::with_capacity(deleted.len()); + let mut retained_sources = Vec::new(); for deleted_row in deleted { - let ready = ready_map + let candidates = ready_map .get(&( deleted_row.ready_slot, deleted_row.ready_generation, @@ -10701,6 +11051,7 @@ impl QueueStorage { deleted_row.enqueue_shard, deleted_row.lane_seq, deleted_row.job_id, + deleted_row.run_lease, )) .ok_or_else(|| { AwaError::Validation(format!( @@ -10708,6 +11059,45 @@ impl QueueStorage { deleted_row.job_id, deleted_row.run_lease )) })?; + let (ready, has_emitted) = + select_claimed_ready_body(candidates, &claimed_lanes, &deleted_row)?; + if !ready.live_body { + retained_sources.push(ReadyRowIdentity { + ready_slot: ready.ready_slot, + ready_generation: ready.ready_generation, + queue: ready.queue.clone(), + priority: ready.priority, + enqueue_shard: ready.enqueue_shard, + lane_seq: ready.lane_seq, + job_id: ready.job_id, + run_lease: ready.run_lease, + }); + } + // Claim-time aging adds this provenance to the returned JobRow, + // not the immutable ready payload. Recovery must reconstruct the + // same default when it reads that retained body after the claim. + // Mutate the cloned JSON so unrelated payload fields survive. + let mut payload = ready.payload.clone(); + if deleted_row.priority < ready.priority { + let metadata = payload + .as_object_mut() + .ok_or_else(|| { + AwaError::Validation( + "queue storage payload must be a JSON object".to_string(), + ) + })? + .entry("metadata") + .or_insert_with(|| serde_json::json!({})) + .as_object_mut() + .ok_or_else(|| { + AwaError::Validation( + "queue storage payload metadata must be a JSON object".to_string(), + ) + })?; + metadata + .entry("_awa_original_priority") + .or_insert_with(|| serde_json::Value::from(i64::from(ready.priority))); + } let attempt = attempt_map.get(&(deleted_row.job_id, deleted_row.run_lease)); hydrated.push(LeaseTransitionRow { @@ -10719,6 +11109,7 @@ impl QueueStorage { args: ready.args.clone(), state: deleted_row.state, priority: deleted_row.priority, + ready_priority: ready.priority, attempt: deleted_row.attempt, run_lease: deleted_row.run_lease, max_attempts: deleted_row.max_attempts, @@ -10729,12 +11120,193 @@ impl QueueStorage { created_at: ready.created_at, unique_key: ready.unique_key.clone(), unique_states: ready.unique_states.clone(), - payload: ready.payload.clone(), + payload, progress: attempt.and_then(|row| row.progress.clone()), + self_contained_terminal: has_emitted + && !ready.live_body + && candidates.iter().any(|row| row.live_body), }); } - Ok(hydrated) + Ok((hydrated, retained_sources)) + } + + async fn retire_retained_attempt_copies_tx( + &self, + tx: &mut sqlx::Transaction<'_, Postgres>, + sources: &[ReadyRowIdentity], + skip_busy: bool, + ) -> Result, HydrateError> { + let schema = self.schema(); + let mut expected = BTreeMap::new(); + for source in sources { + if expected + .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, + }); + } + } + let job_ids: Vec = sources.iter().map(|row| row.job_id).collect(); + let run_leases: Vec = sources.iter().map(|row| row.run_lease).collect(); + if skip_busy { + let busy = sqlx::query_as::<_, (i64, i64)>(&format!( + r#" + WITH refs(job_id, run_lease) AS ( + SELECT * FROM unnest($1::bigint[], $2::bigint[]) + ), locked_heads AS MATERIALIZED ( + SELECT heads.queue, heads.priority, heads.enqueue_shard + FROM {schema}.queue_claim_heads AS heads + WHERE EXISTS ( + SELECT 1 + FROM refs + JOIN {schema}.ready_entries AS ready + ON ready.job_id = refs.job_id AND ready.run_lease = refs.run_lease + WHERE ready.queue = heads.queue AND ready.priority = heads.priority + AND ready.enqueue_shard = heads.enqueue_shard + ) + ORDER BY heads.queue COLLATE "C", heads.priority, heads.enqueue_shard + FOR UPDATE OF heads SKIP LOCKED + ) + SELECT refs.job_id, refs.run_lease + 1 + FROM refs + WHERE EXISTS ( + SELECT 1 + FROM {schema}.ready_entries AS ready + JOIN {schema}.queue_claim_heads AS heads + ON heads.queue = ready.queue AND heads.priority = ready.priority + AND heads.enqueue_shard = ready.enqueue_shard + WHERE ready.job_id = refs.job_id AND ready.run_lease = refs.run_lease + AND NOT EXISTS ( + SELECT 1 FROM locked_heads + WHERE locked_heads.queue = heads.queue + AND locked_heads.priority = heads.priority + AND locked_heads.enqueue_shard = heads.enqueue_shard + ) + ) + "# + )) + .bind(&job_ids) + .bind(&run_leases) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + if !busy.is_empty() { + return Ok(Some(BusyRetained(busy))); + } + } else { + self.lock_claim_heads_of_jobs_with_wait_tx(tx, &job_ids, true) + .await?; + } + let copies: Vec = sqlx::query_as(&format!( + r#" + WITH refs(job_id, run_lease) AS ( + SELECT * FROM unnest($1::bigint[], $2::bigint[]) + ) + SELECT ready.ready_slot, ready.ready_generation, ready.queue, ready.priority, + ready.enqueue_shard, ready.lane_seq, ready.job_id, ready.run_lease + FROM refs + JOIN {schema}.ready_entries AS ready + ON ready.job_id = refs.job_id AND ready.run_lease = refs.run_lease + "# + )) + .bind(&job_ids) + .bind(&run_leases) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let attempt_sources = copies + .iter() + .map(|row| ring_slot_number(row.ready_slot, "queue")) + .collect::, _>>()? + .into_iter() + .map(|slot| { + format!( + "SELECT * FROM {}", + ready_claim_attempt_batch_child_name(schema, slot) + ) + }) + .collect::>() + .join(" UNION ALL "); + if copies.is_empty() { + return Err(AwaError::Validation( + "queue storage ready row missing for deleted lease during retirement".to_string(), + ) + .into()); + } + let slots: Vec = copies.iter().map(|row| row.ready_slot).collect(); + let generations: Vec = copies.iter().map(|row| row.ready_generation).collect(); + let queues: Vec = copies.iter().map(|row| row.queue.clone()).collect(); + let priorities: Vec = copies.iter().map(|row| row.priority).collect(); + let shards: Vec = copies.iter().map(|row| row.enqueue_shard).collect(); + let sequences: Vec = copies.iter().map(|row| row.lane_seq).collect(); + let copy_job_ids: Vec = copies.iter().map(|row| row.job_id).collect(); + let copy_leases: Vec = copies.iter().map(|row| row.run_lease).collect(); + let emitted: Vec = sqlx::query_as(&format!( + r#" + WITH refs(ready_slot, ready_generation, queue, priority, enqueue_shard, + lane_seq, job_id, run_lease) AS ( + SELECT * FROM unnest($1::int[], $2::bigint[], $3::text[], $4::smallint[], + $5::smallint[], $6::bigint[], $7::bigint[], $8::bigint[]) + ) + SELECT * FROM refs + WHERE EXISTS ( + SELECT 1 FROM ({attempt_sources}) AS attempts + WHERE attempts.ready_slot = refs.ready_slot + AND attempts.ready_generation = refs.ready_generation + AND attempts.queue = refs.queue AND attempts.priority = refs.priority + AND attempts.enqueue_shard = refs.enqueue_shard + AND attempts.first_lane_seq <= refs.lane_seq + AND attempts.next_lane_seq > refs.lane_seq + AND attempts.lane_ranges @> int8range(refs.lane_seq, refs.lane_seq + 1, '[)') + ) + "# + )) + .bind(&slots) + .bind(&generations) + .bind(&queues) + .bind(&priorities) + .bind(&shards) + .bind(&sequences) + .bind(©_job_ids) + .bind(©_leases) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let mut seen = BTreeSet::new(); + 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, + }); + } + } + sqlx::query(&format!( + r#" + INSERT INTO {schema}.ready_tombstones ( + ready_slot, ready_generation, queue, priority, enqueue_shard, lane_seq, job_id + ) + SELECT * FROM unnest($1::int[], $2::bigint[], $3::text[], $4::smallint[], + $5::smallint[], $6::bigint[], $7::bigint[]) + ON CONFLICT DO NOTHING + "# + )) + .bind(&slots) + .bind(&generations) + .bind(&queues) + .bind(&priorities) + .bind(&shards) + .bind(&sequences) + .bind(©_job_ids) + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + Ok(None) } async fn close_open_receipt_claim_tx<'a>( @@ -11033,6 +11605,8 @@ impl QueueStorage { &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, cutoff: DateTime, + excluded: &BTreeSet<(i64, i64)>, + selected: Option<&BTreeSet<(i64, i64)>>, ) -> Result, AwaError> { let mut rescued = Vec::new(); let mut remaining = RECEIPT_RESCUE_BATCH_LIMIT; @@ -11040,7 +11614,9 @@ impl QueueStorage { if let Some(slot) = preferred_slot { let mut slot_rescued = self - .rescue_stale_receipt_claims_for_slot_tx(tx, slot, cutoff, remaining) + .rescue_stale_receipt_claims_for_slot_tx( + tx, slot, cutoff, remaining, excluded, selected, + ) .await?; remaining = remaining.saturating_sub(slot_rescued.len() as i64); rescued.append(&mut slot_rescued); @@ -11056,7 +11632,9 @@ impl QueueStorage { } let mut slot_rescued = self - .rescue_stale_receipt_claims_for_slot_tx(tx, slot, cutoff, remaining) + .rescue_stale_receipt_claims_for_slot_tx( + tx, slot, cutoff, remaining, excluded, selected, + ) .await?; remaining = remaining.saturating_sub(slot_rescued.len() as i64); rescued.append(&mut slot_rescued); @@ -11075,7 +11653,14 @@ impl QueueStorage { slot: i32, cutoff: DateTime, rescue_limit: i64, + excluded: &BTreeSet<(i64, i64)>, + selected: Option<&BTreeSet<(i64, i64)>>, ) -> Result, AwaError> { + let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = + excluded.iter().copied().unzip(); + let (selected_job_ids, selected_run_leases): (Vec, Vec) = selected + .map(|pairs| pairs.iter().copied().unzip()) + .unwrap_or_default(); let schema = self.schema(); let claim_child = claim_child_name(schema, slot as usize); let claim_batch_child = claim_batch_child_name(schema, slot as usize); @@ -11260,15 +11845,23 @@ impl QueueStorage { WHERE NOT is_closed AND NOT is_lease_managed AND is_stale + AND NOT EXISTS ( + SELECT 1 FROM unnest($5::bigint[], $6::bigint[]) AS excluded(job_id, run_lease) + WHERE excluded.job_id = candidates.job_id AND excluded.run_lease = candidates.run_lease + ) + AND ($7::boolean OR EXISTS ( + SELECT 1 FROM unnest($8::bigint[], $9::bigint[]) AS selected(job_id, run_lease) + WHERE selected.job_id = candidates.job_id AND selected.run_lease = candidates.run_lease + )) ORDER BY rn LIMIT $4 )"# ); // Closers take the attempt's advisory lock before their statement - // snapshot. Rescue takes the same locks in a statement of its own, so - // the statement that decides and writes the closure starts after any - // closer that held a lock first has committed, and sees its evidence. - // Only attempts locked here may be closed below. + // snapshot. Rescue takes the same locks in a statement of its own, + // so the statement that decides and writes the closure starts after + // any closer that held a lock first has committed and sees its + // evidence. Only attempts locked here may be closed below. let (locked_job_ids, locked_run_leases): (Vec, Vec) = sqlx::query_as::<_, (i64, i64)>(&format!( r#" @@ -11287,6 +11880,11 @@ impl QueueStorage { .bind(cutoff) .bind(RECEIPT_RESCUE_CURSOR_SCAN_LIMIT) .bind(rescue_limit) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)? @@ -11332,7 +11930,7 @@ impl QueueStorage { ) AND (claims.job_id, claims.run_lease) IN ( SELECT locked.job_id, locked.run_lease - FROM unnest($5::bigint[], $6::bigint[]) AS locked(job_id, run_lease) + FROM unnest($10::bigint[], $11::bigint[]) AS locked(job_id, run_lease) ) -- A claim that already materialized into `leases` is -- on the lease-side heartbeat-rescue path (see @@ -11381,7 +11979,7 @@ impl QueueStorage { ) AND (stale_candidates.job_id, stale_candidates.run_lease) IN ( SELECT locked.job_id, locked.run_lease - FROM unnest($5::bigint[], $6::bigint[]) AS locked(job_id, run_lease) + FROM unnest($10::bigint[], $11::bigint[]) AS locked(job_id, run_lease) ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease @@ -11518,6 +12116,11 @@ impl QueueStorage { .bind(cutoff) .bind(RECEIPT_RESCUE_CURSOR_SCAN_LIMIT) .bind(rescue_limit) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) .bind(&locked_job_ids) .bind(&locked_run_leases) .fetch_all(tx.as_mut()) @@ -11546,6 +12149,8 @@ impl QueueStorage { async fn rescue_expired_receipt_deadlines_tx<'a>( &self, tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, + excluded: &BTreeSet<(i64, i64)>, + selected: Option<&BTreeSet<(i64, i64)>>, ) -> Result, AwaError> { let mut rescued = Vec::new(); let mut remaining = RECEIPT_RESCUE_BATCH_LIMIT; @@ -11553,7 +12158,9 @@ impl QueueStorage { if let Some(slot) = preferred_slot { let mut slot_rescued = self - .rescue_expired_receipt_deadlines_for_slot_tx(tx, slot, remaining) + .rescue_expired_receipt_deadlines_for_slot_tx( + tx, slot, remaining, excluded, selected, + ) .await?; remaining = remaining.saturating_sub(slot_rescued.len() as i64); rescued.append(&mut slot_rescued); @@ -11569,7 +12176,9 @@ impl QueueStorage { } let mut slot_rescued = self - .rescue_expired_receipt_deadlines_for_slot_tx(tx, slot, remaining) + .rescue_expired_receipt_deadlines_for_slot_tx( + tx, slot, remaining, excluded, selected, + ) .await?; remaining = remaining.saturating_sub(slot_rescued.len() as i64); rescued.append(&mut slot_rescued); @@ -11611,11 +12220,18 @@ impl QueueStorage { tx: &mut sqlx::Transaction<'a, sqlx::Postgres>, slot: i32, rescue_limit: i64, + excluded: &BTreeSet<(i64, i64)>, + selected: Option<&BTreeSet<(i64, i64)>>, ) -> Result, AwaError> { if rescue_limit <= 0 { return Ok(Vec::new()); } + let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = + excluded.iter().copied().unzip(); + let (selected_job_ids, selected_run_leases): (Vec, Vec) = selected + .map(|pairs| pairs.iter().copied().unzip()) + .unwrap_or_default(); let schema = self.schema(); let claim_child = claim_child_name(schema, slot as usize); let closure_child = closure_child_name(schema, slot as usize); @@ -11723,6 +12339,14 @@ impl QueueStorage { WHERE NOT is_closed AND NOT is_lease_managed AND is_expired + AND NOT EXISTS ( + SELECT 1 FROM unnest($4::bigint[], $5::bigint[]) AS excluded(job_id, run_lease) + WHERE excluded.job_id = candidates.job_id AND excluded.run_lease = candidates.run_lease + ) + AND ($6::boolean OR EXISTS ( + SELECT 1 FROM unnest($7::bigint[], $8::bigint[]) AS selected(job_id, run_lease) + WHERE selected.job_id = candidates.job_id AND selected.run_lease = candidates.run_lease + )) ORDER BY rn LIMIT $3 )"# @@ -11745,6 +12369,11 @@ impl QueueStorage { .bind(slot) .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) .bind(rescue_limit) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)? @@ -11764,7 +12393,7 @@ impl QueueStorage { AND claims.deadline_at < clock_timestamp() AND (claims.job_id, claims.run_lease) IN ( SELECT locked.job_id, locked.run_lease - FROM unnest($4::bigint[], $5::bigint[]) AS locked(job_id, run_lease) + FROM unnest($9::bigint[], $10::bigint[]) AS locked(job_id, run_lease) ) AND NOT EXISTS ( SELECT 1 FROM {schema}.leases AS lease @@ -11857,6 +12486,11 @@ impl QueueStorage { .bind(slot) .bind(RECEIPT_DEADLINE_RESCUE_CURSOR_SCAN_LIMIT) .bind(rescue_limit) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) .bind(&locked_job_ids) .bind(&locked_run_leases) .fetch_all(tx.as_mut()) @@ -11865,6 +12499,87 @@ impl QueueStorage { Ok(rescued) } + async fn active_job_read_row_tx( + &self, + tx: &mut sqlx::Transaction<'_, Postgres>, + row: sqlx::postgres::PgRow, + ) -> Result, AwaError> { + let body_job_id: Option = row.try_get("body_job_id").map_err(map_sqlx_error)?; + if body_job_id.is_some() { + return LeaseJobRow::from_row(&row) + .map_err(map_sqlx_error)? + .into_job_row() + .map(Some); + } + let reference = DeletedLeaseRow::from_row(&row).map_err(map_sqlx_error)?; + let (ready_map, claimed_lanes) = self + .load_ready_snapshot_candidates_tx(tx, std::slice::from_ref(&reference)) + .await?; + let key = ( + reference.ready_slot, + reference.ready_generation, + reference.queue.clone(), + reference.enqueue_shard, + reference.lane_seq, + reference.job_id, + reference.run_lease, + ); + let Some(candidates) = ready_map.get(&key) else { + return Ok(None); + }; + let (ready, _) = select_claimed_ready_body(candidates, &claimed_lanes, &reference)?; + let mut payload = ready.payload.clone(); + if reference.priority < ready.priority { + let metadata = payload + .as_object_mut() + .ok_or_else(|| { + AwaError::Validation("queue storage payload must be a JSON object".into()) + })? + .entry("metadata") + .or_insert_with(|| serde_json::json!({})) + .as_object_mut() + .ok_or_else(|| { + AwaError::Validation( + "queue storage payload metadata must be a JSON object".into(), + ) + })?; + metadata + .entry("_awa_original_priority") + .or_insert_with(|| serde_json::Value::from(i64::from(ready.priority))); + } + LeaseJobRow { + job_id: reference.job_id, + kind: ready.kind.clone(), + queue: ready.queue.clone(), + args: ready.args.clone(), + state: reference.state, + priority: reference.priority, + attempt: reference.attempt, + run_lease: reference.run_lease, + max_attempts: reference.max_attempts, + run_at: ready.run_at, + heartbeat_at: row.try_get("heartbeat_at").map_err(map_sqlx_error)?, + deadline_at: row.try_get("deadline_at").map_err(map_sqlx_error)?, + attempted_at: reference.attempted_at, + finalized_at: None, + created_at: ready.created_at, + unique_key: ready.unique_key.clone(), + callback_id: row.try_get("callback_id").map_err(map_sqlx_error)?, + callback_timeout_at: row.try_get("callback_timeout_at").map_err(map_sqlx_error)?, + callback_filter: row.try_get("callback_filter").map_err(map_sqlx_error)?, + callback_on_complete: row + .try_get("callback_on_complete") + .map_err(map_sqlx_error)?, + callback_on_fail: row.try_get("callback_on_fail").map_err(map_sqlx_error)?, + callback_transform: row.try_get("callback_transform").map_err(map_sqlx_error)?, + payload, + progress: row.try_get("progress").map_err(map_sqlx_error)?, + callback_result: row.try_get("callback_result").map_err(map_sqlx_error)?, + } + .into_job_row() + .map(Some) + } + pub async fn load_job(&self, pool: &PgPool, job_id: i64) -> Result, AwaError> { let schema = self.schema(); let closure_rel = format!("{schema}.lease_claim_closures"); @@ -11875,11 +12590,65 @@ impl QueueStorage { // A job moves between these relations as it is promoted, claimed, // completed or dead-lettered. Read them all from one snapshot so a // move between two reads cannot hide the job or surface a stale row. - let mut tx = pool - .begin_with("BEGIN ISOLATION LEVEL REPEATABLE READ, READ ONLY") + // Keep the located ready-body partitions locked until commit; later + // read savepoints release unrelated relation locks after each leg. + let mut begin_guard = crate::read_transaction::BeginGuard::acquire(pool) .await .map_err(map_sqlx_error)?; + let mut tx = sqlx::Connection::begin_with( + &mut *begin_guard.connection, + "BEGIN ISOLATION LEVEL REPEATABLE READ, READ ONLY", + ) + .await + .map_err(map_sqlx_error)?; + begin_guard.begin_pending = false; + sqlx::query("SAVEPOINT awa_read_leg") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let ready_slots: Vec = sqlx::query_scalar(&format!( + "SELECT DISTINCT ready_slot FROM {schema}.ready_entries WHERE job_id = $1" + )) + .bind(job_id) + .fetch_all(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + sqlx::Executor::execute( + tx.as_mut(), + "ROLLBACK TO SAVEPOINT awa_read_leg; RELEASE SAVEPOINT awa_read_leg", + ) + .await + .map_err(map_sqlx_error)?; + let ready_slots = ready_slots + .into_iter() + .map(|slot| ring_slot_number(slot, "queue")) + .collect::, _>>()?; + let child_sources = |tombstones: bool| { + if ready_slots.is_empty() { + let table = if tombstones { + "ready_tombstones" + } else { + "ready_entries" + }; + format!("SELECT * FROM ONLY {schema}.{table} WHERE FALSE") + } else { + ready_slots + .iter() + .map(|&slot| { + let child = if tombstones { + ready_tombstone_child_name(schema, slot) + } else { + ready_child_name(schema, slot) + }; + format!("SELECT * FROM {child}") + }) + .collect::>() + .join(" UNION ALL ") + } + }; + let ready_sources = child_sources(false); + let tombstone_sources = child_sources(true); let ready_rows: Vec = sqlx::query_as(&format!( r#" SELECT @@ -11897,10 +12666,10 @@ impl QueueStorage { unique_key, unique_states, COALESCE(payload, '{{}}'::jsonb) AS payload - FROM {schema}.ready_entries AS ready + FROM ({ready_sources}) AS ready WHERE ready.job_id = $1 AND NOT EXISTS ( - SELECT 1 FROM {schema}.ready_tombstones AS tomb + SELECT 1 FROM ({tombstone_sources}) AS tomb WHERE tomb.ready_slot = ready.ready_slot AND tomb.ready_generation = ready.ready_generation AND tomb.queue = ready.queue @@ -11919,6 +12688,10 @@ impl QueueStorage { candidates.push(row.into_job_row()?); } + sqlx::query("SAVEPOINT awa_read_leg") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; let deferred_rows: Vec = sqlx::query_as(&format!( r#" SELECT @@ -11946,18 +12719,30 @@ impl QueueStorage { .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; + sqlx::Executor::execute( + tx.as_mut(), + "ROLLBACK TO SAVEPOINT awa_read_leg; RELEASE SAVEPOINT awa_read_leg", + ) + .await + .map_err(map_sqlx_error)?; for row in deferred_rows { candidates.push(row.into_job_row()?); } - let lease_rows: Vec = sqlx::query_as(&format!( + sqlx::query("SAVEPOINT awa_read_leg") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let lease_rows = sqlx::query(&format!( r#" SELECT + ready.job_id AS body_job_id, lease.ready_slot, lease.ready_generation, lease.job_id, + lease.enqueue_shard, ready.kind, - ready.queue, + lease.queue, ready.args, lease.state, lease.priority, @@ -11983,7 +12768,7 @@ impl QueueStorage { attempt.progress, attempt.callback_result FROM {schema}.leases AS lease - JOIN {schema}.ready_entries AS ready + LEFT JOIN ({ready_sources}) AS ready ON ready.ready_slot = lease.ready_slot AND ready.ready_generation = lease.ready_generation AND ready.queue = lease.queue @@ -11993,6 +12778,15 @@ impl QueueStorage { AND ready.run_lease + 1 = lease.run_lease AND ready.enqueue_shard = lease.enqueue_shard AND ready.lane_seq = lease.lane_seq + AND NOT EXISTS ( + SELECT 1 FROM ({tombstone_sources}) AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = lease.job_id AND attempt.run_lease = lease.run_lease @@ -12005,19 +12799,34 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; for row in lease_rows { - candidates.push(row.into_job_row()?); + if let Some(job) = self.active_job_read_row_tx(&mut tx, row).await? { + candidates.push(job); + } } + sqlx::Executor::execute( + tx.as_mut(), + "ROLLBACK TO SAVEPOINT awa_read_leg; RELEASE SAVEPOINT awa_read_leg", + ) + .await + .map_err(map_sqlx_error)?; + // Report receipt-backed attempts as running by anti-joining // lease_claims against every durable closure evidence shape. - let lease_claim_rows: Vec = sqlx::query_as(&format!( + sqlx::query("SAVEPOINT awa_read_leg") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let lease_claim_rows = sqlx::query(&format!( r#" SELECT + ready.job_id AS body_job_id, claims.ready_slot, claims.ready_generation, claims.job_id, + claims.enqueue_shard, ready.kind, - ready.queue, + claims.queue, ready.args, 'running'::awa.job_state AS state, claims.priority, @@ -12043,7 +12852,7 @@ impl QueueStorage { attempt.progress, attempt.callback_result FROM {schema}.lease_claims AS claims - JOIN {schema}.ready_entries AS ready + LEFT JOIN ({ready_sources}) AS ready ON ready.ready_slot = claims.ready_slot AND ready.ready_generation = claims.ready_generation AND ready.queue = claims.queue @@ -12053,6 +12862,15 @@ impl QueueStorage { AND ready.run_lease + 1 = claims.run_lease AND ready.enqueue_shard = claims.enqueue_shard AND ready.lane_seq = claims.lane_seq + AND NOT EXISTS ( + SELECT 1 FROM ({tombstone_sources}) AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = claims.job_id AND attempt.run_lease = claims.run_lease @@ -12074,21 +12892,36 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; for row in lease_claim_rows { - candidates.push(row.into_job_row()?); + if let Some(job) = self.active_job_read_row_tx(&mut tx, row).await? { + candidates.push(job); + } } + sqlx::Executor::execute( + tx.as_mut(), + "ROLLBACK TO SAVEPOINT awa_read_leg; RELEASE SAVEPOINT awa_read_leg", + ) + .await + .map_err(map_sqlx_error)?; + // Zero-deadline receipt claims are stored as compact batches. Expand // them only for this admin read, and report still-open items as // running until durable closure, terminal, or materialized-lease // evidence supersedes the claim. - let lease_claim_batch_rows: Vec = sqlx::query_as(&format!( + sqlx::query("SAVEPOINT awa_read_leg") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; + let lease_claim_batch_rows = sqlx::query(&format!( r#" SELECT + ready.job_id AS body_job_id, claim_batches.ready_slot, claim_batches.ready_generation, items.job_id, + claim_batches.enqueue_shard, ready.kind, - ready.queue, + claim_batches.queue, ready.args, 'running'::awa.job_state AS state, claim_batches.priority, @@ -12122,13 +12955,23 @@ impl QueueStorage { claim_batches.attempts, claim_batches.max_attempts ) AS items(job_id, run_lease, receipt_id, lane_seq, attempt, max_attempts) - JOIN {schema}.ready_entries AS ready + LEFT JOIN ({ready_sources}) AS ready ON ready.ready_slot = claim_batches.ready_slot AND ready.ready_generation = claim_batches.ready_generation AND ready.queue = claim_batches.queue AND ready.enqueue_shard = claim_batches.enqueue_shard AND ready.lane_seq = items.lane_seq AND ready.job_id = items.job_id + AND ready.run_lease + 1 = items.run_lease + AND NOT EXISTS ( + SELECT 1 FROM ({tombstone_sources}) AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = items.job_id AND attempt.run_lease = items.run_lease @@ -12175,9 +13018,22 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; for row in lease_claim_batch_rows { - candidates.push(row.into_job_row()?); + if let Some(job) = self.active_job_read_row_tx(&mut tx, row).await? { + candidates.push(job); + } } + sqlx::Executor::execute( + tx.as_mut(), + "ROLLBACK TO SAVEPOINT awa_read_leg; RELEASE SAVEPOINT awa_read_leg", + ) + .await + .map_err(map_sqlx_error)?; + + sqlx::query("SAVEPOINT awa_read_leg") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; let done_rows: Vec = sqlx::query_as(&format!( r#" SELECT @@ -12210,10 +13066,20 @@ impl QueueStorage { .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; + sqlx::Executor::execute( + tx.as_mut(), + "ROLLBACK TO SAVEPOINT awa_read_leg; RELEASE SAVEPOINT awa_read_leg", + ) + .await + .map_err(map_sqlx_error)?; for row in done_rows { candidates.push(row.into_job_row()?); } + sqlx::query("SAVEPOINT awa_read_leg") + .execute(tx.as_mut()) + .await + .map_err(map_sqlx_error)?; let dlq_rows: Vec = sqlx::query_as(&format!( r#" SELECT @@ -12245,11 +13111,18 @@ impl QueueStorage { .fetch_all(tx.as_mut()) .await .map_err(map_sqlx_error)?; + sqlx::Executor::execute( + tx.as_mut(), + "ROLLBACK TO SAVEPOINT awa_read_leg; RELEASE SAVEPOINT awa_read_leg", + ) + .await + .map_err(map_sqlx_error)?; for row in dlq_rows { candidates.push(row.into_job_row()?); } tx.commit().await.map_err(map_sqlx_error)?; + drop(begin_guard); Ok(candidates.into_iter().max_by_key(|job| { ( job.run_lease, @@ -12538,14 +13411,16 @@ impl QueueStorage { run_lease: i64, ) -> Result, AwaError> { let schema = self.schema(); - let row: Option = sqlx::query_as(&format!( + let row = sqlx::query(&format!( r#" SELECT + ready.job_id AS body_job_id, lease.ready_slot, lease.ready_generation, lease.job_id, + lease.enqueue_shard, ready.kind, - ready.queue, + lease.queue, ready.args, lease.state, lease.priority, @@ -12571,7 +13446,7 @@ impl QueueStorage { attempt.progress, attempt.callback_result FROM {schema}.leases AS lease - JOIN {schema}.ready_entries AS ready + LEFT JOIN {schema}.ready_entries AS ready ON ready.ready_slot = lease.ready_slot AND ready.ready_generation = lease.ready_generation AND ready.queue = lease.queue @@ -12581,6 +13456,15 @@ impl QueueStorage { AND ready.run_lease + 1 = lease.run_lease AND ready.enqueue_shard = lease.enqueue_shard AND ready.lane_seq = lease.lane_seq + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = lease.job_id AND attempt.run_lease = lease.run_lease @@ -12593,7 +13477,10 @@ impl QueueStorage { .fetch_optional(tx.as_mut()) .await .map_err(map_sqlx_error)?; - row.map(LeaseJobRow::into_job_row).transpose() + match row { + Some(row) => self.active_job_read_row_tx(tx, row).await, + None => Ok(None), + } } pub async fn enter_callback_wait( @@ -12737,14 +13624,16 @@ impl QueueStorage { } else { "" }; - let row: Option = sqlx::query_as(&format!( + let row = sqlx::query(&format!( r#" SELECT + ready.job_id AS body_job_id, lease.ready_slot, lease.ready_generation, lease.job_id, + lease.enqueue_shard, ready.kind, - ready.queue, + lease.queue, ready.args, lease.state, lease.priority, @@ -12770,7 +13659,7 @@ impl QueueStorage { attempt.progress, attempt.callback_result FROM {} AS lease - JOIN {schema}.ready_entries AS ready + LEFT JOIN {schema}.ready_entries AS ready ON ready.ready_slot = lease.ready_slot AND ready.ready_generation = lease.ready_generation AND ready.queue = lease.queue @@ -12780,6 +13669,15 @@ impl QueueStorage { AND ready.run_lease + 1 = lease.run_lease AND ready.enqueue_shard = lease.enqueue_shard AND ready.lane_seq = lease.lane_seq + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue + AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq + ) LEFT JOIN {schema}.attempt_state AS attempt ON attempt.job_id = lease.job_id AND attempt.run_lease = lease.run_lease @@ -12799,7 +13697,10 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; - row.map(LeaseJobRow::into_job_row).transpose() + match row { + Some(row) => self.active_job_read_row_tx(tx, row).await, + None => Ok(None), + } } #[tracing::instrument(skip(self, pool, payload), name = "queue_storage.complete_external")] @@ -13921,7 +14822,7 @@ impl QueueStorage { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let done_projection = done_row_projection("done", "ready"); let ready_join = done_ready_join(schema, "done", "ready"); - let moved: Option = sqlx::query_as(&format!( + let moved: Option = sqlx::query_as(&format!( r#" WITH deleted AS ( DELETE FROM {schema}.done_entries @@ -13936,7 +14837,7 @@ impl QueueStorage { ) RETURNING * ) - SELECT {done_projection} + SELECT {done_projection}, COALESCE(ready.priority, done.priority) AS body_priority FROM deleted AS done {ready_join} "# @@ -13946,10 +14847,11 @@ impl QueueStorage { .await .map_err(map_sqlx_error)?; - let Some(moved) = moved else { + let Some(hydrated) = moved else { tx.commit().await.map_err(map_sqlx_error)?; return Ok(None); }; + let mut moved = hydrated.terminal; self.ensure_terminal_removed_receipt_closures_tx(&mut tx, std::slice::from_ref(&moved)) .await?; @@ -13960,6 +14862,7 @@ impl QueueStorage { &Self::done_rows_to_counter_keys(std::slice::from_ref(&moved)), ) .await?; + moved.priority = hydrated.body_priority; let dlq_row = moved .clone() .into_dlq_row(dlq_reason.to_string(), Utc::now()); @@ -13985,7 +14888,7 @@ impl QueueStorage { let mut tx = pool.begin().await.map_err(map_sqlx_error)?; let done_projection = done_row_projection("done", "ready"); let ready_join = done_ready_join(schema, "done", "ready"); - let moved: Vec = sqlx::query_as(&format!( + let moved: Vec = sqlx::query_as(&format!( r#" WITH deleted AS ( DELETE FROM {schema}.done_entries @@ -13994,7 +14897,7 @@ impl QueueStorage { AND ($2::text IS NULL OR queue = $2) RETURNING * ) - SELECT {done_projection} + SELECT {done_projection}, COALESCE(ready.priority, done.priority) AS body_priority FROM deleted AS done {ready_join} "# @@ -14010,6 +14913,10 @@ impl QueueStorage { return Ok(0); } + let (moved, body_priorities): (Vec<_>, Vec<_>) = moved + .into_iter() + .map(|row| (row.terminal, row.body_priority)) + .unzip(); self.ensure_terminal_removed_receipt_closures_tx(&mut tx, &moved) .await?; // Bulk DLQ move deletes from done_entries. Append negative deltas by @@ -14019,7 +14926,11 @@ impl QueueStorage { let dlq_at = Utc::now(); let rows: Vec = moved .into_iter() - .map(|row| row.into_dlq_row(dlq_reason.to_string(), dlq_at)) + .zip(body_priorities) + .map(|(mut row, priority)| { + row.priority = priority; + row.into_dlq_row(dlq_reason.to_string(), dlq_at) + }) .collect(); self.insert_dlq_rows_tx(&mut tx, &rows, Some(JobState::Failed)) .await?; @@ -14411,7 +15322,17 @@ impl QueueStorage { // `(state, callback_timeout_at)` indexes, followed by a heap // recheck of the COALESCE. Bounded by running-lease count and // called at 30s cadence — cheap. - let deleted: Vec = sqlx::query_as(&format!( + let mut excluded = BTreeSet::new(); + let mut selected: Option> = None; + loop { + let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = + excluded.iter().copied().unzip(); + let (selected_job_ids, selected_run_leases): (Vec, Vec) = selected + .as_ref() + .map(|pairs| pairs.iter().copied().unzip()) + .unwrap_or_default(); + let mut attempt = tx.begin().await.map_err(map_sqlx_error)?; + let deleted: Vec = sqlx::query_as(&format!( r#" DELETE FROM {schema}.leases AS target WHERE (target.job_id, target.run_lease) IN ( @@ -14422,6 +15343,14 @@ impl QueueStorage { AND attempt.run_lease = lease.run_lease WHERE lease.state = 'running' AND COALESCE(attempt.heartbeat_at, lease.heartbeat_at) < $1 + AND NOT EXISTS ( + SELECT 1 FROM unnest($2::bigint[], $3::bigint[]) AS excluded(job_id, run_lease) + WHERE excluded.job_id = lease.job_id AND excluded.run_lease = lease.run_lease + ) + AND ($4::boolean OR EXISTS ( + SELECT 1 FROM unnest($5::bigint[], $6::bigint[]) AS selected(job_id, run_lease) + WHERE selected.job_id = lease.job_id AND selected.run_lease = lease.run_lease + )) ORDER BY COALESCE(attempt.heartbeat_at, lease.heartbeat_at) ASC LIMIT 500 FOR UPDATE OF lease SKIP LOCKED @@ -14446,53 +15375,98 @@ impl QueueStorage { "# )) .bind(cutoff) - .fetch_all(tx.as_mut()) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) + .fetch_all(attempt.as_mut()) .await .map_err(map_sqlx_error)?; - let rescued_receipts = if self.lease_claim_receipts() { - self.rescue_stale_receipt_claims_tx(&mut tx, cutoff).await? - } else { - Vec::new() - }; + let receipt_excluded: BTreeSet<_> = excluded + .iter() + .copied() + .chain(deleted.iter().map(|row| (row.job_id, row.run_lease))) + .collect(); + let rescued_receipts = if self.lease_claim_receipts() { + self.rescue_stale_receipt_claims_tx( + &mut attempt, + cutoff, + &receipt_excluded, + selected.as_ref(), + ) + .await? + } else { + Vec::new() + }; - if deleted.is_empty() && rescued_receipts.is_empty() { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(Vec::new()); - } + selected.get_or_insert_with(|| { + deleted + .iter() + .chain(rescued_receipts.iter()) + .map(|row| (row.job_id, row.run_lease)) + .collect() + }); - let moved_leases = self.hydrate_deleted_leases_tx(&mut tx, deleted).await?; - let moved_receipts = self - .hydrate_deleted_leases_tx(&mut tx, rescued_receipts) - .await?; + if deleted.is_empty() && rescued_receipts.is_empty() { + attempt.commit().await.map_err(map_sqlx_error)?; + tx.commit().await.map_err(map_sqlx_error)?; + return Ok(Vec::new()); + } - let mut rescued = Vec::with_capacity(moved_leases.len() + moved_receipts.len()); - for row in moved_leases { - let dlq_reason = dlq_enabled(&row.queue).then_some("heartbeat_stale"); - rescued.push( - self.settle_rescued_attempt_tx( - &mut tx, - row, - "heartbeat stale: worker presumed dead", - dlq_reason, - ) - .await?, - ); - } - for row in moved_receipts { - let dlq_reason = dlq_enabled(&row.queue).then_some("heartbeat_stale"); - rescued.push( - self.settle_rescued_attempt_tx( - &mut tx, - row, - "receipt claim stale: worker presumed dead", - dlq_reason, - ) - .await?, - ); + let moved_leases = match self + .try_hydrate_deleted_leases_tx(&mut attempt, deleted) + .await? + { + Ok(moved) => moved, + Err(busy) => { + attempt.rollback().await.map_err(map_sqlx_error)?; + busy.exclude(&mut excluded)?; + continue; + } + }; + let moved_receipts = match self + .try_hydrate_deleted_leases_tx(&mut attempt, rescued_receipts) + .await? + { + Ok(moved) => moved, + Err(busy) => { + attempt.rollback().await.map_err(map_sqlx_error)?; + busy.exclude(&mut excluded)?; + continue; + } + }; + + let mut rescued = Vec::with_capacity(moved_leases.len() + moved_receipts.len()); + for row in moved_leases { + let dlq_reason = dlq_enabled(&row.queue).then_some("heartbeat_stale"); + rescued.push( + self.settle_rescued_attempt_tx( + &mut attempt, + row, + "heartbeat stale: worker presumed dead", + dlq_reason, + ) + .await?, + ); + } + for row in moved_receipts { + let dlq_reason = dlq_enabled(&row.queue).then_some("heartbeat_stale"); + rescued.push( + self.settle_rescued_attempt_tx( + &mut attempt, + row, + "receipt claim stale: worker presumed dead", + dlq_reason, + ) + .await?, + ); + } + attempt.commit().await.map_err(map_sqlx_error)?; + tx.commit().await.map_err(map_sqlx_error)?; + return Ok(rescued); } - tx.commit().await.map_err(map_sqlx_error)?; - Ok(rescued) } /// Rescue attempts past their hard deadline. Exhausted attempts fail @@ -14517,15 +15491,33 @@ impl QueueStorage { ) -> Result, AwaError> { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; - let deleted: Vec = sqlx::query_as(&format!( + let mut excluded = BTreeSet::new(); + let mut selected: Option> = None; + loop { + let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = + excluded.iter().copied().unzip(); + let (selected_job_ids, selected_run_leases): (Vec, Vec) = selected + .as_ref() + .map(|pairs| pairs.iter().copied().unzip()) + .unwrap_or_default(); + let mut attempt = tx.begin().await.map_err(map_sqlx_error)?; + let deleted: Vec = sqlx::query_as(&format!( r#" DELETE FROM {schema}.leases - WHERE job_id IN ( - SELECT job_id + WHERE (job_id, run_lease) IN ( + SELECT job_id, run_lease FROM {schema}.leases WHERE state = 'running' AND deadline_at IS NOT NULL AND deadline_at < clock_timestamp() + AND NOT EXISTS ( + SELECT 1 FROM unnest($1::bigint[], $2::bigint[]) AS excluded(job_id, run_lease) + WHERE excluded.job_id = leases.job_id AND excluded.run_lease = leases.run_lease + ) + AND ($3::boolean OR EXISTS ( + SELECT 1 FROM unnest($4::bigint[], $5::bigint[]) AS selected(job_id, run_lease) + WHERE selected.job_id = leases.job_id AND selected.run_lease = leases.run_lease + )) ORDER BY deadline_at ASC LIMIT 500 FOR UPDATE SKIP LOCKED @@ -14549,42 +15541,92 @@ impl QueueStorage { callback_timeout_at "# )) - .fetch_all(tx.as_mut()) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) + .fetch_all(attempt.as_mut()) .await .map_err(map_sqlx_error)?; - // Receipts-mode short-path claims hold their deadline on - // `lease_claims.deadline_at` rather than on a `leases` row, so - // the receipt-plane needs its own scan; merge both populations - // into one `moved` set so the maintenance caller observes a - // single rescue batch per tick. - let receipt_deleted = if self.lease_claim_receipts() { - self.rescue_expired_receipt_deadlines_tx(&mut tx).await? - } else { - Vec::new() - }; + // Receipts-mode short-path claims hold their deadline on + // `lease_claims.deadline_at` rather than on a `leases` row, so + // the receipt-plane needs its own scan; merge both populations + // into one `moved` set so the maintenance caller observes a + // single rescue batch per tick. + let receipt_excluded: BTreeSet<_> = excluded + .iter() + .copied() + .chain(deleted.iter().map(|row| (row.job_id, row.run_lease))) + .collect(); + let receipt_deleted = if self.lease_claim_receipts() { + self.rescue_expired_receipt_deadlines_tx( + &mut attempt, + &receipt_excluded, + selected.as_ref(), + ) + .await? + } else { + Vec::new() + }; - if deleted.is_empty() && receipt_deleted.is_empty() { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(Vec::new()); - } + selected.get_or_insert_with(|| { + deleted + .iter() + .chain(receipt_deleted.iter()) + .map(|row| (row.job_id, row.run_lease)) + .collect() + }); - let mut moved = self.hydrate_deleted_leases_tx(&mut tx, deleted).await?; - moved.extend( - self.hydrate_deleted_leases_tx(&mut tx, receipt_deleted) - .await?, - ); + if deleted.is_empty() && receipt_deleted.is_empty() { + attempt.commit().await.map_err(map_sqlx_error)?; + tx.commit().await.map_err(map_sqlx_error)?; + return Ok(Vec::new()); + } - let mut rescued = Vec::with_capacity(moved.len()); - for row in moved { - let dlq_reason = dlq_enabled(&row.queue).then_some("deadline_exceeded"); - rescued.push( - self.settle_rescued_attempt_tx(&mut tx, row, "hard deadline exceeded", dlq_reason) - .await?, + let mut moved = match self + .try_hydrate_deleted_leases_tx(&mut attempt, deleted) + .await? + { + Ok(moved) => moved, + Err(busy) => { + attempt.rollback().await.map_err(map_sqlx_error)?; + busy.exclude(&mut excluded)?; + continue; + } + }; + moved.extend( + match self + .try_hydrate_deleted_leases_tx(&mut attempt, receipt_deleted) + .await? + { + Ok(moved) => moved, + Err(busy) => { + attempt.rollback().await.map_err(map_sqlx_error)?; + busy.exclude(&mut excluded)?; + continue; + } + }, ); + + let mut rescued = Vec::with_capacity(moved.len()); + for row in moved { + let dlq_reason = dlq_enabled(&row.queue).then_some("deadline_exceeded"); + rescued.push( + self.settle_rescued_attempt_tx( + &mut attempt, + row, + "hard deadline exceeded", + dlq_reason, + ) + .await?, + ); + } + attempt.commit().await.map_err(map_sqlx_error)?; + tx.commit().await.map_err(map_sqlx_error)?; + return Ok(rescued); } - tx.commit().await.map_err(map_sqlx_error)?; - Ok(rescued) } /// Rescue attempts whose callback timed out. Exhausted attempts fail @@ -14609,15 +15651,33 @@ impl QueueStorage { ) -> Result, AwaError> { let schema = self.schema(); let mut tx = pool.begin().await.map_err(map_sqlx_error)?; - let deleted: Vec = sqlx::query_as(&format!( + let mut excluded = BTreeSet::new(); + let mut selected: Option> = None; + loop { + let (excluded_job_ids, excluded_run_leases): (Vec<_>, Vec<_>) = + excluded.iter().copied().unzip(); + let (selected_job_ids, selected_run_leases): (Vec, Vec) = selected + .as_ref() + .map(|pairs| pairs.iter().copied().unzip()) + .unwrap_or_default(); + let mut attempt = tx.begin().await.map_err(map_sqlx_error)?; + let deleted: Vec = sqlx::query_as(&format!( r#" DELETE FROM {schema}.leases - WHERE job_id IN ( - SELECT job_id + WHERE (job_id, run_lease) IN ( + SELECT job_id, run_lease FROM {schema}.leases WHERE state = 'waiting_external' AND callback_timeout_at IS NOT NULL AND callback_timeout_at < clock_timestamp() + AND NOT EXISTS ( + SELECT 1 FROM unnest($1::bigint[], $2::bigint[]) AS excluded(job_id, run_lease) + WHERE excluded.job_id = leases.job_id AND excluded.run_lease = leases.run_lease + ) + AND ($3::boolean OR EXISTS ( + SELECT 1 FROM unnest($4::bigint[], $5::bigint[]) AS selected(job_id, run_lease) + WHERE selected.job_id = leases.job_id AND selected.run_lease = leases.run_lease + )) ORDER BY callback_timeout_at ASC LIMIT 500 FOR UPDATE SKIP LOCKED @@ -14641,27 +15701,57 @@ impl QueueStorage { callback_timeout_at "# )) - .fetch_all(tx.as_mut()) + .bind(&excluded_job_ids) + .bind(&excluded_run_leases) + .bind(selected.is_none()) + .bind(&selected_job_ids) + .bind(&selected_run_leases) + .fetch_all(attempt.as_mut()) .await .map_err(map_sqlx_error)?; - if deleted.is_empty() { - tx.commit().await.map_err(map_sqlx_error)?; - return Ok(Vec::new()); - } + selected.get_or_insert_with(|| { + deleted + .iter() + .map(|row| (row.job_id, row.run_lease)) + .collect() + }); - let moved = self.hydrate_deleted_leases_tx(&mut tx, deleted).await?; + if deleted.is_empty() { + attempt.commit().await.map_err(map_sqlx_error)?; + tx.commit().await.map_err(map_sqlx_error)?; + return Ok(Vec::new()); + } - let mut rescued = Vec::with_capacity(moved.len()); - for row in moved { - let dlq_reason = dlq_enabled(&row.queue).then_some("callback_timeout"); - rescued.push( - self.settle_rescued_attempt_tx(&mut tx, row, "callback timed out", dlq_reason) + let moved = match self + .try_hydrate_deleted_leases_tx(&mut attempt, deleted) + .await? + { + Ok(moved) => moved, + Err(busy) => { + attempt.rollback().await.map_err(map_sqlx_error)?; + busy.exclude(&mut excluded)?; + continue; + } + }; + + let mut rescued = Vec::with_capacity(moved.len()); + for row in moved { + let dlq_reason = dlq_enabled(&row.queue).then_some("callback_timeout"); + rescued.push( + self.settle_rescued_attempt_tx( + &mut attempt, + row, + "callback timed out", + dlq_reason, + ) .await?, - ); + ); + } + attempt.commit().await.map_err(map_sqlx_error)?; + tx.commit().await.map_err(map_sqlx_error)?; + return Ok(rescued); } - tx.commit().await.map_err(map_sqlx_error)?; - Ok(rescued) } /// Promote up to `batch_size` due `state` rows into the ready ring. @@ -14752,8 +15842,9 @@ impl QueueStorage { .or_default() .push(job_id); } + let mut promoted_queues = BTreeSet::new(); for (lane, job_ids) in lanes { - match self + let promoted = self .promote_lane( pool, state, @@ -14761,12 +15852,22 @@ impl QueueStorage { DueLaneRows::Selected(&job_ids), lane_lock_timeout, ) - .await? - { - Some(promoted) => outcome.promoted += promoted, - None => outcome.contended.push(lane), + .await; + match promoted { + Ok(Some(promoted)) => { + outcome.promoted += promoted; + if promoted > 0 { + promoted_queues.insert(lane.queue); + } + } + Ok(None) => outcome.contended.push(lane), + Err(err) => { + self.notify_promoted_queues(pool, promoted_queues).await; + return Err(err); + } } } + self.notify_promoted_queues(pool, promoted_queues).await; Ok(outcome) } @@ -14785,32 +15886,34 @@ impl QueueStorage { if !matches!(state, JobState::Scheduled | JobState::Retryable) || batch_size <= 0 { return Ok(PromoteOutcome::default()); } - Ok( - match self - .promote_lane( - pool, - state, - lane, - DueLaneRows::Oldest(batch_size), - lane_lock_timeout, - ) - .await? - { - Some(promoted) => PromoteOutcome { - promoted, - selected: promoted, - contended: Vec::new(), - }, - None => PromoteOutcome { - contended: vec![lane.clone()], - ..PromoteOutcome::default() - }, + let promoted = self + .promote_lane( + pool, + state, + lane, + DueLaneRows::Oldest(batch_size), + lane_lock_timeout, + ) + .await?; + if promoted.is_some_and(|promoted| promoted > 0) { + self.notify_promoted_queues(pool, [lane.queue.clone()]) + .await; + } + Ok(match promoted { + Some(promoted) => PromoteOutcome { + promoted, + selected: promoted, + contended: Vec::new(), }, - ) + None => PromoteOutcome { + contended: vec![lane.clone()], + ..PromoteOutcome::default() + }, + }) } - /// One lane's promotion transaction. `None` when it hit `lock_timeout` - /// and rolled back. + /// One lane's promotion transaction. `None` when its locks hit + /// `lock_timeout` and it rolled back. async fn promote_lane( &self, pool: &PgPool, @@ -14824,11 +15927,10 @@ impl QueueStorage { .promote_lane_tx(&mut tx, state, lane, rows, lock_timeout) .await { - Ok(promoted) => match tx.commit().await { - Ok(()) => Ok(Some(promoted)), - Err(err) if is_lock_contention_error(&err) => Ok(None), - Err(err) => Err(map_sqlx_error(err)), - }, + Ok(promoted) => { + tx.commit().await.map_err(map_sqlx_error)?; + Ok(Some(promoted)) + } Err(AwaError::Database(err)) if is_lock_contention_error(&err) => { tx.rollback().await.map_err(map_sqlx_error)?; Ok(None) @@ -14929,10 +16031,45 @@ impl QueueStorage { .collect(); self.insert_existing_ready_rows_tx(tx, ready_rows, Some(state)) .await?; - self.notify_queues_tx(tx, vec![lane.queue.clone()]).await?; Ok(promoted) } + /// Wake claimers of queues whose promotion has committed. + /// + /// Promotion notifies after its lane transactions instead of inside them. + /// A notifying transaction takes PostgreSQL's notify-queue lock at + /// `COMMIT`, which every notifying transaction in the cluster waits its + /// turn for, and a lane transaction would hold its lane's locks, and run + /// under its lock timeout, while it waited. The wakeup is a hint that + /// claimers also poll for, so a failure is logged and not returned. + async fn notify_promoted_queues( + &self, + pool: &PgPool, + queues: impl IntoIterator, + ) { + let channels: Vec = queues + .into_iter() + .map(|queue| format!("awa:{}", self.logical_queue_name(&queue))) + .collect::>() + .into_iter() + .collect(); + if channels.is_empty() { + return; + } + if let Err(err) = + sqlx::query("SELECT pg_notify(channel, '') FROM unnest($1::text[]) AS channel") + .bind(&channels) + .execute(pool) + .await + { + tracing::warn!( + error = %err, + channels = channels.len(), + "Failed to wake claimers after promotion; they pick the rows up on their next poll" + ); + } + } + async fn relation_has_rows_tx( tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, relation: &str, @@ -15871,8 +17008,12 @@ impl QueueStorage { // never-resetting sequences, so a PK collision means // corrupted terminal state and must abort the prune loudly // rather than silently drop a terminal fact. - let carried = sqlx::query(&format!( + let carry_ready_join = + done_ready_join_from("done", "ready", &ready_child, &tomb_child); + let carry_payload = done_payload_projection("done", "ready"); + let groups: Vec<(String, i16, i16, i16, i64)> = sqlx::query_as(&format!( r#" + WITH carried AS ( INSERT INTO {schema}.done_entries ( ready_slot, ready_generation, job_id, kind, queue, args, state, priority, attempt, run_lease, @@ -15888,7 +17029,7 @@ impl QueueStorage { done.queue, COALESCE(done.args, ready.args, '{{}}'::jsonb), done.state, - done.priority, + COALESCE(ready.priority, done.priority), done.attempt, done.run_lease, COALESCE(done.max_attempts, ready.max_attempts, 25::smallint), @@ -15900,70 +17041,45 @@ impl QueueStorage { COALESCE(done.created_at, ready.created_at, done.finalized_at), COALESCE(done.unique_key, ready.unique_key), COALESCE(done.unique_states, ready.unique_states), - COALESCE(done.payload, ready.payload, '{{}}'::jsonb) + {carry_payload} FROM {done_child} AS done - LEFT JOIN {ready_child} AS ready - ON ready.ready_slot = done.ready_slot - AND ready.ready_generation = done.ready_generation - AND ready.queue = done.queue - AND ready.priority = done.priority - AND ready.enqueue_shard = done.enqueue_shard - AND ready.lane_seq = done.lane_seq + {carry_ready_join} WHERE done.state = 'failed' AND done.finalized_at >= now() - make_interval(secs => $1::bigint) + RETURNING queue, priority, enqueue_shard, job_id + ) + SELECT + queue, + priority, + enqueue_shard, + mod( + mod(job_id, {TERMINAL_COUNTER_BUCKETS}::bigint) + + {TERMINAL_COUNTER_BUCKETS}::bigint, + {TERMINAL_COUNTER_BUCKETS}::bigint + )::smallint AS counter_bucket, + count(*)::bigint + FROM carried + GROUP BY queue, priority, enqueue_shard, counter_bucket + ORDER BY queue, priority, enqueue_shard, counter_bucket "# )) .bind(failed_retention_secs) .bind(state.0) .bind(current_generation) - .execute(tx.as_mut()) + .fetch_all(tx.as_mut()) .await - .map_err(map_sqlx_error)? - .rows_affected(); - - if carried > 0 { - // This slot's live counter rows and delta segment are - // wiped below, so the carried rows' exact-count - // evidence moves with them: re-append positive deltas - // under the current slot, grouped exactly like the - // completion path's delta append. - sqlx::query(&format!( - r#" - INSERT INTO {schema}.queue_terminal_count_deltas ( - ready_slot, - ready_generation, - queue, - priority, - enqueue_shard, - counter_bucket, - terminal_delta - ) - SELECT - $2, - $3, - queue, - priority, - enqueue_shard, - mod( - mod(job_id, {TERMINAL_COUNTER_BUCKETS}::bigint) - + {TERMINAL_COUNTER_BUCKETS}::bigint, - {TERMINAL_COUNTER_BUCKETS}::bigint - )::smallint AS counter_bucket, - count(*)::bigint - FROM {done_child} - WHERE state = 'failed' - AND finalized_at >= now() - make_interval(secs => $1::bigint) - GROUP BY queue, priority, enqueue_shard, counter_bucket - ORDER BY queue, priority, enqueue_shard, counter_bucket - "# - )) - .bind(failed_retention_secs) - .bind(state.0) - .bind(current_generation) - .execute(tx.as_mut()) - .await - .map_err(map_sqlx_error)?; - } + .map_err(map_sqlx_error)?; + + // The replacement delta uses each inserted row's recovered + // priority. The old slot's counter evidence is truncated below. + let carried: u64 = groups.iter().map(|(_, _, _, _, count)| *count as u64).sum(); + let deltas = groups + .into_iter() + .map(|(queue, priority, shard, bucket, count)| { + ((state.0, current_generation, queue, priority, shard, bucket), count) + }) + .collect(); + self.append_terminal_count_deltas_tx(&mut tx, deltas).await?; carried } else { 0 diff --git a/awa-model/src/read_transaction.rs b/awa-model/src/read_transaction.rs new file mode 100644 index 000000000..88155e7fb --- /dev/null +++ b/awa-model/src/read_transaction.rs @@ -0,0 +1,31 @@ +//! Protect owned read connections while SQLx is still acknowledging BEGIN. +//! +//! SQLx records transaction depth after BEGIN completes. Dropping that future +//! earlier can leave a transaction that its normal rollback guard cannot see. + +use sqlx::pool::PoolConnection; +use sqlx::{PgPool, Postgres}; + +pub(crate) struct BeginGuard { + // Callers borrow this field directly so the separate flag can be cleared + // while the resulting transaction retains its connection borrow. + pub(crate) connection: PoolConnection, + pub(crate) begin_pending: bool, +} + +impl BeginGuard { + pub(crate) async fn acquire(pool: &PgPool) -> Result { + Ok(Self { + connection: pool.acquire().await?, + begin_pending: true, + }) + } +} + +impl Drop for BeginGuard { + fn drop(&mut self) { + if self.begin_pending { + self.connection.close_on_drop(); + } + } +} diff --git a/awa-python/Cargo.lock b/awa-python/Cargo.lock index 8922f4d9a..eaf424b46 100644 --- a/awa-python/Cargo.lock +++ b/awa-python/Cargo.lock @@ -95,7 +95,7 @@ checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" [[package]] name = "awa-macros" -version = "0.6.13" +version = "0.6.14" dependencies = [ "proc-macro-crate", "proc-macro2", @@ -105,7 +105,7 @@ dependencies = [ [[package]] name = "awa-metrics" -version = "0.6.13" +version = "0.6.14" dependencies = [ "awa-model", "opentelemetry", @@ -113,7 +113,7 @@ dependencies = [ [[package]] name = "awa-model" -version = "0.6.13" +version = "0.6.14" dependencies = [ "awa-macros", "blake3", @@ -133,7 +133,7 @@ dependencies = [ [[package]] name = "awa-python" -version = "0.6.13" +version = "0.6.14" dependencies = [ "async-trait", "awa-model", @@ -156,7 +156,7 @@ dependencies = [ [[package]] name = "awa-worker" -version = "0.6.13" +version = "0.6.14" dependencies = [ "async-trait", "awa-macros", diff --git a/awa-python/Cargo.toml b/awa-python/Cargo.toml index 373a48e8e..d81c4d54e 100644 --- a/awa-python/Cargo.toml +++ b/awa-python/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "awa-python" -version = "0.6.13" +version = "0.6.14" edition = "2021" [lib] @@ -12,8 +12,8 @@ default = ["cel"] cel = ["awa-model/cel"] [dependencies] -awa-model = { path = "../awa-model", version = "0.6.13" } -awa-worker = { path = "../awa-worker", version = "0.6.13", features = ["__python-bridge"] } +awa-model = { path = "../awa-model", version = "0.6.14" } +awa-worker = { path = "../awa-worker", version = "0.6.14", features = ["__python-bridge"] } pyo3 = { version = "0.29", features = ["macros", "abi3-py310", "chrono"] } pyo3-async-runtimes = { version = "0.29", features = ["tokio-runtime"] } sqlx = { version = "0.8", features = ["runtime-tokio-rustls", "postgres", "chrono", "json"] } diff --git a/awa-python/README.md b/awa-python/README.md index e2aaa7d1e..c30ef007b 100644 --- a/awa-python/README.md +++ b/awa-python/README.md @@ -6,6 +6,8 @@ Python bindings for [awa](https://github.com/hardbyte/awa), a Postgres-native ba pip install awa-pg ``` +CI tests Python 3.14 and 3.15, the newest two stable versions. The package remains installable on Python 3.10–3.13; those versions receive best-effort support. + ## Quick start ```python diff --git a/awa-python/pyproject.toml b/awa-python/pyproject.toml index b20fdc21d..38e7cef76 100644 --- a/awa-python/pyproject.toml +++ b/awa-python/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "maturin" [project] name = "awa-pg" requires-python = ">=3.10" -version = "0.6.13" +version = "0.6.14" description = "Postgres-native background job queue — Python SDK with async/sync workers, transactional enqueue, and progress tracking. Install awa-pg[ui] to add the bundled web dashboard." readme = "README.md" license = {text = "MIT OR Apache-2.0"} @@ -17,6 +17,8 @@ classifiers = [ "Programming Language :: Rust", "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.14", + "Programming Language :: Python :: 3.15", "License :: OSI Approved :: MIT License", "License :: OSI Approved :: Apache Software License", "Topic :: Database", @@ -32,7 +34,7 @@ classifiers = [ # # Kept opt-in so the default `awa-pg` install stays small — workers and # producers don't need the ~10 MB UI bundle. -ui = ["awa-cli==0.6.13"] +ui = ["awa-cli==0.6.14"] [project.urls] Homepage = "https://github.com/hardbyte/awa" diff --git a/awa-python/uv.lock b/awa-python/uv.lock index 8b3cbff7a..c18f79069 100644 --- a/awa-python/uv.lock +++ b/awa-python/uv.lock @@ -88,12 +88,12 @@ wheels = [ [[package]] name = "awa-cli" -version = "0.6.13" +version = "0.6.14" source = { directory = "../awa-cli" } [[package]] name = "awa-pg" -version = "0.6.13" +version = "0.6.14" source = { editable = "." } [package.optional-dependencies] diff --git a/awa-seaorm/Cargo.toml b/awa-seaorm/Cargo.toml index 781c95ab9..336c4de99 100644 --- a/awa-seaorm/Cargo.toml +++ b/awa-seaorm/Cargo.toml @@ -14,6 +14,6 @@ serde_json.workspace = true sqlx.workspace = true [dev-dependencies] -awa-testing = { path = "../awa-testing", version = "0.6.13" } +awa-testing = { path = "../awa-testing", version = "0.6.14" } serde = { workspace = true, features = ["derive"] } tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/awa-worker/src/completion.rs b/awa-worker/src/completion.rs index ac54468ca..0872bc392 100644 --- a/awa-worker/src/completion.rs +++ b/awa-worker/src/completion.rs @@ -367,17 +367,66 @@ impl CompletionWorker { ); } Err(err) => { - self.metrics.record_completion_flush( - self.shard_id, - job_ids.len() as u64, - flush_start.elapsed(), - ); warn!( error = %err, shard = self.shard_id, batch_size = job_ids.len(), "Failed to flush completed job batch" ); + if matches!( + &err, + AwaError::Database(sqlx::Error::Database(error)) + if error.code().as_deref() == Some("55P03") + ) { + if let RuntimeStorage::QueueStorage(runtime) = &self.storage { + let mut committed_pairs = HashSet::new(); + for request in batch { + let pair = (request.job_id, request.run_lease); + let response = if committed_pairs.contains(&pair) { + Ok(true) + } else { + let updated = if let Some(runtime_job) = &request.runtime_job { + runtime + .store + .complete_runtime_batch( + &self.pool, + std::slice::from_ref(runtime_job), + ) + .await + } else if let Some(claim) = &request.claim { + runtime + .store + .complete_claimed_batch( + &self.pool, + std::slice::from_ref(claim), + ) + .await + } else { + runtime + .store + .complete_job_batch_by_id(&self.pool, &[pair]) + .await + }; + updated.map(|updated| updated.contains(&pair)) + }; + if matches!(&response, Ok(true)) { + committed_pairs.insert(pair); + } + let _ = request.response.send(response); + } + self.metrics.record_completion_flush( + self.shard_id, + job_ids.len() as u64, + flush_start.elapsed(), + ); + return; + } + } + self.metrics.record_completion_flush( + self.shard_id, + job_ids.len() as u64, + flush_start.elapsed(), + ); let message = format!("completion batch flush failed: {err}"); for request in batch { let _ = request @@ -1022,4 +1071,203 @@ mod tests { jobs.len() as f64 / elapsed.as_secs_f64() ); } + async fn queue_storage_busy_completion_responses(with_claim: bool) { + let pool = setup(6).await; + let schema = format!("awa_complete_busy_{}", uuid::Uuid::new_v4().simple()); + let runtime = crate::storage::QueueStorageRuntime::new( + awa_model::QueueStorageConfig { + schema: schema.clone(), + queue_slot_count: 4, + lease_slot_count: 2, + claim_slot_count: 2, + queue_stripe_count: 1, + lease_claim_receipts: false, + }, + Duration::from_secs(1), + Duration::from_secs(1), + ) + .unwrap(); + let store = runtime.store.clone(); + store.prepare_schema(&pool).await.unwrap(); + let mut jobs = Vec::new(); + for queue in ["busy", "ordinary", "stale", "carried"] { + store.enqueue_batch(&pool, queue, 4, 1).await.unwrap(); + let mut claimed = store + .claim_runtime_batch(&pool, queue, 1, Duration::from_secs(300)) + .await + .unwrap(); + assert_eq!(claimed.len(), 1); + jobs.push(claimed.remove(0)); + } + let historical = jobs.remove(0); + let ordinary = jobs.remove(0); + let stale = jobs.remove(0); + let carried = jobs.remove(0); + assert_eq!( + store + .complete_runtime_batch(&pool, std::slice::from_ref(&stale)) + .await + .unwrap(), + vec![(stale.job.id, stale.job.run_lease)] + ); + store.enqueue_batch(&pool, "busy", 3, 1).await.unwrap(); + sqlx::query(&format!( + "INSERT INTO {schema}.ready_tombstones + (ready_slot,ready_generation,queue,priority,enqueue_shard,lane_seq,job_id) + SELECT ready_slot,ready_generation,queue,priority,enqueue_shard,lane_seq,job_id + FROM {schema}.ready_entries WHERE job_id=$1;" + )) + .bind(historical.job.id) + .execute(&pool) + .await + .unwrap(); + sqlx::query(&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,3::smallint,attempt,run_lease,max_attempts, + {schema}.reserve_enqueue_seq(queue,3::smallint,enqueue_shard,1::bigint),enqueue_shard, + run_at,attempted_at,created_at,unique_key,unique_states,payload + FROM {schema}.ready_entries WHERE job_id=$1" + )) + .bind(historical.job.id) + .execute(&pool) + .await + .unwrap(); + let mut blocker = pool.begin().await.unwrap(); + sqlx::query(&format!( + "SELECT 1 FROM {schema}.queue_claim_heads + WHERE queue='busy' AND priority=3 AND enqueue_shard=0 FOR UPDATE" + )) + .fetch_one(blocker.as_mut()) + .await + .unwrap(); + let pairs = [ + (historical.job.id, historical.job.run_lease), + (ordinary.job.id, ordinary.job.run_lease), + ]; + let strict = if with_claim { + store + .complete_claimed_batch(&pool, &[historical.claim.clone(), ordinary.claim.clone()]) + .await + } else { + store.complete_job_batch_by_id(&pool, &pairs).await + }; + assert!(matches!(strict, Err(AwaError::Database(ref error)) + if error.as_database_error().is_some_and(|error| error.code().as_deref()==Some("55P03")))); + let mut caller_tx = pool.begin().await.unwrap(); + let strict = store + .complete_job_batch_by_id_in_tx(&mut caller_tx, &pairs) + .await; + assert!(matches!(strict, Err(AwaError::Database(ref error)) + if error.as_database_error().is_some_and(|error| error.code().as_deref()==Some("55P03")))); + caller_tx.rollback().await.unwrap(); + assert_eq!( + store + .load_job(&pool, ordinary.job.id) + .await + .unwrap() + .unwrap() + .state, + awa_model::JobState::Running + ); + + let (mut batcher, _) = CompletionBatcher::new( + pool.clone(), + CancellationToken::new(), + crate::metrics::AwaMetrics::from_global(), + RuntimeStorage::QueueStorage(runtime), + 1, + ); + let worker = batcher.workers.pop().unwrap(); + let mut batch = Vec::new(); + let mut responses = Vec::new(); + for job in [&historical, &ordinary, &stale, &carried] { + let (response, rx) = oneshot::channel(); + batch.push(CompletionRequest { + job_id: job.job.id, + run_lease: job.job.run_lease, + claim: with_claim.then(|| job.claim.clone()), + runtime_job: (job.job.id == carried.job.id).then(|| job.clone()), + response, + }); + responses.push(rx); + } + tokio::time::timeout(Duration::from_secs(5), worker.flush(&mut batch)) + .await + .expect("contended batch isolates singleton errors without waiting"); + let busy = responses.remove(0).await.unwrap(); + assert!( + matches!(busy, Err(AwaError::Database(ref error)) + if error.as_database_error().is_some_and(|error| error.code().as_deref()==Some("55P03"))), + "busy attempt must report its real error, never stale: {busy:?}" + ); + assert!(responses.remove(0).await.unwrap().unwrap()); + assert!(!responses.remove(0).await.unwrap().unwrap()); + assert!(responses.remove(0).await.unwrap().unwrap()); + assert_eq!( + store + .load_job(&pool, historical.job.id) + .await + .unwrap() + .unwrap() + .state, + awa_model::JobState::Running + ); + assert_eq!( + store + .load_job(&pool, ordinary.job.id) + .await + .unwrap() + .unwrap() + .state, + awa_model::JobState::Completed + ); + blocker.rollback().await.unwrap(); + let (response, rx) = oneshot::channel(); + worker + .flush(&mut vec![CompletionRequest { + job_id: historical.job.id, + run_lease: historical.job.run_lease, + claim: with_claim.then(|| historical.claim.clone()), + runtime_job: None, + response, + }]) + .await; + assert!(rx.await.unwrap().unwrap()); + let terminals: Vec<(i64, i64, i64)> = sqlx::query_as(&format!( + "SELECT job_id,run_lease,count(*) FROM {schema}.done_entries + WHERE job_id=ANY($1) GROUP BY job_id,run_lease ORDER BY job_id" + )) + .bind(vec![ + historical.job.id, + ordinary.job.id, + stale.job.id, + carried.job.id, + ]) + .fetch_all(&pool) + .await + .unwrap(); + let mut expected: Vec<_> = [&historical, &ordinary, &stale, &carried] + .into_iter() + .map(|job| (job.job.id, job.job.run_lease, 1_i64)) + .collect(); + expected.sort_unstable(); + assert_eq!(terminals, expected); + sqlx::query(&format!("DROP SCHEMA {schema} CASCADE")) + .execute(&pool) + .await + .unwrap(); + pool.close().await; + } + + #[tokio::test] + async fn queue_storage_id_completion_contention_isolates_responses() { + queue_storage_busy_completion_responses(false).await; + } + + #[tokio::test] + async fn queue_storage_claim_completion_contention_isolates_responses() { + queue_storage_busy_completion_responses(true).await; + } } diff --git a/awa-worker/src/executor.rs b/awa-worker/src/executor.rs index aa6c5da9b..8276a24a5 100644 --- a/awa-worker/src/executor.rs +++ b/awa-worker/src/executor.rs @@ -14,7 +14,7 @@ use std::time::Duration; use tokio::sync::Semaphore; use tokio_util::sync::CancellationToken; use tokio_util::task::TaskTracker; -use tracing::{error, info, info_span, warn, Instrument}; +use tracing::{debug, error, info, info_span, warn, Instrument}; /// Result of executing a job handler. /// @@ -372,6 +372,10 @@ impl JobExecutor { progress: progress_state.clone(), }; in_flight.insert((job_id, job_run_lease), in_flight_state); + let heartbeat_guard = storage.queue_storage().map(|_| FinalizationHeartbeatGuard { + in_flight: in_flight.clone(), + key: (job_id, job_run_lease), + }); if let Some(counter) = queue_in_flight.get(&job_queue) { counter.fetch_add(1, Ordering::SeqCst); } @@ -448,12 +452,9 @@ impl JobExecutor { guard.clone_latest() }; - // Leaving `in_flight` ends the attempt's heartbeat, so only a - // bounded number of finished attempts may wait for finalization. - // While that queue is full the attempt keeps its worker slot and - // its heartbeat, and the dispatcher stops claiming: a completion - // backlog cannot outgrow `heartbeat_staleness` and get finished - // attempts rescued and run again. + // Keep worker capacity and heartbeats until a bounded finalizer slot + // is available. Otherwise returning handlers can create an unbounded + // set of tasks retaining job payloads while completion is stalled. let finalize_permit = match result { Some(_) => tokio::select! { permit = finalize_capacity.acquire_owned() => { @@ -463,7 +464,10 @@ impl JobExecutor { }, None => None, }; - in_flight.remove((job_id, job_run_lease)); + + if heartbeat_guard.is_none() { + in_flight.remove((job_id, job_run_lease)); + } if let Some(counter) = queue_in_flight.get(&job_queue) { counter.fetch_sub(1, Ordering::SeqCst); } @@ -473,17 +477,9 @@ impl JobExecutor { return; }; let dlq_enabled = dlq_policy.enabled_for(&job_queue); - // Graceful shutdown waits for this before it stops the completion - // batcher: the attempt has already left `in_flight`, so nothing - // heartbeats it, and a finalize lost at exit leaves it to rescue. finalizers.spawn(async move { - // Every finalize path is one lease-guarded transaction, and a - // deadlock abort rolls the whole transaction back, so running - // it again is safe. Without the retry the job would sit - // `running` until heartbeat rescue re-ran the handler. - let mut deadlock_retries = 0u32; - let outcome = loop { - let outcome = complete_job( + let outcome = finalize_with_retry( + || complete_job( &pool, &job, queue_storage_claim.as_ref(), @@ -497,29 +493,15 @@ impl JobExecutor { &storage, dlq_enabled, &metrics, - ) - .await; - match outcome { - Err(err) - if err.is_deadlock() - && deadlock_retries < COMPLETION_DEADLOCK_RETRIES => - { - deadlock_retries += 1; - metrics.record_completion_deadlock_retry(&job_kind, &job_queue); - let delay = completion_deadlock_backoff(deadlock_retries); - warn!( - job_id, - retry = deadlock_retries, - max_retries = COMPLETION_DEADLOCK_RETRIES, - delay_ms = delay.as_millis() as u64, - "Completion transaction aborted by the deadlock detector; retrying" - ); - tokio::time::sleep(delay).await; - } - outcome => break outcome, - } - }; - drop(finalize_permit); + ), + &metrics, + job_id, + &job_kind, + &job_queue, + COMPLETION_LOCK_WINDOW, + storage.queue_storage().is_some(), + ) + .await; match &outcome { Ok(CompletionOutcome::Applied { terminal, .. }) => { @@ -590,6 +572,8 @@ impl JobExecutor { } } + drop(heartbeat_guard); + drop(finalize_permit); if let Ok(CompletionOutcome::Applied { event: Some(event), .. }) = outcome @@ -602,17 +586,120 @@ impl JobExecutor { } } +struct FinalizationHeartbeatGuard { + in_flight: InFlightMap, + key: (i64, i64), +} + +impl Drop for FinalizationHeartbeatGuard { + fn drop(&mut self) { + self.in_flight.remove(self.key); + } +} + /// How many times a finalize transaction is re-run after Postgres aborts it /// to break a lock cycle before the failure is surfaced. const COMPLETION_DEADLOCK_RETRIES: u32 = 3; -/// Short exponential backoff between deadlock retries (10ms, 20ms, 40ms, ...). -/// Postgres aborts only one party of a cycle, so the survivor has usually -/// committed by the time the retry runs. -fn completion_deadlock_backoff(retry: u32) -> Duration { +fn completion_contention_backoff(retry: u32) -> Duration { Duration::from_millis(10u64 << retry.saturating_sub(1).min(4)) } +#[derive(Debug, thiserror::Error)] +enum FinalizeError { + #[error(transparent)] + Operation(#[from] AwaError), + #[error("durable finalization lock-retry window expired")] + Deadline, +} + +#[allow(clippy::too_many_arguments)] +async fn finalize_with_retry( + mut finalize: F, + metrics: &crate::metrics::AwaMetrics, + job_id: i64, + kind: &str, + queue: &str, + window: Duration, + retry_lock_contention: bool, +) -> Result +where + F: FnMut() -> Fut, + Fut: std::future::Future>, +{ + let deadline = retry_lock_contention.then(|| tokio::time::Instant::now() + window); + let mut deadlock_retries = 0; + let mut lock_retries = 0; + loop { + if deadline.is_some_and(|deadline| tokio::time::Instant::now() >= deadline) { + return Err(FinalizeError::Deadline); + } + let outcome = finalize().await; + let error = match outcome { + Ok(outcome) => return Ok(outcome), + Err(error) => error, + }; + let delay = if error.is_deadlock() && deadlock_retries < COMPLETION_DEADLOCK_RETRIES { + deadlock_retries += 1; + metrics.record_completion_deadlock_retry(kind, queue); + let delay = completion_contention_backoff(deadlock_retries); + warn!( + job_id, + retry = deadlock_retries, + max_retries = COMPLETION_DEADLOCK_RETRIES, + delay_ms = delay.as_millis() as u64, + "Completion transaction aborted by the deadlock detector; retrying" + ); + delay + } else if retry_lock_contention && is_completion_lock_contention(&error) { + lock_retries += 1; + let delay = completion_contention_backoff(lock_retries); + if lock_retries == 1 { + warn!(job_id, error = %error, delay_ms = delay.as_millis() as u64, "Durable finalization failed; retrying"); + } else { + debug!(job_id, error = %error, delay_ms = delay.as_millis() as u64, "Durable finalization failed; retrying"); + } + delay + } else { + return Err(error.into()); + }; + match deadline { + Some(deadline) => tokio::time::timeout_at(deadline, tokio::time::sleep(delay)) + .await + .map_err(|_| FinalizeError::Deadline)?, + None => tokio::time::sleep(delay).await, + } + } +} + +/// Re-read a job whose transition has committed, for its lifecycle event. +/// The transition stands whatever the read returns, so a failed read is logged +/// and the caller builds the event from the worker's copy instead. +async fn reload_committed_job( + runtime: &QueueStorageRuntime, + pool: &PgPool, + job_id: i64, +) -> Option { + match runtime.store.load_job(pool, job_id).await { + Ok(job) => job, + Err(err) => { + warn!( + job_id, + error = %err, + "Could not re-read a committed job for its lifecycle event; using the worker's copy" + ); + None + } + } +} + +fn is_completion_lock_contention(error: &AwaError) -> bool { + matches!(error, AwaError::Database(sqlx::Error::Database(database)) + if database.code().as_deref() == Some("55P03")) +} + +const COMPLETION_LOCK_WINDOW: Duration = Duration::from_secs(120); + /// Update job state in the database based on handler result. /// /// Returns a `CompletionOutcome` indicating whether the state transition was @@ -1773,18 +1860,12 @@ async fn complete_job_queue_storage( return Ok(CompletionOutcome::IgnoredStale); } if needs_event { - let updated_job = - runtime - .store - .load_job(pool, job.id) - .await? - .unwrap_or_else(|| { - let mut completed_job = job.clone(); - completed_job.state = JobState::Completed; - completed_job.finalized_at = Some(chrono::Utc::now()); - completed_job.progress = None; - completed_job - }); + // A completed attempt changes only these fields, so the event is + // built from the worker's copy 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; Ok(CompletionOutcome::Applied { event: Some(UntypedJobEvent::Completed { job: updated_job, @@ -2124,10 +2205,8 @@ async fn complete_job_queue_storage( } } let event = if needs_event { - let parked_job = runtime - .store - .load_job(pool, job.id) - .await? + let parked_job = reload_committed_job(runtime, pool, job.id) + .await .unwrap_or_else(|| { let mut parked = job.clone(); parked.state = JobState::WaitingExternal; @@ -3086,3 +3165,738 @@ async fn retry_backoff_canonical_with_followups( tx.commit().await?; Ok(Some(updated_job)) } + +#[cfg(test)] +mod tests { + use super::*; + use tokio_util::sync::CancellationToken; + + fn test_metrics() -> crate::metrics::AwaMetrics { + crate::metrics::AwaMetrics::new(&opentelemetry::global::meter("finalize-test")) + } + + #[derive(Debug)] + struct FinalizeDatabaseError(&'static str); + + impl std::fmt::Display for FinalizeDatabaseError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "sqlstate {}", self.0) + } + } + + impl std::error::Error for FinalizeDatabaseError {} + + impl sqlx::error::DatabaseError for FinalizeDatabaseError { + fn message(&self) -> &str { + "finalize transaction rejected" + } + fn code(&self) -> Option> { + Some(std::borrow::Cow::Borrowed(self.0)) + } + fn as_error(&self) -> &(dyn std::error::Error + Send + Sync + 'static) { + self + } + fn as_error_mut(&mut self) -> &mut (dyn std::error::Error + Send + Sync + 'static) { + self + } + fn into_error(self: Box) -> Box { + self + } + fn kind(&self) -> sqlx::error::ErrorKind { + sqlx::error::ErrorKind::Other + } + } + + fn finalize_sqlstate(code: &'static str) -> AwaError { + AwaError::Database(sqlx::Error::Database(Box::new(FinalizeDatabaseError(code)))) + } + + #[tokio::test(start_paused = true)] + async fn finalize_lock_contention_preserves_the_sealed_outcome() { + let calls = AtomicU32::new(0); + let applied = AtomicU32::new(0); + let sealed = serde_json::json!({ + "outcome": "retry_after", "delay_secs": 60, + "progress": {"checkpoint": 7}, "followups": [{"id": 42}], + }); + let outcome = finalize_with_retry( + || { + std::future::ready(if calls.fetch_add(1, Ordering::SeqCst) < 4 { + Err(finalize_sqlstate("55P03")) + } else { + applied.fetch_add(1, Ordering::SeqCst); + Ok(sealed.clone()) + }) + }, + &test_metrics(), + 1, + "test", + "test", + Duration::from_secs(2), + true, + ) + .await + .expect("contention must retry the saved outcome"); + assert_eq!(outcome, sealed); + assert_eq!(calls.load(Ordering::SeqCst), 5); + assert_eq!(applied.load(Ordering::SeqCst), 1); + } + + #[tokio::test(start_paused = true)] + async fn finalize_lock_contention_keeps_the_absolute_deadline() { + let calls = AtomicU32::new(0); + let started = tokio::time::Instant::now(); + let result = finalize_with_retry( + || { + calls.fetch_add(1, Ordering::SeqCst); + std::future::ready(Err::<(), _>(finalize_sqlstate("55P03"))) + }, + &test_metrics(), + 1, + "test", + "test", + Duration::from_millis(100), + true, + ) + .await; + assert!(matches!(result, Err(FinalizeError::Deadline)), "{result:?}"); + assert_eq!(started.elapsed(), Duration::from_millis(100)); + assert!((3..=5).contains(&calls.load(Ordering::SeqCst))); + } + + #[tokio::test(start_paused = true)] + async fn accepted_finalization_observes_its_acknowledgement_after_the_retry_window() { + let calls = Arc::new(AtomicU32::new(0)); + let events = Arc::new(AtomicU32::new(0)); + let started = Arc::new(tokio::sync::Notify::new()); + let (acknowledge, response) = tokio::sync::oneshot::channel::(); + let observed = events.clone(); + let handler: BoxedUntypedEventHandler = Arc::new(move |event| { + assert!(matches!(event, UntypedJobEvent::Completed { job, .. } + if job.id == 1 && job.run_lease == 1)); + observed.fetch_add(1, Ordering::SeqCst); + Box::pin(async {}) + }); + let handlers = Arc::new(HashMap::from([("finalize_parent".into(), vec![handler])])); + let attempts = calls.clone(); + let notification = started.clone(); + let task = tokio::spawn(async move { + let mut response = Some(response); + let outcome = finalize_with_retry( + || { + attempts.fetch_add(1, Ordering::SeqCst); + let response = response.take().expect("accepted work is never restarted"); + let notification = notification.clone(); + async move { + notification.notify_one(); + let event = response.await.map_err(|_| { + AwaError::Validation("completion acknowledgement dropped".into()) + })?; + Ok(CompletionOutcome::Applied { + event: Some(event), + terminal: false, + }) + } + }, + &test_metrics(), + 1, + "finalize_parent", + "test", + Duration::from_millis(100), + true, + ) + .await?; + let CompletionOutcome::Applied { + event: Some(event), .. + } = outcome + else { + panic!("accepted completion must retain its lifecycle event") + }; + dispatch_lifecycle_event(&handlers, "finalize_parent", event).await; + Ok::<(), FinalizeError>(()) + }); + started.notified().await; + tokio::time::advance(Duration::from_millis(110)).await; + tokio::task::yield_now().await; + assert!( + !task.is_finished(), + "the retry window must not cancel accepted work" + ); + let now = chrono::Utc::now(); + let job = serde_json::from_value(serde_json::json!({ + "id": 1, "kind": "finalize_parent", "queue": "test", "args": {}, + "state": "completed", "priority": 2, "attempt": 1, "run_lease": 1, + "max_attempts": 5, "run_at": now, "attempted_at": now, + "created_at": now, "finalized_at": now, "metadata": {}, "tags": [], + })) + .unwrap(); + assert!(acknowledge + .send(UntypedJobEvent::Completed { + job, + duration: Duration::ZERO + }) + .is_ok()); + task.await.unwrap().unwrap(); + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(events.load(Ordering::SeqCst), 1); + } + + #[tokio::test(start_paused = true)] + async fn finalize_lock_contention_does_not_retry_unrelated_statement_errors() { + for code in ["57014", "23505", "55P04"] { + let calls = AtomicU32::new(0); + let result = finalize_with_retry( + || { + calls.fetch_add(1, Ordering::SeqCst); + std::future::ready(Err::<(), _>(finalize_sqlstate(code))) + }, + &test_metrics(), + 1, + "test", + "test", + Duration::from_secs(2), + true, + ) + .await; + assert!( + matches!(result, Err(FinalizeError::Operation(_))), + "{code}: {result:?}" + ); + assert_eq!(calls.load(Ordering::SeqCst), 1, "{code}"); + } + } + + #[tokio::test(start_paused = true)] + async fn canonical_finalization_does_not_retry_lock_contention() { + let calls = AtomicU32::new(0); + let result = finalize_with_retry( + || { + calls.fetch_add(1, Ordering::SeqCst); + std::future::ready(Err::<(), _>(finalize_sqlstate("55P03"))) + }, + &test_metrics(), + 1, + "test", + "test", + Duration::from_secs(2), + false, + ) + .await; + assert!(matches!(result, Err(FinalizeError::Operation(_)))); + assert_eq!(calls.load(Ordering::SeqCst), 1); + } + + #[tokio::test(start_paused = true)] + async fn canonical_finalization_preserves_three_deadlock_retries_without_a_deadline() { + let calls = AtomicU32::new(0); + let started = tokio::time::Instant::now(); + let result = finalize_with_retry( + || { + calls.fetch_add(1, Ordering::SeqCst); + std::future::ready(Err::<(), _>(finalize_sqlstate("40P01"))) + }, + &test_metrics(), + 1, + "test", + "test", + Duration::ZERO, + false, + ) + .await; + assert!(matches!(result, Err(FinalizeError::Operation(_)))); + assert_eq!(calls.load(Ordering::SeqCst), 4); + assert_eq!(started.elapsed(), Duration::from_millis(70)); + } + + struct FinalizeWorker; + + #[async_trait::async_trait] + impl Worker for FinalizeWorker { + fn kind(&self) -> &'static str { + "finalize_parent" + } + + async fn perform(&self, _context: &JobContext) -> Result { + Ok(JobResult::Completed) + } + } + + #[tokio::test(start_paused = true)] + async fn queue_finalization_retains_heartbeat_without_holding_worker_capacity() { + let metrics = test_metrics(); + let now = chrono::Utc::now(); + let job: JobRow = serde_json::from_value(serde_json::json!({ + "id": 1, "kind": "finalize_parent", "queue": "test", "args": {}, + "state": "running", "priority": 2, "attempt": 1, "run_lease": 1, + "max_attempts": 5, "run_at": now, "attempted_at": now, + "created_at": now, "metadata": {}, "tags": [], + })) + .unwrap(); + let key = (job.id, job.run_lease); + let runtime = QueueStorageRuntime::new( + awa_model::QueueStorageConfig::default(), + Duration::from_secs(1), + Duration::from_secs(1), + ) + .unwrap(); + let storage = RuntimeStorage::QueueStorage(runtime); + let pool = sqlx::postgres::PgPoolOptions::new() + .connect_lazy("postgres://postgres:test@localhost/awa_test") + .unwrap(); + let finalizers = TaskTracker::new(); + let in_flight = Arc::new(crate::runtime::InFlightRegistry::default()); + let count = Arc::new(AtomicU32::new(0)); + let (unstarted, handle) = crate::completion::CompletionBatcher::new( + pool.clone(), + CancellationToken::new(), + metrics.clone(), + storage.clone(), + 1, + ); + let executor = JobExecutor::new( + pool.clone(), + Arc::new(HashMap::from([( + "finalize_parent".into(), + Box::new(FinalizeWorker) as BoxedWorker, + )])), + Arc::new(HashMap::new()), + Arc::new(HashMap::new()), + in_flight.clone(), + Arc::new(HashMap::from([("test".into(), count.clone())])), + Arc::new(HashMap::new()), + metrics, + handle, + storage, + DlqPolicy::default(), + finalizers.clone(), + CancellationToken::new(), + 1, + ); + executor + .execute_task( + DispatchedJob { + job, + queue_storage_claim: Some(ClaimedEntry { + queue: "test".into(), + priority: 2, + lane_seq: 1, + ready_slot: 0, + ready_generation: 0, + lease_slot: 0, + lease_generation: 0, + claim_slot: 0, + receipt_id: Some(1), + claim_batch_id: Some(1), + claim_batch_index: Some(1), + lease_claim_receipt: true, + enqueue_shard: 0, + }), + queue_storage_unique_states: None, + abandon_after: None, + }, + Arc::new(AtomicBool::new(false)), + ) + .await; + assert_eq!( + count.load(Ordering::SeqCst), + 0, + "handler capacity must be released" + ); + assert_eq!(finalizers.len(), 1); + assert!( + in_flight.get_cancel(key).is_some(), + "heartbeat ownership must survive handler return" + ); + tokio::task::yield_now().await; + assert!( + in_flight.get_cancel(key).is_some(), + "pending acknowledgement must retain the lease" + ); + tokio::time::advance(COMPLETION_LOCK_WINDOW + Duration::from_secs(1)).await; + tokio::task::yield_now().await; + assert_eq!(count.load(Ordering::SeqCst), 0); + assert_eq!(finalizers.len(), 1); + assert!( + in_flight.get_cancel(key).is_some(), + "accepted completion must retain heartbeat ownership beyond the retry window" + ); + pool.close().await; + drop(unstarted); + finalizers.close(); + tokio::time::timeout(Duration::from_secs(1), finalizers.wait()) + .await + .unwrap(); + assert!( + in_flight.get_cancel(key).is_none(), + "settlement must unregister the exact lease" + ); + } + + #[tokio::test(start_paused = true)] + async fn saturated_finalizer_capacity_retains_worker_until_shutdown() { + for storage in [ + RuntimeStorage::Canonical, + RuntimeStorage::QueueStorage( + QueueStorageRuntime::new( + awa_model::QueueStorageConfig::default(), + Duration::from_secs(1), + Duration::from_secs(1), + ) + .unwrap(), + ), + ] { + let now = chrono::Utc::now(); + let job: JobRow = serde_json::from_value(serde_json::json!({ + "id": 1, "kind": "finalize_parent", "queue": "test", "args": {}, + "state": "running", "priority": 2, "attempt": 1, "run_lease": 1, + "max_attempts": 5, "run_at": now, "attempted_at": now, + "created_at": now, "metadata": {}, "tags": [], + })) + .unwrap(); + let key = (job.id, job.run_lease); + let pool = sqlx::postgres::PgPoolOptions::new() + .connect_lazy("postgres://postgres:test@localhost/awa_test") + .unwrap(); + let metrics = test_metrics(); + let finalizers = TaskTracker::new(); + let finalize_cancel = CancellationToken::new(); + let in_flight = Arc::new(crate::runtime::InFlightRegistry::default()); + let count = Arc::new(AtomicU32::new(0)); + let (_unstarted, handle) = crate::completion::CompletionBatcher::new( + pool.clone(), + CancellationToken::new(), + metrics.clone(), + storage.clone(), + 1, + ); + let executor = JobExecutor::new( + pool.clone(), + Arc::new(HashMap::from([( + "finalize_parent".into(), + Box::new(FinalizeWorker) as BoxedWorker, + )])), + Arc::new(HashMap::new()), + Arc::new(HashMap::new()), + in_flight.clone(), + Arc::new(HashMap::from([("test".into(), count.clone())])), + Arc::new(HashMap::new()), + metrics, + handle, + storage, + DlqPolicy::default(), + finalizers.clone(), + finalize_cancel.clone(), + 1, + ); + // Occupy the entire budget without starting database work. + let capacity = executor.finalize_capacity.clone(); + let held = capacity + .clone() + .acquire_many_owned(capacity.available_permits() as u32) + .await + .unwrap(); + let task = tokio::spawn(executor.execute_task( + DispatchedJob { + job, + queue_storage_claim: None, + queue_storage_unique_states: None, + abandon_after: None, + }, + Arc::new(AtomicBool::new(false)), + )); + tokio::task::yield_now().await; + assert!(!task.is_finished()); + assert_eq!(count.load(Ordering::SeqCst), 1); + assert!(in_flight.get_cancel(key).is_some()); + assert!(finalizers.is_empty()); + + // Cancelling background services at the drain deadline must release + // the capacity waiter without spawning another finalizer. + finalize_cancel.cancel(); + tokio::time::timeout(Duration::from_secs(1), task) + .await + .unwrap() + .unwrap(); + assert_eq!(count.load(Ordering::SeqCst), 0); + assert!(in_flight.get_cancel(key).is_none()); + assert!(finalizers.is_empty()); + drop(held); + pool.close().await; + } + } + + #[tokio::test] + async fn cancelled_finalization_unregisters_only_its_exact_lease() { + let in_flight = Arc::new(crate::runtime::InFlightRegistry::default()); + for key in [(1, 1), (1, 2)] { + in_flight.insert( + key, + InFlightState { + cancel: Arc::new(AtomicBool::new(false)), + progress: Arc::new(std::sync::Mutex::new(ProgressState::new(None))), + }, + ); + } + let guard = FinalizationHeartbeatGuard { + in_flight: in_flight.clone(), + key: (1, 1), + }; + let started = Arc::new(tokio::sync::Notify::new()); + let notification = started.clone(); + let task = tokio::spawn(async move { + let _guard = guard; + notification.notify_one(); + std::future::pending::<()>().await; + }); + started.notified().await; + task.abort(); + assert!(task.await.unwrap_err().is_cancelled()); + assert!(in_flight.get_cancel((1, 1)).is_none()); + assert!(in_flight.get_cancel((1, 2)).is_some()); + } + struct RetainedFinalizeFollowup; + + impl crate::enqueue_specs::EnqueueFollowUp for RetainedFinalizeFollowup { + fn run<'a>( + &'a self, + connection: &'a mut sqlx::PgConnection, + job: &'a JobRow, + _context: Option<&'a crate::enqueue_specs::OutcomeContext>, + ) -> std::pin::Pin> + Send + 'a>> + { + Box::pin(async move { + sqlx::query("INSERT INTO finalize_followups(parent) VALUES ($1)") + .bind(job.id) + .execute(connection) + .await?; + Ok(()) + }) + } + } + + #[sqlx::test(migrations = false)] + async fn retained_head_contention_replays_the_worker_outcome_once(pool: PgPool) { + awa_model::migrations::run(&pool).await.unwrap(); + sqlx::query("CREATE TABLE finalize_followups(parent bigint PRIMARY KEY)") + .execute(&pool) + .await + .unwrap(); + let runtime = QueueStorageRuntime::new( + awa_model::QueueStorageConfig { + queue_stripe_count: 1, + ..Default::default() + }, + Duration::from_secs(1), + Duration::from_secs(1), + ) + .unwrap(); + let queue = "retained_finalize"; + runtime.store.prepare_queue(&pool, queue, 1).await.unwrap(); + runtime + .store + .enqueue_params_batch( + &pool, + &[awa_model::job::InsertParams { + kind: "retained_finalize_parent".into(), + args: serde_json::json!({"body":"retained"}), + opts: awa_model::InsertOpts { + queue: queue.into(), + priority: 4, + max_attempts: 3, + metadata: serde_json::json!({"preserved":true}), + ..Default::default() + }, + }], + ) + .await + .unwrap(); + let claimed = runtime + .store + .claim_runtime_batch(&pool, queue, 1, Duration::ZERO) + .await + .unwrap() + .pop() + .unwrap(); + let pair = (claimed.job.id, claimed.job.run_lease); + let schema = runtime.store.schema(); + let mut historical = pool.begin().await.unwrap(); + sqlx::query(&format!( + "INSERT INTO {schema}.ready_tombstones + (ready_slot,ready_generation,queue,priority,enqueue_shard,lane_seq,job_id) + SELECT ready_slot,ready_generation,queue,priority,enqueue_shard,lane_seq,job_id + FROM {schema}.ready_entries WHERE job_id=$1" + )) + .bind(pair.0) + .execute(historical.as_mut()) + .await + .unwrap(); + sqlx::query(&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,3,attempt,run_lease,max_attempts, + {schema}.reserve_enqueue_seq(queue,3::smallint,enqueue_shard,1::bigint),enqueue_shard, + run_at,attempted_at,created_at,unique_key,unique_states,payload + FROM {schema}.ready_entries WHERE job_id=$1" + )).bind(pair.0).execute(historical.as_mut()).await.unwrap(); + sqlx::query(&format!( + "INSERT INTO {schema}.ready_segments + (ready_slot,ready_generation,queue,priority,enqueue_shard,first_lane_seq,next_lane_seq,first_run_at,created_at) + SELECT ready_slot,ready_generation,queue,priority,enqueue_shard,lane_seq,lane_seq+1,run_at,created_at + FROM {schema}.ready_entries WHERE job_id=$1 AND priority=3 ON CONFLICT DO NOTHING" + )).bind(pair.0).execute(historical.as_mut()).await.unwrap(); + historical.commit().await.unwrap(); + let progress = serde_json::json!({"checkpoint":7}); + runtime + .store + .flush_progress(&pool, pair.0, pair.1, progress.clone()) + .await + .unwrap(); + let mut blocker = pool.begin().await.unwrap(); + sqlx::query(&format!( + "SELECT 1 FROM {schema}.queue_claim_heads + WHERE queue=$1 AND priority=3 AND enqueue_shard=0 FOR UPDATE" + )) + .bind(queue) + .fetch_one(blocker.as_mut()) + .await + .unwrap(); + let metrics = test_metrics(); + let storage = RuntimeStorage::QueueStorage(runtime.clone()); + let (_unstarted, handle) = crate::completion::CompletionBatcher::new( + pool.clone(), + CancellationToken::new(), + metrics.clone(), + storage.clone(), + 1, + ); + let specs: Arc = Arc::new(HashMap::from([( + crate::enqueue_specs::Outcome::Retried, + HashMap::from([( + "retained_finalize_parent".into(), + vec![Arc::new(RetainedFinalizeFollowup) as crate::enqueue_specs::BoxedEnqueueSpec], + )]), + )])); + let intended = Ok(JobResult::RetryAfter(Duration::from_secs(60))); + let calls = AtomicU32::new(0); + let busy = tokio::sync::Notify::new(); + let finalize = finalize_with_retry( + || { + calls.fetch_add(1, Ordering::SeqCst); + let (pool, job, claim, intended, progress, handle, specs, storage, metrics, busy) = ( + &pool, + &claimed.job, + &claimed.claim, + &intended, + &progress, + &handle, + &specs, + &storage, + &metrics, + &busy, + ); + async move { + let result = complete_job( + pool, + job, + Some(claim), + None, + intended, + handle, + Some(progress.clone()), + Duration::ZERO, + true, + specs, + storage, + false, + metrics, + ) + .await; + if let Err(error) = &result { + assert!( + is_completion_lock_contention(error), + "expected actual held-head55P03: {error}" + ); + busy.notify_one(); + } + result + } + }, + &metrics, + pair.0, + &claimed.job.kind, + queue, + Duration::from_secs(5), + true, + ); + let release = async { + tokio::time::timeout(Duration::from_secs(3), busy.notified()) + .await + .expect("worker transition must encounter the held historical head"); + let active = runtime + .store + .load_job(&pool, pair.0) + .await + .unwrap() + .unwrap(); + assert_eq!(active.state, JobState::Running); + assert_eq!(active.run_lease, pair.1); + assert_eq!(active.progress, Some(progress.clone())); + let effects: i64 = sqlx::query_scalar("SELECT count(*) FROM finalize_followups") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(effects, 0, "failed transaction cannot publish its followup"); + blocker.rollback().await.unwrap(); + }; + let (ack, ()) = tokio::join!(finalize, release); + let ack = ack.expect("saved worker outcome must survive transient head contention"); + assert!(calls.load(Ordering::SeqCst) >= 2); + let CompletionOutcome::Applied { + event: Some(UntypedJobEvent::Retried { job, .. }), + terminal: false, + } = ack + else { + panic!("worker must acknowledge its intended retry outcome") + }; + assert_eq!(job.id, pair.0); + assert_eq!(job.run_lease, pair.1); + assert_eq!(job.state, JobState::Retryable); + assert_eq!(job.args, claimed.job.args); + assert_eq!(job.max_attempts, claimed.job.max_attempts); + assert_eq!(job.progress, Some(progress.clone())); + assert_eq!(job.metadata["preserved"], true); + let replay = complete_job( + &pool, + &claimed.job, + Some(&claimed.claim), + None, + &intended, + &handle, + Some(progress), + Duration::ZERO, + true, + &specs, + &storage, + false, + &metrics, + ) + .await + .unwrap(); + assert!(matches!(replay, CompletionOutcome::IgnoredStale)); + let effects: i64 = + sqlx::query_scalar("SELECT count(*) FROM finalize_followups WHERE parent=$1") + .bind(pair.0) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!( + effects, 1, + "the acknowledged outcome publishes exactly one followup" + ); + let deferred:i64 = sqlx::query_scalar(&format!( + "SELECT count(*) FROM {schema}.deferred_jobs WHERE job_id=$1 AND run_lease=$2 AND state='retryable'" + )).bind(pair.0).bind(pair.1).fetch_one(&pool).await.unwrap(); + assert_eq!(deferred, 1); + } +} diff --git a/awa/Cargo.toml b/awa/Cargo.toml index 23ba05b62..60275e5c9 100644 --- a/awa/Cargo.toml +++ b/awa/Cargo.toml @@ -20,7 +20,7 @@ awa-macros.workspace = true awa-worker.workspace = true [dev-dependencies] -awa-testing = { path = "../awa-testing", version = "0.6.13" } +awa-testing = { path = "../awa-testing", version = "0.6.14" } sqlx.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/awa/tests/admin_read_cancel_test.rs b/awa/tests/admin_read_cancel_test.rs new file mode 100644 index 000000000..4ed4a4957 --- /dev/null +++ b/awa/tests/admin_read_cancel_test.rs @@ -0,0 +1,449 @@ +//! Cancel a real public read after PostgreSQL has entered its transaction, +//! before SQLx receives ReadyForQuery. No sleeps or custom BEGIN statements. + +mod ci_timing; + +use awa::model::admin::{self, ListJobsFilter}; +use awa::model::{insert_many, migrations, InsertOpts, InsertParams, QueueStorage}; +use awa::{AwaError, JobState, QueueStorageConfig}; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions, PgSslMode}; +use sqlx::PgPool; +use std::io; +use std::str::FromStr; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::time::Duration; +use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt}; +use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::{watch, Notify}; +use tokio::task::{JoinHandle, JoinSet}; + +type TestResult = Result>; + +fn bound() -> Duration { + ci_timing::scaled_timeout(Duration::from_secs(5)) +} + +#[derive(Default)] +struct BeginGate { + armed: AtomicBool, + held: watch::Sender>, + release: Notify, +} + +impl BeginGate { + async fn held_pid(&self) -> TestResult { + let mut observed = self.held.subscribe(); + loop { + if let Some(pid) = *observed.borrow_and_update() { + return Ok(pid); + } + observed.changed().await?; + } + } +} + +struct Proxy { + port: u16, + gate: Arc, + task: JoinHandle<()>, +} + +impl Proxy { + async fn start(options: &PgConnectOptions) -> TestResult { + // The proxy is local and test-only; credentials and backend secrets are + // forwarded unchanged and never logged. SSL is disabled only on it. + assert!( + options.get_socket().is_none(), + "this TCP fixture needs a TCP DATABASE_URL" + ); + let upstream = (options.get_host().to_string(), options.get_port()); + let listener = TcpListener::bind(("127.0.0.1", 0)).await?; + let port = listener.local_addr()?.port(); + let gate = Arc::new(BeginGate::default()); + let connections_gate = gate.clone(); + let task = tokio::spawn(async move { + let mut connections = JoinSet::new(); + loop { + tokio::select! { + accepted = listener.accept() => { + let Ok((client, _)) = accepted else { break }; + let upstream = upstream.clone(); + let gate = connections_gate.clone(); + connections.spawn(async move { relay(client, upstream, gate).await }); + } + Some(_finished) = connections.join_next(), if !connections.is_empty() => {} + } + } + }); + Ok(Self { port, gate, task }) + } + + async fn stop(self) { + self.gate.release.notify_one(); + self.task.abort(); + let _ = self.task.await; + // Dropping the listener's JoinSet aborts every relay and closes both + // socket halves, including a relay still holding ReadyForQuery. + } +} + +async fn frame(stream: &mut R) -> io::Result)>> { + let tag = match stream.read_u8().await { + Ok(tag) => tag, + Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => return Ok(None), + Err(error) => return Err(error), + }; + let length = stream.read_u32().await?; + if !(4..=1024 * 1024).contains(&length) { + return Err(io::Error::other("invalid test proxy frame length")); + } + let mut payload = vec![0; (length - 4) as usize]; + stream.read_exact(&mut payload).await?; + Ok(Some((tag, payload))) +} + +async fn forward(stream: &mut W, tag: u8, payload: &[u8]) -> io::Result<()> { + stream.write_u8(tag).await?; + stream.write_u32((payload.len() + 4) as u32).await?; + stream.write_all(payload).await +} + +async fn relay( + mut client: TcpStream, + upstream: (String, u16), + gate: Arc, +) -> io::Result<()> { + let mut server = TcpStream::connect((upstream.0.as_str(), upstream.1)).await?; + // Startup packets precede the normal tagged PostgreSQL messages. + let length = client.read_u32().await?; + if !(8..=64 * 1024).contains(&length) { + return Err(io::Error::other("invalid startup packet")); + } + let mut startup = vec![0; (length - 4) as usize]; + client.read_exact(&mut startup).await?; + server.write_u32(length).await?; + server.write_all(&startup).await?; + + let (mut client_read, mut client_write) = client.into_split(); + let (mut server_read, mut server_write) = server.into_split(); + let awaiting_begin = Arc::new(AtomicBool::new(false)); + let frontend_gate = gate.clone(); + let frontend_begin = awaiting_begin.clone(); + let frontend = async move { + while let Some((tag, payload)) = frame(&mut client_read).await? { + if tag == b'Q' { + let sql = String::from_utf8_lossy(&payload); + let is_begin = sql + .trim_end_matches('\0') + .split_whitespace() + .next() + .is_some_and(|first| first.eq_ignore_ascii_case("BEGIN")); + if is_begin && frontend_gate.armed.swap(false, Ordering::SeqCst) { + frontend_begin.store(true, Ordering::SeqCst); + } + } + forward(&mut server_write, tag, &payload).await?; + } + Ok::<_, io::Error>(()) + }; + let backend = async move { + let mut pid = 0; + let mut begin_completed = false; + while let Some((tag, payload)) = frame(&mut server_read).await? { + if tag == b'K' && payload.len() == 8 { + pid = i32::from_be_bytes(payload[..4].try_into().expect("four PID bytes")); + } + if tag == b'C' + && payload.as_slice() == b"BEGIN\0" + && awaiting_begin.load(Ordering::SeqCst) + { + begin_completed = true; + } + if tag == b'Z' && begin_completed && awaiting_begin.swap(false, Ordering::SeqCst) { + if payload.as_slice() != b"T" || pid <= 0 { + return Err(io::Error::other("BEGIN did not enter a server transaction")); + } + // CommandComplete was forwarded, but the real transaction + // status frame stays here until the read future is dropped. + gate.held.send_replace(Some(pid)); + gate.release.notified().await; + begin_completed = false; + } + forward(&mut client_write, tag, &payload).await?; + } + Ok::<_, io::Error>(()) + }; + tokio::select! { result = frontend => result, result = backend => result } +} + +struct Fixture { + admin: PgPool, + observer: PgPool, + reader: PgPool, + store: Arc, + database: String, + job_id: i64, + proxy: Proxy, +} + +impl Fixture { + async fn new() -> TestResult { + let url = awa_testing::setup::database_url(); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect(&url) + .await?; + let database = format!("awa_read_cancel_{}", uuid::Uuid::new_v4().simple()); + sqlx::query(&format!("CREATE DATABASE {database}")) + .execute(&admin) + .await?; + let options = PgConnectOptions::from_str(&url)?.database(&database); + let observer = PgPoolOptions::new() + .max_connections(2) + .connect_with(options.clone()) + .await?; + migrations::run(&observer).await?; + awa_testing::setup::activate_queue_storage(&observer).await; + let store = Arc::new(QueueStorage::new(QueueStorageConfig::default())?); + let rows = insert_many( + &observer, + &[InsertParams { + kind: "admin_read_cancel".to_string(), + args: serde_json::json!({}), + opts: InsertOpts { + queue: "admin_read_cancel".to_string(), + ..Default::default() + }, + }], + ) + .await?; + let proxy = Proxy::start(&options).await?; + let reader = PgPoolOptions::new() + .max_connections(1) + .acquire_timeout(bound()) + .connect_with( + options + .host("127.0.0.1") + .port(proxy.port) + .ssl_mode(PgSslMode::Disable), + ) + .await?; + Ok(Self { + admin, + observer, + reader, + store, + database, + job_id: rows[0].id, + proxy, + }) + } + + async fn close(self) -> TestResult { + self.proxy.gate.release.notify_one(); + // Stop relays even if a broken connection cannot be returned normally. + let closed = tokio::time::timeout(bound(), self.reader.close()).await; + self.proxy.stop().await; + if closed.is_err() { + tokio::time::timeout(bound(), self.reader.close()).await?; + } + self.observer.close().await; + tokio::time::timeout( + bound(), + sqlx::query(&format!("DROP DATABASE {} WITH (FORCE)", self.database)) + .execute(&self.admin), + ) + .await??; + self.admin.close().await; + Ok(()) + } +} + +#[derive(Clone, Copy, Debug)] +enum Read { + Job, + Counts, + Queue, + Kinds, + Timeseries, + Jobs, +} + +async fn public_read(which: Read, store: &QueueStorage, pool: &PgPool, job_id: i64) -> TestResult { + match which { + Read::Job => { + let job = store.load_job(pool, job_id).await?.expect("job exists"); + assert_eq!((job.id, job.state), (job_id, JobState::Available)); + } + Read::Counts => assert_eq!( + admin::state_counts(pool).await?.get(&JobState::Available), + Some(&1) + ), + Read::Queue => assert_eq!( + admin::queue_overview(pool, "admin_read_cancel") + .await? + .unwrap() + .available, + 1 + ), + Read::Kinds => assert_eq!( + admin::job_kind_overviews(pool) + .await? + .iter() + .find(|kind| kind.kind == "admin_read_cancel") + .unwrap() + .job_count, + 1 + ), + Read::Timeseries => assert_eq!( + admin::state_timeseries(pool, 60) + .await? + .iter() + .filter(|bucket| bucket.state == JobState::Available) + .map(|bucket| bucket.count) + .sum::(), + 1 + ), + Read::Jobs => assert_eq!( + admin::list_jobs(pool, &ListJobsFilter::default()) + .await? + .iter() + .map(|job| job.id) + .collect::>(), + vec![job_id] + ), + } + Ok(()) +} + +async fn backend_is_clean(observer: &PgPool, pid: i32) -> TestResult { + tokio::time::timeout(bound(), async { + loop { + let state: Option<(String, bool)> = sqlx::query_as( + "SELECT state, xact_start IS NULL FROM pg_stat_activity WHERE pid = $1", + ) + .bind(pid) + .fetch_optional(observer) + .await?; + if state + .as_ref() + .is_none_or(|(state, no_tx)| state == "idle" && *no_tx) + { + return Ok::<_, sqlx::Error>(()); + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .map_err(|_| { + io::Error::other(format!( + "backend {pid} retained a transaction after read cancellation" + )) + })??; + Ok(()) +} + +async fn cancelled_begin(which: Read) -> TestResult { + let fixture = Fixture::new().await?; + let reader = fixture.reader.clone(); + let observer = fixture.observer.clone(); + let store = fixture.store.clone(); + let gate = fixture.proxy.gate.clone(); + let job_id = fixture.job_id; + let tested = tokio::spawn(async move { + gate.armed.store(true, Ordering::SeqCst); + let mut read = Box::pin(public_read(which, &store, &reader, job_id)); + let pid = tokio::select! { + held = tokio::time::timeout(bound(), gate.held_pid()) => held??, + result = &mut read => return Err(io::Error::other(format!("{which:?} completed before real BEGIN gate: {result:?}")).into()), + }; + let state: (String, bool) = sqlx::query_as( + "SELECT state, xact_start IS NOT NULL FROM pg_stat_activity WHERE pid = $1") + .bind(pid).fetch_one(&observer).await?; + assert_eq!(state, ("idle in transaction".to_string(), true), "real BEGIN reached server"); + drop(read); + gate.release.notify_one(); + backend_is_clean(&observer, pid).await?; + tokio::time::timeout(bound(), public_read(which, &store, &reader, job_id)).await??; + Ok::<_, Box>(()) + }).await; + let cleanup = fixture.close().await; + tested??; + cleanup +} + +#[tokio::test] +async fn cancelled_load_job_begin_does_not_return_open_transaction() -> TestResult { + cancelled_begin(Read::Job).await +} + +#[tokio::test] +async fn cancelled_state_counts_begin_does_not_return_open_transaction() -> TestResult { + cancelled_begin(Read::Counts).await +} + +#[tokio::test] +async fn cancelled_queue_overview_begin_does_not_return_open_transaction() -> TestResult { + cancelled_begin(Read::Queue).await +} + +#[tokio::test] +async fn cancelled_kind_overviews_begin_does_not_return_open_transaction() -> TestResult { + cancelled_begin(Read::Kinds).await +} + +#[tokio::test] +async fn cancelled_timeseries_begin_does_not_return_open_transaction() -> TestResult { + cancelled_begin(Read::Timeseries).await +} + +#[tokio::test] +async fn cancelled_job_list_begin_does_not_return_open_transaction() -> TestResult { + cancelled_begin(Read::Jobs).await +} + +#[tokio::test] +async fn completed_begin_and_query_error_preserve_healthy_connection() -> TestResult { + let fixture = Fixture::new().await?; + let reader = fixture.reader.clone(); + let observer = fixture.observer.clone(); + let store = fixture.store.clone(); + let job_id = fixture.job_id; + let tested = tokio::spawn(async move { + let pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&reader) + .await?; + public_read(Read::Job, &store, &reader, job_id).await?; + public_read(Read::Counts, &store, &reader, job_id).await?; + sqlx::query("ALTER TABLE awa.deferred_jobs RENAME TO cancel_test_deferred") + .execute(&observer) + .await?; + let error = store + .load_job(&reader, job_id) + .await + .expect_err("real post-BEGIN query error"); + assert!( + matches!(error, AwaError::Database(sqlx::Error::Database(ref db)) + if db.code().as_deref() == Some("42P01")), + "expected undefined table, got {error:?}" + ); + backend_is_clean(&observer, pid).await?; + sqlx::query("ALTER TABLE awa.cancel_test_deferred RENAME TO deferred_jobs") + .execute(&observer) + .await?; + public_read(Read::Job, &store, &reader, job_id).await?; + let reused: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&reader) + .await?; + assert_eq!( + reused, pid, + "completed BEGIN and ordinary query errors retain healthy reuse" + ); + Ok::<_, Box>(()) + }) + .await; + let cleanup = fixture.close().await; + tested??; + cleanup +} diff --git a/awa/tests/admin_read_equivalence_test.rs b/awa/tests/admin_read_equivalence_test.rs index 6de0bbb74..cf122a03a 100644 --- a/awa/tests/admin_read_equivalence_test.rs +++ b/awa/tests/admin_read_equivalence_test.rs @@ -26,9 +26,6 @@ struct Fixture { store: QueueStorage, /// Jobs inserted, by kind. inserted: BTreeMap, - /// Jobs left running on open receipt claims, by kind. `current_jobs` - /// on this branch doesn't list them. - on_open_claims: BTreeMap, } impl Fixture { @@ -56,7 +53,6 @@ impl Fixture { pool, store, inserted: BTreeMap::new(), - on_open_claims: BTreeMap::new(), }) } @@ -255,9 +251,6 @@ impl Fixture { .await?; } // The rest stay running on open receipt claims. - for claim in iter { - *self.on_open_claims.entry(claim.job.kind).or_default() += 1; - } // Running on open receipt claims, with a callback registered on some // (which materialises a lease), and one waiting on its callback. @@ -275,12 +268,6 @@ impl Fixture { .await?; } let waiting = &running[2]; - for claim in &running[3..] { - *self - .on_open_claims - .entry(claim.job.kind.clone()) - .or_default() += 1; - } let callback = self .store .register_callback( @@ -299,9 +286,7 @@ impl Fixture { let cancelled = self .enqueue("epsilon", "eq_cancel", 2, 9, &[], None) .await?; - for claim in self.claim("eq_cancel", 2).await? { - *self.on_open_claims.entry(claim.job.kind).or_default() += 1; - } + self.claim("eq_cancel", 2).await?; for job in &cancelled[4..7] { admin::cancel(&self.pool, job.id).await?; } @@ -457,10 +442,9 @@ async fn job_kind_overviews_match_the_current_jobs_definition() -> TestResult { .iter() .find(|overview| &overview.kind == kind) .unwrap_or_else(|| panic!("no overview for {kind}")); - let listed = inserted - fixture.on_open_claims.get(kind).copied().unwrap_or(0); assert_eq!( - overview.job_count, listed, - "job_count for {kind}: every listed job is counted once" + overview.job_count, *inserted, + "job_count for {kind}: every inserted job is counted once, including open receipts" ); } let by_kind: BTreeMap<_, _> = actual diff --git a/awa/tests/fixtures/v0614_alpha1_lane_sequence_repair.sql b/awa/tests/fixtures/v0614_alpha1_lane_sequence_repair.sql new file mode 100644 index 000000000..70bcfefcc --- /dev/null +++ b/awa/tests/fixtures/v0614_alpha1_lane_sequence_repair.sql @@ -0,0 +1,83 @@ +CREATE OR REPLACE FUNCTION awa.repair_lane_sequence(p_schema TEXT, p_sequence TEXT) +RETURNS TABLE ( + sequence_name TEXT, + owner TEXT, + expected_owner TEXT, + outcome TEXT, + missing_grants TEXT[], + reassign_sql TEXT, + repair_roles TEXT[] +) +LANGUAGE plpgsql +SECURITY INVOKER +SET search_path = pg_catalog, pg_temp +SET lock_timeout = '300ms' +AS $repair$ +DECLARE + v_problem RECORD; + v_remaining RECORD; + v_busy BOOLEAN := false; + v_set_mode TEXT := CASE WHEN current_setting('server_version_num')::int >= 160000 + THEN 'SET' ELSE 'MEMBER' END; + v_statement TEXT; +BEGIN + SELECT * INTO v_problem + FROM awa.lane_sequence_problems(p_schema, p_sequence) AS problem; + IF NOT FOUND THEN + RETURN; + END IF; + + BEGIN + -- ALTER ... OWNER TO needs the current owner's privileges, SET ROLE to + -- the new owner, and the new owner's CREATE on the schema (which a + -- schema owner has). Check first so an unrelated owner is skipped + -- without waiting for the lock. + IF v_problem.reassign_sql IS NOT NULL + AND pg_has_role(current_user, to_regrole(quote_ident(v_problem.owner)), 'USAGE') + AND pg_has_role(current_user, to_regrole(quote_ident(v_problem.expected_owner)), v_set_mode) THEN + BEGIN + EXECUTE v_problem.reassign_sql; + EXCEPTION WHEN insufficient_privilege THEN + NULL; + END; + END IF; + + SELECT * INTO v_remaining + FROM awa.lane_sequence_problems(p_schema, p_sequence) AS problem; + IF FOUND + AND cardinality(v_remaining.missing_grants) > 0 + AND pg_has_role(current_user, to_regrole(quote_ident(v_remaining.owner)), 'USAGE') THEN + FOREACH v_statement IN ARRAY v_remaining.missing_grants LOOP + BEGIN + EXECUTE v_statement; + EXCEPTION WHEN insufficient_privilege THEN + NULL; + END; + END LOOP; + END IF; + EXCEPTION WHEN lock_not_available THEN + v_busy := true; + END; + + SELECT * INTO v_remaining + FROM awa.lane_sequence_problems(p_schema, p_sequence) AS problem; + sequence_name := v_problem.sequence_name; + expected_owner := v_problem.expected_owner; + IF FOUND THEN + owner := v_remaining.owner; + missing_grants := v_remaining.missing_grants; + reassign_sql := v_remaining.reassign_sql; + repair_roles := v_remaining.repair_roles; + outcome := CASE WHEN v_busy THEN 'busy' + WHEN cardinality(v_remaining.missing_grants) = 0 THEN 'owner_unchanged' + ELSE 'unrepaired' END; + ELSE + owner := v_problem.expected_owner; + missing_grants := ARRAY[]::TEXT[]; + reassign_sql := NULL; + repair_roles := ARRAY[]::TEXT[]; + outcome := 'repaired'; + END IF; + RETURN NEXT; +END +$repair$; diff --git a/awa/tests/lane_ownership_roles_test.rs b/awa/tests/lane_ownership_roles_test.rs index 026511f77..f7b85eea0 100644 --- a/awa/tests/lane_ownership_roles_test.rs +++ b/awa/tests/lane_ownership_roles_test.rs @@ -776,6 +776,123 @@ async fn migration_repairs_lanes_it_may_and_reports_the_rest() -> TestResult { cluster.teardown().await } +#[tokio::test] +async fn set_only_schema_owner_membership_repairs_legacy_lanes_and_restores_role() -> TestResult { + let cluster = Cluster::new(Schema::BeforeProvisioning, false).await?; + let queue = "set_only_owner_lane"; + assert_eq!(enqueue(&cluster.migrator, queue, 1, 1).await?, 1); + let sequences = cluster.lane_sequence_owners(queue).await?; + assert_eq!(sequences.len(), 2); + for (sequence, owner) in &sequences { + assert_eq!(owner, &cluster.migrator_name); + cluster + .migrator + .execute(sqlx::raw_sql(&sql(format!( + "REVOKE ALL ON SEQUENCE awa.{sequence} FROM {}", + cluster.runtime_name + )))) + .await?; + } + for (_, _, statement) in migrations::schema_patch_sql(migrations::SCHEMA_PATCHES) { + cluster.owner.execute(sqlx::raw_sql(&statement)).await?; + } + cluster + .owner + .execute(sqlx::raw_sql(include_str!( + "fixtures/v0614_alpha1_lane_sequence_repair.sql" + ))) + .await?; + let pending = migrations::pending_schema_patches(&cluster.owner).await?; + assert_eq!(pending.len(), 1, "{pending:?}"); + assert_eq!(pending[0].name, "definer_lane_provisioning"); + let report = migrations::run_with_report(&cluster.owner).await?; + assert!(report.complete(), "{:?}", report.patches); + assert!(migrations::pending_schema_patches(&cluster.owner) + .await? + .is_empty()); + cluster + .owner + .execute(sqlx::raw_sql(&sql(format!( + "GRANT USAGE, CREATE ON SCHEMA awa TO {migrator}; \ + GRANT EXECUTE ON FUNCTION awa.repair_lane_sequence(text, text) TO {migrator}", + migrator = cluster.migrator_name + )))) + .await?; + cluster + .admin + .execute(sqlx::raw_sql(&sql(format!( + "GRANT {owner} TO {migrator} WITH INHERIT FALSE; \ + GRANT {owner} TO {migrator} WITH SET TRUE", + owner = cluster.owner_name, + migrator = cluster.migrator_name + )))) + .await?; + + let mut conn = cluster.migrator.acquire().await?; + let (inherits_owner, can_set_owner, can_create): (bool, bool, bool) = sqlx::query_as( + "SELECT pg_has_role(current_user, $1, 'USAGE'), \ + pg_has_role(current_user, $1, 'SET'), \ + has_schema_privilege(current_user, 'awa', 'USAGE') \ + AND has_schema_privilege(current_user, 'awa', 'CREATE')", + ) + .bind(&cluster.owner_name) + .fetch_one(&mut *conn) + .await?; + assert!(!inherits_owner); + assert!(can_set_owner); + assert!(can_create); + let role_before: (String, String) = + sqlx::query_as("SELECT current_user::text, current_setting('role')") + .fetch_one(&mut *conn) + .await?; + assert_eq!(role_before.0, cluster.migrator_name); + + let problems = storage::lane_sequence_problems(&mut *conn, "awa").await?; + assert_eq!(problems.len(), 2, "{problems:?}"); + for problem in &problems { + assert_eq!(problem.owner, cluster.migrator_name); + assert!(!problem.missing_grants.is_empty(), "{problem:?}"); + assert!( + problem.repair_roles.contains(&cluster.migrator_name), + "{problem:?}" + ); + } + let repaired = storage::repair_lane_sequences(&mut conn, "awa").await?; + assert_eq!(repaired.len(), 2, "{repaired:?}"); + for repair in &repaired { + assert_eq!(repair.outcome, LaneSequenceOutcome::Repaired, "{repair:?}"); + assert_eq!(repair.owner, cluster.owner_name); + assert!(repair.missing_grants.is_empty(), "{repair:?}"); + assert!(repair.reassign_sql.is_none(), "{repair:?}"); + } + let role_after: (String, String) = + sqlx::query_as("SELECT current_user::text, current_setting('role')") + .fetch_one(&mut *conn) + .await?; + assert_eq!(role_after, role_before); + assert!(storage::lane_sequence_problems(&mut *conn, "awa") + .await? + .is_empty()); + drop(conn); + + for (sequence, _) in sequences { + let runtime_grants_complete: bool = sqlx::query_scalar( + "SELECT has_sequence_privilege($1, $2, 'USAGE') \ + AND has_sequence_privilege($1, $2, 'SELECT') \ + AND has_sequence_privilege($1, $2, 'UPDATE')", + ) + .bind(&cluster.runtime_name) + .bind(format!("awa.{sequence}")) + .fetch_one(&cluster.owner) + .await?; + assert!(runtime_grants_complete, "{sequence}"); + } + cluster.assert_lanes_owned_by_owner(queue, 1).await?; + assert_eq!(enqueue(&cluster.app_b, queue, 1, 1).await?, 1); + assert_eq!(drain(&cluster.app_b, queue).await?, 2); + cluster.teardown().await +} + /// Reassigning a lane sequence waits on any open enqueue or claim on the lane, /// so the repair bounds that wait and holds each sequence's lock only for its /// own short transaction. A held sequence is skipped and reported; the rest diff --git a/awa/tests/persona_soak/abandon_oracle.rs b/awa/tests/persona_soak/abandon_oracle.rs new file mode 100644 index 000000000..ba35cb24f --- /dev/null +++ b/awa/tests/persona_soak/abandon_oracle.rs @@ -0,0 +1,1222 @@ +//! Checks for the claim-and-abandon persona: a process holding thousands of claims is +//! SIGKILLed again and again. After each kill the rescue backlog has to clear within a +//! bound, the claim ring has to keep rotating, no job whose owner is alive may be +//! rescued, and attempts and the dead-letter queue have to follow `max_attempts`. + +use super::config::Persona; +use super::shared::Shared; +use serde::Serialize; +use sqlx::PgPool; +use std::collections::HashMap; +use std::io::Write as _; +use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, AtomicI64, Ordering}; +use std::sync::Mutex; +use std::time::{SystemTime, UNIX_EPOCH}; + +const MAX_ROWS: usize = 25; +/// The ring rotates about once a second; this long without a rotation is a stall. +const RING_STALL_AFTER_MS: i64 = 6_000; +/// Owner of attempts run by the surviving in-process runtime, which is never killed. +pub const FINISHER: u32 = u32::MAX; + +pub fn epoch_ms() -> i64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|elapsed| elapsed.as_millis() as i64) + .unwrap_or(0) +} + +#[derive(Debug, Clone, Copy)] +pub struct Attempt { + pub attempt: i16, + pub owner: u32, + pub at_ms: i64, +} + +#[derive(Debug, Clone, Serialize)] +pub struct Kill { + pub number: u32, + pub generation: u32, + /// Milliseconds since the run started, and since the epoch, when SIGKILL was sent. + pub at_run_ms: i64, + pub signalled_epoch_ms: i64, + /// Handlers the process had started, which is what it held when it died. + pub started_by_process: i64, + /// Open claims on the persona's queues right after the kill. + pub open_after_kill: i64, + /// Claims of the process that were rescued before it was killed. + pub rescued_while_alive: i64, + pub expired_peak: i64, + pub oldest_expired_peak_s: f64, + pub cleared_after_s: Option, + pub reported_slow: bool, + pub load_ended: bool, +} + +#[derive(Debug, Clone, Serialize)] +pub struct RingEpisode { + pub start_ms: i64, + pub end_ms: Option, + pub generation: i64, + pub slot: i32, + /// The last time the slot the ring has to enter next held a claim whose owner is + /// alive, which explains a stall. Starts at the beginning of the episode. + pub last_live_ms: i64, + pub expired_in_next_slot: i64, + pub reported: bool, +} + +#[derive(Default)] +struct RingMonitor { + last: Option<(i32, i64)>, + last_change_ms: i64, + episode: Option, + episodes: Vec, +} + +#[derive(Debug, Clone, Copy, Serialize)] +pub struct Sample { + pub at_ms: i64, + pub open: i64, + pub expired: i64, + pub oldest_expired_s: f64, + pub ring_generation: i64, + pub probe_ms: i64, +} + +#[derive(Default)] +pub struct AbandonState { + /// Open claims on the persona's queues at the last probe. + pub open_claims: AtomicI64, + /// A child process is running. + pub child_alive: AtomicBool, + /// Handlers running in the child process at its last report. + pub child_active: AtomicI64, + started_by_process: Mutex>, + /// The running child process and the epoch milliseconds at which it was started. + live: Mutex>, + /// Most claims of a process seen rescued while it was alive, per process. + live_lost: Mutex>, + attempts: Mutex>>, + max_attempts: Mutex>, + /// Epoch milliseconds at which each process generation was killed. + killed_at: Mutex>, + kills: Mutex>, + samples: Mutex>, + ring: Mutex, + csv: Mutex>, + passes_csv: Mutex>, + heartbeat_last: Mutex>, +} + +pub struct Probe { + pub open: i64, + pub expired: i64, + pub oldest_expired_s: f64, + /// Open claims taken at or after the `since` the probe was given. + pub open_since: i64, +} + +/// The process that is alive, and how many handlers it had started before the probe. +#[derive(Debug, Clone, Copy)] +pub struct LiveCheck { + pub generation: u32, + pub started: i64, +} + +impl AbandonState { + pub fn set_csv(&self, path: PathBuf) { + if let Ok(mut file) = std::fs::File::create(&path) { + let _ = writeln!( + file, + "run_s,open_claims,expired_claims,oldest_expired_s,claim_ring_generation,probe_ms" + ); + } + let passes = path.with_file_name("abandon-heartbeat-passes.csv"); + if let Ok(mut file) = std::fs::File::create(&passes) { + let _ = writeln!(file, "run_s,handlers_running,claims_sampled,pass_ms"); + } + *self.passes_csv.lock().expect("passes csv") = Some(passes); + *self.csv.lock().expect("csv") = Some(path); + } + + /// Folds in one reading of the oldest and newest heartbeat (epoch seconds) among the live + /// process's claims. While a pass runs, the oldest heartbeat is still one the previous pass + /// wrote, so it cannot be newer than that pass's newest. The first reading whose oldest + /// heartbeat is newer was taken just after a pass finished, and the spread from oldest to + /// newest is its duration (to within the last chunk). The samples are bucketed by handler + /// count so the report shows how pass time grows with what a process holds. `now` is the + /// database clock at the reading. + pub fn note_heartbeat_sample( + &self, + shared: &Shared, + rows: i64, + oldest: f64, + newest: f64, + now: f64, + active: i64, + ) { + let mut last = self.heartbeat_last.lock().expect("heartbeat last"); + let Some(settled_newest) = *last else { + *last = Some(newest); + return; + }; + if oldest <= settled_newest || active <= 0 { + return; + } + *last = Some(newest); + drop(last); + // A reading taken long after the pass finished may include the next pass's first + // chunks, so only a fresh one measures the pass. + if now - newest > 2.0 { + shared + .stats(Persona::ClaimAbandon) + .add("heartbeat_passes_unmeasured", 1); + return; + } + let ms = ((newest - oldest) * 1000.0).round() as i64; + let stats = shared.stats(Persona::ClaimAbandon); + stats.observe("heartbeat_pass_ms", ms); + let bucket = match active { + ..=4_999 => "under 5k", + 5_000..=9_999 => "5k-10k", + 10_000..=14_999 => "10k-15k", + 15_000..=19_999 => "15k-20k", + _ => "20k and over", + }; + stats.observe(&format!("heartbeat_pass_ms (handlers {bucket})"), ms); + if let Some(path) = self.passes_csv.lock().expect("passes csv").as_ref() { + if let Ok(mut file) = std::fs::OpenOptions::new().append(true).open(path) { + let _ = writeln!( + file, + "{:.1},{active},{rows},{ms}", + shared.now_ms() as f64 / 1000.0 + ); + } + } + } + + pub fn register_job(&self, job_id: i64, max_attempts: i16) { + self.max_attempts + .lock() + .expect("max attempts") + .insert(job_id, max_attempts); + } + + pub fn note_attempt(&self, job_id: i64, attempt: Attempt) { + if attempt.owner != FINISHER { + *self + .started_by_process + .lock() + .expect("started") + .entry(attempt.owner) + .or_insert(0) += 1; + } + self.attempts + .lock() + .expect("attempts") + .entry(job_id) + .or_default() + .push(attempt); + } + + pub fn note_generation_killed(&self, generation: u32, at_epoch_ms: i64) { + self.killed_at + .lock() + .expect("killed at") + .insert(generation, at_epoch_ms); + } + + /// Signal and publish the death before a waiting probe can classify its result. + /// The live-state lock also fences the probe's post-query check; no await occurs + /// between a successful signal and publishing its generation and timestamp. + pub fn signal_generation_killed( + &self, + generation: u32, + signal: impl FnOnce() -> std::io::Result<()>, + ) -> std::io::Result { + let mut live = self.live.lock().expect("live"); + signal()?; + let signalled_at = epoch_ms(); + self.note_generation_killed(generation, signalled_at); + if live.is_some_and(|(current, _)| current == generation) { + *live = None; + } + Ok(signalled_at) + } + + pub fn set_live(&self, live: Option<(u32, i64)>) { + *self.heartbeat_last.lock().expect("heartbeat last") = None; + *self.live.lock().expect("live") = live; + } + + pub fn live(&self) -> Option<(u32, i64)> { + *self.live.lock().expect("live") + } + + pub fn live_lost(&self, generation: u32) -> i64 { + self.live_lost + .lock() + .expect("live lost") + .get(&generation) + .copied() + .unwrap_or(0) + } + + pub fn record_kill(&self, kill: Kill) { + self.kills.lock().expect("kills").push(kill); + } + + pub fn kills(&self) -> Vec { + self.kills.lock().expect("kills").clone() + } + + pub fn started_by(&self, generation: u32) -> i64 { + self.started_by_process + .lock() + .expect("started") + .get(&generation) + .copied() + .unwrap_or(0) + } + + /// Whether the ring has been still for long enough to count as stalled. + pub fn ring_stalled(&self) -> bool { + self.ring.lock().expect("ring").episode.is_some() + } + + pub fn current_ring_episode(&self) -> Option { + self.ring.lock().expect("ring").episode.clone() + } + + pub fn episodes(&self) -> Vec { + let ring = self.ring.lock().expect("ring"); + let mut all = ring.episodes.clone(); + all.extend(ring.episode.clone()); + all + } + + /// Records one probe and advances the per-kill clocks. + pub fn observe( + &self, + shared: &Shared, + probe: &Probe, + live: Option, + ring_generation: i64, + probe_ms: i64, + ) { + let now = shared.now_ms(); + self.open_claims.store(probe.open, Ordering::Relaxed); + let sample = Sample { + at_ms: now, + open: probe.open, + expired: probe.expired, + oldest_expired_s: probe.oldest_expired_s, + ring_generation, + probe_ms, + }; + self.samples.lock().expect("samples").push(sample); + if let Some(path) = self.csv.lock().expect("csv").as_ref() { + if let Ok(mut file) = std::fs::OpenOptions::new().append(true).open(path) { + let _ = writeln!( + file, + "{:.1},{},{},{:.1},{},{}", + now as f64 / 1000.0, + probe.open, + probe.expired, + probe.oldest_expired_s, + ring_generation, + probe_ms + ); + } + } + let stats = shared.stats(Persona::ClaimAbandon); + stats.set_max("peak_open_claims", probe.open); + stats.set_max("peak_expired_claims", probe.expired); + stats.set_max( + "peak_oldest_expired_s", + probe.oldest_expired_s.round() as i64, + ); + stats.observe("probe_ms", probe_ms); + + // A probe can begin while the child is alive and finish after its kill. Only + // a generation still alive after the query proves missing claims were taken + // from a live owner. Keep this check serialized with signaling the kill. + if let Some(live) = live { + let lost = (live.started - probe.open_since).max(0); + let current = self.live.lock().expect("live"); + let still_live = current.is_some_and(|(generation, _)| generation == live.generation) + && !self + .killed_at + .lock() + .expect("killed at") + .contains_key(&live.generation); + if lost > 0 && !still_live { + shared.incomplete( + format!( + "claim probe for process {} crossed its kill or replacement; {lost} missing claim(s) cannot be classified as a live-owner rescue", + live.generation + ), + vec![format!( + "{{\"process\":{},\"started\":{},\"claims_still_open\":{}}}", + live.generation, live.started, probe.open_since + )], + ); + } else { + let mut peaks = self.live_lost.lock().expect("live lost"); + let peak = peaks.entry(live.generation).or_insert(0); + let first = lost > 0 && *peak == 0; + *peak = (*peak).max(lost); + drop(peaks); + stats.set_max("live_claims_rescued_peak", lost); + if first { + shared.violation( + "live_owner_rescued", + format!( + "claims of process {} were rescued while it was alive: it had started {} handler(s) and only {} of its claims were still open ({} claim(s) were eligible for rescue)", + live.generation, live.started, probe.open_since, probe.expired + ), + vec![format!( + "{{\"process\":{},\"started\":{},\"claims_still_open\":{},\"eligible_for_rescue\":{}}}", + live.generation, live.started, probe.open_since, probe.expired + )], + ); + } + } + } + let staleness_ms = shared.config.heartbeat_staleness.as_millis() as i64 + 3_000; + let epoch = epoch_ms(); + let bound = shared.config.abandon_bound; + let mut slow: Vec = Vec::new(); + { + let mut kills = self.kills.lock().expect("kills"); + for kill in kills + .iter_mut() + .filter(|kill| kill.cleared_after_s.is_none()) + { + kill.expired_peak = kill.expired_peak.max(probe.expired); + kill.oldest_expired_peak_s = kill.oldest_expired_peak_s.max(probe.oldest_expired_s); + let since_kill_ms = epoch - kill.signalled_epoch_ms; + // Claims of the dead process only count as stale after the staleness + // window, so a probe sooner than that proves nothing. + if since_kill_ms >= staleness_ms && probe.expired == 0 { + kill.cleared_after_s = Some(since_kill_ms as f64 / 1000.0); + continue; + } + let mut allowed = bound; + if shared.chaos_active() { + allowed += shared.config.chaos_slack; + } + if !kill.reported_slow && since_kill_ms as u64 > allowed.as_millis() as u64 { + kill.reported_slow = true; + slow.push(format!( + "{{\"kill\":{},\"seconds_since_kill\":{:.0},\"open_claims\":{},\"expired_claims\":{},\"oldest_expired_s\":{:.0},\"bound_s\":{}}}", + kill.number, + since_kill_ms as f64 / 1000.0, + probe.open, + probe.expired, + probe.oldest_expired_s, + bound.as_secs() + )); + } + } + } + for row in slow { + shared.violation( + "rescue_backlog_slow", + format!( + "the rescue backlog left by a killed process was not cleared within {}s of the kill ({} expired claim(s) remain, the oldest eligible for {:.0}s)", + bound.as_secs(), + probe.expired, + probe.oldest_expired_s + ), + vec![row], + ); + } + } + + /// Watches the claim ring; called about once a second. + /// + /// `next_slot` is the number of live and expired open claims in the slot the ring + /// enters next, known only while the ring is stalled. + pub fn observe_ring( + &self, + shared: &Shared, + current_slot: i32, + generation: i64, + next_slot: Option<(i64, i64)>, + ) { + let now = shared.now_ms(); + let mut report: Option = None; + { + let mut ring = self.ring.lock().expect("ring"); + let moved = ring.last != Some((current_slot, generation)); + if moved { + ring.last = Some((current_slot, generation)); + ring.last_change_ms = now; + if let Some(mut episode) = ring.episode.take() { + episode.end_ms = Some(now); + ring.episodes.push(episode); + } + return; + } + let mut threshold = crate::ci_timing::scaled_timeout(std::time::Duration::from_millis( + RING_STALL_AFTER_MS as u64, + )) + .as_millis() as i64; + if shared.chaos_active() { + threshold += shared.config.chaos_slack.as_millis() as i64; + } + if now - ring.last_change_ms >= threshold && ring.episode.is_none() { + ring.episode = Some(RingEpisode { + start_ms: ring.last_change_ms, + end_ms: None, + generation, + slot: current_slot, + last_live_ms: now, + expired_in_next_slot: 0, + reported: false, + }); + } + let bound_ms = shared.config.abandon_bound.as_millis() as i64; + if let Some(episode) = ring.episode.as_mut() { + if let Some((live, expired)) = next_slot { + if live > 0 { + episode.last_live_ms = now; + } + episode.expired_in_next_slot = expired; + } + // A stall is expected while a live process holds the next slot; it is a + // violation once nothing alive holds it and the ring still does not move. + if !episode.reported && now - episode.last_live_ms > bound_ms { + episode.reported = true; + report = Some(episode.clone()); + } + } + } + if let Some(episode) = report { + shared.violation( + "claim_ring_stalled", + format!( + "the claim ring has not rotated for more than {}s with no live claim in the slot it has to enter next ({} expired claim(s) there); stuck at slot {}, generation {}", + shared.config.abandon_bound.as_secs(), + episode.expired_in_next_slot, + episode.slot, + episode.generation + ), + vec![format!( + "{{\"stalled_since_s\":{:.1},\"slot\":{},\"generation\":{},\"expired_in_next_slot\":{}}}", + episode.start_ms as f64 / 1000.0, + episode.slot, + episode.generation, + episode.expired_in_next_slot + )], + ); + } + } + + /// Closes the ring episode left open when the run ends. + pub fn finish_ring(&self, shared: &Shared) { + let now = shared.now_ms(); + let mut ring = self.ring.lock().expect("ring"); + if let Some(mut episode) = ring.episode.take() { + episode.end_ms = Some(now); + ring.episodes.push(episode); + } + } + + /// Stall time and resume time of the ring in each kill's window. + fn ring_by_kill(&self) -> Vec<(f64, Option)> { + let kills = self.kills(); + let episodes = self.episodes(); + kills + .iter() + .enumerate() + .map(|(index, kill)| { + let from = kill.at_run_ms; + let until = kills.get(index + 1).map_or(i64::MAX, |next| next.at_run_ms); + let mut stalled = 0i64; + let mut resumed: Option = None; + for episode in &episodes { + let end = episode.end_ms.unwrap_or(i64::MAX); + if end <= from || episode.start_ms >= until { + continue; + } + stalled += end.min(until) - episode.start_ms.max(from); + if let Some(end) = episode.end_ms { + resumed = Some(resumed.map_or(end, |seen| seen.max(end))); + } + } + ( + stalled as f64 / 1000.0, + resumed.map(|end| (end - from) as f64 / 1000.0), + ) + }) + .collect() + } + + /// The per-kill table for the report. + pub fn render(&self) -> String { + let kills = self.kills(); + if kills.is_empty() { + return String::new(); + } + let ring = self.ring_by_kill(); + let samples = self.samples.lock().expect("samples"); + let mut out = String::new(); + out.push_str("### Kills\n\n"); + out.push_str("| kill | held by the process | rescued while alive | open claims after kill | peak expired | oldest expired (s) | time to clear (s) | ring stalled (s) | ring resumed after (s) |\n"); + out.push_str("| -- | -- | -- | -- | -- | -- | -- | -- | -- |\n"); + for (kill, (stalled, resumed)) in kills.iter().zip(ring) { + out.push_str(&format!( + "| {}{} | {} | {} | {} | {} | {:.0} | {} | {:.0} | {} |\n", + kill.number, + if kill.load_ended { " (load end)" } else { "" }, + kill.started_by_process, + kill.rescued_while_alive, + kill.open_after_kill, + kill.expired_peak, + kill.oldest_expired_peak_s, + kill.cleared_after_s + .map_or("not cleared".to_string(), |seconds| format!("{seconds:.0}")), + stalled, + resumed.map_or("-".to_string(), |seconds| format!("{seconds:.0}")), + )); + } + let ring_episodes = self.episodes(); + out.push_str(&format!( + "\nSamples: {}. Ring stall episodes: {}.\n", + samples.len(), + ring_episodes.len() + )); + out + } + + pub fn to_json(&self) -> serde_json::Value { + serde_json::json!({ + "kills": self.kills(), + "ring_episodes": self.episodes(), + "samples": self.samples.lock().expect("samples").len(), + }) + } + + /// Final checks over everything the child processes and the survivors reported. + pub async fn analyse(&self, shared: &Shared, pool: &PgPool, queues: &[String]) { + let killed_at = self.killed_at.lock().expect("killed at").clone(); + let max_attempts = self.max_attempts.lock().expect("max attempts").clone(); + let attempts = self.attempts.lock().expect("attempts").clone(); + let stats = shared.stats(Persona::ClaimAbandon); + + let mut live_rescued: Vec = Vec::new(); + let mut live_rescued_count = 0usize; + let mut duplicates_from_dead = 0i64; + let mut over_max: Vec = Vec::new(); + let mut gaps: Vec = Vec::new(); + let mut exhausted_by_attempts: Vec = Vec::new(); + for (job_id, list) in &attempts { + let mut list = list.clone(); + list.sort_by_key(|attempt| attempt.at_ms); + for pair in list.windows(2) { + let (previous, next) = (pair[0], pair[1]); + let previous_died_at = if previous.owner == FINISHER { + None + } else { + killed_at.get(&previous.owner).copied() + }; + match previous_died_at { + Some(died_at) if next.at_ms >= died_at => duplicates_from_dead += 1, + _ => { + live_rescued_count += 1; + if live_rescued.len() < MAX_ROWS { + live_rescued.push(format!( + "{{\"job_id\":{job_id},\"attempt\":{},\"owner\":\"{}\",\"owner_alive_until_ms\":{},\"next_attempt\":{},\"next_owner\":\"{}\",\"next_started_ms\":{}}}", + previous.attempt, + owner_name(previous.owner), + previous_died_at.map_or("null".to_string(), |at| at.to_string()), + next.attempt, + owner_name(next.owner), + next.at_ms + )); + } + } + } + if next.attempt != previous.attempt + 1 && gaps.len() < MAX_ROWS { + gaps.push(format!( + "{{\"job_id\":{job_id},\"attempt\":{},\"next_attempt\":{}}}", + previous.attempt, next.attempt + )); + } + } + let highest = list + .iter() + .map(|attempt| attempt.attempt) + .max() + .unwrap_or(0); + if let Some(max) = max_attempts.get(job_id) { + if highest > *max && over_max.len() < MAX_ROWS { + over_max.push(format!( + "{{\"job_id\":{job_id},\"attempts_seen\":{highest},\"max_attempts\":{max}}}" + )); + } + let last = list.last().copied(); + if highest == *max + && last.is_some_and(|last| last.owner != FINISHER && last.attempt == *max) + { + exhausted_by_attempts.push(*job_id); + } + } + } + stats.add("duplicate_executions_after_kill", duplicates_from_dead); + stats.add("live_owner_rescues", live_rescued_count as i64); + if live_rescued_count > 0 { + shared.violation( + "live_owner_rescued", + format!("{live_rescued_count} job(s) were run again while the process that held them was still alive"), + live_rescued, + ); + } + if !over_max.is_empty() { + shared.violation( + "attempts_exceeded_max", + format!( + "{} job(s) ran more attempts than max_attempts", + over_max.len() + ), + over_max, + ); + } + if !gaps.is_empty() { + shared.finding( + "abandon_attempt_gap", + format!( + "{} job(s) skipped an attempt number between executions", + gaps.len() + ), + gaps, + ); + } + + // Jobs whose last attempt was abandoned at max_attempts have to be dead-lettered. + let dlq: Result, _> = sqlx::query_as(&format!( + "SELECT job_id, attempt, max_attempts FROM {}.dlq_entries WHERE queue = ANY($1)", + schema(shared) + )) + .bind(queues) + .fetch_all(pool) + .await; + match dlq { + Ok(rows) => { + stats.add("dead_lettered", rows.len() as i64); + let early: Vec = rows + .iter() + .filter(|(_, attempt, max)| attempt < max) + .take(MAX_ROWS) + .map(|(job, attempt, max)| { + format!("{{\"job_id\":{job},\"attempt\":{attempt},\"max_attempts\":{max}}}") + }) + .collect(); + if !early.is_empty() { + shared.violation( + "dead_lettered_early", + format!( + "{} job(s) were dead-lettered before reaching max_attempts", + early.len() + ), + early, + ); + } + let in_dlq: std::collections::HashSet = + rows.iter().map(|(job, _, _)| *job).collect(); + let missing: Vec = exhausted_by_attempts + .iter() + .filter(|job| !in_dlq.contains(job)) + .take(MAX_ROWS) + .map(|job| { + format!( + "{{\"job_id\":{job},\"max_attempts\":{}}}", + max_attempts.get(job).copied().unwrap_or(0) + ) + }) + .collect(); + stats.add("exhausted_after_kill", exhausted_by_attempts.len() as i64); + if !missing.is_empty() && !shared.quiesced.load(Ordering::Relaxed) { + // The rescue was still working when the run ended. + shared.finding( + "abandon_rescue_unfinished", + format!( + "{} job(s) abandoned on their last attempt were not yet dead-lettered when the drain grace ended", + missing.len() + ), + missing, + ); + } else if !missing.is_empty() { + shared.violation( + "exhausted_not_dead_lettered", + format!( + "{} job(s) abandoned on their last attempt are not in the dead-letter queue", + missing.len() + ), + missing, + ); + } + } + Err(error) => shared.incomplete(format!("dead-letter check failed: {error}"), vec![]), + } + } +} + +fn owner_name(owner: u32) -> String { + if owner == FINISHER { + "survivor".to_string() + } else { + format!("process-{owner}") + } +} + +fn schema(shared: &Shared) -> String { + shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".to_string()) +} + +/// Which claims a backlog probe looks at. +pub enum Scope<'a> { + /// Open receipt claims and running leases of these queues. + Queues(&'a [String]), + /// Open receipt identities in one claim-ring slot, including attempts + /// materialized into running or callback-waiting leases. + Slot(i32), +} + +/// Open claims in `scope`, and the ones the rescue is already entitled to take: claims +/// whose last heartbeat (or claim time, when there never was one) is older than +/// `staleness_secs`, or whose deadline has passed. +/// +/// `since_epoch_ms` splits out the open claims taken at or after that time. +pub async fn probe_backlog( + pool: &PgPool, + schema: &str, + scope: Scope<'_>, + staleness_secs: f64, + since_epoch_ms: i64, +) -> Result { + let (batch_filter, row_filter, receipt_lease_filter, lease_leg) = match &scope { + Scope::Queues(_) => ( + "b.queue = ANY($1)", + "c.queue = ANY($1)", + "l.job_id IS NULL", + format!( + "UNION ALL + SELECT COALESCE(a.heartbeat_at, l.heartbeat_at, l.attempted_at), l.deadline_at, + COALESCE(l.attempted_at, l.heartbeat_at), NULL::timestamptz + FROM {schema}.leases l + LEFT JOIN {schema}.attempt_state a ON a.job_id = l.job_id AND a.run_lease = l.run_lease + WHERE l.queue = ANY($1) AND l.state = 'running'" + ), + ), + Scope::Slot(_) => ( + "b.claim_slot = $4", + "c.claim_slot = $4", + "l.job_id IS NULL OR l.state IN ('running', 'waiting_external')", + String::new(), + ), + }; + let sql = format!( + r#" + WITH bounds AS ( + SELECT clock_timestamp() AS now_ts, + clock_timestamp() - make_interval(secs => $2) AS stale_before + ), + batch_claims AS ( + SELECT b.claim_slot, i.receipt_id, i.job_id, i.run_lease, b.queue, b.claimed_at, b.deadline_at + FROM {schema}.lease_claim_batches b + CROSS JOIN LATERAL unnest(b.job_ids, b.run_leases, b.receipt_ids) AS i(job_id, run_lease, receipt_id) + WHERE {batch_filter} + ), + row_claims AS ( + SELECT c.claim_slot, c.receipt_id, c.job_id, c.run_lease, c.queue, c.claimed_at, c.deadline_at + FROM {schema}.lease_claims c + WHERE {row_filter} AND c.closed_at IS NULL + ), + open_receipts AS ( + SELECT x.* FROM (SELECT * FROM batch_claims UNION ALL SELECT * FROM row_claims) x + WHERE NOT EXISTS (SELECT 1 FROM {schema}.lease_claim_closures cl + WHERE cl.claim_slot = x.claim_slot AND cl.job_id = x.job_id AND cl.run_lease = x.run_lease) + AND NOT EXISTS (SELECT 1 FROM {schema}.lease_claim_closure_batches cb + WHERE cb.claim_slot = x.claim_slot AND cb.receipt_ranges @> x.receipt_id + -- Keep membership correlated: joining by slot can multiply + -- receipt members by retained closure batches. + OFFSET 0) + AND NOT EXISTS (SELECT 1 FROM {schema}.done_entries d WHERE d.job_id = x.job_id AND d.run_lease = x.run_lease) + AND NOT EXISTS (SELECT 1 FROM {schema}.deferred_jobs d WHERE d.job_id = x.job_id AND d.run_lease = x.run_lease) + AND NOT EXISTS (SELECT 1 FROM {schema}.dlq_entries d WHERE d.job_id = x.job_id AND d.run_lease = x.run_lease) + ), + holders AS ( + SELECT CASE WHEN l.state = 'waiting_external' THEN NULL + ELSE COALESCE(a.heartbeat_at, l.heartbeat_at, l.attempted_at, x.claimed_at) + END AS beat_at, + CASE WHEN l.job_id IS NULL THEN x.deadline_at + WHEN l.state = 'running' THEN l.deadline_at + END AS deadline_at, + x.claimed_at, + CASE WHEN l.state = 'waiting_external' THEN l.callback_timeout_at END + AS callback_timeout_at + FROM open_receipts x + LEFT JOIN {schema}.leases l ON l.job_id = x.job_id AND l.run_lease = x.run_lease + LEFT JOIN {schema}.attempt_state a ON a.job_id = x.job_id AND a.run_lease = x.run_lease + WHERE {receipt_lease_filter} + {lease_leg} + ), + eligible AS ( + SELECT LEAST( + CASE WHEN beat_at < stale_before THEN beat_at + make_interval(secs => $2) END, + CASE WHEN deadline_at < now_ts THEN deadline_at END, + CASE WHEN callback_timeout_at < now_ts THEN callback_timeout_at END + ) AS eligible_at, + claimed_at >= to_timestamp($3::float8 / 1000.0) AS since_start + FROM holders, bounds + ) + SELECT count(*)::bigint, + count(eligible_at)::bigint, + COALESCE(EXTRACT(EPOCH FROM (clock_timestamp() - min(eligible_at))), 0)::float8, + count(*) FILTER (WHERE since_start)::bigint + FROM eligible + "# + ); + let query = sqlx::query_as::<_, (i64, i64, f64, i64)>(&sql); + let query = match scope { + Scope::Queues(queues) => query + .bind(queues.to_vec()) + .bind(staleness_secs) + .bind(since_epoch_ms as f64) + .bind(0_i32), + Scope::Slot(slot) => query + .bind(Vec::::new()) + .bind(staleness_secs) + .bind(since_epoch_ms as f64) + .bind(slot), + }; + let (open, expired, oldest, open_since) = query.fetch_one(pool).await?; + Ok(Probe { + open, + expired, + oldest_expired_s: oldest.max(0.0), + open_since, + }) +} + +#[cfg(test)] +mod tests { + use super::{epoch_ms, probe_backlog, LiveCheck, Probe, Scope}; + use crate::persona_soak::config::{replace_database_name, SoakConfig}; + use crate::persona_soak::db::{self, Budget, ServerInfo}; + use crate::persona_soak::shared::Shared; + use awa::model::{ClaimedRuntimeJob, QueueStorage, QueueStorageConfig}; + use sqlx::PgPool; + use std::sync::atomic::Ordering; + use std::sync::Arc; + use std::time::Duration; + + fn live_probe_fixture(generation: u32) -> (Arc, LiveCheck) { + let config = SoakConfig::for_test( + "postgres://postgres:test@localhost/awa_soak_live_probe".to_string(), + ); + let info = ServerInfo::default(); + let budget = Budget::new(&info, 0); + let shared = Shared::new(config, info, budget); + shared.abandon.set_live(Some((generation, epoch_ms()))); + let live = LiveCheck { + generation: shared.abandon.live().expect("live generation").0, + started: 1_000, + }; + (shared, live) + } + + fn observe_live_probe(shared: &Shared, live: LiveCheck, open: i64) { + shared.abandon.observe( + shared, + &Probe { + open, + expired: 0, + oldest_expired_s: 0.0, + open_since: open, + }, + Some(live), + 0, + 1, + ); + } + + fn live_rescue_reports(shared: &Shared) -> usize { + shared + .violations + .lock() + .expect("violations") + .iter() + .filter(|violation| violation.kind == "live_owner_rescued") + .count() + } + + #[test] + fn live_probe_crossing_its_generation_kill_is_incomplete() { + let (shared, captured) = live_probe_fixture(1); + // The query was dispatched with this owner alive, but its snapshot + // boundary relative to the signal is unknown when the result arrives. + shared.abandon.note_generation_killed(1, epoch_ms()); + shared.abandon.set_live(None); + observe_live_probe(&shared, captured, 500); + + assert!(shared.oracle_incomplete.load(Ordering::Relaxed)); + assert_eq!(live_rescue_reports(&shared), 0); + assert_eq!(shared.abandon.live_lost(1), 0); + // The usable backlog evidence is retained even though liveness is unknown. + assert_eq!(shared.abandon.open_claims.load(Ordering::Relaxed), 500); + } + + #[test] + fn live_probe_still_alive_missing_claims_remain_a_violation() { + let (shared, captured) = live_probe_fixture(1); + observe_live_probe(&shared, captured, 500); + + assert_eq!(live_rescue_reports(&shared), 1); + assert_eq!(shared.abandon.live_lost(1), 500); + assert!(!shared.oracle_incomplete.load(Ordering::Relaxed)); + } + + #[test] + fn live_probe_other_generation_kill_does_not_hide_live_loss() { + let (shared, captured) = live_probe_fixture(2); + shared.abandon.note_generation_killed(1, epoch_ms()); + observe_live_probe(&shared, captured, 500); + + assert_eq!(live_rescue_reports(&shared), 1); + assert_eq!(shared.abandon.live_lost(2), 500); + assert!(!shared.oracle_incomplete.load(Ordering::Relaxed)); + } + + #[test] + fn live_probe_kill_overlap_preserves_prior_proven_loss() { + let (shared, captured) = live_probe_fixture(1); + observe_live_probe(&shared, captured, 500); + assert_eq!(shared.abandon.live_lost(1), 500); + + shared.abandon.note_generation_killed(1, epoch_ms()); + shared.abandon.set_live(None); + observe_live_probe(&shared, captured, 100); + + assert_eq!(live_rescue_reports(&shared), 1); + assert_eq!(shared.abandon.live_lost(1), 500); + assert!(shared.oracle_incomplete.load(Ordering::Relaxed)); + } + + #[test] + fn live_probe_successful_signal_publishes_death_before_returning() { + let (shared, _) = live_probe_fixture(1); + let at = shared + .abandon + .signal_generation_killed(1, || Ok(())) + .expect("successful signal"); + + assert_eq!(shared.abandon.live(), None); + assert_eq!( + shared.abandon.killed_at.lock().expect("killed at").get(&1), + Some(&at) + ); + } + + #[test] + fn live_probe_failed_signal_does_not_publish_a_death() { + let (shared, _) = live_probe_fixture(1); + let before = shared.abandon.live(); + let result = shared + .abandon + .signal_generation_killed(1, || Err(std::io::Error::other("signal failed"))); + + assert!(result.is_err()); + assert_eq!(shared.abandon.live(), before); + assert!(shared + .abandon + .killed_at + .lock() + .expect("killed at") + .is_empty()); + } + + struct SlotProbe { + config: SoakConfig, + pool: PgPool, + store: QueueStorage, + } + + impl SlotProbe { + async fn create() -> Self { + let base = std::env::var("DATABASE_URL").unwrap_or_else(|_| { + "postgres://postgres:test@localhost:15432/awa_test".to_string() + }); + let name = format!("awa_soak_slot_probe_{}", uuid::Uuid::new_v4().simple()); + let config = SoakConfig::for_test(replace_database_name(&base, &name)); + db::recreate_database(&config) + .await + .expect("isolated database"); + let pool = db::connect_pool(&config.database_url, "slot_probe_regression", 4) + .await + .expect("connect"); + let store = QueueStorage::new(QueueStorageConfig::default()).expect("store"); + Self { + config, + pool, + store, + } + } + + async fn claim(&self, deadline: Duration) -> ClaimedRuntimeJob { + self.store + .enqueue_batch(&self.pool, "slot_probe", 2, 1) + .await + .expect("enqueue"); + self.store + .claim_runtime_batch(&self.pool, "slot_probe", 1, deadline) + .await + .expect("claim") + .pop() + .expect("one claim") + } + + async fn assert_counts(&self, slot: i32, open: i64, expired: i64) { + let probe = probe_backlog(&self.pool, "awa", Scope::Slot(slot), 30.0, 0) + .await + .expect("slot probe"); + assert_eq!((probe.open, probe.expired), (open, expired)); + } + + async fn age_heartbeat(&self, claim: &ClaimedRuntimeJob) { + sqlx::query( + "UPDATE awa.leases + SET heartbeat_at = clock_timestamp() - interval '1 minute' + WHERE job_id = $1 AND run_lease = $2", + ) + .bind(claim.job.id) + .bind(claim.job.run_lease) + .execute(&self.pool) + .await + .expect("age lease heartbeat"); + sqlx::query( + "INSERT INTO awa.attempt_state (job_id, run_lease, heartbeat_at, updated_at) + VALUES ($1, $2, clock_timestamp() - interval '1 minute', clock_timestamp()) + ON CONFLICT (job_id, run_lease) DO UPDATE + SET heartbeat_at = EXCLUDED.heartbeat_at", + ) + .bind(claim.job.id) + .bind(claim.job.run_lease) + .execute(&self.pool) + .await + .expect("age attempt heartbeat"); + } + + async fn drop(self) { + self.pool.close().await; + let admin = db::connect_pool(&self.config.admin_url(), "slot_probe_cleanup", 1) + .await + .expect("cleanup connection"); + sqlx::query(&format!( + "DROP DATABASE {} WITH (FORCE)", + self.config.database_name() + )) + .execute(&admin) + .await + .expect("drop isolated database"); + admin.close().await; + } + } + + #[tokio::test] + async fn slot_probe_keeps_materialized_running_claims_until_durable_completion() { + let run = SlotProbe::create().await; + // Zero-deadline receipts are compact; deadline receipts are row-local. + for deadline in [Duration::ZERO, Duration::from_secs(3600)] { + let claim = run.claim(deadline).await; + run.store + .register_callback( + &run.pool, + claim.job.id, + claim.job.run_lease, + Duration::from_secs(3600), + ) + .await + .expect("materialize the running receipt"); + let slot = claim.claim.claim_slot; + run.assert_counts(slot, 1, 0).await; + run.assert_counts((slot + 1) % 8, 0, 0).await; + run.age_heartbeat(&claim).await; + run.assert_counts(slot, 1, 1).await; + run.store + .heartbeat_batch(&run.pool, &[(claim.job.id, claim.job.run_lease)]) + .await + .expect("refresh authoritative attempt heartbeat"); + // The stale leases.heartbeat_at must not override attempt_state. + run.assert_counts(slot, 1, 0).await; + run.store + .complete_runtime_batch(&run.pool, std::slice::from_ref(&claim)) + .await + .expect("durably complete"); + run.assert_counts(slot, 0, 0).await; + } + run.drop().await; + } + + #[tokio::test] + async fn slot_probe_uses_callback_timeout_for_waiting_materialized_claims() { + let run = SlotProbe::create().await; + for deadline in [Duration::ZERO, Duration::from_secs(3600)] { + let claim = run.claim(deadline).await; + let callback = run + .store + .register_callback( + &run.pool, + claim.job.id, + claim.job.run_lease, + Duration::from_secs(3600), + ) + .await + .expect("materialize the receipt"); + assert!(run + .store + .enter_callback_wait(&run.pool, claim.job.id, claim.job.run_lease, callback) + .await + .expect("wait for callback")); + run.age_heartbeat(&claim).await; + sqlx::query( + "UPDATE awa.leases SET deadline_at = clock_timestamp() - interval '1 second' + WHERE job_id = $1 AND run_lease = $2", + ) + .bind(claim.job.id) + .bind(claim.job.run_lease) + .execute(&run.pool) + .await + .expect("expire the original job deadline"); + let slot = claim.claim.claim_slot; + // Waiting ignores the old heartbeat and deadline until its callback times out. + run.assert_counts(slot, 1, 0).await; + run.assert_counts((slot + 1) % 8, 0, 0).await; + sqlx::query( + "UPDATE awa.leases SET callback_timeout_at = clock_timestamp() - interval '1 second' + WHERE job_id = $1 AND run_lease = $2", + ) + .bind(claim.job.id) + .bind(claim.job.run_lease) + .execute(&run.pool) + .await + .expect("expire callback"); + run.assert_counts(slot, 1, 1).await; + run.store + .complete_external(&run.pool, callback, None, Some(claim.job.run_lease), false) + .await + .expect("durably complete callback"); + run.assert_counts(slot, 0, 0).await; + } + run.drop().await; + } +} diff --git a/awa/tests/persona_soak/adminserver.rs b/awa/tests/persona_soak/adminserver.rs new file mode 100644 index 000000000..79e6b0946 --- /dev/null +++ b/awa/tests/persona_soak/adminserver.rs @@ -0,0 +1,132 @@ +//! The real admin HTTP server (`awa serve`), run as a child process against the soak +//! database so the UI personas exercise the same code the web UI talks to. + +use super::config::SoakConfig; +use std::net::TcpListener; +use std::path::PathBuf; +use std::process::{Child, Command, Stdio}; +use std::time::{Duration, Instant}; + +pub const APPLICATION_NAME: &str = "soak_ui_server"; + +pub struct AdminServer { + child: Child, + pub base_url: String, + pub statement_timeout: Duration, + pub pool_max: u32, +} + +/// Looks for the `awa` binary next to the test executable, where +/// `cargo build --release -p awa-cli` puts it, unless `AWA_SOAK_AWA_BIN` names one. +pub fn locate_binary(config: &SoakConfig) -> Result { + if let Some(path) = &config.awa_bin { + return path + .exists() + .then(|| path.clone()) + .ok_or_else(|| format!("AWA_SOAK_AWA_BIN {} does not exist", path.display())); + } + let exe = std::env::current_exe().map_err(|error| error.to_string())?; + let candidate = exe + .parent() + .and_then(|deps| deps.parent()) + .map(|profile| profile.join("awa")); + match candidate { + Some(path) if path.exists() => Ok(path), + _ => Err( + "the admin UI personas need the awa binary: run `cargo build --release -p awa-cli` \ + or set AWA_SOAK_AWA_BIN" + .to_string(), + ), + } +} + +impl AdminServer { + pub async fn start(config: &SoakConfig) -> Result { + let binary = locate_binary(config)?; + let port = TcpListener::bind("127.0.0.1:0") + .and_then(|listener| listener.local_addr()) + .map_err(|error| format!("pick a port: {error}"))? + .port(); + let separator = if config.database_url.contains('?') { + '&' + } else { + '?' + }; + let url = format!( + "{}{separator}application_name={APPLICATION_NAME}", + config.database_url + ); + let statement_timeout = config.ui_statement_timeout; + let log = std::fs::File::create(config.report_dir.join("awa-serve.log")) + .or_else(|_| { + std::fs::create_dir_all(&config.report_dir)?; + std::fs::File::create(config.report_dir.join("awa-serve.log")) + }) + .map_err(|error| format!("create awa-serve.log: {error}"))?; + let child = Command::new(&binary) + .args([ + "--database-url", + &url, + "serve", + "--host", + "127.0.0.1", + "--port", + &port.to_string(), + ]) + .args([ + "--statement-timeout", + &statement_timeout.as_secs().to_string(), + ]) + .args([ + "--pool-max", + &config.ui_pool_max.to_string(), + "--pool-min", + "2", + ]) + .stdout(Stdio::from( + log.try_clone().map_err(|error| error.to_string())?, + )) + .stderr(Stdio::from(log)) + .spawn() + .map_err(|error| format!("spawn {}: {error}", binary.display()))?; + let mut server = Self { + child, + base_url: format!("http://127.0.0.1:{port}"), + statement_timeout, + pool_max: config.ui_pool_max, + }; + let client = reqwest::Client::new(); + let deadline = Instant::now() + Duration::from_secs(60); + loop { + if let Ok(Some(status)) = server.child.try_wait() { + return Err(format!("awa serve exited during startup: {status}")); + } + if let Ok(response) = client + .get(format!("{}/api/capabilities", server.base_url)) + .send() + .await + { + if response.status().is_success() { + return Ok(server); + } + } + if Instant::now() >= deadline { + return Err("awa serve did not start listening within 60s".to_string()); + } + tokio::time::sleep(Duration::from_millis(300)).await; + } + } + + pub fn is_running(&mut self) -> bool { + matches!(self.child.try_wait(), Ok(None)) + } + + pub fn stop(self) {} +} + +impl Drop for AdminServer { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} diff --git a/awa/tests/persona_soak/config.rs b/awa/tests/persona_soak/config.rs new file mode 100644 index 000000000..aaac98c38 --- /dev/null +++ b/awa/tests/persona_soak/config.rs @@ -0,0 +1,327 @@ +//! Run configuration, read from the environment. + +use crate::ci_timing::{scaled_staleness, scaled_timeout}; +use std::path::PathBuf; +use std::str::FromStr; +use std::time::Duration; + +/// The harness's fast heartbeat cadence, so that rescues happen many times in a short run. +pub const DEFAULT_HEARTBEAT_SECONDS: u64 = 2; +pub const DEFAULT_STALENESS_SECONDS: u64 = 12; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub enum Persona { + Mia, + Tom, + Priya, + Sam, + Alex, + Chaos, + /// Admin UI users, through the real HTTP server. + Watchers, + Support, + UiOps, + Impatient, + /// A separate Python process: producers, async workers and the migrate path. + Python, + /// Handlers that batch-insert unique child jobs into their own queue. + Fanout, + /// An operator changing `enqueue_shards` on hot queues under live load, with the + /// actors running as different non-superuser roles. + ShardOps, + /// A separate process that claims as much as it can and is SIGKILLed again and again. + ClaimAbandon, +} + +impl Persona { + pub const ALL: [Persona; 14] = [ + Persona::Mia, + Persona::Tom, + Persona::Priya, + Persona::Sam, + Persona::Alex, + Persona::Chaos, + Persona::Watchers, + Persona::Support, + Persona::UiOps, + Persona::Impatient, + Persona::Python, + Persona::Fanout, + Persona::ShardOps, + Persona::ClaimAbandon, + ]; + + pub fn name(self) -> &'static str { + match self { + Persona::Mia => "mia", + Persona::Tom => "tom", + Persona::Priya => "priya", + Persona::Sam => "sam", + Persona::Alex => "alex", + Persona::Chaos => "chaos", + Persona::Watchers => "watchers", + Persona::Support => "support", + Persona::UiOps => "uiops", + Persona::Impatient => "impatient", + Persona::Python => "python", + Persona::Fanout => "fanout", + Persona::ShardOps => "shardops", + Persona::ClaimAbandon => "abandon", + } + } + + pub fn index(self) -> u8 { + self as u8 + } + + pub fn from_index(index: u8) -> Persona { + Persona::ALL[(index as usize).min(Persona::ALL.len() - 1)] + } +} + +impl FromStr for Persona { + type Err = String; + + fn from_str(value: &str) -> Result { + Persona::ALL + .into_iter() + .find(|persona| persona.name() == value.trim().to_ascii_lowercase()) + .ok_or_else(|| { + format!("unknown persona {value:?}; expected one of mia,tom,priya,sam,alex,chaos,watchers,support,uiops,impatient,python,fanout,shardops,abandon") + }) + } +} + +#[derive(Debug, Clone)] +pub struct SoakConfig { + /// URL of the disposable soak database. Its database name must start with `awa_soak`. + pub database_url: String, + pub personas: Vec, + /// Length of the load phase. + pub load: Duration, + /// How long the oracle waits for the system to quiesce after load stops. + pub drain_grace: Duration, + pub seed: u64, + /// Multiplies each persona's rates and queue counts. + pub scale: f64, + pub report_dir: PathBuf, + /// Postgres server log file readable from this host. + pub pg_log_file: Option, + /// Docker container running the server, used to read its log. + pub pg_container: Option, + /// Log file path inside `pg_container`; when absent `docker logs` is used. + pub pg_container_log: Option, + /// Seconds a scheduled job may be overdue before the promotion check fails. + pub promotion_overdue: Duration, + /// Seconds a queue may hold available jobs with idle workers before the stall check fails. + pub stall_window: Duration, + /// The `awa` binary that serves the admin UI; located next to the test binary when unset. + pub awa_bin: Option, + /// Python interpreter of a virtualenv with the `awa` extension installed. + pub python_bin: Option, + /// Server-side statement timeout of the admin server, and the matching client timeout. + pub ui_statement_timeout: Duration, + pub ui_client_timeout: Duration, + /// Connections of the admin server's pool, and seconds between the operator's admin reads. + pub ui_pool_max: u32, + pub admin_read_every: Duration, + /// The load phase stops early when the report directory's filesystem has less free space. + pub min_free_disk_gb: i64, + /// Extra tolerance added to the stall and promotion thresholds while chaos is active. + pub chaos_slack: Duration, + /// Multiplies the shard-ops persona's rates independently of `scale`, which is + /// usually lowered to keep the other personas light. + pub shardops_scale: f64, + /// Whether the application roles may create objects in the schema. Without it the + /// operator provisions every lane before the shard count is raised, as the + /// restricted-runtime procedure in `docs/security.md` requires. + pub shardops_schema_create: bool, + /// Claim-and-abandon persona: claims held by each killed process, seconds between + /// kills, how many kills, and the bound within which the rescue backlog has to clear + /// and the claim ring has to resume after a kill. Independent of `scale`. + pub abandon_claims: u64, + pub abandon_kill_every: Duration, + pub abandon_kills: u32, + pub abandon_bound: Duration, + /// Pause between a kill and the replacement process. + pub abandon_restart_delay: Duration, + /// Overrides the heartbeat and deadline rescue intervals of every runtime, to run the + /// soak with production's cadence instead of the harness's fast one. + pub rescue_interval: Option, + /// Heartbeat interval and the heartbeat staleness after which a claim may be rescued, + /// for every runtime. The library's own defaults are 30 s and 90 s. + pub heartbeat_interval: Duration, + pub heartbeat_staleness: Duration, +} + +fn env_string(name: &str) -> Option { + std::env::var(name) + .ok() + .filter(|value| !value.trim().is_empty()) +} + +fn env_parse(name: &str) -> Option { + env_string(name).and_then(|value| value.trim().parse().ok()) +} + +impl SoakConfig { + pub fn from_env() -> Result { + let database_url = env_string("AWA_SOAK_DATABASE_URL").ok_or( + "AWA_SOAK_DATABASE_URL is required (postgres://user:pass@host:port/awa_soak_)", + )?; + Self::from_database_url(database_url) + } + + #[cfg(test)] + pub fn for_test(database_url: String) -> Self { + let mut config = Self::from_database_url(database_url).expect("test database URL"); + config.personas.clear(); + config + } + + fn from_database_url(database_url: String) -> Result { + let database_name = database_name(&database_url); + if !database_name.starts_with("awa_soak") + || !database_name + .chars() + .all(|ch| ch.is_ascii_alphanumeric() || ch == '_') + { + return Err(format!( + "refusing to run: the soak drops and recreates its database, so the name must \ + match awa_soak[A-Za-z0-9_]*, got {database_name:?}" + )); + } + + let personas = match env_string("AWA_SOAK_PERSONAS") { + Some(list) => list + .split(',') + .map(str::parse) + .collect::, _>>()?, + None => Persona::ALL.to_vec(), + }; + let seed = env_parse("AWA_SOAK_SEED").unwrap_or_else(|| { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .map(|elapsed| elapsed.as_nanos() as u64) + .unwrap_or(1) + }); + let report_dir = env_string("AWA_SOAK_REPORT_DIR") + .map(PathBuf::from) + .unwrap_or_else(|| std::env::temp_dir().join(format!("awa-soak-{seed}"))); + + Ok(Self { + database_url, + personas, + load: Duration::from_secs(env_parse("AWA_SOAK_SECONDS").unwrap_or(900)), + drain_grace: scaled_timeout(Duration::from_secs( + env_parse("AWA_SOAK_DRAIN_GRACE_SECONDS").unwrap_or(120), + )), + seed, + scale: env_parse("AWA_SOAK_SCALE").unwrap_or(1.0), + report_dir, + pg_log_file: env_string("AWA_SOAK_PG_LOG").map(PathBuf::from), + pg_container: env_string("AWA_SOAK_PG_CONTAINER"), + pg_container_log: env_string("AWA_SOAK_PG_CONTAINER_LOG"), + awa_bin: env_string("AWA_SOAK_AWA_BIN").map(PathBuf::from), + python_bin: env_string("AWA_SOAK_PYTHON").map(PathBuf::from), + ui_statement_timeout: Duration::from_secs( + env_parse("AWA_SOAK_UI_STATEMENT_TIMEOUT_SECONDS").unwrap_or(30), + ), + ui_client_timeout: Duration::from_secs( + env_parse("AWA_SOAK_UI_CLIENT_TIMEOUT_SECONDS").unwrap_or(45), + ), + ui_pool_max: env_parse("AWA_SOAK_UI_POOL_MAX").unwrap_or(10), + admin_read_every: Duration::from_secs( + env_parse("AWA_SOAK_ADMIN_READ_SECONDS").unwrap_or(2), + ), + min_free_disk_gb: env_parse("AWA_SOAK_MIN_FREE_DISK_GB").unwrap_or(8), + promotion_overdue: scaled_timeout(Duration::from_secs( + env_parse("AWA_SOAK_PROMOTION_OVERDUE_SECONDS").unwrap_or(10), + )), + stall_window: scaled_timeout(Duration::from_secs( + env_parse("AWA_SOAK_STALL_SECONDS").unwrap_or(30), + )), + chaos_slack: scaled_timeout(Duration::from_secs( + env_parse("AWA_SOAK_CHAOS_SLACK_SECONDS").unwrap_or(45), + )), + shardops_scale: env_parse("AWA_SOAK_SHARDOPS_SCALE").unwrap_or(1.0), + shardops_schema_create: env_parse::("AWA_SOAK_SHARDOPS_SCHEMA_CREATE") + .map(|value| value != 0) + .unwrap_or(true), + abandon_claims: env_parse("AWA_SOAK_ABANDON_CLAIMS").unwrap_or(4_000), + abandon_kill_every: Duration::from_secs( + env_parse("AWA_SOAK_ABANDON_KILL_SECONDS").unwrap_or(180), + ), + abandon_kills: env_parse("AWA_SOAK_ABANDON_KILLS").unwrap_or(4), + abandon_bound: scaled_timeout(Duration::from_secs( + env_parse("AWA_SOAK_ABANDON_BOUND_SECONDS").unwrap_or(120), + )), + abandon_restart_delay: Duration::from_secs( + env_parse("AWA_SOAK_ABANDON_RESTART_SECONDS").unwrap_or(3), + ), + rescue_interval: env_parse::("AWA_SOAK_RESCUE_INTERVAL_SECONDS") + .map(Duration::from_secs), + heartbeat_interval: Duration::from_secs( + env_parse("AWA_SOAK_HEARTBEAT_SECONDS").unwrap_or(DEFAULT_HEARTBEAT_SECONDS), + ), + heartbeat_staleness: scaled_staleness(Duration::from_secs( + env_parse("AWA_SOAK_STALENESS_SECONDS").unwrap_or(DEFAULT_STALENESS_SECONDS), + )), + }) + } + + pub fn has(&self, persona: Persona) -> bool { + self.personas.contains(&persona) + } + + /// Scales a count, never below one. + pub fn scaled(&self, base: u64) -> u64 { + ((base as f64 * self.scale).round() as u64).max(1) + } + + pub fn database_name(&self) -> String { + database_name(&self.database_url) + } + + /// The same server, connected to the maintenance database. + pub fn admin_url(&self) -> String { + replace_database_name(&self.database_url, "postgres") + } +} + +pub fn database_name(url: &str) -> String { + let without_query = url.split_once('?').map_or(url, |(prefix, _)| prefix); + without_query + .rsplit_once('/') + .map(|(_, name)| name.to_string()) + .unwrap_or_default() +} + +pub fn replace_database_name(url: &str, database_name: &str) -> String { + let (without_query, query) = match url.split_once('?') { + Some((prefix, query)) => (prefix, Some(query)), + None => (url, None), + }; + let base = without_query + .rsplit_once('/') + .map_or(without_query, |(base, _)| base); + match query { + Some(query) => format!("{base}/{database_name}?{query}"), + None => format!("{base}/{database_name}"), + } +} + +/// `url` with its user and password replaced. +pub fn replace_credentials(url: &str, user: &str, password: &str) -> String { + let Some((scheme, rest)) = url.split_once("://") else { + return url.to_string(); + }; + let authority_end = rest.find('/').unwrap_or(rest.len()); + let host = rest[..authority_end] + .rsplit_once('@') + .map_or(&rest[..authority_end], |(_, host)| host); + format!( + "{scheme}://{user}:{password}@{host}{}", + &rest[authority_end..] + ) +} diff --git a/awa/tests/persona_soak/db.rs b/awa/tests/persona_soak/db.rs new file mode 100644 index 000000000..d195a6c1a --- /dev/null +++ b/awa/tests/persona_soak/db.rs @@ -0,0 +1,293 @@ +//! Database lifecycle, connection budgeting and the per-client runtime host. + +use super::config::{replace_database_name, Persona, SoakConfig}; +use super::roles; +use awa::Client; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; +use sqlx::PgPool; +use std::collections::BTreeMap; +use std::str::FromStr; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +/// Opens a pool whose connections carry `application_name`, so the oracle and the +/// chaos persona can tell the runtimes apart in `pg_stat_activity`. +pub async fn connect_pool( + url: &str, + application_name: &str, + max_connections: u32, +) -> Result { + connect_pool_with_timeout(url, application_name, max_connections, None).await +} + +/// As [`connect_pool`], with a server-side `statement_timeout` on every connection. +pub async fn connect_pool_with_timeout( + url: &str, + application_name: &str, + max_connections: u32, + statement_timeout: Option, +) -> Result { + let mut options = PgConnectOptions::from_str(url)?.application_name(application_name); + if let Some(timeout) = statement_timeout { + options = options.options([("statement_timeout", timeout.as_millis().to_string())]); + } + PgPoolOptions::new() + .max_connections(max_connections) + .acquire_timeout(Duration::from_secs(60)) + .connect_with(options) + .await +} + +/// Drops and recreates the soak database, applies migrations and activates queue +/// storage, the production engine. +pub async fn recreate_database(config: &SoakConfig) -> Result<(), String> { + let name = config.database_name(); + assert!(name.starts_with("awa_soak"), "database name guard"); + let admin = connect_pool(&config.admin_url(), "awa_soak_admin", 2) + .await + .map_err(|error| format!("connect to maintenance database: {error}"))?; + sqlx::query(&format!("DROP DATABASE IF EXISTS {name} WITH (FORCE)")) + .execute(&admin) + .await + .map_err(|error| format!("drop database: {error}"))?; + let restricted_roles = config.has(Persona::ShardOps); + if restricted_roles { + roles::create_roles(&admin).await?; + sqlx::query(&format!("CREATE DATABASE {name} OWNER {}", roles::OWNER)) + .execute(&admin) + .await + .map_err(|error| format!("create database: {error}"))?; + } else { + sqlx::query(&format!("CREATE DATABASE {name}")) + .execute(&admin) + .await + .map_err(|error| format!("create database: {error}"))?; + } + admin.close().await; + + // With the shard-ops persona the schema belongs to a non-superuser owner role, as in + // a production install, and the application roles get the documented grants. + let setup_url = if restricted_roles { + roles::url_for_role(config, roles::OWNER) + } else { + config.database_url.clone() + }; + let pool = connect_pool(&setup_url, "awa_soak_setup", 2) + .await + .map_err(|error| format!("connect to soak database: {error}"))?; + awa::model::migrations::run(&pool) + .await + .map_err(|error| format!("migrate: {error}"))?; + awa_testing::setup::activate_queue_storage(&pool).await; + pool.close().await; + if restricted_roles { + let superuser = connect_pool(&config.database_url, "awa_soak_setup_grants", 1) + .await + .map_err(|error| format!("connect to soak database: {error}"))?; + let mixed_personas = config + .personas + .iter() + .any(|persona| *persona != Persona::ShardOps); + roles::grant_runtime( + &superuser, + &name, + config.shardops_schema_create, + mixed_personas, + ) + .await?; + superuser.close().await; + } + Ok(()) +} + +#[derive(Debug, Clone, Default)] +pub struct ServerInfo { + pub version: String, + pub max_connections: u32, + pub superuser_reserved: u32, + pub connections_at_start: i64, + pub postmaster_start: String, + pub engine: String, + pub queue_storage_schema: Option, + pub awa_version: String, +} + +pub async fn server_info(pool: &PgPool) -> Result { + let version: String = sqlx::query_scalar("SHOW server_version") + .fetch_one(pool) + .await?; + let max_connections: String = sqlx::query_scalar("SHOW max_connections") + .fetch_one(pool) + .await?; + let reserved: String = sqlx::query_scalar("SHOW superuser_reserved_connections") + .fetch_one(pool) + .await?; + let connections: i64 = sqlx::query_scalar( + "SELECT count(*)::bigint FROM pg_stat_activity WHERE datname IS NOT NULL", + ) + .fetch_one(pool) + .await?; + let postmaster_start: String = sqlx::query_scalar("SELECT pg_postmaster_start_time()::text") + .fetch_one(pool) + .await?; + let schema: Option = sqlx::query_scalar("SELECT awa.active_queue_storage_schema()") + .fetch_one(pool) + .await?; + Ok(ServerInfo { + version, + max_connections: max_connections.parse().unwrap_or(100), + superuser_reserved: reserved.parse().unwrap_or(3), + connections_at_start: connections, + postmaster_start, + engine: if schema.is_some() { + "queue_storage" + } else { + "canonical" + } + .to_string(), + queue_storage_schema: schema, + awa_version: env!("CARGO_PKG_VERSION").to_string(), + }) +} + +/// Hands out connections to the personas so that all of them together stay inside +/// the server's `max_connections`. +pub struct Budget { + usable: u32, + allocations: Mutex>, +} + +impl Budget { + /// `headroom` connections are kept free for the oracle, operator probes and chaos + /// sessions on top of the server's superuser reservation. + pub fn new(info: &ServerInfo, headroom: u32) -> Self { + let usable = info + .max_connections + .saturating_sub(info.superuser_reserved) + .saturating_sub(headroom) + .saturating_sub(info.connections_at_start.max(0) as u32); + Self { + usable, + allocations: Mutex::new(BTreeMap::new()), + } + } + + /// 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. + 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 available = self.usable.saturating_sub(used); + assert!( + available >= min, + "connection budget exhausted: {label} needs at least {min}, {available} of {} left; \ + raise max_connections or lower AWA_SOAK_SCALE / persona count. Allocations: {allocations:?}", + self.usable + ); + let granted = want.min(available).max(min); + allocations.insert(label.to_string(), granted); + granted + } + + pub fn release(&self, label: &str) { + self.allocations.lock().expect("budget lock").remove(label); + } + + pub fn usable(&self) -> u32 { + self.usable + } + + pub fn table(&self) -> Vec<(String, u32)> { + self.allocations + .lock() + .expect("budget lock") + .iter() + .map(|(label, count)| (label.clone(), *count)) + .collect() + } +} + +/// A function that builds one runtime's client over its own pool. +pub type ClientFactory = Arc Client + Send + Sync>; + +/// One awa runtime living on its own tokio runtime, so it can be stopped gracefully +/// or killed the way a crashed process dies: tasks vanish mid-job without any +/// shutdown handshake and the connections drop. +pub struct ClientHost { + pub label: String, + runtime: tokio::runtime::Runtime, + client: Arc, + pool: PgPool, +} + +impl ClientHost { + pub async fn spawn( + url: String, + label: String, + pool_size: u32, + worker_threads: usize, + factory: ClientFactory, + ) -> Result { + let thread_label = label.clone(); + tokio::task::spawn_blocking(move || { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(worker_threads) + .thread_name(format!("soak-{thread_label}")) + .enable_all() + .build() + .map_err(|error| format!("build runtime: {error}"))?; + let (pool, client) = runtime.block_on(async { + let pool = connect_pool(&url, &format!("soak_{thread_label}"), pool_size) + .await + .map_err(|error| format!("connect pool: {error}"))?; + let client = factory(pool.clone()); + client + .start() + .await + .map_err(|error| format!("start client: {error}"))?; + Ok::<_, String>((pool, client)) + })?; + Ok(Self { + label: thread_label, + runtime, + client: Arc::new(client), + pool, + }) + }) + .await + .map_err(|error| format!("spawn host task: {error}"))? + } + + pub fn client(&self) -> &Client { + &self.client + } + + pub fn pool(&self) -> &PgPool { + &self.pool + } + + /// Graceful stop: stop claiming, drain in-flight jobs, then close the pool. + pub async fn shutdown(self, timeout: Duration) { + let client = self.client.clone(); + let pool = self.pool.clone(); + let _ = self + .runtime + .handle() + .spawn(async move { + client.shutdown(timeout).await; + pool.close().await; + }) + .await; + self.runtime.shutdown_background(); + } + + /// Abrupt stop: no shutdown handshake, as when a process is killed. + pub fn kill(self) { + self.runtime.shutdown_background(); + } +} + +pub fn url_for_database(config: &SoakConfig, name: &str) -> String { + replace_database_name(&config.database_url, name) +} diff --git a/awa/tests/persona_soak/fleet.rs b/awa/tests/persona_soak/fleet.rs new file mode 100644 index 000000000..b7e709ca1 --- /dev/null +++ b/awa/tests/persona_soak/fleet.rs @@ -0,0 +1,192 @@ +//! A persona's group of runtimes, with the capacity bookkeeping the stall check +//! needs and the lifecycle operations the operator and chaos personas use. + +use super::config::Persona; +use super::db::{ClientFactory, ClientHost}; +use super::registry::QueueInfo; +use super::shared::Shared; +use sqlx::PgPool; +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::Mutex; + +pub struct Fleet { + pub name: String, + pub persona: Persona, + /// Every runtime in the fleet serves these queues with this many workers. + served: Vec<(Arc, i64)>, + factory: ClientFactory, + pool_size: u32, + worker_threads: usize, + url: String, + hosts: Mutex>>, + budget_label: String, +} + +impl Fleet { + #[allow(clippy::too_many_arguments)] + pub async fn start( + shared: &Arc, + name: &str, + persona: Persona, + replicas: usize, + pool_size: u32, + worker_threads: usize, + served: Vec<(Arc, i64)>, + factory: ClientFactory, + ) -> Result, String> { + let url = shared.config.database_url.clone(); + Self::start_with_url( + shared, + url, + name, + persona, + replicas, + pool_size, + worker_threads, + served, + factory, + ) + .await + } + + /// As [`Fleet::start`], with every runtime connecting through `url`, for a fleet that + /// runs as a role other than the one in the run's database URL. + #[allow(clippy::too_many_arguments)] + pub async fn start_with_url( + shared: &Arc, + url: String, + name: &str, + persona: Persona, + replicas: usize, + pool_size: u32, + worker_threads: usize, + served: Vec<(Arc, i64)>, + factory: ClientFactory, + ) -> Result, String> { + let budget_label = name.to_string(); + let granted = shared.budget.reserve( + &budget_label, + pool_size * replicas as u32, + 2 * replicas as u32, + ); + let pool_size = (granted / replicas as u32).max(2); + let fleet = Arc::new(Self { + name: name.to_string(), + persona, + served, + factory, + pool_size, + worker_threads, + url, + hosts: Mutex::new(Vec::new()), + budget_label, + }); + for index in 0..replicas { + let host = fleet.spawn_host(index).await?; + fleet.hosts.lock().await.push(Some(host)); + } + shared + .fleets + .lock() + .expect("fleets") + .insert(name.to_string(), fleet.clone()); + Ok(fleet) + } + + async fn spawn_host(&self, index: usize) -> Result { + let host = ClientHost::spawn( + self.url.clone(), + format!("{}_{index}", self.name), + self.pool_size, + self.worker_threads, + self.factory.clone(), + ) + .await?; + for (queue, workers) in &self.served { + queue.capacity.fetch_add(*workers, Ordering::Relaxed); + } + Ok(host) + } + + fn release_capacity(&self) { + for (queue, workers) in &self.served { + queue.capacity.fetch_sub(*workers, Ordering::Relaxed); + } + } + + pub async fn replicas(&self) -> usize { + self.hosts.lock().await.len() + } + + pub async fn live_replicas(&self) -> Vec { + self.hosts + .lock() + .await + .iter() + .enumerate() + .filter_map(|(index, host)| host.is_some().then_some(index)) + .collect() + } + + pub async fn pool_of(&self, index: usize) -> Option { + self.hosts + .lock() + .await + .get(index)? + .as_ref() + .map(|host| host.pool().clone()) + } + + /// Stops runtime `index` gracefully and starts a fresh one in its place. + pub async fn restart(&self, index: usize, timeout: Duration) -> Result { + let old = self + .hosts + .lock() + .await + .get_mut(index) + .and_then(Option::take); + let Some(old) = old else { return Ok(false) }; + self.release_capacity(); + old.shutdown(timeout).await; + let fresh = self.spawn_host(index).await?; + self.hosts.lock().await[index] = Some(fresh); + Ok(true) + } + + /// Kills runtime `index` with no shutdown handshake, then (optionally) starts a + /// replacement after `downtime`. + pub async fn kill(&self, index: usize, downtime: Option) -> Result { + let old = self + .hosts + .lock() + .await + .get_mut(index) + .and_then(Option::take); + let Some(old) = old else { return Ok(false) }; + self.release_capacity(); + old.kill(); + if let Some(downtime) = downtime { + tokio::time::sleep(downtime).await; + let fresh = self.spawn_host(index).await?; + self.hosts.lock().await[index] = Some(fresh); + } + Ok(true) + } + + pub async fn shutdown_all(&self, timeout: Duration, shared: &Shared) { + let hosts: Vec = self + .hosts + .lock() + .await + .iter_mut() + .filter_map(Option::take) + .collect(); + for host in hosts { + self.release_capacity(); + host.shutdown(timeout).await; + } + shared.budget.release(&self.budget_label); + } +} diff --git a/awa/tests/persona_soak/ledger.rs b/awa/tests/persona_soak/ledger.rs new file mode 100644 index 000000000..a1c8ca236 --- /dev/null +++ b/awa/tests/persona_soak/ledger.rs @@ -0,0 +1,266 @@ +//! The client-side ledger: every job a persona enqueued, what the workers saw, and +//! what the system reported back. The database may prune terminal rows long before +//! the run ends, so the ledger is the oracle's memory of outcomes. + +use super::config::Persona; +use std::collections::HashMap; +use std::sync::Mutex; + +const SHARDS: usize = 64; + +/// NO_EVENTS workloads keep their fast path; durable loss detection covers +/// this deterministic sample, chosen from the immutable assigned job ID. +pub const TERMINAL_PROOF_STRIDE: i64 = 64; + +pub fn samples_terminal_proof(job_id: i64) -> bool { + job_id.rem_euclid(TERMINAL_PROOF_STRIDE) == 0 +} + +/// What a job is allowed to end as. +pub mod expect { + pub const COMPLETED: u8 = 1; + /// Failed, which with the DLQ enabled means a DLQ row. + pub const FAILED: u8 = 2; + pub const CANCELLED: u8 = 4; + pub const ANY_TERMINAL: u8 = COMPLETED | FAILED | CANCELLED; +} + +/// Facts about a job, accumulated as the run proceeds. +pub mod flag { + /// An operator cancel, drain or purge may legitimately have removed the job. + pub const OPERATOR_REMOVED: u32 = 1 << 0; + pub const OPERATOR_CANCEL_OK: u32 = 1 << 1; + pub const OPERATOR_MOVED: u32 = 1 << 2; + pub const DLQ_RETRIED: u32 = 1 << 3; + pub const HANDLER_CANCELLED: u32 = 1 << 4; + pub const RESCUED: u32 = 1 << 5; + pub const EXHAUSTED: u32 = 1 << 6; + pub const COMPLETED_EVENT: u32 = 1 << 7; + pub const RETURNED_COMPLETED: u32 = 1 << 8; + /// The job is served by a firehose queue where lifecycle events are not collected. + pub const NO_EVENTS: u32 = 1 << 9; + /// First seen at execution time (cron-created); not enqueued by a producer. + pub const UNREGISTERED: u32 = 1 << 10; + /// A DLQ row for this job was purged by an operator. + pub const DLQ_PURGED: u32 = 1 << 11; + pub const CHAOS_VICTIM: u32 = 1 << 12; + /// Terminal rows read from the store for jobs whose runtime reports no lifecycle events. + pub const DB_SEEN_COMPLETED: u32 = 1 << 13; + pub const DB_SEEN_FAILED: u32 = 1 << 14; + pub const DB_SEEN_CANCELLED: u32 = 1 << 15; + /// How a firehose job was enqueued. + pub const SRC_BURST: u32 = 1 << 16; + pub const SRC_COPY: u32 = 1 << 17; + pub const SRC_STEADY: u32 = 1 << 18; +} + +#[derive(Debug, Clone)] +pub struct JobRecord { + pub persona: Persona, + pub queue: u32, + pub enqueued_ms: u32, + pub expect: u8, + pub flags: u32, + pub executions: u16, + pub in_flight: u8, + pub overlapped: bool, + pub max_attempt: i16, + pub latest_run_lease: i64, + pub terminal_run_lease: Option, + pub first_start_ms: u32, + pub last_start_ms: u32, + pub completed_events: u8, + pub exhausted_events: u8, + pub cancel_events: u8, + pub completed_event_ms: u32, + /// When an operator cancel on this job committed, 0 if none. + pub operator_cancel_ms: u32, +} + +impl JobRecord { + pub fn new(persona: Persona, queue: u32, enqueued_ms: u32, expect: u8, flags: u32) -> Self { + Self { + persona, + queue, + enqueued_ms, + expect, + flags, + executions: 0, + in_flight: 0, + overlapped: false, + max_attempt: 0, + latest_run_lease: 0, + terminal_run_lease: None, + first_start_ms: 0, + last_start_ms: 0, + completed_events: 0, + exhausted_events: 0, + cancel_events: 0, + completed_event_ms: 0, + operator_cancel_ms: 0, + } + } + + pub fn note_terminal(&mut self, run_lease: i64) { + if self.terminal_run_lease.is_none_or(|old| run_lease > old) { + // DB outcome bits belong to the newest proven attempt. A later + // lifecycle proof also supersedes an older observed DB state. + self.flags &= + !(flag::DB_SEEN_COMPLETED | flag::DB_SEEN_FAILED | flag::DB_SEEN_CANCELLED); + self.terminal_run_lease = Some(run_lease); + } + } + + pub fn note_observed_terminal(&mut self, run_lease: i64, state: &str) { + self.note_terminal(run_lease); + if self.terminal_run_lease == Some(run_lease) { + // The query is unordered: an older row can follow a newer proof. + // Different states at the same lease remain visible as a conflict. + self.flags |= match state { + "completed" => flag::DB_SEEN_COMPLETED, + "failed" => flag::DB_SEEN_FAILED, + "cancelled" => flag::DB_SEEN_CANCELLED, + _ => 0, + }; + } + } + + pub fn has_durable_terminal(&self) -> bool { + self.terminal_run_lease + .is_some_and(|lease| lease >= self.latest_run_lease) + } + + pub fn has(&self, flag: u32) -> bool { + self.flags & flag != 0 + } +} + +pub struct Ledger { + shards: Vec>>, + /// `(cron name, fire time)` -> job ids executed for it; more than one is a fencing failure. + cron_fires: Mutex>>, +} + +impl Default for Ledger { + fn default() -> Self { + Self::new() + } +} + +impl Ledger { + pub fn new() -> Self { + Self { + shards: (0..SHARDS).map(|_| Mutex::new(HashMap::new())).collect(), + cron_fires: Mutex::new(HashMap::new()), + } + } + + fn shard(&self, id: i64) -> &Mutex> { + &self.shards[(id as u64 % SHARDS as u64) as usize] + } + + /// Records an enqueue. A fast worker can run the job before its producer gets here, so + /// an existing record keeps what the workers already reported. + pub fn register(&self, id: i64, record: JobRecord) { + let mut shard = self.shard(id).lock().expect("ledger"); + match shard.get_mut(&id) { + Some(existing) => { + existing.persona = record.persona; + existing.queue = record.queue; + existing.enqueued_ms = record.enqueued_ms; + existing.expect = record.expect; + existing.flags = (existing.flags | record.flags) & !flag::UNREGISTERED; + } + None => { + shard.insert(id, record); + } + } + } + + /// Applies `update` to the job's record, creating an unregistered record when the + /// job was not enqueued by a producer (cron fires, operator moves). + pub fn with( + &self, + id: i64, + default: impl FnOnce() -> JobRecord, + update: impl FnOnce(&mut JobRecord) -> R, + ) -> R { + let mut shard = self.shard(id).lock().expect("ledger"); + let record = shard.entry(id).or_insert_with(|| { + let mut record = default(); + record.flags |= flag::UNREGISTERED; + record + }); + update(record) + } + + pub fn mark(&self, id: i64, flags: u32) { + if let Some(record) = self.shard(id).lock().expect("ledger").get_mut(&id) { + record.flags |= flags; + } + } + + pub fn mark_many(&self, ids: &[i64], flags: u32) { + for id in ids { + self.mark(*id, flags); + } + } + + /// Adds `flags` to every record for which `predicate` holds. + pub fn mark_where(&self, predicate: impl Fn(&JobRecord) -> bool, flags: u32) -> usize { + let mut marked = 0; + for shard in &self.shards { + for record in shard.lock().expect("ledger").values_mut() { + if predicate(record) { + record.flags |= flags; + marked += 1; + } + } + } + marked + } + + pub fn get(&self, id: i64) -> Option { + self.shard(id).lock().expect("ledger").get(&id).cloned() + } + + pub fn len(&self) -> usize { + self.shards + .iter() + .map(|shard| shard.lock().expect("ledger").len()) + .sum() + } + + /// Calls `visit` for every record. Shards are locked one at a time. + pub fn for_each(&self, mut visit: impl FnMut(i64, &JobRecord)) { + for shard in &self.shards { + let shard = shard.lock().expect("ledger"); + for (id, record) in shard.iter() { + visit(*id, record); + } + } + } + + pub fn note_cron_fire(&self, name: String, fire_time: String, job_id: i64) -> bool { + let mut fires = self.cron_fires.lock().expect("cron fires"); + let ids = fires.entry((name, fire_time)).or_default(); + if !ids.contains(&job_id) { + ids.push(job_id); + } + ids.len() == 1 + } + + pub fn cron_duplicates(&self) -> Vec<((String, String), Vec)> { + self.cron_fires + .lock() + .expect("cron fires") + .iter() + .filter(|(_, ids)| ids.len() > 1) + .map(|(key, ids)| (key.clone(), ids.clone())) + .collect() + } + + pub fn cron_fire_count(&self) -> usize { + self.cron_fires.lock().expect("cron fires").len() + } +} diff --git a/awa/tests/persona_soak/mod.rs b/awa/tests/persona_soak/mod.rs new file mode 100644 index 000000000..2bccae9a0 --- /dev/null +++ b/awa/tests/persona_soak/mod.rs @@ -0,0 +1,24 @@ +//! Persona-driven soak harness with an invariant oracle. See the test binary +//! `persona_soak_test.rs` for how to run it. + +pub mod abandon_oracle; +pub mod adminserver; +pub mod config; +pub mod db; +pub mod fleet; +pub mod ledger; +pub mod oracle; +pub mod personas; +pub mod pgserver; +pub mod python_probe; +pub mod registry; +pub mod report; +pub mod ring_locks; +pub mod rng; +pub mod roles; +pub mod run; +pub mod shard_oracle; +pub mod shared; +pub mod stop; +pub mod terminal_proof; +pub mod worker_errors; diff --git a/awa/tests/persona_soak/oracle.rs b/awa/tests/persona_soak/oracle.rs new file mode 100644 index 000000000..aa70b2e5b --- /dev/null +++ b/awa/tests/persona_soak/oracle.rs @@ -0,0 +1,1392 @@ +//! The invariant oracle: periodic checks while the load runs and a final +//! reconciliation of the client ledger against the database. + +use super::ledger::{expect, flag, samples_terminal_proof, JobRecord}; +use super::pgserver::PgWatch; +use super::shared::Shared; +use super::stop::Stop; +use sqlx::PgPool; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; + +const MAX_ROWS: usize = 25; + +/// count, min and max job id, min and max lane sequence. +type ForensicRow = (i64, Option, Option, Option, Option); + +fn schema(shared: &Shared) -> String { + shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".to_string()) +} + +/// Every terminal outcome the store records, one row per `(job, attempt)`. +fn terminal_cte(schema: &str) -> String { + format!( + "terminal AS ( + SELECT job_id, run_lease, 'done:' || state::text AS src FROM {schema}.done_entries + UNION ALL SELECT job_id, run_lease, 'dlq' FROM {schema}.dlq_entries + UNION ALL SELECT b.job_id, b.run_lease, 'receipt_completed' + FROM {schema}.receipt_completion_batches c, + LATERAL unnest(c.job_ids, c.run_leases) AS b(job_id, run_lease) + )" + ) +} + +/// Jobs the lane cursors still consider claimable: at or after the claim cursor and +/// not tombstoned. +fn available_cte(schema: &str) -> String { + format!( + "heads AS MATERIALIZED ( + SELECT queue, priority, enqueue_shard, {schema}.sequence_next_value(seq_name) AS claim_seq + FROM {schema}.queue_claim_heads + ), + available AS MATERIALIZED ( + SELECT r.job_id, r.run_lease, r.queue + FROM {schema}.ready_entries r + JOIN heads h ON h.queue = r.queue AND h.priority = r.priority AND h.enqueue_shard = r.enqueue_shard + WHERE r.lane_seq >= h.claim_seq + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones t + WHERE t.ready_slot = r.ready_slot AND t.ready_generation = r.ready_generation + AND t.queue = r.queue AND t.priority = r.priority + AND t.enqueue_shard = r.enqueue_shard AND t.lane_seq = r.lane_seq + ) + )" + ) +} + +#[derive(Default)] +struct StallTracker { + /// queue -> (claim cursor, millis when the cursor last moved, already reported) + cursors: HashMap, +} + +pub struct Oracle { + shared: Arc, + pool: PgPool, + watch: Arc, + stall: StallTracker, + last_cancel_scan_ms: i64, + reported_duplicates: std::collections::HashSet<(i64, i64)>, + reported_cancel: std::collections::HashSet, + initial_stats_reset: std::sync::Mutex>>, +} + +impl Oracle { + pub fn new(shared: Arc, pool: PgPool, watch: Arc) -> Self { + Self { + shared, + pool, + watch, + stall: StallTracker::default(), + last_cancel_scan_ms: 0, + reported_duplicates: Default::default(), + reported_cancel: Default::default(), + initial_stats_reset: Default::default(), + } + } + + /// Runs the periodic checks until `stop`. + pub async fn run(mut self, stop: Stop) -> Self { + let mut tick = 0u64; + loop { + if stop.sleep(Duration::from_secs(5)).await { + return self; + } + tick += 1; + self.probe_server().await; + self.check_stall().await; + self.check_promotion().await; + if tick.is_multiple_of(2) { + self.check_terminal_uniqueness().await; + self.check_unexplained_cancels().await; + } + if tick.is_multiple_of(3) { + self.sample_resources().await; + } + } + } + + /// Crash recovery discards the cumulative statistics, so a changed reset time means the + /// server restarted its processes even though the postmaster, and its start time, survived. + async fn probe_stats_reset(&self) { + let reset: Result, _> = + sqlx::query_scalar("SELECT stats_reset::text FROM pg_stat_bgwriter") + .fetch_one(&self.pool) + .await; + let reset = match reset { + Ok(reset) => reset, + Err(error) => { + self.shared + .incomplete(format!("crash-recovery probe failed: {error}"), vec![]); + return; + } + }; + let mut initial = self.initial_stats_reset.lock().expect("stats reset"); + match initial.as_ref() { + None => *initial = Some(reset), + Some(first) if *first != reset => { + let message = format!("pg_stat_bgwriter.stats_reset changed from {first:?} to {reset:?}: the server went through crash recovery"); + *initial = Some(reset); + drop(initial); + self.shared + .violation("server_crash_recovery", message, vec![]); + } + Some(_) => {} + } + } + + async fn probe_server(&self) { + self.probe_stats_reset().await; + let mut last_error = String::new(); + for attempt in 0..4 { + match PgWatch::postmaster_start(&self.pool).await { + Ok(start) if start == self.watch.started_at => return, + Ok(start) => { + self.shared.violation( + "server_restart", + format!( + "pg_postmaster_start_time changed from {} to {start}", + self.watch.started_at + ), + vec![], + ); + return; + } + Err(error) => { + last_error = error.to_string(); + tokio::time::sleep(Duration::from_secs(1 + attempt)).await; + } + } + } + self.shared.run.add("server_probe_failures", 1); + self.shared.incomplete( + format!("server probe failed repeatedly: {last_error}"), + vec![], + ); + } + + async fn check_terminal_uniqueness(&mut self) { + let schema = schema(&self.shared); + let terminal = terminal_cte(&schema); + let started = std::time::Instant::now(); + + let duplicates: Result)>, _> = sqlx::query_as(&format!( + "WITH {terminal} + SELECT job_id, run_lease, array_agg(src ORDER BY src) + FROM terminal GROUP BY job_id, run_lease HAVING count(*) > 1 + ORDER BY job_id LIMIT {MAX_ROWS}" + )) + .fetch_all(&self.pool) + .await; + match duplicates { + Ok(rows) => { + let fresh: Vec<_> = rows + .iter() + .filter(|(job_id, run_lease, _)| { + self.reported_duplicates.insert((*job_id, *run_lease)) + }) + .collect(); + if !fresh.is_empty() { + self.shared.violation( + "duplicate_terminal_rows", + format!("{} job attempt(s) have more than one terminal row", fresh.len()), + fresh + .iter() + .map(|(job_id, run_lease, sources)| { + format!("{{\"job_id\":{job_id},\"run_lease\":{run_lease},\"terminal_rows\":{sources:?}}}") + }) + .collect(), + ); + } + } + Err(error) => self + .shared + .incomplete(format!("terminal uniqueness query failed: {error}"), vec![]), + } + + let available = available_cte(&schema); + let overlap: Result, _> = sqlx::query_as(&format!( + "WITH {terminal}, {available} + SELECT t.job_id, 'available'::text, t.run_lease, a.run_lease + FROM terminal t JOIN available a ON a.job_id = t.job_id AND a.run_lease <= t.run_lease + UNION ALL + SELECT t.job_id, 'deferred:' || d.state::text, t.run_lease, d.run_lease + FROM terminal t JOIN {schema}.deferred_jobs d ON d.job_id = t.job_id AND d.run_lease <= t.run_lease + UNION ALL + SELECT t.job_id, 'leased:' || l.state::text, t.run_lease, l.run_lease + FROM terminal t JOIN {schema}.leases l ON l.job_id = t.job_id AND l.run_lease <= t.run_lease + LIMIT {MAX_ROWS}" + )) + .fetch_all(&self.pool) + .await; + match overlap { + Ok(rows) if !rows.is_empty() => { + // Confirm on a second read: a single statement is one snapshot, so a + // repeat that still shows the overlap is not a race with this probe. + tokio::time::sleep(Duration::from_millis(500)).await; + let ids: Vec = rows.iter().map(|row| row.0).collect(); + let still: Result, _> = sqlx::query_as(&format!( + "WITH {terminal}, {available} + SELECT t.job_id, 'available'::text FROM terminal t JOIN available a ON a.job_id = t.job_id AND a.run_lease <= t.run_lease WHERE t.job_id = ANY($1) + UNION ALL SELECT t.job_id, 'deferred' FROM terminal t JOIN {schema}.deferred_jobs d ON d.job_id = t.job_id AND d.run_lease <= t.run_lease WHERE t.job_id = ANY($1) + UNION ALL SELECT t.job_id, 'leased' FROM terminal t JOIN {schema}.leases l ON l.job_id = t.job_id AND l.run_lease <= t.run_lease WHERE t.job_id = ANY($1)" + )) + .bind(&ids) + .fetch_all(&self.pool) + .await; + let confirmed: std::collections::HashSet = match still { + Ok(rows) => rows.into_iter().map(|row| row.0).collect(), + Err(error) => { + self.shared.incomplete( + format!("terminal/live confirmation failed: {error}"), + vec![], + ); + Default::default() + } + }; + let confirmed_rows: Vec = rows + .iter() + .filter(|row| confirmed.contains(&row.0)) + .filter(|row| self.reported_duplicates.insert((row.0, -row.2 - 1))) + .map(|(job_id, place, terminal_lease, live_lease)| { + format!("{{\"job_id\":{job_id},\"also_in\":\"{place}\",\"terminal_run_lease\":{terminal_lease},\"live_run_lease\":{live_lease}}}") + }) + .collect(); + if !confirmed_rows.is_empty() { + self.shared.violation( + "terminal_and_live", + format!( + "{} job(s) are terminal and still available, deferred or leased", + confirmed_rows.len() + ), + confirmed_rows, + ); + } + } + Ok(_) => {} + Err(error) => self + .shared + .incomplete(format!("terminal/live query failed: {error}"), vec![]), + } + self.shared.run.observe( + "oracle_terminal_check_ms", + started.elapsed().as_millis() as i64, + ); + } + + /// A `cancelled` terminal row nobody asked for: no operator cancel, drain or purge + /// touched the job and its handler did not cancel it. + async fn check_unexplained_cancels(&mut self) { + let schema = schema(&self.shared); + let since_ms = self.last_cancel_scan_ms; + let scan_started_ms = self.shared.now_ms(); + let rows: Result, _> = sqlx::query_as(&format!( + "SELECT job_id, queue, run_lease FROM {schema}.done_entries + WHERE state = 'cancelled' AND finalized_at > now() - make_interval(secs => $1) + LIMIT 2000" + )) + .bind(((scan_started_ms - since_ms) as f64 / 1000.0 + 15.0).max(20.0)) + .fetch_all(&self.pool) + .await; + self.last_cancel_scan_ms = scan_started_ms; + let rows = match rows { + Ok(rows) => rows, + Err(error) => { + self.shared + .incomplete(format!("invariant query failed: {error}"), vec![]); + return; + } + }; + let mut unexplained = Vec::new(); + for (job_id, queue_name, run_lease) in rows { + let Some(record) = self.shared.ledger.get(job_id) else { + continue; + }; + let queue = self.shared.queues.get(&queue_name); + let drained = queue.as_ref().is_some_and(|queue| { + let drained_at = queue + .drained_at_ms + .load(std::sync::atomic::Ordering::Relaxed); + drained_at > 0 && i64::from(record.enqueued_ms) <= drained_at + 1000 + }); + let explained = record.has(flag::OPERATOR_REMOVED) + || record.has(flag::HANDLER_CANCELLED) + || record.has(flag::OPERATOR_MOVED) + || drained; + if !explained { + // An operator flag is set just after its cancel commits; look again before accusing. + let lag = if self + .shared + .ledger + .get(job_id) + .is_some_and(|record| record.persona == super::config::Persona::Python) + { + 4500 + } else { + 1500 + }; + tokio::time::sleep(Duration::from_millis(lag)).await; + } + let explained = explained + || self.shared.ledger.get(job_id).is_some_and(|record| { + record.has(flag::OPERATOR_REMOVED) || record.has(flag::HANDLER_CANCELLED) + }); + if !explained && self.reported_cancel.insert(job_id) { + unexplained.push(format!( + "{{\"job_id\":{job_id},\"queue\":\"{queue_name}\",\"run_lease\":{run_lease},\"executions\":{},\"running_now\":{},\"completed_events\":{}}}", + record.executions, record.in_flight, record.completed_events + )); + } + } + if !unexplained.is_empty() { + self.shared.violation( + "unexplained_cancel", + format!( + "{} job(s) recorded as cancelled with no operator or handler cancel", + unexplained.len() + ), + unexplained.into_iter().take(MAX_ROWS).collect(), + ); + } + } + + async fn check_promotion(&self) { + let schema = schema(&self.shared); + let mut threshold = self.shared.config.promotion_overdue; + if self.shared.chaos_active() { + threshold += self.shared.config.chaos_slack; + } + let rows: Result, _> = sqlx::query_as(&format!( + "SELECT job_id, queue, extract(epoch FROM now() - run_at)::float8 + FROM {schema}.deferred_jobs + WHERE state IN ('scheduled', 'retryable') AND run_at < now() - make_interval(secs => $1) + ORDER BY run_at LIMIT {MAX_ROWS}" + )) + .bind(threshold.as_secs_f64()) + .fetch_all(&self.pool) + .await; + let rows = match rows { + Ok(rows) => rows, + Err(error) => { + self.shared + .incomplete(format!("invariant query failed: {error}"), vec![]); + return; + } + }; + if rows.is_empty() { + return; + } + // Re-read so that a job promoted a moment ago is not reported. + tokio::time::sleep(Duration::from_secs(2)).await; + let ids: Vec = rows.iter().map(|row| row.0).collect(); + let still: Result, _> = sqlx::query_scalar(&format!( + "SELECT job_id FROM {schema}.deferred_jobs WHERE job_id = ANY($1) AND state IN ('scheduled','retryable')" + )) + .bind(&ids) + .fetch_all(&self.pool) + .await; + let still = match still { + Ok(still) => still, + Err(error) => { + self.shared + .incomplete(format!("promotion confirmation failed: {error}"), vec![]); + return; + } + }; + let stuck: Vec<_> = rows.iter().filter(|row| still.contains(&row.0)).collect(); + if stuck.is_empty() { + return; + } + let worst = stuck.iter().map(|row| row.2).fold(0.0, f64::max); + self.shared + .run + .set_max("promotion_overdue_max_s", worst as i64); + self.shared.violation( + "promotion_stall", + format!( + "{} scheduled job(s) overdue beyond {}s, worst {worst:.1}s", + stuck.len(), + threshold.as_secs() + ), + stuck + .iter() + .map(|(job_id, queue, overdue)| { + format!( + "{{\"job_id\":{job_id},\"queue\":\"{queue}\",\"overdue_s\":{overdue:.1}}}" + ) + }) + .collect(), + ); + } + + async fn check_stall(&mut self) { + let schema = schema(&self.shared); + let rows: Result, _> = sqlx::query_as(&format!( + "SELECT c.queue, + COALESCE(sum({schema}.sequence_next_value(e.seq_name) - {schema}.sequence_next_value(c.seq_name)), 0)::bigint, + COALESCE(sum({schema}.sequence_next_value(c.seq_name)), 0)::bigint + FROM {schema}.queue_claim_heads c + JOIN {schema}.queue_enqueue_heads e + ON e.queue = c.queue AND e.priority = c.priority AND e.enqueue_shard = c.enqueue_shard + GROUP BY c.queue" + )) + .fetch_all(&self.pool) + .await; + let rows = match rows { + Ok(rows) => rows, + Err(error) => { + self.shared + .incomplete(format!("invariant query failed: {error}"), vec![]); + return; + } + }; + let now = self.shared.now_ms(); + let mut window = self.shared.config.stall_window; + if self.shared.chaos_active() { + window += self.shared.config.chaos_slack; + } + let mut candidates = Vec::new(); + for (queue_name, backlog, cursor) in rows { + let entry = self + .stall + .cursors + .entry(queue_name.clone()) + .or_insert((cursor, now, false)); + if entry.0 != cursor { + *entry = (cursor, now, false); + continue; + } + let Some(queue) = self.shared.queues.get(&queue_name) else { + continue; + }; + let capacity = queue.capacity.load(std::sync::atomic::Ordering::Relaxed); + let in_flight = queue.in_flight.load(std::sync::atomic::Ordering::Relaxed); + let held = queue.is_held() || queue.rate_limited; + let idle_workers = capacity > 0 && in_flight < capacity; + if !entry.2 + && backlog > 0 + && idle_workers + && !held + && (now - entry.1) as u64 >= window.as_millis() as u64 + { + candidates.push((queue_name, backlog, cursor, entry.1, capacity, in_flight)); + } else if !idle_workers || held || backlog <= 0 { + // Busy or deliberately held: the stall clock restarts. + entry.1 = now; + } + } + for (queue_name, backlog, cursor, since, capacity, in_flight) in candidates { + let exact: Result = sqlx::query_scalar(&format!( + "WITH {} SELECT count(*)::bigint FROM available WHERE queue = $1", + available_cte(&schema) + )) + .bind(&queue_name) + .fetch_one(&self.pool) + .await; + let available = match exact { + Ok(available) => available, + Err(error) => { + self.shared + .incomplete(format!("exact claimable-row check failed: {error}"), vec![]); + continue; + } + }; + if available > 0 { + if let Some(entry) = self.stall.cursors.get_mut(&queue_name) { + entry.2 = true; + } + self.shared.run.add("stalls", 1); + self.shared.violation( + "stall", + format!( + "queue {queue_name}: {available} available job(s), claim cursor frozen at {cursor} for {}s with {}/{} workers busy", + (now - since) / 1000, + in_flight, + capacity + ), + vec![format!( + "{{\"queue\":\"{queue_name}\",\"available\":{available},\"lane_backlog\":{backlog},\"claim_cursor\":{cursor},\"workers_busy\":{in_flight},\"worker_capacity\":{capacity}}}" + )], + ); + } else if let Some(entry) = self.stall.cursors.get_mut(&queue_name) { + entry.1 = now; + } + } + } + + /// Stops the run before it can fill the disk it shares with other work. + fn check_disk(&self) { + let output = std::process::Command::new("df") + .args(["--output=avail", "-BM"]) + .arg(&self.shared.config.report_dir) + .output(); + let Ok(output) = output else { return }; + let text = String::from_utf8_lossy(&output.stdout); + let Some(free_mb) = text + .lines() + .nth(1) + .and_then(|line| line.trim().trim_end_matches('M').parse::().ok()) + else { + return; + }; + self.shared.run.observe("free_disk_gb", free_mb / 1024); + if free_mb < self.shared.config.min_free_disk_gb * 1024 + && !self.shared.load_stop.is_stopped() + { + self.shared.violation( + "disk_low", + format!( + "only {} GB free on the report directory's filesystem; stopping the load phase", + free_mb / 1024 + ), + vec![], + ); + self.shared.load_stop.trigger(); + } + } + + async fn sample_resources(&self) { + self.check_disk(); + let connections: Result = sqlx::query_scalar( + "SELECT count(*)::bigint FROM pg_stat_activity WHERE backend_type = 'client backend'", + ) + .fetch_one(&self.pool) + .await; + if let Ok(connections) = connections { + self.shared.run.set_max("peak_connections", connections); + self.shared.run.observe("connections", connections); + } + let size: Result = + sqlx::query_scalar("SELECT pg_database_size(current_database())") + .fetch_one(&self.pool) + .await; + if let Ok(size) = size { + self.shared.run.set_max("peak_db_mb", size / (1024 * 1024)); + self.shared.run.observe("db_mb", size / (1024 * 1024)); + } + if let Ok(status) = std::fs::read_to_string("/proc/self/status") { + if let Some(kb) = status + .lines() + .find_map(|line| line.strip_prefix("VmRSS:")) + .and_then(|rest| { + rest.trim() + .trim_end_matches("kB") + .trim() + .parse::() + .ok() + }) + { + self.shared.run.set_max("peak_harness_rss_mb", kb / 1024); + } + } + if let Ok(load) = std::fs::read_to_string("/proc/loadavg") { + if let Some(one_minute) = load + .split_whitespace() + .next() + .and_then(|v| v.parse::().ok()) + { + self.shared + .run + .observe("loadavg_x100", (one_minute * 100.0) as i64); + } + } + } + + /// Waits for the system to quiesce after load stops. Returns true when no live + /// work remains. + pub async fn wait_quiescent(&mut self, timeout: Duration) -> bool { + let schema = schema(&self.shared); + let deadline = std::time::Instant::now() + timeout; + let mut quiet_polls = 0; + loop { + let live: Result = sqlx::query_scalar(&format!( + "WITH {} SELECT (SELECT count(*) FROM available) + (SELECT count(*) FROM {schema}.leases) + + (SELECT count(*) FROM {schema}.deferred_jobs WHERE run_at < now() + interval '1 second')", + available_cte(&schema) + )) + .fetch_one(&self.pool) + .await; + if let Err(error) = &live { + self.shared + .incomplete(format!("quiescence query failed: {error}"), vec![]); + } + let mut pending_terminal_proof = 0usize; + self.shared.ledger.for_each(|id, record| { + if record.has(flag::RETURNED_COMPLETED) + && !record.has_durable_terminal() + && !record.has(flag::OPERATOR_REMOVED) + && !record.has(flag::DLQ_PURGED) + && (!record.has(flag::NO_EVENTS) || samples_terminal_proof(id)) + { + pending_terminal_proof += 1; + } + }); + let in_flight: i64 = self + .shared + .queues + .all() + .iter() + .map(|queue| queue.in_flight.load(std::sync::atomic::Ordering::Relaxed)) + .sum(); + let python_active = self + .shared + .python_active + .load(std::sync::atomic::Ordering::Relaxed); + // Claims abandoned by a killed process sit in no table the live count reads. + let abandoned_open = self + .shared + .abandon + .open_claims + .load(std::sync::atomic::Ordering::Relaxed); + match live { + Ok(count) + if in_flight == 0 + && python_active == 0 + && abandoned_open == 0 + && pending_terminal_proof == 0 + && count_is_idle(count) => + { + quiet_polls += 1 + } + _ => quiet_polls = 0, + } + if quiet_polls >= 3 { + return true; + } + if std::time::Instant::now() >= deadline { + self.shared.note( + "oracle", + format!( + "not quiescent after {}s: handlers running={in_flight}, live rows={live:?}", + timeout.as_secs() + ), + ); + return false; + } + self.probe_server().await; + self.check_stall().await; + tokio::time::sleep(Duration::from_secs(2)).await; + } + } + + /// Final reconciliation of the ledger against the database. + pub async fn final_check(&mut self, quiesced: bool) { + let schema = schema(&self.shared); + self.check_terminal_uniqueness().await; + self.check_unexplained_cancels().await; + self.check_promotion().await; + self.probe_server().await; + + self.shared.run.add( + "cron_fires_executed", + self.shared.ledger.cron_fire_count() as i64, + ); + for ((name, fire), ids) in self.shared.ledger.cron_duplicates() { + self.shared.violation( + "cron_double_fire", + format!("cron {name} fire {fire} produced jobs {ids:?}"), + vec![], + ); + } + + // Jobs with no terminal evidence on the client side have to be found in the database. + let mut unresolved: Vec = Vec::new(); + let mut unexpected: Vec = Vec::new(); + let mut sampled = 0i64; + let mut confirmed = 0i64; + let mut excluded = 0i64; + self.shared.ledger.for_each(|id, record| { + if record.has(flag::OPERATOR_REMOVED) || record.has(flag::DLQ_PURGED) { + return; + } + let sampled_no_events = record.has(flag::NO_EVENTS) && samples_terminal_proof(id); + sampled += i64::from(sampled_no_events); + if record.has_durable_terminal() { + confirmed += i64::from(sampled_no_events); + for outcome in unexpected_terminal_outcomes(record) { + if unexpected.len() >= MAX_ROWS { + break; + } + let queue = self.shared.queues.by_index(record.queue); + unexpected.push(serde_json::json!({ + "job_id": id, + "outcome": outcome, + "allowed": record.expect, + "persona": record.persona.name(), + "queue": queue.as_ref().map(|queue| queue.name.chars().take(256).collect::()), + "flags": record.flags, + "latest_run_lease": record.latest_run_lease, + "terminal_run_lease": record.terminal_run_lease, + "executions": record.executions, + "in_flight": record.in_flight, + "max_attempt": record.max_attempt, + "completed_events": record.completed_events, + "exhausted_events": record.exhausted_events, + "cancel_events": record.cancel_events, + "enqueued_ms": record.enqueued_ms, + "first_start_ms": record.first_start_ms, + "last_start_ms": record.last_start_ms, + "completed_event_ms": record.completed_event_ms, + "operator_cancel_ms": record.operator_cancel_ms, + "queue_drained_at_ms": queue.map(|queue| queue.drained_at_ms.load(std::sync::atomic::Ordering::Relaxed)), + }).to_string()); + } + } else if record.has(flag::NO_EVENTS) + && record.has(flag::RETURNED_COMPLETED) + && !sampled_no_events + { + // This optimized workload is sampled, not proven per job. + // A return is execution evidence only, never commit evidence. + excluded += 1; + } else { + unresolved.push(id); + } + }); + if !unexpected.is_empty() { + self.shared.violation( + "unexpected_outcome", + format!( + "{} job(s) ended in an outcome their persona did not allow", + unexpected.len() + ), + unexpected, + ); + } + + let mut live_stuck = Vec::new(); + let mut forensic_ids: Vec = Vec::new(); + let mut lost = Vec::new(); + let mut lost_after_run = Vec::new(); + let mut terminal_in_db = 0usize; + for chunk in unresolved.chunks(5000) { + let found: Result)>, _> = sqlx::query_as(&format!( + "SELECT job_id, string_agg(place, ','), max(terminal_run_lease) FROM ( + SELECT job_id, 'done:' || state::text AS place, run_lease AS terminal_run_lease FROM {schema}.done_entries WHERE job_id = ANY($1) + UNION ALL SELECT job_id, 'dlq', run_lease FROM {schema}.dlq_entries WHERE job_id = ANY($1) + UNION ALL SELECT b.job_id, 'receipt_completed', b.run_lease + FROM {schema}.receipt_completion_batches c, LATERAL unnest(c.job_ids, c.run_leases) AS b(job_id, run_lease) + WHERE c.job_ids && $1 AND b.job_id = ANY($1) + UNION ALL SELECT job_id, 'deferred:' || state::text, NULL::bigint FROM {schema}.deferred_jobs WHERE job_id = ANY($1) + UNION ALL SELECT job_id, 'ready', NULL::bigint FROM {schema}.ready_entries WHERE job_id = ANY($1) + UNION ALL SELECT job_id, 'leased:' || state::text, NULL::bigint FROM {schema}.leases WHERE job_id = ANY($1) + ) places GROUP BY job_id" + )) + .bind(chunk) + .fetch_all(&self.pool) + .await; + let found: HashMap)> = match found { + Ok(rows) => rows + .into_iter() + .map(|(id, place, lease)| (id, (place, lease))) + .collect(), + Err(error) => { + self.shared.run.add("final_lookup_errors", 1); + self.shared.incomplete( + format!( + "final lookup failed for {} unresolved jobs: {error}", + chunk.len() + ), + vec![], + ); + continue; + } + }; + for id in chunk { + let record = self.shared.ledger.get(*id); + let describe = |place: &str| { + let record = record.clone(); + format!( + "{{\"job_id\":{id},\"found\":\"{place}\",\"executions\":{},\"persona\":\"{}\",\"flags\":{},\"enqueued_at_s\":{:.1},\"queue\":\"{}\"}}", + record.as_ref().map_or(0, |r| r.executions), + record.as_ref().map_or("?", |r| r.persona.name()), + record.as_ref().map_or(0, |r| r.flags), + record.as_ref().map_or(0.0, |r| f64::from(r.enqueued_ms) / 1000.0), + record + .as_ref() + .and_then(|r| self.shared.queues.by_index(r.queue)) + .map_or("?".to_string(), |queue| queue.name.clone()) + ) + }; + if record + .as_ref() + .is_some_and(|record| record.has_durable_terminal()) + { + confirmed += i64::from( + record + .as_ref() + .is_some_and(|record| record.has(flag::NO_EVENTS)) + && samples_terminal_proof(*id), + ); + continue; + } + match found.get(id) { + Some((place, terminal_lease)) + if terminal_lease.is_some_and(|lease| { + lease >= record.as_ref().map_or(0, |record| record.latest_run_lease) + }) => + { + terminal_in_db += 1; + confirmed += i64::from( + record + .as_ref() + .is_some_and(|record| record.has(flag::NO_EVENTS)) + && samples_terminal_proof(*id), + ); + if place.contains("done:cancelled") { + if let Some(record) = &record { + let queue = self.shared.queues.by_index(record.queue); + let drained = queue.is_some_and(|queue| { + let at = queue + .drained_at_ms + .load(std::sync::atomic::Ordering::Relaxed); + at > 0 && i64::from(record.enqueued_ms) <= at + 1000 + }); + if !(record.has(flag::OPERATOR_REMOVED) || drained) + && record.expect & expect::CANCELLED == 0 + { + lost_after_run.push(describe(place)); + forensic_ids.push(*id); + } + } + } + } + Some((place, _)) => live_stuck.push(describe(place)), + None => { + let executed = record.as_ref().is_some_and(|r| r.executions > 0); + let returned = record + .as_ref() + .is_some_and(|r| r.has(flag::RETURNED_COMPLETED)); + let drained = record.as_ref().is_some_and(|r| { + self.shared.queues.by_index(r.queue).is_some_and(|queue| { + let at = queue + .drained_at_ms + .load(std::sync::atomic::Ordering::Relaxed); + at > 0 && i64::from(r.enqueued_ms) <= at + 1000 + }) + }); + if drained { + continue; + } + if !executed { + lost.push(describe("absent")); + forensic_ids.push(*id); + } else if !returned { + lost_after_run.push(describe("absent")); + forensic_ids.push(*id); + } else { + self.shared.incomplete( + "a returned handler has no durable terminal proof; a missed pruned outcome cannot be distinguished from a lost finalize", + vec![describe("absent without durable terminal evidence")], + ); + } + } + } + } + } + if !forensic_ids.is_empty() { + self.record_forensics(&schema, &forensic_ids).await; + } + self.shared + .run + .add("terminal_proof_no_events_sampled", sampled); + self.shared + .run + .add("terminal_proof_no_events_confirmed", confirmed); + self.shared + .run + .add("terminal_proof_no_events_excluded", excluded); + self.shared + .run + .add("final_unresolved_ids", unresolved.len() as i64); + self.shared + .run + .add("final_terminal_in_db", terminal_in_db as i64); + let report = |kind: &str, what: &str, rows: Vec| { + if !rows.is_empty() { + let count = rows.len(); + self.shared.violation( + kind, + format!("{count} job(s) {what}"), + rows.into_iter().take(MAX_ROWS).collect(), + ); + } + }; + if quiesced { + report( + "stuck_job", + "never reached a terminal state although the system went quiet (available/deferred/leased)", + live_stuck, + ); + } else if !live_stuck.is_empty() { + // Work still queued when the grace ended is backlog, not loss: the run was + // sized beyond what the server could clear. + self.shared.finding( + "backlog_at_end", + format!( + "{} job(s) were still queued when the drain grace ended", + live_stuck.len() + ), + live_stuck.into_iter().take(MAX_ROWS).collect(), + ); + } + report( + "lost_job", + "were enqueued, never ran and are in no table", + lost, + ); + report("lost_job_after_run", "ran without a recorded terminal outcome and are in no table, or ended cancelled without cause", lost_after_run); + } + + /// For jobs that look lost: everywhere the store could still mention them, written to + /// the timeline so the reproducer carries it. + async fn record_forensics(&self, schema: &str, ids: &[i64]) { + let sample: Vec = ids.iter().copied().take(200).collect(); + let tables: [(&str, String); 8] = [ + ("ready_entries", format!("SELECT count(*), min(job_id), max(job_id), min(lane_seq), max(lane_seq) FROM {schema}.ready_entries WHERE job_id = ANY($1)")), + ("ready_tombstones", format!("SELECT count(*), min(job_id), max(job_id), min(lane_seq), max(lane_seq) FROM {schema}.ready_tombstones WHERE job_id = ANY($1)")), + ("lease_claims", format!("SELECT count(*), min(job_id), max(job_id), min(lane_seq), max(lane_seq) FROM {schema}.lease_claims WHERE job_id = ANY($1)")), + ("lease_claim_batches", format!("SELECT count(*), min(j), max(j), 0::bigint, 0::bigint FROM {schema}.lease_claim_batches b, LATERAL unnest(b.job_ids) AS j WHERE b.job_ids && $1 AND j = ANY($1)")), + ("lease_claim_closures", format!("SELECT count(*), min(job_id), max(job_id), 0::bigint, 0::bigint FROM {schema}.lease_claim_closures WHERE job_id = ANY($1)")), + ("attempt_state", format!("SELECT count(*), min(job_id), max(job_id), 0::bigint, 0::bigint FROM {schema}.attempt_state WHERE job_id = ANY($1)")), + ("receipt_completion_tombstones", format!("SELECT count(*), 0::bigint, 0::bigint, 0::bigint, 0::bigint FROM {schema}.receipt_completion_tombstones")), + ("job_unique_claims", "SELECT count(*), min(job_id), max(job_id), 0::bigint, 0::bigint FROM awa.job_unique_claims WHERE job_id = ANY($1)".to_string()), + ]; + let mut lines = Vec::new(); + for (name, sql) in tables { + let row: Result = sqlx::query_as(&sql) + .bind(&sample) + .fetch_one(&self.pool) + .await; + match row { + Ok((count, min_id, max_id, min_seq, max_seq)) => lines.push(format!("{name}: {count} row(s) for the sampled ids (job ids {min_id:?}..{max_id:?}, lane_seq {min_seq:?}..{max_seq:?})")), + Err(error) => lines.push(format!("{name}: query failed: {error}")), + } + } + self.shared.note( + "oracle", + format!( + "forensics for {} lost-looking job(s), first {:?}: {}", + ids.len(), + &ids[..ids.len().min(5)], + lines.join("; ") + ), + ); + } + + pub fn crash_scan(&self) -> Vec { + match self.watch.log_since_start(&self.shared.config) { + Some(log) => super::pgserver::crash_excerpts(&log), + None => Vec::new(), + } + } +} + +/// The final verdict's outcome predicates, shared by deterministic lease-order controls. +fn unexpected_terminal_outcomes(record: &JobRecord) -> Vec<&'static str> { + if !record.has_durable_terminal() { + return Vec::new(); + } + let completed = record.completed_events > 0 + || record.has(flag::RETURNED_COMPLETED) + || record.has(flag::DB_SEEN_COMPLETED); + [ + ( + "cancelled", + record.has(flag::DB_SEEN_CANCELLED) + && record.expect & expect::CANCELLED == 0 + && !record.has(flag::HANDLER_CANCELLED), + ), + ( + "completed", + completed && record.expect & expect::COMPLETED == 0, + ), + ( + "failed", + (record.exhausted_events > 0 || record.has(flag::DB_SEEN_FAILED)) + && record.expect & expect::FAILED == 0, + ), + ] + .into_iter() + .filter_map(|(outcome, unexpected)| unexpected.then_some(outcome)) + .collect() +} + +fn count_is_idle(count: i64) -> bool { + count <= 0 +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::persona_soak::config::{replace_database_name, Persona, SoakConfig}; + use crate::persona_soak::db::{self, Budget}; + use awa::model::{ClaimedRuntimeJob, QueueStorage, QueueStorageConfig}; + use std::sync::atomic::Ordering; + + fn current_terminal_record() -> JobRecord { + let mut record = JobRecord::new(Persona::Tom, 0, 0, expect::COMPLETED, 0); + record.latest_run_lease = 1; + record + } + + #[test] + fn terminal_state_old_cancel_then_new_completion_is_completed() { + let mut record = current_terminal_record(); + record.note_observed_terminal(0, "cancelled"); + record.note_observed_terminal(1, "completed"); + assert!(unexpected_terminal_outcomes(&record).is_empty()); + assert_eq!(record.terminal_run_lease, Some(1)); + assert!(!record.has(flag::DB_SEEN_CANCELLED)); + assert!(record.has(flag::DB_SEEN_COMPLETED)); + } + + #[test] + fn terminal_state_new_completion_then_old_cancel_is_completed() { + let mut record = current_terminal_record(); + record.note_observed_terminal(1, "completed"); + record.note_observed_terminal(0, "cancelled"); + assert!(unexpected_terminal_outcomes(&record).is_empty()); + assert!(!record.has(flag::DB_SEEN_CANCELLED)); + assert!(record.has(flag::DB_SEEN_COMPLETED)); + } + + #[test] + fn terminal_state_lifecycle_completion_supersedes_old_db_cancel() { + let mut record = current_terminal_record(); + record.note_observed_terminal(0, "cancelled"); + record.note_terminal(1); + record.completed_events += 1; + assert!(unexpected_terminal_outcomes(&record).is_empty()); + assert!(!record.has(flag::DB_SEEN_CANCELLED)); + record.note_observed_terminal(0, "cancelled"); + assert!(!record.has(flag::DB_SEEN_CANCELLED)); + assert!(unexpected_terminal_outcomes(&record).is_empty()); + } + + #[test] + fn terminal_state_new_cancel_remains_an_unexpected_outcome() { + let mut record = current_terminal_record(); + record.note_observed_terminal(0, "completed"); + record.note_observed_terminal(1, "cancelled"); + assert!(record.has(flag::DB_SEEN_CANCELLED)); + assert_eq!(unexpected_terminal_outcomes(&record), vec!["cancelled"]); + } + + #[test] + fn terminal_state_same_lease_conflicting_states_remain_visible() { + for states in [["cancelled", "completed"], ["completed", "cancelled"]] { + let mut record = current_terminal_record(); + for state in states { + record.note_observed_terminal(1, state); + } + assert!(record.has(flag::DB_SEEN_CANCELLED)); + assert!(record.has(flag::DB_SEEN_COMPLETED)); + assert_eq!(unexpected_terminal_outcomes(&record), vec!["cancelled"]); + } + } + + struct TestRun { + shared: Arc, + pool: PgPool, + store: QueueStorage, + claim: ClaimedRuntimeJob, + oracle: Oracle, + } + + impl TestRun { + async fn create() -> Self { + Self::create_with(true, true).await + } + + async fn create_with(no_events: bool, sampled: bool) -> Self { + let base = std::env::var("DATABASE_URL").unwrap_or_else(|_| { + "postgres://postgres:test@localhost:15432/awa_test".to_string() + }); + let name = format!("awa_soak_oracle_{}", uuid::Uuid::new_v4().simple()); + let config = SoakConfig::for_test(replace_database_name(&base, &name)); + db::recreate_database(&config) + .await + .expect("create isolated oracle database"); + let pool = db::connect_pool(&config.database_url, "soak_oracle_regression", 4) + .await + .expect("connect oracle"); + let info = db::server_info(&pool).await.expect("server info"); + let watch = Arc::new(PgWatch::begin(&config, &pool).await.expect("watch server")); + let budget = Budget::new(&info, 0); + let shared = Shared::new(config, info, budget); + let queue = shared + .queues + .register("oracle_proof", Persona::Priya, false); + let store = QueueStorage::new(QueueStorageConfig::default()).expect("store"); + store + .enqueue_batch(&pool, "oracle_proof", 2, 64) + .await + .expect("enqueue sample candidates"); + let claimed = store + .claim_runtime_batch(&pool, "oracle_proof", 64, Duration::ZERO) + .await + .expect("claim sample"); + let claim = claimed + .into_iter() + .find(|claim| samples_terminal_proof(claim.job.id) == sampled) + .expect("immutable job ID selects the requested sample status"); + shared.ledger.register( + claim.job.id, + JobRecord::new( + Persona::Priya, + queue.index, + 0, + expect::COMPLETED, + if no_events { flag::NO_EVENTS } else { 0 }, + ), + ); + // The handler returned, but its caller has not finalized the claim. + let execution = shared.begin_execution(Persona::Priya, &claim.job, no_events); + shared.ledger.mark(claim.job.id, flag::RETURNED_COMPLETED); + drop(execution); + let oracle = Oracle::new(shared.clone(), pool.clone(), watch); + Self { + shared, + pool, + store, + claim, + oracle, + } + } + + async fn drop(self) { + self.pool.close().await; + let admin = db::connect_pool(&self.shared.config.admin_url(), "oracle_cleanup", 1) + .await + .expect("cleanup connection"); + sqlx::query(&format!( + "DROP DATABASE {} WITH (FORCE)", + self.shared.config.database_name() + )) + .execute(&admin) + .await + .expect("drop oracle database"); + admin.close().await; + } + } + + #[tokio::test] + async fn returned_handler_without_finalization_cannot_pass_the_oracle() { + let mut run = TestRun::create().await; + run.oracle.final_check(true).await; + assert!( + run.shared + .violations + .lock() + .expect("violations") + .iter() + .any(|violation| violation.kind == "stuck_job"), + "a handler return must not stand in for a terminal commit" + ); + run.drop().await; + } + + #[tokio::test] + async fn sampled_terminal_pruned_before_observation_is_inconclusive() { + let mut run = TestRun::create().await; + run.store + .complete_runtime_batch(&run.pool, &[run.claim.clone()]) + .await + .expect("commit completion"); + run.store + .reset(&run.pool) + .await + .expect("remove retained evidence"); + run.oracle.final_check(true).await; + assert!( + run.shared.oracle_incomplete.load(Ordering::Relaxed), + "an unobserved pruned result cannot establish a pass or a lost-job claim" + ); + run.drop().await; + } + + #[tokio::test] + async fn observed_terminal_survives_pruning_in_the_client_ledger() { + let mut run = TestRun::create().await; + run.store + .complete_runtime_batch(&run.pool, &[run.claim.clone()]) + .await + .expect("commit completion"); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + assert!(run + .shared + .ledger + .get(run.claim.job.id) + .expect("sample record") + .has_durable_terminal()); + run.store + .reset(&run.pool) + .await + .expect("prune retained evidence"); + run.oracle.final_check(true).await; + assert!( + run.shared.violations.lock().expect("violations").is_empty(), + "captured durable proof must survive pruning" + ); + run.drop().await; + } + + #[tokio::test] + async fn eventful_observed_terminal_without_callback_survives_pruning() { + let mut run = TestRun::create_with(false, false).await; + assert!(!samples_terminal_proof(run.claim.job.id)); + let completed = run + .store + .complete_runtime_batch(&run.pool, &[run.claim.clone()]) + .await + .expect("commit completion without a lifecycle callback"); + assert_eq!(completed, vec![(run.claim.job.id, run.claim.job.run_lease)]); + let durable: bool = sqlx::query_scalar( + "SELECT EXISTS ( + SELECT 1 FROM awa.receipt_completion_batches AS batch, + LATERAL unnest(batch.job_ids, batch.run_leases) AS item(job_id, run_lease) + WHERE item.job_id = $1 AND item.run_lease = $2 + )", + ) + .bind(run.claim.job.id) + .bind(run.claim.job.run_lease) + .fetch_one(&run.pool) + .await + .expect("read committed completion independently of the observer"); + assert!(durable); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + let record = run + .shared + .ledger + .get(run.claim.job.id) + .expect("eventful record"); + assert!(record.has_durable_terminal()); + assert!(record.has(flag::DB_SEEN_COMPLETED)); + assert_eq!(record.completed_events, 0); + run.store + .reset(&run.pool) + .await + .expect("prune retained evidence"); + run.oracle.final_check(true).await; + assert!(run.shared.violations.lock().expect("violations").is_empty()); + run.drop().await; + } + + #[tokio::test] + async fn eventful_return_without_finalization_cannot_gain_terminal_proof() { + let mut run = TestRun::create_with(false, false).await; + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + assert!(!run + .shared + .ledger + .get(run.claim.job.id) + .expect("record") + .has_durable_terminal()); + run.oracle.final_check(true).await; + assert!(run + .shared + .violations + .lock() + .expect("violations") + .iter() + .any(|violation| violation.kind == "stuck_job")); + run.drop().await; + } + + #[tokio::test] + async fn eventful_terminal_pruned_before_observation_stays_inconclusive() { + let mut run = TestRun::create_with(false, false).await; + run.store + .complete_runtime_batch(&run.pool, &[run.claim.clone()]) + .await + .expect("commit completion"); + run.store + .reset(&run.pool) + .await + .expect("remove evidence before observation"); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + run.oracle.final_check(true).await; + assert!(run.shared.oracle_incomplete.load(Ordering::Relaxed)); + run.drop().await; + } + + #[tokio::test] + async fn eventful_terminal_observer_query_failure_stays_inconclusive() { + let run = TestRun::create_with(false, false).await; + run.store + .complete_runtime_batch(&run.pool, std::slice::from_ref(&run.claim)) + .await + .expect("commit completion"); + sqlx::query("ALTER TABLE awa.receipt_completion_batches RENAME TO receipt_completion_batches_unavailable") + .execute(&run.pool).await.expect("inject observer query failure"); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + assert!(run.shared.oracle_incomplete.load(Ordering::Relaxed)); + assert!(!run + .shared + .ledger + .get(run.claim.job.id) + .expect("record") + .has_durable_terminal()); + run.drop().await; + } + + #[tokio::test] + async fn eventful_terminal_observer_requires_the_exact_attempt() { + let run = TestRun::create_with(false, false).await; + run.store + .complete_runtime_batch(&run.pool, std::slice::from_ref(&run.claim)) + .await + .expect("complete old attempt"); + run.shared.ledger.with( + run.claim.job.id, + || unreachable!("registered"), + |record| { + record.latest_run_lease += 1; + }, + ); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + let record = run.shared.ledger.get(run.claim.job.id).expect("record"); + assert!(!record.has_durable_terminal()); + assert!(!record.has(flag::DB_SEEN_COMPLETED)); + run.drop().await; + } + + #[tokio::test] + async fn unsampled_no_events_job_does_not_gain_periodic_terminal_proof() { + let run = TestRun::create_with(true, false).await; + run.store + .complete_runtime_batch(&run.pool, std::slice::from_ref(&run.claim)) + .await + .expect("commit completion"); + super::super::terminal_proof::observe(&run.shared, &run.pool).await; + assert!(!run + .shared + .ledger + .get(run.claim.job.id) + .expect("record") + .has_durable_terminal()); + run.drop().await; + } + + #[tokio::test] + async fn failed_final_lookup_cannot_omit_unchecked_jobs_and_pass() { + let mut run = TestRun::create().await; + sqlx::query("ALTER TABLE awa.ready_entries RENAME TO ready_entries_unavailable") + .execute(&run.pool) + .await + .expect("inject a lookup failure"); + run.oracle.final_check(true).await; + assert!(run.shared.oracle_incomplete.load(Ordering::Relaxed)); + assert_eq!( + run.shared.run.get("final_lookup_errors"), + 1, + "the failed lookup chunk must be accounted for" + ); + run.drop().await; + } +} diff --git a/awa/tests/persona_soak/personas/abandon.rs b/awa/tests/persona_soak/personas/abandon.rs new file mode 100644 index 000000000..6ec671bf4 --- /dev/null +++ b/awa/tests/persona_soak/personas/abandon.rs @@ -0,0 +1,1036 @@ +//! Claim and abandon: a runtime that claims as much as it can, holds thousands of jobs +//! with deadlines and heartbeats, and is SIGKILLed again and again. +//! +//! The runtime is a separate process: the test binary started again with only the +//! `claim_abandon_child` test selected, so that a kill leaves its claims genuinely +//! abandoned (no shutdown handshake, connections dropped by the kernel). Between kills a +//! surviving in-process runtime and the other personas keep working on the same +//! database. After the last kill an in-process runtime takes over the abandoned queues +//! and finishes whatever the rescue returns to them. +//! +//! The checks are in `abandon_oracle.rs`: how long the rescue backlog takes to clear +//! after each kill, whether the claim ring keeps rotating, whether any job was rescued +//! while its owner was alive, and whether attempts and the dead-letter queue follow +//! `max_attempts`. + +use super::super::abandon_oracle::{ + epoch_ms, probe_backlog, Attempt, Kill, LiveCheck, Scope, FINISHER, +}; +use super::super::config::{Persona, DEFAULT_HEARTBEAT_SECONDS, DEFAULT_STALENESS_SECONDS}; +use super::super::db::{connect_pool, connect_pool_with_timeout, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag, JobRecord}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{backoff, queue, register_rows, timings, tune}; +use awa::model::{insert_many, insert_many_copy_from_pool}; +use awa::{ + Client, InsertOpts, InsertParams, JobArgs, JobContext, JobError, JobResult, QueueConfig, Worker, +}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::io::{Read, Seek, SeekFrom, Write}; +use std::path::PathBuf; +use std::process::Stdio; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; +use tokio::io::{AsyncBufReadExt, BufReader}; +use tokio::process::{Child, Command}; + +/// Zero-deadline claims: written as compact claim batches. +pub const QUEUE_PLAIN: &str = "awa_soak_abandon_plain"; +/// Claims with a deadline: written as one claim row per job. +pub const QUEUE_DEADLINE: &str = "awa_soak_abandon_deadline"; +/// Short jobs served by the survivors, to see claim latency while the backlog drains. +pub const QUEUE_CANARY: &str = "awa_soak_abandon_canary"; + +/// The name of the test the child process runs. +pub const CHILD_TEST: &str = "claim_abandon_child"; +const ENV_URL: &str = "AWA_SOAK_ABANDON_URL"; +const ENV_GENERATION: &str = "AWA_SOAK_ABANDON_GENERATION"; +const ENV_LEDGER: &str = "AWA_SOAK_ABANDON_LEDGER"; +const ENV_WORKERS: &str = "AWA_SOAK_ABANDON_WORKERS"; +const ENV_POOL: &str = "AWA_SOAK_ABANDON_POOL"; +const ENV_HEARTBEAT: &str = "AWA_SOAK_HEARTBEAT_SECONDS"; +const ENV_STALENESS: &str = "AWA_SOAK_STALENESS_SECONDS"; + +/// Longer than any process lives, so a job is only ever finished by the survivors. +const HOLD: Duration = Duration::from_secs(20 * 60); +const DEADLINE: Duration = Duration::from_secs(15 * 60); +const CHILD_POOL: u32 = 16; +const FINISHER_REPLICAS: usize = 2; +const FINISHER_POOL: u32 = 16; +/// Percentage of jobs that may only be attempted twice, so that repeated kills +/// dead-letter them. +const SHORT_RETRY_PERCENT: u64 = 25; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct AbandonJob { + pub seq: u64, + pub cycle: u32, +} + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct AbandonCanaryJob { + pub seq: u64, +} + +/// The worker in the child process: records every start in the child ledger and then +/// holds the job until the process is killed. +struct ChildWorker { + generation: u32, + ledger: Mutex, + active: Arc, + peak: Arc, +} + +#[async_trait::async_trait] +impl Worker for ChildWorker { + fn kind(&self) -> &'static str { + "abandon_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let line = format!( + "{{\"e\":\"start\",\"id\":{},\"attempt\":{},\"gen\":{},\"q\":\"{}\",\"t\":{}}}\n", + ctx.job.id, + ctx.job.attempt, + self.generation, + ctx.job.queue, + epoch_ms() + ); + // One write per start, straight to the file: nothing is buffered that a kill could + // lose, and the parent never sees a half line it cannot skip. + if let Ok(mut file) = self.ledger.lock() { + let _ = file.write_all(line.as_bytes()); + } + let now = self.active.fetch_add(1, Ordering::Relaxed) + 1; + self.peak.fetch_max(now, Ordering::Relaxed); + tokio::time::sleep(HOLD).await; + self.active.fetch_sub(1, Ordering::Relaxed); + Ok(JobResult::Completed) + } +} + +/// Entry point of the child process; see [`CHILD_TEST`]. Returns at once when the +/// environment is not set up by the parent, so a plain `--ignored` run is unaffected. +pub async fn run_child() { + let (Ok(url), Ok(generation), Ok(ledger), Ok(workers)) = ( + std::env::var(ENV_URL), + std::env::var(ENV_GENERATION), + std::env::var(ENV_LEDGER), + std::env::var(ENV_WORKERS), + ) else { + return; + }; + super::super::worker_errors::install(true).expect("child worker error tracing"); + let generation: u32 = generation.parse().expect("generation"); + let workers: u32 = workers.parse().expect("workers"); + let pool_size: u32 = std::env::var(ENV_POOL) + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(CHILD_POOL); + + // Exit with the harness, so a harness that died leaves no process behind. + let parent = std::os::unix::process::parent_id(); + std::thread::spawn(move || loop { + std::thread::sleep(Duration::from_secs(1)); + if std::os::unix::process::parent_id() != parent { + std::process::exit(0); + } + }); + + let pool = connect_pool(&url, &format!("soak_abandon_{generation}"), pool_size) + .await + .expect("child connect"); + let file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(&ledger) + .expect("child ledger"); + let active = Arc::new(AtomicI64::new(0)); + let peak = Arc::new(AtomicI64::new(0)); + let claiming = |deadline_duration: Duration| QueueConfig { + max_workers: workers, + poll_interval: Duration::from_millis(50), + deadline_duration, + claimers: 2, + claim_batch_size: 1_000, + ..QueueConfig::default() + }; + // The child keeps away from maintenance leadership: the survivors do the rescuing, as + // the other replicas of an application would. + let seconds = |name: &str, default: u64| { + Duration::from_secs( + std::env::var(name) + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(default), + ) + }; + let client = timings( + Client::builder(pool) + .queue(QUEUE_PLAIN, claiming(Duration::ZERO)) + .queue(QUEUE_DEADLINE, claiming(DEADLINE)) + .register_worker(ChildWorker { + generation, + ledger: Mutex::new(file), + active: active.clone(), + peak: peak.clone(), + }), + seconds(ENV_HEARTBEAT, DEFAULT_HEARTBEAT_SECONDS), + seconds(ENV_STALENESS, DEFAULT_STALENESS_SECONDS), + ) + .leader_election_interval(Duration::from_secs(3600)) + .build() + .expect("build child client"); + client.start().await.expect("start child client"); + println!("READY gen={generation}"); + loop { + tokio::time::sleep(Duration::from_secs(2)).await; + println!( + "STAT gen={generation} active={} peak={}", + active.load(Ordering::Relaxed), + peak.load(Ordering::Relaxed) + ); + } +} + +/// Finishes the abandoned jobs once the child processes are gone. +struct FinisherWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for FinisherWorker { + fn kind(&self) -> &'static str { + "abandon_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self + .shared + .begin_execution(Persona::ClaimAbandon, &ctx.job, true); + self.shared.abandon.note_attempt( + ctx.job.id, + Attempt { + attempt: ctx.job.attempt, + owner: FINISHER, + at_ms: epoch_ms(), + }, + ); + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +struct CanaryWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for CanaryWorker { + fn kind(&self) -> &'static str { + "abandon_canary_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self + .shared + .begin_execution(Persona::ClaimAbandon, &ctx.job, true); + let waited = (chrono::Utc::now() - ctx.job.created_at).num_milliseconds(); + self.shared + .stats(Persona::ClaimAbandon) + .observe("canary_wait_ms", waited); + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +/// Copies the running test binary into the report directory. Every process is started +/// from the copy, so a rebuild of the binary during a long run cannot change what is +/// started, or leave nothing to start. +fn copy_executable(directory: &std::path::Path) -> Result { + let copy = directory.join("claim-abandon-child"); + let _ = std::fs::remove_file(©); + let source = std::path::Path::new("/proc/self/exe"); + let source = if source.exists() { + source.to_path_buf() + } else { + std::env::current_exe().map_err(|error| format!("current_exe: {error}"))? + }; + std::fs::copy(&source, ©) + .map_err(|error| format!("copy {} to {}: {error}", source.display(), copy.display()))?; + Ok(copy) +} + +fn spawn_child( + shared: &Arc, + executable: &PathBuf, + generation: u32, + ledger: &PathBuf, + workers_per_queue: u32, + pool_size: u32, +) -> Result { + let url = &shared.config.database_url; + let mut child = Command::new(executable) + .args([ + "--ignored", + "--exact", + CHILD_TEST, + "--nocapture", + "--test-threads=1", + ]) + .env(ENV_URL, url) + .env(ENV_GENERATION, generation.to_string()) + .env(ENV_LEDGER, ledger) + .env(ENV_WORKERS, workers_per_queue.to_string()) + .env(ENV_POOL, pool_size.to_string()) + .env( + ENV_HEARTBEAT, + shared.config.heartbeat_interval.as_secs().to_string(), + ) + .env( + ENV_STALENESS, + shared.config.heartbeat_staleness.as_secs().to_string(), + ) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn() + .map_err(|error| format!("spawn the claiming process: {error}"))?; + for stream in [ + child + .stdout + .take() + .map(|s| Box::new(s) as Box), + child + .stderr + .take() + .map(|s| Box::new(s) as Box), + ] + .into_iter() + .flatten() + { + let shared = shared.clone(); + let output_guard = super::super::worker_errors::child_output_guard(); + tokio::spawn(async move { + let _output_guard = output_guard; + let mut lines = BufReader::new(stream).lines(); + loop { + match lines.next_line().await { + Ok(Some(line)) => absorb_output(&shared, &line), + Ok(None) => break, + Err(error) => { + shared.incomplete( + format!("claiming-child output read failed: {error}"), + vec![], + ); + break; + } + } + } + }); + } + Ok(child) +} + +fn absorb_output(shared: &Shared, line: &str) { + if super::super::worker_errors::absorb_child_marker(line) { + return; + } + let stats = shared.stats(Persona::ClaimAbandon); + let field = |name: &str| { + line.split_whitespace() + .find_map(|part| part.strip_prefix(&format!("{name}="))) + .and_then(|value| value.parse::().ok()) + }; + if line.starts_with("STAT ") { + if let Some(active) = field("active") { + shared.abandon.child_active.store(active, Ordering::Relaxed); + } + if let Some(peak) = field("peak") { + stats.set_max("process_peak_held", peak); + } + } else if line.contains("panicked") || line.contains("Error") { + if stats.get("process_error_lines") < 5 { + shared.note("abandon", line.chars().take(200).collect::()); + } + stats.add("process_error_lines", 1); + } +} + +/// Folds the new lines of the child ledger into the harness ledger. +struct Importer { + path: PathBuf, + offset: u64, +} + +impl Importer { + fn import(&mut self, shared: &Shared) { + let Ok(mut file) = std::fs::File::open(&self.path) else { + return; + }; + if file.seek(SeekFrom::Start(self.offset)).is_err() { + return; + } + let mut text = String::new(); + if file.read_to_string(&mut text).is_err() { + return; + } + for line in text.split_inclusive('\n') { + if !line.ends_with('\n') { + break; + } + self.offset += line.len() as u64; + if let Ok(event) = serde_json::from_str::(line) { + self.apply(shared, &event); + } + } + } + + fn apply(&self, shared: &Shared, event: &Value) { + if event.get("e").and_then(Value::as_str) != Some("start") { + return; + } + let (Some(id), Some(attempt), Some(generation), Some(at_ms)) = ( + event.get("id").and_then(Value::as_i64), + event.get("attempt").and_then(Value::as_i64), + event.get("gen").and_then(Value::as_u64), + event.get("t").and_then(Value::as_i64), + ) else { + return; + }; + let queue = event + .get("q") + .and_then(Value::as_str) + .and_then(|name| shared.queues.get(name)); + let queue_index = queue.as_ref().map_or(u32::MAX, |queue| queue.index); + if let Some(queue) = &queue { + queue.started.fetch_add(1, Ordering::Relaxed); + } + let now = shared.now_ms() as u32; + shared.ledger.with( + id, + || { + JobRecord::new( + Persona::ClaimAbandon, + queue_index, + now, + expect::COMPLETED | expect::FAILED, + flag::NO_EVENTS, + ) + }, + |record| { + record.executions = record.executions.saturating_add(1); + record.max_attempt = record.max_attempt.max(attempt as i16); + if record.first_start_ms == 0 { + record.first_start_ms = now.max(1); + } + record.last_start_ms = now; + }, + ); + shared.abandon.note_attempt( + id, + Attempt { + attempt: attempt as i16, + owner: generation as u32, + at_ms, + }, + ); + shared + .stats(Persona::ClaimAbandon) + .add("process_executions", 1); + } +} + +fn params(rng: &mut Rng, seq: &mut u64, cycle: u32, count: usize) -> Vec { + (0..count) + .map(|index| { + *seq += 1; + awa::insert::params_with( + &AbandonJob { seq: *seq, cycle }, + InsertOpts { + queue: if index % 2 == 0 { + QUEUE_PLAIN + } else { + QUEUE_DEADLINE + } + .to_string(), + max_attempts: if rng.chance(SHORT_RETRY_PERCENT) { + 2 + } else { + 25 + }, + ..Default::default() + }, + ) + .expect("abandon params") + }) + .collect() +} + +/// Enqueues `count` jobs for the next process to claim, in chunks that COPY handles +/// comfortably. +async fn enqueue_cycle( + shared: &Arc, + pool: &sqlx::PgPool, + rng: &mut Rng, + seq: &mut u64, + cycle: u32, + count: usize, +) -> Result { + let mut enqueued = 0; + while enqueued < count { + let chunk = (count - enqueued).min(10_000); + let batch = params(rng, seq, cycle, chunk); + let rows = insert_many_copy_from_pool(pool, &batch) + .await + .map_err(|error| format!("enqueue: {error}"))?; + for row in &rows { + shared.abandon.register_job(row.id, row.max_attempts); + if let Some(queue) = shared.queues.get(&row.queue) { + queue.enqueued.fetch_add(1, Ordering::Relaxed); + } + } + register_rows( + shared, + Persona::ClaimAbandon, + &rows, + |_| expect::COMPLETED | expect::FAILED, + flag::NO_EVENTS, + ); + shared + .stats(Persona::ClaimAbandon) + .add("enqueued", rows.len() as i64); + enqueued += rows.len(); + } + Ok(enqueued) +} + +async fn canary_producer(shared: Arc, pool: sqlx::PgPool) { + let mut seq = 0u64; + while !shared.load_stop.sleep(Duration::from_millis(200)).await { + let batch: Vec = (0..2) + .map(|_| { + seq += 1; + awa::insert::params_with( + &AbandonCanaryJob { seq }, + InsertOpts { + queue: QUEUE_CANARY.to_string(), + max_attempts: 5, + ..Default::default() + }, + ) + .expect("canary params") + }) + .collect(); + match insert_many(&pool, &batch).await { + Ok(rows) => { + shared + .stats(Persona::ClaimAbandon) + .add("canary_enqueued", rows.len() as i64); + register_rows( + &shared, + Persona::ClaimAbandon, + &rows, + |_| expect::COMPLETED, + flag::NO_EVENTS, + ); + } + Err(error) => backoff(&shared, error, Persona::ClaimAbandon).await, + } + } +} + +/// Measures the live process's heartbeat passes from the timestamps its heartbeats write. +/// A pass updates every in-flight claim with the time of its own chunk, so once a pass has +/// finished the oldest heartbeat is newer than anything the previous pass wrote, and the +/// spread from oldest to newest is the pass's duration (to within the last chunk). Only +/// claims taken since the process started are read, so the dead processes' open claims, +/// whose heartbeats stopped, do not count. They are the deadline queue's, which are +/// interleaved with the other queue's in the pass's job-id order and so span the whole pass. +async fn sample_heartbeats( + shared: &Arc, + pool: &sqlx::PgPool, + schema: &str, + since_epoch_ms: i64, +) { + let row: Result<(i64, Option, Option, f64), _> = sqlx::query_as(&format!( + "SELECT count(*)::bigint, + EXTRACT(EPOCH FROM min(a.heartbeat_at))::float8, + EXTRACT(EPOCH FROM max(a.heartbeat_at))::float8, + EXTRACT(EPOCH FROM clock_timestamp())::float8 + FROM {schema}.lease_claims c + JOIN {schema}.attempt_state a ON a.job_id = c.job_id AND a.run_lease = c.run_lease + WHERE c.queue = $1 AND c.closed_at IS NULL + AND c.claimed_at >= to_timestamp($2::float8 / 1000.0)" + )) + .bind(QUEUE_DEADLINE) + .bind(since_epoch_ms as f64) + .fetch_one(pool) + .await; + let Ok((rows, Some(oldest), Some(newest), now)) = row else { + return; + }; + let active = shared.abandon.child_active.load(Ordering::Relaxed); + shared + .abandon + .note_heartbeat_sample(shared, rows, oldest, newest, now, active); +} + +/// Probes the backlog every few seconds and the claim ring every second, until the +/// run is over. +async fn monitor(shared: Arc, queues: Vec) { + let schema = shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".to_string()); + let Ok(pool) = connect_pool_with_timeout( + &shared.config.database_url, + "soak_abandon_monitor", + 2, + Some(Duration::from_secs(60)), + ) + .await + else { + shared.violation( + "persona_failed", + "abandon monitor could not connect", + vec![], + ); + return; + }; + let staleness_secs = shared.config.heartbeat_staleness.as_secs_f64(); + let mut tick = 0u64; + let mut lock_captures = super::super::ring_locks::CaptureBudget::default(); + while !shared.drain_stop.sleep(Duration::from_secs(1)).await { + tick += 1; + let ring: Result<(i32, i64, i32), _> = sqlx::query_as(&format!( + "SELECT current_slot, generation, slot_count FROM {schema}.claim_ring_state WHERE singleton" + )) + .fetch_one(&pool) + .await; + let (slot, ring_generation, slot_count) = match ring { + Ok(ring) => ring, + Err(error) => { + shared.incomplete(format!("claim-ring monitor query failed: {error}"), vec![]); + continue; + } + }; + if let Some((_, spawned)) = shared.abandon.live() { + sample_heartbeats(&shared, &pool, &schema, spawned).await; + } + // While the ring is still, look at the slot it has to enter next: live claims + // there explain the stall. + let next_slot = if shared.abandon.ring_stalled() && tick.is_multiple_of(2) { + probe_backlog( + &pool, + &schema, + Scope::Slot((slot + 1).rem_euclid(slot_count)), + staleness_secs, + 0, + ) + .await + .map(|probe| (probe.open - probe.expired, probe.expired)) + .map_err(|error| { + shared.incomplete(format!("claim-slot monitor query failed: {error}"), vec![]) + }) + .ok() + } else { + None + }; + shared + .abandon + .observe_ring(&shared, slot, ring_generation, next_slot); + super::super::ring_locks::observe( + &shared, + &pool, + &schema, + slot_count, + next_slot, + &mut lock_captures, + ) + .await; + if tick.is_multiple_of(3) { + // Read what the live process has started before the probe, so a claim started + // meanwhile cannot make a claim look missing. + let live = shared.abandon.live(); + let live_check = live.map(|(generation, _)| LiveCheck { + generation, + started: shared.abandon.started_by(generation), + }); + let started = Instant::now(); + match probe_backlog( + &pool, + &schema, + Scope::Queues(&queues), + staleness_secs, + live.map_or(0, |(_, spawned)| spawned), + ) + .await + { + Ok(probe) => shared.abandon.observe( + &shared, + &probe, + live_check, + ring_generation, + started.elapsed().as_millis() as i64, + ), + Err(error) => { + shared.incomplete( + format!("abandoned-claim backlog query failed: {error}"), + vec![], + ); + shared.stats(Persona::ClaimAbandon).add("probe_errors", 1); + if shared.stats(Persona::ClaimAbandon).get("probe_errors") <= 3 { + shared.note("abandon", format!("backlog probe failed: {error}")); + } + } + } + } + } + shared.abandon.finish_ring(&shared); + pool.close().await; +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::ClaimAbandon; + let config = shared.config.clone(); + let staleness_secs = config.heartbeat_staleness.as_secs_f64(); + let plain = shared.queues.register(QUEUE_PLAIN, persona, false); + let deadline = shared.queues.register(QUEUE_DEADLINE, persona, false); + let canary = shared.queues.register(QUEUE_CANARY, persona, false); + let queues = vec![QUEUE_PLAIN.to_string(), QUEUE_DEADLINE.to_string()]; + if config.abandon_kill_every >= DEADLINE { + return Err(format!( + "AWA_SOAK_ABANDON_KILL_SECONDS must stay below the {}s job deadline, or the owner's own deadline rescue would count as a false rescue", + DEADLINE.as_secs() + )); + } + let target = config.abandon_claims.max(2); + let workers_per_queue = (target / 2).max(1) as u32; + + std::fs::create_dir_all(&config.report_dir).map_err(|error| error.to_string())?; + let ledger_path = config.report_dir.join("abandon-child-ledger.jsonl"); + let _ = std::fs::remove_file(&ledger_path); + let executable = copy_executable(&config.report_dir)?; + shared + .abandon + .set_csv(config.report_dir.join("abandon-backlog.csv")); + + let child_pool = shared.budget.reserve("abandon-child", CHILD_POOL, 4); + // Held for the runtime that takes over at the end; Fleet::start replaces it. + shared.budget.reserve( + "abandon-finisher", + FINISHER_POOL * FINISHER_REPLICAS as u32, + FINISHER_REPLICAS as u32 * 2, + ); + + // The survivors: other replicas of the application. They lead maintenance, so they do + // the rescuing, and they serve a canary queue. + let canary_shared = shared.clone(); + let survivor_factory: ClientFactory = Arc::new(move |pool| { + tune( + Client::builder(pool) + .queue(QUEUE_CANARY, queue(4)) + .register_worker(CanaryWorker { + shared: canary_shared.clone(), + }), + &canary_shared, + ) + .build() + .expect("build abandon survivor client") + }); + let survivors = Fleet::start( + &shared, + "abandon-survivor", + persona, + 2, + 8, + 2, + vec![(canary, 4)], + survivor_factory, + ) + .await?; + + let producer_pool = connect_pool( + &config.database_url, + "soak_abandon_producer", + shared.budget.reserve("abandon-producer", 4, 2), + ) + .await + .map_err(|error| error.to_string())?; + let canary_task = tokio::spawn(canary_producer(shared.clone(), producer_pool.clone())); + let monitor_task = tokio::spawn(monitor(shared.clone(), queues.clone())); + let backlog_pool = connect_pool_with_timeout( + &config.database_url, + "soak_abandon_control", + 2, + Some(Duration::from_secs(60)), + ) + .await + .map_err(|error| error.to_string())?; + + let mut rng = Rng::for_task(config.seed, "abandon"); + let mut importer = Importer { + path: ledger_path.clone(), + offset: 0, + }; + let mut seq = 0u64; + let schema = shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".to_string()); + shared.note( + "abandon", + format!( + "armed: {} kill(s) every {}s, each leaving about {target} open claims ({} plain, {} with a deadline); rescue must clear within {}s", + config.abandon_kills, + config.abandon_kill_every.as_secs(), + target / 2, + target - target / 2, + config.abandon_bound.as_secs() + ), + ); + + let mut generation = 0u32; + let mut load_ended = false; + for number in 1..=config.abandon_kills { + if shared.load_stop.is_stopped() { + break; + } + // The first process finds a full queue. Later ones find the rescued jobs of the + // earlier kills and only need fresh jobs to top their capacity up, so that jobs + // are run, and abandoned, more than once. + let fresh = if number == 1 { target } else { target / 2 }; + match enqueue_cycle( + &shared, + &producer_pool, + &mut rng, + &mut seq, + number, + fresh as usize, + ) + .await + { + Ok(count) => shared.note( + "abandon", + format!("enqueued {count} fresh job(s) for process {number}"), + ), + Err(error) => { + backoff(&shared, error, persona).await; + continue; + } + } + generation += 1; + let mut child = spawn_child( + &shared, + &executable, + generation, + &ledger_path, + workers_per_queue, + child_pool, + )?; + shared.abandon.child_alive.store(true, Ordering::Relaxed); + shared.abandon.child_active.store(0, Ordering::Relaxed); + shared + .abandon + .set_live(Some((generation, epoch_ms() - 1_000))); + let spawned = Instant::now(); + let mut ramped = false; + shared.note( + "abandon", + format!("process {generation} started with {workers_per_queue} workers on each queue"), + ); + let mut unexpected_exit = false; + loop { + importer.import(&shared); + if shared.load_stop.sleep(Duration::from_secs(1)).await { + load_ended = true; + break; + } + if let Ok(Some(status)) = child.try_wait() { + shared.violation( + "abandon_process_exited", + format!("the claiming process {generation} exited on its own: {status}"), + vec![], + ); + unexpected_exit = true; + break; + } + if !ramped + && shared.abandon.child_active.load(Ordering::Relaxed) >= (target as i64) * 95 / 100 + { + ramped = true; + shared.note( + "abandon", + format!( + "process {generation} holds {} jobs after {:.0}s", + shared.abandon.child_active.load(Ordering::Relaxed), + spawned.elapsed().as_secs_f64() + ), + ); + shared + .stats(persona) + .observe("ramp_seconds", spawned.elapsed().as_secs() as i64); + } + if spawned.elapsed() >= config.abandon_kill_every { + break; + } + } + + let held = shared.abandon.child_active.load(Ordering::Relaxed); + let signalled_epoch_ms = if unexpected_exit { + let at = epoch_ms(); + shared.abandon.note_generation_killed(generation, at); + shared.abandon.set_live(None); + at + } else { + // SIGKILL: no shutdown handshake, no chance to release a claim. + match shared + .abandon + .signal_generation_killed(generation, || child.start_kill()) + { + Ok(at) => at, + Err(error) => { + shared.incomplete( + format!("could not signal claiming process {generation}: {error}"), + vec![], + ); + // No confirmed signal means no valid death boundary. Stop this + // persona and clean up its owned tasks instead of inventing one. + shared.abandon.set_live(None); + shared.abandon.child_alive.store(false, Ordering::Relaxed); + shared.abandon.child_active.store(0, Ordering::Relaxed); + drop(child); // kill_on_drop remains a best-effort cleanup. + canary_task.abort(); + monitor_task.abort(); + let _ = canary_task.await; + let _ = monitor_task.await; + survivors + .shutdown_all(Duration::from_secs(20), &shared) + .await; + producer_pool.close().await; + backlog_pool.close().await; + return Ok(()); + } + } + }; + let signalled_run_ms = shared.now_ms(); + shared.abandon.child_alive.store(false, Ordering::Relaxed); + shared.abandon.child_active.store(0, Ordering::Relaxed); + if !unexpected_exit { + if let Err(error) = child.wait().await { + shared.incomplete( + format!( + "could not reap claiming process {generation} after signaling it: {error}" + ), + vec![], + ); + } + } + importer.import(&shared); + let open_after_kill = probe_backlog( + &backlog_pool, + &schema, + Scope::Queues(&queues), + staleness_secs, + 0, + ) + .await + .map(|probe| probe.open) + .unwrap_or_else(|error| { + shared.incomplete(format!("post-kill backlog query failed: {error}"), vec![]); + -1 + }); + let started = shared.abandon.started_by(generation); + shared.abandon.record_kill(Kill { + number, + generation, + at_run_ms: signalled_run_ms, + signalled_epoch_ms, + started_by_process: started, + rescued_while_alive: shared.abandon.live_lost(generation), + open_after_kill, + expired_peak: 0, + oldest_expired_peak_s: 0.0, + cleared_after_s: None, + reported_slow: false, + load_ended, + }); + shared.stats(persona).add("kills", 1); + shared.stats(persona).add("claims_abandoned", started); + shared.note( + "abandon", + format!( + "SIGKILL process {generation} (kill {number}): {started} handlers started, {held} running at the last report, {open_after_kill} open claims on its queues" + ), + ); + if load_ended { + break; + } + if shared.load_stop.sleep(config.abandon_restart_delay).await { + break; + } + } + + // Everything the kills left behind is finished by an in-process runtime that is + // never killed. + let finisher_shared = shared.clone(); + let finisher_factory: ClientFactory = Arc::new(move |pool| { + tune( + Client::builder(pool) + .queue(QUEUE_PLAIN, queue(256)) + .queue(QUEUE_DEADLINE, queue(256)) + .register_worker(FinisherWorker { + shared: finisher_shared.clone(), + }), + &finisher_shared, + ) + .build() + .expect("build abandon finisher client") + }); + shared.note( + "abandon", + "kills over; starting the runtime that finishes the rescued jobs", + ); + let finisher = Fleet::start( + &shared, + "abandon-finisher", + persona, + FINISHER_REPLICAS, + FINISHER_POOL, + 4, + vec![(plain, 256), (deadline, 256)], + finisher_factory, + ) + .await?; + + while !shared.drain_stop.sleep(Duration::from_secs(1)).await { + importer.import(&shared); + } + importer.import(&shared); + canary_task.abort(); + let _ = monitor_task.await; + shared + .abandon + .analyse(&shared, &backlog_pool, &queues) + .await; + finisher + .shutdown_all(Duration::from_secs(20), &shared) + .await; + survivors + .shutdown_all(Duration::from_secs(20), &shared) + .await; + producer_pool.close().await; + backlog_pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/alex.rs b/awa/tests/persona_soak/personas/alex.rs new file mode 100644 index 000000000..88498cf04 --- /dev/null +++ b/awa/tests/persona_soak/personas/alex.rs @@ -0,0 +1,756 @@ +//! Alex, the operator: under live load, pauses and resumes queues, drains, cancels, +//! moves and reprioritises jobs, retries and purges DLQ rows, runs the admin read +//! queries the UI issues, and rolls a client restart through another persona's fleet. + +use super::super::config::Persona; +use super::super::db::{connect_pool, connect_pool_with_timeout, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag}; +use super::super::registry::QueueInfo; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{attach_events, backoff, queue, register_rows, tune}; +use awa::model::admin::{self, ListJobsFilter}; +use awa::model::batch_operations::{ + get_batch_operation, submit_batch_operation, BatchOperationFilter, BatchOperationSpec, + BatchOperationState, SubmitBatchOperation, +}; +use awa::model::dlq::{self, ListDlqFilter}; +use awa::model::{insert_many, JobState, RetryFromDlqOpts}; +use awa::{Client, InsertOpts, JobArgs, JobContext, JobError, JobResult, Worker}; +use serde::{Deserialize, Serialize}; +use sqlx::PgPool; +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +pub const QUEUE_A: &str = "awa_soak_alex_a"; +pub const QUEUE_B: &str = "awa_soak_alex_b"; +pub const QUEUE_DRAIN: &str = "awa_soak_alex_drain"; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct AlexJob { + pub seq: u64, + pub work_ms: u64, +} + +struct AlexWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for AlexWorker { + fn kind(&self) -> &'static str { + "alex_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self.shared.begin_execution(Persona::Alex, &ctx.job, false); + let work_ms = ctx.job.args["work_ms"].as_u64().unwrap_or(100); + tokio::time::sleep(Duration::from_millis(work_ms)).await; + if ctx.is_cancelled() { + return Err(JobError::retryable_msg("interrupted")); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn client_factory(shared: Arc) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(QUEUE_A, queue(4)) + .queue(QUEUE_B, queue(4)) + .queue(QUEUE_DRAIN, queue(4)) + .register_worker(AlexWorker { + shared: shared.clone(), + }); + attach_events(tune(builder, &shared), &shared, &["alex_job"]) + .build() + .expect("build alex client") + }) +} + +struct Context { + shared: Arc, + pool: PgPool, + diagnostics: PgPool, + stalled_rows: std::sync::Mutex>, + rng: Rng, +} + +/// pid, application, state, wait event, seconds running, blocking pids, query. +type BackendRow = (i32, String, String, String, i32, String, String); + +/// An admin call this slow under load is a finding in its own right. +const SLOW_CALL: Duration = Duration::from_secs(30); +/// Server-side limit on the operator's statements; hitting it is reported as a finding. +const ADMIN_STATEMENT_TIMEOUT: Duration = Duration::from_secs(150); +const STALLED_CALL: Duration = Duration::from_secs(120); + +impl Context { + fn stats(&self) -> &super::super::shared::Stats { + self.shared.stats(Persona::Alex) + } + + async fn timed( + &self, + name: &str, + call: impl std::future::Future>, + ) -> Option { + let started = Instant::now(); + let mut call = std::pin::pin!(call); + let result = tokio::select! { + result = &mut call => result, + _ = tokio::time::sleep(SLOW_CALL) => { + self.diagnose_slow_call(name).await; + call.await + } + }; + let ms = started.elapsed().as_millis() as i64; + if started.elapsed() > STALLED_CALL { + self.shared.finding( + "admin_call_stalled", + format!("admin call {name} took {ms} ms under load"), + self.stalled_rows.lock().expect("stalled rows").clone(), + ); + } + self.stats().observe(name, ms); + if ms > 5_000 { + self.stats().add("slow_admin_calls_over_5s", 1); + self.shared.note("alex", format!("{name} took {ms} ms")); + } + match result { + Ok(value) => Some(value), + Err(error) => { + if error.to_string().contains("statement timeout") { + self.shared.finding( + "admin_call_timeout", + format!( + "admin call {name} was cancelled by the {}s statement timeout", + ADMIN_STATEMENT_TIMEOUT.as_secs() + ), + self.stalled_rows.lock().expect("stalled rows").clone(), + ); + } + self.stats().add(&format!("{name}_errors"), 1); + if self.stats().get(&format!("{name}_errors")) <= 3 { + self.shared.note("alex", format!("{name} failed: {error}")); + } + None + } + } + } + + /// Records what the operator's backend is waiting on and who blocks it. + async fn diagnose_slow_call(&self, name: &str) { + let rows: Result, _> = sqlx::query_as( + "SELECT a.pid, a.application_name, COALESCE(a.state, ''), + COALESCE(a.wait_event_type, '') || ':' || COALESCE(a.wait_event, ''), + COALESCE(extract(epoch FROM now() - a.query_start)::int, 0), + pg_blocking_pids(a.pid)::text, left(a.query, 200) + FROM pg_stat_activity a + WHERE a.pid <> pg_backend_pid() + AND (a.application_name = 'soak_alex_operator' + OR a.pid IN (SELECT unnest(pg_blocking_pids(op.pid)) FROM pg_stat_activity op + WHERE op.application_name = 'soak_alex_operator'))", + ) + .fetch_all(&self.diagnostics) + .await; + let lines: Vec = match rows { + Ok(rows) => rows + .into_iter() + .map(|(pid, app, state, wait, secs, blockers, query)| { + format!("{{\"pid\":{pid},\"application\":\"{app}\",\"state\":\"{state}\",\"wait\":\"{wait}\",\"running_s\":{secs},\"blocked_by\":\"{blockers}\",\"query\":{query:?}}}") + }) + .collect(), + Err(error) => vec![format!("diagnostic query failed: {error}")], + }; + self.shared.note( + "alex", + format!( + "{name} still running after {}s; backend state: {}", + SLOW_CALL.as_secs(), + lines.join(" | ") + ), + ); + *self.stalled_rows.lock().expect("stalled rows") = lines; + } + + /// A random queue of a persona, other than Alex's own. + fn queue_of(&mut self, persona: Persona) -> Option> { + let queues = self.shared.queues.of(persona); + self.rng.pick(&queues).cloned() + } + + async fn admin_reads(&mut self) { + let pool = self.pool.clone(); + self.timed("admin_queue_overviews", admin::queue_overviews(&pool)) + .await; + self.timed("admin_state_counts", admin::state_counts(&pool)) + .await; + self.timed("admin_job_kind_overviews", admin::job_kind_overviews(&pool)) + .await; + self.timed("admin_runtime_overview", admin::runtime_overview(&pool)) + .await; + self.timed( + "admin_queue_runtime_summary", + admin::queue_runtime_summary(&pool), + ) + .await; + self.timed("admin_dlq_depth_by_queue", dlq::dlq_depth_by_queue(&pool)) + .await; + if let Some(queue) = self.queue_of(Persona::Tom) { + let filter = ListJobsFilter { + queue: Some(queue.name.clone()), + limit: Some(50), + ..Default::default() + }; + self.timed("admin_list_jobs", admin::list_jobs(&pool, &filter)) + .await; + } + } + + /// Resumes a queue, retrying: a queue left paused would look like a stall. + async fn resume_with_retry(&self, queue: &str) { + for attempt in 0..8 { + if self + .timed("resume_queue", admin::resume_queue(&self.pool, queue)) + .await + .is_some() + { + return; + } + tokio::time::sleep(Duration::from_secs(2 + attempt)).await; + } + self.shared.violation( + "resume_failed", + format!("queue {queue} could not be resumed after eight attempts"), + vec![], + ); + } + + async fn pause_resume(&mut self) { + let target = match self.rng.below(3) { + 0 => self.shared.queues.get(QUEUE_B), + 1 => self.queue_of(Persona::Tom), + _ => self + .shared + .queues + .get(super::priya::HOT_QUEUE) + .or_else(|| self.queue_of(Persona::Mia)), + }; + let Some(queue) = target else { return }; + queue.hold(); + let pool = self.pool.clone(); + let paused = self + .timed( + "pause_queue", + admin::pause_queue(&pool, &queue.name, Some("soak")), + ) + .await + .is_some(); + if paused { + self.stats().add("pauses", 1); + self.shared + .stop_sleep(Duration::from_secs(self.rng.range(3, 8))) + .await; + } + // A pause whose answer was lost may still have applied, so always resume. + self.resume_with_retry(&queue.name).await; + // Give the claimers a moment to pick the queue back up before stall checks resume. + self.shared.stop_sleep(Duration::from_secs(3)).await; + queue.release(); + } + + async fn drain(&mut self) { + let target = if self.rng.chance(50) { + self.shared.queues.get(QUEUE_DRAIN) + } else { + self.queue_of(Persona::Tom) + }; + let Some(queue) = target else { return }; + // Until the drain returns, any job enqueued so far may legitimately be cancelled. + queue.drained_at_ms.store(i64::MAX / 4, Ordering::Relaxed); + let pool = self.pool.clone(); + if let Some(count) = self + .timed("drain_queue", admin::drain_queue(&pool, &queue.name)) + .await + { + self.stats().add("drains", 1); + self.shared + .note("alex", format!("drained {} ({count} job(s))", queue.name)); + self.stats().add("drained_jobs", count as i64); + } + queue + .drained_at_ms + .store(self.shared.now_ms(), Ordering::Relaxed); + } + + /// Cancels live jobs of other personas, some of them already running. + async fn cancel(&mut self) { + let persona = *self + .rng + .pick(&[Persona::Mia, Persona::Tom, Persona::Sam, Persona::Priya]) + .expect("personas"); + let Some(queue) = self.queue_of(persona) else { + return; + }; + let states = [ + JobState::Available, + JobState::Running, + JobState::Scheduled, + JobState::Retryable, + ]; + let state = *self.rng.pick(&states).expect("states"); + let pool = self.pool.clone(); + let filter = ListJobsFilter { + queue: Some(queue.name.clone()), + state: Some(state), + limit: Some(30), + ..Default::default() + }; + let Some(rows) = self + .timed("admin_list_jobs", admin::list_jobs(&pool, &filter)) + .await + else { + return; + }; + let ids: Vec = rows.iter().map(|row| row.id).take(8).collect(); + if ids.is_empty() { + return; + } + self.shared.note( + "alex", + format!("cancelling jobs {ids:?} ({state:?}) in {}", queue.name), + ); + if self.rng.chance(50) { + let started_ms = self.shared.now_ms() as u32; + let first = ids[0]; + let completed_before = self + .shared + .ledger + .get(first) + .is_some_and(|record| record.completed_events > 0); + if let Some(Some(job)) = self.timed("cancel", admin::cancel(&pool, first)).await { + self.after_cancel(job.id, completed_before, started_ms); + } + } else if let Some(cancelled) = self + .timed("bulk_cancel", admin::bulk_cancel(&pool, &ids)) + .await + { + let now = self.shared.now_ms() as u32; + for job in &cancelled { + self.after_cancel(job.id, false, now); + } + } + } + + fn after_cancel(&self, job_id: i64, completed_before: bool, call_started_ms: u32) { + self.shared.record_operator_cancel( + Persona::Alex, + job_id, + completed_before, + call_started_ms, + ); + } + + async fn dlq_operations(&mut self) { + let pool = self.pool.clone(); + let filter = ListDlqFilter { + queue: Some(super::mia::INFER_QUEUE.to_string()), + limit: Some(20), + ..Default::default() + }; + let Some(rows) = self.timed("dlq_list", dlq::list_dlq(&pool, &filter)).await else { + return; + }; + let ids: Vec = rows.iter().map(|row| row.job.id).collect(); + self.stats().set_max("dlq_listed_max", ids.len() as i64); + match self.rng.below(4) { + 0 => { + if let Some(&id) = ids.first() { + self.shared.ledger.mark(id, flag::DLQ_RETRIED); + let opts = RetryFromDlqOpts::default(); + if self + .timed("dlq_retry", dlq::retry_from_dlq(&pool, id, &opts)) + .await + .is_some() + { + self.stats().add("dlq_retried", 1); + } + } + } + 1 => { + for id in ids.iter().take(3) { + self.shared.ledger.mark(*id, flag::DLQ_PURGED); + if self + .timed("dlq_purge_job", dlq::purge_dlq_job(&pool, *id)) + .await + == Some(true) + { + self.stats().add("dlq_purged", 1); + } + } + } + 2 => { + // Bulk retry: every exhausted Mia job may come back. + self.shared.ledger.mark_where( + |record| record.persona == Persona::Mia && record.has(flag::EXHAUSTED), + flag::DLQ_RETRIED, + ); + let bulk = ListDlqFilter { + queue: Some(super::mia::INFER_QUEUE.to_string()), + ..Default::default() + }; + if let Some(count) = self + .timed( + "dlq_bulk_retry", + dlq::bulk_retry_from_dlq(&pool, &bulk, false), + ) + .await + { + self.stats().add("dlq_bulk_retried", count as i64); + } + } + _ => { + self.shared.ledger.mark_where( + |record| record.persona == Persona::Mia && record.has(flag::EXHAUSTED), + flag::DLQ_PURGED, + ); + let bulk = ListDlqFilter { + queue: Some(super::mia::TRAIN_QUEUE.to_string()), + ..Default::default() + }; + if let Some(count) = self + .timed("dlq_bulk_purge", dlq::purge_dlq(&pool, &bulk, false)) + .await + { + self.stats().add("dlq_bulk_purged", count as i64); + } + } + } + } + + async fn batch_operations(&mut self) { + let (spec, filter, label, source) = if self.rng.chance(50) { + ( + BatchOperationSpec::MoveQueue { + queue: QUEUE_B.to_string(), + priority: None, + }, + BatchOperationFilter { + queue: Some(QUEUE_A.to_string()), + state: Some(JobState::Available), + ..Default::default() + }, + "move_queue", + QUEUE_A, + ) + } else { + ( + BatchOperationSpec::SetPriority { priority: 1 }, + BatchOperationFilter { + queue: Some(QUEUE_B.to_string()), + state: Some(JobState::Available), + ..Default::default() + }, + "set_priority", + QUEUE_B, + ) + }; + // Hold the source queue so a backlog builds for the operation to act on. + let pool = self.pool.clone(); + let held = self.shared.queues.get(source); + if let Some(queue) = &held { + queue.hold(); + let _ = self + .timed( + "pause_queue", + admin::pause_queue(&pool, &queue.name, Some("soak")), + ) + .await; + self.shared.stop_sleep(Duration::from_secs(6)).await; + } + let request = SubmitBatchOperation { + spec, + filter, + submitted_by: Some("soak".into()), + allow_all: false, + }; + self.run_batch(request, label).await; + if let Some(queue) = &held { + self.resume_with_retry(&queue.name).await; + self.shared.stop_sleep(Duration::from_secs(3)).await; + queue.release(); + } + } + + async fn run_batch(&mut self, request: SubmitBatchOperation, label: &str) { + let pool = self.pool.clone(); + let Some(operation) = self + .timed("batch_submit", submit_batch_operation(&pool, request)) + .await + else { + return; + }; + let started = Instant::now(); + loop { + let Some(current) = self + .timed("batch_get", get_batch_operation(&pool, operation.id)) + .await + else { + return; + }; + match current.state { + BatchOperationState::Completed | BatchOperationState::Cancelled => { + self.stats().add(&format!("batch_{label}"), 1); + self.stats() + .add(&format!("batch_{label}_processed"), current.processed); + self.stats() + .observe("batch_operation_ms", started.elapsed().as_millis() as i64); + return; + } + BatchOperationState::Failed => { + self.shared.violation( + "batch_operation_failed", + format!( + "{label} batch operation {} failed: {:?}", + operation.id, current.last_error + ), + vec![], + ); + return; + } + _ => {} + } + if started.elapsed() > Duration::from_secs(90) { + self.stats().add("batch_operations_incomplete", 1); + self.shared.note( + "alex", + format!( + "{label} batch operation {} still {:?} after 90s", + operation.id, current.state + ), + ); + return; + } + if self.shared.stop_sleep(Duration::from_secs(1)).await { + return; + } + } + } + + async fn rolling_restart(&mut self) { + let names = ["mia", "tom", "priya", "sam", "sam-stripe"]; + let name = *self.rng.pick(&names).expect("fleet names"); + let Some(fleet) = self.shared.fleet(name) else { + return; + }; + let live = fleet.live_replicas().await; + let Some(&index) = self.rng.pick(&live) else { + return; + }; + self.shared + .note("alex", format!("rolling restart: {name} replica {index}")); + self.shared.chaos_window(Duration::from_secs(20)); + let started = Instant::now(); + let restarted = fleet.restart(index, Duration::from_secs(30)).await; + // If the restarted runtime led, the next election can take a full interval. + self.shared.chaos_window(Duration::from_secs(45)); + match restarted { + Ok(_) => { + self.stats().add("rolling_restarts", 1); + self.stats() + .observe("rolling_restart_ms", started.elapsed().as_millis() as i64); + } + Err(error) => { + self.shared.violation( + "restart_failed", + format!("restarting {name} replica {index}: {error}"), + vec![], + ); + } + } + } +} + +impl Shared { + /// Sleeps unless the load phase ends first; true when it ended. + pub async fn stop_sleep(&self, duration: Duration) -> bool { + self.load_stop.sleep(duration).await + } +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Alex; + let queue_a = shared.queues.register(QUEUE_A, persona, false); + let queue_b = shared.queues.register(QUEUE_B, persona, false); + let queue_drain = shared.queues.register(QUEUE_DRAIN, persona, false); + let fleet = Fleet::start( + &shared, + "alex", + persona, + 2, + 10, + 2, + vec![(queue_a, 4), (queue_b, 4), (queue_drain, 4)], + client_factory(shared.clone()), + ) + .await?; + + let pool = connect_pool_with_timeout( + &shared.config.database_url, + "soak_alex_operator", + shared.budget.reserve("alex-operator", 4, 2), + Some(ADMIN_STATEMENT_TIMEOUT), + ) + .await + .map_err(|error| error.to_string())?; + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_alex_producer", + shared.budget.reserve("alex-producer", 2, 1), + ) + .await + .map_err(|error| error.to_string())?; + + // Alex's own traffic: a stream the move and reprioritise operations act on, and a + // stream that gets drained while it is still being produced. + let producer = { + let shared = shared.clone(); + tokio::spawn(async move { + let mut rng = Rng::for_task(shared.config.seed, "alex-producer"); + let mut seq = 0u64; + while !shared.load_stop.sleep(Duration::from_millis(200)).await { + let mut params = Vec::new(); + for queue in [ + QUEUE_A, + QUEUE_A, + QUEUE_A, + QUEUE_A, + QUEUE_B, + QUEUE_DRAIN, + QUEUE_DRAIN, + QUEUE_DRAIN, + ] { + seq += 1; + params.push( + awa::insert::params_with( + &AlexJob { + seq, + work_ms: rng.range(50, 300), + }, + InsertOpts { + queue: queue.into(), + ..Default::default() + }, + ) + .expect("alex params"), + ); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared + .stats(Persona::Alex) + .add("enqueued", rows.len() as i64); + register_rows( + &shared, + Persona::Alex, + &rows, + |_| expect::ANY_TERMINAL, + flag::OPERATOR_MOVED, + ); + } + Err(error) => backoff(&shared, error, Persona::Alex).await, + } + } + producer_pool.close().await; + }) + }; + + let diagnostics = connect_pool( + &shared.config.database_url, + "soak_alex_diagnostics", + shared.budget.reserve("alex-diagnostics", 1, 1), + ) + .await + .map_err(|error| error.to_string())?; + let mut context = Context { + shared: shared.clone(), + pool, + diagnostics, + stalled_rows: Default::default(), + rng: Rng::for_task(shared.config.seed, "alex-operator"), + }; + let load = shared.config.load.as_secs_f64(); + // Scale the cadence down for short runs so each action still happens a few times. + let pace = (load / 900.0).clamp(0.1, 1.0); + let mut schedule: Vec<(&str, f64, Instant)> = vec![ + ( + "reads", + shared.config.admin_read_every.as_secs_f64(), + Instant::now(), + ), + ( + "pause", + 60.0 * pace, + Instant::now() + Duration::from_secs_f64(20.0 * pace), + ), + ( + "drain", + 75.0 * pace, + Instant::now() + Duration::from_secs_f64(35.0 * pace), + ), + ( + "cancel", + 10.0 * pace, + Instant::now() + Duration::from_secs_f64(15.0 * pace), + ), + ( + "dlq", + 45.0 * pace, + Instant::now() + Duration::from_secs_f64(60.0 * pace), + ), + ( + "batch", + 40.0 * pace, + Instant::now() + Duration::from_secs_f64(25.0 * pace), + ), + ( + "restart", + 150.0 * pace, + Instant::now() + Duration::from_secs_f64(90.0 * pace), + ), + ]; + shared.note("alex", "operating: reads, pause/resume, drain, cancel, move/reprioritise, DLQ retry/purge, rolling restarts"); + while !shared.load_stop.is_stopped() { + let now = Instant::now(); + for (name, interval, due) in schedule.iter_mut() { + if *due > now || shared.load_stop.is_stopped() { + continue; + } + *due = Instant::now() + Duration::from_secs_f64(*interval); + match *name { + "reads" => context.admin_reads().await, + "pause" => context.pause_resume().await, + "drain" => context.drain().await, + "cancel" => context.cancel().await, + "dlq" => context.dlq_operations().await, + "batch" => context.batch_operations().await, + "restart" => context.rolling_restart().await, + _ => {} + } + } + shared.load_stop.sleep(Duration::from_millis(250)).await; + } + + let _ = producer.await; + shared.drain_stop.wait().await; + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + context.pool.close().await; + context.diagnostics.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/chaos.rs b/awa/tests/persona_soak/personas/chaos.rs new file mode 100644 index 000000000..fee1c63eb --- /dev/null +++ b/awa/tests/persona_soak/personas/chaos.rs @@ -0,0 +1,242 @@ +//! Chaos: kills a runtime without a shutdown handshake, terminates the maintenance +//! leader's connection, terminates random runtime connections, exhausts a runtime's +//! pool and holds a long idle-in-transaction session. + +use super::super::config::Persona; +use super::super::db::connect_pool; +use super::super::rng::Rng; +use super::super::shared::Shared; +use sqlx::PgPool; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +/// The advisory lock the maintenance leader holds for as long as it leads. +const MAINTENANCE_LEADER_LOCK_KEY: i64 = 0x4157_415f_4d41_494e; +const RUNTIME_FLEETS: [&str; 6] = ["mia", "tom", "priya", "sam", "sam-stripe", "fanout"]; + +async fn kill_runtime(shared: &Arc, rng: &mut Rng) { + let Some(&name) = rng.pick(&RUNTIME_FLEETS) else { + return; + }; + let Some(fleet) = shared.fleet(name) else { + return; + }; + let live = fleet.live_replicas().await; + // Never take the last runtime of a fleet down; the point is a survivor that rescues. + if live.len() < 2 { + return; + } + let Some(&index) = rng.pick(&live) else { + return; + }; + shared.note( + "chaos", + format!("killing {name} replica {index} without shutdown"), + ); + shared.chaos_window(Duration::from_secs(45)); + match fleet.kill(index, Some(Duration::from_secs(4))).await { + Ok(_) => shared.stats(Persona::Chaos).add("runtime_kills", 1), + Err(error) => shared.violation( + "restart_failed", + format!("restarting killed {name} replica {index}: {error}"), + vec![], + ), + } +} + +async fn kill_python(shared: &Arc) { + let handle = shared.python.lock().expect("python handle").clone(); + let Some(handle) = handle else { return }; + shared.note("chaos", "killing the Python process without shutdown"); + shared.chaos_window(Duration::from_secs(45)); + if let Err(error) = handle + .kill_and_restart(shared, Duration::from_secs(4)) + .await + { + shared.violation( + "restart_failed", + format!("restarting the Python process: {error}"), + vec![], + ); + } +} + +async fn terminate_leader(shared: &Arc, pool: &PgPool) { + // Maintenance leadership is an advisory lock held by one runtime connection. + let holders: Result, _> = sqlx::query_as( + "SELECT l.pid, COALESCE(a.application_name, '') + FROM pg_locks l JOIN pg_stat_activity a ON a.pid = l.pid + WHERE l.locktype = 'advisory' AND l.granted AND a.application_name LIKE 'soak\\_%' + AND ((l.classid::bigint << 32) | l.objid::bigint) = $1 + ORDER BY l.pid", + ) + .bind(MAINTENANCE_LEADER_LOCK_KEY) + .fetch_all(pool) + .await; + let Ok(holders) = holders else { return }; + let Some((pid, application)) = holders.first().cloned() else { + shared.note("chaos", "no advisory-lock holder found to terminate"); + return; + }; + shared.note( + "chaos", + format!("terminating maintenance leader backend {pid} ({application})"), + ); + shared.chaos_window(Duration::from_secs(40)); + let _ = sqlx::query("SELECT pg_terminate_backend($1)") + .bind(pid) + .fetch_one(pool) + .await; + shared.stats(Persona::Chaos).add("leader_terminations", 1); +} + +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)", + ) + .bind(count) + .fetch_all(pool) + .await; + let killed = terminated.map(|rows| rows.len()).unwrap_or(0); + shared.note( + "chaos", + format!("pg_terminate_backend on {killed} runtime connection(s)"), + ); + shared.chaos_window(Duration::from_secs(30)); + shared + .stats(Persona::Chaos) + .add("connections_terminated", killed as i64); +} + +/// Takes every connection of one runtime's pool for `hold`, starving its heartbeats +/// and claims the way a connection-pool exhaustion bug would. +async fn exhaust_pool(shared: &Arc, rng: &mut Rng, hold: Duration) { + let Some(&name) = rng.pick(&RUNTIME_FLEETS) else { + return; + }; + let Some(fleet) = shared.fleet(name) else { + return; + }; + let live = fleet.live_replicas().await; + let Some(&index) = rng.pick(&live) else { + return; + }; + let Some(pool) = fleet.pool_of(index).await else { + return; + }; + let size = pool.options().get_max_connections() as usize; + shared.note( + "chaos", + format!( + "exhausting the {size}-connection pool of {name} replica {index} for {}s", + hold.as_secs() + ), + ); + shared.chaos_window(hold + Duration::from_secs(40)); + let mut held = Vec::new(); + let deadline = Instant::now() + Duration::from_secs(20); + while held.len() < size && Instant::now() < deadline { + if let Ok(Ok(connection)) = + tokio::time::timeout(Duration::from_secs(2), pool.acquire()).await + { + held.push(connection); + } + } + shared.stats(Persona::Chaos).add("pool_exhaustions", 1); + shared + .stats(Persona::Chaos) + .add("pool_connections_held", held.len() as i64); + shared.load_stop.sleep(hold).await; + drop(held); +} + +/// One session that begins a transaction, takes a snapshot and sits idle, pinning the +/// xmin horizon so vacuum and prune cannot reclaim anything newer. +async fn idle_in_transaction(shared: Arc, hold: Duration) { + let Ok(pool) = connect_pool(&shared.config.database_url, "soak_chaos_idle", 1).await else { + return; + }; + let Ok(mut transaction) = pool.begin().await else { + return; + }; + if sqlx::query("SELECT txid_current()") + .execute(&mut *transaction) + .await + .is_err() + { + return; + } + shared.note( + "chaos", + format!( + "holding an idle-in-transaction session for {}s", + hold.as_secs() + ), + ); + shared + .stats(Persona::Chaos) + .add("idle_in_transaction_holds", 1); + shared.load_stop.sleep(hold).await; + let _ = transaction.rollback().await; + pool.close().await; +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let pool = connect_pool( + &shared.config.database_url, + "soak_chaos_control", + shared.budget.reserve("chaos", 3, 2), + ) + .await + .map_err(|error| error.to_string())?; + let mut rng = Rng::for_task(shared.config.seed, "chaos"); + let load = shared.config.load.as_secs_f64(); + let pace = (load / 900.0).clamp(0.15, 1.0); + let second = |seconds: f64| Duration::from_secs_f64(seconds * pace); + let mut schedule: Vec<(&str, Duration, Instant)> = vec![ + ("kill", second(120.0), Instant::now() + second(60.0)), + ("python", second(200.0), Instant::now() + second(110.0)), + ("leader", second(150.0), Instant::now() + second(90.0)), + ("terminate", second(75.0), Instant::now() + second(45.0)), + ("pool", second(240.0), Instant::now() + second(150.0)), + ("idle", second(300.0), Instant::now() + second(30.0)), + ]; + shared.note("chaos", "armed: runtime kill, leader termination, connection termination, pool exhaustion, idle transaction"); + while !shared.load_stop.is_stopped() { + let now = Instant::now(); + for (name, interval, due) in schedule.iter_mut() { + if *due > now || shared.load_stop.is_stopped() { + continue; + } + *due = Instant::now() + *interval; + match *name { + "kill" => kill_runtime(&shared, &mut rng).await, + "python" => kill_python(&shared).await, + "leader" => terminate_leader(&shared, &pool).await, + "terminate" => terminate_runtime_connections(&shared, &pool, &mut rng).await, + "pool" => { + exhaust_pool(&shared, &mut rng, Duration::from_secs(20).max(second(20.0))).await + } + "idle" => { + tokio::spawn(idle_in_transaction( + shared.clone(), + second(90.0).max(Duration::from_secs(30)), + )); + } + _ => {} + } + } + shared.load_stop.sleep(Duration::from_millis(500)).await; + } + pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/fanout.rs b/awa/tests/persona_soak/personas/fanout.rs new file mode 100644 index 000000000..8b3146300 --- /dev/null +++ b/awa/tests/persona_soak/personas/fanout.rs @@ -0,0 +1,278 @@ +//! Fanout, replay-style fan-out: each parent job batch-inserts a hundred or so unique +//! child jobs into the queue it runs on, from many concurrent handlers. Some parents +//! then replay a batch that reuses an already-claimed unique key, fail with a unique +//! violation and fall back to inserting that batch row by row. Stresses concurrent +//! batch enqueue into one lane, unique claims and the insert-from-handler path. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{attach_events, backoff, queue, register_rows, tune}; +use awa::model::{insert_many, insert_with}; +use awa::{ + AwaError, Client, InsertOpts, InsertParams, JobArgs, JobContext, JobError, JobResult, JobRow, + UniqueOpts, Worker, +}; +use serde::{Deserialize, Serialize}; +use sqlx::PgPool; +use std::sync::atomic::Ordering; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +pub const QUEUE: &str = "awa_soak_fanout"; +const MIN_CHILDREN: u64 = 100; +const MAX_CHILDREN: u64 = 185; +const REPLAY_PERCENT: u64 = 10; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct FanoutParent { + pub seq: u64, + pub children: u64, + pub replay: bool, +} + +/// Distinct `(parent_seq, index)` per child, so the unique key (queue + args) is +/// shared only by a deliberate replay. +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct FanoutChild { + pub parent_seq: u64, + pub index: u64, +} + +fn child_opts() -> InsertOpts { + InsertOpts { + queue: QUEUE.to_string(), + unique: Some(UniqueOpts { + by_queue: true, + by_args: true, + ..UniqueOpts::default() + }), + ..Default::default() + } +} + +fn is_unique_violation(error: &AwaError) -> bool { + match error { + AwaError::UniqueConflict { .. } => true, + AwaError::Database(sqlx::Error::Database(database)) => { + database.code().as_deref() == Some("23505") + } + _ => false, + } +} + +struct ParentWorker { + shared: Arc, +} + +impl ParentWorker { + /// Records rows the handler enqueued. A child can run before this is reached; the + /// ledger merges the two. + fn record(&self, rows: &[JobRow]) { + let persona = Persona::Fanout; + if let Some(queue) = self.shared.queues.get(QUEUE) { + queue + .enqueued + .fetch_add(rows.len() as i64, Ordering::Relaxed); + } + self.shared + .stats(persona) + .add("enqueued_children", rows.len() as i64); + register_rows( + &self.shared, + persona, + rows, + |_| expect::COMPLETED, + flag::NO_EVENTS, + ); + } + + async fn insert_batch( + &self, + pool: &PgPool, + parent_seq: u64, + indexes: impl Iterator + Clone, + ) -> Result<(), JobError> { + let stats = self.shared.stats(Persona::Fanout); + let params: Vec = indexes + .clone() + .map(|index| { + awa::insert::params_with(&FanoutChild { parent_seq, index }, child_opts()) + .map_err(JobError::retryable) + }) + .collect::>()?; + let started = Instant::now(); + match insert_many(pool, ¶ms).await { + Ok(rows) => { + stats.observe("batch_insert_ms", started.elapsed().as_millis() as i64); + self.record(&rows); + Ok(()) + } + Err(error) if is_unique_violation(&error) => { + stats.add("batch_unique_conflicts", 1); + for index in indexes { + match insert_with(pool, &FanoutChild { parent_seq, index }, child_opts()).await + { + Ok(row) => { + stats.add("fallback_rows_inserted", 1); + self.record(&[row]); + } + Err(error) if is_unique_violation(&error) => { + stats.add("fallback_rows_skipped", 1) + } + Err(error) => return Err(JobError::retryable(error)), + } + } + Ok(()) + } + Err(error) => { + stats.add("batch_insert_errors", 1); + Err(JobError::retryable(error)) + } + } + } +} + +#[async_trait::async_trait] +impl Worker for ParentWorker { + fn kind(&self) -> &'static str { + "fanout_parent" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self + .shared + .begin_execution(Persona::Fanout, &ctx.job, false); + let args: FanoutParent = serde_json::from_value(ctx.job.args.clone()) + .map_err(|error| JobError::terminal(format!("bad args: {error}")))?; + self.insert_batch(ctx.pool(), args.seq, 0..args.children) + .await?; + if args.replay { + // Index 0 is already claimed, so the batch fails as a whole. + let fresh = args.children..args.children * 2 - 1; + self.insert_batch(ctx.pool(), args.seq, std::iter::once(0).chain(fresh)) + .await?; + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +struct ChildWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for ChildWorker { + fn kind(&self) -> &'static str { + "fanout_child" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + // No lifecycle hooks: a hook per child would add a read per job. + let _guard = self.shared.begin_execution(Persona::Fanout, &ctx.job, true); + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn client_factory(shared: Arc) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(QUEUE, queue(16)) + .register_worker(ParentWorker { + shared: shared.clone(), + }) + .register_worker(ChildWorker { + shared: shared.clone(), + }); + attach_events(tune(builder, &shared), &shared, &["fanout_parent"]) + .build() + .expect("build fanout client") + }) +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Fanout; + let queue_info = shared.queues.register(QUEUE, persona, false); + let fleet = Fleet::start( + &shared, + "fanout", + persona, + 2, + 12, + 2, + vec![(queue_info.clone(), 16)], + client_factory(shared.clone()), + ) + .await?; + + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_fanout_producer", + shared.budget.reserve("fanout-producer", 2, 1), + ) + .await + .map_err(|error| error.to_string())?; + let mut rng = Rng::for_task(shared.config.seed, "fanout-producer"); + let tick = Duration::from_millis(500); + let parents_per_tick = shared.config.scaled(2) as usize; + let mut seq = 0u64; + shared.note( + "fanout", + format!( + "{parents_per_tick} parents per {} ms tick, {MIN_CHILDREN}-{MAX_CHILDREN} unique children each, {REPLAY_PERCENT}% replay a used key", + tick.as_millis() + ), + ); + + while !shared.load_stop.sleep(tick).await { + let mut params = Vec::with_capacity(parents_per_tick); + for _ in 0..parents_per_tick { + seq += 1; + params.push( + awa::insert::params_with( + &FanoutParent { + seq, + children: rng.range(MIN_CHILDREN, MAX_CHILDREN), + replay: rng.chance(REPLAY_PERCENT), + }, + InsertOpts { + queue: QUEUE.to_string(), + max_attempts: 5, + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + queue_info + .enqueued + .fetch_add(rows.len() as i64, Ordering::Relaxed); + shared + .stats(persona) + .add("enqueued_parents", rows.len() as i64); + register_rows(&shared, persona, &rows, |_| expect::COMPLETED, 0); + } + Err(error) => backoff(&shared, error, persona).await, + } + } + + shared.note( + "fanout", + format!("load stopped after {seq} parents; draining"), + ); + shared.drain_stop.wait().await; + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + producer_pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/mia.rs b/awa/tests/persona_soak/personas/mia.rs new file mode 100644 index 000000000..07d77f276 --- /dev/null +++ b/awa/tests/persona_soak/personas/mia.rs @@ -0,0 +1,260 @@ +//! Mia, an ML pipeline: jobs from seconds to about a minute (compressed), deadlines +//! that sometimes expire, heartbeats and progress, retries into the DLQ and latency +//! spikes. Stresses deadline and stale rescues racing heartbeats, completions and +//! the DLQ. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{attach_events, backoff, queue, register_rows, tune}; +use awa::model::{insert, insert_many}; +use awa::{Client, InsertOpts, JobArgs, JobContext, JobError, JobResult, Worker}; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; +use std::time::Duration; + +pub const TRAIN_QUEUE: &str = "awa_soak_mia_train"; +pub const INFER_QUEUE: &str = "awa_soak_mia_infer"; +pub const DEADLINE_QUEUE: &str = "awa_soak_mia_deadline"; +/// Deadline given to jobs in the deadline queue; they run longer than this. +const DEADLINE: chrono::Duration = chrono::Duration::seconds(6); + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct MiaJob { + pub seq: u64, + pub behavior: String, + pub work_ms: u64, +} + +struct MiaWorker { + shared: Arc, +} + +/// Sleeps in slices, reporting progress and stopping early when the job is cancelled +/// or the runtime is shutting down. +async fn work(ctx: &JobContext, total_ms: u64) -> bool { + let slice = 500; + let mut elapsed = 0; + while elapsed < total_ms { + if ctx.is_cancelled() { + return false; + } + let step = slice.min(total_ms - elapsed); + tokio::time::sleep(Duration::from_millis(step)).await; + elapsed += step; + if total_ms >= 4_000 && elapsed % 2_000 < slice { + ctx.set_progress(((elapsed * 100) / total_ms).min(99) as u8, "training"); + let _ = ctx.flush_progress().await; + } + } + true +} + +#[async_trait::async_trait] +impl Worker for MiaWorker { + fn kind(&self) -> &'static str { + "mia_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self.shared.begin_execution(Persona::Mia, &ctx.job, false); + let args: MiaJob = serde_json::from_value(ctx.job.args.clone()) + .map_err(|error| JobError::terminal(format!("bad args: {error}")))?; + let attempt = ctx.job.attempt; + let finished = match args.behavior.as_str() { + "flaky" if attempt < 3 => { + work(ctx, args.work_ms.min(1_500)).await; + return Err(JobError::retryable_msg("transient model server error")); + } + "doomed" => { + work(ctx, args.work_ms.min(1_500)).await; + return Err(JobError::retryable_msg("model artefact missing")); + } + "cancel" => { + work(ctx, args.work_ms.min(1_000)).await; + return Ok(JobResult::Cancel("rejected by validation".into())); + } + // The first attempt overruns its deadline; later attempts often recover. + "overrun" if attempt >= 2 && !args.seq.is_multiple_of(3) => work(ctx, 1_000).await, + _ => work(ctx, args.work_ms).await, + }; + if !finished { + return Err(JobError::retryable_msg("interrupted")); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn client_factory(shared: Arc) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(TRAIN_QUEUE, queue(24)) + .queue(INFER_QUEUE, queue(16)) + .queue(DEADLINE_QUEUE, queue(12)) + .register_worker(MiaWorker { + shared: shared.clone(), + }); + attach_events(tune(builder, &shared), &shared, &["mia_job"]) + .build() + .expect("build mia client") + }) +} + +/// One enqueued job: queue, behaviour, how long it works, deadline, max attempts and +/// the outcomes it may end in. +struct Plan { + queue: &'static str, + behavior: &'static str, + work_ms: u64, + max_attempts: i16, + expect: u8, +} + +fn plan(rng: &mut Rng) -> Plan { + let roll = rng.below(100); + if roll < 45 { + Plan { + queue: INFER_QUEUE, + behavior: "fast", + work_ms: rng.range(100, 2_000), + max_attempts: 5, + expect: expect::COMPLETED, + } + } else if roll < 70 { + Plan { + queue: TRAIN_QUEUE, + behavior: "train", + work_ms: rng.range(2_000, 20_000), + max_attempts: 5, + expect: expect::COMPLETED, + } + } else if roll < 78 { + // Latency spike: a long job that heartbeats and reports progress throughout. + Plan { + queue: TRAIN_QUEUE, + behavior: "spike", + work_ms: rng.range(30_000, 60_000), + max_attempts: 5, + expect: expect::COMPLETED, + } + } else if roll < 88 { + Plan { + queue: DEADLINE_QUEUE, + behavior: "overrun", + work_ms: rng.range(10_000, 25_000), + max_attempts: 3, + expect: expect::COMPLETED | expect::FAILED, + } + } else if roll < 94 { + Plan { + queue: INFER_QUEUE, + behavior: "flaky", + work_ms: 500, + max_attempts: 5, + expect: expect::COMPLETED | expect::FAILED, + } + } else if roll < 98 { + Plan { + queue: INFER_QUEUE, + behavior: "doomed", + work_ms: 500, + max_attempts: 3, + expect: expect::FAILED, + } + } else { + Plan { + queue: INFER_QUEUE, + behavior: "cancel", + work_ms: 500, + max_attempts: 3, + expect: expect::CANCELLED, + } + } +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Mia; + let train = shared.queues.register(TRAIN_QUEUE, persona, false); + let infer = shared.queues.register(INFER_QUEUE, persona, false); + let deadline = shared.queues.register(DEADLINE_QUEUE, persona, false); + let fleet = Fleet::start( + &shared, + "mia", + persona, + 2, + 14, + 2, + vec![(train, 24), (infer, 16), (deadline, 12)], + client_factory(shared.clone()), + ) + .await?; + + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_mia_producer", + shared.budget.reserve("mia-producer", 2, 1), + ) + .await + .map_err(|error| error.to_string())?; + let rate_per_second = (6.0 * shared.config.scale).max(0.5); + let tick = Duration::from_millis(250); + let mut rng = Rng::for_task(shared.config.seed, "mia-producer"); + let mut seq = 0u64; + let mut carry = 0.0; + shared.note("mia", format!("producing about {rate_per_second:.1} jobs/s across {TRAIN_QUEUE}, {INFER_QUEUE} and {DEADLINE_QUEUE}")); + + while !shared.load_stop.sleep(tick).await { + carry += rate_per_second * tick.as_secs_f64(); + let count = carry as usize; + carry -= count as f64; + if count == 0 { + continue; + } + let mut params = Vec::with_capacity(count); + let mut expects = Vec::with_capacity(count); + for _ in 0..count { + seq += 1; + let plan = plan(&mut rng); + let args = MiaJob { + seq, + behavior: plan.behavior.to_string(), + work_ms: plan.work_ms, + }; + params.push( + insert::params_with( + &args, + InsertOpts { + queue: plan.queue.to_string(), + max_attempts: plan.max_attempts, + deadline_duration: (plan.queue == DEADLINE_QUEUE).then_some(DEADLINE), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + expects.push(plan.expect); + shared + .stats(persona) + .add(&format!("planned_{}", plan.behavior), 1); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared.stats(persona).add("enqueued", rows.len() as i64); + register_rows(&shared, persona, &rows, |index| expects[index], 0); + } + Err(error) => backoff(&shared, error, persona).await, + } + } + + shared.note("mia", format!("load stopped after {seq} jobs; draining")); + shared.drain_stop.wait().await; + fleet.shutdown_all(Duration::from_secs(30), &shared).await; + producer_pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/mod.rs b/awa/tests/persona_soak/personas/mod.rs new file mode 100644 index 000000000..bfebd3462 --- /dev/null +++ b/awa/tests/persona_soak/personas/mod.rs @@ -0,0 +1,117 @@ +//! The personas. Each one owns its queues, job kinds, runtimes and producers, and +//! reports what it did through [`Shared::stats`]. + +pub mod abandon; +pub mod alex; +pub mod chaos; +pub mod fanout; +pub mod mia; +pub mod priya; +pub mod python; +pub mod sam; +pub mod shardops; +pub mod tom; +pub mod ui; + +use super::config::Persona; +use super::ledger::{expect, JobRecord}; +use super::shared::Shared; +use awa::{ClientBuilder, JobRow, QueueConfig}; +use std::sync::Arc; +use std::time::Duration; + +/// Registers the lifecycle event hook for each kind so the ledger learns about +/// terminal outcomes as they commit. +pub fn attach_events( + mut builder: ClientBuilder, + shared: &Arc, + kinds: &[&str], +) -> ClientBuilder { + for kind in kinds { + let shared = shared.clone(); + builder = builder.on_event_kind(*kind, move |event| { + let shared = shared.clone(); + async move { shared.on_event(&event) } + }); + } + builder +} + +/// Runtime timings shared by every fleet: fast enough that rescues, promotion and +/// rotation happen many times in a short run, but with realistic ratios between them. +pub fn tune(builder: ClientBuilder, shared: &Shared) -> ClientBuilder { + // The maintenance leader evaluates only the schedules its own runtime declares, and + // any runtime can lead, so every fleet declares the same set. + let builder = if shared.config.has(Persona::Tom) { + tom::declare_cron(builder) + } else { + builder + }; + let builder = timings( + builder, + shared.config.heartbeat_interval, + shared.config.heartbeat_staleness, + ); + match shared.config.rescue_interval { + Some(interval) => builder + .heartbeat_rescue_interval(interval) + .deadline_rescue_interval(interval), + None => builder, + } +} + +/// The timings of [`tune`] that do not depend on the run, for runtimes that live in a +/// process of their own. +pub fn timings( + builder: ClientBuilder, + heartbeat_interval: Duration, + heartbeat_staleness: Duration, +) -> ClientBuilder { + builder + .heartbeat_interval(heartbeat_interval) + .heartbeat_staleness(heartbeat_staleness) + .heartbeat_rescue_interval(Duration::from_secs(3)) + .deadline_rescue_interval(Duration::from_secs(2)) + .promote_interval(Duration::from_millis(250)) + .runtime_snapshot_interval(Duration::from_secs(5)) + .dlq_enabled_by_default(true) +} + +pub fn queue(max_workers: u32) -> QueueConfig { + QueueConfig { + max_workers, + ..QueueConfig::default() + } +} + +/// Records freshly enqueued rows in the ledger. +pub fn register_rows( + shared: &Shared, + persona: Persona, + rows: &[JobRow], + expect_for: impl Fn(usize) -> u8, + flags: u32, +) { + let now = shared.now_ms() as u32; + for (index, row) in rows.iter().enumerate() { + let queue = shared + .queues + .get(&row.queue) + .map_or(u32::MAX, |queue| queue.index); + shared.ledger.register( + row.id, + JobRecord::new(persona, queue, now, expect_for(index), flags), + ); + } +} + +pub const ANY: u8 = expect::ANY_TERMINAL; + +/// Short pause used by producers after an error, so a dead connection does not spin. +pub async fn backoff(shared: &Shared, error: impl std::fmt::Display, persona: Persona) { + shared.stats(persona).add("enqueue_errors", 1); + if shared.stats(persona).get("enqueue_errors") <= 5 { + shared.note(persona.name(), format!("enqueue error: {error}")); + } + shared.load_stop.sleep(Duration::from_millis(500)).await; +} diff --git a/awa/tests/persona_soak/personas/priya.rs b/awa/tests/persona_soak/personas/priya.rs new file mode 100644 index 000000000..e674d9a31 --- /dev/null +++ b/awa/tests/persona_soak/personas/priya.rs @@ -0,0 +1,203 @@ +//! Priya, firehose: one hot queue, large batch inserts (above the bind-parameter +//! limit), COPY inserts, tiny jobs and several clients. Stresses bind limits, ring +//! rotation and prune under load and claim contention. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{backoff, queue, register_rows, tune}; +use awa::model::{insert_many, insert_many_copy_from_pool}; +use awa::{Client, InsertOpts, InsertParams, JobArgs, JobContext, JobError, JobResult, Worker}; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +pub const HOT_QUEUE: &str = "awa_soak_priya_hot"; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct PriyaJob { + pub seq: u64, + pub source: String, +} + +struct PriyaWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for PriyaWorker { + fn kind(&self) -> &'static str { + "priya_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self.shared.begin_execution(Persona::Priya, &ctx.job, true); + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn client_factory(shared: Arc) -> ClientFactory { + Arc::new(move |pool| { + // No lifecycle hooks: a hook per completion would add a read per job and + // change the load this persona exists to apply. + tune(Client::builder(pool).queue(HOT_QUEUE, queue(48)), &shared) + .register_worker(PriyaWorker { + shared: shared.clone(), + }) + .build() + .expect("build priya client") + }) +} + +fn batch(seq: &mut u64, count: usize, source: &str) -> Vec { + (0..count) + .map(|_| { + *seq += 1; + awa::insert::params_with( + &PriyaJob { + seq: *seq, + source: source.to_string(), + }, + InsertOpts { + queue: HOT_QUEUE.to_string(), + max_attempts: 5, + ..Default::default() + }, + ) + .expect("priya params") + }) + .collect() +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Priya; + let hot = shared.queues.register(HOT_QUEUE, persona, false); + let fleet = Fleet::start( + &shared, + "priya", + persona, + 3, + 12, + 3, + vec![(hot, 48)], + client_factory(shared.clone()), + ) + .await?; + + let pool_size = shared.budget.reserve("priya-producer", 6, 2); + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_priya_producer", + pool_size, + ) + .await + .map_err(|error| error.to_string())?; + let mut rng = Rng::for_task(shared.config.seed, "priya-producer"); + let burst_size = shared.config.scaled(100_000) as usize; + let copy_size = shared.config.scaled(10_000) as usize; + let steady_per_tick = shared.config.scaled(30) as usize; + let burst_every = Duration::from_secs(180); + let copy_every = Duration::from_secs(45); + let mut next_burst = Instant::now() + Duration::from_secs(10); + let mut next_copy = Instant::now() + Duration::from_secs(25); + let mut seq = 0u64; + shared.note( + "priya", + format!("steady {} jobs/s, {burst_size}-row insert_many burst every {}s, {copy_size}-row COPY every {}s", steady_per_tick * 10, burst_every.as_secs(), copy_every.as_secs()), + ); + + while !shared.load_stop.sleep(Duration::from_millis(100)).await { + let now = Instant::now(); + if now >= next_burst { + next_burst = now + burst_every; + // Vary the size around the target so chunk boundaries move between bursts. + let size = burst_size + rng.below(5_000) as usize; + let params = batch(&mut seq, size, "burst"); + let started = Instant::now(); + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared.stats(persona).add("enqueued", rows.len() as i64); + shared.stats(persona).add("bursts", 1); + shared + .stats(persona) + .observe("burst_insert_ms", started.elapsed().as_millis() as i64); + register_rows( + &shared, + persona, + &rows, + |_| expect::COMPLETED, + flag::NO_EVENTS | flag::SRC_BURST, + ); + shared.note( + "priya", + format!( + "burst of {} rows inserted in {:?}", + rows.len(), + started.elapsed() + ), + ); + } + Err(error) if shared.chaos_active() => { + backoff(&shared, format!("burst of {size}: {error}"), persona).await + } + Err(error) => { + shared.violation( + "insert_many_failed", + format!("insert_many of {size} rows failed: {error}"), + vec![], + ); + } + } + } + if now >= next_copy { + next_copy = now + copy_every; + let params = batch(&mut seq, copy_size, "copy"); + let started = Instant::now(); + match insert_many_copy_from_pool(&producer_pool, ¶ms).await { + Ok(rows) => { + shared.stats(persona).add("enqueued", rows.len() as i64); + shared.stats(persona).add("copy_batches", 1); + shared + .stats(persona) + .observe("copy_insert_ms", started.elapsed().as_millis() as i64); + register_rows( + &shared, + persona, + &rows, + |_| expect::COMPLETED, + flag::NO_EVENTS | flag::SRC_COPY, + ); + } + Err(error) => { + backoff(&shared, format!("copy of {copy_size}: {error}"), persona).await + } + } + } + let params = batch(&mut seq, steady_per_tick, "steady"); + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared.stats(persona).add("enqueued", rows.len() as i64); + register_rows( + &shared, + persona, + &rows, + |_| expect::COMPLETED, + flag::NO_EVENTS | flag::SRC_STEADY, + ); + } + Err(error) => backoff(&shared, error, persona).await, + } + } + + shared.note("priya", format!("load stopped after {seq} jobs; draining")); + shared.drain_stop.wait().await; + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + producer_pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/python.rs b/awa/tests/persona_soak/personas/python.rs new file mode 100644 index 000000000..7ece76691 --- /dev/null +++ b/awa/tests/persona_soak/personas/python.rs @@ -0,0 +1,887 @@ +//! Python users: a separate process built on the `awa-python` bindings, started and +//! stopped by the harness. It produces (single inserts, `insert_many` above the old +//! bind-parameter limit, COPY, and enqueues inside transactions that sometimes roll +//! back), runs async workers (completing, retryable, `RetryAfter`, `Snooze`, progress, +//! deadline overruns and a handler that ignores its deadline) and exercises the +//! migrate-then-query startup path. One queue is shared with a Rust runtime so Rust and +//! Python workers claim from the same lanes. The process writes its own per-job ledger, +//! which this module folds into the harness ledger for the oracle. + +use super::super::config::Persona; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag, JobRecord}; +use super::super::python_probe::{self, Evidence, Phase}; +use super::super::shared::Shared; +use super::{attach_events, queue, tune}; +use awa::{Client, JobArgs, JobContext, JobError, JobResult, Worker}; +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashSet; +use std::io::{BufRead, Seek, SeekFrom}; +use std::path::PathBuf; +use std::process::Stdio; +use std::sync::atomic::{AtomicU8, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::io::{AsyncBufReadExt, BufReader}; +use tokio::process::{Child, Command}; + +pub const QUEUE_WORK: &str = "awa_soak_py_work"; +pub const QUEUE_DEADLINE: &str = "awa_soak_py_deadline"; +pub const QUEUE_SHARED: &str = "awa_soak_py_shared"; +const SCRIPT: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/tests/persona_soak/python/soak_python.py" +); + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct PySharedJob { + pub seq: u64, + pub behavior: String, + pub work_ms: u64, +} + +struct SharedQueueWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for SharedQueueWorker { + fn kind(&self) -> &'static str { + "py_shared_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self + .shared + .begin_execution(Persona::Python, &ctx.job, false); + let work_ms = ctx.job.args["work_ms"].as_u64().unwrap_or(50); + match ctx.job.args["behavior"].as_str().unwrap_or("ok") { + "retry" if ctx.job.attempt < 3 => { + tokio::time::sleep(Duration::from_millis(50)).await; + return Err(JobError::retryable_msg("transient failure")); + } + "progress" => { + for step in 1..=4u8 { + tokio::time::sleep(Duration::from_millis(work_ms / 4)).await; + ctx.set_progress(step * 25, "step"); + let _ = ctx.flush_progress().await; + } + } + _ => tokio::time::sleep(Duration::from_millis(work_ms)).await, + } + if ctx.is_cancelled() { + return Err(JobError::retryable_msg("interrupted")); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +/// Handle to the running Python process, shared with the chaos persona. +pub struct PythonProcess { + child: Option, + ledger_path: PathBuf, + python: PathBuf, + restarts: u32, +} + +fn spawn( + shared: &Arc, + python: &PathBuf, + ledger_path: &PathBuf, + role: &str, +) -> Result { + spawn_process(shared, python, ledger_path, role, "soak_python", None) +} + +fn spawn_process( + shared: &Arc, + python: &PathBuf, + ledger_path: &PathBuf, + role: &str, + application_name: &str, + probe_phase: Option>, +) -> Result { + let mut child = Command::new(python) + .arg(SCRIPT) + .env( + "SOAK_DATABASE_URL", + format!( + "{}{}application_name={}", + shared.config.database_url, + if shared.config.database_url.contains('?') { + '&' + } else { + '?' + }, + application_name + ), + ) + .env("SOAK_LEDGER", ledger_path) + .env("SOAK_SEED", shared.config.seed.to_string()) + .env("SOAK_SCALE", shared.config.scale.to_string()) + .env( + "SOAK_PACE", + (shared.config.load.as_secs_f64() / 900.0) + .clamp(0.1, 1.0) + .to_string(), + ) + .env("SOAK_ROLE", role) + .env( + "SOAK_DRAIN", + if shared.load_stop.is_stopped() { + "1" + } else { + "0" + }, + ) + .env("SOAK_QUEUE_WORK", QUEUE_WORK) + .env("SOAK_QUEUE_DEADLINE", QUEUE_DEADLINE) + .env("SOAK_QUEUE_SHARED", QUEUE_SHARED) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn() + .map_err(|error| format!("spawn {}: {error}", python.display()))?; + for stream in [ + child + .stdout + .take() + .map(|s| Box::new(s) as Box), + child + .stderr + .take() + .map(|s| Box::new(s) as Box), + ] + .into_iter() + .flatten() + { + let shared = shared.clone(); + let probe_phase = probe_phase.clone(); + let application_name = application_name.to_string(); + tokio::spawn(async move { + let mut lines = BufReader::new(stream).lines(); + while let Ok(Some(line)) = lines.next_line().await { + if let (Some(phase), Some(marker)) = (&probe_phase, Phase::marker(&line)) { + // A close marker alone does not prove migrate-and-query + // reached readiness. IDLE also follows that completed call. + if marker != Phase::Closed || phase.load(Ordering::Acquire) > 0 { + phase.fetch_max(marker as u8, Ordering::Release); + } + shared.note("python", format!("{application_name}: {line}")); + } + absorb_output(&shared, &line); + } + }); + } + Ok(child) +} + +/// Turns the interesting stdout lines of the Python process into report counters. +fn absorb_output(shared: &Shared, line: &str) { + let stats = shared.stats(Persona::Python); + let field = |name: &str| { + line.split_whitespace() + .find_map(|part| part.strip_prefix(&format!("{name}="))) + .and_then(|value| value.parse::().ok()) + }; + if line.starts_with("STAT ") { + if let Some(active) = field("active") { + shared + .python_active + .store(active as i64, std::sync::atomic::Ordering::Relaxed); + } + if let Some(max_active) = field("max_active") { + stats.set_max("python_max_concurrent_handlers", max_active as i64); + } + if let Some(errors) = field("errors") { + stats.set_max("python_errors", errors as i64); + } + } else if line.starts_with("BATCH ") { + if let (Some(rows), Some(seconds)) = (field("rows"), field("seconds")) { + stats.add("python_batches", 1); + stats.add("python_batch_rows", rows as i64); + stats.observe("python_batch_insert_ms", (seconds * 1000.0) as i64); + } + } else if line.starts_with("MIGRATED") { + stats.add("python_migrations", 1); + } else if line.starts_with("ERROR") || line.contains("Traceback") { + if stats.get("python_error_lines") < 5 { + shared.note("python", line.chars().take(200).collect::()); + } + stats.add("python_error_lines", 1); + } +} + +pub struct PythonHandle { + process: tokio::sync::Mutex, +} + +// Used in both phases: after load stops a dead process still needs a worker +// replacement, but spawn_process disables its producers through SOAK_DRAIN. +fn restart_if_needed( + child: &mut Option, + start: impl FnOnce() -> Result, +) -> Result, String> { + let reason = match child.as_mut() { + Some(child) => match child.try_wait().map_err(|error| error.to_string())? { + Some(status) => format!("process exited with {status}"), + None => return Ok(None), + }, + None => "process was absent".to_string(), + }; + child.take(); + *child = Some(start()?); + Ok(Some(reason)) +} + +fn finished_probe_errors(probes: &mut tokio::task::JoinSet<()>) -> Vec { + let mut errors = Vec::new(); + while let Some(result) = probes.try_join_next() { + if let Err(error) = result { + errors.push(error.to_string()); + } + } + errors +} + +async fn begin_drain( + stop_producers: impl std::future::Future, + probes: &mut tokio::task::JoinSet<()>, +) -> Vec { + // Startup probes may still be blocked. Stop production before waiting for + // them, and leave their observation tasks running beside worker recovery. + stop_producers.await; + finished_probe_errors(probes) +} + +impl PythonHandle { + async fn ensure_running(&self, shared: &Arc) { + let mut process = self.process.lock().await; + let python = process.python.clone(); + let ledger_path = process.ledger_path.clone(); + match restart_if_needed(&mut process.child, || { + spawn(shared, &python, &ledger_path, "full") + }) { + Ok(Some(reason)) => { + process.restarts += 1; + shared + .stats(Persona::Python) + .add("python_unexpected_exits", 1); + shared.note( + "python", + format!( + "{reason}; restarted (draining={})", + shared.load_stop.is_stopped() + ), + ); + } + Ok(None) => {} + Err(error) => shared.incomplete( + format!("Python worker liveness/restart check failed: {error}"), + vec![], + ), + } + } + /// Kills the process without a shutdown handshake and starts a new one after `downtime`. + pub async fn kill_and_restart( + &self, + shared: &Arc, + downtime: Duration, + ) -> Result<(), String> { + let mut process = self.process.lock().await; + if let Some(mut child) = process.child.take() { + let _ = child.kill().await; + } + tokio::time::sleep(downtime).await; + shared + .python_active + .store(0, std::sync::atomic::Ordering::Relaxed); + process.restarts += 1; + let child = spawn( + shared, + &process.python.clone(), + &process.ledger_path.clone(), + "full", + )?; + process.child = Some(child); + shared.stats(Persona::Python).add("python_kills", 1); + Ok(()) + } +} + +/// Folds new lines of the Python ledger into the harness ledger. +struct Importer { + path: PathBuf, + offset: u64, + rolled_back: HashSet, +} + +impl Importer { + fn import(&mut self, shared: &Shared) { + let Ok(mut file) = std::fs::File::open(&self.path) else { + return; + }; + if file.seek(SeekFrom::Start(self.offset)).is_err() { + return; + } + let mut reader = std::io::BufReader::new(file); + let mut line = String::new(); + loop { + line.clear(); + match reader.read_line(&mut line) { + Ok(0) | Err(_) => break, + Ok(read) => { + if !line.ends_with('\n') { + break; + } + self.offset += read as u64; + if let Ok(event) = serde_json::from_str::(&line) { + self.apply(shared, &event); + } + } + } + } + } + + fn apply(&mut self, shared: &Shared, event: &Value) { + let now = shared.now_ms() as u32; + let id = event.get("id").and_then(Value::as_i64); + let queue_of = |name: Option<&str>| { + name.and_then(|name| shared.queues.get(name)) + .map_or(u32::MAX, |queue| queue.index) + }; + match event.get("e").and_then(Value::as_str) { + Some("enq") => { + let Some(id) = id else { return }; + let allowed = event.get("x").and_then(Value::as_u64).unwrap_or(7) as u8; + let queue = shared + .queues + .get(event.get("q").and_then(Value::as_str).unwrap_or("")); + if let Some(queue) = &queue { + queue + .enqueued + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + let index = queue.map_or(u32::MAX, |queue| queue.index); + let mut record = + JobRecord::new(Persona::Python, index, now, allowed, flag::NO_EVENTS); + record.max_attempt = 0; + // Keep what the Rust side may have seen already (shared queue). + shared.ledger.with( + id, + || record.clone(), + |existing| { + existing.expect = allowed; + existing.flags &= !flag::UNREGISTERED; + existing.persona = Persona::Python; + }, + ); + shared.stats(Persona::Python).add("python_enqueued", 1); + } + Some("start") => { + let Some(id) = id else { return }; + let queue = queue_of(event.get("q").and_then(Value::as_str)); + let attempt = event.get("attempt").and_then(Value::as_i64).unwrap_or(0) as i16; + shared.ledger.with( + id, + || { + JobRecord::new( + Persona::Python, + queue, + now, + expect::ANY_TERMINAL, + flag::NO_EVENTS, + ) + }, + |record| { + record.executions = record.executions.saturating_add(1); + record.max_attempt = record.max_attempt.max(attempt); + if record.first_start_ms == 0 { + record.first_start_ms = now.max(1); + } + record.last_start_ms = now; + }, + ); + if let Some(queue) = event + .get("q") + .and_then(Value::as_str) + .and_then(|name| shared.queues.get(name)) + { + queue + .started + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + shared.stats(Persona::Python).add("python_executions", 1); + } + Some("end") => { + let Some(id) = id else { return }; + let result = event.get("r").and_then(Value::as_str).unwrap_or(""); + shared.ledger.with( + id, + || { + JobRecord::new( + Persona::Python, + u32::MAX, + now, + expect::ANY_TERMINAL, + flag::NO_EVENTS, + ) + }, + |record| match result { + // A completed return is evidence only for jobs that must complete; for + // the rest the outcome is read from the database, since a late return + // after a deadline rescue is discarded by the store. + "completed" if record.expect == expect::COMPLETED => { + record.flags |= flag::RETURNED_COMPLETED + } + "cancel" => { + record.flags |= flag::HANDLER_CANCELLED; + record.cancel_events = record.cancel_events.saturating_add(1); + } + _ => {} + }, + ); + } + Some("rolled_back") => { + if let Some(ids) = event.get("ids").and_then(Value::as_array) { + for id in ids.iter().filter_map(Value::as_i64) { + self.rolled_back.insert(id); + } + shared + .stats(Persona::Python) + .add("python_rolled_back_jobs", ids.len() as i64); + } + } + Some("error") => shared.stats(Persona::Python).add("python_logged_errors", 1), + _ => {} + } + } +} + +/// A migrate-then-query run in a throwaway process; its database sessions must not be +/// left stranded while it idles, and must be gone once it exits. +async fn migrate_probe(shared: &Arc, python: &PathBuf, ledger_path: &PathBuf) { + let stats = shared.stats(Persona::Python); + let application_name = format!("soak_python_probe_{}", uuid::Uuid::new_v4().simple()); + let phase = Arc::new(AtomicU8::new(Phase::Migrating as u8)); + let mut child = match spawn_process( + shared, + python, + ledger_path, + "migrate_probe", + &application_name, + Some(phase.clone()), + ) { + Ok(child) => child, + Err(error) => { + shared.incomplete( + format!("Python migrate probe could not start: {error}"), + vec![], + ); + return; + } + }; + let pool = match super::super::db::connect_pool( + &shared.config.database_url, + "soak_python_probe_watch", + 1, + ) + .await + { + Ok(pool) => pool, + Err(error) => { + shared.incomplete( + format!("Python probe observer could not connect: {error}"), + vec![], + ); + let _ = child.kill().await; + return; + } + }; + let mut evidence = Evidence::default(); + let deadline = Instant::now() + crate::ci_timing::scaled_timeout(Duration::from_secs(45)); + let exit_success = loop { + match child.try_wait() { + Ok(Some(status)) => { + if !status.success() { + evidence.incomplete(format!("Python migrate probe exited with {status}")); + } + break status.success(); + } + Err(error) => { + evidence.incomplete(format!( + "Python migrate probe exit could not be read: {error}" + )); + let _ = child.kill().await; + break false; + } + Ok(None) if Instant::now() >= deadline => { + evidence + .incomplete("Python migrate-then-query probe did not exit within its deadline"); + let _ = child.kill().await; + break false; + } + Ok(None) => {} + } + // Capture readiness before starting the query: a MIGRATED marker that + // arrives during it must not relabel an earlier migration snapshot. + let before_query = Phase::from_byte(phase.load(Ordering::Acquire)); + evidence.observe( + before_query, + python_probe::sessions(&pool, &application_name).await, + ); + tokio::time::sleep(Duration::from_millis(500)).await; + }; + evidence + .finish_exit(&phase, exit_success, || { + python_probe::sessions(&pool, &application_name) + }) + .await; + stats.add("python_migrate_probes", 1); + stats.set_max("python_probe_peak_sessions", evidence.peak_sessions); + stats.set_max( + "python_probe_migrating_idle_transactions", + evidence.migrating_idle_transactions, + ); + stats.add("python_probe_ready_samples", evidence.ready_samples as i64); + if !evidence.complete() { + shared.incomplete( + format!( + "Python migrate probe {application_name} has incomplete evidence: {evidence:?}" + ), + vec![], + ); + } + if evidence.violation() { + shared.violation( + "python_connections_stranded", + format!("Python migrate probe {application_name}: peak={} against pool4, idle-in-transaction after ready={}, sessions after exit={:?}", evidence.peak_sessions, evidence.ready_idle_transactions, evidence.remaining), + vec![], + ); + } + pool.close().await; +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Python; + let python = shared + .config + .python_bin + .clone() + .ok_or("the python persona needs AWA_SOAK_PYTHON, the interpreter of a virtualenv built with `maturin develop` (see scripts/persona-soak.sh)")?; + let shared_queue = shared.queues.register(QUEUE_SHARED, persona, false); + shared.queues.register(QUEUE_WORK, persona, false); + shared.queues.register(QUEUE_DEADLINE, persona, false); + + // The Rust runtime that shares the queue with the Python workers. + let rust_shared = shared.clone(); + let factory: super::super::db::ClientFactory = Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(QUEUE_SHARED, queue(6)) + .register_worker(SharedQueueWorker { + shared: rust_shared.clone(), + }); + attach_events( + tune(builder, &rust_shared), + &rust_shared, + &["py_shared_job"], + ) + .build() + .expect("build python shared-queue client") + }); + let fleet = Fleet::start( + &shared, + "pyrust", + persona, + 1, + 8, + 2, + vec![(shared_queue, 6)], + factory, + ) + .await?; + + std::fs::create_dir_all(&shared.config.report_dir).map_err(|error| error.to_string())?; + let ledger_path = shared.config.report_dir.join("python-ledger.jsonl"); + let _ = std::fs::remove_file(&ledger_path); + let child = spawn(&shared, &python, &ledger_path, "full")?; + let handle = Arc::new(PythonHandle { + process: tokio::sync::Mutex::new(PythonProcess { + child: Some(child), + ledger_path: ledger_path.clone(), + python: python.clone(), + restarts: 0, + }), + }); + shared + .python + .lock() + .expect("python handle") + .replace(handle.clone()); + shared.note("python", format!("{} started; queues {QUEUE_WORK}, {QUEUE_DEADLINE}, {QUEUE_SHARED} (shared with a Rust runtime)", python.display())); + + let tail_pool = + super::super::db::connect_pool(&shared.config.database_url, "soak_python_tail", 1) + .await + .map_err(|error| error.to_string())?; + let mut tail_tick = 0u64; + let mut probes = tokio::task::JoinSet::new(); + let mut importer = Importer { + path: ledger_path.clone(), + offset: 0, + rolled_back: HashSet::new(), + }; + let pace = (shared.config.load.as_secs_f64() / 900.0).clamp(0.1, 1.0); + let mut next_probe = Instant::now() + Duration::from_secs_f64(30.0 * pace.max(0.3)); + let probe_every = Duration::from_secs_f64(150.0 * pace.max(0.3)); + while !shared.load_stop.sleep(Duration::from_secs(1)).await { + importer.import(&shared); + tail_terminal_rows(&shared, &tail_pool, &mut tail_tick).await; + handle.ensure_running(&shared).await; + if Instant::now() >= next_probe { + next_probe = Instant::now() + probe_every; + let (shared, python, ledger_path) = + (shared.clone(), python.clone(), ledger_path.clone()); + probes.spawn(async move { migrate_probe(&shared, &python, &ledger_path).await }); + } + } + + // Stop producing immediately; outstanding migration probes must not delay + // this signal or the liveness/evidence checks during the fixed drain grace. + let errors = begin_drain( + async { + if let Some(pid) = handle + .process + .lock() + .await + .child + .as_ref() + .and_then(Child::id) + { + match std::process::Command::new("kill") + .args(["-USR1", &pid.to_string()]) + .status() + { + Ok(status) if status.success() => {} + result => shared.incomplete( + format!("Python stop-production signal failed: {result:?}"), + vec![], + ), + } + } + }, + &mut probes, + ) + .await; + for error in errors { + shared.incomplete(format!("Python migration probe failed: {error}"), vec![]); + } + shared.note("python", "load stopped; draining"); + while !shared.drain_stop.sleep(Duration::from_secs(1)).await { + handle.ensure_running(&shared).await; + importer.import(&shared); + tail_terminal_rows(&shared, &tail_pool, &mut tail_tick).await; + for error in finished_probe_errors(&mut probes) { + shared.incomplete(format!("Python migration probe failed: {error}"), vec![]); + } + } + // Finished tasks retain their result. Only a task actually cancelled at + // shutdown lacks the required evidence; do not add time to the drain grace. + probes.abort_all(); + while let Some(result) = probes.join_next().await { + if let Err(error) = result { + shared.incomplete( + format!("Python migration probe unfinished or failed at drain shutdown: {error}"), + vec![], + ); + } + } + importer.import(&shared); + if let Some(mut child) = handle.process.lock().await.child.take() { + let _ = child.kill().await; + } + shared.python.lock().expect("python handle").take(); + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + importer.import(&shared); + check_rolled_back(&shared, &importer.rolled_back).await?; + Ok(()) +} + +/// Python jobs report no lifecycle events, and ring rotation prunes their terminal rows +/// within seconds, so their terminal outcomes are read from the store as they appear. +async fn tail_terminal_rows(shared: &Arc, pool: &sqlx::PgPool, tick: &mut u64) { + *tick += 1; + if !(*tick).is_multiple_of(3) { + return; + } + let schema = shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".into()); + let queues = vec![ + QUEUE_WORK.to_string(), + QUEUE_DEADLINE.to_string(), + QUEUE_SHARED.to_string(), + ]; + let rows: Result, _> = sqlx::query_as(&format!( + "SELECT job_id, run_lease, state::text FROM {schema}.done_entries WHERE queue = ANY($1) AND finalized_at > now() - interval '25 seconds' + UNION ALL SELECT job_id, run_lease, 'failed' FROM {schema}.dlq_entries WHERE queue = ANY($1) AND dlq_at > now() - interval '25 seconds' + UNION ALL SELECT b.job_id, b.run_lease, 'completed' FROM {schema}.receipt_completion_batches c, LATERAL unnest(c.job_ids, c.run_leases) AS b(job_id, run_lease) + WHERE c.queue = ANY($1) AND c.finalized_at > now() - interval '25 seconds'" + )) + .bind(&queues) + .fetch_all(pool) + .await; + let rows = match rows { + Ok(rows) => rows, + Err(error) => { + shared.incomplete( + format!("Python terminal evidence query failed: {error}"), + vec![], + ); + shared.stats(Persona::Python).add("python_tail_errors", 1); + if shared.stats(Persona::Python).get("python_tail_errors") <= 3 { + shared.note("python", format!("terminal row tail failed: {error}")); + } + return; + } + }; + shared + .stats(Persona::Python) + .add("python_tail_rows", rows.len() as i64); + let now = shared.now_ms() as u32; + for (id, run_lease, state) in rows { + shared.ledger.with( + id, + || { + JobRecord::new( + Persona::Python, + u32::MAX, + now, + expect::ANY_TERMINAL, + flag::NO_EVENTS, + ) + }, + |record| { + record.note_observed_terminal(run_lease, &state); + }, + ); + } +} + +/// A transaction rolled back by the Python application must leave no job anywhere. +async fn check_rolled_back(shared: &Arc, ids: &HashSet) -> Result<(), String> { + if ids.is_empty() { + return Ok(()); + } + let schema = shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".into()); + let pool = super::super::db::connect_pool(&shared.config.database_url, "soak_python_check", 1) + .await + .map_err(|error| error.to_string())?; + let ids: Vec = ids.iter().copied().collect(); + let mut visible: Vec = Vec::new(); + for chunk in ids.chunks(5000) { + let found: Vec = sqlx::query_scalar(&format!( + "SELECT job_id FROM {schema}.ready_entries WHERE job_id = ANY($1) + UNION SELECT job_id FROM {schema}.deferred_jobs WHERE job_id = ANY($1) + UNION SELECT job_id FROM {schema}.done_entries WHERE job_id = ANY($1) + UNION SELECT job_id FROM {schema}.dlq_entries WHERE job_id = ANY($1)" + )) + .bind(chunk) + .fetch_all(&pool) + .await + .map_err(|error| error.to_string())?; + visible.extend(found); + } + let executed: Vec = ids + .iter() + .copied() + .filter(|id| { + shared + .ledger + .get(*id) + .is_some_and(|record| record.executions > 0) + }) + .collect(); + shared + .stats(Persona::Python) + .add("python_rolled_back_checked", ids.len() as i64); + if !visible.is_empty() || !executed.is_empty() { + shared.violation( + "rolled_back_job_visible", + format!("{} job(s) enqueued in rolled-back Python transactions exist in the store, {} of them executed", visible.len(), executed.len()), + visible.iter().take(10).map(|id| format!("{{\"job_id\":{id}}}")).collect(), + ); + } + pool.close().await; + Ok(()) +} + +#[cfg(test)] +mod lifecycle_tests { + use super::*; + + fn finished_child() -> Result { + Command::new(std::env::current_exe().map_err(|error| error.to_string())?) + .arg("--list") + .stdout(Stdio::null()) + .spawn() + .map_err(|error| error.to_string()) + } + + #[tokio::test] + async fn dead_child_is_restarted_for_the_drain_liveness_loop() { + let mut child = Some(finished_child().unwrap()); + child.as_mut().unwrap().wait().await.unwrap(); + let result = restart_if_needed(&mut child, finished_child).unwrap(); + assert!(result.is_some(), "a dead drain worker needs a replacement"); + assert!(child.is_some()); + child.as_mut().unwrap().wait().await.unwrap(); + } + + #[tokio::test] + async fn outstanding_probe_does_not_delay_stop_production_or_worker_recovery() { + let (release, blocked) = tokio::sync::oneshot::channel(); + let (stopped, mut stop_seen) = tokio::sync::oneshot::channel(); + let mut probes = tokio::task::JoinSet::new(); + probes.spawn(async move { + let _ = blocked.await; + }); + let errors = tokio::time::timeout( + crate::ci_timing::scaled_timeout(Duration::from_secs(1)), + begin_drain( + async { + stopped.send(()).unwrap(); + }, + &mut probes, + ), + ) + .await + .expect("stop production must precede waiting on a blocked probe"); + assert!(errors.is_empty()); + assert!(stop_seen.try_recv().is_ok()); + assert!( + !probes.is_empty(), + "the probe remains blocked while drain proceeds" + ); + release.send(()).unwrap(); + probes.join_next().await.unwrap().unwrap(); + } + + #[test] + fn failed_restart_is_reported_instead_of_losing_the_process_silently() { + let mut child = None; + let result = restart_if_needed(&mut child, || Err("injected spawn failure".into())); + assert_eq!(result.unwrap_err(), "injected spawn failure"); + assert!(child.is_none()); + } +} diff --git a/awa/tests/persona_soak/personas/sam.rs b/awa/tests/persona_soak/personas/sam.rs new file mode 100644 index 000000000..b70876db0 --- /dev/null +++ b/awa/tests/persona_soak/personas/sam.rs @@ -0,0 +1,359 @@ +//! Sam, ordered streams: a queue with enqueue shards and ordering keys, a striped +//! queue (`queue_stripe_count`), scheduled jobs and a rate-limited queue. Stresses +//! routing, promotion and fairness. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{attach_events, backoff, queue, register_rows, tune}; +use awa::model::{insert_many, QueueStorage, QueueStorageConfig}; +use awa::{ + Client, InsertOpts, JobArgs, JobContext, JobError, JobResult, QueueConfig, RateLimit, Worker, +}; +use chrono::Utc; +use serde::{Deserialize, Serialize}; +use std::collections::{HashMap, HashSet}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +pub const ORDERED_QUEUE: &str = "awa_soak_sam_ord"; +pub const RATE_QUEUE: &str = "awa_soak_sam_rate"; +pub const SCHEDULED_QUEUE: &str = "awa_soak_sam_sched"; +pub const STRIPED_QUEUE: &str = "awa_soak_sam_stripe"; +const KEYS: u64 = 64; +const ENQUEUE_SHARDS: i16 = 4; +const STRIPES: usize = 4; +const RATE_PER_SECOND: f64 = 15.0; +/// Starts this far behind the newest start for the same key count as order inversions. +const INVERSION_WINDOW: u64 = 64; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct SamJob { + pub stream: String, + pub key: u64, + pub seq: u64, +} + +#[derive(Default)] +struct Order { + newest_start: HashMap, + inversions: u64, + worst_distance: u64, + striped_seen: HashSet, + striped_repeats: u64, +} + +struct SamWorker { + shared: Arc, + order: Arc>, +} + +#[async_trait::async_trait] +impl Worker for SamWorker { + fn kind(&self) -> &'static str { + "sam_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self.shared.begin_execution(Persona::Sam, &ctx.job, false); + let args: SamJob = serde_json::from_value(ctx.job.args.clone()) + .map_err(|error| JobError::terminal(format!("bad args: {error}")))?; + { + let mut order = self.order.lock().expect("order"); + if args.stream == "ordered" { + let newest = order.newest_start.entry(args.key).or_insert(0); + if args.seq + INVERSION_WINDOW < *newest { + let distance = *newest - args.seq; + order.inversions += 1; + order.worst_distance = order.worst_distance.max(distance); + } else { + *newest = (*newest).max(args.seq); + } + } else if args.stream == "striped" && !order.striped_seen.insert(args.seq) { + order.striped_repeats += 1; + } + } + tokio::time::sleep(Duration::from_millis(5 + args.seq % 40)).await; + if ctx.is_cancelled() { + return Err(JobError::retryable_msg("interrupted")); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn plain_factory(shared: Arc, order: Arc>) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(ORDERED_QUEUE, queue(12)) + .queue(SCHEDULED_QUEUE, queue(8)) + .queue( + RATE_QUEUE, + QueueConfig { + max_workers: 8, + rate_limit: Some(RateLimit { + max_rate: RATE_PER_SECOND, + burst: RATE_PER_SECOND as u32, + }), + ..QueueConfig::default() + }, + ) + .register_worker(SamWorker { + shared: shared.clone(), + order: order.clone(), + }); + attach_events(tune(builder, &shared), &shared, &["sam_job"]) + .build() + .expect("build sam client") + }) +} + +fn striped_config() -> QueueStorageConfig { + QueueStorageConfig { + queue_stripe_count: STRIPES, + ..QueueStorageConfig::default() + } +} + +fn striped_factory(shared: Arc, order: Arc>) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(STRIPED_QUEUE, queue(16)) + .queue_storage( + striped_config(), + Duration::from_millis(1_000), + Duration::from_millis(1_000), + ) + .register_worker(SamWorker { + shared: shared.clone(), + order: order.clone(), + }); + attach_events(tune(builder, &shared), &shared, &["sam_job"]) + .build() + .expect("build sam stripe client") + }) +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Sam; + let ordered = shared.queues.register(ORDERED_QUEUE, persona, false); + let scheduled = shared.queues.register(SCHEDULED_QUEUE, persona, false); + let rated = shared.queues.register(RATE_QUEUE, persona, true); + let striped = shared.queues.register(STRIPED_QUEUE, persona, false); + + let setup_pool = connect_pool(&shared.config.database_url, "soak_sam_setup", 1) + .await + .map_err(|error| error.to_string())?; + 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(ORDERED_QUEUE) + .bind(ENQUEUE_SHARDS) + .execute(&setup_pool) + .await + .map_err(|error| format!("set enqueue_shards: {error}"))?; + setup_pool.close().await; + + let order = Arc::new(Mutex::new(Order::default())); + let plain = Fleet::start( + &shared, + "sam", + persona, + 2, + 14, + 2, + vec![(ordered, 12), (scheduled, 8), (rated, 8)], + plain_factory(shared.clone(), order.clone()), + ) + .await?; + let stripe_fleet = Fleet::start( + &shared, + "sam-stripe", + persona, + 2, + 10, + 2, + vec![(striped, 16)], + striped_factory(shared.clone(), order.clone()), + ) + .await?; + + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_sam_producer", + shared.budget.reserve("sam-producer", 4, 1), + ) + .await + .map_err(|error| error.to_string())?; + let store = QueueStorage::new(striped_config()).map_err(|error| error.to_string())?; + let mut rng = Rng::for_task(shared.config.seed, "sam-producer"); + let scale = shared.config.scale; + let mut key_seq = vec![0u64; KEYS as usize]; + let mut stripe_seq = 0u64; + let mut seq = 0u64; + let mut rate_carry = 0.0_f64; + let tick = Duration::from_millis(100); + shared.note( + "sam", + format!("{ORDERED_QUEUE} ({ENQUEUE_SHARDS} shards, {KEYS} ordering keys), {STRIPED_QUEUE} ({STRIPES} stripes), {SCHEDULED_QUEUE}, {RATE_QUEUE} ({RATE_PER_SECOND}/s)"), + ); + + while !shared.load_stop.sleep(tick).await { + // Ordered stream with ordering keys. + let ordered_count = (4.0 * scale).ceil() as usize; + let mut params = Vec::new(); + for _ in 0..ordered_count { + let key = rng.below(KEYS); + key_seq[key as usize] += 1; + params.push( + awa::insert::params_with( + &SamJob { + stream: "ordered".into(), + key, + seq: key_seq[key as usize], + }, + InsertOpts { + queue: ORDERED_QUEUE.into(), + ordering_key: Some(key.to_be_bytes().to_vec()), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared + .stats(persona) + .add("enqueued_ordered", rows.len() as i64); + register_rows(&shared, persona, &rows, |_| expect::COMPLETED, 0); + } + Err(error) => backoff(&shared, error, persona).await, + } + + // Scheduled jobs a few seconds out, and the rate-limited stream. + let mut params = Vec::new(); + for _ in 0..(1.0 * scale).ceil() as usize { + seq += 1; + params.push( + awa::insert::params_with( + &SamJob { + stream: "scheduled".into(), + key: 0, + seq, + }, + InsertOpts { + queue: SCHEDULED_QUEUE.into(), + run_at: Some( + Utc::now() + + chrono::Duration::milliseconds(rng.range(1_000, 12_000) as i64), + ), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + } + // Stay below the queue's rate limit so a backlog does not build by design. + rate_carry += RATE_PER_SECOND * 0.8 * tick.as_secs_f64(); + while rate_carry >= 1.0 { + rate_carry -= 1.0; + seq += 1; + params.push( + awa::insert::params_with( + &SamJob { + stream: "rated".into(), + key: 0, + seq, + }, + InsertOpts { + queue: RATE_QUEUE.into(), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared + .stats(persona) + .add("enqueued_sched_rate", rows.len() as i64); + register_rows(&shared, persona, &rows, |_| expect::COMPLETED, 0); + } + Err(error) => backoff(&shared, error, persona).await, + } + + // Striped queue, enqueued through a store that shares the workers' routing. + let mut params = Vec::new(); + for _ in 0..(3.0 * scale).ceil() as usize { + stripe_seq += 1; + params.push( + awa::insert::params_with( + &SamJob { + stream: "striped".into(), + key: 0, + seq: stripe_seq, + }, + InsertOpts { + queue: STRIPED_QUEUE.into(), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + } + match store.enqueue_params_batch(&producer_pool, ¶ms).await { + Ok(count) => shared.stats(persona).add("enqueued_striped", count as i64), + Err(error) => backoff(&shared, error, persona).await, + } + } + + shared.note( + "sam", + format!("load stopped: {stripe_seq} striped jobs; draining"), + ); + shared.drain_stop.wait().await; + plain.shutdown_all(Duration::from_secs(20), &shared).await; + stripe_fleet + .shutdown_all(Duration::from_secs(20), &shared) + .await; + producer_pool.close().await; + + let order = order.lock().expect("order"); + let executed_striped = order.striped_seen.len() as u64; + let missing: Vec = (1..=stripe_seq) + .filter(|seq| !order.striped_seen.contains(seq)) + .take(20) + .collect(); + let stats = shared.stats(persona); + stats.add("striped_executed_distinct", executed_striped as i64); + stats.add("striped_repeat_executions", order.striped_repeats as i64); + stats.add("order_inversions", order.inversions as i64); + stats.set_max("order_worst_distance", order.worst_distance as i64); + if executed_striped < stripe_seq { + let message = format!( + "{} of {stripe_seq} striped jobs never executed", + stripe_seq - executed_striped + ); + let rows = vec![format!( + "{{\"queue\":\"{STRIPED_QUEUE}\",\"missing_seq_sample\":{missing:?}}}" + )]; + if shared.quiesced.load(std::sync::atomic::Ordering::Relaxed) { + shared.violation("lost_job", message, rows); + } else { + shared.finding( + "striped_backlog_at_end", + format!("{message}; the system had not quiesced"), + rows, + ); + } + } + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/shardops.rs b/awa/tests/persona_soak/personas/shardops.rs new file mode 100644 index 000000000..a462c632b --- /dev/null +++ b/awa/tests/persona_soak/personas/shardops.rs @@ -0,0 +1,991 @@ +//! Shard operator: an operator changes `awa.queue_meta.enqueue_shards` on two hot queues +//! while SQL-path producers (`insert_many`, `insert_many_copy`), `QueueStorage` producers, +//! promotion (scheduled jobs and retries) and claimers run, as different non-superuser +//! roles (see `roles.rs`). +//! +//! The operator raises a queue from 1 to 4 to 8 shards, lowers it to 2 and back to 1, +//! sometimes running `prepare_queue` first and sometimes not, restarting some runtimes +//! before others so cached shard counts disagree across the fleet, and, for some raises, +//! holding the producers back until the maintenance leader's runtime has created the new +//! lanes so a different role is the next to use them. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::flag; +use super::super::registry::QueueInfo; +use super::super::rng::Rng; +use super::super::roles; +use super::super::shard_oracle; +use super::super::shared::Shared; +use super::super::stop::Stop; +use super::{attach_events, backoff, queue, register_rows, tune, ANY}; +use awa::model::{insert_many, insert_many_copy_from_pool, QueueStorage, QueueStorageConfig}; +use awa::{Client, InsertOpts, InsertParams, JobArgs, JobContext, JobError, JobResult, Worker}; +use chrono::Utc; +use serde::{Deserialize, Serialize}; +use sqlx::PgPool; +use std::sync::atomic::Ordering; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +pub const HOT_QUEUE: &str = "awa_soak_shardops_hot"; +pub const WARM_QUEUE: &str = "awa_soak_shardops_warm"; +const QUEUES: [&str; 2] = [HOT_QUEUE, WARM_QUEUE]; +const KEYS: u64 = 16; +const WORKERS: u32 = 8; +/// The sequence of shard counts each queue is taken through, repeatedly. +const STEPS: [(i16, i16); 4] = [(1, 4), (4, 8), (8, 2), (2, 1)]; +const MAINTENANCE_LEADER_LOCK_KEY: i64 = 0x4157_415f_4d41_494e; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct ShardOpsJob { + pub seq: u64, + pub mode: String, + pub key: Option, + pub key_seq: u64, +} + +struct ShardOpsWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for ShardOpsWorker { + fn kind(&self) -> &'static str { + "shard_ops_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self + .shared + .begin_execution(Persona::ShardOps, &ctx.job, false); + let args: ShardOpsJob = serde_json::from_value(ctx.job.args.clone()) + .map_err(|error| JobError::terminal(format!("bad args: {error}")))?; + self.shared.shardops.note_execution( + args.seq, + ctx.job.id, + args.key.map(|key| (key, args.key_seq)), + ); + tokio::time::sleep(Duration::from_millis(5 + args.seq % 20)).await; + if args.mode == "retry" && ctx.job.attempt <= 1 { + return Ok(JobResult::RetryAfter(Duration::from_millis( + 1_000 + args.seq % 2_000, + ))); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn client_factory(shared: Arc) -> ClientFactory { + Arc::new(move |pool| { + let builder = Client::builder(pool) + .queue(HOT_QUEUE, queue(WORKERS)) + .queue(WARM_QUEUE, queue(WORKERS)) + .register_worker(ShardOpsWorker { + shared: shared.clone(), + }); + attach_events( + // Retrying election quickly lets these runtimes, which run as application + // roles, take maintenance leadership from the other personas' runtimes. + tune(builder, &shared).leader_election_interval(Duration::from_millis(100)), + &shared, + &["shard_ops_job"], + ) + .build() + .expect("build shard-ops client") + }) +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum Mode { + Plain, + Scheduled, + Retry, +} + +impl Mode { + fn roll(rng: &mut Rng) -> Mode { + match rng.below(100) { + 0..=11 => Mode::Retry, + 12..=31 => Mode::Scheduled, + _ => Mode::Plain, + } + } + + fn name(self) -> &'static str { + match self { + Mode::Plain => "plain", + Mode::Scheduled => "scheduled", + Mode::Retry => "retry", + } + } + + fn executions(self) -> u8 { + if self == Mode::Retry { + 2 + } else { + 1 + } + } +} + +struct Built { + params: InsertParams, + seq: u64, + executions: u8, +} + +fn build( + shared: &Shared, + rng: &mut Rng, + queue: &str, + mode: Mode, + key: Option<(u64, u64)>, + run_in: Option, +) -> Result { + let seq = shared.shardops.next_seq(); + let run_at = match (mode, run_in) { + (_, Some(delay)) => { + Some(Utc::now() + chrono::Duration::from_std(delay).unwrap_or_default()) + } + (Mode::Scheduled, None) => { + Some(Utc::now() + chrono::Duration::milliseconds(rng.range(800, 6_000) as i64)) + } + _ => None, + }; + let params = awa::insert::params_with( + &ShardOpsJob { + seq, + mode: mode.name().to_string(), + key: key.map(|(key, _)| key), + key_seq: key.map_or(0, |(_, key_seq)| key_seq), + }, + InsertOpts { + queue: queue.to_string(), + max_attempts: 5, + run_at, + ordering_key: key.map(|(key, _)| key.to_be_bytes().to_vec()), + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?; + Ok(Built { + params, + seq, + executions: mode.executions(), + }) +} + +fn pick_queue<'a>(rng: &mut Rng) -> &'a str { + QUEUES[rng.below(QUEUES.len() as u64) as usize] +} + +fn count_enqueued(shared: &Shared, queues: impl Iterator) { + let mut count = 0; + for name in queues { + count += 1; + if let Some(queue) = shared.queues.get(&name) { + queue.enqueued.fetch_add(1, Ordering::Relaxed); + } + } + shared.stats(Persona::ShardOps).add("enqueued", count); +} + +/// Keeps the operator from changing the shard count while an insert is mid-flight +/// during the leader-first hold. +struct InsertGuard<'a>(&'a Shared); + +impl<'a> InsertGuard<'a> { + fn new(shared: &'a Shared) -> Self { + shared + .shardops + .producers_inserting + .fetch_add(1, Ordering::SeqCst); + Self(shared) + } +} + +impl Drop for InsertGuard<'_> { + fn drop(&mut self) { + self.0 + .shardops + .producers_inserting + .fetch_sub(1, Ordering::SeqCst); + } +} + +async fn enqueue_failed( + shared: &Shared, + role: &str, + operation: &str, + error: impl std::fmt::Display, +) { + let message = error.to_string(); + let short: String = message.chars().take(120).collect(); + shared.shardops.violation_once( + shared, + "shardops_enqueue_error", + &format!("{role}:{operation}:{short}"), + format!("role {role} failed {operation}: {message}"), + vec![format!( + "{{\"role\":\"{role}\",\"statement\":\"{operation}\",\"error\":{}}}", + serde_json::Value::String(message.clone()) + )], + ); + backoff(shared, &message, Persona::ShardOps).await; +} + +/// A producer's `QueueStorage` handle. A runtime's cached shard count is only dropped +/// by a restart, so the operator replaces the handle to restart the producer. +struct Stores { + handles: Mutex>>, +} + +impl Stores { + fn new(count: usize) -> Result { + let handles = (0..count) + .map(|_| fresh_store()) + .collect::, _>>()?; + Ok(Self { + handles: Mutex::new(handles), + }) + } + + fn get(&self, index: usize) -> Arc { + self.handles.lock().expect("stores")[index].clone() + } + + fn renew(&self, index: usize) { + if let Ok(store) = fresh_store() { + self.handles.lock().expect("stores")[index] = store; + } + } + + fn len(&self) -> usize { + self.handles.lock().expect("stores").len() + } +} + +fn fresh_store() -> Result, String> { + QueueStorage::new(QueueStorageConfig::default()) + .map(Arc::new) + .map_err(|error| error.to_string()) +} + +/// A SQL-path producer: `insert_many`, or `insert_many_copy` when `copy` is set. +async fn sql_producer(shared: Arc, pool: PgPool, role: &'static str, copy: bool) { + let label = if copy { + "shardops-copy" + } else { + "shardops-sql" + }; + let mut rng = Rng::for_task(shared.config.seed, label); + let scale = shared.config.shardops_scale; + let (tick, per_tick) = if copy { + (Duration::from_secs(1), (40.0 * scale).ceil() as usize) + } else { + (Duration::from_millis(100), (4.0 * scale).ceil() as usize) + }; + while !shared.load_stop.sleep(tick).await { + if shared.shardops.producers_held() { + continue; + } + let _guard = InsertGuard::new(&shared); + let mut built = Vec::with_capacity(per_tick); + for _ in 0..per_tick { + let queue = pick_queue(&mut rng); + let mode = Mode::roll(&mut rng); + match build(&shared, &mut rng, queue, mode, None, None) { + Ok(job) => built.push(job), + Err(error) => { + shared.violation("shardops_build_failed", error, vec![]); + return; + } + } + } + let params: Vec = built.iter().map(|job| job.params.clone()).collect(); + let result = if copy { + insert_many_copy_from_pool(&pool, ¶ms).await + } else { + insert_many(&pool, ¶ms).await + }; + match result { + Ok(rows) => { + for job in &built { + shared.shardops.note_enqueued(job.seq, job.executions); + } + count_enqueued(&shared, rows.iter().map(|row| row.queue.clone())); + register_rows(&shared, Persona::ShardOps, &rows, |_| ANY, 0); + } + Err(error) => { + enqueue_failed( + &shared, + role, + if copy { + "insert_many_copy" + } else { + "insert_many" + }, + error, + ) + .await + } + } + } +} + +/// A producer that enqueues through `QueueStorage`, the runtime's own path. +async fn store_producer( + shared: Arc, + pool: PgPool, + role: &'static str, + stores: Arc, + store_index: usize, +) { + let mut rng = Rng::for_task(shared.config.seed, "shardops-store"); + let per_tick = (3.0 * shared.config.shardops_scale).ceil() as usize; + while !shared.load_stop.sleep(Duration::from_millis(100)).await { + if shared.shardops.producers_held() { + continue; + } + let _guard = InsertGuard::new(&shared); + let mut built = Vec::with_capacity(per_tick); + for _ in 0..per_tick { + let queue = pick_queue(&mut rng); + let mode = Mode::roll(&mut rng); + match build(&shared, &mut rng, queue, mode, None, None) { + Ok(job) => built.push(job), + Err(error) => { + shared.violation("shardops_build_failed", error, vec![]); + return; + } + } + } + let params: Vec = built.iter().map(|job| job.params.clone()).collect(); + let store = stores.get(store_index); + match store.enqueue_params_batch(&pool, ¶ms).await { + Ok(_) => { + for job in &built { + shared.shardops.note_enqueued(job.seq, job.executions); + } + count_enqueued(&shared, params.iter().map(|job| job.opts.queue.clone())); + } + Err(error) => { + enqueue_failed(&shared, role, "QueueStorage::enqueue_params_batch", error).await + } + } + } +} + +/// Jobs with an ordering key, enqueued strictly in sequence per key, alternating between +/// the SQL path as one role and the `QueueStorage` path as the other. +async fn keyed_producer( + shared: Arc, + sql_pool: PgPool, + store_pool: PgPool, + stores: Arc, + store_index: usize, +) { + let mut rng = Rng::for_task(shared.config.seed, "shardops-keyed"); + let per_tick = (3.0 * shared.config.shardops_scale).ceil() as usize; + let mut key_seq = vec![0u64; KEYS as usize]; + let mut tick = 0u64; + while !shared.load_stop.sleep(Duration::from_millis(100)).await { + if shared.shardops.producers_held() { + continue; + } + tick += 1; + let _guard = InsertGuard::new(&shared); + let mut built = Vec::with_capacity(per_tick); + let mut keys = Vec::with_capacity(per_tick); + for _ in 0..per_tick { + let key = rng.below(KEYS); + key_seq[key as usize] += 1; + let queue = QUEUES[(key % 2) as usize]; + match build( + &shared, + &mut rng, + queue, + Mode::Plain, + Some((key, key_seq[key as usize])), + None, + ) { + Ok(job) => { + built.push(job); + keys.push((key, key_seq[key as usize])); + } + Err(error) => { + shared.violation("shardops_build_failed", error, vec![]); + return; + } + } + } + let params: Vec = built.iter().map(|job| job.params.clone()).collect(); + let window_before = shared.shardops.window(); + let outcome = if tick.is_multiple_of(2) { + match insert_many(&sql_pool, ¶ms).await { + Ok(rows) => { + register_rows(&shared, Persona::ShardOps, &rows, |_| ANY, 0); + Ok(rows.len()) + } + Err(error) => Err(( + roles::APP_A, + "insert_many (ordering_key)", + error.to_string(), + )), + } + } else { + let store = stores.get(store_index); + match store.enqueue_params_batch(&store_pool, ¶ms).await { + Ok(count) => Ok(count), + Err(error) => Err(( + roles::APP_B, + "QueueStorage::enqueue_params_batch (ordering_key)", + error.to_string(), + )), + } + }; + match outcome { + Ok(_) => { + let window_after = shared.shardops.window(); + let window = if window_before == window_after { + window_before + } else { + 0 + }; + for (job, (key, seq)) in built.iter().zip(&keys) { + shared.shardops.note_enqueued(job.seq, job.executions); + shared.shardops.note_keyed(*key, *seq, window); + } + shared + .stats(Persona::ShardOps) + .add("enqueued_keyed", built.len() as i64); + count_enqueued(&shared, params.iter().map(|job| job.opts.queue.clone())); + } + Err((role, operation, error)) => enqueue_failed(&shared, role, operation, error).await, + } + } +} + +/// Everything the operator acts through. +struct Operator { + shared: Arc, + fleets: Vec>, + stores: Arc, + owner_pool: PgPool, + migrator_pool: PgPool, + admin_pool: PgPool, + app_pool: PgPool, + rng: Rng, +} + +impl Operator { + fn note(&self, text: impl Into) { + self.shared.note("shards", text); + } + + fn operator_error(&self, what: &str, error: impl std::fmt::Display) { + let message = error.to_string(); + self.shared.shardops.violation_once( + &self.shared, + "shardops_operator_error", + &format!("{what}:{message}"), + format!("{what} failed: {message}"), + vec![format!( + "{{\"statement\":\"{what}\",\"error\":{}}}", + serde_json::Value::String(message.clone()) + )], + ); + } + + async fn prepare(&mut self, queue: &str, shards: i16) { + let as_migrator = self.rng.chance(50); + let (role, pool) = if as_migrator { + (roles::MIGRATOR, &self.migrator_pool) + } else { + (roles::OWNER, &self.owner_pool) + }; + self.note(format!( + "prepare_queue {queue} for {shards} shards as {role}" + )); + let store = match fresh_store() { + Ok(store) => store, + Err(error) => return self.operator_error("prepare_queue", error), + }; + if let Err(error) = store.prepare_queue(pool, queue, shards).await { + self.operator_error(&format!("prepare_queue as {role}"), error); + } + } + + async fn set_shards(&self, queue: &str, shards: i16) { + let result = sqlx::query( + "INSERT INTO awa.queue_meta (queue, enqueue_shards) VALUES ($1, $2) + ON CONFLICT (queue) DO UPDATE SET enqueue_shards = EXCLUDED.enqueue_shards", + ) + .bind(queue) + .bind(shards) + .execute(&self.owner_pool) + .await; + match result { + Ok(_) => { + let now = self.shared.now_ms(); + self.shared.shardops.plans.lock().expect("plans").insert( + queue.to_string(), + shard_oracle::QueuePlan { + configured: shards, + changed_at_ms: now, + settled_at_ms: None, + }, + ); + } + Err(error) => self.operator_error("set enqueue_shards", error), + } + } + + /// Restarts runtime `(fleet, replica)`, which drops its cached shard counts. + async fn restart(&self, fleet: usize, replica: usize) { + let Some(target) = self.fleets.get(fleet) else { + return; + }; + match target.restart(replica, Duration::from_secs(20)).await { + Ok(true) => self + .shared + .stats(Persona::ShardOps) + .add("runtime_restarts", 1), + Ok(false) => {} + Err(error) => self.operator_error("runtime restart", error), + } + } + + async fn runtimes(&self) -> Vec<(usize, usize)> { + let mut all = Vec::new(); + for (index, fleet) in self.fleets.iter().enumerate() { + for replica in fleet.live_replicas().await { + all.push((index, replica)); + } + } + all + } + + fn settle(&self, queue: &str) { + let now = self.shared.now_ms(); + if let Some(plan) = self + .shared + .shardops + .plans + .lock() + .expect("plans") + .get_mut(queue) + { + plan.settled_at_ms = Some(now); + } + self.shared.shardops.open_window(); + } + + /// The runtime holding maintenance leadership, as `(pid, role, application_name)`. + async fn leader(&self) -> Option<(i32, String, String)> { + sqlx::query_as( + "SELECT l.pid, a.usename::text, a.application_name + FROM pg_locks l JOIN pg_stat_activity a ON a.pid = l.pid + WHERE l.locktype = 'advisory' AND l.granted AND a.application_name LIKE 'soak\\_%' + AND ((l.classid::bigint << 32) | l.objid::bigint) = $1 + ORDER BY l.pid LIMIT 1", + ) + .bind(MAINTENANCE_LEADER_LOCK_KEY) + .fetch_optional(&self.admin_pool) + .await + .ok() + .flatten() + } + + /// Makes one of this persona's runtimes the maintenance leader, ending other + /// runtimes' leadership if needed. Returns its role. + async fn lead_with_own_runtime(&self) -> Option { + for _ in 0..10 { + if let Some((pid, role, application)) = self.leader().await { + if application.starts_with("soak_shardops") { + self.note(format!("maintenance leader is {application} ({role})")); + return Some(role); + } + self.note(format!("ending leadership of {application} ({role})")); + let _ = sqlx::query("SELECT pg_terminate_backend($1)") + .bind(pid) + .fetch_one(&self.admin_pool) + .await; + } + if self.shared.load_stop.sleep(Duration::from_secs(2)).await { + return None; + } + } + self.note("another persona's runtime kept maintenance leadership"); + None + } + + async fn lane_owners(&self, queue: &str, above: i16) -> String { + let owners: Result, _> = sqlx::query_as( + "SELECT pg_get_userbyid(c.relowner)::text, count(*)::bigint + FROM awa.queue_enqueue_heads h + JOIN pg_class c ON c.relname = h.seq_name + JOIN pg_namespace n ON n.oid = c.relnamespace AND n.nspname = 'awa' + WHERE h.queue = $1 AND h.enqueue_shard >= $2 + GROUP BY 1 ORDER BY 1", + ) + .bind(queue) + .bind(above) + .fetch_all(&self.admin_pool) + .await; + match owners { + Ok(owners) if owners.is_empty() => "none yet".to_string(), + Ok(owners) => owners + .into_iter() + .map(|(role, count)| format!("{role}={count}")) + .collect::>() + .join(","), + Err(error) => format!("lookup failed: {error}"), + } + } + + /// Returns false when the load stopped mid-change. + async fn change(&mut self, queue: &str, from: i16, to: i16) -> bool { + let raising = to > from; + let prepare = + raising && (!self.shared.config.shardops_schema_create || self.rng.chance(50)); + let leader_first = raising && !prepare && self.rng.chance(40); + self.note(format!( + "{queue}: {from} -> {to} shards{}{}", + if prepare { ", prepare_queue first" } else { "" }, + if leader_first { + ", leader creates the new lanes first" + } else { + "" + } + )); + self.shared.stats(Persona::ShardOps).add("shard_changes", 1); + if prepare { + self.prepare(queue, to).await; + } + self.shared.shardops.close_window(); + + if leader_first { + self.shared.shardops.hold_producers(true); + for _ in 0..40 { + if self + .shared + .shardops + .producers_inserting + .load(Ordering::SeqCst) + == 0 + { + break; + } + tokio::time::sleep(Duration::from_millis(250)).await; + } + self.set_shards(queue, to).await; + for (fleet, replica) in self.runtimes().await { + self.restart(fleet, replica).await; + } + for index in 0..self.stores.len() { + self.stores.renew(index); + } + let leader_role = self.lead_with_own_runtime().await; + // Scheduled jobs due shortly: the leader promotes them into new lanes before + // any producer has touched those lanes. + for _ in 0..6 { + let mut built = Vec::new(); + for _ in 0..4 { + let delay = Duration::from_millis(self.rng.range(1_500, 4_000)); + if let Ok(job) = build( + &self.shared, + &mut self.rng, + queue, + Mode::Scheduled, + None, + Some(delay), + ) { + built.push(job); + } + } + let params: Vec = + built.iter().map(|job| job.params.clone()).collect(); + match insert_many(&self.app_pool, ¶ms).await { + Ok(rows) => { + for job in &built { + self.shared.shardops.note_enqueued(job.seq, job.executions); + } + count_enqueued(&self.shared, rows.iter().map(|row| row.queue.clone())); + register_rows(&self.shared, Persona::ShardOps, &rows, |_| ANY, 0); + } + Err(error) => self.operator_error("seed scheduled jobs", error), + } + if self.shared.load_stop.sleep(Duration::from_secs(1)).await { + self.shared.shardops.hold_producers(false); + return false; + } + } + for _ in 0..15 { + let created: i64 = sqlx::query_scalar( + "SELECT count(*)::bigint FROM awa.queue_enqueue_heads WHERE queue = $1 AND enqueue_shard >= $2", + ) + .bind(queue) + .bind(from) + .fetch_one(&self.admin_pool) + .await + .unwrap_or(0); + if created > 0 { + break; + } + if self.shared.load_stop.sleep(Duration::from_secs(1)).await { + self.shared.shardops.hold_producers(false); + return false; + } + } + let owners = self.lane_owners(queue, from).await; + self.note(format!( + "new lanes of {queue} created before producers resumed; sequence owners: {owners} (leader role {})", + leader_role.as_deref().unwrap_or("not ours") + )); + self.shared.shardops.hold_producers(false); + self.settle(queue); + return true; + } + + self.set_shards(queue, to).await; + if self.shared.load_stop.sleep(Duration::from_secs(4)).await { + return false; + } + // Restart some runtimes and producer stores now and the rest later, so cached + // shard counts disagree across the fleet for a while. + let runtimes = self.runtimes().await; + let mut restart_now: Vec<(usize, usize)> = runtimes + .iter() + .copied() + .filter(|_| self.rng.chance(50)) + .collect(); + if restart_now.is_empty() { + restart_now.extend(runtimes.first().copied()); + } + if restart_now.len() == runtimes.len() && runtimes.len() > 1 { + restart_now.pop(); + } + let renew_now: Vec = (0..self.stores.len()) + .filter(|_| self.rng.chance(50)) + .collect(); + self.note(format!( + "{queue}: restarting {} of {} runtimes and {} of {} producer stores; the others keep the old shard count", + restart_now.len(), + runtimes.len(), + renew_now.len(), + self.stores.len() + )); + for (fleet, replica) in &restart_now { + self.restart(*fleet, *replica).await; + } + for index in &renew_now { + self.stores.renew(*index); + } + if self.shared.load_stop.sleep(Duration::from_secs(18)).await { + return false; + } + for (fleet, replica) in runtimes + .iter() + .filter(|runtime| !restart_now.contains(runtime)) + { + self.restart(*fleet, *replica).await; + } + for index in 0..self.stores.len() { + if !renew_now.contains(&index) { + self.stores.renew(index); + } + } + self.settle(queue); + let owners = self.lane_owners(queue, from.min(to)).await; + self.note(format!( + "{queue}: settled at {to} shards; lane sequence owners above shard {}: {owners}", + from.min(to) + )); + true + } + + async fn run(&mut self, stop: &Stop, period: Duration) { + if stop.sleep(Duration::from_secs(20)).await { + return; + } + let mut step = [0usize; QUEUES.len()]; + let mut turn = 0usize; + loop { + let which = turn % QUEUES.len(); + turn += 1; + let (from, to) = STEPS[step[which] % STEPS.len()]; + step[which] += 1; + let started = Instant::now(); + if !self.change(QUEUES[which], from, to).await { + return; + } + let remaining = period + .saturating_sub(started.elapsed()) + .max(Duration::from_secs(5)); + if stop.sleep(remaining).await { + return; + } + } + } +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::ShardOps; + let queue_infos: Vec> = QUEUES + .iter() + .map(|name| shared.queues.register(name, persona, false)) + .collect(); + let config = &shared.config; + + let owner_pool = roles::pool_as(config, roles::OWNER, "soak_shardops_owner", 2).await?; + let migrator_pool = + roles::pool_as(config, roles::MIGRATOR, "soak_shardops_migrator", 1).await?; + let admin_pool = connect_pool(&config.database_url, "soak_shardops_admin", 2) + .await + .map_err(|error| error.to_string())?; + + // The schema owner provisions each queue's first shard before any application role runs. + let owner_store = fresh_store()?; + for name in QUEUES { + owner_store + .prepare_queue(&owner_pool, name, 1) + .await + .map_err(|error| format!("prepare_queue {name}: {error}"))?; + shared.shardops.plans.lock().expect("plans").insert( + name.to_string(), + shard_oracle::QueuePlan { + configured: 1, + changed_at_ms: 0, + settled_at_ms: Some(0), + }, + ); + } + shared.shardops.open_window(); + + let served: Vec<(Arc, i64)> = queue_infos + .iter() + .map(|info| (info.clone(), i64::from(WORKERS))) + .collect(); + let mut fleets = Vec::new(); + for (name, role) in [("shardops-a", roles::APP_A), ("shardops-b", roles::APP_B)] { + fleets.push( + Fleet::start_with_url( + &shared, + roles::url_for_role(config, role), + name, + persona, + 2, + 10, + 2, + served.clone(), + client_factory(shared.clone()), + ) + .await?, + ); + } + + let app_a = roles::pool_as( + config, + roles::APP_A, + "soak_shardops_producer_a", + shared.budget.reserve("shardops-producer-a", 6, 2), + ) + .await?; + let app_b = roles::pool_as( + config, + roles::APP_B, + "soak_shardops_producer_b", + shared.budget.reserve("shardops-producer-b", 6, 2), + ) + .await?; + let stores = Arc::new(Stores::new(2)?); + + let period = if config.load >= Duration::from_secs(600) { + Duration::from_secs(60) + } else { + Duration::from_secs(40) + }; + let queue_names: Vec = QUEUES.iter().map(|name| name.to_string()).collect(); + shared.note( + "shards", + format!( + "queues {QUEUES:?} cycle through {STEPS:?} shards every {}s per change; roles: owner {}, migrator {}, applications {} and {}", + period.as_secs(), + roles::OWNER, + roles::MIGRATOR, + roles::APP_A, + roles::APP_B + ), + ); + + let watch_stop = Stop::new(); + let watcher = tokio::spawn(shard_oracle::watch( + shared.clone(), + admin_pool.clone(), + queue_names.clone(), + watch_stop.clone(), + )); + let producers = vec![ + tokio::spawn(sql_producer( + shared.clone(), + app_a.clone(), + roles::APP_A, + false, + )), + tokio::spawn(sql_producer( + shared.clone(), + app_b.clone(), + roles::APP_B, + true, + )), + tokio::spawn(store_producer( + shared.clone(), + app_b.clone(), + roles::APP_B, + stores.clone(), + 0, + )), + tokio::spawn(keyed_producer( + shared.clone(), + app_a.clone(), + app_b.clone(), + stores.clone(), + 1, + )), + ]; + let operator_stop = shared.load_stop.clone(); + let mut operator = Operator { + shared: shared.clone(), + fleets: fleets.clone(), + stores: stores.clone(), + owner_pool: owner_pool.clone(), + migrator_pool, + admin_pool: admin_pool.clone(), + app_pool: app_a.clone(), + rng: Rng::for_task(shared.config.seed, "shardops-operator"), + }; + operator.run(&operator_stop, period).await; + shared.shardops.hold_producers(false); + for producer in producers { + let _ = producer.await; + } + shared.note("shards", "load stopped; draining"); + + shared.drain_stop.wait().await; + shard_oracle::final_check(&shared, &admin_pool, &queue_names).await; + watch_stop.trigger(); + let _ = watcher.await; + for fleet in &fleets { + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + } + for pool in [app_a, app_b, owner_pool, admin_pool] { + pool.close().await; + } + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/tom.rs b/awa/tests/persona_soak/personas/tom.rs new file mode 100644 index 000000000..af963e56c --- /dev/null +++ b/awa/tests/persona_soak/personas/tom.rs @@ -0,0 +1,279 @@ +//! Tom, multi-tenant fan-out: many queues each at a low rate, unique jobs, cron and +//! priority aging. Stresses lane and queue cardinality, per-queue maintenance cost +//! and cron fencing across runtimes. + +use super::super::config::Persona; +use super::super::db::{connect_pool, ClientFactory}; +use super::super::fleet::Fleet; +use super::super::ledger::{expect, flag, JobRecord}; +use super::super::rng::Rng; +use super::super::shared::Shared; +use super::{attach_events, backoff, register_rows, tune}; +use awa::model::{insert_many, insert_with}; +use awa::{ + AwaError, Client, InsertOpts, JobArgs, JobContext, JobError, JobResult, PeriodicJob, + QueueConfig, UniqueOpts, Worker, +}; +use chrono::Utc; +use serde::{Deserialize, Serialize}; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +const PREFIX: &str = "awa_soak_tom_"; +const CRON_JOBS: usize = 4; + +#[derive(Debug, Clone, Serialize, Deserialize, JobArgs)] +pub struct TomJob { + pub seq: u64, + pub work_ms: u64, + pub unique_key: Option, +} + +struct TomWorker { + shared: Arc, +} + +#[async_trait::async_trait] +impl Worker for TomWorker { + fn kind(&self) -> &'static str { + "tom_job" + } + + async fn perform(&self, ctx: &JobContext) -> Result { + let _guard = self.shared.begin_execution(Persona::Tom, &ctx.job, false); + let work_ms = ctx.job.args["work_ms"].as_u64().unwrap_or(50); + tokio::time::sleep(Duration::from_millis(work_ms)).await; + if ctx.is_cancelled() { + return Err(JobError::retryable_msg("interrupted")); + } + self.shared + .ledger + .mark(ctx.job.id, flag::RETURNED_COMPLETED); + Ok(JobResult::Completed) + } +} + +fn queue_name(index: usize) -> String { + format!("{PREFIX}{index:03}") +} + +/// Four schedules firing every ten seconds into the first four tenant queues. +pub fn declare_cron(mut builder: awa::ClientBuilder) -> awa::ClientBuilder { + for cron in 0..CRON_JOBS { + let job = PeriodicJob::builder(format!("soak_tom_cron_{cron}"), "*/10 * * * * *") + .queue(queue_name(cron)) + .build(&TomJob { + seq: 0, + work_ms: 20, + unique_key: None, + }) + .expect("cron job"); + builder = builder.periodic(job); + } + builder +} + +fn client_factory(shared: Arc, queue_count: usize) -> ClientFactory { + Arc::new(move |pool| { + let mut builder = Client::builder(pool); + for index in 0..queue_count { + builder = builder.queue( + queue_name(index), + QueueConfig { + max_workers: 2, + priority_aging_interval: Duration::from_secs(5), + ..QueueConfig::default() + }, + ); + } + builder = builder.register_worker(TomWorker { + shared: shared.clone(), + }); + for cron in 0..CRON_JOBS { + let job = PeriodicJob::builder(format!("soak_tom_cron_{cron}"), "*/10 * * * * *") + .queue(queue_name(cron)) + .build(&TomJob { + seq: 0, + work_ms: 20, + unique_key: None, + }) + .expect("cron job"); + builder = builder.periodic(job); + } + attach_events(tune(builder, &shared), &shared, &["tom_job"]) + .build() + .expect("build tom client") + }) +} + +pub async fn run(shared: Arc) -> Result<(), String> { + let persona = Persona::Tom; + let queue_count = shared.config.scaled(100) as usize; + let mut served = Vec::new(); + for index in 0..queue_count { + served.push(( + shared.queues.register(&queue_name(index), persona, false), + 2, + )); + } + let fleet = Fleet::start( + &shared, + "tom", + persona, + 2, + 24, + 2, + served, + client_factory(shared.clone(), queue_count), + ) + .await?; + + let producer_pool = connect_pool( + &shared.config.database_url, + "soak_tom_producer", + shared.budget.reserve("tom-producer", 4, 1), + ) + .await + .map_err(|error| error.to_string())?; + let tick = Duration::from_millis(100); + let mut rng = Rng::for_task(shared.config.seed, "tom-producer"); + let plain_per_tick = (queue_count as f64 * 0.05).max(1.0) as usize; + shared.note("tom", format!("{queue_count} queues, {plain_per_tick} jobs per 100 ms tick, {CRON_JOBS} cron schedules every 10 s")); + let unique_double_accept = Arc::new(AtomicU64::new(0)); + let mut seq = 0u64; + let mut unique_key = 0u64; + let mut unique_tasks = tokio::task::JoinSet::new(); + + while !shared.load_stop.sleep(tick).await { + let mut params = Vec::with_capacity(plain_per_tick); + for _ in 0..plain_per_tick { + seq += 1; + let scheduled = rng.chance(20); + let run_at = scheduled.then(|| { + Utc::now() + chrono::Duration::milliseconds(rng.range(1_000, 10_000) as i64) + }); + let queue_index = rng.below(queue_count as u64) as usize; + params.push( + awa::insert::params_with( + &TomJob { + seq, + work_ms: rng.range(10, 200), + unique_key: None, + }, + InsertOpts { + queue: queue_name(queue_index), + priority: rng.range(1, 4) as i16, + run_at, + ..Default::default() + }, + ) + .map_err(|error| error.to_string())?, + ); + shared.stats(persona).add( + if scheduled { + "planned_scheduled" + } else { + "planned_immediate" + }, + 1, + ); + } + match insert_many(&producer_pool, ¶ms).await { + Ok(rows) => { + shared.stats(persona).add("enqueued", rows.len() as i64); + register_rows(&shared, persona, &rows, |_| expect::COMPLETED, 0); + } + Err(error) => backoff(&shared, error, persona).await, + } + + // One logical unique job per second: two concurrent identical inserts. + if seq.is_multiple_of(10) && unique_tasks.len() < 16 { + unique_key += 1; + let queue_index = rng.below(queue_count as u64) as usize; + let pool = producer_pool.clone(); + let shared = shared.clone(); + let double = unique_double_accept.clone(); + let key = unique_key; + let work_ms = rng.range(10, 100); + unique_tasks.spawn(async move { + let args = TomJob { + seq: key, + work_ms, + unique_key: Some(key), + }; + let opts = || InsertOpts { + queue: queue_name(queue_index), + unique: Some(UniqueOpts { + by_queue: true, + by_args: true, + by_period: None, + states: 0b0001_1111, + }), + ..Default::default() + }; + let (first, second) = tokio::join!( + insert_with(&pool, &args, opts()), + insert_with(&pool, &args, opts()) + ); + let mut accepted = 0; + for outcome in [first, second] { + match outcome { + Ok(row) => { + accepted += 1; + let queue = shared + .queues + .get(&row.queue) + .map_or(u32::MAX, |queue| queue.index); + shared.ledger.register( + row.id, + JobRecord::new( + Persona::Tom, + queue, + shared.now_ms() as u32, + expect::COMPLETED, + 0, + ), + ); + } + Err(AwaError::UniqueConflict { .. }) => { + shared.stats(Persona::Tom).add("unique_conflicts", 1) + } + Err(error) => shared.stats(Persona::Tom).add( + &format!("unique_error:{error}") + .chars() + .take(60) + .collect::(), + 1, + ), + } + } + shared.stats(Persona::Tom).add("unique_logical", 1); + if accepted > 1 { + double.fetch_add(1, Ordering::Relaxed); + shared.violation( + "unique_double_accept", + format!( + "both concurrent identical unique inserts (key {key}) were accepted" + ), + vec![format!( + "{{\"unique_key\":{key},\"queue\":\"{}\"}}", + queue_name(queue_index) + )], + ); + } + }); + } + while unique_tasks.try_join_next().is_some() {} + } + while unique_tasks.join_next().await.is_some() {} + + shared.note( + "tom", + format!("load stopped after {seq} plain jobs; draining"), + ); + shared.drain_stop.wait().await; + fleet.shutdown_all(Duration::from_secs(20), &shared).await; + producer_pool.close().await; + Ok(()) +} diff --git a/awa/tests/persona_soak/personas/ui.rs b/awa/tests/persona_soak/personas/ui.rs new file mode 100644 index 000000000..d2148fc43 --- /dev/null +++ b/awa/tests/persona_soak/personas/ui.rs @@ -0,0 +1,881 @@ +//! Admin UI users, driving the real `awa serve` HTTP server with the requests the web +//! UI makes (see `awa-ui/frontend/src/lib/api.ts`) at the UI's own refresh intervals. +//! +//! - Watchers: several dashboard tabs polling overview, queues, counts and timeseries. +//! - Support: searching and paging through jobs, opening details, DLQ, runtime and cron. +//! - UiOps: the UI's mutations (retry, cancel, pause, DLQ retry and purge, bulk actions). +//! - Impatient: requests abandoned mid-flight and bursts of refreshes. + +use super::super::adminserver::APPLICATION_NAME; +use super::super::config::Persona; +use super::super::db::connect_pool; +use super::super::ledger::flag; +use super::super::rng::{percentile, Rng}; +use super::super::shared::Shared; +use super::{alex, mia}; +use reqwest::Method; +use serde_json::{json, Value}; +use std::collections::HashMap; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +/// Latency samples per endpoint over the run, and the number of requests in flight. +#[derive(Default)] +pub struct UiMetrics { + samples: Mutex>>, + pub inflight: AtomicI64, + base_url: Mutex>, + reported_errors: Mutex>, +} + +impl UiMetrics { + pub fn set_base_url(&self, url: String) { + *self.base_url.lock().expect("ui base url") = Some(url); + } + + fn base_url(&self) -> String { + self.base_url + .lock() + .expect("ui base url") + .clone() + .unwrap_or_default() + } + + fn record(&self, endpoint: &str, at_ms: i64, latency_ms: i64) { + let mut samples = self.samples.lock().expect("ui samples"); + let list = samples.entry(endpoint.to_string()).or_default(); + if list.len() < 200_000 { + list.push((at_ms, latency_ms)); + } + } + + /// True the first few times an error class is seen for an endpoint. + fn first_errors(&self, key: &str) -> bool { + let mut reported = self.reported_errors.lock().expect("ui errors"); + let count = reported.entry(key.to_string()).or_insert(0); + *count += 1; + *count <= 3 + } +} + +pub enum Outcome { + Json(Value), + Status(u16), + Timeout, + Failed, +} + +impl Outcome { + fn json(self) -> Option { + match self { + Outcome::Json(value) => Some(value), + _ => None, + } + } +} + +pub struct Http { + client: reqwest::Client, + shared: Arc, + persona: Persona, +} + +impl Http { + pub fn new(shared: &Arc, persona: Persona) -> Self { + Self { + client: reqwest::Client::builder() + .timeout(shared.config.ui_client_timeout) + .pool_max_idle_per_host(4) + .build() + .expect("http client"), + shared: shared.clone(), + persona, + } + } + + pub async fn get(&self, endpoint: &str, path: &str) -> Outcome { + self.request(Method::GET, endpoint, path, None).await + } + + pub async fn post(&self, endpoint: &str, path: &str, body: Value) -> Outcome { + self.request(Method::POST, endpoint, path, Some(body)).await + } + + pub async fn request( + &self, + method: Method, + endpoint: &str, + path: &str, + body: Option, + ) -> Outcome { + let name = format!("{method} {endpoint}"); + let started = Instant::now(); + let at_ms = self.shared.now_ms(); + self.shared.ui.inflight.fetch_add(1, Ordering::Relaxed); + let mut request = self + .client + .request(method, format!("{}{path}", self.shared.ui.base_url())); + if let Some(body) = body { + request = request.json(&body); + } + let outcome = match request.send().await { + Ok(response) => { + let status = response.status(); + match response.text().await { + Ok(text) if status.is_success() => { + Outcome::Json(serde_json::from_str(&text).unwrap_or(Value::Null)) + } + Ok(text) => { + if status.is_server_error() && self.shared.ui.first_errors(&name) { + self.shared.finding( + "ui_server_error", + format!("{name} answered {status}"), + vec![text.chars().take(300).collect()], + ); + } + Outcome::Status(status.as_u16()) + } + Err(error) if error.is_timeout() => Outcome::Timeout, + Err(_) => Outcome::Failed, + } + } + Err(error) if error.is_timeout() => Outcome::Timeout, + Err(_) => Outcome::Failed, + }; + self.shared.ui.inflight.fetch_sub(1, Ordering::Relaxed); + let latency = started.elapsed().as_millis() as i64; + self.shared.ui.record(&name, at_ms, latency); + let stats = self.shared.stats(self.persona); + stats.observe(&name, latency); + match &outcome { + Outcome::Timeout => { + stats.add("client_timeouts", 1); + if self.shared.ui.first_errors(&format!("timeout {name}")) { + self.shared.finding( + "ui_request_hang", + format!( + "{name} did not answer within the {}s client timeout", + self.shared.config.ui_client_timeout.as_secs() + ), + vec![format!("{{\"path\":{path:?}}}")], + ); + } + } + Outcome::Failed => stats.add("connection_errors", 1), + Outcome::Status(code) => stats.add(&format!("http_{code}"), 1), + Outcome::Json(_) => {} + } + outcome + } + + /// Sends a GET and abandons it after `after`, closing the connection mid-request. + pub async fn get_abandoned(&self, path: &str, after: Duration) -> bool { + let url = format!("{}{path}", self.shared.ui.base_url()); + let send = self.client.get(url).send(); + tokio::select! { + _ = send => false, + _ = tokio::time::sleep(after) => true, + } + } +} + +fn ids_of(value: &Value) -> Vec { + value + .as_array() + .map(|rows| { + rows.iter() + .filter_map(|row| row.get("id").and_then(Value::as_i64)) + .collect() + }) + .unwrap_or_default() +} + +async fn poll_interval(http: &Http) -> Duration { + let ms = http + .get("/api/capabilities", "/api/capabilities") + .await + .json() + .and_then(|value| value.get("poll_interval_ms").and_then(Value::as_u64)) + .unwrap_or(5_000); + Duration::from_millis(ms) +} + +/// Dashboard tabs: each polls what the overview, the shell and (for some tabs) the +/// queues, kinds and cron pages poll, every server-suggested interval. +pub async fn run_watchers(shared: Arc) -> Result<(), String> { + let tabs = shared.config.scaled(4) as usize; + shared.note( + "watchers", + format!("{tabs} dashboard tabs polling the admin UI"), + ); + let mut tasks = tokio::task::JoinSet::new(); + for tab in 0..tabs { + let shared = shared.clone(); + tasks.spawn(async move { + let http = Http::new(&shared, Persona::Watchers); + let mut rng = Rng::for_task(shared.config.seed, &format!("watcher-{tab}")); + let interval = poll_interval(&http).await; + shared + .load_stop + .sleep(Duration::from_millis( + rng.below(interval.as_millis() as u64), + )) + .await; + while !shared.load_stop.is_stopped() { + // The overview page and the shell, in parallel as the browser does. + tokio::join!( + http.get("/api/stats", "/api/stats"), + http.get("/api/stats/timeseries", "/api/stats/timeseries?minutes=60"), + http.get("/api/queues", "/api/queues"), + http.get("/api/jobs", "/api/jobs?state=failed&limit=10"), + http.get("/api/runtime", "/api/runtime"), + http.get("/api/storage", "/api/storage?history=true"), + ); + match tab % 3 { + 1 => { + tokio::join!( + http.get("/api/queues/runtime", "/api/queues/runtime"), + http.get("/api/cron", "/api/cron") + ); + } + 2 => { + http.get("/api/kinds", "/api/kinds").await; + } + _ => {} + } + let jitter = rng.range(90, 110); + shared.load_stop.sleep(interval * jitter as u32 / 100).await; + } + }); + } + while tasks.join_next().await.is_some() {} + Ok(()) +} + +pub async fn run_support(shared: Arc) -> Result<(), String> { + let http = Http::new(&shared, Persona::Support); + let mut rng = Rng::for_task(shared.config.seed, "support"); + let kinds = [ + "mia_job", + "tom_job", + "priya_job", + "sam_job", + "alex_job", + "py_job", + "py_shared_job", + ]; + let states = [ + "available", + "running", + "scheduled", + "retryable", + "completed", + "failed", + "cancelled", + ]; + let mut known: Vec = Vec::new(); + shared.note( + "support", + "searching, paging and opening jobs, DLQ, runtime and cron pages", + ); + while !shared + .load_stop + .sleep(Duration::from_millis(rng.range(600, 1800))) + .await + { + match rng.below(10) { + 0 | 1 => { + let kind = rng.pick(&kinds).expect("kinds"); + known.extend(ids_of( + &http + .get("/api/jobs", &format!("/api/jobs?kind={kind}&limit=50")) + .await + .json() + .unwrap_or(Value::Null), + )); + } + 2 => { + let queues = shared.queues.all(); + if let Some(queue) = rng.pick(&queues) { + known.extend(ids_of( + &http + .get( + "/api/jobs", + &format!("/api/jobs?queue={}&limit=50", queue.name), + ) + .await + .json() + .unwrap_or(Value::Null), + )); + } + } + 3 => { + let state = rng.pick(&states).expect("states"); + known.extend(ids_of( + &http + .get("/api/jobs", &format!("/api/jobs?state={state}&limit=100")) + .await + .json() + .unwrap_or(Value::Null), + )); + } + 4 if rng.chance(30) => { + http.get("/api/jobs", "/api/jobs?tag=soak&limit=20").await; + } + 4 => {} + 5 => { + // Follow one job the way the detail page does: a fast poll for a while. + if let Some(&id) = rng.pick(&known) { + for _ in 0..5 { + http.get("/api/jobs/{id}", &format!("/api/jobs/{id}")).await; + if shared.load_stop.sleep(Duration::from_secs(2)).await { + break; + } + } + } + http.get( + "/api/jobs/{id}", + &format!("/api/jobs/{}", 9_000_000_000_i64 + rng.below(1000) as i64), + ) + .await; + } + 6 => { + // Page through a large result set with the before_id cursor. + let state = *rng.pick(&["completed", "available"]).expect("states"); + let mut before: Option = None; + for page in 0..10 { + let path = match before { + Some(id) => format!("/api/jobs?state={state}&limit=100&before_id={id}"), + None => format!("/api/jobs?state={state}&limit=100"), + }; + let endpoint = if page == 0 { + "/api/jobs (first page)" + } else { + "/api/jobs (later page)" + }; + let ids = ids_of( + &http + .get(endpoint, &path) + .await + .json() + .unwrap_or(Value::Null), + ); + match ids.iter().min() { + Some(&lowest) if ids.len() >= 100 => before = Some(lowest), + _ => break, + } + } + } + 7 => { + let dlq = http + .get("/api/dlq", "/api/dlq?limit=50") + .await + .json() + .unwrap_or(Value::Null); + let ids = ids_of(&dlq); + if let Some(last) = dlq.as_array().and_then(|rows| rows.last()) { + if let (Some(id), Some(at)) = ( + last.get("id").and_then(Value::as_i64), + last.get("dlq_at").and_then(Value::as_str), + ) { + let at = at.replace('+', "%2B"); + http.get( + "/api/dlq (later page)", + &format!("/api/dlq?limit=50&before_id={id}&before_dlq_at={at}"), + ) + .await; + } + } + if let Some(&id) = rng.pick(&ids) { + http.get("/api/dlq/{id}", &format!("/api/dlq/{id}")).await; + } + http.get("/api/dlq/depth", "/api/dlq/depth").await; + } + 8 => { + tokio::join!( + http.get("/api/runtime", "/api/runtime"), + http.get("/api/queues/runtime", "/api/queues/runtime"), + http.get("/api/cron", "/api/cron"), + ); + } + _ => { + tokio::join!( + http.get("/api/kinds", "/api/kinds"), + http.get("/api/stats/kinds", "/api/stats/kinds"), + http.get("/api/stats/queues", "/api/stats/queues"), + ); + } + } + if known.len() > 500 { + known.drain(..known.len() - 500); + } + } + Ok(()) +} + +/// POSTs a resume until it succeeds; a queue or schedule left paused looks like a stall. +async fn resume_until_ok(http: &Http, shared: &Shared, path: &str, endpoint: &str) { + for attempt in 0..8 { + if http.post(endpoint, path, json!({})).await.json().is_some() { + return; + } + tokio::time::sleep(Duration::from_secs(2 + attempt)).await; + } + shared.violation( + "resume_failed", + format!("{path} did not succeed after eight attempts"), + vec![], + ); +} + +pub async fn run_ui_operator(shared: Arc) -> Result<(), String> { + let persona = Persona::UiOps; + let http = Http::new(&shared, persona); + let mut rng = Rng::for_task(shared.config.seed, "uiops"); + shared.note("uiops", "operating through the UI's mutation endpoints"); + while !shared + .load_stop + .sleep(Duration::from_millis(rng.range(2_000, 5_000))) + .await + { + match rng.below(9) { + 0 => { + let mut candidates = shared.queues.of(Persona::Sam); + candidates.extend(shared.queues.of(Persona::Tom)); + if let Some(queue) = rng.pick(&candidates).cloned() { + queue.hold(); + let paused = http + .post( + "/api/queues/{queue}/pause", + &format!("/api/queues/{}/pause", queue.name), + json!({"paused_by": "soak-ui"}), + ) + .await + .json() + .is_some(); + if paused { + shared.stats(persona).add("pauses", 1); + shared + .stop_sleep_load(Duration::from_secs(rng.range(3, 6))) + .await; + } + // A pause whose answer was lost may still have applied, so always resume. + resume_until_ok( + &http, + &shared, + &format!("/api/queues/{}/resume", queue.name), + "/api/queues/{queue}/resume", + ) + .await; + shared.stop_sleep_load(Duration::from_secs(3)).await; + queue.release(); + } + } + 1 | 2 => { + // Cancel live jobs from a list view, singly or in bulk. + let mut candidates = shared.queues.of(Persona::Tom); + candidates.extend(shared.queues.of(Persona::Sam)); + let Some(queue) = rng.pick(&candidates).cloned() else { + continue; + }; + let state = *rng + .pick(&["available", "scheduled", "running"]) + .expect("states"); + let list = http + .get( + "/api/jobs", + &format!("/api/jobs?queue={}&state={state}&limit=20", queue.name), + ) + .await + .json() + .unwrap_or(Value::Null); + let ids: Vec = ids_of(&list).into_iter().take(6).collect(); + if ids.is_empty() { + continue; + } + let now = shared.now_ms() as u32; + if rng.chance(50) { + let completed_before = shared + .ledger + .get(ids[0]) + .is_some_and(|record| record.completed_events > 0); + if http + .post( + "/api/jobs/{id}/cancel", + &format!("/api/jobs/{}/cancel", ids[0]), + json!({}), + ) + .await + .json() + .is_some() + { + shared.record_operator_cancel(persona, ids[0], completed_before, now); + } + } else if let Some(cancelled) = http + .post( + "/api/jobs/bulk-cancel", + "/api/jobs/bulk-cancel", + json!({"ids": ids}), + ) + .await + .json() + { + for id in ids_of(&cancelled) { + shared.record_operator_cancel(persona, id, false, now); + } + } + } + 3 | 4 => { + let queue = if rng.chance(50) { + mia::INFER_QUEUE + } else { + mia::DEADLINE_QUEUE + }; + let dlq = http + .get("/api/dlq", &format!("/api/dlq?queue={queue}&limit=10")) + .await + .json() + .unwrap_or(Value::Null); + let ids = ids_of(&dlq); + let Some(&id) = rng.pick(&ids) else { continue }; + if rng.chance(60) { + shared.ledger.mark(id, flag::DLQ_RETRIED); + if http + .post( + "/api/dlq/{id}/retry", + &format!("/api/dlq/{id}/retry"), + json!({}), + ) + .await + .json() + .is_some() + { + shared.stats(persona).add("dlq_retried", 1); + } + } else { + shared.ledger.mark(id, flag::DLQ_PURGED); + let outcome = http + .request( + Method::DELETE, + "/api/dlq/{id}", + &format!("/api/dlq/{id}"), + None, + ) + .await; + if outcome.json().is_some() { + shared.stats(persona).add("dlq_purged", 1); + } + } + } + 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, + ); + http.post( + "/api/dlq/bulk-retry", + "/api/dlq/bulk-retry", + json!({"queue": mia::INFER_QUEUE}), + ) + .await; + } else { + shared.ledger.mark_where( + |record| record.persona == Persona::Mia && record.has(flag::EXHAUSTED), + flag::DLQ_PURGED, + ); + http.post( + "/api/dlq/bulk-purge", + "/api/dlq/bulk-purge", + json!({"queue": mia::TRAIN_QUEUE}), + ) + .await; + } + } + 6 => { + let name = format!("soak_tom_cron_{}", rng.below(4)); + let paused = http + .post( + "/api/cron/{name}/pause", + &format!("/api/cron/{name}/pause"), + json!({"paused_by": "soak-ui"}), + ) + .await + .json() + .is_some(); + if paused { + shared.stop_sleep_load(Duration::from_secs(8)).await; + } + resume_until_ok( + &http, + &shared, + &format!("/api/cron/{name}/resume"), + "/api/cron/{name}/resume", + ) + .await; + } + 7 => { + let name = format!("soak_tom_cron_{}", rng.below(4)); + http.post( + "/api/cron/{name}/trigger", + &format!("/api/cron/{name}/trigger"), + json!({}), + ) + .await; + } + _ => { + // Preview then submit a batch operation, as the batch operations page does. + let payload = json!({ + "op_kind": "set_priority", + "spec": { "priority": 2 }, + "filter": { "queue": alex::QUEUE_B, "state": "available" }, + "submitted_by": "soak-ui", + }); + http.post( + "/api/batch-ops/preview", + "/api/batch-ops/preview", + payload.clone(), + ) + .await; + http.post("/api/batch-ops", "/api/batch-ops", payload).await; + http.get("/api/batch-ops", "/api/batch-ops?limit=20").await; + } + } + } + Ok(()) +} + +/// Requests abandoned mid-flight and bursts of refreshes, while sampling the server's +/// database sessions: an abandoned request must not keep holding a connection. +pub async fn run_impatient(shared: Arc) -> Result<(), String> { + let persona = Persona::Impatient; + let http = Arc::new(Http::new(&shared, persona)); + let sampler_pool = connect_pool( + &shared.config.database_url, + "soak_ui_sampler", + shared.budget.reserve("ui-sampler", 1, 1), + ) + .await + .map_err(|error| error.to_string())?; + let statement_timeout = shared.config.ui_statement_timeout; + let sole_ui_persona = [Persona::Watchers, Persona::Support, Persona::UiOps] + .iter() + .all(|other| !shared.config.has(*other)); + let heavy = [ + "/api/kinds", + "/api/queues", + "/api/stats", + "/api/stats/timeseries?minutes=1440", + "/api/jobs?limit=100", + "/api/jobs?state=completed&limit=100", + "/api/dlq?limit=100", + "/api/runtime", + "/api/storage?history=true", + ]; + + let sampler = { + let shared = shared.clone(); + let pool = sampler_pool.clone(); + tokio::spawn(async move { + let mut reported = false; + loop { + if shared.drain_stop.is_stopped() { + return; + } + let row: Result<(i64, i64, f64), _> = sqlx::query_as(&format!( + "SELECT count(*)::bigint, count(*) FILTER (WHERE state = 'active')::bigint, + COALESCE(max(extract(epoch FROM now() - query_start)) FILTER (WHERE state = 'active'), 0)::float8 + FROM pg_stat_activity WHERE application_name = '{APPLICATION_NAME}'" + )) + .fetch_one(&pool) + .await; + if let Ok((sessions, active, oldest)) = row { + shared.run.set_max("ui_server_sessions_peak", sessions); + shared.run.set_max("ui_server_active_sessions_peak", active); + shared + .run + .set_max("ui_server_oldest_active_query_s", oldest as i64); + shared.run.observe("ui_server_sessions", sessions); + if oldest > statement_timeout.as_secs_f64() + 20.0 && !reported { + reported = true; + shared.violation( + "ui_query_outlived_statement_timeout", + format!("an admin server query has been active for {oldest:.0}s with a {}s statement timeout", statement_timeout.as_secs()), + vec![format!("{{\"sessions\":{sessions},\"active\":{active},\"oldest_active_s\":{oldest:.0}}}")], + ); + } + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + }) + }; + + let mut rng = Rng::for_task(shared.config.seed, "impatient"); + let mut next_burst = Instant::now() + Duration::from_secs(20); + shared.note( + "impatient", + "abandoning requests mid-flight and refreshing in bursts", + ); + while !shared + .load_stop + .sleep(Duration::from_millis(rng.range(200, 900))) + .await + { + let path = *rng.pick(&heavy).expect("paths"); + let after = Duration::from_millis(rng.range(2, 150)); + if http.get_abandoned(path, after).await { + shared.stats(persona).add("abandoned_requests", 1); + } else { + shared.stats(persona).add("completed_before_abandon", 1); + } + if Instant::now() >= next_burst { + next_burst = Instant::now() + Duration::from_secs(rng.range(30, 70)); + let mut burst = tokio::task::JoinSet::new(); + let size = shared.config.scaled(12); + for _ in 0..size { + let http = http.clone(); + let path = *rng.pick(&heavy).expect("paths"); + let abandon = rng.chance(70); + let after = Duration::from_millis(rng.range(5, 120)); + burst.spawn(async move { + if abandon { + http.get_abandoned(path, after).await + } else { + http.get(path, path).await; + false + } + }); + } + let mut abandoned = 0; + while let Some(result) = burst.join_next().await { + if result.unwrap_or(false) { + abandoned += 1; + } + } + shared.stats(persona).add("bursts", 1); + shared.stats(persona).add("burst_abandoned", abandoned); + // With no other UI user running, every query still active after the burst + // belongs to an abandoned request; time how long the server keeps running them. + if sole_ui_persona { + let released_started = Instant::now(); + loop { + let active: i64 = sqlx::query_scalar(&format!( + "SELECT count(*)::bigint FROM pg_stat_activity WHERE application_name = '{APPLICATION_NAME}' AND state = 'active'" + )) + .fetch_one(&sampler_pool) + .await + .unwrap_or(0); + if active == 0 + || released_started.elapsed() > Duration::from_secs(90) + || shared.load_stop.is_stopped() + { + break; + } + tokio::time::sleep(Duration::from_millis(500)).await; + } + let seconds = released_started.elapsed().as_secs() as i64; + shared + .stats(persona) + .observe("release_after_burst_s", seconds); + if seconds >= 5 { + shared.finding( + "ui_abandoned_request_runs_on", + format!("after a burst of abandoned requests the server kept {seconds}s of queries running (statement timeout {}s)", statement_timeout.as_secs()), + vec![], + ); + } + } + } + } + + // With every UI client stopped, the server must release its connections once the + // statements it was running reach their timeout. + let quiet_started = Instant::now(); + while shared.ui.inflight.load(Ordering::Relaxed) > 0 + && quiet_started.elapsed() < Duration::from_secs(60) + { + tokio::time::sleep(Duration::from_millis(500)).await; + } + let deadline = Instant::now() + statement_timeout + Duration::from_secs(15); + let mut active = i64::MAX; + while Instant::now() < deadline { + let row: Result<(i64, i64), _> = sqlx::query_as(&format!( + "SELECT count(*)::bigint, count(*) FILTER (WHERE state = 'active')::bigint + FROM pg_stat_activity WHERE application_name = '{APPLICATION_NAME}'" + )) + .fetch_one(&sampler_pool) + .await; + if let Ok((sessions, now_active)) = row { + active = now_active; + shared.run.add("ui_server_sessions_after_quiet", 0); + shared + .run + .set_max("ui_server_sessions_after_quiet", sessions); + if active == 0 { + break; + } + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + if active != 0 { + shared.violation( + "ui_sessions_not_released", + format!("{active} admin server database session(s) still running a query {}s after every UI client stopped", (statement_timeout + Duration::from_secs(15)).as_secs()), + vec![], + ); + } + shared.drain_stop.wait().await; + let _ = sampler.await; + sampler_pool.close().await; + Ok(()) +} + +impl Shared { + /// Sleeps unless the load phase ends first. + pub async fn stop_sleep_load(&self, duration: Duration) -> bool { + self.load_stop.sleep(duration).await + } +} + +/// After the run: latency growth and slow-endpoint findings. +pub fn analyse(shared: &Shared) { + let samples = shared.ui.samples.lock().expect("ui samples").clone(); + let run_ms = shared.config.load.as_millis() as i64; + for (endpoint, mut points) in samples { + if points.len() < 20 { + continue; + } + points.sort_by_key(|point| point.0); + let quarter = points.len() / 4; + let median = |slice: &[(i64, i64)]| { + let mut values: Vec = slice.iter().map(|point| point.1).collect(); + values.sort_unstable(); + percentile(&values, 0.5) + }; + let early = median(&points[..quarter]); + let late = median(&points[points.len() - quarter..]); + let mut all: Vec = points.iter().map(|point| point.1).collect(); + all.sort_unstable(); + let p99 = percentile(&all, 0.99); + let max = all.last().copied().unwrap_or(0); + if endpoint.starts_with("GET") && late >= 500 && late >= early.max(1) * 3 { + shared.finding( + "ui_latency_grows", + format!("{endpoint}: median latency {early} ms in the first quarter of its samples, {late} ms in the last (run {}s)", run_ms / 1000), + vec![format!("{{\"endpoint\":\"{endpoint}\",\"samples\":{},\"p99_ms\":{p99},\"max_ms\":{max}}}", points.len())], + ); + } + if p99 >= 10_000 { + shared.finding( + "ui_endpoint_slow", + format!( + "{endpoint}: p99 {p99} ms, max {max} ms over {} requests", + points.len() + ), + vec![], + ); + } + } +} diff --git a/awa/tests/persona_soak/pgserver.rs b/awa/tests/persona_soak/pgserver.rs new file mode 100644 index 000000000..10b018d99 --- /dev/null +++ b/awa/tests/persona_soak/pgserver.rs @@ -0,0 +1,99 @@ +//! Server crash detection: postmaster start time, and a scan of the server log for +//! the lines a crashed backend leaves behind. + +use super::config::SoakConfig; +use sqlx::PgPool; +use std::process::Command; + +const CRASH_PATTERNS: [&str; 6] = [ + "terminated by signal", + "TRAP:", + "server process (PID", + "PANIC:", + "all server processes terminated", + "database system was not properly shut down", +]; + +pub struct PgWatch { + pub started_at: String, + log_start_offset: u64, +} + +impl PgWatch { + pub async fn begin(config: &SoakConfig, pool: &PgPool) -> Result { + let started_at: String = sqlx::query_scalar("SELECT pg_postmaster_start_time()::text") + .fetch_one(pool) + .await?; + Ok(Self { + started_at, + log_start_offset: current_log_size(config), + }) + } + + pub async fn postmaster_start(pool: &PgPool) -> Result { + sqlx::query_scalar("SELECT pg_postmaster_start_time()::text") + .fetch_one(pool) + .await + } + + /// Log text written since the run began, or None when no log source is configured. + pub fn log_since_start(&self, config: &SoakConfig) -> Option { + if let Some(path) = &config.pg_log_file { + let bytes = std::fs::read(path).ok()?; + let from = (self.log_start_offset as usize).min(bytes.len()); + return Some(String::from_utf8_lossy(&bytes[from..]).into_owned()); + } + let container = config.pg_container.as_ref()?; + let output = match &config.pg_container_log { + Some(path) => Command::new("docker") + .args([ + "exec", + container, + "sh", + "-c", + &format!("tail -c +{} {path}", self.log_start_offset + 1), + ]) + .output() + .ok()?, + None => Command::new("docker") + .args(["logs", container]) + .output() + .ok()?, + }; + let mut text = String::from_utf8_lossy(&output.stdout).into_owned(); + text.push_str(&String::from_utf8_lossy(&output.stderr)); + Some(text) + } +} + +fn current_log_size(config: &SoakConfig) -> u64 { + if let Some(path) = &config.pg_log_file { + return std::fs::metadata(path).map(|meta| meta.len()).unwrap_or(0); + } + let (Some(container), Some(path)) = (&config.pg_container, &config.pg_container_log) else { + return 0; + }; + Command::new("docker") + .args(["exec", container, "sh", "-c", &format!("wc -c < {path}")]) + .output() + .ok() + .and_then(|output| String::from_utf8_lossy(&output.stdout).trim().parse().ok()) + .unwrap_or(0) +} + +/// Lines of `log` that indicate a crashed server process, each with a few lines of +/// context on either side. +pub fn crash_excerpts(log: &str) -> Vec { + let lines: Vec<&str> = log.lines().collect(); + let mut excerpts = Vec::new(); + let mut next_allowed = 0; + for (index, line) in lines.iter().enumerate() { + if index >= next_allowed && CRASH_PATTERNS.iter().any(|pattern| line.contains(pattern)) { + let from = index.saturating_sub(5); + let to = (index + 8).min(lines.len()); + excerpts.push(lines[from..to].join("\n")); + next_allowed = to; + } + } + excerpts +} diff --git a/awa/tests/persona_soak/python/soak_python.py b/awa/tests/persona_soak/python/soak_python.py new file mode 100644 index 000000000..ff6920c3e --- /dev/null +++ b/awa/tests/persona_soak/python/soak_python.py @@ -0,0 +1,359 @@ +"""Python side of the persona soak: producers, async workers and the migrate path. + +Started and stopped by the Rust harness. Every enqueue, handler start and handler end +is appended to a JSON-lines ledger that the harness folds into its own ledger, so the +oracle checks Python jobs like any others. + +Environment: + SOAK_DATABASE_URL database to use (an awa_soak* database) + SOAK_LEDGER ledger file, appended to + SOAK_SEED, SOAK_SCALE, SOAK_PACE + SOAK_ROLE "full" (default) or "migrate_probe" + SOAK_QUEUE_WORK, SOAK_QUEUE_DEADLINE, SOAK_QUEUE_SHARED queue names + +Signals: SIGUSR1 stops producing (workers keep running); SIGTERM exits at once. +""" + +import asyncio +import json +import os +import random +import signal +import sys +import time +from dataclasses import dataclass + +import awa + +URL = os.environ["SOAK_DATABASE_URL"] +ROLE = os.environ.get("SOAK_ROLE", "full") +SEED = int(os.environ.get("SOAK_SEED", "1")) +SCALE = float(os.environ.get("SOAK_SCALE", "1")) +PACE = float(os.environ.get("SOAK_PACE", "1")) +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") +WORK_WORKERS = 8 +DEADLINE_WORKERS = 16 +SHARED_WORKERS = 6 +DEADLINE_MS = 5000 + +# Outcomes a job may end in, matching the harness (1 completed, 2 failed, 4 cancelled). +COMPLETED, FAILED, CANCELLED = 1, 2, 4 + +ledger = open(os.environ["SOAK_LEDGER"], "a", buffering=1) if "SOAK_LEDGER" in os.environ else None + + +def log(**event): + if ledger is not None: + ledger.write(json.dumps(event, separators=(",", ":")) + "\n") + + +@dataclass +class PyJob: + seq: int + behavior: str + work_ms: int + born: float = 0.0 + + +@dataclass +class PySharedJob: + seq: int + behavior: str + work_ms: int + + +active = 0 +max_active = 0 +started = 0 +enqueued = 0 +errors = 0 + + +async def handle(job): + """One handler for every behaviour; see `plan` for what each one is for.""" + global active, max_active, started + args = job.args + attempt = job.attempt + active += 1 + max_active = max(max_active, active) + started += 1 + log(e="start", id=job.id, attempt=attempt, q=job.queue) + result = "completed" + try: + behavior = args.behavior + if behavior == "ok": + await asyncio.sleep(args.work_ms / 1000) + elif behavior == "progress": + for step in range(1, 5): + await asyncio.sleep(args.work_ms / 4000) + job.set_progress(step * 25, f"step {step}") + await job.flush_progress() + elif behavior == "retry": + await asyncio.sleep(0.05) + if attempt < 3: + result = "raise" + raise RuntimeError("transient failure") + elif behavior == "doomed": + result = "raise" + raise RuntimeError("permanent failure") + elif behavior == "retry_after": + if attempt == 1: + result = "retry_after" + return awa.RetryAfter(seconds=1.0) + elif behavior == "snooze": + # A snooze does not count as an attempt, so the wait is keyed on the clock. + if time.time() < args.born + 3.0: + result = "snooze" + return awa.Snooze(seconds=1.0) + elif behavior == "cancel": + result = "cancel" + return awa.Cancel("rejected") + elif behavior == "long": + # Overruns the queue deadline on the first attempt only. + await asyncio.sleep(args.work_ms / 1000 if attempt == 1 else 0.1) + elif behavior == "ignore_deadline": + # Documented behaviour: after the deadline rescue the worker slot is freed + # but this coroutine keeps running, so more handlers than max_workers run. + await asyncio.sleep(args.work_ms / 1000) + else: + await asyncio.sleep(args.work_ms / 1000) + return None + except Exception: + raise + finally: + active -= 1 + log(e="end", id=job.id, attempt=attempt, r=result) + + +def plan(rng, scale_rows=None): + """Returns (queue, behavior, work_ms, max_attempts, allowed_outcomes, kind_class).""" + roll = rng.random() * 100 + if roll < 40: + return QUEUE_WORK, "ok", rng.randint(20, 400), 5, COMPLETED + if roll < 48: + return QUEUE_WORK, "progress", rng.randint(800, 2500), 5, COMPLETED + if roll < 56: + return QUEUE_WORK, "retry", 50, 5, COMPLETED | FAILED + if roll < 60: + return QUEUE_WORK, "doomed", 50, 3, FAILED + if roll < 66: + return QUEUE_WORK, "retry_after", 50, 5, COMPLETED | FAILED + if roll < 72: + return QUEUE_WORK, "snooze", 50, 5, COMPLETED + if roll < 75: + return QUEUE_WORK, "cancel", 50, 3, CANCELLED + if roll < 80: + return QUEUE_DEADLINE, "long", rng.randint(8000, 12000), 3, COMPLETED | FAILED + if roll < 83: + return QUEUE_DEADLINE, "ignore_deadline", rng.randint(9000, 14000), 2, COMPLETED | FAILED + return QUEUE_SHARED, rng.choice(["ok", "ok", "retry", "progress"]), rng.randint(20, 300), 5, COMPLETED | FAILED + + +def kind_for(queue): + return PySharedJob if queue == QUEUE_SHARED else PyJob + + +def make_args(queue, rng, behavior, work_ms): + seq = rng.randint(1, 1 << 40) + if queue == QUEUE_SHARED: + return PySharedJob(seq=seq, behavior=behavior, work_ms=work_ms) + return PyJob(seq=seq, behavior=behavior, work_ms=work_ms, born=time.time()) + + +def kind_name(queue): + return "py_shared_job" if queue == QUEUE_SHARED else "py_job" + + +def record_enqueued(job, queue, behavior, allowed, max_attempts): + global enqueued + enqueued += 1 + log(e="enq", id=job.id, q=queue, b=behavior, x=allowed, ma=max_attempts) + + +async def producer_singles(client, rng, stop): + """Single inserts, one or two at a time.""" + period = 0.25 / max(SCALE, 0.1) + while not stop.is_set(): + for _ in range(rng.randint(1, 2)): + queue, behavior, work_ms, max_attempts, allowed = plan(rng) + args = make_args(queue, rng, behavior, work_ms) + try: + job = await client.insert(args, kind=kind_name(queue), queue=queue, max_attempts=max_attempts) + record_enqueued(job, queue, behavior, allowed, max_attempts) + except Exception as error: # connection killed by chaos, pool timeout + note_error("insert", error) + await asyncio.sleep(period) + + +async def producer_batches(client, rng, stop): + """insert_many above the old 5,461-row bind-parameter limit, and COPY.""" + rows = int(6500 * SCALE) if SCALE >= 1 else 6500 + interval = max(20.0, 75.0 * PACE) + await asyncio.sleep(min(15.0, interval)) + flip = 0 + while not stop.is_set(): + flip += 1 + args = [PyJob(seq=i, behavior="ok", work_ms=0, born=time.time()) for i in range(rows)] + try: + started_at = time.monotonic() + if flip % 3 == 0: + jobs = await client.insert_many_copy(args, kind="py_job", queue=QUEUE_WORK, max_attempts=3) + how = "copy" + else: + # One transaction, one insert_many call far above the old bind-parameter limit. + tx = await client.transaction() + try: + jobs = await tx.insert_many(args, kind="py_job", queue=QUEUE_WORK, max_attempts=3) + await tx.commit() + except Exception: + await tx.rollback() + raise + how = "insert_many" + for job in jobs: + record_enqueued(job, QUEUE_WORK, "ok", COMPLETED, 3) + print(f"BATCH how={how} rows={len(jobs)} seconds={time.monotonic() - started_at:.2f}", flush=True) + except Exception as error: + note_error(f"batch of {rows}", error) + await wait_or_stop(stop, interval) + + +async def producer_transactions(client, rng, stop): + """Enqueue inside an application transaction that sometimes rolls back.""" + while not stop.is_set(): + count = rng.randint(1, 20) + plans = [plan(rng) for _ in range(count)] + rollback = rng.random() < 0.35 + try: + tx = await client.transaction() + try: + created = [] + for queue, behavior, work_ms, max_attempts, allowed in plans: + args = make_args(queue, rng, behavior, work_ms) + job = await tx.insert(args, kind=kind_name(queue), queue=queue, max_attempts=max_attempts) + created.append((job, queue, behavior, allowed, max_attempts)) + ids = [item[0].id for item in created] + if rollback: + await tx.rollback() + log(e="rolled_back", ids=ids) + else: + await tx.commit() + for job, queue, behavior, allowed, max_attempts in created: + record_enqueued(job, queue, behavior, allowed, max_attempts) + except Exception: + try: + await tx.rollback() + except Exception: + pass + raise + except Exception as error: + note_error("transaction", error) + await asyncio.sleep(1.5 / max(SCALE, 0.1)) + + +def note_error(what, error): + global errors + errors += 1 + log(e="error", what=what, error=str(error)[:200]) + if errors <= 5: + print(f"ERROR {what}: {str(error)[:200]}", flush=True) + + +async def wait_or_stop(stop, seconds): + try: + await asyncio.wait_for(stop.wait(), timeout=seconds) + except asyncio.TimeoutError: + pass + + +async def reporter(stop): + while True: + await asyncio.sleep(2) + print(f"STAT active={active} max_active={max_active} started={started} enqueued={enqueued} errors={errors}", flush=True) + + +async def migrate_and_probe(client, label): + """What an application does at startup: migrate, then query.""" + await client.migrate() + stats = await client.queue_stats() + jobs = await client.list_jobs(limit=5) + print(f"MIGRATED {label} queues={len(stats)} sample={len(jobs)}", flush=True) + + +async def migrate_probe_main(): + client = awa.AsyncClient(URL, max_connections=4) + await migrate_and_probe(client, "probe") + # Idle with the pool open: connections stranded by migrate would show up as + # sessions beyond the pool size, or as idle-in-transaction. + await asyncio.sleep(4) + print("IDLE", flush=True) + await asyncio.sleep(4) + await client.close() + print("CLOSED", flush=True) + + +async def exit_with_parent(): + """The harness owns this process; if it dies without cleaning up, so do we.""" + parent = os.getppid() + while True: + await asyncio.sleep(3) + if os.getppid() != parent: + os._exit(1) + + +async def main(): + asyncio.create_task(exit_with_parent()) + if ROLE == "migrate_probe": + await migrate_probe_main() + return + # A chaos restart can still be migrating when the load ends. Remember the + # drain request before startup awaits, then start workers without producers. + stop = asyncio.Event() + if os.environ.get("SOAK_DRAIN") == "1": + stop.set() + loop = asyncio.get_running_loop() + loop.add_signal_handler(signal.SIGUSR1, stop.set) + rng = random.Random(SEED) + client = awa.AsyncClient(URL, max_connections=14) + await migrate_and_probe(client, "startup") + client.task(PyJob, kind="py_job", queue=QUEUE_WORK)(handle) + client.task(PySharedJob, kind="py_shared_job", queue=QUEUE_SHARED)(handle) + await client.start( + [ + {"name": QUEUE_WORK, "max_workers": WORK_WORKERS}, + {"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, + promote_interval_ms=250, + ) + # The deadline queue serves a second kind through the same handler. + print(f"READY pid={os.getpid()}", flush=True) + + tasks = [ + asyncio.create_task(producer_singles(client, rng, stop)), + asyncio.create_task(producer_batches(client, rng, stop)), + asyncio.create_task(producer_transactions(client, rng, stop)), + asyncio.create_task(reporter(stop)), + ] + await asyncio.gather(*tasks, return_exceptions=True) + await asyncio.Event().wait() + + +def exit_now(signum, _frame): + os._exit(128 + signum) + + +if __name__ == "__main__": + # PyO3/tokio teardown can crash while Python unloads modules; exit without finalizers. + signal.signal(signal.SIGTERM, exit_now) + signal.signal(signal.SIGINT, exit_now) + try: + asyncio.run(main()) + finally: + sys.stdout.flush() + os._exit(0) diff --git a/awa/tests/persona_soak/python/test_soak_python.py b/awa/tests/persona_soak/python/test_soak_python.py new file mode 100644 index 000000000..337301719 --- /dev/null +++ b/awa/tests/persona_soak/python/test_soak_python.py @@ -0,0 +1,119 @@ +"""Process-boundary regressions; uses a delayed fake client, never a database.""" + +import os +from pathlib import Path +import select +import signal +import subprocess +import sys +import tempfile +import time +import unittest + + +SCRIPT = Path(os.environ.get("SOAK_PYTHON_SCRIPT", Path(__file__).with_name("soak_python.py"))) +FAKE_AWA = ''' +import asyncio +import os +from pathlib import Path +from types import SimpleNamespace + +class AsyncClient: + def __init__(self, *args, **kwargs): + pass + + async def migrate(self): + print("STARTUP_BLOCKED", flush=True) + while not Path(os.environ["SOAK_TEST_GATE"]).exists(): + await asyncio.sleep(0.01) + + async def queue_stats(self): + return [] + + async def list_jobs(self, **kwargs): + return [] + + def task(self, *args, **kwargs): + return lambda handler: handler + + async def start(self, *args, **kwargs): + Path(os.environ["SOAK_TEST_WORKERS"]).touch() + + async def insert(self, *args, **kwargs): + Path(os.environ["SOAK_TEST_PRODUCED"]).touch() + return SimpleNamespace(id=1) + + async def transaction(self): + Path(os.environ["SOAK_TEST_PRODUCED"]).touch() + raise RuntimeError("unexpected production during drain") +''' + + +class StartupDrainTests(unittest.TestCase): + def setUp(self): + self.directory = tempfile.TemporaryDirectory(prefix="awa-python-drain-") + self.root = Path(self.directory.name) + (self.root / "awa.py").write_text(FAKE_AWA) + self.gate = self.root / "gate" + self.workers = self.root / "workers" + self.produced = self.root / "produced" + self.child = None + + def tearDown(self): + if self.child is not None: + if self.child.poll() is None: + self.child.terminate() + self.child.wait(timeout=10) + self.child.stdout.close() + self.child.stderr.close() + self.directory.cleanup() + + def start(self, *, draining=False): + env = os.environ | { + "PYTHONPATH": str(self.root), + "SOAK_DATABASE_URL": "unused-by-the-fake-client", + "SOAK_ROLE": "full", + "SOAK_DRAIN": "1" if draining else "0", + "SOAK_TEST_GATE": str(self.gate), + "SOAK_TEST_WORKERS": str(self.workers), + "SOAK_TEST_PRODUCED": str(self.produced), + } + self.child = subprocess.Popen( + [sys.executable, str(SCRIPT)], env=env, + stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1, + ) + self.wait_for("STARTUP_BLOCKED") + + def wait_for(self, prefix): + deadline = time.monotonic() + 10 + while time.monotonic() < deadline: + if select.select([self.child.stdout], [], [], 0.1)[0]: + line = self.child.stdout.readline() + if line.startswith(prefix): + return + if not line: + break + self.fail(f"did not see {prefix!r}; child status={self.child.poll()}") + + def assert_draining_workers(self): + self.gate.touch() + self.wait_for("READY pid=") + time.sleep(0.1) + self.assertIsNone(self.child.poll(), "workers must stay alive to drain") + self.assertTrue(self.workers.exists(), "workers must start after migration") + self.assertFalse(self.produced.exists(), "a drain restart must not enqueue new work") + + def test_stop_during_migration_is_remembered_until_workers_start(self): + self.start() + self.child.send_signal(signal.SIGUSR1) + time.sleep(0.05) + self.assertIsNone(self.child.poll(), "SIGUSR1 during startup must not terminate the worker") + self.assert_draining_workers() + + def test_restart_during_drain_starts_workers_without_producers(self): + self.start(draining=True) + self.assert_draining_workers() + + +if __name__ == "__main__": + unittest.main() diff --git a/awa/tests/persona_soak/python_probe.rs b/awa/tests/persona_soak/python_probe.rs new file mode 100644 index 000000000..c91d74247 --- /dev/null +++ b/awa/tests/persona_soak/python_probe.rs @@ -0,0 +1,451 @@ +//! A migrate probe may hold a normal transaction until its ready marker. + +use sqlx::PgPool; +use std::future::Future; +use std::sync::atomic::{AtomicU8, Ordering}; +use std::time::Duration; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +#[repr(u8)] +pub enum Phase { + #[default] + Migrating, + Ready, + Idle, + Closed, +} + +impl Phase { + pub fn marker(line: &str) -> Option { + if line.starts_with("MIGRATED probe ") { + Some(Self::Ready) + } else { + match line { + "IDLE" => Some(Self::Idle), + "CLOSED" => Some(Self::Closed), + _ => None, + } + } + } + + pub fn from_byte(value: u8) -> Self { + match value { + 1 => Self::Ready, + 2 => Self::Idle, + 3 => Self::Closed, + _ => Self::Migrating, + } + } +} + +/// The unique application name isolates this process, including backends older +/// than a stalled probe's deadline. A failed query is never a zero-session fact. +pub async fn sessions(pool: &PgPool, application_name: &str) -> Result<(i64, i64), String> { + tokio::time::timeout( + crate::ci_timing::scaled_timeout(Duration::from_secs(2)), + sqlx::query_as::<_, (i64, i64)>( + "SELECT count(*)::bigint, + count(*) FILTER (WHERE state = 'idle in transaction')::bigint + FROM pg_stat_activity WHERE application_name = $1 AND pid <> pg_backend_pid()", + ) + .bind(application_name) + .fetch_one(pool), + ) + .await + .map_err(|error| format!("session query timed out: {error}"))? + .map_err(|error| format!("session query failed: {error}")) +} + +#[derive(Debug, Default)] +pub struct Evidence { + pub peak_sessions: i64, + pub migrating_idle_transactions: i64, + pub ready_idle_transactions: i64, + pub ready_samples: usize, + pub remaining: Option, + pub last_exit_sample_valid: bool, + pub exit_phase: Option, + pub exit_success: bool, + pub incomplete: Option, +} + +impl Evidence { + pub fn observe(&mut self, phase_before_query: Phase, result: Result<(i64, i64), String>) { + match result { + Ok((sessions, idle)) => { + self.peak_sessions = self.peak_sessions.max(sessions); + if phase_before_query == Phase::Migrating { + self.migrating_idle_transactions = self.migrating_idle_transactions.max(idle); + } else { + if phase_before_query != Phase::Closed { + self.ready_samples += 1; + } + self.ready_idle_transactions = self.ready_idle_transactions.max(idle); + } + } + Err(error) => self.incomplete(error), + } + } + + /// Process exit alone does not prove that an active server query was cancelled. + /// CLOSED is emitted only after the client acknowledged its graceful close. + pub fn record_exit(&mut self, phase: Phase, success: bool) { + self.exit_phase = Some(phase); + self.exit_success = success; + if !self.closed_exit_proven() { + self.incomplete(format!( + "Python migrate probe exited without successful CLOSED evidence (phase={phase:?}, success={success})" + )); + } + } + + fn closed_exit_proven(&self) -> bool { + self.exit_success && self.exit_phase == Some(Phase::Closed) + } + + fn post_exit_observation_complete(&self, phase: Phase, success: bool) -> bool { + self.last_exit_sample_valid + && self.remaining == Some(0) + && (!success || phase == Phase::Closed) + } + + pub fn after_exit(&mut self, result: Result<(i64, i64), String>) { + match result { + Ok((sessions, _)) => { + self.remaining = Some(sessions); + self.last_exit_sample_valid = true; + } + Err(error) => { + self.last_exit_sample_valid = false; + self.incomplete(error); + } + } + } + + /// Reaping may precede stdout consumption. Marker and session observations + /// share the existing post-exit budget, including failed/slow queries. + pub async fn finish_exit(&mut self, phase: &AtomicU8, success: bool, mut sessions: F) + where + F: FnMut() -> Fut, + Fut: Future>, + { + let deadline = + tokio::time::Instant::now() + crate::ci_timing::scaled_timeout(Duration::from_secs(2)); + loop { + if tokio::time::Instant::now() >= deadline { + break; + } + let result = tokio::time::timeout_at(deadline, sessions()) + .await + .unwrap_or_else(|_| { + Err("post-exit session observation exceeded its existing deadline".into()) + }); + let failed = result.is_err(); + self.after_exit(result); + let observed_phase = Phase::from_byte(phase.load(Ordering::Acquire)); + if self.post_exit_observation_complete(observed_phase, success) || failed { + break; + } + tokio::time::sleep_until( + (tokio::time::Instant::now() + Duration::from_millis(100)).min(deadline), + ) + .await; + } + self.record_exit(Phase::from_byte(phase.load(Ordering::Acquire)), success); + } + + pub fn incomplete(&mut self, error: impl Into) { + self.incomplete.get_or_insert_with(|| error.into()); + } + + pub fn violation(&self) -> bool { + self.peak_sessions > 4 + || self.ready_idle_transactions > 0 + || (self.closed_exit_proven() + && self.last_exit_sample_valid + && self.remaining.is_some_and(|sessions| sessions > 0)) + } + + pub fn complete(&self) -> bool { + self.incomplete.is_none() + && self.ready_samples > 0 + && self.closed_exit_proven() + && self.last_exit_sample_valid + && self.remaining.is_some() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use sqlx::postgres::PgPoolOptions; + + async fn pool(application_name: &str, size: u32) -> PgPool { + let url = std::env::var("DATABASE_URL") + .unwrap_or_else(|_| "postgres://postgres:test@localhost:15432/awa_test".into()); + PgPoolOptions::new() + .max_connections(size) + .connect(&format!( + "{url}{}application_name={application_name}", + if url.contains('?') { '&' } else { '?' } + )) + .await + .expect("probe test pool") + } + + fn name() -> String { + format!("probe_test_{}", uuid::Uuid::new_v4().simple()) + } + + #[tokio::test] + async fn healthy_migration_transaction_is_not_a_stranded_ready_connection() { + let application = name(); + let observer = pool(&name(), 1).await; + let client = pool(&application, 1).await; + let mut tx = client.begin().await.expect("migration transaction"); + sqlx::query("SELECT 1").execute(tx.as_mut()).await.unwrap(); + let mut evidence = Evidence::default(); + let migrating = sessions(&observer, &application).await.unwrap(); + assert_eq!( + migrating, + (1, 1), + "a legitimate transaction between statements" + ); + evidence.observe(Phase::Migrating, Ok(migrating)); + tx.commit().await.unwrap(); + evidence.observe(Phase::Ready, sessions(&observer, &application).await); + client.close().await; + evidence.record_exit(Phase::Closed, true); + // An acknowledged close can precede the server removing its backend. + let deadline = + std::time::Instant::now() + crate::ci_timing::scaled_timeout(Duration::from_secs(2)); + loop { + let count = sessions(&observer, &application).await.unwrap(); + evidence.after_exit(Ok(count)); + if count.0 == 0 || std::time::Instant::now() >= deadline { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + assert!(evidence.complete(), "{evidence:?}"); + assert!(!evidence.violation(), "{evidence:?}"); + observer.close().await; + } + + #[tokio::test] + async fn ready_transaction_left_open_is_a_violation() { + let application = name(); + let observer = pool(&name(), 1).await; + let client = pool(&application, 1).await; + let mut tx = client.begin().await.unwrap(); + sqlx::query("SELECT 1").execute(tx.as_mut()).await.unwrap(); + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, sessions(&observer, &application).await); + assert!(evidence.violation(), "{evidence:?}"); + tx.rollback().await.unwrap(); + client.close().await; + observer.close().await; + } + + #[tokio::test] + async fn retained_session_after_exit_is_a_violation_and_other_probes_are_excluded() { + let application = name(); + let observer = pool(&name(), 1).await; + let leaked = pool(&application, 1).await; + let unrelated = pool(&name(), 1).await; + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, sessions(&observer, &application).await); + evidence.record_exit(Phase::Closed, true); + evidence.after_exit(sessions(&observer, &application).await); + assert_eq!(evidence.remaining, Some(1)); + assert!(evidence.complete()); + assert!(evidence.violation(), "{evidence:?}"); + leaked.close().await; + unrelated.close().await; + observer.close().await; + } + + #[tokio::test] + async fn session_query_failure_is_inconclusive_and_not_zero() { + let observer = pool(&name(), 1).await; + observer.close().await; + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, sessions(&observer, &name()).await); + evidence.after_exit(sessions(&observer, &name()).await); + assert!(!evidence.complete()); + assert!(evidence.incomplete.is_some()); + assert_eq!(evidence.remaining, None); + } + + #[test] + fn probe_exit_forced_during_a_read_retains_incomplete_session_evidence() { + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + evidence.record_exit(Phase::Ready, false); + evidence.after_exit(Ok((1, 0))); + assert!( + !evidence.violation(), + "a forced exit does not prove graceful-close failure" + ); + assert!(!evidence.complete()); + assert!(evidence.incomplete.is_some()); + assert_eq!(evidence.remaining, Some(1)); + } + + #[test] + fn probe_exit_delayed_closed_marker_uses_the_existing_session_window() { + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + evidence.after_exit(Ok((0, 0))); + assert!( + !evidence.post_exit_observation_complete(Phase::Ready, true), + "zero sessions must not stop the reader before its flushed CLOSED marker" + ); + // The detached reader consumes the marker during the same post-exit window. + assert!(evidence.post_exit_observation_complete(Phase::Closed, true)); + evidence.record_exit(Phase::Closed, true); + assert!(evidence.complete()); + assert!(!evidence.violation()); + } + + #[tokio::test] + async fn probe_exit_finish_waits_for_delayed_stdout_closed_marker() { + let phase = std::sync::Arc::new(AtomicU8::new(Phase::Ready as u8)); + let reader_phase = phase.clone(); + let reader = tokio::spawn(async move { + tokio::task::yield_now().await; + tokio::time::sleep(Duration::from_millis(10)).await; + reader_phase.store(Phase::Closed as u8, Ordering::Release); + }); + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + let mut calls = 0; + evidence + .finish_exit(&phase, true, || { + calls += 1; + async { Ok((0, 0)) } + }) + .await; + reader.await.unwrap(); + assert!(evidence.complete(), "{evidence:?}"); + assert!(!evidence.violation()); + assert_eq!(evidence.exit_phase, Some(Phase::Closed)); + assert!( + calls >= 2, + "the helper must continue past zero sessions before CLOSED" + ); + } + + #[tokio::test] + async fn probe_exit_finish_query_failure_does_not_prove_a_retained_session_leak() { + let phase = AtomicU8::new(Phase::Closed as u8); + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + let mut calls = 0; + evidence + .finish_exit(&phase, true, || { + calls += 1; + let result = if calls == 1 { + Ok((1, 0)) + } else { + Err("injected final query failure".into()) + }; + async move { result } + }) + .await; + assert!( + !evidence.violation(), + "an earlier positive count is not a successful final observation" + ); + assert!(!evidence.complete()); + assert!(evidence.incomplete.is_some()); + assert_eq!( + evidence.remaining, + Some(1), + "retain the last successful count for diagnosis" + ); + assert!(!evidence.last_exit_sample_valid); + } + + #[test] + fn probe_exit_success_without_closed_marker_is_incomplete() { + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + evidence.record_exit(Phase::Ready, true); + evidence.after_exit(Ok((1, 0))); + assert!( + !evidence.violation(), + "an exit code is not an acknowledged pool close" + ); + assert!(!evidence.complete()); + assert!(evidence.incomplete.is_some()); + assert_eq!(evidence.remaining, Some(1)); + } + + #[test] + fn probe_exit_acknowledged_close_with_retained_session_is_a_violation() { + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + evidence.record_exit(Phase::Closed, true); + evidence.after_exit(Ok((1, 0))); + assert!(evidence.violation()); + assert!( + evidence.complete(), + "complete coverage still reports the failure" + ); + assert_eq!(evidence.remaining, Some(1)); + } + + #[test] + fn probe_exit_acknowledged_close_and_no_session_is_complete() { + let mut evidence = Evidence::default(); + evidence.observe(Phase::Ready, Ok((2, 0))); + evidence.record_exit(Phase::Closed, true); + evidence.after_exit(Ok((0, 0))); + assert!(!evidence.violation()); + assert!(evidence.complete()); + } + + #[test] + fn probe_exit_unproven_close_does_not_hide_pool_or_idle_transaction_violations() { + for (phase, sample) in [(Phase::Migrating, (5, 0)), (Phase::Ready, (2, 1))] { + let mut evidence = Evidence::default(); + evidence.observe(phase, Ok(sample)); + evidence.record_exit(phase, false); + evidence.after_exit(Ok((1, 0))); + assert!(evidence.violation(), "{evidence:?}"); + assert!(!evidence.complete()); + assert_eq!(evidence.remaining, Some(1)); + } + } + + #[test] + fn ready_marker_and_missing_ready_coverage_are_explicit() { + assert_eq!( + Phase::marker("MIGRATED probe queues=1 sample=0"), + Some(Phase::Ready) + ); + assert_eq!(Phase::marker("IDLE"), Some(Phase::Idle)); + assert_eq!(Phase::marker("CLOSED"), Some(Phase::Closed)); + assert_eq!(Phase::marker("MIGRATED startup queues=1 sample=0"), None); + let mut evidence = Evidence::default(); + evidence.observe(Phase::Migrating, Ok((4, 1))); + evidence.after_exit(Ok((0, 0))); + assert!( + !evidence.complete(), + "an exit without ready samples proves no idle invariant" + ); + evidence.observe(Phase::Migrating, Ok((5, 0))); + assert!( + evidence.violation(), + "the pool budget applies during migration too" + ); + // Even valid ready/idle/closed markers may all reach the reader before + // the next snapshot. An already closed pool is not live idle coverage. + let mut closed_only = Evidence::default(); + closed_only.observe(Phase::Closed, Ok((0, 0))); + closed_only.record_exit(Phase::Closed, true); + closed_only.after_exit(Ok((0, 0))); + assert!(!closed_only.complete()); + } +} diff --git a/awa/tests/persona_soak/registry.rs b/awa/tests/persona_soak/registry.rs new file mode 100644 index 000000000..a0cc42835 --- /dev/null +++ b/awa/tests/persona_soak/registry.rs @@ -0,0 +1,131 @@ +//! The set of queues the run owns and the live facts the oracle needs about each: +//! which persona serves it, how many workers can run it, how many are busy, and +//! whether it is intentionally not being served right now. + +use super::config::Persona; +use std::collections::HashMap; +use std::sync::atomic::{AtomicI32, AtomicI64, Ordering}; +use std::sync::{Arc, RwLock}; + +impl QueueInfo { + pub fn hold(&self) { + self.pause_holds.fetch_add(1, Ordering::Relaxed); + } + + pub fn release(&self) { + self.pause_holds.fetch_sub(1, Ordering::Relaxed); + } + + pub fn is_held(&self) -> bool { + self.pause_holds.load(Ordering::Relaxed) > 0 + } + + /// Jobs enqueued and not yet started, as far as the producers and workers know. + pub fn backlog_estimate(&self) -> i64 { + self.enqueued.load(Ordering::Relaxed) - self.started.load(Ordering::Relaxed) + } +} + +pub struct QueueInfo { + pub name: String, + pub persona: Persona, + pub index: u32, + /// Workers currently able to run this queue across all live runtimes. + pub capacity: AtomicI64, + pub in_flight: AtomicI64, + /// Jobs producers enqueued here and executions begun, for rough backlog estimates. + pub enqueued: AtomicI64, + pub started: AtomicI64, + /// An operator (or a rate limit) is deliberately holding jobs back. + /// Operators currently holding the queue paused; the stall check skips it while non-zero. + pub pause_holds: AtomicI32, + pub rate_limited: bool, + /// Milliseconds since run start of the most recent operator drain, 0 if never. + pub drained_at_ms: AtomicI64, +} + +#[derive(Default)] +pub struct QueueRegistry { + queues: RwLock>>, + by_name: RwLock>, +} + +impl QueueRegistry { + pub fn register(&self, name: &str, persona: Persona, rate_limited: bool) -> Arc { + let mut queues = self.queues.write().expect("registry lock"); + let mut by_name = self.by_name.write().expect("registry lock"); + if let Some(index) = by_name.get(name) { + return queues[*index as usize].clone(); + } + let index = queues.len() as u32; + let info = Arc::new(QueueInfo { + name: name.to_string(), + persona, + index, + capacity: AtomicI64::new(0), + in_flight: AtomicI64::new(0), + enqueued: AtomicI64::new(0), + started: AtomicI64::new(0), + pause_holds: AtomicI32::new(0), + rate_limited, + drained_at_ms: AtomicI64::new(0), + }); + by_name.insert(name.to_string(), index); + queues.push(info.clone()); + info + } + + /// 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()) + } + + pub fn by_index(&self, index: u32) -> Option> { + self.queues + .read() + .expect("registry lock") + .get(index as usize) + .cloned() + } + + pub fn all(&self) -> Vec> { + self.queues.read().expect("registry lock").clone() + } + + pub fn of(&self, persona: Persona) -> Vec> { + self.all() + .into_iter() + .filter(|queue| queue.persona == persona) + .collect() + } +} + +/// Counts one running handler against its queue for as long as the guard lives, +/// including when the runtime is killed and the handler future is dropped. +pub struct InFlightGuard { + queue: Option>, +} + +impl InFlightGuard { + pub fn new(queue: Option>) -> Self { + if let Some(queue) = &queue { + queue.in_flight.fetch_add(1, Ordering::Relaxed); + } + Self { queue } + } +} + +impl Drop for InFlightGuard { + fn drop(&mut self) { + if let Some(queue) = &self.queue { + queue.in_flight.fetch_sub(1, Ordering::Relaxed); + } + } +} diff --git a/awa/tests/persona_soak/report.rs b/awa/tests/persona_soak/report.rs new file mode 100644 index 000000000..d309bfc94 --- /dev/null +++ b/awa/tests/persona_soak/report.rs @@ -0,0 +1,341 @@ +//! The run report and, for a failed run, the reproducer. + +use super::config::{Persona, SoakConfig}; +use super::pgserver::crash_excerpts; +use super::shared::Shared; +use serde_json::json; +use std::fmt::Write as _; +use std::path::Path; +use std::process::Command; + +/// Printed at the top of every report; also lets a built binary be matched to this harness. +pub const HEADER: &str = "AWA PERSONA SOAK HARNESS"; + +pub fn git_revision() -> String { + Command::new("git") + .args(["rev-parse", "--short=10", "HEAD"]) + .output() + .ok() + .map(|output| String::from_utf8_lossy(&output.stdout).trim().to_string()) + .filter(|revision| !revision.is_empty()) + .unwrap_or_else(|| "unknown".into()) +} + +pub struct Outcome { + pub quiesced: bool, + pub elapsed_seconds: f64, + pub crash_excerpts: Vec, + pub log_tail: Option, +} + +pub fn render(shared: &Shared, outcome: &Outcome) -> String { + let config = &shared.config; + let info = &shared.info; + let violations = shared.violations.lock().expect("violations"); + let mut out = String::new(); + let _ = writeln!(out, "# {HEADER}"); + let _ = writeln!(out); + let _ = writeln!( + out, + "verdict: {}", + if violations + .iter() + .any(|violation| violation.kind != "oracle_incomplete") + { + "FAIL" + } else if shared + .oracle_incomplete + .load(std::sync::atomic::Ordering::Relaxed) + { + "INCONCLUSIVE" + } else { + "PASS" + } + ); + let _ = writeln!( + out, + "engine: {} (schema {})", + info.engine, + info.queue_storage_schema.as_deref().unwrap_or("-") + ); + let _ = writeln!(out, "server: PostgreSQL {}", info.version); + let _ = writeln!( + out, + "max_connections: {} (usable by personas: {})", + info.max_connections, + shared.budget.usable() + ); + let _ = writeln!(out, "connections at start: {}", info.connections_at_start); + let _ = writeln!(out, "awa: {} @ {}", info.awa_version, git_revision()); + let _ = writeln!( + out, + "seed: {} personas: {} load: {}s scale: {} elapsed: {:.0}s quiesced: {}", + config.seed, + config + .personas + .iter() + .map(|persona| persona.name()) + .collect::>() + .join(","), + config.load.as_secs(), + config.scale, + outcome.elapsed_seconds, + outcome.quiesced + ); + let _ = writeln!(out, "ledger jobs: {}", shared.ledger.len()); + let _ = writeln!(out, "NO_EVENTS durable loss detection: job_id % {} = 0; other returned NO_EVENTS jobs are excluded from terminal-commit proof.", super::ledger::TERMINAL_PROOF_STRIDE); + let _ = writeln!(out); + + let _ = writeln!(out, "## Connection budget"); + for (label, count) in shared.budget.table() { + let _ = writeln!(out, "- {label}: {count}"); + } + let _ = writeln!(out); + + for persona in Persona::ALL { + if !config.has(persona) { + continue; + } + let stats = shared.stats(persona); + let counters = stats.counters(); + let _ = writeln!(out, "## {}", persona.name()); + for (name, value) in counters { + let _ = writeln!(out, "- {name}: {value}"); + } + for (name, summary) in stats.summaries() { + let _ = writeln!( + out, + "- {name}: n={} p50={} p99={} max={}", + summary.count, summary.p50, summary.p99, summary.max + ); + } + if persona == Persona::ClaimAbandon { + let _ = writeln!(out, "{}", shared.abandon.render()); + } + let _ = writeln!(out); + } + + let _ = writeln!(out, "## Run"); + if shared.run.get("claim_ring_lock_capture_records") > 0 { + let _ = writeln!( + out, + "Lock evidence: [claim-ring-locks.jsonl](claim-ring-locks.jsonl)" + ); + } + for (name, value) in shared.run.counters() { + let _ = writeln!(out, "- {name}: {value}"); + } + for (name, summary) in shared.run.summaries() { + let _ = writeln!( + out, + "- {name}: n={} p50={} p99={} max={}", + summary.count, summary.p50, summary.p99, summary.max + ); + } + let _ = writeln!(out); + + let findings = shared.findings.lock().expect("findings"); + let _ = writeln!(out, "## Findings, not fatal ({})", findings.len()); + for finding in findings.iter() { + let _ = writeln!( + out, + "- [{:.1}s] {}: {}", + finding.at_ms as f64 / 1000.0, + finding.kind, + finding.summary + ); + for row in finding.rows.iter().take(6) { + let _ = writeln!(out, " {row}"); + } + } + let _ = writeln!(out); + + let _ = writeln!(out, "## Violations ({})", violations.len()); + for violation in violations.iter() { + let _ = writeln!( + out, + "- [{:.1}s{}] {}: {}", + violation.at_ms as f64 / 1000.0, + if violation.during_chaos { + ", during chaos" + } else { + "" + }, + violation.kind, + violation.summary + ); + for row in violation.rows.iter().take(10) { + let _ = writeln!(out, " {row}"); + } + } + if !outcome.crash_excerpts.is_empty() { + let _ = writeln!(out, "\n## Server log: crash lines"); + for excerpt in &outcome.crash_excerpts { + let _ = writeln!(out, "```\n{excerpt}\n```"); + } + } + out +} + +/// Writes `report.md` and `report.json`, and for a failed run `reproducer.md`. +pub fn write(shared: &Shared, outcome: &Outcome) -> std::io::Result<()> { + let dir = &shared.config.report_dir; + std::fs::create_dir_all(dir)?; + let markdown = render(shared, outcome); + println!("\n{markdown}"); + std::fs::write(dir.join("report.md"), &markdown)?; + + let violations = shared.violations.lock().expect("violations").clone(); + let timeline = shared.timeline.lock().expect("timeline").clone(); + let per_persona: serde_json::Map = Persona::ALL + .into_iter() + .filter(|persona| shared.config.has(*persona)) + .map(|persona| { + ( + persona.name().to_string(), + json!({ "counters": shared.stats(persona).counters(), "samples": shared.stats(persona).summaries() }), + ) + }) + .collect(); + let summary = json!({ + "verdict": if violations.iter().any(|violation| violation.kind != "oracle_incomplete") { "fail" } + else if shared.oracle_incomplete.load(std::sync::atomic::Ordering::Relaxed) { "inconclusive" } else { "pass" }, + "terminal_proof_sample_stride": super::ledger::TERMINAL_PROOF_STRIDE, + "oracle_complete": !shared.oracle_incomplete.load(std::sync::atomic::Ordering::Relaxed), + "seed": shared.config.seed, + "personas": shared.config.personas.iter().map(|p| p.name()).collect::>(), + "engine": shared.info.engine, + "server_version": shared.info.version, + "max_connections": shared.info.max_connections, + "awa_version": shared.info.awa_version, + "git": git_revision(), + "elapsed_seconds": outcome.elapsed_seconds, + "quiesced": outcome.quiesced, + "personas_stats": per_persona, + "abandon": shared.abandon.to_json(), + "run": { "counters": shared.run.counters(), "samples": shared.run.summaries() }, + "violations": violations, + "findings": shared.findings.lock().expect("findings").clone(), + "timeline": timeline, + }); + std::fs::write( + dir.join("report.json"), + serde_json::to_string_pretty(&summary).unwrap_or_default(), + )?; + + if !violations.is_empty() { + write_reproducer(dir, &shared.config, shared, outcome)?; + } + Ok(()) +} + +fn write_reproducer( + dir: &Path, + config: &SoakConfig, + shared: &Shared, + outcome: &Outcome, +) -> std::io::Result<()> { + let violations = shared.violations.lock().expect("violations").clone(); + let timeline = shared.timeline.lock().expect("timeline").clone(); + let mut out = String::new(); + let _ = writeln!(out, "# Reproducer\n"); + let _ = writeln!(out, "Re-run with the same seed, personas and duration:\n"); + let _ = writeln!( + out, + "```\nAWA_SOAK_DATABASE_URL= AWA_SOAK_SEED={} AWA_SOAK_PERSONAS={} AWA_SOAK_SECONDS={} AWA_SOAK_SCALE={} \\\n cargo test --release -p awa --test persona_soak_test -- --ignored --nocapture\n```\n", + config.seed, + config.personas.iter().map(|persona| persona.name()).collect::>().join(","), + config.load.as_secs(), + config.scale + ); + let _ = writeln!(out, "Seed fixes the workload mix and the chaos schedule; timing of races still varies between runs.\n"); + let _ = writeln!( + out, + "Engine: {}; PostgreSQL {}; awa {} @ {}\n", + shared.info.engine, + shared.info.version, + shared.info.awa_version, + git_revision() + ); + let _ = writeln!(out, "## Violations\n"); + for violation in &violations { + let persona = violation.kind.as_str(); + let _ = writeln!( + out, + "### {persona} at {:.1}s{}\n{}\n", + violation.at_ms as f64 / 1000.0, + if violation.during_chaos { + " (during chaos)" + } else { + "" + }, + violation.summary + ); + for row in &violation.rows { + let _ = writeln!(out, " {row}"); + } + let nearby: Vec<_> = timeline + .iter() + .filter(|event| (event.at_ms - violation.at_ms).abs() <= 60_000) + .collect(); + if !nearby.is_empty() { + let _ = writeln!(out, "\nTimeline within 60s:"); + for event in nearby { + let _ = writeln!( + out, + "- {:.1}s [{}] {}", + event.at_ms as f64 / 1000.0, + event.source, + event.text + ); + } + } + let _ = writeln!(out); + } + let _ = writeln!(out, "## Full timeline\n"); + for event in &timeline { + let _ = writeln!( + out, + "- {:.1}s [{}] {}", + event.at_ms as f64 / 1000.0, + event.source, + event.text + ); + } + let _ = writeln!(out, "\n## Server log excerpt\n"); + if !outcome.crash_excerpts.is_empty() { + for excerpt in &outcome.crash_excerpts { + let _ = writeln!(out, "```\n{excerpt}\n```"); + } + } else if let Some(tail) = &outcome.log_tail { + let interesting: Vec<&str> = tail + .lines() + .filter(|line| { + line.contains("ERROR") + || line.contains("FATAL") + || line.contains("PANIC") + || line.contains("WARNING") + }) + .collect(); + let start = interesting.len().saturating_sub(80); + let _ = writeln!( + out, + "No crash lines. Last warnings and errors ({} total):\n```", + interesting.len() + ); + for line in &interesting[start..] { + let _ = writeln!(out, "{line}"); + } + let _ = writeln!(out, "```"); + } else { + let _ = writeln!( + out, + "No server log source configured (set AWA_SOAK_PG_LOG or AWA_SOAK_PG_CONTAINER)." + ); + } + std::fs::write(dir.join("reproducer.md"), out) +} + +pub fn excerpts_of(log: &str) -> Vec { + crash_excerpts(log) +} diff --git a/awa/tests/persona_soak/ring_lock_capture.sql b/awa/tests/persona_soak/ring_lock_capture.sql new file mode 100644 index 000000000..165990fe1 --- /dev/null +++ b/awa/tests/persona_soak/ring_lock_capture.sql @@ -0,0 +1,63 @@ +WITH current_db AS ( + SELECT oid FROM pg_catalog.pg_database WHERE datname = current_database() +), +relations AS MATERIALIZED ( + SELECT c.oid, c.relname::text AS relation + FROM pg_catalog.pg_class c + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = $1 AND c.relname::text = ANY($2::text[]) +), +activity AS MATERIALIZED ( + SELECT pid, leader_pid, backend_start, application_name, usename, state, + wait_event_type, wait_event, xact_start, query_start, + EXTRACT(EPOCH FROM (clock_timestamp() - xact_start))::float8 AS tx_age_s, + EXTRACT(EPOCH FROM (clock_timestamp() - query_start))::float8 AS query_age_s, + left(query, 512) AS query, + query IS NULL OR query = '' AS visibility_limited + FROM pg_catalog.pg_stat_activity + WHERE datid = (SELECT oid FROM current_db) +), +locks AS MATERIALIZED ( + SELECT l.pid, r.oid AS relation_oid, r.relation, l.mode, l.granted, l.fastpath, + l.virtualtransaction, l.waitstart + FROM pg_catalog.pg_locks l + JOIN relations r ON r.oid = l.relation + LEFT JOIN activity a ON a.pid = l.pid + WHERE l.database = (SELECT oid FROM current_db) + AND l.locktype = 'relation' + AND l.pid IS DISTINCT FROM pg_backend_pid() + ORDER BY l.granted, a.xact_start NULLS LAST, a.query_start NULLS LAST, + l.pid NULLS FIRST, r.relation, l.mode + LIMIT 257 --256 retained rows plus an explicit overflow sentinel +), +interested AS MATERIALIZED ( + SELECT DISTINCT l.pid + FROM locks l + JOIN activity a ON a.pid = l.pid + WHERE a.wait_event_type = 'Lock' + ORDER BY l.pid + LIMIT 32 +), +blockers AS MATERIALIZED ( + SELECT i.pid, pg_catalog.pg_blocking_pids(i.pid) AS pids + FROM interested i +), +lock_rows AS MATERIALIZED ( +SELECT clock_timestamp() AS observed_at, l.*, + a.leader_pid, a.backend_start, a.application_name, a.usename, a.state, + a.wait_event_type, a.wait_event, a.xact_start, a.query_start, + a.tx_age_s, a.query_age_s, a.query, a.visibility_limited, + a.pid IS NULL AS activity_missing, + b.pids[1:16] AS blocking_pids, + b.pid IS NOT NULL AS blocking_pids_observed, + COALESCE(cardinality(b.pids) > 16, false) AS blocking_pids_truncated +FROM locks l +LEFT JOIN activity a ON a.pid = l.pid +LEFT JOIN blockers b ON b.pid = l.pid +) +SELECT jsonb_build_object( + 'server_statement_at', statement_timestamp(), + 'server_observed_at', clock_timestamp(), + 'resolved_relations', COALESCE((SELECT jsonb_agg(to_jsonb(r)) FROM relations r), '[]'::jsonb), + 'locks', COALESCE((SELECT jsonb_agg(to_jsonb(l) ORDER BY l.granted, l.xact_start NULLS LAST, l.pid NULLS FIRST, l.relation, l.mode) FROM lock_rows l), '[]'::jsonb) +); diff --git a/awa/tests/persona_soak/ring_locks.rs b/awa/tests/persona_soak/ring_locks.rs new file mode 100644 index 000000000..2045ed1de --- /dev/null +++ b/awa/tests/persona_soak/ring_locks.rs @@ -0,0 +1,543 @@ +//! Supplemental lock evidence; never changes the claim-ring oracle's verdict. + +use super::abandon_oracle::RingEpisode; +use super::shared::Shared; +use serde_json::{json, Value}; +use sqlx::PgPool; +use std::collections::BTreeSet; +use std::io::Write; +use std::time::Duration; + +const SQL: &str = include_str!("ring_lock_capture.sql"); +const LOCK_ROWS: usize = 256; +const BLOCKER_PIDS: usize = 64; +const CLIENT_DEADLINE: Duration = Duration::from_millis(750); + +#[derive(Clone, Copy, Debug, PartialEq)] +pub enum Stage { + Suspected, + Violation, +} + +impl Stage { + fn index(self) -> usize { + match self { + Self::Suspected => 0, + Self::Violation => 1, + } + } + + fn name(self) -> &'static str { + match self { + Self::Suspected => "suspected", + Self::Violation => "violation", + } + } +} + +#[derive(Debug, PartialEq)] +pub enum Decision { + Capture(Stage), + Exhausted(Stage), +} + +/// Four slots per stage reserve half the run budget for actual violations. +#[derive(Default)] +pub struct CaptureBudget { + episode: Option<(i64, i64)>, + seen: [bool; 2], + used: [u8; 2], + exhausted_reported: [bool; 2], +} + +impl CaptureBudget { + pub fn plan( + &mut self, + episode: &RingEpisode, + now_ms: i64, + bound_ms: i64, + zero_live: bool, + ) -> Option { + let key = (episode.generation, episode.start_ms); + if self.episode != Some(key) { + self.episode = Some(key); + self.seen = [false; 2]; + } + let stage = if episode.reported { + Stage::Violation + } else if zero_live && now_ms - episode.start_ms >= bound_ms / 2 { + Stage::Suspected + } else { + return None; + }; + let index = stage.index(); + if self.seen[index] { + return None; + } + self.seen[index] = true; + if self.used[index] < 4 { + self.used[index] += 1; + Some(Decision::Capture(stage)) + } else if !self.exhausted_reported[index] { + self.exhausted_reported[index] = true; + Some(Decision::Exhausted(stage)) + } else { + None + } + } +} + +fn names(slot: i32) -> Vec { + [ + format!("lease_claims_{slot}"), + format!("lease_claim_batches_{slot}"), + format!("lease_claim_closures_{slot}"), + format!("lease_claim_closure_batches_{slot}"), + "claim_ring_state".into(), + "claim_ring_slots".into(), + ] + .into() +} + +fn missing(error: impl std::fmt::Display) -> Value { + json!({"status": "missing", "error": error.to_string().chars().take(512).collect::()}) +} + +fn finish(mut value: Value, expected: &[String]) -> Value { + let mut gaps = Vec::new(); + let resolved: BTreeSet<_> = value["resolved_relations"] + .as_array() + .into_iter() + .flatten() + .filter_map(|row| row["relation"].as_str()) + .collect(); + let unresolved: Vec<_> = expected + .iter() + .filter(|name| !resolved.contains(name.as_str())) + .cloned() + .collect(); + if !unresolved.is_empty() { + gaps.push("unresolved_relations"); + } + if let Some(rows) = value["locks"].as_array_mut() { + if rows.len() > LOCK_ROWS { + rows.truncate(LOCK_ROWS); + gaps.push("lock_rows_truncated"); + } + if rows.iter().any(|row| { + row["blocking_pids_truncated"] == true + || (row["wait_event_type"] == "Lock" && row["blocking_pids_observed"] != true) + || row["activity_missing"] == true + || row["visibility_limited"] == true + }) { + gaps.push("incomplete_activity_or_blocker_graph"); + } + } + if value["blocker_pid_lookup_truncated"] == true { + gaps.push("blocker_pid_lookup_truncated"); + } + if value["missing_blocker_activity"] + .as_array() + .is_some_and(|rows| !rows.is_empty()) + { + gaps.push("missing_blocker_activity"); + } + if value["blocker_activity"] + .as_array() + .is_some_and(|rows| rows.iter().any(|row| row["visibility_limited"] == true)) + { + gaps.push("blocker_activity_visibility_limited"); + } + value["status"] = json!(if gaps.is_empty() { "ok" } else { "partial" }); + value["missing"] = json!(gaps); + value["unresolved_relations"] = json!(unresolved); + value +} + +async fn collect(pool: &PgPool, schema: &str, slot: i32, sql: &str) -> Value { + let expected = names(slot); + let result = tokio::time::timeout(CLIENT_DEADLINE, async { + let mut tx = pool.begin().await?; + sqlx::query("SET TRANSACTION READ ONLY") + .execute(tx.as_mut()) + .await?; + sqlx::query("SET LOCAL statement_timeout = '250ms'") + .execute(tx.as_mut()) + .await?; + sqlx::query("SET LOCAL lock_timeout = '100ms'") + .execute(tx.as_mut()) + .await?; + let mut value: Value = sqlx::query_scalar(sql) + .bind(schema) + .bind(&expected) + .fetch_one(tx.as_mut()) + .await?; + let pids: BTreeSet = value["locks"] + .as_array() + .into_iter() + .flatten() + .flat_map(|row| row["blocking_pids"].as_array().into_iter().flatten()) + .filter_map(|pid| pid.as_i64().and_then(|pid| i32::try_from(pid).ok())) + .filter(|pid| *pid > 0) + .collect(); + value["blocker_pid_lookup_truncated"] = json!(pids.len() > BLOCKER_PIDS); + let requested: Vec<_> = pids.into_iter().take(BLOCKER_PIDS).collect(); + let activity: Value = sqlx::query_scalar( + "SELECT COALESCE(jsonb_agg(to_jsonb(a)), '[]'::jsonb) FROM ( + SELECT pid, leader_pid, backend_start, application_name, usename, state, + wait_event_type, wait_event, xact_start, query_start, + EXTRACT(EPOCH FROM(clock_timestamp()-xact_start))::float8 AS tx_age_s, + EXTRACT(EPOCH FROM(clock_timestamp()-query_start))::float8 AS query_age_s, + left(query,512) AS query, + query IS NULL OR query = '' AS visibility_limited + FROM pg_catalog.pg_stat_activity + WHERE datid=(SELECT oid FROM pg_catalog.pg_database WHERE datname=current_database()) + AND pid=ANY($1::int4[]) ORDER BY pid LIMIT 65 + ) a", + ) + .bind(&requested) + .fetch_one(tx.as_mut()) + .await?; + let found: BTreeSet<_> = activity + .as_array() + .into_iter() + .flatten() + .filter_map(|row| row["pid"].as_i64()) + .collect(); + value["missing_blocker_activity"] = json!(requested + .iter() + .filter(|pid| !found.contains(&i64::from(**pid))) + .collect::>()); + value["blocker_activity"] = activity; + tx.rollback().await?; + Ok::<_, sqlx::Error>(value) + }) + .await; + match result { + Ok(Ok(value)) => finish(value, &expected), + Ok(Err(error)) => missing(error), + Err(_) => missing("claim-ring metadata capture exceeded its 750ms client deadline"), + } +} + +pub async fn observe( + shared: &Shared, + pool: &PgPool, + schema: &str, + slot_count: i32, + next_slot: Option<(i64, i64)>, + budget: &mut CaptureBudget, +) { + let Some(episode) = shared.abandon.current_ring_episode() else { + return; + }; + let now = shared.now_ms(); + let Some(decision) = budget.plan( + &episode, + now, + shared.config.abandon_bound.as_millis() as i64, + next_slot.is_some_and(|(live, _)| live == 0), + ) else { + return; + }; + let stage = match decision { + Decision::Capture(stage) => stage, + Decision::Exhausted(stage) => { + shared.finding( + "claim_ring_lock_capture_missing", + format!("{} lock-capture budget exhausted", stage.name()), + vec![], + ); + return; + } + }; + let incoming = (episode.slot + 1).rem_euclid(slot_count); + let mut value = collect(pool, schema, incoming, SQL).await; + value["stage"] = json!(stage.name()); + value["started_at_ms"] = json!(now); + value["finished_at_ms"] = json!(shared.now_ms()); + value["current_slot"] = json!(episode.slot); + value["next_slot"] = json!(incoming); + value["generation"] = json!(episode.generation); + value["next_slot_probe"] = json!(next_slot); + shared.run.add("claim_ring_lock_capture_attempts", 1); + if value["status"] != "ok" { + shared.finding( + "claim_ring_lock_capture_missing", + "claim-ring lock evidence is missing or partial", + vec![value.to_string()], + ); + } + let path = shared.config.report_dir.join("claim-ring-locks.jsonl"); + let written = (|| { + let mut file = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(path)?; + writeln!(file, "{value}") + })(); + match written { + Ok(()) => shared.run.add("claim_ring_lock_capture_records", 1), + Err(error) => shared.finding( + "claim_ring_lock_capture_missing", + format!("could not write claim-ring lock evidence: {error}"), + vec![], + ), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn episode(generation: i64, reported: bool) -> RingEpisode { + RingEpisode { + generation, + reported, + slot: 7, + start_ms: 0, + end_ms: None, + last_live_ms: 0, + expired_in_next_slot: 0, + } + } + + #[test] + fn ring_lock_budget_waits_half_bound_and_reserves_violation_captures() { + let mut budget = CaptureBudget::default(); + assert_eq!(budget.plan(&episode(0, false), 59_999, 120_000, true), None); + assert_eq!( + budget.plan(&episode(0, false), 60_000, 120_000, false), + None + ); + for generation in 1..=4 { + assert_eq!( + budget.plan(&episode(generation, false), 60_000, 120_000, true), + Some(Decision::Capture(Stage::Suspected)) + ); + assert_eq!( + budget.plan(&episode(generation, false), 70_000, 120_000, true), + None + ); + } + assert_eq!( + budget.plan(&episode(5, false), 60_000, 120_000, true), + Some(Decision::Exhausted(Stage::Suspected)) + ); + assert_eq!(budget.plan(&episode(6, false), 60_000, 120_000, true), None); + for generation in 7..=10 { + assert_eq!( + budget.plan(&episode(generation, true), 120_001, 120_000, true), + Some(Decision::Capture(Stage::Violation)) + ); + assert_eq!( + budget.plan(&episode(generation, true), 130_000, 120_000, true), + None + ); + } + assert_eq!( + budget.plan(&episode(11, true), 120_001, 120_000, true), + Some(Decision::Exhausted(Stage::Violation)) + ); + assert_eq!( + budget.plan(&episode(12, true), 120_001, 120_000, true), + None + ); + assert_eq!(budget.used, [4, 4]); + } + + #[test] + fn ring_lock_result_caps_and_missing_relations_stay_explicit() { + let rows = vec![json!({"blocking_pids_truncated": true}); LOCK_ROWS + 1]; + let value = finish(json!({"resolved_relations": [], "locks": rows}), &names(0)); + assert_eq!(value["status"], "partial"); + assert_eq!(value["locks"].as_array().unwrap().len(), LOCK_ROWS); + assert_eq!(value["unresolved_relations"].as_array().unwrap().len(), 6); + let gaps = value["missing"].as_array().unwrap(); + assert!(gaps.contains(&json!("lock_rows_truncated"))); + assert!(gaps.contains(&json!("incomplete_activity_or_blocker_graph"))); + } + + #[test] + fn ring_lock_hidden_second_lookup_is_partial_even_without_selected_locks() { + let expected = names(0); + let inventory: Vec<_> = expected + .iter() + .map(|name| json!({"relation": name})) + .collect(); + let value = finish( + json!({ + "resolved_relations": inventory, + "locks": [], + "blocker_activity": [{"pid": 42, "query": null, "visibility_limited": true}] + }), + &expected, + ); + assert_eq!(value["status"], "partial"); + assert!(value["missing"] + .as_array() + .unwrap() + .contains(&json!("blocker_activity_visibility_limited"))); + } + + struct Database { + config: super::super::config::SoakConfig, + pool: PgPool, + } + + impl Database { + async fn create() -> Self { + use super::super::config::{replace_database_name, SoakConfig}; + let base = std::env::var("DATABASE_URL").unwrap_or_else(|_| { + "postgres://postgres:test@localhost:15432/awa_test".to_string() + }); + let name = format!("awa_soak_ring_locks_{}", uuid::Uuid::new_v4().simple()); + let config = SoakConfig::for_test(replace_database_name(&base, &name)); + super::super::db::recreate_database(&config) + .await + .expect("isolated database"); + let pool = super::super::db::connect_pool(&config.database_url, "ring_capture_test", 6) + .await + .expect("fixture pool"); + Self { config, pool } + } + + async fn drop(self) { + self.pool.close().await; + let admin = + super::super::db::connect_pool(&self.config.admin_url(), "ring_capture_cleanup", 1) + .await + .expect("cleanup connection"); + sqlx::query(&format!( + "DROP DATABASE {} WITH (FORCE)", + self.config.database_name() + )) + .execute(&admin) + .await + .expect("drop fixture database"); + admin.close().await; + } + } + + #[tokio::test] + async fn ring_lock_capture_identifies_holder_waiter_and_releases_cleanly() { + let db = Database::create().await; + let pool = &db.pool; + let mut holder = pool.begin().await.unwrap(); + sqlx::query("SET LOCAL application_name = 'ring_capture_holder'") + .execute(holder.as_mut()) + .await + .unwrap(); + sqlx::query("LOCK TABLE awa.lease_claims_0 IN ACCESS SHARE MODE") + .execute(holder.as_mut()) + .await + .unwrap(); + let holder_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(holder.as_mut()) + .await + .unwrap(); + let mut waiter = pool.begin().await.unwrap(); + sqlx::query("SET LOCAL application_name = 'ring_capture_waiter'") + .execute(waiter.as_mut()) + .await + .unwrap(); + let waiter_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(waiter.as_mut()) + .await + .unwrap(); + let waiting = tokio::spawn(async move { + sqlx::query("LOCK TABLE awa.lease_claims_0 IN ACCESS EXCLUSIVE MODE") + .execute(waiter.as_mut()) + .await + .unwrap(); + waiter.rollback().await.unwrap(); + }); + let deadline = + tokio::time::Instant::now() + crate::ci_timing::scaled_timeout(Duration::from_secs(5)); + loop { + let waiting: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM pg_locks WHERE pid=$1 AND NOT granted)", + ) + .bind(waiter_pid) + .fetch_one(pool) + .await + .unwrap(); + if waiting { + break; + } + assert!( + tokio::time::Instant::now() < deadline, + "exclusive waiter did not queue" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } + let value = collect(pool, "awa", 0, SQL).await; + assert_eq!(value["status"], "ok", "{value}"); + assert_eq!(value["resolved_relations"].as_array().unwrap().len(), 6); + let rows = value["locks"].as_array().unwrap(); + let held = rows + .iter() + .find(|row| row["pid"] == holder_pid && row["relation"] == "lease_claims_0") + .expect("holder evidence"); + assert_eq!(held["mode"], "AccessShareLock"); + assert_eq!(held["granted"], true); + assert_eq!(held["application_name"], "ring_capture_holder"); + assert!(held["tx_age_s"].as_f64().unwrap() >= 0.0); + let blocked = rows + .iter() + .find(|row| row["pid"] == waiter_pid && row["granted"] == false) + .expect("waiter evidence"); + assert_eq!(blocked["mode"], "AccessExclusiveLock"); + assert_eq!(blocked["application_name"], "ring_capture_waiter"); + assert_eq!(blocked["wait_event_type"], "Lock"); + assert!(blocked["blocking_pids"] + .as_array() + .unwrap() + .contains(&json!(holder_pid))); + assert!(value["blocker_activity"] + .as_array() + .unwrap() + .iter() + .any(|row| row["pid"] == holder_pid)); + holder.rollback().await.unwrap(); + tokio::time::timeout( + crate::ci_timing::scaled_timeout(Duration::from_secs(5)), + waiting, + ) + .await + .unwrap() + .unwrap(); + let clean = collect(pool, "awa", 0, SQL).await; + assert_eq!(clean["status"], "ok", "{clean}"); + assert!(clean["server_statement_at"].as_str().is_some()); + assert!(clean["server_observed_at"].as_str().is_some()); + assert_eq!(clean["resolved_relations"].as_array().unwrap().len(), 6); + assert!(clean["locks"] + .as_array() + .unwrap() + .iter() + .all(|row| row["pid"] != holder_pid && row["pid"] != waiter_pid)); + db.drop().await; + } + + #[tokio::test] + async fn ring_lock_capture_failure_and_unresolved_inventory_are_not_empty_success() { + let db = Database::create().await; + let value = collect(&db.pool, "missing_ring_schema", 0, SQL).await; + assert_eq!(value["status"], "partial"); + assert_eq!(value["resolved_relations"], json!([])); + assert_eq!(value["locks"], json!([])); + assert!(value["server_statement_at"].as_str().is_some()); + assert!(value["server_observed_at"].as_str().is_some()); + assert_eq!(value["unresolved_relations"].as_array().unwrap().len(), 6); + let failed = collect(&db.pool, "awa", 0, + "SELECT jsonb_build_object() FROM pg_sleep(1) WHERE $1::text IS NOT NULL AND $2::text[] IS NOT NULL").await; + assert_eq!(failed["status"], "missing", "{failed}"); + assert!( + failed.get("locks").is_none(), + "failed capture must not look unlocked" + ); + assert!(!failed["error"].as_str().unwrap().is_empty()); + db.drop().await; + } +} diff --git a/awa/tests/persona_soak/rng.rs b/awa/tests/persona_soak/rng.rs new file mode 100644 index 000000000..ca836b27b --- /dev/null +++ b/awa/tests/persona_soak/rng.rs @@ -0,0 +1,64 @@ +//! Deterministic pseudo-random numbers so a failing soak can be replayed from its seed. + +/// SplitMix64: tiny, fast, and good enough for workload shaping. +#[derive(Debug, Clone)] +pub struct Rng { + state: u64, +} + +impl Rng { + pub fn new(seed: u64) -> Self { + Self { state: seed } + } + + /// A generator for one task, derived from the run seed and a stable label so + /// that adding a persona does not shift another persona's sequence. + pub fn for_task(seed: u64, label: &str) -> Self { + let mut hash = seed ^ 0xcbf2_9ce4_8422_2325; + for byte in label.bytes() { + hash ^= u64::from(byte); + hash = hash.wrapping_mul(0x0000_0100_0000_01b3); + } + Self::new(hash) + } + + pub fn next_u64(&mut self) -> u64 { + self.state = self.state.wrapping_add(0x9e37_79b9_7f4a_7c15); + let mut value = self.state; + value = (value ^ (value >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + value = (value ^ (value >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + value ^ (value >> 31) + } + + /// Uniform in `0..bound` (`bound` must be non-zero). + pub fn below(&mut self, bound: u64) -> u64 { + self.next_u64() % bound.max(1) + } + + /// Uniform in `low..=high`. + pub fn range(&mut self, low: u64, high: u64) -> u64 { + low + self.below(high.saturating_sub(low) + 1) + } + + /// True with probability `percent` / 100. + pub fn chance(&mut self, percent: u64) -> bool { + self.below(100) < percent + } + + pub fn pick<'a, T>(&mut self, items: &'a [T]) -> Option<&'a T> { + if items.is_empty() { + None + } else { + Some(&items[self.below(items.len() as u64) as usize]) + } + } +} + +/// Nearest-rank percentile of a sorted slice (`quantile` in `0.0..=1.0`). +pub fn percentile(sorted: &[i64], quantile: f64) -> i64 { + if sorted.is_empty() { + return 0; + } + let index = ((sorted.len() - 1) as f64 * quantile).round() as usize; + sorted[index] +} diff --git a/awa/tests/persona_soak/roles.rs b/awa/tests/persona_soak/roles.rs new file mode 100644 index 000000000..bf819c3a0 --- /dev/null +++ b/awa/tests/persona_soak/roles.rs @@ -0,0 +1,143 @@ +//! The database roles of the shard-ops persona: the shape of a production install that +//! runs its application processes without superuser or schema ownership. +//! +//! * `OWNER` owns the database and the `awa` schema, runs migrations and `prepare_queue`. +//! * `MIGRATOR` is a migration-like role that is a member of `OWNER` and has default +//! privileges for the objects it creates. +//! * `APP_A` and `APP_B` are members of the `RUNTIME_GROUP` role and have no default +//! privileges of their own by default. An explicit comparison profile can share +//! their newly created objects too. They run producers and the maintenance leader. +//! +//! Grants follow the restricted-runtime section of `docs/security.md`. + +use super::config::{replace_credentials, SoakConfig}; +use super::db::connect_pool; +use sqlx::PgPool; + +pub const OWNER: &str = "awa_soak_owner"; +pub const MIGRATOR: &str = "awa_soak_migrator"; +pub const RUNTIME_GROUP: &str = "awa_soak_runtime"; +pub const APP_A: &str = "awa_soak_app_a"; +pub const APP_B: &str = "awa_soak_app_b"; +const PASSWORD: &str = "soak"; + +pub fn url_for_role(config: &SoakConfig, role: &str) -> String { + replace_credentials(&config.database_url, role, PASSWORD) +} + +/// Creates the cluster-level roles. Safe to repeat. +pub async fn create_roles(admin: &PgPool) -> Result<(), String> { + let login = [OWNER, MIGRATOR, APP_A, APP_B]; + let statements: Vec = login + .iter() + .map(|role| { + format!( + "DO $$ BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = '{role}') THEN + CREATE ROLE {role} LOGIN PASSWORD '{PASSWORD}' NOSUPERUSER NOCREATEDB NOCREATEROLE; + END IF; + END $$" + ) + }) + .chain([ + format!( + "DO $$ BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_roles WHERE rolname = '{RUNTIME_GROUP}') THEN + CREATE ROLE {RUNTIME_GROUP} NOLOGIN NOSUPERUSER NOCREATEDB NOCREATEROLE; + END IF; + END $$" + ), + format!("GRANT {OWNER} TO {MIGRATOR}"), + format!("GRANT {RUNTIME_GROUP} TO {APP_A}, {APP_B}"), + ]) + .collect(); + for statement in statements { + sqlx::query(&statement) + .execute(admin) + .await + .map_err(|error| format!("create roles: {error}"))?; + } + Ok(()) +} + +/// Applies the restricted-runtime grants. Run once the schema is migrated, over a +/// superuser connection to the soak database. +pub async fn grant_runtime( + superuser_pool: &PgPool, + database: &str, + schema_create: bool, + mixed_personas: bool, +) -> Result<(), String> { + // Comparison-only setup profile for older binaries whose application roles + // create their own lane sequences. Default off preserves regression coverage. + let shared_creator_grants = match std::env::var("AWA_SOAK_SHARED_CREATOR_GRANTS") { + Ok(value) if value == "1" => true, + Ok(value) if value == "0" => false, + Err(std::env::VarError::NotPresent) => false, + _ => return Err("AWA_SOAK_SHARED_CREATOR_GRANTS must be 0 or 1".into()), + }; + if shared_creator_grants && !schema_create { + return Err("shared creator grants require schema CREATE for the runtime group".into()); + } + let mut statements = vec![ + format!("GRANT USAGE ON SCHEMA awa TO {RUNTIME_GROUP}"), + format!("GRANT USAGE, SELECT, UPDATE ON ALL SEQUENCES IN SCHEMA awa TO {RUNTIME_GROUP}"), + format!("GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON ALL TABLES IN SCHEMA awa TO {RUNTIME_GROUP}"), + format!("GRANT EXECUTE ON ALL FUNCTIONS IN SCHEMA awa TO {RUNTIME_GROUP}"), + format!("REVOKE EXECUTE ON FUNCTION awa.install_queue_storage_substrate(TEXT, INT, INT, INT, BOOLEAN) FROM {RUNTIME_GROUP}"), + format!("GRANT TEMP ON DATABASE {database} TO {RUNTIME_GROUP}"), + ]; + if schema_create { + statements.push(format!("GRANT CREATE ON SCHEMA awa TO {RUNTIME_GROUP}")); + } + let mut creators = vec![OWNER.to_string(), MIGRATOR.to_string()]; + if mixed_personas { + // The other personas keep their superuser runtimes and producers. Their objects + // must stay usable by the application roles, which now lead maintenance for every + // queue. Application-role defaults remain absent unless the explicit + // shared-creator comparison profile below enables them. + let superuser: String = sqlx::query_scalar("SELECT current_user::text") + .fetch_one(superuser_pool) + .await + .map_err(|error| format!("current_user: {error}"))?; + creators.push(superuser); + } + if shared_creator_grants { + creators.extend([APP_A.to_string(), APP_B.to_string()]); + eprintln!("SOAK_SETUP shared_creator_grants=1 schema_create=1 application_creators={APP_A},{APP_B}"); + } + for creator in &creators { + statements.push(format!( + "ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa GRANT SELECT, INSERT, UPDATE, DELETE, TRUNCATE ON TABLES TO {RUNTIME_GROUP}" + )); + statements.push(format!( + "ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa GRANT USAGE, SELECT, UPDATE ON SEQUENCES TO {RUNTIME_GROUP}" + )); + statements.push(format!( + "ALTER DEFAULT PRIVILEGES FOR ROLE {creator} IN SCHEMA awa GRANT EXECUTE ON FUNCTIONS TO {RUNTIME_GROUP}" + )); + } + for statement in statements { + sqlx::query(&statement) + .execute(superuser_pool) + .await + .map_err(|error| format!("grant runtime privileges ({statement}): {error}"))?; + } + Ok(()) +} + +/// A pool that connects as `role`. +pub async fn pool_as( + config: &SoakConfig, + role: &str, + application_name: &str, + max_connections: u32, +) -> Result { + connect_pool( + &url_for_role(config, role), + application_name, + max_connections, + ) + .await + .map_err(|error| format!("connect as {role}: {error}")) +} diff --git a/awa/tests/persona_soak/run.rs b/awa/tests/persona_soak/run.rs new file mode 100644 index 000000000..58992d9e7 --- /dev/null +++ b/awa/tests/persona_soak/run.rs @@ -0,0 +1,303 @@ +//! Orchestration: set up the database, start the personas and the oracle, run the +//! load, let the system quiesce, reconcile, and report. + +use super::adminserver::AdminServer; +use super::config::{Persona, SoakConfig}; +use super::db::{self, Budget}; +use super::oracle::Oracle; +use super::personas; +use super::pgserver::PgWatch; +use super::report::{self, Outcome}; +use super::shared::Shared; +use super::stop::Stop; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +/// Connections kept free of persona pools: oracle, chaos sessions, admin probes. +const HEADROOM: u32 = 14; + +pub async fn run_soak(config: SoakConfig) -> Result, String> { + let worker_errors = super::worker_errors::install(false)?; + eprintln!( + "{}: preparing database {}", + report::HEADER, + config.database_name() + ); + db::recreate_database(&config).await?; + let oracle_pool = db::connect_pool(&config.database_url, "soak_oracle", 4) + .await + .map_err(|error| format!("oracle connect: {error}"))?; + let info = db::server_info(&oracle_pool) + .await + .map_err(|error| format!("server info: {error}"))?; + let watch = Arc::new( + PgWatch::begin(&config, &oracle_pool) + .await + .map_err(|error| format!("pg watch: {error}"))?, + ); + let budget = Budget::new(&info, HEADROOM); + eprintln!( + "{}: engine={} server={} max_connections={} connections_now={} usable={} seed={} personas={:?} load={}s scale={}", + report::HEADER, + info.engine, + info.version, + info.max_connections, + info.connections_at_start, + budget.usable(), + config.seed, + config.personas.iter().map(|persona| persona.name()).collect::>(), + config.load.as_secs(), + config.scale + ); + if info.engine != "queue_storage" { + return Err(format!( + "expected the queue_storage engine, found {}", + info.engine + )); + } + let shared = Shared::new(config, info, budget); + + // The UI personas talk to the real admin server, run as a child process. + let ui_personas = [ + Persona::Watchers, + Persona::Support, + Persona::UiOps, + Persona::Impatient, + ]; + let admin_server = if ui_personas + .iter() + .any(|persona| shared.config.has(*persona)) + { + std::fs::create_dir_all(&shared.config.report_dir).map_err(|error| error.to_string())?; + let server = AdminServer::start(&shared.config).await?; + shared.ui.set_base_url(server.base_url.clone()); + shared.note( + "run", + format!( + "admin server at {} (statement timeout {}s, client timeout {}s)", + server.base_url, + server.statement_timeout.as_secs(), + shared.config.ui_client_timeout.as_secs() + ), + ); + Some(server) + } else { + None + }; + + let terminal_observer = tokio::spawn(super::terminal_proof::run( + shared.clone(), + oracle_pool.clone(), + shared.drain_stop.clone(), + )); + + // Personas start their runtimes and then produce until the load stops. Chaos, the + // operator and the UI users reach other personas' fleets and queues, so they start + // once the rest are up. + let mut tasks: Vec<(Persona, tokio::task::JoinHandle>)> = Vec::new(); + for persona in [ + Persona::Mia, + Persona::Tom, + Persona::Priya, + Persona::Sam, + Persona::Python, + Persona::Fanout, + Persona::ShardOps, + Persona::ClaimAbandon, + ] { + if shared.config.has(persona) { + tasks.push((persona, spawn_persona(persona, shared.clone()))); + } + } + tokio::time::sleep(Duration::from_secs(8)).await; + for persona in [ + Persona::Alex, + Persona::Chaos, + Persona::Watchers, + Persona::Support, + Persona::UiOps, + Persona::Impatient, + ] { + if shared.config.has(persona) { + tasks.push((persona, spawn_persona(persona, shared.clone()))); + } + } + + let oracle_stop = Stop::new(); + let oracle_task = { + let oracle = Oracle::new(shared.clone(), oracle_pool.clone(), watch.clone()); + let stop = oracle_stop.clone(); + tokio::spawn(async move { oracle.run(stop).await }) + }; + shared.note( + "run", + format!("load phase started: {}s", shared.config.load.as_secs()), + ); + shared.started_load(Instant::now()); + + let load_deadline = Instant::now() + shared.config.load; + let mut next_progress = Instant::now() + Duration::from_secs(30); + while Instant::now() < load_deadline && !shared.load_stop.is_stopped() { + tokio::time::sleep(Duration::from_secs(1)).await; + if let Some((persona, _)) = tasks.iter().find(|(_, task)| task.is_finished()) { + // A persona only returns early on a setup failure or a panic. + let persona = *persona; + let index = tasks + .iter() + .position(|(p, _)| *p == persona) + .expect("finished task"); + let (_, task) = tasks.remove(index); + match task.await { + Ok(Ok(())) => shared.note("run", format!("{} finished early", persona.name())), + Ok(Err(error)) => shared.violation( + "persona_failed", + format!("{} stopped: {error}", persona.name()), + vec![], + ), + Err(error) => shared.violation( + "persona_panicked", + format!("{} panicked: {error}", persona.name()), + vec![], + ), + } + } + if Instant::now() >= next_progress { + next_progress = Instant::now() + Duration::from_secs(30); + progress(&shared); + } + } + + shared.note("run", "load phase over; letting the system quiesce"); + shared.load_stop.trigger(); + oracle_stop.trigger(); + let mut oracle = oracle_task + .await + .map_err(|error| format!("oracle task: {error}"))?; + let quiesced = oracle.wait_quiescent(shared.config.drain_grace).await; + shared + .quiesced + .store(quiesced, std::sync::atomic::Ordering::Relaxed); + shared.note("run", format!("quiesced: {quiesced}")); + super::terminal_proof::observe(&shared, &oracle_pool).await; + oracle.final_check(quiesced).await; + let excerpts = oracle.crash_scan(); + for excerpt in &excerpts { + shared.violation( + "server_crash_log", + "the Postgres log contains crash lines", + vec![excerpt.clone()], + ); + } + let log_tail = watch.log_since_start(&shared.config); + if shared.config.has(Persona::ShardOps) { + match &log_tail { + Some(log) => super::shard_oracle::scan_server_log(&shared, &oracle_pool, log).await, + None => shared.note( + "run", + "no Postgres log source is configured, so permission errors raised inside the runtimes are not scanned", + ), + } + } + + shared.drain_stop.trigger(); + for (persona, task) in tasks { + match tokio::time::timeout( + crate::ci_timing::scaled_timeout(Duration::from_secs(180)), + task, + ) + .await + { + Ok(Ok(Ok(()))) => {} + Ok(Ok(Err(error))) => shared.violation( + "persona_failed", + format!("{} stopped: {error}", persona.name()), + vec![], + ), + Ok(Err(error)) => shared.violation( + "persona_panicked", + format!("{} panicked: {error}", persona.name()), + vec![], + ), + Err(_) => shared.violation( + "shutdown_hang", + format!("{} did not shut down within 180s", persona.name()), + vec![], + ), + } + } + + terminal_observer + .await + .map_err(|error| format!("terminal observer task: {error}"))?; + personas::ui::analyse(&shared); + if let Some(mut server) = admin_server { + if !server.is_running() { + shared.violation( + "admin_server_died", + "the awa serve process exited during the run; see awa-serve.log", + vec![], + ); + } + server.stop(); + } + super::worker_errors::reconcile(&shared, &worker_errors).await; + let outcome = Outcome { + quiesced, + elapsed_seconds: shared.started.elapsed().as_secs_f64(), + crash_excerpts: excerpts, + log_tail, + }; + if let Err(error) = report::write(&shared, &outcome) { + eprintln!("could not write the report: {error}"); + } + oracle_pool.close().await; + Ok(shared) +} + +fn spawn_persona( + persona: Persona, + shared: Arc, +) -> tokio::task::JoinHandle> { + tokio::spawn(async move { + match persona { + Persona::Mia => personas::mia::run(shared).await, + Persona::Tom => personas::tom::run(shared).await, + Persona::Priya => personas::priya::run(shared).await, + Persona::Sam => personas::sam::run(shared).await, + Persona::Alex => personas::alex::run(shared).await, + Persona::Chaos => personas::chaos::run(shared).await, + Persona::Watchers => personas::ui::run_watchers(shared).await, + Persona::Support => personas::ui::run_support(shared).await, + Persona::UiOps => personas::ui::run_ui_operator(shared).await, + Persona::Impatient => personas::ui::run_impatient(shared).await, + Persona::Python => personas::python::run(shared).await, + Persona::Fanout => personas::fanout::run(shared).await, + Persona::ShardOps => personas::shardops::run(shared).await, + Persona::ClaimAbandon => personas::abandon::run(shared).await, + } + }) +} + +fn progress(shared: &Arc) { + let mut line = format!("t={:.0}s", shared.now_ms() as f64 / 1000.0); + for persona in Persona::ALL { + if !shared.config.has(persona) { + continue; + } + let counters = shared.stats(persona).counters(); + let enqueued: i64 = counters + .iter() + .filter(|(name, _)| name.starts_with("enqueued")) + .map(|(_, value)| *value) + .sum(); + line.push_str(&format!( + " {}[enq={} exec={}]", + persona.name(), + enqueued, + counters.get("executions").copied().unwrap_or(0) + )); + } + let violations = shared.violations.lock().expect("violations").len(); + line.push_str(&format!(" violations={violations}")); + shared.note("run", line); +} diff --git a/awa/tests/persona_soak/shard_oracle.rs b/awa/tests/persona_soak/shard_oracle.rs new file mode 100644 index 000000000..ef436c04b --- /dev/null +++ b/awa/tests/persona_soak/shard_oracle.rs @@ -0,0 +1,638 @@ +//! Checks specific to live `enqueue_shards` changes (ADR-025): every shard drains, +//! claim cursors keep moving while their shard has backlog, jobs run exactly once, +//! `ordering_key` jobs stay FIFO within a settled window, and no role ever hits a +//! permission or missing-lane error. + +use super::config::Persona; +use super::ledger::flag; +use super::shared::Shared; +use super::stop::Stop; +use sqlx::PgPool; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; +use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +const MAX_ROWS: usize = 25; +/// Same-key starts further behind the newest start than this are an order inversion. +const INVERSION_WINDOW: u64 = 8; +/// Seconds a shard above the configured count may keep backlog once every runtime has +/// observed the new count. +const ABOVE_COUNT_DRAIN_SECONDS: i64 = 45; + +#[derive(Debug, Clone, Default)] +pub struct QueuePlan { + pub configured: i16, + pub changed_at_ms: i64, + /// When every runtime and producer store had observed `configured`; None while the + /// fleet disagrees. + pub settled_at_ms: Option, +} + +/// What the shard-ops persona and the oracle share. +#[derive(Default)] +pub struct ShardOpsState { + pub plans: Mutex>, + /// Id of the current settled window; 0 while a change is propagating. + window: AtomicU64, + window_counter: AtomicU64, + producers_held: AtomicBool, + pub producers_inserting: AtomicI64, + seq: AtomicU64, + start_tick: AtomicU64, + /// Executions each enqueued job should have: 1, or 2 for a job that retries once. + expected: Mutex>, + /// Executions per job sequence number, and the job id that ran. + executions: Mutex>, + keyed_window: Mutex>, + keyed_starts: Mutex>, + reported: Mutex>, +} + +impl ShardOpsState { + pub fn next_seq(&self) -> u64 { + self.seq.fetch_add(1, Ordering::Relaxed) + 1 + } + + pub fn window(&self) -> u64 { + self.window.load(Ordering::SeqCst) + } + + pub fn open_window(&self) -> u64 { + let id = self.window_counter.fetch_add(1, Ordering::SeqCst) + 1; + self.window.store(id, Ordering::SeqCst); + id + } + + pub fn close_window(&self) { + self.window.store(0, Ordering::SeqCst); + } + + pub fn hold_producers(&self, held: bool) { + self.producers_held.store(held, Ordering::SeqCst); + } + + pub fn producers_held(&self) -> bool { + self.producers_held.load(Ordering::SeqCst) + } + + pub fn note_enqueued(&self, seq: u64, executions: u8) { + self.expected + .lock() + .expect("expected") + .insert(seq, executions); + } + + pub fn note_keyed(&self, key: u64, key_seq: u64, window: u64) { + self.keyed_window + .lock() + .expect("keyed window") + .insert((key, key_seq), window); + } + + /// Called by the worker at the start of every execution. + pub fn note_execution(&self, seq: u64, job_id: i64, key: Option<(u64, u64)>) { + let mut executions = self.executions.lock().expect("executions"); + let entry = executions.entry(seq).or_insert((0, job_id)); + entry.0 += 1; + drop(executions); + if let Some((key, key_seq)) = key { + let tick = self.start_tick.fetch_add(1, Ordering::SeqCst); + self.keyed_starts + .lock() + .expect("keyed starts") + .push((key, key_seq, tick)); + } + } + + pub fn configured(&self, queue: &str) -> i16 { + self.plans + .lock() + .expect("plans") + .get(queue) + .map_or(1, |plan| plan.configured.max(1)) + } + + /// Reports `summary` as a violation the first time `key` is seen and only counts + /// repeats, so a persistent failure does not bury the report. + pub fn violation_once( + &self, + shared: &Shared, + kind: &str, + key: &str, + summary: String, + rows: Vec, + ) { + let first = { + let mut reported = self.reported.lock().expect("reported"); + let count = reported.entry(format!("{kind}:{key}")).or_insert(0); + *count += 1; + *count == 1 + }; + if first { + shared.violation(kind, summary, rows); + } + shared + .stats(Persona::ShardOps) + .add(&format!("{kind}_total"), 1); + } +} + +/// One lane's cursors and backlog, as read from the store. +#[derive(Debug, Clone)] +struct ShardRow { + queue: String, + priority: i16, + shard: i16, + claim_seq: i64, + next_seq: i64, + available: i64, +} + +impl ShardRow { + fn json(&self, extra: &str) -> String { + format!( + "{{\"queue\":\"{}\",\"priority\":{},\"enqueue_shard\":{},\"claim_seq\":{},\"next_seq\":{},\"lag\":{},\"available\":{}{extra}}}", + self.queue, + self.priority, + self.shard, + self.claim_seq, + self.next_seq, + self.next_seq - self.claim_seq, + self.available + ) + } +} + +fn schema(shared: &Shared) -> String { + shared + .info + .queue_storage_schema + .clone() + .unwrap_or_else(|| "awa".to_string()) +} + +/// Per-shard claim cursor, enqueue cursor and the rows the cursor has not yet claimed +/// (the lag query from ADR-025, with the unclaimed rows counted exactly because aborted +/// transactions leave gaps in the sequences). +async fn shard_rows( + pool: &PgPool, + schema: &str, + queues: &[String], +) -> Result, sqlx::Error> { + let rows: Vec<(String, i16, i16, i64, i64, i64)> = sqlx::query_as(&format!( + "WITH claims AS MATERIALIZED ( + SELECT queue, priority, enqueue_shard, {schema}.sequence_next_value(seq_name) AS claim_seq + FROM {schema}.queue_claim_heads WHERE queue = ANY($1) + ), + enqueues AS MATERIALIZED ( + SELECT queue, priority, enqueue_shard, {schema}.sequence_next_value(seq_name) AS next_seq + FROM {schema}.queue_enqueue_heads WHERE queue = ANY($1) + ), + available AS ( + SELECT r.queue, r.priority, r.enqueue_shard, count(*)::bigint AS n + FROM {schema}.ready_entries r + JOIN claims c ON c.queue = r.queue AND c.priority = r.priority AND c.enqueue_shard = r.enqueue_shard + WHERE r.queue = ANY($1) AND r.lane_seq >= c.claim_seq + AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones t + WHERE t.ready_slot = r.ready_slot AND t.ready_generation = r.ready_generation + AND t.queue = r.queue AND t.priority = r.priority + AND t.enqueue_shard = r.enqueue_shard AND t.lane_seq = r.lane_seq + ) + GROUP BY r.queue, r.priority, r.enqueue_shard + ) + SELECT c.queue, c.priority, c.enqueue_shard, c.claim_seq::bigint, e.next_seq::bigint, + COALESCE(a.n, 0)::bigint + FROM claims c + JOIN enqueues e USING (queue, priority, enqueue_shard) + LEFT JOIN available a USING (queue, priority, enqueue_shard) + ORDER BY c.queue, c.priority, c.enqueue_shard" + )) + .bind(queues) + .fetch_all(pool) + .await?; + Ok(rows + .into_iter() + .map( + |(queue, priority, shard, claim_seq, next_seq, available)| ShardRow { + queue, + priority, + shard, + claim_seq, + next_seq, + available, + }, + ) + .collect()) +} + +/// Periodic per-shard checks until `stop`: claim cursors advance while their shard has +/// backlog and idle workers, and shards above a lowered count drain. +pub async fn watch(shared: Arc, pool: PgPool, queues: Vec, stop: Stop) { + let schema = schema(&shared); + // (queue, priority, shard) -> (claim cursor, when it last moved, already reported) + let mut cursors: HashMap<(String, i16, i16), (i64, i64, bool)> = HashMap::new(); + // (queue, priority, shard) -> when backlog above the configured count was first seen + let mut above_since: HashMap<(String, i16, i16), (i64, bool)> = HashMap::new(); + while !stop.sleep(Duration::from_secs(5)).await { + let rows = match shard_rows(&pool, &schema, &queues).await { + Ok(rows) => rows, + Err(error) => { + shared.incomplete(format!("shard cursor query failed: {error}"), vec![]); + continue; + } + }; + let now = shared.now_ms(); + let mut window = shared.config.stall_window; + if shared.chaos_active() { + window += shared.config.chaos_slack; + } + let mut shards_seen: BTreeMap = BTreeMap::new(); + for row in rows { + let key = (row.queue.clone(), row.priority, row.shard); + let seen = shards_seen.entry(row.queue.clone()).or_insert(0); + *seen = (*seen).max(i64::from(row.shard) + 1); + + let entry = cursors + .entry(key.clone()) + .or_insert((row.claim_seq, now, false)); + if entry.0 != row.claim_seq { + *entry = (row.claim_seq, now, false); + } else if let Some(queue) = shared.queues.get(&row.queue) { + let capacity = queue.capacity.load(Ordering::Relaxed); + let in_flight = queue.in_flight.load(Ordering::Relaxed); + let idle_workers = capacity > 0 && in_flight < capacity; + if row.available <= 0 || !idle_workers || queue.is_held() { + entry.1 = now; + } else if !entry.2 && (now - entry.1) as u64 >= window.as_millis() as u64 { + entry.2 = true; + shared.violation( + "shard_cursor_stalled", + format!( + "queue {} priority {} shard {}: {} unclaimed job(s), claim cursor frozen at {} for {}s with {in_flight}/{capacity} workers busy", + row.queue, + row.priority, + row.shard, + row.available, + row.claim_seq, + (now - entry.1) / 1000 + ), + vec![row.json(&format!( + ",\"configured_shards\":{}", + shared.shardops.configured(&row.queue) + ))], + ); + } + } + + let plan = shared + .shardops + .plans + .lock() + .expect("plans") + .get(&row.queue) + .cloned(); + let above = plan.as_ref().is_some_and(|plan| { + i32::from(row.shard) >= i32::from(plan.configured) + && plan + .settled_at_ms + .is_some_and(|at| now - at >= ABOVE_COUNT_DRAIN_SECONDS * 1000) + }); + if above && row.available > 0 { + let since = above_since.entry(key).or_insert((now, false)); + if !since.1 && now - since.0 >= ABOVE_COUNT_DRAIN_SECONDS * 1000 { + since.1 = true; + shared.violation( + "shard_above_count_not_draining", + format!( + "queue {} shard {} is above the configured count {} and still holds {} unclaimed job(s) {}s after every runtime observed the change", + row.queue, + row.shard, + plan.as_ref().map_or(0, |plan| plan.configured), + row.available, + (now - since.0) / 1000 + ), + vec![row.json("")], + ); + } + } else { + above_since.remove(&key); + } + } + for (queue, shards) in shards_seen { + shared + .stats(Persona::ShardOps) + .set_max(&format!("highest_shard_count_{queue}"), shards); + } + } +} + +/// After the system quiesced: every shard of every shard-ops queue is empty, every +/// enqueued job ran exactly the expected number of times, and keyed jobs kept their +/// order within each settled window. +pub async fn final_check(shared: &Arc, pool: &PgPool, queues: &[String]) { + let quiesced = shared.quiesced.load(Ordering::Relaxed); + let schema = schema(shared); + match shard_rows(pool, &schema, queues).await { + Ok(rows) => { + let stuck: Vec = rows + .iter() + .filter(|row| row.available > 0) + .map(|row| { + row.json(&format!( + ",\"configured_shards\":{}", + shared.shardops.configured(&row.queue) + )) + }) + .take(MAX_ROWS) + .collect(); + shared.run.add("shardops_lanes_final", rows.len() as i64); + if !stuck.is_empty() { + let message = format!( + "{} shard(s) still hold unclaimed jobs after the load stopped", + stuck.len() + ); + if quiesced { + shared.violation("shard_not_drained", message, stuck); + } else { + shared.finding("shard_backlog_at_end", message, stuck); + } + } + } + Err(error) => shared.incomplete(format!("final shard query failed: {error}"), vec![]), + } + + let state = &shared.shardops; + let expected = state.expected.lock().expect("expected").clone(); + let executions = state.executions.lock().expect("executions").clone(); + let mut missing = Vec::new(); + let mut repeated = Vec::new(); + let mut rescued = Vec::new(); + let mut short = Vec::new(); + for (seq, want) in &expected { + let (ran, job_id) = executions.get(seq).copied().unwrap_or((0, 0)); + let record = shared.ledger.get(job_id); + let row = format!( + "{{\"seq\":{seq},\"job_id\":{job_id},\"expected_executions\":{want},\"executions\":{ran},\"max_attempt\":{},\"ledger_flags\":{}}}", + record.as_ref().map_or(0, |record| record.max_attempt), + record.as_ref().map_or(0, |record| record.flags) + ); + if ran == 0 { + missing.push(row); + } else if ran > u16::from(*want) { + // A rescue re-runs a job whose attempt lost its lease, which is at-least-once + // delivery working as designed; any other repeat is a duplicate. + let was_rescued = record + .as_ref() + .is_some_and(|record| record.has(flag::RESCUED)); + if was_rescued { + rescued.push(row); + } else { + repeated.push(row); + } + } else if ran < u16::from(*want) { + short.push(row); + } + } + shared + .stats(Persona::ShardOps) + .add("tracked_jobs", expected.len() as i64); + let report = |kind: &str, what: &str, rows: Vec, fatal: bool| { + if rows.is_empty() { + return; + } + let summary = format!("{} job(s) {what}", rows.len()); + let rows = rows.into_iter().take(MAX_ROWS).collect(); + if fatal { + shared.violation(kind, summary, rows); + } else { + shared.finding(kind, summary, rows); + } + }; + report( + "lost_job", + "were enqueued across a shard change and never ran", + missing, + quiesced, + ); + report( + "rescued_reexecution", + "ran again after a stale-heartbeat rescue", + rescued, + false, + ); + report( + "duplicate_execution", + "ran more often than their handler asked for", + repeated, + true, + ); + report( + "retry_not_executed", + "ran fewer times than their handler asked for", + short, + quiesced, + ); + + let windows = state.keyed_window.lock().expect("keyed window").clone(); + let mut starts = state.keyed_starts.lock().expect("keyed starts").clone(); + starts.sort_by_key(|(_, _, tick)| *tick); + let mut newest: HashMap<(u64, u64), u64> = HashMap::new(); + let mut inversions = Vec::new(); + let mut checked = 0u64; + let mut worst = 0u64; + let mut seen_per_key: HashSet<(u64, u64)> = HashSet::new(); + for (key, key_seq, _) in starts { + let window = windows.get(&(key, key_seq)).copied().unwrap_or(0); + if window == 0 { + continue; + } + if !seen_per_key.insert((key, key_seq)) { + continue; + } + checked += 1; + let top = newest.entry((key, window)).or_insert(0); + if key_seq + INVERSION_WINDOW < *top { + worst = worst.max(*top - key_seq); + if inversions.len() < MAX_ROWS { + inversions.push(format!( + "{{\"ordering_key\":{key},\"window\":{window},\"started_seq\":{key_seq},\"newest_started_seq\":{top}}}" + )); + } + } else { + *top = (*top).max(key_seq); + } + } + shared + .stats(Persona::ShardOps) + .add("keyed_starts_checked", checked as i64); + shared + .stats(Persona::ShardOps) + .set_max("keyed_worst_inversion", worst as i64); + if !inversions.is_empty() { + shared.violation( + "ordering_key_inversion", + format!( + "keyed jobs started out of order within a settled shard-count window (worst distance {worst})" + ), + inversions, + ); + } +} + +struct LoggedError { + role: String, + application: String, + message: String, + statement: String, +} + +fn is_lane_error(message: &str) -> bool { + message.contains("permission denied") + || message.contains("must be owner") + || message.contains("has not been provisioned") + || message.contains("provision") + || is_lane_race(message) +} + +/// Two sessions creating the same lane sequence at once: `CREATE SEQUENCE IF NOT EXISTS` +/// loses the race with a unique violation on the catalog. +fn is_lane_race(message: &str) -> bool { + message.contains("pg_class_relname_nsp_index") +} + +/// Permission-denied and missing-lane errors in the server log, one violation per role +/// and error message with the lane, the sequence's owner and the statement. +pub async fn scan_server_log(shared: &Arc, pool: &PgPool, log: &str) { + let mut errors: Vec = Vec::new(); + let lines: Vec<&str> = log.lines().collect(); + let is_entry = |line: &str| { + let bytes = line.as_bytes(); + bytes.len() > 20 && bytes[4] == b'-' && bytes[7] == b'-' && bytes[10] == b' ' + }; + let mut index = 0; + while index < lines.len() { + let line = lines[index]; + index += 1; + let Some((head, message)) = line.split_once(" ERROR: ") else { + continue; + }; + if !is_lane_error(message) { + continue; + } + let identity = head + .rsplit_once("] ") + .map_or("?", |(_, identity)| identity.trim()); + let (role, application) = identity.split_once('@').unwrap_or((identity, "")); + let (role, application) = (role.to_string(), application.to_string()); + let mut statement = String::new(); + while index < lines.len() && !is_entry(lines[index]) { + if let Some(text) = lines[index].trim_start().strip_prefix("STATEMENT:") { + statement = text.trim().chars().take(400).collect(); + } else if !statement.is_empty() { + statement.push(' '); + statement.push_str(lines[index].trim()); + } + index += 1; + } + errors.push(LoggedError { + role, + application, + message: message.trim().to_string(), + statement, + }); + } + if errors.is_empty() { + return; + } + let schema = schema(shared); + let mut groups: BTreeMap<(String, String), (u64, String, BTreeSet)> = BTreeMap::new(); + for error in &errors { + let entry = groups + .entry((error.role.clone(), error.message.clone())) + .or_insert((0, error.statement.clone(), BTreeSet::new())); + entry.0 += 1; + entry.2.insert(error.application.clone()); + } + shared + .run + .add("shard_lane_errors_logged", errors.len() as i64); + for ((role, message), (count, statement, applications)) in groups.into_iter().take(MAX_ROWS) { + let (lane, owner) = describe_sequence(pool, &schema, &message).await; + let applications = applications.into_iter().collect::>().join(","); + let kind = if is_lane_race(&message) { + "lane_creation_race" + } else { + "permission_denied" + }; + shared.violation( + kind, + format!("role {role}: {message} ({count}x; lane {lane}; sequence owner {owner})"), + vec![format!( + "{{\"role\":\"{role}\",\"applications\":{},\"error\":{},\"lane\":{},\"sequence_owner\":{},\"occurrences\":{count},\"statement\":{}}}", + json_string(&applications), + json_string(&message), + json_string(&lane), + json_string(&owner), + json_string(&statement) + )], + ); + } +} + +fn json_string(text: &str) -> String { + serde_json::Value::String(text.to_string()).to_string() +} + +/// The lane and owning role of the sequence an error message names, when it names one. +async fn describe_sequence(pool: &PgPool, schema: &str, message: &str) -> (String, String) { + let Some(name) = message + .split("sequence ") + .nth(1) + .map(|rest| { + rest.split_whitespace() + .next() + .unwrap_or("") + .trim_matches('"') + }) + .map(|name| { + name.rsplit('.') + .next() + .unwrap_or(name) + .trim_matches('"') + .to_string() + }) + .filter(|name| !name.is_empty()) + else { + return ("unknown".into(), "unknown".into()); + }; + let owner: Option = sqlx::query_scalar( + "SELECT pg_get_userbyid(c.relowner) FROM pg_class c + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = $1 AND c.relname = $2", + ) + .bind(schema) + .bind(&name) + .fetch_optional(pool) + .await + .ok() + .flatten(); + let lane: Option = sqlx::query_scalar(&format!( + "SELECT queue || ' priority ' || priority || ' shard ' || enqueue_shard + FROM {schema}.queue_enqueue_heads WHERE seq_name = $1 + UNION ALL + SELECT queue || ' priority ' || priority || ' shard ' || enqueue_shard + FROM {schema}.queue_claim_heads WHERE seq_name = $1 + LIMIT 1" + )) + .bind(&name) + .fetch_optional(pool) + .await + .ok() + .flatten(); + ( + lane.unwrap_or_else(|| format!("sequence {name}")), + owner.unwrap_or_else(|| "unknown".into()), + ) +} diff --git a/awa/tests/persona_soak/shared.rs b/awa/tests/persona_soak/shared.rs new file mode 100644 index 000000000..2bd038081 --- /dev/null +++ b/awa/tests/persona_soak/shared.rs @@ -0,0 +1,468 @@ +//! State shared by every persona, the oracle and the report. + +use super::config::{Persona, SoakConfig}; +use super::db::{Budget, ServerInfo}; +use super::fleet::Fleet; +use super::ledger::{expect, flag, JobRecord, Ledger}; +use super::registry::{InFlightGuard, QueueRegistry}; +use super::rng::percentile; +use super::stop::Stop; +use awa::{JobRow, UntypedJobEvent}; +use serde::Serialize; +use std::collections::{BTreeMap, HashMap}; +use std::sync::atomic::{AtomicBool, AtomicI64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Instant; + +#[derive(Debug, Clone, Serialize)] +pub struct TimelineEvent { + pub at_ms: i64, + pub source: String, + pub text: String, +} + +#[derive(Debug, Clone, Serialize)] +pub struct Violation { + pub kind: String, + pub at_ms: i64, + pub during_chaos: bool, + pub summary: String, + /// The violating rows, as JSON text. + pub rows: Vec, +} + +#[derive(Default)] +pub struct Stats { + counters: Mutex>, + samples: Mutex>>, +} + +impl Stats { + pub fn add(&self, name: &str, amount: i64) { + *self + .counters + .lock() + .expect("stats") + .entry(name.to_string()) + .or_insert(0) += amount; + } + + pub fn set_max(&self, name: &str, value: i64) { + let mut counters = self.counters.lock().expect("stats"); + let slot = counters.entry(name.to_string()).or_insert(i64::MIN); + *slot = (*slot).max(value); + } + + pub fn get(&self, name: &str) -> i64 { + self.counters + .lock() + .expect("stats") + .get(name) + .copied() + .unwrap_or(0) + } + + /// Records a latency-style sample; kept bounded so long runs stay small. + pub fn observe(&self, name: &str, value: i64) { + let mut samples = self.samples.lock().expect("stats"); + let list = samples.entry(name.to_string()).or_default(); + if list.len() < 400_000 { + list.push(value); + } + } + + pub fn counters(&self) -> BTreeMap { + self.counters.lock().expect("stats").clone() + } + + pub fn summaries(&self) -> BTreeMap { + self.samples + .lock() + .expect("stats") + .iter() + .map(|(name, values)| { + let mut sorted = values.clone(); + sorted.sort_unstable(); + ( + name.clone(), + SampleSummary { + count: sorted.len() as i64, + p50: percentile(&sorted, 0.5), + p99: percentile(&sorted, 0.99), + max: sorted.last().copied().unwrap_or(0), + }, + ) + }) + .collect() + } +} + +#[derive(Debug, Clone, Serialize)] +pub struct SampleSummary { + pub count: i64, + pub p50: i64, + pub p99: i64, + pub max: i64, +} + +pub struct Shared { + pub config: SoakConfig, + pub info: ServerInfo, + pub started: Instant, + pub ledger: Ledger, + pub queues: QueueRegistry, + pub budget: Budget, + pub load_stop: Stop, + pub drain_stop: Stop, + pub timeline: Mutex>, + pub violations: Mutex>, + /// Required oracle evidence was unavailable; the run cannot establish a pass. + pub oracle_incomplete: AtomicBool, + /// Problems that are not invariant violations (slow admin queries); reported, not fatal. + pub findings: Mutex>, + pub stats: HashMap, + pub ui: super::personas::ui::UiMetrics, + /// Handlers currently running in the Python process, from its own reports. + pub python_active: AtomicI64, + /// Set once the post-load drain finished with no live work left. + pub quiesced: std::sync::atomic::AtomicBool, + /// The running Python process, so the chaos persona can kill it. + pub python: Mutex>>, + /// Run-wide measurements (resource use, oracle probes). + pub run: Stats, + pub fleets: Mutex>>, + /// State of the shard-ops persona and the checks that go with it. + pub shardops: super::shard_oracle::ShardOpsState, + /// State of the claim-and-abandon persona and the checks that go with it. + pub abandon: super::abandon_oracle::AbandonState, + chaos_until_ms: AtomicI64, +} + +impl Shared { + pub fn new(config: SoakConfig, info: ServerInfo, budget: Budget) -> Arc { + let stats = Persona::ALL + .into_iter() + .map(|persona| (persona, Stats::default())) + .collect(); + Arc::new(Self { + config, + info, + started: Instant::now(), + ledger: Ledger::new(), + queues: QueueRegistry::default(), + budget, + load_stop: Stop::new(), + drain_stop: Stop::new(), + timeline: Mutex::new(Vec::new()), + violations: Mutex::new(Vec::new()), + oracle_incomplete: AtomicBool::new(false), + findings: Mutex::new(Vec::new()), + stats, + run: Stats::default(), + ui: Default::default(), + python_active: AtomicI64::new(0), + quiesced: std::sync::atomic::AtomicBool::new(false), + python: Mutex::new(None), + fleets: Mutex::new(HashMap::new()), + shardops: Default::default(), + abandon: Default::default(), + chaos_until_ms: AtomicI64::new(0), + }) + } + + pub fn started_load(&self, _at: Instant) {} + + pub fn now_ms(&self) -> i64 { + self.started.elapsed().as_millis() as i64 + } + + pub fn stats(&self, persona: Persona) -> &Stats { + &self.stats[&persona] + } + + pub fn note(&self, source: &str, text: impl Into) { + let text = text.into(); + let at_ms = self.now_ms(); + eprintln!("[{:>7.1}s] {source:<6} {text}", at_ms as f64 / 1000.0); + self.timeline.lock().expect("timeline").push(TimelineEvent { + at_ms, + source: source.to_string(), + text, + }); + } + + /// Marks a disruptive window (kill, failover, connection termination) so that the + /// time-based checks allow the system its recovery time. + pub fn chaos_window(&self, duration: std::time::Duration) { + let until = self.now_ms() + duration.as_millis() as i64; + self.chaos_until_ms.fetch_max(until, Ordering::Relaxed); + } + + pub fn chaos_active(&self) -> bool { + self.now_ms() < self.chaos_until_ms.load(Ordering::Relaxed) + } + + pub fn violation(&self, kind: &str, summary: impl Into, rows: Vec) { + let violation = Violation { + kind: kind.to_string(), + at_ms: self.now_ms(), + during_chaos: self.chaos_active(), + summary: summary.into(), + rows, + }; + eprintln!( + "[{:>7.1}s] VIOLATION {}: {}", + violation.at_ms as f64 / 1000.0, + violation.kind, + violation.summary + ); + self.violations.lock().expect("violations").push(violation); + } + + pub fn incomplete(&self, summary: impl Into, rows: Vec) { + self.run.add("oracle_incomplete_checks", 1); + if !self.oracle_incomplete.swap(true, Ordering::Relaxed) { + self.violation("oracle_incomplete", summary, rows); + } + } + + /// Records a non-fatal finding: something worth a bug report that is not one of the + /// invariants the soak fails on. + pub fn finding(&self, kind: &str, summary: impl Into, rows: Vec) { + let finding = Violation { + kind: kind.to_string(), + at_ms: self.now_ms(), + during_chaos: self.chaos_active(), + summary: summary.into(), + rows, + }; + eprintln!( + "[{:>7.1}s] FINDING {}: {}", + finding.at_ms as f64 / 1000.0, + finding.kind, + finding.summary + ); + self.findings.lock().expect("findings").push(finding); + } + + /// Records a successful operator cancel: the job may legitimately end cancelled, a + /// handler must not start after it, and it must not overwrite a completion. + pub fn record_operator_cancel( + &self, + persona: Persona, + job_id: i64, + completed_before: bool, + call_started_ms: u32, + ) { + let now = self.now_ms() as u32; + self.stats(persona).add("cancelled", 1); + self.ledger.with( + job_id, + || JobRecord::new(persona, u32::MAX, now, expect::ANY_TERMINAL, 0), + |record| { + record.flags |= flag::OPERATOR_REMOVED | flag::OPERATOR_CANCEL_OK; + record.operator_cancel_ms = now; + }, + ); + if completed_before { + self.violation( + "cancel_overwrote_completed", + format!("operator cancel of job {job_id} succeeded after it was reported completed (call began at {call_started_ms} ms)"), + vec![format!("{{\"job_id\":{job_id}}}")], + ); + } + } + + pub fn fleet(&self, name: &str) -> Option> { + self.fleets.lock().expect("fleets").get(name).cloned() + } + + pub fn fleets_of(&self, persona: Persona) -> Vec> { + self.fleets + .lock() + .expect("fleets") + .values() + .filter(|fleet| fleet.persona == persona) + .cloned() + .collect() + } + + /// Records the start of a handler. Hold the returned guard for the duration of + /// `perform`; dropping it, including when the runtime is killed, ends the execution. + pub fn begin_execution( + self: &Arc, + persona: Persona, + job: &JobRow, + no_events: bool, + ) -> ExecGuard { + let queue = self.queues.get(&job.queue); + let queue_index = queue.as_ref().map_or(u32::MAX, |queue| queue.index); + let now = self.now_ms() as u32; + let default_flags = if no_events { flag::NO_EVENTS } else { 0 }; + let (ran_after_cancel, overlapped) = self.ledger.with( + job.id, + || { + JobRecord::new( + persona, + queue_index, + now, + expect::ANY_TERMINAL, + default_flags, + ) + }, + |record| { + let overlapped = record.in_flight > 0; + record.overlapped |= overlapped; + record.in_flight = record.in_flight.saturating_add(1); + record.executions = record.executions.saturating_add(1); + record.max_attempt = record.max_attempt.max(job.attempt); + record.latest_run_lease = record.latest_run_lease.max(job.run_lease); + if record.first_start_ms == 0 { + record.first_start_ms = now.max(1); + } + record.last_start_ms = now; + ( + record.operator_cancel_ms != 0 && now > record.operator_cancel_ms, + overlapped, + ) + }, + ); + if ran_after_cancel { + self.violation( + "ran_after_cancel", + format!("job {} started after its operator cancel committed", job.id), + vec![format!( + "{{\"job_id\":{},\"queue\":\"{}\",\"attempt\":{}}}", + job.id, job.queue, job.attempt + )], + ); + } + if overlapped { + self.stats(persona).add("overlapping_executions", 1); + } + if let Some(name) = job + .metadata + .get("cron_name") + .and_then(|value| value.as_str()) + { + let fire = job + .metadata + .get("cron_fire_time") + .and_then(|value| value.as_str()) + .unwrap_or_default(); + // Manually triggered jobs carry the schedule's name but no fire time. + if !fire.is_empty() + && !self + .ledger + .note_cron_fire(name.to_string(), fire.to_string(), job.id) + { + self.violation( + "cron_double_fire", + format!( + "cron {name} fire {fire} produced more than one job (latest {})", + job.id + ), + vec![], + ); + } + } + if let Some(queue) = &queue { + queue.started.fetch_add(1, Ordering::Relaxed); + } + self.stats(persona).add("executions", 1); + ExecGuard { + _flight: InFlightGuard::new(queue), + shared: self.clone(), + job_id: job.id, + } + } + + /// Records a lifecycle event emitted by a runtime after the transition committed. + pub fn on_event(&self, event: &UntypedJobEvent) { + let job = event.job(); + let now = self.now_ms() as u32; + let queue_index = self + .queues + .get(&job.queue) + .map_or(u32::MAX, |queue| queue.index); + let persona = self + .queues + .get(&job.queue) + .map_or(Persona::Chaos, |queue| queue.persona); + let mut double_complete = false; + self.ledger.with( + job.id, + || JobRecord::new(persona, queue_index, now, expect::ANY_TERMINAL, 0), + |record| match event { + UntypedJobEvent::Completed { .. } => { + record.note_terminal(job.run_lease); + record.completed_events = record.completed_events.saturating_add(1); + record.completed_event_ms = now; + record.flags |= flag::COMPLETED_EVENT; + double_complete = record.completed_events > 1 && !record.has(flag::DLQ_RETRIED); + } + UntypedJobEvent::Exhausted { .. } => { + record.note_terminal(job.run_lease); + record.exhausted_events = record.exhausted_events.saturating_add(1); + record.flags |= flag::EXHAUSTED; + } + UntypedJobEvent::Cancelled { .. } => { + record.note_terminal(job.run_lease); + record.cancel_events = record.cancel_events.saturating_add(1); + record.flags |= flag::HANDLER_CANCELLED; + } + UntypedJobEvent::Rescued { .. } => record.flags |= flag::RESCUED, + _ => {} + }, + ); + match event { + UntypedJobEvent::Rescued { reason, .. } => { + if persona == Persona::ShardOps { + self.note( + "events", + format!( + "job {} on {} rescued ({}) at attempt {}", + job.id, + job.queue, + reason.as_str(), + job.attempt + ), + ); + } + self.stats(persona) + .add(&format!("rescued_{}", reason.as_str()), 1); + } + UntypedJobEvent::Exhausted { .. } => self.stats(persona).add("exhausted", 1), + UntypedJobEvent::Completed { .. } => self.stats(persona).add("completed_events", 1), + UntypedJobEvent::Cancelled { .. } => self.stats(persona).add("handler_cancelled", 1), + UntypedJobEvent::Retried { .. } => self.stats(persona).add("retried", 1), + _ => {} + } + if double_complete { + self.violation( + "duplicate_terminal_event", + format!("job {} reported completed more than once", job.id), + vec![format!( + "{{\"job_id\":{},\"queue\":\"{}\"}}", + job.id, job.queue + )], + ); + } + } +} + +pub struct ExecGuard { + _flight: InFlightGuard, + shared: Arc, + job_id: i64, +} + +impl Drop for ExecGuard { + fn drop(&mut self) { + self.shared.ledger.with( + self.job_id, + || JobRecord::new(Persona::Chaos, u32::MAX, 0, expect::ANY_TERMINAL, 0), + |record| record.in_flight = record.in_flight.saturating_sub(1), + ); + } +} diff --git a/awa/tests/persona_soak/stop.rs b/awa/tests/persona_soak/stop.rs new file mode 100644 index 000000000..6264c06e2 --- /dev/null +++ b/awa/tests/persona_soak/stop.rs @@ -0,0 +1,50 @@ +//! A cheap cooperative stop signal shared by every task in the run. + +use std::sync::Arc; +use std::time::Duration; +use tokio::sync::watch; + +#[derive(Clone)] +pub struct Stop { + sender: Arc>, +} + +impl Stop { + pub fn new() -> Self { + Self { + sender: Arc::new(watch::channel(false).0), + } + } + + pub fn trigger(&self) { + self.sender.send_replace(true); + } + + pub fn is_stopped(&self) -> bool { + *self.sender.borrow() + } + + /// Sleeps for `duration` or until stopped. Returns true when stopped. + pub async fn sleep(&self, duration: Duration) -> bool { + if self.is_stopped() { + return true; + } + let mut receiver = self.sender.subscribe(); + tokio::select! { + _ = tokio::time::sleep(duration) => false, + _ = receiver.wait_for(|stopped| *stopped) => true, + } + } + + /// Completes once stopped. + pub async fn wait(&self) { + let mut receiver = self.sender.subscribe(); + let _ = receiver.wait_for(|stopped| *stopped).await; + } +} + +impl Default for Stop { + fn default() -> Self { + Self::new() + } +} diff --git a/awa/tests/persona_soak/terminal_proof.rs b/awa/tests/persona_soak/terminal_proof.rs new file mode 100644 index 000000000..1e9773afc --- /dev/null +++ b/awa/tests/persona_soak/terminal_proof.rs @@ -0,0 +1,122 @@ +//! Durable evidence for returned eventful jobs and the NO_EVENTS sample. +//! Handler returns are not commits; lifecycle hooks are best-effort notifications. + +use super::config::Persona; +use super::ledger::{expect, flag, JobRecord, TERMINAL_PROOF_STRIDE}; +use super::shared::Shared; +use super::stop::Stop; +use sqlx::PgPool; +use std::sync::Arc; +use std::time::Duration; + +const MAX_ROWS: usize = 100_000; + +pub async fn run(shared: Arc, pool: PgPool, stop: Stop) { + let mut tick = tokio::time::interval(Duration::from_secs(1)); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + _ = stop.wait() => return, + _ = tick.tick() => observe(&shared, &pool).await, + } + } +} + +/// One statement gives one snapshot. Explicit eventful attempts have bounded +/// input/result counts and a two-second statement timeout. NO_EVENTS keeps its +/// immutable sample; neither a returned handler nor a missed query proves a commit. +pub async fn observe(shared: &Shared, pool: &PgPool) { + let schema = shared.info.queue_storage_schema.as_deref().unwrap_or("awa"); + let mut pending = Vec::new(); + shared.ledger.for_each(|id, record| { + if record.has(flag::RETURNED_COMPLETED) + && !record.has(flag::NO_EVENTS) + && !record.has_durable_terminal() + && !record.has(flag::OPERATOR_REMOVED) + && !record.has(flag::DLQ_PURGED) + && pending.len() <= MAX_ROWS + { + pending.push((id, record.latest_run_lease)); + } + }); + if pending.len() > MAX_ROWS { + shared.incomplete( + "durable terminal observer exceeded its pending job cap", + vec![], + ); + return; + } + let (pending_ids, pending_leases): (Vec<_>, Vec<_>) = pending.into_iter().unzip(); + let result = async { + let mut tx = pool.begin().await?; + sqlx::query("SET LOCAL statement_timeout = '2s'") + .execute(tx.as_mut()) + .await?; + let rows: Vec<(i64, i64, String)> = sqlx::query_as(&format!( + "WITH pending AS ( + SELECT * FROM unnest($3::bigint[], $4::bigint[]) AS p(job_id, run_lease) + ) + SELECT done.job_id, done.run_lease, done.state::text + FROM {schema}.done_entries AS done + JOIN pending USING (job_id, run_lease) + UNION ALL + SELECT dlq.job_id, dlq.run_lease, 'failed' + FROM {schema}.dlq_entries AS dlq + JOIN pending USING (job_id, run_lease) + UNION ALL + SELECT item.job_id, item.run_lease, 'completed' + FROM {schema}.receipt_completion_batches AS batch + CROSS JOIN LATERAL unnest(batch.job_ids, batch.run_leases) AS item(job_id, run_lease) + JOIN pending USING (job_id, run_lease) + WHERE batch.job_ids && $3 + UNION ALL + SELECT job_id, run_lease, state::text FROM {schema}.done_entries + WHERE job_id % $1 = 0 AND finalized_at >= now() - interval '30 seconds' + UNION ALL + SELECT job_id, run_lease, 'failed' FROM {schema}.dlq_entries + WHERE job_id % $1 = 0 AND dlq_at >= now() - interval '30 seconds' + UNION ALL + SELECT item.job_id, item.run_lease, 'completed' + FROM {schema}.receipt_completion_batches AS batch + CROSS JOIN LATERAL unnest(batch.job_ids, batch.run_leases) AS item(job_id, run_lease) + WHERE item.job_id % $1 = 0 AND batch.finalized_at >= now() - interval '30 seconds' + LIMIT $2" + )) + .bind(TERMINAL_PROOF_STRIDE) + .bind(MAX_ROWS as i64 + 1) + .bind(pending_ids) + .bind(pending_leases) + .fetch_all(tx.as_mut()) + .await?; + tx.commit().await?; + Ok::<_, sqlx::Error>(rows) + } + .await; + let rows = match result { + Ok(rows) if rows.len() <= MAX_ROWS => rows, + Ok(_) => { + shared.incomplete("durable terminal observer exceeded its row cap", vec![]); + return; + } + Err(error) => { + shared.incomplete( + format!("durable terminal observer query failed: {error}"), + vec![], + ); + return; + } + }; + shared + .run + .add("terminal_proof_observer_rows", rows.len() as i64); + let now = shared.now_ms() as u32; + for (id, run_lease, state) in rows { + shared.ledger.with( + id, + || JobRecord::new(Persona::Chaos, u32::MAX, now, expect::ANY_TERMINAL, 0), + |record| { + record.note_observed_terminal(run_lease, &state); + }, + ); + } +} diff --git a/awa/tests/persona_soak/worker_errors.rs b/awa/tests/persona_soak/worker_errors.rs new file mode 100644 index 000000000..89dfb5070 --- /dev/null +++ b/awa/tests/persona_soak/worker_errors.rs @@ -0,0 +1,241 @@ +//! A structured gate for a known worker validation failure, independent of log text. +//! +//! This subscriber covers Rust runtimes in this process. The separately spawned +//! Rust claiming child forwards only a fixed marker over its existing output pipe. +//! Python-extension internals have no tracing bridge and are not covered by it. + +use super::shared::Shared; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::sync::{Arc, OnceLock}; +use std::time::Duration; +use tracing::field::{Field, Visit}; +use tracing::{Event, Subscriber}; +use tracing_subscriber::filter::LevelFilter; +use tracing_subscriber::layer::{Context, SubscriberExt}; +use tracing_subscriber::{Layer, Registry}; + +const MISSING_READY: &str = "ready row missing for deleted lease"; +const AMBIGUOUS_READY: &str = "ready body ambiguous for deleted lease"; +const CHILD_MARKER: &str = "AWA_SOAK_WORKER_ERROR ready_body_unavailable"; +static EVIDENCE: OnceLock, String>> = OnceLock::new(); + +#[derive(Default)] +pub struct WorkerErrors { + invalid_ready_body: AtomicU64, + child_readers: AtomicUsize, +} + +#[derive(Default)] +struct ErrorFields { + invalid_ready_body: bool, +} + +impl Visit for ErrorFields { + fn record_str(&mut self, field: &Field, value: &str) { + if matches!(field.name(), "error" | "message") { + self.invalid_ready_body |= + value.contains(MISSING_READY) || value.contains(AMBIGUOUS_READY); + } + } + + fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) { + if matches!(field.name(), "error" | "message") { + self.invalid_ready_body |= { + let formatted = format!("{value:?}"); + formatted.contains(MISSING_READY) || formatted.contains(AMBIGUOUS_READY) + }; + } + } +} + +#[derive(Clone)] +struct ErrorLayer { + evidence: Arc, + child: bool, +} + +impl Layer for ErrorLayer { + fn on_event(&self, event: &Event<'_>, _context: Context<'_, S>) { + if *event.metadata().level() > tracing::Level::WARN { + return; + } + let mut fields = ErrorFields::default(); + event.record(&mut fields); + if fields.invalid_ready_body { + self.evidence + .invalid_ready_body + .fetch_add(1, Ordering::Relaxed); + if self.child { + // One short synchronous line only for this invariant failure. + // The parent's reader is tracked through EOF before reconciliation. + eprintln!("{CHILD_MARKER}"); + } + } + } +} + +/// Install once, before spawning any worker. An existing unrelated subscriber +/// makes setup fail rather than silently dropping the error gate. +pub fn install(child: bool) -> Result, String> { + EVIDENCE + .get_or_init(|| { + let evidence = Arc::new(WorkerErrors::default()); + let subscriber = Registry::default() + .with( + ErrorLayer { + evidence: evidence.clone(), + child, + } + .with_filter(LevelFilter::WARN), + ) + .with( + tracing_subscriber::fmt::layer() + .with_ansi(false) + .with_writer(std::io::stderr) + .with_filter(LevelFilter::WARN), + ); + tracing::subscriber::set_global_default(subscriber) + .map_err(|error| format!("worker error tracing unavailable: {error}"))?; + Ok(evidence) + }) + .clone() +} + +pub struct ChildOutputGuard(Arc); + +impl Drop for ChildOutputGuard { + fn drop(&mut self) { + self.0.child_readers.fetch_sub(1, Ordering::Release); + } +} + +/// Register before spawning the existing output reader, so a fast child exit +/// cannot race the final pending-reader check. +pub fn child_output_guard() -> ChildOutputGuard { + let evidence = EVIDENCE + .get() + .expect("worker error tracing installed before child spawn") + .as_ref() + .expect("worker error tracing installation succeeded") + .clone(); + evidence.child_readers.fetch_add(1, Ordering::Relaxed); + ChildOutputGuard(evidence) +} + +pub fn absorb_child_marker(line: &str) -> bool { + if line != CHILD_MARKER { + return false; + } + EVIDENCE + .get() + .expect("worker error tracing installed before child spawn") + .as_ref() + .expect("worker error tracing installation succeeded") + .invalid_ready_body + .fetch_add(1, Ordering::Relaxed); + true +} + +/// Call after actor shutdown and the admin child stop, before final report. +pub async fn reconcile(shared: &Shared, evidence: &WorkerErrors) { + let drained = tokio::time::timeout( + crate::ci_timing::scaled_timeout(Duration::from_secs(5)), + async { + while evidence.child_readers.load(Ordering::Acquire) != 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }, + ) + .await + .is_ok(); + if !drained { + shared.incomplete("claiming-child error output did not drain", vec![]); + } + let missing = evidence.invalid_ready_body.load(Ordering::Relaxed); + if missing != 0 { + shared + .run + .add("worker_invalid_ready_body_events", missing as i64); + shared.violation( + "worker_ready_body_invalid", + format!("{missing} worker event(s) reported missing or ambiguous ready bodies; see WARN/ERROR output"), + vec![], + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + struct DisplayOnly; + + impl std::fmt::Display for DisplayOnly { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + formatter, + "queue storage {MISSING_READY} job 17 run_lease 2" + ) + } + } + + impl std::fmt::Debug for DisplayOnly { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("debug form omits the error") + } + } + + fn count_events(emit: impl FnOnce()) -> u64 { + let evidence = Arc::new(WorkerErrors::default()); + let subscriber = Registry::default().with(ErrorLayer { + evidence: evidence.clone(), + child: false, + }); + tracing::subscriber::with_default(subscriber, emit); + evidence.invalid_ready_body.load(Ordering::Relaxed) + } + + #[test] + fn display_error_field_is_counted_once_even_when_message_also_matches() { + assert_eq!( + count_events(|| { + tracing::warn!(error = %DisplayOnly, "{MISSING_READY}"); + }), + 1 + ); + } + + #[test] + fn string_error_and_message_only_forms_are_counted() { + assert_eq!( + count_events(|| { + tracing::error!(error = MISSING_READY, "rescue failed"); + tracing::warn!("queue storage {MISSING_READY} job 18 run_lease 3"); + }), + 2 + ); + } + + #[test] + fn ambiguous_ready_body_failure_is_also_gated() { + assert_eq!( + count_events(|| { + tracing::error!(error = AMBIGUOUS_READY, "hydrate failed"); + }), + 1 + ); + } + + #[test] + fn unrelated_chaos_warnings_and_payload_text_are_not_gated() { + assert_eq!( + count_events(|| { + tracing::warn!(error = "canceling statement due to user request", "chaos"); + tracing::error!(error = "lock timeout", "expected contention"); + tracing::warn!(payload = MISSING_READY, "a job argument"); + tracing::info!(error = MISSING_READY, "not a worker failure event"); + }), + 0 + ); + } +} diff --git a/awa/tests/persona_soak_test.rs b/awa/tests/persona_soak_test.rs new file mode 100644 index 000000000..1c33ba709 --- /dev/null +++ b/awa/tests/persona_soak_test.rs @@ -0,0 +1,43 @@ +//! Persona soak: several application-shaped workloads (Mia, Tom, Priya, Sam, Alex and +//! a chaos agent) run concurrently against one disposable database on the queue-storage +//! engine while an invariant oracle checks the system. Ignored by default. +//! +//! ```text +//! AWA_SOAK_DATABASE_URL=postgres://postgres:pw@localhost:5432/awa_soak_run \ +//! AWA_SOAK_SECONDS=900 \ +//! cargo test --release -p awa --test persona_soak_test -- --ignored --nocapture +//! ``` +//! +//! See `docs/test-plan.md` for the environment variables and `scripts/persona-soak.sh` +//! for the engine matrix. + +#![allow(dead_code)] +mod ci_timing; +mod persona_soak; + +use persona_soak::config::SoakConfig; + +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore = "drops and recreates an awa_soak* database; see the module docs"] +async fn persona_soak() { + let config = SoakConfig::from_env().unwrap_or_else(|error| panic!("{error}")); + let shared = persona_soak::run::run_soak(config) + .await + .unwrap_or_else(|error| panic!("soak setup failed: {error}")); + let violations = shared.violations.lock().expect("violations").len(); + assert_eq!( + violations, + 0, + "{violations} invariant violation(s); see {}", + shared.config.report_dir.display() + ); +} + +/// The claiming process of the claim-and-abandon persona. The harness runs this test binary +/// again with only this test selected and kills the process with SIGKILL; started any +/// other way it returns at once. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[ignore = "run by the claim-and-abandon persona as a separate process"] +async fn claim_abandon_child() { + persona_soak::personas::abandon::run_child().await; +} diff --git a/awa/tests/queue_storage_runtime_test.rs b/awa/tests/queue_storage_runtime_test.rs index 241ba2ee7..6c6105fd1 100644 --- a/awa/tests/queue_storage_runtime_test.rs +++ b/awa/tests/queue_storage_runtime_test.rs @@ -2245,12 +2245,9 @@ struct FinalizeBacklogJob { n: i64, } -/// A runtime's finished attempts stop heartbeating once they are handed to -/// finalization, so it must stop claiming when finalization falls behind. -/// Otherwise the backlog of finished, unheartbeated attempts grows with the -/// claim rate, outlives `heartbeat_staleness`, and rescue runs those jobs -/// again. Holding the first jobs' receipt-completion locks stalls the -/// completion batcher, standing in for a batcher that cannot keep up. +/// 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() { // The default completion pipeline: one shard of 512-job batches, one @@ -2375,7 +2372,7 @@ async fn test_queue_storage_runtime_stops_claiming_while_finalization_is_stalled assert!( started_while_stalled <= limit, "{started_while_stalled} jobs started while finalization was stalled; \ - at most {limit} may wait without a heartbeat or hold a worker" + at most {limit} may occupy finalizer slots or hold a worker" ); let runs = runs.lock().expect("runs"); assert_eq!(runs.len() as i64, JOBS); @@ -8500,6 +8497,160 @@ async fn test_queue_storage_promotion_leaves_a_contended_lane_deferred_and_in_or assert_eq!(claimed, expected); } +/// 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. +async fn hold_notify_commit_lock( + pool: &sqlx::PgPool, + commits: usize, +) -> tokio::task::JoinHandle<()> { + sqlx::query("CREATE TABLE IF NOT EXISTS notify_lock_holder (id bigint)") + .execute(pool) + .await + .expect("create notify holder table"); + let holder = PgPoolOptions::new() + .max_connections(1) + .connect_with( + pool.connect_options() + .as_ref() + .clone() + .options([("commit_delay", "100000"), ("commit_siblings", "0")]), + ) + .await + .expect("connect the notify holder"); + let holder_task = tokio::spawn(async move { + for _ in 0..commits { + sqlx::raw_sql( + "BEGIN; INSERT INTO notify_lock_holder VALUES (1); \ + SELECT pg_notify('awa_test_notify_holder', ''); COMMIT;", + ) + .execute(&holder) + .await + .expect("notify holder commit"); + } + holder.close().await; + }); + let deadline = Instant::now() + scaled_timeout(Duration::from_secs(5)); + 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)", + ) + .fetch_one(pool) + .await + .expect("read notify lock"); + if held { + return holder_task; + } + assert!( + Instant::now() < deadline, + "notify holder never took the lock" + ); + tokio::time::sleep(Duration::from_millis(1)).await; + } +} + +/// Promotion's lane transactions do not wait for the notify-queue lock: a lane +/// promotes under its short lock timeout while another notifying transaction +/// holds that lock, and a wakeup that cannot get it, under the connection's own +/// `lock_timeout`, does not undo or fail the promotion. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn test_queue_storage_promotion_does_not_wait_for_the_notify_lock_in_its_lanes() { + const PRIORITY: i16 = 2; + let (_db_guard, pool) = setup_pool(6).await; + let queue = "qs_promote_notify_commit"; + let store = create_store(&pool, "awa_qs_promote_notify_commit").await; + let lane = PromotionLane { + queue: queue.to_string(), + priority: PRIORITY, + }; + let enqueue_due = |id: i64| { + let pool = pool.clone(); + let store = &store; + async move { + enqueue_job( + &pool, + store, + &CompleteJob { id }, + InsertOpts { + queue: queue.to_string(), + priority: PRIORITY, + run_at: Some(Utc::now() - chrono::Duration::seconds(60)), + ..Default::default() + }, + ) + .await + } + }; + for id in 0..3 { + enqueue_due(id).await; + } + + let holder = hold_notify_commit_lock(&pool, 3).await; + let outcome = store + .promote_due_lanes( + &pool, + JobState::Scheduled, + 100, + Some(Duration::from_millis(50)), + &[], + ) + .await + .expect("promote while the notify lock is held"); + holder.await.expect("notify holder"); + assert_eq!(outcome.promoted, 3, "{outcome:?}"); + assert!(outcome.contended.is_empty(), "{outcome:?}"); + + let caller_lock_timeout_pool = PgPoolOptions::new() + .max_connections(1) + .connect_with( + pool.connect_options() + .as_ref() + .clone() + .options([("lock_timeout", "20ms")]), + ) + .await + .expect("connect with a session lock_timeout"); + enqueue_due(3).await; + let holder = hold_notify_commit_lock(&pool, 3).await; + let result = store + .promote_due_lanes( + &caller_lock_timeout_pool, + JobState::Scheduled, + 100, + Some(Duration::from_millis(50)), + &[], + ) + .await; + holder.await.expect("notify holder"); + caller_lock_timeout_pool.close().await; + let outcome = result.expect("a failed wakeup does not fail the promotion"); + assert_eq!(outcome.promoted, 1, "{outcome:?}"); + assert!(outcome.contended.is_empty(), "{outcome:?}"); + let deferred: i64 = sqlx::query_scalar(&format!( + "SELECT count(*) FROM {}.deferred_jobs WHERE queue = $1", + store.schema() + )) + .bind(queue) + .fetch_one(&pool) + .await + .expect("count deferred rows"); + assert_eq!(deferred, 0); + + let claimed = store + .claim_runtime_batch(&pool, queue, 10, Duration::from_secs(60)) + .await + .expect("claim promoted rows"); + assert_eq!(claimed.len(), 4); + assert!(store + .promote_due_lane(&pool, JobState::Scheduled, &lane, 100, None) + .await + .expect("nothing left to promote") + .contended + .is_empty()); +} + /// The maintenance leader must not wait on a lane whose enqueue lock never /// comes free. Two sessions take turns holding the hot lane's lock for /// `HOLD` each, so a lock request queued behind them waits up to two holds. @@ -8985,6 +9136,109 @@ async fn test_queue_storage_admin_queries_cover_running_and_failed_rows() { client.shutdown(Duration::from_secs(5)).await; } +#[tokio::test] +async fn test_row_local_receipt_admin_counts_follow_durable_dispositions() { + let (_db_guard, pool) = setup_pool(6).await; + let queue = "qs_row_receipt_dispositions"; + let schema = "awa_qs_row_receipt_dispositions"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.to_string(), + lease_claim_receipts: true, + ..Default::default() + }, + ) + .await; + + async fn counts(pool: &sqlx::PgPool, queue: &str) -> (i64, i64, i64) { + let overview = admin::queue_overview(pool, queue) + .await + .expect("read receipt disposition counts") + .expect("receipt queue exists"); + ( + overview.running, + overview.retryable, + overview.completed_last_hour, + ) + } + + store.enqueue_batch(&pool, queue, 2, 3).await.unwrap(); + let claimed = store + .claim_runtime_batch(&pool, queue, 3, Duration::from_secs(300)) + .await + .expect("claim three deadline-backed receipts"); + assert_eq!(claimed.len(), 3); + assert_eq!(lease_claim_count(&pool, &store).await, 3); + assert_eq!(lease_count(&pool, &store).await, 0); + assert_eq!(counts(&pool, queue).await, (3, 0, 0)); + + let completed = store + .complete_runtime_batch(&pool, &claimed[..1]) + .await + .unwrap(); + assert_eq!( + completed, + vec![(claimed[0].job.id, claimed[0].job.run_lease)] + ); + let null_fences: i64 = sqlx::query_scalar(&format!("SELECT count(*) FROM {schema}.lease_claims WHERE job_id = $1 AND run_lease = $2 AND closed_at IS NULL")) + .bind(claimed[0].job.id) + .bind(claimed[0].job.run_lease) + .fetch_one(&pool) + .await + .expect("completed row-local claim remains behind compact closure evidence"); + assert_eq!(null_fences, 1); + assert_eq!(lease_claim_closure_batch_count(&pool, &store).await, 1); + assert_eq!(counts(&pool, queue).await, (2, 0, 1)); + + let retried = store + .fail_retryable( + &pool, + claimed[1].job.id, + claimed[1].job.run_lease, + "retry", + None, + ) + .await + .unwrap() + .expect("second attempt moves to retryable"); + assert_eq!(retried.state, JobState::Retryable); + assert_eq!(counts(&pool, queue).await, (1, 1, 1)); + + let callback = store + .register_callback( + &pool, + claimed[2].job.id, + claimed[2].job.run_lease, + Duration::from_secs(300), + ) + .await + .unwrap(); + assert_eq!(lease_count(&pool, &store).await, 1); + assert_eq!( + counts(&pool, queue).await, + (1, 1, 1), + "materialization counts the attempt once" + ); + let finished = store + .complete_external(&pool, callback, None, Some(claimed[2].job.run_lease), false) + .await + .unwrap(); + assert_eq!(finished.state, JobState::Completed); + assert_eq!(counts(&pool, queue).await, (0, 1, 2)); + for (claim, expected) in claimed.iter().zip([ + JobState::Completed, + JobState::Retryable, + JobState::Completed, + ]) { + let job = store.load_job(&pool, claim.job.id).await.unwrap().unwrap(); + assert_eq!( + job.state, expected, + "point reads and aggregate counts agree" + ); + } +} + /// With lease-claim receipts on, a running job is an open receipt claim and /// `leases` stays empty, so every running count must read the claims too. #[tokio::test(flavor = "multi_thread", worker_threads = 4)] @@ -14763,6 +15017,17 @@ async fn test_queue_storage_queue_overview_counts_without_terminal_history() { assert_eq!(overview.total_queued, 2); assert_eq!(overview.running, 0); + let kind = admin::job_kind_overview(&pool, "complete_job") + .await + .expect("read kind overview after reusing the retained lane") + .expect("the completed and newly enqueued jobs have a kind"); + assert_eq!( + kind.job_count, + (completed.len() + waiting.len() - 1) as i64, + "retained completions and later jobs in the same lane count once" + ); + assert_eq!(kind.completed_last_hour, (completed.len() - 1) as i64); + // The dashboard's other two reads, over the same jobs. let counts = admin::state_counts(&pool) .await @@ -16994,6 +17259,3236 @@ async fn test_point_read_uses_live_ready_queue_after_ring_wrap() { ); } +/// Aging moves a waiting job by tombstoning its ready row and re-inserting the +/// same attempt on the next lane. In a quiet queue the new lane hands out the +/// same `lane_seq`, so a lease joined to its ready row by attempt alone matched +/// both rows and the admin reads counted the running job twice. +async fn moved_ready_row_is_joined_once_by_attempt_reads(receipts: bool, row_local: bool) { + let (_db_guard, pool) = setup_pool(6).await; + let queue = "qs_moved_ready_lease_reads"; + let schema = if row_local { + "awa_qs_moved_ready_row_claim_reads" + } else if receipts { + "awa_qs_moved_ready_compact_reads" + } else { + "awa_qs_moved_ready_lease_reads" + }; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + lease_claim_receipts: receipts, + ..unique_cancel_store_config(schema) + }, + ) + .await; + let job_id = enqueue_job( + &pool, + &store, + &unique_cancel_job("moved"), + InsertOpts { + priority: 2, + ..unique_cancel_opts(queue, None) + }, + ) + .await; + backdate_one_lane_ready_job_run_at( + &pool, + schema, + job_id, + Utc::now() - chrono::Duration::minutes(5), + ) + .await; + let aged = store + .age_waiting_priorities(&pool, Duration::from_secs(60), 10) + .await + .expect("age waiting priorities"); + assert_eq!(aged, vec![job_id]); + + let same_position: i64 = sqlx::query_scalar(&format!( + "SELECT count(DISTINCT priority)::bigint + FROM {schema}.ready_entries + WHERE job_id = $1 + GROUP BY ready_slot, ready_generation, enqueue_shard, lane_seq" + )) + .bind(job_id) + .fetch_one(&pool) + .await + .expect("ready rows for the moved job"); + assert_eq!( + same_position, 2, + "the source and destination rows must share a position for this test to mean anything" + ); + + let claimed = store + .claim_runtime_batch( + &pool, + queue, + 1, + if receipts && !row_local { + Duration::ZERO + } else { + Duration::from_secs(30) + }, + ) + .await + .expect("claim moved job"); + assert_eq!(claimed.len(), 1); + assert_eq!(claimed[0].job.id, job_id); + assert_eq!(claimed[0].claim.lease_claim_receipt, receipts); + assert_eq!( + claimed[0].claim.claim_batch_id.is_some(), + receipts && !row_local + ); + let legacy_rows: i64 = sqlx::query_scalar(&format!( + "SELECT count(*) FROM {schema}.lease_claims WHERE job_id = $1" + )) + .bind(job_id) + .fetch_one(&pool) + .await + .expect("verify the claim representation"); + assert_eq!(legacy_rows, i64::from(row_local)); + + // This is the production CTE used by current-job projections, not only + // the point-read candidate list. Every attempt must hydrate exactly once. + let projected: Vec = sqlx::query_scalar(&format!( + "{} SELECT job_id FROM current_jobs WHERE queue = $1 AND state = 'running'", + admin::queue_storage_current_jobs_cte(schema) + )) + .bind(queue) + .fetch_all(&pool) + .await + .expect("project current running jobs"); + assert_eq!(projected, vec![job_id], "hydrate only the live ready row"); + + let running = admin::list_jobs( + &pool, + &admin::ListJobsFilter { + state: Some(JobState::Running), + queue: Some(queue.to_string()), + ..Default::default() + }, + ) + .await + .expect("list running jobs"); + assert_eq!( + running.iter().map(|job| job.id).collect::>(), + vec![job_id], + "the moved running job is listed once" + ); + + let overview = admin::job_kind_overviews(&pool) + .await + .expect("job kind overviews"); + let kind = overview + .iter() + .find(|row| row.kind == UniqueCancelJob::kind()) + .expect("the moved job's kind is listed"); + assert_eq!(kind.job_count, 1, "the moved running job is counted once"); + + let loaded = store + .load_job(&pool, job_id) + .await + .expect("load moved job") + .expect("moved job exists"); + assert_eq!(loaded.state, JobState::Running); + assert_eq!(loaded.priority, 1, "reads the live destination row"); + assert_eq!( + loaded.metadata.get("_awa_original_priority"), + Some(&serde_json::json!(2)), + "payload comes from the aged destination, not its tombstoned source" + ); +} + +#[tokio::test] +async fn test_moved_ready_row_is_joined_once_by_lease_reads() { + moved_ready_row_is_joined_once_by_attempt_reads(false, false).await; +} + +#[tokio::test] +async fn test_moved_ready_row_is_joined_once_by_compact_receipt_reads() { + moved_ready_row_is_joined_once_by_attempt_reads(true, false).await; +} + +#[tokio::test] +async fn test_moved_ready_row_is_joined_once_by_row_local_receipt_reads() { + moved_ready_row_is_joined_once_by_attempt_reads(true, true).await; +} + +async fn receipt_admin_ids(pool: &sqlx::PgPool, filter: admin::ListJobsFilter) -> Vec { + admin::list_jobs(pool, &filter) + .await + .expect("list receipt-held jobs") + .into_iter() + .map(|job| job.id) + .collect() +} + +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) { + (false, _) => "awa_qs_admin_scope_leases", + (true, false) => "awa_qs_admin_scope_compact", + (true, true) => "awa_qs_admin_scope_row_claims", + }; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + lease_claim_receipts: receipts, + ..unique_cancel_store_config(schema) + }, + ) + .await; + let queue = "receipt_admin_scope"; + let other_queue = "receipt_admin_other"; + let first = enqueue_job( + &pool, + &store, + &CompleteJob { id: 1 }, + InsertOpts { + queue: queue.into(), + tags: vec!["selected".into()], + ..Default::default() + }, + ) + .await; + let second = enqueue_job( + &pool, + &store, + &CompleteJob { id: 2 }, + InsertOpts { + queue: queue.into(), + tags: vec!["selected".into()], + ..Default::default() + }, + ) + .await; + let other_kind = enqueue_job( + &pool, + &store, + &DlqJob { id: 3 }, + InsertOpts { + queue: queue.into(), + tags: vec!["selected".into()], + ..Default::default() + }, + ) + .await; + let other = enqueue_job( + &pool, + &store, + &CompleteJob { id: 4 }, + InsertOpts { + queue: other_queue.into(), + tags: vec!["selected".into()], + ..Default::default() + }, + ) + .await; + let deadline = if receipts && !row_local { + Duration::ZERO + } else { + Duration::from_secs(30) + }; + let mut claims = store + .claim_runtime_batch(&pool, queue, 3, deadline) + .await + .unwrap(); + if receipts { + assert!(matches!( + store + .rotate_claims(&pool) + .await + .expect("rotate between queues"), + RotateOutcome::Rotated { slot: 1, .. } + )); + } + claims.extend( + store + .claim_runtime_batch(&pool, other_queue, 1, deadline) + .await + .unwrap(), + ); + assert_eq!(claims.len(), 4); + assert!(claims + .iter() + .all(|claim| claim.claim.lease_claim_receipt == receipts)); + assert!(claims + .iter() + .all(|claim| claim.claim.claim_batch_id.is_some() == (receipts && !row_local))); + let row_claims: i64 = + sqlx::query_scalar(&format!("SELECT count(*) FROM {schema}.lease_claims")) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(row_claims, if row_local { 4 } else { 0 }); + + let scoped = admin::ListJobsFilter { + state: Some(JobState::Running), + kind: Some(CompleteJob::kind().into()), + queue: Some(queue.into()), + tag: Some("selected".into()), + limit: Some(1), + ..Default::default() + }; + assert_eq!(receipt_admin_ids(&pool, scoped.clone()).await, vec![second]); + assert_eq!( + receipt_admin_ids( + &pool, + admin::ListJobsFilter { + before_id: Some(second), + ..scoped.clone() + } + ) + .await, + vec![first] + ); + assert!(receipt_admin_ids( + &pool, + admin::ListJobsFilter { + before_id: Some(first), + ..scoped.clone() + } + ) + .await + .is_empty()); + assert!(receipt_admin_ids( + &pool, + admin::ListJobsFilter { + tag: Some("absent".into()), + ..scoped.clone() + } + ) + .await + .is_empty()); + assert!(receipt_admin_ids( + &pool, + admin::ListJobsFilter { + state: Some(JobState::WaitingExternal), + ..scoped.clone() + } + ) + .await + .is_empty()); + assert_eq!( + receipt_admin_ids( + &pool, + admin::ListJobsFilter { + state: None, + limit: Some(10), + ..scoped.clone() + } + ) + .await, + vec![second, first] + ); + + let projected: Vec = sqlx::query_scalar(&format!( + "{} SELECT job_id FROM current_jobs WHERE state='running' ORDER BY job_id DESC", + admin::queue_storage_current_jobs_cte(schema) + )) + .fetch_all(&pool) + .await + .unwrap(); + assert_eq!(projected, vec![other, other_kind, second, first]); + let kinds = admin::job_kind_overviews(&pool).await.unwrap(); + let complete = kinds + .iter() + .find(|kind| kind.kind == CompleteJob::kind()) + .unwrap(); + assert_eq!((complete.job_count, complete.queue_count), (3, 2)); + let alternate = kinds + .iter() + .find(|kind| kind.kind == DlqJob::kind()) + .unwrap(); + assert_eq!((alternate.job_count, alternate.queue_count), (1, 1)); + + // Materialization keeps the same attempt running. The receipt leg must + // yield to the lease leg rather than report both representations. + let first_claim = claims.iter().find(|claim| claim.job.id == first).unwrap(); + store + .register_callback( + &pool, + first, + first_claim.job.run_lease, + Duration::from_secs(60), + ) + .await + .expect("materialize one running attempt"); + assert_eq!( + receipt_admin_ids( + &pool, + admin::ListJobsFilter { + limit: Some(10), + ..scoped.clone() + } + ) + .await, + vec![second, first] + ); + let running: i64 = sqlx::query_scalar(&format!( + "{} SELECT count(*) FROM current_jobs WHERE state='running'", + admin::queue_storage_current_jobs_cte(schema) + )) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(running, 4, "materialization does not duplicate the claim"); + + let closed = claims + .iter() + .find(|claim| claim.job.id == second) + .unwrap() + .clone(); + assert_eq!( + store + .complete_runtime_batch(&pool, &[closed]) + .await + .unwrap() + .len(), + 1 + ); + assert_eq!( + receipt_admin_ids( + &pool, + admin::ListJobsFilter { + limit: Some(10), + ..scoped + } + ) + .await, + vec![first], + "closed receipts leave running enumeration" + ); + let running: i64 = sqlx::query_scalar(&format!( + "{} SELECT count(*) FROM current_jobs WHERE state='running'", + admin::queue_storage_current_jobs_cte(schema) + )) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(running, 3); + let kinds = admin::job_kind_overviews(&pool).await.unwrap(); + let complete = kinds + .iter() + .find(|kind| kind.kind == CompleteJob::kind()) + .unwrap(); + assert_eq!( + (complete.job_count, complete.queue_count), + (3, 2), + "the completed job is counted once, not as both running and terminal" + ); +} + +#[tokio::test] +async fn test_receipt_admin_scope_compact() { + receipt_admin_reads_preserve_scope_and_disposition(true, false).await; +} + +#[tokio::test] +async fn test_receipt_admin_scope_row_local() { + receipt_admin_reads_preserve_scope_and_disposition(true, true).await; +} + +#[tokio::test] +async fn test_receipt_admin_scope_materialized_control() { + receipt_admin_reads_preserve_scope_and_disposition(false, false).await; +} + +async fn claim_aged_rescue_body_job( + schema: &str, + receipts: bool, + deadline: Duration, + max_attempts: i16, +) -> ( + TestDbGuard, + sqlx::PgPool, + QueueStorage, + awa::model::ClaimedRuntimeJob, +) { + let (guard, pool) = setup_pool(6).await; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.into(), + lease_claim_receipts: receipts, + queue_stripe_count: 1, + ..Default::default() + }, + ) + .await; + let queue = "qs_aged_rescue_body"; + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9027 }, + InsertOpts { + queue: queue.into(), + priority: 4, + max_attempts, + metadata: serde_json::json!({"retained": "aged-rescue-body"}), + ..Default::default() + }, + ) + .await; + // Control elapsed queue age without waiting five minutes; all claim, + // rescue, terminal and retry transitions below use their real APIs. + backdate_one_lane_ready_job_run_at( + &pool, + schema, + job_id, + Utc::now() - chrono::Duration::minutes(5), + ) + .await; + let mut claimed = store + .claim_runtime_batch_with_aging(&pool, queue, 1, deadline, Duration::from_secs(60)) + .await + .expect("claim aged job"); + assert_eq!(claimed.len(), 1); + let original = claimed.remove(0); + assert_eq!(original.job.id, job_id); + assert_eq!(original.claim.priority, 4, "physical ready lane"); + assert_eq!(original.job.priority, 1, "aged effective priority"); + (guard, pool, store, original) +} + +#[derive(Clone, Copy)] +enum AgedFailedConsumer { + Retry, + Carry, + View, +} + +async fn aged_rescued_failure_keeps_body( + schema: &str, + receipts: bool, + deadline: Duration, + consumer: AgedFailedConsumer, +) { + let (_guard, pool, store, original) = + claim_aged_rescue_body_job(schema, receipts, deadline, 1).await; + let job_id = original.job.id; + // Zero staleness makes this committed attempt eligible without sleeps + // or rewriting receipt timestamps. No worker is running in this test. + let rescued = store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .expect("rescue exhausted aged attempt"); + assert_eq!(rescued.len(), 1); + assert_eq!(rescued[0].state, JobState::Failed); + assert_eq!(rescued[0].args, original.job.args); + assert_eq!(rescued[0].max_attempts, 1); + let raw_body = done_body_columns(&pool, &store, job_id).await; + assert!( + raw_body.0.is_none() && raw_body.1.is_none(), + "ready-backed failure stays narrow" + ); + + if matches!(consumer, AgedFailedConsumer::Carry) { + assert!(matches!( + store.rotate(&pool).await.expect("rotate"), + RotateOutcome::Rotated { slot: 1, .. } + )); + let prune_deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(5)); + loop { + match store + .prune_oldest(&pool, Duration::from_secs(3600)) + .await + .expect("carry failed row") + { + PruneOutcome::Pruned { + slot: 0, + carried_failed_rows: 1, + } => break, + PruneOutcome::Blocked { slot: 0 } if Instant::now() < prune_deadline => { + tokio::time::sleep(Duration::from_millis(20)).await; + } + other => panic!("unexpected carry prune outcome: {other:?}"), + } + } + let carried = done_body_columns(&pool, &store, job_id).await; + assert_eq!( + carried.0, + Some(original.job.args.clone()), + "carry must hydrate before truncating ready backing" + ); + assert_eq!(carried.1, Some(original.job.max_attempts)); + assert_eq!(carried.3, Some(original.job.created_at)); + let ready_count: i64 = sqlx::query_scalar(&format!( + "SELECT count(*) FROM {schema}.ready_entries WHERE job_id = $1" + )) + .bind(job_id) + .fetch_one(&pool) + .await + .expect("count pruned ready body"); + assert_eq!( + ready_count, 0, + "carried body must survive without ready backing" + ); + } + + // This is the installed SQL view, independent of the Rust hydration helper. + let view = terminal_view_body_columns(&pool, &store, job_id).await; + if matches!(consumer, AgedFailedConsumer::View) { + assert_eq!( + view.0, original.job.args, + "terminal_jobs must hydrate the aged failure" + ); + assert_eq!(view.1, original.job.max_attempts); + } + let retained = store + .load_job(&pool, job_id) + .await + .expect("load retained failure") + .expect("retained failure"); + let retry = store + .retry_job(&pool, job_id) + .await + .expect("retry aged failure") + .expect("failed job retries"); + assert_eq!( + retry.args, original.job.args, + "manual retry must retain original arguments" + ); + assert_eq!( + retry.max_attempts, original.job.max_attempts, + "manual retry must not default to25" + ); + assert_eq!(retry.created_at, original.job.created_at); + assert_eq!(retry.metadata["retained"], "aged-rescue-body"); + assert_eq!(retry.state, JobState::Available); + assert_eq!(retry.attempt, 0); + assert_eq!(view.0, original.job.args); + assert_eq!(view.1, original.job.max_attempts); + assert_eq!(view.3, original.job.created_at); + assert_eq!(retained.args, original.job.args); + assert_eq!(retained.max_attempts, original.job.max_attempts); + let rerun = store + .claim_runtime_batch(&pool, &original.job.queue, 1, Duration::from_secs(300)) + .await + .expect("claim manual retry"); + assert_eq!(rerun.len(), 1); + assert_eq!(rerun[0].job.args, original.job.args); + assert_eq!(rerun[0].job.max_attempts, original.job.max_attempts); + assert_eq!(rerun[0].job.run_lease, original.job.run_lease + 1); +} + +#[tokio::test] +async fn test_aged_rescued_failure_compact_retry_preserves_body() { + aged_rescued_failure_keeps_body( + "awa_qs_aged_rescue_compact", + true, + Duration::ZERO, + AgedFailedConsumer::Retry, + ) + .await; +} + +#[tokio::test] +async fn test_aged_rescued_failure_deadline_retry_preserves_body() { + aged_rescued_failure_keeps_body( + "awa_qs_aged_rescue_deadline", + true, + Duration::from_secs(300), + AgedFailedConsumer::Retry, + ) + .await; +} + +#[tokio::test] +async fn test_aged_rescued_failure_legacy_retry_preserves_body() { + aged_rescued_failure_keeps_body( + "awa_qs_aged_rescue_legacy", + false, + Duration::from_secs(300), + AgedFailedConsumer::Retry, + ) + .await; +} + +#[tokio::test] +async fn test_aged_rescued_failure_carry_preserves_body() { + aged_rescued_failure_keeps_body( + "awa_qs_aged_rescue_carry", + true, + Duration::ZERO, + AgedFailedConsumer::Carry, + ) + .await; +} + +#[tokio::test] +async fn test_aged_rescued_failure_terminal_view_preserves_body() { + aged_rescued_failure_keeps_body( + "awa_qs_aged_rescue_view", + true, + Duration::ZERO, + AgedFailedConsumer::View, + ) + .await; +} + +#[tokio::test] +async fn test_aged_rescued_failure_expired_deadline_retry_preserves_body() { + let (_guard, pool, store, original) = claim_aged_rescue_body_job( + "awa_qs_aged_rescue_expired_deadline", + true, + Duration::from_micros(1), + 1, + ) + .await; + // The real positive deadline has expired by the time its claim commits. + // This exercises deadline rescue itself, separately from heartbeat rescue. + let rescued = store + .rescue_expired_deadlines(&pool) + .await + .expect("rescue expired aged attempt"); + assert_eq!(rescued.len(), 1); + assert_eq!(rescued[0].id, original.job.id); + assert_eq!(rescued[0].state, JobState::Failed); + let view = terminal_view_body_columns(&pool, &store, original.job.id).await; + let retry = store + .retry_job(&pool, original.job.id) + .await + .expect("retry expired aged failure") + .expect("failed job retries"); + assert_eq!(retry.args, original.job.args); + assert_eq!(retry.max_attempts, original.job.max_attempts); + assert_eq!(retry.created_at, original.job.created_at); + assert_eq!(view.0, original.job.args); + assert_eq!(view.1, original.job.max_attempts); +} + +async fn aged_rescue_keeps_effective_policy_priority(dlq: bool) { + let schema = if dlq { + "awa_qs_aged_rescue_dlq" + } else { + "awa_qs_aged_rescue_retryable" + }; + let max_attempts = if dlq { 1 } else { 3 }; + let (_guard, pool, store, original) = + claim_aged_rescue_body_job(schema, true, Duration::ZERO, max_attempts).await; + let rescued = store + .rescue_stale_heartbeats_with_dlq(&pool, Duration::ZERO, &|_| dlq) + .await + .expect("rescue aged policy job"); + assert_eq!(rescued.len(), 1); + assert_eq!( + rescued[0].priority, 1, + "deferred/DLQ policy keeps effective priority" + ); + assert_eq!(rescued[0].args, original.job.args); + assert_eq!(rescued[0].max_attempts, max_attempts); + let table = if dlq { "dlq_entries" } else { "deferred_jobs" }; + let stored: (i16, serde_json::Value, i16) = sqlx::query_as(&format!( + "SELECT priority, args, max_attempts FROM {schema}.{table} WHERE job_id = $1" + )) + .bind(original.job.id) + .fetch_one(&pool) + .await + .expect("load self-contained policy row"); + assert_eq!(stored, (1, original.job.args, max_attempts)); +} + +#[tokio::test] +async fn test_aged_rescued_failure_dlq_keeps_effective_priority() { + aged_rescue_keeps_effective_policy_priority(true).await; +} + +#[tokio::test] +async fn test_aged_rescued_failure_retryable_keeps_effective_priority() { + aged_rescue_keeps_effective_policy_priority(false).await; +} + +async fn moved_rescued_failure_keeps_live_lane( + initial_priority: i16, + runtime_aging: Duration, + consumer: AgedFailedConsumer, +) { + let (_guard, pool) = setup_pool(4).await; + let schema = "awa_qs_moved_rescue_identity"; + let queue = "qs_moved_rescue_identity"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.into(), + queue_stripe_count: 1, + ..Default::default() + }, + ) + .await; + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9028 }, + InsertOpts { + queue: queue.into(), + priority: initial_priority, + max_attempts: 1, + ..Default::default() + }, + ) + .await; + backdate_one_lane_ready_job_run_at( + &pool, + schema, + job_id, + Utc::now() - chrono::Duration::minutes(5), + ) + .await; + assert_eq!( + store + .age_waiting_priorities(&pool, Duration::from_secs(60), 1) + .await + .expect("move queued job"), + vec![job_id] + ); + let claimed = store + .claim_runtime_batch_with_aging(&pool, queue, 1, Duration::ZERO, runtime_aging) + .await + .expect("claim moved job"); + assert_eq!(claimed.len(), 1); + let physical_priority = initial_priority - 1; + assert_eq!(claimed[0].claim.priority, physical_priority); + assert_eq!(claimed[0].job.priority, 1); + // A physical reorder of the disposable fixture preserves all logical + // tuples while exposing both legal fetch orders of the retained body. + sqlx::raw_sql(&format!( + "CLUSTER {schema}.ready_entries_0 USING idx_{schema}_ready_0_lane_shard" + )) + .execute(&pool) + .await + .expect("order retained fixture rows by physical lane"); + let physical_order: Vec = sqlx::query_scalar(&format!( + "SELECT priority FROM {schema}.ready_entries_0 WHERE job_id = $1 ORDER BY ctid" + )) + .bind(job_id) + .fetch_all(&pool) + .await + .expect("verify fixture row order"); + assert_eq!(physical_order, vec![physical_priority, initial_priority]); + let rescued = store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .expect("rescue moved job"); + assert_eq!(rescued.len(), 1); + if matches!(consumer, AgedFailedConsumer::Carry) { + // Queued aging leaves its tombstoned source lane's cursor behind. + // A real empty claim consumes it before that queue slot can prune. + let remaining = store + .claim_runtime_batch(&pool, queue, 10, Duration::ZERO) + .await + .expect("advance tombstoned source lane"); + assert!( + remaining.is_empty(), + "closed moved attempt must not run again" + ); + assert!(matches!( + store.rotate(&pool).await.expect("rotate"), + RotateOutcome::Rotated { slot: 1, .. } + )); + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(5)); + loop { + match store + .prune_oldest(&pool, Duration::from_secs(3600)) + .await + .expect("carry moved failure") + { + PruneOutcome::Pruned { + slot: 0, + carried_failed_rows: 1, + } => break, + PruneOutcome::Blocked { slot: 0 } if Instant::now() < deadline => { + tokio::time::sleep(Duration::from_millis(20)).await + } + other => panic!("unexpected moved-failure prune outcome: {other:?}"), + } + } + let body = done_body_columns(&pool, &store, job_id).await; + assert_eq!(body.0, Some(claimed[0].job.args.clone())); + assert_eq!(body.1, Some(1)); + } + let terminal: (i16, serde_json::Value, i16) = sqlx::query_as(&format!( + "SELECT priority,args,max_attempts FROM {schema}.terminal_jobs WHERE job_id=$1" + )) + .bind(job_id) + .fetch_one(&pool) + .await + .expect("read moved terminal view"); + if matches!(consumer, AgedFailedConsumer::View) { + assert_eq!( + terminal.0, physical_priority, + "terminal view must retain the live lane" + ); + assert_eq!(terminal.1, claimed[0].job.args); + assert_eq!(terminal.2, 1); + } + let retried = store + .retry_job(&pool, job_id) + .await + .expect("retry moved failure") + .expect("failed job retries"); + assert_eq!( + retried.priority, claimed[0].claim.priority, + "retry must not select the tombstoned source lane" + ); + assert_eq!(retried.args, claimed[0].job.args); + assert_eq!(retried.max_attempts, 1); + assert_eq!(terminal.0, physical_priority); + assert_eq!(terminal.1, claimed[0].job.args); + assert_eq!( + retried.metadata.get("_awa_original_priority"), + Some(&serde_json::json!(initial_priority)) + ); +} + +#[tokio::test] +async fn test_moved_rescued_failure_retry_keeps_live_lane() { + moved_rescued_failure_keeps_live_lane(2, Duration::ZERO, AgedFailedConsumer::Retry).await; +} + +#[tokio::test] +async fn test_moved_rescued_failure_aged_retry_keeps_live_lane() { + moved_rescued_failure_keeps_live_lane(4, Duration::from_secs(60), AgedFailedConsumer::Retry) + .await; +} + +#[tokio::test] +async fn test_moved_rescued_failure_aged_carry_keeps_live_lane() { + moved_rescued_failure_keeps_live_lane(4, Duration::from_secs(60), AgedFailedConsumer::Carry) + .await; +} + +#[tokio::test] +async fn test_moved_rescued_failure_aged_view_keeps_live_lane() { + moved_rescued_failure_keeps_live_lane(4, Duration::from_secs(60), AgedFailedConsumer::View) + .await; +} + +async fn terminal_counts_by_priority( + pool: &sqlx::PgPool, + schema: &str, + queue: &str, +) -> Vec<(i16, i64)> { + sqlx::query_as(&format!( + r#" + SELECT priority, sum(n)::bigint + FROM ( + SELECT priority, live_terminal_count AS n + FROM {schema}.queue_terminal_live_counts WHERE queue = $1 + UNION ALL + SELECT priority, terminal_delta AS n + FROM {schema}.queue_terminal_count_deltas WHERE queue = $1 + ) AS counters + GROUP BY priority HAVING sum(n) <> 0 ORDER BY priority + "# + )) + .bind(queue) + .fetch_all(pool) + .await + .expect("read per-priority terminal accounting") +} + +#[derive(Clone, Copy, Debug)] +enum LegacyFailedConsumer { + Retry, + Carry, + SingleDlq, + BulkDlq, + Discard, +} + +async fn legacy_aged_failure_recovers_body(moved: bool, consumer: LegacyFailedConsumer) { + let (_guard, pool) = setup_pool(4).await; + let schema = "awa_qs_legacy_aged_body"; + let queue = "qs_legacy_aged_body"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.into(), + queue_stripe_count: 1, + ..Default::default() + }, + ) + .await; + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9029 }, + InsertOpts { + queue: queue.into(), + priority: 4, + max_attempts: 1, + metadata: serde_json::json!({"retained": "legacy-aged-body", "ready_only": true, "shared": "ready"}), + ..Default::default() + }, + ) + .await; + backdate_one_lane_ready_job_run_at( + &pool, + schema, + job_id, + Utc::now() - chrono::Duration::minutes(5), + ) + .await; + if moved { + assert_eq!( + store + .age_waiting_priorities(&pool, Duration::from_secs(60), 1) + .await + .expect("move queued job"), + vec![job_id] + ); + } + let claimed = store + .claim_runtime_batch_with_aging(&pool, queue, 1, Duration::ZERO, Duration::from_secs(60)) + .await + .expect("claim aged legacy-shape fixture"); + assert_eq!(claimed.len(), 1); + let original = &claimed[0].job; + let physical_priority = if moved { 3 } else { 4 }; + assert_eq!(claimed[0].claim.priority, physical_priority); + assert_eq!(original.priority, 1); + assert_eq!( + store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .expect("rescue aged fixture") + .len(), + 1 + ); + + // Reproduce the persisted shape of an older writer: its narrow failed + // row and positive terminal delta used the effective attempt priority. + // The ready body and all real claim/closure evidence remain intact. + assert_eq!( + sqlx::query(&format!( + "UPDATE {schema}.done_entries SET priority = $2 WHERE job_id = $1" + )) + .bind(job_id) + .bind(original.priority) + .execute(&pool) + .await + .expect("seed old terminal priority") + .rows_affected(), + 1 + ); + assert_eq!( + sqlx::query(&format!( + "UPDATE {schema}.queue_terminal_count_deltas SET priority = $2 WHERE queue = $1" + )) + .bind(queue) + .bind(original.priority) + .execute(&pool) + .await + .expect("seed matching old terminal delta") + .rows_affected(), + 1 + ); + // An old moved-body lookup could retain the source payload, losing + // live-lane provenance. Keep terminal errors/custom overrides authoritative. + let terminal_payload: serde_json::Value = sqlx::query_scalar(&format!( + r#"UPDATE {schema}.done_entries + SET payload = jsonb_set( + COALESCE(payload, '{{}}'::jsonb), '{{metadata}}', + (COALESCE(payload -> 'metadata', '{{}}'::jsonb) - '_awa_original_priority' - 'ready_only') + || '{{"terminal_only":true,"shared":"terminal"}}'::jsonb + ) || '{{"tags":["terminal-tag"],"progress":{{"terminal":true}}}}'::jsonb + WHERE job_id = $1 RETURNING payload"# + )) + .bind(job_id) + .fetch_one(&pool) + .await + .expect("seed old terminal payload without live provenance"); + assert!(terminal_payload["metadata"] + .get("_awa_original_priority") + .is_none()); + assert!(!terminal_payload["errors"].as_array().unwrap().is_empty()); + let body = done_body_columns(&pool, &store, job_id).await; + assert!( + body.0.is_none() && body.1.is_none(), + "old row requires its ready body" + ); + assert_eq!( + terminal_counts_by_priority(&pool, schema, queue).await, + vec![(1, 1)] + ); + assert_eq!(store.queue_counts(&pool, queue).await.unwrap().terminal, 1); + + if matches!(consumer, LegacyFailedConsumer::Carry) { + let remaining = store + .claim_runtime_batch(&pool, queue, 10, Duration::ZERO) + .await + .expect("advance any tombstoned source cursor"); + assert!(remaining.is_empty()); + assert!(matches!( + store.rotate(&pool).await.expect("rotate"), + RotateOutcome::Rotated { slot: 1, .. } + )); + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(5)); + loop { + match store + .prune_oldest(&pool, Duration::from_secs(3600)) + .await + .expect("carry old failed row") + { + PruneOutcome::Pruned { + slot: 0, + carried_failed_rows: 1, + } => break, + PruneOutcome::Blocked { slot: 0 } if Instant::now() < deadline => { + tokio::time::sleep(Duration::from_millis(20)).await; + } + other => panic!("unexpected legacy carry outcome: {other:?}"), + } + } + let body = done_body_columns(&pool, &store, job_id).await; + assert_eq!(body.0, Some(original.args.clone())); + assert_eq!(body.1, Some(original.max_attempts)); + let ready_count: i64 = sqlx::query_scalar(&format!( + "SELECT count(*) FROM {schema}.ready_entries WHERE job_id = $1" + )) + .bind(job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!( + ready_count, 0, + "carry must preserve the body before its backing is pruned" + ); + assert_eq!( + terminal_counts_by_priority(&pool, schema, queue).await, + vec![(physical_priority, 1)], + "carried delta must match its replacement lane" + ); + assert_eq!(store.queue_counts(&pool, queue).await.unwrap().terminal, 1); + } + let retry = match consumer { + LegacyFailedConsumer::Retry | LegacyFailedConsumer::Carry => store + .retry_job(&pool, job_id) + .await + .expect("retry legacy failure") + .expect("legacy failed row retries"), + LegacyFailedConsumer::SingleDlq | LegacyFailedConsumer::BulkDlq => { + if matches!(consumer, LegacyFailedConsumer::SingleDlq) { + let moved = store + .move_failed_to_dlq(&pool, job_id, "manual-legacy-test") + .await + .expect("move old failed row to DLQ") + .expect("one row moves"); + assert_eq!(moved.id, job_id); + assert_eq!(moved.args, original.args); + assert_eq!(moved.max_attempts, original.max_attempts); + } else { + assert_eq!( + store + .bulk_move_failed_to_dlq( + &pool, + Some(CompleteJob::kind()), + Some(queue), + "manual-legacy-test", + ) + .await + .expect("bulk move old failed row to DLQ"), + 1, + "one terminal row must not multiply through moved ready bodies" + ); + } + let dlq_rows: Vec<(i16, serde_json::Value, i16, serde_json::Value)> = + sqlx::query_as(&format!( + "SELECT priority,args,max_attempts,COALESCE(payload,'{{}}'::jsonb) FROM {schema}.dlq_entries WHERE job_id=$1" + )).bind(job_id).fetch_all(&pool).await.expect("read actual stored DLQ row"); + assert_eq!(dlq_rows.len(), 1); + assert_eq!(dlq_rows[0].0, physical_priority); + assert_eq!(dlq_rows[0].1, original.args); + assert_eq!(dlq_rows[0].2, original.max_attempts); + assert_eq!(dlq_rows[0].3["metadata"]["_awa_original_priority"], 4); + assert_eq!(dlq_rows[0].3["metadata"]["shared"], "terminal"); + assert_eq!(dlq_rows[0].3["errors"], terminal_payload["errors"]); + assert_eq!(done_entries_count(&pool, schema, queue).await, 0); + assert!( + terminal_counts_by_priority(&pool, schema, queue) + .await + .is_empty(), + "DLQ move must decrement the stored priority-1 terminal key" + ); + assert_eq!(store.queue_counts(&pool, queue).await.unwrap().terminal, 0); + store + .retry_from_dlq(&pool, job_id, &awa::model::RetryFromDlqOpts::default()) + .await + .expect("retry recovered DLQ row") + .expect("DLQ row retries") + } + LegacyFailedConsumer::Discard => { + assert_eq!( + store + .discard_failed_by_kind(&pool, CompleteJob::kind()) + .await + .expect("discard old failed row"), + 1, + "one terminal row must not multiply through moved ready bodies" + ); + assert_eq!(done_entries_count(&pool, schema, queue).await, 0); + let dlq_count: i64 = sqlx::query_scalar(&format!( + "SELECT count(*) FROM {schema}.dlq_entries WHERE job_id=$1" + )) + .bind(job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(dlq_count, 0); + assert!( + terminal_counts_by_priority(&pool, schema, queue) + .await + .is_empty(), + "discard must decrement exactly its stored terminal lane" + ); + assert_eq!(store.queue_counts(&pool, queue).await.unwrap().terminal, 0); + return; + } + }; + assert_eq!(retry.args, original.args); + assert_eq!(retry.max_attempts, original.max_attempts); + assert_eq!(retry.created_at, original.created_at); + assert_eq!( + retry.priority, physical_priority, + "manual retry restarts aging from its live lane" + ); + assert_eq!(retry.metadata["retained"], "legacy-aged-body"); + assert_eq!( + retry.metadata["ready_only"], true, + "recover ready metadata defaults" + ); + assert_eq!(retry.metadata["shared"], "terminal", "stored metadata wins"); + assert_eq!(retry.metadata["terminal_only"], true); + if moved { + assert_eq!(retry.metadata["_awa_original_priority"], 4); + } + assert_eq!(retry.tags, vec!["terminal-tag"]); + assert_eq!( + serde_json::to_value(&retry.errors).unwrap(), + terminal_payload["errors"] + ); + let expected_progress = if matches!( + consumer, + LegacyFailedConsumer::SingleDlq | LegacyFailedConsumer::BulkDlq + ) { + None // Existing DLQ retry clears progress before making the job ready. + } else { + Some(serde_json::json!({"terminal": true})) + }; + assert_eq!(retry.progress, expected_progress); + assert_eq!(retry.attempt, 0); + assert_eq!(retry.state, JobState::Available); + assert!( + terminal_counts_by_priority(&pool, schema, queue) + .await + .is_empty(), + "remove the stored terminal lane, not the recovered retry lane" + ); + assert_eq!(store.queue_counts(&pool, queue).await.unwrap().terminal, 0); +} + +#[tokio::test] +async fn test_legacy_aged_failure_retry_recovers_body() { + legacy_aged_failure_recovers_body(false, LegacyFailedConsumer::Retry).await; +} + +#[tokio::test] +async fn test_legacy_aged_moved_failure_retry_recovers_live_body() { + legacy_aged_failure_recovers_body(true, LegacyFailedConsumer::Retry).await; +} + +#[tokio::test] +async fn test_legacy_aged_moved_failure_carry_recovers_body_and_counters() { + legacy_aged_failure_recovers_body(true, LegacyFailedConsumer::Carry).await; +} + +#[tokio::test] +async fn test_legacy_aged_moved_terminal_dlq_and_discard_consumers() { + tokio::task::LocalSet::new() + .run_until(async { + let mut failures = Vec::new(); + for consumer in [ + LegacyFailedConsumer::SingleDlq, + LegacyFailedConsumer::BulkDlq, + LegacyFailedConsumer::Discard, + ] { + // Exercise every parameter even after another panics. These + // database futures do not require Send on a local task. + if let Err(error) = + tokio::task::spawn_local(legacy_aged_failure_recovers_body(true, consumer)) + .await + { + failures.push(format!("{consumer:?}: {error}")); + } + } + assert!( + failures.is_empty(), + "terminal consumer failures: {failures:?}" + ); + }) + .await; +} + +#[tokio::test] +async fn test_legacy_narrow_available_cancel_retry_keeps_tombstoned_body() { + let (_guard, pool) = setup_pool(4).await; + let schema = "awa_qs_legacy_cancel_body"; + let queue = "qs_legacy_cancel_body"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.into(), + queue_stripe_count: 1, + ..Default::default() + }, + ) + .await; + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9030 }, + InsertOpts { + queue: queue.into(), + priority: 3, + max_attempts: 7, + metadata: serde_json::json!({"retained": "legacy-cancel-body"}), + ..Default::default() + }, + ) + .await; + let cancelled = store + .cancel_job(&pool, job_id) + .await + .expect("cancel available job") + .expect("available job cancels"); + assert_eq!(cancelled.state, JobState::Cancelled); + let backing: (i64, i64, bool) = sqlx::query_as(&format!( + r#" + SELECT done.run_lease, ready.run_lease, EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones tomb + WHERE tomb.ready_slot = ready.ready_slot + AND tomb.ready_generation = ready.ready_generation + AND tomb.queue = ready.queue AND tomb.priority = ready.priority + AND tomb.enqueue_shard = ready.enqueue_shard + AND tomb.lane_seq = ready.lane_seq AND tomb.job_id = ready.job_id + ) + FROM {schema}.done_entries done JOIN {schema}.ready_entries ready + ON ready.job_id = done.job_id AND ready.priority = done.priority + WHERE done.job_id = $1 + "# + )) + .bind(job_id) + .fetch_one(&pool) + .await + .expect("inspect cancelled body identity"); + assert_eq!( + backing.0, backing.1, + "unclaimed cancellation keeps the same run lease" + ); + assert!(backing.2, "cancellation tombstones its retained body"); + // Older narrow cancellations legitimately hydrate a same-lease body + // even though cancellation has already tombstoned that exact lane. + assert_eq!(sqlx::query(&format!( + "UPDATE {schema}.done_entries SET args=NULL, max_attempts=NULL, run_at=NULL, created_at=NULL WHERE job_id=$1" + )).bind(job_id).execute(&pool).await.unwrap().rows_affected(), 1); + assert_eq!( + terminal_counts_by_priority(&pool, schema, queue).await, + vec![(3, 1)] + ); + let retry = store + .retry_job(&pool, job_id) + .await + .expect("retry narrow cancellation") + .expect("cancelled job retries"); + assert_eq!(retry.args, cancelled.args); + assert_eq!(retry.max_attempts, 7); + assert_eq!(retry.priority, 3); + assert_eq!(retry.created_at, cancelled.created_at); + assert_eq!(retry.metadata["retained"], "legacy-cancel-body"); + assert!(terminal_counts_by_priority(&pool, schema, queue) + .await + .is_empty()); + assert_eq!(store.queue_counts(&pool, queue).await.unwrap().terminal, 0); +} + +#[tokio::test] +async fn test_receipt_list_refills_pages_past_closed_and_other_kind_claims() { + let (_guard, pool) = setup_pool(6).await; + let queue = "receipt_candidate_refill"; + let store = create_store_with_config( + &pool, + unique_cancel_store_config("awa_qs_receipt_candidate_refill"), + ) + .await; + let mut targets = Vec::new(); + for id in 0..3 { + targets.push( + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.into(), + tags: vec!["selected".into()], + ..Default::default() + }, + ) + .await, + ); + } + let targets_claimed = store + .claim_runtime_batch(&pool, queue, 3, Duration::ZERO) + .await + .expect("claim older targets"); + assert_eq!(targets_claimed.len(), 3); + let first_target = targets_claimed + .iter() + .find(|claim| claim.job.id == targets[0]) + .expect("first target was claimed"); + // The same job appears in both the lease and raw-receipt candidate sources. + store + .register_callback( + &pool, + targets[0], + first_target.job.run_lease, + Duration::from_secs(60), + ) + .await + .expect("materialize one target"); + + let mut all_ids = targets.clone(); + for id in 10..22 { + all_ids.push( + enqueue_job( + &pool, + &store, + &CompleteJob { id }, + InsertOpts { + queue: queue.into(), + ..Default::default() + }, + ) + .await, + ); + } + let finished = store + .claim_runtime_batch(&pool, queue, 12, Duration::ZERO) + .await + .expect("claim newer completed candidates"); + assert_eq!(finished.len(), 12); + assert_eq!( + store + .complete_runtime_batch(&pool, &finished) + .await + .unwrap() + .len(), + 12 + ); + for id in 30..42 { + all_ids.push( + enqueue_job( + &pool, + &store, + &DlqJob { id }, + InsertOpts { + queue: queue.into(), + ..Default::default() + }, + ) + .await, + ); + } + assert_eq!( + store + .claim_runtime_batch(&pool, queue, 12, Duration::ZERO) + .await + .unwrap() + .len(), + 12 + ); + + let filter = admin::ListJobsFilter { + state: Some(JobState::Running), + kind: Some(CompleteJob::kind().into()), + queue: Some(queue.into()), + limit: Some(2), + ..Default::default() + }; + targets.sort_unstable_by(|a, b| b.cmp(a)); + let first = receipt_admin_ids(&pool, filter.clone()).await; + assert_eq!(first, targets[..2]); + let second = receipt_admin_ids( + &pool, + admin::ListJobsFilter { + before_id: first.last().copied(), + ..filter.clone() + }, + ) + .await; + assert_eq!(second, targets[2..]); + assert_eq!( + receipt_admin_ids( + &pool, + admin::ListJobsFilter { + tag: Some("selected".into()), + limit: Some(100), + ..filter + }, + ) + .await, + targets + ); + all_ids.sort_unstable_by(|a, b| b.cmp(a)); + assert_eq!( + receipt_admin_ids( + &pool, + admin::ListJobsFilter { + queue: Some(queue.into()), + limit: Some(100), + ..Default::default() + }, + ) + .await, + all_ids, + "closed receipts and materialized claims must not duplicate results" + ); +} + +struct PointReadControl { + _guard: TestDbGuard, + pool: sqlx::PgPool, + reads: sqlx::PgPool, + store: Arc, + original: JobRow, + reader_pid: i32, +} + +impl PointReadControl { + async fn new(schema: &str) -> Self { + let (guard, pool) = setup_pool(6).await; + let store = Arc::new( + create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.into(), + ..Default::default() + }, + ) + .await, + ); + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9030 }, + InsertOpts { + queue: "qs_point_read_control".into(), + max_attempts: 1, + ..Default::default() + }, + ) + .await; + let reads = PgPoolOptions::new() + .max_connections(1) + .connect_with( + (*pool.connect_options()) + .clone() + .application_name("awa_point_read_control"), + ) + .await + .expect("point reader pool"); + let original = store + .load_job(&reads, job_id) + .await + .expect("warm cached point queries") + .expect("available job"); + let reader_pid = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&reads) + .await + .expect("point reader identity"); + Self { + _guard: guard, + pool, + reads, + store, + original, + reader_pid, + } + } + + async fn block_deferred(&self) -> sqlx::Transaction<'static, sqlx::Postgres> { + let mut locker = self.pool.begin().await.expect("point read blocker"); + sqlx::query(&format!( + "LOCK TABLE {}.deferred_jobs IN ACCESS EXCLUSIVE MODE", + self.store.schema() + )) + .execute(locker.as_mut()) + .await + .expect("hold point reader between ready and deferred legs"); + locker + } + + fn start_read(&self) -> tokio::task::JoinHandle, AwaError>> { + let store = self.store.clone(); + let reads = self.reads.clone(); + let job_id = self.original.id; + tokio::spawn(async move { store.load_job(&reads, job_id).await }) + } + + async fn wait_until_blocked(&self) { + let deadline = Instant::now() + ci_timing::scaled_timeout(Duration::from_secs(5)); + loop { + let waiting: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_locks WHERE pid=$1 AND NOT granted)", + ) + .bind(self.reader_pid) + .fetch_one(&self.pool) + .await + .expect("observe point read barrier"); + if waiting { + return; + } + assert!( + Instant::now() < deadline, + "point read did not reach barrier" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } + } + + async fn assert_pool_has_no_read_transaction(&self) { + let settings: (String, bool) = sqlx::query_as( + "SELECT current_setting('transaction_isolation'), current_setting('transaction_read_only')::boolean", + ) + .fetch_one(&self.reads) + .await + .expect("reuse point reader after interruption"); + assert_eq!(settings, ("read committed".into(), false)); + sqlx::query("CREATE TEMP TABLE point_read_write_control (value int)") + .execute(&self.reads) + .await + .expect("next caller can write outside the read-only transaction"); + } +} + +#[tokio::test] +async fn test_point_read_releases_earlier_locks_but_keeps_its_snapshot() { + let control = PointReadControl::new("awa_qs_point_snapshot").await; + let mut locker = control.block_deferred().await; + let reading = control.start_read(); + control.wait_until_blocked().await; + let held: Vec = sqlx::query_scalar( + "SELECT relation.relname FROM pg_locks AS locks JOIN pg_class AS relation ON relation.oid=locks.relation JOIN pg_namespace AS namespace ON namespace.oid=relation.relnamespace WHERE locks.pid=$1 AND locks.granted AND namespace.nspname=$2 AND relation.relkind IN ('r','p') AND relation.relname LIKE 'ready_entries%' ORDER BY relation.relname", + ) + .bind(control.reader_pid) + .bind(control.store.schema()) + .fetch_all(&control.pool) + .await + .expect("inspect retained ready relation locks"); + assert_eq!( + held, + vec!["ready_entries_0"], + "retain the located body without pinning every ready partition" + ); + let schema = control.store.schema(); + sqlx::query(&format!( + "INSERT INTO {schema}.deferred_jobs (job_id,kind,queue,args,state,priority,attempt,run_lease,max_attempts,run_at,created_at,payload) SELECT job_id,kind,queue,args,'scheduled'::awa.job_state,priority,attempt,run_lease+1,max_attempts,clock_timestamp()+interval '1 hour',created_at,payload FROM {schema}.ready_entries WHERE job_id=$1" + )) + .bind(control.original.id) + .execute(locker.as_mut()) + .await + .expect("move the row after the reader's first leg"); + sqlx::query(&format!( + "DELETE FROM {schema}.ready_entries WHERE job_id=$1" + )) + .bind(control.original.id) + .execute(locker.as_mut()) + .await + .expect("remove old ready row in the same transaction"); + locker.commit().await.expect("commit concurrent move"); + let old = reading + .await + .expect("point read task") + .expect("point read") + .expect("snapshot retained the old row"); + assert_eq!(old.state, JobState::Available); + assert_eq!(old.run_lease, control.original.run_lease); + let new = control + .store + .load_job(&control.reads, control.original.id) + .await + .expect("reuse cached point statements") + .expect("next snapshot sees the successor"); + assert_eq!(new.state, JobState::Scheduled); + assert_eq!(new.run_lease, control.original.run_lease + 1); + assert_eq!(new.args, control.original.args); + assert_eq!(new.max_attempts, control.original.max_attempts); + control.reads.close().await; +} + +#[tokio::test] +async fn test_point_read_cancellation_rolls_back_its_outer_transaction() { + let control = PointReadControl::new("awa_qs_point_cancel").await; + let locker = control.block_deferred().await; + let reading = control.start_read(); + control.wait_until_blocked().await; + reading.abort(); + assert!(reading.await.expect_err("read was aborted").is_cancelled()); + locker.rollback().await.expect("release point read barrier"); + control.assert_pool_has_no_read_transaction().await; + control.reads.close().await; +} + +#[tokio::test] +async fn test_point_read_query_error_rolls_back_its_outer_transaction() { + let control = PointReadControl::new("awa_qs_point_error").await; + sqlx::query("SET lock_timeout='100ms'") + .execute(&control.reads) + .await + .expect("bound blocked point read"); + let locker = control.block_deferred().await; + let result = control + .store + .load_job(&control.reads, control.original.id) + .await; + assert!( + matches!(result, Err(AwaError::Database(ref error)) if error.as_database_error().is_some_and(|error| error.code().as_deref() == Some("55P03"))) + ); + locker.rollback().await.expect("release point read barrier"); + control.assert_pool_has_no_read_transaction().await; + control.reads.close().await; +} + +#[tokio::test] +async fn test_point_read_keeps_located_body_until_failed_carry_can_finish() { + let control = PointReadControl::new("awa_qs_point_carry").await; + let claimed = control + .store + .claim_runtime_batch(&control.pool, &control.original.queue, 1, Duration::ZERO) + .await + .expect("claim carry fixture"); + assert_eq!(claimed.len(), 1); + let failed = control + .store + .rescue_stale_heartbeats(&control.pool, Duration::ZERO) + .await + .expect("exhaust carry fixture"); + assert_eq!(failed.len(), 1); + assert_eq!(failed[0].state, JobState::Failed); + assert!(matches!( + control.store.rotate(&control.pool).await.expect("rotate"), + RotateOutcome::Rotated { slot: 1, .. } + )); + let locker = control.block_deferred().await; + let reading = control.start_read(); + control.wait_until_blocked().await; + assert!(matches!( + control + .store + .prune_oldest(&control.pool, Duration::from_secs(3600)) + .await + .expect("prune while located body is in use"), + PruneOutcome::Blocked { slot: 0 } + )); + locker.rollback().await.expect("release read barrier"); + let before = reading + .await + .expect("point read task") + .expect("read before carry") + .expect("retained failure"); + assert_eq!(before.state, JobState::Failed); + assert_eq!(before.args, control.original.args); + assert_eq!(before.max_attempts, 1); + assert_eq!(before.errors, failed[0].errors); + assert!(matches!( + control + .store + .prune_oldest(&control.pool, Duration::from_secs(3600)) + .await + .expect("carry after point read releases its body"), + PruneOutcome::Pruned { + slot: 0, + carried_failed_rows: 1 + } + )); + let after = control + .store + .load_job(&control.reads, control.original.id) + .await + .expect("read carried failure") + .expect("carried failure remains visible"); + assert_eq!(after.state, JobState::Failed); + assert_eq!(after.args, before.args); + assert_eq!(after.max_attempts, before.max_attempts); + assert_eq!(after.errors, before.errors); + control.reads.close().await; +} + +/// Commit the same installed claim function used by the runtime, deliberately +/// leaving its separate post-commit cursor update unsent. +async fn claim_before_aging_cursor_advance_in_tx( + receipts: bool, + deadline: Duration, + occupy_replacement_lane: bool, +) -> ( + TestDbGuard, + sqlx::PgPool, + QueueStorage, + RawReceiptClaimRow, + sqlx::Transaction<'static, sqlx::Postgres>, +) { + let (guard, pool) = setup_pool(6).await; + let schema = "awa_qs_claimed_aging_body"; + let queue = "qs_claimed_aging_body"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.into(), + lease_claim_receipts: receipts, + queue_stripe_count: 1, + ..Default::default() + }, + ) + .await; + let job_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9041 }, + InsertOpts { + queue: queue.into(), + priority: 4, + max_attempts: 1, + metadata: serde_json::json!({"retained":"claimed-before-aging"}), + ..Default::default() + }, + ) + .await; + backdate_one_lane_ready_job_run_at( + &pool, + schema, + job_id, + Utc::now() - chrono::Duration::minutes(5), + ) + .await; + if occupy_replacement_lane { + // Occupy the replacement lane's first position. Old physical aging must + // append this job at seq2, not accidentally leave a live seq1 body that + // matches the already committed claim's original reference. + enqueue_job( + &pool, + &store, + &CompleteJob { id: 9042 }, + InsertOpts { + queue: queue.into(), + priority: 3, + ..Default::default() + }, + ) + .await; + } + if !occupy_replacement_lane { + store + .prepare_queue(&pool, queue, 1) + .await + .expect("prepare empty replacement lane"); + } + let mut tx = pool.begin().await.expect("begin exact runtime claim stage"); + let claimed: Vec = sqlx::query_as(&format!( + "SELECT ready_slot,ready_generation,job_id,priority,attempt,run_lease,lane_seq,claim_slot + FROM {schema}.claim_ready_runtime($1,1,$2,60::float8)" + )) + .bind(queue) + .bind(deadline.as_secs_f64()) + .fetch_all(tx.as_mut()) + .await + .expect("claim old ready row"); + assert_eq!(claimed.len(), 1); + let original = claimed.into_iter().next().unwrap(); + assert_eq!(original.job_id, job_id); + assert_eq!( + original.priority, 1, + "claim-time effective aging remains enabled" + ); + assert_eq!( + claim_cursor_for(&pool, &store, queue, 4, 0).await, + original.lane_seq + ); + (guard, pool, store, original, tx) +} + +async fn claim_before_aging_cursor_advance( + receipts: bool, + deadline: Duration, +) -> (TestDbGuard, sqlx::PgPool, QueueStorage, RawReceiptClaimRow) { + let (guard, pool, store, claim, tx) = + claim_before_aging_cursor_advance_in_tx(receipts, deadline, true).await; + tx.commit() + .await + .expect("commit claim before advancing cursor"); + (guard, pool, store, claim) +} + +async fn assert_claimed_aging_rescue_body( + pool: &sqlx::PgPool, + store: &QueueStorage, + claim: &RawReceiptClaimRow, + self_contained_terminal: bool, +) { + let rescued = store + .rescue_stale_heartbeats(pool, Duration::ZERO) + .await + .expect("hydrate retained claimed body"); + assert_eq!(rescued.len(), 1); + assert_eq!(rescued[0].id, claim.job_id); + assert_eq!(rescued[0].state, JobState::Failed); + assert_eq!(rescued[0].args, serde_json::json!({"id":9041})); + assert_eq!(rescued[0].max_attempts, 1); + assert_eq!( + rescued[0].priority, 4, + "terminal body keeps its physical lane" + ); + assert_eq!(rescued[0].metadata["retained"], "claimed-before-aging"); + assert_eq!(rescued[0].metadata["_awa_original_priority"], 4); + let schema = store.schema(); + let terminal: (i64, bool) = sqlx::query_as(&format!( + "SELECT lane_seq,args IS NOT NULL FROM {schema}.done_entries WHERE job_id=$1" + )) + .bind(claim.job_id) + .fetch_one(pool) + .await + .expect("terminal backing identity"); + assert_eq!( + terminal, + ( + if self_contained_terminal { + -claim.job_id + } else { + claim.lane_seq + }, + self_contained_terminal, + ), + "only conflicting historical body recovery needs a self-contained terminal row" + ); + for _ in 0..8 { + let available = store + .claim_runtime_batch(pool, "qs_claimed_aging_body", 1, Duration::from_secs(300)) + .await + .expect("drain ready prefixes before explicit retry"); + assert!(available.iter().all(|row| row.job.id != claim.job_id)); + } + let retried = store + .retry_job(pool, claim.job_id) + .await + .expect("retry rescued body") + .expect("failed attempt retries"); + assert_eq!(retried.args, rescued[0].args); + assert_eq!(retried.max_attempts, 1); + assert_eq!(retried.priority, 4); + assert_eq!(retried.metadata["_awa_original_priority"], 4); + let mut claimed_retries = Vec::new(); + for _ in 0..8 { + claimed_retries.extend( + store + .claim_runtime_batch(pool, "qs_claimed_aging_body", 1, Duration::from_secs(300)) + .await + .expect("claim only the new attempt after explicit retry") + .into_iter() + .filter(|row| row.job.id == claim.job_id), + ); + } + assert_eq!(claimed_retries.len(), 1); + assert_eq!(claimed_retries[0].job.run_lease, claim.run_lease + 1); +} + +async fn aging_skips_committed_claim_before_cursor(receipts: bool, deadline: Duration) { + let (_guard, pool, store, claim) = claim_before_aging_cursor_advance(receipts, deadline).await; + let aged = store + .age_waiting_priorities(&pool, Duration::from_secs(60), 1) + .await + .expect("age with lagging claim cursor"); + assert!( + aged.is_empty(), + "a committed claim is not waiting, even before cursor advancement: {aged:?}" + ); + assert_eq!(ready_tombstone_count(&pool, &store).await, 0); + assert_claimed_aging_rescue_body(&pool, &store, &claim, false).await; +} + +async fn historical_tombstoned_claim_fixture( + receipts: bool, + deadline: Duration, + colliding_position: bool, +) -> (TestDbGuard, sqlx::PgPool, QueueStorage, RawReceiptClaimRow) { + let (guard, pool, store, claim, tx) = + claim_before_aging_cursor_advance_in_tx(receipts, deadline, !colliding_position).await; + tx.commit() + .await + .expect("commit claim before old physical aging"); + // Reproduce the retained source and live replacement written by old aging. + let schema = store.schema(); + let mut tx = pool.begin().await.expect("begin historical aging shape"); + sqlx::query(&format!( + "INSERT INTO {schema}.ready_tombstones + (ready_slot,ready_generation,queue,priority,enqueue_shard,lane_seq,job_id) + SELECT ready_slot,ready_generation,queue,priority,enqueue_shard,lane_seq,job_id + FROM {schema}.ready_entries WHERE job_id=$1" + )) + .bind(claim.job_id) + .execute(tx.as_mut()) + .await + .expect("historical source tombstone"); + sqlx::query(&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, + {schema}.reserve_enqueue_seq(queue,(priority-1)::smallint,enqueue_shard,1::bigint),enqueue_shard, + run_at,attempted_at,created_at,unique_key,unique_states, + jsonb_set(COALESCE(payload,'{{}}'::jsonb),'{{metadata}}', + COALESCE(payload->'metadata','{{}}'::jsonb) + || jsonb_build_object('_awa_original_priority',priority)) + FROM {schema}.ready_entries WHERE job_id=$1" + )).bind(claim.job_id).execute(tx.as_mut()).await.expect("historical replacement body"); + sqlx::query(&format!( + "INSERT INTO {schema}.ready_segments + (ready_slot,ready_generation,queue,priority,enqueue_shard,first_lane_seq,next_lane_seq,first_run_at) + SELECT ready_slot,ready_generation,queue,priority,enqueue_shard,lane_seq,lane_seq+1,run_at + FROM {schema}.ready_entries WHERE job_id=$1 AND priority=3" + )) + .bind(claim.job_id) + .execute(tx.as_mut()) + .await + .expect("publish historical replacement segment as physical aging does"); + tx.commit() + .await + .expect("commit historical source and replacement"); + let positions: Vec = sqlx::query_scalar(&format!( + "SELECT lane_seq FROM {schema}.ready_entries WHERE job_id=$1 ORDER BY priority DESC" + )) + .bind(claim.job_id) + .fetch_all(&pool) + .await + .expect("source/replacement positions"); + assert_eq!(positions, vec![1, if colliding_position { 1 } else { 2 }]); + let provenance: Vec<(i16, Option)> = sqlx::query_as(&format!( + "SELECT priority,(payload->'metadata'->>'_awa_original_priority')::bigint + FROM {schema}.ready_entries WHERE job_id=$1 ORDER BY priority DESC" + )) + .bind(claim.job_id) + .fetch_all(&pool) + .await + .expect("old writer provenance placement"); + assert_eq!( + provenance, + vec![(4, None), (3, Some(4))], + "old physical aging adds provenance only to its replacement body" + ); + (guard, pool, store, claim) +} + +async fn historical_tombstoned_claim_hydrates( + receipts: bool, + deadline: Duration, + colliding_position: bool, +) { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(receipts, deadline, colliding_position).await; + assert_claimed_aging_rescue_body(&pool, &store, &claim, colliding_position).await; +} + +#[tokio::test] +async fn test_physical_aging_skips_committed_compact_claim_before_cursor() { + aging_skips_committed_claim_before_cursor(true, Duration::ZERO).await; +} +#[tokio::test] +async fn test_physical_aging_skips_committed_deadline_claim_before_cursor() { + aging_skips_committed_claim_before_cursor(true, Duration::from_secs(300)).await; +} +#[tokio::test] +async fn test_physical_aging_skips_committed_legacy_claim_before_cursor() { + aging_skips_committed_claim_before_cursor(false, Duration::from_secs(300)).await; +} +#[tokio::test] +async fn test_historical_aged_tombstoned_compact_claim_rescues_body() { + historical_tombstoned_claim_hydrates(true, Duration::ZERO, false).await; +} +#[tokio::test] +async fn test_historical_aged_tombstoned_deadline_claim_rescues_body() { + historical_tombstoned_claim_hydrates(true, Duration::from_secs(300), false).await; +} +#[tokio::test] +async fn test_historical_aged_tombstoned_legacy_claim_rescues_body() { + historical_tombstoned_claim_hydrates(false, Duration::from_secs(300), false).await; +} +#[tokio::test] +async fn test_historical_colliding_live_replacement_compact_rescues_claimed_body() { + historical_tombstoned_claim_hydrates(true, Duration::ZERO, true).await; +} +#[tokio::test] +async fn test_historical_colliding_live_replacement_deadline_rescues_claimed_body() { + historical_tombstoned_claim_hydrates(true, Duration::from_secs(300), true).await; +} + +async fn moved_tombstoned_claim_lane( + from: i16, + to: i16, + receipts: bool, + reopen_as_legacy: bool, + discovered_wide: bool, +) { + let (_guard, pool) = setup_pool(6).await; + let schema = "awa_qs_moved_tomb_claim"; + let queue = "qs_moved_tomb_claim"; + let store = create_store_with_config( + &pool, + QueueStorageConfig { + schema: schema.into(), + lease_claim_receipts: receipts, + queue_slot_count: if discovered_wide { + 20 + } else { + QueueStorageConfig::default().queue_slot_count + }, + queue_stripe_count: 1, + ..Default::default() + }, + ) + .await; + if discovered_wide { + enqueue_job( + &pool, + &store, + &CompleteJob { id: 9044 }, + InsertOpts { + queue: "qs_wide_anchor".into(), + ..Default::default() + }, + ) + .await; + for _ in 0..16 { + assert!(matches!( + store.rotate(&pool).await.unwrap(), + RotateOutcome::Rotated { .. } + )); + } + } + let id = enqueue_job(&pool, &store, &CompleteJob { id: 9043 }, InsertOpts { + queue: queue.into(), priority: from, max_attempts: 1, + metadata: serde_json::json!({"retained":"moved-tomb-claim", "_awa_original_priority":from}), + ..Default::default() + }).await; + backdate_one_lane_ready_job_run_at( + &pool, + schema, + id, + Utc::now() - chrono::Duration::minutes(5), + ) + .await; + assert!(store + .set_priority(&pool, id, to) + .await + .expect("move before claim")); + let body_positions: Vec<(i16, i64)> = sqlx::query_as(&format!( + "SELECT priority,lane_seq FROM {schema}.ready_entries WHERE job_id=$1 ORDER BY priority" + )) + .bind(id) + .fetch_all(&pool) + .await + .expect("shared moved positions"); + assert_eq!(body_positions.len(), 2); + assert!(body_positions.iter().all(|row| row.1 == 1)); + let mut claims = Vec::new(); + for _ in 0..3 { + claims = sqlx::query_as::<_, RawReceiptClaimRow>(&format!( + "SELECT ready_slot,ready_generation,job_id,priority,attempt,run_lease,lane_seq,claim_slot + FROM {schema}.claim_ready_runtime($1,1,300::float8,60::float8)" + )).bind(queue).fetch_all(&pool).await.expect("claim moved body before cursor update"); + if !claims.is_empty() { + break; + } + // A higher-priority tombstoned source may first consume its spent + // cursor prefix; the actual target claim still omits cursor advance. + } + assert_eq!(claims.len(), 1); + assert_eq!(claims[0].job_id, id); + if discovered_wide { + assert_eq!(claims[0].ready_slot, 16); + } + assert_eq!( + claims[0].priority, 1, + "effective priority cannot identify either physical lane" + ); + let store = if reopen_as_legacy { + store + .register_callback(&pool, id, claims[0].run_lease, Duration::from_secs(60)) + .await + .expect("materialize receipt before reopening"); + assert_eq!(lease_count(&pool, &store).await, 1); + QueueStorage::new(QueueStorageConfig { + schema: schema.into(), + lease_claim_receipts: false, + queue_stripe_count: 1, + ..Default::default() + }) + .expect("reopen existing custom store in legacy mode") + } else if discovered_wide { + QueueStorage::new(QueueStorageConfig { + schema: schema.into(), + ..Default::default() + }) + .expect("open installed schema without its writer configuration") + } else { + store + }; + // The old source was tombstoned by the real move. Model old physical aging + // tombstoning the claimed source as well; both exact-ref bodies now remain. + tombstone_ready_job(&pool, &store, id).await; + assert_eq!(ready_tombstone_count(&pool, &store).await, 2); + if discovered_wide { + let loaded = store + .load_job(&pool, id) + .await + .expect("read retained body in installed slot") + .unwrap(); + assert_eq!(loaded.state, JobState::Running); + assert_eq!(loaded.args, serde_json::json!({"id":9043})); + assert_eq!(loaded.max_attempts, 1); + } + if !receipts { + let other_queue = "qs_moved_tomb_claim_other"; + let other_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9044 }, + 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 an unrelated job"); + assert_eq!(other_claims.len(), 1); + assert_eq!(lease_count(&pool, &store).await, 2); + + let rescued = store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .expect("an ambiguous attempt does 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, + 1, + "the ambiguous attempt's mutation rolls back and it stays running" + ); + return; + } + let emitted: Vec = sqlx::query_scalar(&format!( + "SELECT priority FROM {schema}.ready_claim_attempt_batches WHERE queue=$1" + )) + .bind(queue) + .fetch_all(&pool) + .await + .expect("physical emitted lane proof"); + assert_eq!(emitted, vec![to]); + let rescued = store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .expect("use emitted physical lane for retained tombstoned body"); + assert_eq!(rescued.len(), 1); + assert_eq!(rescued[0].priority, to); + assert_eq!(rescued[0].args, serde_json::json!({"id":9043})); + assert_eq!(rescued[0].max_attempts, 1); + assert_eq!(rescued[0].metadata["_awa_original_priority"], from); + let retried = store + .retry_job(&pool, id) + .await + .expect("retry selected moved body") + .unwrap(); + assert_eq!(retried.priority, to); + assert_eq!(retried.args, rescued[0].args); + assert_eq!(retried.max_attempts, 1); + assert_eq!(retried.metadata["_awa_original_priority"], from); +} + +#[tokio::test] +async fn test_historical_tombstoned_claim_uses_emitted_lane_low_to_high() { + moved_tombstoned_claim_lane(2, 4, true, false, false).await; +} +#[tokio::test] +async fn test_historical_tombstoned_claim_uses_emitted_lane_high_to_low() { + moved_tombstoned_claim_lane(4, 2, true, false, false).await; +} +#[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; +} + +#[tokio::test] +async fn test_schema_only_store_recovers_retained_body_in_installed_slot() { + moved_tombstoned_claim_lane(2, 4, true, false, true).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_legacy_claim_completes_body() { + let (_guard, pool, store, claim) = + claim_before_aging_cursor_advance(false, Duration::from_secs(300)).await; + tombstone_ready_job(&pool, &store, claim.job_id).await; + let updated = store + .complete_job_batch_by_id(&pool, &[(claim.job_id, claim.run_lease)]) + .await + .expect("complete an already tombstoned claimed body"); + assert_eq!(updated, vec![(claim.job_id, claim.run_lease)]); + let done = store + .load_job(&pool, claim.job_id) + .await + .expect("read completion") + .unwrap(); + assert_eq!(done.state, JobState::Completed); + assert_eq!(done.args, serde_json::json!({"id":9041})); + assert_eq!(done.max_attempts, 1); + assert_eq!(done.priority, 4); + assert_eq!(done.metadata["retained"], "claimed-before-aging"); + assert_eq!(done.metadata["_awa_original_priority"], 4); +} + +async fn tombstoned_claim_cancels_body(receipts: bool, deadline: Duration) { + let (_guard, pool, store, claim) = claim_before_aging_cursor_advance(receipts, deadline).await; + tombstone_ready_job(&pool, &store, claim.job_id).await; + let done = store + .cancel_running( + &pool, + claim.job_id, + claim.run_lease, + "cancel retained body", + None, + ) + .await + .expect("cancel an already tombstoned claimed body") + .unwrap(); + assert_eq!(done.state, JobState::Cancelled); + assert_eq!(done.args, serde_json::json!({"id":9041})); + assert_eq!(done.max_attempts, 1); + assert_eq!(done.priority, 4); + assert_eq!(done.metadata["retained"], "claimed-before-aging"); + assert_eq!(done.metadata["_awa_original_priority"], 4); +} +#[tokio::test] +async fn test_historical_tombstoned_compact_claim_cancels_body() { + tombstoned_claim_cancels_body(true, Duration::ZERO).await; +} +#[tokio::test] +async fn test_historical_tombstoned_deadline_claim_cancels_body() { + tombstoned_claim_cancels_body(true, Duration::from_secs(300)).await; +} +#[tokio::test] +async fn test_historical_tombstoned_legacy_claim_cancels_body() { + tombstoned_claim_cancels_body(false, Duration::from_secs(300)).await; +} + +#[tokio::test] +async fn test_physical_aging_waits_for_claim_head_then_uses_fresh_snapshot() { + let (_guard, pool, store, claim, mut claim_tx) = + claim_before_aging_cursor_advance_in_tx(true, Duration::ZERO, true).await; + let holder: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(claim_tx.as_mut()) + .await + .expect("actual claim holder PID"); + let aging_pool = PgPoolOptions::new() + .max_connections(1) + .connect_with( + (*pool.connect_options()) + .clone() + .application_name("awa_aging_head_wait_control"), + ) + .await + .expect("separate aging connection"); + let age = store.age_waiting_priorities(&aging_pool, Duration::from_secs(60), 1); + let release = async { + tokio::time::timeout(ci_timing::scaled_timeout(Duration::from_secs(5)), async { + loop { + let waiting: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_stat_activity + WHERE application_name='awa_aging_head_wait_control' + AND wait_event_type='Lock' AND $1=ANY(pg_blocking_pids(pid)))", + ) + .bind(holder) + .fetch_one(&pool) + .await + .expect("observe actual claim-head wait"); + if waiting { + break; + } + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("aging must wait for the actual uncommitted claim head"); + claim_tx + .commit() + .await + .expect("publish claim while aging waits"); + }; + let (aged, ()) = tokio::join!(age, release); + assert!( + aged.expect("age after committed claim").is_empty(), + "the statement after the head lock must see the committed claim" + ); + assert_eq!(ready_tombstone_count(&pool, &store).await, 0); + assert_claimed_aging_rescue_body(&pool, &store, &claim, false).await; + aging_pool.close().await; +} + +async fn aging_skips_claimed_candidates_before_limit(receipts: bool, deadline: Duration) { + let (_guard, pool, store, claim) = claim_before_aging_cursor_advance(receipts, deadline).await; + let waiting = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9044 }, + InsertOpts { + queue: "qs_claimed_aging_body".into(), + priority: 4, + ..Default::default() + }, + ) + .await; + backdate_one_lane_ready_job_run_at( + &pool, + store.schema(), + waiting, + Utc::now() - chrono::Duration::minutes(4), + ) + .await; + let aged = store + .age_waiting_priorities(&pool, Duration::from_secs(60), 1) + .await + .expect("age past the oldest recorded claim"); + assert_eq!( + aged, + vec![waiting], + "claimed rows cannot consume the candidate limit" + ); + let waiting = store.load_job(&pool, waiting).await.unwrap().unwrap(); + assert_eq!(waiting.priority, 3); + assert_eq!(waiting.metadata["_awa_original_priority"], 4); + assert_claimed_aging_rescue_body(&pool, &store, &claim, false).await; +} + +#[tokio::test] +async fn test_physical_aging_skips_claimed_candidates_before_limit_compact() { + aging_skips_claimed_candidates_before_limit(true, Duration::ZERO).await; +} +#[tokio::test] +async fn test_physical_aging_skips_claimed_candidates_before_limit_deadline() { + aging_skips_claimed_candidates_before_limit(true, Duration::from_secs(300)).await; +} +#[tokio::test] +async fn test_physical_aging_skips_claimed_candidates_before_limit_legacy() { + aging_skips_claimed_candidates_before_limit(false, Duration::from_secs(300)).await; +} + +fn assert_historical_active_body(job: &JobRow, claim: &RawReceiptClaimRow, state: JobState) { + assert_eq!(job.id, claim.job_id); + assert_eq!(job.state, state); + assert_eq!(job.run_lease, claim.run_lease); + assert_eq!(job.priority, claim.priority); + assert_eq!(job.attempt, claim.attempt); + assert_eq!(job.args, serde_json::json!({"id":9041})); + assert_eq!(job.max_attempts, 1); + assert_eq!(job.metadata["retained"], "claimed-before-aging"); + assert_eq!(job.metadata["_awa_original_priority"], 4); +} + +async fn historical_tombstoned_claim_point_read(receipts: bool, deadline: Duration) { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(receipts, deadline, false).await; + let loaded = store + .load_job(&pool, claim.job_id) + .await + .expect("read historically tombstoned active claim") + .expect("claimed job remains visible"); + assert_historical_active_body(&loaded, &claim, JobState::Running); + assert_claimed_aging_rescue_body(&pool, &store, &claim, false).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_compact_claim_point_read_is_running() { + historical_tombstoned_claim_point_read(true, Duration::ZERO).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_deadline_claim_point_read_is_running() { + historical_tombstoned_claim_point_read(true, Duration::from_secs(300)).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_legacy_lease_point_read_is_running() { + historical_tombstoned_claim_point_read(false, Duration::from_secs(300)).await; +} + +async fn historical_tombstoned_claim_callback_reads(receipts: bool, deadline: Duration) { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(receipts, deadline, false).await; + let callback_id = store + .register_callback( + &pool, + claim.job_id, + claim.run_lease, + Duration::from_secs(300), + ) + .await + .expect("register callback for retained claim"); + let schema = store.schema(); + let progress = serde_json::json!({"phase":"retained-body"}); + sqlx::query(&format!( + "INSERT INTO {schema}.attempt_state + (job_id,run_lease,progress,callback_filter,callback_on_complete,callback_on_fail,callback_transform,updated_at) + VALUES ($1,$2,$3,'payload.ok','payload.done','payload.fail','payload.value',clock_timestamp()) + ON CONFLICT (job_id,run_lease) DO UPDATE SET + progress=EXCLUDED.progress,callback_filter=EXCLUDED.callback_filter, + callback_on_complete=EXCLUDED.callback_on_complete,callback_on_fail=EXCLUDED.callback_on_fail, + callback_transform=EXCLUDED.callback_transform" + )).bind(claim.job_id).bind(claim.run_lease).bind(&progress) + .execute(&pool).await.expect("set mutable callback state"); + let mut tx = pool.begin().await.unwrap(); + let running = store + .load_active_lease_in_tx(&mut tx, claim.job_id, claim.run_lease) + .await + .expect("raw active lease lookup") + .expect("running lease still has retained body"); + assert_historical_active_body(&running, &claim, JobState::Running); + assert_eq!(running.callback_id, Some(callback_id)); + assert_eq!(running.progress, Some(progress.clone())); + assert!(store + .enter_callback_wait_in_tx(&mut tx, claim.job_id, claim.run_lease, callback_id) + .await + .expect("park retained claim")); + let parked = store + .load_active_lease_in_tx(&mut tx, claim.job_id, claim.run_lease) + .await + .expect("read post-park row in caller transaction") + .expect("parked body"); + assert_historical_active_body(&parked, &claim, JobState::WaitingExternal); + assert_eq!(parked.callback_id, Some(callback_id)); + assert_eq!(parked.callback_timeout_at, running.callback_timeout_at); + assert_eq!(parked.heartbeat_at, None); + assert_eq!(parked.deadline_at, None); + for for_update in [false, true] { + let callback = store + .callback_job_in_tx(&mut tx, callback_id, Some(claim.run_lease), for_update) + .await + .expect("lookup retained callback body") + .expect("callback exists"); + assert_historical_active_body(&callback, &claim, JobState::WaitingExternal); + assert_eq!(callback.callback_id, Some(callback_id)); + assert_eq!(callback.callback_timeout_at, parked.callback_timeout_at); + assert_eq!(callback.progress, Some(progress.clone())); + assert_eq!(callback.callback_filter.as_deref(), Some("payload.ok")); + assert_eq!( + callback.callback_on_complete.as_deref(), + Some("payload.done") + ); + assert_eq!(callback.callback_on_fail.as_deref(), Some("payload.fail")); + assert_eq!( + callback.callback_transform.as_deref(), + Some("payload.value") + ); + } + tx.commit().await.expect("commit parked lease"); + let loaded = store + .load_job(&pool, claim.job_id) + .await + .expect("point read parked lease") + .unwrap(); + assert_historical_active_body(&loaded, &claim, JobState::WaitingExternal); + let mut locked = pool.begin().await.unwrap(); + store + .callback_job_in_tx(&mut locked, callback_id, None, true) + .await + .expect("lock callback lease despite missing live body") + .unwrap(); + let mut contender = pool.begin().await.unwrap(); + let err = sqlx::query(&format!( + "SELECT job_id FROM {schema}.leases WHERE job_id=$1 AND run_lease=$2 FOR UPDATE NOWAIT" + )) + .bind(claim.job_id) + .bind(claim.run_lease) + .fetch_one(contender.as_mut()) + .await + .expect_err("callback lookup must retain its lease row lock"); + assert_eq!( + err.as_database_error() + .and_then(|err| err.code()) + .as_deref(), + Some("55P03") + ); + contender.rollback().await.unwrap(); + locked.rollback().await.unwrap(); + let resumed = store + .complete_external( + &pool, + callback_id, + Some(serde_json::json!({"ok":true})), + Some(claim.run_lease), + true, + ) + .await + .expect("resume callback and read retained body inside update transaction"); + assert_historical_active_body(&resumed, &claim, JobState::Running); + assert_eq!(resumed.callback_id, None); + assert_eq!(resumed.progress, Some(progress)); + assert_eq!( + resumed.metadata["_awa_callback_result"], + serde_json::json!({"ok":true}) + ); +} + +#[tokio::test] +async fn test_historical_tombstoned_compact_claim_callback_reads_and_resume() { + historical_tombstoned_claim_callback_reads(true, Duration::ZERO).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_deadline_claim_callback_reads_and_resume() { + historical_tombstoned_claim_callback_reads(true, Duration::from_secs(300)).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_legacy_lease_callback_reads_and_resume() { + historical_tombstoned_claim_callback_reads(false, Duration::from_secs(300)).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_colliding_receipt_point_reads_use_emitted_lane() { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(true, Duration::ZERO, true).await; + tombstone_ready_job(&pool, &store, claim.job_id).await; + let loaded = store + .load_job(&pool, claim.job_id) + .await + .expect("read emitted claimed lane") + .unwrap(); + assert_historical_active_body(&loaded, &claim, JobState::Running); + let callback_id = store + .register_callback( + &pool, + claim.job_id, + claim.run_lease, + Duration::from_secs(300), + ) + .await + .expect("materialize retained emitted claim"); + let mut tx = pool.begin().await.unwrap(); + let active = store + .load_active_lease_in_tx(&mut tx, claim.job_id, claim.run_lease) + .await + .expect("load retained emitted lease") + .unwrap(); + assert_historical_active_body(&active, &claim, JobState::Running); + let callback = store + .callback_job_in_tx(&mut tx, callback_id, None, true) + .await + .expect("callback resolves emitted lane") + .unwrap(); + assert_historical_active_body(&callback, &claim, JobState::Running); + tx.rollback().await.unwrap(); +} + +#[tokio::test] +async fn test_historical_tombstoned_colliding_legacy_point_reads_fail_ambiguous() { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(false, Duration::from_secs(300), true).await; + let callback_id = store + .register_callback( + &pool, + claim.job_id, + claim.run_lease, + Duration::from_secs(300), + ) + .await + .expect("register legacy callback"); + tombstone_ready_job(&pool, &store, claim.job_id).await; + let expected = format!( + "queue storage ready body ambiguous for deleted lease job {} run_lease {}", + claim.job_id, claim.run_lease + ); + let err = store + .load_job(&pool, claim.job_id) + .await + .expect_err("point read cannot choose legacy lane"); + assert!( + matches!(err, AwaError::Validation(ref message) if message == &expected), + "{err}" + ); + let mut tx = pool.begin().await.unwrap(); + let err = store + .load_active_lease_in_tx(&mut tx, claim.job_id, claim.run_lease) + .await + .expect_err("active lookup cannot choose legacy lane"); + assert!( + matches!(err, AwaError::Validation(ref message) if message == &expected), + "{err}" + ); + let err = store + .callback_job_in_tx(&mut tx, callback_id, None, true) + .await + .expect_err("callback lookup cannot choose legacy lane"); + assert!( + matches!(err, AwaError::Validation(ref message) if message == &expected), + "{err}" + ); + tx.rollback().await.unwrap(); + assert_eq!(lease_count(&pool, &store).await, 1); +} + +async fn historical_generic_cancel_live_replacement(deadline: Duration, colliding_position: bool) { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(true, deadline, colliding_position).await; + let schema = store.schema(); + let live_bodies_sql = format!( + "SELECT count(*)::bigint FROM {schema}.ready_entries AS ready + WHERE ready.job_id=$1 AND NOT EXISTS ( + SELECT 1 FROM {schema}.ready_tombstones AS tomb + WHERE tomb.ready_slot=ready.ready_slot AND tomb.ready_generation=ready.ready_generation + AND tomb.queue=ready.queue AND tomb.priority=ready.priority + AND tomb.enqueue_shard=ready.enqueue_shard AND tomb.lane_seq=ready.lane_seq + AND tomb.job_id=ready.job_id)" + ); + let live_before: i64 = sqlx::query_scalar(&live_bodies_sql) + .bind(claim.job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(live_before, 1); + let cancelled = store + .cancel_job(&pool, claim.job_id) + .await + .expect("cancel running receipt ahead of its historical live replacement") + .expect("claimed attempt is cancelled"); + assert_eq!(cancelled.id, claim.job_id); + assert_eq!(cancelled.state, JobState::Cancelled); + assert_eq!(cancelled.run_lease, claim.run_lease); + assert_eq!(cancelled.args, serde_json::json!({"id":9041})); + assert_eq!(cancelled.max_attempts, 1); + assert_eq!(cancelled.priority, 4); + assert_eq!(cancelled.metadata["retained"], "claimed-before-aging"); + assert_eq!(cancelled.metadata["_awa_original_priority"], 4); + let live_after: i64 = sqlx::query_scalar(&live_bodies_sql) + .bind(claim.job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(live_after, 0); + let loaded = store + .load_job(&pool, claim.job_id) + .await + .expect("read cancelled receipt") + .expect("cancelled receipt remains visible"); + assert_eq!(loaded.state, JobState::Cancelled); + assert_eq!(loaded.run_lease, claim.run_lease); + assert_eq!(loaded.args, cancelled.args); + assert_eq!(loaded.priority, 4); + assert_eq!(loaded.metadata["_awa_original_priority"], 4); + assert!(store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .expect("cancelled receipt is closed") + .is_empty()); + for _ in 0..8 { + let before_retry = store + .claim_runtime_batch(&pool, "qs_claimed_aging_body", 1, deadline) + .await + .expect("claim any unrelated ready job after cancel"); + assert!(before_retry + .iter() + .all(|entry| entry.job.id != claim.job_id)); + } + let retried = store + .retry_job(&pool, claim.job_id) + .await + .expect("retry generic receipt cancellation") + .expect("cancelled receipt retries"); + assert_eq!(retried.args, cancelled.args); + assert_eq!(retried.max_attempts, 1); + assert_eq!(retried.priority, 4); + assert_eq!(retried.metadata["_awa_original_priority"], 4); + let mut claimed_retry = Vec::new(); + for _ in 0..8 { + claimed_retry.extend( + store + .claim_runtime_batch(&pool, "qs_claimed_aging_body", 1, deadline) + .await + .expect("claim the explicit retry"), + ); + } + assert_eq!(claimed_retry.len(), 1); + assert_eq!(claimed_retry[0].job.id, claim.job_id); + assert_eq!(claimed_retry[0].job.run_lease, claim.run_lease + 1); + assert_eq!(claimed_retry[0].job.priority, 4); + assert!(store + .claim_runtime_batch(&pool, "qs_claimed_aging_body", 10, deadline) + .await + .expect("old historical copy cannot emit another claim") + .is_empty()); +} + +#[tokio::test] +async fn test_generic_cancel_historical_colliding_live_compact_replacement_keeps_claimed_lane() { + historical_generic_cancel_live_replacement(Duration::ZERO, true).await; +} + +#[tokio::test] +async fn test_generic_cancel_historical_live_deadline_replacement_keeps_claimed_lane() { + historical_generic_cancel_live_replacement(Duration::from_secs(300), false).await; +} + +async fn historical_recovery_head_contention_rolls_back(receipts: bool, deadline: Duration) { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(receipts, deadline, false).await; + let schema = store.schema(); + let mut blocker = pool.begin().await.unwrap(); + sqlx::query(&format!( + "SELECT 1 FROM {schema}.queue_claim_heads + WHERE queue='qs_claimed_aging_body' AND priority=3 AND enqueue_shard=0 FOR UPDATE" + )) + .execute(blocker.as_mut()) + .await + .unwrap(); + let result = tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(2)), + store.rescue_stale_heartbeats(&pool, Duration::ZERO), + ) + .await + .expect("contended recovery head is skipped without waiting") + .expect("busy historical attempt is deferred"); + assert!(result.is_empty(), "busy attempt remains open"); + let done: i64 = sqlx::query_scalar(&format!( + "SELECT count(*) FROM {schema}.done_entries WHERE job_id=$1" + )) + .bind(claim.job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(done, 0); + assert_eq!(ready_tombstone_count(&pool, &store).await, 1); + let active = store.load_job(&pool, claim.job_id).await.unwrap().unwrap(); + assert_historical_active_body(&active, &claim, JobState::Running); + blocker.rollback().await.unwrap(); + assert_claimed_aging_rescue_body(&pool, &store, &claim, false).await; +} + +#[tokio::test] +async fn test_historical_compact_recovery_head_contention_rolls_back() { + historical_recovery_head_contention_rolls_back(true, Duration::ZERO).await; +} + +#[tokio::test] +async fn test_historical_legacy_recovery_head_contention_rolls_back() { + historical_recovery_head_contention_rolls_back(false, Duration::from_secs(300)).await; +} + +async fn historical_tombstoned_claim_completes(receipts: bool, deadline: Duration) { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(receipts, deadline, false).await; + if receipts { + store + .register_callback( + &pool, + claim.job_id, + claim.run_lease, + Duration::from_secs(300), + ) + .await + .expect("materialize retained claim before completion"); + } + let completed = store + .complete_job_batch_by_id(&pool, &[(claim.job_id, claim.run_lease)]) + .await + .expect("complete retained body after old aging race"); + assert_eq!(completed, vec![(claim.job_id, claim.run_lease)]); + let loaded = store.load_job(&pool, claim.job_id).await.unwrap().unwrap(); + assert_eq!(loaded.state, JobState::Completed); + assert_eq!(loaded.run_lease, claim.run_lease); + assert_eq!(loaded.args, serde_json::json!({"id":9041})); + assert_eq!(loaded.max_attempts, 1); + assert_eq!(loaded.priority, 4); + assert_eq!(loaded.metadata["_awa_original_priority"], 4); + for _ in 0..8 { + let available = store + .claim_runtime_batch(&pool, "qs_claimed_aging_body", 10, deadline) + .await + .expect("old aging replacement cannot claim after hydrated completion"); + assert!(available.iter().all(|row| row.job.id != claim.job_id)); + } + assert!(store + .rescue_stale_heartbeats(&pool, Duration::ZERO) + .await + .unwrap() + .iter() + .all(|row| row.id != claim.job_id)); +} + +#[tokio::test] +async fn test_historical_tombstoned_compact_claim_completes_and_retires_copy() { + historical_tombstoned_claim_completes(true, Duration::ZERO).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_deadline_claim_completes_and_retires_copy() { + historical_tombstoned_claim_completes(true, Duration::from_secs(300)).await; +} + +#[tokio::test] +async fn test_historical_tombstoned_legacy_claim_completes_and_retires_copy() { + historical_tombstoned_claim_completes(false, Duration::from_secs(300)).await; +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +enum RetainedBatchTransition { + Heartbeat, + Deadline, + Callback, +} + +async fn retained_batch_transition( + pool: &sqlx::PgPool, + store: &QueueStorage, + transition: RetainedBatchTransition, +) -> Vec<(i64, i64)> { + tokio::time::timeout(ci_timing::scaled_timeout(Duration::from_secs(2)), async { + match transition { + RetainedBatchTransition::Heartbeat => store + .rescue_stale_heartbeats(pool, Duration::ZERO) + .await + .expect("rescue unblocked stale attempts") + .into_iter() + .map(|job| (job.id, job.run_lease)) + .collect(), + RetainedBatchTransition::Deadline => store + .rescue_expired_deadlines(pool) + .await + .expect("rescue unblocked expired attempts") + .into_iter() + .map(|job| (job.id, job.run_lease)) + .collect(), + RetainedBatchTransition::Callback => store + .rescue_expired_callbacks(pool) + .await + .expect("rescue unblocked expired callbacks") + .into_iter() + .map(|job| (job.id, job.run_lease)) + .collect(), + } + }) + .await + .expect("busy historical head cannot delay ordinary transitions") +} + +async fn retained_receipt_cursor( + pool: &sqlx::PgPool, + store: &QueueStorage, + slot: i32, +) -> serde_json::Value { + sqlx::query_scalar(&format!( + "SELECT jsonb_build_array(rescue_cursor_claimed_at, rescue_cursor_job_id, + rescue_cursor_run_lease, deadline_cursor_deadline_at, deadline_cursor_job_id, + deadline_cursor_run_lease) + FROM {}.claim_ring_slots WHERE slot=$1", + store.schema() + )) + .bind(slot) + .fetch_one(pool) + .await + .expect("read receipt rescue cursors") +} + +async fn retained_closed_receipts( + pool: &sqlx::PgPool, + store: &QueueStorage, + pair: (i64, i64), +) -> i64 { + let schema = store.schema(); + sqlx::query_scalar(&format!( + "SELECT + (SELECT count(*) FROM {schema}.lease_claim_closures WHERE job_id=$1 AND run_lease=$2) + + (SELECT count(*) FROM {schema}.lease_claim_closure_batches AS closures + WHERE EXISTS ( + SELECT 1 FROM {schema}.lease_claim_batches AS claims + CROSS JOIN LATERAL unnest(claims.job_ids,claims.run_leases,claims.receipt_ids) + AS members(job_id,run_lease,receipt_id) + WHERE members.job_id=$1 AND members.run_lease=$2 + AND closures.receipt_ranges @> members.receipt_id)) + + (SELECT count(*) FROM {schema}.lease_claims + WHERE job_id=$1 AND run_lease=$2 AND closed_at IS NOT NULL)" + )) + .bind(pair.0) + .bind(pair.1) + .fetch_one(pool) + .await + .expect("count exact receipt closure evidence") +} + +async fn historical_busy_attempt_does_not_block_batch( + transition: RetainedBatchTransition, + materialized: bool, +) { + let deadline = if transition == RetainedBatchTransition::Heartbeat { + Duration::ZERO + } else { + Duration::from_secs(300) + }; + let (_guard, pool, store, historical) = + historical_tombstoned_claim_fixture(true, deadline, false).await; + let queue = "qs_unblocked_batch"; + let ordinary_id = enqueue_job( + &pool, + &store, + &CompleteJob { id: 9042 }, + InsertOpts { + queue: queue.into(), + priority: 4, + max_attempts: 1, + ..Default::default() + }, + ) + .await; + let claimed = store + .claim_runtime_batch(&pool, queue, 1, deadline) + .await + .expect("claim ordinary batch member"); + assert_eq!(claimed.len(), 1); + assert_eq!(claimed[0].job.id, ordinary_id); + let historical_pair = (historical.job_id, historical.run_lease); + let ordinary_pair = (ordinary_id, claimed[0].job.run_lease); + let pairs = [historical_pair, ordinary_pair]; + if materialized || transition == RetainedBatchTransition::Callback { + for &(job_id, run_lease) in &pairs { + let callback = store + .register_callback(&pool, job_id, run_lease, Duration::from_secs(300)) + .await + .expect("materialize attempt"); + if transition == RetainedBatchTransition::Callback { + assert!(store + .enter_callback_wait(&pool, job_id, run_lease, callback) + .await + .expect("park callback attempt")); + } + } + } + let schema = store.schema(); + let progress = serde_json::json!({"phase":"before-retained-head"}); + sqlx::query(&format!( + "INSERT INTO {schema}.attempt_state (job_id,run_lease,progress,heartbeat_at,updated_at) + VALUES ($1,$2,$3,clock_timestamp()-interval '1 hour',clock_timestamp()) + ON CONFLICT (job_id,run_lease) DO UPDATE SET progress=EXCLUDED.progress, + heartbeat_at=EXCLUDED.heartbeat_at" + )) + .bind(historical_pair.0) + .bind(historical_pair.1) + .bind(&progress) + .execute(&pool) + .await + .expect("seed mutable attempt state"); + if transition == RetainedBatchTransition::Deadline { + for table in ["leases", "lease_claims"] { + sqlx::query(&format!( + "UPDATE {schema}.{table} SET deadline_at=clock_timestamp()-interval '1 hour' + WHERE deadline_at IS NOT NULL" + )) + .execute(&pool) + .await + .expect("expire test deadlines"); + } + } + if transition == RetainedBatchTransition::Callback { + sqlx::query(&format!( + "UPDATE {schema}.leases SET callback_timeout_at=clock_timestamp()-interval '1 hour' + WHERE state='waiting_external'" + )) + .execute(&pool) + .await + .expect("expire test callbacks"); + } + let cursor_before = retained_receipt_cursor(&pool, &store, historical.claim_slot).await; + assert_eq!( + retained_closed_receipts(&pool, &store, historical_pair).await, + 0 + ); + let mut blocker = pool.begin().await.expect("begin retained lane blocker"); + sqlx::query(&format!( + "SELECT 1 FROM {schema}.queue_claim_heads + WHERE queue='qs_claimed_aging_body' AND priority=3 AND enqueue_shard=0 FOR UPDATE" + )) + .fetch_one(blocker.as_mut()) + .await + .expect("hold one required historical head"); + + assert_eq!( + retained_batch_transition(&pool, &store, transition).await, + vec![ordinary_pair] + ); + let active = store + .load_job(&pool, historical.job_id) + .await + .unwrap() + .unwrap(); + let expected_state = if transition == RetainedBatchTransition::Callback { + JobState::WaitingExternal + } else { + JobState::Running + }; + assert_historical_active_body(&active, &historical, expected_state); + assert_eq!(active.progress, Some(progress.clone())); + let stored_progress: serde_json::Value = sqlx::query_scalar(&format!( + "SELECT progress FROM {schema}.attempt_state WHERE job_id=$1 AND run_lease=$2" + )) + .bind(historical_pair.0) + .bind(historical_pair.1) + .fetch_one(&pool) + .await + .expect("blocked attempt state survives rollback"); + assert_eq!(stored_progress, progress); + assert_eq!( + retained_closed_receipts(&pool, &store, historical_pair).await, + 0 + ); + if !materialized { + assert_eq!( + retained_receipt_cursor(&pool, &store, historical.claim_slot).await, + cursor_before, + "an open first receipt remains eligible at its original cursor" + ); + } + assert_eq!(ready_tombstone_count(&pool, &store).await, 1); + assert!( + terminal_counts_by_priority(&pool, schema, "qs_claimed_aging_body") + .await + .is_empty() + ); + assert_eq!( + terminal_counts_by_priority(&pool, schema, queue).await, + vec![(4, 1)] + ); + blocker.rollback().await.expect("release retained head"); + assert_eq!( + retained_batch_transition(&pool, &store, transition).await, + vec![historical_pair] + ); + assert!(retained_batch_transition(&pool, &store, transition) + .await + .is_empty()); + let terminals: Vec<(i64, i64, i64)> = sqlx::query_as(&format!( + "SELECT job_id,run_lease,count(*) FROM {schema}.done_entries + WHERE job_id=ANY($1) GROUP BY job_id,run_lease ORDER BY job_id" + )) + .bind(vec![historical.job_id, ordinary_id]) + .fetch_all(&pool) + .await + .expect("read terminal attempt uniqueness"); + let mut expected = vec![ + (historical_pair.0, historical_pair.1, 1), + (ordinary_pair.0, ordinary_pair.1, 1), + ]; + expected.sort_unstable(); + assert_eq!(terminals, expected); + assert_eq!( + terminal_counts_by_priority(&pool, schema, "qs_claimed_aging_body").await, + vec![(4, 1)] + ); + assert_eq!( + terminal_counts_by_priority(&pool, schema, queue).await, + vec![(4, 1)] + ); + for _ in 0..8 { + assert!(store + .claim_runtime_batch(&pool, "qs_claimed_aging_body", 10, deadline) + .await + .expect("retired old replacement stays unclaimable") + .iter() + .all(|row| row.job.id != historical.job_id)); + } +} + +#[tokio::test] +async fn test_busy_retained_head_does_not_block_compact_heartbeat_rescue() { + historical_busy_attempt_does_not_block_batch(RetainedBatchTransition::Heartbeat, false).await; +} + +#[tokio::test] +async fn test_busy_retained_head_does_not_block_receipt_deadline_rescue() { + historical_busy_attempt_does_not_block_batch(RetainedBatchTransition::Deadline, false).await; +} + +#[tokio::test] +async fn test_busy_retained_head_does_not_block_materialized_deadline_rescue() { + historical_busy_attempt_does_not_block_batch(RetainedBatchTransition::Deadline, true).await; +} + +#[tokio::test] +async fn test_busy_retained_head_does_not_block_callback_rescue() { + historical_busy_attempt_does_not_block_batch(RetainedBatchTransition::Callback, true).await; +} + +#[tokio::test] +async fn test_busy_retained_head_cancel_preserves_lock_not_available() { + let (_guard, pool, store, claim) = + historical_tombstoned_claim_fixture(true, Duration::ZERO, false).await; + let mut blocker = pool.begin().await.unwrap(); + sqlx::query(&format!( + "SELECT 1 FROM {}.queue_claim_heads + WHERE queue='qs_claimed_aging_body' AND priority=3 AND enqueue_shard=0 FOR UPDATE", + store.schema() + )) + .fetch_one(blocker.as_mut()) + .await + .unwrap(); + let result = tokio::time::timeout( + ci_timing::scaled_timeout(Duration::from_secs(2)), + store.cancel_running( + &pool, + claim.job_id, + claim.run_lease, + "test cancellation", + None, + ), + ) + .await + .expect("single attempt cancellation fails without waiting"); + assert!( + matches!(result, Err(AwaError::Database(ref error)) if error.as_database_error().is_some_and(|error| error.code().as_deref() == Some("55P03"))), + "single attempt preserves lock_not_available: {result:?}" + ); + assert_eq!( + retained_closed_receipts(&pool, &store, (claim.job_id, claim.run_lease)).await, + 0 + ); + assert_eq!( + store + .load_job(&pool, claim.job_id) + .await + .unwrap() + .unwrap() + .state, + JobState::Running + ); + blocker.rollback().await.unwrap(); +} + +#[tokio::test] +async fn test_busy_retained_head_does_not_block_materialized_heartbeat_rescue() { + historical_busy_attempt_does_not_block_batch(RetainedBatchTransition::Heartbeat, true).await; +} + /// 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)] diff --git a/correctness/README.md b/correctness/README.md index 88dcd9de3..a874863ff 100644 --- a/correctness/README.md +++ b/correctness/README.md @@ -35,6 +35,7 @@ What is intentionally not modeled: - real-time token bucket math for rate limiting - unbounded retry histories; `AwaExtended` caps attempts so TLC can close the state graph - full advisory-lock mechanics; leadership is an abstract exclusive token +- ready-body identity. Job bodies are abstract, so the rescue liveness property `RecoverableAbandoned(j) ~> jobState[j] # "running"` does not cover an attempt that an older release's aging race left with two equally plausible bodies; rescue skips such an attempt and logs it on every pass. - exact permit identity per task attempt; the model is job-centric and approximates task-held capacity closely enough to check the protocol invariants - **priority ordering and starvation.** No model includes a `priority` variable. Priority is a scheduling heuristic, not a safety invariant — all formal invariants (NoDuplicateClaim, TaskLeaseBounded, RunningHasPermit, etc.) are independent of dispatch order. Cross-priority fairness is enforced by the maintenance leader's priority aging task (see ADR-005) - **MVCC, vacuum, and storage reclamation.** The models do not attempt to represent Postgres heap cleanup, autovacuum timing, or horizon pinning from long-lived snapshots. Those are operational/performance concerns rather than protocol-safety invariants, and are covered by runtime benchmarks instead diff --git a/correctness/storage/MAPPING.md b/correctness/storage/MAPPING.md index b1789212f..fb142d098 100644 --- a/correctness/storage/MAPPING.md +++ b/correctness/storage/MAPPING.md @@ -39,7 +39,7 @@ The TLA+ lifecycle model does not represent the completed-history rollup cache, | --- | --- | --- | | `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 | | `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)` | maintenance promote loop in `awa-worker/src/maintenance.rs::promote_due_state` | `DELETE FROM deferred_jobs ... INSERT INTO ready_entries ...` 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. | | `Claim(w, j)` | `QueueStorage::claim_runtime_batch` (`queue_storage.rs`) -> `claim_runtime_batch_with_aging_for_instance` -> dispatcher (`awa-worker/src/dispatcher.rs`) | `claim_ready_runtime()`: lane selection locks `queue_claim_heads` with `FOR UPDATE SKIP LOCKED`, consults `{schema}.ready_segments` (ordered by `next_lane_seq` for an index short-circuit) to identify the committed ready slot/generation for the lane head. It then validates the target ready row, reads `lease_ring_state` and `claim_ring_state`, treats `ready_tombstones` and durable attempt evidence as spent lane evidence for stale-cursor advancement, and in receipts mode INSERTs one `{schema}.ready_claim_attempt_batches_` row for the claimed lane range before inserting compact `lease_claim_batches_` evidence for zero-deadline claims or row-local `lease_claims_` evidence for deadline-backed claims. Legacy mode inserts into `leases_`. The common live-head claim plan filters ready tombstones only; already-emitted attempt ranges are consulted after the head row is proved spent, where they advance a stale claim cursor instead of making every successful claim anti-join the attempt ledger. The broad legacy spent-evidence battery is retained only for non-receipts custom schemas. | | `MaterializeAttemptState(j)` | `QueueStorage::upsert_attempt_state_from_receipts_tx` and `QueueStorage::upsert_attempt_state_progress_from_receipts_tx` | `INSERT INTO attempt_state ... ON CONFLICT (job_id, run_lease) DO NOTHING` | @@ -54,10 +54,10 @@ The TLA+ lifecycle model does not represent the completed-history rollup cache, | `StatefulComplete(w, j)` | `QueueStorage::complete_runtime_batch` + `DELETE FROM attempt_state` | same as above with an observed materialized `attempt_state`; terminal reads hydrate the immutable body through `{schema}.terminal_jobs` and the retained ready row | | `FailToDlq(w, j)` | `QueueStorage::fail_to_dlq` / `QueueStorage::fail_terminal` via executor terminal failure path | `DELETE FROM leases`, `DELETE FROM attempt_state`, `INSERT INTO dlq_entries` in one tx | | `RetryToDeferred(w, j)` | `QueueStorage::retry_after` / `QueueStorage::snooze` on `JobError::RetryAfter` / `Snooze` | `DELETE FROM leases`, `INSERT INTO deferred_jobs` | -| `RescueToReady(j)` | lease-side `rescue_stale_heartbeats` / `rescue_expired_deadlines` in maintenance, plus receipt-side deadline rescue once it has force-closed an open receipt | lease-side paths `DELETE FROM leases ... RETURNING ...; INSERT INTO ready_entries ...`; receipt-side deadline rescue first appends a `deadline_expired` closure row in `lease_claim_closures`, then the same maintenance routing re-appends/dead-letters according to retry policy | +| `RescueToReady(j)` | lease-side `rescue_stale_heartbeats` / `rescue_expired_deadlines` in maintenance, plus receipt-side deadline rescue once it has force-closed an open receipt | lease-side paths `DELETE FROM leases ... RETURNING ...; INSERT INTO ready_entries ...`; receipt-side deadline rescue first appends a `deadline_expired` closure row in `lease_claim_closures`, then the same maintenance routing re-appends/dead-letters according to retry policy. Each candidate is hydrated inside a savepoint: an attempt whose claim heads are busy is rolled back and deferred to a later pass, and an attempt whose retained ready bodies are ambiguous (legacy aging data) is skipped and logged on every pass. | | `CancelWaitingToTerminal(j)` | waiting branch in `QueueStorage::cancel_job_tx` | `DELETE FROM leases WHERE state IN ('running', 'waiting_external')`, hydrate from `ready_entries`, insert `done_entries`, close the matching receipt | | `CancelReadyToTerminal(j)` | available branch in `cancel_job_tx` | `SELECT ready_entries ... FOR UPDATE SKIP LOCKED`, append `ready_tombstones`, insert `done_entries(cancelled)`, and post-commit advance the claim cursor only if the cancelled lane was the current head. The ready row remains until queue prune. | -| `ReprioritizeReady(j)` | `QueueStorage::age_waiting_priorities` and batch `set_priority` via `QueueStorage::set_priority_tx` / `move_ready_batch_fields_tx` | `SELECT ready_entries ... FOR UPDATE SKIP LOCKED`, append `ready_tombstones` for the old lane, append a new `ready_entries` row at the target priority, and notify the queue. Batch operations snapshot job ids in `awa.batch_operation_items`; item status and row mutation commit in one transaction so resume does not replay a completed item. | +| `ReprioritizeReady(j)` | `QueueStorage::age_waiting_priorities` and batch `set_priority` via `QueueStorage::set_priority_tx` / `move_ready_batch_fields_tx` | Lock the affected `queue_claim_heads` rows in sorted order, recheck that no claim has emitted the row, `SELECT ready_entries ... FOR UPDATE SKIP LOCKED`, append `ready_tombstones` for the old lane, append a new `ready_entries` row at the target priority, and notify the queue. Batch operations snapshot job ids in `awa.batch_operation_items`; item status and row mutation commit in one transaction so resume does not replay a completed item. | | `MoveQueueReady(j)` | batch `move_queue` via `QueueStorage::move_queue_tx` / `move_ready_batch_fields_tx` | Same abstract storage transition as `ReprioritizeReady`: tombstone source ready lane, append destination ready row with rewritten queue/optional priority, and notify the destination queue. Queue labels and cross-queue routing are outside `AwaSegmentedStorage`'s one-lane abstraction; lock-order implications are covered by `AwaStorageLockOrder`. | | `StaleCompleteRejected(w, j)` | `complete_runtime_batch` returning `CompletionOutcome::IgnoredStale` | `UPDATE leases ... WHERE run_lease = $2` matching 0 rows | | `MoveFailedToDlq(j)` | `QueueStorage::move_failed_to_dlq`; admin entry in `awa-model/src/dlq.rs` | hydrate from retained ready body, `DELETE FROM done_entries`, append a negative `queue_terminal_count_deltas` row, and `INSERT INTO dlq_entries` guarded by state=failed. The retained ready row remains until queue prune. | diff --git a/docs/configuration.md b/docs/configuration.md index e59a1e665..dae5622ef 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -582,9 +582,13 @@ Use the defaults unless you have a reason not to: - Increase `claim_slot_count` if the rotation cadence (`claim_rotate_interval`) plus the slot count combine to a partition retention window shorter than your longest in-flight zero-deadline short job; running out of empty slots forces `rotate_claims` to return `SkippedBusy` and the receipt-plane churn falls back onto a smaller working set of partitions. - Increase `queue_stripe_count` only for measured workloads where many small replicas contend on the same logical queue. This is queue-storage configuration, not a per-queue `QueueConfig` field: it gives each logical queue the same number of physical stripes. Striping spreads a hot logical queue over `queue#N` physical coordination paths, but it weakens perfect global ordering and can regress calmer shapes if overused. For single hot-queue workloads, benchmark `2` stripes first; higher values should be treated as workload-specific tuning rather than a general default. - Increase `prune_lock_timeout` when `increase(awa_maintenance_prune_attempts_total{awa_ring_outcome="blocked"}[15m])` stays above zero and the rings stop recycling, typically because frequent short reads of the ring partitions (admin or dashboard queries over the queue-storage tables) hold a conflicting lock at the instant prune tries. A stalled ring grows its current slot, which makes those queries slower and the contention worse. Raise it in small steps (for example `50ms`, then `100ms`, then `250ms`) and keep it well below PostgreSQL's `deadlock_timeout` (default `1s`) and below any `statement_timeout` on the pool's connections. A single lock wait that reaches `deadlock_timeout` makes PostgreSQL run its deadlock detector, and if claims or enqueues are part of a lock cycle with it, the detector can cancel them rather than the prune; a short `lock_timeout` abandons the prune first, so the prune is always the one that gives way. A `statement_timeout` at or below the lock timeout cancels the prune statement before the lock timeout fires, so it fails with an error instead of reporting `blocked`. Startup logs a warning when the configured timeout is at or above either setting. While prune waits, new accesses to the same partition also queue behind its lock request, so a longer timeout trades claim latency for reclamation. Shortening the readers is the more durable fix when they run longer than a few hundred milliseconds. Watch `awa_maintenance_prune_attempts_total` by `awa_ring` and `awa_ring_outcome`: `pruned` should advance steadily, `blocked` means the lock timeout expired, `skipped_active` (with `awa_ring_reason`) means live rows kept the slot, `backed_off` means prune was skipped for up to 32 ticks after consecutive `skipped_active` or `blocked` results, so a persistently blocked ring is retried less often over time, and `error` is a database failure. -- Increase rotation intervals to reduce partition churn and metadata activity. +- Consider increasing rotation intervals before raising `prune_lock_timeout`: longer intervals can reduce how often prune requests exclusive locks, while a longer timeout lets each request wait longer. Confirm the effect with workload measurements. - Decrease rotation intervals to tighten dead-tuple bounds at the cost of more frequent rotate/prune work. +Longer rotation intervals accumulate more work in each slot. At a steady enqueue rate, ready rows per queue slot are roughly the enqueue rate multiplied by `queue_rotate_interval`; retries and blocked recycling can increase the retained footprint. Budget for larger partitions and indexes. Queue-ring history becomes coarser and generally spans about `queue_slot_count × queue_rotate_interval`; active or blocked slots can retain it longer. The independent `claim_rotate_interval` changes claim-evidence retention, not completed-job history in the ready/terminal ring. + +A longer interval does not remove the locks required by each prune. `prune_lock_timeout` still applies separately to each acquisition; queue prune can lock up to eight partitions before its later reads and writes. + ### Sharding the enqueue head per queue Each queue's `awa.queue_meta.enqueue_shards` (default `1`, range `1..=64`) governs how many independent enqueue head rows the queue has. With the default value, the queue uses a single head row and the ordering contract is strict FIFO per `(queue, priority)` — identical to v016. diff --git a/docs/queue-storage-substrate.md b/docs/queue-storage-substrate.md index 2f693ab80..c3720b475 100644 --- a/docs/queue-storage-substrate.md +++ b/docs/queue-storage-substrate.md @@ -118,7 +118,7 @@ The split between migration-owned default substrate and helper-installed custom - **Migrations that depend on queue-storage tables can write unconditional DDL.** Operations like the `done_entries` terminal-count delta append inside `awa.delete_job_compat()` need the delta table to exist. Migration ordering guarantees it. - **The default schema cannot be accidentally destroyed.** The reset guard above protects operators from a `DROP SCHEMA awa CASCADE` that would take the canonical migration tables with it. -The helper is `SECURITY INVOKER` so callers need their own DDL privileges on the target schema; the runtime role does not gain DDL through the helper, which keeps the principle-of-least-privilege role model intact. Once each queue is provisioned with `awa storage prepare-queue` and the `prepared_lane_sequences` schema patch is applied, a runtime role with DML, sequence `USAGE`/`SELECT`/`UPDATE`, function `EXECUTE` and `TRUNCATE` needs no schema `CREATE`; see [Provision queues before starting restricted runtimes](security.md#provision-queues-before-starting-restricted-runtimes). +The helper is `SECURITY INVOKER` so callers need their own DDL privileges on the target schema; the runtime role does not gain DDL through the helper, which keeps the principle-of-least-privilege role model intact. Once the `definer_lane_provisioning` schema patch is applied, new lanes are created through schema-owner `SECURITY DEFINER` functions, so a runtime role with DML, sequence `USAGE`/`SELECT`/`UPDATE`, function `EXECUTE` and `TRUNCATE` needs no schema `CREATE`, even to start using a new queue, priority or shard; see [Lane sequences](security.md#lane-sequences). ## See also diff --git a/docs/security.md b/docs/security.md index 2b9848ec7..087203113 100644 --- a/docs/security.md +++ b/docs/security.md @@ -176,7 +176,7 @@ The default `awa` queue-storage substrate is migrated by `awa migrate`. Custom q ### Lane sequences -Each queue lane, `(queue, priority, enqueue_shard)`, has an enqueue and a claim sequence. Once the `definer_lane_provisioning` schema patch is applied, they are created by the schema owner, whichever role first enqueues into, claims from or promotes into the lane: the lane helpers call `provision_lane` and `provision_lane_sequences` in the queue-storage schema. These are `SECURITY DEFINER` functions owned by the schema owner with a pinned `search_path`. Like the other runtime-callable functions they are executable by `PUBLIC`, so schema `USAGE` is what gates them, and they only create lanes, after checking that the queue, priority and shard are set and the shard is 0–63. A runtime role therefore needs neither schema `CREATE` nor default privileges to start using a new queue, priority or shard. Apply the patch with `awa migrate` as the schema owner (`PGOPTIONS='-c role=awa_owner'`), or run the exported `R__definer_lane_provisioning.sql` after `V40`. +Each queue lane, `(queue, priority, enqueue_shard)`, has an enqueue and a claim sequence. Once the `definer_lane_provisioning` schema patch is applied, they are created by the schema owner, whichever role first enqueues into, claims from or promotes into the lane: the lane helpers call `provision_lane` and `provision_lane_sequences` in the queue-storage schema. These are `SECURITY DEFINER` functions owned by the schema owner with a pinned `search_path`. Like the other runtime-callable functions they are executable by `PUBLIC`, so schema `USAGE` is what gates them, and they only create lanes, after checking that the queue, priority and shard are set and the shard is 0–63. A runtime role therefore needs neither schema `CREATE` nor default privileges to start using a new queue, priority or shard. Their inserts into `queue_lanes`, `queue_enqueue_heads` and `queue_claim_heads` fire those tables' triggers as the schema owner, so don't grant `TRIGGER` on them to any other role. Apply the patch with `awa migrate` as the schema owner (`PGOPTIONS='-c role=awa_owner'`), or run the exported `R__definer_lane_provisioning.sql` after `V40`. A new lane sequence gets the privileges of the head table it backs. A role that may `INSERT` or `UPDATE` `queue_enqueue_heads` gets `USAGE, SELECT, UPDATE` on the enqueue sequences, and one that may only `SELECT` it gets `SELECT`; `queue_claim_heads` does the same for the claim sequences. The table grants above are enough. @@ -209,6 +209,18 @@ Which role can repair a sequence: The schema owner holds neither the migration role's nor an application role's privileges, so a repair running as `awa_owner` cannot reassign sequences those roles created. That includes the migration role connected with `PGOPTIONS='-c role=awa_owner'`. Run the repair as the migration role itself, without `PGOPTIONS`, as above. `awa migrate` and `--repair` retry each sequence the effective role could not repair as the login role, so `awa migrate` run with `PGOPTIONS='-c role=awa_owner'`, as documented, still repairs the migration role's sequences. +For a migrator with SET-only membership in `awa_owner`, grant schema access and helper execution explicitly before running the repair. Run as the schema owner: + +```sql +GRANT USAGE ON SCHEMA awa TO awa_migrator; +GRANT EXECUTE ON FUNCTION + awa.lane_sequence_problems(TEXT, TEXT), + awa.repair_lane_sequence(TEXT, TEXT) + TO awa_migrator; +``` + +For a custom queue-storage schema, also grant `USAGE` on that schema. The migrator must still hold the old sequence owner's privileges and be allowed to `SET ROLE` to the schema owner; inherited privileges from the schema owner are not required. + For each sequence that remains, the output lists `repair_roles`, the login roles other than superusers that can repair it, and the statements to run. A busy sequence only needs another `--repair`. When `repair_roles` is empty, which is usual for a sequence an application role created on managed PostgreSQL, either: - run `missing_grants` as the current owner. Every runtime role can then use the lane, although the owner stays unchanged; or @@ -325,7 +337,7 @@ This is additive — no schema changes, no downtime. ## Future: tighter runtime grants -The current model requires broad table grants because compatibility triggers and helper functions still use `SECURITY INVOKER`. A future migration could convert the internal queue-storage helpers to `SECURITY DEFINER` (running as `awa_owner`), which would let the runtime role be restricted to the active queue-storage schema plus the shared control tables (`queue_meta`, `job_unique_claims`, `cron_jobs`, `runtime_instances`, and `runtime_storage_backends`). This is tracked but not yet implemented — the current model is secure for the separation it provides (runtime can't modify schema). +The lane provisioning functions `provision_lane` and `provision_lane_sequences` are `SECURITY DEFINER` since the `definer_lane_provisioning` patch (see [Lane sequences](#lane-sequences)). The rest of the model still requires broad table grants because the other compatibility triggers and helper functions use `SECURITY INVOKER`. A future migration could convert those to `SECURITY DEFINER` (running as `awa_owner`), which would let the runtime role be restricted to the active queue-storage schema plus the shared control tables (`queue_meta`, `job_unique_claims`, `cron_jobs`, `runtime_instances`, and `runtime_storage_backends`). This is tracked but not yet implemented — the current model is secure for the separation it provides (runtime can't modify schema). ## Deployable roles diff --git a/docs/test-plan.md b/docs/test-plan.md index 618ff69a2..dffcb3332 100644 --- a/docs/test-plan.md +++ b/docs/test-plan.md @@ -305,6 +305,76 @@ Planned test matrix for the 0.7 cycle, mapped to the roadmap ([`0.7-roadmap.md`] | V28 | Partitioned-queue preset ≥9k jobs/s e2e on the reference 24-CPU harness | #383 Gate 2 | | V29 | Gate A evidence: allocator bake-off (E1) + WAL/job decomposition (E3) run ids recorded on #295 | #295 | +## Persona soak + +Durable loss detection covers every lifecycle-observed job and a deterministic `job_id % 64 = 0` sample of `NO_EVENTS` workloads. The other 63/64 returned `NO_EVENTS` jobs remain in enqueue/execution accounting, but their terminal commits are not proven. A bulk observer polls sampled terminal rows and eventful attempts whose handlers returned without terminal proof once per second (2 s query timeout, 100,000-row cap) and keeps the run lease in the client ledger, so captured proof survives pruning. A handler return is never terminal proof. Missing sampled proof after pruning or a failed invariant query makes the report `INCONCLUSIVE` and fails the test; it does not assert job loss or establish a pass. Reports include sampled, confirmed and excluded counts. + +On CI, assertion waits and heartbeat staleness use `ci_timing.rs` scaling; load durations, kill schedules and measurement windows retain their configured values. + + +`awa/tests/persona_soak_test.rs` (harness in `awa/tests/persona_soak/`, Python side in `awa/tests/persona_soak/python/`) runs several application-shaped workloads concurrently against one database, the way different applications share an `awa` schema, and fails the run on any invariant violation. + +| Persona | Workload | Stresses | +| --- | --- | --- | +| Mia, ML pipeline | jobs of 0.1-60 s with per-job deadlines, progress updates, retries into the DLQ, handler cancels, latency spikes | deadline and stale-heartbeat rescues racing heartbeats and completions, DLQ | +| Tom, multi-tenant | 100 queues at a low rate each, unique jobs, scheduled jobs, cron every 10 s on two runtimes, mixed priorities with aging | lane and queue cardinality, per-queue maintenance, cron fencing | +| Priya, firehose | one hot queue, 100k-row `insert_many` bursts, COPY batches, steady small batches, tiny jobs, three runtimes | bind-parameter limit, ring rotation and prune, claim contention | +| Sam, ordered streams | a queue with 4 enqueue shards and 64 ordering keys, a 4-stripe queue (`queue_stripe_count`), scheduled jobs, a rate-limited queue | routing, promotion, fairness | +| Alex, operator | pause/resume, drain, cancel and bulk cancel, move and reprioritise batch operations, DLQ retry and purge, the admin read queries, rolling restarts of other personas' runtimes | admin paths under live load | +| Chaos | kills a runtime or the Python process without a shutdown handshake, terminates the maintenance leader's connection and random runtime connections, exhausts a runtime's pool, holds an idle-in-transaction session | recovery paths | +| Watchers | six dashboard tabs polling overview, queues, counts, timeseries, runtime and storage at the UI's refresh interval, some also the kinds, queue-runtime and cron pages | admin read queries under load, through the real `awa serve` HTTP server | +| Support engineer | searches jobs by kind, queue, state and tag, pages through large result sets with `before_id`, opens job detail and follows a job at the detail page's poll rate, pages the DLQ, opens runtime, cron and kind pages | search and paging cost against retained history and backlog | +| UI operator | the UI's mutations over HTTP: pause and resume queues and cron schedules, cancel and bulk-cancel, DLQ retry and purge (single and by filter), batch operation preview and submit | mutation endpoints racing live work | +| Impatient user | requests abandoned after a few milliseconds, bursts of refreshes, while sampling the server's database sessions | an abandoned request must not keep a pool connection past the statement timeout | +| Fanout, replay fan-out | parent jobs whose handlers batch-insert 100-185 unique child jobs (`insert_many`, unique by queue and args) into their own queue from many concurrent handlers; some parents replay a batch with an already-claimed key, hit the unique violation and fall back to row-by-row inserts | concurrent batch enqueue into one lane, unique claims, inserting from inside a handler | +| Shard operator | an operator changes `enqueue_shards` on two hot queues while producers, promotion and claimers run: 1 → 4 → 8, down to 2 and back to 1, sometimes after `prepare_queue` and sometimes not, restarting some runtimes and producer `QueueStorage` handles before others so cached shard counts disagree. Producers use `insert_many`, `insert_many_copy` and `QueueStorage::enqueue_params_batch`, with scheduled jobs, retries (promotion into shards) and `ordering_key` jobs. The actors run as non-superuser roles: a schema owner that migrates and runs `prepare_queue`, a migrator role with default privileges, and two application roles in a runtime group role with no default privileges, granted as in `docs/security.md`. For some raises the producers are held until the maintenance leader's runtime has created the new lanes, so a different role is the next to use them | live shard-count changes (ADR-025), lane provisioning by different roles, stale per-process shard counts, promotion into new shards | +| Claim and abandon | a separate process (the harness binary started again with one test selected) that claims as much as it can: `max_workers` equal to the target on two queues, one with zero deadlines (compact claim batches) and one with a deadline (one claim row per job), jobs held until the process dies, heartbeats running. It is killed with SIGKILL, repeatedly, and replaced after a few seconds by a process that claims the rescued jobs again; some jobs have `max_attempts` 2 so repeated kills dead-letter them. A surviving in-process runtime leads maintenance and serves a canary queue of short jobs, and the other personas keep working; after the last kill an in-process runtime finishes the rescued jobs | rescue of a large abandoned-claim backlog, claim-ring rotation while slots hold abandoned claims, false rescues of live owners, `max_attempts` and the DLQ under repeated rescue | +| Python users | a separate process using `awa-python`: single inserts, a 6,500-row `insert_many` in one transaction and a COPY batch, enqueues in transactions that sometimes roll back; async workers that complete, raise, `RetryAfter`, `Snooze`, report progress, overrun a queue deadline and ignore their deadline; a queue shared with a Rust runtime; a periodic migrate-then-query probe | Python bindings against live load, mixed Rust and Python runtimes on one lane, the documented slot behaviour of an ignored deadline, connections stranded after `migrate()` | + +The oracle checks, while the load runs and again after it stops: + +- a job attempt with more than one terminal row (`done_entries`, `dlq_entries`, compact receipt completions), a job reported completed twice, and a job that is terminal and still available, deferred or leased +- a cancelled job nobody asked to cancel, an operator cancel that overwrote a completion, and a handler that started after its operator cancel committed +- a queue with claimable jobs and idle workers whose claim cursor stops advancing (stall), and scheduled jobs overdue beyond a threshold (promotion stall) +- a lost job: enqueued, never terminal after the system quiesces, and not removed by an operator action; a cron fire that produced more than one job +- shard operations: every job enqueued before, during and after a shard-count change runs and completes exactly once (a repeat is allowed only after a stale-heartbeat rescue and is reported as a finding), every shard drains including those above a lowered count, claim cursors keep advancing on every shard that has backlog while workers are idle (the per-shard lag query from ADR-025), `ordering_key` jobs start in order per key within a window in which every runtime had observed the shard count, and no role raises a permission-denied or missing-lane error (found in the server log, with the role, lane, sequence owner and statement; needs `AWA_SOAK_PG_CONTAINER` or `AWA_SOAK_PG_LOG`, and `log_line_prefix` containing `%u@%a`, as `scripts/persona-soak.sh` sets) +- claim and abandon: after each SIGKILL the rescue backlog (open claims whose last heartbeat or claim time is older than the staleness window, or whose deadline passed, counted every few seconds from the claim tables and leases) clears within `AWA_SOAK_ABANDON_BOUND_SECONDS`, with the age and count of the oldest unrescued expired claim sampled into `abandon-backlog.csv` and the time to clear reported per kill; the claim ring does not stay still for longer than the bound once the slot it must enter next holds no claim whose owner is alive (a ring pinned by a live process is expected); no job is run again while the process that held it is alive (every start is recorded with the process that ran it, and a repeat is allowed only after that process was killed); no attempt exceeds `max_attempts`, no job is dead-lettered early, and a job abandoned on its last attempt ends in the DLQ +- a transaction the Python application rolled back that left a job in the store, and database sessions left stranded by the Python `migrate()` path +- an admin-server query that outlives its statement timeout, and sessions still running a query after every UI client has stopped +- a server crash: `pg_postmaster_start_time()` changing, or `terminated by signal`, `TRAP:`, `server process (PID` and similar lines in the Postgres log + +Slow or failing admin endpoints, latency that grows with retained history or backlog, and admin queries cancelled by the statement timeout are reported as non-fatal findings in the report, next to the violations. + +Terminal rows are pruned by ring rotation within seconds, so the oracle keeps a client-side ledger (job ids per persona, executions per job, lifecycle events, operator actions) and reconciles it with the database. Time-based checks are widened while a chaos action is recovering. + +Environment: + +| Variable | Meaning | +| --- | --- | +| `AWA_SOAK_DATABASE_URL` | required; database name must match `awa_soak[A-Za-z0-9_]*` and is dropped and recreated | +| `AWA_SOAK_PERSONAS` | comma list of `mia,tom,priya,sam,alex,chaos,watchers,support,uiops,impatient,python,fanout,shardops,abandon`, default all | +| `AWA_SOAK_ABANDON_CLAIMS`, `AWA_SOAK_ABANDON_KILL_SECONDS`, `AWA_SOAK_ABANDON_KILLS` | the claim-and-abandon persona's scale, independent of `AWA_SOAK_SCALE`: claims each killed process holds (default 4000; 25,000 to 70,000 reproduces a large backlog), seconds between kills (default 180, long enough for the process to claim everything and heartbeat for a while) and how many processes are killed (default 4; the last kill is the end of the load phase if that comes first) | +| `AWA_SOAK_ABANDON_BOUND_SECONDS`, `AWA_SOAK_ABANDON_RESTART_SECONDS` | the bound after a kill within which the rescue backlog must clear and the claim ring must be free to rotate (default 120), and the pause before the replacement process starts (default 3) | +| `AWA_SOAK_HEARTBEAT_SECONDS`, `AWA_SOAK_STALENESS_SECONDS` | heartbeat interval (default 2) and heartbeat staleness after which a claim may be rescued (default 12) of every runtime, including the claim-and-abandon process; the library's own defaults are 30 and 90. The claim-and-abandon persona also records the duration of each heartbeat pass of its process against the number of handlers running, read from the timestamps the heartbeats write (`abandon-heartbeat-passes.csv`, and percentiles by handler count in the report) | +| `AWA_SOAK_RESCUE_INTERVAL_SECONDS` | heartbeat and deadline rescue interval of every runtime; unset keeps the harness's 3 s and 2 s, the library default for heartbeat rescue is 30 s | +| `AWA_SOAK_SHARDOPS_SCALE`, `AWA_SOAK_SHARDOPS_SCHEMA_CREATE` | the shard operator's rate multiplier, independent of `AWA_SOAK_SCALE` (default 1), and whether its application roles may create objects in the schema (default 1: lanes are created on the fly; 0 withholds schema `CREATE` and the operator provisions every lane before raising a count) | +| `AWA_SOAK_SHARED_CREATOR_GRANTS` | default `0`; `1` shares application-created objects for comparisons with older binaries and requires schema `CREATE`. Keep it off when checking application-role privileges. | +| `AWA_SOAK_UI_POOL_MAX`, `AWA_SOAK_ADMIN_READ_SECONDS` | the admin server's database pool size (default 10) and the seconds between the operator persona's admin read queries (default 2); lower the load on a slow server by shrinking the pool and stretching the interval | +| `AWA_SOAK_AWA_BIN` | the `awa` binary for the UI personas (`cargo build --release -p awa-cli`); defaults to the one next to the test binary | +| `AWA_SOAK_PYTHON` | interpreter of a virtualenv with `awa-python` installed (`uv venv awa-python/.venv`, `uv pip install maturin`, `maturin develop --release` in `awa-python/`); required by the `python` persona | +| `AWA_SOAK_UI_STATEMENT_TIMEOUT_SECONDS`, `AWA_SOAK_UI_CLIENT_TIMEOUT_SECONDS` | the admin server's statement timeout (default 30) and the UI personas' request timeout (default 45); a request that hits either is reported as a finding | +| `AWA_SOAK_SECONDS` | load phase length, default 900 | +| `AWA_SOAK_SEED` | fixes the workload mix and chaos schedule; printed at start | +| `AWA_SOAK_SCALE` | multiplies rates, burst sizes and queue counts, default 1 | +| `AWA_SOAK_DRAIN_GRACE_SECONDS` | how long to wait for the system to quiesce after load stops, default 120 | +| `AWA_SOAK_PG_LOG` / `AWA_SOAK_PG_CONTAINER` / `AWA_SOAK_PG_CONTAINER_LOG` | where to read the server log for the crash scan: a file on this host, or a Docker container (with an optional log path inside it, otherwise `docker logs`) | +| `AWA_SOAK_REPORT_DIR` | where `report.md`, `report.json` and, for a failed run, `reproducer.md` are written | +| `AWA_SOAK_STALL_SECONDS`, `AWA_SOAK_PROMOTION_OVERDUE_SECONDS`, `AWA_SOAK_CHAOS_SLACK_SECONDS` | oracle thresholds | + +The report records the engine, server version, `max_connections`, connections at start and the per-persona connection budget. A failed run leaves `reproducer.md` with the seed, personas and duration, the violating rows, the timeline around each violation and the relevant server log excerpt. `scripts/persona-soak.sh` runs the engine matrix (`postgres:17`, `postgres:18`, and an assert-enabled PostgreSQL 18 image) in private containers on separate ports and removes them afterwards. + +After-exit retained-session findings require a successful `CLOSED` acknowledgement and a valid final session observation. Forced or unproven exits remain incomplete evidence. Pool-budget and idle-transaction checks after readiness remain independent. + ## Running Tests ```bash @@ -329,6 +399,26 @@ DATABASE_URL=postgres://postgres:test@localhost:15432/awa_test \ cargo test --package awa --test chaos_suite_test \ -- --ignored --test-threads=1 --nocapture +# Persona soak: Mia (ML pipeline), Tom (multi-tenant fan-out), Priya (firehose), +# Sam (ordered streams), Alex (operator) and a chaos agent run together against one +# disposable database on the queue-storage engine, checked by an invariant oracle. +# The database name must start with `awa_soak`; it is dropped and recreated. +AWA_SOAK_DATABASE_URL=postgres://postgres:test@localhost:15432/awa_soak_run \ + AWA_SOAK_SECONDS=900 \ + cargo test --release -p awa --test persona_soak_test -- --ignored --nocapture + +# The same soak against throwaway PostgreSQL 17, 18 and assert-enabled 18 containers +scripts/persona-soak.sh # or: scripts/persona-soak.sh pg17 + +# Only the claim-and-abandon persona, with 25,000 claims abandoned by each of four kills; +# the Python bindings and the admin server binary are not needed +SOAK_SKIP_PREPARE=1 AWA_SOAK_PERSONAS=abandon AWA_SOAK_ABANDON_CLAIMS=25000 \ + AWA_SOAK_SECONDS=1500 AWA_SOAK_DRAIN_GRACE_SECONDS=600 scripts/persona-soak.sh pg17 + +# The same harness against another revision of awa (a release tag or a candidate branch): +# a detached worktree of that revision gets this checkout's harness files copied over it +scripts/persona-soak-against.sh v0.6.13 pg17 + # Postgres hot-standby promotion smoke (boots primary + replica + stable proxy endpoint via Docker Compose) cargo test --package awa --test postgres_failover_smoke_test -- --ignored --nocapture diff --git a/docs/troubleshooting.md b/docs/troubleshooting.md index bfa095cd5..46605e5f9 100644 --- a/docs/troubleshooting.md +++ b/docs/troubleshooting.md @@ -81,9 +81,18 @@ ORDER BY heartbeat_at ASC; Interpretation: -- low `heartbeat_age_s`: the worker is still alive +- low `heartbeat_age_s`: the worker is still alive. If its handler has already returned, the worker is retrying finalization after lock contention, for up to 120 s, and its log shows `Durable finalization failed; retrying` - `heartbeat_age_s` well above `heartbeat_staleness` (default `90s`): stale-heartbeat rescue should pick it up on the next scan - very negative `deadline_remaining_s`: deadline rescue should pick it up soon (lease side or receipt side, depending on which table the row lives in) +- stale, but every rescue pass logs `Rescue skipped an attempt: queue storage ready body ambiguous`: an aging race in an older release left two equally plausible bodies for this attempt. Rescue leaves it running and moves on to other attempts. Reads and cancellation of the job fail with the same ambiguity error, so it needs manual repair; open an issue with the job id and run lease from the log + +### Transient Admin Reads During Partition Pruning + +In 0.6, a job-detail read, including list-item hydration, uses several queries within one snapshot. It retains read locks on the located ready body, but does not lock all ring partitions before establishing that snapshot. Prune can truncate a ready partition before its body is located, or a lease or claim partition before that read phase reaches it. A completed job can then briefly appear available or missing. Retry with a fresh read before treating that result as missing durable work. Filtering tombstoned source rows after a queue or priority move does not close these separate windows. + +### Failed Jobs Rescued After Priority Aging + +0.6.14 preserves the physical ready-lane identity when rescue writes a failed terminal row, so manual retry, failed-row carry and `terminal_jobs` retain its original arguments and retry limit. Rust reads, manual retry and failed-row carry also recover older narrow terminal rows by attempt identity while their ready backing remains retained. These paths preserve terminal metadata and fill missing defaults from the live ready row. The installed `terminal_jobs` SQL view and creation-time bucket joins remain unchanged for older rows whose stored priority differs from their ready lane. Already-pruned backing data cannot be recovered by this fix. Manual terminal retry restarts aging from the physical lane. Manual terminal-to-DLQ moves retain that lane for later DLQ retry; automatic retryable rescue and automatic DLQ routing retain effective priority. ### Inspect Runtime Health @@ -139,8 +148,7 @@ In all cases the cancel itself is durable — only the early wake-up is lost. Lo ## Rescue Fails With `idx_awa_jobs_unique` -Symptom: maintenance repeats, every tick (canonical engine — the queue-storage -engine wedges identically but logs `error: unique conflict` instead): +Symptom: maintenance repeats, every tick (canonical engine — the queue-storage engine wedges identically but logs `error: unique conflict` instead): ```text Failed to rescue stale heartbeat jobs, error: ... duplicate key value violates unique constraint "idx_awa_jobs_unique" diff --git a/docs/upgrade-0.6.13-to-0.6.14.md b/docs/upgrade-0.6.13-to-0.6.14.md index b18fa715e..4d78ee395 100644 --- a/docs/upgrade-0.6.13-to-0.6.14.md +++ b/docs/upgrade-0.6.13-to-0.6.14.md @@ -5,17 +5,17 @@ ## Upgrade 1. Install the 0.6.14 CLI. -2. Check lane ownership. This is read-only and exits 2 while lanes need repair: +2. Run `awa migrate` once, with the schema owner as the effective role, while workers keep running: ```bash - awa --database-url "$AWA_DATABASE_URL" storage lane-sequences + PGOPTIONS='-c role=awa_owner' \ + awa --database-url "$AWA_MIGRATOR_DATABASE_URL" migrate ``` -3. Run `awa migrate` once, with the schema owner as the effective role, while workers keep running: +3. Check lane ownership after the patch installs its diagnostic helper. This is read-only and exits 2 while lanes need repair: ```bash - PGOPTIONS='-c role=awa_owner' \ - awa --database-url "$AWA_MIGRATOR_DATABASE_URL" migrate + awa --database-url "$AWA_DATABASE_URL" storage lane-sequences ``` 4. Read the output. The `definer_lane_provisioning` patch reports `applied`, followed by how many lane sequences were repaired. Any left over are listed in warnings: a busy one needs another run of `awa storage lane-sequences --repair` as the migration role (without `PGOPTIONS`), and the others name the roles that can repair them (see [Check and repair lane ownership](security.md#check-and-repair-lane-ownership)). The migration exit status does not depend on them. @@ -27,9 +27,11 @@ The schema stays at v40. The `definer_lane_provisioning` patch installs `provisi Once the patch commits, `awa migrate` repairs existing lane sequences, each in its own short transaction. A sequence not owned by the schema owner is reassigned if the role holds its current owner's privileges and can `SET ROLE` to the schema owner. Missing grants are added if the role holds the owner's privileges. Each sequence is tried as the effective role and then as the login role. Reassigning waits at most 300 ms for the sequence's lock, which an open enqueue or claim on that lane holds, and holds it only for that sequence's transaction. A sequence that is busy, or that neither role may change, is reported and left as it was. +If an earlier build already applied the patch, run `awa migrate` again with the updated CLI. It detects the older installed repair helper and reapplies the repeatable patch before attempting repairs. + The patch: -- only replaces functions; it takes no sequence or table locks; +- only replaces functions, gives the two provisioning functions to the schema owner and grants `EXECUTE` on them to `PUBLIC`; it takes no sequence or table locks; - keeps every lane's sequence position; - leaves storage routing unchanged; - applies as one transaction, so an interrupted run leaves nothing applied; run it again. @@ -38,7 +40,13 @@ The patch: ### External migration runners -`awa migrate --sql --pending` prints the `R__definer_lane_provisioning` patch on a 0.6.13 database; with `--extract-to`, apply `R__definer_lane_provisioning.sql` as a repeatable script after `V40`. Run it as the schema owner. The script does not repair existing lane sequences, so then run `awa storage lane-sequences --repair` as the migration role itself. +`awa migrate --sql --pending` prints the `R__definer_lane_provisioning` patch on a 0.6.13 database. To write it as a file instead, run `awa migrate --extract-to --pending` (without `--sql`, which can't be combined with it) and apply `R__definer_lane_provisioning.sql` as a repeatable script after `V40`. Run it as the role that owns the existing `awa` functions, normally the schema owner; `CREATE OR REPLACE` fails with `must be owner of function` otherwise, and nothing is applied. The script does not repair existing lane sequences, so then run `awa storage lane-sequences --repair` as the migration role itself. + +Each `R__` patch carries the full lane-helper definitions, so an older patch file applied after this one reinstalls the 0.6.13 helpers: + +- Replace every `R__` file with the set the 0.6.14 CLI extracts, together, rather than adding `R__definer_lane_provisioning.sql` beside older ones. +- Don't roll the runner's SQL files back to 0.6.13 when rolling binaries back. Leave the patch in place, as above. +- Keep the `V` files exactly as originally applied. Extract with `--pending`; a full `--extract-to` rewrites `V23`, which fails a checksum-validating runner. ## Roll back diff --git a/scripts/persona-soak-against.sh b/scripts/persona-soak-against.sh new file mode 100755 index 000000000..ef6838238 --- /dev/null +++ b/scripts/persona-soak-against.sh @@ -0,0 +1,37 @@ +#!/usr/bin/env bash +# Runs the persona soak against another revision of awa. +# +# scripts/persona-soak-against.sh [engine ...] +# +# The harness lives entirely in awa/tests, so any revision that builds with the same public +# API can be soaked: a detached worktree of is created, this checkout's harness +# files (including uncommitted changes) are copied over it, scripts/persona-soak.sh runs +# there with the arguments after the ref, and the worktree is removed afterwards. Results go +# to this checkout's soak-results/- unless SOAK_RESULTS_DIR is set. +# +# Environment is passed through to scripts/persona-soak.sh. Build outputs go to +# CARGO_TARGET_DIR when it is set, so builds of several revisions share one cache; the +# harness files are copied with a fresh modification time so cargo rebuilds the test binary +# rather than reusing one built from another revision. +set -euo pipefail + +if [ $# -lt 1 ]; then + echo "usage: $0 [engine ...]" >&2 + exit 2 +fi +ref="$1" +shift + +here="$(cd "$(dirname "$0")/.." && pwd)" +label="$(printf '%s' "$ref" | tr -c 'A-Za-z0-9._-' '_')" +export SOAK_RESULTS_DIR="${SOAK_RESULTS_DIR:-$here/soak-results/$(date +%Y%m%d-%H%M%S)-$label}" +worktree="$(mktemp -d "${TMPDIR:-/tmp}/awa-soak-against.XXXXXX")" +cleanup() { git -C "$here" worktree remove --force "$worktree" >/dev/null 2>&1 || rm -rf "$worktree"; } +trap cleanup EXIT + +git -C "$here" worktree add --detach "$worktree" "$ref" >/dev/null +tar -C "$here" -c awa/tests/persona_soak awa/tests/persona_soak_test.rs awa/tests/ci_timing.rs \ + scripts/persona-soak.sh \ + | tar -C "$worktree" -x -m +echo "== persona soak against $ref ($(git -C "$worktree" rev-parse --short=10 HEAD))" +(cd "$worktree" && scripts/persona-soak.sh "$@") diff --git a/scripts/persona-soak.sh b/scripts/persona-soak.sh new file mode 100755 index 000000000..b1f6a079c --- /dev/null +++ b/scripts/persona-soak.sh @@ -0,0 +1,213 @@ +#!/usr/bin/env bash +# Runs the persona soak (awa/tests/persona_soak_test.rs) against a matrix of +# throwaway PostgreSQL containers, one engine after another. +# +# scripts/persona-soak.sh [engine ...] +# +# Engines: pg17 (postgres:17), pg18 (postgres:18) and pg18-cassert (an +# assert-enabled, -O0 build of PostgreSQL 18; set SOAK_CASSERT_IMAGE). With no +# arguments all three run in that order. +# +# Environment: +# AWA_SOAK_SECONDS load phase length per engine (default 900) +# AWA_SOAK_PERSONAS comma list, default all of mia,tom,priya,sam,alex,chaos, +# watchers,support,uiops,impatient,python,fanout,shardops,abandon +# AWA_SOAK_SEED fixes the workload mix and chaos schedule +# AWA_SOAK_SCALE multiplies rates and queue counts (default 1) +# AWA_SOAK_SHARDOPS_SCALE multiplies the shardops persona's rates on its own (default 1) +# AWA_SOAK_SHARDOPS_SCHEMA_CREATE +# 1 (default) lets the shardops application roles create objects in +# the schema, so lanes are created on the fly; 0 withholds schema +# CREATE and has the operator provision lanes before raising a count +# AWA_SOAK_ABANDON_CLAIMS claims held by each killed process of the abandon persona (default 4000; +# 25000 or more reproduces a large abandoned-claim backlog) +# AWA_SOAK_ABANDON_KILLS, AWA_SOAK_ABANDON_KILL_SECONDS +# how many processes are SIGKILLed (default 4) and the seconds between +# kills (default 180) +# AWA_SOAK_ABANDON_BOUND_SECONDS +# seconds after a kill within which the rescue backlog must clear and the +# claim ring must be free to rotate (default 120) +# AWA_SOAK_RESCUE_INTERVAL_SECONDS +# heartbeat and deadline rescue interval of every runtime (default: the +# harness's 3 s and 2 s; the library default for heartbeat rescue is 30 s) +# AWA_SOAK_HEARTBEAT_SECONDS, AWA_SOAK_STALENESS_SECONDS +# heartbeat interval and heartbeat staleness of every runtime (default 2 s +# and 12 s, so rescues happen often; the library's are 30 s and 90 s) +# SOAK_RESULTS_DIR output directory (default ./soak-results/) +# SOAK_BASE_PORT first host port; engines use consecutive ports (default 15471) +# SOAK_MAX_CONNECTIONS server max_connections (default 400) +# SOAK_PG_CPUS CPU limit for each database container (default: none) +# SOAK_CASSERT_IMAGE image for pg18-cassert (default awa-pg18-cassert:18.6-O0; the +# -epqlog5 build, which logs every partition-pruning setup with its +# query text, is opt-in for diagnosing pruning crashes) +# SOAK_SKIP_PREPARE set to 1 to skip building the awa binary and the Python bindings +# (neither is needed without the UI and python personas) +# SOAK_TEST_BIN a prebuilt harness binary to run instead of `cargo test` +# CARGO_TARGET_DIR cargo target directory, as for any cargo build +# +# The admin UI personas drive the real `awa serve` binary and the python persona runs +# the awa-python bindings in a separate process; both are built once, before the first +# engine, with `cargo build --release -p awa-cli` and `uv` + `maturin develop --release` +# into awa-python/.venv. +# +# Every container is private to the run (name awa-soak--) and is +# removed on exit, including on failure or interruption. Containers that already +# exist are never touched. +set -uo pipefail + +cd "$(dirname "$0")/.." + +SECONDS_PER_ENGINE="${AWA_SOAK_SECONDS:-900}" +BASE_PORT="${SOAK_BASE_PORT:-15471}" +MAX_CONNECTIONS="${SOAK_MAX_CONNECTIONS:-400}" +CASSERT_IMAGE="${SOAK_CASSERT_IMAGE:-awa-pg18-cassert:18.6-O0}" +RESULTS_DIR="${SOAK_RESULTS_DIR:-$PWD/soak-results/$(date +%Y%m%d-%H%M%S)}" +ENGINES=("$@") +if [ ${#ENGINES[@]} -eq 0 ]; then + ENGINES=(pg17 pg18 pg18-cassert) +fi + +mkdir -p "$RESULTS_DIR" + +prepare() { + echo "== building the awa binary and the Python bindings" + cargo build --release -p awa-cli || return 1 + if [ ! -x awa-python/.venv/bin/python ]; then + uv venv --quiet awa-python/.venv || return 1 + uv pip install --quiet --python awa-python/.venv maturin || return 1 + fi + (cd awa-python && VIRTUAL_ENV="$PWD/.venv" PATH="$PWD/.venv/bin:$PATH" maturin develop --release) || return 1 +} +if [ "${SOAK_SKIP_PREPARE:-0}" != 1 ]; then + prepare >"$RESULTS_DIR/prepare.log" 2>&1 || { echo "prepare failed; see $RESULTS_DIR/prepare.log" >&2; exit 2; } +fi +export AWA_SOAK_PYTHON="${AWA_SOAK_PYTHON:-$PWD/awa-python/.venv/bin/python}" +export AWA_SOAK_AWA_BIN="${AWA_SOAK_AWA_BIN:-${CARGO_TARGET_DIR:-$PWD/target}/release/awa}" +CONTAINERS=() +STATS_PID="" + +cleanup() { + [ -n "$STATS_PID" ] && kill "$STATS_PID" 2>/dev/null + for container in "${CONTAINERS[@]:-}"; do + [ -n "$container" ] && docker rm -f -v "$container" >/dev/null 2>&1 + done +} +trap cleanup EXIT +trap 'exit 130' INT TERM + +start_container() { + local engine="$1" name="$2" port="$3" + local cpu_args=() + [ -n "${SOAK_PG_CPUS:-}" ] && cpu_args=(--cpus "$SOAK_PG_CPUS") + local common=(-c "max_connections=$MAX_CONNECTIONS" -c max_locks_per_transaction=256 -c shared_buffers=512MB -c "log_line_prefix=%m [%p] %u@%a ") + case "$engine" in + pg17) + docker run -d --name "$name" --shm-size=1g "${cpu_args[@]}" -e POSTGRES_PASSWORD=test -p "$port:5432" postgres:17 "${common[@]}" >/dev/null ;; + pg18) + docker run -d --name "$name" --shm-size=1g "${cpu_args[@]}" -e POSTGRES_PASSWORD=test -p "$port:5432" postgres:18 "${common[@]}" >/dev/null ;; + pg18-cassert) + # The image ships an initialised, trust-authenticated cluster in /data. Parallel query is + # off: at -O0 a handful of admin reads would otherwise take every core. + local pgargs='-c listen_addresses=0.0.0.0 -c max_connections='"$MAX_CONNECTIONS"' -c max_locks_per_transaction=256 -c shared_buffers=512MB -c track_activity_query_size=65536 -c max_parallel_workers_per_gather=0 -c logging_collector=on -c log_directory=/data/log -c log_filename=pg.log -c log_rotation_size=0' + case "$CASSERT_IMAGE" in + *epqlog*) + # The instrumented build logs an AWAEPQ line, with the full query text, for every + # partition-pruning setup: tens of gigabytes per hour under this workload. The log + # collector writes into a FIFO whose reader drops those messages and keeps the rest + # (assertion failures, crash notices, errors) in /data/log/pg-filtered.log. + SOAK_PG_LOG_FILE=/data/log/pg-filtered.log + docker run -d --name "$name" --shm-size=1g "${cpu_args[@]}" -p "$port:5432" "$CASSERT_IMAGE" sh -c ' + mkdir -p /data/log + rm -f /data/log/pg.log /data/log/pg-filtered.log + mkfifo /data/log/pg.log + awk -W interactive " + /^[0-9][0-9][0-9][0-9]-[0-9][0-9]-[0-9][0-9] /{ + pid=\"\"; if (match(\$0,/\\[[0-9]+\\]/)) pid=substr(\$0,RSTART,RLENGTH) + if (\$0 ~ /AWAEPQ/) {skip=1; epq[pid]=1} + else if (\$0 ~ /STATEMENT:/ && (pid in epq)) {skip=1; delete epq[pid]} + else {skip=0; delete epq[pid]} + } + !skip{print; fflush()}" \ + >/data/log/pg-filtered.log & + exec postgres -D /data '"$pgargs"'' >/dev/null ;; + *) + SOAK_PG_LOG_FILE=/data/log/pg.log + docker run -d --name "$name" --shm-size=1g "${cpu_args[@]}" -p "$port:5432" "$CASSERT_IMAGE" \ + sh -c 'mkdir -p /data/log && exec postgres -D /data '"$pgargs" >/dev/null ;; + esac ;; + *) + echo "unknown engine: $engine (expected pg17, pg18 or pg18-cassert)" >&2 + return 1 ;; + esac +} + +wait_ready() { + local name="$1" deadline=$((SECONDS + 120)) + until docker exec "$name" psql -U postgres -d postgres -Atc 'select 1' >/dev/null 2>&1; do + if [ "$SECONDS" -ge "$deadline" ]; then + echo "container $name did not become ready" >&2 + docker logs --tail 40 "$name" >&2 + return 1 + fi + sleep 1 + done + sleep 2 +} + +sample_stats() { + local name="$1" out="$2" + echo "time,cpu_percent,mem" >"$out" + while true; do + docker stats --no-stream --format '{{.CPUPerc}},{{.MemUsage}}' "$name" 2>/dev/null | + sed "s/^/$(date +%H:%M:%S),/" >>"$out" + sleep 15 + done +} + +if [ -n "${SOAK_TEST_BIN:-}" ]; then + harness_cmd=("$SOAK_TEST_BIN" persona_soak --ignored --nocapture --exact) +else + harness_cmd=(cargo test --release -p awa --test persona_soak_test -- --ignored --nocapture) +fi + +overall=0 +for index in "${!ENGINES[@]}"; do + engine="${ENGINES[$index]}" + port=$((BASE_PORT + index)) + name="awa-soak-$engine-$$" + out="$RESULTS_DIR/$engine" + mkdir -p "$out" + echo "== $engine on port $port (container $name)" + start_container "$engine" "$name" "$port" || exit 2 + CONTAINERS+=("$name") + wait_ready "$name" || exit 2 + + log_env=(AWA_SOAK_PG_CONTAINER="$name") + [ "$engine" = pg18-cassert ] && log_env+=(AWA_SOAK_PG_CONTAINER_LOG="$SOAK_PG_LOG_FILE") + + sample_stats "$name" "$out/container-stats.csv" & + STATS_PID=$! + + env AWA_SOAK_DATABASE_URL="postgres://postgres:test@localhost:$port/awa_soak_${engine//-/_}" \ + AWA_SOAK_SECONDS="$SECONDS_PER_ENGINE" \ + AWA_SOAK_REPORT_DIR="$out" \ + "${log_env[@]}" \ + "${harness_cmd[@]}" >"$out/run.log" 2>&1 + status=$? + + kill "$STATS_PID" 2>/dev/null + STATS_PID="" + if [ $status -ne 0 ]; then + overall=1 + echo " FAILED (exit $status); see $out/report.md and $out/reproducer.md" + # Keep the server log with the reproducer when the run failed. + docker logs "$name" >"$out/server.log" 2>&1 || true + [ "$engine" = pg18-cassert ] && docker exec "$name" cat "$SOAK_PG_LOG_FILE" >"$out/server.log" 2>/dev/null + else + echo " passed; report in $out/report.md" + fi + docker rm -f -v "$name" >/dev/null 2>&1 +done + +echo "results: $RESULTS_DIR" +exit $overall