前言

go 的 goroutine 提供了一種較線程而言更廉價的方式處理并發場景, go 使用二級線程的模式, 將 goroutine 以 M:N 的形式復用到系統線程上, 節省了 cpu 調度的開銷, 也避免了用戶級線程（協程）進行系統調用時阻塞整個系統線程的問題。【1】

但 goroutine 太多仍會導致調度性能下降、GC 頻繁、內存暴漲, 引發一系列問題。在面臨這樣的場景時, 限制 goroutine 的數量、重用 goroutine 顯然很有價值。

本文正是針對上述情況而提供一種簡單的解決方案, 編寫一個協程池（任務池）來實現對 goroutine 的管控。

思路

要解決這個問題, 要思考兩個問題

• goroutine 的數量如何限制, goroutine 如何重用
• 任務如何執行

goroutine 的數量如何限制, goroutine 如何重用

說到限制和重用, 那么最先想到的就是池化。比如 TCP 連接池, 線程池, 都是有效限制、重用資源的最好實踐。所以, 我們可以創建一個 goroutine 池, 用來管理 goroutine。

任務如何執行

在使用原生 goroutine 的場景中, 運行一個任務直接啟動一個 goroutine 來運行, 在池化的場景而言, 任務也是要在 goroutine 中執行, 但是任務需要任務池來放入 goroutine。

生產者消費者模型

在連接池中, 連接在使用時從池中取出, 用完后放入池中。對于 goroutine 而言, goroutine 通過語言關鍵字啟動, 無法像連接一樣操作。那么如何讓 goroutine 可以執行任務, 且執行后可以重新用來執行其它任務呢？這里就需要使用生產者消費者模型了:

生產者 --(生產任務)--> 隊列 --(消費任務)--> 消費者

用來執行任務的 goroutine 可以作為消費者, 操作任務池的 goroutine 作為生產者, 而隊列則可以使用 go 的 buffer channel, 任務池的建模到此結束。

實現

Talk is cheap. Show me the code.

任務的定義

任務要包含需要執行的函數、以及函數要傳的參數, 因為參數類型、個數不確定, 這里使用可變參數和空接口的形式

type Task struct {
    Handler func(v ...interface{})
    Params  []interface{}
}

任務池的定義

任務池的定義包括了池的容量 capacity、當前運行的 worker（goroutine）數量 runningWorkers、任務隊列（channel）chTask 以及任務池的狀態 status（運行中或已關閉, 用于安全關閉任務池）, 最后還有一把互斥鎖 sync.Mutex

type Pool struct {
    capacity       uint64
    runningWorkers uint64
    status          int64
    chTask          chan *Task
    sync.Mutex
}

任務池的構造函數:

var ErrInvalidPoolCap = errors.New("invalid pool cap")

const (
    RUNNING = 1
    STOPED = 0
)

func NewPool(capacity uint64) (*Pool, error) {
    if capacity <= 0 {
        return nil, ErrInvalidPoolCap
    }
    return &Pool{
        capacity: capacity,
        status:    RUNNING,
        // 初始化任務隊列, 隊列長度為容量
        chTask:    make(chan *Task, capacity),
    }, nil
}

啟動 worker

新建 run() 方法作為啟動 worker 的方法:

func (p *Pool) run() {
    p.runningWorkers++ // 運行中的任務加一

    go func() {
        defer func() {
            p.runningWorkers-- // worker 結束, 運行中的任務減一
        }()

        for {
            select { // 阻塞等待任務、結束信號到來
            case task, ok := <-p.chTask: // 從 channel 中消費任務
                if !ok { // 如果 channel 被關閉, 結束 worker 運行
                    return
                }
                // 執行任務
                task.Handler(task.Params...)
            }
        }
    }()
}

上述代碼中, runningWorkers 的加減直接使用了自增運算, 但是考慮到啟動多個 worker 時, runningWorkers 就會有數據競爭, 所以我們使用 sync.atomic 包來保證 runningWorkers 的自增操作是原子的。

對 runningWorkers 的操作進行封裝:

