WaitGroup

WaitGroup 用于等待一组线程的结束，父线程调用 Add 来增加等待的线程数，被等待的线程在结束后调用 Done 来将等待线程数减 1，父线程通过调用 Wait 阻塞等待所有结束（计数器清零）后进行唤醒。

源码位置

WaitGroup 的源码在 SDK 包的路径为src/sync/waitgroup.go。

数据结构

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

1.noCopy noCopy

noCopy 这个主要用来限制不能进行 copy，这里是为了避免 copy 后的 waitGroup 并发使用后，可能会与原 waitGroup 出现异常而 panic 。

2.state1 [3]unit32

数组的三个元素（非顺序）：

counter 通过 Add()设置的子 goroutine 的数量，即被等待线程计数
waiter 通过 Wait()陷入阻塞的等待者计数
semap 信号量，用于唤醒阻塞 waiter

这里需要注意一下 couter 、waiter 、semap 并不是顺序存储的，64bit 操作系统的原子操作需要保证 64bit 的内存对齐，在设计上我们需要保证 couter 和 waiter 的操作原子性。如果数组的首元素地址能被 8 整除，则 counter 和 waiter 刚好可以在同一块原子操作的 64bit 内存上，所以取数组前两个元素分别表示 couter 和 waiter ；如果不能被 8 整除（根据内存对齐的原理，地址必然是 4 的倍数），则取数组后两个。

// 根据内存对齐方式的不同，返回 statep(couter 占用高 32bit 和 waiter 占用低 32bit)和 semap 的地址
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]
	}
}

公共方法

func (wg *WaitGroup) Add(delta int) //增加 waitGroup 子 goruntine 计数值
func (wg *WaitGroup) Done() //当子 goruntine 完成后，将计数器-1
func (wg *WaitGroup) Wait() //调用此方法的 goruntine，阻塞等待计数值为 0

以下方法去除了 race 竞争检查的源代码。

Add

操作 counter 计数值加减。

当 counter 增加时，直接 return
当 counter 减少时, 判断条件：counter > 0 || waiter == 0
- true 时，直接 return
- false （等待线程都完成且有等待者）时，statep 复位为 0，通过 semap 信号量唤醒所有等待者

func (wg *WaitGroup) Add(delta int) {
	//从数组中拿到 stetep （ counter+waiter 的组合）和 semap 信号量的内存地址
	statep, semap := wg.state()
	//stetep 原子加操作，高位 32bit 是 counter,实际 counter+1
	state := atomic.AddUint64(statep, uint64(delta)<<32)
	//state 的高位 32bit，表示 couter 的计数值
	v := int32(state >> 32)
	//state 的低位 32bit，表示 waiter 的等待者数量
	w := uint32(state)
	// couter 不能小于 0
	if v < 0 {
		panic("sync: negative WaitGroup counter")
	}
	// 需要避免错误操作：Add 和 Wait 并发操作，否则会 panic
	if w != 0 && delta > 0 && v == int32(delta) {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}
	// 如果还有等待线程未完成或者并没有等待者，直接 return
	if v > 0 || w == 0 {
		return
	}
	// 需要避免错误操作：Add 和 Wait 并发操作，否则会 panic
	if *statep != state {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}
	// 将 statep 复位为 0 （ counter 和 waiter 都置为 0 ）
	*statep = 0
	// 有多少个等待者就往 semap 循环发信号量（其实就是 semap+1 ），Wait 等待有一个调用	// runtime_Semacquire(semap)就是在等待这个信号量
	for ; w != 0; w-- {
		runtime_Semrelease(semap, false, 0)
	}
}

Done

被等待线程完成后调用 Done，将 counter 计数-1，表示线程结束

func (wg *WaitGroup) Done() {
	wg.Add(-1)
}

Wait

主线程循环对 waiter 原子操作+1 直到成功后，然后阻塞等待 semap 信号量而被唤醒，最后 return

func (wg *WaitGroup) Wait() {
	// 从数组中拿到 stetep （ counter+waiter 的组合）和 semap 信号量的内存地址
	statep, semap := wg.state()
	for {
		//从内存总线中加载最新的 statep 值
		state := atomic.LoadUint64(statep)
		//state 的高位 32bit，表示 couter 的计数值
		v := int32(state >> 32)
		//state 的低位 32bit，表示 waiter 的等待者数量
		w := uint32(state)
		//如果 couter 为 0，表示当前已经没有在运行的等待线程了
		if v == 0 {
			return
		}
		// CAS 操作 statep+1，低位属于 waiter,即 waiter+1
		if atomic.CompareAndSwapUint64(statep, state, state+1) {
			// CAS 操作成功后，阻塞等待 semap 信号为非零，竞争到会将 semap-1，并唤醒线程
			runtime_Semacquire(semap)
			if *statep != 0 {
				panic("sync: WaitGroup is reused before previous Wait has returned")
			}
			return
		}
		// CAS 操作失败了，重新进入循环
	}
}

Go 之 WaitGroup 底层实现