Go协程池gopool源码解析

1、gopool简介

Repository：https://github.com/bytedance/gopkg/tree/develop/util/gopool

gopool is a high-performance goroutine pool which aims to reuse goroutines and limit the number of goroutines.

It is an alternative to the go keyword.

gopool的用法非常简单，将曾经我们经常使用的go func(){...}替换为gopool.Go(func(){...})即可

此时gopool将会使用默认的配置来管理你启动的协程，也可以选择针对业务场景配置池子大小以及扩容上限

old：

go func() {
	// do your job
}()

new：

gopool.Go(func(){
	// do your job
})

2、核心数据结构

1）、Pool

Pool是一个定义了协程池的接口，代码如下：

// util/gopool/pool.go
type Pool interface {
	// Name returns the corresponding pool name.
	// 协程池的名称
	Name() string
	// SetCap sets the goroutine capacity of the pool.
	// 设置协程池内goroutine的容量
	SetCap(cap int32)
	// Go executes f.
	// 执行f函数
	Go(f func())
	// CtxGo executes f and accepts the context.
	// 带ctx,执行f函数
	CtxGo(ctx context.Context, f func())
	// SetPanicHandler sets the panic handler.
	// 设置发生panic时调用的函数
	SetPanicHandler(f func(context.Context, interface{}))
}

gopool提供了Pool这个接口的默认实现pool，代码如下：

// util/gopool/pool.go
type pool struct {
	// The name of the pool
	// 协程池的名字
	name string

	// capacity of the pool, the maximum number of goroutines that are actually working
	// 协程池实际工作的goroutine的最大数量
	cap int32
	// Configuration information
	// 配置信息
	config *Config
	// linked list of tasks
	// task队列的元信息,每一个task代表一个待执行的函数
	taskHead  *task
	taskTail  *task
	taskLock  sync.Mutex
	taskCount int32

	// Record the number of running workers
	// 当前有多少个worker在运行中,每个worker代表一个goroutine
	workerCount int32

	// This method will be called when the worker panic
	// 协程池中的协程引发的panic会由该函数处理
	panicHandler func(context.Context, interface{})
}

pool数据结构如下图：

2）、task

// util/gopool/pool.go
type task struct {
	// 当前task的ctx
	ctx context.Context
	// 当前task需要执行的函数f
	f func()
	// 指向下一个task的指针
	next *task
}

task是一个链表结构，可以把它理解为一个待执行的任务，包含了当前task需要执行的函数f func()以及指向下一个task的指针

一个协程池pool对应了一组task，pool维护了指向链表的头尾的两个指针：taskHead和taskTail以及链表的长度taskCount和对应的锁taskLock

3）、worker

// util/gopool/worker.go
type worker struct {
	pool *pool
}

一个worker就是逻辑上的一个执行器，它对应到一个协程池pool

当一个worker被唤起，将会开启一个goroutine，不断从pool中的task链表获取任务并执行，代码如下：

// util/gopool/worker.go
func (w *worker) run() {
	go func() {
		for {
			var t *task
			// 操作pool中的task链表前,加锁保证并发安全
			w.pool.taskLock.Lock()
			if w.pool.taskHead != nil {
				// 拿到taskHead准备执行
				t = w.pool.taskHead
				// 更新链表的head以及数量
				w.pool.taskHead = w.pool.taskHead.next
				atomic.AddInt32(&w.pool.taskCount, -1)
			}
			if t == nil {
				// if there's no task to do, exit
				// 如果前一步拿到的taskHead为空,说明无任务需要执行,清理后返回(关闭goroutine)
				w.close()
				w.pool.taskLock.Unlock()
				w.Recycle()
				return
			}
			w.pool.taskLock.Unlock()
			// 执行任务,针对panic会recover,并调用配置的handler
			func() {
				defer func() {
					if r := recover(); r != nil {
						msg := fmt.Sprintf("GOPOOL: panic in pool: %s: %v: %s", w.pool.name, r, debug.Stack())
						logger.CtxErrorf(t.ctx, msg)
						if w.pool.panicHandler != nil {
							w.pool.panicHandler(t.ctx, r)
						}
					}
				}()
				t.f()
			}()
			t.Recycle()
		}
	}()
}

3、核心API

