Chat with LLM: Source Code Analysis of Redisson Rate Limiter

Posted on In

Introduction

My project uses Redis as a cache. For user module endpoints, except for login and registration which require a phone number, all other endpoints operate using an ID. The initial code stored both an ID-to-user mapping and a phone-number-to-user mapping in the cache. Storing the same user object twice is obviously unreasonable, not to mention the potential additional consistency problems that could arise.

Since one of the purposes of a cache is to protect the database, naturally we can achieve this by other means, such as rate limiting. If we only apply rate limiting to the login and registration endpoints, then we can store only the ID-to-user mapping in the cache. This is also reasonable from a business perspective, as these two endpoints are not frequently called.

There are at least three specific implementations of rate limiting: the leaky bucket algorithm, the token bucket algorithm, and message queues. I don't want to introduce a message queue just for rate limiting, especially since that component is difficult to operate. The token bucket algorithm is a superior alternative to the leaky bucket algorithm, as it better handles burst traffic. Therefore, the leaky bucket algorithm is rarely used nowadays.

There are many libraries that implement the token bucket algorithm: Guava, Bucket4j, Resilience4j, Sentinel, Redisson, etc. Clearly, this business scenario requires distributed rate limiting. Since Redisson is already used in the project, and my requirements for rate limiting are simple, I decided to use the rate limiter provided by Redisson.

The token bucket algorithm is conceptually simple: it is a non-blocking producer-consumer model. A bucket with a fixed capacity is filled with tokens at a constant rate. When a request arrives, it attempts to acquire a token, and the request is processed only if it successfully obtains one. However, I wasn't clear on the specifics of how this algorithm is implemented, so in this blog post, I leverage an LLM to read the source code of the RedissonRateLimiter class and analyze its implementation.

Source Code Analysis

RedissonRateLimiter is a distributed rate limiter implementation provided by Redisson, based on the token bucket algorithm. It utilizes Redis data structures such as Sorted Sets, Hashes, and Strings, and ensures the atomicity of multiple operations through Lua scripts, thereby achieving cross-JVM rate limiting capabilities.

Key Fields and Methods

  • Name Generation: The rate limiter stores multiple keys in Redis, generated with fixed suffixes:

    • getRawName(): The main key of the rate limiter (stores configuration in a Hash).
    • getValueName(): The key storing the current number of available tokens (String type).
    • getClientValueName(): Client-specific token counter (used in PER_CLIENT mode).
    • getPermitsName(): The Sorted Set storing issued but not yet expired tokens.
    • getClientPermitsName(): Client-specific token set.

  • Core Methods:

    • trySetRateAsync(): Attempts to initialize the rate limiter configuration (atomic operation).
    • tryAcquireAsync(): Attempts to acquire a specified number of tokens, returns the wait time (in milliseconds) or null for immediate success.
    • releaseAsync(): Manually releases tokens (typically for compensation).
    • availablePermitsAsync(): Queries the current number of available tokens.
    • getConfigAsync(): Retrieves the current rate limiter configuration.

Initializing the Rate Limiter Configuration

Java Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
@Override
public RFuture<Boolean> trySetRateAsync(RateType type, long rate, Duration rateInterval, Duration keepAliveTime) {
    // Validate that keepAliveTime is not less than rateInterval, otherwise throw an exception
    if (!keepAliveTime.equals(Duration.ZERO) && keepAliveTime.toMillis() < rateInterval.toMillis()) {
        throw new IllegalArgumentException("The parameter keepAliveTime should be greater than or equal to rateInterval");
    }
    // Execute Lua script using EVAL command, no retry (because HSETNX is idempotent)
    return commandExecutor.evalWriteNoRetryAsync(getRawName(), LongCodec.INSTANCE, RedisCommands.EVAL_BOOLEAN,
            // Lua script content (see below)
            "redis.call('hsetnx', KEYS[1], 'rate', ARGV[1]);"
          + "redis.call('hsetnx', KEYS[1], 'interval', ARGV[2]);"
          + "redis.call('hsetnx', KEYS[1], 'keepAliveTime', ARGV[4]);"
          + "local res = redis.call('hsetnx', KEYS[1], 'type', ARGV[3]);"
          + "if res == 1 and tonumber(ARGV[4]) > 0 then "
                + "redis.call('pexpire', KEYS[1], ARGV[4]); "
          + "end; "
          + "return res;",
            Collections.singletonList(getRawName()),                     // KEYS[1]
            rate, rateInterval.toMillis(), type.ordinal(), keepAliveTime.toMillis()); // ARGV[1..4]
}

