go协程池代码

package main

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

type Task struct {
	ID  int
	Job func()
}

type Pool struct {
	taskQueue chan Task
	wg        sync.WaitGroup
}

// NewPool 创建协程池 size taskQueue消息队列的数量
// size taskQueue消息队列的数量
func NewPool(size int) *Pool {
	return &Pool{
		taskQueue: make(chan Task, size),
	}
}

func (p *Pool) worker() {
	defer p.wg.Done()
	for task := range p.taskQueue {
		fmt.Printf("Worker started task %d\n", task.ID)
		time.Sleep(time.Second)
		task.Job()
	}
}

func (p *Pool) Stop() {
	close(p.taskQueue)
}

func (p *Pool) Wait() {
	p.wg.Wait()
}

// AddWorker 添加协程数量
// num 协程数量
func (p *Pool) AddWorker(num int) {
	p.wg.Add(num)
	for i := 0; i < num; i++ {
		go p.worker()
	}
}

// AddTask 添加任务到协程池
func (p *Pool) AddTask(task Task) {
	p.taskQueue <- task
}

func main() {
	// 创建一个协程池，任务len为3
	pool := NewPool(0)
	pool.AddWorker(3)

	// 添加任务到协程池
	for i := 0; i < 10; i++ {
		taskID := i
		task := Task{
			ID: taskID,
			Job: func() {
				fmt.Printf("Task %d is running\n", taskID)
			},
		}
		pool.AddTask(task)
	}

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