4. Redis 单机令牌桶实现：使用 String + Lua 脚本实现原子性令牌获取

在分布式系统限流场景中，原子性是核心要求。Redis 单机令牌桶的经典实现方案采用 String 数据结构 + Lua 脚本，将令牌的获取、扣减、时间窗口刷新封装为一个原子操作，避免并发竞争导致令牌超发。

4.1 为什么选择 String + Lua

Redis 提供了多种数据结构，但令牌桶场景对原子性和性能有特殊要求：

Lua 脚本在 Redis 中执行时具有事务性：脚本内的所有命令要么全部执行成功，要么全部不执行，且执行期间不会被其他客户端命令打断。这天然解决了令牌桶的并发安全问题。

4.2 令牌桶核心数据结构设计

使用一个 Redis String 键存储令牌桶的当前状态，采用 JSON 或自定义格式序列化：

Key:   rate_limiter:{bucket_name}
Value: {
  "tokens":      10,      // 当前可用令牌数
  "last_refill": 1712345678,  // 上次补充时间戳（秒）
  "capacity":    20,      // 桶容量
  "refill_rate": 5        // 每秒补充令牌数
}

实际生产环境中，为了减少序列化开销，通常将字段拆分为多个 String 键，或使用 Hash 结构。但 String + Lua 方案在单机场景下最为简洁高效。

4.3 Lua 脚本实现原子获取令牌

以下 Lua 脚本实现了令牌桶的核心逻辑：按时间补充令牌 → 判断是否足够 → 扣减并返回成功/失败。

-- KEYS[1]: 令牌桶键名
-- ARGV[1]: 请求的令牌数
-- ARGV[2]: 当前时间戳（秒）
-- ARGV[3]: 桶容量
-- ARGV[4]: 每秒补充速率

local key = KEYS[1]
local requested = tonumber(ARGV[1])
local now = tonumber(ARGV[2])
local capacity = tonumber(ARGV[3])
local refill_rate = tonumber(ARGV[4])

-- 1. 获取当前桶状态
local bucket_str = redis.call('GET', key)
local tokens, last_refill

if bucket_str then
    -- 解析存储的 JSON 字符串
    local bucket = cjson.decode(bucket_str)
    tokens = bucket.tokens
    last_refill = bucket.last_refill
else
    -- 桶不存在，初始化为满桶
    tokens = capacity
    last_refill = now
end

-- 2. 计算时间流逝并补充令牌
local elapsed = now - last_refill
if elapsed > 0 then
    local refill_tokens = math.floor(elapsed * refill_rate)
    tokens = math.min(capacity, tokens + refill_tokens)
    last_refill = now
end

-- 3. 判断是否足够
if tokens >= requested then
    tokens = tokens - requested
    -- 更新 Redis
    local new_bucket = cjson.encode({
        tokens = tokens,
        last_refill = last_refill,
        capacity = capacity,
        refill_rate = refill_rate
    })
    redis.call('SET', key, new_bucket)
    return 1  -- 成功获取
else
    return 0  -- 令牌不足
end

4.4 原子性获取流程

下图展示了客户端调用 Lua 脚本获取令牌的完整流程：

4.5 客户端调用示例

在 Java 客户端（Jedis/Redisson）中调用上述 Lua 脚本：

// 加载脚本，获取 SHA
String script = loadScriptFromFile("token_bucket.lua");
String sha = jedis.scriptLoad(script);

// 每次请求时调用
String key = "rate_limiter:api_gateway";
long now = System.currentTimeMillis() / 1000;
int capacity = 20;
int refillRate = 5;
int requestTokens = 1;

Object result = jedis.evalsha(sha, 
    Arrays.asList(key),
    Arrays.asList(
        String.valueOf(requestTokens),
        String.valueOf(now),
        String.valueOf(capacity),
        String.valueOf(refillRate)
    )
);

boolean allowed = ((Long) result) == 1L;

4.6 方案对比与选型建议