restrainer 1.0.0 → 1.0.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
-SHA1:
-  metadata.gz: c864dec39322b1a42a062119de6bac2a5a9134db
-  data.tar.gz: 2cbc9b6c3b78a4b6bb4fd706bdc645543366a6a6
+SHA256:
+  metadata.gz: ff5a0d8ca6659d4feeb9d6fd8e453e1b1733917c20dbfadb4686173998085903
+  data.tar.gz: 5f3bb001fd3fde135999d50e2383d6f4e826c54a35c4ab02a11d45c7cc29bce3
 SHA512:
-  metadata.gz: 275f22b4379067198e130603461ccd92b0da8e85ef05dd9fe142fa02c6e202b358b9512cbd8146b7c5ded56145380e2ab9b168f729ed8c00223300321217ab18
-  data.tar.gz: 9c30e5358e500d7172a4dd3924e5c64018fb4daed80d8a781786846ae9b7f0069e167af5156fc38b389c991e8a2c7ce08d1262155ed89c43dd6ae39d2fefcd13
+  metadata.gz: e9c097fbdb451fb6de6574427f3b672a09db445cb0df2f0784bd9133adbb97c614a73177c6017f0b6fe57618da24b69734538d127c84b54ec8495514d1c86660
+  data.tar.gz: 0404b8096e47387750cbef9b2ba89d5de6e795f19633c6d613ae202d211b510758db3d1d6d342a2f8c1da4a1d2e1d15727ec6034c38e2de4d6cac23341ffc4d2

data/CHANGE_LOG.md

@@ -0,0 +1,3 @@
+# 1.0.1
+
+* Use Lua script to avoid race conditions and ensure no extra processes slip through.

data/README.md

@@ -13,9 +13,11 @@ restrainer.throttle do
 end
 ```
 
+If the throttle is already full, the block will not be run and a `Restrainer::ThrottledError` will be raised.
+
 You can also override the limit in the `throttle` method. Setting a limit of zero will disable processing entirely. Setting a limit less than zero will remove the limit. Note that the limit set in the throttle is not shared with other processes, but the count of the number of processes is shared. Thus it is possible to have the throttle allow one process but reject another if the limits are different.
 
-Instances of Restrainer do not use any internal state to keep track of number of running processes. All of that information is maintained in redis. Therefore you don't need to worry about maintaining references to Restrainer instances and you can create them as needed as long as they are named consistently. You can create multiple Restrainers for different uses in your application by simply giving them different names.
+Instances of Restrainer do not use any internal state to keep track of the number of running processes. All of that information is maintained in redis. Therefore you don't need to worry about maintaining references to Restrainer instances and you can create them as needed as long as they are named consistently. You can create multiple Restrainers for different uses in your application by simply giving them different names.
 
 ### Configuration
 

data/Rakefile

@@ -1,18 +1,6 @@
 require "bundler/gem_tasks"
+require "rspec/core/rake_task"
 
-desc 'Default: run unit tests.'
-task :default => :test
+RSpec::Core::RakeTask.new(:spec)
 
-desc 'RVM likes to call it tests'
-task :tests => :test
-
-begin
-  require 'rspec'
-  require 'rspec/core/rake_task'
-  desc 'Run the unit tests'
-  RSpec::Core::RakeTask.new(:test)
-rescue LoadError
-  task :test do
-    STDERR.puts "You must have rspec 2.0 installed to run the tests"
-  end
-end
+task :default => :spec

data/VERSION

@@ -1 +1 @@
-1.0.0
+1.0.1

data/lib/restrainer.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require 'redis'
 require 'securerandom'
 
@@ -12,18 +14,51 @@ require 'securerandom'
 # If more than the specified number of processes as identified by the name argument is currently
 # running, then the throttle block will raise an error.
 class Restrainer
-  
+
   attr_reader :name, :limit
-  
+
+  ADD_PROCESS_SCRIPT = <<-LUA
+    -- Parse arguments
+    local sorted_set = ARGV[1]
+    local process_id = ARGV[2]
+    local limit = tonumber(ARGV[3])
+    local ttl = tonumber(ARGV[4])
+    local now = tonumber(ARGV[5])
+
+    -- Get count of current processes. If more than the max, check if any of the processes have timed out
+    -- and try again.
+    local process_count = redis.call('zcard', sorted_set)
+    if process_count >= limit then
+      local max_score = now - ttl
+      local expired_keys = redis.call('zremrangebyscore', sorted_set, '-inf', max_score)
+      if expired_keys > 0 then
+        process_count = redis.call('zcard', sorted_set)
+      end
+    end
+
+    -- Success so add to the list and set a global expiration so the list cleans up after itself.
+    if process_count < limit then
+      redis.call('zadd', sorted_set, now, process_id)
+      redis.call('expire', sorted_set, ttl)
+    end
+    
+    -- Return the number of processes running before the process was added.
+    return process_count
+  LUA
+
+  # This class level variable will be used to load the SHA1 of the script at runtime.
+  @add_process_sha1 = nil
+
+  # This error will be thrown when the throttle block can't be executed.
   class ThrottledError < StandardError
   end
-  
+
   class << self
     # Either configure the redis instance using a block or yield the instance. Configuring with
     # a block allows you to use things like connection pools etc. without hard coding a single
     # instance.
     #
-    # Example: `Restrainer.redis{ redis_pool.instance }
+    # Example: `Restrainer.redis { redis_pool.instance }`
     def redis(&block)
       if block
         @redis = block
