Go 标准库源码分析 - sync 的 WaitGroup

WaitGroup常用于多个goroutine协作，主要功能是阻塞等待一组goroutine完成。 

一、数据结构

type WaitGroup struct {
	noCopy noCopy

	state1 [3]uint32    // 用于存放任务计数器、等待者计数器和信号量 
}

WaitGroup采用64位的值来保存计数器，其中高32位为任务计数器，低32位为等待者计数器，另外用32位的值保存信号量。

WaitGroup在使用时需要64位的计数器进行原子操作，这要求计数器的地址是64位对齐的。如果是64位的操作系统，可直接取数组前两位元素作为计数器；如果是32位操作系统，state1可能不是64位对齐的地址，则将数组第一个元素作为信号量而后两个元素为计数器。

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

 

二、使用方法

1. Add

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")
	}
    // 在Add前执行Wait
	if w != 0 && delta > 0 && v == int32(delta) {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}
	if v > 0 || w == 0 {
		return
	}
	// 执行至此，说明任务数为0，且等待者数不为0
    // 原计数器与现在的不同，可能是Add和Wait同时执行
	if *statep != state {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}
	// 将计数器归零
	*statep = 0
    // 唤醒等待者
	for ; w != 0; w-- {
		runtime_Semrelease(semap, false, 0)
	}
}

2. Done

func (wg *WaitGroup) Done() {
    // 执行Add方法将任务数减1
	wg.Add(-1)
}

3. Wait

func (wg *WaitGroup) Wait() {
    // 获取计数器和信号量地址
	statep, semap := wg.state()
	for {
        // 原子读计数器的值
		state := atomic.LoadUint64(statep)
        // 任务计数器
		v := int32(state >> 32)
        // 等待计数器
		w := uint32(state)
		if v == 0 {
			// 任务数为0，直接返回
			return
		}
		// 任务数不为0时，原子增加等待者数
		if atomic.CompareAndSwapUint64(statep, state, state+1) {
            // 挂起当前goroutine
			runtime_Semacquire(semap)
            // 在Add方法里，唤醒等待者前会将计数器归零，因此若此处的statep不为0，则说明又调用了Add或Wait方法，导致panic
			if *statep != 0 {
				panic("sync: WaitGroup is reused before previous Wait has returned")
			}
			return
		}
	}
}