Go 中的 Worker Pool：限制 goroutine、收集结果和安全取消

1. 概述

Go 中的 goroutine 非常轻量级，所以为每个任务调用 go 是很常见的做法。当任务数达到数万个，且每个任务都调用数据库或 HTTP 时，这种做法容易耗尽连接池或超出外部服务的速率限制。

本文将指导你构建 worker pool，即一定数量的 goroutine 从一个 channel 中接收任务。读完本文，你将学会编写基础 pool、收集结果和错误、用 context 取消、用 errgroup.SetLimit 简化，以及测试 pool 是否真正限制了并发数。

2. 问题：无限制的 goroutine

考虑处理订单列表的代码：

for _, o := range orders {
	go chargeOrder(o) // 50000 个订单 -> 50000 个 goroutine 同时调用 payment API
}

每个 goroutine 只消耗几 KB 内存，所以问题通常不在内存上。问题在于 后端资源：

• 数据库的连接池只有大约 20–50 个连接。
• Payment gateway 限制每秒请求数。
• 进程的文件描述符数量有限制。

Worker pool 通过分离两个概念来解决这个问题。任务数量 可以非常大，但 同时运行的任务数 受 worker 数量的限制。

3. 基础 Worker Pool

该模型由三部分组成：

1. 一个 jobs channel 用于存放要做的任务。
2. N 个 worker goroutine，每个 worker 在 channel 上循环。
3. 一个 sync.WaitGroup 用于等待所有 worker 完成。

package main

import (
	"fmt"
	"sync"
	"time"
)

type Order struct {
	ID     int
	Amount float64
}

func worker(id int, orders <-chan Order, wg *sync.WaitGroup) {
	defer wg.Done()
	for o := range orders {
		time.Sleep(100 * time.Millisecond) // 模拟调用 payment API
		fmt.Printf("worker %d 处理 order %d (%.0f)\n", id, o.ID, o.Amount)
	}
}

func main() {
	orders := make(chan Order)
	var wg sync.WaitGroup

	for i := 1; i <= 3; i++ {
		wg.Add(1)
		go worker(i, orders, &wg)
	}

	start := time.Now()
	for i := 1; i <= 9; i++ {
		orders <- Order{ID: i, Amount: float64(i) * 10}
	}
	close(orders)
	wg.Wait()
	fmt.Println("完成用时", time.Since(start).Round(100*time.Millisecond))
}

代码工作流程如下：

• orders <-chan Order 是一个 只接收 channel。编译器会防止 worker 尝试发送或关闭 channel。
• for o := range orders 在 channel 被 close 且所有元素都被消费后自动退出。
• close(orders) 是信号“工作完成”。如果忘记这行，worker 会一直等待，wg.Wait() 导致死锁。
• wg.Add(1) 必须在 go worker(...) 之前 调用，而不是在 goroutine 内部调用。

运行结果（行的顺序可能不同）：

worker 3 处理 order 1 (10)
worker 1 处理 order 2 (20)
worker 2 处理 order 3 (30)
worker 2 处理 order 4 (40)
...
完成用时 300ms

9 个订单分配给 3 个 worker，每个订单耗时 100ms，所以总时间约 300ms，而不是顺序执行的 900ms。

注意： 从 Go 1.25 开始，你可以写 wg.Go(func() { worker(i, orders) })。这个函数自动调用 Add 和 Done，避免了忘记其中一个的错误。

4. 收集结果和错误

在实际应用中，你通常需要知道哪些订单成功，哪些失败。我们为此添加一个 results channel：

package main

import (
	"errors"
	"fmt"
	"sync"
)

type Order struct {
	ID     int
	Amount float64
}

type Result struct {
	OrderID int
	Fee     float64
	Err     error
}

func chargeOrder(o Order) (float64, error) {
	if o.Amount <= 0 {
		return 0, errors.New("金额无效")
	}
	return o.Amount * 0.02, nil // 手续费 2%
}