Method Explanation:

  • Uses HSETNX commands to sequentially set the rate, interval, keepAliveTime, and type fields. HSETNX succeeds only if the field does not already exist, ensuring the configuration is initialized only once.
  • If this initialization succeeds (res == 1) and keepAliveTime is set, it immediately sets an expiration time on the main key.
  • The return value true indicates successful initialization (i.e., no configuration existed before), false indicates the configuration already exists.

Lua Script

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
-- Attempt to set the rate field, succeeds only if the field does not exist
redis.call('hsetnx', KEYS[1], 'rate', ARGV[1]);

-- Attempt to set the interval field
redis.call('hsetnx', KEYS[1], 'interval', ARGV[2]);

-- Attempt to set the keepAliveTime field
redis.call('hsetnx', KEYS[1], 'keepAliveTime', ARGV[4]);

-- Attempt to set the type field and record the return value
local res = redis.call('hsetnx', KEYS[1], 'type', ARGV[3]);

-- If this initialization succeeded (i.e., none of the fields existed before) and keepAliveTime > 0
if res == 1 and tonumber(ARGV[4]) > 0 then
    -- Set expiration time (milliseconds) on the main rate limiter key
    redis.call('pexpire', KEYS[1], ARGV[4]);
end;

-- Return the initialization result: 1 for success, 0 if configuration already exists
return res;

Attempting to Acquire Tokens

Java Code

tryAcquireAsync has multiple overloads, ultimately calling a private method:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
private <T> RFuture<T> tryAcquireAsync(RedisCommand<T> command, Long value) {
    // Generate a random byte array to identify this request and avoid concurrency conflicts
    byte[] random = getServiceManager().generateIdArray();

    // Execute Lua script
    return commandExecutor.evalWriteAsync(getRawName(), LongCodec.INSTANCE, command,
            // Lua script content (see below)
            "local rate = redis.call('hget', KEYS[1], 'rate');"
          + "local interval = redis.call('hget', KEYS[1], 'interval');"
          + "local type = redis.call('hget', KEYS[1], 'type');"
          + "assert(rate ~= false and interval ~= false and type ~= false, 'RateLimiter is not initialized')"
          
          + "local valueName = KEYS[2];"
          + "local permitsName = KEYS[4];"
          + "if type == '1' then "
              + "valueName = KEYS[3];"
              + "permitsName = KEYS[5];"
          + "end;"

          + "assert(tonumber(rate) >= tonumber(ARGV[1]), 'Requested permits amount cannot exceed defined rate'); "

          + "local currentValue = redis.call('get', valueName); "
          + "local res;"
          + "if currentValue ~= false then "
                 + "local expiredValues = redis.call('zrangebyscore', permitsName, 0, tonumber(ARGV[2]) - interval); "
                 + "local released = 0; "
                 + "for i, v in ipairs(expiredValues) do "
                      + "local random, permits = struct.unpack('Bc0I', v);"
                      + "released = released + permits;"
                 + "end; "

                 + "if released > 0 then "
                      + "redis.call('zremrangebyscore', permitsName, 0, tonumber(ARGV[2]) - interval); "
                      + "if tonumber(currentValue) + released > tonumber(rate) then "
                           + "local values = redis.call('zrange', permitsName, 0, -1); "
                           + "local used = 0; "
                           + "for i, v in ipairs(values) do "
                                + "local random, permits = struct.unpack('Bc0I', v);"
                                + "used = used + permits;"
                           + "end; "
                           + "currentValue = tonumber(rate) - used; "
                      + "else "
                           + "currentValue = tonumber(currentValue) + released; "
                      + "end; "
                      + "redis.call('set', valueName, currentValue);"
                 + "end;"

                 + "if tonumber(currentValue) < tonumber(ARGV[1]) then "
                     + "local firstValue = redis.call('zrange', permitsName, 0, 0, 'withscores'); "
                     + "res = 3 + interval - (tonumber(ARGV[2]) - tonumber(firstValue[2]));"
                 + "else "
                     + "redis.call('zadd', permitsName, ARGV[2], struct.pack('Bc0I', string.len(ARGV[3]), ARGV[3], ARGV[1])); "
                     + "redis.call('decrby', valueName, ARGV[1]); "
                     + "res = nil; "
                 + "end; "
          + "else "
                 + "redis.call('set', valueName, rate); "
                 + "redis.call('zadd', permitsName, ARGV[2], struct.pack('Bc0I', string.len(ARGV[3]), ARGV[3], ARGV[1])); "
                 + "redis.call('decrby', valueName, ARGV[1]); "
                 + "res = nil; "
          + "end;"

          + "local keepAliveTime = redis.call('hget', KEYS[1], 'keepAliveTime'); "
          + "if (keepAliveTime ~= false and tonumber(keepAliveTime) > 0) then "
                + "redis.call('pexpire', KEYS[1], keepAliveTime); "
                + "redis.call('pexpire', valueName, keepAliveTime); "
                + "redis.call('pexpire', permitsName, keepAliveTime); "
          + "else "
                + "local ttl = redis.call('pttl', KEYS[1]); "
                + "if ttl > 0 then "
                    + "redis.call('pexpire', valueName, ttl); "
                    + "redis.call('pexpire', permitsName, ttl); "
                + "end; "
          + "end; "
          + "return res;",
            Arrays.asList(getRawName(), getValueName(), getClientValueName(), getPermitsName(), getClientPermitsName()),
            value, System.currentTimeMillis(), random);
}

