berater 0.3.0 → 0.6.2

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

data/lib/berater/rate_limiter.rb

@@ -3,94 +3,95 @@ module Berater
 
     class Overrated < Overloaded; end
 
-    attr_accessor :count, :interval
+    attr_accessor :interval
 
-    def initialize(key, count, interval, **opts)
-      super(key, **opts)
-
-      self.count = count
+    def initialize(key, capacity, interval, **opts)
       self.interval = interval
+      super(key, capacity, @interval_msec, **opts)
     end
 
-    private def count=(count)
-      unless count.is_a? Integer
-        raise ArgumentError, "expected Integer, found #{count.class}"
-      end
-
-      raise ArgumentError, "count must be >= 0" unless count >= 0
+    private def interval=(interval)
+      @interval = interval
+      @interval_msec = Berater::Utils.to_msec(interval)
 
-      @count = count
+      unless @interval_msec > 0
+        raise ArgumentError, 'interval must be > 0'
+      end
     end
 
-    private def interval=(interval)
-      @interval = interval.dup
-
-      case @interval
-      when Integer
-        raise ArgumentError, "interval must be >= 0" unless @interval >= 0
-        @interval_sec = @interval
-      when String
-        @interval = @interval.to_sym
-      when Symbol
-      else
-        raise ArgumentError, "unexpected interval type: #{interval.class}"
+    LUA_SCRIPT = Berater::LuaScript(<<~LUA
+      local key = KEYS[1]
+      local ts_key = KEYS[2]
+      local ts = tonumber(ARGV[1])
+      local capacity = tonumber(ARGV[2])
+      local interval_msec = tonumber(ARGV[3])
+      local cost = tonumber(ARGV[4])
+      local count = 0
+      local allowed
+      local msec_per_drip = interval_msec / 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) / msec_per_drip)
+        count = math.max(0, count - drips)
       end
 
-      if @interval.is_a? Symbol
-        case @interval
-        when :sec, :second, :seconds
-          @interval = :second
-          @interval_sec = 1
-        when :min, :minute, :minutes
-          @interval = :minute
-          @interval_sec = 60
-        when :hour, :hours
-          @interval = :hour
-          @interval_sec = 60 * 60
-        else
-          raise ArgumentError, "unexpected interval value: #{interval}"
-        end
-      end
-    end
+      if cost == 0 then
+        -- just check limit, ie. for .overlimit?
+        allowed = count < capacity
+      else
+        allowed = (count + cost) <= capacity
 
-    def limit
-      ts = Time.now.to_i
+        if allowed then
+          count = count + cost
 
-      # bucket into time slot
-      rkey = "%s:%d" % [ cache_key(key), ts - ts % @interval_sec ]
+          -- time for bucket to empty, in milliseconds
+          local ttl = math.ceil(count * msec_per_drip)
 
-      count, _ = redis.multi do
-        redis.incr rkey
-        redis.expire rkey, @interval_sec * 2
+          -- update count and last_ts, with expirations
+          redis.call('SET', key, count, 'PX', ttl)
+          redis.call('SET', ts_key, ts, 'PX', ttl)
+        end
       end
 
-      raise Overrated if count > @count
+      return { count, allowed }
+    LUA
+    )
 
-      lock = Lock.new(self, count, count)
+    protected def acquire_lock(capacity, cost)
+      # timestamp in milliseconds
+      ts = (Time.now.to_f * 10**3).to_i
 
-      if block_given?
-        begin
-          yield lock
-        ensure
-          lock.release
-        end
-      else
-        lock
-      end
+      count, allowed = LUA_SCRIPT.eval(
+        redis,
+        [ cache_key(key), cache_key("#{key}-ts") ],
+        [ ts, capacity, @interval_msec, cost ]
+      )
+
+      raise Overrated unless allowed
+
+      Lock.new(capacity, count)
     end
 
+    alias overrated? overloaded?
+
     def to_s
-      msg = if @interval.is_a? Integer
-        if @interval == 1
+      msg = if interval.is_a? Numeric
+        if interval == 1
           "every second"
         else
-          "every #{@interval} seconds"
+          "every #{interval} seconds"
         end
       else
-        "per #{@interval}"
+        "per #{interval}"
       end
 
-      "#<#{self.class}(#{key}: #{count} #{msg})>"
+      "#<#{self.class}(#{key}: #{capacity} #{msg})>"
     end
 
   end

data/lib/berater/rspec.rb

@@ -4,9 +4,11 @@ require 'berater/test_mode'
 require 'rspec'
 
 RSpec.configure do |config|
-  config.include(BeraterMatchers)
+  config.include(Berater::Matchers)
 
   config.after do
     Berater.expunge rescue nil
+    Berater.redis.script(:flush) rescue nil
+    Berater.reset
   end
 end

data/lib/berater/rspec/matchers.rb

@@ -1,60 +1,81 @@
-module BeraterMatchers
-  class Overloaded
-    def initialize(type)
-      @type = type
-    end
+module Berater
+  module Matchers
+    class Overloaded
+      def initialize(type)
+        @type = type
+      end
 
-    def supports_block_expectations?
-      true
-    end
+      def supports_block_expectations?
+        true
+      end
 
-    def matches?(obj)
-      begin
+      def matches?(obj)
         case obj
         when Proc
           # eg. expect { ... }.to be_overrated
           res = obj.call
 
           if res.is_a? Berater::Limiter
