From 1d0186fb06c213db6c6cb8a35360d8c0e73335d9 Mon Sep 17 00:00:00 2001 From: Blake Gentry Date: Tue, 13 Oct 2026 07:45:18 -0700 Subject: [PATCH] [Rust] Let extensions run work hooks on peer jobs An extension that claims and completes peers through `PeerAttempts` works those jobs itself, so River never runs the client's `Hook::work_begin` and `Hook::work_end` hooks for them. Per-job processing that lives in those hooks, such as decoding arguments that were transformed on insertion, is skipped for every peer. Add `PeerAttempts::work_begin` and `PeerAttempts::work_end`, which run the client's configured hooks for a given row exactly as the executor does for a job it works: in registration order, with the coordinator's work context, begin hooks mutating the row and stopping at the first error, and end hooks threading the result from one to the next. The begin helper returns the failing hook's error unchanged so the caller can fail that one peer with a `PeerOutcome` and keep working the rest. `PeerAttempts` is the natural home because it already resolves the coordinator's client and running attempt, and keeps a peer's whole lifecycle (claim, work hooks, completion) on one type. The end helper takes and returns the boxed error type `PeerOutcome` uses. A `WorkError` passed through a box is now unwrapped rather than wrapped a second time when a peer's outcome is recorded, so hooks and the error handler can downcast a peer's error to its original type, as they can for a worker's error. New peer-attempt scenarios on Postgres and SQLite cover hook order, argument changes seen by later hooks and by end hooks, a begin failure failing only its own peer, an end hook replacing a failure before the error handler runs, and rejection outside a running attempt. --- rust/riverqueue/src/__private.rs | 58 ++++++++ rust/riverqueue/src/client/executor.rs | 32 +++- rust/riverqueue/src/extension.rs | 10 ++ rust/riverqueue/tests/peer_attempts.rs | 193 ++++++++++++++++++++++++- 4 files changed, 290 insertions(+), 3 deletions(-) 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; + } }