Skip to content

Instantly share code, notes, and snippets.

@dapangmao
Created November 21, 2017 21:32
Show Gist options
  • Save dapangmao/a5ceca06de26b72517b23a6bc2b84c1b to your computer and use it in GitHub Desktop.
Save dapangmao/a5ceca06de26b72517b23a6bc2b84c1b to your computer and use it in GitHub Desktop.
package main
import (
"io/ioutil"
"path/filepath"
"os"
"fmt"
"sync"
"flag"
"time"
)
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 := filepath.Join(dir, entry.Name())
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 {
sema <- struct{}{} // acquire token
defer func() { <-sema }() // release token
entries, err := ioutil.ReadDir(dir)
if err != nil {
fmt.Fprintf(os.Stderr, "du1: %v\n", err)
return nil
}
return entries
}
var verbose = flag.Bool("v", false, "show verbose progress messages")
func main() {
flag.Parse()
roots := flag.Args()
if len(roots) == 0 {
roots = []string{"."}
}
// Traverse the file tree.
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)
}()
var tick <-chan time.Time
if *verbose {
tick = time.Tick(500 * time.Millisecond)
}
var nfiles, nbytes int64
loop:
for {
select {
case size, ok := <-fileSizes:
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 MB\n", nfiles, float64(nbytes)/1e6)
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment