From 3b6226394b495f861b11676bc965571f9a1e1d08 Mon Sep 17 00:00:00 2001 From: hhs <386998068@qq.com> Date: Sat, 20 Jun 2026 23:53:36 +0800 Subject: [PATCH] =?UTF-8?q?feat:=20=E5=AE=9E=E7=8E=B0=E5=86=85=E5=AD=98?= =?UTF-8?q?=E4=BB=A4=E7=89=8C=E6=A1=B6=E9=99=90=E6=B5=81=E5=99=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Limiter 接口定义（Allow + Stop 方法） - TokenBucket 实现（容量、填充速率、并发安全） - MemoryLimiter 管理多用户令牌桶 - 后台 goroutine 定期清理不活跃桶（10 分钟） - 完整单元测试覆盖（11 个测试用例，全部通过） - 边界情况处理（rate=0、capacity=0、并发安全） --- backend/internal/ratelimit/bucket.go | 172 ++++++++++++++++++ backend/internal/ratelimit/bucket_test.go | 203 ++++++++++++++++++++++ backend/internal/ratelimit/limiter.go | 17 ++ 3 files changed, 392 insertions(+) create mode 100644 backend/internal/ratelimit/bucket.go create mode 100644 backend/internal/ratelimit/bucket_test.go create mode 100644 backend/internal/ratelimit/limiter.go diff --git a/backend/internal/ratelimit/bucket.go b/backend/internal/ratelimit/bucket.go new file mode 100644 index 0000000..308336f --- /dev/null +++ b/backend/internal/ratelimit/bucket.go @@ -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) diff --git a/backend/internal/ratelimit/bucket_test.go b/backend/internal/ratelimit/bucket_test.go new file mode 100644 index 0000000..ebb19e9 --- /dev/null +++ b/backend/internal/ratelimit/bucket_test.go @@ -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() +} diff --git a/backend/internal/ratelimit/limiter.go b/backend/internal/ratelimit/limiter.go new file mode 100644 index 0000000..5fa4e3c --- /dev/null +++ b/backend/internal/ratelimit/limiter.go @@ -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() +}