Verified Commit d63b42bb authored by Elliot Forbes's avatar Elliot Forbes 2️⃣ Committed by GitLab
Browse files

Merge branch 'feat/rate-limit-metrics' into 'master' 

feat: add Prometheus metrics to rate limit evaluator

See merge request !276

Merged-by: Elliot Forbes's avatarElliot Forbes <eforbes@gitlab.com>
Approved-by: Elliot Forbes's avatarElliot Forbes <eforbes@gitlab.com>
Co-authored-by: Bob Van Landuyt's avatarBob Van Landuyt <bob@gitlab.com>
parents 0d42a179 26d1d9a2

lib/labkit/rate_limit.rb

+1 −0
Original line number Diff line number Diff line
@@ -23,6 +23,7 @@ module Labkit 
    autoload :Rule, "labkit/rate_limit/rule"

    autoload :Evaluator, "labkit/rate_limit/evaluator"

    autoload :Limiter, "labkit/rate_limit/limiter"

    autoload :Metrics, "labkit/rate_limit/metrics"



    class << self

      def configure

lib/labkit/rate_limit/evaluator.rb

+37 −2
Original line number Diff line number Diff line
@@ -24,6 +24,7 @@ module Labkit 
      rescue StandardError => e

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

        # timeout, OOM) not only Redis protocol errors.

        report_error_metrics

        log_error(e, identifier)

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

      end
@@ -34,9 +35,12 @@ module Labkit 
        @rules.each do |rule|

          next unless rule_matches?(rule, identifier)



          return evaluate_rule(rule, identifier)

          result = evaluate_rule(rule, identifier)

          report_matched_metrics(result)

          return result

        end



        report_unmatched_metrics

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

      end
@@ -53,7 +57,10 @@ module Labkit 
        exceeded = count > resolved_limit

        action = exceeded ? rule.action : :allow



        Result.new(matched: true, exceeded: exceeded, action: action, rule: rule)

        Result.new(

          matched: true, exceeded: exceeded, action: action, rule: rule,

          resolved_limit: resolved_limit, resolved_period: resolved_period

        )

      end



      def build_redis_key(rule, identifier)
@@ -99,6 +106,34 @@ module Labkit 
          identifier: identifier&.to_h

        )

      end



      def report_matched_metrics(result)

        Metrics.calls_total.increment(

          rate_limiter: @name,

          rule: result.rule.name,

          action: result.action.to_s

        )

        Metrics.limit_gauge.set(

          { rate_limiter: @name, rule: result.rule.name },

          result.resolved_limit

        )

        Metrics.period_gauge.set(

          { rate_limiter: @name, rule: result.rule.name },

          result.resolved_period

        )

      end



      def report_unmatched_metrics

        Metrics.calls_total.increment(

          rate_limiter: @name,

          rule: "unmatched",

          action: "allow"

        )

      end



      def report_error_metrics

        Metrics.errors_total.increment(rate_limiter: @name)

      end

    end

  end

end

lib/labkit/rate_limit/metrics.rb

0 → 100644
+43 −0
Original line number Diff line number Diff line 
# frozen_string_literal: true



module Labkit

  module RateLimit

    module Metrics

      module_function



      def calls_total

        Labkit::Metrics::Client.counter(

          :gitlab_labkit_rate_limiter_calls_total,

          'Total number of successful rate limit checks',

          { rate_limiter: nil, rule: nil, action: nil }

        )

      end



      def errors_total

        Labkit::Metrics::Client.counter(

          :gitlab_labkit_rate_limiter_errors_total,

          'Total number of rate limit check errors',

          { rate_limiter: nil }

        )

      end



      def limit_gauge

        Labkit::Metrics::Client.gauge(

          :gitlab_labkit_rate_limiter_limit,

          'The configured rate limit threshold',

          { rate_limiter: nil, rule: nil },

          :max

        )

      end



      def period_gauge

        Labkit::Metrics::Client.gauge(

          :gitlab_labkit_rate_limiter_period_seconds,

          'The configured rate limit period in seconds',

          { rate_limiter: nil, rule: nil },

          :max

        )

      end

    end

  end

end

lib/labkit/rate_limit/result.rb

