Commit 41323ba9 authored by Max Woolf's avatar Max Woolf
Browse files

feat(rate_limit): rule_extras for per-request callable context 

Adds an optional rule_extras: keyword to Limiter#check / #peek (and
Evaluator + the convenience wrapper). Hash flows through to one-arity
callables on Rule#limit / Rule#period via arity dispatch in
resolve_value:

  - 0-arity:     val.call            (existing behavior, e.g. global settings)
  - 1+arity:     val.call(rule_extras)  (new — receives the caller-supplied
                                         hash, may be nil)

Backwards compatible: zero-arity callables and plain values are
unchanged. Existing tests pass without modification.

This is the second of two MRs needed for cohort 4 (the first is !292,
which added count_distinct: at the rule level). Cohort 4's per-namespace
limits live on namespace.namespace_settings and can't be resolved by
zero-arity callables — the callable can't reach the request context.
rule_extras: is the channel for that context.

Implements gitlab-com/gl-infra/production-engineering#29071

Co-Authored-By: default avatarClaude Opus 4.7 (1M context) <noreply@anthropic.com>
parent b283c27c

lib/labkit/rate_limit.rb

+5 −2
Original line number Diff line number Diff line
@@ -44,9 +44,12 @@ module Labkit 
      # @param redis [Object, nil] Redis client; falls back to config.redis

      # @param logger [Logger, nil] logger; falls back to config.logger

      # @param cost [Numeric] amount to add to the counter; see Limiter#check

      # @param rule_extras [Hash, nil] per-request context for one-arity

      #   callables on rule limit/period; see Limiter#check

      # @return [Result]

      def check(name:, identifier:, rules:, redis: nil, logger: nil, cost: 1)

        Limiter.new(name: name, rules: rules, redis: redis, logger: logger).check(identifier, cost: cost)

      def check(name:, identifier:, rules:, redis: nil, logger: nil, cost: 1, rule_extras: nil)

        Limiter.new(name: name, rules: rules, redis: redis, logger: logger)

          .check(identifier, cost: cost, rule_extras: rule_extras)

      end

    end

  end

lib/labkit/rate_limit/evaluator.rb

+29 −16
Original line number Diff line number Diff line
@@ -65,8 +65,8 @@ module Labkit 
        @logger = logger

      end



      def check(identifier, cost: 1)

        check_rules(identifier, cost)

      def check(identifier, cost: 1, rule_extras: nil)

        check_rules(identifier, cost, rule_extras)

      rescue StandardError => e

        # Intentionally broad: fail-open applies to any unexpected error (network,

        # timeout, OOM) not only Redis protocol errors.
@@ -78,8 +78,8 @@ module Labkit 
      # Read-without-increment counterpart to {#check}. Same matching and Result

      # shape; the underlying Redis counter is not mutated and the TTL is not

      # extended. A missing Redis key is treated as count=0 (matched, not exceeded).

      def peek(identifier)

        peek_rules(identifier)

      def peek(identifier, rule_extras: nil)

        peek_rules(identifier, rule_extras)

      rescue StandardError => e

        report_error_metrics

        log_error(e, identifier)
@@ -95,7 +95,7 @@ module Labkit 
      # is missing the count_distinct key fail open + log + bump errors_total, and

      # the loop continues to the next rule (the rule is treated as not applicable

      # rather than aborting the whole evaluation).

      def check_rules(identifier, cost)

      def check_rules(identifier, cost, rule_extras)

        @rules.each do |rule|

          next unless rule_matches?(rule, identifier)
@@ -105,7 +105,7 @@ module Labkit 
            next

          end



          result = evaluate_rule(rule, identifier, cost)

          result = evaluate_rule(rule, identifier, cost, rule_extras)

          report_matched_metrics(result)

          return result unless rule.action == :log

        end
@@ -120,12 +120,12 @@ module Labkit 
      # peek does not need the count_distinct identifier key - it reads SCARD on

      # the rule-keyed compound key, which contains the cardinality across all

      # members. So missing-key fail-open does not apply here.

      def peek_rules(identifier)

      def peek_rules(identifier, rule_extras)

        @rules.each do |rule|

          next if rule.action == :log

          next unless rule_matches?(rule, identifier)



          return peek_rule(rule, identifier)

          return peek_rule(rule, identifier, rule_extras)

        end



        Result.new(matched: false, action: :allow)
