Implement Repo::HasLocalBranches in shell, also Bench

5 files changed, 73 insertions, 10 deletions

diff --git a/internal/service/repository/repository.go b/internal/service/repository/repository.go
index 726b04f0b..4cb71c2f8 100644
--- a/internal/service/repository/repository.go
+++ b/internal/service/repository/repository.go
@@ -1,13 +1,20 @@
 package repository
 
 import (
+	"io"
+
 	"golang.org/x/net/context"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	pb "gitlab.com/gitlab-org/gitaly-proto/go"
+	"gitlab.com/gitlab-org/gitaly/internal/git"
 	"gitlab.com/gitlab-org/gitaly/internal/helper"
 	"gitlab.com/gitlab-org/gitaly/internal/rubyserver"
 )
 
+var hlbRuby = false
+
 // Deprecated
 func (s *server) Exists(ctx context.Context, in *pb.RepositoryExistsRequest) (*pb.RepositoryExistsResponse, error) {
 	return nil, helper.Unimplemented
@@ -22,7 +29,7 @@ func (s *server) RepositoryExists(ctx context.Context, in *pb.RepositoryExistsRe
 	return &pb.RepositoryExistsResponse{Exists: helper.IsGitDirectory(path)}, nil
 }
 
-func (s *server) HasLocalBranches(ctx context.Context, in *pb.HasLocalBranchesRequest) (*pb.HasLocalBranchesResponse, error) {
+func rubyHasLocalBranches(s *server, ctx context.Context, in *pb.HasLocalBranchesRequest) (*pb.HasLocalBranchesResponse, error) {
 	client, err := s.RepositoryServiceClient(ctx)
 	if err != nil {
 		return nil, err
@@ -35,3 +42,30 @@ func (s *server) HasLocalBranches(ctx context.Context, in *pb.HasLocalBranchesRe
 
 	return client.HasLocalBranches(clientCtx, in)
 }
+
+func (s *server) HasLocalBranches(ctx context.Context, in *pb.HasLocalBranchesRequest) (*pb.HasLocalBranchesResponse, error) {
+	if hlbRuby {
+		return rubyHasLocalBranches(s, ctx, in)
+	}
+
+	args := []string{"for-each-ref", "--count=1", "refs/heads"}
+	cmd, err := git.Command(ctx, in.GetRepository(), args...)
+	if err != nil {
+		if _, ok := status.FromError(err); ok {
+			return nil, err
+		}
+		return nil, status.Errorf(codes.Internal, "HasLocalBranches: gitCommand: %v", err)
+	}
+
+	buf := make([]byte, 1)
+	n, err := cmd.Read(buf)
+	if err != nil && err != io.EOF {
+		return nil, status.Errorf(codes.Internal, "HasLocalBranches: read: %v", err)
+	}
+
+	if err := cmd.Wait(); err != nil {
+		return nil, status.Errorf(codes.Internal, "HasLocalBranches: cmd wait: %v", err)
+	}
+
+	return &pb.HasLocalBranchesResponse{Value: n > 0}, nil
+}
diff --git a/internal/service/repository/repository_test.go b/internal/service/repository/repository_test.go
index 5c03ebd26..5d26f5e3b 100644
--- a/internal/service/repository/repository_test.go
+++ b/internal/service/repository/repository_test.go
@@ -177,6 +177,35 @@ func TestSuccessfulHasLocalBranches(t *testing.T) {
 	}
 }
 
+func BenchmarkHasLocalBranchesShell(b *testing.B) {
+	server, serverSocketPath := runRepoServer(b)
+	defer server.Stop()
+
+	client, conn := newRepositoryClient(b, serverSocketPath)
+	defer conn.Close()
+
+	testRepo, _, cleanupFn := testhelper.NewTestRepo(b)
+	defer cleanupFn()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	for n := 0; n < b.N; n++ {
+		resp, err := client.HasLocalBranches(ctx, &pb.HasLocalBranchesRequest{Repository: testRepo})
+
+		require.NoError(b, err)
+		require.True(b, resp.Value)
+	}
+}
+
+func BenchmarkHasLocalBranchesRuby(b *testing.B) {
+	hlbRuby = true
+	defer func() { hlbRuby = false }()
+
+	BenchmarkHasLocalBranchesShell(b)
+
+}
+
 func TestFailedHasLocalBranches(t *testing.T) {
 	server, serverSocketPath := runRepoServer(t)
 	defer server.Stop()
diff --git a/internal/service/repository/testhelper_test.go b/internal/service/repository/testhelper_test.go
index 7926c015e..a01dc1587 100644
--- a/internal/service/repository/testhelper_test.go
+++ b/internal/service/repository/testhelper_test.go
@@ -28,7 +28,7 @@ var (
 	RubyServer *rubyserver.Server
 )
 
-func newRepositoryClient(t *testing.T, serverSocketPath string) (pb.RepositoryServiceClient, *grpc.ClientConn) {
+func newRepositoryClient(t testing.TB, serverSocketPath string) (pb.RepositoryServiceClient, *grpc.ClientConn) {
 	connOpts := []grpc.DialOption{
 		grpc.WithInsecure(),
 		grpc.WithDialer(func(addr string, timeout time.Duration) (net.Conn, error) {
@@ -46,7 +46,7 @@ func newRepositoryClient(t *testing.T, serverSocketPath string) (pb.RepositorySe
 
 var NewRepositoryClient = newRepositoryClient
 
-func runRepoServer(t *testing.T) (*grpc.Server, string) {
+func runRepoServer(t testing.TB) (*grpc.Server, string) {
 	streamInt := []grpc.StreamServerInterceptor{auth.StreamServerInterceptor()}
 	unaryInt := []grpc.UnaryServerInterceptor{auth.UnaryServerInterceptor()}
 
diff --git a/internal/testhelper/test_hook.go b/internal/testhelper/test_hook.go
index 2e9b9d7a5..cd87665b3 100644
--- a/internal/testhelper/test_hook.go
+++ b/internal/testhelper/test_hook.go
@@ -8,7 +8,7 @@ import (
 )
 
 type testHook struct {
-	t         *testing.T
+	t         testing.TB
 	formatter log.Formatter
 }
 
@@ -43,7 +43,7 @@ func (s testHook) Fire(entry *log.Entry) error {
 }
 
 // NewTestLogger created a logrus hook which can be used with testing logs
-func NewTestLogger(t *testing.T) *log.Logger {
+func NewTestLogger(t testing.TB) *log.Logger {
 	logger := log.New()
 	logger.Out = ioutil.Discard
 	formatter := &log.TextFormatter{}
diff --git a/internal/testhelper/testhelper.go b/internal/testhelper/testhelper.go
index 5b0c4957b..530aaba41 100644
--- a/internal/testhelper/testhelper.go
+++ b/internal/testhelper/testhelper.go
@@ -144,7 +144,7 @@ func AssertGrpcError(t *testing.T, err error, expectedCode codes.Code, containsT
 }
 
 // MustRunCommand runs a command with an optional standard input and returns the standard output, or fails.
-func MustRunCommand(t *testing.T, stdin io.Reader, name string, args ...string) []byte {
+func MustRunCommand(t testing.TB, stdin io.Reader, name string, args ...string) []byte {
 	cmd := exec.Command(name, args...)
 	if stdin != nil {
 		cmd.Stdin = stdin
@@ -269,7 +269,7 @@ func ConfigureRuby() {
 }
 
 // NewTestGrpcServer creates a GRPC Server for testing purposes
-func NewTestGrpcServer(t *testing.T, streamInterceptors []grpc.StreamServerInterceptor, unaryInterceptors []grpc.UnaryServerInterceptor) *grpc.Server {
+func NewTestGrpcServer(t testing.TB, streamInterceptors []grpc.StreamServerInterceptor, unaryInterceptors []grpc.UnaryServerInterceptor) *grpc.Server {
 	logger := NewTestLogger(t)
 	logrusEntry := log.NewEntry(logger).WithField("test", t.Name())
 
@@ -336,7 +336,7 @@ func Context() (context.Context, func()) {
 	return context.WithCancel(context.Background())
 }
 
-func createRepo(t *testing.T, storagePath string) (repo *pb.Repository, repoPath, relativePath string) {
+func createRepo(t testing.TB, storagePath string) (repo *pb.Repository, repoPath, relativePath string) {
 	normalizedPrefix := strings.Replace(t.Name(), "/", "-", -1) //TempDir doesn't like a prefix containing slashes
 
 	repoPath, err := ioutil.TempDir(storagePath, normalizedPrefix)
@@ -377,7 +377,7 @@ func initRepo(t *testing.T, bare bool) (*pb.Repository, string, func()) {
 }
 
 // NewTestRepo creates a bare copy of the test repository.
-func NewTestRepo(t *testing.T) (repo *pb.Repository, repoPath string, cleanup func()) {
+func NewTestRepo(t testing.TB) (repo *pb.Repository, repoPath string, cleanup func()) {
 	return cloneTestRepo(t, true)
 }
 
@@ -387,7 +387,7 @@ func NewTestRepoWithWorktree(t *testing.T) (repo *pb.Repository, repoPath string
 	return cloneTestRepo(t, false)
 }
 
-func cloneTestRepo(t *testing.T, bare bool) (repo *pb.Repository, repoPath string, cleanup func()) {
+func cloneTestRepo(t testing.TB, bare bool) (repo *pb.Repository, repoPath string, cleanup func()) {
 	storagePath := GitlabTestStoragePath()
 	repo, repoPath, relativePath := createRepo(t, storagePath)
 	testRepo := TestRepository()