[AI生成] 基于Redis+go+lua脚本实现qps限流

实现qps限流
每次判断是否要限流时，发送lua脚本给redis执行，控制令牌数量，刚开始令牌桶是满的。
key缓存时间是容量/速率（单位是s）+10s。
value使用hashmap，记录当前令牌数和上次时间，时间是绝对时间的秒数，取整。

package main

import (
	"context"
	"fmt"
	"time"

	"github.com/redis/go-redis/v9"
)

// TokenBucketLimiter 令牌桶限流器
type TokenBucketLimiter struct {
	client   *redis.Client // Redis客户端
	capacity int           // 令牌桶最大容量
	rate     int           // 令牌生成速率：个/秒
	script   *redis.Script // 限流Lua脚本（原子执行）
}

// NewTokenBucketLimiter 创建限流器
// capacity：桶最大令牌数
// rate：每秒生成多少令牌
func NewTokenBucketLimiter(client *redis.Client, capacity, rate int) *TokenBucketLimiter {
	// 加载Lua脚本（核心：原子性执行令牌生成+扣减）
	// redis中保存key是字符串，value是hashmap，包含两个字段：tokens和last_time
	// tokens：当前令牌数
	// last_time：上次刷新时间戳
	// last_time的单位是秒
	script := redis.NewScript(`
		-- KEYS[1]: 限流唯一key
		-- ARGV[1]: 桶容量
		-- ARGV[2]: 生成速率(个/秒)
		-- ARGV[3]: 秒
		local key = KEYS[1]
		local capacity = tonumber(ARGV[1])
		local rate = tonumber(ARGV[2])
		local now = tonumber(ARGV[3])

		-- 1. 桶不存在：初始化满令牌，直接扣减1个
		local exist = redis.call("EXISTS", key)
		if exist == 0 then
			redis.call("HMSET", key, "tokens", capacity - 1, "last_time", now)
			-- 设置过期时间，自动清理无用key
			redis.call("EXPIRE", key, capacity/rate + 10)
			return 1 -- 放行
		end

		-- 2. 获取当前令牌数和上次刷新时间
		local data = redis.call("HMGET", key, "tokens", "last_time")
		local tokens = tonumber(data[1])
		local lastTime = tonumber(data[2])

		-- 3. 计算时间差，生成新令牌 - 只生成整数个令牌
		local elapsed = now - lastTime
		local newTokens = elapsed * rate
		tokens = math.min(capacity, tokens + newTokens)
		lastTime = now

		-- 4. 尝试扣减令牌
		if tokens >= 1 then
			redis.call("HMSET", key, "tokens", tokens - 1, "last_time", lastTime)
			return 1 -- 放行
		else
			redis.call("HMSET", key, "tokens", tokens, "last_time", lastTime)
			return 0 -- 限流
		end
	`)

	return &TokenBucketLimiter{
		client:   client,
		capacity: capacity,
		rate:     rate,
		script:   script,
	}
}

// Allow 判断是否允许请求通过
// key：限流标识（如用户ID、接口URI、IP）
func (l *TokenBucketLimiter) Allow(key string) (bool, error) {
	ctx := context.Background()
	now := time.Now().Second()

	// 执行Lua脚本
	resp, err := l.script.Run(
		ctx,
		l.client,
		[]string{key}, // KEYS
		l.capacity,    // ARGV[1]
		l.rate,        // ARGV[2]
		now,           // ARGV[3]
	).Int()

	if err != nil {
		return false, err
	}

	// 1=放行，0=限流
	return resp == 1, nil
}

