一. 序言
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。