Refactor sidekiq_retry_in impl

Changes.md

@@ -3,6 +3,14 @@ HEAD
 - Adding button to move scheduled job to main queue [guiceolin, #1020]
 - fix i18n support resetting saved locale when job is retried [#1011]
 - log rotation via USR2 now closes the old logger [#1008]
+- Add ability to customize retry schedule, like so [jmazzi, #1027]
+
+```ruby
+class MyWorker
+  include Sidekiq::Worker
+  sidekiq_retry_in { |count| count * 2 }
+end
+```
 
 2.12.4
 -----------

lib/sidekiq/middleware/server/retry_jobs.rb

@@ -78,7 +78,7 @@ module Sidekiq
           end
 
           if count < max_retry_attempts
-            delay = seconds_to_delay(worker, count)
+            delay = delay_for(worker, count)
             logger.debug { "Failure! Retry #{count} in #{delay} seconds" }
             retry_at = Time.now.to_f + delay
             payload = Sidekiq.dump_json(msg)
@@ -109,21 +109,24 @@ module Sidekiq
           end
         end
 
-        # delayed_job uses the same basic formula
-        def seconds_to_delay(worker, count)
-          default_retry_in = (count ** 4) + 15 + (rand(30)*(count+1))
+        def delay_for(worker, count)
+          worker.sidekiq_retry_in_block? && retry_in(worker, count) || seconds_to_delay(count)
+        end
 
-          if worker.sidekiq_retry_with?
-            begin
-              worker.sidekiq_retry_with.call(count)
-            rescue Exception => e
-              logger.error { "Failure scheduling retry using the defined `sidekiq_retry_in`, falling back to the default! #{e.message}"}
-              default_retry_in
-            end
-          else
-            default_retry_in
+        # delayed_job uses the same basic formula
+        def seconds_to_delay(count)
+          (count ** 4) + 15 + (rand(30)*(count+1))
+        end
+
+        def retry_in(worker, count)
+          begin
+            worker.sidekiq_retry_in_block.call(count)
+          rescue Exception => e
+            logger.error { "Failure scheduling retry using the defined `sidekiq_retry_in` in #{worker.class.name}, falling back to default: #{e.message}"}
+            nil
           end
         end
+
       end
     end
   end

lib/sidekiq/worker.rb

@@ -26,7 +26,7 @@ module Sidekiq
     def self.included(base)
       base.extend(ClassMethods)
       base.class_attribute :sidekiq_options_hash
-      base.class_attribute :sidekiq_retry_with
+      base.class_attribute :sidekiq_retry_in_block
     end
 
     def logger
@@ -60,7 +60,7 @@ module Sidekiq
       end
 
       def sidekiq_retry_in(&block)
-        self.sidekiq_retry_with = block
+        self.sidekiq_retry_in_block = block
       end
 
       DEFAULT_OPTIONS = { 'retry' => true, 'queue' => 'default' }

test/test_retry.rb

@@ -240,15 +240,15 @@ class TestRetry < Minitest::Test
       let(:handler) { Sidekiq::Middleware::Server::RetryJobs.new }
 
       it "retries with a default delay" do
-        refute_equal 4, handler.seconds_to_delay(worker, 2)
+        refute_equal 4, handler.delay_for(worker, 2)
       end
 
       it "retries with a custom delay" do
-        assert_equal 4, handler.seconds_to_delay(custom_worker, 2)
+        assert_equal 4, handler.delay_for(custom_worker, 2)
       end
 
       it "falls back to the default retry on exception" do
-        refute_equal 4, handler.seconds_to_delay(error_worker, 2)
+        refute_equal 4, handler.delay_for(error_worker, 2)
         assert_match(/Failure scheduling retry using the defined `sidekiq_retry_in`/,
                      File.read(@tmp_log_path), 'Log entry missing for sidekiq_retry_in')
       end