Git object extraction improvements

Omit extracting core object information before extracting actual object.
If extracted object has type different from expected the NotFound error returns.
Each extracted object represented as instance of Object struct with info
and the reader inside. Based on the info reader could be consumed and
properly parsed.
Moved from pointer to value for ObjectInfo (reduce work for GC).

Relates to: https://gitlab.com/gitlab-org/gitaly/issues/2255

22 files changed, 176 insertions, 46 deletions

diff --git a/changelogs/unreleased/ps-git-object-retrieval.yml b/changelogs/unreleased/ps-git-object-retrieval.yml
new file mode 100644
index 000000000..c254cc0b8
--- /dev/null
+++ b/changelogs/unreleased/ps-git-object-retrieval.yml
@@ -0,0 +1,5 @@
+---
+title: 'Feature flag: look up commits without batch-check'
+merge_request: 1711
+author:
+type: added
diff --git a/internal/git/catfile/batch.go b/internal/git/catfile/batch.go
index 0099ef6dd..aec8fa508 100644
--- a/internal/git/catfile/batch.go
+++ b/internal/git/catfile/batch.go
@@ -69,7 +69,7 @@ func newBatchProcess(ctx context.Context, repo repository.GitRepo) (*batchProces
 	return b, nil
 }
 
-func (b *batchProcess) reader(revspec string, expectedType string) (io.Reader, error) {
+func (b *batchProcess) reader(revspec string, expectedType string) (*Object, error) {
 	b.Lock()
 	defer b.Unlock()
 
@@ -82,7 +82,7 @@ func (b *batchProcess) reader(revspec string, expectedType string) (io.Reader, e
 	}
 
 	if b.n != 0 {
-		return nil, fmt.Errorf("cannot create new reader: batch contains %d unread bytes", b.n)
+		return nil, fmt.Errorf("cannot create new Object: batch contains %d unread bytes", b.n)
 	}
 
 	if _, err := fmt.Fprintln(b.w, revspec); err != nil {
@@ -104,12 +104,15 @@ func (b *batchProcess) reader(revspec string, expectedType string) (io.Reader, e
 		}
 		b.n = 0
 
-		return nil, fmt.Errorf("expected %s to be a %s, got %s", oi.Oid, expectedType, oi.Type)
+		return nil, NotFoundError{error: fmt.Errorf("expected %s to be a %s, got %s", oi.Oid, expectedType, oi.Type)}
 	}
 
-	return &batchReader{
-		batchProcess: b,
-		r:            io.LimitReader(b.r, oi.Size),
+	return &Object{
+		ObjectInfo: *oi,
+		Reader: &batchReader{
+			batchProcess: b,
+			r:            io.LimitReader(b.r, oi.Size),
+		},
 	}, nil
 }
 
diff --git a/internal/git/catfile/catfile.go b/internal/git/catfile/catfile.go
index 1226971bb..a1ce1fbb6 100644
--- a/internal/git/catfile/catfile.go
+++ b/internal/git/catfile/catfile.go
@@ -89,7 +89,7 @@ func (c *Batch) Tree(revspec string) (io.Reader, error) {
 // point to a commit. To prevent this first use Info to resolve the revspec
 // and check the object type. Caller must consume the Reader before
 // making another call on C.
-func (c *Batch) Commit(revspec string) (io.Reader, error) {
+func (c *Batch) Commit(revspec string) (*Object, error) {
 	catfileLookupCounter.WithLabelValues("commit").Inc()
 	return c.batchProcess.reader(revspec, "commit")
 }
diff --git a/internal/git/catfile/object.go b/internal/git/catfile/object.go
new file mode 100644
index 000000000..04363ad21
--- /dev/null
+++ b/internal/git/catfile/object.go
@@ -0,0 +1,13 @@
+package catfile
+
+import (
+	"io"
+)
+
+// Object represents data returned by `git cat-file --batch`
+type Object struct {
+	// ObjectInfo represents main information about object
+	ObjectInfo
+	// Reader provides raw data about object. It differs for each type of object(tag, commit, tree, log, etc.)
+	io.Reader
+}
diff --git a/internal/git/log/commit.go b/internal/git/log/commit.go
index 421bfd59c..282a54404 100644
--- a/internal/git/log/commit.go
+++ b/internal/git/log/commit.go
@@ -13,6 +13,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/internal/git"
 	"gitlab.com/gitlab-org/gitaly/internal/git/catfile"
 	"gitlab.com/gitlab-org/gitaly/internal/helper"
+	"gitlab.com/gitlab-org/gitaly/internal/metadata/featureflag"
 	"gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
 )
 
@@ -24,11 +25,18 @@ func GetCommit(ctx context.Context, repo *gitalypb.Repository, revision string)
 		return nil, err
 	}
 
-	return GetCommitCatfile(c, revision)
+	return GetCommitCatfile(ctx, c, revision)
 }
 
 // GetCommitCatfile looks up a commit by revision using an existing *catfile.Batch instance.
-func GetCommitCatfile(c *catfile.Batch, revision string) (*gitalypb.GitCommit, error) {
+func GetCommitCatfile(ctx context.Context, c *catfile.Batch, revision string) (*gitalypb.GitCommit, error) {
+	if featureflag.IsEnabled(ctx, featureflag.CommitWithoutBatchCheck) {
+		return getCommitCatfileNew(c, revision)
+	}
+	return getCommitCatfileOld(c, revision)
+}
+
+func getCommitCatfileOld(c *catfile.Batch, revision string) (*gitalypb.GitCommit, error) {
 	info, err := c.Info(revision + "^{commit}")
 	if err != nil {
 		return nil, err
@@ -42,6 +50,15 @@ func GetCommitCatfile(c *catfile.Batch, revision string) (*gitalypb.GitCommit, e
 	return parseRawCommit(r, info)
 }
 
+func getCommitCatfileNew(c *catfile.Batch, revision string) (*gitalypb.GitCommit, error) {
+	obj, err := c.Commit(revision + "^{commit}")
+	if err != nil {
+		return nil, err
+	}
+
+	return parseRawCommit(obj.Reader, &obj.ObjectInfo)
+}
+
 // GetCommitMessage looks up a commit message and returns it in its entirety.
 func GetCommitMessage(c *catfile.Batch, repo *gitalypb.Repository, revision string) ([]byte, error) {
 	info, err := c.Info(revision + "^{commit}")
diff --git a/internal/git/log/commit_test.go b/internal/git/log/commit_test.go
index 120604859..8ca3c1ec4 100644
--- a/internal/git/log/commit_test.go
+++ b/internal/git/log/commit_test.go
@@ -7,7 +7,10 @@ import (
 	"github.com/golang/protobuf/ptypes/timestamp"
 	"github.com/stretchr/testify/require"
 	"gitlab.com/gitlab-org/gitaly/internal/git/catfile"
+	"gitlab.com/gitlab-org/gitaly/internal/metadata/featureflag"
+	"gitlab.com/gitlab-org/gitaly/internal/testhelper"
 	"gitlab.com/gitlab-org/gitaly/proto/go/gitalypb"
+	"google.golang.org/grpc/metadata"
 )
 
 func TestParseRawCommit(t *testing.T) {
@@ -102,3 +105,77 @@ func TestParseRawCommit(t *testing.T) {
 		})
 	}
 }
+
+func TestGetCommitCatfile(t *testing.T) {
+	ctx, cancel := testhelper.Context()
+	ctx = metadata.NewIncomingContext(ctx, metadata.MD{})
+	defer cancel()
+
+	testRepo, _, cleanup := testhelper.NewTestRepo(t)
+	defer cleanup()
+
+	const commitSha = "2d1db523e11e777e49377cfb22d368deec3f0793"
+	const commitMsg = "Correct test_env.rb path for adding branch\n"
+	const blobSha = "c60514b6d3d6bf4bec1030f70026e34dfbd69ad5"
+
+	testCases := []struct {
+		desc      string
+		revision  string
+		featureOn bool
+		errStr    string
+	}{
+		{
+			desc:      "feature off | commit",
+			revision:  commitSha,
+			featureOn: false,
+		},
+		{
+			desc:      "feature on | commit",
+			revision:  commitSha,
+			featureOn: true,
+		},
+		{
+			desc:      "feature off | not existing commit",
+			revision:  "not existing revision",
+			featureOn: false,
+			errStr:    "object not found",
+		},
+		{
+			desc:      "feature on | not existing commit",
+			revision:  "not existing revision",
+			featureOn: true,
+			errStr:    "object not found",
+		},
+		{
+			desc:      "feature off | blob sha",
+			revision:  blobSha,
+			featureOn: false,
+			errStr:    "object not found",
+		},
+		{
+			desc:      "feature on | blob sha",
+			revision:  blobSha,
+			featureOn: true,
+			errStr:    "object not found",
+		},
+	}
+
+	c, err := catfile.New(ctx, testRepo)
+	require.NoError(t, err)
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			if tc.featureOn {
+				ctx = featureflag.ContextWithFeatureFlag(ctx, featureflag.CommitWithoutBatchCheck)
+			}
+
+			c, err := GetCommitCatfile(ctx, c, tc.revision)
+
+			if tc.errStr == "" {
+				require.NoError(t, err)
+				require.Equal(t, commitMsg, string(c.Body))
+			} else {
+				require.EqualError(t, err, tc.errStr)
+			}
+		})
+	}
+}
diff --git a/internal/git/log/last_commit.go b/internal/git/log/last_commit.go
index 76b3c9915..763033f30 100644
--- a/internal/git/log/last_commit.go
+++ b/internal/git/log/last_commit.go
@@ -27,7 +27,7 @@ func LastCommitForPath(ctx context.Context, batch *catfile.Batch, repo *gitalypb
 		return nil, err
 	}
 
