go-zero源码阅读-固定时间窗口限流

一. 简介

go-zero框架的PeriodLimit限流是固定时间窗口限流，通过lua脚本操作redis, 对key进行计数并搭配过期时间实现限流

二. 源码

1. 代码路径

core/limit/periodlimit.go

2.  核心struct

type (
	// PeriodOption defines the method to customize a PeriodLimit.
    // PeriodLimit进行参数设置
	PeriodOption func(l *PeriodLimit)

	// A PeriodLimit is used to limit requests during a period of time.
	PeriodLimit struct {
        // 时间窗口大小，单位: 秒
		period     int
        // 限流阈值
		quota      int
        // 依赖redis实现
		limitStore *redis.Redis
        // key前缀
		keyPrefix  string
		align      bool
	}
)

结构体字段围绕redis的set命令设计, align 为false时，时间窗口是固定的值period, 如果为true则是周期性的(如period=3时: 3->2->1->3->2。。。)

3. 限流方法

// Take requests a permit, it returns the permit state.
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) {
// 执行lua脚本进行限流判断
// ARGV为限流阈值和周期
	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
	}
}

核心逻辑就是执行lua脚本进行限流检测，最后校验方法返回值，判断是否限流

4. lua脚本

const periodScript = `local limit = tonumber(ARGV[1]) // 获取阈值
local window = tonumber(ARGV[2]) // 获取限流窗口
local current = redis.call("INCRBY", KEYS[1], 1) // 限流计数, key 加1
// 只有值为1的时候，进行过期时间设置 
// 通过过期时间 控制限流窗口，key过期后，重新计数，进入新的窗口
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`

通过对key进行计数，key的值为1时，设置过期时间，来达到固定时间窗口限流的目的

5. 初始化配置设置

type (
	// PeriodOption defines the method to customize a PeriodLimit.
    // 参数设置 函数
	PeriodOption func(l *PeriodLimit)
)

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
}

// 参数设置
func Align() PeriodOption {
	return func(l *PeriodLimit) {
		l.align = true
	}
}

由初始化方法传入 多个配置设置函数进行配置设置， 由外部控制配置设置，增加了初始化配置的灵活性