Merge pull request #308 from austince/feat/with-idle-conn-timeout

Add WithIdleConnTimeout HTTP client option

Author: Julien Pivotto

config/http_config.go

@@ -47,6 +47,9 @@ var DefaultHTTPClientConfig = HTTPClientConfig{
 var defaultHTTPClientOptions = httpClientOptions{
 	keepAlivesEnabled: true,
 	http2Enabled:      true,
+	// 5 minutes is typically above the maximum sane scrape interval. So we can
+	// use keepalive for all configurations.
+	idleConnTimeout: 5 * time.Minute,
 }
 
 type closeIdler interface {
@@ -283,6 +286,7 @@ type httpClientOptions struct {
 	dialContextFunc   DialContextFunc
 	keepAlivesEnabled bool
 	http2Enabled      bool
+	idleConnTimeout   time.Duration
 }
 
 // HTTPClientOption defines an option that can be applied to the HTTP client.
@@ -309,6 +313,13 @@ func WithHTTP2Disabled() HTTPClientOption {
 	}
 }
 
+// WithIdleConnTimeout allows setting the idle connection timeout.
+func WithIdleConnTimeout(timeout time.Duration) HTTPClientOption {
+	return func(opts *httpClientOptions) {
+		opts.idleConnTimeout = timeout
+	}
+}
+
 // NewClient returns a http.Client using the specified http.RoundTripper.
 func newClient(rt http.RoundTripper) *http.Client {
 	return &http.Client{Transport: rt}
@@ -357,15 +368,13 @@ func NewRoundTripperFromConfig(cfg HTTPClientConfig, name string, optFuncs ...HT
 		// The only timeout we care about is the configured scrape timeout.
 		// It is applied on request. So we leave out any timings here.
 		var rt http.RoundTripper = &http.Transport{
-			Proxy:               http.ProxyURL(cfg.ProxyURL.URL),
-			MaxIdleConns:        20000,
-			MaxIdleConnsPerHost: 1000, // see https://github.com/golang/go/issues/13801
-			DisableKeepAlives:   !opts.keepAlivesEnabled,
-			TLSClientConfig:     tlsConfig,
-			DisableCompression:  true,
-			// 5 minutes is typically above the maximum sane scrape interval. So we can
-			// use keepalive for all configurations.
-			IdleConnTimeout:       5 * time.Minute,
+			Proxy:                 http.ProxyURL(cfg.ProxyURL.URL),
+			MaxIdleConns:          20000,
+			MaxIdleConnsPerHost:   1000, // see https://github.com/golang/go/issues/13801
+			DisableKeepAlives:     !opts.keepAlivesEnabled,
+			TLSClientConfig:       tlsConfig,
+			DisableCompression:    true,
+			IdleConnTimeout:       opts.idleConnTimeout,
 			TLSHandshakeTimeout:   10 * time.Second,
 			ExpectContinueTimeout: 1 * time.Second,
 			DialContext:           dialContext,

config/http_config_test.go

@@ -456,6 +456,25 @@ func TestCustomDialContextFunc(t *testing.T) {
 	}
 }
 
+func TestCustomIdleConnTimeout(t *testing.T) {
+	timeout := time.Second * 5
+
+	cfg := HTTPClientConfig{}
+	rt, err := NewRoundTripperFromConfig(cfg, "test", WithIdleConnTimeout(timeout))
+	if err != nil {
+		t.Fatalf("Can't create a round-tripper from this config: %+v", cfg)
+	}
+
+	transport, ok := rt.(*http.Transport)
+	if !ok {
+		t.Fatalf("Unexpected transport: %+v", transport)
+	}
+
+	if transport.IdleConnTimeout != timeout {
+		t.Fatalf("Unexpected idle connection timeout: %+v", timeout)
+	}
+}
+
 func TestMissingBearerAuthFile(t *testing.T) {
 	cfg := HTTPClientConfig{
 		BearerTokenFile: MissingBearerTokenFile,