berater 0.4.0 → 0.5.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: f57aeaa9319ae38bdd7d2468443fc5c9ea9ba44cecf47bc5a558f18ecd0ab274
-  data.tar.gz: ff14ef90856911e9e18cbfbad6f9799727c999697c5e4d96f47b5fbf665c725c
+  metadata.gz: fa5d51d6eea9d4691e58e8dd7065a1d596a57abdac1227fa19bd22cd2f42c15d
+  data.tar.gz: 55ac15bf6f2f295c525fdea57b3af9b518935a12dd4b644e3c0c9c085429ca50
 SHA512:
-  metadata.gz: c9d3abeeb9d35c608c194eaf505c5af1350a82583dd12257e70b9d4f96bd3a6cd30db802f57c3bf77d371b5deb5e856e9238eadd3990768205fd37ea3e2e3703
-  data.tar.gz: c98078a3c7d3a3eb4e2d523577f9d6dd7f1bbae76c9ecbd8d65961cc7aa1bab6e1a4b8fe84cfb30ee60f4d0ce11dcb58e3b94a3566cf50ee60d1f8e21bab82b1
+  metadata.gz: 52df5781f9c37c7f1c9dde09e579c8da1f2d62f405835b771c6087cda8af046018e85a88629f7c5d3a2be9aafd2680e348d1b2f91bc63fe456e6e51ec5d0fd02
+  data.tar.gz: '0928cb6f2e8e07810dec96f5af302ff9f62d519b9f19ef6f28acb5621569f13edb6d784cbd7084b738df8dfc232ace26e3ce0c6c4a7f2ac3e4391c97789bafae'

data/lib/berater.rb

@@ -1,48 +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 = nil, *args, **opts, &block)
-    if mode.nil?
-      unless args.empty?
-        raise ArgumentError, '0 arguments expected with block'
-      end
-
-      unless block_given?
-        raise ArgumentError, 'expected either mode or block'
-      end
+  def reset
+    @redis = nil
+  end
 
-      mode, *args = DSL.eval(&block)
+  def new(key, capacity, interval = nil, **opts)
+    case capacity
+    when :unlimited, Float::INFINITY
+      Berater::Unlimiter
+    when :inhibited, 0
+      Berater::Inhibitor
     else
-      if block_given?
-        raise ArgumentError, 'expected either mode or block, not both'
+      if interval
+        Berater::RateLimiter
+      else
+        Berater::ConcurrencyLimiter
       end
+    end.yield_self do |klass|
+      args = [ key, capacity, interval ].compact
+      klass.new(*args, **opts)
     end
-
-    klass = MODES[mode.to_sym]
-
-    unless klass
-      raise ArgumentError, "invalid mode: #{mode}"
-    end
-
-    klass.new(key, *args, **opts)
-  end
-
-  def register(mode, klass)
-    MODES[mode.to_sym] = klass
   end
 
   def expunge
@@ -54,20 +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 limiters
-require 'berater/limiter'
 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)
-
-require 'berater/dsl'

data/lib/berater/concurrency_limiter.rb

@@ -3,42 +3,28 @@ module Berater
 
     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_usec = Berater::Utils.to_usec(timeout)
     end
 
-    LUA_SCRIPT = <<~LUA.gsub(/^\s*|\s*--.*/, '')
+    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,69 @@ 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(capacity: nil, cost: 1, &block)
+      capacity ||= @capacity
+      # cost is Integer >= 0
+
+      # timestamp in microseconds
+      ts = (Time.now.to_f * 10**6).to_i
 
