Files
cs-note/hzh/Gen2D/02-协程池.md
T

7.6 KiB
Raw Blame History

02 - 协程池 (Worker Pool)

一句话概括:有界并发协程池,NumCPU*4 workers,per-user 限流,背压保护,优雅关闭。

工作流

graph TB
    subgraph Submit["Submit() 入口"]
        CHECK_CLOSED{"池已关闭?"}
        CHECK_USER{"per-user 限流<br/>active >= max?"}
        CHECK_FULL{"channel 满?"}
    end

    subgraph Channel["有界任务队列"]
        TASK_CHAN["chan Task<br/>capacity = queueSize"]
    end

    subgraph Workers["Worker 协程"]
        W1["Worker 0"]
        W2["Worker 1"]
        W3["Worker ..."]
        WN["Worker N-1"]
    end

    subgraph Execute["任务执行"]
        TIMEOUT["context.WithTimeout<br/>10 min"]
        FN["task.Fn(ctx)"]
        RELEASE["释放用户槽位"]
    end

    CHECK_CLOSED -->|"ErrPoolClosed"| REJECT["返回错误"]
    CHECK_CLOSED -->|否| CHECK_USER
    CHECK_USER -->|"ErrUserLimitReached"| REJECT
    CHECK_USER -->|通过| CHECK_FULL
    CHECK_FULL -->|"ErrPoolFull (HTTP 503)"| REJECT
    CHECK_FULL -->|通过| TASK_CHAN
    TASK_CHAN --> W1 & W2 & W3 & WN
    W1 & W2 & W3 & WN --> TIMEOUT --> FN --> RELEASE

Pool 结构体

type Pool struct {
    workers    int              // worker 数量(默认 NumCPU*4)
    queueSize  int              // 任务队列容量(默认 100)
    maxPerUser int              // 单用户最大并发数(默认 2)
    taskQueue  chan Task         // 有界任务队列
    wg         sync.WaitGroup
    ctx        context.Context
    cancel     context.CancelFunc
    metrics    *Metrics

    mu         sync.Mutex
    userActive map[string]int   // userID -> 当前并发数

    closed   atomic.Bool
    closeCh  chan struct{}
}

核心字段说明:

字段 类型 默认值 说明
workers int runtime.NumCPU() * 4 IO 密集型调参
queueSize int 100 有界缓冲,触发背压
maxPerUser int 2 防止单用户占满池
taskQueue chan Task buffered channel 有界队列,固定容量
userActive map[string]int — 记录每用户活跃任务数

Functional Options 模式

协程池采用 Functional Options 模式进行配置,开箱即用、可选覆盖:

pool := workerpool.New(
    workerpool.WithWorkers(16),       // 覆盖默认的 NumCPU*4
    workerpool.WithQueueSize(200),    // 覆盖默认的 100
    workerpool.WithMaxPerUser(5),     // 覆盖默认的 2
)
Option 默认值 约束 说明
WithWorkers(n) NumCPU * 4 n < 1 时强制为 1 IO 密集型场景推荐 4 倍 CPU 核数
WithQueueSize(n) 100 n < 1 时强制为 1 有界缓冲,满时触发背压
WithMaxPerUser(n) 2 n < 1 时置 0(不限制) 防止单用户占满池

💡 为什么选择 Functional Options? 可选参数天然为零值时保持默认,新增配置项无需修改 New() 签名,调用方按需指定。

有界并发

协程池的核心是 make(chan Task, queueSize) 创建的 有界缓冲 channel。

工作原理:

  1. Submit() 尝试将任务写入 channel
  2. N 个 Worker 协程从 channel 中取任务执行
  3. channel 满时进入背压逻辑(见下文)
// Submit 核心逻辑
select {
case p.taskQueue <- task:      // 成功入队
    return nil
case <-p.ctx.Done():           // 池已关闭
    return ErrPoolClosed
default:                       // 队列满,背压
    return ErrPoolFull
}

Per-user 限流

每个用户同时执行的任务数受到 maxPerUser 限制,防止单用户占满整个池。

实现机制:

Submit() 调用流程:
1. 加锁检查 userActive[userID]
2. 若 active >= maxPerUser → 返回 ErrUserLimitReached
3. 否则 userActive[userID]++(预留槽位)
4. 任务执行完毕后 defer releaseUserSlot(userID)

槽位生命周期:

阶段 操作 说明
Submit userActive[userID]++ 预留槽位,不等到 Worker 取出
Execute — Worker 从 channel 取出后开始执行
Complete defer userActive[userID]-- 任务完成或失败时释放

⚠️ 槽位预留时机:在 Submit() 而非 worker() 中预留,确保 channel 满时不会出现"槽位已分配但任务未入队"的不一致状态。

背压保护

当任务队列已满时,协程池通过 select + default 实现非阻塞拒绝:

队列满(channel 已达 capacity)
    │
    ▼
select 进入 default 分支
    │
    ├── 回滚用户槽位(如果有)
    ├── 原子递增 RejectedTasks
    └── 返回 ErrPoolFull
        │
        ▼
    Handler 层映射为 HTTP 503 Service Unavailable
    响应体:"系统繁忙,请稍后重试"

背压 vs 阻塞:

策略 行为 适用场景
非阻塞拒绝(Gen2D) 立即返回错误 用户交互型 API,快速失败
阻塞等待 阻塞直到有空位 批处理系统,不能丢任务

Gen2D 选择非阻塞拒绝,因为用户期望快速得到反馈,而非无限等待。

优雅关闭

协程池支持优雅关闭,确保正在执行的任务有时间完成:

收到 SIGINT / SIGTERM
    │
    ▼
consumeCancel()              ← 停止消费者,不再接收新任务
    │
    ▼
pool.Shutdown(ctx)           ← 30 秒超时
    │
    ├── closed.Swap(true)    ← 停止接收新任务
    ├── cancel()             ← 通知 Worker 停止取任务
    ├── wg.Wait()            ← 等待所有 Worker 退出
    │
    ├── 成功 → 日志 "workerpool shutdown gracefully"
    └── 超时 → 日志 "workerpool shutdown timeout"

Shutdown 返回值:

返回值 含义
true 所有任务正常完成
false 超时,部分任务可能丢失

Metrics 指标

协程池内置原子计数器,支持运行时观测:

指标 类型 说明
ActiveWorkers atomic.Int32 当前正在执行任务的 Worker 数
QueuedTasks atomic.Int32 当前排队等待的任务数
CompletedTasks atomic.Int64 已完成任务总数
FailedTasks atomic.Int64 失败任务总数
SubmittedTasks atomic.Int64 提交任务总数
RejectedTasks atomic.Int64 被拒绝任务总数(队列满或用户限流)

所有指标通过 Metrics() 方法返回只读快照,同时上报 Prometheus。

Task 结构体

type Task struct {
    ID       string                          // 任务唯一标识
    UserID   string                          // 所属用户 ID
    Fn       func(ctx context.Context) error // 任务执行函数
    Priority int                             // 优先级(预留)
}

每个任务携带超时 context(默认 10 分钟),从池的根 context 派生,确保优雅关闭时能取消正在执行的任务。

与 TaskQueue 的协作

TaskQueue(全局排队)
    │
    ▼
Consumer(消费消息)
    │
    ▼
WorkerPool.Submit()(单机并发控制)
    │
    ▼
Worker 执行 Pipeline
组件 职责 范围
TaskQueue 持久化排队,解耦提交与执行 全局(可跨机器)
WorkerPool 单机并发控制,per-user 限流 单机
Consumer 桥接 TaskQueue 与 WorkerPool 单机

关联文档