介面限流
按賬戶/IP/Token限流
定時加載配置

time/rate 原理

time/rate 是go提供的官方限流工具

原理: 令牌桶，以一個恒定的速度往桶里放入令牌，而如果請求需要被處理，則需要先從桶里獲取一個令牌，當桶里沒有令牌可取時，則拒絕服務

我首先想到的方案初始化定時任務，每隔多長時間往桶里放一令牌，但是相比之下，time/rate的實作方式就是更優雅了

關鍵變數:

令牌放入桶的速度 limit 每秒產生多少token
桶的容量 burst

已知變數：

最后一次消耗令牌的時間 last

簡化公式：

當前可用令牌數，tokens := (now - last).Seconds() * limit

完整代碼是：

// advance calculates and returns an updated state for lim resulting from the passage of time.
// lim is not changed.
func (lim *Limiter) advance(now time.Time) (newNow time.Time, newLast time.Time, newTokens float64) {
  last := lim.last
  if now.Before(last) {
    last = now
  }
?
  // Avoid making delta overflow below when last is very old.
  // 這里限制令牌數不超過burst，lim.tokens 是上一次結余
  maxElapsed := lim.limit.durationFromTokens(float64(lim.burst) - lim.tokens)
  elapsed := now.Sub(last)
  if elapsed > maxElapsed {
    elapsed = maxElapsed
  }
?
  // Calculate the new number of tokens, due to time that passed.
  delta := lim.limit.tokensFromDuration(elapsed)
  tokens := lim.tokens + delta // 上一次結余 + 時間增量
  if burst := float64(lim.burst); tokens > burst {
    tokens = burst
  }
?
  return now, last, tokens
}

比起定時器來說是不是很優雅？這里有個比較繞的點：limit 的單位，n/s, 每秒多少個令牌

// Every converts a minimum time interval between events to a Limit.
func Every(interval time.Duration) Limit {
  if interval <= 0 {
    return Inf
  }
  return 1 / Limit(interval.Seconds())
}

計算可用tokens的公式就更好理解了，用 (now (當前時間) - last(上一次消耗令牌的時間) ).Seconds() * limit 就是增量令牌數了

time/rate 改造

time/rate 只是給出了最基礎的實作，絕大部分的需求

介面 > IP/user 多層限流
動態加載配置
在回應中給出剩余請求次數，下一次重置時間

其實 time/rate 的成本是非常低的，所以我們只需要多層map結構就可以滿足, 每次初始化一個 rate.Limiter 就可以滿足

type ApiKeyRateLimit interface {
  Load()
  GetLimiter(api, key string) (limiter *Limiter, ok bool)
}
?
type ApiKeyRateLimitImpl struct {
  hub map[string]map[string]*rate.Limiter
}

動態加載配置，可以使用etcd訂閱配置，我實作的是定時從資料庫里撈配置

?
type ApiAccountRateLimit struct {
  ctx context.Context
  lim rate.ApiKeyRateLimit
}
?
func NewApiAccountRateLimit(ctx context.Context, limit rate.ApiKeyRateLimit) *ApiAccountRateLimit {
  r := &ApiAccountRateLimit{
    ctx: ctx,
    lim: limit,
  }
  go r.load()
  return r
}
?
func (h *ApiAccountRateLimit) Hook() gin.HandlerFunc {
  return func(ctx *gin.Context) {
    path := ctx.Request.URL.Path
    acctId, ok := ctx.Get(consts.OpenapiAccountId)
    if !ok {
      ctx.JSON(http.StatusOK, gin.H{"code": http.StatusUnauthorized, "message": http.StatusText(http.StatusUnauthorized)})
      ctx.Abort()
      return
    }
    limiter, ok := h.lim.GetLimiter(path, acctId.(string))
    if !ok {
      ctx.JSON(http.StatusOK, gin.H{"code": http.StatusUnauthorized, "message": http.StatusText(http.StatusUnauthorized)})
      ctx.Abort()
    }
    // 這里time/rat沒有放出Advance，需要改一下
    _, _, tokens := limiter.Advance(time.Now())
    burst := limiter.Burst()
    ctx.Header("X-RateLimit-Limit", strconv.Itoa(burst))
    ctx.Header("X-RateLimit-Remaining", strconv.Itoa(int(tokens)))
    if !limiter.Allow() {
      ctx.JSON(http.StatusOK, gin.H{"code": http.StatusTooManyRequests, "message": http.StatusText(http.StatusTooManyRequests)})
      ctx.Abort()
      return
    }
    ctx.Next()
  }
}
?
func (h *ApiAccountRateLimit) load() {
  ticker := time.NewTicker(5 * time.Minute)
  for {
    select {
    case <-h.ctx.Done():
      logger.Errorf("quit: RateLimit load config")
    case <-ticker.C:
      h.lim.Load()
    }
  }

go time/rate 介面-賬戶/IP 限流

time/rate 原理

time/rate 改造

推薦閱讀