Golang Sync.WaitGroup 使用及原理

邱佳飛發表於2022-02-22
GolangAI

Golang Sync.WaitGroup 使用及原理

使用

func main() {
	var wg sync.WaitGroup
	for i := 0; i < 10; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			fmt.Println("Hello WaitGroup!")
		}()
	}
	wg.Wait()
}

實現

首先看 waitgroup 到底是什麼資料結構

type WaitGroup struct {
	noCopy noCopy
	state1 [3]uint32
}

nocopy 避免這個結構體被複制的一個技巧，可以告訴go vet工具違反了複製使用的規則
state1 [3]uint32 欄位中包含了 waitgroup 的所有狀態資訊, 根據標準庫上自帶的註釋簡單翻譯是：state1 由 12 個位元組組成，其中將8位元組看作64位值，其中高32位存放的是 counter 計數器, 代表目前還未完成的 goroutine個數，低32位存放的是 waiter 計數器, 可以理解成下面這個結構體

type WaitGroup struct {
	// 代表目前尚未完成的個數
	// WaitGroup.Add(n) 將會導致 counter += n
	// WaitGroup.Done() 將導致 counter--
	counter uint32

	// 目前已呼叫 WaitGroup.Wait 的 goroutine 的個數
	waiter  uint32

	// 對應於 golang 中 runtime 內部的訊號量的實現
	// runtime_Semacquire 表示增加一個訊號量，並掛起當前 goroutine
	// runtime_Semrelease 表示減少一個訊號量，並喚醒 sema 上其中一個正在等待的 goroutine
	sema    uint32
}

整個使用流程為：

  1. 當呼叫 WaitGroup.Add(n) 時，counter 將會自增: counter += n
  2. 當呼叫 WaitGroup.Wait() 時，會將 waiter++。同時呼叫 runtime_Semacquire(semap), 增加訊號量，並掛起當前 goroutine。
  3. 當呼叫 WaitGroup.Done() 時，將會 counter--。如果自減後的 counter 等於 0，說明 WaitGroup 的等待過程已經結束，則需要呼叫 runtime_Semrelease 釋放訊號量，喚醒正在 WaitGroup.Wait 的 goroutine。

原始碼中是如何拆分 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]  
   }  
}

由於我們能使用到的就是 waitgroup.Add(), waitgroup.Done(), waitgroup.Wait() 這三個方法，就按這三個方法分析

Add(), Done()

Add 方法主要操作的是 state 的計數部分。你可以為計數值增加一個 delta 值，內部通過原子操作把這個值加到計數值上。需要注意的是，這個 delta 也可以是個負數，相當於為計數值減去一個值，Done 方法內部其實就是通過 Add(-1) 實現的。

func (wg *WaitGroup) Add(delta int) {  
   // 獲取拆開後的 state 欄位 
   statep, semap := wg.state()  
   ...
   ...
   ...
   // 在剛剛說的 int64 的高32位上加傷傳進來的 delta 的值， 這一步是原子操作
   state := atomic.AddUint64(statep, uint64(delta)<<32)  
   // 加好後，獲取 counter 也就是 v， 和 waiter 也就是 w 的值
   // 此時 int64 變為兩個 int32
   v := int32(state >> 32)  
   w := uint32(state)  
   // 如果 v 變為負數了，程式異常
   if v < 0 {  
      panic("sync: negative WaitGroup counter")  
   }  
   // 在 wait 沒結束之前， 不允許呼叫 Add 方法
   if w != 0 && delta > 0 && v == int32(delta) {  
      panic("sync: WaitGroup misuse: Add called concurrently with Wait")  
   }

   // 呼叫 add() 之後, 還有正在執行的 goroutine 或者 waiter 等於 0， 正常返回
   if v > 0 || w == 0 {  
      return  
   }  
   // 下面就是非正常返回， 理解到的就是 v 已經等於 0 了，執行釋放操作
   // 首先就是將 counter 和 waiter 全部重置為 0
 *statep = 0  
   // 然後迴圈呼叫還在等待的 waiter， 釋放訊號量
 for ; w != 0; w-- {  
      runtime_Semrelease(semap, false, 0)  
   }  
}

wait()

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

func (wg *WaitGroup) Wait() {
   // 獲取訊號量和兩個計數值
   statep, semap := wg.state()  

   // 不停的迴圈檢查 counter 和 waiter
   for {
      // 先原子性的取出 counter 和 waiter
      state := atomic.LoadUint64(statep)  
      v := int32(state >> 32)  
      w := uint32(state)
      if v == 0 {  
         // counter 已經沒有了，函式可以返回
         return  
 }  
      // 將 waiter 數 + 1
 if atomic.CompareAndSwapUint64(statep, state, state+1) {  
         // 放到訊號量佇列, 並且阻塞住自己
         runtime_Semacquire(semap)
         // 如果被喚醒，檢查 兩個計數是否已經為0 了， 如果不為0 ，則觸發恐慌
         if *statep != 0 {  
            panic("sync: WaitGroup is reused before previous Wait has returned")  
         }
          // 函式返回
         return  
 }  
   }  
}

總結

  1. 保證計數器不能為負值
  2. 保證 Add() 方法全部呼叫完成之後再呼叫 Wait()
  3. waitgroup 可以重複使用
  4. atomic 原子操作代替鎖, 提高併發性
  5. 合併兩個 int32 為一個 int64 提高讀取存入資料效能
  6. 對於不希望被複制的結構體， 可以使用 noCopy 欄位

reference

https://www.cyhone.com/articles/golang-waitgroup/
https://time.geekbang.org/column/intro/100061801?tab=catalog

相關文章