From 64094899b03ce093edfbc93131a66a9aef22a046 Mon Sep 17 00:00:00 2001
From: lsytj0413 <511121939@qq.com>
Date: Sat, 30 May 2020 18:07:23 +0800
Subject: [PATCH] test(e2e_node): Parallelize prepulling all images in
 `e2e_node` tests

---
 test/e2e_node/BUILD         |  1 +
 test/e2e_node/image_list.go | 79 +++++++++++++++++++++++++++----------
 2 files changed, 60 insertions(+), 20 deletions(-)

diff --git a/test/e2e_node/BUILD b/test/e2e_node/BUILD
index dab44a88e2b..da2a052aa84 100644
--- a/test/e2e_node/BUILD
+++ b/test/e2e_node/BUILD
@@ -41,6 +41,7 @@ go_library(
         "//staging/src/k8s.io/api/core/v1:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/api/equality:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/apis/meta/v1:go_default_library",
+        "//staging/src/k8s.io/apimachinery/pkg/util/errors:go_default_library",
         "//staging/src/k8s.io/apimachinery/pkg/util/sets:go_default_library",
         "//staging/src/k8s.io/apiserver/pkg/util/feature:go_default_library",
         "//staging/src/k8s.io/client-go/kubernetes:go_default_library",
diff --git a/test/e2e_node/image_list.go b/test/e2e_node/image_list.go
index b6b1f0f1986..ead5eaeb877 100644
--- a/test/e2e_node/image_list.go
+++ b/test/e2e_node/image_list.go
@@ -17,14 +17,17 @@ limitations under the License.
 package e2enode
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"os/exec"
 	"os/user"
+	"sync"
 	"time"
 
 	"k8s.io/klog"
 
+	utilerrors "k8s.io/apimachinery/pkg/util/errors"
 	"k8s.io/apimachinery/pkg/util/sets"
 	internalapi "k8s.io/cri-api/pkg/apis"
 	runtimeapi "k8s.io/cri-api/pkg/apis/runtime/v1alpha2"
@@ -41,6 +44,8 @@ const (
 	maxImagePullRetries = 5
 	// Sleep duration between image pull retry attempts.
 	imagePullRetryDelay = time.Second
+	// Number of parallel count to pull images.
+	maxParallelImagePullCount = 5
 )
 
 // NodeImageWhiteList is a list of images used in node e2e test. These images will be prepulled
@@ -151,27 +156,61 @@ func PrePullAllImages() error {
 	}
 	images := framework.ImageWhiteList.List()
 	klog.V(4).Infof("Pre-pulling images with %s %+v", puller.Name(), images)
-	for _, image := range images {
-		var (
-			err    error
-			output []byte
-		)
-		for i := 0; i < maxImagePullRetries; i++ {
-			if i > 0 {
-				time.Sleep(imagePullRetryDelay)
-			}
-			if output, err = puller.Pull(image); err == nil {
-				break
-			}
-			klog.Warningf("Failed to pull %s as user %q, retrying in %s (%d of %d): %v",
-				image, usr.Username, imagePullRetryDelay.String(), i+1, maxImagePullRetries, err)
-		}
-		if err != nil {
-			klog.Warningf("Could not pre-pull image %s %v output: %s", image, err, output)
-			return err
-		}
+
+	imageCh := make(chan int, len(images))
+	for i := range images {
+		imageCh <- i
 	}
-	return nil
+	close(imageCh)
+
+	pullErrs := make([]error, len(images))
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	parallelImagePullCount := maxParallelImagePullCount
+	if len(images) < parallelImagePullCount {
+		parallelImagePullCount = len(images)
+	}
+
+	var wg sync.WaitGroup
+	wg.Add(parallelImagePullCount)
+	for i := 0; i < parallelImagePullCount; i++ {
+		go func() {
+			defer wg.Done()
+
+			for i := range imageCh {
+				var (
+					pullErr error
+					output  []byte
+				)
+				for retryCount := 0; retryCount < maxImagePullRetries; retryCount++ {
+					select {
+					case <-ctx.Done():
+						return
+					default:
+					}
+
+					if retryCount > 0 {
+						time.Sleep(imagePullRetryDelay)
+					}
+					if output, pullErr = puller.Pull(images[i]); pullErr == nil {
+						break
+					}
+					klog.Warningf("Failed to pull %s as user %q, retrying in %s (%d of %d): %v",
+						images[i], usr.Username, imagePullRetryDelay.String(), retryCount+1, maxImagePullRetries, pullErr)
+				}
+				if pullErr != nil {
+					klog.Warningf("Could not pre-pull image %s %v output: %s", images[i], pullErr, output)
+					pullErrs[i] = pullErr
+					cancel()
+					return
+				}
+			}
+		}()
+	}
+
+	wg.Wait()
+	return utilerrors.NewAggregate(pullErrs)
 }
 
 // getGPUDevicePluginImage returns the image of GPU device plugin.