Merge branch 'id-reuse-grpc-connections-and-sidechannel' into 'main'

Reuse Gitaly connections and sidechannel

See merge request gitlab-org/gitlab-shell!575

cmd/gitlab-shell/main.go

@@ -68,6 +68,8 @@ func main() {
 	ctx, finished := command.Setup(executable.Name, config)
 	defer finished()
 
+	config.GitalyClient.InitSidechannelRegistry(ctx)
+
 	cmdName := reflect.TypeOf(cmd).String()
 	ctxlog := log.ContextLogger(ctx)
 	ctxlog.WithFields(log.Fields{"env": env, "command": cmdName}).Info("gitlab-shell: main: executing command")

cmd/gitlab-sshd/main.go

@@ -69,6 +69,8 @@ func main() {
 	ctx, finished := command.Setup("gitlab-sshd", cfg)
 	defer finished()
 
+	cfg.GitalyClient.InitSidechannelRegistry(ctx)
+
 	server, err := sshd.NewServer(cfg)
 	if err != nil {
 		log.WithError(err).Fatal("Failed to start GitLab built-in sshd")

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

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

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

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)
 

internal/config/config.go

@@ -12,6 +12,7 @@ import (
 	yaml "gopkg.in/yaml.v2"
 
 	"gitlab.com/gitlab-org/gitlab-shell/client"
+	"gitlab.com/gitlab-org/gitlab-shell/internal/gitaly"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/metrics"
 )
 
@@ -59,6 +60,8 @@ type Config struct {
 	httpClient     *client.HttpClient
 	httpClientErr  error
 	httpClientOnce sync.Once
+
+	GitalyClient gitaly.Client
 }
 
 // The defaults to apply before parsing the config file(s).

internal/gitaly/gitaly.go

+package gitaly
+
+import (
+	"context"
+	"fmt"
+	"sync"
+
+	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
+	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
+	"google.golang.org/grpc"
+
+	gitalyauth "gitlab.com/gitlab-org/gitaly/v14/auth"
+	"gitlab.com/gitlab-org/gitaly/v14/client"
+	gitalyclient "gitlab.com/gitlab-org/gitaly/v14/client"
+	"gitlab.com/gitlab-org/labkit/correlation"
+	grpccorrelation "gitlab.com/gitlab-org/labkit/correlation/grpc"
+	"gitlab.com/gitlab-org/labkit/log"
+	grpctracing "gitlab.com/gitlab-org/labkit/tracing/grpc"
+
+	"gitlab.com/gitlab-org/gitlab-shell/internal/metrics"
+)
+
+type Command struct {
+	ServiceName     string
+	Address         string
+	Token           string
+	DialSidechannel bool
+}
+
+type connectionsCache struct {
+	sync.RWMutex
+
+	connections map[Command]*grpc.ClientConn
+}
+
+type Client struct {
+	SidechannelRegistry *gitalyclient.SidechannelRegistry
+
+	cache connectionsCache
+}
+
+func (c *Client) InitSidechannelRegistry(ctx context.Context) {
+	c.SidechannelRegistry = gitalyclient.NewSidechannelRegistry(log.ContextLogger(ctx))
+}
+
+func (c *Client) GetConnection(ctx context.Context, cmd Command) (*grpc.ClientConn, error) {
+	c.cache.RLock()
+	conn := c.cache.connections[cmd]
+	c.cache.RUnlock()
+
+	if conn != nil {
+		return conn, nil
+	}
+
+	c.cache.Lock()
+	defer c.cache.Unlock()
+
+	if conn := c.cache.connections[cmd]; conn != nil {
+		return conn, nil
+	}
+
+	conn, err := c.newConnection(ctx, cmd)
+	if err != nil {
+		return nil, err
+	}
+
+	if c.cache.connections == nil {
+		c.cache.connections = make(map[Command]*grpc.ClientConn)
+	}
+
+	c.cache.connections[cmd] = conn
+
+	return conn, nil
+}
+
+func (c *Client) newConnection(ctx context.Context, cmd Command) (conn *grpc.ClientConn, err error) {
+	defer func() {
+		label := "ok"
+		if err != nil {
+			label = "fail"
+		}
+		metrics.GitalyConnectionsTotal.WithLabelValues(label).Inc()
+	}()
+
+	if cmd.Address == "" {
+		return nil, fmt.Errorf("no gitaly_address given")
+	}
+
+	serviceName := correlation.ExtractClientNameFromContext(ctx)
+	if serviceName == "" {
+		serviceName = "gitlab-shell-unknown"
+
+		log.WithContextFields(ctx, log.Fields{"service_name": serviceName}).Warn("No gRPC service name specified, defaulting to gitlab-shell-unknown")
+	}
+
+	serviceName = fmt.Sprintf("%s-%s", serviceName, cmd.ServiceName)
+
+	connOpts := client.DefaultDialOpts
+	connOpts = append(
+		connOpts,
+		grpc.WithStreamInterceptor(
+			grpc_middleware.ChainStreamClient(
+				grpctracing.StreamClientTracingInterceptor(),
+				grpc_prometheus.StreamClientInterceptor,
+				grpccorrelation.StreamClientCorrelationInterceptor(
+					grpccorrelation.WithClientName(serviceName),
+				),
+			),
+		),
+
+		grpc.WithUnaryInterceptor(
+			grpc_middleware.ChainUnaryClient(
+				grpctracing.UnaryClientTracingInterceptor(),
+				grpc_prometheus.UnaryClientInterceptor,
+				grpccorrelation.UnaryClientCorrelationInterceptor(
+					grpccorrelation.WithClientName(serviceName),
+				),
+			),
+		),
+	)
+
+	if cmd.Token != "" {
+		connOpts = append(connOpts,
+			grpc.WithPerRPCCredentials(gitalyauth.RPCCredentialsV2(cmd.Token)),
+		)
+	}
+
+	if cmd.DialSidechannel {
+		return client.DialSidechannel(ctx, cmd.Address, c.SidechannelRegistry, connOpts)
+	}
+
+	return client.DialContext(ctx, cmd.Address, connOpts)
+}

internal/gitaly/gitaly_test.go

+package gitaly
+
+import (
+	"context"
+	"testing"
+
+	"github.com/prometheus/client_golang/prometheus/testutil"
+	"github.com/stretchr/testify/require"
+
+	"gitlab.com/gitlab-org/gitlab-shell/internal/metrics"
+)
+
+func TestPrometheusMetrics(t *testing.T) {
+	metrics.GitalyConnectionsTotal.Reset()
+
+	c := &Client{}
+
+	cmd := Command{ServiceName: "git-upload-pack", Address: "tcp://localhost:9999"}
+	c.newConnection(context.Background(), cmd)
+	c.newConnection(context.Background(), cmd)
+
+	require.Equal(t, 1, testutil.CollectAndCount(metrics.GitalyConnectionsTotal))
+	require.InDelta(t, 2, testutil.ToFloat64(metrics.GitalyConnectionsTotal.WithLabelValues("ok")), 0.1)
+	require.InDelta(t, 0, testutil.ToFloat64(metrics.GitalyConnectionsTotal.WithLabelValues("fail")), 0.1)
+
+	cmd = Command{Address: ""}
+	c.newConnection(context.Background(), cmd)
+
+	require.InDelta(t, 2, testutil.ToFloat64(metrics.GitalyConnectionsTotal.WithLabelValues("ok")), 0.1)
+	require.InDelta(t, 1, testutil.ToFloat64(metrics.GitalyConnectionsTotal.WithLabelValues("fail")), 0.1)
+}
+
+func TestCachedConnections(t *testing.T) {
+	c := &Client{}
+
+	require.Len(t, c.cache.connections, 0)
+
+	cmd := Command{ServiceName: "git-upload-pack", Address: "tcp://localhost:9999"}
+
+	conn, err := c.GetConnection(context.Background(), cmd)
+	require.NoError(t, err)
+	require.Len(t, c.cache.connections, 1)
+
+	newConn, err := c.GetConnection(context.Background(), cmd)
+	require.NoError(t, err)
+	require.Len(t, c.cache.connections, 1)
+	require.Equal(t, conn, newConn)
+
+	cmd = Command{ServiceName: "git-upload-pack", Address: "tcp://localhost:9998"}
+	_, err = c.GetConnection(context.Background(), cmd)
+	require.NoError(t, err)
+	require.Len(t, c.cache.connections, 2)
+}

internal/handler/exec.go

@@ -6,56 +6,57 @@ import (
 	"strconv"
 	"strings"
 
-	grpc_middleware "github.com/grpc-ecosystem/go-grpc-middleware"
-	grpc_prometheus "github.com/grpc-ecosystem/go-grpc-prometheus"
 	"google.golang.org/grpc"
 	grpccodes "google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	grpcstatus "google.golang.org/grpc/status"
 
 	"gitlab.com/gitlab-org/gitlab-shell/internal/config"
+	"gitlab.com/gitlab-org/gitlab-shell/internal/gitaly"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/gitlabnet/accessverifier"
 	"gitlab.com/gitlab-org/gitlab-shell/internal/sshenv"
 
-	gitalyauth "gitlab.com/gitlab-org/gitaly/v14/auth"
-	"gitlab.com/gitlab-org/gitaly/v14/client"
 	pb "gitlab.com/gitlab-org/gitaly/v14/proto/go/gitalypb"
-	"gitlab.com/gitlab-org/labkit/correlation"
-	grpccorrelation "gitlab.com/gitlab-org/labkit/correlation/grpc"
 	"gitlab.com/gitlab-org/labkit/log"
-	grpctracing "gitlab.com/gitlab-org/labkit/tracing/grpc"
 )
 
 // GitalyHandlerFunc implementations are responsible for making
 // an appropriate Gitaly call using the provided client and context
 // and returning an error from the Gitaly call.
-type GitalyHandlerFunc func(ctx context.Context, client *grpc.ClientConn, registry *client.SidechannelRegistry) (int32, error)
+type GitalyHandlerFunc func(ctx context.Context, client *grpc.ClientConn) (int32, error)
 
 type GitalyCommand struct {
-	Config          *config.Config
-	ServiceName     string
-	Address         string
-	Token           string
-	Features        map[string]string
-	DialSidechannel bool
+	Config   *config.Config
+	Response *accessverifier.Response
+	Command  gitaly.Command
+}
+
+func NewGitalyCommand(cfg *config.Config, serviceName string, response *accessverifier.Response) *GitalyCommand {
+	gc := gitaly.Command{
+		ServiceName:     serviceName,
+		Address:         response.Gitaly.Address,
+		Token:           response.Gitaly.Token,
+		DialSidechannel: response.Gitaly.UseSidechannel,
+	}
+
+	return &GitalyCommand{Config: cfg, Response: response, Command: gc}
 }
 
 // RunGitalyCommand provides a bootstrap for Gitaly commands executed
 // through GitLab-Shell. It ensures that logging, tracing and other
 // common concerns are configured before executing the `handler`.
 func (gc *GitalyCommand) RunGitalyCommand(ctx context.Context, handler GitalyHandlerFunc) error {
-	registry := client.NewSidechannelRegistry(log.ContextLogger(ctx))
-	conn, err := getConn(ctx, gc, registry)
+	// We leave the connection open for future reuse
+	conn, err := gc.getConn(ctx)
 	if err != nil {
 		log.ContextLogger(ctx).WithError(fmt.Errorf("RunGitalyCommand: %v", err)).Error("Failed to get connection to execute Git command")
 
 		return err
 	}
-	defer conn.Close()
 
-	childCtx := withOutgoingMetadata(ctx, gc.Features)
+	childCtx := withOutgoingMetadata(ctx, gc.Response.Gitaly.Features)
 	ctxlog := log.ContextLogger(childCtx)
-	exitStatus, err := handler(childCtx, conn, registry)
+	exitStatus, err := handler(childCtx, conn)
 
 	if err != nil {
 		if grpcstatus.Convert(err).Code() == grpccodes.Unavailable {
@@ -72,35 +73,35 @@ func (gc *GitalyCommand) RunGitalyCommand(ctx context.Context, handler GitalyHan
 
 // PrepareContext wraps a given context with a correlation ID and logs the command to
 // be run.
-func (gc *GitalyCommand) PrepareContext(ctx context.Context, repository *pb.Repository, response *accessverifier.Response, env sshenv.Env) (context.Context, context.CancelFunc) {
+func (gc *GitalyCommand) PrepareContext(ctx context.Context, repository *pb.Repository, env sshenv.Env) (context.Context, context.CancelFunc) {
 	ctx, cancel := context.WithCancel(ctx)
-	gc.LogExecution(ctx, repository, response, env)
+	gc.LogExecution(ctx, repository, env)
 
 	md, ok := metadata.FromOutgoingContext(ctx)
 	if !ok {
 		md = metadata.New(nil)
 	}
-	md.Append("key_id", strconv.Itoa(response.KeyId))
-	md.Append("key_type", response.KeyType)
-	md.Append("user_id", response.UserId)
-	md.Append("username", response.Username)
+	md.Append("key_id", strconv.Itoa(gc.Response.KeyId))
+	md.Append("key_type", gc.Response.KeyType)
+	md.Append("user_id", gc.Response.UserId)
+	md.Append("username", gc.Response.Username)
 	md.Append("remote_ip", env.RemoteAddr)
 	ctx = metadata.NewOutgoingContext(ctx, md)
 
 	return ctx, cancel
 }
 
-func (gc *GitalyCommand) LogExecution(ctx context.Context, repository *pb.Repository, response *accessverifier.Response, env sshenv.Env) {
+func (gc *GitalyCommand) LogExecution(ctx context.Context, repository *pb.Repository, env sshenv.Env) {
 	fields := log.Fields{
-		"command":         gc.ServiceName,
+		"command":         gc.Command.ServiceName,
 		"gl_project_path": repository.GlProjectPath,
 		"gl_repository":   repository.GlRepository,
-		"user_id":         response.UserId,
-		"username":        response.Username,
+		"user_id":         gc.Response.UserId,
+		"username":        gc.Response.Username,
 		"git_protocol":    env.GitProtocolVersion,
 		"remote_ip":       env.RemoteAddr,
-		"gl_key_type":     response.KeyType,
-		"gl_key_id":       response.KeyId,
+		"gl_key_type":     gc.Response.KeyType,
+		"gl_key_id":       gc.Response.KeyId,
 	}
 
 	log.WithContextFields(ctx, fields).Info("executing git command")
@@ -118,53 +119,6 @@ func withOutgoingMetadata(ctx context.Context, features map[string]string) conte
 	return metadata.NewOutgoingContext(ctx, md)
 }
 
-func getConn(ctx context.Context, gc *GitalyCommand, registry *client.SidechannelRegistry) (*grpc.ClientConn, error) {
-	if gc.Address == "" {
-		return nil, fmt.Errorf("no gitaly_address given")
-	}
-
-	serviceName := correlation.ExtractClientNameFromContext(ctx)
-	if serviceName == "" {
-		serviceName = "gitlab-shell-unknown"
-
-		log.WithContextFields(ctx, log.Fields{"service_name": serviceName}).Warn("No gRPC service name specified, defaulting to gitlab-shell-unknown")
-	}
-
-	serviceName = fmt.Sprintf("%s-%s", serviceName, gc.ServiceName)
-
-	connOpts := client.DefaultDialOpts
-	connOpts = append(
-		connOpts,
-		grpc.WithStreamInterceptor(
-			grpc_middleware.ChainStreamClient(
-				grpctracing.StreamClientTracingInterceptor(),
-				grpc_prometheus.StreamClientInterceptor,
-				grpccorrelation.StreamClientCorrelationInterceptor(
-					grpccorrelation.WithClientName(serviceName),
-				),
-			),
-		),
-
-		grpc.WithUnaryInterceptor(
-			grpc_middleware.ChainUnaryClient(
-				grpctracing.UnaryClientTracingInterceptor(),
-				grpc_prometheus.UnaryClientInterceptor,
-				grpccorrelation.UnaryClientCorrelationInterceptor(
-					grpccorrelation.WithClientName(serviceName),
-				),
-			),
-		),
-	)
-
-	if gc.Token != "" {
-		connOpts = append(connOpts,
-			grpc.WithPerRPCCredentials(gitalyauth.RPCCredentialsV2(gc.Token)),
-		)
-	}
-
-	if gc.DialSidechannel {
-		return client.DialSidechannel(ctx, gc.Address, registry, connOpts)
-	}
-
-	return client.DialContext(ctx, gc.Address, connOpts)
+func (gc *GitalyCommand) getConn(ctx context.Context) (*grpc.ClientConn, error) {
+	return gc.Config.GitalyClient.GetConnection(ctx, gc.Command)
 }

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)

internal/metrics/metrics.go

@@ -9,9 +9,10 @@ import (
 )
 
 const (
-	namespace     = "gitlab_shell"
-	sshdSubsystem = "sshd"
-	httpSubsystem = "http"
+	namespace       = "gitlab_shell"
+	sshdSubsystem   = "sshd"
+	httpSubsystem   = "http"
+	gitalySubsystem = "gitaly"
 
 	httpInFlightRequestsMetricName       = "in_flight_requests"
 	httpRequestsTotalMetricName          = "requests_total"
@@ -20,6 +21,8 @@ const (
 	sshdConnectionsInFlightName = "in_flight_connections"
 	sshdConnectionDuration      = "connection_duration_seconds"
 	sshdHitMaxSessions          = "concurrent_limited_sessions_total"
+
+	gitalyConnectionsTotalName = "connections_total"
 )
 
 var (
@@ -61,6 +64,16 @@ var (
 		},
 	)
 
+	GitalyConnectionsTotal = promauto.NewCounterVec(
+		prometheus.CounterOpts{
+			Namespace: namespace,
+			Subsystem: gitalySubsystem,
+			Name:      gitalyConnectionsTotalName,
+			Help:      "Number of Gitaly connections that have been established",
+		},
+		[]string{"status"},
+	)
+
 	// The metrics and the buckets size are similar to the ones we have for handlers in Labkit
 	// When the MR: https://gitlab.com/gitlab-org/labkit/-/merge_requests/150 is merged,
 	// these metrics can be refactored out of Gitlab Shell code by using the helper function from Labkit