Add timeout for /health endpoint checks (#145)

* Add timeout for /health endpoint checks

* Switch back

* Tests

* Use timeout

* Use timeout server

* Descriptive test case naming

* Error check

Author: cristinaleonr

cmd/heartbeat/health/checker.go

@@ -8,20 +8,23 @@ import (
 type Checker struct {
 	pp  *PortProbe
 	k8s *KubernetesClient
+	ec  *EndpointClient
 }
 
 // NewChecker creates a new Checker.
-func NewChecker(pp *PortProbe) *Checker {
+func NewChecker(pp *PortProbe, ec *EndpointClient) *Checker {
 	return &Checker{
 		pp: pp,
+		ec: ec,
 	}
 }
 
 // NewCheckerK8S creates a new Checker for Kubernetes deployments.
-func NewCheckerK8S(pp *PortProbe, k8s *KubernetesClient) *Checker {
+func NewCheckerK8S(pp *PortProbe, k8s *KubernetesClient, ec *EndpointClient) *Checker {
 	return &Checker{
 		pp:  pp,
 		k8s: k8s,
+		ec:  ec,
 	}
 }
 
@@ -38,7 +41,7 @@ func (hc *Checker) GetHealth(ctx context.Context) float64 {
 	// Some experiments might not support a /health endpoint, so
 	// the result is only taken into account if the request error
 	// is nil.
-	status, err := checkHealthEndpoint()
+	status, err := hc.ec.checkHealthEndpoint()
 	if err == nil && !status {
 		return 0
 	}

cmd/heartbeat/health/checker_test.go

@@ -23,6 +23,7 @@ func TestChecker_getHealth(t *testing.T) {
 				&KubernetesClient{
 					clientset: healthyClientset,
 				},
+				&EndpointClient{},
 			),
 			endpointStatus: 200,
 			want:           1,
@@ -31,6 +32,7 @@ func TestChecker_getHealth(t *testing.T) {
 			name: "health-1-k8s-nil",
 			checker: NewChecker(
 				&PortProbe{},
+				&EndpointClient{},
 			),
 			endpointStatus: 200,
 			want:           1,
@@ -44,6 +46,7 @@ func TestChecker_getHealth(t *testing.T) {
 				&KubernetesClient{
 					clientset: healthyClientset,
 				},
+				&EndpointClient{},
 			),
 			endpointStatus: 200,
 			want:           0,
@@ -55,6 +58,7 @@ func TestChecker_getHealth(t *testing.T) {
 				&KubernetesClient{
 					clientset: fake.NewSimpleClientset(),
 				},
+				&EndpointClient{},
 			),
 			endpointStatus: 200,
 			want:           1,
@@ -79,6 +83,7 @@ func TestChecker_getHealth(t *testing.T) {
 						},
 					),
 				},
+				&EndpointClient{},
 			),
 			endpointStatus: 200,
 			want:           0,
@@ -90,6 +95,7 @@ func TestChecker_getHealth(t *testing.T) {
 				&KubernetesClient{
 					clientset: healthyClientset,
 				},
+				&EndpointClient{},
 			),
 			endpointStatus: 500,
 			want:           0,
@@ -116,6 +122,7 @@ func TestChecker_getHealth(t *testing.T) {
 						},
 					),
 				},
+				&EndpointClient{},
 			),
 			want: 0,
 		},
@@ -125,6 +132,7 @@ func TestChecker_getHealth(t *testing.T) {
 				&PortProbe{
 					ports: map[string]bool{"65536": true},
 				},
+				&EndpointClient{},
 			),
 			want: 0,
 		},

cmd/heartbeat/health/endpoint-check.go

@@ -2,6 +2,7 @@ package health
 
 import (
 	"net/http"
+	"time"
 
 	"github.com/m-lab/locate/metrics"
 )
@@ -10,11 +11,26 @@ var (
 	healthAddress = "http://localhost:8000/health"
 )
 
+// EndpointClient is an http client to check the local /health endpoint.
+type EndpointClient struct {
+	http.Client
+}
+
+// NewEndpointClient returns a new *EndpointClient with the specified request
+// timeout.
+func NewEndpointClient(timeout time.Duration) *EndpointClient {
+	return &EndpointClient{
+		http.Client{
+			Timeout: timeout,
+		},
+	}
+}
+
 // checkHealthEndpoint makes a call to the the local /health endpoint.
 // It returns an error if the HTTP request was not successful.
 // It returns true only if the returned HTTP status code equals 200 (OK).
