上一篇 Go圣经-学习笔记之select多路复用

面试题：输入目录或者文件，求目录或者文件的磁盘占用大小，类似du命令？

编程实现

// 实现思路：
// 1. 获取当前所有目录和当前所有文件
// 2. 计算当前文件大小，并通过channel把数据回传给output
// 3. output读用户输入以及接收channel里的数据，并输出
func dirents(dir string) []os.FileInfo {
    files, err := ioutil.ReadDir(dir)
    if err != nil {
        log.Fatal(err.Error())
        return nil
    }
    return files
}

// 深度遍历
func walkDir(dir string, ch chan<- int64) {
    for _, file := range dirents(dir) {
        if file.IsDir() {
            subDir := filepath.Join(dir, file.Name())
            walkDir(subDir, ch)
        } else {
            ch <- file.Size()
        }
    }
}

func output(ch <-chan int64) {
    var (
        num      int
        fileSize int64
    )
    for size := range ch {
        num++
        fileSize += size
    }
    fmt.Printf("%d files. %.1f GB\n", num, float64(fileSize)/1e9)
    return
}

func main() {
    var (
        ch = make(chan int64)
        wg = new(sync.WaitGroup)
    )
    wg.Add(2) // two goroutines
    // 求所有目录文件的size总和
    go func() {
        for _, dir := range os.Args[1:] {
            walkDir(dir, ch)
        }
        close(ch)  // 发送端别忘记关闭channel了，否则接收端一直阻塞，导致wg.Wait一直等待,死锁。
        wg.Done()
    }()
    // 输出
    go func() {
        output(ch)
        wg.Done()
    }()
    wg.Wait()
    return
}

改进一 实时输出

上面这个DEMO，如果目录或者文件比较大的话，用户会感受到一直卡在那里，没有进度条。不知道当前程序运行是否正常。所以这个版本，运用select多路复用来实时输出数据。改进output方法

func output(ch <-chan int64) {
    var (
        num      int
        fileSize int64
        ticker   = time.NewTicker(1 * time.Second)
    )
Loop:
    for {
        select {
        case size, ok := <-ch:
            if !ok {
                break Loop
            }
            fileSize += size
            num++
        case <-ticker.C:
            fmt.Printf("%d files. %.1f GB\n", num, float64(fileSize)/1e9)
        }
    }
    fmt.Printf("%d files. %.1f GB\n", num, float64(fileSize)/1e9)
    ticker.Stop()
    return
}

说明：当channel发送端关闭后，接收端的goroutine接收完成后，也要退出。用到了goto语句，同时还要注意，防止ticker发送端的goroutine泄露，当output退出时，也要用ticker.Stop触发goroutine关闭。

改进二 并发计算

提高并发计算，每遇到一个目录，开启一个goroutine。也就是在每一个walkDir时，开启一个新的goroutine， 同时也要注意main goroutine需要等待所有的goroutine退出才能退出。需要改造walkDir和main

// 深度遍历
func walkDir(wg *sync.WaitGroup, dir string, ch chan<- int64) {
    for _, file := range dirents(dir) {
        if file.IsDir() {
            wg.Add(1)
            subDir := filepath.Join(dir, file.Name())
            go walkDir(wg, subDir, ch)
        } else {
            ch <- file.Size()
        }
    }
    wg.Done()
}

func main() {
    var (
        ch = make(chan int64)
        wg = new(sync.WaitGroup)
    )
    // 求所有目录文件的size总和
    for _, dir := range os.Args[1:] {
        wg.Add(1)
        go walkDir(wg, dir, ch)
    }
    // 输出
    go func() {
        wg.Wait()
        close(ch)
    }()
    output(ch)
    return
}

这个改进后的DEMO，已和原来的DEMO的设计大不一样，初学者要认真地理解下，主要有以下几点：

• 使用go关键字，只是创建一个goroutine，它不会被GoSched立即调度，所以在执行Go之前要写上wg.Add(1)， 否则，如果你写在walkDir函数的第一行，如下所示的代码：

// 深度遍历
func walkDir(wg *sync.WaitGroup, dir string, ch chan<- int64) {
    wg.Add(1)  // 新添加
    for _, file := range dirents(dir) {
        if file.IsDir() {
            subDir := filepath.Join(dir, file.Name())
            go walkDir(wg, subDir, ch) // 前面去掉了wg.Add(1)
        } else {
            ch <- file.Size()
        }
    }
    wg.Done()
}

func main() {
    var (
        ch = make(chan int64)
        wg = new(sync.WaitGroup)
    )
    // 求所有目录文件的size总和
    for _, dir := range os.Args[1:] {
        go walkDir(wg, dir, ch) // 前面去掉了wg.Add(1)
    }
    // 输出
    go func() {
        wg.Wait()
        close(ch)
    }()
    output(ch)
    return
}

这样程序可能是直接退出的，执行结果：0 files. 0.0 GB

• output放在了main goroutine中执行，主要是我们需要关闭发送端channel，不然接收端一直阻塞会导致死锁。
• wg.Wait和close(ch)从main goroutine中移除，放在了新的goroutine中，它主要是因为不能和output放在同一个goroutine中，否则，要么output只能输出最后一行，要么死锁。

希望初学者认真理解这个DEMO

改进三 限制goroutine数量

由于这个程序在高峰期会创建成百上千的goroutine，我们需要修改dirents函数，用计数信号量来阻止他同时打开太多的文件，就像我们在前面写的并发爬虫一样：

// sema is a counting semaphore for limiting concurrency in dirents.
var sema = make(chan struct{}, 20)

// dirents returns the entries of directory dir.
func dirents(dir string) []os.FileInfo {
    sema <- struct{}{}        // acquire token
    defer func() { <-sema }() // release token
    // ...

说明很重要的一点，控制goroutine的数量，通过全局共享的channel实现，注意点，必须是在goroutine的关键路径上做channel控制，其他还要考虑的一点是不要在递归函数内部做channel控制，很容易死锁，因为总是向channel发送数据，而channel接收端一直不工作，导致channel队列上的数据满，写一直阻塞，读又不工作。最后导致死锁deadlock

更通俗的说法，就是不要在生产goroutine的地方做channel控制，而是在goroutine执行体内的开头和结尾分别写读写channel，这样就goroutine数量就得到了全局的控制