go-zero 两个限流器都踩了坑，最后自行实现了一个分布式令牌桶

背景

项目使用 go-zero 框架，Redis 部署为 Sentinel（一主两从），通过 go-redis 的 UniversalClient 接入。

给文章服务 API 加限流时，我先试了 go-zero 内置的 PeriodLimit，发现限流失效。然后翻文档找到了 TokenLimiter------这个是令牌桶算法，应该靠谱吧？结果一样不兼容。

两个官方限流器都走不通，最终基于令牌桶算法自行实现了一个兼容 Sentinel UniversalClient 的限流中间件。

这篇文章记录完整的踩坑过程、根因分析和最终方案。

---

一、第一次尝试：PeriodLimit（固定窗口计数器）

go-zero 的 PeriodLimit 位于 core/limit/periodlimit.go，核心逻辑：

func (p *PeriodLimit) Take() (bool, error) {
    val, err := p.redis.Incr(p.key)
    if val == 1 {
        p.redis.Expire(p.key, p.period)
    }
    return val <= p.quota, nil
}

固定窗口计数器 ：用 INCR + EXPIRE 组合，在一个时间窗口内统计请求次数。

踩坑点：底层依赖 go-zero 的 redis.Redis，这个类型包装的是 *redis.Client（单节点客户端），不支持 Sentinel 模式。

更严重的是，INCR 和 EXPIRE 是两条独立命令。在 UniversalClient 的连接池下，无法保证它们在同一连接上顺序执行，加上从节点的复制延迟，计数根本不准。

结论：不可用。原因有二------类型不兼容 + 多命令原子性无法保证。

---

二、第二次尝试：TokenLimiter（令牌桶）

PeriodLimit 不行，go-zero 还提供了 TokenLimiter（core/limit/tokenlimit.go）。它用的是 Lua 脚本，算法也是令牌桶，看起来正是我需要的：

-- go-zero tokenscript.lua
local rate = tonumber(ARGV[1])
local capacity = tonumber(ARGV[2])
local now = tonumber(ARGV[3])
local requested = tonumber(ARGV[4])

local fill_time = capacity/rate
local ttl = math.floor(fill_time*2)
local last_tokens = tonumber(redis.call("get", KEYS[1]))
if last_tokens == nil then
    last_tokens = capacity
end
local last_refreshed = tonumber(redis.call("get", KEYS[2]))
if last_refreshed == nil then
    last_refreshed = 0
end

local delta = math.max(0, now-last_refreshed)
local filled_tokens = math.min(capacity, last_tokens+(delta*rate))
local allowed = filled_tokens >= requested
local new_tokens = filled_tokens
if allowed then
    new_tokens = filled_tokens - requested
end

redis.call("setex", KEYS[1], ttl, new_tokens)
redis.call("setex", KEYS[2], ttl, now)
return allowed

算法没问题，但一看 Go 侧的结构体：

type TokenLimiter struct {
    store *redis.Redis  // go-zero 的 Redis 封装，底层是 *redis.Client
    ...
}

func NewTokenLimiter(rate, burst int, store *redis.Redis, key string) *TokenLimiter

又是 *redis.Redis。

go-zero 的 redis.Redis 类型直接包装的是 *redis.Client，不是 redis.UniversalClient 接口。所以 TokenLimiter 和 PeriodLimit 面对的是 同一个根本问题------go-zero 的 Redis 层不支持 Sentinel 模式。

结论：算法对了，但客户端类型绑死了，还是不可用。

---

三、根因分析：为什么 go-zero 的限流器都不兼容 Sentinel

两个限流器的依赖链：

PeriodLimit / TokenLimiter
        ↓
    *redis.Redis          ← go-zero 自己的封装
        ↓
    *redis.Client         ← go-redis 单节点客户端
        ↓
    连接单个 Redis 实例

而我的项目：

redis.UniversalClient     ← go-redis 通用接口
        ↓
*redis.FailoverClusterClient  ← Sentinel 模式
        ↓
通过 Sentinel 发现 Master/Replica，读写分离

go-zero 的整个 Redis 层（core/stores/redis）都是围绕 *redis.Client 设计的。PeriodLimit 和 TokenLimiter 都直接依赖这个类型，无法注入 UniversalClient。

有人可能会想：能不能把 go-zero 的 redis.Redis 包一层，让它指向 Sentinel 的 Master？技术上可以，但：

