golang WaitGroup的使用与底层实现

使用的go版本为 go1.21.2

首先我们写一个简单的WaitGroup的使用代码

package mainimport ("fmt""sync"
)func main() {var wg sync.WaitGroupwg.Add(1)go func() {defer wg.Done()fmt.Println("xiaochuan")}()wg.Wait()
}

WaitGroup的基本使用场景就是等待子协程完毕后，执行主协程，比如我的api需要多个下游api支持开多个协程进行访问，等待耗时最高的api返回过来后执行，这种场景是比较适合WaitGroup的。

我们来看一下WaitGroup构造体相关的底层源码

WaitGroup结构体

//代码位于 GOROOT/src/sync/waitgroup.go L:23type WaitGroup struct {//防止WaitGroup被复制, 君子协议，编译可以通过，某些编辑器会报waring//有兴趣可以看一下这里 https://github.com/golang/go/issues/8005#issuecomment-190753527noCopy noCopy// 高32位表示计数器，低32位表示等待的waiter数量。// 低版本go的state字段类型是[3]uint32，需要进行位数对齐state atomic.Uint64// 信号量sema  uint32
}

编辑器的warning

Add函数

//代码位于 GOROOT/src/sync/waitgroup.go L:43func (wg *WaitGroup) Add(delta int) {if race.Enabled { //使用竞态检查if delta < 0 { //如果传递的数值是负数，递减等待同步// Synchronize decrements with Wait.race.ReleaseMerge(unsafe.Pointer(wg))}race.Disable() //竞态检查 禁用defer race.Enable() //竞态检查 启用}//计算我们要进行add的值，将其加入到比特位上//<< 32 为二进制左位移 32位state := wg.state.Add(uint64(delta) << 32)v := int32(state >> 32) // state变量的高位是计数w := uint32(state) // state变量的低位是waiter计数//使用竞态检查，当前传入的值与v相同，说明当前是第一次调度addif race.Enabled && delta > 0 && v == int32(delta) {// The first increment must be synchronized with Wait.// Need to model this as a read, because there can be// several concurrent wg.counter transitions from 0.race.Read(unsafe.Pointer(&wg.sema))}//如果 计数器小于0 说明了多进行了done操作或者add传递负数，业务代码的出现逻辑错误了if v < 0 {panic("sync: negative WaitGroup counter")}// 如果当前存在等待，而且计数器不为0// 说明当前有地方调度了Wait后，又进行add操作了, 违反了官方的使用设计if w != 0 && delta > 0 && v == int32(delta) {panic("sync: WaitGroup misuse: Add called concurrently with Wait")}// 计数大于0，没有等待，就是单纯的add直接返回if v > 0 || w == 0 {return}// 再做一次检测，防止有并发调度// 比如我有两个goroutine A goroutine 在add, B goroutine 在调度 wait // 刚刚好A加完了计数，B突然wait导致state更变就会触发这个panicif wg.state.Load() != state {panic("sync: WaitGroup misuse: Add called concurrently with Wait")}// 重置waiter为0wg.state.Store(0)for ; w != 0; w-- { // 逐步释放信号量runtime_Semrelease(&wg.sema, false, 0)}
}

Done函数

//代码位于 GOROOT/src/sync/waitgroup.go L:86//这个很简单 调用了一下add函数传了一个-1
func (wg *WaitGroup) Done() {wg.Add(-1)
}

Wait函数

//代码位于 GOROOT/src/sync/waitgroup.go L:91func (wg *WaitGroup) Wait() {if race.Enabled { //使用竞态检查race.Disable() //竞态检查 禁用}for {state := wg.state.Load() // 原子操作读取state字段v := int32(state >> 32) // state变量的高位是计数w := uint32(state) // state变量的低位是waiter计数if v == 0 { // 如果当前计数器为0 就没必要等待直接返回了if race.Enabled {race.Enable() //竞态检查 启用race.Acquire(unsafe.Pointer(wg))}return}// 将waiter计数+1 因为waiter处于低32位所以不需要位移直接加就行了if wg.state.CompareAndSwap(state, state+1) {if race.Enabled && w == 0 { // 使用竞态检查，第一次进行wait操作// Wait must be synchronized with the first Add.// Need to model this is as a write to race with the read in Add.// As a consequence, can do the write only for the first waiter,// otherwise concurrent Waits will race with each other.race.Write(unsafe.Pointer(&wg.sema))}// 获取信号量，这行代码会进行G的阻塞runtime_Semacquire(&wg.sema)//重新获取一下state,正常来讲计数为0, waiter为0//执行判断之前，又有一个协程进行了add操作，会触发panicif wg.state.Load() != 0 {panic("sync: WaitGroup is reused before previous Wait has returned")}if race.Enabled { //使用竞态检查race.Enable() //竞态检查 启用race.Acquire(unsafe.Pointer(wg))}return}}
}