Go同步原语之sync/waitGroup ---- 等待组

基本使用

sync.WaitGroup 是 Go 语言标准库中 sync 包提供的一种同步原语，用于等待一组 goroutine 完成任务，是一个相当常用的数据结构。

它常用于等待一批并发操作完成后再继续执行主程序，以确保并发操作完成后再进行后续的操作。

在 sync.WaitGroup（等待组）类型中，每个 sync.WaitGroup 值在内部维护着一个计数器(此计数的初始默认值为0)，sync.WaitGroup 通过计数器来追踪等待的 goroutine 数量，每个等待的 goroutine 在开始执行之前都会调用 Add 方法来增加计数器。当 goroutine 执行完任务后，应调用 Done 方法来减少计数器。主程序可以调用 Wait 方法来阻塞，直到计数器归零，即所有的 goroutine 都执行完毕。

sync.WaitGroup 提供了如下几个函数可用：

func (wg *WaitGroup) Add(delta int) {...}
func (wg *WaitGroup) Done() {...}
func (wg *WaitGroup) Wait() {...}

• Add(delta int)：增加计数器的值。 delta 表示要增加的计数器值，可以为正数或负数。
• Done()：减少计数器的值，相当于执行完一个任务。
• Wait()：阻塞主程序，直到计数器归零。一般在主程序中调用，用于等待所有任务完成。

以下是一个简单示例，展示了如何使用 WaitGroup 来等待一组并发任务完成：

package main

import (
	"fmt"
	"sync"
	"time"
)

func worker(id int, wg *sync.WaitGroup) {
	defer wg.Done() // 表示任务完成

	fmt.Printf("Worker %d started\n", id)
	time.Sleep(time.Second)
	fmt.Printf("Worker %d completed\n", id)
}

func main() {
	var wg sync.WaitGroup

	for i := 1; i <= 5; i++ {
		wg.Add(1) // 每个任务增加计数器
		go worker(i, &wg)
	}

	wg.Wait() // 等待所有任务完成
	fmt.Println("All workers have completed.")
}

在上面的示例中，我们创建了 5 个 goroutine，每个 goroutine 都执行 worker 函数。每个 worker 函数会通过 Done 方法来减少计数器，表示任务已经完成。主程序中的 wg.Wait() 会阻塞，直到所有任务都完成，然后输出 “All workers have completed.”。这样可以确保所有的任务都完成后再继续主程序的执行。

sync.WaitGroup 在并发编程中的使用场景主要是用于等待一组并发操作完成后再继续执行主程序。它能够有效地管理并发任务的同步和阻塞，避免主程序过早地结束或产生竞态条件。

结构体

WaitGroup 的结构体代码如下：

//go 1.20.3 path: sync/waitgroup.go
type WaitGroup struct {
	noCopy noCopy
	state atomic.Uint64
  sema  uint32
}

下面是对WaitGroup 的结构体每个字段的分析：

• noCopy ， noCopy 是 golang 源码中检测禁止拷贝的技术。如果程序中有 WaitGroup 的赋值行为，使用 go vet 检查程序时，就会发现有报错。但需要注意的是，noCopy 不会影响程序正常的编译和运行。

• sema，对应于 golang 中 runtime 内部的信号量的实现。WaitGroup 中会用到 sema 的两个相关函数：runtime_Semacquire 和 runtime_Semrelease，对应功能如下：

  • runtime_Semacquire 表示增加一个信号量，并挂起 当前 goroutine；
  • runtime_Semrelease 表示减少一个信号量，并唤醒 sema 上其中一个正在等待的 goroutine；

• state, state 是 WaitGroup 的状态数据字段，且是一个无符号64 bit的数据，内容包含了 counter , waiter 信息

  • counter 代表目前尚未完成的个数，WaitGroup.Add(n) 将会导致 counter += n, 而 WaitGroup.Done() 将导致 counter--；
  • waiter 代表目前已调用 WaitGroup.Wait 的 goroutine 的个数；

  state 字段示意图如下:

  

WaitGroup 的整个调用过程可以简单地描述成下面这样：

1. 当调用 WaitGroup.Add(n) 时，counter 将会自增: counter += n;

2. 当调用 WaitGroup.Wait() 时，会将 waiter++。同时调用 runtime_Semacquire(semap), 增加信号量，并挂起当前 goroutine;

