Concurrent Directory Traversal: program that reports the disk usage of one or more directories specified on the command line, like the Unix du command.

readdir.go

// go build -o du && ./du  $HOME /usr /bin
// Root: /bin - 1781 files 0.5 GB
// Root: /usr - 189710 files 7.1 GB
// Root: /home/zucchinidev - 849059 files 39.4 GB

package main

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

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

func main() {

	flag.Parse()
	roots := flag.Args()
	if len(roots) == 0 {
		roots = []string{"."}
	}

	go func() {
		_, _ = os.Stdin.Read(make([]byte, 1)) // read a simple byte
		close(done) // cancelling concurrent operations when a key is pressed
	}()

	var tick <-chan time.Time
	if *verbose {
		tick = time.Tick(500 * time.Microsecond)
	}
	var n sync.WaitGroup

	for _, root := range roots {
		n.Add(1)
		go walk(root, tick, &n)
	}
	n.Wait()
}

func walk(root string, tick <-chan time.Time, generalWaitGroup *sync.WaitGroup) {
	fileSizes := make(chan int64)
	var nfiles, nbytes int64
	var n sync.WaitGroup

	n.Add(1)

	go func() {
		n.Wait()
		close(fileSizes)
	}()

	go walkDir(root, &n, fileSizes)

loop:
	for {
		select {
		case <-done:
			// Drain fileSizes to allow existing goroutines to finish
			for range fileSizes {
			}
			generalWaitGroup.Done()
			break loop
		case size, ok := <-fileSizes:
			nfiles++
			nbytes += size
			if !ok {
				printDiskUsage(root, nfiles, nbytes) // final totals
				generalWaitGroup.Done()
				break loop
			}
		case <-tick:
			printDiskUsage(root, nfiles, nbytes)
		}
	}
}

func printDiskUsage(root string, nfiles int64, nbytes int64) {
	fmt.Printf("Root: %s - %d files %.1f GB\n", root, nfiles, float64(nbytes)/1e9)
}

func walkDir(dir string, n *sync.WaitGroup, fileSizes chan<- int64) {
	defer n.Done()
	if cancelled() {
		return
	}
	for _, entry := range dirEntries(dir) {
		if entry.IsDir() {
			n.Add(1)
			subdir := filepath.Join(dir, entry.Name())
			walkDir(subdir, n, fileSizes)
		} else {
			fileSizes <- entry.Size()
		}
	}
}

// semaphore is a counting semaphore for limiting concurrency in dirents.
var semaphore = make(chan struct{}, 20)

func dirEntries(dirname string) []os.FileInfo {
	select {
	case semaphore <- struct{}{}: // acquire token
	case <-done:
		return nil // cancelled
	}
	defer func() { <-semaphore }() // release token
	files, err := ioutil.ReadDir(dirname)
	if err != nil {
		_, _ = fmt.Fprintf(os.Stderr, "error when in the reading of the directory %s", err)
		return nil
	}
	return files
}

var done = make(chan struct{})

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