Merge branch 'ps-git-object-retrieval' into 'master'

Feature flag: look up commits without batch-check

See merge request gitlab-org/gitaly!1711

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)
 }