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  *tasktaskTail  *tasktaskLock  sync.MutextaskCount int32// Record the number of running workers// 当前有多少个worker在运行中,每个worker代表一个goroutineworkerCount 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的ctxctx context.Context// 当前task需要执行的函数ff 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.nextatomic.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,并调用配置的handlerfunc() {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 = ctxt.f = f// 将task插入pool的链表的尾部,更新链表数量p.taskLock.Lock()if p.taskHead == nil {p.taskHead = tp.taskTail = t} else {p.taskTail.next = tp.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数量+1p.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 Poolfunc 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数量未达到上限,就开始启动新的workerScaleThreshold 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.Poolfunc 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.Poolfunc init() {workerPool.New = newWorker
}func newWorker() interface{} {return &worker{}
}func (w *worker) Recycle() {w.zero()workerPool.Put(w)
}