3. 当调用 WaitGroup.Done() 时，将会 counter--。如果自减后的 counter 等于 0，说明 WaitGroup 的等待过程已经结束，则需要调用 runtime_Semrelease 释放信号量，唤醒正在 WaitGroup.Wait 的 goroutine;

WaitGroup 调用过程示意图:

源码解析

WaitGroup.Add

WaitGroup.Add 函数用于向 WaitGroup 中添加等待的 goroutine 数量。每次调用 Add 函数，计数器的值会增加，表示有多少个并发操作需要等待完成。这个函数通常在启动并发操作之前被调用，以指定需要等待的任务数量。

直接上源码：

//go 1.20.3 path: sync/waitgroup.go

func (wg *WaitGroup) Add(delta int) {
	......
	// 将state的高32位加上delta，低32位不变。即将Counter加上delta
	state := wg.state.Add(uint64(delta) << 32)
	// 取出state的高32位，即 Counter
	v := int32(state >> 32)
	// 取出state的低32位，即 Waiters
	w := uint32(state)
	......
	// 如果Counter小于0，说明Counter溢出了,直接panic
	if v < 0 {
		panic("sync: negative WaitGroup counter")
	}
	// 如果Waiters不为0，说明有Wait在等待，但是Counter又不为0，说明Add和Wait同时在执行，panic
	if w != 0 && delta > 0 && v == int32(delta) {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}
	// 如果Counter大于0或者Waiters为0，说明Counter还没减到0，不需要等待，直接返回
	if v > 0 || w == 0 {
		return
	}

	//如果wg.state.Load() != state，说明Add和Wait同时在执行，panic
	if wg.state.Load() != state {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}

	//重置Waiters为0
	wg.state.Store(0)
	//调用runtime_Semrelease唤醒释放所有的Waiter
	for ; w != 0; w-- {
		runtime_Semrelease(&wg.sema, false, 0)
	}
}

代码很简单，注释很清楚，在此就不多解释了。

WaitGroup.Wait

WaitGroup.Wait 方法用于阻塞主程序，直到计数器归零，也就是所有的等待任务都已经完成。这个方法是一种同步操作，主要用于确保所有等待的 goroutine 都已经执行完毕，然后主程序才能继续执行后续的操作。

直接上源码：

//go 1.20.3 path: sync/waitgroup.go

func (wg *WaitGroup) Wait() {
	......
	// 循环等待counter计数器为0
	for {
		// 读取state
		state := wg.state.Load()
		// 取出高32位,即counter计数器
		v := int32(state >> 32)
		// 取出低32位,即waiters计数器
		w := uint32(state)
		
		// 如果counter计数器为0,则直接返回
		if v == 0 {
			......
			return
		}
		// 如果counter计数器不为0,则将waiters计数器加1
		if wg.state.CompareAndSwap(state, state+1) {
			......
			// 调用runtime_Semacquire函数,该函数会将当前goroutine加入到等待队列中
			runtime_Semacquire(&wg.sema)
			// 如果counter计数器不为0,则panic
			if wg.state.Load() != 0 {
				panic("sync: WaitGroup is reused before previous Wait has returned")
			}
			......
			return
		}
	}
}

代码很简单，注释很清楚，在此就不多解释了。

WaitGroup.Done

WaitGroup.Done 方法用于减少计数器的值，表示一个等待的任务已经完成。每个等待的 goroutine 在执行完任务后，都应该调用这个方法，以通知 WaitGroup 计数器减少一个。主要作用是用于在等待任务完成后进行信号通知。

直接上源码：

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

这个只是WaitGroup.Add 的简单封装，参考WaitGroup.Add 代码注释。

总结

• WaitGroup 可以用于一个 goroutine 等待多个 goroutine 干活完成，也可以多个 goroutine 等待一个 goroutine 干活完成，是一个多对多的关系
• Add(n>0) 方法应该在启动 goroutine 之前调用，然后在 goroution 内部调用 Done 方法
• WaitGroup 必须在 Wait 方法返回之后才能再次使用
• Done 只是 Add 的简单封装，所以实际上是可以通过一次加一个比较大的值减少调用，或者达到快速唤醒的目的。

参考：

cyhone https://zhuanlan.zhihu.com/p/344973865

[mohuishou] (https://lailin.xyz/)https://lailin.xyz/post/go-training-week3-waitgroup.html