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

feat(rate_limit): set-cardinality check_unique / peek_unique primitives 

Adds a SADD+SCARD-backed counter alongside the existing INCR-backed one,
exposed as Limiter#check_unique(identifier, member:) and #peek_unique.
The limit constrains the number of distinct members observed within the
rule's period rather than the number of calls — the labkit-side
equivalent of ApplicationRateLimiter::IncrementPerActionedResource.

Lays the foundation for the cohort 4 rollout
(unique_project_downloads_for_{application,namespace}, EE GitAbuse).
The follow-up consumer-side MR will:

- add `mode: :unique` and `accepts_override: true` to those two
  SupportedRateLimits entries (both call sites always pass
  threshold:/interval: overrides; without accepting them the labkit
  path never runs),
- plumb `resource.id` through LabkitAdapter#run_unique! as `member:`,
- add `rate_limiter_use_labkit_cohort_4{,_enforce}` feature flags.

Design notes:

- Separate public API (check_unique / peek_unique) rather than a `mode:`
  kwarg on #check. SET-mode and INCR-mode have mutually-exclusive
  per-call shapes (cost vs member) and operate on incompatible Redis
  data types — two narrow APIs read more clearly than one polymorphic
  one and avoid WRONGTYPE foot-guns.
- EXPIRE-on-first-write is gated on `added && scard == 1`, matching the
  legacy IncrementPerActionedResource semantic (a re-SADD of an
  existing member should not reset the window).
- Member is to_s-coerced and SHA-256-encoded above CHAR_VALUE_MAX_LENGTH,
  identical to characteristic-value handling, so callers can pass
  arbitrary subjects without worrying about Redis key bloat.
- Pipelined (not Lua), mirroring the current INCR path on master.
  If !291's atomic Lua INCR lands, a follow-up can convert this path
  to a parallel atomic script (SADD + SCARD + EXPIRE in one round-trip).

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

lib/labkit/rate_limit.rb

+14 −0
Original line number Diff line number Diff line
@@ -47,6 +47,20 @@ module Labkit 
      def check(name:, identifier:, rules:, redis: nil, logger: nil)

        Limiter.new(name: name, rules: rules, redis: redis, logger: logger).check(identifier)

      end



      # Convenience wrapper for set-cardinality checks. See {Limiter#check_unique}.

      #

      # @param name [String] call site name

      # @param identifier [Identifier, Hash] caller attributes

      # @param rules [Array<Rule>] ordered list of rules (first match wins)

      # @param member [#to_s] the subject to add to the set

      # @param redis [Object, nil] Redis client; falls back to config.redis

      # @param logger [Logger, nil] logger; falls back to config.logger

      # @return [Result]

      def check_unique(name:, identifier:, rules:, member:, redis: nil, logger: nil)

        Limiter.new(name: name, rules: rules, redis: redis, logger: logger)

               .check_unique(identifier, member: member)

      end

    end

  end

end

lib/labkit/rate_limit/evaluator.rb

+97 −0
Original line number Diff line number Diff line
@@ -40,6 +40,32 @@ module Labkit 
        Result.new(matched: false, error: true, action: :allow)

      end



      # Set-cardinality variant of {#check}: adds `member` to a Redis SET and

      # uses SCARD as the counter, so the limit constrains the number of

      # distinct members observed within the rule's period rather than the

      # number of calls. Members already in the set don't advance the count.

      #

      # Mirrors {#check}'s Result shape (`info.count` is the post-add SCARD)

      # and fail-open behavior.

      def check_unique(identifier, member:)

        check_unique_rules(identifier, member)

      rescue StandardError => e

        report_error_metrics

        log_error(e, identifier)

        Result.new(matched: false, error: true, action: :allow)

      end



      # Read-without-write counterpart to {#check_unique}. Returns the current

      # SCARD without adding to the set or extending the TTL. A missing key

      # reports count=0.

      def peek_unique(identifier)

        peek_unique_rules(identifier)

      rescue StandardError => e

        report_error_metrics

        log_error(e, identifier)

        Result.new(matched: false, error: true, action: :allow)

      end



      private



      # :log rules are non-terminating: they emit metrics and continue,
