berater 0.2.0 → 0.6.1

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,80 +1,97 @@
 module Berater
-  class RateLimiter < BaseLimiter
+  class RateLimiter < Limiter
 
     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
-      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 = 1
-        when :min, :minute, :minutes
-          @interval = 60
-        when :hour, :hours
-          @interval = 60 * 60
-        else
-          raise ArgumentError, "unexpected interval value: #{interval}"
+      if cost == 0 then
+        -- just check limit, ie. for .overlimit?
+        allowed = count < capacity
+      else
+        allowed = (count + cost) <= capacity
+
+        if allowed then
+          count = count + cost
+
+          -- time for bucket to empty, in milliseconds
+          local ttl = math.ceil(count * msec_per_drip)
+
+          -- update count and last_ts, with expirations
+          redis.call('SET', key, count, 'PX', ttl)
+          redis.call('SET', ts_key, ts, 'PX', ttl)
         end
       end
 
-      @interval
-    end
+      return { count, allowed }
+    LUA
+    )
 
-    def limit
-      ts = Time.now.to_i
+    protected def acquire_lock(capacity, cost)
+      # timestamp in milliseconds
+      ts = (Time.now.to_f * 10**3).to_i
 
-      # bucket into time slot
-      rkey = "%s:%d" % [ cache_key(key), ts - ts % @interval ]
+      count, allowed = LUA_SCRIPT.eval(
+        redis,
+        [ cache_key(key), cache_key("#{key}-ts") ],
+        [ ts, capacity, @interval_msec, cost ]
+      )
 
-      count, _ = redis.multi do
-        redis.incr rkey
-        redis.expire rkey, @interval * 2
-      end
+      raise Overrated unless allowed
 
-      raise Overrated if count > @count
+      Lock.new(capacity, count)
+    end
 
-      lock = Lock.new(self, count, count)
+    alias overrated? overloaded?
 
-      if block_given?
-        begin
-          yield lock
-        ensure
-          lock.release
+    def to_s
+      msg = if interval.is_a? Numeric
+        if interval == 1
+          "every second"
+        else
+          "every #{interval} seconds"
         end
       else
-        lock
+        "per #{interval}"
       end
+
+      "#<#{self.class}(#{key}: #{capacity} #{msg})>"
     end
 
   end

data/lib/berater/rspec.rb

@@ -0,0 +1,14 @@
+require 'berater'
+require 'berater/rspec/matchers'
+require 'berater/test_mode'
+require 'rspec'
+
+RSpec.configure do |config|
+  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

@@ -0,0 +1,81 @@
+module Berater
+  module Matchers
+    class Overloaded
+      def initialize(type)
+        @type = type
+      end
+
+      def supports_block_expectations?
+        true
+      end
+
+      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_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_overloaded
+          @limiter = obj
+          obj.overloaded?
+        end
+      rescue @type
+        true
+      end
+
+      def description
+        if @limiter
+          "be #{verb}"
+        else
+          "raise #{@type}"
+        end
+      end
+
+      def failure_message
+        if @limiter
+          "expected to be #{verb}"
+        else
+          "expected #{@type} to be raised"
+        end
+      end
+
+      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
+
+    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)
+    end
+  end
+end

data/lib/berater/test_mode.rb

@@ -0,0 +1,43 @@
+require 'berater'
+
+module Berater
+  extend self
+
+  attr_reader :test_mode
+
+  def test_mode=(mode)
+    unless [ nil, :pass, :fail ].include?(mode)
+      raise ArgumentError, "invalid mode: #{Berater.test_mode}"
+    end
+
+    @test_mode = mode
+  end
+
+  module TestMode
+    def acquire_lock(*)
+      case Berater.test_mode
+      when :pass
+        Lock.new(Float::INFINITY, 0)
+      when :fail
+        # find class specific Overloaded error
+        e = self.class.constants.map do |name|
+          self.class.const_get(name)
+        end.find do |const|
+          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

@@ -1,23 +1,18 @@
 module Berater
-  class Unlimiter < BaseLimiter
+  class Unlimiter < Limiter
 
     def initialize(key = :unlimiter, *args, **opts)
-      super(key, **opts)
+      super(key, Float::INFINITY, **opts)
     end
 
-    def limit
-      count = redis.incr(cache_key('count'))
-      lock = Lock.new(self, count, count)
+    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

data/lib/berater/version.rb

@@ -1,3 +1,3 @@
 module Berater
-  VERSION = '0.2.0'
+  VERSION = '0.6.1'
 end

data/spec/berater_spec.rb

@@ -7,12 +7,12 @@ describe Berater do
   it { is_expected.to respond_to :configure }
 
   describe '.configure' do
