golang-waitgroup

說明

golang通過waitgroup來實現併發控制,用法跟java的CountDownLatch 效果一樣

 

WaitGroup 的使用場景和方法

我們通過goroutine運行一個或者一組任務,需要關心這組任務執行完了進行通知

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

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

package main

import (
    "fmt"
    "sync"
    "time"
)

func worker(wg *sync.WaitGroup) {
    doSomething()
    wg.Done() // 2.1、完成任務
}

func main() {
    var wg sync.WaitGroup
    wg.Add(5) // 1、添加 5 個任務
    for i := 1; i <= 5; i++ {
        go worker(&wg) // 2、每個任務併發執行
    }
    wg.Wait() // 3、阻塞等待所有任務完成

}

WaitGroup 源碼分析

sync/waitgroup.go

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

noCopy
其中，noCopy 表示 WaitGroup 是不可複製的。那麼什麼叫不可複製呢？

舉個例子，當我們對函數參數定義了這個不可複製的類型時，開發者只能通過指針來傳遞函數參數。而規定使用指針傳遞又有什麼好處呢？

好處在於如果有多個函數都定義了這個不可複製的參數時，那麼這多個函數參數就可以共用同一個指針變量，來同步執行結果。而 WaitGroup 就是需要這樣的約束規定。

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

每當有 goroutine 調用了 Wait() 方法阻塞等待時，就會對 waiter 數量 + 1，然後等待信號量的喚起通知。

當我們調用 Add() 方法時，就會對 state1 的 counter 數量 + 1。

當調用 Done() 方法時就會對 counter 數量 -1。

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

ADD方法

func (wg *WaitGroup) Add(delta int) {
    statep, semap := wg.state()
    if race.Enabled { // 此處是 go 的競爭檢測，可以不用關心
        _ = *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 的競爭檢測，可以不用關心
        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：還有任務在執行；waiter == 0 表示沒有在阻塞等待的 goroutine
        return
    }
    if *statep != state {
        panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    }
    // 執行到此處相當於 countr = 0，即所有的任務都已執行完，需要喚起等待的 goroutine了
    *statep = 0
    for ; w != 0; w-- {
        runtime_Semrelease(semap, false, 0)
    }
}

Done

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

Wait方法

func (wg *WaitGroup) Wait() {
    statep, semap := wg.state()
    if race.Enabled { // go 的競爭檢測，可以不用關心
        _ = *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 數目 +1.
        if atomic.CompareAndSwapUint64(statep, state, state+1) {
            if race.Enabled && w == 0 {
                race.Write(unsafe.Pointer(semap)) // go 的競爭檢測，可以不用關心
            }
            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
        }
    }
}