golang 中的 sync.WaitGroup

sync.WaitGroup 是 Go 标准库中的一个同步原语，用于协调多个 goroutine 的执行，确保它们在主线程或其他 goroutine继续执行之前完成任务。

sync.WaitGroup 的用法

1. 创建 WaitGroup 实例

在开始 goroutine 执行之前，需要创建一个 WaitGroup 实例。这个实例将用于跟踪正在运行的 goroutine 的数量。

var wg sync.WaitGroup

2. 添加任务

在启动 goroutine 之前，通过调用 Add 方法设置需要等待的任务数量。例如，如果有两个 goroutine 要执行，就需要调用 wg.Add(2)。

wg.Add(1) // 启动一个新的 goroutine

3. 完成任务

每个 goroutine 在完成其任务时需要调用 Done 方法，以便减少计数器的数量。通常，这个调用在 goroutine 函数的最后。

go func() {
    defer wg.Done() // 在任务完成后调用 Done
    // 执行任务
}()

4. 等待任务完成

主线程或其他 goroutine 可以调用 Wait 方法，阻塞等待直到计数器变为零，即所有被跟踪的 goroutine 都完成了其执行。

wg.Wait() // 阻塞等待所有 goroutine 完成

示例代码

下面是一个示例，演示如何使用 sync.WaitGroup 来等待多个 goroutine 完成其执行：

package main

import (
    "fmt"
    "sync"
)

func main() {
    var wg sync.WaitGroup

    for i := 1; i <= 3; i++ {
        wg.Add(1) // 添加一个新的 goroutine 任务
        go func(i int) {
            defer wg.Done() // 完成任务时调用 Done
            fmt.Printf("Goroutine %d is running\n", i)
        }(i)
    }

    wg.Wait() // 等待所有 goroutine 完成
    fmt.Println("All goroutines completed")
}

展示如何使用 sync.WaitGroup 来处理多个并发的网络请求

package main

import (
    "fmt"
    "net/http"
    "sync"
    "io/ioutil"
)

// fetchData 从指定的 URL 获取数据
func fetchData(url string, wg *sync.WaitGroup, results chan<- string) {
    defer wg.Done() // 任务完成后调用 Done

    resp, err := http.Get(url)
    if err != nil {
        results <- fmt.Sprintf("Error fetching %s: %v", url, err)
        return
    }
    defer resp.Body.Close()

    body, err := ioutil.ReadAll(resp.Body)
    if err != nil {
        results <- fmt.Sprintf("Error reading response body from %s: %v", url, err)
        return
    }

    results <- fmt.Sprintf("Response from %s: %s", url, body)
}

func main() {
    urls := []string{
        "https://jsonplaceholder.typicode.com/posts/1",
        "https://jsonplaceholder.typicode.com/posts/2",
        "https://jsonplaceholder.typicode.com/posts/3",
    }

    var wg sync.WaitGroup
    results := make(chan string, len(urls)) // 缓冲通道，用于存储结果

    for _, url := range urls {
        wg.Add(1) // 添加一个新的 goroutine 任务
        go fetchData(url, &wg, results)
    }

    wg.Wait() // 等待所有 goroutine 完成
    close(results) // 关闭结果通道

    // 处理所有结果
    for result := range results {
        fmt.Println(result)
    }
}

代码解析

1. fetchData 函数:

   • 这是一个并发执行的函数，用于从指定的 URL 获取数据。
   • 使用 defer wg.Done() 来通知 WaitGroup 这个 goroutine 完成了它的任务。
   • 通过 results 通道将获取到的数据或错误信息传递到主 goroutine。

2. main 函数:

   • 定义了一组 URL 列表，程序将并发地从这些 URL 获取数据。
   • 使用 sync.WaitGroup 来确保所有的 fetchData goroutine 在继续之前完成。
   • 使用带缓冲的通道 results 来收集所有的响应结果。
   • 在所有请求完成后，关闭通道并处理所有结果。

这个例子中的 sync.WaitGroup 可以确保在所有的网络请求完成后才进行结果处理。在实际应用中，类似的模式可以用于多个并发的网络请求、并发计算、数据处理等任务，保证所有任务在继续执行后续步骤之前已经完成。

总结

sync.WaitGroup 通过计数器来协调多个 goroutine 的执行，它允许你在所有 goroutine 完成之前阻塞主线程或其他 goroutine。使用 WaitGroup 可以简化并发编程中的同步问题，确保所有并发操作完成后再继续执行后续代码。