Commit 87eceb47 authored by Max Woolf's avatar Max Woolf
Browse files

refactor(rate_limit): atomic SADD_SCRIPT for count_distinct rules 

Mirrors INCR_SCRIPT's shape and atomicity. Closes the pre-existing
crash-window between SADD and the separate EXPIRE, returns the real
post-EXPIRE TTL (no -1 sentinel), and self-heals orphan keys whose
TTL is -1 — all matching the behavior INCR adopted in !291.

Spec rewrite for SET-mode #check / #peek to real Redis (TestRedis),
matching the style master adopted for the rest of the file.

Addresses @reprazent's ttl-on-first-write nit on !292.

Co-Authored-By: default avatarClaude Opus 4.7 (1M context) <noreply@anthropic.com>
parent 2fddab5a

lib/labkit/rate_limit/evaluator.rb

+24 −12
Original line number Diff line number Diff line
@@ -38,6 +38,26 @@ module Labkit 
        return {count, redis.call('TTL', KEYS[1])}

      LUA



      # Atomic SADD + SCARD + conditional EXPIRE. SET-cardinality counterpart

      # of INCR_SCRIPT; same shape (read TTL, mutate, set TTL when missing,

      # return post-state {count, TTL}). count is SCARD, not the SADD return.

      #

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

      # so this also self-heals orphan keys left without TTL.

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

        local ttl = ARGV[1]

        local member = ARGV[2]

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



        redis.call('SADD', KEYS[1], member)

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

        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
@@ -218,21 +238,13 @@ module Labkit 
        end

      end



      # Pipelined SADD + SCARD + TTL. EXPIRE follows as a separate call only

      # when this call both added a new member AND the resulting cardinality

      # is 1, i.e. the key was just created. The pipelined TTL is captured

      # before any EXPIRE; first-write callers see ttl=-1 and build_result

      # falls back to the rule period.

      # Atomic SADD + SCARD + conditional EXPIRE in one Redis operation via Lua.

      # See SADD_SCRIPT for the body. Mirrors incr_with_ttl's shape.

      def sadd_with_ttl(redis_key, member, period)

        member_str = encode_char_value(member.to_s)

        @redis.with do |conn|

          added, count, ttl = conn.pipelined do |pipe|

            pipe.sadd?(redis_key, member_str)

            pipe.scard(redis_key)

            pipe.ttl(redis_key)

          end

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

          [count, ttl]

          raw_count, ttl = SADD_SCRIPT.eval(conn, keys: [redis_key], argv: [period, member_str])

          [Integer(raw_count), ttl]

        end

      end

spec/labkit/rate_limit/evaluator_spec.rb

+77 −59
Original line number Diff line number Diff line
@@ -815,28 +815,25 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
  end



  describe "#check with count_distinct (SET-mode)" do

    let(:unique_pipe) { instance_double(Redis) }

    let(:unique_id) { Labkit::RateLimit::Identifier.new(user: 42, ip: "1.2.3.4", project: 99) }



    before do

      allow(unique_pipe).to receive_messages(sadd?: nil, scard: nil, ttl: nil)

      # Default: added=true, scard=1, ttl=-1 (brand-new key)

      allow(raw_redis).to receive(:pipelined).and_yield(unique_pipe).and_return([true, 1, -1])

    end



    it "SADDs the count_distinct value to the rule-keyed compound key" do

      rule = make_rule(name: "uniq_rule", characteristics: [:user], count_distinct: :project)

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



      expect(unique_pipe).to receive(:sadd?).with("labkit:rl:rack_request:uniq_rule:user:42", "99")

      expect(unique_pipe).to receive(:scard).with("labkit:rl:rack_request:uniq_rule:user:42")

      expect(unique_pipe).to receive(:ttl).with("labkit:rl:rack_request:uniq_rule:user:42")

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



      expect(raw_redis.sismember(key, "99")).to be(true)

      expect(raw_redis.scard(key)).to eq(1)

    end



    it "coerces non-string count_distinct values via to_s" do

      rule = make_rule(name: "uniq_rule", characteristics: [:user], count_distinct: :project)

      expect(unique_pipe).to receive(:sadd?).with(anything, "99")

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

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



      evaluator(rules: [rule]).check(unique_id) # project: 99 (Integer)



      expect(raw_redis.sismember(key, "99")).to be(true)

    end



    it "SHA-256-encodes count_distinct values longer than CHAR_VALUE_MAX_LENGTH" do
