berater 0.4.0 → 0.7.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: f57aeaa9319ae38bdd7d2468443fc5c9ea9ba44cecf47bc5a558f18ecd0ab274
-  data.tar.gz: ff14ef90856911e9e18cbfbad6f9799727c999697c5e4d96f47b5fbf665c725c
+  metadata.gz: 7d25d186bfb9e709986e5c205c9fd2b8107dcf8df1e915316d28a6a59a0d4c0f
+  data.tar.gz: '05380448ad96697b5d3753a93584d4bdfcf3b028c649ca11360a1c92d6afadad'
 SHA512:
-  metadata.gz: c9d3abeeb9d35c608c194eaf505c5af1350a82583dd12257e70b9d4f96bd3a6cd30db802f57c3bf77d371b5deb5e856e9238eadd3990768205fd37ea3e2e3703
-  data.tar.gz: c98078a3c7d3a3eb4e2d523577f9d6dd7f1bbae76c9ecbd8d65961cc7aa1bab6e1a4b8fe84cfb30ee60f4d0ce11dcb58e3b94a3566cf50ee60d1f8e21bab82b1
+  metadata.gz: 6ddcd488d3d1aa293d621f32ad8cd684daffb0ba9c413e8cf6a7a292ae63a307e2cdf8e62803cf2a8c2b3ea7b9424e4874bb773f49620e768895b193385e847f
+  data.tar.gz: cf923f95875da8b7e6c0c0cbe4dccad5a4e3c9f6436da606a03f93c2b9ff86525485c6379ade0f276f30e1aff3bcf0c279fe88d33e6769455818482a5ff1a179

data/lib/berater.rb

@@ -1,48 +1,43 @@
-require 'berater/version'
+require 'berater/limiter'
 require 'berater/lock'
-
+require 'berater/lua_script'
+require 'berater/utils'
+require 'berater/version'
 
 module Berater
   extend self
 
   class Overloaded < StandardError; end
 
-  MODES = {}
-
   attr_accessor :redis
 
   def configure
     yield self
   end
 
-  def new(key, mode = nil, *args, **opts, &block)
-    if mode.nil?
-      unless args.empty?
-        raise ArgumentError, '0 arguments expected with block'
-      end
+  def reset
+    @redis = nil
+  end
 
-      unless block_given?
-        raise ArgumentError, 'expected either mode or block'
-      end
+  def new(key, capacity, **opts)
+    args = []
 
-      mode, *args = DSL.eval(&block)
+    case capacity
+    when Float::INFINITY
+      Berater::Unlimiter
+    when 0
+      Berater::Inhibitor
     else
-      if block_given?
-        raise ArgumentError, 'expected either mode or block, not both'
+      if opts[:interval]
+        args << opts.delete(:interval)
+        Berater::RateLimiter
+      else
+        Berater::ConcurrencyLimiter
       end
+    end.yield_self do |klass|
+      args = [ key, capacity, *args ].compact
+      klass.new(*args, **opts)
     end
-
-    klass = MODES[mode.to_sym]
-
-    unless klass
-      raise ArgumentError, "invalid mode: #{mode}"
-    end
-
-    klass.new(key, *args, **opts)
-  end
-
-  def register(mode, klass)
-    MODES[mode.to_sym] = klass
   end
 
   def expunge
@@ -54,20 +49,17 @@ module Berater
 end
 
 # convenience method
-def Berater(key, mode, *args, **opts, &block)
-  Berater.new(key, mode, *args, **opts).limit(&block)
+def Berater(key, capacity, **opts, &block)
+  limiter = Berater.new(key, capacity, **opts)
+  if block_given?
+    limiter.limit(&block)
+  else
+    limiter
+  end
 end
 
 # load limiters
-require 'berater/limiter'
 require 'berater/concurrency_limiter'
 require 'berater/inhibitor'
 require 'berater/rate_limiter'
 require 'berater/unlimiter'
-
-Berater.register(:concurrency, Berater::ConcurrencyLimiter)
-Berater.register(:inhibited, Berater::Inhibitor)
-Berater.register(:rate, Berater::RateLimiter)
-Berater.register(:unlimited, Berater::Unlimiter)
-
-require 'berater/dsl'

data/lib/berater/concurrency_limiter.rb

@@ -1,88 +1,93 @@
 module Berater
   class ConcurrencyLimiter < Limiter
 
-    class Incapacitated < Overloaded; end
-
-    attr_reader :capacity, :timeout
+    attr_reader :timeout
 
     def initialize(key, capacity, **opts)
