golang?waitgroup輔助并發(fā)控制使用場景和方法解析

更新時間：2023年09月27日 10:16:14 作者：lincoln_hlf1

Golang?提供了簡潔的?go?關(guān)鍵字來讓開發(fā)者更容易的進行并發(fā)編程,同時也提供了?WaitGroup?對象來輔助并發(fā)控制,今天我們就來分析下?WaitGroup?的使用方法,順便瞧一瞧它的底層源碼

WaitGroup 的使用場景和方法

當(dāng)我們有很多任務(wù)要同時進行時，如果并不需要關(guān)心各個任務(wù)的執(zhí)行進度，那直接使用 go 關(guān)鍵字即可。

如果我們需要關(guān)心所有任務(wù)完成后才能往下運行時，則需要 WaitGroup 來阻塞等待這些并發(fā)任務(wù)了。

WaitGroup 如同它的字面意思，就是等待一組 goroutine 運行完成，主要有三個方法組成：

Add(delta int) ：添加任務(wù)數(shù)
Wait()：阻塞等待所有任務(wù)的完成
Done()：完成任務(wù)

下面是它們的具體用法，具體的作用都在注釋上：

package main
import (
    "fmt"
    "sync"
    "time"
)
func worker(wg *sync.WaitGroup) {
    doSomething()
    wg.Done() // 2.1、完成任務(wù)
}
func main() {
    var wg sync.WaitGroup
    wg.Add(5) // 1、添加 5 個任務(wù)
    for i := 1; i <= 5; i++ {
        go worker(&wg) // 2、每個任務(wù)并發(fā)執(zhí)行
    }
    wg.Wait() // 3、阻塞等待所有任務(wù)完成
}

WaitGroup 源碼分析

上面 WaitGroup 的使用很簡單，接下來我們到 src/sync/waitgroup.go 里分析下它的源碼。首先，是 WaitGroup 的結(jié)構(gòu)體：

type WaitGroup struct {
    noCopy noCopy
    state1 [3]uint32
}

noCopy

其中，noCopy 表示 WaitGroup 是不可復(fù)制的。那么什么叫不可復(fù)制呢？

舉個例子，當(dāng)我們對函數(shù)參數(shù)定義了這個不可復(fù)制的類型時，開發(fā)者只能通過指針來傳遞函數(shù)參數(shù)。而規(guī)定使用指針傳遞又有什么好處呢？

好處在于如果有多個函數(shù)都定義了這個不可復(fù)制的參數(shù)時，那么這多個函數(shù)參數(shù)就可以共用同一個指針變量，來同步執(zhí)行結(jié)果。而 WaitGroup 就是需要這樣的約束規(guī)定。

state1 字段

接下來我們來看看 WaitGroup 的 state1 字段。state1 是一個包含了 counter 總數(shù)、 waiter 等待數(shù)、sema 信號量的 uint32 數(shù)組。

每當(dāng)有 goroutine 調(diào)用了 Wait() 方法阻塞等待時，就會對 waiter 數(shù)量 + 1，然后等待信號量的喚起通知。

當(dāng)我們調(diào)用 Add() 方法時，就會對 state1 的 counter 數(shù)量 + 1。

當(dāng)調(diào)用 Done() 方法時就會對 counter 數(shù)量 -1。

直到 counter == 0 時就可以通過信號量喚起對應(yīng) waiter 數(shù)量的 goroutine 了，也就是喚起剛剛阻塞等待的 goroutine 們。

關(guān)于信號量的解釋，可以參考下 golang 重要知識：mutex 里的相關(guān)介紹：

PV 原語解釋：
通過操作信號量 S 來處理進程間的同步與互斥的問題。
S>0：表示有 S 個資源可用；S=0 表示無資源可用；S<0 絕對值表示等待隊列或鏈表中的進程個數(shù)。信號量 S 的初值應(yīng)大于等于 0。
P 原語：表示申請一個資源，對 S 原子性的減 1，若減 1 后仍 S>=0，則該進程繼續(xù)執(zhí)行；若減 1 后 S<0，表示已無資源可用，需要將自己阻塞起來，放到等待隊列上。
V 原語：表示釋放一個資源，對 S 原子性的加 1；若加 1 后 S>0，則該進程繼續(xù)執(zhí)行；若加 1 后 S<=0，表示等待隊列上有等待進程，需要將第一個等待的進程喚醒。

