From 2a41351fcf873da86c9a0bc1401ac3e0f609d935 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 4 Oct 2026 18:02:39 -0500 Subject: [PATCH 01/20] add the Rust migration crate Start a Rust workspace under `rust/` with `riverqueue-migrate`, which applies, lists, previews, and validates River's PostgreSQL and SQLite migration lines from Rust. It shares the `river_migration` history with River Go, so either language can migrate a database the other uses, and it accepts any quotable PostgreSQL schema name. The crate can't read files outside its package, so it carries mirrors of Go's canonical migrations. `syncrustmigrations` writes the mirrors and the hashes in the conformance migration inventories from Go's migration directories, and `make verify/rust-migrations`, which the conformance artifacts CI job now runs, fails when they drift, so Rust never gains an independent schema history. --- .gitignore | 1 + Makefile | 10 + internal/cmd/syncrustmigrations/main.go | 150 ++ rust/Cargo.lock | 1391 +++++++++++++++++ rust/Cargo.toml | 52 + rust/riverqueue-migrate/Cargo.toml | 46 + rust/riverqueue-migrate/LICENSE | 1 + rust/riverqueue-migrate/README.md | 35 + rust/riverqueue-migrate/examples/migrate.rs | 24 + .../main/001_create_river_migration.down.sql | 1 + .../main/001_create_river_migration.up.sql | 8 + .../main/002_initial_schema.down.sql | 5 + .../migrations/main/002_initial_schema.up.sql | 96 ++ .../main/003_river_job_tags_non_null.down.sql | 3 + .../main/003_river_job_tags_non_null.up.sql | 3 + .../main/004_pending_and_more.down.sql | 42 + .../main/004_pending_and_more.up.sql | 45 + .../main/005_migration_unique_client.down.sql | 57 + .../main/005_migration_unique_client.up.sql | 79 + .../migrations/main/006_bulk_unique.down.sql | 16 + .../migrations/main/006_bulk_unique.up.sql | 40 + ...tbox_sqlite_jsonb_and_sql_cleanup.down.sql | 56 + ...outbox_sqlite_jsonb_and_sql_cleanup.up.sql | 44 + .../main/008_job_id_autoincrement.down.sql | 3 + .../main/008_job_id_autoincrement.up.sql | 3 + .../main/001_create_river_migration.down.sql | 1 + .../main/001_create_river_migration.up.sql | 8 + .../sqlite/main/002_initial_schema.down.sql | 8 + .../sqlite/main/002_initial_schema.up.sql | 19 + .../main/003_river_job_tags_non_null.down.sql | 6 + .../main/003_river_job_tags_non_null.up.sql | 6 + .../sqlite/main/004_pending_and_more.down.sql | 26 + .../sqlite/main/004_pending_and_more.up.sql | 33 + .../main/005_migration_unique_client.down.sql | 37 + .../main/005_migration_unique_client.up.sql | 64 + .../sqlite/main/006_bulk_unique.down.sql | 7 + .../sqlite/main/006_bulk_unique.up.sql | 63 + ...tbox_sqlite_jsonb_and_sql_cleanup.down.sql | 255 +++ ...outbox_sqlite_jsonb_and_sql_cleanup.up.sql | 261 ++++ .../main/008_job_id_autoincrement.down.sql | 121 ++ .../main/008_job_id_autoincrement.up.sql | 123 ++ rust/riverqueue-migrate/src/lib.rs | 764 +++++++++ rust/riverqueue-migrate/src/schema.rs | 178 +++ rust/riverqueue-migrate/src/sqlite.rs | 262 ++++ rust/riverqueue-migrate/tests/postgres.rs | 320 ++++ rust/riverqueue-migrate/tests/sqlite.rs | 489 ++++++ 46 files changed, 5262 insertions(+) create mode 100644 internal/cmd/syncrustmigrations/main.go create mode 100644 rust/Cargo.lock create mode 100644 rust/Cargo.toml create mode 100644 rust/riverqueue-migrate/Cargo.toml create mode 120000 rust/riverqueue-migrate/LICENSE create mode 100644 rust/riverqueue-migrate/README.md create mode 100644 rust/riverqueue-migrate/examples/migrate.rs create mode 100644 rust/riverqueue-migrate/migrations/main/001_create_river_migration.down.sql create mode 100644 rust/riverqueue-migrate/migrations/main/001_create_river_migration.up.sql create mode 100644 rust/riverqueue-migrate/migrations/main/002_initial_schema.down.sql create mode 100644 rust/riverqueue-migrate/migrations/main/002_initial_schema.up.sql create mode 100644 rust/riverqueue-migrate/migrations/main/003_river_job_tags_non_null.down.sql create mode 100644 rust/riverqueue-migrate/migrations/main/003_river_job_tags_non_null.up.sql create mode 100644 rust/riverqueue-migrate/migrations/main/004_pending_and_more.down.sql create mode 100644 rust/riverqueue-migrate/migrations/main/004_pending_and_more.up.sql create mode 100644 rust/riverqueue-migrate/migrations/main/005_migration_unique_client.down.sql create mode 100644 rust/riverqueue-migrate/migrations/main/005_migration_unique_client.up.sql create mode 100644 rust/riverqueue-migrate/migrations/main/006_bulk_unique.down.sql create mode 100644 rust/riverqueue-migrate/migrations/main/006_bulk_unique.up.sql create mode 100644 rust/riverqueue-migrate/migrations/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.down.sql create mode 100644 rust/riverqueue-migrate/migrations/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.up.sql create mode 100644 rust/riverqueue-migrate/migrations/main/008_job_id_autoincrement.down.sql create mode 100644 rust/riverqueue-migrate/migrations/main/008_job_id_autoincrement.up.sql create mode 100644 rust/riverqueue-migrate/migrations/sqlite/main/001_create_river_migration.down.sql create mode 100644 rust/riverqueue-migrate/migrations/sqlite/main/001_create_river_migration.up.sql create mode 100644 rust/riverqueue-migrate/migrations/sqlite/main/002_initial_schema.down.sql create mode 100644 rust/riverqueue-migrate/migrations/sqlite/main/002_initial_schema.up.sql create mode 100644 rust/riverqueue-migrate/migrations/sqlite/main/003_river_job_tags_non_null.down.sql create mode 100644 rust/riverqueue-migrate/migrations/sqlite/main/003_river_job_tags_non_null.up.sql create mode 100644 rust/riverqueue-migrate/migrations/sqlite/main/004_pending_and_more.down.sql create mode 100644 rust/riverqueue-migrate/migrations/sqlite/main/004_pending_and_more.up.sql create mode 100644 rust/riverqueue-migrate/migrations/sqlite/main/005_migration_unique_client.down.sql create mode 100644 rust/riverqueue-migrate/migrations/sqlite/main/005_migration_unique_client.up.sql create mode 100644 rust/riverqueue-migrate/migrations/sqlite/main/006_bulk_unique.down.sql create mode 100644 rust/riverqueue-migrate/migrations/sqlite/main/006_bulk_unique.up.sql create mode 100644 rust/riverqueue-migrate/migrations/sqlite/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.down.sql create mode 100644 rust/riverqueue-migrate/migrations/sqlite/main/007_notification_outbox_sqlite_jsonb_and_sql_cleanup.up.sql create mode 100644 rust/riverqueue-migrate/migrations/sqlite/main/008_job_id_autoincrement.down.sql create mode 100644 rust/riverqueue-migrate/migrations/sqlite/main/008_job_id_autoincrement.up.sql create mode 100644 rust/riverqueue-migrate/src/lib.rs create mode 100644 rust/riverqueue-migrate/src/schema.rs create mode 100644 rust/riverqueue-migrate/src/sqlite.rs create mode 100644 rust/riverqueue-migrate/tests/postgres.rs create mode 100644 rust/riverqueue-migrate/tests/sqlite.rs 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..85c30eef2 100644 --- a/Makefile +++ b/Makefile @@ -18,12 +18,17 @@ 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 and hashes to Rust + go run ./internal/cmd/syncrustmigrations + .PHONY: generate/sqlc generate/sqlc: ## Generate sqlc cd riverdriver/riverdatabasesql/internal/dbsqlc && sqlc generate @@ -122,12 +127,17 @@ 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 and protocol hashes + go run ./internal/cmd/syncrustmigrations -check + .PHONY: verify/sqlc verify/sqlc: ## Verify generated sqlc cd riverdriver/riverdatabasesql/internal/dbsqlc && sqlc diff diff --git a/internal/cmd/syncrustmigrations/main.go b/internal/cmd/syncrustmigrations/main.go new file mode 100644 index 000000000..8846a78dc --- /dev/null +++ b/internal/cmd/syncrustmigrations/main.go @@ -0,0 +1,150 @@ +// Command syncrustmigrations mirrors River's canonical database migrations +// into the publishable Rust migration crate and records their hashes for +// cross-language conformance. +package main + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "flag" + "fmt" + "os" + "path/filepath" + "slices" + "strings" +) + +type database struct { + canonicalDir string + manifestPath string + mirrorDir string + name string +} + +type manifest struct { + Database string `json:"database"` + Files []manifestFile `json:"files"` + Line string `json:"line"` +} + +type manifestFile struct { + Path string `json:"path"` + SHA256 string `json:"sha256"` +} + +func main() { + check := flag.Bool("check", false, "check generated files without writing") + flag.Parse() + + databases := []database{ + { + canonicalDir: "riverdriver/riverpgxv5/migration/main", + manifestPath: "conformance/migrations.json", + mirrorDir: "rust/riverqueue-migrate/migrations/main", + name: "postgres", + }, + { + canonicalDir: "riverdriver/riversqlite/migration/main", + manifestPath: "conformance/migrations-sqlite.json", + mirrorDir: "rust/riverqueue-migrate/migrations/sqlite/main", + name: "sqlite", + }, + } + 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) + + generatedManifest := manifest{Database: database.name, Line: "main"} + for _, name := range names { + sourcePath := filepath.Join(database.canonicalDir, name) + contents, err := os.ReadFile(sourcePath) + if err != nil { + fatal(err) + } + hash := sha256.Sum256(contents) + generatedManifest.Files = append(generatedManifest.Files, manifestFile{ + Path: filepath.ToSlash(sourcePath), + SHA256: hex.EncodeToString(hash[:]), + }) + + mirrorPath := filepath.Join(database.mirrorDir, name) + if check { + checkFile(mirrorPath, contents) + } else { + writeFile(mirrorPath, contents) + } + } + removeStaleMirrors(database.mirrorDir, names, check) + + manifestContents, err := json.MarshalIndent(&generatedManifest, "", " ") + if err != nil { + fatal(err) + } + manifestContents = append(manifestContents, '\n') + if check { + checkFile(database.manifestPath, manifestContents) + } else { + writeFile(database.manifestPath, manifestContents) + } +} + +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/Cargo.lock b/rust/Cargo.lock new file mode 100644 index 000000000..5db99fdaf --- /dev/null +++ b/rust/Cargo.lock @@ -0,0 +1,1391 @@ +# 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 = "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 = "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 = "cmov" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" + +[[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", + "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 = "etcetera" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de48cc4d1c1d97a20fd819def54b890cadde72ed3ad0c614822a0a433361be96" +dependencies = [ + "cfg-if", + "windows-sys", +] + +[[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-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-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.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "300e883d756b2e4ec94e02791f39b04b522276138852cfc41d9fb7e904106099" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "rand_core", +] + +[[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 = "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 = "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", +] + +[[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 = "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", + "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 = "riverqueue-migrate" +version = "0.49.0-alpha.1" +dependencies = [ + "serde_json", + "sqlx", + "thiserror", + "tokio", +] + +[[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 = "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 = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[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", +] + +[[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", + "crc", + "crossbeam-queue", + "either", + "event-listener", + "futures-core", + "futures-intrusive", + "futures-io", + "futures-util", + "hashbrown 0.16.1", + "hashlink", + "indexmap", + "log", + "memchr", + "percent-encoding", + "serde", + "serde_json", + "sha2 0.10.9", + "smallvec", + "thiserror", + "tokio", + "tokio-stream", + "tracing", + "url", +] + +[[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", + "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", + "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", + "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 = "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 = "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 = "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", + "socket2", + "tokio-macros", + "windows-sys", +] + +[[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 = "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 = "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 = "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 = "whoami" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "998767ef88740d1f5b0682a9c53c24431453923962269c2db68ee43788c5a40d" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[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 = "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 = "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..da79c57a8 --- /dev/null +++ b/rust/Cargo.toml @@ -0,0 +1,52 @@ +[workspace] +members = [ + "riverqueue-migrate", +] +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.49.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/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"); +} From c48c8f4e5a071fc150d7d15ce73f153bf2159eeb Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 4 Oct 2026 18:02:39 -0500 Subject: [PATCH 02/20] add the Rust JobArgs derive macro Add `riverqueue-macros` with `#[derive(JobArgs)]`, which requires a stable `#[river(kind = "...")]` and can declare kind aliases, the default queue, max attempts, priority, pending state, tags, and default unique options, including the fields that make up a `by_args` key. Invalid attributes fail at compile time with spans that point at them. Applications receive the macro through `riverqueue`. --- rust/Cargo.lock | 9 + rust/Cargo.toml | 1 + rust/riverqueue-macros/Cargo.toml | 23 + rust/riverqueue-macros/LICENSE | 1 + rust/riverqueue-macros/README.md | 41 ++ rust/riverqueue-macros/src/lib.rs | 1090 +++++++++++++++++++++++++++++ 6 files changed, 1165 insertions(+) create mode 100644 rust/riverqueue-macros/Cargo.toml create mode 120000 rust/riverqueue-macros/LICENSE create mode 100644 rust/riverqueue-macros/README.md create mode 100644 rust/riverqueue-macros/src/lib.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 5db99fdaf..1ed1d6d9a 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -745,6 +745,15 @@ dependencies = [ "bitflags", ] +[[package]] +name = "riverqueue-macros" +version = "0.49.0-alpha.1" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "riverqueue-migrate" version = "0.49.0-alpha.1" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index da79c57a8..627219d6c 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -1,5 +1,6 @@ [workspace] members = [ + "riverqueue-macros", "riverqueue-migrate", ] resolver = "3" diff --git a/rust/riverqueue-macros/Cargo.toml b/rust/riverqueue-macros/Cargo.toml new file mode 100644 index 000000000..487cc854e --- /dev/null +++ b/rust/riverqueue-macros/Cargo.toml @@ -0,0 +1,23 @@ +[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 + +[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::")); + } +} From 22e6cee7f35eae0ed11a422d5d602726e982e7ea Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 4 Oct 2026 18:02:39 -0500 Subject: [PATCH 03/20] add River's Rust client Add `riverqueue`, a Rust and Tokio implementation of River that shares the database schema and job protocol with River Go on PostgreSQL and SQLite, so Rust and Go clients can insert and work jobs in the same database. `Client` takes a caller-owned SQLx pool and isn't generic over the database or a driver trait. Workers are typed and async and get a `CancellationToken`; request builders cover insertion (including transactional, batch, unique, scheduled, and pending jobs), job and queue management, and transactional completion from inside a worker. Like River Go's `*Tx` methods, a request given the caller's transaction runs directly in it without a savepoint, so the caller rolls back on an error. The runtime fetches, completes in bounded concurrent batches, retries, snoozes, cancels, and rescues jobs the way Go does, publishes events only after their database update commits, and runs leader election and maintenance services: the rescuer, cleaners, scheduler, reindexer, and periodic jobs with Go-compatible cron parsing. Hooks, middleware, error handlers, retry policies, and a hidden extension module for lockstep add-on crates mirror Go's extension points. Persisted values follow Go exactly where another engine reads them: unique keys, metadata, attempt errors, notifications, and timestamps. Unit tests check unique keys and cron schedules against the conformance fixtures, which the fixture generator now also copies into the crate because published tests can't read files outside it. --- rust/Cargo.lock | 342 ++ rust/Cargo.toml | 1 + rust/riverqueue/Cargo.toml | 52 + rust/riverqueue/LICENSE | 1 + rust/riverqueue/README.md | 379 +++ rust/riverqueue/docs/mixed-deployments.md | 122 + rust/riverqueue/src/__private.rs | 1431 +++++++++ rust/riverqueue/src/client/attempts.rs | 123 + rust/riverqueue/src/client/backoff.rs | 124 + rust/riverqueue/src/client/builder.rs | 872 +++++ rust/riverqueue/src/client/completer.rs | 807 +++++ rust/riverqueue/src/client/executor.rs | 783 +++++ rust/riverqueue/src/client/extension.rs | 367 +++ rust/riverqueue/src/client/insert.rs | 1001 ++++++ rust/riverqueue/src/client/jobs.rs | 453 +++ rust/riverqueue/src/client/local_queues.rs | 178 ++ rust/riverqueue/src/client/mod.rs | 540 ++++ rust/riverqueue/src/client/notifier.rs | 440 +++ rust/riverqueue/src/client/notify_limiter.rs | 96 + rust/riverqueue/src/client/peers.rs | 697 ++++ rust/riverqueue/src/client/producer.rs | 1282 ++++++++ rust/riverqueue/src/client/queues.rs | 265 ++ rust/riverqueue/src/client/record.rs | 287 ++ rust/riverqueue/src/client/request.rs | 170 + rust/riverqueue/src/client/resign.rs | 66 + rust/riverqueue/src/client/run.rs | 755 +++++ rust/riverqueue/src/client/tests.rs | 721 +++++ rust/riverqueue/src/client/validate.rs | 147 + rust/riverqueue/src/database.rs | 1031 ++++++ rust/riverqueue/src/database/erased.rs | 381 +++ .../src/database/postgres_capabilities.rs | 188 ++ rust/riverqueue/src/database/sqlite.rs | 2827 +++++++++++++++++ rust/riverqueue/src/encoding.rs | 586 ++++ rust/riverqueue/src/encoding/finite.rs | 319 ++ rust/riverqueue/src/error.rs | 510 +++ rust/riverqueue/src/event.rs | 403 +++ rust/riverqueue/src/extension.rs | 833 +++++ rust/riverqueue/src/guide.rs | 7 + rust/riverqueue/src/job.rs | 1066 +++++++ rust/riverqueue/src/job/attempt_error.rs | 839 +++++ rust/riverqueue/src/job/metadata.rs | 295 ++ rust/riverqueue/src/lib.rs | 153 + rust/riverqueue/src/maintenance/cleaner.rs | 342 ++ rust/riverqueue/src/maintenance/elector.rs | 787 +++++ rust/riverqueue/src/maintenance/maintainer.rs | 320 ++ rust/riverqueue/src/maintenance/mod.rs | 426 +++ .../src/maintenance/periodic_enqueuer.rs | 42 + rust/riverqueue/src/maintenance/postgres.rs | 160 + rust/riverqueue/src/maintenance/reindexer.rs | 331 ++ rust/riverqueue/src/maintenance/rescuer.rs | 496 +++ rust/riverqueue/src/maintenance/scheduler.rs | 360 +++ rust/riverqueue/src/maintenance/tests.rs | 1095 +++++++ rust/riverqueue/src/periodic/cron.rs | 924 ++++++ rust/riverqueue/src/periodic/mod.rs | 696 ++++ rust/riverqueue/src/pilot.rs | 448 +++ rust/riverqueue/src/protocol.rs | 80 + rust/riverqueue/src/query.rs | 931 ++++++ rust/riverqueue/src/queue.rs | 127 + rust/riverqueue/src/storage.rs | 654 ++++ rust/riverqueue/src/storage/postgres.rs | 502 +++ rust/riverqueue/src/storage/sqlite.rs | 279 ++ rust/riverqueue/src/unique.rs | 1154 +++++++ rust/riverqueue/src/worker.rs | 1537 +++++++++ .../tests/fixtures/maintenance_values.json | 681 ++++ .../tests/fixtures/protocol_values.json | 319 ++ .../tests/fixtures/unique_keys.json | 950 ++++++ 66 files changed, 35581 insertions(+) create mode 100644 rust/riverqueue/Cargo.toml create mode 120000 rust/riverqueue/LICENSE create mode 100644 rust/riverqueue/README.md create mode 100644 rust/riverqueue/docs/mixed-deployments.md create mode 100644 rust/riverqueue/src/__private.rs create mode 100644 rust/riverqueue/src/client/attempts.rs create mode 100644 rust/riverqueue/src/client/backoff.rs create mode 100644 rust/riverqueue/src/client/builder.rs create mode 100644 rust/riverqueue/src/client/completer.rs create mode 100644 rust/riverqueue/src/client/executor.rs create mode 100644 rust/riverqueue/src/client/extension.rs create mode 100644 rust/riverqueue/src/client/insert.rs create mode 100644 rust/riverqueue/src/client/jobs.rs create mode 100644 rust/riverqueue/src/client/local_queues.rs create mode 100644 rust/riverqueue/src/client/mod.rs create mode 100644 rust/riverqueue/src/client/notifier.rs create mode 100644 rust/riverqueue/src/client/notify_limiter.rs create mode 100644 rust/riverqueue/src/client/peers.rs create mode 100644 rust/riverqueue/src/client/producer.rs create mode 100644 rust/riverqueue/src/client/queues.rs create mode 100644 rust/riverqueue/src/client/record.rs create mode 100644 rust/riverqueue/src/client/request.rs create mode 100644 rust/riverqueue/src/client/resign.rs create mode 100644 rust/riverqueue/src/client/run.rs create mode 100644 rust/riverqueue/src/client/tests.rs create mode 100644 rust/riverqueue/src/client/validate.rs create mode 100644 rust/riverqueue/src/database.rs create mode 100644 rust/riverqueue/src/database/erased.rs create mode 100644 rust/riverqueue/src/database/postgres_capabilities.rs create mode 100644 rust/riverqueue/src/database/sqlite.rs create mode 100644 rust/riverqueue/src/encoding.rs create mode 100644 rust/riverqueue/src/encoding/finite.rs create mode 100644 rust/riverqueue/src/error.rs create mode 100644 rust/riverqueue/src/event.rs create mode 100644 rust/riverqueue/src/extension.rs create mode 100644 rust/riverqueue/src/guide.rs create mode 100644 rust/riverqueue/src/job.rs create mode 100644 rust/riverqueue/src/job/attempt_error.rs create mode 100644 rust/riverqueue/src/job/metadata.rs create mode 100644 rust/riverqueue/src/lib.rs create mode 100644 rust/riverqueue/src/maintenance/cleaner.rs create mode 100644 rust/riverqueue/src/maintenance/elector.rs create mode 100644 rust/riverqueue/src/maintenance/maintainer.rs create mode 100644 rust/riverqueue/src/maintenance/mod.rs create mode 100644 rust/riverqueue/src/maintenance/periodic_enqueuer.rs create mode 100644 rust/riverqueue/src/maintenance/postgres.rs create mode 100644 rust/riverqueue/src/maintenance/reindexer.rs create mode 100644 rust/riverqueue/src/maintenance/rescuer.rs create mode 100644 rust/riverqueue/src/maintenance/scheduler.rs create mode 100644 rust/riverqueue/src/maintenance/tests.rs create mode 100644 rust/riverqueue/src/periodic/cron.rs create mode 100644 rust/riverqueue/src/periodic/mod.rs create mode 100644 rust/riverqueue/src/pilot.rs create mode 100644 rust/riverqueue/src/protocol.rs create mode 100644 rust/riverqueue/src/query.rs create mode 100644 rust/riverqueue/src/queue.rs create mode 100644 rust/riverqueue/src/storage.rs create mode 100644 rust/riverqueue/src/storage/postgres.rs create mode 100644 rust/riverqueue/src/storage/sqlite.rs create mode 100644 rust/riverqueue/src/unique.rs create mode 100644 rust/riverqueue/src/worker.rs create mode 100644 rust/riverqueue/tests/fixtures/maintenance_values.json create mode 100644 rust/riverqueue/tests/fixtures/protocol_values.json create mode 100644 rust/riverqueue/tests/fixtures/unique_keys.json diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 1ed1d6d9a..e9c6cd2c1 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -8,6 +8,32 @@ 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" @@ -56,6 +82,12 @@ 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" @@ -95,12 +127,48 @@ dependencies = [ "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" @@ -194,6 +262,7 @@ 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", ] @@ -230,6 +299,16 @@ 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", +] + [[package]] name = "etcetera" version = "0.11.0" @@ -326,6 +405,17 @@ 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" @@ -346,6 +436,7 @@ checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc" dependencies = [ "futures-core", "futures-io", + "futures-macro", "futures-sink", "futures-task", "memchr", @@ -440,6 +531,30 @@ 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" @@ -559,6 +674,23 @@ 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" @@ -674,6 +806,24 @@ 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" @@ -745,6 +895,30 @@ dependencies = [ "bitflags", ] +[[package]] +name = "riverqueue" +version = "0.49.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-macros" version = "0.49.0-alpha.1" @@ -764,6 +938,12 @@ dependencies = [ "tokio", ] +[[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" @@ -846,12 +1026,37 @@ dependencies = [ "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" @@ -908,6 +1113,7 @@ dependencies = [ "base64", "bytes", "cfg-if", + "chrono", "crc", "crossbeam-queue", "either", @@ -980,6 +1186,7 @@ dependencies = [ "bitflags", "byteorder", "bytes", + "chrono", "crc", "digest 0.11.3", "dotenvy", @@ -1007,6 +1214,7 @@ dependencies = [ "base64", "bitflags", "byteorder", + "chrono", "crc", "dotenvy", "etcetera", @@ -1039,6 +1247,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "488e99c397a62007e4229aec669a179816339afc6d2620ca6fa420dbee2e982c" dependencies = [ "atoi", + "chrono", "flume", "form_urlencoded", "futures-channel", @@ -1126,6 +1335,15 @@ dependencies = [ "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" @@ -1161,6 +1379,7 @@ dependencies = [ "libc", "mio", "pin-project-lite", + "signal-hook-registry", "socket2", "tokio-macros", "windows-sys", @@ -1188,6 +1407,20 @@ dependencies = [ "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 = "tracing" version = "0.1.44" @@ -1220,6 +1453,17 @@ 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 = "typenum" version = "1.20.1" @@ -1289,18 +1533,116 @@ 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 = "whoami" version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "998767ef88740d1f5b0682a9c53c24431453923962269c2db68ee43788c5a40d" +[[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.61.2" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 627219d6c..abab0a23e 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -1,5 +1,6 @@ [workspace] members = [ + "riverqueue", "riverqueue-macros", "riverqueue-migrate", ] diff --git a/rust/riverqueue/Cargo.toml b/rust/riverqueue/Cargo.toml new file mode 100644 index 000000000..2049e6830 --- /dev/null +++ b/rust/riverqueue/Cargo.toml @@ -0,0 +1,52 @@ +[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.49.0-alpha.1", default-features = false } +riverqueue-macros = { path = "../riverqueue-macros", version = "=0.49.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.49.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"] } + +[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..926d0a284 --- /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.49.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.49.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..e4ab3b82e --- /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` minor release is matched to the River Go minor release with +the same number: `riverqueue` 0.49 runs alongside River Go 0.49. Patch releases +of either implementation can be mixed within a matched minor release. Upgrade +both implementations together when moving to a new minor release, 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/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..ac9938b14 --- /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_secs(60)), + ] { + 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..c97885e6a --- /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_secs(60)).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..94e344716 --- /dev/null +++ b/rust/riverqueue/src/periodic/cron.rs @@ -0,0 +1,924 @@ +//! 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 generator 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/fixtures/maintenance_values.json b/rust/riverqueue/tests/fixtures/maintenance_values.json new file mode 100644 index 000000000..3e29a6e66 --- /dev/null +++ b/rust/riverqueue/tests/fixtures/maintenance_values.json @@ -0,0 +1,681 @@ +{ + "$schema": "../schema/maintenance-values.schema.json", + "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" + ] + } + ], + "protocol_revision": 1, + "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..a64af845b --- /dev/null +++ b/rust/riverqueue/tests/fixtures/protocol_values.json @@ -0,0 +1,319 @@ +{ + "$schema": "../schema/protocol-values.schema.json", + "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" + } + ], + "protocol_revision": 1, + "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..cbfd04cf5 --- /dev/null +++ b/rust/riverqueue/tests/fixtures/unique_keys.json @@ -0,0 +1,950 @@ +{ + "$schema": "../schema/unique-keys.schema.json", + "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 + } + ], + "protocol_revision": 1, + "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 + } + ] +} From 3300ef4db157d6bb2cdff5d1261ff3634dfefd61 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 4 Oct 2026 18:02:39 -0500 Subject: [PATCH 04/20] test the JobArgs derive against riverqueue Now that `riverqueue` exists, test the derive macro end to end: derived kinds, aliases, insert options, and unique keys, plus compile-fail cases for each attribute error, checked with `trybuild`. --- rust/Cargo.lock | 102 ++++++++++ rust/riverqueue-macros/Cargo.toml | 5 + rust/riverqueue-macros/tests/derive.rs | 181 ++++++++++++++++++ rust/riverqueue-macros/tests/ui.rs | 7 + .../riverqueue-macros/tests/ui/invalid_tag.rs | 10 + .../tests/ui/invalid_tag.stderr | 5 + .../tests/ui/missing_kind.rs | 10 + .../tests/ui/missing_kind.stderr | 5 + .../tests/ui/unique_empty.rs | 10 + .../tests/ui/unique_empty.stderr | 5 + .../tests/ui/unique_field_without_by_args.rs | 11 ++ .../ui/unique_field_without_by_args.stderr | 5 + .../tests/ui/unique_path_numeric.rs | 10 + .../tests/ui/unique_path_numeric.stderr | 5 + .../tests/ui/unique_path_old_syntax.rs | 10 + .../tests/ui/unique_path_old_syntax.stderr | 5 + .../tests/ui/unique_path_unknown_field.rs | 10 + .../tests/ui/unique_path_unknown_field.stderr | 5 + .../tests/ui/unique_period_invalid.rs | 10 + .../tests/ui/unique_period_invalid.stderr | 5 + .../tests/ui/unique_period_too_short.rs | 10 + .../tests/ui/unique_period_too_short.stderr | 5 + .../tests/ui/unique_state_unknown.rs | 10 + .../tests/ui/unique_state_unknown.stderr | 5 + .../ui/unique_states_missing_required.rs | 10 + .../ui/unique_states_missing_required.stderr | 5 + 26 files changed, 461 insertions(+) create mode 100644 rust/riverqueue-macros/tests/derive.rs create mode 100644 rust/riverqueue-macros/tests/ui.rs create mode 100644 rust/riverqueue-macros/tests/ui/invalid_tag.rs create mode 100644 rust/riverqueue-macros/tests/ui/invalid_tag.stderr create mode 100644 rust/riverqueue-macros/tests/ui/missing_kind.rs create mode 100644 rust/riverqueue-macros/tests/ui/missing_kind.stderr create mode 100644 rust/riverqueue-macros/tests/ui/unique_empty.rs create mode 100644 rust/riverqueue-macros/tests/ui/unique_empty.stderr create mode 100644 rust/riverqueue-macros/tests/ui/unique_field_without_by_args.rs create mode 100644 rust/riverqueue-macros/tests/ui/unique_field_without_by_args.stderr create mode 100644 rust/riverqueue-macros/tests/ui/unique_path_numeric.rs create mode 100644 rust/riverqueue-macros/tests/ui/unique_path_numeric.stderr create mode 100644 rust/riverqueue-macros/tests/ui/unique_path_old_syntax.rs create mode 100644 rust/riverqueue-macros/tests/ui/unique_path_old_syntax.stderr create mode 100644 rust/riverqueue-macros/tests/ui/unique_path_unknown_field.rs create mode 100644 rust/riverqueue-macros/tests/ui/unique_path_unknown_field.stderr create mode 100644 rust/riverqueue-macros/tests/ui/unique_period_invalid.rs create mode 100644 rust/riverqueue-macros/tests/ui/unique_period_invalid.stderr create mode 100644 rust/riverqueue-macros/tests/ui/unique_period_too_short.rs create mode 100644 rust/riverqueue-macros/tests/ui/unique_period_too_short.stderr create mode 100644 rust/riverqueue-macros/tests/ui/unique_state_unknown.rs create mode 100644 rust/riverqueue-macros/tests/ui/unique_state_unknown.stderr create mode 100644 rust/riverqueue-macros/tests/ui/unique_states_missing_required.rs create mode 100644 rust/riverqueue-macros/tests/ui/unique_states_missing_required.stderr diff --git a/rust/Cargo.lock b/rust/Cargo.lock index e9c6cd2c1..b8b958cb8 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -466,6 +466,12 @@ dependencies = [ "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" @@ -925,7 +931,10 @@ version = "0.49.0-alpha.1" dependencies = [ "proc-macro2", "quote", + "riverqueue", + "serde", "syn 2.0.119", + "trybuild", ] [[package]] @@ -993,6 +1002,15 @@ dependencies = [ "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" @@ -1315,6 +1333,21 @@ dependencies = [ "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" @@ -1421,6 +1454,45 @@ dependencies = [ "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" @@ -1464,6 +1536,21 @@ dependencies = [ "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" @@ -1584,6 +1671,15 @@ 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", +] + [[package]] name = "windows-core" version = "0.62.2" @@ -1652,6 +1748,12 @@ dependencies = [ "windows-link", ] +[[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" diff --git a/rust/riverqueue-macros/Cargo.toml b/rust/riverqueue-macros/Cargo.toml index 487cc854e..28e922d16 100644 --- a/rust/riverqueue-macros/Cargo.toml +++ b/rust/riverqueue-macros/Cargo.toml @@ -19,5 +19,10 @@ 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/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)))] + | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ From e071c41d064917d91b314bb2d425cbfe2738c576 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 4 Oct 2026 18:02:39 -0500 Subject: [PATCH 05/20] test Rust storage on PostgreSQL and SQLite Add integration tests for the storage layer: the backend contract, job and queue CRUD, list filters and cursors, exact metadata, the protocol fixtures, and parity with the rows River Go writes on each database. A test also fails if a dependency enables `serde_json` features that would change its behavior for the rest of an application. PostgreSQL tests build only with `--cfg river_postgres_tests` and run in a freshly migrated schema with a unique name, so test binaries can share one disposable database. They fail rather than skip when `RIVER_RUST_DATABASE_URL` is unset. --- rust/riverqueue/Cargo.toml | 16 + rust/riverqueue/tests/database_contract.rs | 119 ++ rust/riverqueue/tests/job_list.rs | 347 ++++ rust/riverqueue/tests/metadata_exact.rs | 298 +++ rust/riverqueue/tests/postgres.rs | 2102 ++++++++++++++++++++ rust/riverqueue/tests/postgres_parity.rs | 531 +++++ rust/riverqueue/tests/protocol_fixtures.rs | 126 ++ rust/riverqueue/tests/serde_semantics.rs | 61 + rust/riverqueue/tests/sqlite_parity.rs | 254 +++ rust/riverqueue/tests/sqlite_storage.rs | 748 +++++++ rust/riverqueue/tests/support/mod.rs | 167 ++ 11 files changed, 4769 insertions(+) create mode 100644 rust/riverqueue/tests/database_contract.rs create mode 100644 rust/riverqueue/tests/job_list.rs create mode 100644 rust/riverqueue/tests/metadata_exact.rs create mode 100644 rust/riverqueue/tests/postgres.rs create mode 100644 rust/riverqueue/tests/postgres_parity.rs create mode 100644 rust/riverqueue/tests/protocol_fixtures.rs create mode 100644 rust/riverqueue/tests/serde_semantics.rs create mode 100644 rust/riverqueue/tests/sqlite_parity.rs create mode 100644 rust/riverqueue/tests/sqlite_storage.rs create mode 100644 rust/riverqueue/tests/support/mod.rs diff --git a/rust/riverqueue/Cargo.toml b/rust/riverqueue/Cargo.toml index 2049e6830..6c95d8a43 100644 --- a/rust/riverqueue/Cargo.toml +++ b/rust/riverqueue/Cargo.toml @@ -48,5 +48,21 @@ tracing-subscriber = { version = "0.3.20", default-features = false, features = # Fault-injection tests proxy database connections. tokio = { workspace = true, features = ["io-util", "net", "rt-multi-thread", "signal", "test-util"] } +[[test]] +name = "postgres" +required-features = ["postgres"] + +[[test]] +name = "postgres_parity" +required-features = ["postgres"] + +[[test]] +name = "sqlite_parity" +required-features = ["sqlite"] + +[[test]] +name = "sqlite_storage" +required-features = ["sqlite"] + [lints] workspace = true 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/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/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/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..f58864ab0 --- /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_secs(60)), + ) + .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/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/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_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/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); + } +} From defd7b9dd23da6d93322f3b858744d324524d314 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 4 Oct 2026 18:02:40 -0500 Subject: [PATCH 06/20] test the Rust client runtime Add integration tests for running clients: start and stop lifecycles, runtime configuration, producer lifetimes, stuck jobs, fetching only known kinds, poll-only cancellation, running without leader election, insert notifications, the job and queue handles, and the full SQLite runtime. --- rust/riverqueue/Cargo.toml | 8 + rust/riverqueue/tests/client_handles.rs | 1079 ++++++++++ .../tests/fetch_only_known_kinds.rs | 139 ++ rust/riverqueue/tests/insert_notifications.rs | 457 ++++ .../tests/leader_election_disabled.rs | 384 ++++ rust/riverqueue/tests/lifecycle.rs | 475 +++++ .../tests/poll_only_cancellation.rs | 138 ++ rust/riverqueue/tests/producer_lifetime.rs | 769 +++++++ rust/riverqueue/tests/runtime_config.rs | 786 +++++++ rust/riverqueue/tests/sqlite_runtime.rs | 1856 +++++++++++++++++ rust/riverqueue/tests/stuck_jobs.rs | 141 ++ 11 files changed, 6232 insertions(+) create mode 100644 rust/riverqueue/tests/client_handles.rs create mode 100644 rust/riverqueue/tests/fetch_only_known_kinds.rs create mode 100644 rust/riverqueue/tests/insert_notifications.rs create mode 100644 rust/riverqueue/tests/leader_election_disabled.rs create mode 100644 rust/riverqueue/tests/lifecycle.rs create mode 100644 rust/riverqueue/tests/poll_only_cancellation.rs create mode 100644 rust/riverqueue/tests/producer_lifetime.rs create mode 100644 rust/riverqueue/tests/runtime_config.rs create mode 100644 rust/riverqueue/tests/sqlite_runtime.rs create mode 100644 rust/riverqueue/tests/stuck_jobs.rs diff --git a/rust/riverqueue/Cargo.toml b/rust/riverqueue/Cargo.toml index 6c95d8a43..8745808df 100644 --- a/rust/riverqueue/Cargo.toml +++ b/rust/riverqueue/Cargo.toml @@ -48,6 +48,10 @@ tracing-subscriber = { version = "0.3.20", default-features = false, features = # Fault-injection tests proxy database connections. tokio = { workspace = true, features = ["io-util", "net", "rt-multi-thread", "signal", "test-util"] } +[[test]] +name = "lifecycle" +required-features = ["sqlite"] + [[test]] name = "postgres" required-features = ["postgres"] @@ -60,6 +64,10 @@ required-features = ["postgres"] name = "sqlite_parity" required-features = ["sqlite"] +[[test]] +name = "sqlite_runtime" +required-features = ["sqlite"] + [[test]] name = "sqlite_storage" required-features = ["sqlite"] 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/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/insert_notifications.rs b/rust/riverqueue/tests/insert_notifications.rs new file mode 100644 index 000000000..5ac816bd5 --- /dev/null +++ b/rust/riverqueue/tests/insert_notifications.rs @@ -0,0 +1,457 @@ +//! 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)); + + client.insert(NotificationArgs {}).await.unwrap(); + 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)); + + client.insert(NotificationArgs {}).await.unwrap(); + 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/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/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/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/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/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/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; +} From 7a5c70846530b5d046d9828d5605e8d501a6d679 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 4 Oct 2026 18:02:40 -0500 Subject: [PATCH 07/20] test Rust clients through database faults Add integration tests that run clients against malformed rows and database faults on PostgreSQL and SQLite: proxied connections that are dropped and refused while clients run, after which producers, completion, notifications, and leadership must recover. Clients must also detect a PostgreSQL server that looks like YugabyteDB and fall back to polling. --- rust/riverqueue/tests/resilience_postgres.rs | 1347 ++++++++++++++++++ rust/riverqueue/tests/resilience_sqlite.rs | 882 ++++++++++++ rust/riverqueue/tests/yugabyte.rs | 295 ++++ 3 files changed, 2524 insertions(+) create mode 100644 rust/riverqueue/tests/resilience_postgres.rs create mode 100644 rust/riverqueue/tests/resilience_sqlite.rs create mode 100644 rust/riverqueue/tests/yugabyte.rs diff --git a/rust/riverqueue/tests/resilience_postgres.rs b/rust/riverqueue/tests/resilience_postgres.rs new file mode 100644 index 000000000..d359eb4bf --- /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_secs(60)), + ) + .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/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:#?}" + ); +} From 517846d0a54d19dc36c0c3006d6e6b37d4b525a2 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 4 Oct 2026 18:02:40 -0500 Subject: [PATCH 08/20] test Rust extension points Add integration tests for what add-on crates build on: work middleware, hooks, and error handlers and their ordering; extension services and the client's stop order; producer sessions and the checks River applies to their claims; peer attempts that River owns until each outcome persists; prepared insertion of stored jobs; filtered deletion of finalized jobs; and batched insert interception. Tests also cover requests run in a caller's transaction: they open no savepoint, so an extension step, insert middleware, or decode failure after River's write leaves that write in the transaction for the caller to roll back, a failed statement aborts a PostgreSQL transaction, and every write carries the caller's transaction ID. --- rust/riverqueue/Cargo.toml | 4 + rust/riverqueue/tests/caller_transactions.rs | 706 ++++++++ rust/riverqueue/tests/extension_seams.rs | 632 +++++++ rust/riverqueue/tests/extension_services.rs | 394 +++++ rust/riverqueue/tests/peer_attempts.rs | 1645 ++++++++++++++++++ rust/riverqueue/tests/prepared_insertion.rs | 424 +++++ rust/riverqueue/tests/producer_sessions.rs | 886 ++++++++++ rust/riverqueue/tests/work_extensions.rs | 380 ++++ 8 files changed, 5071 insertions(+) create mode 100644 rust/riverqueue/tests/caller_transactions.rs create mode 100644 rust/riverqueue/tests/extension_seams.rs create mode 100644 rust/riverqueue/tests/extension_services.rs create mode 100644 rust/riverqueue/tests/peer_attempts.rs create mode 100644 rust/riverqueue/tests/prepared_insertion.rs create mode 100644 rust/riverqueue/tests/producer_sessions.rs create mode 100644 rust/riverqueue/tests/work_extensions.rs diff --git a/rust/riverqueue/Cargo.toml b/rust/riverqueue/Cargo.toml index 8745808df..fcfacef49 100644 --- a/rust/riverqueue/Cargo.toml +++ b/rust/riverqueue/Cargo.toml @@ -52,6 +52,10 @@ tokio = { workspace = true, features = ["io-util", "net", "rt-multi-thread", "si name = "lifecycle" required-features = ["sqlite"] +[[test]] +name = "work_extensions" +required-features = ["sqlite"] + [[test]] name = "postgres" required-features = ["postgres"] 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/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/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/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_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/work_extensions.rs b/rust/riverqueue/tests/work_extensions.rs new file mode 100644 index 000000000..32e24acad --- /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_secs(60))) + } + 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; +} From 362d0554be9c627d5e5d5767b70da0de34469c4f Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 4 Oct 2026 18:02:40 -0500 Subject: [PATCH 09/20] add Rust examples Add runnable examples for a basic worker, graceful shutdown, cancellation, transactional completion, unique and periodic jobs, event subscriptions, custom schemas, SQLite, and a deployment where Go and Rust clients share one database. --- rust/riverqueue/Cargo.toml | 40 ++++++ rust/riverqueue/examples/basic_worker.rs | 51 ++++++++ rust/riverqueue/examples/cancellation.rs | 64 ++++++++++ rust/riverqueue/examples/custom_schema.rs | 39 ++++++ rust/riverqueue/examples/events.rs | 94 +++++++++++++++ rust/riverqueue/examples/graceful_shutdown.rs | 63 ++++++++++ rust/riverqueue/examples/mixed_go_rust.rs | 114 ++++++++++++++++++ rust/riverqueue/examples/periodic.rs | 84 +++++++++++++ rust/riverqueue/examples/sqlite.rs | 74 ++++++++++++ rust/riverqueue/examples/transactions.rs | 96 +++++++++++++++ rust/riverqueue/examples/unique.rs | 85 +++++++++++++ 11 files changed, 804 insertions(+) create mode 100644 rust/riverqueue/examples/basic_worker.rs create mode 100644 rust/riverqueue/examples/cancellation.rs create mode 100644 rust/riverqueue/examples/custom_schema.rs create mode 100644 rust/riverqueue/examples/events.rs create mode 100644 rust/riverqueue/examples/graceful_shutdown.rs create mode 100644 rust/riverqueue/examples/mixed_go_rust.rs create mode 100644 rust/riverqueue/examples/periodic.rs create mode 100644 rust/riverqueue/examples/sqlite.rs create mode 100644 rust/riverqueue/examples/transactions.rs create mode 100644 rust/riverqueue/examples/unique.rs diff --git a/rust/riverqueue/Cargo.toml b/rust/riverqueue/Cargo.toml index fcfacef49..cbfd4fa0a 100644 --- a/rust/riverqueue/Cargo.toml +++ b/rust/riverqueue/Cargo.toml @@ -48,6 +48,46 @@ tracing-subscriber = { version = "0.3.20", default-features = false, features = # 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"] 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..66b579960 --- /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_secs(60)) => { + 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(()) +} From 83f9f08d474fc1f968cc15c2875182eec9611d23 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 4 Oct 2026 18:02:40 -0500 Subject: [PATCH 10/20] add Rust test helpers Add `riverqueue-test`: assertions that check which jobs a test's code inserted, with optional expected properties and variants that read through an open transaction, and helpers that run a worker once with or without a database. --- rust/Cargo.lock | 14 + rust/Cargo.toml | 1 + rust/riverqueue-test/Cargo.toml | 35 ++ rust/riverqueue-test/LICENSE | 1 + rust/riverqueue-test/README.md | 109 +++++ rust/riverqueue-test/src/lib.rs | 510 +++++++++++++++++++++ rust/riverqueue-test/src/require.rs | 661 ++++++++++++++++++++++++++++ 7 files changed, 1331 insertions(+) create mode 100644 rust/riverqueue-test/Cargo.toml create mode 120000 rust/riverqueue-test/LICENSE create mode 100644 rust/riverqueue-test/README.md create mode 100644 rust/riverqueue-test/src/lib.rs create mode 100644 rust/riverqueue-test/src/require.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index b8b958cb8..65e0555d9 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -947,6 +947,20 @@ dependencies = [ "tokio", ] +[[package]] +name = "riverqueue-test" +version = "0.49.0-alpha.1" +dependencies = [ + "chrono", + "riverqueue", + "riverqueue-migrate", + "serde", + "serde_json", + "sqlx", + "tokio", + "tokio-util", +] + [[package]] name = "rustversion" version = "1.0.23" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index abab0a23e..849355098 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -3,6 +3,7 @@ members = [ "riverqueue", "riverqueue-macros", "riverqueue-migrate", + "riverqueue-test", ] resolver = "3" diff --git a/rust/riverqueue-test/Cargo.toml b/rust/riverqueue-test/Cargo.toml new file mode 100644 index 000000000..166a76d3f --- /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.49.0-alpha.1", default-features = false } +serde_json.workspace = true +tokio-util.workspace = true + +[dev-dependencies] +riverqueue-migrate = { path = "../riverqueue-migrate", version = "=0.49.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; + } +} From d366e514eb0eb54c4722115b88ecc92b0aaec4d6 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 4 Oct 2026 18:02:41 -0500 Subject: [PATCH 11/20] add the riverqueue command Add `riverqueue-cli`, which installs a `riverqueue` binary that migrates PostgreSQL and SQLite databases like `river migrate-*` and benchmarks worker throughput and end-to-end latency like `river bench`. --- rust/Cargo.lock | 187 +++++++- rust/Cargo.toml | 1 + rust/riverqueue-cli/Cargo.toml | 34 ++ rust/riverqueue-cli/LICENSE | 1 + rust/riverqueue-cli/README.md | 36 ++ rust/riverqueue-cli/src/bench.rs | 741 +++++++++++++++++++++++++++++ rust/riverqueue-cli/src/main.rs | 65 +++ rust/riverqueue-cli/src/migrate.rs | 197 ++++++++ 8 files changed, 1255 insertions(+), 7 deletions(-) create mode 100644 rust/riverqueue-cli/Cargo.toml create mode 120000 rust/riverqueue-cli/LICENSE create mode 100644 rust/riverqueue-cli/README.md create mode 100644 rust/riverqueue-cli/src/bench.rs create mode 100644 rust/riverqueue-cli/src/main.rs create mode 100644 rust/riverqueue-cli/src/migrate.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 65e0555d9..729520912 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -306,7 +306,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -316,7 +316,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "de48cc4d1c1d97a20fd819def54b890cadde72ed3ad0c614822a0a433361be96" dependencies = [ "cfg-if", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -454,6 +454,17 @@ dependencies = [ "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" @@ -759,7 +770,7 @@ checksum = "30d65c71f1ce40ab09135ce117d742b9f8a19ff91a41a8b57ed50bc2de59c427" dependencies = [ "libc", "wasi", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -882,7 +893,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" dependencies = [ "chacha20", - "getrandom", + "getrandom 0.4.3", "rand_core", ] @@ -901,6 +912,20 @@ 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.49.0-alpha.1" @@ -925,6 +950,18 @@ dependencies = [ "tracing-subscriber", ] +[[package]] +name = "riverqueue-cli" +version = "0.49.0-alpha.1" +dependencies = [ + "riverqueue", + "riverqueue-migrate", + "serde", + "sqlx", + "tokio", + "tokio-util", +] + [[package]] name = "riverqueue-macros" version = "0.49.0-alpha.1" @@ -961,6 +998,40 @@ dependencies = [ "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" @@ -1111,7 +1182,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" dependencies = [ "libc", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1160,6 +1231,7 @@ dependencies = [ "log", "memchr", "percent-encoding", + "rustls", "serde", "serde_json", "sha2 0.10.9", @@ -1169,6 +1241,7 @@ dependencies = [ "tokio-stream", "tracing", "url", + "webpki-roots", ] [[package]] @@ -1314,6 +1387,12 @@ dependencies = [ "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" @@ -1429,7 +1508,7 @@ dependencies = [ "signal-hook-registry", "socket2", "tokio-macros", - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1598,6 +1677,12 @@ 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" @@ -1679,6 +1764,15 @@ 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" @@ -1691,7 +1785,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys", + "windows-sys 0.61.2", ] [[package]] @@ -1753,6 +1847,15 @@ 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" @@ -1762,6 +1865,70 @@ 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" @@ -1818,6 +1985,12 @@ dependencies = [ "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" diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 849355098..935266d2f 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -1,6 +1,7 @@ [workspace] members = [ "riverqueue", + "riverqueue-cli", "riverqueue-macros", "riverqueue-migrate", "riverqueue-test", diff --git a/rust/riverqueue-cli/Cargo.toml b/rust/riverqueue-cli/Cargo.toml new file mode 100644 index 000000000..b509e8604 --- /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.49.0-alpha.1", default-features = false } +riverqueue-migrate = { path = "../riverqueue-migrate", version = "=0.49.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]" +} From 83cd495bed21fdf1bd8cab68224615245f8727cf Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 4 Oct 2026 18:02:41 -0500 Subject: [PATCH 12/20] run the Rust crates in CI --- .github/workflows/rust.yaml | 140 ++++++++++++++++++++++++++++++++++++ Makefile | 98 +++++++++++++++++++++++-- rust/deny.toml | 29 ++++++++ 3 files changed, 261 insertions(+), 6 deletions(-) create mode 100644 .github/workflows/rust.yaml create mode 100644 rust/deny.toml diff --git a/.github/workflows/rust.yaml b/.github/workflows/rust.yaml new file mode 100644 index 000000000..a09cd58f5 --- /dev/null +++ b/.github/workflows/rust.yaml @@ -0,0 +1,140 @@ +name: Rust + +on: + push: + branches: + - master + pull_request: + +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 + + msrv: + 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/Makefile b/Makefile index 85c30eef2..e0693717e 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 @@ -31,9 +33,9 @@ generate/rust-migrations: ## Sync database migrations and hashes to Rust .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 @@ -47,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 @@ -77,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 @@ -90,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 @@ -109,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 @@ -140,6 +226,6 @@ verify/rust-migrations: ## Verify Rust migrations and protocol hashes .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/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" From 53f4fa8ad00c3e081f1528d613de90c6d7d2cfa1 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 4 Oct 2026 18:02:41 -0500 Subject: [PATCH 13/20] keep Rust and workflow dependencies updated Have Dependabot update the Rust workspace's dependencies weekly, grouping minor and patch updates after a seven-day cooldown, and the GitHub Actions the workflows use. --- .github/dependabot.yml | 15 +++++++++++++++ 1 file changed, 15 insertions(+) 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: - "**/*" From f22962d06d5e8bea0f094f2611266f1579a832c6 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Sun, 4 Oct 2026 18:02:41 -0500 Subject: [PATCH 14/20] document the Rust workspace Add a workspace README listing the crates and how to run their checks, benchmark, and PostgreSQL tests, and a changelog for the crates, which are versioned and released together. --- rust/CHANGELOG.md | 32 +++++++++++++++++++ rust/README.md | 78 +++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 110 insertions(+) create mode 100644 rust/CHANGELOG.md create mode 100644 rust/README.md 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/README.md b/rust/README.md new file mode 100644 index 000000000..e276f422c --- /dev/null +++ b/rust/README.md @@ -0,0 +1,78 @@ +# 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. Shared cross-language fixtures live in +[`../conformance`](../conformance). + +## 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. +- `riverqueue-conformance`: private verification package. + +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. The conformance performance gate remains the +reproducible Go/Rust comparison across enqueue-only, worker-only, and mixed +workloads. + +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 +``` + +`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. From b069c1f75e464d7a2eb050ecac607c5daefbc958 Mon Sep 17 00:00:00 2001 From: Brandur Date: Mon, 5 Oct 2026 13:15:47 -0500 Subject: [PATCH 15/20] Decouple the Rust port from conformance tooling --- .github/workflows/ci.yaml | 7 ++++ Makefile | 4 +- internal/cmd/syncrustmigrations/main.go | 41 +------------------ rust/README.md | 13 +++--- rust/riverqueue/src/periodic/cron.rs | 3 +- rust/riverqueue/src/unique.rs | 6 +-- .../tests/fixtures/maintenance_values.json | 2 - .../tests/fixtures/protocol_values.json | 2 - .../tests/fixtures/unique_keys.json | 2 - 9 files changed, 19 insertions(+), 61 deletions(-) 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/Makefile b/Makefile index e0693717e..a55a88ee5 100644 --- a/Makefile +++ b/Makefile @@ -28,7 +28,7 @@ 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 and hashes to Rust +generate/rust-migrations: ## Sync database migrations to Rust go run ./internal/cmd/syncrustmigrations .PHONY: generate/sqlc @@ -221,7 +221,7 @@ verify/migrations: ## Verify synced migrations diff -qr riverdriver/riverpgxv5/migration riverdriver/riverdatabasesql/migration .PHONY: verify/rust-migrations -verify/rust-migrations: ## Verify Rust migrations and protocol hashes +verify/rust-migrations: ## Verify Rust migrations match the canonical migrations go run ./internal/cmd/syncrustmigrations -check .PHONY: verify/sqlc diff --git a/internal/cmd/syncrustmigrations/main.go b/internal/cmd/syncrustmigrations/main.go index 8846a78dc..576738863 100644 --- a/internal/cmd/syncrustmigrations/main.go +++ b/internal/cmd/syncrustmigrations/main.go @@ -1,13 +1,9 @@ // Command syncrustmigrations mirrors River's canonical database migrations -// into the publishable Rust migration crate and records their hashes for -// cross-language conformance. +// into the publishable Rust migration crate. package main import ( "bytes" - "crypto/sha256" - "encoding/hex" - "encoding/json" "flag" "fmt" "os" @@ -18,20 +14,7 @@ import ( type database struct { canonicalDir string - manifestPath string mirrorDir string - name string -} - -type manifest struct { - Database string `json:"database"` - Files []manifestFile `json:"files"` - Line string `json:"line"` -} - -type manifestFile struct { - Path string `json:"path"` - SHA256 string `json:"sha256"` } func main() { @@ -41,15 +24,11 @@ func main() { databases := []database{ { canonicalDir: "riverdriver/riverpgxv5/migration/main", - manifestPath: "conformance/migrations.json", mirrorDir: "rust/riverqueue-migrate/migrations/main", - name: "postgres", }, { canonicalDir: "riverdriver/riversqlite/migration/main", - manifestPath: "conformance/migrations-sqlite.json", mirrorDir: "rust/riverqueue-migrate/migrations/sqlite/main", - name: "sqlite", }, } for _, database := range databases { @@ -71,19 +50,12 @@ func syncDatabase(database database, check bool) { } slices.Sort(names) - generatedManifest := manifest{Database: database.name, Line: "main"} for _, name := range names { sourcePath := filepath.Join(database.canonicalDir, name) contents, err := os.ReadFile(sourcePath) if err != nil { fatal(err) } - hash := sha256.Sum256(contents) - generatedManifest.Files = append(generatedManifest.Files, manifestFile{ - Path: filepath.ToSlash(sourcePath), - SHA256: hex.EncodeToString(hash[:]), - }) - mirrorPath := filepath.Join(database.mirrorDir, name) if check { checkFile(mirrorPath, contents) @@ -92,17 +64,6 @@ func syncDatabase(database database, check bool) { } } removeStaleMirrors(database.mirrorDir, names, check) - - manifestContents, err := json.MarshalIndent(&generatedManifest, "", " ") - if err != nil { - fatal(err) - } - manifestContents = append(manifestContents, '\n') - if check { - checkFile(database.manifestPath, manifestContents) - } else { - writeFile(database.manifestPath, manifestContents) - } } func removeStaleMirrors(directory string, expected []string, check bool) { diff --git a/rust/README.md b/rust/README.md index e276f422c..4bf34c0db 100644 --- a/rust/README.md +++ b/rust/README.md @@ -3,8 +3,7 @@ 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. Shared cross-language fixtures live in -[`../conformance`](../conformance). +preview. ## Workspace crates @@ -15,7 +14,6 @@ preview. Shared cross-language fixtures live in - `riverqueue-cli`: the `riverqueue` command-line program for migrations and benchmarks. - `riverqueue-test`: typed fixtures and worker-test helpers. -- `riverqueue-conformance`: private verification package. The API uses a caller-owned SQLx pool, Tokio, typed workers, and `CancellationToken`. `Client` isn't generic over the database: it accepts a @@ -58,9 +56,7 @@ make bench/rust DATABASE_URL=postgres://localhost/river_bench \ 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. The conformance performance gate remains the -reproducible Go/Rust comparison across enqueue-only, worker-only, and mixed -workloads. +--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 @@ -71,6 +67,11 @@ 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 diff --git a/rust/riverqueue/src/periodic/cron.rs b/rust/riverqueue/src/periodic/cron.rs index 94e344716..eb027b280 100644 --- a/rust/riverqueue/src/periodic/cron.rs +++ b/rust/riverqueue/src/periodic/cron.rs @@ -742,8 +742,7 @@ mod tests { fn assert_matches_go(case: &CronCase) { let schedule = CronSchedule::parse(&case.expression) .unwrap_or_else(|error| panic!("{}: {error}", case.name)); - // The generator records five occurrences, stopping early at Go's - // zero time. + // 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 { diff --git a/rust/riverqueue/src/unique.rs b/rust/riverqueue/src/unique.rs index 8f774241f..33e5602cc 100644 --- a/rust/riverqueue/src/unique.rs +++ b/rust/riverqueue/src/unique.rs @@ -479,7 +479,6 @@ mod tests { #[derive(Deserialize)] struct Fixture { cases: Vec, - protocol_revision: u32, typed_only_cases: Vec, } @@ -538,10 +537,7 @@ mod tests { } fn fixture() -> Fixture { - let fixture: Fixture = - serde_json::from_str(include_str!("../tests/fixtures/unique_keys.json")).unwrap(); - assert_eq!(fixture.protocol_revision, 1); - fixture + serde_json::from_str(include_str!("../tests/fixtures/unique_keys.json")).unwrap() } fn golden(name: &str) -> FixtureCase { diff --git a/rust/riverqueue/tests/fixtures/maintenance_values.json b/rust/riverqueue/tests/fixtures/maintenance_values.json index 3e29a6e66..1b69ceb1c 100644 --- a/rust/riverqueue/tests/fixtures/maintenance_values.json +++ b/rust/riverqueue/tests/fixtures/maintenance_values.json @@ -1,5 +1,4 @@ { - "$schema": "../schema/maintenance-values.schema.json", "cron_cases": [ { "expression": "* * * * *", @@ -561,7 +560,6 @@ ] } ], - "protocol_revision": 1, "snooze_counters": [ { "expected_snoozes": 1, diff --git a/rust/riverqueue/tests/fixtures/protocol_values.json b/rust/riverqueue/tests/fixtures/protocol_values.json index a64af845b..15fb9a779 100644 --- a/rust/riverqueue/tests/fixtures/protocol_values.json +++ b/rust/riverqueue/tests/fixtures/protocol_values.json @@ -1,5 +1,4 @@ { - "$schema": "../schema/protocol-values.schema.json", "attempt_error": { "at": "2026-01-02T03:04:05.6789Z", "attempt": 3, @@ -214,7 +213,6 @@ "topic": "river_control" } ], - "protocol_revision": 1, "reserved_metadata_keys": [ { "applicability": "protocol_visible", diff --git a/rust/riverqueue/tests/fixtures/unique_keys.json b/rust/riverqueue/tests/fixtures/unique_keys.json index cbfd04cf5..066f4d86e 100644 --- a/rust/riverqueue/tests/fixtures/unique_keys.json +++ b/rust/riverqueue/tests/fixtures/unique_keys.json @@ -1,5 +1,4 @@ { - "$schema": "../schema/unique-keys.schema.json", "cases": [ { "args": {}, @@ -894,7 +893,6 @@ "selected_unique_paths": null } ], - "protocol_revision": 1, "typed_only_cases": [ { "args": { From a0ff6904d92072d61c97da913ff1842086941a49 Mon Sep 17 00:00:00 2001 From: Brandur Date: Mon, 5 Oct 2026 13:15:47 -0500 Subject: [PATCH 16/20] Use minute durations in Rust tests and examples --- rust/riverqueue/examples/graceful_shutdown.rs | 2 +- rust/riverqueue/src/job.rs | 2 +- rust/riverqueue/src/maintenance/tests.rs | 2 +- rust/riverqueue/tests/postgres_parity.rs | 2 +- rust/riverqueue/tests/resilience_postgres.rs | 2 +- rust/riverqueue/tests/work_extensions.rs | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/rust/riverqueue/examples/graceful_shutdown.rs b/rust/riverqueue/examples/graceful_shutdown.rs index 66b579960..87cff07ad 100644 --- a/rust/riverqueue/examples/graceful_shutdown.rs +++ b/rust/riverqueue/examples/graceful_shutdown.rs @@ -25,7 +25,7 @@ async fn generate_report( // again without using up its attempt. tokio::select! { () = context.cancellation_token().cancelled() => Err(WorkCancelled.into()), - () = tokio::time::sleep(Duration::from_secs(60)) => { + () = tokio::time::sleep(Duration::from_mins(1)) => { println!("generated report {}", job.args.report_id); Ok(WorkOutcome::Complete) } diff --git a/rust/riverqueue/src/job.rs b/rust/riverqueue/src/job.rs index ac9938b14..a099fcf17 100644 --- a/rust/riverqueue/src/job.rs +++ b/rust/riverqueue/src/job.rs @@ -933,7 +933,7 @@ mod tests { for valid in [ kind_only.clone().with_by_args(true), kind_only.clone().with_by_queue(true), - kind_only.with_by_period(Duration::from_secs(60)), + kind_only.with_by_period(Duration::from_mins(1)), ] { assert!(valid.validate().is_ok(), "{valid:?}"); } diff --git a/rust/riverqueue/src/maintenance/tests.rs b/rust/riverqueue/src/maintenance/tests.rs index c97885e6a..9f7993307 100644 --- a/rust/riverqueue/src/maintenance/tests.rs +++ b/rust/riverqueue/src/maintenance/tests.rs @@ -781,7 +781,7 @@ async fn periodic_jobs_use_expected_run_time_and_uniqueness() { let periodic = client.periodic_jobs(); periodic .add(PeriodicJob::conditional( - IntervalSchedule::new(Duration::from_secs(60)).unwrap(), + IntervalSchedule::new(Duration::from_mins(1)).unwrap(), || { Some(( NoTimeoutArgs {}, diff --git a/rust/riverqueue/tests/postgres_parity.rs b/rust/riverqueue/tests/postgres_parity.rs index f58864ab0..9c2b40d60 100644 --- a/rust/riverqueue/tests/postgres_parity.rs +++ b/rust/riverqueue/tests/postgres_parity.rs @@ -337,7 +337,7 @@ async fn schema_names_are_quoted_like_go() { ) .queue( "default", - QueueConfig::new(1).with_fetch_poll_interval(Duration::from_secs(60)), + QueueConfig::new(1).with_fetch_poll_interval(Duration::from_mins(1)), ) .workers(noop_workers()) .build() diff --git a/rust/riverqueue/tests/resilience_postgres.rs b/rust/riverqueue/tests/resilience_postgres.rs index d359eb4bf..21128e1b5 100644 --- a/rust/riverqueue/tests/resilience_postgres.rs +++ b/rust/riverqueue/tests/resilience_postgres.rs @@ -1038,7 +1038,7 @@ async fn client_survives_database_outage_and_catches_up() { "default", QueueConfig::new(4) .with_fetch_cooldown(Duration::from_millis(1)) - .with_fetch_poll_interval(Duration::from_secs(60)), + .with_fetch_poll_interval(Duration::from_mins(1)), ) .build() .unwrap(); diff --git a/rust/riverqueue/tests/work_extensions.rs b/rust/riverqueue/tests/work_extensions.rs index 32e24acad..1b9af082e 100644 --- a/rust/riverqueue/tests/work_extensions.rs +++ b/rust/riverqueue/tests/work_extensions.rs @@ -148,7 +148,7 @@ impl Hook for TracingHook { .downcast_ref::() .is_some() => { - Ok(WorkOutcome::Snooze(Duration::from_secs(60))) + Ok(WorkOutcome::Snooze(Duration::from_mins(1))) } result => result, } From dcfd2e57f2637bd00b4d63cb3aa62fe9161520ad Mon Sep 17 00:00:00 2001 From: Brandur Date: Mon, 5 Oct 2026 14:03:10 -0500 Subject: [PATCH 17/20] Fix flaky tests From Codex: > It assumed two database inserts would finish within a **50 ms > cooldown**. Under CI load, the cooldown expired, so emitting two > notifications was correct. --- rust/riverqueue/tests/insert_notifications.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/rust/riverqueue/tests/insert_notifications.rs b/rust/riverqueue/tests/insert_notifications.rs index 5ac816bd5..cc75475ce 100644 --- a/rust/riverqueue/tests/insert_notifications.rs +++ b/rust/riverqueue/tests/insert_notifications.rs @@ -189,7 +189,8 @@ mod postgres { let mut notifications = Notifications::listen(&schema).await; let client = client(&schema, Duration::from_millis(50)); - client.insert(NotificationArgs {}).await.unwrap(); + // 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; @@ -349,7 +350,8 @@ mod sqlite { let mut notifications = Notifications::default(); let client = client(&pool, Duration::from_millis(50)); - client.insert(NotificationArgs {}).await.unwrap(); + // 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; From f1beb9805ee98d9a2c6f93f7cf83d606d2415bec Mon Sep 17 00:00:00 2001 From: Brandur Date: Mon, 5 Oct 2026 14:15:59 -0500 Subject: [PATCH 18/20] Run Rust workflows conditionally on changes to `rust/` --- .github/workflows/rust.yaml | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/.github/workflows/rust.yaml b/.github/workflows/rust.yaml index a09cd58f5..5d8325749 100644 --- a/.github/workflows/rust.yaml +++ b/.github/workflows/rust.yaml @@ -1,10 +1,26 @@ 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 @@ -68,7 +84,7 @@ jobs: - name: Documentation as docs.rs builds it run: make doc/rust/docsrs - msrv: + rust_versions: runs-on: ubuntu-latest timeout-minutes: 20 strategy: From 75d48c5e85d463e380138aabae44ebd2a0a6363b Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Mon, 5 Oct 2026 14:11:40 -0500 Subject: [PATCH 19/20] keep rust/ out of the root Go module The root module's zip, which is what proxy.golang.org serves for every `go get github.com/riverqueue/river`, includes every file under the module root except nested modules. Without a module boundary in `rust/`, it ships the whole Rust workspace (187 files) to Go users, and `go test ./...` and other `./...` patterns walk into `rust/target`. Add a stub `go.mod` that makes `rust/` its own module so the Go tooling and the published module zip leave it out. --- rust/go.mod | 4 ++++ 1 file changed, 4 insertions(+) create mode 100644 rust/go.mod 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 From 7d794fda3324723c8ad4f727b0a17a67e7d86132 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Mon, 5 Oct 2026 14:11:56 -0500 Subject: [PATCH 20/20] version the Rust crates 0.50.0-alpha.1 Rust and JavaScript previews are versioned as a prerelease of the next River Go release. River v0.49.0 has shipped, so the crates move from `0.49.0-alpha.1` to `0.50.0-alpha.1`, a version that semver orders after the release they're compatible with. Update the workspace version, the exact requirements between workspace crates, `Cargo.lock`, the README install snippets, and the mixed-deployments note on which River Go release a preview works with. --- rust/Cargo.lock | 10 +++++----- rust/Cargo.toml | 2 +- rust/riverqueue-cli/Cargo.toml | 4 ++-- rust/riverqueue-test/Cargo.toml | 4 ++-- rust/riverqueue/Cargo.toml | 6 +++--- rust/riverqueue/README.md | 4 ++-- rust/riverqueue/docs/mixed-deployments.md | 10 +++++----- 7 files changed, 20 insertions(+), 20 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 729520912..c61ec9e09 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -928,7 +928,7 @@ dependencies = [ [[package]] name = "riverqueue" -version = "0.49.0-alpha.1" +version = "0.50.0-alpha.1" dependencies = [ "anyhow", "async-trait", @@ -952,7 +952,7 @@ dependencies = [ [[package]] name = "riverqueue-cli" -version = "0.49.0-alpha.1" +version = "0.50.0-alpha.1" dependencies = [ "riverqueue", "riverqueue-migrate", @@ -964,7 +964,7 @@ dependencies = [ [[package]] name = "riverqueue-macros" -version = "0.49.0-alpha.1" +version = "0.50.0-alpha.1" dependencies = [ "proc-macro2", "quote", @@ -976,7 +976,7 @@ dependencies = [ [[package]] name = "riverqueue-migrate" -version = "0.49.0-alpha.1" +version = "0.50.0-alpha.1" dependencies = [ "serde_json", "sqlx", @@ -986,7 +986,7 @@ dependencies = [ [[package]] name = "riverqueue-test" -version = "0.49.0-alpha.1" +version = "0.50.0-alpha.1" dependencies = [ "chrono", "riverqueue", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 935266d2f..f0a30f705 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -15,7 +15,7 @@ homepage = "https://riverqueue.com" license = "MPL-2.0" repository = "https://github.com/riverqueue/river" rust-version = "1.95" -version = "0.49.0-alpha.1" +version = "0.50.0-alpha.1" [workspace.dependencies] async-trait = "0.1.92" diff --git a/rust/riverqueue-cli/Cargo.toml b/rust/riverqueue-cli/Cargo.toml index b509e8604..34f27a7e3 100644 --- a/rust/riverqueue-cli/Cargo.toml +++ b/rust/riverqueue-cli/Cargo.toml @@ -23,8 +23,8 @@ path = "src/main.rs" doc = false [dependencies] -riverqueue = { path = "../riverqueue", version = "=0.49.0-alpha.1", default-features = false } -riverqueue-migrate = { path = "../riverqueue-migrate", version = "=0.49.0-alpha.1", default-features = false } +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"] } diff --git a/rust/riverqueue-test/Cargo.toml b/rust/riverqueue-test/Cargo.toml index 166a76d3f..ce72e2415 100644 --- a/rust/riverqueue-test/Cargo.toml +++ b/rust/riverqueue-test/Cargo.toml @@ -21,12 +21,12 @@ sqlite = ["riverqueue/sqlite"] [dependencies] chrono.workspace = true -riverqueue = { path = "../riverqueue", version = "=0.49.0-alpha.1", default-features = false } +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.49.0-alpha.1", default-features = false, features = ["sqlite"] } +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 diff --git a/rust/riverqueue/Cargo.toml b/rust/riverqueue/Cargo.toml index cbfd4fa0a..9febe6795 100644 --- a/rust/riverqueue/Cargo.toml +++ b/rust/riverqueue/Cargo.toml @@ -28,8 +28,8 @@ 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.49.0-alpha.1", default-features = false } -riverqueue-macros = { path = "../riverqueue-macros", version = "=0.49.0-alpha.1" } +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 @@ -42,7 +42,7 @@ tracing.workspace = true [dev-dependencies] anyhow = "1.0.104" -riverqueue-migrate = { path = "../riverqueue-migrate", version = "=0.49.0-alpha.1", features = ["postgres", "sqlite"] } +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. diff --git a/rust/riverqueue/README.md b/rust/riverqueue/README.md index 926d0a284..e988e4dfd 100644 --- a/rust/riverqueue/README.md +++ b/rust/riverqueue/README.md @@ -14,7 +14,7 @@ covers running both against one database. ```toml [dependencies] -riverqueue = "0.49.0-alpha.1" +riverqueue = "0.50.0-alpha.1" serde = { version = "1", features = ["derive"] } serde_json = "1" tokio = { version = "1", features = ["macros", "rt-multi-thread", "signal"] } @@ -31,7 +31,7 @@ Rust version is 1.95. | `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.49.0-alpha.1", default-features = false, features = ["sqlite"] }`. +`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 diff --git a/rust/riverqueue/docs/mixed-deployments.md b/rust/riverqueue/docs/mixed-deployments.md index e4ab3b82e..ace84b343 100644 --- a/rust/riverqueue/docs/mixed-deployments.md +++ b/rust/riverqueue/docs/mixed-deployments.md @@ -8,11 +8,11 @@ existing Go deployment and back out again. ## Matching versions -Each `riverqueue` minor release is matched to the River Go minor release with -the same number: `riverqueue` 0.49 runs alongside River Go 0.49. Patch releases -of either implementation can be mixed within a matched minor release. Upgrade -both implementations together when moving to a new minor release, following -the rolling procedure below. +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