berater 0.6.0 → 0.8.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 87d3bdb61a68a8b69dd7cb285dafb174a403acf19ff7586f90ed8b7cb50b7f76
-  data.tar.gz: fff71dd4d8ae28baf7c504efdc27df71f03219b32c1a959dddfd3d035845d058
+  metadata.gz: c32c1de752ccc7a2a71e4628b99f3561db623374dabaf95af39aadcc3b1a2d16
+  data.tar.gz: 33019954de1d79522ed7105e9f0a7da3a204df74cacd25aa4b57e749a81047bb
 SHA512:
-  metadata.gz: c9fb5c0c7ec100d2f50a3b6c6cfa449af052c1b78cff174b545e179c5fcc1324c2a5faafc90e7f7ec64ce200e47571d51b02eb56398d716c487abc7731c0cece
-  data.tar.gz: 148c566413d6593f229143e2871116b45947eae5b101d8bb2f59ce3bdfe89ed4ee1b7aaedfcaedbc80eef8a25a14fab83ad5e540172c41f31f27d7f296161235
+  metadata.gz: 89460a9a81a8d384addb095c425da0eed4fdbe0e15ecd85ee97ffcf4292441deeb34c7e4cd4591adca06096a17f7c131b79c6474d9efb2b635fbe398ad68d00e
+  data.tar.gz: 0a331b939e4c19817c42a8aa8fc5be7fd59d63f6becd5fe6ed1c226690d1524ef040a28d8bcb625d75452ab2f5c8539ddf69093bbd30c4ba9e3353259a3e568d

data/lib/berater.rb

@@ -19,20 +19,25 @@ module Berater
     @redis = nil
   end
 
-  def new(key, capacity, interval = nil, **opts)
+  def new(key, capacity, **opts)
+    args = []
+
     case capacity
-    when :unlimited, Float::INFINITY
+    when Float::INFINITY
       Berater::Unlimiter
-    when :inhibited, 0
+    when 0
       Berater::Inhibitor
     else
-      if interval
+      if opts[:interval]
+        args << opts.delete(:interval)
         Berater::RateLimiter
-      else
+      elsif opts[:timeout]
         Berater::ConcurrencyLimiter
+      else
+        Berater::StaticLimiter
       end
     end.yield_self do |klass|
-      args = [ key, capacity, interval ].compact
+      args = [ key, capacity, *args ].compact
       klass.new(*args, **opts)
     end
   end
@@ -46,17 +51,13 @@ module Berater
 end
 
 # convenience method
-def Berater(key, capacity, interval = nil, **opts, &block)
-  limiter = Berater.new(key, capacity, interval, **opts)
-  if block_given?
-    limiter.limit(&block)
-  else
-    limiter
-  end
+def Berater(*args, **opts, &block)
+  Berater::Utils.convenience_fn(Berater, *args, **opts, &block)
 end
 
 # load limiters
 require 'berater/concurrency_limiter'
 require 'berater/inhibitor'
 require 'berater/rate_limiter'
+require 'berater/static_limiter'
 require 'berater/unlimiter'

data/lib/berater/concurrency_limiter.rb

@@ -1,20 +1,22 @@
 module Berater
   class ConcurrencyLimiter < Limiter
 
-    class Incapacitated < Overloaded; end
-
-    attr_reader :timeout
-
     def initialize(key, capacity, **opts)
       super(key, capacity, **opts)
 
+      # truncate fractional capacity
+      self.capacity = capacity.to_i
+
       self.timeout = opts[:timeout] || 0
     end
 
+    def timeout
+      options[:timeout]
+    end
+
     private def timeout=(timeout)
-      @timeout = timeout
       timeout = 0 if timeout == Float::INFINITY
-      @timeout_msec = Berater::Utils.to_msec(timeout)
+      @timeout = Berater::Utils.to_msec(timeout)
     end
 
     LUA_SCRIPT = Berater::LuaScript(<<~LUA
@@ -28,17 +30,15 @@ module Berater
 
       -- 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 cost == 0 then
-        -- just check limit, ie. for .overlimit?
-        if count < capacity then
-          table.insert(lock_ids, true)
-        end
+        -- 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)
@@ -63,45 +63,27 @@ module Berater
     LUA
     )
 