@@ -140,10 +140,10 @@ module Labkit 
        value.nil? || value.to_s.empty?

      end



      def evaluate_rule(rule, identifier, cost)

      def evaluate_rule(rule, identifier, cost, rule_extras)

        redis_key = build_redis_key(rule, identifier)

        resolved_limit = Integer(resolve_value(rule.limit))

        resolved_period = Integer(resolve_value(rule.period))

        resolved_limit = Integer(resolve_value(rule.limit, rule_extras))

        resolved_period = Integer(resolve_value(rule.period, rule_extras))



        # cost is ignored for count_distinct rules: SADD is binary (a member is

        # either added or not), and the post-add count is SCARD regardless.
@@ -157,10 +157,10 @@ module Labkit 
        build_result(rule, resolved_limit, resolved_period, count, ttl)

      end



      def peek_rule(rule, identifier)

      def peek_rule(rule, identifier, rule_extras)

        redis_key = build_redis_key(rule, identifier)

        resolved_limit = Integer(resolve_value(rule.limit))

        resolved_period = Integer(resolve_value(rule.period))

        resolved_limit = Integer(resolve_value(rule.limit, rule_extras))

        resolved_period = Integer(resolve_value(rule.period, rule_extras))



        count, ttl = rule.count_distinct ? scard_with_ttl(redis_key) : read_with_ttl(redis_key)

        build_result(rule, resolved_limit, resolved_period, count, ttl)
@@ -195,8 +195,21 @@ module Labkit 
        value.to_s

      end



      def resolve_value(val)

        val.respond_to?(:call) ? val.call : val

      # Resolve a limit/period value. Plain values pass through; callables are

      # invoked according to their arity:

      #

      # - Zero-arity callables call with no args (e.g. -> { ApplicationSetting.current.foo }).

      # - One-arity (and variadic) callables receive +rule_extras+, which may be

      #   nil if the caller didn't pass it. One-arity callables must therefore

      #   handle nil - typically with `extras&.[](:key) || default`.

      #

      # This lets rules carry callables that depend on per-request context (e.g.

      # per-namespace settings) without rebuilding the Rule on every call or

      # smuggling state through globals.

      def resolve_value(val, rule_extras = nil)

        return val unless val.respond_to?(:call)



        val.arity.zero? ? val.call : val.call(rule_extras)

      end



      def encode_char_value(value)

lib/labkit/rate_limit/limiter.rb

+9 −4
Original line number Diff line number Diff line
@@ -35,10 +35,14 @@ module Labkit 
      # @param cost [Numeric] amount to add to the counter. Defaults to 1

      #   (count-mode). Pass a non-1 Numeric for cost-mode counters such as

      #   resource-usage limits; passing 0 reads the counter without writing.

      # @param rule_extras [Hash, nil] optional per-request context passed to

      #   one-arity callables on +limit+/+period+. Lets rules resolve dynamic

      #   configuration (e.g. per-namespace settings) without rebuilding the

      #   Rule or doing out-of-band DB queries. Zero-arity callables ignore it.

      # @return [Result]

      def check(identifier, cost: 1)

      def check(identifier, cost: 1, rule_extras: nil)

        id = identifier.is_a?(Identifier) ? identifier : Identifier.new(identifier)

        @evaluator.check(id, cost: cost)

        @evaluator.check(id, cost: cost, rule_extras: rule_extras)

      end



      # Read the current rate-limit state without incrementing the counter.
@@ -54,10 +58,11 @@ module Labkit 
      # open identically to {#check}.

      #

      # @param identifier [Identifier, Hash] caller attributes for this request

      # @param rule_extras [Hash, nil] see {#check}

      # @return [Result]

      def peek(identifier)

      def peek(identifier, rule_extras: nil)

        id = identifier.is_a?(Identifier) ? identifier : Identifier.new(identifier)

        @evaluator.peek(id)

        @evaluator.peek(id, rule_extras: rule_extras)

      end



      private

spec/labkit/rate_limit/evaluator_spec.rb

