基于Go+Redis实现3种限流算法

// AllowFixedWindow 每 window 秒允许 limit 次
func AllowFixedWindow(ctx context.Context, rdb *redis.Client, key string, limit int64, window time.Duration) (bool, error) {
	// INCR + EXPIRE 的经典组合；key 带时间窗口戳天然滚动
	winKey := fmt.Sprintf("rl:%s:%d", key, time.Now().Unix()/int64(window.Seconds()))
	pipe := rdb.Pipeline()
	incr := pipe.Incr(ctx, winKey)
	pipe.Expire(ctx, winKey, window+time.Second) // 多留 1 秒防边界误差
	if _, err := pipe.Exec(ctx); err != nil {
		return false, err
	}
	return incr.Val() <= limit, nil
}

// AllowSlidingWindow 滑动窗口：窗口内请求精确计数
func AllowSlidingWindow(ctx context.Context, rdb *redis.Client, key string, limit int64, window time.Duration) (bool, error) {
	now := time.Now().UnixMilli()
	winStart := now - window.Milliseconds()
	zkey := "rl:sw:" + key

	// 用 WATCH 保证"清理旧记录 + 计数 + 写入新记录"原子（第 13 篇的套路）
	var allowed bool
	_, err := rdb.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
		pipe.ZRemRangeByScore(ctx, zkey, "0", fmt.Sprintf("%d", winStart)) // 清掉窗口外的
		cnt := pipe.ZCard(ctx, zkey)
		if _, err := pipe.Exec(ctx); err != nil { // 占位：go-redis 会在函数返回后统一 EXEC
			return err
		}
		allowed = cnt.Val() < limit
		if allowed {
			pipe.ZAdd(ctx, zkey, redis.Z{Score: float64(now), Member: fmt.Sprintf("%d", now)})
			pipe.Expire(ctx, zkey, window+time.Second)
		}
		return nil
	})
	if err != nil {
		// 并发冲突时 WATCH 事务失败，直接按拒绝处理（限流宁可错杀）
		return false, nil
	}
	return allowed, nil
}

type TokenBucket struct {
	rdb   *redis.Client
	burst int64
	rate  float64 // 每秒补充的令牌数
}

func (b *TokenBucket) Allow(ctx context.Context, key string) (bool, error) {
	hkey := "rl:tb:" + key
	var allowed bool
	_, err := b.rdb.TxPipelined(ctx, func(pipe redis.Pipeliner) error {
		hm := pipe.HGetAll(ctx, hkey)
		if _, err := pipe.Exec(ctx); err != nil {
			return err
		}
		vals := hm.Val()
		tokens := b.burst
		last := time.Now().UnixMilli()
		if t, ok := vals["tokens"]; ok {
			fmt.Sscanf(t, "%d", &tokens)
			fmt.Sscanf(vals["ts"], "%d", &last)
		}
		// 按流逝时间补充令牌，上限 burst
		refill := float64(time.Now().UnixMilli()-last) / 1000 * b.rate
		tokens = int64(math.Min(float64(tokens)+refill, float64(b.burst)))

		if tokens > 0 {
			allowed = true
			tokens--
		}
		pipe.HSet(ctx, hkey, "tokens", tokens, "ts", time.Now().UnixMilli())
		pipe.Expire(ctx, hkey, time.Hour)
		return nil
	})
	if err != nil {
		return false, nil
	}
	return allowed, nil
}

func RateLimitMiddleware(rdb *redis.Client, tb *TokenBucket) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// 按用户/IP 维度限流：key 取登录用户 id 或 X-Forwarded-For
			key := r.Header.Get("X-Forwarded-For")
			ok, err := tb.Allow(r.Context(), key)
			if err != nil || !ok {
				w.Header().Set("Retry-After", "1")
				http.Error(w, `{"code":429,"msg":"请求过于频繁"}`, http.StatusTooManyRequests)
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}