From 2571086578b60bd7976297a1205a54eae0358aac Mon Sep 17 00:00:00 2001
From: Kelley Reynolds <kelley@insidesystems.net>
Date: Fri, 23 Oct 2015 11:54:55 -0400
Subject: [PATCH] Make exception available in #sidekiq_retry_in

---
 lib/sidekiq/middleware/server/retry_jobs.rb | 10 +++----
 test/test_retry.rb                          | 31 +++++++++++++++++----
 2 files changed, 31 insertions(+), 10 deletions(-)

diff --git a/lib/sidekiq/middleware/server/retry_jobs.rb b/lib/sidekiq/middleware/server/retry_jobs.rb
index 990f7c3ee..9b5c807cf 100644
--- a/lib/sidekiq/middleware/server/retry_jobs.rb
+++ b/lib/sidekiq/middleware/server/retry_jobs.rb
@@ -121,7 +121,7 @@ def attempt_retry(worker, msg, queue, exception)
           end
 
           if count < max_retry_attempts
-            delay = delay_for(worker, count)
+            delay = delay_for(worker, count, exception)
             logger.debug { "Failure! Retry #{count} in #{delay} seconds" }
             retry_at = Time.now.to_f + delay
             payload = Sidekiq.dump_json(msg)
@@ -170,8 +170,8 @@ def retry_attempts_from(msg_retry, default)
           end
         end
 
-        def delay_for(worker, count)
-          worker.sidekiq_retry_in_block? && retry_in(worker, count) || seconds_to_delay(count)
+        def delay_for(worker, count, exception)
+          worker.sidekiq_retry_in_block? && retry_in(worker, count, exception) || seconds_to_delay(count)
         end
 
         # delayed_job uses the same basic formula
@@ -179,9 +179,9 @@ def seconds_to_delay(count)
           (count ** 4) + 15 + (rand(30)*(count+1))
         end
 
-        def retry_in(worker, count)
+        def retry_in(worker, count, exception)
           begin
-            worker.sidekiq_retry_in_block.call(count).to_i
+            worker.sidekiq_retry_in_block.call(count, exception).to_i
           rescue Exception => e
             handle_exception(e, { context: "Failure scheduling retry using the defined `sidekiq_retry_in` in #{worker.class.name}, falling back to default" })
             nil
diff --git a/test/test_retry.rb b/test/test_retry.rb
index bffafc360..c5f34a23a 100644
--- a/test/test_retry.rb
+++ b/test/test_retry.rb
@@ -229,7 +229,7 @@ def job(options={})
         File.unlink @tmp_log_path if File.exist?(@tmp_log_path)
       end
 
-      class CustomWorker
+      class CustomWorkerWithoutException
         include Sidekiq::Worker
 
         sidekiq_retry_in do |count|
@@ -237,6 +237,19 @@ class CustomWorker
         end
       end
 
+      class CustomWorkerWithException
+        include Sidekiq::Worker
+
+        sidekiq_retry_in do |count, exception|
+          case exception
+          when ArgumentError
+            count * 4
+          else
+            count * 2
+          end
+        end
+      end
+
       class ErrorWorker
         include Sidekiq::Worker
 
@@ -246,15 +259,23 @@ class ErrorWorker
       end
 
       it "retries with a default delay" do
-        refute_equal 4, handler.__send__(:delay_for, worker, 2)
+        refute_equal 4, handler.__send__(:delay_for, worker, 2, StandardError.new)
+      end
+
+      it "retries with a custom delay and exception 1" do
+        assert_equal 8, handler.__send__(:delay_for, CustomWorkerWithException, 2, ArgumentError.new)
+      end
+
+      it "retries with a custom delay and exception 2" do
+        assert_equal 4, handler.__send__(:delay_for, CustomWorkerWithException, 2, StandardError.new)
       end
 
-      it "retries with a custom delay" do
-        assert_equal 4, handler.__send__(:delay_for, CustomWorker, 2)
+      it "retries with a custom delay without exception" do
+        assert_equal 4, handler.__send__(:delay_for, CustomWorkerWithoutException, 2, StandardError.new)
       end
 
       it "falls back to the default retry on exception" do
-        refute_equal 4, handler.__send__(:delay_for, ErrorWorker, 2)
+        refute_equal 4, handler.__send__(:delay_for, ErrorWorker, 2, StandardError.new)
         assert_match(/Failure scheduling retry using the defined `sidekiq_retry_in`/,
                      File.read(@tmp_log_path), 'Log entry missing for sidekiq_retry_in')
       end