Merge branch 'pks-gitaly-wrapper-refactorings' into 'master'

cmd/gitaly-wrapper: Improve test coverage

See merge request gitlab-org/gitaly!4245

4 files changed, 437 insertions, 39 deletions

diff --git a/cmd/gitaly-wrapper/main.go b/cmd/gitaly-wrapper/main.go
index 3eeb55ad5..247b71f2d 100644
--- a/cmd/gitaly-wrapper/main.go
+++ b/cmd/gitaly-wrapper/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"errors"
 	"fmt"
 	"os"
 	"os/exec"
@@ -34,74 +35,75 @@ func main() {
 		logrus.Fatalf("usage: %s forking_binary [args]", os.Args[0])
 	}
 
-	gitalyBin, gitalyArgs := os.Args[1], os.Args[2:]
+	binary, arguments := os.Args[1], os.Args[2:]
 
 	logger := log.Default().WithField("wrapper", os.Getpid())
 	logger.Info("Wrapper started")
 
-	if pidFile() == "" {
+	pidFilePath := os.Getenv(bootstrap.EnvPidFile)
+	if pidFilePath == "" {
 		logger.Fatalf("missing pid file ENV variable %q", bootstrap.EnvPidFile)
 	}
+	logger.WithField("pid_file", pidFilePath).Info("finding process")
 
-	logger.WithField("pid_file", pidFile()).Info("finding gitaly")
-	gitaly, err := findGitaly()
+	process, err := findProcess(pidFilePath)
 	if err != nil && !isRecoverable(err) {
-		logger.WithError(err).Fatal("find gitaly")
+		logger.WithError(err).Fatal("find process")
 	} else if err != nil {
-		logger.WithError(err).Error("find gitaly")
+		logger.WithError(err).Error("find process")
 	}
 
-	if gitaly != nil && isGitaly(gitaly, gitalyBin) {
+	if process != nil && isExpectedProcess(process, binary) {
 		logger.Info("adopting a process")
 	} else {
 		logger.Info("spawning a process")
 
-		proc, err := spawnGitaly(gitalyBin, gitalyArgs)
+		proc, err := spawnProcess(binary, arguments)
 		if err != nil {
 			logger.WithError(err).Fatal("spawn gitaly")
 		}
 
-		gitaly = proc
+		process = proc
 	}
 
-	logger = logger.WithField("gitaly", gitaly.Pid)
-	logger.Info("monitoring gitaly")
+	logger = logger.WithField("process", process.Pid)
+	logger.Info("monitoring process")
 
-	forwardSignals(gitaly, logger)
+	forwardSignals(process, logger)
 
 	// wait
-	for isAlive(gitaly) {
+	for isProcessAlive(process) {
 		time.Sleep(1 * time.Second)
 	}
 
-	logger.Error("wrapper for gitaly shutting down")
+	logger.Error("wrapper for process shutting down")
 }
 
 func isRecoverable(err error) bool {
-	_, isNumError := err.(*strconv.NumError)
-	return os.IsNotExist(err) || isNumError
+	var numError *strconv.NumError
+	return os.IsNotExist(err) || errors.As(err, &numError)
 }
 
-func findGitaly() (*os.Process, error) {
-	pid, err := getPid()
+func findProcess(pidFilePath string) (*os.Process, error) {
+	pid, err := readPIDFile(pidFilePath)
 	if err != nil {
 		return nil, err
 	}
 
 	// os.FindProcess on unix do not return an error if the process does not exist
-	gitaly, err := os.FindProcess(pid)
+	process, err := os.FindProcess(pid)
 	if err != nil {
 		return nil, err
 	}
 
-	if isAlive(gitaly) {
-		return gitaly, nil
+	if isProcessAlive(process) {
+		return process, nil
 	}
 
 	return nil, nil
 }
 
-func spawnGitaly(bin string, args []string) (*os.Process, error) {
+func spawnProcess(bin string, args []string) (*os.Process, error) {
 	cmd := exec.Command(bin, args...)
 	cmd.Env = append(os.Environ(), fmt.Sprintf("%s=true", bootstrap.EnvUpgradesEnabled))
 
@@ -145,8 +147,8 @@ func forwardSignals(gitaly *os.Process, log *logrus.Entry) {
 	signal.Notify(sigs)
 }
 
-func getPid() (int, error) {
-	data, err := os.ReadFile(pidFile())
+func readPIDFile(pidFilePath string) (int, error) {
+	data, err := os.ReadFile(pidFilePath)
 	if err != nil {
 		return 0, err
 	}
@@ -154,7 +156,7 @@ func getPid() (int, error) {
 	return strconv.Atoi(string(data))
 }
 
-func isAlive(p *os.Process) bool {
+func isProcessAlive(p *os.Process) bool {
 	// After p exits, and after it gets reaped, this p.Signal will fail. It is crucial that p gets reaped.
 	// If p was spawned by the current process, it will get reaped from a goroutine that does cmd.Wait().
 	// If p was spawned by someone else we rely on them to reap it, or on p to become an orphan.
@@ -162,23 +164,19 @@ func isAlive(p *os.Process) bool {
 	return p.Signal(syscall.Signal(0)) == nil
 }
 
-func isGitaly(p *os.Process, gitalyBin string) bool {
+func isExpectedProcess(p *os.Process, binary string) bool {
 	command, err := ps.Comm(p.Pid)
 	if err != nil {
 		return false
 	}
 
-	if filepath.Base(command) == filepath.Base(gitalyBin) {
+	if filepath.Base(command) == filepath.Base(binary) {
 		return true
 	}
 
 	return false
 }
 
-func pidFile() string {
-	return os.Getenv(bootstrap.EnvPidFile)
-}
-
 func jsonLogging() bool {
 	enabled, _ := env.GetBool(envJSONLogging, false)
 	return enabled
diff --git a/cmd/gitaly-wrapper/main_test.go b/cmd/gitaly-wrapper/main_test.go
index abca0aa09..42c584c51 100644
--- a/cmd/gitaly-wrapper/main_test.go
+++ b/cmd/gitaly-wrapper/main_test.go
@@ -1,7 +1,10 @@
 package main
 
 import (
+	"bufio"
 	"errors"
+	"fmt"
+	"io"
 	"os"
 	"os/exec"
 	"path/filepath"
@@ -11,18 +14,18 @@ import (
 	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/bootstrap"
 	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper"
+	"gitlab.com/gitlab-org/gitaly/v14/internal/testhelper/testcfg"
 )
 
 // TestStolenPid tests for regressions in https://gitlab.com/gitlab-org/gitaly/issues/1661
 func TestStolenPid(t *testing.T) {
+	t.Parallel()
+
 	tempDir := testhelper.TempDir(t)
 
 	pidFile, err := os.Create(filepath.Join(tempDir, "pidfile"))
 	require.NoError(t, err)
 
-	cleanup := testhelper.ModifyEnvironment(t, bootstrap.EnvPidFile, pidFile.Name())
-	defer cleanup()
-
 	ctx, cancel := testhelper.Context()
 	defer cancel()
 
@@ -37,21 +40,96 @@ func TestStolenPid(t *testing.T) {
 	require.NoError(t, err)
 	require.NoError(t, pidFile.Close())
 
-	tail, err := findGitaly()
+	tail, err := findProcess(pidFile.Name())
 	require.NoError(t, err)
 	require.NotNil(t, tail)
 	require.Equal(t, cmd.Process.Pid, tail.Pid)
 
 	t.Run("stolen", func(t *testing.T) {
-		require.False(t, isGitaly(tail, "/path/to/gitaly"))
+		require.False(t, isExpectedProcess(tail, "/path/to/gitaly"))
 	})
 
 	t.Run("not stolen", func(t *testing.T) {
-		require.True(t, isGitaly(tail, "/path/to/tail"))
+		require.True(t, isExpectedProcess(tail, "/path/to/tail"))
+	})
+}
+
+func TestFindProcess(t *testing.T) {
+	t.Parallel()
+
+	t.Run("nonexistent PID file", func(t *testing.T) {
+		t.Parallel()
+
+		_, err := findProcess("does-not-exist")
+		require.True(t, os.IsNotExist(err))
+	})
+
+	t.Run("PID file with garbage", func(t *testing.T) {
+		t.Parallel()
+
+		path := filepath.Join(testhelper.TempDir(t), "pid")
+		require.NoError(t, os.WriteFile(path, []byte("garbage"), 0o644))
+
+		_, err := findProcess(path)
+		_, expectedErr := strconv.Atoi("garbage")
+		require.Equal(t, expectedErr, err)
+	})
+
+	t.Run("nonexistent process", func(t *testing.T) {
+		t.Parallel()
+
+		// The below PID can exist, but chances are sufficiently low to hopefully not matter
+		// in practice.
+		path := filepath.Join(testhelper.TempDir(t), "pid")
+		require.NoError(t, os.WriteFile(path, []byte("7777777"), 0o644))
+
+		// The process isn't alive, so we expect neither an error nor a process to be
+		// returned.
+		process, err := findProcess(path)
+		require.Nil(t, process)
+		require.Nil(t, err)
+	})
+
+	t.Run("running process", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := testhelper.Context()
+		defer cancel()
+
+		// The process will exit immediately, but that's not much of a problem given that
+		// `findProcess()` also works with zombies.
+		executable := testhelper.WriteExecutable(t, filepath.Join(testhelper.TempDir(t), "noop"), []byte(
+			`#!/usr/bin/env bash
+			echo ready
+		`))
+
+		cmd := exec.CommandContext(ctx, executable)
+		stdout, err := cmd.StdoutPipe()
+		require.NoError(t, err)
+
+		require.NoError(t, cmd.Start())
+		defer func() {
+			require.NoError(t, cmd.Wait())
+		}()
+
+		// Wait for the process to be ready such that we know it's started up successfully
+		// and is executing the bash shell.
+		_, err = stdout.Read(make([]byte, 10))
+		require.NoError(t, err)
+
+		path := filepath.Join(testhelper.TempDir(t), "pid")
+		require.NoError(t, os.WriteFile(path, []byte(strconv.FormatInt(int64(cmd.Process.Pid), 10)), 0o644))
+
+		process, err := findProcess(path)
+		require.NotNil(t, process)
+		require.Equal(t, cmd.Process.Pid, process.Pid)
+		require.Equal(t, nil, err)
 	})
 }
 
 func TestIsRecoverable(t *testing.T) {
+	t.Parallel()
+
 	_, numericError := strconv.Atoi("")
 
 	tests := []struct {
@@ -83,3 +161,317 @@ func TestIsRecoverable(t *testing.T) {
 		})
 	}
 }
+
+func TestReadPIDFile(t *testing.T) {
+	t.Parallel()
+
+	t.Run("nonexistent", func(t *testing.T) {
+		t.Parallel()
+
+		_, err := readPIDFile("does-not-exist")
+		require.True(t, os.IsNotExist(err))
+	})
+
+	t.Run("empty", func(t *testing.T) {
+		t.Parallel()
+
+		path := filepath.Join(testhelper.TempDir(t), "pid")
+		require.NoError(t, os.WriteFile(path, nil, 0o644))
+		_, err := readPIDFile(path)
+		_, expectedErr := strconv.Atoi("")
+		require.Equal(t, expectedErr, err)
+	})
+
+	t.Run("invalid contents", func(t *testing.T) {
+		t.Parallel()
+
+		path := filepath.Join(testhelper.TempDir(t), "pid")
+		require.NoError(t, os.WriteFile(path, []byte("invalid"), 0o644))
+		_, err := readPIDFile(path)
+		_, expectedErr := strconv.Atoi("invalid")
+		require.Equal(t, expectedErr, err)
+	})
+
+	t.Run("valid", func(t *testing.T) {
+		t.Parallel()
+
+		path := filepath.Join(testhelper.TempDir(t), "pid")
+		require.NoError(t, os.WriteFile(path, []byte("12345"), 0o644))
+		pid, err := readPIDFile(path)
+		require.NoError(t, err)
+		require.Equal(t, 12345, pid)
+	})
+}
+
+func TestIsExpectedProcess(t *testing.T) {
+	t.Parallel()
+
+	ctx, cancel := testhelper.Context()
+	defer cancel()
+
+	executable := testhelper.WriteExecutable(t, filepath.Join(testhelper.TempDir(t), "noop"), []byte(
+		`#!/usr/bin/env bash
+		echo ready
+		read wait_until_killed
+	`))
+
+	cmd := exec.CommandContext(ctx, executable)
+	stdout, err := cmd.StdoutPipe()
+	require.NoError(t, err)
+	stdin, err := cmd.StdinPipe()
+	require.NoError(t, err)
+
+	require.NoError(t, cmd.Start())
+	defer func() {
+		testhelper.MustClose(t, stdin)
+		require.Error(t, cmd.Wait())
+	}()
+
+	// Wait for the process to be ready such that we know it's started up successfully
+	// and is executing the bash shell.
+	_, err = stdout.Read(make([]byte, 10))
+	require.NoError(t, err)
+
+	require.False(t, isExpectedProcess(cmd.Process, "does not match"))
+	require.True(t, isExpectedProcess(cmd.Process, "bash"))
+}
+
+func TestIsProcessAlive(t *testing.T) {
+	t.Parallel()
+
+	t.Run("nonexistent process", func(t *testing.T) {
+		t.Parallel()
+
+		// And now let's check with a nonexistent process. FindProcess never returns an
+		// error on Unix systems even if the process doesn't exist, so this is fine.
+		process, err := os.FindProcess(77777777)
+		require.NoError(t, err)
+		require.False(t, isProcessAlive(process))
+	})
+
+	t.Run("existing process", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := testhelper.Context()
+		defer cancel()
+
+		executable := testhelper.WriteExecutable(t, filepath.Join(testhelper.TempDir(t), "noop"), []byte(
+			`#!/usr/bin/env bash
+			echo ready
+			read wait_until_killed
+		`))
+
+		cmd := exec.CommandContext(ctx, executable)
+		stdout, err := cmd.StdoutPipe()
+		require.NoError(t, err)
+		_, err = cmd.StdinPipe()
+		require.NoError(t, err)
+
+		require.NoError(t, cmd.Start())
+
+		// Wait for the process to be ready such that we know it's started up successfully
+		// and is executing the bash shell.
+		_, err = stdout.Read(make([]byte, 10))
+		require.NoError(t, err)
+
+		t.Run("running", func(t *testing.T) {
+			require.True(t, isProcessAlive(cmd.Process))
+		})
+
+		t.Run("zombie", func(t *testing.T) {
+			// The process will be considered alive as long as it hasn't been reaped yet.
+			require.NoError(t, cmd.Process.Kill())
+			require.True(t, isProcessAlive(cmd.Process))
+		})
+
+		t.Run("reaped", func(t *testing.T) {
+			require.Error(t, cmd.Wait())
+			require.False(t, isProcessAlive(cmd.Process))
+		})
+	})
+}
+
+func TestRun(t *testing.T) {
+	t.Parallel()
+
+	binary := testcfg.BuildGitalyWrapper(t, testcfg.Build(t))
+
+	t.Run("missing arguments", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := testhelper.Context()
+		defer cancel()
+
+		output, err := exec.CommandContext(ctx, binary).CombinedOutput()
+		require.Error(t, err)
+		require.Contains(t, string(output), "usage: ")
+	})
+
+	t.Run("missing PID file envvar", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := testhelper.Context()
+		defer cancel()
+
+		output, err := exec.CommandContext(ctx, binary, "binary").CombinedOutput()
+		require.Error(t, err)
+		require.Contains(t, string(output), "missing pid file ENV variable")
+	})
+
+	t.Run("invalid executable", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := testhelper.Context()
+		defer cancel()
+
+		pidPath := filepath.Join(testhelper.TempDir(t), "pid")
+
+		cmd := exec.CommandContext(ctx, binary, "does-not-exist")
+		cmd.Env = append(os.Environ(), fmt.Sprintf("%s=%s", bootstrap.EnvPidFile, pidPath))
+
+		output, err := cmd.CombinedOutput()
+		require.Error(t, err)
+		require.Contains(t, string(output), "executable file not found in $PATH")
+	})
+
+	t.Run("adopting executable", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := testhelper.Context()
+		defer cancel()
+
+		script := testhelper.WriteExecutable(t, filepath.Join(testhelper.TempDir(t), "script"), []byte(
+			`#!/usr/bin/env bash
+			echo ready
+			read wait_until_closed
+		`))
+
+		scriptCmd := exec.CommandContext(ctx, script)
+		scriptStdout, err := scriptCmd.StdoutPipe()
+		require.NoError(t, err)
+		scriptStdin, err := scriptCmd.StdinPipe()
+		require.NoError(t, err)
+
+		require.NoError(t, scriptCmd.Start())
+
+		// Read the first byte such that we know the process has been spawned.
+		_, err = scriptStdout.Read(make([]byte, 10))
+		require.NoError(t, err)
+
+		// Write the PID of the running process into the PID file. As a result, it should
+		// get adopted by gitaly-wrapper, which means it wouldn't try to execute it anew.
+		pidPath := filepath.Join(testhelper.TempDir(t), "pid")
+		require.NoError(t, os.WriteFile(pidPath, []byte(strconv.FormatInt(int64(scriptCmd.Process.Pid), 10)), 0o644))
+
+		// Run gitaly-script with a binary path whose basename matches, but which ultimately
+		// doesn't exist. This proves that it doesn't try to execute the script again.
+		wrapperCmd := exec.CommandContext(ctx, binary, "/does/not/exist/bash")
+		wrapperCmd.Env = append(os.Environ(), fmt.Sprintf("%s=%s", bootstrap.EnvPidFile, pidPath))
+		wrapperStdout, err := wrapperCmd.StdoutPipe()
+		require.NoError(t, err)
+		require.NoError(t, wrapperCmd.Start())
+
+		// We're now waiting for gitaly-wrapper to adopt the running script.
+		reader := bufio.NewReader(wrapperStdout)
+		for _, expectedLine := range []string{
+			"Wrapper started",
+			"finding process",
+			"adopting a process",
+		} {
+			line, err := reader.ReadBytes('\n')
+			require.NoError(t, err, "reading expected line %q",
+				expectedLine)
+			require.Contains(t, string(line), expectedLine)
+		}
+
+		// The script has been adopted, so we can now close its stdin and thus cause it to
+		// exit.
+		testhelper.MustClose(t, scriptStdin)
+		require.Error(t, scriptCmd.Wait())
+
+		// As a result, the wrapper should also exit successfully.
+		require.NoError(t, wrapperCmd.Wait())
+	})
+
+	t.Run("spawning executable", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := testhelper.Context()
+		defer cancel()
+
+		script := testhelper.WriteExecutable(t, filepath.Join(testhelper.TempDir(t), "script"), []byte(
+			`#!/usr/bin/env bash
+			echo "I have been executed"
+		`))
+
+		pidPath := filepath.Join(testhelper.TempDir(t), "pid")
+
+		cmd := exec.CommandContext(ctx, binary, script)
+		cmd.Env = append(os.Environ(), fmt.Sprintf("%s=%s", bootstrap.EnvPidFile, pidPath))
+		output, err := cmd.CombinedOutput()
+		require.NoError(t, err)
+
+		require.Contains(t, string(output), "spawning a process")
+		require.Contains(t, string(output), "I have been executed")
+
+		require.NoFileExists(t, pidPath)
+	})
+
+	t.Run("spawning executable with missing process", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := testhelper.Context()
+		defer cancel()
+
+		script := testhelper.WriteExecutable(t, filepath.Join(testhelper.TempDir(t), "script"), []byte(
+			`#!/usr/bin/env bash
+			echo "I have been executed"
+		`))
+
+		pidPath := filepath.Join(testhelper.TempDir(t), "pid")
+		require.NoError(t, os.WriteFile(pidPath, []byte("12345"), 0o644))
+
+		cmd := exec.CommandContext(ctx, binary, script)
+		cmd.Env = append(os.Environ(), fmt.Sprintf("%s=%s", bootstrap.EnvPidFile, pidPath))
+
+		output, err := cmd.CombinedOutput()
+		require.NoError(t, err)
+		require.Contains(t, string(output), "spawning a process")
+		require.Contains(t, string(output), "I have been executed")
+	})
+
+	t.Run("spawning executable with zombie process", func(t *testing.T) {
+		t.Parallel()
+
+		ctx, cancel := testhelper.Context()
+		defer cancel()
+
+		script := testhelper.WriteExecutable(t, filepath.Join(testhelper.TempDir(t), "script"), []byte(
+			`#!/usr/bin/env bash
+			echo "I have been executed"
+		`))
+
+		scriptCmd := exec.CommandContext(ctx, script)
+		scriptStdout, err := scriptCmd.StdoutPipe()
+		require.NoError(t, err)
+		require.NoError(t, scriptCmd.Start())
+
+		// Read until we get an EOF, which means that the script has terminated. It's now in
+		// a zombie state because we don't call `Wait()`.
+		_, err = io.ReadAll(scriptStdout)
+		require.NoError(t, err)
+
+		pidPath := filepath.Join(testhelper.TempDir(t), "pid")
+		require.NoError(t, os.WriteFile(pidPath, []byte(strconv.FormatInt(int64(scriptCmd.Process.Pid), 10)), 0o644))
+
+		cmd := exec.CommandContext(ctx, binary, script)
+		cmd.Env = append(os.Environ(), fmt.Sprintf("%s=%s", bootstrap.EnvPidFile, pidPath))
+
+		output, err := cmd.CombinedOutput()
+		require.NoError(t, err)
+		require.Contains(t, string(output), "spawning a process")
+		require.Contains(t, string(output), "I have been executed")
+
+		require.NoError(t, scriptCmd.Wait())
+	})
+}
diff --git a/internal/testhelper/testcfg/build.go b/internal/testhelper/testcfg/build.go
index c18a9fae5..ce958358c 100644
--- a/internal/testhelper/testcfg/build.go
+++ b/internal/testhelper/testcfg/build.go
@@ -33,6 +33,11 @@ func BuildGitalyGit2Go(t testing.TB, cfg config.Cfg) string {
 	return symlinkPath
 }
 
+// BuildGitalyWrapper builds the gitaly-wrapper command and installs it into the binary directory.
+func BuildGitalyWrapper(t *testing.T, cfg config.Cfg) string {
+	return BuildBinary(t, cfg.BinDir, gitalyCommandPath("gitaly-wrapper"))
+}
+
 // BuildGitalyLFSSmudge builds the gitaly-lfs-smudge command and installs it into the binary
 // directory.
 func BuildGitalyLFSSmudge(t *testing.T, cfg config.Cfg) string {
diff --git a/internal/testhelper/testhelper.go b/internal/testhelper/testhelper.go
index b98a7f4d8..a1227fd93 100644
--- a/internal/testhelper/testhelper.go
+++ b/internal/testhelper/testhelper.go
@@ -221,8 +221,9 @@ func TempDir(t testing.TB) string {
 type Cleanup func()
 
 // WriteExecutable ensures that the parent directory exists, and writes an executable with provided
-// content. The executable must not exist previous to writing it.
-func WriteExecutable(t testing.TB, path string, content []byte) {
+// content. The executable must not exist previous to writing it. Returns the path of the written
+// executable.
+func WriteExecutable(t testing.TB, path string, content []byte) string {
 	dir := filepath.Dir(path)
 	require.NoError(t, os.MkdirAll(dir, 0o755))
 
@@ -236,6 +237,8 @@ func WriteExecutable(t testing.TB, path string, content []byte) {
 	t.Cleanup(func() {
 		assert.NoError(t, os.RemoveAll(dir))
 	})
+
+	return path
 }
 
 // ModifyEnvironment will change an environment variable and return a func suitable