Go語言的WaitGroup怎么使用

本篇內容介紹了“Go語言的WaitGroup怎么使用”的有關知識，在實際案例的操作過程中，不少人都會遇到這樣的困境，接下來就讓小編帶領大家學習一下如何處理這些情況吧！希望大家仔細閱讀，能夠學有所成！

小試牛刀

我們先來個簡單的例子，看下 WaitGroup 是怎么使用的。示例中使用 Add(5) 表示我們有 5個 子任務，然后起了 5個 協程去完成任務，主協程使用 Wait() 方法等待 子協程執行完畢，輸出一共等待的時間。

func main() {
    var waitGroup sync.WaitGroup

    start := time.Now()
    waitGroup.Add(5)
    for i := 0; i < 5; i++ {
        go func() {
            defer waitGroup.Done()
            time.Sleep(time.Second)
            fmt.Println("done")
        }()
    }

    waitGroup.Wait()
    fmt.Println(time.Now().Sub(start).Seconds())
}

/*
done
done
done
done
done
1.000306089
*/

總覽

WaitGroup 一共有三個方法：

(wg *WaitGroup) Add(delta int)
(wg *WaitGroup) Done()
(wg *WaitGroup) Wait()

• Add 方法用于設置 WaitGroup 的計數值，可以理解為子任務的數量

• Done 方法用于將 WaitGroup 的計數值減一，可以理解為完成一個子任務

• Wait 方法用于阻塞調用者，直到 WaitGroup 的計數值為0，即所有子任務都完成

正常來說，我們使用的時候，需要先確定子任務的數量，然后調用 Add() 方法傳入相應的數量，在每個子任務的協程中，調用 Done()，需要等待的協程調用 Wait() 方法，狀態流轉如下圖：

底層實現

結構體

type WaitGroup struct {
    noCopy noCopy // noCopy 字段標識，由于 WaitGroup 不能復制，方便工具檢測

    state1 [3]uint32  // 12個字節，8個字節標識 計數值和等待數量，4個字節用于標識信號量
}

state1 是個復合字段，會拆分為兩部分： 64位（8個字節）的 statep 作為一個整體用于原子操作, 其中前面4個字節表示計數值，后面四個字節表示等待數量；剩余 32位（4個字節）semap 用于標識信號量。

Go語言中對于64位的變量進行原子操作，需要保證該變量是64位對齊的，也就是要保證這 8個字節 的首地址是 8 的整數倍。因此當 state1 的首地址是 8 的整數倍時，取前8個字節作為 statep ，后4個字節作為 semap；當 state1 的首地址不是 8 的整數倍時，取后8個字節作為 statep ，前4個字節作為 semap。

func (wg *WaitGroup) state() (statep *uint64, semap *uint32) {
    
    // 首地址是8的倍數時，前8個字節為 statep, 后四個字節為 semap
    if uintptr(unsafe.Pointer(&wg.state1))%8 == 0 {
        return (*uint64)(unsafe.Pointer(&wg.state1)), &wg.state1[2]
    } else { 
        
    // 后8個字節為 statep, 前四個字節為 semap    
        return (*uint64)(unsafe.Pointer(&wg.state1[1])), &wg.state1[0]
    }
}

Add

• Add 方法用于添加一個計數值（負數相當于減），當計數值變為0后， Wait 方法阻塞的所有等待者都會被釋放

• 計數值變為負數是非法操作，產生 panic

• 當計數值為0時（初始狀態），Add 方法不能和 Wait 方法并發調用，需要保證 Add 方法在 Wait 方法之前調用，否則會 panic

func (wg *WaitGroup) Add(delta int) {

    // 拿到計數值等待者變量 statep 和 信號量 semap
    statep, semap := wg.state()

    // 計數值加上 delta: statep 的前四個字節是計數值，因此將 delta 前移 32位
    state := atomic.AddUint64(statep, uint64(delta)<<32)

    // 計數值
    v := int32(state >> 32)

    // 等待者數量
    w := uint32(state)

    // 如果加上 delta 之后，計數值變為負數，不合法，panic
    if v < 0 {
        panic("sync: negative WaitGroup counter")
    }

    // delta > 0 && v == int32(delta) : 表示從 0 開始添加計數值
    // w!=0 ：表示已經有了等待者
    // 說明在添加計數值的時候，同時添加了等待者，非法操作。添加等待者需要在添加計數值之后
    if w != 0 && delta > 0 && v == int32(delta) {
        panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    }

    // v>0 : 計數值不等于0，不需要喚醒等待者，直接返回
    // w==0: 沒有等待者，不需要喚醒，直接返回
    if v > 0 || w == 0 {
        return
    }

    // 再次檢查數據是否一致
    if *statep != state {
        panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    }

    // 到這里說明計數值為0，且等待者大于0，需要喚醒所有的等待者，并把系統置為初始狀態（0狀態）
  
  // 將計數值和等待者數量都置為0
    *statep = 0

    // 喚醒等待者
    for ; w != 0; w-- {
        runtime_Semrelease(semap, false, 0)
    }
}

