rack-berater 0.2.0 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 7db43d04f14f84dd3e2f3f96352226a48d4e36bfe17467fd1e4eb77015ca6148
-  data.tar.gz: 40f79342ca6cd92145a32e2d0b76455850cc4bdd25d701e83157a7acf330bc15
+  metadata.gz: 5840eaff6319d5e24c78c2f50960138ca27fcce25bb3d45354e7eca195b50f0d
+  data.tar.gz: f9cdde9f38c97d82e03d6a1e9cea2e1939eaca13f4202a2a674a6fda4e6f530e
 SHA512:
-  metadata.gz: 1f8b3b29d9d62796a49701340fb53446cfca19258f2083c5973bd17e9a9314abb28c7faf7070a260cd4128bb9028346d18f24c98d96d5d8a9e503e5d113f8456
-  data.tar.gz: 7ee2bef14d6c4334ef2d65b67a60fe432047a0ce8581f62461cf1fd0c65264083833bbd9f3fe45f2909be3b52cb2147db9704a44e2f9cbc9017b83f69da182d4
+  metadata.gz: b0f561957104a0d049d161b720615a97474439c41f16507b3bd34a1f68269fccbea314f54736393a5e030abdf755de8ae30bc5a92724502371728975b6cf7460
+  data.tar.gz: f91a4cb03f24b7c18554c8145375a78f1b61cc66a2d99a0841ddca4b547dc97822bd99dfe9267ae1feeb01d3e2c8eba72d70ea075d367671f3e2077e404a8e36

data/lib/rack/berater/prioritizer.rb

@@ -0,0 +1,90 @@
+require 'rack'
+
+module Rack
+  class Berater
+    class Prioritizer
+      ENV_KEY = 'berater_priority'
+      HEADER = 'X-Berater-Priority'
+
+      def initialize(app, options = {})
+        @app = app
+        @header = options[:header] || HEADER
+      end
+
+      def call(env)
+        priority = env[@header] || env["HTTP_#{@header.upcase.tr('-', '_')}"]
+
+        if priority
+          set_priority(priority)
+          return @app.call(env)
+        end
+
+        cache_key = cache_key_for(env)
+        cached_priority = cache_get(cache_key)
+
+        if cached_priority
+          set_priority(cached_priority)
+        end
+
+        @app.call(env).tap do |status, headers, body|
+          app_priority = headers.delete(@header) if headers
+
+          if app_priority && app_priority != cached_priority
+            # update cache for next time
+            cache_set(cache_key, app_priority)
+          end
+        end
+      ensure
+        Thread.current[ENV_KEY] = nil
+      end
+
+      def self.current_priority
+        Thread.current[ENV_KEY]
+      end
+
+      protected
+
+      def set_priority(priority)
+        Thread.current[ENV_KEY] = priority
+      end
+
+      @@cache_prefix = 'Berater:Rack:Prioritizer'
+      def cache_key_for(env)
+        req = Rack::Request.new(env)
+
+        req_method = env[Rack::REQUEST_METHOD].downcase
+        path = ''
+
+        if defined?(Rails) && Rails.application
+          res = Rails.application.routes.recognize_path(env[Rack::PATH_INFO])
+          path = res.values_at(:controller, :action).compact.join('#')
+        end
+
+        if path.empty?
+          path = env['PATH_INFO'].gsub(%r{/[0-9]+(/|$)}, '/x\1')
+        end
+
+        [
+          @@cache_prefix,
+          req_method,
+          path,
+        ].join(':')
+      end
+
+      @@cache = {}
+
+      def cache_get(key)
+        synchronize { @@cache[key] }
+      end
+
+      def cache_set(key, priority)
+        synchronize { @@cache[key] = priority }
+      end
+
+      @@lock = Thread::Mutex.new
+      def synchronize(&block)
+        @@lock.synchronize(&block)
+      end
+    end
+  end
+end

data/lib/rack/berater/railtie.rb

@@ -1,9 +1,13 @@
-require "rails/railtie"
+require 'rails/railtie'
 
 module Rack
   class Berater
     class Railtie < Rails::Railtie
