20230409@实现Unix中du命令统计文件

在本节中，我们将构建一个程序，根据命令行指定的输入，报告一个或多个目录的磁盘使用情况，类似于 UNIX 的du命令。该程序大多数工作是由下面的 walkDir 函数完成，它使用 dirents 辅助函数来枚举目录中的条目，如下所示：

```go
// wakjDir 递归地遍历以 dir 为根目录的整个文件树，并在 filesizes 上发送每个已找到文件的大小
func walkDir(dir string, fileSizes chan<- int64) {
    for _, entry := range dirents(dir) {
        if entry.IsDir() {
            subdir := filepath.Join(dir, entry.Name())
            walkDir(subdir, fileSizes)
        } else {
            fileSizes <- entry.Size()
        }
    }
}
// dirents 返回 dir 目录中的条目
func dirents(dir string) []os.FileInfo {
    entries, err := ioutil.ReadDir(dir)
    if err != nil {
        fmt.Fprintf(os.Stderr, "du1: %v\n", err)
        return nil
    }
    return entries
}
```

ioutil.ReadDir 函数返回一个 os.FileInfo 类型的 slice，针对单个文件同样的信息可以通过调用 os.Stat 函数来返回。对每一个子目录，walkDir 递归调用它自己，对于每一个文件，walkDir 发送一条消息到 fileSizes 通道，消息的内容为文件所占用的字节数。

>注意⚠️：`ioutil.ReadDir` 函数已过时，最后完整版中会采用`os.readDir`进行替换

程序的完整代码如下所示，代码中 main 函数使用两个 goroutine，后台 goroutine 调用 walkDir 遍历命令行上指定的每一个目录，最后关闭 fileSizes 通道；主 goroutine 计算从通道中接收的文件的大小的和，最后输出总数。

```go
package main
import (
    "flag"
    "fmt"
    "io/ioutil"
    "os"
    "path/filepath"
)
func main() {
    // 确定初始目录
    flag.Parse()
    roots := flag.Args()
    if len(roots) == 0 {
        roots = []string{"."}
    }
    // 遍历文件树
    fileSizes := make(chan int64)
    go func() {
        for _, root := range roots {
            walkDir(root, fileSizes)
        }
        close(fileSizes)
    }()
    // 输出结果
    var nfiles, nbytes int64
    for size := range fileSizes {
        nfiles++
        nbytes += size
    }
    printDiskUsage(nfiles, nbytes)
}
func printDiskUsage(nfiles, nbytes int64) {
    fmt.Printf("%d files  %.1f GB\n", nfiles, float64(nbytes)/1e9)
}
// wakjDir 递归地遍历以 dir 为根目录的整个文件树,并在 filesizes 上发送每个已找到的文件的大小
func walkDir(dir string, fileSizes chan<- int64) {
    for _, entry := range dirents(dir) {
        if entry.IsDir() {
            subdir := filepath.Join(dir, entry.Name())
            walkDir(subdir, fileSizes)
        } else {
            fileSizes <- entry.Size()
        }
    }
}
// dirents 返回 dir 目录中的条目
func dirents(dir string) []os.FileInfo {
    entries, err := ioutil.ReadDir(dir)
    if err != nil {
        fmt.Fprintf(os.Stderr, "du1: %v\n", err)
        return nil
    }
    return entries
}
```
在输出结果前，程序等待较长时间：

```
go run main.go D:/code
18681 files  0.5 GB
```

如果程序可以通知它的进度，将会更友好，但是仅把 printDiskUsage 调用移动到循环内部会使它输出数千行结果，所以这里对上面的程序进行一些调整，在有-v标识的时候周期性的输出当前目录的总和，如果只想看到最终的结果省略-v即可。
```
package main
import (
    "flag"
    "fmt"
    "io/ioutil"
    "os"
    "path/filepath"
    "time"
)
var verbose = flag.Bool("v", false, "显示详细进度")
func main() {
    // ...启动后台 goroutine...
    // 确定初始目录
    flag.Parse()
    roots := flag.Args()
    if len(roots) == 0 {
        roots = []string{"."}
    }
    // 遍历文件树
    fileSizes := make(chan int64)
    go func() {
        for _, root := range roots {
            walkDir(root, fileSizes)
        }
        close(fileSizes)
    }()
    // 定期打印结果
    var tick <-chan time.Time
    if *verbose {
        tick = time.Tick(500 * time.Millisecond)
    }
    var nfiles, nbytes int64
loop:
    for {
        select {
        case size, ok := <-fileSizes:
            if !ok {
                break loop // fileSizes 关闭
            }
            nfiles++
            nbytes += size
        case <-tick:
            printDiskUsage(nfiles, nbytes)
        }
    }
    printDiskUsage(nfiles, nbytes) // 最终总数
}
func printDiskUsage(nfiles, nbytes int64) {
    fmt.Printf("%d files  %.1f GB\n", nfiles, float64(nbytes)/1e9)
}
// wakjDir 递归地遍历以 dir 为根目录的整个文件树,并在 filesizes 上发送每个已找到的文件的大小
func walkDir(dir string, fileSizes chan<- int64) {
    for _, entry := range dirents(dir) {
        if entry.IsDir() {
            subdir := filepath.Join(dir, entry.Name())
            walkDir(subdir, fileSizes)
        } else {
            fileSizes <- entry.Size()
        }
    }
}
// dirents 返回 dir 目录中的条目
func dirents(dir string) []os.FileInfo {
    entries, err := ioutil.ReadDir(dir)
    if err != nil {
        fmt.Fprintf(os.Stderr, "du1: %v\n", err)
        return nil
    }
    return entries
}
```

