From 5f9d0253e6fa4195b8119512674433df30e70fb2 Mon Sep 17 00:00:00 2001 From: James Williams Date: Fri, 7 Jun 2019 10:29:56 -0600 Subject: [PATCH] create system to modify up-front request cost by path closes #CORE-3058 Change-Id: Ie516ce40572cc363a13e613f387319e9c60a97af Reviewed-on: https://gerrit.instructure.com/196909 Tested-by: Jenkins Reviewed-by: Cody Cutrer QA-Review: Cody Cutrer Product-Review: Cody Cutrer --- app/middleware/request_throttle.rb | 41 +++++++++++++++++++++--- spec/lib/canvas/request_throttle_spec.rb | 15 +++++++++ 2 files changed, 51 insertions(+), 5 deletions(-) diff --git a/app/middleware/request_throttle.rb b/app/middleware/request_throttle.rb index 7550a062f81..88b060d4fc6 100644 --- a/app/middleware/request_throttle.rb +++ b/app/middleware/request_throttle.rb @@ -44,15 +44,18 @@ class RequestThrottle starting_cpu = Process.times() request = ActionDispatch::Request.new(env) - # workaround a rails bug where some ActionDispatch::Request methods blow + + # NOTE: calling fullpath was a workaround for a rails bug where some ActionDispatch::Request methods blow # up when using certain servers until fullpath is called once to set env['REQUEST_URI'] - request.fullpath + # so don't remove + path = request.fullpath status, headers, response = nil throttled = false bucket = LeakyBucket.new(client_identifier(request)) - cost = bucket.reserve_capacity do + up_front_cost = bucket.get_up_front_cost_for_path(path) + cost = bucket.reserve_capacity(up_front_cost) do status, headers, response = if !allowed?(request, bucket) throttled = true rate_limit_exceeded @@ -163,7 +166,7 @@ class RequestThrottle end def self.reload! - @whitelist = @blacklist = nil + @whitelist = @blacklist = @dynamic_settings = nil LeakyBucket.reload! end @@ -175,6 +178,10 @@ class RequestThrottle Set.new(Setting.get(key, '').split(',').map(&:strip).reject(&:blank?)) end + def self.dynamic_settings + @dynamic_settings ||= YAML.safe_load(Canvas::DynamicSettings.find(tree: :private)['request_throttle.yml'] || '') || {} + end + def rate_limit_exceeded [403, {'Content-Type' => 'text/plain; charset=utf-8', 'X-Rate-Limit-Remaining' => '0.0'}, @@ -259,14 +266,38 @@ class RequestThrottle end end + def self.up_front_cost_by_path_regex + @up_front_cost_regex_map ||= + begin + hash = RequestThrottle.dynamic_settings['up_front_cost_by_path_regex'] || {} + hash.keys.select{|k| k.is_a?(String)}.map{|k| hash[Regexp.new(k)] = hash.delete(k)} #regexify strings + hash.each do |k, v| + unless k.is_a?(Regexp) && v.is_a?(Numeric) + ::Rails.logger.error("ERROR in request_throttle.yml: up_front_cost_by_path_regex must use Regex => Numeric key-value pairs") + hash.clear + break + end + end + hash + end + end + def self.reload! @custom_settings_hash = nil + @up_front_cost_regex_map = nil end # up_front_cost is a placeholder cost. Essentially it adds some cost to # doing multiple requests in parallel, but that cost is transient -- it # disappears again when the request finishes. - # + def get_up_front_cost_for_path(path) + # if it matches any of the regexes in the setting, return the specified cost + self.class.up_front_cost_by_path_regex.each do |regex, cost| + return cost if regex =~ path + end + self.up_front_cost # otherwise use the default + end + # This method does an initial increment by the up_front_cost, loading the # data out of redis at the same time. It then yields to the block, # expecting the block to return the final cost. It then increments again, diff --git a/spec/lib/canvas/request_throttle_spec.rb b/spec/lib/canvas/request_throttle_spec.rb index 46184a2f2c2..00f05faa3a0 100644 --- a/spec/lib/canvas/request_throttle_spec.rb +++ b/spec/lib/canvas/request_throttle_spec.rb @@ -185,6 +185,7 @@ describe 'RequestThrottle' do it "should throttle if bucket is full" do bucket = throttled_request + expect(bucket).to receive(:get_up_front_cost_for_path).with(base_req['PATH_INFO']).and_return(1) expect(bucket).to receive(:remaining).and_return(-2) expected = rate_limit_exceeded expected[1]['X-Rate-Limit-Remaining'] = "-2" @@ -195,6 +196,7 @@ describe 'RequestThrottle' do allow(RequestThrottle).to receive(:enabled?).and_return(false) bucket = double('Bucket') expect(RequestThrottle::LeakyBucket).to receive(:new).with("user:1").and_return(bucket) + expect(bucket).to receive(:get_up_front_cost_for_path).with(base_req['PATH_INFO']).and_return(1) expect(bucket).to receive(:reserve_capacity).and_yield.and_return(1) expect(bucket).to receive(:remaining).and_return(1) # the cost is still returned anyway @@ -208,6 +210,7 @@ describe 'RequestThrottle' do it "should not throttle, but update, if bucket is not full" do bucket = double('Bucket') expect(RequestThrottle::LeakyBucket).to receive(:new).with("user:1").and_return(bucket) + expect(bucket).to receive(:get_up_front_cost_for_path).with(base_req['PATH_INFO']).and_return(1) expect(bucket).to receive(:reserve_capacity).and_yield.and_return(1) expect(bucket).to receive(:full?).and_return(false) expect(bucket).to receive(:remaining).and_return(599) @@ -357,6 +360,18 @@ describe 'RequestThrottle' do end end + it "uses regexes to predict up front costs by path if set" do + hash = { + /\A\/files\/\d+\/download/ => 1, + "equation_images\/" => 2 + } + expect(RequestThrottle).to receive(:dynamic_settings).and_return({'up_front_cost_by_path_regex' => hash}) + + expect(@bucket.get_up_front_cost_for_path("/files/1/download?frd=1")).to eq 1 + expect(@bucket.get_up_front_cost_for_path("/equation_images/stuff")).to eq 2 + expect(@bucket.get_up_front_cost_for_path("/somethingelse")).to eq @bucket.up_front_cost + end + it "does nothing if disabled" do expect(RequestThrottle).to receive(:enabled?).twice.and_return(false) expect(@bucket).to receive(:increment).never