Verified Commit 0f417f20 authored by Max Woolf's avatar Max Woolf Committed by GitLab
Browse files

Merge branch 'bvl/redis-pool-support' into 'master' 

feat: require Redis pool with `.with` interface

See merge request !278

Merged-by: Max Woolf's avatarMax Woolf <mwoolf@gitlab.com>
Approved-by: Max Woolf's avatarMax Woolf <mwoolf@gitlab.com>
Co-authored-by: Bob Van Landuyt's avatarBob Van Landuyt <bob@gitlab.com>
parents d63b42bb b8f050d2

lib/labkit/rate_limit.rb

+1 −1
Original line number Diff line number Diff line
@@ -6,7 +6,7 @@ module Labkit 
  #

  # @example Configuration (e.g. in a Rails initializer)

  #   Labkit::RateLimit.configure do |c|

  #     c.redis  = Redis.current

  #     c.redis  = ConnectionPool.new { Redis.new }  # must respond to .with { |conn| }

  #     c.logger = Labkit::Logging::JsonLogger.new($stdout)

  #   end

  #

lib/labkit/rate_limit/evaluator.rb

+4 −2
Original line number Diff line number Diff line
@@ -92,11 +92,13 @@ module Labkit 
      end



      def incr_with_ttl(redis_key, period)

        count = @redis.incr(redis_key)

        @redis.with do |conn|

          count = conn.incr(redis_key)

          # Set expiry only on first write to avoid resetting TTL on each call

        @redis.expire(redis_key, period) if count == 1

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

          count

        end

      end



      def log_error(error, identifier)

        @logger.warn(

spec/labkit/rate_limit/configuration_spec.rb

+13 −9
Original line number Diff line number Diff line
@@ -25,7 +25,8 @@ RSpec.describe "Labkit::RateLimit.configure (Scenario M & N)" do 
    Labkit::RateLimit.instance_variable_set(:@config, nil)

  end



  let(:my_redis) { instance_double(Redis) }

  let(:raw_redis) { instance_double(Redis) }

  let(:my_redis) { PooledRedis.new(raw_redis) }

  let(:my_logger) { instance_double(Logger, info: nil, warn: nil) }



  # Scenario M: configure block wires redis and logger
@@ -35,28 +36,31 @@ RSpec.describe "Labkit::RateLimit.configure (Scenario M & N)" do 
      c.logger = my_logger

    end



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

    allow(my_redis).to receive(:expire)

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

    allow(raw_redis).to receive(:expire)



    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(my_redis).to have_received(:incr)

    expect(raw_redis).to have_received(:incr)

  end



  # Scenario N: explicit DI kwargs override configure block

  it "uses explicit redis/logger over the global configure block" do

    global_redis = PooledRedis.new(instance_double(Redis, "global_redis"))



    Labkit::RateLimit.configure do |c|

      c.redis = instance_double(Redis, "global_redis")

      c.redis = global_redis

      c.logger = instance_double(Logger, "global_logger")

    end



    override_redis = instance_double(Redis, "override_redis")

    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(override_redis).to receive(:incr).and_return(1)

    allow(override_redis).to receive(:expire)

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

    allow(raw_override_redis).to receive(:expire)



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

    limiter = Labkit::RateLimit::Limiter.new(
@@ -65,6 +69,6 @@ RSpec.describe "Labkit::RateLimit.configure (Scenario M & N)" do 
    )

    limiter.check({ user: 42 })



    expect(override_redis).to have_received(:incr)

    expect(raw_override_redis).to have_received(:incr)

  end

end

spec/labkit/rate_limit/evaluator_spec.rb

+35 −34
Original line number Diff line number Diff line
@@ -6,7 +6,8 @@ require "redis" 
RSpec.describe Labkit::RateLimit::Evaluator do

  include StubENV



  let(:redis) { instance_double(Redis) }

  let(:raw_redis) { instance_double(Redis) }

  let(:redis) { PooledRedis.new(raw_redis) }

  let(:null_logger) { instance_double(Labkit::Logging::JsonLogger, warn: nil) }

  let(:identifier) { Labkit::RateLimit::Identifier.new(user: 42, ip: "1.2.3.4") }
@@ -30,10 +31,10 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      rule = make_rule(name: "unauthenticated_api", characteristics: [:ip])

      id = Labkit::RateLimit::Identifier.new(ip: "1.2.3.4")



      expect(redis).to receive(:incr)

      expect(raw_redis).to receive(:incr)

        .with("labkit:rl:rack_request:unauthenticated_api:ip:1.2.3.4")

        .and_return(1)

      expect(redis).to receive(:expire).with("labkit:rl:rack_request:unauthenticated_api:ip:1.2.3.4", 60)

      expect(raw_redis).to receive(:expire).with("labkit:rl:rack_request:unauthenticated_api:ip:1.2.3.4", 60)



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

    end
@@ -42,10 +43,10 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      rule = make_rule(name: "auth_api", characteristics: [:user, :ip])

      id = Labkit::RateLimit::Identifier.new(user: 42, ip: "1.2.3.4")



      expect(redis).to receive(:incr)

      expect(raw_redis).to receive(:incr)

        .with("labkit:rl:rack_request:auth_api:user:42:ip:1.2.3.4")

        .and_return(1)

      expect(redis).to receive(:expire)

      expect(raw_redis).to receive(:expire)



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

    end
@@ -56,10 +57,10 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      rule = make_rule(name: "sha_rule", characteristics: [:user])



      expected_hash = OpenSSL::Digest::SHA256.hexdigest(long_value)

      expect(redis).to receive(:incr)

      expect(raw_redis).to receive(:incr)

        .with("labkit:rl:rack_request:sha_rule:user:#{expected_hash}")

        .and_return(1)

      expect(redis).to receive(:expire)

      expect(raw_redis).to receive(:expire)



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

    end
@@ -74,16 +75,16 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
  describe "Scenario R: counter TTL matches rule period" do

    it "sets expire on first write (count == 1)" do

      rule = make_rule(name: "ttl_rule", period: 120)

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

      expect(redis).to receive(:expire).with(anything, 120)

      expect(raw_redis).to receive(:incr).and_return(1)

      expect(raw_redis).to receive(:expire).with(anything, 120)



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

    end



    it "does not reset expire on subsequent writes (count > 1)" do

      rule = make_rule(name: "ttl_rule")

      expect(redis).to receive(:incr).and_return(5)

      expect(redis).not_to receive(:expire)

      expect(raw_redis).to receive(:incr).and_return(5)

      expect(raw_redis).not_to receive(:expire)



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

    end
@@ -94,10 +95,10 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      id = Labkit::RateLimit::Identifier.new(ip: "1.2.3.4")

      rule = make_rule(name: "auth_rule", characteristics: [:user])



      expect(redis).to receive(:incr)

      expect(raw_redis).to receive(:incr)

        .with("labkit:rl:rack_request:auth_rule:user:_unknown_")

        .and_return(1)

      expect(redis).to receive(:expire)

      expect(raw_redis).to receive(:expire)



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

    end
@@ -106,10 +107,10 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      id = Labkit::RateLimit::Identifier.new(user: "", ip: "1.2.3.4")

      rule = make_rule(name: "auth_rule", characteristics: [:user])



      expect(redis).to receive(:incr)

      expect(raw_redis).to receive(:incr)

        .with("labkit:rl:rack_request:auth_rule:user:_unknown_")

        .and_return(1)

      expect(redis).to receive(:expire)

      expect(raw_redis).to receive(:expire)



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

    end
@@ -118,10 +119,10 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      id = Labkit::RateLimit::Identifier.new(ip: "1.2.3.4")

      rule = make_rule(name: "compound", characteristics: [:user, :ip])



      expect(redis).to receive(:incr)

      expect(raw_redis).to receive(:incr)

        .with("labkit:rl:rack_request:compound:user:_unknown_:ip:1.2.3.4")

        .and_return(1)

      expect(redis).to receive(:expire)

      expect(raw_redis).to receive(:expire)



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

    end
@@ -132,10 +133,10 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      rule = make_rule(name: "custom", characteristics: [:custom_field])

      id = Labkit::RateLimit::Identifier.new(custom_field: "abc")



      expect(redis).to receive(:incr)

      expect(raw_redis).to receive(:incr)

        .with("labkit:rl:rack_request:custom:custom_field:abc")

        .and_return(1)

      expect(redis).to receive(:expire)

      expect(raw_redis).to receive(:expire)



      expect { evaluator(rules: [rule]).check(id) }.not_to raise_error

    end
@@ -143,7 +144,7 @@ RSpec.describe Labkit::RateLimit::Evaluator do 


  describe "Scenario S: Redis unavailable" do

    it "returns error Result and logs a warning when Redis is unavailable" do

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

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



      logger = instance_double(Labkit::Logging::JsonLogger)

      expect(logger).to receive(:warn).with(
@@ -162,7 +163,7 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
  describe "non-matching rules" do

    it "does not write to Redis and returns no-match Result" do

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

      expect(redis).not_to receive(:incr)

      expect(raw_redis).not_to receive(:incr)



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

      expect(result.matched?).to be(false)
@@ -178,8 +179,8 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
    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)

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

        allow(raw_redis).to receive(:expire)



        evaluator(rules: [rule]).check(identifier)
@@ -188,8 +189,8 @@ RSpec.describe Labkit::RateLimit::Evaluator do 


      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)

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

        allow(raw_redis).to receive(:expire)



        evaluator(rules: [rule]).check(identifier)
