golang+redis 实现分布式限流

实现分布式滑动窗口限流

复制代码

package main

import (
	"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]) then
            redis.call("ZADD", KEYS[1], ARGV[3], ARGV[3]) -- 添加当前请求时间戳
            redis.call("EXPIRE", KEYS[1], ARGV[4]) -- 设置过期时间，防止 key 长期存在
            return 1
        else
            return 0
        end
    `

	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 main

import (
	"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 令牌桶 Key
	bucketSize = 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 then
            redis.call("SET", KEYS[1], ARGV[2])
            tokens = ARGV[2]
        end
        tokens = tonumber(tokens)
        if tokens > 0 then
            redis.call("DECR", KEYS[1])
            return 1
        else
            return 0
        end
    `
	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) // 模拟请求间隔
	}
}