Golang協程池gopool設計與實現

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

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

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

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{}))
}

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
}

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

	next *task
}

type worker struct {
	pool *pool
}

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()
		}
	}()
}

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()
	}
}

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
}

var taskPool sync.Pool

func init() {
	taskPool.New = newTask
}

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

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

var workerPool sync.Pool

func init() {
	workerPool.New = newWorker
}

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

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