Golang 主協程如何等待子協程

介紹

眾所周知，如果我們使用 go 關鍵字進行協程的操作時，如果 主協程 沒有等待其他協程的話，直接結束會導致 任務無法完全執行。

範例如下：

func main(){
  go sayHi(){
    fmt.Println("say hello......")
  }()
  fmt.Println("main groutine....")
}
// 結果
// main groutine....

一般來說，最暴力解的方式就是 讓主協程睡一會，等待子協程完成再來去完成主協程的任務。

1	time.Sleep(3 * time.Second)

但是這樣做非常沒有效率，且 sleep 的期間子協程可能早就已經結束了。

解決方法

使用 sync.Group
使用 channel

使用 `sync.WaitGroup`

1
2
3

func (wg *WaitGroup) Add(delta int)
func (wg *WaitGroup) Done()
func (wg *WaitGroup) Wait()

sync.WaitGroup 可以提供一個等待子協程完成的功能，主要配合使用 WaitGroup.Add(1)、WaitGroup.Wait()、WaitGroup.Done()。

WaitGroup.Add(1)：sync.WaitGroup 內有計數器，新增需要等待協成的數量。
WaitGroup.Done()：已經完成的程協執行後，可以減少 sync.WaitGroup 當前等待的數量。
WaitGroup.Wait()：用來阻塞當前協程，等待子協程，直到計數器歸零才會繼續執行。

使用方式

func waitByWaitGroup(n int) {
  // 宣告 sync.WaitGroup
  var wg sync.WaitGroup

  for i := 0; i < n; i++ {
    // 協程計數器++
    wg.Add(1)
    go func(i int) {
      // 協程計數器--
      defer wg.Done()
			time.Sleep(1 * time.Nanosecond)
      fmt.Printf("Goroutine %d finished\n", i)
    }(i)
  }

  fmt.Println("Waiting for all goroutines to finish")
  // 阻塞等待計數器為 0
  wg.Wait()
  fmt.Println("All goroutines finished")
}

注意事項

WaitGroup.Add(num) 使用時，num 不能新增負數，不然會報錯 panic: sync: negative WaitGroup counter。
sync.WaitGroup 在使用時是一個 實例 而不是指標，所以在函數裡需要記得以指標方式傳遞才不會產生 deadlock。

使用 `channel`

使用迴圈的方式同步等待 channel 取出值來，邏輯與 sync.WaitGroup 相似。

func waitByChannelCount(n int) {
  done := make(chan bool)
  for i := 0; i < n; i++ {
    go func(done chan bool) {
      fmt.Println("working")
			time.Sleep(1 * time.Nanosecond)
      done <- true
    }(done)
  }

  fmt.Println("waiting")

  i := 0
Loop:
  for {
    select {
    case <-done:
      i++
    default:
      if i == n {
        break Loop
      }
    }
  }
  fmt.Println("done")
}

Benchmark

可以發現使用 sync.WaitGroup 速度約 channel 的兩倍，簡單又好用。

$ go test -v -bench=. -run=none -benchmem .
goos: darwin
goarch: amd64
cpu: VirtualApple @ 2.50GHz
BenchmarkWaitByChannelCount
BenchmarkWaitByChannelCount-8 26499     41924 ns/op      10533 B/op 201 allocs/op
BenchmarkWaitByWaitGroup
BenchmarkWaitByWaitGroup-8    45924     24285 ns/op      12016 B/op 301 allocs/op
BenchmarkSleep
BenchmarkSleep-8              1         1001063750 ns/op 10480 B/op 201 allocs/op
PASS

結論

其實要去等待子協程還是使用 sync.WaitGroup 程式碼看起來才會比較單純且高效，但是如果需要多個不同 channel 同時進行處理，再來考慮要不要使用 channel 配合 for + select 比較好。

本篇文章的程式碼範例。

參考資料

https://learnku.com/articles/35130
https://blog.wu-boy.com/2019/05/handle-multiple-channel-in-15-minutes/

介紹