berater 0.1.4 → 0.6.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: c8ca6fba7e38a014eed325df41dda7372b41ff62d9e78b61565f4d915d8b4677
-  data.tar.gz: 1fae864b93de87ea3dc70531f3f8697e80af526aea53d8ddfe2076eac9bbd61d
+  metadata.gz: 87d3bdb61a68a8b69dd7cb285dafb174a403acf19ff7586f90ed8b7cb50b7f76
+  data.tar.gz: fff71dd4d8ae28baf7c504efdc27df71f03219b32c1a959dddfd3d035845d058
 SHA512:
-  metadata.gz: 1b4b37134edece0657ea9be15119fe25b2208d554b9dbf56edd04f8e5acff77569022be025ec428b81ab1971171ae468e389a9568d7ad0163b0966e622a09f2e
-  data.tar.gz: bd927303eb1475e79a26161a184205008926f95c59a06eee43f453d5e464d15c205a91bffd36228ae28bc5d27f0fa32faebc7c5870ec5d1779bd6c57ecdc0fe2
+  metadata.gz: c9fb5c0c7ec100d2f50a3b6c6cfa449af052c1b78cff174b545e179c5fcc1324c2a5faafc90e7f7ec64ce200e47571d51b02eb56398d716c487abc7731c0cece
+  data.tar.gz: 148c566413d6593f229143e2871116b45947eae5b101d8bb2f59ce3bdfe89ed4ee1b7aaedfcaedbc80eef8a25a14fab83ad5e540172c41f31f27d7f296161235

data/lib/berater.rb

@@ -1,46 +1,40 @@
+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, :mode
+  attr_accessor :redis
 
   def configure
-    self.mode = :unlimited # default
-
     yield self
   end
 
-  def new(mode, *args, **opts)
-    klass = MODES[mode.to_sym]
-
-    unless klass
-      raise ArgumentError, "invalid mode: #{mode}"
-    end
-
-    klass.new(*args, **opts)
-  end
-
-  def register(mode, klass)
-    MODES[mode.to_sym] = klass
+  def reset
+    @redis = nil
   end
 
-  def mode=(mode)
-    unless MODES.include? mode.to_sym
-      raise ArgumentError, "invalid mode: #{mode}"
+  def new(key, capacity, interval = nil, **opts)
+    case capacity
+    when :unlimited, Float::INFINITY
+      Berater::Unlimiter
+    when :inhibited, 0
+      Berater::Inhibitor
+    else
+      if interval
+        Berater::RateLimiter
+      else
+        Berater::ConcurrencyLimiter
+      end
+    end.yield_self do |klass|
+      args = [ key, capacity, interval ].compact
+      klass.new(*args, **opts)
     end
-
-    @mode = mode.to_sym
-  end
-
-  def limit(*args, **opts, &block)
-    mode = opts.delete(:mode) { self.mode }
-    new(mode, *args, **opts).limit(&block)
   end
 
   def expunge
@@ -51,14 +45,18 @@ module Berater
 
 end
 
-# load and register limiters
-require 'berater/base_limiter'
+# 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
+end
+
+# load limiters
 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)

data/lib/berater/concurrency_limiter.rb

@@ -1,153 +1,111 @@
 module Berater
-  class ConcurrencyLimiter < BaseLimiter
+  class ConcurrencyLimiter < Limiter
 
     class Incapacitated < Overloaded; end
 
-    attr_reader :capacity, :timeout
+    attr_reader :timeout
 
-    def initialize(capacity, **opts)
-      super(**opts)
+    def initialize(key, capacity, **opts)
+      super(key, capacity, **opts)
 
-      self.capacity = capacity
       self.timeout = opts[:timeout] || 0
     end
 
-    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
-
-    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
-
+    private def timeout=(timeout)
       @timeout = timeout
+      timeout = 0 if timeout == Float::INFINITY
+      @timeout_msec = Berater::Utils.to_msec(timeout)
     end
 
-    class Lock
-      attr_reader :limiter, :id, :contention
-
-      def initialize(limiter, id, contention)
-        @limiter = limiter
-        @id = id
-        @contention = contention
-        @released = false
-        @locked_at = Time.now
-      end
-
-      def release
-        raise 'lock already released' if released?
-        raise 'lock expired' if expired?
-
-        @released = limiter.release(self)
-      end
-
-      def released?
-        @released
-      end
-
-      def expired?
-        limiter.timeout > 0 && @locked_at + limiter.timeout < Time.now
-      end
-    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 cost = tonumber(ARGV[4])
+      local lock_ids = {}
 
-      local exists
-      local count
-      local lock
-
-      -- check to see if key already exists
-      if ttl == 0 then
-        exists = redis.call('EXISTS', key)
-      else
-        -- and refresh TTL while we're at it
-        exists = redis.call('EXPIRE', key, ttl * 2)
+      -- purge stale hosts
+      if ttl > 0 then
+        redis.call('ZREMRANGEBYSCORE', key, '-inf', ts - ttl)
       end
 
-      if exists == 1 then
-        -- purge stale hosts
-        if ttl > 0 then
-          redis.call('ZREMRANGEBYSCORE', key, '-inf', ts - ttl)
-        end
-
-        -- check capacity (subtract one for next lock entry)
-        count = redis.call('ZCARD', key) - 1
+      -- check capacity
+      local count = redis.call('ZCARD', key)
 
+      if cost == 0 then
+        -- just check limit, ie. for .overlimit?
         if count < capacity then
-          -- yay, grab a lock
+          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 = {}
 
-          -- regenerate next lock entry, which has score inf
-          lock = unpack(redis.call('ZPOPMAX', key))
-          redis.call('ZADD', key, 'inf', (lock + 1) % 2^52)
+        for i = lock_id - cost + 1, lock_id do
+          table.insert(lock_ids, i)
 
-          count = count + 1
+          table.insert(locks, ts)
+          table.insert(locks, i)
         end
-      elseif capacity > 0 then
-        count = 1
-        lock = "1"
 
-        -- create structure to track locks and next id
-        redis.call('ZADD', key, 'inf', lock + 1)
+        redis.call('ZADD', key, unpack(locks))
+        count = count + cost
 
         if ttl > 0 then
-          redis.call('EXPIRE', key, ttl * 2)
+          redis.call('PEXPIRE', key, ttl)
         end
       end
 
-      if lock then
-        -- store lock and timestamp
-        redis.call('ZADD', key, ts, lock)
-      end
-
-      return { count, lock }
+      return { count, unpack(lock_ids) }
     LUA
+    )
+
+    def limit(capacity: nil, cost: 1, &block)
+      capacity ||= @capacity
 
-    def limit(**opts, &block)
-      unless opts.empty?
-        return self.class.new(
-          capacity,
-          **options.merge(opts)
-          # **options.merge(timeout: timeout).merge(opts)
-        ).limit(&block)
+      # since fractional cost is not supported, capacity behaves like int
+      capacity = capacity.to_i
+
+      unless cost.is_a?(Integer) && cost >= 0
+        raise ArgumentError, "invalid cost: #{cost}"
       end
 
-      count, lock_id = redis.eval(
-        LUA_SCRIPT,
-        [ key ],
-        [ capacity, Time.now.to_i, timeout ]
-      )
+      # timestamp in milliseconds
+      ts = (Time.now.to_f * 10**3).to_i
 
-      raise Incapacitated unless lock_id
+      count, *lock_ids = LUA_SCRIPT.eval(
+        redis,
+        [ cache_key(key), cache_key('lock_id') ],
+        [ capacity, ts, @timeout_msec, cost ]
+      )
 
-      lock = Lock.new(self, lock_id, count)
+      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 = 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
     end
 
-    def release(lock)
-      res = redis.zrem(key, lock.id)
-      res == true || res == 1 # depending on which version of Redis
+    def to_s
+      "#<#{self.class}(#{key}: #{capacity} at a time)>"
     end
 
   end

data/lib/berater/dsl.rb

@@ -0,0 +1,68 @@
+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
+      @keywords ||= Class.new do
+        # create a class where DSL keywords are methods
+        KEYWORDS.each do |keyword|
+          define_singleton_method(keyword) { keyword }
+        end
+      end
+
+      install
+      @keywords.class_eval &block
+    ensure
+      uninstall
+    end
+
+    private
+
+    KEYWORDS = [
+      :second, :minute, :hour,
+      :unlimited, :inhibited,
+    ].freeze
+
+    def install
+      Integer.class_eval do
+        def per(unit)
+          [ self, unit ]
+        end
+        alias every per
+
+        def at_once
+          [ self ]
+        end
+        alias concurrently at_once
+        alias at_a_time at_once
+      end
+    end
+
+    def uninstall
+      Integer.remove_method :per
+      Integer.remove_method :every
+
+      Integer.remove_method :at_once
+      Integer.remove_method :concurrently
+      Integer.remove_method :at_a_time
+    end
+  end
+end

data/lib/berater/inhibitor.rb

@@ -1,21 +1,20 @@
 module Berater
-  class Inhibitor < BaseLimiter
+  class Inhibitor < Limiter
 
     class Inhibited < Overloaded; end
 
-    def initialize(*args, **opts)
-      super(**opts)
+    def initialize(key = :inhibitor, *args, **opts)
+      super(key, 0, **opts)
     end
 
-    def limit(**opts, &block)
-      unless opts.empty?
-        return self.class.new(
-          **options.merge(opts)
-        ).limit(&block)
-      end
-
+    def limit(**opts)
       raise Inhibited
     end
 
+    def overloaded?
+      true
+    end
+    alias inhibited? overloaded?
+
   end
 end

data/lib/berater/limiter.rb

@@ -0,0 +1,80 @@
+module Berater
+  class Limiter
+
+    attr_reader :key, :capacity, :options
+
+    def redis
+      options[:redis] || Berater.redis
+    end
+
+    def limit
+      raise NotImplementedError
+    end
+
+    def overloaded?
+      raise NotImplementedError
+    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 NotImplementedError if self == Berater::Limiter
+
+      super
+    end
+
+    protected
+
+    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 cache_key(key)
+      "#{self.class}:#{key}"
+    end
+
+    def yield_lock(lock, &block)
+      if block_given?
+        begin
+          yield lock
+        ensure
+          lock.release
+        end
+      else
+        lock
+      end
+    end
+
+  end
+end