Ruby: Make work_begin/work_end behave more like Go hooks

ruby/CHANGELOG.md

@@ -11,6 +11,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Allow `River::Workers#add` to register a kind and a work block, using the client's default retry and timeout policies. [PR #1486](https://github.com/riverqueue/river/pull/1486).
 
+### Changed
+
+- `work_end(job, error)` hooks now return the attempt's error, matching Go. Return the incoming error to preserve it, a replacement exception to change it, or `nil` to clear a failure. The return value feeds the next hook in registration order. Existing observing hooks must explicitly return `error`. End hooks are skipped if `work_begin` raises or the worker raises an exception outside `StandardError`. [PR #1520](https://github.com/riverqueue/river/pull/1520).
+
 ### Fixed
 
 - A `River::Config` without an explicit `id:` generates its default client ID on first use in each process, so clients built before a fork (as with preloading Puma, Unicorn, or Pitchfork servers) no longer share one ID across forked workers. Explicit IDs are unchanged, and `Config#with` copies share the original's ID. [PR #1498](https://github.com/riverqueue/river/pull/1498).

ruby/docs/README.md

@@ -582,7 +582,7 @@ A plugin may implement any combination of these methods:
 | Hook       | `insert_begin(params)`           | Before each job is inserted; `params` may be modified.                                       |
 | Hook       | `insert_end(result)`             | After each job is inserted.                                                                  |
 | Hook       | `work_begin(job)`                | After a job is claimed, immediately before its worker runs.                                  |
-| Hook       | `work_end(job, error)`           | After the worker returns or raises; `error` is `nil` on success.                             |
+| Hook       | `work_end(job, error)`           | After the worker returns or raises a `StandardError`; returns the attempt's error or `nil`. |
 | Hook       | `job_finalize(job, state)`       | Before successful finalization; returning `:delete` deletes the job instead of retaining it. |
 | Middleware | `insert_many(params, operation)` | Around one insertion call; `params` is an array even for `Client#insert`.                    |
 | Middleware | `work(job, operation)`           | Around the work hooks and worker for one claimed job.                                        |