Method Explanation:

  • The parameter value is the number of requested tokens (permits).
  • Generates a random byte array random used to construct the member in the sorted set (ensuring uniqueness).
  • Executes the Lua script to perform the token acquisition logic:
    1. Retrieves the rate limiter configuration (rate, interval, type).
    2. Selects the corresponding valueName and permitsName based on the mode (OVERALL or PER_CLIENT).
    3. Checks if the requested number of tokens exceeds the rate limit (rate).
    4. Attempts to remove expired token records from permitsName and releases the corresponding tokens.
    5. Determines if enough tokens are currently available; if so, deducts tokens and records the issued tokens in the sorted set; if not, calculates the required wait time.
    6. Updates the expiration time of the keys (if keepAliveTime is configured, refreshes to that time; otherwise, inherits the TTL of the main key).
  • Return result: nil indicates immediate success; a number indicates the wait time in milliseconds (includes a 3ms compensation).

Lua Script

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
-- Retrieve configuration from the main key's hash (KEYS[1])
local rate = redis.call('hget', KEYS[1], 'rate');
local interval = redis.call('hget', KEYS[1], 'interval');
local type = redis.call('hget', KEYS[1], 'type');
-- If any configuration is missing, the rate limiter is not initialized; throw an error
assert(rate ~= false and interval ~= false and type ~= false, 'RateLimiter is not initialized')

-- Default to using the global token bucket (OVERALL)
local valueName = KEYS[2];      -- Key storing the current number of available tokens
local permitsName = KEYS[4];    -- Sorted set storing issued tokens
-- If the type is PER_CLIENT (type == '1'), use client-specific keys
if type == '1' then
    valueName = KEYS[3];        -- Client-specific token counter
    permitsName = KEYS[5];      -- Client-specific token set
end;

-- Validate that the requested number of tokens does not exceed the rate limit
assert(tonumber(rate) >= tonumber(ARGV[1]), 'Requested permits amount cannot exceed defined rate');

-- Get the current number of available tokens
local currentValue = redis.call('get', valueName);
local res;  -- Return value: nil for success, otherwise wait time

