并发编程WaitGroup 基本用法和如何实现以及常见错误

Posted 2021-12-08 了凡

tags:

篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了并发编程WaitGroup 基本用法和如何实现以及常见错误相关的知识，希望对你有一定的参考价值。

我是了凡，微信公众号【了凡银河系】期待你的关注，内有资源相送。未来大家一起加油啊~

前言

文章目录

前言
WaitGroup简单介绍
WaitGroup的基本用法
WaitGroup的实现
- Add 实现过程
- Wait 实现过程
使用WaitGroup时的常见错误
noCopy：辅助vet检查

WaitGroup简单介绍

WaitGroup就是package sync用来做任务编排的一个并发原语。这个要解决的就是并发-等待的问题：现有一个goroutine A在检查点（chaeckpoint）等待一组goroutine全部完成，如果在执行任务的这些goroutine还没有全部完成，那么goroutine A就会阻塞在检查点，直到所有的goroutine都完成后才能继续执行。

要完成一个大的任务，需要使用并行的goroutine执行三个小任务，只有这三个小任务都完成，才能去执行后面的任务。如果通过轮询的方式定时询问三个小任务是否完成，会存在两个问题:
一是，性能比较低，因为三个小任务可能早就完成了，却要等很长时间才被轮询到；二是，会有很多无谓的轮询，空耗CPU资源。所以，这个时候WaitGroup并发原语就比较有效了，它可以阻塞等待的goroutine。等三个小任务都完成了，再即时唤醒它们。

类似这样的并发原语很多操作系统和编程语言都有，例如：Linux中的barrier、Pthread(POSIX线程)中的barrier、C++中的std::barrier、Java中的CyclicBarrier和CountDowmLatch等。

WaitGroup的基本用法

Go标准库中的WaitGroup提供了三个方法，保持了Go简洁的风格。

func (wg *WaitGroup) Add(delta int)
func (wg *WaitGroup) Done()
func (wg *WaitGroup) Wait()

这个三个方法功能分别是：

Add，用来设置WaitGroup的计数值；
Dome，用来将WaitGroup的计数值减1，其实就是调用了Add(-1);
Wait，调用这个方法的goroutine会一直阻塞，直到WaitGroup的计数值变为0。

通过一个实例使用WaitGroup的例子看Add，Done，Wait方法的基本用法。使用计数器struct。启动10个worker，分别对计数值加一，10个worker都完成后，期望输出计数器的值。

// Counter 线程安全的计数器
type Counter struct {
   mu sync.Mutex
   count uint64
}

// Incr 对计数值加一
func (c Counter) Incr()  {
   c.mu.Lock()
   c.count++
   c.mu.Unlock()
}

// Count 获取当前的计数值
func (c Counter) Count() uint64  {
   c.mu.Lock()
   defer c.mu.Unlock()
   return c.count
}

// sleep 1秒，然后计数值加1
func worker(c *Counter, wg *sync.WaitGroup)  {
   defer wg.Done()
   time.Sleep(time.Second)
   c.Incr()
}

func main()  {
   var counter Counter

   var wg sync.WaitGroup
   wg.Add(10) // WaitGroup的值设置为10

   for i := 0; i < 10; i++ { // 启动10个goroutine执行加1任务
      go worker(&counter, &wg)
   }
   // 检查点，等待goroutine都完成任务
   wg.Wait()
   // 输出当前计数器的值
   fmt.Println(counter.Count())
}

使用WaitGroup编排这类任务的常用方式。而“这类任务”指的就是，需要启动多个goroutine执行任务，主goroutine需要等待子goroutine都完成后才能继续执行。

WaitGroup的实现

首先，看看WaitGroup的数据结构。它包括了一个noCopy的辅助字段，一个state1记录WaitGroup状态的数组。

noCopy的辅助字段，主要就是辅助vet工具检查是否通过copy赋值这个WaitGroup实例。
state1，一个具有复合意义的字段，包含WaitGroup的计数、阻塞再检查点的waiter数和信号量。

WaitGroup的数据结构定义以及state信息的获取方法如下：

