goConcurrency/demoWalkDir.go

package goConcurrency

// 并发遍历目录统计信息的Demo
import (
	"fmt"
	"io/fs"
	"log"
	"os"
	"path/filepath"
	"sync"
)

// WalkDir 外部调用的遍历目录统计信息的方法
func WalkDir(dirs ...string) string {
	// 一：保证至少有一个目录需要统计遍历
	// 默认为当前目录
	if len(dirs) == 0 {
		dirs = []string{"."}
	}

	// 二：初始化变量，channel用于完成Size的传递，WaitGroup用于等待调度
	filesizeCh := make(chan int64, 1)
	wg := &sync.WaitGroup{}

	// 三：启动多个Goroutine统计信息，取决于len(dirs)
	for _, dir := range dirs {
		wg.Add(1) // +1
		// 并发的遍历统计每个目录的信息
		go walkDir(dir, filesizeCh, wg)
	}

	// 四：启动累计运算的Goroutine
	// 1.用户关闭filesizeCh
	go func(filesizeCh chan int64, wg *sync.WaitGroup) {
		//等待统计工作完成
		wg.Wait()
		// 关闭filesizeCh
		close(filesizeCh)
	}(filesizeCh, wg)
	// 2.range的方式从filesizeCh中获取文件大小
	// 3.将统计结果，使用channel传递出来
	fileNumCh := make(chan int64, 1)
	sizeTotalCh := make(chan int64, 1)
	go func(filesizeCh <-chan int64, fileNumCh, sizeTotalCh chan<- int64) {
		// 统计文件数，和文件整体大小
		var fileNum, sizeTotal int64
		// 遍历全部的filesizeCh元素，统计文件数量和大小，直到channel被关闭
		for filesize := range filesizeCh {
			// 累计文件数，和统计文件整体大小
			fileNum++
			sizeTotal += filesize
		}
		// 将统计结果发送到对于的Channel中
		fileNumCh <- fileNum
		sizeTotalCh <- sizeTotal
	}(filesizeCh, fileNumCh, sizeTotalCh)

	// 五：整理返回值
	// size 的单位 Byte
	// 需要的单位 MB
	result := fmt.Sprintf("%d files %.2f MB\n", <-fileNumCh, float64(<-sizeTotalCh)/1e6) // 1024*1024
	return result
}

// 遍历并统计某个特定目录的信息
// 核心实现函数，完成递归，统计等
func walkDir(dir string, filesizeCh chan<- int64, wg *sync.WaitGroup) {
	// 一：wg计数器减少
	defer wg.Done()
	// 二：读取dir下的全部文件信息，并遍历
	for _, fileinfo := range fileInfos(dir) {
		// 三：根据dir下的文件信息
		if fileinfo.IsDir() {
			// 1. 如果目录，递归获取信息
			// 子目录地址
			subDir := filepath.Join(dir, fileinfo.Name())
			// 递归调用，也是并发的，也需要wg统计
			wg.Add(1)
			go walkDir(subDir, filesizeCh, wg)
		} else {
			// 2. 如果不是，就是文件，统计文件大小，放入channel
			filesizeCh <- fileinfo.Size() //Byte
		}
	}

}

// 获取某个目录下文件信息列表
func fileInfos(dir string) []fs.FileInfo {
	// 一，读取目录的全部文件
	entries, err := os.ReadDir(dir)
	if err != nil {
		log.Println("WalkDir error:", err)
		return []fs.FileInfo{} // nil
	}

	// 二，获取文件的文件信息
	// DirEntry to FileInfo
	infos := make([]fs.FileInfo, 0, len(entries))
	for _, entry := range entries {
		// 如果获取文件信息无错误，存储到infos中。
		if info, err := entry.Info(); err == nil {
			infos = append(infos, info)
		}
	}

	// 三，返回
	return infos
}