Go實現用戶每日限額的方法(例一天隻能領三次福利)

-- KYES[1]:限流器key
-- ARGV[1]:qos,單位時間內最多請求次數
-- ARGV[2]:單位限流窗口時間
-- 請求最大次數,等於p.quota
local limit = tonumber(ARGV[1])
-- 窗口即一個單位限流周期,這裡用過期模擬窗口效果,等於p.permit
local window = tonumber(ARGV[2])
-- 請求次數+1,獲取請求總數
local current = redis.call("INCRBY",KYES[1],1)
-- 如果是第一次請求,則設置過期時間並返回 成功
if current == 1 then
  redis.call("expire",KYES[1],window)
  return 1
-- 如果當前請求數量小於limit則返回 成功
elseif current < limit then
  return 1
-- 如果當前請求數量==limit則返回 最後一次請求
elseif current == limit then
  return 2
-- 請求數量>limit則返回 失敗
else
  return 0
end

type (
  // PeriodOption defines the method to customize a PeriodLimit.
  // go中常見的option參數模式
  // 如果參數非常多，推薦使用此模式來設置參數
  PeriodOption func(l *PeriodLimit)

  // A PeriodLimit is used to limit requests during a period of time.
  // 固定時間窗口限流器
  PeriodLimit struct {
    // 窗口大小，單位s
    period     int
    // 請求上限
    quota      int
    // 存儲
    limitStore *redis.Redis
    // key前綴
    keyPrefix  string
    // 線性限流，開啟此選項後可以實現周期性的限流
    // 比如quota=5時，quota實際值可能會是5.4.3.2.1呈現出周期性變化
    align      bool
  }
)

// Take requests a permit, it returns the permit state.
// 執行限流
// 註意一下返回值：
// 0：表示錯誤，比如可能是redis故障、過載
// 1：允許
// 2：允許但是當前窗口內已到達上限
// 3：拒絕
func (h *PeriodLimit) Take(key string) (int, error) {
  // 執行lua腳本
  resp, err := h.limitStore.Eval(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
  }
}

// 計算過期時間也就是窗口時間大小
// 如果align==true
// 線性限流，開啟此選項後可以實現周期性的限流
// 比如quota=5時，quota實際值可能會是5.4.3.2.1呈現出周期性變化
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
}