一个小例子，给你讲透典型的 Go 并发操作

在 Go 中，如果任务可以分解为多个子任务并发执行，等到所有子任务执行完毕后再进行下一步处理，使用 sync.WaitGroup 是非常合适的选择。然而，尽管 sync.WaitGroup 用法简单，仍然存在一些容易踩到的坑。本文将通过一个小例子详细介绍 sync.WaitGroup 的正确使用及常见的陷阱。

sync.WaitGroup 的正确使用

我们有一个任务需要执行 3 个子任务，可以这样使用 sync.WaitGroup：

package main

import (
    "fmt"
    "sync"
)

func main() {
    var wg sync.WaitGroup

    wg.Add(3)

    go handlerTask1(&wg)
    go handlerTask2(&wg)
    go handlerTask3(&wg)

    wg.Wait()  // 等待所有子任务执行完毕

    fmt.Println("全部任务执行完毕.")
}

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

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

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

输出结果：

执行任务 3
执行任务 1
执行任务 2
全部任务执行完毕.

sync.WaitGroup 闭坑指南

1. 传递指针

正确做法：传递 sync.WaitGroup 的指针。

go handlerTask1(&wg)  // 正确

错误做法：传递 sync.WaitGroup 的值。

go handlerTask1(wg)   // 错误

2. wg.Add() 位置

错误示例：将 wg.Add() 放在 go handlerTask() 中。

var wg sync.WaitGroup
go handlerTask1(&wg)
wg.Wait()

func handlerTask1(wg *sync.WaitGroup) {
    wg.Add(1)   // 错误位置
    defer wg.Done()
    fmt.Println("执行任务 1")
}

在 wg.Wait() 之前，必须先调用 wg.Add()。因此，wg.Add() 应在启动 goroutine 之前调用。

sync.WaitGroup + Context

当需求是在某个子任务失败时取消其他任务的执行，sync.WaitGroup 本身无法实现此功能，此时可以结合 context 来处理。

代码示例：

package main

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

// handlerTask 处理单个任务
func handlerTask(ctx context.Context, taskId int, wg *sync.WaitGroup, cancel context.CancelFunc) {
    defer wg.Done()

    fmt.Printf("Request %d is processing...\n", taskId)

    // 模拟请求处理，如果 taskId 为1，则模拟失败
    if taskId == 1 {
        fmt.Printf("Request %d failed\n", taskId)
        cancel()  // 取消 context，通知其他请求停止
        return
    }

    // 监听 context.Done() 通道
    select {
    case <-ctx.Done():
        fmt.Printf("Request %d is already cancelled\n", taskId)
        return
    default:
        time.Sleep(3 * time.Second)  // 模拟耗时操作
        fmt.Printf("Request %d succeeded\n", taskId)
    }
}

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

    wg.Add(3)
    go handlerTask(ctx, 1, &wg, cancel)  // 模拟失败任务
    go handlerTask(ctx, 2, &wg, cancel)  // 模拟成功任务
    go handlerTask(ctx, 3, &wg, cancel)  // 模拟成功任务

    // 等待所有任务完成或被取消
    go func() {
        wg.Wait()
        fmt.Println("All requests are finished or cancelled")
    }()

    time.Sleep(5 * time.Second)  // 给一些时间处理，防止主 goroutine 过早退出
}

执行结果：

Request 1 is processing...
Request 1 failed
Request 2 is processing...
Request 3 is processing...
Request 2 is already cancelled
Request 3 is already cancelled
All requests are finished or cancelled

解释：

1. 任务 1 失败：

   • 任务 1 处理开始，发现任务 ID 是 1，输出 "Request 1 failed" 并调用 cancel()，取消 context，通知其他任务停止执行。

2. 任务 2 和任务 3 被取消：

   • 任务 2 和任务 3 监听 ctx.Done() 通道，发现通道已关闭，输出 "Request 2 is already cancelled" 和 "Request 3 is already cancelled"。

3. 等待所有任务结束：

   • 使用 wg.Wait() 等待所有任务完成或取消，输出 "All requests are finished or cancelled"。

总结

通过这两个例子，你可以看到：

• sync.WaitGroup 适用于等待所有子任务执行完毕后再进行下一步处理，使用时注意传递指针并保证 wg.Add() 和 wg.Done() 的调用次数一致。
• sync.WaitGroup + context 组合适用于当某个任务失败时，及时取消其他任务的执行，确保资源的有效利用。

通过合理使用 sync.WaitGroup 和 context，可以简化并发编程中的许多场景，让代码更加健壮。