前言

在前面的文章中,我们使用过 WaitGroup 进行任务编排,Go语言中的 WaitGroup 和 Java 中的 CyclicBarrierCountDownLatch 非常类似。比如我们有一个主任务在执行,执行到某一点时需要并行执行三个子任务,并且需要等到三个子任务都执行完后,再继续执行主任务。那我们就需要设置一个检查点,使主任务一直阻塞在这,等三个子任务执行完后再放行。

说明:本文中的示例,均是基于Go1.17 64位机器

小试牛刀

我们先来个简单的例子,看下 WaitGroup 是怎么使用的。示例中使用 Add(5) 表示我们有 5个 子任务,然后起了 5个 协程去完成任务,主协程使用 Wait() 方法等待 子协程执行完毕,输出一共等待的时间。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
func main() {
	var waitGroup sync.WaitGroup

	start := time.Now()
	waitGroup.Add(5)
	for i := 0; i < 5; i++ {
		go func() {
			defer waitGroup.Done()
			time.Sleep(time.Second)
			fmt.Println("done")
		}()
	}

	waitGroup.Wait()
	fmt.Println(time.Now().Sub(start).Seconds())
}

/*
done
done
done
done
done
1.000306089
*/

总览

WaitGroup 一共有三个方法:

1
2
3
(wg *WaitGroup) Add(delta int)
(wg *WaitGroup) Done()
(wg *WaitGroup) Wait()
  • Add 方法用于设置 WaitGroup 的计数值,可以理解为子任务的数量
  • Done 方法用于将 WaitGroup 的计数值减一,可以理解为完成一个子任务
  • Wait 方法用于阻塞调用者,直到 WaitGroup 的计数值为0,即所有子任务都完成

正常来说,我们使用的时候,需要先确定子任务的数量,然后调用 Add() 方法传入相应的数量,在每个子任务的协程中,调用 Done(),需要等待的协程调用 Wait() 方法,状态流转如下图:

WaitGroup流程图

底层实现

结构体

1
2
3
4
5
type WaitGroup struct {
	noCopy noCopy // noCopy 字段标识,由于 WaitGroup 不能复制,方便工具检测

	state1 [3]uint32  // 12个字节,8个字节标识 计数值和等待数量,4个字节用于标识信号量
}

state1 是个复合字段,会拆分为两部分: 64位(8个字节)的 statep 作为一个整体用于原子操作, 其中前面4个字节表示计数值,后面四个字节表示等待数量;剩余 32位(4个字节)semap 用于标识信号量。

Go语言中对于64位的变量进行原子操作,需要保证该变量是 64位对齐 的,也就是要保证这 8个字节 的首地址是 8 的整数倍。因此当 state1 的首地址是 8 的整数倍时,取前8个字节作为 statep后4个字节作为 semap;当 state1 的首地址不是 8 的整数倍时,取后8个字节作为 statep前4个字节作为 semap

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
func (wg *WaitGroup) state() (statep *uint64, semap *uint32) {
	
	// 首地址是8的倍数时,前8个字节为 statep, 后四个字节为 semap
	if uintptr(unsafe.Pointer(&wg.state1))%8 == 0 {
		return (*uint64)(unsafe.Pointer(&wg.state1)), &wg.state1[2]
	} else { 
		
	// 后8个字节为 statep, 前四个字节为 semap	
		return (*uint64)(unsafe.Pointer(&wg.state1[1])), &wg.state1[0]
	}
}

Add

  • Add 方法用于添加一个计数值(负数相当于是减),当计数值变为0后, Wait 方法阻塞的所有等待者都会被释放
  • 计数值变为负数是非法操作,产生 panic
  • 当计数值为0时(初始状态),Add 方法不能和 Wait 方法并发调用,需要保证 Add 方法在 Wait 方法之前调用,否则会 panic
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
func (wg *WaitGroup) Add(delta int) {

	// 拿到计数值等待者变量 statep 和 信号量 semap
	statep, semap := wg.state()

	// 计数值加上 delta: statep 的前四个字节是计数值,因此将 delta 前移 32位
	state := atomic.AddUint64(statep, uint64(delta)<<32)

	// 计数值
	v := int32(state >> 32)

	// 等待者数量
	w := uint32(state)

	// 如果加上 delta 之后,计数值变为负数,不合法,panic
	if v < 0 {
		panic("sync: negative WaitGroup counter")
	}

	// delta > 0 && v == int32(delta) : 表示从 0 开始添加计数值
	// w!=0 :表示已经有了等待者
	// 说明在添加计数值的时候,同时添加了等待者,非法操作。添加等待者需要在添加计数值之后
	if w != 0 && delta > 0 && v == int32(delta) {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}

	// v>0 : 计数值不等于0,不需要唤醒等待者,直接返回
	// w==0: 没有等待者,不需要唤醒,直接返回
	if v > 0 || w == 0 {
		return
	}

	// 再次检查数据是否一致
	if *statep != state {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}

	// 到这里说明计数值为0,且等待者大于0,需要唤醒所有的等待者,并把系统置为初始状态(0状态)
  
  // 将计数值和等待者数量都置为0
	*statep = 0

	// 唤醒等待者
	for ; w != 0; w-- {
		runtime_Semrelease(semap, false, 0)
	}
}

