go test 私有方法_Go并发编程-WaitGroup的设计实现

在进行服务编程时，一个比较常见的需求是把无依赖的数据，通过批量发出 RPC 或者 HTTP 请求等来拉取，把串行的请求转化为并行的请求，加快程序的处理速度。由于sync.WaitGroup理解和使用简单，在并发编程中被广泛使用。本文从WaitGroup的设计、源码实现和使用的角度出发，讲述WaitGroup的方方面面。

注：源码文件：sync/waitgroup.go

一、sync.WaitGroup的结构与定义

// A WaitGroup must not be copied after first use.
type WaitGroup struct {
    noCopy noCopy    // 一个特殊的结构体，用于禁止对象复制
    // 64-bit waiter-counter-sema
    // 32-bit sema-waiter-counter
    // waiter: 有多少个Goroutine在等待批量运行结束
    // counter: 有多少个正在执行且需要被等待的Goroutine，即Add数
    // sema: 信号量
    state1 [3]uint32
}

1. noCopy

noCopy是一个私有的空结构体，用于保证WaitGroup不会进行值复制，通过go vet命令会报发下的错误：

func main() {
	wg1 := sync.WaitGroup{}
	wg2 := wg1
	wg1.Add(1)
	fmt.Println(wg1, wg2)
}
~ go vet test.go
# command-line-arguments
./test.go:12:9: assignment copies lock value to wg2: sync.WaitGroup contains sync.noCopy
./test.go:14:14: call of fmt.Println copies lock value: sync.WaitGroup contains sync.noCopy
./test.go:14:19: call of fmt.Println copies lock value: sync.WaitGroup contains sync.noCopy

但是使用go build却可以顺利编译并运行，这里先留一个疑问，后续跟进解决，但是不管怎样，WaitGroup都是不应该被复制的。

2.state1

state1是一个3个元素的uint32数组，占12个字节，用于存储WaitGroup的状态和信号量，在32位和64位的机器上，各个元素的含义是不同的：

• Waiter: 有多少个Goroutine在等待批量运行结束
• Counter: 有多少个正在执行且需要被等待的Goroutine，即Add数
• Sema: 信号量，用于唤醒等待中的Goroutine

为了屏蔽不同的位数的系统的差异，WaitGroup提供了state方法，用于统一获取state1中状态值（Waiter+Counter）和信号量，其实现如下：

// 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]
	}
}

该函数首选获取WaitGroup.state1的地址，然后根据该地址是否为8字节的内存对齐来确定系统是否为64位系统，若是则前两个元素为状态信息（Waiter+Counter），最后一个元素为信号量；否则，第一个元素为信息量，后两个元素为状态信息。

二、sync.WaitGroup.Add

sync.WaitGroup.Add 方法用于更新 sync.WaitGroup 中的计数器 counter。当调用后Counter归零，也就是所有任务都执行完成时，就会通过 sync.runtime_Semrelease 唤醒处于等待状态的所有 Goroutine。

sync.WaitGroup.Add 方法传入的参数可以为负数，但是Counter的值只能是非负数，一旦出现负数就会Panic。

下面是其实现的代码（省略了部分与设计实现不相关的代码），相关代码的解释，以注释的形式补充在代码上：

func (wg *WaitGroup) Add(delta int) {
    // statep: *int64，高32位：counter，低32位：waiter
    statep, semap := wg.state()
    // 加上delta，由于counter在高位，需要左移32位
    state := atomic.AddUint64(statep, uint64(delta)<<32)
    // 由于可能被其他的Goroutine修改，需要重新计算状态值
    v := int32(state >> 32) // 计算出counter
    w := uint32(state)      // 计算出waiter
    if v < 0 {
        panic("sync: negative WaitGroup counter")
    }
    if v > 0 || w == 0 {
        return
    }
    // Reset waiters count to 0.
    *statep = 0
    // 唤醒Waiter的GoRoutine
    for ; w != 0; w-- {
        runtime_Semrelease(semap, false, 0)
    }
}

sync.WaitGroup.Done，其实是调用sync.WaitGroup.Add实现的，类似于语法糖：

// Done decrements the WaitGroup counter by one.
func (wg *WaitGroup) Done() {
	wg.Add(-1)
}

三、sync.WaitGroup.Wait

sync.WaitGroup.Wait用于增加Waiter的值，请求并等待信号量，在等待信号量期间，会阻塞当前的Goroutine。当 sync.WaitGroup 的Counter归零时，陷入睡眠状态的 Goroutine 就会wf被唤醒，Wait方法会立刻返回。

其实现如下：

// Wait blocks until the WaitGroup counter is zero.
func (wg *WaitGroup) Wait() {
    // statep: *int64，高32位：counter，低32位：waiter
    statep, semap := wg.state()
    for {
        state := atomic.LoadUint64(statep)
        v := int32(state >> 32) // 计算出counter
        w := uint32(state)      // 计算出waiter
        if v == 0 {
            // Counter is 0, no need to wait.
            return
        }
        // Increment waiters count.
        if atomic.CompareAndSwapUint64(statep, state, state+1) {
            // 请求并等待信号量
            runtime_Semacquire(semap)
            if *statep != 0 {
                panic("sync: WaitGroup is reused before previous Wait has returned")
            }
            return
        }
    }
}

四、总结

通过上述的分析，可以总结出：

• sync.WaitGroup必须在Wait方法返回后，才可以被重新使用
• sync.WaitGroup.Add可以传入任意的整数，但是必须保证Counter为非负数，当Counter为零时，会唤醒正在等待的Goroutine
• sync.WaitGroup.Done只是对Add方法的简单封装，本质上是Add(-1)
• 可以同时有多个 Goroutine 等待当前的sync.WaitGroup的Counter归零，而当Counter归零时，这些等待的Groutine会被同时唤醒