berater 0.8.0 → 0.11.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: c32c1de752ccc7a2a71e4628b99f3561db623374dabaf95af39aadcc3b1a2d16
-  data.tar.gz: 33019954de1d79522ed7105e9f0a7da3a204df74cacd25aa4b57e749a81047bb
+  metadata.gz: f57ab651eeb34e6a0bb7b889fefa9c13d1d1fc5b6e834bc2b9a6e0ba7b74a8e2
+  data.tar.gz: 8112aa91ae48872132d8222a07f03dce42ff6f5ff26a3bf17906a389ba675ec9
 SHA512:
-  metadata.gz: 89460a9a81a8d384addb095c425da0eed4fdbe0e15ecd85ee97ffcf4292441deeb34c7e4cd4591adca06096a17f7c131b79c6474d9efb2b635fbe398ad68d00e
-  data.tar.gz: 0a331b939e4c19817c42a8aa8fc5be7fd59d63f6becd5fe6ed1c226690d1524ef040a28d8bcb625d75452ab2f5c8539ddf69093bbd30c4ba9e3353259a3e568d
+  metadata.gz: 17ae5da54dcf9535d4afe93d7e239ddc1fefaedf3419883cfa22d1f55c0c0aa7a9365986509e8e54e37c08c258068db4a76064dcc6af3f654e170b3efc8c1cd0
+  data.tar.gz: 04b1e7553e86107bbc91fc486b7b0ae099f5575b8654f7800f2a1c281b7b8bfde1e63df2378f03fcb4a9864b15b6a5de610bb6e7a1ce4d4b8e53898d9f4f0869

data/lib/berater/limiter.rb

@@ -7,19 +7,14 @@ module Berater
       options[:redis] || Berater.redis
     end
 
-    def limit(capacity: nil, cost: 1, &block)
-      capacity ||= @capacity
+    def limit(**opts, &block)
+      opts[:capacity] ||= @capacity
+      opts[:cost] ||= 1
 
-      unless capacity.is_a?(Numeric) && capacity >= 0
-        raise ArgumentError, "invalid capacity: #{capacity}"
-      end
-
-      unless cost.is_a?(Numeric) && cost >= 0 && cost < Float::INFINITY
-        raise ArgumentError, "invalid cost: #{cost}"
+      lock = Berater.middleware.call(self, **opts) do |limiter, **opts|
+        limiter.inner_limit(**opts)
       end
 
-      lock = acquire_lock(capacity, cost)
-
       if block_given?
         begin
           yield lock
@@ -31,6 +26,23 @@ module Berater
       end
     end
 
+    protected def inner_limit(capacity:, cost:)
+      unless capacity.is_a?(Numeric) && capacity >= 0
+        raise ArgumentError, "invalid capacity: #{capacity}"
+      end
+
+      unless cost.is_a?(Numeric) && cost >= 0 && cost < Float::INFINITY
+        raise ArgumentError, "invalid cost: #{cost}"
+      end
+
+      acquire_lock(capacity, cost)
+    rescue NoMethodError => e
+      raise unless e.message.include?("undefined method `evalsha' for")
+
+      # repackage error so it's easier to understand
+      raise RuntimeError, "invalid redis connection: #{redis}"
+    end
+
     def utilization
       lock = limit(cost: 0)
 
@@ -52,13 +64,6 @@ module Berater
       self.redis.connection == other.redis.connection
     end
 
-    def self.new(*)
-      # can only call via subclass
-      raise NoMethodError if self == Berater::Limiter
-
-      super
-    end
-
     protected
 
     attr_reader :args
@@ -93,16 +98,33 @@ module Berater
       self.class.cache_key(instance_key)
     end
 
-    def self.cache_key(key)
-      klass = to_s.split(':')[-1]
-      "Berater:#{klass}:#{key}"
-    end
+    class << self
+      def new(*args, **kwargs)
+        # can only call via subclass
+        raise NoMethodError if self == Berater::Limiter
 