-    def limit(capacity: nil, cost: 1, &block)
-      capacity ||= @capacity
-
-      # since fractional cost is not supported, capacity behaves like int
+    protected def acquire_lock(capacity, cost)
+      # round fractional capacity and cost
       capacity = capacity.to_i
-
-      unless cost.is_a?(Integer) && cost >= 0
-        raise ArgumentError, "invalid cost: #{cost}"
-      end
+      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, ts, @timeout_msec, cost ]
+        [ cache_key, self.class.cache_key('lock_id') ],
+        [ capacity, ts, @timeout, cost ]
       )
 
-      raise Incapacitated if lock_ids.empty?
+      raise Overloaded if lock_ids.empty?
 
       release_fn = if cost > 0
-        proc { release(lock_ids) }
+        proc { redis.zrem(cache_key, lock_ids) }
       end
-      lock = Lock.new(capacity, count, release_fn)
-
-      yield_lock(lock, &block)
-    end
-
-    def overloaded?
-      limit(cost: 0) { false }
-    rescue Overloaded
-      true
-    end
-    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
+      Lock.new(capacity, count, release_fn)
     end
 
     def to_s

data/lib/berater/dsl.rb

@@ -1,20 +1,21 @@
 module Berater
   module DSL
     refine Berater.singleton_class do
-      def new(key, mode = nil, *args, **opts, &block)
-        if mode.nil?
+      def new(key, capacity = nil, **opts, &block)
+        if capacity.nil?
           unless block_given?
-            raise ArgumentError, 'expected either mode or block'
+            raise ArgumentError, 'expected either capacity or block'
           end
 
-          mode, *args = DSL.eval(&block)
+          capacity, more_opts = DSL.eval(&block)
+          opts.merge!(more_opts) if more_opts
         else
           if block_given?
-            raise ArgumentError, 'expected either mode or block, not both'
+            raise ArgumentError, 'expected either capacity or block, not both'
           end
         end
 
-        super(key, mode, *args, **opts)
+        super(key, capacity, **opts)
       end
     end
 
@@ -38,13 +39,12 @@ module Berater
 
     KEYWORDS = [
       :second, :minute, :hour,
-      :unlimited, :inhibited,
     ].freeze
 
     def install
       Integer.class_eval do
         def per(unit)
-          [ self, unit ]
+          [ self, interval: unit ]
         end
         alias every per
 

data/lib/berater/inhibitor.rb

@@ -1,20 +1,17 @@
 module Berater
   class Inhibitor < Limiter
 
-    class Inhibited < Overloaded; end
-
     def initialize(key = :inhibitor, *args, **opts)
       super(key, 0, **opts)
     end
 
-    def limit(**opts)
-      raise Inhibited
+    def to_s
+      "#<#{self.class}>"
     end
 
-    def overloaded?
-      true
+    protected def acquire_lock(*)
+      raise Overloaded
     end
-    alias inhibited? overloaded?
 
   end
 end

data/lib/berater/limiter.rb

@@ -7,16 +7,40 @@ module Berater
       options[:redis] || Berater.redis
     end
 
-    def limit
-      raise NotImplementedError
-    end
+    def limit(capacity: nil, cost: 1, &block)
+      capacity ||= @capacity
 
-    def overloaded?
-      raise NotImplementedError
+      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 to_s
-      "#<#{self.class}>"
+    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 ==(other)
@@ -30,7 +54,7 @@ module Berater
 
     def self.new(*)
       # can only call via subclass
-      raise NotImplementedError if self == Berater::Limiter
+      raise NoMethodError if self == Berater::Limiter
 
       super
     end
@@ -60,19 +84,25 @@ module Berater
       @capacity = capacity
     end
 
-    def cache_key(key)
-      "#{self.class}:#{key}"
+    def acquire_lock(capacity, cost)
+      raise NotImplementedError
     end
 
-    def yield_lock(lock, &block)
-      if block_given?
-        begin
-          yield lock
-        ensure
-          lock.release
-        end
-      else
-        lock
+    def cache_key(subkey = nil)
+      instance_key = subkey.nil? ? key : "#{key}:#{subkey}"
+      self.class.cache_key(instance_key)
+    end
+
+    def self.cache_key(key)
+      klass = to_s.split(':')[-1]
+      "Berater:#{klass}:#{key}"
+    end
+
+    def self.inherited(subclass)
+      # automagically create convenience method
+      name = subclass.to_s.split(':')[-1]
+      Berater.define_singleton_method(name) do |*args, **opts, &block|
+        Berater::Utils.convenience_fn(subclass, *args, **opts, &block)
       end
     end
 

