Merge pull request #1920 from weaveworks/probe-client-timeouts

Add timeouts to the probe http client
weaveworks · Oct 13, 2016 · 8da5ca3 · 8da5ca3
2 parents 1d3ae95 + 002770d
commit 8da5ca3
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 26 deletions.
diff --git a/probe/appclient/app_client.go b/probe/appclient/app_client.go
@@ -4,7 +4,6 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
-	"net"
 	"net/http"
 	"net/rpc"
 	"net/url"
@@ -13,6 +12,7 @@ import (
 
 	log "github.com/Sirupsen/logrus"
 	"github.com/gorilla/websocket"
+	"github.com/hashicorp/go-cleanhttp"
 	"github.com/ugorji/go/codec"
 
 	"github.com/weaveworks/scope/common/xfer"
@@ -22,7 +22,6 @@ const (
 	httpClientTimeout = 4 * time.Second
 	initialBackoff    = 1 * time.Second
 	maxBackoff        = 60 * time.Second
-	dialTimeout       = 5 * time.Second
 )
 
 // AppClient is a client to an app for dealing with controls.
@@ -41,7 +40,7 @@ type appClient struct {
 
 	quit     chan struct{}
 	mtx      sync.Mutex
-	client   http.Client
+	client   *http.Client
 	wsDialer websocket.Dialer
 	appID    string
 	hostname string
@@ -63,26 +62,20 @@ type appClient struct {
 
 // NewAppClient makes a new appClient.
 func NewAppClient(pc ProbeConfig, hostname string, target url.URL, control xfer.ControlHandler) (AppClient, error) {
-	httpTransport, err := pc.getHTTPTransport(hostname)
-	if err != nil {
-		return nil, err
-	}
-
-	httpTransport.Dial = func(network, addr string) (net.Conn, error) {
-		return net.DialTimeout(network, addr, dialTimeout)
-	}
+	httpTransport := pc.getHTTPTransport(hostname)
+	httpClient := cleanhttp.DefaultClient()
+	httpClient.Transport = httpTransport
+	httpClient.Timeout = httpClientTimeout
 
 	return &appClient{
 		ProbeConfig: pc,
 		quit:        make(chan struct{}),
 		hostname:    hostname,
 		target:      target,
-		client: http.Client{
-			Transport: httpTransport,
-			Timeout:   httpClientTimeout,
-		},
+		client:      httpClient,
 		wsDialer: websocket.Dialer{
-			TLSClientConfig: httpTransport.TLSClientConfig,
+			TLSClientConfig:  httpTransport.TLSClientConfig,
+			HandshakeTimeout: httpClientTimeout,
 		},
 		conns:   map[string]xfer.Websocket{},
 		readers: make(chan io.Reader, 2),

diff --git a/probe/appclient/probe_config.go b/probe/appclient/probe_config.go
@@ -5,12 +5,20 @@ import (
 	"crypto/x509"
 	"fmt"
 	"io"
+	"net"
 	"net/http"
+	"time"
 
 	"github.com/certifi/gocertifi"
+	"github.com/hashicorp/go-cleanhttp"
+
 	"github.com/weaveworks/scope/common/xfer"
 )
 
+const (
+	dialTimeout = 5 * time.Second
+)
+
 var certPool *x509.CertPool
 
 func init() {
@@ -43,17 +51,19 @@ func (pc ProbeConfig) authorizedRequest(method string, urlStr string, body io.Re
 	return req, err
 }
 
-func (pc ProbeConfig) getHTTPTransport(hostname string) (*http.Transport, error) {
+func (pc ProbeConfig) getHTTPTransport(hostname string) *http.Transport {
+	transport := cleanhttp.DefaultTransport()
+	transport.DialContext = (&net.Dialer{
+		Timeout:   dialTimeout,
+		KeepAlive: 30 * time.Second,
+	}).DialContext
 	if pc.Insecure {
-		return &http.Transport{
-			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
-		}, nil
-	}
-
-	return &http.Transport{
-		TLSClientConfig: &tls.Config{
+		transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
+	} else {
+		transport.TLSClientConfig = &tls.Config{
 			RootCAs:    certPool,
 			ServerName: hostname,
-		},
-	}, nil
+		}
+	}
+	return transport
 }