Golang 中的 WaitGroup｜青训营笔记

这是我参与「第五届青训营 」伴学笔记创作活动的第 4 天

Golang 中的 sync 包提供了一系列用于实现 goroutine 同步的结构体和锁机制，其中包括 mutex、once 和 waitgroup。本文就来看看 Golang 中用于实现 goroutine 同步操作的 WaitGroup 数据结构。

示例

先来看一个使用的例子。

func process(i int, wg *sync.WaitGroup) {
	// ...
	
	wg.Done()
}

func main() {
	n := 3
	var wg sync.WaitGroup
	for i := 0; i < n; i++ {
		wg.Add(1)
		go process(i, &wg)
	}
	wg.Wait()
}

• main 函数运行，开启主 goroutine
• 通过 for 循环每次添加一个 goroutine 执行 process 函数，总计 3 个
• process 函数中逻辑执行完毕后调用 Done 方法发出结束信号
• 当 WaitGroup 通过 Wait 方法等待接收到与内部添加数量相同的结束信号时结束

WaitGroup 等待 goroutine 集合完成，主 goroutine 调用 Add 方法设置需要等待的 goroutine 数量 可以使用 Wait 方法阻塞，等待所有 goroutine 全部完成。

在实际应用中，对于复杂的多模块业务，可以使用 goroutine 并发执行。如果需要等待所有并发的 goroutine 结束才执行的逻辑，可以使用 WaitGroup 来管理 goroutine。

源码分析

本文基于 Go 1.19 版本。关于 WaitGroup 的实现部分在 sync 包下，这里展示的源码只保留了关键的代码片段。

结构体定义

type WaitGroup struct {  
    // 不允许复制
    noCopy noCopy  

    // 两个计数器，高 32 位为 counter，低 32 位为 waiter
    state1 uint64  
}

Add 方法

func (wg *WaitGroup) Add(delta int) {
    // 获取信号量
    statep, semap := wg.state()
    
    // 计算此时的 state
    state := atomic.AddUint64(statep, uint64(delta)<<32)
    v := int32(state >> 32)  // counter
    w := uint32(state)  // waiter
    if v < 0 {
        panic("sync: negative WaitGroup counter")
    }
    if v > 0 || w == 0 {
        return
    }
    
    if *statep != state {
        panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    }
    // v == 0 && w > 0，逐个释放
    *statep = 0
    for ; w != 0; w-- {
        runtime_Semrelease(semap, false, 0)
    }
}

Wait 方法

func (wg *WaitGroup) Wait() {  
    statep, semap := wg.state()  

    for {  
        // 获取 state
        state := atomic.LoadUint64(statep)  
        v := int32(state >> 32)  // counter
        w := uint32(state)  // waiter
        if v == 0 {  
            // 当计数器为 0，则无需等待  
            if race.Enabled {  
                race.Enable()  
                race.Acquire(unsafe.Pointer(wg))  
            }  
            return  
        }  
        // v > 0，陷入阻塞
        if atomic.CompareAndSwapUint64(statep, state, state+1) {  
            runtime_Semacquire(semap)  
        if *statep != 0 {  
            panic("sync: WaitGroup is reused before previous Wait has returned")  
        }
        return  
    }  
}

Done 方法

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

解读

上述源码中，v 就是 counter，w 就是 waiter。

• 执行 Add(n) 时，counter+=n
• 调用 Done() 时，counter-=1
• 当最后一个子 goroutine 结束后，counter 为 0，waiter 大于 0

方法	触发时机	Go 语言操作	系统操作	解释	具体描述
Add	counter == 0 && waiter > 0	runtime_Semrelease()	V	s+1	V 操作将 s 加 1，如果有任何线程阻塞在 P 操作等待 s 变为非零，那么 V 操作会唤醒这些线程中的一个，该线程将 s 减 1，完成 P 操作
Wait	counter > 0	runtime_Semacquire()	P	s-1	如果 s 非零，那么将 s 减 1，并立即返回。如果 s 为 0，则挂起该进程，直至 s 变为非零，等另一个执行 V 操作的线程唤醒该线程。在唤醒后，P 操作将 s 减 1，并将控制权还给调用者