@@ -33,21 +68,23 @@ class Restrainer
         raise "#{self.class.name}.redis not configured"
       end
     end
-    
+
     # Set the redis instance to a specific instance. It is usually preferable to use the block
-    # form for configurating the instance.
+    # form for configurating the instance so that it can be evaluated at runtime.
+    #
+    # Example: `Restrainer.redis = Redis.new`
     def redis=(conn)
       @redis = lambda{ conn }
     end
   end
-  
+
   # Create a new restrainer. The name is used to identify the Restrainer and group processes together.
   # You can create any number of Restrainers with different names.
   #
   # The required limit parameter specifies the maximum number of processes that will be allowed to execute the
   # throttle block at any point in time.
   #
-  # The timeout parameter is used for cleaning up internal data structures so that jobs aren't orphaned 
+  # The timeout parameter is used for cleaning up internal data structures so that jobs aren't orphaned
   # if their process is killed. Processes will automatically be removed from the running jobs list after the
   # specified number of seconds. Note that the Restrainer will not handle timing out any code itself. This
   # value is just used to insure the integrity of internal data structures.
@@ -57,7 +94,7 @@ class Restrainer
     @timeout = timeout
     @key = "#{self.class.name}.#{name.to_s}"
   end
-  
+
   # Wrap a block with this method to throttle concurrent execution. If more than the alotted number
   # of processes (as identified by the name) are currently executing, then a Restrainer::ThrottledError
   # will be raised.
@@ -65,66 +102,67 @@ class Restrainer
   # The limit argument can be used to override the value set in the constructor.
   def throttle(limit: nil)
     limit ||= self.limit
-    
+
     # limit of less zero is no limit; limit of zero is allow none
     return yield if limit < 0
     raise ThrottledError.new("#{self.class}: #{@name} is not allowing any processing") if limit == 0
-    
+
     # Grab a reference to the redis instance to that it will be consistent throughout the method
     redis = self.class.redis
-    check_running_count!(redis, limit)
     process_id = SecureRandom.uuid
+    add_process!(redis, process_id, limit)
+    
     begin
-      add_process!(redis, process_id)
       yield
     ensure
       remove_process!(redis, process_id)
     end
   end
-  
+
   # Get the number of processes currently being executed for this restrainer.
   def current(redis = nil)
     redis ||= self.class.redis
     redis.zcard(key).to_i
   end
-  
+
   private
-  
+
   # Hash key in redis to story a sorted set of current processes.
   def key
     @key
   end
