Merge branch 'jc/add-gitaly-retry-policy' into 'main'

Vasilii Iakliushin · John Cai · GitLab · commit e03830b65844 · 2026-02-17T18:07:10.000+01:00
Add Gitaly retry policy Closes #831 See merge request https://gitlab.com/gitlab-org/gitlab-shell/-/merge_requests/1349 Merged-by: Vasilii Iakliushin <viakliushin@gitlab.com> Approved-by: Fred Reinink <freinink@gitlab.com> Approved-by: Vasilii Iakliushin <viakliushin@gitlab.com> Reviewed-by: Vasilii Iakliushin <viakliushin@gitlab.com> Reviewed-by: Fred Reinink <freinink@gitlab.com> Co-authored-by: John Cai <jcai@gitlab.com>
diff --git a/internal/command/uploadpack/gitalycall_test.go b/internal/command/uploadpack/gitalycall_test.go
@@ -80,3 +80,52 @@ func TestUploadPack(t *testing.T) {
 		})
 	}
 }
+
+func TestUploadPackWithRetryConfig(t *testing.T) {
+	gitalyAddress, testServer := testserver.StartGitalyServer(t, "tcp")
+
+	retryConfig := map[string]interface{}{
+		"maxAttempts":          4,
+		"initialBackoff":       "0.1s",
+		"maxBackoff":           "1s",
+		"backoffMultiplier":    2,
+		"retryableStatusCodes": []string{"UNAVAILABLE"},
+	}
+	requests := requesthandlers.BuildAllowedWithGitalyHandlersAndRetryConfig(t, gitalyAddress, retryConfig)
+	url := testserver.StartHTTPServer(t, requests)
+
+	output := &bytes.Buffer{}
+	input := &bytes.Buffer{}
+
+	repo := "group/repo"
+	env := sshenv.Env{
+		IsSSHConnection: true,
+		OriginalCommand: "git-upload-pack " + repo,
+		RemoteAddr:      "127.0.0.1",
+	}
+
+	args := &commandargs.Shell{
+		GitlabKeyID: "1",
+		CommandType: commandargs.UploadPack,
+		SSHArgs:     []string{"git-upload-pack", repo},
+		Env:         env,
+	}
+
+	ctx := correlation.ContextWithCorrelation(context.Background(), "retry-test")
+	ctx = correlation.ContextWithClientName(ctx, "gitlab-shell-tests")
+
+	cfg := &config.Config{GitlabURL: url}
+	cfg.GitalyClient.InitSidechannelRegistry(ctx)
+
+	cmd := &Command{
+		Config:     cfg,
+		Args:       args,
+		ReadWriter: &readwriter.ReadWriter{ErrOut: output, Out: output, In: input},
+	}
+
+	_, err := cmd.Execute(ctx)
+	require.NoError(t, err)
+
+	require.Equal(t, "SSHUploadPackWithSidechannel: "+repo, output.String())
+	require.Equal(t, "retry-test", testServer.ReceivedMD["x-gitlab-correlation-id"][0])
+}
diff --git a/internal/gitaly/gitaly.go b/internal/gitaly/gitaly.go
@@ -19,17 +19,24 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/metrics"
 )
 
