简单实现一个协程池 | 青训营笔记

这是我参与「第五届青训营 」伴学笔记创作活动的第 5 天
今天通过自己手动实现一个demo级别的协程池来加深对协程池的认识。

package main

import (
   "fmt"
   "sync"
)

type TaskPool struct {
   tasks chan func()
   wg    sync.WaitGroup
}

func NewTaskPool(size int) *TaskPool {
   pool := &TaskPool{tasks: make(chan func(), size)}

   pool.wg.Add(size)
   for i := 0; i < size; i++ {
      go func(idx int) {
         defer pool.wg.Done()
         for task := range pool.tasks {
            fmt.Printf("goroutine%d handling ", idx)
            task()
         }
         
      }(i)
   }
   return pool
}
func (p *TaskPool) SummitTask(task func()) {
   p.tasks <- task
}
func (p *TaskPool) Shutdown() {
   close(p.tasks)
   p.wg.Wait()
}

func main() {
   pool := NewTaskPool(5)
   defer pool.Shutdown()
   for i := 0; i < 10; i++ {
      idx := i
      pool.SummitTask(func() {
         fmt.Printf("task%d work!\n", idx)
      })
   }
}

在这个示例中，我们首先定义了一个 WorkerPool 结构体，其中包含了一个 work 通道和一个 WaitGroup 等待组。在 NewWorkerPool() 函数中，我们创建了一个大小为 size 的协程池，每个协程都会从 work 通道中获取任务并执行。在每个协程执行结束后，我们调用 WaitGroup 的 Done() 方法来标记该协程已经完成了工作，最终使用 Wait() 方法等待所有协程完成工作。

在 Submit() 方法中，我们将任务 task 添加到 work 通道中，由协程池中的协程来处理。在 Shutdown() 方法中，我们关闭 work 通道并调用 Wait() 方法等待所有的协程完成工作。

在 main() 函数中，我们首先创建了一个大小为 3 的协程池。然后，我们将 10 个任务添加到协程池中，每个任务都会调用 Submit() 方法将自己提交到协程池中。在任务执行完成后，我们使用 defer 关键字来调用协程池的 Shutdown() 方法，关闭协程池并等待所有的任务执行完成。

需要实现一个比较完整的协程池，还需要能对协程池的协程进行管理，譬如实现一个排队队列。设置最大协程数、最大存活数、过期时间、定期回收等等。