-      initializer "rack.berater" do |app|
+      initializer 'rack.berater' do |app|
+        if ::Berater.middleware.include?(::Berater::Middleware::LoadShedder)
+          app.middleware.use Rack::Berater::Prioritizer
+        end
+
         app.middleware.use Rack::Berater
       end
     end

data/lib/rack/berater/version.rb

@@ -1,5 +1,5 @@
 module Rack
   class Berater
-    VERSION = "0.2.0"
+    VERSION = "0.3.0"
   end
 end

data/lib/rack/berater.rb

@@ -1,11 +1,12 @@
-require "berater"
-require "rack"
-require "rack/berater/version"
-require "set"
+require 'berater'
+require 'rack'
+require 'rack/berater/version'
+require 'set'
 
 module Rack
   class Berater
-    autoload :Railtie, "rack/berater/railtie"
+    autoload :Prioritizer, 'rack/berater/prioritizer'
+    autoload :Railtie, 'rack/berater/railtie'
 
     ERRORS = Set[ ::Berater::Overloaded ]
 
@@ -27,12 +28,12 @@ module Rack
         when String
           options[:body]
         else
-          raise ArgumentError, "invalid :body option: #{options[:body]}"
+          raise ArgumentError, 'invalid :body option: #{options[:body]}'
         end
 
       # configure headers
       if @options[:body]
-        @options[:headers][Rack::CONTENT_TYPE] = "text/plain"
+        @options[:headers][Rack::CONTENT_TYPE] = 'text/plain'
       end
       @options[:headers].update(options.fetch(:headers, {}))
     end

data/lib/rack-berater.rb

@@ -1 +1 @@
-require "rack/berater"
+require 'rack/berater'

data/spec/limiter_spec.rb

@@ -9,108 +9,108 @@ describe Rack::Berater do
     Rack::Builder.new do
       use Rack::Lint
       run (lambda do |env|
-        [200, {"Content-Type" => "text/plain"}, ["OK"]]
+        [200, {'Content-Type' => 'text/plain'}, ['OK']]
       end)
     end
   end
-  let(:response) { get "/" }
+  let(:response) { get '/' }
 
-  shared_examples "works nominally" do
-    it "has the correct status code" do
+  shared_examples 'works nominally' do
+    it 'has the correct status code' do
       expect(response.status).to eq 200
     end
 
-    it "has the correct headers" do
+    it 'has the correct headers' do
       expect(response.headers).to eq({
-        "Content-Type" => "text/plain",
-        "Content-Length" => "2",
+        'Content-Type' => 'text/plain',
+        'Content-Length' => '2',
       })
     end
 
-    it "has the correct body" do
-      expect(response.body).to eq "OK"
+    it 'has the correct body' do
+      expect(response.body).to eq 'OK'
     end
   end
 
-  context "without a limiter" do
+  context 'without a limiter' do
     before { Berater.test_mode = :fail }
 
-    include_examples "works nominally"
+    include_examples 'works nominally'
   end
 
-  describe "limiter option" do
-    context "when limiter is a limiter" do
+  describe 'limiter option' do
+    context 'when limiter is a limiter' do
       let(:limiter) { ::Berater::Unlimiter.new }
 
-      include_examples "works nominally"
+      include_examples 'works nominally'
 
-      it "calls the limiter" do
+      it 'calls the limiter' do
         expect(limiter).to receive(:limit).and_call_original
         response
       end
 
-      context "when operating beyond limits" do
+      context 'when operating beyond limits' do
         before { Berater.test_mode = :fail }
 
-        it "returns an error" do
+        it 'returns an error' do
           expect(response.status).to eq 429
         end
       end
     end
 
-    context "when limiter is a proc" do
+    context 'when limiter is a proc' do
       let(:limiter_instance) { ::Berater::Unlimiter.new }
       let(:limiter) { Proc.new { limiter_instance } }
 
-      include_examples "works nominally"
+      include_examples 'works nominally'
 
-      it "calls the proc with env" do
+      it 'calls the proc with env' do
         expect(limiter).to receive(:call).with(Hash).and_call_original
         response
       end
 
-      context "when operating beyond limits" do
+      context 'when operating beyond limits' do
         before { Berater.test_mode = :fail }
 
-        it "returns an error" do
+        it 'returns an error' do
           expect(response.status).to eq 429
         end
       end
     end
   end
 
