golang 实现协程池

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

type Task func()

type ThreadPool struct {
	workerCount int
	taskQueue   chan Task
	wg          sync.WaitGroup
}

func NewThreadPool(workerCount, maxTaskNum int) *ThreadPool {
	pool := &ThreadPool{
		workerCount: workerCount,
		taskQueue:   make(chan Task, maxTaskNum),
	}

	for i := 0; i < workerCount; i++ {
		go pool.worker()
	}

	return pool
}

func (p *ThreadPool) worker() {
	for task := range p.taskQueue {
		task()      // 执行任务
		p.wg.Done() // 任务完成，减少等待组计数
	}
}

func (p *ThreadPool) AddTask(task Task) {
	p.wg.Add(1)
	p.taskQueue <- task // 将任务添加到队列
}

func (p *ThreadPool) Wait() {
	p.wg.Wait() // 等待所有任务完成
}

func main() {
	workerCount := 3
	maxTaskNum := 20
	pool := NewThreadPool(workerCount, maxTaskNum)

	for i := 0; i < 50; i++ {
		taskNum := i + 1000
		pool.AddTask(func() {
			fmt.Printf("Task %d is running\n", taskNum)
			time.Sleep(1 * time.Second)
			fmt.Printf("Task %d is done\n", taskNum)
		})
	}

	pool.Wait() // 等待所有任务完成
}