@@ -70,6 +96,30 @@ module Labkit 
        Result.new(matched: false, action: :allow)

      end



      def check_unique_rules(identifier, member)

        @rules.each do |rule|

          next unless rule_matches?(rule, identifier)



          result = evaluate_unique_rule(rule, identifier, member)

          report_matched_metrics(result)

          return result unless rule.action == :log

        end



        report_unmatched_metrics

        Result.new(matched: false, action: :allow)

      end



      def peek_unique_rules(identifier)

        @rules.each do |rule|

          next if rule.action == :log

          next unless rule_matches?(rule, identifier)



          return peek_unique_rule(rule, identifier)

        end



        Result.new(matched: false, action: :allow)

      end



      def rule_matches?(rule, identifier)

        rule.match.all? { |key, matcher| matcher.match?(identifier[key]) }

      end
@@ -92,6 +142,24 @@ module Labkit 
        build_result(rule, resolved_limit, resolved_period, count, ttl)

      end



      def evaluate_unique_rule(rule, identifier, member)

        redis_key = build_redis_key(rule, identifier)

        resolved_limit = Integer(resolve_value(rule.limit))

        resolved_period = Integer(resolve_value(rule.period))



        count, ttl = sadd_with_ttl(redis_key, member, resolved_period)

        build_result(rule, resolved_limit, resolved_period, count, ttl)

      end



      def peek_unique_rule(rule, identifier)

        redis_key = build_redis_key(rule, identifier)

        resolved_limit = Integer(resolve_value(rule.limit))

        resolved_period = Integer(resolve_value(rule.period))



        count, ttl = scard_with_ttl(redis_key)

        build_result(rule, resolved_limit, resolved_period, count, ttl)

      end



      def build_result(rule, resolved_limit, resolved_period, count, ttl)

        exceeded = count > resolved_limit

        action = exceeded ? rule.action : :allow