Done

1
2
3
4
// 完成一个任务,将计数值减一,当计数值减为0时,需要唤醒所有的等待者
func (wg *WaitGroup) Done() {
	wg.Add(-1)
}

Wait

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
// 调用 Wait 方法会被阻塞,直到 计数值 变为0
func (wg *WaitGroup) Wait() {

	// 获取计数、等待数和信号量
	statep, semap := wg.state()

	for {
		state := atomic.LoadUint64(statep)

		// 计数值
		v := int32(state >> 32)

		// 等待者数量
		w := uint32(state)

		// 计数值数量为0,直接返回,无需等待
		if v == 0 {
			return
		}

		// 到这里说明计数值数量大于0
		// 增加等待者数量:这里会有竞争,比如多个 Wait 调用,或者在同时调用 Add 方法,增加不成功会继续 for 循环
		if atomic.CompareAndSwapUint64(statep, state, state+1) {
			// 增加成功后,阻塞在信号量这里,等待被唤醒
			runtime_Semacquire(semap)

			// 被唤醒的时候,应该是0状态。如果重用 WaitGroup,需要等 Wait 返回
			if *statep != 0 {
				panic("sync: WaitGroup is reused before previous Wait has returned")
			}
			return
		}
	}
}

易错点

上面分析源码可以看到几个会产生 panic 的点,这也是我们使用 WaitGroup 需要注意的地方

  1. 计数值变为负数
  • 调用 Add 时参数值传负数
1
2
3
4
5
6
func main() {
var wg sync.WaitGroup
wg.Add(1)
wg.Add(-1)
wg.Add(-1)
}
  • 多次调用 Done 方法
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
func main() {
	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		fmt.Println("test")
		wg.Done()
		wg.Done()
	}()

	time.Sleep(time.Second)
	wg.Wait()
}
  1. Add 和 Wait 并发调用

    AddWait 并发调用,有可能达不到我们预期的效果,甚至 panic。如下示例中,我们想要等待 3 个子任务都执行完后再执行主任务,但实际情况可能是子任务还没起来,主任务就继续往下执行了。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21

func doSomething(wg *sync.WaitGroup) {
	wg.Add(1)
	fmt.Println("do  something")
	defer wg.Done()
}

func main() {
	var wg sync.WaitGroup
	for i := 0; i < 3; i++ {
		go doSomething(&wg)
	}

	wg.Wait()
	fmt.Println("main")
}

//main
//do  something
//do  something

正确的使用方式,应该是在调用 Wait 前先调用 Add

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
func doSomething(wg *sync.WaitGroup) {
	defer wg.Done()
	fmt.Println("do  something")
}

func main() {
	var wg sync.WaitGroup
	wg.Add(3)
	for i := 0; i < 3; i++ {
		go doSomething(&wg)
	}

	wg.Wait()
	fmt.Println("main")
}

//do  something
//do  something
//do  something
//main

  1. 没有等 Wait 返回,就重用 WaitGroup
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
func main() {
	var wg sync.WaitGroup

	wg.Add(1)
	go func() {
		fmt.Println("do something")
		wg.Done()
		wg.Add(1)
	}()

	wg.Wait()
}
  1. 复制使用

我们知道 Go 语言中的参数传递,都是值传递,就会产生复制操作。因此在向函数传递 WaitGroup 时,使用指针进行操作。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

// 错误使用方式,没有使用指针
func doSomething(wg sync.WaitGroup) {
	fmt.Println("do  something")
	defer wg.Done()
}

func main() {
	var wg sync.WaitGroup
	wg.Add(3)
	for i := 0; i < 3; i++ {
		// 这里没使用指针,wg状态一直不会改变,导致 Wait 一直阻塞
		go doSomething(wg)
	}

	wg.Wait()
	fmt.Println("main")
}

总结

本篇文章,我们通过源码+示例的方式,一起学习了 sync.WaitGroup 实现逻辑,同时也给出了一些注意点,只要做到如下操作,就不会出现问题:

  1. 保证 Add 在 Wait 前调用
  2. Add 中不传递负数
  3. 任务完成后不要忘记调用 Done 方法,建议使用 defer wg.Done()
  4. 不要复制使用 WaitGroup,函数传递时使用指针传递
  5. 尽量不复用 WaigGroup,减少出问题的风险

更多

微信公众号:CodePlayer