gitlab-labkit 1.22.0 → 2.0.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 88d5f5a7cd756153b8c69edbf0a887b0d7fbd98c0eab2828c3150cc2fafa1fac
-  data.tar.gz: 0bf07e6f00b14c3a31ba3dab3b3bbd7cd72af81b61dd66992e9d9ef6bd73d6da
+  metadata.gz: 0c6ab89c2e34ad08721cee1e920601a8d6d039bfe8a6552777021714f17d0c74
+  data.tar.gz: 6b9e36ac9ce69866890d6df28132d8abd280cf7d3b4ab7b748470864dbf8a70e
 SHA512:
-  metadata.gz: 737e6ee6b0fc00f716cf3aba4f94bd02c65c01dc5436ca8a4b99ce64ce092efe8bd39af8fe7d6d68ce5e8c2dcd1ba7c991d49b4c0c742c17a4e4ea1d48e50f1b
-  data.tar.gz: 17db3fe540d65f81193adfc8f86d9e984761b76b17f639bad7cd73df46206035690b1c3466bc8d54c26b7c7a3f043c6158edb657137c67b9318098b0eae61f48
+  metadata.gz: af23293cd5d9afe2dd3ba93ca4cd98a6ae5178c04079069e60defff46652266158fc7d74fe335d9345391ee11e89602e0ebf09d1d536fc0f0140426676a23a16
+  data.tar.gz: 3ae857f6940bf2b871b123d88c3caf3e12589efe3b2a3733cb723b73cdca347c3bfa8506ecf3c931d0d83d7dd99882bf25a9abe7efa3e85f76e4958b111be312

data/lib/gitlab-labkit.rb

@@ -25,6 +25,7 @@ module Labkit
   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:

data/lib/labkit/rate_limit/evaluator.rb

@@ -12,6 +12,32 @@ 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 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
@@ -19,8 +45,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 +70,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 +100,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 +159,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 +177,24 @@ 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
 
+      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",

data/lib/labkit/rate_limit/limiter.rb

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

data/lib/labkit/rate_limit/result.rb

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

data/lib/labkit/rate_limit.rb

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

data/lib/labkit/redis/script.rb

@@ -0,0 +1,43 @@
+# frozen_string_literal: true
+
+require "openssl"
+require "redis"
+
+module Labkit
+  module Redis
+    # Wraps a Lua script for EVALSHA-with-NOSCRIPT-fallback execution.
+    # The SHA is computed once at construction. Redis caches the script
+    # body the first time EVAL is invoked; subsequent EVALSHA calls hit
+    # that cache. SCRIPT FLUSH or a Redis restart drops the cache; the
+    # NOSCRIPT recovery re-ships the body and re-populates it.
+    #
+    # @example
+    #   SCRIPT = Labkit::Redis::Script.new(<<~LUA)
+    #     return redis.call('INCRBY', KEYS[1], ARGV[1])
+    #   LUA
+    #
+    #   pool.with { |conn| SCRIPT.eval(conn, keys: ["counter"], argv: [1]) }
+    class Script
+      attr_reader :body, :sha
+
+      def initialize(body)
+        @body = body.freeze
+        # SHA1 is mandated by the Redis EVALSHA wire protocol, not a discretionary hash choice.
+        @sha = OpenSSL::Digest::SHA1.hexdigest(body).freeze # rubocop:disable Fips/SHA1
+        freeze
+      end
+
+      # @param conn a Redis client (the connection checked out of a pool)
+      # @param keys [Array] KEYS arguments to the Lua script
+      # @param argv [Array] ARGV arguments to the Lua script
+      # @return the script's return value
+      def eval(conn, keys:, argv:)
+        conn.evalsha(@sha, keys: keys, argv: argv)
+      rescue ::Redis::CommandError => e
+        raise unless e.message.start_with?("NOSCRIPT")
+
+        conn.eval(@body, keys: keys, argv: argv)
+      end
+    end
+  end
+end

data/lib/labkit/redis.rb

@@ -0,0 +1,8 @@
+# frozen_string_literal: true
+
+module Labkit
+  # Redis utilities shared across Labkit (script execution, key helpers, etc.).
+  module Redis
+    autoload :Script, "labkit/redis/script"
+  end
+end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: gitlab-labkit
 version: !ruby/object:Gem::Version
-  version: 1.22.0
+  version: 2.0.0
 platform: ruby
 authors:
 - Andrew Newdigate
@@ -609,6 +609,8 @@ files:
 - lib/labkit/rate_limit/metrics.rb
 - lib/labkit/rate_limit/result.rb
 - lib/labkit/rate_limit/rule.rb
+- lib/labkit/redis.rb
+- lib/labkit/redis/script.rb
 - lib/labkit/rspec/README.md
 - lib/labkit/rspec/matchers.rb
 - lib/labkit/rspec/matchers/user_experience_matchers.rb