-    def self.inherited(subclass)
-      # automagically create convenience method
-      name = subclass.to_s.split(':')[-1]
-      Berater.define_singleton_method(name) do |*args, **opts, &block|
-        Berater::Utils.convenience_fn(subclass, *args, **opts, &block)
+        if RUBY_VERSION < '3' && kwargs.empty?
+          # avoid ruby 2 problems with empty hashes
+          super(*args)
+        else
+          super
+        end
+      end
+
+      def cache_key(key)
+        klass = to_s.split(':')[-1]
+        "Berater:#{klass}:#{key}"
+      end
+
+      protected
+
+      def inherited(subclass)
+        # automagically create convenience method
+        name = subclass.to_s.split(':')[-1]
+
+        Berater.define_singleton_method(name) do |*args, **opts, &block|
+          Berater::Utils.convenience_fn(subclass, *args, **opts, &block)
+        end
       end
     end
 

data/lib/berater/limiter_set.rb

@@ -0,0 +1,66 @@
+module Berater
+  private
+
+  class LimiterSet
+    include Enumerable
+
+    def initialize
+      @limiters = {}
+    end
+
+    def each(&block)
+      @limiters.each_value(&block)
+    end
+
+    def <<(limiter)
+      key = limiter.key if limiter.respond_to?(:key)
+      send(:[]=, key, limiter)
+    end
+
+    def []=(key, limiter)
+      unless limiter.is_a? Berater::Limiter
+        raise ArgumentError, "expected Berater::Limiter, found: #{limiter}"
+      end
+
+      @limiters[key] = limiter
+    end
+
+    def [](key)
+      @limiters[key]
+    end
+
+    def fetch(key, val = default = true, &block)
+      args = default ? [ key ] : [ key, val ]
+      @limiters.fetch(*args, &block)
+    end
+
+    def include?(key)
+      if key.is_a? Berater::Limiter
+        @limiters.value?(key)
+      else
+        @limiters.key?(key)
+      end
+    end
+
+    def clear
+      @limiters.clear
+    end
+
+    def count
+      @limiters.count
+    end
+
+    def delete(key)
+      if key.is_a? Berater::Limiter
+        @limiters.delete(key.key)
+      else
+        @limiters.delete(key)
+      end
+    end
+    alias remove delete
+
+    def empty?
+      @limiters.empty?
+    end
+  end
+end

data/lib/berater/lock.rb

@@ -20,7 +20,7 @@ module Berater
 
       @released_at = Time.now
       @release_fn&.call
-      nil
+      true
     end
 
   end

data/lib/berater/lua_script.rb

@@ -1,4 +1,5 @@
 require 'digest'
+require 'redis'
 
 module Berater
   class LuaScript
@@ -6,11 +7,11 @@ module Berater
     attr_reader :source
 
     def initialize(source)
-      @source = source
+      @source = source.dup.freeze
     end
 
     def sha
-      @sha ||= Digest::SHA1.hexdigest(minify)
+      @sha ||= Digest::SHA1.hexdigest(minify).freeze
     end
 
     def eval(redis, *args)
@@ -44,7 +45,7 @@ module Berater
     def minify
       # trim comments (whole line and partial)
       # and whitespace (prefix and empty lines)
-      @minify ||= source.gsub(/^\s*--.*\n|\s*--.*|^\s*|^$\n/, '').chomp
+      @minify ||= source.gsub(/^\s*--.*\n|\s*--.*|^\s*|^$\n/, '').chomp.freeze
     end
 
   end

data/lib/berater/middleware/fail_open.rb

@@ -0,0 +1,41 @@
+require 'set'
+
+module Berater
+  module Middleware
+    class FailOpen
+      ERRORS = Set[
+        Redis::BaseConnectionError,
+      ]
+
+      def initialize(errors: nil, on_fail: nil)
+        @errors = errors || ERRORS
+        @on_fail = on_fail
+      end
+
+      def call(*, **opts)
+        yield.tap do |lock|
+          # wrap lock.release so it fails open
+
+          # save reference to original function
+          release_fn = lock.method(:release)
+
+          # make bound variables accessible to block
+          errors = @errors
+          on_fail = @on_fail
+
+          lock.define_singleton_method(:release) do
+            release_fn.call
+          rescue *errors => e
+            on_fail&.call(e)
+            false
+          end
+        end
+      rescue *@errors => e
+        @on_fail&.call(e)
+
+        # fail open by faking a lock
+        Berater::Lock.new(opts[:capacity], -1)
+      end
+    end
+  end
+end

