等待组——WaitGroup-深入浅出Go语言核心编程(四)

当前位置:　首页>> 技术小册>> 深入浅出Go语言核心编程(四)

等待组——WaitGroup

在Go语言的并发编程中，sync.WaitGroup 是一个非常重要的同步机制，它用于等待一组协程（goroutines）的完成。WaitGroup 内部维护了一个计数器，用于记录当前还有多少个协程正在执行。每当一个新的协程启动时，可以通过 Add(delta int) 方法增加计数器的值；当协程执行完毕后，通过调用 Done() 方法（实际上是 Add(-1) 的便捷方式）来减少计数器的值。当计数器的值变为零时，所有等待该 WaitGroup 的协程（通常是通过调用 Wait() 方法实现的）将被唤醒并继续执行。这种机制非常适合于需要等待多个并行任务完成的场景。

一、WaitGroup 的基本用法

1.1 初始化

WaitGroup 可以通过声明一个 sync.WaitGroup 类型的变量来创建，但它不需要显式初始化，零值即可使用。然而，在使用之前，你需要通过调用 Add() 方法来设置或调整计数器的初始值。

var wg sync.WaitGroup

1.2 增加计数器

每当启动一个新的协程去执行某个任务时，都应该通过调用 Add(1) 来增加 WaitGroup 的计数器，表示有一个额外的协程需要被等待。

wg.Add(1)
go func() {
    defer wg.Done() // 或者直接使用 wg.Add(-1)，但通常使用 Done() 更清晰
    // 执行任务...
}()

1.3 减少计数器

每个协程在执行完任务后，应该调用 Done() 方法来减少计数器的值。这是通过延迟函数（defer）来实现的，以确保即使在发生错误或提前返回的情况下，计数器的值也能被正确减少。

defer wg.Done()

1.4 等待所有协程完成

主协程（或任何需要等待所有并发任务完成的协程）可以通过调用 Wait() 方法来阻塞，直到所有通过 Add() 方法增加的协程都通过调用 Done() 方法完成了它们的任务。

wg.Wait()
// 所有协程完成，继续执行后续代码...

二、WaitGroup 的高级用法与注意事项

2.1 计数器为负

如果 Add() 方法的参数是一个负数，并且这个操作会导致计数器的值变为负数，那么 WaitGroup 会引发一个 panic。因此，确保 Add() 的调用是正确且符合逻辑的，特别是在并发环境下。

2.2 重复调用 Add 和 Done

对同一个 WaitGroup 多次调用 Add() 或 Done() 是允许的，但这样做需要谨慎管理以避免计数器出现意外的值。通常，我们会在协程启动前调用 Add(1)，在协程结束时调用 Done()。如果在协程内多次调用 Add() 或 Done()，必须确保计数器的最终值是正确的。

2.3 复制与并发安全

WaitGroup 并不是并发安全的，指的是你不能安全地复制一个 WaitGroup 实例并在不同的协程中使用它的副本。每个 WaitGroup 实例都应该被独立地用于控制一组特定的协程。

2.4 使用场景

并行初始化：在启动多个协程并行地初始化数据或资源后，等待所有初始化完成再继续。
并发处理任务：将一个大任务分解为多个小任务，由不同的协程并行处理，最后等待所有任务完成。
资源清理：在多个协程共享资源（如文件句柄、网络连接等）时，确保所有协程都完成其工作后，再进行资源释放或关闭操作。

三、示例代码

下面是一个使用 WaitGroup 的简单示例，展示了如何并行地下载多个网页内容，并等待所有下载完成后输出一条消息。

package main
import (
    "fmt"
    "net/http"
    "sync"
)
func download(url string, wg *sync.WaitGroup) {
    defer wg.Done()
    // 模拟下载过程
    resp, err := http.Get(url)
    if err != nil {
        fmt.Println("Error downloading:", err)
        return
    }
    defer resp.Body.Close()
    // 这里可以添加处理响应体的代码
    fmt.Println("Downloaded:", url)
}
func main() {
    var wg sync.WaitGroup
    urls := []string{
        "http://example.com",
        "http://golang.org",
        "http://google.com",
    }
    for _, url := range urls {
        wg.Add(1)
        go download(url, &wg)
    }
    wg.Wait()
    fmt.Println("All downloads completed.")
}

在这个例子中，我们创建了一个 WaitGroup 实例 wg，并在一个循环中为每个 URL 启动了一个协程去下载内容。每个协程在开始执行时都会通过 wg.Add(1) 增加计数器的值，并在结束时通过 defer wg.Done() 减少计数器的值。主协程通过调用 wg.Wait() 等待所有下载任务完成，然后输出一条消息。

四、总结

sync.WaitGroup 是 Go 语言并发编程中一个非常实用的工具，它提供了一种简单而有效的方式来等待一组协程的完成。通过合理地使用 Add()、Done() 和 Wait() 方法，我们可以轻松地管理并发任务，确保程序的正确性和效率。然而，在使用 WaitGroup 时也需要注意避免一些常见的错误，如计数器变为负数、重复调用 Add() 或 Done() 等。通过深入理解 WaitGroup 的工作原理和使用场景，我们可以更加灵活地运用它来解决并发编程中的各种问题。