支持run

cmd/saastool/task_make.go

@@ -28,16 +28,20 @@ type makeTaskParams struct {
 
 // 计算任务
 type hashTask struct {
-	chunk []byte
-	index int
+	chunk     *[]byte
+	hash      string
+	blockSize uint64
+	index     int
 }
 
+/*
 // 计算结果
 type hashResult struct {
 	hash      string
 	blockSize uint64
 	index     int
 }
+*/
 
 func RunTaskMake(args ...string) error {
 	fs := flag.NewFlagSet("make", flag.ExitOnError)
@@ -56,7 +60,15 @@ func RunTaskMake(args ...string) error {
 		return nil
 	}
 
-	if blockSize < blockSizeMin || blockSize > blockSizeMax {
+	blockSizeNum, err := ParseByteSize(*blockSize)
+	if err != nil {
+		fmt.Println("Error parsing block size", "err", err)
+		fmt.Println("Using default 50M")
+		blockSizeNum = 50 * 1024 * 1024
+
+	}
+
+	if blockSizeNum < blockSizeMin || blockSizeNum > blockSizeMax {
 		fmt.Println("block size error", "min", blockSizeMin, "max", blockSizeMax)
 		return nil
 	}
@@ -65,7 +77,7 @@ func RunTaskMake(args ...string) error {
 		sourcePath: *sourcePath,
 		hashFile:   *hashFile,
 		task: &saasapi.Task{
-			TaskBlockSize:   blockSize,
+			TaskBlockSize:   blockSizeNum,
 			TaskDescription: *desc,
 		},
 	}
@@ -114,8 +126,15 @@ func doTaskMake(makeTaskParams makeTaskParams) error {
 		return err
 	}
 
-	tasks := make(chan hashTask)
-	results := make(chan hashResult)
+	totalSize := uint64(fi.Size())
+	// 计算读取次数
+	readTimes := int(totalSize / makeTaskParams.task.TaskBlockSize)
+	if totalSize%makeTaskParams.task.TaskBlockSize != 0 {
+		readTimes++
+	}
+
+	tasks := make(chan *hashTask)
+	results := make(chan *hashTask)
 
 	// 启动工作协程
 	hashMaxWorker := runtime.GOMAXPROCS(0)
@@ -123,18 +142,27 @@ func doTaskMake(makeTaskParams makeTaskParams) error {
 		go hashWorker(tasks, results)
 	}
 
+	// 初始化读缓存
+	readBuffers := make([][]byte, hashMaxWorker)
+	for i := range hashMaxWorker {
+		readBuffers[i] = make([]byte, makeTaskParams.task.TaskBlockSize)
+	}
+
 	wg := sync.WaitGroup{}
 	wg.Add(1)
 	go func() {
-		index := 0
-		buffer := make([]byte, makeTaskParams.task.TaskBlockSize)
-		for {
-			n, err := sourceFile.Read(buffer)
+		for index := range readTimes {
+			buffer := &readBuffers[index%hashMaxWorker]
+			n, err := sourceFile.Read(*buffer)
+
 			if n > 0 {
 				wg.Add(1)
-				fmt.Printf("\rhashing file [%v], block [%v]", makeTaskParams.sourcePath, index)
-				tasks <- hashTask{chunk: buffer[:n], index: index}
-				index++
+				//fmt.Printf("\rhashing file [%v], block [%v]\n", makeTaskParams.sourcePath, index)
+				tasks <- &hashTask{
+					chunk:     buffer,
+					index:     index,
+					blockSize: uint64(n),
+				}
 			}
 			if err != nil {
 				break
@@ -144,10 +172,12 @@ func doTaskMake(makeTaskParams makeTaskParams) error {
 		wg.Done()
 	}()
 
-	var allResults []hashResult
+	// 接收结果
+	var allResults []*hashTask
 	go func() {
 		for r := range results {
 			allResults = append(allResults, r)
+			fmt.Printf("\rhashed file [%v], block [%v]", makeTaskParams.sourcePath, r.index)
 			wg.Done()
 		}
 	}()
@@ -178,12 +208,14 @@ func doTaskMake(makeTaskParams makeTaskParams) error {
 }
 
 // hash计算协程
-func hashWorker(tasks <-chan hashTask, results chan<- hashResult) {
+func hashWorker(tasks <-chan *hashTask, results chan<- *hashTask) {
+	h := sha256.New()
 	for t := range tasks {
-		h := sha256.New()
-		h.Write(t.chunk)
-		hash := hex.EncodeToString(h.Sum(nil))
-		results <- hashResult{hash: hash, index: t.index, blockSize: uint64(len(t.chunk))}
+		h.Write((*t.chunk)[:t.blockSize])
+		t.hash = hex.EncodeToString(h.Sum(nil))
+		//fmt.Printf("\rhashing block [%v]\n", t.index)
+		results <- t
+		h.Reset()
 	}
 }