固定窗口

这里的固定窗口是利用 redis 结合 Lua 脚本来实现的

local limit = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local current = redis.call("INCRBY", KEYS[1], 1) // 每次 + 1
if current == 1 then
    redis.call("expire", KEYS[1], window)
end
if current < limit then
    return 1
elseif current == limit then
    return 2
else
    return 0
end

这里的 0， 1， 2 是跟源码的业务代码相结合的

internalOverQuota = 0 // 超配额
internalAllowed   = 1 // 未超过限制次数
internalHitQuota  = 2 // 等于限制次数

初始化：

func NewPeriodLimit(period, quota int, limitStore *redis.Redis, keyPrefix string,
   opts ...PeriodOption) *PeriodLimit {
   limiter := &PeriodLimit{
      period:     period,
      quota:      quota,
      limitStore: limitStore,
      keyPrefix:  keyPrefix,
   }

   for _, opt := range opts {
      opt(limiter)
   }

   return limiter
}

构造函数传入

period
- 时间间隔
quota
- 次数
limitStore
keyPrefix

Take 方法

可以看到这个包对外只提供了 Take 方法，传入 key 返回 int 类型的值进行业务判断：

Unknown = iota
// Allowed means allowed state.
Allowed
// HitQuota means this request exactly hit the quota.
HitQuota
// OverQuota means passed the quota.
OverQuota

现在看看具体实现

// TakeCtx requests a permit with context, it returns the permit state.
func (h *PeriodLimit) TakeCtx(ctx context.Context, key string) (int, error) {
   resp, err := h.limitStore.ScriptRunCtx(ctx, periodScript, []string{h.keyPrefix + key}, []string{
      strconv.Itoa(h.quota),
      strconv.Itoa(h.calcExpireSeconds()),
   })
   if err != nil {
      return Unknown, err
   }

   code, ok := resp.(int64)
   if !ok {
      return Unknown, ErrUnknownCode
   }

   switch code {
   case internalOverQuota:
      return OverQuota, nil
   case internalAllowed:
      return Allowed, nil
   case internalHitQuota:
      return HitQuota, nil
   default:
      return Unknown, ErrUnknownCode
   }
}

这里就直接调用了 lua 脚本还使用了函数 calcExpireSeconds：


func (h *PeriodLimit) calcExpireSeconds() int {
   if h.align {
      now := time.Now()
      _, offset := now.Zone()
      unix := now.Unix() + int64(offset)
      return h.period - int(unix%int64(h.period))
   }

   return h.period
}

代码中的关键部分是unix%int64(h.period)，这部分代码对当前时间戳unix取模h.period。取模运算的结果是一个介于0和h.period-1之间的数，这个数表示当前时间在h.period时间周期中的位置。然后，通过从h.period中减去这个结果，可以得到一个与h.period有关的偏移量。这个偏移量表示了当前时间距离下一个完整的h.period周期还有多少时间。

然后我们来解读 lua 脚本

这里的限流使用的是 incrby + 过期时间 来做的

在第一次调用 incrby 的时候给加一个窗口
每次请求的时候流加1
然后判断 current 与 limit 的关系

备注：Redis 的 incr 命令使用之后会返回当前的值

总感觉这个限流怪怪的🤔

令牌桶限流

在go-zero中，会监控redis的状态，如果redis挂了，则启用本地限流，并启动异步任务探查Redis的存活情况

Redis限流

首先看一下限流器的定义

// A TokenLimiter controls how frequently events are allowed to happen with in one second.
type TokenLimiter struct {
   // 每秒生产速率
   rate           int
   
   // 桶容量
   burst          int
   
   
   store          *redis.Redis
   
   // redis key
   tokenKey       string
   
   // 桶刷新时间
   timestampKey   string
   
   
   rescueLock     sync.Mutex
   
   // redis是否存活
   redisAlive     uint32
   
   // redis监控探测任务标识
   monitorStarted bool
   
   // redis故障时采取进程内令牌桶限流器
   rescueLimiter  *xrate.Limiter
}


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 设为令牌桶的容量
    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

令牌桶容量 cap，令牌生产速率 rate
获取上次剩余令牌数量 last_tokens
获取上次状态保存的时间 last_timestamp
- 如果是第一次保存则设置为 0
计算时间差 delta = max(now_timestamp - last_timestamp)
计算当前令牌桶内令牌数量
- current_tokens = max(cap, delta * rate)
return current_tokens > request_tokens

本地限流

本地限流这里就不讲了，用的是 golang.org/x/time/rate

go-zero限流

固定窗口

令牌桶限流

Redis限流

本地限流