go语言程序设计，练习8.9

------------恢复内容开始------------

package main

import (

    "flag"

    "fmt"

    "io/ioutil"

    "os"

    "path/filepath"

    "sync"

    "time"

)

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

var sema = make(chan struct{}, 20)

func main() {

    var wg sync.WaitGroup

    fileSizes := make(chan map[string]int64)

    nbytes := make(map[string]int64) // 改动点

    nfile := make(map[string]int64)

    flag.Parse()

    roots := flag.Args()

    if len(roots) == 0 {

        roots = []string{"."}

    }

    var tick <-chan time.Time

    if *verbose {

        tick = time.Tick(500 * time.Millisecond)

    }

    for _, root := range roots {

        wg.Add(1)

        go walk(root, fileSizes, &wg)

    }

    go func() {

        wg.Wait()

        close(fileSizes)

    }()

loop:

    for {

        select {

        case <-tick:

            printdir(nfile, nbytes)

            fmt.Println("-------")

        case size, ok := <-fileSizes:

            if !ok {

                break loop

            }

            for key, value := range size { // 改动点

                nbytes[key] += value

                nfile[key]++

            }

        }

    }

    printdir(nfile, nbytes)

}

// 改动点

// 为每一个路径 起一个 walkDir 去统计该路径下的大小，使用map型chan 传递消息

func walk(dir string, dirsize chan map[string]int64, n *sync.WaitGroup) {

    defer n.Done()

    fileSizes := make(chan int64)

    var wg sync.WaitGroup

    wg.Add(1)

    go walkDir(dir, fileSizes, &wg)

    n.Add(1)

    go func() {

        for size := range fileSizes {

            dirsize <- map[string]int64{dir: size}

        }

        n.Done()

    }()

    wg.Wait()

    close(fileSizes)

}

func walkDir(dir string, fileSizes chan int64, n *sync.WaitGroup) {

    defer n.Done()

    for _, entry := range dirents(dir) {

        if entry.IsDir() {

            subdir := filepath.Join(dir, entry.Name())

            n.Add(1)

            go walkDir(subdir, fileSizes, n)

        } else {

            fileSizes <- entry.Size()

        }

    }

}

func dirents(dir string) []os.FileInfo {

    // 限制 打开的文件数

    sema <- struct{}{}

    defer func() { <-sema }()

    entries, err := ioutil.ReadDir(dir)

    if err != nil {

        fmt.Println("read dir faild, ", err)

        return nil

    }

    return entries

}

func printdir(n, dir map[string]int64) {

    for key, size := range dir {

        fmt.Printf("root: %v :", key)

        if size/1e9 >= 1 {

            fmt.Printf("n: %v, size: \t%.1f GB\n", n[key], float64(size)/1e9)

        } else if size/1e6 >= 1 {

            fmt.Printf("n: %v, size: \t%.1f MB\n", n[key], float64(size)/1e6)

        } else if size/1e3 >= 1 {

            fmt.Printf("n: %v, size: \t%.1f KB\n", n[key], float64(size)/1e3)

        } else {

            fmt.Printf("n: %v, size: \t%v B\n", n[key], size)

        }

    }

}

func printdis(n, size int64) {

    if size/1e9 >= 1 {

        fmt.Printf("\rn: \t%v, size: \t%.1f GB  ", n, float64(size)/1e9)

    } else if size/1e6 >= 1 {

        fmt.Printf("\rn: \t%v, size: \t%.1f MB", n, float64(size)/1e6)

    } else if size/1e3 >= 1 {

        fmt.Printf("\rn: %v, size: \t%.1f KB", n, float64(size)/1e3)

    } else {

        fmt.Printf("\rn: %v, size: \t%v B", n, size)

    }

}

------------恢复内容结束------------

------------恢复内容结束------------