data/lib/berater/middleware/load_shedder.rb

@@ -0,0 +1,31 @@
+module Berater
+  module Middleware
+    class LoadShedder
+      PRIORITY_RANGE = 1..5
+
+      def initialize(default_priority: nil)
+        @default_priority = default_priority
+      end
+
+      def call(*args, **opts)
+        if priority = opts.delete(:priority) || @default_priority
+          opts[:capacity] = adjust_capacity(opts[:capacity], priority)
+        end
+
+        yield *args, **opts
+      end
+
+      protected
+
+      def adjust_capacity(capacity, priority)
+        unless PRIORITY_RANGE.include?(priority)
+          return capacity
+        end
+
+        # priority 1 stays at 100%, 2 scales down to 90%, 5 to 60%
+        factor = 1 - (priority - 1) * 0.1
+        (capacity * factor).floor
+      end
+    end
+  end
+end

data/lib/berater/middleware.rb

@@ -0,0 +1,6 @@
+module Berater
+  module Middleware
+    autoload 'FailOpen', 'berater/middleware/fail_open'
+    autoload 'LoadShedder', 'berater/middleware/load_shedder'
+  end
+end

data/lib/berater/utils.rb

@@ -44,12 +44,7 @@ module Berater
 
     def convenience_fn(klass, *args, **opts, &block)
       limiter = klass.new(*args, **opts)
-      if block_given?
-        limiter.limit(&block)
-      else
-        limiter
-      end
+      block ? limiter.limit(&block) : limiter
     end
-
   end
 end

data/lib/berater/version.rb

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

data/lib/berater.rb

@@ -1,8 +1,11 @@
 require 'berater/limiter'
+require 'berater/limiter_set'
 require 'berater/lock'
 require 'berater/lua_script'
+require 'berater/middleware'
 require 'berater/utils'
 require 'berater/version'
+require 'meddleware'
 
 module Berater
   extend self
@@ -15,8 +18,14 @@ module Berater
     yield self
   end
 
-  def reset
-    @redis = nil
+  def limiters
+    @limiters ||= LimiterSet.new
+  end
+
+  def middleware(&block)
+    (@middleware ||= Meddleware.new).tap do
+      @middleware.instance_eval(&block) if block_given?
+    end
   end
 
   def new(key, capacity, **opts)
@@ -48,6 +57,11 @@ module Berater
     end
   end
 
+  def reset
+    @redis = nil
+    limiters.clear
+    middleware.clear
+  end
 end
 
 # convenience method

data/spec/berater_spec.rb

@@ -24,6 +24,24 @@ describe Berater do
     end
   end
 
+  describe '.limiters' do
+    subject { Berater.limiters }
+
+    let(:limiter) { Berater(:key, 1) }
+
+    it 'provides access to predefined limiters' do
+      expect(Berater.limiters).to be_a Berater::LimiterSet
+    end
+
+    it 'resets with Berater' do
+      subject << limiter
+      is_expected.not_to be_empty
+
+      Berater.reset
+      is_expected.to be_empty
+    end
+  end
+
   shared_examples 'a Berater' do |klass, capacity, **opts|
     describe '.new' do
       let(:limiter) { Berater.new(:key, capacity, **opts) }
@@ -41,8 +59,8 @@ describe Berater do
       end
 
       it 'accepts an optional redis parameter' do
-        redis = double('Redis')
-        limiter = Berater.new(:key, capacity, opts.merge(redis: redis))
+        redis = double(Redis)
+        limiter = Berater.new(:key, capacity, **opts.merge(redis: redis))
         expect(limiter.redis).to be redis
       end
     end

data/spec/limiter_set_spec.rb

