fix unexpected exit on download

DiheChen · May 21, 2022 · e624070 · e624070
1 parent 4c2c969
commit e624070
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/main.go b/main.go
@@ -36,4 +36,5 @@ func main() {
 	for _, task := range tasks {
 		c.Download(task)
 	}
+	c.WorkerPool.Wait()
 }
diff --git a/utils/worker.go b/utils/worker.go
@@ -1,6 +1,9 @@
 package utils
 
-import "fmt"
+import (
+	"fmt"
+	"sync"
+)
 
 type TaskQueue chan *MultiThreadDownloader
 
@@ -9,6 +12,8 @@ type Worker struct {
 }
 
 type WorkerPool struct {
+	sync.WaitGroup
+
 	WorkerCount int
 	TaskQueue   TaskQueue
 	WorkerQueue chan TaskQueue
@@ -26,7 +31,7 @@ func NewWorkerPool(WorkerCount int) *WorkerPool {
 	}
 }
 
-func (w *Worker) Run(wq chan TaskQueue) {
+func (w *Worker) Run(wq chan TaskQueue, owner *WorkerPool) {
 	go func() {
 		for {
 			wq <- w.TaskChan
@@ -38,6 +43,7 @@ func (w *Worker) Run(wq chan TaskQueue) {
 					return
 				}
 				fmt.Println("下载完成", t.FullPath)
+				owner.Done()
 			}
 		}
 	}()
@@ -46,12 +52,13 @@ func (w *Worker) Run(wq chan TaskQueue) {
 func (wp *WorkerPool) Start() {
 	for i := 0; i < wp.WorkerCount; i++ {
 		w := NewWorker()
-		w.Run(wp.WorkerQueue)
+		w.Run(wp.WorkerQueue, wp)
 	}
 	go func() {
 		for {
 			select {
 			case t := <-wp.TaskQueue:
+				wp.Add(1)
 				w := <-wp.WorkerQueue
 				w <- t
 			}