[rust] complete extension hooks

rust/riverqueue/src/__private.rs

@@ -60,6 +60,17 @@ impl InsertOptsExt for crate::InsertOpts {
     }
 }
 
+impl InsertOptsExt for crate::InsertParams {
+    fn extension_options(&self) -> &Map<String, Value> {
+        &self.extension_options
+    }
+
+    fn with_extension_option(mut self, key: impl Into<String>, 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<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

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

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::<std::collections::BTreeSet<_>>();
+        let input_metadata = intercepts.then(|| {
+            jobs.iter()
+                .map(|job| job.opts.metadata.clone())
+                .collect::<Vec<_>>()
+        });
         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::<Vec<_>>();
+        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::<Vec<_>>();
-        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

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/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<Vec<ClaimedJob>, 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<Vec<ClaimedJob>, 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)