-            # eg. expect { Berater.new(...) }.to be_overrated
-            res.limit {}
+            # eg. expect { Berater.new(...) }.to be_overloaded
+            @limiter = res
+            res.overloaded?
+          else
+            # eg. expect { Berater(...)  }.to be_overloaded
+            # eg. expect { limiter.limit }.to be_overloaded
+            false
           end
         when Berater::Limiter
-          # eg. expect(Berater.new(...)).to be_overrated
-          obj.limit {}
+          # eg. expect(Berater.new(...)).to be_overloaded
+          @limiter = obj
+          obj.overloaded?
         end
-
-        false
       rescue @type
         true
       end
-    end
 
-    # def description
-    # it { expect { Berater.new(:inhibitor) }.not_to be_overrated }
+      def description
+        if @limiter
+          "be #{verb}"
+        else
+          "raise #{@type}"
+        end
+      end
 
-    def failure_message
-      "expected #{@type} to be raised"
-    end
+      def failure_message
+        if @limiter
+          "expected to be #{verb}"
+        else
+          "expected #{@type} to be raised"
+        end
+      end
 
-    def failure_message_when_negated
-      "did not expect #{@type} to be raised"
+      def failure_message_when_negated
+        if @limiter
+          "expected not to be #{verb}"
+        else
+          "did not expect #{@type} to be raised"
+        end
+      end
+
+      private def verb
+        @type.to_s.split('::')[-1].downcase
+      end
     end
-  end
 
-  def be_overloaded
-    Overloaded.new(Berater::Overloaded)
-  end
+    def be_overloaded
+      Overloaded.new(Berater::Overloaded)
+    end
 
-  def be_overrated
-    Overloaded.new(Berater::RateLimiter::Overrated)
-  end
+    def be_overrated
+      Overloaded.new(Berater::RateLimiter::Overrated)
+    end
 
-  def be_incapacitated
-    Overloaded.new(Berater::ConcurrencyLimiter::Incapacitated)
-  end
+    def be_incapacitated
+      Overloaded.new(Berater::ConcurrencyLimiter::Incapacitated)
+    end
 
-  def be_inhibited
-    Overloaded.new(Berater::Inhibitor::Inhibited)
+    def be_inhibited
+      Overloaded.new(Berater::Inhibitor::Inhibited)
+    end
   end
 end

data/lib/berater/test_mode.rb

@@ -13,31 +13,31 @@ module Berater
     @test_mode = mode
   end
 
-  class Limiter
-    def self.new(*args, **opts)
-      return super unless Berater.test_mode
-
-      # chose a stub class with desired behavior
-      stub_klass = case Berater.test_mode
+  module TestMode
+    def acquire_lock(*)
+      case Berater.test_mode
       when :pass
-        Berater::Unlimiter
+        Lock.new(Float::INFINITY, 0)
       when :fail
-        Berater::Inhibitor
-      end
-
-      # don't stub self
-      return super if self < stub_klass
-
-      # swap out limit method with stub
-      super.tap do |instance|
-        stub = stub_klass.allocate
-        stub.send(:initialize, *args, **opts)
-
-        instance.define_singleton_method(:limit) do |&block|
-          stub.limit(&block)
-        end
+        # find class specific Overloaded error
+        e = self.class.constants.map do |name|
+          self.class.const_get(name)
+        end.find do |const|
+          const.is_a?(Class) && const < Berater::Overloaded
+        end || Berater::Overloaded
+
+        raise e
+      else
+        super
       end
     end
   end
 
 end
+
+# stub each Limiter subclass
+ObjectSpace.each_object(Class).each do |klass|
+  next unless klass < Berater::Limiter
+
+  klass.prepend(Berater::TestMode)
+end

data/lib/berater/unlimiter.rb

@@ -2,21 +2,17 @@ module Berater
   class Unlimiter < Limiter
 
     def initialize(key = :unlimiter, *args, **opts)
-      super(key, **opts)
+      super(key, Float::INFINITY, **opts)
     end
 
-    def limit
-      lock = Lock.new(self, 0, 0)
+    protected
 
-      if block_given?
-        begin
-          yield lock
-        ensure
-          lock.release
-        end
-      else
-        lock
-      end
+    def capacity=(*)
+      @capacity = Float::INFINITY
+    end
+
+    def acquire_lock(*)
+      Lock.new(Float::INFINITY, 0)
     end
 
   end

data/lib/berater/utils.rb

@@ -0,0 +1,46 @@
+module Berater
+  module Utils
+    extend self
+
+    refine Object do
+      def to_msec
+        Berater::Utils.to_msec(self)
+      end
+    end
+
+    def to_msec(val)
+      res = val
+
+      if val.is_a? String
+        # naively attempt casting, otherwise maybe it's a keyword
+        res = Float(val) rescue val.to_sym
+      end
+
+      if res.is_a? Symbol
+        case res
+        when :sec, :second, :seconds
+          res = 1
+        when :min, :minute, :minutes
+          res = 60
+        when :hour, :hours
+          res = 60 * 60
+        end
+      end
+
+      unless res.is_a? Numeric
+        raise ArgumentError, "unexpected value: #{val}"
+      end
+
+      if res < 0
+        raise ArgumentError, "expected value >= 0, found: #{val}"
+      end
+
+      if res == Float::INFINITY
+        raise ArgumentError, "infinite values not allowed"
+      end
+
+      (res * 10**3).to_i
+    end
+
+  end
+end