@@ -0,0 +1,431 @@
# 实施计划：限流 + 协程池 + 消息队列
> 基于简历描述，本项目需要落地两个核心能力：
> 1. 基于 go-redis 的分布式令牌桶限流
> 2. 自定义 IO 密集型协程池 + RabbitMQ 消息队列
## 现状分析
| 维度 | 当前状态 | 目标状态 |
|------|---------|---------|
| 限流 | 无任何限流中间件；API 错误码 429 已预留但未实现 | 令牌桶算法，支持按用户/全局维度限流 |
| 并发控制 | 每个请求裸 `go` 起协程，无上限 | 有界协程池，可控并发度 |
| 任务队列 | 同步阻塞，任务状态通过 DB 轮询 | RabbitMQ 异步队列，生产者-消费者模式 |
| 基础设施 | SQLite + 七牛云，无 Redis / MQ | 新增 Redis + RabbitMQ 依赖 |
---
## 阶段一：基础设施接入
### 1.1 Redis 接入
**新增依赖： ** `github.com/redis/go-redis/v9`
**新增配置： ** `internal/config/redis.go`
``` go
type RedisConfig struct {
Addr string ` yaml:"addr" ` // 默认 "localhost:6379"
Password string ` yaml:"password" `
DB int ` yaml:"db" ` // 默认 0
}
```
**新增初始化： ** `internal/pkg/redis/client.go` — 全局 `*redis.Client` 单例，`InitRedis(cfg)` + `CloseRedis()` ，在 `cmd/main.go` 启动时调用。
**配置文件： ** `config.yaml` 新增 `redis` 节。
### 1.2 RabbitMQ 接入
**新增依赖： ** `github.com/rabbitmq/amqp091-go`
**新增配置： ** `internal/config/rabbitmq.go`
``` go
type RabbitMQConfig struct {
URL string ` yaml:"url" ` // 默认 "amqp://guest:guest@localhost:5672/"
Queue string ` yaml:"queue" ` // 默认 "gen2d:tasks"
}
```
**新增初始化： ** `internal/pkg/mq/rabbitmq.go` — 封装连接、Channel、Queue 声明，提供 `Publish(body)` 和 `Consume(handler)` 方法。
**Docker Compose： ** 新增 `docker-compose.yml` （或更新已有文件），包含 Redis 和 RabbitMQ 服务。
---
## 阶段二：分布式令牌桶限流
### 2.1 算法选型分析
| 方案 | 优点 | 缺点 | 适用场景 |
|------|------|------|---------|
| 固定窗口计数器 | 实现简单 | 窗口边界突发 | 对精度要求不高的场景 |
| 滑动窗口 | 解决边界问题 | 内存开销较大 | 中等精度需求 |
| 漏桶 | 平滑输出 | 无法应对突发 | 流量整形 |
| **令牌桶 ** | 平滑 + 允许突发 | 实现稍复杂 | **API 限流（本项目选用） ** |
**选用令牌桶理由： **
- 用户提交生成任务时存在突发行为（连续点几次生成），令牌桶允许一定的突发消费
- go-redis 社区有成熟的 Lua 脚本实现，原子性有保证
- 天然支持分布式，未来多实例部署无缝扩展
### 2.2 实现方案
**新增包： ** `internal/pkg/ratelimit/`
```
internal/pkg/ratelimit/
├── limiter.go // 接口定义 + 令牌桶实现
├── middleware.go // Gin 限流中间件
└── config.go // 限流配置
```
**核心接口： **
``` go
// limiter.go
type Limiter interface {
// Allow 检查 key 是否允许通过，返回 (是否允许, 剩余令牌数, 重试等待时间)
Allow ( ctx context . Context , key string ) ( bool , int , time . Duration )
}
```
**令牌桶实现要点： **
- 底层使用 Redis + Lua 脚本保证原子性（单次 Redis 请求完成令牌检查 + 扣减）
- Lua 脚本逻辑：计算上次调用到当前的时间差 → 按速率补充令牌 → 判断桶内令牌是否足够 → 扣减并返回
- Key 设计：`ratelimit:{scope}:{identifier}` ，例如 `ratelimit:user:123` 或 `ratelimit:global:generate`
- 支持两级限流：用户级（per-user） ）
**配置： **
``` go
// config.go
type Config struct {
Rate int // 每秒产生的令牌数
Burst int // 桶容量（允许的突发量）
KeyPrefix string // Redis key 前缀
Expiration time . Duration // key 过期时间，防止冷用户占用内存
}
```
**推荐参数（生成接口）： **
- 用户级：Rate=1/s, Burst=3（允许用户连续提交 3 个任务，之后每秒最多 1 个）
- 全局级：Rate=10/s, Burst=20（系统总并发上限）
### 2.3 Gin 中间件集成
**middleware.go： **
``` go
func RateLimit ( limiter Limiter , keyFunc func ( c * gin . Context ) string ) gin . HandlerFunc
```
- `keyFunc` 负责从请求中提取限流 key（如 JWT 中的 userID，或 endpoint 标识）
- 被限流时返回 HTTP 429 + 统一错误格式 `{"code": 429, "message": "请求过于频繁，请稍后再试", "data": nil}`
- 响应头添加 `X-RateLimit-Remaining` 和 `Retry-After`
**路由注册（修改 `internal/router/router.go`）： **
``` go
// 用户级限流 — 基于 JWT userID
authGroup . Use ( ratelimit . RateLimit ( userLimiter , func ( c * gin . Context ) string {
return c . GetString ( "userID" )
} ) )
// 生成接口额外加全局限流
v1 . POST ( "/generate" , ratelimit . RateLimit ( globalLimiter , func ( c * gin . Context ) string {
return "generate"
} ) , generateHandler . Create )
```
---
## 阶段三：自定义协程池
### 3.1 设计思路
当前问题：每个 `POST /api/v1/generate` 裸起 `go RunPipeline()` ，在高并发下可能耗尽系统资源（goroutine 数、文件描述符、DB 连接）。
**方案选型： **
| 方案 | 优点 | 缺点 | 适用场景 |
|------|------|------|---------|
| `ants` 库 | 开箱即用，成熟 | 黑盒，简历无亮点 | 快速实现 |
| **手写协程池 ** | 可控性强，体现设计能力 | 需要维护 | **本项目选用 ** |
| Semaphore | 极简 | 无排队机制 | 临时方案 |
**选用自定义协程池理由： **
- 简历明确写"自定义 IO 密集型协程池"，需要有手写实现
- 可以针对 IO 密集型特征做专门调参（worker 数可设为 CPU 核数的数倍）
- 可以与 RabbitMQ 消费者共用池，统一管理
### 3.2 实现方案
**新增包： ** `internal/pkg/workerpool/`
```
internal/pkg/workerpool/
├── pool.go // 协程池核心实现
├── option.go // 配置选项（函数选项模式）
└── task.go // 任务定义
```
**核心结构： **
``` go
// pool.go
type Pool struct {
workers int // worker 数量
taskQueue chan Task // 有界任务队列
wg sync . WaitGroup
ctx context . Context
cancel context . CancelFunc
metrics * Metrics // 运行指标
}
type Task struct {
ID string
Fn func ( ctx context . Context ) error
Priority int // 预留优先级支持
}
type Metrics struct {
ActiveWorkers atomic . Int32
QueuedTasks atomic . Int32
CompletedTasks atomic . Int64
FailedTasks atomic . Int64
}
```
**关键设计点： **
1. **IO 密集型调参： ** 默认 worker 数 = `runtime.NumCPU() * 4` （而非 CPU 密集型的 1:1），通过 `WithWorkers(n)` 可覆盖
2. **有界队列： ** `taskQueue` channel 容量有上限（默认 100），满时提交阻塞或返回 `ErrPoolFull` ，防止任务无限堆积
3. **优雅关闭： ** `Shutdown(ctx)` 等待所有正在执行的任务完成或 ctx 超时
4. **指标暴露： ** 活跃 worker 数、排队任务数、完成/失败任务数，通过 `GET /api/v1/metrics` 或 Prometheus 格式暴露
**函数选项模式配置： **
``` go
// option.go
type Option func ( * Pool )
func WithWorkers ( n int ) Option // 自定义 worker 数
func WithQueueSize ( n int ) Option // 自定义队列容量
func WithMetrics ( reg prometheus . Registerer ) Option // 注册 Prometheus 指标
```
### 3.3 接入服务层
**修改 `internal/service/generate.go`： **
``` go
// Before:
go svc . runPipeline ( ctx , taskID , input )
// After:
err := svc . pool . Submit ( workerpool . Task {
ID : taskID ,
Fn : func ( ctx context . Context ) error {
return svc . runPipeline ( ctx , taskID , input )
} ,
} )
if err != nil {
// 更新任务状态为 failed，记录原因
return err
}
```
---
## 阶段四：RabbitMQ 消息队列
### 4.1 架构设计
```
POST /api/v1/generate
│
▼
┌─────────────┐ ┌──────────────┐ ┌─────────────────┐
│ Producer │────▶│ RabbitMQ │────▶│ Consumer │
│ (API Handler)│ │ gen2d:tasks │ │ (Worker Goroutine)│
└─────────────┘ └──────────────┘ └────────┬────────┘
│
▼
┌─────────────────┐
│ Worker Pool │
│ (协程池执行 Pipeline)│
└─────────────────┘
```
**同步路径（现有，保留）： ** 简单请求、Prompt 优化等轻量操作仍走 HTTP 同步返回。
**异步路径（新增）： ** `POST /api/v1/generate` → 生产者发送消息到 RabbitMQ → 消费者从队列取出 → 提交到协程池执行 Pipeline。
### 4.2 消息格式
``` go
// internal/pkg/mq/message.go
type TaskMessage struct {
TaskID string ` json:"task_id" `
UserID string ` json:"user_id" `
Input service . PipelineInput ` json:"input" `
CreatedAt time . Time ` json:"created_at" `
RetryCount int ` json:"retry_count" `
}
```
### 4.3 生产者（API Handler 层）
**修改 `internal/handler/generate.go`： **
``` go
func ( h * GenerateHandler ) Create ( c * gin . Context ) {
// 1. 参数绑定、校验
// 2. 去重检查（现有逻辑保留）
// 3. 创建 DB 任务记录（status=pending）
// 4. 发布消息到 RabbitMQ（替代原来的 go runPipeline）
err := h . mq . Publish ( ctx , mq . TaskMessage { ... } )
// 5. 返回 taskId
}
```
### 4.4 消费者（独立启动）
**新增： ** `internal/worker/consumer.go`
``` go
type Consumer struct {
mq * mq . RabbitMQ
pool * workerpool . Pool
service * service . GenerateService
logger * slog . Logger
}
func ( c * Consumer ) Start ( ctx context . Context ) error {
// 持续消费 RabbitMQ 消息
// 每条消息提交到协程池执行
return c . mq . Consume ( ctx , func ( msg mq . TaskMessage ) error {
return c . pool . Submit ( workerpool . Task {
ID : msg . TaskID ,
Fn : func ( ctx context . Context ) error {
return c . service . RunPipeline ( ctx , msg . TaskID , msg . Input )
} ,
} )
} )
}
```
**启动方式： ** `cmd/main.go` 中，API 服务和消费者在同一个进程内启动（通过配置开关控制是否启用消费者），便于开发调试；生产环境可拆分为独立进程。
### 4.5 消息可靠性
| 环节 | 保障措施 |
|------|---------|
| 生产者发送失败 | 降级为同步执行（fallback），记录告警日志 |
| 消费者处理失败 | RabbitMQ NACK + 重新入队（max retry = 3）
| 消费者崩溃 | 手动 ACK 模式，未 ACK 的消息自动重新投递 |
| 消息堆积 | 监控队列深度，超阈值触发告警 |
---
## 阶段五：可观测性与配置
### 5.1 运行指标
新增 `GET /api/v1/metrics` 接口（或接入 Prometheus），暴露：
- **限流指标：** 各 key 被限流次数、当前令牌数
- **协程池指标：** 活跃 worker 数、队列深度、任务完成/失败数、平均执行耗时
- **MQ 指标：** 队列深度、生产/消费速率、NACK 次数
### 5.2 配置汇总
新增配置项（`config.yaml` ）：
``` yaml
redis :
addr : "localhost:6379"
password : ""
db : 0
rabbitmq :
url : "amqp://guest:guest@localhost:5672/"
queue : "gen2d:tasks"
ratelimit :
user :
rate : 1
burst : 3
global :
rate : 10
burst : 20
workerpool :
workers : 16 # NumCPU * 4
queue_size : 100
consumer_enabled : true
```
---
## 文件变更清单
### 新增文件
| 文件路径 | 用途 |
|---------|------|
| `internal/config/redis.go` | Redis 配置结构体 |
| `internal/config/rabbitmq.go` | RabbitMQ 配置结构体 |
| `internal/pkg/redis/client.go` | Redis 客户端初始化 |
| `internal/pkg/mq/rabbitmq.go` | RabbitMQ 连接封装 |
| `internal/pkg/mq/message.go` | 消息体定义 |
| `internal/pkg/ratelimit/limiter.go` | 令牌桶限流器 |
| `internal/pkg/ratelimit/middleware.go` | Gin 限流中间件 |
| `internal/pkg/ratelimit/config.go` | 限流配置 |
| `internal/pkg/workerpool/pool.go` | 协程池核心 |
| `internal/pkg/workerpool/option.go` | 协程池配置选项 |
| `internal/pkg/workerpool/task.go` | 任务定义 |
| `internal/worker/consumer.go` | MQ 消费者 |
| `docker-compose.yml` | Redis + RabbitMQ 服务编排 |
### 修改文件
| 文件路径 | 变更内容 |
|---------|---------|
| `cmd/main.go` | 初始化 Redis、RabbitMQ、WorkerPool；启动 Consumer |
| `internal/config/config.go` | 顶层 Config 新增 Redis / RabbitMQ / Ratelimit / WorkerPool 字段 |
| `internal/router/router.go` | 注册限流中间件 |
| `internal/handler/generate.go` | 将 `go runPipeline` 改为发布 MQ 消息 |
| `internal/service/generate.go` | 接受协程池提交，`RunPipeline` 签名不变 |
| `config.yaml` | 新增 redis / rabbitmq / ratelimit / workerpool 配置节 |
| `go.mod` | 新增 go-redis、amqp091-go 依赖 |
---
## 实施顺序建议
```
Phase 1: 基础设施（Redis + RabbitMQ 接入、配置、初始化）
↓
Phase 2: 令牌桶限流（独立模块，可先接入不影响其他功能）
↓
Phase 3: 协程池（独立模块，替换裸 go 调用）
↓
Phase 4: MQ 异步化（依赖协程池，改造生成流程）
↓
Phase 5: 可观测性（指标暴露、告警配置）
```
Phase 2 和 Phase 3 可并行开发，Phase 4 依赖 Phase 3 完成。
---
## 注意事项
1. **向后兼容： ** 各阶段独立可部署，不依赖全部完成才能上线。限流可以先上线，协程池和 MQ 可以后续接入。
2. **开发环境： ** `docker-compose.yml` 提供 Redis + RabbitMQ，前端开发者无需关心。
3. **降级策略： ** Redis 不可用时限流中间件自动放行（fail-open）；RabbitMQ 不可用时降级为同步执行。
4. **测试覆盖： ** 每个模块需编写单元测试 + 集成测试。限流器可用 `miniredis` 做内存 Redis 测试；协程池测试并发正确性；MQ 测试用 `rabbitmq-testcontainers` 或本地 Docker。
wonder