berater 0.3.0 → 0.6.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: d007dbb664fa57c921047a28dd1e5a591cf5c51fba4931dbc367ab552e4b8a8d
-  data.tar.gz: c710d31a2d8cfdc50af4c2fae87cabad177096d13da52a7f9c48b61aa31b6038
+  metadata.gz: dd61012052d49f83a59057ed5ce4cc46b9c25a2d2bf4b5e494f02e590c2e2249
+  data.tar.gz: 3add1b481830566088eb7a9dce3540716b55dda9bb658a7b26b93a261b403f3a
 SHA512:
-  metadata.gz: 0d9c211bb65e7341b98e637d3f0ce0e31f0a879437c1c309fb01e1efab5e5c6b2d5044158931ab16a67cbbdebd00ddb5f10afc97430230b153d8cc236f7517ab
-  data.tar.gz: 5d8b38bfd7683df641fd6891cd1a493079453ea6c61b2880581eeb2d963fb171365fb3448443e2548b34297d8dcfa352e67cff5688cd24d25422efb5e55d4afd
+  metadata.gz: 476cf552a5fb56498f81a737dba79358d016bedbd76c3c2cc8069c420e82a221c54280781821ef701d25da2551ddc25910818dc400dc7c046cecbe8f6db0ddd2
+  data.tar.gz: a92e64f027425ecdc4a39ee5f88a792f4f2625ec4bcf65fac70e34d591fb136e7544caa80138474ae5f5e6158fc19ec8894892585cf858966bfe5fe13a3d011b

data/lib/berater.rb

@@ -1,48 +1,40 @@
-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
-
-      unless block_given?
-        raise ArgumentError, 'expected either mode or block'
-      end
+  def reset
+    @redis = nil
+  end
 
-      mode, *args = DSL.eval(&block)
+  def new(key, capacity, interval = nil, **opts)
+    case capacity
+    when :unlimited, Float::INFINITY
+      Berater::Unlimiter
+    when :inhibited, 0
+      Berater::Inhibitor
     else
-      if block_given?
-        raise ArgumentError, 'expected either mode or block, not both'
+      if interval
+        Berater::RateLimiter
+      else
+        Berater::ConcurrencyLimiter
       end
+    end.yield_self do |klass|
+      args = [ key, capacity, interval ].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 +46,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, interval = nil, **opts, &block)
+  limiter = Berater.new(key, capacity, interval, **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

@@ -3,42 +3,28 @@ module Berater
 
     class Incapacitated < Overloaded; end
 
-    attr_reader :capacity, :timeout
+    attr_reader :timeout
 
     def initialize(key, capacity, **opts)
-      super(key, **opts)
+      super(key, capacity, **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
-
-      raise ArgumentError, "capacity must be >= 0" unless capacity >= 0
-
-      @capacity = capacity
-    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*(--.*\n)?/, '')
+    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
@@ -48,41 +34,63 @@ module Berater
       -- 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 check limit, ie. for .overlimit?
+        if count < capacity then
+          table.insert(lock_ids, true)
+        end
+      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
+    )
+
+    protected def acquire_lock(capacity, cost)
+      # fractional cost is not supported, but make it work
+      capacity = capacity.to_i
+      cost = cost.ceil
 
-    def limit
-      count, lock_id = redis.eval(
-        LUA_SCRIPT,
+      # 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 Incapacitated 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
+    alias incapacitated? overloaded?
+
+    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,23 @@
 module Berater
   module DSL
+    refine Berater.singleton_class do
+      def new(key, mode = nil, *args, **opts, &block)
+        if mode.nil?
+          unless block_given?
+            raise ArgumentError, 'expected either mode or block'
+          end
+
+          mode, *args = DSL.eval(&block)
+        else
+          if block_given?
+            raise ArgumentError, 'expected either mode or block, not both'
+          end
+        end
+
+        super(key, mode, *args, **opts)
+      end
+    end
+
     extend self
 
     def eval &block
@@ -18,13 +36,6 @@ 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,
@@ -33,12 +44,12 @@ module Berater
     def install
       Integer.class_eval do
         def per(unit)
-          [ :rate, self, unit ]
+          [ self, 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

@@ -4,10 +4,12 @@ module Berater
     class Inhibited < Overloaded; end
 
     def initialize(key = :inhibitor, *args, **opts)
-      super(key, **opts)
+      super(key, 0, **opts)
     end
 
-    def limit
+    alias inhibited? overloaded?
+
+    protected def acquire_lock(*)
       raise Inhibited
     end
 

data/lib/berater/limiter.rb

@@ -1,27 +1,91 @@
 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)
+        raise ArgumentError, "invalid capacity: #{capacity}"
+      end
+
+      unless cost.is_a?(Numeric) && cost >= 0
+        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 overloaded?
+      limit(cost: 0) { false }
+    rescue Overloaded
+      true
     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