Login light
// Copyright © 2016 Alan A. A. Donovan & Brian W. Kernighan.
// License: https://creativecommons.org/licenses/by-nc-sa/4.0/

// See page 251.

// The du4 command computes the disk usage of the files in a directory.
package main

// The du4 variant includes cancellation:
// it terminates quickly when the user hits return.

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

//!+1
var done = make(chan struct{})

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

//!-1

func main() {
    // Determine the initial directories.
    roots := os.Args[1:]
    if len(roots) == 0 {
        roots = []string{"."}
    }

    //!+2
    // Cancel traversal when input is detected.
    go func() {
        os.Stdin.Read(make([]byte, 1)) // read a single byte
        close(done)
    }()
    //!-2

    // Traverse each root of the file tree in parallel.
    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)
    }()

    // Print the results periodically.
    tick := time.Tick(500 * time.Millisecond)
    var nfiles, nbytes int64
loop:
    //!+3
    for {
        select {
        case <-done:
            // Drain fileSizes to allow existing goroutines to finish.
            for range fileSizes {
                // Do nothing.
            }
            return
        case size, ok := <-fileSizes:
            // ...
            //!-3
            if !ok {
                break loop // fileSizes was closed
            }
            nfiles++
            nbytes += size
        case <-tick:
            printDiskUsage(nfiles, nbytes)
        }
    }
    printDiskUsage(nfiles, nbytes) // final totals
}

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

// walkDir recursively walks the file tree rooted at dir
// and sends the size of each found file on fileSizes.
//!+4
func walkDir(dir string, n *sync.WaitGroup, fileSizes chan<- int64) {
    defer n.Done()
    if cancelled() {
        return
    }
    for _, entry := range dirents(dir) {
        // ...
        //!-4
        if entry.IsDir() {
            n.Add(1)
            subdir := filepath.Join(dir, entry.Name())
            go walkDir(subdir, n, fileSizes)
        } else {
            fileSizes <- entry.Size()
        }
        //!+4
    }
}

//!-4

var sema = make(chan struct{}, 20) // concurrency-limiting counting semaphore

// dirents returns the entries of directory dir.
//!+5
func dirents(dir string) []os.FileInfo {
    select {
    case sema <- struct{}{}: // acquire token
    case <-done:
        return nil // cancelled
    }
    defer func() { <-sema }() // release token

    // ...read directory...
    //!-5

    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) // 0 => no limit; read all entries
    if err != nil {
        fmt.Fprintf(os.Stderr, "du: %v\n", err)
        // Don't return: Readdir may return partial results.
    }
    return entries
}