Add 'reset' option to reset limit counter at the end of each interval.

README.md

@@ -66,6 +66,16 @@ sidekiq_options throttle: { threshold: 20, period: 1.day, key: ->{ |user_id| use
 In the above example, jobs are throttled for each user when they exceed 20 in a
 day.
 
+If counters reset at the end of each throttling interval, you can set `reset` parameter to `true`
+(it is 'false' by default). This will reset counter each day at '00:00 UTC':
+
+```ruby
+sidekiq_options throttle: { threshold: 100000, period: 1.day, reset: true }
+```
+
+This will work for any period. For example, `period: 15.minutes` in the above example will reset the counter
+four times per hour at 00:00, 15:00, 30:00, and 45:00.
+
 ## Contributing
 
 1. Fork it
@@ -76,4 +86,4 @@ day.
 
 ## License
 
-MIT Licensed. See LICENSE.txt for details.
+MIT Licensed. See LICENSE.txt for details.

lib/sidekiq/throttler.rb

@@ -31,11 +31,11 @@ def call(worker, msg, queue)
       end
 
       rate_limit.exceeded do |delay|
-        worker.class.perform_in(delay, *msg['args'])
+        worker.class.perform_at(delay, *msg['args'])
       end
 
       rate_limit.execute
     end
 
   end # Throttler
-end # Sidekiq
+end # Sidekiq

lib/sidekiq/throttler/rate_limit.rb

@@ -78,6 +78,31 @@ def period
         @period ||= options['period'].to_f
       end
 
+      ##
+      # Reset throttle counter at the beginning of the next period.
+      #
+      # @return [true, false]
+      def reset?
+        @reset ||= options['reset'] || options['reset'].nil?
+      end
+
+      ##
+      # Get end time of a throttling period.
+      #
+      # @param [Time] time
+      #
+      # @return [Time]
+      #   The end time of the throttling period.
+      def end_of_period(time = Time.now)
+        period_end = if reset?
+          period * ((time + period).to_f / period).floor
+        else
+          time + period
+        end
+
+        Time.at period_end
+      end
+
       ##
       # @return [Symbol]
       #   The key name used when storing counters for jobs.
@@ -137,7 +162,7 @@ def execute
         return @within_bounds.call unless can_throttle?
 
         if exceeded?
-          @exceeded.call(period)
+          @exceeded.call(end_of_period)
         else
           increment
           @within_bounds.call
@@ -193,10 +218,10 @@ def self.executions
       #   The rate limit to prune.
       def self.prune(limiter)
         executions[limiter.key].select! do |execution|
-          (Time.now - execution) < limiter.period
+          limiter.end_of_period(execution) > Time.now
         end
       end
 
     end # RateLimit
   end # Throttler
-end # Sidekiq
+end # Sidekiq

spec/sidekiq/throttler/rate_limit_spec.rb

@@ -234,11 +234,12 @@
 
       before do
         rate_limit.should_receive(:exceeded?).and_return(true)
+        rate_limit.should_receive(:end_of_period).and_return(:time)
       end
 
       it 'calls the exceeded callback with the configured #period' do
         callback = Proc.new {}
-        callback.should_receive(:call).with(rate_limit.period)
+        callback.should_receive(:call).with(:time)
 
         rate_limit.exceeded(&callback)
         rate_limit.execute
@@ -284,7 +285,8 @@
 
     context 'when #period has passed' do
 
-      it 'removes old increments' do
+      it 'removes old increments when counters not reset at the end of period' do
+        rate_limit.options['reset'] = false
         rate_limit.options['period'] = 5
 
         Timecop.freeze
@@ -296,6 +298,59 @@
 
         rate_limit.count.should eq(5)
       end
+
+      it 'removes old increments when counteres reset at the end of period' do
+        rate_limit.options['reset'] = true
+        rate_limit.options['period'] = 1.minute
+
+        Timecop.freeze Time.new(2012, 12, 20, 11, 59, 50)
+
+        15.times do
+          Timecop.travel(1.second.from_now)
+          rate_limit.increment
+        end
+
+        rate_limit.count.should eq(6)
+      end
+    end
+  end
+
+  describe "#end_of_period" do
+    before do
+      rate_limit.options['period'] = 1.minute
+    end
+
+    it "defaults parameter to now" do
+      Timecop.freeze
+      rate_limit.end_of_period.should == rate_limit.end_of_period(Time.now)
+    end
+
+    context "when resetting of counters is disabled" do
+      before do
+        rate_limit.options['reset'] = false
+      end
+
+      it "ends at given time plus period length" do
+        time = Time.now
+        rate_limit.end_of_period(time).should == time + 1.minute
+      end
+    end
+
+    context "when resetting of counters is enabled" do
+      before do
+        rate_limit.options['reset'] = true
+      end
+
+      it "ends at the end of nearest period interval" do
+        time = Time.new(2012, 12, 31, 11, 59, 30)
+        rate_limit.end_of_period(time).should == time + 30
+      end
+
+      it "ends day intervals on UTC time" do
+        rate_limit.options['period'] = 1.day
+        time = Time.new(2012, 12, 31, 12, 00, 00, 3.hours)
+        rate_limit.end_of_period(time).should == Time.utc(2013, 01, 01)
+      end
     end
   end
-end
+end

spec/sidekiq/throttler_spec.rb

@@ -44,9 +44,10 @@
 
       it 'requeues the job with a delay' do
         Sidekiq::Throttler::RateLimit.any_instance.should_receive(:exceeded?).and_return(true)
-        worker.class.should_receive(:perform_in).with(1.minute, *message['args'])
+        Sidekiq::Throttler::RateLimit.any_instance.should_receive(:end_of_period).and_return(:new_time)
+        worker.class.should_receive(:perform_at).with(:new_time, *message['args'])
         throttler.call(worker, message, queue)
       end
     end
   end
-end
+end