berater 0.6.0 → 0.8.0

data/lib/berater/rspec.rb

@@ -1,14 +1,13 @@
 require 'berater'
 require 'berater/rspec/matchers'
 require 'berater/test_mode'
-require 'rspec'
+require 'rspec/core'
 
 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

@@ -1,83 +1,61 @@
 module Berater
   module Matchers
     class Overloaded
-      def initialize(type)
-        @type = type
-      end
-
       def supports_block_expectations?
         true
       end
 
       def matches?(obj)
-        begin
-          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?
+        case obj
+        when Proc
+          # eg. expect { ... }.to be_overloaded
+          res = obj.call
+
+          if res.is_a? Berater::Limiter
+            # 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
-        rescue @type
-          true
+        when Berater::Limiter
+          # eg. expect(Berater.new(...)).to be_overloaded
+          @limiter = obj
+          @limiter.utilization >= 1
         end
+      rescue Berater::Overloaded
+        true
       end
 
       def description
         if @limiter
-          "be #{verb}"
+          "be overloaded"
         else
-          "raise #{@type}"
+          "raise #{Berater::Overloaded}"
         end
       end
 
       def failure_message
         if @limiter
-          "expected to be #{verb}"
+          "expected to be overloaded"
         else
-          "expected #{@type} to be raised"
+          "expected #{Berater::Overloaded} to be raised"
         end
       end
 
       def failure_message_when_negated
         if @limiter
-          "expected not to be #{verb}"
+          "expected not to be overloaded"
         else
-          "did not expect #{@type} to be raised"
+          "did not expect #{Berater::Overloaded} 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)
+      Overloaded.new
     end
   end
 end

data/lib/berater/static_limiter.rb

@@ -0,0 +1,49 @@
+module Berater
+  class StaticLimiter < Limiter
+
+    LUA_SCRIPT = Berater::LuaScript(<<~LUA
+      local key = KEYS[1]
+      local capacity = tonumber(ARGV[1])
+      local cost = tonumber(ARGV[2])
+
+      local count = redis.call('GET', key) or 0
+      local allowed = (count + cost) <= capacity
+
+      if allowed then
+        count = count + cost
+        redis.call('SET', key, count)
+      end
+
+      return { tostring(count), allowed }
+    LUA
+    )
+
+    protected def acquire_lock(capacity, cost)
+      if cost == 0
+        # utilization check
+        count = redis.get(cache_key) || "0"
+        allowed = true
+      else
+        count, allowed = LUA_SCRIPT.eval(
+          redis, [ cache_key ], [ capacity, cost ],
+        )
+      end
+
+      # Redis returns Floats as strings to maintain precision
+      count = count.include?('.') ? count.to_f : count.to_i
+
+      raise Overloaded unless allowed
+
+      release_fn = if cost > 0
+        proc { redis.incrbyfloat(cache_key, -cost) }
+      end
+
+      Lock.new(capacity, count, release_fn)
+    end
+
+    def to_s
+      "#<#{self.class}(#{key}: #{capacity})>"
+    end
+
+  end
+end

data/lib/berater/test_mode.rb

@@ -1,52 +1,47 @@
 require 'berater'
 
 module Berater
-  extend self
 
-  attr_reader :test_mode
+  module TestMode
+    attr_reader :test_mode
 
-  def test_mode=(mode)
-    unless [ nil, :pass, :fail ].include?(mode)
-      raise ArgumentError, "invalid mode: #{Berater.test_mode}"
-    end
+    def test_mode=(mode)
+      unless [ nil, :pass, :fail ].include?(mode)
+        raise ArgumentError, "invalid mode: #{Berater.test_mode}"
+      end
 
-    @test_mode = mode
+      @test_mode = mode
+    end
 
-    # overload class methods
-    unless Berater::Limiter.singleton_class.ancestors.include?(TestMode)
-      Berater::Limiter.singleton_class.prepend(TestMode)
+    def reset
+      super
+      @test_mode = nil
     end
   end
 
-  module TestMode
-    def new(*args, **opts)
-      return super unless Berater.test_mode
-
-      # chose a stub class with desired behavior
-      stub_klass = case Berater.test_mode
-      when :pass
-        Berater::Unlimiter
-      when :fail
-        Berater::Inhibitor
-      end
-
-      # don't stub self
-      return super if self < stub_klass
-
-      # 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 |**opts, &block|
-          stub.limit(**opts, &block)
-        end
-
-        instance.define_singleton_method(:overloaded?) do
-          stub.overloaded?
+  class Limiter
+    module TestMode
+      def acquire_lock(*)
+        case Berater.test_mode
+        when :pass
+          Lock.new(Float::INFINITY, 0)
+        when :fail
+          raise Overloaded
+        else
+          super
         end
       end
     end
   end
 
 end
+
+# prepend class methods
+Berater.singleton_class.prepend Berater::TestMode
+
+# stub each Limiter subclass
+ObjectSpace.each_object(Class).each do |klass|
+  next unless klass < Berater::Limiter
+
+  klass.prepend Berater::Limiter::TestMode
+end

data/lib/berater/unlimiter.rb

