docs: 添加令牌桶限流模块设计文档
- 新建 docs/13-令牌桶限流设计.md,覆盖算法原理、内存/Redis 双实现、配置设计、接入点、测试用例等 - 更新 docs/01-架构设计.md Rate Limiter 模块行链接至新文档 - 更新 docs/README.md 文档索引和推荐阅读顺序
This commit is contained in:
@@ -210,7 +210,7 @@ graph LR
|
||||
| Models | 数据模型定义 |
|
||||
| Migrations | 数据库版本化迁移,嵌入式 SQL 文件自动执行 |
|
||||
| Model Router | 根据请求类型选择 AI 模型(待实现) |
|
||||
| Rate Limiter | 令牌桶限流(待实现) |
|
||||
| Rate Limiter | 令牌桶限流。详细设计见 [令牌桶限流设计](./13-令牌桶限流设计.md) |
|
||||
|
||||
## 前端组件
|
||||
|
||||
|
||||
499
docs/13-令牌桶限流设计.md
Normal file
499
docs/13-令牌桶限流设计.md
Normal file
@@ -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<br/>query 消息"]
|
||||
REST["REST API<br/>login / register"]
|
||||
end
|
||||
|
||||
subgraph LimiterModule["Rate Limiter 模块"]
|
||||
Interface["Limiter 接口<br/>Allow(userID) → (bool, retryAfter)"]
|
||||
MemBucket["TokenBucket<br/>内存令牌桶"]
|
||||
RedisBucket["RedisTokenBucket<br/>Redis 令牌桶(Lua 脚本)"]
|
||||
Middleware["RateLimitMiddleware<br/>Gin 中间件"]
|
||||
end
|
||||
|
||||
subgraph Storage["存储层"]
|
||||
MemSync["sync.RWMutex<br/>进程内 map"]
|
||||
Redis["Redis<br/>分布式计数"]
|
||||
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/)
|
||||
@@ -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-令牌桶限流设计** — 速率限制设计(令牌桶算法、成本控制、防暴力破解)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user