Merge branch 'rate-limit/lua-cost-aware-incr' into 'master'

  autoload :Middleware, "labkit/middleware"

  autoload :Fields, "labkit/fields"

  autoload :RateLimit, "labkit/rate_limit"

  autoload :Redis, "labkit/redis"

  # Publishers to publish notifications whenever a HTTP reqeust is made.

  # A broadcasted notification's payload in topic "request.external_http" includes:

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

      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 also flows through INCRBYFLOAT;

      # Redis treats the result as a no-op on the stored value while

      # still observing the post-state count and TTL we return.

      #

      # 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 = Labkit::Redis::Script.new(<<~LUA)

        local ttl = ARGV[1]

        local cost = tonumber(ARGV[2])

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

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

        if ttl_before < 0 then

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

        end

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

      LUA

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

        @name   = name

        @rules  = rules

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

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

        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

        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

      # 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

      def eval_incr_script(conn, redis_key, period, cost)

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

      end

      def log_error(error, identifier)

        @logger.warn(

          message: "rate_limit_error",

      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.

      # 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

    # 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