From 0f7edd77a0c6df8a85c3c889928b699f45cd9c32 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Tue, 13 Oct 2026 07:45:18 -0700 Subject: [PATCH] [rust] complete extension hooks An extension that coordinates peer jobs works each peer itself, so the client's work hooks never run for peers. Add `PeerAttempts::work_begin` and `work_end`. They run the hooks in registration order and stop at the first failing begin hook, as the executor does. Keep the original error through hooks and outcomes. A boxed `WorkError` is reused instead of wrapped again, so hooks and the error handler can downcast the worker's error. A `JobCancelError` in a peer's error now cancels the peer, as it cancels a job a worker returns it from. Insertion extensions need resolved options in insert hooks and must see unique duplicates after persistence. Expose resolved options through the internal option trait. Pass every returned row to `after_jobs_inserted` in input order, with its requested metadata and a flag that marks existing duplicates. Add `ProducerClaimNext::claim_at_most`, so a producer session can claim fewer jobs than River's limit. The limit caps the requested maximum. A zero maximum runs no query. The continuation is consumed either way. --- rust/riverqueue/src/__private.rs | 98 +++++++++- rust/riverqueue/src/client/executor.rs | 75 ++++++-- rust/riverqueue/src/client/insert.rs | 23 ++- rust/riverqueue/src/extension.rs | 10 + rust/riverqueue/src/pilot.rs | 27 ++- rust/riverqueue/tests/extension_seams.rs | 114 +++++++++++- rust/riverqueue/tests/peer_attempts.rs | 202 ++++++++++++++++++++- rust/riverqueue/tests/producer_sessions.rs | 63 ++++++- 8 files changed, 569 insertions(+), 43 deletions(-) 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;