Golang WaitGroup源码分析

1 WaitGroup功能

WaitGroup对外只提供Add, Done,以及Wait三个方法。我们先假设WaitGroup里面维护了一个Counter计数器, Add和Done都是对这个数字的加减操作。整个WaitGroup并且有以下使用规则：

1. Add可以任意int型数字，但是必须Counter是大于等于0的，否则panic
2. Add(-1)等于Done()
3. Wait()可以调用多次，但是必须要在Counter为0的位置调用。否则一定会panic
4. WaitGroup可以重复使用，但是必须等上一次全部使用结束。否则可能会发生panic
5. Wait()和Add()不能一起并发执行。否则可能会发生panic

2 源码分析

WaitGroup Struct

type WaitGroup struct {
	noCopy noCopy
	state1 uint64
	state2 uint32
}

noCopy暂时不用关注，是golang来判断是否WaitGroup被复制。本身就是一个空struct,不占用任何大小。
state1 state2 两个变量分别一起用来表示三个值，counter(也就是上面说的Counter)、waiter(调用wait()方法的个数)、sema(信号量个数)。
uint64 state1分别可以看做是两个uint32，通过位运算来分割。在早期的实现版本中，state是一个3个大小的uint32的数组，实际上都是一样的。
在实际处理中，counter和waiter是合并到一个uint64变量,但是因为内存对齐原因，这个变量并不一定是state1，sema个数使用uint32大小的变量。这里通过封装内部的state()方法来拿到：

func (wg *WaitGroup) state() (statep *uint64, semap *uint32) {
	if unsafe.Alignof(wg.state1) == 8 || uintptr(unsafe.Pointer(&wg.state1))%8 == 0 {
		return &wg.state1, &wg.state2
	} else {
		state := (*[3]uint32)(unsafe.Pointer(&wg.state1))
		return (*uint64)(unsafe.Pointer(&state[1])), &state[0]
	}
}

因为实现使用了atomic uint64， 所以在32位的平台上，必须手动保证使用的变量uint64值要像是在64位上面一样内存对齐。
在uint32平台上面，内存最小划分是4B（64位是8B），state1是uint64，在32位平台上面，他无论如何都是对齐的（从0B对齐），但是将32位内存映射到64位(8B)上面，state1就有可能是从4B的位置开始，也有可能是0B开始。
就像上面代码，判断条件unsafe.Alignof(wg.state1) == 8是判断，state1是否是64位对齐对齐的。而uintptr(unsafe.Pointer(&wg.state1))%8 == 0则是判断32位上面，state1是0B对齐，还是4B对齐，如果是0B对齐，则直接返回。否则，返回state1前32位表示sema个数，state1后32位加上state2表示counter和waiter。

Add

去掉race相关源代码

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 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")
	}
	*statep = 0
	for ; w != 0; w-- {
		runtime_Semrelease(semap, false, 0)
	}
}

statep：前32位表示counter，后32位表示waiter。
semap：信号量个数。
state：调用Add新增传入state个数（delta）之后，新的statep值
v:调用AddUint64之后的counter
w:调用AddUint64之后的waiter

	if v < 0 {
		panic("sync: negative WaitGroup counter")
	}

校验规则1，内部counter必须任何时候大于等于0

	if w != 0 && delta > 0 && v == int32(delta) {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}

校验规则5，Add(正数)的时候和wait()不能同时调用，但是同时调用时候不是必然校验出来

	if v > 0 || w == 0 {
		return
	}

如果counter大于0，或者waiter等于0，因为只有counter等于0的时候，才回去唤醒waiter，当waiter也等于0，也就没必要唤醒waiter。所以走到下面的逻辑是counter ==0 && waiter > 0

	if *statep != state {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}
	*statep = 0

走到到这里来只有可能是counter等于0，并且waiter大于0
这里直接取statep值运算，其实是会有datarace的，这里只是做一个简单的判断，也有可能判断不出来Add()和Wait()是否正在并发调用。如果没有调用Add()或者Wait(),statep的值应该是和state相等。

	for ; w != 0; w-- {
		runtime_Semrelease(semap, false, 0)
	}

V操作，唤醒waiter

Done

Done()等于Add(-1)

Wait

去掉race相关源代码

// Wait blocks until the WaitGroup counter is zero.
func (wg *WaitGroup) Wait() {
	statep, semap := wg.state()
	for {
		state := atomic.LoadUint64(statep)
		v := int32(state >> 32)
		w := uint32(state)
		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
		}
	}
}

整体实现是一个for循环里面嵌套CAS，去尝试增加waiter。
如果counter为0了，直接退出Wait(),不用增加waiter，也不用P操作。
当CAS成功waiter增加1，则会进行P操作，gorountine挂起。等到Add()里面进行V操作，再次从runtime_Semacquire(semap) 位置唤醒。

	if *statep != 0 {
				panic("sync: WaitGroup is reused before previous Wait has returned")
	}

唤醒的时候也检查下statep，Add()唤醒的时候statep已经被重置为0，如果这里不为0，说明有位置再次调用了Add()。当然这里有data race，也不会百分百检查出是否真的有位置调用了Add()。

3 总结

WaitGroup代码非常少，但涉及的东西比较多。包括内存对齐，PV操作等。
不规范的使用，很多位置会导致panic，但是这个panic并不是确定会发生的。

留下疑问

runtime_Semacquire和runtime_Semrelease在Go里面如何实现的？其本质是什么？

源代码中，race相关代码也有其分析的价值，比如在哪里加race检查最合理，以及race实现原理是什么？

noCopy是如何实现的？这个可能在编译期就能判断出来吗？

atomic底层层是如何实现的？在uint32平台上面又是怎么搬到uint64的原子操作？

uint64能否拆成两个uint32 去表示counter和counter?
拆开似乎是造成counter和counter不实时同步的问题，所以会导致一些判断可能会有问题？

============================
第一次写源码分析，如果有什么位置不对，希望能帮忙指出来。留下的疑问如果知道答案，也希望留言解答一下？