-      super(key, **opts)
-
-      self.capacity = capacity
-      self.timeout = opts[:timeout] || 0
-    end
-
-    private def capacity=(capacity)
-      unless capacity.is_a? Integer
-        raise ArgumentError, "expected Integer, found #{capacity.class}"
-      end
+      super(key, capacity, **opts)
 
-      raise ArgumentError, "capacity must be >= 0" unless capacity >= 0
+      # round fractional capacity
+      self.capacity = capacity.to_i
 
-      @capacity = capacity
+      self.timeout = opts[:timeout] || 0
     end
 
     private def timeout=(timeout)
-      unless timeout.is_a? Integer
-        raise ArgumentError, "expected Integer, found #{timeout.class}"
-      end
-
-      raise ArgumentError, "timeout must be >= 0" unless timeout >= 0
-
       @timeout = timeout
+      timeout = 0 if timeout == Float::INFINITY
+      @timeout_msec = Berater::Utils.to_msec(timeout)
     end
 
-    LUA_SCRIPT = <<~LUA.gsub(/^\s*|\s*--.*/, '')
+    LUA_SCRIPT = Berater::LuaScript(<<~LUA
       local key = KEYS[1]
       local lock_key = KEYS[2]
       local capacity = tonumber(ARGV[1])
       local ts = tonumber(ARGV[2])
       local ttl = tonumber(ARGV[3])
-      local lock
+      local cost = tonumber(ARGV[4])
+      local lock_ids = {}
 
       -- purge stale hosts
       if ttl > 0 then
-        redis.call('ZREMRANGEBYSCORE', key, '-inf', ts - ttl)
+        redis.call('ZREMRANGEBYSCORE', key, 0, ts - ttl)
       end
 
       -- check capacity
       local count = redis.call('ZCARD', key)
 
-      if count < capacity then
-        -- grab a lock
-        lock = redis.call('INCR', lock_key)
-        redis.call('ZADD', key, ts, lock)
-        count = count + 1
+      if cost == 0 then
+        -- just checking count
+        table.insert(lock_ids, true)
+      elseif (count + cost) <= capacity then
+        -- grab locks, one per cost
+        local lock_id = redis.call('INCRBY', lock_key, cost)
+        local locks = {}
+
+        for i = lock_id - cost + 1, lock_id do
+          table.insert(lock_ids, i)
+
+          table.insert(locks, ts)
+          table.insert(locks, i)
+        end
+
+        redis.call('ZADD', key, unpack(locks))
+        count = count + cost
+
+        if ttl > 0 then
+          redis.call('PEXPIRE', key, ttl)
+        end
       end
 
-      return { count, lock }
+      return { count, unpack(lock_ids) }
     LUA
+    )
 
