Merge branch 'feat/retry-on-error' into 'main'

feat: retry on http error

Closes #604

See merge request https://gitlab.com/gitlab-org/gitlab-shell/-/merge_requests/703

Merged-by: Ash McKenzie <amckenzie@gitlab.com>
Approved-by: Alejandro Rodríguez <alejandro@gitlab.com>
Approved-by: Ash McKenzie <amckenzie@gitlab.com>
Reviewed-by: Steve Azzopardi <sazzopardi@gitlab.com>
Reviewed-by: Ash McKenzie <amckenzie@gitlab.com>
Co-authored-by: Steve Azzopardi <sazzopardi@gitlab.com>

11 files changed, 217 insertions, 43 deletions

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index f06353e..2b43481 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -64,6 +64,13 @@ tests:
     - make coverage
   coverage: '/\d+.\d+%/'
 
+tests-integration-retryableHttp:
+  extends: .test
+  variables:
+    FF_GITLAB_SHELL_RETRYABLE_HTTP: '1'
+  script:
+    - make test_ruby
+
 race:
   extends: .test
   script:
diff --git a/client/client_test.go b/client/client_test.go
index aefff33..9bea4fd 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -6,6 +6,8 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"net/http/httptest"
+	"os"
 	"path"
 	"strings"
 	"testing"
@@ -18,9 +20,7 @@ import (
 	"gitlab.com/gitlab-org/gitlab-shell/v14/internal/testhelper"
 )
 
-var (
-	secret = "sssh, it's a secret"
-)
+var secret = "sssh, it's a secret"
 
 func TestClients(t *testing.T) {
 	testhelper.PrepareTestRootDir(t)
@@ -70,6 +70,11 @@ func TestClients(t *testing.T) {
 			},
 			secret: "\n" + secret + "\n",
 		},
+		{
+			desc:   "Retry client",
+			server: testserver.StartRetryHttpServer,
+			secret: secret,
+		},
 	}
 
 	for _, tc := range testCases {
@@ -297,3 +302,47 @@ func buildRequests(t *testing.T, relativeURLRoot string) []testserver.TestReques
 
 	return requests
 }
+
+func TestRetryableHTTPFeatureToggle(t *testing.T) {
+	t.Run("retryable http off", func(t *testing.T) {
+		os.Setenv("FF_GITLAB_SHELL_RETRYABLE_HTTP", "0")
+		reqAttempts := 0
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			reqAttempts++
+			w.WriteHeader(500)
+		}))
+		defer srv.Close()
+
+		httpClient, err := NewHTTPClientWithOpts(srv.URL, "/", "", "", 1, nil)
+		require.NoError(t, err)
+		require.NotNil(t, httpClient.HTTPClient)
+		require.Nil(t, httpClient.RetryableHTTP)
+		client, err := NewGitlabNetClient("", "", "", httpClient)
+		require.NoError(t, err)
+
+		_, err = client.Get(context.Background(), "/")
+		require.EqualError(t, err, "Internal API error (500)")
+		require.Equal(t, 1, reqAttempts)
+	})
+
+	t.Run("retryable http on", func(t *testing.T) {
+		os.Setenv("FF_GITLAB_SHELL_RETRYABLE_HTTP", "1")
+		reqAttempts := 0
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			reqAttempts++
+			w.WriteHeader(500)
+		}))
+		defer srv.Close()
+
+		httpClient, err := NewHTTPClientWithOpts(srv.URL, "/", "", "", 1, nil)
+		require.NoError(t, err)
+		require.Nil(t, httpClient.HTTPClient)
+		require.NotNil(t, httpClient.RetryableHTTP)
+		client, err := NewGitlabNetClient("", "", "", httpClient)
+		require.NoError(t, err)
+
+		_, err = client.Get(context.Background(), "/")
+		require.EqualError(t, err, "Internal API unreachable")
+		require.Equal(t, 3, reqAttempts)
+	})
+}
diff --git a/client/gitlabnet.go b/client/gitlabnet.go
index 24c1d5f..38adf2a 100644
--- a/client/gitlabnet.go
+++ b/client/gitlabnet.go
@@ -7,10 +7,12 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"os"
 	"strings"
 	"time"
 
 	"github.com/golang-jwt/jwt/v4"
+	"github.com/hashicorp/go-retryablehttp"
 
 	"gitlab.com/gitlab-org/labkit/log"
 )
