面试官：谈谈 Go WaitGroup 实现原理

大家好，我是木川

一、基本概念

Go 标准库提供了WaitGroup原语, 可以用它来等待一批 Goroutine 执行完成

二、基本使用

func main() {
    var wg sync.WaitGroup
    for i := 1; i <= 5; i++ {
        wg.Add(1)
        go func() {
            defer wg.Done()
            println("hello")
        }()
    }

    wg.Wait()
}

在WaitGroup里主要有3个方法：

• WaitGroup.Add()：添加需要等待的 goroutine 数量，_Add(n)__ 将会导致 请求计算器 counter += n

• WaitGroup.Done()：通知 WaitGroup goroutine已完成，相当于Add(-1)，Done() 将导致 counter -=1，请求计数器counter为0 时通过信号量调用runtime_Semrelease唤醒waiter线程

• WaitGroup.Wait()：等待所有 goroutine 结束，waiter 线程数 waiter++，同时通过信号量调用 runtime_Semacquire(semap)阻塞当前 goroutine

三、实现原理

WaitGroup 结构体

sync.WaitGroup定义了一个结构体，伪代码如下，包含如下字段。

type WaitGroup struct {
 // 代表目前尚未完成的协程个数
 counter uint32
 // 目前已调用 WaitGroup.Wait 的 goroutine 的个数
 waiter  uint32
 // 对应于 golang 中 runtime 内部的信号量的实现
 // runtime_Semacquire 表示增加一个信号量，并挂起当前 goroutine
 // runtime_Semrelease 表示减少一个信号量，并唤醒 sema 上其中一个正在等待的 goroutine
 sema    uint32
}

Add 方法

当你调用Add方法时，counter的值会增加。这表示有多少个goroutine需要等待。通常，在每个goroutine中，你会调用Add(1)来表示这个goroutine需要被等待。

Done 方法

在每个goroutine的结束处，你应该调用Done方法，以减少counter的值。这表示一个goroutine已经完成了它的工作。

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

Wait 方法

Wait方法用于等待counter的值变为0，当counter的值不为0时，Wait方法会调用runtime_Semacquire函数实现阻塞当前goroutine，直到counter的值减少为0时被 runtime_Semrelease 唤醒。

四、总结

sync.WaitGroup通过counter的增减跟踪goroutine的完成状态，通过信号量实现阻塞和唤醒。这样，它能够确保所有的goroutine都完成后，主程序才能继续执行。这种实现方式简单而高效，是Go语言并发编程中的重要工具之一。

最后给自己的原创 Go 面试小册打个广告，如果你从事 Go 相关开发，欢迎扫码购买，目前 10 元买断，加下面的微信发送支付截图额外赠送一份自己录制的 Go 面试题讲解视频

如果对你有帮助，帮我点一下在看或转发，欢迎关注我的公众号