golang-cookbook

Appearance

WaitGroup ​

Go sync.WaitGroup 详解 ​

sync.WaitGroup 是 Go 并发编程中用于等待一组 Goroutine 完成的同步原语。它主要用于协调多个 Goroutine 在执行结束后继续执行。通常与多个并发的 Goroutine 一起使用，确保主程序等待所有 Goroutine 完成任务后再继续执行。

1. WaitGroup 的基本概念 ​

WaitGroup 是一个计数器，能够跟踪 Goroutine 的数量。它提供了三种主要方法：

  • Add(delta int)：增加或减少计数器的值。
  • Done()：减少计数器的值，通常由每个 Goroutine 调用，表示它已完成工作。
  • Wait()：阻塞直到计数器的值为 0，通常由主 Goroutine 调用，表示等待所有 Goroutine 完成。

2. 常见的使用模式 ​

2.1 启动多个 Goroutine，等待其完成 ​

一个常见的使用场景是启动多个 Goroutine，然后在主程序中使用 WaitGroup 来等待所有 Goroutine 完成。通过 Add() 方法来设置等待的数量，然后每个 Goroutine 在完成时调用 Done() 方法，最后使用 Wait() 方法在主程序中等待。

go
package main

import (
    "fmt"
    "sync"
)

func printMessage(i int, wg *sync.WaitGroup) {
    defer wg.Done()  // 完成后调用 Done，通知 WaitGroup 当前 Goroutine 完成
    fmt.Println("Message", i)
}

func main() {
    var wg sync.WaitGroup

    // 启动多个 Goroutine
    for i := 0; i < 5; i++ {
        wg.Add(1)  // 增加计数器，表示需要等待一个 Goroutine
        go printMessage(i, &wg)
    }

    // 等待所有 Goroutine 完成
    wg.Wait()
    fmt.Println("All Goroutines are done.")
}

输出：

Message 0
Message 1
Message 2
Message 3
Message 4
All Goroutines are done.
2.2 使用 Add() 动态调整计数器 ​

你可以在运行时动态地调整计数器的数量。通常，在启动 Goroutine 之前调用 Add() 来增加计数器。

go
package main

import (
    "fmt"
    "sync"
)

func main() {
    var wg sync.WaitGroup

    // 增加计数器，表示需要等待的 Goroutine 数量
    wg.Add(2)

    go func() {
        defer wg.Done()
        fmt.Println("Goroutine 1 is done")
    }()

    go func() {
        defer wg.Done()
        fmt.Println("Goroutine 2 is done")
    }()

    // 等待所有 Goroutine 完成
    wg.Wait()
    fmt.Println("All Goroutines are done.")
}

输出：

Goroutine 1 is done
Goroutine 2 is done
All Goroutines are done.

3. 工作原理 ​

  • wg.Add(n)：增加计数器的值。n 表示当前要等待的 Goroutine 数量。通常在启动 Goroutine 之前调用 Add()。
  • wg.Done()：每个 Goroutine 在执行完毕后调用 Done()，它会减少 WaitGroup 的计数器。如果计数器为 0，Wait() 会停止阻塞，继续执行。
  • wg.Wait()：主 Goroutine 使用 Wait() 来阻塞，直到所有 Goroutine 完成工作并调用了 Done()。

4. 注意事项 ​

  • 确保调用 Done()：每个 Add() 的计数都必须对应一个 Done() 调用，否则 Wait() 会一直阻塞，导致死锁。通常建议使用 defer 来确保 Done() 被调用。
  • 避免多次调用 Add()：如果在多个 Goroutine 中并发调用 Add()，可能会导致 WaitGroup 计数不一致，建议将 Add() 调用集中到主 Goroutine。
  • 线程安全：sync.WaitGroup 是线程安全的，可以在多个 Goroutine 中安全地使用。

5. 错误示例 ​

以下代码示例中没有正确调用 Done()，会导致死锁，因为 Wait() 会一直阻塞：

go
package main

import (
    "fmt"
    "sync"
)

func printMessage(wg *sync.WaitGroup) {
    // 没有调用 Done()
    fmt.Println("Message")
}

func main() {
    var wg sync.WaitGroup
    wg.Add(1)
    go printMessage(&wg)

    // 死锁，程序永远不会结束
    wg.Wait()
}

6. WaitGroup 在实际应用中的其他场景 ​

  • 协同工作：WaitGroup 可以用于多个 Goroutine 同时执行一个任务，并在所有任务完成后通知主程序。
  • 错误处理：WaitGroup 可以与错误处理结合使用，确保在等待多个任务完成时，可以处理每个任务的错误。
  • 并发计算：在并发计算中，使用 WaitGroup 等待所有计算完成后再进行结果汇总。

7. 与 Channel 一起使用 ​

WaitGroup 和 Channel 可以结合使用，以便在多个 Goroutine 完成后处理结果。

go
package main

import (
    "fmt"
    "sync"
)

func calculate(n int, ch chan int, wg *sync.WaitGroup) {
    defer wg.Done()  // 完成后调用 Done
    ch <- n * 2
}

func main() {
    var wg sync.WaitGroup
    ch := make(chan int, 5)

    for i := 1; i <= 5; i++ {
        wg.Add(1)  // 增加计数器
        go calculate(i, ch, &wg)
    }

    wg.Wait()  // 等待所有 Goroutine 完成
    close(ch)  // 关闭 Channel

    // 处理结果
    for result := range ch {
        fmt.Println("Result:", result)
    }
}

输出：

Result: 2
Result: 4
Result: 6
Result: 8
Result: 10

8. 总结 ​

  • sync.WaitGroup 是 Go 并发编程中一个非常实用的工具，用于等待多个 Goroutine 执行完成。
  • 通过 Add()、Done() 和 Wait() 方法，可以高效地协调多个 Goroutine 的执行。
  • 需要确保在每个 Goroutine 中调用 Done()，以避免死锁。
  • WaitGroup 是线程安全的，可以安全地在多个 Goroutine 中使用。