@@ -0,0 +1,173 @@
+describe Berater::LimiterSet do
+  subject { described_class.new }
+
+  let(:unlimiter) { Berater::Unlimiter.new }
+  let(:inhibitor) { Berater::Inhibitor.new }
+
+  describe '#each' do
+    it 'returns an Enumerator' do
+      expect(subject.each).to be_a Enumerator
+    end
+
+    it 'works with an empty set' do
+      expect(subject.each.to_a).to eq []
+    end
+
+    it 'returns elements' do
+      subject << unlimiter
+      expect(subject.each.to_a).to eq [ unlimiter ]
+    end
+  end
+
+  describe '#<<' do
+    it 'adds a limiter' do
+      subject << unlimiter
+      expect(subject.each.to_a).to eq [ unlimiter ]
+    end
+
+    it 'rejects things that are not limiters' do
+      expect {
+        subject << :foo
+      }.to raise_error(ArgumentError)
+    end
+
+    it 'updates existing keys' do
+      limiter = Berater::Unlimiter.new
+      expect(limiter).to eq unlimiter
+      expect(limiter).not_to be unlimiter
+
+      subject << unlimiter
+      subject << limiter
+
+      expect(subject.each.to_a).to eq [ limiter ]
+    end
+  end
+
+  describe '[]=' do
+    it 'adds a limiter' do
+      subject[:key] = unlimiter
+
+      expect(subject.each.to_a).to eq [ unlimiter ]
+      is_expected.to include :key
+      is_expected.to include unlimiter
+    end
+
+    it 'rejects things that are not limiters' do
+      expect {
+        subject[:key] = :foo
+      }.to raise_error(ArgumentError)
+    end
+  end
+
+  describe '#[]' do
+    it 'returns nil for missing keys' do
+      expect(subject[:key]).to be nil
+      expect(subject[nil]).to be nil
+    end
+
+    it 'retreives limiters' do
+      subject << unlimiter
+      expect(subject[unlimiter.key]).to be unlimiter
+    end
+  end
+
+  describe '#fetch' do
+    it 'raises for missing keys' do
+      expect {
+        subject.fetch(:key)
+      }.to raise_error(KeyError)
+
+      expect {
+        subject.fetch(nil)
+      }.to raise_error(KeyError)
+    end
+
+    it 'returns the default if provided' do
+      expect(subject.fetch(:key, unlimiter)).to be unlimiter
+    end
+
+    it 'calls the default proc if provided' do
+      expect {|block| subject.fetch(:key, &block) }.to yield_control
+    end
+
+    it 'retreives limiters' do
+      subject << unlimiter
+      expect(subject.fetch(unlimiter.key)).to be unlimiter
+      expect(subject.fetch(unlimiter.key, :default)).to be unlimiter
+    end
+  end
+
+  describe '#include?' do
+    before do
+      subject << unlimiter
+    end
+
+    it 'works with keys' do
+      is_expected.to include unlimiter.key
+    end
+
+    it 'works with limiters' do
+      is_expected.to include unlimiter
+    end
+
+    it 'works when target is missing' do
+      is_expected.not_to include inhibitor.key
+      is_expected.not_to include inhibitor
+    end
+  end
+
+  describe '#clear' do
+    it 'works when empty' do
+      subject.clear
+    end
+
+    it 'clears limiters' do
+      subject << unlimiter
+      is_expected.to include unlimiter
+
+      subject.clear
+      is_expected.not_to include unlimiter
+    end
+  end
+
+  describe '#count' do
+    it 'counts' do
+      expect(subject.count).to be 0
+
+      subject << unlimiter
+      expect(subject.count).to be 1
+    end
+  end
+
+  describe '#delete' do
+    it 'works when the target is missing' do
+      subject.delete(unlimiter)
+      subject.delete(unlimiter.key)
+    end
+
+    it 'works with keys' do
+      subject << unlimiter
+      is_expected.to include unlimiter
+
+      subject.delete(unlimiter.key)
+      is_expected.not_to include unlimiter
+    end
+
+    it 'works with limiters' do
+      subject << unlimiter
+      is_expected.to include unlimiter
+
+      subject.delete(unlimiter)
+      is_expected.not_to include unlimiter
+    end
+  end
+
+  describe '#empty?' do
+    it 'works' do
+      is_expected.to be_empty
+
+      subject << unlimiter
+      is_expected.not_to be_empty
+    end
+  end
+end

data/spec/limiter_spec.rb

