Go并发编程：使用sync.WaitGroup同步Goroutine与通道操作

package main

import (
    "fmt"
    "os"
)

type uniprot struct {
    namesInDir chan string
}

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

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

func errorCheck(err error) {
    if err != nil {
        fmt.Fprintf(os.Stderr, "Error: %v\n", err)
        os.Exit(1)
    }
}

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

    for _, n := range names {
        go u.produce(n) // 启动生产者Goroutine
        go u.consume()  // 启动消费者Goroutine
    }
    // 问题：这里缺少等待机制
}

func main() {
    // 假设存在一个名为 "test_dir" 的目录，其中包含一些文件
    // 为了演示，这里不创建实际目录和文件，但实际应用中需要
    // u := &uniprot{}
    // u.readFilenames("test_dir")
}

package main

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

type uniprot struct {
    namesInDir chan string
}

// produce函数现在接受一个WaitGroup指针
func (u *uniprot) produce(n string, wg *sync.WaitGroup) {
    defer wg.Done() // 确保在函数退出时调用wg.Done()
    u.namesInDir <- n
}

// consume函数现在接受一个WaitGroup指针
func (u *uniprot) consume(wg *sync.WaitGroup) {
    defer wg.Done() // 确保在函数退出时调用wg.Done()
    fmt.Println(<-u.namesInDir)
}

func errorCheck(err error) {
    if err != nil {
        fmt.Fprintf(os.Stderr, "Error: %v\n", err)
        os.Exit(1)
    }
}

func (u *uniprot) readFilenames(dirname string) {
    u.namesInDir = make(chan string, 15)
    dir, err := os.Open(dirname)
    errorCheck(err)
    names, err := dir.Readdirnames(0)
    errorCheck(err)

    var wg sync.WaitGroup // 声明一个WaitGroup变量

    for _, n := range names {
        wg.Add(2) // 每次循环增加计数器2，因为启动了两个Goroutine (produce 和 consume)
        go u.produce(n, &wg) // 将WaitGroup的地址传递给Goroutine
        go u.consume(&wg)   // 将WaitGroup的地址传递给Goroutine
    }

    wg.Wait() // 阻塞直到所有Goroutine都调用了Done()，即计数器归零
    close(u.namesInDir) // 所有生产和消费完成后，关闭通道
}

func main() {
    // 为了运行示例，我们需要一个测试目录和文件
    // 假设我们在当前目录下创建一个名为 "test_dir" 的目录，并在其中创建两个文件
    // 例如：
    // mkdir test_dir
    // touch test_dir/file1.txt
    // touch test_dir/file2.txt

    // 确保目录存在
    testDir := "test_dir"
    err := os.MkdirAll(testDir, 0755)
    errorCheck(err)

    // 创建一些测试文件
    for i := 1; i <= 5; i++ {
        filePath := fmt.Sprintf("%s/file%d.txt", testDir, i)
        f, err := os.Create(filePath)
        errorCheck(err)
        f.Close()
    }

    u := &uniprot{}
    u.readFilenames(testDir)

    fmt.Println("所有文件处理完毕。")

    // 清理测试目录 (可选)
    os.RemoveAll(testDir)
}