你是否因使用姿勢不當，而在 WaitGroup 栽了跟頭？

在 Go 中，sync 包下的 WaitGroup 能有助於我們控制協程之間的同步。當需要等待一組協程都執行完各自任務後，才能繼續後續邏輯。這種場景，就非常適合使用它。但是，在使用 WaitGroup 的過程中，你可能會犯錯誤，下文我們將通過示例逐步探討。

任務示例

初始任務

假設我們有以下任務 woker，它執行的任務是將參數 msg 打印出來。

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

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

執行結果如下

worker do task 1
main exit

併發執行

依次執行可以完成所有任務，但由於任務間沒有依賴性，併發執行是更好的選擇。

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")
}

但這樣，我們大概率得到這樣的結果

main exit

使用 WaitGroup

WaitGroup 提供三個 API。

Add(delta int) 函數提供了 WaitGroup 的任務計數，delta 的值可以爲正也可以爲負，通常在添加任務時使用。
Done() 函數其實就是 Add(-1)，在任務完成時調用。
Wait() 函數用於阻塞等待 WaitGroup 的任務們均完成，即阻塞等待至任務數爲 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")
}

執行結果可能

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總會在最後打印輸出，但併發任務未均如願得到執行。

全局變量改爲傳參

也許是我們不應該將 wg 設爲全局變量？那改爲函數傳參試試。

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
main exit

值傳遞改爲指針傳遞

如果去查看 WaitGroup 的這三個 API 函數，你會發現它們的方法接收者都是指針。

我們使用值傳遞 WaitGroup，那就意味着在函數中使用的 wg 是一個複製對象。而 WaitGroup 的定義描述中有提及：使用過程中它不能被複制（詳細原因可以查看菜刀歷史文章 no copy 機制）。

因此，我們需要將 WaitGroup 的參數類型改爲指針。

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

看着好像符合預期了，但是如果多次執行，你發現可能會得到這樣的結果。

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

或者這樣

waiting
main exit

竟然還有問題？！

執行順序

其實問題出在了執行順序。

注意，wg.Add(1)我們是在 worker 函數中執行，而不是在調用方（main函數）。通過 Go 關鍵字讓一個 gotoutine 執行起來存在一小段的滯後時間。而這就會存在問題：當程序執行到了wg.Wait()時，前面的 3 個 goroutine 並不一定都啓動起來了，即它們不一定來得及調用wg.Add(1)。（這個 goroutine 滯後的問題其實也是上文併發執行未能得到預期結果的原因所在。）

例如最後一個結果，每個 worker 都還來不及執行wg.Add(1)，main 函數就已經執行到wg.Wait()，此時它發現任務計數是 0，所以就直接非阻塞執行後續 main 函數邏輯了。

對於這個問題，我們的解決方案是：

在 main 函數調用worker前就應該執行wg.Add(1)來給任務準確計數；
避免潛在複製風險，不再傳遞 WaitGroup 參數；
將wg.Done()從worker中移出，與wg.Add()調用形成對應。

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")
}

這樣，無論執行多少次，結果都能符合預期要求。

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

事實上，上述寫法不夠簡潔。當大量相同子任務通過 goroutine 執行時，我們應該採用 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")
}

總結

我們可以將 WaitGroup 的核心使用姿勢總結爲如下模版

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

在進入 goroutine 之前執行wg.Add(1)，goroutine 中的第一行代碼爲defer wg.Done()。

這樣，我們能讓調用方（例子中的 main 函數）有效地控制任務數，同時既避免了傳遞 WaitGroup 的風險，又能讓子任務YourFunction()只關心自身邏輯。

從本文的例子可以看出，在併發編程時，一定要採用正確的使用姿勢，否則很容易產生讓人困惑的問題。最後，如果讀者想進一步學習 waitGroup，推薦閱讀歷史文章 Go 中看似簡單的 WaitGroup 源碼設計，竟然暗含這麼多知識？

參考文章：https://leangaurav.medium.com/common-mistakes-when-using-golangs-sync-waitgroup-88188556ca54

本文由 Readfog 進行 AMP 轉碼，版權歸原作者所有。
來源：https://mp.weixin.qq.com/s/swLsBgvOCParmyLZcUmjEw