// ==============================================
// 测试示例
// ==============================================
func main() {
	// 1. 连接Redis
	rdb := redis.NewClient(&redis.Options{
		Addr:     "localhost:6379",
		Password: "", // 无密码留空
		DB:       0,
	})

	// 不允许突发流量
	// 2. 初始化限流器：桶容量2，每秒生成2个令牌
	limiter := NewTokenBucketLimiter(rdb, 2, 2)

	// 清理旧的Redis key
	ctx := context.Background()
	key := "api:order:create" // 限流接口
	rdb.Del(ctx, key)

	// 3. 模拟20个并发请求
	for i := 1; i <= 20; i++ {
		allow, err := limiter.Allow(key)
		now := time.Now().Format("2006-01-02 15:04:05.000")

		if err != nil {
			fmt.Printf("[%s] 请求%d：限流错误-%v\n", now, i, err)
			continue
		}
		if allow {
			fmt.Printf("[%s] 请求%d：通过\n", now, i)
		} else {
			fmt.Printf("[%s] 请求%d：被限流\n", now, i)
		}
		time.Sleep(100 * time.Millisecond) // 模拟请求间隔
	}

	// 4. 等待2s
	time.Sleep(2 * time.Second)

	// 5. 模拟20个并发请求
	for i := 1; i <= 20; i++ {
		allow, err := limiter.Allow(key)
		now := time.Now().Format("2006-01-02 15:04:05.000")

		if err != nil {
			fmt.Printf("[%s] 请求%d：限流错误-%v\n", now, i, err)
			continue
		}
		if allow {
			fmt.Printf("[%s] 请求%d：通过\n", now, i)
		} else {
			fmt.Printf("[%s] 请求%d：被限流\n", now, i)
		}
		time.Sleep(100 * time.Millisecond) // 模拟请求间隔
	}
}

 

[2026-04-25 19:58:50.586] 请求1：通过
[2026-04-25 19:58:50.687] 请求2：通过
[2026-04-25 19:58:50.790] 请求3：被限流
[2026-04-25 19:58:50.892] 请求4：被限流
[2026-04-25 19:58:50.994] 请求5：被限流
[2026-04-25 19:58:51.096] 请求6：通过
[2026-04-25 19:58:51.198] 请求7：通过
[2026-04-25 19:58:51.299] 请求8：被限流
[2026-04-25 19:58:51.401] 请求9：被限流
[2026-04-25 19:58:51.503] 请求10：被限流
[2026-04-25 19:58:51.604] 请求11：被限流
[2026-04-25 19:58:51.706] 请求12：被限流
[2026-04-25 19:58:51.808] 请求13：被限流
[2026-04-25 19:58:51.909] 请求14：被限流
[2026-04-25 19:58:52.012] 请求15：通过
[2026-04-25 19:58:52.114] 请求16：通过
[2026-04-25 19:58:52.216] 请求17：被限流
[2026-04-25 19:58:52.317] 请求18：被限流
[2026-04-25 19:58:52.419] 请求19：被限流
[2026-04-25 19:58:52.521] 请求20：被限流
[2026-04-25 19:58:54.623] 请求1：通过
[2026-04-25 19:58:54.726] 请求2：通过
[2026-04-25 19:58:54.827] 请求3：被限流
[2026-04-25 19:58:54.929] 请求4：被限流
[2026-04-25 19:58:55.032] 请求5：通过
[2026-04-25 19:58:55.133] 请求6：通过
[2026-04-25 19:58:55.235] 请求7：被限流
[2026-04-25 19:58:55.337] 请求8：被限流
[2026-04-25 19:58:55.440] 请求9：被限流
[2026-04-25 19:58:55.542] 请求10：被限流
[2026-04-25 19:58:55.644] 请求11：被限流
[2026-04-25 19:58:55.746] 请求12：被限流
[2026-04-25 19:58:55.848] 请求13：被限流
[2026-04-25 19:58:55.950] 请求14：被限流
[2026-04-25 19:58:56.052] 请求15：通过
[2026-04-25 19:58:56.153] 请求16：通过
[2026-04-25 19:58:56.256] 请求17：被限流
[2026-04-25 19:58:56.358] 请求18：被限流
[2026-04-25 19:58:56.459] 请求19：被限流
[2026-04-25 19:58:56.561] 请求20：被限流

为什么绝对时间不取小数的毫秒？
1. 确保同一秒内刚开始的前几个请求成功，后面失败，保证公平性。
2. 小数计算不准确，带来误差。