data/lib/berater/lock.rb

@@ -19,7 +19,8 @@ module Berater
       raise 'lock already released' unless locked?
 
       @released_at = Time.now
-      @release_fn ? @release_fn.call : true
+      @release_fn&.call
+      nil
     end
 
   end

data/lib/berater/rate_limiter.rb

@@ -1,20 +1,19 @@
 module Berater
   class RateLimiter < Limiter
 
-    class Overrated < Overloaded; end
-
-    attr_accessor :interval
-
     def initialize(key, capacity, interval, **opts)
+      super(key, capacity, interval, **opts)
       self.interval = interval
-      super(key, capacity, @interval_msec, **opts)
+    end
+
+    def interval
+      args[0]
     end
 
     private def interval=(interval)
-      @interval = interval
-      @interval_msec = Berater::Utils.to_msec(interval)
+      @interval = Berater::Utils.to_msec(interval)
 
-      unless @interval_msec > 0
+      unless @interval > 0
         raise ArgumentError, 'interval must be > 0'
       end
     end
@@ -26,24 +25,32 @@ module Berater
       local capacity = tonumber(ARGV[2])
       local interval_msec = tonumber(ARGV[3])
       local cost = tonumber(ARGV[4])
-      local count = 0
-      local allowed
-      local msec_per_drip = interval_msec / capacity
-
-      -- timestamp of last update
-      local last_ts = tonumber(redis.call('GET', ts_key))
-
-      if last_ts then
-        count = tonumber(redis.call('GET', key)) or 0
 
-        -- adjust for time passing
-        local drips = math.floor((ts - last_ts) / msec_per_drip)
-        count = math.max(0, count - drips)
+      local allowed -- whether lock was acquired
+      local count -- capacity being utilized
+      local msec_per_drip = interval_msec / capacity
+      local state = redis.call('GET', key)
+
+      if state then
+        local last_ts -- timestamp of last update
+        count, last_ts = string.match(state, '([%d.]+);(%w+)')
+        count = tonumber(count)
+        last_ts = tonumber(last_ts, 16)
+
+        -- adjust for time passing, guarding against clock skew
+        if ts > last_ts then
+          local drips = math.floor((ts - last_ts) / msec_per_drip)
+          count = math.max(0, count - drips)
+        else
+          ts = last_ts
+        end
+      else
+        count = 0
       end
 
       if cost == 0 then
-        -- just check limit, ie. for .overlimit?
-        allowed = count < capacity
+        -- just checking count
+        allowed = true
       else
         allowed = (count + cost) <= capacity
 
@@ -52,41 +59,34 @@ module Berater
 
           -- time for bucket to empty, in milliseconds
           local ttl = math.ceil(count * msec_per_drip)
+          ttl = ttl + 100 -- margin of error, for clock skew
 
-          -- update count and last_ts, with expirations
-          redis.call('SET', key, count, 'PX', ttl)
-          redis.call('SET', ts_key, ts, 'PX', ttl)
+          -- update count and last_ts, with expiration
+          state = string.format('%f;%X', count, ts)
+          redis.call('SET', key, state, 'PX', ttl)
         end
       end
 
-      return { count, allowed }
+      return { tostring(count), allowed }
     LUA
     )
 
-    def limit(capacity: nil, cost: 1, &block)
-      capacity ||= @capacity
-
+    protected def acquire_lock(capacity, cost)
       # timestamp in milliseconds
       ts = (Time.now.to_f * 10**3).to_i
 
       count, allowed = LUA_SCRIPT.eval(
         redis,
-        [ cache_key(key), cache_key("#{key}-ts") ],
-        [ ts, capacity, @interval_msec, cost ]
+        [ cache_key ],
+        [ ts, capacity, @interval, cost ]
       )
 
-      raise Overrated unless allowed
+      count = count.include?('.') ? count.to_f : count.to_i
 
-      lock = Lock.new(capacity, count)
-      yield_lock(lock, &block)
-    end
+      raise Overloaded unless allowed
 
-    def overloaded?
-      limit(cost: 0) { false }
-    rescue Overrated
-      true
+      Lock.new(capacity, count)
     end
-    alias overrated? overloaded?
 
     def to_s
       msg = if interval.is_a? Numeric