?在 Go 中，sync 包下的 WaitGroup 能有助于我們控制協(xié)程之間的同步。當(dāng)需要等待一組協(xié)程都執(zhí)行完各自任務(wù)后，才能繼續(xù)后續(xù)邏輯。這種場景，就非常適合使用它。但是，在使用 WaitGroup 的過程中，你可能會犯錯(cuò)誤，下文我們將通過示例逐步探討。

任務(wù)示例

初始任務(wù)

假設(shè)我們有以下任務(wù) woker，它執(zhí)行的任務(wù)是將參數(shù) msg 打印出來。

func worker(msg string) {
 fmt.Printf("worker do %s\n", msg)
}

func main() {
 worker("task 1")
 fmt.Println("main exit")
}

執(zhí)行結(jié)果如下

worker do task 1
main exit

更多任務(wù)

如果有更多的任務(wù)需要處理

func worker(msg string) {
 fmt.Printf("worker do %s\n", msg)
}

func main() {
 worker("task 1")
 worker("task 2")
 worker("task 3")
 fmt.Println("main exit")
}

它們依次執(zhí)行的結(jié)果

worker do task 1
worker do task 2
worker do task 3
main exit

并發(fā)執(zhí)行

依次執(zhí)行可以完成所有任務(wù)，但由于任務(wù)間沒有依賴性，并發(fā)執(zhí)行是更好的選擇。

func worker(msg string) {
 fmt.Printf("worker do %s\n", msg)
}

func main() {
 go worker("task 1")
 go worker("task 2")
 go worker("task 3")
 fmt.Println("main exit")
}

但這樣，我們大概率得到這樣的結(jié)果

main exit

使用 WaitGroup

WaitGroup 提供三個(gè) API。

• Add(delta int) 函數(shù)提供了 WaitGroup 的任務(wù)計(jì)數(shù)，delta 的值可以為正也可以為負(fù)，通常在添加任務(wù)時(shí)使用。
• Done() 函數(shù)其實(shí)就是 Add(-1)，在任務(wù)完成時(shí)調(diào)用。
• Wait() 函數(shù)用于阻塞等待 WaitGroup 的任務(wù)們均完成，即阻塞等待至任務(wù)數(shù)為 0。

我們將代碼改寫如下

var wg sync.WaitGroup

func worker(msg string) {
 wg.Add(1)
 defer wg.Done()
 fmt.Printf("worker do %s\n", msg)
}

func main() {
 go worker("task 1")
 go worker("task 2")
 go worker("task 3")
  fmt.Println("waiting")
 wg.Wait()
 fmt.Println("main exit")
}

執(zhí)行結(jié)果可能

waiting
worker do task 1
worker do task 3
worker do task 2
main exit

同樣也可能

waiting
worker do task 2
worker do task 1
main exit

還有可能

waiting
main exit

雖然main exit總會在最后打印輸出，但并發(fā)任務(wù)未均如愿得到執(zhí)行。

全局變量改為傳參

也許是我們不應(yīng)該將 wg 設(shè)為全局變量？那改為函數(shù)傳參試試。

func worker(msg string, wg sync.WaitGroup) {
 wg.Add(1)
 defer wg.Done()
 fmt.Printf("worker do %s\n", msg)
}

func main() {
 var wg sync.WaitGroup
 go worker("task 1", wg)
 go worker("task 2", wg)
 go worker("task 3", wg)
 fmt.Println("waiting")
 wg.Wait()
 fmt.Println("main exit")
}

但執(zhí)行結(jié)果顯然更不對了

waiting
main exit

值傳遞改為指針傳遞

如果去查看 WaitGroup 的這三個(gè) API 函數(shù)，你會發(fā)現(xiàn)它們的方法接收者都是指針。

我們使用值傳遞 WaitGroup，那就意味著在函數(shù)中使用的 wg 是一個(gè)復(fù)制對象。而 WaitGroup 的定義描述中有提及：使用過程中它不能被復(fù)制（詳細(xì)原因可以查看菜刀歷史文章no copy 機(jī)制）。

因此，我們需要將 WaitGroup 的參數(shù)類型改為指針。

func worker(msg string, wg *sync.WaitGroup) {
 wg.Add(1)
 defer wg.Done()
 fmt.Printf("worker do %s\n", msg)
}

