sidekiq-throttler 0.2.0 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA1:
-  metadata.gz: f178571c0d6b730e0f2dda4af252fd1c3322219f
-  data.tar.gz: a1ecdbffc68b36c4cd2d879fd7fd10da7e7a1b7b
+  metadata.gz: 0258db1754f26622f1f38b1875774946a6a1c16e
+  data.tar.gz: 151f57ccb677fd3ad6cd25cd30ed605da04acb97
 SHA512:
-  metadata.gz: 3f957bb16f42a8e0024ca3bb961e227ed1f24659cc6863ba5eedb02915ee2065f61224207ffc05af771eccbd8f4bf2a2998b457418063eaee8322f6844c68280
-  data.tar.gz: 43062db0499c028ee9381bb6e6aea06f9409b4023cb0add9969e17c9747da3240badfcb30abbc5590bdc1b811e46256dd29daf38c76e240a51cae061221ef130
+  metadata.gz: 91440b76579924fc449a60c749ce5b88619c5efacfa5fabc45271b7f3f6a72b7408bc2b9b954e99d5aa33ab0b52a18addc957c9753fe4854b521c82738aec389
+  data.tar.gz: 9b57fbf2ae12af64aaf5667d658da9e2fb6af6ba3099373376eedad86c422a13a879a80137f7f9c8d2244dcb91a7db18bbb01555d43c6a57b37f2af999d5ece5

data/.travis.yml

@@ -1,18 +1,17 @@
 language: ruby
 rvm:
+  - 2.0.0
   - 1.9.3
   - jruby-19mode
   - rbx-19mode
-  - 2.0.0
 branches:
   only:
     - master
 notifications:
   email:
     recipients:
-      - gabriel@codeconcoction.com
+      - gabe@ga.be
 matrix:
   allow_failures:
     - rvm: jruby-19mode
     - rvm: rbx-19mode
-    - rvm: 2.0.0

data/CHANGELOG.md

@@ -1,3 +1,15 @@
+## 0.3.0 (October 3, 2013)
+
+* Redis is supported as a storage backend for persistence of job execution
+  counters across multiple Sidekiq processes.
+
+  *Louis Simoneau*
+
+* Only Active Support's `Time` extensions are required. Fixes compatibility with
+  Rails 4.
+
+  *Louis Simoneau*
+
 ## 0.2.0 (June 29, 2013)
 
 * Now supports using a Proc for `:threshold` argument, similar to how the

data/README.md

@@ -8,7 +8,7 @@ limit job execution on a per-worker basis.
 
 ## Compatibility
 
-Sidekiq::Throttler is tested against MRI 1.9.3.
+Sidekiq::Throttler is actively tested against MRI versions 2.0.0 and 1.9.3.
 
 ## Installation
 
