berater 0.4.0 → 0.5.0

data/lib/berater/rate_limiter.rb

@@ -3,63 +3,28 @@ 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
-    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_usec, **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
-
-      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
+      @interval = interval
+      @interval_usec = Berater::Utils.to_usec(interval)
     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 interval_usec = tonumber(ARGV[3])
+      local cost = tonumber(ARGV[4])
       local count = 0
+      local allowed
+      local usec_per_drip = interval_usec / capacity
 
       -- timestamp of last update
       local last_ts = tonumber(redis.call('GET', ts_key))
@@ -72,61 +37,65 @@ module Berater
         count = math.max(0, count - drips)
       end
 
-      local allowed = count + 1 <= capacity
+      if cost == 0 then
+        -- just check limit, ie. for .overlimit?
+        allowed = count < capacity
+      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 * usec_per_drip) / 1000)
 
-        -- 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 expirations
+          redis.call('SET', key, count, 'PX', ttl)
+          redis.call('SET', ts_key, ts, 'PX', ttl)
+        end
       end
 
       return { count, allowed }
     LUA
+    )
 
-    def limit
-      usec_per_drip = (@interval_sec * 10**6) / @count
+    def limit(capacity: nil, cost: 1, &block)
+      capacity ||= @capacity
 
       # timestamp in microseconds
       ts = (Time.now.to_f * 10**6).to_i
 
-      count, allowed = redis.eval(
-        LUA_SCRIPT,
+      count, allowed = LUA_SCRIPT.eval(
+        redis,
         [ cache_key(key), cache_key("#{key}-ts") ],
-        [ ts, @count, usec_per_drip ]
+        [ ts, capacity, @interval_usec, cost ]
       )
 
       raise Overrated unless allowed
 
-      lock = Lock.new(self, "#{ts}-#{count}", count)
+      lock = Lock.new(self, ts, count)
+      yield_lock(lock, &block)
+    end
 
-      if block_given?
-        begin
-          yield lock
-        ensure
-          lock.release
-        end
-      else
-        lock
-      end
+    def overloaded?
+      limit(cost: 0) { false }
+    rescue Overrated
+      true
     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

@@ -8,5 +8,7 @@ RSpec.configure do |config|
 
   config.after do
     Berater.expunge rescue nil
+    Berater.redis.script(:flush) rescue nil
+    Berater.reset
   end
 end

data/lib/berater/rspec/matchers.rb

@@ -16,15 +16,17 @@ module BeraterMatchers
           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
+            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
+          obj.overloaded?
         end
-
-        false
       rescue @type
         true
       end

data/lib/berater/test_mode.rb

@@ -11,10 +11,15 @@ module Berater
     end
 
     @test_mode = mode
+
+    # overload class methods
+    unless Berater::Limiter.singleton_class.ancestors.include?(TestMode)
+      Berater::Limiter.singleton_class.prepend(TestMode)
+    end
   end
 
-  class Limiter
-    def self.new(*args, **opts)
+  module TestMode
+    def new(*args, **opts)
       return super unless Berater.test_mode
 
       # chose a stub class with desired behavior
@@ -28,13 +33,17 @@ module Berater
       # don't stub self
       return super if self < stub_klass
 
-      # swap out limit method with stub
+      # swap out limit and overloaded? methods 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)
+        instance.define_singleton_method(:limit) do |**opts, &block|
+          stub.limit(**opts, &block)
+        end
+
+        instance.define_singleton_method(:overloaded?) do
+          stub.overloaded?
         end
       end
     end

data/lib/berater/unlimiter.rb

@@ -2,21 +2,15 @@ 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)
+    def limit(**opts, &block)
+      yield_lock(Lock.new(self, Float::INFINITY, 0), &block)
+    end
 
-      if block_given?
-        begin
-          yield lock
-        ensure
-          lock.release
-        end
-      else
-        lock
-      end
+    def overloaded?
+      false
     end
 
   end

data/lib/berater/utils.rb

@@ -0,0 +1,46 @@
+module Berater
+  module Utils
+    extend self
+
+    refine Object do
+      def to_usec
+        Berater::Utils.to_usec(self)
+      end
+    end
+
+    def to_usec(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**6).to_i
+    end
+
+  end
+end

data/lib/berater/version.rb

@@ -1,3 +1,3 @@
 module Berater
-  VERSION = '0.4.0'
+  VERSION = '0.5.0'
 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,61 +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
-
-      it 'works with convinience' do
-        expect(Berater).to receive(:new).and_return(limiter)
-        expect {|b| Berater(:key, :concurrency, 1, &b) }.to yield_control
-      end
     end
+  end
 
-    context 'with DSL' do
-      it 'instatiates an Unlimiter' do
-        limiter = Berater.new(:key) { unlimited }
-        expect(limiter).to be_a Berater::Unlimiter
-        expect(limiter.key).to be :key
-      end
-
-      it 'instatiates an Inhibiter' do
-        limiter = Berater.new(:key) { inhibited }
-        expect(limiter).to be_a Berater::Inhibitor
-        expect(limiter.key).to be :key
-      end
-
-      it 'instatiates a RateLimiter' do
-        limiter = Berater.new(:key) { 1.per second }
-        expect(limiter).to be_a Berater::RateLimiter
-        expect(limiter.key).to be :key
-        expect(limiter.count).to be 1
-        expect(limiter.interval).to be :second
-      end
-
-      it 'instatiates a ConcurrencyLimiter' do
-        limiter = Berater.new(:key, timeout: 2) { 1.at_once }
-        expect(limiter).to be_a Berater::ConcurrencyLimiter
-        expect(limiter.key).to be :key
-        expect(limiter.capacity).to be 1
-        expect(limiter.timeout).to be 2
+  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
 
-      it 'does not accept mode/args and dsl block' do
-        expect {
-          Berater.new(:key, :rate) { 1.per second }
-        }.to raise_error(ArgumentError)
+      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
 
-        expect {
-          Berater.new(:key, :concurrency, 2) { 3.at_once }
-        }.to raise_error(ArgumentError)
-      end
-
-      it 'requires either mode or dsl block' do
-        expect {
-          Berater.new(:key)
-        }.to raise_error(ArgumentError)
+          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