From 67362029625d0eade12b5cf249216f5ed95f9d16 Mon Sep 17 00:00:00 2001
From: weileyuan <weileyuan@tencent.com>
Date: Wed, 24 Aug 2022 00:45:32 +0800
Subject: [PATCH] Push image faster by using pigz to compress

The gzip compression used in the docker push has poor performance.
This change switches to trying to use pigz for gzip compression.
If it isn't available, it'll fall back to the default. This code path
can also be disabled by environment variable.

Signed-off-by: weileyuan <weileyuan@tencent.com>
---
 distribution/push.go    | 63 ++++++++++++++++++++++++++++++++++++++++-
 distribution/push_v2.go |  2 +-
 2 files changed, 63 insertions(+), 2 deletions(-)

diff --git a/distribution/push.go b/distribution/push.go
index 808c5ec31681c..0d6cf90505a84 100644
--- a/distribution/push.go
+++ b/distribution/push.go
@@ -2,12 +2,17 @@ package distribution // import "github.com/docker/docker/distribution"
 
 import (
 	"bufio"
+	"bytes"
 	"compress/gzip"
 	"context"
 	"fmt"
 	"io"
+	"os"
+	"os/exec"
+	"strconv"
 
 	"github.com/docker/distribution/reference"
+	"github.com/docker/docker/pkg/ioutils"
 	"github.com/docker/docker/pkg/progress"
 	"github.com/sirupsen/logrus"
 )
@@ -100,7 +105,30 @@ func Push(ctx context.Context, ref reference.Named, config *ImagePushConfig) err
 // is finished. This allows the caller to make sure the goroutine finishes
 // before it releases any resources connected with the reader that was
 // passed in.
-func compress(in io.Reader) (io.ReadCloser, chan struct{}) {
+func compress(ctx context.Context, in io.Reader) (io.ReadCloser, chan struct{}) {
+	var usePigz bool
+	noPigzEnv := os.Getenv("MOBY_DISABLE_PIGZ")
+	if noPigzEnv == "" {
+		usePigz = true
+	} else {
+		noPigz, err := strconv.ParseBool(noPigzEnv)
+		if err != nil {
+			logrus.WithError(err).Warn("invalid value in MOBY_DISABLE_PIGZ env var")
+		}
+		usePigz = !noPigz
+	}
+
+	if usePigz {
+		pigzPath, err := exec.LookPath("pigz")
+		if err != nil {
+			logrus.Debugf("pigz binary not found, falling back to go gzip library")
+		} else {
+			return pigzCompress(ctx, in, pigzPath)
+		}
+	} else {
+		logrus.Debugf("Use of pigz is disabled due to MOBY_DISABLE_PIGZ=%s", noPigzEnv)
+	}
+
 	compressionDone := make(chan struct{})
 
 	pipeReader, pipeWriter := io.Pipe()
@@ -126,3 +154,36 @@ func compress(in io.Reader) (io.ReadCloser, chan struct{}) {
 
 	return pipeReader, compressionDone
 }
+
+func pigzCompress(ctx context.Context, in io.Reader, pigzPath string) (io.ReadCloser, chan struct{}) {
+	cmd := exec.CommandContext(ctx, pigzPath, "-c")
+
+	cmd.Stdin = in
+	pipeR, pipeW := io.Pipe()
+	cmd.Stdout = pipeW
+	var errBuf bytes.Buffer
+	cmd.Stderr = &errBuf
+
+	done := make(chan struct{})
+
+	if err := cmd.Start(); err != nil {
+		pipeW.CloseWithError(fmt.Errorf("could not get compression stream: %v", err))
+		close(done)
+		return pipeR, done
+	}
+
+	go func() {
+		if err := cmd.Wait(); err != nil {
+			pipeW.CloseWithError(fmt.Errorf("%s: %s", err, errBuf.String()))
+		} else {
+			pipeW.Close()
+		}
+		close(done)
+	}()
+
+	return ioutils.NewReadCloserWrapper(pipeR, func() error {
+		err := pipeR.Close()
+		<-done
+		return err
+	}), done
+}
diff --git a/distribution/push_v2.go b/distribution/push_v2.go
index 910123250c7f0..2608e8fa5ed7b 100644
--- a/distribution/push_v2.go
+++ b/distribution/push_v2.go
@@ -469,7 +469,7 @@ func (pd *pushDescriptor) uploadUsingSession(
 
 	switch m := pd.layer.MediaType(); m {
 	case schema2.MediaTypeUncompressedLayer:
-		compressedReader, compressionDone := compress(reader)
+		compressedReader, compressionDone := compress(ctx, reader)
 		defer func(closer io.Closer) {
 			closer.Close()
 			<-compressionDone