berater 0.2.0 → 0.6.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: dec8c5d6d428f795d1489dace166451e2a3a42592d00d71155ed5cd0d5aaa909
-  data.tar.gz: 6956b6c2804b6d616074e239439d4f2a8e551e900db7801226948b9eb764d0fc
+  metadata.gz: 5c5cc50ba02c18bf9ef9f6b9c1dd6ccb8f8af0297df95757ac13406dc4539c40
+  data.tar.gz: f80341dccff30ca3f344f3bdf21de44b1ef1715bfaf4140933a91252ddb7231e
 SHA512:
-  metadata.gz: 327cd2f3a7430c2b81bd1988abd636bfc7d97048c762d3807555e9c1196153570dbe3045fae813852a907ec00b89b8d2c36e9269ed0562419c8d98dd6c1d2b26
-  data.tar.gz: ea1cb39e9201091e4a1357e888ac36f630c326a271e3cfcafa13298b8f54424f776f4dd488e046e2178905d0df9fb5c3b1f0d454e331823cafe5dd338fe9ea6c
+  metadata.gz: ada3fffc841e71498f652dc990eed482bccc942b86280fc5ee736d49fdf99d8ab2f921b5b52d96f3588c268c34343669c00e0be8822c462a3e6f9dcf954a4005
+  data.tar.gz: 0363abef69cfdad8140d89373b5088cfa8a6259f16716011c82ceb1e2e9e948cdb4d44ddd36aaf1d05b121dc7d275d0ee34571a0a3ee9b964f4286375ec356cf

data/lib/berater.rb

@@ -1,32 +1,40 @@
-require 'berater/version'
+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
 
   def configure
     yield self
   end
 
-  def new(key, mode, *args, **opts)
-    klass = MODES[mode.to_sym]
-
-    unless klass
-      raise ArgumentError, "invalid mode: #{mode}"
-    end
-
-    klass.new(key, *args, **opts)
+  def reset
+    @redis = nil
   end
 
-  def register(mode, klass)
-    MODES[mode.to_sym] = klass
+  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
   end
 
   def expunge
@@ -38,18 +46,17 @@ module Berater
 end
 
 # convenience method
-def Berater(key, mode, *args, **opts, &block)
-  Berater.new(key, mode, *args, **opts).limit(&block)
+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 and register limiters
-require 'berater/base_limiter'
+# 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,44 +1,30 @@
 module Berater
-  class ConcurrencyLimiter < BaseLimiter
+  class ConcurrencyLimiter < Limiter
 
     class Incapacitated < Overloaded; end
 
-    attr_reader :capacity, :timeout
+    attr_reader :timeout
 
     def initialize(key, capacity, **opts)
-      super(key, **opts)
+      super(key, capacity, **opts)
 
-      self.capacity = capacity
       self.timeout = opts[:timeout] || 0
     end
 
-    private 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
-
     private 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
-
       @timeout = timeout
+      timeout = 0 if timeout == Float::INFINITY
+      @timeout_msec = Berater::Utils.to_msec(timeout)
     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 lock
+      local cost = tonumber(ARGV[4])
+      local lock_ids = {}
 
       -- purge stale hosts
       if ttl > 0 then
@@ -48,41 +34,67 @@ module Berater
       -- check capacity
       local count = redis.call('ZCARD', key)
 
-      if count < capacity then
-        -- grab a lock
-        lock = redis.call('INCR', lock_key)
-        redis.call('ZADD', key, ts, lock)
-        count = count + 1
+      if cost == 0 then
+        -- just check limit, ie. for .overlimit?
+        if count < capacity then
+          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 = {}
+
+        for i = lock_id - cost + 1, lock_id do
+          table.insert(lock_ids, i)
+
+          table.insert(locks, ts)
+          table.insert(locks, i)
+        end
+
+        redis.call('ZADD', key, unpack(locks))
+        count = count + cost
+
+        if ttl > 0 then
+          redis.call('PEXPIRE', key, ttl)
+        end
       end
 
-      return { count, lock }
+      return { count, unpack(lock_ids) }
     LUA
+    )
 
-    def limit
-      count, lock_id = redis.eval(
-        LUA_SCRIPT,
+    protected def acquire_lock(capacity, cost)
+      # fractional cost is not supported, but make it work
+      capacity = capacity.to_i
+      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, Time.now.to_i, timeout ]
+        [ capacity, ts, @timeout_msec, cost ]
       )
 
-      raise Incapacitated unless lock_id
-
-      lock = Lock.new(self, lock_id, count, -> { release(lock_id) })
+      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.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
 
-    private def release(lock_id)
-      res = redis.zrem(cache_key(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,13 +1,15 @@
 module Berater
-  class Inhibitor < BaseLimiter
+  class Inhibitor < Limiter
 
     class Inhibited < Overloaded; end
 
     def initialize(key = :inhibitor, *args, **opts)
-      super(key, **opts)
+      super(key, 0, **opts)
     end
 
-    def limit
+    alias inhibited? overloaded?
+
+    protected def acquire_lock(*)
       raise Inhibited
     end
 

data/lib/berater/limiter.rb

@@ -0,0 +1,94 @@
+module Berater
+  class Limiter
+
+    attr_reader :key, :capacity, :options
+
+    def redis
+      options[:redis] || Berater.redis
+    end
+
+    def limit(capacity: nil, cost: 1, &block)
+      capacity ||= @capacity
+
+      unless capacity.is_a?(Numeric)
+        raise ArgumentError, "invalid capacity: #{capacity}"
+      end
+
+      unless cost.is_a?(Numeric) && cost >= 0
+        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?
+      limit(cost: 0) { false }
+    rescue Overloaded
+      true
+    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 acquire_lock(capacity, cost)
+      raise NotImplementedError
+    end
+
+    def cache_key(key)
+      "#{self.class}:#{key}"
+    end
+
+  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
@@ -13,24 +12,15 @@ module Berater
     end
 
     def locked?
-      @released_at.nil? && !expired?
-    end
-
-    def expired?
-      timeout > 0 && @locked_at + timeout < Time.now
+      @released_at.nil?
     end
 
     def release
-      raise 'lock expired' if expired?
       raise 'lock already released' unless locked?
 
       @released_at = Time.now
       @release_fn ? @release_fn.call : true
     end
 
-    private def timeout
-      limiter.respond_to?(:timeout) ? limiter.timeout : 0
-    end
-
   end
 end