feat: 实现内存令牌桶限流器
- Limiter 接口定义(Allow + Stop 方法) - TokenBucket 实现(容量、填充速率、并发安全) - MemoryLimiter 管理多用户令牌桶 - 后台 goroutine 定期清理不活跃桶(10 分钟) - 完整单元测试覆盖(11 个测试用例,全部通过) - 边界情况处理(rate=0、capacity=0、并发安全)
This commit is contained in:
172
backend/internal/ratelimit/bucket.go
Normal file
172
backend/internal/ratelimit/bucket.go
Normal file
@@ -0,0 +1,172 @@
|
|||||||
|
package ratelimit
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/hhs/camtalk/internal/config"
|
||||||
|
)
|
||||||
|
|
||||||
|
// TokenBucket 内存令牌桶,适用于单实例部署。
|
||||||
|
type TokenBucket struct {
|
||||||
|
capacity int // 桶容量
|
||||||
|
rate float64 // 每秒填充令牌数
|
||||||
|
tokens float64 // 当前令牌数
|
||||||
|
lastRefill time.Time // 上次填充时间
|
||||||
|
mu sync.Mutex
|
||||||
|
}
|
||||||
|
|
||||||
|
// newTokenBucket 创建令牌桶。
|
||||||
|
func newTokenBucket(capacity int, rate float64) *TokenBucket {
|
||||||
|
return &TokenBucket{
|
||||||
|
capacity: capacity,
|
||||||
|
rate: rate,
|
||||||
|
tokens: float64(capacity), // 初始满桶
|
||||||
|
lastRefill: time.Now(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// allow 尝试消耗一个令牌。
|
||||||
|
func (b *TokenBucket) allow() (bool, time.Duration) {
|
||||||
|
b.mu.Lock()
|
||||||
|
defer b.mu.Unlock()
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
elapsed := now.Sub(b.lastRefill).Seconds()
|
||||||
|
|
||||||
|
// 补充令牌
|
||||||
|
newTokens := elapsed * b.rate
|
||||||
|
b.tokens = min(float64(b.capacity), b.tokens+newTokens)
|
||||||
|
b.lastRefill = now
|
||||||
|
|
||||||
|
// 尝试消耗一个令牌
|
||||||
|
if b.tokens >= 1 {
|
||||||
|
b.tokens -= 1
|
||||||
|
return true, 0
|
||||||
|
}
|
||||||
|
|
||||||
|
// 计算需要等待的时间
|
||||||
|
if b.rate == 0 {
|
||||||
|
// rate=0 时永远无法补充令牌
|
||||||
|
return false, 24 * time.Hour // 返回一个很大的值
|
||||||
|
}
|
||||||
|
retryAfter := time.Duration((1-b.tokens)/b.rate*1000) * time.Millisecond
|
||||||
|
return false, retryAfter
|
||||||
|
}
|
||||||
|
|
||||||
|
// MemoryLimiter 管理多个用户的令牌桶。
|
||||||
|
type MemoryLimiter struct {
|
||||||
|
buckets map[string]*TokenBucket
|
||||||
|
config config.RateLimitConfig
|
||||||
|
mu sync.RWMutex
|
||||||
|
stopOnce sync.Once
|
||||||
|
done chan struct{}
|
||||||
|
}
|
||||||
|
|
||||||
|
// NewMemoryLimiter 创建内存限流器。
|
||||||
|
func NewMemoryLimiter(cfg config.RateLimitConfig) *MemoryLimiter {
|
||||||
|
limiter := &MemoryLimiter{
|
||||||
|
buckets: make(map[string]*TokenBucket),
|
||||||
|
config: cfg,
|
||||||
|
done: make(chan struct{}),
|
||||||
|
}
|
||||||
|
|
||||||
|
// 启动后台清理 goroutine
|
||||||
|
go limiter.cleanup()
|
||||||
|
|
||||||
|
return limiter
|
||||||
|
}
|
||||||
|
|
||||||
|
// Allow 实现 Limiter 接口。
|
||||||
|
func (l *MemoryLimiter) Allow(ctx context.Context, key string) (bool, time.Duration) {
|
||||||
|
bucket := l.getOrCreateBucket(key)
|
||||||
|
return bucket.allow()
|
||||||
|
}
|
||||||
|
|
||||||
|
// Stop 实现 Limiter 接口。
|
||||||
|
func (l *MemoryLimiter) Stop() {
|
||||||
|
l.stopOnce.Do(func() {
|
||||||
|
close(l.done)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
// getOrCreateBucket 获取或创建令牌桶。
|
||||||
|
func (l *MemoryLimiter) getOrCreateBucket(key string) *TokenBucket {
|
||||||
|
// 先尝试读锁
|
||||||
|
l.mu.RLock()
|
||||||
|
bucket, exists := l.buckets[key]
|
||||||
|
l.mu.RUnlock()
|
||||||
|
|
||||||
|
if exists {
|
||||||
|
return bucket
|
||||||
|
}
|
||||||
|
|
||||||
|
// 需要创建新桶,升级为写锁
|
||||||
|
l.mu.Lock()
|
||||||
|
defer l.mu.Unlock()
|
||||||
|
|
||||||
|
// 双重检查(可能其他 goroutine 已创建)
|
||||||
|
bucket, exists = l.buckets[key]
|
||||||
|
if exists {
|
||||||
|
return bucket
|
||||||
|
}
|
||||||
|
|
||||||
|
// 根据 key 确定配置(简化版:假设 key 格式为 "userID:action")
|
||||||
|
cfg := l.getBucketConfig(key)
|
||||||
|
bucket = newTokenBucket(cfg.Capacity, cfg.Rate)
|
||||||
|
l.buckets[key] = bucket
|
||||||
|
|
||||||
|
return bucket
|
||||||
|
}
|
||||||
|
|
||||||
|
// getBucketConfig 根据 key 获取桶配置。
|
||||||
|
func (l *MemoryLimiter) getBucketConfig(key string) config.BucketConfig {
|
||||||
|
// 简化实现:从 key 后缀判断动作类型
|
||||||
|
// 实际使用时调用方会传递正确的 key
|
||||||
|
// 默认使用 query 配置
|
||||||
|
return l.config.Query
|
||||||
|
}
|
||||||
|
|
||||||
|
// cleanup 定期清理不活跃的桶。
|
||||||
|
func (l *MemoryLimiter) cleanup() {
|
||||||
|
ticker := time.NewTicker(10 * time.Minute)
|
||||||
|
defer ticker.Stop()
|
||||||
|
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-ticker.C:
|
||||||
|
l.removeInactiveBuckets()
|
||||||
|
case <-l.done:
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// removeInactiveBuckets 移除超过 10 分钟无活动的桶。
|
||||||
|
func (l *MemoryLimiter) removeInactiveBuckets() {
|
||||||
|
l.mu.Lock()
|
||||||
|
defer l.mu.Unlock()
|
||||||
|
|
||||||
|
now := time.Now()
|
||||||
|
for key, bucket := range l.buckets {
|
||||||
|
bucket.mu.Lock()
|
||||||
|
inactive := now.Sub(bucket.lastRefill) > 10*time.Minute
|
||||||
|
bucket.mu.Unlock()
|
||||||
|
|
||||||
|
if inactive {
|
||||||
|
delete(l.buckets, key)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// min 返回两个 float64 中的较小值。
|
||||||
|
func min(a, b float64) float64 {
|
||||||
|
if a < b {
|
||||||
|
return a
|
||||||
|
}
|
||||||
|
return b
|
||||||
|
}
|
||||||
|
|
||||||
|
// 编译期接口检查
|
||||||
|
var _ Limiter = (*MemoryLimiter)(nil)
|
||||||
203
backend/internal/ratelimit/bucket_test.go
Normal file
203
backend/internal/ratelimit/bucket_test.go
Normal file
@@ -0,0 +1,203 @@
|
|||||||
|
package ratelimit
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"sync"
|
||||||
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
"github.com/hhs/camtalk/internal/config"
|
||||||
|
"github.com/stretchr/testify/assert"
|
||||||
|
"github.com/stretchr/testify/require"
|
||||||
|
)
|
||||||
|
|
||||||
|
func TestTokenBucket_Allow_FirstRequest(t *testing.T) {
|
||||||
|
bucket := newTokenBucket(5, 0.2)
|
||||||
|
|
||||||
|
allowed, retryAfter := bucket.allow()
|
||||||
|
|
||||||
|
assert.True(t, allowed)
|
||||||
|
assert.Equal(t, time.Duration(0), retryAfter)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTokenBucket_Allow_ConsumeUntilEmpty(t *testing.T) {
|
||||||
|
bucket := newTokenBucket(3, 0.2)
|
||||||
|
|
||||||
|
// 连续消耗 3 个令牌
|
||||||
|
for i := 0; i < 3; i++ {
|
||||||
|
allowed, _ := bucket.allow()
|
||||||
|
assert.True(t, allowed, "request %d should be allowed", i+1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 第 4 个请求应被拒绝
|
||||||
|
allowed, retryAfter := bucket.allow()
|
||||||
|
assert.False(t, allowed)
|
||||||
|
assert.Greater(t, retryAfter, time.Duration(0))
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTokenBucket_Allow_RetryAfterCorrect(t *testing.T) {
|
||||||
|
bucket := newTokenBucket(1, 1.0) // 每秒 1 个令牌
|
||||||
|
|
||||||
|
// 消耗唯一的令牌
|
||||||
|
allowed, _ := bucket.allow()
|
||||||
|
require.True(t, allowed)
|
||||||
|
|
||||||
|
// 立即再次请求应被拒绝
|
||||||
|
allowed, retryAfter := bucket.allow()
|
||||||
|
assert.False(t, allowed)
|
||||||
|
// retryAfter 应约为 1 秒(允许一定误差)
|
||||||
|
assert.InDelta(t, 1000, retryAfter.Milliseconds(), 100)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTokenBucket_Allow_RefillAfterWait(t *testing.T) {
|
||||||
|
bucket := newTokenBucket(2, 10.0) // 每秒 10 个令牌(每 100ms 一个)
|
||||||
|
|
||||||
|
// 消耗 2 个令牌
|
||||||
|
bucket.allow()
|
||||||
|
bucket.allow()
|
||||||
|
|
||||||
|
// 等待 150ms,应补充至少 1 个令牌
|
||||||
|
time.Sleep(150 * time.Millisecond)
|
||||||
|
|
||||||
|
allowed, _ := bucket.allow()
|
||||||
|
assert.True(t, allowed)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTokenBucket_Allow_CapacityLimit(t *testing.T) {
|
||||||
|
bucket := newTokenBucket(3, 1.0)
|
||||||
|
|
||||||
|
// 等待足够长时间让桶"溢出"
|
||||||
|
time.Sleep(100 * time.Millisecond)
|
||||||
|
|
||||||
|
// 但最多只能消耗 capacity 个令牌
|
||||||
|
for i := 0; i < 3; i++ {
|
||||||
|
allowed, _ := bucket.allow()
|
||||||
|
assert.True(t, allowed, "request %d should be allowed", i+1)
|
||||||
|
}
|
||||||
|
|
||||||
|
// 第 4 个应被拒绝
|
||||||
|
allowed, _ := bucket.allow()
|
||||||
|
assert.False(t, allowed)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTokenBucket_Allow_ConcurrentSafe(t *testing.T) {
|
||||||
|
bucket := newTokenBucket(100, 10.0)
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
successCount := 0
|
||||||
|
var mu sync.Mutex
|
||||||
|
|
||||||
|
// 100 个并发请求
|
||||||
|
for i := 0; i < 100; i++ {
|
||||||
|
wg.Add(1)
|
||||||
|
go func() {
|
||||||
|
defer wg.Done()
|
||||||
|
allowed, _ := bucket.allow()
|
||||||
|
if allowed {
|
||||||
|
mu.Lock()
|
||||||
|
successCount++
|
||||||
|
mu.Unlock()
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
}
|
||||||
|
|
||||||
|
wg.Wait()
|
||||||
|
|
||||||
|
// 应该正好 100 个成功(桶容量为 100)
|
||||||
|
assert.Equal(t, 100, successCount)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTokenBucket_Allow_ZeroCapacity(t *testing.T) {
|
||||||
|
bucket := newTokenBucket(0, 1.0)
|
||||||
|
|
||||||
|
allowed, retryAfter := bucket.allow()
|
||||||
|
assert.False(t, allowed)
|
||||||
|
assert.Greater(t, retryAfter, time.Duration(0))
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestTokenBucket_Allow_ZeroRate(t *testing.T) {
|
||||||
|
bucket := newTokenBucket(1, 0.0)
|
||||||
|
|
||||||
|
// 第一个通过
|
||||||
|
allowed, _ := bucket.allow()
|
||||||
|
assert.True(t, allowed)
|
||||||
|
|
||||||
|
// 第二个被拒绝,且 retryAfter 应为无限大(实际上会很大)
|
||||||
|
allowed, retryAfter := bucket.allow()
|
||||||
|
assert.False(t, allowed)
|
||||||
|
// rate=0 时,retryAfter 理论上无限大,实际会是一个很大的值
|
||||||
|
assert.Greater(t, retryAfter, 1*time.Hour)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMemoryLimiter_Allow_DifferentKeys(t *testing.T) {
|
||||||
|
cfg := config.RateLimitConfig{
|
||||||
|
Enabled: true,
|
||||||
|
Query: config.BucketConfig{Capacity: 2, Rate: 1.0},
|
||||||
|
}
|
||||||
|
limiter := NewMemoryLimiter(cfg)
|
||||||
|
defer limiter.Stop()
|
||||||
|
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
// user1 消耗 2 个令牌
|
||||||
|
allowed, _ := limiter.Allow(ctx, "user1:query")
|
||||||
|
assert.True(t, allowed)
|
||||||
|
allowed, _ = limiter.Allow(ctx, "user1:query")
|
||||||
|
assert.True(t, allowed)
|
||||||
|
|
||||||
|
// user1 第 3 个被拒绝
|
||||||
|
allowed, _ = limiter.Allow(ctx, "user1:query")
|
||||||
|
assert.False(t, allowed)
|
||||||
|
|
||||||
|
// user2 应该不受影响
|
||||||
|
allowed, _ = limiter.Allow(ctx, "user2:query")
|
||||||
|
assert.True(t, allowed)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMemoryLimiter_Cleanup(t *testing.T) {
|
||||||
|
cfg := config.RateLimitConfig{
|
||||||
|
Enabled: true,
|
||||||
|
Query: config.BucketConfig{Capacity: 1, Rate: 1.0},
|
||||||
|
}
|
||||||
|
limiter := NewMemoryLimiter(cfg)
|
||||||
|
defer limiter.Stop()
|
||||||
|
|
||||||
|
ctx := context.Background()
|
||||||
|
|
||||||
|
// 创建一个桶
|
||||||
|
limiter.Allow(ctx, "user1:query")
|
||||||
|
|
||||||
|
// 验证桶已创建
|
||||||
|
limiter.mu.RLock()
|
||||||
|
initialCount := len(limiter.buckets)
|
||||||
|
limiter.mu.RUnlock()
|
||||||
|
assert.Equal(t, 1, initialCount)
|
||||||
|
|
||||||
|
// 手动触发清理(模拟 10 分钟后)
|
||||||
|
limiter.mu.Lock()
|
||||||
|
for _, bucket := range limiter.buckets {
|
||||||
|
bucket.mu.Lock()
|
||||||
|
bucket.lastRefill = time.Now().Add(-11 * time.Minute)
|
||||||
|
bucket.mu.Unlock()
|
||||||
|
}
|
||||||
|
limiter.mu.Unlock()
|
||||||
|
|
||||||
|
limiter.removeInactiveBuckets()
|
||||||
|
|
||||||
|
// 验证桶已被清理
|
||||||
|
limiter.mu.RLock()
|
||||||
|
finalCount := len(limiter.buckets)
|
||||||
|
limiter.mu.RUnlock()
|
||||||
|
assert.Equal(t, 0, finalCount)
|
||||||
|
}
|
||||||
|
|
||||||
|
func TestMemoryLimiter_Stop(t *testing.T) {
|
||||||
|
cfg := config.RateLimitConfig{
|
||||||
|
Enabled: true,
|
||||||
|
Query: config.BucketConfig{Capacity: 1, Rate: 1.0},
|
||||||
|
}
|
||||||
|
limiter := NewMemoryLimiter(cfg)
|
||||||
|
|
||||||
|
// 多次调用 Stop 不应 panic
|
||||||
|
limiter.Stop()
|
||||||
|
limiter.Stop()
|
||||||
|
}
|
||||||
17
backend/internal/ratelimit/limiter.go
Normal file
17
backend/internal/ratelimit/limiter.go
Normal file
@@ -0,0 +1,17 @@
|
|||||||
|
package ratelimit
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"time"
|
||||||
|
)
|
||||||
|
|
||||||
|
// Limiter 速率限制器接口。
|
||||||
|
type Limiter interface {
|
||||||
|
// Allow 判断 key 是否允许执行一次操作。
|
||||||
|
// key 通常为 "userID:action" 格式。
|
||||||
|
// 返回 (allowed, retryAfter)。retryAfter 表示需要等待的时间。
|
||||||
|
Allow(ctx context.Context, key string) (bool, time.Duration)
|
||||||
|
|
||||||
|
// Stop 停止限流器,清理资源(如后台 goroutine)。
|
||||||
|
Stop()
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user