一. 序言
WaitGroup是Golang應用開發過程中經常使用的併發控制技術。
WaitGroup，可理解為Wait-Goroutine-Group，即等待一組goroutine結束。比如某個goroutine需要等待其他幾個goroutine全部完成，那麼使用WaitGroup可以輕鬆實現。
下面是一段demo.go示例
package main
import (
"fmt"
"sync"
)
func worker(i int) {
fmt.Println("worker: ", i)
}
func main() {
// 例項化一個 wg
var wg sync.WaitGroup
// 啟動10個worker協程
for i := 0; i < 10; i++ {
// 協程執行前 加1
wg.Add(1)
go func(i int) {
// 執行完畢後執行done,相當於計數減1
defer wg.Done()
worker(i)
}(i)
}
// 等待所有的協程執行完畢後，執行其他邏輯
wg.Wait()
}
demo2.go
下面程式展示了一個goroutine等待另外兩個goroutine結束的例子：
package main
import (
"fmt"
"time"
"sync"
)
func main() {
var wg sync.WaitGroup
wg.Add(2) //設定計數器，數值即為goroutine的個數
go func() {
//Do some work
time.Sleep(1*time.Second)
fmt.Println("Goroutine 1 finished!")
wg.Done() //goroutine執行結束後將計數器減1
}()
go func() {
//Do some work
time.Sleep(2*time.Second)
fmt.Println("Goroutine 2 finished!")
wg.Done() //goroutine執行結束後將計數器減1
}()
wg.Wait() //主goroutine阻塞等待計數器變為0
fmt.Printf("All Goroutine finished!")
}
簡單的說，上面程式中wg內部維護了一個計數器：
- 啟動goroutine前將計數器通過Add(2)將計數器設定為待啟動的goroutine個數。
- 啟動goroutine後，使用Wait()方法阻塞自己，等待計數器變為0。
3 . 每個goroutine執行結束通過Done()方法將計數器減1。 - 計數器變為0後，阻塞的goroutine被喚醒。
如何支援多個 goroutine 等待一個 goroutine 完成後再幹活呢？
二.原始碼分析
2.1 訊號量
訊號量是Unix系統提供的一種保護共享資源的機制，用於防止多個執行緒同時訪問某個資源。
可簡單理解為訊號量為一個數值：
- 當訊號量>0時，表示資源可用，獲取訊號量時系統自動將訊號量減1；
- 當訊號量==0時，表示資源暫不可用，獲取訊號量時，當前執行緒會進入睡眠，當訊號量為正時被喚醒；
由於在WaitGroup實現中也是用了訊號量，因此做一個簡單介紹
WaitGroup是一個結構體
type WaitGroup struct {
noCopy noCopy
// 64-bit value: high 32 bits are counter, low 32 bits are waiter count.
// 64-bit atomic operations require 64-bit alignment, but 32-bit
// compilers do not ensure it. So we allocate 12 bytes and then use
// the aligned 8 bytes in them as state, and the other 4 as storage
// for the sema.
state1 [3]uint32
}
結構十分簡單，由 nocopy 和 state1 兩個欄位組成，其中 nocopy 是用來防止複製的.
nocopy是一個空結構體，包含兩個方法
type noCopy struct{}
// Lock is a no-op used by -copylocks checker from `go vet`.
func (*noCopy) Lock() {}
func (*noCopy) Unlock() {}
由於嵌入了 nocopy 所以在執行 go vet 時如果檢查到 WaitGroup 被複制了就會報錯。這樣可以一定程度上保證 WaitGroup 不被複制，對了直接 go run 是不會有錯誤的，所以我們程式碼 push 之前都會強制要求進行 lint 檢查，在 ci/cd 階段也需要先進行 lint 檢查，避免出現這種類似的錯誤。
state1是個長度為3的陣列，其中包含了state和一個訊號量，而state實際上是兩個計數器：
- counter： 當前還未執行結束的goroutine計數器
- waiter count: 等待goroutine-group結束的goroutine數量，即有多少個等候者
- semaphore: 訊號量
state1 的設計非常巧妙，這是一個是十二位元組的資料，這裡面主要包含兩大塊，counter 佔用了 8 位元組用於計數，sema 佔用 4 位元組用做訊號量.
為什麼要這麼搞呢？直接用兩個欄位一個表示 counter，一個表示 sema 不行麼？
不行，我們看看註釋裡面怎麼寫的。
程式碼註釋中大概意思是：在做 64 位的原子操作的時候必須要保證 64 位（8 位元組）對齊，如果沒有對齊的就會有問題，但是 32 位的編譯器並不能保證 64 位對齊所以這裡用一個 12 位元組的 state1 欄位來儲存這兩個狀態，然後根據是否 8 位元組對齊選擇不同的儲存方式。
這個操作巧妙在哪裡呢？
- 如果是 64 位的機器那肯定是 8 位元組對齊了的，即第一種方式
- 如果在 32位的機器上
- 如果恰好 8 位元組對齊，那也是第一種方式取前面的8位元組
- 如果是沒有對其，但是32位4位元組是對齊的，所以只需要後裔四個位元組，那個8個位元組就對齊了
所以通過 sema 訊號量這四個位元組的位置不同，保證了 counter 這個欄位無論在 32 位還是 64 為機器上都是 8 位元組對齊的，後續做 64 位原子操作的時候就沒問題了。
考慮到位元組是否對齊，三者出現的位置不同，為簡單起見，依照位元組已對齊情況下，三者在記憶體中的位置如下所示：
state 方法實現如下
func (wg *WaitGroup) state() (statep *uint64, semap *uint32) {
// 取8的餘數如果餘數為0說明8位元組對齊了
if uintptr(unsafe.Pointer(&wg.state1))%8 == 0 {
return (*uint64)(unsafe.Pointer(&wg.state1)), &wg.state1[2]
} else {
return (*uint64)(unsafe.Pointer(&wg.state1[1])), &wg.state1[0]
}
}
state 方法返回 counter 和訊號量，通過 uintptr(unsafe.Pointer(&wg.state1))%8 == 0 來判斷是否 8 位元組對齊
2.1 Add
func (wg *WaitGroup) Add(delta int) {
// 先從 state 當中把資料和訊號量取出來
statep, semap := wg.state()
// 在 waiter 上加上 delta 值
state := atomic.AddUint64(statep, uint64(delta)<<32)
// 取出當前的 counter
v := int32(state >> 32)
// 取出當前的 waiter，正在等待 goroutine 數量
w := uint32(state)
// counter 不能為負數
if v < 0 {
panic("sync: negative WaitGroup counter")
}
// 這裡屬於防禦性程式設計
// w != 0 說明現在已經有 goroutine 在等待中，說明已經呼叫了 Wait() 方法
// 這時候 delta > 0 && v == int32(delta) 說明在呼叫了 Wait() 方法之後又想加入新的等待者
// 這種操作是不允許的
if w != 0 && delta > 0 && v == int32(delta) {
panic("sync: WaitGroup misuse: Add called concurrently with Wait")
}
// 如果當前沒有人在等待就直接返回，並且 counter > 0
if v > 0 || w == 0 {
return
}
// 這裡也是防禦 主要避免併發呼叫 add 和 wait
if *statep != state {
panic("sync: WaitGroup misuse: Add called concurrently with Wait")
}
// 喚醒所有 waiter，看到這裡就回答了上面的問題了
*statep = 0
for ; w != 0; w-- {
runtime_Semrelease(semap, false, 0)
}
}
2.2 Wait
wait 主要就是等待其他的 goroutine 完事之後喚醒
func (wg *WaitGroup) Wait() {
// 先從 state 當中把資料和訊號量的地址取出來
statep, semap := wg.state()
for {
// 這裡去除 counter 和 waiter 的資料
state := atomic.LoadUint64(statep)
v := int32(state >> 32)
w := uint32(state)
// counter = 0 說明沒有在等的，直接返回就行
if v == 0 {
// Counter is 0, no need to wait.
return
}
// waiter + 1，呼叫一次就多一個等待者，然後休眠當前 goroutine 等待被喚醒
if atomic.CompareAndSwapUint64(statep, state, state+1) {
runtime_Semacquire(semap)
if *statep != 0 {
panic("sync: WaitGroup is reused before previous Wait has returned")
}
return
}
}
}
2.3 Done
Done是對 Add的封裝
func (wg *WaitGroup) Done() {
wg.Add(-1)
}
三. 總結
簡單說來，WaitGroup通常用於等待一組“工作協程”結束的場景，其內部維護兩個計數器，這裡把它們稱為“工作協程”計數器和“坐等協程”計數器，
WaitGroup對外提供的三個方法分工非常明確：
Add(delta int)方法用於增加“工作協程”計數，通常在啟動新的“工作協程”之前呼叫；Done()方法用於減少“工作協程”計數，每次呼叫遞減1，通常在“工作協程”內部且在臨近返回之前呼叫；Wait()方法用於增加“坐等協程”計數，通常在所有”工作協程”全部啟動之後呼叫；
WaitGroup 可以用於一個 goroutine 等待多個 goroutine 幹活完成，也可以多個 goroutine 等待一個 goroutine 幹活完成，是一個多對多的關係
多個等待一個的典型案例是 singleflight，這個在後面將微服務可用性的時候還會再講到，感興趣可以看看原始碼
3.1 注意事項
- Done()方法除了負責遞減“工作協程”計數以外，還會在“工作協程”計數變為0時檢查“坐等協程”計數器並把“坐等協程”喚醒。
需要注意的是，Done()方法遞減“工作協程”計數後，如果“工作協程”計數變成負數時，將會觸發panic，這就要求Add()方法呼叫要早於Done()方法。 - 此外，通過
Add()方法累加的“工作協程”計數要與實際需要等待的“工作協程”數量一致，否則也會觸發panic。 - 當“工作協程”計數多於實際需要等待的“工作協程”數量時，“坐等協程”可能會永遠無法被喚醒而產生列鎖，此時，Go執行時檢測到死鎖會觸發panic
- 當“工作協程”計數小於實際需要等待的“工作協程”數量時，Done()會在“工作協程”計數變為負數時觸發panic。