berater 0.6.1 → 0.9.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 5c5cc50ba02c18bf9ef9f6b9c1dd6ccb8f8af0297df95757ac13406dc4539c40
-  data.tar.gz: f80341dccff30ca3f344f3bdf21de44b1ef1715bfaf4140933a91252ddb7231e
+  metadata.gz: 476b59e3f1e27908f5c8a5097087e6c8c2384531e1284d79b9d89a73fbbb5840
+  data.tar.gz: de50fdcab6ea7dc9520fcf1aefdd07b1f8775c0ce96a30e3f1d45b6ccfb7b00f
 SHA512:
-  metadata.gz: ada3fffc841e71498f652dc990eed482bccc942b86280fc5ee736d49fdf99d8ab2f921b5b52d96f3588c268c34343669c00e0be8822c462a3e6f9dcf954a4005
-  data.tar.gz: 0363abef69cfdad8140d89373b5088cfa8a6259f16716011c82ceb1e2e9e948cdb4d44ddd36aaf1d05b121dc7d275d0ee34571a0a3ee9b964f4286375ec356cf
+  metadata.gz: 2446e64f528a8ef1d37791922d37cddcebdf702f3c5beff2ef7ea9d611ce4cc6d6c40ac42e1ea826987a2eb2ceea3ddf97e420bb7369686b8239b260a1725c55
+  data.tar.gz: 23dccad7e02cc425631ceeb1dfbc6960b94ac92c523d840a991917af02caaf39a7befc5887fe45631f565ddb2eddb4850e557348ab29427d3cbae1db89124c3c

data/lib/berater.rb

@@ -1,8 +1,10 @@
 require 'berater/limiter'
+require 'berater/limiter_set'
 require 'berater/lock'
 require 'berater/lua_script'
 require 'berater/utils'
 require 'berater/version'
+require 'meddleware'
 
 module Berater
   extend self
@@ -15,24 +17,35 @@ module Berater
     yield self
   end
 
-  def reset
-    @redis = nil
+  def limiters
+    @limiters ||= LimiterSet.new
+  end
+
+  def middleware(&block)
+    (@middleware ||= Meddleware.new).tap do
+      @middleware.instance_eval(&block) if block_given?
+    end
   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
@@ -43,20 +56,21 @@ module Berater
     end
   end
 
+  def reset
+    @redis = nil
+    limiters.clear
+    middleware.clear
+  end
 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)
@@ -64,7 +64,7 @@ module Berater
     )
 
     protected def acquire_lock(capacity, cost)
-      # fractional cost is not supported, but make it work
+      # round fractional capacity and cost
       capacity = capacity.to_i
       cost = cost.ceil
 
@@ -73,26 +73,19 @@ module Berater
 
       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.new(capacity, count, release_fn)
     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
-    end
-
     def to_s
       "#<#{self.class}(#{key}: #{capacity} at a time)>"
     end

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

data/lib/berater/limiter.rb

@@ -9,17 +9,12 @@ module Berater
 
     def limit(capacity: nil, cost: 1, &block)
       capacity ||= @capacity
+      lock = nil
 
-      unless capacity.is_a?(Numeric)
-        raise ArgumentError, "invalid capacity: #{capacity}"
-      end
-
-      unless cost.is_a?(Numeric) && cost >= 0
-        raise ArgumentError, "invalid cost: #{cost}"
+      Berater.middleware.call(self, capacity: capacity, cost: cost) do |limiter, **opts|
+        lock = limiter.inner_limit(**opts)
       end
 
-      lock = acquire_lock(capacity, cost)
-
       if block_given?
         begin
           yield lock
@@ -31,14 +26,33 @@ module Berater
       end
     end
 
-    def overloaded?
-      limit(cost: 0) { false }
-    rescue Overloaded
-      true
+    protected def inner_limit(capacity:, cost:)
+      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
+
+      acquire_lock(capacity, cost)
+    rescue NoMethodError => e
+      raise unless e.message.include?("undefined method `evalsha' for")
+
+      # repackage error so it's easier to understand
+      raise RuntimeError, "invalid redis connection: #{redis}"
     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)
@@ -50,13 +64,6 @@ module Berater
       self.redis.connection == other.redis.connection
     end
 
-    def self.new(*)
-      # can only call via subclass
-      raise NotImplementedError if self == Berater::Limiter
-
-      super
-    end
-
     protected
 
     attr_reader :args
@@ -86,8 +93,33 @@ module Berater
       raise NotImplementedError
     end
 
-    def cache_key(key)
-      "#{self.class}:#{key}"
+    def cache_key(subkey = nil)
+      instance_key = subkey.nil? ? key : "#{key}:#{subkey}"
+      self.class.cache_key(instance_key)
+    end
+
+    class << self
+      def new(*)
+        # can only call via subclass
+        raise NoMethodError if self == Berater::Limiter
+
+        super
+      end
+
+      def cache_key(key)
+        klass = to_s.split(':')[-1]
+        "Berater:#{klass}:#{key}"
+      end
+
+      protected
+
+      def 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
     end
 
   end

data/lib/berater/limiter_set.rb

@@ -0,0 +1,66 @@
+module Berater
+  private
+
+  class LimiterSet
+    include Enumerable
+
+    def initialize
+      @limiters = {}
+    end
+
+    def each(&block)
+      @limiters.each_value(&block)
+    end
+
+    def <<(limiter)
+      key = limiter.key if limiter.respond_to?(:key)
+      send(:[]=, key, limiter)
+    end
+
+    def []=(key, limiter)
+      unless limiter.is_a? Berater::Limiter
+        raise ArgumentError, "expected Berater::Limiter, found: #{limiter}"
+      end
+
+      @limiters[key] = limiter
+    end
+
+    def [](key)
+      @limiters[key]
+    end
+
+    def fetch(key, val = default = true, &block)
+      args = default ? [ key ] : [ key, val ]
+      @limiters.fetch(*args, &block)
+    end
+
+    def include?(key)
+      if key.is_a? Berater::Limiter
+        @limiters.value?(key)
+      else
+        @limiters.key?(key)
+      end
+    end
+
+    def clear
+      @limiters.clear
+    end
+
+    def count
+      @limiters.count
+    end
+
+    def delete(key)
+      if key.is_a? Berater::Limiter
+        @limiters.delete(key.key)
+      else
+        @limiters.delete(key)
+      end
+    end
+    alias remove delete
+
+    def empty?
+      @limiters.empty?
+    end
+  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,14 +59,15 @@ 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
     )
 
@@ -69,17 +77,17 @@ module Berater
 
       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
+
+      raise Overloaded unless allowed
 
       Lock.new(capacity, count)
     end
 
-    alias overrated? overloaded?
-
     def to_s
       msg = if interval.is_a? Numeric
         if interval == 1