Skip to content

Commit 4b5eac7

Browse files
nodesire7nodesire7
committed
feat: 优化限流策略：实现滑动窗口和令牌桶算法
- 创建 utils/ratelimit.go：实现滑动窗口和令牌桶限流器 - 更新 middleware/ratelimit.go：优先使用滑动窗口，降级到令牌桶，最终降级到本地限流 - 在 app.go 中初始化限流器 - 实现 redo.md 2.5：限流策略优化（滑动窗口/令牌桶）
1 parent 4413a15 commit 4b5eac7

File tree

3 files changed

+222
-18
lines changed

3 files changed

+222
-18
lines changed

‎internal/app/app.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,9 @@ func Run() error {
3434
utils.LogWarn("Redis初始化失败，缓存功能将不可用: %v", err)
3535
}
3636
defer cache.CloseRedis()
37+
38+
// 初始化限流器（滑动窗口 + 令牌桶）
39+
middleware.InitRateLimiters()
3740

3841
// 设置Gin模式
3942
if cfg.ServerMode == "release" {

‎middleware/ratelimit.go‎

Lines changed: 57 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
/**
22
* 限流中间件
3-
* 提供API限流功能
3+
* 提供API限流功能（滑动窗口 + 令牌桶）
4+
* 实现 redo.md 2.5：限流策略优化
45
*/
56
package middleware
67

@@ -14,39 +15,77 @@ import (
1415
"golang.org/x/time/rate"
1516
)
1617

17-
var limiter = rate.NewLimiter(rate.Every(time.Second), 100) // 每秒100个请求
18+
var (
19+
// 本地限流器（令牌桶，作为降级方案）
20+
localLimiter = rate.NewLimiter(rate.Every(time.Second), 100) // 每秒100个请求
21+
22+
// Redis 滑动窗口限流器（1分钟窗口，100个请求）
23+
slidingWindowLimiter *utils.SlidingWindowLimiter
24+
25+
// Redis 令牌桶限流器（容量100，每秒补充100个令牌）
26+
tokenBucketLimiter *utils.TokenBucketLimiter
27+
)
28+
29+
// InitRateLimiters 初始化限流器（在应用启动时调用）
30+
func InitRateLimiters() {
31+
if cache.RedisClient != nil {
32+
// 滑动窗口：1分钟窗口，100个请求
33+
slidingWindowLimiter = utils.NewSlidingWindowLimiter(cache.RedisClient, time.Minute, 100)
34+
// 令牌桶：容量100，每秒补充100个令牌
35+
tokenBucketLimiter = utils.NewTokenBucketLimiter(cache.RedisClient, 100, 100.0)
36+
}
37+
}
1838

1939
// RateLimitMiddleware 限流中间件
2040
func RateLimitMiddleware() gin.HandlerFunc {
2141
return func(c *gin.Context) {
22-
// 如果Redis可用，使用分布式限流
23-
if cache.RedisClient != nil {
24-
realIP := utils.GetRealIP(c.Request)
25-
key := "rate_limit:" + realIP
26-
count, err := cache.RedisClient.Incr(cache.Ctx, key).Result()
42+
realIP := utils.GetRealIP(c.Request)
43+
key := "rate_limit:" + realIP
44+
45+
var allowed bool
46+
var err error
47+
48+
// 优先使用滑动窗口（更精确）
49+
if slidingWindowLimiter != nil {
50+
allowed, err = slidingWindowLimiter.Allow(key)
2751
if err == nil {
28-
if count == 1 {
29-
cache.RedisClient.Expire(cache.Ctx, key, time.Second)
30-
}
31-
if count > 100 { // 每秒100个请求
52+
if !allowed {
3253
c.JSON(http.StatusTooManyRequests, gin.H{
3354
"error": "请求过于频繁，请稍后再试",
3455
})
3556
c.Abort()
3657
return
3758
}
59+
c.Next()
60+
return
3861
}
39-
} else {
40-
// 使用本地限流
41-
if !limiter.Allow() {
42-
c.JSON(http.StatusTooManyRequests, gin.H{
43-
"error": "请求过于频繁，请稍后再试",
44-
})
45-
c.Abort()
62+
}
63+
64+
// 降级到令牌桶
65+
if tokenBucketLimiter != nil {
66+
allowed, err = tokenBucketLimiter.Allow(key)
67+
if err == nil {
68+
if !allowed {
69+
c.JSON(http.StatusTooManyRequests, gin.H{
70+
"error": "请求过于频繁，请稍后再试",
71+
})
72+
c.Abort()
73+
return
74+
}
75+
c.Next()
4676
return
4777
}
4878
}
4979

80+
// 最终降级到本地限流
81+
if !localLimiter.Allow() {
82+
c.JSON(http.StatusTooManyRequests, gin.H{
83+
"error": "请求过于频繁，请稍后再试",
84+
})
85+
c.Abort()
86+
return
87+
}
88+
5089
c.Next()
5190
}
5291
}

‎utils/ratelimit.go‎

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
/**
2+
* 限流工具
3+
* 实现滑动窗口限流算法（基于 Redis）
4+
* 实现 redo.md 2.5：限流策略优化（滑动窗口/令牌桶）
5+
*/
6+
package utils
7+
8+
import (
9+
"context"
10+
"fmt"
11+
"time"
12+
13+
"github.com/redis/go-redis/v9"
14+
)
15+
16+
// SlidingWindowLimiter 滑动窗口限流器
17+
type SlidingWindowLimiter struct {
18+
client *redis.Client
19+
ctx context.Context
20+
window time.Duration // 时间窗口
21+
limit int // 窗口内允许的最大请求数
22+
}
23+
24+
// NewSlidingWindowLimiter 创建滑动窗口限流器
25+
func NewSlidingWindowLimiter(client *redis.Client, window time.Duration, limit int) *SlidingWindowLimiter {
26+
return &SlidingWindowLimiter{
27+
client: client,
28+
ctx: context.Background(),
29+
window: window,
30+
limit: limit,
31+
}
32+
}
33+
34+
// Allow 检查是否允许请求（滑动窗口算法）
35+
// 使用 Redis Sorted Set 实现滑动窗口
36+
func (l *SlidingWindowLimiter) Allow(key string) (bool, error) {
37+
if l.client == nil {
38+
// Redis 不可用时，允许请求（降级）
39+
return true, nil
40+
}
41+
42+
now := time.Now()
43+
windowStart := now.Add(-l.window)
44+
45+
// 1. 移除窗口外的旧记录
46+
pipe := l.client.Pipeline()
47+
pipe.ZRemRangeByScore(l.ctx, key, "0", fmt.Sprintf("%d", windowStart.UnixNano()))
48+
49+
// 2. 统计当前窗口内的请求数
50+
pipe.ZCard(l.ctx, key)
51+
52+
// 3. 添加当前请求的时间戳
53+
pipe.ZAdd(l.ctx, key, redis.Z{
54+
Score: float64(now.UnixNano()),
55+
Member: fmt.Sprintf("%d", now.UnixNano()),
56+
})
57+
58+
// 4. 设置过期时间（窗口长度 + 1秒缓冲）
59+
pipe.Expire(l.ctx, key, l.window+time.Second)
60+
61+
results, err := pipe.Exec(l.ctx)
62+
if err != nil {
63+
return true, err // 出错时允许请求（降级）
64+
}
65+
66+
// 获取当前窗口内的请求数（在添加当前请求之前）
67+
var count int64
68+
if len(results) >= 2 {
69+
countCmd, ok := results[1].(*redis.IntCmd)
70+
if ok {
71+
count = countCmd.Val()
72+
}
73+
}
74+
75+
// 检查是否超过限制
76+
if count >= int64(l.limit) {
77+
return false, nil
78+
}
79+
80+
return true, nil
81+
}
82+
83+
// TokenBucketLimiter 令牌桶限流器（基于 Redis）
84+
type TokenBucketLimiter struct {
85+
client *redis.Client
86+
ctx context.Context
87+
capacity int // 桶容量
88+
rate float64 // 每秒生成的令牌数
89+
refillKey string // 上次补充时间戳的 key
90+
}
91+
92+
// NewTokenBucketLimiter 创建令牌桶限流器
93+
func NewTokenBucketLimiter(client *redis.Client, capacity int, rate float64) *TokenBucketLimiter {
94+
return &TokenBucketLimiter{
95+
client: client,
96+
ctx: context.Background(),
97+
capacity: capacity,
98+
rate: rate,
99+
refillKey: "token_bucket_refill",
100+
}
101+
}
102+
103+
// Allow 检查是否允许请求（令牌桶算法）
104+
func (l *TokenBucketLimiter) Allow(key string) (bool, error) {
105+
if l.client == nil {
106+
// Redis 不可用时，允许请求（降级）
107+
return true, nil
108+
}
109+
110+
now := time.Now()
111+
tokensKey := key + ":tokens"
112+
lastRefillKey := key + ":last_refill"
113+
114+
// 使用 Lua 脚本保证原子性
115+
luaScript := `
116+
local tokens_key = KEYS[1]
117+
local last_refill_key = KEYS[2]
118+
local capacity = tonumber(ARGV[1])
119+
local rate = tonumber(ARGV[2])
120+
local now = tonumber(ARGV[3])
121+
122+
local tokens = tonumber(redis.call('GET', tokens_key) or capacity)
123+
local last_refill = tonumber(redis.call('GET', last_refill_key) or now)
124+
125+
-- 计算需要补充的令牌数
126+
local elapsed = now - last_refill
127+
local tokens_to_add = math.floor(elapsed * rate / 1000000000) -- 纳秒转秒
128+
129+
if tokens_to_add > 0 then
130+
tokens = math.min(capacity, tokens + tokens_to_add)
131+
redis.call('SET', last_refill_key, now)
132+
end
133+
134+
-- 检查是否有足够的令牌
135+
if tokens >= 1 then
136+
tokens = tokens - 1
137+
redis.call('SET', tokens_key, tokens)
138+
redis.call('EXPIRE', tokens_key, 3600)
139+
redis.call('EXPIRE', last_refill_key, 3600)
140+
return 1
141+
else
142+
redis.call('SET', tokens_key, tokens)
143+
redis.call('EXPIRE', tokens_key, 3600)
144+
redis.call('EXPIRE', last_refill_key, 3600)
145+
return 0
146+
end
147+
`
148+
149+
result, err := l.client.Eval(l.ctx, luaScript, []string{tokensKey, lastRefillKey},
150+
l.capacity, l.rate, now.UnixNano()).Result()
151+
if err != nil {
152+
return true, err // 出错时允许请求（降级）
153+
}
154+
155+
allowed, ok := result.(int64)
156+
if !ok {
157+
return true, nil
158+
}
159+
160+
return allowed == 1, nil
161+
}
162+

0 commit comments

Comments
 (0)