func (p *Pool) incRunning() { // runningWorkers + 1
    atomic.AddUint64(&p.runningWorkers, 1)
}

func (p *Pool) decRunning() { // runningWorkers - 1
    atomic.AddUint64(&p.runningWorkers, ^uint64(0))
}

func (p *Pool) GetRunningWorkers() uint64 {
    return atomic.LoadUint64(&p.runningWorkers)
}

對于 capacity 的操作無需考慮數據競爭, 因為 capacity 在初始化時已經固定。封裝 GetCap() 方法:

func (p *Pool) GetCap() uint64 {
    return p.capacity
}

趁熱打鐵, status 的操作也加鎖封裝為安全操作:

func (p *Pool) setStatus(status int64) bool {
    p.Lock()
    defer p.Unlock()

    if p.status == status {
        return false
    }

    p.status = status

    return true
}

run() 方法改造:

func (p *Pool) run() {
    p.incRunning()

    go func() {
        defer func() {
            p.decRunning()
        }()

        for {
            select {
            case task, ok := <-p.chTask:
                if !ok {
                    return
                }
                task.Handler(task.Params...)
            }
        }
    }()
}

生產任務

新建 Put() 方法用來將任務放入池中:

func (p *Pool) Put(task *Task) {

    // 加鎖防止啟動多個 worker
    p.Lock()
    defer p.Unlock()

    if p.GetRunningWorkers() < p.GetCap() { // 如果任務池滿, 則不再創建 worker
        // 創建啟動一個 worker
        p.run()
    }

    // 將任務推入隊列, 等待消費
    p.chTask <- task
}

任務池安全關閉

當有關閉任務池來節省 goroutine 資源的場景時, 我們需要有一個關閉任務池的方法。

直接銷毀 worker 關閉 channel 并不合適, 因為此時可能還有任務在隊列中沒有被消費掉。要確保所有任務被安全消費后再銷毀掉 worker。

首先, 在關閉任務池時, 需要先關閉掉生產任務的入口。同時, 也要考慮到任務推送到 chTask 時 status 改變的問題。改造 Put() 方法:

var ErrPoolAlreadyClosed = errors.New("pool already closed")

func (p *Pool) Put(task *Task) error {
    p.Lock()
    defer p.Unlock()
    
    if p.status == STOPED { // 如果任務池處于關閉狀態, 再 put 任務會返回 ErrPoolAlreadyClosed 錯誤
        return ErrPoolAlreadyClosed
    }
    
    // run worker
    if p.GetRunningWorkers() < p.GetCap() {
        p.run()
    }

    // send task
    if p.status == RUNNING {
        p.chTask <- task
    }
     
    return nil
}

在 run() 方法中已經對 chTask 的關閉進行了監聽, 銷毀 worker 只需等待任務被消費完后關閉 chTask。Close() 方法如下:

func (p *Pool) Close() {
    p.setStatus(STOPED) // 設置 status 為已停止

    for len(p.chTask) > 0 { // 阻塞等待所有任務被 worker 消費
        time.Sleep(1e6) // 防止等待任務清空 cpu 負載突然變大, 這里小睡一下
    }

    close(p.chTask) // 關閉任務隊列
}

panic handler

每個 worker 都是一個 goroutine, 如果 goroutine 中產生了 panic, 會導致整個程序崩潰。為了保證程序的安全進行, 任務池需要對每個 worker 中的 panic 進行 recover 操作, 并提供可訂制的 panic handler。

更新任務池定義:

type Pool struct {
    capacity       uint64
    runningWorkers uint64
    status          int64
    chTask          chan *Task
    sync.Mutex
    PanicHandler   func(interface{})
}

更新 run() 方法:

func (p *Pool) run() {
    p.incRunning()

    go func() {
        defer func() {
            p.decRunning()
            if r := recover(); r != nil { // 恢復 panic
                if p.PanicHandler != nil { // 如果設置了 PanicHandler, 調用
                    p.PanicHandler(r)
                } else { // 默認處理
                    log.Printf("Worker panic: %s\n", r)
                }
            }
        }()

        for {
            select {
            case task, ok := <-p.chTask:
                if !ok {
                    return
                }
                task.Handler(task.Params...)
            }
        }
    }()
}

