sync包的使用和原理之WaitGroup

161 阅读7分钟

1. 前言

在并发过程中,多个线程或 goroutine 可能同时操作同一内存区域,导致出现竞争问题。为保持内存一致性,Go 的 sync 包提供了常见的并发编程原语。其中包括:Mutex、RWMutex、WaitGroup、Once 和 Pool 等。

2. WaitGroup

WaitGroup 可用于等待一系列 goroutine 执行完成,主 goroutine 调用 Add 方法设置需要等待的 goroutine 数量,并发出去的 goroutine 在执行相关逻辑完成后,调用 Done 方法。同时 Wait 方法可用于阻塞 goroutine 直到所有并发的 goroutine 完成。

2.1 WaitGroup 基本使用

在主 goroutine 初始化 WaitGroup 对象,并使用 Add 方法设置需要等待的 goroutine 数量,并用 Wait 方法阻塞主 goroutine 等待并发的 goroutine 执行完毕,大体如下:

package main

import (
	"fmt"
	"sync"
)

func main() {
	wg := &sync.WaitGroup{}
	wg.Add(1)
	go func() {
		fmt.Println("do something")
		wg.Done()
	}()
	wg.Wait()
        fmt.Println("success")
}

打印输出结果如下:

do something
success

Wait 方法将阻塞主 goroutine, 使得在并发的 goroutine 执行完成后,才打印输出 success。修改程序,移除主程序调用 Wait 方法,修改程序如下:

package main

import (
	"fmt"
	"sync"
)

func main() {
	wg := &sync.WaitGroup{}
	wg.Add(1)
	go func() {
		fmt.Println("do something")
		wg.Done()
	}()
        fmt.Println("success")
}

观察输出如下:

success

主 goroutine 由于没有被 Wait 阻塞,没有等待并发的 goroutine 执行完成便早已打印输出了 success 信息。

2.2 WaitGroup 实现原理

编译生成对应的汇编文件,截取部分能容如下:

  main.go:7             0x5d9                   493b6610                CMPQ 0x10(R14), SP
  main.go:7             0x5dd                   7669                    JBE 0x648
  main.go:7             0x5df                   4883ec28                SUBQ $0x28, SP
  main.go:7             0x5e3                   48896c2420              MOVQ BP, 0x20(SP)
  main.go:7             0x5e8                   488d6c2420              LEAQ 0x20(SP), BP
  main.go:8             0x5ed                   488d0500000000          LEAQ 0(IP), AX          [3:7]R_PCREL:type.sync.WaitGroup
  main.go:8             0x5f4                   0f1f440000              NOPL 0(AX)(AX*1)
  main.go:8             0x5f9                   e800000000              CALL 0x5fe              [1:5]R_CALL:runtime.newobject<1>
  main.go:8             0x5fe                   4889442418              MOVQ AX, 0x18(SP)
  main.go:8             0x603                   48c70000000000          MOVQ $0x0, 0(AX)
  main.go:8             0x60a                   48c7400400000000        MOVQ $0x0, 0x4(AX)
  main.go:8             0x612                   488b442418              MOVQ 0x18(SP), AX
  main.go:8             0x617                   4889442410              MOVQ AX, 0x10(SP)
  main.go:9             0x61c                   bb01000000              MOVL $0x1, BX
  main.go:9             0x621                   e800000000              CALL 0x626              [1:5]R_CALL:sync.(*WaitGroup).Add
  main.go:10            0x626                   488b442410              MOVQ 0x10(SP), AX
  main.go:10            0x62b                   e800000000              CALL 0x630              [1:5]R_CALL:sync.(*WaitGroup).Done
  main.go:11            0x630                   488b442410              MOVQ 0x10(SP), AX
  main.go:11            0x635                   0f1f4000                NOPL 0(AX)
  main.go:11            0x639                   e800000000              CALL 0x63e              [1:5]R_CALL:sync.(*WaitGroup).Wait
  main.go:12            0x63e                   488b6c2420              MOVQ 0x20(SP), BP
  main.go:12            0x643                   4883c428                ADDQ $0x28, SP
  main.go:12            0x647                   c3                      RET
  main.go:7             0x648                   e800000000              CALL 0x64d              [1:5]R_CALL:runtime.morestack_noctxt
  main.go:7             0x64d                   eb8a                    JMP "".main(SB)

编译器将 sync.WaitGroup 转换成 type.sync.WaitGroup 对象,wg.Add,wg.Done,wg.Wait 也被分别编译成sync.(*WaitGroup).Add,sync.(*WaitGroup).Done和sync.(*WaitGroup).Wait

2.2.1 WaitGroup 结构定义

查看文件 /src/sync/waitgroup.go 对 WaitGroup 的结构定义如下所示:

type WaitGroup struct {
	noCopy noCopy

	state1 [3]uint32
}

里边包含两个属性,noCopy 和 state1 ,其中 noCopy 用于辅助帮助vet检查 WaitGroup 对象不可进行拷贝。state1 是一个长度为 3 的数组,里边包含三个信息内容:count、waiter 和 sema。

  • count: 表示还需要等待完成的 goroutine 数量
  • waiter: 表示有多少个 goroutine 在阻塞等待
  • sema: 信号量,用于阻塞 goroutine

