【Go進階—併發程式設計】WaitGroup

與昊發表於2022-03-07
Go程式設計AI

WaitGroup 是開發過程中經常使用的併發控制技術，用來在程式中控制等待一組 goroutine 結束。

實現原理

資料結構

WaitGroup 的資料結構包括了一個 noCopy 的輔助欄位，一個 state1 記錄 WaitGroup 狀態的陣列：

  • noCopy 的輔助欄位；
  • state1，一個具有複合意義的欄位，包含 WaitGroup 的計數、阻塞在檢查點的 waiter 數和訊號量。
type WaitGroup struct {
    // 避免複製使用的一個技巧，可以告訴 vet 工具違反了複製使用的規則
    noCopy noCopy
    // 前 64bit(8bytes) 的值分成兩段，高 32bit 是計數值，低 32bit 是 waiter 的計數
    // 另外 32bit 是用作訊號量的
    // 因為 64bit 值的原子操作需要 64bit 對齊，但是 32bit 編譯器不支援，所以陣列中的元素在不同的架構中不一樣，具體處理看下面的方法
    // 總之，會找到對齊的那 64bit 作為 state，其餘的 32bit 做訊號量
    state1 [3]uint32
}

// 得到state的地址和訊號量的地址
func (wg *WaitGroup) state() (statep *uint64, semap *uint32) {
    if uintptr(unsafe.Pointer(&wg.state1))%8 == 0 {
        // 如果地址是 64bit 對齊的，陣列前兩個元素做 state，後一個元素做訊號量
        return (*uint64)(unsafe.Pointer(&wg.state1)), &wg.state1[2]
    } else {
        // 如果地址是 32bit 對齊的，陣列後兩個元素用來做 state，它可以用來做 64bit 的原子操作，第一個元素 32bit 用來做訊號量
        return (*uint64)(unsafe.Pointer(&wg.state1[1])), &wg.state1[0]
    }
}

在 64 位環境下，state1 的第一個元素是 waiter 數，第二個元素是 WaitGroup 的計數值，第三個元素是訊號量。

image.png

noCopy：輔助 vet 檢查

noCopy 欄位的作用是指示 vet 工具在做檢查的時候，這個資料結構不能做值複製使用。更嚴謹地說，是不能在第一次使用之後複製使用。

vet 會對實現 Locker 介面的資料型別做靜態檢查，一旦程式碼中有複製使用這種資料型別的情況，就會發出警告。但是，WaitGroup 不滿足 Locker 介面，這時就可以通過給 WaitGroup 新增一個 noCopy 欄位來實現 Locker 介面。而且因為 noCopy 欄位是未輸出型別，所以 WaitGroup 不會暴露 Lock/Unlock 方法。

如果你想要自己定義的資料結構不被複制使用，或者說，不能通過 vet 工具檢查出複製使用的報警，就可以通過嵌入 noCopy 這個資料型別來實現。

方法

Add & Done

Add 方法主要操作的是 state 的計數部分，去除 race 檢查和異常檢查的程式碼後，它的實現如下：

func (wg *WaitGroup) Add(delta int) {
    statep, semap := wg.state()
    // 高 32bit 是計數值 v，所以把 delta 左移 32，增加到計數上
    state := atomic.AddUint64(statep, uint64(delta)<<32)
    v := int32(state >> 32) // 當前計數值
    w := uint32(state) // waiter count
    
    if v > 0 || w == 0 {
        return
    }
    
    // 如果計數值 v 為 0 並且 waiter 的數量 w 不為 0，那麼 state 的值就是 waiter 的數量
    // 將waiter的數量設定為 0，因為計數值 v 也是 0,所以它們倆的組合 *statep 直接設定為 0 即可。此時需要並喚醒所有的 waiter
    *statep = 0
    for ; w != 0; w-- {
        runtime_Semrelease(semap, false, 0)
    }
}

// Done 方法實際就是計數器減 1
func (wg *WaitGroup) Done() {
    wg.Add(-1)
}
Wait

Wait 方法的實現邏輯是：不斷檢查 state 的值。如果其中的計數值變為了 0，那麼說明所有的任務已完成，呼叫者不必再等待，直接返回。如果計數值大於 0，說明此時還有任務沒完成，那麼呼叫者就變成了等待者，需要加入 waiter 佇列，並且阻塞住自己。

其主幹實現程式碼如下：

func (wg *WaitGroup) Wait() {
    statep, semap := wg.state()
    
    for {
        state := atomic.LoadUint64(statep)
        v := int32(state >> 32) // 當前計數值
        w := uint32(state) // waiter 的數量
        if v == 0 {
            // 如果計數值為 0, 呼叫這個方法的 goroutine 不必再等待，繼續執行它後面的邏輯即可
            return
        }
        // 否則把 waiter 數量加 1。期間可能有併發呼叫 Wait 的情況，增加可能會失敗，所以最外層使用了一個 for 迴圈
        if atomic.CompareAndSwapUint64(statep, state, state+1) {
            // 阻塞休眠等待
            runtime_Semacquire(semap)
            // 被喚醒，不再阻塞，返回
            return
        }
    }
}

常見錯誤

計數器設定為負值

WaitGroup 的計數器的值必須大於等於 0。我們在更改這個計數值的時候，WaitGroup 會先做檢查，如果計數值被設定為負數，就會導致 panic。

一般情況下，有兩種方法會導致計數器設定為負數：

  1. 呼叫 Add 的時候傳遞一個負數。如果你能保證當前的計數器加上這個負數後還是大於等於 0 的話，也沒有問題，否則就會導致 panic。
  2. 呼叫 Done 方法的次數過多，超過了 WaitGroup 的計數值。

Add 時機錯誤

在使用 WaitGroup 的時候，你一定要遵循的原則就是，等所有的 Add 方法呼叫之後再呼叫 Wait，否則就可能導致 panic 或者不期望的結果。

前一個 Wait 還沒結束就重用 WaitGroup

只要 WaitGroup 的計數值恢復到零值的狀態，那麼它就可以被看作是新建立的 WaitGroup，被重複使用。但是，如果我們在 WaitGroup 的計數值還沒有恢復到零值的時候就重用，就會導致程式 panic。我們看一個例子，初始設定 WaitGroup 的計數值為 1，啟動一個 goroutine 先呼叫 Done 方法，接著就呼叫 Add 方法，Add 方法有可能和主 goroutine 併發執行。

func main() {
    var wg sync.WaitGroup
    wg.Add(1)
    go func() {
        time.Sleep(time.Millisecond)
        wg.Done() // 計數器減 1
        wg.Add(1) // 計數值加 1
    }()
    wg.Wait() // 主 goroutine 等待，有可能和第 7 行併發執行
}

在這個例子中，第 6 行雖然讓 WaitGroup 的計數恢復到 0，但是因為第 9 行有個 waiter 在等待，如果等待 Wait 的 goroutine，剛被喚醒就和 Add 呼叫（第 7 行）有併發執行的衝突，所以就會出現 panic。

WaitGroup 雖然可以重用，但是是有一個前提的，那就是必須等到上一輪的 Wait 完成之後，才能重用 WaitGroup 執行下一輪的 Add/Wait，如果你在 Wait 還沒執行完的時候就呼叫下一輪 Add 方法，就有可能出現 panic。

相關文章