Golang 之 WaitGroup 源码解析
前言
如果我们有一个大的任务要做，我们会尝试将这个任务分解，分解完成之后并发交由 goroutine 去做，并且我需要当全部的任务完成之后再进行下面的步骤，在 sync 包下，就有这样一个东西适合上述情况，WaitGroup，今天我们来看看具体它是怎么实现的。
PS：在下面我统一用 wg 来简称 WaitGroup
使用
它的使用非常简单，如下:
func main () {
wg := sync.WaitGroup {}
for i := 0; i < 10; i++ {
wg.Add (1)
go func (job int) {
defer wg.Done ()
//do something
fmt.Printf ("job % d done\n", job)
}(i)
}
wg.Wait ()
fmt.Println ("all done")
}
输出：
job 9 done
job 1 done
job 0 done
job 8 done
job 7 done
job 3 done
job 6 done
job 2 done
job 4 done
job 5 done
all done
我们可以看到，使用非常简单，每次有一个任务就使用 Add 方法加一个，每次做完任务就使用 Done 方法告诉它已经完成了，而 Wait 就是等着所有的任务完成。
思考问题
在看 wg 的实现之前，首先来问几个问题，来考考自己。
- Wait 方法能否被多次调用，比如再开一个 goroutine 去 wait
- Wait 方法调用后是否还能再继续调用 Add 添加任务
- 每次只能 Done 一个任务，能否一次性 Done 多个任务呢
- wg 能否被拷贝或作为参数传递
- 如果让你自己实现一个，你会如何实现
前几个问题，如果你都能很清楚的回答，那么你对 wg 的了解可以说已经非常熟悉了。首选我来说一下对于最后的一个问题的回答，因为在看源码之前我都会想想如果是我，我会如何去实现，那么我想的也很简单。
- 使用一个变量进行计数
- 每次任务数量变更时使用 atom 原子操作 + 1 或者 - 1
- -1 时判断任务数量是否已经为 0
- 如果为 0 向一个 channel 里面发送消息
- 所有 wait 的地方监听 channel 的消息，收到消息则证明任务全部完成
源码分析
结构
type WaitGroup struct {
noCopy noCopy
// 64-bit value: high 32 bits are counter, low 32 bits are waiter count.
// 64-bit atomic operations require 64-bit alignment, but 32-bit
//compilers do not ensure it. So we allocate 12 bytes and then use
//the aligned 8 bytes in them as state, and the other 4 as storage
//for the sema.
state1 [3] uint32
}
结构非常简单，就只有两个熟悉，一个 noCopy 还有一个 state1（我也很好奇为什么要用 1 来结尾命名，大佬的想法总是很奇妙）
noCopy： sync 包下的一个特殊标记吧，vet 检查，如果有拷贝的变量则会报错
func main () {
wg := sync.WaitGroup {}
w := wg
fmt.Println (w, wg)
}
你 run 肯定没问题的，但是如果你使用 go vet 做个检查就有警告了
➜ go vet main.go
# command-line-arguments
./main.go:10:10: assignment copies lock value to w: sync.WaitGroup contains sync.noCopy
./main.go:11:17: call of fmt.Println copies lock value: sync.WaitGroup contains sync.noCopy
./main.go:11:20: call of fmt.Println copies lock value: sync.WaitGroup contains sync.noCopy
state1：是用来存放任务计数器和等待者计数器的（我一看到这个结构就明白肯定后面又是位操作这样的高端操作了）
| state [0] | state [1] | state [2] | |
|---|---|---|---|
| 64 位 | waiter | counter | sema |
| 32 位 | sema | waiter | counter |
其中 waiter 是等待者计数，counter 是任务计数，sema 是信号量
奇怪的是在 64 位还 32 位操作系统上是不一样的，具体原因以及对于它操作请继续看下去
state
//state returns pointers to the state and sema fields stored within wg.state1.
func (wg *WaitGroup) state () (statep *uint64, semap *uint32) {
if uintptr (unsafe.Pointer (&wg.state1))%8 == 0 {
return (*uint64)(unsafe.Pointer (&wg.state1)), &wg.state1 [2]
} else {
return (*uint64)(unsafe.Pointer (&wg.state1 [1])), &wg.state1 [0]
}
}
这个方法是一个内部方法，就是将 state1 中存储的状态取出来，返回值 statep 就是计数器的状态，semap 是信号量
Done
func (wg *WaitGroup) Done () {
wg.Add (-1)
}
没想到吧～居然 Done 就是调用 Add 并传递一个 - 1
所以其实我们完全可以再外部调用 Add 传递一个 - 3 一次性结束 3 个任务
Add
func (wg *WaitGroup) Add (delta int) {
// 首先获取状态值
statep, semap := wg.state ()
// 对于 statep 中 counter + delta
state := atomic.AddUint64 (statep, uint64 (delta)<<32)
// 获取任务计数器的值
v := int32 (state >> 32)
// 获取等待者计数器的值
w := uint32 (state)
// 任务计数器不能为负数
if v < 0 {
panic ("sync: negative WaitGroup counter")
}
// 已经有人在等待，但是还在添加任务
if w != 0 && delta > 0 && v == int32 (delta) {
panic ("sync: WaitGroup misuse: Add called concurrently with Wait")
}
// 没有等待者或者任务还有没做完的
if v > 0 || w == 0 {
return
}
// 有等待者，但是在这个过程中数据还在变动
if *statep != state {
panic ("sync: WaitGroup misuse: Add called concurrently with Wait")
}
// Reset waiters count to 0.
// 重置状态，并用发出等同于等待者数量的信号量，告诉所有等待者任务已经完成
*statep = 0
for ; w != 0; w-- {
runtime_Semrelease (semap, false, 0)
}
}
这里有几个要点我们其实已经看到了：
- Wait 的 ** 过程中 ** 是不能 Add 的，不然就会 panic，要注意
- 虽然我们可以借助 Add 一个负数来一次性结束多个任务，但是如果任务数量控制的不好，变成负数也会 panic，Done 次数多了也一样
- wg 是通过信号量来通知的，当然可以有很多人在等，wg 它都会一一通知到位的
Wait
func (wg *WaitGroup) Wait () {
// 先获取状态
statep, semap := wg.state ()
for {
// 这里注意要用 atomic 的 Load 来保证一下写操作已经完成
state := atomic.LoadUint64 (statep)
// 同样的，这里是任务计数
v := int32 (state >> 32)
// 这里是等待者计数
w := uint32 (state)
// 如果没有任务，那么直接结束，不用等待了
if v == 0 {
return
}
// 使用 cas 操作，如果不相等，证明中间已经被其他人修改了状态，重新走 for 循环
// 注意这里 if 进去之后等待者的数量就 +1 了
if atomic.CompareAndSwapUint64 (statep, state, state+1) {
// 等待信号量
runtime_Semacquire (semap)
// 如果信号量来了，但是状态还不是 0，则证明 wait 之后还是在人在 add，证明有人想充分利用 wg 但是时机不对
if *statep != 0 {
panic ("sync: WaitGroup is reused before previous Wait has returned")
}
return
}
}
}
其实 wait 虽然简单，也有要点
- 通过 load 和 cas 操作 + 循环来避免了锁，其实这个操作可以学一下
- 其实这里也说明明白了，wg 可以重用，但是你必须等到 wait 全部完成之后再说
其他注意点
func main () {
wg := sync.WaitGroup {}
for i := 0; i < 10; i++ {
wg.Add (1)
go func (job int) {
doJob (job, wg)
}(i)
}
wg.Wait ()
fmt.Println ("all done")
}
func doJob (job int, wg sync.WaitGroup) {
fmt.Printf ("job % d done\n", job)
wg.Done ()
}
上面的代码有问题吗？问题在哪呢？
其实很简单，wg 作为一个参数传递的时候，wg 还是一个普通的结构体，我们在函数中操作的时候还是操作的一个拷贝的变量而已，对于原来的 wg 是不会改变的，所以这里需要传递指针才是正确的
func main () {
wg := &sync.WaitGroup {}
for i := 0; i < 10; i++ {
wg.Add (1)
go func (job int) {
doJob (job, wg)
}(i)
}
wg.Wait ()
fmt.Println ("all done")
}
func doJob (job int, wg *sync.WaitGroup) {
fmt.Printf ("job % d done\n", job)
wg.Done ()
}
但是其实并不推荐这样去传递 wg，因为这样很容易出现问题，一个不好就出问题了，个人还是建议直接在使用 goroutine 之后马上接一个 defer wg.Done () 来的更加靠谱一些
总结
回过头来看看，之前的问题也都有了答案：
- Wait 可以被调用多次，并且每个都会收到完成的通知
- Wait 之后，如果再 Wait 的过程中不能在 Add，否则会 panic，但是 Wait 结束之后可以继续使用 Add 进行重用
- 可以使用 Add 传递负数的方式一次性结束多个任务，但是需要保证任务计数器非负，否则会 panic
- wg作为参数传递的时候需要注意传递指针，或者尽量避免传递
- 官方利用位操作节约了空间，存在在同一个地方；利用信号量来实现任务结束的通知....
总的来说 wg 的实现还是非常简单的，需要注意的就是几个使用上的点不要出现意外即可。
