深入浅出 Go 并发协同等待利器:sync.WaitGroup

2 阅读4分钟

前言

本文接下来要介绍的内容与 Go 语言中的 sync.WaitGroup 并发原语有关,它用于等待一组并发操作完成。如果你面临的场景中,需要将一个复杂任务划分为多个子任务,并等待这些子任务执行完毕后(无先后顺序的限制)才能继续后续操作,那么,sync.WaitGroup 是你解决这类场景问题的理想选择。

准备好了吗?准备一杯你最喜欢的咖啡或茶,随着本文一探究竟吧。

请在此添加图片描述

WaitGroup 的基本介绍

sync.WaitGroupGo 语言中的一个并发原语,它用于等待一组并发操作完成。下面是 sync.WaitGroup 的基本组成部分:

type WaitGroup struct {
    noCopy noCopy

    state atomic.Uint64 
    sema  uint32
}

func (wg *WaitGroup) Add(delta int) { // 省略... }

func (wg *WaitGroup) Done() { // 省略... }

func (wg *WaitGroup) Wait() { // 省略... }

sync.WaitGroup 结构体有三个字段:

  • noCopy noCopynoCopy 是一个辅助字段,它在运行中没有实际的作用。noCopy 字段用于辅助 Go vet 工具判断 WaitGroup 结构体是否被复制。当 WaitGroup 结构体被复制后,可能会导致非预期的错误发生。
  • state atomic.Uint64state 是一个计数字段,用于存储两种类型的数据。高 32 位表示计数器,即待执行完成的协程数量。低 32 位表示等待者数量,即调用 Wait() 方法后等待的协程数量。
  • sema uint32sema 表示信号量,用于阻塞和唤醒等待的 goroutine

sync.WaitGroup 结构体有三个方法:

  • Add(delta int):该方法用于添加或减少 WaitGroup 的计数器。当我们启动一个新协程执行任务时,应调用 Add(1)
  • Done():该方法用于减少 WaitGroup 的计数器,等价于 Add(-1)。当一个协程的任务执行结束后,应调用 Done()
  • Wait():该方法用于阻塞当前协程,直到其他协程都执行完成,即所有通过 Add 方法注册的协程都调用了 Done(),表示它们都已完成。

简单概括就是 WaitGroup 内部维护一个计数器,当调用 Add() 方法并传递正数时,计数器的值会增加,当调用 Done() 方法时,计数器的值会减少。调用 Wait() 方法会阻塞当前的 goroutine,直到计数器为 0,表示其他所有的 goroutine 都执行完成。

WaitGroup 的基本使用

下面通过一个示例展示如何使用 sync.WaitGroup 来管理多个任务的并发执行。

// https://github.com/chenmingyong0423/blog/blob/master/tutorial-code/go/waitgroup/main.go
package main

import (
    "fmt"
    "sync"
)

func main() {
    var wg sync.WaitGroup
    workers := []string{"陈明勇", "Mingyong Chen", "chenmingyong"}
    wg.Add(3)
    for i, worker := range workers {
        go Work(worker, i+1, &wg)
    }
    // 等待所有任务执行完毕
    wg.Wait()
    fmt.Println("所有任务执行完毕!")
}

func Work(name string, no int, wg *sync.WaitGroup) {
    defer wg.Done()
    fmt.Printf("%s 正在执行任务 %d...\n", name, no)
}

上述代码的执行结果如下所示:

chenmingyong 正在执行任务 3...
陈明勇 正在执行任务 1...
Mingyong Chen 正在执行任务 2...
所有任务执行完毕!

这段示例展示了如何使用 sync.WaitGroup 来确保所有并发任务都完成后程序才继续执行。

首先,通过调用 wg.Add(3) 方法为每个任务注册待完成的任务数。

然后启动多个 goroutine 执行任务。每个任务执行完成后调用 wg.Done() 方法标记任务完成,任务数量减一。

最后,主 goroutine 通过 wg.Wait() 等待所有任务完成后继续执行。这样确保了所有并发任务都执行完毕之后,程序才打印最终的消息。

使用 WaitGroup 时的一些注意事项

避免使用未归零的 WaitGroup 实例

WaitGroup 的计数器归零之前,应避免重新使用它。如果需要再次使用 WaitGroup,要等到所有的 Wait() 调用已经返回后才能进行,否则可能会发生不可预测的错误。

正确配对 Add 和 Done

确保每个 Add 调用之后,都有相应的 Done 调用来确保计数器的归零,防止 Wait() 永远阻塞。

正确地添加计数

在启动协程之前确保调用了 Add() 方法,如果是在启动协程之后调用 Add() 方法(即在新协程中调用 Add() 方法),该操作可能在 Wait() 方法执行后才开始,导致 Wait() 方法不会阻塞,而是直接返回。

避免计数为负数

如果调用 Done() 的次数超过了 Add() 的次数,程序将抛出 panic。因此需要保证调用 Add() 方法和 Done() 方法的次数一样。

确保 Done 被调用

我们应该使用 defer wg.Done() 来保证即使协程中途发生错误或提前退出,wg.Done() 方法也会被调用,从而防止死锁的形象发生。

小结

本文深入探讨了 Go 语言中并发协同等待利器 sync.WaitGroup,详细介绍了它的组成部分、基本用法以及在实际开发中需要注意的关键点。

虽然 sync.WaitGroup 的使用相对简单,但如果对计数器管理不当可能会发生不可预测的错误。

作者:陈明勇

个人网站:chenmingyong.cn

文章持续更新,如果本文能让您有所收获,欢迎点赞收藏加关注本掘金号。 微信阅读可搜《程序员陈明勇》。 这篇文章已被收录于 GitHub github.com/chenmingyon…,欢迎大家 Star 催更并持续关注。