type WaitGroup struct {
   // 避免赋值使用的一个技巧，可以告诉vet工具违反了赋值使用的规则
   // noCopy noCopy
   // 64bit(8bytes)的值分成两段，搞32bit是计数值，低32bit是waiter的计数
   // 另外32bit是用信号量的
   // 因为64bit值的原子操作需要64bit对齐，但是32bit编译器不支持，所以数组中的元素在不同的架构中不一样，具体处理看下面的方法
   // 总之，会找到对齐的那64bit作为state，其余的32bit做信号量
   state1 [3]uint32
}

// 得到state的地址和信号量的地址
func (wg *WaitGroup) state() (statep *uint64, semap *uint32) {
   if uintptr(unsafe.Pointer(&wg.state1))%8 == 0 {
      // 如果地址是64bit对齐的，数组前两个元素做state，后一个元素做信号量
      return (*uint64)(unsafe.Pointer(&wg.state1)), &wg.state1[2]
   }else {
      // 如果地址是32bit对齐的，数组后两个元素用来做state，它可以用来做64bit的原子操作，第一个元素32bit用来做信号量
      return (*uint64)(unsafe.Pointer(&wg.state1[1])), &wg.state1[0]
   }
}

因为对64位整数的原子操作要求证书的地址是64位对齐的，所以针对64位和32位环境的state字段的组成是不一样的。
在64位环境下，state1:

第一个元素是waiter数
第二个元素是WaitGroup的计数值
第三个元素是信号量

在32位环境下，如果state1不是64位对齐的地址，那么state1的第一个元素是信号量，后两个元素分别是waiter数和计数值。

继续深入源码，看一下Add、Done和Wait这三个方法的实现。

在这其中有几个检查非常关键，如果检查不通过，会出现panic，现在主要还是先看Add、Wait和Done本身的实现上。

看一下Add方法的逻辑。Add方法主要操作的是state的计数部分。可以为计数值增加一个delta值，内部通过原子操把这个值加到计数值上。但是，这个delta也可以是负数，相当于为计数值减去了一个值，Done方法内部其实就是通过Add(-1) 实现的。

Add 实现过程

func (wg *WaitGroup) Add(delta int) {
   statep, semap := wg.state()
   // 高32bit是计数值v，所以把delta左移32，增加到计数上
   state := atomic.AddUint64(statep, uint64(delta) << 32)
   v := int32(state >> 32) // 当前计数值
   w := uint32(state) // waiter count

   if v > 0 || w == 0 {
      return
   }

   // 如果计数值v为0并且waiter的数量w不为0，那么state的值就是waiter的数量
   // 将waiter的数量设置为0， 一位内计数值v也是0，所以它们俩的组合*statep直接设置为0即可。此时需要并唤醒所有的waiter
   *statep = 0
   for ; w != 0; w-- {
      runtime_Semrelease(semap, false, 0)
   }
}

// Done Done方法实际就是计数器减1
func (wg *WaitGroup) Done() {
   wg.Add(-1)
}

Wait 方法的实现逻辑是：不断检查state的值。如果期中的计数值变了0，那么说明所有的任务已完成，调用者不必再等待，直接返回。如果计数值大于0，说明此时还有任务没完成，那么调用者就变成了等待者，需要假如waiter队列，并且阻塞住自己。

Wait 实现过程

func (wg *WaitGroup) Wait()  {
   statep, semap := wg.state()

   for {
      state := atomic.LoadUint64(statep)
      v := int32(state >> 32) // 当前计数值
      w := uint32(state) // waiter的数量
      if v == 0 {
         // 如果计数值为0，调用这个方法的goroutine不必再等待，继续执行它后面的逻辑即可
         return
      }
      // 否则把waiter数量加1.期间可能有并发调用Wait的情况，所以最外层使用了for循环
      if atomic.CompareAndSwapUint64(statep, state, state+1) {
         // 阻止休眠等待
         runtime_Semrelease(semap)
         // 被唤醒，不再阻塞，返回
         return
      }
   }
}

使用WaitGroup时的常见错误

分析WaitGroup的Add、Done和Wait方法的实现的时候，为避免干扰，删除了异常检查的代码。因为这些异常检查非常有用。

在开发中，经常会遇见或看到误用WaitGroup的场景，究其原因就是没有搞明白这些检查的逻辑。接下来我们就一起来看三个问题。

常见问题一：计数器设置为负值

