diff --git a/docs/01-架构设计.md b/docs/01-架构设计.md index 9f15d3b..8f6740f 100644 --- a/docs/01-架构设计.md +++ b/docs/01-架构设计.md @@ -210,7 +210,7 @@ graph LR | Models | 数据模型定义 | | Migrations | 数据库版本化迁移,嵌入式 SQL 文件自动执行 | | Model Router | 根据请求类型选择 AI 模型(待实现) | -| Rate Limiter | 令牌桶限流(待实现) | +| Rate Limiter | 令牌桶限流。详细设计见 [令牌桶限流设计](./13-令牌桶限流设计.md) | ## 前端组件 diff --git a/docs/13-令牌桶限流设计.md b/docs/13-令牌桶限流设计.md new file mode 100644 index 0000000..94dca6e --- /dev/null +++ b/docs/13-令牌桶限流设计.md @@ -0,0 +1,499 @@ +# 令牌桶限流设计 + +## 概述 + +CamTalk 采用令牌桶(Token Bucket)算法实现按用户维度的速率限制,核心目标是**控制 AI 调用成本**,同时为 REST API 提供防暴力破解保护。 + +**设计原则**: +- **成本优先**:主要限流对象是 WebSocket `query` 消息(每次触发 STT + LLM + TTS 完整调用链) +- **用户隔离**:Per-user 维度限流,单用户超限不影响其他用户 +- **弹性突发**:令牌桶允许合理的突发请求,优于固定窗口的滑动限流 +- **存储适配**:内存 + Redis 双实现,单实例零依赖,多实例分布式一致 + +## 整体架构 + +```mermaid +graph TB + subgraph Entry["入口层"] + WS["WebSocket Handler
query 消息"] + REST["REST API
login / register"] + end + + subgraph LimiterModule["Rate Limiter 模块"] + Interface["Limiter 接口
Allow(userID) → (bool, retryAfter)"] + MemBucket["TokenBucket
内存令牌桶"] + RedisBucket["RedisTokenBucket
Redis 令牌桶(Lua 脚本)"] + Middleware["RateLimitMiddleware
Gin 中间件"] + end + + subgraph Storage["存储层"] + MemSync["sync.RWMutex
进程内 map"] + Redis["Redis
分布式计数"] + end + + WS -->|"限流检查"| Interface + REST -->|"中间件"| Middleware + Middleware --> Interface + Interface --> MemBucket + Interface --> RedisBucket + MemBucket --> MemSync + RedisBucket --> Redis +``` + +## 令牌桶算法 + +### 原理 + +令牌桶以固定速率向桶中添加令牌,桶有最大容量上限。每次请求消耗一个令牌,桶空时拒绝请求。 + +``` +桶容量(capacity) = 允许的突发请求数上限 +填充速率(rate) = 每秒补充的令牌数 + +时间线示例(capacity=5, rate=0.2): + t=0s 桶满 5 令牌 → 用户连续发 5 个 query 全部通过 + t=0s 桶空 → 第 6 个 query 被拒绝,retryAfter=5s + t=5s 桶补充 1 令牌 → 可再发 1 个 query + t=10s 桶补充 1 令牌 → 可再发 1 个 query +``` + +### 算法公式 + +``` +elapsed = now - lastRefill +newTokens = elapsed * rate +currentTokens = min(capacity, lastTokens + newTokens) + +if currentTokens >= 1: + currentTokens -= 1 + allowed = true +else: + allowed = false + retryAfter = (1 - currentTokens) / rate +``` + +## 核心组件 + +### 1. Limiter 接口 + +**文件位置**:`backend/internal/ratelimit/limiter.go` + +```go +// Limiter 速率限制器接口。 +type Limiter interface { + // Allow 判断 key 是否允许执行一次操作。 + // key 通常为 "userID:action" 格式。 + // 返回 (allowed, retryAfter)。retryAfter 表示需要等待的时间。 + Allow(ctx context.Context, key string) (bool, time.Duration) +} +``` + +**设计要点**: +- key 为字符串,不限定格式,由调用方决定维度(用户 ID、IP 地址等) +- 返回 `retryAfter` 供客户端/服务端设置 `Retry-After` header +- 接受 `context.Context` 支持超时和取消(Redis 实现需要) + +### 2. 内存令牌桶(TokenBucket) + +**文件位置**:`backend/internal/ratelimit/bucket.go` + +```go +// TokenBucket 内存令牌桶,适用于单实例部署。 +type TokenBucket struct { + capacity int // 桶容量 + rate float64 // 每秒填充令牌数 + tokens float64 // 当前令牌数 + lastRefill time.Time // 上次填充时间 + mu sync.Mutex +} + +// Limiter 管理多个用户的令牌桶。 +type Limiter struct { + buckets map[string]*TokenBucket + config Config + mu sync.RWMutex + stopOnce sync.Once + done chan struct{} +} +``` + +**并发安全**: +- 每个桶内部用 `sync.Mutex` 保护读写 +- 桶 map 用 `sync.RWMutex` 保护(读多写少场景) +- 用户首次请求时惰性创建桶 + +**内存回收**: +- 后台 goroutine 定期扫描,清理超过 10 分钟无活动的桶 +- 避免长期运行后内存泄漏 + +### 3. Redis 令牌桶(RedisTokenBucket) + +**文件位置**:`backend/internal/ratelimit/redis_bucket.go` + +使用 Redis Lua 脚本保证原子性,避免竞态条件: + +```lua +-- KEYS[1] = 限流 key +-- ARGV[1] = capacity(桶容量) +-- ARGV[2] = rate(每秒填充数) +-- ARGV[3] = now(当前时间戳,秒,浮点) +-- ARGV[4] = ttl(key 过期时间,秒) + +local key = KEYS[1] +local capacity = tonumber(ARGV[1]) +local rate = tonumber(ARGV[2]) +local now = tonumber(ARGV[3]) +local ttl = tonumber(ARGV[4]) + +local data = redis.call('HMGET', key, 'tokens', 'last_refill') +local tokens = tonumber(data[1]) or capacity +local last_refill = tonumber(data[2]) or now + +-- 计算新令牌 +local elapsed = math.max(0, now - last_refill) +tokens = math.min(capacity, tokens + elapsed * rate) + +local allowed = 0 +local retry_after = 0 + +if tokens >= 1 then + tokens = tokens - 1 + allowed = 1 +else + retry_after = (1 - tokens) / rate +end + +-- 回写状态 +redis.call('HMSET', key, 'tokens', tokens, 'last_refill', now) +redis.call('EXPIRE', key, ttl) + +return {allowed, tostring(retry_after)} +``` + +**设计要点**: +- 每个用户的限流状态存储为一个 Redis Hash(`tokens` + `last_refill`) +- TTL 自动过期,无需手动清理 +- Lua 脚本保证"读取-计算-回写"原子执行 + +### 4. Gin 中间件 + +**文件位置**:`backend/internal/ratelimit/middleware.go` + +```go +// RateLimitMiddleware 返回 Gin 中间件,按 key 维度限流。 +// keyFunc 从请求中提取限流 key(如 IP、用户 ID)。 +func RateLimitMiddleware(limiter Limiter, keyFunc func(*gin.Context) string) gin.HandlerFunc +``` + +**使用方式**: + +```go +// 按 IP 限流(登录/注册,未登录用户无 userID) +loginGroup.POST("/login", + ratelimit.Middleware(limiter, func(c *gin.Context) string { + return c.ClientIP() + ":login" + }), + authHandler.Login, +) + +// 按用户 ID 限流(已认证的 API) +authorized.POST("/conversations", + ratelimit.Middleware(limiter, func(c *gin.Context) string { + return c.GetString("user_id") + ":conversation" + }), + convHandler.Create, +) +``` + +**错误响应**: + +REST API 返回 HTTP 429: + +```json +{ + "code": "RATE_LIMITED", + "message": "too many requests, retry after 5s" +} +``` + +同时设置 `Retry-After` header: + +``` +HTTP/1.1 429 Too Many Requests +Retry-After: 5 +``` + +## 限流接入点 + +### WebSocket query 消息(核心) + +在 `ws/handler.go` 的 `case "query"` 分支中,orchestrator 调用前检查: + +```go +case "query": + // ... 解析消息 ... + + // 限流检查 + if limiter != nil { + allowed, retryAfter := limiter.Allow(ctx, userID+":query") + if !allowed { + errors.SendWSError(client, errors.CodeRateLimited, msg.RequestID, + fmt.Errorf("rate limited, retry after %s", retryAfter)) + continue + } + } + + // ... 继续处理 query ... +``` + +### REST API 登录/注册 + +在 `api/auth.go` 的路由注册中添加中间件: + +```go +func (h *AuthHandler) RegisterRoutes(rg *gin.RouterGroup, limiter ratelimit.Limiter) { + auth := rg.Group("/auth") + if limiter != nil { + auth.POST("/register", + ratelimit.Middleware(limiter, ipKeyFunc("register")), + h.Register, + ) + auth.POST("/login", + ratelimit.Middleware(limiter, ipKeyFunc("login")), + h.Login, + ) + } else { + auth.POST("/register", h.Register) + auth.POST("/login", h.Login) + } + auth.POST("/refresh", h.Refresh) + auth.POST("/logout", h.Logout) +} +``` + +### 不限流的端点 + +| 端点 | 原因 | +|------|------| +| `ping` / `pong` | 心跳保活,无 AI 调用成本 | +| `config` | 配置更新,无 AI 调用成本 | +| `interrupt` | 中断请求,取消操作不应被限流 | +| `GET /api/health` | 健康检查,运维必需 | +| `POST /api/auth/refresh` | Token 刷新,限流会导致用户被迫重新登录 | +| `POST /api/auth/logout` | 登出,限流会导致用户无法正常退出 | +| `GET /api/conversations` | 查询列表,无 AI 调用成本 | + +## 配置设计 + +### 配置文件 + +```yaml +# config.yaml 新增 +ratelimit: + enabled: true + # WebSocket query 消息限流(核心,控制 AI 成本) + query: + capacity: 10 # 突发容量:允许连续发 10 个 query + rate: 0.2 # 填充速率:每 5 秒补充 1 个令牌 + # REST API 登录限流(防暴力破解) + login: + capacity: 5 # 突发容量:允许连续 5 次登录尝试 + rate: 0.1 # 填充速率:每 10 秒补充 1 次 + # REST API 注册限流 + register: + capacity: 3 # 突发容量:允许连续 3 次注册 + rate: 0.05 # 填充速率:每 20 秒补充 1 次 +``` + +### 配置结构体 + +```go +// config/config.go 新增 + +type RateLimitConfig struct { + Enabled bool `mapstructure:"enabled"` + Query BucketConfig `mapstructure:"query"` + Login BucketConfig `mapstructure:"login"` + Register BucketConfig `mapstructure:"register"` +} + +type BucketConfig struct { + Capacity int `mapstructure:"capacity"` // 桶容量(突发上限) + Rate float64 `mapstructure:"rate"` // 每秒填充令牌数 +} +``` + +### 默认值 + +```go +// setDefaults 新增 +v.SetDefault("ratelimit.enabled", false) +v.SetDefault("ratelimit.query.capacity", 10) +v.SetDefault("ratelimit.query.rate", 0.2) +v.SetDefault("ratelimit.login.capacity", 5) +v.SetDefault("ratelimit.login.rate", 0.1) +v.SetDefault("ratelimit.register.capacity", 3) +v.SetDefault("ratelimit.register.rate", 0.05) +``` + +### 参数选择建议 + +| 场景 | capacity | rate | 含义 | +|------|----------|------|------| +| WebSocket query | 10 | 0.2 | 突发 10 个,之后每 5 秒 1 个 | +| 登录 | 5 | 0.1 | 突发 5 次,之后每 10 秒 1 次 | +| 注册 | 3 | 0.05 | 突发 3 次,之后每 20 秒 1 次 | + +> **调参原则**:capacity 决定"能忍多少次突发",rate 决定"稳态下多久能再请求一次"。query 的 rate 建议根据 AI 调用成本和目标月预算反推。 + +## 依赖注入 + +### main.go 初始化 + +```go +// 初始化限流器 +var limiter ratelimit.Limiter +if cfg.RateLimit.Enabled { + if rdb != nil { + // 多实例:使用 Redis 令牌桶 + limiter = ratelimit.NewRedisLimiter(rdb, cfg.RateLimit) + logger.Log.Info("rate limiter initialized with Redis backend") + } else { + // 单实例:使用内存令牌桶 + limiter = ratelimit.NewLimiter(cfg.RateLimit) + logger.Log.Info("rate limiter initialized with in-memory backend") + } + defer limiter.Stop() +} +``` + +### 注入到各模块 + +```go +// WebSocket Handler —— 新增 limiter 参数 +r.GET("/ws", ws.ServeWS(sessionMgr, orch, cfg, tokenMgr, limiter)) + +// Auth REST —— 新增 limiter 参数 +authHandler := api.NewAuthHandler(authService, tokenMgr) +authHandler.RegisterRoutes(apiGroup, limiter) +``` + +## 错误码 + +复用已有错误码 `RATE_LIMITED`(`backend/internal/errors/codes.go`): + +| 传输层 | HTTP 状态码 | 错误格式 | +|--------|-----------|---------| +| REST API | 429 Too Many Requests | `{code: "RATE_LIMITED", message: "too many requests, retry after Xs"}` | +| WebSocket | — | `{type: "error", code: "RATE_LIMITED", request_id: "...", message: "..."}` | + +## 文件结构 + +``` +backend/internal/ratelimit/ +├── limiter.go # Limiter 接口 + Config 类型定义 +├── bucket.go # 内存令牌桶实现 +├── bucket_test.go # 内存令牌桶单元测试 +├── redis_bucket.go # Redis 令牌桶实现(Lua 脚本) +├── redis_bucket_test.go# Redis 令牌桶单元测试 +└── middleware.go # Gin 中间件 +``` + +## 测试用例 + +### 单元测试 + +**内存令牌桶**(`bucket_test.go`): +- 首次请求通过 +- 连续消耗至桶空 +- 桶空后拒绝,返回正确 retryAfter +- 等待后令牌补充,请求通过 +- 并发安全性(多个 goroutine 同时 Allow) +- 桶容量边界(capacity=0, capacity=1) +- 填充速率边界(rate=0, rate 极大值) +- 不活跃桶的内存回收 + +**Redis 令牌桶**(`redis_bucket_test.go`): +- 与内存实现行为一致性 +- Lua 脚本原子性 +- key TTL 自动过期 +- 并发安全性(多个客户端同时请求) + +### 集成测试 + +- 限流关闭时不拦截请求 +- 限流开启后,REST API 登录超限返回 429 +- 限流开启后,WebSocket query 超限返回 `RATE_LIMITED` 错误 +- 单实例内存限流 vs 多实例 Redis 限流行为一致 +- 重启后内存限流重置,Redis 限流保持 + +## 扩展点 + +### 1. 多级限流 + +可扩展为多级限流策略: + +``` +全局限流(全用户共享) → 用户级限流(当前实现) → 端点级限流(不同 API 不同限制) +``` + +### 2. 动态调参 + +通过配置热更新或管理 API 动态调整限流参数,无需重启: + +```go +// 预留接口 +type DynamicLimiter interface { + Limiter + UpdateConfig(action string, cfg BucketConfig) error +} +``` + +### 3. 按用户等级差异化 + +不同用户等级使用不同的限流参数: + +```yaml +ratelimit: + query: + capacity: 10 # 免费用户 + rate: 0.2 + query_premium: + capacity: 30 # 付费用户 + rate: 1.0 +``` + +### 4. 滑动窗口限流 + +令牌桶适合允许突发的场景。如果需要更平滑的限流,可增加滑动窗口实现: + +```go +type SlidingWindowLimiter struct { + windowSize time.Duration + maxRequests int +} +``` + +### 5. 分布式全局限流 + +当前 Redis 实现是 Per-Instance 独立计数。如需全局精确限流,可改为 Redis 全局计数器(所有实例共享同一个 key)。 + +## 监控指标 + +### 关键指标 + +- **限流触发率**:被拒绝请求数 / 总请求数 +- **各端点限流分布**:query / login / register 各自的触发率 +- **等待时长分布**:retryAfter 的 P50/P99 +- **桶状态**:各用户桶的平均令牌数(反映使用模式) + +### 告警规则 + +- **限流触发率突增**:可能表示异常流量或攻击 +- **单用户持续被限流**:可能表示客户端 bug(死循环请求) + +## 参考资料 + +- [Token Bucket 算法](https://en.wikipedia.org/wiki/Token_bucket) +- [Redis Rate Limiting](https://redis.io/glossaries/rate-limiting/) +- [Cloudflare - How we built rate limiting capable of scaling to millions of domains](https://blog.cloudflare.com/counting-things-a-lot-of-different-things/) diff --git a/docs/README.md b/docs/README.md index 5a58904..00b707d 100644 --- a/docs/README.md +++ b/docs/README.md @@ -18,6 +18,7 @@ CamTalk 是一款多模态实时 AI 视觉对话助手。用户通过摄像头 | [10-Eino重构方案](10-Eino重构方案.md) | Eino Graph 替换手写 goroutine 管道的设计方案 | | [11-Eino框架技术文档](11-Eino框架技术文档.md) | Eino 框架在 CamTalk 中的使用指南(Graph、Lambda、Callback、State) | | [12-鉴权体系设计](12-鉴权体系设计.md) | JWT 双 token 轮转认证、bcrypt 密码哈希、Refresh Token Rotation、安全机制 | +| [13-令牌桶限流设计](13-令牌桶限流设计.md) | 令牌桶限流算法、内存/Redis 双实现、Gin 中间件、WebSocket query 限流、配置设计 | ## 推荐阅读顺序 @@ -30,4 +31,5 @@ CamTalk 是一款多模态实时 AI 视觉对话助手。用户通过摄像头 6. **09-技术名词解释** — 遇到不熟悉的名词时查阅 7. **10~12** — Eino 重构相关(方案、框架文档、实施记录) 8. **12-鉴权体系设计** — 认证授权机制详细设计(JWT、bcrypt、Refresh Token Rotation) +9. **13-令牌桶限流设计** — 速率限制设计(令牌桶算法、成本控制、防暴力破解)