-  describe "enabled? option" do
+  describe 'enabled? option' do
     after { expect(response.status).to eq 200 }
 
     let(:enabled?) { double }
 
-    context "when there is a limiter" do
+    context 'when there is a limiter' do
       let(:limiter) { ::Berater::Unlimiter.new }
 
-      it "should be called with the env hash" do
+      it 'should be called with the env hash' do
         expect(enabled?).to receive(:call) do |env|
           expect(env).to be_a Hash
-          expect(Rack::Request.new(env).path).to eq "/"
+          expect(Rack::Request.new(env).path).to eq '/'
         end
       end
 
-      context "when enabled" do
-        it "should call the limiter" do
+      context 'when enabled' do
+        it 'should call the limiter' do
           expect(enabled?).to receive(:call).and_return(true)
           expect(limiter).to receive(:limit).and_call_original
         end
       end
 
-      context "when disabled" do
-        it "should not call the limiter" do
+      context 'when disabled' do
+        it 'should not call the limiter' do
           expect(enabled?).to receive(:call).and_return(false)
           expect(limiter).not_to receive(:limit)
         end
       end
     end
 
-    context "when there is no limiter" do
-      it "should not call enabled?" do
+    context 'when there is no limiter' do
+      it 'should not call enabled?' do
         expect(enabled?).not_to receive(:call)
       end
     end

data/spec/prioritizer_spec.rb

