Goroutine Pool 协程池

难度：⭐⭐⭐ 困难

考点

• 限制最大并发 goroutine 数量
• 任务队列与 worker 生命周期管理
• 优雅关闭（等待所有任务完成）

提示

Go 里面所谓的 goroutine pool（协程池），本质上不是“复用 goroutine”（很多人会误解），而是：

提前创建固定数量的 goroutine，让它们不断从任务队列(channel)中取任务执行，从而限制并发数量。

题目描述

实现一个 goroutine 池，限制最大并发数。

要求：

1. NewPool(maxWorkers) — 创建指定大小的协程池
2. Submit(task) — 提交任务到池中执行
3. Wait() — 等待所有已提交的任务完成
4. 同一时刻运行的 goroutine 不超过 maxWorkers
5. Submit 在池未关闭时不应阻塞（除非需要背压，这里不要求）

函数签名

type Pool struct { ... }

func NewPool(maxWorkers int) *Pool
func (p *Pool) Submit(task func())
func (p *Pool) Wait()

提示

1. 用 buffered channel 作为任务队列
2. 启动固定数量的 worker goroutine，从 channel 读取任务执行
3. Wait 时关闭任务 channel，等待所有 worker 退出
4. 注意：Submit 在 Wait 之后调用应该是安全的（不 panic）

---

参考答案（Go）

点击展开参考答案

package goroutine_pool

import (
	"sync"
)

// 如果要做到固定数量的goroutine，让它们不断从任务队列中取任务执行。
// 需要一个ch 存放待执行的任务
// 需考虑竞争锁问题

// Pool 是一个固定大小的 goroutine 池
type Pool struct {
	// TODO: 定义你的字段
	jobs      chan func()    // 任务队列
	taskWG    sync.WaitGroup // 等待任务完成
	workerWG  sync.WaitGroup // 等待worker退出
	closeOnce sync.Once      // 确保只关闭一次
	mu        sync.Mutex     // 保护任务计数器
	closed    bool           // 标记池是否已关闭
}

// NewPool 创建一个最多 maxWorkers 个 worker 的协程池
func NewPool(maxWorkers int) *Pool {
	// TODO: 在这里实现你的代码
	if maxWorkers <= 0 {
		panic("maxWorkers must be greater than 0")
	}
	p := &Pool{
		jobs:   make(chan func(), maxWorkers),
		closed: false,
	}

	for i := 0; i < maxWorkers; i++ {
		p.workerWG.Add(1)
		go func() {
			defer p.workerWG.Done() // 确保worker退出时调用Done

			for task := range p.jobs {
				func() {
					defer p.taskWG.Done()
					defer func() { recover() }() // 捕获panic，防止worker退出 避免没法defer到workerWG.Done导致一直阻塞
					task()
				}()

			}
		}()
	}

	return p
}

// Submit 提交任务到池中执行
// 不应阻塞调用者（除非内部队列满）
func (p *Pool) Submit(task func()) {
	// TODO: 在这里实现你的代码
	if task == nil {
		return
	}
	p.mu.Lock()
	if p.closed {
		p.mu.Unlock()
		return
	}

	// 任务提交时，增加等待计数
	p.taskWG.Add(1)

	// 放入任务队列，由NewPool中的for worker: range p.jobs 取出执行
	p.jobs <- task
	p.mu.Unlock()

}

// Wait 等待所有已提交的任务完成
func (p *Pool) Wait() {
	// TODO: 在这里实现你的代码
	// panic("not implemented")

	p.closeOnce.Do(func() {
		p.mu.Lock()
		p.closed = true
		p.mu.Unlock()

		p.taskWG.Wait()
		close(p.jobs) //Wait结束，直接关闭 jobs，此时range 会退出。
	})

	// 等待所有worker退出
	p.workerWG.Wait()
}