4. 令牌桶算法实战：使用Python实现一个基于Redis的分布式令牌桶，支持高并发原子操作

4.1 为什么分布式限流需要令牌桶 + Redis

在金融压力测试场景中，单机令牌桶无法应对多实例部署的流量均衡问题。当多个API网关实例共享同一套限流策略时，必须依赖外部原子化存储来同步令牌状态。Redis凭借其单线程模型、Lua脚本支持和毫秒级响应，成为分布式令牌桶最合适的载体。

4.2 令牌桶核心算法与Redis数据结构设计

分布式令牌桶在Redis中存储三个关键字段：

• last_refill_time：上次补充令牌的时间戳（毫秒）
• tokens：当前桶内可用令牌数（浮点数）
• max_tokens：桶容量上限（固定配置）
• refill_rate：每秒补充速率（令牌/秒）

算法流程：每次请求到达时，先根据当前时间与last_refill_time的差值计算应补充的令牌数，再判断是否允许消耗一个令牌。

4.3 Lua脚本：原子化令牌获取与补充

Redis Lua脚本保证整个“补充-判断-扣减”过程在单线程中执行，无需分布式锁。以下脚本是生产级实现：

-- 分布式令牌桶Lua脚本
-- KEYS[1] = 桶的Redis键名
-- ARGV[1] = 当前时间戳（毫秒）
-- ARGV[2] = 桶容量 max_tokens
-- ARGV[3] = 补充速率 refill_rate（令牌/秒）
-- ARGV[4] = 每次请求消耗的令牌数（通常为1）

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

-- 获取当前桶状态
local bucket = redis.call('HMGET', key, 'tokens', 'last_refill_time')
local tokens = tonumber(bucket[1])
local last_refill = tonumber(bucket[2])

-- 初始化
if tokens == nil then
    tokens = max_tokens
    last_refill = now
end

-- 计算时间差并补充令牌
local elapsed = now - last_refill
if elapsed > 0 then
    local refill_tokens = elapsed * refill_rate / 1000.0
    tokens = math.min(max_tokens, tokens + refill_tokens)
end

-- 判断是否允许通过
local allowed = 0
if tokens >= cost then
    tokens = tokens - cost
    allowed = 1
end

-- 写回Redis
redis.call('HMSET', key, 'tokens', tokens, 'last_refill_time', now)
-- 设置TTL防止内存泄漏（24小时）
redis.call('EXPIRE', key, 86400)

return {allowed, tokens}

4.4 Python客户端封装：高并发调用

使用redis-py的register_script方法预加载Lua脚本，避免每次请求都传输脚本内容：

import time
import redis

class RedisTokenBucket:
    def __init__(self, redis_client, key, max_tokens, refill_rate):
        self.redis = redis_client
        self.key = key
        self.max_tokens = max_tokens
        self.refill_rate = refill_rate
        # 预加载Lua脚本，返回SHA
        self.script_sha = self._load_script()

    def _load_script(self):
        lua_script = """
        -- 上述Lua脚本内容（略）
        """
        return self.redis.script_load(lua_script)

    def try_acquire(self, cost=1):
        now_ms = int(time.time() * 1000)
        result = self.redis.evalsha(
            self.script_sha,
            1,  # KEYS数量
            self.key,
            now_ms,
            self.max_tokens,
            self.refill_rate,
            cost
        )
        allowed = result[0] == 1
        remaining = result[1]
        return allowed, remaining

# 使用示例
r = redis.StrictRedis(host='localhost', port=6379, decode_responses=True)
bucket = RedisTokenBucket(r, 'api:rate:order', max_tokens=100, refill_rate=10)

for i in range(200):
    allowed, tokens = bucket.try_acquire()
    if allowed:
        print(f"请求{i}: 通过，剩余令牌{tokens:.2f}")
    else:
        print(f"请求{i}: 被限流")
        time.sleep(0.1)

4.5 高并发原子操作验证

在金融压力测试中，使用10个并发协程同时请求令牌桶，验证原子性：

4.6 生产环境关键配置参数