if currentValue ~= false then
    -- Token counter already exists; handle expired token records
    -- Calculate the score range: up to (current time - interval)
    local expiredValues = redis.call('zrangebyscore', permitsName, 0, tonumber(ARGV[2]) - interval);
    local released = 0;
    -- Iterate over expired records and sum the tokens they represent
    for i, v in ipairs(expiredValues) do
        -- struct.unpack('Bc0I', v) parses the previously packed data:
        -- 'B' indicates the length of the random bytes (first byte), 'c0' reads that many random bytes, 'I' reads an integer (token count)
        local random, permits = struct.unpack('Bc0I', v);
        released = released + permits;
    end;

    if released > 0 then
        -- Remove expired records from the sorted set
        redis.call('zremrangebyscore', permitsName, 0, tonumber(ARGV[2]) - interval);
        -- If adding the released tokens would exceed the rate limit, we need to recalculate the actual available count
        if tonumber(currentValue) + released > tonumber(rate) then
            -- Get all currently non-expired records (tokens still within the window)
            local values = redis.call('zrange', permitsName, 0, -1);
            local used = 0;
            for i, v in ipairs(values) do
                local random, permits = struct.unpack('Bc0I', v);
                used = used + permits;
            end;
            -- Available tokens = rate - used tokens
            currentValue = tonumber(rate) - used;
        else
            -- Otherwise, simply add the released tokens back
            currentValue = tonumber(currentValue) + released;
        end;
        -- Update the token counter
        redis.call('set', valueName, currentValue);
    end;

    -- Check if enough tokens are available
    if tonumber(currentValue) < tonumber(ARGV[1]) then
        -- Not enough; need to wait: get the score (timestamp) of the first (earliest) record
        local firstValue = redis.call('zrange', permitsName, 0, 0, 'withscores');
        -- Calculate wait time = 3 + interval - (current time - earliest record's time)
        -- 3 is a conservative extra offset to prevent clock drift
        res = 3 + interval - (tonumber(ARGV[2]) - tonumber(firstValue[2]));
    else
        -- Enough tokens; record this issuance
        -- struct.pack packs the random length, random bytes, and token count into a string, used as the member in the sorted set
        redis.call('zadd', permitsName, ARGV[2], struct.pack('Bc0I', string.len(ARGV[3]), ARGV[3], ARGV[1]));
        -- Deduct the available tokens
        redis.call('decrby', valueName, ARGV[1]);
        -- Success, return nil
        res = nil;
    end;
else
    -- First time acquiring tokens; initialize the counter
    redis.call('set', valueName, rate);
    redis.call('zadd', permitsName, ARGV[2], struct.pack('Bc0I', string.len(ARGV[3]), ARGV[3], ARGV[1]));
    redis.call('decrby', valueName, ARGV[1]);
    res = nil;
end;

-- Handle expiration times
local keepAliveTime = redis.call('hget', KEYS[1], 'keepAliveTime');
if (keepAliveTime ~= false and tonumber(keepAliveTime) > 0) then
    -- If keepAliveTime is configured, refresh the expiration of all relevant keys
    redis.call('pexpire', KEYS[1], keepAliveTime);
    redis.call('pexpire', valueName, keepAliveTime);
    redis.call('pexpire', permitsName, keepAliveTime);
else
    -- Otherwise, attempt to inherit the TTL of the main key
    local ttl = redis.call('pttl', KEYS[1]);
    if ttl > 0 then
        redis.call('pexpire', valueName, ttl);
        redis.call('pexpire', permitsName, ttl);
    end;
end;

-- Return the result: nil or the wait time
return res;

Summary

  • Redisson's rate limiter stores configuration in a Hash, the current available tokens in a String, and issued token timestamps and counts in a Sorted Set.
  • Lua scripts ensure the atomicity of operations such as acquiring tokens, releasing expired tokens, and deducting tokens.
  • Two modes are supported: OVERALL mode and PER_CLIENT mode, where the latter maintains independent counts for each client.
  • The keepAliveTime parameter can automatically clean up inactive rate limiter keys to prevent memory leaks.
  • The wait time returned by tryAcquireAsync includes a 3ms extra offset to compensate for network and clock errors, making the rate limiting more accurate.

The code comments and script analysis above demonstrate how Redisson leverages Redis's atomic features to implement a high-performance distributed token bucket rate limiter.

References

  • LLM