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 11f8dbb..bcd8262 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 @@ -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/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() diff --git a/src/lib.rs b/src/lib.rs index 093a9b2..f660a83 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -259,99 +259,295 @@ 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::env; + use std::collections::HashSet; 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; - - // field 0 is never read, but its drop is important - #[allow(dead_code)] - struct TestGuard(MutexGuard<'static, ()>, T); + use uuid::Uuid; + + /// 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 init_logging() { + static INIT_LOGGER: Once = Once::new(); + INIT_LOGGER.call_once(pretty_env_logger::init); + } - impl Deref for TestGuard { - type Target = T; + /// 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") + } - fn deref(&self) -> &T { - &self.1 + /// 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 } } - async fn test_pool() -> TestGuard> { - static INIT_LOGGER: Once = Once::new(); - static TEST_MUTEX: Mutex<()> = Mutex::const_new(()); - - let guard = TEST_MUTEX.lock().await; - - 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(()) } @@ -360,343 +556,507 @@ 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; - } + assert_eq!(delivered, kept); - #[tokio::test] - async fn it_can_spawn_batch_of_jobs() { - { - let pool = &*test_pool().await; - let (tx, mut rx) = mpsc::unbounded(); - - 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(); - - let (_runner, counter) = test_job_runner(pool, move |job| { - let tx = tx.clone(); - async move { - tx.unbounded_send(job).unwrap(); - } + #[sqlx::test] + async fn it_chains_ordered_jobs_spawned_in_batch(pool: Pool) { + let (mut runner, mut jobs) = test_job_runner(&pool).await; + + 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)); + + let mut job = jobs.next().await; + assert_eq!(job.id(), first); + 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), 1); + assert_eq!(queue_len(&pool).await, 0); + runner.stop().await; + } - let mut job = rx.next().await.unwrap(); + #[sqlx::test] + async fn it_runs_jobs_in_parallel(pool: Pool) { + let (mut runner, mut jobs) = test_job_runner(&pool).await; + + let mut spawned = HashSet::new(); + spawned.insert(JobBuilder::new("foo").spawn(&pool).await.unwrap()); + spawned.insert(JobBuilder::new("bar").spawn(&pool).await.unwrap()); + + // 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 mut job in dispatched { job.complete().await.unwrap(); + } + + assert_eq!(queue_len(&pool).await, 0); + runner.stop().await; + } - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 2); + /// 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(); + + // 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); } - pause().await; + + // 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" + ); + + // `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)); + + // 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); + + runner.stop().await; } - #[tokio::test] - async fn it_runs_jobs_in_parallel() { - { - let pool = &*test_pool().await; - let (tx, mut rx) = mpsc::unbounded(); + /// 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 (_runner, counter) = test_job_runner(pool, move |job| { - let tx = tx.clone(); - async move { - tx.unbounded_send(job).unwrap(); - } - }) - .await; + 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").spawn(pool).await.unwrap(); - JobBuilder::new("bar").spawn(pool).await.unwrap(); + runner.stop().await; - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 2); + 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" + ); + + // The attempt ends without the job being completed, so it is retried. + drop(job); + + 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; + + // 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); + } - for _ in 0..2 { - let mut job = rx.next().await.unwrap(); - job.complete().await.unwrap(); - } + #[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; } - #[tokio::test] - async fn it_retries_failed_jobs() { - { - let pool = &*test_pool().await; - let (_runner, counter) = test_job_runner(pool, move |_| async {}).await; - - let backoff = default_pause() + 300; - - assert_eq!(counter.load(Ordering::SeqCst), 0); - JobBuilder::new("foo") - .set_retry_backoff(Duration::from_millis(backoff)) - .set_retries(2) - .spawn(pool) + /// `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. + #[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) + .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(); + .unwrap() + } - // First attempt - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 1); + let mut tx_a = pool.begin().await.unwrap(); + let mut tx_b = pool.begin().await.unwrap(); - // Second attempt - pause_ms(backoff).await; - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 2); + // `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); - // Third attempt - pause_ms(backoff * 2).await; - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 3); + // `tx_b` must skip what `tx_a` holds rather than waiting for it. + 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(); + 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" + ); + } - // No more attempts - pause_ms(backoff * 5).await; - assert_eq!(counter.load(Ordering::SeqCst), 3); + /// 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. + #[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 |job| { + let _ = tx.unbounded_send((runner_idx, job)); + }) + .set_concurrency(2, 6) + .run() + .await + .unwrap(); + runners.push(runner); } - 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(); - } + // 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` + // 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 }) - .await; - - let backoff = default_pause(); - - 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) + .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, mut job) = tokio::time::timeout(TIMEOUT, rx.next()) .await - .unwrap(); - - // First attempt - pause().await; - assert_eq!(counter.load(Ordering::SeqCst), 1); + .expect("timed out waiting for jobs to be delivered") + .expect("all runners stopped before every job was delivered"); + assert!( + delivered.insert(job.id()), + "job {} was delivered more than once", + job.id() + ); + delivering_runners.insert(runner_idx); + job.complete().await.unwrap(); + } + assert_eq!(delivered, spawned); - // Second attempt - pause_ms(backoff).await; - assert_eq!(counter.load(Ordering::SeqCst), 2); + // Anything already queued beyond the expected count is a redelivery. + assert!(rx.try_recv().is_err(), "a job was delivered more than once"); - // No more attempts - pause_ms(backoff * 3).await; - assert_eq!(counter.load(Ordering::SeqCst), 2); - } - pause().await; - } + // Every job was completed, so the queue must be empty. + assert_eq!(queue_len(&pool).await, 0); - #[tokio::test] - async fn it_can_use_registry() { - { - let pool = &*test_pool().await; - let _runner = named_job_runner(pool).await; + log::info!( + "{} of {} runners took part in delivery", + delivering_runners.len(), + NUM_RUNNERS + ); - 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; + for mut runner in runners { + runner.stop().await; } - pause().await; } } 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); + } +}