berater 0.7.1 → 0.10.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: c12ee8f042d6ebb193f114cd093585c14a4a28b802d7a90f79ba5ce41bbd82e6
-  data.tar.gz: d12ebb1107ea0b24208bfb3ff663a98097cd3367c3972418a139cae9b501e4a1
+  metadata.gz: e691d2d3cdc8f002e0e222c24d2e6f9a52ccb1dac7edaa92ea232ecce17fb77e
+  data.tar.gz: 5d302fff5fd063f34a94b1aae1b9b2e118b18efbe2315b4591405346f1dbc70c
 SHA512:
-  metadata.gz: 491cd616903af917866a41f8b2a4bcc79a469a879f57f1872ce8f4a42c00f1309979c04c142c0083246791aef6cd3847c21197dee0b611ff326ca833cc44ae52
-  data.tar.gz: 7ba3fd90f8b439e4ad2fe1d5a8008c69978a90c206bc6694e9f470796019d2ac849dddaa85039b9c730e3e20edfec3411cc1275598866c63207c209e688f2a3a
+  metadata.gz: 88b8e91557729601336b5235fbcc9710c79aa2f68ff393a09b5c55a1c994cef7371ed4a310c8038b07bdfc5f35d6b367ad7e1fc886c3f7c8579130808af5386b
+  data.tar.gz: 2600478d9761b14abbc2ef6c1a38b2e5b5e7476a40c6955f110746d4a90e31baabbda8d6571dcaf7d4e457d0cc19f120e84cba8915ded01b2a680111806d4400

data/lib/berater/concurrency_limiter.rb

@@ -1,21 +1,22 @@
 module Berater
   class ConcurrencyLimiter < Limiter
 
-    attr_reader :timeout
-
     def initialize(key, capacity, **opts)
       super(key, capacity, **opts)
 
-      # round fractional capacity
+      # 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
@@ -72,24 +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 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
 
-    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/inhibitor.rb

@@ -5,6 +5,10 @@ module Berater
       super(key, 0, **opts)
     end
 
+    def to_s
+      "#<#{self.class}>"
+    end
+
     protected def acquire_lock(*)
       raise Overloaded
     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) && capacity >= 0
-        raise ArgumentError, "invalid capacity: #{capacity}"
+      Berater.middleware.call(self, capacity: capacity, cost: cost) do |limiter, **opts|
+        lock = limiter.inner_limit(**opts)
       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
@@ -31,6 +26,23 @@ module Berater
       end
     end
 
+    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 utilization
       lock = limit(cost: 0)
 
@@ -43,10 +55,6 @@ module Berater
       1.0
     end
 
-    def to_s
-      "#<#{self.class}>"
-    end
-
     def ==(other)
       self.class == other.class &&
       self.key == other.key &&
@@ -56,13 +64,6 @@ module Berater
       self.redis.connection == other.redis.connection
     end
 
-    def self.new(*)
-      # can only call via subclass
-      raise NoMethodError if self == Berater::Limiter
-
-      super
-    end
-
     protected
 
     attr_reader :args
@@ -92,8 +93,39 @@ 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(*args, **kwargs)
+        # can only call via subclass
+        raise NoMethodError if self == Berater::Limiter
+
+        if RUBY_VERSION < '3' && kwargs.empty?
+          # avoid ruby 2 problems with empty hashes
+          super(*args)
+        else
+          super
+        end
+      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
+      true
     end
 
   end

data/lib/berater/lua_script.rb

@@ -1,4 +1,5 @@
 require 'digest'
+require 'redis'
 
 module Berater
   class LuaScript
@@ -6,11 +7,11 @@ module Berater
     attr_reader :source
 
     def initialize(source)
-      @source = source
+      @source = source.dup.freeze
     end
 
     def sha
-      @sha ||= Digest::SHA1.hexdigest(minify)
+      @sha ||= Digest::SHA1.hexdigest(minify).freeze
     end
 
     def eval(redis, *args)
@@ -44,7 +45,7 @@ module Berater
     def minify
       # trim comments (whole line and partial)
       # and whitespace (prefix and empty lines)
-      @minify ||= source.gsub(/^\s*--.*\n|\s*--.*|^\s*|^$\n/, '').chomp
+      @minify ||= source.gsub(/^\s*--.*\n|\s*--.*|^\s*|^$\n/, '').chomp.freeze
     end
 
   end

data/lib/berater/rate_limiter.rb

@@ -1,18 +1,19 @@
 module Berater
   class RateLimiter < Limiter
 
