golang+redis 实现分布式限流

实现分布式滑动窗口限流

package mainimport ("context""fmt""time""github.com/go-redis/redis/v8"
)// Redis 客户端
var ctx = context.Background()
var rdb = redis.NewClient(&redis.Options{Addr:     "localhost:6379", // Redis 服务器地址Password: "",               // 无密码DB:       0,                // 默认数据库
})// 限流配置
const (windowSize  = 10               // 时间窗口（秒）maxRequests = 5                // 允许的最大请求数requestKey  = "sliding_window" // Redis Key
)// 限流检查
func isAllowed() bool {now := time.Now().Unix()           // 当前时间戳minTime := now - int64(windowSize) // 窗口开始时间// 使用 Lua 脚本保证操作的原子性luaScript := `redis.call("ZREMRANGEBYSCORE", KEYS[1], "-inf", ARGV[1]) -- 清除过期请求local reqCount = redis.call("ZCOUNT", KEYS[1], "-inf", "+inf") -- 获取当前窗口内请求数if reqCount < tonumber(ARGV[2]) thenredis.call("ZADD", KEYS[1], ARGV[3], ARGV[3]) -- 添加当前请求时间戳redis.call("EXPIRE", KEYS[1], ARGV[4]) -- 设置过期时间，防止 key 长期存在return 1elsereturn 0end`result, err := rdb.Eval(ctx, luaScript, []string{requestKey}, minTime, maxRequests, now, windowSize).Int()if err != nil {fmt.Println("Redis error:", err)return false}return result == 1
}func main() {for i := 0; i < 10; i++ {if isAllowed() {fmt.Println("Request allowed")} else {fmt.Println("Rate limit exceeded")}time.Sleep(1 * time.Second) // 模拟请求间隔}
}

令牌桶限流

package mainimport ("context""fmt""time""github.com/go-redis/redis/v8"
)// Redis 客户端
var ctx = context.Background()
var rdb = redis.NewClient(&redis.Options{Addr:     "localhost:6379", // Redis 地址Password: "",               // Redis 认证密码（无则留空）DB:       0,                // 使用默认数据库
})// 令牌桶配置
const (bucketKey  = "token_bucket" // Redis 令牌桶 KeybucketSize = 10             // 令牌桶容量refillRate = 2              // 每秒补充的令牌数量
)// 初始化令牌桶
func initBucket() {rdb.Set(ctx, bucketKey, bucketSize, 0) // 初始化令牌桶，存满令牌
}// 获取令牌
func getToken() bool {luaScript := `local tokens = redis.call("GET", KEYS[1])if tokens == false thenredis.call("SET", KEYS[1], ARGV[2])tokens = ARGV[2]endtokens = tonumber(tokens)if tokens > 0 thenredis.call("DECR", KEYS[1])return 1elsereturn 0end`result, err := rdb.Eval(ctx, luaScript, []string{bucketKey}, 1, bucketSize).Int()if err != nil {fmt.Println("Redis Error:", err)return false}return result == 1
}// 定期补充令牌
func refillTokens() {ticker := time.NewTicker(time.Second)defer ticker.Stop()for range ticker.C {rdb.IncrBy(ctx, bucketKey, refillRate)currentTokens, _ := rdb.Get(ctx, bucketKey).Int()if currentTokens > bucketSize {rdb.Set(ctx, bucketKey, bucketSize, 0)}fmt.Printf("Tokens refilled: %d\n", currentTokens)}
}func main() {initBucket()go refillTokens() // 启动令牌补充协程for i := 0; i < 20; i++ {if getToken() {fmt.Println("Request allowed")} else {fmt.Println("Rate limit exceeded")}time.Sleep(50 * time.Millisecond) // 模拟请求间隔}
}