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.

1 files changed, 71 insertions, 31 deletions

diff --git a/internal/handler/exec_test.go b/internal/handler/exec_test.go
index 784cf19..8f1d5b2 100644
--- a/internal/handler/exec_test.go
+++ b/internal/handler/exec_test.go
@@ -11,15 +11,15 @@ import (
 	"google.golang.org/grpc/metadata"
 	grpcstatus "google.golang.org/grpc/status"
 
-	"gitlab.com/gitlab-org/gitaly/v14/client"
 	pb "gitlab.com/gitlab-org/gitaly/v14/proto/go/gitalypb"
+	"gitlab.com/gitlab-org/gitlab-shell/internal/command/commandargs"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet/accessverifier"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/sshenv"
 )
 
-func makeHandler(t *testing.T, err error) func(context.Context, *grpc.ClientConn, *client.SidechannelRegistry) (int32, error) {
-	return func(ctx context.Context, client *grpc.ClientConn, registry *client.SidechannelRegistry) (int32, error) {
+func makeHandler(t *testing.T, err error) func(context.Context, *grpc.ClientConn) (int32, error) {
+	return func(ctx context.Context, client *grpc.ClientConn) (int32, error) {
 		require.NotNil(t, ctx)
 		require.NotNil(t, client)
 
@@ -28,10 +28,13 @@ func makeHandler(t *testing.T, err error) func(context.Context, *grpc.ClientConn
 }
 
 func TestRunGitalyCommand(t *testing.T) {
-	cmd := GitalyCommand{
-		Config:  &config.Config{},
-		Address: "tcp://localhost:9999",
-	}
+	cmd := NewGitalyCommand(
+		&config.Config{},
+		string(commandargs.UploadPack),
+		&accessverifier.Response{
+			Gitaly: accessverifier.Gitaly{Address: "tcp://localhost:9999"},
+		},
+	)
 
 	err := cmd.RunGitalyCommand(context.Background(), makeHandler(t, nil))
 	require.NoError(t, err)
@@ -41,6 +44,32 @@ func TestRunGitalyCommand(t *testing.T) {
 	require.Equal(t, err, expectedErr)
 }
 
+func TestCachingOfGitalyConnections(t *testing.T) {
+	ctx := context.Background()
+	cfg := &config.Config{}
+	cfg.GitalyClient.InitSidechannelRegistry(ctx)
+	response := &accessverifier.Response{
+		Username: "user",
+		Gitaly: accessverifier.Gitaly{
+			Address:        "tcp://localhost:9999",
+			Token:          "token",
+			UseSidechannel: true,
+		},
+	}
+
+	cmd := NewGitalyCommand(cfg, string(commandargs.UploadPack), response)
+
+	conn, err := cmd.getConn(ctx)
+	require.NoError(t, err)
+
+	// Reuses connection for different users
+	response.Username = "another-user"
+	cmd = NewGitalyCommand(cfg, string(commandargs.UploadPack), response)
+	newConn, err := cmd.getConn(ctx)
+	require.NoError(t, err)
+	require.Equal(t, conn, newConn)
+}
+
 func TestMissingGitalyAddress(t *testing.T) {
 	cmd := GitalyCommand{Config: &config.Config{}}
 
@@ -49,10 +78,13 @@ func TestMissingGitalyAddress(t *testing.T) {
 }
 
 func TestUnavailableGitalyErr(t *testing.T) {
-	cmd := GitalyCommand{
-		Config:  &config.Config{},
-		Address: "tcp://localhost:9999",
-	}
+	cmd := NewGitalyCommand(
+		&config.Config{},
+		string(commandargs.UploadPack),
+		&accessverifier.Response{
+			Gitaly: accessverifier.Gitaly{Address: "tcp://localhost:9999"},
+		},
+	)
 
 	expectedErr := grpcstatus.Error(grpccodes.Unavailable, "error")
 	err := cmd.RunGitalyCommand(context.Background(), makeHandler(t, expectedErr))
@@ -68,15 +100,20 @@ func TestRunGitalyCommandMetadata(t *testing.T) {
 	}{
 		{
 			name: "gitaly_feature_flags",
-			gc: &GitalyCommand{
-				Config:  &config.Config{},
-				Address: "tcp://localhost:9999",
-				Features: map[string]string{
-					"gitaly-feature-cache_invalidator":        "true",
-					"other-ff":                                "true",
-					"gitaly-feature-inforef_uploadpack_cache": "false",
+			gc: NewGitalyCommand(
+				&config.Config{},
+				string(commandargs.UploadPack),
+				&accessverifier.Response{
+					Gitaly: accessverifier.Gitaly{
+						Address: "tcp://localhost:9999",
+						Features: map[string]string{
+							"gitaly-feature-cache_invalidator":        "true",
+							"other-ff":                                "true",
+							"gitaly-feature-inforef_uploadpack_cache": "false",
+						},
+					},
 				},
-			},
+			),
 			want: map[string]string{
 				"gitaly-feature-cache_invalidator":        "true",
 				"gitaly-feature-inforef_uploadpack_cache": "false",
@@ -87,7 +124,7 @@ func TestRunGitalyCommandMetadata(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			cmd := tt.gc
 
-			err := cmd.RunGitalyCommand(context.Background(), func(ctx context.Context, _ *grpc.ClientConn, _ *client.SidechannelRegistry) (int32, error) {
+			err := cmd.RunGitalyCommand(context.Background(), func(ctx context.Context, _ *grpc.ClientConn) (int32, error) {
 				md, exists := metadata.FromOutgoingContext(ctx)
 				require.True(t, exists)
 				require.Equal(t, len(tt.want), md.Len())
@@ -117,10 +154,19 @@ func TestPrepareContext(t *testing.T) {
 	}{
 		{
 			name: "client_identity",
-			gc: &GitalyCommand{
-				Config:  &config.Config{},
-				Address: "tcp://localhost:9999",
-			},
+			gc: NewGitalyCommand(
+				&config.Config{},
+				string(commandargs.UploadPack),
+				&accessverifier.Response{
+					KeyId:    1,
+					KeyType:  "key",
+					UserId:   "6",
+					Username: "jane.doe",
+					Gitaly: accessverifier.Gitaly{
+						Address: "tcp://localhost:9999",
+					},
+				},
+			),
 			env: sshenv.Env{
 				GitProtocolVersion: "protocol",
 				IsSSHConnection:    true,
@@ -134,12 +180,6 @@ func TestPrepareContext(t *testing.T) {
 				GlRepository:                  "project-26",
 				GlProjectPath:                 "group/private",
 			},
-			response: &accessverifier.Response{
-				KeyId:    1,
-				KeyType:  "key",
-				UserId:   "6",
-				Username: "jane.doe",
-			},
 			want: map[string]string{
 				"key_id":    "1",
 				"key_type":  "key",
@@ -153,7 +193,7 @@ func TestPrepareContext(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			ctx := context.Background()
 
-			ctx, cancel := tt.gc.PrepareContext(ctx, tt.repo, tt.response, tt.env)
+			ctx, cancel := tt.gc.PrepareContext(ctx, tt.repo, tt.env)
 			defer cancel()
 
 			md, exists := metadata.FromOutgoingContext(ctx)