Commit 3c15ecb6 authored by Max Woolf's avatar Max Woolf
Browse files

refactor(rate_limit): simplify Lua script to always INCRBYFLOAT 

Replace the cost=0 GET short-circuit with a uniform INCRBYFLOAT branch.
The short-circuit was originally there to avoid allocating Redis keys
for cohort 5 zero-usage callers (IncrementResourceUsagePerAction), but
the actual call volume is small enough that the branch's complexity
isn't paying for itself. INCRBYFLOAT with cost=0 is a no-op on the
stored value, with the trade that:

  - Cost=0 against a missing key now allocates a key at value 0 with
    the rule's TTL (was: no allocation).
  - Cost=0 against an existing key issues a write op (was: pure GET).
    The observable count and TTL are unchanged.

Cohort 5 reviewers should be aware that zero-usage workers no longer
get a free pass at the storage layer. At the call rates involved
(~< 100/sec fleet-wide), the additional allocation is well below the
noise floor and the simpler primitive is easier to reason about.

Addresses !291 review thread 148bc41b.
parent b2361e4e

lib/labkit/rate_limit/evaluator.rb

+11 −13
Original line number Diff line number Diff line
@@ -18,25 +18,23 @@ module Labkit 
      #

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

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

      #

      # - cost=0 short-circuits to GET so resource-usage callers that

      #   observed zero usage do not allocate a Redis key.

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

      # 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

        if cost == 0 then

          count = redis.call('GET', KEYS[1]) or '0'

        else

          count = redis.call('INCRBYFLOAT', KEYS[1], cost)



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

        if ttl_before < 0 then

            redis.call('EXPIRE', KEYS[1], ARGV[1])

          end

          redis.call('EXPIRE', KEYS[1], ttl)

        end



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

      LUA

spec/labkit/rate_limit/evaluator_spec.rb

+7 −9
Original line number Diff line number Diff line
@@ -162,24 +162,22 @@ RSpec.describe Labkit::RateLimit::Evaluator do 
      expect(stored_count(key)).to eq(2.5)

    end



    it "with cost=0 reads the counter without writing to Redis" do

    it "with cost=0 leaves the stored count unchanged" do

      ev = evaluator(rules: [rule])

      ev.check(identifier, cost: 3.0)

      initial_count = stored_count(key)

      initial_ttl = raw_redis.ttl(key)

      sleep 1



      result = ev.check(identifier, cost: 0)



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

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

      expect(raw_redis.ttl(key)).to be < initial_ttl

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

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

    end



    it "with cost=0 against a missing key reports count=0 and does not allocate the key" do

    it "with cost=0 against a missing key reports count=0 and allocates the key with the rule's TTL" do

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



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

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

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

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

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

    end

spec/labkit/rate_limit_spec.rb

+4 −3
Original line number Diff line number Diff line
@@ -411,7 +411,7 @@ RSpec.describe Labkit::RateLimit do 
      expect(get_value("labkit:rl:rack_request:cost_frac:user:1")).to be_within(0.0001).of(2.5)

    end



    it "with cost=0 reads without writing to Redis and reports the current count" do

    it "with cost=0 leaves the stored count unchanged and reports the current count" do

      r = rule(name: "cost_zero", limit: 100, characteristics: [:user])

      lim = limiter(rules: [r])

      lim.check({ user: 1 }, cost: 5.0)
@@ -422,14 +422,15 @@ RSpec.describe Labkit::RateLimit do 
      expect(get_value("labkit:rl:rack_request:cost_zero:user:1")).to eq(5.0)

    end



    it "with cost=0 against a missing key does not allocate the key" do

    it "with cost=0 against a missing key reports count=0 and allocates the key at 0" do

      r = rule(name: "cost_zero_missing", limit: 100, characteristics: [:user])

      lim = limiter(rules: [r])



      result = lim.check({ user: 99 }, cost: 0)



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

      expect(raw_redis.exists?("labkit:rl:rack_request:cost_zero_missing:user:99")).to be(false)

      expect(raw_redis.exists?("labkit:rl:rack_request:cost_zero_missing:user:99")).to be(true)

      expect(get_value("labkit:rl:rack_request:cost_zero_missing:user:99")).to eq(0.0)

    end



    it "blocks when accumulated fractional cost exceeds the integer limit" do