diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 12575c1be..387f924b9 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -5,6 +5,21 @@ version: 2 updates: + - package-ecosystem: "cargo" + directory: "/rust" + cooldown: + default-days: 7 + groups: + rust-dependencies: + update-types: + - "minor" + - "patch" + schedule: + interval: "weekly" + - package-ecosystem: "github-actions" + directory: "/" + schedule: + interval: "weekly" - package-ecosystem: "gomod" directories: - "**/*" diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index c0d5af86e..415fb799b 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -418,6 +418,10 @@ jobs: - name: Checkout uses: actions/checkout@v6 + - uses: actions/setup-go@v6 + with: + go-version: "1.27" + - name: Setup sqlc uses: sqlc-dev/setup-sqlc@v5 with: @@ -428,6 +432,9 @@ jobs: echo "Make sure migration directories are the same" make verify/migrations + - name: Verify Rust migrations match + run: make verify/rust-migrations + - name: Verify sqlc generated code run: | echo "Make sure that all sqlc changes are checked in" diff --git a/.github/workflows/rust.yaml b/.github/workflows/rust.yaml new file mode 100644 index 000000000..5d8325749 --- /dev/null +++ b/.github/workflows/rust.yaml @@ -0,0 +1,156 @@ +name: Rust + +on: + # Filter the whole workflow so unrelated changes don't create skipped + # matrix jobs. Keep both events' paths in sync. + push: + branches: + - master + paths: + - ".github/workflows/rust.yaml" + - "Makefile" + - "internal/cmd/syncrustmigrations/**" + - "riverdriver/riverpgxv5/migration/**" + - "riverdriver/riversqlite/migration/**" + - "rust/**" + pull_request: + paths: + - ".github/workflows/rust.yaml" + - "Makefile" + - "internal/cmd/syncrustmigrations/**" + - "riverdriver/riverpgxv5/migration/**" + - "riverdriver/riversqlite/migration/**" + - "rust/**" + +permissions: + contents: read + +env: + # Keep the cross-run cache small; Cargo still reuses compiled dependencies. + CARGO_INCREMENTAL: "0" + +jobs: + quality: + runs-on: ubuntu-latest + timeout-minutes: 20 + + steps: + - uses: actions/checkout@v6 + with: + fetch-depth: 0 + - uses: dtolnay/rust-toolchain@stable + id: rust + with: + components: clippy,rustfmt + - name: Cache Rust dependencies and build artifacts + uses: actions/cache@v5 + with: + path: | + ~/.cargo/registry/index + ~/.cargo/registry/cache + ~/.cargo/git/db + rust/target + key: rust-v1-${{ runner.os }}-${{ runner.arch }}-${{ github.job }}-${{ steps.rust.outputs.cachekey }}-${{ hashFiles('rust/**/Cargo.toml', 'rust/Cargo.lock') }} + restore-keys: rust-v1-${{ runner.os }}-${{ runner.arch }}-${{ github.job }}-${{ steps.rust.outputs.cachekey }}- + - uses: taiki-e/install-action@v2 + with: + tool: cargo-deny@0.20.2 + - uses: taiki-e/install-action@v2 + with: + tool: cargo-semver-checks@0.49.0 + + - name: Lint, including PostgreSQL-only and SQLite-only builds + run: make lint/rust + + - name: Documentation and examples + run: make doc/rust + + - name: Dependency and license policy + run: make check/rust/dependencies + + - name: Package archives + run: make check/rust/package + + # Compare with the latest published Rust release tag, which the full + # checkout above includes. Before the first release there is no + # baseline and the step reports that instead of failing. + - name: Public API compatibility + run: make check/rust/semver + + # docs.rs builds with a nightly toolchain and `--cfg docsrs`, which + # enables the crates' `doc_cfg` feature badges. Last, since installing + # nightly makes it the default toolchain for later steps. + - uses: dtolnay/rust-toolchain@nightly + - name: Documentation as docs.rs builds it + run: make doc/rust/docsrs + + rust_versions: + runs-on: ubuntu-latest + timeout-minutes: 20 + strategy: + fail-fast: false + matrix: + rust-version: ["1.95", "1.96", "1.97"] + + steps: + - uses: actions/checkout@v6 + - uses: dtolnay/rust-toolchain@master + id: rust + with: + toolchain: ${{ matrix.rust-version }} + - name: Cache Rust dependencies and build artifacts + uses: actions/cache@v5 + with: + path: | + ~/.cargo/registry/index + ~/.cargo/registry/cache + ~/.cargo/git/db + rust/target + key: rust-v1-${{ runner.os }}-${{ runner.arch }}-${{ github.job }}-${{ steps.rust.outputs.cachekey }}-${{ hashFiles('rust/**/Cargo.toml', 'rust/Cargo.lock') }} + restore-keys: rust-v1-${{ runner.os }}-${{ runner.arch }}-${{ github.job }}-${{ steps.rust.outputs.cachekey }}- + - name: Check every target and feature + run: cargo check --manifest-path rust/Cargo.toml --workspace --all-targets --all-features --locked + - name: Unit, doc, and SQLite tests + run: make test/rust/sqlite + + postgres: + runs-on: ubuntu-latest + timeout-minutes: 20 + strategy: + fail-fast: false + matrix: + postgres-version: [14, 15, 16, 17, 18] + env: + RIVER_RUST_DATABASE_URL: postgres://postgres:postgres@localhost:5432/river_rust_test?sslmode=disable + + services: + postgres: + image: postgres:${{ matrix.postgres-version }} + env: + POSTGRES_PASSWORD: postgres + options: >- + --health-cmd pg_isready + --health-interval 2s + --health-timeout 5s + --health-retries 5 + ports: + - 5432:5432 + + steps: + - uses: actions/checkout@v6 + - uses: dtolnay/rust-toolchain@stable + id: rust + - name: Cache Rust dependencies and build artifacts + uses: actions/cache@v5 + with: + path: | + ~/.cargo/registry/index + ~/.cargo/registry/cache + ~/.cargo/git/db + rust/target + key: rust-v1-${{ runner.os }}-${{ runner.arch }}-${{ github.job }}-${{ steps.rust.outputs.cachekey }}-${{ hashFiles('rust/**/Cargo.toml', 'rust/Cargo.lock') }} + restore-keys: rust-v1-${{ runner.os }}-${{ runner.arch }}-${{ github.job }}-${{ steps.rust.outputs.cachekey }}- + - name: Create test database + run: PGPASSWORD=postgres createdb -h localhost -U postgres river_rust_test + - name: Rust unit, SQLite, and PostgreSQL tests + run: make test/rust diff --git a/.gitignore b/.gitignore index 94b880868..d09a4cac6 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,4 @@ /river /riverdriver/riverdrivertest/example_libsql_test.libsql /sqlite/ +/rust/**/target/ diff --git a/Makefile b/Makefile index 58cb24008..a55a88ee5 100644 --- a/Makefile +++ b/Makefile @@ -1,5 +1,7 @@ .DEFAULT_GOAL := help +SQLC ?= sqlc + .PHONY: db/reset db/reset: ## Drop, create, and migrate dev and test databases db/reset: db/reset/dev @@ -18,17 +20,22 @@ db/reset/test: ## Drop, create, and migrate test databases .PHONY: generate generate: ## Generate generated artifacts generate: generate/migrations +generate: generate/rust-migrations generate: generate/sqlc .PHONY: generate/migrations generate/migrations: ## Sync changes of pgxv5 migrations to database/sql rsync -au --delete "riverdriver/riverpgxv5/migration/" "riverdriver/riverdatabasesql/migration/" +.PHONY: generate/rust-migrations +generate/rust-migrations: ## Sync database migrations to Rust + go run ./internal/cmd/syncrustmigrations + .PHONY: generate/sqlc generate/sqlc: ## Generate sqlc - cd riverdriver/riverdatabasesql/internal/dbsqlc && sqlc generate - cd riverdriver/riverpgxv5/internal/dbsqlc && sqlc generate - cd riverdriver/riversqlite/internal/dbsqlc && sqlc generate + cd riverdriver/riverdatabasesql/internal/dbsqlc && $(SQLC) generate + cd riverdriver/riverpgxv5/internal/dbsqlc && $(SQLC) generate + cd riverdriver/riversqlite/internal/dbsqlc && $(SQLC) generate # Looks at comments using ## on targets and uses them to produce a help output. .PHONY: help @@ -42,6 +49,8 @@ help: ## Print this message submodules := $(shell go list -f '{{.Dir}}' -m) ITERATIONS ?= 100 +RUST_BENCH_ARGS ?= +RUST_SEMVER_BASELINE_REV ?= $(shell git tag --list 'riverqueue-v*' --sort=-v:refname | head -n 1) TEST_DATABASE ?= all @@ -72,6 +81,16 @@ define lint-target endef $(foreach mod,$(submodules),$(eval $(call lint-target,$(mod)))) +# Rust targets are separate from `lint` and `test` so Go-only contributors and +# the Go CI jobs do not need a Rust toolchain; the Rust workflow runs them. +.PHONY: lint/rust +lint/rust: ## Run Rust formatting and clippy checks, including single-backend builds + cd rust && cargo fmt --all -- --check + cd rust && cargo clippy --workspace --all-targets --all-features --locked -- -D warnings + cd rust && cargo clippy -p riverqueue -p riverqueue-migrate -p riverqueue-cli -p riverqueue-test --no-default-features --features postgres --all-targets --locked -- -D warnings + cd rust && cargo clippy -p riverqueue -p riverqueue-migrate -p riverqueue-cli -p riverqueue-test --no-default-features --features sqlite --all-targets --locked -- -D warnings + cd rust && $(RUST_POSTGRES_TESTS_ENV) cargo clippy -p riverqueue -p riverqueue-migrate --all-targets --all-features --locked -- -D warnings + .PHONY: test test:: ## Run tests (TEST_DATABASE=all, postgres, or sqlite) define test-target @@ -85,6 +104,74 @@ ifneq ($(TEST_DATABASE),sqlite) test:: ; cd ./riverdriver/riverdrivertest && RIVER_USE_LEGACY_SUBTRANSACTIONS=1 go test . -run '^TestDriverRiverPgxV5$$/.*/WithTx$$' -timeout 2m endif +# `--cfg river_postgres_tests` builds the Rust PostgreSQL integration tests. +# It goes to both rustc and rustdoc so any doctest gated on it runs too, and +# into its own target directory so switching it on and off doesn't rebuild +# the ordinary build's artifacts. The default is absolute: trybuild resolves a +# relative target directory from the macros crate's directory. +RUST_POSTGRES_TESTS_ENV = RUSTFLAGS="$$RUSTFLAGS --cfg river_postgres_tests" \ + RUSTDOCFLAGS="$$RUSTDOCFLAGS --cfg river_postgres_tests" \ + CARGO_TARGET_DIR="$${CARGO_TARGET_DIR:-$(CURDIR)/rust/target}/postgres-tests" + +# PostgreSQL integration tests need RIVER_RUST_DATABASE_URL. Without it +# test/rust still runs unit, doc, and SQLite integration tests, and fails in CI +# so a missing URL cannot turn the PostgreSQL suite into a silent pass. +.PHONY: test/rust +test/rust: ## Run Rust unit and SQLite tests, plus PostgreSQL tests when RIVER_RUST_DATABASE_URL is set + @if [ -n "$$RIVER_RUST_DATABASE_URL" ]; then \ + cd rust && $(RUST_POSTGRES_TESTS_ENV) cargo test --workspace --all-features --locked; \ + elif [ -n "$$CI" ]; then \ + echo "RIVER_RUST_DATABASE_URL is required in CI to run the Rust PostgreSQL tests" >&2; exit 1; \ + else \ + echo "RIVER_RUST_DATABASE_URL is unset; skipping Rust PostgreSQL integration tests"; \ + cd rust && cargo test --workspace --features riverqueue/sqlite,riverqueue-migrate/sqlite --locked; \ + fi + +.PHONY: test/rust/postgres +test/rust/postgres: ## Run all Rust tests, including PostgreSQL integration tests (requires RIVER_RUST_DATABASE_URL) + @test -n "$$RIVER_RUST_DATABASE_URL" || { echo "RIVER_RUST_DATABASE_URL is required" >&2; exit 1; } + cd rust && $(RUST_POSTGRES_TESTS_ENV) cargo test --workspace --all-features --locked + +.PHONY: test/rust/sqlite +test/rust/sqlite: ## Run Rust unit, doc, and SQLite integration tests without a PostgreSQL database + cd rust && cargo test --workspace --features riverqueue/sqlite,riverqueue-migrate/sqlite --locked + +.PHONY: doc/rust +doc/rust: ## Build Rust API documentation, compiled examples, and doctests for each backend feature set + cd rust && RUSTDOCFLAGS="-D warnings" cargo doc --workspace --all-features --no-deps --locked + cd rust && RUSTDOCFLAGS="-D warnings" cargo test --workspace --all-features --doc --locked + cd rust && RUSTDOCFLAGS="-D warnings" cargo test -p riverqueue -p riverqueue-migrate -p riverqueue-cli -p riverqueue-test --no-default-features --features postgres --doc --locked + cd rust && RUSTDOCFLAGS="-D warnings" cargo test -p riverqueue -p riverqueue-migrate -p riverqueue-cli -p riverqueue-test --no-default-features --features sqlite --doc --locked + cd rust && cargo check --workspace --examples --all-features --locked + +.PHONY: doc/rust/docsrs +doc/rust/docsrs: ## Build Rust API documentation as docs.rs does (nightly toolchain, `--cfg docsrs`) + cd rust && RUSTDOCFLAGS="--cfg docsrs -D warnings" CARGO_TARGET_DIR="$${CARGO_TARGET_DIR:-target}/docsrs" cargo +nightly doc -p riverqueue -p riverqueue-migrate -p riverqueue-test --all-features --no-deps --locked + +.PHONY: check/rust/dependencies +check/rust/dependencies: ## Audit Rust advisories, licenses, bans, and sources + cd rust && cargo deny check + +.PHONY: check/rust/package +check/rust/package: ## Build and verify publishable crate archives without publishing + cd rust && cargo package --workspace --allow-dirty --locked + +# The baseline is the latest published riverqueue-v* tag, and +# cargo-semver-checks infers the allowed change from the version bump. It +# skips every lint while the workspace version is a pre-release, so +# comparing unreleased revisions with each other checks nothing. Until a +# Rust release is tagged the check reports that there is no baseline. Set +# RUST_SEMVER_BASELINE_REV to compare with another revision. +.PHONY: check/rust/semver +check/rust/semver: ## Check Rust APIs against RUST_SEMVER_BASELINE_REV (default: latest Rust tag) + @if test -z "$(RUST_SEMVER_BASELINE_REV)"; then \ + echo "No published Rust release tag (riverqueue-v*); no public API baseline to compare"; \ + elif ! git cat-file -e "$(RUST_SEMVER_BASELINE_REV):rust/Cargo.toml" 2>/dev/null; then \ + echo "Baseline $(RUST_SEMVER_BASELINE_REV) predates the Rust crates; no public API to compare"; \ + else \ + cd rust && cargo semver-checks --workspace --baseline-rev "$(RUST_SEMVER_BASELINE_REV)"; \ + fi + .PHONY: test/race test/race:: ## Run tests with race detector (TEST_DATABASE=all, postgres, or sqlite) define test-race-target @@ -104,6 +191,10 @@ define bench-target endef $(foreach mod,$(submodules),$(eval $(call bench-target,$(mod)))) +.PHONY: bench/rust +bench/rust: ## Run the destructive Rust PostgreSQL throughput benchmark + cd rust && cargo run --release --locked -p riverqueue-cli --bin riverqueue -- bench $(if $(DATABASE_URL),--database-url "$(DATABASE_URL)") $(RUST_BENCH_ARGS) + .PHONY: tidy tidy:: ## Run `go mod tidy` for all submodules define tidy-target @@ -122,14 +213,19 @@ update-mod-version: ## Update River packages in all submodules to $VERSION .PHONY: verify verify: ## Verify generated artifacts verify: verify/migrations +verify: verify/rust-migrations verify: verify/sqlc .PHONY: verify/migrations verify/migrations: ## Verify synced migrations diff -qr riverdriver/riverpgxv5/migration riverdriver/riverdatabasesql/migration +.PHONY: verify/rust-migrations +verify/rust-migrations: ## Verify Rust migrations match the canonical migrations + go run ./internal/cmd/syncrustmigrations -check + .PHONY: verify/sqlc verify/sqlc: ## Verify generated sqlc - cd riverdriver/riverdatabasesql/internal/dbsqlc && sqlc diff - cd riverdriver/riverpgxv5/internal/dbsqlc && sqlc diff - cd riverdriver/riversqlite/internal/dbsqlc && sqlc diff + cd riverdriver/riverdatabasesql/internal/dbsqlc && $(SQLC) diff + cd riverdriver/riverpgxv5/internal/dbsqlc && $(SQLC) diff + cd riverdriver/riversqlite/internal/dbsqlc && $(SQLC) diff diff --git a/internal/cmd/syncrustmigrations/main.go b/internal/cmd/syncrustmigrations/main.go new file mode 100644 index 000000000..576738863 --- /dev/null +++ b/internal/cmd/syncrustmigrations/main.go @@ -0,0 +1,111 @@ +// Command syncrustmigrations mirrors River's canonical database migrations +// into the publishable Rust migration crate. +package main + +import ( + "bytes" + "flag" + "fmt" + "os" + "path/filepath" + "slices" + "strings" +) + +type database struct { + canonicalDir string + mirrorDir string +} + +func main() { + check := flag.Bool("check", false, "check generated files without writing") + flag.Parse() + + databases := []database{ + { + canonicalDir: "riverdriver/riverpgxv5/migration/main", + mirrorDir: "rust/riverqueue-migrate/migrations/main", + }, + { + canonicalDir: "riverdriver/riversqlite/migration/main", + mirrorDir: "rust/riverqueue-migrate/migrations/sqlite/main", + }, + } + for _, database := range databases { + syncDatabase(database, *check) + } +} + +func syncDatabase(database database, check bool) { + entries, err := os.ReadDir(database.canonicalDir) + if err != nil { + fatal(err) + } + + var names []string + for _, entry := range entries { + if !entry.IsDir() && strings.HasSuffix(entry.Name(), ".sql") { + names = append(names, entry.Name()) + } + } + slices.Sort(names) + + for _, name := range names { + sourcePath := filepath.Join(database.canonicalDir, name) + contents, err := os.ReadFile(sourcePath) + if err != nil { + fatal(err) + } + mirrorPath := filepath.Join(database.mirrorDir, name) + if check { + checkFile(mirrorPath, contents) + } else { + writeFile(mirrorPath, contents) + } + } + removeStaleMirrors(database.mirrorDir, names, check) +} + +func removeStaleMirrors(directory string, expected []string, check bool) { + entries, err := os.ReadDir(directory) + if err != nil { + fatal(err) + } + for _, entry := range entries { + if entry.IsDir() || !strings.HasSuffix(entry.Name(), ".sql") || slices.Contains(expected, entry.Name()) { + continue + } + path := filepath.Join(directory, entry.Name()) + if check { + fatal(fmt.Errorf("generated file is stale: %s (run make generate/rust-migrations)", path)) + } + if err := os.Remove(path); err != nil { + fatal(err) + } + } +} + +func checkFile(path string, expected []byte) { + actual, err := os.ReadFile(path) + if err != nil { + fatal(fmt.Errorf("read generated file %s: %w", path, err)) + } + if !bytes.Equal(actual, expected) { + fatal(fmt.Errorf("generated file is stale: %s (run make generate/rust-migrations)", path)) + } +} + +func fatal(err error) { + fmt.Fprintln(os.Stderr, err) + os.Exit(1) +} + +func writeFile(path string, contents []byte) { + if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { + fatal(err) + } + //nolint:gosec // Generated repository artifacts are intentionally world-readable. + if err := os.WriteFile(path, contents, 0o644); err != nil { + fatal(err) + } +} diff --git a/rust/CHANGELOG.md b/rust/CHANGELOG.md new file mode 100644 index 000000000..83a6a32e5 --- /dev/null +++ b/rust/CHANGELOG.md @@ -0,0 +1,32 @@ +# Changelog + +All notable changes to River's Rust crates are documented in this file. The +workspace crates (`riverqueue`, `riverqueue-macros`, `riverqueue-migrate`, +`riverqueue-cli`, and `riverqueue-test`) are versioned and released together. + +The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), +and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +Changes to River for Go are recorded in the [repository changelog](../CHANGELOG.md). + +## [Unreleased] + +### Added + +- First preview release of River for Rust. `riverqueue` provides a typed, + Tokio-based client for PostgreSQL (through SQLx) and SQLite that shares + River's database schema and job protocol with River for Go, so Rust and Go + clients can insert and work jobs in the same database. It includes typed + workers, transactional inserts and completion, unique, scheduled, periodic, + and resumable jobs, queue management, job cancellation, events, hooks, + middleware, leader election, and maintenance services. +- Requests run with `.tx(...)` use the caller's transaction directly, without + a savepoint or nested transaction, like River for Go's `*Tx` methods. A + request that returns an error may leave partial writes in that transaction, + so roll it back, or open your own savepoint around the request to continue. +- `riverqueue-macros` provides `#[derive(JobArgs)]`, including unique options. +- `riverqueue-migrate` applies and validates River's migration lines on + PostgreSQL and SQLite, sharing migration history with River for Go. +- `riverqueue-cli` installs the `riverqueue` command for migrations and + benchmarks. +- `riverqueue-test` provides fixtures, insertion assertions, and helpers for + running workers in tests. diff --git a/rust/Cargo.lock b/rust/Cargo.lock new file mode 100644 index 000000000..c61ec9e09 --- /dev/null +++ b/rust/Cargo.lock @@ -0,0 +1,2031 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + +[[package]] +name = "android_system_properties" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae221649c9976a6f6c56ae1facf410f3ddb33cc661c4b7b61020a912d4237fbc" +dependencies = [ + "libc", +] + +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + +[[package]] +name = "async-trait" +version = "0.1.92" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "atoi" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f28d99ec8bfea296261ca1af174f24225171fea9664ba9003cbebee704810528" +dependencies = [ + "num-traits", +] + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "bitflags" +version = "2.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b588b76d00fde79687d7646a9b5bdf3cc0f655e0bbd080335a95d7e96f3587da" +dependencies = [ + "serde_core", +] + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "block-buffer" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2f6c7dbe95a6ed67ad9f18e57daf93a2f034c524b99fd2b76d18fdfeb6660aa" +dependencies = [ + "hybrid-array", +] + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + +[[package]] +name = "bytes" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" + +[[package]] +name = "cc" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d262e149917187838d5b42777c8253bcb64500067342904e7d429499a6f277e" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "chacha20" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core", +] + +[[package]] +name = "chrono" +version = "0.4.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" +dependencies = [ + "iana-time-zone", + "js-sys", + "num-traits", + "serde", + "wasm-bindgen", + "windows-link", +] + +[[package]] +name = "chrono-tz" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6139a8597ed92cf816dfb33f5dd6cf0bb93a6adc938f11039f371bc5bcd26c3" +dependencies = [ + "chrono", + "phf", +] + +[[package]] +name = "cmov" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" + +[[package]] +name = "const-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + +[[package]] +name = "crc" +version = "3.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5eb8a2a1cd12ab0d987a5d5e825195d372001a4094a0376319d5a0ad71c1ba0d" +dependencies = [ + "crc-catalog", +] + +[[package]] +name = "crc-catalog" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" + +[[package]] +name = "crossbeam-queue" +version = "0.3.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "803d13fb3b09d88be9f4dbc29062c66b19bf7170867ceb746d2a8689bf6c7a26" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61803da095bee82a81bb1a452ecc25d3b2f1416d1897eb86430c6159ef717c17" + +[[package]] +name = "crypto-common" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "crypto-common" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" +dependencies = [ + "hybrid-array", +] + +[[package]] +name = "ctutils" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +dependencies = [ + "cmov", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer 0.10.4", + "crypto-common 0.1.6", +] + +[[package]] +name = "digest" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +dependencies = [ + "block-buffer 0.12.1", + "const-oid", + "crypto-common 0.2.2", + "ctutils", +] + +[[package]] +name = "displaydoc" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c6232dd377dcc64799954cbd3a9bb882e9cdc1308ccd87b1c098f1fb2eaf82a8" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "dotenvy" +version = "0.15.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" + +[[package]] +name = "either" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e5e8f6c15a24b9a3ee5efec809ccd006d3b30e8b3bb63c39af737c7f87daa1d" +dependencies = [ + "serde", +] + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "etcetera" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de48cc4d1c1d97a20fd819def54b890cadde72ed3ad0c614822a0a433361be96" +dependencies = [ + "cfg-if", + "windows-sys 0.61.2", +] + +[[package]] +name = "event-listener" +version = "5.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a23add41df1562121a9393cb065eab5146a1242410f23a644851e90cfd669d2" +dependencies = [ + "parking", + "pin-project-lite", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26b73573e6edcd2af0cdf47bd6cb58f0b3839491263c314eaad1ccf24430e1de" + +[[package]] +name = "flume" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e139bc46ca777eb5efaf62df0ab8cc5fd400866427e56c68b22e414e53bd3be" +dependencies = [ + "futures-core", + "futures-sink", + "spin", +] + +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "futures-channel" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f9e3d69d39e4862ffed03ed071a76f9a13ba1d9109d355b0f0aa6b15e393c4" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e" + +[[package]] +name = "futures-executor" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-intrusive" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d930c203dd0b6ff06e0201a4a2fe9149b43c684fd4420555b26d21b1a02956f" +dependencies = [ + "futures-core", + "lock_api", + "parking_lot", +] + +[[package]] +name = "futures-io" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed" + +[[package]] +name = "futures-macro" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9fb9654ba8355388abeb8dcb4fc62f511300867002afc858860463bdd9fe0c44" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "futures-sink" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1944426bf7d03f1d14f708785e4b33efd750b36d48a157b836b3efc15ede8e1d" + +[[package]] +name = "futures-task" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd417de3d1d015fc3bfd2b1ea46dfc7bab72ef86f1cc7cc9c78e728b34a6d1fd" + +[[package]] +name = "futures-util" +version = "0.3.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" +dependencies = [ + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4bb6743198531e02858aeaea5398fcc883e71851fcbcb5a2f773e2fb6cb1edf2" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + +[[package]] +name = "getrandom" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "rand_core", +] + +[[package]] +name = "glob" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "hashlink" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824e001ac4f3012dd16a264bec811403a67ca9deb6c102fc5049b32c4574b35f" +dependencies = [ + "hashbrown 0.16.1", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + +[[package]] +name = "hkdf" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4aaa26c720c68b866f2c96ef5c1264b3e6f473fe5d4ce61cd44bbe913e553018" +dependencies = [ + "hmac", +] + +[[package]] +name = "hmac" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6303bc9732ae41b04cb554b844a762b4115a61bfaa81e3e83050991eeb56863f" +dependencies = [ + "digest 0.11.3", +] + +[[package]] +name = "hybrid-array" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "707114b52a152fa7bdb290cd7cd5912d9467273b6d74e21b8d81aca1f8533f6b" +dependencies = [ + "typenum", +] + +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "js-sys" +version = "0.3.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e0c1080212aad755ea003d18543e8768dd432c48819efd73a7bf1e39b7a5a3a" +dependencies = [ + "cfg-if", + "futures-util", + "wasm-bindgen", +] + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "libc" +version = "0.2.189" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" + +[[package]] +name = "libsqlite3-sys" +version = "0.37.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1f111c8c41e7c61a49cd34e44c7619462967221a6443b0ec299e0ac30cfb9b1" +dependencies = [ + "cc", + "pkg-config", + "vcpkg", +] + +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + +[[package]] +name = "md-5" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98" +dependencies = [ + "cfg-if", + "digest 0.11.3", +] + +[[package]] +name = "memchr" +version = "2.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" + +[[package]] +name = "mio" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" +dependencies = [ + "libc", + "wasi", + "windows-sys 0.61.2", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "phf" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "913273894cec178f401a31ec4b656318d95473527be05c0752cc41cdc32be8b7" +dependencies = [ + "phf_shared", +] + +[[package]] +name = "phf_shared" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06005508882fb681fd97892ecff4b7fd0fee13ef1aa569f8695dae7ab9099981" +dependencies = [ + "siphasher", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + +[[package]] +name = "proc-macro2" +version = "1.0.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "985e7ec9bb745e6ce6535b544d84d6cd6f7ad8bd711c398938ae983b91a766d9" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rand" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" +dependencies = [ + "chacha20", + "getrandom 0.4.3", + "rand_core", +] + +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + +[[package]] +name = "riverqueue" +version = "0.50.0-alpha.1" +dependencies = [ + "anyhow", + "async-trait", + "base64", + "chrono", + "chrono-tz", + "futures-util", + "rand", + "riverqueue-macros", + "riverqueue-migrate", + "serde", + "serde_json", + "sha2 0.11.0", + "sqlx", + "thiserror", + "tokio", + "tokio-util", + "tracing", + "tracing-subscriber", +] + +[[package]] +name = "riverqueue-cli" +version = "0.50.0-alpha.1" +dependencies = [ + "riverqueue", + "riverqueue-migrate", + "serde", + "sqlx", + "tokio", + "tokio-util", +] + +[[package]] +name = "riverqueue-macros" +version = "0.50.0-alpha.1" +dependencies = [ + "proc-macro2", + "quote", + "riverqueue", + "serde", + "syn 2.0.119", + "trybuild", +] + +[[package]] +name = "riverqueue-migrate" +version = "0.50.0-alpha.1" +dependencies = [ + "serde_json", + "sqlx", + "thiserror", + "tokio", +] + +[[package]] +name = "riverqueue-test" +version = "0.50.0-alpha.1" +dependencies = [ + "chrono", + "riverqueue", + "riverqueue-migrate", + "serde", + "serde_json", + "sqlx", + "tokio", + "tokio-util", +] + +[[package]] +name = "rustls" +version = "0.23.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" +dependencies = [ + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-pki-types" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96" +dependencies = [ + "zeroize", +] + +[[package]] +name = "rustls-webpki" +version = "0.103.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + +[[package]] +name = "rustversion" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf54715a573b99ac80df0bc206da022bcd442c974952c7b9720069370852e21f" + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "serde" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.229" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "serde_json" +version = "1.0.151" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_spanned" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" +dependencies = [ + "serde_core", +] + +[[package]] +name = "sha1" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aacc4cc499359472b4abe1bf11d0b12e688af9a805fa5e3016f9a386dc2d0214" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures 0.2.17", + "digest 0.10.7", +] + +[[package]] +name = "sha2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", +] + +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + +[[package]] +name = "siphasher" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33f4fe9184a62d842c9ef383018f3306d8ba224fd9d836f56d7288308847c256" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" +dependencies = [ + "serde", +] + +[[package]] +name = "socket2" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "spin" +version = "0.9.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3763264f6b73151db08c50ff20d7d8a0b8796e021cdea7ceedad07b80155fa0e" +dependencies = [ + "lock_api", +] + +[[package]] +name = "sqlx" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "378620ccc25c62c89d8be1c819e76a88d59bdcc3304733330788948e619bfd71" +dependencies = [ + "sqlx-core", + "sqlx-macros", + "sqlx-mysql", + "sqlx-postgres", + "sqlx-sqlite", +] + +[[package]] +name = "sqlx-core" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05b44e85bf579a8eeb4ceaa77a3a523baf2bf0e9bac7e40f405d537b5d2d5ccb" +dependencies = [ + "base64", + "bytes", + "cfg-if", + "chrono", + "crc", + "crossbeam-queue", + "either", + "event-listener", + "futures-core", + "futures-intrusive", + "futures-io", + "futures-util", + "hashbrown 0.16.1", + "hashlink", + "indexmap", + "log", + "memchr", + "percent-encoding", + "rustls", + "serde", + "serde_json", + "sha2 0.10.9", + "smallvec", + "thiserror", + "tokio", + "tokio-stream", + "tracing", + "url", + "webpki-roots", +] + +[[package]] +name = "sqlx-macros" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd2b84f2bc39a5705ef27ec785a11c934a41bbd4a24941e257927cddc26b60bf" +dependencies = [ + "proc-macro2", + "quote", + "sqlx-core", + "sqlx-macros-core", + "syn 2.0.119", +] + +[[package]] +name = "sqlx-macros-core" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb8d96de5fdc85a5c4ec813432b523ec637e80ba98f046555f75f7908ddac7c3" +dependencies = [ + "cfg-if", + "dotenvy", + "either", + "heck", + "hex", + "proc-macro2", + "quote", + "serde", + "serde_json", + "sha2 0.10.9", + "sqlx-core", + "sqlx-mysql", + "sqlx-postgres", + "sqlx-sqlite", + "syn 2.0.119", + "tokio", + "url", +] + +[[package]] +name = "sqlx-mysql" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90b8020fe17c5f2c245bfa2505d7ef59c5604839527c740266ad2214acebea27" +dependencies = [ + "bitflags", + "byteorder", + "bytes", + "chrono", + "crc", + "digest 0.11.3", + "dotenvy", + "either", + "futures-core", + "futures-util", + "generic-array", + "log", + "percent-encoding", + "serde", + "sha1", + "sha2 0.11.0", + "sqlx-core", + "thiserror", + "tracing", +] + +[[package]] +name = "sqlx-postgres" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87a2bdd6e83f6b3ea525ca9fee568030508b58355a43d0b2c1674d5f79dcd65e" +dependencies = [ + "atoi", + "base64", + "bitflags", + "byteorder", + "chrono", + "crc", + "dotenvy", + "etcetera", + "futures-channel", + "futures-core", + "futures-util", + "hex", + "hkdf", + "hmac", + "itoa", + "log", + "md-5", + "memchr", + "rand", + "serde", + "serde_json", + "sha2 0.11.0", + "smallvec", + "sqlx-core", + "stringprep", + "thiserror", + "tracing", + "whoami", +] + +[[package]] +name = "sqlx-sqlite" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488e99c397a62007e4229aec669a179816339afc6d2620ca6fa420dbee2e982c" +dependencies = [ + "atoi", + "chrono", + "flume", + "form_urlencoded", + "futures-channel", + "futures-core", + "futures-executor", + "futures-intrusive", + "futures-util", + "libsqlite3-sys", + "log", + "percent-encoding", + "serde", + "sqlx-core", + "thiserror", + "tracing", + "url", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "stringprep" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b4df3d392d81bd458a8a621b8bffbd2302a12ffe288a9d931670948749463b1" +dependencies = [ + "unicode-bidi", + "unicode-normalization", + "unicode-properties", +] + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + +[[package]] +name = "syn" +version = "2.0.119" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "syn" +version = "3.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "target-tuple" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "876fef147edbcbddc8ac5cbbba92c7b86519e314e86638596c09673b2ed01e7f" + +[[package]] +name = "termcolor" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06794f8f6c5c898b3275aebefa6b8a1cb24cd2c6c79397ab15774837a0bc5755" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "thiserror" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "thread_local" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tinyvec" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb4ebadaa0af04fab11ae01eb5f9fdb5f9c5b875506e210e71c07873528baa7f" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + +[[package]] +name = "tokio" +version = "1.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" +dependencies = [ + "bytes", + "libc", + "mio", + "pin-project-lite", + "signal-hook-registry", + "socket2", + "tokio-macros", + "windows-sys 0.61.2", +] + +[[package]] +name = "tokio-macros" +version = "2.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e" +dependencies = [ + "proc-macro2", + "quote", + "syn 3.0.3", +] + +[[package]] +name = "tokio-stream" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a3d06f0b082ba57c26b79407372e57cf2a1e28124f78e9479fe80322cf53420b" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "tokio-util" +version = "0.7.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "494815d09bf52b5548659851081238f0ca39ff638363907596da739561c62c52" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "futures-util", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "toml" +version = "1.1.6+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "920602543f0911ab71da12c50d59701da54c196d1a2bf5cb4b75667f137a406a" +dependencies = [ + "indexmap", + "serde_core", + "serde_spanned", + "toml_datetime", + "toml_parser", + "toml_writer", + "winnow", +] + +[[package]] +name = "toml_datetime" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_parser" +version = "1.1.3+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56" +dependencies = [ + "winnow", +] + +[[package]] +name = "toml_writer" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d56353a2a665ad0f41a421187180aab746c8c325620617ad883a99a1cbe66d2" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "log", + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "sharded-slab", + "thread_local", + "tracing-core", +] + +[[package]] +name = "trybuild" +version = "1.0.121" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0cabaa10be1917331a313866bd94526343e03c77bcf69144b62b072ad35d47c" +dependencies = [ + "glob", + "serde", + "serde_derive", + "serde_json", + "target-tuple", + "termcolor", + "toml", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-bidi" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c1cb5db39152898a79168971543b1cb5020dff7fe43c8dc468b0885f5e29df5" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-normalization" +version = "0.1.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fd4f6878c9cb28d874b009da9e8d183b5abc80117c40bbd187a1fde336be6e8" +dependencies = [ + "tinyvec", +] + +[[package]] +name = "unicode-properties" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7df058c713841ad818f1dc5d3fd88063241cc61f49f5fbea4b951e8cf5a8d71d" + +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", +] + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasm-bindgen" +version = "0.2.127" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b70935747edd64d89de3efa29d73789b806c15798f8e7dca4d8ac356b50ce70" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.127" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77775f8f3f7217702089053b94958f8f54061a3f663417df76e19cbdcca29bc1" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.127" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e11d33f857dc2fb11b8bc75aee111aa9cbeb12cd9f25efd3d4c2a3dd4e235284" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn 2.0.119", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.127" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ef64dbcc55df09c7e5a46182d181c2cfa3e925f3da937ea764728b4bbb9dcbf" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "webpki-roots" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a" +dependencies = [ + "rustls-pki-types", +] + +[[package]] +name = "whoami" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "998767ef88740d1f5b0682a9c53c24431453923962269c2db68ee43788c5a40d" + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "winnow" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81" + +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" + +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", + "synstructure", +] + +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", + "synstructure", +] + +[[package]] +name = "zeroize" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13c156562582aa81c60cb29407084cdb54c4164760106ab78e6c5b0858cf64e" + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "zmij" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29666d0abbfad1e3dc4dcf6144730dd3a3ab225bbbdac83319345b1b44ccfc1b" diff --git a/rust/Cargo.toml b/rust/Cargo.toml new file mode 100644 index 000000000..f0a30f705 --- /dev/null +++ b/rust/Cargo.toml @@ -0,0 +1,56 @@ +[workspace] +members = [ + "riverqueue", + "riverqueue-cli", + "riverqueue-macros", + "riverqueue-migrate", + "riverqueue-test", +] +resolver = "3" + +[workspace.package] +authors = ["Riverqueue contributors"] +edition = "2024" +homepage = "https://riverqueue.com" +license = "MPL-2.0" +repository = "https://github.com/riverqueue/river" +rust-version = "1.95" +version = "0.50.0-alpha.1" + +[workspace.dependencies] +async-trait = "0.1.92" +base64 = "0.22.1" +chrono = { version = "0.4.45", features = ["serde"] } +futures-util = { version = "0.3.34", default-features = false, features = ["std"] } +proc-macro2 = "1.0.107" +quote = "1.0.47" +rand = "0.10.2" +serde = { version = "1.0.229", features = ["derive"] } +serde_json = { version = "1.0.151", features = ["raw_value"] } +sha2 = "0.11.0" +sqlx = { version = "0.9.0", default-features = false, features = ["runtime-tokio"] } +syn = { version = "2.0", features = ["full"] } +thiserror = "2.0.20" +tokio = { version = "1.53.1", features = ["macros", "rt", "sync", "time"] } +tokio-util = { version = "0.7.19", features = ["rt"] } +tracing = "0.1.44" + +[workspace.lints.rust] +missing_debug_implementations = "warn" +unsafe_code = "forbid" +# `--cfg river_postgres_tests` builds the PostgreSQL integration tests, which +# need `RIVER_RUST_DATABASE_URL`; `make test/rust` sets it when the URL is set. +unexpected_cfgs = { level = "warn", check-cfg = ["cfg(river_postgres_tests)"] } + +[workspace.lints.clippy] +all = { level = "warn", priority = -1 } +pedantic = { level = "warn", priority = -1 } +doc_markdown = "allow" +missing_errors_doc = "warn" +module_name_repetitions = "allow" +must_use_candidate = "allow" + +# Line tables keep file and line numbers in backtraces and panics at a +# fraction of full debug info's size. The `test` profile inherits this. +[profile.dev] +debug = "line-tables-only" diff --git a/rust/README.md b/rust/README.md new file mode 100644 index 000000000..4bf34c0db --- /dev/null +++ b/rust/README.md @@ -0,0 +1,79 @@ +# River for Rust (preview) + +This workspace contains River's Rust implementation. It shares River's +database schema and job protocol with River for Go on PostgreSQL and SQLite, +with an API designed for Rust and Tokio. The crates are a pre-release +preview. + +## Workspace crates + +- `riverqueue`: typed client, worker runtime, CRUD, queues, events, extensions, + periodic/resumable jobs, and maintenance. +- `riverqueue-macros`: `#[derive(JobArgs)]`. +- `riverqueue-migrate`: canonical River migration lines. +- `riverqueue-cli`: the `riverqueue` command-line program for migrations and + benchmarks. +- `riverqueue-test`: typed fixtures and worker-test helpers. + +The API uses a caller-owned SQLx pool, Tokio, typed workers, and +`CancellationToken`. `Client` isn't generic over the database: it accepts a +PostgreSQL or SQLite pool, and there's no driver trait to implement. + +## Quick start + +The [`riverqueue` crate README](riverqueue/README.md) walks through defining +a job, registering a worker, inserting, and starting a client. + +To run Rust clients alongside River Go against one database, including +version matching, queue and kind layout, unique jobs, and rolling deployment +and rollback, see the +[mixed deployment guide](riverqueue/docs/mixed-deployments.md), also published +as `riverqueue::guide::mixed_deployments`. + +Runnable examples in `riverqueue/examples` cover workers and graceful +shutdown, cancellation, transactional completion, unique and periodic jobs, +event subscriptions, custom schemas, SQLite, and a mixed Go and Rust +deployment; `riverqueue-migrate/examples` covers migrations. + +Run the Rust suite from the repository root: + +```sh +make lint/rust +make test/rust +make doc/rust +make check/rust/package +``` + +For basic end-to-end performance figures, the `riverqueue` binary from +`riverqueue-cli` has the Rust equivalent of `river bench`. It truncates the selected River job table, +so use a disposable database: + +```sh +make bench/rust DATABASE_URL=postgres://localhost/river_bench \ + RUST_BENCH_ARGS='--duration 30s' +``` + +The command supports continuous burn, fixed `--num-total-jobs` burn-down, +custom schemas, tunable worker/pool/batch sizes, periodic jobs/sec output, and a +final jobs/sec plus p95 end-to-end latency summary. Use `riverqueue bench +--help` for all options. + +PostgreSQL integration tests require a disposable database. They build only +with `--cfg river_postgres_tests`, which the Makefile targets pass to rustc +and rustdoc, building into `target/postgres-tests`: + +```sh +RIVER_RUST_DATABASE_URL=postgres://localhost/river_rust_test \ + make test/rust/postgres +``` + +CI runs unit, doc, and SQLite tests on each supported Rust version, and +PostgreSQL tests against versions 14 through 18. Rust tests include local +reference fixtures for unique keys, retry bounds, and cron schedules; they +run with Cargo without requiring Go or a separate test harness. + +`make check/rust/package` builds the five publishable crate archives and +verifies that each one builds from its packaged sources, resolving the +exact-version workspace dependencies from the other archives. It does not +publish anything. Release tags use `riverqueue-vX.Y.Z`, independently of Go +module tags. diff --git a/rust/deny.toml b/rust/deny.toml new file mode 100644 index 000000000..4aae2315c --- /dev/null +++ b/rust/deny.toml @@ -0,0 +1,29 @@ +[graph] +all-features = true + +[advisories] +yanked = "deny" + +[licenses] +allow = [ + "Apache-2.0", + "BSD-3-Clause", + "CDLA-Permissive-2.0", + "ISC", + "MIT", + "MPL-2.0", + "Unicode-3.0", + "Zlib", +] +confidence-threshold = 0.8 + +[bans] +highlight = "all" +multiple-versions = "warn" +wildcards = "allow" + +[sources] +allow-git = [] +allow-registry = ["https://github.com/rust-lang/crates.io-index"] +unknown-git = "deny" +unknown-registry = "deny" diff --git a/rust/go.mod b/rust/go.mod new file mode 100644 index 000000000..fd1170a81 --- /dev/null +++ b/rust/go.mod @@ -0,0 +1,4 @@ +// Stub module that only keeps rust/ out of the root Go module and its tools. +module github.com/riverqueue/river/rust + +go 1.26.0 diff --git a/rust/riverqueue-cli/Cargo.toml b/rust/riverqueue-cli/Cargo.toml new file mode 100644 index 000000000..34f27a7e3 --- /dev/null +++ b/rust/riverqueue-cli/Cargo.toml @@ -0,0 +1,34 @@ +[package] +name = "riverqueue-cli" +description = "Command-line tools for River's Rust client: migrations and benchmarks" +keywords = ["background", "jobs", "migrations", "postgres", "queue"] +categories = ["command-line-utilities", "database"] +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +readme = "README.md" +repository.workspace = true +homepage.workspace = true + +[features] +default = ["postgres", "sqlite"] +postgres = ["riverqueue/postgres", "riverqueue-migrate/postgres", "sqlx/postgres"] +sqlite = ["riverqueue/sqlite", "riverqueue-migrate/sqlite", "sqlx/sqlite"] + +[[bin]] +name = "riverqueue" +path = "src/main.rs" +# The binary shares the library crate's name, so its docs would collide. +doc = false + +[dependencies] +riverqueue = { path = "../riverqueue", version = "=0.50.0-alpha.1", default-features = false } +riverqueue-migrate = { path = "../riverqueue-migrate", version = "=0.50.0-alpha.1", default-features = false } +serde.workspace = true +sqlx = { workspace = true, features = ["tls-rustls"] } +tokio = { workspace = true, features = ["macros", "rt-multi-thread", "signal"] } +tokio-util.workspace = true + +[lints] +workspace = true diff --git a/rust/riverqueue-cli/LICENSE b/rust/riverqueue-cli/LICENSE new file mode 120000 index 000000000..30cff7403 --- /dev/null +++ b/rust/riverqueue-cli/LICENSE @@ -0,0 +1 @@ +../../LICENSE \ No newline at end of file diff --git a/rust/riverqueue-cli/README.md b/rust/riverqueue-cli/README.md new file mode 100644 index 000000000..1a27fb7ec --- /dev/null +++ b/rust/riverqueue-cli/README.md @@ -0,0 +1,36 @@ +# riverqueue-cli + +Command-line tools for [River](https://riverqueue.com)'s Rust client. Install +the `riverqueue` binary with: + +```sh +cargo install riverqueue-cli +``` + +## Migrations + +River's schema is managed by versioned migrations shared with every River +implementation. Apply them before starting clients: + +```sh +riverqueue migrate-up --database-url postgres://localhost/app +riverqueue migrate-up --database-url postgres://localhost/app --schema river +riverqueue migrate-up --database-url sqlite://app.sqlite3 +``` + +`migrate-down`, `migrate-list`, and `validate` take the same connection +options. `--target-version N`, `--max-steps N`, and `--dry-run` limit or +preview a migration run. Applications can instead migrate from Rust with the +[`riverqueue-migrate`](https://docs.rs/riverqueue-migrate) crate. + +## Benchmark + +`riverqueue bench` measures worker throughput and end-to-end latency. It +**truncates the River job table** in the selected database, so only point it at +a disposable database: + +```sh +riverqueue bench --database-url postgres://localhost/river_bench --duration 30s +``` + +Run `riverqueue bench --help` for its options. diff --git a/rust/riverqueue-cli/src/bench.rs b/rust/riverqueue-cli/src/bench.rs new file mode 100644 index 000000000..410d04e8e --- /dev/null +++ b/rust/riverqueue-cli/src/bench.rs @@ -0,0 +1,741 @@ +//! The destructive `bench` command. + +use std::{ + convert::Infallible, + error::Error as StdError, + io, + sync::{ + Arc, OnceLock, + atomic::{AtomicU64, Ordering}, + }, + time::{Duration, Instant}, +}; + +use riverqueue::{ + Client, EventKind, EventReceiver, EventRecvError, InsertOpts, Job, JobArgs, QueueConfig, + SubscribeConfig, WorkContext, WorkOutcome, Worker, WorkerRegistry, database::SchemaName, +}; +use serde::{Deserialize, Serialize}; +use sqlx::{ + AssertSqlSafe, PgPool, + postgres::{PgConnectOptions, PgPoolOptions}, +}; +use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; + +const DEFAULT_BACKLOG: u64 = 75_000; +const DEFAULT_BATCH_SIZE: usize = 5_000; +const DEFAULT_MAX_CONNECTIONS: u32 = 50; +const DEFAULT_MAX_WORKERS: usize = 2_000; +const ITERATION_PERIOD: Duration = Duration::from_secs(2); + +pub(crate) const HELP: &str = r"Benchmark River's Rust worker runtime + +Usage: + riverqueue bench [options] + +The benchmark truncates the selected River job table, optionally vacuums it, +then inserts and works no-op jobs while reporting rough throughput and p95 +end-to-end latency. Use only a disposable development or benchmark database. + +Options: + --database-url URL PostgreSQL URL (or set DATABASE_URL) + --schema NAME River schema (default: current schema) + --duration DURATION Stop after a Go-style duration such as 30s or 5m + -n, --num-total-jobs COUNT Insert COUNT jobs, then work them all + --backlog COUNT Target continuous-mode backlog (default: 75000) + --batch-size COUNT Jobs per insertion batch (default: 5000) + --max-connections COUNT SQLx pool size (default: 50) + --max-workers COUNT Concurrent workers (default: 2000) + --skip-vacuum Truncate without VACUUM FULL + -h, --help Print help + +With neither --duration nor --num-total-jobs, the benchmark runs until Ctrl-C. +The two stopping options are mutually exclusive. +"; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "benchmark")] +struct BenchmarkArgs { + num: u64, +} + +struct BenchmarkWorker; + +impl Worker for BenchmarkWorker { + type Error = Infallible; + + fn work( + &self, + _context: WorkContext, + _job: Job, + ) -> impl std::future::Future> + Send { + std::future::ready(Ok(WorkOutcome::Complete)) + } +} + +#[derive(Clone, Debug, Eq, PartialEq)] +struct BenchOptions { + backlog: u64, + batch_size: usize, + database_url: String, + duration: Option, + max_connections: u32, + max_workers: usize, + num_total_jobs: Option, + schema: SchemaName, + vacuum: bool, +} + +enum Command { + Bench(BenchOptions), + Help, +} + +#[derive(Clone, Copy, Debug)] +struct DatabaseStatistics { + failed: i64, + p95_seconds: Option, + worked: u64, +} + +struct Producer { + backlog: u64, + batch_size: usize, + cancellation: CancellationToken, + client: Client, + inserted: Arc, + next_number: u64, + worked: Arc, +} + +#[derive(Clone)] +struct BenchmarkProgress { + completion_notify: Arc, + failed: Arc, + last_worked_elapsed_nanos: Arc, + started_at: Arc>, + worked: Arc, +} + +impl BenchmarkProgress { + fn new() -> Self { + Self { + completion_notify: Arc::new(Notify::new()), + failed: Arc::new(AtomicU64::new(0)), + last_worked_elapsed_nanos: Arc::new(AtomicU64::new(0)), + started_at: Arc::new(OnceLock::new()), + worked: Arc::new(AtomicU64::new(0)), + } + } + + async fn count_events( + &self, + mut events: EventReceiver, + cancellation: CancellationToken, + num_total_jobs: Option, + ) -> Result<(), String> { + loop { + tokio::select! { + () = cancellation.cancelled() => return Ok(()), + event = events.recv() => match event { + Ok(event) if event.kind() == EventKind::JobCompleted => { + if let Some(started_at) = self.started_at.get() { + self.last_worked_elapsed_nanos.store( + u64::try_from(started_at.elapsed().as_nanos()) + .unwrap_or(u64::MAX), + Ordering::Relaxed, + ); + } + let worked_now = self.worked.fetch_add(1, Ordering::Relaxed) + 1; + if num_total_jobs.is_some_and(|total| worked_now >= total) { + self.completion_notify.notify_one(); + } + } + Ok(_) => { + self.failed.fetch_add(1, Ordering::Relaxed); + } + Err(EventRecvError::Lagged(count)) => { + self.failed.store(u64::MAX, Ordering::Relaxed); + return Err(format!("benchmark event receiver lagged by {count} events")); + } + Err(error) => { + self.failed.store(u64::MAX, Ordering::Relaxed); + return Err(format!("benchmark event receiver failed: {error}")); + } + }, + } + } + } +} + +/// Runs the `bench` command with the arguments that follow it. +pub(crate) async fn run( + arguments: impl IntoIterator, + database_url_env: Option, +) -> Result<(), Box> { + match parse_command(arguments, database_url_env)? { + Command::Bench(options) => run_benchmark(options).await?, + Command::Help => print!("{HELP}"), + } + Ok(()) +} + +fn parse_command( + arguments: impl IntoIterator, + database_url_env: Option, +) -> Result { + let mut arguments = arguments.into_iter(); + + let mut backlog = DEFAULT_BACKLOG; + let mut batch_size = DEFAULT_BATCH_SIZE; + let mut database_url = None; + let mut duration = None; + let mut max_connections = DEFAULT_MAX_CONNECTIONS; + let mut max_workers = DEFAULT_MAX_WORKERS; + let mut num_total_jobs = None; + let mut schema = SchemaName::current(); + let mut vacuum = true; + + while let Some(argument) = arguments.next() { + match argument.as_str() { + "-h" | "--help" => return Ok(Command::Help), + "--backlog" => { + backlog = parse_positive(&take_value(&mut arguments, "--backlog")?, "backlog")?; + } + "--batch-size" => { + batch_size = + parse_positive(&take_value(&mut arguments, "--batch-size")?, "batch size")?; + } + "--database-url" => { + database_url = Some(take_value(&mut arguments, "--database-url")?); + } + "--duration" => { + duration = Some(parse_duration(&take_value(&mut arguments, "--duration")?)?); + } + "--max-connections" => { + max_connections = parse_positive( + &take_value(&mut arguments, "--max-connections")?, + "maximum connections", + )?; + } + "--max-workers" => { + max_workers = parse_positive( + &take_value(&mut arguments, "--max-workers")?, + "maximum workers", + )?; + } + "-n" | "--num-total-jobs" => { + num_total_jobs = Some(parse_positive( + &take_value(&mut arguments, "--num-total-jobs")?, + "total jobs", + )?); + } + "--schema" => { + schema = SchemaName::new(take_value(&mut arguments, "--schema")?) + .map_err(|error| invalid_input(error.to_string()))?; + } + "--skip-vacuum" => vacuum = false, + _ => return Err(invalid_input(format!("unknown bench option {argument:?}"))), + } + } + + if duration.is_some() && num_total_jobs.is_some() { + return Err(invalid_input( + "--duration and --num-total-jobs are mutually exclusive", + )); + } + if max_workers > riverqueue::QUEUE_NUM_WORKERS_MAX { + return Err(invalid_input(format!( + "maximum workers cannot exceed {}", + riverqueue::QUEUE_NUM_WORKERS_MAX + ))); + } + let database_url = database_url.or(database_url_env).ok_or_else(|| { + invalid_input("--database-url or the DATABASE_URL environment variable is required") + })?; + + Ok(Command::Bench(BenchOptions { + backlog, + batch_size, + database_url, + duration, + max_connections, + max_workers, + num_total_jobs, + schema, + vacuum, + })) +} + +fn take_value( + arguments: &mut impl Iterator, + option: &str, +) -> Result { + arguments + .next() + .ok_or_else(|| invalid_input(format!("{option} requires a value"))) +} + +fn parse_positive(value: &str, name: &str) -> Result +where + T: TryFrom, +{ + let value = value + .parse::() + .map_err(|error| invalid_input(format!("invalid {name}: {error}")))?; + if value == 0 { + return Err(invalid_input(format!("{name} must be positive"))); + } + T::try_from(value).map_err(|_| invalid_input(format!("{name} is too large"))) +} + +fn parse_duration(value: &str) -> Result { + if value.is_empty() { + return Err(invalid_input("duration cannot be empty")); + } + + let bytes = value.as_bytes(); + let mut index = 0; + let mut total_nanos = 0_u128; + while index < bytes.len() { + let number_start = index; + let mut decimal_seen = false; + while index < bytes.len() + && (bytes[index].is_ascii_digit() || (!decimal_seen && bytes[index] == b'.')) + { + decimal_seen |= bytes[index] == b'.'; + index += 1; + } + if number_start == index { + return Err(invalid_input(format!("invalid duration {value:?}"))); + } + let number = value[number_start..index] + .parse::() + .map_err(|error| invalid_input(format!("invalid duration {value:?}: {error}")))?; + if !number.is_finite() || number < 0.0 { + return Err(invalid_input(format!("invalid duration {value:?}"))); + } + + let units = [ + ("ns", 1_u128), + ("us", 1_000), + ("µs", 1_000), + ("ms", 1_000_000), + ("s", 1_000_000_000), + ("m", 60 * 1_000_000_000), + ("h", 60 * 60 * 1_000_000_000), + ]; + let (unit, multiplier) = units + .into_iter() + .find(|(unit, _)| value[index..].starts_with(unit)) + .ok_or_else(|| invalid_input(format!("invalid duration unit in {value:?}")))?; + index += unit.len(); + #[allow( + clippy::cast_possible_truncation, + clippy::cast_precision_loss, + clippy::cast_sign_loss + )] + let segment_nanos = (number * multiplier as f64).round() as u128; + total_nanos = total_nanos + .checked_add(segment_nanos) + .ok_or_else(|| invalid_input("duration is too large"))?; + } + if total_nanos == 0 { + return Err(invalid_input("duration must be positive")); + } + let seconds = u64::try_from(total_nanos / 1_000_000_000) + .map_err(|_| invalid_input("duration is too large"))?; + let nanos = + u32::try_from(total_nanos % 1_000_000_000).expect("nanosecond remainder always fits u32"); + Ok(Duration::new(seconds, nanos)) +} + +fn invalid_input(message: impl Into) -> io::Error { + io::Error::new(io::ErrorKind::InvalidInput, message.into()) +} + +/// Receiver capacity for a benchmark's completion events. +fn event_capacity(backlog: u64) -> std::num::NonZeroUsize { + usize::try_from(backlog) + .ok() + .and_then(std::num::NonZeroUsize::new) + .unwrap_or(std::num::NonZeroUsize::MAX) +} + +async fn run_benchmark(options: BenchOptions) -> Result<(), Box> { + eprintln!( + "bench: WARNING: truncating {} in the selected database", + options.schema.qualify("river_job") + ); + let pool = PgPoolOptions::new() + .max_connections(options.max_connections) + .connect_with(postgres_connect_options(&options.database_url)?) + .await?; + reset_jobs(&pool, &options.schema, options.vacuum).await?; + let client = benchmark_client(pool.clone(), options.max_workers)?; + let events = client.subscribe_config( + SubscribeConfig::new([ + EventKind::JobCancelled, + EventKind::JobCompleted, + EventKind::JobFailed, + ])? + .with_buffer_capacity(event_capacity(options.backlog)), + )?; + let inserted = Arc::new(AtomicU64::new(0)); + let progress = BenchmarkProgress::new(); + let event_cancel = CancellationToken::new(); + let event_task = tokio::spawn({ + let cancellation = event_cancel.child_token(); + let num_total_jobs = options.num_total_jobs; + let progress = progress.clone(); + async move { + progress + .count_events(events, cancellation, num_total_jobs) + .await + } + }); + let mut next_number = 0_u64; + let initial_jobs = options.num_total_jobs.unwrap_or(options.backlog); + insert_jobs( + &client, + &inserted, + &mut next_number, + initial_jobs, + options.batch_size, + ) + .await?; + + let mut run = client.start()?; + run.wait_ready().await?; + let started_at = Instant::now(); + progress + .started_at + .set(started_at) + .expect("benchmark start time is set once"); + let stop_producer = CancellationToken::new(); + let mut producer = options.num_total_jobs.is_none().then(|| { + tokio::spawn(run_producer(Producer { + backlog: options.backlog, + batch_size: options.batch_size, + cancellation: stop_producer.child_token(), + client: client.clone(), + inserted: Arc::clone(&inserted), + next_number, + worked: Arc::clone(&progress.worked), + })) + }); + + let run_result = monitor_benchmark( + &options, + &inserted, + &progress.worked, + &progress.failed, + &progress.completion_notify, + started_at, + producer.as_mut(), + ) + .await; + stop_producer.cancel(); + if let Some(producer) = producer { + producer.await.map_err(|error| join_error(&error))??; + } + run.shutdown().await?; + event_cancel.cancel(); + event_task.await.map_err(|error| join_error(&error))??; + run_result?; + + let final_stats = database_statistics(&pool, &options.schema).await?; + if final_stats.failed > 0 { + return Err(format!("{} benchmark jobs failed", final_stats.failed).into()); + } + let last_worked_elapsed_nanos = progress.last_worked_elapsed_nanos.load(Ordering::Relaxed); + let elapsed = if last_worked_elapsed_nanos == 0 { + started_at.elapsed() + } else { + Duration::from_nanos(last_worked_elapsed_nanos) + }; + println!( + "bench: total jobs worked [ {:10} ], total jobs inserted [ {:10} ], overall job/sec [ {:10.1} ], p95 [ {:>10} ], running {}", + final_stats.worked, + inserted.load(Ordering::Relaxed), + throughput(final_stats.worked, elapsed), + display_p95(final_stats.p95_seconds), + display_duration(elapsed), + ); + Ok(()) +} + +fn benchmark_client(pool: PgPool, max_workers: usize) -> Result { + let mut workers = WorkerRegistry::new(); + workers.register::(BenchmarkWorker)?; + Client::builder(pool) + .id("riverqueue-benchmark") + .workers(workers) + .queue( + riverqueue::QUEUE_DEFAULT, + QueueConfig::new(max_workers) + .with_fetch_cooldown(Duration::from_millis(2)) + .with_fetch_poll_interval(Duration::from_millis(20)), + ) + .build() +} + +fn postgres_connect_options(database_url: &str) -> Result { + use std::str::FromStr; + + let mut options = PgConnectOptions::from_str(database_url)?; + if !database_url_has_userinfo(database_url) + && let Some(username) = ["PGUSER", "USER", "LOGNAME"] + .into_iter() + .find_map(|name| std::env::var(name).ok().filter(|value| !value.is_empty())) + { + options = options.username(&username); + } + Ok(options) +} + +fn database_url_has_userinfo(database_url: &str) -> bool { + database_url + .split_once("://") + .and_then(|(_, remainder)| remainder.split('/').next()) + .is_some_and(|authority| authority.contains('@')) +} + +async fn reset_jobs(pool: &PgPool, schema: &SchemaName, vacuum: bool) -> Result<(), sqlx::Error> { + let table = schema.qualify("river_job"); + sqlx::query(AssertSqlSafe(format!("TRUNCATE TABLE {table}"))) + .execute(pool) + .await?; + if vacuum { + sqlx::query(AssertSqlSafe(format!("VACUUM FULL {table}"))) + .execute(pool) + .await?; + } + Ok(()) +} + +async fn insert_jobs( + client: &Client, + inserted: &AtomicU64, + next_number: &mut u64, + count: u64, + batch_size: usize, +) -> Result<(), riverqueue::Error> { + let mut remaining = count; + while remaining > 0 { + let batch_size = u64::try_from(batch_size).unwrap_or(u64::MAX); + let current_batch = + usize::try_from(remaining.min(batch_size)).expect("batch size bounds the conversion"); + let jobs = (0..current_batch) + .map(|_| { + *next_number = next_number.wrapping_add(1); + (BenchmarkArgs { num: *next_number }, InsertOpts::default()) + }) + .collect::>(); + let count = u64::try_from(client.insert_many(jobs).await?.len()).unwrap_or(u64::MAX); + inserted.fetch_add(count, Ordering::Relaxed); + remaining -= count; + } + Ok(()) +} + +async fn run_producer(mut producer: Producer) -> Result<(), Box> { + loop { + let jobs_left = producer + .inserted + .load(Ordering::Relaxed) + .saturating_sub(producer.worked.load(Ordering::Relaxed)); + if jobs_left < producer.backlog { + insert_jobs( + &producer.client, + &producer.inserted, + &mut producer.next_number, + producer.backlog - jobs_left, + producer.batch_size, + ) + .await?; + } + tokio::select! { + () = producer.cancellation.cancelled() => return Ok(()), + () = tokio::time::sleep(Duration::from_millis(250)) => {} + } + } +} + +async fn monitor_benchmark( + options: &BenchOptions, + inserted: &AtomicU64, + worked: &AtomicU64, + failed: &AtomicU64, + completion_notify: &Notify, + started_at: Instant, + mut producer: Option<&mut tokio::task::JoinHandle>>>, +) -> Result<(), Box> { + let mut interval = tokio::time::interval(ITERATION_PERIOD); + interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + interval.tick().await; + let deadline = options.duration.map(|duration| started_at + duration); + let mut last_inserted = 0_u64; + let mut last_worked = 0_u64; + + loop { + let failed_now = failed.load(Ordering::Relaxed); + if failed_now > 0 { + return Err(format!("{failed_now} benchmark jobs failed").into()); + } + if options + .num_total_jobs + .is_some_and(|total| worked.load(Ordering::Relaxed) >= total) + { + return Ok(()); + } + if deadline.is_some_and(|deadline| Instant::now() >= deadline) { + return Ok(()); + } + + tokio::select! { + _ = interval.tick() => { + let now = Instant::now(); + let inserted_now = inserted.load(Ordering::Relaxed); + let worked_now = worked.load(Ordering::Relaxed); + let jobs_inserted = inserted_now.saturating_sub(last_inserted); + let jobs_worked = worked_now.saturating_sub(last_worked); + println!( + "bench: jobs worked [ {jobs_worked:10} ], inserted [ {jobs_inserted:10} ], job/sec [ {:10.1} ] [{}]", + throughput(jobs_worked, ITERATION_PERIOD), + display_duration(now.duration_since(started_at)), + ); + last_inserted = inserted_now; + last_worked = worked_now; + } + result = async { + match producer.as_mut() { + Some(producer) => Some(producer.await), + None => std::future::pending().await, + } + } => { + let result = result.expect("producer result is present"); + return match result { + Ok(Ok(())) => Err("benchmark producer stopped unexpectedly".into()), + Ok(Err(error)) => Err(error), + Err(error) => Err(join_error(&error)), + }; + } + result = tokio::signal::ctrl_c() => { + result?; + return Ok(()); + } + () = completion_notify.notified(), if options.num_total_jobs.is_some() => {} + () = async { + match deadline { + Some(deadline) => tokio::time::sleep_until(deadline.into()).await, + None => std::future::pending().await, + } + } => return Ok(()), + } + } +} + +async fn database_statistics( + pool: &PgPool, + schema: &SchemaName, +) -> Result { + let table = schema.qualify("river_job"); + let sql = format!( + "SELECT \ + count(*) FILTER (WHERE state IN ('cancelled', 'discarded'))::bigint, \ + percentile_cont(0.95) WITHIN GROUP (ORDER BY \ + extract(epoch FROM (finalized_at - created_at))::double precision) \ + FILTER (WHERE state = 'completed'), \ + count(*) FILTER (WHERE state = 'completed')::bigint \ + FROM {table}" + ); + let (failed, p95_seconds, worked) = + sqlx::query_as::<_, (i64, Option, i64)>(AssertSqlSafe(sql)) + .fetch_one(pool) + .await?; + Ok(DatabaseStatistics { + failed, + p95_seconds, + worked: u64::try_from(worked).unwrap_or_default(), + }) +} + +#[allow(clippy::cast_precision_loss)] +fn throughput(jobs: u64, duration: Duration) -> f64 { + if duration.is_zero() { + return 0.0; + } + jobs as f64 / duration.as_secs_f64() +} + +fn display_p95(seconds: Option) -> String { + seconds.map_or_else(|| "n/a".to_owned(), |seconds| format!("{seconds:.3}s")) +} + +fn display_duration(duration: Duration) -> String { + format!("{:.1}s", duration.as_secs_f64()) +} + +fn join_error(error: &tokio::task::JoinError) -> Box { + Box::new(io::Error::other(format!( + "benchmark producer task failed: {error}" + ))) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn detects_explicit_postgres_userinfo() { + assert!(!database_url_has_userinfo( + "postgres://localhost/river_bench" + )); + assert!(database_url_has_userinfo( + "postgres://river@localhost/river_bench" + )); + assert!(database_url_has_userinfo( + "postgres://river:secret@localhost/river_bench" + )); + } + + #[test] + fn parses_bench_options_and_composite_duration() { + let command = parse_command( + [ + "--database-url", + "postgres://localhost/river_bench", + "--duration", + "1m30.5s", + "--max-workers", + "32", + "--skip-vacuum", + ] + .into_iter() + .map(str::to_owned), + None, + ) + .unwrap(); + let Command::Bench(options) = command else { + panic!("expected bench command"); + }; + assert_eq!(options.duration, Some(Duration::from_millis(90_500))); + assert_eq!(options.max_workers, 32); + assert!(!options.vacuum); + } + + #[test] + fn rejects_conflicting_stopping_options() { + let error = parse_command( + ["--duration", "1s", "--num-total-jobs", "10"] + .into_iter() + .map(str::to_owned), + Some("postgres://localhost/river_bench".to_owned()), + ) + .err() + .unwrap(); + assert!(error.to_string().contains("mutually exclusive")); + } +} diff --git a/rust/riverqueue-cli/src/main.rs b/rust/riverqueue-cli/src/main.rs new file mode 100644 index 000000000..d132d422c --- /dev/null +++ b/rust/riverqueue-cli/src/main.rs @@ -0,0 +1,65 @@ +//! River's Rust command-line interface. +//! +//! Runs River migrations against PostgreSQL or SQLite and benchmarks the +//! worker runtime: +//! +//! ```text +//! riverqueue migrate-up --database-url postgres://localhost/app +//! riverqueue migrate-list --database-url sqlite://app.sqlite3 +//! riverqueue bench --database-url postgres://localhost/river_bench --duration 30s +//! ``` + +#![forbid(unsafe_code)] + +use std::{env, error::Error, process::ExitCode}; + +#[cfg(feature = "postgres")] +mod bench; +mod migrate; + +const HELP: &str = "River for Rust + +Usage: + riverqueue [options] + +Commands: + migrate-down Run down migrations + migrate-list List applied migration versions + migrate-up Run up migrations + validate Check that all migrations are applied + bench Benchmark job throughput against a disposable database + +Run `riverqueue --help` for a command's options. Commands read the +database URL from --database-url or the DATABASE_URL environment variable. +"; + +#[tokio::main] +async fn main() -> ExitCode { + match run(env::args().skip(1).collect()).await { + Ok(()) => ExitCode::SUCCESS, + Err(error) => { + eprintln!("riverqueue: {error}"); + ExitCode::FAILURE + } + } +} + +async fn run(mut arguments: Vec) -> Result<(), Box> { + if arguments.is_empty() { + print!("{HELP}"); + return Ok(()); + } + let command = arguments.remove(0); + let database_url_env = env::var("DATABASE_URL").ok(); + match command.as_str() { + "-h" | "--help" | "help" => print!("{HELP}"), + "-V" | "--version" | "version" => println!("riverqueue {}", env!("CARGO_PKG_VERSION")), + #[cfg(feature = "postgres")] + "bench" => bench::run(arguments, database_url_env).await?, + "migrate-down" | "migrate-list" | "migrate-up" | "validate" => { + migrate::run(command, arguments, database_url_env).await?; + } + _ => return Err(format!("unknown command {command:?}\n\n{HELP}").into()), + } + Ok(()) +} diff --git a/rust/riverqueue-cli/src/migrate.rs b/rust/riverqueue-cli/src/migrate.rs new file mode 100644 index 000000000..4161425cc --- /dev/null +++ b/rust/riverqueue-cli/src/migrate.rs @@ -0,0 +1,197 @@ +//! Migration commands for River's canonical migration line. + +use std::error::Error; + +#[cfg(feature = "postgres")] +use riverqueue::database::SchemaName; +#[cfg(feature = "postgres")] +use riverqueue_migrate::PostgresMigrator; +#[cfg(feature = "sqlite")] +use riverqueue_migrate::SqliteMigrator; +use riverqueue_migrate::{Direction, MigrateOpts, MigrateResult, ValidateResult}; +#[cfg(feature = "postgres")] +use sqlx::PgPool; +#[cfg(feature = "sqlite")] +use sqlx::SqlitePool; + +#[derive(Default)] +struct Args { + command: String, + database_url: String, + dry_run: bool, + max_steps: Option, + schema: Option, + target_version: Option, +} + +enum CommandMigrator { + #[cfg(feature = "postgres")] + Postgres(PostgresMigrator), + #[cfg(feature = "sqlite")] + Sqlite(SqliteMigrator), +} + +impl CommandMigrator { + async fn existing_versions(&self) -> Result, riverqueue_migrate::Error> { + match self { + #[cfg(feature = "postgres")] + Self::Postgres(migrator) => migrator.existing_versions().await, + #[cfg(feature = "sqlite")] + Self::Sqlite(migrator) => migrator.existing_versions().await, + } + } + + async fn migrate( + &self, + direction: Direction, + opts: MigrateOpts, + ) -> Result { + match self { + #[cfg(feature = "postgres")] + Self::Postgres(migrator) => migrator.migrate(direction, opts).await, + #[cfg(feature = "sqlite")] + Self::Sqlite(migrator) => migrator.migrate(direction, opts).await, + } + } + + async fn validate( + &self, + target_version: Option, + ) -> Result { + match self { + #[cfg(feature = "postgres")] + Self::Postgres(migrator) => migrator.validate(target_version).await, + #[cfg(feature = "sqlite")] + Self::Sqlite(migrator) => migrator.validate(target_version).await, + } + } +} + +/// Runs a migration command with the arguments that follow it. +pub(crate) async fn run( + command: String, + arguments: impl IntoIterator, + database_url_env: Option, +) -> Result<(), Box> { + let Some(args) = parse_args(command, arguments, database_url_env)? else { + println!("{}", usage()); + return Ok(()); + }; + let migrator = if args.database_url.starts_with("sqlite:") { + #[cfg(not(feature = "sqlite"))] + return Err("SQLite support requires the `sqlite` feature".into()); + #[cfg(feature = "sqlite")] + { + if args.schema.is_some() { + return Err("--schema is only supported for PostgreSQL".into()); + } + CommandMigrator::Sqlite(SqliteMigrator::new( + SqlitePool::connect(&args.database_url).await?, + )) + } + } else { + #[cfg(not(feature = "postgres"))] + return Err("PostgreSQL support requires the `postgres` feature".into()); + #[cfg(feature = "postgres")] + { + let pool = PgPool::connect(&args.database_url).await?; + let mut migrator = PostgresMigrator::new(pool); + if let Some(schema) = &args.schema { + migrator = migrator.with_schema(SchemaName::new(schema.clone())?); + } + CommandMigrator::Postgres(migrator) + } + }; + + match args.command.as_str() { + "migrate-down" => { + print_migrations( + migrator + .migrate(Direction::Down, migrate_opts(&args)) + .await?, + ); + } + "migrate-list" => { + for version in migrator.existing_versions().await? { + println!("{version:03}"); + } + } + "migrate-up" => { + print_migrations(migrator.migrate(Direction::Up, migrate_opts(&args)).await?); + } + "validate" => { + let result = migrator.validate(args.target_version).await?; + if !result.is_valid() { + eprintln!("{result}"); + std::process::exit(2); + } + println!("River migrations valid"); + } + command => return Err(format!("unknown command {command:?}\n{}", usage()).into()), + } + Ok(()) +} + +fn migrate_opts(args: &Args) -> MigrateOpts { + let mut opts = MigrateOpts::new().with_dry_run(args.dry_run); + if let Some(max_steps) = args.max_steps { + opts = opts.with_max_steps(max_steps); + } + if let Some(target_version) = args.target_version { + opts = opts.with_target_version(target_version); + } + opts +} + +fn parse_args( + command: String, + arguments: impl IntoIterator, + database_url_env: Option, +) -> Result, Box> { + let mut raw = arguments.into_iter(); + let mut args = Args { + command, + ..Args::default() + }; + while let Some(argument) = raw.next() { + match argument.as_str() { + "--database-url" => { + args.database_url = raw.next().ok_or("--database-url requires a value")?; + } + "--dry-run" => args.dry_run = true, + "--max-steps" => { + args.max_steps = Some(raw.next().ok_or("--max-steps requires a value")?.parse()?); + } + "--schema" => args.schema = Some(raw.next().ok_or("--schema requires a value")?), + "--target-version" => { + args.target_version = Some( + raw.next() + .ok_or("--target-version requires a value")? + .parse()?, + ); + } + "-h" | "--help" => return Ok(None), + _ => return Err(format!("unknown argument {argument:?}\n{}", usage()).into()), + } + } + if args.database_url.is_empty() { + args.database_url = database_url_env.ok_or("--database-url or DATABASE_URL is required")?; + } + Ok(Some(args)) +} + +fn print_migrations(result: MigrateResult) { + for version in result.versions { + println!( + "{:03} {:?} {} ({:?})", + version.version, result.direction, version.name, version.duration + ); + if !version.sql.is_empty() && version.duration.is_zero() { + println!("{}", version.sql); + } + } +} + +fn usage() -> &'static str { + "usage: riverqueue [--database-url URL] [--schema NAME] [--target-version N] [--max-steps N] [--dry-run]" +} diff --git a/rust/riverqueue-macros/Cargo.toml b/rust/riverqueue-macros/Cargo.toml new file mode 100644 index 000000000..28e922d16 --- /dev/null +++ b/rust/riverqueue-macros/Cargo.toml @@ -0,0 +1,28 @@ +[package] +name = "riverqueue-macros" +description = "Derive macros for River's Rust client" +keywords = ["background", "derive", "jobs", "queue"] +categories = ["asynchronous", "database", "development-tools::procedural-macro-helpers"] +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +readme = "README.md" +repository.workspace = true +homepage.workspace = true + +[lib] +proc-macro = true + +[dependencies] +proc-macro2.workspace = true +quote.workspace = true +syn.workspace = true + +[dev-dependencies] +riverqueue = { path = "../riverqueue" } +serde.workspace = true +trybuild = "1.0.121" + +[lints] +workspace = true diff --git a/rust/riverqueue-macros/LICENSE b/rust/riverqueue-macros/LICENSE new file mode 120000 index 000000000..30cff7403 --- /dev/null +++ b/rust/riverqueue-macros/LICENSE @@ -0,0 +1 @@ +../../LICENSE \ No newline at end of file diff --git a/rust/riverqueue-macros/README.md b/rust/riverqueue-macros/README.md new file mode 100644 index 000000000..537a55f83 --- /dev/null +++ b/rust/riverqueue-macros/README.md @@ -0,0 +1,41 @@ +# riverqueue-macros + +Derive macros for River's Rust client. Applications normally receive these +macros through `riverqueue` and do not need to depend on this crate directly. + +`#[derive(JobArgs)]` requires a stable `#[river(kind = "...")]` and can declare +kind aliases, default queue, max attempts, priority, pending state, and tags, +and default uniqueness: + +```rust,ignore +#[derive(Deserialize, JobArgs, Serialize)] +#[river( + kind = "send_invoice", + queue = "billing", + tags("billing"), + unique(by_args, by_period = "24h", by_state(available, pending, running, scheduled)), + insert_opts = invoice_insert_opts, +)] +struct SendInvoice { + #[river(unique)] + invoice_id: i64, + note: String, +} +``` + +`unique(...)` accepts `by_args`, `by_args("nested.path", ...)`, `by_period` +(a duration of at least one second, such as `"90s"` or `"1h30m"`), +`by_queue`, `by_state(...)`, and `exclude_kind`, matching the setters of +`UniqueOpts`. Fields marked +`#[river(unique)]`, plus any `by_args` paths, are the arguments hashed for +uniqueness; with neither, every argument is hashed. Marking a field without +`unique(by_args)` is a compile error, since it would otherwise have no effect. +Unique paths follow Serde's serialization-side `rename` and `rename_all` rules, +including raw Rust identifiers. Conditionally skipped optional fields are +omitted from the unique key when absent; flattened or always skipped unique +fields are rejected because their wire path is ambiguous. + +`insert_opts = path::to_fn` names a `fn() -> InsertOpts` whose options are +overlaid on the attribute defaults. +`crate = "path"` sets the path to `riverqueue` when it is renamed or +re-exported. diff --git a/rust/riverqueue-macros/src/lib.rs b/rust/riverqueue-macros/src/lib.rs new file mode 100644 index 000000000..f09ec0bab --- /dev/null +++ b/rust/riverqueue-macros/src/lib.rs @@ -0,0 +1,1090 @@ +#![doc = include_str!("../README.md")] +#![forbid(unsafe_code)] + +use proc_macro::TokenStream; +use proc_macro2::Span; +use quote::{quote, quote_spanned}; +use syn::{ + Data, DeriveInput, Fields, Ident, Lit, LitBool, LitInt, LitStr, Meta, Token, + ext::IdentExt as _, meta::ParseNestedMeta, parenthesized, parse_macro_input, + punctuated::Punctuated, spanned::Spanned as _, token, +}; + +/// Derives `riverqueue::JobArgs`. +/// +/// The type must declare `#[river(kind = "...")]`. Other type-level options: +/// +/// * `aliases("old_kind", ...)`: former kinds handled by the same worker. +/// * `queue = "..."`, `max_attempts = N`, `priority = N`, `pending = true`, +/// and `tags("a", "b")`: default insertion options. +/// * `unique(...)`: makes the job unique by default. Options are `by_args`, +/// `by_args("nested.path", ...)`, `by_period = "1h"` (a duration of at least +/// one second, such as `"90s"` or `"1h30m"`), `by_queue`, `by_state(available, running, ...)`, and +/// `exclude_kind`. +/// * `insert_opts = path::to_fn`: a `fn() -> InsertOpts` whose options are +/// overlaid on the attribute defaults. +/// * `crate = "path"`: the path to `riverqueue` when it is renamed or +/// re-exported. +/// +/// Fields marked `#[river(unique)]` are the arguments hashed by `unique(by_args)`, +/// together with any `by_args` paths; without either, every argument is +/// hashed. Marking a field requires `unique(by_args)` on the type. Field +/// names follow Serde's serialization-side `rename` and `rename_all`. A +/// unique field may be conditionally omitted with `skip_serializing_if`, but +/// cannot be flattened or unconditionally skipped. +/// `by_args` paths separate nested names with `.`; escape a literal dot or +/// backslash with a backslash (for example, `"user\\.id"` selects the single +/// JSON name `user.id`). Tagged fields use their whole serialized name as one +/// component, even when that name contains a dot. +#[proc_macro_derive(JobArgs, attributes(river))] +pub fn derive_job_args(input: TokenStream) -> TokenStream { + let input = parse_macro_input!(input as DeriveInput); + expand_job_args(&input) + .unwrap_or_else(syn::Error::into_compile_error) + .into() +} + +#[derive(Default)] +struct TypeAttributes { + aliases: Vec, + krate: Option, + insert_opts: Option, + kind: Option, + max_attempts: Option, + pending: Option, + priority: Option, + queue: Option, + tags: Vec, + unique: Option, +} + +struct UniqueAttribute { + by_args: bool, + by_args_paths: Vec, + by_period: Option<(u64, u32)>, + by_queue: bool, + by_state: Option>, + exclude_kind: bool, +} + +const JOB_STATES: [(&str, &str); 8] = [ + ("available", "Available"), + ("cancelled", "Cancelled"), + ("completed", "Completed"), + ("discarded", "Discarded"), + ("pending", "Pending"), + ("retryable", "Retryable"), + ("running", "Running"), + ("scheduled", "Scheduled"), +]; + +const UNIQUE_REQUIRES_OPTION: &str = + "unique(...) requires at least one of by_args, by_period, by_queue, by_state, or exclude_kind"; + +const JOB_STATES_REQUIRED_FOR_UNIQUE: [&str; 4] = ["available", "pending", "running", "scheduled"]; + +fn parse_type_attributes(input: &DeriveInput) -> syn::Result { + let mut parsed = TypeAttributes::default(); + for attribute in &input.attrs { + if !attribute.path().is_ident("river") { + continue; + } + attribute.parse_nested_meta(|meta| { + if meta.path.is_ident("aliases") { + parsed.aliases = parse_string_list(&meta)?; + } else if meta.path.is_ident("crate") { + parsed.krate = Some(meta.value()?.parse::()?.parse()?); + } else if meta.path.is_ident("insert_opts") { + parsed.insert_opts = Some(meta.value()?.parse()?); + } else if meta.path.is_ident("kind") { + parsed.kind = Some(meta.value()?.parse()?); + } else if meta.path.is_ident("max_attempts") { + parsed.max_attempts = Some(meta.value()?.parse()?); + } else if meta.path.is_ident("pending") { + parsed.pending = Some(meta.value()?.parse()?); + } else if meta.path.is_ident("priority") { + parsed.priority = Some(meta.value()?.parse()?); + } else if meta.path.is_ident("queue") { + parsed.queue = Some(meta.value()?.parse()?); + } else if meta.path.is_ident("tags") { + parsed.tags = parse_string_list(&meta)?; + } else if meta.path.is_ident("unique") { + parsed.unique = Some(parse_unique(&meta)?); + } else { + return Err(meta.error("unsupported river type attribute")); + } + Ok(()) + })?; + } + Ok(parsed) +} + +fn parse_string_list(meta: &ParseNestedMeta<'_>) -> syn::Result> { + let content; + parenthesized!(content in meta.input); + Ok(Punctuated::::parse_terminated(&content)? + .into_iter() + .collect()) +} + +fn parse_unique(meta: &ParseNestedMeta<'_>) -> syn::Result { + let span = meta + .path + .get_ident() + .map_or_else(Span::call_site, Ident::span); + { + let lookahead = meta.input.fork(); + let content; + parenthesized!(content in lookahead); + if content.peek(LitStr) { + return Err(content + .error("declare unique argument paths with `unique(by_args(\"path\", ...))`")); + } + if content.is_empty() { + return Err(syn::Error::new(span, UNIQUE_REQUIRES_OPTION)); + } + } + let mut unique = UniqueAttribute { + by_args: false, + by_args_paths: Vec::new(), + by_period: None, + by_queue: false, + by_state: None, + exclude_kind: false, + }; + meta.parse_nested_meta(|option| { + if option.path.is_ident("by_args") { + unique.by_args = true; + if option.input.peek(token::Paren) { + unique.by_args_paths = parse_string_list(&option)?; + } + } else if option.path.is_ident("by_period") { + let period = option.value()?.parse::()?; + let nanos = parse_go_duration(&period.value()) + .map_err(|message| syn::Error::new_spanned(&period, message))?; + if nanos < 1_000_000_000 { + return Err(syn::Error::new_spanned( + &period, + "by_period must be at least one second", + )); + } + let seconds = u64::try_from(nanos / 1_000_000_000) + .map_err(|_| syn::Error::new_spanned(&period, "by_period is too large"))?; + let nanos = u32::try_from(nanos % 1_000_000_000).unwrap_or_default(); + unique.by_period = Some((seconds, nanos)); + } else if option.path.is_ident("by_queue") { + unique.by_queue = true; + } else if option.path.is_ident("by_state") { + let content; + parenthesized!(content in option.input); + let states = Punctuated::::parse_terminated(&content)? + .into_iter() + .collect::>(); + validate_unique_states(&states, &option)?; + unique.by_state = Some(states); + } else if option.path.is_ident("exclude_kind") { + unique.exclude_kind = true; + } else { + return Err(option.error( + "unsupported unique option; expected by_args, by_period, by_queue, by_state, or exclude_kind", + )); + } + Ok(()) + })?; + Ok(unique) +} + +fn validate_unique_states(states: &[Ident], option: &ParseNestedMeta<'_>) -> syn::Result<()> { + let mut seen = Vec::new(); + for state in states { + let name = state.to_string(); + if !JOB_STATES.iter().any(|(known, _)| *known == name) { + return Err(syn::Error::new_spanned( + state, + format!( + "unknown job state `{name}`; expected one of {}", + JOB_STATES + .iter() + .map(|(known, _)| *known) + .collect::>() + .join(", ") + ), + )); + } + if seen.contains(&name) { + return Err(syn::Error::new_spanned( + state, + format!("duplicate job state `{name}`"), + )); + } + seen.push(name); + } + let missing = JOB_STATES_REQUIRED_FOR_UNIQUE + .iter() + .filter(|required| !seen.iter().any(|state| state == *required)) + .copied() + .collect::>(); + if !missing.is_empty() { + return Err(option.error(format!( + "by_state must include available, pending, running, and scheduled; missing {}", + missing.join(", ") + ))); + } + Ok(()) +} + +/// Parses a Go `time.ParseDuration` string without a sign, such as `1h30m` +/// or `1.5h`, into nanoseconds. +fn parse_go_duration(text: &str) -> Result { + let invalid = || format!("invalid duration {text:?}; expected a Go duration such as \"1h\""); + if text.is_empty() { + return Err(invalid()); + } + let mut rest = text; + let mut total = 0_u128; + while !rest.is_empty() { + let number_end = rest + .find(|character: char| !(character.is_ascii_digit() || character == '.')) + .unwrap_or(rest.len()); + let (number, after_number) = rest.split_at(number_end); + let unit_end = after_number + .find(|character: char| character.is_ascii_digit() || character == '.') + .unwrap_or(after_number.len()); + let (unit, after_unit) = after_number.split_at(unit_end); + rest = after_unit; + + let unit_nanos: u128 = match unit { + "ns" => 1, + "us" | "\u{b5}s" | "\u{3bc}s" => 1_000, + "ms" => 1_000_000, + "s" => 1_000_000_000, + "m" => 60_000_000_000, + "h" => 3_600_000_000_000, + _ => return Err(invalid()), + }; + let (whole, fraction) = number.split_once('.').unwrap_or((number, "")); + if whole.is_empty() && fraction.is_empty() { + return Err(invalid()); + } + let whole = if whole.is_empty() { + 0 + } else { + whole.parse::().map_err(|_| invalid())? + }; + let mut value = whole.checked_mul(unit_nanos).ok_or_else(invalid)?; + if !fraction.is_empty() { + let scale = 10_u128 + .checked_pow(u32::try_from(fraction.len()).map_err(|_| invalid())?) + .ok_or_else(invalid)?; + let fraction = fraction.parse::().map_err(|_| invalid())?; + value = value + .checked_add(fraction.checked_mul(unit_nanos).ok_or_else(invalid)? / scale) + .ok_or_else(invalid)?; + } + total = total.checked_add(value).ok_or_else(invalid)?; + } + Ok(total) +} + +#[expect( + clippy::too_many_lines, + reason = "the derive's validation and generated items read the same parsed attributes" +)] +fn expand_job_args(input: &DeriveInput) -> syn::Result { + let attributes = parse_type_attributes(input)?; + let kind = attributes.kind.clone().ok_or_else(|| { + syn::Error::new_spanned(&input.ident, "JobArgs requires #[river(kind = \"...\")]") + })?; + validate_kind(&kind)?; + for alias in &attributes.aliases { + validate_kind(alias)?; + if alias.value() == kind.value() { + return Err(syn::Error::new_spanned( + alias, + "a kind alias cannot equal the primary kind", + )); + } + } + let mut alias_values = attributes + .aliases + .iter() + .map(LitStr::value) + .collect::>(); + alias_values.sort_unstable(); + if alias_values.windows(2).any(|pair| pair[0] == pair[1]) { + return Err(syn::Error::new_spanned( + &input.ident, + "kind aliases must be unique", + )); + } + if let Some(max_attempts) = &attributes.max_attempts + && !(1..=i16::MAX as u64).contains(&max_attempts.base10_parse::()?) + { + return Err(syn::Error::new_spanned( + max_attempts, + "max_attempts must be between 1 and 32767", + )); + } + if let Some(priority) = &attributes.priority + && !(1..=4).contains(&priority.base10_parse::()?) + { + return Err(syn::Error::new_spanned( + priority, + "priority must be between 1 and 4", + )); + } + if let Some(queue) = &attributes.queue { + validate_queue(queue)?; + } + for tag in &attributes.tags { + validate_tag(tag)?; + } + + let fields = match &input.data { + Data::Struct(data) => match &data.fields { + Fields::Named(fields) => &fields.named, + _ => { + return Err(syn::Error::new_spanned( + &input.ident, + "JobArgs can only be derived for a struct with named fields", + )); + } + }, + _ => { + return Err(syn::Error::new_spanned( + &input.ident, + "JobArgs can only be derived for structs", + )); + } + }; + + let rename_all = serde_serialize_rename_all(&input.attrs)?; + let mut available_json_fields = Vec::new(); + let mut unique_fields = Vec::new(); + for field in fields { + let mut unique = false; + for attribute in &field.attrs { + if attribute.path().is_ident("river") { + attribute.parse_nested_meta(|meta| { + if meta.path.is_ident("unique") { + unique = true; + return Ok(()); + } + Err(meta.error("unsupported river field attribute")) + })?; + } + } + let serde = serde_field_attributes(&field.attrs)?; + if unique && (serde.flatten || serde.skip || serde.skip_serializing) { + return Err(syn::Error::new_spanned( + field, + "#[river(unique)] cannot be combined with #[serde(flatten)], #[serde(skip)], or #[serde(skip_serializing)]", + )); + } + let field_ident = field.ident.as_ref().expect("named fields have identifiers"); + let rust_name = field_ident.unraw().to_string(); + let mut json_name = rename_all.as_ref().map_or_else( + || Ok(rust_name.clone()), + |rule| rename_field(&rust_name, rule), + )?; + if let Some(rename) = serde.rename { + json_name = rename.value(); + } + if !serde.flatten && !serde.skip && !serde.skip_serializing { + available_json_fields.push(json_name.clone()); + } + if unique { + if !attributes + .unique + .as_ref() + .is_some_and(|unique| unique.by_args) + { + return Err(syn::Error::new_spanned( + field_ident, + "#[river(unique)] fields are only hashed with `unique(by_args)`; add it to the type's #[river(...)] attribute", + )); + } + unique_fields.push(vec![LitStr::new(&json_name, field_ident.span())]); + } + } + if let Some(unique) = &attributes.unique { + for path in &unique.by_args_paths { + let components = parse_unique_path(path)?; + let first = &components[0]; + if !available_json_fields.iter().any(|field| field == first) { + return Err(syn::Error::new_spanned( + path, + "unique JSON path must start with a serialized field name", + )); + } + unique_fields.push( + components + .into_iter() + .map(|component| LitStr::new(&component, path.span())) + .collect(), + ); + } + } + validate_unique_paths(&unique_fields)?; + let unique_fields = unique_fields.iter().map(|path| quote!(&[#(#path),*])); + + let krate = attributes + .krate + .clone() + .unwrap_or_else(|| syn::parse_quote!(::riverqueue)); + let name = &input.ident; + let aliases = &attributes.aliases; + let (impl_generics, type_generics, where_clause) = input.generics.split_for_impl(); + let set_max_attempts = attributes + .max_attempts + .as_ref() + .map(|value| quote!(.with_max_attempts(#value))); + let set_pending = attributes + .pending + .as_ref() + .map(|value| quote!(.with_pending(#value))); + let set_priority = attributes + .priority + .as_ref() + .map(|value| quote!(.with_priority(#value))); + let set_queue = attributes + .queue + .as_ref() + .map(|value| quote!(.with_queue(#value))); + let set_tags = (!attributes.tags.is_empty()).then(|| { + let tags = &attributes.tags; + quote!(.with_tags([#(#tags),*])) + }); + let set_unique = attributes + .unique + .as_ref() + .map(|unique| expand_unique_opts(&krate, unique)); + // Bind the function's result with the expected type so a mismatched + // function is reported at its path. + let overlay = attributes.insert_opts.as_ref().map(|function| { + let call = quote_spanned!(function.span()=> #function()); + quote! { + .overlay({ + let overrides: #krate::InsertOpts = #call; + overrides + }) + } + }); + Ok(quote! { + impl #impl_generics #krate::JobArgs for #name #type_generics #where_clause { + const KIND: &'static str = #kind; + + fn kind_aliases() -> &'static [&'static str] { + &[#(#aliases),*] + } + + fn default_insert_opts() -> #krate::InsertOpts { + #krate::InsertOpts::default() + #set_max_attempts + #set_pending + #set_priority + #set_queue + #set_tags + #set_unique + #overlay + } + + fn unique_fields() -> &'static [&'static [&'static str]] { + &[#(#unique_fields),*] + } + } + }) +} + +fn expand_unique_opts(krate: &syn::Path, unique: &UniqueAttribute) -> proc_macro2::TokenStream { + let by_args = unique.by_args.then(|| quote!(.with_by_args(true))); + let by_period = unique.by_period.map( + |(seconds, nanos)| quote!(.with_by_period(::core::time::Duration::new(#seconds, #nanos))), + ); + let by_queue = unique.by_queue.then(|| quote!(.with_by_queue(true))); + let by_state = unique.by_state.as_ref().map(|states| { + let variants = states.iter().map(|state| { + let variant = JOB_STATES + .iter() + .find(|(name, _)| state == name) + .map_or("Available", |(_, variant)| variant); + let variant = Ident::new(variant, state.span()); + quote!(#krate::JobState::#variant) + }); + quote!(.with_by_state([#(#variants),*])) + }); + let exclude_kind = unique + .exclude_kind + .then(|| quote!(.with_exclude_kind(true))); + quote! { + .with_unique( + #krate::UniqueOpts::new() + #by_args + #by_period + #by_queue + #by_state + #exclude_kind + ) + } +} + +/// Decode the convenience dotted syntax into literal JSON field names. +fn parse_unique_path(path: &LitStr) -> syn::Result> { + let mut parts = Vec::new(); + let mut current = String::new(); + let mut escaped = false; + for character in path.value().chars() { + if escaped { + current.push(character); + escaped = false; + } else if character == '\\' { + escaped = true; + } else if character == '.' { + parts.push(std::mem::take(&mut current)); + } else { + current.push(character); + } + } + if escaped { + return Err(syn::Error::new_spanned( + path, + "unique JSON path ends in an escape", + )); + } + parts.push(current); + if parts.iter().any(String::is_empty) { + return Err(syn::Error::new_spanned( + path, + "unique JSON path segments cannot be empty", + )); + } + Ok(parts) +} + +fn validate_unique_paths(paths: &[Vec]) -> syn::Result<()> { + for path in paths { + if path.iter().any(|segment| segment.value().is_empty()) { + return Err(syn::Error::new_spanned( + &path[0], + "unique JSON path segments cannot be empty", + )); + } + if let Some(segment) = path.iter().find(|segment| { + let value = segment.value(); + value.bytes().all(|byte| byte.is_ascii_digit()) || value == "-1" + }) { + return Err(syn::Error::new_spanned( + segment, + "numeric unique JSON path segments require array semantics that are not yet supported", + )); + } + } + for path in paths { + let value = path.iter().map(LitStr::value).collect::>(); + if let Some(other) = paths.iter().find(|other| { + other.len() > path.len() + && other + .iter() + .zip(&value) + .all(|(segment, value)| segment.value() == *value) + }) { + return Err(syn::Error::new_spanned( + &other[0], + format!( + "unique JSON path {:?} is inside another unique path {value:?}", + other.iter().map(LitStr::value).collect::>() + ), + )); + } + } + Ok(()) +} + +#[derive(Default)] +struct SerdeFieldAttributes { + flatten: bool, + rename: Option, + skip: bool, + skip_serializing: bool, +} + +fn serde_field_attributes(attributes: &[syn::Attribute]) -> syn::Result { + let mut parsed = SerdeFieldAttributes::default(); + for attribute in attributes { + if !attribute.path().is_ident("serde") { + continue; + } + for meta in serde_metas(attribute)? { + match meta { + Meta::Path(path) if path.is_ident("flatten") => parsed.flatten = true, + Meta::Path(path) if path.is_ident("skip") => parsed.skip = true, + Meta::Path(path) if path.is_ident("skip_serializing") => { + parsed.skip_serializing = true; + } + Meta::NameValue(meta) if meta.path.is_ident("rename") => { + parsed.rename = Some(meta_lit_str(&meta)?); + } + Meta::List(meta) if meta.path.is_ident("rename") => { + if let Some(rename) = serialize_name_from_list(&meta)? { + parsed.rename = Some(rename); + } + } + _ => {} + } + } + } + Ok(parsed) +} + +fn serde_serialize_rename_all(attributes: &[syn::Attribute]) -> syn::Result> { + let mut rename_all = None; + for attribute in attributes { + if !attribute.path().is_ident("serde") { + continue; + } + for meta in serde_metas(attribute)? { + match meta { + Meta::NameValue(meta) if meta.path.is_ident("rename_all") => { + rename_all = Some(meta_lit_str(&meta)?); + } + Meta::List(meta) if meta.path.is_ident("rename_all") => { + if let Some(rename) = serialize_name_from_list(&meta)? { + rename_all = Some(rename); + } + } + _ => {} + } + } + } + Ok(rename_all) +} + +fn serde_metas(attribute: &syn::Attribute) -> syn::Result> { + attribute.parse_args_with(Punctuated::::parse_terminated) +} + +fn serialize_name_from_list(meta: &syn::MetaList) -> syn::Result> { + let nested = meta.parse_args_with(Punctuated::::parse_terminated)?; + nested + .iter() + .find_map(|meta| match meta { + Meta::NameValue(meta) if meta.path.is_ident("serialize") => Some(meta_lit_str(meta)), + _ => None, + }) + .transpose() +} + +fn meta_lit_str(meta: &syn::MetaNameValue) -> syn::Result { + match &meta.value { + syn::Expr::Lit(expression) => match &expression.lit { + Lit::Str(value) => Ok(value.clone()), + _ => Err(syn::Error::new_spanned( + &meta.value, + "expected string literal", + )), + }, + _ => Err(syn::Error::new_spanned( + &meta.value, + "expected string literal", + )), + } +} + +fn rename_field(field: &str, rule: &LitStr) -> syn::Result { + let renamed = match rule.value().as_str() { + "lowercase" | "snake_case" => field.to_owned(), + "UPPERCASE" | "SCREAMING_SNAKE_CASE" => field.to_ascii_uppercase(), + "PascalCase" => rename_field_pascal_case(field), + "camelCase" => { + let pascal = rename_field_pascal_case(field); + let mut characters = pascal.chars(); + characters.next().map_or_else(String::new, |first| { + first.to_ascii_lowercase().to_string() + characters.as_str() + }) + } + "kebab-case" => field.replace('_', "-"), + "SCREAMING-KEBAB-CASE" => field.to_ascii_uppercase().replace('_', "-"), + unsupported => { + return Err(syn::Error::new_spanned( + rule, + format!("unsupported serde rename rule {unsupported:?}"), + )); + } + }; + Ok(renamed) +} + +fn rename_field_pascal_case(field: &str) -> String { + let mut renamed = String::new(); + let mut capitalize = true; + for character in field.chars() { + if character == '_' { + capitalize = true; + } else if capitalize { + renamed.push(character.to_ascii_uppercase()); + capitalize = false; + } else { + renamed.push(character); + } + } + renamed +} + +fn validate_kind(kind: &LitStr) -> syn::Result<()> { + let value = kind.value(); + let mut characters = value.chars(); + if value.len() < 2 + || value.len() >= 128 + || !characters + .next() + .is_some_and(|character| character == '_' || character.is_ascii_alphanumeric()) + || !characters.all(|character| { + character.is_ascii_alphanumeric() + || matches!( + character, + '_' | '-' | '[' | ']' | '<' | '>' | '/' | '.' | '·' | ':' | '+' + ) + }) + { + return Err(syn::Error::new_spanned(kind, "invalid River job kind")); + } + Ok(()) +} + +fn validate_queue(queue: &LitStr) -> syn::Result<()> { + let value = queue.value(); + let mut characters = value.chars(); + let valid = value.len() <= 64 + && characters + .next() + .is_some_and(|character| character.is_ascii_lowercase() || character.is_ascii_digit()) + && value.chars().all(|character| { + character.is_ascii_lowercase() + || character.is_ascii_digit() + || matches!(character, '_' | '-') + }) + && !value.contains("__") + && !value.contains("--") + && !value.contains("_-") + && !value.contains("-_") + && value + .chars() + .next_back() + .is_some_and(|character| character.is_ascii_lowercase() || character.is_ascii_digit()); + if !valid { + return Err(syn::Error::new_spanned(queue, "invalid River queue name")); + } + Ok(()) +} + +/// Validates a tag with River's rules: 3 to 255 bytes of ASCII word +/// characters and `-`, starting and ending with a word character. +fn validate_tag(tag: &LitStr) -> syn::Result<()> { + let value = tag.value(); + let is_word = |character: char| character == '_' || character.is_ascii_alphanumeric(); + let valid = (3..=255).contains(&value.len()) + && value.chars().next().is_some_and(is_word) + && value.chars().next_back().is_some_and(is_word) + && value + .chars() + .all(|character| is_word(character) || character == '-'); + if !valid { + return Err(syn::Error::new_spanned( + tag, + "invalid River tag; tags contain 3 to 255 ASCII letters, digits, `_`, or `-`, and start and end with a letter, digit, or `_`", + )); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn derive(source: &str) -> syn::Result { + expand_job_args(&syn::parse_str(source).expect("valid Rust syntax")) + } + + fn compact(tokens: &proc_macro2::TokenStream) -> String { + tokens.to_string().replace(' ', "") + } + + #[test] + fn accepts_complete_job_args_configuration() { + let expanded = derive( + r#" + #[derive(serde::Serialize)] + #[serde(rename_all = "camelCase")] + #[river( + kind = "email.send", + aliases("email_send_v1"), + max_attempts = 8, + pending = true, + priority = 2, + queue = "email-critical", + tags("email", "outbound-mail"), + unique( + by_args("account.id"), + by_period = "1h30m", + by_queue, + by_state(available, pending, running, scheduled, retryable), + exclude_kind + ), + insert_opts = email_insert_opts + )] + struct EmailArgs { + account: Account, + #[river(unique)] + message_id: String, + } + "#, + ) + .unwrap(); + let expanded = compact(&expanded); + + assert!(expanded.contains("\"email.send\"")); + assert!(expanded.contains("\"email_send_v1\"")); + assert!(expanded.contains(".with_tags([\"email\",\"outbound-mail\"])")); + assert!(expanded.contains(".with_by_args(true)")); + assert!(expanded.contains(".with_by_period(::core::time::Duration::new(5400u64,0u32))")); + assert!(expanded.contains(".with_by_queue(true)")); + assert!(expanded.contains("::riverqueue::JobState::Retryable")); + assert!(expanded.contains(".with_exclude_kind(true)")); + assert!(expanded.contains( + ".overlay({letoverrides:::riverqueue::InsertOpts=email_insert_opts();overrides})" + )); + assert!(expanded.contains("&[&[\"messageId\"],&[\"account\",\"id\"]]")); + } + + #[test] + fn follows_serde_serialization_names() { + let expanded = derive( + r#" + #[serde(rename_all(serialize = "SCREAMING-KEBAB-CASE", deserialize = "camelCase"))] + #[river(kind = "serde_names", unique(by_args))] + struct SerdeNames { + #[river(unique)] + first_value: String, + #[river(unique)] + #[serde(rename(serialize = "wire-name", deserialize = "inputName"))] + second_value: String, + #[river(unique)] + r#type: String, + } + "#, + ) + .unwrap() + .to_string(); + + assert!(expanded.contains("FIRST-VALUE")); + assert!(expanded.contains("wire-name")); + assert!(expanded.contains("TYPE")); + assert!(!expanded.contains("inputName")); + } + + #[test] + fn parses_go_durations() { + for (text, nanos) in [ + ("1s", 1_000_000_000), + ("1h", 3_600_000_000_000), + ("1h30m", 5_400_000_000_000), + ("1.5h", 5_400_000_000_000), + ("90m", 5_400_000_000_000), + ("1500ms", 1_500_000_000), + ("2s500ms", 2_500_000_000), + ("1\u{b5}s", 1_000), + ("7ns", 7), + ] { + assert_eq!(parse_go_duration(text), Ok(nanos), "{text}"); + } + for text in ["", "1", "h", "1x", "-1h", "1.h.", "."] { + assert!(parse_go_duration(text).is_err(), "{text}"); + } + } + + #[test] + fn permits_conditionally_omitted_unique_fields() { + let expanded = derive( + r#" + #[river(kind = "optional_unique", unique(by_args))] + struct OptionalUnique { + #[river(unique)] + #[serde(skip_serializing_if = "Option::is_none")] + optional: Option, + } + "#, + ) + .unwrap() + .to_string(); + + assert!(expanded.contains("optional")); + } + + #[test] + fn rejects_serde_attributes_that_hide_unique_fields() { + for serde_attribute in ["flatten", "skip", "skip_serializing"] { + let source = format!( + r#" + #[river(kind = "invalid_serde", unique(by_args))] + struct InvalidSerde {{ + #[river(unique)] + #[serde({serde_attribute})] + value: String, + }} + "#, + ); + let error = derive(&source).expect_err("attribute combination should be rejected"); + + assert!( + error + .to_string() + .contains("#[river(unique)] cannot be combined"), + "unexpected error for {serde_attribute}: {error}" + ); + } + + for serde_attribute in ["flatten", "skip", "skip_serializing"] { + let source = format!( + r#" + #[river(kind = "invalid_serde_path", unique(by_args("value")))] + struct InvalidSerdePath {{ + #[serde({serde_attribute})] + value: String, + }} + "#, + ); + let error = derive(&source).expect_err("hidden unique path should be rejected"); + + assert!( + error + .to_string() + .contains("unique JSON path must start with a serialized field name"), + "unexpected path error for {serde_attribute}: {error}" + ); + } + } + + #[test] + fn rejects_unsupported_serde_rename_rule() { + let error = derive( + r#" + #[serde(rename_all(serialize = "Title Case"))] + #[river(kind = "invalid_rename")] + struct InvalidRename { value: String } + "#, + ) + .expect_err("rename rule should be rejected"); + + assert!( + error + .to_string() + .contains("unsupported serde rename rule \"Title Case\"") + ); + } + + #[test] + fn rejects_invalid_job_args_configuration() { + let cases = [ + ( + "struct MissingKind { value: String }", + "JobArgs requires #[river(kind = \"...\")]", + ), + ( + r#"#[river(kind = "x")] struct InvalidKind { value: String }"#, + "invalid River job kind", + ), + ( + r#"#[river(kind = "valid", aliases("valid"))] struct DuplicateKind { value: String }"#, + "a kind alias cannot equal the primary kind", + ), + ( + r#"#[river(kind = "valid", priority = 5)] struct InvalidPriority { value: String }"#, + "priority must be between 1 and 4", + ), + ( + r#"#[river(kind = "valid", queue = "Invalid")] struct InvalidQueue { value: String }"#, + "invalid River queue name", + ), + ( + r#"#[river(kind = "valid", tags("x"))] struct InvalidTag { value: String }"#, + "invalid River tag", + ), + ( + r#"#[river(kind = "valid", unique(by_args("missing.id")))] struct InvalidPath { value: String }"#, + "unique JSON path must start with a serialized field name", + ), + ( + r#"#[river(kind = "valid", unique("value"))] struct OldPathSyntax { value: String }"#, + "declare unique argument paths with `unique(by_args(", + ), + ( + r#"#[river(kind = "valid", unique())] struct EmptyUnique { value: String }"#, + "unique(...) requires at least one of", + ), + ( + r#"#[river(kind = "valid", unique(by_arg))] struct UnknownUnique { value: String }"#, + "unsupported unique option", + ), + ( + r#"#[river(kind = "valid", unique(by_period = "500ms"))] struct ShortPeriod { value: String }"#, + "by_period must be at least one second", + ), + ( + r#"#[river(kind = "valid", unique(by_period = "1 hour"))] struct BadPeriod { value: String }"#, + "invalid duration", + ), + ( + r#"#[river(kind = "valid", unique(by_state(available, running)))] struct MissingStates { value: String }"#, + "missing pending, scheduled", + ), + ( + r#"#[river(kind = "valid", unique(by_state(available, pending, running, scheduled, done)))] struct UnknownState { value: String }"#, + "unknown job state `done`", + ), + ( + r#"#[river(kind = "valid", unique(by_state(available, available, pending, running, scheduled)))] struct DuplicateState { value: String }"#, + "duplicate job state `available`", + ), + ( + r#"#[river(kind = "valid", unique(by_queue))] struct UniqueFieldWithoutArgs { #[river(unique)] value: String }"#, + "#[river(unique)] fields are only hashed with `unique(by_args)`", + ), + ( + r#"#[river(kind = "valid")] struct UniqueFieldWithoutUnique { #[river(unique)] value: String }"#, + "#[river(unique)] fields are only hashed with `unique(by_args)`", + ), + ( + r#"#[river(kind = "valid", unique(by_args("value.0")))] struct IndexPath { value: String }"#, + "require array semantics", + ), + ( + r#"#[river(kind = "valid", unique(by_args("value", "value.id")))] struct NestedPath { value: String }"#, + "is inside another unique path", + ), + ]; + + for (source, message) in cases { + let error = derive(source).expect_err("configuration should be rejected"); + assert!( + error.to_string().contains(message), + "unexpected error for {source}: {error}" + ); + } + } + + #[test] + fn uses_configured_crate_path() { + let expanded = derive( + r#" + #[river(kind = "renamed", crate = "::my_app::river", unique(by_state(available, pending, running, scheduled)))] + struct Renamed { value: String } + "#, + ) + .unwrap(); + let expanded = compact(&expanded); + + assert!(expanded.contains("impl::my_app::river::JobArgsforRenamed")); + assert!(expanded.contains("::my_app::river::JobState::Available")); + assert!(!expanded.contains("::riverqueue::")); + } +} diff --git a/rust/riverqueue-macros/tests/derive.rs b/rust/riverqueue-macros/tests/derive.rs new file mode 100644 index 000000000..6b615f0b8 --- /dev/null +++ b/rust/riverqueue-macros/tests/derive.rs @@ -0,0 +1,181 @@ +//! Runtime behavior of `#[derive(JobArgs)]` expansions. + +use std::time::Duration; + +use riverqueue::{InsertOpts, JobArgs, JobState, ScheduleOverride, UniqueOpts}; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, JobArgs, Serialize)] +#[river( + kind = "new_kind", + aliases("old_kind", "older_kind"), + max_attempts = 7, + pending = true, + priority = 3, + queue = "critical_jobs", + tags("billing", "priority-customer") +)] +struct ArgsWithDefaults {} + +#[test] +fn provides_aliases_and_insert_defaults() { + assert_eq!(ArgsWithDefaults::KIND, "new_kind"); + assert_eq!(ArgsWithDefaults::kind_aliases(), ["old_kind", "older_kind"]); + let opts = ArgsWithDefaults::default_insert_opts(); + assert_eq!(opts.max_attempts(), Some(7)); + assert_eq!(opts.pending(), Some(true)); + assert_eq!(opts.priority(), Some(3)); + assert_eq!(opts.queue(), Some("critical_jobs")); + assert_eq!( + opts.tags(), + Some(&["billing".to_owned(), "priority-customer".to_owned()][..]) + ); + assert!(opts.unique().is_none()); + assert_eq!(ArgsWithDefaults::unique_fields(), [] as [&[&str]; 0]); +} + +#[derive(Deserialize, JobArgs, Serialize)] +#[river( + kind = "unique_invoice", + unique( + by_args("customer.id"), + by_period = "1h30m", + by_queue, + by_state(available, pending, running, scheduled, retryable), + exclude_kind + ) +)] +struct UniqueArgs { + customer: Customer, + #[river(unique)] + #[serde(rename = "invoiceNumber")] + invoice_number: String, + note: String, +} + +#[derive(Deserialize, Serialize)] +struct Customer { + id: i64, +} + +#[test] +fn declares_job_type_uniqueness() { + let opts = UniqueArgs::default_insert_opts(); + let unique = opts.unique().expect("unique options"); + assert!(unique.by_args()); + assert_eq!(unique.by_period(), Some(Duration::from_mins(90))); + assert!(unique.by_queue()); + assert_eq!( + unique.by_state(), + Some( + &[ + JobState::Available, + JobState::Pending, + JobState::Running, + JobState::Scheduled, + JobState::Retryable, + ][..] + ) + ); + assert!(unique.exclude_kind()); + assert_eq!( + UniqueArgs::unique_fields(), + [&["invoiceNumber"][..], &["customer", "id"][..]] + ); +} + +#[derive(Deserialize, JobArgs, Serialize)] +#[river(kind = "literal_paths", unique(by_args("user.id", "user\\.id")))] +struct LiteralPaths { + #[river(unique)] + #[serde(rename = "@user")] + at: String, + #[river(unique)] + #[serde(rename = ":id")] + colon: String, + #[serde(rename = "user.id")] + literal: String, + user: Customer, + #[river(unique)] + #[serde(rename = "é")] + unicode: String, +} + +#[test] +fn separates_literal_and_nested_unique_fields() { + assert_eq!( + LiteralPaths::unique_fields(), + [ + &["@user"][..], + &[":id"][..], + &["é"][..], + &["user", "id"][..], + &["user.id"][..], + ] + ); +} + +#[derive(Deserialize, JobArgs, Serialize)] +#[river(kind = "all_args_unique", unique(by_args))] +struct AllArgsUnique { + value: String, +} + +#[test] +fn by_args_without_fields_hashes_every_argument() { + let opts = AllArgsUnique::default_insert_opts(); + assert!(opts.unique().is_some_and(UniqueOpts::by_args)); + assert_eq!(AllArgsUnique::unique_fields(), [] as [&[&str]; 0]); +} + +fn overlaid_insert_opts() -> InsertOpts { + InsertOpts::default() + .with_priority(2) + .with_unique(UniqueOpts::new().with_by_queue(true)) +} + +#[derive(Deserialize, JobArgs, Serialize)] +#[river( + kind = "overlaid", + priority = 4, + queue = "attribute_queue", + unique(by_args), + insert_opts = overlaid_insert_opts +)] +struct OverlaidArgs {} + +#[test] +fn insert_opts_function_overlays_attribute_defaults() { + let opts = OverlaidArgs::default_insert_opts(); + // The function's options win; the attribute's remain where it sets none. + assert_eq!(opts.priority(), Some(2)); + assert_eq!(opts.queue(), Some("attribute_queue")); + assert_eq!(opts.scheduled_at(), ScheduleOverride::Inherit); + let unique = opts.unique().expect("unique options"); + assert!(unique.by_queue()); + assert!(!unique.by_args()); +} + +mod reexport { + pub use riverqueue as river; +} + +#[derive(Deserialize, JobArgs, Serialize)] +#[river( + kind = "renamed_crate", + crate = "reexport::river", + unique(by_state(available, pending, running, scheduled)) +)] +struct RenamedCrateArgs {} + +#[test] +fn crate_attribute_selects_the_riverqueue_path() { + assert_eq!(RenamedCrateArgs::KIND, "renamed_crate"); + assert_eq!( + RenamedCrateArgs::default_insert_opts() + .unique() + .and_then(UniqueOpts::by_state) + .map(<[JobState]>::len), + Some(4) + ); +} diff --git a/rust/riverqueue-macros/tests/ui.rs b/rust/riverqueue-macros/tests/ui.rs new file mode 100644 index 000000000..781b9b70e --- /dev/null +++ b/rust/riverqueue-macros/tests/ui.rs @@ -0,0 +1,7 @@ +//! Compile-time diagnostics of `#[derive(JobArgs)]`. + +#[test] +fn derive_errors() { + let cases = trybuild::TestCases::new(); + cases.compile_fail("tests/ui/*.rs"); +} diff --git a/rust/riverqueue-macros/tests/ui/invalid_tag.rs b/rust/riverqueue-macros/tests/ui/invalid_tag.rs new file mode 100644 index 000000000..f944f063f --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/invalid_tag.rs @@ -0,0 +1,10 @@ +use riverqueue::JobArgs; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, JobArgs, Serialize)] +#[river(kind = "invoice", tags("billing", "x"))] +struct InvoiceArgs { + invoice_number: String, +} + +fn main() {} diff --git a/rust/riverqueue-macros/tests/ui/invalid_tag.stderr b/rust/riverqueue-macros/tests/ui/invalid_tag.stderr new file mode 100644 index 000000000..69e9f09d8 --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/invalid_tag.stderr @@ -0,0 +1,5 @@ +error: invalid River tag; tags contain 3 to 255 ASCII letters, digits, `_`, or `-`, and start and end with a letter, digit, or `_` + --> tests/ui/invalid_tag.rs:5:43 + | +5 | #[river(kind = "invoice", tags("billing", "x"))] + | ^^^ diff --git a/rust/riverqueue-macros/tests/ui/missing_kind.rs b/rust/riverqueue-macros/tests/ui/missing_kind.rs new file mode 100644 index 000000000..c814520bc --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/missing_kind.rs @@ -0,0 +1,10 @@ +use riverqueue::JobArgs; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, JobArgs, Serialize)] +#[river(queue = "invoices")] +struct InvoiceArgs { + invoice_number: String, +} + +fn main() {} diff --git a/rust/riverqueue-macros/tests/ui/missing_kind.stderr b/rust/riverqueue-macros/tests/ui/missing_kind.stderr new file mode 100644 index 000000000..cef0a1cf8 --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/missing_kind.stderr @@ -0,0 +1,5 @@ +error: JobArgs requires #[river(kind = "...")] + --> tests/ui/missing_kind.rs:6:8 + | +6 | struct InvoiceArgs { + | ^^^^^^^^^^^ diff --git a/rust/riverqueue-macros/tests/ui/unique_empty.rs b/rust/riverqueue-macros/tests/ui/unique_empty.rs new file mode 100644 index 000000000..4279f906d --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_empty.rs @@ -0,0 +1,10 @@ +use riverqueue::JobArgs; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, JobArgs, Serialize)] +#[river(kind = "invoice", unique())] +struct InvoiceArgs { + invoice_number: String, +} + +fn main() {} diff --git a/rust/riverqueue-macros/tests/ui/unique_empty.stderr b/rust/riverqueue-macros/tests/ui/unique_empty.stderr new file mode 100644 index 000000000..e179e0beb --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_empty.stderr @@ -0,0 +1,5 @@ +error: unique(...) requires at least one of by_args, by_period, by_queue, by_state, or exclude_kind + --> tests/ui/unique_empty.rs:5:27 + | +5 | #[river(kind = "invoice", unique())] + | ^^^^^^ diff --git a/rust/riverqueue-macros/tests/ui/unique_field_without_by_args.rs b/rust/riverqueue-macros/tests/ui/unique_field_without_by_args.rs new file mode 100644 index 000000000..c77717c0d --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_field_without_by_args.rs @@ -0,0 +1,11 @@ +use riverqueue::JobArgs; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, JobArgs, Serialize)] +#[river(kind = "invoice", unique(by_queue))] +struct InvoiceArgs { + #[river(unique)] + invoice_number: String, +} + +fn main() {} diff --git a/rust/riverqueue-macros/tests/ui/unique_field_without_by_args.stderr b/rust/riverqueue-macros/tests/ui/unique_field_without_by_args.stderr new file mode 100644 index 000000000..afa46812b --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_field_without_by_args.stderr @@ -0,0 +1,5 @@ +error: #[river(unique)] fields are only hashed with `unique(by_args)`; add it to the type's #[river(...)] attribute + --> tests/ui/unique_field_without_by_args.rs:8:5 + | +8 | invoice_number: String, + | ^^^^^^^^^^^^^^ diff --git a/rust/riverqueue-macros/tests/ui/unique_path_numeric.rs b/rust/riverqueue-macros/tests/ui/unique_path_numeric.rs new file mode 100644 index 000000000..3501e5046 --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_path_numeric.rs @@ -0,0 +1,10 @@ +use riverqueue::JobArgs; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, JobArgs, Serialize)] +#[river(kind = "invoice", unique(by_args("lines.0")))] +struct InvoiceArgs { + lines: Vec, +} + +fn main() {} diff --git a/rust/riverqueue-macros/tests/ui/unique_path_numeric.stderr b/rust/riverqueue-macros/tests/ui/unique_path_numeric.stderr new file mode 100644 index 000000000..72c8dec20 --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_path_numeric.stderr @@ -0,0 +1,5 @@ +error: numeric unique JSON path segments require array semantics that are not yet supported + --> tests/ui/unique_path_numeric.rs:5:42 + | +5 | #[river(kind = "invoice", unique(by_args("lines.0")))] + | ^^^^^^^^^ diff --git a/rust/riverqueue-macros/tests/ui/unique_path_old_syntax.rs b/rust/riverqueue-macros/tests/ui/unique_path_old_syntax.rs new file mode 100644 index 000000000..d57744cc8 --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_path_old_syntax.rs @@ -0,0 +1,10 @@ +use riverqueue::JobArgs; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, JobArgs, Serialize)] +#[river(kind = "invoice", unique("invoice_number"))] +struct InvoiceArgs { + invoice_number: String, +} + +fn main() {} diff --git a/rust/riverqueue-macros/tests/ui/unique_path_old_syntax.stderr b/rust/riverqueue-macros/tests/ui/unique_path_old_syntax.stderr new file mode 100644 index 000000000..2bb09f6ab --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_path_old_syntax.stderr @@ -0,0 +1,5 @@ +error: declare unique argument paths with `unique(by_args("path", ...))` + --> tests/ui/unique_path_old_syntax.rs:5:34 + | +5 | #[river(kind = "invoice", unique("invoice_number"))] + | ^^^^^^^^^^^^^^^^ diff --git a/rust/riverqueue-macros/tests/ui/unique_path_unknown_field.rs b/rust/riverqueue-macros/tests/ui/unique_path_unknown_field.rs new file mode 100644 index 000000000..9c807da17 --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_path_unknown_field.rs @@ -0,0 +1,10 @@ +use riverqueue::JobArgs; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, JobArgs, Serialize)] +#[river(kind = "invoice", unique(by_args("customer.id")))] +struct InvoiceArgs { + invoice_number: String, +} + +fn main() {} diff --git a/rust/riverqueue-macros/tests/ui/unique_path_unknown_field.stderr b/rust/riverqueue-macros/tests/ui/unique_path_unknown_field.stderr new file mode 100644 index 000000000..d8e5e3961 --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_path_unknown_field.stderr @@ -0,0 +1,5 @@ +error: unique JSON path must start with a serialized field name + --> tests/ui/unique_path_unknown_field.rs:5:42 + | +5 | #[river(kind = "invoice", unique(by_args("customer.id")))] + | ^^^^^^^^^^^^^ diff --git a/rust/riverqueue-macros/tests/ui/unique_period_invalid.rs b/rust/riverqueue-macros/tests/ui/unique_period_invalid.rs new file mode 100644 index 000000000..6fee92e8c --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_period_invalid.rs @@ -0,0 +1,10 @@ +use riverqueue::JobArgs; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, JobArgs, Serialize)] +#[river(kind = "invoice", unique(by_period = "1 hour"))] +struct InvoiceArgs { + invoice_number: String, +} + +fn main() {} diff --git a/rust/riverqueue-macros/tests/ui/unique_period_invalid.stderr b/rust/riverqueue-macros/tests/ui/unique_period_invalid.stderr new file mode 100644 index 000000000..3c798dc97 --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_period_invalid.stderr @@ -0,0 +1,5 @@ +error: invalid duration "1 hour"; expected a Go duration such as "1h" + --> tests/ui/unique_period_invalid.rs:5:46 + | +5 | #[river(kind = "invoice", unique(by_period = "1 hour"))] + | ^^^^^^^^ diff --git a/rust/riverqueue-macros/tests/ui/unique_period_too_short.rs b/rust/riverqueue-macros/tests/ui/unique_period_too_short.rs new file mode 100644 index 000000000..d319363b8 --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_period_too_short.rs @@ -0,0 +1,10 @@ +use riverqueue::JobArgs; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, JobArgs, Serialize)] +#[river(kind = "invoice", unique(by_period = "500ms"))] +struct InvoiceArgs { + invoice_number: String, +} + +fn main() {} diff --git a/rust/riverqueue-macros/tests/ui/unique_period_too_short.stderr b/rust/riverqueue-macros/tests/ui/unique_period_too_short.stderr new file mode 100644 index 000000000..63fb63bc3 --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_period_too_short.stderr @@ -0,0 +1,5 @@ +error: by_period must be at least one second + --> tests/ui/unique_period_too_short.rs:5:46 + | +5 | #[river(kind = "invoice", unique(by_period = "500ms"))] + | ^^^^^^^ diff --git a/rust/riverqueue-macros/tests/ui/unique_state_unknown.rs b/rust/riverqueue-macros/tests/ui/unique_state_unknown.rs new file mode 100644 index 000000000..4b51a62a4 --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_state_unknown.rs @@ -0,0 +1,10 @@ +use riverqueue::JobArgs; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, JobArgs, Serialize)] +#[river(kind = "invoice", unique(by_state(available, pending, running, scheduled, done)))] +struct InvoiceArgs { + invoice_number: String, +} + +fn main() {} diff --git a/rust/riverqueue-macros/tests/ui/unique_state_unknown.stderr b/rust/riverqueue-macros/tests/ui/unique_state_unknown.stderr new file mode 100644 index 000000000..a59243d9d --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_state_unknown.stderr @@ -0,0 +1,5 @@ +error: unknown job state `done`; expected one of available, cancelled, completed, discarded, pending, retryable, running, scheduled + --> tests/ui/unique_state_unknown.rs:5:83 + | +5 | #[river(kind = "invoice", unique(by_state(available, pending, running, scheduled, done)))] + | ^^^^ diff --git a/rust/riverqueue-macros/tests/ui/unique_states_missing_required.rs b/rust/riverqueue-macros/tests/ui/unique_states_missing_required.rs new file mode 100644 index 000000000..fa2c1647c --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_states_missing_required.rs @@ -0,0 +1,10 @@ +use riverqueue::JobArgs; +use serde::{Deserialize, Serialize}; + +#[derive(Deserialize, JobArgs, Serialize)] +#[river(kind = "invoice", unique(by_state(available, running)))] +struct InvoiceArgs { + invoice_number: String, +} + +fn main() {} diff --git a/rust/riverqueue-macros/tests/ui/unique_states_missing_required.stderr b/rust/riverqueue-macros/tests/ui/unique_states_missing_required.stderr new file mode 100644 index 000000000..f9686f364 --- /dev/null +++ b/rust/riverqueue-macros/tests/ui/unique_states_missing_required.stderr @@ -0,0 +1,5 @@ +error: by_state must include available, pending, running, and scheduled; missing pending, scheduled + --> tests/ui/unique_states_missing_required.rs:5:34 + | +5 | #[river(kind = "invoice", unique(by_state(available, running)))] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/rust/riverqueue-migrate/Cargo.toml b/rust/riverqueue-migrate/Cargo.toml new file mode 100644 index 000000000..536372b6b --- /dev/null +++ b/rust/riverqueue-migrate/Cargo.toml @@ -0,0 +1,46 @@ +[package] +name = "riverqueue-migrate" +description = "Database migrations for River's Rust client" +keywords = ["background", "jobs", "migrations", "postgres", "sqlite"] +categories = ["database"] +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +readme = "README.md" +repository.workspace = true +homepage.workspace = true + +[package.metadata.docs.rs] +all-features = true + +[features] +default = ["postgres"] +postgres = ["sqlx/postgres"] +sqlite = ["sqlx/sqlite"] + +[dependencies] +sqlx.workspace = true +thiserror.workspace = true +tokio.workspace = true + +[[example]] +name = "migrate" +required-features = ["postgres"] + +[[test]] +name = "postgres" +required-features = ["postgres"] + +[[test]] +name = "sqlite" +required-features = ["sqlite"] + +[dev-dependencies] +serde_json.workspace = true +# The PostgreSQL tests decode schema snapshots as JSON. +sqlx = { workspace = true, features = ["json"] } +tokio = { workspace = true, features = ["rt-multi-thread"] } + +[lints] +workspace = true diff --git a/rust/riverqueue-migrate/LICENSE b/rust/riverqueue-migrate/LICENSE new file mode 120000 index 000000000..30cff7403 --- /dev/null +++ b/rust/riverqueue-migrate/LICENSE @@ -0,0 +1 @@ +../../LICENSE \ No newline at end of file diff --git a/rust/riverqueue-migrate/README.md b/rust/riverqueue-migrate/README.md new file mode 100644 index 000000000..83307ef3c --- /dev/null +++ b/rust/riverqueue-migrate/README.md @@ -0,0 +1,35 @@ +# riverqueue-migrate + +River's PostgreSQL and SQLite migrations for Rust, identical to the ones River +for Go applies, so either language can migrate a database the other uses. The +`riverqueue` command from the `riverqueue-cli` crate runs the same migrations +from a shell. + +Use `PostgresMigrator` for PostgreSQL and `SqliteMigrator` for SQLite: + +```rust,no_run +use riverqueue_migrate::{Direction, MigrateOpts, PostgresMigrator, SqliteMigrator}; +use sqlx::{PgPool, SqlitePool}; + +async fn migrate(postgres: PgPool, sqlite: SqlitePool) -> Result<(), Box> { + // Apply every outstanding migration. + let migrator = PostgresMigrator::new(postgres); + let applied = migrator.migrate_up().await?; + println!("applied versions {applied:?}"); + + // Or preview what migrating down one version would run. + let preview = migrator + .migrate(Direction::Down, MigrateOpts::new().with_dry_run(true)) + .await?; + for version in preview.versions { + println!("would revert {:03} {}", version.version, version.name); + } + + // Check that every migration is applied before starting clients. + let validation = SqliteMigrator::new(sqlite).validate(None).await?; + if !validation.is_valid() { + return Err(validation.to_string().into()); + } + Ok(()) +} +``` diff --git a/rust/riverqueue-migrate/examples/migrate.rs b/rust/riverqueue-migrate/examples/migrate.rs new file mode 100644 index 000000000..e8f0e304a --- /dev/null +++ b/rust/riverqueue-migrate/examples/migrate.rs @@ -0,0 +1,24 @@ +use std::error::Error; + +use riverqueue_migrate::{Direction, MigrateOpts, PostgresMigrator}; +use sqlx::PgPool; + +#[tokio::main] +async fn main() -> Result<(), Box> { + let pool = PgPool::connect(&std::env::var("DATABASE_URL")?).await?; + let migrator = PostgresMigrator::new(pool); + + let preview = migrator + .migrate(Direction::Up, MigrateOpts::new().with_dry_run(true)) + .await?; + for migration in preview.versions { + println!("would apply {:03} {}", migration.version, migration.name); + } + + migrator.migrate_up().await?; + let validation = migrator.validate(None).await?; + if !validation.is_valid() { + return Err(validation.to_string().into()); + } + Ok(()) +} diff --git a/rust/riverqueue-migrate/migrations/main/001_create_river_migration.down.sql b/rust/riverqueue-migrate/migrations/main/001_create_river_migration.down.sql new file mode 100644 index 000000000..8bfe82027 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/main/001_create_river_migration.down.sql @@ -0,0 +1 @@ +DROP TABLE /* TEMPLATE: schema */river_migration; \ No newline at end of file diff --git a/rust/riverqueue-migrate/migrations/main/001_create_river_migration.up.sql b/rust/riverqueue-migrate/migrations/main/001_create_river_migration.up.sql new file mode 100644 index 000000000..27006d562 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/main/001_create_river_migration.up.sql @@ -0,0 +1,8 @@ +CREATE TABLE /* TEMPLATE: schema */river_migration( + id bigserial PRIMARY KEY, + created_at timestamptz NOT NULL DEFAULT NOW(), + version bigint NOT NULL, + CONSTRAINT version CHECK (version >= 1) +); + +CREATE UNIQUE INDEX ON /* TEMPLATE: schema */river_migration USING btree(version); \ No newline at end of file diff --git a/rust/riverqueue-migrate/migrations/main/002_initial_schema.down.sql b/rust/riverqueue-migrate/migrations/main/002_initial_schema.down.sql new file mode 100644 index 000000000..d334d8a65 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/main/002_initial_schema.down.sql @@ -0,0 +1,5 @@ +DROP TABLE /* TEMPLATE: schema */river_job; +DROP FUNCTION /* TEMPLATE: schema */river_job_notify; +DROP TYPE /* TEMPLATE: schema */river_job_state; + +DROP TABLE /* TEMPLATE: schema */river_leader; \ No newline at end of file diff --git a/rust/riverqueue-migrate/migrations/main/002_initial_schema.up.sql b/rust/riverqueue-migrate/migrations/main/002_initial_schema.up.sql new file mode 100644 index 000000000..7fbca71b4 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/main/002_initial_schema.up.sql @@ -0,0 +1,96 @@ +CREATE TYPE /* TEMPLATE: schema */river_job_state AS ENUM( + 'available', + 'cancelled', + 'completed', + 'discarded', + 'retryable', + 'running', + 'scheduled' +); + +CREATE TABLE /* TEMPLATE: schema */river_job( + -- 8 bytes + id bigserial PRIMARY KEY, + + -- 8 bytes (4 bytes + 2 bytes + 2 bytes) + -- + -- `state` is kept near the top of the table for operator convenience -- when + -- looking at jobs with `SELECT *` it'll appear first after ID. The other two + -- fields aren't as important but are kept adjacent to `state` for alignment + -- to get an 8-byte block. + state /* TEMPLATE: schema */river_job_state NOT NULL DEFAULT 'available', + attempt smallint NOT NULL DEFAULT 0, + max_attempts smallint NOT NULL, + + -- 8 bytes each (no alignment needed) + attempted_at timestamptz, + created_at timestamptz NOT NULL DEFAULT NOW(), + finalized_at timestamptz, + scheduled_at timestamptz NOT NULL DEFAULT NOW(), + + -- 2 bytes (some wasted padding probably) + priority smallint NOT NULL DEFAULT 1, + + -- types stored out-of-band + args jsonb, + attempted_by text[], + errors jsonb[], + kind text NOT NULL, + metadata jsonb NOT NULL DEFAULT '{}', + queue text NOT NULL DEFAULT 'default', + tags varchar(255)[], + + CONSTRAINT finalized_or_finalized_at_null CHECK ((state IN ('cancelled', 'completed', 'discarded') AND finalized_at IS NOT NULL) OR finalized_at IS NULL), + CONSTRAINT max_attempts_is_positive CHECK (max_attempts > 0), + CONSTRAINT priority_in_range CHECK (priority >= 1 AND priority <= 4), + CONSTRAINT queue_length CHECK (char_length(queue) > 0 AND char_length(queue) < 128), + CONSTRAINT kind_length CHECK (char_length(kind) > 0 AND char_length(kind) < 128) +); + +-- We may want to consider adding another property here after `kind` if it seems +-- like it'd be useful for something. +CREATE INDEX river_job_kind ON /* TEMPLATE: schema */river_job USING btree(kind); + +CREATE INDEX river_job_state_and_finalized_at_index ON /* TEMPLATE: schema */river_job USING btree(state, finalized_at) WHERE finalized_at IS NOT NULL; + +CREATE INDEX river_job_prioritized_fetching_index ON /* TEMPLATE: schema */river_job USING btree(state, queue, priority, scheduled_at, id); + +CREATE INDEX river_job_args_index ON /* TEMPLATE: schema */river_job USING GIN(args); + +CREATE INDEX river_job_metadata_index ON /* TEMPLATE: schema */river_job USING GIN(metadata); + +CREATE OR REPLACE FUNCTION /* TEMPLATE: schema */river_job_notify() + RETURNS TRIGGER + AS $$ +DECLARE + payload json; +BEGIN + IF NEW.state = 'available' THEN + -- Notify will coalesce duplicate notifications within a transaction, so + -- keep these payloads generalized: + payload = json_build_object('queue', NEW.queue); + PERFORM + pg_notify('river_insert', payload::text); + END IF; + RETURN NULL; +END; +$$ +LANGUAGE plpgsql; + +CREATE TRIGGER river_notify + AFTER INSERT ON /* TEMPLATE: schema */river_job + FOR EACH ROW + EXECUTE PROCEDURE /* TEMPLATE: schema */river_job_notify(); + +CREATE UNLOGGED TABLE /* TEMPLATE: schema */river_leader( + -- 8 bytes each (no alignment needed) + elected_at timestamptz NOT NULL, + expires_at timestamptz NOT NULL, + + -- types stored out-of-band + leader_id text NOT NULL, + name text PRIMARY KEY, + + CONSTRAINT name_length CHECK (char_length(name) > 0 AND char_length(name) < 128), + CONSTRAINT leader_id_length CHECK (char_length(leader_id) > 0 AND char_length(leader_id) < 128) +); diff --git a/rust/riverqueue-migrate/migrations/main/003_river_job_tags_non_null.down.sql b/rust/riverqueue-migrate/migrations/main/003_river_job_tags_non_null.down.sql new file mode 100644 index 000000000..acef65cb9 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/main/003_river_job_tags_non_null.down.sql @@ -0,0 +1,3 @@ +ALTER TABLE /* TEMPLATE: schema */river_job + ALTER COLUMN tags DROP NOT NULL, + ALTER COLUMN tags DROP DEFAULT; diff --git a/rust/riverqueue-migrate/migrations/main/003_river_job_tags_non_null.up.sql b/rust/riverqueue-migrate/migrations/main/003_river_job_tags_non_null.up.sql new file mode 100644 index 000000000..0a472dde4 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/main/003_river_job_tags_non_null.up.sql @@ -0,0 +1,3 @@ +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN tags SET DEFAULT '{}'; +UPDATE /* TEMPLATE: schema */river_job SET tags = '{}' WHERE tags IS NULL; +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN tags SET NOT NULL; diff --git a/rust/riverqueue-migrate/migrations/main/004_pending_and_more.down.sql b/rust/riverqueue-migrate/migrations/main/004_pending_and_more.down.sql new file mode 100644 index 000000000..1b7ec7e84 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/main/004_pending_and_more.down.sql @@ -0,0 +1,42 @@ +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN args DROP NOT NULL; + +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN metadata DROP NOT NULL; +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN metadata DROP DEFAULT; + +-- It is not possible to safely remove 'pending' from the river_job_state enum, +-- so leave it in place. + +ALTER TABLE /* TEMPLATE: schema */river_job DROP CONSTRAINT finalized_or_finalized_at_null; +ALTER TABLE /* TEMPLATE: schema */river_job ADD CONSTRAINT finalized_or_finalized_at_null CHECK ( + (state IN ('cancelled', 'completed', 'discarded') AND finalized_at IS NOT NULL) OR finalized_at IS NULL +); + +CREATE OR REPLACE FUNCTION /* TEMPLATE: schema */river_job_notify() + RETURNS TRIGGER + AS $$ +DECLARE + payload json; +BEGIN + IF NEW.state = 'available' THEN + -- Notify will coalesce duplicate notifications within a transaction, so + -- keep these payloads generalized: + payload = json_build_object('queue', NEW.queue); + PERFORM + pg_notify('river_insert', payload::text); + END IF; + RETURN NULL; +END; +$$ +LANGUAGE plpgsql; + +CREATE TRIGGER river_notify + AFTER INSERT ON /* TEMPLATE: schema */river_job + FOR EACH ROW + EXECUTE PROCEDURE /* TEMPLATE: schema */river_job_notify(); + +DROP TABLE /* TEMPLATE: schema */river_queue; + +ALTER TABLE /* TEMPLATE: schema */river_leader + ALTER COLUMN name DROP DEFAULT, + DROP CONSTRAINT name_length, + ADD CONSTRAINT name_length CHECK (char_length(name) > 0 AND char_length(name) < 128); \ No newline at end of file diff --git a/rust/riverqueue-migrate/migrations/main/004_pending_and_more.up.sql b/rust/riverqueue-migrate/migrations/main/004_pending_and_more.up.sql new file mode 100644 index 000000000..9f5e47bb1 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/main/004_pending_and_more.up.sql @@ -0,0 +1,45 @@ +-- The args column never had a NOT NULL constraint or default value at the +-- database level, though we tried to ensure one at the application level. +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN args SET DEFAULT '{}'; +UPDATE /* TEMPLATE: schema */river_job SET args = '{}' WHERE args IS NULL; +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN args SET NOT NULL; +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN args DROP DEFAULT; + +-- The metadata column never had a NOT NULL constraint or default value at the +-- database level, though we tried to ensure one at the application level. +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN metadata SET DEFAULT '{}'; +UPDATE /* TEMPLATE: schema */river_job SET metadata = '{}' WHERE metadata IS NULL; +ALTER TABLE /* TEMPLATE: schema */river_job ALTER COLUMN metadata SET NOT NULL; + +-- The 'pending' job state will be used for upcoming functionality: +ALTER TYPE /* TEMPLATE: schema */river_job_state ADD VALUE IF NOT EXISTS 'pending' AFTER 'discarded'; + +ALTER TABLE /* TEMPLATE: schema */river_job DROP CONSTRAINT finalized_or_finalized_at_null; +ALTER TABLE /* TEMPLATE: schema */river_job ADD CONSTRAINT finalized_or_finalized_at_null CHECK ( + (finalized_at IS NULL AND state NOT IN ('cancelled', 'completed', 'discarded')) OR + (finalized_at IS NOT NULL AND state IN ('cancelled', 'completed', 'discarded')) +); + +DROP TRIGGER river_notify ON /* TEMPLATE: schema */river_job; +DROP FUNCTION /* TEMPLATE: schema */river_job_notify; + +-- +-- Create table `river_queue`. +-- + +CREATE TABLE /* TEMPLATE: schema */river_queue ( + name text PRIMARY KEY NOT NULL, + created_at timestamptz NOT NULL DEFAULT now(), + metadata jsonb NOT NULL DEFAULT '{}' ::jsonb, + paused_at timestamptz, + updated_at timestamptz NOT NULL +); + +-- +-- Alter `river_leader` to add a default value of 'default` to `name`. +-- + +ALTER TABLE /* TEMPLATE: schema */river_leader + ALTER COLUMN name SET DEFAULT 'default', + DROP CONSTRAINT name_length, + ADD CONSTRAINT name_length CHECK (name = 'default'); \ No newline at end of file diff --git a/rust/riverqueue-migrate/migrations/main/005_migration_unique_client.down.sql b/rust/riverqueue-migrate/migrations/main/005_migration_unique_client.down.sql new file mode 100644 index 000000000..b8e041d54 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/main/005_migration_unique_client.down.sql @@ -0,0 +1,57 @@ +-- +-- Revert to migration table based only on `(version)`. +-- +-- If any non-main migrations are present, 005 is considered irreversible. +-- + +DO +$body$ +BEGIN + -- Tolerate users who may be using their own migration system rather than + -- River's. If they are, they will have skipped version 001 containing + -- `CREATE TABLE river_migration`, so this table won't exist. + IF (SELECT to_regclass('/* TEMPLATE: schema */river_migration') IS NOT NULL) THEN + IF EXISTS ( + SELECT * + FROM /* TEMPLATE: schema */river_migration + WHERE line <> 'main' + ) THEN + RAISE EXCEPTION 'Found non-main migration lines in the database; version 005 migration is irreversible because it would result in loss of migration information.'; + END IF; + + ALTER TABLE /* TEMPLATE: schema */river_migration + RENAME TO river_migration_old; + + CREATE TABLE /* TEMPLATE: schema */river_migration( + id bigserial PRIMARY KEY, + created_at timestamptz NOT NULL DEFAULT NOW(), + version bigint NOT NULL, + CONSTRAINT version CHECK (version >= 1) + ); + + CREATE UNIQUE INDEX ON /* TEMPLATE: schema */river_migration USING btree(version); + + INSERT INTO /* TEMPLATE: schema */river_migration + (created_at, version) + SELECT created_at, version + FROM /* TEMPLATE: schema */river_migration_old; + + DROP TABLE /* TEMPLATE: schema */river_migration_old; + END IF; +END; +$body$ +LANGUAGE 'plpgsql'; + +-- +-- Drop `river_job.unique_key`. +-- + +ALTER TABLE /* TEMPLATE: schema */river_job + DROP COLUMN unique_key; + +-- +-- Drop `river_client` and derivative. +-- + +DROP TABLE /* TEMPLATE: schema */river_client_queue; +DROP TABLE /* TEMPLATE: schema */river_client; diff --git a/rust/riverqueue-migrate/migrations/main/005_migration_unique_client.up.sql b/rust/riverqueue-migrate/migrations/main/005_migration_unique_client.up.sql new file mode 100644 index 000000000..e0f1711ec --- /dev/null +++ b/rust/riverqueue-migrate/migrations/main/005_migration_unique_client.up.sql @@ -0,0 +1,79 @@ +-- +-- Rebuild the migration table so it's based on `(line, version)`. +-- + +DO +$body$ +BEGIN + -- Tolerate users who may be using their own migration system rather than + -- River's. If they are, they will have skipped version 001 containing + -- `CREATE TABLE river_migration`, so this table won't exist. + IF (SELECT to_regclass('/* TEMPLATE: schema */river_migration') IS NOT NULL) THEN + ALTER TABLE /* TEMPLATE: schema */river_migration + RENAME TO river_migration_old; + + CREATE TABLE /* TEMPLATE: schema */river_migration( + line TEXT NOT NULL, + version bigint NOT NULL, + created_at timestamptz NOT NULL DEFAULT NOW(), + CONSTRAINT line_length CHECK (char_length(line) > 0 AND char_length(line) < 128), + CONSTRAINT version_gte_1 CHECK (version >= 1), + PRIMARY KEY (line, version) + ); + + INSERT INTO /* TEMPLATE: schema */river_migration + (created_at, line, version) + SELECT created_at, 'main', version + FROM /* TEMPLATE: schema */river_migration_old; + + DROP TABLE /* TEMPLATE: schema */river_migration_old; + END IF; +END; +$body$ +LANGUAGE 'plpgsql'; + +-- +-- Add `river_job.unique_key` and bring up an index on it. +-- + +-- These statements use `IF NOT EXISTS` to allow users with a `river_job` table +-- of non-trivial size to build the index `CONCURRENTLY` out of band of this +-- migration, then follow by completing the migration. +ALTER TABLE /* TEMPLATE: schema */river_job + ADD COLUMN IF NOT EXISTS unique_key bytea; + +CREATE UNIQUE INDEX IF NOT EXISTS river_job_kind_unique_key_idx ON /* TEMPLATE: schema */river_job (kind, unique_key) WHERE unique_key IS NOT NULL; + +-- +-- Create `river_client` and derivative. +-- +-- This feature hasn't quite yet been implemented, but we're taking advantage of +-- the migration to add the schema early so that we can add it later without an +-- additional migration. +-- + +CREATE UNLOGGED TABLE /* TEMPLATE: schema */river_client ( + id text PRIMARY KEY NOT NULL, + created_at timestamptz NOT NULL DEFAULT now(), + metadata jsonb NOT NULL DEFAULT '{}', + paused_at timestamptz, + updated_at timestamptz NOT NULL, + CONSTRAINT name_length CHECK (char_length(id) > 0 AND char_length(id) < 128) +); + +-- Differs from `river_queue` in that it tracks the queue state for a particular +-- active client. +CREATE UNLOGGED TABLE /* TEMPLATE: schema */river_client_queue ( + river_client_id text NOT NULL REFERENCES /* TEMPLATE: schema */river_client (id) ON DELETE CASCADE, + name text NOT NULL, + created_at timestamptz NOT NULL DEFAULT now(), + max_workers bigint NOT NULL DEFAULT 0, + metadata jsonb NOT NULL DEFAULT '{}', + num_jobs_completed bigint NOT NULL DEFAULT 0, + num_jobs_running bigint NOT NULL DEFAULT 0, + updated_at timestamptz NOT NULL, + PRIMARY KEY (river_client_id, name), + CONSTRAINT name_length CHECK (char_length(name) > 0 AND char_length(name) < 128), + CONSTRAINT num_jobs_completed_zero_or_positive CHECK (num_jobs_completed >= 0), + CONSTRAINT num_jobs_running_zero_or_positive CHECK (num_jobs_running >= 0) +); \ No newline at end of file diff --git a/rust/riverqueue-migrate/migrations/main/006_bulk_unique.down.sql b/rust/riverqueue-migrate/migrations/main/006_bulk_unique.down.sql new file mode 100644 index 000000000..26cd84345 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/main/006_bulk_unique.down.sql @@ -0,0 +1,16 @@ + +-- +-- Drop `river_job.unique_states` and its index. +-- + +DROP INDEX /* TEMPLATE: schema */river_job_unique_idx; + +ALTER TABLE /* TEMPLATE: schema */river_job + DROP COLUMN unique_states; + +CREATE UNIQUE INDEX IF NOT EXISTS river_job_kind_unique_key_idx ON /* TEMPLATE: schema */river_job (kind, unique_key) WHERE unique_key IS NOT NULL; + +-- +-- Drop `river_job_state_in_bitmask` function. +-- +DROP FUNCTION /* TEMPLATE: schema */river_job_state_in_bitmask; diff --git a/rust/riverqueue-migrate/migrations/main/006_bulk_unique.up.sql b/rust/riverqueue-migrate/migrations/main/006_bulk_unique.up.sql new file mode 100644 index 000000000..ef96a19f9 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/main/006_bulk_unique.up.sql @@ -0,0 +1,40 @@ +CREATE OR REPLACE FUNCTION /* TEMPLATE: schema */river_job_state_in_bitmask(bitmask BIT(8), state /* TEMPLATE: schema */river_job_state) +RETURNS boolean +LANGUAGE SQL +IMMUTABLE +AS $$ + SELECT CASE state + WHEN 'available' THEN get_bit(bitmask, 7) + WHEN 'cancelled' THEN get_bit(bitmask, 6) + WHEN 'completed' THEN get_bit(bitmask, 5) + WHEN 'discarded' THEN get_bit(bitmask, 4) + WHEN 'pending' THEN get_bit(bitmask, 3) + WHEN 'retryable' THEN get_bit(bitmask, 2) + WHEN 'running' THEN get_bit(bitmask, 1) + WHEN 'scheduled' THEN get_bit(bitmask, 0) + ELSE 0 + END = 1; +$$; + +-- +-- Add `river_job.unique_states` and bring up an index on it. +-- +-- This column may exist already if users manually created the column and index +-- as instructed in the changelog so the index could be created `CONCURRENTLY`. +-- +ALTER TABLE /* TEMPLATE: schema */river_job ADD COLUMN IF NOT EXISTS unique_states BIT(8); + +-- This statement uses `IF NOT EXISTS` to allow users with a `river_job` table +-- of non-trivial size to build the index `CONCURRENTLY` out of band of this +-- migration, then follow by completing the migration. +CREATE UNIQUE INDEX IF NOT EXISTS river_job_unique_idx ON /* TEMPLATE: schema */river_job (unique_key) + WHERE unique_key IS NOT NULL + AND unique_states IS NOT NULL + AND /* TEMPLATE: schema */river_job_state_in_bitmask(unique_states, state); + +-- Remove the old unique index. Users who are actively using the unique jobs +-- feature and who wish to avoid deploy downtime may want od drop this in a +-- subsequent migration once all jobs using the old unique system have been +-- completed (i.e. no more rows with non-null unique_key and null +-- unique_states). +DROP INDEX /* TEMPLATE: schema */river_job_kind_unique_key_idx; diff --git a/rust/riverqueue-migrate/migrations/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.down.sql b/rust/riverqueue-migrate/migrations/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.down.sql new file mode 100644 index 000000000..bed717f87 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.down.sql @@ -0,0 +1,56 @@ +-- +-- SQL cleanup rollback. +-- + +-- +-- Add back unused tables `river_client` and `river_client_queue`. +-- + +CREATE UNLOGGED TABLE /* TEMPLATE: schema */river_client ( + id text PRIMARY KEY NOT NULL, + created_at timestamptz NOT NULL DEFAULT now(), + metadata jsonb NOT NULL DEFAULT '{}', + paused_at timestamptz, + updated_at timestamptz NOT NULL, + CONSTRAINT name_length CHECK (char_length(id) > 0 AND char_length(id) < 128) +); + +CREATE UNLOGGED TABLE /* TEMPLATE: schema */river_client_queue ( + river_client_id text NOT NULL REFERENCES /* TEMPLATE: schema */river_client (id) ON DELETE CASCADE, + name text NOT NULL, + created_at timestamptz NOT NULL DEFAULT now(), + max_workers bigint NOT NULL DEFAULT 0, + metadata jsonb NOT NULL DEFAULT '{}', + num_jobs_completed bigint NOT NULL DEFAULT 0, + num_jobs_running bigint NOT NULL DEFAULT 0, + updated_at timestamptz NOT NULL, + PRIMARY KEY (river_client_id, name), + CONSTRAINT name_length CHECK (char_length(name) > 0 AND char_length(name) < 128), + CONSTRAINT num_jobs_completed_zero_or_positive CHECK (num_jobs_completed >= 0), + CONSTRAINT num_jobs_running_zero_or_positive CHECK (num_jobs_running >= 0) +); + +-- +-- Revert addition of `DEFAULT 25` to `river_job.max_attempts`. +-- + +ALTER TABLE /* TEMPLATE: schema */river_job + ALTER COLUMN max_attempts DROP DEFAULT; + +-- +-- Changes `river_queue.updated_at` to revert the default of `CURRENT_TIMESTAMP`. +-- + +ALTER TABLE /* TEMPLATE: schema */river_queue + ALTER COLUMN updated_at DROP DEFAULT; + +-- +-- SQLite JSONB conversion rollback. +-- +-- No-op. PostgreSQL already stores River JSON columns as jsonb. + +-- +-- Notification outbox rollback. +-- + +DROP TABLE /* TEMPLATE: schema */river_notification; diff --git a/rust/riverqueue-migrate/migrations/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.up.sql b/rust/riverqueue-migrate/migrations/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.up.sql new file mode 100644 index 000000000..39e3249c9 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.up.sql @@ -0,0 +1,44 @@ +-- +-- Notification outbox. +-- + +CREATE TABLE /* TEMPLATE: schema */river_notification ( + id bigserial PRIMARY KEY, + created_at timestamptz NOT NULL DEFAULT now(), + payload text NOT NULL, + topic text NOT NULL, + CONSTRAINT topic_length CHECK (length(topic) > 0 AND length(topic) < 128) +); + +CREATE INDEX river_notification_created_at_idx ON /* TEMPLATE: schema */river_notification (created_at); +CREATE INDEX river_notification_topic_id_idx ON /* TEMPLATE: schema */river_notification (topic, id); + +-- +-- SQLite JSONB conversion. +-- +-- No-op. PostgreSQL already stores River JSON columns as jsonb. + +-- +-- SQL cleanup. +-- + +-- +-- Drop unused tables `river_client` and `river_client_queue`. +-- + +DROP TABLE /* TEMPLATE: schema */river_client_queue; +DROP TABLE /* TEMPLATE: schema */river_client; + +-- +-- Adds `DEFAULT 25` to `river_job.max_attempts`. +-- + +ALTER TABLE /* TEMPLATE: schema */river_job + ALTER COLUMN max_attempts SET DEFAULT 25; + +-- +-- Changes `river_queue.updated_at` to have a default of `CURRENT_TIMESTAMP`. +-- + +ALTER TABLE /* TEMPLATE: schema */river_queue + ALTER COLUMN updated_at SET DEFAULT CURRENT_TIMESTAMP; diff --git a/rust/riverqueue-migrate/migrations/main/008_job_id_autoincrement.down.sql b/rust/riverqueue-migrate/migrations/main/008_job_id_autoincrement.down.sql new file mode 100644 index 000000000..695357bb8 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/main/008_job_id_autoincrement.down.sql @@ -0,0 +1,3 @@ +-- No-op. PostgreSQL sequences already prevent automatically generated job IDs +-- from being reused. +SELECT 1; diff --git a/rust/riverqueue-migrate/migrations/main/008_job_id_autoincrement.up.sql b/rust/riverqueue-migrate/migrations/main/008_job_id_autoincrement.up.sql new file mode 100644 index 000000000..695357bb8 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/main/008_job_id_autoincrement.up.sql @@ -0,0 +1,3 @@ +-- No-op. PostgreSQL sequences already prevent automatically generated job IDs +-- from being reused. +SELECT 1; diff --git a/rust/riverqueue-migrate/migrations/sqlite/main/001_create_river_migration.down.sql b/rust/riverqueue-migrate/migrations/sqlite/main/001_create_river_migration.down.sql new file mode 100644 index 000000000..8bfe82027 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/sqlite/main/001_create_river_migration.down.sql @@ -0,0 +1 @@ +DROP TABLE /* TEMPLATE: schema */river_migration; \ No newline at end of file diff --git a/rust/riverqueue-migrate/migrations/sqlite/main/001_create_river_migration.up.sql b/rust/riverqueue-migrate/migrations/sqlite/main/001_create_river_migration.up.sql new file mode 100644 index 000000000..bdaf09339 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/sqlite/main/001_create_river_migration.up.sql @@ -0,0 +1,8 @@ +CREATE TABLE /* TEMPLATE: schema */river_migration ( + id integer PRIMARY KEY, + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + version integer NOT NULL, + CONSTRAINT version CHECK (version >= 1) +); + +CREATE UNIQUE INDEX /* TEMPLATE: schema */river_migration_version_idx ON river_migration (version); \ No newline at end of file diff --git a/rust/riverqueue-migrate/migrations/sqlite/main/002_initial_schema.down.sql b/rust/riverqueue-migrate/migrations/sqlite/main/002_initial_schema.down.sql new file mode 100644 index 000000000..cbdd56dea --- /dev/null +++ b/rust/riverqueue-migrate/migrations/sqlite/main/002_initial_schema.down.sql @@ -0,0 +1,8 @@ +-- +-- Normally `river_job` and `river_job_notify()` are dropped here, but since +-- SQLite was added well after 002 came about, we push that to version 006 index. +-- + +DROP TABLE /* TEMPLATE: schema */river_job; + +DROP TABLE /* TEMPLATE: schema */river_leader; \ No newline at end of file diff --git a/rust/riverqueue-migrate/migrations/sqlite/main/002_initial_schema.up.sql b/rust/riverqueue-migrate/migrations/sqlite/main/002_initial_schema.up.sql new file mode 100644 index 000000000..043facf29 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/sqlite/main/002_initial_schema.up.sql @@ -0,0 +1,19 @@ +-- +-- Normally `river_job` and `river_job_notify()` are raised here, but since +-- SQLite was added well after 002 came about, we push that to version 006 index. +-- + +-- Dummy `river_job` table so that there's something to truncate in tests when +-- migrated to this version specifically. +CREATE TABLE /* TEMPLATE: schema */river_job ( + id integer PRIMARY KEY +); + +CREATE TABLE /* TEMPLATE: schema */river_leader ( + elected_at timestamp NOT NULL, + expires_at timestamp NOT NULL, + leader_id text NOT NULL, + name text PRIMARY KEY NOT NULL, + CONSTRAINT name_length CHECK (length(name) > 0 AND length(name) < 128), + CONSTRAINT leader_id_length CHECK (length(leader_id) > 0 AND length(leader_id) < 128) +); diff --git a/rust/riverqueue-migrate/migrations/sqlite/main/003_river_job_tags_non_null.down.sql b/rust/riverqueue-migrate/migrations/sqlite/main/003_river_job_tags_non_null.down.sql new file mode 100644 index 000000000..8d314cf06 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/sqlite/main/003_river_job_tags_non_null.down.sql @@ -0,0 +1,6 @@ +-- +-- Normally `river_job.tags` is set back to nullable here, but since SQLite was +-- added well after 003 came about, we push that to version 006 index. +-- + +SELECT 1; diff --git a/rust/riverqueue-migrate/migrations/sqlite/main/003_river_job_tags_non_null.up.sql b/rust/riverqueue-migrate/migrations/sqlite/main/003_river_job_tags_non_null.up.sql new file mode 100644 index 000000000..d4e1e2404 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/sqlite/main/003_river_job_tags_non_null.up.sql @@ -0,0 +1,6 @@ +-- +-- Normally `river_job.tags` is set to `NOT NULL` with a `DEFAULT` here, but since +-- SQLite was added well after 003 came about, we push that to version 006 index. +-- + +SELECT 1; diff --git a/rust/riverqueue-migrate/migrations/sqlite/main/004_pending_and_more.down.sql b/rust/riverqueue-migrate/migrations/sqlite/main/004_pending_and_more.down.sql new file mode 100644 index 000000000..c64554441 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/sqlite/main/004_pending_and_more.down.sql @@ -0,0 +1,26 @@ +-- +-- Normally, args and metadata both become `NOT NULL`, `pending` is added, and +-- the constraint `finalized_at` is changed, but because SQLite was added later, +-- we've just pushed all of this to an initial `river_job` creation in 006. +-- + +-- +-- Drop `river_queue`. +-- + +DROP TABLE /* TEMPLATE: schema */river_queue; + +-- +-- Reverse changes to `river_leader`. +-- + +DROP TABLE /* TEMPLATE: schema */river_leader; + +CREATE TABLE /* TEMPLATE: schema */river_leader ( + elected_at timestamp NOT NULL, + expires_at timestamp NOT NULL, + leader_id text NOT NULL, + name text PRIMARY KEY NOT NULL, + CONSTRAINT name_length CHECK (length(name) > 0 AND length(name) < 128), + CONSTRAINT leader_id_length CHECK (length(leader_id) > 0 AND length(leader_id) < 128) +); \ No newline at end of file diff --git a/rust/riverqueue-migrate/migrations/sqlite/main/004_pending_and_more.up.sql b/rust/riverqueue-migrate/migrations/sqlite/main/004_pending_and_more.up.sql new file mode 100644 index 000000000..254e1f7a3 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/sqlite/main/004_pending_and_more.up.sql @@ -0,0 +1,33 @@ +-- +-- Normally, args and metadata both become `NOT NULL`, `pending` is added, and +-- the constraint `finalized_at` is changed, but because SQLite was added later, +-- we've just pushed all of this to an initial `river_job` creation in 006. +-- + +-- +-- Create table `river_queue`. +-- + +CREATE TABLE /* TEMPLATE: schema */river_queue ( + name text PRIMARY KEY NOT NULL, + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + metadata blob NOT NULL DEFAULT (json('{}')), + paused_at timestamp, + updated_at timestamp NOT NULL +); + +-- +-- Alter `river_leader` to add a default value of 'default` to `name`. SQLite +-- doesn't allow schema modifications, so this redefines the table entirely. +-- + +DROP TABLE /* TEMPLATE: schema */river_leader; + +CREATE TABLE /* TEMPLATE: schema */river_leader ( + elected_at timestamp NOT NULL, + expires_at timestamp NOT NULL, + leader_id text NOT NULL, + name text PRIMARY KEY NOT NULL DEFAULT 'default' CHECK (name = 'default'), + CONSTRAINT name_length CHECK (length(name) > 0 AND length(name) < 128), + CONSTRAINT leader_id_length CHECK (length(leader_id) > 0 AND length(leader_id) < 128) +); \ No newline at end of file diff --git a/rust/riverqueue-migrate/migrations/sqlite/main/005_migration_unique_client.down.sql b/rust/riverqueue-migrate/migrations/sqlite/main/005_migration_unique_client.down.sql new file mode 100644 index 000000000..d94787d3f --- /dev/null +++ b/rust/riverqueue-migrate/migrations/sqlite/main/005_migration_unique_client.down.sql @@ -0,0 +1,37 @@ +-- +-- Revert to migration table based only on `(version)`. +-- +-- If any non-main migrations are present, 005 is considered irreversible. +-- + +ALTER TABLE /* TEMPLATE: schema */river_migration + RENAME TO river_migration_old; + +CREATE TABLE /* TEMPLATE: schema */river_migration ( + id integer PRIMARY KEY, + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + version integer NOT NULL, + CONSTRAINT version CHECK (version >= 1) +); + +CREATE UNIQUE INDEX /* TEMPLATE: schema */river_migration_version_idx ON river_migration (version); + +INSERT INTO /* TEMPLATE: schema */river_migration + (created_at, version) +SELECT created_at, version +FROM /* TEMPLATE: schema */river_migration_old; + +DROP TABLE /* TEMPLATE: schema */river_migration_old; + +-- +-- Normally, `unique_key` and an index are added here, but because SQLite was +-- added later, we've just pushed all of this to an initial `river_job` creation +-- in 006. +-- + +-- +-- Drop `river_client` and derivative. +-- + +DROP TABLE /* TEMPLATE: schema */river_client_queue; +DROP TABLE /* TEMPLATE: schema */river_client; diff --git a/rust/riverqueue-migrate/migrations/sqlite/main/005_migration_unique_client.up.sql b/rust/riverqueue-migrate/migrations/sqlite/main/005_migration_unique_client.up.sql new file mode 100644 index 000000000..dc3273349 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/sqlite/main/005_migration_unique_client.up.sql @@ -0,0 +1,64 @@ +-- +-- Rebuild the migration table so it's based on `(line, version)`. +-- + +DROP INDEX /* TEMPLATE: schema */river_migration_version_idx; + +ALTER TABLE /* TEMPLATE: schema */river_migration + RENAME TO river_migration_old; + +CREATE TABLE /* TEMPLATE: schema */river_migration ( + line text NOT NULL, + version integer NOT NULL, + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + CONSTRAINT line_length CHECK (length(line) > 0 AND length(line) < 128), + CONSTRAINT version_gte_1 CHECK (version >= 1), + PRIMARY KEY (line, version) +); + +INSERT INTO /* TEMPLATE: schema */river_migration + (created_at, line, version) +SELECT created_at, 'main', version +FROM /* TEMPLATE: schema */river_migration_old; + +DROP TABLE /* TEMPLATE: schema */river_migration_old; + +-- +-- Normally, `unique_key` and an index are added here, but because SQLite was +-- added later, we've just pushed all of this to an initial `river_job` creation +-- in 006. +-- + +-- +-- Create `river_client` and derivative. +-- +-- This feature hasn't quite yet been implemented, but we're taking advantage of +-- the migration to add the schema early so that we can add it later without an +-- additional migration. +-- + +CREATE TABLE /* TEMPLATE: schema */river_client ( + id text PRIMARY KEY NOT NULL, + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + metadata blob NOT NULL DEFAULT (json('{}')), + paused_at timestamp, + updated_at timestamp NOT NULL, + CONSTRAINT name_length CHECK (length(id) > 0 AND length(id) < 128) +); + +-- Differs from `river_queue` in that it tracks the queue state for a particular +-- active client. +CREATE TABLE /* TEMPLATE: schema */river_client_queue ( + river_client_id text NOT NULL REFERENCES river_client (id) ON DELETE CASCADE, + name text NOT NULL, + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + max_workers integer NOT NULL DEFAULT 0, + metadata blob NOT NULL DEFAULT (json('{}')), + num_jobs_completed integer NOT NULL DEFAULT 0, + num_jobs_running integer NOT NULL DEFAULT 0, + updated_at timestamp NOT NULL, + PRIMARY KEY (river_client_id, name), + CONSTRAINT name_length CHECK (length(name) > 0 AND length(name) < 128), + CONSTRAINT num_jobs_completed_zero_or_positive CHECK (num_jobs_completed >= 0), + CONSTRAINT num_jobs_running_zero_or_positive CHECK (num_jobs_running >= 0) +); \ No newline at end of file diff --git a/rust/riverqueue-migrate/migrations/sqlite/main/006_bulk_unique.down.sql b/rust/riverqueue-migrate/migrations/sqlite/main/006_bulk_unique.down.sql new file mode 100644 index 000000000..a8d273f84 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/sqlite/main/006_bulk_unique.down.sql @@ -0,0 +1,7 @@ +DROP TABLE /* TEMPLATE: schema */river_job; + +-- Dummy `river_job` table so that there's something to truncate in tests when +-- migrated to this version specifically. +CREATE TABLE /* TEMPLATE: schema */river_job ( + id integer PRIMARY KEY +); diff --git a/rust/riverqueue-migrate/migrations/sqlite/main/006_bulk_unique.up.sql b/rust/riverqueue-migrate/migrations/sqlite/main/006_bulk_unique.up.sql new file mode 100644 index 000000000..528a4680e --- /dev/null +++ b/rust/riverqueue-migrate/migrations/sqlite/main/006_bulk_unique.up.sql @@ -0,0 +1,63 @@ +-- Only drops the trivial `river_job` we created in 002 which puts a placeholder +-- in place so that the right tables exist in the right versions. We don't +-- bother migrating any job data because it's not possible to have had any real +-- jobs by that point because this version (006) preexists the addition of SQLite. +DROP TABLE /* TEMPLATE: schema */river_job; + +CREATE TABLE /* TEMPLATE: schema */river_job ( + id integer PRIMARY KEY, -- SQLite aliases this to ROWID, which may reuse deleted IDs. + args blob NOT NULL DEFAULT '{}', + attempt integer NOT NULL DEFAULT 0, + attempted_at timestamp, + attempted_by blob, -- json + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + errors blob, -- json + finalized_at timestamp, + kind text NOT NULL, + max_attempts integer NOT NULL, + metadata blob NOT NULL DEFAULT (json('{}')), + priority integer NOT NULL DEFAULT 1, + queue text NOT NULL DEFAULT 'default', + state text NOT NULL DEFAULT 'available', + scheduled_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + tags blob NOT NULL DEFAULT (json('[]')), + unique_key blob, + unique_states integer, + CONSTRAINT finalized_or_finalized_at_null CHECK ( + (finalized_at IS NULL AND state NOT IN ('cancelled', 'completed', 'discarded')) OR + (finalized_at IS NOT NULL AND state IN ('cancelled', 'completed', 'discarded')) + ), + CONSTRAINT priority_in_range CHECK (priority >= 1 AND priority <= 4), + CONSTRAINT queue_length CHECK (length(queue) > 0 AND length(queue) < 128), + CONSTRAINT kind_length CHECK (length(kind) > 0 AND length(kind) < 128), + CONSTRAINT state_valid CHECK (state IN ('available', 'cancelled', 'completed', 'discarded', 'pending', 'retryable', 'running', 'scheduled')) +); + +-- All these indexes are normally brought up in version 002. +CREATE INDEX /* TEMPLATE: schema */river_job_kind ON river_job (kind); +CREATE INDEX /* TEMPLATE: schema */river_job_state_and_finalized_at_index ON river_job (state, finalized_at) WHERE finalized_at IS NOT NULL; +CREATE INDEX /* TEMPLATE: schema */river_job_prioritized_fetching_index ON river_job (state, queue, priority, scheduled_at, id); + +-- Not raised because SQLite doesn't support Gin indexes. These aren't used in +-- River anyway. +-- CREATE INDEX river_job_args_index ON /* TEMPLATE: schema */river_job USING GIN(args); +-- CREATE INDEX river_job_metadata_index ON /* TEMPLATE: schema */river_job USING GIN(metadata); + +-- SQLite doesn't support SQL functions, so where the bit extraction logic below +-- goes in the `river_job_state_in_bitmask` function in Postgres, here it's +-- baked right into the index. Use of helpers that don't exist in SQLite like +-- `get_bit` are also dropped by necessity. +CREATE UNIQUE INDEX /* TEMPLATE: schema */river_job_unique_idx ON river_job (unique_key) + WHERE unique_key IS NOT NULL + AND unique_states IS NOT NULL + AND CASE state + WHEN 'available' THEN unique_states & (1 << 0) + WHEN 'cancelled' THEN unique_states & (1 << 1) + WHEN 'completed' THEN unique_states & (1 << 2) + WHEN 'discarded' THEN unique_states & (1 << 3) + WHEN 'pending' THEN unique_states & (1 << 4) + WHEN 'retryable' THEN unique_states & (1 << 5) + WHEN 'running' THEN unique_states & (1 << 6) + WHEN 'scheduled' THEN unique_states & (1 << 7) + ELSE 0 + END >= 1; diff --git a/rust/riverqueue-migrate/migrations/sqlite/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.down.sql b/rust/riverqueue-migrate/migrations/sqlite/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.down.sql new file mode 100644 index 000000000..1e3bcffb0 --- /dev/null +++ b/rust/riverqueue-migrate/migrations/sqlite/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.down.sql @@ -0,0 +1,255 @@ +-- +-- SQL cleanup rollback. +-- + +-- +-- Add back unused tables `river_client` and `river_client_queue`. +-- + +CREATE TABLE /* TEMPLATE: schema */river_client ( + id text PRIMARY KEY NOT NULL, + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + metadata blob NOT NULL DEFAULT (jsonb('{}')), + paused_at timestamp, + updated_at timestamp NOT NULL, + CONSTRAINT name_length CHECK (length(id) > 0 AND length(id) < 128) +); + +CREATE TABLE /* TEMPLATE: schema */river_client_queue ( + river_client_id text NOT NULL REFERENCES river_client (id) ON DELETE CASCADE, + name text NOT NULL, + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + max_workers integer NOT NULL DEFAULT 0, + metadata blob NOT NULL DEFAULT (jsonb('{}')), + num_jobs_completed integer NOT NULL DEFAULT 0, + num_jobs_running integer NOT NULL DEFAULT 0, + updated_at timestamp NOT NULL, + PRIMARY KEY (river_client_id, name), + CONSTRAINT name_length CHECK (length(name) > 0 AND length(name) < 128), + CONSTRAINT num_jobs_completed_zero_or_positive CHECK (num_jobs_completed >= 0), + CONSTRAINT num_jobs_running_zero_or_positive CHECK (num_jobs_running >= 0) +); + +-- +-- SQLite JSONB conversion rollback. +-- +-- Convert JSONB binary columns back to JSON text format and restore json() +-- defaults. The `river_job` rebuild also reverts the addition of `DEFAULT 25` +-- to `river_job.max_attempts`. +-- +-- SQLite doesn't allow `ALTER TABLE ADD COLUMN` with non-constant defaults like +-- `json('{}')`, so rebuild each affected table instead. +-- + +-- +-- river_job +-- + +DROP INDEX /* TEMPLATE: schema */river_job_kind; +DROP INDEX /* TEMPLATE: schema */river_job_state_and_finalized_at_index; +DROP INDEX /* TEMPLATE: schema */river_job_prioritized_fetching_index; +DROP INDEX /* TEMPLATE: schema */river_job_unique_idx; + +ALTER TABLE /* TEMPLATE: schema */river_job RENAME TO river_job_old; + +CREATE TABLE /* TEMPLATE: schema */river_job ( + id integer PRIMARY KEY, -- SQLite aliases this to ROWID, which may reuse deleted IDs. + args blob NOT NULL DEFAULT '{}', + attempt integer NOT NULL DEFAULT 0, + attempted_at timestamp, + attempted_by blob, -- json + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + errors blob, -- json + finalized_at timestamp, + kind text NOT NULL, + max_attempts integer NOT NULL, + metadata blob NOT NULL DEFAULT (json('{}')), + priority integer NOT NULL DEFAULT 1, + queue text NOT NULL DEFAULT 'default', + state text NOT NULL DEFAULT 'available', + scheduled_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + tags blob NOT NULL DEFAULT (json('[]')), + unique_key blob, + unique_states integer, + CONSTRAINT finalized_or_finalized_at_null CHECK ( + (finalized_at IS NULL AND state NOT IN ('cancelled', 'completed', 'discarded')) OR + (finalized_at IS NOT NULL AND state IN ('cancelled', 'completed', 'discarded')) + ), + CONSTRAINT priority_in_range CHECK (priority >= 1 AND priority <= 4), + CONSTRAINT queue_length CHECK (length(queue) > 0 AND length(queue) < 128), + CONSTRAINT kind_length CHECK (length(kind) > 0 AND length(kind) < 128), + CONSTRAINT state_valid CHECK (state IN ('available', 'cancelled', 'completed', 'discarded', 'pending', 'retryable', 'running', 'scheduled')) +); + +INSERT INTO /* TEMPLATE: schema */river_job ( + id, + args, + attempt, + attempted_at, + attempted_by, + created_at, + errors, + finalized_at, + kind, + max_attempts, + metadata, + priority, + queue, + state, + scheduled_at, + tags, + unique_key, + unique_states +) +SELECT + id, + json(args), + attempt, + attempted_at, + CASE WHEN attempted_by IS NULL THEN NULL ELSE json(attempted_by) END, + created_at, + CASE WHEN errors IS NULL THEN NULL ELSE json(errors) END, + finalized_at, + kind, + max_attempts, + json(metadata), + priority, + queue, + state, + scheduled_at, + json(tags), + unique_key, + unique_states +FROM /* TEMPLATE: schema */river_job_old; + +DROP TABLE /* TEMPLATE: schema */river_job_old; + +CREATE INDEX /* TEMPLATE: schema */river_job_kind ON river_job (kind); +CREATE INDEX /* TEMPLATE: schema */river_job_state_and_finalized_at_index ON river_job (state, finalized_at) WHERE finalized_at IS NOT NULL; +CREATE INDEX /* TEMPLATE: schema */river_job_prioritized_fetching_index ON river_job (state, queue, priority, scheduled_at, id); +CREATE UNIQUE INDEX /* TEMPLATE: schema */river_job_unique_idx ON river_job (unique_key) + WHERE unique_key IS NOT NULL + AND unique_states IS NOT NULL + AND CASE state + WHEN 'available' THEN unique_states & (1 << 0) + WHEN 'cancelled' THEN unique_states & (1 << 1) + WHEN 'completed' THEN unique_states & (1 << 2) + WHEN 'discarded' THEN unique_states & (1 << 3) + WHEN 'pending' THEN unique_states & (1 << 4) + WHEN 'retryable' THEN unique_states & (1 << 5) + WHEN 'running' THEN unique_states & (1 << 6) + WHEN 'scheduled' THEN unique_states & (1 << 7) + ELSE 0 + END >= 1; + +-- +-- river_queue +-- + +ALTER TABLE /* TEMPLATE: schema */river_queue RENAME TO river_queue_old; + +CREATE TABLE /* TEMPLATE: schema */river_queue ( + name text PRIMARY KEY NOT NULL, + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + metadata blob NOT NULL DEFAULT (json('{}')), + paused_at timestamp, + updated_at timestamp NOT NULL +); + +INSERT INTO /* TEMPLATE: schema */river_queue ( + name, + created_at, + metadata, + paused_at, + updated_at +) +SELECT + name, + created_at, + json(metadata), + paused_at, + updated_at +FROM /* TEMPLATE: schema */river_queue_old; + +DROP TABLE /* TEMPLATE: schema */river_queue_old; + +-- +-- river_client +-- + +ALTER TABLE /* TEMPLATE: schema */river_client RENAME TO river_client_old; + +CREATE TABLE /* TEMPLATE: schema */river_client ( + id text PRIMARY KEY NOT NULL, + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + metadata blob NOT NULL DEFAULT (json('{}')), + paused_at timestamp, + updated_at timestamp NOT NULL, + CONSTRAINT name_length CHECK (length(id) > 0 AND length(id) < 128) +); + +INSERT INTO /* TEMPLATE: schema */river_client ( + id, + created_at, + metadata, + paused_at, + updated_at +) +SELECT + id, + created_at, + json(metadata), + paused_at, + updated_at +FROM /* TEMPLATE: schema */river_client_old; + +-- +-- river_client_queue +-- + +ALTER TABLE /* TEMPLATE: schema */river_client_queue RENAME TO river_client_queue_old; + +CREATE TABLE /* TEMPLATE: schema */river_client_queue ( + river_client_id text NOT NULL REFERENCES river_client (id) ON DELETE CASCADE, + name text NOT NULL, + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + max_workers integer NOT NULL DEFAULT 0, + metadata blob NOT NULL DEFAULT (json('{}')), + num_jobs_completed integer NOT NULL DEFAULT 0, + num_jobs_running integer NOT NULL DEFAULT 0, + updated_at timestamp NOT NULL, + PRIMARY KEY (river_client_id, name), + CONSTRAINT name_length CHECK (length(name) > 0 AND length(name) < 128), + CONSTRAINT num_jobs_completed_zero_or_positive CHECK (num_jobs_completed >= 0), + CONSTRAINT num_jobs_running_zero_or_positive CHECK (num_jobs_running >= 0) +); + +INSERT INTO /* TEMPLATE: schema */river_client_queue ( + river_client_id, + name, + created_at, + max_workers, + metadata, + num_jobs_completed, + num_jobs_running, + updated_at +) +SELECT + river_client_id, + name, + created_at, + max_workers, + json(metadata), + num_jobs_completed, + num_jobs_running, + updated_at +FROM /* TEMPLATE: schema */river_client_queue_old; + +DROP TABLE /* TEMPLATE: schema */river_client_queue_old; +DROP TABLE /* TEMPLATE: schema */river_client_old; + +-- +-- Notification outbox rollback. +-- + +DROP TABLE /* TEMPLATE: schema */river_notification; diff --git a/rust/riverqueue-migrate/migrations/sqlite/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.up.sql b/rust/riverqueue-migrate/migrations/sqlite/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.up.sql new file mode 100644 index 000000000..b1ca9479e --- /dev/null +++ b/rust/riverqueue-migrate/migrations/sqlite/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.up.sql @@ -0,0 +1,261 @@ +-- +-- Notification outbox. +-- + +CREATE TABLE /* TEMPLATE: schema */river_notification ( + id integer PRIMARY KEY AUTOINCREMENT, + created_at timestamp NOT NULL DEFAULT (datetime('now', 'subsec')), + payload text NOT NULL, + topic text NOT NULL, + CONSTRAINT topic_length CHECK (length(topic) > 0 AND length(topic) < 128) +); + +CREATE INDEX /* TEMPLATE: schema */river_notification_created_at_idx ON river_notification (created_at); +CREATE INDEX /* TEMPLATE: schema */river_notification_topic_id_idx ON river_notification (topic, id); + +-- +-- SQLite JSONB conversion. +-- +-- Convert JSON text columns to JSONB binary format for more efficient storage +-- and processing, and update column defaults from json() to jsonb(). +-- +-- SQLite doesn't allow `ALTER TABLE ADD COLUMN` with non-constant defaults like +-- `jsonb('{}')`, so rebuild each affected table instead. +-- + +-- +-- river_job +-- + +DROP INDEX /* TEMPLATE: schema */river_job_kind; +DROP INDEX /* TEMPLATE: schema */river_job_state_and_finalized_at_index; +DROP INDEX /* TEMPLATE: schema */river_job_prioritized_fetching_index; +DROP INDEX /* TEMPLATE: schema */river_job_unique_idx; + +ALTER TABLE /* TEMPLATE: schema */river_job RENAME TO river_job_old; + +CREATE TABLE /* TEMPLATE: schema */river_job ( + id integer PRIMARY KEY, -- SQLite aliases this to ROWID, which may reuse deleted IDs. + args blob NOT NULL DEFAULT (jsonb('{}')), + attempt integer NOT NULL DEFAULT 0, + attempted_at timestamp, + attempted_by blob, -- json + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + errors blob, -- json + finalized_at timestamp, + kind text NOT NULL, + max_attempts integer NOT NULL, + metadata blob NOT NULL DEFAULT (jsonb('{}')), + priority integer NOT NULL DEFAULT 1, + queue text NOT NULL DEFAULT 'default', + state text NOT NULL DEFAULT 'available', + scheduled_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + tags blob NOT NULL DEFAULT (jsonb('[]')), + unique_key blob, + unique_states integer, + CONSTRAINT finalized_or_finalized_at_null CHECK ( + (finalized_at IS NULL AND state NOT IN ('cancelled', 'completed', 'discarded')) OR + (finalized_at IS NOT NULL AND state IN ('cancelled', 'completed', 'discarded')) + ), + CONSTRAINT priority_in_range CHECK (priority >= 1 AND priority <= 4), + CONSTRAINT queue_length CHECK (length(queue) > 0 AND length(queue) < 128), + CONSTRAINT kind_length CHECK (length(kind) > 0 AND length(kind) < 128), + CONSTRAINT state_valid CHECK (state IN ('available', 'cancelled', 'completed', 'discarded', 'pending', 'retryable', 'running', 'scheduled')) +); + +INSERT INTO /* TEMPLATE: schema */river_job ( + id, + args, + attempt, + attempted_at, + attempted_by, + created_at, + errors, + finalized_at, + kind, + max_attempts, + metadata, + priority, + queue, + state, + scheduled_at, + tags, + unique_key, + unique_states +) +SELECT + id, + jsonb(args), + attempt, + attempted_at, + CASE WHEN attempted_by IS NULL THEN NULL ELSE jsonb(attempted_by) END, + created_at, + CASE WHEN errors IS NULL THEN NULL ELSE jsonb(errors) END, + finalized_at, + kind, + max_attempts, + jsonb(metadata), + priority, + queue, + state, + scheduled_at, + jsonb(tags), + unique_key, + unique_states +FROM /* TEMPLATE: schema */river_job_old; + +DROP TABLE /* TEMPLATE: schema */river_job_old; + +CREATE INDEX /* TEMPLATE: schema */river_job_kind ON river_job (kind); +CREATE INDEX /* TEMPLATE: schema */river_job_state_and_finalized_at_index ON river_job (state, finalized_at) WHERE finalized_at IS NOT NULL; +CREATE INDEX /* TEMPLATE: schema */river_job_prioritized_fetching_index ON river_job (state, queue, priority, scheduled_at, id); +CREATE UNIQUE INDEX /* TEMPLATE: schema */river_job_unique_idx ON river_job (unique_key) + WHERE unique_key IS NOT NULL + AND unique_states IS NOT NULL + AND CASE state + WHEN 'available' THEN unique_states & (1 << 0) + WHEN 'cancelled' THEN unique_states & (1 << 1) + WHEN 'completed' THEN unique_states & (1 << 2) + WHEN 'discarded' THEN unique_states & (1 << 3) + WHEN 'pending' THEN unique_states & (1 << 4) + WHEN 'retryable' THEN unique_states & (1 << 5) + WHEN 'running' THEN unique_states & (1 << 6) + WHEN 'scheduled' THEN unique_states & (1 << 7) + ELSE 0 + END >= 1; + +-- +-- river_queue +-- + +ALTER TABLE /* TEMPLATE: schema */river_queue RENAME TO river_queue_old; + +CREATE TABLE /* TEMPLATE: schema */river_queue ( + name text PRIMARY KEY NOT NULL, + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + metadata blob NOT NULL DEFAULT (jsonb('{}')), + paused_at timestamp, + updated_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +INSERT INTO /* TEMPLATE: schema */river_queue ( + name, + created_at, + metadata, + paused_at, + updated_at +) +SELECT + name, + created_at, + jsonb(metadata), + paused_at, + updated_at +FROM /* TEMPLATE: schema */river_queue_old; + +DROP TABLE /* TEMPLATE: schema */river_queue_old; + +-- +-- river_client +-- + +ALTER TABLE /* TEMPLATE: schema */river_client RENAME TO river_client_old; + +CREATE TABLE /* TEMPLATE: schema */river_client ( + id text PRIMARY KEY NOT NULL, + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + metadata blob NOT NULL DEFAULT (jsonb('{}')), + paused_at timestamp, + updated_at timestamp NOT NULL, + CONSTRAINT name_length CHECK (length(id) > 0 AND length(id) < 128) +); + +INSERT INTO /* TEMPLATE: schema */river_client ( + id, + created_at, + metadata, + paused_at, + updated_at +) +SELECT + id, + created_at, + jsonb(metadata), + paused_at, + updated_at +FROM /* TEMPLATE: schema */river_client_old; + +-- +-- river_client_queue +-- + +ALTER TABLE /* TEMPLATE: schema */river_client_queue RENAME TO river_client_queue_old; + +CREATE TABLE /* TEMPLATE: schema */river_client_queue ( + river_client_id text NOT NULL REFERENCES river_client (id) ON DELETE CASCADE, + name text NOT NULL, + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + max_workers integer NOT NULL DEFAULT 0, + metadata blob NOT NULL DEFAULT (jsonb('{}')), + num_jobs_completed integer NOT NULL DEFAULT 0, + num_jobs_running integer NOT NULL DEFAULT 0, + updated_at timestamp NOT NULL, + PRIMARY KEY (river_client_id, name), + CONSTRAINT name_length CHECK (length(name) > 0 AND length(name) < 128), + CONSTRAINT num_jobs_completed_zero_or_positive CHECK (num_jobs_completed >= 0), + CONSTRAINT num_jobs_running_zero_or_positive CHECK (num_jobs_running >= 0) +); + +INSERT INTO /* TEMPLATE: schema */river_client_queue ( + river_client_id, + name, + created_at, + max_workers, + metadata, + num_jobs_completed, + num_jobs_running, + updated_at +) +SELECT + river_client_id, + name, + created_at, + max_workers, + jsonb(metadata), + num_jobs_completed, + num_jobs_running, + updated_at +FROM /* TEMPLATE: schema */river_client_queue_old; + +DROP TABLE /* TEMPLATE: schema */river_client_queue_old; +DROP TABLE /* TEMPLATE: schema */river_client_old; + +-- +-- SQL cleanup. +-- + +-- +-- Drop unused tables `river_client` and `river_client_queue`. +-- + +DROP TABLE /* TEMPLATE: schema */river_client_queue; +DROP TABLE /* TEMPLATE: schema */river_client; + +-- +-- Adds `DEFAULT 25` to `river_job.max_attempts`. +-- + +-- This may look odd in that we're adding a brand new column, but it's because +-- SQLite doesn't support anything beyond the most trivial DDL. + +ALTER TABLE /* TEMPLATE: schema */river_job + RENAME COLUMN max_attempts TO max_attempts_old; + +ALTER TABLE /* TEMPLATE: schema */river_job + ADD COLUMN max_attempts integer NOT NULL DEFAULT 25; + +UPDATE /* TEMPLATE: schema */river_job +SET max_attempts = max_attempts_old; + +ALTER TABLE /* TEMPLATE: schema */river_job + DROP COLUMN max_attempts_old; diff --git a/rust/riverqueue-migrate/migrations/sqlite/main/008_job_id_autoincrement.down.sql b/rust/riverqueue-migrate/migrations/sqlite/main/008_job_id_autoincrement.down.sql new file mode 100644 index 000000000..aad4f366b --- /dev/null +++ b/rust/riverqueue-migrate/migrations/sqlite/main/008_job_id_autoincrement.down.sql @@ -0,0 +1,121 @@ +-- Rebuild river_job to restore SQLite's default ROWID allocation behavior. + +-- Rebuilding river_job would discard schema installed by River Pro. Check +-- schema objects instead of migration records to also catch manually applied +-- Pro migrations and the legacy workflow migration line. +CREATE TEMP TABLE river_job_pro_schema_guard ( + id integer NOT NULL +); + +CREATE TEMP TRIGGER river_job_pro_schema_guard_enforce + BEFORE INSERT ON river_job_pro_schema_guard + WHEN EXISTS ( + SELECT 1 + FROM /* TEMPLATE: schema */sqlite_master + WHERE name IN ('river_job_sequence', 'river_job_workflow_scheduling', 'river_workflow') + ) +BEGIN + SELECT RAISE(ABORT, 'River SQLite migration 008 cannot run while River Pro schema is installed'); +END; + +INSERT INTO river_job_pro_schema_guard (id) VALUES (1); + +DROP TRIGGER river_job_pro_schema_guard_enforce; +DROP TABLE river_job_pro_schema_guard; + +DROP INDEX /* TEMPLATE: schema */river_job_kind; +DROP INDEX /* TEMPLATE: schema */river_job_state_and_finalized_at_index; +DROP INDEX /* TEMPLATE: schema */river_job_prioritized_fetching_index; +DROP INDEX /* TEMPLATE: schema */river_job_unique_idx; + +ALTER TABLE /* TEMPLATE: schema */river_job RENAME TO river_job_old; + +CREATE TABLE /* TEMPLATE: schema */river_job ( + id integer PRIMARY KEY, + args blob NOT NULL DEFAULT (jsonb('{}')), + attempt integer NOT NULL DEFAULT 0, + attempted_at timestamp, + attempted_by blob, -- json + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + errors blob, -- json + finalized_at timestamp, + kind text NOT NULL, + max_attempts integer NOT NULL DEFAULT 25, + metadata blob NOT NULL DEFAULT (jsonb('{}')), + priority integer NOT NULL DEFAULT 1, + queue text NOT NULL DEFAULT 'default', + state text NOT NULL DEFAULT 'available', + scheduled_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + tags blob NOT NULL DEFAULT (jsonb('[]')), + unique_key blob, + unique_states integer, + CONSTRAINT finalized_or_finalized_at_null CHECK ( + (finalized_at IS NULL AND state NOT IN ('cancelled', 'completed', 'discarded')) OR + (finalized_at IS NOT NULL AND state IN ('cancelled', 'completed', 'discarded')) + ), + CONSTRAINT priority_in_range CHECK (priority >= 1 AND priority <= 4), + CONSTRAINT queue_length CHECK (length(queue) > 0 AND length(queue) < 128), + CONSTRAINT kind_length CHECK (length(kind) > 0 AND length(kind) < 128), + CONSTRAINT state_valid CHECK (state IN ('available', 'cancelled', 'completed', 'discarded', 'pending', 'retryable', 'running', 'scheduled')) +); + +INSERT INTO /* TEMPLATE: schema */river_job ( + id, + args, + attempt, + attempted_at, + attempted_by, + created_at, + errors, + finalized_at, + kind, + max_attempts, + metadata, + priority, + queue, + state, + scheduled_at, + tags, + unique_key, + unique_states +) +SELECT + id, + args, + attempt, + attempted_at, + attempted_by, + created_at, + errors, + finalized_at, + kind, + max_attempts, + metadata, + priority, + queue, + state, + scheduled_at, + tags, + unique_key, + unique_states +FROM /* TEMPLATE: schema */river_job_old; + +DROP TABLE /* TEMPLATE: schema */river_job_old; + +CREATE INDEX /* TEMPLATE: schema */river_job_kind ON river_job (kind); +CREATE INDEX /* TEMPLATE: schema */river_job_state_and_finalized_at_index ON river_job (state, finalized_at) WHERE finalized_at IS NOT NULL; +CREATE INDEX /* TEMPLATE: schema */river_job_prioritized_fetching_index ON river_job (state, queue, priority, scheduled_at, id); +CREATE UNIQUE INDEX /* TEMPLATE: schema */river_job_unique_idx ON river_job (unique_key) + WHERE unique_key IS NOT NULL + AND unique_states IS NOT NULL + AND CASE state + WHEN 'available' THEN unique_states & (1 << 0) + WHEN 'cancelled' THEN unique_states & (1 << 1) + WHEN 'completed' THEN unique_states & (1 << 2) + WHEN 'discarded' THEN unique_states & (1 << 3) + WHEN 'pending' THEN unique_states & (1 << 4) + WHEN 'retryable' THEN unique_states & (1 << 5) + WHEN 'running' THEN unique_states & (1 << 6) + WHEN 'scheduled' THEN unique_states & (1 << 7) + ELSE 0 + END >= 1; diff --git a/rust/riverqueue-migrate/migrations/sqlite/main/008_job_id_autoincrement.up.sql b/rust/riverqueue-migrate/migrations/sqlite/main/008_job_id_autoincrement.up.sql new file mode 100644 index 000000000..c7de15deb --- /dev/null +++ b/rust/riverqueue-migrate/migrations/sqlite/main/008_job_id_autoincrement.up.sql @@ -0,0 +1,123 @@ +-- Rebuild river_job so automatically generated IDs are never reused after the +-- job holding the largest ID is deleted. Unlike PostgreSQL sequences, SQLite's +-- default ROWID allocator may otherwise reuse that deleted ID. + +-- Rebuilding river_job would discard schema installed by River Pro. Check +-- schema objects instead of migration records to also catch manually applied +-- Pro migrations and the legacy workflow migration line. +CREATE TEMP TABLE river_job_pro_schema_guard ( + id integer NOT NULL +); + +CREATE TEMP TRIGGER river_job_pro_schema_guard_enforce + BEFORE INSERT ON river_job_pro_schema_guard + WHEN EXISTS ( + SELECT 1 + FROM /* TEMPLATE: schema */sqlite_master + WHERE name IN ('river_job_sequence', 'river_job_workflow_scheduling', 'river_workflow') + ) +BEGIN + SELECT RAISE(ABORT, 'River SQLite migration 008 cannot run while River Pro schema is installed'); +END; + +INSERT INTO river_job_pro_schema_guard (id) VALUES (1); + +DROP TRIGGER river_job_pro_schema_guard_enforce; +DROP TABLE river_job_pro_schema_guard; + +DROP INDEX /* TEMPLATE: schema */river_job_kind; +DROP INDEX /* TEMPLATE: schema */river_job_state_and_finalized_at_index; +DROP INDEX /* TEMPLATE: schema */river_job_prioritized_fetching_index; +DROP INDEX /* TEMPLATE: schema */river_job_unique_idx; + +ALTER TABLE /* TEMPLATE: schema */river_job RENAME TO river_job_old; + +CREATE TABLE /* TEMPLATE: schema */river_job ( + id integer PRIMARY KEY AUTOINCREMENT, + args blob NOT NULL DEFAULT (jsonb('{}')), + attempt integer NOT NULL DEFAULT 0, + attempted_at timestamp, + attempted_by blob, -- json + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + errors blob, -- json + finalized_at timestamp, + kind text NOT NULL, + max_attempts integer NOT NULL DEFAULT 25, + metadata blob NOT NULL DEFAULT (jsonb('{}')), + priority integer NOT NULL DEFAULT 1, + queue text NOT NULL DEFAULT 'default', + state text NOT NULL DEFAULT 'available', + scheduled_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + tags blob NOT NULL DEFAULT (jsonb('[]')), + unique_key blob, + unique_states integer, + CONSTRAINT finalized_or_finalized_at_null CHECK ( + (finalized_at IS NULL AND state NOT IN ('cancelled', 'completed', 'discarded')) OR + (finalized_at IS NOT NULL AND state IN ('cancelled', 'completed', 'discarded')) + ), + CONSTRAINT priority_in_range CHECK (priority >= 1 AND priority <= 4), + CONSTRAINT queue_length CHECK (length(queue) > 0 AND length(queue) < 128), + CONSTRAINT kind_length CHECK (length(kind) > 0 AND length(kind) < 128), + CONSTRAINT state_valid CHECK (state IN ('available', 'cancelled', 'completed', 'discarded', 'pending', 'retryable', 'running', 'scheduled')) +); + +INSERT INTO /* TEMPLATE: schema */river_job ( + id, + args, + attempt, + attempted_at, + attempted_by, + created_at, + errors, + finalized_at, + kind, + max_attempts, + metadata, + priority, + queue, + state, + scheduled_at, + tags, + unique_key, + unique_states +) +SELECT + id, + args, + attempt, + attempted_at, + attempted_by, + created_at, + errors, + finalized_at, + kind, + max_attempts, + metadata, + priority, + queue, + state, + scheduled_at, + tags, + unique_key, + unique_states +FROM /* TEMPLATE: schema */river_job_old; + +DROP TABLE /* TEMPLATE: schema */river_job_old; + +CREATE INDEX /* TEMPLATE: schema */river_job_kind ON river_job (kind); +CREATE INDEX /* TEMPLATE: schema */river_job_state_and_finalized_at_index ON river_job (state, finalized_at) WHERE finalized_at IS NOT NULL; +CREATE INDEX /* TEMPLATE: schema */river_job_prioritized_fetching_index ON river_job (state, queue, priority, scheduled_at, id); +CREATE UNIQUE INDEX /* TEMPLATE: schema */river_job_unique_idx ON river_job (unique_key) + WHERE unique_key IS NOT NULL + AND unique_states IS NOT NULL + AND CASE state + WHEN 'available' THEN unique_states & (1 << 0) + WHEN 'cancelled' THEN unique_states & (1 << 1) + WHEN 'completed' THEN unique_states & (1 << 2) + WHEN 'discarded' THEN unique_states & (1 << 3) + WHEN 'pending' THEN unique_states & (1 << 4) + WHEN 'retryable' THEN unique_states & (1 << 5) + WHEN 'running' THEN unique_states & (1 << 6) + WHEN 'scheduled' THEN unique_states & (1 << 7) + ELSE 0 + END >= 1; diff --git a/rust/riverqueue-migrate/src/lib.rs b/rust/riverqueue-migrate/src/lib.rs new file mode 100644 index 000000000..497e26de0 --- /dev/null +++ b/rust/riverqueue-migrate/src/lib.rs @@ -0,0 +1,764 @@ +// The README's example uses both backends, so it's the crate documentation +// (and its example runs as a doctest) when both are enabled, as on docs.rs. +#![cfg_attr( + all(feature = "postgres", feature = "sqlite"), + doc = include_str!("../README.md") +)] +#![cfg_attr( + not(all(feature = "postgres", feature = "sqlite")), + doc = "River's PostgreSQL and SQLite migrations for Rust. The full crate documentation, \ + from the README, builds with both the `postgres` and `sqlite` features, as on \ + [docs.rs](https://docs.rs/riverqueue-migrate)." +)] +#![cfg_attr(docsrs, feature(doc_cfg))] +#![forbid(unsafe_code)] +#![warn(missing_docs)] + +#[cfg(not(any(feature = "postgres", feature = "sqlite")))] +compile_error!("riverqueue-migrate requires at least one database feature: `postgres` or `sqlite`"); + +use std::time::Duration; +#[cfg(feature = "postgres")] +use std::time::Instant; + +#[cfg(feature = "postgres")] +use sqlx::{PgPool, Row}; +use thiserror::Error; + +mod schema; +#[cfg(feature = "sqlite")] +mod sqlite; + +pub use schema::{SCHEMA_MAX_LEN, SchemaName, SchemaNameError}; + +#[cfg(feature = "sqlite")] +pub use sqlite::{SQLITE_MIGRATIONS, SqliteMigrator}; + +#[cfg(feature = "postgres")] +const TEMPLATE_SCHEMA: &str = "/* TEMPLATE: schema */"; + +/// River's main migration line. +pub const MIGRATION_LINE_MAIN: &str = "main"; + +/// Latest migration version bundled with this release. +pub const MIGRATION_VERSION_LATEST: i64 = 8; + +#[cfg(feature = "postgres")] +macro_rules! migration { + ($version:literal, $name:literal, $file:literal) => { + Migration { + down_sql: include_str!(concat!("../migrations/main/", $file, ".down.sql")), + name: $name, + up_sql: include_str!(concat!("../migrations/main/", $file, ".up.sql")), + version: $version, + } + }; +} + +/// One canonical River migration. +#[derive(Clone, Copy, Debug)] +#[non_exhaustive] +pub struct Migration { + /// Down migration SQL. + pub down_sql: &'static str, + /// Human-readable migration name. + pub name: &'static str, + /// Up migration SQL. + pub up_sql: &'static str, + /// Monotonically increasing version. + pub version: i64, +} + +/// Migration direction. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Direction { + /// Revert applied migrations. + Down, + /// Apply outstanding migrations. + Up, +} + +/// Controls a migration operation. +#[derive(Clone, Copy, Debug, Default)] +pub struct MigrateOpts { + /// Report SQL without applying it. + dry_run: bool, + /// Maximum number of steps. Down migrations default to one step. + max_steps: Option, + /// Target schema version. Down excludes the target; `-1` removes River. + target_version: Option, +} + +impl MigrateOpts { + /// Creates migration options with no target, step limit, or dry run. + #[must_use] + pub const fn new() -> Self { + Self { + dry_run: false, + max_steps: None, + target_version: None, + } + } + + /// Returns whether SQL is reported without being applied. + #[must_use] + pub const fn dry_run(&self) -> bool { + self.dry_run + } + + /// Returns the maximum number of migration steps. + #[must_use] + pub const fn max_steps(&self) -> Option { + self.max_steps + } + + /// Returns the requested target version. + #[must_use] + pub const fn target_version(&self) -> Option { + self.target_version + } + + /// Reports selected SQL without applying it. + #[must_use] + pub const fn with_dry_run(mut self, dry_run: bool) -> Self { + self.dry_run = dry_run; + self + } + + /// Limits the number of migration steps. + #[must_use] + pub const fn with_max_steps(mut self, maximum: usize) -> Self { + self.max_steps = Some(maximum); + self + } + + /// Migrates toward a target schema version. `-1` removes River. An up + /// migration whose target is already applied does nothing, even if later + /// versions are pending. + #[must_use] + pub const fn with_target_version(mut self, version: i64) -> Self { + self.target_version = Some(version); + self + } +} + +/// One migration selected or applied by an operation. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct MigrateVersion { + /// Database execution time, or zero for a dry run. + pub duration: Duration, + /// Human-readable migration name. + pub name: &'static str, + /// Rendered SQL. + pub sql: String, + /// Migration version. + pub version: i64, +} + +/// Result of a migration operation. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct MigrateResult { + /// Direction requested. + pub direction: Direction, + /// Versions applied or selected. + pub versions: Vec, +} + +/// Result of checking whether required migrations are applied. +/// +/// Its `Display` output describes the unapplied versions, for reporting a +/// failed validation. +#[derive(Clone, Debug, Default, Eq, PartialEq)] +#[non_exhaustive] +pub struct ValidateResult { + /// Required versions that aren't applied, in ascending order. + pub unapplied: Vec, +} + +impl ValidateResult { + /// Returns whether every required migration is applied. + #[must_use] + pub fn is_valid(&self) -> bool { + self.unapplied.is_empty() + } +} + +impl std::fmt::Display for ValidateResult { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if self.is_valid() { + formatter.write_str("all required migrations are applied") + } else { + write!(formatter, "unapplied migrations: {:?}", self.unapplied) + } + } +} + +/// Canonical PostgreSQL migration bundle. +#[cfg(feature = "postgres")] +pub const POSTGRES_MIGRATIONS: [Migration; 8] = [ + migration!(1, "create_river_migration", "001_create_river_migration"), + migration!(2, "initial_schema", "002_initial_schema"), + migration!(3, "river_job_tags_non_null", "003_river_job_tags_non_null"), + migration!(4, "pending_and_more", "004_pending_and_more"), + migration!(5, "migration_unique_client", "005_migration_unique_client"), + migration!(6, "bulk_unique", "006_bulk_unique"), + migration!( + 7, + "notification_outbox_sqlite_jsonb_and_sql_cleanup", + "007_notification_outbox_sqlite_jsonb_and_sql_cleanup" + ), + migration!(8, "job_id_autoincrement", "008_job_id_autoincrement"), +]; + +/// River migration failure. +#[derive(Debug, Error)] +#[non_exhaustive] +pub enum Error { + /// A database operation failed. + #[error(transparent)] + Database(#[from] sqlx::Error), + + /// Migrating down past `version` would delete the migration records of + /// other migration lines, which the down migration can't restore. + #[error( + "found non-main migration lines; version {version:03} is irreversible without losing migration information" + )] + OtherMigrationLines { + /// The migration that would be reverted. + version: i64, + }, + + /// A down migration's target isn't among the applied versions it can + /// revert to within its step limit. + #[error("version {version} is not in target list of valid migrations to apply")] + TargetNotSelected { + /// The requested target version. + version: i64, + }, + + /// A target version isn't one of River's migrations. + #[error("version {version} is not a River migration")] + UnknownVersion { + /// The requested target version. + version: i64, + }, +} + +/// Applies and validates River's PostgreSQL migration history. +#[cfg(feature = "postgres")] +#[derive(Clone, Debug)] +pub struct PostgresMigrator { + pool: PgPool, + schema: SchemaName, +} + +#[cfg(feature = "postgres")] +impl PostgresMigrator { + /// Returns every migration bundled with this crate. + #[must_use] + pub fn all_versions() -> &'static [Migration] { + &POSTGRES_MIGRATIONS + } + + /// Creates a migrator for PostgreSQL's current schema. + #[must_use] + pub fn new(pool: PgPool) -> Self { + Self { + pool, + schema: SchemaName::current(), + } + } + + /// Returns the migrator with `schema` migrated instead of the + /// connection's current schema. + #[must_use] + pub fn with_schema(mut self, schema: SchemaName) -> Self { + self.schema = schema; + self + } + + /// Returns the schema this migrator migrates. + #[must_use] + pub const fn schema(&self) -> &SchemaName { + &self.schema + } + + /// Returns applied main-line versions in ascending order. + /// + /// # Errors + /// + /// Returns [`Error::Database`] when the query fails. + pub async fn existing_versions(&self) -> Result, Error> { + let table = self.schema.qualify("river_migration"); + // Pass the quoted, qualified name through unchanged like Go's + // `TableExists`, so a mixed-case schema is not folded to lowercase. + let exists: bool = sqlx::query_scalar("SELECT to_regclass($1) IS NOT NULL") + .bind(&table) + .fetch_one(&self.pool) + .await?; + if !exists { + return Ok(Vec::new()); + } + + let has_line: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_schema = coalesce($1, current_schema()) AND table_name = 'river_migration' AND column_name = 'line')", + ) + .bind(self.schema.as_deref()) + .fetch_one(&self.pool) + .await?; + let sql = if has_line { + format!("SELECT version FROM {table} WHERE line = $1 ORDER BY version") + } else { + format!("SELECT version FROM {table} ORDER BY version") + }; + let rows = if has_line { + sqlx::query(sqlx::AssertSqlSafe(sql)) + .bind(MIGRATION_LINE_MAIN) + .fetch_all(&self.pool) + .await? + } else { + sqlx::query(sqlx::AssertSqlSafe(sql)) + .fetch_all(&self.pool) + .await? + }; + Ok(rows.iter().map(|row| row.get("version")).collect()) + } + + /// Applies all outstanding up migrations and returns their versions. + /// + /// # Errors + /// + /// Returns [`Error::Database`] when a migration fails. + /// + /// # Cancel safety + /// + /// Each migration and its record in `river_migration` commit together in + /// their own transaction, on a task of their own. Dropping the future + /// stops migrating once the migration in progress finishes: it and + /// every migration before it stay applied, and migrating again + /// continues from there. + pub async fn migrate_up(&self) -> Result, Error> { + Ok(self + .migrate(Direction::Up, MigrateOpts::default()) + .await? + .versions + .into_iter() + .map(|version| version.version) + .collect()) + } + + /// Applies up or down migrations with target, step, and dry-run controls. + /// + /// Each migration runs in its own transaction, so a failure leaves the + /// migrations before it applied. + /// + /// # Errors + /// + /// Returns [`Error::UnknownVersion`] when the target version doesn't exist, + /// [`Error::TargetNotSelected`] when a down target isn't applied or is + /// beyond the step limit, and [`Error::Database`] when a migration fails. + /// + /// # Cancel safety + /// + /// Each migration and its record in `river_migration` commit together in + /// their own transaction, on a task of their own. Dropping the future + /// stops migrating once the migration in progress finishes: it and + /// every migration before it stay applied, and migrating again + /// continues from there. + pub async fn migrate( + &self, + direction: Direction, + opts: MigrateOpts, + ) -> Result { + validate_target(&POSTGRES_MIGRATIONS, opts.target_version, true)?; + let applied = self.existing_versions().await?; + let selected = select_migrations(&POSTGRES_MIGRATIONS, direction, opts, &applied)?; + + let mut versions = Vec::with_capacity(selected.len()); + for migration in selected { + let sql = self.render(match direction { + Direction::Down => migration.down_sql, + Direction::Up => migration.up_sql, + }); + let mut duration = Duration::ZERO; + if !opts.dry_run { + let started_at = Instant::now(); + // Each migration runs to completion on its own task, so + // dropping this future never abandons one partway. + let migrator = self.clone(); + let task_sql = sql.clone(); + run_to_completion( + async move { migrator.apply(direction, migration, &task_sql).await }, + ) + .await?; + duration = started_at.elapsed(); + } + versions.push(MigrateVersion { + duration, + name: migration.name, + sql, + version: migration.version, + }); + } + Ok(MigrateResult { + direction, + versions, + }) + } + + /// Checks that every migration through an optional target is applied. + /// + /// # Errors + /// + /// Returns [`Error::UnknownVersion`] when the target version doesn't exist and + /// [`Error::Database`] when reading the applied versions fails. + pub async fn validate(&self, target_version: Option) -> Result { + validate_target(&POSTGRES_MIGRATIONS, target_version, false)?; + let applied = self.existing_versions().await?; + Ok(validate_migrations( + &POSTGRES_MIGRATIONS, + target_version, + &applied, + )) + } + + async fn apply( + &self, + direction: Direction, + migration: Migration, + sql: &str, + ) -> Result<(), Error> { + let mut transaction = self.pool.begin().await?; + // The only dynamic fragment is a validated and quoted schema name. + sqlx::raw_sql(sqlx::AssertSqlSafe(sql)) + .execute(&mut *transaction) + .await?; + let table = self.schema.qualify("river_migration"); + match direction { + Direction::Down if migration.version == 1 => {} + Direction::Down if migration.version <= 5 => { + sqlx::query(sqlx::AssertSqlSafe(format!( + "DELETE FROM {table} WHERE version = $1" + ))) + .bind(migration.version) + .execute(&mut *transaction) + .await?; + } + Direction::Down => { + sqlx::query(sqlx::AssertSqlSafe(format!( + "DELETE FROM {table} WHERE line = $1 AND version = $2" + ))) + .bind(MIGRATION_LINE_MAIN) + .bind(migration.version) + .execute(&mut *transaction) + .await?; + } + Direction::Up if migration.version >= 5 => { + sqlx::query(sqlx::AssertSqlSafe(format!( + "INSERT INTO {table} (line, version) VALUES ($1, $2)" + ))) + .bind(MIGRATION_LINE_MAIN) + .bind(migration.version) + .execute(&mut *transaction) + .await?; + } + Direction::Up => { + sqlx::query(sqlx::AssertSqlSafe(format!( + "INSERT INTO {table} (version) VALUES ($1)" + ))) + .bind(migration.version) + .execute(&mut *transaction) + .await?; + } + } + transaction.commit().await?; + Ok(()) + } + + fn render(&self, sql: &str) -> String { + sql.replace(TEMPLATE_SCHEMA, &self.schema.migration_prefix()) + } +} + +/// Runs a migration step on its own task and waits for it, so the step +/// finishes even when the caller stops waiting. +async fn run_to_completion( + step: impl std::future::Future> + Send + 'static, +) -> Result<(), Error> { + match tokio::spawn(step).await { + Ok(result) => result, + Err(error) if error.is_panic() => std::panic::resume_unwind(error.into_panic()), + // The runtime is shutting down and cancelled the task. + Err(error) => Err(Error::Database(sqlx::Error::Io(std::io::Error::other( + format!("migration task failed: {error}"), + )))), + } +} + +fn select_migrations( + migrations: &'static [Migration], + direction: Direction, + opts: MigrateOpts, + applied: &[i64], +) -> Result, Error> { + // Versions recorded by a newer River release are ignored like Go: up + // migrations apply only unapplied known versions, and down migrations + // revert only applied known versions. + let mut selected = match direction { + Direction::Up => migrations + .iter() + .filter(|migration| !applied.contains(&migration.version)) + .copied() + .collect::>(), + Direction::Down => migrations + .iter() + .rev() + .filter(|migration| applied.contains(&migration.version)) + .copied() + .collect::>(), + }; + + // An up migration whose target is already applied does nothing, even + // when later versions are pending. Check before the step limit trims + // the list, which can also drop a pending target. + let target_was_pending = opts + .target_version + .is_some_and(|target| selected.iter().any(|migration| migration.version == target)); + + // Go limits steps before locating the target, so a target outside the + // step window is not reached. + let maximum = opts + .max_steps + .or_else(|| (direction == Direction::Down && opts.target_version.is_none()).then_some(1)); + if let Some(maximum) = maximum { + selected.truncate(maximum); + } + + if let Some(target) = opts.target_version.filter(|target| *target != -1) { + match selected + .iter() + .position(|migration| migration.version == target) + { + Some(index) => { + selected.truncate(index + 1); + // A down target is the version that remains applied. + if direction == Direction::Down { + selected.pop(); + } + } + None if direction == Direction::Down => { + return Err(Error::TargetNotSelected { version: target }); + } + // A pending up target beyond the step limit keeps the trimmed + // list; an applied one is a no-op. + None if target_was_pending => {} + None => selected.clear(), + } + } + Ok(selected) +} + +fn validate_migrations( + migrations: &[Migration], + target_version: Option, + applied: &[i64], +) -> ValidateResult { + let missing = migrations + .iter() + .filter(|migration| target_version.is_none_or(|target| migration.version <= target)) + .filter(|migration| !applied.contains(&migration.version)) + .map(|migration| migration.version) + .collect::>(); + ValidateResult { unapplied: missing } +} + +fn validate_target( + migrations: &[Migration], + target_version: Option, + allow_empty: bool, +) -> Result<(), Error> { + if let Some(target) = target_version + && !(allow_empty && target == -1) + && !migrations + .iter() + .any(|migration| migration.version == target) + { + return Err(Error::UnknownVersion { version: target }); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::{Direction, Error, MigrateOpts, Migration, select_migrations}; + + const MIGRATIONS: [Migration; 7] = { + const fn migration(version: i64) -> Migration { + Migration { + down_sql: "", + name: "test", + up_sql: "", + version, + } + } + [ + migration(1), + migration(2), + migration(3), + migration(4), + migration(5), + migration(6), + migration(7), + ] + }; + + fn versions( + direction: Direction, + opts: MigrateOpts, + applied: &[i64], + ) -> Result, Error> { + select_migrations(&MIGRATIONS, direction, opts, applied).map(|selected| { + selected + .into_iter() + .map(|migration| migration.version) + .collect() + }) + } + + #[test] + fn applied_up_targets_are_no_ops() { + // Like Go, even with a gap below the target or later versions + // pending, with or without a step limit. + assert_eq!( + versions( + Direction::Up, + MigrateOpts::new().with_target_version(3), + &[1, 3] + ) + .unwrap(), + Vec::::new() + ); + assert_eq!( + versions( + Direction::Up, + MigrateOpts::new().with_target_version(3).with_max_steps(1), + &[1, 2, 3] + ) + .unwrap(), + Vec::::new() + ); + } + + #[test] + fn selection_matches_go_target_and_step_semantics() { + let all = [1, 2, 3, 4, 5, 6, 7]; + + // Down defaults to one step, and a target is the version left applied. + assert_eq!( + versions(Direction::Down, MigrateOpts::new(), &all).unwrap(), + [7] + ); + assert_eq!( + versions( + Direction::Down, + MigrateOpts::new().with_target_version(4), + &all + ) + .unwrap(), + [7, 6, 5] + ); + assert_eq!( + versions( + Direction::Down, + MigrateOpts::new().with_target_version(4), + &[1, 2, 3, 4] + ) + .unwrap(), + Vec::::new() + ); + assert_eq!( + versions( + Direction::Down, + MigrateOpts::new().with_target_version(-1), + &all + ) + .unwrap(), + [7, 6, 5, 4, 3, 2, 1] + ); + + // Steps limit the list before the target is located, like Go. + assert_eq!( + versions( + Direction::Down, + MigrateOpts::new().with_target_version(5).with_max_steps(3), + &all + ) + .unwrap(), + [7, 6] + ); + assert!(matches!( + versions( + Direction::Down, + MigrateOpts::new().with_target_version(5).with_max_steps(2), + &all + ), + Err(Error::TargetNotSelected { version: 5 }) + )); + + // A down target that is not applied is an error rather than a no-op. + assert!(matches!( + versions( + Direction::Down, + MigrateOpts::new().with_target_version(5), + &[1, 2, 3] + ), + Err(Error::TargetNotSelected { version: 5 }) + )); + + // Up targets stop at the target and are no-ops once applied. + assert_eq!( + versions( + Direction::Up, + MigrateOpts::new().with_target_version(5), + &[1, 2] + ) + .unwrap(), + [3, 4, 5] + ); + assert_eq!( + versions( + Direction::Up, + MigrateOpts::new().with_target_version(2), + &[1, 2, 3] + ) + .unwrap(), + Vec::::new() + ); + assert_eq!( + versions( + Direction::Up, + MigrateOpts::new().with_target_version(6).with_max_steps(2), + &[1] + ) + .unwrap(), + [2, 3] + ); + + // Versions recorded by a newer release are ignored. + assert_eq!( + versions(Direction::Up, MigrateOpts::new(), &[1, 2, 3, 99]).unwrap(), + [4, 5, 6, 7] + ); + assert_eq!( + versions( + Direction::Down, + MigrateOpts::new(), + &[1, 2, 3, 4, 5, 6, 7, 8] + ) + .unwrap(), + [7] + ); + } +} diff --git a/rust/riverqueue-migrate/src/schema.rs b/rust/riverqueue-migrate/src/schema.rs new file mode 100644 index 000000000..92c4cd6ef --- /dev/null +++ b/rust/riverqueue-migrate/src/schema.rs @@ -0,0 +1,178 @@ +//! PostgreSQL schema names. + +use std::fmt; + +use thiserror::Error; + +/// PostgreSQL's maximum identifier length. +const POSTGRES_IDENTIFIER_MAX: usize = 63; + +/// Longest River notification topic, which a schema name must leave room for. +const NOTIFICATION_TOPIC_LONGEST: &str = "river_leadership"; + +/// Maximum schema length after reserving `.river_leadership` for +/// notification channels. +pub const SCHEMA_MAX_LEN: usize = POSTGRES_IDENTIFIER_MAX - NOTIFICATION_TOPIC_LONGEST.len() - 1; + +/// The PostgreSQL schema River's tables live in. +/// +/// [`SchemaName::current`] uses the connection's current schema (normally +/// `public`, following `search_path`). An explicit schema is quoted wherever +/// River renders it, so names such as `river-prod` or `MyRiver` work as +/// written. Pass the same schema to the +/// migrator and the client. +/// +/// ``` +/// use riverqueue_migrate::SchemaName; +/// +/// let schema = SchemaName::new("river")?; +/// assert_eq!(schema.qualify("river_job"), r#""river"."river_job""#); +/// # Ok::<(), riverqueue_migrate::SchemaNameError>(()) +/// ``` +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct SchemaName(Option); + +impl SchemaName { + /// Uses PostgreSQL's current schema. + #[must_use] + pub const fn current() -> Self { + Self(None) + } + + /// Validates an optional explicit schema. + /// + /// Any name is accepted and quoted when rendered, including mixed case + /// and punctuation such as + /// `river-prod`. Names containing NUL are rejected, as are names too long + /// to prefix River's notification topics within PostgreSQL's identifier + /// limit. + /// + /// # Errors + /// + /// Returns an error when the schema is too long or contains NUL. + pub fn new(schema: impl Into) -> Result { + let schema = schema.into(); + if schema.is_empty() { + return Ok(Self::current()); + } + if schema.len() > SCHEMA_MAX_LEN { + return Err(SchemaNameError::TooLong { + length: schema.len(), + maximum: SCHEMA_MAX_LEN, + }); + } + if schema.contains('\0') { + return Err(SchemaNameError::Invalid(schema)); + } + + Ok(Self(Some(schema))) + } + + /// Returns the unquoted explicit schema, if configured. + #[must_use] + pub fn as_deref(&self) -> Option<&str> { + self.0.as_deref() + } + + /// Qualifies and quotes a database object name in this schema, for use in + /// SQL that refers to River's tables. + #[must_use] + pub fn qualify(&self, object: &str) -> String { + match &self.0 { + Some(schema) => format!("{}.{}", quote_identifier(schema), quote_identifier(object)), + None => quote_identifier(object), + } + } + + /// Returns the explicit schema quoted as a PostgreSQL identifier, for + /// SQL that names the schema itself, or `None` for the connection's + /// current schema. + #[must_use] + pub fn quoted_identifier(&self) -> Option { + self.0.as_deref().map(quote_identifier) + } + + /// Prefix used by River's canonical PostgreSQL migration templates. + #[cfg(feature = "postgres")] + pub(crate) fn migration_prefix(&self) -> String { + self.quoted_identifier() + .map_or_else(String::new, |schema| format!("{schema}.")) + } +} + +/// Quotes a PostgreSQL identifier, doubling embedded quotes like Go's +/// `dbutil.SafeIdentifier`. +fn quote_identifier(identifier: &str) -> String { + format!("\"{}\"", identifier.replace('"', "\"\"")) +} + +impl Default for SchemaName { + fn default() -> Self { + Self::current() + } +} + +impl fmt::Display for SchemaName { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(self.as_deref().unwrap_or("")) + } +} + +/// Invalid River schema name. +#[derive(Debug, Error)] +#[non_exhaustive] +pub enum SchemaNameError { + /// Schema contains a NUL character, which PostgreSQL identifiers cannot. + #[error("schema name cannot contain NUL: {0:?}")] + Invalid(String), + + /// Schema is too long to prefix River's notification topics. + #[error("schema length {length} exceeds maximum {maximum}")] + TooLong { + /// Observed byte length. + length: usize, + /// Maximum byte length. + maximum: usize, + }, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn schema_name_validates_and_qualifies() { + let schema = SchemaName::new("river_test").unwrap(); + assert_eq!(schema.qualify("river_job"), "\"river_test\".\"river_job\""); + + // Go quotes any schema with `SafeIdentifier`, so Rust accepts the + // same names and escapes embedded quotes. + let hyphenated = SchemaName::new("river-prod").unwrap(); + assert_eq!( + hyphenated.qualify("river_job"), + "\"river-prod\".\"river_job\"" + ); + assert_eq!( + SchemaName::new("odd\"name").unwrap().quoted_identifier(), + Some("\"odd\"\"name\"".to_owned()) + ); + assert_eq!(SchemaName::current().quoted_identifier(), None); + assert_eq!( + SchemaName::new("odd\"name").unwrap().qualify("river_job"), + "\"odd\"\"name\".\"river_job\"" + ); + assert!(SchemaName::new("1leading_digit").is_ok()); + assert!(SchemaName::new("nul\0byte").is_err()); + assert!(SchemaName::new("a".repeat(SCHEMA_MAX_LEN + 1)).is_err()); + } + + #[cfg(feature = "postgres")] + #[test] + fn migration_prefix_qualifies_templates() { + assert_eq!( + SchemaName::new("MyRiver").unwrap().migration_prefix(), + "\"MyRiver\"." + ); + assert_eq!(SchemaName::current().migration_prefix(), ""); + } +} diff --git a/rust/riverqueue-migrate/src/sqlite.rs b/rust/riverqueue-migrate/src/sqlite.rs new file mode 100644 index 000000000..3a268c47f --- /dev/null +++ b/rust/riverqueue-migrate/src/sqlite.rs @@ -0,0 +1,262 @@ +use std::time::{Duration, Instant}; + +use sqlx::{Row, SqlitePool}; + +use crate::{ + Direction, Error, MIGRATION_LINE_MAIN, MigrateOpts, MigrateResult, MigrateVersion, Migration, + ValidateResult, run_to_completion, select_migrations, validate_migrations, validate_target, +}; + +macro_rules! sqlite_migration { + ($version:literal, $name:literal, $file:literal) => { + Migration { + down_sql: include_str!(concat!("../migrations/sqlite/main/", $file, ".down.sql")), + name: $name, + up_sql: include_str!(concat!("../migrations/sqlite/main/", $file, ".up.sql")), + version: $version, + } + }; +} + +/// Canonical SQLite migration bundle. +pub const SQLITE_MIGRATIONS: [Migration; 8] = [ + sqlite_migration!(1, "create_river_migration", "001_create_river_migration"), + sqlite_migration!(2, "initial_schema", "002_initial_schema"), + sqlite_migration!(3, "river_job_tags_non_null", "003_river_job_tags_non_null"), + sqlite_migration!(4, "pending_and_more", "004_pending_and_more"), + sqlite_migration!(5, "migration_unique_client", "005_migration_unique_client"), + sqlite_migration!(6, "bulk_unique", "006_bulk_unique"), + sqlite_migration!( + 7, + "notification_outbox_sqlite_jsonb_and_sql_cleanup", + "007_notification_outbox_sqlite_jsonb_and_sql_cleanup" + ), + sqlite_migration!(8, "job_id_autoincrement", "008_job_id_autoincrement"), +]; + +/// Applies and validates River's SQLite migration history. +#[derive(Clone, Debug)] +pub struct SqliteMigrator { + pool: SqlitePool, +} + +impl SqliteMigrator { + /// Creates a migrator for a SQLite pool. + #[must_use] + pub const fn new(pool: SqlitePool) -> Self { + Self { pool } + } + + /// Returns every SQLite migration bundled with this crate. + #[must_use] + pub fn all_versions() -> &'static [Migration] { + &SQLITE_MIGRATIONS + } + + /// Returns applied main-line versions in ascending order. + /// + /// # Errors + /// + /// Returns [`Error::Database`] when the query fails. + pub async fn existing_versions(&self) -> Result, Error> { + let exists: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM sqlite_schema WHERE type = 'table' AND name = 'river_migration')", + ) + .fetch_one(&self.pool) + .await + .map_err(Error::Database)?; + if !exists { + return Ok(Vec::new()); + } + + let has_line: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pragma_table_info('river_migration') WHERE name = 'line')", + ) + .fetch_one(&self.pool) + .await + .map_err(Error::Database)?; + let rows = if has_line { + sqlx::query("SELECT version FROM river_migration WHERE line = ?1 ORDER BY version") + .bind(MIGRATION_LINE_MAIN) + .fetch_all(&self.pool) + .await + .map_err(Error::Database)? + } else { + sqlx::query("SELECT version FROM river_migration ORDER BY version") + .fetch_all(&self.pool) + .await + .map_err(Error::Database)? + }; + Ok(rows.iter().map(|row| row.get("version")).collect()) + } + + /// Applies up or down migrations with target, step, and dry-run controls. + /// + /// Each migration runs in its own transaction, so a failure leaves the + /// migrations before it applied. + /// + /// # Errors + /// + /// Returns [`Error::UnknownVersion`] when the target version doesn't exist, + /// [`Error::TargetNotSelected`] when a down target isn't applied or is + /// beyond the step limit, [`Error::OtherMigrationLines`] when reverting + /// version 5 would lose other migration lines' records, and [`Error::Database`] when a migration fails. + /// + /// # Cancel safety + /// + /// Each migration and its record in `river_migration` commit together in + /// their own transaction, on a task of their own. Dropping the future + /// stops migrating once the migration in progress finishes: it and + /// every migration before it stay applied, and migrating again + /// continues from there. + pub async fn migrate( + &self, + direction: Direction, + opts: MigrateOpts, + ) -> Result { + validate_target(&SQLITE_MIGRATIONS, opts.target_version, true)?; + let applied = self.existing_versions().await?; + let selected = select_migrations(&SQLITE_MIGRATIONS, direction, opts, &applied)?; + + let mut versions = Vec::with_capacity(selected.len()); + for migration in selected { + let sql = migration_sql(direction, migration).to_owned(); + let mut duration = Duration::ZERO; + if !opts.dry_run { + let started_at = Instant::now(); + // Each migration runs to completion on its own task, so + // dropping this future never abandons one partway. + let migrator = self.clone(); + let task_sql = sql.clone(); + run_to_completion( + async move { migrator.apply(direction, migration, &task_sql).await }, + ) + .await?; + duration = started_at.elapsed(); + } + versions.push(MigrateVersion { + duration, + name: migration.name, + sql, + version: migration.version, + }); + } + Ok(MigrateResult { + direction, + versions, + }) + } + + /// Applies all outstanding up migrations and returns their versions. + /// + /// # Errors + /// + /// Returns [`Error::Database`] when a migration fails. + /// + /// # Cancel safety + /// + /// Each migration and its record in `river_migration` commit together in + /// their own transaction, on a task of their own. Dropping the future + /// stops migrating once the migration in progress finishes: it and + /// every migration before it stay applied, and migrating again + /// continues from there. + pub async fn migrate_up(&self) -> Result, Error> { + Ok(self + .migrate(Direction::Up, MigrateOpts::default()) + .await? + .versions + .into_iter() + .map(|version| version.version) + .collect()) + } + + /// Checks that every migration through an optional target is applied. + /// + /// # Errors + /// + /// Returns [`Error::UnknownVersion`] when the target version doesn't exist and + /// [`Error::Database`] when reading the applied versions fails. + pub async fn validate(&self, target_version: Option) -> Result { + validate_target(&SQLITE_MIGRATIONS, target_version, false)?; + let applied = self.existing_versions().await?; + Ok(validate_migrations( + &SQLITE_MIGRATIONS, + target_version, + &applied, + )) + } + + async fn apply( + &self, + direction: Direction, + migration: Migration, + sql: &str, + ) -> Result<(), Error> { + let mut transaction = self + .pool + .begin_with("BEGIN IMMEDIATE") + .await + .map_err(Error::Database)?; + if direction == Direction::Down && migration.version == 5 { + let has_other_lines: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM river_migration WHERE line <> ?1)", + ) + .bind(MIGRATION_LINE_MAIN) + .fetch_one(&mut *transaction) + .await + .map_err(Error::Database)?; + if has_other_lines { + return Err(Error::OtherMigrationLines { + version: migration.version, + }); + } + } + + sqlx::raw_sql(sqlx::AssertSqlSafe(sql)) + .execute(&mut *transaction) + .await + .map_err(Error::Database)?; + match direction { + Direction::Down if migration.version == 1 => {} + Direction::Down if migration.version <= 5 => { + sqlx::query("DELETE FROM river_migration WHERE version = ?1") + .bind(migration.version) + .execute(&mut *transaction) + .await + .map_err(Error::Database)?; + } + Direction::Down => { + sqlx::query("DELETE FROM river_migration WHERE line = ?1 AND version = ?2") + .bind(MIGRATION_LINE_MAIN) + .bind(migration.version) + .execute(&mut *transaction) + .await + .map_err(Error::Database)?; + } + Direction::Up if migration.version >= 5 => { + sqlx::query("INSERT INTO river_migration (line, version) VALUES (?1, ?2)") + .bind(MIGRATION_LINE_MAIN) + .bind(migration.version) + .execute(&mut *transaction) + .await + .map_err(Error::Database)?; + } + Direction::Up => { + sqlx::query("INSERT INTO river_migration (version) VALUES (?1)") + .bind(migration.version) + .execute(&mut *transaction) + .await + .map_err(Error::Database)?; + } + } + transaction.commit().await.map_err(Error::Database)?; + Ok(()) + } +} + +const fn migration_sql(direction: Direction, migration: Migration) -> &'static str { + match direction { + Direction::Down => migration.down_sql, + Direction::Up => migration.up_sql, + } +} diff --git a/rust/riverqueue-migrate/tests/postgres.rs b/rust/riverqueue-migrate/tests/postgres.rs new file mode 100644 index 000000000..16013ab36 --- /dev/null +++ b/rust/riverqueue-migrate/tests/postgres.rs @@ -0,0 +1,320 @@ +#![cfg(all(feature = "postgres", river_postgres_tests))] + +use riverqueue_migrate::SchemaName; +use riverqueue_migrate::{Direction, MIGRATION_VERSION_LATEST, MigrateOpts, PostgresMigrator}; +use serde_json::Value; +use sqlx::{AssertSqlSafe, PgPool}; + +#[tokio::test] +async fn upgrades_from_every_historical_version() { + let pool = test_pool().await; + + // Unique names keep concurrent runs against one database apart. + let prefix = unique_schema("rust_migrate"); + let reference_schema = format!("{prefix}_reference"); + recreate_schema(&pool, &reference_schema).await; + let reference = PostgresMigrator::new(pool.clone()) + .with_schema(SchemaName::new(reference_schema.clone()).unwrap()); + reference.migrate_up().await.unwrap(); + let expected = schema_snapshot(&pool, &reference_schema).await; + + for version in 1..=MIGRATION_VERSION_LATEST { + let schema = format!("{prefix}_from_{version}"); + recreate_schema(&pool, &schema).await; + let migrator = PostgresMigrator::new(pool.clone()) + .with_schema(SchemaName::new(schema.clone()).unwrap()); + migrator + .migrate( + Direction::Up, + MigrateOpts::new().with_target_version(version), + ) + .await + .unwrap(); + assert_eq!( + migrator.existing_versions().await.unwrap(), + (1..=version).collect::>() + ); + migrator.migrate_up().await.unwrap(); + assert_eq!(schema_snapshot(&pool, &schema).await, expected); + + if version < MIGRATION_VERSION_LATEST { + migrator + .migrate( + Direction::Down, + MigrateOpts::new().with_target_version(version), + ) + .await + .unwrap(); + assert_eq!( + migrator.existing_versions().await.unwrap(), + (1..=version).collect::>() + ); + migrator.migrate_up().await.unwrap(); + assert_eq!(schema_snapshot(&pool, &schema).await, expected); + } + } + + for version in 1..=MIGRATION_VERSION_LATEST { + let schema = format!("{prefix}_from_{version}"); + drop_schema(&pool, &schema).await; + } + drop_schema(&pool, &reference_schema).await; +} + +// Other migration lines share `river_migration` with River's main line. Main +// line operations must neither read nor remove their rows, and reverting +// version 005, which would drop the `line` column and lose them, must fail +// as it does in Go. +#[tokio::test] +async fn main_line_migrations_preserve_other_lines() { + let pool = test_pool().await; + let schema = unique_schema("rust_migrate_lines"); + recreate_schema(&pool, &schema).await; + let schema_name = SchemaName::new(schema.clone()).unwrap(); + let table = schema_name.qualify("river_migration"); + let migrator = PostgresMigrator::new(pool.clone()).with_schema(schema_name); + migrator.migrate_up().await.unwrap(); + sqlx::query(AssertSqlSafe(format!( + "INSERT INTO {table} (line, version) VALUES ('extension', 1), ('extension', 2)" + ))) + .execute(&pool) + .await + .unwrap(); + let extension_versions = || { + let pool = pool.clone(); + let table = table.clone(); + async move { + sqlx::query_scalar::<_, i64>(AssertSqlSafe(format!( + "SELECT version FROM {table} WHERE line = 'extension' ORDER BY version" + ))) + .fetch_all(&pool) + .await + .unwrap() + } + }; + + let all_main = (1..=MIGRATION_VERSION_LATEST).collect::>(); + assert_eq!(migrator.existing_versions().await.unwrap(), all_main); + assert!(migrator.validate(None).await.unwrap().is_valid()); + assert_eq!(migrator.migrate_up().await.unwrap(), Vec::::new()); + + // Down to 005 and back up only touches the main line. + migrator + .migrate(Direction::Down, MigrateOpts::new().with_target_version(5)) + .await + .unwrap(); + assert_eq!( + migrator.existing_versions().await.unwrap(), + (1..=5).collect::>() + ); + assert_eq!(extension_versions().await, [1, 2]); + migrator.migrate_up().await.unwrap(); + assert_eq!(migrator.existing_versions().await.unwrap(), all_main); + assert_eq!(extension_versions().await, [1, 2]); + + // Reverting 005 would lose the other line, so it fails and changes + // nothing. + migrator + .migrate(Direction::Down, MigrateOpts::new().with_target_version(5)) + .await + .unwrap(); + let error = migrator + .migrate(Direction::Down, MigrateOpts::new().with_target_version(4)) + .await + .unwrap_err(); + assert!( + error.to_string().contains("non-main migration lines"), + "{error}" + ); + assert_eq!( + migrator.existing_versions().await.unwrap(), + (1..=5).collect::>() + ); + assert_eq!(extension_versions().await, [1, 2]); + + // Without other lines, 005 reverts. + sqlx::query(AssertSqlSafe(format!( + "DELETE FROM {table} WHERE line = 'extension'" + ))) + .execute(&pool) + .await + .unwrap(); + migrator + .migrate(Direction::Down, MigrateOpts::new().with_target_version(4)) + .await + .unwrap(); + assert_eq!( + migrator.existing_versions().await.unwrap(), + (1..=4).collect::>() + ); + + drop_schema(&pool, &schema).await; +} + +#[tokio::test] +async fn mixed_case_schema_is_detected_as_migrated() { + let pool = test_pool().await; + let schema = unique_schema("RiverMixedCase"); + recreate_schema(&pool, &schema).await; + let migrator = + PostgresMigrator::new(pool.clone()).with_schema(SchemaName::new(schema.clone()).unwrap()); + + migrator.migrate_up().await.unwrap(); + assert_eq!( + migrator.existing_versions().await.unwrap(), + (1..=MIGRATION_VERSION_LATEST).collect::>() + ); + // A second run must see the applied versions instead of re-running 001. + assert_eq!(migrator.migrate_up().await.unwrap(), Vec::::new()); + assert!(migrator.validate(None).await.unwrap().is_valid()); + + drop_schema(&pool, &schema).await; +} + +#[tokio::test] +async fn unknown_versions_are_ignored_and_unapplied_down_targets_fail() { + let pool = test_pool().await; + let schema = unique_schema("rust_migrate_semantics"); + recreate_schema(&pool, &schema).await; + let schema_name = SchemaName::new(schema.clone()).unwrap(); + let migrator = PostgresMigrator::new(pool.clone()).with_schema(schema_name.clone()); + migrator.migrate_up().await.unwrap(); + + // A newer River release recorded a version this crate does not bundle. + sqlx::query(AssertSqlSafe(format!( + "INSERT INTO {} (line, version) VALUES ('main', $1)", + schema_name.qualify("river_migration") + ))) + .bind(MIGRATION_VERSION_LATEST + 1) + .execute(&pool) + .await + .unwrap(); + assert_eq!(migrator.migrate_up().await.unwrap(), Vec::::new()); + assert!(migrator.validate(None).await.unwrap().is_valid()); + let reverted = migrator + .migrate(Direction::Down, MigrateOpts::new()) + .await + .unwrap(); + assert_eq!( + reverted + .versions + .iter() + .map(|version| version.version) + .collect::>(), + vec![MIGRATION_VERSION_LATEST] + ); + + // Migrating down to a version that is not applied errors like Go. + migrator + .migrate(Direction::Down, MigrateOpts::new().with_target_version(3)) + .await + .unwrap(); + let error = migrator + .migrate(Direction::Down, MigrateOpts::new().with_target_version(5)) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("version 5 is not in target list of valid migrations to apply"), + "{error}" + ); + + drop_schema(&pool, &schema).await; +} + +async fn test_pool() -> PgPool { + let database_url = std::env::var("RIVER_RUST_DATABASE_URL") + .expect("RIVER_RUST_DATABASE_URL must point at a disposable test database"); + PgPool::connect(&database_url).await.unwrap() +} + +fn unique_schema(prefix: &str) -> String { + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .subsec_nanos(); + format!("{prefix}_{:x}_{nanos:x}", std::process::id()) +} + +async fn drop_schema(pool: &PgPool, schema: &str) { + let sql = format!("DROP SCHEMA IF EXISTS \"{schema}\" CASCADE"); + sqlx::raw_sql(AssertSqlSafe(sql)) + .execute(pool) + .await + .unwrap(); +} + +async fn recreate_schema(pool: &PgPool, schema: &str) { + drop_schema(pool, schema).await; + let sql = format!("CREATE SCHEMA \"{schema}\""); + sqlx::raw_sql(AssertSqlSafe(sql)) + .execute(pool) + .await + .unwrap(); +} + +async fn schema_snapshot(pool: &PgPool, schema: &str) -> Value { + let mut snapshot = sqlx::query_scalar::<_, Value>( + r" + WITH objects AS ( + SELECT 'column' AS object_kind, + table_name || '.' || column_name || ':' || data_type || ':' || udt_name || ':' || is_nullable || ':' || coalesce(column_default, '') AS definition + FROM information_schema.columns + WHERE table_schema = $1 + UNION ALL + SELECT 'constraint', c.relname || ':' || pg_get_constraintdef(con.oid, true) + FROM pg_constraint AS con + JOIN pg_class AS c ON c.oid = con.conrelid + WHERE con.connamespace = $1::regnamespace + UNION ALL + SELECT 'function', proname || ':' || pg_get_functiondef(oid) + FROM pg_proc + WHERE pronamespace = $1::regnamespace + UNION ALL + SELECT 'index', tablename || ':' || regexp_replace(indexdef, ' ON [^ ]+\\.', ' ON .') + FROM pg_indexes + WHERE schemaname = $1 + UNION ALL + SELECT 'trigger', event_object_table || ':' || trigger_name || ':' || action_timing || ':' || event_manipulation || ':' || action_statement + FROM information_schema.triggers + WHERE trigger_schema = $1 + UNION ALL + SELECT 'type', t.typname || ':' || string_agg(e.enumlabel, ',' ORDER BY e.enumsortorder) + FROM pg_type AS t + JOIN pg_enum AS e ON e.enumtypid = t.oid + WHERE t.typnamespace = $1::regnamespace + GROUP BY t.typname + ) + SELECT coalesce(jsonb_agg(jsonb_build_array(object_kind, definition) ORDER BY object_kind, definition), '[]'::jsonb) + FROM objects + ", + ) + .bind(schema) + .fetch_one(pool) + .await + .unwrap(); + normalize_schema_name(&mut snapshot, schema); + snapshot +} + +fn normalize_schema_name(value: &mut Value, schema: &str) { + match value { + Value::Array(values) => { + for value in values { + normalize_schema_name(value, schema); + } + } + Value::Object(values) => { + for value in values.values_mut() { + normalize_schema_name(value, schema); + } + } + Value::String(string) => { + *string = string + .replace(&format!("\"{schema}\"."), ".") + .replace(&format!("{schema}."), "."); + } + Value::Bool(_) | Value::Null | Value::Number(_) => {} + } +} diff --git a/rust/riverqueue-migrate/tests/sqlite.rs b/rust/riverqueue-migrate/tests/sqlite.rs new file mode 100644 index 000000000..88f2d5c5f --- /dev/null +++ b/rust/riverqueue-migrate/tests/sqlite.rs @@ -0,0 +1,489 @@ +use std::time::Duration; + +use riverqueue_migrate::{ + Direction, Error, MIGRATION_LINE_MAIN, MIGRATION_VERSION_LATEST, MigrateOpts, + SQLITE_MIGRATIONS, SqliteMigrator, +}; +use sqlx::{Row, SqlitePool, sqlite::SqlitePoolOptions}; + +#[tokio::test] +async fn all_versions_options_and_validation() { + let pool = sqlite_pool().await; + let migrator = SqliteMigrator::new(pool.clone()); + + assert_eq!( + SqliteMigrator::all_versions() + .iter() + .map(|migration| migration.version) + .collect::>(), + (1..=MIGRATION_VERSION_LATEST).collect::>() + ); + assert_eq!( + SqliteMigrator::all_versions().len(), + SQLITE_MIGRATIONS.len() + ); + assert_eq!( + migrator.existing_versions().await.unwrap(), + Vec::::new() + ); + + let dry_run = migrator + .migrate( + Direction::Up, + MigrateOpts::new().with_dry_run(true).with_target_version(3), + ) + .await + .unwrap(); + assert_eq!( + dry_run + .versions + .iter() + .map(|migration| migration.version) + .collect::>(), + vec![1, 2, 3] + ); + assert!( + dry_run + .versions + .iter() + .all(|migration| migration.duration == Duration::ZERO && !migration.sql.is_empty()) + ); + assert_eq!( + migrator.existing_versions().await.unwrap(), + Vec::::new() + ); + + let validation = migrator.validate(Some(3)).await.unwrap(); + assert!(!validation.is_valid()); + assert_eq!(validation.unapplied, [1, 2, 3]); + assert_eq!(validation.to_string(), "unapplied migrations: [1, 2, 3]"); + + migrator + .migrate(Direction::Up, MigrateOpts::new().with_max_steps(2)) + .await + .unwrap(); + assert_eq!(migrator.existing_versions().await.unwrap(), vec![1, 2]); + assert!(migrator.validate(Some(2)).await.unwrap().is_valid()); + + let error = migrator + .migrate( + Direction::Up, + MigrateOpts::new().with_target_version(MIGRATION_VERSION_LATEST + 1), + ) + .await + .unwrap_err(); + assert!(matches!( + error, + Error::UnknownVersion { version } if version == MIGRATION_VERSION_LATEST + 1 + )); + + pool.close().await; +} + +#[tokio::test] +async fn downgrade_preserves_non_main_migration_lines() { + let pool = sqlite_pool().await; + let migrator = SqliteMigrator::new(pool.clone()); + migrate_to(&migrator, 5).await; + sqlx::query("INSERT INTO river_migration (line, version) VALUES ('extension', 1)") + .execute(&pool) + .await + .unwrap(); + + assert_eq!( + migrator.existing_versions().await.unwrap(), + (1..=5).collect::>() + ); + let error = migrator + .migrate(Direction::Down, MigrateOpts::new().with_target_version(4)) + .await + .unwrap_err(); + assert!(matches!(error, Error::OtherMigrationLines { version: 5 })); + assert_eq!( + error.to_string(), + "found non-main migration lines; version 005 is irreversible without losing migration information" + ); + assert_eq!( + migrator.existing_versions().await.unwrap(), + (1..=5).collect::>() + ); + let extension_version: i64 = + sqlx::query_scalar("SELECT version FROM river_migration WHERE line = 'extension'") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(extension_version, 1); + + sqlx::query("DELETE FROM river_migration WHERE line = 'extension'") + .execute(&pool) + .await + .unwrap(); + migrator + .migrate(Direction::Down, MigrateOpts::new().with_target_version(4)) + .await + .unwrap(); + assert_eq!( + migrator.existing_versions().await.unwrap(), + (1..=4).collect::>() + ); + + pool.close().await; +} + +#[tokio::test] +async fn latest_schema_and_json_survive_version_seven_round_trip() { + let pool = sqlite_pool().await; + let migrator = SqliteMigrator::new(pool.clone()); + migrate_to(&migrator, 6).await; + + let job_id: i64 = sqlx::query_scalar( + "INSERT INTO river_job (args, kind, max_attempts, metadata, tags) \ + VALUES (json(?1), 'sqlite_migration_test', 9, json(?2), json(?3)) RETURNING id", + ) + .bind(r#"{"message":"hello"}"#) + .bind(r#"{"source":"test"}"#) + .bind(r#"["one","two"]"#) + .fetch_one(&pool) + .await + .unwrap(); + + migrator.migrate_up().await.unwrap(); + assert_eq!( + migrator.existing_versions().await.unwrap(), + (1..=MIGRATION_VERSION_LATEST).collect::>() + ); + assert_eq!( + column_default(&pool, "river_job", "max_attempts") + .await + .as_deref(), + Some("25") + ); + assert_eq!( + column_default(&pool, "river_queue", "updated_at") + .await + .as_deref(), + Some("CURRENT_TIMESTAMP") + ); + let row = sqlx::query( + "SELECT typeof(args) AS args_type, json(args) AS args, json(metadata) AS metadata, \ + json(tags) AS tags FROM river_job WHERE id = ?1", + ) + .bind(job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(row.get::("args_type"), "blob"); + assert_eq!(row.get::("args"), r#"{"message":"hello"}"#); + assert_eq!(row.get::("metadata"), r#"{"source":"test"}"#); + assert_eq!(row.get::("tags"), r#"["one","two"]"#); + + migrator + .migrate(Direction::Down, MigrateOpts::new().with_target_version(6)) + .await + .unwrap(); + assert_eq!( + column_default(&pool, "river_job", "max_attempts").await, + None + ); + let row = sqlx::query( + "SELECT typeof(args) AS args_type, json(args) AS args, json(metadata) AS metadata, \ + json(tags) AS tags FROM river_job WHERE id = ?1", + ) + .bind(job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(row.get::("args_type"), "text"); + assert_eq!(row.get::("args"), r#"{"message":"hello"}"#); + assert_eq!(row.get::("metadata"), r#"{"source":"test"}"#); + assert_eq!(row.get::("tags"), r#"["one","two"]"#); + + migrator.migrate_up().await.unwrap(); + let args_type: String = sqlx::query_scalar("SELECT typeof(args) FROM river_job WHERE id = ?1") + .bind(job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(args_type, "blob"); + + pool.close().await; +} + +#[tokio::test] +async fn job_ids_are_not_reused_after_version_eight() { + let pool = sqlite_pool().await; + let migrator = SqliteMigrator::new(pool.clone()); + migrate_to(&migrator, 7).await; + let before_upgrade = insert_job(&pool).await; + + migrator.migrate_up().await.unwrap(); + let kept: i64 = sqlx::query_scalar("SELECT id FROM river_job WHERE id = ?1") + .bind(before_upgrade) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(kept, before_upgrade); + + sqlx::query("DELETE FROM river_job WHERE id = ?1") + .bind(before_upgrade) + .execute(&pool) + .await + .unwrap(); + let after_delete = insert_job(&pool).await; + assert!(after_delete > before_upgrade); + + pool.close().await; +} + +#[tokio::test] +async fn version_eight_refuses_to_rebuild_extended_job_tables() { + for (object, sql) in [ + ( + "river_job_workflow_scheduling", + "CREATE INDEX river_job_workflow_scheduling ON river_job (state)", + ), + ( + "river_job_sequence", + "CREATE TABLE river_job_sequence (id integer PRIMARY KEY, key text)", + ), + ( + "river_workflow", + "CREATE TABLE river_workflow (id text PRIMARY KEY)", + ), + ] { + for direction in [Direction::Up, Direction::Down] { + let version = if direction == Direction::Up { 7 } else { 8 }; + let pool = sqlite_pool().await; + let migrator = SqliteMigrator::new(pool.clone()); + migrate_to(&migrator, version).await; + let job_id = insert_job(&pool).await; + sqlx::query(sql).execute(&pool).await.unwrap(); + sqlx::query("ALTER TABLE river_job ADD COLUMN partition_key text") + .execute(&pool) + .await + .unwrap(); + + let error = migrator + .migrate(direction, MigrateOpts::new().with_max_steps(1)) + .await + .unwrap_err(); + assert!( + error + .to_string() + .contains("River SQLite migration 008 cannot run"), + "{object} {direction:?}: {error}" + ); + + assert_eq!( + migrator.existing_versions().await.unwrap(), + (1..=version).collect::>(), + "{object} {direction:?}" + ); + let partition_key_columns: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pragma_table_info('river_job') WHERE name = 'partition_key'", + ) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(partition_key_columns, 1, "{object} {direction:?}"); + let kept: i64 = sqlx::query_scalar("SELECT id FROM river_job WHERE id = ?1") + .bind(job_id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(kept, job_id, "{object} {direction:?}"); + pool.close().await; + } + } +} + +#[tokio::test] +async fn migrates_up_from_every_historical_version() { + let expected = schema_at(MIGRATION_VERSION_LATEST).await; + + for version in 1..=MIGRATION_VERSION_LATEST { + let pool = sqlite_pool().await; + let migrator = SqliteMigrator::new(pool.clone()); + migrate_to(&migrator, version).await; + assert_eq!( + migrator.existing_versions().await.unwrap(), + (1..=version).collect::>() + ); + + migrator.migrate_up().await.unwrap(); + assert_eq!(schema_snapshot(&pool).await, expected, "version {version}"); + assert!(migrator.validate(None).await.unwrap().is_valid()); + pool.close().await; + } +} + +#[tokio::test] +async fn migrates_down_to_every_historical_version_and_empty() { + for target in 1..MIGRATION_VERSION_LATEST { + let expected = schema_structure_at(target).await; + let pool = sqlite_pool().await; + let migrator = SqliteMigrator::new(pool.clone()); + migrator.migrate_up().await.unwrap(); + migrator + .migrate( + Direction::Down, + MigrateOpts::new().with_target_version(target), + ) + .await + .unwrap(); + assert_eq!( + migrator.existing_versions().await.unwrap(), + (1..=target).collect::>() + ); + assert_eq!(schema_structure(&pool).await, expected, "version {target}"); + pool.close().await; + } + + let pool = sqlite_pool().await; + let migrator = SqliteMigrator::new(pool.clone()); + migrator.migrate_up().await.unwrap(); + let result = migrator + .migrate(Direction::Down, MigrateOpts::new().with_target_version(-1)) + .await + .unwrap(); + assert_eq!( + result + .versions + .iter() + .map(|migration| migration.version) + .collect::>(), + (1..=MIGRATION_VERSION_LATEST).rev().collect::>() + ); + assert_eq!( + migrator.existing_versions().await.unwrap(), + Vec::::new() + ); + assert_eq!(schema_snapshot(&pool).await, []); + pool.close().await; +} + +async fn column_default(pool: &SqlitePool, table: &str, column: &str) -> Option { + sqlx::query("SELECT dflt_value FROM pragma_table_info(?1) WHERE name = ?2") + .bind(table) + .bind(column) + .fetch_one(pool) + .await + .unwrap() + .get("dflt_value") +} + +async fn migrate_to(migrator: &SqliteMigrator, version: i64) { + migrator + .migrate( + Direction::Up, + MigrateOpts::new().with_target_version(version), + ) + .await + .unwrap(); +} + +async fn schema_at(version: i64) -> Vec<(String, String, String, String)> { + let pool = sqlite_pool().await; + let migrator = SqliteMigrator::new(pool.clone()); + migrate_to(&migrator, version).await; + let snapshot = schema_snapshot(&pool).await; + pool.close().await; + snapshot +} + +async fn schema_structure_at(version: i64) -> Vec<(String, String, String, String)> { + let pool = sqlite_pool().await; + let migrator = SqliteMigrator::new(pool.clone()); + migrate_to(&migrator, version).await; + let structure = schema_structure(&pool).await; + pool.close().await; + structure +} + +/// Like `schema_snapshot`, but describes each table by its columns in name +/// order instead of its `CREATE TABLE` text. Version 8's down migration +/// rebuilds `river_job` with its columns in declaration order, while +/// migrating up to version 7 leaves `max_attempts` where `ALTER TABLE` +/// appended it. +async fn schema_structure(pool: &SqlitePool) -> Vec<(String, String, String, String)> { + let mut structure = schema_snapshot(pool).await; + for (object_type, name, _, sql) in &mut structure { + if object_type != "table" { + continue; + } + let columns = sqlx::query( + "SELECT name, type, \"notnull\", coalesce(dflt_value, '') AS dflt_value, pk \ + FROM pragma_table_xinfo(?1) ORDER BY name", + ) + .bind(&*name) + .fetch_all(pool) + .await + .unwrap() + .into_iter() + .map(|row| { + format!( + "{} {} notnull={} default={} pk={}", + row.get::("name"), + row.get::("type"), + row.get::("notnull"), + row.get::("dflt_value"), + row.get::("pk"), + ) + }) + .collect::>(); + *sql = columns.join("\n"); + } + structure +} + +async fn insert_job(pool: &SqlitePool) -> i64 { + sqlx::query_scalar("INSERT INTO river_job (kind) VALUES ('sqlite_migration_test') RETURNING id") + .fetch_one(pool) + .await + .unwrap() +} + +async fn schema_snapshot(pool: &SqlitePool) -> Vec<(String, String, String, String)> { + sqlx::query( + "SELECT type, name, tbl_name, coalesce(sql, '') AS sql \ + FROM sqlite_schema WHERE name LIKE 'river_%' ORDER BY type, name", + ) + .fetch_all(pool) + .await + .unwrap() + .into_iter() + .map(|row| { + ( + row.get("type"), + row.get("name"), + row.get("tbl_name"), + row.get("sql"), + ) + }) + .collect() +} + +async fn sqlite_pool() -> SqlitePool { + SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap() +} + +#[test] +fn sqlite_migrations_match_main_line_metadata() { + assert_eq!( + SQLITE_MIGRATIONS.len(), + usize::try_from(MIGRATION_VERSION_LATEST).unwrap() + ); + for (index, migration) in SQLITE_MIGRATIONS.iter().enumerate() { + assert_eq!(migration.version, i64::try_from(index).unwrap() + 1); + assert_ne!(migration.name, ""); + assert_ne!(migration.up_sql, ""); + assert_ne!(migration.down_sql, ""); + assert!(!migration.up_sql.contains("CREATE TYPE")); + assert!(!migration.up_sql.contains("LANGUAGE plpgsql")); + } + assert_eq!(MIGRATION_LINE_MAIN, "main"); +} diff --git a/rust/riverqueue-test/Cargo.toml b/rust/riverqueue-test/Cargo.toml new file mode 100644 index 000000000..ce72e2415 --- /dev/null +++ b/rust/riverqueue-test/Cargo.toml @@ -0,0 +1,35 @@ +[package] +name = "riverqueue-test" +description = "Test helpers for River's Rust client" +keywords = ["background", "jobs", "queue", "testing"] +categories = ["asynchronous", "database", "development-tools::testing"] +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +readme = "README.md" +repository.workspace = true +homepage.workspace = true + +[package.metadata.docs.rs] +all-features = true + +[features] +default = ["postgres"] +postgres = ["riverqueue/postgres"] +sqlite = ["riverqueue/sqlite"] + +[dependencies] +chrono.workspace = true +riverqueue = { path = "../riverqueue", version = "=0.50.0-alpha.1", default-features = false } +serde_json.workspace = true +tokio-util.workspace = true + +[dev-dependencies] +riverqueue-migrate = { path = "../riverqueue-migrate", version = "=0.50.0-alpha.1", default-features = false, features = ["sqlite"] } +serde.workspace = true +sqlx.workspace = true +tokio.workspace = true + +[lints] +workspace = true diff --git a/rust/riverqueue-test/LICENSE b/rust/riverqueue-test/LICENSE new file mode 120000 index 000000000..30cff7403 --- /dev/null +++ b/rust/riverqueue-test/LICENSE @@ -0,0 +1 @@ +../../LICENSE \ No newline at end of file diff --git a/rust/riverqueue-test/README.md b/rust/riverqueue-test/README.md new file mode 100644 index 000000000..be3922202 --- /dev/null +++ b/rust/riverqueue-test/README.md @@ -0,0 +1,109 @@ +# riverqueue-test + +Test helpers for applications using River's Rust client: assertions about +inserted jobs, and ways to run a worker once, with or without a database. + +## Asserting on inserted jobs + +`require_inserted`, `require_many_inserted`, and `require_not_inserted` check +the jobs a test's code inserted. Each +lists jobs of the expected kinds in insertion order and panics with a +descriptive message when the expectation isn't met, failing the test. +The `_with` variants take `RequireInsertedOpts`, which adds expected +properties such as the queue, priority, state, or tags. The `_tx` variants +read through an open transaction, to test code that enqueues jobs +transactionally before it commits. + +```rust,no_run +use riverqueue::{Client, JobArgs, JobState}; +use riverqueue_test::{RequireInsertedOpts, require_inserted_with}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "send_welcome_email")] +struct SendWelcomeEmail { + user_id: i64, +} + +async fn sign_up(client: &Client, user_id: i64) -> Result<(), riverqueue::Error> { + client.insert(SendWelcomeEmail { user_id }).await?; + Ok(()) +} + +async fn test_sign_up(client: &Client) { + sign_up(client, 42).await.unwrap(); + + let job = require_inserted_with::( + client, + &RequireInsertedOpts::new().with_state(JobState::Available), + ) + .await; + assert_eq!(job.args.user_id, 42); +} +``` + +## Running a worker once + +`TestJobBuilder` constructs a realistic `Job` from the argument type's +insertion defaults and lets a test override the persisted ID, attempt, state, +and metadata. `work_once` invokes a typed worker with a detached +`WorkContext`, preserving its concrete error and capturing an immutable +snapshot of recorded output and metadata updates. + +```rust,no_run +use riverqueue::{Job, JobArgs, WorkContext, WorkOutcome, Worker}; +use riverqueue_test::{TestJobBuilder, work_once}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "thumbnail")] +struct Thumbnail { + image_id: i64, +} + +struct ThumbnailWorker; + +impl Worker for ThumbnailWorker { + type Error = serde_json::Error; + + async fn work( + &self, + context: WorkContext, + job: Job, + ) -> Result { + context.record_output(serde_json::json!({"image_id": job.args.image_id}))?; + Ok(WorkOutcome::Complete) + } +} + +#[tokio::test] +async fn thumbnail_records_its_image() { + let job = TestJobBuilder::new(Thumbnail { image_id: 42 }) + .id(100) + .build() + .unwrap(); + let worked = work_once(&ThumbnailWorker, job).await; + + assert_eq!(worked.result.as_ref().unwrap(), &WorkOutcome::Complete); + assert_eq!(worked.output(), Some(&serde_json::json!({"image_id": 42}))); +} +``` + +`work_once` restores and finalizes resumable state, including failures that the +worker catches. Its result distinguishes `TestWorkError::Worker` from +`TestWorkError::Resumable` while preserving the original error source. Pass +`metadata_updates` into the next job's metadata to test a resumed attempt. + +The helper does not run client hooks, middleware, database transactions, +retries, or completion persistence. + +## Running a worker with a client + +`work_with_client` is the database-backed counterpart. It inserts the job with a client, claims it the way a +fetch does, and runs the worker with that client in its `WorkContext`, so a +worker that inserts follow-up jobs through `context.client()` or completes +its job in its own transaction with `context.job_complete_tx` runs as it +would in production. The client doesn't need to be started. River doesn't +record the worker's result, so the job stays running unless the worker +completed it itself, and it also stays running if the test drops the +future partway, for example on a timeout. diff --git a/rust/riverqueue-test/src/lib.rs b/rust/riverqueue-test/src/lib.rs new file mode 100644 index 000000000..836786796 --- /dev/null +++ b/rust/riverqueue-test/src/lib.rs @@ -0,0 +1,510 @@ +#![doc = include_str!("../README.md")] +#![cfg_attr(docsrs, feature(doc_cfg))] +#![forbid(unsafe_code)] +#![warn(missing_docs)] + +use chrono::Utc; +use riverqueue::{ + __private, Client, Error, Job, JobArgs, JobRow, JobState, MAX_ATTEMPTS_DEFAULT, + PRIORITY_DEFAULT, QUEUE_DEFAULT, WorkContext, WorkError, WorkOutcome, Worker, +}; +use serde_json::{Map, Value}; + +mod require; + +pub use require::{ + ExpectedJob, RequireInsertedOpts, require_inserted, require_inserted_tx, + require_inserted_tx_with, require_inserted_with, require_many_inserted, + require_many_inserted_tx, require_not_inserted, require_not_inserted_tx, + require_not_inserted_tx_with, require_not_inserted_with, +}; + +/// Builder for a realistic persisted job value usable in worker unit tests. +#[derive(Debug)] +pub struct TestJobBuilder { + args: A, + attempt: i16, + id: i64, + metadata: Map, + state: JobState, +} + +impl TestJobBuilder { + /// Starts a test job using River's defaults. + #[must_use] + pub fn new(args: A) -> Self { + Self { + args, + attempt: 1, + id: 1, + metadata: Map::new(), + state: JobState::Running, + } + } + + /// Sets the current attempt. + #[must_use] + pub const fn attempt(mut self, attempt: i16) -> Self { + self.attempt = attempt; + self + } + + /// Builds the typed test job. + /// + /// # Errors + /// + /// Returns an error when the arguments can't be encoded as JSON. + pub fn build(self) -> Result, Error> { + let now = Utc::now(); + let encoded_args = riverqueue::encoding::encode_args(&self.args)?; + let defaults = A::default_insert_opts(); + let max_attempts = defaults.max_attempts().unwrap_or(MAX_ATTEMPTS_DEFAULT); + let priority = defaults.priority().unwrap_or(PRIORITY_DEFAULT); + let queue = defaults.queue().unwrap_or(QUEUE_DEFAULT).to_owned(); + let mut row = JobRow::new(self.id, A::KIND, encoded_args, now); + row.attempt = self.attempt; + row.attempted_at = Some(now); + row.attempted_by = vec!["riverqueue-test".to_owned()]; + row.max_attempts = max_attempts; + row.metadata = self.metadata.into(); + row.priority = priority; + row.queue = queue; + row.state = self.state; + Ok(Job::new(self.args, row)) + } + + /// Sets the database ID. + #[must_use] + pub const fn id(mut self, id: i64) -> Self { + self.id = id; + self + } + + /// Replaces metadata. + #[must_use] + pub fn metadata(mut self, metadata: Map) -> Self { + self.metadata = metadata; + self + } + + /// Sets the persisted state. + #[must_use] + pub const fn state(mut self, state: JobState) -> Self { + self.state = state; + self + } +} + +/// Result of running one worker directly in a unit test. +#[derive(Debug)] +#[non_exhaustive] +pub struct TestWorkResult { + /// Context used for the invocation, including output and metadata updates. + pub context: WorkContext, + metadata_updates: Map, + /// Worker outcome, typed worker error, or runtime checkpoint error. + pub result: Result>, +} + +/// Failure from a direct worker invocation or River's resumable coordinator. +#[derive(Debug)] +#[non_exhaustive] +pub enum TestWorkError { + /// Invalid checkpoint metadata or a step error the worker suppresses. + Resumable(WorkError), + /// Original, unerased error returned by the worker. + Worker(E), +} + +impl std::fmt::Display for TestWorkError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Resumable(error) => error.fmt(formatter), + Self::Worker(error) => error.fmt(formatter), + } + } +} + +/// Transparent, like [`WorkError`]: displays and reports the source of the +/// error it holds. +impl std::error::Error for TestWorkError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + Self::Resumable(error) => error.source(), + Self::Worker(error) => error.source(), + } + } +} + +impl TestWorkResult { + /// Returns a snapshot of metadata recorded by the worker. + #[must_use] + pub const fn metadata_updates(&self) -> &Map { + &self.metadata_updates + } + + /// Returns output recorded by the worker, if any. + #[must_use] + pub fn output(&self) -> Option<&Value> { + self.metadata_updates.get(riverqueue::METADATA_KEY_OUTPUT) + } +} + +/// Runs a typed worker once without a database or background runtime. +/// +/// Initializes resumable steps from the job's metadata and records checkpoints +/// on failure, including step errors caught by the worker. This helper does not +/// simulate queue scheduling, timeouts, middleware, or database transactions, +/// and the worker's [`WorkContext::client`] is `None`; use +/// [`work_with_client`] for a worker that enqueues jobs or completes its job +/// in a transaction. +pub async fn work_once(worker: &W, job: Job) -> TestWorkResult +where + A: JobArgs, + W: Worker, +{ + let context = __private::work_context_for_job(&job.row); + run_worker(worker, context, job).await +} + +/// Inserts `args` with `client` and works the job once with `worker`. +/// +/// The job is inserted with its type's default options and claimed for +/// `client` the way a fetch claims it, so it is `running` with its first +/// attempt when the worker starts, and the worker's [`WorkContext::client`] +/// is `client`. A worker can therefore insert follow-up jobs through the +/// client, or complete its job in its own transaction with +/// [`WorkContext::job_complete_tx`], exactly as it would in production. +/// `client` doesn't need to be started, but it must be able to insert jobs of +/// this kind. +/// +/// River doesn't record the worker's result: the job stays `running` unless +/// the worker completed it transactionally. Returns the job as claimed, and +/// the worker's result as [`work_once`] reports it. +/// +/// # Cancel safety +/// +/// Dropping the returned future after the job is claimed, for example when +/// a test times out, drops the worker's future too and leaves the job +/// `running` in the database with this client's attempt. +/// +/// # Errors +/// +/// Returns the error from inserting or claiming the job. +pub async fn work_with_client( + client: &Client, + worker: &W, + args: A, +) -> Result<(JobRow, TestWorkResult), Error> +where + A: JobArgs + Clone, + W: Worker, +{ + let inserted = client.insert(args.clone()).await?; + let (row, context) = __private::claim_job_for_test(client, inserted.id()).await?; + let result = run_worker(worker, context, Job::new(args, row.clone())).await; + Ok((row, result)) +} + +async fn run_worker(worker: &W, context: WorkContext, job: Job) -> TestWorkResult +where + A: JobArgs, + W: Worker, +{ + let mut result = match __private::work_context_resumable_validate(&context) { + Ok(()) => worker + .work(context.clone(), job) + .await + .map_err(TestWorkError::Worker), + Err(error) => Err(TestWorkError::Resumable(error)), + }; + if let Some(error) = __private::work_context_resumable_finish(&context, result.is_err()) + && result.is_ok() + { + result = Err(TestWorkError::Resumable(error)); + } + let metadata_updates = __private::work_context_metadata_updates(&context); + TestWorkResult { + context, + metadata_updates, + result, + } +} + +#[cfg(test)] +mod tests { + use std::{ + convert::Infallible, + sync::atomic::{AtomicUsize, Ordering}, + time::Duration, + }; + + use riverqueue::{InsertOpts, WorkContext}; + use serde::{Deserialize, Serialize}; + + use super::*; + + struct ResumableWorker { + calls: AtomicUsize, + } + + impl Worker for ResumableWorker { + type Error = Infallible; + + async fn work( + &self, + context: WorkContext, + job: Job, + ) -> Result { + let _ = context + .resumable_step("first", || async { + self.calls.fetch_add(1, Ordering::Relaxed); + Ok::<_, std::io::Error>(()) + }) + .await; + let _ = context + .resumable_step("second", || async { + if job.row.attempt == 1 { + Err(std::io::Error::other("try again")) + } else { + Ok(()) + } + }) + .await; + Ok(WorkOutcome::Complete) + } + } + + #[tokio::test] + async fn work_once_checkpoints_suppressed_step_errors_and_resumes() { + let worker = ResumableWorker { + calls: AtomicUsize::new(0), + }; + let args = TestArgs { + message: "resume".to_owned(), + }; + let first = work_once(&worker, TestJobBuilder::new(args.clone()).build().unwrap()).await; + assert!(matches!(first.result, Err(TestWorkError::Resumable(_)))); + assert_eq!( + first.metadata_updates()[riverqueue::METADATA_KEY_RESUMABLE_STEP], + "first" + ); + let second = work_once( + &worker, + TestJobBuilder::new(args) + .attempt(2) + .metadata(first.metadata_updates().clone()) + .build() + .unwrap(), + ) + .await; + assert!(matches!(second.result, Ok(WorkOutcome::Complete))); + assert_eq!(worker.calls.load(Ordering::Relaxed), 1); + } + + #[tokio::test] + async fn work_once_rejects_invalid_checkpoint_before_work() { + let worker = ResumableWorker { + calls: AtomicUsize::new(0), + }; + let job = TestJobBuilder::new(TestArgs { + message: "invalid".to_owned(), + }) + .metadata( + serde_json::json!({ "river:resumable_cursor": [] }) + .as_object() + .unwrap() + .clone(), + ) + .build() + .unwrap(); + let outcome = work_once(&worker, job).await; + assert!(matches!(outcome.result, Err(TestWorkError::Resumable(_)))); + assert_eq!(worker.calls.load(Ordering::Relaxed), 0); + } + + #[derive(Clone, Debug, Deserialize, Serialize)] + struct TestArgs { + message: String, + } + + impl JobArgs for TestArgs { + const KIND: &'static str = "riverqueue_test_helper"; + + fn default_insert_opts() -> InsertOpts { + InsertOpts::default() + .with_max_attempts(7) + .with_priority(3) + .with_queue("testing") + } + } + + static DEFAULT_CALLS: AtomicUsize = AtomicUsize::new(0); + + #[derive(Debug, Deserialize, Serialize)] + struct DefaultsOnceArgs {} + + impl JobArgs for DefaultsOnceArgs { + const KIND: &'static str = "riverqueue_test_defaults_once"; + + fn default_insert_opts() -> InsertOpts { + DEFAULT_CALLS.fetch_add(1, Ordering::Relaxed); + InsertOpts::default() + } + } + + struct TestWorker; + + impl Worker for TestWorker { + type Error = Infallible; + + fn work( + &self, + context: WorkContext, + job: Job, + ) -> impl Future> + Send { + assert_eq!(job.args.message, "work once"); + assert_eq!(job.row.id, 42); + context + .record_output(serde_json::json!({"worked": true})) + .unwrap(); + context.metadata_set("worker_metadata", "set").unwrap(); + std::future::ready(Ok(WorkOutcome::Snooze(Duration::from_secs(30)))) + } + } + + #[test] + fn test_job_builder_applies_overrides_and_argument_defaults() { + let metadata = Map::from_iter([("test".to_owned(), serde_json::json!(true))]); + let job = TestJobBuilder::new(TestArgs { + message: "builder".to_owned(), + }) + .attempt(4) + .id(99) + .metadata(metadata) + .state(JobState::Retryable) + .build() + .unwrap(); + + assert_eq!(job.args.message, "builder"); + assert_eq!(job.row.attempt, 4); + assert_eq!(job.row.attempted_by, ["riverqueue-test"]); + assert_eq!(job.row.encoded_args.get(), r#"{"message":"builder"}"#); + assert_eq!(job.row.id, 99); + assert_eq!(job.row.kind, TestArgs::KIND); + assert_eq!(job.row.max_attempts, 7); + assert_eq!(job.row.metadata.get::("test").unwrap(), Some(true)); + assert_eq!(job.row.priority, 3); + assert_eq!(job.row.queue, "testing"); + assert_eq!(job.row.state, JobState::Retryable); + assert!(job.row.attempted_at.is_some()); + assert!(job.row.finalized_at.is_none()); + } + + #[test] + fn test_job_builder_evaluates_argument_defaults_once() { + DEFAULT_CALLS.store(0, Ordering::Relaxed); + + TestJobBuilder::new(DefaultsOnceArgs {}).build().unwrap(); + + assert_eq!(DEFAULT_CALLS.load(Ordering::Relaxed), 1); + } + + #[tokio::test] + async fn work_once_runs_worker_with_detached_context() { + let job = TestJobBuilder::new(TestArgs { + message: "work once".to_owned(), + }) + .id(42) + .build() + .unwrap(); + + let worked = work_once(&TestWorker, job).await; + + assert_eq!(worked.output(), Some(&serde_json::json!({"worked": true}))); + assert_eq!( + worked.metadata_updates()["worker_metadata"], + serde_json::json!("set") + ); + assert_eq!( + worked.result.unwrap(), + WorkOutcome::Snooze(Duration::from_secs(30)) + ); + assert!(worked.context.client().is_none()); + assert!(!worked.context.cancellation_token().is_cancelled()); + } + + /// Inserts a follow-up job through its context's client and completes + /// its own job in a transaction. + #[cfg(feature = "sqlite")] + struct ChainingWorker { + pool: sqlx::SqlitePool, + } + + #[cfg(feature = "sqlite")] + impl Worker for ChainingWorker { + type Error = Error; + + async fn work( + &self, + context: WorkContext, + job: Job, + ) -> Result { + let client = context.client().expect("worked with a client"); + client + .insert(TestArgs { + message: format!("after {}", job.args.message), + }) + .await?; + let mut transaction = riverqueue::database::begin_sqlite_write(&self.pool).await?; + context.job_complete_tx(&mut transaction).await?; + transaction.commit().await?; + Ok(WorkOutcome::Complete) + } + } + + #[cfg(feature = "sqlite")] + #[tokio::test] + async fn work_with_client_supports_follow_ups_and_transactional_completion() { + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + riverqueue_migrate::SqliteMigrator::new(pool.clone()) + .migrate_up() + .await + .unwrap(); + let client = Client::builder(pool.clone()).build().unwrap(); + let worker = ChainingWorker { pool: pool.clone() }; + + let (claimed, attempt) = work_with_client( + &client, + &worker, + TestArgs { + message: "first".to_owned(), + }, + ) + .await + .unwrap(); + + assert_eq!(claimed.state, JobState::Running); + assert_eq!(claimed.attempt, 1); + assert!(matches!(attempt.result, Ok(WorkOutcome::Complete))); + assert_eq!( + client.jobs().get(claimed.id).await.unwrap().state, + JobState::Completed + ); + let jobs = client + .jobs() + .list(riverqueue::JobListParams::default()) + .await + .unwrap() + .jobs; + assert_eq!(jobs.len(), 2); + assert!( + jobs.iter() + .any(|row| row.encoded_args.get() == r#"{"message":"after first"}"#) + ); + } +} diff --git a/rust/riverqueue-test/src/require.rs b/rust/riverqueue-test/src/require.rs new file mode 100644 index 000000000..fa43ebdfc --- /dev/null +++ b/rust/riverqueue-test/src/require.rs @@ -0,0 +1,661 @@ +//! Assertions about inserted jobs. +//! +//! Each assertion lists the jobs of the expected kinds, in every state, in +//! insertion (ID) order, and panics with a descriptive message when the +//! expectation isn't met, so it fails the calling test. The `_tx` variants +//! look inside a caller's open transaction, which is how to test code that +//! enqueues jobs transactionally before it commits. + +use std::fmt::Write as _; + +use chrono::{DateTime, DurationRound as _, TimeDelta, Utc}; +use riverqueue::{ + Client, Job, JobArgs, JobListOrderBy, JobListParams, JobRow, JobState, + database::DatabaseTransactionExecutor, +}; + +/// The most jobs an assertion reads, which is River's list limit. +const LIST_LIMIT: u32 = 10_000; + +/// Expected properties of an inserted job. +/// +/// Every property that's set must match. [`require_not_inserted`] fails only +/// when a job matches all of them. +#[derive(Clone, Debug, Default)] +#[non_exhaustive] +pub struct RequireInsertedOpts { + max_attempts: Option, + priority: Option, + queue: Option, + scheduled_at: Option>, + state: Option, + tags: Option>, +} + +impl RequireInsertedOpts { + /// Creates expectations that match any job. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Returns the expected maximum number of attempts. + #[must_use] + pub const fn max_attempts(&self) -> Option { + self.max_attempts + } + + /// Returns the expected priority. + #[must_use] + pub const fn priority(&self) -> Option { + self.priority + } + + /// Returns the expected queue. + #[must_use] + pub fn queue(&self) -> Option<&str> { + self.queue.as_deref() + } + + /// Returns the expected scheduled time. + #[must_use] + pub const fn scheduled_at(&self) -> Option> { + self.scheduled_at + } + + /// Returns the expected state. + #[must_use] + pub const fn state(&self) -> Option { + self.state + } + + /// Returns the expected tags. + #[must_use] + pub fn tags(&self) -> Option<&[String]> { + self.tags.as_deref() + } + + /// Expects this maximum number of attempts. + #[must_use] + pub const fn with_max_attempts(mut self, max_attempts: i16) -> Self { + self.max_attempts = Some(max_attempts); + self + } + + /// Expects this priority. + #[must_use] + pub const fn with_priority(mut self, priority: i16) -> Self { + self.priority = Some(priority); + self + } + + /// Expects this queue. + #[must_use] + pub fn with_queue(mut self, queue: impl Into) -> Self { + self.queue = Some(queue.into()); + self + } + + /// Expects this scheduled time, compared at microsecond precision like + /// the database stores it. + #[must_use] + pub const fn with_scheduled_at(mut self, scheduled_at: DateTime) -> Self { + self.scheduled_at = Some(scheduled_at); + self + } + + /// Expects this state. + #[must_use] + pub const fn with_state(mut self, state: JobState) -> Self { + self.state = Some(state); + self + } + + /// Expects exactly these tags, in order. + #[must_use] + pub fn with_tags(mut self, tags: impl IntoIterator>) -> Self { + self.tags = Some(tags.into_iter().map(Into::into).collect()); + self + } + + /// Compares each set property with `job`, returning a description of the + /// differences (for [`require_inserted`]) or of the matches (for + /// [`require_not_inserted`]), or `None` when the assertion holds. + fn compare(&self, job: &JobRow, excluding: bool) -> Option> { + let mut failures = Vec::new(); + let mut check = |matches: bool, matched: String, differs: String| -> bool { + match (matches, excluding) { + // One differing property is enough for a job not to match. + (false, true) => return false, + (true, true) => failures.push(matched), + (false, false) => failures.push(differs), + (true, false) => {} + } + true + }; + if let Some(expected) = self.max_attempts + && !check( + job.max_attempts == expected, + format!("max attempts equal to excluded {expected}"), + format!( + "max attempts {} not equal to expected {expected}", + job.max_attempts + ), + ) + { + return None; + } + if let Some(expected) = self.priority + && !check( + job.priority == expected, + format!("priority equal to excluded {expected}"), + format!("priority {} not equal to expected {expected}", job.priority), + ) + { + return None; + } + if let Some(expected) = &self.queue + && !check( + &job.queue == expected, + format!("queue equal to excluded '{expected}'"), + format!("queue '{}' not equal to expected '{expected}'", job.queue), + ) + { + return None; + } + if let Some(expected) = self.scheduled_at { + let micros = TimeDelta::microseconds(1); + let expected = expected.duration_trunc(micros).unwrap_or(expected); + let actual = job + .scheduled_at + .duration_trunc(micros) + .unwrap_or(job.scheduled_at); + if !check( + actual == expected, + format!("scheduled at equal to excluded {}", micro_time(expected)), + format!( + "scheduled at {} not equal to expected {}", + micro_time(actual), + micro_time(expected) + ), + ) { + return None; + } + } + if let Some(expected) = self.state + && !check( + job.state == expected, + format!("state equal to excluded '{}'", expected.as_str()), + format!( + "state '{}' not equal to expected '{}'", + job.state.as_str(), + expected.as_str() + ), + ) + { + return None; + } + if let Some(expected) = &self.tags + && !check( + &job.tags == expected, + format!("tags equal to excluded {expected:?}"), + format!("tags {:?} not equal to expected {expected:?}", job.tags), + ) + { + return None; + } + (!failures.is_empty()).then_some(failures) + } +} + +fn micro_time(time: DateTime) -> String { + time.to_rfc3339_opts(chrono::SecondsFormat::Micros, true) +} + +/// One job expected by [`require_many_inserted`]. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct ExpectedJob { + kind: &'static str, + opts: Option, +} + +impl ExpectedJob { + /// Expects a job of `A`'s kind. + #[must_use] + pub const fn of() -> Self { + Self { + kind: A::KIND, + opts: None, + } + } + + /// Adds expected properties for this job. + #[must_use] + pub fn opts(mut self, opts: RequireInsertedOpts) -> Self { + self.opts = Some(opts); + self + } +} + +fn params(kinds: impl IntoIterator) -> JobListParams { + JobListParams::default() + .kinds(kinds) + .states(JobState::ALL.iter().copied()) + .order_by(JobListOrderBy::Id) + .limit(LIST_LIMIT) +} + +async fn list(client: &Client, params: JobListParams) -> Vec { + client + .jobs() + .list(params) + .await + .unwrap_or_else(|error| panic!("Internal failure: listing jobs failed: {error}")) + .jobs +} + +async fn list_tx<'t, E>(client: &'t Client, executor: E, params: JobListParams) -> Vec +where + E: DatabaseTransactionExecutor<'t>, +{ + client + .jobs() + .list(params) + .tx(executor) + .await + .unwrap_or_else(|error| panic!("Internal failure: listing jobs failed: {error}")) + .jobs +} + +fn check_inserted(jobs: Vec, opts: Option<&RequireInsertedOpts>) -> Job { + let mut jobs = jobs.into_iter(); + let Some(job) = jobs.next() else { + panic!("No jobs found with kind: {}", A::KIND); + }; + assert!( + jobs.next().is_none(), + "More than one job found with kind: {} (you might want require_many_inserted instead)", + A::KIND + ); + if let Some(failures) = opts.and_then(|opts| opts.compare(&job, false)) { + panic!("Job with kind '{}' {}", job.kind, failures.join(", ")); + } + let args = job + .decode_args::() + .unwrap_or_else(|error| panic!("Internal failure: decoding job args failed: {error}")); + Job::new(args, job) +} + +fn check_not_inserted(kind: &str, jobs: &[JobRow], opts: Option<&RequireInsertedOpts>) { + let Some(opts) = opts else { + assert!( + jobs.is_empty(), + "{} jobs found with kind, but expected to find none: {kind}", + jobs.len() + ); + return; + }; + for job in jobs { + if let Some(failures) = opts.compare(job, true) { + panic!("Job with kind '{}' {}", job.kind, failures.join(", ")); + } + } +} + +fn check_many_inserted(expected: &[ExpectedJob], jobs: Vec) -> Vec { + let expected_kinds = expected.iter().map(|job| job.kind).collect::>(); + let actual_kinds = jobs.iter().map(|job| job.kind.as_str()).collect::>(); + assert!( + expected_kinds == actual_kinds, + "Inserted jobs didn't match expectation; expected: {expected_kinds:?}, actual: {actual_kinds:?}" + ); + for (index, (expected, job)) in expected.iter().zip(&jobs).enumerate() { + if let Some(failures) = expected + .opts + .as_ref() + .and_then(|opts| opts.compare(job, false)) + { + let mut message = format!("Job with kind '{}'", job.kind); + let _ = write!(message, " (expected job slice index {index})"); + panic!("{message} {}", failures.join(", ")); + } + } + jobs +} + +/// Asserts that exactly one job of `A`'s kind was inserted, in any state, +/// and returns it with decoded arguments. +/// +/// # Panics +/// +/// Panics, failing the calling test, when there is no such job, when there +/// is more than one, or when the jobs can't be listed or decoded. +pub async fn require_inserted(client: &Client) -> Job { + check_inserted(list(client, params([A::KIND])).await, None) +} + +/// Like [`require_inserted`], but also requires the job to match every +/// property set in `opts`. +/// +/// # Panics +/// +/// Panics under the same conditions as [`require_inserted`], and when a +/// property in `opts` doesn't match. +pub async fn require_inserted_with( + client: &Client, + opts: &RequireInsertedOpts, +) -> Job { + check_inserted(list(client, params([A::KIND])).await, Some(opts)) +} + +/// Like [`require_inserted`], but reads through `executor`'s open +/// transaction. +/// +/// # Panics +/// +/// Panics under the same conditions as [`require_inserted`]. +pub async fn require_inserted_tx<'t, A, E>(client: &'t Client, executor: E) -> Job +where + A: JobArgs, + E: DatabaseTransactionExecutor<'t>, +{ + check_inserted(list_tx(client, executor, params([A::KIND])).await, None) +} + +/// Like [`require_inserted_with`], but reads through `executor`'s open +/// transaction. +/// +/// # Panics +/// +/// Panics under the same conditions as [`require_inserted_with`]. +pub async fn require_inserted_tx_with<'t, A, E>( + client: &'t Client, + executor: E, + opts: &RequireInsertedOpts, +) -> Job +where + A: JobArgs, + E: DatabaseTransactionExecutor<'t>, +{ + check_inserted( + list_tx(client, executor, params([A::KIND])).await, + Some(opts), + ) +} + +/// Asserts that jobs of exactly the expected kinds were inserted, in this +/// order and number, and returns them. +/// +/// Only jobs of the expected kinds are considered, so a job of any other +/// kind doesn't affect the assertion. Expect a kind once for every job of it. +/// +/// # Panics +/// +/// Panics, failing the calling test, when the inserted kinds differ from the +/// expectation, when a property of an expected job doesn't match, or when +/// the jobs can't be listed. +pub async fn require_many_inserted(client: &Client, expected: &[ExpectedJob]) -> Vec { + let jobs = list(client, params(expected.iter().map(|job| job.kind))).await; + check_many_inserted(expected, jobs) +} + +/// Like [`require_many_inserted`], but reads through `executor`'s open +/// transaction. +/// +/// # Panics +/// +/// Panics under the same conditions as [`require_many_inserted`]. +pub async fn require_many_inserted_tx<'t, E>( + client: &'t Client, + executor: E, + expected: &[ExpectedJob], +) -> Vec +where + E: DatabaseTransactionExecutor<'t>, +{ + let jobs = list_tx( + client, + executor, + params(expected.iter().map(|job| job.kind)), + ) + .await; + check_many_inserted(expected, jobs) +} + +/// Asserts that no job of `A`'s kind was inserted. +/// +/// # Panics +/// +/// Panics, failing the calling test, when such a job exists or the jobs +/// can't be listed. +pub async fn require_not_inserted(client: &Client) { + let jobs = list(client, params([A::KIND])).await; + check_not_inserted(A::KIND, &jobs, None); +} + +/// Asserts that no job of `A`'s kind matches every property set in `opts`. +/// +/// # Panics +/// +/// Panics, failing the calling test, when a matching job exists or the jobs +/// can't be listed. +pub async fn require_not_inserted_with(client: &Client, opts: &RequireInsertedOpts) { + let jobs = list(client, params([A::KIND])).await; + check_not_inserted(A::KIND, &jobs, Some(opts)); +} + +/// Like [`require_not_inserted`], but reads through `executor`'s open +/// transaction. +/// +/// # Panics +/// +/// Panics under the same conditions as [`require_not_inserted`]. +pub async fn require_not_inserted_tx<'t, A, E>(client: &'t Client, executor: E) +where + A: JobArgs, + E: DatabaseTransactionExecutor<'t>, +{ + let jobs = list_tx(client, executor, params([A::KIND])).await; + check_not_inserted(A::KIND, &jobs, None); +} + +/// Like [`require_not_inserted_with`], but reads through `executor`'s open +/// transaction. +/// +/// # Panics +/// +/// Panics under the same conditions as [`require_not_inserted_with`]. +pub async fn require_not_inserted_tx_with<'t, A, E>( + client: &'t Client, + executor: E, + opts: &RequireInsertedOpts, +) where + A: JobArgs, + E: DatabaseTransactionExecutor<'t>, +{ + let jobs = list_tx(client, executor, params([A::KIND])).await; + check_not_inserted(A::KIND, &jobs, Some(opts)); +} + +#[cfg(all(test, feature = "sqlite"))] +mod tests { + use riverqueue::{ + InsertOpts, JobArgs, + migrate::SqliteMigrator, + sqlx::{ + SqlitePool, + sqlite::{SqliteConnectOptions, SqlitePoolOptions}, + }, + }; + use serde::{Deserialize, Serialize}; + + use super::*; + + #[derive(Clone, Debug, Deserialize, JobArgs, PartialEq, Serialize)] + #[river(kind = "require_first")] + struct FirstArgs { + value: i64, + } + + #[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] + #[river(kind = "require_second")] + struct SecondArgs {} + + struct TestBundle { + client: Client, + pool: SqlitePool, + } + + async fn setup() -> TestBundle { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect_with(SqliteConnectOptions::new().filename(":memory:")) + .await + .unwrap(); + SqliteMigrator::new(pool.clone()) + .migrate_up() + .await + .unwrap(); + let client = Client::builder(pool.clone()).build().unwrap(); + TestBundle { client, pool } + } + + #[tokio::test] + async fn require_inserted_returns_the_decoded_job() { + let bundle = setup().await; + bundle + .client + .insert(FirstArgs { value: 7 }) + .opts(InsertOpts::default().with_queue("custom").with_priority(2)) + .await + .unwrap(); + + let job = require_inserted_with::( + &bundle.client, + &RequireInsertedOpts::new() + .with_queue("custom") + .with_priority(2) + .with_state(JobState::Available), + ) + .await; + assert_eq!(job.args, FirstArgs { value: 7 }); + } + + #[tokio::test] + #[should_panic(expected = "No jobs found with kind: require_first")] + async fn require_inserted_fails_without_a_job() { + let bundle = setup().await; + require_inserted::(&bundle.client).await; + } + + #[tokio::test] + #[should_panic(expected = "More than one job found with kind: require_first")] + async fn require_inserted_fails_with_two_jobs() { + let bundle = setup().await; + for value in [1, 2] { + bundle.client.insert(FirstArgs { value }).await.unwrap(); + } + require_inserted::(&bundle.client).await; + } + + #[tokio::test] + #[should_panic( + expected = "Job with kind 'require_first' priority 1 not equal to expected 3, queue 'default' not equal to expected 'other'" + )] + async fn require_inserted_reports_every_mismatch() { + let bundle = setup().await; + bundle.client.insert(FirstArgs { value: 1 }).await.unwrap(); + require_inserted_with::( + &bundle.client, + &RequireInsertedOpts::new() + .with_queue("other") + .with_priority(3), + ) + .await; + } + + #[tokio::test] + async fn require_inserted_tx_sees_uncommitted_jobs() { + let bundle = setup().await; + let mut transaction = bundle.pool.begin().await.unwrap(); + bundle + .client + .insert(FirstArgs { value: 1 }) + .tx(&mut transaction) + .await + .unwrap(); + + require_inserted_tx::(&bundle.client, &mut transaction).await; + require_not_inserted_tx::(&bundle.client, &mut transaction).await; + transaction.rollback().await.unwrap(); + require_not_inserted::(&bundle.client).await; + } + + #[tokio::test] + async fn require_many_inserted_matches_kinds_in_order() { + let bundle = setup().await; + bundle.client.insert(FirstArgs { value: 1 }).await.unwrap(); + bundle.client.insert(SecondArgs {}).await.unwrap(); + bundle.client.insert(FirstArgs { value: 2 }).await.unwrap(); + + let jobs = require_many_inserted( + &bundle.client, + &[ + ExpectedJob::of::(), + ExpectedJob::of::() + .opts(RequireInsertedOpts::new().with_queue("default")), + ExpectedJob::of::(), + ], + ) + .await; + assert_eq!(jobs.len(), 3); + } + + #[tokio::test] + #[should_panic( + expected = "Inserted jobs didn't match expectation; expected: [\"require_second\", \"require_first\"], actual: [\"require_first\", \"require_second\"]" + )] + async fn require_many_inserted_fails_on_a_different_order() { + let bundle = setup().await; + bundle.client.insert(FirstArgs { value: 1 }).await.unwrap(); + bundle.client.insert(SecondArgs {}).await.unwrap(); + + require_many_inserted( + &bundle.client, + &[ + ExpectedJob::of::(), + ExpectedJob::of::(), + ], + ) + .await; + } + + #[tokio::test] + async fn require_not_inserted_passes_when_any_property_differs() { + let bundle = setup().await; + bundle.client.insert(FirstArgs { value: 1 }).await.unwrap(); + + require_not_inserted::(&bundle.client).await; + require_not_inserted_with::( + &bundle.client, + &RequireInsertedOpts::new() + .with_queue("default") + .with_priority(4), + ) + .await; + } + + #[tokio::test] + #[should_panic(expected = "Job with kind 'require_first' queue equal to excluded 'default'")] + async fn require_not_inserted_fails_when_every_property_matches() { + let bundle = setup().await; + bundle.client.insert(FirstArgs { value: 1 }).await.unwrap(); + + require_not_inserted_with::( + &bundle.client, + &RequireInsertedOpts::new().with_queue("default"), + ) + .await; + } +} diff --git a/rust/riverqueue/Cargo.toml b/rust/riverqueue/Cargo.toml new file mode 100644 index 000000000..9febe6795 --- /dev/null +++ b/rust/riverqueue/Cargo.toml @@ -0,0 +1,120 @@ +[package] +name = "riverqueue" +description = "Fast and reliable background jobs for Rust" +keywords = ["background", "jobs", "postgres", "queue", "sqlite"] +categories = ["asynchronous", "database"] +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true +readme = "README.md" +repository.workspace = true +homepage.workspace = true + +[package.metadata.docs.rs] +all-features = true + +[features] +default = ["postgres"] +# Resolves IANA zone names such as `America/New_York` in cron `CRON_TZ=` +# prefixes from a bundled time zone database. +chrono-tz = ["dep:chrono-tz"] +postgres = ["riverqueue-migrate/postgres", "sqlx/postgres"] +sqlite = ["riverqueue-migrate/sqlite", "sqlx/sqlite"] + +[dependencies] +async-trait.workspace = true +base64.workspace = true +chrono.workspace = true +chrono-tz = { version = "0.10.4", optional = true } +futures-util.workspace = true +riverqueue-migrate = { path = "../riverqueue-migrate", version = "=0.50.0-alpha.1", default-features = false } +riverqueue-macros = { path = "../riverqueue-macros", version = "=0.50.0-alpha.1" } +rand = { workspace = true } +serde.workspace = true +serde_json.workspace = true +sha2.workspace = true +sqlx = { workspace = true, features = ["chrono", "derive", "json"] } +thiserror.workspace = true +tokio.workspace = true +tokio-util.workspace = true +tracing.workspace = true + +[dev-dependencies] +anyhow = "1.0.104" +riverqueue-migrate = { path = "../riverqueue-migrate", version = "=0.50.0-alpha.1", features = ["postgres", "sqlite"] } +serde.workspace = true +tracing-subscriber = { version = "0.3.20", default-features = false, features = ["registry", "std"] } +# Fault-injection tests proxy database connections. +tokio = { workspace = true, features = ["io-util", "net", "rt-multi-thread", "signal", "test-util"] } + +[[example]] +name = "basic_worker" +required-features = ["postgres"] + +[[example]] +name = "cancellation" +required-features = ["postgres"] + +[[example]] +name = "custom_schema" +required-features = ["postgres"] + +[[example]] +name = "events" +required-features = ["postgres"] + +[[example]] +name = "graceful_shutdown" +required-features = ["postgres"] + +[[example]] +name = "mixed_go_rust" +required-features = ["postgres"] + +[[example]] +name = "periodic" +required-features = ["postgres"] + +[[example]] +name = "sqlite" +required-features = ["sqlite"] + +[[example]] +name = "transactions" +required-features = ["postgres"] + +[[example]] +name = "unique" +required-features = ["postgres"] + +[[test]] +name = "lifecycle" +required-features = ["sqlite"] + +[[test]] +name = "work_extensions" +required-features = ["sqlite"] + +[[test]] +name = "postgres" +required-features = ["postgres"] + +[[test]] +name = "postgres_parity" +required-features = ["postgres"] + +[[test]] +name = "sqlite_parity" +required-features = ["sqlite"] + +[[test]] +name = "sqlite_runtime" +required-features = ["sqlite"] + +[[test]] +name = "sqlite_storage" +required-features = ["sqlite"] + +[lints] +workspace = true diff --git a/rust/riverqueue/LICENSE b/rust/riverqueue/LICENSE new file mode 120000 index 000000000..30cff7403 --- /dev/null +++ b/rust/riverqueue/LICENSE @@ -0,0 +1 @@ +../../LICENSE \ No newline at end of file diff --git a/rust/riverqueue/README.md b/rust/riverqueue/README.md new file mode 100644 index 000000000..e988e4dfd --- /dev/null +++ b/rust/riverqueue/README.md @@ -0,0 +1,379 @@ +# riverqueue + +`riverqueue` is the Rust and Tokio client for [River](https://riverqueue.com), +a fast and reliable background job system backed by PostgreSQL or SQLite. It +shares River's database schema and job protocol with River for Go, so Rust +and Go services can insert and work jobs in the same database. + +This crate is a pre-release preview. Each release matches the River for Go +release with the same minor version; the +[mixed deployment guide](https://docs.rs/riverqueue/latest/riverqueue/guide/mixed_deployments/index.html) +covers running both against one database. + +## Installation + +```toml +[dependencies] +riverqueue = "0.50.0-alpha.1" +serde = { version = "1", features = ["derive"] } +serde_json = "1" +tokio = { version = "1", features = ["macros", "rt-multi-thread", "signal"] } +``` + +The quick start below needs exactly these. River runs on Tokio, and job +arguments derive Serde's `Serialize` and `Deserialize`. The minimum supported +Rust version is 1.95. + +| Feature | Default | Enables | +|---|---|---| +| `postgres` | yes | PostgreSQL through SQLx | +| `sqlite` | no | SQLite 3.45 or newer through SQLx | +| `chrono-tz` | no | IANA zone names such as `America/New_York` in cron `CRON_TZ=` and `TZ=` prefixes | + +For SQLite alone, use +`riverqueue = { version = "0.50.0-alpha.1", default-features = false, features = ["sqlite"] }`. + +River's API uses types from SQLx (pools and transactions), Chrono +(timestamps), `serde_json` (metadata, outputs, and other JSON values), and +`tokio-util` (the worker's `CancellationToken`). The crate re-exports each one +as `riverqueue::sqlx`, `riverqueue::chrono`, `riverqueue::serde_json`, and +`riverqueue::tokio_util`. Use the re-exports, or depend on versions +compatible with River's (SQLx 0.9, Chrono 0.4, `serde_json` 1, and +`tokio-util` 0.7), so the types match. River doesn't choose a TLS +implementation for SQLx; enable one of SQLx's TLS features in your own SQLx +dependency if your database connections use TLS. + +## Quick start + +Define serializable arguments, register an async function or a [`Worker`], +apply River's migrations, and start a client: + +```rust,no_run +use riverqueue::migrate::PostgresMigrator; +use riverqueue::sqlx::PgPool; +use riverqueue::{ + BoxError, Client, Job, JobArgs, QueueConfig, WorkContext, WorkOutcome, WorkerRegistry, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "send_email")] +struct SendEmail { + address: String, +} + +async fn send_email( + context: WorkContext, + job: Job, +) -> Result { + println!("sending email to {}", job.args.address); + context.record_output(serde_json::json!({"delivered": true}))?; + Ok(WorkOutcome::Complete) +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let pool = PgPool::connect(&std::env::var("DATABASE_URL")?).await?; + // Create or upgrade River's tables. Applications often run + // `riverqueue migrate-up` from `riverqueue-cli` at deploy time instead. + PostgresMigrator::new(pool.clone()).migrate_up().await?; + + let mut workers = WorkerRegistry::new(); + workers.register_fn(send_email)?; + + let client = Client::builder(pool) + .workers(workers) + .queue("default", QueueConfig::new(10)) + .build()?; + // Work jobs until Ctrl-C, then stop fetching and let running jobs finish. + let mut run = client.start_with_graceful_shutdown(async { + let _ = tokio::signal::ctrl_c().await; + })?; + + client + .insert(SendEmail { + address: "person@example.com".to_owned(), + }) + .await?; + + run.wait().await?; + Ok(()) +} +``` + +Apply migrations before any client starts, and start clients inside a Tokio +runtime. `Client::start` returns a `RunHandle`: await `wait`, `shutdown` (a +soft stop that lets running jobs finish), or `shutdown_now` (which cancels +them). `RunHandle::stopper` returns a cloneable `Stopper` for stopping the +client from another task, such as a signal handler. The handle controls the +running client: dropping every `Client` clone doesn't stop it, dropping the +handle requests a hard stop, and `RunHandle::detach` leaves the client running +unsupervised. + +## Inserting jobs + +`Client::insert(args)` inserts a job with its type's default options, which +come from `JobArgs::default_insert_opts`. A call can override only the options +it needs: + +```rust,no_run +use riverqueue::{Client, InsertOpts, JobArgs}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize, JobArgs, Serialize)] +#[river(kind = "send_email")] +struct SendEmail { + address: String, +} + +async fn enqueue_urgent(client: &Client) -> Result<(), riverqueue::Error> { + client + .insert(SendEmail { address: "urgent@example.com".to_owned() }) + .opts( + InsertOpts::default() + .with_queue("critical") + .with_priority(1) + .with_max_attempts(8), + ) + .await?; + Ok(()) +} +``` + +An option set on the call wins over the job type's default, which wins over +the client's default, which wins over River's. `insert_many` inserts many jobs +of one kind atomically and returns results in input order; `insert_batch` does +the same for jobs of different kinds. + +Chain `.tx(&mut transaction)` onto an insertion, or onto any request from +`client.jobs()` or `client.queues()`, to run it in the same SQL transaction as +application writes. Jobs become visible, and clients are notified, only when +the transaction commits. Begin transactions with +`riverqueue::database::begin_postgres(&pool)`, or on SQLite with +`riverqueue::database::begin_sqlite_write(&pool)`, which uses +`BEGIN IMMEDIATE` so a transaction that reads before it writes can't fail with +`SQLITE_BUSY_SNAPSHOT`. Both begin on a separate task, so they're safe to +abandon partway, for example in a `select!` or a timeout; SQLx's own +`pool.begin()` isn't, and can return a connection to the pool still inside a +transaction. + +River runs these requests directly in your transaction, without a savepoint. +When one returns an error it may have already written part of its work there, +so roll the transaction back. To recover from the error and continue the +transaction instead, open your own savepoint before the request. + +## Managing jobs and queues + +`client.jobs()` gets, lists, cancels, retries, updates, and deletes persisted +jobs. `client.queues()` gets, lists, pauses, resumes, and updates the queue +records every client shares, and `client.local_queues()` changes which queues +this client works while it runs. Job and queue requests run when awaited and +take `.tx(&mut transaction)` like insertions: + +```rust,no_run +use riverqueue::sqlx::PgPool; +use riverqueue::{Client, JobListParams, JobState, QueueConfig, QueueSelector}; + +async fn maintain(client: &Client, pool: &PgPool) -> Result<(), riverqueue::Error> { + let page = client + .jobs() + .list(JobListParams::default().states([JobState::Retryable]).limit(50)) + .await?; + for job in &page.jobs { + client.jobs().retry(job.id).await?; + } + + let mut transaction = riverqueue::database::begin_postgres(pool).await?; + client.queues().pause(QueueSelector::All).tx(&mut transaction).await?; + transaction.commit().await?; + + client.local_queues().add("reports", QueueConfig::new(2))?; + Ok(()) +} +``` + +## Worker outcomes and cancellation + +An `Ok(WorkOutcome::Complete)` completes a job. `Snooze` reschedules without +consuming an attempt, `Discard` finalizes without another retry, and `Cancel` +finalizes as cancelled. A worker error is passed through its retry policy until +the maximum attempt count is reached. + +The `WorkContext` cancellation token is triggered by a job timeout, remote job +cancellation, or client stop. Workers should select or check cancellation at +natural await points, and return `WorkCancelled` when they stop because of it: + +```rust +use riverqueue::{Job, JobArgs, WorkCancelled, WorkContext, WorkOutcome}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Deserialize, JobArgs, Serialize)] +#[river(kind = "build_report")] +struct BuildReport {} + +async fn build_report( + context: WorkContext, + _job: Job, +) -> Result { + tokio::select! { + () = context.cancellation_token().cancelled() => Err(WorkCancelled), + () = tokio::time::sleep(std::time::Duration::from_secs(1)) => { + Ok(WorkOutcome::Complete) + } + } +} +``` + +During a client's hard stop (`RunHandle::shutdown_now` or +`Stopper::stop_now`), a job whose worker returns `WorkCancelled`, anywhere in +its error's source chain, becomes available again without using up its +attempt. Any other error is recorded and consumes the attempt, and `Ok` +completes the job. After the configured stuck threshold, River can abort a +Tokio task that yields, which fails its attempt, but it can't stop CPU-bound +work or a blocking call already in progress. + +Implement [`Worker`] when a kind needs a custom timeout or next-retry decision. +Use `WorkerRegistry::register_fn` for an async function or capturing closure. + +## Events + +Subscriptions are local observations, not a durable event stream. Subscribe +before starting a client to see events from its first jobs. Receivers are +bounded and report `EventRecvError::Lagged` with the number of dropped events, +and `EventReceiver` is also a `Stream`. Job events are sent after their +results are persisted, and independent jobs have no global completion order. + +## Reliability features + +- Unique jobs deduplicate by kind, encoded arguments or selected argument + paths, queue, period, and job state. The derive macro follows Serde's + serialization names and omits missing optional fields. +- Periodic jobs run on the elected leader and can be configured when the client + is built or at runtime; stable IDs prevent duplicate registration. + `CronSchedule` accepts standard five-field cron syntax and descriptors such + as `@hourly` and `@every 90s`, evaluated in the process's local time zone + unless another `CronTimeZone` is chosen. `CRON_TZ=` and `TZ=` prefixes + naming IANA zones such as `America/New_York` need the `chrono-tz` feature, + which bundles the time zone database; without it only `UTC`, `Local`, and + `Etc/GMT±N` names parse. +- Resumable steps persist the last completed step and an optional cursor. Use + the transactional checkpoint helpers when progress and business data must + commit together. +- Insertion middleware wraps the insert-begin hooks, and work middleware wraps + the work hooks, argument decoding, and the worker. See [`WorkMiddleware`] + and [`Hook`]. + +## Leadership and maintenance + +One client at a time holds a database lease and runs the leader-owned +services: the job scheduler, the stuck-job rescuer, the job and queue cleaners, +the periodic job enqueuer, the PostgreSQL reindexer, and the SQLite +notification cleaner. Losing the lease or stopping the client stops them +immediately. + +The default client ID combines the host name, the creation time, and a random +suffix; set a stable `id` only when it's unique per process. + +The rescuer considers a job stuck after `rescue_after`, which defaults to one +hour, or to the job timeout plus one hour when a job timeout is configured, and +must not be shorter than the job timeout. The leader discards stuck jobs of +kinds its own worker registry doesn't know, so clients that share a database +should register the same kinds. A client that doesn't know every kind can be +built with `ClientBuilder::without_leader_election`: it works its queues but +never becomes leader, so at least one other client must stay eligible. + +Periodic jobs are scheduled from the time each term begins, and jobs with +`PeriodicJobOpts::with_run_on_start(true)` are inserted once per term gained. +An occurrence whose insert fails is logged and skipped. + +## Database support + +Pass an SQLx pool, or a `PostgresDatabase` or `SqliteDatabase` with options, +to `Client::builder`. Both backends implement the same job and queue behavior, +and `Client` isn't generic over the database, so backend types don't reach +workers, contexts, or extensions. + +SQLite needs the `sqlite` feature, and version 3.45 or newer, since River +stores JSON with SQLite's JSONB functions. River uses the caller's pool as configured and doesn't change +connection pragmas. For a file database, enable WAL and a busy timeout so a +short writer collision waits instead of failing. A private `:memory:` database +belongs to one connection, so limit the pool to one connection or use a +shared-cache URI: + +```rust,no_run +use std::{str::FromStr, time::Duration}; + +use riverqueue::Client; +use riverqueue::sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}; + +async fn sqlite_client() -> Result> { + let options = SqliteConnectOptions::from_str("sqlite://river.db")? + .create_if_missing(true) + .journal_mode(SqliteJournalMode::Wal) + .busy_timeout(Duration::from_secs(5)); + let pool = SqlitePoolOptions::new() + .max_connections(5) + .connect_with(options) + .await?; + Ok(Client::builder(pool).build()?) +} +``` + +## Modules + +- [`job`]: arguments, insertion options, persisted rows, outcomes, and unique + job configuration. +- [`encoding`]: the JSON encoding River uses for job arguments, which keeps + unique keys identical across Rust and Go. +- [`worker`]: typed workers, function registration, cancellation, outputs, + and resumable work. +- [`event`]: event payloads and bounded subscriptions. +- [`queue`] and [`query`]: queue records and job list filters and cursors. +- [`periodic`]: schedules and runtime periodic job registration. +- [`extension`]: hooks, middleware, policies, and metrics. +- [`database`]: PostgreSQL and SQLite database options, and the transactions + River's `.tx` methods accept. +- [`error`]: structured errors that keep their sources. +- [`protocol`]: wire values such as notification topics and unique keys, for + tools that work with River's tables directly. + +Setters follow two conventions. Builders and request parameters, which +exist only to be passed on (`ClientBuilder`, `JobListParams`, +`JobUpdateParams`), take plain setter names such as `queue(..)` and +`limit(..)`. Configuration values that also expose each setting through a +same-named getter (`InsertOpts`, `UniqueOpts`, `PeriodicJobOpts`, +`QueueConfig`, `MaintenanceConfig`, `SubscribeConfig`, `PostgresDatabase`, +`PostgresMigrator`) use `with_*` methods that return the value with one +setting changed, like `PathBuf::with_extension`, so `UniqueOpts::by_args` +reads what `UniqueOpts::with_by_args` sets. Durations that can be disabled +are explicit, as in `ClientBuilder::without_job_timeout` and +`Retention::Keep`. + +The crate's `examples` directory has runnable programs for a basic worker, +graceful shutdown, cancellation, transactional enqueueing and completion, +unique and periodic jobs, events, custom PostgreSQL schemas, SQLite, and a +Rust and Go service sharing one database. The +[River documentation](https://riverqueue.com/docs) explains queueing concepts. + +## Benchmarking + +The [`riverqueue-cli`](https://crates.io/crates/riverqueue-cli) crate provides +`riverqueue bench`, a benchmark for development databases. It truncates the +selected River job table, so use a disposable database, and reports periodic +throughput plus final throughput and p95 latency. Run +`riverqueue bench --help` for its options. + +[`Hook`]: https://docs.rs/riverqueue/latest/riverqueue/trait.Hook.html +[`WorkMiddleware`]: https://docs.rs/riverqueue/latest/riverqueue/trait.WorkMiddleware.html +[`Worker`]: https://docs.rs/riverqueue/latest/riverqueue/trait.Worker.html +[`database`]: https://docs.rs/riverqueue/latest/riverqueue/database/index.html +[`encoding`]: https://docs.rs/riverqueue/latest/riverqueue/encoding/index.html +[`error`]: https://docs.rs/riverqueue/latest/riverqueue/error/index.html +[`event`]: https://docs.rs/riverqueue/latest/riverqueue/event/index.html +[`extension`]: https://docs.rs/riverqueue/latest/riverqueue/extension/index.html +[`job`]: https://docs.rs/riverqueue/latest/riverqueue/job/index.html +[`periodic`]: https://docs.rs/riverqueue/latest/riverqueue/periodic/index.html +[`protocol`]: https://docs.rs/riverqueue/latest/riverqueue/protocol/index.html +[`query`]: https://docs.rs/riverqueue/latest/riverqueue/query/index.html +[`queue`]: https://docs.rs/riverqueue/latest/riverqueue/queue/index.html +[`worker`]: https://docs.rs/riverqueue/latest/riverqueue/worker/index.html diff --git a/rust/riverqueue/docs/mixed-deployments.md b/rust/riverqueue/docs/mixed-deployments.md new file mode 100644 index 000000000..ace84b343 --- /dev/null +++ b/rust/riverqueue/docs/mixed-deployments.md @@ -0,0 +1,122 @@ +# Running Rust and Go River together + +River's Rust and Go clients share one database schema and protocol. A Rust +service can insert jobs that Go workers run, work jobs that Go services insert, +and take part in leader election and maintenance alongside Go clients. This +guide covers what must match between them and how to roll Rust into an +existing Go deployment and back out again. + +## Matching versions + +Each `riverqueue` preview is a prerelease of the next River Go minor release +and works with the River Go release before it: `riverqueue` 0.50.0-alpha.1 +works with River Go 0.49. Patch releases of that River Go release can be mixed +with it. Upgrade both implementations together when moving to a new pairing, +following the rolling procedure below. + +Every River client in a deployment must understand the same schema. Run +migrations once with either implementation, before starting clients of the +new version: + +- Go: `river migrate-up`, or `rivermigrate` from Go code. +- Rust: `riverqueue migrate-up` from `riverqueue-cli`, or `PostgresMigrator`/ + `SqliteMigrator` from `riverqueue-migrate`. + +Both write the same `river_migration` history, so either can migrate a database +the other created. + +## Queues and job kinds + +Clients fetch work by queue, not by kind. A client that fetches a job whose +kind it has no worker for records a retryable "job kind is not registered" +error, and the job is retried until another client works it or it runs out of +attempts. Give each language its own queues for the kinds only it works: + +```text +Go: queues "default", "billing" workers for billing kinds +Rust: queues "rust_default", "images" workers for image kinds +``` + +Either language can insert into any queue. + +To share a queue between clients with different workers instead, build each +of them with `ClientBuilder::fetch_only_known_kinds(true)` in Rust or +`Config.FetchOnlyKnownKinds` in Go. Such a client claims only jobs of kinds it +has a worker for, so jobs of other kinds stay available for the clients that +know them without using attempts. This only affects claiming: pair it with +the leader election advice below, since the leader's rescuer still handles +stuck jobs of every kind. + +The elected leader runs maintenance for the whole database, including the job +rescuer. Like River Go, the rescuer discards stuck jobs whose kind the leader +has no worker for rather than retrying them. If kinds are split between +languages, either register a worker for every kind in every client that can be +elected leader, or keep clients that don't know every kind out of leader +election: `ClientBuilder::without_leader_election` in Rust and +`Config.LeaderElectionDisabled` in Go. Such clients still work their queues, +but never run maintenance, so at least one client in the deployment must +remain eligible to lead. + +## Unique jobs + +Unique keys hash job arguments as encoded JSON, so unique jobs inserted from +both languages must serialize the same way: + +- Use the same JSON field names. Go uses struct tags; Rust follows Serde's + rename rules. +- Without selected fields, every top-level key participates, sorted by its + literal JSON name. With selected fields, use the same JSON-name paths in + both languages. `#[river(unique)]` selects a serialized Rust field name + literally; `unique(by_args("account.id"))` selects a nested field, while + `unique(by_args("account\\.id"))` selects one field named `account.id`. + Manual `JobArgs::unique_fields` implementations pass slices of decoded + path components, such as `&[&["account", "id"], &["account.id"]]`. + Nested objects retain their encoded field order, so a nested Go struct + needs the same field order in Rust, and a Go map needs sorted keys in Rust + (a `BTreeMap`). Selecting individual scalar fields avoids depending on + nested order. Numeric array-index path components are not yet supported + by Rust's unique-field API. +- Rust encodes floating point numbers and escapes strings the way Go's + `encoding/json` does, so `1.0` hashes the same as Go's `1`. +- `ByPeriod` periods are measured in UTC from the job's scheduled time, as in + River Go 0.48 and later. River Go before 0.48 derives periods from the + insertion time in the process's local time zone, which is one more reason + to run matched versions. + +## Periodic jobs + +Periodic jobs are enqueued only by the elected leader. Configure the same +periodic jobs, with the same IDs and schedules, in every client that can become +leader, whichever language it's written in; otherwise a job stops being +enqueued whenever a client without it is elected. Clients kept out of leader +election can't configure periodic jobs, but still work those a leader enqueues +in their queues. + +Cron schedules use Go River's standard five-field syntax in both languages. +A schedule with a `CRON_TZ=` prefix naming an IANA zone parses in Rust only +with the `chrono-tz` feature enabled; without it, Rust rejects the schedule +rather than evaluating it in another zone. + +## SQLite + +SQLite databases shared between processes must use WAL mode and a busy timeout +in every process. Go clients sharing a SQLite file with Rust should open it with +`_txlock=immediate` so that transactions that read before writing don't fail +with `SQLITE_BUSY` when a Rust process commits in between. + +## Rolling deployment + +1. Upgrade River Go to the matched release and run migrations. +2. Deploy a small number of Rust clients alongside the Go clients. +3. Watch queue depth, retries, rescued jobs, leadership changes, and database + connection counts while increasing Rust's share. +4. Keep at least one Go deployment available until you're confident in the + Rust services. + +## Rolling back + +Rolling back doesn't touch the schema. Stop Rust clients gracefully with +`RunHandle::shutdown` and let the Go clients continue. Jobs that Rust inserted +are ordinary River rows that Go workers can run, and any job a stopped Rust +client left running is recovered by the rescuer. Only migrate down as a +separately planned operation once no deployed client needs the newer schema. diff --git a/rust/riverqueue/examples/basic_worker.rs b/rust/riverqueue/examples/basic_worker.rs new file mode 100644 index 000000000..ecbed9e4e --- /dev/null +++ b/rust/riverqueue/examples/basic_worker.rs @@ -0,0 +1,51 @@ +//! Works one job and stops. +//! +//! ```sh +//! DATABASE_URL=postgres://localhost/river_example cargo run -p riverqueue --example basic_worker +//! ``` + +use std::error::Error; + +use riverqueue::sqlx::PgPool; +use riverqueue::{ + BoxError, Client, EventKind, Job, JobArgs, QueueConfig, WorkContext, WorkOutcome, + WorkerRegistry, migrate::PostgresMigrator, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "send_email")] +struct SendEmail { + address: String, +} + +async fn send_email(context: WorkContext, job: Job) -> Result { + println!("sending email to {}", job.args.address); + context.record_output(serde_json::json!({"delivered": true}))?; + Ok(WorkOutcome::Complete) +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let pool = PgPool::connect(&std::env::var("DATABASE_URL")?).await?; + // Apply River's schema before starting a client. + PostgresMigrator::new(pool.clone()).migrate_up().await?; + let mut workers = WorkerRegistry::new(); + workers.register_fn(send_email)?; + let client = Client::builder(pool) + .workers(workers) + .queue("default", QueueConfig::new(10)) + .build()?; + let mut completed = client.subscribe(&[EventKind::JobCompleted])?; + let mut run = client.start()?; + + let inserted = client + .insert(SendEmail { + address: "person@example.com".to_owned(), + }) + .await?; + while completed.recv().await?.as_job().map(|event| event.job.id) != Some(inserted.id()) {} + + run.shutdown().await?; + Ok(()) +} diff --git a/rust/riverqueue/examples/cancellation.rs b/rust/riverqueue/examples/cancellation.rs new file mode 100644 index 000000000..ecd6470f0 --- /dev/null +++ b/rust/riverqueue/examples/cancellation.rs @@ -0,0 +1,64 @@ +//! Cancels a running job from outside its worker. +//! +//! ```sh +//! DATABASE_URL=postgres://localhost/river_example cargo run -p riverqueue --example cancellation +//! ``` +//! +//! Cancelling a running job triggers its worker's cancellation token, on +//! whichever client is working it. + +use std::{error::Error, time::Duration}; + +use riverqueue::sqlx::PgPool; +use riverqueue::{ + Client, Job, JobArgs, QueueConfig, WorkCancelled, WorkContext, WorkOutcome, Worker, + WorkerRegistry, migrate::PostgresMigrator, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "cancellable_report")] +struct CancellableReport { + report_id: i64, +} + +struct CancellableReportWorker; + +impl Worker for CancellableReportWorker { + type Error = WorkCancelled; + + async fn work( + &self, + context: WorkContext, + job: Job, + ) -> Result { + tokio::select! { + // Returning `WorkCancelled` records a remote cancellation as + // cancelled, and during a client's hard stop makes the job + // available again without using up its attempt. + () = context.cancellation_token().cancelled() => Err(WorkCancelled), + () = tokio::time::sleep(Duration::from_secs(30)) => { + println!("generated report {}", job.args.report_id); + Ok(WorkOutcome::Complete) + } + } + } +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let pool = PgPool::connect(&std::env::var("DATABASE_URL")?).await?; + PostgresMigrator::new(pool.clone()).migrate_up().await?; + let mut workers = WorkerRegistry::new(); + workers.register(CancellableReportWorker)?; + let client = Client::builder(pool) + .workers(workers) + .queue("default", QueueConfig::new(1)) + .build()?; + let mut run = client.start()?; + let job = client.insert(CancellableReport { report_id: 42 }).await?; + + client.jobs().cancel(job.id()).await?; + run.shutdown().await?; + Ok(()) +} diff --git a/rust/riverqueue/examples/custom_schema.rs b/rust/riverqueue/examples/custom_schema.rs new file mode 100644 index 000000000..46b85adbb --- /dev/null +++ b/rust/riverqueue/examples/custom_schema.rs @@ -0,0 +1,39 @@ +//! Keeps River's tables in their own PostgreSQL schema. +//! +//! ```sh +//! DATABASE_URL=postgres://localhost/river_example cargo run -p riverqueue --example custom_schema +//! ``` + +use std::error::Error; + +use riverqueue::sqlx::{self, PgPool}; +use riverqueue::{ + Client, + database::{PostgresDatabase, SchemaName}, +}; +use riverqueue_migrate::PostgresMigrator; + +#[tokio::main] +async fn main() -> Result<(), Box> { + let pool = PgPool::connect(&std::env::var("DATABASE_URL")?).await?; + let schema = SchemaName::new("river_jobs")?; + // Migrations create River's tables but not the schema itself. + sqlx::query("CREATE SCHEMA IF NOT EXISTS river_jobs") + .execute(&pool) + .await?; + PostgresMigrator::new(pool.clone()) + .with_schema(schema.clone()) + .migrate_up() + .await?; + + let client = Client::builder(PostgresDatabase::new(pool).with_schema(schema)).build()?; + println!( + "River schema: {}", + client + .database() + .as_postgres() + .expect("client is configured for PostgreSQL") + .schema() + ); + Ok(()) +} diff --git a/rust/riverqueue/examples/events.rs b/rust/riverqueue/examples/events.rs new file mode 100644 index 000000000..21c50c280 --- /dev/null +++ b/rust/riverqueue/examples/events.rs @@ -0,0 +1,94 @@ +//! Observes job outcomes through client events. +//! +//! ```sh +//! DATABASE_URL=postgres://localhost/river_example cargo run -p riverqueue --example events +//! ``` +//! +//! Events describe what this client's workers did. A subscriber opts in to +//! each kind; a slow subscriber drops events rather than blocking workers. + +use std::error::Error; + +use riverqueue::sqlx::PgPool; +use riverqueue::{ + Client, Event, EventKind, InsertOpts, Job, JobArgs, JobEventKind, QueueConfig, WorkContext, + WorkOutcome, WorkerRegistry, migrate::PostgresMigrator, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "charge_card")] +struct ChargeCard { + amount_cents: i64, +} + +#[derive(Debug, thiserror::Error)] +#[error("card declined")] +struct Declined; + +async fn charge_card(_context: WorkContext, job: Job) -> Result { + if job.args.amount_cents > 10_000 { + return Err(Declined); + } + Ok(WorkOutcome::Complete) +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let pool = PgPool::connect(&std::env::var("DATABASE_URL")?).await?; + PostgresMigrator::new(pool.clone()).migrate_up().await?; + + let mut workers = WorkerRegistry::new(); + workers.register_fn(charge_card)?; + let client = Client::builder(pool) + .workers(workers) + .queue("default", QueueConfig::new(4)) + .build()?; + // Subscribe before starting so no event is missed. + let mut events = client.subscribe(&[EventKind::JobCompleted, EventKind::JobFailed])?; + let mut run = client.start()?; + + let small = client.insert(ChargeCard { amount_cents: 500 }).await?; + let large = client + .insert(ChargeCard { + amount_cents: 50_000, + }) + // One attempt, so the failure is final and the example ends quickly. + .opts(InsertOpts::default().with_max_attempts(1)) + .await?; + + let mut seen = 0; + while seen < 2 { + let Event::Job(event) = events.recv().await? else { + continue; + }; + if ![small.id(), large.id()].contains(&event.job.id) { + continue; + } + let run_time = event + .statistics + .map(|statistics| statistics.run_duration) + .unwrap_or_default(); + match event.kind { + JobEventKind::Completed => { + println!("job {} completed in {run_time:?}", event.job.id); + } + JobEventKind::Failed => { + let error = event + .job + .errors + .last() + .map_or("", |error| error.error.as_str()); + println!( + "job {} failed ({:?}): {error}", + event.job.id, event.job.state + ); + } + _ => continue, + } + seen += 1; + } + + run.shutdown().await?; + Ok(()) +} diff --git a/rust/riverqueue/examples/graceful_shutdown.rs b/rust/riverqueue/examples/graceful_shutdown.rs new file mode 100644 index 000000000..87cff07ad --- /dev/null +++ b/rust/riverqueue/examples/graceful_shutdown.rs @@ -0,0 +1,63 @@ +//! Stops a client gracefully on the first Ctrl-C and cancels running jobs on +//! the second. + +use std::{error::Error, time::Duration}; + +use riverqueue::sqlx::PgPool; +use riverqueue::{ + BoxError, Client, Job, JobArgs, QueueConfig, WorkCancelled, WorkContext, WorkOutcome, + WorkerRegistry, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "generate_report")] +struct GenerateReport { + report_id: i64, +} + +async fn generate_report( + context: WorkContext, + job: Job, +) -> Result { + // Long-running work should watch its cancellation token so a hard stop + // can interrupt it. Returning `WorkCancelled` makes the job available + // again without using up its attempt. + tokio::select! { + () = context.cancellation_token().cancelled() => Err(WorkCancelled.into()), + () = tokio::time::sleep(Duration::from_mins(1)) => { + println!("generated report {}", job.args.report_id); + Ok(WorkOutcome::Complete) + } + } +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let pool = PgPool::connect(&std::env::var("DATABASE_URL")?).await?; + let mut workers = WorkerRegistry::new(); + workers.register_fn(generate_report)?; + let client = Client::builder(pool) + .workers(workers) + .queue("default", QueueConfig::new(10)) + // Cancel jobs still running 30 seconds after a graceful stop begins. + .soft_stop_timeout(Duration::from_secs(30)) + .build()?; + client.insert(GenerateReport { report_id: 42 }).await?; + + let mut run = client.start()?; + let stopper = run.stopper(); + tokio::spawn(async move { + if tokio::signal::ctrl_c().await.is_ok() { + println!("stopping; press Ctrl-C again to cancel running jobs"); + stopper.stop(); + } + if tokio::signal::ctrl_c().await.is_ok() { + stopper.stop_now(); + } + }); + + // Returns once the client has stopped and recorded every job's result. + run.wait().await?; + Ok(()) +} diff --git a/rust/riverqueue/examples/mixed_go_rust.rs b/rust/riverqueue/examples/mixed_go_rust.rs new file mode 100644 index 000000000..9ad10ff70 --- /dev/null +++ b/rust/riverqueue/examples/mixed_go_rust.rs @@ -0,0 +1,114 @@ +//! Shares one database between a Go service and a Rust service. +//! +//! This Rust service works `resize_image` jobs, which a Go service inserts, +//! and inserts `send_receipt` jobs, which only the Go service works. Each +//! language fetches from its own queue, so neither claims a kind it can't run. +//! +//! The Go side declares the same kinds, JSON field names, and queues: +//! +//! ```go +//! type ResizeImageArgs struct { +//! Path string `json:"path"` +//! Width int `json:"width"` +//! } +//! +//! func (ResizeImageArgs) Kind() string { return "resize_image" } +//! +//! func (ResizeImageArgs) InsertOpts() river.InsertOpts { +//! return river.InsertOpts{Queue: "rust_images"} +//! } +//! +//! type SendReceiptArgs struct { +//! OrderID int64 `json:"order_id"` +//! } +//! +//! func (SendReceiptArgs) Kind() string { return "send_receipt" } +//! +//! // Go works "default" (including send_receipt) and inserts resize_image: +//! client, _ := river.NewClient(riverpgxv5.New(pool), &river.Config{ +//! Queues: map[string]river.QueueConfig{river.QueueDefault: {MaxWorkers: 10}}, +//! Workers: workers, // registers a SendReceiptArgs worker +//! }) +//! client.Insert(ctx, ResizeImageArgs{Path: "cat.jpg", Width: 640}, nil) +//! ``` +//! +//! Run migrations once, with either implementation, then start both services: +//! +//! ```sh +//! DATABASE_URL=postgres://localhost/river_example cargo run -p riverqueue --example mixed_go_rust +//! ``` +//! +//! The mixed deployment guide (`riverqueue::guide::mixed_deployments`) covers +//! version matching, unique jobs, and rolling deployment. + +use std::error::Error; + +use riverqueue::sqlx::PgPool; +use riverqueue::{ + Client, EventKind, InsertOpts, Job, JobArgs, QueueConfig, WorkContext, WorkOutcome, + WorkerRegistry, migrate::PostgresMigrator, +}; +use serde::{Deserialize, Serialize}; + +/// Inserted by Go, worked here. Field names match the Go struct's JSON tags. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "resize_image", queue = "rust_images")] +struct ResizeImage { + path: String, + width: u32, +} + +/// Inserted here, worked by Go in its `default` queue. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "send_receipt")] +struct SendReceipt { + order_id: i64, +} + +async fn resize_image( + context: WorkContext, + job: Job, +) -> Result { + println!("resizing {} to {}px", job.args.path, job.args.width); + // Enqueue follow-up work for the Go service from inside a Rust worker. + if let Some(client) = context.client() { + client + .insert(SendReceipt { + order_id: job.row.id, + }) + .await?; + } + Ok(WorkOutcome::Complete) +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let pool = PgPool::connect(&std::env::var("DATABASE_URL")?).await?; + PostgresMigrator::new(pool.clone()).migrate_up().await?; + + let mut workers = WorkerRegistry::new(); + workers.register_fn(resize_image)?; + let client = Client::builder(pool) + .workers(workers) + // Only Rust's queue: Rust never fetches the Go service's jobs. + .queue("rust_images", QueueConfig::new(4)) + // Allow inserting kinds only Go works, like Go's SkipUnknownJobCheck. + .allow_unregistered_job_kinds() + .build()?; + let mut completed = client.subscribe(&[EventKind::JobCompleted])?; + let mut run = client.start()?; + + // Stand in for the Go producer so the example runs on its own. + let inserted = client + .insert(ResizeImage { + path: "cat.jpg".to_owned(), + width: 640, + }) + .opts(InsertOpts::default()) + .await?; + while completed.recv().await?.as_job().map(|event| event.job.id) != Some(inserted.id()) {} + println!("send_receipt is waiting in the default queue for the Go service"); + + run.shutdown().await?; + Ok(()) +} diff --git a/rust/riverqueue/examples/periodic.rs b/rust/riverqueue/examples/periodic.rs new file mode 100644 index 000000000..34a0baa7e --- /dev/null +++ b/rust/riverqueue/examples/periodic.rs @@ -0,0 +1,84 @@ +//! Enqueues jobs on a schedule. +//! +//! ```sh +//! DATABASE_URL=postgres://localhost/river_example cargo run -p riverqueue --example periodic +//! ``` +//! +//! Only the elected leader enqueues periodic jobs. Configure the same periodic +//! jobs, with the same IDs, in every client that may become leader. + +use std::{error::Error, time::Duration}; + +use riverqueue::sqlx::PgPool; +use riverqueue::{ + Client, CronSchedule, EventKind, IntervalSchedule, Job, JobArgs, PeriodicJob, PeriodicJobOpts, + QueueConfig, WorkContext, WorkOutcome, WorkerRegistry, migrate::PostgresMigrator, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "refresh_cache")] +struct RefreshCache {} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "nightly_report")] +struct NightlyReport {} + +async fn refresh_cache( + _context: WorkContext, + _job: Job, +) -> Result { + println!("refreshing cache"); + Ok(WorkOutcome::Complete) +} + +async fn nightly_report( + _context: WorkContext, + _job: Job, +) -> Result { + println!("building nightly report"); + Ok(WorkOutcome::Complete) +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let pool = PgPool::connect(&std::env::var("DATABASE_URL")?).await?; + PostgresMigrator::new(pool.clone()).migrate_up().await?; + + let mut workers = WorkerRegistry::new(); + workers.register_fn(refresh_cache)?; + workers.register_fn(nightly_report)?; + let client = Client::builder(pool) + .workers(workers) + .queue("default", QueueConfig::new(4)) + // Every 15 minutes, and once as soon as this client becomes leader. + .periodic_job(PeriodicJob::with_options( + IntervalSchedule::new(Duration::from_mins(15))?, + || RefreshCache {}, + PeriodicJobOpts::new() + .with_id("refresh_cache") + .with_run_on_start(true), + )) + // Standard five-field cron syntax, as in River Go: 02:30 UTC daily. + .periodic_job(PeriodicJob::with_options( + CronSchedule::parse("30 2 * * *")?, + || NightlyReport {}, + PeriodicJobOpts::new().with_id("nightly_report"), + )) + .build()?; + let mut completed = client.subscribe(&[EventKind::JobCompleted])?; + let mut run = client.start()?; + + // Wait for the run-on-start job, then stop. + loop { + let event = completed.recv().await?; + if event + .as_job() + .is_some_and(|event| event.job.kind == RefreshCache::KIND) + { + break; + } + } + run.shutdown().await?; + Ok(()) +} diff --git a/rust/riverqueue/examples/sqlite.rs b/rust/riverqueue/examples/sqlite.rs new file mode 100644 index 000000000..5d5970bc6 --- /dev/null +++ b/rust/riverqueue/examples/sqlite.rs @@ -0,0 +1,74 @@ +//! Runs River on SQLite: migrate a database file, work a job, and stop. +//! +//! ```sh +//! cargo run -p riverqueue --example sqlite --features sqlite +//! ``` +//! +//! Set `SQLITE_PATH` to use a specific database file; otherwise the example +//! uses a temporary one. + +use std::{error::Error, str::FromStr, time::Duration}; + +use riverqueue::sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}; +use riverqueue::{ + Client, EventKind, Job, JobArgs, QueueConfig, WorkContext, WorkOutcome, WorkerRegistry, + migrate::SqliteMigrator, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "resize_image")] +struct ResizeImage { + path: String, + width: u32, +} + +async fn resize_image( + _context: WorkContext, + job: Job, +) -> Result { + println!("resizing {} to {}px", job.args.path, job.args.width); + Ok(WorkOutcome::Complete) +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let directory = std::env::temp_dir().join(format!("river-sqlite-{}", std::process::id())); + std::fs::create_dir_all(&directory)?; + let path = std::env::var("SQLITE_PATH") + .unwrap_or_else(|_| directory.join("river.sqlite3").display().to_string()); + + // Every process sharing a SQLite database needs WAL mode and a busy + // timeout so readers and the single writer don't fail each other. + let options = SqliteConnectOptions::from_str(&format!("sqlite://{path}"))? + .create_if_missing(true) + .journal_mode(SqliteJournalMode::Wal) + .busy_timeout(Duration::from_secs(5)); + let pool = SqlitePoolOptions::new().connect_with(options).await?; + + // Apply River's schema before starting a client. + SqliteMigrator::new(pool.clone()).migrate_up().await?; + + let mut workers = WorkerRegistry::new(); + workers.register_fn(resize_image)?; + let client = Client::builder(pool.clone()) + .workers(workers) + .queue("default", QueueConfig::new(4)) + .build()?; + let mut completed = client.subscribe(&[EventKind::JobCompleted])?; + let mut run = client.start()?; + + let inserted = client + .insert(ResizeImage { + path: "photos/cat.jpg".to_owned(), + width: 640, + }) + .await?; + while completed.recv().await?.as_job().map(|event| event.job.id) != Some(inserted.id()) {} + println!("job {} completed", inserted.id()); + + run.shutdown().await?; + pool.close().await; + std::fs::remove_dir_all(&directory).ok(); + Ok(()) +} diff --git a/rust/riverqueue/examples/transactions.rs b/rust/riverqueue/examples/transactions.rs new file mode 100644 index 000000000..f873fe62f --- /dev/null +++ b/rust/riverqueue/examples/transactions.rs @@ -0,0 +1,96 @@ +//! Enqueues and completes jobs in the same transactions as business writes. +//! +//! ```sh +//! DATABASE_URL=postgres://localhost/river_example cargo run -p riverqueue --example transactions +//! ``` +//! +//! An order and the job that confirms it are inserted in one transaction, so +//! neither exists without the other. The worker records the confirmation and +//! completes its job in one transaction, so a crash between the two can't +//! send a second confirmation. + +use std::error::Error; + +use riverqueue::sqlx::{self, PgPool}; +use riverqueue::{ + Client, EventKind, Job, JobArgs, QueueConfig, WorkContext, WorkOutcome, WorkerRegistry, + migrate::PostgresMigrator, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "confirm_order")] +struct ConfirmOrder { + order_id: i64, +} + +async fn confirm_order( + context: WorkContext, + job: Job, +) -> Result { + let client = context + .client() + .expect("jobs worked by a client have a client"); + let pool = client + .database() + .as_postgres() + .expect("this example uses PostgreSQL") + .pool() + .clone(); + + // River drops a worker that ignores cancellation for too long, so begin + // with River's helper, which can't be abandoned halfway. + let mut transaction = riverqueue::database::begin_postgres(&pool).await?; + sqlx::query("UPDATE example_orders SET confirmed = true WHERE id = $1") + .bind(job.args.order_id) + .execute(&mut *transaction) + .await?; + // The job completes only if this transaction commits. + context.job_complete_tx(&mut transaction).await?; + transaction.commit().await?; + println!("confirmed order {}", job.args.order_id); + Ok(WorkOutcome::Complete) +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let pool = PgPool::connect(&std::env::var("DATABASE_URL")?).await?; + PostgresMigrator::new(pool.clone()).migrate_up().await?; + sqlx::query( + "CREATE TABLE IF NOT EXISTS example_orders \ + (id bigserial PRIMARY KEY, confirmed boolean NOT NULL DEFAULT false)", + ) + .execute(&pool) + .await?; + + let mut workers = WorkerRegistry::new(); + workers.register_fn(confirm_order)?; + let client = Client::builder(pool.clone()) + .workers(workers) + .queue("default", QueueConfig::new(4)) + .build()?; + let mut completed = client.subscribe(&[EventKind::JobCompleted])?; + let mut run = client.start()?; + + // Insert the order and its job together. + let mut transaction = riverqueue::database::begin_postgres(&pool).await?; + let order_id: i64 = + sqlx::query_scalar("INSERT INTO example_orders DEFAULT VALUES RETURNING id") + .fetch_one(&mut *transaction) + .await?; + let inserted = client + .insert(ConfirmOrder { order_id }) + .tx(&mut transaction) + .await?; + transaction.commit().await?; + + while completed.recv().await?.as_job().map(|event| event.job.id) != Some(inserted.id()) {} + let confirmed: bool = sqlx::query_scalar("SELECT confirmed FROM example_orders WHERE id = $1") + .bind(order_id) + .fetch_one(&pool) + .await?; + assert!(confirmed); + + run.shutdown().await?; + Ok(()) +} diff --git a/rust/riverqueue/examples/unique.rs b/rust/riverqueue/examples/unique.rs new file mode 100644 index 000000000..001588d89 --- /dev/null +++ b/rust/riverqueue/examples/unique.rs @@ -0,0 +1,85 @@ +//! Deduplicates jobs with unique options. +//! +//! ```sh +//! DATABASE_URL=postgres://localhost/river_example cargo run -p riverqueue --example unique +//! ``` + +use std::{error::Error, time::Duration}; + +use riverqueue::sqlx::PgPool; +use riverqueue::{Client, InsertOpts, JobArgs, UniqueOpts, migrate::PostgresMigrator}; +use serde::{Deserialize, Serialize}; + +/// Unique by default: at most one reconciliation per account per hour. +/// +/// `#[river(unique)]` selects the arguments that identify a duplicate, so +/// jobs for the same account collide even when `requested_by` differs. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "reconcile_account", unique(by_args, by_period = "1h"))] +struct ReconcileAccount { + #[river(unique)] + account_id: i64, + requested_by: String, +} + +/// Not unique by default; individual insertions opt in with `InsertOpts`. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "send_digest")] +struct SendDigest { + user_id: i64, +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let pool = PgPool::connect(&std::env::var("DATABASE_URL")?).await?; + PostgresMigrator::new(pool.clone()).migrate_up().await?; + // An insert-only client needs no workers or queues. + let client = Client::builder(pool).build()?; + + let first = client + .insert(ReconcileAccount { + account_id: 42, + requested_by: "billing".to_owned(), + }) + .await?; + let second = client + .insert(ReconcileAccount { + account_id: 42, + requested_by: "support".to_owned(), + }) + .await?; + // The second insertion returns the existing job instead of a new one. + assert!(second.unique_skipped_as_duplicate); + assert_eq!(first.id(), second.id()); + println!("reconcile_account deduplicated to job {}", first.id()); + + let opts = || { + InsertOpts::default().with_unique( + UniqueOpts::new() + .with_by_args(true) + .with_by_period(Duration::from_hours(24)), + ) + }; + let digest = client + .insert(SendDigest { user_id: 7 }) + .opts(opts()) + .await?; + let repeat = client + .insert(SendDigest { user_id: 7 }) + .opts(opts()) + .await?; + assert!(repeat.unique_skipped_as_duplicate); + let other = client + .insert(SendDigest { user_id: 8 }) + .opts(opts()) + .await?; + // A different user is a different job. (Running the example again within + // a day finds the jobs from the previous run instead.) + assert_ne!(digest.id(), other.id()); + println!( + "send_digest: user 7 -> job {}, user 8 -> job {}", + digest.id(), + other.id() + ); + Ok(()) +} diff --git a/rust/riverqueue/src/__private.rs b/rust/riverqueue/src/__private.rs new file mode 100644 index 000000000..3fb858260 --- /dev/null +++ b/rust/riverqueue/src/__private.rs @@ -0,0 +1,1431 @@ +//! Unstable extension points for River's own companion crates. +//! +//! Nothing in this module is part of River's public API. It changes without +//! notice between any two versions, so only crates released in lockstep with +//! `riverqueue` may use it. + +#![allow(missing_docs)] + +use std::{fmt, time::Duration}; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use serde_json::{Map, Value, value::RawValue}; + +#[cfg(feature = "postgres")] +use crate::database::SchemaName; +use crate::{AttemptError, InsertResult, Job, JobRow, JobState}; +#[cfg(feature = "postgres")] +use sqlx::{PgConnection, PgPool}; +#[cfg(feature = "sqlite")] +use sqlx::{SqliteConnection, SqlitePool}; +use tokio_util::sync::CancellationToken; + +pub use crate::client::{ExtensionClient, PreparedInsertRequest, RawInsertRequest, WeakClient}; +pub use crate::database::erased::{Database, ErasedExecutor, ErasedTransaction}; +pub use crate::pilot::{ + PilotDatabase, PilotProducer, PilotTransaction, ProducerClaimContext, ProducerClaimNext, + ProducerConfiguration, ProducerKeepAliveContext, ProducerShutdownContext, ProducerStartContext, +}; + +/// Insertion options reserved for River's own companion crates. +/// +/// Extension options travel beside a job's ordinary options, from a job +/// type's `JobArgs::default_insert_opts` or a call's options, to the +/// extension's [`Pilot::before_jobs_insert`] hook as +/// [`JobInsertParams::extension_options`]. River doesn't persist them. They +/// resolve key by key: a call's option replaces the job type's option with +/// the same key, and the job type's other options are kept, so an extension +/// can declare options for a job type that per-call options such as metadata +/// don't disturb. +pub trait InsertOptsExt: Sized { + /// Returns the extension options set on these insertion options. + fn extension_options(&self) -> &Map; + + /// Sets the extension option `key`, replacing any earlier value. + #[must_use] + fn with_extension_option(self, key: impl Into, value: Value) -> Self; +} + +impl InsertOptsExt for crate::InsertOpts { + fn extension_options(&self) -> &Map { + &self.extension_options + } + + fn with_extension_option(mut self, key: impl Into, value: Value) -> Self { + self.extension_options.insert(key.into(), value); + self + } +} + +/// Queue configuration reserved for River's own companion crates. +/// +/// Extension settings travel with a queue's configuration, through +/// `ClientBuilder::queue` or [`LocalQueues`](crate::LocalQueues), to the +/// extension's [`Pilot::validate_queue_settings`] and then its producer +/// session as [`ProducerConfiguration::settings`]. River doesn't persist +/// them. +pub trait QueueConfigExt: Sized { + /// Returns the extension settings on this configuration. + fn extension_settings(&self) -> &Map; + + /// Sets the extension setting `key`, replacing any earlier value. + #[must_use] + fn with_extension_setting(self, key: impl Into, value: Value) -> Self; +} + +impl QueueConfigExt for crate::QueueConfig { + fn extension_settings(&self) -> &Map { + &self.extension_settings + } + + fn with_extension_setting(mut self, key: impl Into, value: Value) -> Self { + self.extension_settings.insert(key.into(), value); + self + } +} + +/// Builder operations reserved for River's own companion crates. +pub trait ClientBuilderExt: Sized { + /// Returns whether the client will stay out of leader election, as set + /// by `ClientBuilder::without_leader_election`. + /// + /// Such a client never runs [`Pilot::maintenance_services`], and River + /// rejects its own periodic jobs when it's built. A companion crate that + /// configures leader-owned work of its own, such as additional periodic + /// jobs, should reject that configuration the same way. + fn leader_election_disabled(&self) -> bool; + + /// Installs a pilot from a companion crate. + #[must_use] + fn pilot(self, pilot: P) -> Self; + + /// Sets how often producers call [`PilotProducer::keep_alive`], 30 + /// seconds by default like River Go's `ProducerReportInterval`. It's a + /// control for tests, not a tuning option. + #[must_use] + fn producer_report_interval(self, interval: Duration) -> Self; +} + +impl ClientBuilderExt for crate::ClientBuilder { + fn leader_election_disabled(&self) -> bool { + self.leader_election_disabled + } + + fn pilot(self, pilot: P) -> Self { + self.with_pilot(pilot) + } + + fn producer_report_interval(mut self, interval: Duration) -> Self { + self.producer_report_interval = interval; + self + } +} + +/// Formats an error and its sources the way River records a job's error, +/// `outer: inner`, with a message that repeats its source's shortened. +/// +/// Add-on crates use it for error text they persist themselves, so it reads +/// the same as the errors River records. +#[must_use] +pub fn error_chain(error: &(dyn std::error::Error + 'static)) -> String { + crate::error::Chain(error).to_string() +} + +/// Decodes one persisted attempt error leniently, like River Go's driver +/// reads, so an element in a shape River doesn't write can't make its row +/// unreadable. [`AttemptError`]'s `Deserialize` is strict like Go's +/// `encoding/json`, so add-on crates reading `errors` from the database use +/// this instead. Only text that isn't valid JSON is an error. +pub fn attempt_error_from_json(json: &str) -> Result { + AttemptError::from_json_lenient(json) +} + +/// Decodes a persisted JSON array of attempt errors leniently, decoding each +/// element like [`attempt_error_from_json`]. `null` is empty, and anything +/// other than an array is an error. +pub fn attempt_errors_from_json(json: &str) -> Result, serde_json::Error> { + AttemptError::from_json_array_lenient(json) +} + +/// Encodes a UTC timestamp in River's canonical SQLite wire format. +/// +/// This keeps companion crates aligned with River and Go's +/// millisecond-rounded, timezone-free SQLite representation. +#[cfg(feature = "sqlite")] +#[must_use] +pub fn sqlite_timestamp(time: DateTime) -> String { + crate::database::sqlite::sqlite_time(time) +} + +/// Adds an add-on crate's indexes to PostgreSQL's default reindexer list. +/// +/// Names already in the list are skipped. A caller who chose index names +/// explicitly with `PostgresReindexConfig::with_index_names`, including an +/// empty list that disables the reindexer, keeps exactly that list. A +/// custom schedule or timeout alone still receives add-on indexes. SQLite +/// sources are returned unchanged. +#[cfg(feature = "postgres")] +#[must_use] +pub fn database_with_default_postgres_reindex_names( + mut database: Database, + names: impl IntoIterator>, +) -> Database { + database.extend_default_postgres_reindex_names(names); + database +} + +/// Creates a detached work context with no client. +#[must_use] +pub fn work_context(cancellation: CancellationToken) -> crate::WorkContext { + crate::WorkContext::new(cancellation) +} + +/// Creates a detached work context for a job, restoring its persisted +/// resumable metadata. +#[must_use] +pub fn work_context_for_job(job: &JobRow) -> crate::WorkContext { + crate::WorkContext::for_test_job(job) +} + +/// Claims an available job for `client` as a fetch would, marking it +/// running with a new attempt, and returns a work context for that attempt +/// whose [`WorkContext::client`](crate::WorkContext::client) is `client`. +/// +/// # Errors +/// +/// Returns [`Error::NotFound`](crate::Error::NotFound) for a missing job, an +/// invalid-job error when the job isn't available, and a database error when +/// the claim fails. +pub async fn claim_job_for_test( + client: &crate::Client, + id: i64, +) -> Result<(JobRow, crate::WorkContext), crate::Error> { + let inner = &client.inner; + let mut session = + crate::storage::Session::begin(&inner.database, crate::storage::Access::Transaction) + .await?; + let row = session.storage(inner).job_claim(id).await?; + session.commit().await?; + let context = crate::WorkContext::for_job( + client.clone(), + CancellationToken::new(), + row.id, + &row.metadata, + ); + Ok((row, context)) +} + +/// Returns a snapshot of metadata recorded during an attempt. +#[must_use] +pub fn work_context_metadata_updates(context: &crate::WorkContext) -> Map { + context.metadata_updates() +} + +/// Validates resumable checkpoint metadata before invoking user work. +/// +/// # Errors +/// +/// Returns the resumable metadata failure recorded for the attempt. +pub fn work_context_resumable_validate( + context: &crate::WorkContext, +) -> Result<(), crate::WorkError> { + context.resumable_validate() +} + +/// Resolves attempt-scoped resumable errors and metadata after user work. +pub fn work_context_resumable_finish( + context: &crate::WorkContext, + worker_failed: bool, +) -> Option { + context.resumable_finish(worker_failed) +} + +/// Notification topic for queue and job control messages. +pub const NOTIFICATION_TOPIC_CONTROL: &str = crate::protocol::NOTIFICATION_TOPIC_CONTROL; + +/// Notification topic for newly available jobs. +pub const NOTIFICATION_TOPIC_INSERT: &str = crate::protocol::NOTIFICATION_TOPIC_INSERT; + +/// Notification topic for leadership changes. +pub const NOTIFICATION_TOPIC_LEADERSHIP: &str = crate::protocol::NOTIFICATION_TOPIC_LEADERSHIP; + +/// A River notification topic, mirroring Go's `notifier.NotificationTopic`. +#[doc(hidden)] +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub enum NotificationTopic { + /// Queue and job control messages. + Control, + /// Newly available jobs. + Insert, + /// Leadership changes. + Leadership, +} + +impl NotificationTopic { + /// Returns the unqualified topic name. + #[must_use] + pub const fn as_str(self) -> &'static str { + match self { + Self::Control => NOTIFICATION_TOPIC_CONTROL, + Self::Insert => NOTIFICATION_TOPIC_INSERT, + Self::Leadership => NOTIFICATION_TOPIC_LEADERSHIP, + } + } +} + +impl fmt::Display for NotificationTopic { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(self.as_str()) + } +} + +/// Sends notifications on a caller's transaction connection, like Go's +/// `riverdriver.Executor.NotifyMany`. +/// +/// PostgreSQL issues `pg_notify` on the schema-qualified channel +/// (`.`, using `current_schema()` when no schema is +/// configured), so delivery happens only when the transaction commits. A +/// server without `LISTEN`/`NOTIFY`, like YugabyteDB by default, gets no +/// notifications; the configuration carries no detected capabilities, so +/// each call checks the server. SQLite appends rows to the durable +/// `river_notification` outbox that River clients poll. An empty payload list +/// does nothing. +/// +/// # Errors +/// +/// Returns an error when the connection and configuration name different +/// backends or when the database rejects the statement. +#[doc(hidden)] +pub async fn notify_many( + connection: DatabaseConnection<'_>, + database: &DatabaseConfig, + topic: NotificationTopic, + payloads: &[String], +) -> Result<(), PilotError> { + if payloads.is_empty() { + return Ok(()); + } + match (connection, database) { + #[cfg(feature = "postgres")] + (DatabaseConnection::Postgres(connection), DatabaseConfig::Postgres { schema }) => { + if !crate::database::postgres_capabilities::PostgresCapabilities::detect( + &mut *connection, + ) + .await? + .supports_listen_notify + { + return Ok(()); + } + sqlx::query( + "SELECT pg_notify(concat(coalesce($1::text, current_schema()), '.', $2::text), payload) \ + FROM unnest($3::text[]) AS payload", + ) + .bind(schema.as_deref()) + .bind(topic.as_str()) + .bind(payloads) + .execute(connection) + .await?; + Ok(()) + } + #[cfg(feature = "sqlite")] + (DatabaseConnection::Sqlite(connection), DatabaseConfig::Sqlite) => { + let mut query = sqlx::QueryBuilder::::new( + "INSERT INTO river_notification (payload, topic) ", + ); + query.push_values(payloads, |mut row, payload| { + row.push_bind(payload).push_bind(topic.as_str()); + }); + query.build().execute(connection).await?; + Ok(()) + } + #[allow(unreachable_patterns)] + (connection, database) => Err(format!( + "notification connection {:?} does not match database {:?}", + connection.kind(), + database.kind() + ) + .into()), + } +} + +/// Filters for [`delete_finalized_jobs`], mirroring the job cleaner's query. +/// +/// Each horizon deletes jobs in that state finalized before it; `None` keeps +/// jobs in that state, however old. +#[doc(hidden)] +#[derive(Clone, Debug, Default)] +#[non_exhaustive] +pub struct FinalizedJobDeleteParams { + /// Delete cancelled jobs finalized before this time. + pub cancelled_before: Option>, + /// Delete completed jobs finalized before this time. + pub completed_before: Option>, + /// Delete discarded jobs finalized before this time. + pub discarded_before: Option>, + /// Maximum jobs to delete, lowest IDs first. + pub limit: i64, + /// Queues whose jobs are kept. + pub queues_excluded: Vec, + /// When set, only jobs in these queues are deleted. + pub queues_included: Option>, +} + +impl FinalizedJobDeleteParams { + /// Creates filters that delete nothing until a horizon is set. + #[must_use] + pub fn new(limit: i64) -> Self { + Self { + limit, + ..Self::default() + } + } +} + +/// Deletes finalized jobs with River's job cleaner query, on a caller's +/// connection, and returns how many were deleted. +/// +/// Add-on crates use it for cleaner passes of their own, such as per-queue +/// retention, so their deletions match River's exactly, including keeping a +/// state whose horizon is `None` on every backend. It runs no timeout or +/// cancellation of its own. +/// +/// # Errors +/// +/// Returns an error when the connection and configuration name different +/// backends or when the database rejects the statement. +#[doc(hidden)] +pub async fn delete_finalized_jobs( + connection: DatabaseConnection<'_>, + database: &DatabaseConfig, + params: &FinalizedJobDeleteParams, +) -> Result { + match (connection, database) { + #[cfg(feature = "postgres")] + (DatabaseConnection::Postgres(connection), DatabaseConfig::Postgres { schema }) => Ok( + crate::maintenance::postgres_delete_finalized_jobs(connection, schema, params).await?, + ), + #[cfg(feature = "sqlite")] + (DatabaseConnection::Sqlite(connection), DatabaseConfig::Sqlite) => { + Ok(crate::maintenance::sqlite_delete_finalized_jobs(connection, params).await?) + } + #[allow(unreachable_patterns)] + (connection, database) => Err(format!( + "deletion connection {:?} does not match database {:?}", + connection.kind(), + database.kind() + ) + .into()), + } +} + +/// Error type used across the exact-version internal pilot seam. +pub type PilotError = Box; + +pub use crate::database::DatabaseKind; + +/// Backend configuration passed through River's exact-version extension seam. +#[doc(hidden)] +#[derive(Clone, Debug)] +#[non_exhaustive] +pub enum DatabaseConfig { + /// PostgreSQL backend configuration. + #[cfg(feature = "postgres")] + Postgres { schema: SchemaName }, + /// SQLite backend configuration. + #[cfg(feature = "sqlite")] + Sqlite, +} + +impl DatabaseConfig { + /// Returns the selected backend. + #[must_use] + pub const fn kind(&self) -> DatabaseKind { + match self { + #[cfg(feature = "postgres")] + Self::Postgres { .. } => DatabaseKind::Postgres, + #[cfg(feature = "sqlite")] + Self::Sqlite => DatabaseKind::Sqlite, + } + } + + /// Returns PostgreSQL's configured schema, if selected. + #[must_use] + #[cfg(feature = "postgres")] + pub const fn postgres_schema(&self) -> Option<&SchemaName> { + match self { + Self::Postgres { schema } => Some(schema), + #[cfg(feature = "sqlite")] + Self::Sqlite => None, + } + } +} + +/// Borrowed transaction connection passed to an exact-version extension. +#[doc(hidden)] +#[non_exhaustive] +pub enum DatabaseConnection<'connection> { + /// PostgreSQL transaction connection. + #[cfg(feature = "postgres")] + Postgres(&'connection mut PgConnection), + /// SQLite transaction connection. + #[cfg(feature = "sqlite")] + Sqlite(&'connection mut SqliteConnection), +} + +impl<'connection> DatabaseConnection<'connection> { + /// Returns the selected backend. + #[must_use] + pub const fn kind(&self) -> DatabaseKind { + match self { + #[cfg(feature = "postgres")] + Self::Postgres(_) => DatabaseKind::Postgres, + #[cfg(feature = "sqlite")] + Self::Sqlite(_) => DatabaseKind::Sqlite, + } + } + + /// Reborrows the connection for one operation, leaving this value + /// usable afterwards. + pub(crate) fn reborrow(&mut self) -> DatabaseConnection<'_> { + match self { + #[cfg(feature = "postgres")] + Self::Postgres(connection) => DatabaseConnection::Postgres(connection), + #[cfg(feature = "sqlite")] + Self::Sqlite(connection) => DatabaseConnection::Sqlite(connection), + } + } + + /// Returns the PostgreSQL connection, if selected. + #[must_use] + #[cfg(feature = "postgres")] + pub fn into_postgres(self) -> Option<&'connection mut PgConnection> { + match self { + Self::Postgres(connection) => Some(connection), + #[cfg(feature = "sqlite")] + Self::Sqlite(_) => None, + } + } + + /// Returns the SQLite connection, if selected. + #[must_use] + #[cfg(feature = "sqlite")] + pub fn into_sqlite(self) -> Option<&'connection mut SqliteConnection> { + match self { + #[cfg(feature = "postgres")] + Self::Postgres(_) => None, + Self::Sqlite(connection) => Some(connection), + } + } +} + +impl fmt::Debug for DatabaseConnection<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("DatabaseConnection") + .field("kind", &self.kind()) + .finish_non_exhaustive() + } +} + +/// Caller-owned pool passed to an exact-version background service. +#[doc(hidden)] +#[derive(Clone)] +#[non_exhaustive] +pub enum DatabasePool { + /// PostgreSQL pool. + #[cfg(feature = "postgres")] + Postgres(PgPool), + /// SQLite pool. + #[cfg(feature = "sqlite")] + Sqlite(SqlitePool), +} + +impl DatabasePool { + /// Returns the selected backend. + #[must_use] + pub const fn kind(&self) -> DatabaseKind { + match self { + #[cfg(feature = "postgres")] + Self::Postgres(_) => DatabaseKind::Postgres, + #[cfg(feature = "sqlite")] + Self::Sqlite(_) => DatabaseKind::Sqlite, + } + } + + /// Returns the caller-owned PostgreSQL pool, if selected. + #[must_use] + #[cfg(feature = "postgres")] + pub const fn postgres(&self) -> Option<&PgPool> { + match self { + Self::Postgres(pool) => Some(pool), + #[cfg(feature = "sqlite")] + Self::Sqlite(_) => None, + } + } + + /// Returns the caller-owned SQLite pool, if selected. + #[must_use] + #[cfg(feature = "sqlite")] + pub const fn sqlite(&self) -> Option<&SqlitePool> { + match self { + #[cfg(feature = "postgres")] + Self::Postgres(_) => None, + Self::Sqlite(pool) => Some(pool), + } + } +} + +impl fmt::Debug for DatabasePool { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("DatabasePool") + .field("kind", &self.kind()) + .finish_non_exhaustive() + } +} + +/// Inputs available while selecting stuck jobs under a rescue transaction. +/// +/// Mirrors Go's `JobGetStuckParams`: selections page by ID after `after_id` +/// and consider only jobs attempted before `stuck_horizon`. +#[derive(Clone, Debug)] +pub struct RescueParams { + /// Only jobs with a greater ID belong to this batch. + pub after_id: i64, + /// Selected database backend configuration. + pub database: DatabaseConfig, + /// Maximum rows to select. + pub maximum: i64, + /// Age at which the OSS runtime considers a running job stuck. + pub rescue_after: Duration, + /// Jobs attempted at or after this time are not stuck. Computed once per + /// rescuer pass. + pub stuck_horizon: DateTime, + /// The limit on the rescuer transaction the selection runs in. An + /// extension that reads through a connection of its own should bound + /// that read the same way. + pub timeout: Duration, +} + +/// One stuck job's transition exactly as the OSS rescuer would persist it. +#[derive(Clone, Debug)] +pub struct RescueJob { + /// Attempt error JSON appended to the job's `errors`. + pub attempt_error: Value, + /// Finalization time, set for `cancelled` and `discarded`. + pub finalized_at: Option>, + /// Job ID. + pub id: i64, + /// Next scheduled time. + pub scheduled_at: DateTime, + /// Target River state string. + pub state: JobState, +} + +/// Inputs of a batched rescue, mirroring Go's `JobRescueManyParams`. +/// +/// OSS only applies each transition to jobs still `running` with +/// `attempted_at` before `stuck_horizon`, so a job completed or claimed again +/// after selection is left untouched. Implementations that handle the rescue +/// themselves should apply the same guard. +#[derive(Clone, Debug)] +pub struct RescueManyParams { + /// Selected database backend configuration. + pub database: DatabaseConfig, + /// Transitions OSS would write, in ID order. + pub jobs: Vec, + /// Horizon the batch was selected with. + pub stuck_horizon: DateTime, +} + +/// Whether the OSS rescuer should perform its normal guarded update. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum RescueAction { + /// Continue through the OSS rescue update. + #[default] + Continue, + /// The extension persisted the rescue itself. + Handled, +} + +/// A job River just cancelled or retried, passed to extension post-hooks in +/// the same transaction as the update. +#[derive(Clone, Debug)] +pub struct JobUpdatedParams { + /// Selected database backend configuration. + pub database: DatabaseConfig, + /// The job after the update. + pub job: JobRow, +} + +/// Rows passed to [`Pilot::after_jobs_set_state`]. +#[derive(Clone, Debug)] +pub struct JobSetStateParams<'a> { + /// Selected database backend configuration. + pub database: DatabaseConfig, + /// The ID of every job in the batch, including jobs deleted while their + /// workers ran, which have no row in `jobs`. Per-attempt resources, such + /// as running counts, are released by [`PilotProducer::job_finished`] + /// instead, which doesn't wait for persistence. + pub job_ids: &'a [i64], + /// Every job in the batch that still exists, as returned by the update, + /// including jobs that were no longer running and so kept their state. + pub jobs: &'a [JobRow], +} + +/// Rows passed to [`Pilot::after_jobs_inserted`]. +#[derive(Clone, Debug)] +pub struct JobsInsertedParams<'a> { + /// Selected database backend configuration. + pub database: DatabaseConfig, + /// Jobs the insertion wrote, excluding unique insertions skipped as + /// duplicates, in input order. + pub jobs: &'a [JobRow], +} + +/// Mutable job insertion fields exposed to an exact-version extension. +/// +/// The references point into River's resolved insertion context. Changes are +/// validated and persisted by the ordinary insertion pipeline after the +/// extension returns. +#[doc(hidden)] +pub struct JobInsertParams<'insert> { + /// Serialized job arguments as exact JSON text. + pub encoded_args: &'insert mut Box, + /// Extension options resolved from the job type's and the call's + /// [`InsertOptsExt`] options, keyed by extension. River doesn't persist + /// them. + pub extension_options: &'insert Map, + /// Stable job kind. + pub kind: &'insert mut String, + /// Arbitrary job metadata. + pub metadata: &'insert mut crate::JobMetadata, + /// Queue in which the job will run. + pub queue: &'insert mut String, + /// Initial state: available, pending, or scheduled. An extension may + /// insert a job as pending, like River Go's insert hooks setting + /// `JobInsertParams.State`. + pub state: &'insert mut JobState, +} + +impl fmt::Debug for JobInsertParams<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("JobInsertParams") + .field("kind", self.kind) + .field("queue", self.queue) + .finish_non_exhaustive() + } +} + +/// A leadership term, handed to [`MaintenanceService::run`]. +/// +/// `token` is cancelled the moment this client stops trusting its +/// leadership: when it resigns, when a renewal fails, or when the trust +/// deadline passes without a renewal. Cancellation is local: it can't fence +/// statements already sent to the database. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct LeaderTerm { + /// When the database recorded this client's election, which identifies + /// the term. + pub elected_at: DateTime, + /// Cancelled when the term ends. + pub token: CancellationToken, +} + +/// Inputs to [`MaintenanceService::run`]. +#[derive(Debug)] +#[non_exhaustive] +pub struct MaintenanceServiceContext { + /// The client running the service, without keeping it alive. + pub client: WeakClient, + /// The client's database. + pub database: PilotDatabase, + /// The leadership term the service runs in. + pub term: LeaderTerm, +} + +/// Inputs to [`RuntimeService::run`]. +#[derive(Debug)] +#[non_exhaustive] +pub struct RuntimeServiceContext { + /// Cancelled when the service should stop, which happens as soon as the + /// client starts stopping. + pub cancellation: CancellationToken, + /// The client running the service, without keeping it alive. + pub client: WeakClient, + /// The client's database. + pub database: PilotDatabase, +} + +/// A leader-owned service supplied by an exact-version extension. +/// +/// River runs each service for every leadership term this client holds and +/// supervises it within the term: a service that returns an error, panics, +/// or returns before its term ends is logged and started again after River's +/// service backoff, which starts over after two minutes of healthy running. +/// A term's services all return before the next term's start. +#[async_trait] +pub trait MaintenanceService: Send + Sync + 'static { + /// A name for the service in River's logs. + fn name(&self) -> &'static str { + "extension maintenance service" + } + + /// Runs until the term's token is cancelled. + /// + /// # Errors + /// + /// Returns an error when the service failed; River restarts it. + async fn run(&self, context: MaintenanceServiceContext) -> Result<(), PilotError>; +} + +/// Per-client service supplied by an exact-version extension. +/// +/// Unlike [`MaintenanceService`], a runtime service runs on every started +/// client rather than only while that client holds River leadership. River +/// starts runtime services before the client's producers, and restarts one +/// that fails, panics, or returns early after its service backoff. +#[async_trait] +pub trait RuntimeService: Send + Sync + 'static { + /// A name for the service in River's logs. + fn name(&self) -> &'static str { + "extension runtime service" + } + + /// Runs until the context's cancellation. + /// + /// # Errors + /// + /// Returns an error when the service failed; River restarts it. + async fn run(&self, context: RuntimeServiceContext) -> Result<(), PilotError>; +} + +/// What a pilot gets when its client is built, from +/// [`Pilot::install`]. +#[derive(Debug)] +#[non_exhaustive] +pub struct PilotInstallContext { + /// The client, without keeping it alive. + pub client: WeakClient, + /// The client's database, for the pilot's own statements. + pub database: PilotDatabase, + /// How often producers report to their sessions, which peers use to tell + /// when a producer has gone stale. + pub producer_report_interval: Duration, +} + +/// Exact-version extension seam for matched companion crates. +/// +/// This trait is intentionally not a stable River API. The internal crate is +/// version-locked to `riverqueue`, allowing the SPI to evolve with both +/// implementations. +#[async_trait] +pub trait Pilot: std::any::Any + Send + Sync + 'static { + /// Binds the pilot to the client being built, like River Go's + /// `PilotInit`. River calls it once per client, before the builder + /// returns the client; a pilot installed on several clients is called + /// once for each. The pilot keeps what it needs from `context` rather + /// than reading the client's public database accessors. + fn install(&self, _context: PilotInstallContext) {} + + /// Queues whose finalized jobs are owned by an extension-specific cleaner + /// and skipped by River's job cleaner, like Go's + /// `Pilot.JobCleanerQueuesExcluded`. Read on every cleaner pass. + fn job_cleaner_queue_exclusions(&self) -> Vec { + Vec::new() + } + + /// Whether job cancellation and retry must run + /// [`Pilot::after_job_cancel`] and [`Pilot::after_job_retry`]. Returning + /// `true` also makes pool-based cancel and retry use a transaction. + fn intercepts_job_cancel_retry(&self) -> bool { + false + } + + /// Whether stuck-job candidate selection must enter the exact-version + /// interception transaction. + fn intercepts_rescue(&self) -> bool { + false + } + + /// Whether job state transitions must run in a transaction that also + /// calls [`Pilot::after_jobs_set_state`], like River Go's + /// `Pilot.JobSetStateIfRunningMany`. + /// + /// Returning `false` keeps River's one-statement completion path. + /// Implementations that override `after_jobs_set_state` return `true`. + fn intercepts_job_set_state(&self) -> bool { + false + } + + /// How many intercepted completion batches may run concurrently, like + /// River Go's `PilotJobCompletionConcurrency`. + /// + /// River never exceeds its backend's own limit (two on PostgreSQL, one on + /// SQLite) and starts a second concurrent batch only when a full batch of + /// completions is waiting. The default allows one batch at a time. + fn job_set_state_concurrency(&self) -> usize { + 1 + } + + /// Whether inserts must enter the exact-version interception transaction. + /// + /// Returning `true` makes pool-based insertion acquire a transaction so + /// [`Pilot::before_job_insert`] can observe backend state on the same + /// connection as the eventual insert. + fn intercepts_insert(&self) -> bool { + false + } + + /// Mutates or validates every job of one insertion call at once, using + /// its transaction connection, like River Go's `Pilot.JobInsertMany` + /// receiving the whole batch. + /// + /// River invokes it once per insertion when it intercepts inserts, after + /// every job's begin hooks and before writing any job. Implementations + /// can share work across the batch, such as reading each distinct queue's + /// configuration once. The default calls [`Pilot::before_job_insert`] + /// for each job in order. + /// + /// `connection` is always the insertion's transaction: River's own, or + /// the caller's. Its owner commits or rolls it back, so implementations + /// must not end it, and must make related writes on it. + async fn before_jobs_insert( + &self, + mut connection: DatabaseConnection<'_>, + jobs: &mut [JobInsertParams<'_>], + ) -> Result<(), PilotError> { + for job in jobs { + self.before_job_insert(connection.reborrow(), job).await?; + } + Ok(()) + } + + /// Mutates or validates a resolved insertion using its transaction + /// connection. Called for each job by the default + /// [`Pilot::before_jobs_insert`]. + /// + /// Insert middleware wraps the whole step: River runs middleware, then + /// ordinary begin hooks, then this method, then the write. The insert and + /// its backend notification remain in the same transaction. In a + /// caller's transaction they all run directly in it without a savepoint, + /// so the caller rolls back when any of them fails. + async fn before_job_insert( + &self, + _connection: DatabaseConnection<'_>, + _params: &mut JobInsertParams<'_>, + ) -> Result<(), PilotError> { + Ok(()) + } + + /// Optionally selects stuck-job candidates, honoring the params' cursor + /// and horizon. Returned IDs are evaluated by the OSS rescuer in the same + /// transaction; returning `maximum` IDs asks for another batch. + async fn select_rescue_job_ids( + &self, + _connection: DatabaseConnection<'_>, + _params: &RescueParams, + ) -> Result>, PilotError> { + Ok(None) + } + + /// Optionally persists a batch of rescues in the rescuer's transaction. + /// + /// Called only when [`Pilot::intercepts_rescue`] returns `true`, after + /// OSS decided each selected job's transition. Returning + /// [`RescueAction::Continue`] lets OSS apply its guarded update. + async fn rescue_jobs( + &self, + _connection: DatabaseConnection<'_>, + _params: &RescueManyParams, + ) -> Result { + Ok(RescueAction::Continue) + } + + /// Runs after River cancels a job, in the same transaction and with the + /// updated row. Called only when [`Pilot::intercepts_job_cancel_retry`] + /// returns `true`; an error rolls back the cancellation. + async fn after_job_cancel( + &self, + _connection: DatabaseConnection<'_>, + _job: &JobUpdatedParams, + ) -> Result<(), PilotError> { + Ok(()) + } + + /// Runs after River retries a job, in the same transaction and with the + /// updated row. Called only when [`Pilot::intercepts_job_cancel_retry`] + /// returns `true`; an error rolls back the retry. + async fn after_job_retry( + &self, + _connection: DatabaseConnection<'_>, + _job: &JobUpdatedParams, + ) -> Result<(), PilotError> { + Ok(()) + } + + /// Observes a batch of job state transitions inside River's transaction. + /// + /// Called only when [`Pilot::intercepts_job_set_state`] returns `true`. + /// River keeps batching completions: each batch runs `BEGIN`, River's + /// set-state-if-running update (which returns full rows), this hook with + /// those rows, then `COMMIT`. The transactional `job_complete_tx` path + /// calls it with its one row inside the caller's transaction. + /// + /// The hook may write further state with the connection, including + /// deleting returned rows; River still reports events from the rows it + /// already holds. Returning an error rolls the batch back, and River + /// retries it like any other failed completion write. + async fn after_jobs_set_state( + &self, + _connection: DatabaseConnection<'_>, + _params: &JobSetStateParams, + ) -> Result<(), PilotError> { + Ok(()) + } + + /// Runs after River writes a batch of inserted jobs, inside the insertion + /// transaction, like the post-insert work in River Go's + /// `Pilot.JobInsertMany`. + /// + /// Called only when [`Pilot::intercepts_insert`] returns `true`, on every + /// insertion path, including batches, caller-managed transactions, and + /// periodic jobs. Unique insertions skipped as duplicates aren't + /// included. Returning an error rolls back the insertion. As in + /// [`Pilot::before_jobs_insert`], `connection` is the insertion's + /// transaction, which the extension must not end. + async fn after_jobs_inserted( + &self, + _connection: DatabaseConnection<'_>, + _params: &JobsInsertedParams<'_>, + ) -> Result<(), PilotError> { + Ok(()) + } + + /// Validates the extension settings of a queue's configuration, set with + /// [`QueueConfigExt::with_extension_setting`]. + /// + /// River calls it when a client is built and when a queue is added or + /// updated through [`LocalQueues`](crate::LocalQueues), before the + /// configuration takes effect. The default accepts only a configuration + /// without extension settings. + /// + /// # Errors + /// + /// Returns an error describing settings the extension doesn't accept. + fn validate_queue_settings( + &self, + queue: &str, + settings: &Map, + ) -> Result<(), PilotError> { + if settings.is_empty() { + return Ok(()); + } + Err(format!("queue {queue:?} has extension settings, but no extension accepts them").into()) + } + + /// Starts the extension's session for a new generation of a queue's + /// producer, like River Go's pilot `ProducerInit`, or returns `None` when + /// the extension doesn't take part in this queue's claims. + /// + /// River calls it once the queue's record exists and before the + /// producer's first claim. When it fails, River logs the error and + /// retries the producer's start with backoff. + /// + /// When the producer stops while this call is still running, River drops + /// its future. An extension that already created shared state by then, + /// such as a producer row, gets no session and so no + /// [`PilotProducer::shutdown`] for it; the same race exists between River + /// Go's `ProducerInit` and a stop. Peers must treat such state like that + /// of a client that exited, for example by letting it go stale. + async fn start_producer( + &self, + _context: ProducerStartContext, + ) -> Result>, PilotError> { + Ok(None) + } + + /// Leader-owned services contributed by the extension. River runs them + /// only while the client is leader, and never on a client built with + /// `ClientBuilder::without_leader_election`. + fn maintenance_services(&self) -> Vec> { + Vec::new() + } + + /// Per-client services contributed by the extension. River runs them on + /// every started client, including one without leader election. + fn runtime_services(&self) -> Vec> { + Vec::new() + } +} + +/// No-op pilot used by River OSS. +#[derive(Clone, Copy, Debug, Default)] +pub struct NoopPilot; + +impl Pilot for NoopPilot {} + +/// Columns River selects to decode a PostgreSQL job row, qualified by +/// `alias`, for use with [`decode_postgres_job_row`]. +#[cfg(feature = "postgres")] +#[must_use] +pub fn postgres_job_projection(alias: &str) -> String { + crate::client::job_projection(alias) +} + +/// Decodes a row selected with [`postgres_job_projection`] exactly as River +/// decodes its own rows. +/// +/// # Errors +/// +/// Returns an error when the row can't be decoded. +#[cfg(feature = "postgres")] +pub fn decode_postgres_job_row(row: &sqlx::postgres::PgRow) -> Result { + crate::client::decode_job_row(row).map_err(|undecodable| undecodable.error.into()) +} + +/// Decodes a claimed row selected with [`postgres_job_projection`] as far +/// as River can, for [`PilotProducer::claim`]. +#[cfg(feature = "postgres")] +#[must_use] +pub fn claimed_postgres_job(row: &sqlx::postgres::PgRow) -> ClaimedJob { + ClaimedJob::from_decoded(crate::client::decode_job_row(row)) +} + +/// Columns River selects to decode a SQLite job row, for use with +/// [`decode_sqlite_job_row`]. +#[cfg(feature = "sqlite")] +pub const SQLITE_JOB_COLUMNS: &str = crate::database::sqlite::JOB_COLUMNS; + +/// Decodes a row selected with [`SQLITE_JOB_COLUMNS`] exactly as River +/// decodes its own rows. +/// +/// # Errors +/// +/// Returns an error when the row can't be decoded. +#[cfg(feature = "sqlite")] +pub fn decode_sqlite_job_row(row: &sqlx::sqlite::SqliteRow) -> Result { + crate::database::sqlite::decode_job_row(row).map_err(|undecodable| undecodable.error.into()) +} + +/// Decodes a claimed row selected with [`SQLITE_JOB_COLUMNS`] as far as +/// River can, for [`PilotProducer::claim`]. +#[cfg(feature = "sqlite")] +#[must_use] +pub fn claimed_sqlite_job(row: &sqlx::sqlite::SqliteRow) -> ClaimedJob { + ClaimedJob::from_decoded(crate::database::sqlite::decode_job_row(row)) +} + +/// A job claimed by [`PilotProducer::claim`]. +/// +/// River works a decoded job normally. Like a row River claims itself, an +/// undecodable one isn't worked: its attempt fails with an error describing +/// the decode failure, before hooks or middleware run, and it's retried or +/// discarded through ordinary error handling. +#[derive(Debug)] +pub struct ClaimedJob(crate::client::DecodedJob); + +impl ClaimedJob { + pub(crate) const fn from_decoded(decoded: crate::client::DecodedJob) -> Self { + Self(decoded) + } + + /// Returns the decoded row, or `None` when some field couldn't be + /// decoded. + #[must_use] + pub fn job(&self) -> Option<&JobRow> { + self.0.as_ref().ok() + } + + /// Returns why the row couldn't be fully decoded, if it couldn't. + #[must_use] + pub fn decode_error(&self) -> Option<&str> { + self.0 + .as_ref() + .err() + .map(|undecodable| undecodable.error.as_str()) + } + + /// Whether `column` of a partly decoded row couldn't be decoded. + pub(crate) fn column_undecodable(&self, column: &str) -> bool { + self.0 + .as_ref() + .err() + .is_some_and(|undecodable| undecodable.columns.iter().any(|name| name == column)) + } + + /// Returns the claimed row, with any field that couldn't be decoded left + /// empty, unless not even the row's identity could be decoded. + pub(crate) fn row(&self) -> Option<&JobRow> { + match &self.0 { + Ok(job) => Some(job), + Err(undecodable) => undecodable.row.as_deref(), + } + } + + pub(crate) fn into_decoded(self) -> crate::client::DecodedJob { + self.0 + } +} + +impl From for ClaimedJob { + fn from(job: JobRow) -> Self { + Self(Ok(job)) + } +} + +/// Type-erased result returned by River's exact-version insertion seam. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct RawInsertResult { + /// Inserted job or the existing matching unique job. + pub job: JobRow, + /// Whether insertion was skipped because a unique job already existed. + pub unique_skipped_as_duplicate: bool, +} + +/// A stored job to insert again, such as one set aside and retried later, +/// with [`ExtensionClient::insert_prepared`]. +/// +/// River inserts it like any other job: insert middleware, begin hooks, the +/// extension's insertion step, and notifications all run once, and they see +/// the stored arguments and metadata. What they return is stored, so the +/// job keeps its identity only when every step leaves a stored job alone. +/// Unique-key calculation doesn't run: the job keeps its unique key and +/// states, creation time, schedule, and metadata, and gets a new ID. +/// `encoded_args` may be any JSON value, including an array or `null`. +#[derive(Clone, Debug)] +pub struct PreparedInsertParams { + /// Original creation time. + pub created_at: DateTime, + /// Serialized job arguments. + pub encoded_args: Box, + /// Stable job kind. + pub kind: String, + /// Maximum attempts, including the first. + pub max_attempts: i16, + /// Arbitrary job metadata. + pub metadata: crate::JobMetadata, + /// Priority from one through four. + pub priority: i16, + /// Queue in which the job runs. + pub queue: String, + /// Earliest time at which the reinserted job may run. + pub scheduled_at: DateTime, + /// Searchable tags. + pub tags: Vec, + /// Existing unique hash, if any. + pub unique_key: Option>, + /// Existing states in which the key is enforced, if any. + pub unique_states: Option>, +} + +/// Inputs to a peer claim's callback, from [`PeerAttempts::claim`]. +#[derive(Debug)] +#[non_exhaustive] +pub struct PeerClaimContext<'c> { + /// The coordinating attempt's cancellation token, cancelled by a hard + /// stop or a remote cancellation of the coordinator's job. A soft stop + /// leaves it alone. + pub cancellation: &'c CancellationToken, + /// This client's identifier, which claimed rows' `attempted_by` must + /// end with. + pub client_id: &'c str, + /// The claim's transaction, which River commits once the callback + /// returns and its rows pass River's checks. + pub connection: DatabaseConnection<'c>, + /// The client's database. + pub database: &'c PilotDatabase, +} + +/// An outcome for one peer, for [`PeerAttempts::complete`]. +#[derive(Debug)] +#[non_exhaustive] +pub struct PeerOutcome { + /// The peer, as [`PeerAttempts::claim`] returned it. + pub job: JobRow, + /// The peer's result, as a worker would return it. + pub result: Result, +} + +impl PeerOutcome { + /// Creates an outcome for `job`. + #[must_use] + pub const fn new(job: JobRow, result: Result) -> Self { + Self { job, result } + } +} + +/// The peers of a running attempt: jobs the attempt, their coordinator, +/// claims and completes alongside its own job, such as a group of related +/// jobs it works together. +/// +/// River owns each peer from the commit of the claim that took it until its +/// outcome persists. Peers take no producer slots and never reach +/// [`PilotProducer::job_finished`]. When the coordinator's attempt ends, +/// River refuses new peer operations, waits for those it accepted, and gives +/// every peer still without an outcome one before the coordinator's own: an +/// interruption when River stopped the coordinator, and a failure otherwise, +/// including when the coordinator's job was cancelled remotely. A peer stops +/// being owned when its outcome persists, before its event, so it can be +/// claimed again at once. +/// +/// A soft stop doesn't end peer operations. A coordinator keeps claiming and +/// completing peers after its producer stops fetching new jobs, until its +/// attempt ends, and the client's stop waits for the attempt and so for +/// every peer it claimed. A hard stop cancels the attempt, which ends its +/// claims. +#[derive(Clone, Copy, Debug)] +pub struct PeerAttempts<'a> { + context: &'a crate::WorkContext, +} + +impl<'a> PeerAttempts<'a> { + /// Returns the peers of the attempt `context` belongs to. + #[must_use] + pub const fn new(context: &'a crate::WorkContext) -> Self { + Self { context } + } + + fn attempt( + self, + ) -> Result< + ( + &'a crate::Client, + &'a std::sync::Arc, + ), + crate::Error, + > { + match (self.context.client(), self.context.peers()) { + (Some(client), Some(peers)) => Ok((client, peers)), + _ => Err(crate::Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "peer attempts", + }, + source: "peer operations require a running attempt".into(), + }), + } + } + + /// Claims peers with `run` in a transaction River opens and commits. + /// + /// `run` must claim on the context's connection, like a producer claim: + /// it moves rows to `running`, increments their attempt, and appends this + /// client to `attempted_by`, and builds each result with + /// [`claimed_postgres_job`] or [`claimed_sqlite_job`]. Before commit, + /// River rejects the whole claim when a row can't be identified, appears + /// twice, is the coordinator's own job, is already owned by an attempt + /// or worked by this client, is at an attempt this coordinator already + /// saw end, or isn't running under this client. A claim whose coordinator + /// is cancelled before commit, by a hard stop or a remote cancellation, + /// rolls back. A soft stop doesn't stop claims: the coordinator may keep + /// claiming until its attempt ends. River doesn't retry a failed claim. + /// + /// Returns the decoded rows River now tracks. A row that couldn't be + /// fully decoded is completed as a failure instead and not returned. + /// + /// # Errors + /// + /// Returns an [`Error::Extension`](crate::Error::Extension) error for a + /// claim River rejected, from `run`, once the coordinator's attempt was + /// cancelled, or once it ended, and a database error when the transaction + /// fails. + pub async fn claim(self, run: F) -> Result, crate::Error> + where + F: for<'c> FnOnce( + PeerClaimContext<'c>, + ) -> futures_util::future::BoxFuture< + 'c, + Result, PilotError>, + > + Send, + { + let (client, peers) = self.attempt()?; + peers.claim(&client.inner, self.context, run).await + } + + /// Completes peers through River's ordinary completion pipeline: the + /// error handler, the coordinator's recorded metadata, retry selection, + /// the extension's set-state step, events, and fenced persistence. + /// Returns once every outcome persisted. + /// + /// Outcomes are accepted all or none: each job must be a peer of this + /// attempt at the attempt it was claimed at, appear once, and have no + /// outcome yet. + /// + /// # Errors + /// + /// Returns an [`Error::Extension`](crate::Error::Extension) error for + /// outcomes River rejected or once the coordinator ended, and a runtime + /// error when an outcome couldn't be handed to the completer or wasn't + /// persisted. An outcome not handed over leaves its peer without one, so + /// River supplies one when the coordinator ends. + pub async fn complete(self, outcomes: Vec) -> Result<(), crate::Error> { + let (client, peers) = self.attempt()?; + peers.complete(&client.inner, self.context, outcomes).await + } +} + +impl RawInsertResult { + /// Converts an exact-version raw result after its arguments are decoded. + #[must_use] + pub fn into_typed(self, args: A) -> InsertResult { + InsertResult { + job: Job::new(args, self.job), + unique_skipped_as_duplicate: self.unique_skipped_as_duplicate, + } + } +} + +/// Complete persisted job fields for exact-version record conversion. +#[derive(Debug)] +pub struct JobRowParts { + pub id: i64, + pub attempt: i16, + pub attempted_at: Option>, + pub attempted_by: Vec, + pub created_at: DateTime, + pub encoded_args: Box, + pub errors: Vec, + pub finalized_at: Option>, + pub kind: String, + pub max_attempts: i16, + pub metadata: crate::JobMetadata, + pub priority: i16, + pub queue: String, + pub scheduled_at: DateTime, + pub state: JobState, + pub tags: Vec, + pub unique_key: Option>, + pub unique_states: Option>, +} + +impl JobRowParts { + /// Converts complete fields from an exact-version database record. + #[must_use] + pub fn into_row(self) -> JobRow { + let parts = self; + JobRow { + attempt: parts.attempt, + attempted_at: parts.attempted_at, + attempted_by: parts.attempted_by, + created_at: parts.created_at, + encoded_args: parts.encoded_args, + errors: parts.errors, + finalized_at: parts.finalized_at, + id: parts.id, + kind: parts.kind, + max_attempts: parts.max_attempts, + metadata: parts.metadata, + priority: parts.priority, + queue: parts.queue, + scheduled_at: parts.scheduled_at, + state: parts.state, + tags: parts.tags, + unique_key: parts.unique_key, + unique_states: parts.unique_states, + } + } +} diff --git a/rust/riverqueue/src/client/attempts.rs b/rust/riverqueue/src/client/attempts.rs new file mode 100644 index 000000000..7483935d5 --- /dev/null +++ b/rust/riverqueue/src/client/attempts.rs @@ -0,0 +1,123 @@ +//! Tracking of running job attempts for cancellation delivery. + +use std::collections::HashMap; +use std::sync::Mutex; +use std::sync::atomic::{AtomicU64, Ordering}; + +use tokio_util::sync::CancellationToken; + +use crate::client::{ClientInner, PENDING_CANCELLATION_LIMIT, PENDING_CANCELLATION_RETENTION}; + +pub(super) fn remove_running_attempt( + running: &Mutex>, + job_id: i64, + cancellation: &CancellationToken, +) { + let mut running = running + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if running + .get(&job_id) + .is_some_and(|active| active == cancellation) + { + running.remove(&job_id); + } +} + +pub(super) fn register_running_attempt( + running: &Mutex>, + pending_cancellations: &Mutex>, + job_id: i64, + cancellation: &CancellationToken, +) { + // Keep the locks in this order here and in `signal_running_attempt` so a + // cancellation cannot fall between checking the active map and recording + // a just-fetched attempt. + let mut running = running + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + running.insert(job_id, cancellation.clone()); + let should_cancel = pending_cancellations + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .remove(&job_id) + .is_some(); + drop(running); + if should_cancel { + cancellation.cancel(); + } +} + +pub(super) fn signal_running_attempt( + running: &Mutex>, + pending_cancellations: &Mutex>, + fetch_registration_windows: &AtomicU64, + job_id: i64, +) { + let running = running + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(cancellation) = running.get(&job_id).cloned() { + drop(running); + cancellation.cancel(); + return; + } + if fetch_registration_windows.load(Ordering::SeqCst) == 0 { + return; + } + + let now = std::time::Instant::now(); + let mut pending = pending_cancellations + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + pending.retain(|_, received_at| { + now.saturating_duration_since(*received_at) <= PENDING_CANCELLATION_RETENTION + }); + if pending.len() >= PENDING_CANCELLATION_LIMIT + && let Some(oldest_job_id) = pending + .iter() + .min_by_key(|(_, received_at)| **received_at) + .map(|(job_id, _)| *job_id) + { + pending.remove(&oldest_job_id); + } + pending.insert(job_id, now); +} + +pub(super) struct FetchRegistrationGuard<'a> { + pub(super) inner: &'a ClientInner, +} + +impl<'a> FetchRegistrationGuard<'a> { + pub(super) fn new(inner: &'a ClientInner) -> Self { + inner + .fetch_registration_windows + .fetch_add(1, Ordering::SeqCst); + Self { inner } + } +} + +impl Drop for FetchRegistrationGuard<'_> { + fn drop(&mut self) { + // Synchronize the last-window transition with + // `signal_running_attempt`, which holds this lock while deciding + // whether to retain an unmatched cancellation. + let _running = self + .inner + .running + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if self + .inner + .fetch_registration_windows + .fetch_sub(1, Ordering::SeqCst) + == 1 + { + self.inner + .pending_cancellations + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clear(); + } + } +} diff --git a/rust/riverqueue/src/client/backoff.rs b/rust/riverqueue/src/client/backoff.rs new file mode 100644 index 000000000..14d0be035 --- /dev/null +++ b/rust/riverqueue/src/client/backoff.rs @@ -0,0 +1,124 @@ +//! Backoff shared by runtime services that retry database operations. + +use std::{ + collections::HashMap, + hash::{BuildHasher, Hash, Hasher}, + time::Duration, +}; + +/// Attempts after which the exponential sequence starts over, matching River +/// Go's `serviceutil.MaxAttemptsBeforeResetDefault`. +const MAX_ATTEMPTS_BEFORE_RESET: u32 = 7; + +/// How long a restarted service or producer must run before a failure counts +/// as a new outage whose backoff starts over, rather than as another failure +/// in a row. It exceeds the longest restart backoff (about 70 seconds). +pub(crate) const SERVICE_RESTART_RESET_AFTER: Duration = Duration::from_mins(2); + +/// Consecutive failures of each restartable service or queue producer. +#[derive(Debug)] +pub(super) struct RestartBackoff { + attempts: HashMap, +} + +impl Default for RestartBackoff { + fn default() -> Self { + Self { + attempts: HashMap::new(), + } + } +} + +impl RestartBackoff { + /// Returns the one-based restart attempt recorded for `key`, if it has + /// failed since it was last forgotten. + pub(super) fn attempt(&self, key: &K) -> Option { + self.attempts.get(key).copied() + } + + /// Records a failure of `key` after it ran for `ran_for`, returning the + /// one-based restart attempt and the backoff before it. Like River Go's + /// services, which reset their error counts once they succeed, one that + /// ran for a while before failing starts its backoff over. + pub(super) fn failed(&mut self, key: K, ran_for: Duration) -> (u32, Duration) { + let attempt = self.attempts.entry(key).or_default(); + if ran_for >= SERVICE_RESTART_RESET_AFTER { + *attempt = 0; + } + *attempt += 1; + (*attempt, exponential_backoff(*attempt)) + } + + /// Forgets the failures of `key`. + pub(super) fn forget(&mut self, key: &K) { + self.attempts.remove(key); + } +} + +/// Returns River's service backoff for a one-based attempt: `2^(attempt - 1)` +/// seconds with ±10% jitter, restarting the sequence every seven attempts so a +/// long outage never sleeps for more than about a minute. +/// +/// This mirrors River Go's `serviceutil.ExponentialBackoff`, which the +/// notifier and completer use. It is intentionally distinct from the job retry +/// policy: services should recover promptly once the database returns. +pub(super) fn exponential_backoff(attempt: u32) -> Duration { + let exponent = attempt.saturating_sub(1) % MAX_ATTEMPTS_BEFORE_RESET; + let seconds = f64::from(1_u32 << exponent); + Duration::from_secs_f64(seconds + seconds * (jitter_unit() * 0.2 - 0.1)) +} + +/// Returns a uniformly distributed value in `[0, 1)` for jitter. +/// +/// Jitter only needs to decorrelate clients, so the standard library's +/// randomly keyed hasher avoids a dedicated random number dependency. +fn jitter_unit() -> f64 { + let mut hasher = std::collections::hash_map::RandomState::new().build_hasher(); + hasher.write_u128( + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(), + ); + #[allow( + clippy::cast_precision_loss, + reason = "53 random bits are plenty for jitter" + )] + let unit = (hasher.finish() >> 11) as f64 / (1_u64 << 53) as f64; + unit +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn exponential_backoff_doubles_with_jitter_and_resets() { + for (attempt, base_seconds) in [(0, 1.0), (1, 1.0), (2, 2.0), (3, 4.0), (7, 64.0), (8, 1.0)] + { + let backoff = exponential_backoff(attempt).as_secs_f64(); + assert!( + (base_seconds * 0.9..=base_seconds * 1.1).contains(&backoff), + "attempt {attempt} slept {backoff}s, expected about {base_seconds}s" + ); + } + } + + #[test] + fn restart_backoff_starts_over_after_a_long_healthy_run() { + let mut restarts = RestartBackoff::default(); + let quick = Duration::from_secs(1); + assert_eq!(restarts.failed("notifier", quick).0, 1); + assert_eq!(restarts.failed("notifier", quick).0, 2); + assert_eq!(restarts.failed("maintenance", quick).0, 1); + assert_eq!(restarts.failed("notifier", quick).0, 3); + assert_eq!(restarts.attempt(&"notifier"), Some(3)); + // A failure after a healthy run is the start of a new outage. + let (attempt, delay) = restarts.failed("notifier", SERVICE_RESTART_RESET_AFTER); + assert_eq!(attempt, 1); + assert!(delay <= Duration::from_millis(1_100), "{delay:?}"); + assert_eq!(restarts.failed("maintenance", quick).0, 2); + restarts.forget(&"maintenance"); + assert_eq!(restarts.attempt(&"maintenance"), None); + } +} diff --git a/rust/riverqueue/src/client/builder.rs b/rust/riverqueue/src/client/builder.rs new file mode 100644 index 000000000..09d1f2a2a --- /dev/null +++ b/rust/riverqueue/src/client/builder.rs @@ -0,0 +1,872 @@ +//! Client configuration and construction. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, AtomicU64}; +use std::sync::{Arc, Mutex, RwLock}; +use std::time::Duration; + +use serde_json::{Map, Value}; +use tokio::sync::{broadcast, watch}; + +use crate::__private::Pilot; +#[cfg(feature = "postgres")] +use crate::SchemaName; +#[cfg(feature = "postgres")] +use crate::client::validate::validate_identifier; +use crate::client::{ClientInner, EVENT_BUFFER_CAPACITY, InsertNotifyLimiter, validate_queue}; +use crate::database::Database; +use crate::periodic::{PeriodicJob, PeriodicJobs}; +use crate::{ + Client, Error, ErrorHandler, FETCH_COOLDOWN_MIN, FETCH_POLL_INTERVAL_DEFAULT, Hook, + InsertMiddleware, Plugin, QUEUE_NUM_WORKERS_MAX, RetryPolicy, WorkMiddleware, WorkerRegistry, +}; + +/// Default age at which running jobs are rescued (Go +/// `JobRescuerRescueAfterDefault`). +const RESCUE_AFTER_DEFAULT: Duration = Duration::from_hours(1); + +/// How long the job cleaner keeps finalized jobs of one state. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum Retention { + /// Delete jobs once they've been finalized for this long. + DeleteAfter(Duration), + /// Never delete these jobs. + Keep, +} + +impl Retention { + const fn from_option(retention: Option) -> Self { + match retention { + Some(duration) => Self::DeleteAfter(duration), + None => Self::Keep, + } + } + + const fn into_option(self) -> Option { + match self { + Self::DeleteAfter(duration) => Some(duration), + Self::Keep => None, + } + } +} + +/// Leader-owned maintenance timing and retention settings. +/// +/// Like River's other configuration values, `MaintenanceConfig` has a getter +/// for each setting and a `with_*` method that returns the configuration +/// with that setting changed. +#[derive(Clone, Debug)] +pub struct MaintenanceConfig { + /// Retention for cancelled jobs; `None` disables deletion. + pub(crate) cancelled_job_retention: Option, + /// Retention for completed jobs; `None` disables deletion. + pub(crate) completed_job_retention: Option, + /// Retention for discarded jobs; `None` disables deletion. + pub(crate) discarded_job_retention: Option, + /// Leader election and renewal interval. + pub(crate) elect_interval: Duration, + /// Job cleaner interval. + pub(crate) job_cleaner_interval: Duration, + /// Timeout for each job-cleaner deletion statement. + pub(crate) job_cleaner_timeout: Duration, + /// Test-only batch sizes of bulk maintenance services. + pub(crate) batch_sizes: crate::maintenance::BatchSizes, + /// Explicit age at which running jobs may be rescued. + pub(crate) rescue_after: Option, + /// Rescue age in effect, resolved against the job timeout at build time. + pub(crate) rescue_after_effective: Duration, + /// Stuck-job rescuer interval. + pub(crate) rescuer_interval: Duration, + /// Retention for inactive queue records. + pub(crate) queue_retention: Duration, + /// Inactive queue cleaner interval. + pub(crate) queue_cleaner_interval: Duration, + /// Due-job scheduler interval. + pub(crate) scheduler_interval: Duration, +} + +macro_rules! maintenance_retention { + ($getter:ident, $setter:ident, $state:literal, $default:literal) => { + #[doc = concat!("Returns how long ", $state, " jobs are kept before the job cleaner deletes them.")] + #[must_use] + pub const fn $getter(&self) -> Retention { + Retention::from_option(self.$getter) + } + + #[doc = concat!("Sets how long ", $state, " jobs are kept before the job cleaner deletes them. Defaults to deleting them after ", $default, ".")] + #[must_use] + pub const fn $setter(mut self, retention: Retention) -> Self { + self.$getter = retention.into_option(); + self + } + }; +} + +macro_rules! maintenance_duration { + ($getter:ident, $setter:ident, $what:literal, $default:literal) => { + #[doc = concat!("Returns ", $what, ".")] + #[must_use] + pub const fn $getter(&self) -> Duration { + self.$getter + } + + #[doc = concat!("Sets ", $what, ". Defaults to ", $default, ".")] + #[must_use] + pub const fn $setter(mut self, value: Duration) -> Self { + self.$getter = value; + self + } + }; +} + +impl MaintenanceConfig { + maintenance_retention!( + cancelled_job_retention, + with_cancelled_job_retention, + "cancelled", + "24 hours" + ); + maintenance_retention!( + completed_job_retention, + with_completed_job_retention, + "completed", + "24 hours" + ); + maintenance_retention!( + discarded_job_retention, + with_discarded_job_retention, + "discarded", + "7 days" + ); + maintenance_duration!( + elect_interval, + with_elect_interval, + "how often the client bids for leadership, or renews it while leader", + "5 seconds" + ); + maintenance_duration!( + job_cleaner_interval, + with_job_cleaner_interval, + "how often the leader deletes finalized jobs past their retention", + "30 seconds" + ); + maintenance_duration!( + job_cleaner_timeout, + with_job_cleaner_timeout, + "the timeout for each batch the job cleaner deletes", + "30 seconds" + ); + + /// Returns the explicitly configured rescue age, if any. + /// + /// When unset, a client rescues jobs running longer than one hour, or + /// than its job timeout plus one hour when a job timeout is configured. + #[must_use] + pub const fn rescue_after(&self) -> Option { + self.rescue_after + } + + /// Sets the age at which running jobs are considered stuck and rescued. + /// It must not be shorter than the client's job timeout. + #[must_use] + pub const fn with_rescue_after(mut self, value: Duration) -> Self { + self.rescue_after = Some(value); + self + } + + pub(crate) const fn effective_rescue_after(&self) -> Duration { + self.rescue_after_effective + } + maintenance_duration!( + rescuer_interval, + with_rescuer_interval, + "how often the leader looks for stuck jobs to rescue", + "30 seconds" + ); + maintenance_duration!( + queue_retention, + with_queue_retention, + "how long a queue record no client has touched is kept before the queue cleaner deletes it", + "24 hours" + ); + maintenance_duration!( + queue_cleaner_interval, + with_queue_cleaner_interval, + "how often the leader deletes queue records past their retention", + "1 hour" + ); + maintenance_duration!( + scheduler_interval, + with_scheduler_interval, + "how often the leader makes due scheduled and retryable jobs available", + "5 seconds" + ); +} + +impl Default for MaintenanceConfig { + fn default() -> Self { + Self { + cancelled_job_retention: Some(Duration::from_hours(24)), + completed_job_retention: Some(Duration::from_hours(24)), + discarded_job_retention: Some(Duration::from_hours(168)), + elect_interval: Duration::from_secs(5), + job_cleaner_interval: Duration::from_secs(30), + job_cleaner_timeout: Duration::from_secs(30), + batch_sizes: crate::maintenance::BatchSizes::default(), + rescue_after: None, + rescue_after_effective: RESCUE_AFTER_DEFAULT, + rescuer_interval: Duration::from_secs(30), + queue_retention: Duration::from_hours(24), + queue_cleaner_interval: Duration::from_hours(1), + scheduler_interval: Duration::from_secs(5), + } + } +} + +/// Queue-specific worker settings. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct QueueConfig { + /// Settings for an add-on crate, which River passes through unchanged. + pub(crate) extension_settings: Map, + /// Minimum delay between fetches, overriding the client's + /// [`ClientBuilder::fetch_cooldown`] when set. + pub(crate) fetch_cooldown: Option, + /// Fallback polling interval. + pub(crate) fetch_poll_interval: Duration, + /// Maximum jobs run concurrently by this client. + pub(crate) max_workers: usize, +} + +impl QueueConfig { + /// Creates queue configuration with River's timing defaults. + #[must_use] + pub fn new(max_workers: usize) -> Self { + Self { + extension_settings: Map::new(), + fetch_cooldown: None, + fetch_poll_interval: FETCH_POLL_INTERVAL_DEFAULT, + max_workers, + } + } + + /// Returns this queue's minimum delay between fetches, or `None` when it + /// uses the client's [`ClientBuilder::fetch_cooldown`]. + #[must_use] + pub const fn fetch_cooldown(&self) -> Option { + self.fetch_cooldown + } + + /// Returns the fallback polling interval. + #[must_use] + pub const fn fetch_poll_interval(&self) -> Duration { + self.fetch_poll_interval + } + + /// Returns the maximum jobs run concurrently. + #[must_use] + pub const fn max_workers(&self) -> usize { + self.max_workers + } + + /// Sets the minimum delay between fetches for this queue, overriding the + /// client's [`ClientBuilder::fetch_cooldown`]. Throughput is limited by + /// this value. It must be at least [`FETCH_COOLDOWN_MIN`](crate::FETCH_COOLDOWN_MIN) and no longer + /// than the fetch poll interval. + /// + /// The override only paces this queue's fetches. Insert notifications + /// are always suppressed for the client's fetch cooldown. + #[must_use] + pub const fn with_fetch_cooldown(mut self, interval: Duration) -> Self { + self.fetch_cooldown = Some(interval); + self + } + + /// Sets how often the queue polls for jobs when no insert notification + /// arrives. Defaults to + /// [`FETCH_POLL_INTERVAL_DEFAULT`](crate::FETCH_POLL_INTERVAL_DEFAULT) + /// (one second), and can't be shorter than the queue's fetch cooldown. + /// River adds up to 10% of jitter to each poll. + #[must_use] + pub const fn with_fetch_poll_interval(mut self, interval: Duration) -> Self { + self.fetch_poll_interval = interval; + self + } + + /// Sets the maximum jobs run concurrently. + #[must_use] + pub const fn with_max_workers(mut self, maximum: usize) -> Self { + self.max_workers = maximum; + self + } + + /// Returns the minimum delay between this queue's fetches, given the + /// client's fetch cooldown. + pub(crate) fn resolved_fetch_cooldown(&self, client_fetch_cooldown: Duration) -> Duration { + self.fetch_cooldown.unwrap_or(client_fetch_cooldown) + } + + /// Validates the queue, given the client's fetch cooldown. + pub(super) fn validate( + &self, + name: &str, + client_fetch_cooldown: Duration, + ) -> Result<(), Error> { + validate_queue(name)?; + if !(1..=QUEUE_NUM_WORKERS_MAX).contains(&self.max_workers) { + return Err(Error::configuration(format!( + "queue {name:?} max_workers must be between 1 and {QUEUE_NUM_WORKERS_MAX}" + ))); + } + if self + .fetch_cooldown + .is_some_and(|cooldown| cooldown < FETCH_COOLDOWN_MIN) + { + return Err(Error::configuration( + "fetch cooldown must be at least one millisecond".to_owned(), + )); + } + if self.fetch_poll_interval < self.resolved_fetch_cooldown(client_fetch_cooldown) { + return Err(Error::configuration( + "fetch poll interval cannot be shorter than fetch cooldown".to_owned(), + )); + } + Ok(()) + } +} + +/// Builder for a River client. +#[allow( + clippy::struct_excessive_bools, + reason = "each flag is an independent configuration option, not a state" +)] +pub struct ClientBuilder { + pub(super) allow_legacy_job_kinds: bool, + pub(super) allow_unregistered_job_kinds: bool, + pub(super) database: Database, + pub(super) default_max_attempts: i16, + pub(super) error_handler: Option>, + pub(super) fetch_cooldown: Duration, + pub(super) fetch_only_known_kinds: bool, + pub(super) hooks: Vec>, + pub(super) id: String, + pub(super) insert_middleware: Vec>, + pub(super) job_stuck_threshold: Duration, + pub(super) job_timeout: Option, + pub(crate) leader_election_disabled: bool, + pub(super) maintenance: MaintenanceConfig, + pub(super) periodic_jobs: Vec, + pub(super) pilot: Arc, + pub(super) poll_only: bool, + pub(crate) producer_report_interval: Duration, + pub(super) queues: HashMap, + pub(super) retry_policy: Arc, + pub(super) soft_stop_timeout: Option, + pub(super) work_middleware: Vec>, + pub(crate) workers: WorkerRegistry, +} + +impl std::fmt::Debug for ClientBuilder { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("ClientBuilder") + .field("database_kind", &self.database.kind()) + .field("id", &self.id) + .field("queue_count", &self.queues.len()) + .field("worker_kinds", &self.workers.kinds()) + .field("hook_count", &self.hooks.len()) + .field("periodic_job_count", &self.periodic_jobs.len()) + .field("leader_election_disabled", &self.leader_election_disabled) + .finish_non_exhaustive() + } +} + +impl ClientBuilder { + /// Temporarily permits inserting legacy job kinds that don't match River's + /// kind format: 2 to 127 bytes, starting with a letter, digit, or `_`, + /// and otherwise made of letters, digits, and `_-[]<>/.·:+`. + #[must_use] + pub fn allow_legacy_job_kinds(mut self) -> Self { + self.allow_legacy_job_kinds = true; + self + } + + /// Allows inserting kinds with no worker in this client's registry. + /// Insert-only clients already permit every kind. + #[must_use] + pub fn allow_unregistered_job_kinds(mut self) -> Self { + self.allow_unregistered_job_kinds = true; + self + } + + /// Sets the maximum attempts used by [`Client::insert`] when the job type + /// does not override it. Defaults to + /// [`MAX_ATTEMPTS_DEFAULT`](crate::MAX_ATTEMPTS_DEFAULT) (25), and must be + /// at least one. + #[must_use] + pub fn default_max_attempts(mut self, maximum: i16) -> Self { + self.default_max_attempts = maximum; + self + } + + /// Installs a worker error and stuck-task handler. + #[must_use] + pub fn error_handler(mut self, handler: H) -> Self { + self.error_handler = Some(Arc::new(handler)); + self + } + + /// Sets the minimum delay between fetches of new jobs. Jobs are fetched + /// at most this often, and when no insert notifications arrive, fetches + /// may wait as long as a queue's fetch poll interval. Throughput is + /// limited by this value. A queue may override it with + /// [`QueueConfig::with_fetch_cooldown`]. + /// + /// It also paces insert notifications. After this client notifies a + /// queue that jobs were inserted, further notifications for that queue + /// are skipped until the cooldown has passed, whichever insertion, + /// transaction, or scheduler pass would send them. The window starts when + /// the notification is written, even if its transaction later rolls + /// back. A job whose notification was skipped is found by the next fetch + /// of its queue, which may wait for the queue's fetch poll interval. + /// + /// Like River Go's `Config.FetchCooldown`, it defaults to + /// [`FETCH_COOLDOWN_DEFAULT`](crate::FETCH_COOLDOWN_DEFAULT) (100 + /// milliseconds) and must be at least + /// [`FETCH_COOLDOWN_MIN`](crate::FETCH_COOLDOWN_MIN) (one millisecond). + /// A queue's fetch poll interval can't be shorter than the cooldown it + /// uses. + #[must_use] + pub fn fetch_cooldown(mut self, cooldown: Duration) -> Self { + self.fetch_cooldown = cooldown; + self + } + + /// Restricts claims to the kinds of registered workers, including their + /// aliases, like River Go's `Config.FetchOnlyKnownKinds`. Jobs of other + /// kinds stay available without using attempts, so clients with + /// different workers can share a queue. + /// + /// It only affects claiming. A leader's rescuer still handles stuck jobs + /// in every queue and discards those whose kinds it doesn't know, so a + /// client with only some of the workers should also be built + /// [`without_leader_election`](Self::without_leader_election), with + /// another client that has every worker eligible to lead. + /// + /// Disabled by default, so a job of an unknown kind is claimed and fails + /// its attempt. + #[must_use] + pub const fn fetch_only_known_kinds(mut self, enabled: bool) -> Self { + self.fetch_only_known_kinds = enabled; + self + } + + /// Adds an ordered lifecycle hook. + #[must_use] + pub fn hook(mut self, hook: H) -> Self { + self.hooks.push(Arc::new(hook)); + self + } + + /// Sets a stable client identifier, 1 to 100 bytes long. It must be + /// unique per running process, since leader election and job attempts + /// record it. Defaults to the host name, the creation time, and a random + /// suffix. + #[must_use] + pub fn id(mut self, id: impl Into) -> Self { + self.id = id.into(); + self + } + + /// Adds ordered insertion middleware. + #[must_use] + pub fn insert_middleware(mut self, middleware: M) -> Self { + self.insert_middleware.push(Arc::new(middleware)); + self + } + + /// Sets how long a job may keep running after its + /// [`WorkContext::cancellation_token`](crate::WorkContext::cancellation_token) + /// is cancelled (by a timeout, a remote cancellation, or a hard stop) + /// before River considers it stuck. A stuck job's + /// [`ErrorHandler::handle_stuck`](crate::ErrorHandler::handle_stuck) + /// runs, its task is aborted, and the attempt fails like any other failed + /// attempt. Defaults to + /// [`JOB_STUCK_THRESHOLD_DEFAULT`](crate::JOB_STUCK_THRESHOLD_DEFAULT) + /// (10 seconds). The threshold must be positive. + #[must_use] + pub fn job_stuck_threshold(mut self, threshold: Duration) -> Self { + self.job_stuck_threshold = threshold; + self + } + + /// Sets how long a job may run before its + /// [`WorkContext::cancellation_token`](crate::WorkContext::cancellation_token) is cancelled and the attempt + /// fails, unless its worker overrides it. Defaults to one minute. The + /// timeout must be positive; use + /// [`without_job_timeout`](Self::without_job_timeout) to let jobs run + /// without a limit. + #[must_use] + pub fn job_timeout(mut self, timeout: Duration) -> Self { + self.job_timeout = Some(timeout); + self + } + + /// Configures leader-owned maintenance services. + /// + /// Has no effect on a client built with + /// [`without_leader_election`](Self::without_leader_election), which + /// never runs them. + #[must_use] + pub fn maintenance(mut self, maintenance: MaintenanceConfig) -> Self { + self.maintenance = maintenance; + self + } + + /// Adds a periodic job to the initial client configuration. + /// + /// Only the elected leader enqueues periodic jobs, so a client built + /// with [`without_leader_election`](Self::without_leader_election) + /// rejects them when built. + #[must_use] + pub fn periodic_job(mut self, job: PeriodicJob) -> Self { + self.periodic_jobs.push(job); + self + } + + /// Installs the hooks and middleware contributed by a plugin, after any + /// registered earlier. + #[must_use] + #[allow( + clippy::needless_pass_by_value, + reason = "taking the plugin by value matches the other registration methods" + )] + pub fn plugin(mut self, plugin: P) -> Self { + let mut extensions = crate::Extensions::default(); + plugin.install(&mut extensions); + self.hooks.extend(extensions.hooks); + self.insert_middleware.extend(extensions.insert_middleware); + self.work_middleware.extend(extensions.work_middleware); + self + } + + /// Adds or replaces a queue. + #[must_use] + pub fn queue(mut self, name: impl Into, config: QueueConfig) -> Self { + self.queues.insert(name.into(), config); + self + } + + /// Replaces River's default retry policy. + #[must_use] + pub fn retry_policy(mut self, retry_policy: P) -> Self { + self.retry_policy = Arc::new(retry_policy); + self + } + + /// Escalates a soft stop to a hard stop after this duration. By default, + /// running jobs finish without a limit. + /// The timeout must be positive. + /// + /// The client starts this timer when fetching stops, however the stop was + /// requested: [`RunHandle::shutdown`](crate::RunHandle::shutdown), + /// [`Stopper::stop`](crate::Stopper::stop), or the signal passed to + /// [`Client::start_with_graceful_shutdown`]. Jobs still running when it + /// expires are cancelled as if by + /// [`Stopper::stop_now`](crate::Stopper::stop_now). + #[must_use] + pub fn soft_stop_timeout(mut self, timeout: Duration) -> Self { + self.soft_stop_timeout = Some(timeout); + self + } + + /// Installs a pilot from a companion crate. + #[must_use] + pub(crate) fn with_pilot(mut self, pilot: P) -> Self { + self.pilot = Arc::new(pilot); + self + } + + /// Lets jobs run without a time limit unless their worker sets one. + #[must_use] + pub fn without_job_timeout(mut self) -> Self { + self.job_timeout = None; + self + } + + /// Keeps this client out of leader election. + /// + /// The client still fetches and works jobs from its queues, sends and + /// receives notifications, and runs extension runtime services, but it + /// never becomes leader, so it never runs leader-owned maintenance: the + /// scheduler, the periodic job enqueuer, the stuck job rescuer, the job + /// and queue cleaners, the reindexer, and extension maintenance + /// services. This suits clients dedicated to particular queues that + /// should spend their resources only on those queues' jobs. + /// + /// At least one other started client using the same database and schema, + /// in any River implementation, must remain eligible to lead. Otherwise + /// scheduled jobs and retries never become available, periodic jobs are + /// never enqueued, stuck jobs are never rescued, and finalized jobs are + /// never deleted. This client stays ineligible even when no other client + /// is running. + /// + /// Such a client can't configure periodic jobs: [`ClientBuilder::build`] + /// fails when any were added with [`periodic_job`](Self::periodic_job), + /// and [`PeriodicJobs::add`] and [`PeriodicJobs::add_many`] fail on its + /// [`Client::periodic_jobs`]. It still works periodic jobs that a leader + /// enqueues in its queues. + #[must_use] + pub fn without_leader_election(mut self) -> Self { + self.leader_election_disabled = true; + self + } + + /// Disables the backend notification channel or outbox poller while + /// retaining queue fetch polling. + /// + /// The client then polls for new jobs every queue's fetch poll interval, + /// and every two seconds for queue changes and for cancellations of its + /// running jobs requested by other clients. A client using a PostgreSQL + /// server without `LISTEN`/`NOTIFY`, like YugabyteDB by default, runs + /// this way on its own. + #[must_use] + pub fn without_notifications(mut self) -> Self { + self.poll_only = true; + self + } + + /// Adds ordered worker middleware. + #[must_use] + pub fn work_middleware(mut self, middleware: M) -> Self { + self.work_middleware.push(Arc::new(middleware)); + self + } + + /// Installs a typed worker registry. + #[must_use] + pub fn workers(mut self, workers: WorkerRegistry) -> Self { + self.workers = workers; + self + } + + /// Validates configuration and builds the client. + #[allow( + clippy::too_many_lines, + reason = "central validation keeps builder failures deterministic before allocating runtime state" + )] + /// + /// # Errors + /// + /// Returns [`Error::Configuration`] when a setting is out of range or + /// settings conflict, such as queues configured without workers, a rescue + /// age shorter than the job timeout, or periodic jobs on a client without + /// leader election. + pub fn build(self) -> Result { + if self.default_max_attempts < 1 { + return Err(Error::configuration( + "default max attempts must be greater than zero".to_owned(), + )); + } + if self.id.is_empty() || self.id.len() > 100 { + return Err(Error::configuration( + "client ID must contain between 1 and 100 bytes".to_owned(), + )); + } + if self + .soft_stop_timeout + .is_some_and(|timeout| timeout.is_zero()) + { + return Err(Error::configuration( + "soft stop timeout must be positive".to_owned(), + )); + } + if self.job_timeout.is_some_and(|timeout| timeout.is_zero()) { + return Err(Error::configuration( + "job timeout must be positive; use without_job_timeout to disable it".to_owned(), + )); + } + if self.job_stuck_threshold.is_zero() { + return Err(Error::configuration( + "job stuck threshold must be positive".to_owned(), + )); + } + if self.fetch_cooldown < FETCH_COOLDOWN_MIN { + return Err(Error::configuration( + "fetch cooldown must be at least one millisecond".to_owned(), + )); + } + for (name, config) in &self.queues { + config.validate(name, self.fetch_cooldown)?; + validate_queue_settings(self.pilot.as_ref(), name, config)?; + } + if self.producer_report_interval.is_zero() { + return Err(Error::configuration( + "producer report interval must be positive".to_owned(), + )); + } + for (name, interval) in [ + ("elect interval", self.maintenance.elect_interval), + ( + "job cleaner interval", + self.maintenance.job_cleaner_interval, + ), + ("job cleaner timeout", self.maintenance.job_cleaner_timeout), + ( + "rescue after", + self.maintenance + .rescue_after + .unwrap_or(RESCUE_AFTER_DEFAULT), + ), + ("rescuer interval", self.maintenance.rescuer_interval), + ( + "queue cleaner interval", + self.maintenance.queue_cleaner_interval, + ), + ("queue retention", self.maintenance.queue_retention), + ("scheduler interval", self.maintenance.scheduler_interval), + ] { + if interval.is_zero() { + return Err(Error::configuration(format!("{name} must be positive"))); + } + } + #[cfg(feature = "postgres")] + let reindex = self.database.postgres_reindex(); + #[cfg(feature = "postgres")] + if reindex.is_some_and(|config| config.timeout().is_some_and(|timeout| timeout.is_zero())) { + return Err(Error::configuration( + "reindexer timeout must be positive; use without_timeout to disable it".to_owned(), + )); + } + #[cfg(feature = "postgres")] + if matches!( + reindex.map(crate::database::PostgresReindexConfig::schedule), + Some(crate::database::PostgresReindexSchedule::Interval(interval)) if interval.is_zero() + ) { + return Err(Error::configuration( + "reindexer interval must be positive".to_owned(), + )); + } + #[cfg(feature = "postgres")] + for index_name in reindex + .into_iter() + .flat_map(crate::database::PostgresReindexConfig::index_names) + { + validate_identifier(index_name, "reindexer index")?; + } + if !self.queues.is_empty() && self.workers.kinds().is_empty() { + return Err(Error::configuration( + "workers must be configured when queues are configured".to_owned(), + )); + } + // Like Go, rescuing jobs before their timeout could run them twice. + if let (Some(rescue_after), Some(job_timeout)) = + (self.maintenance.rescue_after, self.job_timeout) + && rescue_after < job_timeout + { + return Err(Error::configuration( + "rescue after cannot be less than the job timeout".to_owned(), + )); + } + let mut maintenance = self.maintenance; + maintenance.rescue_after_effective = maintenance.rescue_after.unwrap_or_else(|| { + self.job_timeout + .filter(|timeout| !timeout.is_zero()) + .map_or(RESCUE_AFTER_DEFAULT, |timeout| { + timeout + RESCUE_AFTER_DEFAULT + }) + }); + + if self.leader_election_disabled && !self.periodic_jobs.is_empty() { + return Err(Error::configuration( + "periodic jobs must be empty when leader election is disabled".to_owned(), + )); + } + let periodic_jobs = + PeriodicJobs::from_jobs(self.periodic_jobs, self.leader_election_disabled)?; + let fetch_kinds = self.fetch_only_known_kinds.then(|| { + self.workers + .kinds() + .into_iter() + .map(str::to_owned) + .collect::>() + }); + #[cfg(feature = "postgres")] + let schema = self + .database + .postgres_schema() + .cloned() + .unwrap_or_else(SchemaName::current); + let (events, _) = broadcast::channel(EVENT_BUFFER_CAPACITY); + let (queue_changes, _) = watch::channel(0_u64); + let (leadership_wakeups, _) = broadcast::channel(1_024); + let (queue_notifications, _) = broadcast::channel(1_024); + let client = Client { + inner: Arc::new(ClientInner { + allow_legacy_job_kinds: self.allow_legacy_job_kinds, + allow_unregistered_job_kinds: self.allow_unregistered_job_kinds, + completion_sender: Mutex::new(None), + database: self.database, + default_max_attempts: self.default_max_attempts, + error_handler: self.error_handler, + events, + fetch_cooldown: self.fetch_cooldown, + fetch_kinds, + fetch_registration_windows: AtomicU64::new(0), + hooks: self.hooks, + id: self.id, + insert_middleware: self.insert_middleware, + insert_notify_limiter: InsertNotifyLimiter::new(self.fetch_cooldown), + job_stuck_threshold: self.job_stuck_threshold, + job_timeout: self.job_timeout, + leader_election_disabled: self.leader_election_disabled, + leadership_wakeups, + live_queues: watch::channel(std::collections::HashSet::new()).0, + maintenance, + #[cfg(test)] + notifier_start_panics: AtomicU64::new(0), + peer_owners: Mutex::new(HashMap::new()), + pending_cancellations: Mutex::new(HashMap::new()), + periodic_jobs, + pilot: self.pilot, + poll_only: self.poll_only, + producer_report_interval: self.producer_report_interval, + queue_changes, + queue_notifications, + queues: RwLock::new(self.queues), + retry_policy: self.retry_policy, + running: Mutex::new(HashMap::new()), + #[cfg(feature = "postgres")] + schema, + soft_stop_timeout: self.soft_stop_timeout, + started: AtomicBool::new(false), + work_middleware: self.work_middleware, + workers: self.workers, + }), + }; + client + .inner + .pilot + .install(crate::__private::PilotInstallContext { + client: client.downgrade(), + database: client.inner.pilot_database(), + producer_report_interval: client.inner.producer_report_interval, + }); + Ok(client) + } +} + +/// Checks a queue's extension settings with the client's pilot. +pub(super) fn validate_queue_settings( + pilot: &dyn Pilot, + name: &str, + config: &QueueConfig, +) -> Result<(), Error> { + pilot + .validate_queue_settings(name, &config.extension_settings) + .map_err(|source| Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "queue settings", + }, + source, + }) +} diff --git a/rust/riverqueue/src/client/completer.rs b/rust/riverqueue/src/client/completer.rs new file mode 100644 index 000000000..2d96d9b70 --- /dev/null +++ b/rust/riverqueue/src/client/completer.rs @@ -0,0 +1,807 @@ +//! Batched persistence of job completions. +//! +//! This mirrors River Go's `BatchCompleter`: results accumulate briefly and are +//! written with one set-state-if-running statement per batch. A database write +//! is retried with backoff, and a batch that still fails stays queued and is +//! retried again rather than being dropped, so a transient error cannot leave +//! successfully worked jobs `running` until the rescuer runs them again. + +use std::collections::HashMap; +use std::collections::{HashSet, VecDeque}; +use std::sync::Arc; +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use futures_util::FutureExt as _; +use serde_json::{Map, Value}; +#[cfg(feature = "postgres")] +use sqlx::AssertSqlSafe; +#[cfg(feature = "postgres")] +use sqlx::postgres::PgRow; +#[cfg(feature = "postgres")] +use sqlx::types::Json; +use tokio::sync::mpsc; +use tokio::task::JoinSet; +use tokio_util::sync::CancellationToken; +use tracing::{debug, error}; + +use crate::__private::{DatabaseConnection as PilotDatabaseConnection, JobSetStateParams}; +use crate::client::attempts::remove_running_attempt; +use crate::client::backoff::exponential_backoff; +#[cfg(feature = "sqlite")] +use crate::client::sqlite_backend_error; +use crate::client::{ClientInner, RuntimeNotification, peers, tolerant_row}; +#[cfg(feature = "postgres")] +use crate::client::{decode_job_row, job_projection}; +use crate::database::{DatabaseKind, DatabasePool}; +use crate::{AttemptError, Error, Event, JobEventKind, JobRow, JobState, JobStatistics}; + +/// Most updates written by one statement, matching River Go. +pub(super) const COMPLETION_BATCH_SIZE: usize = 5_000; +/// How long sparse results coalesce before they are written. +const COMPLETION_BATCH_DELAY: Duration = Duration::from_millis(10); +/// Ready and deferred updates held before the batcher stops accepting more. +/// The bounded channel then applies backpressure to workers, as River Go's +/// backlog wait does. +pub(super) const COMPLETION_BACKLOG_LIMIT: usize = COMPLETION_BATCH_SIZE * 2; +/// Most concurrent batch writes River OSS issues on PostgreSQL. +#[cfg(feature = "postgres")] +const COMPLETION_POSTGRES_CONCURRENCY: usize = 2; +/// Attempts in one retry cycle, matching River Go's `numRetries`. +const COMPLETION_RETRY_ATTEMPTS: u32 = 3; +/// Per-attempt timeout, matching River Go's `HotOperationTimeout`. +pub(super) const HOT_OPERATION_TIMEOUT: Duration = Duration::from_secs(10); + +/// One requested job state transition. +/// +/// The fields follow River Go's `JobSetStateIfRunningParams`: `None` leaves +/// the corresponding column unchanged, and metadata is merged only when it is +/// not empty. +pub(super) struct CompletionUpdate { + /// Replacement attempt, sent only for snoozes and shutdown interrupts, + /// which return an attempt that should not count. + pub(super) attempt: Option, + pub(super) cancellation: CancellationToken, + pub(super) error: Option, + pub(super) event_kind: JobEventKind, + pub(super) finalized_at: Option>, + pub(super) job_id: i64, + pub(super) metadata: Map, + pub(super) scheduled_at: Option>, + pub(super) state: JobState, + pub(super) timing: CompletionTiming, + /// Set for a peer attempt's outcome, whose ownership ends as it persists. + pub(super) peer: Option>, +} + +/// Maps a persisted row to the event it reports. +/// +/// Returns `None` for `pending` and `running`: the row was moved out of the +/// requested transition by someone else (an operator, an extension, or a +/// newer attempt after a rescue), so reporting a completion would be wrong. +pub(super) fn persisted_completion_event_kind( + state: JobState, + requested: JobEventKind, +) -> Option { + Some(match state { + JobState::Available => match requested { + JobEventKind::Failed | JobEventKind::Interrupted | JobEventKind::Snoozed => requested, + JobEventKind::Cancelled | JobEventKind::Completed => JobEventKind::Failed, + }, + JobState::Cancelled => JobEventKind::Cancelled, + JobState::Completed => JobEventKind::Completed, + JobState::Discarded | JobState::Retryable => JobEventKind::Failed, + JobState::Scheduled => JobEventKind::Snoozed, + JobState::Pending | JobState::Running => return None, + }) +} + +pub(super) struct CompletionAttempt { + pub(super) cancellation: CancellationToken, + pub(super) timing: CompletionTiming, +} + +#[derive(Clone, Copy)] +pub(super) struct CompletionTiming { + pub(super) completion_started: std::time::Instant, + pub(super) queue_wait_duration: Duration, + pub(super) run_duration: Duration, +} + +/// Persists completions until every sender is dropped and all accepted +/// updates were written or abandoned. +pub(super) async fn run_completion_batcher( + inner: Arc, + receiver: mpsc::Receiver, +) -> Result<(), Error> { + CompletionBatcher::new(inner).run(receiver).await; + Ok(()) +} + +type BatchOutcome = (Vec, Result, Error>); + +pub(super) struct CompletionBatcher { + /// Updates accepted while a batch containing the same job is in flight. + /// They replace the in-flight update's successor once it finishes. + deferred: HashMap, + /// Whether ready updates may be written without waiting for a full batch. + flush_due: bool, + in_flight: HashSet, + inner: Arc, + max_concurrency: usize, + ready: HashMap, + ready_order: VecDeque, + /// Set once a batch fails during shutdown; remaining updates are then + /// abandoned without being attempted, like River Go's stop path, which + /// stops flushing at the first error. + stop_retrying: bool, + tasks: JoinSet, + task_ids: HashMap>, +} + +impl CompletionBatcher { + pub(super) fn new(inner: Arc) -> Self { + let backend_concurrency = match inner.database.kind() { + #[cfg(feature = "postgres")] + DatabaseKind::Postgres => COMPLETION_POSTGRES_CONCURRENCY, + #[cfg(feature = "sqlite")] + DatabaseKind::Sqlite => 1, + }; + // Like River Go's `completionConcurrency`, an intercepting extension + // can only lower the backend's limit. + let max_concurrency = if inner.pilot.intercepts_job_set_state() { + backend_concurrency.min(inner.pilot.job_set_state_concurrency().max(1)) + } else { + backend_concurrency + }; + Self { + deferred: HashMap::new(), + flush_due: false, + in_flight: HashSet::new(), + inner, + max_concurrency, + ready: HashMap::new(), + ready_order: VecDeque::new(), + stop_retrying: false, + tasks: JoinSet::new(), + task_ids: HashMap::new(), + } + } + + fn backlog(&self) -> usize { + self.ready.len() + self.deferred.len() + } + + /// Drops an update that will never be written, releasing its attempt. + fn discard(&self, update: &CompletionUpdate) { + remove_running_attempt(&self.inner.running, update.job_id, &update.cancellation); + } + + fn enqueue(&mut self, update: CompletionUpdate) { + if self.in_flight.contains(&update.job_id) { + if let Some(superseded) = self.deferred.insert(update.job_id, update) { + self.discard(&superseded); + } + return; + } + self.enqueue_ready(update); + } + + /// Adds an update to the ready set. A newer update for the same job (from a + /// later attempt after a rescue) supersedes an older unwritten one. + fn enqueue_ready(&mut self, update: CompletionUpdate) { + let job_id = update.job_id; + match self.ready.insert(job_id, update) { + Some(superseded) => self.discard(&superseded), + None => self.ready_order.push_back(job_id), + } + } + + fn finish(&mut self, joined: Result<(tokio::task::Id, BatchOutcome), tokio::task::JoinError>) { + let (task_id, (batch, result)) = match joined { + Ok(joined) => joined, + Err(join_error) => { + // Batch tasks catch panics, so this only happens when the + // runtime is shutting down. Release the batch's jobs so any + // deferred successors are not stranded behind it. + error!(error = %join_error, "River completion batch task stopped"); + for job_id in self.task_ids.remove(&join_error.id()).unwrap_or_default() { + self.release(job_id); + } + return; + } + }; + self.task_ids.remove(&task_id); + match result { + Ok(rows) => { + let mut rows = rows + .into_iter() + .map(|row| (row.id, row)) + .collect::>(); + for update in &batch { + finish_batched_completion(&self.inner, update, rows.remove(&update.job_id)); + } + for update in batch { + self.release(update.job_id); + } + } + Err(error) if self.stop_retrying || is_non_retryable_completion_error(&error) => { + error!( + error = %crate::error::Chain(&error), + num_jobs = batch.len(), + "River could not persist job completions; the rescuer will retry them" + ); + for update in batch { + self.discard(&update); + self.release(update.job_id); + } + } + Err(error) => { + debug!( + error = %crate::error::Chain(&error), + num_jobs = batch.len(), + "requeued River completion batch after repeated errors" + ); + for update in batch { + self.in_flight.remove(&update.job_id); + match self.deferred.remove(&update.job_id) { + Some(newer) => { + self.discard(&update); + self.enqueue_ready(newer); + } + None => self.enqueue_ready(update), + } + } + self.flush_due = true; + } + } + } + + /// Discards every update not yet written: ready ones, and any still + /// queued in `receiver`, whose senders are all gone once shutdown + /// abandons the backlog. Returns whether `receiver` is exhausted. + fn abandon_backlog(&mut self, receiver: &mut mpsc::Receiver) -> bool { + let mut abandoned = 0_usize; + for job_id in std::mem::take(&mut self.ready_order) { + if let Some(update) = self.ready.remove(&job_id) { + self.discard(&update); + abandoned += 1; + } + } + let exhausted = loop { + match receiver.try_recv() { + Ok(update) => { + self.discard(&update); + abandoned += 1; + } + Err(mpsc::error::TryRecvError::Disconnected) => break true, + Err(mpsc::error::TryRecvError::Empty) => break false, + } + }; + if abandoned > 0 { + error!( + num_jobs = abandoned, + "River client stopping after a completion error; abandoning unwritten job completions for the rescuer" + ); + } + exhausted + } + + /// Marks a job's batch finished and promotes a deferred successor. + fn release(&mut self, job_id: i64) { + self.in_flight.remove(&job_id); + if let Some(deferred) = self.deferred.remove(&job_id) { + self.enqueue_ready(deferred); + self.flush_due = true; + } + } + + async fn run(mut self, mut receiver: mpsc::Receiver) { + let mut accepting = true; + let coalesce = tokio::time::sleep(Duration::ZERO); + tokio::pin!(coalesce); + let mut coalescing = false; + loop { + if self.stop_retrying && self.abandon_backlog(&mut receiver) { + accepting = false; + } + self.start_ready_batches(!accepting); + if self.ready.is_empty() { + self.flush_due = false; + } + if !accepting && self.tasks.is_empty() && self.ready.is_empty() { + break; + } + let receiving = accepting && self.backlog() < COMPLETION_BACKLOG_LIMIT; + tokio::select! { + update = receiver.recv(), if receiving => match update { + Some(update) => { + self.enqueue(update); + if !coalescing && !self.flush_due { + coalesce + .as_mut() + .reset(tokio::time::Instant::now() + COMPLETION_BATCH_DELAY); + coalescing = true; + } + } + None => accepting = false, + }, + joined = self.tasks.join_next_with_id(), if !self.tasks.is_empty() => { + if let Some(joined) = joined { + // Every sender is gone once the client is stopping, + // even while a full backlog keeps the batcher from + // receiving the channel's end, which an outage + // would otherwise postpone until the database + // returns. + let stopping = !accepting || receiver.is_closed(); + if stopping && joined.as_ref().is_ok_and(|(_, (_, result))| result.is_err()) { + self.stop_retrying = true; + } + self.finish(joined); + } + } + () = &mut coalesce, if coalescing => { + coalescing = false; + self.flush_due = true; + } + } + } + } + + fn spawn_batch(&mut self) -> bool { + let mut batch = Vec::with_capacity(self.ready.len().min(COMPLETION_BATCH_SIZE)); + while batch.len() < COMPLETION_BATCH_SIZE { + let Some(job_id) = self.ready_order.pop_front() else { + break; + }; + if let Some(update) = self.ready.remove(&job_id) { + self.in_flight.insert(job_id); + batch.push(update); + } + } + if batch.is_empty() { + return false; + } + let job_ids = batch.iter().map(|update| update.job_id).collect(); + let inner = Arc::clone(&self.inner); + let handle = self.tasks.spawn(async move { + // Keep the batch even if persistence panics so its jobs are + // retried rather than silently lost. + let result = std::panic::AssertUnwindSafe(persist_with_retries(&inner, &batch)) + .catch_unwind() + .await + .unwrap_or_else(|panic| { + Err(Error::runtime_context( + "job completion", + format!( + "completion persistence panicked: {}", + crate::error::panic_message(&panic) + ), + )) + }); + if let Ok(rows) = &result { + notify_interrupted_jobs(&inner, &batch, rows).await; + } + (batch, result) + }); + self.task_ids.insert(handle.id(), job_ids); + true + } + + /// Starts as many batches as policy allows. + /// + /// A single writer takes whatever is ready once the coalescing delay has + /// passed. A second concurrent writer is worthwhile only for a full batch, + /// which keeps sparse workloads at one query at a time like River Go. + fn start_ready_batches(&mut self, draining: bool) { + while self.tasks.len() < self.concurrency() && !self.ready.is_empty() { + let full = self.ready.len() >= COMPLETION_BATCH_SIZE; + let may_start = if self.tasks.is_empty() { + full || self.flush_due || draining + } else { + full + }; + if !may_start || !self.spawn_batch() { + return; + } + } + } + + pub(super) const fn concurrency(&self) -> usize { + self.max_concurrency + } +} + +/// Whether a completion error can never succeed on retry, matching River Go's +/// `isNonRetryableCompleterError` for a closed pool. +fn is_non_retryable_completion_error(error: &Error) -> bool { + let mut source: Option<&(dyn std::error::Error + 'static)> = Some(error); + while let Some(mut current) = source { + // A database error is transparent, so its driver error appears in + // the chain only through the variant. + if let Some(Error::Database(inner)) = current.downcast_ref::() { + current = inner; + } + if matches!( + current.downcast_ref::(), + Some(sqlx::Error::PoolClosed) + ) { + return true; + } + source = current.source(); + } + false +} + +/// Runs a completion write up to three times with River's service backoff and +/// a per-attempt timeout, independent of client shutdown, like River Go's +/// `withRetries`. A closed pool is not retried. +pub(super) async fn with_completion_retries( + operation: &'static str, + mut attempt: F, +) -> Result +where + F: FnMut() -> Fut, + Fut: std::future::Future>, +{ + let mut attempt_number = 1; + loop { + let error = match tokio::time::timeout(HOT_OPERATION_TIMEOUT, attempt()).await { + Ok(Ok(value)) => return Ok(value), + Ok(Err(error)) if is_non_retryable_completion_error(&error) => return Err(error), + Ok(Err(error)) => error, + Err(_) => Error::runtime_context( + operation, + format!("attempt timed out after {HOT_OPERATION_TIMEOUT:?}"), + ), + }; + if attempt_number >= COMPLETION_RETRY_ATTEMPTS { + error!( + attempt = attempt_number, + error = %crate::error::Chain(&error), + operation, + "River completer error; too many errors, giving up on this attempt cycle" + ); + return Err(error); + } + let sleep = exponential_backoff(attempt_number); + error!( + attempt = attempt_number, + error = %crate::error::Chain(&error), + operation, + sleep_duration = ?sleep, + "River completer error (will retry after sleep)" + ); + tokio::time::sleep(sleep).await; + attempt_number += 1; + } +} + +async fn persist_with_retries( + inner: &ClientInner, + batch: &[CompletionUpdate], +) -> Result, Error> { + with_completion_retries("job completion", || persist_completion_batch(inner, batch)).await +} + +/// Wakes producers for jobs that client shutdown returned to `available`, so +/// another client picks them up without waiting for its next poll. SQLite +/// writes its durable wakeup inside the completion transaction instead. +#[cfg_attr( + not(feature = "postgres"), + allow(clippy::unused_async, reason = "only PostgreSQL sends a notification") +)] +async fn notify_interrupted_jobs(inner: &ClientInner, batch: &[CompletionUpdate], rows: &[JobRow]) { + let interrupted = batch + .iter() + .filter(|update| update.event_kind == JobEventKind::Interrupted) + .map(|update| update.job_id) + .collect::>(); + if interrupted.is_empty() { + return; + } + let queues = rows + .iter() + .filter(|row| interrupted.contains(&row.id) && row.state == JobState::Available) + .map(|row| row.queue.as_str()) + .collect::>(); + for queue in &queues { + let _ = inner + .queue_notifications + .send(RuntimeNotification::Insert((*queue).to_owned())); + } + match inner.database.pool() { + // SQLite clients wrote these notifications with the completion. + #[cfg(feature = "sqlite")] + DatabasePool::Sqlite(_) => {} + #[cfg(feature = "postgres")] + DatabasePool::Postgres(pool) => { + match inner.postgres_capabilities(pool).await { + Ok(capabilities) if capabilities.supports_listen_notify => {} + Ok(_) => return, + Err(error) => { + debug!(error = %crate::error::Chain(&error), "could not notify peers about interrupted River jobs"); + return; + } + } + for queue in queues { + if let Err(error) = sqlx::query( + "SELECT pg_notify(concat(coalesce($1::text, current_schema()), '.', $2::text), json_build_object('queue', $3::text)::text)", + ) + .bind(inner.schema.as_deref()) + .bind(crate::NOTIFICATION_TOPIC_INSERT) + .bind(queue) + .execute(pool) + .await + { + debug!(error = %crate::error::Chain(&error), queue, "could not notify peers about interrupted River jobs"); + } + } + } + } +} + +/// Applies a batch of state transitions, returning the resulting row for every +/// job that still exists. +/// +/// This is River Go's `JobSetStateIfRunningMany`: a job that is still running +/// takes the requested state (or `cancelled` when a cancellation was attempted +/// and the job would otherwise run again). A job that is no longer running +/// keeps its state, but still receives non-empty metadata updates, and its +/// current row is returned so its event reflects the state that won. +#[allow( + clippy::too_many_lines, + reason = "keeps PostgreSQL batch and transactionally equivalent SQLite completion together" +)] +pub(super) async fn persist_completion_batch( + inner: &ClientInner, + batch: &[CompletionUpdate], +) -> Result, Error> { + match inner.database.pool() { + #[cfg(feature = "sqlite")] + DatabasePool::Sqlite(pool) => { + let mut transaction = crate::database::begin_sqlite_write(pool).await?; + let mut rows = Vec::with_capacity(batch.len()); + let now = Utc::now(); + for update in batch { + let row = crate::database::sqlite::complete_decoded( + &mut transaction, + &crate::database::sqlite::CompleteJob { + attempt: update.attempt, + error: update.error.as_ref(), + finalized_at: update.finalized_at, + id: update.job_id, + metadata_updates: (!update.metadata.is_empty()).then_some(&update.metadata), + now, + scheduled_at: update.scheduled_at, + state: update.state, + }, + ) + .await + .map_err(sqlite_backend_error)?; + let row = match row { + Some(row) => Some(row), + None => crate::database::sqlite::merge_metadata_if_not_running( + &mut transaction, + update.job_id, + &update.metadata, + ) + .await + .map_err(sqlite_backend_error)?, + }; + rows.extend(row.and_then(tolerant_row)); + } + let interrupted_queues = rows + .iter() + .filter(|row| row.state == JobState::Available) + .filter(|row| { + batch.iter().any(|update| { + update.job_id == row.id && update.event_kind == JobEventKind::Interrupted + }) + }) + .map(|row| row.queue.as_str()) + .collect::>(); + for queue in interrupted_queues { + let payload = serde_json::json!({ "queue": queue }).to_string(); + crate::database::sqlite::notification_insert( + &mut transaction, + &[crate::database::sqlite::NotificationInput { + payload: &payload, + topic: crate::NOTIFICATION_TOPIC_INSERT, + }], + ) + .await + .map_err(sqlite_backend_error)?; + } + if inner.pilot.intercepts_job_set_state() { + after_jobs_set_state( + inner, + PilotDatabaseConnection::Sqlite(&mut transaction), + &batch.iter().map(|update| update.job_id).collect::>(), + &rows, + ) + .await?; + } + transaction.commit().await?; + Ok(rows) + } + #[cfg(feature = "postgres")] + DatabasePool::Postgres(pool) => { + let attempt_do_update = batch + .iter() + .map(|update| update.attempt.is_some()) + .collect::>(); + let attempts = batch + .iter() + .map(|update| update.attempt.unwrap_or_default()) + .collect::>(); + let errors = batch + .iter() + .map(|update| update.error.as_ref().map(Json)) + .collect::>(); + let finalized_at = batch + .iter() + .map(|update| update.finalized_at) + .collect::>(); + let ids = batch.iter().map(|update| update.job_id).collect::>(); + let metadata_do_merge = batch + .iter() + .map(|update| !update.metadata.is_empty()) + .collect::>(); + let metadata = batch + .iter() + .map(|update| Json(&update.metadata)) + .collect::>(); + let scheduled_at = batch + .iter() + .map(|update| update.scheduled_at) + .collect::>(); + let states = batch + .iter() + .map(|update| update.state.as_str()) + .collect::>(); + let table = inner.schema.qualify("river_job"); + let state_type = inner.schema.qualify("river_job_state"); + let should_cancel = "(job_input.state IN ('available', 'retryable', 'scheduled') \ + AND job.metadata ? 'cancel_attempted_at')"; + let sql = format!( + "WITH job_input AS (\ + SELECT * FROM unnest(\ + $1::bigint[], $2::boolean[], $3::smallint[], $4::jsonb[], \ + $5::timestamptz[], $6::boolean[], $7::jsonb[], $8::timestamptz[], $9::text[]\ + ) AS job_input(\ + id, attempt_do_update, attempt, errors, finalized_at, \ + metadata_do_merge, metadata_updates, scheduled_at, state)\ + ), updated AS (\ + UPDATE {table} AS job SET \ + attempt = CASE WHEN job.state = 'running' AND NOT {should_cancel} \ + AND job_input.attempt_do_update \ + THEN job_input.attempt ELSE job.attempt END, \ + errors = CASE WHEN job.state = 'running' AND job_input.errors IS NOT NULL \ + THEN array_append(coalesce(job.errors, '{{}}'), job_input.errors) \ + ELSE job.errors END, \ + finalized_at = CASE WHEN job.state = 'running' AND {should_cancel} THEN now() \ + WHEN job.state = 'running' AND job_input.finalized_at IS NOT NULL \ + THEN job_input.finalized_at ELSE job.finalized_at END, \ + metadata = CASE WHEN job_input.metadata_do_merge \ + THEN job.metadata || job_input.metadata_updates ELSE job.metadata END, \ + scheduled_at = CASE WHEN job.state = 'running' AND NOT {should_cancel} \ + AND job_input.scheduled_at IS NOT NULL \ + THEN job_input.scheduled_at ELSE job.scheduled_at END, \ + state = CASE WHEN job.state = 'running' AND {should_cancel} \ + THEN 'cancelled'::{state_type} \ + WHEN job.state = 'running' THEN job_input.state::{state_type} \ + ELSE job.state END \ + FROM job_input \ + WHERE job.id = job_input.id \ + AND (job.state = 'running' OR job_input.metadata_do_merge) \ + RETURNING job.*\ + ) \ + SELECT {projection}, false AS unique_skipped_as_duplicate \ + FROM {table} AS job JOIN job_input ON job.id = job_input.id \ + WHERE NOT EXISTS (SELECT 1 FROM updated WHERE updated.id = job.id) \ + UNION ALL \ + SELECT {projection}, false AS unique_skipped_as_duplicate FROM updated AS job", + projection = job_projection("job"), + ); + let query = sqlx::query(AssertSqlSafe(sql)) + .bind(ids) + .bind(attempt_do_update) + .bind(attempts) + .bind(errors) + .bind(finalized_at) + .bind(metadata_do_merge) + .bind(metadata) + .bind(scheduled_at) + .bind(states); + if inner.pilot.intercepts_job_set_state() { + let mut transaction = crate::database::begin_postgres(pool).await?; + let rows = decode_completion_rows(&query.fetch_all(&mut *transaction).await?); + after_jobs_set_state( + inner, + PilotDatabaseConnection::Postgres(&mut transaction), + &batch.iter().map(|update| update.job_id).collect::>(), + &rows, + ) + .await?; + transaction.commit().await?; + return Ok(rows); + } + Ok(decode_completion_rows(&query.fetch_all(pool).await?)) + } + } +} + +/// Decodes rows returned by a completion. Like River Go's +/// `JobSetStateIfRunningMany`, a row that can't be fully decoded is still +/// returned with its undecodable fields left empty, so its state change is +/// reported like any other. +#[cfg(feature = "postgres")] +fn decode_completion_rows(records: &[PgRow]) -> Vec { + records + .iter() + .filter_map(|row| tolerant_row(decode_job_row(row))) + .collect() +} + +/// Calls the extension hook for rows updated in the current transaction. +pub(crate) async fn after_jobs_set_state( + inner: &ClientInner, + connection: PilotDatabaseConnection<'_>, + job_ids: &[i64], + rows: &[JobRow], +) -> Result<(), Error> { + let params = JobSetStateParams { + database: inner.pilot_database_config(), + job_ids, + jobs: rows, + }; + inner + .pilot + .after_jobs_set_state(connection, ¶ms) + .await + .map_err(|source| Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "job set state", + }, + source, + }) +} + +pub(super) fn finish_batched_completion( + inner: &ClientInner, + update: &CompletionUpdate, + record: Option, +) { + // A stale result settles the peer's ownership just as an applied one + // does, and before its event, so the job can be claimed again at once. + if let Some(peer) = &update.peer { + peer.persisted(inner); + } + if let Some(row) = record { + if let Some(event_kind) = persisted_completion_event_kind(row.state, update.event_kind) { + let event = Event::job_with_statistics( + event_kind, + row, + JobStatistics { + complete_duration: update.timing.completion_started.elapsed(), + queue_wait_duration: update.timing.queue_wait_duration, + run_duration: update.timing.run_duration, + }, + ); + let _ = inner.events.send(event); + } else { + debug!( + job_id = update.job_id, + state = row.state.as_str(), + "job result ignored because the job was moved back to a non-final state" + ); + } + } else { + debug!( + job_id = update.job_id, + "job result ignored because the job no longer exists" + ); + } + remove_running_attempt(&inner.running, update.job_id, &update.cancellation); +} diff --git a/rust/riverqueue/src/client/executor.rs b/rust/riverqueue/src/client/executor.rs new file mode 100644 index 000000000..2aa276516 --- /dev/null +++ b/rust/riverqueue/src/client/executor.rs @@ -0,0 +1,783 @@ +//! Execution of individual job attempts. + +use std::sync::Arc; +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use serde_json::{Map, Value}; +use sha2::{Digest, Sha256}; +use tokio::sync::{mpsc, oneshot}; +use tokio_util::sync::CancellationToken; +use tracing::{Instrument, debug, error, info_span, warn}; + +use crate::client::attempts::remove_running_attempt; +use crate::client::completer::{CompletionAttempt, CompletionTiming, CompletionUpdate}; +use crate::client::{ClientInner, PeerLedger, peers}; +use crate::error::{Chain, panic_message}; +use crate::extension::{WorkEndpoint, WorkNext}; +use crate::{ + AttemptError, BoxError, Client, Error, ErrorHandlerDecision, JobEventKind, JobRow, JobState, + PanicError, WorkCancelled, WorkContext, WorkError, WorkOutcome, WorkResult, WorkerTimeout, +}; + +/// Runs one claimed job's attempt and persists its result. +/// +/// `decode_error` is set for a job whose row couldn't be fully decoded, in +/// which case `row` holds only the fields that could be. Like River Go, such a +/// job isn't worked: its attempt fails with the decode error before hooks or +/// middleware run, the same way as an unknown job kind, and goes through +/// ordinary error handling. The error handler sees the partial row, and the +/// job is retried with the client's retry policy or discarded at its maximum +/// attempts. +#[expect( + clippy::too_many_lines, + reason = "one attempt's decode, work, error handling, and persistence share its state" +)] +pub(super) async fn execute_job( + inner: Arc, + row: JobRow, + decode_error: Option, + hard_cancel: CancellationToken, + cancellation: CancellationToken, + completion_sender: mpsc::Sender, +) { + let span = info_span!("river_job", job_id = row.id, job_kind = %row.kind, queue = %row.queue); + async move { + let queue_wait_duration = row + .attempted_at + .and_then(|attempted_at| { + (attempted_at - row.scheduled_at.max(row.created_at)) + .to_std() + .ok() + }) + .unwrap_or_default(); + // Like River Go's executor start time, which it records as the + // attempt error's `at`. + let attempt_started_at = Utc::now(); + let peers = Arc::new(PeerLedger::new(row.id, attempt_started_at)); + // The worker gets a child of the attempt's token. It observes + // timeouts, remote cancellation, and stops, but cancelling it itself + // (directly or through a drop guard) doesn't look like a remote + // cancellation, just as a Go worker can't cancel its own context. + let context = WorkContext::for_job( + Client { + inner: Arc::clone(&inner), + }, + cancellation.child_token(), + row.id, + &row.metadata, + ) + .with_peers(Arc::clone(&peers)); + let work_started = std::time::Instant::now(); + let mut cancellation_cause = None; + let worked = decode_error.is_none(); + let result = match decode_error { + Some(decode_error) => { + error!(error = %decode_error, "River job row couldn't be decoded; failing attempt without working it"); + Some(Err(worker_failure_from_source( + format!("job row couldn't be decoded: {decode_error}").into(), + ))) + } + None => { + run_worker( + &inner, + &row, + &context, + &hard_cancel, + &cancellation, + &mut cancellation_cause, + ) + .await + } + }; + let Some(result) = result else { + // The task outlived its abort during shutdown and may still be + // running. Leave the row `running` for the rescuer rather than + // making it available to run concurrently with the original. Its + // peers are left to the rescuer as well. + peers.abandon(&inner); + remove_running_attempt(&inner.running, row.id, &cancellation); + return; + }; + + // A cooperative worker can observe cancellation and return before this + // select polls the cancellation branch. Preserve the cancellation cause + // in that race so remote cancellation still gets its canonical outcome. + if cancellation_cause.is_none() && cancellation.is_cancelled() { + cancellation_cause = Some(if hard_cancel.is_cancelled() { + CancellationCause::Shutdown + } else { + CancellationCause::Remote + }); + } + + let run_duration = work_started.elapsed(); + let mut result = result; + if let Some(resumable_failure) = context.resumable_finish(result.is_err()) + && result.is_ok() + { + result = Err(WorkerFailure { + error: Chain(&resumable_failure).to_string(), + kind: WorkerFailureKind::Error, + source: Some(resumable_failure), + trace: String::new(), + }); + } + if cancellation_cause == Some(CancellationCause::Shutdown) + && let Err(failure) = &mut result + && is_soft_stop_failure(failure) + { + failure.error.clear(); + failure.error.push_str("job interrupted by client shutdown"); + failure.kind = WorkerFailureKind::Interrupted; + failure.source = None; + failure.trace.clear(); + } + if cancellation_cause == Some(CancellationCause::Remote) + && !matches!(result, Ok(WorkOutcome::Complete)) + { + result = Err(WorkerFailure { + error: "JobCancelError: job cancelled remotely".to_owned(), + kind: WorkerFailureKind::Cancelled, + source: None, + trace: String::new(), + }); + } + let work_result = public_work_result(&result); + let mut error_handler_result = ErrorHandlerDecision::default(); + if let Some(error_handler) = &inner.error_handler + && matches!( + work_result, + WorkResult::Aborted | WorkResult::Failed(_) | WorkResult::Panicked(_) + ) + { + match error_handler + .handle_error(&context, &row, &work_result) + .await + { + Ok(handler_result) => error_handler_result = handler_result, + Err(handler_error) => { + error!(error = %crate::error::Chain(&handler_error), "River error handler failed"); + } + } + } + // Peers settle before the coordinator's own outcome, so its + // producer's `job_finished` comes after them. + peers + .finish( + &inner, + &context, + cancellation_cause == Some(CancellationCause::Shutdown), + ) + .await; + let metadata_updates = context.metadata_updates(); + let completion = CompletionAttempt { + cancellation: cancellation.clone(), + timing: CompletionTiming { + completion_started: std::time::Instant::now(), + queue_wait_duration, + run_duration, + }, + }; + let persisted = persist_result( + &inner, + &row, + attempt_started_at, + &completion, + result, + metadata_updates, + error_handler_result, + worked, + &completion_sender, + None, + ) + .await; + // Once enqueued, the completer owns the running attempt until the + // result is written. + if let Err(operation_error) = persisted { + error!(error = %crate::error::Chain(&operation_error), "failed to persist River job result"); + remove_running_attempt(&inner.running, row.id, &cancellation); + } + } + .instrument(span) + .await; +} + +/// Works a job, returning its result, or `None` when the worker task outlived +/// its abort during shutdown. +async fn run_worker( + inner: &Arc, + row: &JobRow, + context: &WorkContext, + hard_cancel: &CancellationToken, + cancellation: &CancellationToken, + cancellation_cause: &mut Option, +) -> Option { + let worker_row = row.clone(); + let worker_context = context.clone(); + let worker_inner = Arc::clone(inner); + let (timeout_sender, timeout_receiver) = oneshot::channel(); + // The worker runs in its own task so a panic or an abort can't take the + // executor with it. It stays inside this job's span. + let mut worker_task = AbortOnDrop(tokio::spawn( + async move { + worker_context.resumable_validate()?; + // Like River Go, an unknown kind fails before any middleware or + // hook runs. + worker_inner.workers.check_kind(&worker_row)?; + let hooks_context = worker_context.clone(); + let hooks_inner = Arc::clone(&worker_inner); + let endpoint: WorkEndpoint<'_> = Box::new(move |mut job: JobRow| { + Box::pin(async move { + for hook in &hooks_inner.hooks { + hook.work_begin(&hooks_context, &mut job) + .await + .map_err(WorkError::new)?; + } + // The worker consumes the row, so it's copied only when + // work-end hooks need it afterward. + if hooks_inner.hooks.is_empty() { + return hooks_inner + .workers + .work(hooks_context.clone(), job, timeout_sender) + .await?; + } + let mut result = hooks_inner + .workers + .work(hooks_context.clone(), job.clone(), timeout_sender) + .await?; + for hook in &hooks_inner.hooks { + result = hook.work_end(&hooks_context, &job, result).await; + } + result + }) + }); + WorkNext::new(&worker_inner.work_middleware, &worker_context, endpoint) + .run(worker_row) + .await + } + .in_current_span(), + )); + + // The worker reports its timeout after decoding the job's arguments, + // following any hooks and middleware, so the timeout covers the work + // itself as in River Go. + let timeout_elapsed = async { + let timeout = match timeout_receiver.await { + Ok(WorkerTimeout::After(timeout)) => Some(timeout), + Ok(WorkerTimeout::ClientDefault) => inner.job_timeout, + // Disabled, or the attempt ended before work started. + Ok(WorkerTimeout::Disabled) | Err(_) => None, + }; + match timeout { + Some(timeout) => { + tokio::time::sleep(timeout).await; + timeout + } + None => std::future::pending().await, + } + }; + tokio::select! { + result = &mut worker_task.0 => Some(worker_join_result(result)), + () = cancellation.cancelled() => { + *cancellation_cause = Some(if hard_cancel.is_cancelled() { + CancellationCause::Shutdown + } else { + CancellationCause::Remote + }); + finish_cancelled_task(inner, row, &mut worker_task.0, hard_cancel, None).await + } + timeout = timeout_elapsed => { + *cancellation_cause = Some(CancellationCause::Timeout); + cancellation.cancel(); + finish_cancelled_task(inner, row, &mut worker_task.0, hard_cancel, Some(timeout)).await + } + } +} + +pub(super) type WorkerResult = Result; + +#[derive(Debug)] +pub(super) struct WorkerFailure { + pub(super) error: String, + pub(super) kind: WorkerFailureKind, + pub(super) source: Option, + pub(super) trace: String, +} + +#[derive(Debug)] +pub(super) enum WorkerFailureKind { + Aborted, + Cancelled, + Error, + Interrupted, + Panic, +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(super) enum CancellationCause { + Remote, + Shutdown, + Timeout, +} + +/// Aborts a spawned worker task when the executor that owns it is dropped, so +/// a stopped client never leaves work running detached from its runtime. +pub(super) struct AbortOnDrop(pub(super) tokio::task::JoinHandle); + +impl Drop for AbortOnDrop { + fn drop(&mut self) { + self.0.abort(); + } +} + +/// How long shutdown waits for an aborted task before abandoning it. +const ABORT_GRACE_DURING_SHUTDOWN: Duration = Duration::from_millis(100); + +/// Waits for a cancelled job to return, then treats it as stuck. +/// +/// `timeout` is the timeout that cancelled the job, if one did: the worker's +/// own timeout when it sets one, otherwise the client's. +/// +/// After `job_stuck_threshold`, the stuck handler runs and the task is +/// aborted. The aborted attempt fails with "job aborted after ignoring +/// cancellation", during a hard stop as well, so it counts and is retried or +/// discarded like any failed attempt. Tokio abort only takes effect at the +/// task's next `.await`, so a task blocked in synchronous code keeps its +/// worker slot until it actually ends: the queue never exceeds +/// `max_workers`, and the job is not persisted (and so cannot be fetched +/// again) while the original may still be running. +/// During shutdown the executor gives up after a short grace period and +/// returns `None`, leaving the row `running` for the rescuer. +pub(super) async fn finish_cancelled_task( + inner: &ClientInner, + row: &JobRow, + worker_task: &mut tokio::task::JoinHandle>, + hard_cancel: &CancellationToken, + timeout: Option, +) -> Option { + let stuck_threshold = inner.job_stuck_threshold; + if let Ok(result) = tokio::time::timeout(stuck_threshold, &mut *worker_task).await { + return Some(worker_join_result(result)); + } + // Like Go's stuck job log line, report the timeout that applied to this + // job, which is the worker's own when it sets one. + warn!( + job_id = row.id, + kind = %row.kind, + ?timeout, + ?stuck_threshold, + "River job remained active after cancellation; treating it as stuck and aborting its task" + ); + if let Some(error_handler) = &inner.error_handler + && let Err(handler_error) = error_handler.handle_stuck(row).await + { + error!(error = %crate::error::Chain(&handler_error), "River stuck handler failed"); + } + worker_task.abort(); + let result = tokio::select! { + result = &mut *worker_task => Some(result), + () = async { + hard_cancel.cancelled().await; + tokio::time::sleep(ABORT_GRACE_DURING_SHUTDOWN).await; + } => None, + }; + let Some(result) = result else { + error!( + "River job remained stuck after its task was aborted during shutdown; leaving it running for the rescuer" + ); + return None; + }; + Some(match result { + Err(join_error) if join_error.is_cancelled() => Err(WorkerFailure { + error: "job aborted after ignoring cancellation".to_owned(), + kind: WorkerFailureKind::Aborted, + source: None, + trace: String::new(), + }), + result => worker_join_result(result), + }) +} + +/// Whether a failure during hard shutdown is the job stopping because the +/// client cancelled it, mirroring River Go's `isSoftStopCancelError`. +/// +/// A worker that returns [`WorkCancelled`] (anywhere in its error's source +/// chain) stopped cooperatively. A task River aborted because it still +/// ignored cancellation after the stuck threshold didn't stop on its own: it +/// had that long to respond, so its attempt fails and follows the ordinary +/// retry path like panics and other returned errors. +pub(super) fn is_soft_stop_failure(failure: &WorkerFailure) -> bool { + match failure.kind { + WorkerFailureKind::Error => failure + .source + .as_ref() + .is_some_and(|error| WorkCancelled::is_in_chain(error.get_ref())), + WorkerFailureKind::Aborted + | WorkerFailureKind::Cancelled + | WorkerFailureKind::Interrupted + | WorkerFailureKind::Panic => false, + } +} + +pub(super) fn worker_join_result( + result: Result, tokio::task::JoinError>, +) -> WorkerResult { + match result { + Ok(Ok(outcome)) => Ok(outcome), + Ok(Err(worker_error)) => Err(WorkerFailure { + error: Chain(&worker_error).to_string(), + kind: WorkerFailureKind::Error, + source: Some(worker_error), + trace: String::new(), + }), + // Like Go, a panic is recorded as the panic's value alone. + Err(join_error) if join_error.is_panic() => { + let trace = format!("{join_error:?}"); + Err(WorkerFailure { + error: panic_message(&join_error.into_panic()).to_owned(), + kind: WorkerFailureKind::Panic, + source: None, + trace, + }) + } + Err(join_error) => Err(WorkerFailure { + error: format!("job task cancelled: {join_error}"), + kind: WorkerFailureKind::Aborted, + source: None, + trace: format!("{join_error:?}"), + }), + } +} + +pub(super) fn worker_failure_from_source(error: BoxError) -> WorkerFailure { + let error = WorkError::new(error); + WorkerFailure { + error: Chain(&error).to_string(), + kind: WorkerFailureKind::Error, + source: Some(error), + trace: String::new(), + } +} + +pub(super) fn public_work_result(result: &WorkerResult) -> WorkResult { + match result { + Ok(WorkOutcome::Cancel) => WorkResult::Cancelled, + Ok(WorkOutcome::Complete) => WorkResult::Completed, + Ok(WorkOutcome::Discard) => WorkResult::Discarded, + Ok(WorkOutcome::Snooze(duration)) => WorkResult::Snoozed(*duration), + Err(failure) => match failure.kind { + WorkerFailureKind::Aborted => WorkResult::Aborted, + WorkerFailureKind::Cancelled => WorkResult::Cancelled, + WorkerFailureKind::Error => { + WorkResult::Failed(failure.source.clone().unwrap_or_else(|| { + WorkError::new(Box::new(std::io::Error::other(failure.error.clone()))) + })) + } + WorkerFailureKind::Interrupted => WorkResult::Interrupted, + WorkerFailureKind::Panic => { + WorkResult::Panicked(PanicError::new(failure.error.clone())) + } + }, + } +} + +#[expect( + clippy::too_many_arguments, + clippy::too_many_lines, + reason = "maps every attempt outcome to its completion from the attempt's recorded state" +)] +pub(super) async fn persist_result( + inner: &ClientInner, + row: &JobRow, + attempt_started_at: DateTime, + completion: &CompletionAttempt, + result: WorkerResult, + metadata_updates: Map, + error_handler_result: ErrorHandlerDecision, + worked: bool, + completion_sender: &mpsc::Sender, + peer: Option>, +) -> Result<(), Error> { + let now = Utc::now(); + let (state, finalized_at, scheduled_at, attempt, attempt_error, metadata, event_kind) = + match result { + Ok(WorkOutcome::Complete) => ( + JobState::Completed, + Some(now), + None, + None, + None, + metadata_updates, + JobEventKind::Completed, + ), + Ok(WorkOutcome::Cancel) => ( + JobState::Cancelled, + Some(now), + None, + None, + Some(AttemptError { + at: attempt_started_at, + attempt: row.attempt, + error: "job cancelled by worker".to_owned(), + trace: String::new(), + }), + metadata_updates, + JobEventKind::Cancelled, + ), + Ok(WorkOutcome::Discard) => ( + JobState::Discarded, + Some(now), + None, + None, + Some(AttemptError { + at: attempt_started_at, + attempt: row.attempt, + error: "job discarded by worker".to_owned(), + trace: String::new(), + }), + metadata_updates, + JobEventKind::Failed, + ), + Ok(WorkOutcome::Snooze(duration)) => { + let scheduled_at = scheduled_after(now, duration); + let state = if duration <= inner.maintenance.scheduler_interval { + JobState::Available + } else { + JobState::Scheduled + }; + let mut metadata = metadata_updates; + let snoozes = go_json_int(row.metadata.get_raw("snoozes")).wrapping_add(1); + metadata.insert("snoozes".to_owned(), Value::from(snoozes)); + ( + state, + None, + Some(scheduled_at), + Some(row.attempt - 1), + None, + metadata, + JobEventKind::Snoozed, + ) + } + // River Go's `JobSetStateInterrupted`: make the job available now + // without recording an error or counting the attempt. The + // completer keeps `attempted_at` and still honors a cancellation + // that was attempted while the job ran. + Err(failure) if matches!(failure.kind, WorkerFailureKind::Interrupted) => ( + JobState::Available, + None, + Some(now), + Some((row.attempt - 1).max(0)), + None, + metadata_updates, + JobEventKind::Interrupted, + ), + Err(failure) => { + let retry_error = failure.source.clone().unwrap_or_else(|| { + WorkError::new(Box::new(std::io::Error::other(failure.error.clone()))) + }); + let attempt_error = AttemptError { + at: attempt_started_at, + attempt: row.attempt, + error: failure.error, + trace: failure.trace, + }; + if matches!(failure.kind, WorkerFailureKind::Cancelled) + || error_handler_result == ErrorHandlerDecision::Cancel + { + ( + JobState::Cancelled, + Some(now), + None, + None, + Some(attempt_error), + metadata_updates, + JobEventKind::Cancelled, + ) + } else if row.attempt >= row.max_attempts { + ( + JobState::Discarded, + Some(now), + None, + None, + Some(attempt_error), + metadata_updates, + JobEventKind::Failed, + ) + } else { + // Like River Go, a job that wasn't worked because its row + // couldn't be decoded uses only the client's retry policy. + let worker_retry_after = if worked { + inner + .workers + .next_retry(row, &retry_error, now) + .unwrap_or_else(|retry_error| { + debug!(error = %retry_error, "could not evaluate worker retry override"); + None + }) + } else { + None + }; + let delay = worker_retry_after + .unwrap_or_else(|| inner.retry_policy.next_retry(row, &retry_error, now)); + let scheduled_at = scheduled_after(now, delay); + let state = if delay <= inner.maintenance.scheduler_interval { + JobState::Available + } else { + JobState::Retryable + }; + ( + state, + None, + Some(scheduled_at), + None, + Some(attempt_error), + metadata_updates, + JobEventKind::Failed, + ) + } + } + }; + + completion_sender + .send(CompletionUpdate { + attempt, + cancellation: completion.cancellation.clone(), + error: attempt_error, + event_kind, + finalized_at, + job_id: row.id, + metadata, + scheduled_at, + state, + timing: completion.timing, + peer, + }) + .await + .map_err(|_| Error::runtime_context("job completion", "completion batcher stopped")) +} + +/// Longest delay River schedules ahead, matching Go's `time.Duration` range. +const MAX_SCHEDULE_DELAY: Duration = Duration::from_nanos(i64::MAX.cast_unsigned()); + +/// Adds a snooze or retry delay to `now`, clamping out-of-range delays the way +/// River Go's `time.Duration` arithmetic bounds them instead of failing. +pub(super) fn scheduled_after(now: DateTime, delay: Duration) -> DateTime { + chrono::Duration::from_std(delay.min(MAX_SCHEDULE_DELAY)) + .ok() + .and_then(|delay| now.checked_add_signed(delay)) + .unwrap_or(DateTime::::MAX_UTC) +} + +pub(crate) fn default_retry_delay(row: &JobRow, now: DateTime, seed: u64) -> Duration { + const MAX_RETRY_NANOS: u64 = i64::MAX as u64; + + let error_count = u32::try_from(row.errors.len().saturating_add(1)).unwrap_or(u32::MAX); + let base_seconds = u128::from(error_count).pow(4); + if base_seconds.saturating_mul(1_000_000_000) >= u128::from(MAX_RETRY_NANOS) { + return Duration::from_nanos(MAX_RETRY_NANOS); + } + let base_seconds = u64::try_from(base_seconds).expect("capped retry seconds fit u64"); + let base = Duration::from_secs(base_seconds); + + let mut hasher = Sha256::new(); + hasher.update(seed.to_be_bytes()); + hasher.update(row.id.to_be_bytes()); + hasher.update(error_count.to_be_bytes()); + hasher.update(now.timestamp_nanos_opt().unwrap_or_default().to_be_bytes()); + let hash = hasher.finalize(); + let sample = u32::from_be_bytes(hash[..4].try_into().unwrap()); + let ratio = f64::from(sample) / f64::from(u32::MAX); + // Jitter can push a delay just below the cap past it; Go caps after + // jitter as well. + base.mul_f64(0.9 + ratio * 0.2) + .min(Duration::from_nanos(MAX_RETRY_NANOS)) +} + +/// Coerces a metadata value to an integer exactly like Go's `gjson.Int`, which +/// the Go executor uses to read the `snoozes` counter. Numbers truncate toward +/// zero, numeric strings of optional sign and digits parse, `true` is one, and +/// everything else is zero. +fn go_json_int(value: Option<&serde_json::value::RawValue>) -> i64 { + fn parse_digits(text: &str) -> Option { + let (negative, digits) = text + .strip_prefix('-') + .map_or((false, text), |digits| (true, digits)); + if digits.is_empty() { + return None; + } + let mut number = 0_i64; + for byte in digits.bytes() { + if !byte.is_ascii_digit() { + return None; + } + number = number.wrapping_mul(10).wrapping_add(i64::from(byte - b'0')); + } + Some(if negative { + number.wrapping_neg() + } else { + number + }) + } + + const MAX_SAFE_INTEGER: f64 = 9_007_199_254_740_991.0; + match value.map(serde_json::value::RawValue::get) { + Some("true") => 1, + Some(raw) if raw.starts_with('"') => serde_json::from_str::(raw) + .ok() + .and_then(|text| parse_digits(&text)) + .unwrap_or(0), + Some(raw) if raw.starts_with(['-', '0', '1', '2', '3', '4', '5', '6', '7', '8', '9']) => { + let float = raw.parse::().unwrap_or(0.0); + if (-MAX_SAFE_INTEGER..=MAX_SAFE_INTEGER).contains(&float) { + #[expect( + clippy::cast_possible_truncation, + reason = "Go truncates safe floats toward zero" + )] + return float as i64; + } + #[expect( + clippy::cast_possible_truncation, + reason = "Go falls back to a float conversion for huge numbers" + )] + parse_digits(raw).unwrap_or(float as i64) + } + _ => 0, + } +} + +#[cfg(test)] +mod go_json_int_tests { + use crate::JobMetadata; + use serde::Deserialize; + + use super::go_json_int; + + #[derive(Deserialize)] + struct Fixture { + snooze_counters: Vec, + } + + #[derive(Deserialize)] + struct SnoozeCounter { + expected_snoozes: i64, + metadata: JobMetadata, + name: String, + } + + #[test] + fn snooze_counter_matches_go_fixture() { + let fixture: Fixture = + serde_json::from_str(include_str!("../../tests/fixtures/maintenance_values.json")) + .unwrap(); + assert!(!fixture.snooze_counters.is_empty()); + for case in fixture.snooze_counters { + assert_eq!( + go_json_int(case.metadata.get_raw("snoozes")).wrapping_add(1), + case.expected_snoozes, + "{}", + case.name + ); + } + } +} diff --git a/rust/riverqueue/src/client/extension.rs b/rust/riverqueue/src/client/extension.rs new file mode 100644 index 000000000..e1db13b80 --- /dev/null +++ b/rust/riverqueue/src/client/extension.rs @@ -0,0 +1,367 @@ +//! Unstable extension entry points used by companion crates. + +use std::sync::Arc; +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use serde_json::Map; +use serde_json::value::RawValue; + +use crate::__private::DatabaseConnection as PilotDatabaseConnection; +use crate::__private::{Pilot, PreparedInsertParams, RawInsertResult}; +use crate::client::WeakClient; +use crate::client::request::{Target, request_type}; +use crate::client::validate::validate_insert_parts; +use crate::{ + Client, Error, InsertContext, InsertOpts, InsertParams, JobArgs, JobRow, JobState, WorkError, +}; + +/// Client operations reserved for River's own companion crates. +/// +/// Reached through `riverqueue::__private`, this wrapper keeps these +/// operations off [`Client`]'s public API. +#[derive(Clone, Copy, Debug)] +pub struct ExtensionClient<'client> { + client: &'client Client, +} + +impl<'client> ExtensionClient<'client> { + /// Wraps a client. + #[must_use] + pub const fn new(client: &'client Client) -> Self { + Self { client } + } + + /// Returns the wrapped client. + #[must_use] + pub const fn client(&self) -> &'client Client { + self.client + } + + /// Creates a non-owning handle for an extension service. + #[must_use] + pub fn downgrade(&self) -> WeakClient { + self.client.downgrade() + } + + /// Returns the client's pilot when it's a `P`, so a companion crate can + /// reach the pilot it installed from any handle to the client, such as + /// a worker's [`WorkContext::client`](crate::WorkContext::client). + #[must_use] + pub fn pilot(&self) -> Option> { + let pilot: Arc = Arc::clone(&self.client.inner.pilot); + let pilot: Arc = pilot; + pilot.downcast::

().ok() + } + + /// Resolves typed insertion options the same way a typed insert does. + #[must_use] + pub fn resolve_insert_opts(&self, opts: InsertOpts) -> InsertParams { + self.client.resolve_insert_opts::(opts) + } +} + +impl<'client> ExtensionClient<'client> { + /// Computes the configured retry delay for an exact-version extension. + #[must_use] + pub fn retry_delay(&self, row: &JobRow, error: &WorkError, now: DateTime) -> Duration { + self.client.inner.retry_policy.next_retry(row, error, now) + } + + /// Returns the scheduler horizon used by exact-version completion helpers. + #[must_use] + pub fn scheduler_interval(&self) -> Duration { + self.client.inner.maintenance.scheduler_interval + } + + /// Inserts an encoded job through River's exact-version extension seam, + /// with `opts` resolved against the client's defaults the same way a + /// typed insert resolves them. + /// + /// The request runs when awaited, in a caller-managed transaction with + /// [`tx`](RawInsertRequest::tx). Awaiting it returns the errors of an + /// ordinary insertion: invalid options, an unregistered kind, an + /// extension failure, a transaction from another backend, or a database + /// error. + pub fn insert_raw<'a>( + &self, + kind: &'a str, + unique_fields: &'a [&'a [&'a str]], + encoded_args: Box, + opts: InsertOpts, + ) -> RawInsertRequest<'a> + where + 'client: 'a, + { + self.raw_request( + kind, + unique_fields, + encoded_args, + RawInsertOptions::Opts(opts), + ) + } + + /// Like [`insert_raw`](Self::insert_raw), with already-resolved + /// insertion parameters. + pub fn insert_raw_params<'a>( + &self, + kind: &'a str, + unique_fields: &'a [&'a [&'a str]], + encoded_args: Box, + params: InsertParams, + ) -> RawInsertRequest<'a> + where + 'client: 'a, + { + self.raw_request( + kind, + unique_fields, + encoded_args, + RawInsertOptions::Params(params), + ) + } + + /// Like [`insert_raw_params`](Self::insert_raw_params), inserting the job + /// as the occurrence of a periodic job due at `target`, exactly as + /// River's periodic job enqueuer does. + /// + /// When `params.scheduled_at` is unset, the job is inserted `available` + /// with `scheduled_at` set to `target` so it runs immediately, and a + /// `by_period` unique key uses the target's period. An explicit + /// `scheduled_at` inserts a `scheduled` job, and `pending` is kept. + pub fn insert_periodic<'a>( + &self, + kind: &'a str, + unique_fields: &'a [&'a [&'a str]], + encoded_args: Box, + params: InsertParams, + target: DateTime, + ) -> RawInsertRequest<'a> + where + 'client: 'a, + { + self.raw_request( + kind, + unique_fields, + encoded_args, + RawInsertOptions::Periodic { params, target }, + ) + } + + fn raw_request<'a>( + self, + kind: &'a str, + unique_fields: &'a [&'a [&'a str]], + encoded_args: Box, + options: RawInsertOptions, + ) -> RawInsertRequest<'a> + where + 'client: 'a, + { + RawInsertRequest { + client: self.client, + encoded_args, + kind, + options, + target: Target::Client, + unique_fields, + } + } + + /// Inserts stored jobs again, such as jobs set aside and retried later, + /// the way an ordinary batch insertion runs: insert middleware, begin + /// hooks, the extension's insertion step, and notifications run once, in + /// one transaction. See [`PreparedInsertParams`] for what the jobs keep. + /// + /// The request runs when awaited, in a caller-managed transaction with + /// [`tx`](PreparedInsertRequest::tx). Awaiting it returns the errors of + /// an ordinary insertion: invalid parameters, an extension failure, a + /// transaction from another backend, or a database error. + pub fn insert_prepared<'a>( + &self, + params: Vec, + ) -> PreparedInsertRequest<'a> + where + 'client: 'a, + { + PreparedInsertRequest { + client: self.client, + params, + target: Target::Client, + } + } + + /// Validates stored jobs and turns them into insertions that keep their + /// unique key and states, creation time, and schedule. + pub(crate) fn prepared_jobs( + params: Vec, + ) -> Result, Error> { + params + .into_iter() + .map(|params| { + let unique_states = match (¶ms.unique_key, ¶ms.unique_states) { + (None, None) => None, + (Some(_), Some(states)) => Some( + states + .iter() + .fold(0, |bitmask, state| bitmask | state.unique_bit()), + ), + _ => { + return Err(Error::invalid_job_context( + "prepared insertion", + "unique_key and unique_states must either both be set or both be absent" + .to_owned(), + )); + } + }; + let opts = InsertParams { + extension_options: Map::new(), + max_attempts: params.max_attempts, + metadata: params.metadata, + pending: false, + priority: params.priority, + queue: params.queue, + scheduled_at: Some(params.scheduled_at), + tags: params.tags, + unique: crate::UniqueOpts::default(), + }; + // A stored job's kind was accepted when it was first + // inserted, possibly by an older client, so only its options + // are checked again. + validate_insert_parts(¶ms.kind, &opts, true)?; + Ok(InsertContext { + encoded_args: params.encoded_args, + kind: params.kind, + opts, + state: JobState::Available, + created_at: Some(params.created_at), + unique_key: params.unique_key, + unique_states, + }) + }) + .collect() + } +} + +async fn insert_raw_jobs( + client: &Client, + executor: Option>, + jobs: Vec, +) -> Result, Error> { + client + .run_insert(executor, jobs, |rows| { + Ok(rows + .into_iter() + .map(|row| RawInsertResult { + job: row.job, + unique_skipped_as_duplicate: row.unique_skipped_as_duplicate, + }) + .collect()) + }) + .await +} + +async fn insert_raw_job( + client: &Client, + executor: Option>, + job: InsertContext, +) -> Result { + client + .run_insert(executor, vec![job], |rows| { + let row = rows.into_iter().next().ok_or_else(|| { + Error::runtime_context("exact-version insertion", "insertion returned no row") + })?; + Ok(RawInsertResult { + job: row.job, + unique_skipped_as_duplicate: row.unique_skipped_as_duplicate, + }) + }) + .await +} + +/// How a [`RawInsertRequest`] sets the job's options. +#[derive(Debug)] +enum RawInsertOptions { + /// Options resolved against the client's defaults, like a typed insert. + Opts(InsertOpts), + /// Already-resolved parameters. + Params(InsertParams), + /// A periodic occurrence due at `target`. + Periodic { + params: InsertParams, + target: DateTime, + }, +} + +request_type! { + /// An encoded insertion through River's exact-version extension seam, + /// returned by [`ExtensionClient::insert_raw`], + /// [`ExtensionClient::insert_raw_params`], or + /// [`ExtensionClient::insert_periodic`]. Await it to insert. + write RawInsertRequest { + encoded_args: Box, + kind: &'a str, + options: RawInsertOptions, + unique_fields: &'a [&'a [&'a str]], + } -> RawInsertResult +} + +impl RawInsertRequest<'_> { + async fn run(self) -> Result { + let executor = self.target.into_executor()?; + let client = self.client; + client.validate_known_kind(self.kind)?; + let now = Utc::now(); + let job = match self.options { + RawInsertOptions::Opts(opts) => { + let params = InsertOpts::resolve( + client.inner.default_max_attempts, + InsertOpts::default(), + opts, + ); + client.prepare_encoded( + self.kind, + self.unique_fields, + self.encoded_args, + params, + now, + )? + } + RawInsertOptions::Params(params) => client.prepare_encoded( + self.kind, + self.unique_fields, + self.encoded_args, + params, + now, + )?, + RawInsertOptions::Periodic { params, target } => client.prepare_periodic( + self.kind, + self.unique_fields, + self.encoded_args, + params, + target, + now, + )?, + }; + insert_raw_job(client, executor, job).await + } +} + +request_type! { + /// A reinsertion of stored jobs, returned by + /// [`ExtensionClient::insert_prepared`]. Await it to insert. + write PreparedInsertRequest { + params: Vec, + } -> Vec +} + +impl PreparedInsertRequest<'_> { + async fn run(self) -> Result, Error> { + let executor = self.target.into_executor()?; + if self.params.is_empty() { + return Ok(Vec::new()); + } + let jobs = ExtensionClient::prepared_jobs(self.params)?; + insert_raw_jobs(self.client, executor, jobs).await + } +} diff --git a/rust/riverqueue/src/client/insert.rs b/rust/riverqueue/src/client/insert.rs new file mode 100644 index 000000000..aff2bd35b --- /dev/null +++ b/rust/riverqueue/src/client/insert.rs @@ -0,0 +1,1001 @@ +//! Job insertion. +//! +//! Every insertion, whether a single typed job, a homogeneous or +//! heterogeneous batch, a periodic job, or an extension's raw insert, runs +//! through one pipeline that mirrors River Go's `insertManyShared`: +//! +//! 1. Options are resolved and validated, and the unique key and initial +//! state are computed from the original arguments. +//! 2. Insertion middleware wraps the rest of the operation. +//! 3. Inside the middleware, begin hooks and any extension interception run +//! for each job, the jobs are written, and one insert notification is sent +//! per queue that gained available jobs. +//! 4. Decode hooks run on returned rows. + +#[cfg(feature = "postgres")] +use std::borrow::Cow; +use std::{ + fmt, + future::{Future, IntoFuture}, + pin::Pin, +}; + +use chrono::{DateTime, Utc}; +#[cfg(feature = "postgres")] +use serde_json::Value; +use serde_json::value::RawValue; +#[cfg(feature = "postgres")] +use sqlx::AssertSqlSafe; +#[cfg(feature = "postgres")] +use sqlx::types::Json; + +use crate::__private::{ + DatabaseConnection as PilotDatabaseConnection, JobInsertParams as PilotJobInsertParams, +}; +use crate::client::RuntimeNotification; +use crate::client::request::Target; +#[cfg(feature = "sqlite")] +use crate::client::sqlite_backend_error; +use crate::client::validate::validate_insert_parts; +#[cfg(feature = "postgres")] +use crate::client::{JobRecord, job_projection}; +#[cfg(feature = "postgres")] +use crate::database::postgres_capabilities::UniqueInsertMode; +use crate::database::{DatabasePool, DatabaseTransactionExecutor}; +use crate::extension::{InsertEndpoint, InsertNext, InsertedJob}; +use crate::periodic::PeriodicInsert; +use crate::unique::build_unique_key_parts; +use crate::{ + Client, Error, InsertBatch, InsertBatchResult, InsertContext, InsertOpts, InsertParams, + InsertResult, Job, JobArgs, JobRow, JobState, +}; +#[cfg(feature = "postgres")] +use crate::{JobMetadata, METADATA_KEY_UNIQUE_NONCE}; + +/// One job of a homogeneous [`Client::insert_many`] batch: arguments plus +/// options that override the job type's defaults. +/// +/// Batches accept bare arguments or `(args, opts)` tuples, both of which +/// convert into this type. +#[derive(Debug)] +pub struct InsertManyItem { + args: A, + opts: InsertOpts, +} + +impl InsertManyItem { + /// Pairs job arguments with insertion options. + pub const fn new(args: A, opts: InsertOpts) -> Self { + Self { args, opts } + } +} + +impl From for InsertManyItem { + fn from(args: A) -> Self { + Self::new(args, InsertOpts::default()) + } +} + +impl From<(A, InsertOpts)> for InsertManyItem { + fn from((args, opts): (A, InsertOpts)) -> Self { + Self::new(args, opts) + } +} + +/// A single-job insertion, returned by [`Client::insert`]. Await it to insert +/// the job. +/// +/// The job type's defaults, the client's defaults, and River's defaults apply +/// unless overridden with [`opts`](Self::opts). +#[must_use = "insert requests do nothing unless awaited"] +pub struct InsertRequest<'a, A> { + args: A, + client: &'a Client, + opts: InsertOpts, + target: Target<'a>, +} + +impl<'a, A: JobArgs> InsertRequest<'a, A> { + /// Overrides options for this job. Options not set here fall back to + /// the job type's defaults. + pub fn opts(mut self, opts: InsertOpts) -> Self { + self.opts = opts; + self + } + + /// Inserts the job in a caller-managed transaction. + /// + /// The job becomes visible to workers only when the transaction commits + /// and is discarded if it rolls back. `executor` must be a SQLx + /// transaction for the client's database backend. River runs directly in + /// that transaction without a savepoint, so an insertion that returns an + /// error, including one after the write, may leave partial changes in + /// it: roll the transaction back, or open your own savepoint around the + /// insertion if you need to continue it. + pub fn tx<'t, E>(self, executor: E) -> InsertRequest<'t, A> + where + 'a: 't, + E: DatabaseTransactionExecutor<'t>, + { + InsertRequest { + args: self.args, + client: self.client, + opts: self.opts, + target: Target::transaction(self.client, executor), + } + } +} + +impl fmt::Debug for InsertRequest<'_, A> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("InsertRequest") + .field("opts", &self.opts) + .finish_non_exhaustive() + } +} + +impl<'a, A: JobArgs> IntoFuture for InsertRequest<'a, A> { + type Output = Result, Error>; + type IntoFuture = Pin + Send + 'a>>; + + fn into_future(self) -> Self::IntoFuture { + Box::pin(async move { + let Self { + args, + client, + opts, + target, + } = self; + let job = client.prepare_typed(&args, opts, Utc::now())?; + client + .run_insert(target.into_executor()?, vec![job], |rows| { + let row = rows.into_iter().next().ok_or_else(|| { + Error::runtime_context("job insertion", "insertion returned no row") + })?; + let args = row.job.decode_args()?; + Ok(InsertResult { + job: Job { args, row: row.job }, + unique_skipped_as_duplicate: row.unique_skipped_as_duplicate, + }) + }) + .await + }) + } +} + +/// An atomic homogeneous batch insertion, returned by +/// [`Client::insert_many`]. Await it to insert the jobs and get one result +/// per job, in input order. +#[must_use = "insert requests do nothing unless awaited"] +pub struct InsertManyRequest<'a, A> { + client: &'a Client, + jobs: Vec>, + target: Target<'a>, +} + +impl<'a, A: JobArgs> InsertManyRequest<'a, A> { + /// Inserts the jobs in a caller-managed transaction. + /// + /// The jobs become visible to workers only when the transaction commits. + /// River runs directly in the transaction without a savepoint, so a batch + /// that fails, including after the write, may leave partial changes in + /// it: roll the transaction back, or open your own savepoint around the + /// insertion if you need to continue it. + pub fn tx<'t, E>(self, executor: E) -> InsertManyRequest<'t, A> + where + 'a: 't, + E: DatabaseTransactionExecutor<'t>, + { + InsertManyRequest { + client: self.client, + jobs: self.jobs, + target: Target::transaction(self.client, executor), + } + } + + fn prepare(client: &Client, jobs: Vec>) -> Result, Error> { + let now = Utc::now(); + jobs.into_iter() + .map(|item| client.prepare_typed(&item.args, item.opts, now)) + .collect() + } +} + +impl fmt::Debug for InsertManyRequest<'_, A> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("InsertManyRequest") + .field("jobs", &self.jobs.len()) + .finish_non_exhaustive() + } +} + +impl<'a, A: JobArgs> IntoFuture for InsertManyRequest<'a, A> { + type Output = Result>, Error>; + type IntoFuture = Pin + Send + 'a>>; + + fn into_future(self) -> Self::IntoFuture { + Box::pin(async move { + let Self { + client, + jobs, + target, + } = self; + let jobs = Self::prepare(client, jobs)?; + client + .run_insert(target.into_executor()?, jobs, |rows| { + rows.into_iter() + .map(|row| { + let args = row.job.decode_args()?; + Ok(InsertResult { + job: Job { args, row: row.job }, + unique_skipped_as_duplicate: row.unique_skipped_as_duplicate, + }) + }) + .collect() + }) + .await + }) + } +} + +/// An atomic heterogeneous batch insertion, returned by +/// [`Client::insert_batch`]. Await it to insert the jobs and get one result +/// per job, in input order. +#[must_use = "insert requests do nothing unless awaited"] +pub struct InsertBatchRequest<'a> { + batch: InsertBatch, + client: &'a Client, + target: Target<'a>, +} + +impl<'a> InsertBatchRequest<'a> { + /// Inserts the batch in a caller-managed transaction. + pub fn tx<'t, E>(self, executor: E) -> InsertBatchRequest<'t> + where + 'a: 't, + E: DatabaseTransactionExecutor<'t>, + { + InsertBatchRequest { + batch: self.batch, + client: self.client, + target: Target::transaction(self.client, executor), + } + } +} + +impl fmt::Debug for InsertBatchRequest<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("InsertBatchRequest") + .field("jobs", &self.batch.len()) + .finish_non_exhaustive() + } +} + +impl<'a> IntoFuture for InsertBatchRequest<'a> { + type Output = Result, Error>; + type IntoFuture = Pin + Send + 'a>>; + + fn into_future(self) -> Self::IntoFuture { + Box::pin(async move { + let Self { + batch, + client, + target, + } = self; + let now = Utc::now(); + let mut jobs = Vec::with_capacity(batch.len()); + for item in batch.items { + client.validate_known_kind(item.kind)?; + let opts = InsertOpts::resolve( + client.inner.default_max_attempts, + item.defaults, + item.opts, + ); + jobs.push(client.prepare_encoded( + item.kind, + item.unique_fields, + item.encoded_args?, + opts, + now, + )?); + } + client + .run_insert(target.into_executor()?, jobs, |rows| { + Ok(rows + .into_iter() + .map(|row| InsertBatchResult { + job: row.job, + unique_skipped_as_duplicate: row.unique_skipped_as_duplicate, + }) + .collect()) + }) + .await + }) + } +} + +impl Client { + /// Inserts a job. + /// + /// Await the returned request to insert the job with its job type's + /// defaults, or chain [`opts`](InsertRequest::opts) to override options + /// and [`tx`](InsertRequest::tx) to insert in a caller-managed + /// transaction: + /// + /// ```no_run + /// # use riverqueue::{Client, InsertOpts, JobArgs}; + /// # use serde::{Deserialize, Serialize}; + /// #[derive(Deserialize, JobArgs, Serialize)] + /// #[river(kind = "send_email")] + /// struct SendEmail { + /// address: String, + /// } + /// + /// # #[cfg(feature = "postgres")] + /// # async fn example(client: Client, pool: sqlx::PgPool) -> Result<(), Box> { + /// let inserted = client + /// .insert(SendEmail { address: "user@example.com".to_owned() }) + /// .await?; + /// println!("inserted job {}", inserted.id()); + /// + /// let mut tx = riverqueue::database::begin_postgres(&pool).await?; + /// client + /// .insert(SendEmail { address: "admin@example.com".to_owned() }) + /// .opts(InsertOpts::default().with_queue("email")) + /// .tx(&mut tx) + /// .await?; + /// tx.commit().await?; + /// # Ok(()) + /// # } + /// ``` + /// + /// A unique job whose insertion matches an existing job returns that job + /// with [`InsertResult::unique_skipped_as_duplicate`] set. + /// + /// # Errors + /// + /// Returns [`Error::InvalidJob`] for invalid options, + /// [`Error::Json`] when the arguments can't be encoded, such as a + /// non-finite float, [`Error::UnknownJobKind`] when the client has workers + /// but none for this kind, [`Error::DatabaseMismatch`] for a transaction + /// from another backend, [`Error::Extension`] when a hook or middleware + /// fails, and [`Error::Database`] when the database operation fails. + /// + /// # Cancel safety + /// + /// Dropping the future before it completes rolls back River's own + /// transaction. With [`tx`](InsertRequest::tx), the caller's transaction + /// may contain a partial insertion and should be rolled back. + pub fn insert(&self, args: A) -> InsertRequest<'_, A> { + InsertRequest { + args, + client: self, + opts: InsertOpts::default(), + target: Target::Client, + } + } + + /// Atomically inserts a batch of one or more jobs of one type. + /// + /// Items are job arguments or `(args, opts)` tuples. Await the request to + /// get one [`InsertResult`] per job in input order. + /// + /// ```no_run + /// # use riverqueue::{Client, InsertOpts, JobArgs}; + /// # use serde::{Deserialize, Serialize}; + /// # #[derive(Deserialize, JobArgs, Serialize)] + /// # #[river(kind = "send_email")] + /// # struct SendEmail { address: String } + /// # async fn example(client: Client) -> Result<(), riverqueue::Error> { + /// let results = client + /// .insert_many([ + /// (SendEmail { address: "a@example.com".to_owned() }, InsertOpts::default()), + /// (SendEmail { address: "b@example.com".to_owned() }, InsertOpts::default().with_priority(2)), + /// ]) + /// .await?; + /// assert_eq!(results.len(), 2); + /// # Ok(()) + /// # } + /// ``` + /// + /// # Errors + /// + /// Returns the same errors as [`insert`](Self::insert), and + /// [`Error::InvalidJob`] for an empty batch. Any error rolls back the + /// whole batch. + /// + /// # Cancel safety + /// + /// Same as [`insert`](Self::insert). + pub fn insert_many(&self, jobs: I) -> InsertManyRequest<'_, A> + where + A: JobArgs, + I: IntoIterator, + I::Item: Into>, + { + InsertManyRequest { + client: self, + jobs: jobs.into_iter().map(Into::into).collect(), + target: Target::Client, + } + } + + /// Atomically inserts a batch that can mix job types. + /// + /// Each item keeps its own job type's defaults and uniqueness. Results + /// correspond positionally to the batch items. + /// + /// # Errors + /// + /// Returns the same errors as [`insert_many`](Self::insert_many), plus + /// [`Error::Json`] if an item's arguments failed to serialize when it was + /// added to the batch. + /// + /// # Cancel safety + /// + /// Same as [`insert`](Self::insert). + pub fn insert_batch(&self, batch: InsertBatch) -> InsertBatchRequest<'_> { + InsertBatchRequest { + batch, + client: self, + target: Target::Client, + } + } + + /// Resolves typed insertion options for an extension. + #[must_use] + pub(crate) fn resolve_insert_opts(&self, opts: InsertOpts) -> InsertParams { + InsertOpts::resolve( + self.inner.default_max_attempts, + A::default_insert_opts(), + opts, + ) + } + + pub(crate) async fn insert_periodic( + &self, + insert: PeriodicInsert, + opts: InsertParams, + target: DateTime, + ) -> Result { + let job = self.prepare_periodic( + insert.kind, + insert.unique_fields, + insert.encoded_args, + opts, + target, + Utc::now(), + )?; + self.run_insert(None, vec![job], |rows| { + rows.into_iter().next().map(|row| row.job).ok_or_else(|| { + Error::runtime_context("periodic job insertion", "insertion returned no row") + }) + }) + .await + } + + /// Resolves, validates, and computes the uniqueness of a typed job. + fn prepare_typed( + &self, + args: &A, + opts: InsertOpts, + now: DateTime, + ) -> Result { + self.validate_known_kind(A::KIND)?; + let encoded_args = crate::encoding::encode_args(args)?; + let opts = InsertOpts::resolve( + self.inner.default_max_attempts, + A::default_insert_opts(), + opts, + ); + self.prepare_encoded(A::KIND, A::unique_fields(), encoded_args, opts, now) + } + + /// Prepares a periodic job due at `target`, as River Go's periodic job + /// enqueuer does. + /// + /// When the constructor leaves the schedule unset, the job runs at its + /// target time: it is inserted `available` with `scheduled_at` set to the + /// target, and a `by_period` unique key describes the target's period. + /// An explicit schedule from the constructor keeps the ordinary + /// `scheduled` state, and a pending job stays pending. + pub(super) fn prepare_periodic( + &self, + kind: &str, + unique_fields: &[&[&str]], + encoded_args: Box, + mut opts: InsertParams, + target: DateTime, + now: DateTime, + ) -> Result { + let due_at_target = opts.scheduled_at.is_none(); + opts.scheduled_at.get_or_insert(target); + let mut job = self.prepare_encoded(kind, unique_fields, encoded_args, opts, now)?; + if due_at_target && job.state == JobState::Scheduled { + job.state = JobState::Available; + } + Ok(job) + } + + /// Validates an encoded job and computes its unique key and initial + /// state, as River Go does before insertion middleware and hooks run. + pub(super) fn prepare_encoded( + &self, + kind: &str, + unique_fields: &[&[&str]], + encoded_args: Box, + opts: InsertParams, + now: DateTime, + ) -> Result { + validate_insert_parts(kind, &opts, self.inner.allow_legacy_job_kinds)?; + let unique_key = build_unique_key_parts( + kind, + unique_fields, + &encoded_args, + now, + &opts.unique, + &opts.queue, + opts.scheduled_at, + )? + .map(|key| key.to_vec()); + let unique_states = unique_key.as_ref().map(|_| opts.unique.state_bitmask()); + let state = if opts.pending { + JobState::Pending + } else if opts.scheduled_at.is_some() { + JobState::Scheduled + } else { + JobState::Available + }; + Ok(InsertContext { + encoded_args, + kind: kind.to_owned(), + opts, + state, + created_at: None, + unique_key, + unique_states, + }) + } + + /// Runs the insertion pipeline, decodes the returned rows with decode + /// hooks, and converts them with `finish`. + /// + /// Without a caller transaction, validation has already run, and + /// middleware, begin hooks, extension interception, the write, the insert + /// notification, decode hooks, and `finish` all run in one transaction, + /// like River Go's `Insert` and `InsertMany`: an error anywhere, including + /// after the write, rolls the whole insertion back. In a caller's + /// transaction the same steps run directly in it without a savepoint, + /// like River Go's `InsertTx`, and the caller rolls back on an error. + pub(super) async fn run_insert( + &self, + executor: Option>, + jobs: Vec, + finish: impl FnOnce(Vec) -> Result + Send, + ) -> Result { + if jobs.is_empty() { + return Err(Error::invalid_job("no jobs to insert".to_owned())); + } + let Some(connection) = executor else { + let (inserted, signals) = match self.inner.database.pool() { + #[cfg(feature = "postgres")] + DatabasePool::Postgres(pool) => { + let mut transaction = crate::database::begin_postgres(pool).await?; + let inserted = self + .insert_and_finish( + PilotDatabaseConnection::Postgres(&mut transaction), + jobs, + finish, + ) + .await?; + transaction.commit().await?; + inserted + } + #[cfg(feature = "sqlite")] + DatabasePool::Sqlite(pool) => { + let mut transaction = crate::database::begin_sqlite_write(pool).await?; + let inserted = self + .insert_and_finish( + PilotDatabaseConnection::Sqlite(&mut transaction), + jobs, + finish, + ) + .await?; + transaction.commit().await?; + inserted + } + }; + // Wake local producers for jobs this client committed itself. + for queue in signals { + let _ = self + .inner + .queue_notifications + .send(RuntimeNotification::Insert(queue)); + } + return Ok(inserted); + }; + self.insert_and_finish(connection, jobs, finish) + .await + .map(|(inserted, _)| inserted) + } + + /// Runs middleware and persistence on `connection`, then decode hooks + /// and `finish`, returning `finish`'s result and the queues of newly + /// available jobs. + async fn insert_and_finish( + &self, + connection: PilotDatabaseConnection<'_>, + jobs: Vec, + finish: impl FnOnce(Vec) -> Result + Send, + ) -> Result<(T, Vec), Error> { + let mut rows = self.insert_on_connection(connection, jobs).await?; + for row in &mut rows { + for hook in self.inner.hooks.iter().rev() { + hook.decode_insert_result(&mut row.job).await?; + } + } + let signals = rows + .iter() + .filter(|row| row.job.state == JobState::Available && !row.unique_skipped_as_duplicate) + .map(|row| row.job.queue.clone()) + .collect(); + Ok((finish(rows)?, signals)) + } + + /// Runs insertion middleware around persistence of `jobs`. + async fn insert_on_connection<'c>( + &'c self, + connection: PilotDatabaseConnection<'c>, + jobs: Vec, + ) -> Result, Error> { + let endpoint: InsertEndpoint<'c> = + Box::new(move |jobs| Box::pin(self.persist_jobs(connection, jobs))); + InsertNext::new(&self.inner.insert_middleware, endpoint) + .run(jobs) + .await + } + + /// Runs begin hooks and extension interception for each job, writes the + /// jobs, and notifies the queues of jobs inserted as available. + async fn persist_jobs( + &self, + mut connection: PilotDatabaseConnection<'_>, + mut jobs: Vec, + ) -> Result, Error> { + if jobs.is_empty() { + return Err(Error::invalid_job("no jobs to insert".to_owned())); + } + let intercepts = self.inner.pilot.intercepts_insert(); + for job in &mut jobs { + for hook in &self.inner.hooks { + hook.insert_begin(job).await?; + } + } + if intercepts { + let mut params = extension_insert_params(&mut jobs); + self.inner + .pilot + .before_jobs_insert(connection.reborrow(), &mut params) + .await + .map_err(|source| Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "job insertion", + }, + source, + })?; + } + for job in &jobs { + if !matches!( + job.state, + JobState::Available | JobState::Pending | JobState::Scheduled + ) { + return Err(Error::invalid_job(format!( + "jobs can't be inserted in the {} state", + job.state.as_str() + ))); + } + } + + // A unique key may appear only once among the batch's jobs whose + // state it covers. PostgreSQL rejects River Go's single upsert that + // would affect the same row twice, and River Go checks SQLite + // batches the same way. + let mut unique_keys = std::collections::HashSet::new(); + for job in &jobs { + if let (Some(key), Some(states)) = (&job.unique_key, job.unique_states) + && states & job.state.unique_bit() != 0 + && !unique_keys.insert(key.as_slice()) + { + return Err(Error::invalid_job( + "unique key appears more than once in batch".to_owned(), + )); + } + } + + // Like Go, every job requested as available notifies its queue, + // including one skipped as a unique duplicate. + let queues = jobs + .iter() + .filter(|job| job.state == JobState::Available) + .map(|job| job.opts.queue.clone()) + .collect::>(); + let mut rows = Vec::with_capacity(jobs.len()); + for job in jobs { + rows.push(self.insert_row(connection.reborrow(), job).await?); + } + if intercepts { + self.after_jobs_inserted(connection.reborrow(), &rows) + .await?; + } + self.notify_insert( + connection.reborrow(), + queues.iter().map(String::as_str).collect(), + ) + .await?; + Ok(rows) + } + + /// Runs the extension's post-insert hook on the rows an insertion wrote. + async fn after_jobs_inserted( + &self, + mut connection: PilotDatabaseConnection<'_>, + rows: &[InsertedJob], + ) -> Result<(), Error> { + let inserted = rows + .iter() + .filter(|row| !row.unique_skipped_as_duplicate) + .map(|row| row.job.clone()) + .collect::>(); + if inserted.is_empty() { + return Ok(()); + } + self.inner + .pilot + .after_jobs_inserted( + connection.reborrow(), + &crate::__private::JobsInsertedParams { + database: self.inner.pilot_database_config(), + jobs: &inserted, + }, + ) + .await + .map_err(|source| Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "job insertion", + }, + source, + }) + } + + /// Sends one insert notification per queue, in the insertion's + /// transaction so it's delivered only if the jobs commit, skipping + /// queues this client notified within its fetch cooldown. + async fn notify_insert( + &self, + connection: PilotDatabaseConnection<'_>, + queues: std::collections::BTreeSet<&str>, + ) -> Result<(), Error> { + let queues = self.inner.insert_notify_limiter.due(queues); + if queues.is_empty() { + return Ok(()); + } + match connection { + #[cfg(feature = "postgres")] + PilotDatabaseConnection::Postgres(connection) => { + if !self + .inner + .postgres_capabilities(&mut *connection) + .await? + .supports_listen_notify + { + return Ok(()); + } + let payloads = queues + .into_iter() + .map(crate::protocol::insert_notification_payload) + .collect::>(); + sqlx::query( + "SELECT pg_notify(concat(coalesce($1::text, current_schema()), '.', $2::text), payload) \ + FROM unnest($3::text[]) AS payload", + ) + .bind(self.inner.schema.as_deref()) + .bind(crate::protocol::NOTIFICATION_TOPIC_INSERT) + .bind(payloads) + .execute(connection) + .await?; + } + #[cfg(feature = "sqlite")] + PilotDatabaseConnection::Sqlite(connection) => { + let payloads = queues + .into_iter() + .map(crate::protocol::insert_notification_payload) + .collect::>(); + let notifications = payloads + .iter() + .map(|payload| crate::database::sqlite::NotificationInput { + payload, + topic: crate::protocol::NOTIFICATION_TOPIC_INSERT, + }) + .collect::>(); + crate::database::sqlite::notification_insert(connection, ¬ifications) + .await + .map_err(sqlite_backend_error)?; + } + } + Ok(()) + } + + /// Writes one job, returning it or the existing unique job it matched. + async fn insert_row( + &self, + connection: PilotDatabaseConnection<'_>, + job: InsertContext, + ) -> Result { + let InsertContext { + encoded_args, + kind, + opts, + state, + created_at, + unique_key, + unique_states, + } = job; + match connection { + #[cfg(feature = "postgres")] + PilotDatabaseConnection::Postgres(connection) => { + let table = self.inner.schema.qualify("river_job"); + let state_type = self.inner.schema.qualify("river_job_state"); + let state_function = self.inner.schema.qualify("river_job_state_in_bitmask"); + let mode = self + .inner + .postgres_capabilities(&mut *connection) + .await? + .unique_insert_mode; + let (metadata, nonce) = with_unique_nonce(mode, &opts.metadata)?; + // The no-op update is intentional and matches River Go. `DO + // NOTHING` followed by a select cannot see a conflicting row + // that committed after the statement's snapshot was taken. + // It keeps the existing kind, which may differ under + // `exclude_kind`. + let sql = format!( + "WITH inserted AS (\ + INSERT INTO {table} (args, created_at, kind, max_attempts, metadata, priority, queue, scheduled_at, state, tags, unique_key, unique_states) \ + VALUES ($1, coalesce($2, now()), $3, $4, $5, $6, $7, coalesce($8, now()), $9::text::{state_type}, $10, $11, $12::integer::bit(8)) \ + ON CONFLICT (unique_key) WHERE unique_key IS NOT NULL AND unique_states IS NOT NULL AND {state_function}(unique_states, state) \ + DO UPDATE SET kind = river_job.kind \ + RETURNING *, {} AS unique_skipped_as_duplicate\ + ) \ + SELECT {}, job.unique_skipped_as_duplicate FROM inserted AS job", + mode.sql(), + job_projection("job") + ); + let record = sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) + .bind(Json(&encoded_args)) + .bind(created_at) + .bind(&kind) + .bind(opts.max_attempts) + .bind(Json(&*metadata)) + .bind(opts.priority) + .bind(&opts.queue) + .bind(opts.scheduled_at) + .bind(state.as_str()) + .bind(&opts.tags) + .bind(unique_key) + .bind(unique_states.map(i32::from)) + .fetch_optional(connection) + .await? + .ok_or_else(|| { + Error::invalid_job("unique insert found no conflicting row".to_owned()) + })?; + let duplicate = record.unique_skipped_as_duplicate; + let row = record.into_job_row()?; + let duplicate = nonce.map_or(duplicate, |nonce| lacks_nonce(&row, &nonce)); + Ok(InsertedJob::new(row, duplicate)) + } + #[cfg(feature = "sqlite")] + PilotDatabaseConnection::Sqlite(connection) => { + // Like Go's SQLite driver, every inserted row carries a nonce, + // and times left unset are filled in by SQLite's own clock. + let nonce = unique_insert_nonce(); + let inserted = crate::database::sqlite::insert( + connection, + &crate::database::sqlite::InsertJob { + attempt: 0, + attempted_at: None, + attempted_by: &[], + created_at, + encoded_args: &encoded_args, + errors: &[], + finalized_at: None, + id: None, + kind: &kind, + max_attempts: opts.max_attempts, + metadata: &opts.metadata, + priority: opts.priority, + queue: &opts.queue, + scheduled_at: opts.scheduled_at, + state, + tags: &opts.tags, + unique_key: unique_key.as_deref(), + unique_nonce: Some(&nonce), + unique_states, + }, + ) + .await + .map_err(sqlite_backend_error)?; + Ok(InsertedJob::new( + inserted.job, + inserted.unique_skipped_as_duplicate, + )) + } + } + } +} + +/// Exposes the mutable fields of each job to an extension's insert hook. +fn extension_insert_params(jobs: &mut [InsertContext]) -> Vec> { + jobs.iter_mut() + .map(|job| { + let InsertContext { + encoded_args, + kind, + opts, + state, + .. + } = job; + PilotJobInsertParams { + encoded_args, + extension_options: &opts.extension_options, + kind, + metadata: &mut opts.metadata, + queue: &mut opts.queue, + state, + } + }) + .collect() +} + +/// Returns a nonce that marks an insert as this call's own. +/// +/// A database without `xmax`, like SQLite or YugabyteDB, reports a skipped +/// unique duplicate by checking whether the returned row carries the nonce +/// the insert wrote. The nonce must not +/// repeat across processes: client IDs and counters can (a restarted +/// container keeps its hostname and PID), so it's eight random bytes in +/// lowercase hex, the format of River Go's `randutil.Hex(8)`. +/// +/// Like Go, every row gets its own nonce. +fn unique_insert_nonce() -> String { + format!("{:016x}", rand::random::()) +} + +/// Returns `metadata` with a new unique insert nonce, and the nonce, when +/// `mode` detects duplicates by one. Without `xmax`, as on YugabyteDB, a row +/// carries a nonce like SQLite's, and a returned row without it existed. +#[cfg(feature = "postgres")] +fn with_unique_nonce( + mode: UniqueInsertMode, + metadata: &JobMetadata, +) -> Result<(Cow<'_, JobMetadata>, Option), Error> { + if mode != UniqueInsertMode::MetadataNonce { + return Ok((Cow::Borrowed(metadata), None)); + } + let nonce = unique_insert_nonce(); + let mut metadata = metadata.clone(); + metadata.insert(METADATA_KEY_UNIQUE_NONCE, Value::String(nonce.clone()))?; + Ok((Cow::Owned(metadata), Some(nonce))) +} + +/// Whether a row returned by a unique insert lacks the nonce the insert +/// wrote, so it existed already. +#[cfg(feature = "postgres")] +fn lacks_nonce(row: &JobRow, nonce: &str) -> bool { + row.metadata + .get::(METADATA_KEY_UNIQUE_NONCE) + .ok() + .flatten() + .is_none_or(|stored| stored != nonce) +} diff --git a/rust/riverqueue/src/client/jobs.rs b/rust/riverqueue/src/client/jobs.rs new file mode 100644 index 000000000..a05ec5f7a --- /dev/null +++ b/rust/riverqueue/src/client/jobs.rs @@ -0,0 +1,453 @@ +//! Operations on persisted jobs. + +use std::{ + fmt, + future::{Future, IntoFuture}, + pin::Pin, +}; + +use serde_json::Map; + +use crate::__private::DatabaseConnection as PilotDatabaseConnection; +use crate::client::attempts::signal_running_attempt; +use crate::client::request::{Target, request_type}; +use crate::database::DatabaseTransactionExecutor; +use crate::storage::Access; +use crate::{Client, Error, JobRow}; +use crate::{JobDeleteManyParams, JobListCursor, JobListParams, JobListResult, JobUpdateParams}; + +/// Operations on persisted jobs, returned by [`Client::jobs`]. +/// +/// Each method returns a request that runs on the client's own pool when +/// awaited, or in a caller-managed transaction after `.tx(&mut tx)`: +/// +/// ```no_run +/// # #[cfg(feature = "postgres")] +/// # async fn example(client: riverqueue::Client, pool: sqlx::PgPool) -> Result<(), riverqueue::Error> { +/// let job = client.jobs().get(42).await?; +/// +/// let mut tx = riverqueue::database::begin_postgres(&pool).await?; +/// client.jobs().cancel(job.id).tx(&mut tx).await?; +/// tx.commit().await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// Requests don't run until awaited. Dropping one before it completes rolls +/// back River's own transaction; with `.tx`, the caller's transaction may +/// contain the operation's partial effects and should be rolled back. +#[derive(Clone, Copy, Debug)] +pub struct Jobs<'a> { + client: &'a Client, +} + +impl Client { + /// Returns operations on persisted jobs: getting, listing, cancelling, + /// retrying, updating, and deleting them. + #[must_use] + pub const fn jobs(&self) -> Jobs<'_> { + Jobs { client: self } + } +} + +impl<'a> Jobs<'a> { + /// Cancels a job and returns its current row. + /// + /// A job that is available, scheduled, retryable, or pending is cancelled + /// immediately and won't run again. A running job is marked for + /// cancellation, and the client running it cancels the attempt's + /// [`WorkContext::cancellation_token`](crate::WorkContext::cancellation_token): if the worker then returns an + /// error, the job is cancelled rather than retried, while a job that + /// completes successfully stays completed. A finalized job is returned + /// unchanged. + /// + /// With [`tx`](JobCancelRequest::tx), the cancellation and its + /// notification take effect only when the transaction commits. + /// + /// # Errors + /// + /// Returns [`Error::NotFound`] when the job doesn't exist, + /// [`Error::DatabaseMismatch`] for a transaction from another backend, + /// [`Error::Extension`] when an extension's cancellation hook fails, and + /// [`Error::Database`] when the database operation fails. + pub fn cancel(&self, id: i64) -> JobCancelRequest<'a> { + JobCancelRequest { + client: self.client, + id, + target: Target::Client, + } + } + + /// Completes a running job in a caller-managed transaction, for example + /// alongside business writes that the job performed. + /// + /// The returned request has no effect until it's given the transaction + /// with [`tx`](JobCompleteRequest::tx) and awaited. The job becomes + /// completed only when the transaction commits. If this completes a job + /// that is still being worked, the worker's own result is discarded when + /// it finishes, because the job is no longer running. Workers can use + /// [`WorkContext::job_complete_tx`](crate::WorkContext::job_complete_tx), which also records metadata set on + /// the work context. + /// + /// ```no_run + /// # #[cfg(feature = "postgres")] + /// # async fn example(client: riverqueue::Client, pool: sqlx::PgPool) -> Result<(), riverqueue::Error> { + /// let mut tx = riverqueue::database::begin_postgres(&pool).await?; + /// // ... business writes in `tx` ... + /// client.jobs().complete(42).tx(&mut tx).await?; + /// tx.commit().await?; + /// # Ok(()) + /// # } + /// ``` + pub fn complete(&self, id: i64) -> JobCompleteRequest<'a> { + JobCompleteRequest { + client: self.client, + id, + } + } + + /// Deletes a job that isn't running and returns its former row. + /// + /// # Errors + /// + /// Returns [`Error::NotFound`] when the job doesn't exist, + /// [`Error::JobRunning`] when it's running, [`Error::DatabaseMismatch`] + /// for a transaction from another backend, and [`Error::Database`] when + /// the database operation fails. + pub fn delete(&self, id: i64) -> JobDeleteRequest<'a> { + JobDeleteRequest { + client: self.client, + id, + target: Target::Client, + } + } + + /// Deletes jobs that aren't running and match the parameters, returning + /// the deleted rows in list order. + /// + /// At most the filter's limit of jobs are deleted. Running jobs are + /// skipped before the limit applies, and PostgreSQL also skips jobs + /// locked by another transaction rather than waiting for them. + /// + /// # Errors + /// + /// Returns [`Error::InvalidJob`] for invalid list parameters, + /// [`Error::DatabaseMismatch`] for a transaction from another backend, and + /// [`Error::Database`] when the database operation fails. + pub fn delete_many(&self, params: JobDeleteManyParams) -> JobDeleteManyRequest<'a> { + JobDeleteManyRequest { + client: self.client, + params, + target: Target::Client, + } + } + + /// Gets a job by ID. + /// + /// # Errors + /// + /// Returns [`Error::NotFound`] when the job doesn't exist, + /// [`Error::DatabaseMismatch`] for a transaction from another backend, and + /// [`Error::Database`] when the database operation fails. + pub fn get(&self, id: i64) -> JobGetRequest<'a> { + JobGetRequest { + client: self.client, + id, + target: Target::Client, + } + } + + /// Lists jobs matching the parameters, one page at a time. + /// + /// Pass the result's [`last_cursor`](JobListResult::last_cursor) to + /// [`JobListParams::after`] with otherwise identical parameters to request + /// the next page: + /// + /// ```no_run + /// # use riverqueue::{JobListParams, JobState}; + /// # async fn example(client: riverqueue::Client) -> Result<(), riverqueue::Error> { + /// let params = JobListParams::default() + /// .states([JobState::Completed]) + /// .limit(100); + /// let mut page = client.jobs().list(params.clone()).await?; + /// while let Some(cursor) = page.last_cursor.take() { + /// for job in &page.jobs { + /// println!("{} {:?}", job.id, job.state); + /// } + /// page = client.jobs().list(params.clone().after(cursor)).await?; + /// } + /// # Ok(()) + /// # } + /// ``` + /// + /// # Errors + /// + /// Returns [`Error::InvalidJob`] for invalid parameters, such as a limit + /// outside one through 10,000 or a cursor from a different ordering, + /// [`Error::DatabaseMismatch`] for a transaction from another backend, and + /// [`Error::Database`] when the database operation fails. + pub fn list(&self, params: JobListParams) -> JobListRequest<'a> { + JobListRequest { + client: self.client, + params, + target: Target::Client, + } + } + + /// Makes a job that isn't running available to be worked again and + /// returns its current row. + /// + /// The job's `scheduled_at` moves to now unless it's already available + /// and due, so a waiting job doesn't lose its place, and a job that has + /// used all of its attempts gets one more. A running job is returned + /// unchanged. + /// + /// # Errors + /// + /// Returns [`Error::NotFound`] when the job doesn't exist, + /// [`Error::DatabaseMismatch`] for a transaction from another backend, + /// [`Error::Extension`] when an extension's retry hook fails, and + /// [`Error::Database`] when the database operation fails. + pub fn retry(&self, id: i64) -> JobRetryRequest<'a> { + JobRetryRequest { + client: self.client, + id, + target: Target::Client, + } + } + + /// Merges metadata into a job, optionally setting its recorded output, + /// and returns the updated row. + /// + /// # Errors + /// + /// Returns [`Error::NotFound`] when the job doesn't exist, + /// [`Error::InvalidJob`] when the output's JSON is larger than 32 MB, + /// [`Error::DatabaseMismatch`] for a transaction from + /// another backend, and [`Error::Database`] when the database operation + /// fails. + pub fn update(&self, id: i64, params: JobUpdateParams) -> JobUpdateRequest<'a> { + JobUpdateRequest { + client: self.client, + id, + params, + target: Target::Client, + } + } +} + +request_type! { + /// A job cancellation, returned by [`Jobs::cancel`]. Await it to cancel + /// the job and get its current row. + /// + /// A client without a notification listener, such as a poll-only + /// client, signals its own running attempt of the job directly once the + /// cancellation commits. If the request is dropped while that commit is + /// in flight, the cancellation may commit without the signal; the + /// attempt then keeps running until it ends, and a failure it returns is + /// recorded as a cancellation instead of retried. + write JobCancelRequest { id: i64 } -> JobRow +} + +impl JobCancelRequest<'_> { + async fn run(self) -> Result { + let inner = &self.client.inner; + let own_transaction = !self.target.is_transaction(); + let mut session = self.target.session(inner, Access::Transaction).await?; + let row = session.storage(inner).job_cancel(self.id).await?; + session.commit().await?; + // Without a listener (no backend listener, or a poll-only client), + // wake this client's running attempt directly, like Go's + // `notifyProducerWithoutListenerQueueControlEvent`. Other clients + // observe the committed notification through their own listener or + // outbox poll. + if own_transaction && !inner.listens_for_notifications() { + signal_running_attempt( + &inner.running, + &inner.pending_cancellations, + &inner.fetch_registration_windows, + self.id, + ); + } + Ok(row) + } +} + +request_type! { + /// A job deletion, returned by [`Jobs::delete`]. Await it to delete the + /// job and get its former row. + write JobDeleteRequest { id: i64 } -> JobRow +} + +impl JobDeleteRequest<'_> { + async fn run(self) -> Result { + let inner = &self.client.inner; + let mut session = self.target.session(inner, Access::Transaction).await?; + let row = session.storage(inner).job_delete(self.id).await?; + session.commit().await?; + Ok(row) + } +} + +request_type! { + /// A bulk job deletion, returned by [`Jobs::delete_many`]. Await it to + /// delete the jobs and get their former rows. + write JobDeleteManyRequest { params: JobDeleteManyParams } -> Vec +} + +impl JobDeleteManyRequest<'_> { + async fn run(self) -> Result, Error> { + let inner = &self.client.inner; + let mut session = self.target.session(inner, Access::Transaction).await?; + let rows = session.storage(inner).job_delete_many(&self.params).await?; + session.commit().await?; + Ok(rows) + } +} + +request_type! { + /// A job lookup, returned by [`Jobs::get`]. Await it to get the job. + read JobGetRequest { id: i64 } -> JobRow +} + +impl JobGetRequest<'_> { + async fn run(self) -> Result { + let inner = &self.client.inner; + let mut session = self.target.session(inner, Access::Autocommit).await?; + session.storage(inner).job_get(self.id).await + } +} + +request_type! { + /// A job listing, returned by [`Jobs::list`]. Await it to get a page of + /// jobs. + read JobListRequest { params: JobListParams } -> JobListResult +} + +impl JobListRequest<'_> { + async fn run(self) -> Result { + let inner = &self.client.inner; + let mut session = self.target.session(inner, Access::Autocommit).await?; + let jobs = session.storage(inner).job_list(&self.params).await?; + let last_cursor = jobs + .last() + .map(|job| JobListCursor::after_job(job, &self.params)); + Ok(JobListResult { jobs, last_cursor }) + } +} + +request_type! { + /// A job retry, returned by [`Jobs::retry`]. Await it to make the job + /// available and get its current row. + write JobRetryRequest { id: i64 } -> JobRow +} + +impl JobRetryRequest<'_> { + async fn run(self) -> Result { + let inner = &self.client.inner; + let mut session = self.target.session(inner, Access::Transaction).await?; + let row = session.storage(inner).job_retry(self.id).await?; + session.commit().await?; + Ok(row) + } +} + +request_type! { + /// A job update, returned by [`Jobs::update`]. Await it to update the job + /// and get its new row. + write JobUpdateRequest { id: i64, params: JobUpdateParams } -> JobRow +} + +impl JobUpdateRequest<'_> { + async fn run(self) -> Result { + let inner = &self.client.inner; + let mut session = self.target.session(inner, Access::Autocommit).await?; + session + .storage(inner) + .job_update(self.id, self.params) + .await + } +} + +/// A transactional job completion, returned by [`Jobs::complete`]. +/// +/// It has no effect on its own: pass the transaction to complete the job in +/// with [`tx`](Self::tx) and await the result. +#[must_use = "a completion needs `.tx(&mut tx)` and must be awaited"] +#[derive(Debug)] +pub struct JobCompleteRequest<'a> { + client: &'a Client, + id: i64, +} + +impl<'a> JobCompleteRequest<'a> { + /// Completes the job in a caller-managed transaction. The job becomes + /// completed only when the transaction commits. + /// + /// `executor` must be a SQLx transaction for the client's database + /// backend, begun with + /// [`begin_postgres`](crate::database::begin_postgres) or + /// [`begin_sqlite_write`](crate::database::begin_sqlite_write). + pub fn tx<'t, E>(self, executor: E) -> JobCompleteTxRequest<'t> + where + 'a: 't, + E: DatabaseTransactionExecutor<'t>, + { + JobCompleteTxRequest { + client: self.client, + connection: self.client.inner.transaction_connection(executor), + id: self.id, + } + } +} + +/// A job completion in a caller-managed transaction, returned by +/// [`JobCompleteRequest::tx`]. Await it to complete the job and get its new +/// row. +/// +/// # Errors +/// +/// Awaiting it returns [`Error::NotFound`] when the job doesn't exist, +/// [`Error::JobNotRunning`] when the job isn't running, +/// [`Error::DatabaseMismatch`] for a transaction from another backend, +/// [`Error::Extension`] when an extension's completion hook fails, and +/// [`Error::Database`] when the database operation fails. +/// +/// # Cancel safety +/// +/// Dropping the request before it finishes leaves what it already ran in +/// the caller's transaction; roll that transaction back rather than +/// committing it. +#[must_use = "requests do nothing unless awaited"] +pub struct JobCompleteTxRequest<'a> { + client: &'a Client, + connection: Result, Error>, + id: i64, +} + +impl fmt::Debug for JobCompleteTxRequest<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("JobCompleteTxRequest") + .field("id", &self.id) + .finish_non_exhaustive() + } +} + +impl<'a> IntoFuture for JobCompleteTxRequest<'a> { + type Output = Result; + type IntoFuture = Pin + Send + 'a>>; + + fn into_future(self) -> Self::IntoFuture { + Box::pin(async move { + let connection = self.connection?; + crate::storage::complete_in_caller_transaction( + &self.client.inner, + connection, + self.id, + &Map::new(), + ) + .await + }) + } +} diff --git a/rust/riverqueue/src/client/local_queues.rs b/rust/riverqueue/src/client/local_queues.rs new file mode 100644 index 000000000..90f2dde10 --- /dev/null +++ b/rust/riverqueue/src/client/local_queues.rs @@ -0,0 +1,178 @@ +//! This client's own queue configuration. + +use std::collections::HashMap; +use std::sync::PoisonError; + +use crate::client::builder::validate_queue_settings; +use crate::{Client, Error, QueueConfig}; + +/// The queues this client works, returned by [`Client::local_queues`]. +/// +/// This is the client's runtime configuration, not the shared queue records +/// managed through [`Client::queues`]: adding or removing a queue here +/// changes only which queues this client's producers fetch from. Adding a +/// queue that's already added is an error, and removing one waits for its +/// producer to stop. +/// +/// Changes apply to a running client asynchronously: +/// +/// - An added queue starts fetching jobs shortly after [`add`](Self::add) +/// returns. +/// - An updated queue applies its new configuration while it runs. Lowering +/// `max_workers` stops new fetches until enough running jobs finish; it +/// never cancels them. +/// - A removed queue stops fetching, and [`remove`](Self::remove) waits for +/// the jobs it already fetched to finish. Its persisted jobs and queue +/// record are left for other clients. Its name stays reserved until then, +/// so the queue never runs under two producers at once. +/// +/// ```no_run +/// # use riverqueue::QueueConfig; +/// # async fn example(client: &riverqueue::Client) -> Result<(), riverqueue::Error> { +/// client.local_queues().add("reports", QueueConfig::new(2))?; +/// client.local_queues().update("reports", QueueConfig::new(4))?; +/// assert!(client.local_queues().configs().contains_key("reports")); +/// let removed = client.local_queues().remove("reports").await?; +/// assert_eq!(removed, QueueConfig::new(4)); +/// # Ok(()) +/// # } +/// ``` +#[derive(Clone, Copy, Debug)] +pub struct LocalQueues<'a> { + client: &'a Client, +} + +impl Client { + /// Returns the configuration of the queues this client works, which can + /// change while it runs. + #[must_use] + pub const fn local_queues(&self) -> LocalQueues<'_> { + LocalQueues { client: self } + } +} + +impl LocalQueues<'_> { + /// Adds a queue for this client to work. + /// + /// A running client starts only this queue's producer; other queues keep + /// running. See [`LocalQueues`] for when the change takes effect. + /// + /// # Errors + /// + /// Returns [`Error::QueueAlreadyAdded`] when the queue is already added + /// or a removal of it is still waiting for its producer to stop, + /// [`Error::InvalidJob`] for an invalid queue name, + /// [`Error::Configuration`] for an invalid configuration or when the + /// client has no workers to run the queue's jobs, and + /// [`Error::Extension`] when an add-on crate rejects the configuration's + /// extension settings. + pub fn add(&self, name: impl Into, config: QueueConfig) -> Result<(), Error> { + let name = name.into(); + let inner = &self.client.inner; + self.validate(&name, &config)?; + if inner.workers.kinds().is_empty() { + return Err(Error::configuration( + "workers must be configured when queues are configured".to_owned(), + )); + } + { + let mut queues = inner.queues.write().unwrap_or_else(PoisonError::into_inner); + if queues.contains_key(&name) || inner.live_queues.borrow().contains(&name) { + return Err(Error::QueueAlreadyAdded { name }); + } + queues.insert(name, config); + } + self.changed(); + Ok(()) + } + + /// Returns a snapshot of the queues this client works and their + /// configurations. + #[must_use] + pub fn configs(&self) -> HashMap { + self.client + .inner + .queues + .read() + .unwrap_or_else(PoisonError::into_inner) + .clone() + } + + /// Stops working a queue, waits until its producer has stopped, and + /// returns the queue's configuration. + /// + /// The producer stops fetching at once and then waits for the jobs it + /// fetched to finish. A client that isn't running returns at once. + /// + /// # Cancel safety + /// + /// This method is cancel safe. The queue is removed when the future is + /// first polled; dropping the future afterwards stops only the wait, + /// and the queue's name stays reserved until its producer stops. + /// + /// # Errors + /// + /// Returns [`Error::QueueNotAdded`] when this client doesn't work the + /// queue. + pub async fn remove(&self, name: &str) -> Result { + let inner = &self.client.inner; + let mut live = inner.live_queues.subscribe(); + let config = inner + .queues + .write() + .unwrap_or_else(PoisonError::into_inner) + .remove(name) + .ok_or_else(|| Error::QueueNotAdded { + name: name.to_owned(), + })?; + self.changed(); + // The sender lives as long as the client, so this ends only when the + // producer is gone. + let _ = live.wait_for(|live| !live.contains(name)).await; + Ok(config) + } + + /// Replaces the configuration of a queue this client works. + /// + /// A running producer applies the new configuration without stopping. + /// See [`LocalQueues`] for when the change takes effect. + /// + /// # Errors + /// + /// Returns [`Error::QueueNotAdded`] when this client doesn't work the + /// queue, [`Error::Configuration`] for an invalid configuration, and + /// [`Error::Extension`] when an add-on crate rejects the configuration's + /// extension settings. + pub fn update(&self, name: &str, config: QueueConfig) -> Result<(), Error> { + self.validate(name, &config)?; + { + let mut queues = self + .client + .inner + .queues + .write() + .unwrap_or_else(PoisonError::into_inner); + let Some(current) = queues.get_mut(name) else { + return Err(Error::QueueNotAdded { + name: name.to_owned(), + }); + }; + *current = config; + } + self.changed(); + Ok(()) + } + + /// Tells a running client's queue supervisor to reconcile its producers. + fn changed(self) { + self.client + .inner + .queue_changes + .send_modify(|generation| *generation = generation.wrapping_add(1)); + } + + fn validate(self, name: &str, config: &QueueConfig) -> Result<(), Error> { + config.validate(name, self.client.inner.fetch_cooldown)?; + validate_queue_settings(self.client.inner.pilot.as_ref(), name, config) + } +} diff --git a/rust/riverqueue/src/client/mod.rs b/rust/riverqueue/src/client/mod.rs new file mode 100644 index 000000000..eebe80877 --- /dev/null +++ b/rust/riverqueue/src/client/mod.rs @@ -0,0 +1,540 @@ +//! Database-backed client, insertion, and worker runtime. + +mod attempts; +mod backoff; +mod builder; +mod completer; +mod executor; +mod extension; +mod insert; +mod jobs; +mod local_queues; +mod notifier; +mod notify_limiter; +mod peers; +mod producer; +mod queues; +mod record; +mod request; +mod resign; +mod run; +#[cfg(test)] +mod tests; +mod validate; + +pub(crate) use self::backoff::SERVICE_RESTART_RESET_AFTER; +pub use self::builder::{ClientBuilder, MaintenanceConfig, QueueConfig, Retention}; +pub use self::extension::{ExtensionClient, PreparedInsertRequest, RawInsertRequest}; +pub use self::insert::{InsertBatchRequest, InsertManyItem, InsertManyRequest, InsertRequest}; +pub use self::jobs::{ + JobCancelRequest, JobCompleteRequest, JobCompleteTxRequest, JobDeleteManyRequest, + JobDeleteRequest, JobGetRequest, JobListRequest, JobRetryRequest, JobUpdateRequest, Jobs, +}; +pub use self::local_queues::LocalQueues; +pub(crate) use self::peers::PeerLedger; +pub use self::queues::{ + QueueGetRequest, QueueListRequest, QueuePauseRequest, QueueResumeRequest, QueueUpdateRequest, + Queues, +}; +#[cfg(feature = "sqlite")] +pub(crate) use self::record::FieldErrors; +pub(crate) use self::record::{DecodedJob, UndecodableJob, saturating_i16, tolerant_row}; +#[cfg(feature = "postgres")] +pub(crate) use self::record::{JobRecord, decode_job_row, job_projection}; +pub use self::resign::ResignRequest; +pub use self::run::{RunHandle, Stopper}; +pub(crate) use self::{ + completer::after_jobs_set_state, executor::default_retry_delay, notifier::RuntimeNotification, + producer::standard_claim, validate::validate_queue, +}; +use std::{ + collections::HashMap, + sync::{ + Arc, Mutex, RwLock, Weak, + atomic::{AtomicBool, AtomicU64, Ordering}, + }, + time::Duration, +}; + +use chrono::{DateTime, Utc}; +use tokio::sync::{broadcast, mpsc, watch}; +use tokio_util::sync::CancellationToken; + +use self::completer::CompletionUpdate; +pub(crate) use self::notify_limiter::InsertNotifyLimiter; +use crate::__private::{ + DatabaseConfig as PilotDatabaseConfig, DatabaseConnection as PilotDatabaseConnection, + DatabasePool as PilotDatabasePool, NoopPilot, Pilot, +}; +#[cfg(feature = "postgres")] +use crate::SchemaName; +use crate::maintenance::LeadershipWakeup; +use crate::{ + DefaultRetryPolicy, Error, Event, EventKind, EventReceiver, FETCH_COOLDOWN_DEFAULT, + JOB_STUCK_THRESHOLD_DEFAULT, JOB_TIMEOUT_DEFAULT, MAX_ATTEMPTS_DEFAULT, RetryPolicy, + SubscribeConfig, WorkerRegistry, + database::{ClientDatabase, Database, DatabasePool, DatabaseTransactionExecutor, IntoDatabase}, + periodic::PeriodicJobs, +}; + +pub(crate) const ATTEMPTED_BY_MAX: i32 = 100; +const EVENT_BUFFER_CAPACITY: usize = 10_000; +const PENDING_CANCELLATION_LIMIT: usize = 10_000; +const PENDING_CANCELLATION_RETENTION: Duration = Duration::from_mins(1); +// Large queues otherwise become limited by a single PostgreSQL claim round trip. +// Concurrent `SKIP LOCKED` claims safely divide the available worker slots. +const PARALLEL_FETCH_MINIMUM: usize = 1_000; +const QUEUE_CONFIG_POLL_INTERVAL: Duration = Duration::from_secs(2); +const QUEUE_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(30); +/// How often producers report to an extension's session, like Go's +/// `ProducerReportInterval` default. +const PRODUCER_REPORT_INTERVAL_DEFAULT: Duration = Duration::from_secs(30); +/// How long a producer's report may run, like Go's +/// `reportProducerStatusOnce` timeout. +const PRODUCER_REPORT_TIMEOUT: Duration = Duration::from_secs(10); +/// Peers that haven't reported for this long are stale, like Go's +/// `StaleProducerRetentionPeriod`. +const PRODUCER_STALE_RETENTION: Duration = Duration::from_mins(5); + +#[allow( + clippy::struct_excessive_bools, + reason = "each flag is an independent configuration option, not a state" +)] +pub(crate) struct ClientInner { + allow_legacy_job_kinds: bool, + allow_unregistered_job_kinds: bool, + completion_sender: Mutex>>, + pub(crate) database: Database, + default_max_attempts: i16, + error_handler: Option>, + pub(crate) events: broadcast::Sender, + /// Minimum delay between fetches for queues without their own, and the + /// window in which repeated insert notifications for a queue are + /// skipped. + pub(crate) fetch_cooldown: Duration, + /// Kinds claims are restricted to, including aliases, or `None` to claim + /// every kind. + pub(crate) fetch_kinds: Option>, + fetch_registration_windows: AtomicU64, + pub(crate) hooks: Vec>, + pub(crate) id: String, + insert_middleware: Vec>, + /// Skips a queue's insert notification sent within the fetch cooldown + /// of the previous one. + pub(crate) insert_notify_limiter: InsertNotifyLimiter, + job_stuck_threshold: Duration, + pub(crate) job_timeout: Option, + leader_election_disabled: bool, + /// Leadership notifications for the elector, kept off the busier + /// producer channel so that insert wakeups can't crowd out a resignation + /// request. + pub(crate) leadership_wakeups: broadcast::Sender, + /// Queues whose producers are running or draining, which keep their + /// names reserved until they stop. + live_queues: watch::Sender>, + pub(crate) maintenance: MaintenanceConfig, + /// Notification listener starts that panic before doing anything, so + /// tests can exercise the supervisor's restart path. + #[cfg(test)] + notifier_start_panics: AtomicU64, + /// Peer jobs owned by running attempts, mapped to their ledger. + peer_owners: Mutex>, + pending_cancellations: Mutex>, + pub(crate) periodic_jobs: PeriodicJobs, + pub(crate) pilot: Arc, + poll_only: bool, + producer_report_interval: Duration, + queue_changes: watch::Sender, + queue_notifications: broadcast::Sender, + queues: RwLock>, + pub(crate) retry_policy: Arc, + running: Mutex>, + #[cfg(feature = "postgres")] + pub(crate) schema: SchemaName, + soft_stop_timeout: Option, + started: AtomicBool, + work_middleware: Vec>, + pub(crate) workers: WorkerRegistry, +} + +#[cfg(feature = "sqlite")] +fn sqlite_backend_error(error: crate::database::sqlite::BackendError) -> Error { + Error::Database(error.into()) +} + +impl ClientInner { + /// Borrows a caller-managed transaction's connection, rejecting a + /// transaction from another backend. + pub(crate) fn transaction_connection<'executor, E>( + &self, + transaction: E, + ) -> Result, Error> + where + E: DatabaseTransactionExecutor<'executor>, + { + Ok(self.database.connection(transaction)?) + } + + #[cfg(feature = "postgres")] + pub(crate) const fn database(&self) -> &Database { + &self.database + } + + /// Returns the PostgreSQL server's capabilities, detecting them with + /// `executor` the first time. + #[cfg(feature = "postgres")] + pub(crate) async fn postgres_capabilities<'e>( + &self, + executor: impl sqlx::PgExecutor<'e>, + ) -> Result { + Ok( + crate::database::postgres_capabilities::CapabilitiesCache::load_or_detect( + self.database.postgres_capabilities(), + executor, + ) + .await?, + ) + } + + /// Whether this client hears committed notifications through a notifier, + /// a PostgreSQL listener or SQLite outbox poller, like River Go's client + /// notifier. A poll-only client has none, and neither does a client of a + /// PostgreSQL server without `LISTEN`/`NOTIFY` once that's detected. + pub(crate) fn has_notifier(&self) -> bool { + !self.poll_only && self.database.delivers_notifications() + } + + /// Whether this client receives notifications from other clients, which + /// needs a backend listener and a client that isn't poll-only. When it + /// doesn't, it wakes its own runtime directly after committing a change, + /// like Go's `notifyProducerWithoutListener*` helpers. + pub(crate) fn listens_for_notifications(&self) -> bool { + self.database.supports_listener() && !self.poll_only + } + + pub(crate) fn pilot_database_config(&self) -> PilotDatabaseConfig { + match self.database.pool() { + #[cfg(feature = "postgres")] + DatabasePool::Postgres(_) => PilotDatabaseConfig::Postgres { + #[cfg(feature = "postgres")] + schema: self.schema.clone(), + }, + #[cfg(feature = "sqlite")] + DatabasePool::Sqlite(_) => PilotDatabaseConfig::Sqlite, + } + } + + /// The client's database as an extension sees it. + pub(crate) fn pilot_database(&self) -> crate::__private::PilotDatabase { + crate::__private::PilotDatabase::new( + self.pilot_database_pool(), + self.pilot_database_config(), + ) + } + + pub(crate) fn pilot_database_pool(&self) -> PilotDatabasePool { + match self.database.pool() { + #[cfg(feature = "postgres")] + DatabasePool::Postgres(pool) => PilotDatabasePool::Postgres(pool.clone()), + #[cfg(feature = "sqlite")] + DatabasePool::Sqlite(pool) => PilotDatabasePool::Sqlite(pool.clone()), + } + } +} + +/// A River client backed by a caller-owned pool for a built-in database. +#[derive(Clone)] +pub struct Client { + pub(crate) inner: Arc, +} + +impl std::fmt::Debug for Client { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("Client") + .field("database_kind", &self.database().kind()) + .field("id", &self.id()) + .field("started", &self.inner.started.load(Ordering::Acquire)) + .finish_non_exhaustive() + } +} + +/// Non-owning handle used by extension services. +#[derive(Clone)] +pub struct WeakClient { + inner: Weak, +} + +impl std::fmt::Debug for WeakClient { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WeakClient") + .field("alive", &(self.inner.strong_count() > 0)) + .finish_non_exhaustive() + } +} + +impl WeakClient { + pub(crate) fn new(inner: &Arc) -> Self { + Self { + inner: Arc::downgrade(inner), + } + } + + /// Upgrades the handle while its originating client remains alive. + #[must_use] + pub fn upgrade(&self) -> Option { + self.inner.upgrade().map(|inner| Client { inner }) + } +} + +impl Client { + /// Creates a client builder for `database`, such as an SQLx pool or a + /// [`PostgresDatabase`](crate::database::PostgresDatabase). A client + /// built without queues only inserts and manages jobs; add queues and + /// workers to work them. + #[must_use] + pub fn builder(database: D) -> ClientBuilder + where + D: IntoDatabase, + { + let database = Database::from_source(database); + ClientBuilder { + allow_legacy_job_kinds: false, + allow_unregistered_job_kinds: false, + database, + default_max_attempts: MAX_ATTEMPTS_DEFAULT, + error_handler: None, + fetch_cooldown: FETCH_COOLDOWN_DEFAULT, + fetch_only_known_kinds: false, + hooks: Vec::new(), + id: default_client_id(), + insert_middleware: Vec::new(), + job_stuck_threshold: JOB_STUCK_THRESHOLD_DEFAULT, + job_timeout: Some(JOB_TIMEOUT_DEFAULT), + leader_election_disabled: false, + maintenance: MaintenanceConfig::default(), + periodic_jobs: Vec::new(), + pilot: Arc::new(NoopPilot), + poll_only: false, + producer_report_interval: PRODUCER_REPORT_INTERVAL_DEFAULT, + queues: HashMap::new(), + retry_policy: Arc::new(DefaultRetryPolicy::default()), + soft_stop_timeout: None, + work_middleware: Vec::new(), + workers: WorkerRegistry::new(), + } + } + + /// Creates a non-owning handle for an extension service. + #[must_use] + pub(crate) fn downgrade(&self) -> WeakClient { + WeakClient { + inner: Arc::downgrade(&self.inner), + } + } + + /// Stable identifier recorded in `attempted_by`. + #[must_use] + pub fn id(&self) -> &str { + &self.inner.id + } +} + +impl Client { + /// Returns the dynamic periodic-job bundle for this client. + /// + /// Only the elected leader enqueues periodic jobs, so jobs added here + /// take effect only while this client leads. To fully enable or disable a + /// periodic job, change it on every client eligible for leader election. + /// A client built with + /// [`without_leader_election`](ClientBuilder::without_leader_election) + /// rejects additions. + #[must_use] + pub fn periodic_jobs(&self) -> PeriodicJobs { + self.inner.periodic_jobs.clone() + } + + /// Returns the database this client was built with, with its pool and + /// backend-specific options. + #[must_use] + pub fn database(&self) -> ClientDatabase<'_> { + self.inner.database.client_database() + } + + /// Subscribes to selected local client events with a bounded buffer. + /// + /// # Errors + /// + /// Returns the errors of [`Client::subscribe_config`], and an error when + /// `kinds` is empty. + pub fn subscribe(&self, kinds: &[EventKind]) -> Result { + self.subscribe_config(SubscribeConfig::new(kinds.iter().copied())?) + } + + /// Subscribes with an explicit bounded-buffer capacity. When the receiver + /// falls behind, the next receive reports how many events were dropped. + /// + /// # Errors + /// + /// Returns [`Error::Configuration`] when the client works no queues, and + /// [`Error::RuntimeUnavailable`] outside a Tokio runtime. + pub fn subscribe_config(&self, config: SubscribeConfig) -> Result { + let (buffer_capacity, kinds) = config.into_parts(); + if self + .inner + .queues + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_empty() + { + return Err(Error::configuration( + "event subscriptions require a client configured to work queues".to_owned(), + )); + } + let kinds = crate::event::validate_kinds(&kinds)?; + let mut source = self.inner.events.subscribe(); + let (sender, receiver) = mpsc::channel(buffer_capacity.get()); + let dropped = Arc::new(AtomicU64::new(0)); + let dropped_for_task = Arc::clone(&dropped); + tokio::runtime::Handle::try_current().map_err(|_| Error::RuntimeUnavailable { + operation: "event subscriptions", + })?; + tokio::spawn(async move { + loop { + // Stop forwarding as soon as the subscriber drops its + // receiver, rather than at the next matching event. + let next = tokio::select! { + () = sender.closed() => break, + next = source.recv() => next, + }; + match next { + Ok(event) if kinds.contains(&event.kind()) => match sender.try_send(event) { + Ok(()) => {} + Err(mpsc::error::TrySendError::Full(_)) => { + dropped_for_task.fetch_add(1, Ordering::Relaxed); + } + Err(mpsc::error::TrySendError::Closed(_)) => break, + }, + Ok(_) => {} + Err(broadcast::error::RecvError::Lagged(count)) => { + dropped_for_task.fetch_add(count, Ordering::Relaxed); + } + Err(broadcast::error::RecvError::Closed) => break, + } + } + }); + Ok(EventReceiver::new(dropped, receiver)) + } +} + +impl Client { + pub(crate) fn default_max_attempts(&self) -> i16 { + self.inner.default_max_attempts + } + + pub(crate) fn signal_queue_control(&self, queue: &str) { + let _ = self + .inner + .queue_notifications + .send(RuntimeNotification::QueueControl(queue.to_owned())); + } +} + +impl Client { + fn validate_known_kind(&self, kind: &str) -> Result<(), Error> { + if !self.inner.allow_unregistered_job_kinds + && !self.inner.workers.kinds().is_empty() + && !self.inner.workers.contains_kind(kind) + { + return Err(Error::UnknownJobKind { + kind: kind.to_owned(), + }); + } + Ok(()) + } +} + +/// Formats a time as River Go's `time.Time` JSON (RFC 3339 with nanoseconds +/// and trailing zeros trimmed, in UTC), which River stores for +/// `cancel_attempted_at`. +pub(crate) fn go_time_json(time: DateTime) -> String { + let formatted = time.to_rfc3339_opts(chrono::SecondsFormat::Nanos, true); + let Some((seconds, fraction)) = formatted.trim_end_matches('Z').split_once('.') else { + return formatted; + }; + let fraction = fraction.trim_end_matches('0'); + if fraction.is_empty() { + format!("{seconds}Z") + } else { + format!("{seconds}.{fraction}Z") + } +} + +/// Generates a client ID unique to this `Client` instance. +/// +/// Like Go, the ID combines the host name (dots replaced by underscores and +/// truncated to 60 bytes) with the creation time to the microsecond. A random +/// suffix keeps IDs distinct when several clients start in the same +/// microsecond or containers report identical host names, because a shared ID +/// would let two clients renew one leadership lease. +fn default_client_id() -> String { + default_client_id_with_host(&host_name(), Utc::now(), crate::maintenance::random_u64()) +} + +fn default_client_id_with_host(host: &str, created_at: DateTime, random: u64) -> String { + const MAX_HOST_LENGTH: usize = 60; + + let mut host = host.replace('.', "_"); + if host.len() > MAX_HOST_LENGTH { + let mut end = MAX_HOST_LENGTH; + while !host.is_char_boundary(end) { + end -= 1; + } + host.truncate(end); + } + format!( + "{host}_{}_{:08x}", + created_at.format("%Y_%m_%dT%H_%M_%S_%6f"), + random & 0xffff_ffff + ) +} + +fn host_name() -> String { + std::env::var("HOSTNAME") + .ok() + .or_else(|| std::fs::read_to_string("/proc/sys/kernel/hostname").ok()) + .or_else(|| std::fs::read_to_string("/etc/hostname").ok()) + .map(|host| host.trim().to_owned()) + .filter(|host| !host.is_empty()) + .unwrap_or_else(|| "unknown_host".to_owned()) +} + +#[cfg(test)] +mod default_client_id_tests { + use chrono::{TimeZone, Timelike}; + + use super::*; + + #[test] + fn default_client_id_matches_go_shape_and_is_unique() { + let created_at = Utc + .with_ymd_and_hms(2026, 1, 2, 3, 4, 5) + .unwrap() + .with_nanosecond(678_901_000) + .unwrap(); + assert_eq!( + default_client_id_with_host("worker.example.com", created_at, 0xdead_beef), + "worker_example_com_2026_01_02T03_04_05_678901_deadbeef" + ); + let long = "h".repeat(80); + let id = default_client_id_with_host(&long, created_at, 1); + assert!(id.starts_with(&"h".repeat(60))); + assert!(!id.starts_with(&"h".repeat(61))); + assert!(id.len() <= 100, "client IDs are limited to 100 bytes"); + + let first = default_client_id(); + let second = default_client_id(); + assert_ne!(first, second); + } +} diff --git a/rust/riverqueue/src/client/notifier.rs b/rust/riverqueue/src/client/notifier.rs new file mode 100644 index 000000000..da75fb637 --- /dev/null +++ b/rust/riverqueue/src/client/notifier.rs @@ -0,0 +1,440 @@ +//! Notification listeners and wakeups. +//! +//! Notifications only shorten the time before producers notice new work, +//! cancellations, and leadership changes; producers keep polling regardless. +//! Like River Go's notifier, a listener therefore never stops the client when +//! the database becomes unavailable. It reconnects with exponential backoff, +//! resubscribes to every topic, and wakes all producers after reconnecting in +//! case notifications were missed while it was disconnected. + +use std::sync::{Arc, Mutex}; +#[cfg(feature = "postgres")] +use std::time::Duration; + +use serde::Deserialize; +#[cfg(feature = "postgres")] +use sqlx::postgres::PgListener; +use tokio::sync::{broadcast, oneshot}; +use tokio_util::sync::CancellationToken; +#[cfg(feature = "postgres")] +use tracing::info; +use tracing::{debug, error, warn}; + +use crate::Error; +use crate::client::ClientInner; +use crate::client::attempts::signal_running_attempt; +use crate::client::backoff::exponential_backoff; +#[cfg(feature = "sqlite")] +use crate::client::sqlite_backend_error; +use crate::maintenance::LeadershipWakeup; + +/// Timeout for connecting and subscribing, matching River Go's +/// `listenerTimeout`. +#[cfg(feature = "postgres")] +const LISTENER_TIMEOUT: Duration = Duration::from_secs(10); +/// Idle time after which the listener pings its connection, matching River +/// Go's notifier. A dead connection is then detected even without traffic. +#[cfg(feature = "postgres")] +const LISTENER_PING_INTERVAL: Duration = Duration::from_secs(5); + +#[derive(Deserialize)] +pub(super) struct ControlNotification { + pub(super) action: String, + pub(super) job_id: Option, + pub(super) queue: Option, +} + +#[derive(Deserialize)] +pub(super) struct InsertNotification { + pub(super) queue: String, +} + +#[derive(Deserialize)] +pub(super) struct LeadershipNotification { + pub(super) action: String, + pub(super) leader_id: Option, +} + +/// Wakeups for producers. Leadership events travel separately, on +/// [`ClientInner::leadership_wakeups`], so this busier channel lagging can +/// never drop one. +#[derive(Clone, Debug)] +pub(crate) enum RuntimeNotification { + Insert(String), + QueueControl(String), +} + +/// Readiness reported once by the notification path. +pub(super) type ReadySender = oneshot::Sender<()>; + +/// Holds the notification path's readiness until a listener reports it. The +/// supervisor keeps the slot, so a listener restarted after a panic still +/// reports readiness instead of the client appearing to have stopped. +pub(super) type ReadySlot = Arc>>; + +fn report_ready(ready: &ReadySlot) { + if let Some(ready) = ready + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take() + { + let _ = ready.send(()); + } +} + +/// Wakes every producer to fetch and refresh its queue state. Used after a +/// listener reconnects, when notifications may have been missed. +#[cfg(feature = "postgres")] +fn wake_all_producers(queue_notifications: &broadcast::Sender) { + let _ = queue_notifications.send(RuntimeNotification::Insert("*".to_owned())); + let _ = queue_notifications.send(RuntimeNotification::QueueControl("*".to_owned())); +} + +/// Routes one River notification to local producers and services. +pub(super) fn dispatch_notification( + inner: &ClientInner, + queue_notifications: &broadcast::Sender, + topic: &str, + payload: &str, +) { + match topic { + crate::NOTIFICATION_TOPIC_INSERT => { + if let Ok(payload) = serde_json::from_str::(payload) { + let _ = queue_notifications.send(RuntimeNotification::Insert(payload.queue)); + } + } + crate::NOTIFICATION_TOPIC_LEADERSHIP => { + if let Ok(payload) = serde_json::from_str::(payload) { + if payload.action == "resigned" + && payload.leader_id.as_deref() == Some(inner.id.as_str()) + { + return; + } + let wakeup = if payload.action == "request_resign" { + LeadershipWakeup::RequestResign + } else { + LeadershipWakeup::Changed + }; + let _ = inner.leadership_wakeups.send(wakeup); + } + } + crate::NOTIFICATION_TOPIC_CONTROL => { + let Ok(payload) = serde_json::from_str::(payload) else { + warn!(payload, "ignored invalid River control notification"); + return; + }; + match payload.action.as_str() { + "cancel" => { + if let Some(job_id) = payload.job_id { + signal_running_attempt( + &inner.running, + &inner.pending_cancellations, + &inner.fetch_registration_windows, + job_id, + ); + } + } + // Like Go's producer, a metadata change is handled at once + // rather than at the next queue poll. The producer reloads + // the queue and reports changed metadata to the extension. + "metadata_changed" | "pause" | "resume" => { + if let Some(queue) = payload.queue { + let _ = queue_notifications.send(RuntimeNotification::QueueControl(queue)); + } + } + _ => debug!( + action = payload.action, + "ignored unknown River control action" + ), + } + } + _ => {} + } +} + +/// Listens for PostgreSQL notifications until cancelled. +/// +/// The listener uses a dedicated connection opened with the client pool's +/// connect options, like River Go's hijacked listener connection, so it never +/// occupies one of the caller's pool slots. Connection and subscription +/// failures are retried with River's service backoff for as long as the client +/// runs; `ready` is reported once the first subscription succeeds. +#[cfg(feature = "postgres")] +pub(super) async fn run_notifications( + inner: Arc, + pool: sqlx::PgPool, + cancel: CancellationToken, + queue_notifications: broadcast::Sender, + ready: ReadySlot, +) -> Result<(), Error> { + let mut attempt = 0; + let mut missed_notifications = false; + let mut schema = None; + loop { + let result = tokio::select! { + () = cancel.cancelled() => return Ok(()), + result = listen_until_error( + &inner, + &pool, + &cancel, + &queue_notifications, + &ready, + &mut schema, + &mut attempt, + missed_notifications, + ) => result, + }; + let Err(listener_error) = result else { + return Ok(()); + }; + missed_notifications = true; + attempt += 1; + let sleep = exponential_backoff(attempt); + error!( + attempt, + error = %crate::error::Chain(&listener_error), + sleep_duration = ?sleep, + "River notification listener failed (will reconnect after backoff); producers keep polling" + ); + tokio::select! { + () = cancel.cancelled() => return Ok(()), + () = tokio::time::sleep(sleep) => {} + } + } +} + +/// On a server without `LISTEN`/`NOTIFY`, like YugabyteDB by default, +/// reports the notifier ready and idles until cancelled, returning `true`, +/// so the client polls as without notifications. Detects the server the +/// first time. +#[cfg(feature = "postgres")] +async fn idles_without_listen_notify( + inner: &ClientInner, + pool: &sqlx::PgPool, + cancel: &CancellationToken, + ready: &ReadySlot, +) -> Result { + let capabilities = tokio::time::timeout(LISTENER_TIMEOUT, inner.postgres_capabilities(pool)) + .await + .map_err(|_| { + Error::runtime_context( + "notification listener", + "timed out detecting database capabilities".to_owned(), + ) + })??; + if capabilities.supports_listen_notify { + return Ok(false); + } + info!("River's database does not support LISTEN/NOTIFY; polling instead"); + report_ready(ready); + cancel.cancelled().await; + Ok(true) +} + +/// Connects, subscribes, and dispatches notifications until the connection +/// fails. Returns `Ok` only when cancelled. +#[cfg(feature = "postgres")] +#[expect( + clippy::too_many_arguments, + reason = "the reconnect loop threads its listener state through each attempt" +)] +async fn listen_until_error( + inner: &ClientInner, + pool: &sqlx::PgPool, + cancel: &CancellationToken, + queue_notifications: &broadcast::Sender, + ready: &ReadySlot, + schema: &mut Option, + attempt: &mut u32, + missed_notifications: bool, +) -> Result<(), Error> { + if idles_without_listen_notify(inner, pool, cancel, ready).await? { + return Ok(()); + } + let schema = if let Some(schema) = schema { + schema.clone() + } else { + let resolved = match inner.schema.as_deref() { + Some(schema) => schema.to_owned(), + None => tokio::time::timeout( + LISTENER_TIMEOUT, + sqlx::query_scalar::<_, Option>("SELECT current_schema()").fetch_one(pool), + ) + .await + .map_err(|_| { + Error::runtime_context( + "notification listener", + "timed out resolving the current schema".to_owned(), + ) + })?? + .ok_or_else(|| { + Error::runtime_context( + "notification listener", + "PostgreSQL current_schema() is null", + ) + })?, + }; + schema.insert(resolved).clone() + }; + let topics = [ + crate::NOTIFICATION_TOPIC_CONTROL, + crate::NOTIFICATION_TOPIC_INSERT, + crate::NOTIFICATION_TOPIC_LEADERSHIP, + ] + .map(|topic| (format!("{schema}.{topic}"), topic)); + + // A private one-connection pool lets `PgListener` reconnect by itself + // without borrowing from, or being limited by, the caller's pool. + let listener_pool = sqlx::postgres::PgPoolOptions::new() + .acquire_timeout(LISTENER_TIMEOUT) + .idle_timeout(None) + .max_connections(1) + .max_lifetime(None) + .min_connections(0) + .connect_lazy_with((*pool.connect_options()).clone()); + let mut listener = tokio::time::timeout(LISTENER_TIMEOUT, async { + let mut listener = PgListener::connect_with(&listener_pool).await?; + listener + .listen_all(topics.iter().map(|(channel, _)| channel.as_str())) + .await?; + Ok::<_, sqlx::Error>(listener) + }) + .await + .map_err(|_| { + Error::runtime_context( + "notification listener", + "timed out connecting and subscribing".to_owned(), + ) + })??; + debug!("River notification listener healthy"); + *attempt = 0; + report_ready(ready); + if missed_notifications { + wake_all_producers(queue_notifications); + } + + loop { + let received = tokio::select! { + () = cancel.cancelled() => return Ok(()), + received = tokio::time::timeout(LISTENER_PING_INTERVAL, listener.try_recv()) => received, + }; + match received { + Ok(Ok(Some(notification))) => { + let Some((_, topic)) = topics + .iter() + .find(|(channel, _)| channel == notification.channel()) + else { + continue; + }; + dispatch_notification(inner, queue_notifications, topic, notification.payload()); + } + // The connection dropped and `PgListener` reconnected and + // resubscribed before returning. Anything sent in between was + // lost, so producers must look for themselves. + Ok(Ok(None)) => { + warn!("River notification listener reconnected; waking producers"); + wake_all_producers(queue_notifications); + } + Ok(Err(listener_error)) => return Err(listener_error.into()), + Err(_) => { + tokio::time::timeout( + LISTENER_TIMEOUT, + sqlx::query("SELECT 1").execute(&mut listener), + ) + .await + .map_err(|_| { + Error::runtime_context( + "notification listener", + "health check timed out".to_owned(), + ) + })??; + } + } + } +} + +/// Polls SQLite's notification outbox until cancelled. +/// +/// Like River Go's SQLite listener, the first successful poll connects after +/// the outbox's current maximum ID and subscribes to River's topics, so +/// earlier rows are never delivered. Poll failures (for example `database is +/// locked` while another process holds the write lock) are logged and retried +/// with River's service backoff. The cursor is kept, so no durable +/// notification is skipped. +#[cfg(feature = "sqlite")] +pub(super) async fn run_sqlite_notifications( + inner: Arc, + pool: sqlx::SqlitePool, + cancel: CancellationToken, + queue_notifications: broadcast::Sender, + ready: ReadySlot, +) -> Result<(), Error> { + let mut attempt = 0; + let mut listener = crate::database::sqlite::NotificationListener::default(); + let mut notification_tick = + tokio::time::interval(crate::database::sqlite::DEFAULT_NOTIFICATION_POLL_INTERVAL); + notification_tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + loop { + tokio::select! { + () = cancel.cancelled() => return Ok(()), + _ = notification_tick.tick() => {} + } + // Like the rest of a poll, waiting for a connection during an + // outage must not delay a stop, as Go's cancellable context doesn't. + let polled = async { + let mut connection = tokio::select! { + biased; + () = cancel.cancelled() => return Ok(()), + connection = pool.acquire() => connection?, + }; + if !listener.is_connected() { + listener + .connect(&mut connection) + .await + .map_err(sqlite_backend_error)?; + } + for topic in [ + crate::NOTIFICATION_TOPIC_CONTROL, + crate::NOTIFICATION_TOPIC_INSERT, + crate::NOTIFICATION_TOPIC_LEADERSHIP, + ] { + listener + .listen(&mut connection, topic) + .await + .map_err(sqlite_backend_error)?; + } + report_ready(&ready); + while let Some(notification) = listener + .next(&mut connection) + .await + .map_err(sqlite_backend_error)? + { + dispatch_notification( + &inner, + &queue_notifications, + ¬ification.topic, + ¬ification.payload, + ); + } + Ok::<_, Error>(()) + } + .await; + if let Err(poll_error) = polled { + attempt += 1; + let sleep = exponential_backoff(attempt); + error!( + attempt, + error = %crate::error::Chain(&poll_error), + sleep_duration = ?sleep, + "River notification poll failed (will retry after backoff); producers keep polling" + ); + tokio::select! { + () = cancel.cancelled() => return Ok(()), + () = tokio::time::sleep(sleep) => {} + } + continue; + } + attempt = 0; + } +} diff --git a/rust/riverqueue/src/client/notify_limiter.rs b/rust/riverqueue/src/client/notify_limiter.rs new file mode 100644 index 000000000..8225fd724 --- /dev/null +++ b/rust/riverqueue/src/client/notify_limiter.rs @@ -0,0 +1,96 @@ +//! Paces a client's insert notifications, a port of Go's +//! `notifylimiter.Limiter`. + +use std::collections::HashMap; +use std::sync::Mutex; +use std::time::{Duration, Instant}; + +/// Allows at most one insert notification per queue within each cooldown. +/// +/// Producers fetch at most once per fetch cooldown, so a burst of insertions +/// into one queue needs only its first notification. Like Go, a queue's +/// window starts when a notification is allowed, whether or not the +/// transaction carrying it commits. +#[derive(Debug)] +pub(crate) struct InsertNotifyLimiter { + cooldown: Duration, + last_sent: Mutex>, +} + +impl InsertNotifyLimiter { + pub(crate) fn new(cooldown: Duration) -> Self { + Self { + cooldown, + last_sent: Mutex::new(HashMap::new()), + } + } + + /// Returns the queues among `queues` that are due a notification, + /// recording each as notified now. + pub(crate) fn due<'q>(&self, queues: impl IntoIterator) -> Vec<&'q str> { + let now = Instant::now(); + queues + .into_iter() + .filter(|queue| self.should_trigger_at(queue, now)) + .collect() + } + + /// Returns whether `queue` is due a notification at `now`, recording it + /// as notified then if so. A queue is due once more than the cooldown + /// has passed since its last notification. + fn should_trigger_at(&self, queue: &str, now: Instant) -> bool { + let mut last_sent = self + .last_sent + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if last_sent + .get(queue) + .is_some_and(|last| now.saturating_duration_since(*last) <= self.cooldown) + { + return false; + } + last_sent.insert(queue.to_owned(), now); + true + } +} + +#[cfg(test)] +mod tests { + use std::time::{Duration, Instant}; + + use super::InsertNotifyLimiter; + + #[test] + fn allows_one_notification_per_queue_per_cooldown() { + let limiter = InsertNotifyLimiter::new(Duration::from_millis(100)); + let start = Instant::now(); + + assert!(limiter.should_trigger_at("a", start)); + for _ in 0..10 { + assert!(!limiter.should_trigger_at("a", start)); + } + assert!(!limiter.should_trigger_at("a", start + Duration::from_millis(100))); + assert!(limiter.should_trigger_at("a", start + Duration::from_millis(101))); + assert!(!limiter.should_trigger_at("a", start + Duration::from_millis(150))); + } + + #[test] + fn due_filters_and_records_queues() { + let limiter = InsertNotifyLimiter::new(Duration::from_hours(1)); + + assert_eq!(limiter.due(["a", "b"]), ["a", "b"]); + assert_eq!(limiter.due(["a", "c"]), ["c"]); + assert_eq!(limiter.due(["a", "b", "c"]), Vec::<&str>::new()); + } + + #[test] + fn tracks_queues_independently() { + let limiter = InsertNotifyLimiter::new(Duration::from_millis(100)); + let start = Instant::now(); + + assert!(limiter.should_trigger_at("a", start)); + assert!(limiter.should_trigger_at("b", start)); + assert!(!limiter.should_trigger_at("a", start + Duration::from_millis(50))); + assert!(limiter.should_trigger_at("c", start + Duration::from_millis(50))); + } +} diff --git a/rust/riverqueue/src/client/peers.rs b/rust/riverqueue/src/client/peers.rs new file mode 100644 index 000000000..2e508071d --- /dev/null +++ b/rust/riverqueue/src/client/peers.rs @@ -0,0 +1,697 @@ +//! Peer attempts: jobs a running attempt, their coordinator, claims and +//! completes alongside its own job, such as a group of related jobs it works +//! together. +//! +//! River owns each peer from the commit of the claim that took it until its +//! outcome persists, under the attempt that claimed it. A peer moves from +//! `Claimed` to `Preparing` once an outcome for it is accepted, to +//! `Submitted` once the completer accepts that outcome, and to `Settled` once +//! the completer persisted it; an outcome that fails before the completer +//! accepts it returns the peer to `Claimed`. When the coordinator ends, +//! River stops accepting its peer operations, waits for those it accepted, +//! and completes every peer still `Claimed` with an outcome of its own, all +//! before the coordinator's own outcome. Peers don't take producer slots, and +//! their producer's session never hears about them, like the other jobs of a +//! multi-job result in River for Go. +//! +//! A soft stop doesn't end a coordinator's claims: the producer stops +//! fetching new jobs, but a running coordinator may keep claiming peers until +//! its attempt ends, so it can finish gathering and work the group it has. +//! The stop waits for that attempt, which settles every peer before its own +//! outcome, so the stop waits for the peers too. Only the attempt's +//! cancellation, by a hard stop or a remote cancellation, and the attempt's +//! end refuse claims. + +use std::collections::HashMap; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use futures_util::future::BoxFuture; +use serde_json::{Map, Value}; +use tokio::sync::{Notify, mpsc, oneshot}; +use tokio_util::sync::CancellationToken; +use tracing::error; + +use crate::__private::{ClaimedJob, PilotDatabase, PilotError}; +use crate::client::ClientInner; +use crate::client::completer::{CompletionAttempt, CompletionTiming, CompletionUpdate}; +use crate::client::executor::{ + WorkerFailure, WorkerFailureKind, WorkerResult, persist_result, public_work_result, + worker_failure_from_source, +}; +use crate::{Client, Error, ErrorHandlerDecision, JobRow, JobState, WorkContext, WorkResult}; + +/// Numbers ledgers so the client-wide owner map can tell them apart. +static LEDGER_IDS: AtomicU64 = AtomicU64::new(1); + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum PeerState { + Claimed, + Preparing, + Submitted, + Settled, +} + +struct Peer { + /// The row as claimed, which identifies the peer's attempt. + claimed: JobRow, + state: PeerState, +} + +#[derive(Default)] +struct LedgerState { + /// Set once the coordinator ended; no new operation starts. + closed: bool, + /// Set once River stopped tracking the peers; see `abandon`. + released: bool, + /// Operations accepted and not yet settled. + operations: usize, + /// Peers by job ID, including settled ones. + peers: HashMap, +} + +/// The peers of one coordinating attempt. +pub(crate) struct PeerLedger { + /// The coordinator's job ID. + coordinator: i64, + /// When the coordinator's attempt started, recorded as its peers' + /// attempt errors' time. + started_at: DateTime, + id: u64, + idle: Notify, + state: Mutex, +} + +impl std::fmt::Debug for PeerLedger { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("PeerLedger") + .field("coordinator", &self.coordinator) + .finish_non_exhaustive() + } +} + +/// Ends an accepted operation, waking a coordinator waiting to finish. +struct Operation<'a>(&'a PeerLedger); + +impl Drop for Operation<'_> { + fn drop(&mut self) { + let mut state = self.0.lock(); + state.operations -= 1; + if state.operations == 0 { + self.0.idle.notify_waiters(); + } + } +} + +/// Peer jobs a claim reserved in the client's owner map, released when the +/// guard drops unless the claim recorded them as peers. +struct ReservedOwners<'a> { + ids: Vec, + inner: &'a ClientInner, + ledger: u64, +} + +impl Drop for ReservedOwners<'_> { + fn drop(&mut self) { + for id in &self.ids { + self.inner.release_peer(*id, self.ledger); + } + } +} + +/// Peers whose outcomes a submission is handing to the completer. When the +/// guard drops, a peer still `Preparing`, whose outcome never reached the +/// completer because the operation's future was dropped, returns to +/// `Claimed`, so its coordinator's end still gives it an outcome. +struct Handover<'a> { + ids: Vec, + ledger: &'a PeerLedger, +} + +impl Drop for Handover<'_> { + fn drop(&mut self) { + for id in &self.ids { + self.ledger + .set_state(*id, PeerState::Claimed, PeerState::Preparing); + } + } +} + +/// Tells a peer's ledger that the completer persisted its outcome. +pub(super) struct PeerCompletion { + done: Mutex>>, + job_id: i64, + ledger: Arc, +} + +impl PeerCompletion { + /// Ends ownership as the outcome persists, before its event, so the job + /// can be claimed again at once. + pub(super) fn persisted(&self, inner: &ClientInner) { + if let Some(peer) = self.ledger.lock().peers.get_mut(&self.job_id) { + peer.state = PeerState::Settled; + } + inner.release_peer(self.job_id, self.ledger.id); + if let Some(done) = self + .done + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take() + { + let _ = done.send(()); + } + } +} + +/// An outcome for one peer and the peer it's for. +struct Submission { + job: JobRow, + result: WorkerResult, +} + +fn peer_error(message: impl Into) -> Error { + Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "peer attempts", + }, + source: message.into().into(), + } +} + +fn not_running() -> Error { + peer_error("peer operations require a running attempt") +} + +impl PeerLedger { + pub(super) fn new(coordinator: i64, started_at: DateTime) -> Self { + Self { + coordinator, + started_at, + id: LEDGER_IDS.fetch_add(1, Ordering::Relaxed), + idle: Notify::new(), + state: Mutex::new(LedgerState::default()), + } + } + + fn lock(&self) -> std::sync::MutexGuard<'_, LedgerState> { + self.state + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + /// Accepts an operation unless the coordinator ended. + fn begin(&self) -> Result, Error> { + let mut state = self.lock(); + if state.closed { + return Err(not_running()); + } + state.operations += 1; + Ok(Operation(self)) + } + + /// Claims peers with `run` in a transaction River commits, and returns + /// the decoded rows River now tracks. Rows that couldn't be decoded are + /// completed as failures instead. Claims continue through a soft stop and + /// end with the attempt's cancellation or its end. + pub(crate) async fn claim( + self: &Arc, + inner: &Arc, + context: &WorkContext, + run: F, + ) -> Result, Error> + where + F: for<'c> FnOnce( + crate::__private::PeerClaimContext<'c>, + ) -> BoxFuture<'c, Result, PilotError>> + + Send, + { + let _operation = self.begin()?; + let cancellation = context.cancellation_token(); + if cancellation.is_cancelled() { + return Err(peer_error("peer claim cancelled")); + } + let database: PilotDatabase = inner.pilot_database(); + let mut transaction = tokio::select! { + biased; + () = cancellation.cancelled() => return Err(peer_error("peer claim cancelled")), + transaction = database.begin() => transaction?, + }; + let claimed = run(crate::__private::PeerClaimContext { + cancellation, + client_id: &inner.id, + connection: transaction.connection(), + database: &database, + }) + .await + .map_err(|source| Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "peer attempts", + }, + source, + })?; + // A claim whose coordinator was cancelled before commit rolls back. + if cancellation.is_cancelled() { + return Err(peer_error("peer claim cancelled")); + } + // Released again unless the rows become the coordinator's peers, + // including when this future is dropped during the commit. + let mut reserved = ReservedOwners { + ids: self.reserve_claim(inner, &claimed)?, + inner, + ledger: self.id, + }; + transaction.commit().await?; + // The claim committed: every row is this coordinator's peer now, + // even when the coordinator was cancelled meanwhile, unless River + // stopped tracking its peers while the claim ran, which leaves the + // rows to the rescuer. Checking and recording under one lock keeps + // `abandon` from running in between. + let mut rows = Vec::new(); + let mut failures = Vec::new(); + { + let mut state = self.lock(); + if state.released { + return Err(not_running()); + } + reserved.ids.clear(); + for claimed in claimed { + let decode_error = claimed.decode_error().map(str::to_owned); + let Some(row) = claimed + .into_decoded() + .map_or_else(|undecodable| undecodable.row.map(|row| *row), Some) + else { + continue; + }; + state.peers.insert( + row.id, + Peer { + claimed: row.clone(), + state: if decode_error.is_some() { + PeerState::Preparing + } else { + PeerState::Claimed + }, + }, + ); + match decode_error { + Some(error) => failures.push(Submission { + job: row, + result: Err(worker_failure_from_source( + format!("job row couldn't be decoded: {error}").into(), + )), + }), + None => rows.push(row), + } + } + } + if !failures.is_empty() { + self.submit(inner, context, failures).await?; + } + Ok(rows) + } + + /// Checks a claim's rows before it commits, and reserves them so no other + /// claim of this client can take them meanwhile. + fn reserve_claim( + &self, + inner: &ClientInner, + claimed: &[ClaimedJob], + ) -> Result, Error> { + let fail = |reason: String| peer_error(format!("a peer claim returned {reason}")); + let state = self.lock(); + let running = inner + .running + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let mut owners = inner.lock_peer_owners(); + let mut seen = std::collections::HashSet::new(); + for job in claimed { + let Some(row) = job.row() else { + return Err(fail(format!( + "a row that couldn't be identified: {}", + job.decode_error().unwrap_or_default() + ))); + }; + let id = row.id; + if !seen.insert(id) { + return Err(fail(format!("job {id} twice"))); + } + if id == self.coordinator { + return Err(fail(format!("job {id}, the claiming attempt's own job"))); + } + if owners.contains_key(&id) { + return Err(fail(format!( + "job {id}, which this client already works as a peer" + ))); + } + if running.contains_key(&id) { + return Err(fail(format!("job {id}, which this client already works"))); + } + if row.attempt < 1 { + return Err(fail(format!("job {id}, which has no attempt"))); + } + if let Some(earlier) = state.peers.get(&id) + && row.attempt <= earlier.claimed.attempt + { + return Err(fail(format!( + "job {id} at attempt {}, which already ended here", + row.attempt + ))); + } + if row.state != JobState::Running { + return Err(fail(format!("job {id}, which isn't running"))); + } + let attempted_by_undecodable = job.column_undecodable("attempted_by"); + if !attempted_by_undecodable + && row.attempted_by.last().map(String::as_str) != Some(inner.id.as_str()) + { + return Err(fail(format!("job {id}, which another client claimed"))); + } + } + let reserved = seen.into_iter().collect::>(); + for id in &reserved { + owners.insert(*id, self.id); + } + Ok(reserved) + } + + /// Completes peers through the ordinary completion pipeline, returning + /// once every outcome persisted. + pub(crate) async fn complete( + self: &Arc, + inner: &Arc, + context: &WorkContext, + outcomes: Vec, + ) -> Result<(), Error> { + let _operation = self.begin()?; + let submissions = self.reserve_outcomes(outcomes)?; + self.submit(inner, context, submissions).await + } + + /// Accepts one outcome for each peer, all or none: each job must be a + /// peer of this ledger, identified by its ID, attempt, and attempting + /// client, have no outcome yet, and appear once. + fn reserve_outcomes( + &self, + outcomes: Vec, + ) -> Result, Error> { + let mut state = self.lock(); + let mut seen = std::collections::HashSet::new(); + for outcome in &outcomes { + let id = outcome.job.id; + let Some(peer) = state.peers.get(&id) else { + return Err(peer_error(format!( + "job {id} isn't a peer of the attempt completing it" + ))); + }; + if outcome.job.attempt != peer.claimed.attempt + || outcome.job.attempted_by.last() != peer.claimed.attempted_by.last() + { + return Err(peer_error(format!( + "job {id} attempt {} isn't the peer attempt {} this attempt owns", + outcome.job.attempt, peer.claimed.attempt + ))); + } + if !seen.insert(id) { + return Err(peer_error(format!("job {id} has two outcomes"))); + } + if peer.state != PeerState::Claimed { + return Err(peer_error(format!("job {id} already has an outcome"))); + } + } + let mut submissions = Vec::with_capacity(outcomes.len()); + for outcome in outcomes { + let peer = state + .peers + .get_mut(&outcome.job.id) + .expect("peer checked above"); + peer.state = PeerState::Preparing; + submissions.push(Submission { + job: peer.claimed.clone(), + result: outcome.result.map_err(worker_failure_from_source), + }); + } + Ok(submissions) + } + + /// Hands outcomes to the completer, running the error handler and adding + /// the coordinator's metadata first, like the coordinator's own outcome, + /// and waits for them to persist. An outcome the completer doesn't + /// accept, including one whose submission is dropped first, returns its + /// peer to `Claimed`. + async fn submit( + self: &Arc, + inner: &Arc, + context: &WorkContext, + submissions: Vec, + ) -> Result<(), Error> { + let _handover = Handover { + ids: submissions + .iter() + .map(|submission| submission.job.id) + .collect(), + ledger: self, + }; + let sender = inner.completion_sender(); + let shared_metadata = context.metadata_updates(); + let mut persisting = Vec::with_capacity(submissions.len()); + let mut failure = None; + for submission in submissions { + let id = submission.job.id; + let (done, persisted) = oneshot::channel(); + let completion = Arc::new(PeerCompletion { + done: Mutex::new(Some(done)), + job_id: id, + ledger: Arc::clone(self), + }); + let submitted = match &sender { + Some(sender) => { + self.persist( + inner, + context, + submission, + &shared_metadata, + sender, + completion, + ) + .await + } + None => Err(Error::runtime_context( + "job completion", + "client runtime is not accepting job completions", + )), + }; + // No await separates the completer accepting the outcome from + // this state change, so a dropped submission never mistakes an + // accepted outcome for a missing one. + match submitted { + Ok(()) => { + self.set_state(id, PeerState::Submitted, PeerState::Preparing); + persisting.push((id, persisted)); + } + Err(submit_error) => { + self.set_state(id, PeerState::Claimed, PeerState::Preparing); + failure.get_or_insert(submit_error); + } + } + } + for (id, persisted) in persisting { + if persisted.await.is_err() { + // The completer gave the outcome up; the peer stays owned + // until its coordinator ends and is left to the rescuer. + failure.get_or_insert_with(|| { + Error::runtime_context( + "job completion", + format!("the outcome of peer job {id} was not persisted"), + ) + }); + } + } + failure.map_or(Ok(()), Err) + } + + fn set_state(&self, id: i64, to: PeerState, from: PeerState) { + if let Some(peer) = self.lock().peers.get_mut(&id) + && peer.state == from + { + peer.state = to; + } + } + + async fn persist( + &self, + inner: &Arc, + coordinator: &WorkContext, + Submission { job: row, result }: Submission, + shared_metadata: &Map, + sender: &mpsc::Sender, + completion: Arc, + ) -> Result<(), Error> { + let context = WorkContext::for_job( + Client { + inner: Arc::clone(inner), + }, + coordinator.cancellation_token().clone(), + row.id, + &row.metadata, + ); + for (key, value) in shared_metadata { + context.insert_metadata(key.clone(), value.clone()); + } + let work_result = public_work_result(&result); + let mut decision = ErrorHandlerDecision::default(); + // Like a worker's result, only a failed outcome runs the error + // handler: completions, snoozes, cancellations, discards, and + // interruptions don't. + if let Some(error_handler) = &inner.error_handler + && matches!(work_result, WorkResult::Failed(_)) + { + match error_handler + .handle_error(&context, &row, &work_result) + .await + { + Ok(handler_decision) => decision = handler_decision, + Err(handler_error) => { + error!(error = %crate::error::Chain(&handler_error), "River error handler failed"); + } + } + } + let queue_wait_duration = row + .attempted_at + .and_then(|attempted_at| { + (attempted_at - row.scheduled_at.max(row.created_at)) + .to_std() + .ok() + }) + .unwrap_or_default(); + persist_result( + inner, + &row, + self.started_at, + &CompletionAttempt { + cancellation: CancellationToken::new(), + timing: CompletionTiming { + completion_started: std::time::Instant::now(), + queue_wait_duration, + run_duration: Duration::ZERO, + }, + }, + result, + context.metadata_updates(), + decision, + true, + sender, + Some(completion), + ) + .await + } + + /// Ends the coordinator's peers once its attempt ended: refuses new peer + /// operations, waits for the accepted ones, then completes each peer still + /// without an outcome, as interrupted when River stopped the coordinator + /// and as failed otherwise. + pub(super) async fn finish( + self: &Arc, + inner: &Arc, + context: &WorkContext, + interrupted: bool, + ) { + self.lock().closed = true; + loop { + let idle = self.idle.notified(); + if self.lock().operations == 0 { + break; + } + idle.await; + } + let missing = { + let mut state = self.lock(); + state + .peers + .values_mut() + .filter(|peer| peer.state == PeerState::Claimed) + .map(|peer| { + peer.state = PeerState::Preparing; + peer.claimed.clone() + }) + .collect::>() + }; + if !missing.is_empty() { + let count = missing.len(); + let submissions = missing + .into_iter() + .map(|job| Submission { + job, + result: Err(if interrupted { + WorkerFailure { + error: "job interrupted by client shutdown".to_owned(), + kind: WorkerFailureKind::Interrupted, + source: None, + trace: String::new(), + } + } else { + worker_failure_from_source( + format!( + "the attempt of job {} ended without an outcome for this job", + self.coordinator + ) + .into(), + ) + }), + }) + .collect(); + if let Err(submit_error) = self.submit(inner, context, submissions).await { + error!( + job_id = self.coordinator, + peers = count, + error = %crate::error::Chain(&submit_error), + "River failed to complete peers their attempt left without an outcome" + ); + } + } + self.abandon(inner); + } + + /// Stops tracking the coordinator's peers, leaving any without a + /// persisted outcome to the rescuer. For an attempt that ends without + /// [`finish`](Self::finish). + pub(super) fn abandon(&self, inner: &ClientInner) { + let peers = { + let mut state = self.lock(); + state.closed = true; + state.released = true; + state.peers.keys().copied().collect::>() + }; + for id in peers { + inner.release_peer(id, self.id); + } + } +} + +impl ClientInner { + fn lock_peer_owners(&self) -> std::sync::MutexGuard<'_, HashMap> { + self.peer_owners + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + } + + /// Ends `ledger`'s ownership of peer job `id`, if it still owns it. + fn release_peer(&self, id: i64, ledger: u64) { + let mut owners = self.lock_peer_owners(); + if owners.get(&id) == Some(&ledger) { + owners.remove(&id); + } + } + + /// The completer's sender, while the client runs. + fn completion_sender(&self) -> Option> { + self.completion_sender + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .as_ref() + .and_then(mpsc::WeakSender::upgrade) + } +} diff --git a/rust/riverqueue/src/client/producer.rs b/rust/riverqueue/src/client/producer.rs new file mode 100644 index 000000000..48950f67e --- /dev/null +++ b/rust/riverqueue/src/client/producer.rs @@ -0,0 +1,1282 @@ +//! Queue producers that fetch and dispatch jobs. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; +use std::time::Duration; + +use chrono::Utc; +use futures_util::FutureExt as _; +#[cfg(feature = "postgres")] +use sqlx::postgres::PgRow; +#[cfg(feature = "postgres")] +use sqlx::{AssertSqlSafe, Executor, Postgres}; +use tokio::sync::{broadcast, mpsc, oneshot, watch}; +use tokio::task::JoinSet; +use tokio_util::sync::CancellationToken; +use tracing::{debug, error, warn}; + +use crate::__private::DatabaseConnection as PilotDatabaseConnection; +use crate::client::attempts::{ + FetchRegistrationGuard, register_running_attempt, signal_running_attempt, +}; +use crate::client::backoff::{RestartBackoff, exponential_backoff}; +use crate::client::completer::CompletionUpdate; +use crate::client::executor::{AbortOnDrop, execute_job}; +#[cfg(feature = "sqlite")] +use crate::client::sqlite_backend_error; +use crate::client::{ + ATTEMPTED_BY_MAX, ClientInner, DecodedJob, PARALLEL_FETCH_MINIMUM, PRODUCER_REPORT_TIMEOUT, + PRODUCER_STALE_RETENTION, QUEUE_CONFIG_POLL_INTERVAL, QUEUE_HEARTBEAT_INTERVAL, + RuntimeNotification, UndecodableJob, +}; +#[cfg(feature = "postgres")] +use crate::client::{decode_job_row, job_projection}; +use crate::database::{DatabaseKind, DatabasePool}; +use crate::pilot::{ProducerConfiguration, SharedProducer}; +use crate::{Error, Event, JobRow, Metric, QueueConfig, QueueEventKind}; + +/// Runs one producer per configured queue and reconciles them with runtime +/// queue changes. +/// +/// A producer applies a changed configuration while it runs. A removed +/// queue stops claiming at once, drains its running jobs, and shuts down its +/// extension session before its name can be added again, so a queue never +/// runs under two producers. A producer that stops unexpectedly (for example +/// after a panic) is restarted with backoff. +/// +/// `queues_ready` is sent once every queue configured at startup has created +/// or refreshed its `river_queue` row, as Go's `Client.Start` does before +/// returning, so a peer can pause or inspect those queues right away. +pub(super) async fn run_dynamic_queues( + inner: Arc, + completion_sender: mpsc::Sender, + fetch_cancel: CancellationToken, + work_cancel: CancellationToken, + notifications: broadcast::Sender, + mut changes: watch::Receiver, + queues_ready: oneshot::Sender<()>, +) -> Result<(), Error> { + // Without a notifier, running jobs learn of cancellation requests by + // polling until every producer has drained. Whether the client has one is + // known only once a PostgreSQL server's capabilities are detected, so the + // poll checks each time. + let _cancellation_poll = AbortOnDrop(tokio::spawn(poll_job_cancellations( + Arc::clone(&inner), + fetch_cancel.clone(), + ))); + let (registered_sender, mut registered) = mpsc::unbounded_channel(); + let mut producers = Producers { + active: HashMap::new(), + completion_sender, + draining: HashMap::new(), + fatal: None, + fetch_cancel: fetch_cancel.clone(), + inner, + next_generation: 0, + notifications, + registered: registered_sender, + restarts: RestartBackoff::default(), + task_queues: HashMap::new(), + tasks: JoinSet::new(), + work_cancel, + }; + producers.reconcile(); + let mut startup = Some(( + producers.active.keys().cloned().collect::>(), + queues_ready, + )); + producers.report_startup(&mut startup); + + loop { + if producers.fatal.is_some() { + break; + } + tokio::select! { + () = fetch_cancel.cancelled() => break, + change_result = changes.changed() => { + if change_result.is_err() { + break; + } + producers.reconcile(); + } + joined = producers.tasks.join_next_with_id(), if !producers.tasks.is_empty() => { + if let Some(joined) = joined { + producers.finish(joined); + } + producers.reconcile(); + } + Some(queue) = registered.recv(), if startup.is_some() => { + if let Some((pending, _)) = &mut startup { + pending.remove(&queue); + } + } + } + producers.report_startup(&mut startup); + } + + for active in producers.active.values() { + active.cancel.cancel(); + } + if producers.fatal.is_some() { + producers.work_cancel.cancel(); + } + while let Some(joined) = producers.tasks.join_next_with_id().await { + producers.finish(joined); + } + producers.inner.live_queues.send_replace(HashSet::new()); + producers.fatal.map_or(Ok(()), Err) +} + +/// Most running job IDs checked for cancellation requests in one query, +/// which also bounds SQLite's parameter count. +const JOB_CANCEL_POLL_BATCH_SIZE: usize = 1000; + +/// How long one check for cancellation requests may take. +const JOB_CANCEL_POLL_TIMEOUT: Duration = Duration::from_secs(10); + +/// Returns the fetch poll interval plus random jitter of up to a tenth of it +/// (at least 10 ms), like River Go's `jitteredFetchPollInterval`, so producers +/// that start or recover together don't poll the database in lockstep. +fn jittered_fetch_poll_interval(interval: Duration) -> Duration { + let jitter = (interval / 10).max(Duration::from_millis(10)); + crate::maintenance::random_duration(interval, interval.saturating_add(jitter)) +} + +/// Cancels this client's running attempts whose jobs have a cancellation +/// request, checking every [`QUEUE_CONFIG_POLL_INTERVAL`] while the client +/// has no notifier, like River Go's producers without one. That's checked +/// each time, since a PostgreSQL server's lack of `LISTEN`/`NOTIFY` is +/// detected only once the client runs. +/// It keeps checking after `stopping` fires, so jobs can still be cancelled +/// while producers drain; the caller aborts it once they have. +async fn poll_job_cancellations(inner: Arc, stopping: CancellationToken) { + let mut poll = tokio::time::interval(QUEUE_CONFIG_POLL_INTERVAL); + poll.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + // The first tick completes at once; running jobs start later. + poll.tick().await; + loop { + poll.tick().await; + if inner.has_notifier() { + continue; + } + let mut ids = inner + .running + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .keys() + .copied() + .collect::>(); + ids.sort_unstable(); + for batch in ids.chunks(JOB_CANCEL_POLL_BATCH_SIZE) { + let requested = tokio::time::timeout( + JOB_CANCEL_POLL_TIMEOUT, + crate::storage::job_cancel_requested(&inner, batch), + ) + .await + .map_err(|_| Error::runtime_context("job cancellation poll", "timed out")) + .and_then(|requested| requested); + match requested { + Ok(requested) => { + for id in requested { + signal_running_attempt( + &inner.running, + &inner.pending_cancellations, + &inner.fetch_registration_windows, + id, + ); + } + } + Err(poll_error) => { + if !stopping.is_cancelled() { + error!(error = %crate::error::Chain(&poll_error), "River failed to check for job cancellation requests"); + } + break; + } + } + } + } +} + +type ProducerOutcome = (String, u64, CancellationToken, Result<(), Error>); + +/// A queue's current producer generation. +struct ActiveProducer { + cancel: CancellationToken, + /// The configuration the producer applies while it runs. + config: watch::Sender, + generation: u64, + /// When the producer starts running, after any restart backoff. + started_at: tokio::time::Instant, +} + +struct Producers { + active: HashMap, + completion_sender: mpsc::Sender, + /// Producers of removed queues whose jobs are still finishing. + draining: HashMap, + /// The first producer failure that stops the client. + fatal: Option, + fetch_cancel: CancellationToken, + inner: Arc, + next_generation: u64, + notifications: broadcast::Sender, + /// Receives each queue name once its producer has registered the queue. + registered: mpsc::UnboundedSender, + restarts: RestartBackoff, + task_queues: HashMap, + tasks: JoinSet, + work_cancel: CancellationToken, +} + +impl Producers { + fn finish( + &mut self, + joined: Result<(tokio::task::Id, ProducerOutcome), tokio::task::JoinError>, + ) { + let (task_id, name, generation, failure) = match joined { + Ok((task_id, (name, generation, queue_cancel, result))) => { + let failure = match result { + // A producer returns an error only for a failure that + // stops the client, such as a broken claim protocol. + Err(queue_error) => { + self.fatal.get_or_insert(queue_error); + None + } + Ok(()) if !queue_cancel.is_cancelled() => { + Some("producer exited unexpectedly".to_owned()) + } + Ok(()) => None, + }; + (task_id, name, generation, failure) + } + Err(join_error) => { + let Some((name, generation)) = self.task_queues.get(&join_error.id()).cloned() + else { + error!(error = %join_error, "River queue producer failed"); + return; + }; + ( + join_error.id(), + name, + generation, + Some(join_error.to_string()), + ) + } + }; + self.task_queues.remove(&task_id); + if self.draining.get(&name) == Some(&generation) { + self.draining.remove(&name); + } + if self + .active + .get(&name) + .is_some_and(|active| active.generation == generation) + { + let active = self.active.remove(&name); + if let Some(failure) = failure + && !self.fetch_cancel.is_cancelled() + { + // Like the client's services, a producer that ran for a while + // before failing starts its backoff over, so isolated + // failures days apart don't escalate toward the cap. + let ran_for = active.map_or(Duration::ZERO, |active| active.started_at.elapsed()); + let (attempt, _) = self.restarts.failed(name.clone(), ran_for); + error!( + queue = %name, + error = %failure, + attempt, + "River queue producer failed; restarting it after backoff" + ); + } + } + self.publish_live(); + } + + /// Publishes the queues whose producers are running or draining, which + /// [`LocalQueues`] uses to keep a removed queue's name reserved until its + /// producer stops. + fn publish_live(&self) { + let live = self + .active + .keys() + .chain(self.draining.keys()) + .cloned() + .collect::>(); + self.inner.live_queues.send_if_modified(|current| { + if *current == live { + return false; + } + *current = live; + true + }); + } + + fn reconcile(&mut self) { + if self.fetch_cancel.is_cancelled() { + return; + } + let configured = self + .inner + .queues + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone(); + let removed = self + .active + .keys() + .filter(|name| !configured.contains_key(*name)) + .cloned() + .collect::>(); + for name in removed { + if let Some(active) = self.active.remove(&name) { + active.cancel.cancel(); + self.restarts.forget(&name); + self.draining.insert(name, active.generation); + } + } + + for (name, config) in configured { + if let Some(active) = self.active.get(&name) { + active.config.send_if_modified(|running| { + if *running == config { + return false; + } + *running = config; + true + }); + continue; + } + if self.draining.contains_key(&name) { + continue; + } + let start_delay = self + .restarts + .attempt(&name) + .map_or(Duration::ZERO, exponential_backoff); + let queue_cancel = self.fetch_cancel.child_token(); + self.next_generation = self.next_generation.wrapping_add(1); + let generation = self.next_generation; + let (config_sender, config_receiver) = watch::channel(config); + self.active.insert( + name.clone(), + ActiveProducer { + cancel: queue_cancel.clone(), + config: config_sender, + generation, + started_at: tokio::time::Instant::now() + start_delay, + }, + ); + let inner = Arc::clone(&self.inner); + let completion_sender = self.completion_sender.clone(); + let notifications = self.notifications.subscribe(); + let registered = self.registered.clone(); + let task_cancel = queue_cancel.clone(); + let task_name = name.clone(); + let work_cancel = self.work_cancel.child_token(); + let handle = self.tasks.spawn(async move { + if !start_delay.is_zero() { + tokio::select! { + () = task_cancel.cancelled() => { + return (task_name, generation, task_cancel, Ok(())); + } + () = tokio::time::sleep(start_delay) => {} + } + } + // Boxed: the producer loop's state, including an in-flight + // fetch, is too large to embed in this task's future. + let result = Box::pin(run_queue( + inner, + completion_sender, + task_name.clone(), + config_receiver, + task_cancel.clone(), + work_cancel, + notifications, + registered, + )) + .await; + (task_name, generation, task_cancel, result) + }); + self.task_queues.insert(handle.id(), (name, generation)); + } + self.publish_live(); + } + + /// Reports startup readiness once every startup queue that is still + /// configured has registered. + fn report_startup(&self, startup: &mut Option<(HashSet, oneshot::Sender<()>)>) { + let Some((pending, _)) = startup else { + return; + }; + pending + .retain(|queue| self.active.contains_key(queue) || self.draining.contains_key(queue)); + if pending.is_empty() + && let Some((_, queues_ready)) = startup.take() + { + let _ = queues_ready.send(()); + } + } +} + +/// A queue producer's started generation: its persisted record and the +/// extension's session, if any. +struct Generation { + queue: crate::Queue, + session: Option, +} + +/// The configuration an extension's session sees. +fn producer_configuration(config: &QueueConfig, queue: &crate::Queue) -> ProducerConfiguration { + ProducerConfiguration { + max_workers: config.max_workers, + metadata_text: queue.metadata_text.clone(), + queue: queue.clone(), + settings: config.extension_settings.clone(), + } +} + +/// Creates or refreshes the queue's record and starts the extension's +/// session for this generation. +async fn start_generation( + inner: &ClientInner, + queue: &str, + config: &QueueConfig, +) -> Result { + let queue_row = crate::storage::touch_queue(inner, queue).await?; + let session = inner + .pilot + .start_producer(crate::__private::ProducerStartContext { + client_id: inner.id.clone(), + configuration: producer_configuration(config, &queue_row), + database: inner.pilot_database(), + }) + .await + .map_err(|source| Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "producer", + }, + source, + })?; + Ok(Generation { + queue: queue_row, + session: session.map(SharedProducer::from), + }) +} + +/// Runs one of an extension session's synchronous callbacks, turning a panic +/// into the error that stops the client, so the producer still drains and +/// shuts the session down in order. +fn session_callback(callback_name: &str, callback: impl FnOnce()) -> Result<(), Error> { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(callback)).map_err(|panic| { + Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "producer", + }, + source: format!( + "{callback_name} panicked: {}", + crate::error::panic_message(&panic) + ) + .into(), + } + }) +} + +/// The attempts a producer has running, and the claimed rows it reports to +/// the extension's session as each attempt exits. +struct Attempts { + /// The first panic of the session's `job_finished`, which stops the + /// client. + failure: Option, + rows: HashMap, + session: Option, + tasks: JoinSet<()>, +} + +impl Attempts { + fn len(&self) -> usize { + self.tasks.len() + } + + /// Records that an attempt's task ended, however it ended. + fn exited(&mut self, joined: Result, stopping: bool) { + let task_id = match joined { + Ok(task_id) => task_id, + Err(join_error) => { + if stopping { + error!(error = %join_error, "River queue task failed during shutdown"); + } else { + error!(error = %join_error, "River queue task failed"); + } + join_error.id() + } + }; + if let Some(row) = self.rows.remove(&task_id) + && let Some(session) = &self.session + && let Err(failure) = session_callback("job_finished", || session.job_finished(&row)) + { + error!(error = %crate::error::Chain(&failure), "River extension producer callback panicked; stopping the client"); + self.failure.get_or_insert(failure); + } + } + + /// Spawns an attempt's task. `session_row` is the attempt's row when a + /// session must hear when it finishes. + fn spawn( + &mut self, + session_row: Option, + task: impl Future + Send + 'static, + ) { + let handle = self.tasks.spawn(task); + if let Some(row) = session_row { + self.rows.insert(handle.id(), row); + } + } +} + +/// Keeps a producer's queue record and extension session current until +/// `stop`, which the producer cancels only after its last attempt exits, so +/// reports continue while it drains. +/// +/// Like Go's producer, the two reports run independently, so a slow one +/// never delays the other: the queue record is refreshed after up to a +/// second of jitter and then every [`QUEUE_HEARTBEAT_INTERVAL`], and the +/// session reports after its own jitter and then every producer report +/// interval. Each report runs one at a time and is dropped after +/// [`PRODUCER_REPORT_TIMEOUT`]. +async fn run_reports( + inner: Arc, + queue: String, + session: Option, + stop: CancellationToken, +) { + let jitter = || crate::maintenance::random_duration(Duration::ZERO, Duration::from_secs(1)); + let heartbeat = async { + let mut ticks = tokio::time::interval_at( + tokio::time::Instant::now() + jitter() + QUEUE_HEARTBEAT_INTERVAL, + QUEUE_HEARTBEAT_INTERVAL, + ); + ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + ticks.tick().await; + let touched = tokio::time::timeout( + PRODUCER_REPORT_TIMEOUT, + crate::storage::touch_queue(&inner, &queue), + ) + .await; + match touched { + Ok(Ok(_)) => {} + Ok(Err(queue_error)) => { + error!(queue = %queue, error = %crate::error::Chain(&queue_error), "River queue heartbeat failed; retrying"); + } + Err(_) => { + error!(queue = %queue, timeout = ?PRODUCER_REPORT_TIMEOUT, "River queue heartbeat timed out; retrying"); + } + } + } + }; + let keep_alive = async { + let Some(session) = &session else { + return std::future::pending().await; + }; + let mut ticks = tokio::time::interval_at( + tokio::time::Instant::now() + jitter(), + inner.producer_report_interval, + ); + ticks.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + ticks.tick().await; + let stale_before = Utc::now() + - chrono::Duration::from_std(PRODUCER_STALE_RETENTION) + .unwrap_or(chrono::Duration::MAX); + let report = std::panic::AssertUnwindSafe( + session.keep_alive(crate::__private::ProducerKeepAliveContext { stale_before }), + ) + .catch_unwind(); + match tokio::time::timeout(PRODUCER_REPORT_TIMEOUT, report).await { + Ok(Ok(Ok(()))) => {} + Ok(Ok(Err(report_error))) => { + error!(queue = %queue, error = %crate::error::Chain(&*report_error), "River extension producer report failed; retrying at the next interval"); + } + Ok(Err(panic)) => { + error!(queue = %queue, panic = crate::error::panic_message(&panic), "River extension producer report panicked; retrying at the next interval"); + } + Err(_) => { + error!(queue = %queue, timeout = ?PRODUCER_REPORT_TIMEOUT, "River extension producer report timed out; retrying at the next interval"); + } + } + } + }; + tokio::select! { + () = stop.cancelled() => {} + () = heartbeat => {} + () = keep_alive => {} + } +} + +/// Shuts an extension's session down after its producer stopped, like Go's +/// `finalizeShutdown`: up to four attempts, one at a time, with deadlines of +/// 100 milliseconds growing fivefold. +async fn shut_down_session(queue: &str, session: &dyn crate::__private::PilotProducer) { + const ATTEMPTS: u32 = 4; + const BASE_TIMEOUT: Duration = Duration::from_millis(100); + + let mut timeout = BASE_TIMEOUT; + for attempt in 1..=ATTEMPTS { + let context = crate::__private::ProducerShutdownContext { attempt, timeout }; + let shutdown = std::panic::AssertUnwindSafe(session.shutdown(context)).catch_unwind(); + match tokio::time::timeout(timeout, shutdown).await { + Ok(Ok(Ok(()))) => return, + Ok(Ok(Err(shutdown_error))) => { + error!(queue = %queue, attempt, ?timeout, error = %crate::error::Chain(&*shutdown_error), "River extension producer shutdown failed"); + } + Ok(Err(panic)) => { + error!(queue = %queue, attempt, ?timeout, panic = crate::error::panic_message(&panic), "River extension producer shutdown panicked"); + } + Err(_) => { + error!(queue = %queue, attempt, ?timeout, "River extension producer shutdown timed out"); + } + } + timeout *= 5; + } + warn!(queue = %queue, "River extension producer failed to shut down cleanly after all attempts"); +} + +#[expect( + clippy::too_many_arguments, + clippy::too_many_lines, + reason = "a producer's claim loop, reports, and drain share one queue's state" +)] +pub(super) async fn run_queue( + inner: Arc, + completion_sender: mpsc::Sender, + queue: String, + mut config_changes: watch::Receiver, + fetch_cancel: CancellationToken, + work_cancel: CancellationToken, + mut notifications: broadcast::Receiver, + registered: mpsc::UnboundedSender, +) -> Result<(), Error> { + // Short write contention (common on SQLite) clears quickly. Longer + // outages back off like River's other services; the producer keeps trying + // for as long as the client runs rather than stopping the client. + const START_FAST_RETRY_INTERVAL: Duration = Duration::from_millis(10); + const START_FAST_RETRY_WINDOW: Duration = Duration::from_secs(10); + + let mut config = config_changes.borrow_and_update().clone(); + let start_time = tokio::time::Instant::now(); + let mut start_attempt = 0; + let Generation { + queue: mut queue_row, + session, + } = loop { + let Some(started) = + unless_cancelled(&fetch_cancel, start_generation(&inner, &queue, &config)).await + else { + return Ok(()); + }; + match started { + Ok(generation) => break generation, + Err(queue_error) => { + let sleep = if start_time.elapsed() < START_FAST_RETRY_WINDOW { + debug!(error = %crate::error::Chain(&queue_error), "River queue startup failed; retrying"); + START_FAST_RETRY_INTERVAL + } else { + start_attempt += 1; + let sleep = exponential_backoff(start_attempt); + error!( + queue = %queue, + error = %crate::error::Chain(&queue_error), + sleep_duration = ?sleep, + "River queue startup failed (will retry after backoff)" + ); + sleep + }; + tokio::select! { + () = fetch_cancel.cancelled() => return Ok(()), + () = tokio::time::sleep(sleep) => {} + } + } + } + }; + let _ = registered.send(queue.clone()); + // Reports outlive claiming: they stop only once the last attempt exits. + let reports_stop = CancellationToken::new(); + let reports = AbortOnDrop(tokio::spawn(run_reports( + Arc::clone(&inner), + queue.clone(), + session.clone(), + reports_stop.clone(), + ))); + let mut paused = queue_row.paused_at.is_some(); + let claims_through_session = session + .as_ref() + .is_some_and(|session| session.intercepts_claim()); + let mut attempts = Attempts { + failure: None, + rows: HashMap::new(), + session: session.clone(), + tasks: JoinSet::new(), + }; + // `None` until the first fetch, which needs no cooldown. Subtracting the + // cooldown from now instead would panic for a cooldown longer than the + // monotonic clock's age, as on a freshly booted macOS host. + let mut last_fetch: Option = None; + // The first poll fires at once, so a producer claims as soon as it + // starts; later ones are jittered. + let poll = tokio::time::sleep(Duration::ZERO); + tokio::pin!(poll); + let mut queue_config_poll = tokio::time::interval(QUEUE_CONFIG_POLL_INTERVAL); + queue_config_poll.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + let outcome = loop { + if let Some(failure) = attempts.failure.take() { + break Err(failure); + } + let (mut should_fetch, refresh_queue_state) = tokio::select! { + () = fetch_cancel.cancelled() => break Ok(()), + changed = config_changes.changed() => { + if changed.is_err() { + break Ok(()); + } + let updated = config_changes.borrow_and_update().clone(); + if updated.fetch_poll_interval != config.fetch_poll_interval { + poll.as_mut().reset( + tokio::time::Instant::now() + + jittered_fetch_poll_interval(updated.fetch_poll_interval), + ); + } + config = updated; + if let Some(session) = &session { + let configuration = producer_configuration(&config, &queue_row); + if let Err(failure) = session_callback("configuration_changed", || { + session.configuration_changed(&configuration); + }) { + break Err(failure); + } + } + // More capacity may allow a claim now. + (true, false) + }, + // Like Go's `pollForSettingChanges`, only a client without a + // notifier polls queue state; one with a notifier learns of + // changes from control notifications, and a listener reconnect + // broadcasts a refresh of every queue. Whether a PostgreSQL + // server supports notifications is detected only once the client + // runs, so this is checked on each tick. + _ = queue_config_poll.tick() => (false, !inner.has_notifier()), + () = &mut poll => { + poll.as_mut().reset( + tokio::time::Instant::now() + + jittered_fetch_poll_interval(config.fetch_poll_interval), + ); + (true, false) + }, + Some(joined) = attempts.tasks.join_next_with_id(), if !attempts.tasks.is_empty() => { + attempts.exited(joined.map(|(task_id, ())| task_id), false); + (true, false) + }, + notification = notifications.recv() => match notification { + Ok(RuntimeNotification::Insert(notification_queue)) => ( + notification_queue == "*" || notification_queue == queue, + false, + ), + Ok(RuntimeNotification::QueueControl(notification_queue)) => ( + false, + notification_queue == "*" || notification_queue == queue, + ), + Err(broadcast::error::RecvError::Closed) => (false, false), + Err(broadcast::error::RecvError::Lagged(_)) => (true, true), + }, + }; + + if refresh_queue_state { + let Some(loaded) = + unless_cancelled(&fetch_cancel, crate::storage::load_queue(&inner, &queue)).await + else { + break Ok(()); + }; + match loaded { + Ok(Some(loaded)) => { + // The stored text covers changes the parsed metadata + // hides, such as key order or a number's spelling. + let changed = loaded.metadata_text != queue_row.metadata_text + || loaded.metadata != queue_row.metadata + || loaded.paused_at.is_some() != queue_row.paused_at.is_some(); + queue_row = loaded; + if changed && let Some(session) = &session { + let configuration = producer_configuration(&config, &queue_row); + if let Err(failure) = session_callback("configuration_changed", || { + session.configuration_changed(&configuration); + }) { + break Err(failure); + } + } + let next_paused = queue_row.paused_at.is_some(); + if next_paused != paused { + paused = next_paused; + let event_kind = if paused { + QueueEventKind::Paused + } else { + QueueEventKind::Resumed + }; + let _ = inner + .events + .send(Event::queue(event_kind, queue_row.clone())); + should_fetch |= !paused; + } + } + Ok(None) => {} + Err(queue_error) => { + error!(error = %crate::error::Chain(&queue_error), "River queue state refresh failed; retrying"); + continue; + } + } + } + + if !should_fetch || paused { + continue; + } + if let Some(remaining) = last_fetch.and_then(|last_fetch| { + config + .resolved_fetch_cooldown(inner.fetch_cooldown) + .checked_sub(last_fetch.elapsed()) + }) { + tokio::select! { + () = fetch_cancel.cancelled() => break Ok(()), + () = tokio::time::sleep(remaining) => {} + } + } + // A stop can be requested while another branch above was selected or + // during the cooldown. Go's fetch query fails once its context is + // cancelled, so no jobs are claimed after a stop; match that. + if fetch_cancel.is_cancelled() { + break Ok(()); + } + // Lowering `max_workers` stops claims until enough running jobs + // finish; it never cancels them. + let available = config.max_workers.saturating_sub(attempts.len()); + if available == 0 { + continue; + } + let registration_guard = FetchRegistrationGuard::new(&inner); + let fetched = match (&session, claims_through_session) { + (Some(session), true) => { + match claim_through_session( + &inner, + session.as_ref(), + &queue, + available, + &fetch_cancel, + ) + .await + { + Ok(fetched) => Ok(fetched), + Err(SessionClaimError::Claim(claim_error)) => Err(claim_error), + Err(SessionClaimError::Protocol(protocol_error)) => break Err(protocol_error), + } + } + // Boxed: two concurrent PostgreSQL claims make a large future. + _ => Box::pin(fetch_available(&inner, &queue, available, &fetch_cancel)).await, + }; + last_fetch = Some(tokio::time::Instant::now()); + let FetchedJobs { claimed } = match fetched { + Ok(fetched) => fetched, + Err(fetch_error) => { + error!(error = %crate::error::Chain(&fetch_error), "River job fetch failed; retrying"); + continue; + } + }; + // Like River Go, a claimed job whose row couldn't be fully decoded + // gets an executor that fails its attempt with the decode error + // instead of working it, so it's retried or discarded rather than + // left running. + let claimed = claimed.into_iter().filter_map(|decoded| match decoded { + Ok(row) => Some((row, None)), + Err(UndecodableJob { + error, row: None, .. + }) => { + error!(%error, "claimed River job row couldn't be identified; leaving it for the rescuer"); + None + } + Err(UndecodableJob { + error, + row: Some(row), + .. + }) => Some((*row, Some(error))), + }); + for (row, decode_error) in claimed { + let hard_cancel = work_cancel.child_token(); + let cancellation = hard_cancel.child_token(); + register_running_attempt( + &inner.running, + &inner.pending_cancellations, + row.id, + &cancellation, + ); + let task_inner = Arc::clone(&inner); + let completion_sender = completion_sender.clone(); + let session_row = attempts.session.is_some().then(|| row.clone()); + attempts.spawn(session_row, async move { + execute_job( + task_inner, + row, + decode_error, + hard_cancel, + cancellation, + completion_sender, + ) + .await; + }); + } + drop(registration_guard); + while let Some(joined) = attempts.tasks.try_join_next_with_id() { + attempts.exited(joined.map(|(task_id, ())| task_id), false); + } + }; + + if outcome.is_err() { + // A protocol failure stops the client: cancel this queue's attempts + // like a hard stop, then wait for them. + work_cancel.cancel(); + } + while let Some(joined) = attempts.tasks.join_next_with_id().await { + attempts.exited(joined.map(|(task_id, ())| task_id), true); + } + let outcome = match (outcome, attempts.failure.take()) { + (Ok(()), Some(failure)) => Err(failure), + (outcome, _) => outcome, + }; + reports_stop.cancel(); + let mut reports = reports; + if let Err(join_error) = (&mut reports.0).await { + error!(queue = %queue, error = %join_error, "River producer reports failed"); + } + if let Some(session) = &session { + shut_down_session(&queue, session.as_ref()).await; + } + outcome +} + +/// Why a claim through an extension's session produced no jobs. +enum SessionClaimError { + /// The session reported an error; River tries again later. + Claim(Error), + /// The session returned committed rows River can't accept. + Protocol(Error), +} + +/// Claims through an extension's session and checks what it returned. +async fn claim_through_session( + inner: &ClientInner, + session: &dyn crate::__private::PilotProducer, + queue: &str, + limit: usize, + claim_stop: &CancellationToken, +) -> Result { + let fetch_started = (!inner.hooks.is_empty()).then(std::time::Instant::now); + let database = inner.pilot_database(); + let claimed = session.claim( + crate::__private::ProducerClaimContext { + client_id: &inner.id, + claim_stop, + database: &database, + kinds: inner.fetch_kinds.as_deref(), + limit, + queue, + }, + crate::__private::ProducerClaimNext::new(inner, queue, limit), + ); + let claimed = std::panic::AssertUnwindSafe(claimed) + .catch_unwind() + .await + .map_err(|panic| { + SessionClaimError::Protocol(Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "fetch claim", + }, + source: format!("claim panicked: {}", crate::error::panic_message(&panic)).into(), + }) + })? + .map_err(|source| { + SessionClaimError::Claim(Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "fetch claim", + }, + source, + }) + })?; + if let Err(violation) = crate::pilot::validate_claimed(&claimed, &inner.id, queue, limit) { + error!(queue = %queue, error = %violation, "River extension claim broke the claim protocol; stopping the client"); + return Err(SessionClaimError::Protocol(Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "fetch claim", + }, + source: violation.into(), + })); + } + let rows = claimed + .into_iter() + .map(crate::__private::ClaimedJob::into_decoded) + .collect(); + Ok(finish_fetch(inner, fetch_started, rows).await) +} + +/// Claims available jobs with River's own statements, splitting a large +/// PostgreSQL claim in two. +async fn fetch_available( + inner: &ClientInner, + queue: &str, + available: usize, + cancel: &CancellationToken, +) -> Result { + let use_parallel_fetch = match inner.database.kind() { + #[cfg(feature = "postgres")] + DatabaseKind::Postgres => true, + #[cfg(feature = "sqlite")] + DatabaseKind::Sqlite => false, + }; + if !use_parallel_fetch || available < PARALLEL_FETCH_MINIMUM { + return fetch_jobs(inner, queue, available, cancel).await; + } + let first_maximum = available / 2; + let second_maximum = available - first_maximum; + let (first, second) = tokio::join!( + fetch_jobs(inner, queue, first_maximum, cancel), + fetch_jobs(inner, queue, second_maximum, cancel), + ); + match (first, second) { + (Ok(mut first), Ok(second)) => { + first.extend(second); + Ok(first) + } + (Ok(rows), Err(fetch_error)) | (Err(fetch_error), Ok(rows)) => { + error!( + error = %crate::error::Chain(&fetch_error), + "one parallel River job fetch failed; working the successfully fetched jobs" + ); + Ok(rows) + } + (Err(fetch_error), Err(second_fetch_error)) => { + error!( + secondary_error = %crate::error::Chain(&second_fetch_error), + "the other parallel River job fetch failed too" + ); + Err(fetch_error) + } + } +} + +/// Wraps claimed rows and emits fetch metrics. +async fn finish_fetch( + inner: &ClientInner, + fetch_started: Option, + rows: Vec, +) -> FetchedJobs { + let fetched = FetchedJobs::from_decoded(rows); + if let Some(fetch_started) = fetch_started { + for metric in [ + Metric::JobGetAvailableDuration(fetch_started.elapsed()), + Metric::JobGetAvailableCount(u64::try_from(fetched.len()).unwrap_or(u64::MAX)), + ] { + for hook in &inner.hooks { + if let Err(hook_error) = hook.metric_emit(metric).await { + error!(error = %crate::error::Chain(&hook_error), "River metric hook failed"); + } + } + } + } + fetched +} + +/// Waits for `operation` unless `cancel` fires first. +/// +/// A fetch only abandons connection acquisition and transaction begins, +/// which River's begin helpers make safe to drop. Like Go's fetch, which +/// runs under the fetch context, a stop then doesn't wait out the pool's +/// acquire timeout during a database outage. Nothing is claimed until the +/// claim itself runs, and that always completes. +async fn unless_cancelled( + cancel: &CancellationToken, + operation: impl std::future::Future, +) -> Option { + tokio::select! { + biased; + () = cancel.cancelled() => None, + output = operation => Some(output), + } +} + +/// Runs River's standard claim of up to `limit` available jobs from `queue` +/// on `connection`, the claim a fetch makes without an extension. +pub(crate) async fn standard_claim( + inner: &ClientInner, + connection: PilotDatabaseConnection<'_>, + queue: &str, + limit: usize, +) -> Result, Error> { + let limit = i32::try_from(limit) + .map_err(|_| Error::runtime_context("job fetch", "fetch maximum exceeds i32"))?; + match connection { + #[cfg(feature = "postgres")] + PilotDatabaseConnection::Postgres(connection) => Ok(fetch_oss_records( + connection, + standard_claim_sql(inner), + queue, + limit, + &inner.id, + inner.fetch_kinds.as_deref(), + ) + .await? + .iter() + .map(decode_job_row) + .collect()), + #[cfg(feature = "sqlite")] + PilotDatabaseConnection::Sqlite(connection) => { + let params = crate::database::sqlite::ClaimJobs { + client_id: &inner.id, + kinds: inner.fetch_kinds.as_deref(), + limit, + max_attempted_by: ATTEMPTED_BY_MAX, + now: Utc::now(), + queue, + }; + crate::database::sqlite::claim(connection, ¶ms) + .await + .map_err(sqlite_backend_error) + } + } +} + +/// River's PostgreSQL claim statement for this client's schema. +#[cfg(feature = "postgres")] +fn standard_claim_sql(inner: &ClientInner) -> String { + let table = inner.schema.qualify("river_job"); + let queue_table = inner.schema.qualify("river_queue"); + format!( + "WITH locked AS (\ + SELECT id FROM {table} WHERE state = 'available' AND queue = $1 AND scheduled_at <= now() \ + AND ($5::text[] IS NULL OR kind = any($5::text[])) \ + AND NOT EXISTS (SELECT 1 FROM {queue_table} WHERE name = $1 AND paused_at IS NOT NULL) \ + ORDER BY priority, scheduled_at, id LIMIT $2 FOR UPDATE SKIP LOCKED\ + ) UPDATE {table} AS job \ + SET state = 'running', attempt = job.attempt + 1, attempted_at = now(), \ + attempted_by = array_append(\ + CASE WHEN array_length(job.attempted_by, 1) >= $4 \ + THEN job.attempted_by[array_length(job.attempted_by, 1) + 2 - $4:] \ + ELSE job.attempted_by END, $3) \ + FROM locked WHERE job.id = locked.id \ + RETURNING {}, false AS unique_skipped_as_duplicate", + job_projection("job") + ) +} + +/// Claims up to `maximum` jobs from `queue` with River's own statement on a +/// pooled connection. Returns no jobs when `cancel` fires before a +/// connection is available. +pub(super) async fn fetch_jobs( + inner: &ClientInner, + queue: &str, + maximum: usize, + cancel: &CancellationToken, +) -> Result { + let fetch_started = (!inner.hooks.is_empty()).then(std::time::Instant::now); + let rows = match inner.database.pool() { + #[cfg(feature = "postgres")] + DatabasePool::Postgres(pool) => { + let Some(connection) = unless_cancelled(cancel, pool.acquire()).await else { + return Ok(FetchedJobs::default()); + }; + let mut connection = connection?; + standard_claim( + inner, + PilotDatabaseConnection::Postgres(&mut connection), + queue, + maximum, + ) + .await? + } + #[cfg(feature = "sqlite")] + DatabasePool::Sqlite(pool) => { + let Some(connection) = unless_cancelled(cancel, pool.acquire()).await else { + return Ok(FetchedJobs::default()); + }; + let mut connection = connection?; + standard_claim( + inner, + PilotDatabaseConnection::Sqlite(&mut connection), + queue, + maximum, + ) + .await? + } + }; + Ok(finish_fetch(inner, fetch_started, rows).await) +} + +#[cfg(feature = "postgres")] +pub(super) async fn fetch_oss_records<'executor, E>( + executor: E, + sql: String, + queue: &str, + maximum: i32, + client_id: &str, + kinds: Option<&[String]>, +) -> Result, sqlx::Error> +where + E: Executor<'executor, Database = Postgres>, +{ + sqlx::query(AssertSqlSafe(sql)) + .bind(queue) + .bind(maximum) + .bind(client_id) + .bind(ATTEMPTED_BY_MAX) + .bind(kinds) + .fetch_all(executor) + .await +} + +/// Jobs claimed by one fetch, in the order they were claimed. Claims commit +/// before rows are decoded, so a row that can't be fully decoded is kept with +/// its decode error to have its attempt failed, instead of failing the whole +/// fetch and stranding every claimed job. +#[derive(Default)] +pub(super) struct FetchedJobs { + pub(super) claimed: Vec, +} + +impl FetchedJobs { + pub(super) const fn from_decoded(claimed: Vec) -> Self { + Self { claimed } + } + + pub(super) fn extend(&mut self, other: Self) { + self.claimed.extend(other.claimed); + } + + pub(super) const fn len(&self) -> usize { + self.claimed.len() + } +} + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use super::jittered_fetch_poll_interval; + + #[test] + fn fetch_poll_interval_jitter_is_a_tenth_and_at_least_ten_milliseconds() { + for _ in 0..100 { + let poll = jittered_fetch_poll_interval(Duration::from_secs(1)); + assert!( + (Duration::from_secs(1)..Duration::from_millis(1_100)).contains(&poll), + "{poll:?}" + ); + let poll = jittered_fetch_poll_interval(Duration::from_millis(20)); + assert!( + (Duration::from_millis(20)..Duration::from_millis(30)).contains(&poll), + "{poll:?}" + ); + } + } +} diff --git a/rust/riverqueue/src/client/queues.rs b/rust/riverqueue/src/client/queues.rs new file mode 100644 index 000000000..16dc4e217 --- /dev/null +++ b/rust/riverqueue/src/client/queues.rs @@ -0,0 +1,265 @@ +//! Operations on persisted queue records. + +use crate::client::request::{Target, request_type}; +use crate::storage::Access; +use crate::{Client, Error}; +use crate::{Queue, QueueListParams, QueueSelector, QueueUpdateParams}; + +/// Operations on persisted queue records, returned by [`Client::queues`]. +/// +/// A queue record exists for every queue a client has worked, and is shared +/// by every client of the database, including River clients in other +/// languages. Pausing a queue stops every client from fetching its jobs. To +/// change which queues this client works, use [`Client::local_queues`]. +/// +/// Each method returns a request that runs on the client's own pool when +/// awaited, or in a caller-managed transaction after `.tx(&mut tx)`: +/// +/// ```no_run +/// # use riverqueue::QueueSelector; +/// # #[cfg(feature = "postgres")] +/// # async fn example(client: riverqueue::Client, pool: sqlx::PgPool) -> Result<(), riverqueue::Error> { +/// client.queues().pause("email").await?; +/// +/// let mut tx = riverqueue::database::begin_postgres(&pool).await?; +/// client.queues().resume(QueueSelector::All).tx(&mut tx).await?; +/// tx.commit().await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// Requests don't run until awaited. Dropping one before it completes rolls +/// back River's own transaction; with `.tx`, the caller's transaction may +/// contain the operation's partial effects and should be rolled back. +#[derive(Clone, Copy, Debug)] +pub struct Queues<'a> { + client: &'a Client, +} + +impl Client { + /// Returns operations on persisted queue records: getting, listing, + /// pausing, resuming, and updating them. + #[must_use] + pub const fn queues(&self) -> Queues<'_> { + Queues { client: self } + } +} + +impl<'a> Queues<'a> { + /// Gets a queue record by name. + /// + /// # Errors + /// + /// Returns [`Error::NotFound`] when the queue has no record, + /// [`Error::DatabaseMismatch`] for a transaction from another backend, and + /// [`Error::Database`] when the database operation fails. + pub fn get(&self, name: impl Into) -> QueueGetRequest<'a> { + QueueGetRequest { + client: self.client, + name: name.into(), + target: Target::Client, + } + } + + /// Lists queue records in name order. + /// + /// # Errors + /// + /// Returns [`Error::InvalidJob`] for a limit outside one through 10,000, + /// [`Error::DatabaseMismatch`] for a transaction from another backend, and + /// [`Error::Database`] when the database operation fails. + pub fn list(&self, params: QueueListParams) -> QueueListRequest<'a> { + QueueListRequest { + client: self.client, + params, + target: Target::Client, + } + } + + /// Pauses one queue, or every queue with [`QueueSelector::All`]. + /// + /// Clients stop fetching jobs from a paused queue, while jobs already + /// running finish normally. Clients learn of the pause through a + /// notification when it commits, or on their next poll of the queue's + /// record when they run without notifications. Pausing a paused queue + /// changes nothing. + /// + /// # Errors + /// + /// Returns [`Error::NotFound`] when a named queue has no record, + /// [`Error::DatabaseMismatch`] for a transaction from another backend, and + /// [`Error::Database`] when the database operation fails. + pub fn pause(&self, selector: impl Into) -> QueuePauseRequest<'a> { + QueuePauseRequest { + client: self.client, + selector: selector.into(), + target: Target::Client, + } + } + + /// Resumes one paused queue, or every queue with [`QueueSelector::All`]. + /// + /// Clients learn of the change as they do for [`pause`](Self::pause). + /// Resuming a queue that isn't paused changes nothing. + /// + /// # Errors + /// + /// Returns [`Error::NotFound`] when a named queue has no record, + /// [`Error::DatabaseMismatch`] for a transaction from another backend, and + /// [`Error::Database`] when the database operation fails. + pub fn resume(&self, selector: impl Into) -> QueueResumeRequest<'a> { + QueueResumeRequest { + client: self.client, + selector: selector.into(), + target: Target::Client, + } + } + + /// Updates a queue record and returns it. + /// + /// New metadata is sent to the clients working the queue when the update + /// commits. + /// + /// ```no_run + /// # use riverqueue::QueueUpdateParams; + /// # async fn example(client: riverqueue::Client) -> Result<(), riverqueue::Error> { + /// let mut metadata = serde_json::Map::new(); + /// metadata.insert("owner".to_owned(), "billing".into()); + /// let queue = client + /// .queues() + /// .update("invoices", QueueUpdateParams::new().metadata(metadata)) + /// .await?; + /// assert_eq!(queue.metadata["owner"], "billing"); + /// # Ok(()) + /// # } + /// ``` + /// + /// # Errors + /// + /// Returns [`Error::NotFound`] when the queue has no record, + /// [`Error::DatabaseMismatch`] for a transaction from another backend, and + /// [`Error::Database`] when the database operation fails. + pub fn update( + &self, + name: impl Into, + params: QueueUpdateParams, + ) -> QueueUpdateRequest<'a> { + QueueUpdateRequest { + client: self.client, + name: name.into(), + params, + target: Target::Client, + } + } +} + +request_type! { + /// A queue lookup, returned by [`Queues::get`]. Await it to get the + /// queue record. + read QueueGetRequest { name: String } -> Queue +} + +impl QueueGetRequest<'_> { + async fn run(self) -> Result { + let inner = &self.client.inner; + let mut session = self.target.session(inner, Access::Autocommit).await?; + session + .storage(inner) + .queue_get(&self.name) + .await? + .ok_or_else(|| Error::NotFound(crate::Record::Queue(self.name.clone()))) + } +} + +request_type! { + /// A queue listing, returned by [`Queues::list`]. Await it to get queue + /// records in name order. + read QueueListRequest { params: QueueListParams } -> Vec +} + +impl QueueListRequest<'_> { + async fn run(self) -> Result, Error> { + let inner = &self.client.inner; + let mut session = self.target.session(inner, Access::Autocommit).await?; + session.storage(inner).queue_list(&self.params).await + } +} + +request_type! { + /// A queue pause, returned by [`Queues::pause`]. Await it to pause the + /// selected queues. + write QueuePauseRequest { selector: QueueSelector } -> () +} + +impl QueuePauseRequest<'_> { + async fn run(self) -> Result<(), Error> { + set_paused(self.client, self.target, &self.selector, true).await + } +} + +request_type! { + /// A queue resumption, returned by [`Queues::resume`]. Await it to + /// resume the selected queues. + write QueueResumeRequest { selector: QueueSelector } -> () +} + +impl QueueResumeRequest<'_> { + async fn run(self) -> Result<(), Error> { + set_paused(self.client, self.target, &self.selector, false).await + } +} + +async fn set_paused( + client: &Client, + target: Target<'_>, + selector: &QueueSelector, + paused: bool, +) -> Result<(), Error> { + let Some(name) = selector.protocol_name() else { + return Err(Error::NotFound(crate::Record::Queue( + crate::storage::QUEUE_ALL.to_owned(), + ))); + }; + let inner = &client.inner; + let own_transaction = !target.is_transaction(); + let mut session = target.session(inner, Access::Transaction).await?; + session + .storage(inner) + .queue_set_paused(name, paused) + .await?; + session.commit().await?; + // Wake this client's producers at once rather than at their next + // notification or poll. A caller's transaction may still roll back, so + // it relies on the committed notification alone. + if own_transaction { + client.signal_queue_control(name); + } + Ok(()) +} + +request_type! { + /// A queue update, returned by [`Queues::update`]. Await it to update the + /// queue and get its new record. + write QueueUpdateRequest { name: String, params: QueueUpdateParams } -> Queue +} + +impl QueueUpdateRequest<'_> { + async fn run(self) -> Result { + let inner = &self.client.inner; + let own_transaction = !self.target.is_transaction(); + let metadata = self.params.metadata; + let mut session = self.target.session(inner, Access::Transaction).await?; + let queue = session + .storage(inner) + .queue_update(&self.name, metadata.as_ref()) + .await?; + session.commit().await?; + // Like a pause, a metadata change reaches this client's producers at + // once, including on a poll-only client, as Go's + // `notifyProducerWithoutListenerQueueControlEvent` does. + if own_transaction && metadata.is_some() { + self.client.signal_queue_control(&self.name); + } + Ok(queue) + } +} diff --git a/rust/riverqueue/src/client/record.rs b/rust/riverqueue/src/client/record.rs new file mode 100644 index 000000000..7d616a5fc --- /dev/null +++ b/rust/riverqueue/src/client/record.rs @@ -0,0 +1,287 @@ +//! Decoding of persisted job rows. +//! +//! River Go decodes rows with `encoding/json` and native integers, so a row +//! written by Go (or edited by an operator) can contain values that a strict +//! Rust decoder rejects, like SQLite integers outside `i16`, and on SQLite, +//! JSON columns can be changed to any shape. +//! +//! Like River Go, reads that return a job to a caller decode strictly, while +//! the runtime decodes each row it has claimed, completed, or found stuck on +//! its own and tolerates fields that can't be decoded. Such a row keeps the +//! fields that could be decoded, leaves the others empty, and carries the +//! decode error, so one bad row can't fail or strand the rows read with it. + +use std::fmt::Display; + +#[cfg(feature = "postgres")] +use chrono::{DateTime, Utc}; +#[cfg(feature = "postgres")] +use serde_json::value::RawValue; +#[cfg(feature = "postgres")] +use sqlx::{FromRow, Row, postgres::PgRow, types::Json}; +use tracing::error; + +use crate::JobRow; +#[cfg(feature = "postgres")] +use crate::{AttemptError, Error, JobMetadata, JobState}; + +/// A row with fields that River couldn't decode. +#[derive(Debug)] +pub(crate) struct UndecodableJob { + /// Why the row couldn't be decoded, with a line for each field that + /// couldn't be, like River Go's joined decode errors. + pub(crate) error: String, + /// The columns that couldn't be decoded, when the row could be + /// identified. + pub(crate) columns: Vec, + /// The row with every field that could be decoded and the others left + /// empty, or `None` when not even the columns that identify the job could + /// be. + pub(crate) row: Option>, +} + +/// A row decoded on its own, with any undecodable fields reported +/// separately. +pub(crate) type DecodedJob = Result; + +/// Returns a row for a runtime path that tolerates undecodable fields, like +/// River Go's set-state and stuck-job reads. Only a row that can't be +/// identified at all is dropped, with a log. +pub(crate) fn tolerant_row(decoded: DecodedJob) -> Option { + match decoded { + Ok(row) => Some(row), + Err(UndecodableJob { + error, row: None, .. + }) => { + error!(%error, "River job row couldn't be identified; skipping it"); + None + } + Err(UndecodableJob { row, .. }) => row.map(|row| *row), + } +} + +/// Collects why fields of one row couldn't be decoded. +#[derive(Default)] +pub(crate) struct FieldErrors { + columns: Vec, + messages: Vec, +} + +impl FieldErrors { + /// Returns a decoded field, or records why it couldn't be decoded and + /// leaves it empty. + pub(crate) fn field( + &mut self, + column: &str, + decoded: Result, + ) -> T { + self.field_or_else(column, decoded, T::default) + } + + /// Returns a decoded field, or records why it couldn't be decoded and + /// leaves it as `empty` returns. + pub(crate) fn field_or_else( + &mut self, + column: &str, + decoded: Result, + empty: impl FnOnce() -> T, + ) -> T { + decoded.unwrap_or_else(|error| { + self.columns.push(column.to_owned()); + self.messages + .push(format!("error unmarshaling `{column}`: {error}")); + empty() + }) + } + + /// Finishes decoding `row`, reporting it as undecodable if any of its + /// fields couldn't be decoded. + pub(crate) fn finish(self, row: JobRow) -> DecodedJob { + if self.messages.is_empty() { + Ok(row) + } else { + Err(UndecodableJob { + columns: self.columns, + error: self.messages.join("\n"), + row: Some(Box::new(row)), + }) + } + } +} + +/// Converts a persisted integer to `i16`, saturating at the type bounds. +/// +/// River Go stores `attempt`, `max_attempts`, and `priority` as native +/// integers on SQLite. Values beyond `i16` are only reachable through +/// `max_attempts` in practice; saturating keeps such a job workable with +/// identical retry decisions until its 32,767th attempt. +pub(crate) fn saturating_i16(value: i64) -> i16 { + i16::try_from(value).unwrap_or(if value < 0 { i16::MIN } else { i16::MAX }) +} + +/// A PostgreSQL job row. Columns the database constrains decode strictly, while +/// those that can hold values River can't represent are kept as their decode +/// results. +#[cfg(feature = "postgres")] +pub(crate) struct JobRecord { + attempt: i16, + attempted_at: Option>, + attempted_by: Result>, sqlx::Error>, + created_at: DateTime, + encoded_args: Json>, + errors: Result>>>, sqlx::Error>, + finalized_at: Option>, + id: i64, + kind: String, + max_attempts: i16, + metadata: Json>, + priority: i16, + queue: String, + scheduled_at: DateTime, + state: String, + tags: Result, sqlx::Error>, + unique_key: Option>, + pub(super) unique_skipped_as_duplicate: bool, + unique_states: Option, +} + +#[cfg(feature = "postgres")] +impl<'row> FromRow<'row, PgRow> for JobRecord { + fn from_row(row: &'row PgRow) -> Result { + // `job_projection` fixes the first 18 columns in this order, and every + // JobRecord query appends the insert-only duplicate flag at index 18. + // Positional decoding avoids repeated column-name lookups on hot fetch + // and completion paths. + Ok(Self { + attempt: row.try_get(1)?, + attempted_at: row.try_get(2)?, + attempted_by: row.try_get(3), + created_at: row.try_get(4)?, + encoded_args: row.try_get(5)?, + errors: row.try_get(6), + finalized_at: row.try_get(7)?, + id: row.try_get(0)?, + kind: row.try_get(8)?, + max_attempts: row.try_get(9)?, + metadata: row.try_get(10)?, + priority: row.try_get(11)?, + queue: row.try_get(12)?, + scheduled_at: row.try_get(13)?, + state: row.try_get(14)?, + tags: row.try_get(15), + unique_key: row.try_get(16)?, + unique_skipped_as_duplicate: row.try_get(18)?, + unique_states: row.try_get(17)?, + }) + } +} + +#[cfg(feature = "postgres")] +impl JobRecord { + /// Decodes the row, failing if any field can't be decoded. + pub(crate) fn into_job_row(self) -> Result { + let id = self.id; + self.decode() + .map_err(|job| Error::invalid_job(format!("job {id}: {}", job.error))) + } + + /// Decodes the row, keeping the fields that can be decoded when others + /// can't. River Go reads the metadata as raw JSON, but a [`JobRow`] can + /// only represent an object. + pub(crate) fn decode(self) -> DecodedJob { + let state = JobState::try_from(self.state.as_str()).map_err(|error| UndecodableJob { + columns: Vec::new(), + error: format!("job {}: {error}", self.id), + row: None, + })?; + let mut errors = FieldErrors::default(); + let attempted_by = errors.field("attempted_by", self.attempted_by); + let attempt_errors = errors.field( + "errors", + self.errors + .map_err(|error| error.to_string()) + .and_then(|errors| { + errors + .iter() + .map(|error| match error { + Some(error) => AttemptError::from_json_lenient(error.0.get()) + .map_err(|error| error.to_string()), + None => Err("unexpected SQL NULL element".to_owned()), + }) + .collect() + }), + ); + let metadata = errors.field("metadata", JobMetadata::try_from(self.metadata.0)); + let tags = errors.field("tags", self.tags); + let unique_states = errors.field( + "unique_states", + self.unique_states + .map(|bits| { + u8::from_str_radix(&bits, 2).map(|bitmask| { + JobState::ALL + .iter() + .copied() + .filter(|state| bitmask & state.unique_bit() != 0) + .collect() + }) + }) + .transpose(), + ); + errors.finish(JobRow { + attempt: self.attempt, + attempted_at: self.attempted_at, + attempted_by: attempted_by.unwrap_or_default(), + created_at: self.created_at, + encoded_args: self.encoded_args.0, + errors: attempt_errors, + finalized_at: self.finalized_at, + id: self.id, + kind: self.kind, + max_attempts: self.max_attempts, + metadata, + priority: self.priority, + queue: self.queue, + scheduled_at: self.scheduled_at, + state, + tags, + unique_key: self.unique_key, + unique_states, + }) + } +} + +/// Decodes a row selected with [`job_projection`] on its own. +#[cfg(feature = "postgres")] +pub(crate) fn decode_job_row(row: &PgRow) -> DecodedJob { + JobRecord::from_row(row) + .map_err(|error| UndecodableJob { + columns: Vec::new(), + error: error.to_string(), + row: None, + })? + .decode() +} + +#[cfg(feature = "postgres")] +pub(crate) fn job_projection(alias: &str) -> String { + format!( + "{alias}.id, {alias}.attempt, {alias}.attempted_at, {alias}.attempted_by, \ + {alias}.created_at, {alias}.args AS encoded_args, \ + coalesce({alias}.errors, '{{}}'::jsonb[]) AS errors, \ + {alias}.finalized_at, {alias}.kind, {alias}.max_attempts, {alias}.metadata, \ + {alias}.priority, {alias}.queue, {alias}.scheduled_at, {alias}.state::text AS state, \ + {alias}.tags::text[] AS tags, {alias}.unique_key, {alias}.unique_states::text AS unique_states" + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn integers_saturate_to_i16() { + assert_eq!(saturating_i16(25), 25); + assert_eq!(saturating_i16(40_000), i16::MAX); + assert_eq!(saturating_i16(-40_000), i16::MIN); + } +} diff --git a/rust/riverqueue/src/client/request.rs b/rust/riverqueue/src/client/request.rs new file mode 100644 index 000000000..8b13b5e46 --- /dev/null +++ b/rust/riverqueue/src/client/request.rs @@ -0,0 +1,170 @@ +//! Plumbing shared by the client's request builders. + +use std::fmt; + +use crate::__private::DatabaseConnection as PilotDatabaseConnection; +use crate::client::ClientInner; +use crate::database::DatabaseTransactionExecutor; +use crate::storage::{Access, Session}; +use crate::{Client, Error}; + +/// Where a request runs. +pub(super) enum Target<'a> { + /// The client's own pool, in a transaction River commits when the + /// operation needs one. + Client, + /// A caller-managed transaction. A transaction from another backend is + /// reported when the request is awaited. + Transaction(Result, Error>), +} + +impl<'a> Target<'a> { + pub(super) fn transaction(client: &Client, executor: E) -> Self + where + E: DatabaseTransactionExecutor<'a>, + { + Self::Transaction(client.inner.transaction_connection(executor)) + } + + /// Returns the caller's transaction connection, or `None` to use the + /// client's own pool. + pub(super) fn into_executor(self) -> Result>, Error> { + match self { + Self::Client => Ok(None), + Self::Transaction(connection) => connection.map(Some), + } + } +} + +impl<'a> Target<'a> { + /// Opens the session the request runs in: the caller's transaction, or + /// a connection from the client's own pool with the given access. + pub(super) async fn session( + self, + inner: &ClientInner, + access: Access, + ) -> Result, Error> { + match self { + Self::Client => Session::begin(&inner.database, access).await, + Self::Transaction(connection) => Ok(Session::caller(connection?)), + } + } + + pub(super) const fn is_transaction(&self) -> bool { + matches!(self, Self::Transaction(_)) + } +} + +impl fmt::Debug for Target<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(if self.is_transaction() { + "Transaction" + } else { + "Client" + }) + } +} + +/// Defines a request builder that runs on the client's pool or, with `tx`, +/// in a caller-managed transaction, and that runs when awaited. +/// +/// The request type must implement `async fn run(self) -> Result`. +/// +/// Prefix the name with `read` for a request that only reads, or `write` for +/// one that changes data; each gets the matching cancel safety section. +macro_rules! request_type { + ( + $(#[$attr:meta])* + read $name:ident { $($field:ident: $type:ty),* $(,)? } -> $output:ty + ) => { + $crate::client::request::request_type! { + @define + $(#[$attr])* + /// + /// # Cancel safety + /// + /// The request only reads, so dropping it before it finishes + /// changes nothing. + $name { $($field: $type),* } -> $output + } + }; + ( + $(#[$attr:meta])* + write $name:ident { $($field:ident: $type:ty),* $(,)? } -> $output:ty + ) => { + $crate::client::request::request_type! { + @define + $(#[$attr])* + /// + /// # Cancel safety + /// + /// On the client's own pool, River runs the request as one + /// statement or one transaction, so dropping it before it + /// finishes never leaves it partly applied. A request dropped + /// while its commit is in flight may still have taken effect. + /// With [`tx`](Self::tx), what it already ran stays in the + /// caller's transaction; roll that transaction back rather than + /// committing it. + $name { $($field: $type),* } -> $output + } + }; + ( + @define + $(#[$attr:meta])* + $name:ident { $($field:ident: $type:ty),* $(,)? } -> $output:ty + ) => { + $(#[$attr])* + #[must_use = "requests do nothing unless awaited"] + pub struct $name<'a> { + client: &'a Client, + $($field: $type,)* + target: Target<'a>, + } + + impl<'a> $name<'a> { + /// Runs the request in a caller-managed transaction instead of + /// on the client's own pool. + /// + /// The request sees the transaction's uncommitted changes, and + /// its own changes and notifications take effect only when the + /// caller commits. `executor` must be a SQLx transaction for the + /// client's database backend, begun with + /// [`begin_postgres`](crate::database::begin_postgres) or + /// [`begin_sqlite_write`](crate::database::begin_sqlite_write). + pub fn tx<'t, E>(self, executor: E) -> $name<'t> + where + 'a: 't, + E: crate::database::DatabaseTransactionExecutor<'t>, + { + $name { + client: self.client, + $($field: self.$field,)* + target: Target::transaction(self.client, executor), + } + } + } + + impl std::fmt::Debug for $name<'_> { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct(stringify!($name)) + $(.field(stringify!($field), &self.$field))* + .field("target", &self.target) + .finish_non_exhaustive() + } + } + + impl<'a> std::future::IntoFuture for $name<'a> { + type Output = Result<$output, Error>; + type IntoFuture = + std::pin::Pin + Send + 'a>>; + + fn into_future(self) -> Self::IntoFuture { + Box::pin(self.run()) + } + } + }; +} + +pub(super) use request_type; diff --git a/rust/riverqueue/src/client/resign.rs b/rust/riverqueue/src/client/resign.rs new file mode 100644 index 000000000..f2cd63027 --- /dev/null +++ b/rust/riverqueue/src/client/resign.rs @@ -0,0 +1,66 @@ +//! Leadership resignation requests. + +use crate::client::request::{Target, request_type}; +use crate::maintenance::LeadershipWakeup; +use crate::storage::Access; +use crate::{Client, Error}; + +impl Client { + /// Asks the current leader to resign, so that clients elect a leader + /// again. + /// + /// The request is a notification delivered to every client, which usually + /// makes the leader resign, but has no effect when no leader is elected. + /// With [`tx`](ResignRequest::tx), the notification is sent only when the + /// transaction commits. + /// + /// ```no_run + /// # async fn example(client: riverqueue::Client) -> Result<(), riverqueue::Error> { + /// client.request_resign().await?; + /// # Ok(()) + /// # } + /// ``` + /// + /// # Errors + /// + /// Returns [`Error::DatabaseMismatch`] for a transaction from another + /// backend and [`Error::Database`] when the database operation fails. + pub fn request_resign(&self) -> ResignRequest<'_> { + ResignRequest { + client: self, + target: Target::Client, + } + } +} + +request_type! { + /// A leadership resignation request, returned by + /// [`Client::request_resign`]. Await it to send the request. + /// + /// A client without notifications, poll-only or using a PostgreSQL + /// server without `LISTEN`/`NOTIFY`, hears its own request directly once + /// the request commits, and no other client hears it. If the request is + /// dropped while that commit is in flight, the request may commit without + /// being heard, and such a leader then keeps its term. + write ResignRequest {} -> () +} + +impl ResignRequest<'_> { + async fn run(self) -> Result<(), Error> { + let inner = &self.client.inner; + let own_transaction = !self.target.is_transaction(); + let mut session = self.target.session(inner, Access::Transaction).await?; + session.storage(inner).leader_request_resign().await?; + session.commit().await?; + // A client without a notifier, poll-only or on a server without + // `LISTEN`/`NOTIFY`, learns of its own request directly. Any other + // client receives the committed notification like every other client + // does; also signalling it locally would deliver the request twice. + if own_transaction && !inner.has_notifier() { + let _ = inner + .leadership_wakeups + .send(LeadershipWakeup::RequestResign); + } + Ok(()) + } +} diff --git a/rust/riverqueue/src/client/run.rs b/rust/riverqueue/src/client/run.rs new file mode 100644 index 000000000..eb26532d8 --- /dev/null +++ b/rust/riverqueue/src/client/run.rs @@ -0,0 +1,755 @@ +//! Starting clients and observing their lifecycle. + +use std::collections::HashMap; +use std::sync::atomic::Ordering; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use tokio::sync::{mpsc, oneshot}; +use tokio::task::JoinSet; +use tokio_util::sync::CancellationToken; +use tracing::{debug, error, warn}; + +use crate::client::ClientInner; +use crate::client::backoff::RestartBackoff; +use crate::client::completer::run_completion_batcher; +#[cfg(feature = "postgres")] +use crate::client::notifier::run_notifications; +#[cfg(feature = "sqlite")] +use crate::client::notifier::run_sqlite_notifications; +use crate::client::notifier::{ReadySender, ReadySlot}; +use crate::client::producer::run_dynamic_queues; +use crate::database::DatabasePool; +use crate::{Client, Error}; + +/// A boxed application shutdown signal awaited by a started client. +type ShutdownSignal = std::pin::Pin + Send>>; + +impl Client { + /// Starts configured queues and returns a lifecycle handle. + /// + /// The client supervises its services: a notification listener, SQLite + /// outbox poller, maintenance, or extension service that fails is logged + /// and restarted with backoff, and producers keep polling meanwhile, so a + /// database outage never stops the client. Only a failure of the + /// producers or the completer, which would leave jobs unworked or + /// unpersisted, stops the client: work is then cancelled, every worker is + /// awaited, and [`RunHandle::wait`] returns the error. + /// + /// With notifications enabled on PostgreSQL, the client opens one + /// dedicated listener connection with the pool's connect options. It is + /// not taken from, and does not count against, the pool's + /// `max_connections`. + /// + /// A client can run once at a time. After it stops, it can be started + /// again. + /// + /// # Errors + /// + /// Returns [`Error::RuntimeUnavailable`] when called outside a Tokio + /// runtime, a configuration error when the client has no queues, and + /// [`Error::AlreadyRunning`] when the client is already running. + pub fn start(&self) -> Result { + self.start_inner(None) + } + + /// Starts the client like [`Client::start`] and stops it gracefully once + /// `signal` completes. + /// + /// Completing `signal` has the same effect as [`Stopper::stop`]: the + /// client stops fetching jobs and lets running jobs finish, and the + /// builder's `soft_stop_timeout` escalates to cancelling them when set. + /// Hard stops remain available through [`RunHandle::stopper`] and + /// [`RunHandle::shutdown_now`]. The client drops `signal` without + /// awaiting it further once it stops for any other reason. + /// + /// This mirrors the graceful shutdown hooks of Tokio servers such as + /// axum's `with_graceful_shutdown`, so one application signal can stop an + /// HTTP server and River together. An application + /// [`CancellationToken`] works as a signal through + /// [`CancellationToken::cancelled_owned`]. + /// + /// # Examples + /// + /// ```no_run + /// # async fn example(client: riverqueue::Client) -> Result<(), riverqueue::Error> { + /// let mut run = client.start_with_graceful_shutdown(async { + /// let _ = tokio::signal::ctrl_c().await; + /// })?; + /// // Returns after Ctrl-C once in-flight jobs have finished. + /// run.wait().await + /// # } + /// ``` + /// + /// # Errors + /// + /// Returns the same errors as [`Client::start`]. + pub fn start_with_graceful_shutdown(&self, signal: F) -> Result + where + F: std::future::Future + Send + 'static, + { + self.start_inner(Some(Box::pin(signal))) + } + + fn start_inner(&self, shutdown_signal: Option) -> Result { + let runtime = + tokio::runtime::Handle::try_current().map_err(|_| Error::RuntimeUnavailable { + operation: "starting a client", + })?; + if self + .inner + .queues + .read() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_empty() + { + return Err(Error::configuration( + "at least one queue is required to start a client".to_owned(), + )); + } + if self + .inner + .started + .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire) + .is_err() + { + return Err(Error::AlreadyRunning); + } + let stopper = Stopper { + fetch_cancel: CancellationToken::new(), + work_cancel: CancellationToken::new(), + }; + let inner = Arc::clone(&self.inner); + let (ready_sender, ready) = oneshot::channel(); + let supervisor = Supervisor { + fetch_cancel: stopper.fetch_cancel.clone(), + inner: Arc::clone(&inner), + notifier_ready: Arc::new(Mutex::new(None)), + restarts: RestartBackoff::default(), + services: HashMap::new(), + stop: StopTokens::new(&stopper.fetch_cancel), + tasks: JoinSet::new(), + work_cancel: stopper.work_cancel.clone(), + }; + let join = runtime.spawn(async move { + let result = supervisor.run(ready_sender, shutdown_signal).await; + inner.started.store(false, Ordering::Release); + result + }); + Ok(RunHandle { + join: Some(join), + ready: Readiness::Pending(ready), + stopper, + }) + } +} + +/// A long-running service owned by a started client. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +enum Service { + Completer, + Maintenance, + Notifier, + Queues, + Extension(usize), +} + +impl Service { + /// Services whose failure leaves jobs unworked or unpersisted. Every other + /// service is restarted after a failure. + const fn is_essential(self) -> bool { + matches!(self, Self::Completer | Self::Queues) + } + + const fn name(self) -> &'static str { + match self { + Self::Completer => "completer", + Self::Maintenance => "maintenance", + Self::Notifier => "notifier", + Self::Queues => "producers", + Self::Extension(_) => "extension runtime service", + } + } +} + +/// When each part of a started client stops, matching River Go's order. +/// +/// A stop request cancels the client's fetch token. Claims, leadership with +/// its maintenance services, the notifier, and extension runtime services +/// all stop then, like Go's services started on the fetch context. Running +/// jobs keep going until they finish or `work_cancel` fires, and each +/// producer keeps reporting to its extension session until its last attempt +/// has left, from a token of its own that the stop request doesn't cancel. +struct StopTokens { + /// New claims and queue changes. + claims: CancellationToken, + /// Leader election and leader-owned maintenance. + leadership: CancellationToken, + /// The notifier and extension runtime services. + services: CancellationToken, +} + +impl StopTokens { + fn new(fetch_cancel: &CancellationToken) -> Self { + Self { + claims: fetch_cancel.child_token(), + leadership: fetch_cancel.child_token(), + services: fetch_cancel.child_token(), + } + } +} + +/// Runs a started client's services and restarts the ones that fail. +struct Supervisor { + fetch_cancel: CancellationToken, + inner: Arc, + notifier_ready: ReadySlot, + restarts: RestartBackoff, + /// Each running service and when its current run started. + services: HashMap, + stop: StopTokens, + tasks: JoinSet>, + work_cancel: CancellationToken, +} + +impl Supervisor { + async fn run( + mut self, + ready: ReadySender, + shutdown_signal: Option, + ) -> Result<(), Error> { + self.start_services(ready); + let inner = Arc::clone(&self.inner); + let stop_watch = watch_stop( + self.fetch_cancel.clone(), + self.work_cancel.clone(), + inner.soft_stop_timeout, + shutdown_signal, + ); + tokio::pin!(stop_watch); + let mut stop_watch_done = false; + let mut fatal = None; + loop { + let joined = tokio::select! { + joined = self.tasks.join_next_with_id() => joined, + () = &mut stop_watch, if !stop_watch_done => { + stop_watch_done = true; + continue; + } + }; + let Some(joined) = joined else { + break; + }; + let (task_id, outcome) = match joined { + Ok((task_id, outcome)) => (task_id, outcome), + Err(join_error) => (join_error.id(), Err(Error::from_join(join_error))), + }; + let Some((service, started_at)) = self.services.remove(&task_id) else { + continue; + }; + let stopping = self.fetch_cancel.is_cancelled(); + match outcome { + // The completer ends once every producer has dropped its sender. + Ok(()) if stopping || service == Service::Completer => {} + Err(service_error) if stopping && !service.is_essential() => { + debug!( + service = service.name(), + error = %crate::error::Chain(&service_error), + "River service stopped with an error during shutdown" + ); + } + outcome if service.is_essential() => { + let service_error = outcome.err().unwrap_or_else(|| { + Error::runtime_context(service.name(), "exited unexpectedly".to_owned()) + }); + error!( + service = service.name(), + error = %crate::error::Chain(&service_error), + "River service failed; stopping the client after in-flight work" + ); + fatal.get_or_insert(service_error); + self.fetch_cancel.cancel(); + self.work_cancel.cancel(); + } + outcome => { + let (attempt, delay) = self.restarts.failed(service, started_at.elapsed()); + error!( + service = service.name(), + attempt, + error = %outcome.err().map_or_else(|| "exited unexpectedly".to_owned(), |error| error.to_string()), + sleep_duration = ?delay, + "River service failed; restarting after backoff" + ); + self.spawn_service(service, delay); + } + } + } + fatal.map_or(Ok(()), Err) + } + + /// Starts the client's services and reports readiness to `ready` once + /// its queues are registered and its notification path is active. + fn start_services(&mut self, ready: ReadySender) { + let inner = Arc::clone(&self.inner); + let (completion_sender, completion_receiver) = mpsc::channel(10_000); + *inner + .completion_sender + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = + Some(completion_sender.downgrade()); + // Like the extension design, runtime services start before + // producers. This orders their start, not their readiness. + for index in 0..inner.pilot.runtime_services().len() { + self.spawn_service(Service::Extension(index), Duration::ZERO); + } + let (queues_ready_sender, queues_ready) = oneshot::channel(); + self.spawn_task( + Service::Queues, + run_dynamic_queues( + Arc::clone(&inner), + completion_sender, + self.stop.claims.clone(), + self.work_cancel.child_token(), + inner.queue_notifications.clone(), + inner.queue_changes.subscribe(), + queues_ready_sender, + ), + ); + self.spawn_task( + Service::Completer, + run_completion_batcher(Arc::clone(&inner), completion_receiver), + ); + let (notifier_ready_sender, notifier_ready) = oneshot::channel(); + if inner.poll_only { + let _ = notifier_ready_sender.send(()); + } else { + *self + .notifier_ready + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) = Some(notifier_ready_sender); + self.spawn_service(Service::Notifier, Duration::ZERO); + } + // Not a service: it ends once both parts report, or once either + // part stops first, which leaves the client not ready. The + // supervisor keeps the notifier's readiness across restarts, so the + // client's stop ends the wait. + let stopped = self.fetch_cancel.clone(); + self.tasks.spawn(async move { + let notifier_ready = tokio::select! { + () = stopped.cancelled() => return Ok(()), + notifier_ready = notifier_ready => notifier_ready, + }; + if notifier_ready.is_ok() && queues_ready.await.is_ok() { + let _ = ready.send(()); + } + Ok(()) + }); + if !inner.leader_election_disabled { + self.spawn_service(Service::Maintenance, Duration::ZERO); + } + } + + fn spawn_task(&mut self, service: Service, task: F) + where + F: std::future::Future> + Send + 'static, + { + let handle = self.tasks.spawn(task); + self.services + .insert(handle.id(), (service, tokio::time::Instant::now())); + } + + /// Starts a restartable service after `delay`, unless the client stops + /// first. + fn spawn_service(&mut self, service: Service, delay: Duration) { + let inner = Arc::clone(&self.inner); + let cancel; + let run: std::pin::Pin> + Send>> = + match service { + Service::Maintenance => { + cancel = self.stop.leadership.child_token(); + Box::pin(crate::maintenance::run_maintenance( + Arc::clone(&inner), + cancel.clone(), + inner.leadership_wakeups.subscribe(), + )) + } + Service::Notifier => { + cancel = self.stop.services.child_token(); + let ready = self.notifier_ready.clone(); + #[cfg(test)] + if inner + .notifier_start_panics + .try_update(Ordering::AcqRel, Ordering::Acquire, |remaining| { + remaining.checked_sub(1) + }) + .is_ok() + { + self.spawn_task(service, async move { + let _ready = ready; + panic!("injected notification listener panic") + }); + return; + } + match inner.database.pool() { + #[cfg(feature = "postgres")] + DatabasePool::Postgres(pool) => Box::pin(run_notifications( + Arc::clone(&inner), + pool.clone(), + cancel.clone(), + inner.queue_notifications.clone(), + ready, + )), + #[cfg(feature = "sqlite")] + DatabasePool::Sqlite(pool) => Box::pin(run_sqlite_notifications( + Arc::clone(&inner), + pool.clone(), + cancel.clone(), + inner.queue_notifications.clone(), + ready, + )), + } + } + Service::Extension(index) => { + cancel = self.stop.services.child_token(); + let Some(runtime_service) = + inner.pilot.runtime_services().into_iter().nth(index) + else { + return; + }; + let context = crate::__private::RuntimeServiceContext { + cancellation: cancel.clone(), + client: crate::client::WeakClient::new(&inner), + database: inner.pilot_database(), + }; + Box::pin(async move { + runtime_service.run(context).await.map_err(|service_error| { + Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "runtime service", + }, + source: service_error, + } + }) + }) + } + Service::Completer | Service::Queues => { + unreachable!("essential services are started once") + } + }; + self.spawn_task(service, async move { + if !delay.is_zero() { + tokio::select! { + () = cancel.cancelled() => return Ok(()), + () = tokio::time::sleep(delay) => {} + } + } + run.await + }); + } +} + +/// Requests a soft stop when `shutdown_signal` completes, then escalates any +/// soft stop to a hard stop after `soft_stop_timeout`. +/// +/// The escalation belongs to the client rather than to a caller awaiting +/// [`RunHandle::shutdown`], so it applies however the stop was requested and +/// dropping a shutdown future never changes it. This matches Go's client, +/// which starts its soft stop timer when fetching stops. +async fn watch_stop( + fetch_cancel: CancellationToken, + work_cancel: CancellationToken, + soft_stop_timeout: Option, + shutdown_signal: Option, +) { + match shutdown_signal { + Some(signal) => tokio::select! { + () = fetch_cancel.cancelled() => {} + () = signal => { + tracing::info!("River client received its shutdown signal; stopping gracefully"); + fetch_cancel.cancel(); + } + }, + None => fetch_cancel.cancelled().await, + } + let Some(timeout) = soft_stop_timeout else { + return; + }; + tokio::select! { + () = work_cancel.cancelled() => {} + () = tokio::time::sleep(timeout) => { + warn!( + soft_stop_timeout = ?timeout, + "River client soft stop timed out; cancelling remaining jobs" + ); + work_cancel.cancel(); + } + } +} + +/// Requests that a running client stop. +/// +/// A stopper is a cheap, cloneable trigger obtained from +/// [`RunHandle::stopper`]. It lets any task, such as a signal handler, stop a +/// client while another task owns the [`RunHandle`] and awaits +/// [`RunHandle::wait`]. Its methods only request a stop and return +/// immediately; observe completion through the handle. +/// +/// Requests are idempotent and ordered by severity: calling [`Stopper::stop`] +/// after [`Stopper::stop_now`] does not undo the hard stop, and requests made +/// after the client stopped do nothing. A stopper only affects the run it came +/// from, not a later restart of the same [`Client`]. +/// +/// # Examples +/// +/// ```no_run +/// # async fn example(client: riverqueue::Client) -> Result<(), riverqueue::Error> { +/// let mut run = client.start()?; +/// let stopper = run.stopper(); +/// tokio::spawn(async move { +/// let _ = tokio::signal::ctrl_c().await; +/// // Stop fetching and let running jobs finish. +/// stopper.stop(); +/// let _ = tokio::signal::ctrl_c().await; +/// // A second Ctrl-C cancels jobs that are still running. +/// stopper.stop_now(); +/// }); +/// run.wait().await +/// # } +/// ``` +#[derive(Clone, Debug)] +pub struct Stopper { + fetch_cancel: CancellationToken, + work_cancel: CancellationToken, +} + +impl Stopper { + /// Requests a soft stop. + /// + /// The client stops fetching new jobs and lets running jobs finish before + /// it stops. Leader election, maintenance, and the notification listener + /// stop at once, while each queue's producer keeps reporting its + /// running jobs until they finish. When the builder's `soft_stop_timeout` + /// is set, jobs still running after that timeout are cancelled as if by + /// [`Stopper::stop_now`]. + pub fn stop(&self) { + self.fetch_cancel.cancel(); + } + + /// Requests a hard stop. + /// + /// The client stops fetching new jobs and cancels the + /// [`WorkContext::cancellation_token`](crate::WorkContext::cancellation_token) + /// of every running job, then waits for workers to return. A worker that + /// returns [`WorkCancelled`](crate::WorkCancelled) (anywhere in its + /// error's source chain) promptly after cancellation has its job made + /// available again without using up its attempt. Any other error is + /// recorded and consumes the attempt like an ordinary failure, and + /// `Ok` completes the job. A worker that ignores cancellation for longer + /// than the job stuck threshold is aborted, which also fails its attempt: + /// the attempt counts and the job is retried or discarded like any failed + /// attempt. A job whose cancellation was requested with + /// [`Jobs::cancel`](crate::Jobs::cancel) is cancelled rather than made + /// available. + pub fn stop_now(&self) { + self.fetch_cancel.cancel(); + self.work_cancel.cancel(); + } +} + +/// Controls one running client instance. +/// +/// [`RunHandle::wait`], [`RunHandle::shutdown`], and +/// [`RunHandle::shutdown_now`] take `&mut self`, can be called repeatedly, and +/// are cancel safe: dropping one of their futures, for example from +/// `tokio::time::timeout` or `tokio::select!`, leaves the client and the +/// handle as they were, apart from any stop the method already requested. To +/// stop the client from another task, obtain a [`Stopper`] with +/// [`RunHandle::stopper`] or start the client with +/// [`Client::start_with_graceful_shutdown`]. +/// +/// The client's result is reported to the first call that observes it +/// stopping; later calls return `Ok(())`. +/// +/// Dropping the handle requests a hard stop, like [`Stopper::stop_now`], but +/// cannot wait for in-flight work to be recorded. Use [`RunHandle::shutdown`] +/// or [`RunHandle::shutdown_now`] when shutdown must finish before returning, +/// or [`RunHandle::detach`] to deliberately leave the client running. +/// +/// # Examples +/// +/// Stop gracefully, but cancel jobs that are still running after 30 seconds: +/// +/// ```no_run +/// # async fn example(client: riverqueue::Client) -> Result<(), riverqueue::Error> { +/// use std::time::Duration; +/// +/// let mut run = client.start()?; +/// // ... serve until the application stops ... +/// if tokio::time::timeout(Duration::from_secs(30), run.shutdown()) +/// .await +/// .is_err() +/// { +/// run.shutdown_now().await?; +/// } +/// # Ok(()) +/// # } +/// ``` +/// +/// [`Client`] also offers `soft_stop_timeout` on its builder, which performs +/// this escalation inside the client however the stop was requested. +#[must_use = "dropping the handle requests immediate client shutdown; call detach to run it independently"] +pub struct RunHandle { + join: Option>>, + ready: Readiness, + stopper: Stopper, +} + +/// Whether the client's notification path has become ready. +#[derive(Debug)] +enum Readiness { + Pending(oneshot::Receiver<()>), + Ready, + Stopped, +} + +impl std::fmt::Debug for RunHandle { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("RunHandle") + .field("running", &self.join.is_some()) + .field("ready", &matches!(self.ready, Readiness::Ready)) + .finish_non_exhaustive() + } +} + +impl RunHandle { + /// Leaves the client running independently of this handle. + /// + /// This relinquishes waiting for the client: it runs until a [`Stopper`] + /// obtained earlier from [`RunHandle::stopper`] or the signal passed to + /// [`Client::start_with_graceful_shutdown`] stops it, an essential service + /// fails, or the process exits. Nothing observes its result, and jobs + /// running when the process exits are left `running` for the rescuer. + /// Most applications should keep the handle and await + /// [`RunHandle::shutdown`] instead. + pub fn detach(mut self) { + // Without a join handle, dropping the handle requests no stop. + self.join.take(); + } + + /// Requests a soft stop and waits for the client to stop. + /// + /// This is [`Stopper::stop`] followed by [`RunHandle::wait`]. The stop is + /// requested when the future is first polled. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Dropping the future after its first poll + /// leaves the soft stop in progress, including any `soft_stop_timeout` + /// escalation, and never escalates to a hard stop by itself. The handle + /// remains usable: call [`RunHandle::shutdown_now`] to cancel running jobs + /// or [`RunHandle::wait`] to keep waiting. + /// + /// # Errors + /// + /// Returns the error that stopped the client, as [`RunHandle::wait`] does. + pub async fn shutdown(&mut self) -> Result<(), Error> { + self.stopper.stop(); + self.wait().await + } + + /// Requests a hard stop and waits for the client to stop. + /// + /// This is [`Stopper::stop_now`] followed by [`RunHandle::wait`]. The stop + /// is requested when the future is first polled. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Dropping the future after its first poll + /// leaves the hard stop in progress, and the handle remains usable. + /// + /// # Errors + /// + /// Returns the error that stopped the client, as [`RunHandle::wait`] does. + pub async fn shutdown_now(&mut self) -> Result<(), Error> { + self.stopper.stop_now(); + self.wait().await + } + + /// Returns a [`Stopper`] that can stop this client from any task. + pub fn stopper(&self) -> Stopper { + self.stopper.clone() + } + + /// Waits for the client to stop, without requesting a stop. + /// + /// Returns immediately when the client has already stopped and its + /// result was reported by an earlier call, or when the handle was never + /// attached to a running client. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Dropping the future leaves the client + /// running and the handle usable; a later call keeps waiting. + /// + /// # Errors + /// + /// Returns the error from an essential service, such as producers or the + /// completer, whose failure stopped the client, or + /// [`Error::RuntimeTask`] when the client's supervisor task panicked or was + /// cancelled by its runtime shutting down. + pub async fn wait(&mut self) -> Result<(), Error> { + let Some(join) = self.join.as_mut() else { + return Ok(()); + }; + let result = join.await; + self.join = None; + join_client_result(result) + } + + /// Waits until the client is ready: every queue configured at start has + /// created or refreshed its `river_queue` row, and the selected backend's + /// notification path is active. + /// + /// Once it returns, other clients can rely on the client's queues, for + /// example to pause one of them. Poll-only clients skip + /// the notification path. Once readiness is observed, later calls return + /// the same result immediately. + /// + /// # Cancel safety + /// + /// This method is cancel safe. Dropping the future leaves the client + /// running, and a later call keeps waiting. + /// + /// # Errors + /// + /// Returns [`Error::ClientStopped`] when the client stops before becoming + /// ready. + pub async fn wait_ready(&mut self) -> Result<(), Error> { + if let Readiness::Pending(receiver) = &mut self.ready { + self.ready = match receiver.await { + Ok(()) => Readiness::Ready, + Err(_) => Readiness::Stopped, + }; + } + match &self.ready { + Readiness::Stopped => Err(Error::ClientStopped), + Readiness::Pending(_) | Readiness::Ready => Ok(()), + } + } +} + +impl Drop for RunHandle { + fn drop(&mut self) { + if self.join.is_some() { + self.stopper.stop_now(); + } + } +} + +pub(super) fn join_client_result( + result: Result, tokio::task::JoinError>, +) -> Result<(), Error> { + result.map_err(Error::from_join)??; + Ok(()) +} diff --git a/rust/riverqueue/src/client/tests.rs b/rust/riverqueue/src/client/tests.rs new file mode 100644 index 000000000..929143872 --- /dev/null +++ b/rust/riverqueue/src/client/tests.rs @@ -0,0 +1,721 @@ +#[cfg(feature = "sqlite")] +use serde::Deserialize; +use serde_json::Map; + +use super::attempts::{register_running_attempt, remove_running_attempt, signal_running_attempt}; +#[cfg(feature = "postgres")] +use super::completer::CompletionBatcher; +#[cfg(feature = "sqlite")] +use super::completer::{ + COMPLETION_BACKLOG_LIMIT, COMPLETION_BATCH_SIZE, CompletionTiming, run_completion_batcher, +}; +use super::completer::{persisted_completion_event_kind, with_completion_retries}; +use super::executor::scheduled_after; +#[cfg(feature = "sqlite")] +use super::notifier::dispatch_notification; +use super::*; +use crate::{AttemptError, JobEventKind, JobRow, JobState, WorkError, WorkResult}; +#[cfg(feature = "sqlite")] +use crate::{InsertOpts, InsertParams, WorkContext, WorkOutcome}; +#[cfg(feature = "sqlite")] +use crate::{Job, JobArgs}; + +#[test] +fn completion_events_follow_persisted_state() { + let cases = [ + ( + JobState::Available, + JobEventKind::Failed, + JobEventKind::Failed, + ), + ( + JobState::Available, + JobEventKind::Interrupted, + JobEventKind::Interrupted, + ), + ( + JobState::Available, + JobEventKind::Cancelled, + JobEventKind::Failed, + ), + ( + JobState::Available, + JobEventKind::Completed, + JobEventKind::Failed, + ), + ( + JobState::Available, + JobEventKind::Snoozed, + JobEventKind::Snoozed, + ), + ( + JobState::Cancelled, + JobEventKind::Failed, + JobEventKind::Cancelled, + ), + ( + JobState::Completed, + JobEventKind::Failed, + JobEventKind::Completed, + ), + ( + JobState::Discarded, + JobEventKind::Completed, + JobEventKind::Failed, + ), + ( + JobState::Retryable, + JobEventKind::Completed, + JobEventKind::Failed, + ), + ( + JobState::Scheduled, + JobEventKind::Failed, + JobEventKind::Snoozed, + ), + ]; + + for (state, requested, expected) in cases { + assert_eq!( + persisted_completion_event_kind(state, requested), + Some(expected) + ); + } + // A row moved back to a non-final state by someone else reports nothing + // rather than failing the completer. + for state in [JobState::Pending, JobState::Running] { + assert_eq!( + persisted_completion_event_kind(state, JobEventKind::Completed), + None + ); + } +} + +#[test] +fn completion_cleanup_preserves_newer_attempt() { + let job_id = 42; + let first = CancellationToken::new(); + let second = CancellationToken::new(); + let running = Mutex::new(HashMap::from([(job_id, first.clone())])); + + running + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .insert(job_id, second.clone()); + + remove_running_attempt(&running, job_id, &first); + assert_eq!( + running + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get(&job_id), + Some(&second) + ); + + remove_running_attempt(&running, job_id, &second); + assert!( + running + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .get(&job_id) + .is_none() + ); +} + +#[test] +fn pending_cancellation_reaches_fetched_attempt() { + let job_id = 42; + let cancellation = CancellationToken::new(); + let fetch_registration_windows = AtomicU64::new(1); + let pending_cancellations = Mutex::new(HashMap::new()); + let running = Mutex::new(HashMap::new()); + + signal_running_attempt( + &running, + &pending_cancellations, + &fetch_registration_windows, + job_id, + ); + register_running_attempt(&running, &pending_cancellations, job_id, &cancellation); + + assert!(cancellation.is_cancelled()); + assert!( + pending_cancellations + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_empty() + ); +} + +#[test] +fn unmatched_cancellation_is_not_retained_without_fetch() { + let fetch_registration_windows = AtomicU64::new(0); + let pending_cancellations = Mutex::new(HashMap::new()); + let running = Mutex::new(HashMap::new()); + + signal_running_attempt( + &running, + &pending_cancellations, + &fetch_registration_windows, + 42, + ); + + assert!( + pending_cancellations + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .is_empty() + ); +} + +fn retry_row(error_count: usize) -> JobRow { + let now = DateTime::parse_from_rfc3339("2026-01-02T03:04:05Z") + .unwrap() + .with_timezone(&Utc); + JobRow { + attempt: i16::try_from(error_count).unwrap_or(i16::MAX), + attempted_at: Some(now), + attempted_by: vec!["test".to_owned()], + created_at: now, + encoded_args: serde_json::value::to_raw_value(&serde_json::json!({})).unwrap(), + errors: vec![ + AttemptError { + at: now, + attempt: 1, + error: "failed".to_owned(), + trace: String::new(), + }; + error_count + ], + finalized_at: None, + id: 42, + kind: "retry_test".to_owned(), + max_attempts: 1_000, + metadata: Map::new().into(), + priority: 1, + queue: "default".to_owned(), + scheduled_at: now, + state: JobState::Retryable, + tags: Vec::new(), + unique_key: None, + unique_states: None, + } +} + +#[test] +fn retry_delay_is_seeded_bounded_and_capped() { + let now = Utc::now(); + let row = retry_row(0); + let first = default_retry_delay(&row, now, 123); + assert_eq!(first, default_retry_delay(&row, now, 123)); + assert_ne!(first, default_retry_delay(&row, now, 456)); + assert!(first >= Duration::from_millis(900)); + assert!(first <= Duration::from_millis(1_100)); + + assert_eq!( + default_retry_delay(&retry_row(309), now, 123), + Duration::from_nanos(i64::MAX as u64) + ); + // Just below the cap, upward jitter must not exceed it. + for seed in 0..64 { + assert!( + default_retry_delay(&retry_row(308), now, seed) + <= Duration::from_nanos(i64::MAX as u64) + ); + } +} + +#[tokio::test] +async fn worker_failures_record_the_error_chain_and_panic_value() { + #[derive(Debug, thiserror::Error)] + #[error("charging card")] + struct ChargeError(#[source] std::io::Error); + + let failure = super::executor::worker_join_result(Ok(Err(WorkError::new(ChargeError( + std::io::Error::other("card declined"), + ))))) + .unwrap_err(); + assert_eq!(failure.error, "charging card: card declined"); + + let join_error = tokio::spawn(async { panic!("boom") }).await.unwrap_err(); + let failure = super::executor::worker_join_result(Err(join_error)).unwrap_err(); + assert_eq!(failure.error, "boom"); + let WorkResult::Panicked(panic) = super::executor::public_work_result(&Err(failure)) else { + panic!("expected a panic result"); + }; + assert_eq!(panic.message(), "boom"); + assert_eq!(panic.to_string(), "worker panicked: boom"); +} + +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn extension_client_finds_only_the_installed_pilot() { + #[derive(Debug)] + struct InstalledPilot; + impl Pilot for InstalledPilot {} + + #[derive(Debug)] + struct OtherPilot; + impl Pilot for OtherPilot {} + + let pool = sqlx::SqlitePool::connect_lazy("sqlite::memory:").unwrap(); + let client = Client::builder(pool) + .with_pilot(InstalledPilot) + .build() + .unwrap(); + let extension = crate::__private::ExtensionClient::new(&client); + assert!(extension.pilot::().is_some()); + assert!(extension.pilot::().is_none()); +} + +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn erased_transactions_run_requests_in_the_callers_transaction() { + #[derive(Deserialize, serde::Serialize)] + struct ErasedArgs {} + + impl JobArgs for ErasedArgs { + const KIND: &'static str = "erased"; + } + + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + riverqueue_migrate::SqliteMigrator::new(pool.clone()) + .migrate_up() + .await + .unwrap(); + let client = Client::builder(pool.clone()).build().unwrap(); + let id = client.insert(ErasedArgs {}).await.unwrap().id(); + + let mut transaction = crate::database::begin_sqlite_write(&pool).await.unwrap(); + let mut erased = client.inner.database.transaction(&mut transaction).unwrap(); + let cancelled = client.jobs().cancel(id).tx(&mut erased).await.unwrap(); + assert_eq!(cancelled.state, JobState::Cancelled); + transaction.rollback().await.unwrap(); + assert_eq!( + client.jobs().get(id).await.unwrap().state, + JobState::Available + ); +} + +#[tokio::test] +async fn completion_retries_recover_from_a_transient_error() { + let attempts = AtomicU64::new(0); + let result = with_completion_retries("test completion", || async { + if attempts.fetch_add(1, Ordering::SeqCst) == 0 { + Err(Error::from(sqlx::Error::PoolTimedOut)) + } else { + Ok("persisted") + } + }) + .await; + assert_eq!(result.unwrap(), "persisted"); + assert_eq!(attempts.load(Ordering::SeqCst), 2); +} + +#[tokio::test] +async fn completion_retries_stop_immediately_for_a_closed_pool() { + let attempts = AtomicU64::new(0); + let result = with_completion_retries("test completion", || async { + attempts.fetch_add(1, Ordering::SeqCst); + Err::<(), _>(Error::from(sqlx::Error::PoolClosed)) + }) + .await; + assert!(result.is_err()); + assert_eq!(attempts.load(Ordering::SeqCst), 1); +} + +#[test] +fn go_time_json_matches_go_rfc3339_nano() { + use chrono::TimeZone as _; + + let base = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + assert_eq!(go_time_json(base), "2026-01-02T03:04:05Z"); + assert_eq!( + go_time_json(base + chrono::Duration::nanoseconds(120_000_000)), + "2026-01-02T03:04:05.12Z" + ); + assert_eq!( + go_time_json(base + chrono::Duration::nanoseconds(123_456_789)), + "2026-01-02T03:04:05.123456789Z" + ); +} + +#[test] +fn schedule_delays_clamp_like_go_durations() { + use chrono::TimeZone as _; + + let now = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + assert_eq!( + scheduled_after(now, Duration::from_secs(90)), + now + chrono::Duration::seconds(90) + ); + let clamped = scheduled_after(now, Duration::MAX); + assert_eq!( + clamped, + now + chrono::Duration::nanoseconds(i64::MAX), + "delays saturate at Go's maximum time.Duration" + ); +} + +#[cfg(feature = "sqlite")] +#[tokio::test(flavor = "multi_thread")] +async fn subscription_forwarder_stops_when_the_receiver_drops() { + #[derive(Clone, Debug, serde::Deserialize, crate::JobArgs, serde::Serialize)] + #[river(kind = "subscription_forwarder_test")] + struct ForwarderArgs {} + + let mut workers = WorkerRegistry::new(); + workers + .register_fn(|_context: WorkContext, _job: Job| async { + Ok::<_, std::convert::Infallible>(WorkOutcome::Complete) + }) + .unwrap(); + let pool = sqlx::SqlitePool::connect_lazy("sqlite::memory:").unwrap(); + let client = Client::builder(pool) + .workers(workers) + .queue("default", QueueConfig::new(1)) + .build() + .unwrap(); + let baseline = client.inner.events.receiver_count(); + let receiver = client.subscribe(&[EventKind::JobCompleted]).unwrap(); + assert_eq!(client.inner.events.receiver_count(), baseline + 1); + + drop(receiver); + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + while client.inner.events.receiver_count() > baseline { + assert!( + tokio::time::Instant::now() < deadline, + "forwarder outlived its receiver" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } +} + +#[cfg(feature = "postgres")] +#[tokio::test] +async fn intercepting_extensions_can_only_lower_completion_concurrency() { + #[derive(Clone, Copy)] + struct ConcurrencyPilot(usize); + + #[async_trait::async_trait] + impl crate::__private::Pilot for ConcurrencyPilot { + fn intercepts_job_set_state(&self) -> bool { + true + } + + fn job_set_state_concurrency(&self) -> usize { + self.0 + } + } + + let pool = sqlx::PgPool::connect_lazy("postgres://localhost/unused").unwrap(); + let concurrency = |pilot: Option| { + let builder = Client::builder(pool.clone()); + let client = match pilot { + Some(pilot) => builder.with_pilot(pilot), + None => builder, + } + .build() + .unwrap(); + CompletionBatcher::new(Arc::clone(&client.inner)).concurrency() + }; + assert_eq!(concurrency(None), 2); + assert_eq!(concurrency(Some(ConcurrencyPilot(0))), 1); + assert_eq!(concurrency(Some(ConcurrencyPilot(1))), 1); + assert_eq!(concurrency(Some(ConcurrencyPilot(8))), 2); +} + +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn periodic_jobs_run_at_their_target_unless_scheduled_explicitly() { + let pool = sqlx::SqlitePool::connect_lazy("sqlite::memory:").unwrap(); + let client = Client::builder(pool).build().unwrap(); + let now = Utc::now(); + let target = now - chrono::Duration::milliseconds(5); + let prepare = |opts: InsertParams| { + client + .prepare_periodic( + "periodic_test", + &[], + serde_json::value::to_raw_value(&serde_json::json!({})).unwrap(), + opts, + target, + now, + ) + .unwrap() + }; + let defaults = || { + InsertOpts::resolve( + MAX_ATTEMPTS_DEFAULT, + InsertOpts::default(), + InsertOpts::default(), + ) + }; + + // A due job runs immediately at its target, as Go's enqueuer inserts it. + let due = prepare(defaults()); + assert_eq!(due.state, JobState::Available); + assert_eq!(due.opts.scheduled_at, Some(target)); + + // An explicit schedule from the constructor is kept and waits. + let later = now + chrono::Duration::minutes(5); + let mut explicit = defaults(); + explicit.scheduled_at = Some(later); + let explicit = prepare(explicit); + assert_eq!(explicit.state, JobState::Scheduled); + assert_eq!(explicit.opts.scheduled_at, Some(later)); + + // A pending job stays pending. + let mut pending = defaults(); + pending.pending = true; + assert_eq!(prepare(pending).state, JobState::Pending); +} + +/// Like River Go's completer stop path, a client stopping during an outage +/// gives up on its unwritten completions after the first failed batch, even +/// when the backlog is too full for the batcher to receive the end of its +/// channel, instead of retrying every batch until the database returns. +#[cfg(feature = "sqlite")] +#[tokio::test(flavor = "multi_thread")] +async fn completer_abandons_its_backlog_when_a_batch_fails_during_shutdown() { + // Without River's tables, every completion write fails and is retried. + let pool = sqlx::SqlitePool::connect_lazy("sqlite::memory:").unwrap(); + let client = Client::builder(pool).build().unwrap(); + let inner = Arc::clone(&client.inner); + // More than a full backlog plus the batch in flight. + let updates = COMPLETION_BACKLOG_LIMIT + 2 * COMPLETION_BATCH_SIZE; + let (sender, receiver) = mpsc::channel(updates); + let now = std::time::Instant::now(); + for job_id in 0..i64::try_from(updates).unwrap() { + let cancellation = CancellationToken::new(); + inner + .running + .lock() + .unwrap() + .insert(job_id, cancellation.clone()); + sender + .try_send(CompletionUpdate { + attempt: None, + cancellation, + error: None, + event_kind: JobEventKind::Completed, + finalized_at: Some(Utc::now()), + job_id, + metadata: Map::new(), + peer: None, + scheduled_at: None, + state: JobState::Completed, + timing: CompletionTiming { + completion_started: now, + queue_wait_duration: Duration::ZERO, + run_duration: Duration::ZERO, + }, + }) + .unwrap_or_else(|_| panic!("the channel has room")); + } + // Every producer has stopped. + drop(sender); + + // One batch's retry cycle sleeps about three seconds. + tokio::time::timeout( + Duration::from_secs(30), + run_completion_batcher(Arc::clone(&inner), receiver), + ) + .await + .expect("the completer stops after the first failed batch") + .unwrap(); + assert!(inner.running.lock().unwrap().is_empty()); +} + +/// A notification listener that panics is restarted, and the restarted +/// listener reports the client ready rather than the client appearing to +/// have stopped before becoming ready. +#[cfg(feature = "sqlite")] +#[tokio::test(flavor = "multi_thread")] +async fn readiness_survives_a_notification_listener_panic() { + #[derive(Deserialize, serde::Serialize)] + struct ReadinessArgs {} + + impl JobArgs for ReadinessArgs { + const KIND: &'static str = "readiness"; + } + + let path = std::env::temp_dir().join(format!( + "river-readiness-{}-{}.sqlite", + std::process::id(), + Utc::now().timestamp_nanos_opt().unwrap_or_default() + )); + let pool = sqlx::sqlite::SqlitePoolOptions::new() + .max_connections(4) + .connect_with( + sqlx::sqlite::SqliteConnectOptions::new() + .filename(&path) + .create_if_missing(true) + .journal_mode(sqlx::sqlite::SqliteJournalMode::Wal), + ) + .await + .unwrap(); + riverqueue_migrate::SqliteMigrator::new(pool.clone()) + .migrate_up() + .await + .unwrap(); + let mut workers = WorkerRegistry::new(); + workers + .register_fn(|_context: WorkContext, _job: Job| async { + Ok::<_, std::convert::Infallible>(WorkOutcome::Complete) + }) + .unwrap(); + let client = Client::builder(pool.clone()) + .without_leader_election() + .workers(workers) + .queue("default", QueueConfig::new(1)) + .build() + .unwrap(); + client + .inner + .notifier_start_panics + .store(1, Ordering::Release); + + let mut run = client.start().unwrap(); + tokio::time::timeout(Duration::from_secs(10), run.wait_ready()) + .await + .expect("the restarted listener reports readiness") + .unwrap(); + assert_eq!( + client.inner.notifier_start_panics.load(Ordering::Acquire), + 0 + ); + run.shutdown().await.unwrap(); + pool.close().await; + for suffix in ["", "-shm", "-wal"] { + let mut file = path.as_os_str().to_owned(); + file.push(suffix); + let _ = std::fs::remove_file(file); + } +} + +/// Insert wakeups are far more frequent than leadership events. A burst of +/// them must not push a resignation request out of the elector's channel, as +/// it could when both shared one lagging broadcast channel. +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn resign_requests_survive_a_burst_of_insert_notifications() { + let pool = sqlx::SqlitePool::connect_lazy("sqlite::memory:").unwrap(); + let client = Client::builder(pool).build().unwrap(); + let inner = &client.inner; + // The receiver the supervisor hands to maintenance. + let mut leadership = inner.leadership_wakeups.subscribe(); + let mut producer = inner.queue_notifications.subscribe(); + + dispatch_notification( + inner, + &inner.queue_notifications, + crate::NOTIFICATION_TOPIC_LEADERSHIP, + r#"{"action":"request_resign"}"#, + ); + for _ in 0..4_096 { + dispatch_notification( + inner, + &inner.queue_notifications, + crate::NOTIFICATION_TOPIC_INSERT, + r#"{"queue":"default"}"#, + ); + } + + assert!(matches!( + leadership.try_recv(), + Ok(LeadershipWakeup::RequestResign) + )); + assert!(leadership.try_recv().is_err()); + // The producers' channel lagged, which producers recover from by + // fetching and refreshing everything. + assert!(matches!( + producer.try_recv(), + Err(broadcast::error::TryRecvError::Lagged(_)) + )); +} + +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn typed_timeouts_and_retentions_validate_at_build() { + let pool = sqlx::SqlitePool::connect_lazy("sqlite::memory:").unwrap(); + let builder = || Client::builder(pool.clone()); + + let client = builder().build().unwrap(); + assert_eq!(client.inner.job_timeout, Some(JOB_TIMEOUT_DEFAULT)); + let client = builder().without_job_timeout().build().unwrap(); + assert_eq!(client.inner.job_timeout, None); + let client = builder() + .job_timeout(Duration::from_secs(5)) + .build() + .unwrap(); + assert_eq!(client.inner.job_timeout, Some(Duration::from_secs(5))); + for error in [ + builder().job_timeout(Duration::ZERO).build().unwrap_err(), + builder() + .job_stuck_threshold(Duration::ZERO) + .build() + .unwrap_err(), + builder() + .soft_stop_timeout(Duration::ZERO) + .build() + .unwrap_err(), + ] { + assert!(matches!(error, Error::Configuration(_)), "{error}"); + } + + let defaults = MaintenanceConfig::default(); + assert_eq!( + defaults.completed_job_retention(), + Retention::DeleteAfter(Duration::from_hours(24)) + ); + let keep = defaults.with_completed_job_retention(Retention::Keep); + assert_eq!(keep.completed_job_retention(), Retention::Keep); + assert_eq!(keep.completed_job_retention, None); +} + +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn ending_maintenance_abnormally_cancels_its_terms() { + let pool = sqlx::SqlitePool::connect_lazy("sqlite::memory:").unwrap(); + let client = Client::builder(pool).build().unwrap(); + let cancel = CancellationToken::new(); + let maintenance = tokio::spawn(crate::maintenance::run_maintenance( + Arc::clone(&client.inner), + cancel.clone(), + client.inner.leadership_wakeups.subscribe(), + )); + tokio::task::yield_now().await; + assert!(!cancel.is_cancelled()); + + // As when the task panics: its terms' services run under child tokens, + // and the supervisor restarts it with a new token. + maintenance.abort(); + assert!(maintenance.await.unwrap_err().is_cancelled()); + assert!(cancel.is_cancelled()); +} + +#[cfg(feature = "postgres")] +#[tokio::test] +async fn reindexer_timeout_is_positive_or_disabled() { + use crate::database::{PostgresDatabase, PostgresReindexConfig}; + + let pool = sqlx::PgPool::connect_lazy("postgres://localhost/unused").unwrap(); + let build = |config: PostgresReindexConfig| { + Client::builder(PostgresDatabase::new(pool.clone()).with_reindex(config)).build() + }; + + assert_eq!( + PostgresReindexConfig::default().timeout(), + Some(Duration::from_mins(1)) + ); + let disabled = PostgresReindexConfig::default().without_timeout(); + assert_eq!(disabled.timeout(), None); + build(disabled).unwrap(); + let error = build(PostgresReindexConfig::default().with_timeout(Duration::ZERO)).unwrap_err(); + assert!(matches!(error, Error::Configuration(_)), "{error}"); +} diff --git a/rust/riverqueue/src/client/validate.rs b/rust/riverqueue/src/client/validate.rs new file mode 100644 index 000000000..7906248a6 --- /dev/null +++ b/rust/riverqueue/src/client/validate.rs @@ -0,0 +1,147 @@ +//! Validation of job insertion parameters. + +use crate::{Error, InsertParams}; + +pub(super) fn validate_insert_parts( + kind: &str, + opts: &InsertParams, + allow_legacy_job_kinds: bool, +) -> Result<(), Error> { + let mut kind_characters = kind.chars(); + if !allow_legacy_job_kinds + && (kind.len() < 2 + || kind.len() >= 128 + || !kind_characters.next().is_some_and(is_word) + || !kind_characters.all(valid_kind_character)) + { + return Err(Error::invalid_job(format!("invalid job kind {kind:?}"))); + } + if opts.max_attempts < 1 { + return Err(Error::invalid_job( + "max_attempts must be greater than zero".to_owned(), + )); + } + if !(1..=4).contains(&opts.priority) { + return Err(Error::invalid_job( + "priority must be between one and four".to_owned(), + )); + } + validate_queue(&opts.queue)?; + for tag in &opts.tags { + if tag.len() > 255 || tag.len() < 3 { + return Err(Error::invalid_job( + "tags must contain between 3 and 255 bytes".to_owned(), + )); + } + let mut characters = tag.chars(); + let first = characters.next().unwrap(); + let last = tag.chars().next_back().unwrap(); + if !is_word(first) + || !is_word(last) + || !characters.all(|character| is_word(character) || character == '-') + { + return Err(Error::invalid_job(format!("invalid tag {tag:?}"))); + } + } + opts.unique.validate().map_err(Error::invalid_job) +} + +pub(super) fn valid_kind_character(character: char) -> bool { + character.is_ascii_alphanumeric() + || matches!( + character, + '_' | '-' | '[' | ']' | '<' | '>' | '/' | '.' | '·' | ':' | '+' + ) +} + +pub(crate) fn validate_queue(queue: &str) -> Result<(), Error> { + if queue.is_empty() || queue.len() > 64 { + return Err(Error::invalid_job( + "queue name must contain between 1 and 64 bytes".to_owned(), + )); + } + if !queue + .chars() + .next() + .is_some_and(|character| character.is_ascii_lowercase() || character.is_ascii_digit()) + { + return Err(Error::invalid_job(format!("invalid queue name {queue:?}"))); + } + let mut previous_separator = false; + for character in queue.chars() { + let separator = matches!(character, '_' | '|' | '-'); + if !(character.is_ascii_lowercase() || character.is_ascii_digit() || separator) + || (separator && previous_separator) + { + return Err(Error::invalid_job(format!("invalid queue name {queue:?}"))); + } + previous_separator = separator; + } + if previous_separator { + return Err(Error::invalid_job(format!("invalid queue name {queue:?}"))); + } + Ok(()) +} + +#[cfg(feature = "postgres")] +pub(super) fn validate_identifier(identifier: &str, description: &str) -> Result<(), Error> { + let mut characters = identifier.chars(); + if identifier.is_empty() + || identifier.len() > 63 + || !characters + .next() + .is_some_and(|character| character == '_' || character.is_ascii_alphabetic()) + || !characters.all(|character| character == '_' || character.is_ascii_alphanumeric()) + { + return Err(Error::invalid_job(format!( + "invalid PostgreSQL {description} identifier {identifier:?}" + ))); + } + Ok(()) +} + +pub(super) fn is_word(character: char) -> bool { + character == '_' || character.is_ascii_alphanumeric() +} + +#[cfg(test)] +mod tests { + use super::validate_queue; + + #[test] + fn queue_names_match_go_validation() { + // Mirrors Go's `^(?:[a-z0-9])+(?:[_|\-]?[a-z0-9]+)*$` plus its + // 64-byte limit. + for valid in [ + "0", + "a", + "a-b", + "a_b", + "a|b", + "default", + "tenant|priority_emails-2", + &"a".repeat(64), + ] { + assert!(validate_queue(valid).is_ok(), "{valid:?} should be valid"); + } + for invalid in [ + "", + "-a", + "A", + "_a", + "a b", + "a-", + "a.b", + "a__b", + "a_|b", + "a|", + "|a", + &"a".repeat(65), + ] { + assert!( + validate_queue(invalid).is_err(), + "{invalid:?} should be invalid" + ); + } + } +} diff --git a/rust/riverqueue/src/database.rs b/rust/riverqueue/src/database.rs new file mode 100644 index 000000000..bf8c6d7b5 --- /dev/null +++ b/rust/riverqueue/src/database.rs @@ -0,0 +1,1031 @@ +//! Database sources and executor compatibility contracts. +//! +//! River's database abstraction is closed over its built-in backends. The +//! sealed conversion traits in this module let [`Client`](crate::Client) +//! remain non-generic while preventing an accidental public driver SPI. + +use std::fmt; +#[cfg(feature = "postgres")] +use std::time::Duration; + +#[cfg(feature = "postgres")] +use chrono::NaiveTime; +pub use riverqueue_migrate::{SchemaName, SchemaNameError}; +#[cfg(feature = "postgres")] +use sqlx::{PgConnection, PgPool, Postgres}; +#[cfg(feature = "sqlite")] +use sqlx::{Sqlite, SqliteConnection, SqlitePool}; +use sqlx::{Transaction, pool::PoolConnection}; +use thiserror::Error; + +/// Begins a PostgreSQL transaction that is never abandoned half-started. +/// +/// SQLx 0.9 records a transaction only once the server has answered `BEGIN`. If +/// the future beginning it is dropped after `BEGIN` reaches the server but +/// before that answer arrives, for example because a `select!` or timeout +/// around it fires, SQLx never queues a `ROLLBACK` and the connection goes +/// back to the pool idle in a transaction. This function begins the +/// transaction on its own task instead: if the caller stops waiting, the task +/// still finishes, and dropping the finished transaction rolls it back. +/// +/// River begins its own transactions this way. Use it in place of +/// [`PgPool::begin`] wherever the begin may be cancelled, including inside +/// workers, whose futures River drops when they outlive the job stuck +/// threshold after cancellation. +/// +/// # Examples +/// +/// ```no_run +/// # async fn example(pool: sqlx::PgPool) -> Result<(), sqlx::Error> { +/// let mut transaction = riverqueue::database::begin_postgres(&pool).await?; +/// sqlx::query("SELECT 1").execute(&mut *transaction).await?; +/// transaction.commit().await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// # Errors +/// +/// Returns the error from acquiring a connection or beginning the +/// transaction, or an I/O error when the runtime is shutting down. +/// +/// # Panics +/// +/// Resumes a panic from the task beginning the transaction. +#[cfg(feature = "postgres")] +pub async fn begin_postgres(pool: &PgPool) -> Result, sqlx::Error> { + let pool = pool.clone(); + run_to_completion(async move { pool.begin().await }).await +} + +/// Begins a SQLite transaction that may write, protected like +/// `begin_postgres` against being abandoned half-started. +/// +/// It begins with `BEGIN IMMEDIATE`, which takes the write lock up front, so +/// a transaction that reads before writing can't fail with +/// `SQLITE_BUSY_SNAPSHOT` when another connection commits in between. Use it +/// for transactions passed to River's `.tx` that may write, in place of +/// [`SqlitePool::begin`], whose deferred transactions can. +/// +/// # Examples +/// +/// ```no_run +/// # async fn example(pool: sqlx::SqlitePool) -> Result<(), sqlx::Error> { +/// let mut transaction = riverqueue::database::begin_sqlite_write(&pool).await?; +/// sqlx::query("SELECT 1").execute(&mut *transaction).await?; +/// transaction.commit().await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// # Errors +/// +/// Returns the error from acquiring a connection or beginning the +/// transaction, including `SQLITE_BUSY` when the write lock stays held past +/// the busy timeout, or an I/O error when the runtime is shutting down. +/// +/// # Panics +/// +/// Resumes a panic from the task beginning the transaction. +#[cfg(feature = "sqlite")] +pub async fn begin_sqlite_write( + pool: &SqlitePool, +) -> Result, sqlx::Error> { + let pool = pool.clone(); + run_to_completion(async move { pool.begin_with("BEGIN IMMEDIATE").await }).await +} + +/// Runs `operation` on its own task, so dropping the returned future doesn't +/// cancel it midway. +async fn run_to_completion( + operation: impl Future> + Send + 'static, +) -> Result { + match tokio::spawn(operation).await { + Ok(result) => result, + Err(error) if error.is_panic() => std::panic::resume_unwind(error.into_panic()), + // The runtime is shutting down and cancelled the task. + Err(error) => Err(sqlx::Error::Io(std::io::Error::other(format!( + "transaction begin task failed: {error}" + )))), + } +} + +pub(crate) mod erased; +#[cfg(feature = "postgres")] +pub(crate) mod postgres_capabilities; +#[cfg(feature = "sqlite")] +pub(crate) mod sqlite; + +use crate::__private::DatabaseConnection; +pub(crate) use erased::{ + Database, DatabaseInner, ErasedExecutor, ErasedTransaction, ExecutorInner, +}; +pub(crate) use private::DatabaseExecutorSealed as DatabaseExecutor; + +/// The database a [`Client`](crate::Client) was built with, from +/// [`Client::database`](crate::Client::database). +#[derive(Clone, Copy, Debug)] +#[non_exhaustive] +pub enum ClientDatabase<'a> { + /// A PostgreSQL client's pool and options. + #[cfg(feature = "postgres")] + Postgres(&'a PostgresDatabase), + /// A SQLite client's pool. + #[cfg(feature = "sqlite")] + Sqlite(&'a SqliteDatabase), +} + +impl<'a> ClientDatabase<'a> { + /// Returns the PostgreSQL pool and options, or `None` for another + /// backend. + #[cfg(feature = "postgres")] + #[must_use] + pub const fn as_postgres(&self) -> Option<&'a PostgresDatabase> { + match self { + Self::Postgres(database) => Some(database), + #[cfg(feature = "sqlite")] + Self::Sqlite(_) => None, + } + } + + /// Returns the SQLite pool, or `None` for another backend. + #[cfg(feature = "sqlite")] + #[must_use] + pub const fn as_sqlite(&self) -> Option<&'a SqliteDatabase> { + match self { + #[cfg(feature = "postgres")] + Self::Postgres(_) => None, + Self::Sqlite(database) => Some(database), + } + } + + /// Returns the backend. + #[must_use] + pub const fn kind(&self) -> DatabaseKind { + match self { + #[cfg(feature = "postgres")] + Self::Postgres(_) => DatabaseKind::Postgres, + #[cfg(feature = "sqlite")] + Self::Sqlite(_) => DatabaseKind::Sqlite, + } + } +} + +/// A database backend understood by River. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +#[non_exhaustive] +pub enum DatabaseKind { + /// PostgreSQL. + #[cfg(feature = "postgres")] + Postgres, + /// SQLite and compatible implementations. + #[cfg(feature = "sqlite")] + Sqlite, +} + +impl fmt::Display for DatabaseKind { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(match self { + #[cfg(feature = "postgres")] + Self::Postgres => "PostgreSQL", + #[cfg(feature = "sqlite")] + Self::Sqlite => "SQLite", + }) + } +} + +/// A PostgreSQL source and its backend-specific River options. +/// +/// River detects what the server supports the first time it needs to know +/// and remembers it for this value and its clones. On PostgreSQL 18 and +/// later, a unique insert tells a new row from an existing one with +/// `RETURNING OLD`, and with `xmax` before that. +/// +/// YugabyteDB works as a PostgreSQL server. It has no `xmax`, so a unique +/// insert marks its row with a random `river:unique_nonce` metadata value, +/// as on SQLite. Unless its `yb_enable_listen_notify` setting is on, it has +/// no `LISTEN`/`NOTIFY` either: River then sends no notifications, and a +/// client polls for new jobs every fetch poll interval, and for queue +/// changes and cancellations of its running jobs every two seconds, as if +/// built [`without_notifications`](crate::ClientBuilder::without_notifications). +/// Yugabyte's notifications need version 2025.2.3 or later with +/// `ysql_yb_enable_listen_notify=true` on both masters and tservers. Since +/// the detection is remembered, enabling them takes effect for a new +/// `PostgresDatabase`, such as after a restart. +#[cfg(feature = "postgres")] +#[derive(Clone)] +pub struct PostgresDatabase { + capabilities: postgres_capabilities::CapabilitiesCache, + pool: PgPool, + reindex: PostgresReindexConfig, + schema: SchemaName, +} + +#[cfg(feature = "postgres")] +impl PostgresDatabase { + /// Uses a PostgreSQL pool and the connection's current schema. + #[must_use] + pub fn new(pool: PgPool) -> Self { + Self { + capabilities: postgres_capabilities::CapabilitiesCache::default(), + pool, + reindex: PostgresReindexConfig::default(), + schema: SchemaName::current(), + } + } + + /// Returns the underlying SQLx pool. + #[must_use] + pub const fn pool(&self) -> &PgPool { + &self.pool + } + + /// Returns the server capabilities detected for this database, shared + /// by its clones. + pub(crate) const fn capabilities(&self) -> &postgres_capabilities::CapabilitiesCache { + &self.capabilities + } + + /// Returns the database with PostgreSQL's periodic concurrent index + /// rebuilding configured. + #[must_use] + pub fn with_reindex(mut self, reindex: PostgresReindexConfig) -> Self { + self.reindex = reindex; + self + } + + pub(crate) fn extend_default_reindex_names( + &mut self, + names: impl IntoIterator>, + ) { + self.reindex.extend_default_index_names(names); + } + + /// Returns PostgreSQL reindexer configuration. + #[must_use] + pub const fn reindex(&self) -> &PostgresReindexConfig { + &self.reindex + } + + /// Returns the database with an explicit PostgreSQL schema for River + /// objects and notification channels. + #[must_use] + pub fn with_schema(mut self, schema: SchemaName) -> Self { + self.schema = schema; + self + } + + /// Returns the configured PostgreSQL schema. + #[must_use] + pub const fn schema(&self) -> &SchemaName { + &self.schema + } +} + +#[cfg(feature = "postgres")] +impl fmt::Debug for PostgresDatabase { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("PostgresDatabase") + .field("reindex", &self.reindex) + .field("schema", &self.schema) + .finish_non_exhaustive() + } +} + +/// PostgreSQL-specific concurrent reindexer configuration. +#[cfg(feature = "postgres")] +#[derive(Clone, Debug)] +pub struct PostgresReindexConfig { + index_names: Vec, + /// Whether [`Self::with_index_names`] replaced River's default list. + /// Add-on crates extend only a default list, so a custom schedule or + /// timeout doesn't discard their indexes. + index_names_explicit: bool, + schedule: PostgresReindexSchedule, + timeout: Option, +} + +#[cfg(feature = "postgres")] +impl PostgresReindexConfig { + /// Appends add-on indexes unless the caller chose index names + /// explicitly, including an empty list that disables the service. + fn extend_default_index_names(&mut self, names: impl IntoIterator>) { + if self.index_names_explicit { + return; + } + for name in names { + let name = name.into(); + if !self.index_names.contains(&name) { + self.index_names.push(name); + } + } + } + + /// Returns configured index names. + #[must_use] + pub fn index_names(&self) -> &[String] { + &self.index_names + } + + /// Returns the reindex schedule. + #[must_use] + pub const fn schedule(&self) -> PostgresReindexSchedule { + self.schedule + } + + /// Returns how long one index rebuild may run before River cancels it, + /// or `None` when rebuilds run without a limit. + #[must_use] + pub const fn timeout(&self) -> Option { + self.timeout + } + + /// Replaces indexes rebuilt by River. An empty iterator disables the + /// service. + #[must_use] + pub fn with_index_names(mut self, names: impl IntoIterator>) -> Self { + self.index_names = names.into_iter().map(Into::into).collect(); + self.index_names_explicit = true; + self + } + + /// Replaces the reindex schedule. + #[must_use] + pub const fn with_schedule(mut self, schedule: PostgresReindexSchedule) -> Self { + self.schedule = schedule; + self + } + + /// Sets how long one index rebuild may run before River cancels it. + /// Defaults to one minute. The timeout must be positive; use + /// [`without_timeout`](Self::without_timeout) to let rebuilds run without + /// a limit. + #[must_use] + pub const fn with_timeout(mut self, timeout: Duration) -> Self { + self.timeout = Some(timeout); + self + } + + /// Lets index rebuilds run without a time limit, like River Go's + /// `ReindexerTimeout` of `-1`. + #[must_use] + pub const fn without_timeout(mut self) -> Self { + self.timeout = None; + self + } +} + +#[cfg(feature = "postgres")] +impl Default for PostgresReindexConfig { + fn default() -> Self { + Self { + index_names: vec![ + "river_job_args_index".to_owned(), + "river_job_kind".to_owned(), + "river_job_metadata_index".to_owned(), + "river_job_pkey".to_owned(), + "river_job_prioritized_fetching_index".to_owned(), + "river_job_state_and_finalized_at_index".to_owned(), + "river_job_unique_idx".to_owned(), + ], + index_names_explicit: false, + schedule: PostgresReindexSchedule::default(), + timeout: Some(Duration::from_mins(1)), + } + } +} + +/// Schedule used by PostgreSQL's concurrent reindexer. +#[cfg(feature = "postgres")] +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[non_exhaustive] +pub enum PostgresReindexSchedule { + /// Run each day at the supplied UTC wall-clock time. + DailyUtc(NaiveTime), + /// Run after each elapsed interval from client startup. + Interval(Duration), +} + +#[cfg(feature = "postgres")] +impl Default for PostgresReindexSchedule { + fn default() -> Self { + Self::DailyUtc(NaiveTime::MIN) + } +} + +/// A SQLite source and its backend-specific River options. +#[cfg(feature = "sqlite")] +#[derive(Clone)] +pub struct SqliteDatabase { + pool: SqlitePool, +} + +#[cfg(feature = "sqlite")] +impl SqliteDatabase { + /// Uses a SQLite pool. + #[must_use] + pub const fn new(pool: SqlitePool) -> Self { + Self { pool } + } + + /// Returns the underlying SQLx pool. + #[must_use] + pub const fn pool(&self) -> &SqlitePool { + &self.pool + } +} + +#[cfg(feature = "sqlite")] +impl fmt::Debug for SqliteDatabase { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("SqliteDatabase") + .finish_non_exhaustive() + } +} + +/// Error returned when an operation receives an executor for another backend. +#[derive(Clone, Debug, Eq, Error, PartialEq)] +#[error("database executor mismatch: expected {expected}, received {actual}")] +pub struct DatabaseMismatch { + actual: DatabaseKind, + expected: DatabaseKind, +} + +impl DatabaseMismatch { + /// Returns the executor's backend. + #[must_use] + pub const fn actual(&self) -> DatabaseKind { + self.actual + } + + /// Returns the database backend required by the client. + #[must_use] + pub const fn expected(&self) -> DatabaseKind { + self.expected + } +} + +/// A value accepted as a built-in River database source. +/// +/// This trait has no public methods and is sealed. Applications select a +/// backend by passing a [`PgPool`], [`SqlitePool`], [`PostgresDatabase`], or +/// [`SqliteDatabase`]; they do not implement a River driver trait. +pub trait IntoDatabase: private::IntoDatabaseSealed {} + +#[doc(hidden)] +impl IntoDatabase for Database {} +#[cfg(feature = "postgres")] +impl IntoDatabase for PgPool {} +#[cfg(feature = "postgres")] +impl IntoDatabase for &PgPool {} +#[cfg(feature = "postgres")] +impl IntoDatabase for PostgresDatabase {} +#[cfg(feature = "postgres")] +impl IntoDatabase for &PostgresDatabase {} +#[cfg(feature = "sqlite")] +impl IntoDatabase for SqlitePool {} +#[cfg(feature = "sqlite")] +impl IntoDatabase for &SqlitePool {} +#[cfg(feature = "sqlite")] +impl IntoDatabase for SqliteDatabase {} +#[cfg(feature = "sqlite")] +impl IntoDatabase for &SqliteDatabase {} + +/// A caller-owned SQLx transaction accepted by River's transactional +/// operations. +/// +/// This trait has no public methods and is sealed. It is deliberately not +/// implemented for pools or bare connections so a request's `tx` method +/// cannot accidentally run in autocommit mode. +/// +/// Begin transactions with [`begin_postgres`], or on SQLite with +/// [`begin_sqlite_write`] for transactions that may write. Both are safe to +/// abandon partway, and `begin_sqlite_write` takes SQLite's write lock up +/// front, so a transaction that reads before it writes can't fail with +/// `SQLITE_BUSY_SNAPSHOT` when another connection commits in between. +/// +/// A transaction borrowed mutably is accepted: +/// +/// ```no_run +/// # #[cfg(feature = "postgres")] +/// # async fn example( +/// # client: &riverqueue::Client, +/// # pool: &sqlx::PgPool, +/// # ) -> Result<(), riverqueue::Error> { +/// let mut tx = riverqueue::database::begin_postgres(pool).await?; +/// client.jobs().get(1).tx(&mut tx).await?; +/// tx.commit().await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// Bare connections and pool connections are intentionally rejected: +/// +/// ```compile_fail +/// # async fn example( +/// # client: &riverqueue::Client, +/// # connection: &mut sqlx::PgConnection, +/// # ) -> Result<(), riverqueue::Error> { +/// client.jobs().get(1).tx(connection).await?; +/// # Ok(()) +/// # } +/// ``` +/// +/// ```compile_fail +/// # async fn example( +/// # client: &riverqueue::Client, +/// # connection: &mut sqlx::pool::PoolConnection, +/// # ) -> Result<(), riverqueue::Error> { +/// client.jobs().get(1).tx(connection).await?; +/// # Ok(()) +/// # } +/// ``` +pub trait DatabaseTransactionExecutor<'executor>: + private::DatabaseTransactionExecutorSealed<'executor> +{ +} + +#[cfg(feature = "postgres")] +impl<'executor> DatabaseTransactionExecutor<'executor> + for &'executor mut sqlx::Transaction<'_, sqlx::Postgres> +{ +} +#[cfg(feature = "sqlite")] +impl<'executor> DatabaseTransactionExecutor<'executor> + for &'executor mut sqlx::Transaction<'_, sqlx::Sqlite> +{ +} +/// A transaction a companion crate erased, so it can pass its caller's +/// transaction on to River's requests. +#[doc(hidden)] +impl<'executor> DatabaseTransactionExecutor<'executor> + for &'executor mut crate::__private::ErasedTransaction<'_> +{ +} + +/// Converts a public sealed database source into River's internal erased form. +pub(crate) fn into_database(database: D) -> Database { + private::IntoDatabaseSealed::erase(database) +} + +/// A borrowed built-in pool used by River's internal operation dispatch. +pub(crate) enum DatabasePool<'pool> { + #[cfg(feature = "postgres")] + Postgres(&'pool PgPool), + #[cfg(feature = "sqlite")] + Sqlite(&'pool SqlitePool), +} + +mod private { + use super::{ + Database, DatabaseConnection, DatabaseInner, DatabaseKind, ErasedExecutor, + ErasedTransaction, ExecutorInner, PoolConnection, Transaction, + }; + #[cfg(feature = "postgres")] + use super::{PgConnection, PgPool, Postgres, PostgresDatabase}; + #[cfg(feature = "sqlite")] + use super::{Sqlite, SqliteConnection, SqliteDatabase, SqlitePool}; + + pub trait IntoDatabaseSealed { + fn erase(self) -> Database; + } + + pub trait DatabaseExecutorSealed<'executor> { + fn erase(self) -> ErasedExecutor<'executor>; + } + + /// A caller-managed transaction. Only transactions implement this, so + /// the connection it yields is always inside a transaction River does + /// not commit. + pub trait DatabaseTransactionExecutorSealed<'executor>: + DatabaseExecutorSealed<'executor> + { + fn connection(self) -> DatabaseConnection<'executor>; + } + + const fn connection_executor(connection: DatabaseConnection<'_>) -> ErasedExecutor<'_> { + ErasedExecutor { + inner: ExecutorInner::Connection(connection), + } + } + + const fn pool_executor<'executor>(kind: DatabaseKind) -> ErasedExecutor<'executor> { + ErasedExecutor { + inner: ExecutorInner::Pool(kind), + } + } + + impl IntoDatabaseSealed for Database { + fn erase(self) -> Database { + self + } + } + + #[cfg(feature = "postgres")] + impl IntoDatabaseSealed for PgPool { + fn erase(self) -> Database { + PostgresDatabase::new(self).erase() + } + } + + #[cfg(feature = "postgres")] + impl IntoDatabaseSealed for &PgPool { + fn erase(self) -> Database { + self.clone().erase() + } + } + + #[cfg(feature = "postgres")] + impl IntoDatabaseSealed for PostgresDatabase { + fn erase(self) -> Database { + Database { + inner: DatabaseInner::Postgres(self), + } + } + } + + #[cfg(feature = "postgres")] + impl IntoDatabaseSealed for &PostgresDatabase { + fn erase(self) -> Database { + self.clone().erase() + } + } + + #[cfg(feature = "sqlite")] + impl IntoDatabaseSealed for SqlitePool { + fn erase(self) -> Database { + SqliteDatabase::new(self).erase() + } + } + + #[cfg(feature = "sqlite")] + impl IntoDatabaseSealed for &SqlitePool { + fn erase(self) -> Database { + self.clone().erase() + } + } + + #[cfg(feature = "sqlite")] + impl IntoDatabaseSealed for SqliteDatabase { + fn erase(self) -> Database { + Database { + inner: DatabaseInner::Sqlite(self), + } + } + } + + #[cfg(feature = "sqlite")] + impl IntoDatabaseSealed for &SqliteDatabase { + fn erase(self) -> Database { + self.clone().erase() + } + } + + #[cfg(feature = "postgres")] + impl<'executor> DatabaseExecutorSealed<'executor> for &'executor mut PgConnection { + fn erase(self) -> ErasedExecutor<'executor> { + connection_executor(DatabaseConnection::Postgres(self)) + } + } + + #[cfg(feature = "postgres")] + impl<'executor> DatabaseExecutorSealed<'executor> for &'executor mut PoolConnection { + fn erase(self) -> ErasedExecutor<'executor> { + connection_executor(DatabaseConnection::Postgres(self.as_mut())) + } + } + + #[cfg(feature = "postgres")] + impl<'executor> DatabaseExecutorSealed<'executor> for &'executor PgPool { + fn erase(self) -> ErasedExecutor<'executor> { + pool_executor(DatabaseKind::Postgres) + } + } + + #[cfg(feature = "postgres")] + impl<'executor> DatabaseExecutorSealed<'executor> for &'executor mut Transaction<'_, Postgres> { + fn erase(self) -> ErasedExecutor<'executor> { + connection_executor(DatabaseTransactionExecutorSealed::connection(self)) + } + } + + #[cfg(feature = "postgres")] + impl<'executor> DatabaseTransactionExecutorSealed<'executor> + for &'executor mut Transaction<'_, Postgres> + { + fn connection(self) -> DatabaseConnection<'executor> { + DatabaseConnection::Postgres(self.as_mut()) + } + } + + #[cfg(feature = "sqlite")] + impl<'executor> DatabaseExecutorSealed<'executor> for &'executor mut SqliteConnection { + fn erase(self) -> ErasedExecutor<'executor> { + connection_executor(DatabaseConnection::Sqlite(self)) + } + } + + #[cfg(feature = "sqlite")] + impl<'executor> DatabaseExecutorSealed<'executor> for &'executor mut PoolConnection { + fn erase(self) -> ErasedExecutor<'executor> { + connection_executor(DatabaseConnection::Sqlite(self.as_mut())) + } + } + + #[cfg(feature = "sqlite")] + impl<'executor> DatabaseExecutorSealed<'executor> for &'executor SqlitePool { + fn erase(self) -> ErasedExecutor<'executor> { + pool_executor(DatabaseKind::Sqlite) + } + } + + #[cfg(feature = "sqlite")] + impl<'executor> DatabaseExecutorSealed<'executor> for &'executor mut Transaction<'_, Sqlite> { + fn erase(self) -> ErasedExecutor<'executor> { + connection_executor(DatabaseTransactionExecutorSealed::connection(self)) + } + } + + #[cfg(feature = "sqlite")] + impl<'executor> DatabaseTransactionExecutorSealed<'executor> + for &'executor mut Transaction<'_, Sqlite> + { + fn connection(self) -> DatabaseConnection<'executor> { + DatabaseConnection::Sqlite(self.as_mut()) + } + } + + impl<'executor> DatabaseExecutorSealed<'executor> for &'executor mut ErasedTransaction<'_> { + fn erase(self) -> ErasedExecutor<'executor> { + connection_executor(self.connection()) + } + } + + impl<'executor> DatabaseTransactionExecutorSealed<'executor> + for &'executor mut ErasedTransaction<'_> + { + fn connection(self) -> DatabaseConnection<'executor> { + ErasedTransaction::connection(self) + } + } +} + +#[cfg(all(test, feature = "sqlite"))] +mod tests { + use std::{str::FromStr, sync::Arc, time::Duration}; + + use sqlx::{ + Executor, + sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}, + }; + use tokio::sync::Barrier; + + use super::begin_sqlite_write; + + #[tokio::test] + async fn immediate_writer_avoids_snapshot_upgrade_failure() { + let database_path = std::env::temp_dir().join(format!( + "river-sqlite-write-contention-{}-{}.db", + std::process::id(), + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap() + .as_nanos() + )); + let options = + SqliteConnectOptions::from_str(&format!("sqlite://{}", database_path.display())) + .unwrap() + .create_if_missing(true) + .busy_timeout(Duration::from_secs(2)) + .journal_mode(SqliteJournalMode::Wal); + let pool = SqlitePoolOptions::new() + .max_connections(4) + .connect_with(options) + .await + .unwrap(); + pool.execute("CREATE TABLE contention (value INTEGER NOT NULL)") + .await + .unwrap(); + pool.execute("INSERT INTO contention (value) VALUES (1)") + .await + .unwrap(); + + let mut deferred = pool.begin().await.unwrap(); + let _: i64 = sqlx::query_scalar("SELECT value FROM contention") + .fetch_one(&mut *deferred) + .await + .unwrap(); + pool.execute("UPDATE contention SET value = value + 1") + .await + .unwrap(); + let error = sqlx::query("UPDATE contention SET value = value + 1") + .execute(&mut *deferred) + .await + .unwrap_err(); + assert_eq!( + error + .as_database_error() + .and_then(sqlx::error::DatabaseError::code) + .as_deref(), + Some("517"), + "expected SQLITE_BUSY_SNAPSHOT, received {error}" + ); + deferred.rollback().await.unwrap(); + + let mut immediate = begin_sqlite_write(&pool).await.unwrap(); + let _: i64 = sqlx::query_scalar("SELECT value FROM contention") + .fetch_one(&mut *immediate) + .await + .unwrap(); + let barrier = Arc::new(Barrier::new(2)); + let writer_barrier = Arc::clone(&barrier); + let writer_pool = pool.clone(); + let mut competing_writer = tokio::spawn(async move { + writer_barrier.wait().await; + writer_pool + .execute("UPDATE contention SET value = value + 1") + .await + }); + barrier.wait().await; + assert!( + tokio::time::timeout(Duration::from_millis(50), &mut competing_writer) + .await + .is_err(), + "a competing writer should wait for the immediate transaction" + ); + immediate + .execute("UPDATE contention SET value = value + 1") + .await + .unwrap(); + immediate.commit().await.unwrap(); + competing_writer.await.unwrap().unwrap(); + let value: i64 = sqlx::query_scalar("SELECT value FROM contention") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(value, 4); + + pool.close().await; + std::fs::remove_file(database_path).unwrap(); + } +} + +#[cfg(all(test, feature = "postgres", river_postgres_tests))] +mod postgres_begin_tests { + use std::{sync::Arc, time::Duration}; + + use sqlx::{ + PgPool, + postgres::{PgConnectOptions, PgPoolOptions}, + }; + use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::{TcpListener, TcpStream}, + sync::{Notify, watch}, + }; + + use super::begin_postgres; + + /// A TCP proxy that reports when a client sends `BEGIN` and can hold the + /// server's replies, so a test can stop waiting for a begin that already + /// reached the server. + struct BeginProxy { + address: std::net::SocketAddr, + begin_sent: Arc, + hold_replies: watch::Sender, + } + + impl BeginProxy { + async fn start(upstream: &PgConnectOptions) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let upstream = format!("{}:{}", upstream.get_host(), upstream.get_port()); + let begin_sent = Arc::new(Notify::new()); + let (hold_replies, hold) = watch::channel(false); + let sent = Arc::clone(&begin_sent); + tokio::spawn(async move { + while let Ok((client, _)) = listener.accept().await { + let server = TcpStream::connect(&upstream).await.unwrap(); + let (mut client_read, mut client_write) = client.into_split(); + let (mut server_read, mut server_write) = server.into_split(); + let sent = Arc::clone(&sent); + tokio::spawn(async move { + let mut buffer = vec![0; 8192]; + while let Ok(read) = client_read.read(&mut buffer).await { + if read == 0 || server_write.write_all(&buffer[..read]).await.is_err() { + break; + } + if buffer[..read].windows(5).any(|window| window == b"BEGIN") { + sent.notify_one(); + } + } + }); + let mut hold = hold.clone(); + tokio::spawn(async move { + let mut buffer = vec![0; 8192]; + while let Ok(read) = server_read.read(&mut buffer).await { + if read == 0 { + break; + } + if hold.wait_for(|held| !held).await.is_err() { + break; + } + if client_write.write_all(&buffer[..read]).await.is_err() { + break; + } + } + }); + } + }); + Self { + address, + begin_sent, + hold_replies, + } + } + } + + /// Counts this pool's server connections that are idle inside a + /// transaction, observed through a separate connection. + async fn idle_in_transaction(observer: &PgPool, application_name: &str) -> i64 { + sqlx::query_scalar( + "SELECT count(*) FROM pg_stat_activity \ + WHERE datname = current_database() AND application_name = $1 \ + AND state = 'idle in transaction'", + ) + .bind(application_name) + .fetch_one(observer) + .await + .unwrap() + } + + #[tokio::test(flavor = "multi_thread")] + async fn abandoned_begin_does_not_leave_a_connection_in_a_transaction() { + let url = std::env::var("RIVER_RUST_DATABASE_URL") + .expect("RIVER_RUST_DATABASE_URL must point at a disposable test database"); + let upstream: PgConnectOptions = url.parse().unwrap(); + let observer = PgPool::connect(&url).await.unwrap(); + let proxy = BeginProxy::start(&upstream).await; + let application_name = format!("river-begin-cancel-{}", std::process::id()); + // One connection, used without a liveness query that the held + // replies would stall. + let pool = PgPoolOptions::new() + .max_connections(1) + .test_before_acquire(false) + .connect_with( + upstream + .clone() + .host(&proxy.address.ip().to_string()) + .port(proxy.address.port()) + .application_name(&application_name), + ) + .await + .unwrap(); + + // Stop waiting for a begin after `BEGIN` reached the server but + // before its reply, as a `select!` or timeout around it would. + proxy.hold_replies.send_replace(true); + tokio::select! { + result = begin_postgres(&pool) => panic!("begin finished while replies were held: {result:?}"), + () = proxy.begin_sent.notified() => {} + } + proxy.hold_replies.send_replace(false); + + // The abandoned begin still finishes and rolls back, so the pool's + // only connection ends up idle outside a transaction. + tokio::time::timeout(Duration::from_secs(10), async { + loop { + // Use the connection so any unread reply is consumed, then + // check its server-side state once it's back in the pool. + let mut connection = pool.acquire().await.unwrap(); + sqlx::query("SELECT 1") + .execute(&mut *connection) + .await + .unwrap(); + drop(connection); + if idle_in_transaction(&observer, &application_name).await == 0 + && sqlx::query_scalar::<_, i64>( + "SELECT count(*) FROM pg_stat_activity \ + WHERE datname = current_database() AND application_name = $1 \ + AND state = 'idle'", + ) + .bind(&application_name) + .fetch_one(&observer) + .await + .unwrap() + == 1 + { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("an abandoned begin left its connection inside a transaction"); + + pool.close().await; + observer.close().await; + } +} diff --git a/rust/riverqueue/src/database/erased.rs b/rust/riverqueue/src/database/erased.rs new file mode 100644 index 000000000..f6f2fd264 --- /dev/null +++ b/rust/riverqueue/src/database/erased.rs @@ -0,0 +1,381 @@ +//! Type-erased database sources and executors. +//! +//! These types are reachable only through `riverqueue::__private`. + +use std::fmt; + +#[cfg(feature = "sqlite")] +use super::SqliteDatabase; +use crate::__private::DatabaseConnection; + +use super::{ + ClientDatabase, DatabaseExecutor, DatabaseKind, DatabaseMismatch, DatabasePool, + DatabaseTransactionExecutor, IntoDatabase, SchemaName, into_database, private, +}; +#[cfg(feature = "postgres")] +use super::{PostgresDatabase, PostgresReindexConfig}; + +/// A type-erased built-in River database source. +/// +/// This type is public only so the sealed [`IntoDatabase`] contract can be +/// composed across River's exact-version crates. Its backend representation is +/// intentionally private. +#[doc(hidden)] +#[derive(Clone)] +pub struct Database { + pub(super) inner: DatabaseInner, +} + +impl Database { + /// Erases a sealed built-in database source. + #[must_use] + pub fn from_source(database: D) -> Self { + into_database(database) + } + + #[cfg(feature = "postgres")] + pub(crate) fn extend_default_postgres_reindex_names( + &mut self, + names: impl IntoIterator>, + ) { + match &mut self.inner { + DatabaseInner::Postgres(source) => source.extend_default_reindex_names(names), + #[cfg(feature = "sqlite")] + DatabaseInner::Sqlite(_) => {} + } + } + + /// Returns the client's view of this database. + pub(crate) const fn client_database(&self) -> ClientDatabase<'_> { + match &self.inner { + #[cfg(feature = "postgres")] + DatabaseInner::Postgres(source) => ClientDatabase::Postgres(source), + #[cfg(feature = "sqlite")] + DatabaseInner::Sqlite(source) => ClientDatabase::Sqlite(source), + } + } + + /// Returns the configured backend kind. + #[must_use] + pub const fn kind(&self) -> DatabaseKind { + match &self.inner { + #[cfg(feature = "postgres")] + DatabaseInner::Postgres(_) => DatabaseKind::Postgres, + #[cfg(feature = "sqlite")] + DatabaseInner::Sqlite(_) => DatabaseKind::Sqlite, + } + } + + /// Whether the backend delivers notifications to listeners when a + /// transaction commits, like Go's `SupportsListener`. SQLite clients poll + /// a notification outbox instead, so operations River commits itself also + /// wake the local client directly. A PostgreSQL-compatible server without + /// `LISTEN`/`NOTIFY`, like YugabyteDB by default, has no listener once + /// detected. + pub(crate) fn supports_listener(&self) -> bool { + match &self.inner { + #[cfg(feature = "postgres")] + DatabaseInner::Postgres(source) => source.capabilities().supports_listen_notify(), + #[cfg(feature = "sqlite")] + DatabaseInner::Sqlite(_) => false, + } + } + + /// Whether committed notifications reach other clients: through + /// `LISTEN`/`NOTIFY` on PostgreSQL, assumed until a server without it is + /// detected, and through the notification outbox on SQLite. + pub(crate) fn delivers_notifications(&self) -> bool { + match &self.inner { + #[cfg(feature = "postgres")] + DatabaseInner::Postgres(source) => source.capabilities().supports_listen_notify(), + #[cfg(feature = "sqlite")] + DatabaseInner::Sqlite(_) => true, + } + } + + /// Returns the PostgreSQL server capabilities cache, or `None` for + /// another backend. + #[cfg(feature = "postgres")] + #[cfg_attr( + not(feature = "sqlite"), + expect( + clippy::unnecessary_wraps, + reason = "another backend may be compiled in" + ) + )] + pub(crate) const fn postgres_capabilities( + &self, + ) -> Option<&super::postgres_capabilities::CapabilitiesCache> { + match &self.inner { + DatabaseInner::Postgres(source) => Some(source.capabilities()), + #[cfg(feature = "sqlite")] + DatabaseInner::Sqlite(_) => None, + } + } + + /// Returns the PostgreSQL schema, or `None` for a backend without + /// PostgreSQL schemas. + #[must_use] + pub fn postgres_schema(&self) -> Option<&SchemaName> { + match &self.inner { + #[cfg(feature = "postgres")] + DatabaseInner::Postgres(source) => Some(source.schema()), + #[cfg(feature = "sqlite")] + DatabaseInner::Sqlite(_) => None, + } + } + + #[cfg(feature = "postgres")] + #[cfg_attr( + not(feature = "sqlite"), + expect( + clippy::unnecessary_wraps, + reason = "another backend may be compiled in" + ) + )] + pub(crate) fn postgres_reindex(&self) -> Option<&PostgresReindexConfig> { + match &self.inner { + DatabaseInner::Postgres(source) => Some(source.reindex()), + #[cfg(feature = "sqlite")] + DatabaseInner::Sqlite(_) => None, + } + } + + /// Erases and validates an executor before a backend operation uses it. + /// + /// # Errors + /// + /// Returns [`Error::DatabaseMismatch`] for an executor of another backend. + pub fn executor<'executor, E>( + &self, + executor: E, + ) -> Result, DatabaseMismatch> + where + E: DatabaseExecutor<'executor>, + { + let executor = private::DatabaseExecutorSealed::erase(executor); + self.check_kind(executor.kind())?; + Ok(executor) + } + + /// Erases and validates an actual SQLx transaction while preserving its + /// transaction-only capability for exact-version companion crates. + /// + /// # Errors + /// + /// Returns [`Error::DatabaseMismatch`] for a transaction of another + /// backend. + pub fn transaction<'executor, E>( + &self, + transaction: E, + ) -> Result, DatabaseMismatch> + where + E: DatabaseTransactionExecutor<'executor>, + { + Ok(ErasedTransaction { + connection: self.connection(transaction)?, + }) + } + + /// Borrows a caller-managed transaction's connection after checking that + /// it belongs to this database's backend. + pub(crate) fn connection<'executor, E>( + &self, + transaction: E, + ) -> Result, DatabaseMismatch> + where + E: DatabaseTransactionExecutor<'executor>, + { + let connection = private::DatabaseTransactionExecutorSealed::connection(transaction); + self.check_kind(connection.kind())?; + Ok(connection) + } + + fn check_kind(&self, actual: DatabaseKind) -> Result<(), DatabaseMismatch> { + if self.kind() != actual { + return Err(DatabaseMismatch { + actual, + expected: self.kind(), + }); + } + Ok(()) + } + + /// Returns a backend-specific borrowed pool for internal dispatch. + pub(crate) const fn pool(&self) -> DatabasePool<'_> { + match &self.inner { + #[cfg(feature = "postgres")] + DatabaseInner::Postgres(source) => DatabasePool::Postgres(source.pool()), + #[cfg(feature = "sqlite")] + DatabaseInner::Sqlite(source) => DatabasePool::Sqlite(source.pool()), + } + } +} + +impl fmt::Debug for Database { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("Database") + .field("kind", &self.kind()) + .field("postgres_schema", &self.postgres_schema()) + .finish_non_exhaustive() + } +} + +#[derive(Clone)] +pub(crate) enum DatabaseInner { + #[cfg(feature = "postgres")] + Postgres(PostgresDatabase), + #[cfg(feature = "sqlite")] + Sqlite(SqliteDatabase), +} + +/// A type-erased borrowed SQLx executor. +/// +/// The value is created only through the sealed executor contract. +#[doc(hidden)] +pub struct ErasedExecutor<'executor> { + pub(super) inner: ExecutorInner<'executor>, +} + +impl fmt::Debug for ErasedExecutor<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ErasedExecutor") + .finish_non_exhaustive() + } +} + +/// Transaction-preserving exact-version executor erasure. +/// +/// It holds only a transaction's connection, so an erased transaction can +/// never stand in for a pool. +#[doc(hidden)] +pub struct ErasedTransaction<'executor> { + pub(super) connection: DatabaseConnection<'executor>, +} + +impl fmt::Debug for ErasedTransaction<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ErasedTransaction") + .finish_non_exhaustive() + } +} + +impl ErasedTransaction<'_> { + /// Borrows the backend connection for exact-version SQL while retaining + /// the marker needed to call River's transaction-only methods later. + pub fn connection(&mut self) -> DatabaseConnection<'_> { + self.connection.reborrow() + } +} + +impl<'executor> ErasedExecutor<'executor> { + /// Converts a connection-backed executor for an exact-version extension. + /// Pool-backed executors return `None`. + #[must_use] + pub fn into_connection(self) -> Option> { + match self.inner { + ExecutorInner::Connection(connection) => Some(connection), + ExecutorInner::Pool(_) => None, + } + } + + /// Returns the executor's backend kind. + #[must_use] + pub const fn kind(&self) -> DatabaseKind { + match &self.inner { + ExecutorInner::Connection(connection) => connection.kind(), + ExecutorInner::Pool(kind) => *kind, + } + } +} + +/// What an erased executor borrows. A pool is recorded only by its backend: +/// River never runs an extension's statements on a pool it did not open. +pub(crate) enum ExecutorInner<'executor> { + Connection(DatabaseConnection<'executor>), + Pool(DatabaseKind), +} + +#[cfg(all(test, feature = "postgres"))] +mod tests { + use std::time::Duration; + + use sqlx::postgres::{PgConnectOptions, PgPool, PgPoolOptions}; + + use super::*; + use crate::{ + __private::database_with_default_postgres_reindex_names, database::PostgresReindexSchedule, + }; + + fn lazy_pool() -> PgPool { + PgPoolOptions::new().connect_lazy_with(PgConnectOptions::new()) + } + + #[tokio::test] + async fn add_on_reindex_names_are_kept_by_schedule_and_timeout_changes() { + let pool = lazy_pool(); + for config in [ + PostgresReindexConfig::default().with_timeout(Duration::from_secs(5)), + PostgresReindexConfig::default() + .with_schedule(PostgresReindexSchedule::Interval(Duration::from_hours(1))), + ] { + let database = database_with_default_postgres_reindex_names( + Database::from_source(PostgresDatabase::new(pool.clone()).with_reindex(config)), + ["add_on_hot_index"], + ); + let names = database.postgres_reindex().unwrap().index_names(); + assert!(names.contains(&"river_job_pkey".to_owned())); + assert!(names.contains(&"add_on_hot_index".to_owned())); + } + } + + #[tokio::test] + async fn add_on_reindex_names_extend_defaults_without_duplicates() { + let database = database_with_default_postgres_reindex_names( + Database::from_source(PostgresDatabase::new(lazy_pool())), + ["river_job_pkey", "add_on_hot_index"], + ); + let names = database.postgres_reindex().unwrap().index_names(); + assert!(names.contains(&"add_on_hot_index".to_owned())); + assert_eq!( + names + .iter() + .filter(|name| *name == "river_job_pkey") + .count(), + 1 + ); + } + + #[tokio::test] + async fn add_on_reindex_names_leave_explicit_index_names_unchanged() { + let pool = lazy_pool(); + for explicit in [vec!["custom_index"], vec![]] { + let configured = PostgresDatabase::new(pool.clone()).with_reindex( + PostgresReindexConfig::default() + .with_timeout(Duration::from_secs(5)) + .with_index_names(explicit.clone()), + ); + let database = database_with_default_postgres_reindex_names( + Database::from_source(configured), + ["add_on_hot_index"], + ); + assert_eq!(database.postgres_reindex().unwrap().index_names(), explicit); + } + } + + #[cfg(feature = "sqlite")] + #[tokio::test] + async fn add_on_reindex_names_leave_sqlite_unchanged() { + let sqlite_pool = sqlx::sqlite::SqlitePoolOptions::new() + .connect_lazy_with(sqlx::sqlite::SqliteConnectOptions::new()); + let database = database_with_default_postgres_reindex_names( + Database::from_source(SqliteDatabase::new(sqlite_pool)), + ["add_on_hot_index"], + ); + assert_eq!(database.kind(), DatabaseKind::Sqlite); + assert!(database.postgres_reindex().is_none()); + } +} diff --git a/rust/riverqueue/src/database/postgres_capabilities.rs b/rust/riverqueue/src/database/postgres_capabilities.rs new file mode 100644 index 000000000..9b9c74d3d --- /dev/null +++ b/rust/riverqueue/src/database/postgres_capabilities.rs @@ -0,0 +1,188 @@ +//! Features of a PostgreSQL-compatible server that River adapts to, like +//! River Go's `riverdriver.PostgresCapabilities`. +//! +//! YugabyteDB speaks PostgreSQL's protocol but has no `xmax` system column +//! and, unless configured for it, no `LISTEN`/`NOTIFY`. River detects the +//! server once per database and caches the result, so enabling Yugabyte's +//! notifications takes effect only for a new database value, such as after +//! a restart. + +use std::sync::{Arc, OnceLock}; + +use sqlx::{PgExecutor, Row}; + +/// Reads the server's product, version, and Yugabyte notification setting. +/// Functions are unqualified, as in River Go, so they resolve through the +/// connection's `search_path`. +const DETECT_SQL: &str = "SELECT \ + version()::text AS product, \ + current_setting('server_version_num')::int AS version_num, \ + coalesce(current_setting('yb_enable_listen_notify', true), 'off')::boolean AS yb_listen_notify_enabled"; + +/// How an insert that may conflict on its unique key tells a new row from +/// an existing one it returned instead. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum UniqueInsertMode { + /// The proposed row's metadata carries a random nonce, and a returned + /// row without it is an existing one. Used where `xmax` is unavailable. + MetadataNonce, + /// PostgreSQL 18's `OLD` row in `RETURNING`. + ReturningOld, + /// PostgreSQL's `xmax` system column, nonzero for an updated row. + Xmax, +} + +impl UniqueInsertMode { + /// Returns the SQL expression that is true for a returned existing row. + /// It's always false for [`MetadataNonce`](Self::MetadataNonce), which + /// compares nonces after the insert instead. + pub(crate) const fn sql(self) -> &'static str { + match self { + Self::MetadataNonce => "false", + Self::ReturningOld => "(OLD.id IS NOT NULL)", + Self::Xmax => "(xmax != 0)", + } + } +} + +/// Features detected from a PostgreSQL-compatible server. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct PostgresCapabilities { + /// Whether `pg_notify` delivers notifications to listeners. Without it, + /// River skips notifications and clients poll instead. + pub(crate) supports_listen_notify: bool, + pub(crate) unique_insert_mode: UniqueInsertMode, +} + +impl PostgresCapabilities { + /// Derives capabilities from the server's `version()` text, its + /// `server_version_num`, and Yugabyte's `yb_enable_listen_notify` + /// setting, which is off when absent. + pub(crate) fn new(product: &str, version: i32, yb_listen_notify_enabled: bool) -> Self { + let yugabyte = is_yugabyte(product); + Self { + // Yugabyte's notifications need 2025.2.3 or later with + // `ysql_yb_enable_listen_notify=true` on masters and tservers. + supports_listen_notify: !yugabyte || yb_listen_notify_enabled, + unique_insert_mode: if yugabyte { + UniqueInsertMode::MetadataNonce + } else if version >= 180_000 { + UniqueInsertMode::ReturningOld + } else { + UniqueInsertMode::Xmax + }, + } + } + + /// Detects the capabilities of the server `executor` is connected to. + pub(crate) async fn detect<'e>(executor: impl PgExecutor<'e>) -> Result { + let row = sqlx::query(DETECT_SQL).fetch_one(executor).await?; + Ok(Self::new( + row.try_get("product")?, + row.try_get("version_num")?, + row.try_get("yb_listen_notify_enabled")?, + )) + } +} + +fn is_yugabyte(product: &str) -> bool { + let product = product.to_lowercase(); + product.contains("-yb") || product.contains("yugabyte") +} + +/// Capabilities detected for one database, shared by its clones. +#[derive(Clone, Debug, Default)] +pub(crate) struct CapabilitiesCache(Arc>); + +impl CapabilitiesCache { + /// Returns the cached capabilities, detecting them with `executor` the + /// first time. Concurrent first callers may each detect; the first result + /// stored wins. No lock is held while detecting, since the caller may hold + /// the pool's only connection. + pub(crate) async fn load<'e>( + &self, + executor: impl PgExecutor<'e>, + ) -> Result { + if let Some(capabilities) = self.0.get() { + return Ok(*capabilities); + } + let detected = PostgresCapabilities::detect(executor).await?; + Ok(*self.0.get_or_init(|| detected)) + } + + /// Returns the capabilities from `cache`, or detects them with + /// `executor` each time without one. + pub(crate) async fn load_or_detect<'e>( + cache: Option<&Self>, + executor: impl PgExecutor<'e>, + ) -> Result { + match cache { + Some(cache) => cache.load(executor).await, + None => PostgresCapabilities::detect(executor).await, + } + } + + /// Whether notifications are delivered, assuming they are until the + /// server has been detected, like River Go before `InitDriver`. + pub(crate) fn supports_listen_notify(&self) -> bool { + self.0 + .get() + .is_none_or(|capabilities| capabilities.supports_listen_notify) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn detects_yugabyte_and_postgres_versions() { + let postgres_17 = PostgresCapabilities::new( + "PostgreSQL 17.4 on aarch64-apple-darwin, compiled by clang", + 170_004, + false, + ); + assert!(postgres_17.supports_listen_notify); + assert_eq!(postgres_17.unique_insert_mode, UniqueInsertMode::Xmax); + + let postgres_18 = PostgresCapabilities::new("PostgreSQL 18.1", 180_001, false); + assert!(postgres_18.supports_listen_notify); + assert_eq!( + postgres_18.unique_insert_mode, + UniqueInsertMode::ReturningOld + ); + + for product in [ + "PostgreSQL 15.12-YB-2025.2.1.0-b1 on x86_64", + "YugabyteDB 2025.2.3.0", + ] { + let without = PostgresCapabilities::new(product, 150_012, false); + assert!(!without.supports_listen_notify, "{product}"); + assert_eq!( + without.unique_insert_mode, + UniqueInsertMode::MetadataNonce, + "{product}" + ); + assert!( + PostgresCapabilities::new(product, 150_012, true).supports_listen_notify, + "{product}" + ); + } + } + + #[test] + fn unknown_capabilities_assume_notifications() { + let cache = CapabilitiesCache::default(); + assert!(cache.supports_listen_notify()); + cache + .0 + .set(PostgresCapabilities::new( + "PostgreSQL 15.12-YB-2025.2.1.0", + 150_012, + false, + )) + .unwrap(); + assert!(!cache.supports_listen_notify()); + assert!(cache.clone().0.get().is_some()); + } +} diff --git a/rust/riverqueue/src/database/sqlite.rs b/rust/riverqueue/src/database/sqlite.rs new file mode 100644 index 000000000..856eeb052 --- /dev/null +++ b/rust/riverqueue/src/database/sqlite.rs @@ -0,0 +1,2827 @@ +//! SQLite operation primitives. +//! +//! SQLite deliberately stays behind River's sealed database boundary. These +//! operations mirror River's storage semantics without exposing SQLx executor +//! types through the public API. Multi-step operations are expressed as small +//! primitives so their caller can compose them in one SQLite transaction. + +#![allow( + clippy::needless_raw_string_hashes, + reason = "consistent SQL delimiters make large dialect-specific statements easier to audit" +)] +#![allow( + clippy::struct_field_names, + reason = "leader_id is the cross-language River protocol field name" +)] + +use std::{ + collections::{HashMap, VecDeque}, + time::Duration, +}; + +use chrono::{DateTime, SubsecRound, Utc}; +use serde_json::{Map, Value, value::RawValue}; +use sqlx::{AssertSqlSafe, FromRow, QueryBuilder, Sqlite, SqliteConnection}; + +use sqlx::sqlite::SqliteRow; + +use crate::{ + AttemptError, JobMetadata, JobRow, JobState, METADATA_KEY_UNIQUE_NONCE, Queue, + client::{DecodedJob, FieldErrors, UndecodableJob, go_time_json, saturating_i16, tolerant_row}, + query::{JobListKeyset, JobListSqlPart}, +}; + +/// Job columns as River reads them. A JSON column can be changed out of band +/// to text that isn't valid JSON, which makes SQLite's JSON functions fail +/// with "malformed JSON", including the `json()` that converts the stored +/// JSONB to text. Like River Go, such a value is returned as is instead, so +/// the row decodes as undecodable rather than failing the whole statement, +/// and a fetched job with one has its attempt failed. River writes these +/// columns as JSONB, so only text values are checked. +pub(crate) const JOB_COLUMNS: &str = r#" + id, + attempt, + attempted_at, + CASE WHEN typeof(attempted_by) = 'text' AND NOT json_valid(attempted_by) THEN attempted_by ELSE json(attempted_by) END AS attempted_by, + created_at, + CASE WHEN typeof(args) = 'text' AND NOT json_valid(args) THEN args ELSE json(args) END AS encoded_args, + CASE WHEN typeof(errors) = 'text' AND NOT json_valid(errors) THEN errors ELSE json(errors) END AS errors, + finalized_at, + kind, + max_attempts, + CASE WHEN typeof(metadata) = 'text' AND NOT json_valid(metadata) THEN metadata ELSE json(metadata) END AS metadata, + priority, + queue, + scheduled_at, + state, + CASE WHEN typeof(tags) = 'text' AND NOT json_valid(tags) THEN tags ELSE json(tags) END AS tags, + unique_key, + unique_states +"#; + +/// Whether a stored JSON column holds text that isn't valid JSON. See +/// [`JOB_COLUMNS`]. +macro_rules! invalid_json { + ($column:literal) => { + concat!( + "(typeof(", + $column, + ") = 'text' AND NOT json_valid(", + $column, + "))" + ) + }; +} + +/// A job's `cancel_attempted_at`, treating metadata that isn't valid JSON as +/// not having one. Like River Go, a key with a JSON `null` value still counts, +/// which `->` distinguishes from a missing key. +const CANCEL_ATTEMPTED_AT: &str = concat!( + "(CASE WHEN NOT ", + invalid_json!("metadata"), + " THEN metadata -> 'cancel_attempted_at' END)" +); + +/// Appends the attempt error bound next to `errors`. `errors` is always an +/// array unless it's been changed out of band. Like River Go, any other value +/// is wrapped in an array, as a string if it isn't valid JSON, so the new +/// error is still appended without losing it. +const ERRORS_APPENDED: &str = concat!( + "CASE WHEN ", + invalid_json!("errors"), + " THEN jsonb(json_array(errors, json(?))) \ + WHEN coalesce(json_type(errors), 'array') <> 'array' \ + THEN jsonb(json_array(json(errors), json(?))) \ + ELSE jsonb(json_insert(json(coalesce(errors, jsonb('[]'))), '$[#]', json(?))) END" +); + +const QUEUE_COLUMNS: &str = r#" + created_at, + json(metadata) AS metadata, + name, + paused_at, + updated_at +"#; + +/// A short poll interval keeps local wakeups responsive while queue fetch +/// polling remains the durable recovery path. +pub(crate) const DEFAULT_NOTIFICATION_POLL_INTERVAL: Duration = Duration::from_millis(100); + +/// Maximum number of outbox rows a notification listener reads at once, +/// matching River Go's SQLite listener. +pub(crate) const NOTIFICATION_BATCH_SIZE: i64 = 256; + +#[derive(Debug, thiserror::Error)] +pub(crate) enum BackendError { + #[error("invalid SQLite River row: {0}")] + InvalidRow(String), + #[error(transparent)] + Sqlx(#[from] sqlx::Error), + #[error("invalid SQLite River JSON")] + Json(#[from] serde_json::Error), +} + +/// Reports SQLite failures the way PostgreSQL's reach callers: as SQLx's own +/// error, with stored values River can't decode as [`sqlx::Error::Decode`]. +impl From for sqlx::Error { + fn from(error: BackendError) -> Self { + match error { + BackendError::Sqlx(error) => error, + BackendError::InvalidRow(message) => { + Self::Decode(format!("invalid SQLite River row: {message}").into()) + } + BackendError::Json(error) => Self::Decode(Box::new(error)), + } + } +} + +#[derive(Clone, Debug)] +pub(crate) struct InsertJob<'a> { + pub attempted_at: Option>, + pub attempted_by: &'a [String], + pub attempt: i16, + /// When unset, SQLite's `datetime('now', 'subsec')`, like Go's driver. + pub created_at: Option>, + pub encoded_args: &'a serde_json::value::RawValue, + pub errors: &'a [AttemptError], + pub finalized_at: Option>, + pub id: Option, + pub kind: &'a str, + pub max_attempts: i16, + pub metadata: &'a JobMetadata, + pub priority: i16, + pub queue: &'a str, + /// When unset, SQLite's `datetime('now', 'subsec')`, like Go's driver. + pub scheduled_at: Option>, + pub state: JobState, + pub tags: &'a [String], + pub unique_key: Option<&'a [u8]>, + pub unique_nonce: Option<&'a str>, + pub unique_states: Option, +} + +#[derive(Clone, Debug)] +pub(crate) struct InsertedJob { + pub job: JobRow, + pub unique_skipped_as_duplicate: bool, +} + +#[derive(Clone, Debug)] +pub(crate) struct ClaimJobs<'a> { + pub client_id: &'a str, + /// Kinds to claim, or `None` for every kind. Ignored when claiming + /// selected IDs. + pub kinds: Option<&'a [String]>, + pub limit: i32, + pub max_attempted_by: i32, + pub now: DateTime, + pub queue: &'a str, +} + +#[derive(Clone, Debug, Default)] +pub(crate) struct ListJobs<'a> { + /// Excludes running jobs before applying the limit, as bulk deletion does. + pub exclude_running: bool, + pub ids: &'a [i64], + pub keyset: JobListKeyset, + pub kinds: &'a [&'a str], + pub limit: i32, + pub metadata: Option<&'a Map>, + pub priorities: &'a [i16], + pub queues: &'a [&'a str], + pub states: &'a [JobState], + pub tags_all: &'a [&'a str], + pub tags_any: &'a [&'a str], +} + +#[derive(Clone, Debug)] +pub(crate) struct CompleteJob<'a> { + pub attempt: Option, + pub error: Option<&'a AttemptError>, + pub finalized_at: Option>, + pub id: i64, + pub metadata_updates: Option<&'a Map>, + pub now: DateTime, + pub scheduled_at: Option>, + pub state: JobState, +} + +#[derive(Clone, Debug)] +pub(crate) struct RescueJob<'a> { + pub error: &'a AttemptError, + pub finalized_at: Option>, + pub id: i64, + pub scheduled_at: DateTime, + pub state: JobState, + /// The rescue applies only to a job still running from before this time. + pub stuck_horizon: DateTime, +} + +/// Job cleaner deletion horizons; `None` keeps that state indefinitely. +#[derive(Clone, Debug)] +pub(crate) struct CleanupJobs<'a> { + pub cancelled_before: Option>, + pub completed_before: Option>, + pub discarded_before: Option>, + pub limit: i32, + pub queues_excluded: &'a [&'a str], + pub queues_included: Option<&'a [&'a str]>, +} + +#[derive(Clone, Debug, PartialEq)] +pub(crate) struct Leader { + pub elected_at: DateTime, + pub expires_at: DateTime, + pub leader_id: String, +} + +#[derive(Clone, Debug, FromRow, PartialEq)] +pub(crate) struct Notification { + pub id: i64, + pub payload: String, + pub topic: String, +} + +#[derive(Clone, Copy, Debug)] +pub(crate) struct NotificationInput<'a> { + pub payload: &'a str, + pub topic: &'a str, +} + +// River Go stores `attempt`, `max_attempts`, and `priority` as native +// integers, so they decode as `i64` and saturate into `JobRow`'s fields. +// JSON columns are read as bytes, since a value changed out of band might not +// even be UTF-8. +#[derive(Clone, Debug, FromRow)] +struct JobRecord { + attempt: i64, + attempted_at: Option>, + attempted_by: Option>, + created_at: DateTime, + encoded_args: Vec, + errors: Option>, + finalized_at: Option>, + id: i64, + kind: String, + max_attempts: i64, + metadata: Vec, + priority: i64, + queue: String, + scheduled_at: DateTime, + state: String, + tags: Vec, + unique_key: Option>, + unique_states: Option, +} + +impl JobRecord { + /// Decodes the row, failing if any field can't be decoded. + fn into_job(self) -> Result { + let id = self.id; + self.decode() + .map_err(|job| BackendError::InvalidRow(format!("job {id}: {}", job.error))) + } + + /// Decodes the row, keeping the fields that can be decoded when others + /// can't, like River Go. JSON columns can be changed to any shape, so + /// each decodes on its own. River Go reads the metadata as raw JSON, but a + /// [`JobRow`] can only represent an object. + fn decode(self) -> DecodedJob { + let unidentifiable = |error: String| UndecodableJob { + columns: Vec::new(), + error: format!("job {}: {error}", self.id), + row: None, + }; + let state = JobState::try_from(self.state.as_str()) + .map_err(|error| unidentifiable(error.to_string()))?; + + let mut errors = FieldErrors::default(); + // Like River Go, args that aren't valid JSON are left empty, as + // `null`, so the job isn't worked with them. + let encoded_args = errors.field_or_else( + "args", + json_text(self.encoded_args) + .and_then(|text| RawValue::from_string(text).map_err(|error| error.to_string())), + || RawValue::from_string("null".to_owned()).expect("null is valid JSON"), + ); + let attempted_by = errors.field( + "attempted_by", + self.attempted_by + .map(json_text) + .transpose() + .and_then(|text| { + decode_json_strings(text.as_deref()).map_err(|error| error.to_string()) + }), + ); + let attempt_errors = errors.field( + "errors", + self.errors.map(json_text).transpose().and_then(|text| { + text.as_deref() + .map_or_else(|| Ok(Vec::new()), AttemptError::from_json_array_lenient) + .map_err(|error| error.to_string()) + }), + ); + let metadata = errors.field( + "metadata", + json_text(self.metadata).and_then(|text| { + text.parse::() + .map_err(|error| error.to_string()) + }), + ); + let tags = errors.field( + "tags", + json_text(self.tags).and_then(|text| { + decode_json_strings(Some(&text)).map_err(|error| error.to_string()) + }), + ); + let unique_states = errors.field( + "unique_states", + self.unique_states.map(decode_unique_states).transpose(), + ); + errors.finish(JobRow { + id: self.id, + attempt: saturating_i16(self.attempt), + attempted_at: self.attempted_at, + attempted_by, + created_at: self.created_at, + encoded_args, + errors: attempt_errors, + finalized_at: self.finalized_at, + kind: self.kind, + max_attempts: saturating_i16(self.max_attempts), + metadata, + priority: saturating_i16(self.priority), + queue: self.queue, + scheduled_at: self.scheduled_at, + state, + tags, + unique_key: self.unique_key, + unique_states, + }) + } +} + +/// Decodes a row selected with [`JOB_COLUMNS`] on its own. +pub(crate) fn decode_job_row(row: &SqliteRow) -> DecodedJob { + JobRecord::from_row(row) + .map_err(|error| UndecodableJob { + columns: Vec::new(), + error: error.to_string(), + row: None, + })? + .decode() +} + +#[derive(Clone, Debug, FromRow)] +struct QueueRecord { + created_at: DateTime, + metadata: String, + name: String, + paused_at: Option>, + updated_at: DateTime, +} + +impl QueueRecord { + fn into_queue(self) -> Result { + let metadata: Value = serde_json::from_str(&self.metadata)?; + let metadata = metadata.as_object().cloned().ok_or_else(|| { + BackendError::InvalidRow(format!("queue {:?} metadata is not an object", self.name)) + })?; + Ok(Queue { + created_at: self.created_at, + metadata, + metadata_text: self.metadata, + name: self.name, + paused_at: self.paused_at, + updated_at: self.updated_at, + }) + } +} + +#[derive(Clone, Debug, FromRow)] +struct LeaderRecord { + elected_at: DateTime, + expires_at: DateTime, + leader_id: String, +} + +impl From for Leader { + fn from(record: LeaderRecord) -> Self { + Self { + elected_at: record.elected_at, + expires_at: record.expires_at, + leader_id: record.leader_id, + } + } +} + +/// Converts a JSON column read as bytes to text. +fn json_text(bytes: Vec) -> Result { + String::from_utf8(bytes).map_err(|error| error.to_string()) +} + +fn decode_json_or_default(encoded: Option<&str>) -> Result +where + T: serde::de::DeserializeOwned + Default, +{ + encoded.map_or_else(|| Ok(T::default()), serde_json::from_str) +} + +/// Decodes a JSON array of strings like Go's `encoding/json` decodes a +/// `[]string`: `null` is empty, as is a `null` element. +fn decode_json_strings(encoded: Option<&str>) -> Result, serde_json::Error> { + Ok( + decode_json_or_default::>>>(encoded)? + .unwrap_or_default() + .into_iter() + .map(Option::unwrap_or_default) + .collect(), + ) +} + +fn decode_unique_states(bits: i64) -> Result, String> { + let bits = u8::try_from(bits).map_err(|_| format!("value out of range for byte: {bits}"))?; + Ok(JobState::ALL + .iter() + .copied() + .filter(|state| bits & state.unique_bit() != 0) + .collect()) +} + +pub(crate) fn sqlite_time(time: DateTime) -> String { + time.round_subsecs(3) + .format("%Y-%m-%d %H:%M:%S%.3f") + .to_string() +} + +fn sqlite_time_optional(time: Option>) -> Option { + time.map(sqlite_time) +} + +fn sqlite_ttl(ttl: Duration) -> String { + format!("{:.3} seconds", ttl.as_secs_f64()) +} + +pub(crate) async fn insert( + connection: &mut SqliteConnection, + params: &InsertJob<'_>, +) -> Result { + let mut metadata = params.metadata.clone(); + if let Some(nonce) = params.unique_nonce { + metadata.insert(METADATA_KEY_UNIQUE_NONCE, Value::String(nonce.to_owned()))?; + } + let attempted_by = serde_json::to_string(params.attempted_by)?; + let encoded_args = params.encoded_args.get(); + let errors = serde_json::to_string(params.errors)?; + let metadata = serde_json::to_string(&metadata)?; + let tags = serde_json::to_string(params.tags)?; + let sql = format!( + r#" + INSERT INTO river_job ( + id, args, attempt, attempted_at, attempted_by, created_at, errors, + finalized_at, kind, max_attempts, metadata, priority, queue, + scheduled_at, state, tags, unique_key, unique_states + ) VALUES ( + ?, jsonb(?), ?, ?, CASE WHEN ? = '[]' THEN NULL ELSE jsonb(?) END, + coalesce(?, datetime('now', 'subsec')), + CASE WHEN ? = '[]' THEN NULL ELSE jsonb(?) END, ?, ?, ?, + jsonb(?), ?, ?, coalesce(?, datetime('now', 'subsec')), ?, jsonb(?), ?, ? + ) + ON CONFLICT (unique_key) + WHERE unique_key IS NOT NULL + AND unique_states IS NOT NULL + AND CASE state + WHEN 'available' THEN unique_states & (1 << 0) + WHEN 'cancelled' THEN unique_states & (1 << 1) + WHEN 'completed' THEN unique_states & (1 << 2) + WHEN 'discarded' THEN unique_states & (1 << 3) + WHEN 'pending' THEN unique_states & (1 << 4) + WHEN 'retryable' THEN unique_states & (1 << 5) + WHEN 'running' THEN unique_states & (1 << 6) + WHEN 'scheduled' THEN unique_states & (1 << 7) + ELSE 0 + END >= 1 + -- Keep the existing kind, which may differ under `exclude_kind`. + DO UPDATE SET kind = river_job.kind + RETURNING {JOB_COLUMNS} + "# + ); + let record = sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) + .bind(params.id) + .bind(encoded_args) + .bind(params.attempt) + .bind(sqlite_time_optional(params.attempted_at)) + .bind(&attempted_by) + .bind(&attempted_by) + .bind(sqlite_time_optional(params.created_at)) + .bind(&errors) + .bind(&errors) + .bind(sqlite_time_optional(params.finalized_at)) + .bind(params.kind) + .bind(params.max_attempts) + .bind(metadata) + .bind(params.priority) + .bind(params.queue) + .bind(sqlite_time_optional(params.scheduled_at)) + .bind(params.state.as_str()) + .bind(tags) + .bind(params.unique_key) + .bind(params.unique_states.map(i64::from)) + .fetch_one(&mut *connection) + .await?; + let job = record.into_job()?; + let unique_skipped_as_duplicate = params.unique_nonce.is_some_and(|nonce| { + job.metadata + .get::(METADATA_KEY_UNIQUE_NONCE) + .ok() + .flatten() + .as_deref() + != Some(nonce) + }); + Ok(InsertedJob { + job, + unique_skipped_as_duplicate, + }) +} + +/// Claims due jobs. The claim commits even when a row cannot be decoded, so +/// each row is decoded separately and the caller records a failed attempt for +/// any undecodable row instead of stranding the whole batch as running. +pub(crate) async fn claim( + connection: &mut SqliteConnection, + params: &ClaimJobs<'_>, +) -> Result, BackendError> { + if params.limit <= 0 { + return Ok(Vec::new()); + } + + let now = sqlite_time(params.now); + let mut query = QueryBuilder::::new( + r#" + UPDATE river_job + SET + attempt = attempt + 1, + attempted_at = "#, + ); + query.push_bind(&now); + query.push(", attempted_by = "); + push_attempted_by_appended(&mut query, params); + query.push( + r#", + state = 'running' + WHERE id IN ( + SELECT river_job.id + FROM river_job + WHERE queue = "#, + ); + query.push_bind(params.queue); + query.push(" AND scheduled_at <= "); + query.push_bind(&now); + query.push( + r#" + AND state = 'available' + AND NOT EXISTS ( + SELECT 1 + FROM river_queue + WHERE river_queue.name = river_job.queue + AND river_queue.paused_at IS NOT NULL + )"#, + ); + if let Some(kinds) = params.kinds { + query.push(" AND kind IN (SELECT value FROM json_each("); + query.push_bind(serde_json::to_string(kinds)?); + query.push("))"); + } + query.push( + r#" + ORDER BY priority ASC, scheduled_at ASC, id ASC + LIMIT "#, + ); + query.push_bind(params.limit); + query.push(format!(") RETURNING {JOB_COLUMNS}")); + + let rows = query.build().fetch_all(&mut *connection).await?; + Ok(rows.iter().map(decode_job_row).collect()) +} + +/// Pushes `attempted_by` with this client appended, keeping at most +/// `max_attempted_by` entries. Like River Go, an `attempted_by` that isn't +/// valid JSON is left in place, since `json_each` would fail on it and with +/// it the whole claim; the job is then claimed as undecodable. `json_each`'s +/// input is guarded too, as River Go guards it for SQLite implementations that +/// evaluate a `CASE` branch that isn't taken. +fn push_attempted_by_appended(query: &mut QueryBuilder, params: &ClaimJobs<'_>) { + query.push(concat!( + "CASE WHEN ", + invalid_json!("attempted_by"), + r#" THEN attempted_by ELSE jsonb(json_insert( + ( + SELECT jsonb_group_array(value) + FROM ( + SELECT value FROM ( + SELECT key, value + FROM json_each(CASE WHEN "#, + invalid_json!("attempted_by"), + r#" THEN jsonb('[]') ELSE coalesce(attempted_by, jsonb('[]')) END) + ORDER BY key DESC + LIMIT "#, + )); + query.push_bind(params.max_attempted_by.saturating_sub(1)); + query.push( + r#" + ) ORDER BY key ASC + ) + ), + '$[#]', + "#, + ); + query.push_bind(params.client_id.to_owned()); + query.push(")) END"); +} + +/// Claims exactly the IDs selected by an exact-version extension. +/// +/// The caller keeps selection and this update in one transaction. Eligibility +/// beyond the final running-state guard is deliberately the selector's +/// responsibility, matching the PostgreSQL interception path. +pub(crate) async fn claim_selected( + connection: &mut SqliteConnection, + params: &ClaimJobs<'_>, + ids: &[i64], +) -> Result, BackendError> { + if params.limit <= 0 || ids.is_empty() { + return Ok(Vec::new()); + } + + let now = sqlite_time(params.now); + let mut query = QueryBuilder::::new( + r#" + UPDATE river_job + SET + attempt = attempt + 1, + attempted_at = "#, + ); + query.push_bind(&now); + query.push(", attempted_by = "); + push_attempted_by_appended(&mut query, params); + query.push( + r#", + state = 'running' + WHERE state = 'available' AND id IN ("#, + ); + { + let mut separated = query.separated(", "); + for id in ids { + separated.push_bind(id); + } + } + query.push(format!(") RETURNING {JOB_COLUMNS}")); + + let rows = query.build().fetch_all(&mut *connection).await?; + Ok(rows.iter().map(decode_job_row).collect()) +} + +pub(crate) async fn get( + connection: &mut SqliteConnection, + id: i64, +) -> Result, BackendError> { + let sql = format!("SELECT {JOB_COLUMNS} FROM river_job WHERE id = ? LIMIT 1"); + sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) + .bind(id) + .fetch_optional(&mut *connection) + .await? + .map(JobRecord::into_job) + .transpose() +} + +#[allow( + clippy::too_many_lines, + reason = "the bound-only query builder keeps all list filters and keyset ordering auditable" +)] +pub(crate) async fn list( + connection: &mut SqliteConnection, + params: &ListJobs<'_>, +) -> Result, BackendError> { + if params.limit <= 0 { + return Ok(Vec::new()); + } + + let mut query = + QueryBuilder::::new(format!("SELECT {JOB_COLUMNS} FROM river_job WHERE true")); + if let Some(after) = params.keyset.after_sql() { + query.push(" AND "); + for part in after { + match part { + JobListSqlPart::AfterId => { + query.push_bind(params.keyset.after_id().expect("cursor has an ID")); + } + JobListSqlPart::AfterTime => { + query.push_bind(sqlite_time( + params.keyset.after_time().expect("cursor has a time"), + )); + } + JobListSqlPart::Sql(sql) => { + query.push(sql); + } + } + } + } + if params.exclude_running { + query.push(" AND state != 'running'"); + } + if !params.ids.is_empty() { + query.push(" AND id IN ("); + let mut separated = query.separated(", "); + for id in params.ids { + separated.push_bind(id); + } + separated.push_unseparated(")"); + } + if !params.kinds.is_empty() { + query.push(" AND kind IN ("); + let mut separated = query.separated(", "); + for kind in params.kinds { + separated.push_bind(kind); + } + separated.push_unseparated(")"); + } + if !params.priorities.is_empty() { + query.push(" AND priority IN ("); + let mut separated = query.separated(", "); + for priority in params.priorities { + separated.push_bind(priority); + } + separated.push_unseparated(")"); + } + if !params.queues.is_empty() { + query.push(" AND queue IN ("); + let mut separated = query.separated(", "); + for queue in params.queues { + separated.push_bind(queue); + } + separated.push_unseparated(")"); + } + if !params.states.is_empty() { + query.push(" AND state IN ("); + let mut separated = query.separated(", "); + for state in params.states { + separated.push_bind(state.as_str()); + } + separated.push_unseparated(")"); + } + if let Some(metadata) = params.metadata { + push_contains_object(&mut query, "metadata", metadata, &mut 0); + } + for tag in params.tags_all { + query + .push(" AND EXISTS (SELECT 1 FROM json_each(json(tags)) WHERE value = ") + .push_bind(tag) + .push(")"); + } + if !params.tags_any.is_empty() { + query.push(" AND EXISTS (SELECT 1 FROM json_each(json(tags)) WHERE value IN ("); + let mut separated = query.separated(", "); + for tag in params.tags_any { + separated.push_bind(tag); + } + separated.push_unseparated("))"); + } + query.push(" ORDER BY ").push(params.keyset.order_sql()); + query.push(" LIMIT ").push_bind(params.limit); + + let records = query + .build_query_as::() + .fetch_all(&mut *connection) + .await?; + records.into_iter().map(JobRecord::into_job).collect() +} + +/// Pushes conditions that the JSON object `source` (an SQL expression) +/// contains `fragment`, with PostgreSQL's `jsonb @>` semantics: each key must +/// be present with a contained value. Keys and scalars are bound, so no JSON +/// path quoting is involved. +fn push_contains_object( + query: &mut QueryBuilder, + source: &str, + fragment: &Map, + aliases: &mut usize, +) { + for (key, value) in fragment { + let alias = format!("contains_{aliases}"); + *aliases += 1; + query.push(format!( + " AND EXISTS (SELECT 1 FROM json_each({source}) AS {alias} WHERE {alias}.key = " + )); + query.push_bind(key.clone()); + query.push(" AND "); + push_contains_value(query, &alias, value, aliases); + query.push(")"); + } +} + +/// Pushes a condition that the `json_each` row `alias` contains `fragment`. +/// Like PostgreSQL, scalars match only scalars of the same type (numbers +/// compare numerically), objects match objects containing every key, and an +/// array matches an array containing each of its elements. +fn push_contains_value( + query: &mut QueryBuilder, + alias: &str, + fragment: &Value, + aliases: &mut usize, +) { + match fragment { + Value::Null => { + query.push(format!("{alias}.type = 'null'")); + } + Value::Bool(true) => { + query.push(format!("{alias}.type = 'true'")); + } + Value::Bool(false) => { + query.push(format!("{alias}.type = 'false'")); + } + Value::Number(number) => { + query.push(format!( + "{alias}.type IN ('integer', 'real') AND {alias}.value = " + )); + if let Some(integer) = number.as_i64() { + query.push_bind(integer); + } else { + query.push_bind(number.as_f64().unwrap_or(f64::NAN)); + } + } + Value::String(string) => { + query.push(format!("{alias}.type = 'text' AND {alias}.value = ")); + query.push_bind(string.clone()); + } + Value::Object(object) => { + query.push(format!("{alias}.type = 'object'")); + push_contains_object(query, &format!("{alias}.value"), object, aliases); + } + Value::Array(elements) => { + query.push(format!("{alias}.type = 'array'")); + for element in elements { + let element_alias = format!("contains_{aliases}"); + *aliases += 1; + query.push(format!( + " AND EXISTS (SELECT 1 FROM json_each({alias}.value) AS {element_alias} \ + WHERE " + )); + push_contains_value(query, &element_alias, element, aliases); + query.push(")"); + } + } + } +} + +/// Marks cancellation intent. A running job stays running so its worker can +/// observe cancellation; any other non-finalized job is finalized immediately. +pub(crate) async fn cancel( + connection: &mut SqliteConnection, + id: i64, + now: DateTime, +) -> Result, BackendError> { + let cancel_attempted_at = go_time_json(now); + let sql = format!( + r#" + UPDATE river_job + SET + state = CASE WHEN state = 'running' THEN state ELSE 'cancelled' END, + finalized_at = CASE WHEN state = 'running' THEN finalized_at ELSE ? END, + metadata = jsonb_set(metadata, '$.cancel_attempted_at', ?) + WHERE id = ? + AND state NOT IN ('cancelled', 'completed', 'discarded') + AND finalized_at IS NULL + RETURNING {JOB_COLUMNS} + "# + ); + sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) + .bind(sqlite_time(now)) + .bind(cancel_attempted_at) + .bind(id) + .fetch_optional(&mut *connection) + .await? + .map(JobRecord::into_job) + .transpose() +} + +/// Deletes a non-running job. A missing result is intentionally ambiguous; +/// callers distinguish not-found from running inside the same transaction. +pub(crate) async fn delete( + connection: &mut SqliteConnection, + id: i64, +) -> Result, BackendError> { + let sql = format!( + "DELETE FROM river_job WHERE id = ? AND state != 'running' RETURNING {JOB_COLUMNS}" + ); + sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) + .bind(id) + .fetch_optional(&mut *connection) + .await? + .map(JobRecord::into_job) + .transpose() +} + +/// Makes a non-running job immediately available. A missing result is +/// intentionally ambiguous; callers fetch in the same transaction to +/// distinguish running, already-available, and not-found jobs. +pub(crate) async fn retry( + connection: &mut SqliteConnection, + id: i64, + now: DateTime, +) -> Result, BackendError> { + let sql = format!( + r#" + UPDATE river_job + SET + state = 'available', + max_attempts = CASE + WHEN attempt = max_attempts THEN max_attempts + 1 + ELSE max_attempts + END, + finalized_at = NULL, + scheduled_at = ? + WHERE id = ? + AND state != 'running' + AND (state != 'available' OR scheduled_at > ?) + RETURNING {JOB_COLUMNS} + "# + ); + sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) + .bind(sqlite_time(now)) + .bind(id) + .bind(sqlite_time(now)) + .fetch_optional(&mut *connection) + .await? + .map(JobRecord::into_job) + .transpose() +} + +pub(crate) async fn complete( + connection: &mut SqliteConnection, + params: &CompleteJob<'_>, +) -> Result, BackendError> { + complete_decoded(connection, params) + .await? + .map(|row| row.map_err(|job| BackendError::InvalidRow(job.error))) + .transpose() +} + +/// Sets a running job's state, decoding the returned row separately so a +/// malformed row cannot fail the surrounding completion transaction. +pub(crate) async fn complete_decoded( + connection: &mut SqliteConnection, + params: &CompleteJob<'_>, +) -> Result, BackendError> { + let error = params + .error + .map(serde_json::to_string) + .transpose()? + .unwrap_or_else(|| "{}".to_owned()); + let metadata = params + .metadata_updates + .map(serde_json::to_string) + .transpose()? + .unwrap_or_else(|| "{}".to_owned()); + let should_cancel = format!( + "((? IN ('available', 'retryable', 'scheduled')) AND {CANCEL_ATTEMPTED_AT} IS NOT NULL)" + ); + let metadata_valid = concat!("NOT ", invalid_json!("metadata")); + let sql = format!( + r#" + UPDATE river_job + SET + attempt = CASE + WHEN NOT {should_cancel} AND ? THEN ? + ELSE attempt + END, + errors = CASE WHEN ? THEN {ERRORS_APPENDED} ELSE errors END, + finalized_at = CASE + WHEN {should_cancel} THEN ? + WHEN ? THEN ? + ELSE finalized_at + END, + -- Like River Go, metadata that isn't valid JSON is left in place. + metadata = CASE + WHEN ? AND {metadata_valid} THEN jsonb_patch(json(metadata), json(?)) + ELSE metadata + END, + scheduled_at = CASE + WHEN NOT {should_cancel} AND ? THEN ? + ELSE scheduled_at + END, + state = CASE WHEN {should_cancel} THEN 'cancelled' ELSE ? END + WHERE id = ? AND state = 'running' + RETURNING {JOB_COLUMNS} + "# + ); + let state = params.state.as_str(); + let row = sqlx::query(AssertSqlSafe(sql)) + .bind(state) + .bind(params.attempt.is_some()) + .bind(params.attempt.unwrap_or_default()) + .bind(params.error.is_some()) + .bind(&error) + .bind(&error) + .bind(&error) + .bind(state) + .bind(sqlite_time(params.now)) + .bind(params.finalized_at.is_some()) + .bind(sqlite_time_optional(params.finalized_at)) + .bind(params.metadata_updates.is_some()) + .bind(metadata) + .bind(state) + .bind(params.scheduled_at.is_some()) + .bind(sqlite_time_optional(params.scheduled_at)) + .bind(state) + .bind(state) + .bind(params.id) + .fetch_optional(&mut *connection) + .await?; + Ok(row.as_ref().map(decode_job_row)) +} + +/// Applies completion metadata after another actor has already moved a job +/// out of `running`, preserving the winning terminal state. +pub(crate) async fn merge_metadata_if_not_running( + connection: &mut SqliteConnection, + id: i64, + metadata_updates: &Map, +) -> Result, BackendError> { + let metadata = serde_json::to_string(metadata_updates)?; + // Like River Go, metadata that isn't valid JSON is left in place. + let metadata_valid = concat!("NOT ", invalid_json!("metadata")); + let sql = format!( + r#" + UPDATE river_job + SET metadata = CASE + WHEN {metadata_valid} THEN jsonb_patch(json(metadata), json(?)) + ELSE metadata + END + WHERE id = ? AND state != 'running' + RETURNING {JOB_COLUMNS} + "# + ); + let row = sqlx::query(AssertSqlSafe(sql)) + .bind(metadata) + .bind(id) + .fetch_optional(&mut *connection) + .await?; + Ok(row.as_ref().map(decode_job_row)) +} + +pub(crate) async fn update( + connection: &mut SqliteConnection, + id: i64, + metadata_updates: &Map, +) -> Result, BackendError> { + let metadata = serde_json::to_string(metadata_updates)?; + let sql = format!( + r#" + UPDATE river_job + SET metadata = jsonb_patch(json(metadata), json(?)) + WHERE id = ? + RETURNING {JOB_COLUMNS} + "# + ); + sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) + .bind(metadata) + .bind(id) + .fetch_optional(&mut *connection) + .await? + .map(JobRecord::into_job) + .transpose() +} + +pub(crate) async fn queue_upsert( + connection: &mut SqliteConnection, + name: &str, + metadata: &Map, + paused_at: Option>, + now: DateTime, +) -> Result { + let metadata = serde_json::to_string(metadata)?; + let sql = format!( + r#" + INSERT INTO river_queue (created_at, metadata, name, paused_at, updated_at) + VALUES (?, jsonb(?), ?, ?, ?) + ON CONFLICT (name) DO UPDATE SET updated_at = excluded.updated_at + RETURNING {QUEUE_COLUMNS} + "# + ); + sqlx::query_as::<_, QueueRecord>(AssertSqlSafe(sql)) + .bind(sqlite_time(now)) + .bind(metadata) + .bind(name) + .bind(sqlite_time_optional(paused_at)) + .bind(sqlite_time(now)) + .fetch_one(&mut *connection) + .await? + .into_queue() +} + +pub(crate) async fn queue_get( + connection: &mut SqliteConnection, + name: &str, +) -> Result, BackendError> { + let sql = format!("SELECT {QUEUE_COLUMNS} FROM river_queue WHERE name = ?"); + sqlx::query_as::<_, QueueRecord>(AssertSqlSafe(sql)) + .bind(name) + .fetch_optional(&mut *connection) + .await? + .map(QueueRecord::into_queue) + .transpose() +} + +pub(crate) async fn queue_list( + connection: &mut SqliteConnection, + limit: i32, +) -> Result, BackendError> { + if limit <= 0 { + return Ok(Vec::new()); + } + let sql = format!("SELECT {QUEUE_COLUMNS} FROM river_queue ORDER BY name ASC LIMIT ?"); + let records = sqlx::query_as::<_, QueueRecord>(AssertSqlSafe(sql)) + .bind(limit) + .fetch_all(&mut *connection) + .await?; + records.into_iter().map(QueueRecord::into_queue).collect() +} + +pub(crate) async fn queue_pause( + connection: &mut SqliteConnection, + name: &str, + now: DateTime, +) -> Result, BackendError> { + queue_set_paused(connection, name, Some(now), now).await +} + +pub(crate) async fn queue_resume( + connection: &mut SqliteConnection, + name: &str, + now: DateTime, +) -> Result, BackendError> { + queue_set_paused(connection, name, None, now).await +} + +async fn queue_set_paused( + connection: &mut SqliteConnection, + name: &str, + paused_at: Option>, + now: DateTime, +) -> Result, BackendError> { + let (paused_at_sql, changed_sql) = if paused_at.is_some() { + ("coalesce(paused_at, ?)", "paused_at IS NULL") + } else { + ("NULL", "paused_at IS NOT NULL") + }; + let sql = format!( + r#" + UPDATE river_queue + SET + paused_at = {paused_at_sql}, + updated_at = CASE WHEN {changed_sql} THEN ? ELSE updated_at END + WHERE (? = '*' OR name = ?) + RETURNING {QUEUE_COLUMNS} + "# + ); + let mut query = sqlx::query_as::<_, QueueRecord>(AssertSqlSafe(sql)); + if let Some(paused_at) = paused_at { + query = query.bind(sqlite_time(paused_at)); + } + let records = query + .bind(sqlite_time(now)) + .bind(name) + .bind(name) + .fetch_all(&mut *connection) + .await?; + records.into_iter().map(QueueRecord::into_queue).collect() +} + +pub(crate) async fn queue_update( + connection: &mut SqliteConnection, + name: &str, + metadata: Option<&Map>, + now: DateTime, +) -> Result, BackendError> { + let metadata = metadata.map(serde_json::to_string).transpose()?; + // Like Go, a queue update without metadata leaves it alone. + let sql = format!( + r#" + UPDATE river_queue + SET metadata = CASE WHEN ? THEN jsonb(?) ELSE metadata END, updated_at = ? + WHERE name = ? + RETURNING {QUEUE_COLUMNS} + "# + ); + sqlx::query_as::<_, QueueRecord>(AssertSqlSafe(sql)) + .bind(metadata.is_some()) + .bind(metadata) + .bind(sqlite_time(now)) + .bind(name) + .fetch_optional(&mut *connection) + .await? + .map(QueueRecord::into_queue) + .transpose() +} + +pub(crate) async fn queue_delete_expired( + connection: &mut SqliteConnection, + updated_before: DateTime, + limit: i32, +) -> Result, BackendError> { + if limit <= 0 { + return Ok(Vec::new()); + } + let records = sqlx::query_scalar::<_, String>( + r#" + DELETE FROM river_queue + WHERE name IN ( + SELECT name + FROM river_queue + WHERE updated_at < ? + ORDER BY name ASC + LIMIT ? + ) + RETURNING name + "#, + ) + .bind(sqlite_time(updated_before)) + .bind(limit) + .fetch_all(&mut *connection) + .await?; + Ok(records) +} + +pub(crate) async fn leader_elect( + connection: &mut SqliteConnection, + leader_id: &str, + now: DateTime, + ttl: Duration, +) -> Result, BackendError> { + sqlx::query_as::<_, LeaderRecord>( + r#" + INSERT INTO river_leader (leader_id, elected_at, expires_at) + VALUES (?, ?, datetime(?, 'subsec', ?)) + ON CONFLICT (name) DO NOTHING + RETURNING elected_at, expires_at, leader_id + "#, + ) + .bind(leader_id) + .bind(sqlite_time(now)) + .bind(sqlite_time(now)) + .bind(sqlite_ttl(ttl)) + .fetch_optional(&mut *connection) + .await + .map(|record| record.map(Leader::from)) + .map_err(BackendError::from) +} + +/// Extends a lease only for the same leader *and* term (`elected_at`), like +/// Go's `LeaderAttemptReelect`. +pub(crate) async fn leader_reelect( + connection: &mut SqliteConnection, + leader_id: &str, + elected_at: DateTime, + now: DateTime, + ttl: Duration, +) -> Result, BackendError> { + sqlx::query_as::<_, LeaderRecord>( + r#" + UPDATE river_leader + SET expires_at = datetime(?, 'subsec', ?) + WHERE unixepoch(elected_at, 'subsec') = unixepoch(?, 'subsec') + AND unixepoch(expires_at, 'subsec') >= unixepoch(?, 'subsec') + AND leader_id = ? + RETURNING elected_at, expires_at, leader_id + "#, + ) + .bind(sqlite_time(now)) + .bind(sqlite_ttl(ttl)) + .bind(sqlite_time(elected_at)) + .bind(sqlite_time(now)) + .bind(leader_id) + .fetch_optional(&mut *connection) + .await + .map(|record| record.map(Leader::from)) + .map_err(BackendError::from) +} + +/// Deletes a lease only for the same leader *and* term, like Go's +/// `LeaderResign`. +pub(crate) async fn leader_resign( + connection: &mut SqliteConnection, + leader_id: &str, + elected_at: DateTime, +) -> Result { + let result = sqlx::query( + r#" + DELETE FROM river_leader + WHERE unixepoch(elected_at, 'subsec') = unixepoch(?, 'subsec') + AND leader_id = ? + "#, + ) + .bind(sqlite_time(elected_at)) + .bind(leader_id) + .execute(&mut *connection) + .await?; + Ok(result.rows_affected() > 0) +} + +pub(crate) async fn leader_delete_expired( + connection: &mut SqliteConnection, + now: DateTime, +) -> Result { + let result = sqlx::query("DELETE FROM river_leader WHERE expires_at < ?") + .bind(sqlite_time(now)) + .execute(&mut *connection) + .await?; + Ok(result.rows_affected()) +} + +pub(crate) async fn notification_insert( + connection: &mut SqliteConnection, + notifications: &[NotificationInput<'_>], +) -> Result { + if notifications.is_empty() { + return Ok(0); + } + let mut query = QueryBuilder::::new("INSERT INTO river_notification (payload, topic) "); + query.push_values(notifications, |mut row, notification| { + row.push_bind(notification.payload) + .push_bind(notification.topic); + }); + Ok(query + .build() + .execute(&mut *connection) + .await? + .rows_affected()) +} + +/// Reads up to `limit` notifications on `topics` written after `after_id`, in +/// ID order. +pub(crate) async fn notification_get_after( + connection: &mut SqliteConnection, + after_id: i64, + topics: &[&str], + limit: i64, +) -> Result, BackendError> { + if limit <= 0 || topics.is_empty() { + return Ok(Vec::new()); + } + let topics = serde_json::to_string(topics)?; + Ok(sqlx::query_as::<_, Notification>( + r#" + SELECT id, payload, topic + FROM river_notification + WHERE id > ? + AND topic IN (SELECT value FROM json_each(?)) + ORDER BY id ASC + LIMIT ? + "#, + ) + .bind(after_id) + .bind(topics) + .bind(limit) + .fetch_all(&mut *connection) + .await?) +} + +pub(crate) async fn notification_last_id( + connection: &mut SqliteConnection, +) -> Result { + sqlx::query_scalar("SELECT coalesce(max(id), 0) FROM river_notification") + .fetch_one(&mut *connection) + .await + .map_err(BackendError::from) +} + +/// Reads River notifications from SQLite's `river_notification` outbox, +/// mirroring River Go's SQLite listener. +/// +/// Connecting starts after the outbox's current maximum ID so historical rows +/// aren't replayed. Each topic also records the maximum ID when it's +/// subscribed, so rows written before a subscription are never delivered, +/// including rows buffered from an earlier subscription to the same topic. +/// Rows are read in batches of [`NOTIFICATION_BATCH_SIZE`] and delivered one +/// at a time. +#[derive(Debug, Default)] +pub(crate) struct NotificationListener { + /// ID of the last row read, or `None` before connecting. + last_id: Option, + /// Rows read but not yet delivered. + pending: VecDeque, + /// Subscribed topics and the ID after which each one's rows are + /// delivered. + topics: HashMap, +} + +impl NotificationListener { + /// Disconnects, dropping subscriptions and undelivered rows. + /// + /// A client keeps its listener connected for its whole run, keeping the + /// cursor across failed reads so no durable notification is skipped. + #[cfg(test)] + pub(crate) fn close(&mut self) { + self.last_id = None; + self.pending.clear(); + self.topics.clear(); + } + + /// Starts reading after the outbox's current maximum ID. + pub(crate) async fn connect( + &mut self, + connection: &mut SqliteConnection, + ) -> Result<(), BackendError> { + self.last_id = Some(notification_last_id(connection).await?); + Ok(()) + } + + /// Whether [`connect`](Self::connect) has succeeded since the last + /// [`close`](Self::close). + pub(crate) const fn is_connected(&self) -> bool { + self.last_id.is_some() + } + + /// Subscribes to `topic`, delivering only rows written from now on. + /// Subscribing to a topic that's already subscribed changes nothing. + pub(crate) async fn listen( + &mut self, + connection: &mut SqliteConnection, + topic: &str, + ) -> Result<(), BackendError> { + if self.topics.contains_key(topic) { + return Ok(()); + } + // Only this topic's starting ID advances. Advancing the shared cursor + // could skip undelivered rows of topics that are already subscribed. + let last_id = notification_last_id(connection).await?; + // Cleanup may have deleted rows that are still buffered locally. + let start_id = last_id.max(self.last_id.unwrap_or_default()); + self.topics.insert(topic.to_owned(), start_id); + Ok(()) + } + + /// Returns the next notification on a subscribed topic, or `None` once + /// the outbox has no newer rows on subscribed topics. + pub(crate) async fn next( + &mut self, + connection: &mut SqliteConnection, + ) -> Result, BackendError> { + loop { + if self.pending.is_empty() { + let Some(after_id) = self.last_id else { + return Ok(None); + }; + let topics = self.topics.keys().map(String::as_str).collect::>(); + let notifications = + notification_get_after(connection, after_id, &topics, NOTIFICATION_BATCH_SIZE) + .await?; + let Some(last) = notifications.last() else { + return Ok(None); + }; + self.last_id = Some(last.id); + self.pending.extend(notifications); + } + // Subscriptions can change between reading a batch and delivering + // it. + while let Some(notification) = self.pending.pop_front() { + if self + .topics + .get(¬ification.topic) + .is_some_and(|start_id| notification.id > *start_id) + { + return Ok(Some(notification)); + } + } + } + } + + /// Unsubscribes from `topic`, dropping its undelivered rows. + /// + /// A client subscribes to its topics once for its whole run. + #[cfg(test)] + pub(crate) fn unlisten(&mut self, topic: &str) { + self.topics.remove(topic); + } +} + +pub(crate) async fn notification_cleanup( + connection: &mut SqliteConnection, + created_before: DateTime, + limit: i64, +) -> Result { + if limit <= 0 { + return Ok(0); + } + Ok(sqlx::query( + "DELETE FROM river_notification WHERE id IN (\ + SELECT id FROM river_notification WHERE created_at < ? \ + ORDER BY created_at, id LIMIT ?)", + ) + .bind(sqlite_time(created_before)) + .bind(limit) + .execute(&mut *connection) + .await? + .rows_affected()) +} + +pub(crate) async fn stuck_jobs( + connection: &mut SqliteConnection, + after_id: i64, + attempted_before: DateTime, + limit: i32, +) -> Result, BackendError> { + if limit <= 0 { + return Ok(Vec::new()); + } + let sql = format!( + r#" + SELECT {JOB_COLUMNS} + FROM river_job + WHERE state = 'running' + AND id > ? + AND attempted_at < ? + ORDER BY id ASC + LIMIT ? + "# + ); + let records = sqlx::query(AssertSqlSafe(sql)) + .bind(after_id) + .bind(sqlite_time(attempted_before)) + .bind(limit) + .fetch_all(&mut *connection) + .await?; + Ok(tolerant_rows(&records)) +} + +/// Loads exactly the rescue candidates selected by an exact-version +/// extension. The enclosing transaction protects the subsequent state +/// transition. +pub(crate) async fn jobs_by_ids( + connection: &mut SqliteConnection, + ids: &[i64], +) -> Result, BackendError> { + if ids.is_empty() { + return Ok(Vec::new()); + } + let mut query = + QueryBuilder::::new(format!("SELECT {JOB_COLUMNS} FROM river_job WHERE id IN (")); + { + let mut separated = query.separated(", "); + for id in ids { + separated.push_bind(id); + } + } + query.push(") ORDER BY id ASC"); + let records = query.build().fetch_all(&mut *connection).await?; + Ok(tolerant_rows(&records)) +} + +/// Decodes rows for the rescuer, which like River Go's `JobGetStuck` +/// tolerates undecodable fields so a job stranded by one can be rescued. +fn tolerant_rows(records: &[SqliteRow]) -> Vec { + records + .iter() + .filter_map(|record| tolerant_row(decode_job_row(record))) + .collect() +} + +pub(crate) async fn rescue( + connection: &mut SqliteConnection, + params: &RescueJob<'_>, +) -> Result, BackendError> { + let error = serde_json::to_string(params.error)?; + // Like River Go, the rescue error is appended to any `errors` value, and + // metadata that isn't valid JSON is left in place. + let metadata_invalid = invalid_json!("metadata"); + let sql = format!( + r#" + UPDATE river_job + SET + errors = {ERRORS_APPENDED}, + finalized_at = ?, + scheduled_at = ?, + metadata = CASE WHEN {metadata_invalid} THEN metadata ELSE jsonb_set( + metadata, + '$."river:rescue_count"', + coalesce( + CASE json_type(metadata, '$."river:rescue_count"') + WHEN 'integer' THEN json_extract(metadata, '$."river:rescue_count"') + WHEN 'real' THEN json_extract(metadata, '$."river:rescue_count"') + END, + 0 + ) + 1 + ) END, + state = ? + WHERE id = ? AND state = 'running' AND attempted_at < ? + RETURNING {JOB_COLUMNS} + "# + ); + let row = sqlx::query(AssertSqlSafe(sql)) + .bind(&error) + .bind(&error) + .bind(&error) + .bind(sqlite_time_optional(params.finalized_at)) + .bind(sqlite_time(params.scheduled_at)) + .bind(params.state.as_str()) + .bind(params.id) + .bind(sqlite_time(params.stuck_horizon)) + .fetch_optional(&mut *connection) + .await?; + // Like River Go, a rescued job whose row can't be fully decoded is still + // rescued rather than failing the rescuer's transaction. + Ok(row + .as_ref() + .and_then(|row| tolerant_row(decode_job_row(row)))) +} + +pub(crate) async fn cleanup_jobs( + connection: &mut SqliteConnection, + params: &CleanupJobs<'_>, +) -> Result { + let horizons = [ + ("cancelled", params.cancelled_before), + ("completed", params.completed_before), + ("discarded", params.discarded_before), + ]; + if params.limit <= 0 || horizons.iter().all(|(_, before)| before.is_none()) { + return Ok(0); + } + let mut query = QueryBuilder::::new( + "DELETE FROM river_job WHERE id IN (SELECT id FROM river_job WHERE (", + ); + let mut first = true; + for (state, before) in horizons { + let Some(before) = before else { + continue; + }; + if !first { + query.push(" OR "); + } + first = false; + query + .push("(state = ") + .push_bind(state) + .push(" AND finalized_at < ") + .push_bind(sqlite_time(before)) + .push(")"); + } + query.push(")"); + if !params.queues_excluded.is_empty() { + query.push(" AND queue NOT IN ("); + let mut separated = query.separated(", "); + for queue in params.queues_excluded { + separated.push_bind(queue); + } + separated.push_unseparated(")"); + } + if let Some(queues_included) = params.queues_included { + if queues_included.is_empty() { + return Ok(0); + } + query.push(" AND queue IN ("); + let mut separated = query.separated(", "); + for queue in queues_included { + separated.push_bind(queue); + } + separated.push_unseparated(")"); + } + query + .push(" ORDER BY id ASC LIMIT ") + .push_bind(params.limit) + .push(")"); + Ok(query + .build() + .execute(&mut *connection) + .await? + .rows_affected()) +} + +/// A due job the scheduler may make available. +#[derive(Clone, Debug, FromRow)] +pub(crate) struct ScheduleCandidate { + pub id: i64, + pub unique_key: Option>, +} + +/// Selects due retryable/scheduled jobs in scheduler order. Scheduling is a +/// multi-step SQLite operation: callers keep a write transaction open while +/// checking unique collisions and applying the transitions below. +/// +/// Like River Go, only the columns scheduling needs are selected, so a job +/// whose row can't be decoded, like one whose attempt failed because a JSON +/// column isn't valid JSON, doesn't fail scheduling for every other job. The +/// transitions below return such a job's row with the undecodable fields left +/// empty. +pub(crate) async fn schedule_candidates( + connection: &mut SqliteConnection, + now: DateTime, + limit: i32, +) -> Result, BackendError> { + if limit <= 0 { + return Ok(Vec::new()); + } + Ok(sqlx::query_as::<_, ScheduleCandidate>( + r#" + SELECT id, unique_key + FROM river_job + WHERE state IN ('retryable', 'scheduled') AND scheduled_at <= ? + ORDER BY priority ASC, scheduled_at ASC, id ASC + LIMIT ? + "#, + ) + .bind(sqlite_time(now)) + .bind(limit) + .fetch_all(&mut *connection) + .await?) +} + +pub(crate) async fn schedule_has_unique_collision( + connection: &mut SqliteConnection, + id: i64, + unique_key: &[u8], +) -> Result { + let exists = sqlx::query_scalar::<_, bool>( + r#" + SELECT EXISTS ( + SELECT 1 + FROM river_job + WHERE id != ? + AND unique_key = ? + AND unique_states IS NOT NULL + AND CASE state + WHEN 'available' THEN unique_states & (1 << 0) + WHEN 'cancelled' THEN unique_states & (1 << 1) + WHEN 'completed' THEN unique_states & (1 << 2) + WHEN 'discarded' THEN unique_states & (1 << 3) + WHEN 'pending' THEN unique_states & (1 << 4) + WHEN 'retryable' THEN unique_states & (1 << 5) + WHEN 'running' THEN unique_states & (1 << 6) + WHEN 'scheduled' THEN unique_states & (1 << 7) + ELSE 0 + END >= 1 + ) + "#, + ) + .bind(id) + .bind(unique_key) + .fetch_one(&mut *connection) + .await?; + Ok(exists) +} + +pub(crate) async fn schedule_set_available( + connection: &mut SqliteConnection, + ids: &[i64], +) -> Result, BackendError> { + if ids.is_empty() { + return Ok(Vec::new()); + } + let mut query = + QueryBuilder::::new("UPDATE river_job SET state = 'available' WHERE id IN ("); + { + let mut separated = query.separated(", "); + for id in ids { + separated.push_bind(id); + } + separated.push_unseparated(")"); + } + query.push(format!(" RETURNING {JOB_COLUMNS}")); + let records = query.build().fetch_all(&mut *connection).await?; + Ok(tolerant_rows(&records)) +} + +pub(crate) async fn schedule_discard_conflicts( + connection: &mut SqliteConnection, + ids: &[i64], + now: DateTime, +) -> Result, BackendError> { + if ids.is_empty() { + return Ok(Vec::new()); + } + let mut query = QueryBuilder::::new( + r#" + UPDATE river_job + SET + metadata = CASE WHEN "#, + ); + // Like River Go, metadata that isn't valid JSON is left in place. + query.push(concat!( + invalid_json!("metadata"), + r#" THEN metadata ELSE jsonb_patch( + json(metadata), + json('{"unique_key_conflict":"scheduler_discarded"}') + ) END, + finalized_at = "#, + )); + query + .push_bind(sqlite_time(now)) + .push(", state = 'discarded' WHERE id IN ("); + { + let mut separated = query.separated(", "); + for id in ids { + separated.push_bind(id); + } + separated.push_unseparated(")"); + } + query.push(format!(" RETURNING {JOB_COLUMNS}")); + let records = query.build().fetch_all(&mut *connection).await?; + Ok(tolerant_rows(&records)) +} + +#[cfg(test)] +mod tests { + use chrono::{TimeDelta, TimeZone, Timelike}; + use serde_json::json; + use sqlx::{SqlitePool, sqlite::SqlitePoolOptions}; + + use super::*; + + const SCHEMA: &str = r#" + CREATE TABLE river_job ( + id integer PRIMARY KEY, + args jsonb NOT NULL DEFAULT (jsonb('{}')), + attempt integer NOT NULL DEFAULT 0, + attempted_at timestamp, + attempted_by jsonb, + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + errors jsonb, + finalized_at timestamp, + kind text NOT NULL, + max_attempts integer NOT NULL DEFAULT 25, + metadata jsonb NOT NULL DEFAULT (jsonb('{}')), + priority integer NOT NULL DEFAULT 1, + queue text NOT NULL DEFAULT 'default', + state text NOT NULL DEFAULT 'available', + scheduled_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + tags jsonb NOT NULL DEFAULT (jsonb('[]')), + unique_key blob, + unique_states integer + ); + CREATE UNIQUE INDEX river_job_unique_idx ON river_job (unique_key) + WHERE unique_key IS NOT NULL + AND unique_states IS NOT NULL + AND CASE state + WHEN 'available' THEN unique_states & (1 << 0) + WHEN 'cancelled' THEN unique_states & (1 << 1) + WHEN 'completed' THEN unique_states & (1 << 2) + WHEN 'discarded' THEN unique_states & (1 << 3) + WHEN 'pending' THEN unique_states & (1 << 4) + WHEN 'retryable' THEN unique_states & (1 << 5) + WHEN 'running' THEN unique_states & (1 << 6) + WHEN 'scheduled' THEN unique_states & (1 << 7) + ELSE 0 + END >= 1; + CREATE TABLE river_queue ( + name text PRIMARY KEY NOT NULL, + created_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP, + metadata jsonb NOT NULL DEFAULT (jsonb('{}')), + paused_at timestamp, + updated_at timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP + ); + CREATE TABLE river_leader ( + elected_at timestamp NOT NULL, + expires_at timestamp NOT NULL, + leader_id text NOT NULL, + name text PRIMARY KEY NOT NULL DEFAULT 'default' CHECK (name = 'default') + ); + CREATE TABLE river_notification ( + id integer PRIMARY KEY AUTOINCREMENT, + created_at timestamp NOT NULL DEFAULT (datetime('now', 'subsec')), + payload text NOT NULL, + topic text NOT NULL + ); + "#; + + async fn setup() -> SqlitePool { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + sqlx::raw_sql(SCHEMA).execute(&pool).await.unwrap(); + pool + } + + #[allow( + clippy::too_many_lines, + reason = "one scenario verifies the complete persisted SQLite job lifecycle" + )] + #[tokio::test] + async fn job_lifecycle_and_unique_insert() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let now = Utc + .with_ymd_and_hms(2026, 1, 2, 3, 4, 5) + .unwrap() + .with_nanosecond(123_800_000) + .unwrap(); + let args = serde_json::value::to_raw_value(&json!({"message": "hello"})).unwrap(); + let metadata = JobMetadata::default(); + let tags = vec!["mail".to_owned()]; + let unique_key = [7_u8; 32]; + let insert_params = InsertJob { + attempt: 0, + attempted_at: None, + attempted_by: &[], + created_at: Some(now), + encoded_args: &args, + errors: &[], + finalized_at: None, + id: None, + kind: "send_mail", + max_attempts: 25, + metadata: &metadata, + priority: 1, + queue: "default", + scheduled_at: Some(now), + state: JobState::Available, + tags: &tags, + unique_key: Some(&unique_key), + unique_nonce: Some("first"), + unique_states: Some( + JobState::UNIQUE_DEFAULT + .iter() + .fold(0, |bits, state| bits | state.unique_bit()), + ), + }; + + let inserted = insert(&mut connection, &insert_params).await.unwrap(); + assert!(!inserted.unique_skipped_as_duplicate); + assert_eq!(inserted.job.encoded_args.get(), args.get()); + assert_eq!(inserted.job.tags, tags); + let (created_at, scheduled_at): (String, String) = + sqlx::query_as("SELECT created_at, scheduled_at FROM river_job WHERE id = ?") + .bind(inserted.job.id) + .fetch_one(&mut *connection) + .await + .unwrap(); + assert_eq!(created_at, "2026-01-02 03:04:05.124"); + assert_eq!(scheduled_at, "2026-01-02 03:04:05.124"); + + let duplicate = InsertJob { + unique_nonce: Some("second"), + ..insert_params + }; + let duplicate = insert(&mut connection, &duplicate).await.unwrap(); + assert!(duplicate.unique_skipped_as_duplicate); + assert_eq!(duplicate.job.id, inserted.job.id); + + queue_upsert(&mut connection, "default", &Map::new(), None, now) + .await + .unwrap(); + queue_pause(&mut connection, "default", now).await.unwrap(); + let claim_params = ClaimJobs { + client_id: "client-1", + kinds: None, + limit: 10, + max_attempted_by: 100, + now, + queue: "default", + }; + assert!( + claim(&mut connection, &claim_params) + .await + .unwrap() + .is_empty() + ); + + queue_resume(&mut connection, "default", now).await.unwrap(); + let claimed = claim(&mut connection, &claim_params) + .await + .unwrap() + .into_iter() + .collect::, _>>() + .unwrap(); + assert_eq!(claimed.len(), 1); + assert_eq!(claimed[0].attempt, 1); + assert_eq!(claimed[0].attempted_by, ["client-1"]); + assert_eq!(claimed[0].state, JobState::Running); + + let mut output = Map::new(); + output.insert(crate::METADATA_KEY_OUTPUT.to_owned(), json!({"sent": true})); + let completed = complete( + &mut connection, + &CompleteJob { + attempt: None, + error: None, + finalized_at: Some(now), + id: claimed[0].id, + metadata_updates: Some(&output), + now, + scheduled_at: None, + state: JobState::Completed, + }, + ) + .await + .unwrap() + .unwrap(); + assert_eq!(completed.state, JobState::Completed); + assert_eq!( + completed.output().map(serde_json::value::RawValue::get), + Some(r#"{"sent":true}"#) + ); + + let listed = list( + &mut connection, + &ListJobs { + states: &[JobState::Completed], + tags_all: &["mail"], + limit: 10, + ..ListJobs::default() + }, + ) + .await + .unwrap(); + assert_eq!(listed.len(), 1); + assert!( + delete(&mut connection, completed.id) + .await + .unwrap() + .is_some() + ); + assert!(get(&mut connection, completed.id).await.unwrap().is_none()); + } + + #[tokio::test] + #[expect( + clippy::too_many_lines, + reason = "one scenario walks every maintenance query in order" + )] + async fn maintenance_leadership_and_wakeup_outbox() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let now = Utc::now(); + let args = serde_json::value::to_raw_value(&json!({})).unwrap(); + let metadata = JobMetadata::default(); + let scheduled = insert( + &mut connection, + &InsertJob { + attempt: 0, + attempted_at: None, + attempted_by: &[], + created_at: Some(now), + encoded_args: &args, + errors: &[], + finalized_at: None, + id: None, + kind: "scheduled", + max_attempts: 3, + metadata: &metadata, + priority: 1, + queue: "default", + scheduled_at: Some(now), + state: JobState::Scheduled, + tags: &[], + unique_key: None, + unique_nonce: None, + unique_states: None, + }, + ) + .await + .unwrap() + .job; + + let candidates = schedule_candidates(&mut connection, now, 10).await.unwrap(); + assert_eq!(candidates.len(), 1); + let available = schedule_set_available(&mut connection, &[scheduled.id]) + .await + .unwrap(); + assert_eq!(available[0].state, JobState::Available); + + let leader = leader_elect(&mut connection, "leader-1", now, Duration::from_secs(30)) + .await + .unwrap() + .unwrap(); + assert!( + leader_elect(&mut connection, "leader-2", now, Duration::from_secs(30)) + .await + .unwrap() + .is_none() + ); + // A same-ID lease from another term is neither renewed nor resigned. + let other_term = leader.elected_at - TimeDelta::seconds(1); + assert!( + leader_reelect( + &mut connection, + "leader-1", + other_term, + now + TimeDelta::seconds(1), + Duration::from_secs(30), + ) + .await + .unwrap() + .is_none() + ); + assert!( + !leader_resign(&mut connection, "leader-1", other_term) + .await + .unwrap() + ); + let renewed = leader_reelect( + &mut connection, + "leader-1", + leader.elected_at, + now + TimeDelta::seconds(1), + Duration::from_secs(30), + ) + .await + .unwrap() + .unwrap(); + assert!(renewed.expires_at > leader.expires_at); + assert!( + leader_resign(&mut connection, "leader-1", renewed.elected_at) + .await + .unwrap() + ); + + let inserted = notification_insert( + &mut connection, + &[ + NotificationInput { + payload: "1", + topic: "insert_many", + }, + NotificationInput { + payload: "default", + topic: "queue_pause", + }, + ], + ) + .await + .unwrap(); + assert_eq!(inserted, 2); + let notifications = + notification_get_after(&mut connection, 0, &["insert_many", "queue_pause"], 10) + .await + .unwrap(); + assert_eq!(notifications.len(), 2); + assert_eq!(notifications[0].topic, "insert_many"); + assert_eq!(notification_last_id(&mut connection).await.unwrap(), 2); + } + + #[tokio::test] + async fn notification_cleanup_is_bounded_and_preserves_recent_rows() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let now = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + // The oldest expired row has the highest ID, so cleanup by age differs + // from cleanup by ID. + sqlx::query( + "INSERT INTO river_notification (created_at, payload, topic) VALUES \ + (?, 'old', 'test'), (?, 'recent', 'test'), (?, 'oldest', 'test')", + ) + .bind(sqlite_time( + now - TimeDelta::minutes(5) - TimeDelta::seconds(1), + )) + .bind(sqlite_time(now - TimeDelta::minutes(4))) + .bind(sqlite_time(now - TimeDelta::minutes(6))) + .execute(&mut *connection) + .await + .unwrap(); + let remaining = async |connection: &mut SqliteConnection| -> Vec { + sqlx::query_scalar("SELECT payload FROM river_notification ORDER BY id") + .fetch_all(&mut *connection) + .await + .unwrap() + }; + + assert_eq!( + notification_cleanup(&mut connection, now - TimeDelta::minutes(5), 1) + .await + .unwrap(), + 1 + ); + assert_eq!(remaining(&mut connection).await, ["old", "recent"]); + assert_eq!( + notification_cleanup(&mut connection, now - TimeDelta::minutes(5), 10) + .await + .unwrap(), + 1 + ); + assert_eq!(remaining(&mut connection).await, ["recent"]); + } + + async fn notify(connection: &mut SqliteConnection, topic: &str, payloads: &[&str]) { + let notifications = payloads + .iter() + .map(|payload| NotificationInput { payload, topic }) + .collect::>(); + notification_insert(connection, ¬ifications) + .await + .unwrap(); + } + + async fn next_notification( + listener: &mut NotificationListener, + connection: &mut SqliteConnection, + ) -> Option<(String, String)> { + listener + .next(connection) + .await + .unwrap() + .map(|notification| (notification.topic, notification.payload)) + } + + fn notification(topic: &str, payload: &str) -> (String, String) { + (topic.to_owned(), payload.to_owned()) + } + + async fn connected_listener(connection: &mut SqliteConnection) -> NotificationListener { + let mut listener = NotificationListener::default(); + listener.connect(connection).await.unwrap(); + listener + } + + #[tokio::test] + async fn notification_listener_does_not_replay_before_subscription() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + notify(&mut connection, "topic", &["historical"]).await; + let mut listener = connected_listener(&mut connection).await; + + notify(&mut connection, "topic", &["old"]).await; + listener.listen(&mut connection, "topic").await.unwrap(); + notify(&mut connection, "topic", &["new"]).await; + + assert_eq!( + next_notification(&mut listener, &mut connection).await, + Some(notification("topic", "new")) + ); + assert_eq!( + next_notification(&mut listener, &mut connection).await, + None + ); + } + + #[tokio::test] + async fn notification_listener_new_topic_preserves_other_topics() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let mut listener = connected_listener(&mut connection).await; + + listener.listen(&mut connection, "topic1").await.unwrap(); + notify(&mut connection, "topic1", &["pending"]).await; + notify(&mut connection, "topic2", &["old"]).await; + listener.listen(&mut connection, "topic2").await.unwrap(); + notify(&mut connection, "topic2", &["new"]).await; + + assert_eq!( + next_notification(&mut listener, &mut connection).await, + Some(notification("topic1", "pending")) + ); + assert_eq!( + next_notification(&mut listener, &mut connection).await, + Some(notification("topic2", "new")) + ); + assert_eq!( + next_notification(&mut listener, &mut connection).await, + None + ); + } + + #[tokio::test] + async fn notification_listener_reads_multiple_batches() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let mut listener = connected_listener(&mut connection).await; + listener.listen(&mut connection, "topic").await.unwrap(); + + // More rows than one read returns, behind a run of unsubscribed rows. + let payloads = (0..600) + .map(|index| format!("payload_{index}")) + .collect::>(); + let payloads = payloads.iter().map(String::as_str).collect::>(); + notify(&mut connection, "ignored", &payloads).await; + notify(&mut connection, "topic", &payloads).await; + + for payload in payloads { + assert_eq!( + next_notification(&mut listener, &mut connection).await, + Some(notification("topic", payload)) + ); + } + assert_eq!( + next_notification(&mut listener, &mut connection).await, + None + ); + } + + #[tokio::test] + async fn notification_listener_reconnect_discards_buffered_notifications() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let mut listener = connected_listener(&mut connection).await; + listener.listen(&mut connection, "topic").await.unwrap(); + notify(&mut connection, "topic", &["first", "buffered"]).await; + assert_eq!( + next_notification(&mut listener, &mut connection).await, + Some(notification("topic", "first")) + ); + + listener.close(); + listener.connect(&mut connection).await.unwrap(); + listener.listen(&mut connection, "topic").await.unwrap(); + notify(&mut connection, "topic", &["new"]).await; + + assert_eq!( + next_notification(&mut listener, &mut connection).await, + Some(notification("topic", "new")) + ); + assert_eq!( + next_notification(&mut listener, &mut connection).await, + None + ); + } + + #[tokio::test] + async fn notification_listener_repeated_listen_preserves_pending_notifications() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let mut listener = connected_listener(&mut connection).await; + + listener.listen(&mut connection, "topic").await.unwrap(); + notify(&mut connection, "topic", &["pending"]).await; + listener.listen(&mut connection, "topic").await.unwrap(); + + assert_eq!( + next_notification(&mut listener, &mut connection).await, + Some(notification("topic", "pending")) + ); + assert_eq!( + next_notification(&mut listener, &mut connection).await, + None + ); + } + + #[tokio::test] + async fn notification_listener_resubscribe_after_cleanup() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let mut listener = connected_listener(&mut connection).await; + listener.listen(&mut connection, "topic").await.unwrap(); + notify(&mut connection, "topic", &["first", "buffered"]).await; + assert_eq!( + next_notification(&mut listener, &mut connection).await, + Some(notification("topic", "first")) + ); + + // Cleanup deletes rows that are still buffered, so the outbox's + // maximum ID no longer covers them. + listener.unlisten("topic"); + sqlx::query("DELETE FROM river_notification") + .execute(&mut *connection) + .await + .unwrap(); + listener.listen(&mut connection, "topic").await.unwrap(); + notify(&mut connection, "topic", &["new"]).await; + + assert_eq!( + next_notification(&mut listener, &mut connection).await, + Some(notification("topic", "new")) + ); + assert_eq!( + next_notification(&mut listener, &mut connection).await, + None + ); + } + + #[tokio::test] + async fn notification_listener_resubscribe_discards_buffered_notifications() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let mut listener = connected_listener(&mut connection).await; + listener.listen(&mut connection, "topic").await.unwrap(); + notify(&mut connection, "topic", &["first", "buffered"]).await; + assert_eq!( + next_notification(&mut listener, &mut connection).await, + Some(notification("topic", "first")) + ); + + listener.unlisten("topic"); + notify(&mut connection, "topic", &["gap"]).await; + listener.listen(&mut connection, "topic").await.unwrap(); + notify(&mut connection, "topic", &["new"]).await; + + assert_eq!( + next_notification(&mut listener, &mut connection).await, + Some(notification("topic", "new")) + ); + assert_eq!( + next_notification(&mut listener, &mut connection).await, + None + ); + } + + #[tokio::test] + async fn notification_listener_resubscribe_skips_gap() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let mut listener = connected_listener(&mut connection).await; + + listener.listen(&mut connection, "topic").await.unwrap(); + listener.unlisten("topic"); + notify(&mut connection, "topic", &["gap"]).await; + listener.listen(&mut connection, "topic").await.unwrap(); + notify(&mut connection, "topic", &["new"]).await; + + assert_eq!( + next_notification(&mut listener, &mut connection).await, + Some(notification("topic", "new")) + ); + assert_eq!( + next_notification(&mut listener, &mut connection).await, + None + ); + } + + #[tokio::test] + async fn notification_listener_unlisten_discards_buffered_notifications() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let mut listener = connected_listener(&mut connection).await; + listener.listen(&mut connection, "topic1").await.unwrap(); + listener.listen(&mut connection, "topic2").await.unwrap(); + notify(&mut connection, "topic1", &["first"]).await; + notify(&mut connection, "topic2", &["buffered"]).await; + assert_eq!( + next_notification(&mut listener, &mut connection).await, + Some(notification("topic1", "first")) + ); + + listener.unlisten("topic2"); + notify(&mut connection, "topic1", &["new"]).await; + + assert_eq!( + next_notification(&mut listener, &mut connection).await, + Some(notification("topic1", "new")) + ); + assert_eq!( + next_notification(&mut listener, &mut connection).await, + None + ); + } + + #[tokio::test] + async fn claim_caps_attempted_by_without_reversing_history() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let now = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + let inserted = insert( + &mut connection, + &InsertJob { + attempt: 0, + attempted_at: None, + attempted_by: &[], + created_at: Some(now), + encoded_args: &serde_json::value::to_raw_value(&json!({})).unwrap(), + errors: &[], + finalized_at: None, + id: None, + kind: "attempt_history", + max_attempts: 25, + metadata: &JobMetadata::default(), + priority: 1, + queue: "default", + scheduled_at: Some(now), + state: JobState::Available, + tags: &[], + unique_key: None, + unique_nonce: None, + unique_states: None, + }, + ) + .await + .unwrap() + .job; + for client_id in ["one", "two", "three", "four", "five"] { + let claimed = claim( + &mut connection, + &ClaimJobs { + client_id, + kinds: None, + limit: 1, + max_attempted_by: 3, + now, + queue: "default", + }, + ) + .await + .unwrap(); + assert_eq!(claimed.len(), 1); + sqlx::query("UPDATE river_job SET state = 'available' WHERE id = ?") + .bind(inserted.id) + .execute(&mut *connection) + .await + .unwrap(); + } + let row = get(&mut connection, inserted.id).await.unwrap().unwrap(); + assert_eq!(row.attempted_by, ["three", "four", "five"]); + } + + /// Inserts a job for the undecodable row tests. + async fn insert_test_job( + connection: &mut SqliteConnection, + state: JobState, + now: DateTime, + ) -> JobRow { + let running = state == JobState::Running; + insert( + connection, + &InsertJob { + attempt: i16::from(running), + attempted_at: running.then_some(now - TimeDelta::hours(2)), + attempted_by: &[], + created_at: Some(now), + encoded_args: &serde_json::value::to_raw_value(&json!({})).unwrap(), + errors: &[], + finalized_at: None, + id: None, + kind: "undecodable", + max_attempts: 25, + metadata: &JobMetadata::default(), + priority: 1, + queue: "default", + scheduled_at: Some(now), + state, + tags: &["tag".to_owned()], + unique_key: None, + unique_nonce: None, + unique_states: None, + }, + ) + .await + .unwrap() + .job + } + + /// Overwrites a JSON column the way a row changed out of band would be, + /// into a shape River can't decode. + async fn set_json_column(connection: &mut SqliteConnection, id: i64, column: &str, json: &str) { + sqlx::query(AssertSqlSafe(format!( + "UPDATE river_job SET {column} = jsonb(?) WHERE id = ?" + ))) + .bind(json) + .bind(id) + .execute(&mut *connection) + .await + .unwrap(); + } + + fn retryable_completion(id: i64, error: &AttemptError, now: DateTime) -> CompleteJob<'_> { + CompleteJob { + attempt: None, + error: Some(error), + finalized_at: None, + id, + metadata_updates: None, + now, + scheduled_at: Some(now + TimeDelta::hours(1)), + state: JobState::Retryable, + } + } + + // A claimed job whose row can't be decoded is returned separately with the + // fields that could be decoded, without failing the others. + #[tokio::test] + async fn claim_returns_undecodable_rows_separately() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let now = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + let job1 = insert_test_job(&mut connection, JobState::Available, now).await; + let job2 = insert_test_job(&mut connection, JobState::Available, now).await; + let job3 = insert_test_job(&mut connection, JobState::Available, now).await; + set_json_column(&mut connection, job2.id, "errors", r#"{"not":"an array"}"#).await; + set_json_column(&mut connection, job2.id, "tags", r#"{"not":"an array"}"#).await; + + let claimed = claim( + &mut connection, + &ClaimJobs { + client_id: "client", + kinds: None, + limit: 10, + max_attempted_by: 100, + now, + queue: "default", + }, + ) + .await + .unwrap(); + assert_eq!(claimed.len(), 3); + let decoded = claimed + .iter() + .filter_map(|job| job.as_ref().ok().map(|job| job.id)) + .collect::>(); + assert_eq!(decoded, [job1.id, job3.id]); + + let undecodable = claimed + .into_iter() + .find_map(Result::err) + .expect("undecodable job"); + assert!( + undecodable.error.contains("error unmarshaling `errors`"), + "{}", + undecodable.error + ); + assert!( + undecodable.error.contains("error unmarshaling `tags`"), + "{}", + undecodable.error + ); + let row = undecodable.row.expect("partially decoded row"); + assert_eq!(row.id, job2.id); + assert_eq!(row.attempt, 1); + assert_eq!(row.attempted_by, ["client"]); + assert_eq!(row.kind, "undecodable"); + assert_eq!(row.state, JobState::Running); + assert_eq!(row.errors, []); + assert_eq!(row.tags, Vec::::new()); + } + + // A job whose row can't be fully decoded still has its state set and is + // returned with the fields that could be decoded, so it doesn't fail the + // other jobs completed in the same transaction. The undecodable value is + // left as it was. + #[tokio::test] + async fn complete_returns_undecodable_rows_partially() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let now = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + let job1 = insert_test_job(&mut connection, JobState::Running, now).await; + let job2 = insert_test_job(&mut connection, JobState::Running, now).await; + set_json_column(&mut connection, job2.id, "tags", r#"{"not":"an array"}"#).await; + let error = AttemptError::new(now, 1, "fake error"); + + let row1 = complete_decoded(&mut connection, &retryable_completion(job1.id, &error, now)) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(row1.state, JobState::Retryable); + let undecodable = + complete_decoded(&mut connection, &retryable_completion(job2.id, &error, now)) + .await + .unwrap() + .unwrap() + .unwrap_err(); + let row2 = undecodable.row.expect("partially decoded row"); + assert_eq!(row2.id, job2.id); + assert_eq!(row2.state, JobState::Retryable); + assert_eq!(row2.errors, [error]); + assert_eq!(row2.tags, Vec::::new()); + + let strict = get(&mut connection, job2.id).await.unwrap_err(); + assert!( + strict.to_string().contains("error unmarshaling `tags`"), + "{strict}" + ); + let tags: String = sqlx::query_scalar("SELECT json(tags) FROM river_job WHERE id = ?") + .bind(job2.id) + .fetch_one(&mut *connection) + .await + .unwrap(); + assert_eq!(tags, r#"{"not":"an array"}"#); + } + + // `errors` that isn't an array is wrapped in one so the new error can be + // appended without losing the existing value, like River Go. + #[tokio::test] + async fn complete_wraps_a_non_array_errors_value() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let now = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + let job = insert_test_job(&mut connection, JobState::Running, now).await; + set_json_column( + &mut connection, + job.id, + "errors", + r#"{"error":"existing value"}"#, + ) + .await; + let error = AttemptError::new(now, 1, "fake error"); + + let row = complete_decoded(&mut connection, &retryable_completion(job.id, &error, now)) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(row.state, JobState::Retryable); + assert_eq!( + row.errors + .iter() + .map(|error| error.error.as_str()) + .collect::>(), + ["existing value", "fake error"] + ); + } + + // A stuck job whose row can't be fully decoded is still returned so that + // it can be rescued. + #[tokio::test] + async fn stuck_jobs_include_undecodable_rows() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let now = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + let job1 = insert_test_job(&mut connection, JobState::Running, now).await; + let job2 = insert_test_job(&mut connection, JobState::Running, now).await; + set_json_column(&mut connection, job1.id, "tags", r#"{"not":"an array"}"#).await; + + let stuck = stuck_jobs(&mut connection, 0, now, 10).await.unwrap(); + assert_eq!( + stuck.iter().map(|job| job.id).collect::>(), + [job1.id, job2.id] + ); + assert_eq!(stuck[0].tags, Vec::::new()); + assert_eq!(stuck[1].tags, ["tag"]); + } + + #[tokio::test] + async fn late_completion_merges_metadata_without_changing_terminal_state() { + let pool = setup().await; + let mut connection = pool.acquire().await.unwrap(); + let now = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + let inserted = insert( + &mut connection, + &InsertJob { + attempt: 1, + attempted_at: Some(now), + attempted_by: &["client".to_owned()], + created_at: Some(now), + encoded_args: &serde_json::value::to_raw_value(&json!({})).unwrap(), + errors: &[], + finalized_at: None, + id: None, + kind: "late_completion", + max_attempts: 25, + metadata: &JobMetadata::from(Map::from_iter([("winner".to_owned(), json!(true))])), + priority: 1, + queue: "default", + scheduled_at: Some(now), + state: JobState::Running, + tags: &[], + unique_key: None, + unique_nonce: None, + unique_states: None, + }, + ) + .await + .unwrap() + .job; + for terminal in [JobState::Completed, JobState::Discarded] { + sqlx::query("UPDATE river_job SET state = ?, metadata = jsonb('{\"winner\":true}') WHERE id = ?") + .bind(terminal.as_str()) + .bind(inserted.id) + .execute(&mut *connection) + .await + .unwrap(); + let completion = complete( + &mut connection, + &CompleteJob { + attempt: Some(1), + error: None, + finalized_at: Some(now), + id: inserted.id, + metadata_updates: Some(&Map::from_iter([("stale".to_owned(), json!(true))])), + now, + scheduled_at: None, + state: JobState::Completed, + }, + ) + .await + .unwrap(); + assert!(completion.is_none()); + let completion = merge_metadata_if_not_running( + &mut connection, + inserted.id, + &Map::from_iter([("stale".to_owned(), json!(true))]), + ) + .await + .unwrap() + .unwrap() + .unwrap(); + assert_eq!(completion.state, terminal); + let row = get(&mut connection, inserted.id).await.unwrap().unwrap(); + assert_eq!(row.state, terminal); + assert_eq!(row.metadata.get::("winner").unwrap(), Some(true)); + assert_eq!(row.metadata.get::("stale").unwrap(), Some(true)); + } + } + + #[test] + fn wire_encoders_match_go_sqlite() { + let base = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + assert_eq!( + sqlite_time(base.with_nanosecond(123_400_000).unwrap()), + "2026-01-02 03:04:05.123" + ); + assert_eq!( + sqlite_time(base.with_nanosecond(123_800_000).unwrap()), + "2026-01-02 03:04:05.124" + ); + assert_eq!( + sqlite_time(base.with_nanosecond(999_800_000).unwrap()), + "2026-01-02 03:04:06.000" + ); + assert_eq!(sqlite_ttl(Duration::from_millis(3_255)), "3.255 seconds"); + } +} diff --git a/rust/riverqueue/src/encoding.rs b/rust/riverqueue/src/encoding.rs new file mode 100644 index 000000000..bfe7f2768 --- /dev/null +++ b/rust/riverqueue/src/encoding.rs @@ -0,0 +1,586 @@ +//! Go-compatible JSON encoding for job arguments. +//! +//! River stores job arguments as the exact JSON bytes produced when they were +//! inserted, and River Go hashes those bytes to build unique keys. For a Rust +//! job and an equivalent Go job to share a unique key, they must encode to the +//! same bytes. [`encode_args`] serializes with [`serde_json`] using the output +//! rules of Go's `encoding/json`: +//! +//! * Floats use Go's shortest round-trip digits, in plain decimal notation +//! when `1e-6 <= |x| < 1e21` and exponent notation otherwise (`1`, `0.1`, +//! `100000000000000000000`, `1e+21`, `1e-7`). Negative zero encodes as `-0`. +//! * Strings escape `<`, `>`, `&`, U+2028, and U+2029 as `\u003c`, `\u003e`, +//! `\u0026`, `\u2028`, and `\u2029`. Control characters use `\b`, `\f`, +//! `\n`, `\r`, and `\t` where available and lowercase `\u00XX` otherwise. +//! * Struct fields keep their declaration order, as in Go. +//! +//! Some differences come from how types serialize rather than from the JSON +//! encoder and must be handled in the argument type: +//! +//! * Go sorts map keys. Use an ordered map such as +//! [`BTreeMap`](std::collections::BTreeMap) with string keys to match; a +//! `HashMap` serializes in an unspecified order. +//! * Go encodes `[]byte` as a base64 string; a Rust `Vec` encodes as an +//! array of numbers. +//! * Go encodes `time.Time` with RFC 3339 and the shortest fractional +//! seconds, while `chrono` pads fractional seconds to 3, 6, or 9 digits. Use +//! [`go_time`] for `DateTime` fields that participate in unique keys. +//! +//! Like Go, encoding fails for non-finite floats (`NaN` and infinities) +//! instead of writing them as `null` the way [`serde_json`] does. +//! +//! These rules match Go 1.22 and later, which escape backspace and form feed +//! as `\b` and `\f`. + +use std::{fmt, fmt::Write as _, io}; + +use chrono::{DateTime, SecondsFormat, Timelike, Utc}; +use serde::Serialize; +use serde_json::{ + ser::{CharEscape, Formatter}, + value::RawValue, +}; + +mod finite; + +pub(crate) use finite::check as check_finite; + +/// Encodes job arguments to JSON bytes identical to those Go's +/// `encoding/json` produces for an equivalent Go value. +/// +/// River uses this encoding for every inserted job, so it is only needed +/// directly when constructing [`JobRow`](crate::JobRow) values by hand, such +/// as in tests. +/// +/// # Errors +/// +/// Returns an error when the value's [`Serialize`] implementation fails, for +/// example because a map has non-string keys, or when it contains a +/// non-finite float, which Go's `encoding/json` rejects as an unsupported +/// value. +pub fn encode_args(args: &T) -> Result, serde_json::Error> { + RawValue::from_string(to_go_string(args)?) +} + +/// Serde helpers that encode a `DateTime` the way Go's `encoding/json` +/// encodes a UTC `time.Time`: RFC 3339 with the shortest fractional seconds +/// (`2026-01-02T03:04:05.5Z`). +/// +/// Use it with `#[serde(with = "riverqueue::encoding::go_time")]` on argument +/// fields that participate in Go-compatible unique keys. Deserialization +/// accepts any RFC 3339 timestamp. +pub mod go_time { + use chrono::{DateTime, Utc}; + use serde::{Deserialize, Deserializer, Serializer}; + + /// Serializes a timestamp in Go's `time.RFC3339Nano` form. + /// + /// # Errors + /// + /// Returns the serializer's error. + pub fn serialize( + timestamp: &DateTime, + serializer: S, + ) -> Result { + serializer.serialize_str(&super::format_rfc3339_nano(*timestamp)) + } + + /// Deserializes an RFC 3339 timestamp. + /// + /// # Errors + /// + /// Returns an error when the input is not an RFC 3339 timestamp. + pub fn deserialize<'de, D: Deserializer<'de>>( + deserializer: D, + ) -> Result, D::Error> { + DateTime::::deserialize(deserializer) + } +} + +/// Formats a UTC timestamp like Go's `time.RFC3339Nano`. +pub(crate) fn format_rfc3339_nano(timestamp: DateTime) -> String { + let mut formatted = timestamp.to_rfc3339_opts(SecondsFormat::Secs, true); + let nanos = timestamp.nanosecond() % 1_000_000_000; + if nanos > 0 { + let fraction = format!("{nanos:09}"); + formatted.insert(formatted.len() - 1, '.'); + formatted.insert_str(formatted.len() - 1, fraction.trim_end_matches('0')); + } + formatted +} + +/// A [`serde_json`] formatter that writes compact JSON with Go's +/// `encoding/json` number formatting and string escaping. +#[derive(Clone, Copy, Debug, Default)] +pub(crate) struct GoFormatter; + +impl Formatter for GoFormatter { + fn write_f32(&mut self, writer: &mut W, value: f32) -> io::Result<()> + where + W: ?Sized + io::Write, + { + let absolute = value.abs(); + let exponent = absolute != 0.0 && !(1e-6..1e21).contains(&absolute); + writer.write_all(format_go_float(value, exponent).as_bytes()) + } + + fn write_f64(&mut self, writer: &mut W, value: f64) -> io::Result<()> + where + W: ?Sized + io::Write, + { + let absolute = value.abs(); + let exponent = absolute != 0.0 && !(1e-6..1e21).contains(&absolute); + writer.write_all(format_go_float(value, exponent).as_bytes()) + } + + fn write_string_fragment(&mut self, writer: &mut W, fragment: &str) -> io::Result<()> + where + W: ?Sized + io::Write, + { + let mut start = 0; + for (index, character) in fragment.char_indices() { + let escaped = match character { + '<' => "\\u003c", + '>' => "\\u003e", + '&' => "\\u0026", + '\u{2028}' => "\\u2028", + '\u{2029}' => "\\u2029", + _ => continue, + }; + writer.write_all(&fragment.as_bytes()[start..index])?; + writer.write_all(escaped.as_bytes())?; + start = index + character.len_utf8(); + } + writer.write_all(&fragment.as_bytes()[start..]) + } + + /// Embeds raw JSON (a [`RawValue`] nested in job arguments) the way Go + /// embeds a `json.RawMessage`: compacted, with the same HTML-safe string + /// escaping, and every other token byte for byte, because unique keys + /// hash these bytes. + fn write_raw_fragment(&mut self, writer: &mut W, fragment: &str) -> io::Result<()> + where + W: ?Sized + io::Write, + { + writer.write_all(go_compact(fragment).as_bytes()) + } + + fn write_char_escape(&mut self, writer: &mut W, char_escape: CharEscape) -> io::Result<()> + where + W: ?Sized + io::Write, + { + let escaped: &[u8] = match char_escape { + CharEscape::Quote => b"\\\"", + CharEscape::ReverseSolidus => b"\\\\", + CharEscape::Solidus => b"/", + CharEscape::Backspace => b"\\b", + CharEscape::FormFeed => b"\\f", + CharEscape::LineFeed => b"\\n", + CharEscape::CarriageReturn => b"\\r", + CharEscape::Tab => b"\\t", + CharEscape::AsciiControl(byte) => { + return writer.write_all(&control_escape(byte)); + } + }; + writer.write_all(escaped) + } +} + +/// Formats a float with Go's `strconv.FormatFloat(value, 'f' or 'e', -1)` +/// followed by `encoding/json`'s exponent cleanup (`e-07` becomes `e-7`). +fn format_go_float(value: F, exponent_notation: bool) -> String { + let wide = value.to_f64(); + if wide == 0.0 { + return if wide.is_sign_negative() { "-0" } else { "0" }.to_owned(); + } + + // Rust's `LowerExp` writes the shortest round-trip digits, as Go does, + // but breaks exact ties differently; see `round_tie_to_even`. + let scientific = format!("{value:e}"); + let (negative, scientific) = match scientific.strip_prefix('-') { + Some(rest) => (true, rest), + None => (false, scientific.as_str()), + }; + let Some((mantissa, power)) = scientific.split_once('e') else { + return scientific.to_owned(); + }; + let mut digits = mantissa.replace('.', ""); + let mut power = power.parse::().unwrap_or_default(); + round_tie_to_even(value, &mut digits, &mut power); + + let mut output = String::with_capacity(digits.len() + 8); + if negative { + output.push('-'); + } + if exponent_notation { + output.push_str(&digits[..1]); + if digits.len() > 1 { + output.push('.'); + output.push_str(&digits[1..]); + } + if power < 0 { + let _ = write!(output, "e-{}", -power); + } else { + let _ = write!(output, "e+{power:02}"); + } + return output; + } + // `'f'` notation: place the decimal point `power + 1` digits into the + // significant digits, padding with zeros on either side. + let integer_digits = power + 1; + if integer_digits <= 0 { + output.push_str("0."); + output.extend(std::iter::repeat_n( + '0', + usize::try_from(-integer_digits).unwrap_or(0), + )); + output.push_str(&digits); + } else { + let integer_digits = usize::try_from(integer_digits).unwrap_or(0); + if integer_digits >= digits.len() { + output.push_str(&digits); + output.extend(std::iter::repeat_n('0', integer_digits - digits.len())); + } else { + output.push_str(&digits[..integer_digits]); + output.push('.'); + output.push_str(&digits[integer_digits..]); + } + } + output +} + +/// Adjusts shortest digits for an exact tie. When a float lies exactly +/// halfway between the two nearest decimals with the shortest round-trip +/// digit count, Go chooses the one with an even last digit while Rust rounds +/// away from zero (`472476.125_f32` is `472476.12` in Go and `472476.13` in +/// Rust). +fn round_tie_to_even(value: F, digits: &mut String, power: &mut i32) { + // Decompose |value| exactly as `mantissa * 2^exponent` with an odd + // mantissa. Only negative exponents have a fractional decimal expansion + // that can end in the 5 of an exact tie. + let bits = value.to_f64().abs().to_bits(); + let biased = i32::try_from((bits >> 52) & 0x7ff).unwrap_or(0); + let fraction = bits & ((1_u64 << 52) - 1); + let (mut mantissa, mut exponent) = if biased == 0 { + (fraction, -1074) + } else { + (fraction | (1_u64 << 52), biased - 1075) + }; + let trailing = mantissa.trailing_zeros(); + mantissa >>= trailing; + exponent += i32::try_from(trailing).unwrap_or(0); + let Ok(scale) = u32::try_from(-exponent) else { + return; + }; + // The exact value is `mantissa * 5^scale * 10^-scale`. A tie at the + // shortest length means it has exactly one more significant digit. + let Some(exact) = 5_u128 + .checked_pow(scale) + .and_then(|factor| factor.checked_mul(u128::from(mantissa))) + else { + return; + }; + let exact_digits = exact.to_string(); + if exact_digits.len() != digits.len() + 1 { + return; + } + let below = exact / 10; + let even = if below % 2 == 0 { below } else { below + 1 }; + let candidate = even.to_string(); + if candidate == *digits { + return; + } + let scale = i32::try_from(scale).unwrap_or(i32::MAX); + let candidate_power = i32::try_from(candidate.len()).unwrap_or(0) - 1 - scale + 1; + if F::parses_to(&format!("{candidate}e{}", 1 - scale), value) { + candidate.trim_end_matches('0').clone_into(digits); + *power = candidate_power; + } +} + +/// Float widths formatted with Go's rules. +trait GoFloat: Copy + fmt::LowerExp { + fn to_f64(self) -> f64; + fn parses_to(text: &str, value: Self) -> bool; +} + +impl GoFloat for f32 { + fn to_f64(self) -> f64 { + f64::from(self) + } + + fn parses_to(text: &str, value: Self) -> bool { + text.parse::() + .is_ok_and(|parsed| parsed.abs().to_bits() == value.abs().to_bits()) + } +} + +impl GoFloat for f64 { + fn to_f64(self) -> f64 { + self + } + + fn parses_to(text: &str, value: Self) -> bool { + text.parse::() + .is_ok_and(|parsed| parsed.abs().to_bits() == value.abs().to_bits()) + } +} + +fn control_escape(byte: u8) -> [u8; 6] { + const HEX: &[u8; 16] = b"0123456789abcdef"; + [ + b'\\', + b'u', + b'0', + b'0', + HEX[usize::from(byte >> 4)], + HEX[usize::from(byte & 0xf)], + ] +} + +/// Compacts valid JSON like Go's `json.Compact` after `json.HTMLEscape`: +/// whitespace between tokens is removed, and inside strings `<`, `>`, `&`, +/// U+2028, and U+2029 are escaped. Numbers, key order, and existing escapes +/// are kept byte for byte. Only raw JSON embedded in job arguments goes +/// through it, since unique keys hash the argument bytes. +fn go_compact(json: &str) -> String { + let mut output = String::with_capacity(json.len()); + let mut in_string = false; + let mut escaped = false; + for character in json.chars() { + if !in_string { + match character { + ' ' | '\t' | '\n' | '\r' => {} + '"' => { + in_string = true; + output.push(character); + } + _ => output.push(character), + } + continue; + } + if escaped { + escaped = false; + output.push(character); + continue; + } + match character { + '\\' => { + escaped = true; + output.push(character); + } + '"' => { + in_string = false; + output.push(character); + } + '<' => output.push_str("\\u003c"), + '>' => output.push_str("\\u003e"), + '&' => output.push_str("\\u0026"), + '\u{2028}' => output.push_str("\\u2028"), + '\u{2029}' => output.push_str("\\u2029"), + _ => output.push(character), + } + } + output +} + +/// Serializes `value` to JSON text with Go's `encoding/json` output rules, +/// as [`encode_args`] does, including its rejection of non-finite floats. +/// Other stored JSON only needs to decode to the same value, so it uses plain +/// [`serde_json`]. +pub(crate) fn to_go_string(value: &T) -> Result { + check_finite(value)?; + let mut buffer = Vec::with_capacity(128); + value.serialize(&mut serde_json::Serializer::with_formatter( + &mut buffer, + GoFormatter, + ))?; + String::from_utf8(buffer).map_err(::custom) +} + +/// Appends `value` as a JSON string with Go's `encoding/json` escaping. +pub(crate) fn write_go_string(value: &str, output: &mut String) { + output.push('"'); + for character in value.chars() { + match character { + '"' => output.push_str("\\\""), + '\\' => output.push_str("\\\\"), + '\u{8}' => output.push_str("\\b"), + '\u{c}' => output.push_str("\\f"), + '\n' => output.push_str("\\n"), + '\r' => output.push_str("\\r"), + '\t' => output.push_str("\\t"), + '<' => output.push_str("\\u003c"), + '>' => output.push_str("\\u003e"), + '&' => output.push_str("\\u0026"), + '\u{2028}' => output.push_str("\\u2028"), + '\u{2029}' => output.push_str("\\u2029"), + character if u32::from(character) < 0x20 => { + let escaped = control_escape(u8::try_from(u32::from(character)).unwrap_or(0)); + output.extend(escaped.iter().map(|&byte| char::from(byte))); + } + character => output.push(character), + } + } + output.push('"'); +} + +#[cfg(test)] +mod tests { + use std::collections::BTreeMap; + + use chrono::TimeZone; + + use super::*; + + fn encoded(value: &T) -> String { + encode_args(value).unwrap().get().to_owned() + } + + #[test] + fn encodes_floats_like_go() { + // Expected values produced by Go's `encoding/json`. + let cases: [(f64, &str); 19] = [ + (0.0, "0"), + (-0.0, "-0"), + (1.0, "1"), + (0.1, "0.1"), + (100.0, "100"), + (1e6, "1000000"), + (12_345.678, "12345.678"), + (1e-6, "0.000001"), + (0.000_001_234, "0.000001234"), + (1e-7, "1e-7"), + (-1.5e-9, "-1.5e-9"), + (1e20, "100000000000000000000"), + (123_456_789_012_345_680_000.0, "123456789012345680000"), + (1e21, "1e+21"), + (1.5e300, "1.5e+300"), + (5e-324, "5e-324"), + (f64::MAX, "1.7976931348623157e+308"), + // Exact ties between two shortest candidates round to even. + (1_357_346_946_266_522.2, "1357346946266522.2"), + (3_371_836_896_475.031_2, "3371836896475.0312"), + ]; + for (value, expected) in cases { + assert_eq!(encoded(&value), expected, "{value:e}"); + } + + let cases: [(f32, &str); 10] = [ + (1.0, "1"), + (0.1, "0.1"), + (1.1, "1.1"), + (1e-7, "1e-7"), + (1e20, "100000000000000000000"), + (1e21, "1e+21"), + (16_777_216.0, "16777216"), + (472_476.12, "472476.12"), + (-368.140_62, "-368.14062"), + (2_569_406.2, "2569406.2"), + ]; + for (value, expected) in cases { + assert_eq!(encoded(&value), expected, "{value:e}"); + } + assert_eq!(encoded(&f32::MAX), "3.4028235e+38"); + } + + #[test] + fn rejects_non_finite_floats_like_go() { + #[derive(Serialize)] + struct Args { + nested: Vec>>, + value: f64, + } + + // Messages match Go's `json: unsupported value: ...` after the prefix. + for (value, message) in [ + (f64::NAN, "unsupported value: NaN"), + (f64::INFINITY, "unsupported value: +Inf"), + (f64::NEG_INFINITY, "unsupported value: -Inf"), + ] { + assert_eq!(encode_args(&value).unwrap_err().to_string(), message); + let args = Args { + nested: Vec::new(), + value, + }; + assert_eq!(encode_args(&args).unwrap_err().to_string(), message); + } + let nested = Args { + nested: vec![None, Some(BTreeMap::from([("x", f32::NAN)]))], + value: 1.0, + }; + assert_eq!( + encode_args(&nested).unwrap_err().to_string(), + "unsupported value: NaN" + ); + assert!(to_go_string(&serde_json::json!({"x": [1.5, null]})).is_ok()); + assert_eq!(encoded(&Option::::None), "null"); + } + + #[test] + fn escapes_strings_like_go() { + let value = "<>&\u{2028}\u{2029}\u{8}\u{c}\n\r\t\u{1}\u{1f}\u{7f}\"\\/é😀"; + let expected = r#""\u003c\u003e\u0026\u2028\u2029\b\f\n\r\t\u0001\u001f"#.to_owned() + + "\u{7f}\\\"\\\\/é😀\""; + assert_eq!(encoded(value), expected); + + let mut direct = String::new(); + write_go_string(value, &mut direct); + assert_eq!(direct, expected); + } + + #[test] + fn embeds_raw_json_like_go_raw_messages() { + let raw = RawValue::from_string( + "{ \"b\" : \"a&\u{2028}\\u003c\\\"<\" ,\n \"n\": 1.50e0, \"z\":[ 1 , 2 ] }" + .to_owned(), + ) + .unwrap(); + assert_eq!( + encoded(&raw), + r#"{"b":"a\u003cb\u003e\u0026\u2028\u003c\"\u003c","n":1.50e0,"z":[1,2]}"# + ); + } + + #[test] + fn escapes_keys_and_preserves_struct_order() { + #[derive(Serialize)] + struct Args { + zulu: u8, + #[serde(rename = "a")] + angle: u8, + map: BTreeMap<&'static str, u8>, + } + + let args = Args { + zulu: 1, + angle: 2, + map: BTreeMap::from([("é&", 3), ("", 4)]), + }; + assert_eq!( + encoded(&args), + r#"{"zulu":1,"a\u003cb\u003e":2,"map":{"\u003ck\u003e":4,"é\u0026":3}}"# + ); + } + + #[test] + fn go_time_trims_fractional_seconds() { + #[derive(Serialize)] + struct Times { + #[serde(with = "go_time")] + at: DateTime, + } + + let whole = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + for (nanos, expected) in [ + (0, r#"{"at":"2026-01-02T03:04:05Z"}"#), + (500_000_000, r#"{"at":"2026-01-02T03:04:05.5Z"}"#), + (120_000_000, r#"{"at":"2026-01-02T03:04:05.12Z"}"#), + (123_456_000, r#"{"at":"2026-01-02T03:04:05.123456Z"}"#), + (1, r#"{"at":"2026-01-02T03:04:05.000000001Z"}"#), + ] { + let at = whole + chrono::Duration::nanoseconds(nanos); + assert_eq!(encoded(&Times { at }), expected); + } + } +} diff --git a/rust/riverqueue/src/encoding/finite.rs b/rust/riverqueue/src/encoding/finite.rs new file mode 100644 index 000000000..b41a50bb7 --- /dev/null +++ b/rust/riverqueue/src/encoding/finite.rs @@ -0,0 +1,319 @@ +//! A serializer that only checks a value for non-finite floats. +//! +//! [`serde_json`] writes `NaN` and infinities as `null` through the same +//! formatter call as a real `null`, so the check can't happen in +//! [`GoFormatter`](super::GoFormatter). Go's `encoding/json` rejects these +//! values with an `UnsupportedValueError`, so River walks the value once with +//! this serializer before encoding it. + +use std::fmt; + +use serde::{Serialize, ser}; + +/// Returns an error like Go's `json: unsupported value: NaN` when `value` +/// contains a non-finite float anywhere. +pub(crate) fn check(value: &T) -> Result<(), serde_json::Error> { + value + .serialize(FiniteCheck) + .map_err(|error| ::custom(error.0)) +} + +#[derive(Debug)] +pub(crate) struct CheckError(String); + +impl fmt::Display for CheckError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(&self.0) + } +} + +impl std::error::Error for CheckError {} + +impl ser::Error for CheckError { + fn custom(message: T) -> Self { + Self(message.to_string()) + } +} + +fn check_float(value: f64) -> Result<(), CheckError> { + if value.is_finite() { + return Ok(()); + } + // Go formats the value with `strconv.FormatFloat(v, 'g', -1, bits)`. + let formatted = if value.is_nan() { + "NaN" + } else if value.is_sign_positive() { + "+Inf" + } else { + "-Inf" + }; + Err(CheckError(format!("unsupported value: {formatted}"))) +} + +struct FiniteCheck; + +impl ser::Serializer for FiniteCheck { + type Error = CheckError; + type Ok = (); + type SerializeMap = Self; + type SerializeSeq = Self; + type SerializeStruct = Self; + type SerializeStructVariant = Self; + type SerializeTuple = Self; + type SerializeTupleStruct = Self; + type SerializeTupleVariant = Self; + + fn serialize_bool(self, _: bool) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_i8(self, _: i8) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_i16(self, _: i16) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_i32(self, _: i32) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_i64(self, _: i64) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_i128(self, _: i128) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_u8(self, _: u8) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_u16(self, _: u16) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_u32(self, _: u32) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_u64(self, _: u64) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_u128(self, _: u128) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_f32(self, value: f32) -> Result<(), CheckError> { + check_float(f64::from(value)) + } + + fn serialize_f64(self, value: f64) -> Result<(), CheckError> { + check_float(value) + } + + fn serialize_char(self, _: char) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_str(self, _: &str) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_bytes(self, _: &[u8]) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_none(self) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_some(self, value: &T) -> Result<(), CheckError> { + value.serialize(self) + } + + fn serialize_unit(self) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_unit_struct(self, _: &'static str) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_unit_variant( + self, + _: &'static str, + _: u32, + _: &'static str, + ) -> Result<(), CheckError> { + Ok(()) + } + + fn serialize_newtype_struct( + self, + _: &'static str, + value: &T, + ) -> Result<(), CheckError> { + value.serialize(self) + } + + fn serialize_newtype_variant( + self, + _: &'static str, + _: u32, + _: &'static str, + value: &T, + ) -> Result<(), CheckError> { + value.serialize(self) + } + + fn serialize_seq(self, _: Option) -> Result { + Ok(self) + } + + fn serialize_tuple(self, _: usize) -> Result { + Ok(self) + } + + fn serialize_tuple_struct(self, _: &'static str, _: usize) -> Result { + Ok(self) + } + + fn serialize_tuple_variant( + self, + _: &'static str, + _: u32, + _: &'static str, + _: usize, + ) -> Result { + Ok(self) + } + + fn serialize_map(self, _: Option) -> Result { + Ok(self) + } + + fn serialize_struct(self, _: &'static str, _: usize) -> Result { + Ok(self) + } + + fn serialize_struct_variant( + self, + _: &'static str, + _: u32, + _: &'static str, + _: usize, + ) -> Result { + Ok(self) + } + + fn collect_str(self, _: &T) -> Result<(), CheckError> { + Ok(()) + } +} + +impl ser::SerializeSeq for FiniteCheck { + type Error = CheckError; + type Ok = (); + + fn serialize_element(&mut self, value: &T) -> Result<(), CheckError> { + value.serialize(Self) + } + + fn end(self) -> Result<(), CheckError> { + Ok(()) + } +} + +impl ser::SerializeTuple for FiniteCheck { + type Error = CheckError; + type Ok = (); + + fn serialize_element(&mut self, value: &T) -> Result<(), CheckError> { + value.serialize(Self) + } + + fn end(self) -> Result<(), CheckError> { + Ok(()) + } +} + +impl ser::SerializeTupleStruct for FiniteCheck { + type Error = CheckError; + type Ok = (); + + fn serialize_field(&mut self, value: &T) -> Result<(), CheckError> { + value.serialize(Self) + } + + fn end(self) -> Result<(), CheckError> { + Ok(()) + } +} + +impl ser::SerializeTupleVariant for FiniteCheck { + type Error = CheckError; + type Ok = (); + + fn serialize_field(&mut self, value: &T) -> Result<(), CheckError> { + value.serialize(Self) + } + + fn end(self) -> Result<(), CheckError> { + Ok(()) + } +} + +impl ser::SerializeMap for FiniteCheck { + type Error = CheckError; + type Ok = (); + + fn serialize_key(&mut self, key: &T) -> Result<(), CheckError> { + key.serialize(Self) + } + + fn serialize_value(&mut self, value: &T) -> Result<(), CheckError> { + value.serialize(Self) + } + + fn end(self) -> Result<(), CheckError> { + Ok(()) + } +} + +impl ser::SerializeStruct for FiniteCheck { + type Error = CheckError; + type Ok = (); + + fn serialize_field( + &mut self, + _: &'static str, + value: &T, + ) -> Result<(), CheckError> { + value.serialize(Self) + } + + fn end(self) -> Result<(), CheckError> { + Ok(()) + } +} + +impl ser::SerializeStructVariant for FiniteCheck { + type Error = CheckError; + type Ok = (); + + fn serialize_field( + &mut self, + _: &'static str, + value: &T, + ) -> Result<(), CheckError> { + value.serialize(Self) + } + + fn end(self) -> Result<(), CheckError> { + Ok(()) + } +} diff --git a/rust/riverqueue/src/error.rs b/rust/riverqueue/src/error.rs new file mode 100644 index 000000000..4e39d9f2a --- /dev/null +++ b/rust/riverqueue/src/error.rs @@ -0,0 +1,510 @@ +//! Public errors. +//! +//! Every error in this crate either describes itself and exposes what caused +//! it through [`std::error::Error::source`], or is transparent and delegates +//! both its message and its source to the error it wraps. None does both, so +//! a report that prints the whole chain, such as `anyhow`'s `{:#}`, never +//! repeats a message. +//! +//! Match [`enum@Error`]'s variants for the cases a caller handles, and keep a +//! fallback arm, because the enum is `#[non_exhaustive]`: +//! +//! ```no_run +//! use riverqueue::{Client, Error, Record, sqlx}; +//! +//! # async fn example(client: Client, id: i64) -> Result<(), Error> { +//! match client.jobs().get(id).await { +//! Ok(job) => println!("job {id} is {}", job.state), +//! Err(Error::NotFound(Record::Job(_))) => println!("job {id} doesn't exist"), +//! // Database failures hold SQLx's error on either backend. +//! Err(Error::Database(sqlx::Error::PoolTimedOut)) => println!("database busy"), +//! Err(error) => return Err(error), +//! } +//! # Ok(()) +//! # } +//! ``` +//! +//! An error from a hook, middleware, or add-on crate is an +//! [`Error::Extension`] whose [`source`](Error::Extension::source) is the +//! extension's own error; downcast it to that error's type. Errors a worker +//! returns reach hooks, error handlers, and retry policies as a +//! [`WorkError`](crate::WorkError), which downcasts the same way. + +use std::fmt; + +use thiserror::Error; + +use crate::JobState; + +/// A thread-safe error source whose concrete type can be inspected by callers. +pub type BoxError = Box; + +macro_rules! context_error { + ($(#[$meta:meta])* $name:ident, $format:literal) => { + $(#[$meta])* + #[derive(Debug, Error)] + #[error($format)] + pub struct $name { + context: &'static str, + message: String, + #[source] + source: Option, + } + + impl $name { + pub(crate) fn new(context: &'static str, message: impl Into) -> Self { + Self { + context, + message: message.into(), + source: None, + } + } + + /// Returns the operation or field the error is about. + #[must_use] + pub const fn context(&self) -> &'static str { + self.context + } + + /// Returns the specific failure message. + #[must_use] + pub fn message(&self) -> &str { + &self.message + } + } + }; +} + +context_error!( + /// Invalid client, queue, subscription, or maintenance configuration. + ConfigurationError, + "invalid {context} configuration: {message}" +); +context_error!( + /// Invalid job arguments or insertion options. + JobValidationError, + "invalid {context}: {message}" +); +context_error!( + /// A failure of River's runtime, with the operation it happened in. + RuntimeError, + "{context}: {message}" +); + +/// The kind of record an [`Error::NotFound`] refers to, with the key that +/// was looked up. +#[derive(Clone, Debug, Eq, PartialEq)] +#[non_exhaustive] +pub enum Record { + /// A job, by ID. + Job(i64), + /// A queue, by name. + Queue(String), +} + +impl fmt::Display for Record { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::Job(id) => write!(formatter, "job {id}"), + Self::Queue(name) => write!(formatter, "queue {name:?}"), + } + } +} + +/// Where an extension that failed was running. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +#[non_exhaustive] +pub enum ExtensionPhase { + /// An add-on crate's extension of River's runtime, such as its own + /// insertion, claim, or rescue step. + AddOn { + /// The add-on step that failed, such as `"rescue"`, for diagnostics. + operation: &'static str, + }, + /// An [`ErrorHandler::handle_error`](crate::ErrorHandler::handle_error) + /// call. + ErrorHandler, + /// A [`Hook::insert_begin`](crate::Hook::insert_begin) hook. + InsertBeginHook, + /// An [`InsertMiddleware`](crate::InsertMiddleware). + InsertMiddleware, + /// A [`Hook::decode_insert_result`](crate::Hook::decode_insert_result) + /// hook. + InsertResultDecodeHook, + /// A [`Hook::metric_emit`](crate::Hook::metric_emit) hook. + MetricEmitHook, + /// A [`Hook::periodic_jobs_start`](crate::Hook::periodic_jobs_start) + /// hook. + PeriodicJobsStartHook, + /// An [`ErrorHandler::handle_stuck`](crate::ErrorHandler::handle_stuck) + /// call. + StuckJobHandler, +} + +impl fmt::Display for ExtensionPhase { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(match self { + Self::AddOn { operation } => return write!(formatter, "add-on {operation}"), + Self::ErrorHandler => "error handler", + Self::InsertBeginHook => "insert begin hook", + Self::InsertMiddleware => "insert middleware", + Self::InsertResultDecodeHook => "insert result decode hook", + Self::MetricEmitHook => "metric hook", + Self::PeriodicJobsStartHook => "periodic jobs start hook", + Self::StuckJobHandler => "stuck job handler", + }) + } +} + +/// Error returned by River operations. +#[derive(Debug, Error)] +#[non_exhaustive] +pub enum Error { + /// The client is already running; a client runs once at a time. + #[error("client is already running")] + AlreadyRunning, + + /// The client stopped before reaching the state being waited for, such as + /// readiness in [`RunHandle::wait_ready`](crate::RunHandle::wait_ready). + #[error("client stopped")] + ClientStopped, + + /// Client, queue, subscription, or maintenance configuration is invalid. + #[error(transparent)] + Configuration(ConfigurationError), + + /// A database operation failed. + /// + /// The payload is SQLx's error on either backend, so a caller can match a + /// failure such as [`sqlx::Error::PoolTimedOut`] or inspect + /// [`sqlx::Error::Database`] for a constraint violation. A stored row + /// River can't decode is reported as [`sqlx::Error::Decode`] or + /// [`sqlx::Error::ColumnDecode`]. + #[error(transparent)] + Database(#[from] sqlx::Error), + + /// A transactional executor belongs to another database backend. + #[error(transparent)] + DatabaseMismatch(#[from] crate::database::DatabaseMismatch), + + /// A hook, middleware, or add-on crate failed. + #[error("{phase} failed")] + Extension { + /// Where the extension was running. + phase: ExtensionPhase, + /// Original extension error. + #[source] + source: BoxError, + }, + + /// Job arguments or options are invalid. + #[error(transparent)] + InvalidJob(JobValidationError), + + /// JSON encoding or decoding failed. + #[error(transparent)] + Json(#[from] serde_json::Error), + + /// The operation needs a running job, such as completing it in a + /// transaction, but the job is in another state. + #[error("job is {state}, not running")] + JobNotRunning { + /// The job's current state. + state: JobState, + }, + + /// A running job cannot be deleted. + #[error("running jobs cannot be deleted")] + JobRunning, + + /// The requested record does not exist. + #[error("{0} not found")] + NotFound(Record), + + /// A queue is already added to this client, or its removal is still + /// waiting for its producer to stop. + #[error("queue {name:?} is already added")] + QueueAlreadyAdded { + /// The queue's name. + name: String, + }, + + /// This client doesn't work the queue. + #[error("queue {name:?} is not added to this client")] + QueueNotAdded { + /// The queue's name. + name: String, + }, + + /// A user-provided resumable step returned an error. + #[error("resumable step {name:?} failed")] + ResumableStep { + /// Name of the step that failed. + name: String, + /// Original step error. + #[source] + source: BoxError, + }, + + /// River's runtime failed. + #[error(transparent)] + Runtime(RuntimeError), + + /// An operation that spawns tasks was called outside Tokio. + #[error("{operation} requires an active Tokio runtime")] + RuntimeUnavailable { + /// Operation that requires Tokio task spawning. + operation: &'static str, + }, + + /// A spawned runtime task panicked or was cancelled. + #[error(transparent)] + RuntimeTask(tokio::task::JoinError), + + /// A client with workers cannot insert an unregistered kind by default. + #[error("job kind {kind:?} has no registered worker")] + UnknownJobKind { + /// The job kind without a registered worker. + kind: String, + }, +} + +impl Error { + pub(crate) fn configuration(message: impl Into) -> Self { + Self::configuration_context("client", message) + } + + pub(crate) fn configuration_context(context: &'static str, message: impl Into) -> Self { + Self::Configuration(ConfigurationError::new(context, message)) + } + + /// Wraps an error raised by a hook, middleware, or add-on crate. + /// + /// Insert middleware uses this to fail an insertion with its own error: + /// + /// ``` + /// use riverqueue::{Error, ExtensionPhase}; + /// + /// let error = Error::extension(ExtensionPhase::InsertMiddleware, "quota exceeded"); + /// assert_eq!(error.to_string(), "insert middleware failed"); + /// ``` + pub fn extension(phase: ExtensionPhase, source: impl Into) -> Self { + Self::Extension { + phase, + source: source.into(), + } + } + + pub(crate) fn invalid_job(message: impl Into) -> Self { + Self::invalid_job_context("job", message) + } + + pub(crate) fn invalid_job_context(context: &'static str, message: impl Into) -> Self { + Self::InvalidJob(JobValidationError::new(context, message)) + } + + pub(crate) fn runtime_context(context: &'static str, message: impl Into) -> Self { + Self::Runtime(RuntimeError::new(context, message)) + } + + pub(crate) fn runtime_source( + context: &'static str, + message: impl Into, + source: impl Into, + ) -> Self { + Self::Runtime(RuntimeError { + context, + message: message.into(), + source: Some(source.into()), + }) + } + + pub(crate) const fn from_join(error: tokio::task::JoinError) -> Self { + Self::RuntimeTask(error) + } +} + +/// Formats an error with its whole source chain, `outer: inner: innermost`, +/// for recorded job errors and River's own log lines. +/// +/// An error whose message already ends with `: {source}`, as some libraries' +/// errors (including SQLx's) do, is shortened so the source's message +/// appears once; snafu's `CleanedErrorText` applies the same rule. +pub(crate) struct Chain<'a>(pub(crate) &'a (dyn std::error::Error + 'static)); + +impl fmt::Display for Chain<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + let mut error = self.0; + let mut message = error.to_string(); + let mut separator = ""; + loop { + let Some(source) = error.source() else { + return write!(formatter, "{separator}{message}"); + }; + let source_message = source.to_string(); + let own = message + .strip_suffix(source_message.as_str()) + .and_then(|own| { + if own.is_empty() { + Some(own) + } else { + own.strip_suffix(": ") + } + }) + .unwrap_or(&message); + if !own.is_empty() { + write!(formatter, "{separator}{own}")?; + separator = ": "; + } + error = source; + message = source_message; + } + } +} + +/// Returns a caught panic's message, when it has one. +pub(crate) fn panic_message(panic: &Box) -> &str { + panic + .downcast_ref::<&str>() + .copied() + .or_else(|| panic.downcast_ref::().map(String::as_str)) + .unwrap_or("non-string panic payload") +} + +#[cfg(test)] +mod tests { + use std::error::Error as _; + + use super::*; + + /// Renders an error the way `anyhow`'s `{:#}` does. + fn report(error: &(dyn std::error::Error + 'static)) -> String { + Chain(error).to_string() + } + + #[test] + fn chains_show_each_message_once() { + #[derive(Debug, Error)] + #[error("fetching user")] + struct Outer(#[source] Repeats); + + /// Prints its source itself, as `sqlx::Error::Database` does. + #[derive(Debug, Error)] + #[error("query failed: {0}")] + struct Repeats(#[source] std::io::Error); + + #[derive(Debug, Error)] + #[error("{0}")] + struct Same(#[source] std::io::Error); + + #[derive(Debug, Error)] + #[error("retry 5")] + struct Suffix(#[source] std::io::Error); + + let error = Outer(Repeats(std::io::Error::other("connection reset"))); + assert_eq!( + report(&error), + "fetching user: query failed: connection reset" + ); + + let error = Same(std::io::Error::other("connection reset")); + assert_eq!(report(&error), "connection reset"); + + // A source's message that merely ends the outer message isn't a + // repeat. + let error = Suffix(std::io::Error::other("5")); + assert_eq!(report(&error), "retry 5: 5"); + } + + #[test] + fn database_errors_are_transparent() { + let error = Error::from(sqlx::Error::RowNotFound); + assert!(matches!(error, Error::Database(sqlx::Error::RowNotFound))); + assert_eq!(report(&error), sqlx::Error::RowNotFound.to_string()); + } + + #[test] + fn extension_preserves_concrete_source() { + let error = Error::extension( + ExtensionPhase::InsertMiddleware, + std::io::Error::other("failed"), + ); + let source = error.source().unwrap(); + + assert!(source.downcast_ref::().is_some()); + assert_eq!(report(&error), "insert middleware failed: failed"); + } + + #[test] + fn reports_never_repeat_a_message() { + let cases = [ + ( + Error::configuration("bad"), + "invalid client configuration: bad", + ), + (Error::invalid_job("bad kind"), "invalid job: bad kind"), + (Error::NotFound(Record::Job(42)), "job 42 not found"), + ( + Error::NotFound(Record::Queue("default".to_owned())), + r#"queue "default" not found"#, + ), + ( + Error::JobNotRunning { + state: JobState::Completed, + }, + "job is completed, not running", + ), + ( + Error::ResumableStep { + name: "second".to_owned(), + source: "step failed".into(), + }, + r#"resumable step "second" failed: step failed"#, + ), + ]; + for (error, expected) in cases { + assert_eq!(report(&error), expected); + } + } + + #[test] + fn work_errors_are_transparent() { + let inner = Error::extension(ExtensionPhase::ErrorHandler, "handler failed"); + let error = crate::WorkError::new(inner); + + assert_eq!(report(&error), "error handler failed: handler failed"); + assert!(matches!( + error.downcast_ref::(), + Some(Error::Extension { + phase: ExtensionPhase::ErrorHandler, + .. + }) + )); + } + + #[test] + fn structured_runtime_error_preserves_context_and_source() { + let error = Error::runtime_source( + "resumable cursor", + "cannot decode cursor", + std::io::Error::other("bad JSON"), + ); + let Error::Runtime(runtime) = &error else { + panic!("expected runtime error"); + }; + + assert_eq!(runtime.context(), "resumable cursor"); + assert_eq!(runtime.message(), "cannot decode cursor"); + assert!( + error + .source() + .unwrap() + .downcast_ref::() + .is_some() + ); + assert_eq!( + report(&error), + "resumable cursor: cannot decode cursor: bad JSON" + ); + } +} diff --git a/rust/riverqueue/src/event.rs b/rust/riverqueue/src/event.rs new file mode 100644 index 000000000..9ecc24d52 --- /dev/null +++ b/rust/riverqueue/src/event.rs @@ -0,0 +1,403 @@ +//! Bounded local client event subscriptions. + +use std::{ + collections::HashSet, + num::NonZeroUsize, + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, +}; + +use thiserror::Error; +use tokio::sync::mpsc; + +use crate::{Error, JobRow, Queue}; + +/// A client event kind. Callers must opt in to each kind explicitly. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +#[non_exhaustive] +pub enum EventKind { + /// A job reached the cancelled state. + JobCancelled, + /// A job completed successfully. + JobCompleted, + /// A job failed, whether retryable or terminal. + JobFailed, + /// A running job was interrupted during shutdown. + JobInterrupted, + /// A job was snoozed. + JobSnoozed, + /// A queue was paused. + QueuePaused, + /// A queue was resumed. + QueueResumed, +} + +/// An event emitted by this client instance. +/// +/// The enum separates job and queue payloads so an event can never contain an +/// invalid combination such as a queue event with job statistics. +#[derive(Clone, Debug)] +#[non_exhaustive] +#[allow( + clippy::large_enum_variant, + reason = "job events dominate and boxing every event would add an allocation" +)] +pub enum Event { + /// A job lifecycle event. + Job(JobEvent), + /// A queue lifecycle event. + Queue(QueueEvent), +} + +impl Event { + pub(crate) fn queue(kind: QueueEventKind, queue: Queue) -> Self { + Self::Queue(QueueEvent { kind, queue }) + } + + /// Returns this event's subscription discriminator. + #[must_use] + pub const fn kind(&self) -> EventKind { + match self { + Self::Job(event) => event.kind.as_event_kind(), + Self::Queue(event) => event.kind.as_event_kind(), + } + } + + /// Returns the job event payload, if this is a job event. + #[must_use] + pub const fn as_job(&self) -> Option<&JobEvent> { + match self { + Self::Job(event) => Some(event), + Self::Queue(_) => None, + } + } + + /// Returns the queue event payload, if this is a queue event. + #[must_use] + pub const fn as_queue(&self) -> Option<&QueueEvent> { + match self { + Self::Job(_) => None, + Self::Queue(event) => Some(event), + } + } + + pub(crate) fn job_with_statistics( + kind: JobEventKind, + job: JobRow, + statistics: JobStatistics, + ) -> Self { + Self::Job(JobEvent { + job, + kind, + statistics: Some(statistics), + }) + } +} + +/// A job lifecycle event and its valid payload. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct JobEvent { + /// Job snapshot after its state transition committed. + pub job: JobRow, + /// Job event discriminator derived from the persisted job state. + /// + /// An `available` row keeps the worker's requested retry, snooze, or + /// interruption reason because that state alone is ambiguous. Terminal, + /// retryable, and scheduled rows always determine the emitted kind. + pub kind: JobEventKind, + /// Timing information for the corresponding execution, when applicable. + pub statistics: Option, +} + +/// A job event kind. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +#[non_exhaustive] +pub enum JobEventKind { + /// A job reached the cancelled state. + Cancelled, + /// A job completed successfully. + Completed, + /// A job failed, whether retryable or terminal. + Failed, + /// A running job was interrupted during shutdown. + Interrupted, + /// A job was snoozed. + Snoozed, +} + +impl JobEventKind { + const fn as_event_kind(self) -> EventKind { + match self { + Self::Cancelled => EventKind::JobCancelled, + Self::Completed => EventKind::JobCompleted, + Self::Failed => EventKind::JobFailed, + Self::Interrupted => EventKind::JobInterrupted, + Self::Snoozed => EventKind::JobSnoozed, + } + } +} + +impl From for EventKind { + fn from(kind: JobEventKind) -> Self { + kind.as_event_kind() + } +} + +/// A queue lifecycle event and its valid payload. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct QueueEvent { + /// Queue event discriminator. + pub kind: QueueEventKind, + /// Queue snapshot after its observed state transition committed. + /// + /// Queue events are best-effort wakeups rather than a durable transition + /// log. Rapid pause/resume transitions may coalesce before a client reads + /// the persisted queue state; use storage operations when authoritative + /// current state is required. + pub queue: Queue, +} + +/// A queue event kind. +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +#[non_exhaustive] +pub enum QueueEventKind { + /// A queue was paused. + Paused, + /// A queue was resumed. + Resumed, +} + +impl QueueEventKind { + const fn as_event_kind(self) -> EventKind { + match self { + Self::Paused => EventKind::QueuePaused, + Self::Resumed => EventKind::QueueResumed, + } + } +} + +impl From for EventKind { + fn from(kind: QueueEventKind) -> Self { + kind.as_event_kind() + } +} + +/// Timing information for one execution of a job. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +#[non_exhaustive] +pub struct JobStatistics { + /// Time spent persisting the worker result. + pub complete_duration: Duration, + /// Time between the job becoming eligible and beginning work. + pub queue_wait_duration: Duration, + /// Time spent running the worker and its work extensions. + pub run_duration: Duration, +} + +/// Receiver capacity of a subscription that doesn't set one. +const DEFAULT_BUFFER_CAPACITY: NonZeroUsize = NonZeroUsize::new(1_000).unwrap(); + +/// Configuration for one event subscription. +#[derive(Clone, Debug)] +pub struct SubscribeConfig { + buffer_capacity: NonZeroUsize, + kinds: Vec, +} + +impl SubscribeConfig { + /// Creates a subscription for at least one event kind, with a receiver + /// buffer of 1,000 events. + /// + /// # Errors + /// + /// Returns an error when `kinds` is empty. + pub fn new(kinds: impl IntoIterator) -> Result { + let kinds = kinds.into_iter().collect::>(); + validate_kinds(&kinds)?; + Ok(Self { + buffer_capacity: DEFAULT_BUFFER_CAPACITY, + kinds, + }) + } + + /// Returns the configuration with a receiver buffer of `capacity` + /// events. A receiver that falls further behind loses the oldest events + /// and learns how many on its next receive. + #[must_use] + pub const fn with_buffer_capacity(mut self, capacity: NonZeroUsize) -> Self { + self.buffer_capacity = capacity; + self + } + + /// Returns the bounded receiver capacity. + #[must_use] + pub const fn buffer_capacity(&self) -> NonZeroUsize { + self.buffer_capacity + } + + /// Returns the requested event kinds. + #[must_use] + pub fn kinds(&self) -> &[EventKind] { + &self.kinds + } + + pub(crate) fn into_parts(self) -> (NonZeroUsize, Vec) { + (self.buffer_capacity, self.kinds) + } +} + +/// Error returned while receiving client events. +#[derive(Debug, Error)] +#[non_exhaustive] +pub enum EventRecvError { + /// The client dropped the contained number of events because the receiver + /// lagged its bounded buffer. The next call resumes at the oldest retained + /// event. + #[error("event receiver lagged by {0} events")] + Lagged(u64), + /// The client event channel closed. + #[error("event channel closed")] + Closed, +} + +/// A filtered receiver for locally generated client events. +/// +/// Job events are emitted only after their state transition commits. Concurrent +/// jobs and completion batches have no global event-ordering guarantee; use the +/// job ID and persisted timestamps when an application needs stable ordering. +pub struct EventReceiver { + dropped: Arc, + receiver: mpsc::Receiver, +} + +impl std::fmt::Debug for EventReceiver { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("EventReceiver") + .field("dropped", &self.dropped.load(Ordering::Acquire)) + .field("closed", &self.receiver.is_closed()) + .finish_non_exhaustive() + } +} + +impl EventReceiver { + pub(crate) fn new(dropped: Arc, receiver: mpsc::Receiver) -> Self { + Self { dropped, receiver } + } + + /// Receives the next requested event. + /// + /// # Errors + /// + /// Returns [`EventRecvError::Lagged`] with the number of events dropped + /// because the receiver fell behind, after which receiving resumes, and + /// [`EventRecvError::Closed`] once the client is gone. + /// + /// # Cancel safety + /// + /// This method is cancel safe: dropping its future before it completes + /// loses no event, and the next call receives it. + pub async fn recv(&mut self) -> Result { + let dropped = self.dropped.swap(0, Ordering::AcqRel); + if dropped > 0 { + return Err(EventRecvError::Lagged(dropped)); + } + self.receiver.recv().await.ok_or(EventRecvError::Closed) + } +} + +/// Yields what [`EventReceiver::recv`] returns, including +/// [`EventRecvError::Lagged`], and ends once the client is gone instead of +/// yielding [`EventRecvError::Closed`]. +impl futures_util::Stream for EventReceiver { + type Item = Result; + + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + context: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + let dropped = self.dropped.swap(0, Ordering::AcqRel); + if dropped > 0 { + return std::task::Poll::Ready(Some(Err(EventRecvError::Lagged(dropped)))); + } + self.receiver.poll_recv(context).map(|event| event.map(Ok)) + } +} + +pub(crate) fn validate_kinds(kinds: &[EventKind]) -> Result, Error> { + if kinds.is_empty() { + return Err(Error::configuration_context( + "event subscription", + "event subscription requires at least one event kind".to_owned(), + )); + } + Ok(kinds.iter().copied().collect()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn typed_event_kinds_map_to_subscription_kinds() { + assert_eq!( + EventKind::from(JobEventKind::Completed), + EventKind::JobCompleted + ); + assert_eq!( + EventKind::from(QueueEventKind::Resumed), + EventKind::QueueResumed + ); + } + + #[tokio::test] + async fn receiver_streams_lags_events_and_ends_when_closed() { + use futures_util::StreamExt as _; + + let dropped = Arc::new(AtomicU64::new(2)); + let (sender, receiver) = mpsc::channel(1); + let mut events = EventReceiver::new(Arc::clone(&dropped), receiver); + let now = chrono::Utc::now(); + sender + .send(Event::queue( + QueueEventKind::Paused, + Queue { + created_at: now, + metadata: serde_json::Map::new(), + metadata_text: "{}".to_owned(), + name: "default".to_owned(), + paused_at: Some(now), + updated_at: now, + }, + )) + .await + .unwrap(); + drop(sender); + + assert!(matches!( + events.next().await, + Some(Err(EventRecvError::Lagged(2))) + )); + assert_eq!( + events.next().await.unwrap().unwrap().kind(), + EventKind::QueuePaused + ); + assert!(events.next().await.is_none()); + } + + #[test] + fn subscription_is_valid_by_construction() { + assert!(SubscribeConfig::new([]).is_err()); + let capacity = NonZeroUsize::new(42).unwrap(); + let config = SubscribeConfig::new([EventKind::JobCompleted]) + .unwrap() + .with_buffer_capacity(capacity); + assert_eq!(config.buffer_capacity(), capacity); + assert_eq!(config.kinds(), [EventKind::JobCompleted]); + } +} diff --git a/rust/riverqueue/src/extension.rs b/rust/riverqueue/src/extension.rs new file mode 100644 index 000000000..bee8a944b --- /dev/null +++ b/rust/riverqueue/src/extension.rs @@ -0,0 +1,833 @@ +//! Ordered hooks, middleware, and plugin registration. + +use std::{fmt, future::Future, pin::Pin, sync::Arc, time::Duration}; + +use chrono::{DateTime, Utc}; + +use crate::{BoxError, Error, InsertParams, JobRow, PeriodicJobs, WorkContext, WorkOutcome}; + +/// Cloneable worker error passed to hooks, error handlers, and retry +/// policies. +/// +/// A `WorkError` is transparent: it displays as the error it wraps and +/// reports that error's source as its own, so a report of the whole chain +/// shows each message once. Use [`WorkError::get_ref`] or +/// [`WorkError::downcast_ref`] to inspect the wrapped error itself. +#[derive(Clone)] +pub struct WorkError { + inner: Arc, +} + +impl WorkError { + /// Wraps an error, for example one a [`WorkMiddleware`] or + /// [`Hook::work_end`] returns in place of the worker's result. + pub fn new(error: impl Into) -> Self { + Self { + inner: Arc::from(error.into()), + } + } + + /// Returns the wrapped worker error for inspection. + #[must_use] + pub fn get_ref(&self) -> &(dyn std::error::Error + Send + Sync + 'static) { + self.inner.as_ref() + } + + /// Returns the wrapped worker error as a `T`, if it is one. + #[must_use] + pub fn downcast_ref(&self) -> Option<&T> { + self.inner.downcast_ref() + } +} + +impl fmt::Debug for WorkError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Debug::fmt(&self.inner, formatter) + } +} + +impl fmt::Display for WorkError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt::Display::fmt(&self.inner, formatter) + } +} + +impl std::error::Error for WorkError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + self.inner.source() + } +} + +/// A worker panic, reported as [`WorkResult::Panicked`]. +#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)] +#[error("worker panicked: {message}")] +pub struct PanicError { + message: String, +} + +impl PanicError { + pub(crate) fn new(message: impl Into) -> Self { + Self { + message: message.into(), + } + } + + /// Returns the panic's message, or a placeholder for a panic whose + /// payload isn't a string. + #[must_use] + pub fn message(&self) -> &str { + &self.message + } +} + +/// Error a worker returns when it stops because its +/// [`WorkContext::cancellation_token`](crate::WorkContext::cancellation_token) +/// was cancelled. +/// +/// When a client's hard shutdown cancels a job, a worker that returns this +/// error (directly or +/// anywhere in its error's source chain) is treated as interrupted: the job +/// becomes available again with the attempt refunded and no error recorded. +/// Any other error returned during shutdown is recorded and retried like an +/// ordinary failure, so a job that genuinely fails while the client stops +/// still consumes its attempt. +/// +/// ``` +/// use riverqueue::{WorkCancelled, WorkContext, WorkOutcome}; +/// +/// async fn work(context: WorkContext) -> Result { +/// tokio::select! { +/// () = context.cancellation_token().cancelled() => Err(WorkCancelled), +/// () = tokio::time::sleep(std::time::Duration::from_secs(1)) => Ok(WorkOutcome::Complete), +/// } +/// } +/// ``` +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, thiserror::Error)] +#[error("job work cancelled")] +pub struct WorkCancelled; + +impl WorkCancelled { + /// Whether `error` or any error in its source chain is [`WorkCancelled`]. + /// + /// Errors wrapped by `std::io::Error::other` are inspected as well, + /// because `io::Error` does not expose its payload as a source. + #[must_use] + pub fn is_in_chain(error: &(dyn std::error::Error + 'static)) -> bool { + let mut current = Some(error); + while let Some(error) = current { + if error.is::() { + return true; + } + if let Some(payload) = error + .downcast_ref::() + .and_then(std::io::Error::get_ref) + && Self::is_in_chain(payload) + { + return true; + } + current = error.source(); + } + false + } +} + +/// Name of an internal runtime metric emitted to hooks. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +#[non_exhaustive] +pub enum MetricName { + /// Duration of one successful available-job fetch. + JobGetAvailableDuration, + /// Number of rows claimed by one successful available-job fetch. + JobGetAvailableCount, +} + +/// Strongly typed metric emitted by River without installing a recorder. +#[derive(Clone, Copy, Debug)] +#[non_exhaustive] +pub enum Metric { + /// Duration of one successful available-job fetch. + JobGetAvailableDuration(Duration), + /// Number of rows claimed by one successful available-job fetch. + JobGetAvailableCount(u64), +} + +impl Metric { + /// Stable metric name. + #[must_use] + pub const fn name(self) -> MetricName { + match self { + Self::JobGetAvailableDuration(_) => MetricName::JobGetAvailableDuration, + Self::JobGetAvailableCount(_) => MetricName::JobGetAvailableCount, + } + } +} + +/// A job about to be inserted, as seen by hooks and insertion middleware. +/// +/// River resolves options, validates them, and computes the unique key before +/// any extension runs, so changing the arguments, queue, or schedule here +/// doesn't change the job's uniqueness. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct InsertContext { + /// Serialized arguments that will be persisted, as exact JSON text. + /// Replace them with [`encode_args`](crate::encoding::encode_args) to keep + /// River's argument encoding. + pub encoded_args: Box, + /// Stable job kind. + pub kind: String, + /// Resolved insertion options. + pub opts: InsertParams, + /// State the job is inserted in: available, pending, or scheduled. + pub state: crate::JobState, + /// Creation time to persist instead of the database's current time. + pub(crate) created_at: Option>, + /// Unique key hash computed from the original insertion. + pub(crate) unique_key: Option>, + /// Bitmask of states in which the unique key is enforced. + pub(crate) unique_states: Option, +} + +/// Public summary of a worker result passed to extensions. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub enum WorkResult { + /// Worker requested cancellation. + Cancelled, + /// Worker completed successfully. + Completed, + /// Worker requested terminal discard. + Discarded, + /// Worker returned an error. + Failed(WorkError), + /// Worker panicked. + Panicked(PanicError), + /// Worker was aborted after ignoring cancellation. + Aborted, + /// Worker returned because its client was shutting down. + Interrupted, + /// Worker requested a snooze. + Snoozed(Duration), +} + +/// A thread-safe boxed future, used where River erases extension types. +pub(crate) type BoxFuture<'a, T> = Pin + Send + 'a>>; + +/// Lifecycle hooks. Hooks run in registration order and observe or adjust +/// jobs without wrapping River's operations; use [`InsertMiddleware`] or +/// [`WorkMiddleware`] to wrap them. +/// +/// Every method has a default no-op implementation, so implement only the +/// ones you need. Methods are ordinary `async fn`s: +/// +/// ``` +/// use riverqueue::{BoxError, Hook, InsertContext}; +/// +/// struct TagEverything; +/// +/// impl Hook for TagEverything { +/// async fn insert_begin(&self, insert: &mut InsertContext) -> Result<(), BoxError> { +/// insert.opts.tags.push("tagged".to_owned()); +/// Ok(()) +/// } +/// } +/// ``` +/// +/// An error returned from a hook fails the operation it observes and is +/// reported as [`Error::Extension`] with the hook's error as its source. +pub trait Hook: Send + Sync + 'static { + /// Decodes a persisted row before River returns it from an insertion. + /// + /// This is the inverse of any storage transformation performed by + /// [`Hook::insert_begin`]. Decode hooks run in reverse registration order + /// so that nested transformations compose. + fn decode_insert_result( + &self, + job: &mut JobRow, + ) -> impl Future> + Send { + let _ = job; + std::future::ready(Ok(())) + } + + /// Runs for each job inside insertion middleware, before the job is + /// written. It may change the job's arguments, options, or initial + /// state. + fn insert_begin( + &self, + insert: &mut InsertContext, + ) -> impl Future> + Send { + let _ = insert; + std::future::ready(Ok(())) + } + + /// Observes a runtime metric. Failures and panics are logged and don't + /// affect the operation that produced the metric. + fn metric_emit(&self, metric: Metric) -> impl Future> + Send { + let _ = metric; + std::future::ready(Ok(())) + } + + /// Runs when this client's periodic job enqueuer starts, which happens + /// each time the client is elected leader. + fn periodic_jobs_start( + &self, + jobs: &PeriodicJobs, + ) -> impl Future> + Send { + let _ = jobs; + std::future::ready(Ok(())) + } + + /// Runs inside work middleware, before the job's arguments are decoded + /// and the worker runs. It may change the job, for example to decode + /// arguments another hook or middleware transformed on insertion. + /// + /// An error fails the attempt with that error; the worker and + /// [`Hook::work_end`] don't run. + fn work_begin( + &self, + context: &WorkContext, + job: &mut JobRow, + ) -> impl Future> + Send { + let _ = (context, job); + std::future::ready(Ok(())) + } + + /// Runs inside work middleware, after the worker returns, and returns + /// the attempt's result. + /// + /// The returned result replaces the + /// worker's, so a hook should return `result` unchanged unless it means + /// to change the outcome, for example to turn a specific error into a + /// snooze. Hooks run in registration order, each receiving the previous + /// hook's result. It doesn't run when the worker panics. + fn work_end( + &self, + context: &WorkContext, + job: &JobRow, + result: Result, + ) -> impl Future> + Send { + let _ = (context, job); + std::future::ready(result) + } +} + +/// Object-safe form of [`Hook`] that River stores after registration. +pub(crate) trait DynHook: Send + Sync + 'static { + fn decode_insert_result<'a>(&'a self, job: &'a mut JobRow) -> BoxFuture<'a, Result<(), Error>>; + fn insert_begin<'a>( + &'a self, + insert: &'a mut InsertContext, + ) -> BoxFuture<'a, Result<(), Error>>; + fn metric_emit(&self, metric: Metric) -> BoxFuture<'_, Result<(), Error>>; + fn periodic_jobs_start<'a>( + &'a self, + jobs: &'a PeriodicJobs, + ) -> BoxFuture<'a, Result<(), Error>>; + fn work_begin<'a>( + &'a self, + context: &'a WorkContext, + job: &'a mut JobRow, + ) -> BoxFuture<'a, Result<(), BoxError>>; + fn work_end<'a>( + &'a self, + context: &'a WorkContext, + job: &'a JobRow, + result: Result, + ) -> BoxFuture<'a, Result>; +} + +fn hook_error(phase: crate::ExtensionPhase) -> impl FnOnce(BoxError) -> Error { + move |source| Error::Extension { phase, source } +} + +impl DynHook for H { + fn decode_insert_result<'a>(&'a self, job: &'a mut JobRow) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(async move { + Hook::decode_insert_result(self, job) + .await + .map_err(hook_error(crate::ExtensionPhase::InsertResultDecodeHook)) + }) + } + + fn insert_begin<'a>( + &'a self, + insert: &'a mut InsertContext, + ) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(async move { + Hook::insert_begin(self, insert) + .await + .map_err(hook_error(crate::ExtensionPhase::InsertBeginHook)) + }) + } + + fn metric_emit(&self, metric: Metric) -> BoxFuture<'_, Result<(), Error>> { + // A panicking metric hook would otherwise unwind the queue's producer + // and abort every job it's working. + Box::pin(recover_extension_panic( + crate::ExtensionPhase::MetricEmitHook, + async move { Hook::metric_emit(self, metric).await }, + )) + } + + fn periodic_jobs_start<'a>( + &'a self, + jobs: &'a PeriodicJobs, + ) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(async move { + Hook::periodic_jobs_start(self, jobs) + .await + .map_err(hook_error(crate::ExtensionPhase::PeriodicJobsStartHook)) + }) + } + + fn work_begin<'a>( + &'a self, + context: &'a WorkContext, + job: &'a mut JobRow, + ) -> BoxFuture<'a, Result<(), BoxError>> { + Box::pin(Hook::work_begin(self, context, job)) + } + + fn work_end<'a>( + &'a self, + context: &'a WorkContext, + job: &'a JobRow, + result: Result, + ) -> BoxFuture<'a, Result> { + Box::pin(Hook::work_end(self, context, job, result)) + } +} + +/// Middleware wrapping each insertion, including batches. +/// +/// Middleware sees every job in an insertion at once and decides whether and how to continue by calling +/// [`InsertNext::run`]. Middleware registered first is outermost. It can +/// change jobs before passing them on, observe or change the results, wrap +/// the insertion in a span or timer, or return early without inserting. +/// +/// ``` +/// use riverqueue::{Error, InsertContext, InsertMiddleware, InsertNext, InsertedJob}; +/// +/// struct CountInserts; +/// +/// impl InsertMiddleware for CountInserts { +/// async fn insert_many( +/// &self, +/// jobs: Vec, +/// next: InsertNext<'_>, +/// ) -> Result, Error> { +/// let count = jobs.len(); +/// let inserted = next.run(jobs).await?; +/// println!("inserted {count} jobs"); +/// Ok(inserted) +/// } +/// } +/// ``` +pub trait InsertMiddleware: Send + Sync + 'static { + /// Wraps the insertion of `jobs`. + fn insert_many( + &self, + jobs: Vec, + next: InsertNext<'_>, + ) -> impl Future, Error>> + Send; +} + +/// Object-safe form of [`InsertMiddleware`]. +pub(crate) trait DynInsertMiddleware: Send + Sync + 'static { + fn insert_many<'a>( + &'a self, + jobs: Vec, + next: InsertNext<'a>, + ) -> BoxFuture<'a, Result, Error>>; +} + +impl DynInsertMiddleware for M { + fn insert_many<'a>( + &'a self, + jobs: Vec, + next: InsertNext<'a>, + ) -> BoxFuture<'a, Result, Error>> { + Box::pin(InsertMiddleware::insert_many(self, jobs, next)) + } +} + +pub(crate) type InsertEndpoint<'a> = Box< + dyn FnOnce(Vec) -> BoxFuture<'a, Result, Error>> + Send + 'a, +>; + +/// The remainder of an insertion: any inner middleware followed by River's +/// persistence of the jobs. +pub struct InsertNext<'a> { + endpoint: InsertEndpoint<'a>, + remaining: &'a [Arc], +} + +impl<'a> InsertNext<'a> { + pub(crate) fn new( + middleware: &'a [Arc], + endpoint: InsertEndpoint<'a>, + ) -> Self { + Self { + endpoint, + remaining: middleware, + } + } + + /// Continues the insertion with `jobs`. + /// + /// # Errors + /// + /// Returns the error of any inner middleware, hook, or extension, or of + /// the database insertion. + pub async fn run(self, jobs: Vec) -> Result, Error> { + match self.remaining.split_first() { + Some((middleware, remaining)) => { + middleware + .insert_many( + jobs, + InsertNext { + endpoint: self.endpoint, + remaining, + }, + ) + .await + } + None => (self.endpoint)(jobs).await, + } + } +} + +impl fmt::Debug for InsertNext<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("InsertNext") + .field("remaining_middleware", &self.remaining.len()) + .finish_non_exhaustive() + } +} + +/// One job written by an insertion, as seen by [`InsertMiddleware`], which +/// receives one per inserted job, in input order. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct InsertedJob { + /// The inserted job, or the existing job a unique insertion matched. + pub job: JobRow, + /// Whether the insertion was skipped because a matching unique job + /// already existed. + pub unique_skipped_as_duplicate: bool, +} + +impl InsertedJob { + pub(crate) const fn new(job: JobRow, unique_skipped_as_duplicate: bool) -> Self { + Self { + job, + unique_skipped_as_duplicate, + } + } +} + +/// Middleware wrapping each job attempt. +/// +/// Middleware decides whether and how to continue by calling [`WorkNext::run`], and returns the attempt's result. +/// Middleware registered first is outermost. [`Hook::work_begin`], argument +/// decoding, the worker, and [`Hook::work_end`] all run inside the innermost +/// middleware, so middleware can change the job before it's decoded, wrap +/// the attempt in a span or timer, or change its result. A job whose kind has +/// no registered worker fails before any middleware runs. +/// +/// When the worker panics, the panic unwinds through middleware, and River +/// records it as a failed attempt. +/// +/// ``` +/// use riverqueue::{JobRow, WorkContext, WorkError, WorkMiddleware, WorkNext, WorkOutcome}; +/// +/// struct TimeJobs; +/// +/// impl WorkMiddleware for TimeJobs { +/// async fn work( +/// &self, +/// _context: &WorkContext, +/// job: JobRow, +/// next: WorkNext<'_>, +/// ) -> Result { +/// let kind = job.kind.clone(); +/// let started = std::time::Instant::now(); +/// let result = next.run(job).await; +/// println!("{kind} took {:?}", started.elapsed()); +/// result +/// } +/// } +/// ``` +pub trait WorkMiddleware: Send + Sync + 'static { + /// Wraps one attempt of `job`. + fn work( + &self, + context: &WorkContext, + job: JobRow, + next: WorkNext<'_>, + ) -> impl Future> + Send; +} + +/// Object-safe form of [`WorkMiddleware`]. +pub(crate) trait DynWorkMiddleware: Send + Sync + 'static { + fn work<'a>( + &'a self, + context: &'a WorkContext, + job: JobRow, + next: WorkNext<'a>, + ) -> BoxFuture<'a, Result>; +} + +impl DynWorkMiddleware for M { + fn work<'a>( + &'a self, + context: &'a WorkContext, + job: JobRow, + next: WorkNext<'a>, + ) -> BoxFuture<'a, Result> { + Box::pin(WorkMiddleware::work(self, context, job, next)) + } +} + +pub(crate) type WorkEndpoint<'a> = + Box BoxFuture<'a, Result> + Send + 'a>; + +/// The remainder of a job attempt: any inner middleware followed by River's +/// work hooks and the worker. +pub struct WorkNext<'a> { + context: &'a WorkContext, + endpoint: WorkEndpoint<'a>, + remaining: &'a [Arc], +} + +impl<'a> WorkNext<'a> { + pub(crate) fn new( + middleware: &'a [Arc], + context: &'a WorkContext, + endpoint: WorkEndpoint<'a>, + ) -> Self { + Self { + context, + endpoint, + remaining: middleware, + } + } + + /// Continues the attempt with `job`. + /// + /// # Errors + /// + /// Returns the worker's error, or the error of an inner middleware, a + /// work hook, or argument decoding. + pub async fn run(self, job: JobRow) -> Result { + match self.remaining.split_first() { + Some((middleware, remaining)) => { + middleware + .work( + self.context, + job, + WorkNext { + context: self.context, + endpoint: self.endpoint, + remaining, + }, + ) + .await + } + None => (self.endpoint)(job).await, + } + } +} + +impl fmt::Debug for WorkNext<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("WorkNext") + .field("remaining_middleware", &self.remaining.len()) + .finish_non_exhaustive() + } +} + +/// Retry scheduling policy for ordinary worker errors and panics. +pub trait RetryPolicy: Send + Sync + 'static { + /// Returns the delay before another attempt of `job`, whose attempt + /// failed with `error`. + fn next_retry(&self, job: &JobRow, error: &WorkError, now: DateTime) -> Duration; +} + +/// River's quartic retry policy with compatibility jitter. +#[derive(Clone, Copy, Debug, Default)] +pub struct DefaultRetryPolicy { + seed: u64, +} + +impl DefaultRetryPolicy { + /// Uses a deterministic jitter seed, primarily for reproducible tests. + #[must_use] + pub const fn with_seed(seed: u64) -> Self { + Self { seed } + } +} + +impl RetryPolicy for DefaultRetryPolicy { + fn next_retry(&self, job: &JobRow, _error: &WorkError, now: DateTime) -> Duration { + crate::client::default_retry_delay(job, now, self.seed) + } +} + +/// Result override returned by an error handler. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +#[non_exhaustive] +pub enum ErrorHandlerDecision { + /// Continue normal retry or discard handling. + #[default] + Continue, + /// Cancel immediately regardless of remaining attempts. + Cancel, +} + +/// Handler invoked for worker errors, panics, and stuck jobs. +/// +/// Both methods have default implementations. Handler errors and panics are +/// logged and don't change how River handles the job. +pub trait ErrorHandler: Send + Sync + 'static { + /// Called when a worker returns an error, panics, or is aborted. Returning + /// [`ErrorHandlerDecision::Cancel`] cancels the job regardless of its + /// remaining attempts. + fn handle_error( + &self, + context: &WorkContext, + job: &JobRow, + result: &WorkResult, + ) -> impl Future> + Send { + let _ = (context, job, result); + std::future::ready(Ok(ErrorHandlerDecision::default())) + } + + /// Called when a job keeps running past its cancellation grace period. + fn handle_stuck(&self, job: &JobRow) -> impl Future> + Send { + let _ = job; + std::future::ready(Ok(())) + } +} + +/// Object-safe form of [`ErrorHandler`]. +pub(crate) trait DynErrorHandler: Send + Sync + 'static { + fn handle_error<'a>( + &'a self, + context: &'a WorkContext, + job: &'a JobRow, + result: &'a WorkResult, + ) -> BoxFuture<'a, Result>; + fn handle_stuck<'a>(&'a self, job: &'a JobRow) -> BoxFuture<'a, Result<(), Error>>; +} + +impl DynErrorHandler for H { + fn handle_error<'a>( + &'a self, + context: &'a WorkContext, + job: &'a JobRow, + result: &'a WorkResult, + ) -> BoxFuture<'a, Result> { + Box::pin(recover_extension_panic( + crate::ExtensionPhase::ErrorHandler, + async move { ErrorHandler::handle_error(self, context, job, result).await }, + )) + } + + fn handle_stuck<'a>(&'a self, job: &'a JobRow) -> BoxFuture<'a, Result<(), Error>> { + Box::pin(recover_extension_panic( + crate::ExtensionPhase::StuckJobHandler, + async move { ErrorHandler::handle_stuck(self, job).await }, + )) + } +} + +/// Awaits an extension whose failures River only logs, treating a panic +/// like a returned error. For error handlers this matches Go's +/// `invokeErrorHandler`: the job's result is still persisted rather than the +/// panic unwinding the executor and leaving the job running. +async fn recover_extension_panic( + phase: crate::ExtensionPhase, + handler: impl Future>, +) -> Result { + use futures_util::FutureExt as _; + + match std::panic::AssertUnwindSafe(handler).catch_unwind().await { + Ok(result) => result.map_err(hook_error(phase)), + Err(panic) => Err(hook_error(phase)( + format!("panicked: {}", crate::error::panic_message(&panic)).into(), + )), + } +} + +/// A set of extensions installed together, such as a tracing integration +/// that needs a hook and middleware. +/// +/// ``` +/// use riverqueue::{BoxError, Extensions, Hook, InsertContext, Plugin}; +/// +/// struct AuditHook; +/// +/// impl Hook for AuditHook { +/// async fn insert_begin(&self, insert: &mut InsertContext) -> Result<(), BoxError> { +/// println!("inserting {}", insert.kind); +/// Ok(()) +/// } +/// } +/// +/// struct Audit; +/// +/// impl Plugin for Audit { +/// fn install(&self, extensions: &mut Extensions) { +/// extensions.hook(AuditHook); +/// } +/// } +/// ``` +pub trait Plugin: Send + Sync + 'static { + /// Registers the plugin's hooks and middleware. + fn install(&self, extensions: &mut Extensions); +} + +/// Registrar through which a [`Plugin`] adds hooks and middleware. +/// +/// Extensions are appended after any registered earlier, in the order the +/// plugin adds them. +#[derive(Default)] +pub struct Extensions { + pub(crate) hooks: Vec>, + pub(crate) insert_middleware: Vec>, + pub(crate) work_middleware: Vec>, +} + +impl Extensions { + /// Adds a lifecycle hook. + pub fn hook(&mut self, hook: H) -> &mut Self { + self.hooks.push(Arc::new(hook)); + self + } + + /// Adds insertion middleware. + pub fn insert_middleware(&mut self, middleware: M) -> &mut Self { + self.insert_middleware.push(Arc::new(middleware)); + self + } + + /// Adds worker middleware. + pub fn work_middleware(&mut self, middleware: M) -> &mut Self { + self.work_middleware.push(Arc::new(middleware)); + self + } +} + +impl fmt::Debug for Extensions { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("Extensions") + .field("hooks", &self.hooks.len()) + .field("insert_middleware", &self.insert_middleware.len()) + .field("work_middleware", &self.work_middleware.len()) + .finish() + } +} diff --git a/rust/riverqueue/src/guide.rs b/rust/riverqueue/src/guide.rs new file mode 100644 index 000000000..553519bd3 --- /dev/null +++ b/rust/riverqueue/src/guide.rs @@ -0,0 +1,7 @@ +//! Guides for running River in production. +//! +//! These pages contain no API items; they document how River behaves when +//! deployed alongside other River implementations. + +#[doc = include_str!("../docs/mixed-deployments.md")] +pub mod mixed_deployments {} diff --git a/rust/riverqueue/src/job.rs b/rust/riverqueue/src/job.rs new file mode 100644 index 000000000..a099fcf17 --- /dev/null +++ b/rust/riverqueue/src/job.rs @@ -0,0 +1,1066 @@ +//! Persisted and typed job values. + +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize, de::DeserializeOwned}; +use serde_json::{Map, Value, value::RawValue}; + +use crate::{PRIORITY_DEFAULT, QUEUE_DEFAULT}; + +mod attempt_error; +mod metadata; + +pub use metadata::JobMetadata; + +/// Arguments for a typed River job. +pub trait JobArgs: DeserializeOwned + Send + Serialize + Sync + 'static { + /// Stable job kind stored with each job. + const KIND: &'static str; + + /// Former kind names handled by the same worker during safe renames. + fn kind_aliases() -> &'static [&'static str] { + &[] + } + + /// Job-type insertion defaults. Options set for a single insertion + /// replace these. Per-call metadata replaces the defaults' metadata + /// entirely, as in River Go, so it doesn't keep any of their keys. + fn default_insert_opts() -> InsertOpts { + InsertOpts::default() + } + + /// Literal JSON path components selected for argument-scoped uniqueness. + fn unique_fields() -> &'static [&'static [&'static str]] { + &[] + } +} + +/// An atomic insertion batch that can contain multiple job argument types. +/// +/// Items retain their [`JobArgs`] insertion defaults and may additionally set +/// per-item [`InsertOpts`]. Results are returned in the same order. +#[derive(Debug, Default)] +pub struct InsertBatch { + pub(crate) items: Vec, +} + +impl InsertBatch { + /// Creates an empty batch. + #[must_use] + pub const fn new() -> Self { + Self { items: Vec::new() } + } + + /// Returns whether the batch contains no jobs. + #[must_use] + pub fn is_empty(&self) -> bool { + self.items.is_empty() + } + + /// Returns the number of jobs in the batch. + #[must_use] + pub fn len(&self) -> usize { + self.items.len() + } + + /// Appends a job using its job-type defaults. + /// + /// Arguments are encoded immediately. If encoding fails, the error is + /// returned when the batch is inserted and no job in it is inserted. + pub fn push(&mut self, args: A) -> &mut Self { + self.push_with(args, InsertOpts::default()) + } + + /// Appends a job with options overlaid on its job-type defaults. + /// + /// Arguments are encoded immediately. If encoding fails, the error is + /// returned when the batch is inserted and no job in it is inserted. + #[expect( + clippy::needless_pass_by_value, + reason = "the batch takes ownership of its jobs" + )] + pub fn push_with(&mut self, args: A, opts: InsertOpts) -> &mut Self { + self.items.push(InsertBatchItem { + defaults: A::default_insert_opts(), + encoded_args: crate::encoding::encode_args(&args), + kind: A::KIND, + opts, + unique_fields: A::unique_fields(), + }); + self + } +} + +#[derive(Debug)] +pub(crate) struct InsertBatchItem { + pub(crate) defaults: InsertOpts, + pub(crate) encoded_args: Result, serde_json::Error>, + pub(crate) kind: &'static str, + pub(crate) opts: InsertOpts, + pub(crate) unique_fields: &'static [&'static [&'static str]], +} + +/// A failed job attempt persisted in `river_job.errors`. +/// +/// Attempt errors deserialize from JSON like River Go's `encoding/json` +/// decodes them, which requires a JSON deserializer such as `serde_json`'s. +/// Job rows read from the database decode their attempt errors leniently +/// instead: an element written by another tool or edited by hand decodes on a +/// best effort basis rather than making its row unreadable. +#[derive(Clone, Debug, Eq, PartialEq, Serialize)] +#[non_exhaustive] +pub struct AttemptError { + /// Time at which the error occurred. It serializes as RFC 3339 with the + /// shortest fractional seconds. + #[serde(serialize_with = "crate::encoding::go_time::serialize")] + pub at: DateTime, + /// Attempt number on which the error occurred. + pub attempt: i16, + /// Stringified worker error or panic value. + pub error: String, + /// Backtrace for a panic, otherwise empty. + pub trace: String, +} + +impl AttemptError { + /// Creates a persisted attempt error without a panic trace. + #[must_use] + pub fn new(at: DateTime, attempt: i16, error: impl Into) -> Self { + Self { + at, + attempt, + error: error.into(), + trace: String::new(), + } + } + + /// Sets the captured panic trace. + #[must_use] + pub fn with_trace(mut self, trace: impl Into) -> Self { + self.trace = trace.into(); + self + } +} + +/// Partial options applied while inserting a job. +/// +/// Options declared by [`JobArgs::default_insert_opts`] are overlaid on River +/// and client defaults. Options supplied to an insertion call are then overlaid +/// on the job-type options. A value is therefore never treated as "unset" +/// merely because it happens to equal River's default. +#[derive(Clone, Debug, Default)] +pub struct InsertOpts { + /// Options for an exact-version extension, which River carries to the + /// extension's insert hook without persisting them. + pub(crate) extension_options: Map, + max_attempts: Option, + metadata: Option, + pending: Option, + priority: Option, + queue: Option, + scheduled_at: ScheduleOverride, + tags: Option>, + unique: Option, +} + +/// How an [`InsertOpts`] layer affects a job's schedule. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum ScheduleOverride { + /// Schedule the job no earlier than this time. + At(DateTime), + /// Make the job immediately eligible, replacing any schedule from a lower + /// layer. + Immediate, + /// Keep the schedule from a lower layer, or run immediately when none + /// sets one. + #[default] + Inherit, +} + +impl InsertOpts { + /// Returns the configured maximum attempts override. + #[must_use] + pub const fn max_attempts(&self) -> Option { + self.max_attempts + } + + /// Returns the configured metadata replacement. + #[must_use] + pub const fn metadata(&self) -> Option<&JobMetadata> { + self.metadata.as_ref() + } + + /// Returns the configured pending-state override. + #[must_use] + pub const fn pending(&self) -> Option { + self.pending + } + + /// Returns the configured priority override. + #[must_use] + pub const fn priority(&self) -> Option { + self.priority + } + + /// Returns the configured queue override. + #[must_use] + pub fn queue(&self) -> Option<&str> { + self.queue.as_deref() + } + + /// Returns how these options affect the job's schedule. + #[must_use] + pub const fn scheduled_at(&self) -> ScheduleOverride { + self.scheduled_at + } + + /// Returns the configured tags replacement. + #[must_use] + pub fn tags(&self) -> Option<&[String]> { + self.tags.as_deref() + } + + /// Returns the configured uniqueness replacement. + #[must_use] + pub const fn unique(&self) -> Option<&UniqueOpts> { + self.unique.as_ref() + } + + /// Overrides the maximum number of attempts, including the first. + #[must_use] + pub const fn with_max_attempts(mut self, maximum: i16) -> Self { + self.max_attempts = Some(maximum); + self + } + + /// Replaces arbitrary JSON object metadata. + /// + /// Per-call metadata replaces the job type's default metadata as a + /// whole. + /// + /// Accepts a [`JobMetadata`] or a `serde_json::Map`. Build a + /// [`JobMetadata`] from JSON text (for example with `str::parse`) to keep + /// number tokens, such as `1e400` or integers wider than 64 bits, that + /// `serde_json::Value` can't represent exactly. + #[must_use] + pub fn with_metadata(mut self, metadata: impl Into) -> Self { + self.metadata = Some(metadata.into()); + self + } + + /// Selects whether to insert in the pending state. + #[must_use] + pub const fn with_pending(mut self, pending: bool) -> Self { + self.pending = Some(pending); + self + } + + /// Overrides priority from one (highest) through four (lowest). + #[must_use] + pub const fn with_priority(mut self, priority: i16) -> Self { + self.priority = Some(priority); + self + } + + /// Overrides the queue in which the job runs. + #[must_use] + pub fn with_queue(mut self, queue: impl Into) -> Self { + self.queue = Some(queue.into()); + self + } + + /// Schedules the job no earlier than `scheduled_at`. + #[must_use] + pub const fn with_scheduled_at(mut self, scheduled_at: DateTime) -> Self { + self.scheduled_at = ScheduleOverride::At(scheduled_at); + self + } + + /// Explicitly overrides a job-type schedule to make the job immediately + /// eligible. + #[must_use] + pub const fn without_schedule(mut self) -> Self { + self.scheduled_at = ScheduleOverride::Immediate; + self + } + + /// Replaces searchable tags. + #[must_use] + pub fn with_tags(mut self, tags: impl IntoIterator>) -> Self { + self.tags = Some(tags.into_iter().map(Into::into).collect()); + self + } + + /// Replaces unique-job options. + #[must_use] + pub fn with_unique(mut self, unique: UniqueOpts) -> Self { + self.unique = Some(unique); + self + } + + /// Returns these options with every option set in `overrides` replacing + /// the corresponding option here. Options `overrides` leaves unset are + /// kept. + /// + /// This is how River layers call-site options over job-type defaults, and + /// how `#[river(insert_opts = ...)]` layers a function's options over the + /// derive's attribute defaults. + #[must_use] + pub fn overlay(mut self, overrides: Self) -> Self { + let Self { + extension_options, + max_attempts, + metadata, + pending, + priority, + queue, + scheduled_at, + tags, + unique, + } = overrides; + self.extension_options.extend(extension_options); + self.max_attempts = max_attempts.or(self.max_attempts); + self.metadata = metadata.or(self.metadata); + self.pending = pending.or(self.pending); + self.priority = priority.or(self.priority); + self.queue = queue.or(self.queue); + if scheduled_at != ScheduleOverride::Inherit { + self.scheduled_at = scheduled_at; + } + self.tags = tags.or(self.tags); + self.unique = unique.or(self.unique); + self + } + + pub(crate) fn resolve( + client_max_attempts: i16, + job_defaults: Self, + call_overrides: Self, + ) -> InsertParams { + let mut resolved = InsertParams { + extension_options: Map::new(), + max_attempts: client_max_attempts, + metadata: JobMetadata::default(), + pending: false, + priority: PRIORITY_DEFAULT, + queue: QUEUE_DEFAULT.to_owned(), + scheduled_at: None, + tags: Vec::new(), + unique: UniqueOpts::default(), + }; + resolved.apply(job_defaults); + resolved.apply(call_overrides); + resolved + } +} + +/// Fully resolved insertion parameters visible to insertion extensions. +/// +/// River resolves and validates these from call, job-type, client, and +/// library defaults before invoking hooks or middleware, and computes the +/// unique key from them. Extensions may still change them, but River doesn't +/// validate the changes again before persisting them. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct InsertParams { + /// Options for an exact-version extension; see [`InsertOpts`]. + pub(crate) extension_options: Map, + /// Maximum number of attempts, including the first. + pub max_attempts: i16, + /// Arbitrary JSON object metadata. + pub metadata: JobMetadata, + /// Insert in the pending state. + pub pending: bool, + /// Priority from one (highest) through four (lowest). + pub priority: i16, + /// Queue in which the job runs. + pub queue: String, + /// Earliest time the job may run. + pub scheduled_at: Option>, + /// Searchable tags. + pub tags: Vec, + /// Unique job options. + pub unique: UniqueOpts, +} + +impl InsertParams { + fn apply(&mut self, options: InsertOpts) { + self.extension_options.extend(options.extension_options); + if let Some(value) = options.max_attempts { + self.max_attempts = value; + } + if let Some(value) = options.metadata { + self.metadata = value; + } + if let Some(value) = options.pending { + self.pending = value; + } + if let Some(value) = options.priority { + self.priority = value; + } + if let Some(value) = options.queue { + self.queue = value; + } + match options.scheduled_at { + ScheduleOverride::At(value) => self.scheduled_at = Some(value), + ScheduleOverride::Immediate => self.scheduled_at = None, + ScheduleOverride::Inherit => {} + } + if let Some(value) = options.tags { + self.tags = value; + } + if let Some(value) = options.unique { + self.unique = value; + } + } +} + +/// Result of inserting a job. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct InsertResult { + /// Inserted job or the existing matching unique job. + pub job: Job, + /// Whether insertion was skipped because a unique job already existed. + pub unique_skipped_as_duplicate: bool, +} + +impl InsertResult { + /// Returns the ID of the inserted job, or of the existing job when a + /// unique insertion was skipped. + #[must_use] + pub const fn id(&self) -> i64 { + self.job.id() + } +} + +/// Type-erased result from inserting an item in an [`InsertBatch`]. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct InsertBatchResult { + /// Inserted job or the existing matching unique job. + pub job: JobRow, + /// Whether insertion was skipped because a unique job already existed. + pub unique_skipped_as_duplicate: bool, +} + +impl InsertBatchResult { + /// Returns the ID of the inserted job, or of the existing job when a + /// unique insertion was skipped. + #[must_use] + pub const fn id(&self) -> i64 { + self.job.id + } +} + +/// Typed job passed to a worker. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct Job { + /// Decoded arguments. + pub args: A, + /// Persisted job fields. + pub row: JobRow, +} + +impl Job { + /// Creates a typed job from decoded arguments and a persisted row. + /// This is primarily useful for worker unit tests. + #[must_use] + pub const fn new(args: A, row: JobRow) -> Self { + Self { args, row } + } + + /// Returns the job's database ID. + #[must_use] + pub const fn id(&self) -> i64 { + self.row.id + } +} + +/// Persisted River job fields. +/// +/// Arguments are kept as the exact JSON text stored with the job, so values +/// written by other River clients (including numbers beyond `f64` precision +/// and member order) are preserved when a row is read and passed along. Use +/// [`JobRow::decode_args`] to decode them into a typed value. +/// +/// Metadata retains the database's JSON text, including numbers outside the +/// range of [`serde_json::Value`]. Read individual fields through +/// [`JobMetadata::get_raw`] or [`JobMetadata::get`]. +#[derive(Clone, Debug, Deserialize, Serialize)] +#[non_exhaustive] +pub struct JobRow { + /// Database-generated ID. + pub id: i64, + /// Current attempt number. + pub attempt: i16, + /// Last attempt time. + pub attempted_at: Option>, + /// IDs of clients that attempted the job. + pub attempted_by: Vec, + /// Creation time. + pub created_at: DateTime, + /// Encoded job arguments as the exact JSON text stored with the job. + pub encoded_args: Box, + /// Failed attempts in chronological order. + pub errors: Vec, + /// Terminal-state time. + pub finalized_at: Option>, + /// Stable job kind. + pub kind: String, + /// Maximum attempts. + pub max_attempts: i16, + /// Arbitrary and River-reserved metadata. + pub metadata: JobMetadata, + /// Priority from one through four. + pub priority: i16, + /// Queue name. + pub queue: String, + /// Earliest run time. + pub scheduled_at: DateTime, + /// Current state. + pub state: JobState, + /// Searchable tags. + pub tags: Vec, + /// Unique hash, if any. + pub unique_key: Option>, + /// States in which this job's unique key is enforced, if any. + pub unique_states: Option>, +} + +impl JobRow { + /// Creates a minimal persisted row suitable for tests and adapters. + /// + /// Use [`encode_args`](crate::encoding::encode_args) to encode typed + /// arguments the same way River does when inserting them. + #[must_use] + pub fn new( + id: i64, + kind: impl Into, + encoded_args: Box, + now: DateTime, + ) -> Self { + Self { + attempt: 0, + attempted_at: None, + attempted_by: Vec::new(), + created_at: now, + encoded_args, + errors: Vec::new(), + finalized_at: None, + id, + kind: kind.into(), + max_attempts: crate::MAX_ATTEMPTS_DEFAULT, + metadata: JobMetadata::default(), + priority: crate::PRIORITY_DEFAULT, + queue: crate::QUEUE_DEFAULT.to_owned(), + scheduled_at: now, + state: JobState::Available, + tags: Vec::new(), + unique_key: None, + unique_states: None, + } + } + + /// Decodes the job's arguments into `T`. + /// + /// # Errors + /// + /// Returns an error when the stored arguments do not deserialize as `T`. + pub fn decode_args(&self) -> Result { + serde_json::from_str(self.encoded_args.get()) + } + + /// Decodes recorded output from metadata into a caller-selected type. + /// + /// # Errors + /// + /// Returns an error if the stored output cannot deserialize as `T`. + pub fn decode_output(&self) -> Result, serde_json::Error> { + self.metadata.get(crate::METADATA_KEY_OUTPUT) + } + + /// Returns recorded output from metadata. + #[must_use] + pub fn output(&self) -> Option<&RawValue> { + self.metadata.get_raw(crate::METADATA_KEY_OUTPUT) + } +} + +/// Persisted River job state. +/// +/// Displays as its database string, like [`JobState::as_str`]. River has +/// added states before, so matches outside this crate need a wildcard arm. +#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, PartialEq, Serialize)] +#[serde(rename_all = "lowercase")] +#[non_exhaustive] +pub enum JobState { + /// Eligible to run now. + Available, + /// Manually cancelled. + Cancelled, + /// Successfully completed. + Completed, + /// Exhausted retries. + Discarded, + /// Parked pending external action. + Pending, + /// Failed and scheduled for retry. + Retryable, + /// Actively running. + Running, + /// Scheduled for the future. + Scheduled, +} + +impl JobState { + /// All states in River's canonical bit order. + pub const ALL: &'static [Self] = &[ + Self::Available, + Self::Cancelled, + Self::Completed, + Self::Discarded, + Self::Pending, + Self::Retryable, + Self::Running, + Self::Scheduled, + ]; + + /// States required for a custom unique-state set. + pub const UNIQUE_REQUIRED: &'static [Self] = &[ + Self::Available, + Self::Pending, + Self::Running, + Self::Scheduled, + ]; + + /// Default states that enforce uniqueness. + pub const UNIQUE_DEFAULT: &'static [Self] = &[ + Self::Available, + Self::Completed, + Self::Pending, + Self::Retryable, + Self::Running, + Self::Scheduled, + ]; + + /// Canonical database string representation. + #[must_use] + pub const fn as_str(self) -> &'static str { + match self { + Self::Available => "available", + Self::Cancelled => "cancelled", + Self::Completed => "completed", + Self::Discarded => "discarded", + Self::Pending => "pending", + Self::Retryable => "retryable", + Self::Running => "running", + Self::Scheduled => "scheduled", + } + } + + /// Bit used by `river_job.unique_states`. + pub(crate) const fn unique_bit(self) -> u8 { + match self { + Self::Available => 0b0000_0001, + Self::Cancelled => 0b0000_0010, + Self::Completed => 0b0000_0100, + Self::Discarded => 0b0000_1000, + Self::Pending => 0b0001_0000, + Self::Retryable => 0b0010_0000, + Self::Running => 0b0100_0000, + Self::Scheduled => 0b1000_0000, + } + } +} + +impl std::fmt::Display for JobState { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str(self.as_str()) + } +} + +/// Failure to parse a River job state. +#[derive(Clone, Debug, Eq, PartialEq, thiserror::Error)] +#[error("unknown River job state {value:?}")] +pub struct JobStateParseError { + value: String, +} + +impl JobStateParseError { + /// Returns the unrecognized state value. + #[must_use] + pub fn value(&self) -> &str { + &self.value + } +} + +impl std::str::FromStr for JobState { + type Err = JobStateParseError; + + fn from_str(value: &str) -> Result { + match value { + "available" => Ok(Self::Available), + "cancelled" => Ok(Self::Cancelled), + "completed" => Ok(Self::Completed), + "discarded" => Ok(Self::Discarded), + "pending" => Ok(Self::Pending), + "retryable" => Ok(Self::Retryable), + "running" => Ok(Self::Running), + "scheduled" => Ok(Self::Scheduled), + _ => Err(JobStateParseError { + value: value.to_owned(), + }), + } + } +} + +impl TryFrom<&str> for JobState { + type Error = JobStateParseError; + + fn try_from(value: &str) -> Result { + value.parse() + } +} + +/// Dimensions used to deduplicate a job. +#[derive(Clone, Debug, Default)] +pub struct UniqueOpts { + /// Include encoded arguments. + pub(crate) by_args: bool, + /// Include the lower bound of this period. + pub(crate) by_period: Option, + /// Include the queue. + pub(crate) by_queue: bool, + /// States in which the key is unique. + pub(crate) by_state: Option>, + /// Exclude the job kind. + pub(crate) exclude_kind: bool, +} + +impl UniqueOpts { + /// Creates disabled uniqueness options. + #[must_use] + pub const fn new() -> Self { + Self { + by_args: false, + by_period: None, + by_queue: false, + by_state: None, + exclude_kind: false, + } + } + + /// Returns whether encoded arguments are included in the unique key. + #[must_use] + pub const fn by_args(&self) -> bool { + self.by_args + } + + /// Returns the period whose lower bound is included in the unique key. + #[must_use] + pub const fn by_period(&self) -> Option { + self.by_period + } + + /// Returns whether the queue is included in the unique key. + #[must_use] + pub const fn by_queue(&self) -> bool { + self.by_queue + } + + /// Returns the custom set of states in which the key is unique. + #[must_use] + pub fn by_state(&self) -> Option<&[JobState]> { + self.by_state.as_deref() + } + + /// Returns whether the job kind is excluded from the unique key. + #[must_use] + pub const fn exclude_kind(&self) -> bool { + self.exclude_kind + } + + /// Returns the options with encoded arguments included in the unique + /// key, or not. + /// + /// The arguments must encode to a JSON object. An empty array is treated + /// as `{}`, and inserting any other non-object arguments fails. + #[must_use] + pub const fn with_by_args(mut self, by_args: bool) -> Self { + self.by_args = by_args; + self + } + + /// Returns the options with the lower bound of `period` included in the + /// unique key. + #[must_use] + pub const fn with_by_period(mut self, period: Duration) -> Self { + self.by_period = Some(period); + self + } + + /// Returns the options with the queue included in the unique key, or + /// not. + #[must_use] + pub const fn with_by_queue(mut self, by_queue: bool) -> Self { + self.by_queue = by_queue; + self + } + + /// Returns the options with a custom set of states in which the key is + /// unique. + #[must_use] + pub fn with_by_state(mut self, states: impl IntoIterator) -> Self { + self.by_state = Some(states.into_iter().collect()); + self + } + + /// Returns the options with the job kind excluded from the unique key, + /// or not. + /// + /// Excluding the kind requires arguments, queue, or period in the key; + /// inserting with only the kind excluded fails. + #[must_use] + pub const fn with_exclude_kind(mut self, exclude_kind: bool) -> Self { + self.exclude_kind = exclude_kind; + self + } + + /// Whether no uniqueness dimension is enabled. + #[must_use] + pub fn is_empty(&self) -> bool { + !self.by_args + && self.by_period.is_none() + && !self.by_queue + && self.by_state.is_none() + && !self.exclude_kind + } + + /// Canonical persisted bitmask for the configured states. Like Go, an + /// empty custom set means the default states. + pub(crate) fn state_bitmask(&self) -> u8 { + self.by_state + .as_deref() + .filter(|states| !states.is_empty()) + .unwrap_or(JobState::UNIQUE_DEFAULT) + .iter() + .fold(0, |mask, state| mask | state.unique_bit()) + } + + /// Validates River's uniqueness invariants. + pub(crate) fn validate(&self) -> Result<(), String> { + if let Some(period) = self.by_period + && period < Duration::from_secs(1) + { + return Err("unique period must be at least one second".to_owned()); + } + // Without the kind, arguments, queue, or period, every job would + // share one key. + if self.exclude_kind && !self.by_args && !self.by_queue && self.by_period.is_none() { + return Err( + "excluding the kind from a unique key requires arguments, queue, or period" + .to_owned(), + ); + } + // Like Go, an empty custom set means the default states, which + // include every required one. + if let Some(states) = &self.by_state + && !states.is_empty() + { + let missing = JobState::UNIQUE_REQUIRED + .iter() + .filter(|state| !states.contains(state)) + .map(|state| state.as_str()) + .collect::>(); + if !missing.is_empty() { + return Err(format!( + "unique states must contain required states: {}", + missing.join(", ") + )); + } + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::MAX_ATTEMPTS_DEFAULT; + + #[test] + fn unique_states_must_include_the_required_states() { + let error = UniqueOpts::new() + .with_by_state([JobState::Available, JobState::Completed]) + .validate() + .unwrap_err(); + assert_eq!( + error, + "unique states must contain required states: pending, running, scheduled" + ); + let required = UniqueOpts::new().with_by_state(JobState::UNIQUE_REQUIRED.iter().copied()); + assert!(required.validate().is_ok()); + + // As in Go, where a non-nil empty `ByState` enables uniqueness with + // the default states. + let empty = UniqueOpts::new().with_by_state([]); + assert!(!empty.is_empty()); + assert!(empty.validate().is_ok()); + assert_eq!( + empty.state_bitmask(), + UniqueOpts::new().with_by_args(true).state_bitmask() + ); + } + + #[test] + fn excluding_the_kind_requires_another_dimension() { + let kind_only = UniqueOpts::new().with_exclude_kind(true); + assert!(!kind_only.is_empty()); + assert_eq!( + kind_only.validate().unwrap_err(), + "excluding the kind from a unique key requires arguments, queue, or period" + ); + assert!( + kind_only + .clone() + .with_by_state(JobState::UNIQUE_REQUIRED.iter().copied()) + .validate() + .is_err() + ); + for valid in [ + kind_only.clone().with_by_args(true), + kind_only.clone().with_by_queue(true), + kind_only.with_by_period(Duration::from_mins(1)), + ] { + assert!(valid.validate().is_ok(), "{valid:?}"); + } + } + + #[test] + fn insertion_options_resolve_by_layer_without_sentinels() { + let job_defaults = InsertOpts::default() + .with_max_attempts(9) + .with_priority(3) + .with_queue("job_queue") + .with_scheduled_at(Utc::now()); + let resolved = InsertOpts::resolve( + 7, + job_defaults, + InsertOpts::default().with_priority(2).without_schedule(), + ); + + assert_eq!(resolved.max_attempts, 9); + assert_eq!(resolved.priority, 2); + assert_eq!(resolved.queue, "job_queue"); + assert_eq!(resolved.scheduled_at, None); + } + + #[test] + fn per_call_metadata_replaces_default_metadata_wholesale() { + let defaults = InsertOpts::default().with_metadata( + r#"{"team":"a","river:addon":{"key":1e400},"river:shared":"default"}"# + .parse::() + .unwrap(), + ); + // Like Go, even reserved `river:` keys from the defaults are replaced. + let resolved = InsertOpts::resolve( + 7, + defaults.clone(), + InsertOpts::default().with_metadata( + r#"{"call":true,"river:shared":"call"}"#.parse::().unwrap(), + ), + ); + assert_eq!( + resolved.metadata.as_raw().get(), + r#"{"call":true,"river:shared":"call"}"# + ); + + // Without per-call metadata, the defaults apply unchanged. + let resolved = InsertOpts::resolve(7, defaults, InsertOpts::default()); + assert_eq!( + resolved.metadata.as_raw().get(), + r#"{"team":"a","river:addon":{"key":1e400},"river:shared":"default"}"# + ); + } + + #[test] + fn overlay_replaces_only_options_set_in_overrides() { + let scheduled_at = Utc::now(); + let base = InsertOpts::default() + .with_max_attempts(9) + .with_priority(3) + .with_queue("base_queue") + .with_scheduled_at(scheduled_at) + .with_tags(["base"]); + + let kept = base.clone().overlay(InsertOpts::default()); + assert_eq!(kept.max_attempts(), Some(9)); + assert_eq!(kept.queue(), Some("base_queue")); + assert_eq!(kept.scheduled_at(), ScheduleOverride::At(scheduled_at)); + assert_eq!(kept.tags(), Some(&["base".to_owned()][..])); + + let overlaid = base.overlay( + InsertOpts::default() + .with_priority(2) + .with_tags(Vec::::new()) + .with_unique(UniqueOpts::new().with_by_queue(true)) + .without_schedule(), + ); + assert_eq!(overlaid.max_attempts(), Some(9)); + assert_eq!(overlaid.priority(), Some(2)); + assert_eq!(overlaid.queue(), Some("base_queue")); + assert_eq!(overlaid.scheduled_at(), ScheduleOverride::Immediate); + assert_eq!(overlaid.tags(), Some(&[][..])); + assert!(overlaid.unique().is_some_and(UniqueOpts::by_queue)); + } + + #[test] + fn job_and_insert_result_expose_ids() { + let row = JobRow::new( + 42, + "id_test", + crate::encoding::encode_args(&serde_json::json!({})).unwrap(), + Utc::now(), + ); + let job = Job::new((), row.clone()); + assert_eq!(job.id(), 42); + assert_eq!( + InsertResult { + job, + unique_skipped_as_duplicate: false, + } + .id(), + 42 + ); + assert_eq!( + InsertBatchResult { + job: row, + unique_skipped_as_duplicate: true, + } + .id(), + 42 + ); + } + + #[test] + fn explicit_river_default_overrides_a_job_default() { + let resolved = InsertOpts::resolve( + 7, + InsertOpts::default() + .with_max_attempts(9) + .with_priority(3) + .with_queue("job_queue"), + InsertOpts::default() + .with_max_attempts(MAX_ATTEMPTS_DEFAULT) + .with_priority(PRIORITY_DEFAULT) + .with_queue(QUEUE_DEFAULT), + ); + + assert_eq!(resolved.max_attempts, MAX_ATTEMPTS_DEFAULT); + assert_eq!(resolved.priority, PRIORITY_DEFAULT); + assert_eq!(resolved.queue, QUEUE_DEFAULT); + } +} diff --git a/rust/riverqueue/src/job/attempt_error.rs b/rust/riverqueue/src/job/attempt_error.rs new file mode 100644 index 000000000..e47a0fd94 --- /dev/null +++ b/rust/riverqueue/src/job/attempt_error.rs @@ -0,0 +1,839 @@ +//! Decoding of attempt errors. +//! +//! [`AttemptError`] deserializes the way River Go's `encoding/json` decodes +//! `rivertype.AttemptError`. Persisted attempt errors are decoded leniently +//! instead, like River Go's driver reads: River always writes them in the +//! shape [`AttemptError`] serializes to, but elements written by other tools +//! or edited by hand might not match it, and a job row can't be read or worked +//! unless every one of its attempt errors decodes. + +use std::{borrow::Cow, fmt}; + +use chrono::{DateTime, NaiveDate, Utc}; +use serde::{ + Deserialize, Deserializer, + de::{self, MapAccess, Visitor}, +}; +use serde_json::value::RawValue; + +use super::AttemptError; +use crate::client::saturating_i16; + +impl<'de> Deserialize<'de> for AttemptError { + /// Decodes an attempt error like Go's `encoding/json`: fields match + /// case-insensitively, a repeated field takes its last value, missing, + /// `null`, and unknown fields are accepted, and `at` must be an RFC 3339 + /// timestamp. A field of any other type is an error. + /// + /// Only a JSON deserializer (such as [`serde_json`]'s) can decode an + /// attempt error. + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let raw = Box::::deserialize(deserializer)?; + Self::from_json_strict(raw.get()).map_err(de::Error::custom) + } +} + +impl AttemptError { + /// Decodes an attempt error like Go's `encoding/json`. See + /// [`AttemptError`]'s `Deserialize` implementation. + fn from_json_strict(json: &str) -> Result { + let json = json.trim_matches(is_json_whitespace); + if json == "null" { + return Ok(Self::new(go_zero_time(), 0, "")); + } + if !json.starts_with('{') { + return Err(format!("cannot decode {json} into an attempt error")); + } + let fields: Fields = serde_json::from_str(json).map_err(|error| error.to_string())?; + if let Some(invalid) = fields.invalid { + return Err(invalid); + } + Ok(Self { + at: strict_time(raw_or_null(fields.at.as_deref()))?, + attempt: strict_attempt(raw_or_null(fields.attempt.as_deref()))?, + error: strict_string(raw_or_null(fields.error.as_deref()))?, + trace: strict_string(raw_or_null(fields.trace.as_deref()))?, + }) + } + + /// Decodes one persisted attempt error exactly like River Go's + /// `riverdriver.UnmarshalAttemptError`. + /// + /// Elements in the shape River writes decode as they would with Go's + /// `encoding/json` defaults. Any other valid JSON decodes on a best + /// effort basis: + /// + /// * `at` accepts only what Go's `time.Time` does, RFC 3339 timestamps. + /// Any other value leaves Go's zero time. + /// * `attempt` accepts integers, numbers with an integral value, and + /// strings containing either. Any other value leaves zero. + /// * `error` and `trace` accept strings. Any other non-null value is kept + /// as its compacted JSON text. + /// * An element that's a JSON string instead of an object is used as + /// `error`, and any other element that isn't an object is kept as its + /// JSON text in `error`. + /// + /// Only text that isn't valid JSON is an error. + pub(crate) fn from_json_lenient(json: &str) -> Result { + let json = json.trim_matches(is_json_whitespace); + if json.starts_with('{') { + let fields: Fields = serde_json::from_str(json)?; + return Ok(Self { + at: fields.at.map_or_else(go_zero_time, |raw| { + strict_time(raw.get()).unwrap_or_else(|_| go_zero_time()) + }), + attempt: fields.attempt.map_or(0, |raw| lenient_attempt(raw.get())), + error: fields + .error + .map(|raw| lenient_string(raw.get())) + .unwrap_or_default(), + trace: fields + .trace + .map(|raw| lenient_string(raw.get())) + .unwrap_or_default(), + }); + } + + // Valid JSON, but not an object. `null` leaves every field empty. + let raw: Box = serde_json::from_str(json)?; + Ok(Self::new(go_zero_time(), 0, lenient_string(raw.get()))) + } + + /// Decodes a persisted JSON array of attempt errors like River Go's + /// `riverdriver.UnmarshalAttemptErrors`: each element decodes with + /// [`from_json_lenient`](Self::from_json_lenient), `null` is empty, and + /// anything other than an array is an error. + pub(crate) fn from_json_array_lenient(json: &str) -> Result, serde_json::Error> { + serde_json::from_str::>>>(json)? + .unwrap_or_default() + .iter() + .map(|raw| Self::from_json_lenient(raw.get())) + .collect() + } +} + +/// Go's zero `time.Time`, which Go leaves in an attempt error without a +/// usable `at`. +fn go_zero_time() -> DateTime { + NaiveDate::from_ymd_opt(1, 1, 1) + .and_then(|date| date.and_hms_opt(0, 0, 0)) + .expect("Go's zero time is a valid date") + .and_utc() +} + +/// An attempt error object's fields as raw JSON, matched the way Go's +/// `encoding/json` matches struct fields. +#[derive(Default)] +struct Fields { + at: Option>, + attempt: Option>, + error: Option>, + trace: Option>, + /// Why Go's `encoding/json` would reject a value, including one later + /// replaced by a repeated field, which Go still reports. + invalid: Option, +} + +impl<'de> Deserialize<'de> for Fields { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + struct FieldsVisitor; + + impl<'de> Visitor<'de> for FieldsVisitor { + type Value = Fields; + + fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("an attempt error object") + } + + fn visit_map(self, mut map: A) -> Result + where + A: MapAccess<'de>, + { + let mut fields = Fields::default(); + while let Some(key) = map.next_key::()? { + let value = map.next_value::>()?; + // Go folds case when matching field names. Only ASCII + // letters fold to the letters of these four names. + let (field, valid) = if key.eq_ignore_ascii_case("at") { + (&mut fields.at, strict_time(value.get()).map(drop)) + } else if key.eq_ignore_ascii_case("attempt") { + (&mut fields.attempt, strict_attempt(value.get()).map(drop)) + } else if key.eq_ignore_ascii_case("error") { + (&mut fields.error, strict_string(value.get()).map(drop)) + } else if key.eq_ignore_ascii_case("trace") { + (&mut fields.trace, strict_string(value.get()).map(drop)) + } else { + continue; + }; + *field = Some(value); + if let Err(invalid) = valid { + fields.invalid.get_or_insert(invalid); + } + } + Ok(fields) + } + } + + deserializer.deserialize_map(FieldsVisitor) + } +} + +/// A field's raw JSON, with a missing field read as `null` like Go does. +fn raw_or_null(field: Option<&RawValue>) -> &str { + field.map_or("null", RawValue::get) +} + +fn is_json_whitespace(character: char) -> bool { + matches!(character, ' ' | '\t' | '\n' | '\r') +} + +/// Decodes `at` like Go's `time.Time.UnmarshalJSON`: `null` is Go's zero +/// time, and a string must hold an RFC 3339 timestamp. Like Go 1.26, the +/// string isn't unescaped first. +fn strict_time(raw: &str) -> Result, String> { + if raw == "null" { + return Ok(go_zero_time()); + } + raw.strip_prefix('"') + .and_then(|text| text.strip_suffix('"')) + .and_then(parse_rfc3339) + .ok_or_else(|| format!("attempt error time {raw} isn't an RFC 3339 timestamp")) +} + +/// Decodes `attempt` like Go's `encoding/json` decodes an `int`: an integer +/// literal or `null`. `attempt` is narrower than Go's `int`, so values beyond +/// `i16` saturate like other persisted attempt counts. +fn strict_attempt(raw: &str) -> Result { + if raw == "null" { + return Ok(0); + } + raw.parse::() + .map(saturating_i16) + .map_err(|_| format!("attempt error attempt {raw} isn't an integer")) +} + +/// Decodes `error` or `trace` like Go's `encoding/json` decodes a `string`: +/// a string or `null`. +fn strict_string(raw: &str) -> Result { + if raw == "null" { + return Ok(String::new()); + } + if raw.starts_with('"') { + return serde_json::from_str(raw).map_err(|error| error.to_string()); + } + Err(format!("attempt error value {raw} isn't a string")) +} + +/// The largest magnitude up to which every integer is exactly representable +/// as an `f64`, 2^53. +const MAX_EXACT_FLOAT_INTEGER: f64 = 9_007_199_254_740_992.0; + +/// Decodes `attempt` like Go: integers, and numbers or numeric strings with an +/// integral value no larger in magnitude than 2^53. `attempt` is narrower +/// than Go's `int`, so values beyond `i16` saturate like other persisted +/// attempt counts. Unlike Go, a string in hexadecimal floating point notation +/// (such as `"0x1p4"`) isn't recognized and decodes as zero. +#[allow( + clippy::float_cmp, + reason = "an exact comparison checks for an integral value" +)] +fn lenient_attempt(raw: &str) -> i16 { + let number: Cow<'_, str> = if raw.starts_with('"') { + match serde_json::from_str::(raw) { + Ok(text) => Cow::Owned(text.trim().to_owned()), + Err(_) => return 0, + } + } else if raw.starts_with(|character: char| character == '-' || character.is_ascii_digit()) { + Cow::Borrowed(raw) + } else { + return 0; + }; + + if let Ok(integer) = number.parse::() { + return saturating_i16(integer); + } + match number.parse::() { + Ok(float) if float == float.trunc() && float.abs() <= MAX_EXACT_FLOAT_INTEGER => + { + #[expect( + clippy::cast_possible_truncation, + reason = "the float is integral and within the exact integer range" + )] + saturating_i16(float as i64) + } + _ => 0, + } +} + +/// Decodes `error` and `trace` like Go: a string is used as is, `null` is +/// empty, and any other value is kept as its compacted JSON text. +fn lenient_string(raw: &str) -> String { + if raw == "null" { + return String::new(); + } + if raw.starts_with('"') + && let Ok(text) = serde_json::from_str::(raw) + { + return text; + } + compact_json(raw) +} + +/// Removes insignificant whitespace from valid JSON text without otherwise +/// changing it, like Go's `json.Compact`. +fn compact_json(raw: &str) -> String { + let mut compacted = String::with_capacity(raw.len()); + let mut in_string = false; + let mut escaped = false; + for character in raw.chars() { + if in_string { + compacted.push(character); + if escaped { + escaped = false; + } else if character == '\\' { + escaped = true; + } else if character == '"' { + in_string = false; + } + } else if !is_json_whitespace(character) { + in_string = character == '"'; + compacted.push(character); + } + } + compacted +} + +/// Parses a timestamp exactly as Go's `time.Time.UnmarshalJSON` does, which +/// is with `time.Parse` and Go's RFC 3339 layout: `YYYY-MM-DD`, `T`, a one or +/// two digit hour, `:MM:SS` with valid ranges and no leap second, an optional +/// fraction introduced by `.` or `,` (digits past nanoseconds are ignored), +/// and `Z` or a `±hh:mm` offset of up to 24 hours and 60 minutes. +fn parse_rfc3339(text: &str) -> Option> { + let mut parser = TimeParser(text.as_bytes()); + let year = parser.digits(4)?; + parser.expect(b"-")?; + let month = parser.digits(2)?; + parser.expect(b"-")?; + let day = parser.digits(2)?; + parser.expect(b"T")?; + // Go's `15` hour takes one digit when a second one doesn't follow. + let hour = parser.digits(2).or_else(|| parser.digits(1))?; + parser.expect(b":")?; + let minute = parser.digits(2)?; + parser.expect(b":")?; + let second = parser.digits(2)?; + let nanosecond = parser.fraction(); + let offset_seconds = parser.offset()?; + if !parser.0.is_empty() || hour > 23 || minute > 59 || second > 59 { + return None; + } + + let local = NaiveDate::from_ymd_opt(i32::try_from(year).ok()?, month, day)? + .and_hms_nano_opt(hour, minute, second, nanosecond)? + .and_utc(); + local.checked_sub_signed(chrono::Duration::seconds(offset_seconds)) +} + +/// The unparsed remainder of a timestamp. +struct TimeParser<'a>(&'a [u8]); + +impl TimeParser<'_> { + /// Consumes exactly `count` ASCII digits. + fn digits(&mut self, count: usize) -> Option { + let digits = self.0.get(..count)?; + if !digits.iter().all(u8::is_ascii_digit) { + return None; + } + self.0 = &self.0[count..]; + Some( + digits + .iter() + .fold(0, |value, digit| value * 10 + u32::from(digit - b'0')), + ) + } + + fn expect(&mut self, literal: &[u8]) -> Option<()> { + self.0 = self.0.strip_prefix(literal)?; + Some(()) + } + + /// Consumes an optional fractional second, returning nanoseconds. + fn fraction(&mut self) -> u32 { + let [b'.' | b',', first, ..] = self.0 else { + return 0; + }; + if !first.is_ascii_digit() { + return 0; + } + let digit_count = self.0[1..] + .iter() + .take_while(|byte| byte.is_ascii_digit()) + .count(); + let digits = &self.0[1..=digit_count]; + self.0 = &self.0[1 + digit_count..]; + digits + .iter() + .chain(std::iter::repeat(&b'0')) + .take(9) + .fold(0, |nanoseconds, digit| { + nanoseconds * 10 + u32::from(digit - b'0') + }) + } + + /// Consumes a `Z` or `±hh:mm` UTC offset, returning it in seconds east of + /// UTC. + fn offset(&mut self) -> Option { + let sign = match self.0.first()? { + b'Z' => { + self.0 = &self.0[1..]; + return Some(0); + } + b'+' => 1, + b'-' => -1, + _ => return None, + }; + self.0 = &self.0[1..]; + let hours = self.digits(2)?; + self.expect(b":")?; + let minutes = self.digits(2)?; + if hours > 24 || minutes > 60 { + return None; + } + Some(sign * (i64::from(hours) * 3_600 + i64::from(minutes) * 60)) + } +} + +#[cfg(test)] +mod tests { + use chrono::{DateTime, NaiveDate, Utc}; + + use super::{AttemptError, go_zero_time}; + + fn attempt_at() -> DateTime { + NaiveDate::from_ymd_opt(2024, 1, 2) + .and_then(|date| date.and_hms_micro_opt(3, 4, 5, 123_456)) + .unwrap() + .and_utc() + } + + fn attempt_error(at: DateTime, attempt: i16, error: &str, trace: &str) -> AttemptError { + AttemptError::new(at, attempt, error).with_trace(trace) + } + + fn assert_lenient(cases: [(&str, &str, AttemptError); N]) { + for (name, json, expected) in cases { + assert_eq!( + AttemptError::from_json_lenient(json).unwrap(), + expected, + "{name}" + ); + } + } + + fn whole_second() -> DateTime { + NaiveDate::from_ymd_opt(2024, 1, 2) + .and_then(|date| date.and_hms_opt(3, 4, 5)) + .unwrap() + .and_utc() + } + + #[test] + fn invalid_json_is_an_error() { + assert!(AttemptError::from_json_lenient(r#"{"at":"#).is_err()); + assert!(AttemptError::from_json_array_lenient(r#"[{"at":"#).is_err()); + assert!(serde_json::from_str::(r#"{"at":"#).is_err()); + } + + // The cases of River Go's `TestUnmarshalAttemptError`. + #[test] + fn lenient_like_go() { + let zero = go_zero_time(); + assert_lenient([ + ( + "AtInvalid", + r#"{"at":"not a time","attempt":2,"error":"err"}"#, + attempt_error(zero, 2, "err", ""), + ), + ( + "AtNoOffset", + r#"{"at":"2024-01-02T03:04:05.123456","attempt":2}"#, + attempt_error(zero, 2, "", ""), + ), + ( + "AtNumber", + r#"{"at":1704164645,"attempt":2}"#, + attempt_error(zero, 2, "", ""), + ), + ( + "AtPostgresText", + r#"{"at":"2024-01-02 03:04:05.123456+00","attempt":2}"#, + attempt_error(zero, 2, "", ""), + ), + ( + "AtRFC3339WithOtherInvalidField", + r#"{"at":"2024-01-02T03:04:05.123456Z","attempt":"2"}"#, + attempt_error(attempt_at(), 2, "", ""), + ), + ( + "AtSpaceNoOffset", + r#"{"at":"2024-01-02 03:04:05.123456","attempt":2}"#, + attempt_error(zero, 2, "", ""), + ), + ( + "AttemptFloat", + r#"{"attempt":3.0,"error":"err"}"#, + attempt_error(zero, 3, "err", ""), + ), + ( + "AttemptFractional", + r#"{"attempt":3.5,"error":"err"}"#, + attempt_error(zero, 0, "err", ""), + ), + ( + "AttemptObject", + r#"{"attempt":{},"error":"err"}"#, + attempt_error(zero, 0, "err", ""), + ), + ( + "AttemptString", + r#"{"attempt":" 3 ","error":"err"}"#, + attempt_error(zero, 3, "err", ""), + ), + ( + "AttemptStringInvalid", + r#"{"attempt":"three","error":"err"}"#, + attempt_error(zero, 0, "err", ""), + ), + ( + "ElementArray", + r#"[1, "two"]"#, + attempt_error(zero, 0, r#"[1,"two"]"#, ""), + ), + ("ElementNumber", "123", attempt_error(zero, 0, "123", "")), + ( + "ElementString", + r#""job failed""#, + attempt_error(zero, 0, "job failed", ""), + ), + ( + "ErrorObject", + r#"{"attempt":1,"error":{"message": "boom", "code": 7}}"#, + attempt_error(zero, 1, r#"{"message":"boom","code":7}"#, ""), + ), + ( + "TraceArray", + r#"{"attempt":1,"error":"err","trace":["frame1", "frame2"]}"#, + attempt_error(zero, 1, "err", r#"["frame1","frame2"]"#), + ), + ( + "TraceNullWithInvalidField", + r#"{"attempt":"x","error":null,"trace":null}"#, + attempt_error(zero, 0, "", ""), + ), + ( + "StrictShapeUnchanged", + r#"{"attempt":2,"error":null}"#, + attempt_error(zero, 2, "", ""), + ), + ]); + } + + // Behavior of Go's `encoding/json` that the Go test cases don't reach. + #[test] + fn lenient_fields_edges_like_go() { + let zero = go_zero_time(); + let whole = whole_second(); + assert_lenient([ + ( + "CaseInsensitiveLastWins", + r#"{"error":"first","ERROR":"second","Attempt":2}"#, + attempt_error(zero, 2, "second", ""), + ), + ( + "EscapedKey", + concat!(r#"{""#, "\\", r#"u0061t":"2024-01-02T03:04:05Z"}"#), + attempt_error(whole, 0, "", ""), + ), + ("ElementNull", "null", attempt_error(zero, 0, "", "")), + ("ElementTrue", "true", attempt_error(zero, 0, "true", "")), + ( + "ErrorKeepsNumberAndEscapeText", + r#"{"error":{"n": 1.50, "s": "a \"b\"\n c"}}"#, + attempt_error(zero, 0, r#"{"n":1.50,"s":"a \"b\"\n c"}"#, ""), + ), + ( + "AttemptExponentString", + r#"{"attempt":"1e1"}"#, + attempt_error(zero, 10, "", ""), + ), + ( + "AttemptSignedString", + r#"{"attempt":"+4"}"#, + attempt_error(zero, 4, "", ""), + ), + ( + "AttemptBeyondExactFloat", + r#"{"attempt":1e20}"#, + attempt_error(zero, 0, "", ""), + ), + ( + "AttemptSaturates", + r#"{"attempt":40000}"#, + attempt_error(zero, i16::MAX, "", ""), + ), + ( + "AttemptBool", + r#"{"attempt":true}"#, + attempt_error(zero, 0, "", ""), + ), + ]); + } + + // Behavior of Go's `time.Time.UnmarshalJSON` that the Go test cases don't + // reach: it accepts what `time.Parse` does with Go's RFC 3339 layout, and + // anything else leaves zero. + #[test] + #[allow(clippy::too_many_lines)] + fn lenient_time_edges_like_go() { + let zero = go_zero_time(); + let whole = whole_second(); + assert_lenient([ + ( + "AtCommaFraction", + r#"{"at":"2024-01-02T03:04:05,123456Z"}"#, + attempt_error(attempt_at(), 0, "", ""), + ), + ( + "AtFractionBeyondNanoseconds", + r#"{"at":"2024-01-02T03:04:05.1234560009Z"}"#, + attempt_error(attempt_at(), 0, "", ""), + ), + ( + "AtOffset", + r#"{"at":"2024-01-02T00:04:05-03:00"}"#, + attempt_error(whole, 0, "", ""), + ), + ( + "AtOffsetLargestGoAccepts", + r#"{"at":"2024-01-03T04:04:05+24:60"}"#, + attempt_error(whole, 0, "", ""), + ), + ( + "AtOneDigitHour", + r#"{"at":"2024-01-02T3:04:05Z"}"#, + attempt_error(whole, 0, "", ""), + ), + ( + "AtYearZero", + r#"{"at":"0000-01-01T00:00:00Z"}"#, + attempt_error( + NaiveDate::from_ymd_opt(0, 1, 1) + .and_then(|date| date.and_hms_opt(0, 0, 0)) + .unwrap() + .and_utc(), + 0, + "", + "", + ), + ), + ( + "AtEmptyFraction", + r#"{"at":"2024-01-02T03:04:05.Z"}"#, + attempt_error(zero, 0, "", ""), + ), + ( + "AtEscaped", + concat!(r#"{"at":"2024-01-02T03:04:05"#, "\\", r#"u005a"}"#), + attempt_error(zero, 0, "", ""), + ), + ( + "AtHourOutOfRange", + r#"{"at":"2024-01-02T24:04:05Z"}"#, + attempt_error(zero, 0, "", ""), + ), + ( + "AtInvalidDay", + r#"{"at":"2023-02-29T03:04:05Z"}"#, + attempt_error(zero, 0, "", ""), + ), + ( + "AtLeapSecond", + r#"{"at":"2024-01-02T03:04:60Z"}"#, + attempt_error(zero, 0, "", ""), + ), + ( + "AtLowercaseSeparator", + r#"{"at":"2024-01-02t03:04:05Z"}"#, + attempt_error(zero, 0, "", ""), + ), + ( + "AtOffsetHourOutOfRange", + r#"{"at":"2024-01-02T03:04:05+25:00"}"#, + attempt_error(zero, 0, "", ""), + ), + ( + "AtOffsetHoursOnly", + r#"{"at":"2024-01-02T08:04:05+05"}"#, + attempt_error(zero, 0, "", ""), + ), + ( + "AtOffsetMinuteOutOfRange", + r#"{"at":"2024-01-02T03:04:05+05:61"}"#, + attempt_error(zero, 0, "", ""), + ), + ( + "AtOffsetWithoutColon", + r#"{"at":"2024-01-02T05:34:05+0230"}"#, + attempt_error(zero, 0, "", ""), + ), + ( + "AtSurroundingSpace", + r#"{"at":" 2024-01-02T03:04:05Z "}"#, + attempt_error(zero, 0, "", ""), + ), + ( + "AtTrailingText", + r#"{"at":"2024-01-02T03:04:05Zjunk"}"#, + attempt_error(zero, 0, "", ""), + ), + ]); + } + + // The cases of River Go's `TestUnmarshalAttemptErrors`. + #[test] + fn lenient_array_like_go() { + assert_eq!( + AttemptError::from_json_array_lenient("[]").unwrap(), + Vec::new() + ); + assert_eq!( + AttemptError::from_json_array_lenient("null").unwrap(), + Vec::new() + ); + assert!(AttemptError::from_json_array_lenient(r#"{"error":"not an array"}"#).is_err()); + // One unexpected element doesn't prevent decoding the others. + assert_eq!( + AttemptError::from_json_array_lenient( + r#"[{"at":"2024-01-02T03:04:05.123456Z","attempt":1,"error":"err1","trace":""},"err2"]"# + ) + .unwrap(), + vec![ + attempt_error(attempt_at(), 1, "err1", ""), + attempt_error(go_zero_time(), 0, "err2", ""), + ] + ); + assert_eq!( + AttemptError::from_json_array_lenient( + r#"[{"at":"invalid","attempt":"2","error":"err"},{"error":"next"}]"# + ) + .unwrap(), + vec![ + attempt_error(go_zero_time(), 2, "err", ""), + attempt_error(go_zero_time(), 0, "next", ""), + ] + ); + } + + // Deserializing an attempt error is strict like Go's `encoding/json`. + #[test] + fn deserializes_like_go_json() { + let zero = go_zero_time(); + for (name, json, expected) in [ + ( + "Full", + r#"{"at":"2024-01-02T03:04:05.123456Z","attempt":3,"error":"err","trace":"t","extra":[1]}"#, + attempt_error(attempt_at(), 3, "err", "t"), + ), + ( + "CaseInsensitiveLastWins", + r#"{"error":"first","ERROR":"second","Attempt":2}"#, + attempt_error(zero, 2, "second", ""), + ), + ( + "MissingAndNull", + r#"{"at":null,"attempt":null,"error":"err","trace":null}"#, + attempt_error(zero, 0, "err", ""), + ), + ("Null", "null", attempt_error(zero, 0, "", "")), + ( + "AttemptSaturates", + r#"{"attempt":-40000}"#, + attempt_error(zero, i16::MIN, "", ""), + ), + ] { + assert_eq!( + serde_json::from_str::(json).unwrap(), + expected, + "{name}" + ); + } + + for (name, json) in [ + ("AtNotRFC3339", r#"{"at":"2024-01-02 03:04:05+00"}"#), + ("AtNumber", r#"{"at":1704164645}"#), + ("AttemptFloat", r#"{"attempt":3.0}"#), + ("AttemptString", r#"{"attempt":"3"}"#), + ("AttemptOverflow", r#"{"attempt":9223372036854775808}"#), + ("ErrorObject", r#"{"error":{"message":"boom"}}"#), + ("TraceArray", r#"{"trace":["frame"]}"#), + ("RepeatedInvalidField", r#"{"error":1,"error":"err"}"#), + ("ElementString", r#""job failed""#), + ("ElementArray", "[]"), + ] { + assert!( + serde_json::from_str::(json).is_err(), + "{name}" + ); + } + } + + #[test] + fn serializes_at_like_go() { + let at = NaiveDate::from_ymd_opt(2024, 1, 2) + .unwrap() + .and_hms_opt(3, 4, 5) + .unwrap() + .and_utc() + + chrono::Duration::nanoseconds(678_900_000); + let encoded = serde_json::to_string(&attempt_error(at, 1, "", "")).unwrap(); + assert_eq!( + encoded, + r#"{"at":"2024-01-02T03:04:05.6789Z","attempt":1,"error":"","trace":""}"# + ); + } + + #[test] + fn round_trip() { + let attempt_error = attempt_error(attempt_at(), 3, "job failed", "frame one"); + let encoded = serde_json::to_string(&attempt_error).unwrap(); + assert_eq!( + serde_json::from_str::(&encoded).unwrap(), + serde_json::json!({ + "at": "2024-01-02T03:04:05.123456Z", + "attempt": 3, + "error": "job failed", + "trace": "frame one", + }) + ); + assert_eq!( + serde_json::from_str::(&encoded).unwrap(), + attempt_error + ); + assert_eq!( + AttemptError::from_json_lenient(&encoded).unwrap(), + attempt_error + ); + // Values decoded from `serde_json::Value` work the same way. + assert_eq!( + serde_json::from_value::(serde_json::to_value(&attempt_error).unwrap()) + .unwrap(), + attempt_error + ); + } +} diff --git a/rust/riverqueue/src/job/metadata.rs b/rust/riverqueue/src/job/metadata.rs new file mode 100644 index 000000000..81f2c90bb --- /dev/null +++ b/rust/riverqueue/src/job/metadata.rs @@ -0,0 +1,295 @@ +//! Exact persisted job metadata. + +use std::{collections::BTreeMap, fmt, str::FromStr}; + +use serde::{Deserialize, Deserializer, Serialize, Serializer, de::DeserializeOwned}; +use serde_json::{Map, Value, value::RawValue}; + +/// A JSON object whose original value tokens remain intact. +/// +/// PostgreSQL may store numbers beyond the range of `serde_json::Value`, so +/// decoding the entire object into a map can make an otherwise valid job +/// unreadable. This type keeps the persisted text and decodes only requested +/// fields. Database-side metadata merges preserve unrelated value tokens. +#[derive(Clone)] +pub struct JobMetadata(Box); + +impl JobMetadata { + /// Accepts `raw` only when it holds a JSON object. Every fallible + /// conversion goes through here so they reject the same inputs with the + /// same error. + fn from_raw_object(raw: Box) -> Result { + if raw.get().trim_start().starts_with('{') { + Ok(Self(raw)) + } else { + Err(::custom( + "job metadata must be a JSON object", + )) + } + } + + /// Returns the exact stored JSON object. + #[must_use] + pub fn as_raw(&self) -> &RawValue { + &self.0 + } + + /// Returns whether the object contains a field, without decoding values. + #[must_use] + pub fn contains_key(&self, key: &str) -> bool { + self.get_raw(key).is_some() + } + + /// Decodes one field into a caller-selected type. + /// + /// # Errors + /// + /// Returns an error if the selected field cannot deserialize as `T`. + pub fn get(&self, key: &str) -> Result, serde_json::Error> { + self.get_raw(key) + .map(|raw| serde_json::from_str(raw.get())) + .transpose() + } + + /// Borrows one field's original JSON value, resolving duplicate names to + /// their last occurrence. + #[must_use] + pub fn get_raw(&self, key: &str) -> Option<&RawValue> { + // Borrowed RawValue skips the number parser, including numbers much + // larger than f64. The map owns only field names, not value bytes. + let fields: BTreeMap = serde_json::from_str(self.0.get()).ok()?; + fields.get(key).copied() + } + + /// Returns whether this object has no fields. + /// + /// # Panics + /// + /// Panics only if the internally validated JSON object becomes malformed. + #[must_use] + pub fn is_empty(&self) -> bool { + let fields: BTreeMap = + serde_json::from_str(self.0.get()).expect("validated metadata object"); + fields.is_empty() + } + + /// Sets a field while preserving every other field's original value text + /// and relative order. + /// + /// # Errors + /// + /// Returns an error if `value` cannot be serialized as JSON. + /// + /// # Panics + /// + /// Panics only if the internally validated JSON object becomes malformed. + pub fn insert(&mut self, key: &str, value: impl Serialize) -> Result<(), serde_json::Error> { + let encoded_key = serde_json::to_string(key)?; + let encoded_value = serde_json::to_string(&value)?; + let members = crate::unique::object_members(self.0.get()).expect("validated object"); + let last = members.iter().rposition(|member| member.key == key); + let mut result = String::from("{"); + for (index, member) in members.iter().enumerate() { + if index > 0 { + result.push(','); + } + result.push_str(member.raw_key); + result.push(':'); + result.push_str(if Some(index) == last { + &encoded_value + } else { + member.value + }); + } + if last.is_none() { + if !members.is_empty() { + result.push(','); + } + result.push_str(&encoded_key); + result.push(':'); + result.push_str(&encoded_value); + } + result.push('}'); + self.0 = RawValue::from_string(result).expect("valid object update"); + Ok(()) + } + + /// Consumes the metadata and returns its exact JSON object. + #[must_use] + pub fn into_raw(self) -> Box { + self.0 + } + + /// Removes all occurrences of a field. Returns whether it was present. + /// + /// # Panics + /// + /// Panics only if the internally validated JSON object becomes malformed. + pub fn remove(&mut self, key: &str) -> bool { + let members = crate::unique::object_members(self.0.get()).expect("validated object"); + let mut result = String::from("{"); + let mut removed = false; + for member in &members { + if member.key == key { + removed = true; + continue; + } + if result.len() > 1 { + result.push(','); + } + result.push_str(member.raw_key); + result.push(':'); + result.push_str(member.value); + } + if removed { + result.push('}'); + self.0 = RawValue::from_string(result).expect("valid object removal"); + } + removed + } + + /// Decodes the complete object when every number is representable by + /// `serde_json::Value`. + /// + /// # Errors + /// + /// Returns an error for out-of-range numbers such as `1e400`. + pub fn to_map(&self) -> Result, serde_json::Error> { + serde_json::from_str(self.0.get()) + } +} + +impl Default for JobMetadata { + fn default() -> Self { + Self(RawValue::from_string("{}".to_owned()).expect("valid empty object")) + } +} + +impl fmt::Debug for JobMetadata { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_tuple("JobMetadata") + .field(&self.0.get()) + .finish() + } +} + +impl fmt::Display for JobMetadata { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(self.0.get()) + } +} + +impl From> for JobMetadata { + fn from(map: Map) -> Self { + Self( + RawValue::from_string(serde_json::to_string(&map).expect("JSON map serializes")) + .expect("serialized JSON map is valid"), + ) + } +} + +impl FromStr for JobMetadata { + type Err = serde_json::Error; + + fn from_str(text: &str) -> Result { + Self::from_raw_object(serde_json::from_str(text)?) + } +} + +impl PartialEq for JobMetadata { + fn eq(&self, other: &Self) -> bool { + self.0.get() == other.0.get() + } +} + +impl Eq for JobMetadata {} + +impl Serialize for JobMetadata { + fn serialize(&self, serializer: S) -> Result { + self.0.serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for JobMetadata { + fn deserialize>(deserializer: D) -> Result { + Self::from_raw_object(Box::::deserialize(deserializer)?) + .map_err(serde::de::Error::custom) + } +} + +impl TryFrom> for JobMetadata { + type Error = serde_json::Error; + + fn try_from(raw: Box) -> Result { + Self::from_raw_object(raw) + } +} + +impl TryFrom for JobMetadata { + type Error = serde_json::Error; + + /// Accepts a JSON object value. Because `Value` has already parsed its + /// numbers, use [`FromStr`] or `TryFrom>` to keep number + /// text that `Value` can't represent exactly. + fn try_from(value: Value) -> Result { + match value { + Value::Object(map) => Ok(map.into()), + _ => Err(::custom( + "job metadata must be a JSON object", + )), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn keeps_large_number_tokens_and_last_duplicate() { + let metadata: JobMetadata = r#"{"n":1,"large":1e400,"n":2}"#.parse().unwrap(); + assert_eq!(metadata.as_raw().get(), r#"{"n":1,"large":1e400,"n":2}"#); + assert_eq!(metadata.get_raw("large").unwrap().get(), "1e400"); + assert_eq!(metadata.get::("n").unwrap(), Some(2)); + assert!(metadata.to_map().is_err()); + } + + #[test] + fn rejects_non_objects_through_every_conversion() { + for text in ["[1]", "null", " 1"] { + let parsed = text.parse::().unwrap_err(); + let deserialized = serde_json::from_str::(text).unwrap_err(); + let raw = + JobMetadata::try_from(RawValue::from_string(text.to_owned()).unwrap()).unwrap_err(); + let value = + JobMetadata::try_from(serde_json::from_str::(text).unwrap()).unwrap_err(); + for error in [parsed, deserialized, raw, value] { + assert!( + error + .to_string() + .contains("job metadata must be a JSON object"), + "{error}" + ); + } + } + } + + #[test] + fn converts_json_object_values() { + let metadata = JobMetadata::try_from(serde_json::json!({"a": 1})).unwrap(); + assert_eq!(metadata.as_raw().get(), r#"{"a":1}"#); + } + + #[test] + fn changes_one_field_without_round_tripping_other_numbers() { + let mut metadata: JobMetadata = r#"{"z":1e400,"n":0.1000000000000000055511151231257827}"# + .parse() + .unwrap(); + metadata.insert("flag", true).unwrap(); + metadata.insert("n", 2).unwrap(); + assert_eq!(metadata.as_raw().get(), r#"{"z":1e400,"n":2,"flag":true}"#); + assert!(metadata.remove("flag")); + assert_eq!(metadata.as_raw().get(), r#"{"z":1e400,"n":2}"#); + } +} diff --git a/rust/riverqueue/src/lib.rs b/rust/riverqueue/src/lib.rs new file mode 100644 index 000000000..e84ec151c --- /dev/null +++ b/rust/riverqueue/src/lib.rs @@ -0,0 +1,153 @@ +// The README's examples use both backends, so it's the crate documentation +// (and its examples run as doctests) when both are enabled, as on docs.rs. +#![cfg_attr( + all(feature = "postgres", feature = "sqlite"), + doc = include_str!("../README.md") +)] +#![cfg_attr( + not(all(feature = "postgres", feature = "sqlite")), + doc = "The Rust and Tokio client for [River](https://riverqueue.com). The full crate \ + documentation, from the README, builds with both the `postgres` and `sqlite` \ + features, as on [docs.rs](https://docs.rs/riverqueue)." +)] +#![cfg_attr(docsrs, feature(doc_cfg))] +#![forbid(unsafe_code)] +#![warn(missing_docs)] + +#[cfg(not(any(feature = "postgres", feature = "sqlite")))] +compile_error!("riverqueue requires at least one database feature: `postgres` or `sqlite`"); + +extern crate self as riverqueue; + +#[doc(hidden)] +pub mod __private; +mod client; +pub mod database; +pub mod encoding; +pub mod error; +pub mod event; +pub mod extension; +pub mod guide; +pub mod job; +mod maintenance; +pub mod periodic; +mod pilot; +pub mod protocol; +pub mod query; +pub mod queue; +mod storage; +mod unique; +pub mod worker; + +/// The Chrono version of the timestamps in River's API, such as +/// [`JobRow::scheduled_at`] and [`InsertOpts::with_scheduled_at`]. +pub use chrono; +pub use client::{ + Client, ClientBuilder, InsertBatchRequest, InsertManyItem, InsertManyRequest, InsertRequest, + JobCancelRequest, JobCompleteRequest, JobCompleteTxRequest, JobDeleteManyRequest, + JobDeleteRequest, JobGetRequest, JobListRequest, JobRetryRequest, JobUpdateRequest, Jobs, + LocalQueues, MaintenanceConfig, QueueConfig, QueueGetRequest, QueueListRequest, + QueuePauseRequest, QueueResumeRequest, QueueUpdateRequest, Queues, ResignRequest, Retention, + RunHandle, Stopper, +}; +#[cfg(feature = "postgres")] +pub(crate) use database::SchemaName; +pub use error::{ + BoxError, ConfigurationError, Error, ExtensionPhase, JobValidationError, Record, RuntimeError, +}; +pub use event::{ + Event, EventKind, EventReceiver, EventRecvError, JobEvent, JobEventKind, JobStatistics, + QueueEvent, QueueEventKind, SubscribeConfig, +}; +pub use extension::{ + DefaultRetryPolicy, ErrorHandler, ErrorHandlerDecision, Extensions, Hook, InsertContext, + InsertMiddleware, InsertNext, InsertedJob, Metric, MetricName, PanicError, Plugin, RetryPolicy, + WorkCancelled, WorkError, WorkMiddleware, WorkNext, WorkResult, +}; +pub use job::{ + AttemptError, InsertBatch, InsertBatchResult, InsertOpts, InsertParams, InsertResult, Job, + JobArgs, JobMetadata, JobRow, JobState, JobStateParseError, ScheduleOverride, UniqueOpts, +}; +pub use periodic::{ + CronSchedule, CronScheduleParseError, CronTimeZone, IntervalSchedule, NeverSchedule, + PeriodicJob, PeriodicJobHandle, PeriodicJobOpts, PeriodicJobs, PeriodicSchedule, +}; +#[allow(unused_imports, reason = "backend-specific modules use a subset")] +pub(crate) use protocol::{ + NOTIFICATION_TOPIC_CONTROL, NOTIFICATION_TOPIC_INSERT, NOTIFICATION_TOPIC_LEADERSHIP, +}; +pub use query::{ + JobDeleteManyParams, JobListCursor, JobListCursorError, JobListOrderBy, JobListParams, + JobListResult, JobUpdateParams, SortDirection, +}; +pub use queue::{Queue, QueueListParams, QueueSelector, QueueUpdateParams}; +pub use riverqueue_macros::JobArgs; +/// Database migrations, re-exported from `riverqueue-migrate`. +/// +/// Apply migrations before starting clients, for example with +/// [`migrate::PostgresMigrator`]. +pub use riverqueue_migrate as migrate; +/// The `serde_json` version of the JSON values in River's API, such as job +/// metadata and outputs. +pub use serde_json; +/// The SQLx version River's pools and transactions come from. +/// +/// River accepts SQLx pools and transactions directly, so applications must +/// use the same SQLx major version. Depend on SQLx through this re-export, or +/// pin the same version, to avoid mismatched `PgPool`/`SqlitePool` types. River +/// doesn't choose a TLS implementation; enable one of SQLx's TLS features, such +/// as `tls-rustls` or `tls-native-tls`, in your own dependency on SQLx if your +/// database connections use TLS. +pub use sqlx; +/// The `tokio-util` version of [`WorkContext::cancellation_token`]'s +/// [`CancellationToken`](tokio_util::sync::CancellationToken). +pub use tokio_util; +pub use worker::{WorkContext, WorkOutcome, Worker, WorkerRegistry, WorkerTimeout}; + +/// Default maximum number of attempts for a job. +pub const MAX_ATTEMPTS_DEFAULT: i16 = 25; + +/// Default minimum delay between queue fetches. +pub const FETCH_COOLDOWN_DEFAULT: std::time::Duration = std::time::Duration::from_millis(100); + +/// Minimum supported queue fetch cooldown. +pub const FETCH_COOLDOWN_MIN: std::time::Duration = std::time::Duration::from_millis(1); + +/// Default polling interval used as notification-loss recovery. +pub const FETCH_POLL_INTERVAL_DEFAULT: std::time::Duration = std::time::Duration::from_secs(1); + +/// Minimum supported queue polling interval. +pub const FETCH_POLL_INTERVAL_MIN: std::time::Duration = std::time::Duration::from_millis(1); + +/// Default delay before a cancelled worker is considered stuck. +pub const JOB_STUCK_THRESHOLD_DEFAULT: std::time::Duration = std::time::Duration::from_secs(10); + +/// Default per-job execution timeout. +pub const JOB_TIMEOUT_DEFAULT: std::time::Duration = std::time::Duration::from_mins(1); + +/// Maximum worker concurrency allowed for one queue. +pub const QUEUE_NUM_WORKERS_MAX: usize = 10_000; + +/// Default job priority, where one is highest and four is lowest. +pub const PRIORITY_DEFAULT: i16 = 1; + +/// Default queue name. +pub const QUEUE_DEFAULT: &str = "default"; + +/// Reserved metadata key containing recorded job output. +pub const METADATA_KEY_OUTPUT: &str = "output"; + +/// Reserved metadata key containing a periodic job identifier. +pub const METADATA_KEY_PERIODIC_JOB_ID: &str = "river:periodic_job_id"; + +/// Reserved metadata key containing resumable cursors. +pub const METADATA_KEY_RESUMABLE_CURSOR: &str = "river:resumable_cursor"; + +/// Reserved metadata key containing the completed resumable step. +pub const METADATA_KEY_RESUMABLE_STEP: &str = "river:resumable_step"; + +/// Reserved metadata key counting rescues. +pub const METADATA_KEY_RESCUE_COUNT: &str = "river:rescue_count"; + +/// Reserved metadata key used to distinguish unique upserts. +pub const METADATA_KEY_UNIQUE_NONCE: &str = "river:unique_nonce"; diff --git a/rust/riverqueue/src/maintenance/cleaner.rs b/rust/riverqueue/src/maintenance/cleaner.rs new file mode 100644 index 000000000..4a0514ddc --- /dev/null +++ b/rust/riverqueue/src/maintenance/cleaner.rs @@ -0,0 +1,342 @@ +//! Deletes expired finalized jobs, idle queues, and (on SQLite) delivered +//! notifications, ports of Go's `JobCleaner`, `QueueCleaner`, and +//! `SQLiteNotificationCleaner`. + +use std::time::Duration; + +use chrono::{DateTime, Utc}; +#[cfg(feature = "postgres")] +use sqlx::AssertSqlSafe; + +use crate::database::DatabasePool; +#[cfg(feature = "sqlite")] +use crate::database::sqlite; +use crate::{__private::FinalizedJobDeleteParams, Error}; + +use super::{ + MaintenanceError, TIMEOUT_DEFAULT, batch_backoff, batch_size, maintainer::ServiceContext, + record_batch, +}; + +/// Interval of the SQLite notification cleaner. +#[cfg(feature = "sqlite")] +pub(super) const NOTIFICATION_CLEANER_INTERVAL: Duration = Duration::from_mins(1); + +/// Age after which SQLite notification outbox rows are deleted. +#[cfg(feature = "sqlite")] +const NOTIFICATION_RETENTION: Duration = Duration::from_mins(5); + +fn horizon( + now: DateTime, + retention: Option, +) -> Result>, Error> { + retention + .map(|retention| { + chrono::Duration::from_std(retention) + .map(|retention| now - retention) + .map_err(|error| Error::configuration_context("maintenance", error.to_string())) + }) + .transpose() +} + +/// Deletes cancelled, completed, and discarded jobs past their retention in +/// batches. Queues named by an extension's `job_cleaner_queue_exclusions` +/// are skipped; the exclusion list is read on every pass. +pub(super) async fn clean_jobs(context: &ServiceContext) -> Result<(), MaintenanceError> { + let maintenance = &context.inner.maintenance; + // Like Go, skip the query entirely when every retention is indefinite. + if maintenance.cancelled_job_retention.is_none() + && maintenance.completed_job_retention.is_none() + && maintenance.discarded_job_retention.is_none() + { + return Ok(()); + } + let queues_excluded = context.inner.pilot.job_cleaner_queue_exclusions(); + loop { + let now = Utc::now(); + let limit = batch_size(&context.breakers.job_cleaner); + let mut params = FinalizedJobDeleteParams::new(limit); + params.cancelled_before = horizon(now, maintenance.cancelled_job_retention)?; + params.completed_before = horizon(now, maintenance.completed_job_retention)?; + params.discarded_before = horizon(now, maintenance.discarded_job_retention)?; + params.queues_excluded.clone_from(&queues_excluded); + let result = clean_jobs_batch(context, ¶ms).await; + record_batch(&context.breakers.job_cleaner, &result); + if i64::try_from(result?).unwrap_or(i64::MAX) < limit { + return Ok(()); + } + batch_backoff(&context.cancel).await?; + } +} + +async fn clean_jobs_batch( + context: &ServiceContext, + params: &FinalizedJobDeleteParams, +) -> Result { + let timeout = context.inner.maintenance.job_cleaner_timeout; + match context.inner.database.pool() { + #[cfg(feature = "sqlite")] + DatabasePool::Sqlite(pool) => { + let operation = async { + let mut transaction = crate::database::begin_sqlite_write(pool).await?; + let count = sqlite_delete_finalized_jobs(&mut transaction, params).await?; + transaction.commit().await?; + Ok::<_, MaintenanceError>(count) + }; + super::sqlite_cancellable(&context.cancel, timeout, operation).await + } + #[cfg(feature = "postgres")] + DatabasePool::Postgres(pool) => { + use super::postgres::{MaintenanceTransaction, cancellable}; + + let mut transaction = + MaintenanceTransaction::begin(pool, &context.cancel, timeout).await?; + let backend_pid = transaction.backend_pid; + let count = cancellable( + pool, + backend_pid, + &context.cancel, + timeout, + postgres_delete_finalized_jobs( + &mut transaction.transaction, + &context.inner.schema, + params, + ), + ) + .await?; + transaction.commit(pool, &context.cancel).await?; + Ok(count) + } + } +} + +/// Runs the job cleaner's deletion on PostgreSQL. +#[cfg(feature = "postgres")] +pub(crate) async fn postgres_delete_finalized_jobs( + connection: &mut sqlx::PgConnection, + schema: &crate::database::SchemaName, + params: &FinalizedJobDeleteParams, +) -> Result { + if params.limit <= 0 + || (params.cancelled_before.is_none() + && params.completed_before.is_none() + && params.discarded_before.is_none()) + { + return Ok(0); + } + let table = schema.qualify("river_job"); + let now = Utc::now(); + let result = sqlx::query(AssertSqlSafe(format!( + "DELETE FROM {table} WHERE id IN (\ + SELECT id FROM {table} WHERE (\ + (state = 'cancelled' AND $1 AND finalized_at < $2) OR \ + (state = 'completed' AND $3 AND finalized_at < $4) OR \ + (state = 'discarded' AND $5 AND finalized_at < $6)\ + ) AND NOT (queue = ANY($7::text[])) \ + AND ($8::text[] IS NULL OR queue = ANY($8::text[])) \ + ORDER BY id LIMIT $9\ + )" + ))) + .bind(params.cancelled_before.is_some()) + .bind(params.cancelled_before.unwrap_or(now)) + .bind(params.completed_before.is_some()) + .bind(params.completed_before.unwrap_or(now)) + .bind(params.discarded_before.is_some()) + .bind(params.discarded_before.unwrap_or(now)) + .bind(¶ms.queues_excluded) + .bind(params.queues_included.as_deref()) + .bind(params.limit) + .execute(connection) + .await?; + Ok(result.rows_affected()) +} + +/// Runs the job cleaner's deletion on SQLite. +#[cfg(feature = "sqlite")] +pub(crate) async fn sqlite_delete_finalized_jobs( + connection: &mut sqlx::SqliteConnection, + params: &FinalizedJobDeleteParams, +) -> Result { + let queues_excluded = params + .queues_excluded + .iter() + .map(String::as_str) + .collect::>(); + let queues_included = params + .queues_included + .as_ref() + .map(|queues| queues.iter().map(String::as_str).collect::>()); + sqlite::cleanup_jobs( + connection, + &sqlite::CleanupJobs { + cancelled_before: params.cancelled_before, + completed_before: params.completed_before, + discarded_before: params.discarded_before, + limit: i32::try_from(params.limit).unwrap_or(i32::MAX), + queues_excluded: &queues_excluded, + queues_included: queues_included.as_deref(), + }, + ) + .await +} + +/// Deletes queue records that no client has touched within the retention. +/// Active producers refresh `updated_at`, so their queues survive. +pub(super) async fn clean_queues(context: &ServiceContext) -> Result<(), MaintenanceError> { + loop { + let updated_before = Utc::now() + - chrono::Duration::from_std(context.inner.maintenance.queue_retention) + .map_err(|error| Error::configuration_context("maintenance", error.to_string()))?; + let limit = batch_size(&context.breakers.queue_cleaner); + let result = clean_queues_batch(context, updated_before, limit).await; + record_batch(&context.breakers.queue_cleaner, &result); + if i64::try_from(result?).unwrap_or(i64::MAX) < limit { + return Ok(()); + } + batch_backoff(&context.cancel).await?; + } +} + +async fn clean_queues_batch( + context: &ServiceContext, + updated_before: DateTime, + limit: i64, +) -> Result { + match context.inner.database.pool() { + #[cfg(feature = "sqlite")] + DatabasePool::Sqlite(pool) => { + let operation = async { + let mut transaction = crate::database::begin_sqlite_write(pool).await?; + let deleted = sqlite::queue_delete_expired( + &mut transaction, + updated_before, + i32::try_from(limit).unwrap_or(i32::MAX), + ) + .await?; + transaction.commit().await?; + Ok::<_, MaintenanceError>(deleted.len()) + }; + super::sqlite_cancellable(&context.cancel, TIMEOUT_DEFAULT, operation).await + } + #[cfg(feature = "postgres")] + DatabasePool::Postgres(pool) => { + use super::postgres::{MaintenanceTransaction, cancellable}; + + let table = context.inner.schema.qualify("river_queue"); + let mut transaction = + MaintenanceTransaction::begin(pool, &context.cancel, TIMEOUT_DEFAULT).await?; + let backend_pid = transaction.backend_pid; + let deleted = cancellable( + pool, + backend_pid, + &context.cancel, + TIMEOUT_DEFAULT, + sqlx::query_scalar::<_, String>(AssertSqlSafe(format!( + "DELETE FROM {table} WHERE name IN (\ + SELECT name FROM {table} WHERE updated_at < $1 ORDER BY name LIMIT $2\ + ) RETURNING name" + ))) + .bind(updated_before) + .bind(limit) + .fetch_all(&mut *transaction.transaction), + ) + .await?; + transaction.commit(pool, &context.cancel).await?; + Ok(deleted.len()) + } + } +} + +/// Deletes SQLite notification outbox rows old enough that every poller has +/// consumed them, oldest first, in batches like River Go's +/// `SQLiteNotificationCleaner`. +/// +/// Each batch commits on its own, and a pause between batches yields SQLite's +/// write lock to job inserts and updates. The horizon is fixed for the pass +/// so rows expiring meanwhile don't extend it. Repeated timeouts switch to +/// the reduced batch size. +#[cfg(feature = "sqlite")] +pub(super) async fn clean_notifications( + context: &ServiceContext, + pool: &sqlx::SqlitePool, +) -> Result { + let retention = chrono::Duration::from_std(NOTIFICATION_RETENTION) + .map_err(|error| Error::configuration_context("maintenance", error.to_string()))?; + let created_before = Utc::now() - retention; + let mut deleted = 0; + loop { + let limit = batch_size(&context.breakers.notification_cleaner); + let operation = async { + let mut connection = pool.acquire().await?; + Ok::<_, MaintenanceError>( + sqlite::notification_cleanup(&mut connection, created_before, limit).await?, + ) + }; + let result = super::sqlite_cancellable(&context.cancel, TIMEOUT_DEFAULT, operation).await; + record_batch(&context.breakers.notification_cleaner, &result); + let count = result?; + deleted += count; + if count < u64::try_from(limit).unwrap_or(u64::MAX) { + return Ok(deleted); + } + batch_backoff(&context.cancel).await?; + } +} + +#[cfg(all(test, feature = "sqlite"))] +mod sqlite_tests { + use std::sync::Arc; + + use chrono::Utc; + use riverqueue_migrate::SqliteMigrator; + use sqlx::sqlite::SqlitePoolOptions; + use tokio_util::sync::CancellationToken; + + use super::{super::BatchSizes, super::Breakers, clean_notifications}; + use crate::{Client, database::sqlite::sqlite_time, maintenance::maintainer::ServiceContext}; + + #[tokio::test] + async fn notification_cleaner_deletes_expired_rows_in_batches() { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + SqliteMigrator::new(pool.clone()) + .migrate_up() + .await + .unwrap(); + let expired = sqlite_time(Utc::now() - chrono::Duration::hours(2)); + for _ in 0..5 { + sqlx::query( + "INSERT INTO river_notification (created_at, payload, topic) \ + VALUES (?, 'expired', 'topic')", + ) + .bind(&expired) + .execute(&pool) + .await + .unwrap(); + } + sqlx::query("INSERT INTO river_notification (payload, topic) VALUES ('recent', 'topic')") + .execute(&pool) + .await + .unwrap(); + let client = Client::builder(pool.clone()).build().unwrap(); + let context = ServiceContext { + // Two full batches followed by a partial one. + breakers: Arc::new(Breakers::new(BatchSizes { + default: 2, + reduced: 1, + })), + cancel: CancellationToken::new(), + inner: Arc::clone(&client.inner), + }; + + assert_eq!(clean_notifications(&context, &pool).await.unwrap(), 5); + let remaining: Vec = sqlx::query_scalar("SELECT payload FROM river_notification") + .fetch_all(&pool) + .await + .unwrap(); + assert_eq!(remaining, ["recent"]); + } +} diff --git a/rust/riverqueue/src/maintenance/elector.rs b/rust/riverqueue/src/maintenance/elector.rs new file mode 100644 index 000000000..c3905c8b6 --- /dev/null +++ b/rust/riverqueue/src/maintenance/elector.rs @@ -0,0 +1,787 @@ +//! Leader election, a port of Go's `internal/leadership` elector. + +use std::{sync::Arc, time::Duration}; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +#[cfg(feature = "postgres")] +use sqlx::AssertSqlSafe; +use tokio::{sync::mpsc, time::Instant}; +use tokio_util::sync::CancellationToken; +use tracing::{debug, error, info, warn}; + +use crate::database::DatabasePool; +#[cfg(feature = "sqlite")] +use crate::database::sqlite; +use crate::{Error, client::ClientInner}; + +use super::{LeadershipWakeup, exponential_backoff, random_duration, sleep_cancellable}; + +/// Lease padding added to the elect interval (Go `electIntervalTTLPaddingDefault`). +pub(crate) const LEADER_TTL_PADDING: Duration = Duration::from_secs(10); + +/// Maximum duration of one election or renewal attempt (Go `deadlineTimeout`). +const ATTEMPT_TIMEOUT: Duration = Duration::from_secs(5); + +/// Margin subtracted from the lease when deciding how long this client trusts +/// that it still leads (Go `leaderLocalDeadlineSafetyMargin`). +const LOCAL_DEADLINE_SAFETY_MARGIN: Duration = Duration::from_secs(1); + +/// Go's `ElectIntervalJitter` default. Rust scales it down for short elect +/// intervals so that tests with millisecond intervals stay fast. +const ELECT_INTERVAL_JITTER_MAX: Duration = Duration::from_secs(1); + +/// Upper bound on the random delay before bidding after another client +/// resigns, so that followers do not all bid at once. +const RESIGNED_WAKEUP_JITTER: Duration = Duration::from_millis(50); + +/// Number of resignation attempts on step-down (Go `attemptResignLoop`). +const RESIGN_ATTEMPTS: u32 = 3; + +/// Default exponential backoff reset (Go `MaxAttemptsBeforeResetDefault`). +const BACKOFF_RESET: u32 = 7; + +/// A leadership term held by this client. +#[derive(Clone, Debug)] +pub(crate) struct Term { + /// Database `elected_at` identifying this term. + pub(crate) elected_at: DateTime, + /// Cancelled the moment this client stops trusting the term. + pub(crate) token: CancellationToken, +} + +/// Local view of a held lease (Go `leadershipTerm`). +#[derive(Clone, Copy, Debug)] +struct Lease { + elected_at: DateTime, + trusted_until: Instant, +} + +impl Lease { + /// Trusts the lease until `ttl - 1s` after the attempt *started*, so time + /// spent waiting on a slow database never extends local trust. + fn new(elected_at: DateTime, attempt_started: Instant, ttl: Duration) -> Self { + Self { + elected_at, + trusted_until: attempt_started + ttl.saturating_sub(LOCAL_DEADLINE_SAFETY_MARGIN), + } + } + + fn remaining(&self, now: Instant) -> Duration { + self.trusted_until.saturating_duration_since(now) + } + + fn reelect_attempt_timeout(&self, now: Instant) -> Duration { + self.remaining(now).min(ATTEMPT_TIMEOUT) + } +} + +/// Database operations behind the elector, separated so tests can inject +/// slow or failing renewals. +#[async_trait] +pub(crate) trait LeaderStore: Send + Sync + 'static { + /// Deletes an expired lease and inserts this client's lease if none + /// exists (`ON CONFLICT DO NOTHING`). Returns the new term's `elected_at`. + async fn elect(&self, ttl: Duration) -> Result>, Error>; + + /// Extends the lease only when it still belongs to this client *and* term. + /// Returns the renewed term's `elected_at`, or `None` if it was lost. + async fn reelect( + &self, + elected_at: DateTime, + ttl: Duration, + ) -> Result>, Error>; + + /// Deletes this client's lease for exactly this term and announces it. + async fn resign(&self, elected_at: DateTime) -> Result; +} + +/// Events observed by tests. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum ElectorEvent { + Denied, + Gained, + Lost, + Maintained, + Resigned, +} + +/// Why the leader state ended. +enum StepDown { + /// The client is shutting down; resign before exiting. + Shutdown, + /// Resign and return to the follower state. + Resign { requested: bool }, + /// The lease is gone already; do not resign. + Lost, +} + +pub(crate) struct Elector { + client_id: String, + elect_interval: Duration, + events: Option>, + store: Arc, +} + +impl Elector { + pub(crate) fn new( + store: Arc, + client_id: String, + elect_interval: Duration, + ) -> Self { + Self { + client_id, + elect_interval, + events: None, + store, + } + } + + #[cfg(test)] + pub(crate) fn with_events(mut self, events: mpsc::UnboundedSender) -> Self { + self.events = Some(events); + self + } + + fn ttl(&self) -> Duration { + self.elect_interval + LEADER_TTL_PADDING + } + + fn signal(&self, event: ElectorEvent) { + if let Some(events) = &self.events { + let _ = events.send(event); + } + } + + /// Runs the follower/leader state machine until `cancel` fires. Each + /// gained term is published to the maintainer; its token is cancelled when + /// the term ends. + pub(crate) async fn run( + mut self, + cancel: CancellationToken, + mut wakeups: mpsc::UnboundedReceiver, + terms: mpsc::UnboundedSender, + ) { + let mut cooldown = false; + loop { + let Some(lease) = self.run_follower(&cancel, &mut wakeups, cooldown).await else { + return; + }; + let term = Term { + elected_at: lease.elected_at, + token: cancel.child_token(), + }; + debug!(client_id = %self.client_id, "River client gained leadership"); + self.signal(ElectorEvent::Gained); + let _ = terms.send(term.clone()); + // Like Go, which honors a resignation request only when it + // arrives while this client leads, drop requests that were + // queued while it was bidding. + if !drain_wakeups(&mut wakeups) { + term.token.cancel(); + self.resign(lease.elected_at).await; + return; + } + + let (step_down, lease) = self.run_leader(&cancel, &mut wakeups, lease).await; + term.token.cancel(); + cooldown = false; + match step_down { + StepDown::Shutdown => { + self.resign(lease.elected_at).await; + return; + } + StepDown::Resign { requested } => { + self.resign(lease.elected_at).await; + // Unlike Go, which bids again immediately, give peers one + // elect interval to take the lease after an explicit + // resignation request. Wakeups do not shorten it. + cooldown = requested; + } + StepDown::Lost => {} + } + if cancel.is_cancelled() { + return; + } + } + } + + /// Bids for leadership until elected or cancelled. + async fn run_follower( + &mut self, + cancel: &CancellationToken, + wakeups: &mut mpsc::UnboundedReceiver, + cooldown: bool, + ) -> Option { + if cooldown + && !self + .sleep_ignoring_wakeups(cancel, wakeups, self.elect_interval) + .await + { + return None; + } + let mut attempt = 0_u32; + loop { + attempt += 1; + let attempt_started = Instant::now(); + let result = tokio::select! { + biased; + () = cancel.cancelled() => return None, + result = tokio::time::timeout(ATTEMPT_TIMEOUT, self.store.elect(self.ttl())) => result, + }; + match result { + Ok(Ok(Some(elected_at))) => { + return Some(Lease::new(elected_at, attempt_started, self.ttl())); + } + Ok(Ok(None)) => {} + Ok(Err(elect_error)) => { + let backoff = exponential_backoff(attempt, BACKOFF_RESET); + error!(error = %crate::error::Chain(&elect_error), attempt, ?backoff, "River leader election failed"); + if !sleep_cancellable(cancel, backoff).await { + return None; + } + continue; + } + Err(_) => { + let backoff = exponential_backoff(attempt, BACKOFF_RESET); + error!(attempt, ?backoff, "River leader election timed out"); + if !sleep_cancellable(cancel, backoff).await { + return None; + } + continue; + } + } + + attempt = 0; + self.signal(ElectorEvent::Denied); + let jitter_max = ELECT_INTERVAL_JITTER_MAX.min(self.elect_interval / 5); + let wait = self.elect_interval + random_duration(Duration::ZERO, jitter_max); + tokio::select! { + biased; + () = cancel.cancelled() => return None, + () = tokio::time::sleep(wait) => {} + wakeup = wakeups.recv() => match wakeup { + None => return None, + // A follower ignores resignation requests. + Some(LeadershipWakeup::RequestResign | LeadershipWakeup::ResignTerm(_)) => {} + Some(LeadershipWakeup::Changed) => { + // Somebody resigned; bid soon, but not all at once. + if !sleep_cancellable( + cancel, + random_duration(Duration::ZERO, RESIGNED_WAKEUP_JITTER), + ) + .await + { + return None; + } + } + }, + } + } + } + + /// Keeps renewing the lease until it is lost, its trust window elapses, a + /// resignation is requested, or the client stops. + async fn run_leader( + &mut self, + cancel: &CancellationToken, + wakeups: &mut mpsc::UnboundedReceiver, + mut lease: Lease, + ) -> (StepDown, Lease) { + let mut wait = self.elect_interval; + let mut errors = 0_u32; + loop { + let deadline = tokio::time::sleep(wait); + tokio::pin!(deadline); + loop { + tokio::select! { + biased; + () = cancel.cancelled() => return (StepDown::Shutdown, lease), + wakeup = wakeups.recv() => match wakeup { + None => return (StepDown::Shutdown, lease), + Some(LeadershipWakeup::RequestResign) => { + info!(client_id = %self.client_id, "River leader received a resignation request"); + return (StepDown::Resign { requested: true }, lease); + } + Some(LeadershipWakeup::ResignTerm(elected_at)) if elected_at == lease.elected_at => { + info!(client_id = %self.client_id, "River leader resigning at its maintenance's request"); + return (StepDown::Resign { requested: true }, lease); + } + // A request for an earlier term, or a change a leader + // needn't act on. + Some(LeadershipWakeup::ResignTerm(_) | LeadershipWakeup::Changed) => {} + }, + () = &mut deadline => break, + } + } + + let attempt_started = Instant::now(); + let attempt_timeout = lease.reelect_attempt_timeout(attempt_started); + if attempt_timeout.is_zero() { + warn!( + client_id = %self.client_id, + "River leader stepping down because its renewal deadline elapsed" + ); + self.signal(ElectorEvent::Lost); + return (StepDown::Resign { requested: false }, lease); + } + let result = tokio::select! { + biased; + () = cancel.cancelled() => return (StepDown::Shutdown, lease), + result = tokio::time::timeout( + attempt_timeout, + self.store.reelect(lease.elected_at, self.ttl()), + ) => result, + }; + match result { + Ok(Ok(Some(elected_at))) => { + errors = 0; + lease = Lease::new(elected_at, attempt_started, self.ttl()); + self.signal(ElectorEvent::Maintained); + wait = self.elect_interval; + } + Ok(Ok(None)) => { + info!(client_id = %self.client_id, "River leader lost its lease"); + self.signal(ElectorEvent::Lost); + return (StepDown::Lost, lease); + } + failure => { + errors += 1; + let remaining = lease.remaining(Instant::now()); + if remaining.is_zero() { + warn!( + client_id = %self.client_id, + "River leader stepping down because its renewal deadline elapsed after an error" + ); + self.signal(ElectorEvent::Lost); + return (StepDown::Resign { requested: false }, lease); + } + let backoff = exponential_backoff(errors, 3).min(remaining); + if let Ok(Err(renew_error)) = failure { + error!(error = %crate::error::Chain(&renew_error), attempt = errors, ?backoff, "River leader renewal failed"); + } else { + error!(attempt = errors, ?backoff, "River leader renewal timed out"); + } + if !sleep_cancellable(cancel, backoff).await { + return (StepDown::Shutdown, lease); + } + // Retry immediately because the failed attempt already + // consumed part of this lease's trust window. + wait = Duration::ZERO; + } + } + } + } + + /// Makes a bounded, good-faith attempt to give up the lease even during + /// shutdown. The TTL is the backstop if every attempt fails. + async fn resign(&self, elected_at: DateTime) { + for attempt in 1..=RESIGN_ATTEMPTS { + let timeout = Duration::from_secs(u64::from(attempt)); + match tokio::time::timeout(timeout, self.store.resign(elected_at)).await { + Ok(Ok(resigned)) => { + if resigned { + debug!(client_id = %self.client_id, "River leader resigned"); + self.signal(ElectorEvent::Resigned); + } + return; + } + Ok(Err(resign_error)) => { + error!(error = %crate::error::Chain(&resign_error), attempt, "River leader resignation failed"); + } + Err(_) => error!(attempt, "River leader resignation timed out"), + } + if attempt < RESIGN_ATTEMPTS { + tokio::time::sleep(exponential_backoff(attempt, RESIGN_ATTEMPTS)).await; + } + } + } + + /// Sleeps while draining wakeups, returning `false` on cancellation. + async fn sleep_ignoring_wakeups( + &self, + cancel: &CancellationToken, + wakeups: &mut mpsc::UnboundedReceiver, + duration: Duration, + ) -> bool { + let deadline = tokio::time::sleep(duration); + tokio::pin!(deadline); + loop { + tokio::select! { + biased; + () = cancel.cancelled() => return false, + () = &mut deadline => return true, + wakeup = wakeups.recv() => if wakeup.is_none() { + return false; + }, + } + } + } +} + +/// Discards queued wakeups, returning `false` once the channel is closed. +fn drain_wakeups(wakeups: &mut mpsc::UnboundedReceiver) -> bool { + loop { + match wakeups.try_recv() { + Ok(_) => {} + Err(mpsc::error::TryRecvError::Empty) => return true, + Err(mpsc::error::TryRecvError::Disconnected) => return false, + } + } +} + +/// Leader persistence for the client's configured backend. +pub(crate) struct DatabaseLeaderStore { + inner: Arc, +} + +impl DatabaseLeaderStore { + pub(crate) fn new(inner: Arc) -> Self { + Self { inner } + } +} + +#[async_trait] +impl LeaderStore for DatabaseLeaderStore { + async fn elect(&self, ttl: Duration) -> Result>, Error> { + match self.inner.database.pool() { + #[cfg(feature = "sqlite")] + DatabasePool::Sqlite(pool) => { + let mut transaction = crate::database::begin_sqlite_write(pool).await?; + let now = Utc::now(); + sqlite::leader_delete_expired(&mut transaction, now) + .await + .map_err(sqlite_error)?; + let leader = sqlite::leader_elect(&mut transaction, &self.inner.id, now, ttl) + .await + .map_err(sqlite_error)?; + transaction.commit().await?; + Ok(leader.map(|leader| leader.elected_at)) + } + #[cfg(feature = "postgres")] + DatabasePool::Postgres(pool) => { + let table = self.inner.schema.qualify("river_leader"); + let mut transaction = crate::database::begin_postgres(pool).await?; + sqlx::query(AssertSqlSafe(format!( + "DELETE FROM {table} WHERE expires_at < now()" + ))) + .execute(&mut *transaction) + .await?; + let elected_at = sqlx::query_scalar::<_, DateTime>(AssertSqlSafe(format!( + "INSERT INTO {table} (leader_id, elected_at, expires_at) \ + VALUES ($1, now(), now() + make_interval(secs => $2)) \ + ON CONFLICT (name) DO NOTHING RETURNING elected_at" + ))) + .bind(&self.inner.id) + .bind(ttl.as_secs_f64()) + .fetch_optional(&mut *transaction) + .await?; + transaction.commit().await?; + Ok(elected_at) + } + } + } + + async fn reelect( + &self, + elected_at: DateTime, + ttl: Duration, + ) -> Result>, Error> { + match self.inner.database.pool() { + #[cfg(feature = "sqlite")] + DatabasePool::Sqlite(pool) => { + let mut transaction = crate::database::begin_sqlite_write(pool).await?; + let leader = sqlite::leader_reelect( + &mut transaction, + &self.inner.id, + elected_at, + Utc::now(), + ttl, + ) + .await + .map_err(sqlite_error)?; + transaction.commit().await?; + Ok(leader.map(|leader| leader.elected_at)) + } + #[cfg(feature = "postgres")] + DatabasePool::Postgres(pool) => { + let table = self.inner.schema.qualify("river_leader"); + Ok( + sqlx::query_scalar::<_, DateTime>(AssertSqlSafe(format!( + "UPDATE {table} SET expires_at = now() + make_interval(secs => $1) \ + WHERE elected_at = $2 AND expires_at >= now() AND leader_id = $3 \ + RETURNING elected_at" + ))) + .bind(ttl.as_secs_f64()) + .bind(elected_at) + .bind(&self.inner.id) + .fetch_optional(pool) + .await?, + ) + } + } + } + + async fn resign(&self, elected_at: DateTime) -> Result { + match self.inner.database.pool() { + #[cfg(feature = "sqlite")] + DatabasePool::Sqlite(pool) => { + let mut transaction = crate::database::begin_sqlite_write(pool).await?; + let resigned = sqlite::leader_resign(&mut transaction, &self.inner.id, elected_at) + .await + .map_err(sqlite_error)?; + // Like Go's SQLite driver, a resignation writes no outbox + // row; other clients bid at their next election attempt. + transaction.commit().await?; + Ok(resigned) + } + #[cfg(feature = "postgres")] + DatabasePool::Postgres(pool) => { + let notify = self + .inner + .postgres_capabilities(pool) + .await? + .supports_listen_notify; + let table = self.inner.schema.qualify("river_leader"); + let result = sqlx::query(AssertSqlSafe(format!( + "WITH currently_held_leaders AS (\ + SELECT * FROM {table} WHERE elected_at = $1 AND leader_id = $2 FOR UPDATE\ + ), notified_resignations AS (\ + SELECT CASE WHEN $5::boolean THEN pg_notify(\ + concat(coalesce($3::text, current_schema()), '.', $4::text), \ + json_build_object('leader_id', leader_id, 'action', 'resigned')::text\ + ) END FROM currently_held_leaders\ + ) \ + DELETE FROM {table} USING notified_resignations" + ))) + .bind(elected_at) + .bind(&self.inner.id) + .bind(self.inner.schema.as_deref()) + .bind(crate::NOTIFICATION_TOPIC_LEADERSHIP) + .bind(notify) + .execute(pool) + .await?; + Ok(result.rows_affected() > 0) + } + } + } +} + +#[cfg(feature = "sqlite")] +fn sqlite_error(error: sqlite::BackendError) -> Error { + Error::Database(error.into()) +} + +#[cfg(test)] +mod unit_tests { + use std::sync::Mutex; + + use super::*; + + /// A lease store whose renewals can be slowed or failed. + struct ScriptedStore { + elected_at: DateTime, + reelect: Mutex>, + resigned: Mutex>>, + } + + enum Reelect { + /// Succeeds only after local time passes the whole TTL. + SlowSuccess(Duration), + /// Fails after local time passes the whole TTL. + SlowFailure(Duration), + } + + #[async_trait] + impl LeaderStore for ScriptedStore { + async fn elect(&self, _ttl: Duration) -> Result>, Error> { + Ok(Some(self.elected_at)) + } + + async fn reelect( + &self, + elected_at: DateTime, + _ttl: Duration, + ) -> Result>, Error> { + let step = self.reelect.lock().unwrap().pop(); + match step { + Some(Reelect::SlowSuccess(delay)) => { + tokio::time::advance(delay).await; + Ok(Some(elected_at)) + } + Some(Reelect::SlowFailure(delay)) => { + tokio::time::advance(delay).await; + Err(Error::runtime_context("leader election", "renewal failed")) + } + // Later terms only need to stay alive until the test stops. + None => std::future::pending().await, + } + } + + async fn resign(&self, elected_at: DateTime) -> Result { + self.resigned.lock().unwrap().push(elected_at); + Ok(true) + } + } + + async fn run_scripted( + reelect: Reelect, + ) -> (Vec, Vec>, DateTime) { + let elect_interval = Duration::from_millis(100); + let elected_at = Utc::now(); + let store = Arc::new(ScriptedStore { + elected_at, + reelect: Mutex::new(vec![reelect]), + resigned: Mutex::new(Vec::new()), + }); + let (events_sender, mut events) = mpsc::unbounded_channel(); + let (_wakeup_sender, wakeups) = mpsc::unbounded_channel(); + let (terms_sender, mut terms) = mpsc::unbounded_channel(); + let cancel = CancellationToken::new(); + let elector = Elector::new( + Arc::clone(&store) as Arc, + "scripted".to_owned(), + elect_interval, + ) + .with_events(events_sender); + let run = tokio::spawn(elector.run(cancel.clone(), wakeups, terms_sender)); + + let term = terms.recv().await.unwrap(); + // Wait for the term to end, then stop before the next election. + term.token.cancelled().await; + let mut observed = Vec::new(); + while let Some(event) = events.recv().await { + observed.push(event); + if event == ElectorEvent::Resigned { + break; + } + } + cancel.cancel(); + run.await.unwrap(); + let resigned = store.resigned.lock().unwrap().clone(); + (observed, resigned, elected_at) + } + + #[tokio::test(start_paused = true)] + async fn ignores_resign_requests_queued_before_gaining_leadership() { + let elected_at = Utc::now(); + let store = Arc::new(ScriptedStore { + elected_at, + reelect: Mutex::new(Vec::new()), + resigned: Mutex::new(Vec::new()), + }); + let (wakeup_sender, wakeups) = mpsc::unbounded_channel(); + let (terms_sender, mut terms) = mpsc::unbounded_channel(); + let cancel = CancellationToken::new(); + // Queued while the client is still a follower, as when a request + // arrives during its election attempt. + wakeup_sender.send(LeadershipWakeup::RequestResign).unwrap(); + let elector = Elector::new( + Arc::clone(&store) as Arc, + "stale-resign".to_owned(), + Duration::from_millis(100), + ); + let run = tokio::spawn(elector.run(cancel.clone(), wakeups, terms_sender)); + + let term = terms.recv().await.unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(!term.token.is_cancelled()); + assert!(store.resigned.lock().unwrap().is_empty()); + + // A request that arrives while the client leads is honored. + wakeup_sender.send(LeadershipWakeup::RequestResign).unwrap(); + term.token.cancelled().await; + cancel.cancel(); + run.await.unwrap(); + assert_eq!(store.resigned.lock().unwrap().first(), Some(&elected_at)); + } + + #[tokio::test(start_paused = true)] + async fn resigns_only_the_named_term() { + let elected_at = Utc::now(); + let store = Arc::new(ScriptedStore { + elected_at, + reelect: Mutex::new(Vec::new()), + resigned: Mutex::new(Vec::new()), + }); + let (wakeup_sender, wakeups) = mpsc::unbounded_channel(); + let (terms_sender, mut terms) = mpsc::unbounded_channel(); + let cancel = CancellationToken::new(); + let elector = Elector::new( + Arc::clone(&store) as Arc, + "term-resign".to_owned(), + Duration::from_millis(100), + ); + let run = tokio::spawn(elector.run(cancel.clone(), wakeups, terms_sender)); + + let term = terms.recv().await.unwrap(); + // A request for an earlier term, like one from a start that failed + // after that term ended, is ignored. + wakeup_sender + .send(LeadershipWakeup::ResignTerm( + elected_at - chrono::Duration::seconds(1), + )) + .unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(!term.token.is_cancelled()); + assert!(store.resigned.lock().unwrap().is_empty()); + + wakeup_sender + .send(LeadershipWakeup::ResignTerm(elected_at)) + .unwrap(); + term.token.cancelled().await; + cancel.cancel(); + run.await.unwrap(); + assert_eq!(store.resigned.lock().unwrap().first(), Some(&elected_at)); + } + + #[tokio::test(start_paused = true)] + async fn slow_successful_reelect_does_not_extend_trust_window() { + let ttl = Duration::from_millis(100) + LEADER_TTL_PADDING; + let (events, resigned, elected_at) = run_scripted(Reelect::SlowSuccess(ttl)).await; + // The renewal succeeded, but it started a full TTL ago, so the next + // renewal is never attempted and the term is resigned instead. + assert_eq!( + events, + [ + ElectorEvent::Gained, + ElectorEvent::Maintained, + ElectorEvent::Lost, + ElectorEvent::Resigned + ] + ); + assert_eq!(resigned.first(), Some(&elected_at)); + } + + #[tokio::test(start_paused = true)] + async fn resigns_current_term_after_reelect_errors_exhaust_trust() { + let ttl = Duration::from_millis(100) + LEADER_TTL_PADDING; + let (events, resigned, elected_at) = run_scripted(Reelect::SlowFailure(ttl)).await; + assert_eq!( + events, + [ + ElectorEvent::Gained, + ElectorEvent::Lost, + ElectorEvent::Resigned + ] + ); + assert_eq!(resigned.first(), Some(&elected_at)); + } + + #[test] + fn lease_trust_is_measured_from_attempt_start() { + let started = Instant::now(); + let lease = Lease::new(Utc::now(), started, Duration::from_secs(15)); + assert_eq!(lease.remaining(started), Duration::from_secs(14)); + assert_eq!(lease.reelect_attempt_timeout(started), ATTEMPT_TIMEOUT); + assert_eq!( + lease.reelect_attempt_timeout(started + Duration::from_secs(12)), + Duration::from_secs(2) + ); + assert!( + lease + .reelect_attempt_timeout(started + Duration::from_secs(20)) + .is_zero() + ); + } +} diff --git a/rust/riverqueue/src/maintenance/maintainer.rs b/rust/riverqueue/src/maintenance/maintainer.rs new file mode 100644 index 000000000..c707c101d --- /dev/null +++ b/rust/riverqueue/src/maintenance/maintainer.rs @@ -0,0 +1,320 @@ +//! Starts and stops leader-owned services on leadership transitions, a port +//! of Go's `QueueMaintainerLeader` and `QueueMaintainer`. + +use std::{sync::Arc, time::Duration}; + +use chrono::{DateTime, Utc}; +use tokio::{sync::mpsc, task::JoinSet}; +use tokio_util::sync::CancellationToken; +use tracing::{debug, error}; + +use crate::database::DatabasePool; +use crate::{Client, client::ClientInner}; + +use super::{ + Breakers, LeadershipWakeup, MaintenanceError, STAGGER_MAX, cleaner, elector::Term, + exponential_backoff, periodic_enqueuer, random_duration, rescuer, scheduler, sleep_cancellable, +}; + +/// Attempts to start the maintainer before requesting resignation (Go +/// `queueMaintainerMaxStartAttempts`). +const START_ATTEMPTS: u32 = 3; + +/// Shared inputs of every service in one term. +pub(super) struct ServiceContext { + pub(super) breakers: Arc, + pub(super) cancel: CancellationToken, + pub(super) inner: Arc, +} + +impl ServiceContext { + pub(super) fn client(&self) -> Client { + Client { + inner: Arc::clone(&self.inner), + } + } +} + +pub(super) struct Maintainer { + breakers: Arc, + inner: Arc, + /// Asks this client's elector to resign a term. + resign: mpsc::UnboundedSender, +} + +impl Maintainer { + pub(super) fn new( + inner: Arc, + resign: mpsc::UnboundedSender, + ) -> Self { + Self { + breakers: Arc::new(Breakers::new(inner.maintenance.batch_sizes)), + inner, + resign, + } + } + + /// Runs one term at a time. A new term starts only after every service of + /// the previous term has returned, so services never overlap across terms. + pub(super) async fn run( + self, + cancel: CancellationToken, + mut terms: mpsc::UnboundedReceiver, + ) { + let mut previous: Option> = None; + loop { + let term = tokio::select! { + biased; + () = cancel.cancelled() => break, + term = terms.recv() => match term { + Some(term) => term, + None => break, + }, + }; + if let Some(handle) = previous.take() { + join_term(handle).await; + } + if term.token.is_cancelled() { + continue; + } + previous = Some(tokio::spawn(run_term( + Arc::clone(&self.inner), + Arc::clone(&self.breakers), + self.resign.clone(), + term, + ))); + } + if let Some(handle) = previous { + join_term(handle).await; + } + } +} + +async fn join_term(handle: tokio::task::JoinHandle<()>) { + if let Err(join_error) = handle.await { + error!(error = %join_error, "River maintenance term task failed"); + } +} + +/// Starts the term's services, retrying start failures and requesting +/// resignation once retries are exhausted, then waits for every service. +async fn run_term( + inner: Arc, + breakers: Arc, + resign: mpsc::UnboundedSender, + term: Term, +) { + let Term { + elected_at, + token: cancel, + } = term; + if !start_or_resign(&inner, &cancel, &resign, elected_at).await { + return; + } + + let context = Arc::new(ServiceContext { + breakers, + cancel: cancel.clone(), + inner: Arc::clone(&inner), + }); + let mut services = JoinSet::new(); + services.spawn(periodic_enqueuer::run(Arc::clone(&context))); + services.spawn(run_periodically( + Arc::clone(&context), + "job scheduler", + inner.maintenance.scheduler_interval, + |context| Box::pin(async move { scheduler::run_once(&context).await }), + )); + services.spawn(run_periodically( + Arc::clone(&context), + "job rescuer", + inner.maintenance.rescuer_interval, + |context| Box::pin(async move { rescuer::run_once(&context).await }), + )); + services.spawn(run_periodically( + Arc::clone(&context), + "job cleaner", + inner.maintenance.job_cleaner_interval, + |context| Box::pin(async move { cleaner::clean_jobs(&context).await }), + )); + services.spawn(run_periodically( + Arc::clone(&context), + "queue cleaner", + inner.maintenance.queue_cleaner_interval, + |context| Box::pin(async move { cleaner::clean_queues(&context).await }), + )); + match inner.database.pool() { + #[cfg(feature = "sqlite")] + DatabasePool::Sqlite(pool) => { + let pool = pool.clone(); + services.spawn(run_periodically( + Arc::clone(&context), + "SQLite notification cleaner", + cleaner::NOTIFICATION_CLEANER_INTERVAL, + move |context| { + let pool = pool.clone(); + Box::pin(async move { + cleaner::clean_notifications(&context, &pool) + .await + .map(|_| ()) + }) + }, + )); + } + #[cfg(feature = "postgres")] + DatabasePool::Postgres(pool) => { + services.spawn(super::reindexer::run(Arc::clone(&context), pool.clone())); + } + } + let term = crate::__private::LeaderTerm { + elected_at, + token: cancel.clone(), + }; + for service in inner.pilot.maintenance_services() { + services.spawn(supervise_extension_service( + service, + crate::client::WeakClient::new(&inner), + inner.pilot_database(), + term.clone(), + )); + } + + while let Some(result) = services.join_next().await { + if let Err(join_error) = result { + error!(error = %join_error, "River maintenance service task failed"); + } + } + debug!("River maintenance services stopped"); +} + +/// Starts the term's maintenance, retrying start failures and asking this +/// client's elector to resign the term once retries are exhausted. Returns +/// whether the term's services should run. +async fn start_or_resign( + inner: &Arc, + cancel: &CancellationToken, + resign: &mpsc::UnboundedSender, + elected_at: DateTime, +) -> bool { + for attempt in 1..=START_ATTEMPTS { + match start(inner, cancel).await { + Ok(()) => return true, + Err(MaintenanceError::Cancelled) => return false, + Err(start_error) => { + error!(error = %crate::error::Chain(&start_error), attempt, "River maintenance start failed"); + if attempt < START_ATTEMPTS + && !sleep_cancellable(cancel, exponential_backoff(attempt, 7)).await + { + return false; + } + } + } + } + if cancel.is_cancelled() { + return false; + } + // Resign locally rather than through a notification, which a client + // without notifications wouldn't hear. Naming the term keeps a late + // failure from resigning a newer one. + error!("River maintenance failed to start after all attempts; resigning leadership"); + let _ = resign.send(LeadershipWakeup::ResignTerm(elected_at)); + false +} + +/// Runs an extension's maintenance service for the whole term, restarting it +/// after River's service backoff when it fails, panics, or returns before the +/// term ends. Each run settles before the next starts, and the backoff starts +/// over after a long healthy run. +async fn supervise_extension_service( + service: Arc, + client: crate::client::WeakClient, + database: crate::__private::PilotDatabase, + term: crate::__private::LeaderTerm, +) { + let mut attempt = 0; + loop { + let started_at = tokio::time::Instant::now(); + // A task of its own, so a panic ends only this run. + let mut run = JoinSet::new(); + let context = crate::__private::MaintenanceServiceContext { + client: client.clone(), + database: database.clone(), + term: term.clone(), + }; + let task_service = Arc::clone(&service); + run.spawn(async move { task_service.run(context).await }); + let outcome = run.join_next().await; + if term.token.is_cancelled() { + return; + } + if started_at.elapsed() >= crate::client::SERVICE_RESTART_RESET_AFTER { + attempt = 0; + } + attempt += 1; + let delay = exponential_backoff(attempt, 7); + let failure = match outcome { + Some(Ok(Ok(()))) | None => "returned before its leadership term ended".to_owned(), + Some(Ok(Err(service_error))) => crate::error::Chain(&*service_error).to_string(), + Some(Err(join_error)) => join_error.to_string(), + }; + error!( + service = service.name(), + attempt, + error = %failure, + sleep_duration = ?delay, + "River extension maintenance service failed; restarting after backoff" + ); + if !sleep_cancellable(&term.token, delay).await { + return; + } + } +} + +/// Mirrors the only fallible part of Go's `QueueMaintainer.Start`: the +/// periodic job enqueuer runs start hooks on every leadership gain. +async fn start(inner: &ClientInner, cancel: &CancellationToken) -> Result<(), MaintenanceError> { + for hook in &inner.hooks { + tokio::select! { + biased; + () = cancel.cancelled() => return Err(MaintenanceError::Cancelled), + result = hook.periodic_jobs_start(&inner.periodic_jobs) => result?, + } + } + Ok(()) +} + +type RunOnceFuture = std::pin::Pin> + Send>>; + +/// Runs a service on an interval with an initial staggered tick, like Go's +/// `StaggerStart` plus `NewTickerWithInitialTick`. The stagger is capped by +/// the interval so short test intervals stay responsive. +async fn run_periodically( + context: Arc, + name: &'static str, + interval: Duration, + run_once: impl Fn(Arc) -> RunOnceFuture + Send + 'static, +) { + let stagger = random_duration(Duration::ZERO, STAGGER_MAX.min(interval)); + if !sleep_cancellable(&context.cancel, stagger).await { + return; + } + let mut ticker = tokio::time::interval(interval); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + loop { + tokio::select! { + biased; + () = context.cancel.cancelled() => return, + _ = ticker.tick() => {} + } + match run_once(Arc::clone(&context)).await { + Ok(()) => {} + Err(MaintenanceError::Cancelled) => return, + Err(run_error) => { + if context.cancel.is_cancelled() { + return; + } + error!(error = %crate::error::Chain(&run_error), service = name, "River maintenance service failed"); + } + } + } +} diff --git a/rust/riverqueue/src/maintenance/mod.rs b/rust/riverqueue/src/maintenance/mod.rs new file mode 100644 index 000000000..27a6cfa73 --- /dev/null +++ b/rust/riverqueue/src/maintenance/mod.rs @@ -0,0 +1,426 @@ +//! Leader election and leader-owned database maintenance. +//! +//! The structure mirrors Go's `internal/leadership` and +//! `internal/maintenance` packages: +//! +//! - An elector task renews leadership on its own schedule, bounds every +//! attempt by a deadline and by the locally trusted remainder of the current +//! term, and guards renewal and resignation with the term's `elected_at`. +//! - A maintainer task (Go's `QueueMaintainerLeader`) starts every +//! leader-owned service in its own task under a per-term cancellation token. +//! The elector cancels that token the moment leadership is lost or its trust +//! window elapses, and a new term starts only after the previous term's +//! services have stopped. +//! - Every maintenance database call is selected against the term token. On +//! PostgreSQL, cancellation and timeouts are enforced server-side with +//! `SET LOCAL statement_timeout` and `pg_cancel_backend`, so abandoned work +//! does not keep holding locks. + +mod cleaner; +mod elector; +mod maintainer; +mod periodic_enqueuer; +#[cfg(feature = "postgres")] +mod postgres; +#[cfg(feature = "postgres")] +mod reindexer; +mod rescuer; +mod scheduler; +#[cfg(all(test, feature = "postgres", river_postgres_tests))] +mod tests; + +#[cfg(feature = "postgres")] +pub(crate) use cleaner::postgres_delete_finalized_jobs; +#[cfg(feature = "sqlite")] +pub(crate) use cleaner::sqlite_delete_finalized_jobs; + +use std::{ + hash::{BuildHasher, Hasher}, + sync::{Arc, Mutex}, + time::Duration, +}; + +use tokio::sync::{broadcast, mpsc}; +use tokio_util::sync::CancellationToken; +use tracing::warn; + +use crate::{Error, client::ClientInner}; + +/// Batch size used by bulk maintenance services (Go `BatchSizeDefault`). +pub(crate) const BATCH_SIZE_DEFAULT: i64 = 10_000; + +/// Batch size used after repeated timeouts (Go `BatchSizeReduced`). +pub(crate) const BATCH_SIZE_REDUCED: i64 = 1_000; + +/// Bounds of the pause between maintenance batches (Go `BatchBackoffMin/Max`). +const BATCH_BACKOFF_MIN: Duration = Duration::from_millis(50); +const BATCH_BACKOFF_MAX: Duration = Duration::from_secs(1); + +/// Timeout for one maintenance batch (Go `riversharedmaintenance.TimeoutDefault`). +const TIMEOUT_DEFAULT: Duration = Duration::from_secs(30); + +/// Maximum random delay before a service's first run (Go `StaggerStart`). +const STAGGER_MAX: Duration = Duration::from_secs(1); + +/// Runs election and leader-owned maintenance until `cancel` fires. +pub(crate) async fn run_maintenance( + inner: Arc, + cancel: CancellationToken, + notifications: broadcast::Receiver, +) -> Result<(), Error> { + // Each term's services run in detached tasks under child tokens of + // `cancel`. Should this task panic or be aborted, the supervisor + // restarts it with a new token, so cancel this one on the way out to + // stop the old term's services rather than leave a second leader's + // maintenance running. + let _cancel_on_exit = cancel.clone().drop_guard(); + let (wakeup_sender, wakeup_receiver) = mpsc::unbounded_channel(); + let (term_sender, term_receiver) = mpsc::unbounded_channel(); + let elector = elector::Elector::new( + Arc::new(elector::DatabaseLeaderStore::new(Arc::clone(&inner))), + inner.id.clone(), + inner.maintenance.elect_interval, + ); + let maintainer = maintainer::Maintainer::new(Arc::clone(&inner), wakeup_sender.clone()); + tokio::join!( + forward_leadership_notifications(notifications, wakeup_sender, cancel.clone()), + elector.run(cancel.clone(), wakeup_receiver, term_sender), + maintainer.run(cancel, term_receiver), + ); + Ok(()) +} + +/// Leadership events the elector reacts to. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum LeadershipWakeup { + /// Another client resigned, or leadership may otherwise have changed. + Changed, + /// Some client asked the current leader to resign. + RequestResign, + /// This client's maintenance asks it to resign the term elected at this + /// time, and no later one. It never passes through the database, so it + /// works without notifications. + ResignTerm(chrono::DateTime), +} + +/// Moves leadership notifications from the client's broadcast channel into +/// the elector's unbounded queue. +/// +/// The channel carries only leadership events, which are rare, so it lagging +/// would take over a thousand of them arriving at once. Should it ever lag, +/// the relay still emits a wakeup so the elector re-checks the lease rather +/// than silently missing a transition. +async fn forward_leadership_notifications( + mut notifications: broadcast::Receiver, + wakeups: mpsc::UnboundedSender, + cancel: CancellationToken, +) { + loop { + let notification = tokio::select! { + biased; + () = cancel.cancelled() => return, + notification = notifications.recv() => notification, + }; + let wakeup = match notification { + Ok(wakeup) => wakeup, + Err(broadcast::error::RecvError::Lagged(count)) => { + warn!( + skipped = count, + "River leadership relay lagged; re-checking leadership" + ); + LeadershipWakeup::Changed + } + Err(broadcast::error::RecvError::Closed) => return, + }; + if wakeups.send(wakeup).is_err() { + return; + } + } +} + +/// Failure of one maintenance operation. +#[derive(Debug)] +pub(crate) enum MaintenanceError { + /// The term or client was cancelled; nothing needs to be reported. + Cancelled, + /// A database call exceeded its deadline. + TimedOut, + /// Any other failure. + Failed(Error), +} + +impl std::fmt::Display for MaintenanceError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Cancelled => formatter.write_str("cancelled"), + Self::TimedOut => formatter.write_str("timed out"), + Self::Failed(error) => error.fmt(formatter), + } + } +} + +impl std::error::Error for MaintenanceError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + match self { + // Displayed transparently, so its source is the error's own. + Self::Failed(error) => error.source(), + Self::Cancelled | Self::TimedOut => None, + } + } +} + +impl From for MaintenanceError { + fn from(error: Error) -> Self { + Self::Failed(error) + } +} + +impl From for MaintenanceError { + fn from(error: sqlx::Error) -> Self { + // PostgreSQL reports a `statement_timeout` expiry as `query_canceled`. + // Explicit cancellation through `pg_cancel_backend` is reported as + // `Cancelled` before a database error is ever mapped. + if let sqlx::Error::Database(database_error) = &error + && database_error.code().as_deref() == Some("57014") + { + return Self::TimedOut; + } + Self::Failed(error.into()) + } +} + +#[cfg(feature = "sqlite")] +impl From for MaintenanceError { + fn from(error: crate::database::sqlite::BackendError) -> Self { + Self::Failed(Error::Database(error.into())) + } +} + +/// Batch sizes of one maintenance service. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) struct BatchSizes { + pub(crate) default: i64, + pub(crate) reduced: i64, +} + +impl Default for BatchSizes { + fn default() -> Self { + Self { + default: BATCH_SIZE_DEFAULT, + reduced: BATCH_SIZE_REDUCED, + } + } +} + +/// Go's reduced-batch circuit breaker: three timeouts within ten minutes +/// switch a service to its reduced batch size for the life of the client. +#[derive(Debug)] +pub(crate) struct ReducedBatchBreaker { + open: bool, + sizes: BatchSizes, + trips: Vec, +} + +impl ReducedBatchBreaker { + const LIMIT: usize = 3; + const WINDOW: Duration = Duration::from_mins(10); + + fn new(sizes: BatchSizes) -> Self { + Self { + open: false, + sizes, + trips: Vec::new(), + } + } + + fn batch_size(&self) -> i64 { + if self.open { + self.sizes.reduced + } else { + self.sizes.default + } + } + + fn reset_if_not_open(&mut self) { + if !self.open { + self.trips.clear(); + } + } + + fn trip(&mut self) { + if self.open { + return; + } + let now = tokio::time::Instant::now(); + self.trips + .retain(|trip| now.saturating_duration_since(*trip) <= Self::WINDOW); + self.trips.push(now); + if self.trips.len() >= Self::LIMIT { + self.open = true; + } + } +} + +/// Breakers shared by every term of one running client, like Go's services +/// which keep their breakers across leadership changes. +#[derive(Debug)] +pub(crate) struct Breakers { + pub(crate) job_cleaner: Mutex, + #[cfg(feature = "sqlite")] + pub(crate) notification_cleaner: Mutex, + pub(crate) queue_cleaner: Mutex, + pub(crate) rescuer: Mutex, + pub(crate) scheduler: Mutex, +} + +impl Breakers { + fn new(sizes: BatchSizes) -> Self { + Self { + job_cleaner: Mutex::new(ReducedBatchBreaker::new(sizes)), + #[cfg(feature = "sqlite")] + notification_cleaner: Mutex::new(ReducedBatchBreaker::new(sizes)), + queue_cleaner: Mutex::new(ReducedBatchBreaker::new(sizes)), + rescuer: Mutex::new(ReducedBatchBreaker::new(sizes)), + scheduler: Mutex::new(ReducedBatchBreaker::new(sizes)), + } + } +} + +/// Runs one batch-oriented maintenance operation, feeding its outcome into +/// the service's breaker. +fn batch_size(breaker: &Mutex) -> i64 { + breaker + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .batch_size() +} + +fn record_batch(breaker: &Mutex, result: &Result) { + let mut breaker = breaker + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + match result { + Ok(_) => breaker.reset_if_not_open(), + Err(MaintenanceError::TimedOut) => breaker.trip(), + Err(_) => {} + } +} + +/// Returns a random `u64` for jitter and identifiers. Not cryptographic. +pub(crate) fn random_u64() -> u64 { + // `RandomState` is seeded randomly per process and advanced per instance. + let mut hasher = std::collections::hash_map::RandomState::new().build_hasher(); + hasher.write_u128( + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_nanos(), + ); + hasher.finish() +} + +/// Returns a uniformly distributed duration in `[minimum, maximum)`. +pub(crate) fn random_duration(minimum: Duration, maximum: Duration) -> Duration { + if maximum <= minimum { + return minimum; + } + let span_nanos = u64::try_from(maximum.saturating_sub(minimum).as_nanos()).unwrap_or(u64::MAX); + minimum + Duration::from_nanos(random_u64() % span_nanos.max(1)) +} + +/// Go's `serviceutil.ExponentialBackoff`: `2^(attempt-1)` seconds with +/-10% +/// jitter, restarting after `max_attempts_before_reset` attempts. +pub(crate) fn exponential_backoff(attempt: u32, max_attempts_before_reset: u32) -> Duration { + let attempt = attempt.saturating_sub(1) % max_attempts_before_reset.max(1); + let seconds = f64::from(2_u32.saturating_pow(attempt)); + let jitter = random_duration(Duration::ZERO, Duration::from_secs(1)).as_secs_f64(); + Duration::from_secs_f64(seconds + seconds * (jitter * 0.2 - 0.1)) +} + +/// Sleeps unless cancellation happens first. Returns `false` when cancelled. +pub(crate) async fn sleep_cancellable(cancel: &CancellationToken, duration: Duration) -> bool { + tokio::select! { + biased; + () = cancel.cancelled() => false, + () = tokio::time::sleep(duration) => true, + } +} + +/// Pauses between batches of a large maintenance backlog. +async fn batch_backoff(cancel: &CancellationToken) -> Result<(), MaintenanceError> { + if sleep_cancellable( + cancel, + random_duration(BATCH_BACKOFF_MIN, BATCH_BACKOFF_MAX), + ) + .await + { + Ok(()) + } else { + Err(MaintenanceError::Cancelled) + } +} + +/// Runs a SQLite operation under the term token and a client-side deadline. +/// +/// SQLite executes statements on the client, so dropping the future cannot +/// leave server-side work holding locks the way it can on PostgreSQL. +#[cfg(feature = "sqlite")] +async fn sqlite_cancellable( + cancel: &CancellationToken, + timeout: Duration, + operation: impl Future>, +) -> Result +where + MaintenanceError: From, +{ + tokio::select! { + biased; + () = cancel.cancelled() => Err(MaintenanceError::Cancelled), + result = tokio::time::timeout(timeout, operation) => match result { + Ok(result) => result.map_err(MaintenanceError::from), + Err(_) => Err(MaintenanceError::TimedOut), + }, + } +} + +#[cfg(test)] +mod unit_tests { + use super::{BatchSizes, ReducedBatchBreaker, exponential_backoff}; + + #[tokio::test(start_paused = true)] + async fn reduced_batch_breaker_opens_after_three_timeouts_in_ten_minutes() { + let sizes = BatchSizes { + default: 10, + reduced: 2, + }; + let mut breaker = ReducedBatchBreaker::new(sizes); + breaker.trip(); + breaker.trip(); + // A success between failures resets the count. + breaker.reset_if_not_open(); + breaker.trip(); + breaker.trip(); + assert_eq!(breaker.batch_size(), 10); + // Trips older than the window no longer count. + tokio::time::advance(std::time::Duration::from_mins(11)).await; + breaker.trip(); + assert_eq!(breaker.batch_size(), 10); + breaker.trip(); + breaker.trip(); + assert_eq!(breaker.batch_size(), 2); + // Once open, the breaker stays open. + breaker.reset_if_not_open(); + assert_eq!(breaker.batch_size(), 2); + } + + #[test] + fn exponential_backoff_matches_go_schedule() { + for (attempt, seconds) in [(1, 1.0), (2, 2.0), (3, 4.0), (7, 64.0), (8, 1.0)] { + let backoff = exponential_backoff(attempt, 7).as_secs_f64(); + assert!( + (seconds * 0.9..=seconds * 1.1).contains(&backoff), + "attempt {attempt}: {backoff}" + ); + } + } +} diff --git a/rust/riverqueue/src/maintenance/periodic_enqueuer.rs b/rust/riverqueue/src/maintenance/periodic_enqueuer.rs new file mode 100644 index 000000000..4a3dd456e --- /dev/null +++ b/rust/riverqueue/src/maintenance/periodic_enqueuer.rs @@ -0,0 +1,42 @@ +//! Inserts periodic jobs while this client leads, a port of Go's +//! `PeriodicJobEnqueuer` run loop. + +use std::{sync::Arc, time::Duration}; + +use chrono::Utc; + +use super::maintainer::ServiceContext; + +/// Sleep used when no periodic job is scheduled (Go's "very long duration"). +const IDLE_WAIT: Duration = Duration::from_hours(24); + +/// Runs periodic jobs for one leadership term. +/// +/// Every term starts from a fresh schedule computed from the time leadership +/// began, like Go's enqueuer `Start`, and inserts `run_on_start` jobs once per +/// gained term. +pub(super) async fn run(context: Arc) { + let periodic_jobs = context.inner.periodic_jobs.clone(); + periodic_jobs.reset_for_leadership(); + let client = context.client(); + loop { + // Subscribe before scheduling so a job added while this pass runs + // still wakes the next wait. + let changed = periodic_jobs.changed(); + tokio::pin!(changed); + changed.as_mut().enable(); + + periodic_jobs + .run_due(&client, Utc::now(), &context.cancel) + .await; + let wait = periodic_jobs.next_run_at().map_or(IDLE_WAIT, |next| { + (next - Utc::now()).to_std().unwrap_or_default() + }); + tokio::select! { + biased; + () = context.cancel.cancelled() => return, + () = &mut changed => {} + () = tokio::time::sleep(wait) => {} + } + } +} diff --git a/rust/riverqueue/src/maintenance/postgres.rs b/rust/riverqueue/src/maintenance/postgres.rs new file mode 100644 index 000000000..44ecd1f7e --- /dev/null +++ b/rust/riverqueue/src/maintenance/postgres.rs @@ -0,0 +1,160 @@ +//! PostgreSQL helpers that make maintenance statements cancellable and +//! bounded on the server, not just abandoned by the client. + +use std::time::Duration; + +use sqlx::{PgPool, Postgres, Transaction}; +use tokio_util::sync::CancellationToken; +use tracing::debug; + +use super::MaintenanceError; + +/// Extra client-side patience beyond a server-side `statement_timeout` +/// before the client gives up on an unresponsive connection. +const CLIENT_TIMEOUT_GRACE: Duration = Duration::from_secs(5); + +/// Upper bound on waiting for a cancelled statement to report back. +const CANCEL_GRACE: Duration = Duration::from_secs(5); + +/// A maintenance transaction whose statements are bounded by +/// `SET LOCAL statement_timeout` and can be cancelled server-side. +pub(super) struct MaintenanceTransaction { + pub(super) backend_pid: i32, + pub(super) timeout: Duration, + pub(super) transaction: Transaction<'static, Postgres>, +} + +impl MaintenanceTransaction { + /// Begins a transaction, records its backend PID, and applies a + /// transaction-local statement timeout. Nothing outlives the transaction, + /// so pooled connections keep their role and database defaults. + pub(super) async fn begin( + pool: &PgPool, + cancel: &CancellationToken, + timeout: Duration, + ) -> Result { + // Dropping the begin when cancellation wins is safe: River begins on + // its own task, which rolls the transaction back once it starts. + let mut transaction = tokio::select! { + biased; + () = cancel.cancelled() => return Err(MaintenanceError::Cancelled), + transaction = crate::database::begin_postgres(pool) => transaction?, + }; + let backend_pid = tokio::select! { + biased; + () = cancel.cancelled() => return Err(MaintenanceError::Cancelled), + row = sqlx::query_as::<_, (i32, String)>( + "SELECT pg_backend_pid(), set_config('statement_timeout', $1, true)", + ) + .bind(timeout_setting(timeout)) + .fetch_one(&mut *transaction) => row?.0, + }; + Ok(Self { + backend_pid, + timeout, + transaction, + }) + } + + /// Commits under the same cancellation and deadline rules. + pub(super) async fn commit( + self, + pool: &PgPool, + cancel: &CancellationToken, + ) -> Result<(), MaintenanceError> { + let backend_pid = self.backend_pid; + let timeout = self.timeout; + cancellable( + pool, + backend_pid, + cancel, + timeout, + self.transaction.commit(), + ) + .await + } +} + +/// Formats a duration for `statement_timeout`, in whole milliseconds. +fn timeout_setting(timeout: Duration) -> String { + timeout.as_millis().max(1).to_string() +} + +/// Runs one statement on the connection identified by `backend_pid`. +/// +/// When `cancel` fires first, the statement is cancelled server-side with +/// `pg_cancel_backend` and awaited briefly so the connection is idle again. +/// A server that never answers is abandoned after the statement timeout plus a +/// grace period. +pub(super) async fn cancellable( + pool: &PgPool, + backend_pid: i32, + cancel: &CancellationToken, + timeout: Duration, + operation: impl Future>, +) -> Result { + cancellable_within( + pool, + backend_pid, + cancel, + Some(timeout.saturating_add(CLIENT_TIMEOUT_GRACE)), + operation, + ) + .await +} + +/// Like [`cancellable`], but for a statement without a server-side +/// `statement_timeout`: the statement is cancelled server-side once +/// `deadline` passes, or never when it's `None`. +pub(super) async fn cancellable_within( + pool: &PgPool, + backend_pid: i32, + cancel: &CancellationToken, + deadline: Option, + operation: impl Future>, +) -> Result { + tokio::pin!(operation); + let deadline = async { + match deadline { + Some(deadline) => tokio::time::sleep(deadline).await, + None => std::future::pending().await, + } + }; + tokio::select! { + biased; + result = &mut operation => result.map_err(MaintenanceError::from), + () = cancel.cancelled() => { + cancel_backend(pool, backend_pid).await; + let _ = tokio::time::timeout(CANCEL_GRACE, operation).await; + Err(MaintenanceError::Cancelled) + } + () = deadline => { + cancel_backend(pool, backend_pid).await; + let _ = tokio::time::timeout(CANCEL_GRACE, operation).await; + Err(MaintenanceError::TimedOut) + } + } +} + +/// Asks PostgreSQL to cancel the statement running on `backend_pid`. +pub(super) async fn cancel_backend(pool: &PgPool, backend_pid: i32) { + let result = tokio::time::timeout( + CANCEL_GRACE, + sqlx::query("SELECT pg_cancel_backend($1)") + .bind(backend_pid) + .execute(pool), + ) + .await; + match result { + Ok(Ok(_)) => {} + Ok(Err(error)) => { + debug!(error = %crate::error::Chain(&error), backend_pid, "River could not cancel a maintenance statement"); + } + Err(_) => { + debug!( + backend_pid, + "River timed out cancelling a maintenance statement" + ); + } + } +} diff --git a/rust/riverqueue/src/maintenance/reindexer.rs b/rust/riverqueue/src/maintenance/reindexer.rs new file mode 100644 index 000000000..e13a2a087 --- /dev/null +++ b/rust/riverqueue/src/maintenance/reindexer.rs @@ -0,0 +1,331 @@ +//! Periodically rebuilds River's hot indexes with `REINDEX CONCURRENTLY`, a +//! port of Go's `Reindexer`. + +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use sqlx::{AssertSqlSafe, PgConnection, PgPool}; +use tokio_util::sync::CancellationToken; +use tracing::{debug, error, info, warn}; + +use crate::{SchemaName, database::PostgresReindexSchedule}; + +use super::{ + MaintenanceError, + maintainer::ServiceContext, + postgres::{cancel_backend, cancellable_within}, + sleep_cancellable, +}; + +/// Time allowed to drop concurrent-build artifacts after a cancelled rebuild +/// (Go uses 15 seconds). +const ARTIFACT_CLEANUP_TIMEOUT: Duration = Duration::from_secs(15); + +/// Timeout for catalog queries that decide what to rebuild. +const CATALOG_TIMEOUT: Duration = Duration::from_secs(30); + +/// Returns the first scheduled run strictly after `after`. +pub(crate) fn next_run(schedule: PostgresReindexSchedule, after: DateTime) -> DateTime { + match schedule { + PostgresReindexSchedule::DailyUtc(time) => { + let today = after.date_naive().and_time(time).and_utc(); + if today > after { + today + } else { + after + .date_naive() + .succ_opt() + .expect("UTC date has a following day") + .and_time(time) + .and_utc() + } + } + PostgresReindexSchedule::Interval(interval) => { + after + + chrono::Duration::from_std(interval) + .expect("validated reindexer interval fits chrono duration") + } + } +} + +/// Runs the reindexer for one leadership term. +/// +/// Each term schedules from its own start time, so a client that becomes +/// leader after another leader already ran today's reindex does not run it +/// again immediately. Later runs advance from the previous scheduled time, not +/// from when a run finished. +pub(super) async fn run(context: std::sync::Arc, pool: sqlx::PgPool) { + let Some(config) = context.inner.database().postgres_reindex().cloned() else { + return; + }; + if config.index_names().is_empty() { + return; + } + let mut scheduled = next_run(config.schedule(), Utc::now()); + debug!(next_run_at = %scheduled, "River reindexer scheduled its first run"); + loop { + let wait = (scheduled - Utc::now()).to_std().unwrap_or_default(); + if !sleep_cancellable(&context.cancel, wait).await { + return; + } + match reindexable_index_names( + &pool, + &context.inner.schema, + &context.cancel, + config.index_names(), + ) + .await + { + Ok(index_names) => { + for index_name in index_names { + match reindex_one( + &pool, + &context.inner.schema, + &context.cancel, + &index_name, + config.timeout(), + ) + .await + { + Ok(true) => info!(index_name, "River reindexer rebuilt an index"), + Ok(false) => {} + Err(MaintenanceError::Cancelled) => return, + Err(reindex_error) => { + error!(error = %crate::error::Chain(&reindex_error), index_name, "River reindexer failed"); + } + } + } + } + Err(MaintenanceError::Cancelled) => return, + Err(list_error) => { + error!(error = %crate::error::Chain(&list_error), "River reindexer could not list indexes"); + } + } + scheduled = next_run(config.schedule(), scheduled); + } +} + +/// Returns configured indexes that exist in the River schema, warning about +/// missing ones (Go `reindexableIndexNames`). +async fn reindexable_index_names( + pool: &PgPool, + schema: &SchemaName, + cancel: &CancellationToken, + index_names: &[String], +) -> Result, MaintenanceError> { + let rows = tokio::select! { + biased; + () = cancel.cancelled() => return Err(MaintenanceError::Cancelled), + rows = tokio::time::timeout(CATALOG_TIMEOUT, sqlx::query_as::<_, (String, bool)>( + "SELECT index_name::text, EXISTS (\ + SELECT 1 FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = coalesce($2::text, current_schema()) \ + AND c.relname = index_name AND c.relkind = 'i'\ + ) FROM unnest($1::text[]) AS index_name", + ) + .bind(index_names) + .bind(schema.as_deref()) + .fetch_all(pool)) => rows.map_err(|_| MaintenanceError::TimedOut)??, + }; + let mut existing = Vec::with_capacity(rows.len()); + let mut missing = Vec::new(); + for (index_name, exists) in rows { + if exists { + existing.push(index_name); + } else { + missing.push(index_name); + } + } + if !missing.is_empty() { + warn!( + ?missing, + "River reindexer indexes do not exist; run migrations or update the reindexer configuration" + ); + } + Ok(existing) +} + +/// Lists leftovers of an interrupted `REINDEX CONCURRENTLY`: indexes named +/// like the target with a `_ccnew`/`_ccold` suffix and optional digits. +async fn reindex_artifacts( + connection: &mut PgConnection, + schema: &SchemaName, + index_name: &str, +) -> Result, sqlx::Error> { + sqlx::query_scalar( + "SELECT index_name FROM (\ + SELECT c.relname::text AS index_name, \ + substring(c.relname FROM length($1::text) + 1) AS suffix \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = coalesce($2::text, current_schema()) \ + AND c.relkind = 'i' AND left(c.relname, length($1::text)) = $1::text\ + ) AS index_artifacts \ + WHERE suffix ~ '^_cc(new|old)[0-9]*$' ORDER BY index_name", + ) + .bind(index_name) + .bind(schema.as_deref()) + .fetch_all(connection) + .await +} + +/// Rebuilds one index unless a previous attempt left artifacts behind. +/// +/// Like Go, an existing artifact means an earlier rebuild timed out, and +/// retrying would likely fail the same way, so the index is skipped with a +/// warning. When the term is cancelled mid-build, the running statement is +/// cancelled server-side and the fresh artifacts are dropped so that future +/// runs are not skipped forever. +async fn reindex_one( + pool: &PgPool, + schema: &SchemaName, + cancel: &CancellationToken, + index_name: &str, + timeout: Option, +) -> Result { + let mut connection = tokio::select! { + biased; + () = cancel.cancelled() => return Err(MaintenanceError::Cancelled), + connection = pool.acquire() => connection?, + }; + let artifacts = cancellable_catalog( + cancel, + reindex_artifacts(&mut connection, schema, index_name), + ) + .await?; + if !artifacts.is_empty() { + warn!( + index_name, + ?artifacts, + "River reindexer found artifacts of a previous partially completed rebuild; skipping" + ); + return Ok(false); + } + + // `REINDEX CONCURRENTLY` cannot run inside a transaction, so its timeout + // can't be a `SET LOCAL statement_timeout`. A session setting could + // outlive the rebuild on a pooled server connection behind a + // transaction-pooling proxy, so like Go, which bounds the rebuild with + // its context alone, the timeout is enforced client-side by cancelling + // the statement. + let backend_pid: i32 = cancellable_catalog( + cancel, + sqlx::query_scalar("SELECT pg_backend_pid()").fetch_one(&mut *connection), + ) + .await?; + let result = cancellable_within( + pool, + backend_pid, + cancel, + timeout, + sqlx::raw_sql(AssertSqlSafe(format!( + "REINDEX INDEX CONCURRENTLY {}", + schema.qualify(index_name) + ))) + .execute(&mut *connection), + ) + .await; + + if result.is_err() && cancel.is_cancelled() { + drop_artifacts(pool, &mut connection, schema, index_name, backend_pid).await; + } + result.map(|_| true) +} + +#[cfg(all(test, feature = "postgres", river_postgres_tests))] +pub(super) async fn reindex_one_for_test( + pool: &PgPool, + schema: &SchemaName, + cancel: &CancellationToken, + index_name: &str, +) -> Result { + reindex_one( + pool, + schema, + cancel, + index_name, + Some(Duration::from_mins(1)), + ) + .await +} + +/// Drops concurrent-build artifacts left by a cancelled rebuild, bounded so +/// shutdown cannot hang on it. +async fn drop_artifacts( + pool: &PgPool, + connection: &mut PgConnection, + schema: &SchemaName, + index_name: &str, + backend_pid: i32, +) { + info!( + index_name, + "River reindexer stopped mid-build; dropping concurrent artifacts" + ); + let cleanup = async { + let artifacts = reindex_artifacts(&mut *connection, schema, index_name).await?; + for artifact in artifacts { + if let Err(drop_error) = sqlx::raw_sql(AssertSqlSafe(format!( + "DROP INDEX CONCURRENTLY IF EXISTS {}", + schema.qualify(&artifact) + ))) + .execute(&mut *connection) + .await + { + error!(error = %crate::error::Chain(&drop_error), artifact, "River reindexer could not drop an artifact"); + } + } + Ok::<_, sqlx::Error>(()) + }; + match tokio::time::timeout(ARTIFACT_CLEANUP_TIMEOUT, cleanup).await { + Ok(Ok(())) => {} + Ok(Err(list_error)) => { + error!(error = %crate::error::Chain(&list_error), "River reindexer could not list artifacts"); + } + Err(_) => { + cancel_backend(pool, backend_pid).await; + error!(index_name, "River reindexer timed out dropping artifacts"); + } + } +} + +async fn cancellable_catalog( + cancel: &CancellationToken, + operation: impl Future>, +) -> Result { + tokio::select! { + biased; + () = cancel.cancelled() => Err(MaintenanceError::Cancelled), + result = tokio::time::timeout(CATALOG_TIMEOUT, operation) => match result { + Ok(result) => result.map_err(MaintenanceError::from), + Err(_) => Err(MaintenanceError::TimedOut), + }, + } +} + +#[cfg(test)] +mod unit_tests { + use chrono::{NaiveTime, TimeZone, Utc}; + + use super::{PostgresReindexSchedule, next_run}; + + #[test] + fn schedule_advances_from_the_previous_run() { + let before_midnight = Utc.with_ymd_and_hms(2026, 8, 11, 23, 59, 0).unwrap(); + let midnight = Utc.with_ymd_and_hms(2026, 8, 12, 0, 0, 0).unwrap(); + let daily = PostgresReindexSchedule::DailyUtc(NaiveTime::MIN); + assert_eq!(next_run(daily, before_midnight), midnight); + // A run exactly at midnight schedules the following midnight, like + // Go's `t.Add(24h).Truncate(24h)`. + assert_eq!( + next_run(daily, midnight), + Utc.with_ymd_and_hms(2026, 8, 13, 0, 0, 0).unwrap() + ); + let interval = PostgresReindexSchedule::Interval(std::time::Duration::from_secs(30)); + assert_eq!( + next_run(interval, before_midnight), + before_midnight + chrono::Duration::seconds(30) + ); + } +} diff --git a/rust/riverqueue/src/maintenance/rescuer.rs b/rust/riverqueue/src/maintenance/rescuer.rs new file mode 100644 index 000000000..9030f4c78 --- /dev/null +++ b/rust/riverqueue/src/maintenance/rescuer.rs @@ -0,0 +1,496 @@ +//! Rescues jobs stuck in `running`, a port of Go's `JobRescuer`. + +use chrono::{DateTime, Utc}; +#[cfg(feature = "postgres")] +use sqlx::AssertSqlSafe; +use tracing::{debug, error}; + +use crate::__private::{ + DatabaseConnection as PilotDatabaseConnection, RescueAction, RescueJob, RescueManyParams, + RescueParams, +}; + +#[cfg(feature = "postgres")] +use crate::client::{JobRecord, job_projection, tolerant_row}; +use crate::database::DatabasePool; +#[cfg(feature = "sqlite")] +use crate::database::sqlite; +use crate::{AttemptError, Error, JobRow, JobState, WorkerTimeout, client::ClientInner}; + +use super::{ + MaintenanceError, TIMEOUT_DEFAULT, batch_backoff, batch_size, maintainer::ServiceContext, + record_batch, +}; + +/// Error recorded on every rescued attempt, identical to Go's. +pub(crate) const RESCUE_ERROR: &str = "Stuck job rescued by JobRescuer"; + +/// Rescues stuck jobs in `id` order. +/// +/// Like Go, the stuck horizon is computed once per pass and paging continues +/// after the last selected ID, so a full batch of running jobs whose worker +/// timeout has not elapsed cannot livelock the rescuer. Updates only apply to +/// rows that are still `running` with `attempted_at` before the horizon, which +/// leaves jobs that completed or were claimed again after selection untouched. +pub(super) async fn run_once(context: &ServiceContext) -> Result<(), MaintenanceError> { + let stuck_horizon = Utc::now() + - chrono::Duration::from_std(context.inner.maintenance.effective_rescue_after()) + .map_err(|error| Error::configuration_context("maintenance", error.to_string()))?; + let mut after_id = 0_i64; + loop { + let limit = batch_size(&context.breakers.rescuer); + let result = rescue_batch(context, after_id, limit, stuck_horizon).await; + record_batch(&context.breakers.rescuer, &result); + let batch = result?; + if let Some(last_id) = batch.last_id { + after_id = last_id; + } + if i64::try_from(batch.selected).unwrap_or(i64::MAX) < limit { + return Ok(()); + } + batch_backoff(&context.cancel).await?; + } +} + +struct Batch { + last_id: Option, + selected: usize, +} + +async fn rescue_batch( + context: &ServiceContext, + after_id: i64, + limit: i64, + stuck_horizon: DateTime, +) -> Result { + match context.inner.database.pool() { + #[cfg(feature = "sqlite")] + DatabasePool::Sqlite(pool) => { + super::sqlite_cancellable( + &context.cancel, + TIMEOUT_DEFAULT, + rescue_batch_sqlite(&context.inner, pool, after_id, limit, stuck_horizon), + ) + .await + } + #[cfg(feature = "postgres")] + DatabasePool::Postgres(pool) => { + rescue_batch_postgres(context, pool, after_id, limit, stuck_horizon).await + } + } +} + +fn rescue_params( + inner: &ClientInner, + after_id: i64, + limit: i64, + stuck_horizon: DateTime, +) -> RescueParams { + RescueParams { + after_id, + database: inner.pilot_database_config(), + maximum: limit, + rescue_after: inner.maintenance.effective_rescue_after(), + stuck_horizon, + timeout: TIMEOUT_DEFAULT, + } +} + +/// Decides what OSS writes for each selected job, skipping jobs whose worker +/// timeout has not elapsed yet. +fn rescue_jobs( + inner: &ClientInner, + rows: &[JobRow], + now: DateTime, +) -> Result, Error> { + let mut jobs = Vec::with_capacity(rows.len()); + for row in rows { + let Some((state, finalized_at, scheduled_at)) = decide(inner, row, now)? else { + continue; + }; + let attempt_error = serde_json::to_value(AttemptError { + at: now, + attempt: row.attempt.max(0), + error: RESCUE_ERROR.to_owned(), + trace: String::new(), + })?; + jobs.push(RescueJob { + attempt_error, + finalized_at, + id: row.id, + scheduled_at, + state, + }); + } + Ok(jobs) +} + +type Decision = Option<(JobState, Option>, DateTime)>; + +/// Go's `makeRetryDecision`, preceded by its cancellation check. +fn decide(inner: &ClientInner, row: &JobRow, now: DateTime) -> Result { + if cancel_attempted(row.metadata.get_raw("cancel_attempted_at")) { + return Ok(Some((JobState::Cancelled, Some(now), row.scheduled_at))); + } + if !inner.workers.contains_kind(&row.kind) { + error!( + job_id = row.id, + job_kind = row.kind, + "River rescuer discarding a stuck job of an unhandled kind" + ); + return Ok(Some((JobState::Discarded, Some(now), row.scheduled_at))); + } + let retry_or_discard = |retry_at: DateTime| { + if row.attempt < row.max_attempts.max(0) { + (JobState::Retryable, None, retry_at) + } else { + (JobState::Discarded, Some(now), row.scheduled_at) + } + }; + // Retry policies get the error the rescue records as the attempt's + // error. As in Go, whose policy sees the job before the rescue appends + // that error, the row's `errors` don't include it yet. + let rescued_error = crate::WorkError::new(RESCUE_ERROR); + let client_retry = |row: &JobRow| -> Result, Error> { + let delay = inner.retry_policy.next_retry(row, &rescued_error, now); + Ok(now + + chrono::Duration::from_std(delay) + .map_err(|error| Error::invalid_job_context("maintenance", error.to_string()))?) + }; + + // A worker that cannot evaluate the job, for example because its args no + // longer decode, is retried with the client policy like Go's unmarshal + // failure path, without consulting the worker timeout. + let timeout = match inner.workers.timeout(row) { + Ok(WorkerTimeout::After(timeout)) => Some(timeout), + Ok(WorkerTimeout::ClientDefault) => inner.job_timeout, + Ok(WorkerTimeout::Disabled) => None, + Err(timeout_error) => { + debug!(error = %timeout_error, job_id = row.id, "River rescuer could not evaluate a stuck job"); + return Ok(Some(retry_or_discard(client_retry(row)?))); + } + }; + let Some(timeout) = timeout else { + // A disabled timeout means the job may legitimately run forever. + return Ok(None); + }; + let elapsed = row + .attempted_at + .and_then(|attempted_at| now.signed_duration_since(attempted_at).to_std().ok()) + .unwrap_or_default(); + if !timeout.is_zero() && elapsed < timeout { + return Ok(None); + } + + let retry_at = match inner.workers.next_retry(row, &rescued_error, now) { + Ok(Some(delay)) => { + now + chrono::Duration::from_std(delay) + .map_err(|error| Error::invalid_job_context("maintenance", error.to_string()))? + } + Ok(None) => client_retry(row)?, + Err(retry_error) => { + debug!(error = %retry_error, job_id = row.id, "River rescuer used the client retry policy"); + client_retry(row)? + } + }; + Ok(Some(retry_or_discard(retry_at))) +} + +/// Go decodes `cancel_attempted_at` as a `time.Time` and cancels only when it +/// is a non-zero timestamp; absent, null, or unparsable values do not cancel. +fn cancel_attempted(value: Option<&serde_json::value::RawValue>) -> bool { + let go_zero_time = + chrono::NaiveDate::from_ymd_opt(1, 1, 1).and_then(|date| date.and_hms_opt(0, 0, 0)); + value + .and_then(|raw| serde_json::from_str::(raw.get()).ok()) + .and_then(|value| DateTime::parse_from_rfc3339(&value).ok()) + .is_some_and(|time| Some(time.naive_utc()) != go_zero_time) +} + +#[cfg(feature = "postgres")] +#[expect( + clippy::too_many_lines, + reason = "selection, extension interception, and the guarded update share one transaction" +)] +async fn rescue_batch_postgres( + context: &ServiceContext, + pool: &sqlx::PgPool, + after_id: i64, + limit: i64, + stuck_horizon: DateTime, +) -> Result { + use super::postgres::{MaintenanceTransaction, cancellable}; + + let inner = &context.inner; + let table = inner.schema.qualify("river_job"); + let mut transaction = + MaintenanceTransaction::begin(pool, &context.cancel, TIMEOUT_DEFAULT).await?; + let backend_pid = transaction.backend_pid; + + let selected_ids = if inner.pilot.intercepts_rescue() { + inner + .pilot + .select_rescue_job_ids( + PilotDatabaseConnection::Postgres(&mut transaction.transaction), + &rescue_params(inner, after_id, limit, stuck_horizon), + ) + .await + .map_err(|source| Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "rescue selection", + }, + source, + })? + } else { + None + }; + // Like Go's `JobGetStuck`, selection takes no row locks; the guarded + // update below is what keeps a stale selection from rescuing a job. + let records = match selected_ids { + Some(ids) => { + cancellable( + pool, + backend_pid, + &context.cancel, + TIMEOUT_DEFAULT, + sqlx::query_as::<_, JobRecord>(AssertSqlSafe(format!( + "SELECT {}, false AS unique_skipped_as_duplicate FROM {table} AS job \ + WHERE id = ANY($1::bigint[]) ORDER BY id", + job_projection("job") + ))) + .bind(ids) + .fetch_all(&mut *transaction.transaction), + ) + .await? + } + None => { + cancellable( + pool, + backend_pid, + &context.cancel, + TIMEOUT_DEFAULT, + sqlx::query_as::<_, JobRecord>(AssertSqlSafe(format!( + "SELECT {}, false AS unique_skipped_as_duplicate FROM {table} AS job \ + WHERE state = 'running' AND id > $1 AND attempted_at < $2 \ + ORDER BY id LIMIT $3", + job_projection("job") + ))) + .bind(after_id) + .bind(stuck_horizon) + .bind(limit) + .fetch_all(&mut *transaction.transaction), + ) + .await? + } + }; + // Like River Go's `JobGetStuck`, a row that can't be fully decoded is + // still returned with its undecodable fields left empty, so a job + // stranded by such a row can be rescued. + let rows = records + .into_iter() + .filter_map(|record| tolerant_row(record.decode())) + .collect::>(); + let batch = Batch { + last_id: rows.last().map(|row| row.id), + selected: rows.len(), + }; + let jobs = rescue_jobs(inner, &rows, Utc::now())?; + if jobs.is_empty() { + transaction.commit(pool, &context.cancel).await?; + return Ok(batch); + } + + let params = RescueManyParams { + database: inner.pilot_database_config(), + jobs, + stuck_horizon, + }; + let action = if inner.pilot.intercepts_rescue() { + inner + .pilot + .rescue_jobs( + PilotDatabaseConnection::Postgres(&mut transaction.transaction), + ¶ms, + ) + .await + .map_err(|source| Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "rescue", + }, + source, + })? + } else { + RescueAction::Continue + }; + if action == RescueAction::Continue { + let state_type = inner.schema.qualify("river_job_state"); + let rescue_count = crate::METADATA_KEY_RESCUE_COUNT; + cancellable( + pool, + backend_pid, + &context.cancel, + TIMEOUT_DEFAULT, + sqlx::query(AssertSqlSafe(format!( + "UPDATE {table} AS job SET \ + errors = array_append(job.errors, updated_job.error), \ + finalized_at = updated_job.finalized_at, \ + scheduled_at = updated_job.scheduled_at, \ + metadata = job.metadata || jsonb_build_object('{rescue_count}', \ + coalesce(CASE WHEN jsonb_typeof(job.metadata -> '{rescue_count}') = 'number' \ + THEN (job.metadata ->> '{rescue_count}')::int END, 0) + 1), \ + state = updated_job.state \ + FROM (\ + SELECT unnest($1::bigint[]) AS id, unnest($2::jsonb[]) AS error, \ + unnest($3::timestamptz[]) AS finalized_at, \ + unnest($4::timestamptz[]) AS scheduled_at, \ + unnest($5::text[])::{state_type} AS state\ + ) AS updated_job \ + WHERE job.id = updated_job.id AND job.state = 'running' \ + AND job.attempted_at < $6" + ))) + .bind(params.jobs.iter().map(|job| job.id).collect::>()) + .bind( + params + .jobs + .iter() + .map(|job| job.attempt_error.clone()) + .collect::>(), + ) + .bind( + params + .jobs + .iter() + .map(|job| job.finalized_at) + .collect::>(), + ) + .bind( + params + .jobs + .iter() + .map(|job| job.scheduled_at) + .collect::>(), + ) + .bind( + params + .jobs + .iter() + .map(|job| job.state.as_str()) + .collect::>(), + ) + .bind(stuck_horizon) + .execute(&mut *transaction.transaction), + ) + .await?; + } + transaction.commit(pool, &context.cancel).await?; + Ok(batch) +} + +#[cfg(feature = "sqlite")] +async fn rescue_batch_sqlite( + inner: &ClientInner, + pool: &sqlx::SqlitePool, + after_id: i64, + limit: i64, + stuck_horizon: DateTime, +) -> Result { + let mut transaction = crate::database::begin_sqlite_write(pool).await?; + let selected_ids = if inner.pilot.intercepts_rescue() { + inner + .pilot + .select_rescue_job_ids( + PilotDatabaseConnection::Sqlite(&mut transaction), + &rescue_params(inner, after_id, limit, stuck_horizon), + ) + .await + .map_err(|source| Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "rescue selection", + }, + source, + })? + } else { + None + }; + let rows = match selected_ids { + Some(ids) => sqlite::jobs_by_ids(&mut transaction, &ids).await?, + None => { + sqlite::stuck_jobs( + &mut transaction, + after_id, + stuck_horizon, + i32::try_from(limit).unwrap_or(i32::MAX), + ) + .await? + } + }; + let batch = Batch { + last_id: rows.last().map(|row| row.id), + selected: rows.len(), + }; + let jobs = rescue_jobs(inner, &rows, Utc::now())?; + if !jobs.is_empty() { + let params = RescueManyParams { + database: inner.pilot_database_config(), + jobs, + stuck_horizon, + }; + let action = if inner.pilot.intercepts_rescue() { + inner + .pilot + .rescue_jobs(PilotDatabaseConnection::Sqlite(&mut transaction), ¶ms) + .await + .map_err(|source| Error::Extension { + phase: crate::ExtensionPhase::AddOn { + operation: "rescue", + }, + source, + })? + } else { + RescueAction::Continue + }; + if action == RescueAction::Continue { + for job in ¶ms.jobs { + let state = job.state; + let error: AttemptError = + serde_json::from_value(job.attempt_error.clone()).map_err(Error::from)?; + sqlite::rescue( + &mut transaction, + &sqlite::RescueJob { + error: &error, + finalized_at: job.finalized_at, + id: job.id, + scheduled_at: job.scheduled_at, + state, + stuck_horizon, + }, + ) + .await?; + } + } + } + transaction.commit().await?; + Ok(batch) +} + +#[cfg(test)] +mod unit_tests { + use serde_json::json; + + use super::cancel_attempted; + + #[test] + fn cancel_attempted_requires_a_non_zero_timestamp_like_go() { + let cancel = |value: Option| { + let raw = value.map(|value| serde_json::value::to_raw_value(&value).unwrap()); + cancel_attempted(raw.as_deref()) + }; + assert!(cancel(Some(json!("2026-01-02T03:04:05Z")))); + assert!(cancel(Some(json!("2026-01-02T03:04:05.123456+00:00")))); + assert!(!cancel(None)); + assert!(!cancel(Some(json!(null)))); + assert!(!cancel(Some(json!("0001-01-01T00:00:00Z")))); + assert!(!cancel(Some(json!("not a time")))); + assert!(!cancel(Some(json!(true)))); + } +} diff --git a/rust/riverqueue/src/maintenance/scheduler.rs b/rust/riverqueue/src/maintenance/scheduler.rs new file mode 100644 index 000000000..03dad9571 --- /dev/null +++ b/rust/riverqueue/src/maintenance/scheduler.rs @@ -0,0 +1,360 @@ +//! Moves due `scheduled` and `retryable` jobs to `available`, a port of Go's +//! `JobScheduler`. + +use std::{collections::BTreeSet, time::Duration}; + +use chrono::{DateTime, Utc}; +#[cfg(feature = "postgres")] +use sqlx::{AssertSqlSafe, Row}; + +#[cfg(feature = "sqlite")] +use crate::client::InsertNotifyLimiter; +use crate::database::DatabasePool; +#[cfg(feature = "postgres")] +use crate::database::postgres_capabilities::CapabilitiesCache; +#[cfg(feature = "sqlite")] +use crate::database::sqlite; + +use super::{ + MaintenanceError, TIMEOUT_DEFAULT, batch_backoff, batch_size, maintainer::ServiceContext, + record_batch, +}; + +/// Jobs due within this margin of the scheduling pass are announced to +/// producers; later look-ahead jobs are left to fetch polling (Go uses 5ms). +const NOTIFICATION_HORIZON: Duration = Duration::from_millis(5); + +/// A job the scheduler transitioned. +struct Scheduled { + queue: String, + scheduled_at: DateTime, +} + +/// Runs scheduling batches until a batch is smaller than the batch size. +/// +/// Like Go, jobs due within one scheduler interval are made available now, so +/// they can be fetched as soon as they are due instead of waiting for the next +/// pass. Only queues with jobs due by the end of the pass are notified. +pub(super) async fn run_once(context: &ServiceContext) -> Result<(), MaintenanceError> { + loop { + let limit = batch_size(&context.breakers.scheduler); + let result = schedule_batch(context, limit).await; + record_batch(&context.breakers.scheduler, &result); + if i64::try_from(result?).unwrap_or(i64::MAX) < limit { + return Ok(()); + } + batch_backoff(&context.cancel).await?; + } +} + +async fn schedule_batch(context: &ServiceContext, limit: i64) -> Result { + let now = Utc::now(); + let look_ahead = now + + chrono::Duration::from_std(context.inner.maintenance.scheduler_interval) + .unwrap_or_default(); + match context.inner.database.pool() { + #[cfg(feature = "sqlite")] + DatabasePool::Sqlite(pool) => { + super::sqlite_cancellable( + &context.cancel, + TIMEOUT_DEFAULT, + schedule_batch_sqlite( + pool, + &context.inner.insert_notify_limiter, + look_ahead, + limit, + ), + ) + .await + } + #[cfg(feature = "postgres")] + DatabasePool::Postgres(pool) => { + schedule_batch_postgres(context, pool, look_ahead, limit).await + } + } +} + +fn notified_queues(scheduled: &[Scheduled]) -> BTreeSet { + let horizon = Utc::now() + chrono::Duration::from_std(NOTIFICATION_HORIZON).unwrap_or_default(); + scheduled + .iter() + .filter(|job| job.scheduled_at <= horizon) + .map(|job| job.queue.clone()) + .collect() +} + +#[cfg(feature = "postgres")] +async fn schedule_batch_postgres( + context: &ServiceContext, + pool: &sqlx::PgPool, + look_ahead: DateTime, + limit: i64, +) -> Result { + use super::postgres::{MaintenanceTransaction, cancellable}; + + let inner = &context.inner; + let table = inner.schema.qualify("river_job"); + let state_function = inner.schema.qualify("river_job_state_in_bitmask"); + let state_type = inner.schema.qualify("river_job_state"); + // Mirrors Go's `JobSchedule`, including the index-friendly predicates and + // using the look-ahead time for both eligibility and conflict finalization. + let sql = format!( + "WITH jobs_to_schedule AS (\ + SELECT id, unique_key, unique_states, priority, scheduled_at FROM {table} \ + WHERE state IN ('retryable', 'scheduled') AND priority >= 0 AND queue IS NOT NULL \ + AND scheduled_at <= $2 \ + ORDER BY priority, scheduled_at, id LIMIT $1 FOR UPDATE\ + ), jobs_with_rownum AS (\ + SELECT *, CASE WHEN unique_key IS NOT NULL AND unique_states IS NOT NULL THEN \ + row_number() OVER (PARTITION BY unique_key ORDER BY priority, scheduled_at, id) END AS row_num \ + FROM jobs_to_schedule\ + ), unique_conflicts AS (\ + SELECT job.unique_key FROM {table} AS job JOIN jobs_with_rownum AS candidate \ + ON job.unique_key = candidate.unique_key AND job.id != candidate.id \ + WHERE job.unique_key IS NOT NULL AND job.unique_states IS NOT NULL \ + AND {state_function}(job.unique_states, job.state)\ + ), job_updates AS (\ + SELECT candidate.id, CASE \ + WHEN candidate.row_num IS NULL THEN 'available'::{state_type} \ + WHEN conflict.unique_key IS NOT NULL THEN 'discarded'::{state_type} \ + WHEN candidate.row_num = 1 THEN 'available'::{state_type} \ + ELSE 'discarded'::{state_type} END AS new_state \ + FROM jobs_with_rownum AS candidate LEFT JOIN unique_conflicts AS conflict \ + ON candidate.unique_key = conflict.unique_key\ + ), updated AS (\ + UPDATE {table} AS job SET state = job_updates.new_state, \ + finalized_at = CASE WHEN job_updates.new_state = 'discarded' THEN $2 ELSE job.finalized_at END, \ + metadata = CASE WHEN job_updates.new_state = 'discarded' \ + THEN job.metadata || '{{\"unique_key_conflict\": \"scheduler_discarded\"}}'::jsonb \ + ELSE job.metadata END \ + FROM job_updates WHERE job.id = job_updates.id \ + RETURNING job.queue, job.scheduled_at, job.state::text\ + ) SELECT queue, scheduled_at, state FROM updated" + ); + let mut transaction = + MaintenanceTransaction::begin(pool, &context.cancel, TIMEOUT_DEFAULT).await?; + let backend_pid = transaction.backend_pid; + let rows = cancellable( + pool, + backend_pid, + &context.cancel, + TIMEOUT_DEFAULT, + sqlx::query(AssertSqlSafe(sql)) + .bind(limit) + .bind(look_ahead) + .fetch_all(&mut *transaction.transaction), + ) + .await?; + let count = rows.len(); + let scheduled = rows + .iter() + .filter(|row| row.get::("state") == "available") + .map(|row| Scheduled { + queue: row.get("queue"), + scheduled_at: row.get("scheduled_at"), + }) + .collect::>(); + let notified = notified_queues(&scheduled); + let queues = context + .inner + .insert_notify_limiter + .due(notified.iter().map(String::as_str)); + if !queues.is_empty() && delivers_notifications(context, pool, &mut transaction).await? { + cancellable( + pool, + backend_pid, + &context.cancel, + TIMEOUT_DEFAULT, + sqlx::query( + "SELECT pg_notify(concat(coalesce($1::text, current_schema()), '.', $2::text), payload) \ + FROM unnest($3::text[]) AS payload", + ) + .bind(context.inner.schema.as_deref()) + .bind(crate::NOTIFICATION_TOPIC_INSERT) + .bind( + queues + .iter() + .map(|queue| crate::protocol::insert_notification_payload(queue)) + .collect::>(), + ) + .execute(&mut *transaction.transaction), + ) + .await?; + } + transaction.commit(pool, &context.cancel).await?; + Ok(count) +} + +/// Whether the server delivers notifications, detected on the batch's +/// transaction. A server without `LISTEN`/`NOTIFY` gets none. +#[cfg(feature = "postgres")] +async fn delivers_notifications( + context: &ServiceContext, + pool: &sqlx::PgPool, + transaction: &mut super::postgres::MaintenanceTransaction, +) -> Result { + Ok(super::postgres::cancellable( + pool, + transaction.backend_pid, + &context.cancel, + TIMEOUT_DEFAULT, + CapabilitiesCache::load_or_detect( + context.inner.database.postgres_capabilities(), + &mut *transaction.transaction, + ), + ) + .await? + .supports_listen_notify) +} + +#[cfg(feature = "sqlite")] +async fn schedule_batch_sqlite( + pool: &sqlx::SqlitePool, + notify_limiter: &InsertNotifyLimiter, + look_ahead: DateTime, + limit: i64, +) -> Result { + let mut transaction = crate::database::begin_sqlite_write(pool).await?; + let candidates = sqlite::schedule_candidates( + &mut transaction, + look_ahead, + i32::try_from(limit).unwrap_or(i32::MAX), + ) + .await?; + let count = candidates.len(); + let mut available_ids = Vec::new(); + let mut conflict_ids = Vec::new(); + let mut scheduled = Vec::new(); + for candidate in candidates { + let Some(unique_key) = candidate.unique_key.as_deref() else { + available_ids.push(candidate.id); + continue; + }; + if sqlite::schedule_has_unique_collision(&mut transaction, candidate.id, unique_key).await? + { + conflict_ids.push(candidate.id); + } else { + // Transition unique jobs one at a time so that a later duplicate + // in the same batch observes the earlier one as a collision. + let available = + sqlite::schedule_set_available(&mut transaction, &[candidate.id]).await?; + scheduled.extend(available.into_iter().map(|job| Scheduled { + queue: job.queue, + scheduled_at: job.scheduled_at, + })); + } + } + if !available_ids.is_empty() { + let available = sqlite::schedule_set_available(&mut transaction, &available_ids).await?; + scheduled.extend(available.into_iter().map(|job| Scheduled { + queue: job.queue, + scheduled_at: job.scheduled_at, + })); + } + if !conflict_ids.is_empty() { + sqlite::schedule_discard_conflicts(&mut transaction, &conflict_ids, look_ahead).await?; + } + let notified = notified_queues(&scheduled); + for queue in notify_limiter.due(notified.iter().map(String::as_str)) { + let payload = crate::protocol::insert_notification_payload(queue); + sqlite::notification_insert( + &mut transaction, + &[sqlite::NotificationInput { + payload: &payload, + topic: crate::NOTIFICATION_TOPIC_INSERT, + }], + ) + .await?; + } + transaction.commit().await?; + Ok(count) +} + +#[cfg(all(test, feature = "sqlite"))] +mod sqlite_tests { + use std::{sync::Arc, time::Duration}; + + use chrono::Utc; + use riverqueue_migrate::SqliteMigrator; + use serde::{Deserialize, Serialize}; + use sqlx::{SqlitePool, sqlite::SqlitePoolOptions}; + use tokio_util::sync::CancellationToken; + + use super::super::{BatchSizes, Breakers, maintainer::ServiceContext}; + use crate::{Client, InsertOpts, JobArgs, database::sqlite::sqlite_time}; + + #[derive(Debug, Deserialize, JobArgs, Serialize)] + #[river(kind = "scheduler_notification")] + struct NotificationArgs {} + + async fn insert_due_job(pool: &SqlitePool, queue: &str) { + sqlx::query( + "INSERT INTO river_job (args, kind, max_attempts, metadata, queue, scheduled_at, state) \ + VALUES (jsonb('{}'), 'scheduler_notification', 25, jsonb('{}'), ?, ?, 'scheduled')", + ) + .bind(queue) + .bind(sqlite_time(Utc::now() - chrono::Duration::hours(1))) + .execute(pool) + .await + .unwrap(); + } + + async fn insert_notifications(pool: &SqlitePool, queue: &str) -> i64 { + sqlx::query_scalar( + "SELECT count(*) FROM river_notification WHERE topic = 'river_insert' \ + AND json_extract(payload, '$.queue') = ?", + ) + .bind(queue) + .fetch_one(pool) + .await + .unwrap() + } + + /// Like Go, the scheduler notifies through the client's insert + /// notification limiter, which insertions share. + #[tokio::test] + async fn scheduler_notifications_wait_for_the_fetch_cooldown() { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + SqliteMigrator::new(pool.clone()) + .migrate_up() + .await + .unwrap(); + let client = Client::builder(pool.clone()) + .fetch_cooldown(Duration::from_hours(1)) + .build() + .unwrap(); + let context = ServiceContext { + breakers: Arc::new(Breakers::new(BatchSizes::default())), + cancel: CancellationToken::new(), + inner: Arc::clone(&client.inner), + }; + + insert_due_job(&pool, "scheduled").await; + super::run_once(&context).await.unwrap(); + assert_eq!(insert_notifications(&pool, "scheduled").await, 1); + + insert_due_job(&pool, "scheduled").await; + super::run_once(&context).await.unwrap(); + assert_eq!(insert_notifications(&pool, "scheduled").await, 1); + + client + .insert(NotificationArgs {}) + .opts(InsertOpts::default().with_queue("inserted")) + .await + .unwrap(); + insert_due_job(&pool, "inserted").await; + super::run_once(&context).await.unwrap(); + assert_eq!(insert_notifications(&pool, "inserted").await, 1); + + let available: i64 = + sqlx::query_scalar("SELECT count(*) FROM river_job WHERE state = 'available'") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(available, 4); + } +} diff --git a/rust/riverqueue/src/maintenance/tests.rs b/rust/riverqueue/src/maintenance/tests.rs new file mode 100644 index 000000000..9f7993307 --- /dev/null +++ b/rust/riverqueue/src/maintenance/tests.rs @@ -0,0 +1,1095 @@ +//! PostgreSQL tests of individual maintenance services, ported from Go's +//! `internal/maintenance` and `internal/leadership` suites. Each test uses its +//! own freshly migrated schema and fails when `RIVER_RUST_DATABASE_URL` is +//! unset. + +use std::{ + convert::Infallible, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use sqlx::{AssertSqlSafe, PgPool, postgres::PgPoolOptions}; +use tokio::sync::mpsc; +use tokio_util::sync::CancellationToken; + +use crate::__private::{ + DatabaseConnection, JobUpdatedParams, Pilot, PilotError, RescueAction, RescueManyParams, +}; +use riverqueue_migrate::PostgresMigrator; + +use super::{ + BatchSizes, Breakers, cleaner, + elector::{DatabaseLeaderStore, Elector, ElectorEvent}, + maintainer::ServiceContext, + rescuer, scheduler, +}; +use crate::{ + Client, Job, JobArgs, JobState, MaintenanceConfig, QueueConfig, SchemaName, UniqueOpts, + WorkContext, WorkOutcome, Worker, WorkerRegistry, WorkerTimeout, + database::{PostgresDatabase, PostgresReindexConfig}, +}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "maintenance_no_timeout")] +struct NoTimeoutArgs {} + +struct NoTimeoutWorker; + +impl Worker for NoTimeoutWorker { + type Error = Infallible; + + fn work( + &self, + _context: WorkContext, + _job: Job, + ) -> impl Future> + Send { + std::future::ready(Ok(WorkOutcome::Complete)) + } + + fn timeout(&self, _job: &Job) -> WorkerTimeout { + WorkerTimeout::Disabled + } +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "maintenance_short_timeout")] +struct ShortTimeoutArgs {} + +struct ShortTimeoutWorker; + +impl Worker for ShortTimeoutWorker { + type Error = Infallible; + + fn work( + &self, + _context: WorkContext, + _job: Job, + ) -> impl Future> + Send { + std::future::ready(Ok(WorkOutcome::Complete)) + } + + fn timeout(&self, _job: &Job) -> WorkerTimeout { + WorkerTimeout::After(Duration::from_millis(1)) + } +} + +fn workers() -> WorkerRegistry { + let mut workers = WorkerRegistry::new(); + workers + .register::(NoTimeoutWorker) + .unwrap(); + workers + .register::(ShortTimeoutWorker) + .unwrap(); + workers +} + +/// A migrated schema owned by one test. +struct TestDatabase { + pool: PgPool, + schema: SchemaName, + name: String, +} + +impl TestDatabase { + async fn new(prefix: &str) -> Self { + static NONCE: AtomicUsize = AtomicUsize::new(0); + let url = std::env::var("RIVER_RUST_DATABASE_URL") + .expect("RIVER_RUST_DATABASE_URL must point at a disposable test database"); + let pool = PgPoolOptions::new() + .max_connections(16) + .connect(&url) + .await + .expect("connect to RIVER_RUST_DATABASE_URL"); + let mut name = format!( + "{prefix}_{:x}_{:x}_{:x}", + std::process::id(), + NONCE.fetch_add(1, Ordering::Relaxed), + Utc::now().timestamp_subsec_nanos() + ); + name.truncate(riverqueue_migrate::SCHEMA_MAX_LEN); + sqlx::raw_sql(AssertSqlSafe(format!("CREATE SCHEMA \"{name}\""))) + .execute(&pool) + .await + .unwrap(); + let schema = SchemaName::new(name.clone()).unwrap(); + PostgresMigrator::new(pool.clone()) + .with_schema(schema.clone()) + .migrate_up() + .await + .unwrap(); + Self { pool, schema, name } + } + + fn table(&self, table: &str) -> String { + self.schema.qualify(table) + } + + fn client(&self) -> crate::ClientBuilder { + Client::builder( + PostgresDatabase::new(self.pool.clone()) + .with_schema(self.schema.clone()) + .with_reindex(PostgresReindexConfig::default().with_index_names([] as [&str; 0])), + ) + .workers(workers()) + } + + /// Inserts a raw job row and returns its ID. + async fn insert_job(&self, job: RawJob<'_>) -> i64 { + sqlx::query_scalar(AssertSqlSafe(format!( + "INSERT INTO {} (args, attempt, attempted_at, finalized_at, kind, max_attempts, \ + metadata, queue, scheduled_at, state, unique_key, unique_states) \ + VALUES ('{{}}', $1, $2, $3, $4, $5, $6, $7, coalesce($8, now()), $9::text::{}, $10, \ + CASE WHEN $10 IS NULL THEN NULL ELSE $11::int::bit(8) END) RETURNING id", + self.table("river_job"), + self.schema.qualify("river_job_state"), + ))) + .bind(job.attempt) + .bind(job.attempted_at) + .bind(job.finalized_at) + .bind(job.kind) + .bind(job.max_attempts) + .bind(sqlx::types::Json(job.metadata)) + .bind(job.queue) + .bind(job.scheduled_at) + .bind(job.state) + .bind(job.unique_key) + .bind(job.unique_states) + .fetch_one(&self.pool) + .await + .unwrap() + } + + async fn job(&self, id: i64) -> Option<(String, i32, serde_json::Value)> { + sqlx::query_as(AssertSqlSafe(format!( + "SELECT state::text, coalesce(array_length(errors, 1), 0), metadata FROM {} WHERE id = $1", + self.table("river_job") + ))) + .bind(id) + .fetch_optional(&self.pool) + .await + .unwrap() + } + + async fn state(&self, id: i64) -> Option { + self.job(id).await.map(|(state, _, _)| state) + } + + async fn cleanup(self) { + sqlx::raw_sql(AssertSqlSafe(format!( + "DROP SCHEMA \"{}\" CASCADE", + self.name + ))) + .execute(&self.pool) + .await + .unwrap(); + self.pool.close().await; + } +} + +struct RawJob<'a> { + attempt: i16, + attempted_at: Option>, + finalized_at: Option>, + kind: &'a str, + max_attempts: i16, + metadata: serde_json::Value, + queue: &'a str, + scheduled_at: Option>, + state: &'a str, + unique_key: Option>, + /// Unique-state bitmask; defaults to every state when a key is set. + unique_states: i32, +} + +impl Default for RawJob<'_> { + fn default() -> Self { + Self { + attempt: 0, + attempted_at: None, + finalized_at: None, + kind: NoTimeoutArgs::KIND, + max_attempts: 25, + metadata: serde_json::json!({}), + queue: "default", + scheduled_at: None, + state: "available", + unique_key: None, + unique_states: 0xFF, + } + } +} + +fn hours_ago(hours: i64) -> DateTime { + Utc::now() - chrono::Duration::hours(hours) +} + +fn stuck(kind: &str) -> RawJob<'_> { + RawJob { + attempt: 1, + attempted_at: Some(hours_ago(3)), + kind, + state: "running", + ..RawJob::default() + } +} + +fn context(client: &Client, batch: i64) -> ServiceContext { + ServiceContext { + breakers: Arc::new(Breakers::new(BatchSizes { + default: batch, + reduced: 1, + })), + cancel: CancellationToken::new(), + inner: Arc::clone(&client.inner), + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn rescuer_rescues_past_full_batch_of_jobs_with_no_timeout() { + let database = TestDatabase::new("rmt_rescue_batch").await; + let client = database.client().build().unwrap(); + + // A full batch of stuck jobs whose timeout is disabled must not stop the + // rescuer from reaching later eligible jobs (Go + // `RescuesPastFullBatchOfJobsWithNoTimeout`). + let mut ignored = Vec::new(); + for _ in 0..6 { + ignored.push(database.insert_job(stuck(NoTimeoutArgs::KIND)).await); + } + let eligible = database.insert_job(stuck(ShortTimeoutArgs::KIND)).await; + + tokio::time::timeout( + Duration::from_secs(20), + rescuer::run_once(&context(&client, 3)), + ) + .await + .expect("rescuer must not livelock on a full batch of ignored jobs") + .unwrap(); + + assert_eq!(database.state(eligible).await.as_deref(), Some("retryable")); + for id in ignored { + let (state, errors, metadata) = database.job(id).await.unwrap(); + assert_eq!(state, "running"); + assert_eq!(errors, 0); + assert!(metadata.get("river:rescue_count").is_none()); + } + database.cleanup().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn rescuer_rescues_undecodable_stuck_jobs() { + let database = TestDatabase::new("rmt_rescue_undecodable").await; + let client = database.client().build().unwrap(); + + // Like River Go's `JobGetStuck`, a stuck job whose row can't be fully + // decoded is still read, so neither it nor the jobs read with it are + // stranded. + let undecodable = database + .insert_job(RawJob { + metadata: serde_json::json!([1]), + ..stuck(ShortTimeoutArgs::KIND) + }) + .await; + let decodable = database.insert_job(stuck(ShortTimeoutArgs::KIND)).await; + + rescuer::run_once(&context(&client, 100)).await.unwrap(); + + for id in [undecodable, decodable] { + let (state, errors, _) = database.job(id).await.unwrap(); + assert_eq!((state.as_str(), errors), ("retryable", 1), "job {id}"); + } + database.cleanup().await; +} + +/// Mutates selected jobs inside the rescue transaction to simulate workers +/// that finish or re-claim jobs after selection, then lets OSS continue. +struct StaleSnapshotPilot { + completed: Arc>>, + handled: bool, + reclaimed: Arc>>, + schema: SchemaName, +} + +#[async_trait] +impl Pilot for StaleSnapshotPilot { + fn intercepts_rescue(&self) -> bool { + true + } + + async fn rescue_jobs( + &self, + connection: DatabaseConnection<'_>, + params: &RescueManyParams, + ) -> Result { + let connection = connection.into_postgres().expect("PostgreSQL connection"); + let completed = self.completed.lock().unwrap().clone(); + let reclaimed = self.reclaimed.lock().unwrap().clone(); + assert!(params.jobs.iter().any(|job| completed.contains(&job.id))); + let table = self.schema.qualify("river_job"); + sqlx::query(AssertSqlSafe(format!( + "UPDATE {table} SET state = 'completed', finalized_at = now() WHERE id = ANY($1)" + ))) + .bind(&completed) + .execute(&mut *connection) + .await?; + sqlx::query(AssertSqlSafe(format!( + "UPDATE {table} SET attempt = attempt + 1, attempted_at = now() WHERE id = ANY($1)" + ))) + .bind(&reclaimed) + .execute(&mut *connection) + .await?; + Ok(if self.handled { + RescueAction::Handled + } else { + RescueAction::Continue + }) + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn rescuer_update_is_guarded_against_stale_selection() { + for handled in [false, true] { + let database = TestDatabase::new("rmt_rescue_stale").await; + let completed = database.insert_job(stuck(ShortTimeoutArgs::KIND)).await; + let reclaimed = database.insert_job(stuck(ShortTimeoutArgs::KIND)).await; + let eligible = database.insert_job(stuck(ShortTimeoutArgs::KIND)).await; + let client = database + .client() + .with_pilot(StaleSnapshotPilot { + completed: Arc::new(std::sync::Mutex::new(vec![completed])), + handled, + reclaimed: Arc::new(std::sync::Mutex::new(vec![reclaimed])), + schema: database.schema.clone(), + }) + .build() + .unwrap(); + + rescuer::run_once(&context(&client, 100)).await.unwrap(); + + // Jobs completed or claimed again after selection keep their new + // state, errors, and metadata. + let (state, errors, metadata) = database.job(completed).await.unwrap(); + assert_eq!((state.as_str(), errors), ("completed", 0)); + assert!(metadata.get("river:rescue_count").is_none()); + let (state, errors, _) = database.job(reclaimed).await.unwrap(); + assert_eq!((state.as_str(), errors), ("running", 0)); + // An extension that handles the rescue suppresses the OSS update. + let expected = if handled { "running" } else { "retryable" }; + assert_eq!(database.state(eligible).await.as_deref(), Some(expected)); + database.cleanup().await; + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn scheduler_look_ahead_and_unique_conflicts() { + let database = TestDatabase::new("rmt_scheduler").await; + let client = database + .client() + .maintenance(MaintenanceConfig::default().with_scheduler_interval(Duration::from_secs(5))) + .build() + .unwrap(); + + let due = database + .insert_job(RawJob { + scheduled_at: Some(hours_ago(1)), + state: "scheduled", + ..RawJob::default() + }) + .await; + let look_ahead = database + .insert_job(RawJob { + scheduled_at: Some(Utc::now() + chrono::Duration::seconds(2)), + state: "retryable", + ..RawJob::default() + }) + .await; + let later = database + .insert_job(RawJob { + scheduled_at: Some(Utc::now() + chrono::Duration::minutes(5)), + state: "scheduled", + ..RawJob::default() + }) + .await; + // Retryable jobs whose unique key is held by another job in a unique + // state are discarded instead of made available (Go + // `MovesUniqueKeyConflictingJobsToDiscarded`). Like Go's test, unique + // states are the defaults without `retryable` (available, completed, + // pending, running, scheduled) so the duplicates can exist while waiting. + let without_retryable = 0b1101_0101; + let mut non_conflicting = Vec::new(); + for key in [1_u8, 2] { + non_conflicting.push( + database + .insert_job(RawJob { + scheduled_at: Some(hours_ago(1)), + state: "retryable", + unique_key: Some(vec![key; 32]), + unique_states: without_retryable, + ..RawJob::default() + }) + .await, + ); + } + let mut conflicting = Vec::new(); + for (key, holder_state) in [ + (3_u8, "available"), + (4, "completed"), + (5, "pending"), + (6, "running"), + (7, "scheduled"), + ] { + conflicting.push( + database + .insert_job(RawJob { + scheduled_at: Some(hours_ago(1)), + state: "retryable", + unique_key: Some(vec![key; 32]), + unique_states: without_retryable, + ..RawJob::default() + }) + .await, + ); + database + .insert_job(RawJob { + attempt: i16::from(holder_state == "running"), + attempted_at: (holder_state == "running").then(Utc::now), + finalized_at: (holder_state == "completed").then(Utc::now), + scheduled_at: (holder_state == "scheduled") + .then(|| Utc::now() + chrono::Duration::hours(1)), + state: holder_state, + unique_key: Some(vec![key; 32]), + unique_states: without_retryable, + ..RawJob::default() + }) + .await; + } + + scheduler::run_once(&context(&client, 100)).await.unwrap(); + + assert_eq!(database.state(due).await.as_deref(), Some("available")); + assert_eq!( + database.state(look_ahead).await.as_deref(), + Some("available") + ); + assert_eq!(database.state(later).await.as_deref(), Some("scheduled")); + for id in non_conflicting { + assert_eq!(database.state(id).await.as_deref(), Some("available")); + } + for id in conflicting { + let (state, _, metadata) = database.job(id).await.unwrap(); + assert_eq!(state, "discarded"); + assert_eq!(metadata["unique_key_conflict"], "scheduler_discarded"); + } + database.cleanup().await; +} + +struct ExcludingPilot; + +impl Pilot for ExcludingPilot { + fn job_cleaner_queue_exclusions(&self) -> Vec { + vec!["extension_owned".to_owned()] + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn job_cleaner_retention_exclusions_and_batches() { + let retentions = [ + (None, None, None), + (None, Some(1), Some(1)), + (Some(1), None, Some(1)), + (Some(1), Some(1), None), + ]; + for (cancelled, completed, discarded) in retentions { + let database = TestDatabase::new("rmt_job_cleaner").await; + let hours = |retention: Option| { + retention.map_or(crate::Retention::Keep, |hours| { + crate::Retention::DeleteAfter(Duration::from_hours(hours)) + }) + }; + let client = database + .client() + .with_pilot(ExcludingPilot) + .maintenance( + MaintenanceConfig::default() + .with_cancelled_job_retention(hours(cancelled)) + .with_completed_job_retention(hours(completed)) + .with_discarded_job_retention(hours(discarded)), + ) + .build() + .unwrap(); + let mut expired = Vec::new(); + for state in ["cancelled", "completed", "discarded"] { + for _ in 0..3 { + expired.push(( + state, + database + .insert_job(RawJob { + finalized_at: Some(hours_ago(2)), + state, + ..RawJob::default() + }) + .await, + )); + } + } + let recent = database + .insert_job(RawJob { + finalized_at: Some(Utc::now()), + state: "completed", + ..RawJob::default() + }) + .await; + let excluded = database + .insert_job(RawJob { + finalized_at: Some(hours_ago(2)), + queue: "extension_owned", + state: "completed", + ..RawJob::default() + }) + .await; + let running = database + .insert_job(RawJob { + attempt: 1, + attempted_at: Some(hours_ago(2)), + state: "running", + ..RawJob::default() + }) + .await; + + // A batch size of two forces several batches. + cleaner::clean_jobs(&context(&client, 2)).await.unwrap(); + + for (state, id) in expired { + let retention = match state { + "cancelled" => cancelled, + "completed" => completed, + _ => discarded, + }; + assert_eq!( + database.state(id).await.is_none(), + retention.is_some(), + "{state} job with retention {retention:?}" + ); + } + assert!(database.state(recent).await.is_some()); + assert!(database.state(excluded).await.is_some()); + assert!(database.state(running).await.is_some()); + database.cleanup().await; + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn queue_cleaner_keeps_active_queues() { + let database = TestDatabase::new("rmt_queue_cleaner").await; + let client = database.client().build().unwrap(); + for (name, age_hours) in [ + ("stale_a", 25), + ("stale_b", 30), + ("stale_c", 48), + ("active", 0), + ] { + sqlx::query(AssertSqlSafe(format!( + "INSERT INTO {} (name, created_at, metadata, updated_at) \ + VALUES ($1, now(), '{{}}', now() - make_interval(hours => $2))", + database.table("river_queue") + ))) + .bind(name) + .bind(age_hours) + .execute(&database.pool) + .await + .unwrap(); + } + + cleaner::clean_queues(&context(&client, 2)).await.unwrap(); + + let remaining: Vec = sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT name FROM {} ORDER BY name", + database.table("river_queue") + ))) + .fetch_all(&database.pool) + .await + .unwrap(); + assert_eq!(remaining, ["active"]); + database.cleanup().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn elector_loses_leadership_when_same_id_term_is_replaced() { + let database = TestDatabase::new("rmt_term_replaced").await; + let client = database + .client() + .id("shared-leader-id") + .maintenance(MaintenanceConfig::default().with_elect_interval(Duration::from_millis(25))) + .build() + .unwrap(); + let (events_sender, mut events) = mpsc::unbounded_channel(); + let (_wakeups_sender, wakeups) = mpsc::unbounded_channel(); + let (terms_sender, mut terms) = mpsc::unbounded_channel(); + let cancel = CancellationToken::new(); + let elector = Elector::new( + Arc::new(DatabaseLeaderStore::new(Arc::clone(&client.inner))), + client.id().to_owned(), + Duration::from_millis(25), + ) + .with_events(events_sender); + let run = tokio::spawn(elector.run(cancel.clone(), wakeups, terms_sender)); + let term = terms.recv().await.unwrap(); + + // Another instance with the same ID replaces the term: same leader ID, + // newer `elected_at`. + let table = database.table("river_leader"); + let replaced_elected_at: DateTime = sqlx::query_scalar(AssertSqlSafe(format!( + "WITH removed AS (DELETE FROM {table} RETURNING leader_id, elected_at) \ + INSERT INTO {table} (leader_id, elected_at, expires_at) \ + SELECT leader_id, elected_at + interval '1 second', now() + interval '1 hour' FROM removed \ + RETURNING elected_at" + ))) + .fetch_one(&database.pool) + .await + .unwrap(); + assert!(replaced_elected_at > term.elected_at); + + tokio::time::timeout(Duration::from_secs(10), term.token.cancelled()) + .await + .expect("the replaced term must be given up"); + let mut observed = Vec::new(); + tokio::time::timeout(Duration::from_secs(10), async { + while let Some(event) = events.recv().await { + observed.push(event); + if event == ElectorEvent::Denied { + break; + } + } + }) + .await + .unwrap(); + assert!(observed.contains(&ElectorEvent::Lost)); + assert!(!observed.contains(&ElectorEvent::Resigned)); + + // The replacement term was neither renewed nor deleted by this client. + let (elected_at, expires_in_minutes): (DateTime, f64) = + sqlx::query_as(AssertSqlSafe(format!( + "SELECT elected_at, (extract(epoch FROM expires_at - now()) / 60)::float8 FROM {table}" + ))) + .fetch_one(&database.pool) + .await + .unwrap(); + assert_eq!(elected_at, replaced_elected_at); + assert!(expires_in_minutes > 50.0); + + cancel.cancel(); + run.await.unwrap(); + database.cleanup().await; +} + +#[derive(Clone, Default)] +struct HookPilot { + cancels: Arc>>, + retries: Arc>>, + fail: bool, +} + +#[async_trait] +impl Pilot for HookPilot { + fn intercepts_job_cancel_retry(&self) -> bool { + true + } + + async fn after_job_cancel( + &self, + _connection: DatabaseConnection<'_>, + job: &JobUpdatedParams, + ) -> Result<(), PilotError> { + self.cancels + .lock() + .unwrap() + .push((job.job.id, job.job.state.as_str().to_owned())); + if self.fail { + return Err(std::io::Error::other("cancel hook failed").into()); + } + Ok(()) + } + + async fn after_job_retry( + &self, + _connection: DatabaseConnection<'_>, + job: &JobUpdatedParams, + ) -> Result<(), PilotError> { + self.retries + .lock() + .unwrap() + .push((job.job.id, job.job.state.as_str().to_owned())); + if self.fail { + return Err(std::io::Error::other("retry hook failed").into()); + } + Ok(()) + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn cancel_and_retry_post_hooks_share_the_transaction() { + let database = TestDatabase::new("rmt_cancel_retry").await; + let pilot = HookPilot::default(); + let client = database.client().with_pilot(pilot.clone()).build().unwrap(); + let id = database.insert_job(RawJob::default()).await; + + let cancelled = client.jobs().cancel(id).await.unwrap(); + assert_eq!(cancelled.state, JobState::Cancelled); + assert_eq!( + *pilot.cancels.lock().unwrap(), + [(id, "cancelled".to_owned())] + ); + let retried = client.jobs().retry(id).await.unwrap(); + assert_eq!(retried.state, JobState::Available); + assert_eq!( + *pilot.retries.lock().unwrap(), + [(id, "available".to_owned())] + ); + + // A failing hook rolls back the operation it follows. + let failing = database + .client() + .with_pilot(HookPilot { + fail: true, + ..HookPilot::default() + }) + .build() + .unwrap(); + assert!(failing.jobs().cancel(id).await.is_err()); + assert_eq!(database.state(id).await.as_deref(), Some("available")); + let mut transaction = database.pool.begin().await.unwrap(); + assert!(client.jobs().cancel(id).tx(&mut transaction).await.is_ok()); + transaction.rollback().await.unwrap(); + assert_eq!(database.state(id).await.as_deref(), Some("available")); + database.cleanup().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn periodic_jobs_use_expected_run_time_and_uniqueness() { + use crate::{InsertOpts, IntervalSchedule, PeriodicJob}; + + let database = TestDatabase::new("rmt_periodic").await; + let client = database.client().build().unwrap(); + let periodic = client.periodic_jobs(); + periodic + .add(PeriodicJob::conditional( + IntervalSchedule::new(Duration::from_mins(1)).unwrap(), + || { + Some(( + NoTimeoutArgs {}, + InsertOpts::default().with_unique(UniqueOpts::new().with_by_args(true)), + )) + }, + )) + .unwrap(); + + let start = Utc::now(); + periodic.reset_for_leadership(); + periodic + .run_due(&client, start, &CancellationToken::new()) + .await; + let target = periodic.next_run_at().unwrap(); + // Running a little before the target still inserts it, scheduled at the + // expected run time rather than when the enqueuer woke up (Go + // `SetsScheduledAtAccordingToExpectedNextRunAt`). + periodic + .run_due( + &client, + target - chrono::Duration::milliseconds(50), + &CancellationToken::new(), + ) + .await; + let rows: Vec> = sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT scheduled_at FROM {}", + database.table("river_job") + ))) + .fetch_all(&database.pool) + .await + .unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!( + rows[0].timestamp_micros(), + target.timestamp_micros(), + "periodic job scheduled at its expected run time" + ); + + // The next occurrence is unique by args with the first and is skipped + // (Go `RespectsJobUniqueness`). + let next = periodic.next_run_at().unwrap(); + assert!(next > target); + periodic + .run_due(&client, next, &CancellationToken::new()) + .await; + let count: i64 = sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT count(*) FROM {}", + database.table("river_job") + ))) + .fetch_one(&database.pool) + .await + .unwrap(); + assert_eq!(count, 1); + database.cleanup().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn reindexer_skips_artifacts_and_drops_artifacts_when_cancelled() { + let database = TestDatabase::new("rmt_reindexer").await; + let table = database.table("river_job"); + for index in ["maint_reindex_artifact_idx", "maint_reindex_cancel_idx"] { + sqlx::raw_sql(AssertSqlSafe(format!( + "CREATE INDEX \"{index}\" ON {table} (kind)" + ))) + .execute(&database.pool) + .await + .unwrap(); + } + // A leftover of an earlier interrupted rebuild causes a skip. + sqlx::raw_sql(AssertSqlSafe(format!( + "CREATE INDEX \"maint_reindex_artifact_idx_ccnew1\" ON {table} (kind)" + ))) + .execute(&database.pool) + .await + .unwrap(); + let filenode = |index: &'static str| { + let pool = database.pool.clone(); + let schema = database.name.clone(); + async move { + sqlx::query_scalar::<_, i64>( + "SELECT c.relfilenode::bigint FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = $1 AND c.relname = $2", + ) + .bind(schema) + .bind(index) + .fetch_one(&pool) + .await + .unwrap() + } + }; + let artifact_before = filenode("maint_reindex_artifact_idx").await; + let cancel = CancellationToken::new(); + assert!( + !super::reindexer::reindex_one_for_test( + &database.pool, + &database.schema, + &cancel, + "maint_reindex_artifact_idx", + ) + .await + .unwrap() + ); + assert_eq!( + filenode("maint_reindex_artifact_idx").await, + artifact_before + ); + + // A rebuild blocked behind an old snapshot is cancelled when the term + // ends, and the concurrent-build artifact it created is dropped (Go + // `ReindexDeletesArtifactsWhenCancelledWithStop`). + let mut snapshot = database.pool.begin().await.unwrap(); + sqlx::raw_sql(AssertSqlSafe(format!( + "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ; SELECT count(*) FROM {table}" + ))) + .execute(&mut *snapshot) + .await + .unwrap(); + let rebuild = tokio::spawn({ + let pool = database.pool.clone(); + let schema = database.schema.clone(); + let cancel = cancel.clone(); + async move { + super::reindexer::reindex_one_for_test( + &pool, + &schema, + &cancel, + "maint_reindex_cancel_idx", + ) + .await + } + }); + let artifact_count = || { + let pool = database.pool.clone(); + let schema = database.name.clone(); + async move { + sqlx::query_scalar::<_, i64>( + "SELECT count(*) FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = $1 AND c.relname LIKE 'maint\\_reindex\\_cancel\\_idx\\_cc%'", + ) + .bind(schema) + .fetch_one(&pool) + .await + .unwrap() + } + }; + tokio::time::timeout(Duration::from_secs(10), async { + while artifact_count().await == 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("the rebuild should create its concurrent artifact"); + cancel.cancel(); + snapshot.rollback().await.unwrap(); + assert!(rebuild.await.unwrap().is_err()); + assert_eq!(artifact_count().await, 0); + database.cleanup().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn periodic_start_hooks_and_run_on_start_follow_each_leadership_gain() { + use crate::{Hook, IntervalSchedule, PeriodicJob, PeriodicJobOpts, PeriodicJobs}; + + struct CountingHook(Arc); + + #[allow( + clippy::unused_async_trait_impl, + reason = "the hook only counts starts" + )] + impl Hook for CountingHook { + async fn periodic_jobs_start(&self, _jobs: &PeriodicJobs) -> Result<(), crate::BoxError> { + self.0.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + } + + let database = TestDatabase::new("rmt_periodic_gain").await; + let starts = Arc::new(AtomicUsize::new(0)); + let client = database + .client() + .hook(CountingHook(Arc::clone(&starts))) + .maintenance(MaintenanceConfig::default().with_elect_interval(Duration::from_millis(50))) + .periodic_job(PeriodicJob::with_options( + IntervalSchedule::new(Duration::from_hours(1)).unwrap(), + || NoTimeoutArgs {}, + PeriodicJobOpts::new() + .with_id("gain") + .with_run_on_start(true), + )) + .queue("default", QueueConfig::new(1)) + .build() + .unwrap(); + let periodic_count = || { + let pool = database.pool.clone(); + let table = database.table("river_job"); + async move { + sqlx::query_scalar::<_, i64>(AssertSqlSafe(format!( + "SELECT count(*) FROM {table} WHERE metadata ->> 'river:periodic_job_id' = 'gain'" + ))) + .fetch_one(&pool) + .await + .unwrap() + } + }; + let wait_for = |expected: usize| { + let starts = Arc::clone(&starts); + async move { + tokio::time::timeout(Duration::from_secs(10), async { + while starts.load(Ordering::SeqCst) < expected + || periodic_count().await < i64::try_from(expected).unwrap() + { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("expected {expected} leadership gains")); + } + }; + let mut handle = client.start().unwrap(); + // The resignation request is a notification, so it's lost unless the + // listener has subscribed by the time it's sent. + handle.wait_ready().await.unwrap(); + wait_for(1).await; + client.request_resign().await.unwrap(); + wait_for(2).await; + handle.shutdown().await.unwrap(); + assert_eq!(starts.load(Ordering::SeqCst), 2); + assert_eq!(periodic_count().await, 2); + database.cleanup().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn maintenance_start_retries_then_requests_resignation() { + // The leader resigns without a notification, so a client without + // notifications resigns too. + for poll_only in [false, true] { + maintenance_start_retries_then_resigns(poll_only).await; + } +} + +async fn maintenance_start_retries_then_resigns(poll_only: bool) { + use crate::{Hook, PeriodicJobs}; + + // Fails the first three start attempts of every client, like Go's + // `QueueMaintainerStartRetriesAndResigns`. + struct FlakyHook(Arc); + + #[allow( + clippy::unused_async_trait_impl, + reason = "the hook only counts starts" + )] + impl Hook for FlakyHook { + async fn periodic_jobs_start(&self, _jobs: &PeriodicJobs) -> Result<(), crate::BoxError> { + if self.0.fetch_add(1, Ordering::SeqCst) < 3 { + return Err("start failed".into()); + } + Ok(()) + } + } + + let database = TestDatabase::new("rmt_start_retry").await; + let attempts = Arc::new(AtomicUsize::new(0)); + let mut builder = database + .client() + .hook(FlakyHook(Arc::clone(&attempts))) + .maintenance(MaintenanceConfig::default().with_elect_interval(Duration::from_millis(50))) + .queue("default", QueueConfig::new(1)); + if poll_only { + builder = builder.without_notifications(); + } + let client = builder.build().unwrap(); + let elected_at = || { + let pool = database.pool.clone(); + let table = database.table("river_leader"); + async move { + sqlx::query_scalar::<_, DateTime>(AssertSqlSafe(format!( + "SELECT elected_at FROM {table}" + ))) + .fetch_optional(&pool) + .await + .unwrap() + } + }; + let mut handle = client.start().unwrap(); + let first_term = tokio::time::timeout(Duration::from_secs(10), async { + loop { + if let Some(elected_at) = elected_at().await { + return elected_at; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + + // After three failed attempts the leader asks to resign, and the next + // term's start succeeds. + tokio::time::timeout(Duration::from_secs(20), async { + while attempts.load(Ordering::SeqCst) < 4 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("maintenance start should be retried in a new term"); + let second_term = elected_at().await.unwrap(); + assert_ne!(second_term, first_term); + handle.shutdown().await.unwrap(); + assert_eq!(attempts.load(Ordering::SeqCst), 4); + database.cleanup().await; +} diff --git a/rust/riverqueue/src/periodic/cron.rs b/rust/riverqueue/src/periodic/cron.rs new file mode 100644 index 000000000..eb027b280 --- /dev/null +++ b/rust/riverqueue/src/periodic/cron.rs @@ -0,0 +1,923 @@ +//! Standard cron schedules with River Go's semantics. +//! +//! River Go documents periodic cron schedules parsed by robfig/cron's +//! `ParseStandard`. This module ports that parser and its `Next` algorithm so +//! that one expression string fires at the same times from either language: +//! five fields (minute, hour, day of month, month, day of week), weekdays +//! numbered 0-6 from Sunday, case-insensitive month and weekday names, +//! ranges, steps, lists, `*` and `?`, Vixie cron's day-of-month *or* +//! day-of-week rule, the `@yearly`, `@annually`, `@monthly`, `@weekly`, +//! `@daily`, `@midnight`, `@hourly`, and `@every ` descriptors, and +//! `CRON_TZ=`/`TZ=` prefixes. + +use std::{fmt, str::FromStr, time::Duration}; + +use chrono::{ + DateTime, Datelike, FixedOffset, Local, NaiveDate, NaiveDateTime, Offset as _, TimeZone, + Timelike, Utc, +}; +use thiserror::Error as ThisError; + +use super::PeriodicSchedule; + +/// Set when a field was written as `*` or `?` (robfig's `starBit`). +const STAR_BIT: u64 = 1 << 63; + +struct Bounds { + maximum: u32, + minimum: u32, + names: &'static [(&'static str, u32)], +} + +const MINUTES: Bounds = Bounds { + maximum: 59, + minimum: 0, + names: &[], +}; +const HOURS: Bounds = Bounds { + maximum: 23, + minimum: 0, + names: &[], +}; +const DAYS_OF_MONTH: Bounds = Bounds { + maximum: 31, + minimum: 1, + names: &[], +}; +const MONTHS: Bounds = Bounds { + maximum: 12, + minimum: 1, + names: &[ + ("jan", 1), + ("feb", 2), + ("mar", 3), + ("apr", 4), + ("may", 5), + ("jun", 6), + ("jul", 7), + ("aug", 8), + ("sep", 9), + ("oct", 10), + ("nov", 11), + ("dec", 12), + ], +}; +const DAYS_OF_WEEK: Bounds = Bounds { + maximum: 6, + minimum: 0, + names: &[ + ("sun", 0), + ("mon", 1), + ("tue", 2), + ("wed", 3), + ("thu", 4), + ("fri", 5), + ("sat", 6), + ], +}; + +/// Time zone in which a [`CronSchedule`] is evaluated. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +#[non_exhaustive] +pub enum CronTimeZone { + /// The process's local time zone, the default. Containers usually run in + /// UTC; set an explicit zone when clients in different zones share a + /// schedule. + #[default] + Local, + /// Coordinated Universal Time. + Utc, + /// A fixed offset from UTC, without daylight saving time. + Fixed(FixedOffset), + /// A zone from the IANA time zone database, with its daylight saving + /// rules. `CRON_TZ=` and `TZ=` prefixes naming a zone such as + /// `America/New_York` parse to this variant. + #[cfg(feature = "chrono-tz")] + Named(chrono_tz::Tz), +} + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum Spec { + Every(Duration), + Fields { + dom: u64, + dow: u64, + hour: u64, + minute: u64, + month: u64, + }, +} + +/// A periodic schedule written in standard five-field cron syntax. +/// +/// Parsing accepts exactly the expressions every River client accepts, and +/// [`CronSchedule::next_after`] returns the same occurrences, so a schedule +/// fires at the same times whichever client leads. +/// Without the `chrono-tz` feature, which bundles the IANA time zone +/// database, named `CRON_TZ=`/`TZ=` zones are limited to `UTC`, `Local`, and +/// `Etc/GMT±N`, and other names are rejected; use +/// [`CronSchedule::with_time_zone`] for other fixed offsets. +/// +/// Across daylight saving transitions, a daily schedule at a skipped +/// wall-clock time doesn't fire that day, and one at a repeated time can +/// fire twice. +#[derive(Clone, Debug, Eq, PartialEq)] +pub struct CronSchedule { + expression: String, + spec: Spec, + zone: Option, +} + +impl CronSchedule { + /// Parses a standard five-field cron expression or descriptor. + /// + /// # Errors + /// + /// Returns an error for an expression robfig/cron's `ParseStandard` + /// rejects, and for a `CRON_TZ=` zone this build can't resolve. + pub fn parse(expression: &str) -> Result { + expression.parse() + } + + /// Evaluates the schedule in `zone`, overriding any `CRON_TZ=` prefix. + #[must_use] + pub fn with_time_zone(mut self, zone: CronTimeZone) -> Self { + self.zone = Some(zone); + self + } + + /// Returns the zone in which periodic occurrences are computed. + #[must_use] + pub fn time_zone(&self) -> CronTimeZone { + self.zone.unwrap_or_default() + } + + /// Returns the first occurrence strictly after `after`, or `None` when + /// the schedule never matches within five years (robfig's zero time). + /// + /// Without an explicit zone, the occurrence is computed in `after`'s own + /// time zone, exactly like robfig's `Next` for a schedule without + /// `CRON_TZ`. + #[must_use] + pub fn next_after(&self, after: &DateTime) -> Option> { + let zone = after.timezone(); + match self.zone { + None => next_in(self.spec, after.clone()), + Some(CronTimeZone::Local) => next_in(self.spec, after.with_timezone(&Local)) + .map(|next| next.with_timezone(&zone)), + Some(CronTimeZone::Utc) => { + next_in(self.spec, after.with_timezone(&Utc)).map(|next| next.with_timezone(&zone)) + } + Some(CronTimeZone::Fixed(offset)) => next_in(self.spec, after.with_timezone(&offset)) + .map(|next| next.with_timezone(&zone)), + #[cfg(feature = "chrono-tz")] + Some(CronTimeZone::Named(named)) => next_in(self.spec, after.with_timezone(&named)) + .map(|next| next.with_timezone(&zone)), + } + } +} + +impl FromStr for CronSchedule { + type Err = CronScheduleParseError; + + fn from_str(expression: &str) -> Result { + let (zone, spec) = parse(expression).map_err(|message| CronScheduleParseError { + expression: expression.to_owned(), + message, + })?; + Ok(Self { + expression: expression.to_owned(), + spec, + zone, + }) + } +} + +impl fmt::Display for CronSchedule { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(&self.expression) + } +} + +impl PeriodicSchedule for CronSchedule { + fn next(&self, current: DateTime) -> Option> { + match self.time_zone() { + CronTimeZone::Local => next_in(self.spec, current.with_timezone(&Local)) + .map(|next| next.with_timezone(&Utc)), + CronTimeZone::Utc => next_in(self.spec, current), + CronTimeZone::Fixed(offset) => next_in(self.spec, current.with_timezone(&offset)) + .map(|next| next.with_timezone(&Utc)), + #[cfg(feature = "chrono-tz")] + CronTimeZone::Named(named) => next_in(self.spec, current.with_timezone(&named)) + .map(|next| next.with_timezone(&Utc)), + } + } +} + +/// Error returned when parsing a [`CronSchedule`]. +#[derive(Clone, Debug, Eq, PartialEq, ThisError)] +#[error("invalid periodic cron expression {expression:?}: {message}")] +pub struct CronScheduleParseError { + expression: String, + message: String, +} + +fn parse(expression: &str) -> Result<(Option, Spec), String> { + if expression.is_empty() { + return Err("empty spec string".to_owned()); + } + let mut spec = expression; + let mut zone = None; + if spec.starts_with("TZ=") || spec.starts_with("CRON_TZ=") { + let space = spec + .find(' ') + .ok_or_else(|| "time zone prefix must be followed by a schedule".to_owned())?; + let equals = spec.find('=').expect("prefix contains '='"); + zone = Some(parse_zone(&spec[equals + 1..space])?); + spec = spec[space..].trim(); + } + if spec.starts_with('@') { + return parse_descriptor(spec).map(|spec| (zone, spec)); + } + let fields = spec.split_whitespace().collect::>(); + if fields.len() != 5 { + return Err(format!( + "expected exactly 5 fields, found {}: {fields:?}", + fields.len() + )); + } + Ok(( + zone, + Spec::Fields { + minute: field(fields[0], &MINUTES)?, + hour: field(fields[1], &HOURS)?, + dom: field(fields[2], &DAYS_OF_MONTH)?, + month: field(fields[3], &MONTHS)?, + dow: field(fields[4], &DAYS_OF_WEEK)?, + }, + )) +} + +fn parse_zone(name: &str) -> Result { + match name { + // Go's `time.LoadLocation` maps "" to UTC. + "" | "UTC" | "Etc/UTC" => return Ok(CronTimeZone::Utc), + "Local" => return Ok(CronTimeZone::Local), + _ => {} + } + // POSIX-style `Etc/GMT+5` means five hours *behind* UTC. + if let Some(offset) = name.strip_prefix("Etc/GMT") + && let Some((sign, hours)) = offset + .strip_prefix('+') + .map(|hours| (-1, hours)) + .or_else(|| offset.strip_prefix('-').map(|hours| (1, hours))) + && let Ok(hours) = hours.parse::() + && (0..=14).contains(&hours) + && let Some(offset) = FixedOffset::east_opt(sign * hours * 3_600) + { + return Ok(CronTimeZone::Fixed(offset)); + } + #[cfg(feature = "chrono-tz")] + if let Ok(named) = name.parse::() { + return Ok(CronTimeZone::Named(named)); + } + #[cfg(feature = "chrono-tz")] + return Err(format!("provided bad location {name}: unknown time zone")); + #[cfg(not(feature = "chrono-tz"))] + Err(format!( + "provided bad location {name}: only UTC, Local, and Etc/GMT offsets are supported \ + without River's chrono-tz feature; use CronSchedule::with_time_zone for other zones" + )) +} + +fn parse_descriptor(descriptor: &str) -> Result { + let all = |bounds: &Bounds| bits(bounds.minimum, bounds.maximum, 1) | STAR_BIT; + let spec = match descriptor { + "@yearly" | "@annually" => Spec::Fields { + dom: 1 << 1, + dow: all(&DAYS_OF_WEEK), + hour: 1, + minute: 1, + month: 1 << 1, + }, + "@monthly" => Spec::Fields { + dom: 1 << 1, + dow: all(&DAYS_OF_WEEK), + hour: 1, + minute: 1, + month: all(&MONTHS), + }, + "@weekly" => Spec::Fields { + dom: all(&DAYS_OF_MONTH), + dow: 1, + hour: 1, + minute: 1, + month: all(&MONTHS), + }, + "@daily" | "@midnight" => Spec::Fields { + dom: all(&DAYS_OF_MONTH), + dow: all(&DAYS_OF_WEEK), + hour: 1, + minute: 1, + month: all(&MONTHS), + }, + "@hourly" => Spec::Fields { + dom: all(&DAYS_OF_MONTH), + dow: all(&DAYS_OF_WEEK), + hour: all(&HOURS), + minute: 1, + month: all(&MONTHS), + }, + _ => { + let Some(duration) = descriptor.strip_prefix("@every ") else { + return Err(format!("unrecognized descriptor: {descriptor}")); + }; + let nanos = parse_go_duration(duration) + .map_err(|message| format!("failed to parse duration {descriptor}: {message}"))?; + // robfig's `Every` rounds up to one second and drops subseconds. + let nanos = nanos.max(1_000_000_000); + let nanos = nanos - nanos % 1_000_000_000; + Spec::Every(Duration::from_nanos( + u64::try_from(nanos).expect("positive duration"), + )) + } + }; + Ok(spec) +} + +/// Parses a comma-separated list of ranges, skipping empty items like Go's +/// `strings.FieldsFunc`. +fn field(field: &str, bounds: &Bounds) -> Result { + let mut result = 0; + for expression in field.split(',').filter(|expression| !expression.is_empty()) { + result |= range(expression, bounds)?; + } + Ok(result) +} + +fn range(expression: &str, bounds: &Bounds) -> Result { + let range_and_step = expression.split('/').collect::>(); + let low_and_high = range_and_step[0].split('-').collect::>(); + let single = low_and_high.len() == 1; + let (start, mut end, mut extra) = if low_and_high[0] == "*" || low_and_high[0] == "?" { + (bounds.minimum, bounds.maximum, STAR_BIT) + } else { + let start = int_or_name(low_and_high[0], bounds)?; + let end = match low_and_high.len() { + 1 => start, + 2 => int_or_name(low_and_high[1], bounds)?, + _ => return Err(format!("too many hyphens: {expression}")), + }; + (start, end, 0) + }; + let step = match range_and_step.len() { + 1 => 1, + 2 => { + let step = go_atoi(range_and_step[1])?; + // "N/step" means "N-max/step". + if single { + end = bounds.maximum; + } + if step > 1 { + extra = 0; + } + step + } + _ => return Err(format!("too many slashes: {expression}")), + }; + if start < bounds.minimum { + return Err(format!( + "beginning of range ({start}) below minimum ({}): {expression}", + bounds.minimum + )); + } + if end > bounds.maximum { + return Err(format!( + "end of range ({end}) above maximum ({}): {expression}", + bounds.maximum + )); + } + if start > end { + return Err(format!( + "beginning of range ({start}) beyond end of range ({end}): {expression}" + )); + } + if step == 0 { + return Err(format!( + "step of range should be a positive number: {expression}" + )); + } + Ok(bits(start, end, step) | extra) +} + +fn int_or_name(expression: &str, bounds: &Bounds) -> Result { + let lower = expression.to_ascii_lowercase(); + if let Some((_, value)) = bounds.names.iter().find(|(name, _)| *name == lower) { + return Ok(*value); + } + go_atoi(expression) +} + +/// Go's `strconv.Atoi` followed by robfig's non-negative check. +fn go_atoi(expression: &str) -> Result { + let digits = expression.strip_prefix(['+', '-']).unwrap_or(expression); + if digits.is_empty() || !digits.bytes().all(|byte| byte.is_ascii_digit()) { + return Err(format!( + "failed to parse int from {expression}: invalid syntax" + )); + } + let number = expression + .parse::() + .map_err(|_| format!("failed to parse int from {expression}: value out of range"))?; + if number < 0 { + return Err(format!( + "negative number ({number}) not allowed: {expression}" + )); + } + // Anything above 63 is out of every field's range and rejected later. + Ok(u32::try_from(number.min(i64::from(u32::MAX))).expect("clamped")) +} + +fn bits(minimum: u32, maximum: u32, step: u32) -> u64 { + let mut bits = 0; + let mut value = minimum; + while value <= maximum { + bits |= 1 << value; + value += step; + } + bits +} + +/// Parses Go's `time.ParseDuration` syntax into signed nanoseconds. +fn parse_go_duration(text: &str) -> Result { + let invalid = || format!("time: invalid duration {text:?}"); + let (negative, mut rest) = match text.as_bytes().first() { + Some(b'-') => (true, &text[1..]), + Some(b'+') => (false, &text[1..]), + _ => (false, text), + }; + if rest == "0" { + return Ok(0); + } + if rest.is_empty() { + return Err(invalid()); + } + let mut total: i128 = 0; + while !rest.is_empty() { + let integer_length = rest.bytes().take_while(u8::is_ascii_digit).count(); + let integer = &rest[..integer_length]; + rest = &rest[integer_length..]; + let mut fraction = ""; + if let Some(after_dot) = rest.strip_prefix('.') { + let fraction_length = after_dot.bytes().take_while(u8::is_ascii_digit).count(); + fraction = &after_dot[..fraction_length]; + rest = &after_dot[fraction_length..]; + if integer.is_empty() && fraction.is_empty() { + return Err(invalid()); + } + } else if integer.is_empty() { + return Err(invalid()); + } + let unit_length = rest + .char_indices() + .find(|(_, character)| *character == '.' || character.is_ascii_digit()) + .map_or(rest.len(), |(index, _)| index); + let unit = &rest[..unit_length]; + rest = &rest[unit_length..]; + let unit_nanos: i128 = match unit { + "ns" => 1, + "us" | "\u{b5}s" | "\u{3bc}s" => 1_000, + "ms" => 1_000_000, + "s" => 1_000_000_000, + "m" => 60_000_000_000, + "h" => 3_600_000_000_000, + "" => return Err(format!("time: missing unit in duration {text:?}")), + _ => return Err(format!("time: unknown unit {unit:?} in duration {text:?}")), + }; + let integer = if integer.is_empty() { + 0 + } else { + integer.parse::().map_err(|_| invalid())? + }; + let mut value = integer * unit_nanos; + let mut scale = unit_nanos; + for digit in fraction.bytes() { + scale /= 10; + if scale == 0 { + break; + } + value += i128::from(digit - b'0') * scale; + } + total += value; + if total > i128::from(i64::MAX) { + return Err(invalid()); + } + } + Ok(if negative { -total } else { total }) +} + +/// robfig/cron's `SpecSchedule.Next` or `ConstantDelaySchedule.Next`. +#[expect( + clippy::too_many_lines, + reason = "a line-for-line port keeps robfig's field loops auditable" +)] +fn next_in(spec: Spec, after: DateTime) -> Option> { + let subsecond = chrono::Duration::nanoseconds(i64::from(after.nanosecond())); + let (minute, hour, dom, month, dow) = match spec { + Spec::Every(delay) => { + return Some(after + chrono::Duration::from_std(delay).ok()? - subsecond); + } + Spec::Fields { + dom, + dow, + hour, + minute, + month, + } => (minute, hour, dom, month, dow), + }; + let zone = after.timezone(); + let one_hour = chrono::Duration::hours(1); + let one_minute = chrono::Duration::minutes(1); + + // Start at the earliest possible time (the upcoming second). + let mut time = after + chrono::Duration::seconds(1) - subsecond; + let mut added = false; + let year_limit = time.year() + 5; + + 'wrap: loop { + if time.year() > year_limit { + return None; + } + + while (1 << time.month()) & month == 0 { + if !added { + added = true; + time = go_date(&zone, time.year(), time.month(), 1, 0, 0, 0)?; + } + // Go's `AddDate(0, 1, 0)`, normalizing overflowing days. + time = go_date( + &zone, + time.year(), + time.month() + 1, + time.day(), + time.hour(), + time.minute(), + time.second(), + )?; + if time.month() == 1 { + continue 'wrap; + } + } + + while !day_matches(dom, dow, &time) { + if !added { + added = true; + time = go_date(&zone, time.year(), time.month(), time.day(), 0, 0, 0)?; + } + time = go_date( + &zone, + time.year(), + time.month(), + time.day() + 1, + time.hour(), + time.minute(), + time.second(), + )?; + // Midnight may not exist on a daylight saving transition. + if time.hour() != 0 { + let hour = i64::from(time.hour()); + time = if hour > 12 { + time + chrono::Duration::hours(24 - hour) + } else { + time - chrono::Duration::hours(hour) + }; + } + if time.day() == 1 { + continue 'wrap; + } + } + + while (1 << time.hour()) & hour == 0 { + if !added { + added = true; + time = go_date( + &zone, + time.year(), + time.month(), + time.day(), + time.hour(), + 0, + 0, + )?; + } + time += one_hour; + if time.hour() == 0 { + continue 'wrap; + } + } + + while (1 << time.minute()) & minute == 0 { + if !added { + added = true; + time = truncate_to_minute(time); + } + time += one_minute; + if time.minute() == 0 { + continue 'wrap; + } + } + + // Standard specs always fire at second zero. robfig's seconds loop + // steps one second at a time until the minute rolls over and then + // re-validates every field, which is a single jump here. + let second = i64::from(time.second()); + if second != 0 { + added = true; + time += chrono::Duration::seconds(60 - second); + continue 'wrap; + } + return Some(time); + } +} + +fn day_matches(dom: u64, dow: u64, time: &DateTime) -> bool { + let day_of_month = (1 << time.day()) & dom != 0; + let weekday = (1 << time.weekday().num_days_from_sunday()) & dow != 0; + if dom & STAR_BIT != 0 || dow & STAR_BIT != 0 { + day_of_month && weekday + } else { + day_of_month || weekday + } +} + +/// Go's `time.Truncate(time.Minute)`, which rounds absolute time. +fn truncate_to_minute(time: DateTime) -> DateTime { + let seconds = time.timestamp().rem_euclid(60); + let nanos = i64::from(time.nanosecond()); + time - chrono::Duration::seconds(seconds) - chrono::Duration::nanoseconds(nanos) +} + +/// Go's `time.Date` in `zone`: overflowing months and days roll forward, and +/// a wall-clock time that a transition skips or repeats resolves as Go +/// resolves it (see [`resolve_local`]). +fn go_date( + zone: &Z, + year: i32, + month: u32, + day: u32, + hour: u32, + minute: u32, + second: u32, +) -> Option> { + let months = i64::from(year) * 12 + i64::from(month) - 1; + let year = i32::try_from(months.div_euclid(12)).ok()?; + let month = u32::try_from(months.rem_euclid(12)).ok()? + 1; + let date = NaiveDate::from_ymd_opt(year, month, 1)? + .checked_add_days(chrono::Days::new(u64::from(day.checked_sub(1)?)))?; + let naive = date.and_hms_opt(hour, minute, second)?; + resolve_local(zone, naive) +} + +/// Converts a wall-clock time in `zone` to an instant exactly like Go's +/// `time.Date`, which robfig's schedule arithmetic relies on. +/// +/// Go takes the offset in effect at the instant whose UTC reading equals the +/// wall-clock time, and keeps it if subtracting it lands in the same offset +/// period; otherwise it uses the offset in effect at that result. A repeated +/// wall-clock time therefore usually resolves to its earlier instant, and a +/// skipped one to the instant an hour before the transition's end, which +/// reads as the hour before it: `02:30` on a New York spring-forward day is +/// `01:30` EST, so a daily `30 2 * * *` schedule skips that day. Like Go, a +/// zero first offset is used as is. +fn resolve_local(zone: &Z, naive: NaiveDateTime) -> Option> { + let offset_at = |instant: NaiveDateTime| { + i64::from( + zone.offset_from_utc_datetime(&instant) + .fix() + .local_minus_utc(), + ) + }; + let mut offset = offset_at(naive); + if offset != 0 { + let utc = naive.checked_sub_signed(chrono::Duration::seconds(offset))?; + let corrected = offset_at(utc); + if corrected != offset { + offset = corrected; + } + } + let utc = naive.checked_sub_signed(chrono::Duration::seconds(offset))?; + Some(zone.from_utc_datetime(&utc)) +} + +#[cfg(test)] +mod tests { + use chrono::{DateTime, FixedOffset}; + use serde::Deserialize; + + use super::{CronSchedule, CronTimeZone, PeriodicSchedule}; + + #[derive(Deserialize)] + #[expect( + clippy::struct_field_names, + reason = "the fields mirror the fixture's keys" + )] + struct Fixture { + cron_cases: Vec, + cron_invalid: Vec, + cron_named_zone_cases: Vec, + } + + #[derive(Deserialize)] + struct CronCase { + expression: String, + from: DateTime, + name: String, + next: Vec>, + } + + fn fixture() -> Fixture { + serde_json::from_str(include_str!("../../tests/fixtures/maintenance_values.json")).unwrap() + } + + fn assert_matches_go(case: &CronCase) { + let schedule = CronSchedule::parse(&case.expression) + .unwrap_or_else(|error| panic!("{}: {error}", case.name)); + // The fixture records five occurrences, stopping early at Go's zero time. + let mut current = case.from; + let mut observed = Vec::new(); + while observed.len() < 5 { + let Some(next) = schedule.next_after(¤t) else { + break; + }; + observed.push(next); + current = next; + } + assert_eq!(observed, case.next, "{}", case.name); + for (observed, expected) in observed.iter().zip(&case.next) { + assert_eq!(observed.offset(), expected.offset(), "{}", case.name); + } + } + + #[test] + fn cron_schedules_match_go_fixture() { + let fixture = fixture(); + assert!(!fixture.cron_cases.is_empty()); + for case in &fixture.cron_cases { + assert_matches_go(case); + } + } + + // Go's `time.Date` resolves wall-clock times that a daylight saving + // transition skips or repeats in a particular way, which decides whether + // a schedule fires that day. Fixed offsets never exercise it, so check it + // against the local zones Go would use, rebuilt from the named zone + // cases' own offsets. + #[test] + fn transitions_resolve_like_go_time_date() { + use chrono::{NaiveDate, TimeZone as _}; + + // A zone at UTC-5 that springs forward to UTC-4 at 07:00 UTC on + // 2026-03-08 and falls back at 06:00 UTC on 2026-11-01, like New + // York. + #[derive(Clone, Copy, Debug)] + struct NewYork2026; + + impl chrono::TimeZone for NewYork2026 { + type Offset = FixedOffset; + + fn from_offset(_: &FixedOffset) -> Self { + Self + } + + fn offset_from_local_date(&self, _: &NaiveDate) -> chrono::LocalResult { + unimplemented!("River resolves wall-clock times itself") + } + + fn offset_from_local_datetime( + &self, + _: &chrono::NaiveDateTime, + ) -> chrono::LocalResult { + unimplemented!("River resolves wall-clock times itself") + } + + fn offset_from_utc_date(&self, _: &NaiveDate) -> FixedOffset { + unimplemented!("River only looks up instants") + } + + fn offset_from_utc_datetime(&self, utc: &chrono::NaiveDateTime) -> FixedOffset { + let spring = NaiveDate::from_ymd_opt(2026, 3, 8) + .unwrap() + .and_hms_opt(7, 0, 0) + .unwrap(); + let fall = NaiveDate::from_ymd_opt(2026, 11, 1) + .unwrap() + .and_hms_opt(6, 0, 0) + .unwrap(); + let hours = if (spring..fall).contains(utc) { 4 } else { 5 }; + FixedOffset::west_opt(hours * 3_600).unwrap() + } + } + + let wall = |month, day, hour, minute| { + NaiveDate::from_ymd_opt(2026, month, day) + .unwrap() + .and_hms_opt(hour, minute, 0) + .unwrap() + }; + let resolved = |naive| { + super::resolve_local(&NewYork2026, naive) + .unwrap() + .naive_utc() + }; + // A skipped time resolves to the hour before the transition. + assert_eq!(resolved(wall(3, 8, 2, 30)), wall(3, 8, 6, 30)); + // A repeated time resolves to its earlier instant. + assert_eq!(resolved(wall(11, 1, 1, 30)), wall(11, 1, 5, 30)); + assert_eq!(resolved(wall(7, 1, 9, 0)), wall(7, 1, 13, 0)); + assert_eq!(resolved(wall(12, 1, 9, 0)), wall(12, 1, 14, 0)); + + // So, like Go, a daily schedule at a skipped time misses that day, + // and one at a repeated time fires twice. + let from = NewYork2026.from_utc_datetime(&wall(3, 7, 12, 0)); + let skipped = CronSchedule::parse("30 2 * * *").unwrap(); + let first = skipped.next_after(&from).unwrap(); + assert_eq!(first.naive_utc(), wall(3, 9, 6, 30)); + let from = NewYork2026.from_utc_datetime(&wall(10, 31, 12, 0)); + let repeated = CronSchedule::parse("30 1 * * *").unwrap(); + let first = repeated.next_after(&from).unwrap(); + let second = repeated.next_after(&first).unwrap(); + assert_eq!( + (first.naive_utc(), second.naive_utc()), + (wall(11, 1, 5, 30), wall(11, 1, 6, 30)) + ); + } + + #[cfg(feature = "chrono-tz")] + #[test] + fn named_zone_schedules_match_go_fixture() { + let fixture = fixture(); + assert!(!fixture.cron_named_zone_cases.is_empty()); + for case in &fixture.cron_named_zone_cases { + assert_matches_go(case); + } + } + + #[cfg(not(feature = "chrono-tz"))] + #[test] + fn named_zones_need_the_time_zone_database() { + for case in fixture().cron_named_zone_cases { + let error = CronSchedule::parse(&case.expression).unwrap_err(); + assert!(error.to_string().contains("chrono-tz"), "{error}"); + } + } + + #[test] + fn cron_rejects_what_go_rejects() { + for expression in fixture().cron_invalid { + assert!( + CronSchedule::parse(&expression).is_err(), + "{expression:?} should be rejected" + ); + } + } + + #[test] + fn explicit_time_zones_override_the_reference_zone() { + let schedule = CronSchedule::parse("0 9 * * *").unwrap(); + assert_eq!(schedule.time_zone(), CronTimeZone::Local); + let eastern = FixedOffset::west_opt(5 * 3_600).unwrap(); + let utc = schedule.clone().with_time_zone(CronTimeZone::Utc); + let from = DateTime::parse_from_rfc3339("2026-03-07T08:00:00-05:00").unwrap(); + assert_eq!( + utc.next_after(&from).unwrap().to_rfc3339(), + "2026-03-08T04:00:00-05:00" + ); + let fixed = schedule.with_time_zone(CronTimeZone::Fixed(eastern)); + assert_eq!( + fixed + .next( + DateTime::parse_from_rfc3339("2026-03-07T13:00:00Z") + .unwrap() + .to_utc() + ) + .unwrap() + .to_rfc3339(), + "2026-03-07T14:00:00+00:00" + ); + assert_eq!( + CronSchedule::parse("CRON_TZ=Etc/GMT+5 0 9 * * *") + .unwrap() + .time_zone(), + CronTimeZone::Fixed(eastern) + ); + #[cfg(feature = "chrono-tz")] + assert_eq!( + CronSchedule::parse("CRON_TZ=America/New_York 0 9 * * *") + .unwrap() + .time_zone(), + CronTimeZone::Named(chrono_tz::America::New_York) + ); + #[cfg(not(feature = "chrono-tz"))] + assert!(CronSchedule::parse("CRON_TZ=America/New_York 0 9 * * *").is_err()); + } +} diff --git a/rust/riverqueue/src/periodic/mod.rs b/rust/riverqueue/src/periodic/mod.rs new file mode 100644 index 000000000..69ffe8ffb --- /dev/null +++ b/rust/riverqueue/src/periodic/mod.rs @@ -0,0 +1,696 @@ +//! Leader-owned periodic job scheduling. + +mod cron; + +pub use self::cron::{CronSchedule, CronScheduleParseError, CronTimeZone}; + +use std::{ + collections::{HashMap, HashSet}, + fmt, + sync::{Arc, Mutex, MutexGuard, PoisonError}, + time::Duration, +}; + +use chrono::{DateTime, Utc}; +use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; + +use crate::{Client, Error, InsertOpts, JobArgs}; + +/// A schedule that calculates the next periodic run after a UTC instant. +pub trait PeriodicSchedule: Send + Sync + 'static { + /// Returns the next run time, or `None` to disable future runs. + fn next(&self, current: DateTime) -> Option>; +} + +/// A fixed-duration periodic schedule. +#[derive(Clone, Copy, Debug)] +pub struct IntervalSchedule(Duration); + +impl IntervalSchedule { + /// Creates a fixed schedule. Like River Go's `PeriodicInterval`, any + /// positive interval is accepted, including sub-second ones, so the same + /// periodic jobs can be configured in Rust and Go clients that share a + /// database. + /// + /// # Errors + /// + /// Returns [`Error::Configuration`] for a zero interval. + pub fn new(interval: Duration) -> Result { + if interval.is_zero() { + return Err(Error::configuration( + "periodic interval must be positive".to_owned(), + )); + } + Ok(Self(interval)) + } + + /// Returns the fixed interval between occurrences. + #[must_use] + pub const fn interval(&self) -> Duration { + self.0 + } +} + +impl PeriodicSchedule for IntervalSchedule { + fn next(&self, current: DateTime) -> Option> { + chrono::Duration::from_std(self.0) + .ok() + .and_then(|interval| current.checked_add_signed(interval)) + } +} + +/// A schedule that never runs. +#[derive(Clone, Copy, Debug, Default)] +pub struct NeverSchedule; + +impl PeriodicSchedule for NeverSchedule { + fn next(&self, _current: DateTime) -> Option> { + None + } +} + +/// Options for a periodic job. +#[derive(Clone, Debug, Default)] +pub struct PeriodicJobOpts { + /// Optional identifier, unique within one client. + pub(crate) id: Option, + /// Inserts once whenever this client becomes leader. + pub(crate) run_on_start: bool, +} + +impl PeriodicJobOpts { + /// Creates periodic-job options with no ID and no initial run. + #[must_use] + pub const fn new() -> Self { + Self { + id: None, + run_on_start: false, + } + } + + /// Returns the optional identifier. + #[must_use] + pub fn id(&self) -> Option<&str> { + self.id.as_deref() + } + + /// Returns whether an occurrence is inserted whenever leadership begins. + #[must_use] + pub const fn run_on_start(&self) -> bool { + self.run_on_start + } + + /// Returns the options with an occurrence inserted, or not, whenever + /// this client becomes leader. + #[must_use] + pub const fn with_run_on_start(mut self, run_on_start: bool) -> Self { + self.run_on_start = run_on_start; + self + } + + /// Sets an identifier unique within one client. + #[must_use] + pub fn with_id(mut self, id: impl Into) -> Self { + self.id = Some(id.into()); + self + } +} + +#[derive(Clone)] +pub(crate) struct PeriodicInsert { + pub(crate) defaults: InsertOpts, + pub(crate) encoded_args: Box, + pub(crate) kind: &'static str, + pub(crate) opts: InsertOpts, + pub(crate) unique_fields: &'static [&'static [&'static str]], +} + +/// Type-erased periodic job definition. +#[derive(Clone)] +pub struct PeriodicJob { + pub(crate) constructor: Arc Result, Error> + Send + Sync>, + pub(crate) opts: PeriodicJobOpts, + pub(crate) schedule: Arc, +} + +impl PeriodicJob { + /// Creates a periodic job that inserts one job using its type defaults on + /// every scheduled occurrence. + pub fn new(schedule: S, constructor: F) -> Self + where + A: JobArgs, + F: Fn() -> A + Send + Sync + 'static, + S: PeriodicSchedule, + { + Self::with_options(schedule, constructor, PeriodicJobOpts::new()) + } + + /// Creates a periodic job with registration options. + pub fn with_options(schedule: S, constructor: F, opts: PeriodicJobOpts) -> Self + where + A: JobArgs, + F: Fn() -> A + Send + Sync + 'static, + S: PeriodicSchedule, + { + Self::conditional_with_options( + schedule, + move || Some((constructor(), InsertOpts::default())), + opts, + ) + } + + /// Creates a periodic job whose constructor can skip an occurrence or set + /// per-occurrence insertion options. + pub fn conditional(schedule: S, constructor: F) -> Self + where + A: JobArgs, + F: Fn() -> Option<(A, InsertOpts)> + Send + Sync + 'static, + S: PeriodicSchedule, + { + Self::conditional_with_options(schedule, constructor, PeriodicJobOpts::new()) + } + + /// Creates a conditional periodic job with registration options. + pub fn conditional_with_options( + schedule: S, + constructor: F, + opts: PeriodicJobOpts, + ) -> Self + where + A: JobArgs, + F: Fn() -> Option<(A, InsertOpts)> + Send + Sync + 'static, + S: PeriodicSchedule, + { + Self { + constructor: Arc::new(move || { + let Some((args, opts)) = constructor() else { + return Ok(None); + }; + Ok(Some(PeriodicInsert { + defaults: A::default_insert_opts(), + encoded_args: crate::encoding::encode_args(&args)?, + kind: A::KIND, + opts, + unique_fields: A::unique_fields(), + })) + }), + opts, + schedule: Arc::new(schedule), + } + } +} + +impl fmt::Debug for PeriodicJob { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("PeriodicJob") + .field("opts", &self.opts) + .finish_non_exhaustive() + } +} + +/// Opaque handle used to remove a dynamically configured periodic job. +#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] +pub struct PeriodicJobHandle(u64); + +pub(crate) struct PeriodicEntry { + pub(crate) job: PeriodicJob, + pub(crate) next_run: Option>, + pub(crate) needs_initialization: bool, +} + +#[derive(Default)] +pub(crate) struct PeriodicRegistry { + pub(crate) entries: HashMap, + next_handle: u64, +} + +/// Dynamically configurable periodic jobs for a client. +#[derive(Clone)] +pub struct PeriodicJobs { + changed: Arc, + leader_election_disabled: bool, + pub(crate) registry: Arc>, +} + +/// Jobs due within this margin are inserted in the current pass, like Go's +/// enqueuer, which also keeps each occurrence's original scheduled time. +const DUE_MARGIN: chrono::Duration = chrono::Duration::milliseconds(100); + +/// Bound on one periodic insert (Go `riversharedmaintenance.TimeoutDefault`). +const INSERT_TIMEOUT: Duration = Duration::from_secs(30); + +impl fmt::Debug for PeriodicJobs { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("PeriodicJobs") + .field("len", &self.lock().entries.len()) + .finish_non_exhaustive() + } +} + +impl PeriodicJobs { + pub(crate) fn from_jobs( + jobs: Vec, + leader_election_disabled: bool, + ) -> Result { + validate_jobs(&jobs, &HashSet::new())?; + let mut registry = PeriodicRegistry::default(); + for job in jobs { + registry.insert(job); + } + Ok(Self { + changed: Arc::new(Notify::new()), + leader_election_disabled, + registry: Arc::new(Mutex::new(registry)), + }) + } + + /// Locks the registry. It's never held across an await, and every update + /// leaves it consistent, so a poisoned lock is still usable. + fn lock(&self) -> MutexGuard<'_, PeriodicRegistry> { + self.registry.lock().unwrap_or_else(PoisonError::into_inner) + } + + /// Completes after the next registration change. + pub(crate) fn changed(&self) -> tokio::sync::futures::Notified<'_> { + self.changed.notified() + } + + fn notify_changed(&self) { + self.changed.notify_waiters(); + } + + /// Rejects additions to a client that never leads, which would never + /// enqueue them. + fn ensure_electable(&self) -> Result<(), Error> { + if self.leader_election_disabled { + return Err(Error::configuration( + "periodic jobs can't be added when leader election is disabled".to_owned(), + )); + } + Ok(()) + } + + /// Returns the earliest scheduled occurrence, or now when a job still + /// needs its first schedule computed. + pub(crate) fn next_run_at(&self) -> Option> { + let registry = self.lock(); + let mut next: Option> = None; + for entry in registry.entries.values() { + let candidate = if entry.needs_initialization { + Some(Utc::now()) + } else { + entry.next_run + }; + if let Some(candidate) = candidate { + next = Some(next.map_or(candidate, |next| next.min(candidate))); + } + } + next + } + + /// Adds one periodic job and returns its removal handle. + /// + /// Adding or removing periodic jobs affects only this client, which + /// enqueues them only while it's the elected leader. To make sure a + /// periodic job is fully enabled or disabled, change it on every client + /// eligible for leader election across all processes. + /// + /// # Errors + /// + /// Returns a configuration error when the job's identifier is invalid or + /// already configured, or when the client was built with + /// [`ClientBuilder::without_leader_election`](crate::ClientBuilder::without_leader_election). + pub fn add(&self, job: PeriodicJob) -> Result { + self.ensure_electable()?; + let mut registry = self.lock(); + let ids = registry + .entries + .values() + .filter_map(|entry| entry.job.opts.id.clone()) + .collect(); + validate_jobs(std::slice::from_ref(&job), &ids)?; + let handle = registry.insert(job); + drop(registry); + self.notify_changed(); + Ok(handle) + } + + /// Adds many jobs atomically after validating their identifiers. + /// + /// Like [`PeriodicJobs::add`], this affects only this client. + /// + /// # Errors + /// + /// Returns the same errors as [`PeriodicJobs::add`]; no job is added + /// when any is rejected. + pub fn add_many( + &self, + jobs: impl IntoIterator, + ) -> Result, Error> { + let jobs = jobs.into_iter().collect::>(); + self.ensure_electable()?; + let mut registry = self.lock(); + let ids = registry + .entries + .values() + .filter_map(|entry| entry.job.opts.id.clone()) + .collect(); + validate_jobs(&jobs, &ids)?; + let handles = jobs.into_iter().map(|job| registry.insert(job)).collect(); + drop(registry); + self.notify_changed(); + Ok(handles) + } + + /// Removes all configured periodic jobs. + pub fn clear(&self) { + self.lock().entries.clear(); + self.notify_changed(); + } + + /// Removes a job by handle. + pub fn remove(&self, handle: PeriodicJobHandle) -> bool { + let removed = self.lock().entries.remove(&handle).is_some(); + self.notify_changed(); + removed + } + + /// Removes a job by identifier. + pub fn remove_by_id(&self, id: &str) -> bool { + let mut registry = self.lock(); + let handle = registry.entries.iter().find_map(|(handle, entry)| { + (entry.job.opts.id.as_deref() == Some(id)).then_some(*handle) + }); + let removed = handle.is_some_and(|handle| registry.entries.remove(&handle).is_some()); + drop(registry); + self.notify_changed(); + removed + } + + pub(crate) fn reset_for_leadership(&self) { + for entry in self.lock().entries.values_mut() { + entry.needs_initialization = true; + entry.next_run = None; + } + } + + /// Inserts every occurrence due by `now` plus a small margin. Newly added + /// jobs are scheduled from `now` and, when configured, inserted once. + /// + /// Like Go's enqueuer, an occurrence whose insert fails or times out is + /// logged and skipped rather than retried, and the pass stops as soon as + /// `cancel` fires so a client that lost leadership stops inserting. + pub(crate) async fn run_due( + &self, + client: &Client, + now: DateTime, + cancel: &CancellationToken, + ) { + struct DueJob { + advance_handle: Option, + handle: PeriodicJobHandle, + job: PeriodicJob, + target: DateTime, + } + + let due = { + let mut registry = self.lock(); + let mut due = Vec::new(); + for (handle, entry) in &mut registry.entries { + if entry.needs_initialization { + entry.needs_initialization = false; + entry.next_run = entry.job.schedule.next(now); + if entry.job.opts.run_on_start { + due.push(DueJob { + advance_handle: None, + handle: *handle, + job: entry.job.clone(), + target: now, + }); + } + continue; + } + if let Some(target) = entry.next_run + && target < now + DUE_MARGIN + { + due.push(DueJob { + advance_handle: Some(*handle), + handle: *handle, + job: entry.job.clone(), + target, + }); + } + } + // Insert in the order the jobs were added, as Go's enqueuer does. + due.sort_by_key(|due_job| due_job.handle); + due + }; + + for due_job in due { + let result = (due_job.job.constructor)(); + match result { + Ok(Some(insert)) => { + let mut opts = InsertOpts::resolve( + client.default_max_attempts(), + insert.defaults.clone(), + insert.opts.clone(), + ); + opts.metadata + .insert("periodic", true) + .expect("boolean metadata serializes"); + if let Some(id) = &due_job.job.opts.id { + opts.metadata + .insert(crate::METADATA_KEY_PERIODIC_JOB_ID, id) + .expect("string metadata serializes"); + } + let insert = tokio::time::timeout( + INSERT_TIMEOUT, + client.insert_periodic(insert, opts, due_job.target), + ); + let result = tokio::select! { + biased; + () = cancel.cancelled() => return, + result = insert => result, + }; + match result { + Ok(Ok(_)) => {} + Ok(Err(error)) => { + tracing::error!(error = %crate::error::Chain(&error), "River periodic job insertion failed"); + } + Err(_) => tracing::error!("River periodic job insertion timed out"), + } + } + Ok(None) => {} + Err(error) => { + tracing::error!(error = %crate::error::Chain(&error), "River periodic job constructor failed"); + } + } + + if let Some(handle) = due_job.advance_handle { + let mut registry = self.lock(); + if let Some(entry) = registry.entries.get_mut(&handle) + && !entry.needs_initialization + && entry.next_run == Some(due_job.target) + { + entry.next_run = entry.job.schedule.next(due_job.target); + } + } + } + } +} + +impl PeriodicRegistry { + fn insert(&mut self, job: PeriodicJob) -> PeriodicJobHandle { + self.next_handle = self.next_handle.wrapping_add(1); + let handle = PeriodicJobHandle(self.next_handle); + self.entries.insert( + handle, + PeriodicEntry { + job, + needs_initialization: true, + next_run: None, + }, + ); + handle + } +} + +fn validate_jobs(jobs: &[PeriodicJob], existing_ids: &HashSet) -> Result<(), Error> { + let mut ids = existing_ids.clone(); + for job in jobs { + if let Some(id) = &job.opts.id { + if id.is_empty() { + return Err(Error::invalid_job_context( + "periodic job", + "periodic job ID cannot be empty".to_owned(), + )); + } + if !ids.insert(id.clone()) { + return Err(Error::invalid_job_context( + "periodic job", + format!("periodic job with ID already registered: {id}"), + )); + } + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + #[cfg(feature = "postgres")] + use std::sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }; + + use super::*; + use serde::{Deserialize, Serialize}; + #[cfg(feature = "postgres")] + use sqlx::postgres::PgPoolOptions; + + #[derive(Clone, Deserialize, Serialize)] + struct TestArgs; + + #[test] + fn interval_schedules_accept_any_positive_interval() { + let schedule = IntervalSchedule::new(Duration::from_millis(500)).unwrap(); + let now = Utc::now(); + assert_eq!( + schedule.next(now), + Some(now + chrono::Duration::milliseconds(500)) + ); + let error = IntervalSchedule::new(Duration::ZERO).unwrap_err(); + assert!(matches!(error, Error::Configuration(_)), "{error}"); + } + + impl JobArgs for TestArgs { + const KIND: &'static str = "periodic_test"; + } + + fn job(id: &str) -> PeriodicJob { + PeriodicJob::with_options( + NeverSchedule, + || TestArgs, + PeriodicJobOpts::new().with_id(id), + ) + } + + #[test] + fn dynamic_registration_is_atomic_and_removable() { + let jobs = PeriodicJobs::from_jobs(Vec::new(), false).unwrap(); + let first = jobs.add(job("first")).unwrap(); + let added = jobs.add_many([job("second"), job("third")]).unwrap(); + assert_eq!(added.len(), 2); + assert_eq!(jobs.lock().entries.len(), 3); + + assert!(jobs.add_many([job("fourth"), job("second")]).is_err()); + assert_eq!(jobs.lock().entries.len(), 3); + + assert!(jobs.remove(first)); + assert!(!jobs.remove(first)); + assert!(jobs.remove_by_id("second")); + assert!(!jobs.remove_by_id("missing")); + + jobs.clear(); + assert!(jobs.lock().entries.is_empty()); + } + + #[cfg(feature = "postgres")] + #[tokio::test] + async fn insert_failure_skips_the_occurrence_like_go() { + let attempts = Arc::new(AtomicUsize::new(0)); + let constructed = Arc::clone(&attempts); + let jobs = PeriodicJobs::from_jobs( + vec![PeriodicJob::new( + IntervalSchedule::new(Duration::from_secs(1)).unwrap(), + move || { + constructed.fetch_add(1, Ordering::SeqCst); + TestArgs + }, + )], + false, + ) + .unwrap(); + // Nothing listens on this port, so every insertion fails. + let pool = PgPoolOptions::new() + .acquire_timeout(Duration::from_millis(100)) + .connect_lazy("postgres://127.0.0.1:1/river_periodic_test") + .unwrap(); + let client = Client::builder(pool).build().unwrap(); + let now = Utc::now(); + let target = now + chrono::Duration::seconds(1); + let cancel = CancellationToken::new(); + + jobs.run_due(&client, now, &cancel).await; + jobs.run_due(&client, target, &cancel).await; + jobs.run_due(&client, target, &cancel).await; + + // The failed occurrence isn't retried; the schedule moves on. + assert_eq!(attempts.load(Ordering::SeqCst), 1); + let registry = jobs.lock(); + assert_eq!( + registry.entries.values().next().unwrap().next_run, + Some(target + chrono::Duration::seconds(1)) + ); + } + + #[cfg(feature = "postgres")] + #[tokio::test] + async fn cancellation_stops_a_pass_without_advancing() { + let jobs = PeriodicJobs::from_jobs( + vec![PeriodicJob::new( + IntervalSchedule::new(Duration::from_secs(1)).unwrap(), + || TestArgs, + )], + false, + ) + .unwrap(); + // A server that accepts connections and never answers makes the + // insert hang while connecting, whatever the network does. + let server = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = server.local_addr().unwrap(); + let server = tokio::spawn(async move { + let mut connections = Vec::new(); + while let Ok((connection, _)) = server.accept().await { + connections.push(connection); + } + }); + let pool = PgPoolOptions::new() + .acquire_timeout(Duration::from_hours(1)) + .connect_lazy(&format!("postgres://{address}/river_periodic_test")) + .unwrap(); + let client = Client::builder(pool).build().unwrap(); + let now = Utc::now(); + let target = now + chrono::Duration::seconds(1); + let cancel = CancellationToken::new(); + jobs.run_due(&client, now, &cancel).await; + + let pass = jobs.run_due(&client, target, &cancel); + let cancel_soon = async { + tokio::time::sleep(Duration::from_millis(50)).await; + cancel.cancel(); + }; + tokio::time::timeout(Duration::from_secs(5), async { + tokio::join!(pass, cancel_soon); + }) + .await + .expect("a cancelled pass returns promptly"); + + server.abort(); + let registry = jobs.lock(); + assert_eq!( + registry.entries.values().next().unwrap().next_run, + Some(target) + ); + } + + #[test] + fn static_registration_rejects_invalid_identifiers() { + assert!(PeriodicJobs::from_jobs(vec![job("duplicate"), job("duplicate")], false).is_err()); + assert!(PeriodicJobs::from_jobs(vec![job("")], false).is_err()); + } +} diff --git a/rust/riverqueue/src/pilot.rs b/rust/riverqueue/src/pilot.rs new file mode 100644 index 000000000..3406ca9b5 --- /dev/null +++ b/rust/riverqueue/src/pilot.rs @@ -0,0 +1,448 @@ +//! Producer sessions and database bindings for River's own companion crates. +//! +//! Everything here is re-exported from [`crate::__private`] and shares its +//! stability rules: it changes without notice between any two versions. + +use std::{fmt, sync::Arc, time::Duration}; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use serde_json::{Map, Value}; +#[cfg(feature = "postgres")] +use sqlx::Postgres; +#[cfg(feature = "sqlite")] +use sqlx::Sqlite; +use sqlx::Transaction; +use tokio_util::sync::CancellationToken; + +use crate::__private::{ClaimedJob, DatabaseConfig, DatabaseConnection, DatabasePool, PilotError}; +use crate::client::ClientInner; +use crate::database::DatabaseKind; +use crate::{Error, JobRow, Queue}; + +/// The client's database as seen by an extension: its caller-owned pool and +/// backend configuration. +/// +/// Cloning is cheap. An extension opens its own transactions with +/// [`PilotDatabase::begin`], which, like River's own, can be dropped at any +/// point without leaking an open transaction. +#[derive(Clone)] +pub struct PilotDatabase { + config: DatabaseConfig, + pool: DatabasePool, +} + +impl PilotDatabase { + pub(crate) const fn new(pool: DatabasePool, config: DatabaseConfig) -> Self { + Self { config, pool } + } + + /// Begins a transaction that may write. SQLite transactions take the + /// write lock up front with `BEGIN IMMEDIATE`. + /// + /// # Errors + /// + /// Returns the database error when the transaction can't begin. + pub async fn begin(&self) -> Result { + let transaction = match &self.pool { + #[cfg(feature = "postgres")] + DatabasePool::Postgres(pool) => { + PilotTransactionInner::Postgres(crate::database::begin_postgres(pool).await?) + } + #[cfg(feature = "sqlite")] + DatabasePool::Sqlite(pool) => { + PilotTransactionInner::Sqlite(crate::database::begin_sqlite_write(pool).await?) + } + }; + Ok(PilotTransaction(transaction)) + } + + /// Returns the backend configuration. + #[must_use] + pub const fn config(&self) -> &DatabaseConfig { + &self.config + } + + /// Returns the selected backend. + #[must_use] + pub const fn kind(&self) -> DatabaseKind { + self.pool.kind() + } + + /// Returns the caller-owned pool. + #[must_use] + pub const fn pool(&self) -> &DatabasePool { + &self.pool + } +} + +impl fmt::Debug for PilotDatabase { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("PilotDatabase") + .field("config", &self.config) + .finish_non_exhaustive() + } +} + +/// A transaction an extension opened with [`PilotDatabase::begin`]. Dropping +/// it without committing rolls it back. +pub struct PilotTransaction(PilotTransactionInner); + +enum PilotTransactionInner { + #[cfg(feature = "postgres")] + Postgres(Transaction<'static, Postgres>), + #[cfg(feature = "sqlite")] + Sqlite(Transaction<'static, Sqlite>), +} + +impl PilotTransaction { + /// Commits the transaction. + /// + /// # Errors + /// + /// Returns the database error when the commit fails. + pub async fn commit(self) -> Result<(), Error> { + match self.0 { + #[cfg(feature = "postgres")] + PilotTransactionInner::Postgres(transaction) => transaction.commit().await?, + #[cfg(feature = "sqlite")] + PilotTransactionInner::Sqlite(transaction) => transaction.commit().await?, + } + Ok(()) + } + + /// Borrows the transaction's connection. + pub fn connection(&mut self) -> DatabaseConnection<'_> { + match &mut self.0 { + #[cfg(feature = "postgres")] + PilotTransactionInner::Postgres(transaction) => { + DatabaseConnection::Postgres(transaction) + } + #[cfg(feature = "sqlite")] + PilotTransactionInner::Sqlite(transaction) => DatabaseConnection::Sqlite(transaction), + } + } +} + +impl fmt::Debug for PilotTransaction { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + let kind = match &self.0 { + #[cfg(feature = "postgres")] + PilotTransactionInner::Postgres(_) => DatabaseKind::Postgres, + #[cfg(feature = "sqlite")] + PilotTransactionInner::Sqlite(_) => DatabaseKind::Sqlite, + }; + formatter + .debug_struct("PilotTransaction") + .field("kind", &kind) + .finish_non_exhaustive() + } +} + +/// A queue producer's configuration, as an extension's +/// [`PilotProducer`] sees it. +#[derive(Clone, Debug, PartialEq)] +#[non_exhaustive] +pub struct ProducerConfiguration { + /// Most jobs this client runs from the queue at once. + pub max_workers: usize, + /// The text of the queue's persisted metadata as the database renders + /// it: PostgreSQL's `metadata::text` or SQLite's `json(metadata)`. Unlike + /// the parsed `queue.metadata`, it keeps key order, duplicate keys, and + /// number literals, for decoding the metadata exactly like River Go's + /// `encoding/json`. + pub metadata_text: String, + /// The queue's persisted record, including its metadata and pause state. + pub queue: Queue, + /// The extension's settings for this queue, as configured with + /// [`QueueConfigExt::with_extension_setting`](crate::__private::QueueConfigExt::with_extension_setting) + /// and accepted by + /// [`Pilot::validate_queue_settings`](crate::__private::Pilot::validate_queue_settings). + pub settings: Map, +} + +/// Inputs to [`Pilot::start_producer`](crate::__private::Pilot::start_producer). +#[derive(Debug)] +#[non_exhaustive] +pub struct ProducerStartContext { + /// This client's identifier, recorded in `attempted_by` by claims. + pub client_id: String, + /// The producer's initial configuration. + pub configuration: ProducerConfiguration, + /// The client's database. + pub database: PilotDatabase, +} + +/// Inputs to one [`PilotProducer::keep_alive`]. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct ProducerKeepAliveContext { + /// Peers that haven't reported since this time are stale, like River + /// Go's `StaleUpdatedAtHorizon`. + pub stale_before: DateTime, +} + +/// Inputs to one [`PilotProducer::shutdown`] attempt. +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct ProducerShutdownContext { + /// One-based attempt number, up to four. + pub attempt: u32, + /// How long River waits for this attempt before dropping it and trying + /// again with a longer deadline. + pub timeout: Duration, +} + +/// Inputs to one [`PilotProducer::claim`]. +#[derive(Debug)] +#[non_exhaustive] +pub struct ProducerClaimContext<'a> { + /// This client's identifier. Claimed rows must end their `attempted_by` + /// with it. + pub client_id: &'a str, + /// Cancelled when the producer stops claiming. It ends retries and + /// backoff; a claim that already committed must still be returned. + pub claim_stop: &'a CancellationToken, + /// The client's database. + pub database: &'a PilotDatabase, + /// Kinds the claim is restricted to, including aliases, when the client + /// fetches only known kinds, or `None` to claim every kind. River's + /// standard claim applies it; a session's own claim must too. + pub kinds: Option<&'a [String]>, + /// Most jobs the claim may return. + pub limit: usize, + /// The queue being claimed from. + pub queue: &'a str, +} + +/// River's standard claim, handed to [`PilotProducer::claim`]. +/// +/// [`claim`](Self::claim) consumes it, so a session runs River's claim at most +/// once per call, on the transaction of its choosing. +pub struct ProducerClaimNext<'a> { + inner: &'a ClientInner, + limit: usize, + queue: &'a str, +} + +impl<'a> ProducerClaimNext<'a> { + pub(crate) const fn new(inner: &'a ClientInner, queue: &'a str, limit: usize) -> Self { + Self { + inner, + limit, + queue, + } + } + + /// Claims up to the claim's limit of available jobs on `connection`, + /// exactly as River claims them without an extension. The claim takes + /// effect when the connection's transaction commits. + /// + /// # Errors + /// + /// Returns the database error when the claim fails. + pub async fn claim(self, connection: DatabaseConnection<'_>) -> Result, Error> { + Ok( + crate::client::standard_claim(self.inner, connection, self.queue, self.limit) + .await? + .into_iter() + .map(ClaimedJob::from_decoded) + .collect(), + ) + } +} + +impl fmt::Debug for ProducerClaimNext<'_> { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter + .debug_struct("ProducerClaimNext") + .field("limit", &self.limit) + .field("queue", &self.queue) + .finish_non_exhaustive() + } +} + +/// An extension's state for one generation of one queue's producer, like +/// River Go's pilot `ProducerState`. +/// +/// River creates a session with +/// [`Pilot::start_producer`](crate::__private::Pilot::start_producer) when a +/// producer starts, before its first claim, and drops it once the producer +/// has stopped. A session is never shared between generations: removing and +/// adding a queue again, or restarting a failed producer, starts a new one. +/// +/// Calls into one session follow these rules: +/// +/// - At most one [`claim`](Self::claim) is in flight at a time. +/// - [`configuration_changed`](Self::configuration_changed) runs between +/// claims, never during one. +/// - [`job_finished`](Self::job_finished) may run at any time, including while +/// a claim is in flight, and runs once for every row a claim returned that +/// River accepted. +/// - [`keep_alive`](Self::keep_alive) runs at River's producer report +/// interval, never overlapping another report, and may overlap a claim. It +/// keeps running while the producer drains after it stops claiming. +/// - [`shutdown`](Self::shutdown) runs once the last attempt has left the +/// producer and reporting has stopped. No other call follows it. +/// +/// A producer stops claiming when the client stops or the queue is removed. +/// It then drains its running attempts and reports until they finish, so +/// peers keep counting them, before it shuts the session down. +/// +/// A panic in [`claim`](Self::claim), +/// [`configuration_changed`](Self::configuration_changed), or +/// [`job_finished`](Self::job_finished) stops the client like a broken claim: +/// the producer cancels and drains its attempts, still calls `job_finished` +/// for each, stops reporting, and shuts the session down. A panic in +/// [`keep_alive`](Self::keep_alive) or [`shutdown`](Self::shutdown) is +/// logged and handled like an error from it. +#[async_trait] +pub trait PilotProducer: Send + Sync + 'static { + /// Whether River claims through [`PilotProducer::claim`]. When `false`, + /// River claims with its own statement and no transaction. + fn intercepts_claim(&self) -> bool { + false + } + + /// Claims up to `context.limit` jobs, like River Go's pilot + /// `JobGetAvailable`. + /// + /// The session owns the transaction boundary: it opens each transaction, + /// runs [`ProducerClaimNext::claim`] or its own claim on it, and commits. + /// It may retry a failed attempt in a new transaction, releasing the + /// connection between attempts, until `context.claim_stop` is cancelled. + /// It returns only committed rows, and must undo any tentative + /// bookkeeping itself when it returns an error or when its future is + /// dropped, which River does only when its runtime shuts down. + /// + /// River checks the returned rows before working them: they must + /// identify distinct jobs, at most `context.limit` of them, running, in + /// this queue, and last attempted by this client. A partly decoded row is + /// checked with the fields that could be decoded. A result that breaks those + /// rules is a protocol error that stops the client; its rows are left for + /// the rescuer. + /// + /// The default opens one transaction and runs River's claim in it. + /// + /// # Errors + /// + /// Returns an error when nothing was claimed. River logs it and tries + /// again after the queue's fetch cooldown. + async fn claim( + &self, + context: ProducerClaimContext<'_>, + next: ProducerClaimNext<'_>, + ) -> Result, PilotError> { + let mut transaction = context.database.begin().await?; + let jobs = next.claim(transaction.connection()).await?; + transaction.commit().await?; + Ok(jobs) + } + + /// Replaces the session's configuration, between claims. + /// + /// River calls it when the queue's persisted record changes, such as its + /// metadata, including a change only to the metadata's stored text, or + /// its pause state, and when this client's configuration of the + /// queue changes through + /// [`LocalQueues::update`](crate::LocalQueues::update). It must not block + /// or perform I/O. + fn configuration_changed(&self, _configuration: &ProducerConfiguration) {} + + /// Reports that a claimed job's attempt left the producer, like River Go's + /// `ProducerState.JobFinish`. + /// + /// River calls it once for each accepted claimed row, with that row as + /// claimed, when its attempt exits: after its result is handed to the + /// completer, or when its attempt is abandoned, for example because its + /// worker outlived an abort during shutdown or its result couldn't be + /// handed off. It doesn't wait for the result to be persisted. It must not + /// block or perform I/O. + fn job_finished(&self, _job: &JobRow) {} + + /// Reports that the producer is alive, like River Go's pilot + /// `ProducerKeepAlive`. + /// + /// River calls it after a random delay of up to a second, then at the + /// client's producer report interval, 30 seconds by default, including + /// while the producer drains. River drops a call that runs longer than + /// ten seconds. + /// + /// # Errors + /// + /// Returns an error when the report failed. River logs it and reports + /// again at the next interval. + async fn keep_alive(&self, _context: ProducerKeepAliveContext) -> Result<(), PilotError> { + Ok(()) + } + + /// Releases the session's shared state once the producer has stopped, + /// like River Go's pilot `ProducerShutdown`. + /// + /// River makes up to four attempts, one at a time, with deadlines of + /// 100 milliseconds, 500 milliseconds, 2.5 seconds, and 12.5 seconds, + /// dropping an attempt when its deadline passes, and logs the failure + /// when every attempt fails. + /// + /// # Errors + /// + /// Returns an error when this attempt failed and another may succeed. + async fn shutdown(&self, _context: ProducerShutdownContext) -> Result<(), PilotError> { + Ok(()) + } +} + +/// Checks the rows a session claimed before River works them. +pub(crate) fn validate_claimed( + claimed: &[ClaimedJob], + client_id: &str, + queue: &str, + limit: usize, +) -> Result<(), String> { + if claimed.len() > limit { + return Err(format!( + "claim returned {} jobs, more than its limit of {limit}", + claimed.len() + )); + } + let mut ids = std::collections::HashSet::with_capacity(claimed.len()); + for job in claimed { + // A row River can't identify could never be finished, so the + // extension's accounting for it would leak. + let Some(row) = job.row() else { + return Err(format!( + "claim returned a row that couldn't be identified: {}", + job.decode_error().unwrap_or_default() + )); + }; + let id = row.id; + if !ids.insert(id) { + return Err(format!("claim returned job {id} more than once")); + } + // A partly decoded row still has its state and queue, which River + // always decodes. Its `attempted_by` is empty when that column is + // what couldn't be decoded, and then can't be checked. + if row.state != crate::JobState::Running { + return Err(format!("claim returned job {id} in state {}", row.state)); + } + if row.queue != queue { + return Err(format!( + "claim for queue {queue:?} returned job {id} from queue {:?}", + row.queue + )); + } + let attempted_by_undecodable = job.column_undecodable("attempted_by"); + if !attempted_by_undecodable + && row.attempted_by.last().map(String::as_str) != Some(client_id) + { + return Err(format!( + "claim returned job {id} not last attempted by this client" + )); + } + } + Ok(()) +} + +/// A producer session River runs for a queue generation. +pub(crate) type SharedProducer = Arc; diff --git a/rust/riverqueue/src/protocol.rs b/rust/riverqueue/src/protocol.rs new file mode 100644 index 000000000..82f89e6ec --- /dev/null +++ b/rust/riverqueue/src/protocol.rs @@ -0,0 +1,80 @@ +//! Wire-protocol values shared by every River client. +//! +//! Applications don't need these to insert or work jobs. They're useful for +//! tools that interoperate with River's tables directly, such as computing +//! the unique key River would assign to a job or listening for River's +//! notifications. + +use chrono::{DateTime, Utc}; +use serde_json::value::RawValue; + +use crate::{Error, JobState, UniqueOpts}; + +/// Notification topic for queue and job control messages. +pub const NOTIFICATION_TOPIC_CONTROL: &str = "river_control"; + +/// Notification topic for newly available jobs. +pub const NOTIFICATION_TOPIC_INSERT: &str = "river_insert"; + +/// Notification topic for leadership changes. +pub const NOTIFICATION_TOPIC_LEADERSHIP: &str = "river_leadership"; + +/// Inputs used to compute a job's unique key. +#[derive(Clone, Copy, Debug)] +pub struct UniqueKeyInput<'a> { + /// Encoded arguments exactly as they will be stored, for example from + /// [`encode_args`](crate::encoding::encode_args). + pub encoded_args: &'a RawValue, + /// Job kind. + pub kind: &'a str, + /// Current time, used for period-scoped uniqueness when `scheduled_at` is + /// absent. + pub now: DateTime, + /// Uniqueness options. + pub opts: &'a UniqueOpts, + /// Queue name. + pub queue: &'a str, + /// Scheduled time, if the job is scheduled. + pub scheduled_at: Option>, + /// Literal argument path components selected for uniqueness, such as + /// [`JobArgs::unique_fields`](crate::JobArgs::unique_fields). For example, + /// `&[&["user.id"], &["user", "id"]]` selects two distinct fields. + /// When empty, every top-level argument participates. + pub unique_fields: &'a [&'a [&'a str]], +} + +/// Computes the SHA-256 unique key River stores in `river_job.unique_key` +/// for these inputs. Returns `None` when `opts` enables no uniqueness +/// dimension. +/// +/// # Errors +/// +/// Returns an error when the options or selected argument paths are invalid. +pub fn unique_key(input: &UniqueKeyInput<'_>) -> Result, Error> { + crate::unique::build_unique_key_parts( + input.kind, + input.unique_fields, + input.encoded_args, + input.now, + input.opts, + input.queue, + input.scheduled_at, + ) +} + +/// Bit representing `state` in `river_job.unique_states`. +#[must_use] +pub const fn unique_state_bit(state: JobState) -> u8 { + state.unique_bit() +} + +/// Value River stores in `river_job.unique_states` for `opts`. +#[must_use] +pub fn unique_states_bitmask(opts: &UniqueOpts) -> u8 { + opts.state_bitmask() +} + +/// An insert notification's payload, `{"queue":""}`. +pub(crate) fn insert_notification_payload(queue: &str) -> String { + serde_json::json!({ "queue": queue }).to_string() +} diff --git a/rust/riverqueue/src/query.rs b/rust/riverqueue/src/query.rs new file mode 100644 index 000000000..69e79a4ad --- /dev/null +++ b/rust/riverqueue/src/query.rs @@ -0,0 +1,931 @@ +//! Job querying and update parameters. + +use std::{fmt, str::FromStr}; + +use base64::{ + Engine as _, alphabet, + engine::{DecodePaddingMode, GeneralPurpose, GeneralPurposeConfig, general_purpose}, +}; +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; +use thiserror::Error; + +use crate::{JobRow, JobState}; + +/// Stable keyset cursor for job-list pagination. +#[derive(Clone, Debug)] +pub struct JobListCursor { + pub(crate) id: i64, + kind: String, + order_by: JobListOrderBy, + queue: String, + /// Value of the time field the list is ordered by, or `None` when the + /// list is ordered by ID alone or the job's value is null. + pub(crate) sort_time: Option>, +} + +impl JobListCursor { + /// Builds a cursor from a returned row and the parameters used to list it. + /// + /// [`JobListResult::last_cursor`] already holds the cursor after a page's + /// last job. + /// + /// # Errors + /// + /// Returns [`JobListCursorError::InvalidListParams`] when `params` are + /// invalid. + pub fn from_job(job: &JobRow, params: &JobListParams) -> Result { + params + .validate() + .map_err(JobListCursorError::InvalidListParams)?; + Ok(Self::after_job(job, params)) + } + + /// Builds the cursor after `job` for already validated parameters. + pub(crate) fn after_job(job: &JobRow, params: &JobListParams) -> Self { + // The time comes from the field the list is ordered by. For time + // ordering over several states, that's the first state's field for + // every job, which may differ from the field of the job's own state. + Self { + id: job.id, + kind: job.kind.clone(), + order_by: params.order_by, + queue: job.queue.clone(), + sort_time: params.time_field().and_then(|field| field.value(job)), + } + } + + /// Decodes an opaque cursor emitted by either matched implementation. + /// + /// # Errors + /// + /// Returns an error when `encoded` isn't a cursor River emitted. + pub fn decode(encoded: &str) -> Result { + encoded.parse() + } + + /// Encodes this cursor for storage in an API pagination token. + /// + /// The text is the cursor format every River client reads and writes: + /// padded URL-safe Base64 of the cursor's JSON, in the + /// [`encoding`](crate::encoding) River uses for job arguments. A cursor + /// without a time, because its list is ordered by ID or its job's time + /// field is null, carries `0001-01-01T00:00:00Z`. + /// + /// # Panics + /// + /// Panics only if Serde cannot serialize River's fixed, internally + /// constructed cursor representation. Its fields have no fallible custom + /// serializers, so this indicates a River implementation bug. + pub fn encode(&self) -> String { + let value = JobListCursorValue { + id: self.id, + kind: self.kind.clone(), + queue: self.queue.clone(), + sort_field: self.order_by.as_str().to_owned(), + time: self.sort_time.unwrap_or_else(go_zero_time), + }; + let json = crate::encoding::encode_args(&value) + .expect("fixed job-list cursor value always serializes"); + general_purpose::URL_SAFE.encode(json.get()) + } +} + +impl fmt::Display for JobListCursor { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(&self.encode()) + } +} + +impl FromStr for JobListCursor { + type Err = JobListCursorError; + + /// Accepts URL-safe or standard Base64, with or without padding. + fn from_str(encoded: &str) -> Result { + const TOLERANT: GeneralPurposeConfig = + GeneralPurposeConfig::new().with_decode_padding_mode(DecodePaddingMode::Indifferent); + const URL_SAFE: GeneralPurpose = GeneralPurpose::new(&alphabet::URL_SAFE, TOLERANT); + const STANDARD: GeneralPurpose = GeneralPurpose::new(&alphabet::STANDARD, TOLERANT); + + let bytes = URL_SAFE + .decode(encoded) + .or_else(|_| STANDARD.decode(encoded)) + .map_err(JobListCursorError::Base64)?; + let value: JobListCursorValue = + serde_json::from_slice(&bytes).map_err(JobListCursorError::Json)?; + let order_by = value.sort_field.parse()?; + Ok(Self { + id: value.id, + kind: value.kind, + order_by, + queue: value.queue, + // The zero time stands for no time. No job's time field holds it. + sort_time: (order_by != JobListOrderBy::Id && value.time != go_zero_time()) + .then_some(value.time), + }) + } +} + +/// Failure to build or decode a job-list cursor. +#[derive(Debug, Error)] +#[non_exhaustive] +pub enum JobListCursorError { + /// Cursor text is neither URL-safe nor standard Base64. + #[error("invalid job-list cursor base64")] + Base64(#[source] base64::DecodeError), + /// List parameters cannot produce a valid cursor. + #[error("invalid job-list parameters: {0}")] + InvalidListParams(String), + /// Cursor contents are not valid JSON. + #[error("invalid job-list cursor JSON")] + Json(#[source] serde_json::Error), + /// Cursor names an unsupported ordering field. + #[error("unknown job-list cursor sort field {0:?}")] + UnknownSortField(String), +} + +#[derive(Deserialize, Serialize)] +struct JobListCursorValue { + id: i64, + kind: String, + queue: String, + sort_field: String, + #[serde(with = "crate::encoding::go_time")] + time: DateTime, +} + +/// Field used for stable job-list ordering. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum JobListOrderBy { + /// Finalization time followed by ID. + FinalizedAt, + /// Database ID only. + #[default] + Id, + /// Scheduled time followed by ID. + ScheduledAt, + /// State-appropriate time followed by ID. + /// + /// The time field is chosen by the first state given to + /// [`JobListParams::states`], and jobs in every listed state are ordered + /// by that field: + /// + /// * `available`, `pending`, `retryable`, and `scheduled` use + /// `scheduled_at`. + /// * `running` uses `attempted_at`. + /// * `cancelled`, `completed`, and `discarded` use `finalized_at`. + /// + /// Without a state filter, jobs are ordered by `scheduled_at`. Jobs whose + /// field is null, like `finalized_at` for jobs that haven't finalized, + /// come after all others in ascending order and before all others in + /// descending order. + Time, +} + +impl JobListOrderBy { + /// Returns the cross-language wire value. + #[must_use] + pub const fn as_str(self) -> &'static str { + match self { + Self::FinalizedAt => "finalized_at", + Self::Id => "id", + Self::ScheduledAt => "scheduled_at", + Self::Time => "time", + } + } +} + +impl fmt::Display for JobListOrderBy { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str(self.as_str()) + } +} + +impl FromStr for JobListOrderBy { + type Err = JobListCursorError; + + fn from_str(value: &str) -> Result { + match value { + "finalized_at" => Ok(Self::FinalizedAt), + "id" => Ok(Self::Id), + "scheduled_at" => Ok(Self::ScheduledAt), + "time" => Ok(Self::Time), + _ => Err(JobListCursorError::UnknownSortField(value.to_owned())), + } + } +} + +/// Direction used for job-list ordering and cursor comparison. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum SortDirection { + /// Ascending order. + #[default] + Ascending, + /// Descending order. + Descending, +} + +/// Safe filters for bulk job deletion. +#[derive(Clone, Debug)] +pub struct JobDeleteManyParams { + pub(crate) all: bool, + pub(crate) filter: JobListParams, +} + +impl JobDeleteManyParams { + /// Explicitly selects every non-running job, subject to the filter limit. + #[must_use] + pub fn all() -> Self { + Self { + all: true, + filter: JobListParams::default(), + } + } + + /// Selects jobs matching a nonempty filter. + #[must_use] + pub const fn matching(filter: JobListParams) -> Self { + Self { all: false, filter } + } + + /// Returns whether this operation explicitly selects every job. + #[must_use] + pub const fn deletes_all(&self) -> bool { + self.all + } + + /// Returns the row filter and deletion limit. + #[must_use] + pub const fn filter(&self) -> &JobListParams { + &self.filter + } +} + +/// A page of jobs returned by [`Jobs::list`](crate::Jobs::list). +#[derive(Clone, Debug)] +#[non_exhaustive] +pub struct JobListResult { + /// Jobs in the requested order. + pub jobs: Vec, + /// Cursor after the last job, to request the next page by passing it + /// to [`JobListParams::after`] with otherwise identical parameters. + /// `None` when the page is empty. + pub last_cursor: Option, +} + +/// Filters and pagination for listing jobs. +/// +/// Filters combine with AND, and each list-valued filter matches any of its +/// values. Unset filters match every job: +/// +/// ``` +/// # use riverqueue::{JobListOrderBy, JobListParams, JobState, SortDirection}; +/// let params = JobListParams::default() +/// .queues(["email"]) +/// .states([JobState::Completed]) +/// .order_by(JobListOrderBy::Time) +/// .direction(SortDirection::Descending) +/// .limit(50); +/// ``` +#[derive(Clone, Debug)] +pub struct JobListParams { + pub(crate) direction: SortDirection, + pub(crate) ids: Vec, + pub(crate) kinds: Vec, + pub(crate) limit: u32, + pub(crate) metadata: Option>, + pub(crate) order_by: JobListOrderBy, + pub(crate) priorities: Vec, + pub(crate) queues: Vec, + pub(crate) start: Option, + pub(crate) states: Vec, + pub(crate) tags_all: Vec, + pub(crate) tags_any: Vec, +} + +/// Where a job listing starts. A listing continues either from a keyset +/// cursor or after an ID, never both. +#[derive(Clone, Debug)] +pub(crate) enum JobListStart { + Cursor(JobListCursor), + Id(i64), +} + +impl Default for JobListParams { + fn default() -> Self { + Self { + direction: SortDirection::Ascending, + ids: Vec::new(), + kinds: Vec::new(), + limit: 100, + metadata: None, + order_by: JobListOrderBy::Id, + priorities: Vec::new(), + queues: Vec::new(), + start: None, + states: Vec::new(), + tags_all: Vec::new(), + tags_any: Vec::new(), + } + } +} + +impl JobListParams { + /// Returns jobs after a cursor from a previous page, usually that page's + /// [`JobListResult::last_cursor`]. The cursor must come from a listing + /// with the same ordering. Replaces any [`after_id`](Self::after_id). + #[must_use] + pub fn after(mut self, cursor: JobListCursor) -> Self { + self.start = Some(JobListStart::Cursor(cursor)); + self + } + + /// Returns jobs whose ID comes after `id` in the sort direction. + /// Replaces any [`after`](Self::after) cursor. + #[must_use] + pub fn after_id(mut self, id: i64) -> Self { + self.start = Some(JobListStart::Id(id)); + self + } + + /// Sets the sort and cursor direction. Defaults to ascending. + #[must_use] + pub const fn direction(mut self, direction: SortDirection) -> Self { + self.direction = direction; + self + } + + /// Matches jobs with any of these IDs. + #[must_use] + pub fn ids(mut self, ids: impl IntoIterator) -> Self { + self.ids = ids.into_iter().collect(); + self + } + + /// Matches jobs of any of these kinds. + #[must_use] + pub fn kinds(mut self, kinds: impl IntoIterator>) -> Self { + self.kinds = kinds.into_iter().map(Into::into).collect(); + self + } + + /// Sets the maximum number of jobs returned, from one through 10,000. + /// Defaults to 100. Listing fails with a limit outside that range. + #[must_use] + pub const fn limit(mut self, limit: u32) -> Self { + self.limit = limit; + self + } + + /// Matches jobs whose metadata contains this JSON object, like + /// PostgreSQL's `@>` operator. + #[must_use] + pub fn metadata(mut self, metadata: Map) -> Self { + self.metadata = Some(metadata); + self + } + + /// Sets the field jobs are ordered by, with ID breaking ties. Defaults + /// to [`JobListOrderBy::Id`]. + #[must_use] + pub const fn order_by(mut self, order_by: JobListOrderBy) -> Self { + self.order_by = order_by; + self + } + + /// Matches jobs with any of these priorities. + #[must_use] + pub fn priorities(mut self, priorities: impl IntoIterator) -> Self { + self.priorities = priorities.into_iter().collect(); + self + } + + /// Matches jobs in any of these queues. + #[must_use] + pub fn queues(mut self, queues: impl IntoIterator>) -> Self { + self.queues = queues.into_iter().map(Into::into).collect(); + self + } + + /// Matches jobs in any of these states. + #[must_use] + pub fn states(mut self, states: impl IntoIterator) -> Self { + self.states = states.into_iter().collect(); + self + } + + /// Matches jobs that have every one of these tags. + #[must_use] + pub fn tags_all(mut self, tags: impl IntoIterator>) -> Self { + self.tags_all = tags.into_iter().map(Into::into).collect(); + self + } + + /// Matches jobs that have at least one of these tags. + #[must_use] + pub fn tags_any(mut self, tags: impl IntoIterator>) -> Self { + self.tags_any = tags.into_iter().map(Into::into).collect(); + self + } + + /// Whether at least one narrowing predicate was supplied. + #[must_use] + pub fn has_filter(&self) -> bool { + self.start.is_some() + || !self.ids.is_empty() + || !self.kinds.is_empty() + || self.metadata.is_some() + || !self.priorities.is_empty() + || !self.queues.is_empty() + || !self.states.is_empty() + || !self.tags_all.is_empty() + || !self.tags_any.is_empty() + } + + /// Returns the keyset cursor to continue after, if any. + pub(crate) const fn cursor(&self) -> Option<&JobListCursor> { + match &self.start { + Some(JobListStart::Cursor(cursor)) => Some(cursor), + Some(JobListStart::Id(_)) | None => None, + } + } + + /// Returns how the list is ordered and where it resumes, for the storage + /// backends to render as SQL. + pub(crate) fn keyset(&self) -> JobListKeyset { + let time_field = self.time_field(); + let nullable = time_field.is_some_and(|field| self.time_field_nullable(field)); + let after = self.start.as_ref().map(|start| match start { + JobListStart::Id(id) => JobListAfter::Id(*id), + JobListStart::Cursor(cursor) => match (time_field, cursor.sort_time) { + (Some(_), Some(time)) => JobListAfter::Time { + id: cursor.id, + time, + }, + (Some(_), None) if nullable => JobListAfter::NullTime { id: cursor.id }, + // Like Go, a cursor without a time for a field that can't be + // null resumes by ID. + _ => JobListAfter::Id(cursor.id), + }, + }); + JobListKeyset { + after, + direction: self.direction, + nullable, + time_field, + } + } + + /// Returns the time field jobs are ordered by before ID, or `None` when + /// ordering by ID alone. + fn time_field(&self) -> Option { + match self.order_by { + JobListOrderBy::FinalizedAt => Some(JobListTimeField::Finalized), + JobListOrderBy::Id => None, + JobListOrderBy::ScheduledAt => Some(JobListTimeField::Scheduled), + // No state filter lists every state. Like Go, whose default + // states start with `available`, use that state's field. + JobListOrderBy::Time => Some( + self.states + .first() + .map_or(JobListTimeField::Scheduled, |state| { + JobListTimeField::for_state(*state) + }), + ), + } + } + + /// Returns whether `field` may be null for listed jobs. + fn time_field_nullable(&self, field: JobListTimeField) -> bool { + match field { + JobListTimeField::Attempted => true, + // The schema requires `finalized_at` for exactly the finalized + // states. Unlike Go's custom SQL conditions, no Rust filter can + // widen the state filter, so other filters don't matter here. + JobListTimeField::Finalized => { + self.states.is_empty() + || self.states.iter().any(|state| { + !matches!( + state, + JobState::Cancelled | JobState::Completed | JobState::Discarded + ) + }) + } + JobListTimeField::Scheduled => false, + } + } + + pub(crate) fn validate(&self) -> Result<(), String> { + if !(1..=10_000).contains(&self.limit) { + return Err("job list limit must be between 1 and 10000".to_owned()); + } + if let Some(cursor) = self.cursor() + && cursor.order_by != self.order_by + { + return Err("job list cursor sort field does not match list ordering".to_owned()); + } + if self.order_by == JobListOrderBy::FinalizedAt + && (self.states.is_empty() + || self.states.iter().any(|state| { + !matches!( + state, + JobState::Cancelled | JobState::Completed | JobState::Discarded + ) + })) + { + return Err( + "finalized_at ordering requires only cancelled, completed, or discarded states" + .to_owned(), + ); + } + Ok(()) + } +} + +/// Where a job list resumes, relative to its ordering. +#[derive(Clone, Copy, Debug, PartialEq)] +pub(crate) enum JobListAfter { + /// Jobs whose ID follows this one in the list direction. + Id(i64), + /// Jobs following a cursor job whose time field is null. Nulls sort + /// last ascending and first descending. + NullTime { id: i64 }, + /// Jobs following a cursor job with this time. + Time { id: i64, time: DateTime }, +} + +/// How a validated job list is ordered and where it resumes. Each storage +/// backend renders it as SQL, so they order and page identically. +#[derive(Clone, Copy, Debug, Default, PartialEq)] +pub(crate) struct JobListKeyset { + pub(crate) after: Option, + pub(crate) direction: SortDirection, + /// Whether the time field may be null for listed jobs. Nulls then sort + /// explicitly last ascending and first descending, the PostgreSQL + /// default, so that every backend agrees and cursors can match them. + pub(crate) nullable: bool, + /// Time field ordered before ID, or `None` to order by ID alone. + pub(crate) time_field: Option, +} + +impl JobListKeyset { + /// Returns the ID of the job the list resumes after. + pub(crate) const fn after_id(&self) -> Option { + match self.after { + Some( + JobListAfter::Id(id) + | JobListAfter::NullTime { id } + | JobListAfter::Time { id, .. }, + ) => Some(id), + None => None, + } + } + + /// Returns the time of the job the list resumes after, if that job has + /// one. + pub(crate) const fn after_time(&self) -> Option> { + match self.after { + Some(JobListAfter::Time { time, .. }) => Some(time), + Some(JobListAfter::Id(_) | JobListAfter::NullTime { .. }) | None => None, + } + } + + /// Returns the condition selecting rows after the cursor, or `None` to + /// start from the beginning. Backends replace the placeholders with + /// their parameters for [`after_time`](Self::after_time) and + /// [`after_id`](Self::after_id). + pub(crate) fn after_sql(&self) -> Option> { + use JobListSqlPart::{AfterId, AfterTime, Sql}; + + let comparison = match self.direction { + SortDirection::Ascending => ">", + SortDirection::Descending => "<", + }; + let id_only = || vec![Sql(format!("id {comparison} ")), AfterId]; + let after = self.after?; + let Some(field) = self.time_field.map(JobListTimeField::column) else { + return Some(id_only()); + }; + Some(match (after, self.direction) { + (JobListAfter::Id(_), _) => id_only(), + // After a null time, only nulls with a later ID follow ascending, + // and every non-null time also follows descending. + (JobListAfter::NullTime { .. }, SortDirection::Ascending) => { + vec![ + Sql(format!("({field} IS NULL AND id > ")), + AfterId, + Sql(")".to_owned()), + ] + } + (JobListAfter::NullTime { .. }, SortDirection::Descending) => vec![ + Sql(format!("({field} IS NOT NULL OR id < ")), + AfterId, + Sql(")".to_owned()), + ], + (JobListAfter::Time { .. }, direction) => { + // Nulls follow every time ascending and precede every time + // descending. + let or_null = if self.nullable && direction == SortDirection::Ascending { + format!(" OR {field} IS NULL") + } else { + String::new() + }; + vec![ + Sql(format!("({field} {comparison} ")), + AfterTime, + Sql(format!(" OR ({field} = ")), + AfterTime, + Sql(format!(" AND id {comparison} ")), + AfterId, + Sql(format!("){or_null})")), + ] + } + }) + } + + /// Returns the `ORDER BY` clause, without the keyword. + pub(crate) fn order_sql(&self) -> String { + let direction = match self.direction { + SortDirection::Ascending => "ASC", + SortDirection::Descending => "DESC", + }; + let Some(field) = self.time_field else { + return format!("id {direction}"); + }; + let nulls = match (self.nullable, self.direction) { + (false, _) => "", + (true, SortDirection::Ascending) => " NULLS LAST", + (true, SortDirection::Descending) => " NULLS FIRST", + }; + format!("{} {direction}{nulls}, id {direction}", field.column()) + } +} + +/// A piece of a job list cursor condition. +#[derive(Clone, Debug, Eq, PartialEq)] +pub(crate) enum JobListSqlPart { + /// Placeholder for the ID of the job the list resumes after. + AfterId, + /// Placeholder for the time of the job the list resumes after. + AfterTime, + /// SQL text. + Sql(String), +} + +/// Time field a job list orders by before its ID tiebreaker. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum JobListTimeField { + Attempted, + Finalized, + Scheduled, +} + +impl JobListTimeField { + /// Returns the field that [`JobListOrderBy::Time`] uses for `state`. + const fn for_state(state: JobState) -> Self { + match state { + JobState::Available | JobState::Pending | JobState::Retryable | JobState::Scheduled => { + Self::Scheduled + } + JobState::Running => Self::Attempted, + JobState::Cancelled | JobState::Completed | JobState::Discarded => Self::Finalized, + } + } + + /// Returns the job table column. + pub(crate) const fn column(self) -> &'static str { + match self { + Self::Attempted => "attempted_at", + Self::Finalized => "finalized_at", + Self::Scheduled => "scheduled_at", + } + } + + /// Returns the field's value on `job`, or `None` if it's null. + const fn value(self, job: &JobRow) -> Option> { + match self { + Self::Attempted => job.attempted_at, + Self::Finalized => job.finalized_at, + Self::Scheduled => Some(job.scheduled_at), + } + } +} + +fn go_zero_time() -> DateTime { + DateTime::parse_from_rfc3339("0001-01-01T00:00:00Z") + .expect("Go zero time is valid RFC 3339") + .with_timezone(&Utc) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Texts from River Go's `JobListCursor.MarshalText`. + const GO_CURSORS: [&str; 3] = [ + "eyJpZCI6NDIsImtpbmQiOiJzZW5kX2VtYWlsIiwicXVldWUiOiJwcmlvcml0eSIsInNvcnRfZmllbGQiOiJzY2hlZHVsZWRfYXQiLCJ0aW1lIjoiMjAyNi0wMS0wMlQwMzowNDowNS42Nzg5WiJ9", + "eyJpZCI6Nywia2luZCI6ImFcdTAwM2NiXHUwMDNlXHUwMDI2Y1x1MjAyOCIsInF1ZXVlIjoiZGVmYXVsdCIsInNvcnRfZmllbGQiOiJpZCIsInRpbWUiOiIwMDAxLTAxLTAxVDAwOjAwOjAwWiJ9", + "eyJpZCI6OTAwNzE5OTI1NDc0MDk5Mywia2luZCI6ImNvbmZvcm1hbmNlX2N1cnNvcn5-fiIsInF1ZXVlIjoiZGVmYXVsdCIsInNvcnRfZmllbGQiOiJmaW5hbGl6ZWRfYXQiLCJ0aW1lIjoiMjAyNi0wMS0wMlQwMzowNDowNS4xMloifQ==", + ]; + + fn utc(text: &str) -> DateTime { + DateTime::parse_from_rfc3339(text) + .unwrap() + .with_timezone(&Utc) + } + + #[test] + fn cursor_encodes_go_text_byte_for_byte() { + let cursors = [ + JobListCursor { + id: 42, + kind: "send_email".to_owned(), + order_by: JobListOrderBy::ScheduledAt, + queue: "priority".to_owned(), + sort_time: Some(utc("2026-01-02T03:04:05.6789Z")), + }, + JobListCursor { + id: 7, + kind: "a&c\u{2028}".to_owned(), + order_by: JobListOrderBy::Id, + queue: "default".to_owned(), + sort_time: None, + }, + JobListCursor { + id: 9_007_199_254_740_993, + kind: "conformance_cursor~~~".to_owned(), + order_by: JobListOrderBy::FinalizedAt, + queue: "default".to_owned(), + sort_time: Some(utc("2026-01-02T03:04:05.12Z")), + }, + ]; + for (cursor, go_text) in cursors.iter().zip(GO_CURSORS) { + assert_eq!(cursor.encode(), go_text); + let decoded = JobListCursor::decode(go_text).unwrap(); + assert_eq!(decoded.encode(), go_text); + assert_eq!(decoded.kind, cursor.kind); + assert_eq!(decoded.sort_time, cursor.sort_time); + } + } + + #[test] + fn cursor_decodes_either_alphabet_with_or_without_padding() { + let go_text = GO_CURSORS[2]; + assert!(go_text.contains('-') && go_text.ends_with("==")); + let standard = go_text.replace('-', "+").replace('_', "/"); + for text in [ + go_text, + go_text.trim_end_matches('='), + &standard, + standard.trim_end_matches('='), + ] { + let decoded = JobListCursor::decode(text).unwrap(); + assert_eq!(decoded.encode(), go_text, "{text}"); + } + assert!(matches!( + JobListCursor::decode("not base64!"), + Err(JobListCursorError::Base64(_)) + )); + } + + #[test] + fn cursor_round_trips_go_compatible_text() { + let cursor = JobListCursor { + id: 42, + kind: "send_email".to_owned(), + order_by: JobListOrderBy::ScheduledAt, + queue: "priority".to_owned(), + sort_time: Some( + DateTime::parse_from_rfc3339("2026-01-02T03:04:05.6789Z") + .unwrap() + .with_timezone(&Utc), + ), + }; + + let encoded = cursor.encode(); + let decoded = JobListCursor::decode(&encoded).unwrap(); + assert_eq!(decoded.id, cursor.id); + assert_eq!(decoded.kind, cursor.kind); + assert_eq!(decoded.order_by, cursor.order_by); + assert_eq!(decoded.queue, cursor.queue); + assert_eq!(decoded.sort_time, cursor.sort_time); + } + + #[test] + fn zero_time_cursor_keeps_nullable_time_ordering() { + let cursor = JobListCursor { + id: 42, + kind: "send_email".to_owned(), + order_by: JobListOrderBy::Time, + queue: "priority".to_owned(), + sort_time: None, + }; + let decoded = JobListCursor::decode(&cursor.encode()).unwrap(); + assert_eq!(decoded.sort_time, None); + + for (direction, expected_order) in [ + ( + SortDirection::Ascending, + "attempted_at ASC NULLS LAST, id ASC", + ), + ( + SortDirection::Descending, + "attempted_at DESC NULLS FIRST, id DESC", + ), + ] { + let keyset = JobListParams::default() + .states([JobState::Running, JobState::Available]) + .order_by(JobListOrderBy::Time) + .direction(direction) + .after(decoded.clone()) + .keyset(); + assert_eq!(keyset.after, Some(JobListAfter::NullTime { id: 42 })); + assert_eq!(keyset.order_sql(), expected_order); + assert!(keyset.after_sql().is_some()); + } + } + + #[test] + fn a_later_start_replaces_an_earlier_one() { + let cursor = JobListCursor::decode( + &JobListCursor { + id: 7, + kind: "kind".to_owned(), + order_by: JobListOrderBy::Id, + queue: "default".to_owned(), + sort_time: None, + } + .encode(), + ) + .unwrap(); + + let params = JobListParams::default().after(cursor.clone()).after_id(3); + assert_eq!(params.keyset().after_id(), Some(3)); + assert!(params.cursor().is_none()); + assert!(params.validate().is_ok()); + + let params = JobListParams::default().after_id(3).after(cursor); + assert_eq!(params.keyset().after_id(), Some(7)); + assert!(params.cursor().is_some()); + assert!(params.has_filter()); + } + + #[test] + fn limits_outside_the_supported_range_are_rejected() { + for limit in [0, 10_001, u32::MAX] { + assert_eq!( + JobListParams::default() + .limit(limit) + .validate() + .unwrap_err(), + "job list limit must be between 1 and 10000" + ); + } + assert!(JobListParams::default().limit(10_000).validate().is_ok()); + } + + #[test] + fn defaults_do_not_filter_states() { + let params = JobListParams::default(); + + assert_eq!(params.states, []); + assert!(!params.has_filter()); + } + + #[test] + fn finalized_order_requires_terminal_states() { + let params = JobListParams::default().order_by(JobListOrderBy::FinalizedAt); + + assert_eq!( + params.validate().unwrap_err(), + "finalized_at ordering requires only cancelled, completed, or discarded states" + ); + } +} + +/// Mutable public job fields. +#[derive(Clone, Debug, Default)] +#[non_exhaustive] +pub struct JobUpdateParams { + /// Merge these metadata keys into the existing object. + pub metadata: Map, + /// Set or replace the reserved job output value. + pub output: Option, +} + +impl JobUpdateParams { + /// Merges metadata keys into the existing object. + #[must_use] + pub fn metadata(mut self, metadata: Map) -> Self { + self.metadata = metadata; + self + } + + /// Sets or replaces the reserved output value. + #[must_use] + pub fn output(mut self, output: Value) -> Self { + self.output = Some(output); + self + } +} diff --git a/rust/riverqueue/src/queue.rs b/rust/riverqueue/src/queue.rs new file mode 100644 index 000000000..987648b46 --- /dev/null +++ b/rust/riverqueue/src/queue.rs @@ -0,0 +1,127 @@ +//! Persisted queue configuration. + +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +/// A queue currently or recently operated by a River client. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[non_exhaustive] +pub struct Queue { + /// Time at which this active queue record was created. + pub created_at: DateTime, + /// Reserved queue metadata. + pub metadata: Map, + /// The metadata's text as the database renders it, keeping the key + /// order, duplicate keys, and number literals the parsed map loses. + #[serde(skip)] + pub(crate) metadata_text: String, + /// Stable queue name. + pub name: String, + /// Time at which the queue was paused. + pub paused_at: Option>, + /// Last client heartbeat or configuration update. + pub updated_at: DateTime, +} + +/// Parameters for listing queues. +#[derive(Clone, Debug)] +pub struct QueueListParams { + pub(crate) limit: u32, +} + +impl Default for QueueListParams { + fn default() -> Self { + Self { limit: 100 } + } +} + +impl QueueListParams { + /// Sets the maximum number of queues returned, from one through 10,000. + /// Defaults to 100. Listing fails with a limit outside that range. + #[must_use] + pub const fn limit(mut self, limit: u32) -> Self { + self.limit = limit; + self + } +} + +/// The persisted queues that [`Queues::pause`](crate::Queues::pause) and +/// [`Queues::resume`](crate::Queues::resume) act on. +/// +/// Strings convert into [`Named`](Self::Named), so a queue can be passed by +/// name: +/// +/// ```no_run +/// # use riverqueue::QueueSelector; +/// # async fn example(client: riverqueue::Client) -> Result<(), riverqueue::Error> { +/// client.queues().pause("email").await?; +/// client.queues().resume(QueueSelector::All).await?; +/// # Ok(()) +/// # } +/// ``` +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +pub enum QueueSelector { + /// Every queue that has a persisted record. Succeeds even when there are + /// none. + All, + /// The queue with this name, which must have a persisted record. A name + /// is matched literally, so `Named("*")` names no queue. + Named(String), +} + +impl QueueSelector { + /// Returns the queue name River's storage and notification protocol use + /// for this selection, or `None` for a name no queue can have. + pub(crate) fn protocol_name(&self) -> Option<&str> { + match self { + Self::All => Some(crate::storage::QUEUE_ALL), + Self::Named(name) if name == crate::storage::QUEUE_ALL => None, + Self::Named(name) => Some(name), + } + } +} + +impl From<&str> for QueueSelector { + fn from(name: &str) -> Self { + Self::Named(name.to_owned()) + } +} + +impl From<&String> for QueueSelector { + fn from(name: &String) -> Self { + Self::Named(name.clone()) + } +} + +impl From for QueueSelector { + fn from(name: String) -> Self { + Self::Named(name) + } +} + +/// Changes applied by [`Queues::update`](crate::Queues::update). +/// +/// Fields left unset keep their current value. The queue's `updated_at` is +/// refreshed either way. +#[derive(Clone, Debug, Default, PartialEq)] +#[non_exhaustive] +pub struct QueueUpdateParams { + pub(crate) metadata: Option>, +} + +impl QueueUpdateParams { + /// Creates parameters that change nothing. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Replaces the queue's metadata object. Clients working the queue are + /// notified of the new metadata. + #[must_use] + pub fn metadata(mut self, metadata: Map) -> Self { + self.metadata = Some(metadata); + self + } +} diff --git a/rust/riverqueue/src/storage.rs b/rust/riverqueue/src/storage.rs new file mode 100644 index 000000000..0d0e617d6 --- /dev/null +++ b/rust/riverqueue/src/storage.rs @@ -0,0 +1,654 @@ +//! River's storage operations and their per-backend implementations. +//! +//! Each backend-specific statement is a method of the private [`Backend`] +//! trait, implemented once per built-in database. [`Storage`] binds a backend +//! to one connection and adds the semantics every backend shares: parameter +//! validation, not-found and state errors, extension hooks, and control +//! notifications. [`Session`] supplies that connection, either borrowed from +//! a caller-managed transaction or owned by River for one operation. +//! +//! Adding a backend means implementing [`Backend`] and adding a variant for +//! it to the connection, [`AnyBackend`], and [`Session`] enums; the +//! operations themselves don't change. + +#[cfg(feature = "postgres")] +mod postgres; +#[cfg(feature = "sqlite")] +mod sqlite; + +use serde_json::{Map, Value}; +#[cfg(feature = "postgres")] +use sqlx::Postgres; +#[cfg(feature = "sqlite")] +use sqlx::Sqlite; +use sqlx::{Transaction, pool::PoolConnection}; + +use crate::__private::DatabaseConnection; +use crate::client::{ClientInner, after_jobs_set_state}; +use crate::database::{Database, DatabasePool}; +use crate::{ + Error, JobDeleteManyParams, JobListParams, JobRow, JobUpdateParams, Queue, QueueListParams, +}; + +/// Queue name that addresses every persisted queue in pause and resume. It is +/// part of River's cross-language storage and notification protocol, not the +/// public API. +pub(crate) const QUEUE_ALL: &str = "*"; + +/// One built-in backend's implementation of River's storage statements, +/// bound to a connection. +/// +/// Methods run on the bound connection only; whether that connection is in a +/// transaction is the caller's concern. +pub(crate) trait Backend { + /// Borrows the bound connection for an extension hook. + fn connection(&mut self) -> DatabaseConnection<'_>; + + /// Cancels a job, notifying the client running it, and returns its + /// current row. `None` means the job doesn't exist. + async fn job_cancel(&mut self, id: i64) -> Result, Error>; + + /// Returns the IDs among `ids` of running jobs with a cancellation + /// request, in ascending order. + async fn job_cancel_requested(&mut self, ids: &[i64]) -> Result, Error>; + + /// Completes a running job, merging metadata updates. + async fn job_complete( + &mut self, + id: i64, + metadata_updates: &Map, + ) -> Result; + + /// Deletes a non-running job. + async fn job_delete(&mut self, id: i64) -> Result; + + /// Claims an available job for `client_id` like a fetch does. `None` + /// means the job doesn't exist or isn't available. + async fn job_claim( + &mut self, + id: i64, + client_id: &str, + max_attempted_by: i32, + ) -> Result, Error>; + + /// Deletes non-running jobs matching a validated filter. + async fn job_delete_many(&mut self, filter: &JobListParams) -> Result, Error>; + + async fn job_get(&mut self, id: i64) -> Result, Error>; + + /// Lists jobs matching validated parameters. + async fn job_list(&mut self, params: &JobListParams) -> Result, Error>; + + /// Makes a job available again, notifying its queue, and returns its + /// current row. `None` means the job doesn't exist. + async fn job_retry(&mut self, id: i64) -> Result, Error>; + + /// Merges metadata into a job. `None` means the job doesn't exist. + async fn job_update( + &mut self, + id: i64, + metadata: &Map, + ) -> Result, Error>; + + /// Sends a notification that is delivered when the connection's + /// transaction commits. + async fn notify(&mut self, topic: &str, payload: &str) -> Result<(), Error>; + + async fn queue_get(&mut self, name: &str) -> Result, Error>; + + async fn queue_list(&mut self, limit: u32) -> Result, Error>; + + /// Pauses or resumes the named queue, or every queue for + /// [`QUEUE_ALL`], returning how many queues matched. + async fn queue_set_paused(&mut self, name: &str, paused: bool) -> Result; + + /// Creates a queue record or refreshes its `updated_at`. + async fn queue_touch(&mut self, name: &str) -> Result; + + /// Refreshes a queue's `updated_at`, replacing its metadata when given. + /// `None` means the queue doesn't exist. + async fn queue_update( + &mut self, + name: &str, + metadata: Option<&Map>, + ) -> Result, Error>; +} + +/// The backend selected by a connection. +enum AnyBackend<'c> { + #[cfg(feature = "postgres")] + Postgres(postgres::PostgresBackend<'c>), + #[cfg(feature = "sqlite")] + Sqlite(sqlite::SqliteBackend<'c>), +} + +/// Runs `$call` with `$backend` bound to the concrete backend. +macro_rules! dispatch { + ($any:expr, $backend:ident => $call:expr) => { + match $any { + #[cfg(feature = "postgres")] + AnyBackend::Postgres($backend) => $call, + #[cfg(feature = "sqlite")] + AnyBackend::Sqlite($backend) => $call, + } + }; +} + +impl Backend for AnyBackend<'_> { + fn connection(&mut self) -> DatabaseConnection<'_> { + dispatch!(self, backend => backend.connection()) + } + + async fn job_cancel(&mut self, id: i64) -> Result, Error> { + dispatch!(self, backend => backend.job_cancel(id).await) + } + + async fn job_cancel_requested(&mut self, ids: &[i64]) -> Result, Error> { + dispatch!(self, backend => backend.job_cancel_requested(ids).await) + } + + async fn job_complete( + &mut self, + id: i64, + metadata_updates: &Map, + ) -> Result { + dispatch!(self, backend => backend.job_complete(id, metadata_updates).await) + } + + async fn job_delete(&mut self, id: i64) -> Result { + dispatch!(self, backend => backend.job_delete(id).await) + } + + async fn job_delete_many(&mut self, filter: &JobListParams) -> Result, Error> { + dispatch!(self, backend => backend.job_delete_many(filter).await) + } + + async fn job_claim( + &mut self, + id: i64, + client_id: &str, + max_attempted_by: i32, + ) -> Result, Error> { + dispatch!(self, backend => backend.job_claim(id, client_id, max_attempted_by).await) + } + + async fn job_get(&mut self, id: i64) -> Result, Error> { + dispatch!(self, backend => backend.job_get(id).await) + } + + async fn job_list(&mut self, params: &JobListParams) -> Result, Error> { + dispatch!(self, backend => backend.job_list(params).await) + } + + async fn job_retry(&mut self, id: i64) -> Result, Error> { + dispatch!(self, backend => backend.job_retry(id).await) + } + + async fn job_update( + &mut self, + id: i64, + metadata: &Map, + ) -> Result, Error> { + dispatch!(self, backend => backend.job_update(id, metadata).await) + } + + async fn notify(&mut self, topic: &str, payload: &str) -> Result<(), Error> { + dispatch!(self, backend => backend.notify(topic, payload).await) + } + + async fn queue_get(&mut self, name: &str) -> Result, Error> { + dispatch!(self, backend => backend.queue_get(name).await) + } + + async fn queue_list(&mut self, limit: u32) -> Result, Error> { + dispatch!(self, backend => backend.queue_list(limit).await) + } + + async fn queue_set_paused(&mut self, name: &str, paused: bool) -> Result { + dispatch!(self, backend => backend.queue_set_paused(name, paused).await) + } + + async fn queue_touch(&mut self, name: &str) -> Result { + dispatch!(self, backend => backend.queue_touch(name).await) + } + + async fn queue_update( + &mut self, + name: &str, + metadata: Option<&Map>, + ) -> Result, Error> { + dispatch!(self, backend => backend.queue_update(name, metadata).await) + } +} + +/// River's storage operations on one connection. +pub(crate) struct Storage<'c> { + backend: AnyBackend<'c>, + inner: &'c ClientInner, +} + +impl<'c> Storage<'c> { + pub(crate) fn new(inner: &'c ClientInner, connection: DatabaseConnection<'c>) -> Self { + let backend = match connection { + #[cfg(feature = "postgres")] + DatabaseConnection::Postgres(connection) => { + AnyBackend::Postgres(postgres::PostgresBackend { + capabilities: inner.database.postgres_capabilities(), + connection, + schema: &inner.schema, + }) + } + #[cfg(feature = "sqlite")] + DatabaseConnection::Sqlite(connection) => { + AnyBackend::Sqlite(sqlite::SqliteBackend { connection }) + } + }; + Self { backend, inner } + } + + /// Cancels a job and returns its current row. The running client is + /// notified when the transaction commits. + pub(crate) async fn job_cancel(&mut self, id: i64) -> Result { + let row = self + .backend + .job_cancel(id) + .await? + .ok_or(Error::NotFound(crate::Record::Job(id)))?; + self.after_job_update(&row, JobUpdate::Cancel).await?; + Ok(row) + } + + /// Completes a running job, merging metadata updates, and runs the + /// extension's set-state hook in the same transaction. + pub(crate) async fn job_complete( + &mut self, + id: i64, + metadata_updates: &Map, + ) -> Result { + let row = self.backend.job_complete(id, metadata_updates).await?; + if self.inner.pilot.intercepts_job_set_state() { + after_jobs_set_state( + self.inner, + self.backend.connection(), + &[row.id], + std::slice::from_ref(&row), + ) + .await?; + } + Ok(row) + } + + /// Deletes a non-running job and returns its former row. + pub(crate) async fn job_delete(&mut self, id: i64) -> Result { + self.backend.job_delete(id).await + } + + /// Deletes matching non-running jobs with an explicit safety guard. + pub(crate) async fn job_delete_many( + &mut self, + params: &JobDeleteManyParams, + ) -> Result, Error> { + if !params.all && !params.filter.has_filter() { + return Err(Error::invalid_job_context( + "storage parameters", + "bulk delete requires a filter or all=true".to_owned(), + )); + } + params.filter.validate().map_err(Error::invalid_job)?; + self.backend.job_delete_many(¶ms.filter).await + } + + /// Claims one available job for this client, as a fetch would. + pub(crate) async fn job_claim(&mut self, id: i64) -> Result { + if let Some(row) = self + .backend + .job_claim(id, &self.inner.id, crate::client::ATTEMPTED_BY_MAX) + .await? + { + return Ok(row); + } + let row = self.job_get(id).await?; + Err(Error::invalid_job(format!( + "job {id} must be available to be claimed; state is {}", + row.state + ))) + } + + pub(crate) async fn job_get(&mut self, id: i64) -> Result { + self.backend + .job_get(id) + .await? + .ok_or(Error::NotFound(crate::Record::Job(id))) + } + + pub(crate) async fn job_list(&mut self, params: &JobListParams) -> Result, Error> { + params.validate().map_err(Error::invalid_job)?; + self.backend.job_list(params).await + } + + /// Makes a non-running job available again and returns its current row. + pub(crate) async fn job_retry(&mut self, id: i64) -> Result { + let row = self + .backend + .job_retry(id) + .await? + .ok_or(Error::NotFound(crate::Record::Job(id)))?; + self.after_job_update(&row, JobUpdate::Retry).await?; + Ok(row) + } + + /// Merges job metadata and optionally sets recorded output. + pub(crate) async fn job_update( + &mut self, + id: i64, + params: JobUpdateParams, + ) -> Result { + let mut metadata = params.metadata; + if let Some(output) = params.output { + crate::worker::check_output_size(&output).map_err(Error::invalid_job)?; + metadata.insert(crate::METADATA_KEY_OUTPUT.to_owned(), output); + } + self.backend + .job_update(id, &metadata) + .await? + .ok_or(Error::NotFound(crate::Record::Job(id))) + } + + /// Asks the current leader to resign once the transaction commits, with + /// the payload Go writes, including its empty `leader_id`. + pub(crate) async fn leader_request_resign(&mut self) -> Result<(), Error> { + self.backend + .notify( + crate::NOTIFICATION_TOPIC_LEADERSHIP, + r#"{"action":"request_resign","leader_id":""}"#, + ) + .await + } + + pub(crate) async fn queue_get(&mut self, name: &str) -> Result, Error> { + self.backend.queue_get(name).await + } + + pub(crate) async fn queue_list( + &mut self, + params: &QueueListParams, + ) -> Result, Error> { + if !(1..=10_000).contains(¶ms.limit) { + return Err(Error::invalid_job_context( + "storage parameters", + "queue list limit must be between 1 and 10000".to_owned(), + )); + } + self.backend.queue_list(params.limit).await + } + + /// Pauses or resumes one queue, or every queue for [`QUEUE_ALL`], and + /// notifies clients when the transaction commits. + /// + /// Like Go, naming a queue that has no persisted record is an error, + /// while [`QUEUE_ALL`] succeeds even when no queues exist yet. + pub(crate) async fn queue_set_paused(&mut self, name: &str, paused: bool) -> Result<(), Error> { + let updated = self.backend.queue_set_paused(name, paused).await?; + if updated == 0 && name != QUEUE_ALL { + return Err(Error::NotFound(crate::Record::Queue(name.to_owned()))); + } + let payload = serde_json::json!({ + "action": if paused { "pause" } else { "resume" }, + "queue": name, + }); + self.backend + .notify(crate::NOTIFICATION_TOPIC_CONTROL, &payload.to_string()) + .await + } + + pub(crate) async fn queue_touch(&mut self, name: &str) -> Result { + self.backend.queue_touch(name).await + } + + /// Updates a queue, refreshing its `updated_at`, and notifies clients of + /// new metadata when the transaction commits. + pub(crate) async fn queue_update( + &mut self, + name: &str, + metadata: Option<&Map>, + ) -> Result { + let queue = self + .backend + .queue_update(name, metadata) + .await? + .ok_or_else(|| Error::NotFound(crate::Record::Queue(name.to_owned())))?; + // Like Go, only a metadata change notifies clients. + if let Some(metadata) = metadata { + let payload = serde_json::json!({ + "action": "metadata_changed", + "metadata": metadata, + "queue": name, + }); + self.backend + .notify(crate::NOTIFICATION_TOPIC_CONTROL, &payload.to_string()) + .await?; + } + Ok(queue) + } + + /// Runs the extension's cancel or retry post-hook in the operation's + /// transaction when it intercepts those operations. An error rolls back + /// the caller's transaction along with the update. + async fn after_job_update(&mut self, row: &JobRow, update: JobUpdate) -> Result<(), Error> { + let pilot = &self.inner.pilot; + if !pilot.intercepts_job_cancel_retry() { + return Ok(()); + } + let params = crate::__private::JobUpdatedParams { + database: self.inner.pilot_database_config(), + job: row.clone(), + }; + let connection = self.backend.connection(); + let (phase, result) = match update { + JobUpdate::Cancel => ( + crate::ExtensionPhase::AddOn { + operation: "job cancel", + }, + pilot.after_job_cancel(connection, ¶ms).await, + ), + JobUpdate::Retry => ( + crate::ExtensionPhase::AddOn { + operation: "job retry", + }, + pilot.after_job_retry(connection, ¶ms).await, + ), + }; + result.map_err(|source| Error::Extension { phase, source }) + } +} + +/// Which job operation an extension post-hook follows. +#[derive(Clone, Copy, Debug)] +enum JobUpdate { + Cancel, + Retry, +} + +/// Error for transactional completion of a job that isn't running. +fn job_not_running(state: crate::JobState) -> Error { + Error::JobNotRunning { state } +} + +/// How an operation River runs on its own pool uses its connection. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum Access { + /// A pooled connection in autocommit mode, for operations of a single + /// statement. + Autocommit, + /// A transaction River commits, for operations of several statements or + /// with notifications. SQLite transactions take the write lock up front. + Transaction, +} + +/// The connection one storage operation runs on. +pub(crate) struct Session<'a> { + connection: SessionConnection<'a>, +} + +enum SessionConnection<'a> { + /// A caller-managed transaction, which River never commits. + Caller(DatabaseConnection<'a>), + #[cfg(feature = "postgres")] + PostgresConnection(PoolConnection), + #[cfg(feature = "postgres")] + PostgresTransaction(Transaction<'static, Postgres>), + #[cfg(feature = "sqlite")] + SqliteConnection(PoolConnection), + #[cfg(feature = "sqlite")] + SqliteTransaction(Transaction<'static, Sqlite>), +} + +impl<'a> Session<'a> { + /// Runs on a caller-managed transaction. + pub(crate) const fn caller(connection: DatabaseConnection<'a>) -> Self { + Self { + connection: SessionConnection::Caller(connection), + } + } + + /// Acquires a connection from the client's own pool. + pub(crate) async fn begin(database: &Database, access: Access) -> Result { + let connection = match (database.pool(), access) { + #[cfg(feature = "postgres")] + (DatabasePool::Postgres(pool), Access::Autocommit) => { + SessionConnection::PostgresConnection(pool.acquire().await?) + } + #[cfg(feature = "postgres")] + (DatabasePool::Postgres(pool), Access::Transaction) => { + SessionConnection::PostgresTransaction(crate::database::begin_postgres(pool).await?) + } + #[cfg(feature = "sqlite")] + (DatabasePool::Sqlite(pool), Access::Autocommit) => { + SessionConnection::SqliteConnection(pool.acquire().await?) + } + #[cfg(feature = "sqlite")] + (DatabasePool::Sqlite(pool), Access::Transaction) => { + SessionConnection::SqliteTransaction( + crate::database::begin_sqlite_write(pool).await?, + ) + } + }; + Ok(Self { connection }) + } + + /// Returns storage operations bound to this session's connection. + pub(crate) fn storage<'s>(&'s mut self, inner: &'s ClientInner) -> Storage<'s> { + let connection = match &mut self.connection { + SessionConnection::Caller(connection) => connection.reborrow(), + #[cfg(feature = "postgres")] + SessionConnection::PostgresConnection(connection) => { + DatabaseConnection::Postgres(connection) + } + #[cfg(feature = "postgres")] + SessionConnection::PostgresTransaction(transaction) => { + DatabaseConnection::Postgres(transaction) + } + #[cfg(feature = "sqlite")] + SessionConnection::SqliteConnection(connection) => { + DatabaseConnection::Sqlite(connection) + } + #[cfg(feature = "sqlite")] + SessionConnection::SqliteTransaction(transaction) => { + DatabaseConnection::Sqlite(transaction) + } + }; + Storage::new(inner, connection) + } + + /// Commits a transaction River owns. A caller-managed transaction is + /// left for the caller to commit, and an autocommit connection has + /// nothing to commit. Dropping a session instead rolls River's own + /// transaction back. + pub(crate) async fn commit(self) -> Result<(), Error> { + match self.connection { + #[cfg(feature = "postgres")] + SessionConnection::PostgresTransaction(transaction) => transaction.commit().await?, + #[cfg(feature = "sqlite")] + SessionConnection::SqliteTransaction(transaction) => transaction.commit().await?, + SessionConnection::Caller(_) => {} + #[cfg(feature = "postgres")] + SessionConnection::PostgresConnection(_) => {} + #[cfg(feature = "sqlite")] + SessionConnection::SqliteConnection(_) => {} + } + Ok(()) + } +} + +/// Creates the client's queue record or refreshes its `updated_at`. +pub(crate) async fn touch_queue(inner: &ClientInner, name: &str) -> Result { + let mut session = Session::begin(&inner.database, Access::Autocommit).await?; + session.storage(inner).queue_touch(name).await +} + +/// Returns the IDs among `ids` of running jobs with a cancellation request. +pub(crate) async fn job_cancel_requested( + inner: &ClientInner, + ids: &[i64], +) -> Result, Error> { + let mut session = Session::begin(&inner.database, Access::Autocommit).await?; + session + .storage(inner) + .backend + .job_cancel_requested(ids) + .await +} + +/// Loads a queue record, if one exists. +pub(crate) async fn load_queue(inner: &ClientInner, name: &str) -> Result, Error> { + let mut session = Session::begin(&inner.database, Access::Autocommit).await?; + session.storage(inner).queue_get(name).await +} + +/// Transactional job operations for a [`WorkContext`](crate::WorkContext), +/// whose generic executor can't be reborrowed into a request builder. +impl crate::Client { + /// Completes a running job in a caller-managed transaction, merging + /// metadata recorded on the work context. + pub(crate) async fn job_complete_tx_with_metadata<'executor, E>( + &self, + executor: E, + id: i64, + metadata_updates: Map, + ) -> Result + where + E: crate::database::DatabaseTransactionExecutor<'executor>, + { + let connection = self.inner.transaction_connection(executor)?; + complete_in_caller_transaction(&self.inner, connection, id, &metadata_updates).await + } + + /// Updates a job in a caller-managed transaction. + pub(crate) async fn job_update_tx<'executor, E>( + &self, + executor: E, + id: i64, + params: JobUpdateParams, + ) -> Result + where + E: crate::database::DatabaseTransactionExecutor<'executor>, + { + let connection = self.inner.transaction_connection(executor)?; + Storage::new(&self.inner, connection) + .job_update(id, params) + .await + } +} + +/// Completes a running job in a caller-managed transaction, running directly +/// in it without a savepoint. +pub(crate) async fn complete_in_caller_transaction( + inner: &ClientInner, + connection: DatabaseConnection<'_>, + id: i64, + metadata_updates: &Map, +) -> Result { + Storage::new(inner, connection) + .job_complete(id, metadata_updates) + .await +} diff --git a/rust/riverqueue/src/storage/postgres.rs b/rust/riverqueue/src/storage/postgres.rs new file mode 100644 index 000000000..1ae5ae0f0 --- /dev/null +++ b/rust/riverqueue/src/storage/postgres.rs @@ -0,0 +1,502 @@ +//! PostgreSQL implementation of River's storage operations. + +use chrono::{DateTime, Utc}; +use serde_json::{Map, Value}; +use sqlx::{AssertSqlSafe, FromRow, PgConnection, Postgres, types::Json}; + +use super::Backend; +use crate::__private::DatabaseConnection; +use crate::client::{JobRecord, go_time_json, job_projection}; +use crate::database::postgres_capabilities::CapabilitiesCache; +use crate::query::{JobListSqlPart, JobListTimeField}; +use crate::{Error, JobListParams, JobRow, JobState, Queue, SchemaName}; + +/// PostgreSQL storage bound to one connection. +pub(super) struct PostgresBackend<'c> { + /// The database's detected server capabilities, or `None` to detect + /// them for each statement that needs them. + pub(super) capabilities: Option<&'c CapabilitiesCache>, + pub(super) connection: &'c mut PgConnection, + pub(super) schema: &'c SchemaName, +} + +impl PostgresBackend<'_> { + /// Whether `pg_notify` reaches listeners on this server. + async fn supports_listen_notify(&mut self) -> Result { + Ok( + CapabilitiesCache::load_or_detect(self.capabilities, &mut *self.connection) + .await? + .supports_listen_notify, + ) + } +} + +impl Backend for PostgresBackend<'_> { + fn connection(&mut self) -> DatabaseConnection<'_> { + DatabaseConnection::Postgres(self.connection) + } + + // The fallback arm of `job_cancel` and `job_retry` returns the row when the + // update matched nothing, as when a concurrent cancel or retry won. Like + // River Go, it locks the row so it reads the winner's committed version + // rather than this statement's older snapshot. + async fn job_cancel(&mut self, id: i64) -> Result, Error> { + let notify = self.supports_listen_notify().await?; + let table = self.schema.qualify("river_job"); + let sql = format!( + "WITH locked AS (\ + SELECT id, queue, state, finalized_at FROM {table} WHERE id = $1 FOR UPDATE\ + ), notified AS (\ + SELECT id, CASE WHEN $5::boolean THEN pg_notify(concat(coalesce($2::text, current_schema()), '.', $3::text), json_build_object('action', 'cancel', 'job_id', id, 'queue', queue)::text) END \ + FROM locked WHERE state NOT IN ('cancelled', 'completed', 'discarded') AND finalized_at IS NULL\ + ), updated AS (\ + UPDATE {table} AS job SET \ + state = CASE WHEN state = 'running' THEN state ELSE 'cancelled' END, \ + finalized_at = CASE WHEN state = 'running' THEN finalized_at ELSE now() END, \ + metadata = jsonb_set(metadata, '{{cancel_attempted_at}}'::text[], to_jsonb($4::text), true) \ + FROM notified WHERE job.id = notified.id RETURNING job.*\ + ) \ + SELECT {}, false AS unique_skipped_as_duplicate FROM updated AS job \ + UNION ALL \ + SELECT {}, false AS unique_skipped_as_duplicate \ + FROM (SELECT * FROM {table} WHERE id = $1 FOR UPDATE) AS job \ + WHERE NOT EXISTS (SELECT 1 FROM updated) LIMIT 1", + job_projection("job"), + job_projection("job") + ); + sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) + .bind(id) + .bind(self.schema.as_deref()) + .bind(crate::NOTIFICATION_TOPIC_CONTROL) + .bind(go_time_json(Utc::now())) + .bind(notify) + .fetch_optional(&mut *self.connection) + .await? + .map(JobRecord::into_job_row) + .transpose() + } + + async fn job_cancel_requested(&mut self, ids: &[i64]) -> Result, Error> { + let table = self.schema.qualify("river_job"); + let sql = format!( + "SELECT id FROM {table} \ + WHERE id = any($1) AND metadata ? 'cancel_attempted_at' AND state = 'running' \ + ORDER BY id" + ); + Ok(sqlx::query_scalar(AssertSqlSafe(sql)) + .bind(ids) + .fetch_all(&mut *self.connection) + .await?) + } + + async fn job_complete( + &mut self, + id: i64, + metadata_updates: &Map, + ) -> Result { + let table = self.schema.qualify("river_job"); + let state: Option = sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT state::text FROM {table} WHERE id = $1 FOR UPDATE" + ))) + .bind(id) + .fetch_optional(&mut *self.connection) + .await?; + match state.as_deref() { + None => return Err(Error::NotFound(crate::Record::Job(id))), + Some("running") => {} + Some(state) => { + return Err(super::job_not_running( + state + .parse() + .map_err(|error| sqlx::Error::Decode(Box::new(error)))?, + )); + } + } + let sql = format!( + "UPDATE {table} AS job SET state = 'completed', finalized_at = now(), \ + metadata = metadata || $2::jsonb \ + WHERE id = $1 AND state = 'running' \ + RETURNING {}, false AS unique_skipped_as_duplicate", + job_projection("job") + ); + sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) + .bind(id) + .bind(Json(metadata_updates)) + .fetch_optional(&mut *self.connection) + .await? + .ok_or(Error::NotFound(crate::Record::Job(id)))? + .into_job_row() + } + + async fn job_delete(&mut self, id: i64) -> Result { + let table = self.schema.qualify("river_job"); + let state: Option = sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT state::text FROM {table} WHERE id = $1 FOR UPDATE" + ))) + .bind(id) + .fetch_optional(&mut *self.connection) + .await?; + match state.as_deref() { + None => return Err(Error::NotFound(crate::Record::Job(id))), + Some("running") => return Err(Error::JobRunning), + Some(_) => {} + } + let sql = format!( + "DELETE FROM {table} AS job WHERE id = $1 RETURNING {}, false AS unique_skipped_as_duplicate", + job_projection("job") + ); + sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) + .bind(id) + .fetch_one(&mut *self.connection) + .await? + .into_job_row() + } + + async fn job_delete_many(&mut self, filter: &JobListParams) -> Result, Error> { + let table = self.schema.qualify("river_job"); + let parts = job_list_sql_parts(self.schema, filter, false); + // Mirrors Go's `JobDeleteMany`: running jobs are excluded before the + // limit applies, candidates already locked by another transaction are + // skipped rather than waited on, and rows come back in the list order. + let sql = format!( + "WITH jobs_to_delete AS (\ + SELECT id FROM {table} AS job WHERE {where_sql} AND state != 'running' \ + ORDER BY {order_sql} LIMIT $11 FOR UPDATE SKIP LOCKED\ + ), deleted AS (\ + DELETE FROM {table} WHERE id IN (SELECT id FROM jobs_to_delete) RETURNING *\ + ) \ + SELECT {}, false AS unique_skipped_as_duplicate FROM deleted AS job ORDER BY {order_sql}", + job_projection("job"), + where_sql = parts.where_sql, + order_sql = parts.order_sql, + ); + let records = bind_job_list(sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)), filter) + .fetch_all(&mut *self.connection) + .await?; + records.into_iter().map(JobRecord::into_job_row).collect() + } + + async fn job_get(&mut self, id: i64) -> Result, Error> { + let table = self.schema.qualify("river_job"); + let sql = format!( + "SELECT {}, false AS unique_skipped_as_duplicate FROM {table} AS job WHERE id = $1 LIMIT 1", + job_projection("job") + ); + sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) + .bind(id) + .fetch_optional(&mut *self.connection) + .await? + .map(JobRecord::into_job_row) + .transpose() + } + + async fn job_list(&mut self, params: &JobListParams) -> Result, Error> { + let table = self.schema.qualify("river_job"); + let parts = job_list_sql_parts(self.schema, params, true); + let sql = format!( + "SELECT {}, false AS unique_skipped_as_duplicate FROM {table} AS job \ + WHERE {} ORDER BY {} LIMIT $11", + job_projection("job"), + parts.where_sql, + parts.order_sql, + ); + let records = bind_job_list(sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)), params) + .fetch_all(&mut *self.connection) + .await?; + records.into_iter().map(JobRecord::into_job_row).collect() + } + + async fn job_claim( + &mut self, + id: i64, + client_id: &str, + max_attempted_by: i32, + ) -> Result, Error> { + let table = self.schema.qualify("river_job"); + let sql = format!( + "UPDATE {table} AS job SET state = 'running', attempt = job.attempt + 1, \ + attempted_at = now(), attempted_by = array_append(\ + CASE WHEN array_length(job.attempted_by, 1) >= $3 \ + THEN job.attempted_by[array_length(job.attempted_by, 1) + 2 - $3:] \ + ELSE job.attempted_by END, $2) \ + WHERE id = $1 AND state = 'available' \ + RETURNING {}, false AS unique_skipped_as_duplicate", + job_projection("job") + ); + sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) + .bind(id) + .bind(client_id) + .bind(max_attempted_by) + .fetch_optional(&mut *self.connection) + .await? + .map(JobRecord::into_job_row) + .transpose() + } + + async fn job_retry(&mut self, id: i64) -> Result, Error> { + let table = self.schema.qualify("river_job"); + let sql = format!( + "WITH locked AS (SELECT id FROM {table} WHERE id = $1 FOR UPDATE), \ + updated AS (UPDATE {table} AS job SET state = 'available', \ + max_attempts = CASE WHEN attempt = max_attempts THEN max_attempts + 1 ELSE max_attempts END, \ + finalized_at = NULL, scheduled_at = now() \ + FROM locked WHERE job.id = locked.id AND job.state != 'running' \ + AND NOT (job.state = 'available' AND job.scheduled_at < now()) RETURNING job.*) \ + SELECT {}, false AS unique_skipped_as_duplicate FROM updated AS job \ + UNION ALL SELECT {}, false AS unique_skipped_as_duplicate \ + FROM (SELECT * FROM {table} WHERE id = $1 FOR UPDATE) AS job \ + WHERE NOT EXISTS (SELECT 1 FROM updated) LIMIT 1", + job_projection("job"), + job_projection("job") + ); + // Like Go's `JobRetry`, a retry sends no insert notification; + // producers find the job on their next poll. + sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) + .bind(id) + .fetch_optional(&mut *self.connection) + .await? + .map(JobRecord::into_job_row) + .transpose() + } + + async fn job_update( + &mut self, + id: i64, + metadata: &Map, + ) -> Result, Error> { + let table = self.schema.qualify("river_job"); + let sql = format!( + "UPDATE {table} AS job SET metadata = metadata || $2::jsonb WHERE id = $1 \ + RETURNING {}, false AS unique_skipped_as_duplicate", + job_projection("job") + ); + sqlx::query_as::<_, JobRecord>(AssertSqlSafe(sql)) + .bind(id) + .bind(Json(metadata)) + .fetch_optional(&mut *self.connection) + .await? + .map(JobRecord::into_job_row) + .transpose() + } + + async fn notify(&mut self, topic: &str, payload: &str) -> Result<(), Error> { + if !self.supports_listen_notify().await? { + return Ok(()); + } + sqlx::query( + "SELECT pg_notify(concat(coalesce($1::text, current_schema()), '.', $2::text), $3::text)", + ) + .bind(self.schema.as_deref()) + .bind(topic) + .bind(payload) + .execute(&mut *self.connection) + .await?; + Ok(()) + } + + async fn queue_get(&mut self, name: &str) -> Result, Error> { + let table = self.schema.qualify("river_queue"); + sqlx::query_as::<_, QueueRecord>(AssertSqlSafe(format!( + "SELECT {QUEUE_COLUMNS} FROM {table} WHERE name = $1" + ))) + .bind(name) + .fetch_optional(&mut *self.connection) + .await? + .map(QueueRecord::into_queue) + .transpose() + } + + async fn queue_list(&mut self, limit: u32) -> Result, Error> { + let table = self.schema.qualify("river_queue"); + sqlx::query_as::<_, QueueRecord>(AssertSqlSafe(format!( + "SELECT {QUEUE_COLUMNS} FROM {table} ORDER BY name LIMIT $1" + ))) + .bind(i64::from(limit)) + .fetch_all(&mut *self.connection) + .await? + .into_iter() + .map(QueueRecord::into_queue) + .collect() + } + + async fn queue_set_paused(&mut self, name: &str, paused: bool) -> Result { + let table = self.schema.qualify("river_queue"); + let sql = if paused { + format!( + "UPDATE {table} SET paused_at = coalesce(paused_at, now()), \ + updated_at = CASE WHEN paused_at IS NULL THEN now() ELSE updated_at END \ + WHERE $1 = '*' OR name = $1" + ) + } else { + format!( + "UPDATE {table} SET updated_at = CASE WHEN paused_at IS NOT NULL THEN now() ELSE updated_at END, \ + paused_at = NULL WHERE $1 = '*' OR name = $1" + ) + }; + Ok(sqlx::query(AssertSqlSafe(sql)) + .bind(name) + .execute(&mut *self.connection) + .await? + .rows_affected()) + } + + async fn queue_touch(&mut self, name: &str) -> Result { + let table = self.schema.qualify("river_queue"); + let sql = format!( + "INSERT INTO {table} (name, metadata, updated_at) VALUES ($1, '{{}}'::jsonb, now()) \ + ON CONFLICT (name) DO UPDATE SET updated_at = excluded.updated_at RETURNING {QUEUE_COLUMNS}" + ); + sqlx::query_as::<_, QueueRecord>(AssertSqlSafe(sql)) + .bind(name) + .fetch_one(&mut *self.connection) + .await? + .into_queue() + } + + async fn queue_update( + &mut self, + name: &str, + metadata: Option<&Map>, + ) -> Result, Error> { + let table = self.schema.qualify("river_queue"); + let sql = format!( + "UPDATE {table} SET metadata = CASE WHEN $2::boolean THEN $3::jsonb ELSE metadata END, \ + updated_at = now() WHERE name = $1 RETURNING {QUEUE_COLUMNS}" + ); + sqlx::query_as::<_, QueueRecord>(AssertSqlSafe(sql)) + .bind(name) + .bind(metadata.is_some()) + .bind(metadata.map(Json)) + .fetch_optional(&mut *self.connection) + .await? + .map(QueueRecord::into_queue) + .transpose() + } +} + +/// The columns of a queue row, with its metadata's stored text. +const QUEUE_COLUMNS: &str = + "created_at, metadata, metadata::text AS metadata_text, name, paused_at, updated_at"; + +#[derive(FromRow)] +struct QueueRecord { + created_at: DateTime, + metadata: Json, + metadata_text: String, + name: String, + paused_at: Option>, + updated_at: DateTime, +} + +impl QueueRecord { + fn into_queue(self) -> Result { + Ok(Queue { + created_at: self.created_at, + metadata: self.metadata.0.as_object().cloned().ok_or_else(|| { + Error::invalid_job_context( + "storage parameters", + format!("queue {:?} metadata is not an object", self.name), + ) + })?, + metadata_text: self.metadata_text, + name: self.name, + paused_at: self.paused_at, + updated_at: self.updated_at, + }) + } +} + +/// SQL fragments shared by job listing and bulk deletion. Both bind the same +/// eleven positional parameters through [`bind_job_list`]. +struct JobListSqlParts { + order_sql: String, + where_sql: String, +} + +fn job_list_sql_parts( + schema: &SchemaName, + params: &JobListParams, + optimize_single_state: bool, +) -> JobListSqlParts { + let keyset = params.keyset(); + let cursor_predicate = keyset.after_sql().map_or_else( + || "true".to_owned(), + |parts| { + parts + .into_iter() + .map(|part| match part { + JobListSqlPart::AfterId => "$10".to_owned(), + JobListSqlPart::AfterTime => "$9".to_owned(), + JobListSqlPart::Sql(sql) => sql, + }) + .collect::() + }, + ); + let state_type = schema.qualify("river_job_state"); + // Like Go (upstream 35c4eab8), a single-state list without metadata + // predicates compares state with equality so PostgreSQL can use the + // `(state, &"#.to_owned(), + keys: BTreeMap::from( + [("", 1), ("a&b", 2), ("é", 3), ("é<", 4)] + .map(|(key, value)| (key.to_owned(), value)), + ), + separators: "line\u{2028}paragraph\u{2029}end".to_owned(), + unicode: "é😀/\\".to_owned(), + unicode_amp: "unicode key".to_owned(), + }, + ); + } + + #[derive(Deserialize, JobArgsDerive, Serialize)] + #[river(kind = "conformance_all_args")] + struct NestedOrderArgs { + nested: NestedOrder, + } + + #[derive(Deserialize, Serialize)] + struct NestedOrder { + z: i64, + a: i64, + } + + #[test] + fn typed_args_match_go_nested_order_golden() { + assert_typed_golden( + "nested_struct_wire_order", + &NestedOrderArgs { + nested: NestedOrder { z: 1, a: 2 }, + }, + ); + } + + #[derive(Deserialize, JobArgsDerive, Serialize)] + #[river(kind = "conformance_numeric_boundaries")] + struct NumericBoundaryArgs { + exponent: f64, + fraction: f64, + maximum: i64, + minimum: i64, + unsigned_maximum: u64, + } + + #[test] + fn typed_args_match_go_numeric_boundaries_golden() { + assert_typed_golden( + "numeric_boundaries", + &NumericBoundaryArgs { + exponent: 1e100, + fraction: 1.25, + maximum: i64::MAX, + minimum: i64::MIN, + unsigned_maximum: u64::MAX, + }, + ); + } + + #[derive(Default, Deserialize, Serialize)] + struct SelectedAccount { + #[serde(skip_serializing_if = "String::is_empty")] + id: String, + #[serde(skip_serializing_if = "String::is_empty")] + ignored: String, + #[serde(skip_serializing_if = "String::is_empty")] + region: String, + } + + impl SelectedAccount { + fn is_zero(&self) -> bool { + self.id.is_empty() && self.ignored.is_empty() && self.region.is_empty() + } + } + + #[derive(Deserialize, JobArgsDerive, Serialize)] + #[river( + kind = "conformance_selected_args", + unique(by_args("account.id", "account.region", "label", "path/key")) + )] + struct SelectedArgs { + #[serde(skip_serializing_if = "SelectedAccount::is_zero")] + account: SelectedAccount, + #[serde(skip_serializing_if = "std::ops::Not::not")] + ignored: bool, + #[serde(skip_serializing_if = "String::is_empty")] + label: String, + #[serde(rename = "path/key", skip_serializing_if = "String::is_empty")] + path_key: String, + } + + #[test] + fn typed_args_match_go_selected_goldens() { + assert_typed_golden( + "all_selected_fields_omitted", + &SelectedArgs { + account: SelectedAccount::default(), + ignored: false, + label: String::new(), + path_key: String::new(), + }, + ); + assert_typed_golden( + "selected_siblings_and_slash_key", + &SelectedArgs { + account: SelectedAccount { + id: "acct".to_owned(), + ignored: "irrelevant".to_owned(), + region: "west".to_owned(), + }, + ignored: false, + label: String::new(), + path_key: "slash".to_owned(), + }, + ); + assert_typed_golden( + "selected_nested_args", + &SelectedArgs { + account: SelectedAccount { + id: "acct-123".to_owned(), + ignored: "not selected".to_owned(), + region: String::new(), + }, + ignored: true, + label: "selected".to_owned(), + path_key: String::new(), + }, + ); + } + + #[derive(Deserialize, JobArgsDerive, Serialize)] + #[river( + kind = "conformance_selected_args", + unique(by_args("account.id", "account.region", "label", "path/key")) + )] + struct SelectedNullArgs { + #[serde(skip_serializing_if = "SelectedAccount::is_zero")] + account: SelectedAccount, + label: Option, + #[serde(rename = "path/key", skip_serializing_if = "String::is_empty")] + path_key: String, + } + + #[test] + fn typed_args_match_go_selected_explicit_null_golden() { + assert_typed_golden( + "selected_explicit_null", + &SelectedNullArgs { + account: SelectedAccount::default(), + label: None, + path_key: String::new(), + }, + ); + } + + #[derive(Deserialize, JobArgsDerive, Serialize)] + #[river(kind = "conformance_all_args")] + struct TimeArgs { + #[serde(with = "crate::encoding::go_time")] + fraction: DateTime, + #[serde(with = "crate::encoding::go_time")] + micros: DateTime, + #[serde(with = "crate::encoding::go_time")] + millis: DateTime, + whole: DateTime, + } + + #[test] + fn typed_args_match_go_time_golden() { + let whole = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + assert_typed_golden( + "typed_time_values", + &TimeArgs { + fraction: whole + chrono::Duration::milliseconds(500), + micros: whole + chrono::Duration::microseconds(123_456), + millis: whole + chrono::Duration::milliseconds(120), + // Whole seconds need no helper: chrono and Go agree. + whole, + }, + ); + } + + #[derive(Deserialize, JobArgsDerive, Serialize)] + #[river(kind = "conformance_all_args")] + struct TypedFloatArgs { + below_large: f64, + large: f64, + large_boundary: f64, + largest: f64, + negative: f64, + negative_zero: f64, + one: f64, + single: f32, + single_large: f32, + single_small: f32, + small: f64, + small_boundary: f64, + smallest: f64, + tenth: f64, + } + + #[test] + fn typed_args_match_go_float_golden() { + assert_typed_golden( + "typed_float_formatting", + &TypedFloatArgs { + below_large: f64::from_bits(1e21_f64.to_bits() - 1), + large: 1e20, + large_boundary: 1e21, + largest: f64::MAX, + negative: -1.5e-9, + negative_zero: -0.0, + one: 1.0, + single: 1.1, + single_large: 1e21, + single_small: 1e-7, + small: 1e-7, + small_boundary: 1e-6, + smallest: 5e-324, + tenth: 0.1, + }, + ); + } + + #[derive(Deserialize, Serialize)] + #[serde(transparent)] + struct MapArgs(BTreeMap); + + impl JobArgs for MapArgs { + const KIND: &'static str = "conformance_all_args"; + } + + #[test] + fn typed_map_args_hash_like_go_despite_member_order() { + // Go encoded these members in a custom order; the all-arguments hash + // sorts top-level keys, so a sorted Rust map hashes identically. + let case = golden("map_order_and_negative_zero"); + let args = MapArgs(BTreeMap::from( + [ + ("2", 2.0), + ("10", 10.0), + ("zero", -0.0), + ("😀", 1.0), + ("\u{e000}", 2.0), + ] + .map(|(key, value)| (key.to_owned(), value)), + )); + let encoded = encode_args(&args).unwrap(); + assert_ne!(encoded.get(), case.compact_args().get()); + assert_eq!( + key_for(MapArgs::KIND, &[], &encoded, &case.unique_opts()).unwrap(), + case.expected_key() + ); + } + + #[test] + fn hashes_raw_bytes_without_reinterpreting_values() { + let opts = UniqueOpts::new().with_by_args(true); + let raw = |json: &str| RawValue::from_string(json.to_owned()).unwrap(); + let key = |json: &str| key_for("raw", &[], &raw(json), &opts).unwrap(); + let expected = + |text: &str| -> [u8; 32] { Sha256::digest(format!("&kind=raw&args={text}")).into() }; + + // Number tokens, nested whitespace, and nested order are hashed as + // written; top-level whitespace is not. + assert_eq!( + key(r#"{"b":1.0,"a":1e2}"#), + expected(r#"{"a":1e2,"b":1.0}"#) + ); + assert_eq!( + key(r#" { "a" : [1, {"z":1, "y":2}] , "b":null } "#), + expected(r#"{"a":[1, {"z":1, "y":2}],"b":null}"#) + ); + // Duplicate keys resolve to their first value. + assert_eq!(key(r#"{"a":1,"a":2}"#), expected(r#"{"a":1}"#)); + // Escaped printable ASCII keys are written verbatim, like sjson; other + // keys are re-encoded with Go's escaping. + assert_eq!( + key(r#"{"\u003ck\u003e":1,"a\"b":2,"é\u0026":3,"line\n":4}"#), + expected(r#"{"":1,"a\"b":2,"line\n":4,"é\u0026":3}"#) + ); + } + + #[test] + fn selected_paths_follow_sjson_assembly() { + let raw = RawValue::from_string( + r#"{"b":{"y":2,"x":1},"a-b":3,"a":{"c":null},"ignored":true}"#.to_owned(), + ) + .unwrap(); + let mut output = String::new(); + write_selected_args( + raw.get(), + &[ + &["b", "x"], + &["a", "c"], + &["a-b"], + &["b", "y"], + &["missing", "path"], + ], + &mut output, + ) + .unwrap(); + // Paths are applied in sorted order ("a-b" < "a.c" < "b.x" < "b.y"), + // explicit nulls are kept, and missing paths are omitted. + assert_eq!(output, r#"{"a-b":3,"a":{"c":null},"b":{"x":1,"y":2}}"#); + + let mut output = String::new(); + write_selected_args(raw.get(), &[&["missing"]], &mut output).unwrap(); + assert_eq!(output, ""); + } + + #[test] + fn hashes_literal_top_level_names() { + let opts = UniqueOpts::new().with_by_args(true); + for json in [r#"{"a.b":1}"#, r#"{"@this":1}"#, r#"{"":1}"#] { + let raw = RawValue::from_string(json.to_owned()).unwrap(); + assert!(key_for("raw", &[], &raw, &opts).is_ok(), "{json}"); + } + + let raw = RawValue::from_string(r#"{"a":{"b":1}}"#.to_owned()).unwrap(); + assert!(key_for("raw", &[&["a"], &["a", "b"]], &raw, &opts).is_err()); + assert!(key_for("raw", &[&["a", ""]], &raw, &opts).is_err()); + assert!(key_for("raw", &[&["a", "0"]], &raw, &opts).is_err()); + assert!(key_for("raw", &[&["-1"]], &raw, &opts).is_err()); + } + + #[test] + fn empty_array_all_args_hash_an_empty_object() { + let opts = UniqueOpts::new().with_by_args(true); + let expected: [u8; 32] = Sha256::digest(b"&kind=raw&args={}").into(); + for json in ["[]", " [ \n] "] { + let raw = RawValue::from_string(json.to_owned()).unwrap(); + assert_eq!( + key_for("raw", &[], &raw, &opts).unwrap(), + expected, + "{json}" + ); + } + } + + #[test] + fn non_object_all_args_are_rejected() { + let opts = UniqueOpts::new().with_by_args(true); + for json in ["[1]", "[[]]", "[{}]", "1", "true", "null", r#""text""#] { + let raw = RawValue::from_string(json.to_owned()).unwrap(); + let error = key_for("raw", &[], &raw, &opts).unwrap_err(); + assert!( + error + .to_string() + .contains("unique args must encode a JSON object"), + "{json}: {error}" + ); + assert!(key_for("raw", &[&["a"]], &raw, &opts).is_err(), "{json}"); + } + } + + #[test] + fn state_bitmask_matches_postgres_function() { + assert_eq!(UniqueOpts::default().state_bitmask(), 0b1111_0101); + } + + #[test] + fn truncates_from_go_time_zero() { + let timestamp = Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 5).unwrap(); + let truncated = truncate_period(timestamp, Duration::from_mins(1)).unwrap(); + assert_eq!( + truncated, + Utc.with_ymd_and_hms(2026, 1, 2, 3, 4, 0).unwrap() + ); + } +} diff --git a/rust/riverqueue/src/worker.rs b/rust/riverqueue/src/worker.rs new file mode 100644 index 000000000..234db9906 --- /dev/null +++ b/rust/riverqueue/src/worker.rs @@ -0,0 +1,1537 @@ +//! Typed worker interfaces and registration. + +use std::{ + collections::{HashMap, HashSet}, + error::Error as StdError, + future::Future, + sync::{Arc, Mutex, MutexGuard, PoisonError}, + time::Duration, +}; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use serde::Serialize; +use serde::de::DeserializeOwned; +use serde_json::{Map, Value}; +use tokio::sync::oneshot; +use tokio_util::sync::CancellationToken; + +use crate::{ + BoxError, Client, Error, Job, JobArgs, JobMetadata, JobRow, JobUpdateParams, WorkError, + database::DatabaseTransactionExecutor, +}; + +/// Context available while a job is running. +#[derive(Clone)] +pub struct WorkContext { + cancellation: CancellationToken, + client: Option, + job_id: Option, + metadata_updates: Arc>>, + /// The attempt's peers, for an attempt a client runs. + peers: Option>, + resumable: Arc>, +} + +impl WorkContext { + /// Creates a detached context with no client. + #[must_use] + pub(crate) fn new(cancellation: CancellationToken) -> Self { + Self { + cancellation, + client: None, + job_id: None, + metadata_updates: Arc::new(Mutex::new(Map::new())), + peers: None, + resumable: Arc::new(Mutex::new(ResumableState::default())), + } + } + + /// Cancellation token triggered by timeout, remote cancellation, or stop. + /// + /// A worker that stops because the token was cancelled should return + /// [`WorkCancelled`](crate::WorkCancelled): during a client's hard stop, + /// that makes its job available again without using up the attempt, + /// while any other error consumes it. + /// + /// The token belongs to this attempt's worker. Cancelling it, for + /// example through a drop guard over subtasks, only cancels whatever + /// observes it and doesn't change how River records the attempt. + #[must_use] + pub fn cancellation_token(&self) -> &CancellationToken { + &self.cancellation + } + + /// Returns the River client supervising this job. Contexts constructed by + /// test helpers such as `riverqueue-test` are detached and return `None`. + #[must_use] + pub fn client(&self) -> Option<&Client> { + self.client.as_ref() + } + + /// Completes the running job in a caller-managed transaction, for example + /// alongside business writes the job performed. + /// + /// Metadata recorded on this context, including output, is merged into + /// the job. The job becomes completed only when the transaction commits, + /// and River then leaves the completed row unchanged when the worker + /// returns. [`Jobs::complete`](crate::Jobs::complete) completes a running + /// job by ID outside a worker. + /// + /// # Errors + /// + /// Returns [`Error::Runtime`] when this context doesn't belong to a + /// running worker, [`Error::JobNotRunning`] when the job is no longer + /// running, [`Error::DatabaseMismatch`] for a transaction from another + /// backend, [`Error::Extension`] when an extension's completion hook + /// fails, and [`Error::Database`] when the database operation fails. + /// + /// # Cancel safety + /// + /// Dropping the future before it finishes leaves what it already ran in + /// the caller's transaction; roll that transaction back rather than + /// committing it. The job stays running until a committed completion. + pub async fn job_complete_tx<'executor, E>(&self, connection: E) -> Result + where + E: DatabaseTransactionExecutor<'executor>, + { + let (client, job_id) = self.current_job()?; + client + .job_complete_tx_with_metadata(connection, job_id, self.metadata_updates()) + .await + } + + /// Sets a metadata key that River merges into the job's metadata when it + /// records the attempt's result. Setting a key again replaces its value. + /// + /// # Errors + /// + /// Returns an error when `value` can't be serialized to JSON, including + /// when it contains a non-finite float, which JSON can't represent. + pub fn metadata_set( + &self, + key: impl Into, + value: impl Serialize, + ) -> Result<(), serde_json::Error> { + crate::encoding::check_finite(&value)?; + let value = serde_json::to_value(value)?; + self.insert_metadata(key.into(), value); + Ok(()) + } + + /// Records the job's output under River's reserved output metadata key, + /// where [`JobRow::output`] and River UI read it. Output is limited to + /// 32 MB of JSON, but should be kept much smaller. + /// + /// # Errors + /// + /// Returns an error when `output` can't be serialized to JSON (including + /// when it contains a non-finite float) or its JSON is larger + /// than 32 MB. + pub fn record_output(&self, output: impl Serialize) -> Result<(), serde_json::Error> { + crate::encoding::check_finite(&output)?; + let output = serde_json::to_value(output)?; + check_output_size(&output).map_err(::custom)?; + self.insert_metadata(crate::METADATA_KEY_OUTPUT.to_owned(), output); + Ok(()) + } + + pub(crate) fn insert_metadata(&self, key: String, value: Value) { + self.lock_metadata().insert(key, value); + } + + /// Metadata updates are never held across an await, and each update + /// leaves the map consistent, so a poisoned lock is still usable. + fn lock_metadata(&self) -> MutexGuard<'_, Map> { + self.metadata_updates + .lock() + .unwrap_or_else(PoisonError::into_inner) + } + + /// Like metadata, resumable state is never locked across an await, and + /// each update leaves it consistent, so a poisoned lock is still usable. + fn lock_resumable(&self) -> MutexGuard<'_, ResumableState> { + self.resumable + .lock() + .unwrap_or_else(PoisonError::into_inner) + } + + /// Runs a named resumable step, skipping work completed by an earlier + /// failed attempt. + /// + /// Await steps sequentially. Nested steps are supported, but concurrent + /// steps do not define a checkpoint order. A step may fail with any error + /// convertible into [`BoxError`], including `anyhow::Error`; it is + /// returned as the source of [`Error::ResumableStep`]. + /// + /// # Errors + /// + /// Returns [`Error::ResumableStep`] with the step's error, and a runtime + /// error when the context doesn't belong to a job being worked. + /// + /// # Cancel safety + /// + /// Dropping the future drops the step's future with it. The step isn't + /// recorded as completed or as started, so it can run again, in this + /// attempt or a later one. + pub async fn resumable_step(&self, name: &str, step: F) -> Result<(), Error> + where + E: Into, + F: FnOnce() -> Fut, + Fut: Future>, + { + let guard = match self.begin_resumable_step(name, false)? { + StepAction::Run(previous) => RunningStep::new(self, name, previous), + StepAction::Skip => return Ok(()), + }; + + let result = step().await; + let mut state = guard.finish(); + match result { + Ok(()) => { + state.completed_step = Some(name.to_owned()); + Ok(()) + } + Err(error) => Err(state.fail_step(name, error.into())), + } + } + + /// Runs a named resumable step with the last cursor recorded for that step. + /// + /// Errors are handled as in [`WorkContext::resumable_step`]. + /// + /// # Errors + /// + /// Returns the errors of [`WorkContext::resumable_step`], and a runtime + /// error when the recorded cursor can't be decoded as `T`. + /// + /// # Cancel safety + /// + /// As with [`WorkContext::resumable_step`], a dropped step isn't recorded + /// as completed. A cursor it set is kept, so a later attempt resumes the + /// step from that cursor. + pub async fn resumable_step_with_cursor( + &self, + name: &str, + step: F, + ) -> Result<(), Error> + where + E: Into, + F: FnOnce(T) -> Fut, + Fut: Future>, + T: Default + DeserializeOwned, + { + let guard = match self.begin_resumable_step(name, true)? { + StepAction::Run(previous) => RunningStep::new(self, name, previous), + StepAction::Skip => return Ok(()), + }; + + let cursor = { + let state = self.lock_resumable(); + state + .cursors + .get(name) + .cloned() + .map(serde_json::from_value) + .transpose() + }; + let cursor = match cursor { + Ok(cursor) => cursor.unwrap_or_default(), + Err(error) => { + let mut state = guard.finish(); + return Err(state.fail_step(name, Box::new(error))); + } + }; + let result = step(cursor).await; + let mut state = guard.finish(); + match result { + Ok(()) => { + state.completed_step = Some(name.to_owned()); + state.cursors.remove(name); + Ok(()) + } + Err(error) => Err(state.fail_step(name, error.into())), + } + } + + /// Records progress for the currently running resumable cursor step. + /// + /// # Errors + /// + /// Returns an error when called outside a resumable step or when `cursor` + /// can't be serialized to JSON. + pub fn resumable_set_cursor(&self, cursor: &T) -> Result<(), Error> { + let cursor = serde_json::to_value(cursor)?; + let mut state = self.lock_resumable(); + let step_name = state.step_name.clone().ok_or_else(|| { + Error::runtime_context( + "worker context", + "resumable cursor can only be set inside a resumable cursor step".to_owned(), + ) + })?; + state.cursors.insert(step_name, cursor); + Ok(()) + } + + /// Persists the current resumable step in a caller-managed transaction. + /// + /// # Errors + /// + /// Returns a runtime error outside a resumable step, and the database + /// error when the update fails. + /// + /// # Cancel safety + /// + /// Dropping the future before it finishes leaves what it already ran in + /// the caller's transaction; roll that transaction back rather than + /// committing it. + pub async fn resumable_set_step_tx<'executor, E>(&self, connection: E) -> Result + where + E: DatabaseTransactionExecutor<'executor>, + { + self.resumable_checkpoint_tx::(connection, None) + .await + } + + /// Persists the current resumable step and cursor in a caller-managed + /// transaction. + /// + /// # Errors + /// + /// Returns a runtime error outside a resumable cursor step, a JSON error + /// when the cursor can't be encoded, and the database error when the update + /// fails. + /// + /// # Cancel safety + /// + /// Like [`WorkContext::resumable_set_step_tx`], a dropped future leaves + /// what it already ran in the caller's transaction. + pub async fn resumable_set_step_cursor_tx<'executor, T, E>( + &self, + connection: E, + cursor: &T, + ) -> Result + where + E: DatabaseTransactionExecutor<'executor>, + T: Serialize, + { + self.resumable_checkpoint_tx(connection, Some(cursor)).await + } + + async fn resumable_checkpoint_tx<'executor, T, E>( + &self, + connection: E, + cursor: Option<&T>, + ) -> Result + where + E: DatabaseTransactionExecutor<'executor>, + T: Serialize, + { + let metadata = self.resumable_checkpoint(cursor)?; + let (client, job_id) = self.current_job()?; + client + .job_update_tx( + connection, + job_id, + JobUpdateParams { + metadata, + output: None, + }, + ) + .await + } + + /// Marks the current step complete, records `cursor` for it, and returns + /// the checkpoint metadata to persist. + fn resumable_checkpoint( + &self, + cursor: Option<&T>, + ) -> Result, Error> { + let mut state = self.lock_resumable(); + let step_name = state.step_name.clone().ok_or_else(|| { + Error::runtime_context( + "worker context", + "resumable checkpoint must be set inside a resumable step".to_owned(), + ) + })?; + state.completed_step = Some(step_name.clone()); + if let Some(cursor) = cursor { + state + .cursors + .insert(step_name.clone(), serde_json::to_value(cursor)?); + } + let mut metadata = Map::new(); + metadata.insert( + crate::METADATA_KEY_RESUMABLE_STEP.to_owned(), + step_name.into(), + ); + if !state.cursors.is_empty() { + metadata.insert( + crate::METADATA_KEY_RESUMABLE_CURSOR.to_owned(), + Value::Object(state.cursors.clone()), + ); + } + Ok(metadata) + } + + fn begin_resumable_step(&self, name: &str, cursor_step: bool) -> Result { + if name.is_empty() { + return Err(Error::runtime_context( + "worker context", + "resumable step name cannot be empty".to_owned(), + )); + } + let mut state = self.lock_resumable(); + if let Some(failure) = &state.failure { + return Err(Error::runtime_source( + "worker context", + failure.to_string(), + failure.clone(), + )); + } + if !state.all_step_names.insert(name.to_owned()) { + let message = format!("duplicate resumable step name {name:?}"); + state.failure = Some(WorkError::new(Box::new(Error::runtime_context( + "worker context", + message.clone(), + )))); + return Err(Error::runtime_context("worker context", message)); + } + if !state.resume_matched { + if state.resume_step.as_deref() == Some(name) { + state.completed_step = Some(name.to_owned()); + state.resume_matched = true; + if !cursor_step || !state.cursors.contains_key(name) { + return Ok(StepAction::Skip); + } + } else { + return Ok(StepAction::Skip); + } + } + let previous = state.step_name.replace(name.to_owned()); + Ok(StepAction::Run(previous)) + } + + fn current_job(&self) -> Result<(&Client, i64), Error> { + self.client + .as_ref() + .zip(self.job_id) + .ok_or_else(|| { + Error::runtime_context( + "worker context", + "transactional context operation requires a WorkContext supplied to a running River worker" + .to_owned(), + ) + }) + } + + /// Returns a snapshot of metadata recorded during this attempt. + pub(crate) fn metadata_updates(&self) -> Map { + self.lock_metadata().clone() + } + + pub(crate) fn for_job( + client: Client, + cancellation: CancellationToken, + job_id: i64, + metadata: &JobMetadata, + ) -> Self { + let state = ResumableState::from_metadata(metadata); + Self { + cancellation, + client: Some(client), + job_id: Some(job_id), + metadata_updates: Arc::new(Mutex::new(Map::new())), + peers: None, + resumable: Arc::new(Mutex::new(state)), + } + } + + /// Attaches the attempt's peer ledger. + pub(crate) fn with_peers(mut self, peers: Arc) -> Self { + self.peers = Some(peers); + self + } + + /// Returns the attempt's peer ledger, for an attempt a client runs. + pub(crate) const fn peers(&self) -> Option<&Arc> { + self.peers.as_ref() + } + + /// Creates a detached attempt context using persisted resumable metadata. + #[must_use] + pub(crate) fn for_test_job(job: &JobRow) -> Self { + let mut context = Self::new(CancellationToken::new()); + context.resumable = Arc::new(Mutex::new(ResumableState::from_metadata(&job.metadata))); + context + } + + /// Validates checkpoint metadata before invoking user work. + pub(crate) fn resumable_validate(&self) -> Result<(), WorkError> { + match &self.lock_resumable().failure { + Some(error) => Err(error.clone()), + None => Ok(()), + } + } + + /// Resolves attempt-scoped resumable errors and metadata for runtime/test parity. + pub(crate) fn resumable_finish(&self, worker_failed: bool) -> Option { + let state = self.lock_resumable(); + let failure = state.failure.clone().or_else(|| { + (!worker_failed && !state.resume_matched).then(|| { + WorkError::new(Box::new(Error::runtime_context( + "worker context", + format!( + "resumable step {:?} not found in worker", + state.resume_step.as_deref().unwrap_or_default() + ), + ))) + }) + }); + if (worker_failed || failure.is_some()) + && let Some(completed_step) = &state.completed_step + { + let mut updates = self.lock_metadata(); + updates.insert( + crate::METADATA_KEY_RESUMABLE_STEP.to_owned(), + completed_step.clone().into(), + ); + if state.cursors.is_empty() { + if state.had_cursors { + updates.insert(crate::METADATA_KEY_RESUMABLE_CURSOR.to_owned(), Value::Null); + } + } else { + updates.insert( + crate::METADATA_KEY_RESUMABLE_CURSOR.to_owned(), + Value::Object(state.cursors.clone()), + ); + } + } + failure + } +} + +impl std::fmt::Debug for WorkContext { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WorkContext") + .field("job_id", &self.job_id) + .field("cancelled", &self.cancellation.is_cancelled()) + .field("metadata_updates", &*self.lock_metadata()) + .finish_non_exhaustive() + } +} + +#[derive(Debug)] +struct ResumableState { + all_step_names: HashSet, + completed_step: Option, + cursors: Map, + failure: Option, + had_cursors: bool, + resume_matched: bool, + resume_step: Option, + step_name: Option, +} + +impl Default for ResumableState { + fn default() -> Self { + Self { + all_step_names: HashSet::new(), + completed_step: None, + cursors: Map::new(), + failure: None, + had_cursors: false, + resume_matched: true, + resume_step: None, + step_name: None, + } + } +} + +impl ResumableState { + fn from_metadata(metadata: &JobMetadata) -> Self { + let mut state = Self::default(); + state.resume_step = metadata + .get::(crate::METADATA_KEY_RESUMABLE_STEP) + .ok() + .flatten() + .filter(|step| !step.is_empty()); + state.resume_matched = state.resume_step.is_none(); + match metadata.get_raw(crate::METADATA_KEY_RESUMABLE_CURSOR) { + Some(raw) if raw.get().starts_with('{') => { + if let Ok(cursors) = serde_json::from_str::>(raw.get()) { + state.had_cursors = !cursors.is_empty(); + state.cursors = cursors; + } else { + state.failure = Some(WorkError::new(Box::new(Error::invalid_job( + "river:resumable_cursor cannot be decoded", + )))); + } + } + Some(raw) if raw.get().starts_with('[') => { + state.failure = Some(WorkError::new(Box::new(Error::invalid_job( + "river:resumable_cursor must be an object when present", + )))); + } + _ => {} + } + state + } + + fn fail_step(&mut self, name: &str, error: BoxError) -> Error { + let source = WorkError::new(error); + self.failure = Some(WorkError::new(Box::new(Error::ResumableStep { + name: name.to_owned(), + source: Box::new(source.clone()), + }))); + Error::ResumableStep { + name: name.to_owned(), + source: Box::new(source), + } + } +} + +enum StepAction { + Run(Option), + Skip, +} + +/// A resumable step that has started. Finishing it restores the enclosing +/// step. Dropping it unfinished, when the step's future is dropped, also +/// forgets that the step started, so the step can run again in this attempt +/// and no cursor can be set for it outside the step. +struct RunningStep<'a> { + context: &'a WorkContext, + finished: bool, + name: &'a str, + previous: Option, +} + +impl<'a> RunningStep<'a> { + const fn new(context: &'a WorkContext, name: &'a str, previous: Option) -> Self { + Self { + context, + finished: false, + name, + previous, + } + } + + /// Restores the enclosing step and returns the locked state. + fn finish(mut self) -> MutexGuard<'a, ResumableState> { + self.finished = true; + let mut state = self.context.lock_resumable(); + state.step_name = self.previous.take(); + state + } +} + +impl Drop for RunningStep<'_> { + fn drop(&mut self) { + if !self.finished { + let mut state = self.context.lock_resumable(); + state.step_name = self.previous.take(); + state.all_step_names.remove(self.name); + } + } +} + +/// Successful control outcome returned by a worker. +#[derive(Clone, Debug, Default, Eq, PartialEq)] +#[non_exhaustive] +pub enum WorkOutcome { + /// Mark the job cancelled. + Cancel, + /// Mark the job complete. + #[default] + Complete, + /// Discard without another attempt. + Discard, + /// Reschedule without consuming an attempt. + Snooze(Duration), +} + +/// Per-worker timeout selection. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +#[non_exhaustive] +pub enum WorkerTimeout { + /// Use the client-wide timeout. + #[default] + ClientDefault, + /// Do not time out this kind of job. + Disabled, + /// Cancel the job after this duration. + After(Duration), +} + +/// A typed asynchronous job worker. +pub trait Worker: Send + Sync + 'static +where + A: JobArgs, +{ + /// Worker-specific error type. Errors use River's retry policy. + /// + /// The job's recorded error is the error's message followed by each of + /// its sources, `outer: inner`, so the cause of a failure stays visible + /// on the job. + /// + /// Any error convertible into [`BoxError`] works, including concrete + /// error types, `Box`, and report types such as + /// `anyhow::Error` or `eyre::Report`. Hooks and error handlers receive it + /// as a [`WorkError`], whose [`downcast_ref`](WorkError::downcast_ref) + /// returns the concrete error type. A report type converts into its own + /// wrapper, which keeps its message and source chain but can't be + /// downcast to the type it wraps; return a concrete error type when an + /// extension needs to downcast it. + type Error: Into; + + /// Overrides the client retry delay for this job. Returning `None` uses the + /// client policy. + fn next_retry( + &self, + _job: &Job, + _error: &WorkError, + _now: DateTime, + ) -> Option { + None + } + + /// Overrides the client timeout for this job. + fn timeout(&self, _job: &Job) -> WorkerTimeout { + WorkerTimeout::ClientDefault + } + + /// Executes a job. + /// + /// Implementations can use `async fn`; the explicit return type guarantees + /// that the resulting future can run on River's multithreaded Tokio + /// runtime without requiring each implementation to box its future. + fn work( + &self, + context: WorkContext, + job: Job, + ) -> impl Future> + Send; +} + +/// Type-erased adapter from persisted rows to a typed [`Worker`]. +#[async_trait] +trait ErasedWorker: Send + Sync { + fn next_retry( + &self, + row: &JobRow, + error: &WorkError, + now: DateTime, + ) -> Result, Box>; + + fn timeout(&self, row: &JobRow) -> Result>; + + /// Runs one attempt. Arguments are decoded once, and the worker's + /// timeout for the decoded job is reported through `timeout` before + /// work starts. The outer error reports arguments that couldn't be + /// decoded, in which case the worker didn't run. + async fn work( + &self, + context: WorkContext, + row: JobRow, + timeout: oneshot::Sender, + ) -> Result, WorkError>; +} + +struct FunctionWorker { + function: F, +} + +impl Worker for FunctionWorker +where + A: JobArgs, + E: Into, + F: Fn(WorkContext, Job) -> Fut + Send + Sync + 'static, + Fut: Future> + Send, +{ + type Error = E; + + fn work( + &self, + context: WorkContext, + job: Job, + ) -> impl Future> + Send { + (self.function)(context, job) + } +} + +struct RegisteredWorker { + worker: W, + _args: std::marker::PhantomData, +} + +#[async_trait] +impl ErasedWorker for RegisteredWorker +where + A: JobArgs, + W: Worker, +{ + fn next_retry( + &self, + row: &JobRow, + error: &WorkError, + now: DateTime, + ) -> Result, Box> { + // `work` consumes its job, so a failed attempt decodes once more to + // consult the worker's retry override. + let job = Job { + args: row.decode_args()?, + row: row.clone(), + }; + Ok(Worker::::next_retry(&self.worker, &job, error, now)) + } + + fn timeout(&self, row: &JobRow) -> Result> { + let job = Job { + args: row.decode_args()?, + row: row.clone(), + }; + Ok(Worker::::timeout(&self.worker, &job)) + } + + async fn work( + &self, + context: WorkContext, + row: JobRow, + timeout: oneshot::Sender, + ) -> Result, WorkError> { + let job = Job { + args: row.decode_args().map_err(WorkError::new)?, + row, + }; + // The supervisor may have stopped waiting for a timeout; that is not + // an error for the attempt. + let _ = timeout.send(Worker::::timeout(&self.worker, &job)); + // Sending wakes the supervisor on this worker thread, where Tokio + // may hold it in a slot other threads can't steal. Yield once so it + // starts the timeout before a worker that blocks the thread (which + // the supervisor exists to detect) can delay it. + tokio::task::yield_now().await; + Ok(self + .worker + .work(context, job) + .await + .map_err(|error| WorkError::new(error.into()))) + } +} + +/// Type-erased collection of workers keyed by job kind. +#[derive(Clone, Default)] +pub struct WorkerRegistry { + workers: HashMap<&'static str, Arc>, +} + +impl std::fmt::Debug for WorkerRegistry { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("WorkerRegistry") + .field("kinds", &self.kinds()) + .finish_non_exhaustive() + } +} + +impl WorkerRegistry { + /// Creates an empty worker registry. + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Returns registered kinds in stable order. + #[must_use] + pub fn kinds(&self) -> Vec<&'static str> { + let mut kinds = self.workers.keys().copied().collect::>(); + kinds.sort_unstable(); + kinds + } + + pub(crate) fn contains_kind(&self, kind: &str) -> bool { + self.workers.contains_key(kind) + } + + pub(crate) fn next_retry( + &self, + row: &JobRow, + error: &WorkError, + now: DateTime, + ) -> Result, Box> { + self.worker_for(row)?.next_retry(row, error, now) + } + + /// Evaluates the worker timeout for a persisted row outside an attempt, + /// such as when rescuing stuck jobs. + pub(crate) fn timeout( + &self, + row: &JobRow, + ) -> Result> { + self.worker_for(row)?.timeout(row) + } + + /// Registers one worker, rejecting duplicate kinds. + /// + /// # Errors + /// + /// Returns [`Error::InvalidJob`] when the kind or one of its aliases is + /// empty or 128 bytes or longer, or when a worker is already registered + /// for it. + pub fn register(&mut self, worker: W) -> Result<&mut Self, Error> + where + A: JobArgs, + W: Worker, + { + if A::KIND.is_empty() || A::KIND.len() >= 128 { + return Err(Error::invalid_job_context( + "worker registration", + format!( + "job kind must contain between 1 and 127 bytes: {:?}", + A::KIND + ), + )); + } + let mut kinds = vec![A::KIND]; + for alias in A::kind_aliases() { + if alias.is_empty() || alias.len() >= 128 { + return Err(Error::invalid_job_context( + "worker registration", + format!("job kind alias must contain between 1 and 127 bytes: {alias:?}"), + )); + } + if kinds.contains(alias) || self.workers.contains_key(alias) { + return Err(Error::invalid_job_context( + "worker registration", + format!("worker already registered for kind {alias:?}"), + )); + } + kinds.push(alias); + } + if self.workers.contains_key(A::KIND) { + return Err(Error::invalid_job_context( + "worker registration", + format!("worker already registered for kind {:?}", A::KIND), + )); + } + let worker: Arc = Arc::new(RegisteredWorker:: { + worker, + _args: std::marker::PhantomData, + }); + for kind in kinds { + self.workers.insert(kind, Arc::clone(&worker)); + } + Ok(self) + } + + /// Registers an asynchronous function or closure as a worker. + /// + /// The function may return any error convertible into [`BoxError`], such + /// as `anyhow::Result`; see [`Worker::Error`]. Use [`Worker`] + /// instead when a job kind needs to override its timeout or retry + /// schedule. + /// + /// # Errors + /// + /// Returns an error when the job kind or one of its aliases is invalid or + /// already registered. + pub fn register_fn(&mut self, function: F) -> Result<&mut Self, Error> + where + A: JobArgs, + E: Into, + F: Fn(WorkContext, Job) -> Fut + Send + Sync + 'static, + Fut: Future> + Send, + { + self.register::(FunctionWorker { function }) + } + + /// Returns an error for a row whose kind has no registered worker, which + /// River fails before running any hook or middleware, as River Go does. + pub(crate) fn check_kind(&self, row: &JobRow) -> Result<(), WorkError> { + self.worker_for(row).map(|_| ()).map_err(WorkError::new) + } + + /// Runs one attempt of `row`, decoding its arguments once. The worker's + /// timeout for the job is sent on `timeout` before work starts; the + /// sender is dropped without a value when the attempt fails first. The + /// outer error reports an unknown kind or arguments that couldn't be + /// decoded, in which case the worker didn't run. + pub(crate) async fn work( + &self, + context: WorkContext, + row: JobRow, + timeout: oneshot::Sender, + ) -> Result, WorkError> { + let worker = self.worker_for(&row).map_err(WorkError::new)?; + worker.work(context, row, timeout).await + } + + fn worker_for( + &self, + row: &JobRow, + ) -> Result<&Arc, Box> { + self.workers.get(row.kind.as_str()).ok_or_else(|| { + Box::new(UnregisteredKind(row.kind.clone())) as Box + }) + } +} + +/// The attempt error River records for a job whose kind has no worker. Its +/// text is part of the protocol: every implementation records River Go's +/// wording, so a job's errors read the same whichever client worked it. +#[derive(Debug, thiserror::Error)] +#[error("job kind is not registered in the client's Workers bundle: {0}")] +struct UnregisteredKind(String); + +/// Maximum encoded size of recorded output (Go `maxOutputSizeBytes`). +const MAX_OUTPUT_BYTES: usize = 32 * 1024 * 1024; + +/// Rejects output whose JSON is larger than River Go allows. +pub(crate) fn check_output_size(output: &Value) -> Result<(), String> { + let size = crate::encoding::to_go_string(output) + .map_err(|error| error.to_string())? + .len(); + if size > MAX_OUTPUT_BYTES { + return Err(format!( + "output is too large: {size} bytes (max {} MB)", + MAX_OUTPUT_BYTES / 1024 / 1024 + )); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use std::{ + fmt, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + }; + + use chrono::Utc; + use serde::{Deserialize, Serialize}; + use serde_json::json; + + use super::*; + use crate::JobState; + + #[derive(Debug)] + struct FunctionError; + + impl fmt::Display for FunctionError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + formatter.write_str("function worker failed") + } + } + + impl StdError for FunctionError {} + + #[derive(Debug, Deserialize, Serialize)] + struct FunctionJobArgs { + fail: bool, + } + + impl JobArgs for FunctionJobArgs { + const KIND: &'static str = "function_worker"; + + fn kind_aliases() -> &'static [&'static str] { + &["function_worker_v1"] + } + } + + async fn function_worker( + context: WorkContext, + job: Job, + ) -> Result { + if job.args.fail { + return Err(FunctionError); + } + context.record_output(json!({"function": true})).unwrap(); + Ok(WorkOutcome::Complete) + } + + fn job_row(kind: &str, fail: bool) -> JobRow { + let now = Utc::now(); + JobRow { + attempt: 1, + attempted_at: Some(now), + attempted_by: vec!["test".to_owned()], + created_at: now, + encoded_args: serde_json::value::to_raw_value(&json!({"fail": fail})).unwrap(), + errors: Vec::new(), + finalized_at: None, + id: 1, + kind: kind.to_owned(), + max_attempts: 25, + metadata: JobMetadata::default(), + priority: 1, + queue: "default".to_owned(), + scheduled_at: now, + state: JobState::Running, + tags: Vec::new(), + unique_key: None, + unique_states: None, + } + } + + #[tokio::test] + async fn resumable_context_runs_without_a_checkpoint() { + for metadata in [json!({}), json!({"river:resumable_step": ""})] { + let mut row = job_row(FunctionJobArgs::KIND, false); + row.metadata = metadata.as_object().unwrap().clone().into(); + let context = WorkContext::for_test_job(&row); + let mut ran = false; + context + .resumable_step("first", || async { + ran = true; + Ok::<_, FunctionError>(()) + }) + .await + .unwrap(); + assert!(ran); + assert!(context.resumable_finish(false).is_none()); + } + } + + #[tokio::test] + async fn resumable_cursor_decode_failure_is_sticky() { + let mut row = job_row(FunctionJobArgs::KIND, false); + row.metadata = json!({ + "river:resumable_step": "first", + "river:resumable_cursor": { "second": "not a number" } + }) + .as_object() + .unwrap() + .clone() + .into(); + let context = WorkContext::for_test_job(&row); + context + .resumable_step("first", || async { + panic!("already completed"); + #[allow(unreachable_code)] + Ok::<_, FunctionError>(()) + }) + .await + .unwrap(); + let error = context + .resumable_step_with_cursor("second", |_: i64| async { + panic!("invalid cursor must not reach worker"); + #[allow(unreachable_code)] + Ok::<_, FunctionError>(()) + }) + .await + .unwrap_err(); + assert!( + crate::error::Chain(&error) + .to_string() + .contains("invalid type"), + "{error}" + ); + assert!(context.resumable_set_cursor(&1).is_err()); + assert!(context.resumable_finish(false).is_some()); + assert_eq!( + context.metadata_updates()[crate::METADATA_KEY_RESUMABLE_STEP], + "first" + ); + } + + #[tokio::test] + async fn resumable_dropped_step_can_run_again() { + let context = WorkContext::new(CancellationToken::new()); + // Drop the step's future while the step is running. + let dropped = tokio::time::timeout( + std::time::Duration::from_millis(10), + context.resumable_step_with_cursor("first", |_: i64| { + let context = context.clone(); + async move { + context.resumable_set_cursor(&1)?; + std::future::pending::<()>().await; + Ok::<_, Error>(()) + } + }), + ) + .await; + assert!(dropped.is_err()); + + // The step no longer runs, so no cursor can be set outside it, and + // the same step runs again with the cursor it recorded. + assert!(context.resumable_set_cursor(&2).is_err()); + let mut seen = None; + context + .resumable_step_with_cursor("first", |cursor: i64| { + seen = Some(cursor); + async { Ok::<_, Error>(()) } + }) + .await + .unwrap(); + assert_eq!(seen, Some(1)); + context + .resumable_step("second", || async { Ok::<_, Error>(()) }) + .await + .unwrap(); + } + + #[tokio::test] + async fn resumable_nested_steps_restore_parent_and_error_sources() { + let context = WorkContext::new(CancellationToken::new()); + let error = context + .resumable_step_with_cursor("outer", |_: i64| async { + context + .resumable_step("inner", || async { Ok::<_, Error>(()) }) + .await?; + context.resumable_set_cursor(&7)?; + Err::<(), _>(Error::ResumableStep { + name: "source".to_owned(), + source: Box::new(FunctionError), + }) + }) + .await + .unwrap_err(); + let mut source: &(dyn StdError + 'static) = &error; + while !source.is::() { + source = source.source().expect("preserved source"); + } + let finished = context.resumable_finish(false).unwrap(); + let mut source: &(dyn StdError + 'static) = &finished; + while !source.is::() { + source = source.source().expect("preserved suppressed source"); + } + assert_eq!( + context.metadata_updates(), + json!({ + "river:resumable_step": "inner", + "river:resumable_cursor": {"outer": 7} + }) + .as_object() + .unwrap() + .clone() + ); + } + + #[tokio::test] + async fn register_fn_accepts_capturing_closure() { + let calls = Arc::new(AtomicUsize::new(0)); + let calls_for_worker = Arc::clone(&calls); + let mut workers = WorkerRegistry::new(); + workers + .register_fn(move |_context: WorkContext, _job: Job| { + let calls = Arc::clone(&calls_for_worker); + async move { + calls.fetch_add(1, Ordering::Relaxed); + Ok::<_, FunctionError>(WorkOutcome::Snooze(Duration::from_secs(1))) + } + }) + .unwrap(); + + let (timeout_sender, _timeout_receiver) = oneshot::channel(); + let outcome = workers + .work( + WorkContext::new(CancellationToken::new()), + job_row(FunctionJobArgs::KIND, false), + timeout_sender, + ) + .await + .unwrap() + .unwrap(); + + assert_eq!(outcome, WorkOutcome::Snooze(Duration::from_secs(1))); + assert_eq!(calls.load(Ordering::Relaxed), 1); + } + + #[tokio::test] + async fn register_fn_handles_aliases_and_typed_errors() { + let mut workers = WorkerRegistry::new(); + workers.register_fn(function_worker).unwrap(); + + assert_eq!( + workers.kinds(), + [FunctionJobArgs::KIND, "function_worker_v1"] + ); + + let (timeout_sender, _timeout_receiver) = oneshot::channel(); + let error = workers + .work( + WorkContext::new(CancellationToken::new()), + job_row("function_worker_v1", true), + timeout_sender, + ) + .await + .unwrap() + .unwrap_err(); + assert!(error.get_ref().downcast_ref::().is_some()); + } + + static COUNTED_DECODES: AtomicUsize = AtomicUsize::new(0); + + #[derive(Serialize)] + struct CountedArgs { + timeout_ms: u64, + } + + impl<'de> Deserialize<'de> for CountedArgs { + fn deserialize>(deserializer: D) -> Result { + #[derive(Deserialize)] + struct Fields { + timeout_ms: u64, + } + COUNTED_DECODES.fetch_add(1, Ordering::SeqCst); + let fields = Fields::deserialize(deserializer)?; + Ok(Self { + timeout_ms: fields.timeout_ms, + }) + } + } + + impl JobArgs for CountedArgs { + const KIND: &'static str = "counted_args"; + } + + struct CountedWorker; + + impl Worker for CountedWorker { + type Error = FunctionError; + + fn timeout(&self, job: &Job) -> WorkerTimeout { + WorkerTimeout::After(Duration::from_millis(job.args.timeout_ms)) + } + + fn work( + &self, + _context: WorkContext, + job: Job, + ) -> impl Future> + Send { + assert_eq!(job.row.kind, CountedArgs::KIND); + std::future::ready(Ok(WorkOutcome::Complete)) + } + } + + #[tokio::test] + async fn work_decodes_args_once_and_reports_timeout_first() { + let mut workers = WorkerRegistry::new(); + workers.register(CountedWorker).unwrap(); + let mut row = job_row(CountedArgs::KIND, false); + row.encoded_args = serde_json::value::to_raw_value(&json!({"timeout_ms": 1234})).unwrap(); + + let (timeout_sender, timeout_receiver) = oneshot::channel(); + let outcome = workers + .work( + WorkContext::new(CancellationToken::new()), + row.clone(), + timeout_sender, + ) + .await + .unwrap() + .unwrap(); + + assert_eq!(outcome, WorkOutcome::Complete); + assert_eq!( + timeout_receiver.await.unwrap(), + WorkerTimeout::After(Duration::from_millis(1234)) + ); + assert_eq!(COUNTED_DECODES.load(Ordering::SeqCst), 1); + } + + #[tokio::test] + async fn work_with_undecodable_args_fails_without_reporting_timeout() { + let mut workers = WorkerRegistry::new(); + workers.register_fn(function_worker).unwrap(); + let mut row = job_row(FunctionJobArgs::KIND, false); + row.encoded_args = serde_json::value::to_raw_value(&json!({"fail": "no"})).unwrap(); + + let (timeout_sender, timeout_receiver) = oneshot::channel(); + let error = workers + .work( + WorkContext::new(CancellationToken::new()), + row.clone(), + timeout_sender, + ) + .await + .unwrap_err(); + + assert!(error.get_ref().is::()); + assert!(timeout_receiver.await.is_err()); + } + + async fn run_once(workers: &WorkerRegistry, row: &JobRow) -> Result { + let (timeout_sender, _timeout_receiver) = oneshot::channel(); + workers + .work( + WorkContext::new(CancellationToken::new()), + row.clone(), + timeout_sender, + ) + .await + .and_then(|result| result) + } + + async fn anyhow_function_worker( + _context: WorkContext, + job: Job, + ) -> anyhow::Result { + use anyhow::Context as _; + + if job.args.fail { + return Err(std::io::Error::other("disk full")).context("writing report"); + } + Ok(WorkOutcome::Complete) + } + + #[tokio::test] + async fn register_fn_accepts_anyhow_results() { + let mut workers = WorkerRegistry::new(); + workers.register_fn(anyhow_function_worker).unwrap(); + + assert_eq!( + run_once(&workers, &job_row(FunctionJobArgs::KIND, false)) + .await + .unwrap(), + WorkOutcome::Complete + ); + let error = run_once(&workers, &job_row(FunctionJobArgs::KIND, true)) + .await + .unwrap_err(); + assert_eq!(error.to_string(), "writing report"); + // The report's source chain is preserved for inspection. + let root = error + .get_ref() + .source() + .and_then(|source| source.downcast_ref::()) + .expect("anyhow context source"); + assert_eq!(root.to_string(), "disk full"); + } + + struct BoxedErrorWorker; + + impl Worker for BoxedErrorWorker { + type Error = BoxError; + + async fn work( + &self, + _context: WorkContext, + _job: Job, + ) -> Result { + tokio::task::yield_now().await; + Err(Box::new(FunctionError)) + } + } + + #[tokio::test] + async fn boxed_worker_errors_remain_downcastable() { + let mut workers = WorkerRegistry::new(); + workers.register(BoxedErrorWorker).unwrap(); + + let error = run_once(&workers, &job_row(FunctionJobArgs::KIND, false)) + .await + .unwrap_err(); + assert!(error.get_ref().downcast_ref::().is_some()); + } + + #[tokio::test] + async fn resumable_steps_accept_anyhow_errors() { + let context = WorkContext::new(CancellationToken::new()); + let error = context + .resumable_step("first", || async { + Err::<(), _>(anyhow::anyhow!("step failed")) + }) + .await + .unwrap_err(); + + assert!(matches!(&error, Error::ResumableStep { name, .. } if name == "first")); + assert_eq!(error.source().unwrap().to_string(), "step failed"); + + let context = WorkContext::new(CancellationToken::new()); + let error = context + .resumable_step_with_cursor("second", |_: i64| async { + Err::<(), _>(anyhow::anyhow!("cursor step failed")) + }) + .await + .unwrap_err(); + assert!(matches!(&error, Error::ResumableStep { name, .. } if name == "second")); + assert_eq!(error.source().unwrap().to_string(), "cursor step failed"); + } + + #[test] + fn metadata_set_and_record_output_serialize_values() { + #[derive(Serialize)] + struct Receipt { + delivered: bool, + } + + let context = WorkContext::new(CancellationToken::new()); + context.metadata_set("attempts", 3).unwrap(); + context.metadata_set("attempts", 4).unwrap(); + context.record_output(Receipt { delivered: true }).unwrap(); + + let bad_output = std::collections::BTreeMap::from([((1, 2), true)]); + assert!(context.record_output(&bad_output).is_err()); + assert!(context.record_output(vec![1.0, f64::NAN]).is_err()); + assert!(context.metadata_set("ratio", f64::INFINITY).is_err()); + assert_eq!( + context.metadata_updates(), + json!({"attempts": 4, "output": {"delivered": true}}) + .as_object() + .unwrap() + .clone() + ); + } + + #[test] + fn work_context_debug_shows_attempt_state() { + let cancellation = CancellationToken::new(); + let context = WorkContext::new(cancellation.clone()); + context.metadata_set("attempts", 3).unwrap(); + cancellation.cancel(); + + assert_eq!( + format!("{context:?}"), + r#"WorkContext { job_id: None, cancelled: true, metadata_updates: {"attempts": Number(3)}, .. }"# + ); + } + + #[test] + fn register_fn_rejects_duplicate_kinds() { + let mut workers = WorkerRegistry::new(); + workers.register_fn(function_worker).unwrap(); + + let Err(error) = workers.register_fn(function_worker) else { + panic!("duplicate registration should fail"); + }; + + assert!(error.to_string().contains("already registered")); + } + + #[test] + fn registry_debug_lists_kinds_without_worker_internals() { + let mut registry = WorkerRegistry::new(); + registry + .register_fn( + |_context: WorkContext, _job: Job| async move { + Ok::<_, std::io::Error>(WorkOutcome::Complete) + }, + ) + .unwrap(); + + let debug = format!("{registry:?}"); + assert!(debug.contains(FunctionJobArgs::KIND)); + assert!(!debug.contains("dyn ErasedWorker")); + } + + #[test] + fn recorded_output_is_limited_like_go() { + let context = crate::__private::work_context(tokio_util::sync::CancellationToken::new()); + let limit = super::MAX_OUTPUT_BYTES; + // A JSON string's two quotes count toward the limit. + context.record_output("x".repeat(limit - 2)).unwrap(); + let error = context.record_output("x".repeat(limit - 1)).unwrap_err(); + assert!(error.to_string().contains("output is too large"), "{error}"); + assert_eq!( + context.metadata_updates()[crate::METADATA_KEY_OUTPUT] + .as_str() + .map(str::len), + Some(limit - 2) + ); + } +} diff --git a/rust/riverqueue/tests/caller_transactions.rs b/rust/riverqueue/tests/caller_transactions.rs new file mode 100644 index 000000000..7d6db181f --- /dev/null +++ b/rust/riverqueue/tests/caller_transactions.rs @@ -0,0 +1,706 @@ +//! Operations on a caller-managed transaction run directly in it. +//! +//! River opens no savepoint or nested transaction in a caller's +//! transaction, like River Go. When an operation fails, including in an +//! extension step or insert middleware after River's own write, whatever it +//! already wrote stays in the caller's transaction, which the caller rolls +//! back. Without a caller transaction, River's own transaction rolls the +//! whole operation back. +//! +//! PostgreSQL scenarios run in a unique schema and fail rather than skip when +//! `RIVER_RUST_DATABASE_URL` is unset; SQLite scenarios use temporary files. + +#![cfg(any(all(feature = "postgres", river_postgres_tests), feature = "sqlite"))] + +mod support; + +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, +}; + +use async_trait::async_trait; +use riverqueue::__private::{ + ClientBuilderExt, DatabaseConnection, JobSetStateParams, JobUpdatedParams, JobsInsertedParams, + Pilot, PilotError, +}; +use riverqueue::{ + BoxError, Client, Error, Hook, InsertContext, InsertMiddleware, InsertNext, InsertOpts, + InsertedJob, JobArgs, JobRow, JobState, +}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "caller_transaction")] +struct ScopeArgs { + name: String, +} + +fn args(name: &str) -> ScopeArgs { + ScopeArgs { + name: name.to_owned(), + } +} + +/// Writes a row of its own after each intercepted operation and then fails +/// while `fail` is set. With `fail_in_database` it instead runs a statement +/// the database rejects. +struct EffectPilot { + effect_table: String, + fail: Arc, + fail_in_database: bool, +} + +impl EffectPilot { + async fn effect(&self, connection: DatabaseConnection<'_>, id: i64) -> Result<(), PilotError> { + let sql = sqlx::AssertSqlSafe(if self.fail_in_database { + "SELECT * FROM river_nonexistent_table".to_owned() + } else { + format!("INSERT INTO {} (job_id) VALUES ({id})", self.effect_table) + }); + match connection { + #[cfg(feature = "postgres")] + DatabaseConnection::Postgres(connection) => { + sqlx::query(sql).execute(connection).await?; + } + #[cfg(feature = "sqlite")] + DatabaseConnection::Sqlite(connection) => { + sqlx::query(sql).execute(connection).await?; + } + #[allow(unreachable_patterns)] + _ => unreachable!("built-in backends only"), + } + if self.fail.load(Ordering::SeqCst) { + return Err("extension step failed after its write".into()); + } + Ok(()) + } +} + +#[async_trait] +impl Pilot for EffectPilot { + fn intercepts_insert(&self) -> bool { + true + } + + fn intercepts_job_cancel_retry(&self) -> bool { + true + } + + fn intercepts_job_set_state(&self) -> bool { + true + } + + async fn after_jobs_inserted( + &self, + connection: DatabaseConnection<'_>, + params: &JobsInsertedParams<'_>, + ) -> Result<(), PilotError> { + self.effect(connection, params.jobs[0].id).await + } + + async fn after_job_cancel( + &self, + connection: DatabaseConnection<'_>, + params: &JobUpdatedParams, + ) -> Result<(), PilotError> { + self.effect(connection, params.job.id).await + } + + async fn after_job_retry( + &self, + connection: DatabaseConnection<'_>, + params: &JobUpdatedParams, + ) -> Result<(), PilotError> { + self.effect(connection, params.job.id).await + } + + async fn after_jobs_set_state( + &self, + connection: DatabaseConnection<'_>, + params: &JobSetStateParams, + ) -> Result<(), PilotError> { + self.effect(connection, params.job_ids[0]).await + } +} + +/// Fails every insertion after River wrote it. +struct FailAfterWrite; + +impl InsertMiddleware for FailAfterWrite { + async fn insert_many( + &self, + jobs: Vec, + next: InsertNext<'_>, + ) -> Result, Error> { + next.run(jobs).await?; + Err(Error::extension( + riverqueue::ExtensionPhase::InsertMiddleware, + std::io::Error::other("middleware failed after the write"), + )) + } +} + +/// Fails to decode every inserted row. +struct FailingDecode; + +impl Hook for FailingDecode { + fn decode_insert_result( + &self, + _job: &mut JobRow, + ) -> impl Future> + Send { + std::future::ready(Err("decode failed on purpose".into())) + } +} + +/// Arguments that serialize but never deserialize, so decoding the inserted +/// row's arguments fails after the write. +#[derive(Clone, Debug, Serialize, JobArgs)] +#[river(kind = "caller_transaction_undecodable")] +struct UndecodableArgs { + name: String, +} + +impl<'de> Deserialize<'de> for UndecodableArgs { + fn deserialize>(_deserializer: D) -> Result { + Err(serde::de::Error::custom("never decodes")) + } +} + +fn assert_extension_error(error: &Error) { + assert!(matches!(error, Error::Extension { .. }), "{error}"); +} + +/// Defines each scenario for one backend's `Fixture`. +macro_rules! scenarios { + () => { + #[tokio::test(flavor = "multi_thread")] + async fn failed_insert_steps_stay_in_caller_transaction() { + let fixture = Fixture::new().await; + let fail = Arc::new(AtomicBool::new(true)); + let client = fixture.client(&fail); + let plain = fixture.builder().build().unwrap(); + + let mut tx = fixture.begin().await; + let prior = plain.insert(args("prior")).tx(&mut tx).await.unwrap().id(); + let error = client.insert(args("single")).tx(&mut tx).await.unwrap_err(); + assert_extension_error(&error); + let error = client + .insert_many([args("many_1"), args("many_2")]) + .tx(&mut tx) + .await + .unwrap_err(); + assert_extension_error(&error); + + // Without a savepoint, the failed insertions' writes and their + // extension steps' writes remain in the caller's transaction + // next to its earlier work. + let in_tx = fixture.job_ids_in(&mut tx).await; + assert_eq!(in_tx.len(), 4, "{in_tx:?}"); + assert!(in_tx.contains(&prior)); + assert_eq!(fixture.effect_ids_in(&mut tx).await.len(), 2); + tx.rollback().await.unwrap(); + + assert!(fixture.job_ids().await.is_empty()); + assert!(fixture.effect_ids().await.is_empty()); + + // Without a caller transaction, River's own transaction rolls + // the whole insertion back. + let error = client.insert(args("pool")).await.unwrap_err(); + assert_extension_error(&error); + assert!(fixture.job_ids().await.is_empty()); + assert!(fixture.effect_ids().await.is_empty()); + fixture.cleanup().await; + } + + // Decode hooks and argument decoding run after the write. Without a + // caller transaction an error from either means nothing was + // written; in a caller's transaction the row stays until the caller + // rolls back. + #[tokio::test(flavor = "multi_thread")] + async fn failed_result_decoding_after_the_write() { + let fixture = Fixture::new().await; + let hooked = fixture.builder().hook(FailingDecode).build().unwrap(); + let plain = fixture.builder().build().unwrap(); + + let error = hooked.insert(args("pool")).await.unwrap_err(); + assert_extension_error(&error); + let error = plain + .insert(UndecodableArgs { + name: "pool".to_owned(), + }) + .await + .unwrap_err(); + assert!(matches!(error, Error::Json(_)), "{error}"); + assert!(fixture.job_ids().await.is_empty()); + + let mut tx = fixture.begin().await; + let error = hooked.insert(args("tx")).tx(&mut tx).await.unwrap_err(); + assert_extension_error(&error); + let error = plain + .insert(UndecodableArgs { + name: "tx".to_owned(), + }) + .tx(&mut tx) + .await + .unwrap_err(); + assert!(matches!(error, Error::Json(_)), "{error}"); + assert_eq!(fixture.job_ids_in(&mut tx).await.len(), 2); + tx.rollback().await.unwrap(); + + assert!(fixture.job_ids().await.is_empty()); + fixture.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn failed_insert_middleware_after_the_write() { + let fixture = Fixture::new().await; + let client = fixture + .builder() + .insert_middleware(FailAfterWrite) + .build() + .unwrap(); + + let error = client.insert(args("pool")).await.unwrap_err(); + assert_extension_error(&error); + assert!(fixture.job_ids().await.is_empty()); + + let mut tx = fixture.begin().await; + let error = client.insert(args("single")).tx(&mut tx).await.unwrap_err(); + assert_extension_error(&error); + assert_eq!(fixture.job_ids_in(&mut tx).await.len(), 1); + tx.rollback().await.unwrap(); + + assert!(fixture.job_ids().await.is_empty()); + fixture.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn failed_cancel_and_retry_steps_stay_in_caller_transaction() { + let fixture = Fixture::new().await; + let fail = Arc::new(AtomicBool::new(false)); + let client = fixture.client(&fail); + let available = client.insert(args("available")).await.unwrap().id(); + let scheduled = client + .insert(args("scheduled")) + .opts( + InsertOpts::default() + .with_scheduled_at(chrono::Utc::now() + chrono::Duration::hours(1)), + ) + .await + .unwrap() + .id(); + fixture.clear_effects().await; + fail.store(true, Ordering::SeqCst); + + let mut tx = fixture.begin().await; + let error = client + .jobs() + .cancel(available) + .tx(&mut tx) + .await + .unwrap_err(); + assert_extension_error(&error); + let error = client + .jobs() + .retry(scheduled) + .tx(&mut tx) + .await + .unwrap_err(); + assert_extension_error(&error); + assert_eq!( + fixture.state_in(&mut tx, available).await, + JobState::Cancelled + ); + assert_eq!( + fixture.state_in(&mut tx, scheduled).await, + JobState::Available + ); + assert_eq!( + fixture.effect_ids_in(&mut tx).await, + vec![available.min(scheduled), available.max(scheduled)] + ); + tx.rollback().await.unwrap(); + + assert_eq!(fixture.state(available).await, JobState::Available); + assert_eq!(fixture.state(scheduled).await, JobState::Scheduled); + assert!(fixture.effect_ids().await.is_empty()); + fixture.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn failed_completion_steps_stay_in_caller_transaction() { + let fixture = Fixture::new().await; + let fail = Arc::new(AtomicBool::new(false)); + let client = fixture.client(&fail); + let id = client.insert(args("running")).await.unwrap().id(); + let (_, context) = riverqueue::__private::claim_job_for_test(&client, id) + .await + .unwrap(); + fixture.clear_effects().await; + fail.store(true, Ordering::SeqCst); + + let mut tx = fixture.begin().await; + let error = client.jobs().complete(id).tx(&mut tx).await.unwrap_err(); + assert_extension_error(&error); + assert_eq!(fixture.state_in(&mut tx, id).await, JobState::Completed); + assert_eq!(fixture.effect_ids_in(&mut tx).await, vec![id]); + tx.rollback().await.unwrap(); + + let mut tx = fixture.begin().await; + let error = context.job_complete_tx(&mut tx).await.unwrap_err(); + assert_extension_error(&error); + assert_eq!(fixture.state_in(&mut tx, id).await, JobState::Completed); + tx.rollback().await.unwrap(); + + assert_eq!(fixture.state(id).await, JobState::Running); + assert!(fixture.effect_ids().await.is_empty()); + + fail.store(false, Ordering::SeqCst); + let mut tx = fixture.begin().await; + client.jobs().complete(id).tx(&mut tx).await.unwrap(); + tx.commit().await.unwrap(); + assert_eq!(fixture.state(id).await, JobState::Completed); + assert_eq!(fixture.effect_ids().await, vec![id]); + fixture.cleanup().await; + } + }; +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +mod postgres { + use riverqueue::database::PostgresDatabase; + use sqlx::{Postgres, Transaction}; + + use super::*; + use crate::support::PostgresSchema; + + struct Fixture { + schema: PostgresSchema, + } + + impl Fixture { + async fn new() -> Self { + let schema = PostgresSchema::new("river_caller_tx").await; + sqlx::raw_sql(sqlx::AssertSqlSafe(format!( + "CREATE TABLE {} (job_id bigint NOT NULL)", + schema.table("scope_effect") + ))) + .execute(&schema.pool) + .await + .unwrap(); + Self { schema } + } + + fn builder(&self) -> riverqueue::ClientBuilder { + Client::builder( + PostgresDatabase::new(self.schema.pool.clone()) + .with_schema(self.schema.schema.clone()), + ) + } + + fn client(&self, fail: &Arc) -> Client { + self.pilot_client(fail, false) + } + + fn pilot_client(&self, fail: &Arc, fail_in_database: bool) -> Client { + self.builder() + .pilot(EffectPilot { + effect_table: self.schema.table("scope_effect"), + fail: Arc::clone(fail), + fail_in_database, + }) + .build() + .unwrap() + } + + async fn begin(&self) -> Transaction<'static, Postgres> { + self.schema.pool.begin().await.unwrap() + } + + async fn ids(&self, table: &str, column: &str) -> Vec { + sqlx::query_scalar(sqlx::AssertSqlSafe(format!( + "SELECT {column} FROM {} ORDER BY {column}", + self.schema.table(table) + ))) + .fetch_all(&self.schema.pool) + .await + .unwrap() + } + + async fn job_ids(&self) -> Vec { + self.ids("river_job", "id").await + } + + async fn ids_in( + &self, + tx: &mut Transaction<'static, Postgres>, + table: &str, + column: &str, + ) -> Vec { + sqlx::query_scalar(sqlx::AssertSqlSafe(format!( + "SELECT {column} FROM {} ORDER BY {column}", + self.schema.table(table) + ))) + .fetch_all(&mut **tx) + .await + .unwrap() + } + + async fn job_ids_in(&self, tx: &mut Transaction<'static, Postgres>) -> Vec { + self.ids_in(tx, "river_job", "id").await + } + + async fn effect_ids_in(&self, tx: &mut Transaction<'static, Postgres>) -> Vec { + self.ids_in(tx, "scope_effect", "job_id").await + } + + async fn state_in(&self, tx: &mut Transaction<'static, Postgres>, id: i64) -> JobState { + let state: String = sqlx::query_scalar(sqlx::AssertSqlSafe(format!( + "SELECT state::text FROM {} WHERE id = $1", + self.schema.table("river_job") + ))) + .bind(id) + .fetch_one(&mut **tx) + .await + .unwrap(); + state.parse().unwrap() + } + + async fn effect_ids(&self) -> Vec { + self.ids("scope_effect", "job_id").await + } + + async fn clear_effects(&self) { + sqlx::raw_sql(sqlx::AssertSqlSafe(format!( + "DELETE FROM {}", + self.schema.table("scope_effect") + ))) + .execute(&self.schema.pool) + .await + .unwrap(); + } + + async fn state(&self, id: i64) -> JobState { + let state: String = sqlx::query_scalar(sqlx::AssertSqlSafe(format!( + "SELECT state::text FROM {} WHERE id = $1", + self.schema.table("river_job") + ))) + .bind(id) + .fetch_one(&self.schema.pool) + .await + .unwrap(); + state.parse().unwrap() + } + + async fn cleanup(self) { + self.schema.cleanup().await; + } + } + + scenarios!(); + + // A statement the database rejects aborts a PostgreSQL transaction. + // River doesn't hide that behind a savepoint, so the caller's + // transaction can only be rolled back, along with its earlier work. + #[tokio::test(flavor = "multi_thread")] + async fn database_error_aborts_caller_transaction() { + let fixture = Fixture::new().await; + let client = fixture.pilot_client(&Arc::new(AtomicBool::new(false)), true); + let plain = fixture.builder().build().unwrap(); + + let mut tx = fixture.begin().await; + plain.insert(args("prior")).tx(&mut tx).await.unwrap(); + assert!(client.insert(args("failed")).tx(&mut tx).await.is_err()); + assert!(sqlx::query("SELECT 1").execute(&mut *tx).await.is_err()); + tx.rollback().await.unwrap(); + + assert_eq!(fixture.job_ids().await, Vec::::new()); + fixture.cleanup().await; + } + + // Every write River makes in a caller's transaction, including an + // intercepting extension's, carries the caller's transaction ID. A + // savepoint would give its writes their own subtransaction ID. + #[tokio::test(flavor = "multi_thread")] + async fn writes_use_the_callers_transaction_id() { + let fixture = Fixture::new().await; + let fail = Arc::new(AtomicBool::new(false)); + let client = fixture.client(&fail); + let plain = fixture.builder().build().unwrap(); + let cancellable = plain.insert(args("cancellable")).await.unwrap().id(); + let retryable = plain + .insert(args("retryable")) + .opts( + InsertOpts::default() + .with_scheduled_at(chrono::Utc::now() + chrono::Duration::hours(1)), + ) + .await + .unwrap() + .id(); + + let mut tx = fixture.begin().await; + // A write directly in the caller's transaction, so that even one + // savepoint around all of River's writes would be detected. + plain.insert(args("direct")).tx(&mut tx).await.unwrap(); + // More than PostgreSQL's cached subtransaction ID limit. + for index in 0..70 { + client + .insert(args(&format!("single_{index}"))) + .tx(&mut tx) + .await + .unwrap(); + } + client + .insert_many([args("many_1"), args("many_2")]) + .tx(&mut tx) + .await + .unwrap(); + client.jobs().cancel(cancellable).tx(&mut tx).await.unwrap(); + client.jobs().retry(retryable).tx(&mut tx).await.unwrap(); + + let (rows, transactions): (i64, i64) = sqlx::query_as(sqlx::AssertSqlSafe(format!( + "SELECT count(*), count(DISTINCT xmin::text) FROM (\ + SELECT xmin FROM {} UNION ALL SELECT xmin FROM {}) AS written", + fixture.schema.table("river_job"), + fixture.schema.table("scope_effect") + ))) + .fetch_one(&mut *tx) + .await + .unwrap(); + // 75 jobs, plus an effect row for each of 71 intercepted insertions + // and the cancellation and retry. + assert_eq!(rows, 75 + 73); + assert_eq!( + transactions, 1, + "all writes use the caller's transaction ID" + ); + tx.rollback().await.unwrap(); + fixture.cleanup().await; + } +} + +#[cfg(feature = "sqlite")] +mod sqlite { + use sqlx::{Sqlite, SqlitePool, Transaction}; + + use super::*; + use crate::support::{sqlite_cleanup, sqlite_file_pool}; + + struct Fixture { + path: std::path::PathBuf, + pool: SqlitePool, + } + + impl Fixture { + async fn new() -> Self { + let (pool, path) = sqlite_file_pool(4).await; + sqlx::raw_sql("CREATE TABLE scope_effect (job_id INTEGER NOT NULL)") + .execute(&pool) + .await + .unwrap(); + Self { path, pool } + } + + fn builder(&self) -> riverqueue::ClientBuilder { + Client::builder(self.pool.clone()) + } + + fn client(&self, fail: &Arc) -> Client { + self.pilot_client(fail, false) + } + + fn pilot_client(&self, fail: &Arc, fail_in_database: bool) -> Client { + self.builder() + .pilot(EffectPilot { + effect_table: "scope_effect".to_owned(), + fail: Arc::clone(fail), + fail_in_database, + }) + .build() + .unwrap() + } + + async fn begin(&self) -> Transaction<'static, Sqlite> { + self.pool.begin_with("BEGIN IMMEDIATE").await.unwrap() + } + + async fn ids(&self, sql: &'static str) -> Vec { + sqlx::query_scalar(sql).fetch_all(&self.pool).await.unwrap() + } + + async fn job_ids(&self) -> Vec { + self.ids("SELECT id FROM river_job ORDER BY id").await + } + + async fn effect_ids(&self) -> Vec { + self.ids("SELECT job_id FROM scope_effect ORDER BY job_id") + .await + } + + async fn job_ids_in(&self, tx: &mut Transaction<'static, Sqlite>) -> Vec { + sqlx::query_scalar("SELECT id FROM river_job ORDER BY id") + .fetch_all(&mut **tx) + .await + .unwrap() + } + + async fn effect_ids_in(&self, tx: &mut Transaction<'static, Sqlite>) -> Vec { + sqlx::query_scalar("SELECT job_id FROM scope_effect ORDER BY job_id") + .fetch_all(&mut **tx) + .await + .unwrap() + } + + async fn state_in(&self, tx: &mut Transaction<'static, Sqlite>, id: i64) -> JobState { + let state: String = sqlx::query_scalar("SELECT state FROM river_job WHERE id = ?") + .bind(id) + .fetch_one(&mut **tx) + .await + .unwrap(); + state.parse().unwrap() + } + + async fn clear_effects(&self) { + sqlx::raw_sql("DELETE FROM scope_effect") + .execute(&self.pool) + .await + .unwrap(); + } + + async fn state(&self, id: i64) -> JobState { + let state: String = sqlx::query_scalar("SELECT state FROM river_job WHERE id = ?") + .bind(id) + .fetch_one(&self.pool) + .await + .unwrap(); + state.parse().unwrap() + } + + async fn cleanup(self) { + sqlite_cleanup(self.pool, self.path).await; + } + } + + scenarios!(); + // A statement SQLite rejects doesn't abort the transaction, so the + // caller's transaction keeps River's write before the failure and its + // own earlier work; the caller still rolls back on the error. + #[tokio::test(flavor = "multi_thread")] + async fn database_error_keeps_caller_transaction_writes() { + let fixture = Fixture::new().await; + let client = fixture.pilot_client(&Arc::new(AtomicBool::new(false)), true); + let plain = fixture.builder().build().unwrap(); + + let mut tx = fixture.begin().await; + let prior = plain.insert(args("prior")).tx(&mut tx).await.unwrap().id(); + assert!(client.insert(args("failed")).tx(&mut tx).await.is_err()); + let in_tx = fixture.job_ids_in(&mut tx).await; + assert_eq!(in_tx.len(), 2, "{in_tx:?}"); + assert!(in_tx.contains(&prior)); + tx.rollback().await.unwrap(); + + assert_eq!(fixture.job_ids().await, Vec::::new()); + fixture.cleanup().await; + } +} diff --git a/rust/riverqueue/tests/client_handles.rs b/rust/riverqueue/tests/client_handles.rs new file mode 100644 index 000000000..9e0623632 --- /dev/null +++ b/rust/riverqueue/tests/client_handles.rs @@ -0,0 +1,1079 @@ +//! Behavior of the client's scoped operation handles on every backend. +//! +//! Each scenario runs against PostgreSQL (in a unique schema, failing rather +//! than skipping when `RIVER_RUST_DATABASE_URL` is unset) and SQLite (in a +//! temporary file). PostgreSQL scenarios build only with `--cfg river_postgres_tests`. + +#![cfg(any(all(feature = "postgres", river_postgres_tests), feature = "sqlite"))] + +mod support; + +use std::{convert::Infallible, time::Duration}; + +use riverqueue::{ + Client, Error, EventKind, InsertBatch, InsertContext, InsertMiddleware, InsertNext, InsertOpts, + InsertedJob, Job, JobArgs, JobDeleteManyParams, JobListParams, JobState, JobUpdateParams, + QueueConfig, QueueListParams, QueueSelector, QueueUpdateParams, UniqueOpts, WorkContext, + WorkOutcome, WorkerRegistry, +}; +use serde::{Deserialize, Serialize}; + +/// Blocks until its attempt is cancelled. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "client_handles_blocking")] +struct BlockingArgs {} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "client_handles")] +struct HandleArgs { + name: String, +} + +/// Encodes like `HandleArgs` under another kind. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "client_handles_other_kind")] +struct OtherKindArgs { + name: String, +} + +fn args(name: &str) -> HandleArgs { + HandleArgs { + name: name.to_owned(), + } +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "client_handles_float")] +struct FloatArgs { + value: f64, +} + +/// Fails every insertion after River has written its jobs. +struct FailAfterWrite; + +impl InsertMiddleware for FailAfterWrite { + async fn insert_many( + &self, + jobs: Vec, + next: InsertNext<'_>, + ) -> Result, Error> { + let inserted = next.run(jobs).await?; + assert!(!inserted.is_empty()); + Err(Error::extension( + riverqueue::ExtensionPhase::InsertMiddleware, + std::io::Error::other("failed after the write"), + )) + } +} + +/// Defines each scenario for one backend's `Fixture`. +macro_rules! scenarios { + () => { + // Like Go, a claim appends the client to at most the 100 most + // recent `attempted_by` entries. + #[tokio::test(flavor = "multi_thread")] + async fn attempted_by_keeps_the_most_recent_hundred_clients() { + let fixture = Fixture::new().await; + let id = fixture + .client + .insert(args("attempted_by")) + .await + .unwrap() + .id(); + let previous = (1..=100) + .map(|index| format!("client-{index}")) + .collect::>(); + fixture.set_attempted_by(id, &previous).await; + let client = fixture + .builder() + .id("attempted-by-worker") + .workers(workers()) + .queue("default", fast_queue()) + .build() + .unwrap(); + let mut run = client.start().unwrap(); + wait_for_completion(&client, id).await; + run.shutdown().await.unwrap(); + + let attempted_by = client.jobs().get(id).await.unwrap().attempted_by; + let mut expected = previous[1..].to_vec(); + expected.push("attempted-by-worker".to_owned()); + assert_eq!(attempted_by, expected); + fixture.cleanup().await; + } + + // A worker that returns successfully after its job is cancelled + // completes the job, as in Go. + #[tokio::test(flavor = "multi_thread")] + async fn cancelled_job_that_succeeds_is_completed() { + let fixture = Fixture::new().await; + let (started_sender, mut started) = tokio::sync::mpsc::unbounded_channel(); + let mut workers = WorkerRegistry::new(); + workers + .register_fn(move |context: WorkContext, job: Job| { + let started_sender = started_sender.clone(); + async move { + let _ = started_sender.send(job.id()); + context.cancellation_token().cancelled().await; + Ok::<_, Infallible>(WorkOutcome::Complete) + } + }) + .unwrap(); + let client = fixture + .builder() + .without_notifications() + .workers(workers) + .queue("default", fast_queue()) + .build() + .unwrap(); + let mut run = client.start().unwrap(); + let id = client.insert(BlockingArgs {}).await.unwrap().id(); + tokio::time::timeout(Duration::from_secs(10), started.recv()) + .await + .expect("job starts") + .unwrap(); + client.jobs().cancel(id).await.unwrap(); + wait_for_completion(&client, id).await; + run.shutdown().await.unwrap(); + fixture.cleanup().await; + } + + // Like Go's `Insert` and `InsertMany`, an insertion without a + // caller transaction runs middleware, hooks, and the write in one + // transaction, so middleware failing after the write rolls it back. + #[tokio::test(flavor = "multi_thread")] + async fn insert_middleware_error_after_write_rolls_back() { + let fixture = Fixture::new().await; + let client = fixture + .builder() + .insert_middleware(FailAfterWrite) + .build() + .unwrap(); + + let error = client.insert(args("single")).await.unwrap_err(); + assert!(matches!(error, Error::Extension { .. }), "{error}"); + let error = client + .insert_many([args("many_1"), args("many_2")]) + .await + .unwrap_err(); + assert!(matches!(error, Error::Extension { .. }), "{error}"); + let mut batch = InsertBatch::new(); + batch.push(args("batch_1")).push(BlockingArgs {}); + let error = client.insert_batch(batch).await.unwrap_err(); + assert!(matches!(error, Error::Extension { .. }), "{error}"); + + assert_eq!(fixture.job_count().await, 0); + fixture.cleanup().await; + } + + // Go's `encoding/json` can't encode NaN or infinities, so River Go + // refuses such arguments. River Rust refuses them too rather than + // storing `null`, which a Go worker would decode as a different + // value. + #[tokio::test(flavor = "multi_thread")] + async fn insert_rejects_non_finite_float_args() { + let fixture = Fixture::new().await; + let client = fixture.builder().build().unwrap(); + + for value in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] { + let error = client.insert(FloatArgs { value }).await.unwrap_err(); + assert!(matches!(error, Error::Json(_)), "{error}"); + let error = client + .insert_many([FloatArgs { value: 1.0 }, FloatArgs { value }]) + .await + .unwrap_err(); + assert!(matches!(error, Error::Json(_)), "{error}"); + let mut batch = InsertBatch::new(); + batch.push(args("batch")).push(FloatArgs { value }); + let error = client.insert_batch(batch).await.unwrap_err(); + assert!(matches!(error, Error::Json(_)), "{error}"); + } + + assert_eq!(fixture.job_count().await, 0); + fixture.cleanup().await; + } + + // Like Go, a custom unique state set must include the states a job + // passes through while it's being worked; an empty set means the + // default states. + #[tokio::test(flavor = "multi_thread")] + async fn insert_requires_unique_states_to_include_required_states() { + let fixture = Fixture::new().await; + let client = fixture.builder().build().unwrap(); + + let missing = InsertOpts::default().with_unique( + UniqueOpts::new().with_by_state([JobState::Available, JobState::Completed]), + ); + let error = client + .insert(args("missing_states")) + .opts(missing) + .await + .unwrap_err(); + assert!(matches!(error, Error::InvalidJob(_)), "{error}"); + assert!( + error.to_string().contains("pending, running, scheduled"), + "{error}" + ); + assert_eq!(fixture.job_count().await, 0); + + let empty = InsertOpts::default().with_unique(UniqueOpts::new().with_by_state([])); + let first = client + .insert(args("empty_states")) + .opts(empty.clone()) + .await + .unwrap(); + let duplicate = client + .insert(args("empty_states")) + .opts(empty) + .await + .unwrap(); + assert!(duplicate.unique_skipped_as_duplicate); + assert_eq!(duplicate.id(), first.id()); + fixture.cleanup().await; + } + + // Like Go, with `exclude_kind` jobs of different kinds share a unique + // key, and an insertion skipped as a duplicate of a job of another + // kind returns that job as it is rather than rewriting its kind. + #[tokio::test(flavor = "multi_thread")] + async fn insert_unique_skip_keeps_the_existing_jobs_kind() { + let fixture = Fixture::new().await; + let unique = InsertOpts::default().with_unique( + UniqueOpts::new() + .with_by_args(true) + .with_exclude_kind(true), + ); + let first = fixture + .client + .insert(args("exclude_kind")) + .opts(unique.clone()) + .await + .unwrap(); + assert!(!first.unique_skipped_as_duplicate); + + let other = OtherKindArgs { + name: "exclude_kind".to_owned(), + }; + let duplicate = fixture + .client + .insert(other.clone()) + .opts(unique.clone()) + .await + .unwrap(); + assert!(duplicate.unique_skipped_as_duplicate); + assert_eq!(duplicate.id(), first.id()); + assert_eq!(duplicate.job.row.kind, HandleArgs::KIND); + + let batch = fixture + .client + .insert_many([(other, unique)]) + .await + .unwrap(); + assert!(batch[0].unique_skipped_as_duplicate); + assert_eq!(batch[0].id(), first.id()); + assert_eq!(batch[0].job.row.kind, HandleArgs::KIND); + + let stored = fixture.client.jobs().get(first.id()).await.unwrap(); + assert_eq!(stored.kind, HandleArgs::KIND); + assert_eq!(fixture.job_count().await, 1); + fixture.cleanup().await; + } + + // Each delete-many filter deletes exactly the jobs it matches. + #[tokio::test(flavor = "multi_thread")] + async fn delete_many_filters_by_kind_queue_priority_and_state() { + let fixture = Fixture::new().await; + let client = fixture.builder().build().unwrap(); + let jobs = client.jobs(); + + let plain = client.insert(args("plain")).await.unwrap().id(); + let other_kind = client.insert(FloatArgs { value: 1.0 }).await.unwrap().id(); + let other_queue = client + .insert(args("other_queue")) + .opts(InsertOpts::default().with_queue("delete_many_other")) + .await + .unwrap() + .id(); + let urgent = client + .insert(args("urgent")) + .opts(InsertOpts::default().with_priority(2)) + .await + .unwrap() + .id(); + let cancelled = client.insert(args("cancelled")).await.unwrap().id(); + jobs.cancel(cancelled).await.unwrap(); + + let deleted_ids = |rows: Vec| { + let mut ids = rows.into_iter().map(|row| row.id).collect::>(); + ids.sort_unstable(); + ids + }; + let delete = + |params: JobListParams| jobs.delete_many(JobDeleteManyParams::matching(params)); + assert_eq!( + deleted_ids( + delete(JobListParams::default().kinds([FloatArgs::KIND])) + .await + .unwrap() + ), + [other_kind] + ); + assert_eq!( + deleted_ids( + delete(JobListParams::default().queues(["delete_many_other"])) + .await + .unwrap() + ), + [other_queue] + ); + assert_eq!( + deleted_ids( + delete(JobListParams::default().priorities([2])) + .await + .unwrap() + ), + [urgent] + ); + assert_eq!( + deleted_ids( + delete(JobListParams::default().states([JobState::Cancelled])) + .await + .unwrap() + ), + [cancelled] + ); + // Combined filters must all match. + assert!( + delete( + JobListParams::default() + .kinds([HandleArgs::KIND]) + .states([JobState::Cancelled]) + ) + .await + .unwrap() + .is_empty() + ); + let remaining = jobs.list(JobListParams::default()).await.unwrap().jobs; + assert_eq!( + remaining.iter().map(|row| row.id).collect::>(), + [plain] + ); + fixture.cleanup().await; + } + + // Jobs are fetched by priority, then scheduled time, then ID, as in + // Go. + #[tokio::test(flavor = "multi_thread")] + async fn fetches_by_priority_then_schedule_then_id() { + let fixture = Fixture::new().await; + let base = chrono::Utc::now() - chrono::Duration::minutes(10); + let mut ids = Vec::new(); + for (name, priority, minutes) in [ + ("p2_early", 2, 0), + ("p1_late", 1, 5), + ("p1_early_first", 1, 1), + ("p1_early_second", 1, 1), + ("p4_earliest", 4, -5), + ] { + let id = fixture + .client + .insert(args(name)) + .opts( + InsertOpts::default() + .with_priority(priority) + .with_scheduled_at(base + chrono::Duration::minutes(minutes)), + ) + .await + .unwrap() + .id(); + ids.push(id); + } + let worked = std::sync::Arc::new(std::sync::Mutex::new(Vec::new())); + let mut workers = WorkerRegistry::new(); + let recorder = std::sync::Arc::clone(&worked); + workers + .register_fn(move |_context: WorkContext, job: Job| { + let recorder = std::sync::Arc::clone(&recorder); + async move { + recorder.lock().unwrap().push(job.id()); + Ok::<_, Infallible>(WorkOutcome::Complete) + } + }) + .unwrap(); + let client = fixture + .builder() + .workers(workers) + .queue("default", fast_queue()) + .build() + .unwrap(); + let mut run = client.start().unwrap(); + wait_for_completion(&client, ids[4]).await; + run.shutdown().await.unwrap(); + + assert_eq!( + *worked.lock().unwrap(), + [ids[2], ids[3], ids[1], ids[0], ids[4]] + ); + fixture.cleanup().await; + } + + // Port of Go's `CancelRunningJobPollOnly`: with no listener, the + // cancelling client must wake its own running attempt. + #[tokio::test(flavor = "multi_thread")] + async fn cancel_reaches_a_running_job_on_a_poll_only_client() { + let fixture = Fixture::new().await; + let (started_sender, mut started) = tokio::sync::mpsc::unbounded_channel(); + let mut workers = WorkerRegistry::new(); + workers + .register_fn(move |context: WorkContext, job: Job| { + let started_sender = started_sender.clone(); + async move { + let _ = started_sender.send(job.id()); + context.cancellation_token().cancelled().await; + Err::(std::io::Error::other("cancelled")) + } + }) + .unwrap(); + let client = fixture + .builder() + .without_notifications() + .workers(workers) + .queue("default", fast_queue()) + .build() + .unwrap(); + let mut events = client.subscribe(&[EventKind::JobCancelled]).unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + let id = client.insert(BlockingArgs {}).await.unwrap().id(); + let started_id = tokio::time::timeout(Duration::from_secs(10), started.recv()) + .await + .expect("job starts") + .unwrap(); + assert_eq!(started_id, id); + + let row = client.jobs().cancel(id).await.unwrap(); + assert_eq!(row.state, JobState::Running); + + let event = tokio::time::timeout(Duration::from_secs(10), events.recv()) + .await + .expect("cancellation reaches the running attempt") + .unwrap(); + let job = &event.as_job().unwrap().job; + assert_eq!(job.id, id); + assert_eq!(job.state, JobState::Cancelled); + let finalized_at = job.finalized_at.unwrap(); + assert!((chrono::Utc::now() - finalized_at).num_seconds().abs() < 2); + + run.shutdown().await.unwrap(); + fixture.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn job_requests_take_effect_only_when_the_transaction_commits() { + let fixture = Fixture::new().await; + let client = &fixture.client; + let jobs = client.jobs(); + let cancelled = client.insert(args("cancel")).await.unwrap().id(); + let deleted = client.insert(args("delete")).await.unwrap().id(); + let updated = client.insert(args("update")).await.unwrap().id(); + + // Every write is rolled back with the caller's transaction. + let mut tx = fixture.begin().await; + let row = jobs.cancel(cancelled).tx(&mut tx).await.unwrap(); + assert_eq!(row.state, JobState::Cancelled); + let row = jobs.delete(deleted).tx(&mut tx).await.unwrap(); + assert_eq!(row.id, deleted); + let row = jobs + .update( + updated, + JobUpdateParams::default().output(serde_json::json!("rolled back")), + ) + .tx(&mut tx) + .await + .unwrap(); + assert_eq!( + row.metadata.get::("output").unwrap().as_deref(), + Some("rolled back") + ); + // Reads in the transaction see its uncommitted writes. + assert!(matches!( + jobs.get(deleted).tx(&mut tx).await, + Err(Error::NotFound(_)) + )); + let inserted = client + .insert(args("uncommitted")) + .tx(&mut tx) + .await + .unwrap() + .id(); + let listed = jobs + .list(JobListParams::default().ids([inserted])) + .tx(&mut tx) + .await + .unwrap(); + assert_eq!(listed.jobs.len(), 1); + tx.rollback().await.unwrap(); + + assert_eq!( + jobs.get(cancelled).await.unwrap().state, + JobState::Available + ); + assert!(jobs.get(deleted).await.is_ok()); + assert!( + !jobs + .get(updated) + .await + .unwrap() + .metadata + .contains_key("output") + ); + let listed = jobs + .list(JobListParams::default().ids([inserted])) + .await + .unwrap(); + assert!(listed.jobs.is_empty()); + assert!(listed.last_cursor.is_none()); + + // The same requests persist once the transaction commits. + let mut tx = fixture.begin().await; + jobs.cancel(cancelled).tx(&mut tx).await.unwrap(); + jobs.delete_many(JobDeleteManyParams::matching( + JobListParams::default().ids([deleted]), + )) + .tx(&mut tx) + .await + .unwrap(); + let retried = jobs.retry(cancelled).tx(&mut tx).await.unwrap(); + assert_eq!(retried.state, JobState::Available); + tx.commit().await.unwrap(); + + assert_eq!( + jobs.get(cancelled).await.unwrap().state, + JobState::Available + ); + assert!(matches!(jobs.get(deleted).await, Err(Error::NotFound(_)))); + + fixture.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn job_requests_run_on_the_pool_without_a_transaction() { + let fixture = Fixture::new().await; + let jobs = fixture.client.jobs(); + let first = fixture.client.insert(args("first")).await.unwrap().id(); + let second = fixture.client.insert(args("second")).await.unwrap().id(); + + let page = jobs + .list(JobListParams::default().ids([first, second]).limit(1)) + .await + .unwrap(); + assert_eq!( + page.jobs.iter().map(|job| job.id).collect::>(), + [first] + ); + let cursor = page.last_cursor.expect("a nonempty page has a cursor"); + let page = jobs + .list( + JobListParams::default() + .ids([first, second]) + .limit(1) + .after(cursor), + ) + .await + .unwrap(); + assert_eq!( + page.jobs.iter().map(|job| job.id).collect::>(), + [second] + ); + + assert_eq!(jobs.cancel(first).await.unwrap().state, JobState::Cancelled); + assert_eq!(jobs.retry(first).await.unwrap().state, JobState::Available); + assert_eq!(jobs.delete(second).await.unwrap().id, second); + assert!(matches!(jobs.delete(second).await, Err(Error::NotFound(_)))); + assert!(matches!( + jobs.cancel(i64::MAX).await, + Err(Error::NotFound(_)) + )); + assert!(matches!( + jobs.retry(i64::MAX).await, + Err(Error::NotFound(_)) + )); + + fixture.cleanup().await; + } + + // SQLite must match PostgreSQL's `@>` containment exactly. + #[tokio::test(flavor = "multi_thread")] + async fn metadata_filters_match_postgres_containment() { + let fixture = Fixture::new().await; + let client = &fixture.client; + let mut ids = std::collections::HashMap::new(); + for (name, metadata) in [ + ("null", r#"{"a":null}"#), + ("missing", "{}"), + ("integer", r#"{"a":1}"#), + ("float", r#"{"a":1.0}"#), + ("string", r#"{"a":"1"}"#), + ("array", r#"{"a":[1,2,{"b":"x"}],"s":"a = + serde_json::from_str(fragment).unwrap(); + let mut want = want.iter().map(|name| ids[name]).collect::>(); + want.sort_unstable(); + let listed = client + .jobs() + .list(JobListParams::default().metadata(fragment.clone())) + .await + .unwrap(); + let got = listed.jobs.iter().map(|job| job.id).collect::>(); + assert_eq!(got, want, "{fragment:?}"); + } + + fixture.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn queue_requests_take_effect_only_when_the_transaction_commits() { + let fixture = Fixture::new().await; + let queues = fixture.client.queues(); + fixture.insert_queue("alpha").await; + fixture.insert_queue("beta").await; + let owner = serde_json::Map::from_iter([("owner".to_owned(), "rust".into())]); + + let mut tx = fixture.begin().await; + queues.pause(QueueSelector::All).tx(&mut tx).await.unwrap(); + assert!( + queues + .get("alpha") + .tx(&mut tx) + .await + .unwrap() + .paused_at + .is_some() + ); + let beta = queues + .update("beta", QueueUpdateParams::new().metadata(owner.clone())) + .tx(&mut tx) + .await + .unwrap(); + assert_eq!(beta.metadata, owner); + tx.rollback().await.unwrap(); + assert!(paused(&fixture.client).await.is_empty()); + assert!(queues.get("beta").await.unwrap().metadata.is_empty()); + + let mut tx = fixture.begin().await; + queues.pause(QueueSelector::All).tx(&mut tx).await.unwrap(); + queues + .update("beta", QueueUpdateParams::new().metadata(owner.clone())) + .tx(&mut tx) + .await + .unwrap(); + tx.commit().await.unwrap(); + assert_eq!(paused(&fixture.client).await, ["alpha", "beta"]); + assert_eq!(queues.get("beta").await.unwrap().metadata, owner); + + queues.resume("alpha").await.unwrap(); + assert_eq!(paused(&fixture.client).await, ["beta"]); + queues.resume(QueueSelector::All).await.unwrap(); + assert!(paused(&fixture.client).await.is_empty()); + + fixture.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn queue_selectors_match_names_literally() { + let fixture = Fixture::new().await; + let queues = fixture.client.queues(); + + // Selecting every queue succeeds when there are none, like Go. + queues.pause(QueueSelector::All).await.unwrap(); + queues.resume(QueueSelector::All).await.unwrap(); + + fixture.insert_queue("alpha").await; + assert_eq!( + QueueSelector::from("alpha"), + QueueSelector::Named("alpha".to_owned()) + ); + assert_eq!( + QueueSelector::from("*".to_owned()), + QueueSelector::Named("*".to_owned()) + ); + // `*` is only a name, and no queue can have it. + assert!(matches!(queues.pause("*").await, Err(Error::NotFound(_)))); + assert!(matches!(queues.resume("*").await, Err(Error::NotFound(_)))); + assert!(paused(&fixture.client).await.is_empty()); + assert!(matches!( + queues.pause("missing").await, + Err(Error::NotFound(_)) + )); + assert!(matches!( + queues.get("missing").await, + Err(Error::NotFound(_)) + )); + assert!(matches!( + queues.update("missing", QueueUpdateParams::new()).await, + Err(Error::NotFound(_)) + )); + + // Updating without metadata keeps it while refreshing the record. + let owner = serde_json::Map::from_iter([("owner".to_owned(), "rust".into())]); + let before = queues + .update("alpha", QueueUpdateParams::new().metadata(owner.clone())) + .await + .unwrap(); + let after = queues + .update("alpha", QueueUpdateParams::new()) + .await + .unwrap(); + assert_eq!(after.metadata, owner); + assert!(after.updated_at >= before.updated_at); + + fixture.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn local_queues_change_configuration_without_lock_errors() { + let fixture = Fixture::new().await; + let client = fixture + .builder() + .workers(workers()) + .queue("default", QueueConfig::new(1)) + .build() + .unwrap(); + let local = client.local_queues(); + assert_eq!( + local.configs(), + [("default".to_owned(), QueueConfig::new(1))].into() + ); + + local.add("second", QueueConfig::new(2)).unwrap(); + // Like Go, adding a queue twice is an error; update reconfigures. + assert!(matches!( + local.add("default", QueueConfig::new(3)), + Err(Error::QueueAlreadyAdded { name }) if name == "default" + )); + local.update("default", QueueConfig::new(3)).unwrap(); + assert!(matches!( + local.update("missing", QueueConfig::new(3)), + Err(Error::QueueNotAdded { name }) if name == "missing" + )); + assert_eq!( + local.configs(), + [ + ("default".to_owned(), QueueConfig::new(3)), + ("second".to_owned(), QueueConfig::new(2)), + ] + .into() + ); + assert!(matches!( + local.add("not a queue name", QueueConfig::new(1)), + Err(Error::InvalidJob(_)) + )); + assert!(matches!( + local.add("third", QueueConfig::new(0)), + Err(Error::Configuration(_)) + )); + // A client that isn't running has no producer to wait for. + assert_eq!(local.remove("second").await.unwrap(), QueueConfig::new(2)); + assert!(matches!( + local.remove("second").await, + Err(Error::QueueNotAdded { .. }) + )); + assert_eq!(local.configs().len(), 1); + + // A client without workers can't run any queue. + assert!(matches!( + fixture + .client + .local_queues() + .add("default", QueueConfig::new(1)), + Err(Error::Configuration(_)) + )); + + fixture.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn local_queues_start_producers_while_running() { + let fixture = Fixture::new().await; + let client = fixture + .builder() + .workers(workers()) + .queue("default", fast_queue()) + .build() + .unwrap(); + let mut run = client.start().unwrap(); + + client.local_queues().add("dynamic", fast_queue()).unwrap(); + let job = client + .insert(args("dynamic")) + .opts(InsertOpts::default().with_queue("dynamic")) + .await + .unwrap(); + wait_for_completion(&client, job.id()).await; + assert_eq!( + client.local_queues().remove("dynamic").await.unwrap(), + fast_queue() + ); + assert!(!client.local_queues().configs().contains_key("dynamic")); + + run.shutdown().await.unwrap(); + fixture.cleanup().await; + } + }; +} + +fn workers() -> WorkerRegistry { + let mut workers = WorkerRegistry::new(); + workers + .register_fn(|_context: WorkContext, _job: Job| async { + Ok::<_, Infallible>(WorkOutcome::Complete) + }) + .unwrap(); + workers +} + +fn fast_queue() -> QueueConfig { + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)) +} + +/// Waits for a job to complete, failing after ten seconds. +async fn wait_for_completion(client: &Client, id: i64) { + let deadline = tokio::time::Instant::now() + Duration::from_secs(10); + while client.jobs().get(id).await.unwrap().state != JobState::Completed { + assert!( + tokio::time::Instant::now() < deadline, + "job {id} did not complete" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } +} + +/// Returns the names of paused queues. +async fn paused(client: &Client) -> Vec { + client + .queues() + .list(QueueListParams::default()) + .await + .unwrap() + .into_iter() + .filter(|queue| queue.paused_at.is_some()) + .map(|queue| queue.name) + .collect() +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +mod postgres { + use sqlx::{PgPool, Postgres, Transaction}; + + use super::*; + use crate::support::PostgresSchema; + + struct Fixture { + client: Client, + pool: PgPool, + schema: PostgresSchema, + } + + impl Fixture { + async fn new() -> Self { + let schema = PostgresSchema::new("river_handles").await; + let client = builder(&schema).build().unwrap(); + Self { + client, + pool: schema.pool.clone(), + schema, + } + } + + fn builder(&self) -> riverqueue::ClientBuilder { + builder(&self.schema) + } + + async fn begin(&self) -> Transaction<'static, Postgres> { + self.pool.begin().await.unwrap() + } + + async fn job_count(&self) -> i64 { + sqlx::query_scalar(sqlx::AssertSqlSafe(format!( + "SELECT count(*) FROM {}", + self.schema.table("river_job") + ))) + .fetch_one(&self.pool) + .await + .unwrap() + } + + async fn insert_queue(&self, name: &str) { + sqlx::query(sqlx::AssertSqlSafe(format!( + "INSERT INTO {} (name, created_at, metadata, updated_at) \ + VALUES ($1, now(), '{{}}', now())", + self.schema.table("river_queue") + ))) + .bind(name) + .execute(&self.pool) + .await + .unwrap(); + } + + async fn set_attempted_by(&self, id: i64, attempted_by: &[String]) { + sqlx::query(sqlx::AssertSqlSafe(format!( + "UPDATE {} SET attempted_by = $2 WHERE id = $1", + self.schema.table("river_job") + ))) + .bind(id) + .bind(attempted_by) + .execute(&self.pool) + .await + .unwrap(); + } + + async fn cleanup(self) { + self.schema.cleanup().await; + } + } + + fn builder(schema: &PostgresSchema) -> riverqueue::ClientBuilder { + Client::builder( + riverqueue::database::PostgresDatabase::new(schema.pool.clone()) + .with_schema(schema.schema.clone()), + ) + } + + scenarios!(); +} + +#[cfg(feature = "sqlite")] +mod sqlite { + use sqlx::{Sqlite, SqlitePool, Transaction}; + + use super::*; + use crate::support::{sqlite_cleanup, sqlite_file_pool}; + + struct Fixture { + client: Client, + path: std::path::PathBuf, + pool: SqlitePool, + } + + impl Fixture { + async fn new() -> Self { + let (pool, path) = sqlite_file_pool(4).await; + let client = Client::builder(pool.clone()).build().unwrap(); + Self { client, path, pool } + } + + fn builder(&self) -> riverqueue::ClientBuilder { + Client::builder(self.pool.clone()) + } + + async fn begin(&self) -> Transaction<'static, Sqlite> { + self.pool.begin_with("BEGIN IMMEDIATE").await.unwrap() + } + + async fn job_count(&self) -> i64 { + sqlx::query_scalar("SELECT count(*) FROM river_job") + .fetch_one(&self.pool) + .await + .unwrap() + } + + async fn insert_queue(&self, name: &str) { + sqlx::query("INSERT INTO river_queue (name, metadata) VALUES (?, jsonb('{}'))") + .bind(name) + .execute(&self.pool) + .await + .unwrap(); + } + + async fn set_attempted_by(&self, id: i64, attempted_by: &[String]) { + sqlx::query("UPDATE river_job SET attempted_by = jsonb(?) WHERE id = ?") + .bind(serde_json::to_string(attempted_by).unwrap()) + .bind(id) + .execute(&self.pool) + .await + .unwrap(); + } + + async fn cleanup(self) { + sqlite_cleanup(self.pool, self.path).await; + } + } + + scenarios!(); + + /// SQLite delivers notifications through an outbox table, which shows + /// that a resignation request is sent only when its transaction commits. + #[tokio::test(flavor = "multi_thread")] + async fn resign_requests_are_sent_when_the_transaction_commits() { + let fixture = Fixture::new().await; + let requests = async || -> i64 { + sqlx::query_scalar("SELECT count(*) FROM river_notification WHERE topic = ?") + .bind(riverqueue::protocol::NOTIFICATION_TOPIC_LEADERSHIP) + .fetch_one(&fixture.pool) + .await + .unwrap() + }; + + let mut tx = fixture.begin().await; + fixture.client.request_resign().tx(&mut tx).await.unwrap(); + tx.rollback().await.unwrap(); + assert_eq!(requests().await, 0); + + let mut tx = fixture.begin().await; + fixture.client.request_resign().tx(&mut tx).await.unwrap(); + tx.commit().await.unwrap(); + assert_eq!(requests().await, 1); + + fixture.client.request_resign().await.unwrap(); + assert_eq!(requests().await, 2); + + fixture.cleanup().await; + } +} diff --git a/rust/riverqueue/tests/database_contract.rs b/rust/riverqueue/tests/database_contract.rs new file mode 100644 index 000000000..e1caa7772 --- /dev/null +++ b/rust/riverqueue/tests/database_contract.rs @@ -0,0 +1,119 @@ +use riverqueue::__private::Database; +#[cfg(feature = "postgres")] +use riverqueue::database::PostgresDatabase; +#[cfg(feature = "postgres")] +use riverqueue::database::SchemaName; +#[cfg(feature = "sqlite")] +use riverqueue::database::SqliteDatabase; +use riverqueue::database::{DatabaseKind, DatabaseTransactionExecutor, IntoDatabase}; +use sqlx::Transaction; +#[cfg(feature = "postgres")] +use sqlx::{ + PgPool, Postgres, + postgres::{PgConnectOptions, PgPoolOptions}, +}; +#[cfg(feature = "sqlite")] +use sqlx::{ + Sqlite, SqlitePool, + sqlite::{SqliteConnectOptions, SqlitePoolOptions}, +}; + +#[cfg(feature = "postgres")] +fn postgres_pool() -> PgPool { + PgPoolOptions::new().connect_lazy_with(PgConnectOptions::new()) +} + +#[cfg(feature = "sqlite")] +fn sqlite_pool() -> SqlitePool { + SqlitePoolOptions::new().connect_lazy_with(SqliteConnectOptions::new()) +} + +fn require_executor<'executor, E: DatabaseTransactionExecutor<'executor>>( + database: &Database, + executor: E, +) -> DatabaseKind { + database.executor(executor).unwrap().kind() +} + +fn require_source(database: D) -> Database { + Database::from_source(database) +} + +#[test] +fn transactions_implement_the_executor_contract() { + #[cfg(feature = "postgres")] + fn postgres_transaction(database: &Database, transaction: &mut Transaction<'_, Postgres>) { + assert_eq!( + require_executor(database, transaction), + DatabaseKind::Postgres + ); + } + + #[cfg(feature = "sqlite")] + fn sqlite_transaction(database: &Database, transaction: &mut Transaction<'_, Sqlite>) { + assert_eq!( + require_executor(database, transaction), + DatabaseKind::Sqlite + ); + } + + // These function-pointer assignments compile only while the sealed public + // contract accepts actual SQLx transactions. No live server is needed. + #[cfg(feature = "postgres")] + let _: for<'executor, 'transaction> fn( + &Database, + &'executor mut Transaction<'transaction, Postgres>, + ) = postgres_transaction; + #[cfg(feature = "sqlite")] + let _: for<'executor, 'transaction> fn( + &Database, + &'executor mut Transaction<'transaction, Sqlite>, + ) = sqlite_transaction; +} + +#[tokio::test] +async fn pool_sources_preserve_backend_options() { + #[cfg(feature = "postgres")] + let postgres_pool = postgres_pool(); + #[cfg(feature = "sqlite")] + let sqlite_pool = sqlite_pool(); + #[cfg(feature = "postgres")] + let schema = SchemaName::new("river_other").unwrap(); + + #[cfg(feature = "postgres")] + let postgres = require_source(PostgresDatabase::new(postgres_pool.clone()).with_schema(schema)); + #[cfg(feature = "postgres")] + assert_eq!(postgres.kind(), DatabaseKind::Postgres); + #[cfg(feature = "postgres")] + assert_eq!( + postgres.postgres_schema().and_then(SchemaName::as_deref), + Some("river_other") + ); + #[cfg(feature = "sqlite")] + let sqlite = require_source(SqliteDatabase::new(sqlite_pool.clone())); + #[cfg(feature = "sqlite")] + assert_eq!(sqlite.kind(), DatabaseKind::Sqlite); + #[cfg(feature = "sqlite")] + assert_eq!(sqlite.postgres_schema(), None); + #[cfg(feature = "postgres")] + assert_eq!(require_source(postgres_pool).kind(), DatabaseKind::Postgres); + #[cfg(feature = "sqlite")] + assert_eq!(require_source(sqlite_pool).kind(), DatabaseKind::Sqlite); +} + +#[tokio::test] +#[cfg(all(feature = "postgres", feature = "sqlite"))] +async fn rejects_an_executor_from_another_backend() { + let postgres = require_source(postgres_pool()); + let sqlite_pool = sqlite_pool(); + + let Err(error) = postgres.executor(&sqlite_pool) else { + panic!("SQLite executor should not be accepted by a PostgreSQL client"); + }; + assert_eq!(error.expected(), DatabaseKind::Postgres); + assert_eq!(error.actual(), DatabaseKind::Sqlite); + assert_eq!( + error.to_string(), + "database executor mismatch: expected PostgreSQL, received SQLite" + ); +} diff --git a/rust/riverqueue/tests/extension_seams.rs b/rust/riverqueue/tests/extension_seams.rs new file mode 100644 index 000000000..6842fb945 --- /dev/null +++ b/rust/riverqueue/tests/extension_seams.rs @@ -0,0 +1,632 @@ +//! Exact-version seams that add-on crates build on: filtered finalized-job +//! deletion, batched insertion interception, and extension insert options. +//! +//! PostgreSQL scenarios run in a unique schema and fail rather than skip when +//! `RIVER_RUST_DATABASE_URL` is unset; SQLite scenarios use temporary files. + +#![cfg(any(all(feature = "postgres", river_postgres_tests), feature = "sqlite"))] + +mod support; + +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; +use riverqueue::__private::{ + ClientBuilderExt, DatabaseConnection, InsertOptsExt, Pilot, PilotError, +}; +use riverqueue::{Client, InsertOpts, JobArgs}; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value, json}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "extension_seams")] +struct SeamArgs { + value: i64, +} + +/// Finalized jobs seeded for deletion: `(queue, state)`, all finalized an +/// hour ago. +const FINALIZED_SEEDS: [(&str, &str); 6] = [ + ("alpha", "completed"), + ("alpha", "cancelled"), + ("alpha", "discarded"), + ("beta", "completed"), + ("gamma", "completed"), + ("gamma", "discarded"), +]; + +/// Deletion filters and the seeds (indexes into [`FINALIZED_SEEDS`]) each +/// leaves in place, applied in order to the same rows. +fn finalized_deletions() -> Vec<(riverqueue::__private::FinalizedJobDeleteParams, Vec)> { + use riverqueue::__private::FinalizedJobDeleteParams; + + let before = chrono::Utc::now(); + let mut none = FinalizedJobDeleteParams::new(100); + none.queues_included = Some(vec!["alpha".to_owned()]); + + let mut alpha_completed = FinalizedJobDeleteParams::new(100); + alpha_completed.completed_before = Some(before); + alpha_completed.queues_included = Some(vec!["alpha".to_owned()]); + + let mut not_gamma = FinalizedJobDeleteParams::new(100); + not_gamma.completed_before = Some(before); + not_gamma.discarded_before = Some(before); + not_gamma.queues_excluded = vec!["gamma".to_owned()]; + + let mut limited = FinalizedJobDeleteParams::new(1); + limited.completed_before = Some(before); + limited.discarded_before = Some(before); + + vec![ + // No horizon keeps everything, like the cleaner's `None` retention. + (none, vec![0, 1, 2, 3, 4, 5]), + (alpha_completed, vec![1, 2, 3, 4, 5]), + (not_gamma, vec![1, 4, 5]), + // The lowest ID goes first. + (limited, vec![1, 5]), + ] +} + +/// Queues of the jobs seeded for each [`QueueFilterCase`], in ID order. +/// Retained jobs in `kept1`/`kept2` come first and outnumber a batch of 2, so +/// a query that limits candidates before applying queue filters selects only +/// retained jobs, deletes nothing, and stops the cleaner from making progress. +/// Seeds rotate through the finalized states so every state's branch of the +/// query is exercised. +const QUEUE_FILTER_SEEDS: [&str; 11] = [ + "kept1", "kept2", "kept1", "kept2", "kept1", "kept2", "deleted1", "deleted2", "deleted1", + "deleted2", "deleted1", +]; + +/// Finalized state of the seed at each index of [`QUEUE_FILTER_SEEDS`]. +fn queue_filter_seed_state(index: usize) -> &'static str { + ["cancelled", "completed", "discarded"][index % 3] +} + +/// Repeated cleaner batches over [`QUEUE_FILTER_SEEDS`] with one set of +/// queue filters. +struct QueueFilterCase { + /// Jobs deleted by each successive batch. + batches: Vec, + /// Queues whose jobs are eligible for deletion. + deleted_queues: Vec<&'static str>, + name: &'static str, + params: riverqueue::__private::FinalizedJobDeleteParams, +} + +impl QueueFilterCase { + /// Seed IDs left after `deleted` jobs are gone: batches delete the oldest + /// eligible jobs first. + fn remaining(&self, ids: &[i64], deleted: u64) -> Vec { + let mut deleted = usize::try_from(deleted).unwrap(); + ids.iter() + .zip(QUEUE_FILTER_SEEDS) + .filter(|(_, queue)| { + if deleted > 0 && self.deleted_queues.contains(queue) { + deleted -= 1; + return false; + } + true + }) + .map(|(&id, _)| id) + .collect() + } + + /// Number of seeds eligible for deletion. + fn eligible(&self) -> u64 { + QUEUE_FILTER_SEEDS + .iter() + .filter(|queue| self.deleted_queues.contains(queue)) + .count() + .try_into() + .unwrap() + } +} + +fn queue_filter_cases() -> Vec { + let case = |name: &'static str, + queues_excluded: &[&str], + queues_included: Option<&[&str]>, + batches: Vec, + deleted_queues: Vec<&'static str>| { + let horizon = chrono::Utc::now(); + let mut params = riverqueue::__private::FinalizedJobDeleteParams::new(2); + params.cancelled_before = Some(horizon); + params.completed_before = Some(horizon); + params.discarded_before = Some(horizon); + params.queues_excluded = queues_excluded + .iter() + .map(|&queue| queue.to_owned()) + .collect(); + params.queues_included = + queues_included.map(|queues| queues.iter().map(|&queue| queue.to_owned()).collect()); + QueueFilterCase { + batches, + deleted_queues, + name, + params, + } + }; + let all = vec!["deleted1", "deleted2", "kept1", "kept2"]; + vec![ + // `kept1` appears in both lists; exclusion takes precedence. + case( + "both", + &["kept1", "kept2"], + Some(&["deleted1", "deleted2", "kept1"]), + vec![2, 2, 1, 0], + vec!["deleted1", "deleted2"], + ), + // An empty inclusion list matches no queues, unlike `None`. + case("empty_included", &[], Some(&[]), vec![0], vec![]), + case( + "excluded", + &["kept1", "kept2"], + None, + vec![2, 2, 1, 0], + vec!["deleted1", "deleted2"], + ), + case( + "included", + &[], + Some(&["deleted1", "deleted2"]), + vec![2, 2, 1, 0], + vec!["deleted1", "deleted2"], + ), + case("missing_included", &[], Some(&["missing"]), vec![0], vec![]), + case("no_filters", &[], None, vec![2, 2, 2, 2, 2, 1, 0], all), + ] +} + +/// Records what River binds the pilot to when its client is built. +#[derive(Clone, Default)] +struct InstallPilot { + installs: Arc>>, +} + +impl Pilot for InstallPilot { + fn install(&self, context: riverqueue::__private::PilotInstallContext) { + self.installs.lock().unwrap().push(context); + } +} + +/// River binds its pilot to the client once, as it's built: to the client +/// itself without keeping it alive, its database, and the report interval. +async fn assert_pilot_is_installed_once(builder: impl Fn() -> riverqueue::ClientBuilder) { + let pilot = InstallPilot::default(); + let client = builder() + .pilot(pilot.clone()) + .producer_report_interval(std::time::Duration::from_secs(7)) + .build() + .unwrap(); + let installs = std::mem::take(&mut *pilot.installs.lock().unwrap()); + assert_eq!(installs.len(), 1); + let install = &installs[0]; + assert_eq!(install.database.kind(), client.database().kind()); + assert_eq!( + install.producer_report_interval, + std::time::Duration::from_secs(7) + ); + assert_eq!(install.client.upgrade().unwrap().id(), client.id()); + drop(client); + assert!(install.client.upgrade().is_none()); + // The database works on its own. + install + .database + .begin() + .await + .unwrap() + .commit() + .await + .unwrap(); +} + +/// Records each batch [`Pilot::before_jobs_insert`] receives and tags its +/// jobs. +struct BatchInsertPilot { + batches: Arc>>, +} + +#[async_trait] +impl Pilot for BatchInsertPilot { + fn intercepts_insert(&self) -> bool { + true + } + + async fn before_jobs_insert( + &self, + _connection: DatabaseConnection<'_>, + jobs: &mut [riverqueue::__private::JobInsertParams<'_>], + ) -> Result<(), PilotError> { + self.batches.lock().unwrap().push(jobs.len()); + for job in jobs { + job.metadata.insert("batched", true)?; + } + Ok(()) + } +} + +/// Checks that each insertion call reaches the extension as one batch. +async fn assert_batched_insert_interception(builder: impl Fn() -> riverqueue::ClientBuilder) { + let batches = Arc::new(Mutex::new(Vec::new())); + let client = builder() + .pilot(BatchInsertPilot { + batches: Arc::clone(&batches), + }) + .build() + .unwrap(); + + let single = client.insert(SeamArgs { value: 1 }).await.unwrap(); + assert_eq!( + single.job.row.metadata.get::("batched").unwrap(), + Some(true) + ); + let many = client + .insert_many((2..=4).map(|value| SeamArgs { value })) + .await + .unwrap(); + assert!( + many.iter() + .all(|job| job.job.row.metadata.contains_key("batched")) + ); + assert_eq!(*batches.lock().unwrap(), [1, 3]); +} + +/// A job type that declares extension options and default metadata. +#[derive(Clone, Debug, Deserialize, Serialize)] +struct DeclaredArgs { + value: i64, +} + +impl JobArgs for DeclaredArgs { + const KIND: &'static str = "extension_seams_declared"; + + fn default_insert_opts() -> InsertOpts { + InsertOpts::default() + .with_metadata(json!({"team": "a"}).as_object().unwrap().clone()) + .with_extension_option("declared", json!({"type": true})) + .with_extension_option("shared", json!("type")) + } +} + +/// The extension options and metadata JSON a job reached the insert hook +/// with. +type SeenInsert = (Map, String); + +/// Records what each inserted job reaches the insert hook with. +#[derive(Clone, Default)] +struct OptionsPilot { + seen: Arc>>, +} + +#[async_trait] +impl Pilot for OptionsPilot { + fn intercepts_insert(&self) -> bool { + true + } + + async fn before_jobs_insert( + &self, + _connection: DatabaseConnection<'_>, + jobs: &mut [riverqueue::__private::JobInsertParams<'_>], + ) -> Result<(), PilotError> { + let mut seen = self.seen.lock().unwrap(); + for job in jobs { + seen.push(( + job.extension_options.clone(), + job.metadata.as_raw().get().to_owned(), + )); + } + Ok(()) + } +} + +/// Checks that extension options reach the insert hook resolved key by key +/// and aren't persisted, while per-call metadata replaces the job type's +/// default metadata wholesale, as in Go. +async fn assert_extension_options_reach_the_insert_hook( + builder: impl Fn() -> riverqueue::ClientBuilder, +) { + let pilot = OptionsPilot::default(); + let client = builder().pilot(pilot.clone()).build().unwrap(); + + let defaults = client.insert(DeclaredArgs { value: 1 }).await.unwrap(); + let overridden = client + .insert(DeclaredArgs { value: 2 }) + .opts( + InsertOpts::default() + .with_metadata(json!({"call": 1}).as_object().unwrap().clone()) + .with_extension_option("shared", json!("call")), + ) + .await + .unwrap(); + + let seen = pilot.seen.lock().unwrap().clone(); + assert_eq!(seen.len(), 2); + assert_eq!( + Value::Object(seen[0].0.clone()), + json!({"declared": {"type": true}, "shared": "type"}) + ); + assert_eq!(seen[0].1, r#"{"team":"a"}"#); + // A call's metadata replaces the defaults, but not the declared + // extension options, which the call overrides key by key. + assert_eq!( + Value::Object(seen[1].0.clone()), + json!({"declared": {"type": true}, "shared": "call"}) + ); + assert_eq!(seen[1].1, r#"{"call":1}"#); + // Nothing about the extension options is persisted. (SQLite adds its + // insert nonce.) + for (id, expected) in [ + (defaults.id(), json!({"team": "a"})), + (overridden.id(), json!({"call": 1})), + ] { + let mut metadata: Map = + serde_json::from_str(client.jobs().get(id).await.unwrap().metadata.as_raw().get()) + .unwrap(); + metadata.remove("river:unique_nonce"); + assert_eq!(Value::Object(metadata), expected); + } +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +mod postgres { + use riverqueue::database::PostgresDatabase; + use sqlx::AssertSqlSafe; + + use super::*; + use crate::support::PostgresSchema; + + fn builder(schema: &PostgresSchema) -> riverqueue::ClientBuilder { + Client::builder( + PostgresDatabase::new(schema.pool.clone()).with_schema(schema.schema.clone()), + ) + } + + #[tokio::test(flavor = "multi_thread")] + async fn deletes_finalized_jobs_with_the_cleaner_filters() { + use riverqueue::__private::{DatabaseConfig, delete_finalized_jobs}; + + let schema = PostgresSchema::new("seam_finalized_delete").await; + let mut ids = Vec::new(); + for (queue, state) in FINALIZED_SEEDS { + let id: i64 = sqlx::query_scalar(AssertSqlSafe(format!( + "INSERT INTO {} (args, finalized_at, kind, max_attempts, queue, state) \ + VALUES ('{{}}', now() - interval '1 hour', 'extension_seams', 25, $1, \ + $2::text::{}) RETURNING id", + schema.table("river_job"), + schema.table("river_job_state"), + ))) + .bind(queue) + .bind(state) + .fetch_one(&schema.pool) + .await + .unwrap(); + ids.push(id); + } + let database = DatabaseConfig::Postgres { + schema: schema.schema.clone(), + }; + for (params, kept) in finalized_deletions() { + let mut connection = schema.pool.acquire().await.unwrap(); + delete_finalized_jobs( + DatabaseConnection::Postgres(&mut connection), + &database, + ¶ms, + ) + .await + .unwrap(); + let remaining: Vec = sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT id FROM {} ORDER BY id", + schema.table("river_job") + ))) + .fetch_all(&schema.pool) + .await + .unwrap(); + let expected: Vec = kept.iter().map(|&index| ids[index]).collect(); + assert_eq!(remaining, expected, "{params:?}"); + } + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn filters_queues_before_the_batch_limit() { + use riverqueue::__private::{DatabaseConfig, delete_finalized_jobs}; + + let schema = PostgresSchema::new("seam_queue_filters").await; + let table = schema.table("river_job"); + let database = DatabaseConfig::Postgres { + schema: schema.schema.clone(), + }; + for case in queue_filter_cases() { + sqlx::query(AssertSqlSafe(format!("DELETE FROM {table}"))) + .execute(&schema.pool) + .await + .unwrap(); + let mut ids = Vec::new(); + for (index, queue) in QUEUE_FILTER_SEEDS.into_iter().enumerate() { + let id: i64 = sqlx::query_scalar(AssertSqlSafe(format!( + "INSERT INTO {table} (args, finalized_at, kind, max_attempts, queue, state) \ + VALUES ('{{}}', now() - interval '1 hour', 'extension_seams', 25, $1, \ + $2::text::{}) RETURNING id", + schema.table("river_job_state"), + ))) + .bind(queue) + .bind(queue_filter_seed_state(index)) + .fetch_one(&schema.pool) + .await + .unwrap(); + ids.push(id); + } + let mut deleted_total = 0; + for (batch, &want) in case.batches.iter().enumerate() { + let mut connection = schema.pool.acquire().await.unwrap(); + let deleted = delete_finalized_jobs( + DatabaseConnection::Postgres(&mut connection), + &database, + &case.params, + ) + .await + .unwrap(); + drop(connection); + assert_eq!(deleted, want, "{} batch {batch}", case.name); + deleted_total += deleted; + let remaining: Vec = sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT id FROM {table} ORDER BY id" + ))) + .fetch_all(&schema.pool) + .await + .unwrap(); + assert_eq!( + remaining, + case.remaining(&ids, deleted_total), + "{} batch {batch}", + case.name + ); + } + assert_eq!(deleted_total, case.eligible(), "{}", case.name); + } + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn pilot_is_installed_once() { + let schema = PostgresSchema::new("seam_install").await; + assert_pilot_is_installed_once(|| builder(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn extension_options_reach_the_insert_hook() { + let schema = PostgresSchema::new("seam_extension_options").await; + assert_extension_options_reach_the_insert_hook(|| builder(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn insertions_reach_the_extension_as_one_batch() { + let schema = PostgresSchema::new("seam_batch_insert").await; + assert_batched_insert_interception(|| builder(&schema)).await; + schema.cleanup().await; + } +} + +#[cfg(feature = "sqlite")] +mod sqlite { + use super::*; + use crate::support::{sqlite_cleanup, sqlite_file_pool}; + + #[tokio::test(flavor = "multi_thread")] + async fn deletes_finalized_jobs_with_the_cleaner_filters() { + use riverqueue::__private::{DatabaseConfig, delete_finalized_jobs}; + + let (pool, path) = sqlite_file_pool(4).await; + let mut ids = Vec::new(); + for (queue, state) in FINALIZED_SEEDS { + let id: i64 = sqlx::query_scalar( + "INSERT INTO river_job (args, finalized_at, kind, max_attempts, queue, state) \ + VALUES (jsonb('{}'), strftime('%Y-%m-%d %H:%M:%f', 'now', '-1 hour'), \ + 'extension_seams', 25, ?, ?) RETURNING id", + ) + .bind(queue) + .bind(state) + .fetch_one(&pool) + .await + .unwrap(); + ids.push(id); + } + for (params, kept) in finalized_deletions() { + let mut connection = pool.acquire().await.unwrap(); + delete_finalized_jobs( + DatabaseConnection::Sqlite(&mut connection), + &DatabaseConfig::Sqlite, + ¶ms, + ) + .await + .unwrap(); + drop(connection); + let remaining: Vec = sqlx::query_scalar("SELECT id FROM river_job ORDER BY id") + .fetch_all(&pool) + .await + .unwrap(); + let expected: Vec = kept.iter().map(|&index| ids[index]).collect(); + assert_eq!(remaining, expected, "{params:?}"); + } + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn filters_queues_before_the_batch_limit() { + use riverqueue::__private::{DatabaseConfig, delete_finalized_jobs}; + + let (pool, path) = sqlite_file_pool(4).await; + for case in queue_filter_cases() { + sqlx::query("DELETE FROM river_job") + .execute(&pool) + .await + .unwrap(); + let mut ids = Vec::new(); + for (index, queue) in QUEUE_FILTER_SEEDS.into_iter().enumerate() { + let id: i64 = sqlx::query_scalar( + "INSERT INTO river_job (args, finalized_at, kind, max_attempts, queue, state) \ + VALUES (jsonb('{}'), strftime('%Y-%m-%d %H:%M:%f', 'now', '-1 hour'), \ + 'extension_seams', 25, ?, ?) RETURNING id", + ) + .bind(queue) + .bind(queue_filter_seed_state(index)) + .fetch_one(&pool) + .await + .unwrap(); + ids.push(id); + } + let mut deleted_total = 0; + for (batch, &want) in case.batches.iter().enumerate() { + let mut connection = pool.acquire().await.unwrap(); + let deleted = delete_finalized_jobs( + DatabaseConnection::Sqlite(&mut connection), + &DatabaseConfig::Sqlite, + &case.params, + ) + .await + .unwrap(); + drop(connection); + assert_eq!(deleted, want, "{} batch {batch}", case.name); + deleted_total += deleted; + let remaining: Vec = + sqlx::query_scalar("SELECT id FROM river_job ORDER BY id") + .fetch_all(&pool) + .await + .unwrap(); + assert_eq!( + remaining, + case.remaining(&ids, deleted_total), + "{} batch {batch}", + case.name + ); + } + assert_eq!(deleted_total, case.eligible(), "{}", case.name); + } + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn pilot_is_installed_once() { + let (pool, path) = sqlite_file_pool(4).await; + assert_pilot_is_installed_once(|| Client::builder(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn extension_options_reach_the_insert_hook() { + let (pool, path) = sqlite_file_pool(4).await; + assert_extension_options_reach_the_insert_hook(|| Client::builder(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn insertions_reach_the_extension_as_one_batch() { + let (pool, path) = sqlite_file_pool(4).await; + assert_batched_insert_interception(|| Client::builder(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } +} diff --git a/rust/riverqueue/tests/extension_services.rs b/rust/riverqueue/tests/extension_services.rs new file mode 100644 index 000000000..d3d0b670f --- /dev/null +++ b/rust/riverqueue/tests/extension_services.rs @@ -0,0 +1,394 @@ +//! Extension services and the client's stop order. +//! +//! Maintenance services get their leadership term and are supervised within +//! it. On a graceful stop, claims, leadership with its maintenance services, +//! and runtime services end at once, like River Go's services started on its +//! fetch context, while each producer keeps reporting to its extension +//! session until its running jobs finish. +//! +//! PostgreSQL scenarios run in a unique schema and fail rather than skip when +//! `RIVER_RUST_DATABASE_URL` is unset; SQLite scenarios use temporary files. + +#![cfg(any(all(feature = "postgres", river_postgres_tests), feature = "sqlite"))] + +mod support; + +use std::{ + convert::Infallible, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use riverqueue::__private::{ + ClientBuilderExt, ExtensionClient, MaintenanceService, MaintenanceServiceContext, Pilot, + PilotError, PilotProducer, ProducerKeepAliveContext, ProducerShutdownContext, + ProducerStartContext, RuntimeService, RuntimeServiceContext, +}; +use riverqueue::{ + Client, Job, JobArgs, JobRow, JobState, MaintenanceConfig, QueueConfig, WorkContext, + WorkOutcome, WorkerRegistry, +}; +use serde::{Deserialize, Serialize}; +use tokio::sync::{Notify, Semaphore}; + +const WAIT: Duration = Duration::from_secs(10); + +/// Bounds a whole scenario, so a stop that never finishes fails the test +/// instead of hanging it. +const SCENARIO: Duration = Duration::from_mins(1); + +/// A job whose worker holds its slot until the test releases it. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "extension_services_gated")] +struct GatedArgs {} + +/// Everything River told the extension, in order. +#[derive(Clone, Debug, PartialEq)] +enum Call { + Finished(i64), + KeepAlive, + MaintenanceStarted(DateTime), + MaintenanceStopped, + RuntimeStarted, + RuntimeStopped, + Shutdown, +} + +#[derive(Default)] +struct Calls { + changed: Notify, + calls: Mutex>, +} + +impl Calls { + fn push(&self, call: Call) { + self.calls.lock().unwrap().push(call); + self.changed.notify_waiters(); + } + + fn snapshot(&self) -> Vec { + self.calls.lock().unwrap().clone() + } + + fn position(&self, call: &Call) -> Option { + self.snapshot().iter().position(|recorded| recorded == call) + } + + async fn wait_until(&self, what: &str, done: impl Fn(&[Call]) -> bool) { + tokio::time::timeout(WAIT, async { + loop { + let changed = self.changed.notified(); + if done(&self.calls.lock().unwrap()) { + return; + } + changed.await; + } + }) + .await + .unwrap_or_else(|_| panic!("timed out waiting for {what}")); + } +} + +/// How the maintenance service's runs end. +#[derive(Clone, Copy, Debug)] +enum Runs { + /// Every run waits for its term to end. + UntilTermEnds, + /// The first run fails, the second panics, and later runs wait for the + /// term to end. + FailThenPanic, +} + +#[derive(Clone)] +struct ServicePilot { + calls: Arc, + maintenance_runs: Arc, + runs: Runs, +} + +impl ServicePilot { + fn new(runs: Runs) -> Self { + Self { + calls: Arc::default(), + maintenance_runs: Arc::default(), + runs, + } + } +} + +#[async_trait] +impl Pilot for ServicePilot { + fn maintenance_services(&self) -> Vec> { + vec![Arc::new(self.clone())] + } + + fn runtime_services(&self) -> Vec> { + vec![Arc::new(self.clone())] + } + + async fn start_producer( + &self, + _context: ProducerStartContext, + ) -> Result>, PilotError> { + Ok(Some(Box::new(self.clone()))) + } +} + +#[async_trait] +impl MaintenanceService for ServicePilot { + fn name(&self) -> &'static str { + "test maintenance" + } + + async fn run(&self, context: MaintenanceServiceContext) -> Result<(), PilotError> { + // The service reaches its client and, through it, this pilot. A + // failure here restarts the service, which the recorded calls show. + let client = context.client.upgrade().ok_or("client is gone")?; + ExtensionClient::new(&client) + .pilot::() + .ok_or("pilot not found")?; + let run = self.maintenance_runs.fetch_add(1, Ordering::SeqCst); + self.calls + .push(Call::MaintenanceStarted(context.term.elected_at)); + match (self.runs, run) { + (Runs::FailThenPanic, 0) => return Err("maintenance failed on purpose".into()), + (Runs::FailThenPanic, 1) => panic!("maintenance panicked on purpose"), + _ => {} + } + // The service can use the client's database. + context.database.begin().await?.commit().await?; + context.term.token.cancelled().await; + self.calls.push(Call::MaintenanceStopped); + Ok(()) + } +} + +#[async_trait] +impl RuntimeService for ServicePilot { + async fn run(&self, context: RuntimeServiceContext) -> Result<(), PilotError> { + let client = context.client.upgrade().ok_or("client is gone")?; + ExtensionClient::new(&client) + .pilot::() + .ok_or("pilot not found")?; + self.calls.push(Call::RuntimeStarted); + context.database.begin().await?.commit().await?; + context.cancellation.cancelled().await; + self.calls.push(Call::RuntimeStopped); + Ok(()) + } +} + +#[async_trait] +impl PilotProducer for ServicePilot { + fn job_finished(&self, job: &JobRow) { + self.calls.push(Call::Finished(job.id)); + } + + async fn keep_alive(&self, _context: ProducerKeepAliveContext) -> Result<(), PilotError> { + self.calls.push(Call::KeepAlive); + Ok(()) + } + + async fn shutdown(&self, _context: ProducerShutdownContext) -> Result<(), PilotError> { + self.calls.push(Call::Shutdown); + Ok(()) + } +} + +fn client( + builder: riverqueue::ClientBuilder, + pilot: &ServicePilot, + release: &Arc, +) -> Client { + let mut workers = WorkerRegistry::new(); + let release = Arc::clone(release); + workers + .register_fn(move |_context: WorkContext, _job: Job| { + let release = Arc::clone(&release); + async move { + release.acquire().await.unwrap().forget(); + Ok::<_, Infallible>(WorkOutcome::Complete) + } + }) + .unwrap(); + builder + .pilot(pilot.clone()) + .producer_report_interval(Duration::from_millis(20)) + .maintenance(MaintenanceConfig::default().with_elect_interval(Duration::from_millis(100))) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .workers(workers) + .build() + .unwrap() +} + +/// A graceful stop ends leadership, maintenance services, and runtime +/// services at once, while the producer keeps reporting until its running +/// job finishes and only then shuts its session down. +async fn assert_stop_order_matches_go(builder: riverqueue::ClientBuilder) { + let pilot = ServicePilot::new(Runs::UntilTermEnds); + let release = Arc::new(Semaphore::new(0)); + let client = client(builder, &pilot, &release); + let id = client.insert(GatedArgs {}).await.unwrap().id(); + let mut run = client.start().unwrap(); + pilot + .calls + .wait_until("leadership and a running job", |calls| { + calls + .iter() + .any(|call| matches!(call, Call::MaintenanceStarted(_))) + && calls.contains(&Call::RuntimeStarted) + }) + .await; + tokio::time::timeout(WAIT, async { + while client.jobs().get(id).await.unwrap().state != JobState::Running { + tokio::task::yield_now().await; + } + }) + .await + .expect("job starts"); + + run.stopper().stop(); + pilot + .calls + .wait_until("services to stop", |calls| { + calls.contains(&Call::MaintenanceStopped) && calls.contains(&Call::RuntimeStopped) + }) + .await; + let stopped = pilot.calls.snapshot().len(); + pilot + .calls + .wait_until("reports after the services stopped", |calls| { + calls[stopped..] + .iter() + .filter(|call| **call == Call::KeepAlive) + .count() + >= 2 + }) + .await; + assert_eq!(pilot.calls.position(&Call::Shutdown), None); + assert_eq!( + client.jobs().get(id).await.unwrap().state, + JobState::Running + ); + + release.add_permits(1); + tokio::time::timeout(WAIT, run.wait()) + .await + .expect("client stops") + .unwrap(); + let finished = pilot.calls.position(&Call::Finished(id)).unwrap(); + let shutdown = pilot.calls.position(&Call::Shutdown).unwrap(); + assert!(finished < shutdown); + assert_eq!( + client.jobs().get(id).await.unwrap().state, + JobState::Completed + ); +} + +/// A maintenance service that fails or panics is restarted within its +/// leadership term, and every run gets the same term. +async fn assert_maintenance_services_restart_within_their_term(builder: riverqueue::ClientBuilder) { + let pilot = ServicePilot::new(Runs::FailThenPanic); + let release = Arc::new(Semaphore::new(0)); + let client = client(builder, &pilot, &release); + let mut run = client.start().unwrap(); + pilot + .calls + .wait_until("the third maintenance run", |calls| { + calls + .iter() + .filter(|call| matches!(call, Call::MaintenanceStarted(_))) + .count() + >= 3 + }) + .await; + run.shutdown().await.unwrap(); + + let terms = pilot + .calls + .snapshot() + .into_iter() + .filter_map(|call| match call { + Call::MaintenanceStarted(elected_at) => Some(elected_at), + _ => None, + }) + .collect::>(); + assert_eq!(terms.len(), 3, "{terms:?}"); + assert!(terms.iter().all(|term| *term == terms[0]), "{terms:?}"); + assert!(pilot.calls.snapshot().contains(&Call::MaintenanceStopped)); +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +mod postgres { + use riverqueue::database::PostgresDatabase; + + use super::*; + use crate::support::PostgresSchema; + + fn builder(schema: &PostgresSchema) -> riverqueue::ClientBuilder { + Client::builder( + PostgresDatabase::new(schema.pool.clone()).with_schema(schema.schema.clone()), + ) + } + + #[tokio::test(flavor = "multi_thread")] + async fn maintenance_services_restart_within_their_term() { + let schema = PostgresSchema::new("services_restart").await; + tokio::time::timeout( + SCENARIO, + assert_maintenance_services_restart_within_their_term(builder(&schema)), + ) + .await + .expect("scenario finishes"); + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn stop_order_matches_go() { + let schema = PostgresSchema::new("services_stop_order").await; + tokio::time::timeout(SCENARIO, assert_stop_order_matches_go(builder(&schema))) + .await + .expect("scenario finishes"); + schema.cleanup().await; + } +} + +#[cfg(feature = "sqlite")] +mod sqlite { + use super::*; + use crate::support::{sqlite_cleanup, sqlite_file_pool}; + + #[tokio::test(flavor = "multi_thread")] + async fn maintenance_services_restart_within_their_term() { + let (pool, path) = sqlite_file_pool(4).await; + tokio::time::timeout( + SCENARIO, + assert_maintenance_services_restart_within_their_term(Client::builder(pool.clone())), + ) + .await + .expect("scenario finishes"); + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn stop_order_matches_go() { + let (pool, path) = sqlite_file_pool(4).await; + tokio::time::timeout( + SCENARIO, + assert_stop_order_matches_go(Client::builder(pool.clone())), + ) + .await + .expect("scenario finishes"); + sqlite_cleanup(pool, path).await; + } +} diff --git a/rust/riverqueue/tests/fetch_only_known_kinds.rs b/rust/riverqueue/tests/fetch_only_known_kinds.rs new file mode 100644 index 000000000..4fd737d8c --- /dev/null +++ b/rust/riverqueue/tests/fetch_only_known_kinds.rs @@ -0,0 +1,139 @@ +//! Clients built with `fetch_only_known_kinds`, on every backend. +//! +//! PostgreSQL tests run in a unique schema and fail rather than skip when +//! `RIVER_RUST_DATABASE_URL` is unset; SQLite tests use a temporary file. + +#![cfg(any(all(feature = "postgres", river_postgres_tests), feature = "sqlite"))] + +mod support; + +use std::{collections::HashSet, convert::Infallible, time::Duration}; + +use riverqueue::{ + __private::Database, Client, EventKind, Job, JobArgs, JobState, QueueConfig, WorkContext, + WorkOutcome, WorkerRegistry, +}; +use serde::{Deserialize, Serialize}; + +/// Every wait in these tests is bounded by this timeout. +const TIMEOUT: Duration = Duration::from_secs(10); + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_fetch_known", aliases("rust_fetch_known_old"))] +struct KnownArgs {} + +/// Inserts jobs under the known kind's alias. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_fetch_known_old")] +struct AliasArgs {} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_fetch_unknown")] +struct UnknownArgs {} + +/// A migrated database on one backend. +enum Backend { + #[cfg(all(feature = "postgres", river_postgres_tests))] + Postgres(support::PostgresSchema), + #[cfg(feature = "sqlite")] + Sqlite(sqlx::SqlitePool, std::path::PathBuf), +} + +impl Backend { + fn database(&self) -> Database { + match self { + #[cfg(all(feature = "postgres", river_postgres_tests))] + Self::Postgres(schema) => Database::from_source( + riverqueue::database::PostgresDatabase::new(schema.pool.clone()) + .with_schema(schema.schema.clone()), + ), + #[cfg(feature = "sqlite")] + Self::Sqlite(pool, _) => Database::from_source(pool.clone()), + } + } + + async fn cleanup(self) { + match self { + #[cfg(all(feature = "postgres", river_postgres_tests))] + Self::Postgres(schema) => schema.cleanup().await, + #[cfg(feature = "sqlite")] + Self::Sqlite(pool, path) => support::sqlite_cleanup(pool, path).await, + } + } +} + +/// Works jobs of the registered kind and its alias, and leaves a job of +/// another kind available without using an attempt, even though it's +/// first in the queue. +async fn claims_only_registered_kinds(backend: Backend) { + let inserter = Client::builder(backend.database()).build().unwrap(); + let unknown = inserter.insert(UnknownArgs {}).await.unwrap().job.row.id; + let known = inserter.insert(KnownArgs {}).await.unwrap().job.row.id; + let alias = inserter.insert(AliasArgs {}).await.unwrap().job.row.id; + + let mut workers = WorkerRegistry::new(); + workers + .register_fn(|_context: WorkContext, _job: Job| async { + Ok::<_, Infallible>(WorkOutcome::Complete) + }) + .unwrap(); + let client = Client::builder(backend.database()) + .fetch_only_known_kinds(true) + .queue(riverqueue::QUEUE_DEFAULT, QueueConfig::new(10)) + .without_leader_election() + .workers(workers) + .build() + .unwrap(); + let mut events = client + .subscribe(&[EventKind::JobCompleted, EventKind::JobFailed]) + .unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + let mut completed = HashSet::new(); + while completed.len() < 2 { + let event = tokio::time::timeout(TIMEOUT, events.recv()) + .await + .expect("known jobs should complete") + .unwrap(); + let job = &event.as_job().expect("a job event").job; + assert_eq!(job.state, JobState::Completed, "{job:?}"); + completed.insert(job.id); + } + assert_eq!(completed, HashSet::from([known, alias])); + + tokio::time::timeout(TIMEOUT, run.shutdown()) + .await + .expect("the client should stop") + .unwrap(); + let unknown = client.jobs().get(unknown).await.unwrap(); + assert_eq!(unknown.state, JobState::Available); + assert_eq!(unknown.attempt, 0); + assert_eq!(unknown.attempted_by, Vec::::new()); + + backend.cleanup().await; +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +mod postgres { + use super::*; + + #[tokio::test] + async fn claims_only_registered_kinds() { + super::claims_only_registered_kinds(Backend::Postgres( + support::PostgresSchema::new("river_known_kinds").await, + )) + .await; + } +} + +#[cfg(feature = "sqlite")] +mod sqlite { + use super::*; + + #[tokio::test] + async fn claims_only_registered_kinds() { + let (pool, path) = support::sqlite_file_pool(4).await; + super::claims_only_registered_kinds(Backend::Sqlite(pool, path)).await; + } +} diff --git a/rust/riverqueue/tests/fixtures/maintenance_values.json b/rust/riverqueue/tests/fixtures/maintenance_values.json new file mode 100644 index 000000000..1b69ceb1c --- /dev/null +++ b/rust/riverqueue/tests/fixtures/maintenance_values.json @@ -0,0 +1,679 @@ +{ + "cron_cases": [ + { + "expression": "* * * * *", + "from": "2026-01-02T03:04:05.6789Z", + "name": "every_minute", + "next": [ + "2026-01-02T03:05:00Z", + "2026-01-02T03:06:00Z", + "2026-01-02T03:07:00Z", + "2026-01-02T03:08:00Z", + "2026-01-02T03:09:00Z" + ] + }, + { + "expression": "30 * * * *", + "from": "2026-01-02T03:04:05.6789Z", + "name": "half_past_every_hour", + "next": [ + "2026-01-02T03:30:00Z", + "2026-01-02T04:30:00Z", + "2026-01-02T05:30:00Z", + "2026-01-02T06:30:00Z", + "2026-01-02T07:30:00Z" + ] + }, + { + "expression": "0 9 * * 1", + "from": "2026-01-02T03:04:05.6789Z", + "name": "monday_numeric_weekday", + "next": [ + "2026-01-05T09:00:00Z", + "2026-01-12T09:00:00Z", + "2026-01-19T09:00:00Z", + "2026-01-26T09:00:00Z", + "2026-02-02T09:00:00Z" + ] + }, + { + "expression": "0 9 * * mon", + "from": "2026-01-02T03:04:05.6789Z", + "name": "monday_named_weekday", + "next": [ + "2026-01-05T09:00:00Z", + "2026-01-12T09:00:00Z", + "2026-01-19T09:00:00Z", + "2026-01-26T09:00:00Z", + "2026-02-02T09:00:00Z" + ] + }, + { + "expression": "0 0 * * 0", + "from": "2026-01-02T03:04:05.6789Z", + "name": "sunday_is_zero", + "next": [ + "2026-01-04T00:00:00Z", + "2026-01-11T00:00:00Z", + "2026-01-18T00:00:00Z", + "2026-01-25T00:00:00Z", + "2026-02-01T00:00:00Z" + ] + }, + { + "expression": "0 0 * * SUN", + "from": "2026-01-02T03:04:05.6789Z", + "name": "weekday_names_ignore_case", + "next": [ + "2026-01-04T00:00:00Z", + "2026-01-11T00:00:00Z", + "2026-01-18T00:00:00Z", + "2026-01-25T00:00:00Z", + "2026-02-01T00:00:00Z" + ] + }, + { + "expression": "*/15 9-17 * * mon-fri", + "from": "2026-01-02T03:04:05.6789Z", + "name": "business_hours_steps", + "next": [ + "2026-01-02T09:00:00Z", + "2026-01-02T09:15:00Z", + "2026-01-02T09:30:00Z", + "2026-01-02T09:45:00Z", + "2026-01-02T10:00:00Z" + ] + }, + { + "expression": "0 0 1 * *", + "from": "2026-01-02T03:04:05.6789Z", + "name": "first_of_month", + "next": [ + "2026-02-01T00:00:00Z", + "2026-03-01T00:00:00Z", + "2026-04-01T00:00:00Z", + "2026-05-01T00:00:00Z", + "2026-06-01T00:00:00Z" + ] + }, + { + "expression": "0 0 1 jan,JUL *", + "from": "2026-01-02T03:04:05.6789Z", + "name": "named_months", + "next": [ + "2026-07-01T00:00:00Z", + "2027-01-01T00:00:00Z", + "2027-07-01T00:00:00Z", + "2028-01-01T00:00:00Z", + "2028-07-01T00:00:00Z" + ] + }, + { + "expression": "0 0 29 2 *", + "from": "2026-01-02T03:04:05.6789Z", + "name": "leap_day", + "next": [ + "2028-02-29T00:00:00Z", + "2032-02-29T00:00:00Z", + "2036-02-29T00:00:00Z", + "2040-02-29T00:00:00Z", + "2044-02-29T00:00:00Z" + ] + }, + { + "expression": "0 0 30 2 *", + "from": "2026-01-02T03:04:05.6789Z", + "name": "impossible_date_never_runs", + "next": [] + }, + { + "expression": "0 12 1,15 * 5", + "from": "2026-01-02T03:04:05.6789Z", + "name": "day_of_month_or_weekday", + "next": [ + "2026-01-02T12:00:00Z", + "2026-01-09T12:00:00Z", + "2026-01-15T12:00:00Z", + "2026-01-16T12:00:00Z", + "2026-01-23T12:00:00Z" + ] + }, + { + "expression": "0 12 * * 5", + "from": "2026-01-02T03:04:05.6789Z", + "name": "wildcard_day_of_month_and_weekday", + "next": [ + "2026-01-02T12:00:00Z", + "2026-01-09T12:00:00Z", + "2026-01-16T12:00:00Z", + "2026-01-23T12:00:00Z", + "2026-01-30T12:00:00Z" + ] + }, + { + "expression": "0 12 ? * 5", + "from": "2026-01-02T03:04:05.6789Z", + "name": "question_mark_wildcard", + "next": [ + "2026-01-02T12:00:00Z", + "2026-01-09T12:00:00Z", + "2026-01-16T12:00:00Z", + "2026-01-23T12:00:00Z", + "2026-01-30T12:00:00Z" + ] + }, + { + "expression": "0 12 */2 * 5", + "from": "2026-01-02T03:04:05.6789Z", + "name": "stepped_day_of_month_or_weekday", + "next": [ + "2026-01-02T12:00:00Z", + "2026-01-03T12:00:00Z", + "2026-01-05T12:00:00Z", + "2026-01-07T12:00:00Z", + "2026-01-09T12:00:00Z" + ] + }, + { + "expression": "0 12 */1 * 5", + "from": "2026-01-02T03:04:05.6789Z", + "name": "unit_step_keeps_wildcard", + "next": [ + "2026-01-02T12:00:00Z", + "2026-01-09T12:00:00Z", + "2026-01-16T12:00:00Z", + "2026-01-23T12:00:00Z", + "2026-01-30T12:00:00Z" + ] + }, + { + "expression": "5/15 * * * *", + "from": "2026-01-02T03:04:05.6789Z", + "name": "start_with_step", + "next": [ + "2026-01-02T03:05:00Z", + "2026-01-02T03:20:00Z", + "2026-01-02T03:35:00Z", + "2026-01-02T03:50:00Z", + "2026-01-02T04:05:00Z" + ] + }, + { + "expression": "0-10/5 * * * *", + "from": "2026-01-02T03:04:05.6789Z", + "name": "range_with_step", + "next": [ + "2026-01-02T03:05:00Z", + "2026-01-02T03:10:00Z", + "2026-01-02T04:00:00Z", + "2026-01-02T04:05:00Z", + "2026-01-02T04:10:00Z" + ] + }, + { + "expression": "59 23 31 12 *", + "from": "2026-01-02T03:04:05.6789Z", + "name": "year_end", + "next": [ + "2026-12-31T23:59:00Z", + "2027-12-31T23:59:00Z", + "2028-12-31T23:59:00Z", + "2029-12-31T23:59:00Z", + "2030-12-31T23:59:00Z" + ] + }, + { + "expression": "@hourly", + "from": "2026-01-02T03:04:05.6789Z", + "name": "descriptor_hourly", + "next": [ + "2026-01-02T04:00:00Z", + "2026-01-02T05:00:00Z", + "2026-01-02T06:00:00Z", + "2026-01-02T07:00:00Z", + "2026-01-02T08:00:00Z" + ] + }, + { + "expression": "@daily", + "from": "2026-01-02T03:04:05.6789Z", + "name": "descriptor_daily", + "next": [ + "2026-01-03T00:00:00Z", + "2026-01-04T00:00:00Z", + "2026-01-05T00:00:00Z", + "2026-01-06T00:00:00Z", + "2026-01-07T00:00:00Z" + ] + }, + { + "expression": "@midnight", + "from": "2026-01-02T03:04:05.6789Z", + "name": "descriptor_midnight", + "next": [ + "2026-01-03T00:00:00Z", + "2026-01-04T00:00:00Z", + "2026-01-05T00:00:00Z", + "2026-01-06T00:00:00Z", + "2026-01-07T00:00:00Z" + ] + }, + { + "expression": "@weekly", + "from": "2026-01-02T03:04:05.6789Z", + "name": "descriptor_weekly", + "next": [ + "2026-01-04T00:00:00Z", + "2026-01-11T00:00:00Z", + "2026-01-18T00:00:00Z", + "2026-01-25T00:00:00Z", + "2026-02-01T00:00:00Z" + ] + }, + { + "expression": "@monthly", + "from": "2026-01-02T03:04:05.6789Z", + "name": "descriptor_monthly", + "next": [ + "2026-02-01T00:00:00Z", + "2026-03-01T00:00:00Z", + "2026-04-01T00:00:00Z", + "2026-05-01T00:00:00Z", + "2026-06-01T00:00:00Z" + ] + }, + { + "expression": "@yearly", + "from": "2026-01-02T03:04:05.6789Z", + "name": "descriptor_yearly", + "next": [ + "2027-01-01T00:00:00Z", + "2028-01-01T00:00:00Z", + "2029-01-01T00:00:00Z", + "2030-01-01T00:00:00Z", + "2031-01-01T00:00:00Z" + ] + }, + { + "expression": "@annually", + "from": "2026-01-02T03:04:05.6789Z", + "name": "descriptor_annually", + "next": [ + "2027-01-01T00:00:00Z", + "2028-01-01T00:00:00Z", + "2029-01-01T00:00:00Z", + "2030-01-01T00:00:00Z", + "2031-01-01T00:00:00Z" + ] + }, + { + "expression": "@every 1h30m", + "from": "2026-01-02T03:04:05.6789Z", + "name": "every_compound_duration", + "next": [ + "2026-01-02T04:34:05Z", + "2026-01-02T06:04:05Z", + "2026-01-02T07:34:05Z", + "2026-01-02T09:04:05Z", + "2026-01-02T10:34:05Z" + ] + }, + { + "expression": "@every 1.5h", + "from": "2026-01-02T03:04:05.6789Z", + "name": "every_fractional_duration", + "next": [ + "2026-01-02T04:34:05Z", + "2026-01-02T06:04:05Z", + "2026-01-02T07:34:05Z", + "2026-01-02T09:04:05Z", + "2026-01-02T10:34:05Z" + ] + }, + { + "expression": "@every 90s", + "from": "2026-01-02T03:04:05.6789Z", + "name": "every_seconds", + "next": [ + "2026-01-02T03:05:35Z", + "2026-01-02T03:07:05Z", + "2026-01-02T03:08:35Z", + "2026-01-02T03:10:05Z", + "2026-01-02T03:11:35Z" + ] + }, + { + "expression": "@every 500ms", + "from": "2026-01-02T03:04:05.6789Z", + "name": "every_rounds_up_to_one_second", + "next": [ + "2026-01-02T03:04:06Z", + "2026-01-02T03:04:07Z", + "2026-01-02T03:04:08Z", + "2026-01-02T03:04:09Z", + "2026-01-02T03:04:10Z" + ] + }, + { + "expression": "@every 1500ms", + "from": "2026-01-02T03:04:05.6789Z", + "name": "every_truncates_subseconds", + "next": [ + "2026-01-02T03:04:06Z", + "2026-01-02T03:04:07Z", + "2026-01-02T03:04:08Z", + "2026-01-02T03:04:09Z", + "2026-01-02T03:04:10Z" + ] + }, + { + "expression": "0 9 * * *", + "from": "2026-03-07T08:00:00-05:00", + "name": "reference_time_offset", + "next": [ + "2026-03-07T09:00:00-05:00", + "2026-03-08T09:00:00-05:00", + "2026-03-09T09:00:00-05:00", + "2026-03-10T09:00:00-05:00", + "2026-03-11T09:00:00-05:00" + ] + }, + { + "expression": "30 0 * * *", + "from": "2026-03-07T23:45:00+05:30", + "name": "reference_time_half_hour_offset", + "next": [ + "2026-03-08T00:30:00+05:30", + "2026-03-09T00:30:00+05:30", + "2026-03-10T00:30:00+05:30", + "2026-03-11T00:30:00+05:30", + "2026-03-12T00:30:00+05:30" + ] + }, + { + "expression": "CRON_TZ=UTC 0 9 * * *", + "from": "2026-03-07T08:00:00-05:00", + "name": "cron_tz_utc_prefix", + "next": [ + "2026-03-08T04:00:00-05:00", + "2026-03-09T04:00:00-05:00", + "2026-03-10T04:00:00-05:00", + "2026-03-11T04:00:00-05:00", + "2026-03-12T04:00:00-05:00" + ] + }, + { + "expression": "TZ=UTC 0 9 * * *", + "from": "2026-03-07T08:00:00-05:00", + "name": "tz_utc_prefix", + "next": [ + "2026-03-08T04:00:00-05:00", + "2026-03-09T04:00:00-05:00", + "2026-03-10T04:00:00-05:00", + "2026-03-11T04:00:00-05:00", + "2026-03-12T04:00:00-05:00" + ] + }, + { + "expression": " 0 9 * * 1 ", + "from": "2026-01-02T03:04:05.6789Z", + "name": "extra_whitespace", + "next": [ + "2026-01-05T09:00:00Z", + "2026-01-12T09:00:00Z", + "2026-01-19T09:00:00Z", + "2026-01-26T09:00:00Z", + "2026-02-02T09:00:00Z" + ] + } + ], + "cron_invalid": [ + "", + "* * * *", + "* * * * * *", + "0 9 * * 7", + "60 * * * *", + "* 24 * * *", + "* * 0 * *", + "* * 32 * *", + "* * * 0 *", + "* * * 13 *", + "-1 * * * *", + "5-1 * * * *", + "1-2-3 * * * *", + "1/2/3 * * * *", + "*/0 * * * *", + "*/x * * * *", + "0 9 * * funday", + "@every", + "@every 5x", + "@reboot", + "CRON_TZ=Nowhere/Invalid 0 9 * * *" + ], + "cron_named_zone_cases": [ + { + "expression": "CRON_TZ=America/New_York 0 9 * * *", + "from": "2026-03-06T12:00:00Z", + "name": "new_york_across_dst_start", + "next": [ + "2026-03-06T14:00:00Z", + "2026-03-07T14:00:00Z", + "2026-03-08T13:00:00Z", + "2026-03-09T13:00:00Z", + "2026-03-10T13:00:00Z" + ] + }, + { + "expression": "CRON_TZ=America/New_York 30 2 * * *", + "from": "2026-03-06T12:00:00Z", + "name": "new_york_skipped_wall_time", + "next": [ + "2026-03-07T07:30:00Z", + "2026-03-09T06:30:00Z", + "2026-03-10T06:30:00Z", + "2026-03-11T06:30:00Z", + "2026-03-12T06:30:00Z" + ] + }, + { + "expression": "CRON_TZ=America/New_York 30 1 * * *", + "from": "2026-10-30T12:00:00Z", + "name": "new_york_repeated_wall_time", + "next": [ + "2026-10-31T05:30:00Z", + "2026-11-01T05:30:00Z", + "2026-11-01T06:30:00Z", + "2026-11-02T06:30:00Z", + "2026-11-03T06:30:00Z" + ] + }, + { + "expression": "CRON_TZ=America/New_York 0 * * * *", + "from": "2026-11-01T04:30:00Z", + "name": "new_york_hourly_across_dst_end", + "next": [ + "2026-11-01T05:00:00Z", + "2026-11-01T06:00:00Z", + "2026-11-01T07:00:00Z", + "2026-11-01T08:00:00Z", + "2026-11-01T09:00:00Z" + ] + }, + { + "expression": "CRON_TZ=Europe/London 0 0 * * *", + "from": "2026-10-23T12:00:00Z", + "name": "london_across_dst_end", + "next": [ + "2026-10-23T23:00:00Z", + "2026-10-24T23:00:00Z", + "2026-10-26T00:00:00Z", + "2026-10-27T00:00:00Z", + "2026-10-28T00:00:00Z" + ] + }, + { + "expression": "CRON_TZ=America/Santiago 0 0 * * *", + "from": "2026-09-03T12:00:00Z", + "name": "santiago_skipped_midnight", + "next": [ + "2026-09-04T04:00:00Z", + "2026-09-05T04:00:00Z", + "2026-09-07T03:00:00Z", + "2026-09-08T03:00:00Z", + "2026-09-09T03:00:00Z" + ] + }, + { + "expression": "CRON_TZ=America/Santiago 0 12 * * *", + "from": "2026-09-03T12:00:00Z", + "name": "santiago_day_after_skipped_midnight", + "next": [ + "2026-09-03T16:00:00Z", + "2026-09-04T16:00:00Z", + "2026-09-05T16:00:00Z", + "2026-09-06T15:00:00Z", + "2026-09-07T15:00:00Z" + ] + }, + { + "expression": "CRON_TZ=America/Santiago 30 23 * * *", + "from": "2026-04-02T12:00:00Z", + "name": "santiago_repeated_hour_before_midnight", + "next": [ + "2026-04-03T02:30:00Z", + "2026-04-04T02:30:00Z", + "2026-04-05T02:30:00Z", + "2026-04-05T03:30:00Z", + "2026-04-06T03:30:00Z" + ] + }, + { + "expression": "TZ=Asia/Kolkata 0 9 * * mon", + "from": "2026-01-02T03:04:05-05:00", + "name": "kolkata_tz_prefix", + "next": [ + "2026-01-04T22:30:00-05:00", + "2026-01-11T22:30:00-05:00", + "2026-01-18T22:30:00-05:00", + "2026-01-25T22:30:00-05:00", + "2026-02-01T22:30:00-05:00" + ] + } + ], + "snooze_counters": [ + { + "expected_snoozes": 1, + "metadata": {}, + "name": "absent" + }, + { + "expected_snoozes": 3, + "metadata": { + "snoozes": 2 + }, + "name": "integer" + }, + { + "expected_snoozes": 3, + "metadata": { + "snoozes": 2.9 + }, + "name": "fraction_truncates" + }, + { + "expected_snoozes": -1, + "metadata": { + "snoozes": -2.5 + }, + "name": "negative_fraction_truncates_toward_zero" + }, + { + "expected_snoozes": 1001, + "metadata": { + "snoozes": 1e3 + }, + "name": "exponent" + }, + { + "expected_snoozes": 9007199254740994, + "metadata": { + "snoozes": 9007199254740993 + }, + "name": "beyond_float_precision" + }, + { + "expected_snoozes": 5, + "metadata": { + "snoozes": "4" + }, + "name": "numeric_string" + }, + { + "expected_snoozes": -6, + "metadata": { + "snoozes": "-7" + }, + "name": "negative_numeric_string" + }, + { + "expected_snoozes": 1, + "metadata": { + "snoozes": "4.5" + }, + "name": "fractional_string_is_zero" + }, + { + "expected_snoozes": 1, + "metadata": { + "snoozes": " 5" + }, + "name": "padded_string_is_zero" + }, + { + "expected_snoozes": 1, + "metadata": { + "snoozes": "abc" + }, + "name": "non_numeric_string_is_zero" + }, + { + "expected_snoozes": 2, + "metadata": { + "snoozes": true + }, + "name": "true_is_one" + }, + { + "expected_snoozes": 1, + "metadata": { + "snoozes": false + }, + "name": "false_is_zero" + }, + { + "expected_snoozes": 1, + "metadata": { + "snoozes": null + }, + "name": "null_is_zero" + }, + { + "expected_snoozes": 1, + "metadata": { + "snoozes": [ + 3 + ] + }, + "name": "array_is_zero" + }, + { + "expected_snoozes": 1, + "metadata": { + "snoozes": { + "count": 3 + } + }, + "name": "object_is_zero" + } + ] +} diff --git a/rust/riverqueue/tests/fixtures/protocol_values.json b/rust/riverqueue/tests/fixtures/protocol_values.json new file mode 100644 index 000000000..15fb9a779 --- /dev/null +++ b/rust/riverqueue/tests/fixtures/protocol_values.json @@ -0,0 +1,317 @@ +{ + "attempt_error": { + "at": "2026-01-02T03:04:05.6789Z", + "attempt": 3, + "error": "worker failed: escaped \"detail\"", + "trace": "frame one\nframe two" + }, + "job_states": [ + { + "state": "available", + "unique_bit": 1 + }, + { + "state": "cancelled", + "unique_bit": 2 + }, + { + "state": "completed", + "unique_bit": 4 + }, + { + "state": "discarded", + "unique_bit": 8 + }, + { + "state": "pending", + "unique_bit": 16 + }, + { + "state": "retryable", + "unique_bit": 32 + }, + { + "state": "running", + "unique_bit": 64 + }, + { + "state": "scheduled", + "unique_bit": 128 + } + ], + "metadata_keys": { + "output": "output", + "periodic_job_id": "river:periodic_job_id", + "rescue_count": "river:rescue_count", + "resumable_cursor": "river:resumable_cursor", + "resumable_step": "river:resumable_step", + "unique_nonce": "river:unique_nonce" + }, + "notifications": [ + { + "fields": [ + { + "name": "action", + "omitempty": false + }, + { + "name": "job_id", + "omitempty": true + }, + { + "name": "metadata", + "omitempty": true + }, + { + "name": "queue", + "omitempty": false + } + ], + "name": "cancel", + "payload": { + "action": "cancel", + "job_id": 42, + "queue": "priority" + }, + "source": "producer.go:controlEventPayload; riverdriver/riverpgxv5/internal/dbsqlc/river_job.sql:JobCancel", + "topic": "river_control" + }, + { + "fields": [ + { + "name": "queue", + "omitempty": false + } + ], + "name": "insert", + "payload": { + "queue": "priority" + }, + "source": "producer.go:insertPayload", + "topic": "river_insert" + }, + { + "fields": [ + { + "name": "action", + "omitempty": false + }, + { + "name": "job_id", + "omitempty": true + }, + { + "name": "metadata", + "omitempty": true + }, + { + "name": "queue", + "omitempty": false + } + ], + "name": "metadata_changed", + "payload": { + "action": "metadata_changed", + "metadata": { + "owner": "candidate" + }, + "queue": "priority" + }, + "source": "producer.go:controlEventPayload", + "topic": "river_control" + }, + { + "fields": [ + { + "name": "action", + "omitempty": false + }, + { + "name": "job_id", + "omitempty": true + }, + { + "name": "metadata", + "omitempty": true + }, + { + "name": "queue", + "omitempty": false + } + ], + "name": "pause", + "payload": { + "action": "pause", + "queue": "priority" + }, + "source": "producer.go:controlEventPayload", + "topic": "river_control" + }, + { + "fields": [ + { + "name": "action", + "omitempty": false + }, + { + "name": "leader_id", + "omitempty": false + } + ], + "name": "request_resign", + "payload": { + "action": "request_resign", + "leader_id": "" + }, + "source": "internal/leadership/elector.go:DBNotification", + "topic": "river_leadership" + }, + { + "fields": [ + { + "name": "action", + "omitempty": false + }, + { + "name": "leader_id", + "omitempty": false + } + ], + "name": "resigned", + "payload": { + "action": "resigned", + "leader_id": "client-1" + }, + "source": "internal/leadership/elector.go:DBNotification; riverdriver/riverpgxv5/internal/dbsqlc/river_leader.sql:LeaderResign", + "topic": "river_leadership" + }, + { + "fields": [ + { + "name": "action", + "omitempty": false + }, + { + "name": "job_id", + "omitempty": true + }, + { + "name": "metadata", + "omitempty": true + }, + { + "name": "queue", + "omitempty": false + } + ], + "name": "resume", + "payload": { + "action": "resume", + "queue": "priority" + }, + "source": "producer.go:controlEventPayload", + "topic": "river_control" + } + ], + "reserved_metadata_keys": [ + { + "applicability": "protocol_visible", + "key": "cancel_attempted_at" + }, + { + "applicability": "protocol_visible", + "key": "output" + }, + { + "applicability": "protocol_visible", + "key": "periodic" + }, + { + "applicability": "protocol_visible", + "key": "river:log" + }, + { + "applicability": "protocol_visible", + "key": "river:periodic_job_id" + }, + { + "applicability": "protocol_visible", + "key": "river:rescue_count" + }, + { + "applicability": "protocol_visible", + "key": "river:resumable_cursor" + }, + { + "applicability": "protocol_visible", + "key": "river:resumable_step" + }, + { + "applicability": "protocol_visible", + "key": "river:unique_nonce" + }, + { + "applicability": "protocol_visible", + "key": "snoozes" + }, + { + "applicability": "protocol_visible", + "key": "unique_key_conflict" + } + ], + "retry_cases": [ + { + "error_count": 1, + "job_id": 42, + "max_delay_ns": 1100000000, + "min_delay_ns": 900000000, + "now": "2026-01-02T03:04:05.6789Z", + "seed": 0 + }, + { + "error_count": 2, + "job_id": 42, + "max_delay_ns": 17600000000, + "min_delay_ns": 14400000000, + "now": "2026-01-02T03:04:05.6789Z", + "seed": 123 + }, + { + "error_count": 3, + "job_id": 9007199254740991, + "max_delay_ns": 89100000000, + "min_delay_ns": 72900000000, + "now": "2026-01-02T03:04:05.6789Z", + "seed": 18446744073709551615 + }, + { + "error_count": 11, + "job_id": 1, + "max_delay_ns": 16105100000000, + "min_delay_ns": 13176900000000, + "now": "2026-01-02T03:04:05.6789Z", + "seed": 456 + }, + { + "error_count": 309, + "job_id": 42, + "max_delay_ns": 9223372036854775807, + "min_delay_ns": 8204959224899999744, + "now": "2026-01-02T03:04:05.6789Z", + "seed": 789 + }, + { + "error_count": 310, + "job_id": 42, + "max_delay_ns": 9223372036854775807, + "min_delay_ns": 9223372036854775807, + "now": "2026-01-02T03:04:05.6789Z", + "seed": 123 + } + ], + "topics": { + "control": "river_control", + "insert": "river_insert", + "leadership": "river_leadership" + } +} diff --git a/rust/riverqueue/tests/fixtures/unique_keys.json b/rust/riverqueue/tests/fixtures/unique_keys.json new file mode 100644 index 000000000..066f4d86e --- /dev/null +++ b/rust/riverqueue/tests/fixtures/unique_keys.json @@ -0,0 +1,948 @@ +{ + "cases": [ + { + "args": {}, + "expected_sha256": "23aa86692d9807ab10e433e378f1c0804573f5e345818461b919322dd381b4c3", + "expected_state_mask": 245, + "kind": "conformance_selected_args", + "name": "all_selected_fields_omitted", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_components": [ + [ + "account", + "id" + ], + [ + "account", + "region" + ], + [ + "label" + ], + [ + "path/key" + ] + ], + "selected_unique_paths": [ + "account.id", + "account.region", + "label", + "path/key" + ] + }, + { + "args": { + "account": { + "id": "acct", + "ignored": "irrelevant", + "region": "west" + }, + "path/key": "slash" + }, + "expected_sha256": "7d62e81ac25cfa2dec69ad5a41e0b78188ee1b299bed329b453da6b3abca70bd", + "expected_state_mask": 245, + "kind": "conformance_selected_args", + "name": "selected_siblings_and_slash_key", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_components": [ + [ + "account", + "id" + ], + [ + "account", + "region" + ], + [ + "label" + ], + [ + "path/key" + ] + ], + "selected_unique_paths": [ + "account.id", + "account.region", + "label", + "path/key" + ] + }, + { + "args": { + "nested": { + "z": 1, + "a": 2 + } + }, + "expected_sha256": "46ff499cb031d0458bb00ef87e7b83321eefb6e7534497c279c4ae0f474cdde0", + "expected_state_mask": 245, + "kind": "conformance_all_args", + "name": "nested_struct_wire_order", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": { + "zeta": "quoted \\\"value\\\" and \\\\ slash", + "alpha": "\u003calpha\u003e\u0026\u2028line", + "maximum": 9007199254740991 + }, + "expected_sha256": "7a84c62c8d470ca388a0a1e41c311b9eb1ea21f7b88157ceb876fc82e698b6af", + "expected_state_mask": 245, + "kind": "conformance_all_args", + "name": "all_args_sorted_and_escaped", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": { + "2": 2, + "10": 10, + "zero": -0, + "😀": 1, + "": 2 + }, + "expected_sha256": "fcdf33e0c39c1fc7e956876345a985f2418bd69c6e4d6a5c794abf1e78cdfdb6", + "expected_state_mask": 245, + "kind": "conformance_all_args", + "name": "map_order_and_negative_zero", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": { + "": 0, + "a.b": 1, + "@x": 2, + ":lead": 3, + "!bang": 4, + "[open": 5, + "{brace": 6, + "a\\b": 7 + }, + "expected_sha256": "1d254dda1efe1009ffb205ede791d481545d66abb8542e85f0d895415b05cdf9", + "expected_state_mask": 245, + "kind": "conformance_all_args", + "name": "all_args_literal_path_syntax", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": { + "a\"b": 1, + "line\n": 2, + "é": 3, + "a\u003cb": 4, + "a\u0026b": 5, + "a\u2028b": 6 + }, + "expected_sha256": "bab84635792449d758b18950f6f014bfb16d26430e798d84afee9fc9cb793163", + "expected_state_mask": 245, + "kind": "conformance_all_args", + "name": "all_args_escaped_key_encoding", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": [], + "expected_sha256": "fe05a58ddb79a8d4544da962582d9a290d59788c920afd3597da3a62e3c1b0ac", + "expected_state_mask": 245, + "kind": "conformance_all_args", + "name": "all_args_empty_array", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": [ + 1 + ], + "expected_error": "rejected", + "expected_state_mask": 245, + "kind": "conformance_all_args", + "name": "all_args_array_rejected", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": null, + "expected_error": "rejected", + "expected_state_mask": 245, + "kind": "conformance_all_args", + "name": "all_args_null_rejected", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": "args", + "expected_error": "rejected", + "expected_state_mask": 245, + "kind": "conformance_all_args", + "name": "all_args_scalar_rejected", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": { + "exponent": 1e+100, + "fraction": 1.25, + "maximum": 9223372036854775807, + "minimum": -9223372036854775808, + "unsigned_maximum": 18446744073709551615 + }, + "expected_sha256": "2c1533b3ab43068407d14e82ddb34a295a51375ae3a27fef6931123f07677f38", + "expected_state_mask": 245, + "kind": "conformance_numeric_boundaries", + "name": "numeric_boundaries", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": { + "account": { + "id": "acct-123", + "ignored": "not selected" + }, + "ignored": true, + "label": "selected" + }, + "expected_sha256": "6130dc4f753402d1faeb6bbc3e6c21415245bb282ad1fd16bcbfeebde525e726", + "expected_state_mask": 245, + "kind": "conformance_selected_args", + "name": "selected_nested_args", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_components": [ + [ + "account", + "id" + ], + [ + "account", + "region" + ], + [ + "label" + ], + [ + "path/key" + ] + ], + "selected_unique_paths": [ + "account.id", + "account.region", + "label", + "path/key" + ] + }, + { + "args": { + "user.id": "literal", + "user": {} + }, + "expected_sha256": "7d478fa6978b3fbb5c326d90fd10c2eab663cf2c9c2e1e3ddcce032989a05cdb", + "expected_state_mask": 245, + "kind": "conformance_dotted_selected_args", + "name": "selected_literal_dotted_name", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_components": [ + [ + "@user" + ], + [ + "!x" + ], + [ + "{x}" + ], + [ + "[x]" + ], + [ + ":id" + ], + [ + "user", + "id" + ], + [ + "user.id" + ], + [ + "a*b?c#d|e" + ], + [ + "é" + ] + ], + "selected_unique_paths": [ + "\\@user", + "\\!x", + "\\{x\\}", + "\\[x\\]", + "\\:id", + "user.id", + "user\\.id", + "a\\*b\\?c\\#d\\|e", + "é" + ] + }, + { + "args": { + "user": { + "id": "nested" + } + }, + "expected_sha256": "6fd34aa5a46274e4f7d159063be43c2c02e90438b84dab208c8c5209f94dda25", + "expected_state_mask": 245, + "kind": "conformance_dotted_selected_args", + "name": "selected_nested_dotted_path", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_components": [ + [ + "@user" + ], + [ + "!x" + ], + [ + "{x}" + ], + [ + "[x]" + ], + [ + ":id" + ], + [ + "user", + "id" + ], + [ + "user.id" + ], + [ + "a*b?c#d|e" + ], + [ + "é" + ] + ], + "selected_unique_paths": [ + "\\@user", + "\\!x", + "\\{x\\}", + "\\[x\\]", + "\\:id", + "user.id", + "user\\.id", + "a\\*b\\?c\\#d\\|e", + "é" + ] + }, + { + "args": { + "user": {}, + "é": "café" + }, + "expected_sha256": "28513f484784e6b0fe8aed6cc1fadb04498f43305b74619aa56e701a2feff578", + "expected_state_mask": 245, + "kind": "conformance_dotted_selected_args", + "name": "selected_unicode_field_name", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_components": [ + [ + "@user" + ], + [ + "!x" + ], + [ + "{x}" + ], + [ + "[x]" + ], + [ + ":id" + ], + [ + "user", + "id" + ], + [ + "user.id" + ], + [ + "a*b?c#d|e" + ], + [ + "é" + ] + ], + "selected_unique_paths": [ + "\\@user", + "\\!x", + "\\{x\\}", + "\\[x\\]", + "\\:id", + "user.id", + "user\\.id", + "a\\*b\\?c\\#d\\|e", + "é" + ] + }, + { + "args": { + "@user": "at", + "!x": "bang", + "{x}": "brace", + "[x]": "bracket", + ":id": "colon", + "a*b?c#d|e": "symbols", + "user": {} + }, + "expected_sha256": "d00ff085218024d7059a4556b24af92cbe57744935e2d37ff644b149fae4c2f3", + "expected_state_mask": 245, + "kind": "conformance_dotted_selected_args", + "name": "selected_punctuation_field_names", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_components": [ + [ + "@user" + ], + [ + "!x" + ], + [ + "{x}" + ], + [ + "[x]" + ], + [ + ":id" + ], + [ + "user", + "id" + ], + [ + "user.id" + ], + [ + "a*b?c#d|e" + ], + [ + "é" + ] + ], + "selected_unique_paths": [ + "\\@user", + "\\!x", + "\\{x\\}", + "\\[x\\]", + "\\:id", + "user.id", + "user\\.id", + "a\\*b\\?c\\#d\\|e", + "é" + ] + }, + { + "args": { + "empty": [], + "labels": { + "alpha": "first", + "k10": "ten", + "k2": "two", + "zulu": "last" + }, + "matrix": [ + [ + 3, + 1 + ], + [], + [ + 2 + ] + ], + "missing": null, + "objects": [ + { + "zulu": "z", + "alpha": 1 + }, + { + "zulu": "y", + "alpha": null + } + ], + "pointer": null + }, + "expected_sha256": "66d457888b4b71f0a0041251f75d494e1ac543717283d84633a274648559306d", + "expected_state_mask": 245, + "kind": "conformance_all_args", + "name": "typed_collections_and_nulls", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": {}, + "expected_sha256": "fe05a58ddb79a8d4544da962582d9a290d59788c920afd3597da3a62e3c1b0ac", + "expected_state_mask": 245, + "kind": "conformance_all_args", + "name": "typed_empty_args", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": { + "a\u003cb\u003e": "\u003cangle\u003e", + "controls": "\b\f\n\r\t\u0000\u0001\u001f", + "html": "\u003ca href=\"x\"\u003e\u0026amp;\u003c/a\u003e", + "keys": { + "\u003ck\u003e": 1, + "a\u0026b": 2, + "é": 3, + "é\u003c": 4 + }, + "separators": "line\u2028paragraph\u2029end", + "unicode": "é😀/\\", + "é\u0026": "unicode key" + }, + "expected_sha256": "a4c2a164225cba2e3ae56edbc11258e49a7325d73a9b70e5c9b47a51963d00c5", + "expected_state_mask": 245, + "kind": "conformance_all_args", + "name": "typed_escaping", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": { + "label": null + }, + "expected_sha256": "d137d7c4f1e3f8369037b1890357655b4fa1978a329bcf51435f9f0b54abcab7", + "expected_state_mask": 245, + "kind": "conformance_selected_args", + "name": "selected_explicit_null", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_components": [ + [ + "account", + "id" + ], + [ + "account", + "region" + ], + [ + "label" + ], + [ + "path/key" + ] + ], + "selected_unique_paths": [ + "account.id", + "account.region", + "label", + "path/key" + ] + }, + { + "args": { + "fraction": "2026-01-02T03:04:05.5Z", + "micros": "2026-01-02T03:04:05.123456Z", + "millis": "2026-01-02T03:04:05.12Z", + "whole": "2026-01-02T03:04:05Z" + }, + "expected_sha256": "18b38780de3019cc75d49ff24a74eecdcff99fda0bb36c0244a35f262b15322b", + "expected_state_mask": 245, + "kind": "conformance_all_args", + "name": "typed_time_values", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": { + "below_large": 999999999999999900000, + "large": 100000000000000000000, + "large_boundary": 1e+21, + "largest": 1.7976931348623157e+308, + "negative": -1.5e-9, + "negative_zero": -0, + "one": 1, + "single": 1.1, + "single_large": 1e+21, + "single_small": 1e-7, + "small": 1e-7, + "small_boundary": 0.000001, + "smallest": 5e-324, + "tenth": 0.1 + }, + "expected_sha256": "dc330477ebe8bf2bc3402476365fcfe22c82253f744481c49a23f760f059cad4", + "expected_state_mask": 245, + "kind": "conformance_all_args", + "name": "typed_float_formatting", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": { + "id": 42 + }, + "expected_sha256": "5396f06a082abd7a929915135ebd363a9a47d800176b03ce7736f93a5ba9e22e", + "expected_state_mask": 245, + "kind": "conformance_simple", + "name": "period_from_now", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": false, + "by_period_nanos": 5400000000000, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": { + "id": 42 + }, + "expected_sha256": "b7f3c49952996b760b8b3ff6cf48f426e03a6ef0f004fb6faa51725365cf309a", + "expected_state_mask": 245, + "kind": "conformance_simple", + "name": "period_from_schedule", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": false, + "by_period_nanos": 3600000000000, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": "2026-01-02T05:21:05.6789Z", + "selected_unique_paths": null + }, + { + "args": { + "id": 42 + }, + "expected_sha256": "5396f06a082abd7a929915135ebd363a9a47d800176b03ce7736f93a5ba9e22e", + "expected_state_mask": 245, + "kind": "conformance_simple", + "name": "period_from_non_utc_now", + "now": "2026-01-01T22:04:05.6789-05:00", + "options": { + "by_args": false, + "by_period_nanos": 3600000000000, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": { + "id": 42 + }, + "expected_sha256": "b7f3c49952996b760b8b3ff6cf48f426e03a6ef0f004fb6faa51725365cf309a", + "expected_state_mask": 245, + "kind": "conformance_simple", + "name": "period_from_non_utc_schedule", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": false, + "by_period_nanos": 3600000000000, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": "2026-01-02T10:51:05.6789+05:30", + "selected_unique_paths": null + }, + { + "args": { + "id": 42 + }, + "expected_sha256": "d20ce47da8e8015e68b020bbca2b17494139519ae05f8985e5992db4d8dd8a09", + "expected_state_mask": 245, + "kind": "conformance_simple", + "name": "queue_without_kind", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": false, + "by_period_nanos": 0, + "by_queue": true, + "exclude_kind": true + }, + "queue": "priority_emails", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": { + "id": 42 + }, + "expected_sha256": "6f20262e7b1fa9beaf98255f23030800376636484a98b20c2c04a303bab5a8d5", + "expected_state_mask": 213, + "kind": "conformance_simple", + "name": "all_dimensions_custom_states", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 60000000000, + "by_queue": true, + "by_state": [ + "available", + "completed", + "pending", + "running", + "scheduled" + ], + "exclude_kind": false + }, + "queue": "priority_emails", + "scheduled_at": "2026-01-02T05:21:05.6789Z", + "selected_unique_paths": null + } + ], + "typed_only_cases": [ + { + "args": { + "a": 1, + "a": 2, + "b": 3 + }, + "expected_sha256": "f3568e94e18a68ce633bd444fe9ec448382740bcf02b83b210cb68e74220fb5f", + "expected_state_mask": 245, + "kind": "conformance_all_args", + "name": "typed_duplicate_top_level_keys", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + }, + { + "args": { + "empty": [], + "labels": { + "10": "ten", + "2": "two", + "alpha": "first", + "zulu": "last" + }, + "matrix": [], + "missing": null, + "objects": [], + "pointer": null + }, + "expected_sha256": "38406019aea1ea67f81186d845e983b197b2d919ded5eca871e4fef0c8b0ad30", + "expected_state_mask": 245, + "kind": "conformance_all_args", + "name": "typed_integer_like_map_keys", + "now": "2026-01-02T03:04:05.6789Z", + "options": { + "by_args": true, + "by_period_nanos": 0, + "by_queue": false, + "exclude_kind": false + }, + "queue": "default", + "scheduled_at": null, + "selected_unique_paths": null + } + ] +} diff --git a/rust/riverqueue/tests/insert_notifications.rs b/rust/riverqueue/tests/insert_notifications.rs new file mode 100644 index 000000000..cc75475ce --- /dev/null +++ b/rust/riverqueue/tests/insert_notifications.rs @@ -0,0 +1,459 @@ +//! A client notifies a queue of inserted jobs at most once per fetch +//! cooldown, like River Go's insert notification limiter: every insertion +//! path shares one window per queue, a rolled-back transaction still starts +//! it, and each client keeps its own. +//! +//! PostgreSQL scenarios run in a unique schema and fail rather than skip when +//! `RIVER_RUST_DATABASE_URL` is unset; SQLite scenarios use temporary files. + +#![cfg(any(all(feature = "postgres", river_postgres_tests), feature = "sqlite"))] + +mod support; + +use std::time::Duration; + +use riverqueue::{Client, InsertBatch, InsertManyItem, InsertOpts, JobArgs}; +use serde::{Deserialize, Serialize}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "insert_notification")] +struct NotificationArgs {} + +fn queue(name: &str) -> InsertOpts { + InsertOpts::default().with_queue(name) +} + +fn unique() -> InsertOpts { + InsertOpts::default().with_unique(riverqueue::UniqueOpts::new().with_by_args(true)) +} + +fn batch(queues: &[&str]) -> InsertBatch { + let mut batch = InsertBatch::new(); + for name in queues { + batch.push_with(NotificationArgs {}, queue(name)); + } + batch +} + +fn many(queues: &[&str]) -> Vec> { + queues + .iter() + .map(|name| InsertManyItem::new(NotificationArgs {}, queue(name))) + .collect() +} + +/// The client's fetch cooldown is validated like Go's and is the default for +/// queues without their own. +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn fetch_cooldown_validates_against_queue_poll_intervals() { + use riverqueue::{Error, QueueConfig, WorkContext, WorkOutcome, WorkerRegistry}; + + let (pool, path) = support::sqlite_file_pool(1).await; + let workers = || { + let mut workers = WorkerRegistry::new(); + workers + .register_fn( + |_context: WorkContext, _job: riverqueue::Job| async { + Ok::<_, std::convert::Infallible>(WorkOutcome::Complete) + }, + ) + .unwrap(); + workers + }; + let builder = || Client::builder(pool.clone()).workers(workers()); + + assert_eq!(QueueConfig::new(1).fetch_cooldown(), None); + assert_eq!( + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(5)) + .fetch_cooldown(), + Some(Duration::from_millis(5)) + ); + for builder in [ + builder().fetch_cooldown(Duration::ZERO), + builder().fetch_cooldown(Duration::from_micros(999)), + // A queue's poll interval (one second by default) can't be + // shorter than the client's cooldown it inherits... + builder() + .fetch_cooldown(Duration::from_secs(2)) + .queue("default", QueueConfig::new(1)), + // ...or than its own. + builder().queue( + "default", + QueueConfig::new(1).with_fetch_cooldown(Duration::from_secs(2)), + ), + builder().queue( + "default", + QueueConfig::new(1).with_fetch_cooldown(Duration::ZERO), + ), + ] { + let error = builder.build().unwrap_err(); + assert!(matches!(error, Error::Configuration(_)), "{error}"); + } + + builder() + .fetch_cooldown(Duration::from_millis(1)) + .build() + .unwrap(); + builder() + .fetch_cooldown(Duration::from_secs(2)) + .queue( + "overridden", + QueueConfig::new(1).with_fetch_cooldown(Duration::from_millis(100)), + ) + .queue( + "slow", + QueueConfig::new(1).with_fetch_poll_interval(Duration::from_secs(2)), + ) + .build() + .unwrap(); + + // Queues added at runtime are checked against the client's cooldown too. + let client = builder() + .fetch_cooldown(Duration::from_secs(2)) + .build() + .unwrap(); + let error = client + .local_queues() + .add("default", QueueConfig::new(1)) + .unwrap_err(); + assert!(matches!(error, Error::Configuration(_)), "{error}"); + client + .local_queues() + .add( + "default", + QueueConfig::new(1).with_fetch_poll_interval(Duration::from_secs(2)), + ) + .unwrap(); + support::sqlite_cleanup(pool, path).await; +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +mod postgres { + use std::time::Duration; + + use riverqueue::Client; + use riverqueue::database::PostgresDatabase; + use sqlx::postgres::PgListener; + + use super::support::PostgresSchema; + use super::{NotificationArgs, batch, many, queue, unique}; + + /// Listens to a schema's insert channel. + struct Notifications { + channel: String, + listener: PgListener, + } + + impl Notifications { + async fn listen(schema: &PostgresSchema) -> Self { + let channel = format!("{}.river_insert", schema.schema.as_deref().unwrap()); + let mut listener = PgListener::connect_with(&schema.pool).await.unwrap(); + listener.listen(&channel).await.unwrap(); + Self { channel, listener } + } + + /// Returns the queues notified since the last call, in order. + async fn next(&mut self, schema: &PostgresSchema) -> Vec { + sqlx::query("SELECT pg_notify($1, 'marker')") + .bind(&self.channel) + .execute(&schema.pool) + .await + .unwrap(); + let mut queues = Vec::new(); + loop { + let notification = self.listener.recv().await.unwrap(); + if notification.payload() == "marker" { + return queues; + } + let payload: serde_json::Value = + serde_json::from_str(notification.payload()).unwrap(); + queues.push(payload["queue"].as_str().unwrap().to_owned()); + } + } + } + + fn client(schema: &PostgresSchema, cooldown: Duration) -> Client { + Client::builder( + PostgresDatabase::new(schema.pool.clone()).with_schema(schema.schema.clone()), + ) + .fetch_cooldown(cooldown) + .build() + .unwrap() + } + + #[tokio::test(flavor = "multi_thread")] + async fn insert_notifications_resume_after_the_fetch_cooldown() { + let schema = PostgresSchema::new("insert_notify_resume").await; + let mut notifications = Notifications::listen(&schema).await; + let client = client(&schema, Duration::from_millis(50)); + + // Suppression is covered separately with a long cooldown; database + // round trips can exceed this short window on a busy runner. + client.insert(NotificationArgs {}).await.unwrap(); + assert_eq!(notifications.next(&schema).await, ["default"]); + tokio::time::sleep(Duration::from_millis(60)).await; + client.insert(NotificationArgs {}).await.unwrap(); + assert_eq!(notifications.next(&schema).await, ["default"]); + // The listener holds a pool connection, which closing the pool awaits. + drop(notifications); + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn insert_notifications_wait_for_the_fetch_cooldown() { + let schema = PostgresSchema::new("insert_notify_cooldown").await; + let mut notifications = Notifications::listen(&schema).await; + let client = client(&schema, Duration::from_hours(1)); + + client + .insert(NotificationArgs {}) + .opts(queue("a")) + .await + .unwrap(); + assert_eq!(notifications.next(&schema).await, ["a"]); + + client + .insert(NotificationArgs {}) + .opts(queue("a")) + .await + .unwrap(); + client.insert_many(many(&["a", "b", "b"])).await.unwrap(); + client.insert_batch(batch(&["b", "c"])).await.unwrap(); + assert_eq!(notifications.next(&schema).await, ["b", "c"]); + + // Scheduled jobs send nothing and leave the queue's window alone. + client + .insert(NotificationArgs {}) + .opts(queue("d").with_scheduled_at(chrono::Utc::now() + chrono::Duration::hours(1))) + .await + .unwrap(); + assert_eq!(notifications.next(&schema).await, Vec::::new()); + + let mut transaction = schema.pool.begin().await.unwrap(); + client + .insert(NotificationArgs {}) + .opts(queue("d")) + .tx(&mut transaction) + .await + .unwrap(); + client + .insert_many(many(&["d", "e"])) + .tx(&mut transaction) + .await + .unwrap(); + transaction.commit().await.unwrap(); + assert_eq!(notifications.next(&schema).await, ["d", "e"]); + + // A rolled-back transaction delivers nothing but still starts its + // queue's window. + let mut transaction = schema.pool.begin().await.unwrap(); + client + .insert_batch(batch(&["f"])) + .tx(&mut transaction) + .await + .unwrap(); + transaction.rollback().await.unwrap(); + client + .insert(NotificationArgs {}) + .opts(queue("f")) + .await + .unwrap(); + assert_eq!(notifications.next(&schema).await, Vec::::new()); + + // Each client has its own windows. + let other = self::client(&schema, Duration::from_hours(1)); + other + .insert(NotificationArgs {}) + .opts(queue("a")) + .await + .unwrap(); + assert_eq!(notifications.next(&schema).await, ["a"]); + // The listener holds a pool connection, which closing the pool awaits. + drop(notifications); + schema.cleanup().await; + } + + /// Like Go, a job skipped as a unique duplicate still notifies its queue. + #[tokio::test(flavor = "multi_thread")] + async fn unique_duplicates_notify_their_queue() { + let schema = PostgresSchema::new("insert_notify_duplicate").await; + let mut notifications = Notifications::listen(&schema).await; + let client = client(&schema, Duration::from_millis(1)); + + let first = client + .insert(NotificationArgs {}) + .opts(unique()) + .await + .unwrap(); + assert!(!first.unique_skipped_as_duplicate); + assert_eq!(notifications.next(&schema).await, ["default"]); + tokio::time::sleep(Duration::from_millis(2)).await; + let duplicate = client + .insert(NotificationArgs {}) + .opts(unique()) + .await + .unwrap(); + assert!(duplicate.unique_skipped_as_duplicate); + assert_eq!(notifications.next(&schema).await, ["default"]); + // The listener holds a pool connection, which closing the pool awaits. + drop(notifications); + schema.cleanup().await; + } +} + +#[cfg(feature = "sqlite")] +mod sqlite { + use std::time::Duration; + + use riverqueue::Client; + use sqlx::SqlitePool; + + use super::support::{sqlite_cleanup, sqlite_file_pool}; + use super::{NotificationArgs, batch, many, queue, unique}; + + /// Reads the insert notifications written to the outbox. + #[derive(Default)] + struct Notifications { + after_id: i64, + } + + impl Notifications { + /// Returns the queues notified since the last call, in order. + async fn next(&mut self, pool: &SqlitePool) -> Vec { + let rows: Vec<(i64, String)> = sqlx::query_as( + "SELECT id, json_extract(payload, '$.queue') FROM river_notification \ + WHERE topic = 'river_insert' AND id > ? ORDER BY id", + ) + .bind(self.after_id) + .fetch_all(pool) + .await + .unwrap(); + if let Some((id, _)) = rows.last() { + self.after_id = *id; + } + rows.into_iter().map(|(_, queue)| queue).collect() + } + } + + fn client(pool: &SqlitePool, cooldown: Duration) -> Client { + Client::builder(pool.clone()) + .fetch_cooldown(cooldown) + .build() + .unwrap() + } + + #[tokio::test(flavor = "multi_thread")] + async fn insert_notifications_resume_after_the_fetch_cooldown() { + let (pool, path) = sqlite_file_pool(4).await; + let mut notifications = Notifications::default(); + let client = client(&pool, Duration::from_millis(50)); + + // Suppression is covered separately with a long cooldown; database + // round trips can exceed this short window on a busy runner. + client.insert(NotificationArgs {}).await.unwrap(); + assert_eq!(notifications.next(&pool).await, ["default"]); + tokio::time::sleep(Duration::from_millis(60)).await; + client.insert(NotificationArgs {}).await.unwrap(); + assert_eq!(notifications.next(&pool).await, ["default"]); + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn insert_notifications_wait_for_the_fetch_cooldown() { + let (pool, path) = sqlite_file_pool(4).await; + let mut notifications = Notifications::default(); + let client = client(&pool, Duration::from_hours(1)); + + client + .insert(NotificationArgs {}) + .opts(queue("a")) + .await + .unwrap(); + assert_eq!(notifications.next(&pool).await, ["a"]); + + client + .insert(NotificationArgs {}) + .opts(queue("a")) + .await + .unwrap(); + client.insert_many(many(&["a", "b", "b"])).await.unwrap(); + client.insert_batch(batch(&["b", "c"])).await.unwrap(); + assert_eq!(notifications.next(&pool).await, ["b", "c"]); + + // Scheduled jobs send nothing and leave the queue's window alone. + client + .insert(NotificationArgs {}) + .opts(queue("d").with_scheduled_at(chrono::Utc::now() + chrono::Duration::hours(1))) + .await + .unwrap(); + assert_eq!(notifications.next(&pool).await, Vec::::new()); + + let mut transaction = pool.begin_with("BEGIN IMMEDIATE").await.unwrap(); + client + .insert(NotificationArgs {}) + .opts(queue("d")) + .tx(&mut transaction) + .await + .unwrap(); + client + .insert_many(many(&["d", "e"])) + .tx(&mut transaction) + .await + .unwrap(); + transaction.commit().await.unwrap(); + assert_eq!(notifications.next(&pool).await, ["d", "e"]); + + // A rolled-back transaction delivers nothing but still starts its + // queue's window. + let mut transaction = pool.begin_with("BEGIN IMMEDIATE").await.unwrap(); + client + .insert_batch(batch(&["f"])) + .tx(&mut transaction) + .await + .unwrap(); + transaction.rollback().await.unwrap(); + client + .insert(NotificationArgs {}) + .opts(queue("f")) + .await + .unwrap(); + assert_eq!(notifications.next(&pool).await, Vec::::new()); + + // Each client has its own windows. + let other = self::client(&pool, Duration::from_hours(1)); + other + .insert(NotificationArgs {}) + .opts(queue("a")) + .await + .unwrap(); + assert_eq!(notifications.next(&pool).await, ["a"]); + sqlite_cleanup(pool, path).await; + } + + /// Like Go, a job skipped as a unique duplicate still notifies its queue. + #[tokio::test(flavor = "multi_thread")] + async fn unique_duplicates_notify_their_queue() { + let (pool, path) = sqlite_file_pool(4).await; + let mut notifications = Notifications::default(); + let client = client(&pool, Duration::from_millis(1)); + + let first = client + .insert(NotificationArgs {}) + .opts(unique()) + .await + .unwrap(); + assert!(!first.unique_skipped_as_duplicate); + assert_eq!(notifications.next(&pool).await, ["default"]); + tokio::time::sleep(Duration::from_millis(2)).await; + let duplicate = client + .insert(NotificationArgs {}) + .opts(unique()) + .await + .unwrap(); + assert!(duplicate.unique_skipped_as_duplicate); + assert_eq!(notifications.next(&pool).await, ["default"]); + sqlite_cleanup(pool, path).await; + } +} diff --git a/rust/riverqueue/tests/job_list.rs b/rust/riverqueue/tests/job_list.rs new file mode 100644 index 000000000..0806c48e3 --- /dev/null +++ b/rust/riverqueue/tests/job_list.rs @@ -0,0 +1,347 @@ +//! Job list ordering and keyset pagination on every backend. +//! +//! Each scenario runs against PostgreSQL (in a unique schema, failing rather +//! than skipping when `RIVER_RUST_DATABASE_URL` is unset) and SQLite (in a +//! temporary file). PostgreSQL scenarios build only with `--cfg river_postgres_tests`. + +#![cfg(any(all(feature = "postgres", river_postgres_tests), feature = "sqlite"))] + +mod support; + +use chrono::{DateTime, Duration, TimeZone, Utc}; +use riverqueue::{Client, JobListCursor, JobListOrderBy, JobListParams, JobState, SortDirection}; + +/// A job to seed, with the time fields job lists order by. +#[derive(Clone, Copy)] +struct Seed { + attempted_at: Option>, + finalized_at: Option>, + scheduled_at: DateTime, + state: JobState, +} + +impl Seed { + fn new(state: JobState) -> Self { + Self { + attempted_at: None, + finalized_at: None, + scheduled_at: now(), + state, + } + } + + const fn attempted(mut self, at: DateTime) -> Self { + self.attempted_at = Some(at); + self + } + + const fn finalized(mut self, at: DateTime) -> Self { + self.finalized_at = Some(at); + self + } + + const fn scheduled(mut self, at: DateTime) -> Self { + self.scheduled_at = at; + self + } +} + +/// A fixed time with millisecond precision, which both backends store +/// exactly. +fn now() -> DateTime { + Utc.with_ymd_and_hms(2026, 9, 9, 12, 0, 0).unwrap() + Duration::milliseconds(123) +} + +fn at(offset: Duration) -> DateTime { + now() + offset +} + +/// Lists jobs with `order_by` and `states` in both directions, and checks +/// the full listing and a one-job-at-a-time pagination against `want_order` +/// (indexes into `seeds`, ascending). Pages alternate serialized cursors and +/// cursors built from the page's job. +async fn assert_pagination( + client: &Client, + ids: &[i64], + states: &[JobState], + order_by: JobListOrderBy, + want_order: &[usize], +) { + for direction in [SortDirection::Ascending, SortDirection::Descending] { + let mut want = want_order + .iter() + .map(|&index| ids[index]) + .collect::>(); + if direction == SortDirection::Descending { + want.reverse(); + } + let params = JobListParams::default() + .states(states.iter().copied()) + .order_by(order_by) + .direction(direction); + + let listed = client.jobs().list(params.clone()).await.unwrap(); + assert_eq!( + listed.jobs.iter().map(|job| job.id).collect::>(), + want, + "{direction:?} listing" + ); + + let mut got = Vec::new(); + let mut page_params = params.clone().limit(1); + for page in 0.. { + assert!( + page <= want.len(), + "{direction:?}: too many pages; got IDs so far: {got:?}" + ); + let result = client.jobs().list(page_params).await.unwrap(); + let Some(job) = result.jobs.first() else { + break; + }; + got.push(job.id); + let cursor = if page % 2 == 0 { + let encoded = result.last_cursor.expect("nonempty page").encode(); + JobListCursor::decode(&encoded).unwrap() + } else { + JobListCursor::from_job(job, ¶ms).unwrap() + }; + page_params = params.clone().limit(1).after(cursor); + } + assert_eq!(got, want, "{direction:?} pagination"); + } +} + +/// Defines each scenario for one backend's `Fixture`. +macro_rules! scenarios { + () => { + /// Time ordering over running and available jobs orders every job + /// by `attempted_at`, which is null for jobs never run. + #[tokio::test(flavor = "multi_thread")] + async fn mixed_states_page_by_attempted_at() { + let fixture = Fixture::new().await; + let ids = fixture + .insert(&[ + Seed::new(JobState::Running).attempted(at(Duration::seconds(2))), + Seed::new(JobState::Available), + Seed::new(JobState::Running).attempted(at(Duration::seconds(1))), + Seed::new(JobState::Available) + .attempted(at(-Duration::hours(1))) + .scheduled(at(Duration::hours(1))), + Seed::new(JobState::Running).attempted(at(Duration::seconds(1))), + Seed::new(JobState::Available), + ]) + .await; + + assert_pagination( + &fixture.client, + &ids, + &[JobState::Running, JobState::Available], + JobListOrderBy::Time, + &[3, 2, 4, 0, 1, 5], + ) + .await; + + fixture.cleanup().await; + } + + /// Time ordering over completed and available jobs orders every job + /// by `finalized_at`, which is null for the available ones. + #[tokio::test(flavor = "multi_thread")] + async fn mixed_states_page_by_finalized_at() { + let fixture = Fixture::new().await; + let ids = fixture + .insert(&[ + Seed::new(JobState::Available).scheduled(at(-Duration::hours(2))), + Seed::new(JobState::Completed) + .finalized(at(Duration::seconds(1))) + .scheduled(at(-Duration::hours(3))), + Seed::new(JobState::Completed) + .finalized(now()) + .scheduled(at(-Duration::hours(1))), + Seed::new(JobState::Available).scheduled(at(-Duration::hours(4))), + Seed::new(JobState::Completed) + .finalized(at(Duration::seconds(1))) + .scheduled(at(-Duration::hours(5))), + ]) + .await; + + assert_pagination( + &fixture.client, + &ids, + &[JobState::Completed, JobState::Available], + JobListOrderBy::Time, + &[2, 1, 4, 0, 3], + ) + .await; + + fixture.cleanup().await; + } + + /// Time ordering over available and cancelled jobs orders every job + /// by `scheduled_at`, even though a cancelled job's own time field + /// is `finalized_at`. + #[tokio::test(flavor = "multi_thread")] + async fn mixed_states_page_by_scheduled_at() { + let fixture = Fixture::new().await; + let ids = fixture + .insert(&[ + Seed::new(JobState::Available).scheduled(at(Duration::seconds(1))), + Seed::new(JobState::Cancelled) + .finalized(at(-Duration::hours(1))) + .scheduled(at(Duration::seconds(2))), + Seed::new(JobState::Available).scheduled(at(Duration::seconds(3))), + Seed::new(JobState::Cancelled) + .finalized(at(Duration::hours(1))) + .scheduled(at(Duration::seconds(2))), + Seed::new(JobState::Available).scheduled(at(Duration::seconds(4))), + ]) + .await; + + assert_pagination( + &fixture.client, + &ids, + &[JobState::Available, JobState::Cancelled], + JobListOrderBy::Time, + &[0, 1, 3, 2, 4], + ) + .await; + + fixture.cleanup().await; + } + + /// Time ordering without a state filter lists every state by + /// `scheduled_at`, the field for available jobs. + #[tokio::test(flavor = "multi_thread")] + async fn time_order_without_states_pages_by_scheduled_at() { + let fixture = Fixture::new().await; + let ids = fixture + .insert(&[ + Seed::new(JobState::Available).scheduled(at(Duration::seconds(3))), + Seed::new(JobState::Running) + .attempted(at(-Duration::hours(1))) + .scheduled(at(Duration::seconds(1))), + Seed::new(JobState::Completed) + .finalized(at(-Duration::hours(1))) + .scheduled(at(Duration::seconds(2))), + Seed::new(JobState::Scheduled).scheduled(at(Duration::seconds(1))), + ]) + .await; + + assert_pagination( + &fixture.client, + &ids, + &[], + JobListOrderBy::Time, + &[1, 3, 2, 0], + ) + .await; + + fixture.cleanup().await; + } + }; +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +mod postgres { + use super::*; + use crate::support::PostgresSchema; + + struct Fixture { + client: Client, + schema: PostgresSchema, + } + + impl Fixture { + async fn new() -> Self { + let schema = PostgresSchema::new("river_job_list").await; + let client = Client::builder( + riverqueue::database::PostgresDatabase::new(schema.pool.clone()) + .with_schema(schema.schema.clone()), + ) + .build() + .unwrap(); + Self { client, schema } + } + + async fn insert(&self, seeds: &[Seed]) -> Vec { + let sql = format!( + "INSERT INTO {} (args, attempted_at, finalized_at, kind, max_attempts, \ + scheduled_at, state) \ + VALUES ('{{}}', $1, $2, 'job_list', 25, $3, $4::text::{}) RETURNING id", + self.schema.table("river_job"), + self.schema.table("river_job_state"), + ); + let mut ids = Vec::with_capacity(seeds.len()); + for seed in seeds { + let id: i64 = sqlx::query_scalar(sqlx::AssertSqlSafe(sql.clone())) + .bind(seed.attempted_at) + .bind(seed.finalized_at) + .bind(seed.scheduled_at) + .bind(seed.state.as_str()) + .fetch_one(&self.schema.pool) + .await + .unwrap(); + ids.push(id); + } + ids + } + + async fn cleanup(self) { + self.schema.cleanup().await; + } + } + + scenarios!(); +} + +#[cfg(feature = "sqlite")] +mod sqlite { + use sqlx::SqlitePool; + + use super::*; + use crate::support::{sqlite_cleanup, sqlite_file_pool}; + + struct Fixture { + client: Client, + path: std::path::PathBuf, + pool: SqlitePool, + } + + impl Fixture { + async fn new() -> Self { + let (pool, path) = sqlite_file_pool(4).await; + let client = Client::builder(pool.clone()).build().unwrap(); + Self { client, path, pool } + } + + async fn insert(&self, seeds: &[Seed]) -> Vec { + let time = |time: DateTime| time.format("%Y-%m-%d %H:%M:%S%.3f").to_string(); + let mut ids = Vec::with_capacity(seeds.len()); + for seed in seeds { + let id: i64 = sqlx::query_scalar( + "INSERT INTO river_job (args, attempt, attempted_at, attempted_by, \ + created_at, errors, finalized_at, kind, max_attempts, metadata, priority, \ + queue, scheduled_at, state, tags) \ + VALUES (jsonb('{}'), 0, ?, jsonb('[]'), ?, jsonb('[]'), ?, 'job_list', 25, \ + jsonb('{}'), 1, 'default', ?, ?, jsonb('[]')) RETURNING id", + ) + .bind(seed.attempted_at.map(time)) + .bind(time(now())) + .bind(seed.finalized_at.map(time)) + .bind(time(seed.scheduled_at)) + .bind(seed.state.as_str()) + .fetch_one(&self.pool) + .await + .unwrap(); + ids.push(id); + } + ids + } + + async fn cleanup(self) { + sqlite_cleanup(self.pool, self.path).await; + } + } + + scenarios!(); +} diff --git a/rust/riverqueue/tests/leader_election_disabled.rs b/rust/riverqueue/tests/leader_election_disabled.rs new file mode 100644 index 000000000..16523ff3a --- /dev/null +++ b/rust/riverqueue/tests/leader_election_disabled.rs @@ -0,0 +1,384 @@ +//! Clients built with `without_leader_election`, on every backend. +//! +//! PostgreSQL tests run in a unique schema and fail rather than skip when +//! `RIVER_RUST_DATABASE_URL` is unset; SQLite tests use a temporary file. + +#![cfg(any(all(feature = "postgres", river_postgres_tests), feature = "sqlite"))] + +mod support; + +use std::{ + convert::Infallible, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use async_trait::async_trait; +use riverqueue::{ + __private::{ + ClientBuilderExt, Database, MaintenanceService, MaintenanceServiceContext, Pilot, + PilotError, RuntimeService, RuntimeServiceContext, + }, + Client, ClientBuilder, EventKind, EventReceiver, Job, JobArgs, JobRow, JobState, + MaintenanceConfig, NeverSchedule, PeriodicJob, PeriodicJobOpts, QueueConfig, WorkContext, + WorkOutcome, WorkerRegistry, +}; +use serde::{Deserialize, Serialize}; +use tokio_util::sync::CancellationToken; + +/// Every wait in these tests is bounded by this timeout. +const TIMEOUT: Duration = Duration::from_secs(10); + +/// A short election interval, so that a client wrongly taking part in +/// elections would become leader well within these tests. +const ELECT_INTERVAL: Duration = Duration::from_millis(50); + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_leader_election_disabled_noop")] +struct NoopArgs {} + +fn noop_workers() -> WorkerRegistry { + let mut workers = WorkerRegistry::new(); + workers + .register_fn(|_context: WorkContext, _job: Job| async { + Ok::<_, Infallible>(WorkOutcome::Complete) + }) + .unwrap(); + workers +} + +fn run_on_start_job() -> PeriodicJob { + PeriodicJob::with_options( + NeverSchedule, + || NoopArgs {}, + PeriodicJobOpts::new().with_run_on_start(true), + ) +} + +/// Counts starts of an extension's leader-owned and per-client services. +#[derive(Clone, Default)] +struct ServicePilot { + maintenance_services_calls: Arc, + maintenance_starts: Arc, + runtime_starts: Arc, +} + +#[async_trait] +impl Pilot for ServicePilot { + fn maintenance_services(&self) -> Vec> { + self.maintenance_services_calls + .fetch_add(1, Ordering::SeqCst); + vec![Arc::new(CountingService(Arc::clone( + &self.maintenance_starts, + )))] + } + + fn runtime_services(&self) -> Vec> { + vec![Arc::new(CountingService(Arc::clone(&self.runtime_starts)))] + } +} + +struct CountingService(Arc); + +impl CountingService { + async fn run_until_cancelled(&self, cancellation: CancellationToken) { + self.0.fetch_add(1, Ordering::SeqCst); + cancellation.cancelled().await; + } +} + +#[async_trait] +impl MaintenanceService for CountingService { + async fn run(&self, context: MaintenanceServiceContext) -> Result<(), PilotError> { + let cancellation = context.term.token; + self.run_until_cancelled(cancellation).await; + Ok(()) + } +} + +#[async_trait] +impl RuntimeService for CountingService { + async fn run(&self, context: RuntimeServiceContext) -> Result<(), PilotError> { + let cancellation = context.cancellation; + self.run_until_cancelled(cancellation).await; + Ok(()) + } +} + +/// A migrated database on one backend. +enum Backend { + #[cfg(all(feature = "postgres", river_postgres_tests))] + Postgres(support::PostgresSchema), + #[cfg(feature = "sqlite")] + Sqlite(sqlx::SqlitePool, std::path::PathBuf), +} + +impl Backend { + #[cfg(all(feature = "postgres", river_postgres_tests))] + async fn postgres() -> Self { + Self::Postgres(support::PostgresSchema::new("river_no_election").await) + } + + #[cfg(feature = "sqlite")] + async fn sqlite() -> Self { + let (pool, path) = support::sqlite_file_pool(4).await; + Self::Sqlite(pool, path) + } + + fn database(&self) -> Database { + match self { + #[cfg(all(feature = "postgres", river_postgres_tests))] + Self::Postgres(schema) => Database::from_source( + riverqueue::database::PostgresDatabase::new(schema.pool.clone()) + .with_schema(schema.schema.clone()), + ), + #[cfg(feature = "sqlite")] + Self::Sqlite(pool, _) => Database::from_source(pool.clone()), + } + } + + /// Returns a builder for a worker client with a short election interval. + fn builder(&self, id: &str, queue: &str) -> ClientBuilder { + Client::builder(self.database()) + .id(id) + .maintenance(MaintenanceConfig::default().with_elect_interval(ELECT_INTERVAL)) + .queue(queue, QueueConfig::new(1)) + .workers(noop_workers()) + } + + /// Returns the elected leader's client ID, if any. + async fn leader_id(&self) -> Option { + match self { + #[cfg(all(feature = "postgres", river_postgres_tests))] + Self::Postgres(schema) => sqlx::query_scalar(sqlx::AssertSqlSafe(format!( + "SELECT leader_id FROM {}", + schema.table("river_leader") + ))) + .fetch_optional(&schema.pool) + .await + .unwrap(), + #[cfg(feature = "sqlite")] + Self::Sqlite(pool, _) => sqlx::query_scalar("SELECT leader_id FROM river_leader") + .fetch_optional(pool) + .await + .unwrap(), + } + } + + async fn wait_for_leader(&self, id: &str) { + tokio::time::timeout(TIMEOUT, async { + while self.leader_id().await.as_deref() != Some(id) { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("{id} was not elected leader")); + } + + async fn cleanup(self) { + match self { + #[cfg(all(feature = "postgres", river_postgres_tests))] + Self::Postgres(schema) => schema.cleanup().await, + #[cfg(feature = "sqlite")] + Self::Sqlite(pool, path) => support::sqlite_cleanup(pool, path).await, + } + } +} + +async fn next_completed(events: &mut EventReceiver) -> JobRow { + let event = tokio::time::timeout(TIMEOUT, events.recv()) + .await + .expect("a job should complete") + .unwrap(); + let job = event.as_job().expect("a job event").clone(); + assert_eq!(job.job.state, JobState::Completed); + job.job +} + +/// Works jobs over two runs of the same client without ever electing a +/// leader, including stopping with no leadership to resign. +async fn works_jobs_without_electing(backend: Backend, poll_only: bool) { + let mut builder = backend + .builder("no_election", "default") + .without_leader_election(); + if poll_only { + builder = builder.without_notifications(); + } + let client = builder.build().unwrap(); + for _ in 0..2 { + let mut events = client.subscribe(&[EventKind::JobCompleted]).unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + let inserted = client.insert(NoopArgs {}).await.unwrap(); + assert_eq!(next_completed(&mut events).await.id, inserted.job.row.id); + assert_eq!(backend.leader_id().await, None); + + tokio::time::timeout(Duration::from_secs(5), run.shutdown()) + .await + .expect("the client should stop") + .unwrap(); + } + backend.cleanup().await; +} + +/// Works a periodic job another client enqueues as leader, runs no +/// leader-owned extension services, and stays ineligible once that leader +/// stops. +async fn stays_ineligible_after_leader_stops(backend: Backend) { + let pilot = ServicePilot::default(); + let client = backend + .builder("no_election", "default") + .without_leader_election() + .pilot(pilot.clone()) + .build() + .unwrap(); + let mut events = client.subscribe(&[EventKind::JobCompleted]).unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + // The leader enqueues a periodic job on the default queue, which only + // the client without leader election works. + let leader = backend + .builder("eligible_leader", "leader") + .periodic_job(run_on_start_job()) + .build() + .unwrap(); + let mut leader_run = leader.start().unwrap(); + leader_run.wait_ready().await.unwrap(); + backend.wait_for_leader("eligible_leader").await; + + let periodic = next_completed(&mut events).await; + assert_eq!(periodic.attempted_by, ["no_election"]); + assert_eq!( + backend.leader_id().await.as_deref(), + Some("eligible_leader") + ); + + tokio::time::timeout(Duration::from_secs(5), leader_run.shutdown()) + .await + .expect("the leader should stop") + .unwrap(); + + let inserted = client.insert(NoopArgs {}).await.unwrap(); + let worked = next_completed(&mut events).await; + assert_eq!(worked.id, inserted.job.row.id); + assert_eq!(worked.attempted_by, ["no_election"]); + assert_eq!(backend.leader_id().await, None); + + assert_eq!(pilot.maintenance_services_calls.load(Ordering::SeqCst), 0); + assert_eq!(pilot.maintenance_starts.load(Ordering::SeqCst), 0); + assert_eq!(pilot.runtime_starts.load(Ordering::SeqCst), 1); + + run.shutdown().await.unwrap(); + backend.cleanup().await; +} + +#[cfg(feature = "sqlite")] +mod configuration { + use riverqueue::IntervalSchedule; + + use super::*; + + #[tokio::test] + async fn build_rejects_periodic_jobs() { + let backend = Backend::sqlite().await; + let error = backend + .builder("no_election", "default") + .without_leader_election() + .periodic_job(run_on_start_job()) + .build() + .unwrap_err(); + assert!( + error + .to_string() + .contains("periodic jobs must be empty when leader election is disabled"), + "{error}" + ); + backend.cleanup().await; + } + + #[tokio::test] + async fn builder_reports_leader_election_disabled() { + let backend = Backend::sqlite().await; + let builder = backend.builder("default_client", "default"); + assert!(!builder.leader_election_disabled()); + assert!(builder.without_leader_election().leader_election_disabled()); + backend.cleanup().await; + } + + #[tokio::test] + async fn periodic_jobs_reject_additions() { + let backend = Backend::sqlite().await; + let client = backend + .builder("no_election", "default") + .without_leader_election() + .build() + .unwrap(); + let periodic = client.periodic_jobs(); + let job = || { + PeriodicJob::new( + IntervalSchedule::new(Duration::from_mins(1)).unwrap(), + || NoopArgs {}, + ) + }; + for error in [ + periodic.add(job()).unwrap_err(), + periodic.add_many(vec![job()]).unwrap_err(), + ] { + assert!( + error + .to_string() + .contains("periodic jobs can't be added when leader election is disabled"), + "{error}" + ); + } + // Removing jobs from the always-empty bundle is harmless. + periodic.clear(); + assert!(!periodic.remove_by_id("missing")); + backend.cleanup().await; + } +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +mod postgres { + use super::*; + + #[tokio::test(flavor = "multi_thread")] + async fn stays_ineligible_after_leader_stops() { + super::stays_ineligible_after_leader_stops(Backend::postgres().await).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn works_jobs_without_electing() { + super::works_jobs_without_electing(Backend::postgres().await, false).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn works_jobs_without_electing_poll_only() { + super::works_jobs_without_electing(Backend::postgres().await, true).await; + } +} + +#[cfg(feature = "sqlite")] +mod sqlite { + use super::*; + + #[tokio::test(flavor = "multi_thread")] + async fn stays_ineligible_after_leader_stops() { + super::stays_ineligible_after_leader_stops(Backend::sqlite().await).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn works_jobs_without_electing() { + super::works_jobs_without_electing(Backend::sqlite().await, false).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn works_jobs_without_electing_poll_only() { + super::works_jobs_without_electing(Backend::sqlite().await, true).await; + } +} diff --git a/rust/riverqueue/tests/lifecycle.rs b/rust/riverqueue/tests/lifecycle.rs new file mode 100644 index 000000000..b5fd3427d --- /dev/null +++ b/rust/riverqueue/tests/lifecycle.rs @@ -0,0 +1,475 @@ +//! Client lifecycle: stopping from other tasks, graceful shutdown signals, +//! cancel safety, and idempotency. +//! +//! Lifecycle behavior doesn't depend on the backend, so these tests use +//! temporary SQLite databases and need no external services. + +use std::{ + path::PathBuf, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use riverqueue::{ + Client, Job, JobArgs, JobState, QueueConfig, RunHandle, WorkCancelled, WorkContext, + WorkOutcome, WorkerRegistry, +}; +use riverqueue_migrate::SqliteMigrator; +use serde::{Deserialize, Serialize}; +use sqlx::{ + SqlitePool, + sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}, +}; +use tokio::sync::{Semaphore, oneshot}; +use tokio_util::sync::CancellationToken; + +/// A job that runs until the test releases it or the client cancels it. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_lifecycle_gated")] +struct GatedArgs {} + +/// How a gated job's worker ended. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum Ending { + Cancelled, + Released, +} + +/// Holds gated jobs inside their workers and records how each ended. +#[derive(Clone)] +struct Gate { + endings: Arc>>, + release: Arc, + started: Arc, +} + +impl Gate { + fn new() -> Self { + Self { + endings: Arc::new(Mutex::new(Vec::new())), + release: Arc::new(Semaphore::new(0)), + started: Arc::new(Semaphore::new(0)), + } + } + + fn ending(&self, id: i64) -> Option { + self.endings + .lock() + .unwrap() + .iter() + .find_map(|(job_id, ending)| (*job_id == id).then_some(*ending)) + } + + fn release(&self) { + self.release.add_permits(1); + } + + async fn wait_started(&self) { + tokio::time::timeout(Duration::from_secs(10), self.started.acquire()) + .await + .expect("gated job did not start") + .unwrap() + .forget(); + } + + fn workers(&self) -> WorkerRegistry { + let gate = self.clone(); + let mut workers = WorkerRegistry::new(); + workers + .register_fn(move |context: WorkContext, job: Job| { + let gate = gate.clone(); + async move { + gate.started.add_permits(1); + // Cancellation wins when both are ready, so a stop that + // cancelled work is always observed. + let ending = tokio::select! { + biased; + () = context.cancellation_token().cancelled() => Ending::Cancelled, + permit = gate.release.acquire() => { + permit.unwrap().forget(); + Ending::Released + } + }; + gate.endings.lock().unwrap().push((job.id(), ending)); + match ending { + Ending::Cancelled => Err(WorkCancelled), + Ending::Released => Ok(WorkOutcome::Complete), + } + } + }) + .unwrap(); + workers + } +} + +/// A migrated WAL database file that is removed when the test finishes. +struct TestDatabase { + path: PathBuf, + pool: SqlitePool, +} + +impl TestDatabase { + async fn new() -> Self { + static DATABASE_NONCE: AtomicUsize = AtomicUsize::new(0); + let path = std::env::temp_dir().join(format!( + "river-lifecycle-{}-{}.sqlite", + std::process::id(), + DATABASE_NONCE.fetch_add(1, Ordering::Relaxed) + )); + let pool = SqlitePoolOptions::new() + .max_connections(4) + .connect_with( + SqliteConnectOptions::new() + .filename(&path) + .create_if_missing(true) + .journal_mode(SqliteJournalMode::Wal) + .busy_timeout(Duration::from_secs(5)), + ) + .await + .unwrap(); + SqliteMigrator::new(pool.clone()) + .migrate_up() + .await + .unwrap(); + Self { path, pool } + } + + fn client(&self, gate: &Gate, max_workers: usize) -> Client { + self.client_with(gate, max_workers, None) + } + + fn client_with( + &self, + gate: &Gate, + max_workers: usize, + soft_stop_timeout: Option, + ) -> Client { + let builder = Client::builder(self.pool.clone()); + let builder = match soft_stop_timeout { + Some(timeout) => builder.soft_stop_timeout(timeout), + None => builder, + }; + builder + .id("rust-lifecycle-test") + .workers(gate.workers()) + .queue( + "default", + QueueConfig::new(max_workers) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(20)), + ) + .build() + .unwrap() + } +} + +impl Drop for TestDatabase { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.path); + for suffix in ["-shm", "-wal"] { + let mut path = self.path.as_os_str().to_owned(); + path.push(suffix); + let _ = std::fs::remove_file(path); + } + } +} + +async fn insert_gated(client: &Client) -> i64 { + client.insert(GatedArgs {}).await.unwrap().id() +} + +/// Waits for the client to stop, failing the test if it doesn't. +async fn wait_stopped(run: &mut RunHandle) { + tokio::time::timeout(Duration::from_secs(10), run.wait()) + .await + .expect("client did not stop") + .unwrap(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn dropping_a_shutdown_future_keeps_the_soft_stop() { + let database = TestDatabase::new().await; + let gate = Gate::new(); + let client = database.client(&gate, 1); + let running = insert_gated(&client).await; + let unfetched = insert_gated(&client).await; + + let mut run = client.start().unwrap(); + gate.wait_started().await; + // Poll `shutdown` once, which requests a soft stop, then drop it while + // the gated job keeps it pending, as `tokio::time::timeout` would. + tokio::select! { + biased; + result = run.shutdown() => panic!("shutdown finished while a job was held: {result:?}"), + () = std::future::ready(()) => {} + } + gate.release(); + wait_stopped(&mut run).await; + + assert_eq!(gate.ending(running), Some(Ending::Released)); + assert_eq!( + client.jobs().get(running).await.unwrap().state, + JobState::Completed + ); + // The soft stop requested by the dropped future stopped fetching. + assert_eq!( + client.jobs().get(unfetched).await.unwrap().state, + JobState::Available + ); + // The handle is still usable after the dropped future. + run.shutdown_now().await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn dropping_the_handle_requests_a_hard_stop() { + let database = TestDatabase::new().await; + let gate = Gate::new(); + let client = database.client(&gate, 1); + let running = insert_gated(&client).await; + + let run = client.start().unwrap(); + gate.wait_started().await; + drop(run); + + let deadline = tokio::time::Instant::now() + Duration::from_secs(10); + while client.jobs().get(running).await.unwrap().state != JobState::Available { + assert!( + tokio::time::Instant::now() < deadline, + "the job was not interrupted" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } + assert_eq!(gate.ending(running), Some(Ending::Cancelled)); +} + +#[tokio::test(flavor = "multi_thread")] +async fn graceful_shutdown_signal_accepts_a_cancellation_token() { + let database = TestDatabase::new().await; + let client = database.client(&Gate::new(), 1); + let token = CancellationToken::new(); + + let mut run = client + .start_with_graceful_shutdown(token.clone().cancelled_owned()) + .unwrap(); + run.wait_ready().await.unwrap(); + token.cancel(); + wait_stopped(&mut run).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn graceful_shutdown_signal_stops_softly() { + let database = TestDatabase::new().await; + let gate = Gate::new(); + let client = database.client(&gate, 1); + let running = insert_gated(&client).await; + let (signal_sender, signal) = oneshot::channel::<()>(); + + let mut run = client + .start_with_graceful_shutdown(async move { + let _ = signal.await; + }) + .unwrap(); + gate.wait_started().await; + signal_sender.send(()).unwrap(); + // The stop is applied asynchronously, so tests that also check that + // fetching stops use a `Stopper`, which is equivalent and synchronous. + gate.release(); + wait_stopped(&mut run).await; + + assert_eq!(gate.ending(running), Some(Ending::Released)); + assert_eq!( + client.jobs().get(running).await.unwrap().state, + JobState::Completed + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn lifecycle_methods_are_idempotent() { + let database = TestDatabase::new().await; + let gate = Gate::new(); + let client = database.client(&gate, 1); + + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + let stopper = run.stopper(); + stopper.stop(); + stopper.stop(); + run.shutdown().await.unwrap(); + run.shutdown().await.unwrap(); + run.shutdown_now().await.unwrap(); + run.wait().await.unwrap(); + run.wait_ready().await.unwrap(); + stopper.stop_now(); + stopper.stop(); + + // A stopper affects only its own run, not a restart of the client. + let job = insert_gated(&client).await; + let mut restarted = client.start().unwrap(); + gate.wait_started().await; + stopper.stop_now(); + gate.release(); + let completed = async { + while client.jobs().get(job).await.unwrap().state != JobState::Completed { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }; + tokio::time::timeout(Duration::from_secs(10), completed) + .await + .expect("the restarted client did not complete its job"); + assert_eq!(gate.ending(job), Some(Ending::Released)); + restarted.shutdown().await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn soft_stop_timeout_escalates_a_stop_from_a_stopper() { + let database = TestDatabase::new().await; + let gate = Gate::new(); + let client = database.client_with(&gate, 1, Some(Duration::from_millis(50))); + let running = insert_gated(&client).await; + + let mut run = client.start().unwrap(); + gate.wait_started().await; + run.stopper().stop(); + wait_stopped(&mut run).await; + + assert_eq!(gate.ending(running), Some(Ending::Cancelled)); + let job = client.jobs().get(running).await.unwrap(); + assert_eq!(job.state, JobState::Available); + assert_eq!(job.attempt, 0, "an interrupted job keeps its attempt"); + assert_eq!(job.errors, []); +} + +#[tokio::test(flavor = "multi_thread")] +async fn stop_from_another_task_while_waiting() { + let database = TestDatabase::new().await; + let gate = Gate::new(); + let client = database.client(&gate, 1); + let running = insert_gated(&client).await; + let unfetched = insert_gated(&client).await; + + let mut run = client.start().unwrap(); + let stopper = run.stopper(); + let stop_gate = gate.clone(); + let stop_task = tokio::spawn(async move { + stop_gate.wait_started().await; + stopper.stop(); + stop_gate.release(); + }); + wait_stopped(&mut run).await; + stop_task.await.unwrap(); + + assert_eq!(gate.ending(running), Some(Ending::Released)); + assert_eq!( + client.jobs().get(running).await.unwrap().state, + JobState::Completed + ); + assert_eq!( + client.jobs().get(unfetched).await.unwrap().state, + JobState::Available + ); +} + +#[tokio::test(flavor = "multi_thread")] +async fn stop_now_from_another_task_interrupts_running_jobs() { + let database = TestDatabase::new().await; + let gate = Gate::new(); + let client = database.client(&gate, 2); + let interrupted = insert_gated(&client).await; + let cancel_attempted = insert_gated(&client).await; + + let mut run = client.start().unwrap(); + let stopper = run.stopper(); + let stop_gate = gate.clone(); + let pool = database.pool.clone(); + let stop_task = tokio::spawn(async move { + stop_gate.wait_started().await; + stop_gate.wait_started().await; + // A cancellation whose notification never reached this client, as + // `job_cancel` records it on a running job. + sqlx::query( + "UPDATE river_job SET metadata = jsonb_set(metadata, '$.cancel_attempted_at', \ + '2026-01-02T03:04:05Z') WHERE id = ?", + ) + .bind(cancel_attempted) + .execute(&pool) + .await + .unwrap(); + stopper.stop_now(); + }); + wait_stopped(&mut run).await; + stop_task.await.unwrap(); + + assert_eq!(gate.ending(interrupted), Some(Ending::Cancelled)); + let interrupted = client.jobs().get(interrupted).await.unwrap(); + assert_eq!(interrupted.state, JobState::Available); + assert_eq!( + interrupted.attempt, 0, + "an interrupted job keeps its attempt" + ); + assert_eq!(interrupted.errors, []); + + // Like Go, a hard stop finalizes a job whose cancellation was requested + // instead of making it available again. + assert_eq!(gate.ending(cancel_attempted), Some(Ending::Cancelled)); + let cancel_attempted = client.jobs().get(cancel_attempted).await.unwrap(); + assert_eq!(cancel_attempted.state, JobState::Cancelled); + assert!(cancel_attempted.finalized_at.is_some()); +} + +// Like Go, whose fetches and notification polls run under a context the +// stop cancels, a stop during a database outage doesn't wait for the pool to +// hand out a connection. +#[tokio::test(flavor = "multi_thread")] +async fn stop_during_an_outage_does_not_wait_for_a_connection() { + let database = TestDatabase::new().await; + // A leader's resignation deliberately outlives a stop, with the bounded + // retries Go uses, so this client never leads. + let client = Client::builder(database.pool.clone()) + .without_leader_election() + .workers(Gate::new().workers()) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(20)), + ) + .build() + .unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + // Taking every connection leaves the client's producer, notification + // poller, and heartbeat waiting on the pool for its 30-second acquire + // timeout. + let mut held = Vec::new(); + for _ in 0..4 { + held.push(database.pool.acquire().await.unwrap()); + } + // Give polls and fetches time to start waiting; a stop must end them + // wherever they are. + tokio::time::sleep(Duration::from_millis(300)).await; + tokio::time::timeout(Duration::from_secs(5), run.shutdown()) + .await + .expect("the stop waited for a pool connection") + .unwrap(); + drop(held); +} + +#[tokio::test(flavor = "multi_thread")] +async fn wait_ready_waits_for_queue_registration() { + let database = TestDatabase::new().await; + let client = database.client(&Gate::new(), 1); + + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + // Like Go's `Client.Start`, readiness means peers can manage the queue. + let queues: i64 = sqlx::query_scalar("SELECT count(*) FROM river_queue WHERE name = 'default'") + .fetch_one(&database.pool) + .await + .unwrap(); + assert_eq!(queues, 1); + run.shutdown().await.unwrap(); +} diff --git a/rust/riverqueue/tests/metadata_exact.rs b/rust/riverqueue/tests/metadata_exact.rs new file mode 100644 index 000000000..72121d22c --- /dev/null +++ b/rust/riverqueue/tests/metadata_exact.rs @@ -0,0 +1,298 @@ +//! Persisted metadata stays readable when JSON numbers exceed `f64`. + +#![cfg(any(all(feature = "postgres", river_postgres_tests), feature = "sqlite"))] + +mod support; + +use riverqueue::{Client, InsertOpts, JobArgs, JobMetadata, JobRow, JobUpdateParams}; +#[cfg(feature = "sqlite")] +use riverqueue::{Job, JobState, QueueConfig, WorkContext, WorkOutcome, WorkerRegistry}; +use serde::{Deserialize, Serialize}; +#[cfg(feature = "sqlite")] +use std::{convert::Infallible, time::Duration}; + +const METADATA: &str = r#"{"zeta":"first","big_integer":123456789012345678901234567890,"beyond_float":1e400,"long_decimal":0.1000000000000000055511151231257827}"#; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "metadata_exact_insert")] +struct InsertArgs {} + +#[cfg(feature = "sqlite")] +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "metadata_exact_snooze")] +struct SnoozeArgs {} + +/// Inserts one job with `insert` and one with `insert_many`, both carrying +/// [`METADATA`] through [`InsertOpts`], and returns the stored rows. +async fn insert_with_exact_metadata(client: &Client) -> Vec { + let metadata: JobMetadata = METADATA.parse().unwrap(); + let opts = InsertOpts::default().with_metadata(metadata); + assert_eq!(opts.metadata().unwrap().as_raw().get(), METADATA); + + let single = client + .insert(InsertArgs {}) + .opts(opts.clone()) + .await + .unwrap(); + let many = client.insert_many([(InsertArgs {}, opts)]).await.unwrap(); + let mut rows = Vec::new(); + for id in std::iter::once(single.id()).chain(many.iter().map(riverqueue::InsertResult::id)) { + rows.push(client.jobs().get(id).await.unwrap()); + } + rows +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +fn raw_field<'a>(metadata: &'a JobMetadata, key: &str) -> &'a str { + metadata.get_raw(key).unwrap().get() +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +#[tokio::test] +async fn postgres_reads_metadata_with_large_numbers() { + use riverqueue::__private::{ExtensionClient, PreparedInsertParams}; + use riverqueue::database::PostgresDatabase; + use sqlx::AssertSqlSafe; + + let schema = support::PostgresSchema::new("meta_exact").await; + let table = schema.table("river_job"); + let id: i64 = sqlx::query_scalar(AssertSqlSafe(format!( + "INSERT INTO {table} (args, kind, max_attempts, metadata) \ + VALUES ('{{}}', 'metadata_exact', 25, $1::jsonb) RETURNING id" + ))) + .bind(METADATA) + .fetch_one(&schema.pool) + .await + .unwrap(); + let client = Client::builder( + PostgresDatabase::new(schema.pool.clone()).with_schema(schema.schema.clone()), + ) + .build() + .unwrap(); + + let original = client.jobs().get(id).await.unwrap(); + let big = raw_field(&original.metadata, "big_integer").to_owned(); + let beyond = raw_field(&original.metadata, "beyond_float").to_owned(); + assert_eq!(big, "123456789012345678901234567890"); + assert!(beyond.len() > 400); // PostgreSQL expands `1e400` in jsonb. + + let updated = client + .jobs() + .update(id, JobUpdateParams::default().output("done".into())) + .await + .unwrap(); + assert_eq!(raw_field(&updated.metadata, "big_integer"), big); + assert_eq!(raw_field(&updated.metadata, "beyond_float"), beyond); + assert_eq!( + updated.decode_output::().unwrap().as_deref(), + Some("done") + ); + let reread = client.jobs().get(id).await.unwrap(); + assert_eq!(raw_field(&reread.metadata, "beyond_float"), beyond); + + sqlx::query(AssertSqlSafe(format!( + "UPDATE {table} SET state = 'running' WHERE id = $1" + ))) + .bind(id) + .execute(&schema.pool) + .await + .unwrap(); + let mut complete_tx = schema.pool.begin().await.unwrap(); + client + .jobs() + .complete(id) + .tx(&mut complete_tx) + .await + .unwrap(); + complete_tx.commit().await.unwrap(); + let completed = client.jobs().get(id).await.unwrap(); + assert_eq!(raw_field(&completed.metadata, "beyond_float"), beyond); + + let mut transaction = schema.pool.begin().await.unwrap(); + let reinserted = ExtensionClient::new(&client) + .insert_prepared(vec![PreparedInsertParams { + created_at: reread.created_at, + encoded_args: reread.encoded_args.clone(), + kind: reread.kind.clone(), + max_attempts: reread.max_attempts, + metadata: reread.metadata.clone(), + priority: reread.priority, + queue: reread.queue.clone(), + scheduled_at: reread.scheduled_at, + tags: reread.tags.clone(), + unique_key: reread.unique_key.clone(), + unique_states: reread.unique_states.clone(), + }]) + .tx(&mut transaction) + .await + .unwrap() + .remove(0); + transaction.commit().await.unwrap(); + assert_eq!(raw_field(&reinserted.job.metadata, "beyond_float"), beyond); + + schema.cleanup().await; +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +#[tokio::test] +async fn postgres_insert_opts_keep_metadata_number_tokens() { + use riverqueue::database::PostgresDatabase; + + let schema = support::PostgresSchema::new("meta_exact_insert").await; + let client = Client::builder( + PostgresDatabase::new(schema.pool.clone()).with_schema(schema.schema.clone()), + ) + .build() + .unwrap(); + + for row in insert_with_exact_metadata(&client).await { + assert_eq!( + raw_field(&row.metadata, "big_integer"), + "123456789012345678901234567890" + ); + assert_eq!( + raw_field(&row.metadata, "long_decimal"), + "0.1000000000000000055511151231257827" + ); + // PostgreSQL expands `1e400` in jsonb rather than rejecting it. + let beyond = raw_field(&row.metadata, "beyond_float"); + assert!(beyond.starts_with('1') && beyond.len() > 400, "{beyond}"); + } + + schema.cleanup().await; +} + +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn sqlite_insert_opts_keep_metadata_number_tokens() { + let (pool, path) = support::sqlite_file_pool(4).await; + let client = Client::builder(pool.clone()).build().unwrap(); + + for row in insert_with_exact_metadata(&client).await { + for (key, token) in [ + ("big_integer", "123456789012345678901234567890"), + ("beyond_float", "1e400"), + ("long_decimal", "0.1000000000000000055511151231257827"), + ] { + assert_eq!(row.metadata.get_raw(key).unwrap().get(), token); + } + } + + support::sqlite_cleanup(pool, path).await; +} + +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn sqlite_reads_metadata_with_large_numbers() { + let (pool, path) = support::sqlite_file_pool(4).await; + let id: i64 = sqlx::query_scalar( + "INSERT INTO river_job (args, kind, max_attempts, metadata) \ + VALUES (jsonb('{}'), 'metadata_exact', 25, jsonb(?)) RETURNING id", + ) + .bind(METADATA) + .fetch_one(&pool) + .await + .unwrap(); + let client = Client::builder(pool.clone()).build().unwrap(); + + let original = client.jobs().get(id).await.unwrap(); + assert_eq!( + original.metadata.get_raw("beyond_float").unwrap().get(), + "1e400" + ); + assert_eq!( + original.metadata.get_raw("long_decimal").unwrap().get(), + "0.1000000000000000055511151231257827" + ); + + let updated = client + .jobs() + .update(id, JobUpdateParams::default().output("done".into())) + .await + .unwrap(); + assert_eq!( + updated.metadata.get_raw("beyond_float").unwrap().get(), + "1e400" + ); + assert_eq!( + updated.decode_output::().unwrap().as_deref(), + Some("done") + ); + let reread = client.jobs().get(id).await.unwrap(); + assert_eq!( + reread.metadata.get_raw("beyond_float").unwrap().get(), + "1e400" + ); + + sqlx::query("UPDATE river_job SET state = 'running' WHERE id = ?") + .bind(id) + .execute(&pool) + .await + .unwrap(); + let mut complete_tx = pool.begin_with("BEGIN IMMEDIATE").await.unwrap(); + client + .jobs() + .complete(id) + .tx(&mut complete_tx) + .await + .unwrap(); + complete_tx.commit().await.unwrap(); + let completed = client.jobs().get(id).await.unwrap(); + assert_eq!( + completed.metadata.get_raw("beyond_float").unwrap().get(), + "1e400" + ); + + support::sqlite_cleanup(pool, path).await; +} + +#[cfg(feature = "sqlite")] +#[tokio::test] +async fn sqlite_snooze_preserves_large_metadata_numbers() { + let (pool, path) = support::sqlite_file_pool(4).await; + let mut workers = WorkerRegistry::new(); + workers + .register_fn(|_context: WorkContext, _job: Job| async { + Ok::<_, Infallible>(WorkOutcome::Snooze(Duration::from_hours(1))) + }) + .unwrap(); + let client = Client::builder(pool.clone()) + .workers(workers) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let inserted = client.insert(SnoozeArgs {}).await.unwrap(); + sqlx::query("UPDATE river_job SET metadata = jsonb(?) WHERE id = ?") + .bind(METADATA) + .bind(inserted.job.row.id) + .execute(&pool) + .await + .unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + let row = tokio::time::timeout(Duration::from_secs(5), async { + loop { + let row = client.jobs().get(inserted.job.row.id).await.unwrap(); + if row.state == JobState::Scheduled + && row.metadata.get::("snoozes").unwrap() == Some(1) + { + break row; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + run.shutdown().await.unwrap(); + assert_eq!(row.metadata.get_raw("beyond_float").unwrap().get(), "1e400"); + assert_eq!( + row.metadata.get_raw("long_decimal").unwrap().get(), + "0.1000000000000000055511151231257827" + ); + support::sqlite_cleanup(pool, path).await; +} diff --git a/rust/riverqueue/tests/peer_attempts.rs b/rust/riverqueue/tests/peer_attempts.rs new file mode 100644 index 000000000..25f12ac67 --- /dev/null +++ b/rust/riverqueue/tests/peer_attempts.rs @@ -0,0 +1,1645 @@ +//! Peer attempts: jobs a running attempt claims and completes alongside its +//! own, owned by River from the claim's commit until each outcome persists. +//! +//! Every scenario uses generic claim statements written here, on PostgreSQL +//! (in a unique schema, failing rather than skipping when +//! `RIVER_RUST_DATABASE_URL` is unset) and SQLite (in a temporary file). + +#![cfg(any(all(feature = "postgres", river_postgres_tests), feature = "sqlite"))] + +mod support; + +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::Duration, +}; + +use async_trait::async_trait; +use futures_util::future::BoxFuture; +use riverqueue::__private::{ + ClaimedJob, ClientBuilderExt, DatabaseConnection, JobSetStateParams, PeerAttempts, + PeerClaimContext, PeerOutcome, Pilot, PilotError, PilotProducer, ProducerStartContext, +}; +use riverqueue::{ + BoxError, Client, ErrorHandler, ErrorHandlerDecision, Event, EventKind, InsertOpts, Job, + JobArgs, JobRow, JobState, QueueConfig, WorkContext, WorkOutcome, WorkResult, WorkerRegistry, +}; +use serde::{Deserialize, Serialize}; +use tokio::sync::{Notify, mpsc}; + +const WAIT: Duration = Duration::from_secs(10); + +/// The client ID every scenario's client uses. +const PEER_CLIENT: &str = "peer-client"; + +/// A job this client works as an ordinary attempt, held until released. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "peer_busy")] +struct BusyArgs {} + +/// The job whose attempt coordinates peers, running a scenario's script. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "peer_coordinator")] +struct CoordinatorArgs {} + +/// A peer job. Peers go to a queue no client works, so only claims take +/// them. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "peer_job")] +struct PeerArgs { + value: i64, +} + +/// Runs raw statements on either backend. +#[derive(Clone)] +enum Db { + #[cfg(all(feature = "postgres", river_postgres_tests))] + /// The pool, the qualified job table, and the qualified name for a + /// test function. + Postgres(sqlx::PgPool, String, String), + #[cfg(feature = "sqlite")] + Sqlite(sqlx::SqlitePool), +} + +impl Db { + /// Runs `postgres` (with `{table}` for the job table) or `sqlite` on job + /// `id`, bound as the only parameter. + async fn exec(&self, postgres: &str, sqlite: &str, id: i64) { + let _ = (postgres, sqlite); + match self { + #[cfg(all(feature = "postgres", river_postgres_tests))] + Self::Postgres(pool, table, _) => { + sqlx::query(sqlx::AssertSqlSafe(postgres.replace("{table}", table))) + .bind(id) + .execute(pool) + .await + .unwrap(); + } + #[cfg(feature = "sqlite")] + Self::Sqlite(pool) => { + sqlx::query(sqlx::AssertSqlSafe(sqlite.to_owned())) + .bind(id) + .execute(pool) + .await + .unwrap(); + } + } + } + + /// Runs statements that take no parameters. + async fn raw(&self, postgres: &str, sqlite: &str) { + let _ = (postgres, sqlite); + match self { + #[cfg(all(feature = "postgres", river_postgres_tests))] + Self::Postgres(pool, table, function) => { + sqlx::raw_sql(sqlx::AssertSqlSafe( + postgres + .replace("{table}", table) + .replace("{function}", function), + )) + .execute(pool) + .await + .unwrap(); + } + #[cfg(feature = "sqlite")] + Self::Sqlite(pool) => { + sqlx::raw_sql(sqlx::AssertSqlSafe(sqlite.to_owned())) + .execute(pool) + .await + .unwrap(); + } + } + } + + /// Makes every commit that leaves a peer job running fail. + async fn fail_peer_commits(&self) { + self.raw( + "CREATE FUNCTION {function}() RETURNS trigger LANGUAGE plpgsql AS $$ \ + BEGIN IF NEW.state = 'running' AND NEW.kind = 'peer_job' THEN \ + RAISE EXCEPTION 'peer commit failed on purpose'; END IF; RETURN NULL; END $$; \ + CREATE CONSTRAINT TRIGGER fail_peer_commit AFTER UPDATE ON {table} \ + DEFERRABLE INITIALLY DEFERRED FOR EACH ROW EXECUTE FUNCTION {function}();", + "CREATE TABLE guard_parent (id INTEGER PRIMARY KEY); \ + CREATE TABLE peer_guard (parent_id INTEGER REFERENCES guard_parent (id) \ + DEFERRABLE INITIALLY DEFERRED); \ + CREATE TRIGGER fail_peer_commit AFTER UPDATE OF state ON river_job \ + WHEN NEW.state = 'running' AND NEW.kind = 'peer_job' \ + BEGIN INSERT INTO peer_guard VALUES (-1); END;", + ) + .await; + } + + async fn allow_peer_commits(&self) { + self.raw( + "DROP TRIGGER fail_peer_commit ON {table}", + "DROP TRIGGER fail_peer_commit", + ) + .await; + } + + /// Marks job `id` completed at attempt 1 by this scenario's client. + async fn complete_behind_the_back(&self, id: i64) { + self.exec( + &format!( + "UPDATE {{table}} SET state = 'completed', attempt = 1, finalized_at = now(), \ + attempted_by = ARRAY['{PEER_CLIENT}'] WHERE id = $1" + ), + &format!( + "UPDATE river_job SET state = 'completed', attempt = 1, \ + finalized_at = strftime('%Y-%m-%d %H:%M:%f', 'now'), \ + attempted_by = jsonb('[\"{PEER_CLIENT}\"]') WHERE id = ?" + ), + id, + ) + .await; + } + + /// Makes job `id` undecodable while keeping its identity. + async fn corrupt(&self, id: i64) { + self.exec( + "UPDATE {table} SET metadata = '[1]'::jsonb WHERE id = $1", + "UPDATE river_job SET tags = jsonb('{\"not\":\"an array\"}') WHERE id = ?", + id, + ) + .await; + } + + /// Finalizes a running job behind River's back. + async fn discard(&self, id: i64) { + self.exec( + "UPDATE {table} SET state = 'discarded', finalized_at = now() WHERE id = $1", + "UPDATE river_job SET state = 'discarded', \ + finalized_at = strftime('%Y-%m-%d %H:%M:%f', 'now') WHERE id = ?", + id, + ) + .await; + } +} + +/// How a scripted claim takes its rows. +#[derive(Clone, Copy, Debug)] +enum Take { + /// Claims the rows for this client, as a real peer claim does. + Claim, + /// Claims the rows for another client. + Foreign, + /// Claims the rows, then returns the first one twice. + Duplicate, + /// Reads the rows without claiming them. + Read, + /// Returns a row River can't identify. + Unidentifiable, + /// Claims the rows, then waits for the coordinator's cancellation before + /// returning, so it's cancelled before commit. + ClaimUntilCancelled(&'static str), + /// Claims the rows, then waits for the scenario's gate before returning. + ClaimAfterGate, + /// Claims the rows, then raises a signal as it returns. + ClaimAndSignal(&'static str), +} + +/// Forces the claim callback's signature. +fn claim_callback(callback: F) -> F +where + F: for<'c> FnOnce(PeerClaimContext<'c>) -> BoxFuture<'c, Result, PilotError>> + + Send, +{ + callback +} + +/// A peer claim of `ids`, taken as `take` says. +async fn scripted_claim( + context: PeerClaimContext<'_>, + take: Take, + ids: Vec, + gate: Arc, + signals: Signals, +) -> Result, PilotError> { + let client_id = match take { + Take::Foreign => "another-client".to_owned(), + _ => context.client_id.to_owned(), + }; + let claims = !matches!(take, Take::Read | Take::Unidentifiable); + let mut rows = match context.connection { + #[cfg(feature = "postgres")] + DatabaseConnection::Postgres(connection) => { + let projection = riverqueue::__private::postgres_job_projection("job"); + let table = context + .database + .config() + .postgres_schema() + .unwrap() + .qualify("river_job"); + let rows = if matches!(take, Take::Unidentifiable) { + sqlx::query("SELECT 1 AS id").fetch_all(connection).await? + } else if claims { + sqlx::query(sqlx::AssertSqlSafe(format!( + "UPDATE {table} AS job SET state = 'running', attempt = job.attempt + 1, \ + attempted_at = now(), attempted_by = array_append(job.attempted_by, $1) \ + WHERE id = ANY($2) RETURNING {projection}, false AS unique_skipped_as_duplicate" + ))) + .bind(&client_id) + .bind(&ids) + .fetch_all(connection) + .await? + } else { + sqlx::query(sqlx::AssertSqlSafe(format!( + "SELECT {projection}, false AS unique_skipped_as_duplicate FROM {table} AS job \ + WHERE id = ANY($1)" + ))) + .bind(&ids) + .fetch_all(connection) + .await? + }; + rows.iter() + .map(riverqueue::__private::claimed_postgres_job) + .collect::>() + } + #[cfg(feature = "sqlite")] + DatabaseConnection::Sqlite(connection) => { + let columns = riverqueue::__private::SQLITE_JOB_COLUMNS; + let ids_json = serde_json::to_string(&ids)?; + let rows = if matches!(take, Take::Unidentifiable) { + sqlx::query("SELECT 1 AS id").fetch_all(connection).await? + } else if claims { + sqlx::query(sqlx::AssertSqlSafe(format!( + "UPDATE river_job SET state = 'running', attempt = attempt + 1, \ + attempted_at = ?, \ + attempted_by = jsonb_insert(coalesce(attempted_by, jsonb('[]')), '$[#]', ?) \ + WHERE id IN (SELECT value FROM json_each(?)) RETURNING {columns}" + ))) + .bind(riverqueue::__private::sqlite_timestamp(chrono::Utc::now())) + .bind(&client_id) + .bind(&ids_json) + .fetch_all(connection) + .await? + } else { + sqlx::query(sqlx::AssertSqlSafe(format!( + "SELECT {columns} FROM river_job WHERE id IN (SELECT value FROM json_each(?))" + ))) + .bind(&ids_json) + .fetch_all(connection) + .await? + }; + rows.iter() + .map(riverqueue::__private::claimed_sqlite_job) + .collect::>() + } + #[allow(unreachable_patterns)] + _ => unreachable!("built-in backends only"), + }; + match take { + Take::Duplicate => { + let first = rows[0].job().unwrap().clone(); + rows.push(first.into()); + } + Take::ClaimUntilCancelled(signal) => { + signals.raise(signal); + context.cancellation.cancelled().await; + } + Take::ClaimAfterGate => { + signals.raise("claim holds its rows"); + gate.notified().await; + } + Take::ClaimAndSignal(signal) => signals.raise(signal), + _ => {} + } + Ok(rows) +} + +/// Named signals a script raises for the test to wait on. +#[derive(Clone, Default)] +struct Signals { + changed: Arc, + raised: Arc>>, +} + +impl Signals { + fn raise(&self, signal: &'static str) { + self.raised.lock().unwrap().push(signal); + self.changed.notify_waiters(); + } + + async fn wait(&self, signal: &'static str) { + tokio::time::timeout(WAIT, async { + loop { + let changed = self.changed.notified(); + if self.raised.lock().unwrap().contains(&signal) { + return; + } + changed.await; + } + }) + .await + .unwrap_or_else(|_| panic!("timed out waiting for {signal}")); + } +} + +/// What a scenario's script observed, by name, checked by the test. +type Checks = Vec<(String, bool)>; + +/// Everything a script can use. +#[derive(Clone)] +struct Env { + checks: mpsc::UnboundedSender, + db: Db, + gate: Arc, + /// The coordinator's context, kept for operations after it ended. + kept: Arc>>, + /// Releases [`BusyArgs`] jobs. + busy: Arc, + peers: Vec, + signals: Signals, +} + +impl Env { + async fn claim( + &self, + context: &WorkContext, + take: Take, + ids: &[i64], + ) -> Result, riverqueue::Error> { + let (ids, gate, signals) = (ids.to_vec(), Arc::clone(&self.gate), self.signals.clone()); + PeerAttempts::new(context) + .claim(claim_callback(move |claim| { + Box::pin(scripted_claim(claim, take, ids, gate, signals)) + })) + .await + } +} + +type Script = + Arc BoxFuture<'static, WorkOutcome> + Send + Sync>; + +fn peer_error(error: &riverqueue::Error, text: &str) -> bool { + matches!( + error, + riverqueue::Error::Extension { + phase: riverqueue::ExtensionPhase::AddOn { + operation: "peer attempts" + }, + .. + } + ) && riverqueue::__private::error_chain(error).contains(text) +} + +/// Records what River tells the coordinator's producer session and the +/// completion step, to show peers take neither. +#[derive(Clone, Default)] +struct PeerPilot { + finished: Arc>>, + set_state: Arc>>, +} + +#[async_trait] +impl Pilot for PeerPilot { + fn intercepts_job_set_state(&self) -> bool { + true + } + + async fn after_jobs_set_state( + &self, + _connection: DatabaseConnection<'_>, + params: &JobSetStateParams, + ) -> Result<(), PilotError> { + self.set_state + .lock() + .unwrap() + .extend_from_slice(params.job_ids); + Ok(()) + } + + async fn start_producer( + &self, + _context: ProducerStartContext, + ) -> Result>, PilotError> { + Ok(Some(Box::new(self.clone()))) + } +} + +impl PilotProducer for PeerPilot { + fn job_finished(&self, job: &JobRow) { + self.finished.lock().unwrap().push(job.id); + } +} + +/// Cancels failed peers whose `value` is 2 and records every job it sees. +#[derive(Clone, Default)] +struct PeerErrorHandler(Arc>>); + +impl ErrorHandler for PeerErrorHandler { + async fn handle_error( + &self, + _context: &WorkContext, + job: &JobRow, + _result: &WorkResult, + ) -> Result { + tokio::task::yield_now().await; + self.0.lock().unwrap().push(job.id); + let args: PeerArgs = job.decode_args()?; + Ok(if args.value == 2 { + ErrorHandlerDecision::Cancel + } else { + ErrorHandlerDecision::default() + }) + } +} + +/// A started scenario. +struct Run { + client: Client, + coordinator: i64, + events: riverqueue::EventReceiver, + handle: riverqueue::RunHandle, + env: Env, + checks: mpsc::UnboundedReceiver, + handler: PeerErrorHandler, + pilot: PeerPilot, +} + +impl Run { + async fn start( + builder: riverqueue::ClientBuilder, + db: Db, + peers: usize, + script: Script, + ) -> Self { + let pilot = PeerPilot::default(); + let handler = PeerErrorHandler::default(); + let (checks_sender, checks) = mpsc::unbounded_channel(); + let env = Env { + checks: checks_sender, + db, + gate: Arc::new(Notify::new()), + kept: Arc::default(), + busy: Arc::new(Notify::new()), + peers: Vec::new(), + signals: Signals::default(), + }; + let env_slot = Arc::new(Mutex::new(None::)); + let mut workers = WorkerRegistry::new(); + let worker_env = Arc::clone(&env_slot); + workers + .register_fn(move |context: WorkContext, job: Job| { + let env = worker_env + .lock() + .unwrap() + .clone() + .expect("scenario started"); + let run = script(context, job.row, env); + async move { Ok::<_, BoxError>(run.await) } + }) + .unwrap() + .register_fn(|_context: WorkContext, _job: Job| async { + Ok::<_, BoxError>(WorkOutcome::Complete) + }) + .unwrap(); + let busy_env = Arc::clone(&env_slot); + workers + .register_fn(move |_context: WorkContext, _job: Job| { + let env = busy_env.lock().unwrap().clone().expect("scenario started"); + async move { + let released = env.busy.notified(); + env.signals.raise("busy"); + released.await; + Ok::<_, BoxError>(WorkOutcome::Complete) + } + }) + .unwrap(); + let client = builder + .id(PEER_CLIENT) + .pilot(pilot.clone()) + .error_handler(handler.clone()) + .queue( + "default", + QueueConfig::new(2) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .workers(workers) + .build() + .unwrap(); + let mut env = env; + for value in 1..=i64::try_from(peers).unwrap() { + env.peers.push( + client + .insert(PeerArgs { value }) + .opts(InsertOpts::default().with_queue("peers")) + .await + .unwrap() + .id(), + ); + } + *env_slot.lock().unwrap() = Some(env.clone()); + let events = client + .subscribe(&[ + EventKind::JobCancelled, + EventKind::JobCompleted, + EventKind::JobFailed, + EventKind::JobSnoozed, + ]) + .unwrap(); + let coordinator = client.insert(CoordinatorArgs {}).await.unwrap().id(); + let handle = client.start().unwrap(); + Self { + client, + coordinator, + events, + handle, + env, + checks, + handler, + pilot, + } + } + + /// Waits for the script's checks and asserts every one. + async fn assert_checks(&mut self) { + let checks = tokio::time::timeout(WAIT, self.checks.recv()) + .await + .expect("script reports its checks") + .unwrap(); + let failed = checks + .iter() + .filter(|(_, passed)| !passed) + .map(|(name, _)| name.as_str()) + .collect::>(); + assert!(failed.is_empty(), "failed checks: {failed:?}"); + } + + /// Waits for job `id` to reach a state other than `running`. + async fn settled(&self, id: i64) -> JobRow { + tokio::time::timeout(WAIT, async { + loop { + let row = self.client.jobs().get(id).await.unwrap(); + if row.state != JobState::Running { + return row; + } + tokio::task::yield_now().await; + } + }) + .await + .unwrap_or_else(|_| panic!("job {id} never settled")) + } + + /// Collects events until one for each of `ids` arrived, in order. + async fn events_for(&mut self, ids: &[i64]) -> Vec { + let mut seen = Vec::new(); + tokio::time::timeout(WAIT, async { + while !ids.iter().all(|id| { + seen.iter() + .any(|event: &Event| event.as_job().unwrap().job.id == *id) + }) { + seen.push(self.events.recv().await.unwrap()); + } + }) + .await + .expect("events arrive"); + seen + } + + async fn stop(mut self) { + tokio::time::timeout(WAIT, self.handle.shutdown()) + .await + .expect("client stops") + .unwrap(); + } +} + +fn check(checks: &mut Checks, name: &str, passed: bool) { + checks.push((name.to_owned(), passed)); +} + +/// Peer outcomes go through River's ordinary completion pipeline: the error +/// handler, the coordinator's metadata, the completion step, fencing, and +/// events. A cancellation requested while a peer ran wins, a stale result +/// leaves the row alone, and peers never reach the producer session. +async fn assert_outcomes_use_the_completion_pipeline(builder: riverqueue::ClientBuilder, db: Db) { + let script: Script = Arc::new(|context, _row, env| { + Box::pin(async move { + let mut checks = Checks::new(); + context.metadata_set("shared", true).unwrap(); + let rows = env.claim(&context, Take::Claim, &env.peers).await.unwrap(); + check(&mut checks, "claims all four", rows.len() == 4); + let row = |id: i64| rows.iter().find(|row| row.id == id).unwrap().clone(); + let [p1, p2, p3, p4] = [env.peers[0], env.peers[1], env.peers[2], env.peers[3]]; + env.db.discard(p3).await; + let client = context.client().unwrap(); + client.jobs().cancel(p4).await.unwrap(); + let completed = PeerAttempts::new(&context) + .complete(vec![ + PeerOutcome::new(row(p1), Ok(WorkOutcome::Complete)), + PeerOutcome::new(row(p2), Err("peer failed".into())), + PeerOutcome::new(row(p3), Ok(WorkOutcome::Complete)), + PeerOutcome::new(row(p4), Ok(WorkOutcome::Snooze(Duration::from_hours(1)))), + ]) + .await; + check(&mut checks, "complete succeeds", completed.is_ok()); + let _ = env.checks.send(checks); + WorkOutcome::Complete + }) + }); + let mut run = Run::start(builder, db, 4, script).await; + run.assert_checks().await; + let [p1, p2, p3, p4] = [ + run.env.peers[0], + run.env.peers[1], + run.env.peers[2], + run.env.peers[3], + ]; + let events = run.events_for(&[p1, p2, p4, run.coordinator]).await; + let states = events + .iter() + .map(|event| { + ( + event.as_job().unwrap().job.id, + event.as_job().unwrap().job.state, + ) + }) + .collect::>(); + assert_eq!(states[&p1], JobState::Completed); + // The error handler cancelled the failed peer. + assert_eq!(states[&p2], JobState::Cancelled); + // The cancellation requested while the peer ran wins over its snooze. + assert_eq!(states[&p4], JobState::Cancelled); + let p1_event = events + .iter() + .find(|event| event.as_job().unwrap().job.id == p1) + .unwrap(); + assert_eq!( + p1_event + .as_job() + .unwrap() + .job + .metadata + .get::("shared") + .unwrap(), + Some(true) + ); + // The stale result left the finalized row alone. + assert_eq!( + run.client.jobs().get(p3).await.unwrap().state, + JobState::Discarded + ); + assert_eq!(*run.handler.0.lock().unwrap(), [p2]); + let set_state = run.pilot.set_state.lock().unwrap().clone(); + for id in [p1, p2, p3, p4] { + assert!(set_state.contains(&id), "{id} in {set_state:?}"); + } + // Peers take no producer slot, so the session never hears of them. + run.settled(run.coordinator).await; + let coordinator = run.coordinator; + tokio::time::timeout(WAIT, async { + while !run.pilot.finished.lock().unwrap().contains(&coordinator) { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + assert_eq!(*run.pilot.finished.lock().unwrap(), [coordinator]); + run.stop().await; +} + +/// A claim River can't accept rolls back whole: duplicate rows, the +/// coordinator's own job, another client's attempt, a row that isn't +/// running, one River can't identify, one already owned, and a stale +/// attempt. A peer is owned only until its outcome persists, after which the +/// same coordinator can claim it again at a new attempt. +#[allow( + clippy::too_many_lines, + reason = "one coordinator walks through every rejected claim in turn" +)] +async fn assert_claims_are_checked(builder: riverqueue::ClientBuilder, db: Db) { + let script: Script = Arc::new(|context, row, env| { + Box::pin(async move { + let mut checks = Checks::new(); + let [p1, p2, p3] = [env.peers[0], env.peers[1], env.peers[2]]; + let client = context.client().unwrap().clone(); + let job_state = |id: i64| { + let client = client.clone(); + async move { client.jobs().get(id).await.unwrap().state } + }; + for (name, take, ids, reason) in [ + ("duplicate", Take::Duplicate, vec![p1], "twice"), + ("coordinator", Take::Claim, vec![row.id], "own job"), + ("foreign", Take::Foreign, vec![p1], "another client"), + ("not running", Take::Read, vec![p1], "has no attempt"), + ( + "unidentifiable", + Take::Unidentifiable, + vec![p1], + "couldn't be identified", + ), + ] { + let claimed = env.claim(&context, take, &ids).await; + check( + &mut checks, + &format!("{name} rejected: {claimed:?}"), + claimed + .as_ref() + .is_err_and(|error| peer_error(error, reason)), + ); + check( + &mut checks, + &format!("{name} rolled back"), + job_state(p1).await == JobState::Available, + ); + } + // A row at an attempt of this client that finished elsewhere. + let p4 = env.peers[3]; + env.db.complete_behind_the_back(p4).await; + let finished = env.claim(&context, Take::Read, &[p4]).await; + check( + &mut checks, + "not running rejected", + finished.is_err_and(|error| peer_error(&error, "isn't running")), + ); + // A job this client works as an ordinary attempt. + let busy = client.insert(BusyArgs {}).await.unwrap().id(); + env.signals.wait("busy").await; + let worked = env.claim(&context, Take::Read, &[busy]).await; + check( + &mut checks, + "worked here rejected", + worked.is_err_and(|error| peer_error(&error, "which this client already works")), + ); + env.busy.notify_one(); + let owned = env.claim(&context, Take::Claim, &[p2]).await.unwrap(); + let again = env.claim(&context, Take::Read, &[p2]).await; + check( + &mut checks, + "owned rejected", + again.is_err_and(|error| peer_error(&error, "already works as a peer")), + ); + // A snooze gives the attempt back, so claiming the job again + // reaches the same attempt number, which already ended here. + PeerAttempts::new(&context) + .complete(vec![PeerOutcome::new( + owned[0].clone(), + Ok(WorkOutcome::Snooze(Duration::ZERO)), + )]) + .await + .unwrap(); + let stale = env.claim(&context, Take::Claim, &[p2]).await; + check( + &mut checks, + "stale rejected", + stale.is_err_and(|error| peer_error(&error, "already ended here")), + ); + // A failure keeps the attempt, and ownership ends once the + // outcome persisted, so the job can be claimed again at once. + let first = env.claim(&context, Take::Claim, &[p3]).await.unwrap(); + PeerAttempts::new(&context) + .complete(vec![PeerOutcome::new( + first[0].clone(), + Err("retry now".into()), + )]) + .await + .unwrap(); + let second = env.claim(&context, Take::Claim, &[p3]).await; + check( + &mut checks, + "re-claimed at the next attempt", + second + .as_ref() + .is_ok_and(|rows| rows.len() == 1 && rows[0].attempt == 2), + ); + if let Ok(rows) = second { + PeerAttempts::new(&context) + .complete(vec![PeerOutcome::new( + rows[0].clone(), + Ok(WorkOutcome::Complete), + )]) + .await + .unwrap(); + } + let _ = env.checks.send(checks); + WorkOutcome::Complete + }) + }); + let mut run = Run::start(builder, db, 4, script).await; + run.assert_checks().await; + run.settled(run.coordinator).await; + run.stop().await; +} + +/// Outcomes are accepted all or none: each must be for one of this +/// coordinator's peers at its claimed attempt, once, without an earlier +/// outcome. Of two concurrent submissions for one peer, exactly one wins. +async fn assert_outcomes_are_checked(builder: riverqueue::ClientBuilder, db: Db) { + let script: Script = Arc::new(|context, _row, env| { + Box::pin(async move { + let mut checks = Checks::new(); + let [p1, p2, p3] = [env.peers[0], env.peers[1], env.peers[2]]; + let rows = env.claim(&context, Take::Claim, &[p1, p2]).await.unwrap(); + let row = |id: i64| rows.iter().find(|row| row.id == id).unwrap().clone(); + let peers = PeerAttempts::new(&context); + let outsider = context.client().unwrap().jobs().get(p3).await.unwrap(); + let not_peer = peers + .complete(vec![ + PeerOutcome::new(row(p1), Ok(WorkOutcome::Complete)), + PeerOutcome::new(outsider, Ok(WorkOutcome::Complete)), + ]) + .await; + check( + &mut checks, + "not a peer rejected", + not_peer.is_err_and(|error| peer_error(&error, "isn't a peer")), + ); + let mut wrong_attempt = row(p1); + wrong_attempt.attempt += 1; + let wrong = peers + .complete(vec![PeerOutcome::new( + wrong_attempt, + Ok(WorkOutcome::Complete), + )]) + .await; + check( + &mut checks, + "wrong attempt rejected", + wrong.is_err_and(|error| peer_error(&error, "isn't the peer attempt")), + ); + let mut other_client = row(p1); + other_client.attempted_by.push("another-client".to_owned()); + let foreign = peers + .complete(vec![PeerOutcome::new( + other_client, + Ok(WorkOutcome::Complete), + )]) + .await; + check( + &mut checks, + "another client's attempt rejected", + foreign.is_err_and(|error| peer_error(&error, "isn't the peer attempt")), + ); + let twice = peers + .complete(vec![ + PeerOutcome::new(row(p1), Ok(WorkOutcome::Complete)), + PeerOutcome::new(row(p1), Ok(WorkOutcome::Complete)), + ]) + .await; + check( + &mut checks, + "two outcomes rejected", + twice.is_err_and(|error| peer_error(&error, "two outcomes")), + ); + // Nothing above was accepted, so p1 still takes an outcome. + let (first, second) = tokio::join!( + peers.complete(vec![PeerOutcome::new(row(p2), Ok(WorkOutcome::Complete))]), + peers.complete(vec![PeerOutcome::new(row(p2), Ok(WorkOutcome::Complete))]), + ); + check( + &mut checks, + "exactly one concurrent outcome wins", + first.is_ok() != second.is_ok(), + ); + let late = [first, second] + .into_iter() + .find_map(Result::err) + .is_some_and(|error| peer_error(&error, "already has an outcome")); + check(&mut checks, "the other already had one", late); + check( + &mut checks, + "p1 still takes an outcome", + peers + .complete(vec![PeerOutcome::new(row(p1), Ok(WorkOutcome::Complete))]) + .await + .is_ok(), + ); + let _ = env.checks.send(checks); + WorkOutcome::Complete + }) + }); + let mut run = Run::start(builder, db, 3, script).await; + run.assert_checks().await; + for id in [run.env.peers[0], run.env.peers[1]] { + assert_eq!(run.settled(id).await.state, JobState::Completed); + } + run.stop().await; +} + +/// A peer the coordinator left without an outcome fails when the +/// coordinator ends on its own, before the coordinator's own outcome, and a +/// row River couldn't decode fails inside the claim. +async fn assert_missing_outcomes_fail(builder: riverqueue::ClientBuilder, db: Db) { + let script: Script = Arc::new(|context, _row, env| { + Box::pin(async move { + let mut checks = Checks::new(); + env.db.corrupt(env.peers[1]).await; + let rows = env.claim(&context, Take::Claim, &env.peers).await; + check( + &mut checks, + "only the decodable peer returned", + rows.as_ref() + .is_ok_and(|rows| rows.len() == 1 && rows[0].id == env.peers[0]), + ); + let _ = env.checks.send(checks); + WorkOutcome::Complete + }) + }); + let mut run = Run::start(builder, db.clone(), 2, script).await; + run.assert_checks().await; + let (p1, coordinator) = (run.env.peers[0], run.coordinator); + let events = run.events_for(&[p1, coordinator]).await; + let position = |id: i64| { + events + .iter() + .position(|event| event.as_job().unwrap().job.id == id) + .unwrap() + }; + assert!(position(p1) < position(coordinator), "peers settle first"); + let peer = run.client.jobs().get(p1).await.unwrap(); + assert!(matches!( + peer.state, + JobState::Available | JobState::Retryable + )); + assert!( + peer.errors[0].error.contains("ended without an outcome"), + "{:?}", + peer.errors + ); + // The undecodable peer's attempt failed as well; it can't be read back. + db.exec( + "UPDATE {table} SET metadata = '{}'::jsonb WHERE id = $1", + "UPDATE river_job SET tags = jsonb('[]') WHERE id = ?", + run.env.peers[1], + ) + .await; + let undecodable = run.client.jobs().get(run.env.peers[1]).await.unwrap(); + assert!( + undecodable.errors[0].error.contains("couldn't be decoded"), + "{:?}", + undecodable.errors + ); + run.stop().await; +} + +/// A hard stop interrupts peers without an outcome, like the coordinator, +/// while a remote cancellation of the coordinator fails them. +async fn assert_stops_interrupt_and_cancellations_fail( + builder: impl Fn() -> riverqueue::ClientBuilder, + db: Db, +) { + for remote in [false, true] { + let script: Script = Arc::new(|context, _row, env| { + Box::pin(async move { + env.claim(&context, Take::Claim, &env.peers).await.unwrap(); + env.signals.raise("claimed"); + context.cancellation_token().cancelled().await; + WorkOutcome::Complete + }) + }); + let mut run = Run::start(builder(), db.clone(), 1, script).await; + run.env.signals.wait("claimed").await; + let peer = run.env.peers[0]; + if remote { + // The coordinator hears of its cancellation through the + // notification listener, which must be listening first. + run.handle.wait_ready().await.unwrap(); + run.client.jobs().cancel(run.coordinator).await.unwrap(); + let row = run.settled(peer).await; + assert!( + matches!(row.state, JobState::Available | JobState::Retryable), + "{row:?}" + ); + assert!(row.errors[0].error.contains("ended without an outcome")); + run.stop().await; + } else { + tokio::time::timeout(WAIT, run.handle.shutdown_now()) + .await + .expect("client stops") + .unwrap(); + let row = run.client.jobs().get(peer).await.unwrap(); + assert_eq!(row.state, JobState::Available, "{row:?}"); + assert_eq!(row.attempt, 0); + assert!(row.errors.is_empty(), "{:?}", row.errors); + } + } +} + +/// A soft stop doesn't end a coordinator's claims: a coordinator that starts +/// claiming after its producer stopped fetching still claims and completes +/// its peers, and the stop resolves only after they persisted. A claim after +/// a hard stop or a remote cancellation of the coordinator is refused and +/// leaves the peers untouched. +async fn assert_soft_stops_keep_claims_open( + builder: impl Fn() -> riverqueue::ClientBuilder, + db: Db, +) { + let script: Script = Arc::new(|context, _row, env| { + Box::pin(async move { + let mut checks = Checks::new(); + env.signals.raise("running"); + env.gate.notified().await; + let claimed = env.claim(&context, Take::Claim, &env.peers).await; + check( + &mut checks, + "claim during a soft stop succeeds", + claimed + .as_ref() + .is_ok_and(|rows| rows.len() == env.peers.len()), + ); + let outcomes = claimed + .unwrap_or_default() + .into_iter() + .map(|row| PeerOutcome::new(row, Ok(WorkOutcome::Complete))) + .collect(); + check( + &mut checks, + "peers complete during a soft stop", + PeerAttempts::new(&context).complete(outcomes).await.is_ok(), + ); + let _ = env.checks.send(checks); + WorkOutcome::Complete + }) + }); + let mut run = Run::start(builder(), db.clone(), 2, script).await; + run.env.signals.wait("running").await; + // The producer stops fetching at once; only then does the coordinator + // start claiming. + run.handle.stopper().stop(); + run.env.gate.notify_one(); + tokio::time::timeout(WAIT, run.handle.wait()) + .await + .expect("client stops") + .unwrap(); + run.assert_checks().await; + for id in run.env.peers.iter().copied().chain([run.coordinator]) { + let row = run.client.jobs().get(id).await.unwrap(); + assert_eq!(row.state, JobState::Completed, "{row:?}"); + } + assert_eq!(*run.pilot.finished.lock().unwrap(), vec![run.coordinator]); + + for remote in [false, true] { + let script: Script = Arc::new(|context, _row, env| { + Box::pin(async move { + let mut checks = Checks::new(); + env.signals.raise("running"); + context.cancellation_token().cancelled().await; + check( + &mut checks, + "claim after cancellation is refused", + env.claim(&context, Take::Claim, &env.peers) + .await + .is_err_and(|error| peer_error(&error, "cancelled")), + ); + let _ = env.checks.send(checks); + WorkOutcome::Complete + }) + }); + let mut run = Run::start(builder(), db.clone(), 1, script).await; + run.env.signals.wait("running").await; + if remote { + // The coordinator hears of its cancellation through the + // notification listener, which must be listening first. + run.handle.wait_ready().await.unwrap(); + run.client.jobs().cancel(run.coordinator).await.unwrap(); + run.assert_checks().await; + run.settled(run.coordinator).await; + } else { + run.handle.stopper().stop_now(); + run.assert_checks().await; + } + let peer = run.client.jobs().get(run.env.peers[0]).await.unwrap(); + assert_eq!(peer.state, JobState::Available, "{peer:?}"); + assert_eq!(peer.attempt, 0); + run.stop().await; + } +} + +/// Once the coordinator's attempt ended, its peer operations are refused. +/// A claim still in flight when it ends is waited for: its rows become +/// peers and then fail like any peer left without an outcome. A claim +/// whose coordinator is cancelled before commit rolls back. +async fn assert_coordinator_lifetime_bounds_operations( + builder: impl Fn() -> riverqueue::ClientBuilder, + db: Db, +) { + // In-flight claim at exit, then late operations. + let script: Script = Arc::new(|context, _row, env| { + Box::pin(async move { + *env.kept.lock().unwrap() = Some(context.clone()); + let claim_env = env.clone(); + let claim_context = context.clone(); + tokio::spawn(async move { + let claimed = claim_env + .claim(&claim_context, Take::ClaimAfterGate, &claim_env.peers) + .await; + claim_env.signals.raise(if claimed.is_ok() { + "claim committed" + } else { + "claim failed" + }); + }); + env.signals.wait("claim holds its rows").await; + env.signals.raise("coordinator returns"); + WorkOutcome::Complete + }) + }); + let run = Run::start(builder(), db.clone(), 1, script).await; + run.env.signals.wait("coordinator returns").await; + // The coordinator's end waits for the claim it accepted. + run.env.gate.notify_one(); + run.env.signals.wait("claim committed").await; + let peer = run.settled(run.env.peers[0]).await; + assert!(peer.errors[0].error.contains("ended without an outcome")); + assert_eq!( + run.settled(run.coordinator).await.state, + JobState::Completed + ); + let kept = run.env.kept.lock().unwrap().clone().unwrap(); + let late = run.env.claim(&kept, Take::Claim, &run.env.peers).await; + assert!(late.is_err_and(|error| peer_error(&error, "running attempt"))); + let late = PeerAttempts::new(&kept).complete(Vec::new()).await; + assert!(late.is_err_and(|error| peer_error(&error, "running attempt"))); + run.stop().await; + + // Cancelled before commit. + let script: Script = Arc::new(|context, _row, env| { + Box::pin(async move { + let mut checks = Checks::new(); + let claimed = env + .claim(&context, Take::ClaimUntilCancelled("claiming"), &env.peers) + .await; + check( + &mut checks, + "cancelled claim fails", + claimed.is_err_and(|error| peer_error(&error, "cancelled")), + ); + let _ = env.checks.send(checks); + WorkOutcome::Complete + }) + }); + let mut run = Run::start(builder(), db, 1, script).await; + run.env.signals.wait("claiming").await; + // A hard stop cancels the coordinator without a write, which SQLite + // couldn't take while the claim holds its write lock. + run.handle.stopper().stop_now(); + run.assert_checks().await; + let peer = run.client.jobs().get(run.env.peers[0]).await.unwrap(); + assert_eq!(peer.state, JobState::Available); + assert_eq!(peer.attempt, 0); + run.stop().await; +} + +/// An operation whose future is dropped before its outcomes reached the +/// completer, as by a `select!` or `timeout` around it or an aborted task, +/// leaves its peers without an outcome, so the coordinator's end still gives +/// them one instead of leaving them running. +async fn assert_dropped_operations_leave_peers_without_outcomes( + builder: riverqueue::ClientBuilder, + db: Db, +) { + let script: Script = Arc::new(|context, _row, env| { + Box::pin(async move { + let mut checks = Checks::new(); + let rows = env.claim(&context, Take::Claim, &env.peers).await.unwrap(); + // The error handler yields, so the completion is pending when the + // other branch wins and the completion's future is dropped. + let dropped = tokio::select! { + biased; + _ = PeerAttempts::new(&context) + .complete(vec![PeerOutcome::new(rows[0].clone(), Err("dropped".into()))]) => false, + () = std::future::ready(()) => true, + }; + check(&mut checks, "completion dropped", dropped); + let _ = env.checks.send(checks); + WorkOutcome::Complete + }) + }); + let mut run = Run::start(builder, db, 1, script).await; + run.assert_checks().await; + let peer = run.settled(run.env.peers[0]).await; + assert!( + peer.errors + .last() + .is_some_and(|error| error.error.contains("ended without an outcome")), + "{peer:?}" + ); + run.stop().await; +} + +/// A claim whose commit fails, or whose future is dropped while it commits, +/// gives its rows back, so the same coordinator can claim them again. +async fn assert_failed_commits_release_reservations(builder: riverqueue::ClientBuilder, db: Db) { + db.fail_peer_commits().await; + let script: Script = Arc::new(|context, _row, env| { + Box::pin(async move { + let mut checks = Checks::new(); + let failed = env.claim(&context, Take::Claim, &env.peers).await; + check( + &mut checks, + "commit failed", + matches!(failed, Err(riverqueue::Error::Database(_))), + ); + env.db.allow_peer_commits().await; + let again = env.claim(&context, Take::Claim, &env.peers).await; + check( + &mut checks, + "claimed again after the failed commit", + again.as_ref().is_ok_and(|rows| rows.len() == 1), + ); + if let Ok(rows) = again { + PeerAttempts::new(&context) + .complete(vec![PeerOutcome::new( + rows[0].clone(), + Ok(WorkOutcome::Complete), + )]) + .await + .unwrap(); + } + let _ = env.checks.send(checks); + WorkOutcome::Complete + }) + }); + let mut run = Run::start(builder, db, 1, script).await; + run.assert_checks().await; + assert_eq!( + run.settled(run.env.peers[0]).await.state, + JobState::Completed + ); + run.stop().await; +} + +/// An attempt abandoned after its worker outlived an abort leaves its peers +/// to the rescuer but stops owning them, so a later attempt of the same +/// client can claim them. With `racing_claim`, a claim still committing when +/// the attempt is abandoned is refused once it commits and gives its rows +/// back too; that needs a commit slow enough to abandon the attempt during +/// it. +async fn assert_abandoned_attempts_release_peers( + builder: riverqueue::ClientBuilder, + db: Db, + racing_claim: bool, +) { + let (unblock, blocked) = std::sync::mpsc::channel::<()>(); + let blocked = Arc::new(Mutex::new(Some(blocked))); + let runs = Arc::new(std::sync::atomic::AtomicUsize::new(0)); + let script: Script = Arc::new(move |context, _row, env| { + let blocked = Arc::clone(&blocked); + let first = runs.fetch_add(1, std::sync::atomic::Ordering::SeqCst) == 0; + Box::pin(async move { + if first { + env.claim(&context, Take::Claim, &env.peers[..1]) + .await + .unwrap(); + if racing_claim { + let claim_env = env.clone(); + let claim_context = context.clone(); + tokio::spawn(async move { + let claimed = claim_env + .claim( + &claim_context, + Take::ClaimAndSignal("claim commits"), + &claim_env.peers[1..], + ) + .await; + claim_env.signals.raise( + if claimed.is_err_and(|error| peer_error(&error, "running attempt")) { + "late claim refused" + } else { + "late claim accepted" + }, + ); + }); + env.signals.wait("claim commits").await; + } + env.signals.raise("coordinator blocks"); + // Blocks the worker's thread, so neither cancellation nor an + // abort can end it. + let receiver = blocked.lock().unwrap().take().unwrap(); + let _ = receiver.recv(); + return WorkOutcome::Complete; + } + let mut checks = Checks::new(); + let rows = env.claim(&context, Take::Read, &env.peers).await; + check( + &mut checks, + &format!("claimed again by a later attempt: {rows:?}"), + rows.as_ref() + .is_ok_and(|rows| rows.len() == env.peers.len()), + ); + if let Ok(rows) = rows { + PeerAttempts::new(&context) + .complete( + rows.into_iter() + .map(|row| PeerOutcome::new(row, Ok(WorkOutcome::Complete))) + .collect(), + ) + .await + .unwrap(); + } + let _ = env.checks.send(checks); + WorkOutcome::Complete + }) + }); + let peers = if racing_claim { 2 } else { 1 }; + let mut run = Run::start( + builder.job_stuck_threshold(Duration::from_millis(50)), + db, + peers, + script, + ) + .await; + run.env.signals.wait("coordinator blocks").await; + tokio::time::timeout(WAIT, run.handle.shutdown_now()) + .await + .expect("client stops") + .unwrap(); + if racing_claim { + run.env.signals.wait("late claim refused").await; + } + unblock.send(()).unwrap(); + + run.client.insert(CoordinatorArgs {}).await.unwrap(); + run.handle = run.client.start().unwrap(); + run.assert_checks().await; + for id in run.env.peers.clone() { + assert_eq!(run.settled(id).await.state, JobState::Completed); + } + run.stop().await; +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +mod postgres { + use riverqueue::database::PostgresDatabase; + + use super::*; + use crate::support::PostgresSchema; + + fn builder(schema: &PostgresSchema) -> riverqueue::ClientBuilder { + Client::builder( + PostgresDatabase::new(schema.pool.clone()).with_schema(schema.schema.clone()), + ) + } + + fn db(schema: &PostgresSchema) -> Db { + Db::Postgres( + schema.pool.clone(), + schema.table("river_job"), + schema.table("peer_test_function"), + ) + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn abandoned_attempts_release_peers() { + let schema = PostgresSchema::new("peer_abandoned").await; + assert_abandoned_attempts_release_peers(builder(&schema), db(&schema), false).await; + schema.cleanup().await; + } + + /// Only PostgreSQL can hold a commit open long enough to abandon the + /// attempt while its claim commits. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn claims_committing_while_abandoned_release_peers() { + let schema = PostgresSchema::new("peer_abandon_race").await; + let db = db(&schema); + slow_peer_commits(&db).await; + assert_abandoned_attempts_release_peers(builder(&schema), db, true).await; + schema.cleanup().await; + } + + /// Makes each commit that changes a peer job take a second. + async fn slow_peer_commits(db: &Db) { + db.raw( + "CREATE FUNCTION {function}() RETURNS trigger LANGUAGE plpgsql AS $$ \ + BEGIN IF NEW.kind = 'peer_job' AND NEW.state = 'running' THEN \ + PERFORM pg_sleep(1); END IF; RETURN NULL; END $$; \ + CREATE CONSTRAINT TRIGGER slow_peer_commit AFTER UPDATE ON {table} \ + DEFERRABLE INITIALLY DEFERRED FOR EACH ROW EXECUTE FUNCTION {function}();", + "", + ) + .await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn dropped_operations_leave_peers_without_outcomes() { + let schema = PostgresSchema::new("peer_dropped").await; + assert_dropped_operations_leave_peers_without_outcomes(builder(&schema), db(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn failed_commits_release_reservations() { + let schema = PostgresSchema::new("peer_commit").await; + assert_failed_commits_release_reservations(builder(&schema), db(&schema)).await; + schema.cleanup().await; + } + + /// Only PostgreSQL can hold a commit open long enough to drop the claim + /// while it commits. + #[tokio::test(flavor = "multi_thread")] + async fn claims_dropped_while_committing_release_reservations() { + let schema = PostgresSchema::new("peer_commit_drop").await; + let db = db(&schema); + slow_peer_commits(&db).await; + let script: Script = Arc::new(|context, _row, env| { + Box::pin(async move { + let mut checks = Checks::new(); + let dropped = tokio::time::timeout( + Duration::from_millis(300), + env.claim(&context, Take::Claim, &env.peers), + ) + .await; + check( + &mut checks, + "claim dropped while committing", + dropped.is_err(), + ); + // Waits for the commit the dropped claim started. + env.db + .raw("DROP TRIGGER slow_peer_commit ON {table}", "") + .await; + let state = context + .client() + .unwrap() + .jobs() + .get(env.peers[0]) + .await + .unwrap() + .state; + let take = if state == JobState::Running { + Take::Read + } else { + Take::Claim + }; + let again = env.claim(&context, take, &env.peers).await; + check( + &mut checks, + &format!("claimed again: {again:?}"), + again.as_ref().is_ok_and(|rows| rows.len() == 1), + ); + if let Ok(rows) = again { + PeerAttempts::new(&context) + .complete(vec![PeerOutcome::new( + rows[0].clone(), + Ok(WorkOutcome::Complete), + )]) + .await + .unwrap(); + } + let _ = env.checks.send(checks); + WorkOutcome::Complete + }) + }); + let mut run = Run::start(builder(&schema), db, 1, script).await; + run.assert_checks().await; + run.stop().await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn claims_are_checked() { + let schema = PostgresSchema::new("peer_claims").await; + assert_claims_are_checked(builder(&schema), db(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn coordinator_lifetime_bounds_operations() { + let schema = PostgresSchema::new("peer_lifetime").await; + assert_coordinator_lifetime_bounds_operations(|| builder(&schema), db(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn missing_outcomes_fail() { + let schema = PostgresSchema::new("peer_missing").await; + assert_missing_outcomes_fail(builder(&schema), db(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn outcomes_are_checked() { + let schema = PostgresSchema::new("peer_outcomes").await; + assert_outcomes_are_checked(builder(&schema), db(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn outcomes_use_the_completion_pipeline() { + let schema = PostgresSchema::new("peer_pipeline").await; + assert_outcomes_use_the_completion_pipeline(builder(&schema), db(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn soft_stops_keep_claims_open() { + let schema = PostgresSchema::new("peer_soft_stop").await; + assert_soft_stops_keep_claims_open(|| builder(&schema), db(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn stops_interrupt_and_cancellations_fail() { + let schema = PostgresSchema::new("peer_stop").await; + assert_stops_interrupt_and_cancellations_fail(|| builder(&schema), db(&schema)).await; + schema.cleanup().await; + } +} + +#[cfg(feature = "sqlite")] +mod sqlite { + use super::*; + use crate::support::{sqlite_cleanup, sqlite_file_pool}; + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn abandoned_attempts_release_peers() { + let (pool, path) = sqlite_file_pool(4).await; + assert_abandoned_attempts_release_peers( + Client::builder(pool.clone()), + Db::Sqlite(pool.clone()), + false, + ) + .await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn dropped_operations_leave_peers_without_outcomes() { + let (pool, path) = sqlite_file_pool(4).await; + assert_dropped_operations_leave_peers_without_outcomes( + Client::builder(pool.clone()), + Db::Sqlite(pool.clone()), + ) + .await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn failed_commits_release_reservations() { + let (pool, path) = sqlite_file_pool(4).await; + assert_failed_commits_release_reservations( + Client::builder(pool.clone()), + Db::Sqlite(pool.clone()), + ) + .await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn claims_are_checked() { + let (pool, path) = sqlite_file_pool(4).await; + assert_claims_are_checked(Client::builder(pool.clone()), Db::Sqlite(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn coordinator_lifetime_bounds_operations() { + let (pool, path) = sqlite_file_pool(4).await; + assert_coordinator_lifetime_bounds_operations( + || Client::builder(pool.clone()), + Db::Sqlite(pool.clone()), + ) + .await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn missing_outcomes_fail() { + let (pool, path) = sqlite_file_pool(4).await; + assert_missing_outcomes_fail(Client::builder(pool.clone()), Db::Sqlite(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn outcomes_are_checked() { + let (pool, path) = sqlite_file_pool(4).await; + assert_outcomes_are_checked(Client::builder(pool.clone()), Db::Sqlite(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn outcomes_use_the_completion_pipeline() { + let (pool, path) = sqlite_file_pool(4).await; + assert_outcomes_use_the_completion_pipeline( + Client::builder(pool.clone()), + Db::Sqlite(pool.clone()), + ) + .await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn soft_stops_keep_claims_open() { + let (pool, path) = sqlite_file_pool(4).await; + assert_soft_stops_keep_claims_open( + || Client::builder(pool.clone()), + Db::Sqlite(pool.clone()), + ) + .await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn stops_interrupt_and_cancellations_fail() { + let (pool, path) = sqlite_file_pool(4).await; + assert_stops_interrupt_and_cancellations_fail( + || Client::builder(pool.clone()), + Db::Sqlite(pool.clone()), + ) + .await; + sqlite_cleanup(pool, path).await; + } +} diff --git a/rust/riverqueue/tests/poll_only_cancellation.rs b/rust/riverqueue/tests/poll_only_cancellation.rs new file mode 100644 index 000000000..bb6724fc0 --- /dev/null +++ b/rust/riverqueue/tests/poll_only_cancellation.rs @@ -0,0 +1,138 @@ +//! Cancellation of running jobs in clients without notifications, on every +//! backend. +//! +//! PostgreSQL tests run in a unique schema and fail rather than skip when +//! `RIVER_RUST_DATABASE_URL` is unset; SQLite tests use a temporary file. + +#![cfg(any(all(feature = "postgres", river_postgres_tests), feature = "sqlite"))] + +mod support; + +use std::{convert::Infallible, sync::Arc, time::Duration}; + +use riverqueue::{ + __private::Database, Client, EventKind, Job, JobArgs, JobState, QueueConfig, WorkContext, + WorkOutcome, WorkerRegistry, +}; +use serde::{Deserialize, Serialize}; +use tokio::sync::Semaphore; + +/// Every wait in these tests is bounded by this timeout. It covers a few of +/// the two-second polls for cancellation requests. +const TIMEOUT: Duration = Duration::from_secs(10); + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_poll_only_cancellation")] +struct WaitArgs {} + +/// A migrated database on one backend. +enum Backend { + #[cfg(all(feature = "postgres", river_postgres_tests))] + Postgres(support::PostgresSchema), + #[cfg(feature = "sqlite")] + Sqlite(sqlx::SqlitePool, std::path::PathBuf), +} + +impl Backend { + fn database(&self) -> Database { + match self { + #[cfg(all(feature = "postgres", river_postgres_tests))] + Self::Postgres(schema) => Database::from_source( + riverqueue::database::PostgresDatabase::new(schema.pool.clone()) + .with_schema(schema.schema.clone()), + ), + #[cfg(feature = "sqlite")] + Self::Sqlite(pool, _) => Database::from_source(pool.clone()), + } + } + + async fn cleanup(self) { + match self { + #[cfg(all(feature = "postgres", river_postgres_tests))] + Self::Postgres(schema) => schema.cleanup().await, + #[cfg(feature = "sqlite")] + Self::Sqlite(pool, path) => support::sqlite_cleanup(pool, path).await, + } + } +} + +/// A client without notifications cancels its running job once it polls the +/// cancellation another client requested, including while it's stopping and +/// waiting for that job. +async fn polls_for_remote_cancellation(backend: Backend) { + for while_stopping in [false, true] { + let started = Arc::new(Semaphore::new(0)); + let mut workers = WorkerRegistry::new(); + let worker_started = Arc::clone(&started); + workers + .register_fn(move |context: WorkContext, _job: Job| { + let started = Arc::clone(&worker_started); + async move { + started.add_permits(1); + context.cancellation_token().cancelled().await; + // Any outcome but completion becomes the cancellation. + Ok::<_, Infallible>(WorkOutcome::Snooze(Duration::from_hours(1))) + } + }) + .unwrap(); + let client = Client::builder(backend.database()) + .queue(riverqueue::QUEUE_DEFAULT, QueueConfig::new(1)) + .without_leader_election() + .without_notifications() + .workers(workers) + .build() + .unwrap(); + let other = Client::builder(backend.database()).build().unwrap(); + let mut events = client.subscribe(&[EventKind::JobCancelled]).unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + let id = other.insert(WaitArgs {}).await.unwrap().job.row.id; + tokio::time::timeout(TIMEOUT, started.acquire()) + .await + .expect("the job should start") + .unwrap() + .forget(); + if while_stopping { + run.stopper().stop(); + } + other.jobs().cancel(id).await.unwrap(); + + let event = tokio::time::timeout(TIMEOUT, events.recv()) + .await + .expect("the job should be cancelled") + .unwrap(); + let job = &event.as_job().expect("a job event").job; + assert_eq!(job.id, id); + assert_eq!(job.state, JobState::Cancelled, "{job:?}"); + tokio::time::timeout(TIMEOUT, run.shutdown()) + .await + .expect("the client should stop") + .unwrap(); + } + backend.cleanup().await; +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +mod postgres { + use super::*; + + #[tokio::test] + async fn polls_for_remote_cancellation() { + super::polls_for_remote_cancellation(Backend::Postgres( + support::PostgresSchema::new("river_poll_cancel").await, + )) + .await; + } +} + +#[cfg(feature = "sqlite")] +mod sqlite { + use super::*; + + #[tokio::test] + async fn polls_for_remote_cancellation() { + let (pool, path) = support::sqlite_file_pool(4).await; + super::polls_for_remote_cancellation(Backend::Sqlite(pool, path)).await; + } +} diff --git a/rust/riverqueue/tests/postgres.rs b/rust/riverqueue/tests/postgres.rs new file mode 100644 index 000000000..81e5204b5 --- /dev/null +++ b/rust/riverqueue/tests/postgres.rs @@ -0,0 +1,2102 @@ +#![cfg(all(feature = "postgres", river_postgres_tests))] + +mod support; + +use std::{ + convert::Infallible, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use async_trait::async_trait; +use riverqueue::__private::ClientBuilderExt; +use riverqueue::__private::{ + ClaimedJob, DatabaseConnection, JobSetStateParams, MaintenanceService, + MaintenanceServiceContext, Pilot, PilotError, PilotProducer, ProducerClaimContext, + ProducerClaimNext, ProducerStartContext, RuntimeService, RuntimeServiceContext, +}; +use riverqueue::{ + Client, EventKind, InsertBatch, InsertOpts, IntervalSchedule, Job, JobArgs, JobListOrderBy, + JobListParams, JobRow, JobState, JobUpdateParams, MaintenanceConfig, PeriodicJob, + PeriodicJobOpts, QueueConfig, QueueListParams, UniqueOpts, WorkContext, WorkError, WorkOutcome, + Worker, WorkerRegistry, WorkerTimeout, + database::{PostgresDatabase, PostgresReindexConfig, PostgresReindexSchedule}, +}; +use riverqueue_migrate::{Direction, MigrateOpts}; +use riverqueue_migrate::{MIGRATION_VERSION_LATEST, PostgresMigrator}; +use serde::{Deserialize, Serialize}; +use sqlx::{AssertSqlSafe, PgPool}; +use tokio_util::sync::CancellationToken; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_conformance_echo")] +struct EchoArgs { + message: String, +} + +struct EchoWorker; + +impl Worker for EchoWorker { + type Error = Infallible; + + fn work( + &self, + context: WorkContext, + job: Job, + ) -> impl Future> + Send { + assert_ne!(job.args.message, ""); + context + .record_output(serde_json::json!({"message": job.args.message})) + .unwrap(); + std::future::ready(Ok(WorkOutcome::Complete)) + } +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_conformance_cancel")] +struct CancelArgs {} + +struct CancelWorker; + +impl Worker for CancelWorker { + type Error = Infallible; + + async fn work( + &self, + context: WorkContext, + _job: Job, + ) -> Result { + context.cancellation_token().cancelled().await; + Ok(WorkOutcome::Cancel) + } +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_conformance_fail")] +struct FailArgs {} + +struct FailWorker; + +impl Worker for FailWorker { + type Error = std::io::Error; + + fn work( + &self, + _context: WorkContext, + _job: Job, + ) -> impl std::future::Future> + Send { + std::future::ready(Err(std::io::Error::other("intentional failure"))) + } +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_conformance_ignores_cancel")] +struct IgnoresCancelArgs {} + +struct IgnoresCancelWorker; + +impl Worker for IgnoresCancelWorker { + type Error = Infallible; + + async fn work( + &self, + _context: WorkContext, + _job: Job, + ) -> Result { + std::future::pending().await + } +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_rescue_default_timeout")] +struct RescueDefaultTimeoutArgs {} + +struct RescueDefaultTimeoutWorker; + +impl Worker for RescueDefaultTimeoutWorker { + type Error = Infallible; + + fn work( + &self, + _context: WorkContext, + _job: Job, + ) -> impl std::future::Future> + Send { + std::future::ready(Ok(WorkOutcome::Complete)) + } +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_rescue_disabled_timeout")] +struct RescueDisabledTimeoutArgs {} + +struct RescueDisabledTimeoutWorker; + +impl Worker for RescueDisabledTimeoutWorker { + type Error = Infallible; + + fn timeout(&self, _job: &Job) -> WorkerTimeout { + WorkerTimeout::Disabled + } + + fn work( + &self, + _context: WorkContext, + _job: Job, + ) -> impl std::future::Future> + Send { + std::future::ready(Ok(WorkOutcome::Complete)) + } +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_rescue_long_timeout")] +struct RescueLongTimeoutArgs {} + +struct RescueLongTimeoutWorker; + +impl Worker for RescueLongTimeoutWorker { + type Error = Infallible; + + fn timeout(&self, _job: &Job) -> WorkerTimeout { + WorkerTimeout::After(Duration::from_hours(1)) + } + + fn work( + &self, + _context: WorkContext, + _job: Job, + ) -> impl std::future::Future> + Send { + std::future::ready(Ok(WorkOutcome::Complete)) + } +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_rescue_retry_override")] +struct RescueRetryOverrideArgs {} + +struct RescueRetryOverrideWorker; + +impl Worker for RescueRetryOverrideWorker { + type Error = Infallible; + + fn next_retry( + &self, + _job: &Job, + _error: &WorkError, + _now: chrono::DateTime, + ) -> Option { + Some(Duration::from_hours(2)) + } + + fn work( + &self, + _context: WorkContext, + _job: Job, + ) -> impl std::future::Future> + Send { + std::future::ready(Ok(WorkOutcome::Complete)) + } +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_resumable_checkpoint")] +struct ResumableCheckpointArgs { + mode: String, +} + +#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)] +struct ResumableCursor { + offset: i64, +} + +struct ResumableCheckpointWorker { + cursor_values: Arc>>, + pool: PgPool, + validate_runs: Arc, +} + +impl Worker for ResumableCheckpointWorker { + type Error = riverqueue::BoxError; + + fn next_retry( + &self, + job: &Job, + _error: &WorkError, + _now: chrono::DateTime, + ) -> Option { + (job.args.mode == "cursor_retry").then_some(Duration::from_millis(500)) + } + + async fn work( + &self, + context: WorkContext, + job: Job, + ) -> Result { + match job.args.mode.as_str() { + "cursor_retry" => { + context + .resumable_step("validate", || async { + self.validate_runs.fetch_add(1, Ordering::SeqCst); + Ok::<_, riverqueue::Error>(()) + }) + .await?; + let cursor_context = context.clone(); + let cursor_values = Arc::clone(&self.cursor_values); + let attempt = job.row.attempt; + context + .resumable_step_with_cursor( + "process", + move |cursor: ResumableCursor| async move { + cursor_values.lock().unwrap().push(cursor.clone()); + if attempt == 1 { + cursor_context + .resumable_set_cursor(&ResumableCursor { offset: 42 })?; + return Err("intentional resumable cursor failure".into()); + } + Ok::<(), riverqueue::BoxError>(()) + }, + ) + .await?; + } + "commit_cursor" | "rollback_cursor" => { + let checkpoint_context = context.clone(); + let mode = job.args.mode.clone(); + let pool = self.pool.clone(); + context + .resumable_step_with_cursor("tx_cursor", move |_: ResumableCursor| async move { + let mut transaction = pool.begin().await?; + checkpoint_context + .resumable_set_step_cursor_tx( + &mut transaction, + &ResumableCursor { offset: 7 }, + ) + .await?; + if mode == "commit_cursor" { + transaction.commit().await?; + } else { + transaction.rollback().await?; + } + Ok::<_, riverqueue::Error>(()) + }) + .await?; + } + "commit_step" | "rollback_step" => { + let checkpoint_context = context.clone(); + let mode = job.args.mode.clone(); + let pool = self.pool.clone(); + context + .resumable_step("tx_step", move || async move { + let mut transaction = pool.begin().await?; + checkpoint_context + .resumable_set_step_tx(&mut transaction) + .await?; + if mode == "commit_step" { + transaction.commit().await?; + } else { + transaction.rollback().await?; + } + Ok::<_, riverqueue::Error>(()) + }) + .await?; + } + mode => { + return Err(format!("unknown test mode {mode}").into()); + } + } + Ok(WorkOutcome::Complete) + } +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_conformance_resumable")] +struct ResumableArgs {} + +#[derive(Clone, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_transactional")] +struct TransactionalArgs {} + +#[derive(Default)] +struct ResumableWorker { + first_runs: Arc, + second_runs: Arc, +} + +struct TransactionalWorker { + pool: PgPool, +} + +#[derive(Clone, Default)] +struct TestPilot { + completions: Arc, + fetches: Arc, + maintenance_starts: Arc, + maintenance_stops: Arc, + runtime_starts: Arc, + runtime_stops: Arc, +} + +#[async_trait] +impl Pilot for TestPilot { + fn intercepts_job_set_state(&self) -> bool { + true + } + + async fn start_producer( + &self, + _context: ProducerStartContext, + ) -> Result>, PilotError> { + Ok(Some(Box::new(CountingProducer { + fetches: Arc::clone(&self.fetches), + }))) + } + + async fn after_jobs_set_state( + &self, + connection: DatabaseConnection<'_>, + params: &JobSetStateParams, + ) -> Result<(), PilotError> { + self.completions + .fetch_add(params.jobs.len(), Ordering::SeqCst); + let table = params + .database + .postgres_schema() + .unwrap() + .qualify("river_job"); + let completed = params + .jobs + .iter() + .filter(|job| job.state == JobState::Completed) + .map(|job| job.id) + .collect::>(); + let sql = format!( + "UPDATE {table} SET metadata = metadata || '{{\"extension_handled\": true}}'::jsonb \ + WHERE id = ANY($1)" + ); + sqlx::query(AssertSqlSafe(sql)) + .bind(completed) + .execute(connection.into_postgres().unwrap()) + .await?; + Ok(()) + } + + fn maintenance_services(&self) -> Vec> { + vec![Arc::new(TestMaintenance { + starts: Arc::clone(&self.maintenance_starts), + stops: Arc::clone(&self.maintenance_stops), + })] + } + + fn runtime_services(&self) -> Vec> { + vec![Arc::new(TestRuntime { + starts: Arc::clone(&self.runtime_starts), + stops: Arc::clone(&self.runtime_stops), + })] + } +} + +/// Claims with River's standard claim in its own transaction and counts +/// claims. +struct CountingProducer { + fetches: Arc, +} + +#[async_trait] +impl PilotProducer for CountingProducer { + fn intercepts_claim(&self) -> bool { + true + } + + async fn claim( + &self, + context: ProducerClaimContext<'_>, + next: ProducerClaimNext<'_>, + ) -> Result, PilotError> { + self.fetches.fetch_add(1, Ordering::SeqCst); + let mut transaction = context.database.begin().await?; + let jobs = next.claim(transaction.connection()).await?; + transaction.commit().await?; + Ok(jobs) + } +} + +struct TestMaintenance { + starts: Arc, + stops: Arc, +} + +#[async_trait] +impl MaintenanceService for TestMaintenance { + async fn run(&self, context: MaintenanceServiceContext) -> Result<(), PilotError> { + let cancellation = context.term.token; + self.starts.fetch_add(1, Ordering::SeqCst); + cancellation.cancelled().await; + self.stops.fetch_add(1, Ordering::SeqCst); + Ok(()) + } +} + +struct TestRuntime { + starts: Arc, + stops: Arc, +} + +#[async_trait] +impl RuntimeService for TestRuntime { + async fn run(&self, context: RuntimeServiceContext) -> Result<(), PilotError> { + let cancellation = context.cancellation; + self.starts.fetch_add(1, Ordering::SeqCst); + cancellation.cancelled().await; + self.stops.fetch_add(1, Ordering::SeqCst); + Ok(()) + } +} + +impl Worker for ResumableWorker { + type Error = riverqueue::Error; + + async fn work( + &self, + context: WorkContext, + job: Job, + ) -> Result { + context + .resumable_step("first", || async { + self.first_runs.fetch_add(1, Ordering::SeqCst); + Ok::<_, std::io::Error>(()) + }) + .await?; + context + .resumable_step("second", || async { + self.second_runs.fetch_add(1, Ordering::SeqCst); + if job.row.attempt == 1 { + Err(std::io::Error::other("fail second step once")) + } else { + Ok(()) + } + }) + .await?; + Ok(WorkOutcome::Complete) + } +} + +impl Worker for TransactionalWorker { + type Error = riverqueue::Error; + + async fn work( + &self, + context: WorkContext, + _job: Job, + ) -> Result { + assert_eq!(context.client().unwrap().id(), "rust-maintenance-client"); + context + .metadata_set("transactional_completion", true) + .unwrap(); + let mut transaction = self.pool.begin().await?; + let completed = context.job_complete_tx(&mut transaction).await?; + assert_eq!(completed.state, JobState::Completed); + assert_eq!( + completed + .metadata + .get::("transactional_completion") + .unwrap(), + Some(true) + ); + transaction.commit().await?; + Ok(WorkOutcome::Complete) + } +} + +#[tokio::test] +async fn client_cancels_a_running_job() { + let database = support::PostgresSchema::current("rs_cancel_job").await; + let client = worker_client(&database.pool, ResumableWorker::default()); + let mut run_handle = client.start().unwrap(); + // Remote cancellation arrives by notification, so the listener must be + // subscribed before the job is cancelled. + run_handle.wait_ready().await.unwrap(); + + let cancelling = client.insert(CancelArgs {}).await.unwrap(); + wait_for_state(&client, cancelling.job.row.id, JobState::Running).await; + client.jobs().cancel(cancelling.job.row.id).await.unwrap(); + let cancelled = wait_for_state(&client, cancelling.job.row.id, JobState::Cancelled).await; + assert!(cancelled.finalized_at.is_some()); + assert!(cancelled.metadata.contains_key("cancel_attempted_at")); + + run_handle.shutdown().await.unwrap(); + database.cleanup().await; +} + +#[tokio::test] +async fn client_completes_a_job_with_output_and_event() { + let database = support::PostgresSchema::current("rs_complete_job").await; + let client = worker_client(&database.pool, ResumableWorker::default()); + + let inserted = client + .insert(EchoArgs { + message: "from Rust".to_owned(), + }) + .await + .unwrap(); + assert_eq!(inserted.job.row.state, JobState::Available); + + let mut completed_events = client.subscribe(&[EventKind::JobCompleted]).unwrap(); + let mut run_handle = client.start().unwrap(); + let row = wait_for_state(&client, inserted.job.row.id, JobState::Completed).await; + assert_eq!(row.attempt, 1); + assert_eq!(row.attempted_by, ["rust-conformance-client"]); + assert_eq!( + row.decode_output::().unwrap(), + Some(serde_json::json!({"message": "from Rust"})) + ); + loop { + let event = tokio::time::timeout(Duration::from_secs(1), completed_events.recv()) + .await + .unwrap() + .unwrap(); + if event.as_job().unwrap().job.id == inserted.job.row.id { + break; + } + } + + run_handle.shutdown().await.unwrap(); + database.cleanup().await; +} + +#[tokio::test] +async fn client_discards_a_failing_job_after_max_attempts() { + let database = support::PostgresSchema::current("rs_discard_fail").await; + let client = worker_client(&database.pool, ResumableWorker::default()); + let mut run_handle = client.start().unwrap(); + + let failed = client + .insert(FailArgs {}) + .opts(InsertOpts::default().with_max_attempts(1)) + .await + .unwrap(); + let failed = wait_for_state(&client, failed.job.row.id, JobState::Discarded).await; + assert_eq!(failed.errors.len(), 1); + assert_eq!(failed.errors[0].error, "intentional failure"); + + run_handle.shutdown().await.unwrap(); + database.cleanup().await; +} + +#[tokio::test] +async fn client_discards_a_job_of_an_unregistered_kind() { + let database = support::PostgresSchema::current("rs_unknown_kind").await; + let client = worker_client(&database.pool, ResumableWorker::default()); + let mut run_handle = client.start().unwrap(); + + let unknown_kind_id: i64 = sqlx::query_scalar( + "INSERT INTO river_job (args, kind, max_attempts) \ + VALUES ('{}'::jsonb, 'rust_unregistered_kind', 1) RETURNING id", + ) + .fetch_one(&database.pool) + .await + .unwrap(); + let unknown_kind = wait_for_state(&client, unknown_kind_id, JobState::Discarded).await; + assert_eq!(unknown_kind.attempt, 1); + assert_eq!(unknown_kind.errors.len(), 1); + assert_eq!( + unknown_kind.errors[0].error, + "job kind is not registered in the client's Workers bundle: rust_unregistered_kind" + ); + + run_handle.shutdown().await.unwrap(); + database.cleanup().await; +} + +#[tokio::test] +async fn client_restarts_after_shutdown() { + let database = support::PostgresSchema::current("rs_restart").await; + let client = worker_client(&database.pool, ResumableWorker::default()); + + let mut run_handle = client.start().unwrap(); + run_handle.wait_ready().await.unwrap(); + run_handle.shutdown().await.unwrap(); + let mut restarted_handle = client.start().unwrap(); + restarted_handle.wait_ready().await.unwrap(); + restarted_handle.shutdown().await.unwrap(); + + database.cleanup().await; +} + +#[tokio::test] +async fn client_resumes_resumable_steps_on_retry() { + let database = support::PostgresSchema::current("rs_resumable").await; + let resumable_worker = ResumableWorker::default(); + let resumable_first_runs = Arc::clone(&resumable_worker.first_runs); + let resumable_second_runs = Arc::clone(&resumable_worker.second_runs); + let client = worker_client(&database.pool, resumable_worker); + let mut run_handle = client.start().unwrap(); + + let resumable = client + .insert(ResumableArgs {}) + .opts(InsertOpts::default().with_max_attempts(2)) + .await + .unwrap(); + let resumable = wait_for_state(&client, resumable.job.row.id, JobState::Completed).await; + assert_eq!( + resumable + .metadata + .get::("river:resumable_step") + .unwrap() + .as_deref(), + Some("first") + ); + assert_eq!(resumable_first_runs.load(Ordering::SeqCst), 1); + assert_eq!(resumable_second_runs.load(Ordering::SeqCst), 2); + + run_handle.shutdown().await.unwrap(); + database.cleanup().await; +} + +#[tokio::test] +async fn complete_tx_requires_a_running_job_and_rolls_back() { + let database = support::PostgresSchema::current("rs_complete_tx").await; + let pool = database.pool.clone(); + let client = worker_client(&pool, ResumableWorker::default()); + + let non_running = client + .insert(EchoArgs { + message: "not running".to_owned(), + }) + .await + .unwrap(); + let mut transaction = pool.begin().await.unwrap(); + let error = client + .jobs() + .complete(non_running.job.row.id) + .tx(&mut transaction) + .await + .unwrap_err(); + assert!( + matches!( + error, + riverqueue::Error::JobNotRunning { + state: JobState::Available + } + ), + "{error}" + ); + assert!(matches!( + client.jobs().complete(i64::MAX).tx(&mut transaction).await, + Err(riverqueue::Error::NotFound(riverqueue::Record::Job( + i64::MAX + ))) + )); + transaction.rollback().await.unwrap(); + + sqlx::query("UPDATE river_job SET state = 'running' WHERE id = $1") + .bind(non_running.job.row.id) + .execute(&pool) + .await + .unwrap(); + let mut transaction = pool.begin().await.unwrap(); + let completed_then_rolled_back = client + .jobs() + .complete(non_running.job.row.id) + .tx(&mut transaction) + .await + .unwrap(); + assert_eq!(completed_then_rolled_back.state, JobState::Completed); + transaction.rollback().await.unwrap(); + assert_eq!( + client + .jobs() + .get(non_running.job.row.id) + .await + .unwrap() + .state, + JobState::Running + ); + + database.cleanup().await; +} + +#[tokio::test] +#[allow(clippy::too_many_lines)] +async fn concurrent_unique_inserts_return_the_conflicting_job() { + const INSERT_COUNT: usize = 32; + + let database = support::PostgresSchema::new("rs_unique_conc").await; + let pool = database.pool.clone(); + let schema = database.schema.clone(); + let client = Client::builder(PostgresDatabase::new(pool.clone()).with_schema(schema)) + .build() + .unwrap(); + + let all_states = vec![ + JobState::Available, + JobState::Cancelled, + JobState::Completed, + JobState::Discarded, + JobState::Pending, + JobState::Retryable, + JobState::Running, + JobState::Scheduled, + ]; + let fixed_scheduled_at = chrono::Utc::now() - chrono::Duration::minutes(1); + let cases = [ + ( + "by_args", + InsertOpts::default().with_unique(UniqueOpts::new().with_by_args(true)), + ), + ( + "by_args_and_queue", + InsertOpts::default() + .with_queue("unique_queue") + .with_unique(UniqueOpts::new().with_by_args(true).with_by_queue(true)), + ), + ( + "by_args_and_states", + InsertOpts::default().with_unique( + UniqueOpts::new() + .with_by_args(true) + .with_by_state(all_states), + ), + ), + ( + "by_args_and_period", + InsertOpts::default() + .with_scheduled_at(fixed_scheduled_at) + .with_unique( + UniqueOpts::new() + .with_by_args(true) + .with_by_period(Duration::from_mins(1)), + ), + ), + ]; + + for (message, opts) in cases { + let barrier = Arc::new(tokio::sync::Barrier::new(INSERT_COUNT)); + let mut tasks = tokio::task::JoinSet::new(); + for _ in 0..INSERT_COUNT { + let barrier = Arc::clone(&barrier); + let client = client.clone(); + let message = message.to_owned(); + let opts = opts.clone(); + tasks.spawn(async move { + barrier.wait().await; + client.insert(EchoArgs { message }).opts(opts).await + }); + } + + let mut results = Vec::with_capacity(INSERT_COUNT); + while let Some(result) = tasks.join_next().await { + results.push(result.unwrap().unwrap()); + } + let job_id = results[0].job.row.id; + assert!(results.iter().all(|result| result.job.row.id == job_id)); + assert_eq!( + results + .iter() + .filter(|result| !result.unique_skipped_as_duplicate) + .count(), + 1, + "unique case {message} should insert exactly one job" + ); + assert_eq!( + results + .iter() + .filter(|result| result.unique_skipped_as_duplicate) + .count(), + INSERT_COUNT - 1, + "unique case {message} should return the winner to every conflicting insert" + ); + } + + database.cleanup().await; +} + +#[tokio::test] +#[allow(clippy::too_many_lines)] +async fn insert_many_variants_preserve_order_and_transactionality() { + let database = support::PostgresSchema::new("rs_insert_many").await; + let pool = database.pool.clone(); + let schema = database.schema.clone(); + let client = Client::builder(PostgresDatabase::new(pool.clone()).with_schema(schema.clone())) + .build() + .unwrap(); + let table = schema.qualify("river_job"); + + let empty_many = client + .insert_many(Vec::::new()) + .await + .unwrap_err(); + assert_eq!(empty_many.to_string(), "invalid job: no jobs to insert"); + let empty_batch = client.insert_batch(InsertBatch::new()).await.unwrap_err(); + assert_eq!(empty_batch.to_string(), "invalid job: no jobs to insert"); + let mut empty_transaction = pool.begin().await.unwrap(); + let empty_many_tx = client + .insert_many(Vec::::new()) + .tx(&mut empty_transaction) + .await + .unwrap_err(); + assert_eq!(empty_many_tx.to_string(), "invalid job: no jobs to insert"); + let empty_batch_tx = client + .insert_batch(InsertBatch::new()) + .tx(&mut empty_transaction) + .await + .unwrap_err(); + assert_eq!(empty_batch_tx.to_string(), "invalid job: no jobs to insert"); + empty_transaction.commit().await.unwrap(); + + let past_scheduled_at = chrono::Utc::now() - chrono::Duration::minutes(1); + let ordered = client + .insert_many([ + ( + EchoArgs { + message: "ordered-one".to_owned(), + }, + InsertOpts::default(), + ), + ( + EchoArgs { + message: "ordered-two".to_owned(), + }, + InsertOpts::default(), + ), + ( + EchoArgs { + message: "ordered-past-scheduled".to_owned(), + }, + InsertOpts::default().with_scheduled_at(past_scheduled_at), + ), + ]) + .await + .unwrap(); + assert_eq!( + ordered + .iter() + .map(|result| result.job.args.message.as_str()) + .collect::>(), + ["ordered-one", "ordered-two", "ordered-past-scheduled"] + ); + assert!( + ordered + .windows(2) + .all(|pair| pair[0].job.row.id < pair[1].job.row.id) + ); + assert_eq!(ordered[2].job.row.state, JobState::Scheduled); + + let defaults = client + .insert_many([ + EchoArgs { + message: "default-one".to_owned(), + }, + EchoArgs { + message: "default-two".to_owned(), + }, + ]) + .await + .unwrap(); + assert_eq!(defaults.len(), 2); + + let mut heterogeneous = InsertBatch::new(); + heterogeneous + .push(EchoArgs { + message: "heterogeneous".to_owned(), + }) + .push_with( + CancelArgs {}, + InsertOpts::default().with_queue("heterogeneous-queue"), + ); + let heterogeneous = client.insert_batch(heterogeneous).await.unwrap(); + assert_eq!(heterogeneous.len(), 2); + assert_eq!(heterogeneous[0].job.kind, EchoArgs::KIND); + assert_eq!(heterogeneous[1].job.kind, CancelArgs::KIND); + assert_eq!(heterogeneous[1].job.queue, "heterogeneous-queue"); + + let time_without_states = client + .jobs() + .list( + JobListParams::default() + .ids(ordered.iter().map(|result| result.job.row.id)) + .order_by(JobListOrderBy::Time), + ) + .await + .unwrap() + .jobs; + assert_eq!( + time_without_states + .iter() + .map(|row| row.id) + .collect::>(), + [ + ordered[2].job.row.id, + ordered[0].job.row.id, + ordered[1].job.row.id, + ] + ); + let finalized_without_states = client + .jobs() + .list(JobListParams::default().order_by(JobListOrderBy::FinalizedAt)) + .await; + assert!(matches!( + finalized_without_states, + Err(riverqueue::Error::InvalidJob(_)) + )); + + // Like River Go, whose single upsert PostgreSQL refuses, a batch may not + // repeat a unique key, and nothing of it is written. + let unique_opts = InsertOpts::default().with_unique(UniqueOpts::new().with_by_args(true)); + let repeated = client + .insert_many([ + ( + EchoArgs { + message: "unique-batch".to_owned(), + }, + unique_opts.clone(), + ), + ( + EchoArgs { + message: "unique-batch".to_owned(), + }, + unique_opts.clone(), + ), + ]) + .await + .unwrap_err(); + assert!( + matches!(repeated, riverqueue::Error::InvalidJob(_)), + "{repeated:?}" + ); + let unique = client + .insert_many([ + ( + EchoArgs { + message: "unique-batch".to_owned(), + }, + unique_opts.clone(), + ), + ( + EchoArgs { + message: "unique-batch-other".to_owned(), + }, + unique_opts.clone(), + ), + ]) + .await + .unwrap(); + assert!(!unique[0].unique_skipped_as_duplicate); + assert!(!unique[1].unique_skipped_as_duplicate); + let duplicate = client + .insert(EchoArgs { + message: "unique-batch".to_owned(), + }) + .opts(unique_opts.clone()) + .await + .unwrap(); + assert_eq!(duplicate.job.row.id, unique[0].job.row.id); + assert!(duplicate.unique_skipped_as_duplicate); + + let mut transaction = pool.begin().await.unwrap(); + let rolled_back = client + .insert_many(["tx-rollback-one", "tx-rollback-two"].map(|message| { + ( + EchoArgs { + message: message.to_owned(), + }, + InsertOpts::default().with_tags(["tx-rollback"]), + ) + })) + .tx(&mut transaction) + .await + .unwrap(); + assert_eq!(rolled_back[0].job.args.message, "tx-rollback-one"); + assert_eq!(rolled_back[1].job.args.message, "tx-rollback-two"); + transaction.rollback().await.unwrap(); + let rolled_back_count: i64 = sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT count(*) FROM {table} WHERE 'tx-rollback' = ANY(tags)" + ))) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(rolled_back_count, 0); + + let mut transaction = pool.begin().await.unwrap(); + client + .insert_many(["tx-commit-one", "tx-commit-two"].map(|message| { + ( + EchoArgs { + message: message.to_owned(), + }, + InsertOpts::default().with_tags(["tx-commit"]), + ) + })) + .tx(&mut transaction) + .await + .unwrap(); + transaction.commit().await.unwrap(); + let committed_messages: Vec = sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT args ->> 'message' FROM {table} WHERE 'tx-commit' = ANY(tags) ORDER BY id" + ))) + .fetch_all(&pool) + .await + .unwrap(); + assert_eq!(committed_messages, ["tx-commit-one", "tx-commit-two"]); + + let invalid_batch = client + .insert_many([ + ( + EchoArgs { + message: "atomic-valid".to_owned(), + }, + InsertOpts::default().with_tags(["atomic-ordinary"]), + ), + ( + EchoArgs { + message: "atomic-invalid".to_owned(), + }, + InsertOpts::default().with_priority(0), + ), + ]) + .await; + assert!(invalid_batch.is_err()); + let atomic_ordinary_count: i64 = sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT count(*) FROM {table} WHERE 'atomic-ordinary' = ANY(tags)" + ))) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(atomic_ordinary_count, 0); + + // A batch rejected by validation fails before River writes anything, so + // the caller's transaction keeps its earlier work and can still commit. + let mut transaction = pool.begin().await.unwrap(); + client + .insert(EchoArgs { + message: "tx-validation-control".to_owned(), + }) + .tx(&mut transaction) + .await + .unwrap(); + let tx_validation = client + .insert_many([ + ( + EchoArgs { + message: "tx-validation-prefix".to_owned(), + }, + InsertOpts::default().with_tags(["tx-validation-batch"]), + ), + ( + EchoArgs { + message: "tx-validation-invalid".to_owned(), + }, + InsertOpts::default().with_priority(0), + ), + ]) + .tx(&mut transaction) + .await; + assert!(tx_validation.is_err()); + transaction.commit().await.unwrap(); + let ordinary_control_count: i64 = sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT count(*) FROM {table} WHERE args ->> 'message' = 'tx-validation-control'" + ))) + .fetch_one(&pool) + .await + .unwrap(); + let ordinary_batch_count: i64 = sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT count(*) FROM {table} WHERE 'tx-validation-batch' = ANY(tags)" + ))) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(ordinary_control_count, 1); + assert_eq!(ordinary_batch_count, 0); + + database.cleanup().await; +} + +#[tokio::test] +async fn insert_unique_by_args_returns_the_existing_job() { + let database = support::PostgresSchema::current("rs_unique_insert").await; + let client = worker_client(&database.pool, ResumableWorker::default()); + + let unique_options = InsertOpts::default().with_unique(UniqueOpts::new().with_by_args(true)); + let unique_first = client + .insert(EchoArgs { + message: "unique".to_owned(), + }) + .opts(unique_options.clone()) + .await + .unwrap(); + let unique_second = client + .insert(EchoArgs { + message: "unique".to_owned(), + }) + .opts(unique_options) + .await + .unwrap(); + assert_eq!(unique_first.job.row.id, unique_second.job.row.id); + assert!(unique_second.unique_skipped_as_duplicate); + + database.cleanup().await; +} + +#[tokio::test] +async fn job_admin_lists_updates_retries_and_deletes() { + let database = support::PostgresSchema::current("rs_job_admin").await; + let client = worker_client(&database.pool, ResumableWorker::default()); + let mut run_handle = client.start().unwrap(); + + let inserted = client + .insert(EchoArgs { + message: "from Rust".to_owned(), + }) + .await + .unwrap(); + let failed = client + .insert(FailArgs {}) + .opts(InsertOpts::default().with_max_attempts(1)) + .await + .unwrap(); + wait_for_state(&client, inserted.job.row.id, JobState::Completed).await; + let failed = wait_for_state(&client, failed.job.row.id, JobState::Discarded).await; + run_handle.shutdown().await.unwrap(); + + let listed = client + .jobs() + .list(JobListParams::default().kinds([EchoArgs::KIND])) + .await + .unwrap() + .jobs; + assert!(listed.iter().any(|row| row.id == inserted.job.row.id)); + let updated = client + .jobs() + .update( + inserted.job.row.id, + JobUpdateParams::default().output(serde_json::json!({"ok": true})), + ) + .await + .unwrap(); + assert_eq!( + updated.decode_output::().unwrap(), + Some(serde_json::json!({"ok": true})) + ); + + let retried = client.jobs().retry(failed.id).await.unwrap(); + assert_eq!(retried.state, JobState::Available); + assert_eq!(retried.max_attempts, 2); + let deleted = client.jobs().delete(retried.id).await.unwrap(); + assert_eq!(deleted.id, retried.id); + assert!(matches!( + client.jobs().get(retried.id).await, + Err(riverqueue::Error::NotFound(_)) + )); + + database.cleanup().await; +} + +#[tokio::test] +async fn local_queue_added_at_runtime_works_jobs() { + let database = support::PostgresSchema::current("rs_dynamic_queue").await; + let client = worker_client(&database.pool, ResumableWorker::default()); + let mut run_handle = client.start().unwrap(); + run_handle.wait_ready().await.unwrap(); + + client + .local_queues() + .add( + "dynamic", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .unwrap(); + let dynamic = client + .insert(EchoArgs { + message: "dynamic queue".to_owned(), + }) + .opts(InsertOpts::default().with_queue("dynamic")) + .await + .unwrap(); + wait_for_state(&client, dynamic.job.row.id, JobState::Completed).await; + client.local_queues().remove("dynamic").await.unwrap(); + run_handle.shutdown().await.unwrap(); + + // The removed queue's row stays behind alongside the configured one. + let queues = client + .queues() + .list(QueueListParams::default()) + .await + .unwrap(); + assert_eq!(queues.len(), 2); + assert!(queues.iter().any(|queue| queue.name == "dynamic")); + + database.cleanup().await; +} + +#[tokio::test] +async fn maintenance_cleans_old_jobs_and_queues_and_reindexes() { + let database = support::PostgresSchema::current("rs_cleanup").await; + let pool = database.pool.clone(); + + sqlx::raw_sql("CREATE INDEX rust_maintenance_reindex_idx ON river_job (id)") + .execute(&pool) + .await + .unwrap(); + let cleanup_job_ids = sqlx::query_scalar::<_, i64>( + "INSERT INTO river_job (args, finalized_at, kind, state) VALUES \ + ('{}'::jsonb, now() - interval '1 hour', 'cleanup_cancelled', 'cancelled'), \ + ('{}'::jsonb, now() - interval '1 hour', 'cleanup_completed', 'completed'), \ + ('{}'::jsonb, now() - interval '1 hour', 'cleanup_discarded', 'discarded') \ + RETURNING id", + ) + .fetch_all(&pool) + .await + .unwrap(); + sqlx::query( + "INSERT INTO river_queue (name, created_at, updated_at) \ + VALUES ('stale_cleanup_queue', now() - interval '2 hours', now() - interval '1 hour')", + ) + .execute(&pool) + .await + .unwrap(); + let reindex_file_node_before: i64 = sqlx::query_scalar( + "SELECT pg_relation_filenode('rust_maintenance_reindex_idx'::regclass)::bigint", + ) + .fetch_one(&pool) + .await + .unwrap(); + + let mut cleanup_workers = WorkerRegistry::new(); + cleanup_workers.register::(EchoWorker).unwrap(); + let cleanup_client = Client::builder( + PostgresDatabase::new(pool.clone()).with_reindex( + PostgresReindexConfig::default() + .with_index_names(["rust_maintenance_reindex_idx"]) + .with_schedule(PostgresReindexSchedule::Interval(Duration::from_millis(50))), + ), + ) + .id("rust-cleanup-client") + .maintenance( + MaintenanceConfig::default() + .with_cancelled_job_retention(riverqueue::Retention::DeleteAfter( + Duration::from_millis(1), + )) + .with_completed_job_retention(riverqueue::Retention::DeleteAfter( + Duration::from_millis(1), + )) + .with_discarded_job_retention(riverqueue::Retention::DeleteAfter( + Duration::from_millis(1), + )) + .with_elect_interval(Duration::from_millis(20)) + .with_job_cleaner_interval(Duration::from_millis(20)) + .with_queue_cleaner_interval(Duration::from_millis(20)) + .with_queue_retention(Duration::from_millis(1)), + ) + .workers(cleanup_workers) + .queue("cleanup_active", QueueConfig::new(1)) + .build() + .unwrap(); + let mut cleanup_handle = cleanup_client.start().unwrap(); + let cleanup_deadline = tokio::time::Instant::now() + Duration::from_secs(5); + loop { + let old_job_count: i64 = + sqlx::query_scalar("SELECT count(*) FROM river_job WHERE id = ANY($1::bigint[])") + .bind(&cleanup_job_ids) + .fetch_one(&pool) + .await + .unwrap(); + let stale_queue_count: i64 = sqlx::query_scalar( + "SELECT count(*) FROM river_queue WHERE name = 'stale_cleanup_queue'", + ) + .fetch_one(&pool) + .await + .unwrap(); + let reindex_file_node_after: i64 = sqlx::query_scalar( + "SELECT pg_relation_filenode('rust_maintenance_reindex_idx'::regclass)::bigint", + ) + .fetch_one(&pool) + .await + .unwrap(); + if old_job_count == 0 + && stale_queue_count == 0 + && reindex_file_node_after != reindex_file_node_before + { + break; + } + assert!( + tokio::time::Instant::now() < cleanup_deadline, + "maintenance did not clean old jobs/queues and reindex in time: \ + old_job_count={old_job_count}, stale_queue_count={stale_queue_count}, \ + reindex_file_node_before={reindex_file_node_before}, \ + reindex_file_node_after={reindex_file_node_after}" + ); + tokio::time::sleep(Duration::from_millis(20)).await; + } + cleanup_handle.shutdown().await.unwrap(); + + database.cleanup().await; +} + +#[tokio::test] +async fn maintenance_client_runs_pilot_periodic_scheduled_and_transactional_jobs() { + let database = support::PostgresSchema::current("rs_pilot").await; + let pilot = TestPilot::default(); + let maintenance_client = maintenance_client(&database.pool, pilot.clone()); + + let scheduled = maintenance_client + .insert(EchoArgs { + message: "scheduled by leader".to_owned(), + }) + .opts( + InsertOpts::default() + .with_scheduled_at(chrono::Utc::now() + chrono::Duration::milliseconds(100)), + ) + .await + .unwrap(); + let transactional = maintenance_client + .insert(TransactionalArgs {}) + .await + .unwrap(); + let mut maintenance_handle = maintenance_client.start().unwrap(); + let periodic = wait_for_job_matching(&maintenance_client, |row| { + row.metadata + .get::("river:periodic_job_id") + .ok() + .flatten() + .as_deref() + == Some("rust-periodic") + }) + .await; + assert_eq!( + periodic.metadata.get::("periodic").unwrap(), + Some(true) + ); + wait_for_state( + &maintenance_client, + scheduled.job.row.id, + JobState::Completed, + ) + .await; + let transactional = wait_for_state( + &maintenance_client, + transactional.job.row.id, + JobState::Completed, + ) + .await; + assert_eq!( + transactional + .metadata + .get::("transactional_completion") + .unwrap(), + Some(true) + ); + assert_eq!( + transactional + .metadata + .get::("extension_handled") + .unwrap(), + Some(true) + ); + assert!(pilot.fetches.load(Ordering::SeqCst) > 0); + assert!(pilot.completions.load(Ordering::SeqCst) > 0); + assert_eq!( + maintenance_client + .jobs() + .get(scheduled.job.row.id) + .await + .unwrap() + .metadata + .get::("extension_handled") + .unwrap(), + Some(true) + ); + assert_eq!(pilot.maintenance_starts.load(Ordering::SeqCst), 1); + assert_eq!(pilot.runtime_starts.load(Ordering::SeqCst), 1); + maintenance_handle.shutdown().await.unwrap(); + assert_eq!(pilot.maintenance_stops.load(Ordering::SeqCst), 1); + assert_eq!(pilot.runtime_stops.load(Ordering::SeqCst), 1); + + database.cleanup().await; +} + +#[tokio::test] +async fn maintenance_leader_rescues_stuck_jobs_and_resigns_on_shutdown() { + let database = support::PostgresSchema::current("rs_rescue_leader").await; + let pool = database.pool.clone(); + let maintenance_client = maintenance_client(&pool, TestPilot::default()); + + let stuck_id: i64 = sqlx::query_scalar( + "INSERT INTO river_job (args, attempt, attempted_at, attempted_by, kind, max_attempts, state) \ + VALUES ('{}'::jsonb, 1, now() - interval '2 hours', ARRAY['dead-client'], \ + 'unregistered_stuck_kind', 2, 'running') RETURNING id", + ) + .fetch_one(&pool) + .await + .unwrap(); + let mut maintenance_handle = maintenance_client.start().unwrap(); + let rescued = wait_for_state(&maintenance_client, stuck_id, JobState::Discarded).await; + assert_eq!( + rescued.metadata.get::("river:rescue_count").unwrap(), + Some(1) + ); + assert_eq!( + rescued.errors.last().unwrap().error, + "Stuck job rescued by JobRescuer" + ); + let leader_id: String = sqlx::query_scalar("SELECT leader_id FROM river_leader") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(leader_id, "rust-maintenance-client"); + maintenance_handle.shutdown().await.unwrap(); + let leader_count: i64 = sqlx::query_scalar("SELECT count(*) FROM river_leader") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(leader_count, 0); + + database.cleanup().await; +} + +#[tokio::test] +async fn migrator_migrates_the_current_schema() { + // The pool's `search_path` points at a fresh schema, so the default + // migrator targets the connection's current schema without touching + // `public`. + let database = support::PostgresSchema::current_unmigrated("rs_migrate_current").await; + + let migrator = PostgresMigrator::new(database.pool.clone()); + migrator.migrate_up().await.unwrap(); + assert_eq!( + migrator.existing_versions().await.unwrap(), + (1..=MIGRATION_VERSION_LATEST).collect::>() + ); + + database.cleanup().await; +} + +#[tokio::test] +async fn migrator_steps_a_custom_schema_up_and_down() { + let database = support::PostgresSchema::unmigrated("rs_migrate_custom").await; + + let custom_migrator = + PostgresMigrator::new(database.pool.clone()).with_schema(database.schema.clone()); + let first_up = custom_migrator + .migrate(Direction::Up, MigrateOpts::new().with_target_version(4)) + .await + .unwrap(); + assert_eq!( + first_up + .versions + .iter() + .map(|version| version.version) + .collect::>(), + vec![1, 2, 3, 4] + ); + assert!(!custom_migrator.validate(None).await.unwrap().is_valid()); + custom_migrator.migrate_up().await.unwrap(); + assert!(custom_migrator.validate(None).await.unwrap().is_valid()); + custom_migrator + .migrate(Direction::Down, MigrateOpts::new().with_target_version(3)) + .await + .unwrap(); + assert_eq!( + custom_migrator.existing_versions().await.unwrap(), + vec![1, 2, 3] + ); + custom_migrator.migrate_up().await.unwrap(); + let dry_run = custom_migrator + .migrate( + Direction::Down, + MigrateOpts::new().with_dry_run(true).with_max_steps(2), + ) + .await + .unwrap(); + assert_eq!(dry_run.versions.len(), 2); + assert_eq!( + custom_migrator.existing_versions().await.unwrap(), + (1..=MIGRATION_VERSION_LATEST).collect::>() + ); + custom_migrator + .migrate(Direction::Down, MigrateOpts::new().with_target_version(-1)) + .await + .unwrap(); + assert_eq!( + custom_migrator.existing_versions().await.unwrap(), + Vec::::new() + ); + + database.cleanup().await; +} + +#[tokio::test] +async fn queue_admin_gets_pauses_resumes_and_updates() { + let database = support::PostgresSchema::current("rs_queue_admin").await; + let client = worker_client(&database.pool, ResumableWorker::default()); + // Starting the client records its configured queue. + let mut run_handle = client.start().unwrap(); + run_handle.wait_ready().await.unwrap(); + run_handle.shutdown().await.unwrap(); + + let queue = client.queues().get("default").await.unwrap(); + assert_eq!(queue.name, "default"); + assert!(queue.paused_at.is_none()); + client.queues().pause("default").await.unwrap(); + assert!( + client + .queues() + .get("default") + .await + .unwrap() + .paused_at + .is_some() + ); + client.queues().resume("default").await.unwrap(); + assert!( + client + .queues() + .get("default") + .await + .unwrap() + .paused_at + .is_none() + ); + let queue = client + .queues() + .update( + "default", + riverqueue::QueueUpdateParams::new().metadata(serde_json::Map::from_iter([( + "owner".to_owned(), + serde_json::json!("rust"), + )])), + ) + .await + .unwrap(); + assert_eq!(queue.metadata["owner"], "rust"); + + database.cleanup().await; +} + +#[tokio::test] +#[allow( + clippy::too_many_lines, + reason = "one end-to-end rescuer scenario compares all worker timeout and retry overrides" +)] +async fn rescuer_honors_worker_timeout_and_retry_overrides() { + let database = support::PostgresSchema::new("rs_rescue_timeout").await; + let pool = database.pool.clone(); + let schema = database.schema.clone(); + + let table = schema.qualify("river_job"); + let insert_sql = format!( + "INSERT INTO {table} \ + (args, attempt, attempted_at, attempted_by, kind, max_attempts, state) \ + VALUES ('{{}}'::jsonb, 1, now() - interval '1 second', ARRAY['dead-client'], $1, $2, 'running') \ + RETURNING id" + ); + let default_timeout_id: i64 = sqlx::query_scalar(AssertSqlSafe(insert_sql.clone())) + .bind(RescueDefaultTimeoutArgs::KIND) + .bind(1_i16) + .fetch_one(&pool) + .await + .unwrap(); + let disabled_timeout_id: i64 = sqlx::query_scalar(AssertSqlSafe(insert_sql.clone())) + .bind(RescueDisabledTimeoutArgs::KIND) + .bind(1_i16) + .fetch_one(&pool) + .await + .unwrap(); + let long_timeout_id: i64 = sqlx::query_scalar(AssertSqlSafe(insert_sql.clone())) + .bind(RescueLongTimeoutArgs::KIND) + .bind(1_i16) + .fetch_one(&pool) + .await + .unwrap(); + let retry_override_id: i64 = sqlx::query_scalar(AssertSqlSafe(insert_sql)) + .bind(RescueRetryOverrideArgs::KIND) + .bind(2_i16) + .fetch_one(&pool) + .await + .unwrap(); + + let mut workers = WorkerRegistry::new(); + workers + .register::(RescueDefaultTimeoutWorker) + .unwrap(); + workers + .register::(RescueDisabledTimeoutWorker) + .unwrap(); + workers + .register::(RescueLongTimeoutWorker) + .unwrap(); + workers + .register::(RescueRetryOverrideWorker) + .unwrap(); + let client = Client::builder( + PostgresDatabase::new(pool.clone()) + .with_schema(schema) + .with_reindex(PostgresReindexConfig::default().with_index_names([] as [&str; 0])), + ) + .id("rust-rescuer-timeout-client") + .job_timeout(Duration::from_millis(100)) + .maintenance( + MaintenanceConfig::default() + .with_elect_interval(Duration::from_millis(20)) + .with_rescue_after(Duration::from_millis(100)) + .with_rescuer_interval(Duration::from_millis(20)), + ) + .queue("default", QueueConfig::new(1)) + .workers(workers) + .build() + .unwrap(); + let mut handle = client.start().unwrap(); + + let default_timeout = wait_for_state(&client, default_timeout_id, JobState::Discarded).await; + assert_eq!( + default_timeout + .metadata + .get::("river:rescue_count") + .unwrap(), + Some(1) + ); + assert_eq!(default_timeout.errors.len(), 1); + let retry_override = wait_for_state(&client, retry_override_id, JobState::Retryable).await; + assert_eq!( + retry_override + .metadata + .get::("river:rescue_count") + .unwrap(), + Some(1) + ); + assert_eq!(retry_override.errors.len(), 1); + assert!( + retry_override.scheduled_at > chrono::Utc::now() + chrono::Duration::minutes(90), + "worker retry override was not applied: {:?}", + retry_override.scheduled_at + ); + + // Every job was stuck from the start, so the rescuer pass that rescued + // the job with the highest ID also looked at these two and left them + // running. + for id in [disabled_timeout_id, long_timeout_id] { + let row = client.jobs().get(id).await.unwrap(); + assert_eq!(row.state, JobState::Running); + assert_eq!(row.errors, []); + assert!(!row.metadata.contains_key("river:rescue_count")); + } + + handle.shutdown().await.unwrap(); + database.cleanup().await; +} + +#[tokio::test] +#[allow(clippy::too_many_lines)] +async fn resumable_cursor_and_transactional_checkpoints() { + let detached_context = riverqueue::__private::work_context(CancellationToken::new()); + let cursor_error = detached_context + .resumable_set_cursor(&ResumableCursor { offset: 1 }) + .unwrap_err(); + assert!( + cursor_error + .to_string() + .contains("resumable cursor can only be set inside a resumable cursor step") + ); + + let database = support::PostgresSchema::new("rs_resumable_ckpt").await; + let pool = database.pool.clone(); + let schema = database.schema.clone(); + + let cursor_values = Arc::new(Mutex::new(Vec::new())); + let validate_runs = Arc::new(AtomicUsize::new(0)); + let mut workers = WorkerRegistry::new(); + workers + .register::(ResumableCheckpointWorker { + cursor_values: Arc::clone(&cursor_values), + pool: pool.clone(), + validate_runs: Arc::clone(&validate_runs), + }) + .unwrap(); + let client = Client::builder( + PostgresDatabase::new(pool.clone()) + .with_schema(schema) + .with_reindex(PostgresReindexConfig::default().with_index_names([] as [&str; 0])), + ) + .id("rust-resumable-checkpoint-test") + .maintenance( + MaintenanceConfig::default() + .with_elect_interval(Duration::from_millis(20)) + .with_scheduler_interval(Duration::from_millis(20)), + ) + .without_notifications() + .queue( + "default", + QueueConfig::new(5) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .workers(workers) + .build() + .unwrap(); + let mut job_ids = std::collections::HashMap::new(); + for mode in [ + "commit_cursor", + "commit_step", + "cursor_retry", + "rollback_cursor", + "rollback_step", + ] { + let inserted = client + .insert(ResumableCheckpointArgs { + mode: mode.to_owned(), + }) + .opts(InsertOpts::default().with_max_attempts(2)) + .await + .unwrap(); + job_ids.insert(mode, inserted.job.row.id); + } + let mut handle = client.start().unwrap(); + + let first_failure = wait_for_state(&client, job_ids["cursor_retry"], JobState::Retryable).await; + assert_eq!( + first_failure + .metadata + .get::("river:resumable_step") + .unwrap() + .as_deref(), + Some("validate") + ); + assert_eq!( + first_failure + .metadata + .get::("river:resumable_cursor") + .unwrap() + .unwrap()["process"], + serde_json::json!({"offset": 42}) + ); + assert_eq!(first_failure.errors.len(), 1); + + let resumed = wait_for_state(&client, job_ids["cursor_retry"], JobState::Completed).await; + assert_eq!(resumed.attempt, 2); + assert_eq!(validate_runs.load(Ordering::SeqCst), 1); + assert_eq!( + *cursor_values.lock().unwrap(), + [ResumableCursor::default(), ResumableCursor { offset: 42 }] + ); + + let committed_cursor = + wait_for_state(&client, job_ids["commit_cursor"], JobState::Completed).await; + assert_eq!( + committed_cursor + .metadata + .get::("river:resumable_step") + .unwrap() + .as_deref(), + Some("tx_cursor") + ); + assert_eq!( + committed_cursor + .metadata + .get::("river:resumable_cursor") + .unwrap() + .unwrap()["tx_cursor"], + serde_json::json!({"offset": 7}) + ); + let committed_step = wait_for_state(&client, job_ids["commit_step"], JobState::Completed).await; + assert_eq!( + committed_step + .metadata + .get::("river:resumable_step") + .unwrap() + .as_deref(), + Some("tx_step") + ); + assert!( + !committed_step + .metadata + .contains_key("river:resumable_cursor") + ); + + for mode in ["rollback_cursor", "rollback_step"] { + let rolled_back = wait_for_state(&client, job_ids[mode], JobState::Completed).await; + assert!(!rolled_back.metadata.contains_key("river:resumable_step")); + assert!(!rolled_back.metadata.contains_key("river:resumable_cursor")); + } + + handle.shutdown().await.unwrap(); + database.cleanup().await; +} + +/// A hard stop gives a job ignoring cancellation the stuck threshold to +/// return, then aborts it and fails its attempt: the attempt counts, the +/// error is recorded, and the job is retried or, at its maximum attempts, +/// discarded. +#[tokio::test] +async fn shutdown_now_fails_a_job_ignoring_cancellation() { + let database = support::PostgresSchema::current("rs_interrupt").await; + + let mut workers = WorkerRegistry::new(); + workers + .register::(IgnoresCancelWorker) + .unwrap(); + let client = Client::builder(database.pool.clone()) + .id("rust-interrupt-client") + .job_stuck_threshold(Duration::from_millis(10)) + .workers(workers) + .queue("interrupt", QueueConfig::new(2)) + .build() + .unwrap(); + let retried = client + .insert(IgnoresCancelArgs {}) + .opts(InsertOpts::default().with_queue("interrupt")) + .await + .unwrap() + .job + .row + .id; + let discarded = client + .insert(IgnoresCancelArgs {}) + .opts( + InsertOpts::default() + .with_queue("interrupt") + .with_max_attempts(1), + ) + .await + .unwrap() + .job + .row + .id; + let mut events = client + .subscribe(&[EventKind::JobFailed, EventKind::JobInterrupted]) + .unwrap(); + let mut handle = client.start().unwrap(); + wait_for_state(&client, retried, JobState::Running).await; + wait_for_state(&client, discarded, JobState::Running).await; + handle.shutdown_now().await.unwrap(); + + for (id, states) in [ + (retried, &[JobState::Available, JobState::Retryable][..]), + (discarded, &[JobState::Discarded][..]), + ] { + let row = client.jobs().get(id).await.unwrap(); + assert_eq!(row.attempt, 1); + assert!(states.contains(&row.state), "{row:?}"); + assert_eq!(row.errors.len(), 1); + assert_eq!(row.errors[0].attempt, 1); + assert_eq!( + row.errors[0].error, + "job aborted after ignoring cancellation" + ); + } + let mut failed = Vec::new(); + for _ in 0..2 { + let event = tokio::time::timeout(Duration::from_secs(1), events.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(event.kind(), EventKind::JobFailed); + failed.push(event.as_job().unwrap().job.id); + } + failed.sort_unstable(); + assert_eq!(failed, [retried, discarded]); + + database.cleanup().await; +} + +#[tokio::test] +async fn transactional_get_and_update_roll_back() { + let database = support::PostgresSchema::current("rs_tx_update").await; + let pool = database.pool.clone(); + let client = worker_client(&pool, ResumableWorker::default()); + let inserted = client + .insert(EchoArgs { + message: "from Rust".to_owned(), + }) + .await + .unwrap(); + + let mut transaction = pool.begin().await.unwrap(); + let tx_row = client + .jobs() + .get(inserted.job.row.id) + .tx(&mut transaction) + .await + .unwrap(); + assert_eq!(tx_row.id, inserted.job.row.id); + client + .jobs() + .update( + tx_row.id, + JobUpdateParams::default().output(serde_json::json!("transactional")), + ) + .tx(&mut transaction) + .await + .unwrap(); + transaction.rollback().await.unwrap(); + assert!( + client + .jobs() + .get(inserted.job.row.id) + .await + .unwrap() + .output() + .is_none() + ); + + database.cleanup().await; +} + +#[tokio::test] +async fn transactional_inserts_become_visible_on_commit() { + let database = support::PostgresSchema::current("rs_tx_insert").await; + let pool = database.pool.clone(); + let client = worker_client(&pool, ResumableWorker::default()); + + let mut transaction = pool.begin().await.unwrap(); + let transaction_insert = client + .insert(EchoArgs { + message: "from Rust".to_owned(), + }) + .tx(&mut transaction) + .await + .unwrap(); + let raw_transaction_insert = riverqueue::__private::ExtensionClient::new(&client) + .insert_raw( + EchoArgs::KIND, + &[], + serde_json::value::to_raw_value(&serde_json::json!({"message": "raw from Rust"})) + .unwrap(), + InsertOpts::default(), + ) + .tx(&mut transaction) + .await + .unwrap(); + assert!(matches!( + client.jobs().get(transaction_insert.job.row.id).await, + Err(riverqueue::Error::NotFound(_)) + )); + assert!(matches!( + client.jobs().get(raw_transaction_insert.job.id).await, + Err(riverqueue::Error::NotFound(_)) + )); + transaction.commit().await.unwrap(); + assert_eq!( + client + .jobs() + .get(transaction_insert.job.row.id) + .await + .unwrap() + .state, + JobState::Available + ); + assert_eq!( + client + .jobs() + .get(raw_transaction_insert.job.id) + .await + .unwrap() + .decode_args::() + .unwrap()["message"], + "raw from Rust" + ); + let pool_connection = client + .database() + .as_postgres() + .expect("client is configured for PostgreSQL") + .pool() + .acquire() + .await + .unwrap(); + // Return the connection so closing the pool at cleanup doesn't wait on it. + drop(pool_connection); + + database.cleanup().await; +} + +/// Builds the maintenance client shared by the pilot and rescuer tests. +fn maintenance_client(pool: &PgPool, pilot: TestPilot) -> Client { + let mut maintenance_workers = WorkerRegistry::new(); + maintenance_workers + .register::(EchoWorker) + .unwrap(); + maintenance_workers + .register::(TransactionalWorker { pool: pool.clone() }) + .unwrap(); + Client::builder(pool.clone()) + .id("rust-maintenance-client") + .maintenance( + MaintenanceConfig::default() + .with_elect_interval(Duration::from_millis(20)) + // Like Go, the rescue age cannot be shorter than the default + // one-minute job timeout. + .with_rescue_after(Duration::from_mins(1)) + .with_rescuer_interval(Duration::from_millis(20)) + .with_scheduler_interval(Duration::from_millis(20)), + ) + .periodic_job(PeriodicJob::with_options( + IntervalSchedule::new(Duration::from_mins(1)).unwrap(), + || EchoArgs { + message: "periodic run on start".to_owned(), + }, + PeriodicJobOpts::new() + .with_id("rust-periodic") + .with_run_on_start(true), + )) + .pilot(pilot) + .workers(maintenance_workers) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap() +} + +async fn wait_for_job_matching(client: &Client, predicate: impl Fn(&JobRow) -> bool) -> JobRow { + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + loop { + let rows = client + .jobs() + .list(JobListParams::default().limit(10_000)) + .await + .unwrap() + .jobs; + if let Some(row) = rows.into_iter().find(&predicate) { + return row; + } + assert!( + tokio::time::Instant::now() < deadline, + "matching job was not inserted" + ); + tokio::time::sleep(Duration::from_millis(20)).await; + } +} + +async fn wait_for_state(client: &Client, id: i64, expected: JobState) -> JobRow { + let deadline = tokio::time::Instant::now() + Duration::from_secs(5); + loop { + let row = client.jobs().get(id).await.unwrap(); + if row.state == expected { + return row; + } + assert!( + tokio::time::Instant::now() < deadline, + "job did not reach {expected:?}; last state: {:?}", + row.state + ); + tokio::time::sleep(Duration::from_millis(20)).await; + } +} + +/// Builds a client on the pool's current schema that works the basic +/// conformance job kinds. +fn worker_client(pool: &PgPool, resumable: ResumableWorker) -> Client { + let mut workers = WorkerRegistry::new(); + workers.register::(CancelWorker).unwrap(); + workers.register::(EchoWorker).unwrap(); + workers.register::(FailWorker).unwrap(); + workers.register::(resumable).unwrap(); + Client::builder(pool.clone()) + .id("rust-conformance-client") + .workers(workers) + .queue("default", QueueConfig::new(2)) + .build() + .unwrap() +} diff --git a/rust/riverqueue/tests/postgres_parity.rs b/rust/riverqueue/tests/postgres_parity.rs new file mode 100644 index 000000000..9c2b40d60 --- /dev/null +++ b/rust/riverqueue/tests/postgres_parity.rs @@ -0,0 +1,531 @@ +//! PostgreSQL parity tests for maintenance, leadership, and storage semantics +//! that mirror the Go implementation. + +#![cfg(all(feature = "postgres", river_postgres_tests))] + +mod support; + +use std::{convert::Infallible, time::Duration}; + +use riverqueue::{ + Client, Error, Job, JobArgs, JobState, QueueConfig, WorkContext, WorkOutcome, WorkerRegistry, + database::PostgresDatabase, +}; +use serde::{Deserialize, Serialize}; + +use support::PostgresSchema; + +fn insert_only_client(database: &PostgresSchema) -> Client { + Client::builder( + PostgresDatabase::new(database.pool.clone()).with_schema(database.schema.clone()), + ) + .build() + .unwrap() +} + +#[tokio::test(flavor = "multi_thread")] +async fn queue_pause_and_resume() { + let database = PostgresSchema::new("rpp_queue_pause").await; + let client = insert_only_client(&database); + + // An unknown named queue is reported like Go's `ErrNotFound`, while `*` + // succeeds even with no persisted queues. + assert!(matches!( + client.queues().pause("missing").await, + Err(Error::NotFound(_)) + )); + assert!(matches!( + client.queues().resume("missing").await, + Err(Error::NotFound(_)) + )); + client + .queues() + .pause(riverqueue::QueueSelector::All) + .await + .unwrap(); + client + .queues() + .resume(riverqueue::QueueSelector::All) + .await + .unwrap(); + + // Go accepts `|` as a queue-name separator. + sqlx::query(sqlx::AssertSqlSafe(format!( + "INSERT INTO {} (name, created_at, metadata, updated_at) VALUES ('tenant|emails', now(), '{{}}', now())", + database.table("river_queue") + ))) + .execute(&database.pool) + .await + .unwrap(); + client.queues().pause("tenant|emails").await.unwrap(); + assert!( + client + .queues() + .get("tenant|emails") + .await + .unwrap() + .paused_at + .is_some() + ); + // Pausing an already paused queue still addresses an existing row. + client.queues().pause("tenant|emails").await.unwrap(); + client.queues().resume("tenant|emails").await.unwrap(); + client.queues().resume("tenant|emails").await.unwrap(); + assert!( + client + .queues() + .get("tenant|emails") + .await + .unwrap() + .paused_at + .is_none() + ); + + database.cleanup().await; +} + +/// Inserts a raw job row and returns its ID. +async fn insert_raw_job( + database: &PostgresSchema, + state: &str, + finalized_ago_secs: Option, +) -> i64 { + let sql = format!( + "INSERT INTO {} (args, kind, max_attempts, state, attempt, attempted_at, finalized_at) \ + VALUES ('{{}}', 'parity_raw', 25, $1::text::{}, \ + CASE WHEN $1 = 'running' THEN 1 ELSE 0 END, \ + CASE WHEN $1 = 'running' THEN now() END, \ + now() - make_interval(secs => $2::bigint)) \ + RETURNING id", + database.table("river_job"), + database.schema.qualify("river_job_state"), + ); + sqlx::query_scalar(sqlx::AssertSqlSafe(sql)) + .bind(state) + .bind(finalized_ago_secs) + .fetch_one(&database.pool) + .await + .unwrap() +} + +#[tokio::test(flavor = "multi_thread")] +async fn concurrent_cancels_return_the_committed_row() { + let database = PostgresSchema::new("rpp_cancel_race").await; + let client = insert_only_client(&database); + let id = insert_raw_job(&database, "scheduled", None).await; + + // Like Go's `JobCancelConcurrentRaceFreshReturn`: the loser of each race + // must return the winner's committed row, not its own snapshot's. + let mut first_finalized_at = None; + for _ in 0..20 { + let (first, second) = tokio::join!(client.jobs().cancel(id), client.jobs().cancel(id)); + let (first, second) = (first.unwrap(), second.unwrap()); + assert_eq!(first.state, JobState::Cancelled); + assert_eq!(second.state, JobState::Cancelled); + assert_eq!(first.finalized_at, second.finalized_at); + let finalized_at = *first_finalized_at.get_or_insert(first.finalized_at); + assert_eq!( + first.finalized_at, finalized_at, + "later cancels must not change finalized_at" + ); + } + + database.cleanup().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn concurrent_retry_loser_returns_the_committed_row() { + let database = PostgresSchema::new("rpp_retry_race").await; + let client = insert_only_client(&database); + let id = insert_raw_job(&database, "cancelled", Some(1)).await; + + // The winner retries in an open transaction, holding the row lock until + // the loser waits on it, so the loser's statement snapshot predates the + // winner's commit. + let mut winner = riverqueue::database::begin_postgres(&database.pool) + .await + .unwrap(); + let retried = client.jobs().retry(id).tx(&mut winner).await.unwrap(); + assert_eq!(retried.state, JobState::Available); + + let mut loser = riverqueue::database::begin_postgres(&database.pool) + .await + .unwrap(); + let loser_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *loser) + .await + .unwrap(); + let loser_client = client.clone(); + let loser_retry = tokio::spawn(async move { + let row = loser_client.jobs().retry(id).tx(&mut loser).await; + loser.rollback().await.unwrap(); + row + }); + tokio::time::timeout(Duration::from_secs(10), async { + loop { + let wait_event_type: Option = + sqlx::query_scalar("SELECT wait_event_type FROM pg_stat_activity WHERE pid = $1") + .bind(loser_pid) + .fetch_one(&database.pool) + .await + .unwrap(); + if wait_event_type.as_deref() == Some("Lock") { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("the losing retry should wait on the winner's row lock"); + winner.commit().await.unwrap(); + + let loser = tokio::time::timeout(Duration::from_secs(10), loser_retry) + .await + .expect("the losing retry should finish once the winner commits") + .unwrap() + .unwrap(); + assert_eq!(loser.state, JobState::Available); + assert_eq!(loser.finalized_at, None); + + database.cleanup().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn job_delete_many() { + let database = PostgresSchema::new("rpp_delete_many").await; + let client = insert_only_client(&database); + + let running = insert_raw_job(&database, "running", None).await; + let first = insert_raw_job(&database, "available", None).await; + let locked = insert_raw_job(&database, "completed", Some(1)).await; + let last = insert_raw_job(&database, "cancelled", Some(1)).await; + + // Running jobs are excluded before the limit, so a limit of two deletes two + // non-running rows even though the lowest ID is running. + let mut blocker = database.pool.begin().await.unwrap(); + sqlx::query(sqlx::AssertSqlSafe(format!( + "SELECT id FROM {} WHERE id = $1 FOR UPDATE", + database.table("river_job") + ))) + .bind(locked) + .execute(&mut *blocker) + .await + .unwrap(); + + // A row locked by another transaction is skipped rather than waited on. + let deleted = tokio::time::timeout( + std::time::Duration::from_secs(5), + client + .jobs() + .delete_many(riverqueue::JobDeleteManyParams::matching( + riverqueue::JobListParams::default() + .ids([running, first, locked, last]) + .limit(2), + )), + ) + .await + .expect("bulk delete must skip locked rows instead of blocking") + .unwrap(); + assert_eq!( + deleted.iter().map(|job| job.id).collect::>(), + vec![first, last] + ); + blocker.rollback().await.unwrap(); + + let remaining = client + .jobs() + .delete_many(riverqueue::JobDeleteManyParams::all()) + .await + .unwrap(); + assert_eq!( + remaining.iter().map(|job| job.id).collect::>(), + vec![locked] + ); + assert_eq!( + client.jobs().get(running).await.unwrap().state, + riverqueue::JobState::Running + ); + + database.cleanup().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn job_list_single_finalized_state_by_time() { + use riverqueue::{JobListOrderBy, JobListParams, JobState, SortDirection}; + + let database = PostgresSchema::new("rpp_list_final").await; + let client = insert_only_client(&database); + + let oldest = insert_raw_job(&database, "completed", Some(30)).await; + let newest = insert_raw_job(&database, "completed", Some(10)).await; + let middle = insert_raw_job(&database, "completed", Some(20)).await; + let _other_state = insert_raw_job(&database, "discarded", Some(15)).await; + + for (direction, expected) in [ + (SortDirection::Ascending, vec![oldest, middle, newest]), + (SortDirection::Descending, vec![newest, middle, oldest]), + ] { + let params = JobListParams::default() + .order_by(JobListOrderBy::Time) + .limit(2) + .states([JobState::Completed]) + .direction(direction); + let first_page = client.jobs().list(params.clone()).await.unwrap(); + assert_eq!( + first_page.jobs.iter().map(|job| job.id).collect::>(), + expected[..2] + ); + let cursor = first_page.last_cursor.unwrap(); + let second_page = client + .jobs() + .list(params.clone().after(cursor)) + .await + .unwrap() + .jobs; + assert_eq!( + second_page.iter().map(|job| job.id).collect::>(), + expected[2..] + ); + } + + // Multiple states keep the generic predicate and still filter correctly. + let params = JobListParams::default() + .order_by(JobListOrderBy::FinalizedAt) + .states([JobState::Completed, JobState::Discarded]); + let both = client.jobs().list(params).await.unwrap().jobs; + assert_eq!(both.len(), 4); + + database.cleanup().await; +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "parity_noop")] +struct NoopArgs {} + +fn noop_workers() -> WorkerRegistry { + let mut workers = WorkerRegistry::new(); + workers + .register_fn(|_context: WorkContext, _job: Job| async { + Ok::<_, Infallible>(WorkOutcome::Complete) + }) + .unwrap(); + workers +} + +async fn wait_for_job_state(client: &Client, id: i64, state: JobState) -> riverqueue::JobRow { + tokio::time::timeout(Duration::from_secs(10), async { + loop { + let row = client.jobs().get(id).await.unwrap(); + if row.state == state { + return row; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap_or_else(|_| panic!("job {id} did not reach {state:?}")) +} + +#[tokio::test(flavor = "multi_thread")] +async fn schema_names_are_quoted_like_go() { + // Go quotes any schema with `SafeIdentifier`; a hyphenated mixed-case + // schema must migrate, notify, elect, and work jobs from Rust. + let database = PostgresSchema::new("Rpp-Mixed-Schema").await; + assert!(database.schema.as_deref().unwrap().contains('-')); + let client = Client::builder( + PostgresDatabase::new(database.pool.clone()).with_schema(database.schema.clone()), + ) + .queue( + "default", + QueueConfig::new(1).with_fetch_poll_interval(Duration::from_mins(1)), + ) + .workers(noop_workers()) + .build() + .unwrap(); + let mut handle = client.start().unwrap(); + handle.wait_ready().await.unwrap(); + + let inserted = client.insert(NoopArgs {}).await.unwrap(); + wait_for_job_state(&client, inserted.job.row.id, JobState::Completed).await; + + handle.shutdown().await.unwrap(); + database.cleanup().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn extension_notify_many_is_transactional() { + use riverqueue::__private::{ + DatabaseConfig, DatabaseConnection, NotificationTopic, notify_many, + }; + + let database = PostgresSchema::new("rpp_notify_many").await; + let config = DatabaseConfig::Postgres { + schema: database.schema.clone(), + }; + let mut listener = sqlx::postgres::PgListener::connect_with(&database.pool) + .await + .unwrap(); + listener + .listen(&format!( + "{}.river_insert", + database.schema.as_deref().unwrap() + )) + .await + .unwrap(); + + // A rolled-back transaction delivers nothing; the committed batch that + // follows is therefore the first thing the listener receives. + let mut rolled_back = database.pool.begin().await.unwrap(); + notify_many( + DatabaseConnection::Postgres(&mut rolled_back), + &config, + NotificationTopic::Insert, + &[r#"{"queue":"rolled_back"}"#.to_owned()], + ) + .await + .unwrap(); + rolled_back.rollback().await.unwrap(); + + let mut committed = database.pool.begin().await.unwrap(); + notify_many( + DatabaseConnection::Postgres(&mut committed), + &config, + NotificationTopic::Insert, + &[ + r#"{"queue":"first"}"#.to_owned(), + r#"{"queue":"second"}"#.to_owned(), + ], + ) + .await + .unwrap(); + committed.commit().await.unwrap(); + + for expected in [r#"{"queue":"first"}"#, r#"{"queue":"second"}"#] { + let notification = tokio::time::timeout(Duration::from_secs(5), listener.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(notification.payload(), expected); + } + // The listener holds a pooled connection that must be returned first. + drop(listener); + database.cleanup().await; +} + +#[tokio::test] +async fn rescue_after_defaults_and_validation_match_go() { + use riverqueue::MaintenanceConfig; + + let pool = sqlx::postgres::PgPoolOptions::new() + .connect_lazy("postgres://localhost/unused") + .unwrap(); + // Like Go's `RescueStuckJobsAfter`, a rescue age shorter than the job + // timeout is rejected, while an equal one is accepted. + let error = Client::builder(pool.clone()) + .job_timeout(Duration::from_mins(5)) + .maintenance(MaintenanceConfig::default().with_rescue_after(Duration::from_mins(4))) + .build() + .unwrap_err(); + assert!(error.to_string().contains("rescue after"), "{error}"); + Client::builder(pool.clone()) + .job_timeout(Duration::from_mins(5)) + .maintenance(MaintenanceConfig::default().with_rescue_after(Duration::from_mins(5))) + .build() + .unwrap(); + assert_eq!(MaintenanceConfig::default().rescue_after(), None); +} + +#[tokio::test(flavor = "multi_thread")] +async fn leader_renews_while_maintenance_is_blocked() { + use riverqueue::MaintenanceConfig; + + let database = PostgresSchema::new("rpp_slow_maintenance").await; + let expired = insert_raw_job(&database, "completed", Some(48 * 3_600)).await; + + // Hold the expired row so the job cleaner's delete blocks on it. + let mut blocker = database.pool.begin().await.unwrap(); + sqlx::query(sqlx::AssertSqlSafe(format!( + "SELECT id FROM {} WHERE id = $1 FOR UPDATE", + database.table("river_job") + ))) + .bind(expired) + .execute(&mut *blocker) + .await + .unwrap(); + + let client = Client::builder( + PostgresDatabase::new(database.pool.clone()).with_schema(database.schema.clone()), + ) + .maintenance( + MaintenanceConfig::default() + .with_elect_interval(Duration::from_millis(50)) + .with_job_cleaner_interval(Duration::from_millis(50)), + ) + .queue("default", QueueConfig::new(1)) + .workers(noop_workers()) + .build() + .unwrap(); + let mut handle = client.start().unwrap(); + + let blocked_deletes = || { + let pool = database.pool.clone(); + async move { + sqlx::query_scalar::<_, i64>( + "SELECT count(*) FROM pg_stat_activity \ + WHERE datname = current_database() AND wait_event_type = 'Lock' \ + AND query LIKE 'DELETE FROM%river_job%'", + ) + .fetch_one(&pool) + .await + .unwrap() + } + }; + let lease = || { + let pool = database.pool.clone(); + let table = database.table("river_leader"); + async move { + sqlx::query_as::<_, (chrono::DateTime, chrono::DateTime)>( + sqlx::AssertSqlSafe(format!("SELECT elected_at, expires_at FROM {table}")), + ) + .fetch_optional(&pool) + .await + .unwrap() + } + }; + tokio::time::timeout(Duration::from_secs(10), async { + while blocked_deletes().await == 0 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("the job cleaner should block on the locked row"); + + // The lease keeps being renewed within the same term while maintenance is + // stuck, rather than waiting for the blocked service. + let (elected_at, mut expires_at) = lease().await.unwrap(); + for _ in 0..3 { + tokio::time::timeout(Duration::from_secs(10), async { + loop { + let (current_elected_at, current_expires_at) = lease().await.unwrap(); + assert_eq!(current_elected_at, elected_at); + if current_expires_at > expires_at { + expires_at = current_expires_at; + return; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("the leader should renew while maintenance is blocked"); + } + assert!(blocked_deletes().await > 0); + + // Shutdown cancels the blocked statement server-side instead of waiting. + tokio::time::timeout(Duration::from_secs(10), handle.shutdown()) + .await + .expect("shutdown should cancel blocked maintenance") + .unwrap(); + assert_eq!(blocked_deletes().await, 0); + blocker.rollback().await.unwrap(); + database.cleanup().await; +} diff --git a/rust/riverqueue/tests/prepared_insertion.rs b/rust/riverqueue/tests/prepared_insertion.rs new file mode 100644 index 000000000..25ab47c27 --- /dev/null +++ b/rust/riverqueue/tests/prepared_insertion.rs @@ -0,0 +1,424 @@ +//! Prepared insertion inserts stored jobs again, such as jobs set aside and +//! retried later, like an ordinary insertion of them: begin hooks and +//! middleware run once and see the stored arguments, nothing decodes the +//! stored row first, and the job keeps its identity with a new ID. +//! +//! PostgreSQL scenarios run in a unique schema and fail rather than skip when +//! `RIVER_RUST_DATABASE_URL` is unset; SQLite scenarios use temporary files. + +#![cfg(any(all(feature = "postgres", river_postgres_tests), feature = "sqlite"))] + +mod support; + +use std::sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, +}; + +use riverqueue::__private::{ExtensionClient, PreparedInsertParams}; +use riverqueue::{ + BoxError, Client, Error, Hook, InsertContext, InsertMiddleware, InsertNext, InsertOpts, + InsertedJob, JobArgs, JobRow, UniqueOpts, +}; +use serde::{Deserialize, Serialize}; +use serde_json::value::RawValue; + +/// The fetch cooldown of clients whose insertions must each notify. +const COOLDOWN: std::time::Duration = std::time::Duration::from_millis(1); + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "prepared_insertion")] +struct PreparedArgs { + value: i64, +} + +/// Wraps arguments in an envelope, keeping one that's already there as a +/// hook that wraps arguments must, and records what it saw. +#[derive(Clone, Default)] +struct EnvelopeHook { + decoded: Arc, + seen: Arc>>, +} + +impl Hook for EnvelopeHook { + async fn insert_begin(&self, insert: &mut InsertContext) -> Result<(), BoxError> { + tokio::task::yield_now().await; + let args = insert.encoded_args.get().to_owned(); + self.seen.lock().unwrap().push(args.clone()); + if !args.starts_with(r#"{"envelope":"#) { + insert.encoded_args = RawValue::from_string(format!(r#"{{"envelope":{args}}}"#))?; + } + Ok(()) + } + + async fn decode_insert_result(&self, job: &mut JobRow) -> Result<(), BoxError> { + tokio::task::yield_now().await; + self.decoded.fetch_add(1, Ordering::SeqCst); + let mut outer: std::collections::HashMap> = job.decode_args()?; + if let Some(inner) = outer.remove("envelope") { + job.encoded_args = inner; + } + Ok(()) + } +} + +/// Counts insertions. +#[derive(Clone, Default)] +struct CountingMiddleware(Arc); + +impl InsertMiddleware for CountingMiddleware { + async fn insert_many( + &self, + jobs: Vec, + next: InsertNext<'_>, + ) -> Result, Error> { + self.0.fetch_add(jobs.len(), Ordering::SeqCst); + next.run(jobs).await + } +} + +fn prepared(row: &JobRow, encoded_args: Box) -> PreparedInsertParams { + PreparedInsertParams { + created_at: row.created_at, + encoded_args, + kind: row.kind.clone(), + max_attempts: row.max_attempts, + metadata: row.metadata.clone(), + priority: row.priority, + queue: row.queue.clone(), + scheduled_at: row.scheduled_at, + tags: row.tags.clone(), + unique_key: row.unique_key.clone(), + unique_states: row.unique_states.clone(), + } +} + +/// Inserts a job, deletes it, and inserts it again from its stored row. +async fn assert_stored_jobs_insert_like_ordinary_jobs( + builder: impl Fn() -> riverqueue::ClientBuilder, + delete: impl AsyncFn(i64), + stored_args: impl AsyncFn(i64) -> String, +) { + let hook = EnvelopeHook::default(); + let middleware = CountingMiddleware::default(); + let client = builder() + .hook(hook.clone()) + .insert_middleware(middleware.clone()) + .build() + .unwrap(); + let scheduled_at = chrono::Utc::now() + chrono::Duration::hours(1); + let original = client + .insert(PreparedArgs { value: 7 }) + .opts( + InsertOpts::default() + .with_metadata( + serde_json::json!({"source": true}) + .as_object() + .unwrap() + .clone(), + ) + .with_scheduled_at(scheduled_at) + .with_tags(["prepared"]) + .with_unique(UniqueOpts::new().with_by_args(true)), + ) + .await + .unwrap() + .job + .row; + let stored = stored_args(original.id).await; + let row = client.jobs().get(original.id).await.unwrap(); + delete(original.id).await; + // SQLite may reuse the highest deleted ID, so take a newer one first. + let sentinel = client.insert(PreparedArgs { value: 8 }).await.unwrap().id(); + let (seen, decoded, inserted) = ( + hook.seen.lock().unwrap().len(), + hook.decoded.load(Ordering::SeqCst), + middleware.0.load(Ordering::SeqCst), + ); + + let reinserted = ExtensionClient::new(&client) + .insert_prepared(vec![prepared( + &row, + RawValue::from_string(stored.clone()).unwrap(), + )]) + .await + .unwrap() + .remove(0); + + // Each step ran once, the hook saw the stored arguments, and nothing + // decoded the stored row before inserting it. + assert_eq!(middleware.0.load(Ordering::SeqCst), inserted + 1); + let seen_now = hook.seen.lock().unwrap().clone(); + assert_eq!(seen_now.len(), seen + 1); + assert_eq!(seen_now.last().unwrap(), &stored); + assert_eq!(hook.decoded.load(Ordering::SeqCst), decoded + 1); + // The hook kept its envelope, so the stored arguments are unchanged. + assert_eq!(stored_args(reinserted.job.id).await, stored); + assert!(reinserted.job.id > sentinel); + assert_eq!(reinserted.job.created_at, original.created_at); + assert_eq!(reinserted.job.scheduled_at, original.scheduled_at); + assert_eq!(reinserted.job.unique_key, original.unique_key); + assert_eq!(reinserted.job.unique_states, original.unique_states); + assert_eq!(reinserted.job.tags, original.tags); + assert_eq!( + reinserted.job.metadata.get::("source").unwrap(), + Some(true) + ); + assert_eq!(reinserted.job.attempt, 0); + assert_eq!(reinserted.job.errors, []); + + // Arguments may be any JSON value. + for args in ["[1,2]", "null"] { + let mut params = prepared(&row, RawValue::from_string(args.to_owned()).unwrap()); + params.unique_key = None; + params.unique_states = None; + let inserted = ExtensionClient::new(&client) + .insert_prepared(vec![params]) + .await + .unwrap(); + assert_eq!(inserted.len(), 1, "{args}"); + } + assert!( + ExtensionClient::new(&client) + .insert_prepared(Vec::new()) + .await + .unwrap() + .is_empty() + ); +} + +/// Two available stored jobs to insert again, without uniqueness. +fn available_params() -> Vec { + let now = chrono::Utc::now(); + (1..=2) + .map(|value| PreparedInsertParams { + created_at: now, + encoded_args: RawValue::from_string(format!(r#"{{"value":{value}}}"#)).unwrap(), + kind: PreparedArgs::KIND.to_owned(), + max_attempts: 25, + metadata: riverqueue::JobMetadata::default(), + priority: 1, + queue: "prepared".to_owned(), + scheduled_at: now - chrono::Duration::seconds(1), + tags: Vec::new(), + unique_key: None, + unique_states: None, + }) + .collect() +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +mod postgres { + use riverqueue::database::PostgresDatabase; + + use super::*; + use crate::support::PostgresSchema; + + #[tokio::test(flavor = "multi_thread")] + async fn stored_jobs_insert_like_ordinary_jobs() { + let schema = PostgresSchema::new("prepared_insert").await; + let table = schema.table("river_job"); + let pool = schema.pool.clone(); + assert_stored_jobs_insert_like_ordinary_jobs( + || { + Client::builder( + PostgresDatabase::new(schema.pool.clone()).with_schema(schema.schema.clone()), + ) + }, + async |id| { + sqlx::query(sqlx::AssertSqlSafe(format!( + "DELETE FROM {table} WHERE id = $1" + ))) + .bind(id) + .execute(&pool) + .await + .unwrap(); + }, + async |id| { + sqlx::query_scalar(sqlx::AssertSqlSafe(format!( + "SELECT args::text FROM {table} WHERE id = $1" + ))) + .bind(id) + .fetch_one(&pool) + .await + .unwrap() + }, + ) + .await; + schema.cleanup().await; + } + + /// Stored jobs inserted again notify their queue once committed, on the + /// client's pool or in a caller's transaction, and a rolled-back + /// transaction keeps neither the jobs nor the notification. + #[tokio::test(flavor = "multi_thread")] + async fn prepared_insertions_notify_on_commit_only() { + let schema = PostgresSchema::new("prepared_notify").await; + // Each insertion below waits out the fetch cooldown, within which a + // client skips a queue's repeated notification. + let client = Client::builder( + PostgresDatabase::new(schema.pool.clone()).with_schema(schema.schema.clone()), + ) + .fetch_cooldown(COOLDOWN) + .build() + .unwrap(); + let channel = format!("{}.river_insert", schema.schema.as_deref().unwrap()); + let mut listener = sqlx::postgres::PgListener::connect_with(&schema.pool) + .await + .unwrap(); + listener.listen(&channel).await.unwrap(); + // Returns the queues notified before a marker sent now. + let mut notified = async || { + sqlx::query("SELECT pg_notify($1, 'marker')") + .bind(&channel) + .execute(&schema.pool) + .await + .unwrap(); + let mut payloads = Vec::new(); + loop { + let notification = listener.recv().await.unwrap(); + if notification.payload() == "marker" { + return payloads; + } + payloads.push(notification.payload().to_owned()); + } + }; + let extension = ExtensionClient::new(&client); + let count = async || -> i64 { + sqlx::query_scalar(sqlx::AssertSqlSafe(format!( + "SELECT count(*) FROM {}", + schema.table("river_job") + ))) + .fetch_one(&schema.pool) + .await + .unwrap() + }; + + assert_eq!( + extension + .insert_prepared(available_params()) + .await + .unwrap() + .len(), + 2 + ); + assert_eq!(notified().await, [r#"{"queue":"prepared"}"#]); + + tokio::time::sleep(COOLDOWN * 2).await; + let mut transaction = schema.pool.begin().await.unwrap(); + extension + .insert_prepared(available_params()) + .tx(&mut transaction) + .await + .unwrap(); + assert!(notified().await.is_empty(), "notified before commit"); + transaction.commit().await.unwrap(); + assert_eq!(notified().await, [r#"{"queue":"prepared"}"#]); + assert_eq!(count().await, 4); + + tokio::time::sleep(COOLDOWN * 2).await; + let mut transaction = schema.pool.begin().await.unwrap(); + extension + .insert_prepared(available_params()) + .tx(&mut transaction) + .await + .unwrap(); + transaction.rollback().await.unwrap(); + assert!(notified().await.is_empty(), "rolled back but notified"); + assert_eq!(count().await, 4); + + drop(listener); + schema.cleanup().await; + } +} + +#[cfg(feature = "sqlite")] +mod sqlite { + use super::*; + use crate::support::{sqlite_cleanup, sqlite_file_pool}; + + #[tokio::test(flavor = "multi_thread")] + async fn stored_jobs_insert_like_ordinary_jobs() { + let (pool, path) = sqlite_file_pool(4).await; + assert_stored_jobs_insert_like_ordinary_jobs( + || Client::builder(pool.clone()), + async |id| { + sqlx::query("DELETE FROM river_job WHERE id = ?") + .bind(id) + .execute(&pool) + .await + .unwrap(); + }, + async |id| { + sqlx::query_scalar("SELECT json(args) FROM river_job WHERE id = ?") + .bind(id) + .fetch_one(&pool) + .await + .unwrap() + }, + ) + .await; + sqlite_cleanup(pool, path).await; + } + + /// Stored jobs inserted again write their queue's notification with the + /// jobs, on the client's pool or in a caller's transaction, and a + /// rolled-back transaction keeps neither. + #[tokio::test(flavor = "multi_thread")] + async fn prepared_insertions_notify_on_commit_only() { + let (pool, path) = sqlite_file_pool(4).await; + // Each insertion below waits out the fetch cooldown, within which a + // client skips a queue's repeated notification. + let client = Client::builder(pool.clone()) + .fetch_cooldown(COOLDOWN) + .build() + .unwrap(); + let extension = ExtensionClient::new(&client); + let counts = async || -> (i64, i64) { + let jobs = sqlx::query_scalar("SELECT count(*) FROM river_job") + .fetch_one(&pool) + .await + .unwrap(); + let notifications = sqlx::query_scalar( + "SELECT count(*) FROM river_notification WHERE topic = 'river_insert' \ + AND json_extract(payload, '$.queue') = 'prepared'", + ) + .fetch_one(&pool) + .await + .unwrap(); + (jobs, notifications) + }; + + assert_eq!( + extension + .insert_prepared(available_params()) + .await + .unwrap() + .len(), + 2 + ); + assert_eq!(counts().await, (2, 1)); + + tokio::time::sleep(COOLDOWN * 2).await; + let mut transaction = pool.begin_with("BEGIN IMMEDIATE").await.unwrap(); + extension + .insert_prepared(available_params()) + .tx(&mut transaction) + .await + .unwrap(); + transaction.commit().await.unwrap(); + assert_eq!(counts().await, (4, 2)); + + tokio::time::sleep(COOLDOWN * 2).await; + let mut transaction = pool.begin_with("BEGIN IMMEDIATE").await.unwrap(); + extension + .insert_prepared(available_params()) + .tx(&mut transaction) + .await + .unwrap(); + transaction.rollback().await.unwrap(); + assert_eq!(counts().await, (4, 2)); + + sqlite_cleanup(pool, path).await; + } +} diff --git a/rust/riverqueue/tests/producer_lifetime.rs b/rust/riverqueue/tests/producer_lifetime.rs new file mode 100644 index 000000000..4e117f1bc --- /dev/null +++ b/rust/riverqueue/tests/producer_lifetime.rs @@ -0,0 +1,769 @@ +//! A producer's lifetime as an extension session sees it: reports that +//! continue while the producer drains, the serial shutdown after its last +//! attempt, queue removal that waits for that shutdown and keeps the name +//! reserved, live reconfiguration, and extension queue settings. +//! +//! PostgreSQL scenarios run in a unique schema and fail rather than skip when +//! `RIVER_RUST_DATABASE_URL` is unset; SQLite scenarios use temporary files. + +#![cfg(any(all(feature = "postgres", river_postgres_tests), feature = "sqlite"))] + +mod support; + +use std::{ + convert::Infallible, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use async_trait::async_trait; +use riverqueue::__private::{ + ClientBuilderExt, Pilot, PilotError, PilotProducer, ProducerConfiguration, + ProducerKeepAliveContext, ProducerShutdownContext, ProducerStartContext, QueueConfigExt, +}; +use riverqueue::{ + Client, Error, ExtensionPhase, InsertOpts, Job, JobArgs, JobRow, JobState, QueueConfig, + WorkContext, WorkOutcome, WorkerRegistry, +}; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value, json}; +use tokio::sync::{Notify, Semaphore}; + +const WAIT: Duration = Duration::from_secs(10); + +/// A job whose worker holds its slot until the test releases it. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "producer_lifetime_gated")] +struct GatedArgs {} + +/// Starts and releases [`GatedArgs`] jobs. +#[derive(Clone)] +struct Gate { + release: Arc, + started: Arc>, +} + +impl Gate { + fn new() -> Self { + Self { + release: Arc::new(Semaphore::new(0)), + started: Arc::default(), + } + } + + fn release(&self, jobs: usize) { + self.release.add_permits(jobs); + } + + fn workers(&self) -> WorkerRegistry { + let gate = self.clone(); + let mut workers = WorkerRegistry::new(); + workers + .register_fn(move |_context: WorkContext, job: Job| { + let gate = gate.clone(); + async move { + gate.started.push(job.id()); + gate.release.acquire().await.unwrap().forget(); + Ok::<_, Infallible>(WorkOutcome::Complete) + } + }) + .unwrap(); + workers + } +} + +fn fast_queue(max_workers: usize) -> QueueConfig { + QueueConfig::new(max_workers) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)) +} + +/// Values recorded by a session, with a notification on every change. +struct Recorder { + changed: Notify, + values: Mutex>, +} + +impl Default for Recorder { + fn default() -> Self { + Self { + changed: Notify::new(), + values: Mutex::new(Vec::new()), + } + } +} + +impl Recorder { + fn push(&self, value: T) { + self.values.lock().unwrap().push(value); + self.changed.notify_waiters(); + } + + fn snapshot(&self) -> Vec { + self.values.lock().unwrap().clone() + } + + /// Waits until `done` holds for the recorded values. + async fn wait_until(&self, what: &str, done: impl Fn(&[T]) -> bool) { + tokio::time::timeout(WAIT, async { + loop { + let changed = self.changed.notified(); + if done(&self.values.lock().unwrap()) { + return; + } + changed.await; + } + }) + .await + .unwrap_or_else(|_| panic!("timed out waiting for {what}")); + } +} + +/// What a session was told, in order. +#[derive(Clone, Debug, PartialEq)] +enum Call { + Configured(ProducerConfiguration), + Finished(i64), + KeepAlive, + Shutdown { attempt: u32, timeout: Duration }, + Started(ProducerConfiguration), +} + +/// How a session's shutdown attempts behave. +#[derive(Clone, Copy, Debug)] +enum ShutdownBehavior { + /// Every attempt succeeds. + Succeed, + /// The first attempt never finishes, the second fails, and the third + /// succeeds. + HangThenFail, +} + +/// Records every call River makes to its sessions. Accepts extension +/// settings that are JSON objects. +#[derive(Clone)] +struct LifetimePilot { + calls: Arc>, + /// The first keep-alive never finishes. + hang_first_keep_alive: bool, + /// The session callback that panics, if any. + panic_in: Option<&'static str>, + shutdown: ShutdownBehavior, + starts: Arc, +} + +impl LifetimePilot { + fn new(shutdown: ShutdownBehavior) -> Self { + Self { + calls: Arc::default(), + hang_first_keep_alive: false, + panic_in: None, + shutdown, + starts: Arc::default(), + } + } + + fn keep_alives(&self) -> usize { + self.calls + .snapshot() + .iter() + .filter(|call| matches!(call, Call::KeepAlive)) + .count() + } + + fn shutdowns(calls: &[Call]) -> Vec<(u32, Duration)> { + calls + .iter() + .filter_map(|call| match call { + Call::Shutdown { attempt, timeout } => Some((*attempt, *timeout)), + _ => None, + }) + .collect() + } +} + +#[async_trait] +impl Pilot for LifetimePilot { + fn validate_queue_settings( + &self, + _queue: &str, + settings: &Map, + ) -> Result<(), PilotError> { + match settings.get("limit") { + None | Some(Value::Number(_)) => Ok(()), + Some(other) => Err(format!("limit must be a number, not {other}").into()), + } + } + + async fn start_producer( + &self, + context: ProducerStartContext, + ) -> Result>, PilotError> { + self.starts.fetch_add(1, Ordering::SeqCst); + self.calls.push(Call::Started(context.configuration)); + Ok(Some(Box::new(self.clone()))) + } +} + +#[async_trait] +impl PilotProducer for LifetimePilot { + fn configuration_changed(&self, configuration: &ProducerConfiguration) { + self.calls.push(Call::Configured(configuration.clone())); + assert_ne!( + self.panic_in, + Some("configuration_changed"), + "panicked on purpose" + ); + } + + fn job_finished(&self, job: &JobRow) { + self.calls.push(Call::Finished(job.id)); + assert_ne!(self.panic_in, Some("job_finished"), "panicked on purpose"); + } + + async fn keep_alive(&self, context: ProducerKeepAliveContext) -> Result<(), PilotError> { + assert!(context.stale_before < chrono::Utc::now()); + self.calls.push(Call::KeepAlive); + if self.hang_first_keep_alive && self.keep_alives() == 1 { + std::future::pending::<()>().await; + } + Ok(()) + } + + async fn shutdown(&self, context: ProducerShutdownContext) -> Result<(), PilotError> { + self.calls.push(Call::Shutdown { + attempt: context.attempt, + timeout: context.timeout, + }); + match (self.shutdown, context.attempt) { + (ShutdownBehavior::HangThenFail, 1) => std::future::pending().await, + (ShutdownBehavior::HangThenFail, 2) => Err("shutdown failed on purpose".into()), + _ => Ok(()), + } + } +} + +fn builder_with( + builder: riverqueue::ClientBuilder, + pilot: &LifetimePilot, + gate: &Gate, +) -> riverqueue::ClientBuilder { + builder + .pilot(pilot.clone()) + .producer_report_interval(Duration::from_millis(20)) + .workers(gate.workers()) +} + +/// A stopping producer keeps reporting while its attempts drain, so peers +/// keep counting them, and shuts its session down only after the last one +/// finished. Nothing reaches the session after shutdown. +async fn assert_reports_continue_through_the_drain(builder: riverqueue::ClientBuilder) { + let pilot = LifetimePilot::new(ShutdownBehavior::Succeed); + let gate = Gate::new(); + let client = builder_with(builder, &pilot, &gate) + .queue("default", fast_queue(1)) + .build() + .unwrap(); + let id = client.insert(GatedArgs {}).await.unwrap().id(); + let mut run = client.start().unwrap(); + gate.started + .wait_until("the job to start", |started| started.contains(&id)) + .await; + + run.stopper().stop(); + let before = pilot.keep_alives(); + pilot + .calls + .wait_until("two reports during the drain", |calls| { + calls + .iter() + .filter(|call| matches!(call, Call::KeepAlive)) + .count() + >= before + 2 + }) + .await; + assert_eq!(LifetimePilot::shutdowns(&pilot.calls.snapshot()), []); + assert_eq!( + client.jobs().get(id).await.unwrap().state, + JobState::Running + ); + + gate.release(1); + tokio::time::timeout(WAIT, run.wait()) + .await + .expect("client stops") + .unwrap(); + let calls = pilot.calls.snapshot(); + let finished = calls + .iter() + .position(|call| *call == Call::Finished(id)) + .expect("job finished"); + let shutdown = calls + .iter() + .position(|call| matches!(call, Call::Shutdown { .. })) + .expect("session shut down"); + assert!(finished < shutdown, "{calls:?}"); + assert_eq!(shutdown, calls.len() - 1, "{calls:?}"); + assert_eq!( + LifetimePilot::shutdowns(&calls), + [(1, Duration::from_millis(100))] + ); + assert_eq!( + client.jobs().get(id).await.unwrap().state, + JobState::Completed + ); +} + +/// Shutdown attempts run one at a time with Go's growing deadlines: an +/// attempt that doesn't finish in time is dropped, and a failed one is +/// retried. +async fn assert_shutdown_retries_with_growing_deadlines(builder: riverqueue::ClientBuilder) { + let pilot = LifetimePilot::new(ShutdownBehavior::HangThenFail); + let gate = Gate::new(); + let client = builder_with(builder, &pilot, &gate) + .queue("default", fast_queue(1)) + .build() + .unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + tokio::time::timeout(WAIT, run.shutdown()) + .await + .expect("client stops") + .unwrap(); + + assert_eq!( + LifetimePilot::shutdowns(&pilot.calls.snapshot()), + [ + (1, Duration::from_millis(100)), + (2, Duration::from_millis(500)), + (3, Duration::from_millis(2_500)), + ] + ); +} + +/// A keep-alive that never finishes is dropped after Go's ten seconds, and +/// the next report follows. +async fn assert_stuck_keep_alives_time_out(builder: riverqueue::ClientBuilder) { + let pilot = LifetimePilot { + hang_first_keep_alive: true, + ..LifetimePilot::new(ShutdownBehavior::Succeed) + }; + let gate = Gate::new(); + let client = builder_with(builder, &pilot, &gate) + .queue("default", fast_queue(1)) + .build() + .unwrap(); + let mut run = client.start().unwrap(); + pilot + .calls + .wait_until("the first report", |calls| calls.contains(&Call::KeepAlive)) + .await; + let started = std::time::Instant::now(); + tokio::time::timeout(Duration::from_secs(20), async { + while pilot.keep_alives() < 2 { + pilot.calls.changed.notified().await; + } + }) + .await + .expect("a report after the stuck one"); + let waited = started.elapsed(); + assert!( + (Duration::from_secs(9)..Duration::from_secs(15)).contains(&waited), + "{waited:?}" + ); + run.shutdown().await.unwrap(); +} + +/// A panic in `job_finished` or `configuration_changed` stops the client +/// with an extension error, after the producer drains and shuts the session +/// down. +async fn assert_callback_panics_stop_the_client_in_order( + builder: impl Fn() -> riverqueue::ClientBuilder, + callback: &'static str, +) { + let pilot = LifetimePilot { + panic_in: Some(callback), + ..LifetimePilot::new(ShutdownBehavior::Succeed) + }; + let gate = Gate::new(); + let client = builder_with(builder(), &pilot, &gate) + .queue("default", fast_queue(1)) + .build() + .unwrap(); + let id = client.insert(GatedArgs {}).await.unwrap().id(); + let mut run = client.start().unwrap(); + gate.started + .wait_until("the job to start", |started| started.contains(&id)) + .await; + if callback == "configuration_changed" { + client + .local_queues() + .update("default", fast_queue(2)) + .unwrap(); + pilot + .calls + .wait_until("the configuration change", |calls| { + calls.iter().any(|call| matches!(call, Call::Configured(_))) + }) + .await; + } + // The job ignores cancellation, so the drain waits for it. + gate.release(1); + let error = tokio::time::timeout(WAIT, run.wait()) + .await + .expect("client stops") + .unwrap_err(); + assert!( + matches!( + error, + Error::Extension { + phase: ExtensionPhase::AddOn { + operation: "producer" + }, + .. + } + ), + "{callback}: {error}" + ); + let calls = pilot.calls.snapshot(); + assert!(calls.contains(&Call::Finished(id)), "{callback}: {calls:?}"); + assert!( + matches!(calls.last(), Some(Call::Shutdown { attempt: 1, .. })), + "{callback}: {calls:?}" + ); +} + +/// Like Go's `QueueBundle.Remove`, removing a queue waits for its producer +/// to drain and shut its session down, and the name stays reserved until +/// then. Adding the queue again starts a new session. +async fn assert_removal_waits_for_the_drain(builder: riverqueue::ClientBuilder) { + let pilot = LifetimePilot::new(ShutdownBehavior::Succeed); + let gate = Gate::new(); + let client = builder_with(builder, &pilot, &gate) + .queue("default", fast_queue(1)) + .build() + .unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + client.local_queues().add("gated", fast_queue(1)).unwrap(); + let id = client + .insert(GatedArgs {}) + .opts(InsertOpts::default().with_queue("gated")) + .await + .unwrap() + .id(); + gate.started + .wait_until("the job to start", |started| started.contains(&id)) + .await; + + let removing = tokio::spawn({ + let client = client.clone(); + async move { client.local_queues().remove("gated").await } + }); + // Reports keep arriving while the removed queue drains, and the removal + // stays pending until the job finishes. + let before = pilot.keep_alives(); + pilot + .calls + .wait_until("reports during the removal", |calls| { + calls + .iter() + .filter(|call| matches!(call, Call::KeepAlive)) + .count() + >= before + 4 + }) + .await; + assert!(!removing.is_finished()); + assert!(!client.local_queues().configs().contains_key("gated")); + assert!(matches!( + client.local_queues().add("gated", fast_queue(1)), + Err(Error::QueueAlreadyAdded { name }) if name == "gated" + )); + + gate.release(1); + let removed = tokio::time::timeout(WAIT, removing) + .await + .expect("removal finishes") + .unwrap() + .unwrap(); + assert_eq!(removed, fast_queue(1)); + // Both the default queue's session and the removed one's may report, but + // the removed session shut down before the removal returned. + assert_eq!(LifetimePilot::shutdowns(&pilot.calls.snapshot()).len(), 1); + + let starts = pilot.starts.load(Ordering::SeqCst); + client.local_queues().add("gated", fast_queue(1)).unwrap(); + pilot + .calls + .wait_until("a new session", |_| { + pilot.starts.load(Ordering::SeqCst) > starts + }) + .await; + run.shutdown().await.unwrap(); +} + +/// A running producer applies an updated configuration without restarting: +/// more workers start another job at once, and fewer workers never cancel +/// running jobs. The session sees each configuration. +async fn assert_updates_apply_while_running(builder: riverqueue::ClientBuilder) { + let pilot = LifetimePilot::new(ShutdownBehavior::Succeed); + let gate = Gate::new(); + let client = builder_with(builder, &pilot, &gate) + .queue("default", fast_queue(1)) + .build() + .unwrap(); + let first = client.insert(GatedArgs {}).await.unwrap().id(); + let second = client.insert(GatedArgs {}).await.unwrap().id(); + let mut run = client.start().unwrap(); + gate.started + .wait_until("the first job", |started| started == [first]) + .await; + + client + .local_queues() + .update("default", fast_queue(2)) + .unwrap(); + gate.started + .wait_until("the second job", |started| started.contains(&second)) + .await; + client + .local_queues() + .update("default", fast_queue(1)) + .unwrap(); + pilot + .calls + .wait_until("both configurations", |calls| { + calls + .iter() + .filter_map(|call| match call { + Call::Configured(configuration) => Some(configuration.max_workers), + _ => None, + }) + .eq([2, 1]) + }) + .await; + + gate.release(2); + for id in [first, second] { + tokio::time::timeout(WAIT, async { + while client.jobs().get(id).await.unwrap().state != JobState::Completed { + tokio::task::yield_now().await; + } + }) + .await + .expect("jobs complete"); + } + run.shutdown().await.unwrap(); + assert_eq!(pilot.starts.load(Ordering::SeqCst), 1, "no restart"); +} + +/// Extension queue settings reach the session and are validated by the +/// extension when the client is built and when a queue is added or +/// updated. +async fn assert_queue_settings_reach_the_session(builder: impl Fn() -> riverqueue::ClientBuilder) { + let gate = Gate::new(); + let error = builder() + .workers(gate.workers()) + .queue( + "default", + fast_queue(1).with_extension_setting("limit", json!(1)), + ) + .build() + .unwrap_err(); + assert!( + matches!( + error, + Error::Extension { + phase: ExtensionPhase::AddOn { + operation: "queue settings" + }, + .. + } + ), + "{error}" + ); + + let pilot = LifetimePilot::new(ShutdownBehavior::Succeed); + let client = builder_with(builder(), &pilot, &gate) + .queue( + "default", + fast_queue(1).with_extension_setting("limit", json!(1)), + ) + .build() + .unwrap(); + for rejected in [ + client.local_queues().add( + "other", + fast_queue(1).with_extension_setting("limit", json!("x")), + ), + client.local_queues().update( + "default", + fast_queue(1).with_extension_setting("limit", json!("x")), + ), + ] { + assert!( + matches!( + rejected, + Err(Error::Extension { + phase: ExtensionPhase::AddOn { + operation: "queue settings" + }, + .. + }) + ), + "{rejected:?}" + ); + } + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + client + .local_queues() + .update( + "default", + fast_queue(1).with_extension_setting("limit", json!(2)), + ) + .unwrap(); + pilot + .calls + .wait_until("the updated settings", |calls| { + calls.iter().any(|call| { + matches!(call, Call::Configured(configuration) + if configuration.settings.get("limit") == Some(&json!(2))) + }) + }) + .await; + run.shutdown().await.unwrap(); + let Some(Call::Started(started)) = pilot.calls.snapshot().first().cloned() else { + panic!("session never started"); + }; + assert_eq!(started.settings.get("limit"), Some(&json!(1))); +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +mod postgres { + use riverqueue::database::PostgresDatabase; + + use super::*; + use crate::support::PostgresSchema; + + fn builder(schema: &PostgresSchema) -> riverqueue::ClientBuilder { + Client::builder( + PostgresDatabase::new(schema.pool.clone()).with_schema(schema.schema.clone()), + ) + } + + #[tokio::test(flavor = "multi_thread")] + async fn callback_panics_stop_the_client_in_order() { + for callback in ["configuration_changed", "job_finished"] { + let schema = PostgresSchema::new("lifetime_panic").await; + assert_callback_panics_stop_the_client_in_order(|| builder(&schema), callback).await; + schema.cleanup().await; + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn stuck_keep_alives_time_out() { + let schema = PostgresSchema::new("lifetime_stuck_report").await; + assert_stuck_keep_alives_time_out(builder(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn queue_settings_reach_the_session() { + let schema = PostgresSchema::new("lifetime_settings").await; + assert_queue_settings_reach_the_session(|| builder(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn removal_waits_for_the_drain() { + let schema = PostgresSchema::new("lifetime_removal").await; + assert_removal_waits_for_the_drain(builder(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn reports_continue_through_the_drain() { + let schema = PostgresSchema::new("lifetime_drain").await; + assert_reports_continue_through_the_drain(builder(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn shutdown_retries_with_growing_deadlines() { + let schema = PostgresSchema::new("lifetime_shutdown").await; + assert_shutdown_retries_with_growing_deadlines(builder(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn updates_apply_while_running() { + let schema = PostgresSchema::new("lifetime_update").await; + assert_updates_apply_while_running(builder(&schema)).await; + schema.cleanup().await; + } +} + +#[cfg(feature = "sqlite")] +mod sqlite { + use super::*; + use crate::support::{sqlite_cleanup, sqlite_file_pool}; + + #[tokio::test(flavor = "multi_thread")] + async fn callback_panics_stop_the_client_in_order() { + for callback in ["configuration_changed", "job_finished"] { + let (pool, path) = sqlite_file_pool(4).await; + assert_callback_panics_stop_the_client_in_order( + || Client::builder(pool.clone()), + callback, + ) + .await; + sqlite_cleanup(pool, path).await; + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn stuck_keep_alives_time_out() { + let (pool, path) = sqlite_file_pool(4).await; + assert_stuck_keep_alives_time_out(Client::builder(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn queue_settings_reach_the_session() { + let (pool, path) = sqlite_file_pool(4).await; + assert_queue_settings_reach_the_session(|| Client::builder(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn removal_waits_for_the_drain() { + let (pool, path) = sqlite_file_pool(4).await; + assert_removal_waits_for_the_drain(Client::builder(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn reports_continue_through_the_drain() { + let (pool, path) = sqlite_file_pool(4).await; + assert_reports_continue_through_the_drain(Client::builder(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn shutdown_retries_with_growing_deadlines() { + let (pool, path) = sqlite_file_pool(4).await; + assert_shutdown_retries_with_growing_deadlines(Client::builder(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn updates_apply_while_running() { + let (pool, path) = sqlite_file_pool(4).await; + assert_updates_apply_while_running(Client::builder(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } +} diff --git a/rust/riverqueue/tests/producer_sessions.rs b/rust/riverqueue/tests/producer_sessions.rs new file mode 100644 index 000000000..695f1eb46 --- /dev/null +++ b/rust/riverqueue/tests/producer_sessions.rs @@ -0,0 +1,886 @@ +//! Extension producer sessions: claims the session owns, the protocol checks +//! River applies to what it returns, per-attempt accounting through +//! `job_finished`, and configuration changes. +//! +//! PostgreSQL scenarios run in a unique schema and fail rather than skip when +//! `RIVER_RUST_DATABASE_URL` is unset; SQLite scenarios use temporary files. + +#![cfg(any(all(feature = "postgres", river_postgres_tests), feature = "sqlite"))] + +mod support; + +use std::{ + collections::HashMap, + convert::Infallible, + sync::{Arc, Mutex}, + time::Duration, +}; + +use async_trait::async_trait; +use riverqueue::__private::{ + ClaimedJob, ClientBuilderExt, DatabaseConnection, Pilot, PilotError, PilotProducer, + ProducerClaimContext, ProducerClaimNext, ProducerConfiguration, ProducerStartContext, +}; +use riverqueue::{ + Client, Error, ExtensionPhase, InsertOpts, Job, JobArgs, JobRow, JobState, QueueConfig, + QueueUpdateParams, WorkContext, WorkOutcome, WorkerRegistry, +}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use tokio::sync::Notify; + +const WAIT: Duration = Duration::from_secs(10); + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "producer_session")] +struct SessionArgs { + fail: bool, +} + +/// A job whose worker blocks its thread until released, ignoring +/// cancellation. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "producer_session_blocking")] +struct BlockingArgs {} + +fn fast_queue(max_workers: usize) -> QueueConfig { + QueueConfig::new(max_workers) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)) +} + +fn workers() -> WorkerRegistry { + let mut workers = WorkerRegistry::new(); + workers + .register_fn(|_context: WorkContext, job: Job| async move { + if job.args.fail { + return Err(std::io::Error::other("failed on purpose")); + } + Ok(WorkOutcome::Complete) + }) + .unwrap(); + workers +} + +/// Values recorded by a session, with a notification on every change. +struct Recorder { + changed: Notify, + values: Mutex>, +} + +impl Default for Recorder { + fn default() -> Self { + Self { + changed: Notify::new(), + values: Mutex::new(Vec::new()), + } + } +} + +impl Recorder { + fn push(&self, value: T) { + self.values.lock().unwrap().push(value); + self.changed.notify_waiters(); + } + + fn snapshot(&self) -> Vec { + self.values.lock().unwrap().clone() + } + + /// Waits until `done` holds for the recorded values. + async fn wait_until(&self, what: &str, done: impl Fn(&[T]) -> bool) { + tokio::time::timeout(WAIT, async { + loop { + let changed = self.changed.notified(); + if done(&self.values.lock().unwrap()) { + return; + } + changed.await; + } + }) + .await + .unwrap_or_else(|_| panic!("timed out waiting for {what}")); + } +} + +/// How [`SessionPilot`]'s sessions claim. +#[derive(Clone, Copy, Debug)] +enum Claim { + /// River's claim in the session's own transaction. + Standard, + /// River's claim, returning its first job twice. + Duplicate, + /// River's claim, reporting its first job as still available. + NotRunning, + /// River's claim, reporting its first job in another queue. + WrongQueue, + /// River's claim, reporting its first job as last attempted by another + /// client. + ForeignClient, + /// River's claim, padded with made-up jobs past the claim's limit. + OverLimit, + /// Available rows selected with River's projection but never claimed, + /// decoded as far as possible. + Unclaimed, + /// A row without River's columns, which can't be identified. + Unidentifiable, +} + +/// Starts a session for every producer generation and records what River +/// tells it. +#[derive(Clone)] +struct SessionPilot { + claim: Claim, + configurations: Arc>, + finished: Arc>, +} + +impl SessionPilot { + fn new(claim: Claim) -> Self { + Self { + claim, + configurations: Arc::default(), + finished: Arc::default(), + } + } +} + +#[async_trait] +impl Pilot for SessionPilot { + async fn start_producer( + &self, + context: ProducerStartContext, + ) -> Result>, PilotError> { + self.configurations.push(context.configuration); + Ok(Some(Box::new(self.clone()))) + } +} + +#[async_trait] +impl PilotProducer for SessionPilot { + fn intercepts_claim(&self) -> bool { + true + } + + async fn claim( + &self, + context: ProducerClaimContext<'_>, + next: ProducerClaimNext<'_>, + ) -> Result, PilotError> { + let mut transaction = context.database.begin().await?; + if matches!(self.claim, Claim::Unclaimed | Claim::Unidentifiable) { + let jobs = raw_claim(self.claim, &mut transaction).await?; + transaction.commit().await?; + return Ok(jobs); + } + let mut jobs = next.claim(transaction.connection()).await?; + transaction.commit().await?; + let first = jobs.first().and_then(ClaimedJob::job).cloned(); + match (self.claim, first) { + (Claim::Duplicate, Some(first)) => jobs.push(first.into()), + (Claim::NotRunning, Some(mut first)) => { + first.state = JobState::Available; + jobs[0] = first.into(); + } + (Claim::WrongQueue, Some(mut first)) => { + "elsewhere".clone_into(&mut first.queue); + jobs[0] = first.into(); + } + (Claim::ForeignClient, Some(mut first)) => { + first.attempted_by.push("another-client".to_owned()); + jobs[0] = first.into(); + } + (Claim::OverLimit, Some(first)) => { + for offset in 1..=i64::try_from(context.limit)? { + let mut extra = first.clone(); + extra.id += offset * 1_000_000; + jobs.push(extra.into()); + } + } + _ => {} + } + Ok(jobs) + } + + fn configuration_changed(&self, configuration: &ProducerConfiguration) { + self.configurations.push(configuration.clone()); + } + + fn job_finished(&self, job: &JobRow) { + self.finished.push(job.id); + } +} + +/// Selects rows for [`Claim::Unclaimed`] or [`Claim::Unidentifiable`] +/// without claiming anything. +async fn raw_claim( + claim: Claim, + transaction: &mut riverqueue::__private::PilotTransaction, +) -> Result, PilotError> { + let unclaimed = matches!(claim, Claim::Unclaimed); + match transaction.connection() { + #[cfg(feature = "postgres")] + DatabaseConnection::Postgres(connection) => { + let sql = if unclaimed { + format!( + "SELECT {}, false AS unique_skipped_as_duplicate FROM river_job AS job \ + WHERE state = 'available'", + riverqueue::__private::postgres_job_projection("job") + ) + } else { + "SELECT 1 AS id".to_owned() + }; + let rows = sqlx::query(sqlx::AssertSqlSafe(sql)) + .fetch_all(connection) + .await?; + Ok(rows + .iter() + .map(riverqueue::__private::claimed_postgres_job) + .collect()) + } + #[cfg(feature = "sqlite")] + DatabaseConnection::Sqlite(connection) => { + let sql = if unclaimed { + format!( + "SELECT {} FROM river_job WHERE state = 'available'", + riverqueue::__private::SQLITE_JOB_COLUMNS + ) + } else { + "SELECT 1 AS id".to_owned() + }; + let rows = sqlx::query(sqlx::AssertSqlSafe(sql)) + .fetch_all(connection) + .await?; + Ok(rows + .iter() + .map(riverqueue::__private::claimed_sqlite_job) + .collect()) + } + #[allow(unreachable_patterns)] + _ => unreachable!("built-in backends only"), + } +} + +fn counts(ids: &[i64]) -> HashMap { + let mut counts = HashMap::new(); + for id in ids { + *counts.entry(*id).or_default() += 1; + } + counts +} + +/// Every accepted claimed row reaches `job_finished` exactly once, whatever +/// its attempt's outcome. +async fn assert_every_attempt_finishes_once(builder: riverqueue::ClientBuilder) { + let pilot = SessionPilot::new(Claim::Standard); + let client = builder + .pilot(pilot.clone()) + .queue("default", fast_queue(3)) + .workers(workers()) + .build() + .unwrap(); + let mut ids = Vec::new(); + for fail in [false, true, false, false, true] { + // A failed job is discarded, so it can't run twice. + ids.push( + client + .insert(SessionArgs { fail }) + .opts(InsertOpts::default().with_max_attempts(1)) + .await + .unwrap() + .id(), + ); + } + let mut run = client.start().unwrap(); + pilot + .finished + .wait_until("every attempt to finish", |finished| { + ids.iter().all(|id| finished.contains(id)) + }) + .await; + run.shutdown().await.unwrap(); + + let finished = counts(&pilot.finished.snapshot()); + assert_eq!(finished.len(), ids.len(), "{finished:?}"); + assert!(finished.values().all(|count| *count == 1), "{finished:?}"); + for (index, id) in ids.iter().enumerate() { + let expected = if matches!(index, 1 | 4) { + JobState::Discarded + } else { + JobState::Completed + }; + assert_eq!(client.jobs().get(*id).await.unwrap().state, expected); + } +} + +/// An attempt whose worker outlives its abort during shutdown leaves its job +/// running for the rescuer, but still finishes in the session, so the +/// extension doesn't count it against the queue for the rest of the run. +async fn assert_abandoned_attempts_finish(builder: riverqueue::ClientBuilder) { + let pilot = SessionPilot::new(Claim::Standard); + let (release, blocked) = std::sync::mpsc::channel::<()>(); + let blocked = Arc::new(Mutex::new(blocked)); + let started = Arc::new(Notify::new()); + let mut workers = WorkerRegistry::new(); + let worker_started = Arc::clone(&started); + workers + .register_fn(move |_context: WorkContext, _job: Job| { + let blocked = Arc::clone(&blocked); + let started = Arc::clone(&worker_started); + async move { + started.notify_one(); + // Blocks the worker's thread, so neither cancellation nor + // an abort can end it. + let _ = blocked.lock().unwrap().recv(); + Ok::<_, Infallible>(WorkOutcome::Complete) + } + }) + .unwrap(); + let client = builder + .pilot(pilot.clone()) + .job_stuck_threshold(Duration::from_millis(50)) + .queue("default", fast_queue(1)) + .workers(workers) + .build() + .unwrap(); + let id = client.insert(BlockingArgs {}).await.unwrap().id(); + let mut run = client.start().unwrap(); + tokio::time::timeout(WAIT, started.notified()) + .await + .expect("worker starts"); + + tokio::time::timeout(WAIT, run.shutdown_now()) + .await + .expect("client stops") + .unwrap(); + assert_eq!(pilot.finished.snapshot(), [id]); + assert_eq!( + client.jobs().get(id).await.unwrap().state, + JobState::Running + ); + release.send(()).unwrap(); +} + +/// Panics while handling a failed job, which ends its attempt's task. +struct PanickingErrorHandler; + +impl riverqueue::ErrorHandler for PanickingErrorHandler { + async fn handle_error( + &self, + _context: &WorkContext, + _job: &JobRow, + _result: &riverqueue::WorkResult, + ) -> Result { + tokio::task::yield_now().await; + panic!("error handler panicked on purpose") + } +} + +/// An attempt whose task panics, here in the client's error handler, still +/// finishes in the session. +async fn assert_panicked_attempts_finish(builder: riverqueue::ClientBuilder) { + let pilot = SessionPilot::new(Claim::Standard); + let client = builder + .pilot(pilot.clone()) + .error_handler(PanickingErrorHandler) + .queue("default", fast_queue(1)) + .workers(workers()) + .build() + .unwrap(); + let id = client + .insert(SessionArgs { fail: true }) + .await + .unwrap() + .id(); + let mut run = client.start().unwrap(); + pilot + .finished + .wait_until("the panicked attempt to finish", |finished| { + finished.contains(&id) + }) + .await; + run.shutdown().await.unwrap(); + assert_eq!(pilot.finished.snapshot(), [id]); +} + +/// A session result River can't accept stops the client with a protocol +/// error, and its rows never reach `job_finished`. +/// +/// `corrupt` runs on the inserted job before the client starts, to make it +/// undecodable for [`Claim::Unclaimed`]. +async fn assert_broken_claims_stop_the_client( + builder: impl Fn() -> riverqueue::ClientBuilder, + claim: Claim, + corrupt: F, +) where + F: FnOnce(i64) -> Fut, + Fut: std::future::Future, +{ + let pilot = SessionPilot::new(claim); + let client = builder() + .pilot(pilot.clone()) + .queue("default", fast_queue(2)) + .workers(workers()) + .build() + .unwrap(); + let id = client + .insert(SessionArgs { fail: false }) + .await + .unwrap() + .id(); + if matches!(claim, Claim::Unclaimed) { + corrupt(id).await; + } + let mut run = client.start().unwrap(); + let error = tokio::time::timeout(WAIT, run.wait()) + .await + .expect("client stops") + .unwrap_err(); + assert!( + matches!( + error, + Error::Extension { + phase: ExtensionPhase::AddOn { + operation: "fetch claim" + }, + .. + } + ), + "{claim:?}: {error}" + ); + assert!(pilot.finished.snapshot().is_empty(), "{claim:?}"); + // A committed claim is left for the rescuer. (River can't read a + // corrupted row back.) + if matches!(claim, Claim::Unclaimed) { + return; + } + let expected = if matches!(claim, Claim::Unidentifiable) { + JobState::Available + } else { + JobState::Running + }; + assert_eq!( + client.jobs().get(id).await.unwrap().state, + expected, + "{claim:?}" + ); +} + +/// Every way a claim result can break the protocol. +const BROKEN_CLAIMS: [Claim; 7] = [ + Claim::Duplicate, + Claim::NotRunning, + Claim::WrongQueue, + Claim::ForeignClient, + Claim::OverLimit, + Claim::Unclaimed, + Claim::Unidentifiable, +]; + +/// Queue record changes reach the session at once, as Go's producer handles +/// `metadata_changed`, rather than at the next queue poll two seconds +/// later. A listening client learns of another client's update through the +/// control notification; a poll-only client learns of its own update through +/// a local signal. +/// +/// Two consecutive updates must each arrive within 1.5 seconds. A poll could +/// catch the first by chance, but the second is made right after that poll, +/// so it would wait almost the whole interval. +async fn assert_queue_changes_reach_the_session( + builder: impl Fn() -> riverqueue::ClientBuilder, + poll_only: bool, +) { + let pilot = SessionPilot::new(Claim::Standard); + let queue = if poll_only { "poll_only" } else { "listening" }; + let mut client_builder = builder() + .pilot(pilot.clone()) + .queue(queue, fast_queue(4)) + .workers(workers()); + if poll_only { + client_builder = client_builder.without_notifications(); + } + let client = client_builder.build().unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + let started = pilot.configurations.snapshot(); + assert_eq!(started.len(), 1); + assert_eq!(started[0].max_workers, 4); + assert_eq!(started[0].queue.name, queue); + + let updater = if poll_only { + client.clone() + } else { + builder().build().unwrap() + }; + for value in 1..=2 { + let metadata = json!({"value": value}); + let Value::Object(map) = metadata.clone() else { + unreachable!() + }; + updater + .queues() + .update(queue, QueueUpdateParams::new().metadata(map.clone())) + .await + .unwrap(); + tokio::time::timeout( + Duration::from_millis(1500), + pilot + .configurations + .wait_until("the metadata change", |configurations| { + configurations + .last() + .is_some_and(|configuration| configuration.queue.metadata == map) + }), + ) + .await + .unwrap_or_else(|_| panic!("metadata {metadata} not reported in time")); + } + run.shutdown().await.unwrap(); +} + +/// A job whose worker reports whether its attempt started cancelled. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "producer_session_cancel_probe")] +struct CancelProbeArgs {} + +/// Holds each claim that returned jobs, after it committed, until released. +#[derive(Clone, Default)] +struct GatedPilot { + claimed: Arc, + release: Arc, +} + +#[async_trait] +impl Pilot for GatedPilot { + async fn start_producer( + &self, + _context: ProducerStartContext, + ) -> Result>, PilotError> { + Ok(Some(Box::new(self.clone()))) + } +} + +#[async_trait] +impl PilotProducer for GatedPilot { + fn intercepts_claim(&self) -> bool { + true + } + + async fn claim( + &self, + context: ProducerClaimContext<'_>, + next: ProducerClaimNext<'_>, + ) -> Result, PilotError> { + let mut transaction = context.database.begin().await?; + let jobs = next.claim(transaction.connection()).await?; + transaction.commit().await?; + if !jobs.is_empty() { + self.claimed.notify_one(); + self.release.notified().await; + } + Ok(jobs) + } +} + +/// A cancellation that arrives after a job is claimed but before its attempt +/// is registered still reaches the attempt, like Go's producer keeping +/// cancellations received during a fetch. The client is poll-only, so +/// cancelling through it signals its producer directly and the cancellation +/// is handled before the claim is released. +async fn assert_cancellation_during_claim_reaches_the_attempt(builder: riverqueue::ClientBuilder) { + let pilot = GatedPilot::default(); + let mut workers = WorkerRegistry::new(); + workers + .register_fn( + |context: WorkContext, _job: Job| async move { + if context.cancellation_token().is_cancelled() { + return Err(std::io::Error::other("started cancelled")); + } + Ok(WorkOutcome::Complete) + }, + ) + .unwrap(); + let client = builder + .pilot(pilot.clone()) + .without_notifications() + .queue("default", fast_queue(1)) + .workers(workers) + .build() + .unwrap(); + let id = client.insert(CancelProbeArgs {}).await.unwrap().id(); + let mut run = client.start().unwrap(); + tokio::time::timeout(WAIT, pilot.claimed.notified()) + .await + .expect("job claimed"); + + let requested = client.jobs().cancel(id).await.unwrap(); + assert_eq!(requested.state, JobState::Running); + pilot.release.notify_one(); + tokio::time::timeout(WAIT, async { + while client.jobs().get(id).await.unwrap().state == JobState::Running { + tokio::task::yield_now().await; + } + }) + .await + .expect("job finishes"); + run.shutdown().await.unwrap(); + + assert_eq!( + client.jobs().get(id).await.unwrap().state, + JobState::Cancelled + ); +} + +/// The session gets the queue metadata's stored text, as the database +/// renders it, at start and whenever it changes, including a change the +/// parsed metadata can't show. `store` writes the stored metadata from JSON +/// text; `first` and `second` parse to the same map but render differently. +async fn assert_sessions_see_metadata_text( + builder: riverqueue::ClientBuilder, + store: F, + first: &str, + second: &str, +) where + F: Fn(&'static str) -> Fut, + Fut: std::future::Future, +{ + let pilot = SessionPilot::new(Claim::Standard); + // The store writes no control notification, so only a client without a + // notifier, which polls the queue's record, sees the change. + let client = builder + .pilot(pilot.clone()) + .without_notifications() + .queue("texted", fast_queue(1)) + .workers(workers()) + .build() + .unwrap(); + store("first").await; + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + let started = pilot.configurations.snapshot(); + assert_eq!(started[0].metadata_text, first); + + store("second").await; + // Without a notifier, the queue's record is read again every two + // seconds. + pilot + .configurations + .wait_until("the new metadata text", |configurations| { + configurations + .last() + .is_some_and(|configuration| configuration.metadata_text == second) + }) + .await; + let configurations = pilot.configurations.snapshot(); + assert_eq!( + configurations.first().unwrap().queue.metadata, + configurations.last().unwrap().queue.metadata, + "only the text changed" + ); + run.shutdown().await.unwrap(); +} + +#[cfg(all(feature = "postgres", river_postgres_tests))] +mod postgres { + use riverqueue::database::PostgresDatabase; + + use super::*; + use crate::support::PostgresSchema; + + fn builder(schema: &PostgresSchema) -> riverqueue::ClientBuilder { + Client::builder( + PostgresDatabase::new(schema.pool.clone()).with_schema(schema.schema.clone()), + ) + } + + #[tokio::test(flavor = "multi_thread")] + async fn panicked_attempts_finish() { + let schema = PostgresSchema::new("session_panicked").await; + assert_panicked_attempts_finish(builder(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn abandoned_attempts_finish() { + let schema = PostgresSchema::new("session_abandoned").await; + assert_abandoned_attempts_finish(builder(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn cancellation_during_claim_reaches_the_attempt() { + let schema = PostgresSchema::new("session_claim_cancel").await; + assert_cancellation_during_claim_reaches_the_attempt(builder(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn broken_claims_stop_the_client() { + for claim in BROKEN_CLAIMS { + // The current schema, so the unqualified raw claims find it. + let schema = PostgresSchema::current("session_broken").await; + // Metadata that isn't an object can't be decoded. + let (pool, table) = (schema.pool.clone(), schema.table("river_job")); + assert_broken_claims_stop_the_client( + || builder(&schema), + claim, + |id| async move { + sqlx::query(sqlx::AssertSqlSafe(format!( + "UPDATE {table} SET metadata = '[1]'::jsonb WHERE id = $1" + ))) + .bind(id) + .execute(&pool) + .await + .unwrap(); + }, + ) + .await; + schema.cleanup().await; + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn every_attempt_finishes_once() { + let schema = PostgresSchema::new("session_finished").await; + assert_every_attempt_finishes_once(builder(&schema)).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn queue_changes_reach_the_session() { + let schema = PostgresSchema::new("session_config").await; + assert_queue_changes_reach_the_session(|| builder(&schema), false).await; + assert_queue_changes_reach_the_session(|| builder(&schema), true).await; + schema.cleanup().await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn sessions_see_metadata_text() { + let schema = PostgresSchema::new("session_metadata_text").await; + let table = schema.table("river_queue"); + let pool = schema.pool.clone(); + assert_sessions_see_metadata_text( + builder(&schema), + |which| { + let (table, pool) = (table.clone(), pool.clone()); + async move { + let metadata = if which == "first" { + r#"{"n": 1.0}"# + } else { + r#"{"n": 1.00}"# + }; + sqlx::query(sqlx::AssertSqlSafe(format!( + "INSERT INTO {table} (name, created_at, metadata, updated_at) \ + VALUES ('texted', now(), $1::jsonb, now()) \ + ON CONFLICT (name) DO UPDATE SET metadata = excluded.metadata" + ))) + .bind(metadata) + .execute(&pool) + .await + .unwrap(); + } + }, + r#"{"n": 1.0}"#, + r#"{"n": 1.00}"#, + ) + .await; + schema.cleanup().await; + } +} + +#[cfg(feature = "sqlite")] +mod sqlite { + use super::*; + use crate::support::{sqlite_cleanup, sqlite_file_pool}; + + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn abandoned_attempts_finish() { + let (pool, path) = sqlite_file_pool(4).await; + assert_abandoned_attempts_finish(Client::builder(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn cancellation_during_claim_reaches_the_attempt() { + let (pool, path) = sqlite_file_pool(4).await; + assert_cancellation_during_claim_reaches_the_attempt(Client::builder(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn broken_claims_stop_the_client() { + for claim in BROKEN_CLAIMS { + let (pool, path) = sqlite_file_pool(4).await; + // Tags that aren't an array can't be decoded. + let corrupt_pool = pool.clone(); + assert_broken_claims_stop_the_client( + || Client::builder(pool.clone()), + claim, + |id| async move { + sqlx::query( + "UPDATE river_job SET tags = jsonb('{\"not\":\"an array\"}') WHERE id = ?", + ) + .bind(id) + .execute(&corrupt_pool) + .await + .unwrap(); + }, + ) + .await; + sqlite_cleanup(pool, path).await; + } + } + + #[tokio::test(flavor = "multi_thread")] + async fn panicked_attempts_finish() { + let (pool, path) = sqlite_file_pool(4).await; + assert_panicked_attempts_finish(Client::builder(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn every_attempt_finishes_once() { + let (pool, path) = sqlite_file_pool(4).await; + assert_every_attempt_finishes_once(Client::builder(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn queue_changes_reach_the_session() { + let (pool, path) = sqlite_file_pool(4).await; + assert_queue_changes_reach_the_session(|| Client::builder(pool.clone()), false).await; + assert_queue_changes_reach_the_session(|| Client::builder(pool.clone()), true).await; + sqlite_cleanup(pool, path).await; + } + + #[tokio::test(flavor = "multi_thread")] + async fn sessions_see_metadata_text() { + let (pool, path) = sqlite_file_pool(4).await; + let store_pool = pool.clone(); + assert_sessions_see_metadata_text( + Client::builder(pool.clone()), + |which| { + let pool = store_pool.clone(); + async move { + let metadata = if which == "first" { + r#"{"b":1,"a":2}"# + } else { + r#"{"a":2,"b":1}"# + }; + sqlx::query( + "INSERT INTO river_queue (name, metadata) VALUES ('texted', jsonb(?)) \ + ON CONFLICT (name) DO UPDATE SET metadata = excluded.metadata", + ) + .bind(metadata) + .execute(&pool) + .await + .unwrap(); + } + }, + r#"{"b":1,"a":2}"#, + r#"{"a":2,"b":1}"#, + ) + .await; + sqlite_cleanup(pool, path).await; + } +} diff --git a/rust/riverqueue/tests/protocol_fixtures.rs b/rust/riverqueue/tests/protocol_fixtures.rs new file mode 100644 index 000000000..82f645aa5 --- /dev/null +++ b/rust/riverqueue/tests/protocol_fixtures.rs @@ -0,0 +1,126 @@ +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use riverqueue::{ + AttemptError, DefaultRetryPolicy, JobRow, JobState, METADATA_KEY_OUTPUT, + METADATA_KEY_PERIODIC_JOB_ID, METADATA_KEY_RESCUE_COUNT, METADATA_KEY_RESUMABLE_CURSOR, + METADATA_KEY_RESUMABLE_STEP, METADATA_KEY_UNIQUE_NONCE, RetryPolicy, + protocol::{ + NOTIFICATION_TOPIC_CONTROL, NOTIFICATION_TOPIC_INSERT, NOTIFICATION_TOPIC_LEADERSHIP, + unique_state_bit, + }, +}; +use serde::Deserialize; +use serde_json::{Map, Value}; + +#[derive(Deserialize)] +struct Fixture { + attempt_error: AttemptError, + job_states: Vec, + metadata_keys: Map, + notifications: Vec, + retry_cases: Vec, + topics: Map, +} + +#[derive(Deserialize)] +struct NotificationFixture { + name: String, + payload: Map, + topic: String, +} + +#[derive(Deserialize)] +struct RetryFixture { + error_count: usize, + job_id: i64, + max_delay_ns: u64, + min_delay_ns: u64, + now: DateTime, + seed: u64, +} + +#[derive(Deserialize)] +struct StateFixture { + state: JobState, + unique_bit: u8, +} + +#[test] +fn go_protocol_values_match_rust() { + let fixture: Fixture = + serde_json::from_str(include_str!("fixtures/protocol_values.json")).unwrap(); + + assert_eq!(fixture.attempt_error.attempt, 3); + assert!(fixture.attempt_error.error.contains("escaped")); + assert_eq!(fixture.job_states.len(), JobState::ALL.len()); + for state in fixture.job_states { + assert_eq!(state.unique_bit, unique_state_bit(state.state)); + } + for (name, expected) in [ + ("output", METADATA_KEY_OUTPUT), + ("periodic_job_id", METADATA_KEY_PERIODIC_JOB_ID), + ("rescue_count", METADATA_KEY_RESCUE_COUNT), + ("resumable_cursor", METADATA_KEY_RESUMABLE_CURSOR), + ("resumable_step", METADATA_KEY_RESUMABLE_STEP), + ("unique_nonce", METADATA_KEY_UNIQUE_NONCE), + ] { + assert_eq!(fixture.metadata_keys[name], expected); + } + assert_eq!(fixture.topics["control"], NOTIFICATION_TOPIC_CONTROL); + assert_eq!(fixture.topics["insert"], NOTIFICATION_TOPIC_INSERT); + assert_eq!(fixture.topics["leadership"], NOTIFICATION_TOPIC_LEADERSHIP); + for notification in fixture.notifications { + assert_ne!(notification.name, ""); + assert!(notification.payload.contains_key("action") || notification.name == "insert"); + assert!( + [ + NOTIFICATION_TOPIC_CONTROL, + NOTIFICATION_TOPIC_INSERT, + NOTIFICATION_TOPIC_LEADERSHIP, + ] + .contains(¬ification.topic.as_str()) + ); + } + + for test_case in fixture.retry_cases { + let row = retry_row(test_case.job_id, test_case.now, test_case.error_count - 1); + let delay = DefaultRetryPolicy::with_seed(test_case.seed).next_retry( + &row, + &riverqueue::WorkError::new("fixture failure"), + test_case.now, + ); + let delay = delay.as_nanos(); + assert!( + (u128::from(test_case.min_delay_ns)..=u128::from(test_case.max_delay_ns)) + .contains(&delay), + "error count {} delay {delay}ns outside Go's bounds", + test_case.error_count + ); + } +} + +fn retry_row(id: i64, now: DateTime, previous_errors: usize) -> JobRow { + let mut row = JobRow::new( + id, + "fixture_retry", + riverqueue::encoding::encode_args(&serde_json::json!({})).unwrap(), + now, + ); + row.attempt = i16::try_from(previous_errors + 1).unwrap(); + row.attempted_at = Some(now); + row.attempted_by = vec!["fixture".to_owned()]; + row.errors = vec![AttemptError::new(now, 1, "previous failure"); previous_errors]; + row.max_attempts = 1_000; + row.metadata = Map::new().into(); + row.state = JobState::Retryable; + row +} + +#[test] +fn retry_duration_cap_matches_go_time_duration() { + assert_eq!( + Duration::from_nanos(i64::MAX as u64).as_nanos(), + 9_223_372_036_854_775_807 + ); +} diff --git a/rust/riverqueue/tests/resilience_postgres.rs b/rust/riverqueue/tests/resilience_postgres.rs new file mode 100644 index 000000000..21128e1b5 --- /dev/null +++ b/rust/riverqueue/tests/resilience_postgres.rs @@ -0,0 +1,1347 @@ +//! PostgreSQL runtime behavior under malformed rows and database faults. +//! +//! These tests require `RIVER_RUST_DATABASE_URL` and fail when it is missing. +//! Each test migrates a uniquely named schema so concurrent runs against one +//! database cannot interfere. + +#![cfg(all(feature = "postgres", river_postgres_tests))] + +mod support; + +use std::{ + convert::Infallible, + future::Future, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use riverqueue::__private::ClientBuilderExt; +use riverqueue::{ + __private::{DatabaseConnection, JobSetStateParams, Pilot, PilotError}, + AttemptError, BoxError, Client, ErrorHandler, ErrorHandlerDecision, EventKind, InsertOpts, Job, + JobArgs, JobEventKind, JobRow, JobState, QueueConfig, RetryPolicy, WorkCancelled, WorkContext, + WorkOutcome, WorkResult, WorkerRegistry, + database::{PostgresDatabase, SchemaName}, +}; +use riverqueue_migrate::PostgresMigrator; +use serde::{Deserialize, Serialize}; +use sqlx::{ + AssertSqlSafe, PgPool, + postgres::{PgConnectOptions, PgPoolOptions}, +}; +use tokio::{net::TcpListener, sync::Semaphore, task::AbortHandle}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_postgres_resilience")] +struct ResilienceArgs {} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_postgres_resilience_gated")] +struct GatedArgs {} + +/// A job that waits for client shutdown and then stops in the given way. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_postgres_resilience_shutdown")] +struct ShutdownArgs { + behavior: String, +} + +#[derive(Debug, thiserror::Error)] +enum ShutdownError { + #[error("stopped for shutdown")] + Cancelled(#[source] WorkCancelled), + #[error("real failure during shutdown")] + Real, +} + +/// A job that blocks its thread without yielding, so Tokio cannot abort it +/// until the blocking section ends. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_postgres_resilience_blocking")] +struct BlockingArgs { + block_ms: u64, +} + +/// Records when blocking work finished and when later work started. +#[derive(Default)] +struct BlockingTimeline { + blocking_finished: std::sync::Mutex>, + later_started: std::sync::Mutex>, +} + +/// Signals every stuck job. +#[derive(Clone)] +struct StuckSignal(Arc); + +#[allow( + clippy::unused_async_trait_impl, + reason = "these extensions only record state synchronously" +)] +impl ErrorHandler for StuckSignal { + async fn handle_stuck(&self, _job: &JobRow) -> Result<(), BoxError> { + self.0.add_permits(1); + Ok(()) + } +} + +fn blocking_workers(gate: &Gate, timeline: &Arc) -> WorkerRegistry { + let mut workers = WorkerRegistry::new(); + let blocking_gate = gate.clone(); + let blocking_timeline = Arc::clone(timeline); + workers + .register_fn(move |_context: WorkContext, job: Job| { + let gate = blocking_gate.clone(); + let timeline = Arc::clone(&blocking_timeline); + async move { + gate.started.add_permits(1); + // Deliberately ignore cancellation without an await point. + std::thread::sleep(Duration::from_millis(job.args.block_ms)); + *timeline.blocking_finished.lock().unwrap() = Some(std::time::Instant::now()); + tokio::task::yield_now().await; + Ok::<_, Infallible>(WorkOutcome::Complete) + } + }) + .unwrap(); + let later_timeline = Arc::clone(timeline); + workers + .register_fn(move |_context: WorkContext, _job: Job| { + let timeline = Arc::clone(&later_timeline); + async move { + timeline + .later_started + .lock() + .unwrap() + .get_or_insert_with(std::time::Instant::now); + Ok::<_, Infallible>(WorkOutcome::Complete) + } + }) + .unwrap(); + workers +} + +/// A job that snoozes for longer than any representable schedule. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_postgres_resilience_snooze")] +struct SnoozeForeverArgs {} + +/// Lets a test hold a gated job inside its worker until released. +#[derive(Clone)] +struct Gate { + active: Arc, + max_active: Arc, + release: Arc, + started: Arc, +} + +impl Default for Gate { + fn default() -> Self { + Self { + active: Arc::new(AtomicUsize::new(0)), + max_active: Arc::new(AtomicUsize::new(0)), + release: Arc::new(Semaphore::new(0)), + started: Arc::new(Semaphore::new(0)), + } + } +} + +impl Gate { + async fn wait_started(&self) { + tokio::time::timeout(Duration::from_secs(10), self.started.acquire()) + .await + .expect("gated job did not start") + .unwrap() + .forget(); + } + + /// Fails if a gated job starts within `window`. + async fn assert_none_started(&self, window: Duration) { + assert!( + tokio::time::timeout(window, self.started.acquire()) + .await + .is_err(), + "a gated job started" + ); + } + + fn release(&self) { + self.release.add_permits(1); + } +} + +fn database_url() -> String { + std::env::var("RIVER_RUST_DATABASE_URL") + .expect("RIVER_RUST_DATABASE_URL must point at a disposable PostgreSQL test database") +} + +/// A migrated schema with a unique name, dropped by [`TestSchema::drop`]. +struct TestSchema { + name: String, + pool: PgPool, + schema: SchemaName, +} + +impl TestSchema { + async fn new(label: &str) -> Self { + let mut name = format!("river_res_{label}_{}", support::unique_suffix()); + name.truncate(riverqueue::migrate::SCHEMA_MAX_LEN); + let pool = PgPool::connect(&database_url()).await.unwrap(); + sqlx::raw_sql(AssertSqlSafe(format!("CREATE SCHEMA {name}"))) + .execute(&pool) + .await + .unwrap(); + let schema = SchemaName::new(name.clone()).unwrap(); + PostgresMigrator::new(pool.clone()) + .with_schema(schema.clone()) + .migrate_up() + .await + .unwrap(); + Self { name, pool, schema } + } + + fn database(&self) -> PostgresDatabase { + PostgresDatabase::new(self.pool.clone()).with_schema(self.schema.clone()) + } + + fn table(&self) -> String { + format!("{}.river_job", self.name) + } + + async fn execute(&self, sql: String) { + sqlx::raw_sql(AssertSqlSafe(sql)) + .execute(&self.pool) + .await + .unwrap(); + } + + async fn job_attempt(&self, id: i64) -> i16 { + sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT attempt FROM {} WHERE id = $1", + self.table() + ))) + .bind(id) + .fetch_one(&self.pool) + .await + .unwrap() + } + + async fn job_state(&self, id: i64) -> String { + sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT state::text FROM {} WHERE id = $1", + self.table() + ))) + .bind(id) + .fetch_one(&self.pool) + .await + .unwrap() + } + + async fn drop(self) { + self.execute(format!("DROP SCHEMA {} CASCADE", self.name)) + .await; + self.pool.close().await; + } +} + +/// A TCP proxy between a client and PostgreSQL that can make the database +/// unavailable: it resets open connections and refuses new ones until +/// restored. Unlike terminating backends, this keeps the database down for +/// the client while other connections still work. +struct FaultProxy { + accept_task: AbortHandle, + connections: Arc>>, + options: PgConnectOptions, + rejected: Arc, + up: Arc, +} + +impl FaultProxy { + async fn start() -> Self { + let upstream: PgConnectOptions = database_url().parse().unwrap(); + let upstream_host = upstream.get_host().to_owned(); + let upstream_port = upstream.get_port(); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let connections = Arc::new(std::sync::Mutex::new(Vec::::new())); + let rejected = Arc::new(AtomicUsize::new(0)); + let up = Arc::new(std::sync::atomic::AtomicBool::new(true)); + let accept_connections = Arc::clone(&connections); + let accept_rejected = Arc::clone(&rejected); + let accept_up = Arc::clone(&up); + let accept_task = tokio::spawn(async move { + loop { + let Ok((mut client, _)) = listener.accept().await else { + continue; + }; + if !accept_up.load(Ordering::SeqCst) { + accept_rejected.fetch_add(1, Ordering::SeqCst); + drop(client); + continue; + } + let host = upstream_host.clone(); + let connection = tokio::spawn(async move { + if host.starts_with('/') { + let path = format!("{host}/.s.PGSQL.{upstream_port}"); + if let Ok(mut server) = tokio::net::UnixStream::connect(path).await { + let _ = tokio::io::copy_bidirectional(&mut client, &mut server).await; + } + } else if let Ok(mut server) = + tokio::net::TcpStream::connect((host.as_str(), upstream_port)).await + { + let _ = tokio::io::copy_bidirectional(&mut client, &mut server).await; + } + }); + accept_connections + .lock() + .unwrap() + .push(connection.abort_handle()); + } + }) + .abort_handle(); + let options = upstream.host("127.0.0.1").port(address.port()); + Self { + accept_task, + connections, + options, + rejected, + up, + } + } + + /// A pool that reaches PostgreSQL only through the proxy. A short acquire + /// timeout keeps operations failing quickly while the database is down. + fn pool(&self, max_connections: u32) -> PgPool { + PgPoolOptions::new() + .acquire_timeout(Duration::from_millis(500)) + .max_connections(max_connections) + .connect_lazy_with(self.options.clone()) + } + + fn take_down(&self) { + self.up.store(false, Ordering::SeqCst); + for connection in self.connections.lock().unwrap().drain(..) { + connection.abort(); + } + } + + fn restore(&self) { + self.up.store(true, Ordering::SeqCst); + } + + async fn wait_for_rejections(&self, count: usize) { + wait_until(Duration::from_secs(30), "reconnection attempts", || async { + self.rejected.load(Ordering::SeqCst) >= count + }) + .await; + } +} + +impl Drop for FaultProxy { + fn drop(&mut self) { + self.accept_task.abort(); + for connection in self.connections.lock().unwrap().drain(..) { + connection.abort(); + } + } +} + +fn completing_workers() -> WorkerRegistry { + gated_workers(&Gate::default()) +} + +fn gated_workers(gate: &Gate) -> WorkerRegistry { + let mut workers = WorkerRegistry::new(); + workers + .register_fn(|_context: WorkContext, _job: Job| async { + Ok::<_, Infallible>(WorkOutcome::Complete) + }) + .unwrap(); + workers + .register_fn( + |_context: WorkContext, _job: Job| async { + Ok::<_, Infallible>(WorkOutcome::Snooze(Duration::MAX)) + }, + ) + .unwrap(); + let shutdown_gate = gate.clone(); + workers + .register_fn(move |context: WorkContext, job: Job| { + let gate = shutdown_gate.clone(); + async move { + gate.started.add_permits(1); + context.cancellation_token().cancelled().await; + match job.args.behavior.as_str() { + "cooperative" => Err(ShutdownError::Cancelled(WorkCancelled)), + "error" => Err(ShutdownError::Real), + "panic" => panic!("panic during shutdown"), + behavior => unreachable!("unknown shutdown behavior {behavior}"), + } + } + }) + .unwrap(); + let gate = gate.clone(); + workers + .register_fn(move |_context: WorkContext, _job: Job| { + let gate = gate.clone(); + async move { + let active = gate.active.fetch_add(1, Ordering::SeqCst) + 1; + gate.max_active.fetch_max(active, Ordering::SeqCst); + gate.started.add_permits(1); + gate.release.acquire().await.unwrap().forget(); + gate.active.fetch_sub(1, Ordering::SeqCst); + Ok::<_, Infallible>(WorkOutcome::Complete) + } + }) + .unwrap(); + workers +} + +fn fast_queue() -> QueueConfig { + QueueConfig::new(4) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(20)) +} + +/// Polls `condition` until it holds, failing the test after `timeout`. +async fn wait_until(timeout: Duration, description: &str, mut condition: F) +where + F: FnMut() -> Fut, + Fut: Future, +{ + let deadline = tokio::time::Instant::now() + timeout; + while !condition().await { + assert!( + tokio::time::Instant::now() < deadline, + "timed out waiting for {description}" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } +} + +/// Records every failed attempt the error handler sees. +#[derive(Clone, Default)] +struct RecordingErrorHandler(Arc>>); + +#[allow( + clippy::unused_async_trait_impl, + reason = "these extensions only record state synchronously" +)] +impl ErrorHandler for RecordingErrorHandler { + async fn handle_error( + &self, + _context: &WorkContext, + job: &JobRow, + result: &WorkResult, + ) -> Result { + let error = match result { + WorkResult::Failed(error) => error.to_string(), + other => format!("{other:?}"), + }; + self.0.lock().unwrap().push((job.clone(), error)); + Ok(ErrorHandlerDecision::Continue) + } +} + +/// Schedules every retry an hour out so a failed job stays `retryable`. +struct RetryAnHourLater; + +impl RetryPolicy for RetryAnHourLater { + fn next_retry( + &self, + _job: &JobRow, + _error: &riverqueue::WorkError, + _now: DateTime, + ) -> Duration { + Duration::from_hours(1) + } +} + +#[tokio::test(flavor = "multi_thread")] +#[allow(clippy::too_many_lines)] +async fn claimed_rows_decode_individually_and_leniently() { + let schema = TestSchema::new("decode").await; + let error_handler = RecordingErrorHandler::default(); + let client = Client::builder(schema.database()) + .id("postgres-resilience-decode") + .without_notifications() + .error_handler(error_handler.clone()) + .retry_policy(RetryAnHourLater) + .workers(completing_workers()) + .queue("default", fast_queue()) + .build() + .unwrap(); + let mut events = client + .subscribe(&[EventKind::JobCompleted, EventKind::JobFailed]) + .unwrap(); + + // River Go decodes attempt errors with `encoding/json`, which tolerates + // missing and unknown fields. + let sparse_errors = client.insert(ResilienceArgs {}).await.unwrap(); + schema + .execute(format!( + "UPDATE {} SET errors = ARRAY['{{\"error\": \"go\", \"extra\": 1}}'::jsonb] \ + WHERE id = {}", + schema.table(), + sparse_errors.job.row.id + )) + .await; + // Attempt errors in a shape River doesn't write decode leniently like + // River Go's, so the job is still worked. + let odd_errors = client.insert(ResilienceArgs {}).await.unwrap(); + schema + .execute(format!( + "UPDATE {} SET errors = ARRAY[\ + '{{\"at\": \"2024-01-02 03:04:05+00\", \"attempt\": \"1\", \ + \"error\": {{\"message\": \"boom\"}}, \"trace\": [\"frame\"]}}'::jsonb, \ + '42'::jsonb] \ + WHERE id = {}", + schema.table(), + odd_errors.job.row.id + )) + .await; + // Array metadata can't become a `JobRow`. Claimed with the others, such a + // job isn't worked, and its attempt fails like any other: retried with + // the client's retry policy, or discarded at its maximum attempts. + let malformed_retried = client.insert(ResilienceArgs {}).await.unwrap(); + let malformed_discarded = client + .insert(ResilienceArgs {}) + .opts(InsertOpts::default().with_max_attempts(1)) + .await + .unwrap(); + let malformed_ids = [malformed_retried.job.row.id, malformed_discarded.job.row.id]; + for id in malformed_ids { + schema + .execute(format!( + "UPDATE {} SET metadata = '[1]'::jsonb WHERE id = {id}", + schema.table() + )) + .await; + } + let ordinary = client.insert(ResilienceArgs {}).await.unwrap(); + let decodable_ids = [ + sparse_errors.job.row.id, + odd_errors.job.row.id, + ordinary.job.row.id, + ]; + + let mut run = client.start().unwrap(); + let mut events_by_id = std::collections::HashMap::new(); + while events_by_id.len() < decodable_ids.len() + malformed_ids.len() { + let event = tokio::time::timeout(Duration::from_secs(10), events.recv()) + .await + .expect("job events") + .unwrap(); + let event = event.as_job().unwrap().clone(); + events_by_id.insert(event.job.id, event); + } + run.shutdown().await.unwrap(); + + for id in decodable_ids { + assert_eq!(events_by_id[&id].kind, JobEventKind::Completed); + } + let sparse_errors = client.jobs().get(sparse_errors.job.row.id).await.unwrap(); + assert_eq!(sparse_errors.errors.len(), 1); + assert_eq!(sparse_errors.errors[0].error, "go"); + assert_eq!(sparse_errors.errors[0].attempt, 0); + let odd_errors = client.jobs().get(odd_errors.job.row.id).await.unwrap(); + assert_eq!( + odd_errors.errors, + [ + // Like Go's `time.Time`, `at` accepts only RFC 3339. + AttemptError::new( + "0001-01-01T00:00:00Z".parse().unwrap(), + 1, + r#"{"message":"boom"}"# + ) + .with_trace(r#"["frame"]"#), + AttemptError::new("0001-01-01T00:00:00Z".parse().unwrap(), 0, "42"), + ] + ); + + // Failed events and the error handler carry the fields that could be + // decoded, with the metadata left empty. + let handled = error_handler.0.lock().unwrap().clone(); + assert_eq!(handled.len(), malformed_ids.len()); + for (id, state) in [ + (malformed_retried.job.row.id, JobState::Retryable), + (malformed_discarded.job.row.id, JobState::Discarded), + ] { + let event = &events_by_id[&id]; + assert_eq!(event.kind, JobEventKind::Failed); + assert_eq!(event.job.state, state); + assert_eq!(event.job.kind, ResilienceArgs::KIND); + assert!(event.job.metadata.is_empty()); + assert!( + event.job.errors[0] + .error + .starts_with("job row couldn't be decoded: error unmarshaling `metadata`: "), + "{:?}", + event.job.errors + ); + + let (job, error) = handled.iter().find(|(job, _)| job.id == id).unwrap(); + assert_eq!(job.attempt, 1); + assert!(job.metadata.is_empty()); + assert_eq!(error, &event.job.errors[0].error); + + // The attempt error is appended without rewriting the undecodable + // metadata. + let (state, attempt, errors, metadata, scheduled_at): ( + String, + i16, + Vec, + serde_json::Value, + DateTime, + ) = sqlx::query_as(AssertSqlSafe(format!( + "SELECT state::text, attempt, errors, metadata, scheduled_at FROM {} WHERE id = $1", + schema.table() + ))) + .bind(id) + .fetch_one(&schema.pool) + .await + .unwrap(); + assert_eq!(state, event.job.state.as_str()); + assert_eq!(attempt, 1); + assert_eq!(errors.len(), 1); + assert_eq!(errors[0]["attempt"], 1); + assert_eq!(errors[0]["error"], event.job.errors[0].error.as_str()); + assert_eq!(metadata, serde_json::json!([1])); + if state == "retryable" { + assert!( + scheduled_at > Utc::now() + chrono::Duration::minutes(50), + "the client retry policy wasn't used: {scheduled_at}" + ); + } + } + + schema.drop().await; +} + +fn gated_client(schema: &TestSchema, id: &str, gate: &Gate) -> Client { + Client::builder(schema.database()) + .id(id) + .without_notifications() + .workers(gated_workers(gate)) + .queue("default", fast_queue()) + .build() + .unwrap() +} + +#[tokio::test(flavor = "multi_thread")] +async fn completion_retries_a_transient_database_error() { + let schema = TestSchema::new("retry").await; + // Fail the first running-to-completed transition with a serialization + // failure. A sequence records the injection outside the aborted statement. + schema + .execute(format!( + "CREATE SEQUENCE {name}.completion_fault; \ + CREATE FUNCTION {name}.fail_completion_once() RETURNS trigger \ + LANGUAGE plpgsql AS $$ BEGIN \ + IF OLD.state = 'running' AND NEW.state = 'completed' \ + AND nextval('{name}.completion_fault') = 1 THEN \ + RAISE EXCEPTION 'injected completion failure' USING ERRCODE = '40001'; \ + END IF; \ + RETURN NEW; \ + END $$; \ + CREATE TRIGGER fail_completion_once BEFORE UPDATE ON {name}.river_job \ + FOR EACH ROW EXECUTE FUNCTION {name}.fail_completion_once()", + name = schema.name + )) + .await; + let client = gated_client(&schema, "postgres-resilience-retry", &Gate::default()); + let job = client.insert(ResilienceArgs {}).await.unwrap(); + + let mut run = client.start().unwrap(); + wait_until( + Duration::from_secs(10), + "completion after retry", + || async { schema.job_state(job.job.row.id).await == "completed" }, + ) + .await; + run.shutdown().await.unwrap(); + + let injected: i64 = sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT last_value FROM {}.completion_fault", + schema.name + ))) + .fetch_one(&schema.pool) + .await + .unwrap(); + assert!(injected >= 2, "the injected failure never fired"); + let job = client.jobs().get(job.job.row.id).await.unwrap(); + assert_eq!(job.attempt, 1); + assert_eq!(job.errors, []); + + schema.drop().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn completion_waits_for_a_row_lock() { + let schema = TestSchema::new("lock").await; + let gate = Gate::default(); + let client = gated_client(&schema, "postgres-resilience-lock", &gate); + let job = client.insert(GatedArgs {}).await.unwrap(); + + let mut run = client.start().unwrap(); + gate.wait_started().await; + let mut locker = schema.pool.begin().await.unwrap(); + sqlx::query(AssertSqlSafe(format!( + "SELECT 1 FROM {} WHERE id = $1 FOR UPDATE", + schema.table() + ))) + .bind(job.job.row.id) + .execute(&mut *locker) + .await + .unwrap(); + gate.release(); + wait_until( + Duration::from_secs(10), + "completion to wait on the lock", + || async { + let waiting: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pg_locks WHERE NOT granted AND locktype = 'transactionid'", + ) + .fetch_one(&schema.pool) + .await + .unwrap(); + waiting > 0 + }, + ) + .await; + locker.commit().await.unwrap(); + wait_until( + Duration::from_secs(10), + "completion after unlock", + || async { schema.job_state(job.job.row.id).await == "completed" }, + ) + .await; + run.shutdown().await.unwrap(); + + schema.drop().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn completion_leaves_rows_moved_out_of_running_and_keeps_working() { + let schema = TestSchema::new("moved").await; + let gate = Gate::default(); + let client = gated_client(&schema, "postgres-resilience-moved", &gate); + let pending = client.insert(GatedArgs {}).await.unwrap(); + + let mut run = client.start().unwrap(); + gate.wait_started().await; + // An operator or extension moves the running job back to `pending`. + schema + .execute(format!( + "UPDATE {} SET state = 'pending' WHERE id = {}", + schema.table(), + pending.job.row.id + )) + .await; + gate.release(); + let later = client.insert(ResilienceArgs {}).await.unwrap(); + wait_until( + Duration::from_secs(10), + "a later job to complete", + || async { schema.job_state(later.job.row.id).await == "completed" }, + ) + .await; + run.shutdown().await.unwrap(); + + assert_eq!(schema.job_state(pending.job.row.id).await, "pending"); + + schema.drop().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn completion_does_not_rewrite_a_newer_attempt_number() { + let schema = TestSchema::new("attempt").await; + let gate = Gate::default(); + let client = gated_client(&schema, "postgres-resilience-attempt", &gate); + let job = client.insert(GatedArgs {}).await.unwrap(); + + let mut run = client.start().unwrap(); + gate.wait_started().await; + // Another client rescued and refetched the job while this attempt ran. + schema + .execute(format!( + "UPDATE {} SET attempt = 5 WHERE id = {}", + schema.table(), + job.job.row.id + )) + .await; + gate.release(); + wait_until(Duration::from_secs(10), "stale completion", || async { + schema.job_state(job.job.row.id).await == "completed" + }) + .await; + run.shutdown().await.unwrap(); + + assert_eq!(schema.job_attempt(job.job.row.id).await, 5); + + schema.drop().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn hard_shutdown_interrupts_only_cooperative_cancellations() { + let schema = TestSchema::new("shutdown").await; + let gate = Gate::default(); + let client = gated_client(&schema, "postgres-resilience-shutdown", &gate); + let insert = |behavior: &str| { + client.insert(ShutdownArgs { + behavior: behavior.to_owned(), + }) + }; + let cooperative = insert("cooperative").await.unwrap(); + let cancel_attempted = insert("cooperative").await.unwrap(); + let real_error = insert("error").await.unwrap(); + let panicked = insert("panic").await.unwrap(); + + let mut listener = sqlx::postgres::PgListener::connect_with(&schema.pool) + .await + .unwrap(); + listener + .listen(&format!("{}.river_insert", schema.name)) + .await + .unwrap(); + let mut run = client.start().unwrap(); + for _ in 0..4 { + gate.wait_started().await; + } + // A cancellation whose notification never reached this client. + schema + .execute(format!( + "UPDATE {} SET metadata = jsonb_set(metadata, '{{cancel_attempted_at}}', \ + to_jsonb('2026-01-02T03:04:05Z'::text)) WHERE id = {}", + schema.table(), + cancel_attempted.job.row.id + )) + .await; + tokio::time::timeout(Duration::from_secs(10), run.shutdown_now()) + .await + .unwrap() + .unwrap(); + + let cooperative = client.jobs().get(cooperative.job.row.id).await.unwrap(); + assert_eq!(cooperative.state, JobState::Available); + assert_eq!(cooperative.attempt, 0); + assert!(cooperative.attempted_at.is_some(), "attempted_at is kept"); + assert_eq!(cooperative.errors, []); + let notification = tokio::time::timeout(Duration::from_secs(5), listener.recv()) + .await + .expect("interrupted job did not notify peers") + .unwrap(); + assert_eq!(notification.payload(), r#"{"queue" : "default"}"#); + drop(listener); + + let cancel_attempted = client + .jobs() + .get(cancel_attempted.job.row.id) + .await + .unwrap(); + assert_eq!(cancel_attempted.state, JobState::Cancelled); + assert!(cancel_attempted.finalized_at.is_some()); + + for (job, error) in [ + (real_error, "real failure during shutdown"), + (panicked, "panic during shutdown"), + ] { + let job = client.jobs().get(job.job.row.id).await.unwrap(); + assert!( + matches!(job.state, JobState::Available | JobState::Retryable), + "{:?}", + job.state + ); + assert_eq!(job.attempt, 1, "a genuine failure consumes its attempt"); + assert_eq!(job.errors.len(), 1); + assert!(job.errors[0].error.contains(error), "{:?}", job.errors); + } + + schema.drop().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn stuck_job_keeps_its_worker_slot_until_it_ends() { + let schema = TestSchema::new("stuck").await; + let gate = Gate::default(); + let timeline = Arc::new(BlockingTimeline::default()); + let stuck = StuckSignal(Arc::new(Semaphore::new(0))); + let client = Client::builder(schema.database()) + .id("postgres-resilience-stuck") + .error_handler(stuck.clone()) + .job_stuck_threshold(Duration::from_millis(50)) + .job_timeout(Duration::from_millis(100)) + .without_notifications() + .workers(blocking_workers(&gate, &timeline)) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(20)), + ) + .build() + .unwrap(); + let blocked = client + .insert(BlockingArgs { block_ms: 1_500 }) + .await + .unwrap(); + let later = client + .insert(ResilienceArgs {}) + .opts(InsertOpts::default().with_priority(2)) + .await + .unwrap(); + + let mut run = client.start().unwrap(); + gate.wait_started().await; + tokio::time::timeout(Duration::from_secs(5), stuck.0.acquire()) + .await + .expect("stuck handler was not invoked") + .unwrap() + .forget(); + // The aborted task is still blocking its thread, so its row must stay + // `running` rather than becoming retryable while the original runs. + assert_eq!(schema.job_state(blocked.job.row.id).await, "running"); + wait_until(Duration::from_secs(10), "the later job", || async { + schema.job_state(later.job.row.id).await == "completed" + }) + .await; + run.shutdown().await.unwrap(); + + let blocking_finished = timeline.blocking_finished.lock().unwrap().unwrap(); + let later_started = timeline.later_started.lock().unwrap().unwrap(); + assert!( + later_started >= blocking_finished, + "the stuck job's worker slot was released while it still ran" + ); + let blocked = client.jobs().get(blocked.job.row.id).await.unwrap(); + assert_eq!(blocked.attempt, 1); + assert_eq!(blocked.errors.len(), 1); + assert_eq!( + blocked.errors[0].error, + "job aborted after ignoring cancellation" + ); + + schema.drop().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn shutdown_leaves_a_job_still_stuck_after_abort_running() { + let schema = TestSchema::new("abandon").await; + let gate = Gate::default(); + let timeline = Arc::new(BlockingTimeline::default()); + let client = Client::builder(schema.database()) + .id("postgres-resilience-abandon") + .job_stuck_threshold(Duration::from_millis(50)) + .without_notifications() + .workers(blocking_workers(&gate, &timeline)) + .queue("default", fast_queue()) + .build() + .unwrap(); + let blocked = client + .insert(BlockingArgs { block_ms: 1_500 }) + .await + .unwrap(); + + let mut run = client.start().unwrap(); + gate.wait_started().await; + let started = std::time::Instant::now(); + tokio::time::timeout(Duration::from_secs(1), run.shutdown_now()) + .await + .expect("shutdown waited for a task that cannot be aborted") + .unwrap(); + assert!(started.elapsed() < Duration::from_secs(1)); + assert!(timeline.blocking_finished.lock().unwrap().is_none()); + assert_eq!(schema.job_state(blocked.job.row.id).await, "running"); + + schema.drop().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn out_of_range_snooze_is_clamped_and_cancel_time_matches_go() { + let schema = TestSchema::new("wire").await; + let client = gated_client(&schema, "postgres-resilience-wire", &Gate::default()); + let snoozed = client.insert(SnoozeForeverArgs {}).await.unwrap(); + + let mut run = client.start().unwrap(); + wait_until(Duration::from_secs(10), "the snooze", || async { + schema.job_state(snoozed.job.row.id).await == "scheduled" + }) + .await; + run.shutdown().await.unwrap(); + let snoozed = client.jobs().get(snoozed.job.row.id).await.unwrap(); + assert_eq!(snoozed.attempt, 0); + assert!(snoozed.scheduled_at > chrono::Utc::now() + chrono::Duration::days(365 * 200)); + + // River Go writes `cancel_attempted_at` as `time.Time` JSON. + let cancelled = client.jobs().cancel(snoozed.id).await.unwrap(); + let cancel_attempted_at = cancelled + .metadata + .get::("cancel_attempted_at") + .unwrap() + .unwrap(); + assert!(cancel_attempted_at.ends_with('Z'), "{cancel_attempted_at}"); + if let Some((_, fraction)) = cancel_attempted_at.trim_end_matches('Z').split_once('.') { + assert!( + !fraction.ends_with('0'), + "trailing zeros are trimmed: {cancel_attempted_at}" + ); + } + chrono::DateTime::parse_from_rfc3339(&cancel_attempted_at).unwrap(); + + schema.drop().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn snooze_preserves_metadata_numbers_beyond_float_range() { + let schema = TestSchema::new("snooze_meta").await; + let client = gated_client(&schema, "postgres-resilience-snooze-meta", &Gate::default()); + let snoozed = client.insert(SnoozeForeverArgs {}).await.unwrap(); + sqlx::query(AssertSqlSafe(format!( + "UPDATE {} SET metadata = metadata || $1::jsonb WHERE id = $2", + schema.table() + ))) + .bind(r#"{"unrelated":1e400}"#) + .bind(snoozed.job.row.id) + .execute(&schema.pool) + .await + .unwrap(); + + let mut run = client.start().unwrap(); + wait_until(Duration::from_secs(10), "the snooze", || async { + schema.job_state(snoozed.job.row.id).await == "scheduled" + }) + .await; + run.shutdown().await.unwrap(); + let snoozed = client.jobs().get(snoozed.job.row.id).await.unwrap(); + // PostgreSQL expands `1e400` in jsonb; the snooze must keep it intact. + assert!(snoozed.metadata.get_raw("unrelated").unwrap().get().len() > 400); + assert_eq!(snoozed.metadata.get::("snoozes").unwrap(), Some(1)); + + schema.drop().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn client_survives_database_outage_and_catches_up() { + let schema = TestSchema::new("outage").await; + let proxy = FaultProxy::start().await; + let gate = Gate::default(); + let client = + Client::builder(PostgresDatabase::new(proxy.pool(4)).with_schema(schema.schema.clone())) + .id("postgres-resilience-outage") + .workers(gated_workers(&gate)) + // Notifications, not polling, must deliver work inserted during the + // outage once the listener reconnects. + .queue( + "default", + QueueConfig::new(4) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_mins(1)), + ) + .build() + .unwrap(); + // A second client inserts directly while the first is cut off. Its + // short fetch cooldown lets every insertion here notify the queue. + let inserter = Client::builder(schema.database()) + .fetch_cooldown(Duration::from_millis(1)) + .build() + .unwrap(); + + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + let before = inserter.insert(ResilienceArgs {}).await.unwrap(); + wait_until( + Duration::from_secs(10), + "work before the outage", + || async { schema.job_state(before.job.row.id).await == "completed" }, + ) + .await; + let in_flight = inserter.insert(GatedArgs {}).await.unwrap(); + gate.wait_started().await; + + proxy.take_down(); + // The in-flight job finishes while its completion cannot be written. + gate.release(); + let during = inserter.insert(ResilienceArgs {}).await.unwrap(); + proxy.wait_for_rejections(3).await; + proxy.restore(); + + for (id, description) in [ + (in_flight.job.row.id, "in-flight job completion"), + (during.job.row.id, "job inserted during the outage"), + ] { + wait_until(Duration::from_secs(30), description, || async { + schema.job_state(id).await == "completed" + }) + .await; + } + tokio::time::timeout(Duration::from_secs(10), run.shutdown()) + .await + .unwrap() + .unwrap(); + + let in_flight = inserter.jobs().get(in_flight.job.row.id).await.unwrap(); + assert_eq!(in_flight.attempt, 1, "the in-flight job was not rescued"); + assert_eq!(in_flight.errors, []); + + drop(proxy); + schema.drop().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn client_started_during_an_outage_becomes_ready_after_recovery() { + let schema = TestSchema::new("startdown").await; + let proxy = FaultProxy::start().await; + proxy.take_down(); + let client = + Client::builder(PostgresDatabase::new(proxy.pool(2)).with_schema(schema.schema.clone())) + .id("postgres-resilience-start-outage") + .workers(completing_workers()) + .queue("default", fast_queue()) + .build() + .unwrap(); + let job = Client::builder(schema.database()) + .build() + .unwrap() + .insert(ResilienceArgs {}) + .await + .unwrap(); + + let mut run = client.start().unwrap(); + // Outlast the producer's fast startup retries, after which it backs off + // instead of stopping the client. + tokio::time::sleep(Duration::from_secs(11)).await; + proxy.restore(); + tokio::time::timeout(Duration::from_secs(30), run.wait_ready()) + .await + .expect("listener did not recover") + .unwrap(); + wait_until(Duration::from_secs(30), "work after recovery", || async { + schema.job_state(job.job.row.id).await == "completed" + }) + .await; + run.shutdown().await.unwrap(); + + drop(proxy); + schema.drop().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn listener_does_not_occupy_a_pool_connection() { + let schema = TestSchema::new("listener").await; + let pool = PgPoolOptions::new() + .max_connections(1) + .connect(&database_url()) + .await + .unwrap(); + let client = + Client::builder(PostgresDatabase::new(pool.clone()).with_schema(schema.schema.clone())) + .id("postgres-resilience-listener-pool") + .workers(completing_workers()) + .queue("default", fast_queue()) + .build() + .unwrap(); + + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + let job = client.insert(ResilienceArgs {}).await.unwrap(); + wait_until( + Duration::from_secs(10), + "work through a one-connection pool", + || async { schema.job_state(job.job.row.id).await == "completed" }, + ) + .await; + run.shutdown().await.unwrap(); + pool.close().await; + + schema.drop().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn queue_update_keeps_the_producer_within_max_workers() { + let schema = TestSchema::new("reconfig").await; + let gate = Gate::default(); + let one_worker = |poll_interval_ms| { + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(poll_interval_ms)) + }; + let client = Client::builder(schema.database()) + .id("postgres-resilience-reconfigure") + .without_notifications() + .workers(gated_workers(&gate)) + .queue("default", one_worker(20)) + .build() + .unwrap(); + // Both jobs exist before the client starts. The producer has one worker + // slot, so it can't take the second job while the first runs, no matter + // when it applies the new configuration. + let first = client.insert(GatedArgs {}).await.unwrap(); + let second = client.insert(GatedArgs {}).await.unwrap(); + + let mut run = client.start().unwrap(); + gate.wait_started().await; + client + .local_queues() + .update("default", one_worker(10)) + .unwrap(); + // A second producer, or one that lost count of its running job, would + // take the second job while the first is still held. + gate.assert_none_started(Duration::from_millis(500)).await; + gate.release(); + gate.wait_started().await; + gate.release(); + for id in [first.id(), second.id()] { + wait_until(Duration::from_secs(10), "both jobs", || async { + schema.job_state(id).await == "completed" + }) + .await; + } + run.shutdown().await.unwrap(); + assert_eq!( + gate.max_active.load(Ordering::SeqCst), + 1, + "reconfiguration exceeded max_workers" + ); + + schema.drop().await; +} + +/// An extension that observes completions inside River's transaction, +/// failing its first call and deleting jobs marked for deletion. +#[derive(Clone, Default)] +struct SetStatePilot { + calls: Arc, + seen: Arc>>, +} + +#[async_trait] +impl Pilot for SetStatePilot { + fn intercepts_job_set_state(&self) -> bool { + true + } + + async fn after_jobs_set_state( + &self, + connection: DatabaseConnection<'_>, + params: &JobSetStateParams, + ) -> Result<(), PilotError> { + let connection = connection.into_postgres().unwrap(); + let schema = params.database.postgres_schema().unwrap(); + let ids = params.jobs.iter().map(|job| job.id).collect::>(); + sqlx::query(AssertSqlSafe(format!( + "INSERT INTO {} (job_id) SELECT unnest($1::bigint[])", + schema.qualify("hook_effect") + ))) + .bind(&ids) + .execute(&mut *connection) + .await?; + if self.calls.fetch_add(1, Ordering::SeqCst) == 0 { + return Err(std::io::Error::other("first set-state hook call fails").into()); + } + self.seen.lock().unwrap().extend( + params + .jobs + .iter() + .map(|job| (job.id, job.state.as_str().to_owned())), + ); + let deleted = params + .jobs + .iter() + .filter(|job| job.metadata.contains_key("delete_me")) + .map(|job| job.id) + .collect::>(); + sqlx::query(AssertSqlSafe(format!( + "DELETE FROM {} WHERE id = ANY($1)", + schema.qualify("river_job") + ))) + .bind(deleted) + .execute(&mut *connection) + .await?; + Ok(()) + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn extension_set_state_hook_runs_in_the_completion_transaction() { + let schema = TestSchema::new("hook").await; + schema + .execute(format!( + "CREATE TABLE {}.hook_effect (job_id bigint NOT NULL)", + schema.name + )) + .await; + let pilot = SetStatePilot::default(); + let gate = Gate::default(); + let client = Client::builder(schema.database()) + .id("postgres-resilience-hook") + .pilot(pilot.clone()) + .without_notifications() + .workers(gated_workers(&gate)) + .queue("default", fast_queue()) + .build() + .unwrap(); + let mut events = client.subscribe(&[EventKind::JobCompleted]).unwrap(); + let kept = client.insert(GatedArgs {}).await.unwrap(); + let deleted = client + .insert(GatedArgs {}) + .opts( + InsertOpts::default().with_metadata(serde_json::Map::from_iter([( + "delete_me".to_owned(), + serde_json::json!(true), + )])), + ) + .await + .unwrap(); + + let mut run = client.start().unwrap(); + gate.wait_started().await; + gate.wait_started().await; + // Release both together so they share one batch. + gate.release(); + gate.release(); + let mut completed = Vec::new(); + for _ in 0..2 { + let event = tokio::time::timeout(Duration::from_secs(10), events.recv()) + .await + .expect("completion events") + .unwrap(); + completed.push(event.as_job().unwrap().job.id); + } + run.shutdown().await.unwrap(); + + completed.sort_unstable(); + let mut expected = vec![kept.job.row.id, deleted.job.row.id]; + expected.sort_unstable(); + assert_eq!( + completed, expected, + "events come from rows River already holds" + ); + assert_eq!(schema.job_state(kept.job.row.id).await, "completed"); + assert!(matches!( + client.jobs().get(deleted.job.row.id).await, + Err(riverqueue::Error::NotFound(_)) + )); + let mut seen = pilot.seen.lock().unwrap().clone(); + seen.sort_unstable(); + assert_eq!( + seen, + expected + .iter() + .map(|id| (*id, "completed".to_owned())) + .collect::>() + ); + // The failed first call's writes rolled back with its batch. + let effects: i64 = sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT count(*) FROM {}.hook_effect", + schema.name + ))) + .fetch_one(&schema.pool) + .await + .unwrap(); + assert_eq!(effects, 2); + + schema.drop().await; +} diff --git a/rust/riverqueue/tests/resilience_sqlite.rs b/rust/riverqueue/tests/resilience_sqlite.rs new file mode 100644 index 000000000..32d04a4a9 --- /dev/null +++ b/rust/riverqueue/tests/resilience_sqlite.rs @@ -0,0 +1,882 @@ +//! SQLite runtime behavior under malformed rows and database faults. + +#![cfg(feature = "sqlite")] + +use std::{ + convert::Infallible, + future::Future, + path::{Path, PathBuf}, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use chrono::{DateTime, Utc}; +use riverqueue::{ + AttemptError, BoxError, Client, ErrorHandler, ErrorHandlerDecision, EventKind, InsertOpts, Job, + JobArgs, JobEventKind, JobRow, JobState, MaintenanceConfig, QueueConfig, RetryPolicy, + UniqueOpts, WorkCancelled, WorkContext, WorkOutcome, WorkResult, WorkerRegistry, +}; +use riverqueue_migrate::SqliteMigrator; +use serde::{Deserialize, Serialize}; +use sqlx::{ + SqlitePool, + sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}, +}; +use tokio::sync::Semaphore; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_sqlite_resilience")] +struct ResilienceArgs {} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_sqlite_resilience_gated")] +struct GatedArgs { + fail: bool, +} + +#[derive(Debug, thiserror::Error)] +#[error("gated job failed")] +struct GatedError; + +/// A job that waits for client shutdown, then stops cooperatively or fails. +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_sqlite_resilience_shutdown")] +struct ShutdownArgs { + cooperative: bool, +} + +#[derive(Debug, thiserror::Error)] +enum ShutdownError { + #[error("stopped for shutdown")] + Cancelled(#[source] WorkCancelled), + #[error("real failure during shutdown")] + Real, +} + +/// Lets a test hold a gated job inside its worker until released. +#[derive(Clone)] +struct Gate { + release: Arc, + started: Arc, +} + +impl Default for Gate { + fn default() -> Self { + Self { + release: Arc::new(Semaphore::new(0)), + started: Arc::new(Semaphore::new(0)), + } + } +} + +impl Gate { + async fn wait_started(&self) { + tokio::time::timeout(Duration::from_secs(10), self.started.acquire()) + .await + .expect("gated job did not start") + .unwrap() + .forget(); + } + + fn release(&self) { + self.release.add_permits(1); + } +} + +/// A migrated WAL database file that is removed when the test finishes. +struct TestDatabase { + path: PathBuf, + pool: SqlitePool, +} + +impl TestDatabase { + async fn new(busy_timeout: Duration) -> Self { + static DATABASE_NONCE: AtomicUsize = AtomicUsize::new(0); + let path = std::env::temp_dir().join(format!( + "river-sqlite-resilience-{}-{}.sqlite", + std::process::id(), + DATABASE_NONCE.fetch_add(1, Ordering::Relaxed) + )); + let pool = Self::connect(&path, busy_timeout, 4).await; + SqliteMigrator::new(pool.clone()) + .migrate_up() + .await + .unwrap(); + Self { path, pool } + } + + async fn connect(path: &Path, busy_timeout: Duration, connections: u32) -> SqlitePool { + SqlitePoolOptions::new() + .max_connections(connections) + .connect_with( + SqliteConnectOptions::new() + .filename(path) + .create_if_missing(true) + .journal_mode(SqliteJournalMode::Wal) + .busy_timeout(busy_timeout), + ) + .await + .unwrap() + } +} + +impl Drop for TestDatabase { + fn drop(&mut self) { + let _ = std::fs::remove_file(&self.path); + for suffix in ["-shm", "-wal"] { + let mut path = self.path.as_os_str().to_owned(); + path.push(suffix); + let _ = std::fs::remove_file(path); + } + } +} + +fn completing_workers() -> WorkerRegistry { + gated_workers(&Gate::default()) +} + +fn gated_workers(gate: &Gate) -> WorkerRegistry { + let mut workers = WorkerRegistry::new(); + workers + .register_fn(|_context: WorkContext, _job: Job| async { + Ok::<_, Infallible>(WorkOutcome::Complete) + }) + .unwrap(); + let shutdown_gate = gate.clone(); + workers + .register_fn(move |context: WorkContext, job: Job| { + let gate = shutdown_gate.clone(); + async move { + gate.started.add_permits(1); + context.cancellation_token().cancelled().await; + if job.args.cooperative { + Err::(ShutdownError::Cancelled(WorkCancelled)) + } else { + Err(ShutdownError::Real) + } + } + }) + .unwrap(); + let gate = gate.clone(); + workers + .register_fn(move |_context: WorkContext, job: Job| { + let gate = gate.clone(); + async move { + gate.started.add_permits(1); + gate.release.acquire().await.unwrap().forget(); + if job.args.fail { + Err(GatedError) + } else { + Ok(WorkOutcome::Complete) + } + } + }) + .unwrap(); + workers +} + +fn fast_queue() -> QueueConfig { + QueueConfig::new(4) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(20)) +} + +/// Polls `condition` until it holds, failing the test after `timeout`. +async fn wait_until(timeout: Duration, description: &str, mut condition: F) +where + F: FnMut() -> Fut, + Fut: Future, +{ + let deadline = tokio::time::Instant::now() + timeout; + while !condition().await { + assert!( + tokio::time::Instant::now() < deadline, + "timed out waiting for {description}" + ); + tokio::time::sleep(Duration::from_millis(10)).await; + } +} + +async fn job_state(pool: &SqlitePool, id: i64) -> String { + sqlx::query_scalar("SELECT state FROM river_job WHERE id = ?") + .bind(id) + .fetch_one(pool) + .await + .unwrap() +} + +/// Records every failed attempt the error handler sees. +#[derive(Clone, Default)] +struct RecordingErrorHandler(Arc>>); + +#[allow( + clippy::unused_async_trait_impl, + reason = "these extensions only record state synchronously" +)] +impl ErrorHandler for RecordingErrorHandler { + async fn handle_error( + &self, + _context: &WorkContext, + job: &JobRow, + result: &WorkResult, + ) -> Result { + let error = match result { + WorkResult::Failed(error) => error.to_string(), + other => format!("{other:?}"), + }; + self.0.lock().unwrap().push((job.clone(), error)); + Ok(ErrorHandlerDecision::Continue) + } +} + +/// Schedules every retry an hour out so a failed job stays `retryable`. +struct RetryAnHourLater; + +impl RetryPolicy for RetryAnHourLater { + fn next_retry( + &self, + _job: &JobRow, + _error: &riverqueue::WorkError, + _now: DateTime, + ) -> Duration { + Duration::from_hours(1) + } +} + +async fn set_json_column(pool: &SqlitePool, id: i64, column: &str, json: &str) { + sqlx::query(sqlx::AssertSqlSafe(format!( + "UPDATE river_job SET {column} = jsonb(?) WHERE id = ?" + ))) + .bind(json) + .bind(id) + .execute(pool) + .await + .unwrap(); +} + +#[tokio::test(flavor = "multi_thread")] +#[allow(clippy::too_many_lines)] +async fn claimed_rows_decode_individually_and_accept_go_integer_ranges() { + let database = TestDatabase::new(Duration::from_secs(5)).await; + let error_handler = RecordingErrorHandler::default(); + let client = Client::builder(database.pool.clone()) + .id("sqlite-resilience-decode") + .maintenance( + MaintenanceConfig::default().with_scheduler_interval(Duration::from_millis(50)), + ) + .error_handler(error_handler.clone()) + .retry_policy(RetryAnHourLater) + .workers(completing_workers()) + .queue("default", fast_queue()) + .build() + .unwrap(); + let mut events = client + .subscribe(&[EventKind::JobCompleted, EventKind::JobFailed]) + .unwrap(); + + // River Go stores native integers on SQLite, so `max_attempts` can exceed + // Rust's `i16`. Such a job must still be worked. + let wide = client.insert(ResilienceArgs {}).await.unwrap(); + sqlx::query("UPDATE river_job SET max_attempts = 40000 WHERE id = ?") + .bind(wide.job.row.id) + .execute(&database.pool) + .await + .unwrap(); + // Attempt errors in a shape River doesn't write decode leniently like + // River Go's, so the job is still worked. + let odd_errors = client.insert(ResilienceArgs {}).await.unwrap(); + set_json_column( + &database.pool, + odd_errors.job.row.id, + "errors", + r#"[{"attempt": "1", "error": {"message": "boom"}}, 42]"#, + ) + .await; + // A row whose tags aren't an array can't become a `JobRow`. Claimed with + // the others, such a job isn't worked, and its attempt fails like any + // other: retried with the client's retry policy, or discarded at its + // maximum attempts. An `errors` value that isn't an array is wrapped in + // one so the attempt error can still be appended. + let malformed_retried = client.insert(ResilienceArgs {}).await.unwrap(); + set_json_column( + &database.pool, + malformed_retried.job.row.id, + "errors", + r#"{"not":"an array"}"#, + ) + .await; + let malformed_discarded = client + .insert(ResilienceArgs {}) + .opts(InsertOpts::default().with_max_attempts(1)) + .await + .unwrap(); + let malformed_ids = [malformed_retried.job.row.id, malformed_discarded.job.row.id]; + for id in malformed_ids { + set_json_column(&database.pool, id, "tags", r#"{"not":"an array"}"#).await; + } + let ordinary = client.insert(ResilienceArgs {}).await.unwrap(); + let decodable_ids = [wide.job.row.id, odd_errors.job.row.id, ordinary.job.row.id]; + + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + let mut events_by_id = std::collections::HashMap::new(); + while events_by_id.len() < decodable_ids.len() + malformed_ids.len() { + let event = tokio::time::timeout(Duration::from_secs(10), events.recv()) + .await + .expect("job events") + .unwrap(); + let event = event.as_job().unwrap().clone(); + events_by_id.insert(event.job.id, event); + } + run.shutdown().await.unwrap(); + + for id in decodable_ids { + assert_eq!(events_by_id[&id].kind, JobEventKind::Completed); + } + let (max_attempts, errors): (i64, Option) = + sqlx::query_as("SELECT max_attempts, json(errors) FROM river_job WHERE id = ?") + .bind(wide.job.row.id) + .fetch_one(&database.pool) + .await + .unwrap(); + assert_eq!( + max_attempts, 40_000, + "decoding must not rewrite the stored value" + ); + assert!(errors.is_none()); + let wide = client.jobs().get(wide.job.row.id).await.unwrap(); + assert_eq!(wide.max_attempts, i16::MAX); + + let odd_errors = client.jobs().get(odd_errors.job.row.id).await.unwrap(); + let zero_time = "0001-01-01T00:00:00Z".parse().unwrap(); + assert_eq!( + odd_errors.errors, + [ + AttemptError::new(zero_time, 1, r#"{"message":"boom"}"#), + AttemptError::new(zero_time, 0, "42"), + ] + ); + + // Failed events and the error handler carry the fields that could be + // decoded, with the others left empty. + let handled = error_handler.0.lock().unwrap().clone(); + assert_eq!(handled.len(), malformed_ids.len()); + for (id, state) in [ + (malformed_retried.job.row.id, JobState::Retryable), + (malformed_discarded.job.row.id, JobState::Discarded), + ] { + let event = &events_by_id[&id]; + assert_eq!(event.kind, JobEventKind::Failed); + assert_eq!(event.job.state, state); + assert_eq!(event.job.kind, ResilienceArgs::KIND); + assert_eq!(event.job.tags, Vec::::new()); + + let (job, error) = handled.iter().find(|(job, _)| job.id == id).unwrap(); + assert_eq!(job.attempt, 1); + assert_eq!(job.tags, Vec::::new()); + assert!( + error.starts_with("job row couldn't be decoded: "), + "{error}" + ); + assert!(error.contains("error unmarshaling `tags`: "), "{error}"); + + // The attempt error is appended without rewriting the undecodable + // tags, and the undecodable row still can't be read. + assert!(client.jobs().get(id).await.is_err()); + let (stored_state, attempt, errors, tags, scheduled_at): ( + String, + i64, + String, + String, + DateTime, + ) = sqlx::query_as( + "SELECT state, attempt, json(errors), json(tags), scheduled_at FROM river_job WHERE id = ?", + ) + .bind(id) + .fetch_one(&database.pool) + .await + .unwrap(); + assert_eq!(stored_state, state.as_str()); + assert_eq!(attempt, 1); + assert_eq!(tags, r#"{"not":"an array"}"#); + let errors: Vec = serde_json::from_str(&errors).unwrap(); + let appended = errors.last().unwrap(); + assert_eq!(appended["attempt"], 1); + assert_eq!(appended["error"], error.as_str()); + if state == JobState::Retryable { + assert_eq!(errors[0], serde_json::json!({"not": "an array"})); + assert!( + scheduled_at > Utc::now() + chrono::Duration::minutes(50), + "the client retry policy wasn't used: {scheduled_at}" + ); + } else { + assert_eq!(errors.len(), 1); + } + } +} + +// Like River Go's `JobGetStuck`, the rescuer reads a stuck job whose row +// can't be fully decoded, so it can recover it along with the others. +#[tokio::test(flavor = "multi_thread")] +async fn rescuer_recovers_undecodable_stuck_jobs() { + let database = TestDatabase::new(Duration::from_secs(5)).await; + let client = Client::builder(database.pool.clone()) + .id("sqlite-resilience-rescue") + .job_timeout(Duration::from_millis(100)) + .maintenance( + MaintenanceConfig::default() + .with_elect_interval(Duration::from_millis(20)) + .with_rescue_after(Duration::from_millis(100)) + .with_rescuer_interval(Duration::from_millis(20)), + ) + .retry_policy(RetryAnHourLater) + .workers(completing_workers()) + .queue("default", fast_queue()) + .build() + .unwrap(); + let mut stuck = Vec::new(); + for _ in 0..2 { + let job = client.insert(ResilienceArgs {}).await.unwrap(); + sqlx::query( + "UPDATE river_job SET state = 'running', attempt = 1, \ + attempted_at = datetime('now', '-1 hour') WHERE id = ?", + ) + .bind(job.job.row.id) + .execute(&database.pool) + .await + .unwrap(); + stuck.push(job.job.row.id); + } + set_json_column(&database.pool, stuck[0], "tags", r#"{"not":"an array"}"#).await; + // Like River Go, the rescue error is appended to `errors` that aren't + // valid JSON by wrapping them in an array, and metadata that isn't valid + // JSON is left in place. + let invalid_json = client.insert(ResilienceArgs {}).await.unwrap().job.row.id; + sqlx::query( + "UPDATE river_job SET state = 'running', attempt = 1, \ + attempted_at = datetime('now', '-1 hour'), errors = 'not json', \ + metadata = 'not json' WHERE id = ?", + ) + .bind(invalid_json) + .execute(&database.pool) + .await + .unwrap(); + stuck.push(invalid_json); + + let mut run = client.start().unwrap(); + for id in &stuck { + wait_until(Duration::from_secs(10), "stuck job rescue", || async { + job_state(&database.pool, *id).await == "retryable" + }) + .await; + } + run.shutdown().await.unwrap(); + + let tags: String = sqlx::query_scalar("SELECT json(tags) FROM river_job WHERE id = ?") + .bind(stuck[0]) + .fetch_one(&database.pool) + .await + .unwrap(); + assert_eq!(tags, r#"{"not":"an array"}"#); + + let (errors, metadata): (String, String) = + sqlx::query_as("SELECT json(errors), metadata FROM river_job WHERE id = ?") + .bind(invalid_json) + .fetch_one(&database.pool) + .await + .unwrap(); + let errors: Vec = serde_json::from_str(&errors).unwrap(); + assert_eq!(errors.len(), 2); + assert_eq!(errors[0], "not json"); + assert_eq!(metadata, "not json"); +} + +/// Sets a JSON column to text that isn't valid JSON, as an out of band +/// change could. +async fn set_invalid_json_column(pool: &SqlitePool, id: i64, column: &str) { + sqlx::query(sqlx::AssertSqlSafe(format!( + "UPDATE river_job SET {column} = 'not json' WHERE id = ?" + ))) + .bind(id) + .execute(pool) + .await + .unwrap(); +} + +// Like River Go, a JSON column holding text that isn't valid JSON doesn't +// fail fetches from its queue, completion, or scheduling. A job with one is +// claimed as undecodable and its attempt fails, the value is left in place, +// and the other jobs are worked. +#[tokio::test(flavor = "multi_thread")] +#[allow(clippy::too_many_lines)] +async fn invalid_json_columns_fail_attempts_without_stalling_the_queue() { + const COLUMNS: [&str; 5] = ["args", "attempted_by", "errors", "metadata", "tags"]; + + let database = TestDatabase::new(Duration::from_secs(5)).await; + let error_handler = RecordingErrorHandler::default(); + let client = Client::builder(database.pool.clone()) + .id("sqlite-resilience-invalid-json") + .maintenance( + MaintenanceConfig::default() + .with_elect_interval(Duration::from_millis(20)) + .with_scheduler_interval(Duration::from_millis(50)), + ) + .error_handler(error_handler.clone()) + .retry_policy(RetryAnHourLater) + .workers(completing_workers()) + .queue("default", fast_queue()) + .build() + .unwrap(); + let mut events = client + .subscribe(&[EventKind::JobCompleted, EventKind::JobFailed]) + .unwrap(); + + let mut invalid = Vec::new(); + for column in COLUMNS { + let job = client.insert(ResilienceArgs {}).await.unwrap(); + set_invalid_json_column(&database.pool, job.job.row.id, column).await; + invalid.push((column, job.job.row.id)); + } + let ordinary = client.insert(ResilienceArgs {}).await.unwrap(); + + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + let mut events_by_id = std::collections::HashMap::new(); + while events_by_id.len() < invalid.len() + 1 { + let event = tokio::time::timeout(Duration::from_secs(10), events.recv()) + .await + .expect("job events") + .unwrap(); + let event = event.as_job().unwrap().clone(); + events_by_id.insert(event.job.id, event); + } + assert_eq!( + events_by_id[&ordinary.job.row.id].kind, + JobEventKind::Completed + ); + + let handled = error_handler.0.lock().unwrap().clone(); + for &(column, id) in &invalid { + assert_eq!(events_by_id[&id].kind, JobEventKind::Failed, "{column}"); + let (_, error) = handled.iter().find(|(job, _)| job.id == id).unwrap(); + assert!( + error.starts_with("job row couldn't be decoded: "), + "{column}: {error}" + ); + assert!( + error.contains(&format!("error unmarshaling `{column}`: ")), + "{column}: {error}" + ); + + // The invalid value is left in place, except that the attempt error + // is appended to an invalid `errors` value wrapped in an array. + let (state, attempt, stored_type, stored): (String, i64, String, String) = + sqlx::query_as(sqlx::AssertSqlSafe(format!( + "SELECT state, attempt, typeof({column}), \ + CASE WHEN typeof({column}) = 'text' THEN {column} ELSE json({column}) END \ + FROM river_job WHERE id = ?" + ))) + .bind(id) + .fetch_one(&database.pool) + .await + .unwrap(); + assert_eq!(state, "retryable", "{column}"); + assert_eq!(attempt, 1, "{column}"); + if column == "errors" { + let errors: Vec = serde_json::from_str(&stored).unwrap(); + assert_eq!(errors.len(), 2); + assert_eq!(errors[0], "not json"); + assert_eq!(errors[1]["error"], error.as_str()); + } else { + assert_eq!( + (stored_type.as_str(), stored.as_str()), + ("text", "not json"), + "{column}" + ); + } + } + + // The scheduler makes the failed jobs available again without failing on + // their invalid values, along with an ordinary scheduled job. + let scheduled = client + .insert(ResilienceArgs {}) + .opts(InsertOpts::default().with_scheduled_at(Utc::now() + chrono::Duration::hours(1))) + .await + .unwrap(); + sqlx::query("UPDATE river_job SET scheduled_at = datetime('now', '-1 second') WHERE state IN ('retryable', 'scheduled')") + .execute(&database.pool) + .await + .unwrap(); + wait_until( + Duration::from_secs(10), + "scheduled job completion", + || async { job_state(&database.pool, scheduled.job.row.id).await == "completed" }, + ) + .await; + // Wrapping the invalid `errors` in an array made that job decodable, so + // it's worked this time. The others fail again. + for &(column, id) in &invalid { + let expected = if column == "errors" { + "completed" + } else { + "retryable" + }; + wait_until(Duration::from_secs(10), column, || async { + let (state, attempt): (String, i64) = + sqlx::query_as("SELECT state, attempt FROM river_job WHERE id = ?") + .bind(id) + .fetch_one(&database.pool) + .await + .unwrap(); + state == expected && attempt == 2 + }) + .await; + } + run.shutdown().await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn completion_retries_while_a_foreign_writer_holds_the_lock() { + // The client gives up on a busy database after 50 ms, far less than the + // foreign transaction below holds the write lock. + let database = TestDatabase::new(Duration::from_millis(50)).await; + let gate = Gate::default(); + let client = Client::builder(database.pool.clone()) + .id("sqlite-resilience-writer-lock") + .without_notifications() + .workers(gated_workers(&gate)) + .queue("default", fast_queue()) + .build() + .unwrap(); + let job = client.insert(GatedArgs { fail: false }).await.unwrap(); + + let mut run = client.start().unwrap(); + gate.wait_started().await; + let foreign = TestDatabase::connect(&database.path, Duration::from_secs(5), 1).await; + let writer = foreign.begin_with("BEGIN IMMEDIATE").await.unwrap(); + gate.release(); + // Hold the lock across several busy timeouts and the first retry. + tokio::time::sleep(Duration::from_millis(1_500)).await; + writer.rollback().await.unwrap(); + wait_until( + Duration::from_secs(15), + "completion after unlock", + || async { job_state(&database.pool, job.job.row.id).await == "completed" }, + ) + .await; + run.shutdown().await.unwrap(); + foreign.close().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn completion_cancels_on_a_null_cancel_attempted_at_key() { + let database = TestDatabase::new(Duration::from_secs(5)).await; + let gate = Gate::default(); + let client = Client::builder(database.pool.clone()) + .id("sqlite-resilience-cancel-key") + .without_notifications() + .workers(gated_workers(&gate)) + .queue("default", fast_queue()) + .build() + .unwrap(); + let job = client.insert(GatedArgs { fail: true }).await.unwrap(); + + let mut run = client.start().unwrap(); + gate.wait_started().await; + // River Go treats the key's presence, not its value, as a cancellation. + sqlx::query( + "UPDATE river_job SET metadata = jsonb_set(metadata, '$.cancel_attempted_at', json('null')) \ + WHERE id = ?", + ) + .bind(job.job.row.id) + .execute(&database.pool) + .await + .unwrap(); + gate.release(); + wait_until(Duration::from_secs(10), "cancellation", || async { + job_state(&database.pool, job.job.row.id).await == "cancelled" + }) + .await; + run.shutdown().await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn hard_shutdown_interrupts_only_cooperative_cancellations() { + let database = TestDatabase::new(Duration::from_secs(5)).await; + let gate = Gate::default(); + let client = Client::builder(database.pool.clone()) + .id("sqlite-resilience-shutdown") + .without_notifications() + .workers(gated_workers(&gate)) + .queue("default", fast_queue()) + .build() + .unwrap(); + let cooperative = client + .insert(ShutdownArgs { cooperative: true }) + .await + .unwrap(); + let cancel_attempted = client + .insert(ShutdownArgs { cooperative: true }) + .await + .unwrap(); + let real_error = client + .insert(ShutdownArgs { cooperative: false }) + .await + .unwrap(); + + let mut run = client.start().unwrap(); + for _ in 0..3 { + gate.wait_started().await; + } + let notifications_before: i64 = + sqlx::query_scalar("SELECT count(*) FROM river_notification WHERE topic = 'river_insert'") + .fetch_one(&database.pool) + .await + .unwrap(); + sqlx::query( + "UPDATE river_job SET metadata = jsonb_set(metadata, '$.cancel_attempted_at', \ + '2026-01-02T03:04:05Z') WHERE id = ?", + ) + .bind(cancel_attempted.job.row.id) + .execute(&database.pool) + .await + .unwrap(); + tokio::time::timeout(Duration::from_secs(10), run.shutdown_now()) + .await + .unwrap() + .unwrap(); + + let cooperative = client.jobs().get(cooperative.job.row.id).await.unwrap(); + assert_eq!(cooperative.state, JobState::Available); + assert_eq!(cooperative.attempt, 0); + assert!(cooperative.attempted_at.is_some(), "attempted_at is kept"); + assert_eq!(cooperative.errors, []); + let notifications_after: i64 = + sqlx::query_scalar("SELECT count(*) FROM river_notification WHERE topic = 'river_insert'") + .fetch_one(&database.pool) + .await + .unwrap(); + assert!(notifications_after > notifications_before); + + let cancel_attempted = client + .jobs() + .get(cancel_attempted.job.row.id) + .await + .unwrap(); + assert_eq!(cancel_attempted.state, JobState::Cancelled); + assert!(cancel_attempted.finalized_at.is_some()); + + let real_error = client.jobs().get(real_error.job.row.id).await.unwrap(); + assert!(matches!( + real_error.state, + JobState::Available | JobState::Retryable + )); + assert_eq!(real_error.attempt, 1); + assert_eq!(real_error.errors.len(), 1); + assert_eq!(real_error.errors[0].error, "real failure during shutdown"); +} + +#[tokio::test(flavor = "multi_thread")] +async fn cancel_attempted_at_matches_go_time_json() { + let database = TestDatabase::new(Duration::from_secs(5)).await; + let client = Client::builder(database.pool.clone()) + .id("sqlite-resilience-cancel-time") + .build() + .unwrap(); + let job = client.insert(ResilienceArgs {}).await.unwrap(); + let cancelled = client.jobs().cancel(job.job.row.id).await.unwrap(); + let cancel_attempted_at = cancelled + .metadata + .get::("cancel_attempted_at") + .unwrap() + .unwrap(); + assert!(cancel_attempted_at.ends_with('Z'), "{cancel_attempted_at}"); + if let Some((_, fraction)) = cancel_attempted_at.trim_end_matches('Z').split_once('.') { + assert!( + !fraction.ends_with('0'), + "trailing zeros are trimmed: {cancel_attempted_at}" + ); + } + chrono::DateTime::parse_from_rfc3339(&cancel_attempted_at).unwrap(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn notification_poll_failures_do_not_stop_the_client() { + // In rollback-journal mode an exclusive writer blocks readers, so the + // outbox poll fails with `database is locked` while the lock is held. + static DATABASE_NONCE: AtomicUsize = AtomicUsize::new(0); + let path = std::env::temp_dir().join(format!( + "river-sqlite-resilience-journal-{}-{}.sqlite", + std::process::id(), + DATABASE_NONCE.fetch_add(1, Ordering::Relaxed) + )); + let connect = |busy_timeout| { + SqlitePoolOptions::new().max_connections(2).connect_with( + SqliteConnectOptions::new() + .filename(&path) + .create_if_missing(true) + .journal_mode(SqliteJournalMode::Delete) + .busy_timeout(busy_timeout), + ) + }; + let pool = connect(Duration::from_millis(20)).await.unwrap(); + SqliteMigrator::new(pool.clone()) + .migrate_up() + .await + .unwrap(); + let database = TestDatabase { + path: path.clone(), + pool, + }; + let client = Client::builder(database.pool.clone()) + .id("sqlite-resilience-outbox") + .workers(completing_workers()) + .queue("default", fast_queue()) + .build() + .unwrap(); + + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + let foreign = connect(Duration::from_secs(5)).await.unwrap(); + let writer = foreign.begin_with("BEGIN EXCLUSIVE").await.unwrap(); + // Hold the lock across several 100 ms outbox polls. + tokio::time::sleep(Duration::from_millis(500)).await; + writer.rollback().await.unwrap(); + foreign.close().await; + + let job = client.insert(ResilienceArgs {}).await.unwrap(); + wait_until(Duration::from_secs(15), "work after the lock", || async { + job_state(&database.pool, job.job.row.id).await == "completed" + }) + .await; + run.shutdown().await.unwrap(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn unique_duplicates_are_detected_across_clients_with_the_same_id() { + let database = TestDatabase::new(Duration::from_secs(5)).await; + // Default client IDs repeat across restarted containers, so duplicate + // detection must not depend on them. + let first = Client::builder(database.pool.clone()) + .id("sqlite-resilience-same-id") + .build() + .unwrap(); + let second = Client::builder(database.pool.clone()) + .id("sqlite-resilience-same-id") + .build() + .unwrap(); + let opts = InsertOpts::default().with_unique(UniqueOpts::new().with_by_args(true)); + + let inserted = first + .insert(ResilienceArgs {}) + .opts(opts.clone()) + .await + .unwrap(); + assert!(!inserted.unique_skipped_as_duplicate); + let duplicate = second.insert(ResilienceArgs {}).opts(opts).await.unwrap(); + assert!(duplicate.unique_skipped_as_duplicate); + assert_eq!(duplicate.job.row.id, inserted.job.row.id); +} diff --git a/rust/riverqueue/tests/runtime_config.rs b/rust/riverqueue/tests/runtime_config.rs new file mode 100644 index 000000000..365b96f32 --- /dev/null +++ b/rust/riverqueue/tests/runtime_config.rs @@ -0,0 +1,786 @@ +#![cfg(all(feature = "postgres", river_postgres_tests))] + +mod support; + +use std::{ + collections::HashSet, + convert::Infallible, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use riverqueue::{ + BoxError, Client, EventKind, EventReceiver, EventRecvError, Extensions, Hook, InsertContext, + InsertMiddleware, InsertNext, InsertedJob, Job, JobArgs, JobRow, JobState, Metric, + PeriodicJobs, Plugin, QueueConfig, SubscribeConfig, WorkContext, WorkError, WorkMiddleware, + WorkNext, WorkOutcome, Worker, WorkerRegistry, database::PostgresDatabase, +}; +use serde::{Deserialize, Serialize}; +use sqlx::AssertSqlSafe; +use tokio::sync::Semaphore; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_runtime_config")] +struct RuntimeArgs {} + +struct RuntimeWorker; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_runtime_burst")] +struct BurstArgs {} + +struct BurstWorker; + +impl Worker for BurstWorker { + type Error = Infallible; + + fn work( + &self, + _context: WorkContext, + _job: Job, + ) -> impl std::future::Future> + Send { + std::future::ready(Ok(WorkOutcome::Complete)) + } +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_runtime_cancel_snooze")] +struct CancelSnoozeArgs {} + +struct CancelSnoozeWorker { + started: Arc, +} + +impl Worker for CancelSnoozeWorker { + type Error = Infallible; + + async fn work( + &self, + context: WorkContext, + _job: Job, + ) -> Result { + self.started.add_permits(1); + context.cancellation_token().cancelled().await; + Ok(WorkOutcome::Snooze(Duration::from_hours(1))) + } +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_runtime_terminal_race")] +struct TerminalRaceArgs {} + +struct TerminalRaceWorker { + finish: Arc, + started: Arc, +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_runtime_shutdown")] +struct ShutdownArgs { + ignore_cancellation: bool, +} + +struct ShutdownWorker { + finish: Arc, + started: Arc, +} + +impl Worker for ShutdownWorker { + type Error = Infallible; + + async fn work( + &self, + _context: WorkContext, + job: Job, + ) -> Result { + self.started.add_permits(1); + if job.args.ignore_cancellation { + std::future::pending::<()>().await; + } else { + self.finish.acquire().await.unwrap().forget(); + } + Ok(WorkOutcome::Complete) + } +} + +impl Worker for TerminalRaceWorker { + type Error = Infallible; + + async fn work( + &self, + context: WorkContext, + _job: Job, + ) -> Result { + self.started.add_permits(1); + self.finish.acquire().await.unwrap().forget(); + context.metadata_set("worker_completion", true).unwrap(); + Ok(WorkOutcome::Complete) + } +} + +#[derive(Clone)] +struct RuntimeHook { + counts: Arc, +} + +#[derive(Default)] +struct RuntimeCounts { + insert_after: AtomicUsize, + insert_before: AtomicUsize, + metrics: AtomicUsize, + periodic_starts: AtomicUsize, + work_after: AtomicUsize, + work_before: AtomicUsize, +} + +#[allow( + clippy::unused_async_trait_impl, + reason = "these extensions only record state synchronously" +)] +impl Hook for RuntimeHook { + async fn insert_begin(&self, _insert: &mut InsertContext) -> Result<(), BoxError> { + self.counts.insert_before.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + + async fn metric_emit(&self, _metric: Metric) -> Result<(), BoxError> { + self.counts.metrics.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + + async fn periodic_jobs_start(&self, _jobs: &PeriodicJobs) -> Result<(), BoxError> { + self.counts.periodic_starts.fetch_add(1, Ordering::SeqCst); + Ok(()) + } + + async fn work_begin(&self, _context: &WorkContext, job: &mut JobRow) -> Result<(), BoxError> { + self.counts.work_before.fetch_add(1, Ordering::SeqCst); + let mut args: serde_json::Value = job.decode_args()?; + args["hook_decrypted"] = true.into(); + job.encoded_args = riverqueue::encoding::encode_args(&args)?; + Ok(()) + } + + async fn work_end( + &self, + _context: &WorkContext, + _job: &JobRow, + result: Result, + ) -> Result { + self.counts.work_after.fetch_add(1, Ordering::SeqCst); + result + } +} + +#[derive(Clone)] +struct RuntimeInsertMiddleware(Arc); + +impl InsertMiddleware for RuntimeInsertMiddleware { + async fn insert_many( + &self, + mut jobs: Vec, + next: InsertNext<'_>, + ) -> Result, riverqueue::Error> { + for job in &mut jobs { + job.opts + .metadata + .insert("middleware", true) + .expect("boolean metadata serializes"); + } + let inserted = next.run(jobs).await?; + self.0 + .insert_after + .fetch_add(inserted.len(), Ordering::SeqCst); + Ok(inserted) + } +} + +struct RuntimePlugin { + counts: Arc, +} + +impl Plugin for RuntimePlugin { + fn install(&self, extensions: &mut Extensions) { + extensions + .hook(RuntimeHook { + counts: Arc::clone(&self.counts), + }) + .insert_middleware(RuntimeInsertMiddleware(Arc::clone(&self.counts))) + .work_middleware(RuntimeWorkMiddleware(Arc::clone(&self.counts))); + } +} + +#[derive(Clone)] +struct RuntimeWorkMiddleware(Arc); + +impl WorkMiddleware for RuntimeWorkMiddleware { + async fn work( + &self, + _context: &WorkContext, + job: JobRow, + next: WorkNext<'_>, + ) -> Result { + // Like River Go, work hooks run inside middleware, so the hook + // hasn't transformed the arguments yet. + let args = job + .decode_args::() + .map_err(WorkError::new)?; + assert!(args.get("hook_decrypted").is_none()); + self.0.work_before.fetch_add(1, Ordering::SeqCst); + let result = next.run(job).await; + self.0.work_after.fetch_add(1, Ordering::SeqCst); + result + } +} + +impl Worker for RuntimeWorker { + type Error = Infallible; + + async fn work( + &self, + _context: WorkContext, + job: Job, + ) -> Result { + assert_eq!( + job.row.decode_args::().unwrap()["hook_decrypted"], + true + ); + tokio::time::sleep(Duration::from_millis(5)).await; + Ok(WorkOutcome::Complete) + } +} + +async fn setup_runtime() -> (Client, Arc, support::PostgresSchema) { + let database = support::PostgresSchema::new("rt_config").await; + let pool = database.pool.clone(); + let schema = database.schema.clone(); + + let mut workers = WorkerRegistry::new(); + workers.register::(RuntimeWorker).unwrap(); + let counts = Arc::new(RuntimeCounts::default()); + let client = Client::builder(PostgresDatabase::new(pool).with_schema(schema)) + .default_max_attempts(7) + .plugin(RuntimePlugin { + counts: Arc::clone(&counts), + }) + .id("rust-runtime-config-test") + .without_notifications() + .workers(workers) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + (client, counts, database) +} + +#[tokio::test] +async fn completion_burst_does_not_lag_large_subscription() { + const JOB_COUNT: usize = 6_000; + + let database = support::PostgresSchema::new("rt_burst").await; + let pool = database.pool.clone(); + let schema = database.schema.clone(); + + let mut workers = WorkerRegistry::new(); + workers.register::(BurstWorker).unwrap(); + let client = Client::builder(PostgresDatabase::new(pool.clone()).with_schema(schema.clone())) + .id("rust-runtime-burst-test") + .without_notifications() + .workers(workers) + .queue( + "default", + QueueConfig::new(1_000) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let mut completed = client + .subscribe_config( + SubscribeConfig::new([EventKind::JobCompleted]) + .unwrap() + .with_buffer_capacity(std::num::NonZeroUsize::new(JOB_COUNT).unwrap()), + ) + .unwrap(); + let jobs = (0..JOB_COUNT).map(|_| (BurstArgs {}, riverqueue::InsertOpts::default())); + assert_eq!(client.insert_many(jobs).await.unwrap().len(), JOB_COUNT); + let expected_ids = sqlx::query_scalar::<_, i64>(AssertSqlSafe(format!( + "SELECT id FROM {}", + schema.qualify("river_job") + ))) + .fetch_all(&pool) + .await + .unwrap() + .into_iter() + .collect::>(); + assert_eq!(expected_ids.len(), JOB_COUNT); + + let mut run_handle = client.start().unwrap(); + run_handle.wait_ready().await.unwrap(); + let received_ids = tokio::time::timeout(Duration::from_secs(10), async { + let mut received_ids = HashSet::with_capacity(JOB_COUNT); + for _ in 0..JOB_COUNT { + let event = completed.recv().await.unwrap(); + assert_eq!(event.kind(), EventKind::JobCompleted); + let id = event.as_job().expect("completion event has a job").job.id; + assert!( + received_ids.insert(id), + "duplicate completion event for job {id}" + ); + } + received_ids + }) + .await + .unwrap(); + assert_eq!(received_ids, expected_ids); + assert!( + tokio::time::timeout(Duration::from_millis(50), completed.recv()) + .await + .is_err(), + "unexpected extra completion event" + ); + run_handle.shutdown().await.unwrap(); + + let completed_count: i64 = sqlx::query_scalar(AssertSqlSafe(format!( + "SELECT count(*) FROM {} WHERE state = 'completed'", + schema.qualify("river_job") + ))) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(completed_count, i64::try_from(JOB_COUNT).unwrap()); + database.cleanup().await; +} + +#[tokio::test] +async fn external_terminal_state_wins_worker_completion_race() { + let database = support::PostgresSchema::new("rt_terminal_race").await; + let pool = database.pool.clone(); + let schema = database.schema.clone(); + + let finish = Arc::new(Semaphore::new(0)); + let started = Arc::new(Semaphore::new(0)); + let mut workers = WorkerRegistry::new(); + workers + .register::(TerminalRaceWorker { + finish: Arc::clone(&finish), + started: Arc::clone(&started), + }) + .unwrap(); + let client = Client::builder(PostgresDatabase::new(pool.clone()).with_schema(schema.clone())) + .id("rust-runtime-terminal-race-test") + .without_notifications() + .workers(workers) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let mut terminal_events = client + .subscribe(&[ + EventKind::JobCancelled, + EventKind::JobCompleted, + EventKind::JobFailed, + ]) + .unwrap(); + let mut run_handle = client.start().unwrap(); + let table = schema.qualify("river_job"); + let state_type = schema.qualify("river_job_state"); + + for external_state in [ + JobState::Cancelled, + JobState::Completed, + JobState::Discarded, + ] { + let inserted = client.insert(TerminalRaceArgs {}).await.unwrap(); + started.acquire().await.unwrap().forget(); + sqlx::query(AssertSqlSafe(format!( + "UPDATE {table} SET finalized_at = now(), \ + metadata = metadata || '{{\"external_terminal\":true}}'::jsonb, \ + state = $2::text::{state_type} \ + WHERE id = $1 AND state = 'running'" + ))) + .bind(inserted.job.row.id) + .bind(external_state.as_str()) + .execute(&pool) + .await + .unwrap(); + finish.add_permits(1); + let event = tokio::time::timeout(Duration::from_secs(5), terminal_events.recv()) + .await + .unwrap() + .unwrap(); + let expected_event = match external_state { + JobState::Cancelled => EventKind::JobCancelled, + JobState::Completed => EventKind::JobCompleted, + JobState::Discarded => EventKind::JobFailed, + _ => unreachable!("test uses terminal external states"), + }; + assert_eq!(event.kind(), expected_event); + let event = event.as_job().unwrap(); + assert_eq!(event.job.id, inserted.job.row.id); + assert_eq!(event.job.state, external_state); + assert_eq!( + event.job.metadata.get::("worker_completion").unwrap(), + Some(true) + ); + + let row = client.jobs().get(inserted.job.row.id).await.unwrap(); + assert_eq!(row.state, external_state); + assert_eq!( + row.metadata.get::("external_terminal").unwrap(), + Some(true) + ); + assert_eq!( + row.metadata.get::("worker_completion").unwrap(), + Some(true) + ); + } + + run_handle.shutdown().await.unwrap(); + database.cleanup().await; +} + +#[tokio::test] +async fn remote_cancellation_overrides_worker_snooze() { + let database = support::PostgresSchema::new("rt_cancel_snooze").await; + let pool = database.pool.clone(); + let schema = database.schema.clone(); + + let started = Arc::new(Semaphore::new(0)); + let mut workers = WorkerRegistry::new(); + workers + .register::(CancelSnoozeWorker { + started: Arc::clone(&started), + }) + .unwrap(); + let client = Client::builder(PostgresDatabase::new(pool.clone()).with_schema(schema.clone())) + .id("rust-runtime-cancel-snooze-test") + .workers(workers) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let mut cancelled_events = client.subscribe(&[EventKind::JobCancelled]).unwrap(); + let mut run_handle = client.start().unwrap(); + run_handle.wait_ready().await.unwrap(); + let inserted = client.insert(CancelSnoozeArgs {}).await.unwrap(); + started.acquire().await.unwrap().forget(); + client.jobs().cancel(inserted.job.row.id).await.unwrap(); + + let event = tokio::time::timeout(Duration::from_secs(5), cancelled_events.recv()) + .await + .unwrap() + .unwrap(); + let row = &event.as_job().expect("cancellation event has a job").job; + assert_eq!(row.id, inserted.job.row.id); + assert_eq!(row.state, JobState::Cancelled); + assert_eq!( + row.errors.last().unwrap().error, + "JobCancelError: job cancelled remotely" + ); + assert_eq!(row.attempt, 1); + + run_handle.shutdown().await.unwrap(); + database.cleanup().await; +} + +#[tokio::test] +#[allow(clippy::too_many_lines)] +async fn shutdown_waits_for_active_work_and_soft_stop_escalates() { + let database = support::PostgresSchema::new("rt_shutdown").await; + let pool = database.pool.clone(); + let schema = database.schema.clone(); + + let graceful_finish = Arc::new(Semaphore::new(0)); + let graceful_started = Arc::new(Semaphore::new(0)); + let mut graceful_workers = WorkerRegistry::new(); + graceful_workers + .register::(ShutdownWorker { + finish: Arc::clone(&graceful_finish), + started: Arc::clone(&graceful_started), + }) + .unwrap(); + let graceful_client = + Client::builder(PostgresDatabase::new(pool.clone()).with_schema(schema.clone())) + .id("rust-runtime-graceful-shutdown-test") + .without_notifications() + .workers(graceful_workers) + .queue( + "graceful", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let active = graceful_client + .insert(ShutdownArgs { + ignore_cancellation: false, + }) + .opts(riverqueue::InsertOpts::default().with_queue("graceful")) + .await + .unwrap(); + let mut graceful_handle = graceful_client.start().unwrap(); + graceful_handle.wait_ready().await.unwrap(); + graceful_started.acquire().await.unwrap().forget(); + let unfetched = graceful_client + .insert(ShutdownArgs { + ignore_cancellation: false, + }) + .opts(riverqueue::InsertOpts::default().with_queue("graceful")) + .await + .unwrap(); + // Request the stop before releasing the worker, then check when the + // shutdown returns that the worker had taken its release: a shutdown that + // didn't wait for active work would return with the permit unclaimed. + graceful_handle.stopper().stop(); + let finish = Arc::clone(&graceful_finish); + let graceful_shutdown = tokio::spawn(async move { + graceful_handle.shutdown().await.unwrap(); + finish.available_permits() + }); + graceful_finish.add_permits(1); + let unclaimed = tokio::time::timeout(Duration::from_secs(2), graceful_shutdown) + .await + .unwrap() + .unwrap(); + assert_eq!( + unclaimed, 0, + "graceful shutdown returned while barrier work was active" + ); + assert_eq!( + graceful_client + .jobs() + .get(active.job.row.id) + .await + .unwrap() + .state, + JobState::Completed + ); + assert_eq!( + graceful_client + .jobs() + .get(unfetched.job.row.id) + .await + .unwrap() + .state, + JobState::Available + ); + + let escalation_started = Arc::new(Semaphore::new(0)); + let mut escalation_workers = WorkerRegistry::new(); + escalation_workers + .register::(ShutdownWorker { + finish: Arc::new(Semaphore::new(0)), + started: Arc::clone(&escalation_started), + }) + .unwrap(); + let escalation_client = + Client::builder(PostgresDatabase::new(pool.clone()).with_schema(schema.clone())) + .id("rust-runtime-soft-stop-escalation-test") + .job_stuck_threshold(Duration::from_millis(10)) + .without_notifications() + .soft_stop_timeout(Duration::from_millis(50)) + .workers(escalation_workers) + .queue( + "escalation", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let mut failed_events = escalation_client + .subscribe(&[EventKind::JobFailed, EventKind::JobInterrupted]) + .unwrap(); + let stuck = escalation_client + .insert(ShutdownArgs { + ignore_cancellation: true, + }) + .opts(riverqueue::InsertOpts::default().with_queue("escalation")) + .await + .unwrap(); + let mut escalation_handle = escalation_client.start().unwrap(); + escalation_handle.wait_ready().await.unwrap(); + escalation_started.acquire().await.unwrap().forget(); + let shutdown_started = tokio::time::Instant::now(); + tokio::time::timeout(Duration::from_secs(2), escalation_handle.shutdown()) + .await + .unwrap() + .unwrap(); + assert!(shutdown_started.elapsed() >= Duration::from_millis(50)); + // The escalated stop aborted the job after the stuck threshold, which + // fails its attempt. + let aborted = escalation_client + .jobs() + .get(stuck.job.row.id) + .await + .unwrap(); + assert_eq!(aborted.attempt, 1); + assert!( + matches!(aborted.state, JobState::Available | JobState::Retryable), + "{aborted:?}" + ); + assert_eq!(aborted.errors.len(), 1); + assert_eq!( + aborted.errors[0].error, + "job aborted after ignoring cancellation" + ); + let event = tokio::time::timeout(Duration::from_secs(1), failed_events.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(event.kind(), EventKind::JobFailed); + assert_eq!(event.as_job().unwrap().job.id, stuck.job.row.id); + + database.cleanup().await; +} + +async fn next_queue_event(receiver: &mut EventReceiver) -> EventKind { + tokio::time::timeout(Duration::from_secs(2), receiver.recv()) + .await + .unwrap() + .unwrap() + .kind() +} + +#[tokio::test] +async fn poll_only_and_subscription_configuration() { + let (client, counts, database) = setup_runtime().await; + let mut completed = client + .subscribe_config( + SubscribeConfig::new([EventKind::JobCompleted]) + .unwrap() + .with_buffer_capacity(std::num::NonZeroUsize::new(4).unwrap()), + ) + .unwrap(); + let mut run_handle = client.start().unwrap(); + run_handle.wait_ready().await.unwrap(); + assert_eq!( + client + .insert_many([ + ( + RuntimeArgs {}, + riverqueue::InsertOpts::default().with_pending(true), + ), + ( + RuntimeArgs {}, + riverqueue::InsertOpts::default().with_pending(true), + ), + ]) + .await + .unwrap() + .len(), + 2 + ); + let inserted = client.insert(RuntimeArgs {}).await.unwrap(); + assert_eq!(inserted.job.row.max_attempts, 7); + assert!( + inserted + .job + .row + .metadata + .get::("middleware") + .unwrap() + .unwrap() + ); + let event = tokio::time::timeout(Duration::from_secs(2), completed.recv()) + .await + .unwrap() + .unwrap(); + let job_event = event.as_job().unwrap(); + assert_eq!(job_event.job.id, inserted.job.row.id); + let statistics = job_event.statistics.unwrap(); + assert!(statistics.run_duration >= Duration::from_millis(5)); + assert!(statistics.complete_duration > Duration::ZERO); + assert!(counts.metrics.load(Ordering::SeqCst) >= 2); + assert_eq!(counts.periodic_starts.load(Ordering::SeqCst), 1); + assert_eq!(counts.insert_before.load(Ordering::SeqCst), 3); + assert_eq!(counts.insert_after.load(Ordering::SeqCst), 3); + assert_eq!(counts.work_before.load(Ordering::SeqCst), 2); + assert_eq!(counts.work_after.load(Ordering::SeqCst), 2); + + let mut lagged = client + .subscribe_config( + SubscribeConfig::new([EventKind::QueuePaused, EventKind::QueueResumed]) + .unwrap() + .with_buffer_capacity(std::num::NonZeroUsize::new(1).unwrap()), + ) + .unwrap(); + let mut transitions = client + .subscribe(&[EventKind::QueuePaused, EventKind::QueueResumed]) + .unwrap(); + client.queues().pause("default").await.unwrap(); + assert_eq!( + next_queue_event(&mut transitions).await, + EventKind::QueuePaused + ); + client.queues().resume("default").await.unwrap(); + assert_eq!( + next_queue_event(&mut transitions).await, + EventKind::QueueResumed + ); + client.queues().pause("default").await.unwrap(); + assert_eq!( + next_queue_event(&mut transitions).await, + EventKind::QueuePaused + ); + assert!(matches!( + lagged.recv().await, + Err(EventRecvError::Lagged(2)) + )); + assert_eq!(lagged.recv().await.unwrap().kind(), EventKind::QueuePaused); + + run_handle.shutdown().await.unwrap(); + assert_eq!( + client.jobs().get(inserted.job.row.id).await.unwrap().state, + JobState::Completed + ); + database.cleanup().await; +} + +#[test] +fn start_without_runtime_returns_error_and_is_restartable() { + let runtime = tokio::runtime::Runtime::new().unwrap(); + let database = runtime.block_on(support::PostgresSchema::new("rt_missing")); + let mut workers = WorkerRegistry::new(); + workers.register::(BurstWorker).unwrap(); + let client = Client::builder( + PostgresDatabase::new(database.pool.clone()).with_schema(database.schema.clone()), + ) + .queue("default", QueueConfig::new(1)) + .workers(workers) + .build() + .unwrap(); + + let Err(error) = client.start() else { + panic!("start should require Tokio"); + }; + assert!(error.to_string().contains("active Tokio runtime")); + + runtime.block_on(async { + let mut run = client + .start() + .expect("failed start must not poison the client"); + run.wait_ready().await.unwrap(); + run.shutdown_now().await.unwrap(); + database.cleanup().await; + }); + drop(client); + runtime.shutdown_background(); +} diff --git a/rust/riverqueue/tests/serde_semantics.rs b/rust/riverqueue/tests/serde_semantics.rs new file mode 100644 index 000000000..a3bfa39ae --- /dev/null +++ b/rust/riverqueue/tests/serde_semantics.rs @@ -0,0 +1,61 @@ +//! River must not change `serde_json` semantics for the rest of an +//! application. Cargo unifies features across a build, so enabling +//! `serde_json`'s `arbitrary_precision` or `preserve_order` anywhere in the +//! workspace would silently alter unrelated user code. These tests fail if +//! either feature is enabled. + +use riverqueue::{JobArgs, JobRow, encoding::encode_args}; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +#[derive(Debug, Deserialize, JobArgs, PartialEq, Serialize)] +#[river(kind = "flattened_float")] +struct FlattenedArgs { + #[serde(flatten)] + inner: Inner, + label: String, +} + +#[derive(Debug, Deserialize, PartialEq, Serialize)] +struct Inner { + amount: f64, +} + +#[test] +fn flattened_float_deserializes_with_river_in_the_graph() { + // `arbitrary_precision` breaks numbers inside `#[serde(flatten)]` with + // "invalid type: map, expected f64". + let decoded: FlattenedArgs = serde_json::from_str(r#"{"amount":1.5,"label":"x"}"#).unwrap(); + assert_eq!(decoded.inner.amount.to_bits(), 1.5_f64.to_bits()); +} + +#[test] +fn flattened_job_args_round_trip_through_a_job_row() { + let args = FlattenedArgs { + inner: Inner { amount: 0.25 }, + label: "invoice".to_owned(), + }; + let row = JobRow::new( + 1, + FlattenedArgs::KIND, + encode_args(&args).unwrap(), + chrono::Utc::now(), + ); + + assert_eq!( + row.encoded_args.get(), + r#"{"amount":0.25,"label":"invoice"}"# + ); + assert_eq!(row.decode_args::().unwrap(), args); +} + +#[test] +fn serde_json_defaults_are_unchanged() { + // Without `arbitrary_precision`, numbers are parsed into `f64`. + let value: Value = serde_json::from_str("1.10").unwrap(); + assert_eq!(value.to_string(), "1.1"); + + // Without `preserve_order`, `Map` iterates in key order. + let map: Map = serde_json::from_str(r#"{"b":1,"a":2}"#).unwrap(); + assert_eq!(map.keys().collect::>(), ["a", "b"]); +} diff --git a/rust/riverqueue/tests/sqlite_parity.rs b/rust/riverqueue/tests/sqlite_parity.rs new file mode 100644 index 000000000..e23912ddf --- /dev/null +++ b/rust/riverqueue/tests/sqlite_parity.rs @@ -0,0 +1,254 @@ +//! SQLite parity tests for maintenance, leadership, and storage semantics that +//! mirror the Go implementation. + +#![cfg(feature = "sqlite")] + +mod support; + +use riverqueue::__private::ClientBuilderExt; +use riverqueue::{Client, Error}; + +#[tokio::test(flavor = "multi_thread")] +async fn queue_pause_and_resume() { + let (pool, path) = support::sqlite_file_pool(2).await; + let client = Client::builder(pool.clone()).build().unwrap(); + + assert!(matches!( + client.queues().pause("missing").await, + Err(Error::NotFound(_)) + )); + assert!(matches!( + client.queues().resume("missing").await, + Err(Error::NotFound(_)) + )); + client + .queues() + .pause(riverqueue::QueueSelector::All) + .await + .unwrap(); + client + .queues() + .resume(riverqueue::QueueSelector::All) + .await + .unwrap(); + + sqlx::query( + "INSERT INTO river_queue (name, created_at, metadata, updated_at) \ + VALUES ('tenant|emails', datetime('now', 'subsec'), jsonb('{}'), datetime('now', 'subsec'))", + ) + .execute(&pool) + .await + .unwrap(); + client.queues().pause("tenant|emails").await.unwrap(); + assert!( + client + .queues() + .get("tenant|emails") + .await + .unwrap() + .paused_at + .is_some() + ); + client.queues().pause("tenant|emails").await.unwrap(); + client.queues().resume("tenant|emails").await.unwrap(); + client.queues().resume("tenant|emails").await.unwrap(); + assert!( + client + .queues() + .get("tenant|emails") + .await + .unwrap() + .paused_at + .is_none() + ); + + support::sqlite_cleanup(pool, path).await; +} + +async fn insert_raw_job(pool: &sqlx::SqlitePool, state: &str) -> i64 { + sqlx::query_scalar( + "INSERT INTO river_job (args, kind, max_attempts, state, attempt, attempted_at, finalized_at) \ + VALUES (jsonb('{}'), 'parity_raw', 25, ?1, \ + CASE WHEN ?1 = 'running' THEN 1 ELSE 0 END, \ + CASE WHEN ?1 = 'running' THEN datetime('now', 'subsec') END, \ + CASE WHEN ?1 IN ('cancelled', 'completed', 'discarded') THEN datetime('now', 'subsec') END) \ + RETURNING id", + ) + .bind(state) + .fetch_one(pool) + .await + .unwrap() +} + +#[tokio::test(flavor = "multi_thread")] +async fn job_delete_many() { + let (pool, path) = support::sqlite_file_pool(2).await; + let client = Client::builder(pool.clone()).build().unwrap(); + + let running = insert_raw_job(&pool, "running").await; + let first = insert_raw_job(&pool, "available").await; + let second = insert_raw_job(&pool, "completed").await; + let third = insert_raw_job(&pool, "cancelled").await; + + // Running jobs are excluded before the limit applies. + let deleted = client + .jobs() + .delete_many(riverqueue::JobDeleteManyParams::matching( + riverqueue::JobListParams::default() + .ids([running, first, second, third]) + .limit(2), + )) + .await + .unwrap(); + assert_eq!( + deleted.iter().map(|job| job.id).collect::>(), + vec![first, second] + ); + let remaining = client + .jobs() + .delete_many(riverqueue::JobDeleteManyParams::all()) + .await + .unwrap(); + assert_eq!( + remaining.iter().map(|job| job.id).collect::>(), + vec![third] + ); + assert_eq!( + client.jobs().get(running).await.unwrap().state, + riverqueue::JobState::Running + ); + + support::sqlite_cleanup(pool, path).await; +} + +/// Hook invocations as `(operation, job ID, state)`. +type HookCalls = std::sync::Arc>>; + +#[derive(Clone, Default)] +struct HookPilot { + calls: HookCalls, + fail: bool, +} + +#[async_trait::async_trait] +impl riverqueue::__private::Pilot for HookPilot { + fn intercepts_job_cancel_retry(&self) -> bool { + true + } + + async fn after_job_cancel( + &self, + connection: riverqueue::__private::DatabaseConnection<'_>, + job: &riverqueue::__private::JobUpdatedParams, + ) -> Result<(), riverqueue::__private::PilotError> { + // The hook sees the update inside the same transaction. + let connection = connection.into_sqlite().expect("SQLite connection"); + let state: String = sqlx::query_scalar("SELECT state FROM river_job WHERE id = ?") + .bind(job.job.id) + .fetch_one(connection) + .await?; + self.calls + .lock() + .unwrap() + .push(("cancel", job.job.id, state)); + if self.fail { + return Err(std::io::Error::other("cancel hook failed").into()); + } + Ok(()) + } + + async fn after_job_retry( + &self, + _connection: riverqueue::__private::DatabaseConnection<'_>, + job: &riverqueue::__private::JobUpdatedParams, + ) -> Result<(), riverqueue::__private::PilotError> { + self.calls + .lock() + .unwrap() + .push(("retry", job.job.id, job.job.state.as_str().to_owned())); + if self.fail { + return Err(std::io::Error::other("retry hook failed").into()); + } + Ok(()) + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn cancel_and_retry_post_hooks_share_the_transaction() { + let (pool, path) = support::sqlite_file_pool(2).await; + let pilot = HookPilot::default(); + let client = Client::builder(pool.clone()) + .pilot(pilot.clone()) + .build() + .unwrap(); + let id = insert_raw_job(&pool, "available").await; + + client.jobs().cancel(id).await.unwrap(); + client.jobs().retry(id).await.unwrap(); + assert_eq!( + *pilot.calls.lock().unwrap(), + [ + ("cancel", id, "cancelled".to_owned()), + ("retry", id, "available".to_owned()) + ] + ); + + let failing = Client::builder(pool.clone()) + .pilot(HookPilot { + fail: true, + ..HookPilot::default() + }) + .build() + .unwrap(); + assert!(failing.jobs().cancel(id).await.is_err()); + assert_eq!( + client.jobs().get(id).await.unwrap().state, + riverqueue::JobState::Available + ); + + support::sqlite_cleanup(pool, path).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn extension_notify_many_writes_the_outbox() { + use riverqueue::__private::{ + DatabaseConfig, DatabaseConnection, NotificationTopic, notify_many, + }; + + let (pool, path) = support::sqlite_file_pool(2).await; + let mut rolled_back = pool.begin().await.unwrap(); + notify_many( + DatabaseConnection::Sqlite(&mut rolled_back), + &DatabaseConfig::Sqlite, + NotificationTopic::Control, + &["rolled back".to_owned()], + ) + .await + .unwrap(); + rolled_back.rollback().await.unwrap(); + + let mut committed = pool.begin().await.unwrap(); + notify_many( + DatabaseConnection::Sqlite(&mut committed), + &DatabaseConfig::Sqlite, + NotificationTopic::Control, + &["first".to_owned(), "second".to_owned()], + ) + .await + .unwrap(); + committed.commit().await.unwrap(); + + let rows: Vec<(String, String)> = + sqlx::query_as("SELECT topic, payload FROM river_notification ORDER BY id") + .fetch_all(&pool) + .await + .unwrap(); + assert_eq!( + rows, + [ + ("river_control".to_owned(), "first".to_owned()), + ("river_control".to_owned(), "second".to_owned()) + ] + ); + support::sqlite_cleanup(pool, path).await; +} diff --git a/rust/riverqueue/tests/sqlite_runtime.rs b/rust/riverqueue/tests/sqlite_runtime.rs new file mode 100644 index 000000000..e24ccea83 --- /dev/null +++ b/rust/riverqueue/tests/sqlite_runtime.rs @@ -0,0 +1,1856 @@ +use std::{ + convert::Infallible, + sync::{ + Arc, Mutex, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use async_trait::async_trait; +use riverqueue::__private::{ + ClaimedJob, DatabaseConnection, JobInsertParams, JobSetStateParams, Pilot, PilotError, + PilotProducer, ProducerClaimContext, ProducerClaimNext, ProducerStartContext, RescueParams, +}; +use riverqueue::__private::{ClientBuilderExt, ExtensionClient, PreparedInsertParams}; +use riverqueue::__private::{MaintenanceService, MaintenanceServiceContext}; +use riverqueue::{ + BoxError, Client, EventKind, Hook, InsertBatch, InsertOpts, Job, JobArgs, JobRow, JobState, + MaintenanceConfig, QueueConfig, UniqueOpts, WorkContext, WorkOutcome, WorkerRegistry, + database::DatabaseKind, +}; +use riverqueue_migrate::SqliteMigrator; +use serde::{Deserialize, Serialize}; +use sqlx::sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}; +use tokio::sync::Semaphore; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_sqlite_runtime")] +struct RuntimeArgs { + value: i64, +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_sqlite_cancel")] +struct CancelArgs {} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_sqlite_cancel_ignored")] +struct CancelIgnoredArgs {} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_sqlite_unknown")] +struct UnknownArgs {} + +struct WrapperTransformHook(&'static str); + +#[allow( + clippy::unused_async_trait_impl, + reason = "these extensions only record state synchronously" +)] +impl Hook for WrapperTransformHook { + async fn decode_insert_result(&self, job: &mut JobRow) -> Result<(), BoxError> { + // Unwrap without reparsing the inner arguments so their exact bytes + // are preserved. + let mut outer: std::collections::HashMap> = + job.decode_args()?; + job.encoded_args = outer + .remove(self.0) + .ok_or_else(|| format!("missing outer insertion wrapper {:?}", self.0))?; + Ok(()) + } + + async fn insert_begin(&self, insert: &mut riverqueue::InsertContext) -> Result<(), BoxError> { + // As a hook that wraps arguments must, leave arguments that already + // carry the outermost wrapper alone, so a stored job inserted again + // keeps its arguments. + if let Ok(outer) = serde_json::from_str::< + std::collections::HashMap>, + >(insert.encoded_args.get()) + && outer.len() == 1 + && outer.contains_key("B") + { + return Ok(()); + } + insert.encoded_args = serde_json::value::RawValue::from_string(format!( + "{{{}:{}}}", + serde_json::to_string(self.0)?, + insert.encoded_args.get() + ))?; + Ok(()) + } +} + +#[derive(Clone, Copy)] +enum CompletionBehavior { + Continue, + Fail, + Mark, +} + +#[derive(Clone, Copy)] +enum SelectionBehavior { + Fail, + FailFirst, + Success, +} + +#[derive(Clone)] +struct SqlitePilot { + completion: Option, + completion_calls: Arc, + fetch: Option, + fetch_calls: Arc, + insert: Option, + insert_calls: Arc, + maintenance_service: Option>, + rescue: Option, + rescue_calls: Arc, +} + +impl SqlitePilot { + fn new() -> Self { + Self { + completion: None, + completion_calls: Arc::new(AtomicUsize::new(0)), + fetch: None, + fetch_calls: Arc::new(AtomicUsize::new(0)), + insert: None, + insert_calls: Arc::new(AtomicUsize::new(0)), + maintenance_service: None, + rescue: None, + rescue_calls: Arc::new(AtomicUsize::new(0)), + } + } +} + +#[async_trait] +impl Pilot for SqlitePilot { + fn intercepts_insert(&self) -> bool { + self.insert.is_some() + } + + fn intercepts_job_set_state(&self) -> bool { + self.completion.is_some() + } + + fn intercepts_rescue(&self) -> bool { + self.rescue.is_some() + } + + fn maintenance_services(&self) -> Vec> { + self.maintenance_service + .as_ref() + .map(|state| { + vec![Arc::new(LeadershipService(Arc::clone(state))) as Arc] + }) + .unwrap_or_default() + } + + async fn after_jobs_set_state( + &self, + connection: DatabaseConnection<'_>, + params: &JobSetStateParams, + ) -> Result<(), PilotError> { + self.completion_calls.fetch_add(1, Ordering::SeqCst); + let connection = connection + .into_sqlite() + .ok_or_else(|| std::io::Error::other("expected SQLite completion connection"))?; + for job in params.jobs { + sqlx::query("INSERT INTO pilot_effect (operation, job_id) VALUES ('completion', ?)") + .bind(job.id) + .execute(&mut *connection) + .await?; + } + match self.completion.expect("completion interception is enabled") { + CompletionBehavior::Fail => { + Err(std::io::Error::other("completion interception failed").into()) + } + CompletionBehavior::Continue => Ok(()), + CompletionBehavior::Mark => { + for job in params.jobs { + sqlx::query( + "UPDATE river_job SET \ + metadata = jsonb_set(metadata, '$.pilot_handled', jsonb('true')) \ + WHERE id = ?", + ) + .bind(job.id) + .execute(&mut *connection) + .await?; + } + Ok(()) + } + } + } + + async fn before_job_insert( + &self, + connection: DatabaseConnection<'_>, + params: &mut JobInsertParams<'_>, + ) -> Result<(), PilotError> { + self.insert_calls.fetch_add(1, Ordering::SeqCst); + let connection = connection + .into_sqlite() + .ok_or_else(|| std::io::Error::other("expected SQLite insertion connection"))?; + let marker: String = + sqlx::query_scalar("SELECT marker FROM pilot_insert_config WHERE queue = ?") + .bind(&*params.queue) + .fetch_one(&mut *connection) + .await?; + params.metadata.insert("pilot_insert", marker)?; + sqlx::query("INSERT INTO pilot_effect (operation, job_id) VALUES ('insert', 0)") + .execute(&mut *connection) + .await?; + if matches!(self.insert, Some(SelectionBehavior::Fail)) { + return Err(std::io::Error::other("insert interception failed").into()); + } + Ok(()) + } + + async fn start_producer( + &self, + _context: ProducerStartContext, + ) -> Result>, PilotError> { + Ok(self + .fetch + .map(|_| Box::new(SqliteFetchSession(self.clone())) as Box)) + } + + async fn select_rescue_job_ids( + &self, + connection: DatabaseConnection<'_>, + params: &RescueParams, + ) -> Result>, PilotError> { + // River bounds the selection like its own rescuer reads. + if params.timeout != Duration::from_secs(30) { + return Err(std::io::Error::other("unexpected rescue timeout").into()); + } + self.rescue_calls.fetch_add(1, Ordering::SeqCst); + let connection = connection + .into_sqlite() + .ok_or_else(|| std::io::Error::other("expected SQLite rescue connection"))?; + let ids = sqlx::query_scalar( + "SELECT id FROM river_job WHERE state = 'running' ORDER BY id LIMIT ?", + ) + .bind(params.maximum) + .fetch_all(&mut *connection) + .await?; + for id in &ids { + sqlx::query("INSERT INTO pilot_effect (operation, job_id) VALUES ('rescue', ?)") + .bind(id) + .execute(&mut *connection) + .await?; + } + if matches!(self.rescue, Some(SelectionBehavior::Fail)) { + return Err(std::io::Error::other("rescue interception failed").into()); + } + Ok(Some(ids)) + } +} + +#[derive(Default)] +struct LeadershipServiceState { + starts: AtomicUsize, + stops: AtomicUsize, +} + +struct LeadershipService(Arc); + +/// Claims for [`SqlitePilot`]: records a side effect in each claim's +/// transaction, fails as configured, and otherwise runs River's claim. +struct SqliteFetchSession(SqlitePilot); + +#[async_trait] +impl PilotProducer for SqliteFetchSession { + fn intercepts_claim(&self) -> bool { + true + } + + async fn claim( + &self, + context: ProducerClaimContext<'_>, + next: ProducerClaimNext<'_>, + ) -> Result, PilotError> { + let pilot = &self.0; + pilot.fetch_calls.fetch_add(1, Ordering::SeqCst); + let mut transaction = context.database.begin().await?; + let connection = transaction + .connection() + .into_sqlite() + .ok_or_else(|| std::io::Error::other("expected SQLite fetch connection"))?; + sqlx::query("INSERT INTO pilot_effect (operation, job_id) VALUES ('fetch', 0)") + .execute(&mut *connection) + .await?; + if matches!(pilot.fetch, Some(SelectionBehavior::Fail)) + || matches!(pilot.fetch, Some(SelectionBehavior::FailFirst)) + && pilot.fetch_calls.load(Ordering::SeqCst) == 1 + { + return Err(std::io::Error::other("fetch interception failed").into()); + } + let jobs = next.claim(transaction.connection()).await?; + transaction.commit().await?; + Ok(jobs) + } +} + +#[async_trait] +impl MaintenanceService for LeadershipService { + async fn run(&self, context: MaintenanceServiceContext) -> Result<(), PilotError> { + let cancellation = context.term.token; + self.0.starts.fetch_add(1, Ordering::SeqCst); + cancellation.cancelled().await; + self.0.stops.fetch_add(1, Ordering::SeqCst); + Ok(()) + } +} + +/// The current leader's lease expiry, which each successful renewal moves. +#[tokio::test] +async fn sqlite_database_errors_are_sqlx_errors() { + let pool = setup().await; + let client = Client::builder(pool.clone()).build().unwrap(); + let inserted = client.insert(UnknownArgs {}).await.unwrap(); + sqlx::query("UPDATE river_job SET metadata = jsonb('[]') WHERE id = ?") + .bind(inserted.job.row.id) + .execute(&pool) + .await + .unwrap(); + + let error = client.jobs().get(inserted.job.row.id).await.unwrap_err(); + assert!( + matches!(&error, riverqueue::Error::Database(sqlx::Error::Decode(_))), + "{error:?}" + ); + + pool.close().await; + let error = client.jobs().get(inserted.job.row.id).await.unwrap_err(); + assert!( + matches!(error, riverqueue::Error::Database(sqlx::Error::PoolClosed)), + "{error:?}" + ); +} + +async fn leader_expires_at(pool: &sqlx::SqlitePool) -> Option { + sqlx::query_scalar("SELECT CAST(expires_at AS TEXT) FROM river_leader") + .fetch_optional(pool) + .await + .unwrap() +} + +async fn setup() -> sqlx::SqlitePool { + let options = SqliteConnectOptions::new() + .filename(":memory:") + .create_if_missing(true) + .journal_mode(SqliteJournalMode::Wal); + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect_with(options) + .await + .unwrap(); + SqliteMigrator::new(pool.clone()) + .migrate_up() + .await + .unwrap(); + sqlx::query("CREATE TABLE pilot_effect (operation TEXT NOT NULL, job_id INTEGER NOT NULL)") + .execute(&pool) + .await + .unwrap(); + sqlx::query("CREATE TABLE pilot_insert_config (queue TEXT PRIMARY KEY, marker TEXT NOT NULL)") + .execute(&pool) + .await + .unwrap(); + sqlx::query("INSERT INTO pilot_insert_config (queue, marker) VALUES ('default', 'default')") + .execute(&pool) + .await + .unwrap(); + pool +} + +async fn setup_file_pool(busy_timeout: Duration) -> (sqlx::SqlitePool, std::path::PathBuf) { + static DATABASE_NONCE: AtomicUsize = AtomicUsize::new(0); + let database_path = std::env::temp_dir().join(format!( + "river-sqlite-runtime-{}-{}.sqlite", + std::process::id(), + DATABASE_NONCE.fetch_add(1, Ordering::Relaxed) + )); + let options = SqliteConnectOptions::new() + .filename(&database_path) + .create_if_missing(true) + .journal_mode(SqliteJournalMode::Wal) + .busy_timeout(busy_timeout); + let pool = SqlitePoolOptions::new() + .max_connections(4) + .connect_with(options) + .await + .unwrap(); + SqliteMigrator::new(pool.clone()) + .migrate_up() + .await + .unwrap(); + (pool, database_path) +} + +fn remove_sqlite_files(database_path: &std::path::Path) { + let _ = std::fs::remove_file(database_path); + for suffix in ["-shm", "-wal"] { + let mut path = database_path.as_os_str().to_owned(); + path.push(suffix); + let _ = std::fs::remove_file(path); + } +} + +fn runtime_workers(worked: Arc) -> WorkerRegistry { + let mut workers = WorkerRegistry::new(); + workers + .register_fn(move |_context: WorkContext, _job: Job| { + let worked = Arc::clone(&worked); + async move { + worked.add_permits(1); + Ok::<_, Infallible>(WorkOutcome::Complete) + } + }) + .unwrap(); + workers +} + +#[tokio::test] +async fn sqlite_pilot_completion_continue_and_mark_are_atomic() { + for behavior in [CompletionBehavior::Continue, CompletionBehavior::Mark] { + let pool = setup().await; + let worked = Arc::new(Semaphore::new(0)); + let mut pilot = SqlitePilot::new(); + pilot.completion = Some(behavior); + pilot.fetch = Some(SelectionBehavior::Success); + let client = Client::builder(pool.clone()) + .id("sqlite-pilot-completion") + .pilot(pilot.clone()) + .workers(runtime_workers(Arc::clone(&worked))) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let mut events = client.subscribe(&[EventKind::JobCompleted]).unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + let inserted = client.insert(RuntimeArgs { value: 1 }).await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), worked.acquire()) + .await + .unwrap() + .unwrap() + .forget(); + let event = tokio::time::timeout(Duration::from_secs(5), events.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!( + event.as_job().map(|event| event.job.id), + Some(inserted.job.row.id) + ); + + let row = client.jobs().get(inserted.job.row.id).await.unwrap(); + assert_eq!(row.state, JobState::Completed); + assert_eq!( + row.metadata.get::("pilot_handled").unwrap(), + matches!(behavior, CompletionBehavior::Mark).then_some(true) + ); + let effects: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pilot_effect WHERE operation = 'completion' AND job_id = ?", + ) + .bind(inserted.job.row.id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(effects, 1); + assert!(pilot.fetch_calls.load(Ordering::SeqCst) >= 1); + assert_eq!(pilot.completion_calls.load(Ordering::SeqCst), 1); + + run.shutdown().await.unwrap(); + } +} + +#[tokio::test] +async fn sqlite_pilot_completion_error_rolls_back_side_effects() { + let pool = setup().await; + let worked = Arc::new(Semaphore::new(0)); + let mut pilot = SqlitePilot::new(); + pilot.completion = Some(CompletionBehavior::Fail); + let client = Client::builder(pool.clone()) + .id("sqlite-pilot-completion-error") + .pilot(pilot.clone()) + .workers(runtime_workers(Arc::clone(&worked))) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + let inserted = client.insert(RuntimeArgs { value: 1 }).await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), worked.acquire()) + .await + .unwrap() + .unwrap() + .forget(); + tokio::time::timeout(Duration::from_secs(5), async { + while pilot.completion_calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let effects: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pilot_effect WHERE operation = 'completion' AND job_id = ?", + ) + .bind(inserted.job.row.id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(effects, 0); + assert_eq!( + client.jobs().get(inserted.job.row.id).await.unwrap().state, + JobState::Running + ); + + run.shutdown_now().await.unwrap(); +} + +#[tokio::test] +async fn sqlite_pilot_fetch_error_rolls_back_selection_side_effects() { + let pool = setup().await; + let mut pilot = SqlitePilot::new(); + pilot.fetch = Some(SelectionBehavior::Fail); + let client = Client::builder(pool.clone()) + .id("sqlite-pilot-fetch-error") + .pilot(pilot.clone()) + .workers(runtime_workers(Arc::new(Semaphore::new(0)))) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let inserted = client.insert(RuntimeArgs { value: 1 }).await.unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + tokio::time::timeout(Duration::from_secs(5), async { + while pilot.fetch_calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let effects: i64 = + sqlx::query_scalar("SELECT count(*) FROM pilot_effect WHERE operation = 'fetch'") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(effects, 0); + assert_eq!( + client.jobs().get(inserted.job.row.id).await.unwrap().state, + JobState::Available + ); + + let _ = run.shutdown_now().await; +} + +#[tokio::test] +async fn sqlite_pilot_fetch_transient_error_retries_without_stopping_the_queue() { + let pool = setup().await; + let worked = Arc::new(Semaphore::new(0)); + let mut pilot = SqlitePilot::new(); + pilot.fetch = Some(SelectionBehavior::FailFirst); + let client = Client::builder(pool.clone()) + .id("sqlite-pilot-fetch-retry") + .pilot(pilot.clone()) + .workers(runtime_workers(Arc::clone(&worked))) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let inserted = client.insert(RuntimeArgs { value: 1 }).await.unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + tokio::time::timeout(Duration::from_secs(5), worked.acquire()) + .await + .unwrap() + .unwrap() + .forget(); + tokio::time::timeout(Duration::from_secs(5), async { + while client.jobs().get(inserted.job.row.id).await.unwrap().state != JobState::Completed { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + run.shutdown().await.unwrap(); + + let fetch_calls = pilot.fetch_calls.load(Ordering::SeqCst); + let effects: i64 = + sqlx::query_scalar("SELECT count(*) FROM pilot_effect WHERE operation = 'fetch'") + .fetch_one(&pool) + .await + .unwrap(); + assert!(fetch_calls >= 2); + assert_eq!(usize::try_from(effects).unwrap() + 1, fetch_calls); +} + +#[tokio::test] +async fn sqlite_queue_start_retries_transient_write_contention() { + let (pool, database_path) = setup_file_pool(Duration::from_millis(1)).await; + let worked = Arc::new(Semaphore::new(0)); + let client = Client::builder(pool.clone()) + .id("sqlite-queue-start-retry") + .workers(runtime_workers(Arc::clone(&worked))) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let inserted = client.insert(RuntimeArgs { value: 1 }).await.unwrap(); + let writer = pool.begin_with("BEGIN IMMEDIATE").await.unwrap(); + let release_writer = tokio::spawn(async move { + tokio::time::sleep(Duration::from_millis(50)).await; + writer.rollback().await.unwrap(); + }); + + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + release_writer.await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), worked.acquire()) + .await + .unwrap() + .unwrap() + .forget(); + tokio::time::timeout(Duration::from_secs(5), async { + while client.jobs().get(inserted.job.row.id).await.unwrap().state != JobState::Completed { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + + run.shutdown().await.unwrap(); + pool.close().await; + remove_sqlite_files(&database_path); +} + +#[tokio::test] +async fn sqlite_transient_renewal_contention_preserves_leadership_services() { + let (pool, database_path) = setup_file_pool(Duration::from_millis(5)).await; + let service = Arc::new(LeadershipServiceState::default()); + let mut pilot = SqlitePilot::new(); + pilot.maintenance_service = Some(Arc::clone(&service)); + let client = Client::builder(pool.clone()) + .id("sqlite-contention-leader") + .maintenance(MaintenanceConfig::default().with_elect_interval(Duration::from_millis(10))) + .pilot(pilot) + .queue( + "default", + QueueConfig::new(1).with_fetch_poll_interval(Duration::from_mins(1)), + ) + .workers(runtime_workers(Arc::new(Semaphore::new(0)))) + .build() + .unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + // The client is the only one bidding, so it wins the first election. + // Requesting resignations here would churn leadership and leave + // requests in the outbox that could arrive during the contention below. + let startup_deadline = tokio::time::Instant::now() + Duration::from_secs(5); + while service.starts.load(Ordering::SeqCst) == 0 { + assert!( + tokio::time::Instant::now() < startup_deadline, + "leadership services never started" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } + let starts_before_contention = service.starts.load(Ordering::SeqCst); + let stops_before_contention = service.stops.load(Ordering::SeqCst); + let lease_before_contention = leader_expires_at(&pool).await; + let mut writer = pool.begin_with("BEGIN IMMEDIATE").await.unwrap(); + sqlx::query("UPDATE river_queue SET updated_at = updated_at WHERE name = 'default'") + .execute(&mut *writer) + .await + .unwrap(); + // Hold the write lock across several 10 ms renewal attempts. + tokio::time::sleep(Duration::from_millis(100)).await; + writer.rollback().await.unwrap(); + + // Once a renewal after the contention has succeeded, the services have + // been through the whole contention; the counters show whether they + // stopped at any point. + let renewal_deadline = tokio::time::Instant::now() + Duration::from_secs(5); + while leader_expires_at(&pool).await == lease_before_contention { + assert!( + tokio::time::Instant::now() < renewal_deadline, + "leadership was never renewed after the contention" + ); + tokio::time::sleep(Duration::from_millis(5)).await; + } + assert_eq!( + service.starts.load(Ordering::SeqCst), + starts_before_contention + ); + assert_eq!( + service.stops.load(Ordering::SeqCst), + stops_before_contention + ); + + // A renewal that failed during the contention backs off for about a + // second before the leader reads its next wakeup. + client.request_resign().await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), async { + while service.stops.load(Ordering::SeqCst) == stops_before_contention { + tokio::time::sleep(Duration::from_millis(5)).await; + } + }) + .await + .expect("leader resigns on request"); + + run.shutdown_now().await.unwrap(); + assert!(service.stops.load(Ordering::SeqCst) >= 1); + pool.close().await; + remove_sqlite_files(&database_path); +} + +#[tokio::test] +async fn sqlite_pilot_insert_uses_the_insertion_transaction() { + let pool = setup().await; + let mut pilot = SqlitePilot::new(); + pilot.insert = Some(SelectionBehavior::Success); + let client = Client::builder(pool.clone()) + .pilot(pilot.clone()) + .build() + .unwrap(); + + let mut transaction = pool.begin_with("BEGIN IMMEDIATE").await.unwrap(); + sqlx::query("UPDATE pilot_insert_config SET marker = 'uncommitted' WHERE queue = 'default'") + .execute(&mut *transaction) + .await + .unwrap(); + let inserted = client + .insert(RuntimeArgs { value: 31 }) + .tx(&mut transaction) + .await + .unwrap(); + assert_eq!( + inserted + .job + .row + .metadata + .get::("pilot_insert") + .unwrap() + .as_deref(), + Some("uncommitted") + ); + transaction.commit().await.unwrap(); + + assert_eq!(pilot.insert_calls.load(Ordering::SeqCst), 1); + let effects: i64 = + sqlx::query_scalar("SELECT count(*) FROM pilot_effect WHERE operation = 'insert'") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(effects, 1); + + let mut failing_pilot = SqlitePilot::new(); + failing_pilot.insert = Some(SelectionBehavior::Fail); + let failing_client = Client::builder(pool.clone()) + .pilot(failing_pilot.clone()) + .build() + .unwrap(); + let error = failing_client + .insert(RuntimeArgs { value: 32 }) + .await + .unwrap_err(); + assert!( + matches!( + &error, + riverqueue::Error::Extension { + phase: riverqueue::ExtensionPhase::AddOn { operation: "job insertion" }, + source, + } if source.to_string().contains("insert interception failed") + ), + "{error:?}" + ); + let failed_jobs: i64 = sqlx::query_scalar( + "SELECT count(*) FROM river_job WHERE json_extract(args, '$.value') = 32", + ) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(failed_jobs, 0); + let effects_after_failure: i64 = + sqlx::query_scalar("SELECT count(*) FROM pilot_effect WHERE operation = 'insert'") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(effects_after_failure, 1); + assert_eq!(failing_pilot.insert_calls.load(Ordering::SeqCst), 1); +} + +#[tokio::test] +#[allow( + clippy::too_many_lines, + reason = "one exact-version insertion regression verifies every preserved/reset wire field and hook phase" +)] +async fn sqlite_reinsert_preserves_wire_fields_and_runs_the_canonical_pipeline() { + let pool = setup().await; + let producer = Client::builder(pool.clone()) + .hook(WrapperTransformHook("A")) + .hook(WrapperTransformHook("B")) + .build() + .unwrap(); + let mut pilot = SqlitePilot::new(); + pilot.insert = Some(SelectionBehavior::Success); + let mut workers = WorkerRegistry::new(); + workers + .register_fn(|_context: WorkContext, _job: Job| async move { + Ok::<_, Infallible>(WorkOutcome::Complete) + }) + .unwrap(); + let client = Client::builder(pool.clone()) + .hook(WrapperTransformHook("A")) + .hook(WrapperTransformHook("B")) + .pilot(pilot.clone()) + .workers(workers) + .build() + .unwrap(); + let scheduled_at = chrono::Utc::now() + chrono::Duration::hours(2); + let raw = riverqueue::__private::ExtensionClient::new(&client) + .insert_raw( + CancelArgs::KIND, + &[], + serde_json::value::to_raw_value(&serde_json::json!({"raw": true})).unwrap(), + InsertOpts::default().with_pending(true), + ) + .await + .unwrap(); + assert_eq!(raw.job.encoded_args.get(), r#"{"raw":true}"#); + let stored_raw_args: String = + sqlx::query_scalar("SELECT json(args) FROM river_job WHERE id = ?") + .bind(raw.job.id) + .fetch_one(&pool) + .await + .unwrap(); + let stored_raw_args: serde_json::Value = serde_json::from_str(&stored_raw_args).unwrap(); + assert_eq!(stored_raw_args["B"]["A"]["raw"], true); + let mut original = producer + .insert(RuntimeArgs { value: 41 }) + .opts( + InsertOpts::default() + .with_metadata(serde_json::Map::from_iter([( + "source".to_owned(), + serde_json::json!(true), + )])) + .with_scheduled_at(scheduled_at) + .with_tags(["reinserted"]) + .with_unique(UniqueOpts::new().with_by_args(true)), + ) + .await + .unwrap() + .job + .row; + sqlx::query("UPDATE river_job SET metadata = jsonb_patch(metadata, jsonb('{\"beyond_float\":1e400}')) WHERE id = ?") + .bind(original.id) + .execute(&pool) + .await + .unwrap(); + original.metadata = producer.jobs().get(original.id).await.unwrap().metadata; + assert_eq!( + original.metadata.get_raw("beyond_float").unwrap().get(), + "1e400" + ); + assert_eq!(original.encoded_args.get(), r#"{"value":41}"#); + let stored_source_args: String = + sqlx::query_scalar("SELECT json(args) FROM river_job WHERE id = ?") + .bind(original.id) + .fetch_one(&pool) + .await + .unwrap(); + let stored_source_args = serde_json::value::RawValue::from_string(stored_source_args).unwrap(); + assert_eq!( + serde_json::from_str::(stored_source_args.get()).unwrap()["B"]["A"]["value"], + 41 + ); + let sentinel = producer + .insert(RuntimeArgs { value: 42 }) + .opts(InsertOpts::default().with_scheduled_at(scheduled_at)) + .await + .unwrap(); + + let mut transaction = pool.begin_with("BEGIN IMMEDIATE").await.unwrap(); + sqlx::query("DELETE FROM river_job WHERE id = ?") + .bind(original.id) + .execute(&mut *transaction) + .await + .unwrap(); + let reinserted = ExtensionClient::new(&client) + .insert_prepared(vec![PreparedInsertParams { + created_at: original.created_at, + encoded_args: stored_source_args, + kind: "x".to_owned(), + max_attempts: original.max_attempts, + metadata: original.metadata.clone(), + priority: original.priority, + queue: original.queue.clone(), + scheduled_at: original.scheduled_at, + tags: original.tags.clone(), + unique_key: original.unique_key.clone(), + unique_states: original.unique_states.clone(), + }]) + .tx(&mut transaction) + .await + .unwrap() + .remove(0); + transaction.commit().await.unwrap(); + + assert_ne!(reinserted.job.id, original.id); + assert!(reinserted.job.id > sentinel.job.row.id); + assert_eq!(reinserted.job.attempt, 0); + assert!(reinserted.job.attempted_at.is_none()); + assert_eq!(reinserted.job.attempted_by, Vec::::new()); + assert_eq!(reinserted.job.created_at, original.created_at); + assert_eq!(reinserted.job.errors, []); + assert!(reinserted.job.finalized_at.is_none()); + assert_eq!(reinserted.job.scheduled_at, original.scheduled_at); + assert_eq!(reinserted.job.state, JobState::Available); + assert_eq!(reinserted.job.kind, "x"); + assert_eq!(reinserted.job.unique_key, original.unique_key); + assert_eq!(reinserted.job.unique_states, original.unique_states); + assert_eq!( + reinserted + .job + .metadata + .get_raw("beyond_float") + .unwrap() + .get(), + "1e400" + ); + assert_eq!( + reinserted.job.metadata.get::("source").unwrap(), + Some(true) + ); + assert_eq!( + reinserted + .job + .metadata + .get::("pilot_insert") + .unwrap() + .as_deref(), + Some("default") + ); + assert!(!reinserted.unique_skipped_as_duplicate); + assert_eq!(pilot.insert_calls.load(Ordering::SeqCst), 2); + + assert_eq!(reinserted.job.encoded_args.get(), r#"{"value":41}"#); + let stored_args: String = sqlx::query_scalar("SELECT json(args) FROM river_job WHERE id = ?") + .bind(reinserted.job.id) + .fetch_one(&pool) + .await + .unwrap(); + let stored_args: serde_json::Value = serde_json::from_str(&stored_args).unwrap(); + assert_eq!(stored_args["B"]["A"]["value"], 41); + + let notifications: i64 = + sqlx::query_scalar("SELECT count(*) FROM river_notification WHERE topic = 'river_insert'") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(notifications, 1); +} + +#[tokio::test] +async fn sqlite_pilot_rescue_selection_and_update_share_a_transaction() { + let pool = setup().await; + let mut pilot = SqlitePilot::new(); + pilot.rescue = Some(SelectionBehavior::Success); + let client = Client::builder(pool.clone()) + .id("sqlite-pilot-rescue") + .maintenance( + MaintenanceConfig::default() + .with_elect_interval(Duration::from_millis(10)) + .with_rescue_after(Duration::from_mins(1)) + .with_rescuer_interval(Duration::from_millis(10)), + ) + .pilot(pilot.clone()) + .queue( + "default", + QueueConfig::new(1).with_fetch_poll_interval(Duration::from_mins(1)), + ) + .workers(runtime_workers(Arc::new(Semaphore::new(0)))) + .build() + .unwrap(); + let inserted = client.insert(RuntimeArgs { value: 1 }).await.unwrap(); + sqlx::query( + "UPDATE river_job SET state = 'running', attempt = 1, \ + attempted_at = '2000-01-01 00:00:00.000' WHERE id = ?", + ) + .bind(inserted.job.row.id) + .execute(&pool) + .await + .unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + // Wait for the rescue's recorded error rather than the `retryable` state, + // which the scheduler ends about a second later and a loaded poll can miss. + let rescued = tokio::time::timeout(Duration::from_secs(5), async { + loop { + let row = client.jobs().get(inserted.job.row.id).await.unwrap(); + if !row.errors.is_empty() { + break row; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + assert_eq!(rescued.errors[0].error, "Stuck job rescued by JobRescuer"); + let effects: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pilot_effect WHERE operation = 'rescue' AND job_id = ?", + ) + .bind(inserted.job.row.id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(effects, 1); + assert!(pilot.rescue_calls.load(Ordering::SeqCst) >= 1); + + run.shutdown().await.unwrap(); +} + +#[tokio::test] +async fn sqlite_pilot_rescue_error_rolls_back_selection_side_effects() { + let pool = setup().await; + let mut pilot = SqlitePilot::new(); + pilot.rescue = Some(SelectionBehavior::Fail); + let client = Client::builder(pool.clone()) + .id("sqlite-pilot-rescue-error") + .maintenance( + MaintenanceConfig::default() + .with_elect_interval(Duration::from_millis(10)) + .with_rescue_after(Duration::from_mins(1)) + .with_rescuer_interval(Duration::from_millis(10)), + ) + .pilot(pilot.clone()) + .queue( + "default", + QueueConfig::new(1).with_fetch_poll_interval(Duration::from_mins(1)), + ) + .workers(runtime_workers(Arc::new(Semaphore::new(0)))) + .build() + .unwrap(); + let inserted = client.insert(RuntimeArgs { value: 1 }).await.unwrap(); + sqlx::query( + "UPDATE river_job SET state = 'running', attempt = 1, \ + attempted_at = '2000-01-01 00:00:00.000' WHERE id = ?", + ) + .bind(inserted.job.row.id) + .execute(&pool) + .await + .unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + tokio::time::timeout(Duration::from_secs(5), async { + while pilot.rescue_calls.load(Ordering::SeqCst) == 0 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + let effects: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pilot_effect WHERE operation = 'rescue' AND job_id = ?", + ) + .bind(inserted.job.row.id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(effects, 0); + assert_eq!( + client.jobs().get(inserted.job.row.id).await.unwrap().state, + JobState::Running + ); + + run.shutdown_now().await.unwrap(); +} + +#[tokio::test] +async fn sqlite_queue_events_are_emitted_once_per_transition() { + let pool = setup().await; + let client = Client::builder(pool) + .workers(runtime_workers(Arc::new(Semaphore::new(0)))) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let mut events = client + .subscribe(&[EventKind::QueuePaused, EventKind::QueueResumed]) + .unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + client.queues().pause("default").await.unwrap(); + let paused = tokio::time::timeout(Duration::from_secs(5), events.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(paused.kind(), EventKind::QueuePaused); + assert!(paused.as_queue().unwrap().queue.paused_at.is_some()); + assert!( + tokio::time::timeout(Duration::from_millis(300), events.recv()) + .await + .is_err() + ); + + client.queues().pause("default").await.unwrap(); + assert!( + tokio::time::timeout(Duration::from_millis(300), events.recv()) + .await + .is_err() + ); + + client.queues().resume("default").await.unwrap(); + let resumed = tokio::time::timeout(Duration::from_secs(5), events.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!(resumed.kind(), EventKind::QueueResumed); + assert!(resumed.as_queue().unwrap().queue.paused_at.is_none()); + assert!( + tokio::time::timeout(Duration::from_millis(300), events.recv()) + .await + .is_err() + ); + + client.queues().resume("default").await.unwrap(); + assert!( + tokio::time::timeout(Duration::from_millis(300), events.recv()) + .await + .is_err() + ); + + run.shutdown().await.unwrap(); +} + +#[tokio::test] +async fn sqlite_runs_jobs_and_persists_output() { + let pool = setup().await; + let worked = Arc::new(Semaphore::new(0)); + let worked_for_worker = Arc::clone(&worked); + let mut workers = WorkerRegistry::new(); + workers + .register_fn(move |context: WorkContext, job: Job| { + let worked = Arc::clone(&worked_for_worker); + async move { + context + .record_output(serde_json::json!({"doubled": job.args.value * 2})) + .unwrap(); + worked.add_permits(1); + Ok::<_, Infallible>(WorkOutcome::Complete) + } + }) + .unwrap(); + let client = Client::builder(pool.clone()) + .id("sqlite-runtime") + .workers(workers) + .queue( + "default", + QueueConfig::new(4) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + assert_eq!(client.database().kind(), DatabaseKind::Sqlite); + let mut events = client.subscribe(&[EventKind::JobCompleted]).unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + let inserted = client.insert(RuntimeArgs { value: 21 }).await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), worked.acquire()) + .await + .unwrap() + .unwrap() + .forget(); + let event = tokio::time::timeout(Duration::from_secs(5), events.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!( + event.as_job().map(|job_event| job_event.job.id), + Some(inserted.job.row.id) + ); + let row = client.jobs().get(inserted.job.row.id).await.unwrap(); + assert_eq!(row.state, JobState::Completed); + assert_eq!( + row.output().map(serde_json::value::RawValue::get), + Some(r#"{"doubled":42}"#) + ); + + run.shutdown().await.unwrap(); +} + +#[tokio::test] +async fn sqlite_worker_cancelling_its_own_token_fails_the_attempt_normally() { + #[derive(Debug, thiserror::Error)] + #[error("worker gave up")] + struct GaveUp; + + let pool = setup().await; + let mut workers = WorkerRegistry::new(); + workers + .register_fn(|context: WorkContext, _job: Job| async move { + // Like a worker that cancels its subtasks through a drop guard + // on its own token before returning. + drop(context.cancellation_token().clone().drop_guard()); + Err::(GaveUp) + }) + .unwrap(); + let client = Client::builder(pool.clone()) + .workers(workers) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let mut events = client + .subscribe(&[EventKind::JobCancelled, EventKind::JobFailed]) + .unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + let inserted = client.insert(CancelArgs {}).await.unwrap(); + let event = tokio::time::timeout(Duration::from_secs(5), events.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!( + event.as_job().map(|job_event| job_event.kind), + Some(riverqueue::JobEventKind::Failed) + ); + let row = client.jobs().get(inserted.job.row.id).await.unwrap(); + // The first retry is due soon enough that, like Go, the row goes + // straight back to `available`. + assert!( + matches!(row.state, JobState::Available | JobState::Retryable), + "{:?}", + row.state + ); + assert_eq!(row.errors.len(), 1); + assert_eq!(row.errors[0].error, "worker gave up"); + + run.shutdown().await.unwrap(); +} + +#[tokio::test] +async fn sqlite_attempt_errors_record_when_the_attempt_started() { + #[derive(Debug, thiserror::Error)] + #[error("worker failed late")] + struct FailedLate; + + let pool = setup().await; + let worker_started = Arc::new(Mutex::new(None)); + let worker_started_for_worker = Arc::clone(&worker_started); + let mut workers = WorkerRegistry::new(); + workers + .register_fn(move |_context: WorkContext, _job: Job| { + let worker_started = Arc::clone(&worker_started_for_worker); + async move { + *worker_started.lock().unwrap() = Some(chrono::Utc::now()); + tokio::time::sleep(Duration::from_millis(200)).await; + Err::(FailedLate) + } + }) + .unwrap(); + let client = Client::builder(pool.clone()) + .workers(workers) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let mut events = client.subscribe(&[EventKind::JobFailed]).unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + let inserted = client.insert(CancelArgs {}).await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), events.recv()) + .await + .unwrap() + .unwrap(); + let row = client.jobs().get(inserted.job.row.id).await.unwrap(); + let worker_started = worker_started.lock().unwrap().unwrap(); + // Like Go's executor, the error's time is when the attempt started, not + // when the worker returned. SQLite stores milliseconds. + let at = row.errors[0].at; + assert!( + at <= worker_started && at > worker_started - chrono::Duration::milliseconds(100), + "error at {at}, worker started at {worker_started}" + ); + + run.shutdown().await.unwrap(); +} + +#[tokio::test] +async fn sqlite_returning_insert_nonce_matches_go() { + let pool = setup().await; + let client = Client::builder(pool).build().unwrap(); + + let first = client.insert(RuntimeArgs { value: 1 }).await.unwrap(); + let second = client.insert(RuntimeArgs { value: 2 }).await.unwrap(); + let first_nonce = first + .job + .row + .metadata + .get::("river:unique_nonce") + .unwrap() + .unwrap(); + let second_nonce = second + .job + .row + .metadata + .get::("river:unique_nonce") + .unwrap() + .unwrap(); + assert_eq!(first_nonce.len(), 16); + assert!( + first_nonce + .bytes() + .all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase()) + ); + assert_ne!(first_nonce, second_nonce); + assert_eq!( + client + .jobs() + .get(first.job.row.id) + .await + .unwrap() + .metadata + .get::("river:unique_nonce") + .unwrap(), + Some(first_nonce) + ); +} + +#[tokio::test] +async fn sqlite_transaction_insert_respects_rollback() { + let pool = setup().await; + let client = Client::builder(pool.clone()).build().unwrap(); + let mut transaction = pool.begin().await.unwrap(); + let inserted = client + .insert(RuntimeArgs { value: 1 }) + .tx(&mut transaction) + .await + .unwrap(); + transaction.rollback().await.unwrap(); + + let error = client.jobs().get(inserted.job.row.id).await.unwrap_err(); + assert!(matches!(error, riverqueue::Error::NotFound(_))); +} + +#[tokio::test] +async fn sqlite_inserts_heterogeneous_batch_in_order() { + let pool = setup().await; + let client = Client::builder(pool).build().unwrap(); + let mut batch = InsertBatch::new(); + batch.push(RuntimeArgs { value: 7 }).push_with( + CancelArgs {}, + InsertOpts::default().with_queue("heterogeneous-queue"), + ); + + let results = client.insert_batch(batch).await.unwrap(); + + assert_eq!(results.len(), 2); + assert_eq!(results[0].job.kind, RuntimeArgs::KIND); + assert_eq!(results[1].job.kind, CancelArgs::KIND); + assert_eq!(results[1].job.queue, "heterogeneous-queue"); + assert!(results[0].job.id < results[1].job.id); +} + +#[tokio::test] +async fn sqlite_transaction_batches_roll_back_only_the_failed_batch() { + let pool = setup().await; + let client = Client::builder(pool.clone()).build().unwrap(); + + let mut transaction = pool.begin().await.unwrap(); + client + .insert(RuntimeArgs { value: 100 }) + .tx(&mut transaction) + .await + .unwrap(); + let result = client + .insert_many([ + ( + RuntimeArgs { value: 101 }, + InsertOpts::default().with_tags(["ordinary-failed-batch"]), + ), + ( + RuntimeArgs { value: 102 }, + InsertOpts::default().with_priority(0), + ), + ]) + .tx(&mut transaction) + .await; + assert!(result.is_err()); + transaction.commit().await.unwrap(); + + let control_count: i64 = sqlx::query_scalar( + "SELECT count(*) FROM river_job WHERE json_extract(args, '$.value') = 100", + ) + .fetch_one(&pool) + .await + .unwrap(); + let batch_count: i64 = sqlx::query_scalar( + "SELECT count(*) FROM river_job WHERE json_extract(args, '$.value') = 101", + ) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(control_count, 1); + assert_eq!(batch_count, 0); +} + +#[tokio::test] +async fn sqlite_long_fetch_cooldown_still_fetches_first() { + let pool = setup().await; + let worked = Arc::new(Semaphore::new(0)); + // Longer than the monotonic clock has run on any host. + let ten_years = Duration::from_hours(24 * 365 * 10); + let client = Client::builder(pool) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(ten_years) + .with_fetch_poll_interval(ten_years), + ) + .workers(runtime_workers(Arc::clone(&worked))) + .build() + .unwrap(); + client.insert(RuntimeArgs { value: 1 }).await.unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + tokio::time::timeout(Duration::from_secs(5), worked.acquire()) + .await + .expect("the first fetch doesn't wait for a cooldown") + .unwrap() + .forget(); + run.shutdown().await.unwrap(); +} + +#[tokio::test] +async fn sqlite_outbox_cancels_work_from_another_client() { + let pool = setup().await; + let started = Arc::new(Semaphore::new(0)); + let started_for_worker = Arc::clone(&started); + let mut workers = WorkerRegistry::new(); + workers + .register_fn(move |context: WorkContext, _job: Job| { + let started = Arc::clone(&started_for_worker); + async move { + started.add_permits(1); + context.cancellation_token().cancelled().await; + Ok::<_, Infallible>(WorkOutcome::Snooze(Duration::from_mins(1))) + } + }) + .unwrap(); + let worker_client = Client::builder(pool.clone()) + .id("sqlite-worker") + .workers(workers) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let cancelling_client = Client::builder(pool) + .id("sqlite-canceller") + .build() + .unwrap(); + let mut cancelled = worker_client.subscribe(&[EventKind::JobCancelled]).unwrap(); + let mut run = worker_client.start().unwrap(); + run.wait_ready().await.unwrap(); + + let inserted = worker_client.insert(CancelArgs {}).await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), started.acquire()) + .await + .unwrap() + .unwrap() + .forget(); + let running = worker_client.jobs().get(inserted.job.row.id).await.unwrap(); + assert_eq!(running.state, JobState::Running); + cancelling_client.jobs().cancel(running.id).await.unwrap(); + let event = tokio::time::timeout(Duration::from_secs(5), cancelled.recv()) + .await + .unwrap() + .unwrap(); + assert_eq!( + event.as_job().map(|job_event| job_event.job.id), + Some(running.id) + ); + assert_eq!( + worker_client.jobs().get(running.id).await.unwrap().state, + JobState::Cancelled + ); + + run.shutdown().await.unwrap(); +} + +#[tokio::test] +async fn sqlite_completion_event_follows_external_cancelled_state() { + let pool = setup().await; + let finish = Arc::new(Semaphore::new(0)); + let started = Arc::new(Semaphore::new(0)); + let mut workers = WorkerRegistry::new(); + workers + .register_fn({ + let finish = Arc::clone(&finish); + let started = Arc::clone(&started); + move |_context: WorkContext, _job: Job| { + let finish = Arc::clone(&finish); + let started = Arc::clone(&started); + async move { + started.add_permits(1); + finish.acquire().await.unwrap().forget(); + Ok::<_, Infallible>(WorkOutcome::Complete) + } + } + }) + .unwrap(); + let worker_client = Client::builder(pool.clone()) + .id("sqlite-worker-cancel-ignored") + .workers(workers) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let mut cancelled = worker_client.subscribe(&[EventKind::JobCancelled]).unwrap(); + let mut run = worker_client.start().unwrap(); + run.wait_ready().await.unwrap(); + + let inserted = worker_client.insert(CancelIgnoredArgs {}).await.unwrap(); + tokio::time::timeout(Duration::from_secs(5), started.acquire()) + .await + .unwrap() + .unwrap() + .forget(); + sqlx::query( + "UPDATE river_job SET state = 'cancelled', \ + finalized_at = strftime('%Y-%m-%d %H:%M:%f', 'now') WHERE id = ?", + ) + .bind(inserted.job.row.id) + .execute(&pool) + .await + .unwrap(); + finish.add_permits(1); + let event = tokio::time::timeout(Duration::from_secs(5), cancelled.recv()) + .await + .unwrap() + .unwrap(); + let event = event.as_job().unwrap(); + assert_eq!(event.job.id, inserted.job.row.id); + assert_eq!(event.job.state, JobState::Cancelled); + + run.shutdown().await.unwrap(); +} + +#[tokio::test] +async fn sqlite_fetches_and_discards_unregistered_kinds() { + let pool = setup().await; + let producer = Client::builder(pool.clone()).build().unwrap(); + let inserted = producer + .insert(UnknownArgs {}) + .opts(InsertOpts::default().with_max_attempts(1)) + .await + .unwrap(); + let mut workers = WorkerRegistry::new(); + workers + .register_fn(|_context: WorkContext, _job: Job| async move { + Ok::<_, Infallible>(WorkOutcome::Complete) + }) + .unwrap(); + let worker = Client::builder(pool) + .workers(workers) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .build() + .unwrap(); + let mut run = worker.start().unwrap(); + + tokio::time::timeout(Duration::from_secs(5), async { + loop { + if worker.jobs().get(inserted.job.row.id).await.unwrap().state == JobState::Discarded { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .unwrap(); + + run.shutdown().await.unwrap(); +} + +/// An extension that claims every available job itself and records the jobs +/// its set-state hook sees and the attempts that finished. +#[derive(Clone, Default)] +struct ClaimingPilot { + finished: Arc>>, + set_state_ids: Arc>>, + set_state_rows: Arc>>, +} + +#[async_trait] +impl Pilot for ClaimingPilot { + fn intercepts_job_set_state(&self) -> bool { + true + } + + async fn start_producer( + &self, + _context: ProducerStartContext, + ) -> Result>, PilotError> { + Ok(Some(Box::new(self.clone()))) + } + + async fn after_jobs_set_state( + &self, + _connection: DatabaseConnection<'_>, + params: &JobSetStateParams, + ) -> Result<(), PilotError> { + self.set_state_ids + .lock() + .unwrap() + .extend_from_slice(params.job_ids); + self.set_state_rows + .lock() + .unwrap() + .extend(params.jobs.iter().map(|job| job.id)); + Ok(()) + } +} + +#[async_trait] +impl PilotProducer for ClaimingPilot { + fn intercepts_claim(&self) -> bool { + true + } + + async fn claim( + &self, + context: ProducerClaimContext<'_>, + _next: ProducerClaimNext<'_>, + ) -> Result, PilotError> { + use sqlx::Row as _; + + let mut transaction = context.database.begin().await?; + let connection = transaction + .connection() + .into_sqlite() + .ok_or_else(|| std::io::Error::other("expected a SQLite connection"))?; + let ids: Vec = sqlx::query_scalar( + "SELECT id FROM river_job WHERE state = 'available' AND queue = ? \ + ORDER BY priority, scheduled_at, id LIMIT ?", + ) + .bind(context.queue) + .bind(i64::try_from(context.limit)?) + .fetch_all(&mut *connection) + .await?; + let mut claimed = Vec::new(); + for id in ids { + sqlx::query( + "UPDATE river_job SET state = 'running', attempt = attempt + 1, \ + attempted_at = ?, \ + attempted_by = jsonb_insert(coalesce(attempted_by, jsonb('[]')), '$[#]', ?) \ + WHERE id = ?", + ) + .bind(riverqueue::__private::sqlite_timestamp(chrono::Utc::now())) + .bind(context.client_id) + .bind(id) + .execute(&mut *connection) + .await?; + let row = sqlx::query(sqlx::AssertSqlSafe(format!( + "SELECT {} FROM river_job WHERE id = ?", + riverqueue::__private::SQLITE_JOB_COLUMNS + ))) + .bind(id) + .fetch_one(&mut *connection) + .await?; + assert_eq!(row.get::("id"), id); + claimed.push(riverqueue::__private::claimed_sqlite_job(&row)); + } + transaction.commit().await?; + Ok(claimed) + } + + fn job_finished(&self, job: &JobRow) { + self.finished.lock().unwrap().push(job.id); + } +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_sqlite_self_deleting")] +struct SelfDeletingArgs {} + +#[tokio::test(flavor = "multi_thread")] +async fn extension_claimed_rows_fail_undecodable_attempts_like_river_claims() { + let (pool, database_path) = setup_file_pool(Duration::from_secs(5)).await; + let pilot = ClaimingPilot::default(); + let client = Client::builder(pool.clone()) + .id("sqlite-extension-claimer") + .pilot(pilot.clone()) + .queue( + "default", + QueueConfig::new(2) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .workers(runtime_workers(Arc::new(Semaphore::new(0)))) + .build() + .unwrap(); + let good = client.insert(RuntimeArgs { value: 1 }).await.unwrap(); + let bad = client + .insert(RuntimeArgs { value: 2 }) + .opts(InsertOpts::default().with_max_attempts(1)) + .await + .unwrap(); + sqlx::query("UPDATE river_job SET tags = jsonb('{\"not\":\"an array\"}') WHERE id = ?") + .bind(bad.id()) + .execute(&pool) + .await + .unwrap(); + + let mut failed = client.subscribe(&[EventKind::JobFailed]).unwrap(); + let mut run = client.start().unwrap(); + let event = tokio::time::timeout(Duration::from_secs(10), failed.recv()) + .await + .expect("undecodable claimed job did not fail") + .unwrap(); + let event = event.as_job().unwrap(); + assert_eq!(event.job.id, bad.id()); + let (state, errors): (String, String) = + sqlx::query_as("SELECT state, json(errors) FROM river_job WHERE id = ?") + .bind(bad.id()) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(state, "discarded"); + assert!( + errors.contains("job row couldn't be decoded: "), + "unexpected attempt errors {errors}" + ); + tokio::time::timeout(Duration::from_secs(10), async { + loop { + let state: String = sqlx::query_scalar("SELECT state FROM river_job WHERE id = ?") + .bind(good.id()) + .fetch_one(&pool) + .await + .unwrap(); + if state == "completed" { + break; + } + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("decodable claimed job did not complete"); + run.shutdown().await.unwrap(); + // Both attempts, including the undecodable row's, finished once. + let mut finished = pilot.finished.lock().unwrap().clone(); + finished.sort_unstable(); + let mut expected = vec![good.id(), bad.id()]; + expected.sort_unstable(); + assert_eq!(finished, expected); + pool.close().await; + remove_sqlite_files(&database_path); +} + +#[tokio::test(flavor = "multi_thread")] +async fn set_state_extension_sees_jobs_deleted_while_worked() { + let (pool, database_path) = setup_file_pool(Duration::from_secs(5)).await; + let pilot = ClaimingPilot::default(); + let mut workers = WorkerRegistry::new(); + let delete_pool = pool.clone(); + workers + .register_fn(move |_context: WorkContext, job: Job| { + let pool = delete_pool.clone(); + async move { + sqlx::query("DELETE FROM river_job WHERE id = ?") + .bind(job.row.id) + .execute(&pool) + .await?; + Ok::<_, sqlx::Error>(WorkOutcome::Complete) + } + }) + .unwrap(); + let client = Client::builder(pool.clone()) + .id("sqlite-set-state-deleted") + .pilot(pilot.clone()) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .workers(workers) + .build() + .unwrap(); + let deleted = client.insert(SelfDeletingArgs {}).await.unwrap(); + + let mut run = client.start().unwrap(); + tokio::time::timeout(Duration::from_secs(10), async { + while !pilot.set_state_ids.lock().unwrap().contains(&deleted.id()) { + tokio::time::sleep(Duration::from_millis(10)).await; + } + }) + .await + .expect("set-state extension never saw the deleted job"); + run.shutdown().await.unwrap(); + + // The job's row is gone, so only its ID reaches the extension. + assert!(!pilot.set_state_rows.lock().unwrap().contains(&deleted.id())); + pool.close().await; + remove_sqlite_files(&database_path); +} diff --git a/rust/riverqueue/tests/sqlite_storage.rs b/rust/riverqueue/tests/sqlite_storage.rs new file mode 100644 index 000000000..9fa610fbd --- /dev/null +++ b/rust/riverqueue/tests/sqlite_storage.rs @@ -0,0 +1,748 @@ +use chrono::{DateTime, Duration, SubsecRound, Utc}; +use riverqueue::{ + Client, Error, InsertBatch, InsertOpts, JobArgs, JobDeleteManyParams, JobListCursor, + JobListOrderBy, JobListParams, JobState, JobUpdateParams, QueueListParams, QueueSelector, + QueueUpdateParams, SortDirection, UniqueOpts, +}; +use riverqueue_migrate::SqliteMigrator; +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value, json}; +use sqlx::{Row, SqlitePool, sqlite::SqlitePoolOptions}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "sqlite_empty_batch")] +struct EmptyBatchArgs { + value: i32, +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "sqlite_unencodable")] +struct UnencodableArgs { + // JSON object keys must be strings. + pairs: std::collections::BTreeMap<(i32, i32), i32>, +} + +#[tokio::test] +async fn batch_encoding_errors_surface_at_insert_without_inserting() { + let (client, pool) = setup().await; + + let mut batch = InsertBatch::new(); + batch + .push(EmptyBatchArgs { value: 1 }) + .push(UnencodableArgs { + pairs: std::collections::BTreeMap::from([((1, 2), 3)]), + }); + let error = client.insert_batch(batch).await.unwrap_err(); + + assert!(matches!(error, Error::Json(_)), "{error:?}"); + let count: i64 = sqlx::query_scalar("SELECT count(*) FROM river_job") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(count, 0); + + pool.close().await; +} + +#[tokio::test] +async fn control_and_resign_notifications_carry_their_fields() { + let (client, pool) = setup().await; + insert_queue(&pool, "alpha").await; + + client + .queues() + .update( + "alpha", + QueueUpdateParams::new().metadata(Map::from_iter([("note".to_owned(), json!("<&>"))])), + ) + .await + .unwrap(); + client.queues().pause("alpha").await.unwrap(); + client.request_resign().await.unwrap(); + + let notifications: Vec<(String, String)> = + sqlx::query_as("SELECT topic, payload FROM river_notification ORDER BY id") + .fetch_all(&pool) + .await + .unwrap(); + let notifications = notifications + .into_iter() + .map(|(topic, payload)| (topic, serde_json::from_str::(&payload).unwrap())) + .collect::>(); + assert_eq!( + notifications, + [ + ( + "river_control".to_owned(), + json!({"action": "metadata_changed", "metadata": {"note": "<&>"}, "queue": "alpha"}), + ), + ( + "river_control".to_owned(), + json!({"action": "pause", "queue": "alpha"}), + ), + ( + "river_leadership".to_owned(), + json!({"action": "request_resign", "leader_id": ""}), + ), + ] + ); + + pool.close().await; +} + +#[tokio::test] +async fn empty_batches_are_rejected_before_database_work() { + let (client, pool) = setup().await; + + let empty_many = client + .insert_many(Vec::::new()) + .await + .unwrap_err(); + assert_eq!(empty_many.to_string(), "invalid job: no jobs to insert"); + let empty_batch = client.insert_batch(InsertBatch::new()).await.unwrap_err(); + assert_eq!(empty_batch.to_string(), "invalid job: no jobs to insert"); + + let mut transaction = pool.begin().await.unwrap(); + let empty_many_tx = client + .insert_many(Vec::::new()) + .tx(&mut transaction) + .await + .unwrap_err(); + assert_eq!(empty_many_tx.to_string(), "invalid job: no jobs to insert"); + let empty_batch_tx = client + .insert_batch(InsertBatch::new()) + .tx(&mut transaction) + .await + .unwrap_err(); + assert_eq!(empty_batch_tx.to_string(), "invalid job: no jobs to insert"); + transaction.commit().await.unwrap(); + + pool.close().await; +} + +/// Like River Go, a batch may not repeat a unique key among jobs whose +/// state it covers, and rejecting it writes nothing. Uniqueness without the +/// kind needs another dimension. +#[tokio::test] +async fn batches_repeating_a_unique_key_fail_without_inserting() { + let (client, pool) = setup().await; + let unique = InsertOpts::default().with_unique(UniqueOpts::new().with_by_args(true)); + + let error = client + .insert_many([ + (EmptyBatchArgs { value: 1 }, unique.clone()), + (EmptyBatchArgs { value: 2 }, unique.clone()), + (EmptyBatchArgs { value: 1 }, unique.clone()), + ]) + .await + .unwrap_err(); + assert_eq!( + error.to_string(), + "invalid job: unique key appears more than once in batch" + ); + let count: i64 = sqlx::query_scalar("SELECT count(*) FROM river_job") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(count, 0); + + let inserted = client + .insert_many([ + (EmptyBatchArgs { value: 1 }, unique.clone()), + (EmptyBatchArgs { value: 2 }, unique), + ]) + .await + .unwrap(); + assert!( + inserted + .iter() + .all(|result| !result.unique_skipped_as_duplicate) + ); + + let error = client + .insert(EmptyBatchArgs { value: 3 }) + .opts(InsertOpts::default().with_unique(UniqueOpts::new().with_exclude_kind(true))) + .await + .unwrap_err(); + assert!(matches!(error, Error::InvalidJob(_)), "{error:?}"); + + pool.close().await; +} + +/// Like Go's SQLite driver, an insert binds no time it wasn't given: SQLite +/// stamps `created_at` on every row, and `scheduled_at` on unscheduled rows, +/// with its own `datetime('now', 'subsec')`, while an explicit schedule is +/// stored as given. +#[tokio::test] +async fn inserts_leave_unset_times_to_sqlite() { + let (client, pool) = setup().await; + let sqlite_now = || async { + sqlx::query_scalar::<_, String>("SELECT datetime('now', 'subsec')") + .fetch_one(&pool) + .await + .unwrap() + }; + let stored = |id: i64| { + let pool = pool.clone(); + async move { + sqlx::query_as::<_, (String, String)>( + "SELECT created_at, scheduled_at FROM river_job WHERE id = ?", + ) + .bind(id) + .fetch_one(&pool) + .await + .unwrap() + } + }; + let is_sqlite_time = |text: &str| { + text.len() == 23 + && chrono::NaiveDateTime::parse_from_str(text, "%Y-%m-%d %H:%M:%S%.3f").is_ok() + }; + + let before = sqlite_now().await; + let unscheduled = client + .insert(EmptyBatchArgs { value: 1 }) + .await + .unwrap() + .id(); + let mut batch = InsertBatch::new(); + batch + .push(EmptyBatchArgs { value: 2 }) + .push(EmptyBatchArgs { value: 3 }); + let batched = client + .insert_batch(batch) + .await + .unwrap() + .iter() + .map(riverqueue::InsertBatchResult::id) + .collect::>(); + let scheduled_for = Utc::now() + Duration::hours(1); + let scheduled = client + .insert(EmptyBatchArgs { value: 4 }) + .opts(riverqueue::InsertOpts::default().with_scheduled_at(scheduled_for)) + .await + .unwrap() + .id(); + let after = sqlite_now().await; + + for id in std::iter::once(unscheduled).chain(batched) { + let (created_at, scheduled_at) = stored(id).await; + assert!(is_sqlite_time(&created_at), "{created_at}"); + // SQLite's `now` holds for one statement, so both columns agree. + assert_eq!(created_at, scheduled_at); + assert!( + before <= created_at && created_at <= after, + "{before} <= {created_at} <= {after}" + ); + } + let (created_at, scheduled_at) = stored(scheduled).await; + assert!(is_sqlite_time(&created_at), "{created_at}"); + assert!(before <= created_at && created_at <= after); + assert_eq!(scheduled_at, sqlite_time(scheduled_for)); + + pool.close().await; +} + +#[tokio::test] +async fn job_list_time_without_states_uses_schedule_and_finalized_requires_states() { + let (client, pool) = setup().await; + let now = Utc::now(); + let first = insert_job( + &pool, + JobSeed { + scheduled_at: now + Duration::hours(1), + ..JobSeed::default() + }, + ) + .await; + let second = insert_job( + &pool, + JobSeed { + scheduled_at: now - Duration::hours(1), + ..JobSeed::default() + }, + ) + .await; + + let params = JobListParams::default() + .ids([first, second]) + .order_by(JobListOrderBy::Time); + let rows = client.jobs().list(params.clone()).await.unwrap().jobs; + assert_eq!( + rows.iter().map(|row| row.id).collect::>(), + [second, first] + ); + let cursor = JobListCursor::from_job(&rows[0], ¶ms).unwrap(); + let page = client + .jobs() + .list(params.clone().after(cursor)) + .await + .unwrap() + .jobs; + assert_eq!(page.iter().map(|row| row.id).collect::>(), [first]); + + let error = client + .jobs() + .list(JobListParams::default().order_by(JobListOrderBy::FinalizedAt)) + .await + .unwrap_err(); + assert!(matches!(error, Error::InvalidJob(_))); +} + +#[tokio::test] +#[allow( + clippy::too_many_lines, + reason = "one lifecycle test keeps ordered CRUD state transitions and rollback assertions together" +)] +async fn job_crud_preserves_sqlite_semantics() { + let (client, pool) = setup().await; + let now = Utc::now(); + let delete_id = insert_job( + &pool, + JobSeed { + metadata: json!({"original": true}), + scheduled_at: now - Duration::minutes(3), + tags: json!(["delete-me", "shared"]), + ..JobSeed::default() + }, + ) + .await; + let retry_id = insert_job( + &pool, + JobSeed { + attempt: 3, + max_attempts: 3, + scheduled_at: now - Duration::minutes(2), + state: JobState::Retryable, + tags: json!(["retry-me", "shared"]), + ..JobSeed::default() + }, + ) + .await; + let running_id = insert_job( + &pool, + JobSeed { + scheduled_at: now - Duration::minutes(1), + state: JobState::Running, + tags: json!(["running", "shared"]), + ..JobSeed::default() + }, + ) + .await; + + let running = client.jobs().get(running_id).await.unwrap(); + assert_eq!(running.state, JobState::Running); + assert_eq!(running.tags, ["running", "shared"]); + + let list_params = JobListParams::default() + .ids([delete_id, retry_id, running_id]) + .limit(2) + .order_by(JobListOrderBy::ScheduledAt) + .direction(SortDirection::Descending); + let first_page = client.jobs().list(list_params.clone()).await.unwrap(); + assert_eq!( + first_page.jobs.iter().map(|job| job.id).collect::>(), + [running_id, retry_id] + ); + let last_cursor = first_page.last_cursor.unwrap(); + assert_eq!( + last_cursor.encode(), + JobListCursor::from_job(&first_page.jobs[1], &list_params) + .unwrap() + .encode() + ); + let second_page = client + .jobs() + .list(list_params.after(last_cursor)) + .await + .unwrap() + .jobs; + assert_eq!( + second_page.iter().map(|job| job.id).collect::>(), + [delete_id] + ); + + let updated = client + .jobs() + .update( + delete_id, + JobUpdateParams::default() + .metadata(Map::from_iter([("added".to_owned(), json!(42))])) + .output(json!({"ok": true})), + ) + .await + .unwrap(); + assert_eq!( + updated.metadata.get::("original").unwrap(), + Some(true) + ); + assert_eq!(updated.metadata.get::("added").unwrap(), Some(42)); + assert_eq!( + updated.output().map(serde_json::value::RawValue::get), + Some(r#"{"ok":true}"#) + ); + + let mut transaction = pool.begin().await.unwrap(); + let transaction_update = client + .jobs() + .update( + delete_id, + JobUpdateParams::default() + .metadata(Map::from_iter([("rolled_back".to_owned(), json!(true))])), + ) + .tx(&mut transaction) + .await + .unwrap(); + assert_eq!( + transaction_update + .metadata + .get::("rolled_back") + .unwrap(), + Some(true) + ); + assert_eq!( + client + .jobs() + .get(delete_id) + .tx(&mut transaction) + .await + .unwrap() + .metadata + .get::("rolled_back") + .unwrap(), + Some(true) + ); + transaction.rollback().await.unwrap(); + assert!( + !client + .jobs() + .get(delete_id) + .await + .unwrap() + .metadata + .contains_key("rolled_back") + ); + + let mut transaction = pool.begin().await.unwrap(); + let completed = client + .jobs() + .complete(running_id) + .tx(&mut transaction) + .await + .unwrap(); + assert_eq!(completed.state, JobState::Completed); + transaction.rollback().await.unwrap(); + assert_eq!( + client.jobs().get(running_id).await.unwrap().state, + JobState::Running + ); + let mut transaction = pool.begin().await.unwrap(); + client + .jobs() + .complete(running_id) + .tx(&mut transaction) + .await + .unwrap(); + transaction.commit().await.unwrap(); + let completed = client.jobs().get(running_id).await.unwrap(); + assert_eq!(completed.state, JobState::Completed); + assert!(completed.finalized_at.is_some()); + + let retried = client.jobs().retry(retry_id).await.unwrap(); + assert_eq!(retried.state, JobState::Available); + assert_eq!(retried.max_attempts, 4); + // Like Go, a retry sends no insert notification. + let insert_notifications: i64 = + sqlx::query_scalar("SELECT count(*) FROM river_notification WHERE topic = 'river_insert'") + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(insert_notifications, 0); + + let deleted = client.jobs().delete(delete_id).await.unwrap(); + assert_eq!(deleted.id, delete_id); + assert!(matches!( + client.jobs().get(delete_id).await, + Err(Error::NotFound(_)) + )); + + let still_running_id = insert_job( + &pool, + JobSeed { + state: JobState::Running, + ..JobSeed::default() + }, + ) + .await; + assert!(matches!( + client.jobs().delete(still_running_id).await, + Err(Error::JobRunning) + )); + + pool.close().await; +} + +#[tokio::test] +async fn job_delete_many_is_atomic_and_skips_running_jobs() { + let (client, pool) = setup().await; + let first = insert_job(&pool, JobSeed::default()).await; + let second = insert_job( + &pool, + JobSeed { + state: JobState::Pending, + ..JobSeed::default() + }, + ) + .await; + let running = insert_job( + &pool, + JobSeed { + state: JobState::Running, + ..JobSeed::default() + }, + ) + .await; + + assert!(matches!( + client + .jobs() + .delete_many(JobDeleteManyParams::matching(JobListParams::default())) + .await, + Err(Error::InvalidJob(_)) + )); + let params = + JobDeleteManyParams::matching(JobListParams::default().ids([first, second, running])); + let mut transaction = pool.begin().await.unwrap(); + let rolled_back = client + .jobs() + .delete_many(params.clone()) + .tx(&mut transaction) + .await + .unwrap(); + assert_eq!( + rolled_back.iter().map(|job| job.id).collect::>(), + [first, second] + ); + transaction.rollback().await.unwrap(); + assert!(client.jobs().get(first).await.is_ok()); + assert!(client.jobs().get(second).await.is_ok()); + + let deleted = client.jobs().delete_many(params).await.unwrap(); + assert_eq!( + deleted.iter().map(|job| job.id).collect::>(), + [first, second] + ); + assert_eq!( + client.jobs().get(running).await.unwrap().state, + JobState::Running + ); + + pool.close().await; +} + +#[tokio::test] +async fn queue_crud_and_notifications_share_the_caller_transaction() { + let (client, pool) = setup().await; + insert_queue(&pool, "alpha").await; + insert_queue(&pool, "beta").await; + + assert_eq!(client.queues().get("alpha").await.unwrap().name, "alpha"); + assert_eq!( + client + .queues() + .list(QueueListParams::default()) + .await + .unwrap() + .iter() + .map(|queue| queue.name.as_str()) + .collect::>(), + ["alpha", "beta"] + ); + + client.queues().pause("alpha").await.unwrap(); + assert!( + client + .queues() + .get("alpha") + .await + .unwrap() + .paused_at + .is_some() + ); + let initial_notification_count = notification_count(&pool).await; + + let mut transaction = pool.begin().await.unwrap(); + client + .queues() + .resume("alpha") + .tx(&mut transaction) + .await + .unwrap(); + assert!( + client + .queues() + .get("alpha") + .tx(&mut transaction) + .await + .unwrap() + .paused_at + .is_none() + ); + transaction.rollback().await.unwrap(); + assert!( + client + .queues() + .get("alpha") + .await + .unwrap() + .paused_at + .is_some() + ); + assert_eq!(notification_count(&pool).await, initial_notification_count); + + let mut transaction = pool.begin().await.unwrap(); + let updated = client + .queues() + .update( + "alpha", + QueueUpdateParams::new() + .metadata(Map::from_iter([("owner".to_owned(), json!("rust"))])), + ) + .tx(&mut transaction) + .await + .unwrap(); + assert_eq!(updated.metadata["owner"], "rust"); + transaction.commit().await.unwrap(); + assert_eq!( + client.queues().get("alpha").await.unwrap().metadata["owner"], + "rust" + ); + assert_eq!( + notification_count(&pool).await, + initial_notification_count + 1 + ); + + client.queues().resume("alpha").await.unwrap(); + client.queues().pause(QueueSelector::All).await.unwrap(); + assert!( + client + .queues() + .list(QueueListParams::default()) + .await + .unwrap() + .iter() + .all(|queue| queue.paused_at.is_some()) + ); + + let payload: String = sqlx::query_scalar( + "SELECT payload FROM river_notification WHERE topic = 'river_control' ORDER BY id DESC LIMIT 1", + ) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!( + serde_json::from_str::(&payload).unwrap(), + json!({"action": "pause", "queue": "*"}) + ); + + pool.close().await; +} + +#[derive(Clone)] +struct JobSeed { + attempt: i16, + kind: &'static str, + max_attempts: i16, + metadata: Value, + queue: &'static str, + scheduled_at: DateTime, + state: JobState, + tags: Value, +} + +impl Default for JobSeed { + fn default() -> Self { + Self { + attempt: 0, + kind: "sqlite_storage_test", + max_attempts: 25, + metadata: json!({}), + queue: "default", + scheduled_at: Utc::now(), + state: JobState::Available, + tags: json!([]), + } + } +} + +async fn insert_job(pool: &SqlitePool, seed: JobSeed) -> i64 { + let now = Utc::now(); + let finalized_at = matches!( + seed.state, + JobState::Cancelled | JobState::Completed | JobState::Discarded + ) + .then_some(now); + sqlx::query_scalar( + "INSERT INTO river_job (args, attempt, attempted_at, attempted_by, created_at, errors, \ + finalized_at, kind, max_attempts, metadata, priority, queue, scheduled_at, state, tags) \ + VALUES (jsonb(?), ?, ?, jsonb(?), ?, jsonb(?), ?, ?, ?, jsonb(?), 1, ?, ?, ?, jsonb(?)) \ + RETURNING id", + ) + .bind(r#"{"message":"hello"}"#) + .bind(seed.attempt) + .bind( + (seed.state == JobState::Running) + .then_some(now) + .map(sqlite_time), + ) + .bind("[]") + .bind(sqlite_time(now)) + .bind("[]") + .bind(finalized_at.map(sqlite_time)) + .bind(seed.kind) + .bind(seed.max_attempts) + .bind(seed.metadata.to_string()) + .bind(seed.queue) + .bind(sqlite_time(seed.scheduled_at)) + .bind(seed.state.as_str()) + .bind(seed.tags.to_string()) + .fetch_one(pool) + .await + .unwrap() +} + +fn sqlite_time(time: DateTime) -> String { + time.round_subsecs(3) + .format("%Y-%m-%d %H:%M:%S%.3f") + .to_string() +} + +async fn insert_queue(pool: &SqlitePool, name: &str) { + sqlx::query("INSERT INTO river_queue (name, metadata) VALUES (?, jsonb('{}'))") + .bind(name) + .execute(pool) + .await + .unwrap(); +} + +async fn notification_count(pool: &SqlitePool) -> i64 { + sqlx::query("SELECT count(*) AS count FROM river_notification") + .fetch_one(pool) + .await + .unwrap() + .get("count") +} + +async fn setup() -> (Client, SqlitePool) { + let pool = SqlitePoolOptions::new() + .max_connections(1) + .connect("sqlite::memory:") + .await + .unwrap(); + SqliteMigrator::new(pool.clone()) + .migrate_up() + .await + .unwrap(); + let client = Client::builder(pool.clone()).build().unwrap(); + (client, pool) +} diff --git a/rust/riverqueue/tests/stuck_jobs.rs b/rust/riverqueue/tests/stuck_jobs.rs new file mode 100644 index 000000000..906efb88a --- /dev/null +++ b/rust/riverqueue/tests/stuck_jobs.rs @@ -0,0 +1,141 @@ +//! The stuck job log line reports the timeout that applied to the job, the +//! worker's own when it sets one, like River Go's executor. + +#![cfg(feature = "sqlite")] + +mod support; + +use std::{ + convert::Infallible, + sync::{Arc, Mutex}, + time::Duration, +}; + +use riverqueue::{ + Client, Job, JobArgs, QueueConfig, WorkContext, WorkOutcome, Worker, WorkerRegistry, + WorkerTimeout, +}; +use serde::{Deserialize, Serialize}; +use tokio::sync::Notify; +use tracing::field::{Field, Visit}; +use tracing_subscriber::{ + Layer, + layer::{Context, SubscriberExt}, +}; + +use crate::support::{sqlite_cleanup, sqlite_file_pool}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "stuck_job")] +struct StuckArgs {} + +/// Ignores cancellation, with a timeout of its own. +struct StuckWorker { + started: Arc, +} + +impl Worker for StuckWorker { + type Error = Infallible; + + fn timeout(&self, _job: &Job) -> WorkerTimeout { + WorkerTimeout::After(Duration::from_millis(5)) + } + + async fn work( + &self, + _context: WorkContext, + _job: Job, + ) -> Result { + self.started.notify_one(); + std::future::pending().await + } +} + +/// The `timeout` field of each stuck job log line. +#[derive(Clone, Default)] +struct StuckLines { + changed: Arc, + timeouts: Arc>>, +} + +impl Layer for StuckLines { + fn on_event(&self, event: &tracing::Event<'_>, _context: Context<'_, S>) { + #[derive(Default)] + struct Fields { + message: String, + timeout: Option, + } + + impl Visit for Fields { + fn record_debug(&mut self, field: &Field, value: &dyn std::fmt::Debug) { + match field.name() { + "message" => self.message = format!("{value:?}"), + "timeout" => self.timeout = Some(format!("{value:?}")), + _ => {} + } + } + } + + let mut fields = Fields::default(); + event.record(&mut fields); + if fields.message.contains("treating it as stuck") { + self.timeouts + .lock() + .unwrap() + .push(fields.timeout.unwrap_or_default()); + self.changed.notify_waiters(); + } + } +} + +// Current-thread runtime: the subscriber set for this thread sees every +// task the client spawns. +#[tokio::test] +async fn stuck_log_line_reports_the_worker_timeout() { + let lines = StuckLines::default(); + let _subscriber = + tracing::subscriber::set_default(tracing_subscriber::registry().with(lines.clone())); + let (pool, path) = sqlite_file_pool(4).await; + let started = Arc::new(Notify::new()); + let mut workers = WorkerRegistry::new(); + workers + .register(StuckWorker { + started: Arc::clone(&started), + }) + .unwrap(); + // A client timeout long enough that the job can only have been + // cancelled by the worker's timeout. + let client = Client::builder(pool.clone()) + .job_timeout(Duration::from_mins(1)) + .job_stuck_threshold(Duration::from_millis(10)) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .workers(workers) + .build() + .unwrap(); + client.insert(StuckArgs {}).await.unwrap(); + let mut run = client.start().unwrap(); + tokio::time::timeout(Duration::from_secs(10), started.notified()) + .await + .expect("worker starts"); + + tokio::time::timeout(Duration::from_secs(10), async { + loop { + let changed = lines.changed.notified(); + if !lines.timeouts.lock().unwrap().is_empty() { + return; + } + changed.await; + } + }) + .await + .expect("stuck job logged"); + run.shutdown().await.unwrap(); + + assert_eq!(*lines.timeouts.lock().unwrap(), ["Some(5ms)"]); + sqlite_cleanup(pool, path).await; +} diff --git a/rust/riverqueue/tests/support/mod.rs b/rust/riverqueue/tests/support/mod.rs new file mode 100644 index 000000000..4355769c3 --- /dev/null +++ b/rust/riverqueue/tests/support/mod.rs @@ -0,0 +1,167 @@ +//! Shared fixtures for maintenance, leadership, and storage parity tests. +//! +//! PostgreSQL tests run in a freshly migrated schema with a unique name so that +//! concurrent test binaries sharing one disposable database never clobber each +//! other. They fail rather than skip when `RIVER_RUST_DATABASE_URL` is unset. + +#![allow(dead_code, reason = "each test binary uses a different subset")] + +use std::sync::atomic::{AtomicUsize, Ordering}; + +#[cfg(feature = "postgres")] +use riverqueue::database::SchemaName; +#[cfg(feature = "postgres")] +use riverqueue_migrate::PostgresMigrator; +#[cfg(feature = "sqlite")] +use riverqueue_migrate::SqliteMigrator; +#[cfg(feature = "postgres")] +use sqlx::{ + AssertSqlSafe, PgPool, + postgres::{PgConnectOptions, PgPoolOptions}, +}; +#[cfg(feature = "sqlite")] +use sqlx::{ + SqlitePool, + sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}, +}; + +static NONCE: AtomicUsize = AtomicUsize::new(0); + +/// Returns a process-unique suffix for schema and file names. +pub fn unique_suffix() -> String { + let nanos = std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .expect("system clock is after the Unix epoch") + .subsec_nanos(); + format!( + "{:x}_{:x}_{:x}", + std::process::id(), + NONCE.fetch_add(1, Ordering::Relaxed), + nanos + ) +} + +/// A uniquely named PostgreSQL schema owned by one test, migrated unless +/// created by an `unmigrated` constructor. +#[cfg(feature = "postgres")] +pub struct PostgresSchema { + pub pool: PgPool, + pub schema: SchemaName, + name: String, +} + +#[cfg(feature = "postgres")] +impl PostgresSchema { + /// Creates and migrates a uniquely named schema. + /// + /// # Panics + /// + /// Panics when `RIVER_RUST_DATABASE_URL` is unset so an explicitly + /// selected database test can never pass vacuously. + pub async fn new(prefix: &str) -> Self { + Self::create(prefix, false, true).await + } + + /// Creates and migrates a uniquely named schema and makes it the current + /// schema of every pool connection through `search_path`. + /// + /// Use it for tests that exercise unqualified access through the + /// connection's current schema rather than an explicit schema, without + /// touching `public`. + pub async fn current(prefix: &str) -> Self { + Self::create(prefix, true, true).await + } + + /// Like [`PostgresSchema::current`], but leaves the schema empty. + pub async fn current_unmigrated(prefix: &str) -> Self { + Self::create(prefix, true, false).await + } + + /// Creates a uniquely named schema without migrating it. + pub async fn unmigrated(prefix: &str) -> Self { + Self::create(prefix, false, false).await + } + + async fn create(prefix: &str, search_path: bool, migrate: bool) -> Self { + let url = std::env::var("RIVER_RUST_DATABASE_URL") + .expect("RIVER_RUST_DATABASE_URL must point at a disposable test database"); + let mut name = format!("{prefix}_{}", unique_suffix()); + name.truncate(riverqueue::migrate::SCHEMA_MAX_LEN); + let mut options: PgConnectOptions = url + .parse() + .expect("parse RIVER_RUST_DATABASE_URL as PostgreSQL connect options"); + if search_path { + // Test schema names are lowercase identifiers, so no quoting is + // needed in the startup parameter. + options = options.options([("search_path", name.as_str())]); + } + let pool = PgPoolOptions::new() + .max_connections(16) + .connect_with(options) + .await + .expect("connect to RIVER_RUST_DATABASE_URL"); + sqlx::raw_sql(AssertSqlSafe(format!("CREATE SCHEMA \"{name}\""))) + .execute(&pool) + .await + .expect("create test schema"); + let schema = SchemaName::new(name.clone()).expect("valid test schema name"); + if migrate { + PostgresMigrator::new(pool.clone()) + .with_schema(schema.clone()) + .migrate_up() + .await + .expect("migrate test schema"); + } + Self { pool, schema, name } + } + + /// Qualifies a River table in this schema. + pub fn table(&self, table: &str) -> String { + self.schema.qualify(table) + } + + /// Drops the schema and closes the pool. + pub async fn cleanup(self) { + sqlx::raw_sql(AssertSqlSafe(format!( + "DROP SCHEMA \"{}\" CASCADE", + self.name + ))) + .execute(&self.pool) + .await + .expect("drop test schema"); + self.pool.close().await; + } +} + +/// Opens a migrated SQLite database in a unique temporary file. +#[cfg(feature = "sqlite")] +pub async fn sqlite_file_pool(max_connections: u32) -> (SqlitePool, std::path::PathBuf) { + let path = std::env::temp_dir().join(format!("river-maint-{}.sqlite", unique_suffix())); + let options = SqliteConnectOptions::new() + .filename(&path) + .create_if_missing(true) + .journal_mode(SqliteJournalMode::Wal) + .busy_timeout(std::time::Duration::from_secs(5)); + let pool = SqlitePoolOptions::new() + .max_connections(max_connections) + .connect_with(options) + .await + .expect("open SQLite test database"); + SqliteMigrator::new(pool.clone()) + .migrate_up() + .await + .expect("migrate SQLite test database"); + (pool, path) +} + +/// Closes a SQLite pool and removes its database files. +#[cfg(feature = "sqlite")] +pub async fn sqlite_cleanup(pool: SqlitePool, path: std::path::PathBuf) { + pool.close().await; + let _ = std::fs::remove_file(&path); + for suffix in ["-shm", "-wal"] { + let mut sidecar = path.as_os_str().to_owned(); + sidecar.push(suffix); + let _ = std::fs::remove_file(sidecar); + } +} diff --git a/rust/riverqueue/tests/work_extensions.rs b/rust/riverqueue/tests/work_extensions.rs new file mode 100644 index 000000000..1b9af082e --- /dev/null +++ b/rust/riverqueue/tests/work_extensions.rs @@ -0,0 +1,380 @@ +//! Work middleware, work hooks, and error handlers: ordering, result +//! replacement, handler panics, unknown kinds, and the job span around +//! worker code. +//! +//! Extension behavior doesn't depend on the backend, so these tests use +//! temporary SQLite databases and need no external services. + +use std::{ + path::PathBuf, + sync::{ + Arc, Mutex, OnceLock, + atomic::{AtomicUsize, Ordering}, + }, + time::Duration, +}; + +use riverqueue::{ + Client, ErrorHandler, ErrorHandlerDecision, EventKind, Hook, InsertOpts, Job, JobArgs, JobRow, + JobState, Metric, QueueConfig, WorkContext, WorkError, WorkMiddleware, WorkNext, WorkOutcome, + WorkerRegistry, +}; +use riverqueue_migrate::SqliteMigrator; +use serde::{Deserialize, Serialize}; +use sqlx::{ + SqlitePool, + sqlite::{SqliteConnectOptions, SqliteJournalMode, SqlitePoolOptions}, +}; + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_work_extensions")] +struct ExtensionArgs { + fail: bool, +} + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_work_extensions_unknown")] +struct UnknownArgs {} + +#[derive(Debug, thiserror::Error)] +#[error("worker failed on purpose")] +struct PurposefulFailure; + +/// A migrated WAL database file that is removed when the test finishes. +struct TestDatabase { + path: PathBuf, + pool: SqlitePool, +} + +impl TestDatabase { + async fn new() -> Self { + static DATABASE_NONCE: AtomicUsize = AtomicUsize::new(0); + let path = std::env::temp_dir().join(format!( + "river-work-extensions-{}-{}.sqlite", + std::process::id(), + DATABASE_NONCE.fetch_add(1, Ordering::Relaxed) + )); + let pool = SqlitePoolOptions::new() + .max_connections(4) + .connect_with( + SqliteConnectOptions::new() + .filename(&path) + .create_if_missing(true) + .journal_mode(SqliteJournalMode::Wal) + .busy_timeout(Duration::from_secs(5)), + ) + .await + .unwrap(); + SqliteMigrator::new(pool.clone()) + .migrate_up() + .await + .unwrap(); + Self { path, pool } + } + + async fn close(self) { + self.pool.close().await; + for suffix in ["", "-shm", "-wal"] { + let mut path = self.path.as_os_str().to_owned(); + path.push(suffix); + let _ = std::fs::remove_file(path); + } + } +} + +/// Records the order in which extensions and the worker run. +#[derive(Clone, Default)] +struct Trace(Arc>>); + +impl Trace { + fn push(&self, entry: impl Into) { + self.0.lock().unwrap().push(entry.into()); + } + + fn entries(&self) -> Vec { + self.0.lock().unwrap().clone() + } +} + +struct TracingMiddleware(Trace, &'static str); + +impl WorkMiddleware for TracingMiddleware { + async fn work( + &self, + _context: &WorkContext, + job: JobRow, + next: WorkNext<'_>, + ) -> Result { + self.0.push(format!("middleware {} before", self.1)); + let result = next.run(job).await; + self.0.push(format!("middleware {} after", self.1)); + result + } +} + +/// Records work hooks and replaces a failure with a snooze, like a Go +/// `HookWorkEnd` that returns a different error. +struct TracingHook { + snooze_failures: bool, + trace: Trace, +} + +#[allow( + clippy::unused_async_trait_impl, + reason = "the hook only records state synchronously" +)] +impl Hook for TracingHook { + async fn work_begin( + &self, + _context: &WorkContext, + _job: &mut JobRow, + ) -> Result<(), riverqueue::BoxError> { + self.trace.push("hook begin"); + Ok(()) + } + + async fn work_end( + &self, + _context: &WorkContext, + _job: &JobRow, + result: Result, + ) -> Result { + self.trace.push("hook end"); + match result { + Err(error) + if self.snooze_failures + && error + .get_ref() + .downcast_ref::() + .is_some() => + { + Ok(WorkOutcome::Snooze(Duration::from_mins(1))) + } + result => result, + } + } +} + +/// Panics on every metric, like a buggy metrics integration. +struct PanickingMetricHook; + +impl Hook for PanickingMetricHook { + async fn metric_emit(&self, _metric: Metric) -> Result<(), riverqueue::BoxError> { + panic!("metric hook panicked on purpose"); + } +} + +/// Panics instead of handling the error, like a buggy Go `ErrorHandler`. +struct PanickingErrorHandler; + +impl ErrorHandler for PanickingErrorHandler { + async fn handle_error( + &self, + _context: &WorkContext, + _job: &JobRow, + _result: &riverqueue::WorkResult, + ) -> Result { + panic!("error handler panicked on purpose"); + } +} + +fn workers(trace: &Trace) -> WorkerRegistry { + let trace = trace.clone(); + let mut workers = WorkerRegistry::new(); + workers + .register_fn(move |_context: WorkContext, job: Job| { + let trace = trace.clone(); + async move { + trace.push("worker"); + // The worker runs inside River's span for the job. + trace.push(format!( + "span {}", + tracing::Span::current() + .metadata() + .map_or("none", |metadata| metadata.name()) + )); + if job.args.fail { + return Err(PurposefulFailure); + } + Ok(WorkOutcome::Complete) + } + }) + .unwrap(); + workers +} + +fn client(database: &TestDatabase, trace: &Trace, snooze_failures: bool) -> Client { + Client::builder(database.pool.clone()) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .workers(workers(trace)) + .work_middleware(TracingMiddleware(trace.clone(), "outer")) + .work_middleware(TracingMiddleware(trace.clone(), "inner")) + .hook(TracingHook { + snooze_failures, + trace: trace.clone(), + }) + .build() + .unwrap() +} + +/// Installs a subscriber that records spans, so `Span::current` identifies +/// River's job span inside workers on any runtime thread. +fn install_span_subscriber() { + static INSTALLED: OnceLock<()> = OnceLock::new(); + INSTALLED.get_or_init(|| { + tracing::subscriber::set_global_default(tracing_subscriber::registry()) + .expect("no other global subscriber in this test binary"); + }); +} + +/// Works `job_id` and returns its row from the first event of `kind`. +async fn work_until(client: &Client, kind: EventKind, job_id: i64) -> JobRow { + // Subscribe before starting, so the event can't be missed. + let mut events = client.subscribe(&[kind]).unwrap(); + let mut run = client.start().unwrap(); + let row = tokio::time::timeout(Duration::from_secs(10), async { + loop { + let event = events.recv().await.unwrap(); + if let Some(event) = event.as_job() + && event.job.id == job_id + { + return event.job.clone(); + } + } + }) + .await + .expect("job event"); + run.shutdown().await.unwrap(); + row +} + +#[tokio::test(flavor = "multi_thread")] +async fn hooks_run_inside_middleware_around_the_worker_like_go() { + install_span_subscriber(); + let database = TestDatabase::new().await; + let trace = Trace::default(); + let client = client(&database, &trace, false); + let job = client.insert(ExtensionArgs { fail: false }).await.unwrap(); + let completed = work_until(&client, EventKind::JobCompleted, job.id()).await; + + assert_eq!(completed.state, JobState::Completed); + assert_eq!( + trace.entries(), + [ + "middleware outer before", + "middleware inner before", + "hook begin", + "worker", + "span river_job", + "hook end", + "middleware inner after", + "middleware outer after", + ] + ); + database.close().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn work_end_hooks_replace_the_workers_result() { + let database = TestDatabase::new().await; + let trace = Trace::default(); + let client = client(&database, &trace, true); + let job = client + .insert(ExtensionArgs { fail: true }) + .opts(InsertOpts::default().with_max_attempts(1)) + .await + .unwrap(); + let snoozed = work_until(&client, EventKind::JobSnoozed, job.id()).await; + + // The hook turned the final failure into a snooze, which neither + // records an error nor consumes the attempt. + assert_eq!(snoozed.state, JobState::Scheduled); + assert_eq!(snoozed.errors, []); + database.close().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn panicking_error_handlers_still_persist_the_result() { + let database = TestDatabase::new().await; + let trace = Trace::default(); + let client = Client::builder(database.pool.clone()) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .workers(workers(&trace)) + .error_handler(PanickingErrorHandler) + .build() + .unwrap(); + let job = client + .insert(ExtensionArgs { fail: true }) + .opts(InsertOpts::default().with_max_attempts(1)) + .await + .unwrap(); + let failed = work_until(&client, EventKind::JobFailed, job.id()).await; + + // Like Go, the panic is treated as a handler failure: the worker's + // error is still recorded and the job leaves `running`. + assert_eq!(failed.state, JobState::Discarded); + assert_eq!(failed.errors.len(), 1); + assert_eq!( + client.jobs().get(job.id()).await.unwrap().state, + JobState::Discarded + ); + database.close().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn panicking_metric_hooks_dont_stop_the_queue() { + let database = TestDatabase::new().await; + let trace = Trace::default(); + let client = Client::builder(database.pool.clone()) + .queue( + "default", + QueueConfig::new(1) + .with_fetch_cooldown(Duration::from_millis(1)) + .with_fetch_poll_interval(Duration::from_millis(10)), + ) + .workers(workers(&trace)) + .hook(PanickingMetricHook) + .build() + .unwrap(); + let job = client.insert(ExtensionArgs { fail: false }).await.unwrap(); + let completed = work_until(&client, EventKind::JobCompleted, job.id()).await; + + assert_eq!(completed.state, JobState::Completed); + database.close().await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn unknown_kinds_fail_before_middleware_and_hooks() { + let database = TestDatabase::new().await; + let trace = Trace::default(); + let client = client(&database, &trace, false); + // Insert with a client that knows the kind, then work it with one that + // doesn't. + let job = Client::builder(database.pool.clone()) + .build() + .unwrap() + .insert(UnknownArgs {}) + .await + .unwrap(); + let failed = work_until(&client, EventKind::JobFailed, job.id()).await; + + // A short first retry stays available, as in River Go. + assert!( + matches!(failed.state, JobState::Available | JobState::Retryable), + "{:?}", + failed.state + ); + assert_eq!(failed.errors.len(), 1); + assert!(trace.entries().is_empty(), "{:?}", trace.entries()); + database.close().await; +} diff --git a/rust/riverqueue/tests/yugabyte.rs b/rust/riverqueue/tests/yugabyte.rs new file mode 100644 index 000000000..ad68f9418 --- /dev/null +++ b/rust/riverqueue/tests/yugabyte.rs @@ -0,0 +1,295 @@ +//! PostgreSQL-compatible servers without `xmax` or `LISTEN`/`NOTIFY`, like +//! YugabyteDB, simulated on PostgreSQL the way River Go's tests do. +//! +//! A test schema shadows `version()` and `current_setting(text, boolean)` +//! ahead of `pg_catalog` on the connections' `search_path`, so River detects +//! a Yugabyte version and notification setting. When notifications are off, +//! it also shadows `pg_notify` with a function that raises, so any attempt +//! to notify fails. This exercises detection and River's fallbacks, not +//! Yugabyte's storage or transaction semantics. +//! +//! These tests fail rather than skip when `RIVER_RUST_DATABASE_URL` is unset. + +#![cfg(all(feature = "postgres", river_postgres_tests))] + +mod support; + +use std::{convert::Infallible, sync::Arc, time::Duration}; + +use riverqueue::{ + Client, EventKind, InsertOpts, Job, JobArgs, JobState, QueueConfig, QueueSelector, UniqueOpts, + WorkContext, WorkOutcome, WorkerRegistry, database::PostgresDatabase, +}; +use serde::{Deserialize, Serialize}; +use sqlx::{ + AssertSqlSafe, PgPool, + postgres::{PgConnectOptions, PgPoolOptions}, +}; +use tokio::sync::Semaphore; + +/// Every wait in these tests is bounded by this timeout. It covers a few of +/// the two-second polls for cancellation requests. +const TIMEOUT: Duration = Duration::from_secs(10); + +#[derive(Clone, Debug, Deserialize, JobArgs, Serialize)] +#[river(kind = "rust_yugabyte")] +struct YugabyteArgs { + value: i32, +} + +/// Which server a test schema simulates. +#[derive(Clone, Copy, Debug)] +enum Server { + /// PostgreSQL 17, before `RETURNING OLD`. + Postgres17, + /// YugabyteDB before 2025.2.3, without `yb_enable_listen_notify`. + YugabyteUnavailable, + /// YugabyteDB with `yb_enable_listen_notify` off. + YugabyteDisabled, + /// YugabyteDB with `yb_enable_listen_notify` on. + YugabyteEnabled, +} + +impl Server { + const ALL: [Self; 4] = [ + Self::Postgres17, + Self::YugabyteUnavailable, + Self::YugabyteDisabled, + Self::YugabyteEnabled, + ]; + + const fn listen_notify(self) -> bool { + matches!(self, Self::Postgres17 | Self::YugabyteEnabled) + } + + const fn yugabyte(self) -> bool { + !matches!(self, Self::Postgres17) + } +} + +/// A migrated test schema and a pool whose connections see the simulated +/// server. +struct Simulated { + pool: PgPool, + schema: support::PostgresSchema, +} + +impl Simulated { + async fn new(server: Server) -> Self { + let schema = support::PostgresSchema::new("river_yugabyte").await; + let name = schema.schema.as_deref().unwrap().to_owned(); + let functions = match server { + Server::Postgres17 => format!( + "CREATE FUNCTION \"{name}\".current_setting(setting_name text) RETURNS text \ + LANGUAGE sql AS $$ SELECT CASE WHEN setting_name = 'server_version_num' \ + THEN '170004' ELSE pg_catalog.current_setting(setting_name) END $$;" + ), + Server::YugabyteUnavailable | Server::YugabyteDisabled | Server::YugabyteEnabled => { + let (version, setting) = match server { + Server::YugabyteUnavailable => ("2025.2.1.0", "NULL::text"), + Server::YugabyteDisabled => ("2025.2.3.0", "'off'::text"), + _ => ("2025.2.3.0", "'on'::text"), + }; + format!( + "CREATE FUNCTION \"{name}\".version() RETURNS text LANGUAGE sql AS $$ \ + SELECT 'PostgreSQL 15.12-YB-{version}-b1'::text $$; \ + CREATE FUNCTION \"{name}\".current_setting(setting_name text, missing_ok boolean) \ + RETURNS text LANGUAGE sql AS $$ SELECT CASE WHEN setting_name = \ + 'yb_enable_listen_notify' THEN {setting} ELSE \ + pg_catalog.current_setting(setting_name, missing_ok) END $$;" + ) + } + }; + sqlx::raw_sql(AssertSqlSafe(functions)) + .execute(&schema.pool) + .await + .unwrap(); + if !server.listen_notify() { + sqlx::raw_sql(AssertSqlSafe(format!( + "CREATE FUNCTION \"{name}\".pg_notify(text, text) RETURNS void LANGUAGE plpgsql \ + AS $$ BEGIN RAISE EXCEPTION 'LISTEN/NOTIFY is unavailable'; END $$;" + ))) + .execute(&schema.pool) + .await + .unwrap(); + } + let url = std::env::var("RIVER_RUST_DATABASE_URL").unwrap(); + let options = url + .parse::() + .unwrap() + .options([("search_path", format!("{name},pg_catalog"))]); + let pool = PgPoolOptions::new() + .max_connections(8) + .connect_with(options) + .await + .unwrap(); + Self { pool, schema } + } + + fn database(&self) -> PostgresDatabase { + PostgresDatabase::new(self.pool.clone()).with_schema(self.schema.schema.clone()) + } + + async fn cleanup(self) { + self.pool.close().await; + self.schema.cleanup().await; + } +} + +/// An insert-only client detects the server as it goes: unique inserts tell +/// a duplicate from a new row with a nonce on Yugabyte and `xmax` before +/// PostgreSQL 18, and notifications, cancellation, queue changes, and +/// resignation requests work without `pg_notify` when it's unavailable. +#[tokio::test] +async fn detects_the_server_without_starting() { + for server in Server::ALL { + let simulated = Simulated::new(server).await; + let client = Client::builder(simulated.database()).build().unwrap(); + let unique = InsertOpts::default().with_unique(UniqueOpts::new().with_by_args(true)); + + let first = client + .insert(YugabyteArgs { value: 1 }) + .opts(unique.clone()) + .await + .unwrap(); + assert!(!first.unique_skipped_as_duplicate, "{server:?}"); + let second = client + .insert(YugabyteArgs { value: 1 }) + .opts(unique.clone()) + .await + .unwrap(); + assert!(second.unique_skipped_as_duplicate, "{server:?}"); + assert_eq!(second.job.row.id, first.job.row.id, "{server:?}"); + let other = client + .insert(YugabyteArgs { value: 2 }) + .opts(unique) + .await + .unwrap(); + assert!(!other.unique_skipped_as_duplicate, "{server:?}"); + // Like River Go, the nonce stays in the stored metadata. + let nonce = first + .job + .row + .metadata + .get::("river:unique_nonce") + .unwrap(); + assert_eq!(nonce.is_some(), server.yugabyte(), "{server:?}"); + + let cancelled = client.jobs().cancel(other.job.row.id).await.unwrap(); + assert_eq!(cancelled.state, JobState::Cancelled, "{server:?}"); + client.queues().pause(QueueSelector::All).await.unwrap(); + client.request_resign().await.unwrap(); + + simulated.cleanup().await; + } +} + +/// A client of a server without `LISTEN`/`NOTIFY` works jobs and hears a +/// cancellation from another client by polling, without being configured +/// as poll-only. +#[tokio::test] +async fn polls_without_listen_notify() { + for server in [Server::YugabyteUnavailable, Server::YugabyteDisabled] { + let simulated = Simulated::new(server).await; + let started = Arc::new(Semaphore::new(0)); + let worker_started = Arc::clone(&started); + let mut workers = WorkerRegistry::new(); + workers + .register_fn(move |context: WorkContext, job: Job| { + let started = Arc::clone(&worker_started); + async move { + if job.args.value == 0 { + return Ok::<_, Infallible>(WorkOutcome::Complete); + } + started.add_permits(1); + context.cancellation_token().cancelled().await; + // Any outcome but completion becomes the cancellation. + Ok(WorkOutcome::Snooze(Duration::from_hours(1))) + } + }) + .unwrap(); + let client = Client::builder(simulated.database()) + .queue( + riverqueue::QUEUE_DEFAULT, + QueueConfig::new(2).with_fetch_poll_interval(Duration::from_millis(100)), + ) + .workers(workers) + .build() + .unwrap(); + let other = Client::builder(simulated.database()).build().unwrap(); + let mut events = client + .subscribe(&[EventKind::JobCompleted, EventKind::JobCancelled]) + .unwrap(); + let mut run = client.start().unwrap(); + run.wait_ready().await.unwrap(); + + let completed = other.insert(YugabyteArgs { value: 0 }).await.unwrap(); + let event = tokio::time::timeout(TIMEOUT, events.recv()) + .await + .expect("the job should complete") + .unwrap(); + let job = &event.as_job().expect("a job event").job; + assert_eq!(job.id, completed.job.row.id, "{server:?}"); + assert_eq!(job.state, JobState::Completed, "{server:?}"); + + let cancellable = other.insert(YugabyteArgs { value: 1 }).await.unwrap(); + tokio::time::timeout(TIMEOUT, started.acquire()) + .await + .expect("the job should start") + .unwrap() + .forget(); + other.jobs().cancel(cancellable.job.row.id).await.unwrap(); + let event = tokio::time::timeout(TIMEOUT, events.recv()) + .await + .expect("the job should be cancelled") + .unwrap(); + let job = &event.as_job().expect("a job event").job; + assert_eq!(job.id, cancellable.job.row.id, "{server:?}"); + assert_eq!(job.state, JobState::Cancelled, "{server:?}"); + + tokio::time::timeout(TIMEOUT, run.shutdown()) + .await + .expect("the client should stop") + .unwrap(); + simulated.cleanup().await; + } +} + +/// Like River Go's check that YugabyteDB-incompatible system columns only +/// appear where the unique insert mode replaces them. +#[test] +fn system_columns_appear_only_in_unique_insert_modes() { + let source = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("src"); + let mut violations = Vec::new(); + let mut directories = vec![source]; + while let Some(directory) = directories.pop() { + for entry in std::fs::read_dir(&directory).unwrap() { + let path = entry.unwrap().path(); + if path.is_dir() { + directories.push(path); + continue; + } + if path.extension().is_none_or(|extension| extension != "rs") + || path.ends_with("database/postgres_capabilities.rs") + { + continue; + } + let contents = std::fs::read_to_string(&path).unwrap(); + for (index, line) in contents.lines().enumerate() { + if line.trim_start().starts_with("//") { + continue; + } + let has_column = line + .split(|character: char| !character.is_ascii_alphanumeric() && character != '_') + .any(|word| ["cmax", "cmin", "ctid", "xmax", "xmin"].contains(&word)); + if has_column { + violations.push(format!("{}:{}: {line}", path.display(), index + 1)); + } + } + } + } + assert!( + violations.is_empty(), + "system columns YugabyteDB lacks must only appear in the unique insert modes: {violations:#?}" + ); +}