gen2d/docs/PLAN.md

实施计划：限流 + 协程池 + 消息队列

基于简历描述，本项目需要落地两个核心能力：

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

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

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        // 限流配置

核心接口：

// 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）+ 全局级（per-endpoint）

配置：

// 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：

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）：

// 用户级限流 — 基于 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       // 任务定义

核心结构：

// 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 格式暴露

函数选项模式配置：

// 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：

// 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 消息格式

// 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：

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

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）：

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。