第25课 - WaitGroup 与并发控制

学习目标

理解 WaitGroup 的作用与工作原理：明白 WaitGroup 是如何协调多个并发任务，实现“等待所有任务完成”这一关键同步操作的。
掌握 WaitGroup 的核心方法：熟练使用 Add、Done 和 Wait 方法来管理一组 goroutine。
能够编写正确的并发控制代码：避免因 Add 调用时机不当或计数器不匹配导致的程序崩溃（panic）。
识别并实践 WaitGroup 的典型应用场景：在需要并行执行多个独立任务并汇总结果时，正确应用 WaitGroup。

核心概念

在之前的课程中，我们学习了使用 sync.Mutex 来保护共享资源。今天，我们来学习另一个关键的同步原语：sync.WaitGroup。

想象一下，你是一个团队项目的经理。你需要团队中的 5 名成员（goroutines）各自去完成一项独立的任务（例如，查询不同数据库）。作为经理，你不能在所有人完成工作前就宣布项目结束。你需要一个机制来“等待”所有成员的工作报告，当所有人都完成后，你才能进行下一步（例如，汇总结果）。

sync.WaitGroup 就是 Go 语言中实现这种“等待所有并发任务完成”模式的标准工具。它维护一个内部计数器，这个计数器代表需要等待的“工作任务”的数量。

它的核心工作流程如下：

启动任务前：使用 Add(n) 方法，将计数器增加 n（n 通常等于你将要启动的 goroutine 数量）。
每个任务完成时：在每个 goroutine 函数体内，使用 Done() 方法，将计数器减一。Done() 等价于 Add(-1)。
等待所有任务完成：在主 goroutine 中调用 Wait() 方法。该方法会阻塞，直到内部计数器归零（即所有任务都调用了 Done()）。

下面我们通过代码来具体学习。

代码示例

示例1：基本用法 - 并发打印数字

package main

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

func main() {
	// 1. 创建一个 WaitGroup 实例
	var wg sync.WaitGroup

	// 我们要启动 5 个并发任务
	numTasks := 5

	// 2. 在启动任何 goroutine 之前，调用 Add 方法设置计数器
	wg.Add(numTasks)

	for i := 0; i < numTasks; i++ {
		// 启动一个 goroutine
		go func(id int) {
			// 3. 在 goroutine 函数体的最后，务必调用 Done
			// 使用 defer 是一个好习惯，确保即使函数中途 return 也会执行
			defer wg.Done()

			// 模拟耗时的工作
			fmt.Printf("Worker %d starting\n", id)
			time.Sleep(time.Second * time.Duration(id)) // 模拟不同耗时
			fmt.Printf("Worker %d done\n", id)
		}(i) // 注意：将循环变量 i 作为参数传入闭包
	}

	// 4. 主 goroutine 在此处阻塞，等待所有任务完成
	wg.Wait()
	fmt.Println("All workers have finished.")
}

预期输出 (顺序可能略有不同，但 “All workers have finished” 一定在最后):

Worker 0 starting
Worker 1 starting
Worker 2 starting
Worker 3 starting
Worker 4 starting
Worker 0 done
Worker 1 done
Worker 2 done
Worker 3 done
Worker 4 done
All workers have finished.

示例2：错误示范与修复

一个常见的错误是在 goroutine 内部调用 Add。这会导致主 goroutine 的 Wait() 可能在计数器增加前就返回，或者产生竞态条件。

package main

import (
	"fmt"
	"sync"
)

func main() {
	var wg sync.WaitGroup

	for i := 0; i < 5; i++ {
		go func(id int) {
			// 错误：在 goroutine 内部调用 Add
			// 此时主 goroutine 可能已经执行到 Wait()，且计数器为0，导致提前返回。
			wg.Add(1)
			defer wg.Done()
			fmt.Printf("Worker %d\n", id)
		}(i)
	}

	wg.Wait() // 可能这里 goroutine 还没开始 Add，Wait 就返回了。
	fmt.Println("Finished.")
}