在不同操作系统中,state1 每位表示的内容不一样,具体表示含义参考 state 方法:

func (wg *WaitGroup) state() (statep *uint64, semap *uint32) {
	if uintptr(unsafe.Pointer(&wg.state1))%8 == 0 {
		return (*uint64)(unsafe.Pointer(&wg.state1)), &wg.state1[2]
	} else {
		return (*uint64)(unsafe.Pointer(&wg.state1[1])), &wg.state1[0]
	}
}

在 64 位系统中,取数组前两位作为 WaitGroup 状态,最后一位作为信号量 sema,如图所示:

image.png

在 32 位系统中,取数组末两位作为 WaitGroup 状态,第一位作为信号量 sema,如图所示:

image.png

2.2.2 WaitGroup 的 Add 方法

func (wg *WaitGroup) Add(delta int) {
	statep, semap := wg.state()
	state := atomic.AddUint64(statep, uint64(delta)<<32)
	v := int32(state >> 32)
	w := uint32(state)
	// count 小于0,panic
	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")
	}
	if v > 0 || w == 0 {
		return
	}
	
	// 当 waiter 大于 0 时,此时 goroutine 已将 counter 置为 0
	// 现在不能有并发更改 state
	// * Add 不能和 Wait 同时发生
	// * 如果 count == 0,那么 Wait 不能增加
	// 做一个简单的健全检测,避免 WaitGroup 滥用
	if *statep != state {
		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
	}
	// 重置 waiter 和 count 为 0
	// 此时 v ==0 ,w != 0
	*statep = 0
	// 释放所有 waiter
	for ; w != 0; w-- {
		runtime_Semrelease(semap, false, 0)
	}
}

流程描述如下:

image.png

  • Add 方法添加 delta 到 count 可能为负数,对于 WaitGroup, count 不能为负数,否则会 panic;
  • 不能并发同时执行 Add 和 Wait 方法
  • 当 count 为 0 时,所有调用 Wait 被阻塞的 goroutine 将会释放。

2.2.3 WaitGroup 的 Done 方法

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

Done 方法只是调用 Add 方法将 count 数减一

WaitGroup.Done() 方法直接调用 WaitGroup.Add(-1)方法

2.2.4 WaitGroup 的 Wait 方法

// Wait 会阻塞 goroutine 直到 count 为 0
func (wg *WaitGroup) Wait() {
	statep, semap := wg.state()
	for {
		state := atomic.LoadUint64(statep)
		v := int32(state >> 32)
		w := uint32(state)
		// count == 0, 返回
		if v == 0 {
			return
		}
		// 增加 waiter 数
		if atomic.CompareAndSwapUint64(statep, state, state+1) {
			// 阻塞等待信号释放
			runtime_Semacquire(semap)
			// 在 Wait 返回之前被重用了,会引起 panic
			if *statep != 0 {
				panic("sync: WaitGroup is reused before previous Wait has returned")
			}
			return
		}
	}
}

调用流程如图所示:

image.png

  • 调用 WaitGroup.Wait 方法,会阻塞当前 goroutine,直到 WaitGroup 的 count 为 0
  • 在 Wait 信号释放之前,不能重用 WaitGroup
  • 通过循环 & CAS 的方式,进行 wait 数增加

2.3 非合理使用 WaitGroup 示例

2.3.1 count 被置为负数,引发 panic

方式一:调用 Add 方法多次增加负数,导致 count 小于零,操作如下:

package main

import (
	"fmt"
	"sync"
)

func main() {
	wg := &sync.WaitGroup{}
	wg.Add(2)
	fmt.Println("add 2 success")
	wg.Add(-2)
	fmt.Println("add -2 success")
	wg.Add(-2)
	fmt.Println("add -2 success")
}

打印输出如下:

add 2 success
add -2 success
panic: sync: negative WaitGroup counter

goroutine 1 [running]:
sync.(*WaitGroup).Add(0x10bfe00, 0xc00000e018)
        /usr/local/go/src/sync/waitgroup.go:74 +0x105
main.main()
        /Users/liangjinsi/Documents/workspace/wenote/go/source/sync/main.go:14 +0xd6
exit status 2

我们可以给 Add 方法传递负参数,第一次 wg.Add(-2) 后 count 不小于 0,程序能正常运行。在第二次 wg.Add(-2) 后,由于 count < 0, 此时引发 panic: sync: negative WaitGroup counter。

方式二:调用 Done 方法次数比调用 Add 方法次数多,操作如下:

package main

import (
	"fmt"
	"sync"
)

func main() {
	wg := &sync.WaitGroup{}
	wg.Add(1)
	fmt.Println("wg add success")
	wg.Done()
	fmt.Println("wg done success")
	wg.Done()
	fmt.Println("wg done twice success")
}

打印输出结果如下:

wg add success
wg done success
panic: sync: negative WaitGroup counter

