最清晰易懂的 Go WaitGroup 原始碼剖析
本篇主要介紹 WaitGroup 的一些特性,讓我們從本質上去了解 WaitGroup。關於 WaitGroup 的基本用法這裡就不做過多介紹了。相對於《這可能是最容易理解的 Go Mutex 原始碼剖析》來說,WaitGroup 就簡單的太多了。
原始碼剖析
Add()
Wait()
type WaitGroup struct {
noCopy noCopy
state1 [3]uint32
}
WaitGroup 底層結構看起來簡單,但 WaitGroup.state1 其實代表三個欄位:counter,waiter,sema。
- counter :可以理解為一個計數器,計算經過 wg.Add(N), wg.Done() 後的值。
- waiter :當前等待 WaitGroup 任務結束的等待者數量。其實就是呼叫 wg.Wait() 的次數,所以通常這個值是 1 。
- sema : 訊號量,用來喚醒 Wait() 函式。
為什麼要將 counter 和 waiter 放在一起 ?
其實是為了保證 WaitGroup 狀態的完整性。舉個例子,看下面的一段原始碼
// sync/waitgroup.go:L79 --> Add()
if v > 0 || w == 0 { // v => counter, w => waiter
return
}
// ...
*statep = 0
for ; w != 0; w-- {
runtime_Semrelease(semap, false, 0)
}
當同時發現 wg.counter <= 0 && wg.waiter != 0 時,才會去喚醒等待的 waiters,讓等待的協程繼續執行。但是使用 WaitGroup 的呼叫方一般都是併發操作,如果不同時獲取的 counter 和 waiter 的話,就會造成獲取到的 counter 和 waiter 可能不匹配,造成程式 deadlock 或者程式提前結束等待。
如何獲取 counter 和 waiter ?
對於 wg.state 的狀態變更,WaitGroup 的 Add(),Wait() 是使用 atomic 來做原子計算的(為了避免鎖競爭)。但是由於 atomic 需要使用者保證其 64 位對齊,所以將 counter 和 waiter 都設定成 uint32,同時作為一個變數,即滿足了 atomic 的要求,同時也保證了獲取 waiter 和 counter 的狀態完整性。但這也就導致了 32位,64位機器上獲取 state 的方式並不相同。如下圖:
簡單解釋下:
因為 64 位機器上本身就能保證 64 位對齊,所以按照 64 位對齊來取資料,拿到 state1[0], state1[1] 本身就是64 位對齊的。但是 32 位機器上並不能保證 64 位對齊,因為 32 位機器是 4 位元組對齊,如果也按照 64 位機器取 state[0],state[1] 就有可能會造成 atmoic 的使用錯誤。
於是 32 位機器上空出第一個 32 位,也就使後面 64 位天然滿足 64 位對齊,第一個 32 位放入 sema 剛好合適。早期 WaitGroup 的實現 sema 是和 state1 分開的,也就造成了使用 WaitGroup 就會造成 4 個位元組浪費,不過 go1.11 之後就是現在的結構了。
為什麼流程圖裡缺少了 Done ?
其實並不是,是因為 Done 的實現就是 Add. 只不過我們常規用法 wg.Add(1) 是加 1 ,wg.Done() 是減 1,即 wg.Done() 可以用 wg.Add(-1) 來代替。 儘管我們知道 wg.Add 可以傳遞負數當 wg.Done 使用,但是還是別這麼用。
退出waitgroup的條件
其實就一個條件, WaitGroup.counter 等於 0
日常開發中特殊需求
1. 控制超時/錯誤控制
雖說 WaitGroup 能夠讓主 Goroutine 等待子 Goroutine 退出,但是 WaitGroup 遇到一些特殊的需求,如:超時,錯誤控制,並不能很好的滿足,需要做一些特殊的處理。
使用者在電商平臺中購買某個貨物,為了計算使用者能優惠的金額,需要去獲取 A 系統(權益系統),B 系統(角色系統),C 系統(商品系統),D 系統(xx系統)。為了提高程式效能,可能會同時發起多個 Goroutine 去訪問這些系統,必然會使用 WaitGroup 等待資料的返回,但是存在一些問題:
- 當某個系統發生錯誤,等待的 Goroutine 如何感知這些錯誤?
- 當某個系統響應過慢,等待的 Goroutine 如何控制訪問超時?
這些問題都是直接使用 WaitGroup 沒法處理的。如果直接使用 channel 配合 WaitGroup 來控制超時和錯誤返回的話,封裝起來並不簡單,而且還容易出錯。我們可以採用 ErrGroup 來代替 WaitGroup。
有關 ErrGroup 的用法這裡就不再闡述。golang.org/x/sync/errgroup
package main
import (
"context"
"fmt"
"golang.org/x/sync/errgroup"
"time"
)
func main() {
ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
defer cancel()
errGroup, newCtx := errgroup.WithContext(ctx)
done := make(chan struct{})
go func() {
for i := 0; i < 10; i++ {
errGroup.Go(func() error {
time.Sleep(time.Second * 10)
return nil
})
}
if err := errGroup.Wait(); err != nil {
fmt.Printf("do err:%v\n", err)
return
}
done <- struct{}{}
}()
select {
case <-newCtx.Done():
fmt.Printf("err:%v ", newCtx.Err())
return
case <-done:
}
fmt.Println("success")
}
2. 控制 Goroutine 數量
場景模擬:
大概有 2000 - 3000 萬個資料需要處理,根據對伺服器的測試,當啟動 200 個 Goroutine 處理時效能最佳。如何控制?
遇到諸如此類的問題時,單純使用 WaitGroup 是不行的。既要保證所有的資料都能被處理,同時也要保證同時最多隻有 200 個 Goroutine。這種問題需要 WaitGroup 配合 Channel 一塊使用。
package main
import (
"fmt"
"sync"
"time"
)
func main() {
var wg = sync.WaitGroup{}
manyDataList := []int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
ch := make(chan bool, 3)
for _, v := range manyDataList {
wg.Add(1)
go func(data int) {
defer wg.Done()
ch <- true
fmt.Printf("go func: %d, time: %d\n", data, time.Now().Unix())
time.Sleep(time.Second)
<-ch
}(v)
}
wg.Wait()
}
使用注意點
使用 WaitGroup 同樣不能被複制。具體例子就不再分析了。具體分析過程可以參見《這可能是最容易理解的 Go Mutex 原始碼剖析》
WaitGroup 的剖析到這裡基本就結束了。有什麼想跟我交流的,歡迎評論區留言。
歡迎關注我的公眾號:HHFCodeRV,一起學習一起進步