WaitGroup

WaitGroup

总结

• WaitGroup 提供Add、Wait、done三个方法,Add负责增加计数器,done负责发送信号,然后Add进行减一操作,等到计数器为0的时候会唤醒全部的休眠goroutine
• Add 不能在和 Wait 方法在 Goroutine 中并发调用，一旦出现就会造成程序崩溃；
• WaitGroup 必须在 Wait 方法返回之后才能被重新使用；
• Done 只是对 Add 方法的简单封装，我们可以向 Add 方法传入任意负数（需要保证计数器非负）快速将计数器归零以唤醒其他等待的 Goroutine；
• 可以同时有多个 Goroutine 等待当前 WaitGroup 计数器的归零，这些 Goroutine 也会被『同时』唤醒；

WaitGroup 是 Go 语言 sync 包中比较常见的同步机制，它可以用于等待一系列的 Goroutine 的返回，一个比较常见的使用场景是批量执行 RPC 或者调用外部服务：

requests := []*Request{...}
wg := &sync.WaitGroup{}
wg.Add(len(requests))
for _, request := range requests {
  go func(r *Request) {
    defer wg.Done()
    // res, err := service.call(r)
  }(request)
}
wg.Wait()

通过 WaitGroup 我们可以在多个 Goroutine 之间非常轻松地同步信息，原本顺序执行的代码也可以在多个 Goroutine 中并发执行，加快了程序处理的速度，在上述代码中只有在所有的 Goroutine 都执行完毕之后 Wait 方法才会返回，程序可以继续执行其他的逻辑。

总而言之，它的作用就像它的名字一样，，通过 Done 来传递任务完成的信号，比较常用于等待一组 Goroutine 中并发执行的任务全部结束。

结构体

WaitGroup 结构体中的成员变量非常简单，其中的 noCopy 的主要作用就是保证 WaitGroup 不会被开发者通过再赋值的方式进行拷贝，进而导致一些诡异的行为：

type WaitGroup struct {   
	noCopy noCopy    
	state1 [3]uint32
}

copylock 包就是一个用于检查类似错误的分析器，它的原理就是在 编译期间 检查被拷贝的变量中是否包含 noCopy 或者 sync 关键字，如果包含当前关键字就会报出以下的错误：

package main
import (
    "fmt"
    "sync"
)
func main() {
    wg := sync.Mutex{}
    yawg := wg
    fmt.Println(wg, yawg)
}
$ go run proc.go
./prog.go:10:10: assignment copies lock value to yawg: sync.Mutex
./prog.go:11:14: call of fmt.Println copies lock value: sync.Mutex
./prog.go:11:18: call of fmt.Println copies lock value: sync.Mutex

除了 noCopy 之外，WaitGroup 结构体中还包含一个总共占用 12 字节大小的数组，这个数组中会存储当前结构体持有的状态和信号量，在 64 位与 32 位的机器上表现也非常不同。

WaitGroup 提供了私有方法 state 能够帮助我们从 state1 字段中取出它的状态和信号量。

操作

WaitGroup 对外暴露的接口只有三个 Add、Wait 和 Done，其中 Done 方法只是调用了 wg.Add(-1) 本身并没有什么特殊的逻辑，我们来了解一下剩余的两个方法：

func (wg *WaitGroup) Add(delta int) {
    statep, semap := wg.state()
    state := atomic.AddUint64(statep, uint64(delta)<<32)
    v := int32(state >> 32)
    w := uint32(state)
    if v < 0 {
        panic("sync: negative WaitGroup counter")
    }
    if v > 0 || w == 0 {
        return
    }
    *statep = 0
    for ; w != 0; w-- {
        runtime_Semrelease(semap, false, 0)
    }
}

Add 方法的主要作用就是更新 WaitGroup 中持有的计数器 counter，64 位状态的高 32 位，虽然 Add 方法传入的参数可以为负数，但是一个 WaitGroup 的计数器只能是非负数，当调用 Add 方法导致计数器归零并且还有等待的 Goroutine 时，就会通过 runtime_Semrelease 唤醒处于等待状态的所有 Goroutine。

另一个 WaitGroup 的方法 Wait 就会在当前计数器中保存的数据大于 0 时修改等待 Goroutine 的个数 waiter 并调用 runtime_Semacquire 陷入睡眠状态。

func (wg *WaitGroup) Wait() {
    statep, semap := wg.state()
    for {
        state := atomic.LoadUint64(statep)
        v := int32(state >> 32)
        if v == 0 {
            return
        }
        if atomic.CompareAndSwapUint64(statep, state, state+1) {
            runtime_Semacquire(semap)
            if +statep != 0 {
                panic("sync: WaitGroup is reused before previous Wait has returned")
            }
            return
        }
    }
}

陷入睡眠的 Goroutine 就会等待 Add 方法在计数器为 0 时唤醒。