可用 worker 檢查

recover 后，gorotine 退出，當池的容量為 1 時，此時會有一個問題，觀察 Put() 方法：

if p.GetRunningWorkers() < p.GetCap() {
    p.run() // 此時有一個 task (上一次 Put) panic，worker 退出了
}


if p.status == RUNNING {
    p.chTask <- task // 當前的 task 推送到 chTask，但是沒有一個 worker 可以消費到，deadlock!
}

感謝提出這個場景的朋友，詳細參考 issue 極端情況 #4，此問題已經在 release v1.5 中修復

為了解決這個 bug，我們需要在 worker 退出時檢查當前是否還有運行著的 worker，如果沒有，則創建一個，保證 task 可以被正常消費，checkWorker() 方法如下：

func (p *Pool) checkWorker() {
    p.Lock()
    defer p.Unlock()

    // 當前沒有 worker 且有任務存在，運行一個 worker 消費任務
    // 沒有任務無需考慮 (當前 Put 不會阻塞，下次 Put 會啟動 worker)
    if p.runningWorkers == 0 && len(p.chTask) > 0 {
        p.run()
    }
}

改造 run() 方法：

func (p *Pool) run() {
    p.incRunning()

    go func() {
        defer func() {
            p.decRunning()
            if r := recover(); r != nil {
                if p.PanicHandler != nil {
                    p.PanicHandler(r)
                } else {
                    log.Printf("Worker panic: %s\n", r)
                }
            }
            p.checkWorker() // worker 退出時檢測是否有可運行的 worker
        }()

        for {
            select {
            case task, ok := <-p.chTask:
                if !ok {
                    return
                }
                task.Handler(task.Params...)
            }
        }
    }()
}

使用

OK, 我們的任務池就這么簡單的寫好了, 試試:

func main() {
    // 創建任務池
    pool, err := NewPool(10)
    if err != nil {
        panic(err)
    }

    for i := 0; i < 20; i++ {
        // 任務放入池中
        pool.Put(&Task{
            Handler: func(v ...interface{}) {
                fmt.Println(v)
            },
            Params: []interface{}{i},
        })
    }

    time.Sleep(1e9) // 等待執行
}

詳細例子見 mortar/examples

benchmark

作為協程池, 性能和內存占用的指標測試肯定是少不了的, 測試數據才是最有說服力的

測試流程

100w 次執行，原子增量操作

測試任務:

var wg = sync.WaitGroup{}

var sum int64

func demoTask(v ...interface{}) {
    defer wg.Done()
    for i := 0; i < 100; i++ {
        atomic.AddInt64(&sum, 1)
    }
}

測試方法:

var runTimes = 1000000
// 原生 goroutine
func BenchmarkGoroutineTimeLifeSetTimes(b *testing.B) {

    for i := 0; i < runTimes; i++ {
        wg.Add(1)
        go demoTask2()
    }
    wg.Wait() // 等待執行完畢
}

// 使用協程池
func BenchmarkPoolTimeLifeSetTimes(b *testing.B) {
    pool, err := NewPool(20)
    if err != nil {
        b.Error(err)
    }

    task := &Task{
        Handler: demoTask2,
    }

    for i := 0; i < runTimes; i++ {
        wg.Add(1)
        pool.Put(task)
    }

    wg.Wait() // 等待執行完畢
}

對比結果

使用任務池和原生 goroutine 性能相近（略好于原生）

使用任務池比直接 goroutine 內存分配節省 7000 倍左右, 內存分配次數減少 2700 倍左右

tips: 當任務為耗時任務時, 防止任務堆積（消費不過來）可以結合業務調整容量, 或根據業務控制每個任務的超時時間

源碼地址

該項目的全部源碼詳見 mortar

參考文章:

【1】線程的 3 種實現方式