-- KYES[1]:限流器key
-- ARGV[1]:qos,單位時(shí)間內(nèi)最多請(qǐng)求次數(shù)
-- ARGV[2]:單位限流窗口時(shí)間
-- 請(qǐng)求最大次數(shù),等于p.quota
local limit = tonumber(ARGV[1])
-- 窗口即一個(gè)單位限流周期,這里用過(guò)期模擬窗口效果,等于p.permit
local window = tonumber(ARGV[2])
-- 請(qǐng)求次數(shù)+1,獲取請(qǐng)求總數(shù)
local current = redis.call("INCRBY",KYES[1],1)
-- 如果是第一次請(qǐng)求,則設(shè)置過(guò)期時(shí)間并返回 成功
if current == 1 then
  redis.call("expire",KYES[1],window)
  return 1
-- 如果當(dāng)前請(qǐng)求數(shù)量小于limit則返回 成功
elseif current < limit then
  return 1
-- 如果當(dāng)前請(qǐng)求數(shù)量==limit則返回 最后一次請(qǐng)求
elseif current == limit then
  return 2
-- 請(qǐng)求數(shù)量>limit則返回 失敗
else
  return 0
end

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

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

// Take requests a permit, it returns the permit state.
// 執(zhí)行限流
// 注意一下返回值：
// 0：表示錯(cuò)誤，比如可能是redis故障、過(guò)載
// 1：允許
// 2：允許但是當(dāng)前窗口內(nèi)已到達(dá)上限
// 3：拒絕
func (h *PeriodLimit) Take(key string) (int, error) {
? // 執(zhí)行l(wèi)ua腳本
? 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
? }
}

// 計(jì)算過(guò)期時(shí)間也就是窗口時(shí)間大小
// 如果align==true
// 線性限流，開(kāi)啟此選項(xiàng)后可以實(shí)現(xiàn)周期性的限流
// 比如quota=5時(shí)，quota實(shí)際值可能會(huì)是5.4.3.2.1呈現(xiàn)出周期性變化
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
}