From f0ce69aa7d46cd12f75f27ac5eb615099b411df4 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 1 Aug 2026 12:57:46 +0000 Subject: [PATCH 1/5] Fix the workspace build and make CI cover it `sqlxmq_stress` has not compiled since the sqlx 0.9 bump: it still pinned sqlx 0.8, so its `Pool` was a different type from the one `sqlxmq` expects and none of the executor bounds resolved. CI did not notice because `cargo check` and `cargo clippy --all-targets` at a root that is both a package and a workspace only build the root package. Both now run with `--workspace`, as does `cargo fmt`. The stress crate also carried its own copy of the migrations, which had drifted to three of the seven currently in the repository, so it was exercising a schema several fixes behind. Replace the copy with a symlink to the canonical directory. Verified by running the harness end to end against a fresh database: 10000 jobs, ~649 jobs/s. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01CmYdXrFbwtMHHHm1tREDUJ --- .github/workflows/toolchain.yml | 6 +- sqlxmq_stress/Cargo.toml | 2 +- sqlxmq_stress/migrations | 1 + .../migrations/20210316025847_setup.down.sql | 12 - .../migrations/20210316025847_setup.up.sql | 289 ------------------ .../migrations/20210921115907_clear.down.sql | 2 - .../migrations/20210921115907_clear.up.sql | 21 -- ...11013151757_fix_mq_latest_message.down.sql | 15 - ...0211013151757_fix_mq_latest_message.up.sql | 19 -- sqlxmq_stress/src/main.rs | 2 +- 10 files changed, 6 insertions(+), 363 deletions(-) create mode 120000 sqlxmq_stress/migrations delete mode 100644 sqlxmq_stress/migrations/20210316025847_setup.down.sql delete mode 100644 sqlxmq_stress/migrations/20210316025847_setup.up.sql delete mode 100644 sqlxmq_stress/migrations/20210921115907_clear.down.sql delete mode 100644 sqlxmq_stress/migrations/20210921115907_clear.up.sql delete mode 100644 sqlxmq_stress/migrations/20211013151757_fix_mq_latest_message.down.sql delete mode 100644 sqlxmq_stress/migrations/20211013151757_fix_mq_latest_message.up.sql diff --git a/.github/workflows/toolchain.yml b/.github/workflows/toolchain.yml index 11f8dbb..bcb9b8b 100644 --- a/.github/workflows/toolchain.yml +++ b/.github/workflows/toolchain.yml @@ -10,7 +10,7 @@ jobs: - uses: actions/checkout@v4 - uses: dtolnay/rust-toolchain@stable - uses: Swatinem/rust-cache@v2 - - run: cargo check + - run: cargo check --workspace --all-targets fmt: name: Rustfmt @@ -20,7 +20,7 @@ jobs: - uses: dtolnay/rust-toolchain@stable with: components: rustfmt - - run: cargo fmt -- --check + - run: cargo fmt --all -- --check clippy: name: Clippy @@ -31,7 +31,7 @@ jobs: with: components: clippy - uses: Swatinem/rust-cache@v2 - - run: cargo clippy --all-targets -- -D warnings + - run: cargo clippy --workspace --all-targets -- -D warnings test: name: Test diff --git a/sqlxmq_stress/Cargo.toml b/sqlxmq_stress/Cargo.toml index 90a8f12..23c137e 100644 --- a/sqlxmq_stress/Cargo.toml +++ b/sqlxmq_stress/Cargo.toml @@ -10,7 +10,7 @@ edition = "2018" sqlxmq = { path = ".." } tokio = { version = "1.4.0", features = ["full"] } dotenvy = "0.15" -sqlx = "0.8" +sqlx = { version = "0.9", default-features = false, features = ["postgres"] } serde = "1.0.125" lazy_static = "1.4.0" futures = "0.3.13" diff --git a/sqlxmq_stress/migrations b/sqlxmq_stress/migrations new file mode 120000 index 0000000..f0dcf84 --- /dev/null +++ b/sqlxmq_stress/migrations @@ -0,0 +1 @@ +../migrations \ No newline at end of file diff --git a/sqlxmq_stress/migrations/20210316025847_setup.down.sql b/sqlxmq_stress/migrations/20210316025847_setup.down.sql deleted file mode 100644 index 1aa472e..0000000 --- a/sqlxmq_stress/migrations/20210316025847_setup.down.sql +++ /dev/null @@ -1,12 +0,0 @@ -DROP FUNCTION mq_checkpoint; -DROP FUNCTION mq_keep_alive; -DROP FUNCTION mq_delete; -DROP FUNCTION mq_commit; -DROP FUNCTION mq_insert; -DROP FUNCTION mq_poll; -DROP FUNCTION mq_active_channels; -DROP FUNCTION mq_latest_message; -DROP TABLE mq_payloads; -DROP TABLE mq_msgs; -DROP FUNCTION mq_uuid_exists; -DROP TYPE mq_new_t; diff --git a/sqlxmq_stress/migrations/20210316025847_setup.up.sql b/sqlxmq_stress/migrations/20210316025847_setup.up.sql deleted file mode 100644 index bf7f8f8..0000000 --- a/sqlxmq_stress/migrations/20210316025847_setup.up.sql +++ /dev/null @@ -1,289 +0,0 @@ -CREATE EXTENSION IF NOT EXISTS "uuid-ossp"; - --- The UDT for creating messages -CREATE TYPE mq_new_t AS ( - -- Unique message ID - id UUID, - -- Delay before message is processed - delay INTERVAL, - -- Number of retries if initial processing fails - retries INT, - -- Initial backoff between retries - retry_backoff INTERVAL, - -- Name of channel - channel_name TEXT, - -- Arguments to channel - channel_args TEXT, - -- Interval for two-phase commit (or NULL to disable two-phase commit) - commit_interval INTERVAL, - -- Whether this message should be processed in order with respect to other - -- ordered messages. - ordered BOOLEAN, - -- Name of message - name TEXT, - -- JSON payload - payload_json TEXT, - -- Binary payload - payload_bytes BYTEA -); - --- Small, frequently updated table of messages -CREATE TABLE mq_msgs ( - id UUID PRIMARY KEY, - created_at TIMESTAMPTZ DEFAULT NOW(), - attempt_at TIMESTAMPTZ DEFAULT NOW(), - attempts INT NOT NULL DEFAULT 5, - retry_backoff INTERVAL NOT NULL DEFAULT INTERVAL '1 second', - channel_name TEXT NOT NULL, - channel_args TEXT NOT NULL, - commit_interval INTERVAL, - after_message_id UUID DEFAULT uuid_nil() REFERENCES mq_msgs(id) ON DELETE SET DEFAULT -); - --- Insert dummy message so that the 'nil' UUID can be referenced -INSERT INTO mq_msgs (id, channel_name, channel_args, after_message_id) VALUES (uuid_nil(), '', '', NULL); - --- Internal helper function to check that a UUID is neither NULL nor NIL -CREATE FUNCTION mq_uuid_exists( - id UUID -) RETURNS BOOLEAN AS $$ - SELECT id IS NOT NULL AND id != uuid_nil() -$$ LANGUAGE SQL IMMUTABLE; - --- Index for polling -CREATE INDEX ON mq_msgs(channel_name, channel_args, attempt_at) WHERE id != uuid_nil() AND NOT mq_uuid_exists(after_message_id); --- Index for adding messages -CREATE INDEX ON mq_msgs(channel_name, channel_args, created_at, id) WHERE id != uuid_nil() AND after_message_id IS NOT NULL; - --- Index for ensuring strict message order -CREATE UNIQUE INDEX mq_msgs_channel_name_channel_args_after_message_id_idx ON mq_msgs(channel_name, channel_args, after_message_id); - - --- Large, less frequently updated table of message payloads -CREATE TABLE mq_payloads( - id UUID PRIMARY KEY, - name TEXT NOT NULL, - payload_json JSONB, - payload_bytes BYTEA -); - --- Internal helper function to return the most recently added message in a queue. -CREATE FUNCTION mq_latest_message(from_channel_name TEXT, from_channel_args TEXT) -RETURNS UUID AS $$ - SELECT COALESCE( - ( - SELECT id FROM mq_msgs - WHERE channel_name = from_channel_name - AND channel_args = from_channel_args - AND after_message_id IS NOT NULL - AND id != uuid_nil() - ORDER BY created_at DESC, id DESC - LIMIT 1 - ), - uuid_nil() - ) -$$ LANGUAGE SQL STABLE; - --- Internal helper function to randomly select a set of channels with "ready" messages. -CREATE FUNCTION mq_active_channels(channel_names TEXT[], batch_size INT) -RETURNS TABLE(name TEXT, args TEXT) AS $$ - SELECT channel_name, channel_args - FROM mq_msgs - WHERE id != uuid_nil() - AND attempt_at <= NOW() - AND (channel_names IS NULL OR channel_name = ANY(channel_names)) - AND NOT mq_uuid_exists(after_message_id) - GROUP BY channel_name, channel_args - ORDER BY RANDOM() - LIMIT batch_size -$$ LANGUAGE SQL STABLE; - --- Main entry-point for job runner: pulls a batch of messages from the queue. -CREATE FUNCTION mq_poll(channel_names TEXT[], batch_size INT DEFAULT 1) -RETURNS TABLE( - id UUID, - is_committed BOOLEAN, - name TEXT, - payload_json TEXT, - payload_bytes BYTEA, - retry_backoff INTERVAL, - wait_time INTERVAL -) AS $$ -BEGIN - RETURN QUERY UPDATE mq_msgs - SET - attempt_at = CASE WHEN mq_msgs.attempts = 1 THEN NULL ELSE NOW() + mq_msgs.retry_backoff END, - attempts = mq_msgs.attempts - 1, - retry_backoff = mq_msgs.retry_backoff * 2 - FROM ( - SELECT - msgs.id - FROM mq_active_channels(channel_names, batch_size) AS active_channels - INNER JOIN LATERAL ( - SELECT * FROM mq_msgs - WHERE mq_msgs.id != uuid_nil() - AND mq_msgs.attempt_at <= NOW() - AND mq_msgs.channel_name = active_channels.name - AND mq_msgs.channel_args = active_channels.args - AND NOT mq_uuid_exists(mq_msgs.after_message_id) - ORDER BY mq_msgs.attempt_at ASC - LIMIT batch_size - ) AS msgs ON TRUE - LIMIT batch_size - ) AS messages_to_update - LEFT JOIN mq_payloads ON mq_payloads.id = messages_to_update.id - WHERE mq_msgs.id = messages_to_update.id - RETURNING - mq_msgs.id, - mq_msgs.commit_interval IS NULL, - mq_payloads.name, - mq_payloads.payload_json::TEXT, - mq_payloads.payload_bytes, - mq_msgs.retry_backoff / 2, - interval '0' AS wait_time; - - IF NOT FOUND THEN - RETURN QUERY SELECT - NULL::UUID, - NULL::BOOLEAN, - NULL::TEXT, - NULL::TEXT, - NULL::BYTEA, - NULL::INTERVAL, - MIN(mq_msgs.attempt_at) - NOW() - FROM mq_msgs - WHERE mq_msgs.id != uuid_nil() - AND NOT mq_uuid_exists(mq_msgs.after_message_id) - AND (channel_names IS NULL OR mq_msgs.channel_name = ANY(channel_names)); - END IF; -END; -$$ LANGUAGE plpgsql; - --- Creates new messages -CREATE FUNCTION mq_insert(new_messages mq_new_t[]) -RETURNS VOID AS $$ -BEGIN - PERFORM pg_notify(CONCAT('mq_', channel_name), '') - FROM unnest(new_messages) AS new_msgs - GROUP BY channel_name; - - IF FOUND THEN - PERFORM pg_notify('mq', ''); - END IF; - - INSERT INTO mq_payloads ( - id, - name, - payload_json, - payload_bytes - ) SELECT - id, - name, - payload_json::JSONB, - payload_bytes - FROM UNNEST(new_messages); - - INSERT INTO mq_msgs ( - id, - attempt_at, - attempts, - retry_backoff, - channel_name, - channel_args, - commit_interval, - after_message_id - ) - SELECT - id, - NOW() + delay + COALESCE(commit_interval, INTERVAL '0'), - retries + 1, - retry_backoff, - channel_name, - channel_args, - commit_interval, - CASE WHEN ordered - THEN - LAG(id, 1, mq_latest_message(channel_name, channel_args)) - OVER (PARTITION BY channel_name, channel_args, ordered ORDER BY id) - ELSE - NULL - END - FROM UNNEST(new_messages); -END; -$$ LANGUAGE plpgsql; - --- Commits messages previously created with a non-NULL commit interval. -CREATE FUNCTION mq_commit(msg_ids UUID[]) -RETURNS VOID AS $$ -BEGIN - UPDATE mq_msgs - SET - attempt_at = attempt_at - commit_interval, - commit_interval = NULL - WHERE id = ANY(msg_ids) - AND commit_interval IS NOT NULL; -END; -$$ LANGUAGE plpgsql; - - --- Deletes messages from the queue. This occurs when a message has been --- processed, or when it expires without being processed. -CREATE FUNCTION mq_delete(msg_ids UUID[]) -RETURNS VOID AS $$ -BEGIN - PERFORM pg_notify(CONCAT('mq_', channel_name), '') - FROM mq_msgs - WHERE id = ANY(msg_ids) - AND after_message_id = uuid_nil() - GROUP BY channel_name; - - IF FOUND THEN - PERFORM pg_notify('mq', ''); - END IF; - - DELETE FROM mq_msgs WHERE id = ANY(msg_ids); - DELETE FROM mq_payloads WHERE id = ANY(msg_ids); -END; -$$ LANGUAGE plpgsql; - - --- Can be called during the initial commit interval, or when processing --- a message. Indicates that the caller is still active and will prevent either --- the commit interval elapsing or the message being retried for the specified --- interval. -CREATE FUNCTION mq_keep_alive(msg_ids UUID[], duration INTERVAL) -RETURNS VOID AS $$ - UPDATE mq_msgs - SET - attempt_at = NOW() + duration, - commit_interval = commit_interval + ((NOW() + duration) - attempt_at) - WHERE id = ANY(msg_ids) - AND attempt_at < NOW() + duration; -$$ LANGUAGE SQL; - - --- Called during lengthy processing of a message to checkpoint the progress. --- As well as behaving like `mq_keep_alive`, the message payload can be --- updated. -CREATE FUNCTION mq_checkpoint( - msg_id UUID, - duration INTERVAL, - new_payload_json TEXT, - new_payload_bytes BYTEA, - extra_retries INT -) -RETURNS VOID AS $$ - UPDATE mq_msgs - SET - attempt_at = GREATEST(attempt_at, NOW() + duration), - attempts = attempts + COALESCE(extra_retries, 0) - WHERE id = msg_id; - - UPDATE mq_payloads - SET - payload_json = COALESCE(new_payload_json::JSONB, payload_json), - payload_bytes = COALESCE(new_payload_bytes, payload_bytes) - WHERE - id = msg_id; -$$ LANGUAGE SQL; - diff --git a/sqlxmq_stress/migrations/20210921115907_clear.down.sql b/sqlxmq_stress/migrations/20210921115907_clear.down.sql deleted file mode 100644 index e15638d..0000000 --- a/sqlxmq_stress/migrations/20210921115907_clear.down.sql +++ /dev/null @@ -1,2 +0,0 @@ -DROP FUNCTION mq_clear; -DROP FUNCTION mq_clear_all; diff --git a/sqlxmq_stress/migrations/20210921115907_clear.up.sql b/sqlxmq_stress/migrations/20210921115907_clear.up.sql deleted file mode 100644 index bd1c1f6..0000000 --- a/sqlxmq_stress/migrations/20210921115907_clear.up.sql +++ /dev/null @@ -1,21 +0,0 @@ --- Deletes all messages from a list of channel names. -CREATE FUNCTION mq_clear(channel_names TEXT[]) -RETURNS VOID AS $$ -BEGIN - WITH deleted_ids AS ( - DELETE FROM mq_msgs WHERE channel_name = ANY(channel_names) RETURNING id - ) - DELETE FROM mq_payloads WHERE id IN (SELECT id FROM deleted_ids); -END; -$$ LANGUAGE plpgsql; - --- Deletes all messages. -CREATE FUNCTION mq_clear_all() -RETURNS VOID AS $$ -BEGIN - WITH deleted_ids AS ( - DELETE FROM mq_msgs RETURNING id - ) - DELETE FROM mq_payloads WHERE id IN (SELECT id FROM deleted_ids); -END; -$$ LANGUAGE plpgsql; diff --git a/sqlxmq_stress/migrations/20211013151757_fix_mq_latest_message.down.sql b/sqlxmq_stress/migrations/20211013151757_fix_mq_latest_message.down.sql deleted file mode 100644 index d09bd4a..0000000 --- a/sqlxmq_stress/migrations/20211013151757_fix_mq_latest_message.down.sql +++ /dev/null @@ -1,15 +0,0 @@ -CREATE OR REPLACE FUNCTION mq_latest_message(from_channel_name TEXT, from_channel_args TEXT) -RETURNS UUID AS $$ - SELECT COALESCE( - ( - SELECT id FROM mq_msgs - WHERE channel_name = from_channel_name - AND channel_args = from_channel_args - AND after_message_id IS NOT NULL - AND id != uuid_nil() - ORDER BY created_at DESC, id DESC - LIMIT 1 - ), - uuid_nil() - ) -$$ LANGUAGE SQL STABLE; diff --git a/sqlxmq_stress/migrations/20211013151757_fix_mq_latest_message.up.sql b/sqlxmq_stress/migrations/20211013151757_fix_mq_latest_message.up.sql deleted file mode 100644 index b987c5e..0000000 --- a/sqlxmq_stress/migrations/20211013151757_fix_mq_latest_message.up.sql +++ /dev/null @@ -1,19 +0,0 @@ -CREATE OR REPLACE FUNCTION mq_latest_message(from_channel_name TEXT, from_channel_args TEXT) -RETURNS UUID AS $$ - SELECT COALESCE( - ( - SELECT id FROM mq_msgs - WHERE channel_name = from_channel_name - AND channel_args = from_channel_args - AND after_message_id IS NOT NULL - AND id != uuid_nil() - AND NOT EXISTS( - SELECT * FROM mq_msgs AS mq_msgs2 - WHERE mq_msgs2.after_message_id = mq_msgs.id - ) - ORDER BY created_at DESC - LIMIT 1 - ), - uuid_nil() - ) -$$ LANGUAGE SQL STABLE; \ No newline at end of file diff --git a/sqlxmq_stress/src/main.rs b/sqlxmq_stress/src/main.rs index 8daad76..6b83cc4 100644 --- a/sqlxmq_stress/src/main.rs +++ b/sqlxmq_stress/src/main.rs @@ -49,7 +49,7 @@ async fn start_job( pool: Pool, seed: usize, ) -> Result<(), Box> { - let channel_name = if seed % 3 == 0 { "foo" } else { "bar" }; + let channel_name = if seed.is_multiple_of(3) { "foo" } else { "bar" }; let channel_args = format!("{}", seed / 32); example_job .builder() From bb6a06358d7a1e76f32fea150405af53e13d319b Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 1 Aug 2026 12:57:53 +0000 Subject: [PATCH 2/5] Add unit tests for to_duration The pure helpers in this crate had no coverage at all: every test needed a live PostgreSQL connection, so nothing exercised the plain arithmetic. `to_duration` converts a `PgInterval` into the sleep duration the runner waits before its next poll. Cover the month/day/microsecond components and, in particular, the negative case: `Duration` cannot represent a negative value, so an overdue `attempt_at` must collapse to zero and be polled immediately rather than waited on. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01CmYdXrFbwtMHHHm1tREDUJ --- src/runner.rs | 67 +++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/src/runner.rs b/src/runner.rs index 056a67a..1bc6d28 100644 --- a/src/runner.rs +++ b/src/runner.rs @@ -512,3 +512,70 @@ async fn keep_job_alive(id: Uuid, pool: Pool, mut interval: Duration) } } } + +/// Unit tests for the pure helpers in this module. These do not require a +/// database connection. +#[cfg(test)] +mod tests { + use super::*; + + fn interval(months: i32, days: i32, microseconds: i64) -> PgInterval { + PgInterval { + months, + days, + microseconds, + } + } + + const DAY: Duration = Duration::from_secs(24 * 60 * 60); + + #[test] + fn to_duration_converts_zero() { + assert_eq!(to_duration(interval(0, 0, 0)), Duration::ZERO); + } + + #[test] + fn to_duration_converts_microseconds() { + assert_eq!(to_duration(interval(0, 0, 1)), Duration::from_micros(1)); + assert_eq!( + to_duration(interval(0, 0, 1_500_000)), + Duration::from_millis(1500) + ); + } + + #[test] + fn to_duration_converts_days() { + assert_eq!(to_duration(interval(0, 1, 0)), DAY); + assert_eq!(to_duration(interval(0, 3, 0)), 3 * DAY); + } + + /// Postgres intervals do not carry an anchor date, so a month is treated + /// as exactly 30 days. + #[test] + fn to_duration_treats_a_month_as_thirty_days() { + assert_eq!(to_duration(interval(1, 0, 0)), 30 * DAY); + assert_eq!(to_duration(interval(2, 0, 0)), 60 * DAY); + } + + #[test] + fn to_duration_sums_all_components() { + assert_eq!( + to_duration(interval(1, 2, 500_000)), + 32 * DAY + Duration::from_millis(500) + ); + } + + /// `Duration` cannot represent negative values, so an interval that is + /// negative in any component collapses to zero. This matters because the + /// runner uses the result as a sleep duration: a message whose + /// `attempt_at` is already in the past must be polled immediately, not + /// waited on. + #[test] + fn to_duration_clamps_negative_components_to_zero() { + assert_eq!(to_duration(interval(0, 0, -1)), Duration::ZERO); + assert_eq!(to_duration(interval(0, -1, 0)), Duration::ZERO); + assert_eq!(to_duration(interval(-1, 0, 0)), Duration::ZERO); + // A negative component wins even when the total would be positive. + assert_eq!(to_duration(interval(1, 0, -1)), Duration::ZERO); + } +} From 8374030f61f07091fdaeba744bee724f8cb21c4d Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 1 Aug 2026 12:58:18 +0000 Subject: [PATCH 3/5] Add unit tests for should_retry `should_retry` tells callers which database errors are worth retrying, but nothing exercised it. The cases it distinguishes are easy to get wrong and expensive when wrong: a constraint violation on one of the two ordered-channel constraints is a lost race and should be retried, whereas the same SQLSTATE raised by the caller's own schema is a genuine conflict that will just recur. The errors are built from a stub `DatabaseError`, so these run without a database. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01CmYdXrFbwtMHHHm1tREDUJ --- src/lib.rs | 120 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 120 insertions(+) diff --git a/src/lib.rs b/src/lib.rs index 093a9b2..12f7660 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -259,11 +259,131 @@ pub fn should_retry(error: &sqlx::Error) -> bool { } } +/// Unit tests for [`should_retry`]. These do not require a database +/// connection: the errors are constructed from a stub `DatabaseError`. +#[cfg(test)] +mod should_retry_tests { + use super::*; + + use std::borrow::Cow; + use std::error::Error as StdError; + use std::fmt::{self, Display}; + + use sqlx::error::{DatabaseError, ErrorKind}; + + #[derive(Debug)] + struct StubDbError { + code: Option<&'static str>, + constraint: Option<&'static str>, + } + + impl Display for StubDbError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("stub database error") + } + } + + impl StdError for StubDbError {} + + impl DatabaseError for StubDbError { + fn message(&self) -> &str { + "stub database error" + } + fn code(&self) -> Option> { + self.code.map(Cow::Borrowed) + } + fn constraint(&self) -> Option<&str> { + self.constraint + } + fn kind(&self) -> ErrorKind { + ErrorKind::Other + } + fn as_error(&self) -> &(dyn StdError + Send + Sync + 'static) { + self + } + fn as_error_mut(&mut self) -> &mut (dyn StdError + Send + Sync + 'static) { + self + } + fn into_error(self: Box) -> Box { + self + } + } + + fn db_error(code: Option<&'static str>, constraint: Option<&'static str>) -> sqlx::Error { + sqlx::Error::Database(Box::new(StubDbError { code, constraint })) + } + + #[test] + fn retries_serialization_failure() { + assert!(should_retry(&db_error(Some("40001"), None))); + } + + #[test] + fn retries_deadlock() { + assert!(should_retry(&db_error(Some("40P01"), None))); + } + + /// Two ordered messages racing to chain onto the same predecessor. + #[test] + fn retries_ordered_channel_unique_violation() { + assert!(should_retry(&db_error( + Some("23505"), + Some("mq_msgs_channel_name_channel_args_after_message_id_idx"), + ))); + } + + /// The predecessor of an ordered message was deleted concurrently. + #[test] + fn retries_ordered_channel_foreign_key_violation() { + assert!(should_retry(&db_error( + Some("23503"), + Some("mq_msgs_after_message_id_fkey"), + ))); + } + + /// The constraint violations are only retryable on the ordered-channel + /// constraints: the same SQLSTATE raised by the caller's own schema means + /// a genuine conflict, and retrying would just raise it again. + #[test] + fn does_not_retry_constraint_violations_on_other_constraints() { + assert!(!should_retry(&db_error( + Some("23505"), + Some("users_email_key") + ))); + assert!(!should_retry(&db_error( + Some("23503"), + Some("orders_user_id_fkey") + ))); + assert!(!should_retry(&db_error(Some("23505"), None))); + } + + #[test] + fn does_not_retry_other_database_errors() { + // Undefined table. + assert!(!should_retry(&db_error(Some("42P01"), None))); + // Syntax error. + assert!(!should_retry(&db_error(Some("42601"), None))); + assert!(!should_retry(&db_error(None, None))); + } + + /// Non-database errors (connection loss, decoding failures, pool + /// timeouts) are not retryable by this helper. + #[test] + fn does_not_retry_non_database_errors() { + assert!(!should_retry(&sqlx::Error::PoolTimedOut)); + assert!(!should_retry(&sqlx::Error::RowNotFound)); + assert!(!should_retry(&sqlx::Error::Io(std::io::Error::other( + "connection reset" + )))); + } +} + #[cfg(test)] mod tests { use super::*; use crate as sqlxmq; + use std::collections::HashSet; use std::env; use std::error::Error; use std::future::Future; From 4e429463f4ff9243f473321e107f8c5690a7b06e Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 1 Aug 2026 12:58:28 +0000 Subject: [PATCH 4/5] Add concurrency tests for job delivery The SKIP LOCKED change in #2 exists so that concurrent pollers claim disjoint sets of messages instead of blocking on each other's row locks, but nothing tested it: no test ever started more than one runner. Two tests, at different levels: `it_polls_disjoint_messages_from_concurrent_transactions` calls mq_poll from two explicit transactions so the first still holds its row locks when the second runs. It asserts the second poll returns promptly, that the two claims are disjoint, and that together they cover every spawned job. Confirmed to fail against the pre-SKIP-LOCKED mq_poll, where the second poll blocks until the test's timeout fires. `it_delivers_each_job_exactly_once_with_concurrent_runners` covers the same ground end to end: four runners with deliberately low concurrency limits, jobs spread over several channels, asserting exactly-once delivery and an empty queue afterwards. It does not single out SKIP LOCKED - it guards the delivery contract the locking strategy exists to uphold, so it stays valid if that strategy changes again. Both are event-driven rather than sleep-based, and complete in about a second. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01CmYdXrFbwtMHHHm1tREDUJ --- src/lib.rs | 179 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 179 insertions(+) diff --git a/src/lib.rs b/src/lib.rs index 12f7660..81b6e83 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -397,6 +397,7 @@ mod tests { use sqlx::{Pool, Postgres}; use tokio::sync::{Mutex, MutexGuard}; use tokio::task; + use uuid::Uuid; // field 0 is never read, but its drop is important #[allow(dead_code)] @@ -819,4 +820,182 @@ mod tests { } pause().await; } + + /// `mq_poll` claims candidate rows with `FOR UPDATE SKIP LOCKED`, so a + /// poller must never block on, or hand back, rows another poller has + /// already claimed. + /// + /// Both polls run inside explicit transactions: row locks are held until + /// the transaction ends, so the second poll genuinely contends with the + /// first. Without `SKIP LOCKED` the second poll blocks on the first + /// transaction's row locks and this test times out; without the + /// `MATERIALIZED` CTE the locking subquery can be re-executed and return + /// rows already claimed, and the disjointness assertion fails. + #[tokio::test] + async fn it_polls_disjoint_messages_from_concurrent_transactions() { + let pool = &*test_pool().await; + + const NUM_JOBS: usize = 4; + const BATCH_SIZE: i32 = 2; + + let builders: Vec<_> = (0..NUM_JOBS) + .map(|_| JobBuilder::new("concurrent")) + .collect(); + let spawned: HashSet = spawn_batch(pool, &builders) + .await + .unwrap() + .into_iter() + .collect(); + + async fn poll_ids(tx: &mut sqlx::Transaction<'_, Postgres>, batch_size: i32) -> Vec { + sqlx::query_scalar::<_, Uuid>("SELECT id FROM mq_poll($1, $2) WHERE id IS NOT NULL") + .bind(Option::>::None) + .bind(batch_size) + .fetch_all(&mut **tx) + .await + .unwrap() + } + + let mut tx_a = pool.begin().await.unwrap(); + let mut tx_b = pool.begin().await.unwrap(); + + // `tx_a` claims a batch and holds the locks: it is not committed yet. + let claimed_a = poll_ids(&mut tx_a, BATCH_SIZE).await; + assert_eq!(claimed_a.len(), BATCH_SIZE as usize); + + // `tx_b` must skip what `tx_a` holds rather than waiting for it. + let claimed_b = + tokio::time::timeout(Duration::from_secs(10), poll_ids(&mut tx_b, BATCH_SIZE)) + .await + .expect("concurrent poll blocked on the first transaction's row locks"); + assert_eq!(claimed_b.len(), BATCH_SIZE as usize); + + tx_a.commit().await.unwrap(); + tx_b.commit().await.unwrap(); + + let set_a: HashSet = claimed_a.iter().copied().collect(); + let set_b: HashSet = claimed_b.iter().copied().collect(); + assert_eq!( + set_a.len(), + claimed_a.len(), + "a poll returned duplicate ids" + ); + assert_eq!( + set_b.len(), + claimed_b.len(), + "a poll returned duplicate ids" + ); + assert!( + set_a.is_disjoint(&set_b), + "concurrent polls claimed overlapping messages: {:?} and {:?}", + set_a, + set_b + ); + assert_eq!( + &set_a | &set_b, + spawned, + "the two polls together should have claimed every spawned job" + ); + } + + /// End-to-end counterpart to the test above: several independent job + /// runners sharing one database must between them deliver each job + /// exactly once, and must all make progress rather than deadlocking or + /// starving each other. + /// + /// Concurrency limits are kept low so that no single runner can drain the + /// queue in one batch, forcing the runners to contend for the same rows. + /// + /// Unlike the test above this one does not single out `SKIP LOCKED` — it + /// guards the delivery contract that the locking strategy exists to + /// uphold, and so stays valid if that strategy is changed again. + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] + async fn it_delivers_each_job_exactly_once_with_concurrent_runners() { + let pool = &*test_pool().await; + + const NUM_RUNNERS: usize = 4; + const NUM_JOBS: usize = 24; + const NUM_CHANNELS: usize = 4; + + let (tx, mut rx) = mpsc::unbounded(); + + let mut runners = Vec::new(); + for runner_idx in 0..NUM_RUNNERS { + let tx = tx.clone(); + let runner = JobRunnerOptions::new(pool, move |mut job| { + let tx = tx.clone(); + task::spawn(async move { + let id = job.id(); + job.complete().await.unwrap(); + tx.unbounded_send((runner_idx, id)).unwrap(); + }); + }) + .set_concurrency(2, 6) + .run() + .await + .unwrap(); + runners.push(runner); + } + // Drop the extra sender so the receiver terminates if every runner + // goes away. + drop(tx); + + // Spread the jobs over several channels so that `mq_active_channels` + // is exercised as well as the per-channel batching. + let channel_args: Vec = (0..NUM_JOBS) + .map(|i| (i % NUM_CHANNELS).to_string()) + .collect(); + let builders: Vec<_> = channel_args + .iter() + .map(|args| { + let mut builder = JobBuilder::new("concurrent"); + builder.set_channel_args(args); + builder + }) + .collect(); + let spawned: HashSet = spawn_batch(pool, &builders) + .await + .unwrap() + .into_iter() + .collect(); + assert_eq!(spawned.len(), NUM_JOBS); + + let mut delivered = HashSet::new(); + let mut delivering_runners = HashSet::new(); + for _ in 0..NUM_JOBS { + let (runner_idx, id) = tokio::time::timeout(Duration::from_secs(60), rx.next()) + .await + .expect("timed out waiting for jobs to be delivered") + .expect("all runners stopped before every job was delivered"); + assert!( + delivered.insert(id), + "job {} was delivered more than once", + id + ); + delivering_runners.insert(runner_idx); + } + assert_eq!(delivered, spawned); + + // Anything already queued beyond the expected count is a redelivery. + assert!(rx.try_recv().is_err(), "a job was delivered more than once"); + + // Every job was completed, so the queue must be empty. + for id in &spawned { + assert!( + !exists(pool, *id).await.unwrap(), + "job {} was not deleted", + id + ); + } + + log::info!( + "{} of {} runners took part in delivery", + delivering_runners.len(), + NUM_RUNNERS + ); + + for mut runner in runners { + runner.stop().await; + } + } } From 60f12d973b12bc6a8539766af006665a2a578be4 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 1 Aug 2026 13:34:43 +0000 Subject: [PATCH 5/5] Give each test its own database and make the tests event-driven Two changes to the same harness. Isolation: replace the hand-rolled pool with `#[sqlx::test]`, which gives each test a fresh database with the migrations already applied. That removes three problems at once. The global `TEST_MUTEX` is gone, so tests run in parallel instead of queueing for one shared database. The `TRUNCATE mq_payloads` / `DELETE FROM mq_msgs` reset is gone, and with it the risk of `cargo test` wiping the job tables of whatever database an ambient `DATABASE_URL` happened to point at - `dotenvy::dotenv()` does not override an already-set variable, so `.env` was not the safeguard it looked like. And migrations are now exercised by every test run rather than applied out of band, so CI no longer builds sqlx-cli: the workflow just creates an empty database. Synchronisation: the tests asserted on a counter after sleeping a fixed interval, which made them a race between the runner and the clock - `it_can_checkpoint_jobs` lost that race consistently on a slower machine, because the automatic keep-alive pushes `attempt_at` out from under a test that assumes a retry lands exactly one backoff after the attempt. Runners now hand each dispatched job to the test, which waits for the events it cares about. Where the old tests waited to see whether something further would happen, the new ones assert on the state that makes it impossible: a cleared channel has no rows left, an exhausted message has `attempt_at IS NULL` and can never be polled again, and a completed one is deleted. Remaining timeouts are hang detectors set far above any real duration, not timing assumptions. The assertions got stronger in the process. The ordered-batch test now checks the chain in `mq_msgs` directly, which is what the ordinality migration actually changed; the checkpoint test checks that the payload was replaced rather than only counting attempts; and `it_can_use_registry`, which previously asserted nothing at all, now confirms all three jobs ran by having them report through a context. No `sleep` remains in the test module. The suite is green 15 runs in a row and 5 more under 2x CPU oversubscription, in about 4 seconds against 11.5 locally and roughly 37 seconds of sleeping on CI. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01CmYdXrFbwtMHHHm1tREDUJ --- .github/actions/postgres/action.yml | 8 +- .github/workflows/toolchain.yml | 2 - Cargo.toml | 4 + src/lib.rs | 773 +++++++++++++++------------- 4 files changed, 427 insertions(+), 360 deletions(-) diff --git a/.github/actions/postgres/action.yml b/.github/actions/postgres/action.yml index c056531..86be199 100644 --- a/.github/actions/postgres/action.yml +++ b/.github/actions/postgres/action.yml @@ -13,6 +13,10 @@ runs: pg_isready sudo -u postgres psql -c "ALTER USER postgres PASSWORD 'password'" - - name: Setup database + # `#[sqlx::test]` gives each test its own database and applies the + # migrations to it, so this only has to exist for sqlx to connect to and + # keep its bookkeeping in. No migrations are run here, and sqlx-cli is not + # needed. + - name: Create database shell: bash - run: cargo sqlx database setup + run: sudo -u postgres createdb sqlxmq diff --git a/.github/workflows/toolchain.yml b/.github/workflows/toolchain.yml index bcb9b8b..bcd8262 100644 --- a/.github/workflows/toolchain.yml +++ b/.github/workflows/toolchain.yml @@ -42,7 +42,6 @@ jobs: - uses: actions/checkout@v4 - uses: dtolnay/rust-toolchain@stable - uses: Swatinem/rust-cache@v2 - - run: cargo install sqlx-cli --locked - uses: ./.github/actions/postgres - run: cargo test -- --nocapture @@ -55,7 +54,6 @@ jobs: - uses: actions/checkout@v4 - uses: dtolnay/rust-toolchain@nightly - uses: Swatinem/rust-cache@v2 - - run: cargo install sqlx-cli --locked - uses: ./.github/actions/postgres - run: cargo test -- --nocapture diff --git a/Cargo.toml b/Cargo.toml index dd4d358..35015a0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,3 +36,7 @@ dotenvy = "0.15.3" pretty_env_logger = "0.4.0" futures = "0.3.13" tokio = { version = "1", features = ["full"] } +# `migrate` is only needed by `#[sqlx::test]`, which gives each test its own +# freshly migrated database. Keeping it in dev-dependencies means it is not +# enabled for consumers of the library. +sqlx = { version = "0.9", default-features = false, features = ["migrate"] } diff --git a/src/lib.rs b/src/lib.rs index 81b6e83..f660a83 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -384,95 +384,170 @@ mod tests { use crate as sqlxmq; use std::collections::HashSet; - use std::env; use std::error::Error; - use std::future::Future; - use std::ops::Deref; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::{Arc, Once}; - use std::time::Duration; + use std::sync::Once; + use std::time::{Duration, Instant}; use futures::channel::mpsc; use futures::StreamExt; use sqlx::{Pool, Postgres}; - use tokio::sync::{Mutex, MutexGuard}; - use tokio::task; use uuid::Uuid; - // field 0 is never read, but its drop is important - #[allow(dead_code)] - struct TestGuard(MutexGuard<'static, ()>, T); - - impl Deref for TestGuard { - type Target = T; + /// Upper bound on every wait in these tests. + /// + /// This is a hang detector, not a timing assumption: a working + /// implementation dispatches as soon as the work is ready, so the only + /// way to reach this limit is for something to be genuinely stuck. It is + /// deliberately far longer than any operation here should take, so that a + /// slow or loaded machine cannot turn it into a flaky failure. + const TIMEOUT: Duration = Duration::from_secs(30); + + /// Initial retry backoff for the tests that exercise the retry schedule. + /// + /// It can be short because those tests assert on the schedule recorded in + /// the database rather than on how long the test waited. + const BACKOFF: Duration = Duration::from_millis(100); - fn deref(&self) -> &T { - &self.1 - } + fn init_logging() { + static INIT_LOGGER: Once = Once::new(); + INIT_LOGGER.call_once(pretty_env_logger::init); } - async fn test_pool() -> TestGuard> { - static INIT_LOGGER: Once = Once::new(); - static TEST_MUTEX: Mutex<()> = Mutex::const_new(()); + /// Jobs dispatched by a test runner, in the order the runner handed them + /// out. + /// + /// Tests take jobs from here and drive them explicitly: completing a job + /// finishes it, dropping one without completing it leaves it to be + /// retried, exactly as a failed job would be. + struct JobStream(mpsc::UnboundedReceiver); + + impl JobStream { + /// Wait for the next job to be dispatched. + async fn next(&mut self) -> CurrentJob { + tokio::time::timeout(TIMEOUT, self.0.next()) + .await + .expect("timed out waiting for a job to be dispatched") + .expect("the job runner stopped before dispatching a job") + } - let guard = TEST_MUTEX.lock().await; + /// Wait for the next `n` jobs to be dispatched. + async fn next_n(&mut self, n: usize) -> Vec { + let mut jobs = Vec::with_capacity(n); + for _ in 0..n { + jobs.push(self.next().await); + } + jobs + } + } - let _ = dotenvy::dotenv(); + /// Start a job runner which forwards every dispatched job to the returned + /// [`JobStream`] instead of running it. + /// + /// Handing jobs to the test rather than to a handler is what lets these + /// tests wait for the events they care about instead of sleeping for long + /// enough that the events have probably happened. + async fn test_job_runner(pool: &Pool) -> (JobRunnerHandle, JobStream) { + configured_job_runner(pool, |options| options).await + } - INIT_LOGGER.call_once(pretty_env_logger::init); + /// As [`test_job_runner`], with an opportunity to change the runner + /// options first. + async fn configured_job_runner( + pool: &Pool, + configure: impl FnOnce(&mut JobRunnerOptions) -> &mut JobRunnerOptions, + ) -> (JobRunnerHandle, JobStream) { + init_logging(); + + let (tx, rx) = mpsc::unbounded(); + let mut options = JobRunnerOptions::new(pool, move |job| { + // If the test has dropped the receiver the job is dropped too, + // and will be retried like any other unfinished job. + let _ = tx.unbounded_send(job); + }); + configure(&mut options); + let runner = options.run().await.unwrap(); + + (runner, JobStream(rx)) + } - let pool = Pool::connect(&env::var("DATABASE_URL").unwrap()) + /// Number of messages in the queue, excluding the nil sentinel row. + async fn queue_len(pool: &Pool) -> i64 { + sqlx::query_scalar("SELECT COUNT(*) FROM mq_msgs WHERE id != uuid_nil()") + .fetch_one(pool) .await - .unwrap(); + .unwrap() + } - sqlx::query("TRUNCATE TABLE mq_payloads") - .execute(&pool) + /// The message this one is chained behind: the nil UUID for the head of + /// an ordered chain, and `None` for an unordered message. + async fn after_message_id(pool: &Pool, id: Uuid) -> Option { + sqlx::query_scalar("SELECT after_message_id FROM mq_msgs WHERE id = $1") + .bind(id) + .fetch_one(pool) .await - .unwrap(); - sqlx::query("DELETE FROM mq_msgs WHERE id != uuid_nil()") - .execute(&pool) + .unwrap() + } + + /// Attempts remaining, and whether a further attempt is scheduled. + /// + /// `mq_poll` clears `attempt_at` as it hands out the final attempt, so + /// `(0, false)` means the message is exhausted and can never be polled + /// again. Asserting on this beats waiting to see whether another attempt + /// shows up. + async fn attempt_state(pool: &Pool, id: Uuid) -> (i32, bool) { + sqlx::query_as("SELECT attempts, attempt_at IS NOT NULL FROM mq_msgs WHERE id = $1") + .bind(id) + .fetch_one(pool) .await - .unwrap(); + .unwrap() + } - TestGuard(guard, pool) + /// The raw JSON payload currently stored for a message. + async fn stored_payload(pool: &Pool, id: Uuid) -> Option { + sqlx::query_scalar("SELECT payload_json::TEXT FROM mq_payloads WHERE id = $1") + .bind(id) + .fetch_one(pool) + .await + .unwrap() } - async fn test_job_runner( - pool: &Pool, - f: impl (Fn(CurrentJob) -> F) + Send + Sync + 'static, - ) -> (JobRunnerHandle, Arc) - where - F::Output: Send + 'static, - { - let counter = Arc::new(AtomicUsize::new(0)); - let counter2 = counter.clone(); - let runner = JobRunnerOptions::new(pool, move |job| { - counter2.fetch_add(1, Ordering::SeqCst); - task::spawn(f(job)); - }) - .run() + /// The retry backoff currently recorded for a message. `mq_poll` doubles + /// it on every attempt. + async fn retry_backoff(pool: &Pool, id: Uuid) -> Duration { + let micros: i64 = sqlx::query_scalar( + "SELECT (EXTRACT(EPOCH FROM retry_backoff) * 1000000)::BIGINT FROM mq_msgs WHERE id = $1", + ) + .bind(id) + .fetch_one(pool) .await .unwrap(); - (runner, counter) + Duration::from_micros(micros as u64) } fn job_proto<'a, 'b>(builder: &'a mut JobBuilder<'b>) -> &'a mut JobBuilder<'b> { builder.set_channel_name("bar") } + /// Context which lets the registry-based jobs report that they ran. + type RanJobs = mpsc::UnboundedSender<&'static str>; + #[job(channel_name = "foo", ordered, retries = 3, backoff_secs = 2.0)] async fn example_job1( mut current_job: CurrentJob, + ran: RanJobs, ) -> Result<(), Box> { current_job.complete().await?; + ran.unbounded_send("example_job1")?; Ok(()) } #[job(proto(job_proto))] async fn example_job2( mut current_job: CurrentJob, + ran: RanJobs, ) -> Result<(), Box> { current_job.complete().await?; + ran.unbounded_send("example_job2")?; Ok(()) } @@ -481,344 +556,344 @@ mod tests { mut current_job: CurrentJob, ctx1: i32, ctx2: &'static str, + ran: RanJobs, ) -> Result<(), Box> { assert_eq!(ctx1, 42); assert_eq!(ctx2, "Hello, world!"); current_job.complete().await?; + ran.unbounded_send("example_job_with_ctx")?; Ok(()) } - async fn named_job_runner(pool: &Pool) -> JobRunnerHandle { + async fn named_job_runner(pool: &Pool, ran: RanJobs) -> JobRunnerHandle { + init_logging(); + let mut registry = JobRegistry::new(&[example_job1, example_job2, example_job_with_ctx]); - registry.set_context(42).set_context("Hello, world!"); + registry + .set_context(42) + .set_context("Hello, world!") + .set_context(ran); registry.runner(pool).run().await.unwrap() } - fn is_ci() -> bool { - std::env::var("CI").ok().is_some() - } + #[sqlx::test] + async fn it_can_spawn_job(pool: Pool) { + let (mut runner, mut jobs) = test_job_runner(&pool).await; - fn default_pause() -> u64 { - if is_ci() { - 1000 - } else { - 200 - } - } + let id = JobBuilder::new("foo").spawn(&pool).await.unwrap(); - async fn pause() { - pause_ms(default_pause()).await; - } + let mut job = jobs.next().await; + assert_eq!(job.id(), id); + assert_eq!(job.name(), "foo"); + job.complete().await.unwrap(); - async fn pause_ms(ms: u64) { - tokio::time::sleep(Duration::from_millis(ms)).await; + assert_eq!(queue_len(&pool).await, 0); + runner.stop().await; } - #[tokio::test] - async fn it_can_spawn_job() { - { - let pool = &*test_pool().await; - let (_runner, counter) = - test_job_runner(pool, |mut job| async move { job.complete().await }).await; - - assert_eq!(counter.load(Ordering::SeqCst), 0); - JobBuilder::new("foo").spawn(pool).await.unwrap(); - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 1); + #[sqlx::test] + async fn it_can_clear_jobs(pool: Pool) { + let mut kept = HashSet::new(); + for channel_name in ["foo", "bar", "baz"] { + for _ in 0..2 { + let id = JobBuilder::new("foo") + .set_channel_name(channel_name) + .spawn(&pool) + .await + .unwrap(); + if channel_name == "bar" { + kept.insert(id); + } + } } - pause().await; - } + assert_eq!(queue_len(&pool).await, 6); - #[tokio::test] - async fn it_can_clear_jobs() { - { - let pool = &*test_pool().await; - JobBuilder::new("foo") - .set_channel_name("foo") - .spawn(pool) - .await - .unwrap(); - JobBuilder::new("foo") - .set_channel_name("foo") - .spawn(pool) - .await - .unwrap(); - JobBuilder::new("foo") - .set_channel_name("bar") - .spawn(pool) - .await - .unwrap(); - JobBuilder::new("foo") - .set_channel_name("bar") - .spawn(pool) - .await - .unwrap(); - JobBuilder::new("foo") - .set_channel_name("baz") - .spawn(pool) - .await - .unwrap(); - JobBuilder::new("foo") - .set_channel_name("baz") - .spawn(pool) - .await - .unwrap(); + sqlxmq::clear(&pool, &["foo", "baz"]).await.unwrap(); - sqlxmq::clear(pool, &["foo", "baz"]).await.unwrap(); + // Clearing deletes the messages outright, so what survived can be + // checked directly rather than inferred from what fails to run. + assert_eq!(queue_len(&pool).await, 2); - let (_runner, counter) = - test_job_runner(pool, |mut job| async move { job.complete().await }).await; + let (mut runner, mut jobs) = test_job_runner(&pool).await; - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 2); + let mut delivered = HashSet::new(); + for mut job in jobs.next_n(2).await { + delivered.insert(job.id()); + job.complete().await.unwrap(); } - pause().await; - } - - #[tokio::test] - async fn it_can_spawn_batch_of_jobs() { - { - let pool = &*test_pool().await; - let (tx, mut rx) = mpsc::unbounded(); + assert_eq!(delivered, kept); - let (_runner, counter) = test_job_runner(pool, move |mut job| { - let tx = tx.clone(); - async move { - let payload: Option = job.json().unwrap(); - tx.unbounded_send((job.name().to_owned(), payload)).unwrap(); - job.complete().await.unwrap(); - } - }) - .await; + assert_eq!(queue_len(&pool).await, 0); + runner.stop().await; + } - let mut job_a = JobBuilder::new("a"); - job_a.set_json(&"first").unwrap(); - let mut job_b = JobBuilder::new("b"); - job_b.set_json(&"second").unwrap(); - let job_c = JobBuilder::new("c"); + #[sqlx::test] + async fn it_can_spawn_batch_of_jobs(pool: Pool) { + let (mut runner, mut jobs) = test_job_runner(&pool).await; - let ids = spawn_batch(pool, &[job_a, job_b, job_c]).await.unwrap(); - assert_eq!(ids.len(), 3); + let mut job_a = JobBuilder::new("a"); + job_a.set_json(&"first").unwrap(); + let mut job_b = JobBuilder::new("b"); + job_b.set_json(&"second").unwrap(); + let job_c = JobBuilder::new("c"); - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 3); + let ids = spawn_batch(&pool, &[job_a, job_b, job_c]).await.unwrap(); + assert_eq!(ids.len(), 3); - let mut received = Vec::new(); - for _ in 0..3 { - received.push(rx.next().await.unwrap()); - } - received.sort(); - assert_eq!( - received, - vec![ - ("a".to_owned(), Some("first".to_owned())), - ("b".to_owned(), Some("second".to_owned())), - ("c".to_owned(), None), - ] - ); + let mut received = Vec::new(); + for mut job in jobs.next_n(3).await { + let payload: Option = job.json().unwrap(); + received.push((job.name().to_owned(), payload)); + job.complete().await.unwrap(); } - pause().await; + received.sort(); + assert_eq!( + received, + vec![ + ("a".to_owned(), Some("first".to_owned())), + ("b".to_owned(), Some("second".to_owned())), + ("c".to_owned(), None), + ] + ); + + assert_eq!(queue_len(&pool).await, 0); + runner.stop().await; } - #[tokio::test] - async fn it_chains_ordered_jobs_spawned_in_batch() { - { - let pool = &*test_pool().await; - let (tx, mut rx) = mpsc::unbounded(); + #[sqlx::test] + async fn it_chains_ordered_jobs_spawned_in_batch(pool: Pool) { + let (mut runner, mut jobs) = test_job_runner(&pool).await; - let (_runner, counter) = test_job_runner(pool, move |job| { - let tx = tx.clone(); - async move { - tx.unbounded_send(job).unwrap(); - } + let builders: Vec<_> = ["a", "b", "c"] + .iter() + .copied() + .map(|name| { + let mut builder = JobBuilder::new(name); + builder.set_ordered(true); + builder }) - .await; - - let mut job_a = JobBuilder::new("a"); - job_a.set_ordered(true); - let mut job_b = JobBuilder::new("b"); - job_b.set_ordered(true); - let mut job_c = JobBuilder::new("c"); - job_c.set_ordered(true); - - spawn_batch(pool, &[job_a, job_b, job_c]).await.unwrap(); - - // Only the first job in the chain should be delivered, and jobs - // must be delivered in the order they appeared in the batch. - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 1); - - for (i, &expected_name) in ["a", "b", "c"].iter().enumerate() { - let mut job = rx.next().await.unwrap(); - assert_eq!(job.name(), expected_name); - job.complete().await.unwrap(); - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), (i + 2).min(3)); - } + .collect(); + let ids = spawn_batch(&pool, &builders).await.unwrap(); + + // The chain must follow the order of the batch, not the order of the + // randomly generated message ids. + assert_eq!(after_message_id(&pool, ids[0]).await, Some(Uuid::nil())); + assert_eq!(after_message_id(&pool, ids[1]).await, Some(ids[0])); + assert_eq!(after_message_id(&pool, ids[2]).await, Some(ids[1])); + + // `mq_poll` only considers messages at the head of a chain, so the + // assertions above are also what stops `b` and `c` running early. + for (i, expected_name) in ["a", "b", "c"].iter().copied().enumerate() { + let mut job = jobs.next().await; + assert_eq!(job.id(), ids[i]); + assert_eq!(job.name(), expected_name); + job.complete().await.unwrap(); } - pause().await; + + assert_eq!(queue_len(&pool).await, 0); + runner.stop().await; } - #[tokio::test] - async fn it_runs_jobs_in_order() { - { - let pool = &*test_pool().await; - let (tx, mut rx) = mpsc::unbounded(); + #[sqlx::test] + async fn it_runs_jobs_in_order(pool: Pool) { + let (mut runner, mut jobs) = test_job_runner(&pool).await; - let (_runner, counter) = test_job_runner(pool, move |job| { - let tx = tx.clone(); - async move { - tx.unbounded_send(job).unwrap(); - } - }) - .await; + let first = JobBuilder::new("foo") + .set_ordered(true) + .spawn(&pool) + .await + .unwrap(); + let second = JobBuilder::new("bar") + .set_ordered(true) + .spawn(&pool) + .await + .unwrap(); - assert_eq!(counter.load(Ordering::SeqCst), 0); - JobBuilder::new("foo") - .set_ordered(true) - .spawn(pool) - .await - .unwrap(); - JobBuilder::new("bar") - .set_ordered(true) - .spawn(pool) - .await - .unwrap(); + // Spawned one at a time, `bar` still chains behind `foo`. + assert_eq!(after_message_id(&pool, first).await, Some(Uuid::nil())); + assert_eq!(after_message_id(&pool, second).await, Some(first)); - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 1); + let mut job = jobs.next().await; + assert_eq!(job.id(), first); + job.complete().await.unwrap(); - let mut job = rx.next().await.unwrap(); - job.complete().await.unwrap(); + let mut job = jobs.next().await; + assert_eq!(job.id(), second); + job.complete().await.unwrap(); - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 2); - } - pause().await; + assert_eq!(queue_len(&pool).await, 0); + runner.stop().await; } - #[tokio::test] - async fn it_runs_jobs_in_parallel() { - { - let pool = &*test_pool().await; - let (tx, mut rx) = mpsc::unbounded(); - - let (_runner, counter) = test_job_runner(pool, move |job| { - let tx = tx.clone(); - async move { - tx.unbounded_send(job).unwrap(); - } - }) - .await; + #[sqlx::test] + async fn it_runs_jobs_in_parallel(pool: Pool) { + let (mut runner, mut jobs) = test_job_runner(&pool).await; - assert_eq!(counter.load(Ordering::SeqCst), 0); - JobBuilder::new("foo").spawn(pool).await.unwrap(); - JobBuilder::new("bar").spawn(pool).await.unwrap(); + let mut spawned = HashSet::new(); + spawned.insert(JobBuilder::new("foo").spawn(&pool).await.unwrap()); + spawned.insert(JobBuilder::new("bar").spawn(&pool).await.unwrap()); - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 2); + // Neither job is completed before the other is dispatched: unordered + // jobs must not wait for each other. + let dispatched = jobs.next_n(2).await; + assert_eq!( + dispatched + .iter() + .map(|job| job.id()) + .collect::>(), + spawned + ); - for _ in 0..2 { - let mut job = rx.next().await.unwrap(); - job.complete().await.unwrap(); - } + for mut job in dispatched { + job.complete().await.unwrap(); } - pause().await; + + assert_eq!(queue_len(&pool).await, 0); + runner.stop().await; } - #[tokio::test] - async fn it_retries_failed_jobs() { - { - let pool = &*test_pool().await; - let (_runner, counter) = test_job_runner(pool, move |_| async {}).await; + /// A job which is never completed is retried, with the backoff doubling + /// each time, until its attempts run out. + /// + /// Keep-alive is switched off here: its whole purpose is to postpone the + /// retry of a job that is still running, which is precisely what this + /// test measures. With it on, the retry schedule depends on how quickly + /// the dispatched job happens to be dropped. + #[sqlx::test] + async fn it_retries_failed_jobs(pool: Pool) { + let (mut runner, mut jobs) = + configured_job_runner(&pool, |options| options.set_keep_alive(false)).await; + + let start = Instant::now(); + let id = JobBuilder::new("foo") + .set_retry_backoff(BACKOFF) + .set_retries(2) + .spawn(&pool) + .await + .unwrap(); - let backoff = default_pause() + 300; + // The initial attempt plus two retries. Dropping a job without + // completing it is what marks the attempt as failed. + for _ in 0..3 { + let job = jobs.next().await; + assert_eq!(job.id(), id); + drop(job); + } - assert_eq!(counter.load(Ordering::SeqCst), 0); - JobBuilder::new("foo") - .set_retry_backoff(Duration::from_millis(backoff)) - .set_retries(2) - .spawn(pool) - .await - .unwrap(); + // Each attempt is scheduled a backoff ahead of the previous one, and + // `start` precedes the first attempt, so the elapsed time is a sound + // lower bound: it fails if `mq_poll` ever hands out a message before + // its `attempt_at`. + assert!( + start.elapsed() >= BACKOFF + BACKOFF * 2, + "retries were delivered earlier than their backoff allows" + ); - // First attempt - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 1); + // `mq_poll` clears `attempt_at` as it hands out the last attempt, so + // the message can never be polled again. That makes the state below + // final, and establishes that there is no fourth attempt without + // having to wait to see whether one turns up. + assert_eq!(attempt_state(&pool, id).await, (0, false)); - // Second attempt - pause_ms(backoff).await; - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 2); + // Three attempts, so the backoff has been doubled three times. This + // is only read once the message is exhausted: reading it between + // attempts would race the next poll. + assert_eq!(retry_backoff(&pool, id).await, BACKOFF * 8); - // Third attempt - pause_ms(backoff * 2).await; - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 3); + runner.stop().await; + } - // No more attempts - pause_ms(backoff * 5).await; - assert_eq!(counter.load(Ordering::SeqCst), 3); - } - pause().await; - } - - #[tokio::test] - async fn it_can_checkpoint_jobs() { - { - let pool = &*test_pool().await; - let (_runner, counter) = test_job_runner(pool, move |mut current_job| async move { - let state: bool = current_job.json().unwrap().unwrap(); - if state { - current_job.complete().await.unwrap(); - } else { - current_job - .checkpoint(Checkpoint::new().set_json(&true).unwrap()) - .await - .unwrap(); - } - }) - .await; + /// Checkpointing replaces the payload used by the next attempt without + /// consuming one, so the retry picks up where the first attempt left off. + /// + /// The runner is stopped for the duration of the checkpoint and restarted + /// afterwards. Nothing about checkpointing requires that, but it removes + /// the one window in which this test could race the implementation: a + /// checkpoint carries no keep-alive of its own, so a runner left polling + /// may hand out the retry before the new payload is committed and read + /// the old one. Stopping the runner closes the window outright rather + /// than making the backoff long enough to hide it. + /// + /// Keep-alive is switched off for the same reason as in + /// `it_retries_failed_jobs`. + #[sqlx::test] + async fn it_can_checkpoint_jobs(pool: Pool) { + let (mut runner, mut jobs) = + configured_job_runner(&pool, |options| options.set_keep_alive(false)).await; + + let id = JobBuilder::new("foo") + .set_retry_backoff(BACKOFF) + .set_retries(5) + .set_json(&false) + .unwrap() + .spawn(&pool) + .await + .unwrap(); - let backoff = default_pause(); + let mut job = jobs.next().await; + assert_eq!(job.id(), id); + assert_eq!(job.json::().unwrap(), Some(false)); - assert_eq!(counter.load(Ordering::SeqCst), 0); - JobBuilder::new("foo") - .set_retry_backoff(Duration::from_millis(backoff)) - .set_retries(5) - .set_json(&false) - .unwrap() - .spawn(pool) - .await - .unwrap(); + runner.stop().await; - // First attempt - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 1); + job.checkpoint(Checkpoint::new().set_json(&true).unwrap()) + .await + .unwrap(); + assert_eq!( + stored_payload(&pool, id).await.as_deref(), + Some("true"), + "the checkpoint should have replaced the stored payload" + ); - // Second attempt - pause_ms(backoff).await; - assert_eq!(counter.load(Ordering::SeqCst), 2); + // The attempt ends without the job being completed, so it is retried. + drop(job); - // No more attempts - pause_ms(backoff * 3).await; - assert_eq!(counter.load(Ordering::SeqCst), 2); - } - pause().await; - } + let (mut runner, mut jobs) = + configured_job_runner(&pool, |options| options.set_keep_alive(false)).await; + + let mut job = jobs.next().await; + assert_eq!(job.id(), id); + assert_eq!( + job.json::().unwrap(), + Some(true), + "the retry should see the checkpointed payload" + ); + job.complete().await.unwrap(); + + runner.stop().await; - #[tokio::test] - async fn it_can_use_registry() { - { - let pool = &*test_pool().await; - let _runner = named_job_runner(pool).await; + // Completing deletes the message, so there is no third attempt to + // wait for. + assert!(!exists(&pool, id).await.unwrap()); + assert_eq!(queue_len(&pool).await, 0); + } - example_job1.builder().spawn(pool).await.unwrap(); - example_job2.builder().spawn(pool).await.unwrap(); - example_job_with_ctx.builder().spawn(pool).await.unwrap(); - pause().await; + #[sqlx::test] + async fn it_can_use_registry(pool: Pool) { + let (ran_tx, mut ran_rx) = mpsc::unbounded(); + let mut runner = named_job_runner(&pool, ran_tx).await; + + example_job1.builder().spawn(&pool).await.unwrap(); + example_job2.builder().spawn(&pool).await.unwrap(); + example_job_with_ctx.builder().spawn(&pool).await.unwrap(); + + let mut ran = Vec::new(); + for _ in 0..3 { + ran.push( + tokio::time::timeout(TIMEOUT, ran_rx.next()) + .await + .expect("timed out waiting for the registered jobs to run") + .expect("the job runner stopped before every job ran"), + ); } - pause().await; + ran.sort_unstable(); + assert_eq!( + ran, + ["example_job1", "example_job2", "example_job_with_ctx"] + ); + + assert_eq!(queue_len(&pool).await, 0); + runner.stop().await; } /// `mq_poll` claims candidate rows with `FOR UPDATE SKIP LOCKED`, so a @@ -831,17 +906,15 @@ mod tests { /// transaction's row locks and this test times out; without the /// `MATERIALIZED` CTE the locking subquery can be re-executed and return /// rows already claimed, and the disjointness assertion fails. - #[tokio::test] - async fn it_polls_disjoint_messages_from_concurrent_transactions() { - let pool = &*test_pool().await; - + #[sqlx::test] + async fn it_polls_disjoint_messages_from_concurrent_transactions(pool: Pool) { const NUM_JOBS: usize = 4; const BATCH_SIZE: i32 = 2; let builders: Vec<_> = (0..NUM_JOBS) .map(|_| JobBuilder::new("concurrent")) .collect(); - let spawned: HashSet = spawn_batch(pool, &builders) + let spawned: HashSet = spawn_batch(&pool, &builders) .await .unwrap() .into_iter() @@ -864,10 +937,9 @@ mod tests { assert_eq!(claimed_a.len(), BATCH_SIZE as usize); // `tx_b` must skip what `tx_a` holds rather than waiting for it. - let claimed_b = - tokio::time::timeout(Duration::from_secs(10), poll_ids(&mut tx_b, BATCH_SIZE)) - .await - .expect("concurrent poll blocked on the first transaction's row locks"); + let claimed_b = tokio::time::timeout(TIMEOUT, poll_ids(&mut tx_b, BATCH_SIZE)) + .await + .expect("concurrent poll blocked on the first transaction's row locks"); assert_eq!(claimed_b.len(), BATCH_SIZE as usize); tx_a.commit().await.unwrap(); @@ -909,26 +981,21 @@ mod tests { /// Unlike the test above this one does not single out `SKIP LOCKED` — it /// guards the delivery contract that the locking strategy exists to /// uphold, and so stays valid if that strategy is changed again. - #[tokio::test(flavor = "multi_thread", worker_threads = 4)] - async fn it_delivers_each_job_exactly_once_with_concurrent_runners() { - let pool = &*test_pool().await; - + #[sqlx::test] + async fn it_delivers_each_job_exactly_once_with_concurrent_runners(pool: Pool) { const NUM_RUNNERS: usize = 4; const NUM_JOBS: usize = 24; const NUM_CHANNELS: usize = 4; + init_logging(); + let (tx, mut rx) = mpsc::unbounded(); let mut runners = Vec::new(); for runner_idx in 0..NUM_RUNNERS { let tx = tx.clone(); - let runner = JobRunnerOptions::new(pool, move |mut job| { - let tx = tx.clone(); - task::spawn(async move { - let id = job.id(); - job.complete().await.unwrap(); - tx.unbounded_send((runner_idx, id)).unwrap(); - }); + let runner = JobRunnerOptions::new(&pool, move |job| { + let _ = tx.unbounded_send((runner_idx, job)); }) .set_concurrency(2, 6) .run() @@ -936,8 +1003,7 @@ mod tests { .unwrap(); runners.push(runner); } - // Drop the extra sender so the receiver terminates if every runner - // goes away. + // Drop the spare sender so the receiver ends if every runner stops. drop(tx); // Spread the jobs over several channels so that `mq_active_channels` @@ -953,7 +1019,7 @@ mod tests { builder }) .collect(); - let spawned: HashSet = spawn_batch(pool, &builders) + let spawned: HashSet = spawn_batch(&pool, &builders) .await .unwrap() .into_iter() @@ -963,16 +1029,17 @@ mod tests { let mut delivered = HashSet::new(); let mut delivering_runners = HashSet::new(); for _ in 0..NUM_JOBS { - let (runner_idx, id) = tokio::time::timeout(Duration::from_secs(60), rx.next()) + let (runner_idx, mut job) = tokio::time::timeout(TIMEOUT, rx.next()) .await .expect("timed out waiting for jobs to be delivered") .expect("all runners stopped before every job was delivered"); assert!( - delivered.insert(id), + delivered.insert(job.id()), "job {} was delivered more than once", - id + job.id() ); delivering_runners.insert(runner_idx); + job.complete().await.unwrap(); } assert_eq!(delivered, spawned); @@ -980,13 +1047,7 @@ mod tests { assert!(rx.try_recv().is_err(), "a job was delivered more than once"); // Every job was completed, so the queue must be empty. - for id in &spawned { - assert!( - !exists(pool, *id).await.unwrap(), - "job {} was not deleted", - id - ); - } + assert_eq!(queue_len(&pool).await, 0); log::info!( "{} of {} runners took part in delivery",