+6 −2
Original line number Diff line number Diff line
@@ -13,8 +13,12 @@ module Labkit 
    #                    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)

    Result = Data.define(:matched, :exceeded, :action, :rule, :error) do

      def initialize(matched:, action:, exceeded: false, rule: nil, error: false)

    # resolved_limit   - the resolved limit value as Integer (nil when matched? is false or error)

    # resolved_period  - the resolved period value as Integer (nil when matched? is false or error)

    Result = Data.define(:matched, :exceeded, :action, :rule, :error, :resolved_limit, :resolved_period) do

      def initialize(

        matched:, action:, exceeded: false, rule: nil, error: false,

        resolved_limit: nil, resolved_period: nil)

        super

      end

spec/labkit/rate_limit/evaluator_spec.rb

+115 −0
Original line number Diff line number Diff line
@@ -168,4 +168,119 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      expect(result.matched?).to be(false)

    end

  end



  describe "metrics emission", :with_metrics_config do

    let(:calls_total) { Labkit::RateLimit::Metrics.calls_total }

    let(:errors_total) { Labkit::RateLimit::Metrics.errors_total }

    let(:limit_gauge) { Labkit::RateLimit::Metrics.limit_gauge }

    let(:period_gauge) { Labkit::RateLimit::Metrics.period_gauge }



    context "when a rule matches and is not exceeded" do

      it "increments calls_total with action: allow" do

        rule = make_rule(name: "api_rule", limit: 100, period: 60)

        allow(redis).to receive(:incr).and_return(1)

        allow(redis).to receive(:expire)



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



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

      end



      it "sets the limit gauge with the resolved value" do

        rule = make_rule(name: "api_rule", limit: 100, period: 60)

        allow(redis).to receive(:incr).and_return(1)

        allow(redis).to receive(:expire)



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



        expect(limit_gauge.get(rate_limiter: "rack_request", rule: "api_rule")).to eq(100.0)

      end



      it "sets the period gauge with the resolved value" do

        rule = make_rule(name: "api_rule", limit: 100, period: 120)

        allow(redis).to receive(:incr).and_return(1)

        allow(redis).to receive(:expire)



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



        expect(period_gauge.get(rate_limiter: "rack_request", rule: "api_rule")).to eq(120.0)

      end

    end



    context "when a rule matches and is exceeded with action: :block" do

      it "increments calls_total with action: block" do

        rule = make_rule(name: "api_rule", limit: 5, action: :block)

        allow(redis).to receive(:incr).and_return(6)

        allow(redis).to receive(:expire)



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



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

      end

    end



    context "when a rule matches and is exceeded with action: :log" do

      it "increments calls_total with action: log" do

        rule = make_rule(name: "api_rule", limit: 5, action: :log)

        allow(redis).to receive(:incr).and_return(6)

        allow(redis).to receive(:expire)



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



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

      end

    end



    context "when no rule matches" do

      it "increments calls_total with rule: unmatched" do

        rule = make_rule(match: { user: 999 })



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



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

      end



      it "does not set the limit or period gauges" do

        rule = make_rule(match: { user: 999 })



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



        expect(limit_gauge.get(rate_limiter: "rack_request", rule: "default")).to eq(0.0)

      end

    end



    context "when Redis fails" do

      it "increments errors_total and does not increment calls_total" do

        rule = make_rule(name: "err_rule")

        allow(redis).to receive(:incr).and_raise(RuntimeError, "connection refused")



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



        expect(errors_total.get(rate_limiter: "rack_request")).to eq(1.0)

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

        expect(calls_total.get(rate_limiter: "rack_request", rule: "err_rule", action: "block")).to eq(0.0)

      end

    end



    context "with callable limit and period" do

      it "sets gauges with the resolved integer values" do

        rule = make_rule(name: "callable_rule", limit: -> { 42 }, period: -> { 300 })

        allow(redis).to receive(:incr).and_return(1)

        allow(redis).to receive(:expire)



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



        expect(limit_gauge.get(rate_limiter: "rack_request", rule: "callable_rule")).to eq(42.0)

        expect(period_gauge.get(rate_limiter: "rack_request", rule: "callable_rule")).to eq(300.0)

      end

    end



    context "with empty rules array" do

      it "increments calls_total with unmatched" do

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



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

      end

    end

  end

end