berater 0.5.0 → 0.7.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: fa5d51d6eea9d4691e58e8dd7065a1d596a57abdac1227fa19bd22cd2f42c15d
-  data.tar.gz: 55ac15bf6f2f295c525fdea57b3af9b518935a12dd4b644e3c0c9c085429ca50
+  metadata.gz: c12ee8f042d6ebb193f114cd093585c14a4a28b802d7a90f79ba5ce41bbd82e6
+  data.tar.gz: d12ebb1107ea0b24208bfb3ff663a98097cd3367c3972418a139cae9b501e4a1
 SHA512:
-  metadata.gz: 52df5781f9c37c7f1c9dde09e579c8da1f2d62f405835b771c6087cda8af046018e85a88629f7c5d3a2be9aafd2680e348d1b2f91bc63fe456e6e51ec5d0fd02
-  data.tar.gz: '0928cb6f2e8e07810dec96f5af302ff9f62d519b9f19ef6f28acb5621569f13edb6d784cbd7084b738df8dfc232ace26e3ce0c6c4a7f2ac3e4391c97789bafae'
+  metadata.gz: 491cd616903af917866a41f8b2a4bcc79a469a879f57f1872ce8f4a42c00f1309979c04c142c0083246791aef6cd3847c21197dee0b611ff326ca833cc44ae52
+  data.tar.gz: 7ba3fd90f8b439e4ad2fe1d5a8008c69978a90c206bc6694e9f470796019d2ac849dddaa85039b9c730e3e20edfec3411cc1275598866c63207c209e688f2a3a

data/lib/berater.rb

@@ -19,20 +19,23 @@ 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
         Berater::ConcurrencyLimiter
       end
     end.yield_self do |klass|
-      args = [ key, capacity, interval ].compact
+      args = [ key, capacity, *args ].compact
       klass.new(*args, **opts)
     end
   end
@@ -46,8 +49,8 @@ module Berater
 end
 
 # convenience method
-def Berater(key, capacity, interval = nil, **opts, &block)
-  limiter = Berater.new(key, capacity, interval, **opts)
+def Berater(key, capacity, **opts, &block)
+  limiter = Berater.new(key, capacity, **opts)
   if block_given?
     limiter.limit(&block)
   else

data/lib/berater/concurrency_limiter.rb

@@ -1,20 +1,21 @@
 module Berater
   class ConcurrencyLimiter < Limiter
 
-    class Incapacitated < Overloaded; end
-
     attr_reader :timeout
 
     def initialize(key, capacity, **opts)
       super(key, capacity, **opts)
 
+      # round fractional capacity
+      self.capacity = capacity.to_i
+
       self.timeout = opts[:timeout] || 0
     end
 
     private def timeout=(timeout)
       @timeout = timeout
       timeout = 0 if timeout == Float::INFINITY
-      @timeout_usec = Berater::Utils.to_usec(timeout)
+      @timeout_msec = Berater::Utils.to_msec(timeout)
     end
 
     LUA_SCRIPT = Berater::LuaScript(<<~LUA
@@ -28,17 +29,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,36 +62,28 @@ module Berater
     LUA
     )
 
-    def limit(capacity: nil, cost: 1, &block)
-      capacity ||= @capacity
-      # cost is Integer >= 0
+    protected def acquire_lock(capacity, cost)
+      # round fractional capacity and cost
+      capacity = capacity.to_i
+      cost = cost.ceil
 
-      # timestamp in microseconds
-      ts = (Time.now.to_f * 10**6).to_i
+      # 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_usec, cost ]
+        [ capacity, ts, @timeout_msec, cost ]
       )
 
-      raise Incapacitated if lock_ids.empty?
+      raise Overloaded if lock_ids.empty?
 
-      if cost == 0
-        lock = Lock.new(self, nil, count)
-      else
-        lock = Lock.new(self, lock_ids[0], count, -> { release(lock_ids) })
+      release_fn = if cost > 0
+        proc { release(lock_ids) }
       end
 
-      yield_lock(lock, &block)
-    end
-
-    def overloaded?
-      limit(cost: 0) { false }
-    rescue Overloaded
-      true
+      Lock.new(capacity, count, release_fn)
     end
-    alias incapacitated? overloaded?
 
     private def release(lock_ids)
       res = redis.zrem(cache_key(key), lock_ids)

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,13 @@
 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
-    end
-
-    def overloaded?
-      true
+    protected def acquire_lock(*)
+      raise Overloaded
     end
-    alias inhibited? overloaded?
 
   end
 end

data/lib/berater/limiter.rb

@@ -7,12 +7,40 @@ module Berater
       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 overloaded?
-      raise NotImplementedError
+    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
@@ -30,7 +58,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
@@ -47,29 +75,25 @@ module Berater
     end
 
     def capacity=(capacity)
-      unless capacity.is_a?(Integer) || capacity == Float::INFINITY
-        raise ArgumentError, "expected Integer, found #{capacity.class}"
+      unless capacity.is_a?(Numeric)
+        raise ArgumentError, "expected Numeric, found #{capacity.class}"
       end
 
-      raise ArgumentError, "capacity must be >= 0" unless capacity >= 0
+      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 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
-      end
+    def cache_key(key)
+      "#{self.class}:#{key}"
     end
 
   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

data/lib/berater/rate_limiter.rb

@@ -1,18 +1,20 @@
 module Berater
   class RateLimiter < Limiter
 
-    class Overrated < Overloaded; end
-
     attr_accessor :interval
 
     def initialize(key, capacity, interval, **opts)
       self.interval = interval
-      super(key, capacity, @interval_usec, **opts)
+      super(key, capacity, @interval_msec, **opts)
     end
 
     private def interval=(interval)
       @interval = interval
-      @interval_usec = Berater::Utils.to_usec(interval)
+      @interval_msec = Berater::Utils.to_msec(interval)
+
+      unless @interval_msec > 0
+        raise ArgumentError, 'interval must be > 0'
+      end
     end
 
     LUA_SCRIPT = Berater::LuaScript(<<~LUA
@@ -20,26 +22,34 @@ module Berater
       local ts_key = KEYS[2]
       local ts = tonumber(ARGV[1])
       local capacity = tonumber(ARGV[2])
-      local interval_usec = tonumber(ARGV[3])
+      local interval_msec = tonumber(ARGV[3])
       local cost = tonumber(ARGV[4])
-      local count = 0
-      local allowed
-      local usec_per_drip = interval_usec / 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) / usec_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
 
@@ -47,42 +57,35 @@ module Berater
           count = count + cost
 
           -- time for bucket to empty, in milliseconds
-          local ttl = math.ceil((count * usec_per_drip) / 1000)
+          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
-
-      # timestamp in microseconds
-      ts = (Time.now.to_f * 10**6).to_i
+    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_usec, cost ]
+        [ cache_key(key) ],
+        [ ts, capacity, @interval_msec, cost ]
       )
 
-      raise Overrated unless allowed
+      count = count.include?('.') ? count.to_f : count.to_i
 
-      lock = Lock.new(self, ts, 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