-  
-  # Raise an error if there are too many running processes.
-  def check_running_count!(redis, limit)
-    running_count = current(redis)
-    if running_count >= limit
-      running_count = current(redis) if cleanup!(redis)
-      if running_count >= limit
-        raise ThrottledError.new("#{self.class}: #{@name} already has #{running_count} processes running")
-      end
-    end
-  end
-  
+
   # Add a process to the currently run set.
-  def add_process!(redis, process_id)
-    redis.multi do |conn|
-      conn.zadd(key, Time.now.to_i, process_id)
-      conn.expire(key, @timeout)
+  def add_process!(redis, process_id, throttle_limit)
+    process_count = eval_script(redis, process_id, throttle_limit)    
+    if process_count >= throttle_limit
+      raise ThrottledError.new("#{self.class}: #{@name} already has #{process_count} processes running")
     end
   end
-  
+
   # Remove a process to the currently run set.
   def remove_process!(redis, process_id)
     redis.zrem(key, process_id)
   end
-  
-  # Protect against kill -9 which can cause processes to not be removed from the lists.
-  # Processes will be assumed to have finished by a specified timeout (in seconds).
-  def cleanup!(redis)
-    max_score = Time.now.to_i - @timeout
-    expired = redis.zremrangebyscore(key, "-inf", max_score)
-    expired > 0 ? true : false
+
+  # Evaluate and execute a Lua script on the redis server.
+  def eval_script(redis, process_id, throttle_limit)
+    sha1 = @add_process_sha1
+    if sha1 == nil
+      sha1 = redis.script(:load, ADD_PROCESS_SCRIPT)
+      @add_process_sha1 = sha1
+    end
+
+    begin
+      redis.evalsha(sha1, [], [key, process_id, throttle_limit, @timeout, Time.now.to_i])
+    rescue Redis::CommandError => e
+      if e.message.include?('NOSCRIPT')
+        sha1 = redis.script(:load, ADD_PROCESS_SCRIPT)
+        @add_process_sha1 = sha1
+        retry
+      else
+        raise e
+      end
+    end
   end
-  
 end

data/spec/restrainer_spec.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require 'spec_helper'
 
 describe Restrainer do
@@ -17,10 +19,18 @@ describe Restrainer do
   end
   
   it "should throw an error if too many processes are already running" do
-    restrainer = Restrainer.new(:restrainer_test, limit: 1)
+    restrainer = Restrainer.new(:restrainer_test, limit: 5)
     x = nil
     restrainer.throttle do
-      expect(lambda{restrainer.throttle{ x = 1 }}).to raise_error(Restrainer::ThrottledError)
+      restrainer.throttle do
+        restrainer.throttle do
+          restrainer.throttle do
+            restrainer.throttle do
+              expect(lambda{restrainer.throttle{ x = 1 }}).to raise_error(Restrainer::ThrottledError)
+            end
+          end
+        end
+      end
     end
     expect(x).to eq(nil)
   end

data/spec/spec_helper.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require File.expand_path('../../lib/restrainer', __FILE__)
 require 'timecop'
 
@@ -10,6 +12,6 @@ RSpec.configure do |config|
   # the seed, which is printed after each run.
   #     --seed 1234
   config.order = 'random'
-  
+
   Restrainer.redis = Redis.new
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: restrainer
 version: !ruby/object:Gem::Version
-  version: 1.0.0
+  version: 1.0.1
 platform: ruby
 authors:
 - We Heart It
@@ -9,7 +9,7 @@ authors:
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2015-10-30 00:00:00.000000000 Z
+date: 2019-07-03 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: redis
@@ -90,6 +90,7 @@ extensions: []
 extra_rdoc_files: []
 files:
 - ".gitignore"
+- CHANGE_LOG.md
 - MIT_LICENSE.txt
 - README.md
 - Rakefile
@@ -117,8 +118,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubyforge_project: 
-rubygems_version: 2.4.5
+rubygems_version: 3.0.3
 signing_key: 
 specification_version: 4
 summary: Code for throttling workloads so as not to overwhelm external services