因为这个程序没有使用 range 循环，所以第一个 select 情况必须显式判断 fileSizes 通道是否已经关闭，使用两个返回值的形式进行接收操作。如果通道已经关闭，程序退出循环。标签化的 break 语句将跳出 select 和 for 循环的逻辑。没有标签的 break 只能跳出 select 的逻辑，导致循环的下一次迭代。

运行结果如下所示：
```
go run main.go -v D:\
296077 files  57.9 GB
302142 files  58.0 GB
306669 files  58.1 GB
314725 files  58.2 GB
320050 files  58.3 GB
341713 files  58.6 GB
346102 files  64.2 GB
```

此程序的弊端也很明显，它依然会耗费太长的时间。

所以，下面为每一个 walkDir 的调用创建一个新的 goroutine。它使用 sync.WaitGroup 来为当前存活的 walkDir 调用计数，一个 goroutine 在计数器减为 0 的时候关闭 fileSizes 通道。
```go
package main

import (
	"flag"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"sync"
	"time"
)

var verbose = flag.Bool("v", false, "显示详细信息")

// dirents 返回目录下的列表，包含文件和目录
func dirents(dir string) []fs.DirEntry {

/*
		ioutil.ReadDir方法已过时
		采用os.ReadDir读取文件列表并返回文件信息
	*/
	entries, err := os.ReadDir(dir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "du: %v\n", err)
	}
	return entries
}

/*
walkDir 根据输入的目录，返回目录下所有条目的文件大小
如果是文件就直接返回文件大小
如果是目录，则再调用函数自身进行子目录下文件的计算
*/
func walkDir(dir string, fileSize chan<- int64, n *sync.WaitGroup) {

// 线程结束
	defer n.Done()
	for _, entry := range dirents(dir) {
		if entry.IsDir() {
			// 添加新的等待组
			n.Add(1)
			subdir := filepath.Join(dir, entry.Name())
			// 新建线程进行统计
			go walkDir(subdir, fileSize, n)
		} else {
			fileInfo, err := entry.Info()
			if err != nil {
				fmt.Fprintf(os.Stderr, "%v\n", err)
			}
			fileSize <- fileInfo.Size()
		}
	}
}

func printDiskUsage(files, bytes int64) {
	fmt.Printf("%d files  %.1f GB\n", files, float64(bytes)/1e9)
}

func main() {
	//启动主goroutine

//解析输入的目录
	flag.Parse()
	roots := flag.Args()
	if len(roots) == 0 {
		//输入为空，没有指定统计的目录，则直接统计当前目录
		roots = []string{"."}
	}

fileSize := make(chan int64)

var n sync.WaitGroup

//遍历文件树

for _, root := range roots {
		n.Add(1)
		go walkDir(root, fileSize, &n)
	}

// 等所有线程结束后关闭
	go func() {
		n.Wait()
		close(fileSize)
	}()

// 周期性打印统计结果
	var tick <-chan time.Time
	// 如果verbose值为真，就打印
	if *verbose {
		tick = time.Tick(1 * time.Second)
	}

// 声明变量用于统计文件个数和文件大小
	var files, bytes int64

loop:
	for {
		select {
		case size, ok := <-fileSize:
			if !ok {
				break loop
			}
			files++
			bytes += size
		case <-tick:
			printDiskUsage(files, bytes)
		}

}

// 循环结束，打印最终大小
	printDiskUsage(files, bytes)
}
```
尽管系统与系统之间有很多的不同，但是这个版本的速度比前一个版本快几倍。

![](/media/202304/2023-04-09_164704_6203810.5461977614754076.png)