Go语言的协程池| 豆包MarsCode AI刷题

Go语言协程的使用可以说是非常简单了，定义一个函数然后go就可以开启协程了。但是，这样我们怎么管理这个协程呢？ 接下来，咱们就来简单实现一下go的线程池的demo。

1. 协程池的创建

在JAVA里面，线程池是一堆线程去一个阻塞队列里面取对象出来调用它的run方法。同样的，go虽然没有线程对象了，但是协程我们可以直接go，而阻塞队列我们一样可以用channel去实现多个协程同步的按顺序取任务。

我们需要一个channel可以存放我们给协程池的任务，同时需要开启很多个协程，让他们去channel中取任务，同时，我们需要一个方法提交任务，以及开启线程池与关闭线程池，实现如下：

type (
	Task          func()
	GoroutinePool struct {
		task_queue   chan Task
		poolCoreSize int
		wg           *sync.WaitGroup
		ctx          context.Context
	}
)
func NewGoroutinePool(poolCoreSize int, poolmaxSize int, queueSize int) *GoroutinePool {
	return &GoroutinePool{
		task_queue:   make(chan Task, queueSize),
		poolCoreSize: poolCoreSize,
		wg:           &sync.WaitGroup{},
		ctx:          ctx,
	}
}
func (p *GoroutinePool) Submit(task Task) {
	p.task_queue <- task
}
func (p *GoroutinePool) Start() {
	for i := 0; i < p.poolCoreSize; i++ {
		p.wg.Add(1)
		go func() {
			defer p.wg.Done()
			for {
				select {
				case task, ok := <-p.task_queue:
					if !ok {
						fmt.Println("协程退出")
						return
					}
					task()
			}
		}()
	}
}
func (p *GoroutinePool) Stop() {
	close(p.task_queue)
	p.wg.Wait()
	fmt.Print("协程池结束")
}

不过，上面的实现虽然可以实现多个协程的复用，但是另一个关键的动态扩容与收缩的能力却还是没有，关于收缩，我们可以设置一个定时器，要是超过这个时间还是没有得到task我们就让这个协程结束（不会小于核心协程数）。而动态扩容就会简单一点了，我们判断一下当前的任务数，要是提交任务的时候，任务数小于核心协程的个数，我们就让它再go一个协程。如果当前的channel满了，同时协程数小于max，我们就也让它go一个新协程。

2. 线程的动态伸缩

要实现这种逻辑，我们需要对协程池进行一些调整。我们不能在启动协程池时就创建所有协程，而是要在提交任务时根据需要动态创建协程。同时，创建协程的行为必须是同步的，以确保不会创建过多的协程。

type (
	Task          func()
	GoroutinePool struct {
		task_queue   chan Task
		poolCoreSize int
		queueSize    int
		poolMaxSize  int
		size         int
		wg           *sync.WaitGroup
		lock         *sync.Mutex
		ctx          context.Context
		timeout      time.Duration
	}
)
func NewGoroutinePool(poolCoreSize int, poolmaxSize int, queueSize int, timeout time.Duration) *GoroutinePool {
	return &GoroutinePool{
		task_queue:   make(chan Task, queueSize),
		poolCoreSize: poolCoreSize,
		poolMaxSize:  poolmaxSize,
		queueSize:    queueSize,
		size:         0,
		wg:           &sync.WaitGroup{},
		lock:         &sync.Mutex{},
		timeout:      timeout,
		ctx:          context.Background(),
	}
}
func (p *GoroutinePool) Submit(task Task) {
	if p.size < p.poolCoreSize {
		p.addGoroutine()
	}
	if p.trySubmit(task) {
		return
	} else if p.size < p.poolMaxSize {
		p.addGoroutine()
	}
	p.mustSubmit(task)
}
func (p *GoroutinePool) trySubmit(task Task) bool {
	select {
	case p.task_queue <- task:
		return true
	default:
		return false
	}
}
func (p *GoroutinePool) mustSubmit(task Task) {
	p.task_queue <- task
}
func (p *GoroutinePool) addGoroutine() {
	p.lock.Lock()
	defer p.lock.Unlock()
	if p.size >= p.poolMaxSize {
		return
	}
	p.size++
	fmt.Println("协程启动", p.size)
	p.wg.Add(1)
	index := p.size
	go func() {
		defer p.wg.Done()
		ticker := time.NewTicker(p.timeout)
		for {
			select {
			case task, ok := <-p.task_queue:
				if !ok {
					fmt.Println("协程退出", index)
					return
				}
				ticker.Stop()
				task()
				ticker.Reset(p.timeout)
			case <-ticker.C:
				if index > p.poolCoreSize {
					fmt.Println("协程超时退出", index, len(p.task_queue))
					return
				}
			}
		}
	}()
}
func (p *GoroutinePool) Stop() {
	close(p.task_queue)
	p.wg.Wait()
	fmt.Print("协程池结束")
}

需要注意的是，ticker是不管你协程有没有运行task都会按时发消息的，所以我们得在拿到任务之后就把它给停了，不然如果过task本身的时间超过了倒计时，它也会发信号，导致协程退出的。

以上，我们的线程池demo就算是完成了，通过线程池的使用，我们可以有效地管理协程的生命周期，确保任务能够按顺序执行，同时避免创建过多的协程导致资源耗尽。这种模式在处理大量短生命周期任务时特别有用，可以显著提高程序的性能和响应速度。