berater 0.4.0 → 0.7.0

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

@@ -1,132 +1,104 @@
 module Berater
   class RateLimiter < Limiter
 
-    class Overrated < Overloaded; end
+    attr_accessor :interval
 
-    attr_accessor :count, :interval
-
-    def initialize(key, count, interval, **opts)
-      super(key, **opts)
-
-      self.count = count
+    def initialize(key, capacity, interval, **opts)
       self.interval = interval
-    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
-
-      @count = count
+      super(key, capacity, @interval_msec, **opts)
     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}"
-      end
+      @interval = interval
+      @interval_msec = Berater::Utils.to_msec(interval)
 
-      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
+      unless @interval_msec > 0
+        raise ArgumentError, 'interval must be > 0'
       end
     end
 
-    LUA_SCRIPT = <<~LUA.gsub(/^\s*|\s*--.*/, '')
+    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 usec_per_drip = tonumber(ARGV[3])
-      local count = 0
-
-      -- 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) / usec_per_drip)
-        count = math.max(0, count - drips)
+      local interval_msec = tonumber(ARGV[3])
+      local cost = tonumber(ARGV[4])
+
+      local allowed -- whether lock was acquired
+      local count -- capacity being utilized
+      local msec_per_drip = interval_msec / capacity
+      local state = redis.call('GET', key)
+
+      if state then
+        local last_ts -- timestamp of last update
+        count, last_ts = string.match(state, '([%d.]+);(%w+)')
+        count = tonumber(count)
+        last_ts = tonumber(last_ts, 16)
+
+        -- adjust for time passing, guarding against clock skew
+        if ts > last_ts then
+          local drips = math.floor((ts - last_ts) / msec_per_drip)
+          count = math.max(0, count - drips)
+        else
+          ts = last_ts
+        end
+      else
+        count = 0
       end
 
-      local allowed = count + 1 <= capacity
+      if cost == 0 then
+        -- just checking count
+        allowed = true
+      else
+        allowed = (count + cost) <= capacity
 
-      if allowed then
-        count = count + 1
+        if allowed then
+          count = count + cost
 
-        -- time for bucket to empty, in milliseconds
-        local ttl = math.ceil((count * usec_per_drip) / 1000)
+          -- time for bucket to empty, in milliseconds
+          local ttl = math.ceil(count * msec_per_drip)
+          ttl = ttl + 100 -- margin of error, for clock skew
 
-        -- update count and last_ts, with expirations
-        redis.call('SET', key, count, 'PX', ttl)
-        redis.call('SET', ts_key, ts, 'PX', ttl)
+          -- update count and last_ts, with expiration
+          state = string.format('%f;%X', count, ts)
+          redis.call('SET', key, state, 'PX', ttl)
+        end
       end
 
-      return { count, allowed }
+      return { tostring(count), allowed }
     LUA
+    )
 
-    def limit
-      usec_per_drip = (@interval_sec * 10**6) / @count
-
-      # timestamp in microseconds
-      ts = (Time.now.to_f * 10**6).to_i
+    protected def acquire_lock(capacity, cost)
+      # timestamp in milliseconds
+      ts = (Time.now.to_f * 10**3).to_i
 
-      count, allowed = redis.eval(
-        LUA_SCRIPT,
-        [ cache_key(key), cache_key("#{key}-ts") ],
-        [ ts, @count, usec_per_drip ]
+      count, allowed = LUA_SCRIPT.eval(
+        redis,
+        [ cache_key(key) ],
+        [ ts, capacity, @interval_msec, cost ]
       )
 
-      raise Overrated unless allowed
+      count = count.include?('.') ? count.to_f : count.to_i
 
-      lock = Lock.new(self, "#{ts}-#{count}", count)
+      raise Overloaded unless allowed
 
-      if block_given?
-        begin
-          yield lock
-        ensure
-          lock.release
-        end
-      else
-        lock
-      end
+      Lock.new(capacity, count)
     end
 
     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,61 @@
-module BeraterMatchers
-  class Overloaded
-    def initialize(type)
-      @type = type
-    end
-
-    def supports_block_expectations?
-      true
-    end
+module Berater
+  module Matchers
+    class Overloaded
+      def supports_block_expectations?
+        true
+      end
 
-    def matches?(obj)
-      begin
+      def matches?(obj)
         case obj
         when Proc
-          # eg. expect { ... }.to be_overrated
+          # eg. expect { ... }.to be_overloaded
           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
+            @limiter.utilization >= 1
+          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
+          @limiter.utilization >= 1
         end
-
-        false
-      rescue @type
+      rescue Berater::Overloaded
         true
       end
-    end
 
-    # def description
-    # it { expect { Berater.new(:inhibitor) }.not_to be_overrated }
+      def description
+        if @limiter
+          "be overloaded"
+        else
+          "raise #{Berater::Overloaded}"
+        end
+      end
 
-    def failure_message
-      "expected #{@type} to be raised"
-    end
+      def failure_message
+        if @limiter
+          "expected to be overloaded"
+        else
+          "expected #{Berater::Overloaded} 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 overloaded"
+        else
+          "did not expect #{Berater::Overloaded} to be raised"
+        end
+      end
     end
-  end
-
-  def be_overloaded
-    Overloaded.new(Berater::Overloaded)
-  end
 
-  def be_overrated
-    Overloaded.new(Berater::RateLimiter::Overrated)
-  end
-
-  def be_incapacitated
-    Overloaded.new(Berater::ConcurrencyLimiter::Incapacitated)
-  end
-
-  def be_inhibited
-    Overloaded.new(Berater::Inhibitor::Inhibited)
+    def be_overloaded
+      Overloaded.new
+    end
   end
 end

data/lib/berater/test_mode.rb

@@ -13,31 +13,24 @@ 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
+        raise Overloaded
+      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