@@ -162,6 +230,35 @@ 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.

      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]

        end

      end



      # Pipelined SCARD + TTL. SCARD on a missing key returns 0, so no

      # explicit nil handling is needed (unlike GET in read_with_ttl).

      def scard_with_ttl(redis_key)

        @redis.with do |conn|

          conn.pipelined do |pipe|

            pipe.scard(redis_key)

            pipe.ttl(redis_key)

          end

        end

      end



      def log_error(error, identifier)

        @logger.warn(

          message: "rate_limit_error",

lib/labkit/rate_limit/limiter.rb

+26 −0
Original line number Diff line number Diff line
@@ -57,6 +57,32 @@ module Labkit 
        @evaluator.peek(id)

      end



      # Set-cardinality check: adds +member+ to a Redis SET keyed by the

      # matched rule's identifier and uses SCARD as the counter. Use this

      # when the limit constrains the number of distinct subjects observed

      # within the rule's period (e.g. "max N unique projects downloaded

      # per user per hour"), as opposed to {#check}'s "max N calls" shape.

      #

      # @param identifier [Identifier, Hash] caller attributes for this request

      # @param member [#to_s] the subject to add to the set. Must be non-nil.

      # @return [Result]

      def check_unique(identifier, member:)

        raise ArgumentError, "member must be non-nil" if member.nil?



        id = identifier.is_a?(Identifier) ? identifier : Identifier.new(identifier)

        @evaluator.check_unique(id, member: member)

      end



      # Read-without-write counterpart to {#check_unique}: returns the current

      # SCARD without adding to the set or extending the TTL.

      #

      # @param identifier [Identifier, Hash] caller attributes for this request

      # @return [Result]

      def peek_unique(identifier)

        id = identifier.is_a?(Identifier) ? identifier : Identifier.new(identifier)

        @evaluator.peek_unique(id)

      end



      private



      def validate_name!(name)

spec/labkit/rate_limit/evaluator_spec.rb

+168 −0
Original line number Diff line number Diff line
@@ -950,4 +950,172 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      end

    end

  end



  describe "#check_unique (SET-mode)" do

    let(:unique_pipe) { instance_double(Redis) }



    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 member to the rule-keyed compound key" do

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



      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(identifier, member: 99)

    end



    it "coerces non-string members via to_s" do

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

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

      evaluator(rules: [rule]).check_unique(identifier, member: 99)

    end



    it "SHA-256-encodes members longer than CHAR_VALUE_MAX_LENGTH" do

      long_member = "x" * 201

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

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



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

      evaluator(rules: [rule]).check_unique(identifier, member: long_member)

    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])

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



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

      evaluator(rules: [rule]).check_unique(identifier, member: 1)

    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")



      expect(raw_redis).not_to receive(:expire)

      evaluator(rules: [rule]).check_unique(identifier, member: 1)

    end



    it "does not EXPIRE when member 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")



      expect(raw_redis).not_to receive(:expire)

      evaluator(rules: [rule]).check_unique(identifier, member: 1)

    end



    it "reports exceeded when scard > resolved_limit" do

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

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



      result = evaluator(rules: [rule]).check_unique(identifier, member: 99)



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

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

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

    end



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

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

      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")

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

        .check_unique(identifier, member: 1)



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

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

    end



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

      # added=false because the member 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])

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



      result = evaluator(rules: [rule]).check_unique(identifier, member: 99)



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

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

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

    end



    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)

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

      # 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(identifier, member: 1)



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

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

    end



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

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

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



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

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

      evaluator(rules: [rule]).check_unique(id, member: 99)

    end

  end



  describe "#peek_unique (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

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



      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_unique(identifier)

    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)



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



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

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

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

      expect(result.info.remaining).to eq(5)

    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)



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



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

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

    end



    it "fails open on Redis error" do

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

      rule = make_rule(name: "uniq_rule")



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



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

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

    end

  end

end

spec/labkit/rate_limit/limiter_spec.rb

+65 −0
Original line number Diff line number Diff line
@@ -236,4 +236,69 @@ RSpec.describe Labkit::RateLimit::Limiter do 
      end

    end

  end



  describe "#check_unique" 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, 1, -1])

    end



    it "delegates to the evaluator's check_unique with the member" do

      lim = limiter

      expect(lim.instance_variable_get(:@evaluator))

        .to receive(:check_unique).with(instance_of(Labkit::RateLimit::Identifier), member: 7)

      lim.check_unique({ user: 42 }, member: 7)

    end



    it "raises ArgumentError when member is nil" do

      expect { limiter.check_unique({ user: 42 }, member: nil) }

        .to raise_error(ArgumentError, /member must be non-nil/)

    end



    it "accepts an Identifier instance directly without re-wrapping" do

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

      expect(Labkit::RateLimit::Identifier).not_to receive(:new)

      limiter.check_unique(id, member: 7)

    end



    it "returns a Result with count equal to SCARD" do

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

      r = rule(limit: 10)



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

        .check_unique({ user: 42 }, member: 99)



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

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

      expect(result.info.remaining).to eq(6)

    end

  end



  describe "#peek_unique" 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])

    end



    it "delegates to the evaluator's peek_unique path" do

      lim = limiter

      expect(lim.instance_variable_get(:@evaluator))

        .to receive(:peek_unique).with(instance_of(Labkit::RateLimit::Identifier))

      lim.peek_unique({ user: 42 })

    end



    it "accepts an Identifier instance directly without re-wrapping" do

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

      expect(Labkit::RateLimit::Identifier).not_to receive(:new)

      limiter.peek_unique(id)

    end



    it "does not call SADD or EXPIRE" do

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

      expect(raw_redis).not_to receive(:expire)

      limiter.peek_unique({ user: 42 })

    end

  end

end