@@ -844,38 +841,54 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      expected_hash = OpenSSL::Digest::SHA256.hexdigest(long_value)

      rule = make_rule(name: "uniq_rule", characteristics: [:user], count_distinct: :project)

      id = Labkit::RateLimit::Identifier.new(user: 42, project: long_value)

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



      expect(unique_pipe).to receive(:sadd?).with(anything, expected_hash)

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



      expect(raw_redis.sismember(key, expected_hash)).to be(true)

      expect(raw_redis.sismember(key, long_value)).to be(false)

    end



    it "EXPIREs on first write (added=true AND scard=1)" do

      allow(raw_redis).to receive(:pipelined).and_yield(unique_pipe).and_return([true, 1, -1])

    it "sets the TTL on first write" do

      rule = make_rule(name: "uniq_rule", period: 120, count_distinct: :project)

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



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

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



      expect(raw_redis.ttl(key)).to be_between(1, 120)

    end



    it "does not EXPIRE when scard > 1 (existing key, new member)" do

      allow(raw_redis).to receive(:pipelined).and_yield(unique_pipe).and_return([true, 5, 30])

      rule = make_rule(name: "uniq_rule", count_distinct: :project)

    it "does not reset TTL on subsequent writes to a key that already has one" do

      rule = make_rule(name: "uniq_rule", period: 120, count_distinct: :project)

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

      ev = evaluator(rules: [rule])



      expect(raw_redis).not_to receive(:expire)

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

      ev.check(unique_id)

      # Force a distinctive TTL the rule's period would never produce; a subsequent

      # check that mistakenly EXPIRE'd would clobber it back to ~120.

      raw_redis.expire(key, 7)



      ev.check(Labkit::RateLimit::Identifier.new(user: 42, project: 100))



      expect(raw_redis.ttl(key)).to be_between(1, 7)

      expect(raw_redis.scard(key)).to eq(2)

    end



    it "does not EXPIRE when the value was already in the set (added=false)" do

      allow(raw_redis).to receive(:pipelined).and_yield(unique_pipe).and_return([false, 1, 30])

      rule = make_rule(name: "uniq_rule", count_distinct: :project)

    it "self-heals a key that exists without expiry (TTL = -1)" do

      rule = make_rule(name: "heal_rule", period: 60, count_distinct: :project)

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

      raw_redis.sadd(key, %w[1 2]) # no TTL; simulates an orphan from a prior bug



      expect(raw_redis).not_to receive(:expire)

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



      expect(raw_redis.ttl(key)).to be_between(1, 60)

      expect(raw_redis.scard(key)).to eq(3)

    end



    it "reports exceeded when scard > resolved_limit" do

      allow(raw_redis).to receive(:pipelined).and_yield(unique_pipe).and_return([true, 6, 55])

    it "reports exceeded when the post-add cardinality exceeds the limit" do

      rule = make_rule(name: "uniq_rule", limit: 5, action: :block, count_distinct: :project)

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

      raw_redis.sadd(key, %w[a b c d e]) # already at the limit; project 99 tips us over



      result = evaluator(rules: [rule]).check(unique_id)
