?在 Go 中，sync 包下的 WaitGroup 能有助于我們控制協(xié)程之間的同步。當需要等待一組協(xié)程都執(zhí)行完各自任務后，才能繼續(xù)后續(xù)邏輯。這種場景，就非常適合使用它。但是，在使用 WaitGroup 的過程中，你可能會犯錯誤，下文我們將通過示例逐步探討。
任務示例
初始任務
假設我們有以下任務 woker，它執(zhí)行的任務是將參數 msg 打印出來。
func worker(msg string) {
fmt.Printf("worker do %s\n", msg)
}
func main() {
worker("task 1")
fmt.Println("main exit")
}執(zhí)行結果如下
worker do task 1
main exit
更多任務
如果有更多的任務需要處理
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í)行的結果
worker do task 1
worker do task 2
worker do task 3
main exit
并發(fā)執(zhí)行
依次執(zhí)行可以完成所有任務，但由于任務間沒有依賴性，并發(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")
}但這樣，我們大概率得到這樣的結果
使用 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")
}
執(zhí)行結果可能
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
還有可能
雖然main exit總會在最后打印輸出，但并發(fā)任務未均如愿得到執(zhí)行。
全局變量改為傳參
也許是我們不應該將 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")
}但執(zhí)行結果顯然更不對了
值傳遞改為指針傳遞
如果去查看 WaitGroup 的這三個 API 函數，你會發(fā)現它們的方法接收者都是指針。
我們使用值傳遞 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
看著好像符合預期了，但是如果多次執(zhí)行，你發(fā)現可能會得到這樣的結果。
worker do task 2
waiting
worker do task 1
worker do task 3
main exit
或者這樣
竟然還有問題？！
執(zhí)行順序
其實問題出在了執(zhí)行順序。
注意，wg.Add(1)?我們是在 worker 函數中執(zhí)行，而不是在調用方（main?函數）。通過 Go 關鍵字讓一個 gotoutine 執(zhí)行起來存在一小段的滯后時間。而這就會存在問題：當程序執(zhí)行到了wg.Wait()?時，前面的 3 個goroutine 并不一定都啟動起來了，即它們不一定來得及調用wg.Add(1)。（這個 goroutine 滯后的問題其實也是上文并發(fā)執(zhí)行未能得到預期結果的原因所在。）
例如最后一個結果，每個 worker 都還來不及執(zhí)行wg.Add(1)?，main 函數就已經執(zhí)行到wg.Wait()，此時它發(fā)現任務計數是0，所以就直接非阻塞執(zhí)行后續(xù) main 函數邏輯了。
對于這個問題，我們的解決方案是：
- 在 main 函數調用worker前就應該執(zhí)行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")
}這樣，無論執(zhí)行多少次，結果都能符合預期要求。
waiting
worker do task 3
worker do task 2
worker do task 1
main exit
事實上，上述寫法不夠簡潔。當大量相同子任務通過 goroutine 執(zhí)行時，我們應該采用 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 之前執(zhí)行wg.Add(1)?，goroutine 中的第一行代碼為defer wg.Done()。
這樣，我們能讓調用方（例子中的main函數）有效地控制任務數，同時既避免了傳遞 WaitGroup 的風險，又能讓子任務YourFunction()只關心自身邏輯。
從本文的例子可以看出，在并發(fā)編程時，一定要采用正確的使用姿勢，否則很容易產生讓人困惑的問題。?
