Go语言sync.WaitGroup如何正确使用才能避免并发编程中的陷阱？

在Go语言中，sync.WaitGroup 是一种用于等待一组 goroutine 完成的同步机制。正确使用 sync.WaitGroup 可以避免并发编程中的一些常见陷阱。以下是一些关键点和最佳实践：

1. 初始化 WaitGroup

在使用 WaitGroup 之前，必须先初始化它。通常使用 var 声明一个 WaitGroup 变量，然后直接使用它。

var wg sync.WaitGroup

2. 使用 Add 方法

在启动每个 goroutine 之前，调用 wg.Add(1) 来增加 WaitGroup 的计数器。这个计数器表示有多少个 goroutine 需要等待。

wg.Add(1)
go func() {
    defer wg.Done()
    // 你的代码
}()

3. 使用 Done 方法

在每个 goroutine 完成时，调用 wg.Done() 来减少 WaitGroup 的计数器。通常使用 defer 来确保 Done 一定会被调用，即使在 goroutine 中发生了 panic。

go func() {
    defer wg.Done()
    // 你的代码
}()

4. 使用 Wait 方法

在主 goroutine 中调用 wg.Wait() 来阻塞，直到所有 goroutine 都完成（即 WaitGroup 的计数器变为 0）。

wg.Wait()

5. 避免计数器为负

确保 Add 和 Done 的调用次数匹配，否则会导致 WaitGroup 的计数器变为负数，从而引发 panic。

6. 避免在 goroutine 中调用 Add

不要在 goroutine 内部调用 wg.Add(1)，因为这样可能会导致主 goroutine 在调用 Wait 时，WaitGroup 的计数器还没有被正确增加。

7. 使用 WaitGroup 的典型模式

以下是一个典型的使用 WaitGroup 的模式：

package main

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

func worker(id int, wg *sync.WaitGroup) {
    defer wg.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

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

    wg.Wait()
    fmt.Println("All workers done")
}

8. 避免在 Wait 之后修改 WaitGroup

一旦调用了 Wait，WaitGroup 的计数器应该已经变为 0，此时不应该再对 WaitGroup 进行任何操作。

9. 使用 WaitGroup 与 context 结合

在某些情况下，你可能需要结合 context 来管理 goroutine 的生命周期，以便在超时或取消时优雅地停止 goroutine。

package main

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

func worker(ctx context.Context, id int, wg *sync.WaitGroup) {
    defer wg.Done()
    select {
    case <-time.After(time.Second):
        fmt.Printf("Worker %d done\n", id)
    case <-ctx.Done():
        fmt.Printf("Worker %d canceled\n", id)
    }
}

func main() {
    var wg sync.WaitGroup
    ctx, cancel := context.WithTimeout(context.Background(), 500*time.Millisecond)
    defer cancel()

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

    wg.Wait()
    fmt.Println("All workers done")
}

总结

正确使用 sync.WaitGroup 的关键在于确保 Add 和 Done 的调用次数匹配，并且在适当的地方调用 Wait。遵循这些最佳实践可以避免并发编程中的常见陷阱，确保程序的正确性和稳定性。