-	return GetCommitCatfile(batch, text.ChompBytes(commitID))
+	return GetCommitCatfile(ctx, batch, text.ChompBytes(commitID))
 }
 
 // GitLogCommand returns a Command that executes git log with the given the arguments
diff --git a/internal/git/log/log.go b/internal/git/log/log.go
index c591f36bf..e6de9c37f 100644
--- a/internal/git/log/log.go
+++ b/internal/git/log/log.go
@@ -19,6 +19,7 @@ type Parser struct {
 	currentCommit *gitalypb.GitCommit
 	err           error
 	c             *catfile.Batch
+	ctx           context.Context
 }
 
 // NewLogParser returns a new Parser
@@ -31,6 +32,7 @@ func NewLogParser(ctx context.Context, repo *gitalypb.Repository, src io.Reader)
 	parser := &Parser{
 		scanner: bufio.NewScanner(src),
 		c:       c,
+		ctx:     ctx,
 	}
 
 	return parser, nil
@@ -46,7 +48,7 @@ func (parser *Parser) Parse() bool {
 
 	commitID := parser.scanner.Text()
 
-	commit, err := GetCommitCatfile(parser.c, commitID)
+	commit, err := GetCommitCatfile(parser.ctx, parser.c, commitID)
 	if err != nil {
 		parser.err = err
 		return false
diff --git a/internal/git/log/tag.go b/internal/git/log/tag.go
index 6e78cb0db..ae82afd74 100644
--- a/internal/git/log/tag.go
+++ b/internal/git/log/tag.go
@@ -3,6 +3,7 @@ package log
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"io/ioutil"
@@ -23,7 +24,7 @@ const (
 // When 'trimRightNewLine' is 'true', the tag message will be trimmed to remove all '\n' characters from right.
 // note: we pass in the tagName because the tag name from refs/tags may be different
 // than the name found in the actual tag object. We want to use the tagName found in refs/tags
-func GetTagCatfile(c *catfile.Batch, tagID, tagName string, trimLen, trimRightNewLine bool) (*gitalypb.Tag, error) {
+func GetTagCatfile(ctx context.Context, c *catfile.Batch, tagID, tagName string, trimLen, trimRightNewLine bool) (*gitalypb.Tag, error) {
 	r, err := c.Tag(tagID)
 	if err != nil {
 		return nil, err
@@ -35,7 +36,7 @@ func GetTagCatfile(c *catfile.Batch, tagID, tagName string, trimLen, trimRightNe
 	}
 
 	// the tagID is the oid of the tag object
-	tag, err := buildAnnotatedTag(c, tagID, tagName, header, body, trimLen, trimRightNewLine)
+	tag, err := buildAnnotatedTag(ctx, c, tagID, tagName, header, body, trimLen, trimRightNewLine)
 	if err != nil {
 		return nil, err
 	}
@@ -92,7 +93,7 @@ func splitRawTag(r io.Reader, trimRightNewLine bool) (*tagHeader, []byte, error)
 	return &header, body, nil
 }
 
-func buildAnnotatedTag(b *catfile.Batch, tagID, name string, header *tagHeader, body []byte, trimLen, trimRightNewLine bool) (*gitalypb.Tag, error) {
+func buildAnnotatedTag(ctx context.Context, b *catfile.Batch, tagID, name string, header *tagHeader, body []byte, trimLen, trimRightNewLine bool) (*gitalypb.Tag, error) {
 	tag := &gitalypb.Tag{
 		Id:          tagID,
 		Name:        []byte(name),
@@ -107,13 +108,13 @@ func buildAnnotatedTag(b *catfile.Batch, tagID, name string, header *tagHeader,
 	var err error
 	switch header.tagType {
 	case "commit":
-		tag.TargetCommit, err = GetCommitCatfile(b, header.oid)
+		tag.TargetCommit, err = GetCommitCatfile(ctx, b, header.oid)
 		if err != nil {
 			return nil, fmt.Errorf("buildAnnotatedTag error when getting target commit: %v", err)
 		}
 
 	case "tag":
-		tag.TargetCommit, err = dereferenceTag(b, header.oid)
+		tag.TargetCommit, err = dereferenceTag(ctx, b, header.oid)
 		if err != nil {
 			return nil, fmt.Errorf("buildAnnotatedTag error when dereferencing tag: %v", err)
 		}
@@ -137,7 +138,7 @@ func buildAnnotatedTag(b *catfile.Batch, tagID, name string, header *tagHeader,
 // we also protect against circular tag references. Even though this is not possible in git,
 // we still want to protect against an infinite looop
 
-func dereferenceTag(b *catfile.Batch, Oid string) (*gitalypb.GitCommit, error) {
+func dereferenceTag(ctx context.Context, b *catfile.Batch, Oid string) (*gitalypb.GitCommit, error) {
 	for depth := 0; depth < MaxTagReferenceDepth; depth++ {
 		i, err := b.Info(Oid)
 		if err != nil {
@@ -159,7 +160,7 @@ func dereferenceTag(b *catfile.Batch, Oid string) (*gitalypb.GitCommit, error) {
 			Oid = header.oid
 			continue
 		case "commit":
-			return GetCommitCatfile(b, Oid)
+			return GetCommitCatfile(ctx, b, Oid)
 		default: // This current tag points to a tree or a blob
 			return nil, nil
 		}
diff --git a/internal/git/log/tag_test.go b/internal/git/log/tag_test.go
index 2aadb8511..ab975ee25 100644
--- a/internal/git/log/tag_test.go
+++ b/internal/git/log/tag_test.go
@@ -52,7 +52,7 @@ func TestGetTag(t *testing.T) {
 		t.Run(testCase.tagName, func(t *testing.T) {
 			tagID := testhelper.CreateTag(t, testRepoPath, testCase.tagName, testCase.rev, &testhelper.CreateTagOpts{Message: testCase.message})
 
-			tag, err := GetTagCatfile(c, tagID, testCase.tagName, testCase.trim, true)
+			tag, err := GetTagCatfile(ctx, c, tagID, testCase.tagName, testCase.trim, true)
 			require.NoError(t, err)
 			if testCase.trim && len(testCase.message) >= helper.MaxCommitOrTagMessageSize {
 				testCase.message = testCase.message[:helper.MaxCommitOrTagMessageSize]
diff --git a/internal/metadata/featureflag/feature_flags.go b/internal/metadata/featureflag/feature_flags.go
index c2178e554..a11d2cbee 100644
--- a/internal/metadata/featureflag/feature_flags.go
+++ b/internal/metadata/featureflag/feature_flags.go
@@ -12,6 +12,10 @@ const (
 	// annotations (i.e. accessor and mutator RPC's). This enables cache
 	// invalidation by changing state when the repo is modified.
 	CacheInvalidator = "cache_invalidator"
+	// CommitWithoutBatchCheck controls which implementation of the GetCommitCatfile needs to be used.
+	// The old one with query fot Info before fetching info about Commit
+	// or the new one that skips Info call and checks object type in Commit method call.
+	CommitWithoutBatchCheck = "commit_without_batch_check"
 )
 
 const (
diff --git a/internal/metadata/featureflag/test_utils.go b/internal/metadata/featureflag/test_utils.go
index c0711a92d..8131b82a4 100644
--- a/internal/metadata/featureflag/test_utils.go
+++ b/internal/metadata/featureflag/test_utils.go
@@ -6,7 +6,7 @@ import (
 	"google.golang.org/grpc/metadata"
 )
 
-// ContextWithFeatureFlag is used in tests to enablea a feature flag in the context metadata
+// ContextWithFeatureFlag is used in tests to enable a feature flag in the context metadata
 func ContextWithFeatureFlag(ctx context.Context, flag string) context.Context {
 	md, ok := metadata.FromIncomingContext(ctx)
 	if !ok {
diff --git a/internal/service/commit/filter_shas_with_signatures.go b/internal/service/commit/filter_shas_with_signatures.go
index 1a6def272..6af3475d0 100644
--- a/internal/service/commit/filter_shas_with_signatures.go
+++ b/internal/service/commit/filter_shas_with_signatures.go
@@ -1,6 +1,7 @@
 package commit
 
 import (
+	"context"
 	"errors"
 	"io"
 
@@ -34,14 +35,15 @@ func validateFirstFilterShasWithSignaturesRequest(in *gitalypb.FilterShasWithSig
 }
 
 func (s *server) filterShasWithSignatures(bidi gitalypb.CommitService_FilterShasWithSignaturesServer, firstRequest *gitalypb.FilterShasWithSignaturesRequest) error {
-	c, err := catfile.New(bidi.Context(), firstRequest.GetRepository())
+	ctx := bidi.Context()
+	c, err := catfile.New(ctx, firstRequest.GetRepository())
 	if err != nil {
 		return err
 	}
 
 	var request = firstRequest
 	for {
-		shas, err := filterCommitShasWithSignatures(c, request.GetShas())
+		shas, err := filterCommitShasWithSignatures(ctx, c, request.GetShas())
 		if err != nil {
 			return err
 		}
@@ -61,10 +63,10 @@ func (s *server) filterShasWithSignatures(bidi gitalypb.CommitService_FilterShas
 	}
 }
 
-func filterCommitShasWithSignatures(c *catfile.Batch, shas [][]byte) ([][]byte, error) {
+func filterCommitShasWithSignatures(ctx context.Context, c *catfile.Batch, shas [][]byte) ([][]byte, error) {
 	var foundShas [][]byte
 	for _, sha := range shas {
-		commit, err := log.GetCommitCatfile(c, string(sha))
+		commit, err := log.GetCommitCatfile(ctx, c, string(sha))
 		if catfile.IsNotFound(err) {
 			continue
 		}
diff --git a/internal/service/commit/find_commits.go b/internal/service/commit/find_commits.go
index 3fb3e395c..b31e44bf8 100644
--- a/internal/service/commit/find_commits.go
+++ b/internal/service/commit/find_commits.go
@@ -59,7 +59,7 @@ func findCommits(ctx context.Context, req *gitalypb.FindCommitsRequest, stream g
 		return fmt.Errorf("creating catfile: %v", err)
 	}
 
-	getCommits := NewGetCommits(logCmd, batch)
+	getCommits := NewGetCommits(ctx, logCmd, batch)
 
 	if calculateOffsetManually(req) {
 		getCommits.Offset(int(req.GetOffset()))
@@ -79,13 +79,15 @@ func calculateOffsetManually(req *gitalypb.FindCommitsRequest) bool {
 type GetCommits struct {
 	scanner *bufio.Scanner
 	batch   *catfile.Batch
+	ctx     context.Context
 }
 
 // NewGetCommits returns a new GetCommits object
-func NewGetCommits(cmd *command.Command, batch *catfile.Batch) *GetCommits {
+func NewGetCommits(ctx context.Context, cmd *command.Command, batch *catfile.Batch) *GetCommits {
 	return &GetCommits{
 		scanner: bufio.NewScanner(cmd),
 		batch:   batch,
+		ctx:     ctx,
 	}
 }
 
@@ -112,7 +114,7 @@ func (g *GetCommits) Offset(offset int) error {
 // Commit returns the current commit
 func (g *GetCommits) Commit() (*gitalypb.GitCommit, error) {
 	revision := strings.TrimSpace(g.scanner.Text())
-	commit, err := log.GetCommitCatfile(g.batch, revision)
+	commit, err := log.GetCommitCatfile(g.ctx, g.batch, revision)
 	if err != nil {
 		return nil, fmt.Errorf("cat-file get commit %q: %v", revision, err)
 	}
diff --git a/internal/service/commit/list_commits_by_oid.go b/internal/service/commit/list_commits_by_oid.go
index 9d024566d..caa08a90a 100644
--- a/internal/service/commit/list_commits_by_oid.go
+++ b/internal/service/commit/list_commits_by_oid.go
@@ -37,7 +37,7 @@ func (s *server) ListCommitsByOid(in *gitalypb.ListCommitsByOidRequest, stream g
 	listCommitsbyOidHistogram.Observe(float64(len(in.Oid)))
 
 	for _, oid := range in.Oid {
-		commit, err := gitlog.GetCommitCatfile(c, oid)
+		commit, err := gitlog.GetCommitCatfile(ctx, c, oid)
 		if catfile.IsNotFound(err) {
 			continue
 		}
diff --git a/internal/service/commit/list_commits_by_ref_name.go b/internal/service/commit/list_commits_by_ref_name.go
index baf281214..24379008b 100644
--- a/internal/service/commit/list_commits_by_ref_name.go
+++ b/internal/service/commit/list_commits_by_ref_name.go
@@ -19,7 +19,7 @@ func (s *server) ListCommitsByRefName(in *gitalypb.ListCommitsByRefNameRequest,
 	sender := chunk.New(&commitsByRefNameSender{stream: stream})
 
 	for _, refName := range in.RefNames {
-		commit, err := gitlog.GetCommitCatfile(c, string(refName))
+		commit, err := gitlog.GetCommitCatfile(ctx, c, string(refName))
 		if catfile.IsNotFound(err) {
 			continue
 		}
diff --git a/internal/service/ref/list_new_commits.go b/internal/service/ref/list_new_commits.go
index d465dca05..e5533e42a 100644
--- a/internal/service/ref/list_new_commits.go
+++ b/internal/service/ref/list_new_commits.go
@@ -45,7 +45,7 @@ func listNewCommits(in *gitalypb.ListNewCommitsRequest, stream gitalypb.RefServi
 	for scanner.Scan() {
 		line := scanner.Text()
 
-		commit, err := log.GetCommitCatfile(batch, line)
+		commit, err := log.GetCommitCatfile(ctx, batch, line)
 		if err != nil {
 			return err
 		}
diff --git a/internal/service/ref/refs.go b/internal/service/ref/refs.go
index f7de0d344..0d1ee5de6 100644
--- a/internal/service/ref/refs.go
+++ b/internal/service/ref/refs.go
@@ -100,7 +100,7 @@ func parseAndReturnTags(ctx context.Context, repo *gitalypb.Repository, stream g
 
 	scanner := bufio.NewScanner(tagsCmd)
 	for scanner.Scan() {
-		tag, err := parseTagLine(c, scanner.Text())
+		tag, err := parseTagLine(ctx, c, scanner.Text())
 		if err != nil {
 			return fmt.Errorf("parsing tag: %v", err)
 		}
@@ -360,7 +360,7 @@ func (s *server) FindTag(ctx context.Context, in *gitalypb.FindTagRequest) (*git
 }
 
 // parseTagLine parses a line of text with the output format %(objectname) %(objecttype) %(refname:lstrip=2)
-func parseTagLine(c *catfile.Batch, tagLine string) (*gitalypb.Tag, error) {
+func parseTagLine(ctx context.Context, c *catfile.Batch, tagLine string) (*gitalypb.Tag, error) {
 	fields := strings.SplitN(tagLine, " ", 3)
 	if len(fields) != 3 {
 		return nil, fmt.Errorf("invalid output from for-each-ref command: %v", tagLine)
@@ -376,13 +376,13 @@ func parseTagLine(c *catfile.Batch, tagLine string) (*gitalypb.Tag, error) {
 	switch refType {
 	// annotated tag
 	case "tag":
-		tag, err := gitlog.GetTagCatfile(c, tagID, refName, true, true)
+		tag, err := gitlog.GetTagCatfile(ctx, c, tagID, refName, true, true)
 		if err != nil {
 			return nil, fmt.Errorf("getting annotated tag: %v", err)
 		}
 		return tag, nil
 	case "commit":
-		commit, err := gitlog.GetCommitCatfile(c, tagID)
+		commit, err := gitlog.GetCommitCatfile(ctx, c, tagID)
 		if err != nil {
 			return nil, fmt.Errorf("getting commit catfile: %v", err)
 		}
@@ -414,7 +414,7 @@ func findTag(ctx context.Context, repository *gitalypb.Repository, tagName []byt
 
 	scanner := bufio.NewScanner(tagCmd)
 	if scanner.Scan() {
-		tag, err = parseTagLine(c, scanner.Text())
+		tag, err = parseTagLine(ctx, c, scanner.Text())
 		if err != nil {
 			return nil, err
 		}
diff --git a/internal/service/ref/refs_test.go b/internal/service/ref/refs_test.go
index 247825ba1..38e63637b 100644
--- a/internal/service/ref/refs_test.go
+++ b/internal/service/ref/refs_test.go
@@ -744,7 +744,7 @@ func TestFindAllTagNestedTags(t *testing.T) {
 
 				// only expect the TargetCommit to be populated if it is a commit and if its less than 10 tags deep
 				if info.Type == "commit" && depth < log.MaxTagReferenceDepth {
-					commit, err := log.GetCommitCatfile(batch, tc.originalOid)
+					commit, err := log.GetCommitCatfile(ctx, batch, tc.originalOid)
 					require.NoError(t, err)
 					expectedTag.TargetCommit = commit
 				}
@@ -1665,7 +1665,7 @@ func TestFindTagNestedTag(t *testing.T) {
 			}
 			// only expect the TargetCommit to be populated if it is a commit and if its less than 10 tags deep
 			if info.Type == "commit" && tc.depth < log.MaxTagReferenceDepth {
-				commit, err := log.GetCommitCatfile(batch, tc.originalOid)
+				commit, err := log.GetCommitCatfile(ctx, batch, tc.originalOid)
 				require.NoError(t, err)
 				expectedTag.TargetCommit = commit
 			}
diff --git a/internal/service/ref/tag_messages.go b/internal/service/ref/tag_messages.go
index e632dd7f8..d4ec3e07e 100644
--- a/internal/service/ref/tag_messages.go
+++ b/internal/service/ref/tag_messages.go
@@ -42,7 +42,7 @@ func (s *server) getAndStreamTagMessages(request *gitalypb.GetTagMessagesRequest
 	}
 
 	for _, tagID := range request.GetTagIds() {
-		tag, err := log.GetTagCatfile(c, tagID, "", false, false)
+		tag, err := log.GetTagCatfile(ctx, c, tagID, "", false, false)
 		if err != nil {
 			return fmt.Errorf("failed to get tag: %v", err)
 		}
diff --git a/internal/service/ref/util.go b/internal/service/ref/util.go
index eecc0c49d..f648bc5e2 100644
--- a/internal/service/ref/util.go
+++ b/internal/service/ref/util.go
@@ -2,6 +2,7 @@ package ref
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 
 	"gitlab.com/gitlab-org/gitaly/internal/git/catfile"
@@ -53,8 +54,8 @@ func buildLocalBranch(name []byte, target *gitalypb.GitCommit) *gitalypb.FindLoc
 	return response
 }
 
-func buildAllBranchesBranch(c *catfile.Batch, elements [][]byte) (*gitalypb.FindAllBranchesResponse_Branch, error) {
-	target, err := log.GetCommitCatfile(c, string(elements[1]))
+func buildAllBranchesBranch(ctx context.Context, c *catfile.Batch, elements [][]byte) (*gitalypb.FindAllBranchesResponse_Branch, error) {
+	target, err := log.GetCommitCatfile(ctx, c, string(elements[1]))
 	if err != nil {
 		return nil, err
 	}
@@ -65,8 +66,8 @@ func buildAllBranchesBranch(c *catfile.Batch, elements [][]byte) (*gitalypb.Find
 	}, nil
 }
 
-func buildBranch(c *catfile.Batch, elements [][]byte) (*gitalypb.Branch, error) {
-	target, err := log.GetCommitCatfile(c, string(elements[1]))
+func buildBranch(ctx context.Context, c *catfile.Batch, elements [][]byte) (*gitalypb.Branch, error) {
+	target, err := log.GetCommitCatfile(ctx, c, string(elements[1]))
 	if err != nil {
 		return nil, err
 	}
@@ -78,6 +79,7 @@ func buildBranch(c *catfile.Batch, elements [][]byte) (*gitalypb.Branch, error)
 }
 
 func newFindLocalBranchesWriter(stream gitalypb.RefService_FindLocalBranchesServer, c *catfile.Batch) lines.Sender {
+	ctx := stream.Context()
 	return func(refs [][]byte) error {
 		var branches []*gitalypb.FindLocalBranchResponse
 
@@ -87,7 +89,7 @@ func newFindLocalBranchesWriter(stream gitalypb.RefService_FindLocalBranchesServ
 				return err
 			}
 
-			target, err := log.GetCommitCatfile(c, string(elements[1]))
+			target, err := log.GetCommitCatfile(ctx, c, string(elements[1]))
 			if err != nil {
 				return err
 			}
@@ -99,6 +101,7 @@ func newFindLocalBranchesWriter(stream gitalypb.RefService_FindLocalBranchesServ
 }
 
 func newFindAllBranchesWriter(stream gitalypb.RefService_FindAllBranchesServer, c *catfile.Batch) lines.Sender {
+	ctx := stream.Context()
 	return func(refs [][]byte) error {
 		var branches []*gitalypb.FindAllBranchesResponse_Branch
 
@@ -107,7 +110,7 @@ func newFindAllBranchesWriter(stream gitalypb.RefService_FindAllBranchesServer,
 			if err != nil {
 				return err
 			}
-			branch, err := buildAllBranchesBranch(c, elements)
+			branch, err := buildAllBranchesBranch(ctx, c, elements)
 			if err != nil {
 				return err
 			}
@@ -118,6 +121,7 @@ func newFindAllBranchesWriter(stream gitalypb.RefService_FindAllBranchesServer,
 }
 
 func newFindAllRemoteBranchesWriter(stream gitalypb.RefService_FindAllRemoteBranchesServer, c *catfile.Batch) lines.Sender {
+	ctx := stream.Context()
 	return func(refs [][]byte) error {
 		var branches []*gitalypb.Branch
 
@@ -126,7 +130,7 @@ func newFindAllRemoteBranchesWriter(stream gitalypb.RefService_FindAllRemoteBran
 			if err != nil {
 				return err
 			}
-			branch, err := buildBranch(c, elements)
+			branch, err := buildBranch(ctx, c, elements)
 			if err != nil {
 				return err
 			}
diff --git a/internal/testhelper/testhelper.go b/internal/testhelper/testhelper.go
index af2673d85..805aa0a5f 100644
--- a/internal/testhelper/testhelper.go
+++ b/internal/testhelper/testhelper.go
@@ -443,7 +443,7 @@ func NewTestRepo(t testing.TB) (repo *gitalypb.Repository, repoPath string, clea
 
 // NewTestRepoWithWorktree creates a copy of the test repository with a
 // worktree. This is allows you to run normal 'non-bare' Git commands.
-func NewTestRepoWithWorktree(t *testing.T) (repo *gitalypb.Repository, repoPath string, cleanup func()) {
+func NewTestRepoWithWorktree(t testing.TB) (repo *gitalypb.Repository, repoPath string, cleanup func()) {
 	return cloneTestRepo(t, false)
 }