From c7b23133515131d2debdd1c4fe3af0fe08b18a9d Mon Sep 17 00:00:00 2001 From: wonder Date: Fri, 29 May 2026 16:55:37 +0800 Subject: [PATCH] =?UTF-8?q?docs:=20=E6=96=B0=E5=A2=9E=E9=99=90=E6=B5=81?= =?UTF-8?q?=E4=B8=8E=E5=8D=8F=E7=A8=8B=E6=B1=A0=E5=AE=9E=E6=96=BD=E8=AE=A1?= =?UTF-8?q?=E5=88=92?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/PLAN.md | 431 +++++++++++++++++++++++++++++++++++++++++++++++++ docs/_index.md | 1 + 2 files changed, 432 insertions(+) create mode 100644 docs/PLAN.md diff --git a/docs/PLAN.md b/docs/PLAN.md new file mode 100644 index 0000000..53f5140 --- /dev/null +++ b/docs/PLAN.md @@ -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）+ 全局级（per-endpoint） + +**配置：** + +```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。 diff --git a/docs/_index.md b/docs/_index.md index bde8805..494d0b6 100755 --- a/docs/_index.md +++ b/docs/_index.md @@ -17,6 +17,7 @@ Go + Gin + Eino / Vite + React + TypeScript + zustand / SQLite / 七牛云对象 | [异步任务](async-tasks.md) | 任务生命周期、并发控制、重试策略、进度推送 | | [预设风格键](style-keys.md) | 美术风格、色调、线条等风格键分类与可选值 | | [架构图提示词](prompts/) | 全局/后端/前端架构图生成提示词，供多模态 AI 生成架构图 | +| [实施计划](PLAN.md) | 限流（令牌桶）、协程池、RabbitMQ 异步化的详细设计方案与实施步骤 | ## 系统架构