Merge branch '495-ttfb-timeout' into 'master'

Add more client timeouts

Closes #495

See merge request gitlab-org/gitlab-pages!382

4 files changed, 88 insertions, 25 deletions

diff --git a/internal/httprange/http_reader.go b/internal/httprange/http_reader.go
index 5dc0f693..467256a0 100644
--- a/internal/httprange/http_reader.go
+++ b/internal/httprange/http_reader.go
@@ -62,6 +62,7 @@ var httpClient = &http.Client{
 		metrics.HTTPRangeTraceDuration,
 		metrics.HTTPRangeRequestDuration,
 		metrics.HTTPRangeRequestsTotal,
+		httptransport.DefaultTTFBTimeout,
 	),
 }
 
diff --git a/internal/httptransport/transport.go b/internal/httptransport/transport.go
index bc871ea7..d8e6a3fe 100644
--- a/internal/httptransport/transport.go
+++ b/internal/httptransport/transport.go
@@ -1,6 +1,7 @@
 package httptransport
 
 import (
+	"context"
 	"crypto/tls"
 	"crypto/x509"
 	"net"
@@ -15,6 +16,13 @@ import (
 	log "github.com/sirupsen/logrus"
 )
 
+const (
+	// DefaultTTFBTimeout is the timeout used in the meteredRoundTripper
+	// when calling http.Transport.RoundTrip. The request will be cancelled
+	// if the response takes longer than this.
+	DefaultTTFBTimeout = 15 * time.Second
+)
+
 var (
 	sysPoolOnce = &sync.Once{}
 	sysPool     *x509.CertPool
@@ -26,11 +34,12 @@ var (
 )
 
 type meteredRoundTripper struct {
-	next      http.RoundTripper
-	name      string
-	tracer    *prometheus.HistogramVec
-	durations *prometheus.HistogramVec
-	counter   *prometheus.CounterVec
+	next        http.RoundTripper
+	name        string
+	tracer      *prometheus.HistogramVec
+	durations   *prometheus.HistogramVec
+	counter     *prometheus.CounterVec
+	ttfbTimeout time.Duration
 }
 
 func newInternalTransport() *http.Transport {
@@ -43,19 +52,24 @@ func newInternalTransport() *http.Transport {
 		MaxIdleConns:        100,
 		MaxIdleConnsPerHost: 100,
 		IdleConnTimeout:     90 * time.Second,
+		// Set more timeouts https://gitlab.com/gitlab-org/gitlab-pages/-/issues/495
+		TLSHandshakeTimeout:   10 * time.Second,
+		ResponseHeaderTimeout: 15 * time.Second,
+		ExpectContinueTimeout: 15 * time.Second,
 	}
 }
 
 // NewTransportWithMetrics will create a custom http.RoundTripper that can be used with an http.Client.
 // The RoundTripper will report metrics based on the collectors passed.
 func NewTransportWithMetrics(name string, tracerVec, durationsVec *prometheus.
-	HistogramVec, counterVec *prometheus.CounterVec) http.RoundTripper {
+	HistogramVec, counterVec *prometheus.CounterVec, ttfbTimeout time.Duration) http.RoundTripper {
 	return &meteredRoundTripper{
-		next:      InternalTransport,
-		name:      name,
-		tracer:    tracerVec,
-		durations: durationsVec,
-		counter:   counterVec,
+		next:        InternalTransport,
+		name:        name,
+		tracer:      tracerVec,
+		durations:   durationsVec,
+		counter:     counterVec,
+		ttfbTimeout: ttfbTimeout,
 	}
 }
 
@@ -88,7 +102,13 @@ func loadPool() {
 func (mrt *meteredRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
 	start := time.Now()
 
-	r = r.WithContext(httptrace.WithClientTrace(r.Context(), mrt.newTracer(start)))
+	ctx := httptrace.WithClientTrace(r.Context(), mrt.newTracer(start))
+	ctx, cancel := context.WithCancel(ctx)
+
+	timer := time.AfterFunc(mrt.ttfbTimeout, cancel)
+	defer timer.Stop()
+
+	r = r.WithContext(ctx)
 
 	resp, err := mrt.next.RoundTrip(r)
 	if err != nil {
diff --git a/internal/httptransport/transport_test.go b/internal/httptransport/transport_test.go
index a4105bef..9059ea15 100644
--- a/internal/httptransport/transport_test.go
+++ b/internal/httptransport/transport_test.go
@@ -1,6 +1,8 @@
 package httptransport
 
 import (
+	"context"
+	"errors"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
@@ -43,22 +45,17 @@ func Test_withRoundTripper(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			histVec := prometheus.NewHistogramVec(prometheus.HistogramOpts{
-				Name: t.Name(),
-			}, []string{"status_code"})
-
-			counterVec := prometheus.NewCounterVec(prometheus.CounterOpts{
-				Name: t.Name(),
-			}, []string{"status_code"})
+			histVec, counterVec := newTestMetrics(t)
 
 			next := &mockRoundTripper{
 				res: &http.Response{
 					StatusCode: tt.statusCode,
 				},
-				err: tt.err,
+				err:     tt.err,
+				timeout: time.Nanosecond,
 			}
 
-			mtr := &meteredRoundTripper{next: next, durations: histVec, counter: counterVec}
+			mtr := &meteredRoundTripper{next: next, durations: histVec, counter: counterVec, ttfbTimeout: DefaultTTFBTimeout}
 			r := httptest.NewRequest("GET", "/", nil)
 
 			res, err := mtr.RoundTrip(r)
@@ -78,13 +75,53 @@ func Test_withRoundTripper(t *testing.T) {
 	}
 }
 
+func TestRoundTripTTFBTimeout(t *testing.T) {
+	histVec, counterVec := newTestMetrics(t)
+
+	next := &mockRoundTripper{
+		res: &http.Response{
+			StatusCode: http.StatusOK,
+		},
+		timeout: time.Millisecond,
+		err:     nil,
+	}
+
+	mtr := &meteredRoundTripper{next: next, durations: histVec, counter: counterVec, ttfbTimeout: time.Nanosecond}
+	req, err := http.NewRequest("GET", "https://gitlab.com", nil)
+	require.NoError(t, err)
+
+	res, err := mtr.RoundTrip(req)
+	require.Nil(t, res)
+	require.True(t, errors.Is(err, context.Canceled), "context must have been canceled after ttfb timeout")
+}
+
+func newTestMetrics(t *testing.T) (*prometheus.HistogramVec, *prometheus.CounterVec) {
+	t.Helper()
+
+	histVec := prometheus.NewHistogramVec(prometheus.HistogramOpts{
+		Name: t.Name(),
+	}, []string{"status_code"})
+
+	counterVec := prometheus.NewCounterVec(prometheus.CounterOpts{
+		Name: t.Name(),
+	}, []string{"status_code"})
+
+	return histVec, counterVec
+}
+
 type mockRoundTripper struct {
-	res *http.Response
-	err error
+	res     *http.Response
+	err     error
+	timeout time.Duration
 }
 
 func (mrt *mockRoundTripper) RoundTrip(r *http.Request) (*http.Response, error) {
-	return mrt.res, mrt.err
+	select {
+	case <-r.Context().Done():
+		return nil, r.Context().Err()
+	case <-time.After(mrt.timeout):
+		return mrt.res, mrt.err
+	}
 }
 
 func TestInternalTransportShouldHaveCustomConnectionPoolSettings(t *testing.T) {
@@ -92,4 +129,7 @@ func TestInternalTransportShouldHaveCustomConnectionPoolSettings(t *testing.T) {
 	require.EqualValues(t, 100, InternalTransport.MaxIdleConnsPerHost)
 	require.EqualValues(t, 0, InternalTransport.MaxConnsPerHost)
 	require.EqualValues(t, 90*time.Second, InternalTransport.IdleConnTimeout)
+	require.EqualValues(t, 10*time.Second, InternalTransport.TLSHandshakeTimeout)
+	require.EqualValues(t, 15*time.Second, InternalTransport.ResponseHeaderTimeout)
+	require.EqualValues(t, 15*time.Second, InternalTransport.ExpectContinueTimeout)
 }
diff --git a/internal/source/gitlab/client/client.go b/internal/source/gitlab/client/client.go
index 0e8235c0..b11ea2cb 100644
--- a/internal/source/gitlab/client/client.go
+++ b/internal/source/gitlab/client/client.go
@@ -60,7 +60,9 @@ func NewClient(baseURL string, secretKey []byte, connectionTimeout, jwtTokenExpi
 				"gitlab_internal_api",
 				metrics.DomainsSourceAPITraceDuration,
 				metrics.DomainsSourceAPICallDuration,
-				metrics.DomainsSourceAPIReqTotal),
+				metrics.DomainsSourceAPIReqTotal,
+				httptransport.DefaultTTFBTimeout,
+			),
 		},
 		jwtTokenExpiry: jwtTokenExpiry,
 	}, nil