上一篇 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函数，用计数信号量来阻止他同时打开太多的文件，就像我们在8.7节中的并发爬虫一样：

// 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     // ...