Go 业务中 Worker Pool 的设计实践：Goroutine 防泄漏与生命周期管理

package pool

import (
	"context"
	"fmt"
	"log"
	"sync"
)

type Task func(ctx context.Context)

type WorkerPool struct {
	workerNum int
	taskQueue chan Task
	wg        sync.WaitGroup
	ctx       context.Context
	cancel    context.CancelFunc
}

func NewWorkerPool(workerNum, queueSize int) *WorkerPool {
	ctx, cancel := context.WithCancel(context.Background())
	return &WorkerPool{
		workerNum: workerNum,
		taskQueue: make(chan Task, queueSize),
		ctx:       ctx,
		cancel:    cancel,
	}
}

func (p *WorkerPool) Start() {
	for i := 0; i < p.workerNum; i++ {
		p.wg.Add(1)
		go p.worker()
	}
}

func (p *WorkerPool) worker() {
	defer p.wg.Done()

	for {
		select {
		case <-p.ctx.Done():
			// 收到全局停止信号,退出
			return
		case task, ok := <-p.taskQueue:
			if !ok {
				// 任务队列已关闭
				return
			}
			p.runSafe(task)
		}
	}
}

// 防泄漏核心 1: 兜底 panic,避免单一任务炸掉整个 worker 或进程
func (p *WorkerPool) runSafe(task Task) {
	defer func() {
		if r := recover(); r != nil {
			log.Printf("[WorkerPool] task panic recovered: %v", r)
		}
	}()
	task(p.ctx)
}

// 提交任务: 包含防死锁/背压机制
func (p *WorkerPool) Submit(ctx context.Context, task Task) error {
	select {
	case <-p.ctx.Done():
		return fmt.Errorf("worker pool is stopped")
	case <-ctx.Done():
		return ctx.Err() // 提交方超时放弃,不堵死调用方
	case p.taskQueue <- task:
		return nil
	}
}

// 防泄漏核心 2: 优雅退出,保证已提交的任务消费完,且不漏协程
func (p *WorkerPool) Stop() {
	close(p.taskQueue) // 停止接收新任务,worker 会消费完剩余队列
	p.wg.Wait()        // 等待所有 worker 退出
	p.cancel()         // 兜底释放资源
}

goos: darwin
goarch: arm64
pkg: bench_test
cpu: Apple M2 Pro
BenchmarkDirectGoroutine-10       100000     11842 ns/op     2840 B/op     19 allocs/op
BenchmarkWorkerPool-10            100000      3951 ns/op      320 B/op      3 allocs/op