Golang中 WaitGroup的實現原理是什么

這篇文章給大家介紹Golang中 WaitGroup的實現原理是什么，內容非常詳細，感興趣的小伙伴們可以參考借鑒，希望對大家能有所幫助。

1 前言

WaitGroup是Golang應用開發過程中經常使用的并發控制技術。

WaitGroup，可理解為Wait-Goroutine-Group，即等待一組goroutine結束。比如某個goroutine需要等待其他幾個goroutine全部完成，那么使用WaitGroup可以輕松實現。

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內部維護了一個計數器：

1. 啟動goroutine前將計數器通過Add(2)將計數器設置為待啟動的goroutine個數。

2. 啟動goroutine后，使用Wait()方法阻塞自己，等待計數器變為0。

3. 每個goroutine執行結束通過Done()方法將計數器減1。

4. 計數器變為0后，阻塞的goroutine被喚醒。

其實WaitGroup也可以實現一組goroutine等待另一組goroutine，這有點像玩雜技，很容出錯，如果不了解其實現原理更是如此。實際上，WaitGroup的實現源碼非常簡單。

2 基礎知識

2.1 信號量

信號量是Unix系統提供的一種保護共享資源的機制，用于防止多個線程同時訪問某個資源。

可簡單理解為信號量為一個數值：

• 當信號量>0時，表示資源可用，獲取信號量時系統自動將信號量減1；

• 當信號量==0時，表示資源暫不可用，獲取信號量時，當前線程會進入睡眠，當信號量為正時被喚醒；

由于WaitGroup實現中也使用了信號量，在此做個簡單介紹。

3 WaitGroup

3.1 數據結構

源碼包中src/sync/waitgroup.go:WaitGroup定義了其數據結構：

type WaitGroup struct {
	state1 [3]uint32
}

state1是個長度為3的數組，其中包含了state和一個信號量，而state實際上是兩個計數器：

• counter： 當前還未執行結束的goroutine計數器

• waiter count: 等待goroutine-group結束的goroutine數量，即有多少個等候者

• semaphore: 信號量

考慮到字節是否對齊，三者出現的位置不同，為簡單起見，依照字節已對齊情況下，三者在內存中的位置如下所示：

WaitGroup對外提供三個接口：

• Add(delta int): 將delta值加到counter中

• Wait()： waiter遞增1，并阻塞等待信號量semaphore

• Done()： counter遞減1，按照waiter數值釋放相應次數信號量

下面分別介紹這三個函數的實現細節。

3.2 Add(delta int)

Add()做了兩件事，一是把delta值累加到counter中，因為delta可以為負值，也就是說counter有可能變成0或負值，所以第二件事就是當counter值變為0時，跟據waiter數值釋放等量的信號量，把等待的goroutine全部喚醒，如果counter變為負值，則panic.

Add()偽代碼如下：

func (wg *WaitGroup) Add(delta int) {
    statep, semap := wg.state() //獲取state和semaphore地址指針
    
    state := atomic.AddUint64(statep, uint64(delta)<<32) //把delta左移32位累加到state，即累加到counter中
    v := int32(state >> 32) //獲取counter值
    w := uint32(state)      //獲取waiter值
    
    if v < 0 {              //經過累加后counter值變為負值，panic
        panic("sync: negative WaitGroup counter")
    }

    //經過累加后，此時，counter >= 0
    //如果counter為正，說明不需要釋放信號量，直接退出
    //如果waiter為零，說明沒有等待者，也不需要釋放信號量，直接退出
    if v > 0 || w == 0 {
        return
    }

    //此時，counter一定等于0，而waiter一定大于0（內部維護waiter，不會出現小于0的情況），
    //先把counter置為0，再釋放waiter個數的信號量
    *statep = 0
    for ; w != 0; w-- {
        runtime_Semrelease(semap, false) //釋放信號量，執行一次釋放一個，喚醒一個等待者
    }
}

3.3 Wait()

Wait()方法也做了兩件事，一是累加waiter, 二是阻塞等待信號量

func (wg *WaitGroup) Wait() {
    statep, semap := wg.state() //獲取state和semaphore地址指針
    for {
        state := atomic.LoadUint64(statep) //獲取state值
        v := int32(state >> 32)            //獲取counter值
        w := uint32(state)                 //獲取waiter值
        if v == 0 {                        //如果counter值為0，說明所有goroutine都退出了，不需要待待，直接返回
            return
        }
        
        // 使用CAS（比較交換算法）累加waiter，累加可能會失敗，失敗后通過for loop下次重試
        if atomic.CompareAndSwapUint64(statep, state, state+1) {
            runtime_Semacquire(semap) //累加成功后，等待信號量喚醒自己
            return
        }
    }
}

這里用到了CAS算法保證有多個goroutine同時執行Wait()時也能正確累加waiter。

3.4 Done()

Done()只做一件事，即把counter減1，我們知道Add()可以接受負值，所以Done實際上只是調用了Add(-1)。

源碼如下：

func (wg *WaitGroup) Done() {
	wg.Add(-1)
}

Done()的執行邏輯就轉到了Add()，實際上也正是最后一個完成的goroutine把等待者喚醒的。

4 編程Tips

• Add()操作必須早于Wait(), 否則會panic

• Add()設置的值必須與實際等待的goroutine個數一致，否則會panic

關于Golang中 WaitGroup的實現原理是什么就分享到這里了，希望以上內容可以對大家有一定的幫助，可以學到更多知識。如果覺得文章不錯，可以把它分享出去讓更多的人看到。