Prevent catfile error return resource leak

4 files changed, 144 insertions, 2 deletions

diff --git a/internal/git/catfile/batch.go b/internal/git/catfile/batch.go
index 4c7f68193..5bb4e5c65 100644
--- a/internal/git/catfile/batch.go
+++ b/internal/git/catfile/batch.go
@@ -51,6 +51,11 @@ func newBatchProcess(ctx context.Context, repoPath string, env []string) (*batch
 		currentCatfileProcesses.Dec()
 	}()
 
+	if injectSpawnErrors {
+		// Testing only: intentionally leak process
+		return nil, &simulatedBatchSpawnError{}
+	}
+
 	return b, nil
 }
 
diff --git a/internal/git/catfile/batchcheck.go b/internal/git/catfile/batchcheck.go
index 77f49a05b..199341656 100644
--- a/internal/git/catfile/batchcheck.go
+++ b/internal/git/catfile/batchcheck.go
@@ -35,6 +35,11 @@ func newBatchCheck(ctx context.Context, repoPath string, env []string) (*batchCh
 		bc.w.Close()
 	}()
 
+	if injectSpawnErrors {
+		// Testing only: intentionally leak process
+		return nil, &simulatedBatchSpawnError{}
+	}
+
 	return bc, nil
 }
 
diff --git a/internal/git/catfile/catfile.go b/internal/git/catfile/catfile.go
index 1b29ec4af..897282836 100644
--- a/internal/git/catfile/catfile.go
+++ b/internal/git/catfile/catfile.go
@@ -45,6 +45,9 @@ const (
 	// CacheFeatureFlagKey is the feature flag key for catfile batch caching. This should match
 	// what is in gitlab-ce
 	CacheFeatureFlagKey = "catfile-cache"
+
+	// SessionIDField is the gRPC metadata field we use to store the gitaly session ID.
+	SessionIDField = "gitaly-session-id"
 )
 
 func init() {
@@ -146,7 +149,7 @@ func New(ctx context.Context, repo *gitalypb.Repository) (*Batch, error) {
 		return nil, err
 	}
 
-	sessionID := metadata.GetValue(ctx, "gitaly-session-id")
+	sessionID := metadata.GetValue(ctx, SessionIDField)
 	if sessionID == "" {
 		return newBatch(ctx, repoPath, env)
 	}
@@ -189,7 +192,20 @@ func returnWhenDone(done <-chan struct{}, bc *batchCache, cacheKey key, c *Batch
 	bc.Add(cacheKey, c)
 }
 
-func newBatch(ctx context.Context, repoPath string, env []string) (*Batch, error) {
+var injectSpawnErrors = false
+
+type simulatedBatchSpawnError struct{}
+
+func (simulatedBatchSpawnError) Error() string { return "simulated spawn error" }
+
+func newBatch(_ctx context.Context, repoPath string, env []string) (_ *Batch, err error) {
+	ctx, cancel := context.WithCancel(_ctx)
+	defer func() {
+		if err != nil {
+			cancel()
+		}
+	}()
+
 	batch, err := newBatchProcess(ctx, repoPath, env)
 	if err != nil {
 		return nil, err
diff --git a/internal/git/catfile/catfile_test.go b/internal/git/catfile/catfile_test.go
index b0ec2064c..8825afbd1 100644
--- a/internal/git/catfile/catfile_test.go
+++ b/internal/git/catfile/catfile_test.go
@@ -1,12 +1,21 @@
 package catfile
 
 import (
+	"bytes"
+	"context"
 	"io"
 	"io/ioutil"
+	"os"
+	"os/exec"
+	"strconv"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/require"
+	"gitlab.com/gitlab-org/gitaly/internal/command"
+	"gitlab.com/gitlab-org/gitaly/internal/helper/text"
 	"gitlab.com/gitlab-org/gitaly/internal/testhelper"
+	"google.golang.org/grpc/metadata"
 )
 
 func TestInfo(t *testing.T) {
@@ -226,3 +235,110 @@ func TestRepeatedCalls(t *testing.T) {
 
 	require.Equal(t, string(treeBytes), string(tree2))
 }
+
+func TestSpawnFailure(t *testing.T) {
+	defer func() { injectSpawnErrors = false }()
+
+	defer func(bc *batchCache) {
+		// reset global cache
+		cache = bc
+	}(cache)
+
+	// Use very high values to effectively disable auto-expiry
+	testCache := newCache(1*time.Hour, 1000)
+	cache = testCache
+	defer testCache.EvictAll()
+
+	require.True(
+		t,
+		waitTrue(func() bool { return numGitChildren(t) == 0 }),
+		"test setup: wait for there to be 0 git children",
+	)
+	require.Equal(t, 0, cacheSize(testCache), "sanity check: cache empty")
+
+	ctx1, cancel1 := testhelper.Context()
+	defer cancel1()
+
+	injectSpawnErrors = false
+	_, err := catfileWithFreshSessionID(ctx1)
+	require.NoError(t, err, "catfile spawn should succeed in normal circumstances")
+	require.Equal(t, 2, numGitChildren(t), "there should be 2 git child processes")
+
+	// cancel request context: this should asynchronously move the processes into the cat-file cache
+	cancel1()
+
+	require.True(
+		t,
+		waitTrue(func() bool { return cacheSize(testCache) == 1 }),
+		"1 cache entry, meaning 2 processes, should be in the cache now",
+	)
+
+	require.Equal(t, 2, numGitChildren(t), "there should still be 2 git child processes")
+
+	testCache.EvictAll()
+	require.Equal(t, 0, cacheSize(testCache), "the cache should be empty now")
+
+	require.True(
+		t,
+		waitTrue(func() bool { return numGitChildren(t) == 0 }),
+		"number of git processes should drop to 0 again",
+	)
+
+	ctx2, cancel2 := testhelper.Context()
+	defer cancel2()
+
+	injectSpawnErrors = true
+	_, err = catfileWithFreshSessionID(ctx2)
+	require.Error(t, err, "expect simulated error")
+	require.IsType(t, &simulatedBatchSpawnError{}, err)
+
+	require.True(
+		t,
+		waitTrue(func() bool { return numGitChildren(t) == 0 }),
+		"there should be no git children after spawn failure scenario",
+	)
+}
+
+func catfileWithFreshSessionID(ctx context.Context) (*Batch, error) {
+	id, err := text.RandomHex(4)
+	if err != nil {
+		return nil, err
+	}
+
+	md := metadata.New(map[string]string{
+		SessionIDField: id,
+	})
+
+	return New(metadata.NewIncomingContext(ctx, md), testhelper.TestRepository())
+}
+
+func waitTrue(callback func() bool) bool {
+	for start := time.Now(); time.Since(start) < 1*time.Second; time.Sleep(1 * time.Millisecond) {
+		if callback() {
+			return true
+		}
+	}
+
+	return false
+}
+
+func numGitChildren(t *testing.T) int {
+	out, err := exec.Command("pgrep", "-x", "-P", strconv.Itoa(os.Getpid()), "git").Output()
+
+	if err != nil {
+		if code, ok := command.ExitStatus(err); ok && code == 1 {
+			// pgrep exit code 1 means: no processes found
+			return 0
+		}
+
+		t.Fatal(err)
+	}
+
+	return bytes.Count(out, []byte("\n"))
+}
+
+func cacheSize(bc *batchCache) int {
+	bc.Lock()
+	defer bc.Unlock()
+	return bc.len()
+}