156. Go手写一个简易的协程池

代码地址：https://gitee.com/lymgoforIT/golang-trick/tree/master/45-gpool

前言

在工作中，我们经常需要开启不少的协程去处理一些事情，当然协程池是一个不错的选择，并且也有一些优秀的开源库可以直接使用。但是授人以鱼不如授人以渔，知其然更知其所以然，才能让我们更好的使用它，本文就实现一个简单的协程池，对于一些简单的业务，可以直接使用。

一、直观版

假如我们有一百个（实际可能非常大）任务需要处理，但是挨个处理会太慢，不能满足业务诉求，但是每个任务都开启一个协程去处理，则可能导致协程爆炸，而且每次开启和关闭协程也是一笔开销，会影响性能。因此我们想到，使用五个常驻协程去处理，代码如下：

package main

import (
	"fmt"
	"sync"
)

func task(id int) {
	fmt.Println(id)
}

func main() {
	taskNums := 100
	concurrency := 5
	taskCh := make(chan int, concurrency)
	var wg sync.WaitGroup
	// 开启多个协程执行任务
	for i := 0; i < concurrency; i++ {
		wg.Add(1)
		go func() {
			defer func() {
				wg.Done()
				if err := recover(); err != nil {
					fmt.Println(err)
				}
			}()
			// 从channel中获取任务并执行
			for taskID := range taskCh {
				task(taskID)
			}
		}()
	}

	// 向channel中发送任务
	for i := 0; i < taskNums; i++ {
		taskCh <- i
	}
	close(taskCh)
	wg.Wait()
}

执行结果就是从0打印到99。

上面这份代码模式在工作中是非常常见的，但是每次有相关诉求的时候，都重新写一份，似乎有点冗余不优雅，且每次的任务可能还是不同的，所以我们应该做一点封装，让其具有通用性，那就是简易版的协程池了。

二、封装版

package main

import (
	"fmt"
	"sync"
)

// ITask 任务接口,实现了该接口，也即实现了Run方法，便可以作为协程池所需的任务进行提交
// 工作中，有时候不是定义接口作为协程池的参数，而是将函数作为协程池的参数
type ITask interface {
	Run()
}

type IGoPool interface {
	Start()
	Submit(task ITask)
	WaitAndStop()
}

type gPool struct {
	workerNum int
	taskCh    chan ITask
	wg        sync.WaitGroup
}

func NewGoPool(workerNum int) IGoPool {
	return &gPool{
		workerNum: workerNum,
		taskCh:    make(chan ITask, workerNum),
	}
}

// Start 开启协程池，协程池的大小为workerNum，协程池的协程会从taskCh中获取任务并执行
func (g *gPool) Start() {
	for i := 0; i < g.workerNum; i++ {
	    g.wg.Add(1)
		go func() {
			defer func() {
				g.wg.Done()
				if err := recover(); err != nil {
					fmt.Println(err)
				}
			}()
			for task := range g.taskCh {
				task.Run()
			}
		}()
	}
}

// Submit 提交任务到协程池中
func (g *gPool) Submit(task ITask) {
	g.taskCh <- task
}

// WaitAndStop 等待所有任务提交完毕后关闭通道，然后等待所有协程退出
func (g *gPool) WaitAndStop() {
	close(g.taskCh)
	g.wg.Wait()
}

代码说明：

• 定义了ITask接口，实现该接口的结构体就能作为线程池的任务提交。
  定义协程池接口IGoPool以及实现类gPool，具有启动协程池，提交任务，关闭协程池等方法。

使用示例：

package main

type Task struct {
	ID int
}

func (t *Task) Run() {
	fmt.Println(t.ID)
}

func main() {
	gPool := NewGoPool(5)
	gPool.Start()
	for i := 0; i < 100; i++ {
		gPool.Submit(&Task{
			ID: i,
		})
	}
	gPool.WaitAndStop()
}

• 定义Task结构体，实现了ITask接口，从而可以作为协程池的任务提交。
• 创建一个协程池并启动。
• 提交100个任务。
• 任务提交完毕后，等待协程池中所有任务执行完毕后关闭协程池。