来看下使用gopool的核心APIGo(f func())，实现如下：

// util/gopool/gopool.go
func Go(f func()) {
	CtxGo(context.Background(), f)
}

func CtxGo(ctx context.Context, f func()) {
	defaultPool.CtxGo(ctx, f)
}

// util/gopool/pool.go
func (p *pool) CtxGo(ctx context.Context, f func()) {
	// 创建一个task对象,将ctx和待执行的函数赋值
	t := taskPool.Get().(*task)
	t.ctx = ctx
	t.f = f
	// 将task插入pool的链表的尾部,更新链表数量
	p.taskLock.Lock()
	if p.taskHead == nil {
		p.taskHead = t
		p.taskTail = t
	} else {
		p.taskTail.next = t
		p.taskTail = t
	}
	p.taskLock.Unlock()
	atomic.AddInt32(&p.taskCount, 1)
	// The following two conditions are met:
	// 1. the number of tasks is greater than the threshold.
	// 2. The current number of workers is less than the upper limit p.cap.
	// or there are currently no workers.
	// 以下任意条件满足时,创建新的worker并唤起执行
	// 1.待执行的task超过了扩容阈值(默认值为1)且当前运行的worker数量小于上限(默认值为10000)
	// 2.无worker运行
	if (atomic.LoadInt32(&p.taskCount) >= p.config.ScaleThreshold && p.WorkerCount() < atomic.LoadInt32(&p.cap)) || p.WorkerCount() == 0 {
		// worker数量+1
		p.incWorkerCount()
		// 创建一个新的worker,并把当前pool赋值
		w := workerPool.Get().(*worker)
		w.pool = p
		// 唤起worker执行
		w.run()
	}
}

以下任意条件满足时，会扩容worker：

待执行的task超过了扩容阈值（默认值为1）且当前运行的worker数量小于上限（默认值为10000）
无worker运行

gopool自行维护一个defaultPool，这是一个默认的pool结构体，在引入包的时候就进行初始化。当我们直接调用gopool.Go()时，本质上是调用了defaultPool的同名方法

// util/gopool/gopool.go
var defaultPool Pool

func init() {
	defaultPool = NewPool("gopool.DefaultPool", 10000, NewConfig())
}

// util/gopool/config.go
const (
	defaultScalaThreshold = 1
)

type Config struct {
	// threshold for scale.
	// new goroutine is created if len(task chan) > ScaleThreshold.
	// defaults to defaultScalaThreshold.
	// 控制扩容的阈值,一旦待执行的task超过此值,且worker数量未达到上限,就开始启动新的worker
	ScaleThreshold int32
}

func NewConfig() *Config {
	c := &Config{
		ScaleThreshold: defaultScalaThreshold,
	}
	return c
}

defaultPool的名称为gopool.DefaultPool，池子容量一万，扩容阈值为1

当调用gopool.Go()时，gopool就会更新维护的任务链表，并且判断是否需要扩容worker：

若此时已经有很多worker启动（底层一个worker对应一个goroutine），不需要扩容，就直接返回
若判断需要扩容，就创建一个新的worker，并调用worker.run()方法启动，各个worker会异步地检查pool里面的任务链表是否还有待执行的任务，如果有就执行

gopool中三个角色的定位：

task是一个待执行的任务节点，同时还包含了指向下一个任务的指针，链表结构
worker是一个实际执行任务的执行器，它会异步启动一个goroutine执行协程池里面未执行的task
pool是一个逻辑上的协程池，对应了一个task链表，同时负责维护task状态的更新，以及在需要的时候创建新的worker

gopool核心实现原理如下图：

4、使用sync.Pool进行性能优化

gopool中多次使用了sync.Pool来池化对象的创建，复用woker和task对象

task池化：

// util/gopool/pool.go
var taskPool sync.Pool

func init() {
	taskPool.New = newTask
}

func newTask() interface{} {
	return &task{}
}

func (t *task) Recycle() {
	t.zero()
	taskPool.Put(t)
}

worker池化：

// util/gopool/worker.go
var workerPool sync.Pool

func init() {
	workerPool.New = newWorker
}

func newWorker() interface{} {
	return &worker{}
}

func (w *worker) Recycle() {
	w.zero()
	workerPool.Put(w)
}