在Go语言中,给我们提供了用于线程同步的sync.WaitGroup,简单来讲,WaitGroup就是指等待一组,等待一个系列执行完成后才会继续向下执行。
WaitGroup数据结构
type WaitGroup struct {
noCopy noCopy
state atomic.Uint64 // 高 32 bit 是计数值, 低 32 bit 是 waiter 的计数。
sema uint32
}
- noCopy 是一个辅助字段,用于辅助 vet 工具检查这个 WaitGroup 实例是否被复制。具体来说,如果在执行 go vet 时如果检测到 WaitGroup 实例被复制就会报错。但是执行 go run 是不会报错的。
- state 是一个复合字段,高 32 位是计数值,低 32 位是 waiter 的计数。
- sema 是信号量字段。
三个方法
func (wg *WaitGroup) Add(delta int)
func (wg *WaitGroup) Done()
func (wg *WaitGroup) Wait()
Add()
- Add方法,将传入的 delta 值通过原子操作加到 state 的计数部分。delta 值也可以为负数,相当于 state 值的计数部分减去 delta,但是 WaitGroup 会检查计数值,如果为负数,则会 panic;
- 如果内部计数器变为0,Wait方法阻塞等待的所有线程都会释放,如果计数器小于0,方法panic。
- 注意Add加上正数的调用应在Wait之前,否则Wait可能只会等待很少的线程。
Done()
Done方法减少WaitGroup计数器的值,应在线程的最后执行。实际上,Done 方法就是通过 Add(-1) 实现的。
Wait()
Wait方法阻塞直到WaitGroup计数器减为0,不断检查 state 值,如果其中的计数值为 0,说明所有的任务已完成,调用者不必阻塞等待,可以继续执行后面的逻辑;如果计数值大于 0,说明此时还有任务未完成,阻塞等待。
注意:
- WaitGroup 必须通过指针传递给函数;
- Add和Done函数一定要配对,否则可能发生死锁;
- 如果协程之间需要进行消息传递,或者一个协程需要等待另一个协程全部正确完成,不推荐这种用法。
实例
package main
import (
"fmt"
"sync"
"time"
)
// 注意,WaitGroup 必须通过指针传递给函数。
func worker(id int, wg *sync.WaitGroup) {
fmt.Printf("Worker %d starting\n", id)
// 睡眠一秒钟,以此来模拟耗时的任务。
time.Sleep(time.Second)
fmt.Printf("Worker %d done\n", id)
// 通知 WaitGroup ,当前协程的工作已经完成。
wg.Done()
}
func main() {
var wg sync.WaitGroup
// 开启几个协程,并为其递增 WaitGroup 的计数器。
for i := 1; i <= 5; i++ {
wg.Add(1)
go worker(i, &wg)
}
// 阻塞,直到 WaitGroup 计数器恢复为 0,即所有协程的工作都已经完成。
wg.Wait()
}