Skip to content

4.5 学习 Go 协程:WaitGroup

在前两篇文章里,我们学习了 协程信道 的内容,里面有很多例子,当时为了保证 main goroutine 在所有的 goroutine 都执行完毕后再退出,我使用了 time.Sleep 这种简单的方式。由于写的 demo 都是比较简单的, sleep 个 1 秒,我们主观上认为是够用的。

但在实际开发中,开发人员是无法预知,所有的 goroutine 需要多长的时间才能执行完毕,sleep 多了吧主程序就阻塞了, sleep 少了吧有的子协程的任务就没法完成。因此,使用time.Sleep 是一种极不推荐的方式,今天主要就要来介绍 一下如何优雅的处理这种情况。

一、WaitGroup 简介

sync.WaitGroup(等待组)主要用于并发控制,它通过一个计数器来实现对一组 goroutine 的同步。主 goroutine 使用 WaitGroup 来等待所有子 goroutine 完成任务后再继续执行或退出。

核心思想:

  1. 启动 goroutine 前: 设置或增加计数器的值。
  2. 子 goroutine 中: 任务完成后,减少计数器的值。
  3. 主 goroutine 中: 阻塞等待,直到计数器归零。

二、三个核心方法

WaitGroup 并不是一个结构体指针,可以直接作为值类型使用,但通常通过传递其指针 *sync.WaitGroup 来避免不必要的拷贝,尤其是在函数传参时。

方法作用计数器变化描述
Add(delta int)增加/减少计数计数器 += delta设置或增加等待组的计数器。在启动 goroutine 之前调用。如果 delta 是负数,则减少计数。
Done()减少计数计数器 -= 1相当于 Add(-1)。通常放在子 goroutine 的 defer 语句中,确保任务完成后计数器减一。
Wait()阻塞等待保持不变阻塞当前的 goroutine(通常是主 goroutine),直到等待组的计数器归零。

三、基本使用示例

以下是一个完整的示例,展示了如何使用 WaitGroup 来确保主程序等待所有子 goroutine 打印完数字后再退出。

go
package main

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

// 1. 定义一个 WaitGroup 变量
var wg sync.WaitGroup

// 工作函数:模拟一个耗时操作
func worker(id int) {
    // 3. 任务完成后,调用 Done() 减少计数器
    // defer 确保 worker 函数无论如何退出(正常或 panic),Done() 都会被调用
    defer wg.Done() 

    fmt.Printf("Worker %d 正在开始工作...\n", id)
    time.Sleep(time.Second) // 模拟工作耗时 1 秒
    fmt.Printf("Worker %d 完成工作。\n", id)
}

func main() {
    // 假设我们要启动 5 个 worker goroutine
    numWorkers := 5

    // 2. 设置计数器的值:启动 n 个 goroutine,就 Add(n)
    wg.Add(numWorkers)

    // 启动 goroutine
    for i := 1; i <= numWorkers; i++ {
        go worker(i) 
    }

    fmt.Println("主 goroutine 正在等待所有 worker 完成...")

    // 4. 阻塞等待:直到计数器变为 0 
    wg.Wait() 

    // 此时所有 worker 都已完成
    fmt.Println("所有 worker 已完成,主 goroutine 退出。")
}

输出结果:

主 goroutine 正在等待所有 worker 完成...
Worker 1 正在开始工作...
Worker 2 正在开始工作...
Worker 3 正在开始工作...
Worker 4 正在开始工作...
Worker 5 正在开始工作...
Worker 5 完成工作。
Worker 4 完成工作。
Worker 3 完成工作。
Worker 2 完成工作。
Worker 1 完成工作。
所有 worker 已完成,主 goroutine 退出。

四、使用注意事项

1. Add() 必须在 Wait() 之前调用

你必须在启动 goroutine 之前或在确定 goroutine 已启动但尚未调用 Done() 之前调用 Add()

  • 错误情况: 如果你在 Wait() 之后才调用 Add(),可能会导致:
    • 如果 Wait() 已经返回,后续的 Add() 将是无效的。
    • 如果 Wait() 正在等待,同时 Add() 被调用,可能会导致死锁(如果计数器变为非零后,没有 goroutine 会调用 Done())。

2. Done() 必须与 Add() 匹配

Done() 调用必须严格匹配 Add() 设置的计数。

  • 计数器小于零: 如果 Done()Add(-1) 使得计数器小于零,程序会触发 panic

3. 推荐使用 defer 调用 Done()

为了保证 goroutine 无论成功或失败都能正确通知 WaitGroup,强烈推荐将 wg.Done() 放在 goroutine 函数的开头,使用 defer 关键字修饰。

go
func worker(id int, wg *sync.WaitGroup) {
    defer wg.Done() // 确保无论发生什么,计数器都会减一
    // ... 任务代码
}

4. 传参时使用指针 *sync.WaitGroup

虽然 WaitGroup 是一个结构体,但在函数传参时,应传入其指针 *sync.WaitGroup

原因: 传入值拷贝会导致每个 goroutine 接收到的是 WaitGroup 的副本,对副本调用的 Done() 无法影响到主 goroutine 中 Wait() 所依赖的原始 WaitGroup 计数器,从而造成程序提前退出或永久阻塞(死锁)。

正确写法:

go
func main() {
    var wg sync.WaitGroup
    wg.Add(1)
    go worker(&wg) // 传递指针
    wg.Wait()
}

func worker(wg *sync.WaitGroup) { // 接收指针
    defer wg.Done()
}

五、WaitGroupChannel 的区别

WaitGroupChannel 都可以实现并发同步,但它们的作用和关注点不同:

特性sync.WaitGroupchannel(通道)
主要作用同步等待一组任务完成。数据通信,在 goroutine 间安全地传递数据。
功能计数器模型,仅用于通知任务完成状态。队列模型,用于传输数据,也具备同步能力(阻塞)。
通信方向单向:从子 goroutine 通知主 goroutine。双向:可在任意两个 goroutine 间发送和接收数据。
关注点任务数量,确保所有任务都执行完成。数据流转,确保数据的安全和有序传输。

总结:

  • 如果你的需求是等待所有并发任务完成,而不涉及任务之间的数据传递,请使用 sync.WaitGroup
  • 如果你的需求是在并发任务之间安全地交换数据,请使用 channel