Go并发编程：使用WaitGroup实现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完成
}

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.")
}