WaitGroup的计数器的值必须大于等于0。在更改这个计数值的时候，WaitGroup会先做检查，如果计数值被设置为负数，就会导致panic。
这里一般两种方法会导致计数器设置为负数。

第一个方法是：调用Add的时候传递一个负数。如果你能保证当前的计数器加上这个负数后还是大于等于0的化，也没有问题，否则就会对导致panic。

func main() {
    var wg sync.WaitGroup
    wg.Add(10)
    
    wg.Add(-10)// 将-10作为参数调用Add，计数值被设置为0
    
    wg.Add(-1)// 将-1作为参数调用Add，如果加上-1计数值就会变为负数。这是不对的，所以会触发panic
}

以上的问题就是计数器的初始值为10，当第一次传入-10的时候，计数值被设置为0，不会有啥问题。但是，紧接着传入-1以后，计数值就被设置为负数了，程序就会出现panic。

第二个方法是：调用Done方法的次数过多，超过了WaitGroup的计数值。

使用WaitGroup的正确姿势是，预先确定好WaitGroup的计数值，然后调用相同次数的Done完成相应的任务。例如，在WaitGroup变量声明之后，就立即设置它的计数值，或者在goroutine启动之前增加1，然后在goroutine中调用Done。

如果没有遵循这些规则，就很可能会导致Done方法调用的次数和计数值不一致，进而造成死锁（Done调用次数比计数值少）或者panic（Done调用次数比计数值多）。

不然就像以下情况，多调用了一次Donef方法后，会导致计数值为负，所以程序运行到这一行出现panic。

func main() {
    var wg sync.WaitGroup
    wg.Add(1)
    
    wg.Done()
    
    wg.Done()
}

运行结果：

常见问题二：不期望的Add时机

使用WaitGroup一定要遵守的原则就是，等所有的Add方法调用之后再调用Wait，否则就可能导致panic或者不期望的结果。

构造一个场景：只有部分的Add/Done执行完后，Wait就返回。这里我们启动四个goroutine，每个goroutine内部调用Add(1)然后调用Done()，主goroutine调用Wait等待任务完成。

func main() {
    var wg sync.WaitGroup
    go dosomething(100, &wg) // 启动第一个goroutine
    go dosomething(110, &wg) // 启动第二个goroutine
    go dosomething(120, &wg) // 启动第三个goroutine
    go dosomething(130, &wg) // 启动第四个goroutine
    
    wg.Wait() // 主goroutine等待完成
    fmt.Println("Done")
}

func dosomething(millisecs time.Duration, wg *sync.WaitGroup) {
    duration := millisecs * time.Millisecond
    time.Sleep(duration) // 故意sleep一段时间
    
    wg.Add(1)
    fmt.Println("后台执行，duration："，duration)
    wg.Done()
}

原本期望的是，等四个goroutine都执行接受后输出Done的信息，但是它的错误之处在于，将WaitGroup.Add方法的调用放在了子gorotuine中。等主goroutine调用Wait的时候，因为四个任务goroutine一开始都休眠，所以可能WaitGroup的Add方法还没有被调用，WaitGroup的计数还是0，所以它并没有等待四个子goroutine执行完毕才继续执行，而是立刻执行了下一步。

导致这个错误的原因就是，没有遵循先完成所有的Add 之后才Wait。要解决这个问题，一个方法是，预先设置计数值：

func main() {
    var wg sync.WaitGroup
    wg.Add(4) // 预先设定WaitGroup的计数值
    
    go dosomething(100, &wg) // 启动第一个goroutine
    go dosomething(110, &wg) // 启动第二个goroutine
    go dosomething(120, &wg) // 启动第三个goroutine
    go dosomething(130, &wg) // 启动第四个goroutine
    
    wg.Wait() // 主goroutine等待完成
    fmt.Println("Done")
}

func dosomething(millisecs time.Duration, wg *sync.WaitGroup) {
    duration := millisecs * time.Millisecond
    time.Sleep(duration) // 故意sleep一段时间
    
    fmt.Println("后台执行，duration："，duration)
    wg.Done()
}

第二个方法在启动子goroutine之前才调用Add：

