Commit 99204787 authored by Max Woolf's avatar Max Woolf
Browse files

feat(rate_limit)!: cost-aware atomic increment via Lua script 

Consolidate INCR/EXPIRE/TTL into a single EVAL call and extend the
counter primitive to INCRBYFLOAT with a per-call cost parameter. The
script runs atomically from Redis's perspective, eliminating the
latent race between increment and EXPIRE in the prior pipelined
implementation.

Adds `cost:` keyword to Limiter#check and Labkit::RateLimit.check;
default 1 preserves existing call-site semantics. cost=0
short-circuits to a GET so resource-usage callers that observed zero
usage do not allocate a Redis key. ttl_before < 0 self-heals keys
left without expiry by any prior bug. EVALSHA with NOSCRIPT fallback
covers Redis restarts that wipe the script cache.

spec/labkit/rate_limit/evaluator_spec.rb is rewritten from mock-based
to TestRedis-backed, matching rate_limit_spec.rb's idiom now that real
Redis test infrastructure is available. Recommend reading the new
file as a standalone, not as a line-by-line diff.

BREAKING CHANGE: Result::Info#count is now Float (previously Integer).
INCRBYFLOAT returns a string-encoded number that the evaluator parses
as Float for both integer-valued and fractional counters.
to_response_headers coerces remaining/limit back to Integer per the
RateLimit header spec. Numeric comparisons against Integer thresholds
work unchanged via Ruby coercion; pattern-matching on Integer type
will break and must be updated.

Refs gitlab-com/gl-infra/production-engineering#28827
parent a51a33d4

lib/labkit/rate_limit.rb

+3 −2
Original line number Diff line number Diff line
@@ -43,9 +43,10 @@ module Labkit 
      # @param rules [Array<Rule>] ordered list of rules (first match wins)

      # @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

      # @return [Result]

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

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

      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)

      end

    end

  end

lib/labkit/rate_limit/evaluator.rb

+61 −17
Original line number Diff line number Diff line
@@ -12,6 +12,35 @@ module Labkit 
      CHAR_VALUE_MAX_LENGTH = 200

      MISSING_VALUE_SENTINEL = "_unknown_"



      # Atomic increment-with-TTL Lua script. The whole script runs as one

      # operation from Redis's perspective, so there is no window between

      # the increment and EXPIRE that can leak a key without TTL.

      #

      # INCRBYFLOAT serves both count-mode (cost=1, equivalent to INCR for

      # integer-encoded keys) and cost-mode callers, so a single script

      # handles every rule shape.

      #

      # - cost=0 short-circuits to GET so resource-usage callers that

      #   observed zero usage do not allocate a Redis key.

      # - ttl_before < 0 covers TTL=-2 (key missing) and TTL=-1 (no expiry).

      #   The -1 case shouldn't arise with the atomic script, but

      #   self-healing recovers keys left without TTL by any prior bug.

      INCR_SCRIPT = <<~LUA.freeze

        local cost = tonumber(ARGV[2])

        local ttl_before = redis.call('TTL', KEYS[1])

        local count

        if cost == 0 then

          count = redis.call('GET', KEYS[1]) or '0'

        else

          count = redis.call('INCRBYFLOAT', KEYS[1], cost)

          if ttl_before < 0 then

            redis.call('EXPIRE', KEYS[1], ARGV[1])

          end

        end

        return {count, redis.call('TTL', KEYS[1])}

      LUA

      INCR_SCRIPT_SHA = OpenSSL::Digest::SHA1.hexdigest(INCR_SCRIPT).freeze



      def initialize(name:, rules:, redis:, logger:)

        @name   = name

        @rules  = rules
@@ -19,8 +48,8 @@ module Labkit 
        @logger = logger

      end



      def check(identifier)

        check_rules(identifier)

      def check(identifier, cost: 1)

        check_rules(identifier, cost)

      rescue StandardError => e

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

        # timeout, OOM) not only Redis protocol errors.