1. 需要魔改 go-zero 源码或写适配层，维护成本高
2. 即使指向了 Master，go-zero 的连接池管理和 Sentinel 故障转移的交互未经验证
3. 升级 go-zero 版本时适配层可能失效

正路是绕开 go-zero 的 Redis 层，直接基于 UniversalClient 自行实现。

---

四、自行实现：基于 UniversalClient 的分布式令牌桶

既然 go-zero 的两个限流器都被 *redis.Redis 绑死了，那就直接用 go-redis 的 UniversalClient 接口，参考 TokenLimiter 的令牌桶算法自行实现。

4.1 为什么选令牌桶

策略	优点	缺点
固定窗口	实现简单	窗口边界突发（两个窗口交界处可放过 2x 请求）
滑动窗口	平滑	需要存储每个请求的时间戳，内存开销大
令牌桶	允许突发 + 平滑限流	实现稍复杂

go-zero 的 TokenLimiter 选的也是令牌桶，说明这个算法确实是限流场景的主流选择。我的工作是把它从 *redis.Redis 上适配到了 UniversalClient 上，并补充了一些工程细节。

4.2 Lua 脚本

local tokens_key = KEYS[1]       -- 当前剩余令牌数
local timestamp_key = KEYS[2]    -- 上次刷新时间戳
local rate = tonumber(ARGV[1])   -- 每秒补充速率
local capacity = tonumber(ARGV[2]) -- 桶容量（突发上限）
local now = tonumber(ARGV[3])    -- 当前时间（毫秒）
local requested = tonumber(ARGV[4]) -- 本次请求消耗的令牌数

-- 计算 TTL：桶填满所需时间的 2 倍，确保 key 不会过早淘汰
local fill_time = capacity / rate
local ttl = math.floor(fill_time * 2)
if ttl < 1 then
    ttl = 1
end

-- 惰性计算：读取上次剩余令牌，按时间差补充
local last_tokens = tonumber(redis.call("get", tokens_key))
if last_tokens == nil then
    last_tokens = capacity    -- 首次访问，满桶
end

local last_refreshed = tonumber(redis.call("get", timestamp_key))
if last_refreshed == nil then
    last_refreshed = 0
end

-- 核心：按时间差计算应补充的令牌数
local delta_ms = math.max(0, now - last_refreshed)
local delta_sec = delta_ms / 1000.0
local filled_tokens = math.min(capacity, last_tokens + (delta_sec * rate))

-- 判断并扣减
local allowed = filled_tokens >= requested

if allowed then
    redis.call("setex", tokens_key, ttl, filled_tokens - requested)
    redis.call("setex", timestamp_key, ttl, now)
    return {1, 0}         -- {允许, 无需等待}
else
    local wait_sec = math.ceil((requested - filled_tokens) / rate)
    return {0, wait_sec}  -- {拒绝, 建议等待秒数}
end

和 go-zero 的 tokenscript.lua 对比，核心算法相同（都是惰性令牌桶），主要差异：

维度	go-zero TokenLimiter	自行实现
时间精度	now.Unix()（秒）	now.UnixMilli()（毫秒）
拒绝响应	返回 false（bool）	返回 {0, wait_sec}，带建议等待时间
Redis 客户端	*redis.Redis	redis.UniversalClient

4.3 为什么单 Lua 脚本能解决 Sentinel 兼容问题

go-zero TokenLimiter 的 Lua 脚本本身没问题，问题出在它通过 *redis.Redis 执行，而这个客户端不支持 Sentinel。

自行实现的版本直接使用 redis.UniversalClient：

type RateLimitMiddleware struct {
    Redis     redis.UniversalClient   // ← 关键区别
    Rate      float64
    Burst     int
    scriptSHA string
}

UniversalClient 发送 EVALSHA 命令时，因为脚本内包含写操作（SETEX），命令会被路由到 Master 节点。整个 Lua 脚本在 Master 上原子执行，不存在跨节点问题。

而 go-zero 的 PeriodLimit（INCR + EXPIRE 两条独立命令）在 UniversalClient 的连接池下无法保证原子性。TokenLimiter 的 Lua 脚本虽然原子，但 *redis.Redis 类型根本连不上 Sentinel。

一个 Lua 脚本 + UniversalClient，同时解决了原子性和 Sentinel 兼容两个问题。

