Go并发：WaitGroup同步与等待技巧

本文探讨了Go语言并发编程中一个常见问题：主goroutine过早退出导致子goroutine无法完成任务。通过分析原始代码的潜在问题，我们引入了`sync.WaitGroup`这一强大的同步原语，详细阐述了其工作原理及在生产者-消费者模型中的应用，旨在帮助开发者正确地等待所有并发任务完成，确保程序的健壮性与预期行为。

在Go语言中，利用goroutine实现并发操作是其核心优势之一。然而，初学者常会遇到一个陷阱：主程序（或调用goroutine的函数）在子goroutine完成其工作之前就退出了，导致子goroutine的输出或副作用没有被观察到。这通常发生在没有适当同步机制的情况下。

理解并发执行中的常见问题

考虑一个典型的生产者-消费者模型，其中一个goroutine负责生产数据并发送到通道（channel），另一个goroutine负责从通道接收并处理数据。以下是一个简化但存在问题的示例代码：

package main

import (
    "fmt"
    "os"
    "time" // 仅用于演示问题，非解决方案
)

type uniprot struct {
    namesInDir chan string
}

func errorCheck(err error) {
    if err != nil {
        fmt.Printf("Error: %v\n", err)
        // 实际应用中应有更完善的错误处理
    }
}

func (u *uniprot) produce(n string) {
    u.namesInDir <- n
}

func (u *uniprot) consume() {
    fmt.Println(<-u.namesInDir)
}

func (u *uniprot) readFilenames(dirname string) {
    u.namesInDir = make(chan string, 15) // 创建一个带缓冲的通道
    dir, err := os.Open(dirname)
    errorCheck(err)
    defer dir.Close() // 确保文件句柄关闭

    names, err := dir.Readdirnames(0) // 读取所有文件名
    errorCheck(err)

    for _, n := range names {
        go u.produce(n) // 启动生产者goroutine
        go u.consume()  // 启动消费者goroutine
    }
    // 问题所在：readFilenames函数可能在此处立即返回
    // 主goroutine可能在produce和consume goroutine完成之前退出
    // time.Sleep(1 * time.Second) // 临时演示问题，不推荐作为解决方案
}

func main() {
    u := &uniprot{}
    // 假设当前目录下有一个test_dir，包含一些文件
    // 例如：mkdir test_dir && touch test_dir/file1.txt test_dir/file2.txt
    u.readFilenames("./test_dir")
    // 如果没有同步机制，main函数可能会在readFilenames返回后立即退出
    // 导致produce和consume goroutine没有足够时间执行
    time.Sleep(100 * time.Millisecond) // 即使这样也可能无法保证所有goroutine完成
}

在上述代码中，readFilenames函数通过循环为每个文件名启动了一个生产者goroutine和一个消费者goroutine。然而，由于Go运行时调度goroutine的特性，readFilenames函数本身可能会在这些新启动的goroutine有机会执行之前就完成并返回。如果main函数随后也立即退出，那么这些生产者和消费者goroutine可能永远不会打印任何内容，因为它们的宿主进程已经终止。

为了验证这个问题，有时会临时在readFilenames函数的末尾添加一个time.Sleep来强制主goroutine等待一段时间。虽然这能“暂时”解决问题，但它是一种不确定且不可靠的解决方案，因为你无法预知所有goroutine完成所需的确切时间。

解决方案：使用 sync.WaitGroup 进行同步

Go语言标准库提供了sync.WaitGroup类型，它是解决这类问题的理想工具。WaitGroup用于等待一组goroutine完成。它有三个主要方法：

• Add(delta int)：将内部计数器增加delta。通常在启动goroutine之前调用。
• Done()：将内部计数器减1。通常在goroutine完成其工作时调用，通常通过defer语句确保执行。
• Wait()：阻塞直到内部计数器归零。

下面是使用sync.WaitGroup改进后的代码示例：

package main

import (
    "fmt"
    "os"
    "sync" // 导入sync包
)

type uniprot struct {
    namesInDir chan string
}

func errorCheck(err error) {
    if err != nil {
        fmt.Printf("Error: %v\n", err)
        // 实际应用中应有更完善的错误处理，例如panic或返回错误
    }
}

// produce函数现在接收一个WaitGroup指针
func (u *uniprot) produce(n string, wg *sync.WaitGroup) {
    defer wg.Done() // goroutine完成时，通知WaitGroup
    u.namesInDir <- n
    fmt.Printf("Produced: %s\n", n) // 添加日志以便观察
}

// consume函数现在接收一个WaitGroup指针
func (u *uniprot) consume(wg *sync.WaitGroup) {
    defer wg.Done() // goroutine完成时，通知WaitGroup
    data := <-u.namesInDir
    fmt.Printf("Consumed: %s\n", data)
}

func (u *uniprot) readFilenames(dirname string) {
    u.namesInDir = make(chan string, 15) // 创建一个带缓冲的通道
    dir, err := os.Open(dirname)
    errorCheck(err)
    defer dir.Close()

    names, err := dir.Readdirnames(0)
    errorCheck(err)

    var wg sync.WaitGroup // 创建一个WaitGroup实例

    for _, n := range names {
        wg.Add(2) // 每次启动一对生产者和消费者goroutine，计数器加2
        go u.produce(n, &wg)
        go u.consume(&wg)
    }

    wg.Wait() // 阻塞直到所有Add的goroutine都调用了Done()
    close(u.namesInDir) // 所有生产者完成后关闭通道，通知消费者没有更多数据
}

func main() {
    u := &uniprot{}
    // 确保当前目录下存在一个名为"test_dir"的文件夹，并包含一些文件
    // 例如：
    // mkdir -p test_dir
    // touch test_dir/file1.txt test_dir/file2.txt test_dir/file3.txt
    u.readFilenames("./test_dir")
    fmt.Println("All goroutines finished and main function exiting.")
}

代码详解与注意事项

1. var wg sync.WaitGroup: 在readFilenames函数中创建了一个WaitGroup实例。
2. wg.Add(2): 在每次循环迭代中，我们启动了一个生产者和一个消费者goroutine，所以需要将WaitGroup的计数器增加2。这发生在启动goroutine之前，以确保在Wait()被调用时，计数器已经被正确设置。
3. defer wg.Done(): 在produce和consume这两个goroutine的内部，我们使用defer wg.Done()。defer确保了无论goroutine是正常完成还是因为panic而退出，Done()方法都会被调用，从而将WaitGroup的计数器减1。这是确保所有goroutine都被正确计数的关键。
4. wg.Wait(): 在readFilenames函数的末尾，wg.Wait()方法会被调用。这将阻塞readFilenames函数，直到WaitGroup的内部计数器变为零（即所有通过Add(2)增加的goroutine都通过Done()完成了任务）。一旦计数器归零，readFilenames函数才能继续执行并最终返回。
5. 通道的关闭: 在wg.Wait()之后，我们添加了close(u.namesInDir)。这在生产者-消费者模型中是一个最佳实践。当所有生产者都完成并确认不会再向通道发送数据时，关闭通道可以向消费者发出信号，表明不会再有新的数据。如果消费者使用for range循环从通道读取数据，通道关闭后循环会自动结束。

总结

sync.WaitGroup是Go语言中一个非常重要的并发原语，它提供了一种简单而有效的方式来等待一组goroutine完成。正确使用WaitGroup可以避免主goroutine过早退出，确保所有并发任务都能按预期执行，从而使并发程序更加健壮和可预测。在设计Go并发程序时，特别是涉及到多个goroutine协作完成一项任务时，务必考虑使用WaitGroup或其他同步机制来协调它们的生命周期。