Verified Commit 0834e03f authored by Bob Van Landuyt's avatar Bob Van Landuyt 💬
Browse files

perf(rate_limit): read the counter TTL once per check

Both Lua scripts read the key's TTL twice: once before mutating, to decide
whether to EXPIRE, and once after, to return the window's remaining time.
The pre-read is redundant. INCRBYFLOAT and SADD both preserve an existing
key's TTL and create a missing key with no expiry, so reading the TTL after
the mutation gives what the pre-read gave. It still tells apart the two
cases the script has to handle: -2 for a missing key, -1 for a key left
without an expiry by some earlier bug.

So move the read after the mutation and drop the second one. INCR_SCRIPT
goes from four unconditional Redis calls to three, SADD_SCRIPT from five to
four. Atomicity is unchanged, TTL-less keys still self-heal, and the
{count, ttl} return contract is the same, so the existing specs for the
three TTL states of each script pass without modification.

One thing does change shape. On the first write of a window the script now
returns the period it just set rather than reading it back from Redis. Same
value, and reset_at still derives from it.

Worth being honest about the size of this. TTL is the most-called command
on gprd's redis-cluster-ratelimiting, roughly 115k/s out of 566k/s, since
the rack shadow went to 100% on 2026-07-13 and every check spends two.
Halving the count only buys about 1% of node CPU though: command execution
is some 18% of process CPU on the busiest primary, and this removes no
round trips. Dropping whole EVALSHA calls is what would matter for
capacity, and that is what :skip does.

Related to gitlab-com/gl-infra/production-engineering#28807
parent fa04d0a3
Loading
Loading
Loading
Loading
+31 −15
Original line number Diff line number Diff line
@@ -15,7 +15,7 @@ flowchart LR
    App[Application code] -->|"check(identifier)"| Limiter
    Limiter -->|delegates| Evaluator
    Evaluator -->|iterates ordered| Rules[Rule list]
    Evaluator <-->|INCR / TTL / EXPIRE| Redis[(Redis)]
    Evaluator <-->|EVALSHA / GET / TTL| Redis[(Redis)]
    Evaluator -->|emits| Metrics[Prometheus metrics]
    Evaluator -->|returns| Result
    Result --> App
@@ -236,9 +236,14 @@ labkit:rl:<limiter_name>:<rule_name>:<char>:<value>[:<char>:<value>...]

Characteristic values longer than 200 bytes are replaced with a SHA-256
hexdigest to bound key length. Missing or empty characteristic values are
encoded as `_unknown_`. The TTL is set on the first write of each window
(`count == 1`) and is not extended on subsequent INCRs, so the window is a
true fixed window starting at the first request, not a sliding window.
encoded as `_unknown_`. The TTL is set on the first write of each window and
is not extended on subsequent increments, so the window is a true fixed
window starting at the first request, not a sliding window.

A check is a single `EVALSHA` of `INCR_SCRIPT` (or `SADD_SCRIPT` for
`count_distinct` rules). Doing the whole read-modify-write inside Lua means
there is no window between the increment and the `EXPIRE` in which a key
could be left without a TTL.

```mermaid
sequenceDiagram
@@ -246,24 +251,35 @@ sequenceDiagram
    participant E as Evaluator
    participant P as Connection pool
    participant R as Redis
    participant L as INCR_SCRIPT (Lua)

    E->>P: pool.with { |conn| ... }
    P-->>E: conn
    E->>R: PIPELINE { INCR key, TTL key }
    R-->>E: [count, ttl]
    alt count == 1 (first write of window)
        E->>R: EXPIRE key period
        R-->>E: 1
        Note over E: ttl returned is -1 here;<br/>build_result falls back to<br/>resolved_period for reset_at.
    else count > 1
        Note over E: TTL is not extended:<br/>fixed window from first write.
    E->>R: EVALSHA INCR_SCRIPT key, [period, cost]
    R->>L: run script
    L->>L: INCRBYFLOAT key cost
    L->>L: TTL key
    alt TTL < 0 (key was missing, or had no expiry)
        L->>L: EXPIRE key period
        Note over L: Returns period as the TTL:<br/>the window starts now.
    else TTL >= 0 (window already running)
        Note over L: TTL is not extended:<br/>fixed window from first write.
    end
    L-->>R: {count, ttl}
    R-->>E: [count, ttl]
    E-->>P: release conn
```

`peek` follows the same shape but uses `GET` instead of `INCR` and never
issues `EXPIRE`. A missing key (`GET → nil`, `TTL → -2`) is reported as
`count = 0` and the window is treated as not-yet-started.
Because `INCRBYFLOAT` preserves an existing key's TTL, and creates a missing
key with no expiry, the TTL only needs reading once — after the increment.
That single read distinguishes both cases the script must handle, so there is
no reason to read it again before mutating.

`peek` does not use a script: it pipelines `GET` + `TTL` (or `SCARD` + `TTL`)
and never issues `EXPIRE`, so it cannot start or extend a window. A missing
key (`GET → nil`, `TTL → -2`) is reported as `count = 0`, and `build_result`
falls back to the rule's period for `reset_at` since there is no Redis-side
window to read.

## Result

+12 −10
Original line number Diff line number Diff line
@@ -22,40 +22,42 @@ module Labkit
      # Redis treats the result as a no-op on the stored value while
      # still observing the post-state count and TTL we return.
      #
      # ttl_before < 0 covers TTL=-2 (key missing) and TTL=-1 (no expiry).
      # ttl_after < 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 = redis.call('INCRBYFLOAT', KEYS[1], cost)
        if ttl_before < 0 then
        local ttl_after = redis.call('TTL', KEYS[1])
        if ttl_after < 0 then
          redis.call('EXPIRE', KEYS[1], ttl)
          ttl_after = tonumber(ttl)
        end

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

      # Atomic SADD + SCARD + conditional EXPIRE. SET-cardinality counterpart
      # of INCR_SCRIPT; same shape (read TTL, mutate, set TTL when missing,
      # of INCR_SCRIPT; same shape (mutate, read TTL, 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),
      # ttl_after < 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
        local ttl_after = redis.call('TTL', KEYS[1])
        if ttl_after < 0 then
          redis.call('EXPIRE', KEYS[1], ttl)
          ttl_after = tonumber(ttl)
        end

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

      def initialize(name:, rules:, redis:, logger:)
@@ -247,7 +249,7 @@ module Labkit
      end

      # Atomically increments the counter by `cost`, sets the TTL on first
      # write, and reads back the post-increment TTL, all in one Redis
      # write, and returns the window's remaining TTL, all in one Redis
      # operation via Lua. See INCR_SCRIPT for the script body.
      #
      # count is parsed as Float because INCRBYFLOAT returns a string-encoded