Go语言sync.WaitGroup计数机制：如何避免Wait方法提前结束？

在Go语言中，sync.WaitGroup 用于等待一组 goroutine 完成。它通过一个计数器来跟踪 goroutine 的数量。Add 方法增加计数器，Done 方法减少计数器，而 Wait 方法会阻塞，直到计数器归零。

为了避免 Wait 方法提前结束，你需要确保以下几点：

1. 正确使用 Add 和 Done

• Add 方法：在启动 goroutine 之前调用 Add 方法，增加计数器的值。这个值应该与你要启动的 goroutine 数量一致。
• Done 方法：在每个 goroutine 的末尾调用 Done 方法，减少计数器的值。

2. 确保 Add 和 Done 的调用次数匹配

• 如果你调用了 Add(1)，那么必须确保在某个地方调用 Done()，否则 Wait 方法会一直阻塞。
• 如果你调用了 Add(n)，那么必须确保 Done() 被调用 n 次。

3. 避免在 goroutine 中调用 Add

• 不要在 goroutine 内部调用 Add，因为这可能导致 Wait 方法在 Add 之前被调用，从而导致 Wait 提前结束。

4. 使用 defer 调用 Done

• 在 goroutine 中使用 defer 调用 Done，以确保即使 goroutine 发生 panic，Done 也会被调用。

示例代码

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

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

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

解释

• 在 main 函数中，我们启动了 3 个 goroutine，每个 goroutine 都会调用 worker 函数。
• 在启动每个 goroutine 之前，我们调用 wg.Add(1) 来增加计数器。
• 在 worker 函数中，我们使用 defer wg.Done() 来确保在函数结束时减少计数器。
• 最后，wg.Wait() 会阻塞，直到所有 goroutine 都调用了 Done，即计数器归零。

常见错误

• 忘记调用 Add：如果忘记调用 Add，计数器不会增加，Wait 会立即返回。
• 忘记调用 Done：如果忘记调用 Done，计数器不会减少，Wait 会一直阻塞。
• 在 goroutine 中调用 Add：这可能导致 Wait 在 Add 之前被调用，从而导致 Wait 提前结束。

通过遵循上述规则，你可以确保 Wait 方法不会提前结束，并且能够正确等待所有 goroutine 完成。