Done

// 完成一個任務，將計數值減一，當計數值減為0時，需要喚醒所有的等待者
func (wg *WaitGroup) Done() {
    wg.Add(-1)
}

Wait

// 調用 Wait 方法會被阻塞，直到 計數值 變為0
func (wg *WaitGroup) Wait() {

    // 獲取計數、等待數和信號量
    statep, semap := wg.state()

    for {
        state := atomic.LoadUint64(statep)

        // 計數值
        v := int32(state >> 32)

        // 等待者數量
        w := uint32(state)

        // 計數值數量為0，直接返回，無需等待
        if v == 0 {
            return
        }

        // 到這里說明計數值數量大于0
        // 增加等待者數量：這里會有競爭，比如多個 Wait 調用，或者在同時調用 Add 方法，增加不成功會繼續 for 循環
        if atomic.CompareAndSwapUint64(statep, state, state+1) {
            // 增加成功后，阻塞在信號量這里，等待被喚醒
            runtime_Semacquire(semap)

            // 被喚醒的時候，應該是0狀態。如果重用 WaitGroup，需要等 Wait 返回
            if *statep != 0 {
                panic("sync: WaitGroup is reused before previous Wait has returned")
            }
            return
        }
    }
}

易錯點

上面分析源碼可以看到幾個會產生 panic 的點，這也是我們使用 WaitGroup 需要注意的地方

1.計數值變為負數

調用 Add 時參數值傳負數

func main() {
var wg sync.WaitGroup
wg.Add(1)
wg.Add(-1)
wg.Add(-1)
}

多次調用 Done 方法

func main() {
    var wg sync.WaitGroup

    wg.Add(1)
    go func() {
        fmt.Println("test")
        wg.Done()
        wg.Done()
    }()

    time.Sleep(time.Second)
    wg.Wait()
}

2.Add 和 Wait 并發調用

Add 和 Wait 并發調用，有可能達不到我們預期的效果，甚至 panic。如下示例中，我們想要等待 3 個子任務都執行完后再執行主任務，但實際情況可能是子任務還沒起來，主任務就繼續往下執行了。

func doSomething(wg *sync.WaitGroup) {
    wg.Add(1)
    fmt.Println("do  something")
    defer wg.Done()
}

func main() {
    var wg sync.WaitGroup
    for i := 0; i < 3; i++ {
        go doSomething(&wg)
    }

    wg.Wait()
    fmt.Println("main")
}

//main
//do  something
//do  something

正確的使用方式，應該是在調用 Wait 前先調用 Add

func doSomething(wg *sync.WaitGroup) {
    defer wg.Done()
    fmt.Println("do  something")
}

func main() {
    var wg sync.WaitGroup
    wg.Add(3)
    for i := 0; i < 3; i++ {
        go doSomething(&wg)
    }

    wg.Wait()
    fmt.Println("main")
}

//do  something
//do  something
//do  something
//main

3.沒有等 Wait 返回，就重用 WaitGroup

func main() {
    var wg sync.WaitGroup

    wg.Add(1)
    go func() {
        fmt.Println("do something")
        wg.Done()
        wg.Add(1)
    }()

    wg.Wait()
}

4.復制使用

我們知道 Go 語言中的參數傳遞，都是值傳遞，就會產生復制操作。因此在向函數傳遞 WaitGroup 時，使用指針進行操作。

// 錯誤使用方式，沒有使用指針
func doSomething(wg sync.WaitGroup) {
    fmt.Println("do  something")
    defer wg.Done()
}

func main() {
    var wg sync.WaitGroup
    wg.Add(3)
    for i := 0; i < 3; i++ {
        // 這里沒使用指針，wg狀態一直不會改變，導致 Wait 一直阻塞
        go doSomething(wg)
    }

    wg.Wait()
    fmt.Println("main")
}

“Go語言的WaitGroup怎么使用”的內容就介紹到這里了，感謝大家的閱讀。如果想了解更多行業相關的知識可以關注億速云網站，小編將為大家輸出更多高質量的實用文章！