-// Command represents a gRPC service command with its address and token.
-type Command struct {
+// CacheKey contains the fields used for connection caching.
+type CacheKey struct {
 	ServiceName string
 	Address     string
 	Token       string
 }
 
+// Command represents a gRPC service command with its address and token.
+type Command struct {
+	CacheKey
+
+	RetryPolicy *gitalyclient.RetryPolicy
+}
+
 type connectionsCache struct {
 	sync.RWMutex
 
-	connections map[Command]*grpc.ClientConn
+	connections map[CacheKey]*grpc.ClientConn
 }
 
 // Client manages connections to Gitaly services and handles sidechannel communication.
@@ -46,8 +53,10 @@ func (c *Client) InitSidechannelRegistry(ctx context.Context) {
 
 // GetConnection returns a gRPC connection for the given command, using a cached connection if available.
 func (c *Client) GetConnection(ctx context.Context, cmd Command) (*grpc.ClientConn, error) {
+	key := CacheKey{ServiceName: cmd.ServiceName, Address: cmd.Address, Token: cmd.Token}
+
 	c.cache.RLock()
-	existingConn := c.cache.connections[cmd]
+	existingConn := c.cache.connections[key]
 	c.cache.RUnlock()
 
 	if existingConn != nil {
@@ -57,7 +66,7 @@ func (c *Client) GetConnection(ctx context.Context, cmd Command) (*grpc.ClientCo
 	c.cache.Lock()
 	defer c.cache.Unlock()
 
-	if cachedConn := c.cache.connections[cmd]; cachedConn != nil {
+	if cachedConn := c.cache.connections[key]; cachedConn != nil {
 		return cachedConn, nil
 	}
 
@@ -67,10 +76,10 @@ func (c *Client) GetConnection(ctx context.Context, cmd Command) (*grpc.ClientCo
 	}
 
 	if c.cache.connections == nil {
-		c.cache.connections = make(map[Command]*grpc.ClientConn)
+		c.cache.connections = make(map[CacheKey]*grpc.ClientConn)
 	}
 
-	c.cache.connections[cmd] = newConn
+	c.cache.connections[key] = newConn
 
 	return newConn, nil
 }
@@ -134,5 +143,9 @@ func (c *Client) newConnection(ctx context.Context, cmd Command) (conn *grpc.Cli
 		gitalyclient.WithGrpcOptions(grpcOpts),
 	}
 
+	if cmd.RetryPolicy != nil {
+		connOpts = append(connOpts, gitalyclient.WithRetryPolicy(cmd.RetryPolicy))
+	}
+
 	return gitalyclient.DialSidechannel(ctx, cmd.Address, c.SidechannelRegistry, connOpts...)
 }
diff --git a/internal/gitaly/gitaly_test.go b/internal/gitaly/gitaly_test.go
@@ -15,15 +15,15 @@ func TestPrometheusMetrics(t *testing.T) {
 
 	c := newClient()
 
-	cmd := Command{ServiceName: "git-upload-pack", Address: "tcp://localhost:9999"}
+	cmd := Command{CacheKey: CacheKey{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: ""}
+	cmd = Command{CacheKey: CacheKey{Address: ""}}
 	c.newConnection(context.Background(), cmd)
 
 	require.InDelta(t, 2, testutil.ToFloat64(metrics.GitalyConnectionsTotal.WithLabelValues("ok")), 0.1)
@@ -35,7 +35,7 @@ func TestCachedConnections(t *testing.T) {
 
 	require.Empty(t, c.cache.connections)
 
-	cmd := Command{ServiceName: "git-upload-pack", Address: "tcp://localhost:9999"}
+	cmd := Command{CacheKey: CacheKey{ServiceName: "git-upload-pack", Address: "tcp://localhost:9999"}}
 
 	conn, err := c.GetConnection(context.Background(), cmd)
 	require.NoError(t, err)
@@ -46,7 +46,7 @@ func TestCachedConnections(t *testing.T) {
 	require.Len(t, c.cache.connections, 1)
 	require.Equal(t, conn, newConn)
 
-	cmd = Command{ServiceName: "git-upload-pack", Address: "tcp://localhost:9998"}
+	cmd = Command{CacheKey: CacheKey{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)
diff --git a/internal/gitlabnet/accessverifier/client.go b/internal/gitlabnet/accessverifier/client.go
@@ -3,6 +3,7 @@ package accessverifier
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"net/http"
 
@@ -82,7 +83,8 @@ type Response struct {
 	Who              string
 	StatusCode       int
 	// NeedAudit indicates whether git event should be audited to rails.
-	NeedAudit bool `json:"need_audit"`
+	NeedAudit   bool            `json:"need_audit"`
+	RetryConfig json.RawMessage `json:"retry_config,omitempty"`
 }
 
 // NewClient creates a new instance of Client
diff --git a/internal/handler/exec.go b/internal/handler/exec.go
@@ -3,6 +3,7 @@ package handler
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"strconv"
 	"strings"
@@ -11,13 +12,14 @@ import (
 	grpccodes "google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
 	grpcstatus "google.golang.org/grpc/status"
+	"google.golang.org/protobuf/encoding/protojson"
 
+	gitalyclient "gitlab.com/gitlab-org/gitaly/v18/client"
+	pb "gitlab.com/gitlab-org/gitaly/v18/proto/go/gitalypb"
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/config"
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/gitaly"
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/gitlabnet/accessverifier"
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/sshenv"
-
-	pb "gitlab.com/gitlab-org/gitaly/v18/proto/go/gitalypb"
 	"gitlab.com/gitlab-org/labkit/log"
 )
 
@@ -36,14 +38,31 @@ type GitalyCommand struct {
 // NewGitalyCommand creates a new GitalyCommand instance
 func NewGitalyCommand(cfg *config.Config, serviceName string, response *accessverifier.Response) *GitalyCommand {
 	gc := gitaly.Command{
-		ServiceName: serviceName,
-		Address:     response.Gitaly.Address,
-		Token:       response.Gitaly.Token,
+		CacheKey: gitaly.CacheKey{
+			ServiceName: serviceName,
+			Address:     response.Gitaly.Address,
+			Token:       response.Gitaly.Token,
+		},
+		RetryPolicy: parseRetryConfig(response.RetryConfig),
 	}
 
 	return &GitalyCommand{Config: cfg, Response: response, Command: gc}
 }
 
+func parseRetryConfig(rawConfig json.RawMessage) *gitalyclient.RetryPolicy {
+	if len(rawConfig) == 0 {
+		return nil
+	}
+
+	var policy gitalyclient.RetryPolicy
+	if err := protojson.Unmarshal(rawConfig, &policy); err != nil {
+		log.WithError(err).Error("failed to unmarshal retry policy")
+		return nil
+	}
+
+	return &policy
+}
+
 // processGitalyError handles errors that come back from Gitaly that may be a
 // LimitError. A LimitError is returned by Gitaly when it is at its limit in
 // handling requests. Since this is a known error, we should print a sensible
diff --git a/internal/handler/exec_test.go b/internal/handler/exec_test.go
@@ -2,6 +2,7 @@ package handler
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"testing"
 
@@ -242,3 +243,33 @@ func errWithDetail(t *testing.T, detail proto.Message) error {
 
 	return grpcstatus.ErrorProto(proto)
 }
+
+func TestNewGitalyCommandWithRetryConfig(t *testing.T) {
+	retryConfig := json.RawMessage(`{"maxAttempts":4,"initialBackoff":"0.1s","maxBackoff":"1s","backoffMultiplier":2,"retryableStatusCodes":["UNAVAILABLE"]}`)
+
+	cmd := NewGitalyCommand(
+		newConfig(),
+		string(commandargs.UploadPack),
+		&accessverifier.Response{
+			Gitaly:      accessverifier.Gitaly{Address: "tcp://localhost:9999"},
+			RetryConfig: retryConfig,
+		},
+	)
+
+	require.NotNil(t, cmd.Command.RetryPolicy)
+	require.Equal(t, uint32(4), cmd.Command.RetryPolicy.MaxAttempts)
+	require.InEpsilon(t, float32(2), cmd.Command.RetryPolicy.BackoffMultiplier, 0.0001)
+	require.Equal(t, []string{"UNAVAILABLE"}, cmd.Command.RetryPolicy.RetryableStatusCodes)
+}
+
+func TestNewGitalyCommandWithoutRetryConfig(t *testing.T) {
+	cmd := NewGitalyCommand(
+		newConfig(),
+		string(commandargs.UploadPack),
+		&accessverifier.Response{
+			Gitaly: accessverifier.Gitaly{Address: "tcp://localhost:9999"},
+		},
+	)
+
+	require.Nil(t, cmd.Command.RetryPolicy)
+}
diff --git a/internal/testhelper/requesthandlers/requesthandlers.go b/internal/testhelper/requesthandlers/requesthandlers.go
@@ -33,40 +33,47 @@ func BuildDisallowedByAPIHandlers(t *testing.T) []testserver.TestRequestHandler
 
 // BuildAllowedWithGitalyHandlers returns test request handlers for allowed API calls with Gitaly.
 func BuildAllowedWithGitalyHandlers(t *testing.T, gitalyAddress string) []testserver.TestRequestHandler {
-	requests := []testserver.TestRequestHandler{
+	return BuildAllowedWithGitalyHandlersAndRetryConfig(t, gitalyAddress, nil)
+}
+
+// BuildAllowedWithGitalyHandlersAndRetryConfig returns test request handlers for allowed API calls with Gitaly and retry config.
+func BuildAllowedWithGitalyHandlersAndRetryConfig(t *testing.T, gitalyAddress string, retryConfig map[string]interface{}) []testserver.TestRequestHandler {
+	body := map[string]interface{}{
+		"status":      true,
+		"gl_id":       "1",
+		"gl_key_type": "key",
+		"gl_key_id":   123,
+		"gl_username": "alex-doe",
+		"gitaly": map[string]interface{}{
+			"repository": map[string]interface{}{
+				"storage_name":                     "storage_name",
+				"relative_path":                    "relative_path",
+				"git_object_directory":             "path/to/git_object_directory",
+				"git_alternate_object_directories": []string{"path/to/git_alternate_object_directory"},
+				"gl_repository":                    "group/repo",
+				"gl_project_path":                  "group/project-path",
+			},
+			"address": gitalyAddress,
+			"token":   "token",
+			"features": map[string]string{
+				"gitaly-feature-cache_invalidator":        "true",
+				"gitaly-feature-inforef_uploadpack_cache": "false",
+			},
+		},
+	}
+
+	if retryConfig != nil {
+		body["retry_config"] = retryConfig
+	}
+
+	return []testserver.TestRequestHandler{
 		{
 			Path: "/api/v4/internal/allowed",
 			Handler: func(w http.ResponseWriter, _ *http.Request) {
-				body := map[string]interface{}{
-					"status":      true,
-					"gl_id":       "1",
-					"gl_key_type": "key",
-					"gl_key_id":   123,
-					"gl_username": "alex-doe",
-					"gitaly": map[string]interface{}{
-						"repository": map[string]interface{}{
-							"storage_name":                     "storage_name",
-							"relative_path":                    "relative_path",
-							"git_object_directory":             "path/to/git_object_directory",
-							"git_alternate_object_directories": []string{"path/to/git_alternate_object_directory"},
-							"gl_repository":                    "group/repo",
-							"gl_project_path":                  "group/project-path",
-						},
-						"address": gitalyAddress,
-						"token":   "token",
-						"features": map[string]string{
-							"gitaly-feature-cache_invalidator":        "true",
-							"gitaly-feature-inforef_uploadpack_cache": "false",
-							"some-other-ff":                           "true",
-						},
-					},
-				}
 				assert.NoError(t, json.NewEncoder(w).Encode(body))
 			},
 		},
 	}
-
-	return requests
 }
 
 // BuildAllowedWithCustomActionsHandlers returns test request handlers for allowed API calls with custom actions.
