Improve layer decompression speed by using pigz

Signed-off-by: Maksym Pavlenko <makpav@amazon.com>

archive/compression/compression.go

@@ -20,9 +20,15 @@ import (
 	"bufio"
 	"bytes"
 	"compress/gzip"
+	"context"
 	"fmt"
 	"io"
+	"os"
+	"os/exec"
+	"strconv"
 	"sync"
+
+	"github.com/containerd/containerd/log"
 )
 
 type (
@@ -37,6 +43,16 @@ const (
 	Gzip
 )
 
+const disablePigzEnv = "CONTAINERD_DISABLE_PIGZ"
+
+var unpigzPath string
+
+func init() {
+	if unpigzPath = detectPigz(); unpigzPath != "" {
+		log.L.Debug("using pigz for decompression")
+	}
+}
+
 var (
 	bufioReader32KPool = &sync.Pool{
 		New: func() interface{} { return bufio.NewReaderSize(nil, 32*1024) },
@@ -120,11 +136,18 @@ func DecompressStream(archive io.Reader) (DecompressReadCloser, error) {
 		readBufWrapper := &readCloserWrapper{buf, compression, closer}
 		return readBufWrapper, nil
 	case Gzip:
-		gzReader, err := gzip.NewReader(buf)
+		ctx, cancel := context.WithCancel(context.Background())
+		gzReader, err := gzipDecompress(ctx, buf)
 		if err != nil {
+			cancel()
 			return nil, err
 		}
-		readBufWrapper := &readCloserWrapper{gzReader, compression, closer}
+
+		readBufWrapper := &readCloserWrapper{gzReader, compression, func() error {
+			cancel()
+			return closer()
+		}}
+
 		return readBufWrapper, nil
 	default:
 		return nil, fmt.Errorf("unsupported compression format %s", (&compression).Extension())
@@ -151,3 +174,61 @@ func (compression *Compression) Extension() string {
 	}
 	return ""
 }
+
+func gzipDecompress(ctx context.Context, buf io.Reader) (io.ReadCloser, error) {
+	if unpigzPath == "" {
+		return gzip.NewReader(buf)
+	}
+
+	return cmdStream(exec.CommandContext(ctx, unpigzPath, "-d", "-c"), buf)
+}
+
+func cmdStream(cmd *exec.Cmd, in io.Reader) (io.ReadCloser, error) {
+	reader, writer := io.Pipe()
+
+	cmd.Stdin = in
+	cmd.Stdout = writer
+
+	var errBuf bytes.Buffer
+	cmd.Stderr = &errBuf
+
+	if err := cmd.Start(); err != nil {
+		return nil, err
+	}
+
+	go func() {
+		if err := cmd.Wait(); err != nil {
+			writer.CloseWithError(fmt.Errorf("%s: %s", err, errBuf.String()))
+		} else {
+			writer.Close()
+		}
+	}()
+
+	return reader, nil
+}
+
+func detectPigz() string {
+	path, err := exec.LookPath("unpigz")
+	if err != nil {
+		log.L.WithError(err).Debug("unpigz not found, falling back to go gzip")
+		return ""
+	}
+
+	// Check if pigz disabled via CONTAINERD_DISABLE_PIGZ env variable
+	value := os.Getenv(disablePigzEnv)
+	if value == "" {
+		return path
+	}
+
+	disable, err := strconv.ParseBool(value)
+	if err != nil {
+		log.L.WithError(err).Warnf("could not parse %s: %s", disablePigzEnv, value)
+		return path
+	}
+
+	if disable {
+		return ""
+	}
+
+	return path
+}