Go 语言系列26：WaitGroup

之前写的协程有关的例子，为了保证所有协程都执行完毕后退出，我们都使用了 ​​time.Sleep(time.Millisecond)​​ 使主协程休眠，但在实际开发中我们并不能保证要休眠多久才能执行完毕，用多了主程序就阻塞了，用少了又会导致子协程的任务可能没法完成。下面我们介绍处理这种情况的方式。

使用信道

信道可以实现多个协程间的通信，于是乎我们可以定义一个信道，在任务执行完成后，往信道中写入 ​​true​​​ ，然后在主协程中获取到 ​​true​​ ，就可以认为子协程已经执行完毕。

package main

import "fmt"

func main() {
 isDone := make(chan bool)
 go func() {
  for i := 0; i < 5; i++{
   fmt.Println(i)
  }
  isDone <- true
 }()
 <- isDone
}

运行上面的程序，主协程就会等待创建的协程执行完毕后退出，运行后输出如下：

0
1
2
3
4

使用 WaitGroup

使用上面的信道方法，虽然可行，但在你程序中使用很多协程的话，你的代码就会看起来很复杂，这里就要介绍一种更好的方法，那就是使用 ​​sync​​ 包中提供的 WaitGroup 类型。​​WaitGroup​​​ 用于等待一批 Go 协程执行结束。程序控制会一直阻塞，直到这些协程全部执行完毕。当然 ​​WaitGroup​​ 也可以用于实现工作池。

​​WaitGroup​​ 实例化后就能使用：

var name sync.WaitGroup

​​WaitGroup​​ 有几个方法：

​​Add​​​：初始值为 ​​0​​​ ，这里直接传入子协程的数量，你传入的值会往计数器上加。​​Done​​​：当某个子协程完成后，可调用此方法，会从计数器上减一，即子协程的数量减一，通常使用 ​​defer​​​ 来调用。​​Wait​​：阻塞当前协程，直到实例里的计数器归零。

package main

import (
 "fmt"
 "sync"
)

func task(taskNum int, waitGroup *sync.WaitGroup) {
 // 延迟调用 执行完子协程计数器减一
 defer waitGroup.Done()
 // 输出任务号
 for i := 0; i < 3; i++ {
  fmt.Printf("task %d: %d\n", taskNum, i)
 }
}

func main() {
 // 实例化 sync.WaitGroup
 var waitGroup sync.WaitGroup
 // 传入子协程的数量
 waitGroup.Add(3)
 // 开启一个子协程 协程 1 以及 实例 waitGroup
 go task(1, &waitGroup)
 // 开启一个子协程 协程 2 以及 实例 waitGroup
 go task(2, &waitGroup)
 // 开启一个子协程 协程 3 以及 实例 waitGroup
 go task(3, &waitGroup)
 // 实例 waitGroup 阻塞当前协程 等待所有子协程执行完
 waitGroup.Wait()
}

运行上面的程序一种输出如下：

task 3: 0
task 3: 1
task 3: 2
task 1: 0
task 1: 1
task 1: 2
task 2: 0
task 2: 1
task 2: 2

参考文献：

[1] Alan A. A. Donovan; Brian W. Kernighan, Go 程序设计语言, Translated by 李道兵, 高博, 庞向才, 金鑫鑫 and 林齐斌, 机械工业出版社, 2017.