Add WaitGroup for HashingWorker

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)
+}

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 <path>")
@@ -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
 	}