@@ -53,7 +55,6 @@ func NewGitlabNetClient(
 	secret string,
 	httpClient *HttpClient,
 ) (*GitlabNetClient, error) {
-
 	if httpClient == nil {
 		return nil, fmt.Errorf("Unsupported protocol")
 	}
@@ -107,6 +108,25 @@ func newRequest(ctx context.Context, method, host, path string, data interface{}
 	return request, nil
 }
 
+func newRetryableRequest(ctx context.Context, method, host, path string, data interface{}) (*retryablehttp.Request, error) {
+	var jsonReader io.Reader
+	if data != nil {
+		jsonData, err := json.Marshal(data)
+		if err != nil {
+			return nil, err
+		}
+
+		jsonReader = bytes.NewReader(jsonData)
+	}
+
+	request, err := retryablehttp.NewRequestWithContext(ctx, method, appendPath(host, path), jsonReader)
+	if err != nil {
+		return nil, err
+	}
+
+	return request, nil
+}
+
 func parseError(resp *http.Response) error {
 	if resp.StatusCode >= 200 && resp.StatusCode <= 399 {
 		return nil
@@ -136,9 +156,15 @@ func (c *GitlabNetClient) DoRequest(ctx context.Context, method, path string, da
 		return nil, err
 	}
 
+	retryableRequest, err := newRetryableRequest(ctx, method, c.httpClient.Host, path, data)
+	if err != nil {
+		return nil, err
+	}
+
 	user, password := c.user, c.password
 	if user != "" && password != "" {
 		request.SetBasicAuth(user, password)
+		retryableRequest.SetBasicAuth(user, password)
 	}
 
 	claims := jwt.RegisteredClaims{
@@ -152,18 +178,31 @@ func (c *GitlabNetClient) DoRequest(ctx context.Context, method, path string, da
 		return nil, err
 	}
 	request.Header.Set(apiSecretHeaderName, tokenString)
+	retryableRequest.Header.Set(apiSecretHeaderName, tokenString)
 
 	originalRemoteIP, ok := ctx.Value(OriginalRemoteIPContextKey{}).(string)
 	if ok {
 		request.Header.Add("X-Forwarded-For", originalRemoteIP)
+		retryableRequest.Header.Add("X-Forwarded-For", originalRemoteIP)
 	}
 
 	request.Header.Add("Content-Type", "application/json")
+	retryableRequest.Header.Add("Content-Type", "application/json")
 	request.Header.Add("User-Agent", c.userAgent)
+	retryableRequest.Header.Add("User-Agent", c.userAgent)
 	request.Close = true
+	retryableRequest.Close = true
 
 	start := time.Now()
-	response, err := c.httpClient.Do(request)
+
+	var response *http.Response
+	var respErr error
+	if c.httpClient.HTTPClient != nil {
+		response, respErr = c.httpClient.HTTPClient.Do(request)
+	}
+	if os.Getenv("FF_GITLAB_SHELL_RETRYABLE_HTTP") == "1" && c.httpClient.RetryableHTTP != nil {
+		response, respErr = c.httpClient.RetryableHTTP.Do(retryableRequest)
+	}
 	fields := log.Fields{
 		"method":      method,
 		"url":         request.URL.String(),
@@ -171,8 +210,8 @@ func (c *GitlabNetClient) DoRequest(ctx context.Context, method, path string, da
 	}
 	logger := log.WithContextFields(ctx, fields)
 
-	if err != nil {
-		logger.WithError(err).Error("Internal API unreachable")
+	if respErr != nil {
+		logger.WithError(respErr).Error("Internal API unreachable")
 		return nil, &ApiError{"Internal API unreachable"}
 	}
 
diff --git a/client/httpclient.go b/client/httpclient.go
index bd00b6b..fe6e2e5 100644
--- a/client/httpclient.go
+++ b/client/httpclient.go
@@ -13,6 +13,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/hashicorp/go-retryablehttp"
 	"gitlab.com/gitlab-org/labkit/correlation"
 	"gitlab.com/gitlab-org/labkit/tracing"
 )
@@ -25,13 +26,12 @@ const (
 	defaultReadTimeoutSeconds = 300
 )
 
-var (
-	ErrCafileNotFound = errors.New("cafile not found")
-)
+var ErrCafileNotFound = errors.New("cafile not found")
 
 type HttpClient struct {
-	*http.Client
-	Host string
+	HTTPClient    *http.Client
+	RetryableHTTP *retryablehttp.Client
+	Host          string
 }
 
 type httpClientCfg struct {
@@ -106,7 +106,22 @@ func NewHTTPClientWithOpts(gitlabURL, gitlabRelativeURLRoot, caFile, caPath stri
 		Timeout:   readTimeout(readTimeoutSeconds),
 	}
 
-	client := &HttpClient{Client: c, Host: host}
+	client := &HttpClient{HTTPClient: c, Host: host}
+
+	if os.Getenv("FF_GITLAB_SHELL_RETRYABLE_HTTP") == "1" {
+		c := retryablehttp.NewClient()
+		c.RetryMax = 2
+		c.RetryWaitMax = 15 * time.Second
+		c.Logger = nil
+		c.HTTPClient.Transport = correlation.NewInstrumentedRoundTripper(tracing.NewRoundTripper(transport))
+		c.HTTPClient.Timeout = readTimeout(readTimeoutSeconds)
+
+		client = &HttpClient{RetryableHTTP: c, Host: host}
+	}
+
+	if client.HTTPClient == nil && client.RetryableHTTP == nil {
+		panic("client/httpclient.go did not set http client")
+	}
 
 	return client, nil
 }
@@ -132,7 +147,6 @@ func buildSocketTransport(gitlabURL, gitlabRelativeURLRoot string) (*http.Transp
 
 func buildHttpsTransport(hcc httpClientCfg, gitlabURL string) (*http.Transport, string, error) {
 	certPool, err := x509.SystemCertPool()
-
 	if err != nil {
 		certPool = x509.NewCertPool()
 	}
diff --git a/client/httpclient_test.go b/client/httpclient_test.go
index ea05a88..ee2c6fd 100644
--- a/client/httpclient_test.go
+++ b/client/httpclient_test.go
@@ -21,7 +21,13 @@ func TestReadTimeout(t *testing.T) {
 	require.NoError(t, err)
 
 	require.NotNil(t, client)
-	require.Equal(t, time.Duration(expectedSeconds)*time.Second, client.Client.Timeout)
+	if client.HTTPClient != nil {
+		require.Equal(t, time.Duration(expectedSeconds)*time.Second, client.HTTPClient.Timeout)
+	}
+
+	if client.RetryableHTTP != nil {
+		require.Equal(t, time.Duration(expectedSeconds)*time.Second, client.RetryableHTTP.HTTPClient.Timeout)
+	}
 }
 
 const (
@@ -117,7 +123,6 @@ func TestRequestWithUserAgent(t *testing.T) {
 	client.SetUserAgent(gitalyUserAgent)
 	_, err = client.Get(context.Background(), "/override_user_agent")
 	require.NoError(t, err)
-
 }
 
 func setup(t *testing.T, username, password string, requests []testserver.TestRequestHandler) *GitlabNetClient {
diff --git a/client/testserver/testserver.go b/client/testserver/testserver.go
index b09c7ba..105e7a0 100644
--- a/client/testserver/testserver.go
+++ b/client/testserver/testserver.go
@@ -59,6 +59,35 @@ func StartHttpServer(t *testing.T, handlers []TestRequestHandler) string {
 	return server.URL
 }
 
+func StartRetryHttpServer(t *testing.T, handlers []TestRequestHandler) string {
+	os.Setenv("FF_GITLAB_SHELL_RETRYABLE_HTTP", "1")
+	attempts := map[string]int{}
+
+	retryMiddileware := func(next func(w http.ResponseWriter, r *http.Request)) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			attempts[r.URL.String()+r.Method]++
+			if attempts[r.URL.String()+r.Method] == 1 {
+				w.WriteHeader(500)
+				return
+			}
+
+			http.HandlerFunc(next).ServeHTTP(w, r)
+		})
+	}
+	t.Helper()
+
+	h := http.NewServeMux()
+
+	for _, handler := range handlers {
+		h.Handle(handler.Path, retryMiddileware(handler.Handler))
+	}
+
+	server := httptest.NewServer(h)
+	t.Cleanup(func() { server.Close() })
+
+	return server.URL
+}
+
 func StartHttpsServer(t *testing.T, handlers []TestRequestHandler, clientCAPath string) string {
 	t.Helper()
 
diff --git a/go.mod b/go.mod
index 23eecbb..4a1a5cf 100644
--- a/go.mod
+++ b/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/golang-jwt/jwt/v4 v4.4.1
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
 	github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0
+	github.com/hashicorp/go-retryablehttp v0.7.1
 	github.com/mattn/go-shellwords v1.0.11
 	github.com/mikesmitty/edkey v0.0.0-20170222072505-3356ea4e686a
 	github.com/otiai10/copy v1.4.2
@@ -48,6 +49,7 @@ require (
 	github.com/google/pprof v0.0.0-20210804190019-f964ff605595 // indirect
 	github.com/google/uuid v1.3.0 // indirect
 	github.com/googleapis/gax-go/v2 v2.2.0 // indirect
+	github.com/hashicorp/go-cleanhttp v0.5.1 // indirect
 	github.com/hashicorp/yamux v0.1.1 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/kr/text v0.2.0 // indirect
diff --git a/go.sum b/go.sum
index 29c9560..468372d 100644
--- a/go.sum
+++ b/go.sum
@@ -227,6 +227,12 @@ github.com/grpc-ecosystem/go-grpc-middleware v1.3.0/go.mod h1:z0ButlSOZa5vEBq9m2
 github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0 h1:Ovs26xHkKqVztRpIrF/92BcuyuQ/YW4NSIpoGtfXNho=
 github.com/grpc-ecosystem/go-grpc-prometheus v1.2.0/go.mod h1:8NvIoxWQoOIhqOTXgfV/d3M/q6VIi02HzZEHgUlZvzk=
 github.com/grpc-ecosystem/grpc-gateway v1.16.0/go.mod h1:BDjrQk3hbvj6Nolgz8mAMFbcEtjT1g+wF4CSlocrBnw=
+github.com/hashicorp/go-cleanhttp v0.5.1 h1:dH3aiDG9Jvb5r5+bYHsikaOUIpcM0xvgMXVoDkXMzJM=
+github.com/hashicorp/go-cleanhttp v0.5.1/go.mod h1:JpRdi6/HCYpAwUzNwuwqhbovhLtngrth3wmdIIUrZ80=
+github.com/hashicorp/go-hclog v0.9.2 h1:CG6TE5H9/JXsFWJCfoIVpKFIkFe6ysEuHirp4DxCsHI=
+github.com/hashicorp/go-hclog v0.9.2/go.mod h1:5CU+agLiy3J7N7QjHK5d05KxGsuXiQLrjA0H7acj2lQ=
+github.com/hashicorp/go-retryablehttp v0.7.1 h1:sUiuQAnLlbvmExtFQs72iFW/HXeUn8Z1aJLQ4LJJbTQ=
+github.com/hashicorp/go-retryablehttp v0.7.1/go.mod h1:vAew36LZh98gCBJNLH42IQ1ER/9wtLZZ8meHqQvEYWY=
 github.com/hashicorp/golang-lru v0.5.0/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/golang-lru v0.5.1/go.mod h1:/m3WP610KZHVQ1SGc6re/UDhFvYD7pJ4Ao+sR/qLZy8=
 github.com/hashicorp/yamux v0.1.1 h1:yrQxtgseBDrq9Y652vSRDvsKCJKOUD+GzTS4Y0Y8pvE=
diff --git a/internal/config/config.go b/internal/config/config.go
index 35d8e74..b52f6f7 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -134,8 +134,14 @@ func (c *Config) HttpClient() (*client.HttpClient, error) {
 			return
 		}
 
-		tr := client.Transport
-		client.Transport = metrics.NewRoundTripper(tr)
+		if client.HTTPClient != nil {
+			tr := client.HTTPClient.Transport
+			client.HTTPClient.Transport = metrics.NewRoundTripper(tr)
+		}
+		if os.Getenv("FF_GITLAB_SHELL_RETRYABLE_HTTP") == "1" && client.RetryableHTTP != nil {
+			tr := client.RetryableHTTP.HTTPClient.Transport
+			client.RetryableHTTP.HTTPClient.Transport = metrics.NewRoundTripper(tr)
+		}
 
 		c.httpClient = client
 	})
diff --git a/internal/config/config_test.go b/internal/config/config_test.go
index c5fc219..4d0531b 100644
--- a/internal/config/config_test.go
+++ b/internal/config/config_test.go
@@ -28,36 +28,48 @@ func TestConfigApplyGlobalState(t *testing.T) {
 }
 
 func TestCustomPrometheusMetrics(t *testing.T) {
-	url := testserver.StartHttpServer(t, []testserver.TestRequestHandler{})
+	for _, ffValue := range []string{"0", "1"} {
+		t.Run("FF_GITLAB_SHELL_RETRYABLE_HTTP="+ffValue, func(t *testing.T) {
+			os.Setenv("FF_GITLAB_SHELL_RETRYABLE_HTTP", ffValue)
+			url := testserver.StartHttpServer(t, []testserver.TestRequestHandler{})
 
-	config := &Config{GitlabUrl: url}
-	client, err := config.HttpClient()
-	require.NoError(t, err)
+			config := &Config{GitlabUrl: url}
+			client, err := config.HttpClient()
+			require.NoError(t, err)
 
-	_, err = client.Get(url)
-	require.NoError(t, err)
+			if client.HTTPClient != nil {
+				_, err = client.HTTPClient.Get(url)
+				require.NoError(t, err)
+			}
 
-	ms, err := prometheus.DefaultGatherer.Gather()
-	require.NoError(t, err)
+			if os.Getenv("FF_GITLAB_SHELL_RETRYABLE_HTTP") == "1" && client.RetryableHTTP != nil {
+				_, err = client.RetryableHTTP.Get(url)
+				require.NoError(t, err)
+			}
 
-	var actualNames []string
-	for _, m := range ms[0:9] {
-		actualNames = append(actualNames, m.GetName())
-	}
+			ms, err := prometheus.DefaultGatherer.Gather()
+			require.NoError(t, err)
 
-	expectedMetricNames := []string{
-		"gitlab_shell_http_in_flight_requests",
-		"gitlab_shell_http_request_duration_seconds",
-		"gitlab_shell_http_requests_total",
-		"gitlab_shell_sshd_concurrent_limited_sessions_total",
-		"gitlab_shell_sshd_in_flight_connections",
-		"gitlab_shell_sshd_session_duration_seconds",
-		"gitlab_shell_sshd_session_established_duration_seconds",
-		"gitlab_sli:shell_sshd_sessions:errors_total",
-		"gitlab_sli:shell_sshd_sessions:total",
+			var actualNames []string
+			for _, m := range ms[0:9] {
+				actualNames = append(actualNames, m.GetName())
+			}
+
+			expectedMetricNames := []string{
+				"gitlab_shell_http_in_flight_requests",
+				"gitlab_shell_http_request_duration_seconds",
+				"gitlab_shell_http_requests_total",
+				"gitlab_shell_sshd_concurrent_limited_sessions_total",
+				"gitlab_shell_sshd_in_flight_connections",
+				"gitlab_shell_sshd_session_duration_seconds",
+				"gitlab_shell_sshd_session_established_duration_seconds",
+				"gitlab_sli:shell_sshd_sessions:errors_total",
+				"gitlab_sli:shell_sshd_sessions:total",
+			}
+
+			require.Equal(t, expectedMetricNames, actualNames)
+		})
 	}
-
-	require.Equal(t, expectedMetricNames, actualNames)
 }
 
 func TestNewFromDir(t *testing.T) {
diff --git a/spec/gitlab_shell_discover_spec.rb b/spec/gitlab_shell_discover_spec.rb
index 225d6b9..fe18b53 100644
--- a/spec/gitlab_shell_discover_spec.rb
+++ b/spec/gitlab_shell_discover_spec.rb
@@ -124,7 +124,12 @@ describe 'bin/gitlab-shell' do
     it 'returns an error message when the API call fails without a message' do
       _, stderr, status = run!(["-c/usr/share/webapps/gitlab-shell/bin/gitlab-shell", "username-broken"])
 
-      expect(stderr).to match(/Failed to get username: Internal API error \(500\)/)
+      stderr_output = if ENV['FF_GITLAB_SHELL_RETRYABLE_HTTP'] == '1'
+                        /Failed to get username: Internal API unreachable/
+                      else
+                        /Failed to get username: Internal API error \(500\)/
+                      end
+      expect(stderr).to match(stderr_output)
       expect(status).not_to be_success
     end
   end