1. 程式人生 > >Go基礎系列:Go實現工作池的兩種方式

Go基礎系列:Go實現工作池的兩種方式

worker pool簡介

worker pool其實就是執行緒池thread pool。對於go來說,直接使用的是goroutine而非執行緒,不過這裡仍然以執行緒來解釋執行緒池。

線上程池模型中,有2個佇列一個池子:任務佇列、已完成任務佇列和執行緒池。其中已完成任務佇列可能存在也可能不存在,依據實際需求而定。

只要有任務進來,就會放進任務佇列中。只要執行緒執行完了一個任務,就將任務放進已完成任務佇列,有時候還會將任務的處理結果也放進已完成佇列中。

worker pool中包含了一堆的執行緒(worker,對go而言每個worker就是一個goroutine),這些執行緒嗷嗷待哺,等待著為它們分配任務,或者自己去任務佇列中取任務。取得任務後更新任務佇列,然後執行任務,並將執行完成的任務放進已完成佇列。

下圖來自wiki:

在Go中有兩種方式可以實現工作池:傳統的互斥鎖、channel。

傳統互斥鎖機制的工作池

假設Go中的任務的定義形式為:

type Task struct {
    ...
}

每次有任務進來時,都將任務放在任務佇列中。

使用傳統的互斥鎖方式實現,任務佇列的定義結構大概如下:

type Queue struct{
    M     sync.Mutex
    Tasks []Task
}

然後在執行任務的函式中加上Lock()和Unlock()。例如:

func Worker(queue *Queue) {
    for {
        // Lock()和Unlock()之間的是critical section
        queue.M.Lock()
        // 取出任務
        task := queue.Tasks[0]
        // 更新任務佇列
        queue.Tasks = queue.Tasks[1:]
        queue.M.Unlock()
        // 在此goroutine中執行任務
        process(task)
    }
}

假如線上程池中激活了100個goroutine來執行Worker()。Lock()和Unlock()保證了在同一時間點只能有一個goroutine取得任務並隨之更新任務列表,取任務和更新任務佇列都是critical section中的程式碼,它們是具有原子性。然後這個goroutine可以執行自己取得的任務。於此同時,其它goroutine可以爭奪互斥鎖,只要爭搶到互斥鎖,就可以取得任務並更新任務列表。當某個goroutine執行完process(task),它將因為for迴圈再次參與互斥鎖的爭搶。

上面只是給出了一點主要的程式碼段,要實現完整的執行緒池,還有很多額外的程式碼。

通過互斥鎖,上面的一切操作都是執行緒安全的。但問題在於加鎖/解鎖的機制比較重量級,當worker(即goroutine)的數量足夠多,鎖機制的實現將出現瓶頸。

通過buffered channel實現工作池

在Go中,也能用buffered channel實現工作池。

示例程式碼很長,所以這裡先拆分解釋每一部分,最後給出完整的程式碼段。

在下面的示例中,每個worker的工作都是計算每個數值的位數相加之和。例如給定一個數值234,worker則計算2+3+4=9。這裡交給worker的數值是隨機生成的[0,999)範圍內的數值。

這個示例有幾個核心功能需要先解釋,也是通過channel實現執行緒池的一般功能:

  • 建立一個task buffered channel,並通過allocate()函式將生成的任務存放到task buffered channel中
  • 建立一個goroutine pool,每個goroutine監聽task buffered channel,並從中取出任務
  • goroutine執行任務後,將結果寫入到result buffered channel中
  • 從result buffered channel中取出計算結果並輸出

首先,建立Task和Result兩個結構,並建立它們的通道:

type Task struct {
    ID      int
    randnum int
}

type Result struct {
    task    Task
    result  int
}

var tasks = make(chan Task, 10)
var results = make(chan Result, 10)

這裡,每個Task都有自己的ID,以及該任務將要被worker計算的隨機數。每個Result都包含了worker的計算結果result以及這個結果對應的task,這樣從Result中就可以取出任務資訊以及計算結果。

另外,兩個通道都是buffered channel,容量都是10。每個worker都會監聽tasks通道,並取出其中的任務進行計算,然後將計算結果和任務自身放進results通道中。

然後是計算位數之和的函式process(),它將作為worker的工作任務之一。

func process(num int) int {
    sum := 0
    for num != 0 {
        digit := num % 10
        sum += digit
        num /= 10
    }
    time.Sleep(2 * time.Second)
    return sum
}

這個計算過程其實很簡單,但隨後還睡眠了2秒,用來假裝執行一個計算任務是需要一點時間的。

然後是worker(),它監聽tasks通道並取出任務進行計算,並將結果放進results通道。

func worker(wg *WaitGroup){
    defer wg.Done()
    for task := range tasks {
        result := Result{task, process(task.randnum)}
        results <- result
    }
}

上面的程式碼很容易理解,只要tasks channel不關閉,就會一直監聽該channel。需要注意的是,該函式使用指標型別的*WaitGroup作為引數,不能直接使用值型別的WaitGroup作為引數,這樣會使得每個worker都有一個自己的WaitGroup。

然後是建立工作池的函式createWorkerPool(),它有一個數值引數,表示要建立多少個worker。