@@ -198,8 +199,8 @@ RSpec.describe Labkit::RateLimit::Evaluator do 


      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)

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

        allow(raw_redis).to receive(:expire)



        evaluator(rules: [rule]).check(identifier)
@@ -210,8 +211,8 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
    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)

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

        allow(raw_redis).to receive(:expire)



        evaluator(rules: [rule]).check(identifier)
@@ -222,8 +223,8 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
    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)

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

        allow(raw_redis).to receive(:expire)



        evaluator(rules: [rule]).check(identifier)
@@ -252,7 +253,7 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
    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")

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



        evaluator(rules: [rule]).check(identifier)
@@ -265,8 +266,8 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
    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)

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

        allow(raw_redis).to receive(:expire)



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

spec/labkit/rate_limit/limiter_spec.rb

+4 −3
Original line number Diff line number Diff line
@@ -6,7 +6,8 @@ require "redis" 
RSpec.describe Labkit::RateLimit::Limiter do

  include StubENV



  let(:redis) { instance_double(Redis, incr: 1, expire: true) }

  let(:raw_redis) { instance_double(Redis, incr: 1, expire: true) }

  let(:redis) { PooledRedis.new(raw_redis) }

  let(:logger) { instance_double(Logger, info: nil, warn: nil) }



  def rule(name: "default", match: {}, limit: 100, period: 60, action: :block, characteristics: [:user])
@@ -74,14 +75,14 @@ RSpec.describe Labkit::RateLimit::Limiter do 


    it "uses the lambda return value as the operative limit" do

      r = rule(limit: -> { 5 }, action: :block)

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

      allow(raw_redis).to receive(:incr).and_return(5)



      result = described_class.new(name: "rack_request", rules: [r], redis: redis, logger: logger)

        .check({ user: 1 })



      expect(result.exceeded?).to be(false)



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

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

      result2 = described_class.new(name: "rack_request", rules: [r], redis: redis, logger: logger)

        .check({ user: 1 })