@@ -1,6 +1,8 @@
 describe Berater::Limiter do
-  it 'can not be initialized' do
-    expect { described_class.new }.to raise_error(NoMethodError)
+  describe '.new' do
+    it 'can only be called on subclasses' do
+      expect { described_class.new }.to raise_error(NoMethodError)
+    end
   end
 
   describe 'abstract methods' do
@@ -50,6 +52,43 @@ describe Berater::Limiter do
         }.to raise_error(ArgumentError)
       end
     end
+
+    context 'when Berater.redis is nil' do
+      let!(:redis) { Berater.redis }
+
+      before { Berater.redis = nil }
+
+      it 'works with Unlimiter since redis is not used' do
+        expect(subject.redis).to be nil
+        expect {|b| subject.limit(&b) }.to yield_control
+      end
+
+      it 'raises when redis is needed' do
+        limiter = Berater::RateLimiter.new(:key, 1, :second)
+        expect(limiter.redis).to be nil
+        expect { limiter.limit }.to raise_error(RuntimeError)
+      end
+
+      it 'works when redis is passed in' do
+        limiter = Berater::RateLimiter.new(:key, 1, :second, redis: redis)
+        expect {|b| limiter.limit(&b) }.to yield_control
+      end
+
+      it 'raises when redis is bogus' do
+        limiter = Berater::RateLimiter.new(:key, 1, :second, redis: :stub)
+        expect { limiter.limit }.to raise_error(RuntimeError)
+      end
+    end
+
+    it 'releases the lock even when limited code raises an error' do
+      lock = Berater::Lock.new(Float::INFINITY, 0)
+      expect(subject).to receive(:acquire_lock).and_return(lock)
+      expect(lock).to receive(:release)
+
+      expect {
+        subject.limit { raise 'fail' }
+      }.to raise_error(RuntimeError)
+    end
   end
 
   describe '#==' do

data/spec/lua_script_spec.rb

@@ -93,5 +93,4 @@ describe Berater::LuaScript do
       it { subject }
     end
   end
-
 end

data/spec/middleware/fail_open_spec.rb