@@ -0,0 +1,295 @@
+describe Rack::Berater::Prioritizer do
+  after do
+    cache.clear
+    Thread.current[described_class::ENV_KEY] = nil
+  end
+
+  let(:cache) { described_class.class_variable_get(:@@cache) }
+
+  describe '#call' do
+    subject { described_class.new(app) }
+
+    let(:app) { ->(*) { [200, {'Content-Type' => 'text/plain'}, ['OK']] } }
+    let(:cache_key) { subject.send(:cache_key_for, env) }
+    let(:env) { Rack::MockRequest.env_for('/') }
+
+    specify 'sanity check' do
+      expect(app).to receive(:call).and_call_original
+
+      Rack::Lint.new(subject).call(env)
+    end
+
+    it 'checks the cache' do
+      is_expected.to receive(:cache_get)
+
+      subject.call(env)
+    end
+
+    context 'with a cached priority' do
+      before do
+        allow(subject).to receive(:cache_get).with(cache_key).and_return(priority)
+      end
+
+      let(:priority) { '3' }
+
+      after { subject.call(env) }
+
+      it 'sets the priority accordingly' do
+        is_expected.to receive(:set_priority).with(priority)
+      end
+
+      it 'updates the global priority during the request' do
+        expect(app).to receive(:call) do
+          expect(described_class.current_priority).to eq priority
+        end
+      end
+
+      it 'resets the priority after the request completes' do
+        subject.call(env)
+        expect(described_class.current_priority).to be nil
+      end
+    end
+
+    context 'with an incoming priority header' do
+      let(:env) do
+        Rack::MockRequest.env_for(
+          '/',
+          described_class::HEADER => priority,
+        )
+      end
+      let(:priority) { '2' }
+
+      after { subject.call(env) }
+
+      it 'uses the header' do
+        is_expected.to receive(:set_priority).with(priority)
+      end
+
+      it 'ignores any cached value' do
+        allow(subject).to receive(:cache_get).with(cache_key).and_return('123')
+        is_expected.to receive(:set_priority).with(priority)
+      end
+
+      it 'resets the priority after the request completes' do
+        subject.call(env)
+        expect(described_class.current_priority).to be nil
+      end
+    end
+
+    context 'when the app returns a priority header' do
+      let(:app) do
+        ->(*) { [200, { described_class::HEADER => priority }, ['OK']] }
+      end
+
+      let(:priority) { '5' }
+
+      after { subject.call(env) }
+
+      it 'caches the priority' do
+        is_expected.to receive(:cache_set).with(cache_key, priority)
+      end
+
+      it 'removes the header' do
+        _, headers, _ = subject.call(env)
+        expect(headers).not_to include described_class::HEADER
+      end
+
+      it 'updates the cache when a different priority is returned' do
+        expect(subject).to receive(:cache_get).and_return('123')
+        is_expected.to receive(:cache_set).with(cache_key, priority)
+      end
+
+      it 'does not update the cache when the same priority is returned' do
+        expect(subject).to receive(:cache_get).and_return(priority)
+        is_expected.not_to receive(:cache_set)
+      end
+    end
+
+    it 'does not update the cache when no priority is returned' do
+      is_expected.not_to receive(:cache_set)
+      subject.call(env)
+    end
+  end
+
+  describe '#cache_key_for' do
+    subject{ described_class.new(nil).send(:cache_key_for, env) }
+
+    context 'with a basic env' do
+      let(:env) { Rack::MockRequest.env_for('/') }
+
+      it 'has a Berater prefix' do
+        is_expected.to match /^Berater:/
+      end
+
+      it 'combines the verb and path' do
+        is_expected.to match %r{:get:/$}
+      end
+    end
+
+    context 'with a different verb' do
+      let(:env) { Rack::MockRequest.env_for('/', method: 'PUT') }
+
+      it 'combines the verb and path' do
+        is_expected.to match %r{:put:/$}
+      end
+    end
+
+    context 'with a RESTful path' do
+      let(:env) { Rack::MockRequest.env_for('/user/123') }
+
+      it 'normalizes the id' do
+        is_expected.to match %r{:/user/x$}
+      end
+    end
+
+    context 'with a RESTful path and trailing slash' do
+      let(:env) { Rack::MockRequest.env_for('/user/123/') }
+
+      it 'normalizes the id and keeps the trailing slash' do
+        is_expected.to match %r{:/user/x/$}
+      end
+    end
+
+    context 'with a very RESTful path' do
+      let(:env) { Rack::MockRequest.env_for('/user/123/friend/456') }
+
+      it 'normalizes both ids' do
+        is_expected.to match %r{:/user/x/friend/x$}
+      end
+    end
+  end
+
+  context 'as Rack middleware' do
+    def call(path = '/')
+      get(path).body
+    end
+
+    let(:app) do
+      headers = {
+        'Content-Type' => 'text/plain',
+        described_class::HEADER => app_priority,
+      }.compact
+
+      Rack::Builder.new do
+        use Rack::Lint
+        use Rack::Berater::Prioritizer
+
+        run (lambda do |env|
+          [200, headers, [ Rack::Berater::Prioritizer.current_priority.to_s ]]
+        end)
+      end
+    end
+
+    let(:app_priority) { nil }
+
+    it 'starts empty' do
+      expect(call).to be_empty
+    end
+
+    it 'parses incoming priority header' do
+      header described_class::HEADER, '7'
+
+      expect(call).to eq '7'
+    end
+
+    context 'when app returns a priority header' do
+      let(:app_priority) { '8' }
+
+      it 'parses the priority returned from the app' do
+        expect(call).to be_empty
+        expect(cache.values).to include app_priority
+      end
+
+      it 'uses the cached priority for subsequent calls' do
+        expect(call).to be_empty
+        expect(call).to eq app_priority
+      end
+    end
+
+    # context 'when two different endpoints are called' do
+    #   fit 'parses and caches each priority' do
+    #     @app_priority = '6'
+    #     expect(call('/six')).to be_empty
+
+    #     expect(call('/six')).to eq '6'
+
+    #     @app_priority = '9'
+    #     expect(call('/nine')).to be_empty
+    #     expect(call('/nine')).to '9'
+    #   end
+    # end
+  end
+
+  context 'as Rails middleware' do
+    before do
+      class EchoController < ActionController::Base
+        def index
+          render plain: Rack::Berater::Prioritizer.current_priority
+        end
+
+        def six
+          response.set_header(Rack::Berater::Prioritizer::HEADER, '6')
+          index
+        end
+
+        def nine
+          response.set_header(Rack::Berater::Prioritizer::HEADER, '9')
+          index
+        end
+      end
+
+      Rails.application = Class.new(Rails::Application) do
+        config.eager_load = false
+        config.hosts.clear # disable hostname filtering
+      end
+      Rails.application.middleware.use described_class
+      Rails.initialize!
+
+      Rails.application.routes.draw do
+        get '/' => 'echo#index'
+        get '/six' => 'echo#six'
+        get '/nine' => 'echo#nine'
+      end
+    end
+
+    let(:app) { Rails.application }
+    let(:middleware) { described_class.new(app) }
+
+    after { Rails.application = nil }
+
+    describe '#cache_key_for' do
+      subject { described_class.new(app).send(:cache_key_for, env) }
+
+      let(:env) { Rack::MockRequest.env_for('/') }
+
+      it 'uses the controller and action name' do
+        is_expected.to match %r{:echo#index$}
+      end
+    end
+
+    context 'when a priority header is sent' do
+      before { header described_class::HEADER, priority }
+
+      let(:priority) { '6' }
+
+      it 'sets the priority' do
+        expect(get('/six').body).to eq priority
+      end
+    end
+
+    context 'when the app returns a priority' do
+      it 'does not know the first time the controller is called' do
+        expect(get('/six').body).to be_empty
+        expect(get('/nine').body).to be_empty
+      end
+
+      it 'caches the repsonses for the second time' do
+        expect(get('/six').body).to be_empty
+        expect(get('/nine').body).to be_empty
+
+        expect(get('/six').body).to eq '6'
+        expect(get('/nine').body).to eq '9'
+      end
+    end
+  end
+end

data/spec/railtie_spec.rb

@@ -1,5 +1,5 @@
-require "rails"
-require "rack/berater/railtie"
+require 'rails'
+require 'rack/berater/railtie'
 
 RSpec.describe Rack::Berater::Railtie do
   subject { Rails.initialize! }
@@ -7,11 +7,10 @@ RSpec.describe Rack::Berater::Railtie do
   before do
     Rails.application = Class.new(Rails::Application) do
       config.eager_load = false
-      config.logger = ActiveSupport::Logger.new($stdout)
     end
   end
 
-  it "adds middleware automatically" do
+  it 'adds middleware automatically' do
     expect(subject.middleware).to include(Rack::Berater)
   end
 end

data/spec/rescuer_spec.rb

@@ -4,34 +4,34 @@ describe Rack::Berater do
       use Rack::Lint
       run (lambda do |env|
         Berater::Unlimiter() do
-          [200, {"Content-Type" => "text/plain"}, ["OK"]]
+          [200, {'Content-Type' => 'text/plain'}, ['OK']]
         end
       end)
     end
   end
-  let(:response) { get "/" }
+  let(:response) { get '/' }
 
-  shared_examples "works nominally" do
-    it "has the correct status code" do
+  shared_examples 'works nominally' do
+    it 'has the correct status code' do
       expect(response.status).to eq 200
     end
 
-    it "has the correct headers" do
+    it 'has the correct headers' do
       expect(response.headers).to eq({
-        "Content-Type" => "text/plain",
-        "Content-Length" => "2",
+        'Content-Type' => 'text/plain',
+        'Content-Length' => '2',
       })
     end
 
-    it "has the correct body" do
-      expect(response.body).to eq "OK"
+    it 'has the correct body' do
+      expect(response.body).to eq 'OK'
     end
   end
 
-  context "without middleware" do
-    include_examples "works nominally"
+  context 'without middleware' do
+    include_examples 'works nominally'
 
-    it "does not catch limit errors" do
+    it 'does not catch limit errors' do
       Berater.test_mode = :fail
       expect {
         response
@@ -39,55 +39,55 @@ describe Rack::Berater do
     end
   end
 
-  context "with middleware using default settings" do
+  context 'with middleware using default settings' do
     before { app.use described_class }
 
-    include_examples "works nominally"
+    include_examples 'works nominally'
 
-    it "catches and transforms limit errors" do
+    it 'catches and transforms limit errors' do
       Berater.test_mode = :fail
       expect(response.status).to eq 429
-      expect(response.body).to eq "Too Many Requests"
+      expect(response.body).to eq 'Too Many Requests'
     end
   end
 
-  context "with middleware using custom settings" do
+  context 'with middleware using custom settings' do
     before do
       app.use described_class, options
       Berater.test_mode = :fail
     end
 
-    context "with a custom body" do
-      context "with body nil" do
+    context 'with a custom body' do
+      context 'with body nil' do
         let(:options) { { body: nil } }
 
-        it "falls back to the default" do
-          expect(response.body).to eq "Too Many Requests"
+        it 'falls back to the default' do
+          expect(response.body).to eq 'Too Many Requests'
         end
       end
 
-      context "with body disabled" do
+      context 'with body disabled' do
         let(:options) { { body: false } }
 
-        it "should not send a body" do
+        it 'should not send a body' do
           expect(response.body).to be_empty
         end
 
-        it "should not send the Content-Type header" do
+        it 'should not send the Content-Type header' do
           expect(response.headers.keys).not_to include(Rack::CONTENT_TYPE)
         end
       end
 
-      context "with a string" do
-        let(:body) { "none shall pass!" }
+      context 'with a string' do
+        let(:body) { 'none shall pass!' }
         let(:options) { { body: body } }
 
-        it "should send the custom string" do
+        it 'should send the custom string' do
           expect(response.body).to eq body
         end
       end
 
-      context "with an erroneous value" do
+      context 'with an erroneous value' do
         let(:options) { { body: 123 } }
 
         it 'should raise an error' do
@@ -98,56 +98,56 @@ describe Rack::Berater do
       end
     end
 
-    context "with custom headers" do
-      context "with an extra header" do
-        let(:options) { { headers: { Rack::CACHE_CONTROL => "no-cache" } } }
+    context 'with custom headers' do
+      context 'with an extra header' do
+        let(:options) { { headers: { Rack::CACHE_CONTROL => 'no-cache' } } }
 
-        it "should contain the default headers" do
+        it 'should contain the default headers' do
           expect(response.headers.keys).to include(Rack::CONTENT_TYPE)
         end
 
-        it "should also contain the custom header" do
+        it 'should also contain the custom header' do
           expect(response.headers).to include(options[:headers])
         end
       end
 
-      context "with a new content type" do
-        let(:options) { { headers: { Rack::CONTENT_TYPE => "application/json" } } }
+      context 'with a new content type' do
+        let(:options) { { headers: { Rack::CONTENT_TYPE => 'application/json' } } }
 
-        it "should override the Content-Type header" do
+        it 'should override the Content-Type header' do
           expect(response.headers).to include(options[:headers])
         end
       end
     end
 
-    context "with custom status code" do
+    context 'with custom status code' do
       let(:options) { { status_code: 503 } }
 
-      it "catches and transforms limit errors" do
+      it 'catches and transforms limit errors' do
         expect(response.status).to eq 503
-        expect(response.body).to eq "Service Unavailable"
+        expect(response.body).to eq 'Service Unavailable'
       end
     end
   end
 
-  context "with custom error type" do
+  context 'with custom error type' do
     before do
       app.use described_class
       expect(Berater::Limiter).to receive(:new).and_raise(IOError)
     end
 
-    it "normally crashes the app" do
+    it 'normally crashes the app' do
       expect { response }.to raise_error(IOError)
     end
 
-    context "when an error type is registered with middleware" do
+    context 'when an error type is registered with middleware' do
       around do |example|
         Rack::Berater::ERRORS << IOError
         example.run
         Rack::Berater::ERRORS.delete(IOError)
       end
 
-      it "catches and transforms limit errors" do
+      it 'catches and transforms limit errors' do
         expect(response.status).to eq 429
       end
     end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: rack-berater
 version: !ruby/object:Gem::Version
-  version: 0.2.0
+  version: 0.3.0
 platform: ruby
 authors:
 - Daniel Pepper
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2021-10-18 00:00:00.000000000 Z
+date: 2021-11-17 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: berater
@@ -130,9 +130,11 @@ extra_rdoc_files: []
 files:
 - lib/rack-berater.rb
 - lib/rack/berater.rb
+- lib/rack/berater/prioritizer.rb
 - lib/rack/berater/railtie.rb
 - lib/rack/berater/version.rb
 - spec/limiter_spec.rb
+- spec/prioritizer_spec.rb
 - spec/railtie_spec.rb
 - spec/rescuer_spec.rb
 homepage: https://github.com/dpep/rack-berater
@@ -161,4 +163,5 @@ summary: Rack::Berater
 test_files:
 - spec/railtie_spec.rb
 - spec/rescuer_spec.rb
+- spec/prioritizer_spec.rb
 - spec/limiter_spec.rb