-    def limit
-      count, lock_id = redis.eval(
-        LUA_SCRIPT,
+      count, *lock_ids = LUA_SCRIPT.eval(
+        redis,
         [ cache_key(key), cache_key('lock_id') ],
-        [ capacity, Time.now.to_i, timeout ]
+        [ capacity, ts, @timeout_usec, cost ]
       )
 
-      raise Incapacitated unless lock_id
+      raise Incapacitated if lock_ids.empty?
 
-      lock = Lock.new(self, lock_id, count, -> { release(lock_id) })
-
-      if block_given?
-        begin
-          yield lock
-        ensure
-          lock.release
-        end
+      if cost == 0
+        lock = Lock.new(self, nil, count)
       else
-        lock
+        lock = Lock.new(self, lock_ids[0], count, -> { release(lock_ids) })
       end
+
+      yield_lock(lock, &block)
+    end
+
+    def overloaded?
+      limit(cost: 0) { false }
+    rescue Overloaded
+      true
     end
+    alias incapacitated? overloaded?
 
-    private def release(lock_id)
-      res = redis.zrem(cache_key(key), lock_id)
-      res == true || res == 1 # depending on which version of Redis
+    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

data/lib/berater/dsl.rb

@@ -1,5 +1,23 @@
 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
@@ -18,13 +36,6 @@ module Berater
 
     private
 
-    def each &block
-      Berater::MODES.map do |mode, limiter|
-        next unless limiter.const_defined?(:DSL, false)
-        limiter.const_get(:DSL)
-      end.compact.each(&block)
-    end
-
     KEYWORDS = [
       :second, :minute, :hour,
       :unlimited, :inhibited,
@@ -33,12 +44,12 @@ module Berater
     def install
       Integer.class_eval do
         def per(unit)
-          [ :rate, self, unit ]
+          [ self, unit ]
         end
         alias every per
 
         def at_once
-          [ :concurrency, self ]
+          [ self ]
         end
         alias concurrently at_once
         alias at_a_time at_once

data/lib/berater/inhibitor.rb

@@ -4,12 +4,17 @@ module Berater
     class Inhibited < Overloaded; end
 
     def initialize(key = :inhibitor, *args, **opts)
-      super(key, **opts)
+      super(key, 0, **opts)
     end
 
-    def limit
+    def limit(**opts)
       raise Inhibited
     end
 
+    def overloaded?
+      true
+    end
+    alias inhibited? overloaded?
+
   end
 end

data/lib/berater/limiter.rb

@@ -1,7 +1,7 @@
 module Berater
   class Limiter
 
-    attr_reader :key, :options
+    attr_reader :key, :capacity, :options
 
     def redis
       options[:redis] || Berater.redis
@@ -11,20 +11,66 @@ module Berater
       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
 
-    def initialize(key, **opts)
+    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?(Integer) || capacity == Float::INFINITY
+        raise ArgumentError, "expected Integer, found #{capacity.class}"
+      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

data/lib/berater/lock.rb

@@ -13,24 +13,15 @@ module Berater
     end
 
     def locked?
-      @released_at.nil? && !expired?
-    end
-
-    def expired?
-      timeout ? @locked_at + timeout < Time.now : false
+      @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
 
-    def timeout
-      limiter.respond_to?(:timeout) ? limiter.timeout : nil
-    end
-
   end
 end

data/lib/berater/lua_script.rb

@@ -0,0 +1,55 @@
+require 'digest'
+
+module Berater
+  class LuaScript
+
+    attr_reader :source
+
+    def initialize(source)
+      @source = source
+    end
+
+    def sha
+      @sha ||= Digest::SHA1.hexdigest(minify)
+    end
+
+    def eval(redis, *args)
+      redis.evalsha(sha, *args)
+    rescue Redis::CommandError => e
+      raise unless e.message.include?('NOSCRIPT')
+
+      # fall back to regular eval, which will trigger
+      # script to be cached for next time
+      redis.eval(minify, *args)
+    end
+
+    def load(redis)
+      redis.script(:load, minify).tap do |sha|
+        unless sha == self.sha
+          raise "unexpected script SHA: expected #{self.sha}, got #{sha}"
+        end
+      end
+    end
+
+    def loaded?(redis)
+      redis.script(:exists, sha)
+    end
+
+    def to_s
+      source
+    end
+
+    private
+
+    def minify
+      # trim comments (whole line and partial)
+      # and whitespace (prefix and empty lines)
+      @minify ||= source.gsub(/^\s*--.*\n|\s*--.*|^\s*|^$\n/, '').chomp
+    end
+
+  end
+
+  def LuaScript(source)
+    LuaScript.new(source)
+  end
+end