@@ -0,0 +1,184 @@
+describe Berater::Middleware::FailOpen do
+  let(:limiter) { Berater::Unlimiter.new }
+  let(:lock) { limiter.limit }
+  let(:error) { Redis::TimeoutError }
+
+  describe '.call' do
+    let(:instance) { described_class.new(errors: errors, on_fail: on_fail) }
+    let(:errors) { nil }
+    let(:on_fail) { nil }
+
+    it 'returns the blocks value' do
+      expect(instance.call { lock }).to be lock
+    end
+
+    context 'when there is an error during lock acquisition' do
+      subject { instance.call { raise error } }
+
+      it 'still returns a lock' do
+        expect(subject).to be_a Berater::Lock
+      end
+
+      it 'creates a new, fake lock' do
+        expect(Berater::Lock).to receive(:new)
+        subject
+      end
+
+      it 'returns a lock that is releasable' do
+        expect(subject.release).to be true
+      end
+
+      context 'when an on_fail handler is defined' do
+        let(:on_fail) { double(Proc) }
+
+        it 'calls the handler' do
+          expect(on_fail).to receive(:call).with(error)
+          subject
+        end
+      end
+
+      context 'when the error is an IOError' do
+        let(:error) { IOError }
+
+        it 'would normally not catch the error' do
+          expect { subject }.to raise_error(error)
+        end
+
+        context 'and errors option is set' do
+          let(:errors) { [ error ] }
+
+          it 'catches the error' do
+            expect { subject }.not_to raise_error
+          end
+        end
+      end
+    end
+
+    context 'when there is an error during lock release' do
+      subject { instance.call { lock }.release }
+
+      before do
+        expect(lock).to receive(:release).and_raise(error)
+      end
+
+      it 'handles the exception' do
+        expect { subject }.not_to raise_error
+      end
+
+      it 'returns false since lock was not released' do
+        is_expected.to be false
+      end
+
+      context 'when an on_fail handler is defined' do
+        let(:on_fail) { double(Proc) }
+
+        it 'calls the handler' do
+          expect(on_fail).to receive(:call).with(Exception)
+          subject
+        end
+      end
+
+      context 'when the error is an IOError' do
+        let(:error) { IOError }
+
+        it 'would normally not catch the error' do
+          expect { subject }.to raise_error(error)
+        end
+
+        context 'and errors option is set' do
+          let(:errors) { [ error ] }
+
+          it 'catches the error' do
+            expect { subject }.not_to raise_error
+          end
+        end
+      end
+    end
+  end
+
+  context 'when there is an error during lock acquisition' do
+    before do
+      expect(limiter).to receive(:acquire_lock).and_raise(error)
+    end
+
+    it 'raises an exception for the caller' do
+      expect { limiter.limit }.to raise_error(error)
+    end
+
+    context 'when FailOpen middleware is enabled' do
+      before do
+        Berater.middleware.use described_class
+      end
+
+      it 'fails open' do
+        expect(limiter.limit).to be_a Berater::Lock
+      end
+
+      it 'returns the intended result' do
+        expect(limiter.limit { 123 }).to be 123
+      end
+    end
+
+    context 'when FailOpen middleware is enabled with callback' do
+      before do
+        Berater.middleware.use described_class, on_fail: on_fail
+      end
+      let(:on_fail) { double(Proc) }
+
+      it 'calls the callback' do
+        expect(on_fail).to receive(:call).with(Exception)
+        limiter.limit
+      end
+    end
+  end
+
+  context 'when there is an error during lock release' do
+    before do
+      allow(limiter).to receive(:acquire_lock).and_return(lock)
+      allow(lock).to receive(:release).and_raise(error)
+    end
+
+    it 'acquires a lock' do
+      expect(limiter.limit).to be_a Berater::Lock
+      expect(limiter.limit).to be lock
+    end
+
+    it 'raises an exception when lock is released' do
+      expect {
+        limiter.limit.release
+      }.to raise_error(error)
+    end
+
+    it 'raises an exception when lock is auto released' do
+      expect {
+        limiter.limit {}
+      }.to raise_error(error)
+    end
+
+    context 'when FailOpen middleware is enabled' do
+      before do
+        Berater.middleware.use described_class
+      end
+
+      it 'fails open' do
+        expect { limiter.limit.release }.not_to raise_error
+      end
+
+      it 'returns the intended result' do
+        expect(limiter.limit { 123 }).to be 123
+      end
+    end
+
+    context 'when FailOpen middleware is enabled with callback' do
+      before do
+        Berater.middleware.use described_class, on_fail: on_fail
+      end
+      let(:on_fail) { double(Proc) }
+
+      it 'calls the callback' do
+        expect(on_fail).to receive(:call).with(Exception)
+        limiter.limit {}
+      end
+    end
+  end
+end

data/spec/middleware/load_shedder_spec.rb

