多线程Hash文件

hash_files.go

package main

import (
	"crypto/sha256"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strconv"
	"sync"
	"time"

	log "github.com/Sirupsen/logrus"
)

type FileInfo struct {
	Path string
	Hash string
	Size int64
}

type walkArgs struct {
	fp   string
	info os.FileInfo
	e    error
}

type Incer struct {
	sync.Mutex
	count int
}

func (i *Incer) Inc() {
	i.Lock()
	i.count++
	i.Unlock()
}

func (i *Incer) String() string {
	return strconv.Itoa(i.count)
}

func main() {
	var walkc, hashc, outc Incer
	start := time.Now()
	defer func() {
		log.Infof("use %s ", time.Now().Sub(start).String())
	}()
	buf := make(chan *FileInfo, 100)
	wksbuf := make(chan *FileInfo, 100)
	done := make(chan struct{})

	MaxRunc := 10
	wg := new(sync.WaitGroup)

	for i := 0; i < MaxRunc; i++ {
		go func() {
			for {
				select {
				case file := <-wksbuf:
					f, err := os.Open(file.Path)
					if err != nil {
						log.Error(err)
						wg.Done()
						return
					}

					hash := sha256.New()
					_, err = io.Copy(hash, f)
					if err != nil {
						log.Error(err)
						wg.Done()
						f.Close()
						return
					}

					file.Hash = fmt.Sprintf("%x", hash.Sum(nil))
					// hashc.Inc()
					// log.Infof("Hash: %s", file.Path)
					// log.Infof("file: %+v", file)
					f.Close()
					buf <- file

				case <-done:
					return
				}
			}
		}()
	}

	go func() {
		for {
			select {
			case fi := <-buf:
				outc.Inc()
				wg.Done()
				log.Infof("Out File: %+v", fi)
			case <-done:
				return
			}
		}
	}()

	filepath.Walk("/", func(fp string, info os.FileInfo, e error) error {
		if e != nil {
			log.Error(e)
			return e
		}
		if info.IsDir() {
			return nil
		}

		wg.Add(1)
		// walkc.Inc()
		// log.Infof("Walk: %s", fp)
		wksbuf <- &FileInfo{
			Path: fp,
			Size: info.Size(),
		}
		return nil
	})
	wg.Wait()
	close(done)

	log.Infof("Walk dir done. walk: %v, hash: %v, out: %v", walkc, hashc, outc)
}