-    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
@@ -76,8 +77,8 @@ module Berater
 
       count, allowed = LUA_SCRIPT.eval(
         redis,
-        [ cache_key(key) ],
-        [ ts, capacity, @interval_msec, cost ]
+        [ cache_key ],
+        [ ts, capacity, @interval, cost ]
       )
 
       count = count.include?('.') ? count.to_f : count.to_i

data/lib/berater/rspec.rb

@@ -8,7 +8,6 @@ RSpec.configure do |config|
 
   config.after do
     Berater.expunge rescue nil
-    Berater.redis.script(:flush) rescue nil
     Berater.reset
   end
 end

data/lib/berater/static_limiter.rb

@@ -0,0 +1,49 @@
+module Berater
+  class StaticLimiter < Limiter
+
+    LUA_SCRIPT = Berater::LuaScript(<<~LUA
+      local key = KEYS[1]
+      local capacity = tonumber(ARGV[1])
+      local cost = tonumber(ARGV[2])
+
+      local count = redis.call('GET', key) or 0
+      local allowed = (count + cost) <= capacity
+
+      if allowed then
+        count = count + cost
+        redis.call('SET', key, count)
+      end
+
+      return { tostring(count), allowed }
+    LUA
+    )
+
+    protected def acquire_lock(capacity, cost)
+      if cost == 0
+        # utilization check
+        count = redis.get(cache_key) || "0"
+        allowed = true
+      else
+        count, allowed = LUA_SCRIPT.eval(
+          redis, [ cache_key ], [ capacity, cost ],
+        )
+      end
+
+      # Redis returns Floats as strings to maintain precision
+      count = count.include?('.') ? count.to_f : count.to_i
+
+      raise Overloaded unless allowed
+
+      release_fn = if cost > 0
+        proc { redis.incrbyfloat(cache_key, -cost) }
+      end
+
+      Lock.new(capacity, count, release_fn)
+    end
+
+    def to_s
+      "#<#{self.class}(#{key}: #{capacity})>"
+    end
+
+  end
+end

data/lib/berater/test_mode.rb

@@ -1,36 +1,47 @@
 require 'berater'
 
 module Berater
-  extend self
 
-  attr_reader :test_mode
+  module TestMode
+    attr_reader :test_mode
+
+    def test_mode=(mode)
+      unless [ nil, :pass, :fail ].include?(mode)
+        raise ArgumentError, "invalid mode: #{Berater.test_mode}"
+      end
 
-  def test_mode=(mode)
-    unless [ nil, :pass, :fail ].include?(mode)
-      raise ArgumentError, "invalid mode: #{Berater.test_mode}"
+      @test_mode = mode
     end
 
-    @test_mode = mode
+    def reset
+      super
+      @test_mode = nil
+    end
   end
 
-  module TestMode
-    def acquire_lock(*)
-      case Berater.test_mode
-      when :pass
-        Lock.new(Float::INFINITY, 0)
-      when :fail
-        raise Overloaded
-      else
-        super
+  class Limiter
+    module TestMode
+      def acquire_lock(*)
+        case Berater.test_mode
+        when :pass
+          Lock.new(Float::INFINITY, 0)
+        when :fail
+          raise Overloaded
+        else
+          super
+        end
       end
     end
   end
 
 end
 
+# prepend class methods
+Berater.singleton_class.prepend Berater::TestMode
+
 # stub each Limiter subclass
 ObjectSpace.each_object(Class).each do |klass|
   next unless klass < Berater::Limiter
 
-  klass.prepend(Berater::TestMode)
+  klass.prepend Berater::Limiter::TestMode
 end

data/lib/berater/unlimiter.rb

@@ -5,6 +5,10 @@ module Berater
       super(key, Float::INFINITY, **opts)
     end
 
+    def to_s
+      "#<#{self.class}>"
+    end
+
     protected
 
     def capacity=(*)

data/lib/berater/utils.rb

@@ -42,5 +42,9 @@ module Berater
       (res * 10**3).to_i
     end
 
+    def convenience_fn(klass, *args, **opts, &block)
+      limiter = klass.new(*args, **opts)
+      block ? limiter.limit(&block) : limiter
+    end
   end
 end

data/lib/berater/version.rb

@@ -1,3 +1,3 @@
 module Berater
-  VERSION = '0.7.1'
+  VERSION = "0.10.1"
 end

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,8 +17,14 @@ 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, **opts)
@@ -31,8 +39,10 @@ module Berater
       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, *args ].compact
@@ -46,20 +56,21 @@ module Berater
     end
   end
 
+  def reset
+    @redis = nil
+    limiters.clear
+    middleware.clear
+  end
 end
 
 # convenience method
-def Berater(key, capacity, **opts, &block)
-  limiter = Berater.new(key, capacity, **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'