@@ -0,0 +1,130 @@
+describe Berater::Middleware::LoadShedder do
+  describe '#call' do
+    subject { described_class.new }
+
+    before { Berater.test_mode = :pass }
+
+    it 'yields' do
+      expect {|b| subject.call(&b) }.to yield_control
+    end
+
+    it 'passes through capacity and cost options' do
+      opts = {
+        capacity: 1,
+        cost: 2,
+      }
+
+      subject.call(**opts) do |**passed_opts|
+        expect(passed_opts).to eq(opts)
+      end
+    end
+
+    it 'strips out priority from options' do
+      opts = {
+        capacity: 1,
+        priority: 3,
+      }
+
+      subject.call(**opts) do |**passed_opts|
+        expect(passed_opts.keys).not_to include(:priority)
+      end
+    end
+
+    it 'keeps full capacity for priority 1' do
+      subject.call(capacity: 100, priority: 1) do |capacity:|
+        expect(capacity).to eq 100
+      end
+    end
+
+    it 'adjusts the capactiy according to priority' do
+      subject.call(capacity: 100, priority: 2) do |capacity:|
+        expect(capacity).to be < 100
+      end
+
+      subject.call(capacity: 100, priority: 5) do |capacity:|
+        expect(capacity).to eq 60
+      end
+    end
+
+    it 'ignores bogus priority options' do
+      subject.call(capacity: 100, priority: 50) do |capacity:|
+        expect(capacity).to eq 100
+      end
+
+      subject.call(capacity: 100, priority: 'abc') do |capacity:|
+        expect(capacity).to eq 100
+      end
+
+      subject.call(capacity: 100, priority: '123') do |capacity:|
+        expect(capacity).to eq 100
+      end
+    end
+
+    it 'works with a fractional priority' do
+      subject.call(capacity: 100, priority: 1.5) do |capacity:|
+        expect(capacity).to be < 100
+      end
+    end
+
+    context 'with a default priority' do
+      subject { described_class.new(default_priority: 5) }
+
+      it 'keeps full capacity for priority 1' do
+        subject.call(capacity: 100, priority: 1) do |capacity:|
+          expect(capacity).to eq 100
+        end
+      end
+
+      it 'uses the default priority' do
+        subject.call(capacity: 100) do |capacity:|
+          expect(capacity).to eq 60
+        end
+      end
+    end
+  end
+
+  context 'with a limiter' do
+    before do
+      Berater.middleware.use Berater::Middleware::LoadShedder
+    end
+
+    shared_examples 'limiter load shedding' do |limiter|
+      it 'passes through the capactiy properly' do
+        expect(limiter).to receive(:inner_limit).with(
+          hash_including(capacity: 100)
+        ).and_call_original
+
+        limiter.limit
+      end
+
+      it 'scales the capactiy with priority' do
+        expect(limiter).to receive(:inner_limit).with(
+          hash_including(capacity: 60)
+        ).and_call_original
+
+        limiter.limit(priority: 5)
+      end
+
+      it 'overloads properly' do
+        60.times { limiter.limit(priority: 5) }
+
+        expect {
+          limiter.limit(priority: 5)
+        }.to be_overloaded
+
+        expect {
+          limiter.limit(priority: 4)
+        }.not_to be_overloaded
+
+        39.times { limiter.limit(priority: 1) }
+
+        expect {
+          limiter.limit(priority: 1)
+        }.to be_overloaded
+      end
+    end
+
+    include_examples 'limiter load shedding', Berater::ConcurrencyLimiter.new(:key, 100)
+    include_examples 'limiter load shedding', Berater::RateLimiter.new(:key, 100, :second)
+  end
+end

data/spec/middleware_spec.rb

@@ -0,0 +1,120 @@
+class Meddler
+  def call(*args, **kwargs)
+    yield
+  end
+end
+
+describe 'Berater.middleware' do
+  subject { Berater.middleware }
+
+  describe 'adding middleware' do
+    after { is_expected.to include Meddler }
+
+    it 'can be done inline' do
+      Berater.middleware.use Meddler
+    end
+
+    it 'can be done with a block' do
+      Berater.middleware do
+        use Meddler
+      end
+    end
+  end
+
+  it 'resets along with Berater' do
+    Berater.middleware.use Meddler
+    is_expected.to include Meddler
+
+    Berater.reset
+    is_expected.to be_empty
+  end
+
+  describe 'Berater::Limiter#limit' do
+    let(:middleware) { Meddler.new }
+    let(:limiter) { Berater::ConcurrencyLimiter.new(:key, 1) }
+
+    before do
+      expect(Meddler).to receive(:new).and_return(middleware).at_least(1)
+      Berater.middleware.use Meddler
+    end
+
+    it 'calls the middleware' do
+      expect(middleware).to receive(:call)
+      limiter.limit
+    end
+
+    it 'calls the middleware, passing the limiter and options' do
+      expect(middleware).to receive(:call).with(
+        limiter,
+        hash_including(:capacity, :cost)
+      )
+
+      limiter.limit
+    end
+
+    context 'when used per ususual' do
+      before do
+        expect(middleware).to receive(:call).and_call_original.at_least(1)
+      end
+
+      it 'still works inline' do
+        expect(limiter.limit).to be_a Berater::Lock
+      end
+
+      it 'still works in block mode' do
+        expect(limiter.limit { 123 }).to be 123
+      end
+
+      it 'still has limits' do
+        limiter.limit
+        expect(limiter).to be_overloaded
+      end
+    end
+
+    context 'when middleware meddles' do
+      it 'can change the capacity' do
+        expect(middleware).to receive(:call) do |limiter, **opts, &block|
+          opts[:capacity] = 0
+          block.call(limiter, **opts)
+        end
+
+        expect { limiter.limit }.to be_overloaded
+      end
+
+      it 'can change the cost' do
+        expect(middleware).to receive(:call) do |limiter, **opts, &block|
+          opts[:cost] = 2
+          block.call(limiter, **opts)
+        end
+
+        expect { limiter.limit }.to be_overloaded
+      end
+
+      it 'can change the limiter' do
+        other_limiter = Berater::Inhibitor.new
+
+        expect(middleware).to receive(:call) do |limiter, **opts, &block|
+          block.call(other_limiter, **opts)
+        end
+        expect(other_limiter).to receive(:acquire_lock).and_call_original
+
+        expect { limiter.limit }.to be_overloaded
+      end
+
+      it 'can abort by not yielding' do
+        expect(middleware).to receive(:call)
+        expect(limiter.limit).to be nil
+      end
+
+      it 'can intercept the lock' do
+        expect(middleware).to receive(:call) do |&block|
+          lock = block.call
+          expect(lock).to be_a Berater::Lock
+          expect(lock.capacity).to eq limiter.capacity
+        end
+
+        limiter.limit
+      end
+    end
+  end
+end