@@ -37,6 +37,18 @@ Sidekiq.configure_server do |config|
 end
 ```
 
+Sidekiq::Throttler defaults to in-memory storage of job execution times. If
+you have multiple worker processes, or frequently restart your processes, this
+will be unreliable. Instead, specify the `:redis` storage option:
+
+```ruby
+Sidekiq.configure_server do |config|
+  config.server_middleware do |chain|
+    chain.add Sidekiq::Throttler, storage: :redis
+  end
+end
+```
+
 ## Usage
 
 In a worker, specify a threshold (maximum jobs) and period for throttling:
@@ -66,6 +78,15 @@ 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.
 
+Thresholds can be configured based on the arguments passed to your worker's `perform` method,
+similar to how the `key` option works:
+
+```ruby
+sidekiq_options throttle: { threshold: -> {|user_id, rate_limit| rate_limit }, period: 1.hour, key: ->{ |user_id, rate_limit| user_id } }
+```
+
+In the above example, jobs are throttled for each user when they exceed the rate limit provided in the message. This is useful in cases where each user may have a different rate limit (ex: interacting with external APIs)
+
 ## Contributing
 
 1. Fork it
@@ -76,4 +97,4 @@ day.
 
 ## License
 
-MIT Licensed. See LICENSE.txt for details.
+MIT Licensed. See LICENSE.txt for details.

data/lib/sidekiq/throttler/rate_limit.rb

@@ -32,10 +32,18 @@ module Sidekiq
       #
       # @param [String] queue
       #   The queue to rate limit.
-      def initialize(worker, payload, queue)
+      #
+      # @option [Symbol] :storage
+      #   Either :memory or :redis, the storage backend to use
+      def initialize(worker, payload, queue, options = {})
         @worker = worker
         @payload = payload
         @queue = queue
+
+        unless @storage_class = lookup_storage(options.fetch(:storage, :memory))
+          raise ArgumentError,
+            "Unrecognized storage backend: #{options[:storage].inspect}"
+        end
       end
 
       ##
@@ -68,7 +76,7 @@ module Sidekiq
       # @return [Integer]
       #   The number of jobs that are allowed within the `period`.
       def threshold
-        @threshold ||= options['threshold'].to_i
+        @threshold ||= (options['threshold'].respond_to?(:call) ? options['threshold'].call(*payload) : options['threshold']).to_i
       end
 
       ##
@@ -79,7 +87,7 @@ module Sidekiq
       end
 
       ##
-      # @return [Symbol]
+      # @return [String]
       #   The key name used when storing counters for jobs.
       def key
         @key ||= if options['key']
@@ -146,8 +154,14 @@ module Sidekiq
 
       ##
       # Reset the tracking of job executions.
-      def self.reset!
-        @executions = Hash.new { |hash, key| hash[key] = [] }
+      def reset!
+        executions.reset
+      end
+
+      ##
+      # Get the storage backend.
+      def executions
+        @storage_class.instance
       end
 
       private
@@ -162,7 +176,7 @@ module Sidekiq
       def self.count(limiter)
         Thread.exclusive do
           prune(limiter)
-          executions[limiter.key].length
+          limiter.executions.count(limiter.key)
         end
       end
 
@@ -175,28 +189,31 @@ module Sidekiq
       #   The current number of jobs executed.
       def self.increment(limiter)
         Thread.exclusive do
-          executions[limiter.key] << Time.now
+          limiter.executions.append(limiter.key, Time.now)
         end
         count(limiter)
       end
 
-      ##
-      # A hash storing job executions as timestamps for each throttled worker.
-      def self.executions
-        @executions || reset!
-      end
-
       ##
       # Remove old entries for the provided `RateLimit`.
       #
       # @param [RateLimit] limiter
       #   The rate limit to prune.
       def self.prune(limiter)
-        executions[limiter.key].select! do |execution|
-          (Time.now - execution) < limiter.period
-        end
+        limiter.executions.prune(limiter.key, Time.now - limiter.period)
       end
 
+      ##
+      # Lookup storage class for a given options key
+      #
+      # @param [Symbol] key
+      #   The options key, :memory or :redis
+      #
+      # @return [Class]
+      #   The storage backend class, or nil if the key is not found
+      def lookup_storage(key)
+        { memory: Storage::Memory, redis: Storage::Redis }[key]
+      end
     end # RateLimit
   end # Throttler
-end # Sidekiq
+end # Sidekiq

data/lib/sidekiq/throttler/storage/memory.rb

@@ -0,0 +1,55 @@
+module Sidekiq
+  class Throttler
+    module Storage
+      ##
+      # Stores job executions in a Hash of Arrays.
+      class Memory
+        include Singleton
+
+        def initialize
+          @hash = Hash.new { |hash, key| hash[key] = [] }
+        end
+
+        ##
+        # Number of executions for +key+.
+        #
+        # @param [String]
+        #   Key to fetch count for
+        #
+        # @return [Fixnum]
+        #   Execution count
+        def count(key)
+          @hash[key].length
+        end
+
+        ##
+        # Remove entries older than +cutoff+.
+        #
+        # @param [String] key
+        #   The key to prune
+        #
+        # @param [Time] cutoff
+        #   Oldest allowable time
+        def prune(key, cutoff)
+          @hash[key].reject! { |time| time <= cutoff }
+        end
+
+        ##
+        # Add a new entry to the hash.
+        #
+        # @param [String] key
+        #   The key to append to
+        #
+        # @param [Time]
+        #   The time to insert
+        def append(key, time)
+          @hash[key] << time
+        end
+
+        def reset
+          @hash.clear
+        end
+      end
+    end # Storage
+  end # Throttler
+end # Sidekiq

data/lib/sidekiq/throttler/storage/redis.rb

@@ -0,0 +1,89 @@
+module Sidekiq
+  class Throttler
+    module Storage
+      ##
+      # Stores job executions in Redis lists.
+      #
+      # Timestamps in each list are ordered with the oldest on the right.
+      # Values are inserted on the left (LPUSH) & pruned from the right (RPOP)
+      class Redis
+        include Singleton
+
+        ##
+        # Number of executions for +key+.
+        #
+        # @param [String]
+        #   Key to fetch count for
+        #
+        # @return [Fixnum]
+        #   Execution count
+        def count(key)
+          Sidekiq.redis do |conn|
+            conn.llen(namespace_key(key))
+          end
+        end
+
+        ##
+        # Remove entries older than +cutoff+.
+        #
+        # @param [String] key
+        #   The key to prune
+        #
+        # @param [Time] cutoff
+        #   Oldest allowable time
+        def prune(key, cutoff)
+          # Repeatedly pop from the right of the list until we encounter
+          # a value greater than the cutoff.
+          #
+          # We compare popped values to account for race conditions,
+          # pushing them back if they don't match.
+          Sidekiq.redis do |conn|
+            prune_one = ->(timestamp) {
+              if timestamp && timestamp.to_i <= cutoff.to_i
+                last = conn.rpop(namespace_key(key))
+                if last == timestamp
+                  true
+                else
+                  conn.rpush(namespace_key(key), last)
+                  nil
+                end
+              end
+            }
+
+            loop while prune_one.call(conn.lindex(namespace_key(key), -1))
+          end
+        end
+
+        ##
+        # Add a new entry to the hash.
+        #
+        # @param [String] key
+        #   The key to append to
+        #
+        # @param [Time]
+        #   The time to insert
+        def append(key, time)
+          Sidekiq.redis do |conn|
+            conn.lpush(namespace_key(key), time.to_i)
+          end
+        end
+
+        ##
+        # Clear all data from storage.
+        def reset
+          Sidekiq.redis do |conn|
+            conn.keys(namespace_key("*")).each do |key|
+              conn.del(key)
+            end
+          end
+        end
+
+        private
+
+        def namespace_key(key)
+          "throttled:#{key}"
+        end
+      end
+    end # Storage
+  end # Throttler
+end # Sidekiq

data/lib/sidekiq/throttler/version.rb

@@ -1,5 +1,5 @@
 module Sidekiq
   class Throttler
-    VERSION = '0.2.0'
+    VERSION = '0.3.0'
   end
 end

data/lib/sidekiq/throttler.rb

@@ -1,15 +1,20 @@
 require 'sidekiq'
-require 'active_support'
-require 'active_support/core_ext'
+require 'active_support/core_ext/numeric/time'
 
 require 'sidekiq/throttler/version'
 require 'sidekiq/throttler/rate_limit'
 
+require 'sidekiq/throttler/storage/memory'
+require 'sidekiq/throttler/storage/redis'
+
 module Sidekiq
   ##
   # Sidekiq server middleware. Throttles jobs when they exceed limits specified
   # on the worker. Jobs that exceed the limit are requeued with a delay.
   class Throttler
+    def initialize(options = {})
+      @options = options.dup
+    end
 
     ##
     # Passes the worker, arguments, and queue to {RateLimit} and either yields
@@ -24,7 +29,7 @@ module Sidekiq
     # @param [String] queue
     #   The current queue.
     def call(worker, msg, queue)
-      rate_limit = RateLimit.new(worker, msg['args'], queue)
+      rate_limit = RateLimit.new(worker, msg['args'], queue, @options)
 
       rate_limit.within_bounds do
         yield
@@ -38,4 +43,4 @@ module Sidekiq
     end
 
   end # Throttler
-end # Sidekiq
+end # Sidekiq

data/spec/app/workers/proc_threshold_worker.rb

@@ -0,0 +1,10 @@
+class ProcThresholdWorker
+  include Sidekiq::Worker
+
+  sidekiq_options throttle: { threshold: Proc.new { |user_id, limit| limit }, period: 1.minute }
+
+  def perform(user_id, limit)
+    puts user_id
+    puts limit
+  end
+end

data/spec/sidekiq/throttler/rate_limit_spec.rb

@@ -1,5 +1,29 @@
 require 'spec_helper'
 
+shared_examples "incrementing" do
+  it 'increments #count by one' do
+    Timecop.freeze do
+      expect { rate_limit.increment }.to change{ rate_limit.count }.by(1)
+    end
+  end
+
+  context 'when #period has passed' do
+
+    it 'removes old increments' do
+      rate_limit.options['period'] = 5
+
+      Timecop.freeze
+
+      20.times do
+        Timecop.travel(1.second.from_now)
+        rate_limit.increment
+      end
+
+      rate_limit.count.should eq(5)
+    end
+  end
+end
+
 describe Sidekiq::Throttler::RateLimit do
 
   let(:worker_class) do
@@ -22,6 +46,10 @@ describe Sidekiq::Throttler::RateLimit do
     described_class.new(worker, payload, 'meow')
   end
 
+  before(:each) do
+    rate_limit.reset!
+  end
+
   describe '.new' do
 
     it 'initializes with a provided worker' do
@@ -35,6 +63,14 @@ describe Sidekiq::Throttler::RateLimit do
     it 'initializes with a provided queue' do
       rate_limit.queue.should eq('meow')
     end
+
+    context "with an invalid storage backend" do
+      it "raises an ArgumentError" do
+        expect {
+          described_class.new(worker, payload, 'meow', storage: :blarg)
+        }.to raise_error(ArgumentError)
+      end
+    end
   end
 
   describe '#options' do
@@ -69,6 +105,20 @@ describe Sidekiq::Throttler::RateLimit do
 
   describe '#threshold' do
 
+    context 'when threshold is a Proc' do
+      let(:worker_class) do
+        ProcThresholdWorker
+      end
+
+      let(:payload) do
+        [1, 500]
+      end
+
+      it 'returns the result of the called Proc' do
+        rate_limit.threshold.should eq(500)
+      end
+    end
+
     it 'retrieves the threshold from #options' do
       rate_limit.options['threshold'] = 26
       rate_limit.threshold.should eq(26)
@@ -275,27 +325,14 @@ describe Sidekiq::Throttler::RateLimit do
   end
 
   describe '#increment' do
+    include_examples "incrementing"
+  end
 
-    it 'increments #count by one' do
-      Timecop.freeze do
-        expect { rate_limit.increment }.to change{ rate_limit.count }.by(1)
-      end
+  context "with a :redis storage backend" do
+    subject(:rate_limit) do
+      described_class.new(worker, payload, 'meow', storage: :redis)
     end
 
-    context 'when #period has passed' do
-
-      it 'removes old increments' do
-        rate_limit.options['period'] = 5
-
-        Timecop.freeze
-
-        20.times do
-          Timecop.travel(1.second.from_now)
-          rate_limit.increment
-        end
-
-        rate_limit.count.should eq(5)
-      end
-    end
+    include_examples "incrementing"
   end
-end
+end

data/spec/sidekiq/throttler/storage/redis_spec.rb

@@ -0,0 +1,37 @@
+require "spec_helper"
+
+describe Sidekiq::Throttler::Storage::Redis do
+  let(:storage) { described_class.instance }
+
+  before(:each) do
+    @sidekiq = double()
+    Sidekiq.stub(:redis).and_yield(@sidekiq)
+  end
+
+  describe "#prune" do
+    it "pops the last item off the list if it's lower than the cutoff" do
+      @sidekiq.stub(:lindex).and_return(100, nil)
+      @sidekiq.should_receive(:rpop).with("throttled:fake").and_return(100)
+      storage.prune("fake", 200)
+    end
+
+    it "leaves the last item on the list if it's higher than the cutoff" do
+      @sidekiq.stub(:lindex).and_return(200, nil)
+      @sidekiq.should_not_receive(:rpop)
+      storage.prune("fake", 100)
+    end
+
+    context "when another job has concurrently removed a timestamp" do
+
+      before(:each) do
+        @sidekiq.stub(:lindex) { 100 }
+        @sidekiq.stub(:rpop)   { 200 }
+      end
+
+      it "pushes the value back onto the list" do
+        @sidekiq.should_receive(:rpush).with("throttled:fake", 200)
+        storage.prune("fake", 1000)
+      end
+    end
+  end
+end

data/spec/sidekiq/throttler_spec.rb

@@ -3,13 +3,17 @@ require 'spec_helper'
 describe Sidekiq::Throttler do
 
   subject(:throttler) do
-    described_class.new
+    described_class.new(options)
   end
 
   let(:worker) do
     LolzWorker.new
   end
 
+  let(:options) do
+    { storage: :memory }
+  end
+
   let(:message) do
     {
       args: 'Clint Eastwood'
@@ -23,10 +27,9 @@ describe Sidekiq::Throttler do
   describe '#call' do
 
     it 'instantiates a rate limit with the worker, args, and queue' do
-      rate_limit = Sidekiq::Throttler::RateLimit.new(worker, message['args'], queue)
       Sidekiq::Throttler::RateLimit.should_receive(:new).with(
-        worker, message['args'], queue
-      ).and_return(rate_limit)
+        worker, message['args'], queue, options
+      ).and_call_original
 
       throttler.call(worker, message, queue) {}
     end
@@ -49,4 +52,4 @@ describe Sidekiq::Throttler do
       end
     end
   end
-end
+end

data/spec/spec_helper.rb

@@ -28,12 +28,8 @@ RSpec.configure do |config|
   # the seed, which is printed after each run.
   #     --seed 1234
   config.order = 'random'
-
-  config.before(:each) do
-    Sidekiq::Throttler::RateLimit.reset!
-  end
 end
 
 # Requires supporting files with custom matchers and macros, etc,
 # in ./support/ and its subdirectories.
-Dir["#{File.dirname(__FILE__)}/support/**/*.rb"].each { |f| require f }
+Dir["#{File.dirname(__FILE__)}/support/**/*.rb"].each { |f| require f }

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: sidekiq-throttler
 version: !ruby/object:Gem::Version
-  version: 0.2.0
+  version: 0.3.0
 platform: ruby
 authors:
 - Gabe Evans
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2013-06-29 00:00:00.000000000 Z
+date: 2013-10-03 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: activesupport
@@ -274,13 +274,17 @@ files:
 - lib/sidekiq-throttler.rb
 - lib/sidekiq/throttler.rb
 - lib/sidekiq/throttler/rate_limit.rb
+- lib/sidekiq/throttler/storage/memory.rb
+- lib/sidekiq/throttler/storage/redis.rb
 - lib/sidekiq/throttler/version.rb
 - sidekiq-throttler.gemspec
 - spec/app/workers/custom_key_worker.rb
 - spec/app/workers/lolz_worker.rb
+- spec/app/workers/proc_threshold_worker.rb
 - spec/app/workers/proc_worker.rb
 - spec/app/workers/regular_worker.rb
 - spec/sidekiq/throttler/rate_limit_spec.rb
+- spec/sidekiq/throttler/storage/redis_spec.rb
 - spec/sidekiq/throttler_spec.rb
 - spec/spec.opts
 - spec/spec_helper.rb
@@ -304,7 +308,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
       version: '0'
 requirements: []
 rubyforge_project: 
-rubygems_version: 2.0.3
+rubygems_version: 2.1.5
 signing_key: 
 specification_version: 4
 summary: Sidekiq::Throttler is a middleware for Sidekiq that adds the ability to rate
@@ -312,9 +316,11 @@ summary: Sidekiq::Throttler is a middleware for Sidekiq that adds the ability to
 test_files:
 - spec/app/workers/custom_key_worker.rb
 - spec/app/workers/lolz_worker.rb
+- spec/app/workers/proc_threshold_worker.rb
 - spec/app/workers/proc_worker.rb
 - spec/app/workers/regular_worker.rb
 - spec/sidekiq/throttler/rate_limit_spec.rb
+- spec/sidekiq/throttler/storage/redis_spec.rb
 - spec/sidekiq/throttler_spec.rb
 - spec/spec.opts
 - spec/spec_helper.rb