Verified Commit 2b17b156 authored by Max Woolf's avatar Max Woolf Committed by GitLab
Browse files

Merge branch 'rate-limit/single-ttl-read' into 'master'

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

See merge request !333

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 fa04d0a3 0834e03f
Loading
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