metadata

@@ -1,29 +1,43 @@
 --- !ruby/object:Gem::Specification
 name: berater
 version: !ruby/object:Gem::Version
-  version: 0.8.0
+  version: 0.11.0
 platform: ruby
 authors:
 - Daniel Pepper
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2021-03-22 00:00:00.000000000 Z
+date: 2021-10-25 00:00:00.000000000 Z
 dependencies:
+- !ruby/object:Gem::Dependency
+  name: meddleware
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0.2'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0.2'
 - !ruby/object:Gem::Dependency
   name: redis
   requirement: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '0'
+        version: '3'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - ">="
       - !ruby/object:Gem::Version
-        version: '0'
+        version: '3'
 - !ruby/object:Gem::Dependency
   name: benchmark
   requirement: !ruby/object:Gem::Requirement
@@ -133,8 +147,12 @@ files:
 - lib/berater/dsl.rb
 - lib/berater/inhibitor.rb
 - lib/berater/limiter.rb
+- lib/berater/limiter_set.rb
 - lib/berater/lock.rb
 - lib/berater/lua_script.rb
+- lib/berater/middleware.rb
+- lib/berater/middleware/fail_open.rb
+- lib/berater/middleware/load_shedder.rb
 - lib/berater/rate_limiter.rb
 - lib/berater/rspec.rb
 - lib/berater/rspec/matchers.rb
@@ -148,9 +166,13 @@ files:
 - spec/dsl_refinement_spec.rb
 - spec/dsl_spec.rb
 - spec/inhibitor_spec.rb
+- spec/limiter_set_spec.rb
 - spec/limiter_spec.rb
 - spec/lua_script_spec.rb
 - spec/matchers_spec.rb
+- spec/middleware/fail_open_spec.rb
+- spec/middleware/load_shedder_spec.rb
+- spec/middleware_spec.rb
 - spec/rate_limiter_spec.rb
 - spec/riddle_spec.rb
 - spec/static_limiter_spec.rb
@@ -176,19 +198,23 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.0.8
+rubygems_version: 3.1.6
 signing_key:
 specification_version: 4
 summary: Berater
 test_files:
 - spec/rate_limiter_spec.rb
+- spec/middleware/load_shedder_spec.rb
+- spec/middleware/fail_open_spec.rb
 - spec/matchers_spec.rb
 - spec/dsl_refinement_spec.rb
 - spec/test_mode_spec.rb
+- spec/middleware_spec.rb
 - spec/dsl_spec.rb
 - spec/lua_script_spec.rb
 - spec/concurrency_limiter_spec.rb
 - spec/riddle_spec.rb
+- spec/limiter_set_spec.rb
 - spec/utils_spec.rb
 - spec/berater_spec.rb
 - spec/limiter_spec.rb