func main() {
	orders := []Order{{1, 100}, {2, 0}, {3, 250}, {4, 80}}

	jobs := make(chan Order)
	results := make(chan Result)
	var wg sync.WaitGroup

	for range 3 {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for o := range jobs {
				fee, err := chargeOrder(o)
				results <- Result{OrderID: o.ID, Fee: fee, Err: err}
			}
		}()
	}

	go func() {
		for _, o := range orders {
			jobs <- o
		}
		close(jobs)
	}()

	go func() {
		wg.Wait()
		close(results)
	}()

	for r := range results {
		if r.Err != nil {
			fmt.Printf("order %d 错误: %v\n", r.OrderID, r.Err)
			continue
		}
		fmt.Printf("order %d 手续费 %.2f\n", r.OrderID, r.Fee)
	}
}

输出形式如下：

order 1 手续费 2.00
order 3 手续费 5.00
order 2 错误: 金额无效
order 4 手续费 1.60

有三个关键点需要注意：

• 发送 job 在 单独的 goroutine 中运行。如果 main 既发送 job 又读取结果，worker 会在 results <- 处阻塞，而 main 在 jobs <- 处阻塞，导致死锁。
• 第三个 goroutine 等待 wg.Wait() 后才 close(results)。只有当所有 worker 完成后，才能确保没有人再向 results 发送。通常规则：只有发送方才能关闭 channel。
• 错误被包装在 Result 中，而不是停止整个 pool。一个订单的错误不应该阻止其他订单处理。

结果的顺序 不 与输入顺序一致。如果需要保持顺序，可以在 job 中携带索引 (index)，或使用第 6 节的方法。

5. 使用 Context 取消 Pool

当 HTTP 请求被取消或超时时，pool 必须停止，而不是继续调用 payment API。你可以通过让 worker 监听 ctx.Done() 来实现：

func worker(ctx context.Context, jobs <-chan Order, results chan<- Result) {
	for {
		select {
		case <-ctx.Done():
			return
		case o, ok := <-jobs:
			if !ok {
				return
			}
			fee, err := chargeOrder(o)
			select {
			case results <- Result{OrderID: o.ID, Fee: fee, Err: err}:
			case <-ctx.Done():
				return
			}
		}
	}
}

Producer 端也需要类似的 select：

go func() {
	defer close(jobs)
	for _, o := range orders {
		select {
		case jobs <- o:
		case <-ctx.Done():
			return
		}
	}
}()

在调用处创建带超时的 context：

ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()

需要在 两个发送方向 都做 select 的原因是：如果 worker 只在接收 job 时检查 ctx，当读取端已离开时，它仍可能被永久阻塞在 results <- 处。这是 goroutine leak 的常见来源。

注意： select 只能帮助 worker 在任务 之间 退出。如果 chargeOrder 进行 HTTP 调用，你需要将 ctx 传给它（http.NewRequestWithContext）才能取消正在进行的请求。

6. 使用 errgroup.SetLimit 简化

手动编写 channel、WaitGroup 和 context 容易出错。errgroup 包提供了 SetLimit，只需几行就能创建受限的 pool。

首先，添加依赖：

go get golang.org/x/sync/errgroup

下面的泛型函数用最多 limit 个 goroutine 处理一个 slice，保持结果顺序，并在第一个错误时停止：

package pool

import (
	"context"

	"golang.org/x/sync/errgroup"
)

func Process[T, R any](
	ctx context.Context,
	limit int,
	items []T,
	fn func(context.Context, T) (R, error),
) ([]R, error) {
	results := make([]R, len(items))
	g, ctx := errgroup.WithContext(ctx)
	g.SetLimit(limit)

	for i, item := range items {
		g.Go(func() error {
			r, err := fn(ctx, item)
			if err != nil {
				return err
			}
			results[i] = r
			return nil
		})
	}

	if err := g.Wait(); err != nil {
		return nil, err
	}
	return results, nil
}

代码工作流程如下：

• g.SetLimit(limit) 使 g.Go 在已有 limit 个 goroutine 运行时被 阻塞。在效率上相当于一个 worker pool。
• errgroup.WithContext 返回的 ctx 会在任何函数返回错误时立即取消。使用 ctx 的其他任务会收到停止信号。
• results[i] = r 在并发运行时是安全的，因为每个 goroutine 写入 自己的索引。这种情况下不需要 mutex。
• 从 Go 1.22 开始，循环变量 i 和 item 在每次迭代时都会创建新的副本，所以闭包不会再共享变量，如旧版本那样。

调用方式：

fees, err := pool.Process(ctx, 5, orders, func(ctx context.Context, o Order) (float64, error) {
	return paymentService.Charge(ctx, o)
})

与自写 pool 的区别：errgroup 为每个 item 创建一个新的 goroutine，而不是复用 N 个 worker。对于大多数应用，这个差异可以忽略。如果你需要 worker 保持自己的状态，比如每个 worker 有一个连接，那就使用第 3-5 节的自写 pool。

7. 测试 Worker Pool

你需要测试两个关键行为。首先，pool 永远不会 超过 limit。其次，错误被正确返回。为了测量并发级别，我们使用 atomic.Int32：

package pool

import (
	"context"
	"errors"
	"sync/atomic"
	"testing"
	"time"
)

func TestProcess_LimitsConcurrencyAndKeepsOrder(t *testing.T) {
	var running, maxSeen atomic.Int32
	items := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}

	got, err := Process(context.Background(), 3, items, func(_ context.Context, n int) (int, error) {
		cur := running.Add(1)
		defer running.Add(-1)
		for {
			m := maxSeen.Load()
			if cur <= m || maxSeen.CompareAndSwap(m, cur) {
				break
			}
		}
		time.Sleep(10 * time.Millisecond)
		return n * 2, nil
	})

	if err != nil {
		t.Fatalf("不期望错误: %v", err)
	}
	if maxSeen.Load() > 3 {
		t.Errorf("并发运行 %d，限制是 3", maxSeen.Load())
	}
	want := []int{2, 4, 6, 8, 10, 12, 14, 16, 18, 20}
	for i := range want {
		if got[i] != want[i] {
			t.Fatalf("got[%d] = %d, want %d", i, got[i], want[i])
		}
	}
}

func TestProcess_ReturnsFirstError(t *testing.T) {
	errDeclined := errors.New("卡被拒绝")

	_, err := Process(context.Background(), 2, []int{1, 2, 3}, func(_ context.Context, n int) (int, error) {
		if n == 2 {
			return 0, errDeclined
		}
		return n, nil
	})

	if !errors.Is(err, errDeclined) {
		t.Fatalf("期望 errDeclined，收到 %v", err)
	}
}

运行带 race detector 的测试：

go test -race ./...

ok  	example.com/shop/pool	0.061s

CompareAndSwap 循环在不需要锁的情况下更新 maxSeen。-race 标志对于并发代码是必需的，因为它可以检测出容易被忽视的数据竞争。

8. 常见错误和何时不应该使用

常见错误：

• 忘记 close(jobs)：worker 永远不会退出，程序报告 fatal error: all goroutines are asleep - deadlock!。
• 在接收端关闭 channel，或关闭两次：导致 panic: close of closed channel。
• 在 goroutine 内部调用 wg.Add：Wait 可能在 goroutine 还没来得及 Add 时就返回了。
• 在同一 goroutine 上发送 job 和读取结果，使用无缓冲 channel：如第 4 节分析的那样导致死锁。
• 在发送结果时不 select ctx.Done()：当读取端提前退出时导致 goroutine leak。
• 任意选择 worker 数量：对于 CPU-bound 任务，使用 runtime.NumCPU()。对于 I/O-bound 任务，根据后端资源限制，例如连接池大小。

何时 不应该 使用 worker pool：

• 任务数量小且已知，例如并行调用 3 个服务。此时不需要 SetLimit 的 errgroup 就足够了。
• 任务极短，只有几微秒。通过 channel 的成本可能大于任务本身，所以顺序处理或批处理会更高效。
• 你需要限制 速率（请求/秒）而不是 并发数。这种情况需要 rate limiter，例如 golang.org/x/time/rate。Worker pool 无法替代 rate limiter。

9. 总结

Worker pool 帮助你控制 Go 中的并发级别：任务数量可以很大，但后端资源仍然受到保护。本文经历了以下步骤：

• 使用 channel 和 sync.WaitGroup 的基础 pool。
• 通过 results channel 收集结果和错误，正确关闭 channel。
• 在两个发送方向都用 context 安全取消。
• 用泛型 Process 函数和 errgroup.SetLimit 简化。
• 用 go test -race 测试并发限制和错误。

在大多数情况下，errgroup.SetLimit 是最简单、最不容易出错的选择。只有当 worker 需要保持自己的状态，或需要精细控制 job 和结果流时，才应该自写 pool。