Golang 之 WaitGroup 源码解析
如果我们有一个大的任务要做,我们会尝试将这个任务分解,分解完成之后并发交由 goroutine 去做,并且我需要当全部的任务完成之后再进行下面的步骤,在 sync 包下,就有这样一个东西适合上述情况,WaitGroup,今天我们来看看具体它是怎么实现的。
PS:在下面我统一用 wg 来简称 WaitGroup
使用
它的使用非常简单,如下:
1 | func main () { |
输出:
1 | job 9 done |
我们可以看到,使用非常简单,每次有一个任务就使用 Add 方法加一个,每次做完任务就使用 Done 方法告诉它已经完成了,而 Wait 就是等着所有的任务完成。
思考问题
在看 wg 的实现之前,首先来问几个问题,来考考自己。
- Wait 方法能否被多次调用,比如再开一个 goroutine 去 wait
- Wait 方法调用后是否还能再继续调用 Add 添加任务
- 每次只能 Done 一个任务,能否一次性 Done 多个任务呢
- wg 能否被拷贝或作为参数传递
- 如果让你自己实现一个,你会如何实现
前几个问题,如果你都能很清楚的回答,那么你对 wg 的了解可以说已经非常熟悉了。首选我来说一下对于最后的一个问题的回答,因为在看源码之前我都会想想如果是我,我会如何去实现,那么我想的也很简单。
- 使用一个变量进行计数
- 每次任务数量变更时使用 atom 原子操作 + 1 或者 - 1
- -1 时判断任务数量是否已经为 0
- 如果为 0 向一个 channel 里面发送消息
- 所有 wait 的地方监听 channel 的消息,收到消息则证明任务全部完成
源码分析
结构
1 | type WaitGroup struct { |
结构非常简单,就只有两个熟悉,一个 noCopy
还有一个 state1
(我也很好奇为什么要用 1 来结尾命名,大佬的想法总是很奇妙)
noCopy: sync 包下的一个特殊标记吧,vet 检查,如果有拷贝的变量则会报错
1 | func main () { |
你 run 肯定没问题的,但是如果你使用 go vet 做个检查就有警告了
1 | ➜ go vet main.go |
state1:是用来存放任务计数器和等待者计数器的(我一看到这个结构就明白肯定后面又是位操作这样的高端操作了)
state [0] | state [1] | state [2] | |
---|---|---|---|
64 位 | waiter | counter | sema |
32 位 | sema | waiter | counter |
其中 waiter 是等待者计数,counter 是任务计数,sema 是信号量
奇怪的是在 64 位还 32 位操作系统上是不一样的,具体原因以及对于它操作请继续看下去
state
1 | //state returns pointers to the state and sema fields stored within wg.state1. |
这个方法是一个内部方法,就是将 state1 中存储的状态取出来,返回值 statep 就是计数器的状态,semap 是信号量
Done
1 | func (wg *WaitGroup) Done () { |
没想到吧~居然 Done 就是调用 Add 并传递一个 - 1
所以其实我们完全可以再外部调用 Add 传递一个 - 3 一次性结束 3 个任务
Add
1 | func (wg *WaitGroup) Add (delta int) { |
这里有几个要点我们其实已经看到了:
- Wait 的 ** 过程中 ** 是不能 Add 的,不然就会 panic,要注意
- 虽然我们可以借助 Add 一个负数来一次性结束多个任务,但是如果任务数量控制的不好,变成负数也会 panic,Done 次数多了也一样
- wg 是通过信号量来通知的,当然可以有很多人在等,wg 它都会一一通知到位的
Wait
1 | func (wg *WaitGroup) Wait () { |
其实 wait 虽然简单,也有要点
- 通过 load 和 cas 操作 + 循环来避免了锁,其实这个操作可以学一下
- 其实这里也说明明白了,wg 可以重用,但是你必须等到 wait 全部完成之后再说
其他注意点
1 | func main () { |
上面的代码有问题吗?问题在哪呢?
其实很简单,wg 作为一个参数传递的时候,wg 还是一个普通的结构体,我们在函数中操作的时候还是操作的一个拷贝的变量而已,对于原来的 wg 是不会改变的,所以这里需要传递指针才是正确的
1 | func main () { |
但是其实并不推荐这样去传递 wg,因为这样很容易出现问题,一个不好就出问题了,个人还是建议直接在使用 goroutine 之后马上接一个 defer wg.Done () 来的更加靠谱一些
总结
回过头来看看,之前的问题也都有了答案:
- Wait 可以被调用多次,并且每个都会收到完成的通知
- Wait 之后,如果再 Wait 的过程中不能在 Add,否则会 panic,但是 Wait 结束之后可以继续使用 Add 进行重用
- 可以使用 Add 传递负数的方式一次性结束多个任务,但是需要保证任务计数器非负,否则会 panic
- wg作为参数传递的时候需要注意传递指针,或者尽量避免传递
- 官方利用位操作节约了空间,存在在同一个地方;利用信号量来实现任务结束的通知….
总的来说 wg 的实现还是非常简单的,需要注意的就是几个使用上的点不要出现意外即可。