技術實踐——教你用100行寫一個 go 的協程池 (任務池)！！！

有隻黑白貓發表於2020-01-10
Go

點選這裡，檢視封裝 GetCap() 方法等重要內容

實現
Talk is cheap. Show me the code.

任務的定義
任務要包含需要執行的函式、以及函式要傳的引數, 因為引數型別、個數不確定, 這裡使用可變引數和空介面的形式

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

任務池的定義
任務池的定義包括了池的容量 capacity、當前執行的 worker（goroutine）數量 runningWorkers、任務佇列（channel）taskC、關閉任務池的 channel closeC 以及任務池的狀態 state（執行中或已關閉, 用於安全關閉任務池）

type Pool struct {
    capacity       uint64
    runningWorkers uint64
    state          int64
    taskC          chan *Task
    closeC         chan bool
}

任務池的建構函式:

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,
        state:    RUNNING,
        // 初始化任務佇列, 佇列長度為容量
        taskC:    make(chan *Task, capacity),
        closeC:   make(chan bool),
    }, nil
}

啟動 worker

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

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

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

        for {
            select { // 阻塞等待任務、結束訊號到來
            case task, ok :=

上述程式碼中, 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 的操作也考慮資料競爭, 封裝 GetCap() 方法:


關鍵字：網路協議 Java Go

相關文章