@@ -44,11 +73,11 @@ module Labkit 


      # :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)

      def check_rules(identifier, cost)

        @rules.each do |rule|

          next unless rule_matches?(rule, identifier)



          result = evaluate_rule(rule, identifier)

          result = evaluate_rule(rule, identifier, cost)

          report_matched_metrics(result)

          return result unless rule.action == :log

        end
@@ -74,12 +103,12 @@ module Labkit 
        rule.match.all? { |key, matcher| matcher.match?(identifier[key]) }

      end



      def evaluate_rule(rule, identifier)

      def evaluate_rule(rule, identifier, cost)

        redis_key = build_redis_key(rule, identifier)

        resolved_limit = Integer(resolve_value(rule.limit))

        resolved_period = Integer(resolve_value(rule.period))



        count, ttl = incr_with_ttl(redis_key, resolved_period)

        count, ttl = incr_with_ttl(redis_key, resolved_period, cost)

        build_result(rule, resolved_limit, resolved_period, count, ttl)

      end
@@ -133,17 +162,17 @@ module Labkit 
        end

      end



      # Pipelines INCR and TTL so both are fetched in a single round-trip.

      # EXPIRE follows as a separate call only on first write (count == 1).

      # On first write TTL will be -1 (expiry not yet set); callers fall back to period.

      def incr_with_ttl(redis_key, period)

      # Atomically increments the counter by `cost`, sets the TTL on first

      # write, and reads back the post-increment TTL, all in one Redis

      # operation via Lua. See INCR_SCRIPT for the script body.

      #

      # count is parsed as Float because INCRBYFLOAT returns a string-encoded

      # number; the Float is integer-valued when cost is 1, fractional

      # otherwise.

      def incr_with_ttl(redis_key, period, cost)

        @redis.with do |conn|

          count, ttl = conn.pipelined do |pipe|

            pipe.incr(redis_key)

            pipe.ttl(redis_key)

          end

          conn.expire(redis_key, period) if count == 1

          [count, ttl]

          raw_count, ttl = eval_incr_script(conn, redis_key, period, cost)

          [Float(raw_count), ttl]

        end

      end
