Skip to content

Instantly share code, notes, and snippets.

@g14a
Last active Dec 5, 2021
Embed
What would you like to do?
package main
import (
"flag"
"fmt"
"os"
"path/filepath"
"sync"
"time"
)
var maxConcurrency int
func main() {
flag.IntVar(&maxConcurrency, "p", 1, "")
flag.Parse()
roots := flag.Args()
if len(roots) == 0 {
roots = []string{"."}
}
now := time.Now()
fileSizes := make(chan int64)
var wg sync.WaitGroup
var lock = make(chan struct{}, maxConcurrency)
for _, root := range roots {
wg.Add(1)
go walkDirectory(root, &wg, fileSizes, lock)
}
go func() {
wg.Wait()
close(fileSizes)
}()
var nFiles, nBytes int64
for size := range fileSizes {
nFiles++
nBytes += size
}
fmt.Println("total time: ", time.Since(now))
printUsage(nFiles, nBytes)
}
func printUsage(files, bytes int64) {
fmt.Printf("%d files %.1f GB\n", files, float64(bytes)/1e9)
}
func walkDirectory(dir string, wg *sync.WaitGroup, fileSize chan<- int64, lock chan struct{}) {
defer wg.Done()
for _, entry := range getEntries(dir, lock) {
if entry.IsDir() {
wg.Add(1)
subdir := filepath.Join(dir, entry.Name())
go walkDirectory(subdir, wg, fileSize, lock)
} else {
info, _ := entry.Info()
fileSize <- info.Size()
}
}
}
func getEntries(dir string, lock chan struct{}) []os.DirEntry {
lock <- struct{}{}
defer func() { <-lock }()
entries, err := os.ReadDir(dir)
if err != nil {
panic(err)
}
return entries
}
package main
import (
"flag"
"fmt"
"os"
"path/filepath"
"time"
)
func main() {
flag.Parse()
roots := flag.Args()
if len(roots) == 0 {
roots = []string{"."}
}
now := time.Now()
var totalFiles, totalBytes int64
for _, root := range roots {
files, bytes := walkDirectory(root)
totalFiles += files
totalBytes += bytes
}
fmt.Println("total time: ", time.Since(now))
printUsage(totalFiles, totalBytes)
}
func printUsage(files, bytes int64) {
fmt.Printf("%d files %.1f GB\n", files, float64(bytes)/1e9)
}
func walkDirectory(dir string) (noOfFiles, totalSize int64) {
for _, entry := range getEntries(dir) {
if entry.IsDir() {
subDir := filepath.Join(dir, entry.Name())
files, bytes := walkDirectory(subDir)
noOfFiles += files
totalSize += bytes
} else {
noOfFiles++
e, err := entry.Info()
if err != nil {
panic(err)
}
totalSize += e.Size()
}
}
return
}
func getEntries(dir string) []os.DirEntry {
entries, err := os.ReadDir(dir)
if err != nil {
return nil
}
return entries
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment