Merge branch 'refactor-refnames-chunker' into 'master'

Refactor refnames RPC's to use chunker

See merge request gitlab-org/gitaly!1041

9 files changed, 182 insertions, 94 deletions

diff --git a/changelogs/unreleased/refactor-refnames-chunker.yml b/changelogs/unreleased/refactor-refnames-chunker.yml
new file mode 100644
index 000000000..56962b510
--- /dev/null
+++ b/changelogs/unreleased/refactor-refnames-chunker.yml
@@ -0,0 +1,5 @@
+---
+title: Refactor refnames RPC's to use chunker
+merge_request: 1041
+author:
+type: other
diff --git a/internal/helper/chunker/chunker.go b/internal/helper/chunk/chunker.go
index f8bf5dc1b..48d99f03f 100644
--- a/internal/helper/chunker/chunker.go
+++ b/internal/helper/chunk/chunker.go
@@ -1,4 +1,4 @@
-package chunker
+package chunk
 
 // Item could be e.g. a commit in an RPC that returns a chunked stream of
 // commits.
diff --git a/internal/service/commit/list_commits_by_oid.go b/internal/service/commit/list_commits_by_oid.go
index 3d29beacc..6d7361419 100644
--- a/internal/service/commit/list_commits_by_oid.go
+++ b/internal/service/commit/list_commits_by_oid.go
@@ -4,7 +4,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly-proto/go/gitalypb"
 	"gitlab.com/gitlab-org/gitaly/internal/git/catfile"
 	gitlog "gitlab.com/gitlab-org/gitaly/internal/git/log"
-	"gitlab.com/gitlab-org/gitaly/internal/helper/chunker"
+	"gitlab.com/gitlab-org/gitaly/internal/helper/chunk"
 )
 
 func (s *server) ListCommitsByOid(in *gitalypb.ListCommitsByOidRequest, stream gitalypb.CommitService_ListCommitsByOidServer) error {
@@ -15,7 +15,7 @@ func (s *server) ListCommitsByOid(in *gitalypb.ListCommitsByOidRequest, stream g
 		return err
 	}
 
-	sender := chunker.New(&commitsByOidSender{stream: stream})
+	sender := chunk.New(&commitsByOidSender{stream: stream})
 
 	for _, oid := range in.Oid {
 		commit, err := gitlog.GetCommitCatfile(c, oid)
@@ -39,7 +39,7 @@ type commitsByOidSender struct {
 	stream   gitalypb.CommitService_ListCommitsByOidServer
 }
 
-func (c *commitsByOidSender) Append(it chunker.Item) {
+func (c *commitsByOidSender) Append(it chunk.Item) {
 	c.response.Commits = append(c.response.Commits, it.(*gitalypb.GitCommit))
 }
 
diff --git a/internal/service/commit/list_files.go b/internal/service/commit/list_files.go
index f5a100e8c..bab30138c 100644
--- a/internal/service/commit/list_files.go
+++ b/internal/service/commit/list_files.go
@@ -10,7 +10,7 @@ import (
 	"gitlab.com/gitlab-org/gitaly/internal/git"
 	"gitlab.com/gitlab-org/gitaly/internal/git/lstree"
 	"gitlab.com/gitlab-org/gitaly/internal/helper"
-	"gitlab.com/gitlab-org/gitaly/internal/helper/chunker"
+	"gitlab.com/gitlab-org/gitaly/internal/helper/chunk"
 	"google.golang.org/grpc/codes"
 )
 
@@ -52,7 +52,7 @@ func listFiles(repo *gitalypb.Repository, revision string, stream gitalypb.Commi
 		return err
 	}
 
-	sender := chunker.New(&listFilesSender{stream: stream})
+	sender := chunk.New(&listFilesSender{stream: stream})
 
 	for parser := lstree.NewParser(cmd); ; {
 		entry, err := parser.NextEntry()
@@ -82,6 +82,6 @@ type listFilesSender struct {
 
 func (s *listFilesSender) Reset()      { s.response = &gitalypb.ListFilesResponse{} }
 func (s *listFilesSender) Send() error { return s.stream.Send(s.response) }
-func (s *listFilesSender) Append(it chunker.Item) {
+func (s *listFilesSender) Append(it chunk.Item) {
 	s.response.Paths = append(s.response.Paths, it.([]byte))
 }
diff --git a/internal/service/commit/tree_entries.go b/internal/service/commit/tree_entries.go
index f1456a71f..f46bfd910 100644
--- a/internal/service/commit/tree_entries.go
+++ b/internal/service/commit/tree_entries.go
@@ -7,7 +7,7 @@ import (
 	log "github.com/sirupsen/logrus"
 	"gitlab.com/gitlab-org/gitaly-proto/go/gitalypb"
 	"gitlab.com/gitlab-org/gitaly/internal/git/catfile"
-	"gitlab.com/gitlab-org/gitaly/internal/helper/chunker"
+	"gitlab.com/gitlab-org/gitaly/internal/helper/chunk"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 )
@@ -63,7 +63,7 @@ func sendTreeEntries(stream gitalypb.CommitService_GetTreeEntriesServer, c *catf
 		}
 	}
 
-	sender := chunker.New(&treeEntriesSender{stream: stream})
+	sender := chunk.New(&treeEntriesSender{stream: stream})
 	for _, e := range entries {
 		sender.Send(e)
 	}
@@ -76,7 +76,7 @@ type treeEntriesSender struct {
 	stream   gitalypb.CommitService_GetTreeEntriesServer
 }
 
-func (c *treeEntriesSender) Append(it chunker.Item) {
+func (c *treeEntriesSender) Append(it chunk.Item) {
 	c.response.Entries = append(c.response.Entries, it.(*gitalypb.TreeEntry))
 }
 
diff --git a/internal/service/ref/refnames.go b/internal/service/ref/refnames.go
new file mode 100644
index 000000000..6c50c84ef
--- /dev/null
+++ b/internal/service/ref/refnames.go
@@ -0,0 +1,81 @@
+package ref
+
+import (
+	"bufio"
+	"context"
+
+	"gitlab.com/gitlab-org/gitaly-proto/go/gitalypb"
+	"gitlab.com/gitlab-org/gitaly/internal/git"
+	"gitlab.com/gitlab-org/gitaly/internal/helper/chunk"
+)
+
+// FindAllBranchNames creates a stream of ref names for all branches in the given repository
+func (s *server) FindAllBranchNames(in *gitalypb.FindAllBranchNamesRequest, stream gitalypb.RefService_FindAllBranchNamesServer) error {
+	chunker := chunk.New(&findAllBranchNamesSender{stream: stream})
+
+	return listRefNames(stream.Context(), chunker, "refs/heads", in.Repository, nil)
+}
+
+type findAllBranchNamesSender struct {
+	stream      gitalypb.RefService_FindAllBranchNamesServer
+	branchNames [][]byte
+}
+
+func (ts *findAllBranchNamesSender) Reset() { ts.branchNames = nil }
+func (ts *findAllBranchNamesSender) Append(it chunk.Item) {
+	ts.branchNames = append(ts.branchNames, it.([]byte))
+}
+func (ts *findAllBranchNamesSender) Send() error {
+	return ts.stream.Send(&gitalypb.FindAllBranchNamesResponse{Names: ts.branchNames})
+}
+
+// FindAllTagNames creates a stream of ref names for all tags in the given repository
+func (s *server) FindAllTagNames(in *gitalypb.FindAllTagNamesRequest, stream gitalypb.RefService_FindAllTagNamesServer) error {
+	chunker := chunk.New(&findAllTagNamesSender{stream: stream})
+
+	return listRefNames(stream.Context(), chunker, "refs/tags", in.Repository, nil)
+}
+
+type findAllTagNamesSender struct {
+	stream   gitalypb.RefService_FindAllTagNamesServer
+	tagNames [][]byte
+}
+
+func (ts *findAllTagNamesSender) Reset() { ts.tagNames = nil }
+func (ts *findAllTagNamesSender) Append(it chunk.Item) {
+	ts.tagNames = append(ts.tagNames, it.([]byte))
+}
+func (ts *findAllTagNamesSender) Send() error {
+	return ts.stream.Send(&gitalypb.FindAllTagNamesResponse{Names: ts.tagNames})
+}
+
+func listRefNames(ctx context.Context, chunker *chunk.Chunker, prefix string, repo *gitalypb.Repository, extraArgs []string) error {
+	args := []string{
+		"for-each-ref",
+		"--format=%(refname)",
+	}
+	args = append(args, extraArgs...)
+	args = append(args, prefix)
+
+	cmd, err := git.Command(ctx, repo, args...)
+	if err != nil {
+		return err
+	}
+
+	scanner := bufio.NewScanner(cmd)
+	for scanner.Scan() {
+		if err := chunker.Send(scanner.Bytes()); err != nil {
+			return err
+		}
+	}
+
+	if err := cmd.Wait(); err != nil {
+		return err
+	}
+
+	if err := scanner.Err(); err != nil {
+		return err
+	}
+
+	return chunker.Flush()
+}
diff --git a/internal/service/ref/refnames_containing.go b/internal/service/ref/refnames_containing.go
new file mode 100644
index 000000000..0ec9e7554
--- /dev/null
+++ b/internal/service/ref/refnames_containing.go
@@ -0,0 +1,86 @@
+package ref
+
+import (
+	"bytes"
+	"fmt"
+
+	"gitlab.com/gitlab-org/gitaly-proto/go/gitalypb"
+	"gitlab.com/gitlab-org/gitaly/internal/git"
+	"gitlab.com/gitlab-org/gitaly/internal/helper"
+	"gitlab.com/gitlab-org/gitaly/internal/helper/chunk"
+)
+
+// ListBranchNamesContainingCommit returns a maximum of in.GetLimit() Branch names
+// which contain the SHA1 passed as argument
+func (*server) ListBranchNamesContainingCommit(in *gitalypb.ListBranchNamesContainingCommitRequest, stream gitalypb.RefService_ListBranchNamesContainingCommitServer) error {
+	if err := git.ValidateCommitID(in.GetCommitId()); err != nil {
+		return helper.ErrInvalidArgument(err)
+	}
+
+	chunker := chunk.New(&branchNamesContainingCommitSender{stream: stream})
+	ctx := stream.Context()
+	if err := listRefNames(ctx, chunker, "refs/heads", in.Repository, containingArgs(in)); err != nil {
+		return helper.ErrInternal(err)
+	}
+
+	return nil
+}
+
+type containingRequest interface {
+	GetCommitId() string
+	GetLimit() uint32
+}
+
+func containingArgs(req containingRequest) []string {
+	args := []string{fmt.Sprintf("--contains=%s", req.GetCommitId())}
+	if limit := req.GetLimit(); limit != 0 {
+		args = append(args, fmt.Sprintf("--count=%d", limit))
+	}
+	return args
+}
+
+type branchNamesContainingCommitSender struct {
+	stream      gitalypb.RefService_ListBranchNamesContainingCommitServer
+	branchNames [][]byte
+}
+
+func (bs *branchNamesContainingCommitSender) Reset() { bs.branchNames = nil }
+func (bs *branchNamesContainingCommitSender) Append(it chunk.Item) {
+	bs.branchNames = append(bs.branchNames, stripPrefix(it, "refs/heads/"))
+}
+func (bs *branchNamesContainingCommitSender) Send() error {
+	return bs.stream.Send(&gitalypb.ListBranchNamesContainingCommitResponse{BranchNames: bs.branchNames})
+}
+
+// ListTagNamesContainingCommit returns a maximum of in.GetLimit() Tag names
+// which contain the SHA1 passed as argument
+func (*server) ListTagNamesContainingCommit(in *gitalypb.ListTagNamesContainingCommitRequest, stream gitalypb.RefService_ListTagNamesContainingCommitServer) error {
+	if err := git.ValidateCommitID(in.GetCommitId()); err != nil {
+		return helper.ErrInvalidArgument(err)
+	}
+
+	chunker := chunk.New(&tagNamesContainingCommitSender{stream: stream})
+	ctx := stream.Context()
+	if err := listRefNames(ctx, chunker, "refs/tags", in.Repository, containingArgs(in)); err != nil {
+		return helper.ErrInternal(err)
+	}
+
+	return nil
+}
+
+type tagNamesContainingCommitSender struct {
+	stream   gitalypb.RefService_ListTagNamesContainingCommitServer
+	tagNames [][]byte
+}
+
+func (ts *tagNamesContainingCommitSender) Reset() { ts.tagNames = nil }
+func (ts *tagNamesContainingCommitSender) Append(it chunk.Item) {
+	ts.tagNames = append(ts.tagNames, stripPrefix(it, "refs/tags/"))
+}
+func (ts *tagNamesContainingCommitSender) Send() error {
+	return ts.stream.Send(&gitalypb.ListTagNamesContainingCommitResponse{TagNames: ts.tagNames})
+}
+
+func stripPrefix(it chunk.Item, prefix string) []byte {
+	return bytes.TrimPrefix(it.([]byte), []byte(prefix))
+}
diff --git a/internal/service/ref/refs.go b/internal/service/ref/refs.go
index 08334be86..4f5e9c569 100644
--- a/internal/service/ref/refs.go
+++ b/internal/service/ref/refs.go
@@ -52,16 +52,6 @@ func findRefs(ctx context.Context, writer lines.Sender, repo *gitalypb.Repositor
 	return cmd.Wait()
 }
 
-// FindAllBranchNames creates a stream of ref names for all branches in the given repository
-func (s *server) FindAllBranchNames(in *gitalypb.FindAllBranchNamesRequest, stream gitalypb.RefService_FindAllBranchNamesServer) error {
-	return findRefs(stream.Context(), newFindAllBranchNamesWriter(stream), in.Repository, []string{"refs/heads"}, &findRefsOpts{})
-}
-
-// FindAllTagNames creates a stream of ref names for all tags in the given repository
-func (s *server) FindAllTagNames(in *gitalypb.FindAllTagNamesRequest, stream gitalypb.RefService_FindAllTagNamesServer) error {
-	return findRefs(stream.Context(), newFindAllTagNamesWriter(stream), in.Repository, []string{"refs/tags"}, &findRefsOpts{})
-}
-
 func (s *server) FindAllTags(in *gitalypb.FindAllTagsRequest, stream gitalypb.RefService_FindAllTagsServer) error {
 	ctx := stream.Context()
 
@@ -280,65 +270,3 @@ func findAllBranches(in *gitalypb.FindAllBranchesRequest, stream gitalypb.RefSer
 
 	return findRefs(ctx, writer, in.Repository, patterns, opts)
 }
-
-// ListBranchNamesContainingCommit returns a maximum of in.GetLimit() Branch names
-// which contain the SHA1 passed as argument
-func (*server) ListBranchNamesContainingCommit(in *gitalypb.ListBranchNamesContainingCommitRequest, stream gitalypb.RefService_ListBranchNamesContainingCommitServer) error {
-	if err := git.ValidateCommitID(in.GetCommitId()); err != nil {
-		return helper.ErrInvalidArgument(err)
-	}
-
-	if err := listBranchNamesContainingCommit(in, stream); err != nil {
-		return helper.ErrInternal(err)
-	}
-
-	return nil
-}
-
-func listBranchNamesContainingCommit(in *gitalypb.ListBranchNamesContainingCommitRequest, stream gitalypb.RefService_ListBranchNamesContainingCommitServer) error {
-	args := []string{fmt.Sprintf("--contains=%s", in.GetCommitId()), "--format=%(refname:strip=2)"}
-	if in.GetLimit() != 0 {
-		args = append(args, fmt.Sprintf("--count=%d", in.GetLimit()))
-	}
-
-	writer := func(refs [][]byte) error {
-		return stream.Send(&gitalypb.ListBranchNamesContainingCommitResponse{BranchNames: refs})
-	}
-
-	return findRefs(stream.Context(), writer, in.Repository, []string{"refs/heads"},
-		&findRefsOpts{
-			cmdArgs: args,
-			delim:   []byte("\n"),
-		})
-}
-
-// ListTagNamesContainingCommit returns a maximum of in.GetLimit() Tag names
-// which contain the SHA1 passed as argument
-func (*server) ListTagNamesContainingCommit(in *gitalypb.ListTagNamesContainingCommitRequest, stream gitalypb.RefService_ListTagNamesContainingCommitServer) error {
-	if err := git.ValidateCommitID(in.GetCommitId()); err != nil {
-		return helper.ErrInvalidArgument(err)
-	}
-
-	if err := listTagNamesContainingCommit(in, stream); err != nil {
-		return helper.ErrInternal(err)
-	}
-
-	return nil
-}
-
-func listTagNamesContainingCommit(in *gitalypb.ListTagNamesContainingCommitRequest, stream gitalypb.RefService_ListTagNamesContainingCommitServer) error {
-	args := []string{fmt.Sprintf("--contains=%s", in.GetCommitId()), "--format=%(refname:strip=2)"}
-	if in.GetLimit() != 0 {
-		args = append(args, fmt.Sprintf("--count=%d", in.GetLimit()))
-	}
-
-	writer := func(refs [][]byte) error {
-		return stream.Send(&gitalypb.ListTagNamesContainingCommitResponse{TagNames: refs})
-	}
-
-	return findRefs(stream.Context(), writer, in.Repository, []string{"refs/tags"},
-		&findRefsOpts{
-			cmdArgs: args,
-			delim:   []byte("\n"),
-		})
-}
diff --git a/internal/service/ref/util.go b/internal/service/ref/util.go
index 978d3b914..b2f0161a6 100644
--- a/internal/service/ref/util.go
+++ b/internal/service/ref/util.go
@@ -75,18 +75,6 @@ func buildBranch(c *catfile.Batch, elements [][]byte) (*gitalypb.Branch, error)
 	}, nil
 }
 
-func newFindAllBranchNamesWriter(stream gitalypb.RefService_FindAllBranchNamesServer) lines.Sender {
-	return func(refs [][]byte) error {
-		return stream.Send(&gitalypb.FindAllBranchNamesResponse{Names: refs})
-	}
-}
-
-func newFindAllTagNamesWriter(stream gitalypb.RefService_FindAllTagNamesServer) lines.Sender {
-	return func(refs [][]byte) error {
-		return stream.Send(&gitalypb.FindAllTagNamesResponse{Names: refs})
-	}
-}
-
 func newFindLocalBranchesWriter(stream gitalypb.RefService_FindLocalBranchesServer, c *catfile.Batch) lines.Sender {
 	return func(refs [][]byte) error {
 		var branches []*gitalypb.FindLocalBranchResponse