15分钟学 Go 第 25 天：使用WaitGroup

第25天：使用WaitGroup

目标

理解Go语言中的WaitGroup的用法，掌握其在并发编程中的重要性以及使用指南。

什么是WaitGroup？

WaitGroup是Go语言中的一个同步原语，用于等待一组Goroutine完成。它能够控制多个Goroutine的执行，确保在所有工作完成之前主程序不会退出。

WaitGroup的基本原理

WaitGroup包含了三个主要的方法：

• Add(int)：增加等待的Goroutine数量
• Done()：减少等待的Goroutine数量
• Wait()：阻塞直到等待的Goroutine完成

WaitGroup的使用流程

在使用WaitGroup时，一般遵循以下流程：

1. 创建一个WaitGroup实例
2. 在启动每个Goroutine之前调用Add()方法增加计数
3. 在Goroutine完成时调用Done()方法减少计数
4. 在主线程中调用Wait()方法，阻塞直到所有Goroutine完成

WaitGroup的基本示例

以下是一个简单的示例，通过WaitGroup来等待多个Goroutine完成任务。

package main

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

func worker(id int, wg *sync.WaitGroup) {
    defer wg.Done() // 确保Goroutine完成时调用Done
    fmt.Printf("Worker %d is starting\n", id)
    time.Sleep(time.Second) // 模拟工作
    fmt.Printf("Worker %d is done\n", id)
}

func main() {
    var wg sync.WaitGroup

    // 启动多个Goroutine
    for i := 1; i <= 3; i++ {
        wg.Add(1) // 增加等待的Goroutine数量
        go worker(i, &wg) // 启动Goroutine
    }

    wg.Wait() // 阻塞直到所有Goroutine完成
    fmt.Println("All workers are done")
}

代码说明

1. worker函数：每个Goroutine的工作函数，接受Goroutine的ID和WaitGroup指针。在完成工作后调用wg.Done()以表明任务完成。
2. main函数：
   • 创建WaitGroup实例wg。
   • 使用循环启动多个Goroutine，每次调用wg.Add(1)来增加计数。
   • 最后调用wg.Wait()，等待所有Goroutine完成。

代码运行流程图

下面是代码执行的流程图，帮助你理解WaitGroup的工作原理：

+--------------------------+
|        Main Goroutine    |
+--------------------------+
            |
            v
+--------------------------+
|        wg.Add(1)        |
+--------------------------+
            |
            v
+--------------------------+
|    Go Create Goroutine   |——> w := worker(i, &wg)
+--------------------------+
            |                      |
            v                      |
+--------------------------+       |
|       worker函数        |       |
+--------------------------+       |
|         wg.Done()       |———> *wg计数减少
|  (任务完成, 计数-1)     |       |
+--------------------------+       |
            |                      |
            v                      |
+--------------------------+       |
|         wg.Wait()       |<———–+
|  (阻塞等所有Goroutine)   |
+--------------------------+
            |
            v
+--------------------------+
|  All workers are done    |
+--------------------------+

WaitGroup的注意事项

在使用WaitGroup时，以下几点是需要特别注意的：

1. 避免并发调用：在多个Goroutine之间共享同一个WaitGroup时，务必确保对Add()、Done()和Wait()的调用是安全的。一般情况下，每个Goroutine都只调用自己的Done()，而主Goroutine则在所有其他Goroutine完成后调用Wait()。

2. 不应在Wait()调用之前调用Done()：如果在Wait()之前调用Done()，可能会导致程序异常结束。确保每次调用Add()时，都会有对应的Done()来平衡。

如果Goroutine中出现错误

如果某个Goroutine中的任务发生错误，你可以使用defer和recover对错误进行处理，确保Done()被调用。在下面的示例中，我们将模拟错误处理。

package main

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

func workerWithError(id int, wg *sync.WaitGroup) {
    defer wg.Done()

    if id == 2 { // 在id为2的工作中模拟错误
        fmt.Printf("Worker %d encountered an error\n", id)
        return
    }

    fmt.Printf("Worker %d is starting\n", id)
    time.Sleep(time.Second)
    fmt.Printf("Worker %d is done\n", id)
}

func main() {
    var wg sync.WaitGroup

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

    wg.Wait()
    fmt.Println("All workers are done with error handling")
}

代码说明

在这个示例中，当Goroutine的ID为2时，它将不会正常完成任务，而是输出错误信息。尽管如此，Done()仍将被调用，确保其他Goroutine可以继续执行。

使用WaitGroup处理不同的任务

有时，我们可能希望在同一个程序中使用WaitGroup来处理不同类型的任务。例如，发送请求和处理数据库记录的异步操作。这是一个更复杂的例子，包含多种任务。

package main

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

func taskA(wg *sync.WaitGroup) {
    defer wg.Done()
    fmt.Println("Task A is running...")
    time.Sleep(2 * time.Second)
    fmt.Println("Task A is completed")
}

func taskB(wg *sync.WaitGroup) {
    defer wg.Done()
    fmt.Println("Task B is running...")
    time.Sleep(1 * time.Second)
    fmt.Println("Task B is completed")
}

func main() {
    var wg sync.WaitGroup

    wg.Add(2) // 两个任务
    go taskA(&wg)
    go taskB(&wg)

    wg.Wait() // 等待所有任务完成
    fmt.Println("All tasks are done")
}

WaitGroup的总结

WaitGroup是Go语言中处理并发任务的重要工具，它能够有效地组织和协调多个Goroutine，确保它们在程序退出前完成。理解WaitGroup的使用能帮助开发者更好地管理并发操作，提高程序的效率与可维护性。

通过前面的例子和流程图，我们已经了解了WaitGroup的基本用法和一些注意事项。在实际项目中，合理利用WaitGroup能够为代码结构带来极大的便利和清晰度。

表格总结

那么，现在你掌握了WaitGroup的基本用法，接下来可以尝试用WaitGroup来完成一些更复杂的并发任务练习。你可以尝试组合多个Goroutine任务，或者在Goroutine中引入更多的错误处理逻辑，以提高你的编程能力。

怎么样今天的内容还满意吗？再次感谢观众老爷的观看，关注GZH：凡人的AI工具箱，回复666，送您价值199的AI大礼包。最后，祝您早日实现财务自由，还请给个赞，谢谢！