From ca3d460278e7e0df31428e349aa1a761dd68f826 Mon Sep 17 00:00:00 2001
From: Aaron Lehmann <aaron.lehmann@docker.com>
Date: Wed, 19 Aug 2015 17:57:18 -0700
Subject: [PATCH] Add a TCP health checker

Also, add timeout and status code parameters to the HTTP checker, and
remove the threshold parameter for the file checker.

Signed-off-by: Aaron Lehmann <aaron.lehmann@docker.com>
---
 docs/handlers/app.go         | 38 ++++++++++++++++------
 docs/handlers/health_test.go | 63 ++++++++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+), 9 deletions(-)

diff --git a/docs/handlers/app.go b/docs/handlers/app.go
index 91f4e1a3..24f43f37 100644
--- a/docs/handlers/app.go
+++ b/docs/handlers/app.go
@@ -266,13 +266,8 @@ func (app *App) RegisterHealthChecks(healthRegistries ...*health.Registry) {
 		if interval == 0 {
 			interval = defaultCheckInterval
 		}
-		if fileChecker.Threshold != 0 {
-			ctxu.GetLogger(app).Infof("configuring file health check path=%s, interval=%d, threshold=%d", fileChecker.File, interval/time.Second, fileChecker.Threshold)
-			healthRegistry.Register(fileChecker.File, health.PeriodicThresholdChecker(checks.FileChecker(fileChecker.File), interval, fileChecker.Threshold))
-		} else {
-			ctxu.GetLogger(app).Infof("configuring file health check path=%s, interval=%d", fileChecker.File, interval/time.Second)
-			healthRegistry.Register(fileChecker.File, health.PeriodicChecker(checks.FileChecker(fileChecker.File), interval))
-		}
+		ctxu.GetLogger(app).Infof("configuring file health check path=%s, interval=%d", fileChecker.File, interval/time.Second)
+		healthRegistry.Register(fileChecker.File, health.PeriodicChecker(checks.FileChecker(fileChecker.File), interval))
 	}
 
 	for _, httpChecker := range app.Config.Health.HTTPCheckers {
@@ -280,12 +275,37 @@ func (app *App) RegisterHealthChecks(healthRegistries ...*health.Registry) {
 		if interval == 0 {
 			interval = defaultCheckInterval
 		}
+
+		statusCode := httpChecker.StatusCode
+		if statusCode == 0 {
+			statusCode = 200
+		}
+
+		checker := checks.HTTPChecker(httpChecker.URI, statusCode, httpChecker.Timeout)
+
 		if httpChecker.Threshold != 0 {
 			ctxu.GetLogger(app).Infof("configuring HTTP health check uri=%s, interval=%d, threshold=%d", httpChecker.URI, interval/time.Second, httpChecker.Threshold)
-			healthRegistry.Register(httpChecker.URI, health.PeriodicThresholdChecker(checks.HTTPChecker(httpChecker.URI), interval, httpChecker.Threshold))
+			healthRegistry.Register(httpChecker.URI, health.PeriodicThresholdChecker(checker, interval, httpChecker.Threshold))
 		} else {
 			ctxu.GetLogger(app).Infof("configuring HTTP health check uri=%s, interval=%d", httpChecker.URI, interval/time.Second)
-			healthRegistry.Register(httpChecker.URI, health.PeriodicChecker(checks.HTTPChecker(httpChecker.URI), interval))
+			healthRegistry.Register(httpChecker.URI, health.PeriodicChecker(checker, interval))
+		}
+	}
+
+	for _, tcpChecker := range app.Config.Health.TCPCheckers {
+		interval := tcpChecker.Interval
+		if interval == 0 {
+			interval = defaultCheckInterval
+		}
+
+		checker := checks.TCPChecker(tcpChecker.Addr, tcpChecker.Timeout)
+
+		if tcpChecker.Threshold != 0 {
+			ctxu.GetLogger(app).Infof("configuring TCP health check addr=%s, interval=%d, threshold=%d", tcpChecker.Addr, interval/time.Second, tcpChecker.Threshold)
+			healthRegistry.Register(tcpChecker.Addr, health.PeriodicThresholdChecker(checker, interval, tcpChecker.Threshold))
+		} else {
+			ctxu.GetLogger(app).Infof("configuring TCP health check addr=%s, interval=%d", tcpChecker.Addr, interval/time.Second)
+			healthRegistry.Register(tcpChecker.Addr, health.PeriodicChecker(checker, interval))
 		}
 	}
 }
diff --git a/docs/handlers/health_test.go b/docs/handlers/health_test.go
index de2b71cc..bb460b47 100644
--- a/docs/handlers/health_test.go
+++ b/docs/handlers/health_test.go
@@ -2,6 +2,7 @@ package handlers
 
 import (
 	"io/ioutil"
+	"net"
 	"net/http"
 	"net/http/httptest"
 	"os"
@@ -61,6 +62,68 @@ func TestFileHealthCheck(t *testing.T) {
 	}
 }
 
+func TestTCPHealthCheck(t *testing.T) {
+	interval := time.Second
+
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("could not create listener: %v", err)
+	}
+	addrStr := ln.Addr().String()
+
+	// Start accepting
+	go func() {
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				// listener was closed
+				return
+			}
+			defer conn.Close()
+		}
+	}()
+
+	config := configuration.Configuration{
+		Storage: configuration.Storage{
+			"inmemory": configuration.Parameters{},
+		},
+		Health: configuration.Health{
+			TCPCheckers: []configuration.TCPChecker{
+				{
+					Interval: interval,
+					Addr:     addrStr,
+					Timeout:  500 * time.Millisecond,
+				},
+			},
+		},
+	}
+
+	ctx := context.Background()
+
+	app := NewApp(ctx, config)
+	healthRegistry := health.NewRegistry()
+	app.RegisterHealthChecks(healthRegistry)
+
+	// Wait for health check to happen
+	<-time.After(2 * interval)
+
+	if len(healthRegistry.CheckStatus()) != 0 {
+		t.Fatal("expected 0 items in health check results")
+	}
+
+	ln.Close()
+	<-time.After(2 * interval)
+
+	// Health check should now fail
+	status := healthRegistry.CheckStatus()
+	if len(status) != 1 {
+		t.Fatal("expected 1 item in health check results")
+	}
+	if status[addrStr] != "connection to "+addrStr+" failed" {
+		t.Fatal(`did not get "connection failed" result for health check`)
+	}
+}
+
 func TestHTTPHealthCheck(t *testing.T) {
 	interval := time.Second
 	threshold := 3