func main() {
 var wg sync.WaitGroup
 go worker("task 1", &wg)
 go worker("task 2", &wg)
 go worker("task 3", &wg)
 fmt.Println("waiting")
 wg.Wait()
 fmt.Println("main exit")
}

那這樣是不是就可以了呢？

waiting
worker do task 3
worker do task 2
worker do task 1
main exit

看著好像符合預(yù)期了，但是如果多次執(zhí)行，你發(fā)現(xiàn)可能會得到這樣的結(jié)果。

worker do task 2
waiting
worker do task 1
worker do task 3
main exit

或者這樣

waiting
main exit

竟然還有問題？！

執(zhí)行順序

其實(shí)問題出在了執(zhí)行順序。

注意，wg.Add(1)?我們是在 worker 函數(shù)中執(zhí)行，而不是在調(diào)用方（main?函數(shù)）。通過 Go 關(guān)鍵字讓一個(gè) gotoutine 執(zhí)行起來存在一小段的滯后時(shí)間。而這就會存在問題：當(dāng)程序執(zhí)行到了wg.Wait()?時(shí)，前面的 3 個(gè)goroutine 并不一定都啟動起來了，即它們不一定來得及調(diào)用wg.Add(1)。（這個(gè) goroutine 滯后的問題其實(shí)也是上文并發(fā)執(zhí)行未能得到預(yù)期結(jié)果的原因所在。）

例如最后一個(gè)結(jié)果，每個(gè) worker 都還來不及執(zhí)行wg.Add(1)?，main 函數(shù)就已經(jīng)執(zhí)行到wg.Wait()，此時(shí)它發(fā)現(xiàn)任務(wù)計(jì)數(shù)是0，所以就直接非阻塞執(zhí)行后續(xù) main 函數(shù)邏輯了。

對于這個(gè)問題，我們的解決方案是：

• 在 main 函數(shù)調(diào)用worker前就應(yīng)該執(zhí)行wg.Add(1)來給任務(wù)準(zhǔn)確計(jì)數(shù)；
• 避免潛在復(fù)制風(fēng)險(xiǎn)，不再傳遞 WaitGroup 參數(shù)；
• 將wg.Done()從worker中移出，與wg.Add()調(diào)用形成對應(yīng)。

func worker(msg string) {
 fmt.Printf("worker do %s\n", msg)
}

func main() {
 var wg sync.WaitGroup

 wg.Add(1)
 go func() {
  defer wg.Done()
  worker("task 1")
 }()

 wg.Add(1)
 go func() {
  defer wg.Done()
  worker("task 2")
 }()

 wg.Add(1)
 go func() {
  defer wg.Done()
  worker("task 3")
 }()

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

這樣，無論執(zhí)行多少次，結(jié)果都能符合預(yù)期要求。

waiting
worker do task 3
worker do task 2
worker do task 1
main exit

事實(shí)上，上述寫法不夠簡潔。當(dāng)大量相同子任務(wù)通過 goroutine 執(zhí)行時(shí)，我們應(yīng)該采用 for 語句來編寫代碼。

func worker(msg string) {
 fmt.Printf("worker do %s\n", msg)
}

func main() {
 var wg sync.WaitGroup

 for i := 0; i < 3; i++ {
  wg.Add(1)
  go func(i int) {
   defer wg.Done()
   worker(fmt.Sprintf("task %d", i+1))
  }(i)
 }
 fmt.Println("waiting")
 wg.Wait()
 fmt.Println("main exit")
}

總結(jié)

我們可以將 WaitGroup 的核心使用姿勢總結(jié)為如下模版

wg.Add(1)
 go func() {
  defer wg.Done()
  YourFunction()
 }()

在進(jìn)入 goroutine 之前執(zhí)行wg.Add(1)?，goroutine 中的第一行代碼為defer wg.Done()。

這樣，我們能讓調(diào)用方（例子中的main函數(shù)）有效地控制任務(wù)數(shù)，同時(shí)既避免了傳遞 WaitGroup 的風(fēng)險(xiǎn)，又能讓子任務(wù)YourFunction()只關(guān)心自身邏輯。

從本文的例子可以看出，在并發(fā)編程時(shí)，一定要采用正確的使用姿勢，否則很容易產(chǎn)生讓人困惑的問題。?