基本介绍
WaitGroup是go用来做任务编排的一个并发原语，它要解决的就是并发 - 等待的问题：
当有一个 goroutine A 在检查点（checkpoint）等待一组 goroutine 全部完成，如果这些 goroutine 还没全部完成，goroutine A 就会阻塞在检查点，直到所有 goroutine 都完成后才能继续执行
试想如果没有WaitGroup，想要在协程A等到其他协程执行完成后能立马执行，只能不断轮询其他协程是否执行完毕，这样的问题是：
- 及时性差：轮询间隔越高，及时性越差
- 无谓的空轮训，浪费系统资源
而用WaitGroup时，协程A只用阻塞，直到其他协程执行完毕后，再通知协程A
其他语言也提供了类似的工具，例如Java的CountDownLatch
使用
Waitgroup提供了3个方法：
func (wg *WaitGroup) Add(delta int)
func (wg *WaitGroup) Done()
func (wg *WaitGroup) Wait()
复制代码- Add：增加计数值
- Done：减少计数值
-
Wait：调用这个方法的 goroutine 会一直阻塞，直到 WaitGroup 的计数值变为 0
源码分析
type WaitGroup struct {
// 避免复制
noCopy noCopy
// 64位环境下，高32位是计数值，低32位记录waiter的数量
state1 uint64
// 用于信号量
state2 uint32
}
复制代码Add
func (wg *WaitGroup) Add(delta int) {
// 获取状态值，信号量
statep, semap := wg.state()
// 将参数delta左32位，加到statep中，即给计数值加上delta
state := atomic.AddUint64(statep, uint64(delta)<<32)
// 加后的计数值
v := int32(state >> 32)
// waiter的数量
w := uint32(state)
// 加后不能是负值
if v < 0 {
panic( "sync: negative WaitGroup counter" )
}
// 有waiter的情况下，当前协程又加了计数值，panic
// 即有waiter的情况下，不能再给waitgroup增加计数值了
if w != 0 && delta > 0 && v == int32(delta) {
panic( "sync: WaitGroup misuse: Add called concurrently with Wait" )
}
// 如果加完后v大于0，或者加完后v等于0，但没有等待者，直接返回
if v > 0 || w == 0 {
return
}
// 接下来就是v等于0，且w大于0的情况
// 再次检查是否有Add和Wait并发调用的情况
if *statep != state {
panic( "sync: WaitGroup misuse: Add called concurrently with Wait" )
}
// 将计数值和waiter数量清0
*statep = 0
// 唤醒所有的waiter
for ; w != 0; w-- {
runtime_Semrelease(semap, false, 0)
}
}
复制代码- 因为state高32位保存计数值，因此需要将参数delta左移32位后加到state上才正确
-
如果加完后v大于0，或者加完后v等于0，但没有等待者，直接返回
- v大于0：表示自己不是最后一个调用Done的协程，不用自己来释放waiter，直接返回
- v等于0，但没有等待者：因为没有等待者，也就不用释放等待者，也直接返回
-
否则就是v等于0，且w大于0的情况：
- 自己是最后一个调用Done的，且还有等待者，那就唤醒所有等待者
Done
Done内部调用Add，只是参数传-1，表示减少计数值
func (wg *WaitGroup) Done() {
wg.Add(-1)
}
复制代码Wait
func (wg *WaitGroup) Wait() {
statep, semap := wg.state()
for {
state := atomic.LoadUint64(statep)
// v：计数值
v := int32(state >> 32)
w := uint32(state)
// 如果计数值为0，自己不需要等到，直接返回
if v == 0 {
return
}
// 增加waiter计数值
if atomic.CompareAndSwapUint64(statep, state, state+1) {
// 自己在信号量上阻塞
runtime_Semacquire(semap)
// 检查Waitgroup是否在wait返回前被重用
if *statep != 0 {
panic( "sync: WaitGroup is reused before previous Wait has returned" )
}
return
}
}
}
复制代码- 如果计数值为0，当前不需要阻塞，直接返回
- 否则将waiter数量加1，如果添加成功，就把自己阻塞到信号量上
-
被唤醒时，如果statep不为0，表示该waitgroup是否在wait返回前被重用了，panic
注意事项
通过源码分析可以看出，Waitgroup有以下使用注意事项：
-
计数器的值必须大于等于0：
- 一开始调用Add时，不能传负数
- 调用Done的次数不能过多，导致超过了 WaitGroup 的计数值
- 因此使用 WaitGroup 的正确姿势是，预先确定好 WaitGroup 的计数值，然后调用相同次数的 Done 完成相应的任务
-
要保证在期望的Add调用完成后，再调用Wait，否则Wait发现计数值为0时不会阻塞
- 最好在一个协程中，按顺序先调Add，再调Wait
-
需要重用时，需要在前一组调用Wait结束后，再开始新一轮的使用
- WaitGroup 是可以重用的。只要 WaitGroup 的计值恢复到零值的状态，那么它就可以被看作是新创建的 WaitGroup，被重复使用，而不能在前一组没使用完的情况下又使用