一、介紹
WaitGroup是多個goroutine之間協作的一種實現方式，主要功能就是阻塞等待一組goroutine執行完成。
常用的使用場景：主goroutine調用Add函數設置需要等待的goroutine的數量,當每個goroutine執行完成後調用Done函數(將counter減1)，Wait函數用於阻塞等待直到該組中的所有goroutine都執行完成。
源碼中主要設計到的三個概念：counter、waiter和semaphore
counter： 當前還未執行結束的goroutine計數器
waiter : 等待goroutine-group結束的goroutine數量，即有多少個等候者
semaphore: 信號量
信號量是Unix系統提供的一種保護共享資源的機制，用於防止多個線程同時訪問某個資源。
可簡單理解爲信號量爲一個數值：
當信號量>0時，表示資源可用，獲取信號量時系統自動將信號量減1；
當信號量=0時，表示資源暫不可用，獲取信號量時，當前線程會進入睡眠，當信號量爲正時被喚醒。
二、源碼分析
Golang源碼版本 ：1.10.3
1.結構體
type WaitGroup struct {
noCopy noCopy //該WaitGroup對象不允許拷貝使用,只能用指針傳遞
// 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.
//用於存儲計數器(counter)和waiter的值
// 只需要64位,即8個字節,其中高32位是counter值,低32位值是waiter值
// 不直接使用uint64,是因爲uint64的原子操作需要64位系統,而32位系統下,可能會出現崩潰
// 所以這裏用byte數組來實現,32位系統下4字節對齊,64位系統下8字節對齊,所以申請12個字節,其中必定有8個字節是符合8字節對齊的,下面的state()函數中有進行判斷
state1 [12]byte
sema uint32 //信號量
}
從結構體中我們看到
state1是一個12位長度的byte數組，用於存儲counter和waiter的值
sema就是傳說中的信號量
2.state函數
state是一個內部函數，用於獲取counter和 waiter的值
//獲取counter 、 waiter的值 (counter是uint64的高32位，waiter是uint64的低32位)
func (wg *WaitGroup) state() *uint64 {
// 根據state1的起始地址分析,若是8字節對齊的,則直接用前8個字節作爲*uint64類型
// 若不是,說明是4字節對齊,則後移4個字節後,這樣必爲8字節對齊,然後取後面8個字節作爲*uint64類型
if uintptr(unsafe.Pointer(&wg.state1))%8 == 0 {
return (*uint64)(unsafe.Pointer(&wg.state1))
} else {
return (*uint64)(unsafe.Pointer(&wg.state1[4]))
}
}
3.Add方法
//用於增加或減少計數器(counter)的值
//如果計數器爲0，則釋放調用Wait方法時的阻塞，如果計數器爲負，則panic
//Add()方法應該在Wait()方法調用之前
func (wg *WaitGroup) Add(delta int) {
//獲取當前counter和 waiter的值
statep := wg.state()
if race.Enabled {
_ = *statep // trigger nil deref early
if delta < 0 {
// Synchronize decrements with Wait.
race.ReleaseMerge(unsafe.Pointer(wg))
}
race.Disable()
defer race.Enable()
}
//將delta的值添加到counter上
state := atomic.AddUint64(statep, uint64(delta)<<32)
v := int32(state >> 32) //counter值
w := uint32(state) //waiter值
if race.Enabled && delta > 0 && v == int32(delta) {
// The first increment must be synchronized with Wait.
// Need to model this as a read, because there can be
// several concurrent wg.counter transitions from 0.
race.Read(unsafe.Pointer(&wg.sema))
}
//counter爲負數，則觸發panic
if v < 0 {
panic("sync: negative WaitGroup counter")
}
// waiter值不爲0,累加後的counter值和delta相等,說明Wait()方法沒有在Add()方法之後調用,觸發panic,因爲正確的做法是先Add()後Wait()
if w != 0 && delta > 0 && v == int32(delta) {
panic("sync: WaitGroup misuse: Add called concurrently with Wait")
}
//Add()添加正常返回
//1.counter > 0,說明還不需要釋放信號量，可以直接返回
//2. waiter = 0 ,說明沒有等待的goroutine，也不需要釋放信號量，可以直接返回
if v > 0 || w == 0 {
return
}
// This goroutine has set counter to 0 when waiters > 0.
// Now there can't be concurrent mutations of state:
// - Adds must not happen concurrently with Wait,
// - Wait does not increment waiters if it sees counter == 0.
// Still do a cheap sanity check to detect WaitGroup misuse.
//下面是 counter == 0 並且 waiter > 0的情況
//現在若原state和新的state不等，則有以下兩種可能
//1. Add 和 Wait方法同時調用
//2. counter已經爲0，但waiter值有增加，這種情況永遠不會觸發信號量了
// 以上兩種情況都是錯誤的，所以觸發異常
//注：state := atomic.AddUint64(statep, uint64(delta)<<32) 這一步調用之後，state和*statep的值應該是相等的，除非有以上兩種情況發生
if *statep != state {
panic("sync: WaitGroup misuse: Add called concurrently with Wait")
}
// Reset waiters count to 0.
//將waiter 和 counter都置爲0
*statep = 0
//原子遞減信號量，並通知等待的goroutine
for ; w != 0; w-- {
runtime_Semrelease(&wg.sema, false)
}
}
4.Done方法
// Done decrements the WaitGroup counter by one.
//將計數器(counter)的值減1
func (wg *WaitGroup) Done() {
wg.Add(-1)
}
5.Wait方法
// Wait blocks until the WaitGroup counter is zero.
//調用Wait方法會阻塞當前調用的goroutine直到 counter的值爲0
//也會增加waiter的值
func (wg *WaitGroup) Wait() {
//獲取當前counter和 waiter的值
statep := wg.state()
if race.Enabled {
_ = *statep // trigger nil deref early
race.Disable()
}
//一直等待，直到無需等待或信號量觸發，才返回
for {
state := atomic.LoadUint64(statep)
v := int32(state >> 32) //counter值
w := uint32(state) //waiter值
//如果counter值爲0，則說明所有goroutine都退出了，無需等待，直接退出
if v == 0 {
// Counter is 0, no need to wait.
if race.Enabled {
race.Enable()
race.Acquire(unsafe.Pointer(wg))
}
return
}
// Increment waiters count.
//原子增加waiter的值，CAS方法，外面for循環會一直嘗試，保證多個goroutine同時調用Wait()也能正常累加waiter
if atomic.CompareAndSwapUint64(statep, state, state+1) {
if race.Enabled && w == 0 {
// Wait must be synchronized with the first Add.
// Need to model this is as a write to race with the read in Add.
// As a consequence, can do the write only for the first waiter,
// otherwise concurrent Waits will race with each other.
race.Write(unsafe.Pointer(&wg.sema))
}
//一直等待信號量sema，直到信號量觸發，
runtime_Semacquire(&wg.sema)
//從上面的Add()方法看到，觸發信號量之前會將seatep置爲0(即counter和waiter都置爲0)，所以此時應該也爲0
//如果不爲0，說明WaitGroup此時又執行了Add()或者Wait()操作，所以會觸發panic
if *statep != 0 {
panic("sync: WaitGroup is reused before previous Wait has returned")
}
if race.Enabled {
race.Enable()
race.Acquire(unsafe.Pointer(wg))
}
return
}
}
}
三、注意點
1.Add()必須在Wait()前調用
2.Add()設置的值必須與實際等待的goroutine個數一致，如果設置的值大於實際的goroutine數量，可能會一直阻塞。如果小於會觸發panic
3. WaitGroup不可拷貝，可以通過指針傳遞，否則很容易造成BUG
以下爲值拷貝引起的Bug示例
demo1：因爲值拷貝引起的死鎖
func main() {
var wg sync.WaitGroup
wg.Add(5)
for i := 0 ; i < 5 ; i++ {
test(wg)
}
wg.Wait()
}
func test(wg sync.WaitGroup) {
go func() {
fmt.Println("hello")
wg.Done()
}()
}
demo2:因爲值拷貝引起的不會阻塞等待現象
func main() {
var wg sync.WaitGroup
for i := 0 ; i < 5 ; i++ {
test(wg)
}
wg.Wait()
}
func test(wg sync.WaitGroup) {
go func() {
wg.Add(1)
fmt.Println("hello")
time.Sleep(time.Second*5)
wg.Done()
}()
}
demo3:因爲值拷貝引發的panic
type person struct {
wg sync.WaitGroup
}
func (t *person) say() {
go func() {
fmt.Println("say Hello!")
time.Sleep(time.Second*5)
t.wg.Done()
}()
}
func main() {
var wg sync.WaitGroup
t := person{wg:wg}
wg.Add(5)
for i := 0 ; i< 5 ;i++ {
t.say()
}
wg.Wait()
}
感謝：https://blog.csdn.net/yzf279533105/article/details/97302666