multithread support

This commit is contained in:
GenZmeY 2020-11-02 14:19:16 +03:00
parent 3ba52e60cf
commit 7915b63f88

75
main.go
View File

@ -31,9 +31,6 @@ const (
var ( var (
Version string = "dev" Version string = "dev"
hashes map[string]string
names []string
) )
func main() { func main() {
@ -50,7 +47,7 @@ func main() {
os.Exit(EXIT_ARG_ERR) os.Exit(EXIT_ARG_ERR)
} }
files, err := ioutil.ReadDir(ArgInput) allFiles, err := ioutil.ReadDir(ArgInput)
if err != nil { if err != nil {
output.Errorln("Read dir error") output.Errorln("Read dir error")
os.Exit(EXIT_DIR_READ_ERR) os.Exit(EXIT_DIR_READ_ERR)
@ -61,18 +58,11 @@ func main() {
} }
runtime.GOMAXPROCS(ArgJobs) runtime.GOMAXPROCS(ArgJobs)
hashes, names := calcHashes(files) pngFiles := pngList(allFiles)
hashes := calcHashes(pngFiles)
ranges := calcRanges(hashes, names, Threshold) ranges := calcRanges(hashes, Threshold)
writeRanges(ranges) writeRanges(ranges)
wg := new(sync.WaitGroup)
//wg.Add(1)
wg.Wait()
os.Exit(EXIT_SUCCESS) os.Exit(EXIT_SUCCESS)
} }
@ -86,43 +76,54 @@ func closeHandler() {
}() }()
} }
func calcHashes(files []os.FileInfo) (map[string]string, []string) { func pngList(files []os.FileInfo) []os.FileInfo {
var pngFiles []os.FileInfo
for _, file := range files {
if !file.IsDir() && strings.HasSuffix(strings.ToLower(file.Name()), ".png") {
pngFiles = append(pngFiles, file)
}
}
return pngFiles
}
func calcHashes(files []os.FileInfo) map[string]string {
var hashes map[string]string var hashes map[string]string
var keys []string var mutex sync.Mutex
var err error
wg := new(sync.WaitGroup)
wg.Add(len(files))
hashes = make(map[string]string) hashes = make(map[string]string)
for _, file := range files { for _, file := range files {
if !file.IsDir() && strings.HasSuffix(file.Name(), ".png") { go calcHash(ArgInput, file, &hashes, &mutex, wg)
keys = append(keys, file.Name())
hashes[file.Name()], err = calcHash(filepath.Join(ArgInput, file.Name()))
if err != nil {
output.Errorln(err)
}
}
} }
sort.Strings(keys) wg.Wait()
return hashes, keys
return hashes
} }
func calcHash(filepath string) (string, error) { func calcHash(path string, fileinfo os.FileInfo, hashes *map[string]string, mutex *sync.Mutex, wg *sync.WaitGroup) {
file, err := os.Open(filepath) file, err := os.Open(filepath.Join(path, fileinfo.Name()))
if err != nil { if err != nil {
return "", err output.Errorln(err)
} }
defer file.Close() defer file.Close()
image, _, err := image.Decode(file) image, _, err := image.Decode(file)
if err != nil { if err != nil {
return "", err output.Errorln(err)
} }
hash := blockhash.NewBlockhash(image, 16).Hexdigest()
return blockhash.NewBlockhash(image, 16).Hexdigest(), nil mutex.Lock()
(*hashes)[fileinfo.Name()] = hash
mutex.Unlock()
wg.Done()
} }
func calcRanges(hashes map[string]string, names []string, Threshold int) string { func calcRanges(hashes map[string]string, Threshold int) string {
var ranges strings.Builder var ranges strings.Builder
var prevHash string = "" var prevHash string = ""
var dist int = 0 var dist int = 0
@ -133,6 +134,12 @@ func calcRanges(hashes map[string]string, names []string, Threshold int) string
rangeNoise = "\t" + string(ArgDefaultNoiseLevel) rangeNoise = "\t" + string(ArgDefaultNoiseLevel)
} }
names := []string{}
for key := range hashes {
names = append(names, key)
}
sort.Strings(names)
for i := 0; i < len(names); i++ { for i := 0; i < len(names); i++ {
name := names[i] name := names[i]
if startName == "" { if startName == "" {