From 68b471c4dda22f74057f43acc92542e23aa738b6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rafael=20Mendon=C3=A7a=20Fran=C3=A7a?= <rafael@franca.dev>
Date: Mon, 28 Dec 2020 05:30:30 +0000
Subject: [PATCH] Make sure job instrumentation keep return value

The implementaiton of `instrument` in `ActiveJob::Instrumentation` was
not keeping the API of `ActiveSupport::Notification.instrument` and
returning the value of the block.

Fixes #40931.
---
 activejob/lib/active_job/instrumentation.rb | 5 ++++-
 activejob/test/cases/rescue_test.rb         | 2 +-
 activejob/test/jobs/rescue_job.rb           | 4 +++-
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/activejob/lib/active_job/instrumentation.rb b/activejob/lib/active_job/instrumentation.rb
index ff1833dd67..c5f4e017e6 100644
--- a/activejob/lib/active_job/instrumentation.rb
+++ b/activejob/lib/active_job/instrumentation.rb
@@ -18,11 +18,14 @@ module ActiveJob
     private
       def instrument(operation, payload = {}, &block)
         enhanced_block = ->(event_payload) do
-          block.call if block
+          value = block.call if block
+
           if defined?(@_halted_callback_hook_called) && @_halted_callback_hook_called
             event_payload[:aborted] = true
             @_halted_callback_hook_called = nil
           end
+
+          value
         end
 
         ActiveSupport::Notifications.instrument \
diff --git a/activejob/test/cases/rescue_test.rb b/activejob/test/cases/rescue_test.rb
index da9c87dbf0..f4a6a3a065 100644
--- a/activejob/test/cases/rescue_test.rb
+++ b/activejob/test/cases/rescue_test.rb
@@ -12,7 +12,7 @@ class RescueTest < ActiveSupport::TestCase
   test "rescue perform exception with retry" do
     job = RescueJob.new("david")
     job.perform_now
-    assert_equal [ "rescued from ArgumentError", "performed beautifully" ], JobBuffer.values
+    assert_equal [ "rescued from ArgumentError", "performed beautifully", "Retried job DIFFERENT!" ], JobBuffer.values
   end
 
   test "let through unhandled perform exception" do
diff --git a/activejob/test/jobs/rescue_job.rb b/activejob/test/jobs/rescue_job.rb
index d142cec1ea..049a8d9adf 100644
--- a/activejob/test/jobs/rescue_job.rb
+++ b/activejob/test/jobs/rescue_job.rb
@@ -8,7 +8,9 @@ class RescueJob < ActiveJob::Base
   rescue_from(ArgumentError) do
     JobBuffer.add("rescued from ArgumentError")
     arguments[0] = "DIFFERENT!"
-    retry_job
+    job = retry_job
+    JobBuffer.add("Retried job #{job.arguments[0]}")
+    job
   end
 
   rescue_from(ActiveJob::DeserializationError) do |e|