A Go script to count folder size.

du.go

package main

import (
	"os"
	"runtime"
	"log"
	"time"
)

var dir string

type SubDirsAndSize struct {
	Path    string
	Size    int64
	SubDirs []string
}

func getTotalAndSubDirs(d string) *SubDirsAndSize {
	f, err := os.Open(d)
	if err != nil {
		log.Fatal("Open dir failed: ", err)
	}
	files, err := f.Readdir(-1)
	f.Close()
	if err != nil {
		log.Fatal("Read dir failed: ", err)
	}

	total := int64(0)
	subdirs := []string{}
	for _, f := range files {
		if f.IsDir() {
			subdirs = append(subdirs, d+"/"+f.Name())
		} else {
			total += f.Size()
		}
	}
	return &SubDirsAndSize{
		Path:    d,
		Size:    total,
		SubDirs: subdirs,
	}
}

const N = 250

func count(dir string) int64 {
	root := getTotalAndSubDirs(dir)
	total := root.Size
	dirList := root.SubDirs
	semaphore := make(chan bool, N)
	for n := len(dirList); n > 0; {
		tasks := make(chan *SubDirsAndSize, n)
		for _, sub := range dirList {
			go func(sub string) {
				semaphore <- true
				tasks <- getTotalAndSubDirs(sub)
				<-semaphore
			}(sub)
		}
		dirList = []string{}
		for i := 0; i < n; i++ {
			c := <-tasks
			total += c.Size
			if len(c.SubDirs) > 0 {
				dirList = append(dirList, c.SubDirs...)
			}
		}
		close(tasks)
		n = len(dirList)
	}
	return total
}

func init() {
	if len(os.Args) != 2 {
		log.Println("Usage: du <dir>")
		os.Exit(1)
	}
	dir = os.Args[1]
	runtime.GOMAXPROCS(runtime.NumCPU())
}

func main() {
	t1 := time.Now().UnixNano()

	total := count(dir)

	t2 := time.Now().UnixNano()

	log.Println("time cost:", (t2-t1)/1000000)
	log.Println("total:", total)
}