-func checkHealthEndpoint() (bool, error) {
-	resp, err := http.Get(healthAddress)
+func (ec *EndpointClient) checkHealthEndpoint() (bool, error) {
+	resp, err := ec.Get(healthAddress)
 	if err != nil {
 		metrics.HealthEndpointChecksTotal.WithLabelValues("HTTP request error").Inc()
 		return false, err

cmd/heartbeat/health/endpoint-check_test.go

@@ -3,6 +3,7 @@ package health
 import (
 	"net/http"
 	"testing"
+	"time"
 
 	"github.com/m-lab/locate/cmd/heartbeat/health/healthtest"
 )
@@ -27,7 +28,7 @@ func Test_checkHealthEndpoint(t *testing.T) {
 			wantErr: false,
 		},
 		{
-			name:    "error",
+			name:    "error-server-not-running",
 			want:    false,
 			wantErr: true,
 		},
@@ -40,7 +41,8 @@ func Test_checkHealthEndpoint(t *testing.T) {
 				defer srv.Close()
 			}
 
-			got, err := checkHealthEndpoint()
+			hc := NewEndpointClient(time.Second)
+			got, err := hc.checkHealthEndpoint()
 			if (err != nil) != tt.wantErr {
 				t.Errorf("checkHealthEndpoint() error = %v, wantErr %v", err, tt.wantErr)
 				return
@@ -52,3 +54,20 @@ func Test_checkHealthEndpoint(t *testing.T) {
 		})
 	}
 }
+
+func Test_checkHealthEndpoint_timeout(t *testing.T) {
+	srv := healthtest.TestTimeoutServer(2 * time.Second)
+	healthAddress = srv.URL + "/health"
+	defer srv.Close()
+
+	hc := NewEndpointClient(time.Second)
+	got, err := hc.checkHealthEndpoint()
+	if err == nil {
+		t.Errorf("checkHealthEndpoint() error = %v, wantErr %s", err, "context deadline error")
+		return
+	}
+
+	if got != false {
+		t.Errorf("checkHealthEndpoint() = %v, want %v", got, false)
+	}
+}

cmd/heartbeat/health/healthtest/healthtest.go

@@ -3,6 +3,7 @@ package healthtest
 import (
 	"net/http"
 	"net/http/httptest"
+	"time"
 )
 
 func TestHealthServer(code int) *httptest.Server {
@@ -13,3 +14,12 @@ func TestHealthServer(code int) *httptest.Server {
 	s := httptest.NewServer(mux)
 	return s
 }
+
+func TestTimeoutServer(timeout time.Duration) *httptest.Server {
+	mux := http.NewServeMux()
+	mux.Handle("/health", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(timeout)
+	}))
+	s := httptest.NewServer(mux)
+	return s
+}

cmd/heartbeat/main.go

@@ -77,12 +77,13 @@ func main() {
 	rtx.Must(err, "failed to establish a websocket connection with %s", heartbeatURL)
 
 	probe := health.NewPortProbe(svcs)
+	ec := health.NewEndpointClient(static.HealthEndpointTimeout)
 	hc := &health.Checker{}
 	if kubernetesURL.URL == nil {
-		hc = health.NewChecker(probe)
+		hc = health.NewChecker(probe, ec)
 	} else {
 		k8s := health.MustNewKubernetesClient(kubernetesURL.URL, pod, node, namespace, kubernetesAuth)
-		hc = health.NewCheckerK8S(probe, k8s)
+		hc = health.NewCheckerK8S(probe, k8s, ec)
 	}
 
 	write(conn, hc, ldr)

static/configs.go

@@ -20,6 +20,7 @@ const (
 	BackoffMultiplier          = 2
 	BackoffMaxInterval         = 5 * time.Minute
 	BackoffMaxElapsedTime      = 0
+	HealthEndpointTimeout      = 5 * time.Second
 	HeartbeatPeriod            = 10 * time.Second
 	MemorystoreExportPeriod    = 10 * time.Second
 	PrometheusCheckPeriod      = time.Minute