[Rust] Let extensions run work hooks on peer jobs

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<crate::WorkOutcome, crate::BoxError>,
+    ) -> Result<crate::WorkOutcome, crate::BoxError> {
+        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 {

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::<std::io::Error>().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::<std::io::Error>()
+                .is_some()
+        );
+    }
+}

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::<Self>() {
+            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) {

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<Mutex<Vec<String>>>,
+}
+
+#[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<WorkOutcome, WorkError>,
+    ) -> Result<WorkOutcome, WorkError> {
+        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::<PeerFailure>().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::<riverqueue::Error>()
+                        .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::<std::io::Error>().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::<PeerArgs>().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;
+    }
 }