func createWorkerPool(numOfWorkers int) {
    var wg sync.WaitGroup
    for i := 0; i < numOfWorkers; i++ {
        wg.Add(1)
        go worker(&wg)
    }
    wg.Wait()
    close(results)
}

建立工作池時,首先建立一個WaitGroup的值wg,這個wg被工作池中的所有goroutine共享,每建立一個goroutine都wg.Add(1)。建立完所有的goroutine後等待所有的groutine都執行完它們的任務,只要有一個任務還沒有執行完,這個函式就會被Wait()阻塞。當所有任務都執行完成後,關閉results通道,因為沒有結果再需要向該通道寫了。

當然,這裡是否需要關閉results通道,是由稍後的range迭代這個通道決定的,不關閉這個通道會一直阻塞range,最終導致死鎖。

工作池部分已經完成了。現在需要使用allocate()函式分配任務:生成一大堆的隨機數,然後將Task放進tasks通道。該函式有一個代表建立任務數量的數值引數:

func allocate(numOfTasks int) {
    for i := 0; i < numOfTasks; i++ {
        randnum := rand.Intn(999)
        task := Task{i, randnum}
        tasks <- task
    }
    close(tasks)
}

注意,最後需要關閉tasks通道,因為所有任務都分配完之後,沒有任務再需要分配。當然,這裡之所以需要關閉tasks通道,是因為worker()中使用了range迭代tasks通道,如果不關閉這個通道,worker將在取完所有任務後一直阻塞,最終導致死鎖。

再接著的是取出results通道中的結果進行輸出,函式名為getResult():

func getResult(done chan bool) {
    for result := range results {
        fmt.Printf("Task id %d, randnum %d , sum %d\n", result.task.id, result.task.randnum, result.result)
    }
    done <- true
}

getResult()中使用了一個done引數,這個引數是一個訊號通道,用來表示results中的所有結果都取出來並處理完成了,這個通道不一定要用bool型別,任何型別皆可,它不用來傳資料,僅用來返回可讀,所以上面直接close(done)的效果也一樣。通過下面的main()函式,就能理解done訊號通道的作用。

最後還差main()函式:

func main() {
    // 記錄起始終止時間,用來測試完成所有任務耗費時長
    startTime := time.Now()
    
    numOfWorkers := 20
    numOfTasks := 100
    // 建立任務到任務佇列中
    go allocate(numOfTasks)
    // 建立工作池
    go createWorkerPool(numOfWorkers)
    // 取得結果
    var done = make(chan bool)
    go getResult(done)

    // 如果results中還有資料,將阻塞在此
    // 直到傳送了訊號給done通道
    <- done
    endTime := time.Now()
    diff := endTime.Sub(startTime)
    fmt.Println("total time taken ", diff.Seconds(), "seconds")
}

上面分配了20個worker,這20個worker總共需要處理的任務數量為100。但注意,無論是tasks還是results通道,容量都是10,意味著任務佇列最長只能是10個任務。

下面是完整的程式碼段:

package main

import (
    "fmt"
    "math/rand"
    "sync"
    "time"
)

type Task struct {
    id      int
    randnum int
}
type Result struct {
    task   Task
    result int
}

var tasks = make(chan Task, 10)
var results = make(chan Result, 10)

func process(num int) int {
    sum := 0
    for num != 0 {
        digit := num % 10
        sum += digit
        num /= 10
    }
    time.Sleep(2 * time.Second)
    return sum
}
func worker(wg *sync.WaitGroup) {
    defer wg.Done()
    for task := range tasks {
        result := Result{task, process(task.randnum)}
        results <- result
    }
}
func createWorkerPool(numOfWorkers int) {
    var wg sync.WaitGroup
    for i := 0; i < numOfWorkers; i++ {
        wg.Add(1)
        go worker(&wg)
    }
    wg.Wait()
    close(results)
}
func allocate(numOfTasks int) {
    for i := 0; i < numOfTasks; i++ {
        randnum := rand.Intn(999)
        task := Task{i, randnum}
        tasks <- task
    }
    close(tasks)
}
func getResult(done chan bool) {
    for result := range results {
        fmt.Printf("Task id %d, randnum %d , sum %d\n", result.task.id, result.task.randnum, result.result)
    }
    done <- true
}
func main() {
    startTime := time.Now()
    numOfWorkers := 20
    numOfTasks := 100

    var done = make(chan bool)
    go getResult(done)
    go allocate(numOfTasks)
    go createWorkerPool(numOfWorkers)
    // 必須在allocate()和getResult()之後建立工作池
    <-done
    endTime := time.Now()
    diff := endTime.Sub(startTime)
    fmt.Println("total time taken ", diff.Seconds(), "seconds")
}

執行結果:

Task id 19, randnum 914 , sum 14
Task id 9, randnum 150 , sum 6
Task id 15, randnum 215 , sum 8
............
Task id 97, randnum 315 , sum 9
Task id 99, randnum 641 , sum 11
total time taken  10.0174705 seconds

總共花費10秒。

可以試著將任務數量、worker數量修改修改,看看它們的效能比例情況。例如,將worker數量設定為99,將需要4秒,將worker數量設定為10,將需要20秒。