-    def limit
-      count, lock_id = redis.eval(
-        LUA_SCRIPT,
+    protected def acquire_lock(capacity, cost)
+      # round fractional capacity and cost
+      capacity = capacity.to_i
+      cost = cost.ceil
+
+      # timestamp in milliseconds
+      ts = (Time.now.to_f * 10**3).to_i
+
+      count, *lock_ids = LUA_SCRIPT.eval(
+        redis,
         [ cache_key(key), cache_key('lock_id') ],
-        [ capacity, Time.now.to_i, timeout ]
+        [ capacity, ts, @timeout_msec, cost ]
       )
 
-      raise Incapacitated unless lock_id
-
-      lock = Lock.new(self, lock_id, count, -> { release(lock_id) })
+      raise Overloaded if lock_ids.empty?
 
-      if block_given?
-        begin
-          yield lock
-        ensure
-          lock.release
-        end
-      else
-        lock
+      release_fn = if cost > 0
+        proc { release(lock_ids) }
       end
+
+      Lock.new(capacity, count, release_fn)
     end
 
-    private def release(lock_id)
-      res = redis.zrem(cache_key(key), lock_id)
-      res == true || res == 1 # depending on which version of Redis
+    private def release(lock_ids)
+      res = redis.zrem(cache_key(key), lock_ids)
+      res == true || res == lock_ids.count # depending on which version of Redis
     end
 
     def to_s

data/lib/berater/dsl.rb

@@ -1,5 +1,24 @@
 module Berater
   module DSL
+    refine Berater.singleton_class do
+      def new(key, capacity = nil, **opts, &block)
+        if capacity.nil?
+          unless block_given?
+            raise ArgumentError, 'expected either capacity or block'
+          end
+
+          capacity, more_opts = DSL.eval(&block)
+          opts.merge!(more_opts) if more_opts
+        else
+          if block_given?
+            raise ArgumentError, 'expected either capacity or block, not both'
+          end
+        end
+
+        super(key, capacity, **opts)
+      end
+    end
+
     extend self
 
     def eval &block
@@ -18,27 +37,19 @@ module Berater
 
     private
 
-    def each &block
-      Berater::MODES.map do |mode, limiter|
-        next unless limiter.const_defined?(:DSL, false)
-        limiter.const_get(:DSL)
-      end.compact.each(&block)
-    end
-
     KEYWORDS = [
       :second, :minute, :hour,
-      :unlimited, :inhibited,
     ].freeze
 
     def install
       Integer.class_eval do
         def per(unit)
-          [ :rate, self, unit ]
+          [ self, interval: unit ]
         end
         alias every per
 
         def at_once
-          [ :concurrency, self ]
+          [ self ]
         end
         alias concurrently at_once
         alias at_a_time at_once

data/lib/berater/inhibitor.rb

@@ -1,14 +1,12 @@
 module Berater
   class Inhibitor < Limiter
 
-    class Inhibited < Overloaded; end
-
     def initialize(key = :inhibitor, *args, **opts)
-      super(key, **opts)
+      super(key, 0, **opts)
     end
 
-    def limit
-      raise Inhibited
+    protected def acquire_lock(*)
+      raise Overloaded
     end
 
   end

data/lib/berater/limiter.rb

@@ -1,27 +1,97 @@
 module Berater
   class Limiter
 
-    attr_reader :key, :options
+    attr_reader :key, :capacity, :options
 
     def redis
       options[:redis] || Berater.redis
     end
 
-    def limit
-      raise NotImplementedError
+    def limit(capacity: nil, cost: 1, &block)
+      capacity ||= @capacity
+
+      unless capacity.is_a?(Numeric) && capacity >= 0
+        raise ArgumentError, "invalid capacity: #{capacity}"
+      end
+
+      unless cost.is_a?(Numeric) && cost >= 0 && cost < Float::INFINITY
+        raise ArgumentError, "invalid cost: #{cost}"
+      end
+
+      lock = acquire_lock(capacity, cost)
+
+      if block_given?
+        begin
+          yield lock
+        ensure
+          lock.release
+        end
+      else
+        lock
+      end
+    end
+
+    def utilization
+      lock = limit(cost: 0)
+
+      if lock.capacity == 0
+        1.0
+      else
+        lock.contention.to_f / lock.capacity
+      end
+    rescue Berater::Overloaded
+      1.0
     end
 
     def to_s
       "#<#{self.class}>"
     end
 
+    def ==(other)
+      self.class == other.class &&
+      self.key == other.key &&
+      self.capacity == other.capacity &&
+      self.args == other.args &&
+      self.options == other.options &&
+      self.redis.connection == other.redis.connection
+    end
+
+    def self.new(*)
+      # can only call via subclass
+      raise NoMethodError if self == Berater::Limiter
+
+      super
+    end
+
     protected
 
-    def initialize(key, **opts)
+    attr_reader :args
+
+    def initialize(key, capacity, *args, **opts)
       @key = key
+      self.capacity = capacity
+      @args = args
       @options = opts
     end
 
+    def capacity=(capacity)
+      unless capacity.is_a?(Numeric)
+        raise ArgumentError, "expected Numeric, found #{capacity.class}"
+      end
+
+      if capacity == Float::INFINITY
+        raise ArgumentError, 'infinite capacity not supported, use Unlimiter'
+      end
+
+      raise ArgumentError, 'capacity must be >= 0' unless capacity >= 0
+
+      @capacity = capacity
+    end
+
+    def acquire_lock(capacity, cost)
+      raise NotImplementedError
+    end
+
     def cache_key(key)
       "#{self.class}:#{key}"
     end

data/lib/berater/lock.rb

@@ -1,11 +1,10 @@
 module Berater
   class Lock
 
-    attr_reader :limiter, :id, :contention
+    attr_reader :capacity, :contention
 
-    def initialize(limiter, id, contention, release_fn = nil)
-      @limiter = limiter
-      @id = id
+    def initialize(capacity, contention, release_fn = nil)
+      @capacity = capacity
       @contention = contention
       @locked_at = Time.now
       @release_fn = release_fn
@@ -13,24 +12,15 @@ module Berater
     end
 
     def locked?
-      @released_at.nil? && !expired?
-    end
-
-    def expired?
-      timeout ? @locked_at + timeout < Time.now : false
+      @released_at.nil?
     end
 
     def release
-      raise 'lock expired' if expired?
       raise 'lock already released' unless locked?
 
       @released_at = Time.now
       @release_fn ? @release_fn.call : true
     end
 
-    def timeout
-      limiter.respond_to?(:timeout) ? limiter.timeout : nil
-    end
-
   end
 end