@@ -885,26 +898,28 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
    end



    it "fails open on Redis error", :aggregate_failures do

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

      broken = instance_double(Redis)

      allow(broken).to receive(:evalsha).and_raise(RuntimeError, "connection refused")

      broken_pool = PooledRedis.new(broken)

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

      expect(logger).to receive(:warn).with(

        hash_including(message: "rate_limit_error", error: "RuntimeError")

      )



      rule = make_rule(name: "uniq_rule", count_distinct: :project)

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

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

        .check(unique_id)



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

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

    end



    it "still reports exceeded on re-SADD of an existing value when scard is already over the limit" do

      # added=false because the value was already in the set; scard unchanged at 6 (over limit=5)

      allow(raw_redis).to receive(:pipelined).and_yield(unique_pipe).and_return([false, 6, 30])

    it "still reports exceeded on re-add when the set is already over the limit" do

      rule = make_rule(name: "uniq_rule", limit: 5, action: :block, count_distinct: :project)

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

      raw_redis.sadd(key, %w[a b c d e 99]) # 6 members; 99 already present



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

      result = evaluator(rules: [rule]).check(unique_id) # re-adds 99, scard unchanged



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

      expect(result.action).to eq(:block)
@@ -914,22 +929,23 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
    it "treats a :log rule as non-terminating and continues to the next rule" do

      log_r = make_rule(name: "log_r", action: :log, limit: 1, count_distinct: :project)

      block_r = make_rule(name: "block_r", action: :block, limit: 100, count_distinct: :project)

      # both rules will SADD; second one returns the same shape

      allow(raw_redis).to receive(:pipelined).and_yield(unique_pipe).and_return([true, 2, 30])



      result = evaluator(rules: [log_r, block_r]).check(unique_id)



      expect(result.rule).to eq(block_r)

      expect(result.action).to eq(:allow)

      expect(raw_redis.scard("labkit:rl:rack_request:log_r:user:42")).to eq(1)

      expect(raw_redis.scard("labkit:rl:rack_request:block_r:user:42")).to eq(1)

    end



    it "uses _unknown_ sentinel for missing characteristics in the rule key (count_distinct is separate)" do

      rule = make_rule(name: "uniq_rule", characteristics: [:user, :ip], count_distinct: :project)

      id = Labkit::RateLimit::Identifier.new(user: 42, project: 99)

      key = "labkit:rl:rack_request:uniq_rule:user:42:ip:_unknown_"



      expect(unique_pipe).to receive(:sadd?)

        .with("labkit:rl:rack_request:uniq_rule:user:42:ip:_unknown_", "99")

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



      expect(raw_redis.sismember(key, "99")).to be(true)

    end



    describe "fail-open + log when count_distinct identifier key is missing" do
@@ -967,13 +983,14 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
        evaluator(rules: [rule]).check(id)

      end



      it "does not call SADD when the count_distinct key is missing" do

      it "does not create the Redis key when the count_distinct value is missing" do

        rule = make_rule(name: "uniq_rule", count_distinct: :project)

        id = Labkit::RateLimit::Identifier.new(user: 42)

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



        expect(unique_pipe).not_to receive(:sadd?)

        expect(raw_redis).not_to receive(:expire)

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



        expect(raw_redis.exists?(key)).to be(false)

      end



      it "falls through to a following matching rule when the SET-mode rule's key is missing" do
@@ -1004,25 +1021,21 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
  end



  describe "#peek with count_distinct (SET-mode)" do

    let(:unique_pipe) { instance_double(Redis) }



    before do

      allow(unique_pipe).to receive_messages(scard: nil, ttl: nil)

      allow(raw_redis).to receive(:pipelined).and_yield(unique_pipe).and_return([3, 30])

    end



    it "reads SCARD and TTL without SADD or EXPIRE" do

    it "reads SCARD without mutating the set or touching TTL" do

      rule = make_rule(name: "uniq_rule", characteristics: [:user], count_distinct: :project)

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

      raw_redis.sadd(key, %w[a b c])

      raw_redis.expire(key, 7)



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



      expect(unique_pipe).not_to receive(:sadd?)

      expect(raw_redis).not_to receive(:expire)

      expect(unique_pipe).to receive(:scard).with("labkit:rl:rack_request:uniq_rule:user:42")

      expect(unique_pipe).to receive(:ttl).with("labkit:rl:rack_request:uniq_rule:user:42")

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

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

      expect(result.info.count).to eq(3)

      expect(raw_redis.scard(key)).to eq(3) # no SADD happened

      expect(raw_redis.ttl(key)).to be_between(1, 7) # TTL not extended

    end



    it "reports count=0 when key is missing (SCARD returns 0)" do

      allow(raw_redis).to receive(:pipelined).and_yield(unique_pipe).and_return([0, -2])

      rule = make_rule(name: "uniq_rule", limit: 5, period: 60, count_distinct: :project)



      result = evaluator(rules: [rule]).peek(identifier)
