Golang -- WaitGroup 解析

Hud.

于 2022-07-08 13:34:34 发布

阅读量539

点赞数

分类专栏： Golang 文章标签： golang 开发语言后端

原文链接：https://zhuanlan.zhihu.com/p/400106071

版权

Golang 专栏收录该内容

6 篇文章 0 订阅

订阅专栏

摘要

Golang 提供了简洁的 go 关键字来让开发者更容易的进行并发编程，同时也提供了 WaitGroup 对象来辅助并发控制。今天我们就来分析下 WaitGroup 的使用方法，顺便瞧一瞧它的底层源码。

WaitGroup 的使用场景和方法

当我们有很多任务要同时进行时，如果并不需要关心各个任务的执行进度，那直接使用 go 关键字即可。

如果我们需要关心所有任务完成后才能往下运行时，则需要 WaitGroup 来阻塞等待这些并发任务了。

WaitGroup 如同它的字面意思，就是等待一组 goroutine 运行完成，主要有三个方法组成：

Add(delta int) ：添加任务数
Wait()：阻塞等待所有任务的完成
Done()：完成任务

简单举例看下用法

package main

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

func worker(wg *sync.WaitGroup) {
 doSomething()
    wg.Done() // 2.1、完成任务
}

func main() {
    var wg sync.WaitGroup
    wg.Add(5) // 1、添加 5 个任务
    for i := 1; i <= 5; i++ {
        go worker(&wg) // 2、每个任务并发执行
    }
    wg.Wait() // 3、阻塞等待所有任务完成
}

WaitGroup 源码分析

上面 WaitGroup 的使用很简单，接下来我们到 src/sync/waitgroup.go 里分析下它的源码。首先，是 WaitGroup 的结构体：

type WaitGroup struct {
 noCopy noCopy
 state1 [3]uint32
}

noCopy

其中，noCopy 表示 WaitGroup 是不可复制的。那么什么叫不可复制呢？

举个例子，当我们对函数参数定义了这个不可复制的类型时，开发者只能通过指针来传递函数参数。而规定使用指针传递又有什么好处呢？

好处在于如果有多个函数都定义了这个不可复制的参数时，那么这多个函数参数就可以共用同一个指针变量，来同步执行结果。而 WaitGroup 就是需要这样的约束规定。

state1 字段

接下来我们来看看 WaitGroup 的 state1 字段。state1 是一个包含了 counter 总数、 waiter 等待数、sema 信号量的 uint32 数组。

每当有 goroutine 调用了 Wait() 方法阻塞等待时，就会对 waiter 数量 + 1，然后等待信号量的唤起通知。

当我们调用 Add() 方法时，就会对 state1 的 counter 数量 + 1。

当调用 Done() 方法时就会对 counter 数量 -1。

直到 counter == 0 时就可以通过信号量唤起对应 waiter 数量的 goroutine 了，也就是唤起刚刚阻塞等待的 goroutine 们。

关于信号量的解释，可以参考下 golang 重要知识：mutex 里的相关介绍：

PV 原语解释：
通过操作信号量 S 来处理进程间的同步与互斥的问题。
S>0：表示有 S 个资源可用；S=0 表示无资源可用；S<0 绝对值表示等待队列或链表中的进程个数。信号量 S 的初值应大于等于 0。
P 原语：表示申请一个资源，对 S 原子性的减 1，若减 1 后仍 S>=0，则该进程继续执行；若减 1 后 S<0，表示已无资源可用，需要将自己阻塞起来，放到等待队列上。
V 原语：表示释放一个资源，对 S 原子性的加 1；若加 1 后 S>0，则该进程继续执行；若加 1 后 S<=0，表示等待队列上有等待进程，需要将第一个等待的进程唤醒。

此处操作系统可以理解为 Go 的运行时 runtime，进程可以理解为协程。

源码解释

最后，我们来深入 WaitGroup 的三个方法，进行源码分析。大家感兴趣的可以继续往下看，主要是对源码的分析注释。

Add(delta int) 方法

func (wg *WaitGroup) Add(delta int) {
 statep, semap := wg.state()
 if race.Enabled { // 此处是 go 的竞争检测，可以不用关心
  _ = *statep
  if delta < 0 {
   race.ReleaseMerge(unsafe.Pointer(wg))
  }
  race.Disable()
  defer race.Enable()
 }
 state := atomic.AddUint64(statep, uint64(delta)<<32)
 v := int32(state >> 32) // 获取 counter
 w := uint32(state) // 获取 waiter
 if race.Enabled && delta > 0 && v == int32(delta) { // go 的竞争检测，可以不用关心
  race.Read(unsafe.Pointer(semap))
 }
 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 { // counter > 0：还有任务在执行；waiter == 0 表示没有在阻塞等待的 goroutine
  return
 }
 if *statep != state {
  panic("sync: WaitGroup misuse: Add called concurrently with Wait")
 }
 // 执行到此处相当于 countr = 0，即所有的任务都已执行完，需要唤起等待的 goroutine了
 *statep = 0
 for ; w != 0; w-- {
  runtime_Semrelease(semap, false, 0)
 }
}

Done 方法

func (wg *WaitGroup) Done() {
 wg.Add(-1) // 直接调用 Add 方法 对 counter -1
}

Wait 方法

func (wg *WaitGroup) Wait() {
 statep, semap := wg.state()
 if race.Enabled { // go 的竞争检测，可以不用关心
  _ = *statep
  race.Disable()
 }
 for {
  state := atomic.LoadUint64(statep)
  v := int32(state >> 32)
  w := uint32(state)
  if v == 0 {
   // counter 为 0， 不需要再等待了。
   if race.Enabled {
    race.Enable()
    race.Acquire(unsafe.Pointer(wg))
   }
   return
  }
  // waiters 数目 +1.
  if atomic.CompareAndSwapUint64(statep, state, state+1) {
   if race.Enabled && w == 0 {
    race.Write(unsafe.Pointer(semap)) // go 的竞争检测，可以不用关心
   }
   runtime_Semacquire(semap) // 阻塞等待唤起
   if *statep != 0 {
    panic("sync: WaitGroup is reused before previous Wait has returned")
   }
   if race.Enabled {
    race.Enable()
    race.Acquire(unsafe.Pointer(wg))
   }
   return
  }
 }
}

从这几个方法的源码，我们可以看出，Go 并没有使用 mutex 等锁去做字段值修改，而是采用了 atomic 原子操作来进行修改的。这是在底层硬件上支持的，所以性能更好。