Another Go script to count folder size

du.go

package main

import (
	"fmt"
	"os"
	"runtime"
	"sync"
	"sync/atomic"
	"time"
)

func getTotalAndSubDirs(d string) {
	f, err := os.Open(d)
	if err != nil {
		panic(err)
	}
	files, err := f.Readdir(-1)
	f.Close()

	if err != nil {
		panic(err)
	}

	for _, f := range files {
		if f.IsDir() {
			select {
			case dir_chan <- d + "/" + f.Name():
			default:
				panic("Buffer not big enough")
			}
		} else {
			atomic.AddInt64(&total, f.Size())
		}
	}
}

const N = 200

func count(dir string) {
	getTotalAndSubDirs(dir)
	wait := &sync.WaitGroup{}
	for i := 0; i < N; i++ {
		wait.Add(1)
		go func() {
			for done := false; !done; {
				select {
				case f := <-dir_chan:
					getTotalAndSubDirs(f)
				default:
					done = true
				}
			}
			wait.Done()
		}()
	}
	wait.Wait()
}

var dir string
var total int64
var dir_chan chan string

func init() {
	if len(os.Args) != 2 {
		fmt.Println("Usage: du <dir>")
		os.Exit(1)
	}
	dir = os.Args[1]
	dir_chan = make(chan string, 10000) //这里写死了10000。。。
        runtime.GOMAXPROCS(runtime.NumCPU())
	total = 0
}

func main() {
	t1 := time.Now().UnixNano()
	count(dir)
	fmt.Println("total:", total)
	t2 := time.Now().UnixNano()
	fmt.Println("time cost:", (t2-t1)/1000000)
}