Go 工作池与上下文取消

功能说明

示例实现一个有界工作池。run 接收任务切片和 worker 数量，创建无缓冲的任务通道与容量等于任务数的结果通道。每个 worker 从任务通道取任务，遇到负值就写入带错误的结果，否则写入平方值。调用方持续读取结果通道，直到所有 worker 退出并且通道关闭，最后检查 context 是否已取消或超时，从而同时拿到部分结果和停止原因。

关键代码

worker 使用两层 select：外层同时等待 ctx.Done 和任务通道，发送结果时再次等待 ctx.Done，避免下游不读取时永久阻塞。run 用 sync.WaitGroup 等待全部 worker，再由单独的 goroutine 关闭结果通道。任务发送方同样在取消时返回并关闭任务通道。workers 小于 1 会直接报错，空任务则提前返回。结果通道带缓冲，减少了发送方与汇总方之间的耦合。

使用注意

context 超时只表示停止领取新任务，已经写入结果通道的值仍会被收集；如果需要全有或全无，业务层必须丢弃部分结果。worker 数量应受下游连接池和 CPU 限制，并非越大越好。错误结果与 context 错误要分开处理，前者是单个任务失败，后者是整体中止。该实现不保证结果顺序，排序应在汇总后完成。

验证方式

直接运行 main 可看到 5 条平方结果和 1 条负数错误。把 context 超时改为 1 纳秒，可观察 stopped early 输出以及不完整的结果列表。编写单元测试时传入已取消的 context，断言 run 返回 context.Canceled，并可用 goleak 检查 worker goroutine 是否全部退出。

代码示例

package main

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

type Job struct {
	ID    int
	Value int
}

type Result struct {
	JobID int
	Out   int
	Err   error
}

func worker(ctx context.Context, jobs <-chan Job, results chan<- Result, wg *sync.WaitGroup) {
	defer wg.Done()
	for {
		select {
		case <-ctx.Done():
			return
		case job, ok := <-jobs:
			if !ok {
				return
			}
			if job.Value < 0 {
				results <- Result{JobID: job.ID, Err: fmt.Errorf("job %d: negative value %d", job.ID, job.Value)}
				continue
			}
			select {
			case <-ctx.Done():
				return
			case results <- Result{JobID: job.ID, Out: job.Value * job.Value}:
			}
		}
	}
}

func run(ctx context.Context, jobs []Job, workers int) ([]Result, error) {
	if workers < 1 {
		return nil, fmt.Errorf("workers must be at least 1, got %d", workers)
	}
	if len(jobs) == 0 {
		return nil, nil
	}

	jobCh := make(chan Job)
	resCh := make(chan Result, len(jobs))

	var wg sync.WaitGroup
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go worker(ctx, jobCh, resCh, &wg)
	}

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

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

	results := make([]Result, 0, len(jobs))
	for result := range resCh {
		results = append(results, result)
	}
	if err := ctx.Err(); err != nil {
		return results, err
	}
	return results, nil
}

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

	jobs := []Job{
		{ID: 1, Value: 4},
		{ID: 2, Value: 9},
		{ID: 3, Value: -1},
		{ID: 4, Value: 16},
		{ID: 5, Value: 25},
		{ID: 6, Value: 36},
	}

	results, err := run(ctx, jobs, 3)
	if err != nil {
		fmt.Println("stopped early:", err)
	}

	failed := 0
	for _, result := range results {
		if result.Err != nil {
			failed++
			fmt.Println("error:", result.Err)
			continue
		}
		fmt.Printf("job %d -> %d\n", result.JobID, result.Out)
	}
	fmt.Printf("%d results, %d failed\n", len(results), failed)
}