[Daily] Go Language Bible - Example: Concurrent Directory Traversal Exercises

Exercise 8.9: Write a du tool that calculates and displays the size of directories in the root directory at regular intervals.

package main

import (
        //      "filepath"
        "flag"
        "fmt"
        "io/ioutil"
        "os"
        "path"
        "sync"
        "time"
)
/*
Exercise 8.9: Write a du tool that calculates and displays the size of directories in the root directory at regular intervals.
*/

//Receive command line parameter -v
var verbose = flag.Bool("v", false, "show verbose progress messages")

func main() {
        //Receive command line parameters, multiple paths
        flag.Parse()
        roots := flag.Args()
        //If no path is passed, give default value
        if len(roots) == 0 {
                roots = []string{"/"}
        }   

        for {
                sumFileSize(roots)
                time.Sleep(20 * time.Second)
        }   
}

func sumFileSize(roots []string) {
        //Send and receive channel of file size in bytes
        fileSizes := make(chan int64)
        // counter for goroutine
        var n sync.WaitGroup
        //The path passed by the loop command line
        for _, root := range roots {
                n.Add(1)
                //Start goroutine calculation
                go walkDir(root, &n, fileSizes)
        }   
        //Start goroutine and wait for all goroutines that compute directories to end
        go func() {
                n.Wait()
                close(fileSizes)
        }()
        //Timely display the channel sent by the directory progress
        var tick <-chan time.Time
        if *verbose {
                tick = time.Tick(500 * time.Millisecond)
        }   
        var nfiles, nbytes int64
        //select and loop loop, multiplexing
loop:
        for {
                select {
                case size, ok := <-fileSizes:
                        if !ok {
                                break loop // fileSizes was closed
                        }
                        / / Calculate the number of directories, calculate the size in bytes
                        nfiles++
                        nbytes += size
                case <-tick:
                        //Receive timing channel print progress
                        printDiskUsage(nfiles, nbytes)
                }
        }
        //last print total
        printDiskUsage(nfiles, nbytes) // final totals
}

func walkDir(dir string, n *sync.WaitGroup, fileSizes chan<- int64) {
        defer n.Done()
        for _, entry := range dirents(dir) {
                if entry.IsDir() {
                        n.Add(1)
                        subdir := path.Join(dir, entry.Name())
                        //Open multiple goroutines for recursion
                        go walkDir(subdir, n, fileSizes)
                } else {
                        fileSizes <- entry.Size()
                }
        }
}

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

// dirents returns the entries of directory dir.
func dirents (dir string) [] os.FileInfo {
        // Limit too much concurrency with counting semaphore logic
        sema <- struct{}{}
        entries, err := ioutil.ReadDir(dir)
        <-sema
        if err != nil {
                fmt.Fprintf(os.Stderr, "du1: %v\n", err)
                return nil
        }
        return entries
}
func printDiskUsage(nfiles, nbytes int64) {
        fmt.Printf("%d files  %.1f GB\n", nfiles, float64(nbytes)/1e9)
}