4.4 Go 侧实现

初始化时预加载脚本：

func NewRateLimitMiddleware(rate float64, burst int, redisClient redis.UniversalClient) (*RateLimitMiddleware, error) {
    sha, err := redisClient.ScriptLoad(context.Background(), rateLimitScript).Result()
    if err != nil {
        return nil, fmt.Errorf("load rate limit script: %w", err)
    }
    return &RateLimitMiddleware{
        Redis:     redisClient,
        Rate:      rate,
        Burst:     burst,
        scriptSHA: sha,
    }, nil
}

通过 SCRIPT LOAD 预先将 Lua 脚本加载到 Redis 并缓存 SHA 值，后续调用使用 EVALSHA 减少网络传输。

执行时优先 EvalSha，NOSCRIPT 时回退 Eval：

func (m *RateLimitMiddleware) evalScript(ctx context.Context, keys []string, args ...interface{}) (interface{}, error) {
    result, err := m.Redis.EvalSha(ctx, m.scriptSHA, keys, args...).Result()
    if err != nil && strings.HasPrefix(err.Error(), "NOSCRIPT") {
        return m.Redis.Eval(ctx, rateLimitScript, keys, args...).Result()
    }
    return result, err
}

Redis 在内存压力下可能驱逐已缓存的脚本，此时 EVALSHA 返回 NOSCRIPT，捕获后回退到 EVAL 保证可用性。

---

五、生产环境考量

5.1 Fail-Closed 策略

if err != nil {
    logx.WithContext(r.Context()).Errorf("rate limit eval error (fail closed): %v", err)
    w.WriteHeader(http.StatusInternalServerError)
    w.Write([]byte(`{"code":500, "msg":"系统繁忙，请稍后再试"}`))
    return
}

Redis 不可用时拒绝请求而非降级放行。这是写操作（创建/更新文章）的限流，降级放行等于在最脆弱时放弃保护。

对比 go-zero TokenLimiter 的策略：它在 Redis 不可用时会降级到本地 golang.org/x/time/rate 限流器。这在多实例部署下有个问题------每个实例独立限流，N 个实例的实际阈值就是配置值的 N 倍。Fail-Closed 虽然严格，但限流行为是可预期的。

5.2 限流粒度

基于 JWT 中解析出的 userId 进行限流，而非 IP。因为：

• API 已经挂了 JWT 认证，userId 是可信的
• 同一 NAT 下的多个用户不会被误限
• 用户无法通过换 IP 绕过限流

5.3 Key 设计

req_limit:tokens:user:{userId}   -- 剩余令牌数
req_limit:ts:user:{userId}       -- 上次刷新时间戳

每个用户两个 key，TTL 为 2 * (burst / rate) 秒。以默认配置 rate=10, burst=20 计算，TTL 为 4 秒。空闲用户的 key 会自动过期，不会无限膨胀。

---

六、配置示例

RateLimit:
  Rate: 10     # 每秒补充 10 个令牌
  Burst: 20    # 桶容量 20（允许短时突发）

这意味着：

• 稳态下每秒允许 10 个请求
• 短时突发最多允许 20 个请求
• 超过后返回 429 + Retry-After 头

---

七、总结

维度	go-zero PeriodLimit	go-zero TokenLimiter	自行实现
算法	固定窗口计数	令牌桶	令牌桶
Lua 脚本	无（INCR+EXPIRE）	有	有
Redis 客户端	*redis.Redis	*redis.Redis	redis.UniversalClient
Sentinel 兼容	不兼容	不兼容	兼容
原子性	多命令，无法保证	单脚本，但连不上 Sentinel	单脚本，路由到 Master
Redis 不可用	无降级	本地 x/time/rate 兜底	Fail-Closed 拒绝
时间精度	秒级	秒级	毫秒级
拒绝响应	无	返回 bool	返回建议等待秒数

核心收获：go-zero 的两个限流器算法都没问题，但它们都被 *redis.Redis 类型绑死了，无法接入 Sentinel 模式的 UniversalClient。这项工作的本质不是发明新算法，而是 把令牌桶算法从 *redis.Redis 适配到了 UniversalClient 上 ，同时补充了 EvalSha 预加载与 NOSCRIPT 回退、毫秒精度、Retry-After 响应等工程细节。