Allow setting custom behavior when rate limit exceeded

.gitignore

@@ -17,3 +17,4 @@ spec/reports
 test/tmp
 test/version_tmp
 tmp
+.DS_Store

lib/sidekiq/throttler.rb

@@ -36,8 +36,28 @@ def call(worker, msg, queue)
         yield
       end
 
-      rate_limit.exceeded do |delay|
-        worker.class.perform_in(delay, *msg['args'])
+      # We now allow for explicitly setting the exceeded behavior.
+      # By default (or with a :retry value), the exceeded behavior follows the
+      # previous behavior and retries the job.
+      #
+      # Otherwise, it attempts to use the behavior specified in the class options
+      # The specified behavior should be a a proc that takes up to 4 arguments:
+      # 1: the period of the rate limiter
+      # 2: the worker for the job
+      # 3: the message payload
+      # 4: the queue the job was pulled from
+      # NB: The exceeded behavior passed *MUST* be a proc if you are using less
+      # then 4 arguments. Because the rate limiter itself always passes four
+      # arguments it doesn't work with a lambda.
+
+      worker_options = (worker.class.get_sidekiq_options['throttle'] || {}).stringify_keys
+
+      if worker_options['exceeded'].nil? || worker_options['exceeded'] == :retry
+        rate_limit.exceeded do |delay|
+          worker.class.perform_in(delay, *msg['args'])
+        end
+      else
+        rate_limit.exceeded &worker_options['exceeded']
       end
 
       rate_limit.execute

lib/sidekiq/throttler/rate_limit.rb

@@ -149,7 +149,8 @@ def execute
         return @within_bounds.call unless can_throttle?
 
         if exceeded?
-          @exceeded.call(period)
+          # passing a bunch of variables for use in the exceeded behavior
+          @exceeded.call(period, worker, payload, queue)
         else
           increment
           @within_bounds.call

spec/app/workers/lolz_worker2.rb

@@ -0,0 +1,20 @@
+class LolzWorker2
+  include Sidekiq::Worker
+
+  # immediately requeue
+  sidekiq_options throttle: { 
+    threshold: 10,
+    period: 1.minute,
+    # specify an exceeded behavior which immediately requeues the job instead
+    # of retrying it.
+    exceeded: Proc.new {
+      |delay, worker, payload|
+      worker.class.perform_async(payload)
+      }
+    }
+
+  def perform(name)
+    puts "LolzWorker2!!!"
+    puts "OHAI #{name.upcase}!"
+  end
+end

spec/sidekiq/throttler/rate_limit_spec.rb

@@ -302,7 +302,9 @@
 
       it 'calls the exceeded callback with the configured #period' do
         callback = Proc.new {}
-        expect(callback).to receive(:call).with(rate_limit.period)
+        #expect(callback).to receive(:call).with(rate_limit.period)
+        expect(callback).to receive(:call).with(rate_limit.period, worker, payload, queue)
+
 
         rate_limit.exceeded(&callback)
         rate_limit.execute

spec/sidekiq/throttler_spec.rb

@@ -10,6 +10,10 @@
     LolzWorker.new
   end
 
+  let(:worker2) do
+    LolzWorker2.new
+  end
+
   let(:options) do
     { storage: :memory }
   end
@@ -50,6 +54,12 @@
         expect(worker.class).to receive(:perform_in).with(1.minute, *message['args'])
         throttler.call(worker, message, queue)
       end
+
+      it 'properly performs the behavior specificed in exceeded option' do
+        Sidekiq::Throttler::RateLimit.any_instance.should_receive(:exceeded?).and_return(true)
+        worker2.class.should_receive(:perform_async).with(*message['args'])
+        throttler.call(worker2, message, queue)
+      end
     end
   end
 end