golang：併發目錄遍歷

《GO程序設計語言》設計中案例，僅作爲筆記進行收藏。併發遍歷目錄實現了統計文件個數及文件總大小。

package main

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

var vFlag = flag.Bool("v", false, "show verbose progress messages")

func main() {
	flag.Parse()

	// 確定初始目錄
	roots := os.Args[1:]
	if len(roots) == 0 {
		roots = []string{"."}
	}

	// 當檢測到輸入時取消遍歷
	go func() {
		// 讀一個字節
		os.Stdin.Read(make([]byte, 1))
		close(done)
	}()

	// 並行遍歷文件樹的每個根
	fileSizes := make(chan int64)
	var n sync.WaitGroup
	for _, root := range roots {
		n.Add(1)
		go walkDir(root, &n, fileSizes)
	}
	go func() {
		n.Wait()
		close(fileSizes)
	}()

	// 定期打印結果
	// tick := time.Tick(500 * time.Millisecond)
	var tick <-chan time.Time
	if *vFlag {
		tick = time.Tick(500 * time.Millisecond)
	}

	var nfiles, nbytes int64
loop:
	for {
		select {
		case <-done:
			// 耗盡 fileSizes 以允許已有的  goroutine 結束
			for range fileSizes {
				// 不執行任何操作
			}
			return
		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)
}

// walkDir 遞歸地遍歷以dir爲根目錄的整個文件樹
// 並在 fileSizes 上發送每個已找到的文件的大小
func walkDir(dir string, n *sync.WaitGroup, fileSizes chan<- int64) {
	defer n.Done()
	if cancelled() {
		return
	}

	for _, entry := range dirents(dir) {
		if entry.IsDir() {
			n.Add(1)
			subdir := filepath.Join(dir, entry.Name())
			go walkDir(subdir, n, fileSizes)
		} else {
			fileSizes <- entry.Size()
		}
	}
}

// sema 是一個用於限制目錄併發數的計數信號量
var sema = make(chan struct{}, 20)

// dirents 返回 dir 目錄中的條目
func dirents(dir string) []os.FileInfo {
	select {
	case sema <- struct{}{}: // 獲取令牌
	case <-done:
		return nil // 取消
	}
	defer func() { <-sema }() // 釋放令牌

	// 讀取目錄
	f, err := os.Open(dir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "du：%v\n", err)
		return nil
	}
	defer f.Close()

	// 不做限制，讀取所有條目
	entries, err := f.Readdir(0)
	if err != nil {
		fmt.Fprintf(os.Stderr, "du：%v\n", err)
	}
	return entries
}

// 創建取消通道
var done = make(chan struct{})

func cancelled() bool {
	select {
	case <-done:
		return true
	default:
		return false
	}
}

併發遍歷目錄代碼，實現了 命令行參數 -v 和輸入取消 。

1.如果帶有 -v ，則定時打印遍歷結果；否則只打印最終結果。

2.在遍歷目錄時，輸入任意一個字符，即將取消遍歷。