Reuse Gitaly conns and Sidechannel

When gitlab-sshd has been introduced we've started running our
own SSH server. In this case we're able to cache and reuse
Gitaly connections and Registry.

It helps to reduce memory usage.

4 files changed, 19 insertions, 34 deletions

diff --git a/internal/command/receivepack/gitalycall.go b/internal/command/receivepack/gitalycall.go
index 6bd9f49..c2dd2de 100644
--- a/internal/command/receivepack/gitalycall.go
+++ b/internal/command/receivepack/gitalycall.go
@@ -13,13 +13,7 @@ import (
 )
 
 func (c *Command) performGitalyCall(ctx context.Context, response *accessverifier.Response) error {
-	gc := &handler.GitalyCommand{
-		Config:      c.Config,
-		ServiceName: string(commandargs.ReceivePack),
-		Address:     response.Gitaly.Address,
-		Token:       response.Gitaly.Token,
-		Features:    response.Gitaly.Features,
-	}
+	gc := handler.NewGitalyCommand(c.Config, string(commandargs.ReceivePack), response)
 
 	request := &pb.SSHReceivePackRequest{
 		Repository:       &response.Gitaly.Repo,
@@ -30,8 +24,8 @@ func (c *Command) performGitalyCall(ctx context.Context, response *accessverifie
 		GitConfigOptions: response.GitConfigOptions,
 	}
 
-	return gc.RunGitalyCommand(ctx, func(ctx context.Context, conn *grpc.ClientConn, registry *client.SidechannelRegistry) (int32, error) {
-		ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, c.Args.Env)
+	return gc.RunGitalyCommand(ctx, func(ctx context.Context, conn *grpc.ClientConn) (int32, error) {
+		ctx, cancel := gc.PrepareContext(ctx, request.Repository, c.Args.Env)
 		defer cancel()
 
 		rw := c.ReadWriter
diff --git a/internal/command/uploadarchive/gitalycall.go b/internal/command/uploadarchive/gitalycall.go
index 33f0b71..6051d2c 100644
--- a/internal/command/uploadarchive/gitalycall.go
+++ b/internal/command/uploadarchive/gitalycall.go
@@ -13,18 +13,12 @@ import (
 )
 
 func (c *Command) performGitalyCall(ctx context.Context, response *accessverifier.Response) error {
-	gc := &handler.GitalyCommand{
-		Config:      c.Config,
-		ServiceName: string(commandargs.UploadArchive),
-		Address:     response.Gitaly.Address,
-		Token:       response.Gitaly.Token,
-		Features:    response.Gitaly.Features,
-	}
+	gc := handler.NewGitalyCommand(c.Config, string(commandargs.UploadArchive), response)
 
 	request := &pb.SSHUploadArchiveRequest{Repository: &response.Gitaly.Repo}
 
-	return gc.RunGitalyCommand(ctx, func(ctx context.Context, conn *grpc.ClientConn, registry *client.SidechannelRegistry) (int32, error) {
-		ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, c.Args.Env)
+	return gc.RunGitalyCommand(ctx, func(ctx context.Context, conn *grpc.ClientConn) (int32, error) {
+		ctx, cancel := gc.PrepareContext(ctx, request.Repository, c.Args.Env)
 		defer cancel()
 
 		rw := c.ReadWriter
diff --git a/internal/command/uploadpack/gitalycall.go b/internal/command/uploadpack/gitalycall.go
index 43ae4f8..96dd823 100644
--- a/internal/command/uploadpack/gitalycall.go
+++ b/internal/command/uploadpack/gitalycall.go
@@ -13,26 +13,20 @@ import (
 )
 
 func (c *Command) performGitalyCall(ctx context.Context, response *accessverifier.Response) error {
-	gc := &handler.GitalyCommand{
-		Config:      c.Config,
-		ServiceName: string(commandargs.UploadPack),
-		Address:     response.Gitaly.Address,
-		Token:       response.Gitaly.Token,
-		Features:    response.Gitaly.Features,
-	}
+	gc := handler.NewGitalyCommand(c.Config, string(commandargs.UploadPack), response)
 
 	if response.Gitaly.UseSidechannel {
-		gc.DialSidechannel = true
 		request := &pb.SSHUploadPackWithSidechannelRequest{
 			Repository:       &response.Gitaly.Repo,
 			GitProtocol:      c.Args.Env.GitProtocolVersion,
 			GitConfigOptions: response.GitConfigOptions,
 		}
 
-		return gc.RunGitalyCommand(ctx, func(ctx context.Context, conn *grpc.ClientConn, registry *client.SidechannelRegistry) (int32, error) {
-			ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, c.Args.Env)
+		return gc.RunGitalyCommand(ctx, func(ctx context.Context, conn *grpc.ClientConn) (int32, error) {
+			ctx, cancel := gc.PrepareContext(ctx, request.Repository, c.Args.Env)
 			defer cancel()
 
+			registry := c.Config.GitalyClient.SidechannelRegistry
 			rw := c.ReadWriter
 			return client.UploadPackWithSidechannel(ctx, conn, registry, rw.In, rw.Out, rw.ErrOut, request)
 		})
@@ -44,8 +38,8 @@ func (c *Command) performGitalyCall(ctx context.Context, response *accessverifie
 		GitConfigOptions: response.GitConfigOptions,
 	}
 
-	return gc.RunGitalyCommand(ctx, func(ctx context.Context, conn *grpc.ClientConn, registry *client.SidechannelRegistry) (int32, error) {
-		ctx, cancel := gc.PrepareContext(ctx, request.Repository, response, c.Args.Env)
+	return gc.RunGitalyCommand(ctx, func(ctx context.Context, conn *grpc.ClientConn) (int32, error) {
+		ctx, cancel := gc.PrepareContext(ctx, request.Repository, c.Args.Env)
 		defer cancel()
 
 		rw := c.ReadWriter
diff --git a/internal/command/uploadpack/gitalycall_test.go b/internal/command/uploadpack/gitalycall_test.go
index 213b6f9..e0a15ee 100644
--- a/internal/command/uploadpack/gitalycall_test.go
+++ b/internal/command/uploadpack/gitalycall_test.go
@@ -97,15 +97,18 @@ func TestUploadPack_withSidechannel(t *testing.T) {
 		Env:         env,
 	}
 
+	ctx := correlation.ContextWithCorrelation(context.Background(), "a-correlation-id")
+	ctx = correlation.ContextWithClientName(ctx, "gitlab-shell-tests")
+
+	cfg := &config.Config{GitlabUrl: url}
+	cfg.GitalyClient.InitSidechannelRegistry(ctx)
+
 	cmd := &Command{
-		Config:     &config.Config{GitlabUrl: url},
+		Config:     cfg,
 		Args:       args,
 		ReadWriter: &readwriter.ReadWriter{ErrOut: output, Out: output, In: input},
 	}
 
-	ctx := correlation.ContextWithCorrelation(context.Background(), "a-correlation-id")
-	ctx = correlation.ContextWithClientName(ctx, "gitlab-shell-tests")
-
 	err := cmd.Execute(ctx)
 	require.NoError(t, err)