Analyze the design and implementation of Golang coroutine pool gopool

Get into the habit of writing together! This is the sixth day of my participation in the "Nuggets Daily New Plan · April Update Challenge", click to view the details of the event .

Goroutine

Goroutine is a lightweight thread provided by Golang. We usually call it "coroutine". Compared with threads, the cost of creating a coroutine is very low. So you will often see thousands of coroutines concurrently appear in applications developed by Golang.

Advantages of goroutines:

• Goroutines are cheap compared to threads.

Their stack size is only a few kilobytes, the stack can grow and shrink according to the needs of the application, the context switch is also fast, and in the case of threads, the stack size must be specified and fixed.

• Goroutines are multiplexed to a smaller number of OS threads.

A program with thousands of goroutines may have only one thread. If any of the goroutines in this thread block waiting for user input, create another OS thread and move the remaining goroutines to the new OS thread. All of this is handled by the runtime, and as a developer, we don't need to worry about it, which also allows us to have a very clean API to support concurrency.

• Goroutines use channels to communicate.

The design of channels effectively prevents race conditions from occurring when using Goroutines to access shared memory. A channel can be thought of as a channel through which Goroutines communicate.

In the following we will refer to Goroutine as "coroutine".

coroutine pool

In high concurrency scenarios, we may start a large number of coroutines to process business logic. Coroutine pooling is a technology that uses pooling technology to reuse objects, reduce the frequency of memory allocation and the overhead of coroutine creation, thereby improving the efficiency of coroutine execution.

最近抽空了解了字节官方开源的 gopkg 库提供的 gopool 协程池实现，感觉还是很高质量的，代码也非常简洁清晰，而且 Kitex 底层也在使用 gopool 来管理协程，这里我们梳理一下设计和实现。

gopool

Repository：github.com/bytedance/g…

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.

了解官方 README 就会发现gopool的用法其实非常简单，将曾经我们经常使用的 go func(){...} 替换为 gopool.Go(func(){...}) 即可。

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

old:

go func() {
	// do your job
}()
复制代码

new:

import (
    "github.com/bytedance/gopkg/util/gopool"
)

gopool.Go(func(){
	/// do your job
})
复制代码

核心实现

下面我们来看看gopool是怎样实现协程池管理的。

Pool

Pool 是一个定义了协程池能力的接口。

type Pool interface {
	// 池子的名称
	Name() string
        
	// 设置池子内Goroutine的容量
	SetCap(cap int32)
        
	// 执行 f 函数
	Go(f func())
        
	// 带 ctx，执行 f 函数
	CtxGo(ctx context.Context, f func())
        
	// 设置发生panic时调用的函数
	SetPanicHandler(f func(context.Context, interface{}))
}
复制代码

gopool 提供了这个接口的默认实现（即下面即将介绍的pool)，当我们直接调用 gopool.CtxGo 时依赖的就是这个。

这样的设计模式在 Kitex 中也经常出现，所有的依赖均设计为接口，便于随后扩展，底层提供一个默认的实现暴露出去，这样对调用方也很友好。

type pool struct {
	// 池子名称
	name string

	// 池子的容量, 即最大并发工作的 goroutine 的数量
	cap int32
        
	// 池子配置
	config *Config
        
	// task 链表
	taskHead  *task
	taskTail  *task
	taskLock  sync.Mutex
	taskCount int32

	// 记录当前正在运行的 worker 的数量
	workerCount int32

	// 当 worker 出现panic时被调用
	panicHandler func(context.Context, interface{})
}

// NewPool 创建一个新的协程池，初始化名称，容量，配置
func NewPool(name string, cap int32, config *Config) Pool {
	p := &pool{
		name:   name,
		cap:    cap,
		config: config,
	}
	return p
}
复制代码

调用 NewPool 获取了以 Pool 的形式返回的 pool 结构体。

Task

type task struct {
	ctx context.Context
	f   func()

	next *task
}
复制代码

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

综合前一节 pool 的定义，我们可以看到，一个协程池 pool 对应了一组task。

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

Worker

type worker struct {
	pool *pool
}
复制代码

一个 worker 就是逻辑上的一个执行器，它唯一对应到一个协程池 pool。当一个worker被唤起，将会开启一个goroutine ，不断地从 pool 中的 task链表获取任务并执行。

func (w *worker) run() {
	go func() {
		for {
                        // 声明即将执行的 task
			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)
			}
                        // 如果前一步拿到的 taskHead 为空，说明无任务需要执行，清理后返回
			if t == nil {
				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()
		}
	}()
}
复制代码

整体来看

看到这里，其实就能把整个流程串起来了。我们来看看对外的接口 CtxGo(context.Context, f func()) 到底做了什么？

func Go(f func()) {
	CtxGo(context.Background(), f)
}

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

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)
        
        
	// 以下两个条件满足时，创建新的 worker 并唤起执行：
	// 1. task的数量超过了配置的限制 
	// 2. 当前运行的worker数量小于上限（或无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()
	}
}
复制代码

相信看了代码注释，大家就能理解发生了什么。

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

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

const (
	defaultScalaThreshold = 1
)

// Config is used to config pool.
type Config struct {
	// 控制扩容的门槛，一旦待执行的 task 超过此值，且 worker 数量未达到上限，就开始启动新的 worker
	ScaleThreshold int32
}

// NewConfig creates a default Config.
func NewConfig() *Config {
	c := &Config{
		ScaleThreshold: defaultScalaThreshold,
	}
	return c
}

复制代码

defaultPool 的名称为 gopool.DefaultPool，池子容量一万，扩容下限为 1。

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

• 若此时已经有很多 worker 启动（底层一个 worker 对应一个 goroutine），不需要扩容，就直接返回。
• If it is determined that expansion is required, a new one is created workerand the worker.run()method start. Each workerwill asynchronously poolcheck whether there are still tasks to be executed in the task list, and execute if there are any.

Positioning of the three roles

• taskIt is a task node to be executed, and also contains a pointer to the next task, a linked list structure;
• workeris an executor that actually executes tasks, it will asynchronously start an goroutineexecution coroutine pool that is not executed task;
• poolIt is a logical coroutine pool, which corresponds to a tasklinked list, and is responsible for maintaining taskstate updates and creating new ones when needed worker.

Performance optimization with sync.Pool

In fact, at this place, gopoolit is already a coroutine pool library with concise and clear code, but there is obviously room for improvement in performance, so gopoolthe author has applied many times sync.Poolto create pooled objects and reuse worker and task objects.

It is recommended that you look at the source code directly. In fact, it has already been involved in the above code.

• task pooling

var taskPool sync.Pool

func init() {
	taskPool.New = newTask
}

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

func (t *task) Recycle() {
	t.zero()
	taskPool.Put(t)
}
复制代码

• worker pooling

var workerPool sync.Pool

func init() {
	workerPool.New = newWorker
}

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

func (w *worker) Recycle() {
	w.zero()
	workerPool.Put(w)
}
复制代码