-    it 'can be set via configure' do
+    it 'is used with a block' do
       Berater.configure do |c|
         c.redis = :redis
       end
 
-      expect(Berater.redis).to eq :redis
+      expect(Berater.redis).to be :redis
     end
   end
 
@@ -26,7 +26,7 @@ describe Berater do
 
   describe '.new' do
     context 'unlimited mode' do
-      let(:limiter) { Berater.new(:key, :unlimited) }
+      let(:limiter) { Berater.new(:key, Float::INFINITY) }
 
       it 'instantiates an Unlimiter' do
         expect(limiter).to be_a Berater::Unlimiter
@@ -39,18 +39,13 @@ describe Berater do
 
       it 'accepts options' do
         redis = double('Redis')
-        limiter = Berater.new(:key, :unlimited, redis: redis)
+        limiter = Berater.new(:key, Float::INFINITY, redis: redis)
         expect(limiter.redis).to be redis
       end
-
-      it 'works with convinience' do
-        expect(Berater).to receive(:new).and_return(limiter)
-        expect {|b| Berater(:key, :unlimited, &b) }.to yield_control
-      end
     end
 
     context 'inhibited mode' do
-      let(:limiter) { Berater.new(:key, :inhibited) }
+      let(:limiter) { Berater.new(:key, 0) }
 
       it 'instantiates an Inhibitor' do
         expect(limiter).to be_a Berater::Inhibitor
@@ -63,18 +58,13 @@ describe Berater do
 
       it 'accepts options' do
         redis = double('Redis')
-        limiter = Berater.new(:key, :inhibited, redis: redis)
+        limiter = Berater.new(:key, 0, redis: redis)
         expect(limiter.redis).to be redis
       end
-
-      it 'works with convinience' do
-        expect(Berater).to receive(:new).and_return(limiter)
-        expect { Berater(:key, :inhibited) }.to be_inhibited
-      end
     end
 
     context 'rate mode' do
-      let(:limiter) { Berater.new(:key, :rate, 1, :second) }
+      let(:limiter) { Berater.new(:key, 1, :second) }
 
       it 'instantiates a RateLimiter' do
         expect(limiter).to be_a Berater::RateLimiter
@@ -87,18 +77,13 @@ describe Berater do
 
       it 'accepts options' do
         redis = double('Redis')
-        limiter = Berater.new(:key, :rate, 1, :second, redis: redis)
+        limiter = Berater.new(:key, 1, :second, redis: redis)
         expect(limiter.redis).to be redis
       end
-
-      it 'works with convinience' do
-        expect(Berater).to receive(:new).and_return(limiter)
-        expect {|b| Berater(:key, :rate, 1, :second, &b) }.to yield_control
-      end
     end
 
     context 'concurrency mode' do
-      let(:limiter) { Berater.new(:key, :concurrency, 1) }
+      let(:limiter) { Berater.new(:key, 1) }
 
       it 'instantiates a ConcurrencyLimiter' do
         expect(limiter).to be_a Berater::ConcurrencyLimiter
@@ -111,15 +96,39 @@ describe Berater do
 
       it 'accepts options' do
         redis = double('Redis')
-        limiter = Berater.new(:key, :concurrency, 1, redis: redis)
+        limiter = Berater.new(:key, 1, redis: redis)
         expect(limiter.redis).to be redis
       end
+    end
+  end
+
+  describe 'Berater() - convenience method' do
+    RSpec.shared_examples 'test convenience' do |klass, *args|
+      it 'creates a limiter' do
+        limiter = Berater(:key, *args)
+        expect(limiter).to be_a klass
+      end
+
+      context 'with a block' do
+        it 'creates a limiter and calls limit' do
+          limiter = Berater(:key, *args)
+          expect(klass).to receive(:new).and_return(limiter)
+          expect(limiter).to receive(:limit).and_call_original
 
-      it 'works with convinience' do
-        expect(Berater).to receive(:new).and_return(limiter)
-        expect {|b| Berater(:key, :concurrency, 1, &b) }.to yield_control
+          begin
+            res = Berater(:key, *args) { true }
+            expect(res).to be true
+          rescue Berater::Overloaded
+            expect(klass).to be Berater::Inhibitor
+          end
+        end
       end
     end
+
+    include_examples 'test convenience', Berater::Unlimiter, Float::INFINITY
+    include_examples 'test convenience', Berater::Inhibitor, 0
+    include_examples 'test convenience', Berater::RateLimiter, 1, :second
+    include_examples 'test convenience', Berater::ConcurrencyLimiter, 1
   end
 
 end