@@ -151,17 +180,32 @@ module Labkit 
      # A missing key (GET => nil, TTL => -2) is reported as count=0; the

      # build_result fallback then derives reset_at from the rule period

      # since there is no Redis-side window to read.

      #

      # Float parsing accepts both INCR-stored ("5") and INCRBYFLOAT-stored

      # ("5.7") values uniformly.

      def read_with_ttl(redis_key)

        @redis.with do |conn|

          raw_count, ttl = conn.pipelined do |pipe|

            pipe.get(redis_key)

            pipe.ttl(redis_key)

          end

          count = raw_count.nil? ? 0 : Integer(raw_count)

          count = raw_count.nil? ? 0.0 : Float(raw_count)

          [count, ttl]

        end

      end



      # EVALSHA with NOSCRIPT fallback. The fallback EVAL ships the script

      # body, which Redis caches; subsequent calls hit EVALSHA again. Redis

      # may drop the script cache on restart or via SCRIPT FLUSH, so the

      # fallback is part of the steady-state contract, not a one-off.

      def eval_incr_script(conn, redis_key, period, cost)

        conn.evalsha(INCR_SCRIPT_SHA, keys: [redis_key], argv: [period, cost])

      rescue ::Redis::CommandError => e

        raise unless e.message.start_with?("NOSCRIPT")



        conn.eval(INCR_SCRIPT, keys: [redis_key], argv: [period, cost])

      end



      def log_error(error, identifier)

        @logger.warn(

          message: "rate_limit_error",

lib/labkit/rate_limit/limiter.rb

+5 −2
Original line number Diff line number Diff line
@@ -32,10 +32,13 @@ module Labkit 
      end



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

      # @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.

      # @return [Result]

      def check(identifier)

      def check(identifier, cost: 1)

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

        @evaluator.check(id)

        @evaluator.check(id, cost: cost)

      end



      # Read the current rate-limit state without incrementing the counter.

lib/labkit/rate_limit/result.rb

+10 −5
Original line number Diff line number Diff line
@@ -33,13 +33,14 @@ module Labkit 


      # Returns RFC-compliant rate limit response headers, or {} when no rule matched or an error occurred.

      # Keys: RateLimit-Limit, RateLimit-Remaining, RateLimit-Reset (Unix timestamp).

      # reset_at is advisory only - derived from a pipelined redis.ttl call, not fully atomic.

      # remaining is coerced to Integer for header output even when info.remaining is fractional;

      # the RateLimit header spec requires integer values.

      def to_response_headers

        return {} unless matched? && !error? && info



        {

          "RateLimit-Limit" => info.resolved_limit.to_s,

          "RateLimit-Remaining" => info.remaining.to_s,

          "RateLimit-Limit" => info.resolved_limit.to_i.to_s,

          "RateLimit-Remaining" => info.remaining.to_i.to_s,

          "RateLimit-Reset" => info.reset_at.to_i.to_s

        }

      end
@@ -48,8 +49,12 @@ module Labkit 
    # Per-window counter data attached to a matched Result.

    # resolved_limit  - the evaluated limit Integer for this rule

    # resolved_period - the evaluated period Integer (seconds) for this rule

    # count           - the raw INCR value; useful for utilization-ratio metrics

    # remaining       - requests remaining before the limit is hit (floors at 0)

    # count           - the post-increment counter value as a Float; integer-valued

    #                   for default cost=1 callers, fractional for cost-mode callers.

    #                   Pre-2.x releases exposed this as Integer; see the migration

    #                   note in the cost-aware Lua script change.

    # remaining       - requests remaining before the limit is hit (floors at 0).

    #                   Inherits Float typing from count when count is fractional.

    # reset_at        - best-effort UTC Time when the counter window resets

    Result::Info = Data.define(:resolved_limit, :resolved_period, :count, :remaining, :reset_at)

  end

spec/labkit/rate_limit/configuration_spec.rb

+4 −6
Original line number Diff line number Diff line
@@ -36,14 +36,13 @@ RSpec.describe "Labkit::RateLimit.configure (Scenario M & N)" do 
      c.logger = my_logger

    end



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

    allow(raw_redis).to receive(:expire)

    allow(raw_redis).to receive(:evalsha).and_return(["1", 55])

    rule = Labkit::RateLimit::Rule.new(name: "test", limit: 10, period: 60, characteristics: [:user])

    limiter = Labkit::RateLimit::Limiter.new(name: "test_limiter", rules: [rule])

    result = limiter.check({ user: 42 })



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

    expect(raw_redis).to have_received(:pipelined)

    expect(raw_redis).to have_received(:evalsha)

  end



  # Scenario N: explicit DI kwargs override configure block
@@ -58,8 +57,7 @@ RSpec.describe "Labkit::RateLimit.configure (Scenario M & N)" do 
    raw_override_redis = instance_double(Redis, "override_redis")

    override_redis = PooledRedis.new(raw_override_redis)

    override_logger = instance_double(Logger, "override_logger", info: nil, warn: nil)

    allow(raw_override_redis).to receive(:pipelined).and_return([1, 55])

    allow(raw_override_redis).to receive(:expire)

    allow(raw_override_redis).to receive(:evalsha).and_return(["1", 55])

    rule = Labkit::RateLimit::Rule.new(name: "test", limit: 10, period: 60, characteristics: [:user])

    limiter = Labkit::RateLimit::Limiter.new(

      name: "test_limiter", rules: [rule],
@@ -67,6 +65,6 @@ RSpec.describe "Labkit::RateLimit.configure (Scenario M & N)" do 
    )

    limiter.check({ user: 42 })



    expect(raw_override_redis).to have_received(:pipelined)

    expect(raw_override_redis).to have_received(:evalsha)

  end

end