修复后 (正确做法)：

// 将 Add 放在循环和 go 语句之前
wg.Add(5) // 或者放在循环内部，但必须在 go 语句之前
for i := 0; i < 5; i++ {
    go func(id int) {
        defer wg.Done()
        fmt.Printf("Worker %d\n", id)
    }(i)
}

实践练习

练习1：基础 - 并发计算平方和

编写一个程序，启动 5 个 goroutine，分别计算 1 到 5 的平方（1², 2², ..., 5²），并将结果发送到同一个 channel。主 goroutine 从 channel 中接收所有结果并求和。使用 WaitGroup 确保所有计算 goroutine 完成后再关闭 channel。

要求：

使用 WaitGroup 协调计算 goroutine 和关闭 channel 的动作。
最终输出总和：55 (因为 1+4+9+16+25=55)。

练习2：进阶 - 带错误处理的并发任务

模拟一个函数，它并发地从 3 个不同的“数据源”（使用随机睡眠模拟）获取数据。如果任何一个数据源在 500 毫秒内没有返回，则视为获取失败。

// 模拟从数据源获取数据
func fetchData(sourceID int) (string, error) {
    // 模拟不同的延迟
    time.Sleep(time.Duration(rand.Intn(1000)) * time.Millisecond)
    return fmt.Sprintf("Data from source %d", sourceID), nil
}

要求：

使用 WaitGroup 等待所有数据获取尝试完成。
使用 channel 收集成功的结果或错误信息。
最后，打印成功获取的数据和遇到的错误数量。

练习3：综合 - 并发文件行数统计

编写一个函数，接收一个字符串切片（代表文件路径），并发地统计每个文件的行数。要求：

对每个文件路径启动一个 goroutine。
使用 WaitGroup 等待所有 goroutine 完成。
使用一个并发安全的结构（例如带锁的 map 或 channel）来存储 文件路径 -> 行数 的映射。
最后，打印每个文件的行数。

常见错误

在 goroutine 内部调用 Add：这是最典型的错误。必须在启动 goroutine 之前调用 Add，或者至少确保 Wait() 调用之前 Add 已经执行。否则，Wait() 可能立即返回，因为计数器初始为 0。
Add 和 Done 计数不匹配：调用 Add(n) 后，必须确保最终会有恰好 n 次 Done() 被调用。如果 Done() 调用少了，Wait() 会永远阻塞；如果 Done() 调用多了，计数器会变为负数，引发 panic。使用 defer wg.Done() 是避免遗漏的好方法。
复制 WaitGroup：sync.WaitGroup 是一个结构体，通常通过指针传递。在函数间传递时，确保传递的是同一个 WaitGroup 实例，而不是副本。副本的状态是独立的，会导致同步失效。
在 Wait() 之后操作 WaitGroup：Wait() 返回后，WaitGroup 的内部状态可能不明确。官方文档建议，当 Wait 返回后，WaitGroup 可以被重用，但必须确保之前的 Add 调用全部匹配了 Done。为了代码清晰，通常建议一个 WaitGroup 实例只用于一组任务。

小结

在本课中，我们学习了 Go 并发控制中至关重要的 sync.WaitGroup：

核心作用：它用于等待一组 goroutine 完成。是协调并发任务生命周期的利器。
三大金刚方法：Add(delta) 增加等待计数，Done() 减少计数（通常用 defer 调用），Wait() 阻塞直到计数归零。
黄金法则：在启动 goroutine 之前调用 Add，并在 goroutine 内第一行使用 defer wg.Done()。
典型场景：适用于“启动一批任务，等它们全部完成后再继续”的场景，如并发网络请求、并行数据处理等。

掌握了 WaitGroup，你就拥有了构建可靠并发程序的核心组件之一。下一课，我们将学习 sync.Once 和 atomic 包，它们分别用于确保函数只执行一次和进行高效的原子操作，进一步丰富你的并发工具箱。

WaitGroup 与并发控制