Created
November 21, 2017 21:32
-
-
Save dapangmao/a5ceca06de26b72517b23a6bc2b84c1b to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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