当前位置： 技术文章>> Go语言中的sync.WaitGroup如何使用？

文章标题：Go语言中的sync.WaitGroup如何使用？
  • 文章分类: 后端
  • 8757 阅读

在Go语言编程中，sync.WaitGroup 是一个非常重要的并发控制工具，它允许我们等待一组协程（goroutines）的完成。WaitGroup 通过计数器的方式工作，每当一个协程启动时，我们就增加计数器的值；当协程完成时，就减少计数器的值。当计数器的值变为零时，表示所有的协程都已经完成了它们的工作，此时主协程（或等待的协程）可以继续执行。这种机制非常适合于处理并行任务，确保所有子任务都完成后才进行下一步操作。

基本使用

要使用 sync.WaitGroup，首先需要引入 sync 包。WaitGroup 类型提供了三个方法：Add(delta int)、Done() 和 Wait()。

  • Add(delta int)：用于设置或增加等待协程的数量。如果 delta 为正数，则增加等待协程的数量；如果 delta 为负数，并且绝对值大于当前计数器的值，则会引发 panic。通常，在启动协程之前调用 Add(1) 来增加计数。
  • Done()：用于减少等待协程的数量，实际上是调用 Add(-1) 的简便方法。每个协程结束时应该调用 Done() 来表明自己已完成。
  • Wait()：阻塞调用它的协程，直到所有通过 Add 方法增加的协程都调用了 Done() 方法，即计数器的值变为零。

示例：使用 WaitGroup 等待多个协程完成

以下是一个使用 sync.WaitGroup 的基本示例，该示例中我们启动了多个协程来模拟一些耗时的操作，并使用 WaitGroup 等待它们全部完成。

package main

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

func worker(id int, wg *sync.WaitGroup) {
    defer wg.Done() // 确保协程结束时调用 Done 方法

    fmt.Printf("Worker %d starting\n", id)

    // 模拟耗时操作
    time.Sleep(time.Second)
    fmt.Printf("Worker %d done\n", id)
}

func main() {
    var wg sync.WaitGroup

    // 假设我们要启动 5 个协程
    for i := 1; i <= 5; i++ {
        wg.Add(1) // 为每个协程增加计数
        go worker(i, &wg) // 启动协程
    }

    wg.Wait() // 等待所有协程完成
    fmt.Println("All workers have finished their jobs")
}

在上面的示例中，main 函数首先创建了一个 sync.WaitGroup 实例 wg。然后，它通过一个循环启动了 5 个协程，每个协程都执行 worker 函数。在启动每个协程之前，wg.Add(1) 被调用以增加等待协程的计数。每个协程在其函数体的开始处通过 defer wg.Done() 语句确保在函数结束时（无论是正常返回还是由于 panic 退出）都会调用 Done() 方法来减少计数。最后，main 函数中的 wg.Wait() 调用会阻塞，直到所有协程都完成了它们的工作（即所有协程都调用了 Done()，导致计数器归零）。

进阶用法

嵌套 WaitGroup

有时，我们可能会遇到需要在协程中启动更多协程的情况，这时可以使用嵌套的 WaitGroup。

func nestedWorker(id int, wg, parentWg *sync.WaitGroup) {
    defer wg.Done() // 对于当前 WaitGroup 的 Done
    defer parentWg.Done() // 对于外层 WaitGroup 的 Done

    // 启动子协程
    wgChild := &sync.WaitGroup{}
    wgChild.Add(2) // 假设每个嵌套协程需要启动两个子协程
    for i := 0; i < 2; i++ {
        go func(i int) {
            defer wgChild.Done()
            fmt.Printf("Nested worker %d of %d\n", i, id)
            time.Sleep(time.Second)
        }(i)
    }

    wgChild.Wait() // 等待所有子协程完成
    fmt.Printf("Nested worker %d done\n", id)
}

// 在 main 中调用 nestedWorker 时，需要为 nestedWorker 和它的父协程都调用 Add

注意，在嵌套 WaitGroup 的情况下，确保每个层级的 WaitGroup 都正确管理其生命周期是非常重要的。

结合通道（Channel）使用

虽然 sync.WaitGroup 非常适合于等待一组协程的完成，但在某些情况下，我们可能还需要从协程中收集结果。这时，可以结合使用通道（Channel）和 WaitGroup。

func resultWorker(id int, wg *sync.WaitGroup, results chan<- int) {
    defer wg.Done()

    // 假设这是某种计算的结果
    result := id * 2
    results <- result // 将结果发送到通道
    fmt.Printf("Worker %d sent result %d\n", id, result)
}

func main() {
    var wg sync.WaitGroup
    results := make(chan int, 5) // 缓冲通道，大小为5

    for i := 1; i <= 5; i++ {
        wg.Add(1)
        go resultWorker(i, &wg, results)
    }

    go func() {
        wg.Wait()
        close(results) // 所有协程完成后关闭通道
    }()

    // 从通道中接收结果
    for result := range results {
        fmt.Printf("Received result: %d\n", result)
    }

    fmt.Println("All results have been processed")
}

在这个示例中，我们创建了一个缓冲通道 results 来接收从协程中发送的结果。每个协程计算一个结果并将其发送到通道中。我们还启动了一个额外的协程来等待所有工作协程完成（通过 wg.Wait()），并在所有工作协程完成后关闭通道。然后，我们使用 range 循环从通道中接收并处理所有结果。

总结

sync.WaitGroup 是 Go 语言中处理并发时非常重要的一个工具，它提供了一种简单而有效的方式来等待一组协程的完成。通过结合使用 Add、Done 和 Wait 方法，我们可以轻松地管理协程的生命周期，确保主协程（或任何等待的协程）在所有子协程都完成其任务之前不会继续执行。此外，WaitGroup 还可以与通道等其他并发原语结合使用，以构建更复杂、更灵活的并发模式。

希望这个详细的介绍和示例能够帮助你更好地理解和使用 sync.WaitGroup。如果你在探索 Go 语言的并发编程时遇到了任何问题，或者想要更深入地了解其他并发原语，不妨访问我的网站“码小课”，那里有更多关于 Go 语言及其生态的优质内容和教程等待着你。

推荐文章