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

refactor(rate_limit): extract EVALSHA caching into Labkit::Redis::Script 

The Evaluator's EVALSHA-with-NOSCRIPT-fallback dispatch is a generic Lua
primitive that has no rate-limit-specific knowledge. Pull it out into
Labkit::Redis::Script so:

  - Other Labkit callers can reuse the pattern (the monolith currently
    has unguarded eval/evalsha call sites that could benefit).
  - EVALSHA-vs-EVAL dispatch correctness can be tested in isolation,
    via a spy connection that counts calls — the previous coverage in
    evaluator_spec could only observe \"didn't raise\", which doesn't
    distinguish a cached hit from a NOSCRIPT recovery.

Behavior is unchanged: the Lua body, the SHA1 digest, the NOSCRIPT
detection on Redis::CommandError.message, and the EVAL fallback shape
all carry over verbatim. The Evaluator keeps a single Script instance
on the class as INCR_SCRIPT (per-process), and the per-call eval just
threads the checked-out connection through to Script#eval.

The new spec covers happy path + SCRIPT FLUSH recovery against real
Redis, plus EVALSHA/EVAL dispatch counts and WRONGTYPE propagation via
a spy conn. The two equivalent tests in evaluator_spec drop in favor
of one wiring-smoke test that the Evaluator still survives a flush.

Addresses !291 review threads
5e1f3e92 and 3058e1cd.
parent 3d256250

lib/gitlab-labkit.rb

+1 −0
Original line number Diff line number Diff line
@@ -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:

lib/labkit/rate_limit/evaluator.rb

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

      INCR_SCRIPT = Labkit::Redis::Script.new(<<~LUA)

        local cost = tonumber(ARGV[2])

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

        local count
@@ -39,8 +39,6 @@ module Labkit 
        end

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

      LUA

      # SHA1 is mandated by the Redis EVALSHA wire protocol, not a discretionary hash choice.

      INCR_SCRIPT_SHA = OpenSSL::Digest::SHA1.hexdigest(INCR_SCRIPT).freeze # rubocop:disable Fips/SHA1



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

        @name   = name
@@ -195,16 +193,8 @@ module Labkit 
        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])

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

      end



      def log_error(error, identifier)

lib/labkit/redis.rb

0 → 100644
+8 −0
Original line number Diff line number Diff line 
# 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

lib/labkit/redis/script.rb

0 → 100644
+43 −0
Original line number Diff line number Diff line 
# 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

spec/labkit/rate_limit/evaluator_spec.rb

+6 −40
Original line number Diff line number Diff line
@@ -200,11 +200,13 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
    end

  end



  describe "EVALSHA / NOSCRIPT fallback" do

    let(:rule) { make_rule(name: "noscript_rule") }

    let(:key) { "labkit:rl:rack_request:noscript_rule:user:42" }



  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

    # proves the Evaluator is wired to the Script correctly.

    it "recovers via EVAL when Redis's script cache has been flushed" do

      rule = make_rule(name: "noscript_rule")

      key = "labkit:rl:rack_request:noscript_rule:user:42"

      raw_redis.script(:flush)



      result = evaluator(rules: [rule]).check(identifier)
@@ -212,42 +214,6 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      expect(result.matched?).to be(true)

      expect(stored_count(key)).to eq(1.0)

    end



    it "uses EVALSHA on the subsequent call after a fallback (script is recached)" do

      raw_redis.script(:flush)

      ev = evaluator(rules: [rule])

      ev.check(identifier)



      # Now the script should be cached again; this call hits EVALSHA without fallback.

      expect { ev.check(identifier) }.not_to raise_error

      expect(stored_count(key)).to eq(2.0)

    end



    it "propagates non-NOSCRIPT Redis errors instead of swallowing them" do

      faulty_redis = Class.new do

        def with

          yield self

        end



        def evalsha(*)

          raise Redis::CommandError, "WRONGTYPE Operation against a key holding the wrong kind of value"

        end



        def get(*)

          nil

        end



        def pipelined

          yield self

          [nil, -2]

        end

      end.new



      ev = described_class.new(name: "rack_request", rules: [rule], redis: faulty_redis, logger: null_logger)

      result = ev.check(identifier)



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

    end

  end



  describe "Peek (read-only)" do