feat: reuse connections and limit the number of connections for prehe…

…ating

Signed-off-by: Gaius <[email protected]>

manager/job/preheat.go

@@ -65,6 +65,14 @@ const (
 	PreheatFileType PreheatType = "file"
 )
 
+// defaultHTTPTransport is the default http transport.
+var defaultHTTPTransport = &http.Transport{
+	MaxIdleConns:        400,
+	MaxIdleConnsPerHost: 20,
+	MaxConnsPerHost:     50,
+	IdleConnTimeout:     120 * time.Second,
+}
+
 // accessURLPattern is the pattern of access url.
 var accessURLPattern, _ = regexp.Compile("^(.*)://(.*)/v2/(.*)/manifests/(.*)")
 
@@ -77,20 +85,34 @@ type Preheat interface {
 // preheat is an implementation of Preheat.
 type preheat struct {
 	job                *internaljob.Job
-	registryTimeout    time.Duration
-	rootCAs            *x509.CertPool
 	certificateChain   [][]byte
 	insecureSkipVerify bool
+	httpClient         *http.Client
 }
 
 // newPreheat creates a new Preheat.
 func newPreheat(job *internaljob.Job, registryTimeout time.Duration, rootCAs *x509.CertPool, insecureSkipVerify bool) (Preheat, error) {
-	var certificateChain [][]byte
+	p := &preheat{
+		job:                job,
+		insecureSkipVerify: insecureSkipVerify,
+		httpClient: &http.Client{
+			Timeout: registryTimeout,
+			Transport: &http.Transport{
+				DialContext:         nethttp.NewSafeDialer().DialContext,
+				TLSClientConfig:     &tls.Config{RootCAs: rootCAs, InsecureSkipVerify: insecureSkipVerify},
+				MaxIdleConns:        defaultHTTPTransport.MaxIdleConns,
+				MaxIdleConnsPerHost: defaultHTTPTransport.MaxIdleConnsPerHost,
+				MaxConnsPerHost:     defaultHTTPTransport.MaxConnsPerHost,
+				IdleConnTimeout:     defaultHTTPTransport.IdleConnTimeout,
+			},
+		},
+	}
+
 	if rootCAs != nil {
-		certificateChain = rootCAs.Subjects()
+		p.certificateChain = rootCAs.Subjects()
 	}
 
-	return &preheat{job, registryTimeout, rootCAs, certificateChain, insecureSkipVerify}, nil
+	return p, nil
 }
 
 // CreatePreheat creates a preheat job.
@@ -191,26 +213,20 @@ func (p *preheat) getImageLayers(ctx context.Context, args types.PreheatArgs) ([
 		return nil, err
 	}
 
-	opts := []imageAuthClientOption{
-		withHTTPClient(&http.Client{
-			Timeout: p.registryTimeout,
-			Transport: &http.Transport{
-				DialContext:     nethttp.NewSafeDialer().DialContext,
-				TLSClientConfig: &tls.Config{RootCAs: p.rootCAs, InsecureSkipVerify: p.insecureSkipVerify},
-			},
-		}),
+	options := []imageAuthClientOption{
+		withHTTPClient(p.httpClient),
 		withBasicAuth(args.Username, args.Password),
 	}
 	// Background:
-	//   Harbor uses the V1 preheat request and will carry the auth info in the headers.
+	// Harbor uses the V1 preheat request and will carry the auth info in the headers.
 	header := nethttp.MapToHeader(args.Headers)
 	if token := header.Get("Authorization"); len(token) > 0 {
-		opts = append(opts, withIssuedToken(token))
+		options = append(options, withIssuedToken(token))
 		header.Set("Authorization", token)
 	}
 
 	// Init docker auth client.
-	client, err := newImageAuthClient(image, opts...)
+	client, err := newImageAuthClient(image, options...)
 	if err != nil {
 		return nil, err
 	}
@@ -395,8 +411,11 @@ type imageAuthClient struct {
 
 // newImageAuthClient creates a new imageAuthClient.
 func newImageAuthClient(image *preheatImage, opts ...imageAuthClientOption) (*imageAuthClient, error) {
+	httpClient := http.DefaultClient
+	httpClient.Transport = defaultHTTPTransport
+
 	d := &imageAuthClient{
-		httpClient:              http.DefaultClient,
+		httpClient:              httpClient,
 		interceptorTokenHandler: newInterceptorTokenHandler(),
 	}
 

manager/job/preheat_test.go

@@ -2,6 +2,7 @@ package job
 
 import (
 	"context"
+	"net/http"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -42,7 +43,7 @@ func TestPreheat_getImageLayers(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
-			p := &preheat{}
+			p := &preheat{httpClient: http.DefaultClient}
 			layers, err := p.getImageLayers(context.Background(), tc.args)
 			if err != nil {
 				t.Fatal(err)