+102 −0
Original line number Diff line number Diff line
@@ -199,6 +199,108 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
    end

  end



  describe "rule_extras: per-request callable context" do

    it "passes rule_extras to a one-arity limit callable" do

      received = nil

      rule = make_rule(

        name: "extras_limit",

        limit: lambda { |extras|

          received = extras

          10

        },

        period: 60

      )



      evaluator(rules: [rule]).check(identifier, rule_extras: { limit: 99, period: 100 })



      expect(received).to eq({ limit: 99, period: 100 })

    end



    it "passes rule_extras to a one-arity period callable" do

      received = nil

      rule = make_rule(

        name: "extras_period",

        limit: 10,

        period: lambda { |extras|

          received = extras

          60

        }

      )



      evaluator(rules: [rule]).check(identifier, rule_extras: { period: 100 })



      expect(received).to eq({ period: 100 })

    end



    it "passes nil to a one-arity callable when rule_extras is omitted" do

      received = :unset

      rule = make_rule(name: "extras_nil", limit: lambda { |extras|

        received = extras

        10

      })



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



      expect(received).to be_nil

    end



    it "still calls a zero-arity callable with no args even when rule_extras is provided" do

      call_count = 0

      rule = make_rule(name: "zero_arity", limit: lambda {

        call_count += 1

        10

      })



      evaluator(rules: [rule]).check(identifier, rule_extras: { limit: 99 })



      expect(call_count).to eq(1)

    end



    it "resolves limit from rule_extras for an exceeded check" do

      rule = make_rule(

        name: "extras_exceeded", action: :block,

        limit: ->(extras) { extras&.dig(:limit) || 0 },

        period: 60

      )

      # Pre-populate the counter so a low override limit is exceeded immediately.

      raw_redis.set("labkit:rl:rack_request:extras_exceeded:user:42", "5")

      raw_redis.expire("labkit:rl:rack_request:extras_exceeded:user:42", 60)



      result = evaluator(rules: [rule]).check(identifier, rule_extras: { limit: 3 })



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

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

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

    end



    it "passes rule_extras through peek as well" do

      received = nil

      rule = make_rule(

        name: "extras_peek",

        limit: lambda { |extras|

          received = extras

          10

        }

      )



      evaluator(rules: [rule]).peek(identifier, rule_extras: { source: "peek" })



      expect(received).to eq({ source: "peek" })

    end



    it "treats a variadic callable as receiving rule_extras (single arg passed)" do

      received = :unset

      rule = make_rule(name: "variadic", limit: lambda { |*args|

        received = args

        10

      })



      evaluator(rules: [rule]).check(identifier, rule_extras: { limit: 7 })



      expect(received).to eq([{ limit: 7 }])

    end

  end



  describe "EVALSHA / NOSCRIPT fallback (wiring smoke)" do

    # Dispatch correctness (EVALSHA vs EVAL, NOSCRIPT recovery, non-NOSCRIPT

    # propagation) lives in spec/labkit/redis/script_spec.rb. This one test

spec/labkit/rate_limit/limiter_spec.rb

+20 −1
Original line number Diff line number Diff line
@@ -117,7 +117,8 @@ RSpec.describe Labkit::RateLimit::Limiter do 


    it "delegates to the evaluator's peek path" do

      lim = limiter

      expect(lim.instance_variable_get(:@evaluator)).to receive(:peek).with(instance_of(Labkit::RateLimit::Identifier))

      expect(lim.instance_variable_get(:@evaluator))

        .to receive(:peek).with(instance_of(Labkit::RateLimit::Identifier), rule_extras: nil)

      lim.peek({ user: 42 })

    end
@@ -236,6 +237,24 @@ RSpec.describe Labkit::RateLimit::Limiter do 
    end

  end



  describe "rule_extras: kwarg passthrough" do

    it "threads rule_extras: through to the evaluator's check" do

      lim = limiter

      expect(lim.instance_variable_get(:@evaluator))

        .to receive(:check)

        .with(instance_of(Labkit::RateLimit::Identifier), cost: 1, rule_extras: { limit: 7 })

      lim.check({ user: 42 }, rule_extras: { limit: 7 })

    end



    it "threads rule_extras: through to the evaluator's peek" do

      lim = limiter

      expect(lim.instance_variable_get(:@evaluator))

        .to receive(:peek)

        .with(instance_of(Labkit::RateLimit::Identifier), rule_extras: { source: "peek" })

      lim.peek({ user: 42 }, rule_extras: { source: "peek" })

    end

  end



  describe "#check with a count_distinct rule" do

    before do

      # SADD_SCRIPT.eval returns [scard, ttl]