Merge pull request #51 from appleboy/patch2

chore(executor): only wait number of go routine

Author: popcornylu

internal/executor/executor.go

@@ -13,36 +13,41 @@ func ExecuteAll(numCPU int, tasks ...TaskFunc) error {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
-	wg := sync.WaitGroup{}
-	wg.Add(len(tasks))
-
 	if numCPU == 0 {
 		numCPU = runtime.NumCPU()
 	}
-	queue := make(chan TaskFunc, numCPU)
+
+	wg := sync.WaitGroup{}
+	wg.Add(numCPU)
+
+	queue := make(chan TaskFunc, len(tasks))
+	// Add tasks to queue
+	for _, task := range tasks {
+		queue <- task
+	}
+	close(queue)
 
 	// Spawn the executer
 	for i := 0; i < numCPU; i++ {
 		go func() {
-			for task := range queue {
-				if err == nil {
-					taskErr := task(ctx)
-					if taskErr != nil {
-						err = taskErr
+			defer wg.Done()
+			for {
+				select {
+				case task, ok := <-queue:
+					if ctx.Err() != nil || !ok {
+						return
+					}
+					if e := task(ctx); e != nil {
+						err = e
 						cancel()
 					}
+				case <-ctx.Done():
+					return
 				}
-				wg.Done()
 			}
 		}()
 	}
 
-	// Add tasks to queue
-	for _, task := range tasks {
-		queue <- task
-	}
-	close(queue)
-
 	// wait for all task done
 	wg.Wait()
 	return err