goroutine 1 [running]:
sync.(*WaitGroup).Add(0x10bfe60, 0xc00000e018)
        /usr/local/go/src/sync/waitgroup.go:74 +0x105
sync.(*WaitGroup).Done(...)
        /usr/local/go/src/sync/waitgroup.go:99
main.main()
        /Users/liangjinsi/Documents/workspace/wenote/go/source/sync/main.go:14 +0xd7
exit status 2

可以发现,当第二次调用 wg.Done 后,此时 count 小于 0, 引发 panic: sync: negative WaitGroup counter 。

我们在实际操作过程中,应避免调用 Done 方法的次数超过其前边调用 Add 方法的次数。

4.3.2 Wait 执行后再 Add

WaitGroup 的使用原则是,在所有 Add 方法调用完之后再调用 Wait, 否则可能导致程序 panic 或不期望的结果。

package main

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

func main() {
	wg := &sync.WaitGroup{}
	go sleepThenAdd(10, wg)
	go sleepThenAdd(11, wg)
	go sleepThenAdd(12, wg)
	time.Sleep(11 * time.Millisecond)
	wg.Wait()
}

func sleepThenAdd(t int32, wg *sync.WaitGroup) {
	time.Sleep(time.Duration(t) * time.Millisecond)
	wg.Add(1)
	fmt.Println("add success ttl ", t)
	wg.Done()
}

输出结果如下:

add success ttl  10
add success ttl  11

第三次调用 sleepThenAdd 没有按预期输出。我们需要在调用 Wait 前调用完所有 Add 的方法,解决方案可通过预设 count 数,如下:

package main

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

func main() {
	wg := &sync.WaitGroup{}
	// 预设并发数
	wg.Add(3)
	go sleepThenAdd(10, wg)
	go sleepThenAdd(11, wg)
	go sleepThenAdd(12, wg)
	wg.Wait()
}

func sleepThenAdd(t int32, wg *sync.WaitGroup) {
	time.Sleep(time.Duration(t) * time.Millisecond)
	fmt.Println("add success ttl ", t)
	wg.Done()
}

或在调用并发前先调用Add,如下:

package main

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

func main() {
	wg := &sync.WaitGroup{}
	// 在并发前先调用 Add 方法
	wg.Add(1)
	go sleepThenAdd(10, wg)
	wg.Add(1)
	go sleepThenAdd(11, wg)
	wg.Add(1)
	go sleepThenAdd(12, wg)
	wg.Wait()
}

func sleepThenAdd(t int32, wg *sync.WaitGroup) {
	time.Sleep(time.Duration(t) * time.Millisecond)
	fmt.Println("add success ttl ", t)
	wg.Done()
}

最终打印结果如下所示:

add success ttl  10
add success ttl  11
add success ttl  12

4.3.3 前一个 Wait 还没结束,就重用 WaitGroup

WaitGroup 的重用,需在前一次 Wait 全部释放后才可以重用,否则会导致程序 panic。

package main

import (
	"fmt"
	"sync"
)

func main() {
	wg := &sync.WaitGroup{}
	wg.Add(1)
	fmt.Println("wg add success")
	go func() {
		wg.Done()
		fmt.Println("wg done success")
		wg.Add(1)
		fmt.Println("wg add again success")
	}()
	wg.Wait()
	fmt.Println("wg wait success")
}

打印输出如下:

wg add success
wg done success
panic: sync: WaitGroup is reused before previous Wait has returned

goroutine 1 [running]:
sync.(*WaitGroup).Wait(0x10c0080)
        /usr/local/go/src/sync/waitgroup.go:132 +0xa5
main.main()
        /Users/liangjinsi/Documents/workspace/wenote/go/source/sync/main.go:18 +0xbe
exit status 2
liangjinsi@liangjinsideMacBook-Pro sync % go run main.go
wg add success
wg done success
wg add again success
panic: sync: WaitGroup is reused before previous Wait has returned

goroutine 1 [running]:
sync.(*WaitGroup).Wait(0x10c0100)
        /usr/local/go/src/sync/waitgroup.go:132 +0xa5
main.main()
        /Users/liangjinsi/Documents/workspace/wenote/go/source/sync/main.go:19 +0xbe
exit status 2

可以发现,在 Wait 释放前重用的 WaitGroup, 引发 panic: sync: WaitGroup is reused before previous Wait has returned

修改程序在重用 WaitGroup 前等待 2 秒,以保证 wg.Wait 释放后再重用 WaitGroup, 修改如下:

package main

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

func main() {
	wg := &sync.WaitGroup{}
	wg.Add(1)
	fmt.Println("wg add success")
	go func() {
		wg.Done()
		fmt.Println("wg done success")
		time.Sleep(2 * time.Second)
		wg.Add(1)
		fmt.Println("wg add again success")
	}()
	wg.Wait()
	fmt.Println("wg wait success")
}

打印输出如下:

wg add success
wg done success
wg wait success

WaitGroup 的重用在 Wait 释放后,程序能够正常执行。