深入理解 Go 并发操作：通过示例讲解 sync.WaitGroup 和 Context 的使用

在 Go 语言中，并发编程是一个重要的特性。当一个任务可以被拆分为多个相互独立的子任务，并且这些子任务没有先后执行顺序的限制，需要等到所有子任务执行完毕后再进行下一步处理时，适合使用 sync.WaitGroup。

然而，虽然 sync.WaitGroup 使用起来相对简单，但如果不小心，可能会踩到一些坑。本文将通过示例深入讲解如何正确使用 sync.WaitGroup，并介绍在更复杂的并发场景下如何结合 context 来处理任务取消。

正确使用 sync.WaitGroup

假设有一个任务需要执行 3 个子任务，可以使用 sync.WaitGroup 来等待所有子任务完成：

package main

import (
    "fmt"
    "sync"
)

func main() {
    var wg sync.WaitGroup

    wg.Add(3)

    go handleTask1(&wg)
    go handleTask2(&wg)
    go handleTask3(&wg)

    wg.Wait()

    fmt.Println("所有任务执行完毕.")
}

func handleTask1(wg *sync.WaitGroup) {
    defer wg.Done()
    fmt.Println("执行任务 1")
}

func handleTask2(wg *sync.WaitGroup) {
    defer wg.Done()
    fmt.Println("执行任务 2")
}

func handleTask3(wg *sync.WaitGroup) {
    defer wg.Done()
    fmt.Println("执行任务 3")
}

执行输出：

执行任务 1
执行任务 3
执行任务 2
所有任务执行完毕.

注意，任务的执行顺序可能会有所不同，因为 Goroutine 是并发运行的。

sync.WaitGroup 使用注意事项

1. 传递指针

正确做法：

go handleTask1(&wg)

错误做法：

go handleTask1(wg)

在启动子任务时，必须传递 *sync.WaitGroup 的指针类型。如果传值，会导致每个子任务都有自己的副本，无法正确同步。

2. 确保 `wg.Add()` 在启动 Goroutine 之前执行

错误示例：

var wg sync.WaitGroup

go handleTask1(&wg)

wg.Wait()

// ...

func handleTask1(wg *sync.WaitGroup) {
    wg.Add(1)
    defer wg.Done()
    fmt.Println("执行任务 1")
}

在上述代码中，wg.Add(1) 在子 Goroutine 内部执行，但主 Goroutine 已经调用了 wg.Wait()，这可能导致主 Goroutine 提前结束。应当确保 wg.Add() 在启动 Goroutine 之前执行。

3. 保持计数器一致

注意 wg.Add() 和 wg.Done() 的调用次数应当一致。wg.Done() 实际上等价于 wg.Add(-1)。

更复杂的场景：任务失败时取消其他任务

sync.WaitGroup 适用于等待所有子任务完成的场景，但如果需要在某个子任务失败时，立即取消其他正在执行的子任务，就需要结合 context 来实现。

使用 sync.WaitGroup 和 Context

以下示例展示了如何在一个子任务失败时，取消其他所有子任务：

package main

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

func handleTask(ctx context.Context, taskID int, wg *sync.WaitGroup, cancel context.CancelFunc) {
    defer wg.Done()

    fmt.Printf("任务 %d 正在处理...\n", taskID)

    // 模拟任务处理，任务 1 将模拟失败
    if taskID == 1 {
        fmt.Printf("任务 %d 失败\n", taskID)
        cancel() // 取消上下文，通知其他任务停止
        return
    }

    // 检查上下文是否已取消
    select {
    case <-ctx.Done():
        fmt.Printf("任务 %d 已被取消\n", taskID)
        return
    default:
        // 模拟耗时操作
        time.Sleep(3 * time.Second)
        fmt.Printf("任务 %d 成功完成\n", taskID)
    }
}

func main() {
    ctx, cancel := context.WithCancel(context.Background())
    var wg sync.WaitGroup

    wg.Add(3)
    go handleTask(ctx, 1, &wg, cancel)
    go handleTask(ctx, 2, &wg, cancel)
    go handleTask(ctx, 3, &wg, cancel)

    wg.Wait()
    fmt.Println("所有任务已结束")
}

输出示例：

任务 1 正在处理...
任务 2 正在处理...
任务 3 正在处理...
任务 1 失败
任务 2 已被取消
任务 3 已被取消
所有任务已结束

解释

任务 1 开始处理：
- 输出 任务 1 正在处理...。
- 由于 taskID == 1，模拟任务失败，输出 任务 1 失败。
- 调用 cancel()，取消上下文 ctx，通知其他任务停止。
任务 2 和任务 3 开始处理：
- 输出 任务 2 正在处理... 和 任务 3 正在处理...。
- 检查上下文 ctx，发现已被取消。
- 输出 任务 2 已被取消 和 任务 3 已被取消，然后退出。
等待所有任务结束：
- 主 Goroutine 调用 wg.Wait()，等待所有子任务完成。
- 所有子任务结束后，输出 所有任务已结束。

通过这种方式，可以确保在任一子任务失败时，其他子任务能够及时检测到取消信号并停止执行，避免浪费资源。

小结

使用 sync.WaitGroup 时，确保传递的是指针，且在启动 Goroutine 之前调用 wg.Add()。
在需要控制子任务的执行，如在某个任务失败时取消其他任务的场景下，结合 context 使用可以更好地管理并发任务。
sync.WaitGroup 适用于等待所有子任务完成的简单场景；对于更复杂的并发控制，context 提供了强大的功能。

通过对 sync.WaitGroup 和 context 的灵活运用，可以更有效地管理 Go 程序中的并发操作，提高程序的健壮性和效率。

复制全文生成海报 Go语言并发编程编程技巧

编程 深入理解 Go 并发操作：通过示例讲解 sync.WaitGroup 和 Context 的使用