在 Go 语言中，Goroutine 是轻量级的协程，能够高效地实现并发编程。
然而，当我们在程序中启动多个 Goroutine 后，如何确保主线程在所有子任务完成后再退出？这就是 sync.WaitGroup 的用武之地。

为什么需要 WaitGroup？

假设我们需要同时执行多个任务，例如并行下载多个文件或批量处理数据。如果直接启动 Goroutine 后就退出主线程，主线程可能会在子任务完成前终止，导致程序异常退出。此时，WaitGroup 可以帮助我们同步多个 Goroutine 的执行状态，确保主线程等待所有任务完成后再继续执行。

WaitGroup 的核心方法

WaitGroup 是 Go 标准库 sync 包中的一个结构体，其核心方法如下：

示例代码：并行任务的等待

场景：模拟并行下载多个文件

package main

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

func downloadFile(name string, wg *sync.WaitGroup) {
    defer wg.Done() // 标记任务完成（必须放在 defer 中确保执行）

    fmt.Printf("Started downloading %s...\n", name)
    time.Sleep(2 * time.Second) // 模拟下载耗时
    fmt.Printf("Finished downloading %s\n", name)
}

func main() {
    var wg sync.WaitGroup

    files := []string{"file1.txt", "file2.txt", "file3.txt"}
    for _, file := range files {
        wg.Add(1) // 每启动一个 Goroutine，计数器+1
        go downloadFile(file, &wg)
    }

    wg.Wait() // 阻塞主线程，直到所有任务完成
    fmt.Println("All downloads completed!")
}

运行结果：

Started downloading file1.txt...
Started downloading file2.txt...
Started downloading file3.txt...
Finished downloading file1.txt
Finished downloading file2.txt
Finished downloading file3.txt
All downloads completed!

关键点解析

1. Add(1) 的时机：

   • 在启动 Goroutine 之前 调用 wg.Add(1)，确保计数器正确增加。
   • 如果 Goroutine 启动失败（如栈分配问题），可能导致计数器不一致，但这种情况极少发生。

2. Done() 的使用：

   • 建议将 wg.Done() 放在函数的 defer 中，确保即使发生 panic 也能正确减少计数器。

3. Wait() 的位置：

   • wg.Wait() 必须在所有 Goroutine 启动之后调用，否则可能阻塞主线程过早。

常见错误与注意事项

错误示例 1：忘记调用 Add()

func main() {
    var wg sync.WaitGroup
    go downloadFile("file.txt", &wg) // 没有调用 wg.Add(1)
    wg.Wait() // 死锁！因为计数器初始值为0
}

错误示例 2：多次调用 Wait()

wg.Wait() // 第一次等待，此时计数器可能已归零
wg.Wait() // 第二次调用会立即返回，但逻辑可能不正确

注意事项：

• 避免竞态条件：如果多个 Goroutine 同时修改 WaitGroup，需确保线程安全。
• 内存泄漏：如果 WaitGroup 的计数器无法归零（如忘记调用 Done()），程序会无限阻塞。

实际应用场景

• 批量数据处理：例如并行处理数据库查询结果。
• 微服务中的异步任务：在 HTTP 处理函数中启动多个后台任务，并等待关键任务完成后再返回响应。
• 测试并行操作：在测试中确保所有 Goroutine 执行完毕后再验证结果。

总结

sync.WaitGroup 是 Go 语言中实现并发控制的利器，能够帮助开发者优雅地管理 Goroutine 的生命周期。通过合理使用 Add()、Done() 和 Wait()，可以避免因 Goroutine 提前退出导致的程序崩溃或数据不一致问题。在实际开发中，结合 WaitGroup 和通道（Channel）等工具，可以进一步实现更复杂的并发逻辑。

希望这篇文章能帮助你更好地理解 WaitGroup 的核心思想和使用场景！如果有任何疑问，欢迎留言讨论。