func main() {
    var wg sync.WaitGroup
    
    go dosomething(100, &wg) // 调用方法，把计数值加1，并启动任务goroutine
    go dosomething(110, &wg) // 调用方法，把计数值加1，并启动任务goroutine
    go dosomething(120, &wg) // 调用方法，把计数值加1，并启动任务goroutine
    go dosomething(130, &wg) // 调用方法，把计数值加1，并启动任务goroutine
    
    wg.Wait() // 主goroutine等待，代码逻辑保证四次Add(1)都已经执行完了
    fmt.Println("Done")
}

func dosomething(millisecs time.Duration, wg *sync.WaitGroup) {
    wg.Add(1) // 计数值加1，再启动goroutine
    
    go func() {
        duration := millisecs * time.Millisecond
        time.Sleep(duration) // 故意sleep一段时间
        fmt.Println("后台执行，duration："，duration)
        wg.Done()
    }()
}

结论：无论怎么修复，都要保证所有的Add方法是再Wait方法之前被调用的。

常见问题三：前一个Wait还没有结束就重用WaitGroup

"前一个Wait还没有结束就重用WaitGroup"这一点不好理解。举个例子也许就懂了，再田径比赛的百米小组赛中，需要把选手分为计组，一组选手比赛之后，就可以进行下一组了。为了确保两组比赛上没有冲突，我们在模型化这个场景的时候，可以使用WaitGroup。

WaitGroup 等一组比赛的所有选手都跑完后 5 分钟，才开始下一组比赛。下一组比赛还可以使用这个 WaitGroup 来控制，因为 WaitGroup 是可以重用的。只要 WaitGroup 的计数值恢复到零值的状态，那么它就可以被看作是新创建的 WaitGroup，被重复使用。

但是，如果我们在WaitGroup的计数值还没有恢复到零值的时候就重用，就会导致程序panic。假如初始设置WaitGroup的计数值为1，启动一个goroutine先调用Done方法，接着就调用Add方法，Add方法有可能和主goroutine并发执行。

func main() {
    var wg sync.WaitGroup
    wg.Add(1)
    go func() {
        time.Sleep(time.Millisecond)
        wg.Done() // 计数器减1
        wg.Add(1) // 计数值加1
    }()
    wg.Wait() // 主goroutine等待，有可能和第7行并发执行
}

在第6行虽然让WaitGroup的计数恢复到0，但是因为第9行有个waiter在等待，如果等待Wait的goroutine，刚被唤醒就和Add调用（第7行）由并发执行的冲突，所以就会出现panic。

总结：WaitGroup虽然可以重用，但是是有一个前提的，那就是必须等到上一轮的Wait完成之后，才能重用WaitGroup执行下一轮的Add/Wait，如果在Wait还没执行完的时候就调用下一轮Add方法，就有可能出现panic。

noCopy：辅助vet检查

vet检查，功能和原理在：【并发编程】并发中互斥锁常见问题总结

vet会对实现Locker接口的数据类型做静态检查，一旦代码中有复制使用这种数据类型的情况，就会发出警告。但是WaitGroup同步原语不就是Add、Done和Wait方法吗？vet能检查出来么？

当然可以，但是需要通过WaitGroup添加一个noCopy字段，就可以为WaitGroup实现Locker接口，这样vet工具就可以做复制检查了。而且因为noCopy字段是未输出类型，所以WaitGroup不会暴露Lock/Unlock方法。

noCopy字段的类型noCopy是一个辅助帮助vet检查用的类型：

type noCopy struct{}

// Lock is a no-op used by -copylocks checker from `go vet`.
func (*noCopy) Lock() {}
func (*noCopy) Unlock() {}

如果想要自己定义的数据结构不被复制使用，或不能通过vet工具检查出复制使用的报警，就可以通过嵌入noCopy这个数据类型来实现。

这次就先讲到这里，如果想要了解更多的golang语言内容一键三连后序每周持续更新！

以上是关于并发编程WaitGroup 基本用法和如何实现以及常见错误的主要内容，如果未能解决你的问题，请参考以下文章

并发编程Cond 基本用法和如何实现以及常见错误

并发编程map 基本用法和常见错误以及如何实现线程安全的map类型

Golang 基础：底层并发原语 Mutex RWMutex Cond WaitGroup Once等使用和基本实现