diff --git a/rust/riverqueue/src/__private.rs b/rust/riverqueue/src/__private.rs index 1c0c8ffb1..6e01bace5 100644 --- a/rust/riverqueue/src/__private.rs +++ b/rust/riverqueue/src/__private.rs @@ -1273,6 +1273,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 @@ -1366,6 +1370,60 @@ 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. + /// + /// # 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. An error keeps its type + /// through the hooks and the outcome, so hooks and the error handler can + /// downcast it as they would a worker's error. A panicking hook panics + /// the caller. + /// + /// # Errors + /// + /// Returns the hooks' final error, and 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 07520ee97..48ffb991e 100644 --- a/rust/riverqueue/src/client/executor.rs +++ b/rust/riverqueue/src/client/executor.rs @@ -452,7 +452,7 @@ pub(super) fn worker_join_result( } pub(super) fn worker_failure_from_source(error: BoxError) -> WorkerFailure { - let error = WorkError::new(error); + let error = WorkError::from_boxed(error); WorkerFailure { error: Chain(&error).to_string(), kind: WorkerFailureKind::Error, @@ -765,3 +765,33 @@ mod snooze_counter_tests { } } } + +#[cfg(test)] +mod worker_failure_tests { + use crate::{BoxError, WorkError}; + + use super::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"); + } + + #[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/extension.rs b/rust/riverqueue/src/extension.rs index bee8a944b..b17ed7008 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/tests/peer_attempts.rs b/rust/riverqueue/tests/peer_attempts.rs index bca12d849..561768cac 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,179 @@ 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, + } + } +} + +/// 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 mut detached_row = row.clone(); + let detached = PeerAttempts::new(&riverqueue::__private::work_context_for_job(&row)) + .work_begin(&mut detached_row) + .await; + check( + &mut checks, + "work hooks require a running attempt", + detached.is_err_and(|error| { + error + .downcast_ref::() + .is_some_and(|error| peer_error(error, "require a running attempt")) + }), + ); + + 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 +1759,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 +1870,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; + } }