此處操作系統(tǒng)可以理解為 Go 的運行時 runtime，進程可以理解為協(xié)程。

方法解釋

最后，我們來深入 WaitGroup 的三個方法，進行源碼分析。大家感興趣的可以繼續(xù)往下看，主要是對源碼的分析注釋。

Add(delta int) 方法

func (wg *WaitGroup) Add(delta int) {
    statep, semap := wg.state()
    if race.Enabled { // 此處是 go 的競爭檢測，可以不用關(guān)心
        _ = *statep
        if delta < 0 {
            race.ReleaseMerge(unsafe.Pointer(wg))
        }
        race.Disable()
        defer race.Enable()
    }
    state := atomic.AddUint64(statep, uint64(delta)<<32)
    v := int32(state >> 32) // 獲取 counter
    w := uint32(state) // 獲取 waiter
    if race.Enabled && delta > 0 && v == int32(delta) { // go 的競爭檢測，可以不用關(guān)心
        race.Read(unsafe.Pointer(semap))
    }
    if v < 0 {
        panic("sync: negative WaitGroup counter")
    }
    if w != 0 && delta > 0 && v == int32(delta) {
        panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    }
    if v > 0 || w == 0 { // counter > 0：還有任務(wù)在執(zhí)行；waiter == 0 表示沒有在阻塞等待的 goroutine
        return
    }
    if *statep != state {
        panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    }
    // 執(zhí)行到此處相當(dāng)于 countr = 0，即所有的任務(wù)都已執(zhí)行完，需要喚起等待的 goroutine了
    *statep = 0
    for ; w != 0; w-- {
        runtime_Semrelease(semap, false, 0)
    }
}

Done 方法

func (wg *WaitGroup) Done() {
    wg.Add(-1) // 直接調(diào)用 Add 方法 對 counter -1
}

Wait 方法

func (wg *WaitGroup) Wait() {
    statep, semap := wg.state()
    if race.Enabled { // go 的競爭檢測，可以不用關(guān)心
        _ = *statep
        race.Disable()
    }
    for {
        state := atomic.LoadUint64(statep)
        v := int32(state >> 32)
        w := uint32(state)
        if v == 0 {
            // counter 為 0， 不需要再等待了。
            if race.Enabled {
                race.Enable()
                race.Acquire(unsafe.Pointer(wg))
            }
            return
        }
        // waiters 數(shù)目 +1.
        if atomic.CompareAndSwapUint64(statep, state, state+1) {
            if race.Enabled && w == 0 {
                race.Write(unsafe.Pointer(semap)) // go 的競爭檢測，可以不用關(guān)心
            }
            runtime_Semacquire(semap) // 阻塞等待喚起
            if *statep != 0 {
                panic("sync: WaitGroup is reused before previous Wait has returned")
            }
            if race.Enabled {
                race.Enable()
                race.Acquire(unsafe.Pointer(wg))
            }
            return
        }
    }
}

從這幾個方法的源碼，我們可以看出，Go 并沒有使用 mutex 等鎖去做字段值修改，而是采用了 atomic 原子操作來進行修改的。這是在底層硬件上支持的，所以性能更好。

總結(jié)

WaitGroup 比較簡單，就是一些計數(shù)值的維護和 goroutine 的阻塞喚起。它的運用也簡單，Add、Done、Wait 這三個方法經(jīng)常是同時出現(xiàn)的。相信大伙深入到源碼也能瞧出個大概，更多關(guān)于golang waitgroup并發(fā)控制的資料請關(guān)注腳本之家其它相關(guān)文章！

您可能感興趣的文章: