Go 语言结合 Redis 实现固定窗口、滑动窗口、令牌桶和漏桶限流算法的示例代码

固定窗口算法

• 原理：将时间划分为固定大小的窗口，在每个窗口内对请求进行计数。如果请求数超过设定的阈值，则拒绝后续请求，直到进入下一个窗口。
• 代码：

package mainimport ("fmt""time""github.com/go-redis/redis/v8"
)// rdb 是全局的 Redis 客户端实例，用于与 Redis 服务器进行交互
var rdb = redis.NewClient(&redis.Options{Addr:     "localhost:6379", // Redis 服务器地址和端口Password: "",               // Redis 密码，这里为空表示无密码DB:       0,                // 使用默认的数据库编号
})// 定义固定窗口算法的常量
const (windowSize = 60 // 窗口大小，单位为秒limit      = 100 // 窗口内允许的最大请求数
)// fixedWindowRateLimit 函数实现了固定窗口限流算法
func fixedWindowRateLimit(key string) bool {// 获取当前时间的 Unix 时间戳currentTime := time.Now().Unix()// 计算当前窗口的起始时间windowStart := currentTime - (currentTime % windowSize)// 构建当前窗口对应的 Redis keywindowKey := fmt.Sprintf("%s:%d", key, windowStart)// 对 Redis 中的当前窗口计数器进行自增操作，并获取自增后的计数值count, err := rdb.Incr(rdb.Context(), windowKey).Result()if err != nil {// 如果自增操作出现错误，打印错误信息并返回 falsefmt.Println("Error incrementing counter:", err)return false}// 如果计数值为 1，说明是该窗口的第一个请求，为该窗口设置过期时间if count == 1 {rdb.Expire(rdb.Context(), windowKey, time.Duration(windowSize)*time.Second)}// 如果计数值超过了允许的最大请求数，返回 false 表示请求被限流if count > limit {return false}// 否则，返回 true 表示请求通过return true
}func main() {// 调用 fixedWindowRateLimit 函数进行限流判断if fixedWindowRateLimit("api_request") {fmt.Println("请求通过")} else {fmt.Println("请求被限流")}
}

• 优缺点：实现简单，但可能会出现临界问题，即在窗口切换的瞬间可能会出现流量高峰。

滑动窗口算法

• 原理：将固定窗口进一步细分，通过多个小窗口来统计请求数，随着时间的推移，窗口不断滑动，从而更精确地控制流量。
• 代码：

package mainimport ("fmt""time""github.com/go-redis/redis/v8"
)// rdb 是全局的 Redis 客户端实例，用于与 Redis 服务器进行交互
var rdb = redis.NewClient(&redis.Options{Addr:     "localhost:6379", // Redis 服务器地址和端口Password: "",               // Redis 密码，这里为空表示无密码DB:       0,                // 使用默认的数据库编号
})// 定义滑动窗口算法的常量
const (windowSize = 60 // 窗口大小，单位为秒limit      = 100 // 窗口内允许的最大请求数
)// slidingWindowRateLimit 函数实现了滑动窗口限流算法
func slidingWindowRateLimit(key string) bool {// 获取当前时间的 Unix 时间戳currentTime := time.Now().Unix()// 移除 Redis 有序集合中时间戳小于当前窗口起始时间的元素rdb.ZRemRangeByScore(rdb.Context(), key, "0", fmt.Sprintf("%d", currentTime-windowSize))// 向 Redis 有序集合中添加当前时间戳，分数为当前时间戳rdb.ZAdd(rdb.Context(), key, &redis.Z{Score:  float64(currentTime),Member: currentTime,})// 获取 Redis 有序集合中当前窗口内的元素数量count, err := rdb.ZCard(rdb.Context(), key).Result()if err != nil {// 如果获取元素数量操作出现错误，打印错误信息并返回 falsefmt.Println("Error getting count:", err)return false}// 如果元素数量超过了允许的最大请求数，返回 false 表示请求被限流if count > limit {return false}// 否则，返回 true 表示请求通过return true
}func main() {// 调用 slidingWindowRateLimit 函数进行限流判断if slidingWindowRateLimit("api_request") {fmt.Println("请求通过")} else {fmt.Println("请求被限流")}
}

• 优缺点：比固定窗口算法更精确地控制流量，但实现相对复杂，且需要更多的 Redis 操作。

令牌桶算法

• 原理：系统以固定的速率向令牌桶中添加令牌，每个请求需要从令牌桶中获取一个或多个令牌才能被处理。如果令牌桶中没有足够的令牌，则请求被拒绝。
• 代码：

package mainimport ("fmt""time""github.com/go-redis/redis/v8"
)// rdb 是全局的 Redis 客户端实例，用于与 Redis 服务器进行交互
var rdb = redis.NewClient(&redis.Options{Addr:     "localhost:6379", // Redis 服务器地址和端口Password: "",               // Redis 密码，这里为空表示无密码DB:       0,                // 使用默认的数据库编号
})// 定义令牌桶算法的常量
const (capacity        = 100        // 令牌桶的容量rate            = 1          // 令牌生成速率，每秒生成 1 个令牌requiredTokens  = 1          // 每个请求所需的令牌数
)// tokenBucketScript 是用于执行令牌桶算法的 Lua 脚本
var tokenBucketScript = `
-- 获取传入的令牌桶 key
local tokens_key = KEYS[1]
-- 构建存储上次更新时间的 key
local last_time_key = tokens_key .. ":last_time"
-- 获取令牌桶容量
local capacity = tonumber(ARGV[2])
-- 获取令牌生成速率
local rate = tonumber(ARGV[3])
-- 获取当前时间
local now = tonumber(ARGV[1])
-- 获取每个请求所需的令牌数
local required_tokens = tonumber(ARGV[4])-- 从 Redis 中获取上次更新时间
local last_time = tonumber(redis.call('get', last_time_key))
-- 从 Redis 中获取当前令牌数
local tokens = tonumber(redis.call('get', tokens_key))-- 如果上次更新时间为空，说明是首次请求，初始化相关值
if last_time == nil thenlast_time = nowtokens = capacity
end-- 计算从上次更新到现在生成的令牌数
local generated_tokens = (now - last_time) * rate
-- 更新令牌数，确保不超过令牌桶容量
tokens = math.min(capacity, tokens + generated_tokens)-- 如果令牌数小于请求所需的令牌数，说明令牌不足，拒绝请求
if tokens < required_tokens thenredis.call('set', last_time_key, last_time)redis.call('set', tokens_key, tokens)return 0
else-- 否则，处理请求，扣除相应的令牌数tokens = tokens - required_tokensredis.call('set', last_time_key, now)redis.call('set', tokens_key, tokens)return 1
end
`// tokenBucketRateLimit 函数调用 Lua 脚本来实现令牌桶限流算法
func tokenBucketRateLimit(key string) bool {// 获取当前时间的 Unix 时间戳currentTime := time.Now().Unix()// 执行 Lua 脚本，并获取执行结果result, err := rdb.Eval(rdb.Context(), tokenBucketScript, []string{key}, currentTime, capacity, rate, requiredTokens).Int64()if err != nil {// 如果执行脚本出现错误，打印错误信息并返回 falsefmt.Println("Error executing script:", err)return false}// 如果结果为 1，表示请求通过；否则，表示请求被限流return result == 1
}func main() {// 调用 tokenBucketRateLimit 函数进行限流判断if tokenBucketRateLimit("api_request") {fmt.Println("请求通过")} else {fmt.Println("请求被限流")}
}

• 优缺点：能够平滑地控制流量，允许一定程度的突发流量，但实现相对复杂。

漏桶限流算法

• 原理：请求就像水一样注入漏桶，漏桶以固定的速率处理请求。如果请求的注入速率超过漏桶的处理速率，多余的请求将被丢弃。
• 代码：

package mainimport ("fmt""time""github.com/go-redis/redis/v8"
)// rdb 是全局的 Redis 客户端实例，用于与 Redis 服务器进行交互
var rdb = redis.NewClient(&redis.Options{Addr:     "localhost:6379", // Redis 服务器地址和端口Password: "",               // Redis 密码，这里为空表示无密码DB:       0,                // 使用默认的数据库编号
})// 定义漏桶算法的常量
const (capacity     = 100        // 漏桶的容量rate         = 1          // 漏桶的流出速率，每秒流出 1 个单位requestSize  = 1          // 每个请求的大小
)// leakyBucketScript 是用于执行漏桶算法的 Lua 脚本
var leakyBucketScript = `
-- 获取传入的漏桶 key
local bucket_key = KEYS[1]
-- 构建存储上次更新时间的 key
local last_time_key = bucket_key .. ":last_time"
-- 获取漏桶容量
local capacity = tonumber(ARGV[2])
-- 获取漏桶流出速率
local rate = tonumber(ARGV[3])
-- 获取当前时间
local now = tonumber(ARGV[1])
-- 获取每个请求的大小
local request_size = tonumber(ARGV[4])-- 从 Redis 中获取上次更新时间
local last_time = tonumber(redis.call('get', last_time_key))
-- 从 Redis 中获取当前漏桶中的水量
local water = tonumber(redis.call('get', bucket_key))-- 如果上次更新时间为空，说明是首次请求，初始化相关值
if last_time == nil thenlast_time = nowwater = 0
end-- 计算从上次更新到现在流出的水量
local outflow = (now - last_time) * rate
-- 更新漏桶中的水量，确保不小于 0
water = math.max(0, water - outflow)-- 如果漏桶中的水量加上当前请求的大小超过了漏桶容量，说明漏桶已满，拒绝请求
if water + request_size > capacity thenredis.call('set', last_time_key, last_time)redis.call('set', bucket_key, water)return 0
else-- 否则，处理请求，增加漏桶中的水量water = water + request_sizeredis.call('set', last_time_key, now)redis.call('set', bucket_key, water)return 1
end
`// leakyBucketRateLimit 函数调用 Lua 脚本来实现漏桶限流算法
func leakyBucketRateLimit(key string) bool {// 获取当前时间的 Unix 时间戳currentTime := time.Now().Unix()// 执行 Lua 脚本，并获取执行结果result, err := rdb.Eval(rdb.Context(), leakyBucketScript, []string{key}, currentTime, capacity, rate, requestSize).Int64()if err != nil {// 如果执行脚本出现错误，打印错误信息并返回 falsefmt.Println("Error executing script:", err)return false}// 如果结果为 1，表示请求通过；否则，表示请求被限流return result == 1
}func main() {// 调用 leakyBucketRateLimit 函数进行限流判断if leakyBucketRateLimit("api_request") {fmt.Println("请求通过")} else {fmt.Println("请求被限流")}
}

• 优缺点：可以平滑地处理请求，保证请求以固定的速率被处理，但无法应对突发流量。

总结

算法	实现复杂度	限流精度	突发流量处理	适用场景
固定窗口算法	简单	低	差	流量平稳、精度要求低的场景
滑动窗口算法	较复杂	高	一般	流量波动大、精度要求高的场景
令牌桶算法	较复杂	高	好	可容忍突发流量、需平滑限流的场景
漏桶算法	较复杂	高	差	对流量稳定性要求极高的场景