【源码解剖】Go任务编排助手sync.WaitGroup

WaitGroup定义

Go官方在 sync 包中提供的任务编排工具，用于为一个任务拆分为N个任务时，可以创建N个协程来处理任务，并阻塞线程直到这N个协程处理完毕

WaitGroup使用

go 官方提供了三个方法分别是

• Add(delta int) 为WaitGroup的计数器添加指定数值，计数器为0时 Wait()方法则放行goroutine不会阻塞，如果计数器的值为负数时，会panic() sync: negative WaitGroup counter
• Done() 为WaitGroup的计数器值减少1，表示当前协程完成了任务
• Wait() wait方法在调用时会阻塞，直到WaitGroup计数器的值减少至0，才会停止阻塞

Wait Group的基本使用方式

1. 推荐写法

package main
func main(){
	var wg sync.WaitGroup
	var strs = make([]string, 5)
	wg.Add(5) // 此时WaitGroup的计数器值为5
	for i := 0; i < 5; i++ {
		i := i
		go func() {
			defer wg.Done() // 当一个goroutine完成任务时，调用Done方法，计数器-1
			strs[i] = fmt.Sprintf("%v togo ..", i+1)
		}()
	}
	wg.Wait() // 阻塞，直至全部goroutine完成任务(计数器归零)
	fmt.Println(strs) // [1 togo .. 2 togo .. 3 togo .. 4 togo .. 5 togo ..]
}

2. 在for时添加计数器

var wg sync.WaitGroup
	var strs = make([]string, 5)
	for i := 0; i < 5; i++ {
		wg.Add(1) // for循环过程中，为WaitGroup计数器累加1
		i := i
		go func() {
			// wg.Add(1) 切忌不能将Add方法写至goroutine中。假设第一个goroutine执行完毕，其它goroutine还未开始任务，计数器会直接归零，Wait放行，则其余4个任务无法执行
			defer wg.Done()
			strs[i] = fmt.Sprintf("%v togo ..", i+1)
		}()
	}
	wg.Wait()
	fmt.Println(strs)

WaitGroup实现原理

首先看一下WaitGroup的结构体定义

type WaitGroup struct {
	noCopy noCopy // 辅助字段，辅助vet工具检查是否复制WaitGroup实例的，忽略
	state atomic.Uint64 // 计数器本器
	sema  uint32 // 信号量，用于唤醒waiter
}

Add方法的实现

func (wg *WaitGroup) Add(delta int) {
	state := wg.state.Add(uint64(delta) << 32) // 累加delta的值
	v := int32(state >> 32) // 右移32为，拿到计数器的值
	w := uint32(state) // waiter的数量
	if v < 0 {
		panic("sync: negative WaitGroup counter") // 计数器的值不应该为负数，不存在负的任务数量
	}
	if w != 0 && delta > 0 && v == int32(delta) {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait") // waiter在等待的情况下不应该并发的调取.Add方法，注: 必须要同时，很难复现这个panic
	}
	if v > 0 || w == 0 { // 添加成功，函数结束
		return
	}
	if wg.state.Load() != state { // waiter在等待不能并发add
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}
	// Reset waiters count to 0. 计数器清零，唤醒waiter
	wg.state.Store(0)
	for ; w != 0; w-- {
		runtime_Semrelease(&wg.sema, false, 0)
	}

1. 如果计数器的值大于0，或者waiter的数量是0，则不需要做处理直接返回成功
2. 如果计数器在清零时，还存在被阻塞的waiter，则唤醒它

Done方法

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

1. Done方法的实现非常简单，只是传入了-1让计数器去累加，属于辅助方法

Wait方法实现

func (wg *WaitGroup) Wait() {
	if race.Enabled {
		race.Disable() // 禁用race检测，忽略
	}
	for {
		state := wg.state.Load()// 拿到计数器状态
		v := int32(state >> 32) // 得到计数器具体的值
		w := uint32(state) // waiter数量
		if v == 0 { // 计数器归零，则全部goroutine都任务结束
			return
		}
		if wg.state.CompareAndSwap(state, state+1) {
			runtime_Semacquire(&wg.sema)
			if wg.state.Load() != 0 {
				panic("sync: WaitGroup is reused before previous Wait has returned")
			}
			return
		}
	}
}

1. Wait方法通过尝试检查state的状态，如果计数器的值是0，则直接返回，不会继续阻塞协程
2. 如果不为0，则把当前的goroutine放置到waiter中，如果CompareAndSwap成功，则进入runtime_Semacquite方法，等待被唤醒
3. 不为零，而且也无法放置成功，就进行循环检查，因为可能同时有多个Waiter调用Wait方法

需要注意的点

1. 不要在goroutine中调用.Add方法，因为可能会存在第一个goroutine运行过快，而导致Wait函数提前结束，纵使让其他goroutine无法执行
2. 计数器的值为负数的时候，则会报panic (sync: negative WaitGroup counter)
   计数器为负数的情况有两种

• 第一种 Add时放入了负数，从而panic

package main
import "sync"
func main(){
	var wg = sync.WaitGroup
	wg.Add(-1) // Add为负数
}

• 第二种 Done的次数和Add的次数不同，导致计数器为负数，从而panic

package main
import "sync"
func main(){
	var wg sync.WaitGroup
	wg.Add(1) // Add了一次
	wg.Done()
	wg.Done() 
}

完