@@ -609,10 +609,33 @@ end
 config = River::Config.new(plugins: [AuditPlugin.new, TimingPlugin.new])
 ```
 
-Plugins earlier in the list are the outermost wrappers. Begin callbacks run in
-configuration order, while `insert_end` callbacks run in reverse order. In
-effect, work execution is nested as: middleware before, `work_begin`, worker,
-`work_end`, middleware after.
+Plugins earlier in the list are the outermost wrappers. Begin callbacks and
+`work_end` callbacks run in configuration order, while `insert_end` callbacks
+run in reverse order. In effect, work execution is nested as: middleware before,
+`work_begin`, worker, `work_end`, middleware after.
+
+Each `work_end(job, error)` receives `nil` on success or the worker's exception
+on failure. Return `error` to preserve that outcome, a different exception to
+replace it, or `nil` to clear a failure. Each subsequent end hook receives the
+previous hook's return value. After all end hooks run, River raises the final
+error through the middleware, or completes the operation if it is `nil`.
+Returning `River.job_cancel(...)` or `River.job_snooze(...)` selects the
+corresponding cancellation or snooze behavior.
+
+An observing hook must explicitly return its incoming error:
+
+```ruby
+def work_end(job, error)
+  Metrics.record_attempt(job.kind, error: error)
+  error
+end
+```
+
+If a `work_begin` hook raises, River skips the worker and all end hooks. If an
+end hook raises instead of returning an error, it stops the end-hook chain and
+propagates that exception through the middleware. Fatal exceptions outside
+`StandardError` also bypass end hooks. Use middleware with `ensure` for cleanup
+that must run when hooks or work raise.
 
 ### [Subscriptions](https://riverqueue.com/docs/subscriptions)
 

ruby/lib/client_runtime.rb

@@ -425,15 +425,22 @@ def wake
       return perform_work(worker, job) if @config.plugins.empty?
 
       operation = -> do
+        invoke_plugins(:work_begin, job)
+
         error = nil
-        begin
-          invoke_plugins(:work_begin, job)
+        result = begin
           perform_work(worker, job)
-        rescue => error
-          raise
-        ensure
-          invoke_plugins(:work_end, job, error)
+        rescue => work_error
+          error = work_error
+          nil
+        end
+
+        @config.plugins.each do |plugin|
+          error = plugin.work_end(job, error) if plugin.respond_to?(:work_end)
         end
+        raise error if error
+
+        result
       end
       @config.plugins.reverse_each do |plugin|
         next unless plugin.respond_to?(:work)

ruby/spec/runtime_feature_spec.rb

@@ -419,10 +419,16 @@ def next_retry(_job, _error)
     calls = []
     plugin_one = Object.new
     plugin_one.define_singleton_method(:work_begin) { |_job| calls << :one_begin }
-    plugin_one.define_singleton_method(:work_end) { |_job, _error| calls << :one_end }
+    plugin_one.define_singleton_method(:work_end) do |_job, error|
+      calls << :one_end
+      error
+    end
     plugin_two = Object.new
     plugin_two.define_singleton_method(:work_begin) { |_job| calls << :two_begin }
-    plugin_two.define_singleton_method(:work_end) { |_job, _error| calls << :two_end }
+    plugin_two.define_singleton_method(:work_end) do |_job, error|
+      calls << :two_end
+      error
+    end
     client = build_client(Class.new {
       def work(_job)
       end

ruby/spec/runtime_spec.rb

@@ -416,7 +416,10 @@ def work(_job)
     plugin.define_singleton_method(:insert_begin) { |_params| calls << :insert_begin }
     plugin.define_singleton_method(:insert_end) { |_result| calls << :insert_end }
     plugin.define_singleton_method(:work_begin) { |_job| calls << :work_begin }
-    plugin.define_singleton_method(:work_end) { |_job, error| calls << (error ? :work_error : :work_end) }
+    plugin.define_singleton_method(:work_end) do |_job, error|
+      calls << (error ? :work_error : :work_end)
+      error
+    end
     plugin.define_singleton_method(:work) do |_job, operation|
       calls << :middleware
       operation.call

ruby/spec/testing_spec.rb

@@ -35,7 +35,10 @@ def insert(value = 1, **options)
     calls = []
     failure = RuntimeError.new("step failed")
     plugin = Object.new
-    plugin.define_singleton_method(:work_end) { |_job, error| calls << error }
+    plugin.define_singleton_method(:work_end) do |_job, error|
+      calls << error
+      error
+    end
     plugins << plugin
     worker.define_singleton_method(:work) do |job|
       job.resumable_step(:download) { raise failure }
@@ -313,7 +316,10 @@ def insert(value = 1, **options)
     callbacks = []
     plugins << Object.new.tap do |plugin|
       plugin.define_singleton_method(:work_begin) { |_job| callbacks << :begin }
-      plugin.define_singleton_method(:work_end) { |_job, error| callbacks << error }
+      plugin.define_singleton_method(:work_end) do |_job, error|
+        callbacks << error
+        error
+      end
       plugin.define_singleton_method(:work) do |_job, operation|
         callbacks << :before
         operation.call

ruby/spec/work_hooks_spec.rb

@@ -0,0 +1,254 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+require "riverqueue-sequel"
+require "riverqueue/testing"
+require_relative "support/river_sqlite_schema_fixture"
+
+RSpec.describe "River work hooks", database: :sqlite do
+  let(:calls) { [] }
+  let(:database) do
+    Sequel.sqlite.tap { |db| db.synchronize { |connection| RiverSQLiteSchemaFixture.load(connection) } }
+  end
+  let(:handled_errors) { [] }
+  let(:plugins) { [] }
+  let(:worker) { Object.new }
+  let(:client) do
+    River::Client.new(River::Driver::Sequel.new(database), config: River::Config.new(
+      error_handler: ->(error, _job) {
+        handled_errors << error
+        false
+      },
+      plugins: plugins,
+      workers: River::Workers.new.add("work_hooks", worker)
+    ))
+  end
+
+  after { database.disconnect }
+
+  def add_end_hook(&block)
+    Object.new.tap do |plugin|
+      plugin.define_singleton_method(:work_end, &block)
+      plugins << plugin
+    end
+  end
+
+  def perform
+    row = client.insert(River::JobArgsHash.new("work_hooks", {})).job
+    River::Testing.perform_job(client, row.id)
+  end
+
+  it "allows an end hook to fail a successful attempt" do
+    failure = RuntimeError.new("end hook failed")
+    worker.define_singleton_method(:work) { |_job| :success }
+    received = []
+    add_end_hook do |_job, error|
+      received << error
+      failure
+    end
+
+    result = perform
+
+    expect(received).to eq([nil])
+    expect(result).to have_attributes(error: failure, outcome: :retried)
+    expect(result.job.errors.last.error).to eq("end hook failed")
+    expect(handled_errors).to eq([failure])
+  end
+
+  it "chains replacement errors and nil through end hooks in registration order" do
+    original = RuntimeError.new("worker failed")
+    replacement = RuntimeError.new("replacement")
+    final = RuntimeError.new("final")
+    worker.define_singleton_method(:work) { |_job| raise original }
+    received = []
+    [replacement, nil, final].each do |next_error|
+      add_end_hook do |job, error|
+        received << [job.id, error]
+        next_error
+      end
+    end
+    # Plugins without this hook must not change the accumulated error.
+    plugins.insert(1, Object.new)
+    subscription = client.subscribe(River::EVENT_JOB_FAILED)
+
+    result = perform
+
+    expect(received).to eq([[result.id, original], [result.id, replacement], [result.id, nil]])
+    expect(result).to have_attributes(error: final, outcome: :retried)
+    expect(result.job.errors.map(&:error)).to eq(["final"])
+    expect(handled_errors).to eq([final])
+    expect(subscription.pop(true).job.errors.last.error).to eq("final")
+  end
+
+  it "clears a worker failure before returning successfully through middleware" do
+    original = RuntimeError.new("worker failed")
+    received = []
+    worker.define_singleton_method(:work) { |_job| raise original }
+    plugin = add_end_hook do |_job, error|
+      received << error
+      nil
+    end
+    plugin.define_singleton_method(:work) do |_job, operation|
+      received << :before
+      received << operation.call
+      received << :after
+    end
+    subscription = client.subscribe(River::EVENT_JOB_COMPLETED)
+
+    result = perform
+
+    expect(received).to eq([:before, original, nil, :after])
+    expect(result).to have_attributes(error: nil, outcome: :completed)
+    expect(result.job).to have_attributes(errors: [], state: "completed")
+    expect(handled_errors).to be_empty
+    expect(subscription.pop(true).kind).to eq(River::EVENT_JOB_COMPLETED)
+  end
+
+  it "preserves a successful worker return value through hooks and middleware" do
+    received = calls
+    worker.define_singleton_method(:work) do |_job|
+      received << :work
+      :worker_result
+    end
+    first = add_end_hook do |_job, error|
+      received << [:first_end, error]
+      error
+    end
+    first.define_singleton_method(:work_begin) { |_job| received << :begin }
+    first.define_singleton_method(:work) do |_job, operation|
+      received << :before
+      received << operation.call
+      received << :after
+    end
+    add_end_hook do |_job, error|
+      received << [:second_end, error]
+      error
+    end
+
+    result = perform
+
+    expect(result).to have_attributes(error: nil, outcome: :completed)
+    expect(calls).to eq([:before, :begin, :work, [:first_end, nil], [:second_end, nil], :worker_result, :after])
+  end
+
+  it "preserves the original error and backtrace when an end hook returns it" do
+    failure = RuntimeError.new("worker failed")
+    failure.set_backtrace(["worker.rb:12:in 'work'"])
+    worker.define_singleton_method(:work) { |_job| raise failure }
+    received = []
+    plugin = add_end_hook do |_job, error|
+      received << [:end, error]
+      error
+    end
+    plugin.define_singleton_method(:work) do |_job, operation|
+      operation.call
+    rescue => error
+      received << [:middleware, error]
+      raise
+    end
+
+    result = perform
+
+    expect(received).to eq([[:end, failure], [:middleware, failure]])
+    expect(result.error).to equal(failure)
+    expect(result.outcome).to eq(:retried)
+    expect(result.job.errors.last).to have_attributes(error: "worker failed", trace: "worker.rb:12:in 'work'")
+    expect(handled_errors).to eq([failure])
+  end
+
+  {
+    cancelled: [River.job_cancel("cancelled by end hook"), "cancelled"],
+    interrupted: [River::ClientRuntime::Interrupted.new, "available"],
+    snoozed: [River.job_snooze(60), "scheduled"]
+  }.each do |outcome, (replacement, state)|
+    it "replaces a worker failure with a #{outcome} outcome" do
+      worker.define_singleton_method(:work) { |_job| raise "worker failed" }
+      add_end_hook { |_job, _error| replacement }
+
+      result = perform
+
+      expect(result.error).to equal(replacement)
+      expect(result.outcome).to eq(outcome)
+      expect(result.job.state).to eq(state)
+      expect(result.job.errors.map(&:error)).to eq((outcome == :cancelled) ? [replacement.message] : [])
+    end
+  end
+
+  it "skips end hooks when the worker raises outside StandardError" do
+    fatal = Exception.new("fatal worker error")
+    worker.define_singleton_method(:work) { |_job| raise fatal }
+    received = []
+    plugin = add_end_hook do |_job, error|
+      received << :end
+      error
+    end
+    plugin.define_singleton_method(:work) do |_job, operation|
+      operation.call
+    ensure
+      received << :cleanup
+    end
+
+    expect { perform }.to raise_error { |error| expect(error).to equal(fatal) }
+    expect(received).to eq([:cleanup])
+  end
+
+  it "skips the worker and all end hooks when a begin hook fails" do
+    failure = RuntimeError.new("begin failed")
+    received = []
+    worker.define_singleton_method(:work) { |_job| received << :work }
+    first = add_end_hook do |_job, error|
+      received << :first_end
+      error
+    end
+    first.define_singleton_method(:work_begin) do |_job|
+      received << :first_begin
+      raise failure
+    end
+    first.define_singleton_method(:work) do |_job, operation|
+      operation.call
+    rescue => error
+      received << [:middleware, error]
+      raise
+    end
+    second = add_end_hook do |_job, error|
+      received << :second_end
+      error
+    end
+    second.define_singleton_method(:work_begin) { |_job| received << :second_begin }
+
+    result = perform
+
+    expect(result.error).to equal(failure)
+    expect(result.outcome).to eq(:retried)
+    expect(received).to eq([:first_begin, [:middleware, failure]])
+  end
+
+  [nil, RuntimeError.new("worker failed")].each do |worker_error|
+    it "stops the end-hook chain when a hook raises after #{worker_error ? "failure" : "success"}" do
+      failure = RuntimeError.new("end hook raised")
+      worker.define_singleton_method(:work) { |_job| raise worker_error if worker_error }
+      received = []
+      plugin = add_end_hook do |_job, error|
+        received << error
+        raise failure
+      end
+      plugin.define_singleton_method(:work) do |_job, operation|
+        operation.call
+      rescue => error
+        received << [:middleware, error]
+        raise
+      end
+      add_end_hook do |_job, error|
+        received << :must_not_run
+        error
+      end
+
+      result = perform
+
+      expect(result.error).to equal(failure)
+      expect(result.outcome).to eq(:retried)
+      expect(received).to eq([worker_error, [:middleware, failure]])
+      expect(result.job.errors.last.error).to eq("end hook raised")
+    end
+  end
+end