Verified Commit 1e71e212 authored by Sam Wiskow's avatar Sam Wiskow Committed by GitLab
Browse files

feat(rate_limit): :log rules continue evaluation; add :allow rule action (Spec 13)

parent fd51ce57

lib/labkit/rate_limit/evaluator.rb

+6 −1
Original line number Diff line number Diff line
@@ -42,21 +42,26 @@ module Labkit 


      private



      # :log rules are non-terminating: they emit metrics and continue,

      # so a shadow :log rule cannot disable a following :block rule.

      def check_rules(identifier)

        @rules.each do |rule|

          next unless rule_matches?(rule, identifier)



          result = evaluate_rule(rule, identifier)

          report_matched_metrics(result)

          return result

          return result unless rule.action == :log

        end



        report_unmatched_metrics

        Result.new(matched: false, action: :allow)

      end



      # Mirror of check_rules without metrics: peek skips :log rules (their state

      # is unobservable through peek).

      def peek_rules(identifier)

        @rules.each do |rule|

          next if rule.action == :log

          next unless rule_matches?(rule, identifier)



          return peek_rule(rule, identifier)

lib/labkit/rate_limit/metrics.rb

+3 −0
Original line number Diff line number Diff line
@@ -5,6 +5,9 @@ module Labkit 
    module Metrics

      module_function



      # :log rules are non-terminating: a check that matched only :log rules

      # increments calls_total once per matched :log rule AND once with

      # rule="unmatched", action="allow", since no terminating decision was made.

      def calls_total

        Labkit::Metrics::Client.counter(

          :gitlab_labkit_rate_limiter_calls_total,

lib/labkit/rate_limit/result.rb

+2 −2
Original line number Diff line number Diff line
@@ -8,9 +8,9 @@ module Labkit 
    # action    - the outcome: what the caller should do

    #             :block = rule matched, exceeded, rule configured to block

    #             :log   = rule matched, exceeded, rule configured to log only

    #             :allow = rule matched but count within limit, or

    #             :allow = rule matched but count within limit, rule configured to allow,

    #                      no rule matched, or error (fail-open)

    #             The rule's configured action is available via rule.action

    #             The rule's configured action is available via rule.action.

    # rule      - the matched Rule object (nil when matched? is false)

    # error?    - true if Redis was unavailable; result fails open (exceeded? is false)

    # info      - Result::Info with per-window counters; nil when matched? is false or error?

lib/labkit/rate_limit/rule.rb

+4 −2
Original line number Diff line number Diff line
@@ -2,7 +2,7 @@ 


module Labkit

  module RateLimit

    KNOWN_ACTIONS = [:block, :log].freeze

    KNOWN_ACTIONS = %i[block log allow].freeze

    RULE_NAME_PATTERN = /\A[a-z0-9_]+\z/

    RULE_NAME_MAX_LENGTH = 64
@@ -12,7 +12,9 @@ module Labkit 
    #                   the rule to apply; empty hash matches any identifier

    # limit           - request threshold; may be a callable (resolved per check)

    # period          - window in seconds; may be a callable (resolved per check)

    # action          - :block (enforce) or :log (count and log, but do not block)

    # action          - :block (enforce), :log (count and log only, do not block,

    #                   evaluation continues to subsequent rules), or :allow

    #                   (bypass: short-circuit evaluation with no Redis writes)

    # characteristics - identifier keys used to build the compound Redis counter key

    #

    # +name+ must be a lowercase alphanumeric-and-underscore string of at most 64

spec/labkit/rate_limit/evaluator_spec.rb

+228 −8
Original line number Diff line number Diff line
@@ -8,7 +8,7 @@ RSpec.describe Labkit::RateLimit::Evaluator do 


  let(:raw_redis) { instance_double(Redis) }

  let(:redis) { PooledRedis.new(raw_redis) }

  let(:null_logger) { instance_double(Labkit::Logging::JsonLogger, warn: nil) }

  let(:null_logger) { instance_double(Labkit::Logging::JsonLogger, warn: nil, error: nil) }

  let(:identifier) { Labkit::RateLimit::Identifier.new(user: 42, ip: "1.2.3.4") }

  let(:pipe) { instance_double(Redis) }
@@ -137,8 +137,8 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
    end

  end



  describe "Scenario S: Redis unavailable" do

    it "returns error Result and logs a warning when Redis is unavailable" do

  describe "when Redis raises during evaluate_rule" do

    it "fails open with an error Result and logs a warning", :aggregate_failures do

      allow(raw_redis).to receive(:pipelined).and_raise(RuntimeError, "connection refused")

      logger = instance_double(Labkit::Logging::JsonLogger)

      expect(logger).to receive(:warn).with(
@@ -347,7 +347,7 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      evaluator(rules: [rule]).check(identifier)

    end



    it "fails open and logs when pipelined call raises after incr (Scenario N)" do

    it "fails open and logs when pipelined call raises after incr" do

      rule = make_rule(name: "r", limit: 10, period: 60)

      allow(raw_redis).to receive(:pipelined).and_raise(RuntimeError, "connection lost")
@@ -369,7 +369,7 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      expect(result.info).to be_nil

    end



    it "returns nil for new fields on Redis error (Scenario M)" do

    it "returns nil for new fields on Redis error" do

      rule = make_rule(name: "r")

      allow(raw_redis).to receive(:pipelined).and_raise(RuntimeError, "down")
@@ -481,8 +481,8 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      end

    end



    describe "error path" do

      it "returns an error Result and fails open when Redis raises" do

    describe "when peek_rule raises" do

      it "fails open with an error Result and logs a warning", :aggregate_failures do

        allow(raw_redis).to receive(:pipelined).and_raise(RuntimeError, "connection refused")

        logger = instance_double(Labkit::Logging::JsonLogger)

        expect(logger).to receive(:warn).with(
@@ -688,7 +688,7 @@ RSpec.describe Labkit::RateLimit::Evaluator do 


      it "falls through to a subsequent rule when the cast value does not match the first" do

        regex_rule = make_rule(name: "regex_first", match: { user: { re: "^9\\d+$" } }, action: :block)

        eq_rule    = make_rule(name: "eq_second",   match: { user: 42 }, action: :log)

        eq_rule    = make_rule(name: "eq_second",   match: { user: 42 }, action: :block)

        int_id = Labkit::RateLimit::Identifier.new(user: 42)



        expect(pipe).to receive(:incr).and_return(nil)
@@ -730,4 +730,224 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      expect(evaluator(rules: [rule]).check(id).matched?).to be(false)

    end

  end



  # Multi-rule evaluation behaviour. Initially specified in production-engineering#28890.

  describe "multi-rule evaluation with :log/:block/:allow actions", :with_metrics_config do

    let(:metrics) { Labkit::RateLimit::Metrics }



    describe "with a single :log rule that matches and exceeds its limit" do

      it "increments the Redis counter, emits both the :log and unmatched metrics, and returns the fall-through", :aggregate_failures do

        rule = make_rule(name: "log_rule_a", action: :log, limit: 1, characteristics: [:user])

        allow(raw_redis).to receive(:pipelined).and_return([2, 30])



        result = evaluator(rules: [rule]).check(identifier)



        expect(result).to eq(Labkit::RateLimit::Result.new(matched: false, action: :allow))

        expect(metrics.calls_total.get(rate_limiter: "rack_request", rule: "log_rule_a", action: "log")).to eq(1.0)

        expect(metrics.calls_total.get(rate_limiter: "rack_request", rule: "unmatched", action: "allow")).to eq(1.0)

      end

    end



    describe "with a :log rule preceding a :block rule, both matching, neither exceeded" do

      it "increments both Redis counters and returns the last evaluated rule's resolved_limit", :aggregate_failures do

        log_r = make_rule(name: "log_rule_b", action: :log, limit: 1000, characteristics: [:user])

        block_r = make_rule(name: "block_rule_b", action: :block, limit: 50, characteristics: [:user])



        call_count = 0

        allow(raw_redis).to receive(:pipelined) do

          call_count += 1

          [call_count, 30]

        end



        result = evaluator(rules: [log_r, block_r]).check(identifier)



        expect(result.matched?).to be(true)

        expect(result.action).to eq(:allow)

        expect(result.rule).to eq(block_r)

        expect(result.info.resolved_limit).to eq(50)



        expect(metrics.calls_total.get(rate_limiter: "rack_request", rule: "log_rule_b", action: "allow")).to eq(1.0)

        expect(metrics.calls_total.get(rate_limiter: "rack_request", rule: "block_rule_b", action: "allow")).to eq(1.0)

      end

    end



    describe "with a :log rule preceding a :block rule, only :block exceeded" do

      it "increments both Redis counters and returns the :block result", :aggregate_failures do

        log_r = make_rule(name: "log_rule_c", action: :log, limit: 100, characteristics: [:user])

        block_r = make_rule(name: "block_rule_c", action: :block, limit: 1, characteristics: [:user])



        results = [[3, 30], [2, 30]]

        allow(raw_redis).to receive(:pipelined) { results.shift }



        result = evaluator(rules: [log_r, block_r]).check(identifier)



        expect(result.matched?).to be(true)

        expect(result.action).to eq(:block)

        expect(result.rule).to eq(block_r)



        expect(metrics.calls_total.get(rate_limiter: "rack_request", rule: "log_rule_c", action: "allow")).to eq(1.0)

        expect(metrics.calls_total.get(rate_limiter: "rack_request", rule: "block_rule_c", action: "block")).to eq(1.0)

      end

    end



    describe "with a :block rule preceding a :log rule, :block matching and exceeded" do

      it "returns on the :block match without evaluating the trailing :log rule", :aggregate_failures do

        block_r = make_rule(name: "block_rule_d", action: :block, limit: 1, characteristics: [:user])

        log_r = make_rule(name: "log_rule_d", action: :log, limit: 1, characteristics: [:user])



        expect(raw_redis).to receive(:pipelined).once.and_return([2, 30])



        result = evaluator(rules: [block_r, log_r]).check(identifier)



        expect(result.action).to eq(:block)

        expect(result.rule).to eq(block_r)

      end

    end



    describe "with an :allow rule whose match: is satisfied" do

      it "evaluates against Redis, returns :allow with populated info, and emits headers", :aggregate_failures do

        allow_r = make_rule(name: "allow_rule_e", action: :allow, match: { bypass: true }, limit: 5, period: 60)

        bypass_id = Labkit::RateLimit::Identifier.new(bypass: true)



        allow(raw_redis).to receive(:pipelined).and_return([1, 30])



        result = evaluator(rules: [allow_r]).check(bypass_id)



        expect(result.matched?).to be(true)

        expect(result.action).to eq(:allow)

        expect(result.rule).to eq(allow_r)

        expect(result.exceeded?).to be(false)

        expect(result.error?).to be(false)

        expect(result.info.count).to eq(1)

        expect(result.info.resolved_limit).to eq(5)

        expect(result.to_response_headers).to include(

          "RateLimit-Limit" => "5",

          "RateLimit-Remaining" => "4"

        )



        expect(metrics.calls_total.get(rate_limiter: "rack_request", rule: "allow_rule_e", action: "allow")).to eq(1.0)

      end

    end



    describe "with an :allow rule preceding a :block rule, identifier matches the :allow" do

      it "increments the :allow rule's counter and never evaluates the :block rule", :aggregate_failures do

        allow_r = make_rule(name: "allow_rule_f", action: :allow, match: { bypass: true }, limit: 1, period: 60)

        block_r = make_rule(name: "block_rule_f", action: :block, limit: 1, characteristics: [:user])

        bypass_id = Labkit::RateLimit::Identifier.new(bypass: true, user: 42)



        expect(raw_redis).to receive(:pipelined).once.and_return([1, 30])



        result = evaluator(rules: [allow_r, block_r]).check(bypass_id)



        expect(result.action).to eq(:allow)

        expect(result.rule).to eq(allow_r)

      end

    end



    describe "when no rule's match: predicate is satisfied" do

      it "increments the unmatched metric and returns the fall-through Result", :aggregate_failures do

        rule = make_rule(name: "unmatched_g", action: :block, limit: 1, match: { user: 999 })



        expect(raw_redis).not_to receive(:pipelined)



        result = evaluator(rules: [rule]).check(identifier)



        expect(result).to eq(Labkit::RateLimit::Result.new(matched: false, action: :allow))

        expect(metrics.calls_total.get(rate_limiter: "rack_request", rule: "unmatched", action: "allow")).to eq(1.0)

      end

    end



    describe "with two :log rules (distinct names), both matching" do

      it "increments both Redis counters and emits the unmatched metric (no terminating decision was made)", :aggregate_failures do

        log_a = make_rule(name: "log_rule_h_a", action: :log, limit: 100, characteristics: [:user])

        log_b = make_rule(name: "log_rule_h_b", action: :log, limit: 100, characteristics: [:user])



        allow(raw_redis).to receive(:pipelined).and_return([1, 30])



        result = evaluator(rules: [log_a, log_b]).check(identifier)



        expect(result).to eq(Labkit::RateLimit::Result.new(matched: false, action: :allow))

        expect(metrics.calls_total.get(rate_limiter: "rack_request", rule: "log_rule_h_a", action: "allow")).to eq(1.0)

        expect(metrics.calls_total.get(rate_limiter: "rack_request", rule: "log_rule_h_b", action: "allow")).to eq(1.0)

        expect(metrics.calls_total.get(rate_limiter: "rack_request", rule: "unmatched", action: "allow")).to eq(1.0)

      end

    end



    describe "with an :allow rule whose match: is NOT satisfied" do

      it "skips the :allow rule and evaluates the subsequent :block rule", :aggregate_failures do

        allow_r = make_rule(name: "allow_rule_m", action: :allow, match: { bypass: true }, limit: 1, period: 60)

        block_r = make_rule(name: "block_rule_m", action: :block, limit: 1, characteristics: [:user])

        non_bypass_id = Labkit::RateLimit::Identifier.new(user: 1)



        expect(raw_redis).to receive(:pipelined).once.and_return([2, 30])



        result = evaluator(rules: [allow_r, block_r]).check(non_bypass_id)



        expect(result.action).to eq(:block)

        expect(result.rule).to eq(block_r)

        expect(metrics.calls_total.get(rate_limiter: "rack_request", rule: "allow_rule_m", action: "allow")).to eq(0.0)

      end

    end



    describe "with an :allow rule using match: {} as a universal bypass" do

      it "matches every identifier and never falls through to a subsequent :block rule", :aggregate_failures do

        allow_r = make_rule(name: "allow_rule_n", action: :allow, match: {}, limit: 1, period: 60)

        block_r = make_rule(name: "block_rule_n", action: :block, limit: 1, characteristics: [:user])



        call_count = 0

        allow(raw_redis).to receive(:pipelined) do

          call_count += 1

          [call_count, 30]

        end



        ev = evaluator(rules: [allow_r, block_r])

        results = Array.new(10) do |i|

          ev.check(Labkit::RateLimit::Identifier.new(user: i))

        end



        expect(results).to all(have_attributes(matched?: true, action: :allow, rule: allow_r))

        expect(call_count).to eq(10)

        expect(metrics.calls_total.get(rate_limiter: "rack_request", rule: "allow_rule_n", action: "allow")).to eq(10.0)

      end

    end



    describe "peek with a :log rule preceding a :block rule" do

      it "skips the :log rule and returns the :block rule's read-only state with no metrics", :aggregate_failures do

        log_r = make_rule(name: "log_rule_k", action: :log, limit: 100, characteristics: [:user])

        block_r = make_rule(name: "block_rule_k", action: :block, limit: 5, characteristics: [:user])



        allow(raw_redis).to receive(:pipelined).and_yield(pipe).and_return(["2", 30])

        allow(pipe).to receive(:get)

        allow(pipe).to receive(:ttl)

        expect(raw_redis).not_to receive(:expire)



        result = evaluator(rules: [log_r, block_r]).peek(identifier)



        expect(result.matched?).to be(true)

        expect(result.rule).to eq(block_r)

        expect(metrics.calls_total.get(rate_limiter: "rack_request", rule: "log_rule_k", action: "log")).to eq(0.0)

        expect(metrics.calls_total.get(rate_limiter: "rack_request", rule: "block_rule_k", action: "allow")).to eq(0.0)

      end

    end



    describe "peek with an :allow rule whose match: is satisfied" do

      it "reads through Redis on the :allow rule and never evaluates the :block rule, no metrics", :aggregate_failures do

        allow_r = make_rule(name: "allow_rule_l", action: :allow, match: { bypass: true }, limit: 5, period: 60)

        block_r = make_rule(name: "block_rule_l", action: :block, limit: 1, characteristics: [:user])

        bypass_id = Labkit::RateLimit::Identifier.new(bypass: true, user: 42)



        allow(pipe).to receive_messages(get: nil, ttl: nil)

        expect(raw_redis).to receive(:pipelined).once.and_yield(pipe).and_return([nil, -2])



        result = evaluator(rules: [allow_r, block_r]).peek(bypass_id)



        expect(result.matched?).to be(true)

        expect(result.action).to eq(:allow)

        expect(result.rule).to eq(allow_r)

        expect(result.info.count).to eq(0)

        expect(result.info.resolved_limit).to eq(5)

        expect(metrics.calls_total.get(rate_limiter: "rack_request", rule: "allow_rule_l", action: "allow")).to eq(0.0)

      end

    end

  end

end