@@ -1034,8 +1047,9 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
    end



    it "reports exceeded when the current cardinality is over the limit" do

      allow(raw_redis).to receive(:pipelined).and_yield(unique_pipe).and_return([10, 30])

      rule = make_rule(name: "uniq_rule", limit: 5, action: :block, count_distinct: :project)

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

      raw_redis.sadd(key, %w[a b c d e f g h i j])



      result = evaluator(rules: [rule]).peek(identifier)
@@ -1044,10 +1058,13 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
    end



    it "fails open on Redis error" do

      allow(raw_redis).to receive(:pipelined).and_raise(RuntimeError, "down")

      broken = instance_double(Redis)

      allow(broken).to receive(:pipelined).and_raise(RuntimeError, "down")

      broken_pool = PooledRedis.new(broken)

      rule = make_rule(name: "uniq_rule", count_distinct: :project)



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

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

        .peek(identifier)



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

      expect(result.exceeded?).to be(false)
@@ -1055,8 +1072,9 @@ RSpec.describe Labkit::RateLimit::Evaluator do 


    it "does not require the count_distinct identifier key (reads the bucket cardinality)" do

      rule = make_rule(name: "uniq_rule", characteristics: [:user], count_distinct: :project)

      id = Labkit::RateLimit::Identifier.new(user: 42)

      allow(raw_redis).to receive(:pipelined).and_yield(unique_pipe).and_return([3, 30])

      id = Labkit::RateLimit::Identifier.new(user: 42) # no :project

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

      raw_redis.sadd(key, %w[a b c])



      expect(null_logger).not_to receive(:warn)

      result = evaluator(rules: [rule]).peek(id)

spec/labkit/rate_limit/limiter_spec.rb

+5 −8
Original line number Diff line number Diff line
@@ -237,10 +237,9 @@ RSpec.describe Labkit::RateLimit::Limiter do 
  end



  describe "#check with a count_distinct rule" do

    let(:unique_pipe) { instance_double(Redis, sadd?: true, scard: 1, ttl: -1) }



    before do

      allow(raw_redis).to receive(:pipelined).and_yield(unique_pipe).and_return([true, 4, 30])

      # SADD_SCRIPT.eval returns [scard, ttl]

      allow(raw_redis).to receive(:evalsha).and_return([4, 30])

    end



    it "returns a Result whose count is the SCARD post-add" do
@@ -259,20 +258,18 @@ RSpec.describe Labkit::RateLimit::Limiter do 
  end



  describe "#peek with a count_distinct rule" do

    let(:unique_pipe) { instance_double(Redis, scard: 0, ttl: -2) }



    before do

      allow(raw_redis).to receive(:pipelined).and_yield(unique_pipe).and_return([0, -2])

      allow(raw_redis).to receive(:pipelined).and_return([0, -2])

    end



    it "does not call SADD or EXPIRE" do

    it "does not invoke the SADD script" do

      r = Labkit::RateLimit::Rule.new(

        name: "uniq", limit: 10, period: 60,

        characteristics: [:user], count_distinct: :project

      )

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



      expect(unique_pipe).not_to receive(:sadd?)

      expect(raw_redis).not_to receive(:evalsha)

      expect(raw_redis).not_to receive(:expire)

      lim.peek({ user: 42 })

    end