diff --git a/rust/riverqueue/src/__private.rs b/rust/riverqueue/src/__private.rs index 50d18e6ac..5c3f36930 100644 --- a/rust/riverqueue/src/__private.rs +++ b/rust/riverqueue/src/__private.rs @@ -60,6 +60,17 @@ impl InsertOptsExt for crate::InsertOpts { } } +impl InsertOptsExt for crate::InsertParams { + fn extension_options(&self) -> &Map { + &self.extension_options + } + + fn with_extension_option(mut self, key: impl Into, value: Value) -> Self { + self.extension_options.insert(key.into(), value); + self + } +} + /// Queue configuration reserved for River's own companion crates. /// /// Extension settings travel with a queue's configuration, through @@ -679,12 +690,20 @@ pub struct JobSetStateParams<'a> { /// Rows passed to [`Pilot::after_jobs_inserted`]. #[derive(Clone, Debug)] +#[non_exhaustive] pub struct JobsInsertedParams<'a> { /// Selected database backend configuration. pub database: DatabaseConfig, - /// Jobs the insertion wrote, excluding unique insertions skipped as - /// duplicates, in input order. + /// Metadata of every requested job, in input order, after insert hooks + /// and [`Pilot::before_jobs_insert`]. A unique duplicate keeps its + /// requested metadata here. + pub input_metadata: &'a [crate::JobMetadata], + /// Every returned job, including existing unique duplicates, in input + /// order. A duplicate keeps the stored row's metadata here. pub jobs: &'a [JobRow], + /// Whether each returned job is an existing unique duplicate rather + /// than a new row, in input order. + pub unique_skipped_as_duplicate: &'a [bool], } /// Mutable job insertion fields exposed to an exact-version extension. @@ -1070,15 +1089,15 @@ pub trait Pilot: std::any::Any + Send + Sync + 'static { } /// Runs after River writes a batch of inserted jobs, inside the insertion - /// transaction, like the post-insert work in River Go's - /// `Pilot.JobInsertMany`. + /// transaction. /// /// Called only when [`Pilot::intercepts_insert`] returns `true`, on every /// insertion path, including batches, caller-managed transactions, and - /// periodic jobs. Unique insertions skipped as duplicates aren't - /// included. Returning an error rolls back the insertion. As in - /// [`Pilot::before_jobs_insert`], `connection` is the insertion's - /// transaction, which the extension must not end. + /// periodic jobs. The rows include existing unique duplicates, flagged + /// as such, alongside the requested metadata. Returning an error rolls + /// back the insertion. As in [`Pilot::before_jobs_insert`], + /// `connection` is the insertion's transaction, which the extension must + /// not end. async fn after_jobs_inserted( &self, _connection: DatabaseConnection<'_>, @@ -1355,6 +1374,10 @@ impl PeerOutcome { /// being owned when its outcome persists, before its event, so it can be /// claimed again at once. /// +/// River doesn't run work hooks for peers on its own, because it doesn't +/// work them. A coordinator that works a peer runs them with +/// [`PeerAttempts::work_begin`] and [`PeerAttempts::work_end`]. +/// /// A soft stop doesn't end peer operations. A coordinator keeps claiming and /// completing peers after its producer stops fetching new jobs, until its /// attempt ends, and the client's stop waits for the attempt and so for @@ -1448,6 +1471,65 @@ impl<'a> PeerAttempts<'a> { let (client, peers) = self.attempt()?; peers.complete(&client.inner, self.context, outcomes).await } + + /// Runs the client's [`Hook::work_begin`](crate::Hook::work_begin) hooks + /// for `job` as River does before working a job itself: in registration + /// order, with the coordinator's work context, each seeing the changes + /// earlier hooks made to `job`. The first failing hook stops the rest. + /// + /// Call it on a copy of a claimed row before decoding its arguments, and + /// give the changed row to [`PeerAttempts::work_end`]. River completes a + /// peer from the row its claim returned, so changes made here aren't + /// persisted. A panicking hook panics the caller. + /// + /// Hooks see the coordinator's work context. Metadata a hook records + /// goes to the coordinator, and River merges it into the coordinator + /// and into every peer completed after that. + /// + /// # Errors + /// + /// Returns the failing hook's error unchanged, for the caller to fail + /// that job alone with a [`PeerOutcome`], and an + /// [`Error::Extension`](crate::Error::Extension) error, boxed, when + /// there's no running attempt. + pub async fn work_begin(self, job: &mut JobRow) -> Result<(), crate::BoxError> { + let (client, _) = self.attempt()?; + for hook in &client.inner.hooks { + hook.work_begin(self.context, job).await?; + } + Ok(()) + } + + /// Runs the client's [`Hook::work_end`](crate::Hook::work_end) hooks for + /// `job` and its result as River does after working a job itself: in + /// registration order, with the coordinator's work context, each + /// receiving the previous hook's result. Returns the final result, which + /// replaces `result` in the job's [`PeerOutcome`]. + /// + /// `job` should be the row [`PeerAttempts::work_begin`] changed, as + /// River passes the same row to both hooks. Work middleware and the job + /// timeout don't apply to peers. A panicking hook panics the caller. + /// + /// # Errors + /// + /// Returns the hooks' final error as a boxed + /// [`WorkError`](crate::WorkError). Give it to the job's + /// [`PeerOutcome`] unchanged: River unwraps it, so hooks and the error + /// handler can downcast the original error as they would a worker's. + /// Returns an [`Error::Extension`](crate::Error::Extension) error, + /// boxed, in place of `result` when there's no running attempt. + pub async fn work_end( + self, + job: &JobRow, + result: Result, + ) -> Result { + let (client, _) = self.attempt()?; + let mut result = result.map_err(crate::WorkError::from_boxed); + for hook in &client.inner.hooks { + result = hook.work_end(self.context, job, result).await; + } + result.map_err(|error| Box::new(error) as crate::BoxError) + } } impl RawInsertResult { diff --git a/rust/riverqueue/src/client/executor.rs b/rust/riverqueue/src/client/executor.rs index 1e71f8669..0cebc104f 100644 --- a/rust/riverqueue/src/client/executor.rs +++ b/rust/riverqueue/src/client/executor.rs @@ -233,7 +233,7 @@ async fn run_worker( for hook in &hooks_inner.hooks { hook.work_begin(&hooks_context, &mut job) .await - .map_err(WorkError::new)?; + .map_err(WorkError::from_boxed)?; } // The worker consumes the row, so it's copied only when // work-end hooks need it afterward. @@ -426,18 +426,7 @@ pub(super) fn worker_join_result( ) -> WorkerResult { match result { Ok(Ok(outcome)) => Ok(outcome), - // Like Go's `JobCancel`, a cancellation the worker returns cancels - // the job and records its reason. - Ok(Err(worker_error)) => Err(WorkerFailure { - error: Chain(&worker_error).to_string(), - kind: if JobCancelError::is_in_chain(worker_error.get_ref()) { - WorkerFailureKind::Cancelled - } else { - WorkerFailureKind::Error - }, - source: Some(worker_error), - trace: String::new(), - }), + Ok(Err(worker_error)) => Err(worker_failure(worker_error)), // Like Go, a panic is recorded as the panic's value alone. Err(join_error) if join_error.is_panic() => { let trace = format!("{join_error:?}"); @@ -457,16 +446,26 @@ pub(super) fn worker_join_result( } } -pub(super) fn worker_failure_from_source(error: BoxError) -> WorkerFailure { - let error = WorkError::new(error); +/// A `JobCancelError` anywhere in the error's chain cancels the job and +/// records its reason, whether a worker returned it or an extension reported +/// it for a peer. +fn worker_failure(error: WorkError) -> WorkerFailure { WorkerFailure { error: Chain(&error).to_string(), - kind: WorkerFailureKind::Error, + kind: if JobCancelError::is_in_chain(error.get_ref()) { + WorkerFailureKind::Cancelled + } else { + WorkerFailureKind::Error + }, source: Some(error), trace: String::new(), } } +pub(super) fn worker_failure_from_source(error: BoxError) -> WorkerFailure { + worker_failure(WorkError::from_boxed(error)) +} + pub(super) fn public_work_result(result: &WorkerResult) -> WorkResult { match result { Ok(WorkOutcome::Cancel) => WorkResult::Cancelled, @@ -771,3 +770,47 @@ mod snooze_counter_tests { } } } + +#[cfg(test)] +mod worker_failure_tests { + use crate::{BoxError, JobCancelError, WorkError}; + + use super::{WorkerFailureKind, worker_failure_from_source}; + + /// A boxed `WorkError`, as an extension returns after running work-end + /// hooks, keeps its original error rather than being wrapped again. + #[test] + fn boxed_work_errors_keep_their_original_error() { + let original = WorkError::new(std::io::Error::other("worker failed")); + let failure = worker_failure_from_source(Box::new(original) as BoxError); + let source = failure.source.unwrap(); + assert!(source.downcast_ref::().is_some()); + assert_eq!(failure.error, "worker failed"); + } + + /// An extension reporting a peer's cancellation, boxed directly or + /// after work-end hooks, cancels the peer like a worker's cancellation. + #[test] + fn job_cancel_errors_cancel() { + for error in [ + Box::new(JobCancelError::new("card expired")) as BoxError, + Box::new(WorkError::new(JobCancelError::new("card expired"))), + ] { + let failure = worker_failure_from_source(error); + assert!(matches!(failure.kind, WorkerFailureKind::Cancelled)); + assert_eq!(failure.error, "JobCancelError: card expired"); + } + } + + #[test] + fn other_errors_are_wrapped() { + let failure = worker_failure_from_source(Box::new(std::io::Error::other("worker failed"))); + assert!( + failure + .source + .unwrap() + .downcast_ref::() + .is_some() + ); + } +} diff --git a/rust/riverqueue/src/client/insert.rs b/rust/riverqueue/src/client/insert.rs index 81c136fb5..5cf507b43 100644 --- a/rust/riverqueue/src/client/insert.rs +++ b/rust/riverqueue/src/client/insert.rs @@ -722,12 +722,17 @@ impl Client { .filter(|job| job.state == JobState::Available) .map(|job| job.opts.queue.clone()) .collect::>(); + let input_metadata = intercepts.then(|| { + jobs.iter() + .map(|job| job.opts.metadata.clone()) + .collect::>() + }); let mut rows = Vec::with_capacity(jobs.len()); for job in jobs { rows.push(self.insert_row(connection.reborrow(), job).await?); } - if intercepts { - self.after_jobs_inserted(connection.reborrow(), &rows) + if let Some(input_metadata) = input_metadata { + self.after_jobs_inserted(connection.reborrow(), &rows, &input_metadata) .await?; } self.notify_insert( @@ -743,22 +748,22 @@ impl Client { &self, mut connection: PilotDatabaseConnection<'_>, rows: &[InsertedJob], + input_metadata: &[crate::JobMetadata], ) -> Result<(), Error> { - let inserted = rows + let returned = rows.iter().map(|row| row.job.clone()).collect::>(); + let unique_skipped_as_duplicate = rows .iter() - .filter(|row| !row.unique_skipped_as_duplicate) - .map(|row| row.job.clone()) + .map(|row| row.unique_skipped_as_duplicate) .collect::>(); - if inserted.is_empty() { - return Ok(()); - } self.inner .pilot .after_jobs_inserted( connection.reborrow(), &crate::__private::JobsInsertedParams { database: self.inner.pilot_database_config(), - jobs: &inserted, + input_metadata, + jobs: &returned, + unique_skipped_as_duplicate: &unique_skipped_as_duplicate, }, ) .await diff --git a/rust/riverqueue/src/extension.rs b/rust/riverqueue/src/extension.rs index 7b62a3353..26691ef96 100644 --- a/rust/riverqueue/src/extension.rs +++ b/rust/riverqueue/src/extension.rs @@ -27,6 +27,16 @@ impl WorkError { } } + /// Wraps a boxed error, reusing it when it's already a `WorkError` so a + /// result passed through a box keeps its original error for + /// [`WorkError::downcast_ref`]. + pub(crate) fn from_boxed(error: BoxError) -> Self { + match error.downcast::() { + Ok(error) => *error, + Err(error) => Self::new(error), + } + } + /// Returns the wrapped worker error for inspection. #[must_use] pub fn get_ref(&self) -> &(dyn std::error::Error + Send + Sync + 'static) { diff --git a/rust/riverqueue/src/pilot.rs b/rust/riverqueue/src/pilot.rs index 90d8835e1..4f5487477 100644 --- a/rust/riverqueue/src/pilot.rs +++ b/rust/riverqueue/src/pilot.rs @@ -366,8 +366,9 @@ pub struct ProducerClaimContext<'a> { /// River's standard claim, handed to [`PilotProducer::claim`]. /// -/// [`claim`](Self::claim) consumes it, so a session runs River's claim at most -/// once per call, on the transaction of its choosing. +/// [`claim`](Self::claim) and [`claim_at_most`](Self::claim_at_most) consume +/// it, so a session runs River's claim at most once per call, on the +/// transaction of its choosing. pub struct ProducerClaimNext<'a> { inner: &'a ClientInner, limit: usize, @@ -391,8 +392,28 @@ impl<'a> ProducerClaimNext<'a> { /// /// Returns the database error when the claim fails. pub async fn claim(self, connection: DatabaseConnection<'_>) -> Result, Error> { + let limit = self.limit; + self.claim_at_most(connection, limit).await + } + + /// Claims up to `maximum` jobs, capped by River's original claim limit, + /// on `connection`. A zero maximum performs no query. This consumes the + /// continuation just as [`claim`](Self::claim) does. + /// + /// # Errors + /// + /// Returns the database error when the claim fails. + pub async fn claim_at_most( + self, + connection: DatabaseConnection<'_>, + maximum: usize, + ) -> Result, Error> { + let limit = self.limit.min(maximum); + if limit == 0 { + return Ok(Vec::new()); + } Ok( - crate::client::standard_claim(self.inner, connection, self.queue, self.limit) + crate::client::standard_claim(self.inner, connection, self.queue, limit) .await? .into_iter() .map(ClaimedJob::from_decoded) diff --git a/rust/riverqueue/tests/extension_seams.rs b/rust/riverqueue/tests/extension_seams.rs index 053bc11d0..8c9fe4fd8 100644 --- a/rust/riverqueue/tests/extension_seams.rs +++ b/rust/riverqueue/tests/extension_seams.rs @@ -299,6 +299,19 @@ struct OptionsPilot { seen: Arc>>, } +impl riverqueue::Hook for OptionsPilot { + fn insert_begin( + &self, + insert: &mut riverqueue::InsertContext, + ) -> impl Future> + Send { + self.seen.lock().unwrap().push(( + insert.opts.extension_options().clone(), + insert.opts.metadata.as_raw().get().to_owned(), + )); + std::future::ready(Ok(())) + } +} + #[async_trait] impl Pilot for OptionsPilot { fn intercepts_insert(&self) -> bool { @@ -328,7 +341,11 @@ async fn assert_extension_options_reach_the_insert_hook( builder: impl Fn() -> riverqueue::ClientBuilder, ) { let pilot = OptionsPilot::default(); - let client = builder().pilot(pilot.clone()).build().unwrap(); + let client = builder() + .hook(pilot.clone()) + .pilot(pilot.clone()) + .build() + .unwrap(); let defaults = client.insert(DeclaredArgs { value: 1 }).await.unwrap(); let overridden = client @@ -342,7 +359,9 @@ async fn assert_extension_options_reach_the_insert_hook( .unwrap(); let seen = pilot.seen.lock().unwrap().clone(); - assert_eq!(seen.len(), 2); + assert_eq!(seen.len(), 4); + assert_eq!(seen[0], seen[1]); + assert_eq!(seen[2], seen[3]); assert_eq!( Value::Object(seen[0].0.clone()), json!({"declared": {"type": true}, "shared": "type"}) @@ -351,10 +370,10 @@ async fn assert_extension_options_reach_the_insert_hook( // A call's metadata replaces the defaults, but not the declared // extension options, which the call overrides key by key. assert_eq!( - Value::Object(seen[1].0.clone()), + Value::Object(seen[2].0.clone()), json!({"declared": {"type": true}, "shared": "call"}) ); - assert_eq!(seen[1].1, r#"{"call":1}"#); + assert_eq!(seen[2].1, r#"{"call":1}"#); // Nothing about the extension options is persisted. (SQLite adds its // insert nonce.) for (id, expected) in [ @@ -369,6 +388,79 @@ async fn assert_extension_options_reach_the_insert_hook( } } +type SeenResults = (Vec, Vec, Vec); + +/// Records every returned row, including an existing unique row. +#[derive(Clone, Default)] +struct InsertResultsPilot { + seen: Arc>>, +} + +#[async_trait] +impl Pilot for InsertResultsPilot { + fn intercepts_insert(&self) -> bool { + true + } + + async fn after_jobs_inserted( + &self, + _connection: DatabaseConnection<'_>, + params: &riverqueue::__private::JobsInsertedParams<'_>, + ) -> Result<(), PilotError> { + self.seen.lock().unwrap().push(( + params.jobs.iter().map(|job| job.id).collect(), + params + .input_metadata + .iter() + .map(|metadata| metadata.get::("input").unwrap().unwrap()) + .collect(), + params.unique_skipped_as_duplicate.to_vec(), + )); + Ok(()) + } +} + +async fn assert_insert_results_include_unique_duplicates( + builder: impl Fn() -> riverqueue::ClientBuilder, +) { + let pilot = InsertResultsPilot::default(); + let client = builder().pilot(pilot.clone()).build().unwrap(); + let opts = |input| { + InsertOpts::default() + .with_unique(riverqueue::UniqueOpts::new().with_by_args(true)) + .with_metadata(json!({"input": input}).as_object().unwrap().clone()) + }; + let first = client + .insert(SeamArgs { value: 1 }) + .opts(opts(1)) + .await + .unwrap(); + let duplicate = client + .insert(SeamArgs { value: 1 }) + .opts(opts(2)) + .await + .unwrap(); + assert!(duplicate.unique_skipped_as_duplicate); + let mut batch = riverqueue::InsertBatch::new(); + batch.push_with(SeamArgs { value: 1 }, opts(3)); + batch.push_with(SeamArgs { value: 2 }, opts(4)); + let mixed = client.insert_batch(batch).await.unwrap(); + assert!(mixed[0].unique_skipped_as_duplicate); + assert!(!mixed[1].unique_skipped_as_duplicate); + assert_eq!( + *pilot.seen.lock().unwrap(), + [ + (vec![first.id()], vec![1], vec![false]), + (vec![first.id()], vec![2], vec![true]), + ( + vec![first.id(), mixed[1].job.id], + vec![3, 4], + vec![true, false] + ) + ], + ); +} + #[cfg(all(feature = "postgres", river_postgres_tests))] mod postgres { use riverqueue::database::PostgresDatabase; @@ -503,6 +595,13 @@ mod postgres { schema.cleanup().await; } + #[tokio::test(flavor = "multi_thread")] + async fn insert_results_include_unique_duplicates() { + let schema = PostgresSchema::new("seam_unique_results").await; + assert_insert_results_include_unique_duplicates(|| builder(&schema)).await; + schema.cleanup().await; + } + #[tokio::test(flavor = "multi_thread")] async fn insertions_reach_the_extension_as_one_batch() { let schema = PostgresSchema::new("seam_batch_insert").await; @@ -623,6 +722,13 @@ mod sqlite { sqlite_cleanup(pool, path).await; } + #[tokio::test(flavor = "multi_thread")] + async fn insert_results_include_unique_duplicates() { + let (pool, path) = sqlite_file_pool(4).await; + assert_insert_results_include_unique_duplicates(|| Client::builder(pool.clone())).await; + sqlite_cleanup(pool, path).await; + } + #[tokio::test(flavor = "multi_thread")] async fn insertions_reach_the_extension_as_one_batch() { let (pool, path) = sqlite_file_pool(4).await; diff --git a/rust/riverqueue/tests/peer_attempts.rs b/rust/riverqueue/tests/peer_attempts.rs index bca12d849..1ece92bf6 100644 --- a/rust/riverqueue/tests/peer_attempts.rs +++ b/rust/riverqueue/tests/peer_attempts.rs @@ -22,8 +22,9 @@ use riverqueue::__private::{ PeerClaimContext, PeerOutcome, Pilot, PilotError, PilotProducer, ProducerStartContext, }; use riverqueue::{ - BoxError, Client, ErrorHandler, ErrorHandlerDecision, Event, EventKind, InsertOpts, Job, - JobArgs, JobRow, JobState, QueueConfig, WorkContext, WorkOutcome, WorkResult, Workers, + BoxError, Client, ErrorHandler, ErrorHandlerDecision, Event, EventKind, Hook, InsertOpts, Job, + JobArgs, JobRow, JobState, QueueConfig, WorkContext, WorkError, WorkOutcome, WorkResult, + Workers, }; use serde::{Deserialize, Serialize}; use tokio::sync::{Notify, mpsc}; @@ -1366,6 +1367,188 @@ async fn assert_abandoned_peers_reclaimed(mut run: Run, unblock: std::sync::mpsc run.stop().await; } +/// A peer worker's own error, which work-end hooks recognize by type. +#[derive(Debug)] +struct PeerFailure; + +impl std::fmt::Display for PeerFailure { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("peer failed on purpose") + } +} + +impl std::error::Error for PeerFailure {} + +/// Records work hooks for peer jobs. The hook named `first` fails peers +/// whose `value` is 1 and multiplies the others' by ten, like a hook that +/// decodes arguments. The one named `second` turns a [`PeerFailure`] into a +/// discard. +struct PeerHook { + name: &'static str, + trace: Arc>>, +} + +#[allow( + clippy::unused_async_trait_impl, + reason = "the hook only records state synchronously" +)] +impl Hook for PeerHook { + async fn work_begin(&self, _context: &WorkContext, job: &mut JobRow) -> Result<(), BoxError> { + if job.kind != PeerArgs::KIND { + return Ok(()); + } + let mut args: PeerArgs = job.decode_args()?; + self.trace + .lock() + .unwrap() + .push(format!("{} begin {}", self.name, args.value)); + if self.name == "first" { + if args.value == 1 { + return Err(Box::new(std::io::Error::other("begin failed on purpose"))); + } + args.value *= 10; + job.encoded_args = serde_json::value::to_raw_value(&args)?; + } + Ok(()) + } + + async fn work_end( + &self, + _context: &WorkContext, + job: &JobRow, + result: Result, + ) -> Result { + if job.kind != PeerArgs::KIND { + return result; + } + let args: PeerArgs = job.decode_args().unwrap(); + self.trace + .lock() + .unwrap() + .push(format!("{} end {}", self.name, args.value)); + match result { + Err(error) + if self.name == "second" && error.downcast_ref::().is_some() => + { + Ok(WorkOutcome::Discard) + } + result => result, + } + } +} + +/// Whether a peer operation failed because no attempt is running. +fn requires_running_attempt(result: Result) -> bool { + result.is_err_and(|error| { + error + .downcast_ref::() + .is_some_and(|error| peer_error(error, "require a running attempt")) + }) +} + +/// A coordinator runs the client's work hooks for its peers as River does +/// for a job it works itself: begin hooks in order, each seeing earlier +/// changes and the first failure stopping the rest, and end hooks in order +/// on the changed row, each replacing the result. A begin failure fails +/// only its own peer. +async fn assert_work_hooks_run_for_peers(builder: riverqueue::ClientBuilder, db: Db) { + let trace = Arc::new(Mutex::new(Vec::new())); + let builder = builder + .hook(PeerHook { + name: "first", + trace: Arc::clone(&trace), + }) + .hook(PeerHook { + name: "second", + trace: Arc::clone(&trace), + }); + let script: Script = Arc::new(|context, row, env| { + Box::pin(async move { + let mut checks = Checks::new(); + let detached_context = riverqueue::__private::work_context_for_job(&row); + let detached = PeerAttempts::new(&detached_context); + let mut detached_row = row.clone(); + check( + &mut checks, + "work-begin hooks require a running attempt", + requires_running_attempt(detached.work_begin(&mut detached_row).await), + ); + check( + &mut checks, + "work-end hooks require a running attempt", + requires_running_attempt(detached.work_end(&row, Ok(WorkOutcome::Complete)).await), + ); + + let mut rows = env.claim(&context, Take::Claim, &env.peers).await.unwrap(); + rows.sort_by_key(|row| row.id); + let peers = PeerAttempts::new(&context); + let mut outcomes = Vec::new(); + for row in rows { + let mut job = row.clone(); + let result = match peers.work_begin(&mut job).await { + Ok(()) => { + let args: PeerArgs = job.decode_args().unwrap(); + let result = if args.value == 20 { + Err(Box::new(PeerFailure) as BoxError) + } else { + Ok(WorkOutcome::Complete) + }; + peers.work_end(&job, result).await + } + Err(error) => { + check( + &mut checks, + "begin returns the hook's error", + error.downcast_ref::().is_some(), + ); + Err(error) + } + }; + outcomes.push(PeerOutcome::new(row, result)); + } + check( + &mut checks, + "complete succeeds", + peers.complete(outcomes).await.is_ok(), + ); + let _ = env.checks.send(checks); + WorkOutcome::Complete + }) + }); + let mut run = Run::start(builder, db, 3, script).await; + run.assert_checks().await; + let [p1, p2, p3] = [run.env.peers[0], run.env.peers[1], run.env.peers[2]]; + assert_eq!( + *trace.lock().unwrap(), + [ + "first begin 1", + "first begin 2", + "second begin 20", + "first end 20", + "second end 20", + "first begin 3", + "second begin 30", + "first end 30", + "second end 30", + ] + ); + // The begin failure failed only its own peer's attempt. + let failed = run.settled(p1).await; + assert_eq!( + failed.errors.last().unwrap().error, + "begin failed on purpose" + ); + // River completes peers from their claimed rows, so hook changes to + // arguments aren't persisted. + assert_eq!(failed.decode_args::().unwrap().value, 1); + assert_eq!(run.settled(p2).await.state, JobState::Discarded); + assert_eq!(run.settled(p3).await.state, JobState::Completed); + // The end hook replaced the second peer's failure before the error + // handler would have seen it. + assert_eq!(*run.handler.0.lock().unwrap(), [p1]); + run.stop().await; +} + #[cfg(all(feature = "postgres", river_postgres_tests))] mod postgres { use riverqueue::database::PostgresDatabase; @@ -1585,6 +1768,13 @@ mod postgres { assert_stops_interrupt_and_cancellations_fail(|| builder(&schema), db(&schema)).await; schema.cleanup().await; } + + #[tokio::test(flavor = "multi_thread")] + async fn work_hooks_run_for_peers() { + let schema = PostgresSchema::new("peer_work_hooks").await; + assert_work_hooks_run_for_peers(builder(&schema), db(&schema)).await; + schema.cleanup().await; + } } #[cfg(feature = "sqlite")] @@ -1689,4 +1879,12 @@ mod sqlite { .await; sqlite_cleanup(pool, path).await; } + + #[tokio::test(flavor = "multi_thread")] + async fn work_hooks_run_for_peers() { + let (pool, path) = sqlite_file_pool(4).await; + assert_work_hooks_run_for_peers(Client::builder(pool.clone()), Db::Sqlite(pool.clone())) + .await; + sqlite_cleanup(pool, path).await; + } } diff --git a/rust/riverqueue/tests/producer_sessions.rs b/rust/riverqueue/tests/producer_sessions.rs index c9da6397d..2ba240675 100644 --- a/rust/riverqueue/tests/producer_sessions.rs +++ b/rust/riverqueue/tests/producer_sessions.rs @@ -108,6 +108,8 @@ impl Recorder { enum Claim { /// River's claim in the session's own transaction. Standard, + /// River's claim with an extension's smaller maximum. + Capped(usize), /// River's claim, returning its first job twice. Duplicate, /// River's claim, reporting its first job as still available. @@ -131,6 +133,7 @@ enum Claim { #[derive(Clone)] struct SessionPilot { claim: Claim, + claims: Arc>, configurations: Arc>, finished: Arc>, } @@ -139,6 +142,7 @@ impl SessionPilot { fn new(claim: Claim) -> Self { Self { claim, + claims: Arc::default(), configurations: Arc::default(), finished: Arc::default(), } @@ -173,8 +177,15 @@ impl PilotProducer for SessionPilot { transaction.commit().await?; return Ok(jobs); } - let mut jobs = next.claim(transaction.connection()).await?; + let mut jobs = match self.claim { + Claim::Capped(maximum) => { + next.claim_at_most(transaction.connection(), maximum) + .await? + } + _ => next.claim(transaction.connection()).await?, + }; transaction.commit().await?; + self.claims.push(jobs.len()); let first = jobs.first().and_then(ClaimedJob::job).cloned(); match (self.claim, first) { (Claim::Duplicate, Some(first)) => jobs.push(first.into()), @@ -313,6 +324,38 @@ async fn assert_every_attempt_finishes_once(builder: riverqueue::ClientBuilder) } } +/// A companion extension can narrow the ordinary claim, but never widen +/// River's original limit. Zero leaves every job available. +async fn assert_standard_claim_cap(builder: riverqueue::ClientBuilder, maximum: usize) { + let pilot = SessionPilot::new(Claim::Capped(maximum)); + let client = builder + .pilot(pilot.clone()) + .without_leader_election() + .queue("default", fast_queue(3)) + .workers(workers()) + .build() + .unwrap(); + let inserted = client + .insert_many((0..3).map(|_| SessionArgs { fail: false })) + .await + .unwrap(); + let mut run = client.start().unwrap(); + pilot + .claims + .wait_until("the first capped claim", |claims| !claims.is_empty()) + .await; + run.stop().await.unwrap(); + assert_eq!(pilot.claims.snapshot()[0], maximum.min(3)); + if maximum == 0 { + for inserted in inserted { + assert_eq!( + client.jobs().get(inserted.id()).await.unwrap().state, + JobState::Available + ); + } + } +} + /// An attempt whose worker outlives its abort during shutdown leaves its job /// running for the rescuer, but still finishes in the session, so the /// extension doesn't count it against the queue for the rest of the run. @@ -707,6 +750,15 @@ mod postgres { schema.cleanup().await; } + #[tokio::test(flavor = "multi_thread")] + async fn caps_standard_claims() { + for maximum in [0, 1, usize::MAX] { + let schema = PostgresSchema::new("session_claim_cap").await; + assert_standard_claim_cap(builder(&schema), maximum).await; + schema.cleanup().await; + } + } + #[tokio::test(flavor = "multi_thread")] async fn cancellation_during_claim_reaches_the_attempt() { let schema = PostgresSchema::new("session_claim_cancel").await; @@ -800,6 +852,15 @@ mod sqlite { sqlite_cleanup(pool, path).await; } + #[tokio::test(flavor = "multi_thread")] + async fn caps_standard_claims() { + for maximum in [0, 1, usize::MAX] { + let (pool, path) = sqlite_file_pool(4).await; + assert_standard_claim_cap(Client::builder(pool.clone()), maximum).await; + sqlite_cleanup(pool, path).await; + } + } + #[tokio::test(flavor = "multi_thread")] async fn cancellation_during_claim_reaches_the_attempt() { let (pool, path) = sqlite_file_pool(4).await;