Refactor Git archive creation

internal/git/archive.go

@@ -10,10 +10,8 @@ import (
 	"io/ioutil"
 	"net/http"
 	"os"
-	"os/exec"
 	"path"
 	"path/filepath"
-	"syscall"
 	"time"
 
 	"gitlab.com/gitlab-org/gitlab-workhorse/internal/helper"
@@ -52,7 +50,6 @@ func (a *archive) Inject(w http.ResponseWriter, r *http.Request, sendData string
 		return
 	}
 
-	var format string
 	urlPath := r.URL.Path
 	format, ok := parseBasename(filepath.Base(urlPath))
 	if !ok {
@@ -87,64 +84,21 @@ func (a *archive) Inject(w http.ResponseWriter, r *http.Request, sendData string
 	defer tempFile.Close()
 	defer os.Remove(tempFile.Name())
 
-	compressCmd, archiveFormat := parseArchiveFormat(format)
-
-	archiveCmd := gitCommand("", "", "git", "--git-dir="+params.RepoPath, "archive", "--format="+archiveFormat, "--prefix="+params.ArchivePrefix+"/", params.CommitId)
-	archiveStdout, err := archiveCmd.StdoutPipe()
+	archiveReader, err := newArchiveReader(r.Context(), params.RepoPath, format, params.ArchivePrefix, params.CommitId)
 	if err != nil {
-		helper.Fail500(w, r, fmt.Errorf("SendArchive: archive stdout: %v", err))
-		return
-	}
-	defer archiveStdout.Close()
-	if err := archiveCmd.Start(); err != nil {
-		helper.Fail500(w, r, fmt.Errorf("SendArchive: start %v: %v", archiveCmd.Args, err))
+		helper.Fail500(w, r, err)
 		return
 	}
-	defer helper.CleanUpProcessGroup(archiveCmd) // Ensure brute force subprocess clean-up
 
-	var stdout io.ReadCloser
-	if compressCmd == nil {
-		stdout = archiveStdout
-	} else {
-		compressCmd.Stdin = archiveStdout
-		compressCmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
-
-		stdout, err = compressCmd.StdoutPipe()
-		if err != nil {
-			helper.Fail500(w, r, fmt.Errorf("SendArchive: compress stdout: %v", err))
-			return
-		}
-		defer stdout.Close()
-
-		if err := compressCmd.Start(); err != nil {
-			helper.Fail500(w, r, fmt.Errorf("SendArchive: start %v: %v", compressCmd.Args, err))
-			return
-		}
-		defer helper.CleanUpProcessGroup(compressCmd)
-
-		archiveStdout.Close()
-	}
-	// Every Read() from stdout will be synchronously written to tempFile
-	// before it comes out the TeeReader.
-	archiveReader := io.TeeReader(stdout, tempFile)
+	reader := io.TeeReader(archiveReader, tempFile)
 
 	// Start writing the response
 	setArchiveHeaders(w, format, archiveFilename)
 	w.WriteHeader(200) // Don't bother with HTTP 500 from this point on, just return
-	if _, err := io.Copy(w, archiveReader); err != nil {
+	if _, err := io.Copy(w, reader); err != nil {
 		helper.LogError(r, &copyError{fmt.Errorf("SendArchive: copy 'git archive' output: %v", err)})
 		return
 	}
-	if err := archiveCmd.Wait(); err != nil {
-		helper.LogError(r, fmt.Errorf("SendArchive: archiveCmd: %v", err))
-		return
-	}
-	if compressCmd != nil {
-		if err := compressCmd.Wait(); err != nil {
-			helper.LogError(r, fmt.Errorf("SendArchive: compressCmd: %v", err))
-			return
-		}
-	}
 
 	if err := finalizeCachedArchive(tempFile, params.ArchivePath); err != nil {
 		helper.LogError(r, fmt.Errorf("SendArchive: finalize cached archive: %v", err))
@@ -152,10 +106,10 @@ func (a *archive) Inject(w http.ResponseWriter, r *http.Request, sendData string
 	}
 }
 
-func setArchiveHeaders(w http.ResponseWriter, format string, archiveFilename string) {
+func setArchiveHeaders(w http.ResponseWriter, format ArchiveFormat, archiveFilename string) {
 	w.Header().Del("Content-Length")
 	w.Header().Set("Content-Disposition", fmt.Sprintf(`attachment; filename="%s"`, archiveFilename))
-	if format == "zip" {
+	if format == ZipFormat {
 		w.Header().Set("Content-Type", "application/zip")
 	} else {
 		w.Header().Set("Content-Type", "application/octet-stream")
@@ -164,20 +118,6 @@ func setArchiveHeaders(w http.ResponseWriter, format string, archiveFilename str
 	w.Header().Set("Cache-Control", "private")
 }
 
-func parseArchiveFormat(format string) (*exec.Cmd, string) {
-	switch format {
-	case "tar":
-		return nil, "tar"
-	case "tar.gz":
-		return exec.Command("gzip", "-c", "-n"), "tar"
-	case "tar.bz2":
-		return exec.Command("bzip2", "-c"), "tar"
-	case "zip":
-		return nil, "zip"
-	}
-	return nil, "unknown"
-}
-
 func prepareArchiveTempfile(dir string, prefix string) (*os.File, error) {
 	if err := os.MkdirAll(dir, 0700); err != nil {
 		return nil, err
@@ -196,20 +136,20 @@ func finalizeCachedArchive(tempFile *os.File, archivePath string) error {
 	return nil
 }
 
-func parseBasename(basename string) (string, bool) {
-	var format string
+func parseBasename(basename string) (ArchiveFormat, bool) {
+	var format ArchiveFormat
 
 	switch basename {
 	case "archive.zip":
-		format = "zip"
+		format = ZipFormat
 	case "archive.tar":
-		format = "tar"
+		format = TarFormat
 	case "archive", "archive.tar.gz", "archive.tgz", "archive.gz":
-		format = "tar.gz"
+		format = TarGzFormat
 	case "archive.tar.bz2", "archive.tbz", "archive.tbz2", "archive.tb2", "archive.bz2":
-		format = "tar.bz2"
+		format = TarBz2Format
 	default:
-		return "", false
+		return InvalidFormat, false
 	}
 
 	return format, true

internal/git/archive_test.go

@@ -9,16 +9,19 @@ import (
 )
 
 func TestParseBasename(t *testing.T) {
-	for _, testCase := range []struct{ in, out string }{
-		{"", "tar.gz"},
-		{".tar.gz", "tar.gz"},
-		{".tgz", "tar.gz"},
-		{".gz", "tar.gz"},
-		{".tar.bz2", "tar.bz2"},
-		{".tbz", "tar.bz2"},
-		{".tbz2", "tar.bz2"},
-		{".tb2", "tar.bz2"},
-		{".bz2", "tar.bz2"},
+	for _, testCase := range []struct {
+		in  string
+		out ArchiveFormat
+	}{
+		{"", TarGzFormat},
+		{".tar.gz", TarGzFormat},
+		{".tgz", TarGzFormat},
+		{".gz", TarGzFormat},
+		{".tar.bz2", TarBz2Format},
+		{".tbz", TarBz2Format},
+		{".tbz2", TarBz2Format},
+		{".tb2", TarBz2Format},
+		{".bz2", TarBz2Format},
 	} {
 		basename := "archive" + testCase.in
 		out, ok := parseBasename(basename)
@@ -47,11 +50,13 @@ func TestFinalizeArchive(t *testing.T) {
 }
 
 func TestSetArchiveHeaders(t *testing.T) {
-	for _, testCase := range []struct{ in, out string }{
-		{"zip", "application/zip"},
-		{"zippy", "application/octet-stream"},
-		{"rezip", "application/octet-stream"},
-		{"_anything_", "application/octet-stream"},
+	for _, testCase := range []struct {
+		in  ArchiveFormat
+		out string
+	}{
+		{ZipFormat, "application/zip"},
+		{TarFormat, "application/octet-stream"},
+		{InvalidFormat, "application/octet-stream"},
 	} {
 		w := httptest.NewRecorder()
 

internal/git/archivereader.go

+package git
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"os/exec"
+	"syscall"
+
+	"gitlab.com/gitlab-org/gitlab-workhorse/internal/helper"
+)
+
+type ArchiveFormat int
+
+const (
+	InvalidFormat ArchiveFormat = iota
+	ZipFormat
+	TarFormat
+	TarGzFormat
+	TarBz2Format
+)
+
+func parseArchiveFormat(format ArchiveFormat) (*exec.Cmd, string) {
+	switch format {
+	case TarFormat:
+		return nil, "tar"
+	case TarGzFormat:
+		return exec.Command("gzip", "-c", "-n"), "tar"
+	case TarBz2Format:
+		return exec.Command("bzip2", "-c"), "tar"
+	case ZipFormat:
+		return nil, "zip"
+	default:
+		return nil, "invalid format"
+	}
+}
+
+type archiveReader struct {
+	waitCmds []*exec.Cmd
+	stdout   io.Reader
+}
+
+func (a *archiveReader) Read(p []byte) (int, error) {
+	n, err := a.stdout.Read(p)
+
+	if err != io.EOF {
+		return n, err
+	}
+
+	err = a.wait()
+	if err == nil {
+		err = io.EOF
+	}
+	return n, err
+}
+
+func (a *archiveReader) wait() error {
+	var waitErrors []error
+
+	// Must call Wait() on _all_ commands
+	for _, cmd := range a.waitCmds {
+		waitErrors = append(waitErrors, cmd.Wait())
+	}
+
+	for _, err := range waitErrors {
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func newArchiveReader(ctx context.Context, repoPath string, format ArchiveFormat, archivePrefix string, commitId string) (a *archiveReader, err error) {
+	a = &archiveReader{}
+
+	compressCmd, formatArg := parseArchiveFormat(format)
+	archiveCmd := gitCommand("", "", "git", "--git-dir="+repoPath, "archive", "--format="+formatArg, "--prefix="+archivePrefix+"/", commitId)
+
+	var archiveStdout io.ReadCloser
+	archiveStdout, err = archiveCmd.StdoutPipe()
+	if err != nil {
+		return nil, fmt.Errorf("SendArchive: archive stdout: %v", err)
+	}
+	defer func() {
+		if err != nil {
+			archiveStdout.Close()
+		}
+	}()
+
+	a.stdout = archiveStdout
+
+	if compressCmd != nil {
+		compressCmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}
+		compressCmd.Stdin = archiveStdout
+
+		var compressStdout io.ReadCloser
+		compressStdout, err = compressCmd.StdoutPipe()
+		if err != nil {
+			return nil, fmt.Errorf("SendArchive: compress stdout: %v", err)
+		}
+		defer func() {
+			if err != nil {
+				compressStdout.Close()
+			}
+		}()
+
+		if err := compressCmd.Start(); err != nil {
+			return nil, fmt.Errorf("SendArchive: start %v: %v", compressCmd.Args, err)
+		}
+
+		go ctxKill(ctx, compressCmd)
+		a.waitCmds = append(a.waitCmds, compressCmd)
+
+		a.stdout = compressStdout
+		archiveStdout.Close()
+	}
+
+	if err := archiveCmd.Start(); err != nil {
+		return nil, fmt.Errorf("SendArchive: start %v: %v", archiveCmd.Args, err)
+	}
+
+	go ctxKill(ctx, archiveCmd)
+	a.waitCmds = append(a.waitCmds, archiveCmd)
+
+	return a, nil
+}
+
+func ctxKill(ctx context.Context, cmd *exec.Cmd) {
+	<-ctx.Done()
+	helper.CleanUpProcessGroup(cmd)
+	cmd.Wait()
+}