详解go-zero如何实现计数器限流

// core/limit/periodlimit.go
const periodScript = `local limit = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local current = redis.call("INCRBY", KEYS[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`

type (
    // PeriodOption defines the method to customize a PeriodLimit.
    PeriodOption func(l *PeriodLimit)
    // A PeriodLimit is used to limit requests during a period of time.
    PeriodLimit struct {
        period     int  // 窗口大小，单位 s
        quota      int  // 请求上限
        limitStore *redis.Redis
        keyPrefix  string   // key 前缀
        align      bool
    }
)
// NewPeriodLimit returns a PeriodLimit with given parameters.
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
}

// key 就是需要被限制的资源标识
func (h *PeriodLimit) Take(key string) (int, error) {
    return h.TakeCtx(context.Background(), key)
}
// 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.EvalCtx(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
    }
}