@@ -5,17 +5,19 @@ module Berater
       super(key, Float::INFINITY, **opts)
     end
 
-    def limit(**opts, &block)
-      yield_lock(Lock.new(Float::INFINITY, 0), &block)
+    def to_s
+      "#<#{self.class}>"
     end
 
-    def overloaded?
-      false
-    end
+    protected
 
-    protected def capacity=(*)
+    def capacity=(*)
       @capacity = Float::INFINITY
     end
 
+    def acquire_lock(*)
+      Lock.new(Float::INFINITY, 0)
+    end
+
   end
 end

data/lib/berater/utils.rb

@@ -42,5 +42,14 @@ module Berater
       (res * 10**3).to_i
     end
 
+    def convenience_fn(klass, *args, **opts, &block)
+      limiter = klass.new(*args, **opts)
+      if block_given?
+        limiter.limit(&block)
+      else
+        limiter
+      end
+    end
+
   end
 end

data/lib/berater/version.rb

@@ -1,3 +1,3 @@
 module Berater
-  VERSION = '0.6.0'
+  VERSION = '0.8.0'
 end

data/spec/berater_spec.rb

@@ -24,50 +24,15 @@ describe Berater do
     end
   end
 
-  describe '.new' do
-    context 'unlimited mode' do
-      let(:limiter) { Berater.new(:key, Float::INFINITY) }
+  shared_examples 'a Berater' do |klass, capacity, **opts|
+    describe '.new' do
+      let(:limiter) { Berater.new(:key, capacity, **opts) }
 
-      it 'instantiates an Unlimiter' do
-        expect(limiter).to be_a Berater::Unlimiter
-        expect(limiter.key).to be :key
-      end
-
-      it 'inherits redis' do
-        expect(limiter.redis).to be Berater.redis
-      end
-
-      it 'accepts options' do
-        redis = double('Redis')
-        limiter = Berater.new(:key, Float::INFINITY, redis: redis)
-        expect(limiter.redis).to be redis
-      end
-    end
-
-    context 'inhibited mode' do
-      let(:limiter) { Berater.new(:key, 0) }
-
-      it 'instantiates an Inhibitor' do
-        expect(limiter).to be_a Berater::Inhibitor
-        expect(limiter.key).to be :key
-      end
-
-      it 'inherits redis' do
-        expect(limiter.redis).to be Berater.redis
-      end
-
-      it 'accepts options' do
-        redis = double('Redis')
-        limiter = Berater.new(:key, 0, redis: redis)
-        expect(limiter.redis).to be redis
+      it 'instantiates the right class' do
+        expect(limiter).to be_a klass
       end
-    end
 
-    context 'rate mode' do
-      let(:limiter) { Berater.new(:key, 1, :second) }
-
-      it 'instantiates a RateLimiter' do
-        expect(limiter).to be_a Berater::RateLimiter
+      it 'sets the key' do
         expect(limiter.key).to be :key
       end
 
@@ -75,60 +40,45 @@ describe Berater do
         expect(limiter.redis).to be Berater.redis
       end
 
-      it 'accepts options' do
+      it 'accepts an optional redis parameter' do
         redis = double('Redis')
-        limiter = Berater.new(:key, 1, :second, redis: redis)
+        limiter = Berater.new(:key, capacity, opts.merge(redis: redis))
         expect(limiter.redis).to be redis
       end
     end
 
-    context 'concurrency mode' do
-      let(:limiter) { Berater.new(:key, 1) }
+    describe 'Berater() convenience method' do
+      let(:limiter) { Berater(:key, capacity, **opts) }
 
-      it 'instantiates a ConcurrencyLimiter' do
-        expect(limiter).to be_a Berater::ConcurrencyLimiter
-        expect(limiter.key).to be :key
+      it 'creates a limiter' do
+        expect(limiter).to be_a klass
       end
 
-      it 'inherits redis' do
-        expect(limiter.redis).to be Berater.redis
+      it 'creates an equivalent limiter' do
+        expect(limiter).to eq Berater.new(:key, capacity, **opts)
       end
 
-      it 'accepts options' do
-        redis = double('Redis')
-        limiter = Berater.new(:key, 1, redis: redis)
-        expect(limiter.redis).to be redis
-      end
-    end
-  end
+      context 'with a block' do
+        before { Berater.test_mode = :pass }
 
-  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
+        subject { Berater(:key, capacity, **opts) { 123 } }
 
-      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
+          subject
+        end
 
-          begin
-            res = Berater(:key, *args) { true }
-            expect(res).to be true
-          rescue Berater::Overloaded
-            expect(klass).to be Berater::Inhibitor
-          end
+        it 'yields' do
+          is_expected.to be 123
         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
 
+  include_examples 'a Berater', Berater::ConcurrencyLimiter, 1, timeout: 1
+  include_examples 'a Berater', Berater::Inhibitor, 0
+  include_examples 'a Berater', Berater::RateLimiter, 1, interval: :second
+  include_examples 'a Berater', Berater::StaticLimiter, 1
+  include_examples 'a Berater', Berater::Unlimiter, Float::INFINITY
 end