diff --git a/filesmap.go b/filesmap.go index 13d0359..2bff489 100644 --- a/filesmap.go +++ b/filesmap.go @@ -1,8 +1,11 @@ package main import ( + "flag" "fmt" "log" + "os" + "path/filepath" "sync" ) @@ -53,7 +56,7 @@ func (fm *FilesMap) IncomingWorker() { close(fm.FilesHashing) } -func (fm *FilesMap) HashingWorker() { +func (fm *FilesMap) HashingWorker(wg *sync.WaitGroup) { for file := range fm.FilesHashing { if *verbose { fmt.Println("Hashing", file.path) @@ -69,7 +72,7 @@ func (fm *FilesMap) HashingWorker() { file.hash = hash fm.FilesHashed <- file } - close(fm.FilesHashed) + wg.Done() } func (fm *FilesMap) HashedWorker(done chan bool) { @@ -89,3 +92,18 @@ func (fm *FilesMap) HashedWorker(done chan bool) { done <- true } + +func (fm *FilesMap) WalkDirectories() { + for _, path := range flag.Args() { + filepath.Walk(path, func(path string, info os.FileInfo, err error) error { + if info.IsDir() { + return nil + } + + fm.FilesIncoming <- fileEntry{path, info.Size(), ""} + return nil + }) + } + + close(fm.FilesIncoming) +} diff --git a/main.go b/main.go index bef6647..bc66df9 100644 --- a/main.go +++ b/main.go @@ -8,8 +8,10 @@ import ( "io/ioutil" "os" "path/filepath" + "runtime" "strconv" "strings" + "sync" ) var fromFile = flag.String("from-file", "", "Load results file from ") @@ -36,26 +38,20 @@ func main() { } } else { done := make(chan bool) - //for i := 0; i < runtime.GOMAXPROCS(0); i++ { - go filesMap.HashingWorker() - //} + wg := sync.WaitGroup{} + for i := 0; i < runtime.GOMAXPROCS(0); i++ { + wg.Add(1) + go filesMap.HashingWorker(&wg) + } go filesMap.IncomingWorker() go filesMap.HashedWorker(done) - for _, path := range flag.Args() { - filepath.Walk(path, func(path string, info os.FileInfo, err error) error { - if info.IsDir() { - return nil - } + filesMap.WalkDirectories() - filesMap.FilesIncoming <- fileEntry{path, info.Size(), ""} - return nil - }) - } - - close(filesMap.FilesIncoming) + wg.Wait() + close(filesMap.FilesHashed) <-done }