go-redis限流算法：令牌桶与漏桶实现
项目地址: https://gitcode.com/GitHub_Trending/go/go-redis 为什么需要限流？
在现代分布式系统中，限流（Rate Limiting）是保护服务稳定性的关键技术。当面对突发流量、异常访问或资源竞争时，有效的限流策略可以：
- ✅ 防止系统过载崩溃
- ✅ 保证服务质量（QoS）
- ✅ 公平分配资源
- ✅ 防御异常流量
Redis作为高性能内存数据库，是实现分布式限流的理想选择。本文将深入探讨两种经典限流算法在go-redis中的实现：令牌桶算法（Token Bucket）和漏桶算法（Leaky Bucket）。
限流算法核心概念
令牌桶算法 (Token Bucket)
工作原理：
- 系统以固定速率向桶中添加令牌
- 每个请求需要消耗一个令牌才能通过
- 桶有最大容量限制，防止令牌无限累积
- 允许突发流量（桶中有积累的令牌时）
漏桶算法 (Leaky Bucket)
工作原理：
- 请求以任意速率进入桶中
- 桶以固定速率向外漏出请求进行处理
- 桶有最大容量，超出的请求会被丢弃
- 平滑流量，防止突发请求
基于go-redis的令牌桶实现
数据结构设计
type TokenBucket struct {
client *redis.Client
key string // Redis键名
capacity int64 // 桶容量
rate time.Duration // 令牌生成速率
lastUpdate time.Time // 最后更新时间
}
核心实现代码
package main
import (
"context"
"fmt"
"math"
"time"
"github.com/redis/go-redis/v9"
)
type TokenBucketLimiter struct {
client *redis.Client
key string
capacity int64
rate time.Duration // 每个令牌的时间间隔
}
func NewTokenBucketLimiter(client *redis.Client, key string, capacity int64, rate time.Duration) *TokenBucketLimiter {
return &TokenBucketLimiter{
client: client,
key: key,
capacity: capacity,
rate: rate,
}
}
func (l *TokenBucketLimiter) Allow(ctx context.Context) (bool, error) {
now := time.Now().UnixNano()
window := int64(l.rate) // 时间窗口纳秒数
// 使用Lua脚本保证原子性操作
script := `
local key = KEYS[1]
local now = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local capacity = tonumber(ARGV[3])
-- 获取当前令牌数和最后更新时间
local data = redis.call('HMGET', key, 'tokens', 'last_update')
local tokens = tonumber(data[1] or capacity)
local last_update = tonumber(data[2] or now)
-- 计算应该补充的令牌数
local elapsed = now - last_update
local refill_tokens = math.floor(elapsed / window)
-- 更新令牌数，不超过容量
tokens = math.min(tokens + refill_tokens, capacity)
-- 如果有令牌可用
if tokens >= 1 then
tokens = tokens - 1
redis.call('HMSET', key, 'tokens', tokens, 'last_update', now)
return 1
else
-- 更新最后更新时间但不消耗令牌
redis.call('HSET', key, 'last_update', now)
return 0
end
`
result, err := l.client.Eval(ctx, script, []string{l.key}, now, window, l.capacity).Result()
if err != nil {
return false, err
}
return result.(int64) == 1, nil
}
func (l *TokenBucketLimiter) AllowN(ctx context.Context, n int64) (bool, error) {
if n <= 0 {
return true, nil
}
if n > l.capacity {
return false, nil
}
now := time.Now().UnixNano()
window := int64(l.rate)
script := `
local key = KEYS[1]
local now = tonumber(ARGV[1])
local window = tonumber(ARGV[2])
local capacity = tonumber(ARGV[3])
local n = tonumber(ARGV[4])
local data = redis.call('HMGET', key, 'tokens', 'last_update')
local tokens = tonumber(data[1] or capacity)
local last_update = tonumber(data[2] or now)
local elapsed = now - last_update
local refill_tokens = math.floor(elapsed / window)
tokens = math.min(tokens + refill_tokens, capacity)
if tokens >= n then
tokens = tokens - n
redis.call('HMSET', key, 'tokens', tokens, 'last_update', now)
return 1
else
redis.call('HSET', key, 'last_update', now)
return 0
end
`
result, err := l.client.Eval(ctx, script, []string{l.key}, now, window, l.capacity, n).Result()
if err != nil {
return false, err
}
return result.(int64) == 1, nil
}
使用示例
func main() {
// 初始化Redis客户端
rdb := redis.NewClient(&redis.Options{
Addr: "localhost:6379",
Password: "", // 无密码
DB: 0, // 默认数据库
})
ctx := context.Background()
// 创建限流器：容量10个令牌，每秒生成1个令牌
limiter := NewTokenBucketLimiter(rdb, "api:rate_limit", 10, time.Second)
// 模拟请求
for i := 0; i < 15; i++ {
allowed, err := limiter.Allow(ctx)
if err != nil {
fmt.Printf("Error: %v\n", err)
continue
}
if allowed {
fmt.Printf("Request %d: ✅ Allowed\n", i+1)
} else {
fmt.Printf("Request %d: ❌ Rate limited\n", i+1)
}
time.Sleep(100 * time.Millisecond)
}
}
基于go-redis的漏桶算法实现
数据结构设计
type LeakyBucket struct {
client *redis.Client
key string
capacity int64 // 桶容量
rate time.Duration // 处理速率
}
核心实现代码
type LeakyBucketLimiter struct {
client *redis.Client
key string
capacity int64
rate time.Duration
}
func NewLeakyBucketLimiter(client *redis.Client, key string, capacity int64, rate time.Duration) *LeakyBucketLimiter {
return &LeakyBucketLimiter{
client: client,
key: key,
capacity: capacity,
rate: rate,
}
}
func (l *LeakyBucketLimiter) TryAcquire(ctx context.Context) (bool, time.Duration, error) {
now := time.Now().UnixNano()
leakInterval := int64(l.rate) // 漏出间隔
script := `
local key = KEYS[1]
local now = tonumber(ARGV[1])
local leak_interval = tonumber(ARGV[2])
local capacity = tonumber(ARGV[3])
-- 获取桶状态
local data = redis.call('HMGET', key, 'water_level', 'last_leak_time')
local water_level = tonumber(data[1] or 0)
local last_leak_time = tonumber(data[2] or now)
-- 计算应该漏出的水量
local elapsed = now - last_leak_time
local leaks = math.floor(elapsed / leak_interval)
-- 更新水位
water_level = math.max(0, water_level - leaks)
last_leak_time = now - (elapsed % leak_interval)
-- 尝试加水
if water_level < capacity then
water_level = water_level + 1
redis.call('HMSET', key, 'water_level', water_level, 'last_leak_time', last_leak_time)
return {1, 0} -- 允许通过，无等待时间
else
-- 计算需要等待的时间
local next_leak_time = last_leak_time + leak_interval
local wait_time = next_leak_time - now
redis.call('HMSET', key, 'last_leak_time', last_leak_time)
return {0, wait_time} -- 需要等待
end
`
result, err := l.client.Eval(ctx, script, []string{l.key}, now, leakInterval, l.capacity).Result()
if err != nil {
return false, 0, err
}
results := result.([]interface{})
allowed := results[0].(int64) == 1
waitNanos := results[1].(int64)
return allowed, time.Duration(waitNanos), nil
}
使用示例
func main() {
rdb := redis.NewClient(&redis.Options{
Addr: "localhost:6379",
Password: "",
DB: 0,
})
ctx := context.Background()
// 创建漏桶限流器：容量5，每秒处理2个请求
limiter := NewLeakyBucketLimiter(rdb, "leaky:bucket", 5, time.Second/2)
for i := 0; i < 10; i++ {
allowed, waitTime, err := limiter.TryAcquire(ctx)
if err != nil {
fmt.Printf("Error: %v\n", err)
continue
}
if allowed {
fmt.Printf("Request %d: ✅ Immediate processing\n", i+1)
} else {
fmt.Printf("Request %d: ⏳ Need to wait %v\n", i+1, waitTime)
time.Sleep(waitTime)
// 重试
allowed, _, _ = limiter.TryAcquire(ctx)
if allowed {
fmt.Printf("Request %d: ✅ Processed after wait\n", i+1)
}
}
time.Sleep(100 * time.Millisecond)
}
}
性能优化技巧
1. 管道化操作 (Pipelining)
func BatchAllow(ctx context.Context, limiter *TokenBucketLimiter, n int) ([]bool, error) {
pipe := limiter.client.Pipeline()
results := make([]*redis.Cmd, n)
for i := 0; i < n; i++ {
results[i] = pipe.Eval(ctx, tokenBucketScript, []string{limiter.key},
time.Now().UnixNano(), int64(limiter.rate), limiter.capacity, 1)
}
if _, err := pipe.Exec(ctx); err != nil {
return nil, err
}
allowed := make([]bool, n)
for i, cmd := range results {
result, _ := cmd.Result()
allowed[i] = result.(int64) == 1
}
return allowed, nil
}
2. 本地缓存优化
type CachedLimiter struct {
redisLimiter *TokenBucketLimiter
localTokens int64
lastSync time.Time
mu sync.Mutex
}
func (c *CachedLimiter) Allow(ctx context.Context) (bool, error) {
c.mu.Lock()
defer c.mu.Unlock()
// 先尝试本地令牌
if c.localTokens > 0 {
c.localTokens--
return true, nil
}
// 本地无令牌，从Redis获取一批
allowed, err := c.redisLimiter.AllowN(ctx, 10) // 一次获取10个
if err != nil {
return false, err
}
if allowed {
c.localTokens = 9 // 使用1个，缓存9个
return true, nil
}
return false, nil
}
算法对比与选择指南
| 特性 | 令牌桶算法 | 漏桶算法 |
|---|---|---|
| 突发流量处理 | ✅ 支持 | ❌ 不支持 |
| 流量平滑 | ❌ 不支持 | ✅ 支持 |
| 实现复杂度 | 中等 | 简单 |
| 内存使用 | 低 | 低 |
| 适用场景 | API限流、用户配额 | 流量整形、队列处理 |
选择建议：
-
选择令牌桶当：
- 需要允许合理的突发流量
- 用户API调用限制
- 需要精确控制请求速率
-
选择漏桶当：
- 需要绝对平滑的流量输出
- 消息队列处理
- 防止流量突刺
生产环境最佳实践
1. 监控与告警
type MonitoringLimiter struct {
limiter RateLimiter
metrics *prometheus.CounterVec
redis *redis.Client
}
func (m *MonitoringLimiter) Allow(ctx context.Context) (bool, error) {
allowed, err := m.limiter.Allow(ctx)
// 记录指标
if allowed {
m.metrics.WithLabelValues("allowed").Inc()
} else {
m.metrics.WithLabelValues("denied").Inc()
}
// 检查Redis健康状态
if err != nil && errors.Is(err, redis.ErrClosed) {
log.Error("Redis connection lost")
}
return allowed, err
}
2. 降级策略
func WithFallback(primary, fallback RateLimiter) RateLimiter {
return &fallbackLimiter{primary: primary, fallback: fallback}
}
type fallbackLimiter struct {
primary, fallback RateLimiter
}
func (f *fallbackLimiter) Allow(ctx context.Context) (bool, error) {
allowed, err := f.primary.Allow(ctx)
if err != nil {
// 主限流器失败，使用备用
return f.fallback.Allow(ctx)
}
return allowed, nil
}
3. 多级限流架构
常见问题与解决方案
Q1: Redis宕机怎么办？
A: 实现降级策略，在Redis不可用时切换到本地内存限流
Q2: 如何避免Race Condition？
A: 使用Lua脚本保证原子性操作，避免多客户端竞争
Q3: 性能瓶颈在哪里？
A: Redis网络IO是主要瓶颈，可通过管道化和本地缓存优化
Q4: 如何测试限流效果？
A: 使用压力测试工具模拟并发请求，验证限流准确性
总结
通过go-redis实现令牌桶和漏桶限流算法，我们获得了：
🎯 高可用性：基于Redis的分布式特性 🎯 精确控制：Lua脚本保证原子性操作
🎯 灵活配置：支持多种限流策略和参数 🎯 易于监控：与现有监控体系无缝集成
选择适合业务场景的限流算法，结合go-redis的强大功能，可以为你的分布式系统提供可靠的流量保护机制。
💡 提示: 在实际生产环境中，建议结合具体业务需求进行参数调优和压力测试，确保限流策略既保护系统又不过度限制正常流量。
创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
