From f4e97be78e4688e32c139c4efaded4c3476a1523 Mon Sep 17 00:00:00 2001
From: Brendan Burns <bburns@google.com>
Date: Tue, 23 Jun 2015 21:47:24 -0700
Subject: [PATCH] Switch to using the official etcd health check.

---
 pkg/apiserver/validator.go      |  8 ++++++++
 pkg/apiserver/validator_test.go | 20 +++++++++++++++++---
 pkg/master/master.go            |  2 +-
 pkg/tools/etcd_helper.go        | 17 +++++++++++++++++
 pkg/tools/etcd_helper_test.go   | 29 +++++++++++++++++++++++++++++
 5 files changed, 72 insertions(+), 4 deletions(-)

diff --git a/pkg/apiserver/validator.go b/pkg/apiserver/validator.go
index cd23efaa5bb..93f45b14135 100644
--- a/pkg/apiserver/validator.go
+++ b/pkg/apiserver/validator.go
@@ -33,11 +33,14 @@ type httpGet interface {
 	Get(url string) (*http.Response, error)
 }
 
+type ValidatorFn func([]byte) error
+
 type Server struct {
 	Addr        string
 	Port        int
 	Path        string
 	EnableHTTPS bool
+	Validate    ValidatorFn
 }
 
 type ServerStatus struct {
@@ -85,5 +88,10 @@ func (server *Server) DoServerCheck(rt http.RoundTripper) (probe.Result, string,
 		return probe.Failure, string(data),
 			fmt.Errorf("unhealthy http status code: %d (%s)", resp.StatusCode, resp.Status)
 	}
+	if server.Validate != nil {
+		if err := server.Validate(data); err != nil {
+			return probe.Failure, string(data), err
+		}
+	}
 	return probe.Success, string(data), nil
 }
diff --git a/pkg/apiserver/validator_test.go b/pkg/apiserver/validator_test.go
index 582109a6720..c359f5735bc 100644
--- a/pkg/apiserver/validator_test.go
+++ b/pkg/apiserver/validator_test.go
@@ -18,6 +18,7 @@ package apiserver
 
 import (
 	"bytes"
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"net/http"
@@ -37,6 +38,15 @@ func (f *fakeRoundTripper) RoundTrip(req *http.Request) (*http.Response, error)
 	return f.resp, f.err
 }
 
+func alwaysError([]byte) error { return errors.New("test error") }
+
+func matchError(data []byte) error {
+	if string(data) == "bar" {
+		return errors.New("match error")
+	}
+	return nil
+}
+
 func TestValidate(t *testing.T) {
 	tests := []struct {
 		err            error
@@ -44,10 +54,13 @@ func TestValidate(t *testing.T) {
 		expectedStatus probe.Result
 		code           int
 		expectErr      bool
+		validator      ValidatorFn
 	}{
-		{fmt.Errorf("test error"), "", probe.Unknown, 500 /*ignored*/, true},
-		{nil, "foo", probe.Success, 200, false},
-		{nil, "foo", probe.Failure, 500, true},
+		{fmt.Errorf("test error"), "", probe.Unknown, 500 /*ignored*/, true, nil},
+		{nil, "foo", probe.Success, 200, false, nil},
+		{nil, "foo", probe.Failure, 500, true, nil},
+		{nil, "foo", probe.Failure, 200, true, alwaysError},
+		{nil, "foo", probe.Success, 200, false, matchError},
 	}
 
 	s := Server{Addr: "foo.com", Port: 8080, Path: "/healthz"}
@@ -60,6 +73,7 @@ func TestValidate(t *testing.T) {
 				StatusCode: test.code,
 			},
 		}
+		s.Validate = test.validator
 		status, data, err := s.DoServerCheck(fakeRT)
 		expect := fmt.Sprintf("http://%s:%d/healthz", s.Addr, s.Port)
 		if fakeRT.url != expect {
diff --git a/pkg/master/master.go b/pkg/master/master.go
index 26c549c6973..a15b7a3daad 100644
--- a/pkg/master/master.go
+++ b/pkg/master/master.go
@@ -715,7 +715,7 @@ func (m *Master) getServersToValidate(c *Config) map[string]apiserver.Server {
 			addr = etcdUrl.Host
 			port = 4001
 		}
-		serversToValidate[fmt.Sprintf("etcd-%d", ix)] = apiserver.Server{Addr: addr, Port: port, Path: "/v2/keys/"}
+		serversToValidate[fmt.Sprintf("etcd-%d", ix)] = apiserver.Server{Addr: addr, Port: port, Path: "/health", Validate: tools.EtcdHealthCheck}
 	}
 	return serversToValidate
 }
diff --git a/pkg/tools/etcd_helper.go b/pkg/tools/etcd_helper.go
index d60cf68b43a..05e2892eecd 100644
--- a/pkg/tools/etcd_helper.go
+++ b/pkg/tools/etcd_helper.go
@@ -17,6 +17,7 @@ limitations under the License.
 package tools
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io/ioutil"
@@ -712,3 +713,19 @@ func NewEtcdClientStartServerIfNecessary(server string) (EtcdClient, error) {
 	servers := []string{server}
 	return etcd.NewClient(servers), nil
 }
+
+type etcdHealth struct {
+	// Note this has to be public so the json library can modify it.
+	Health string `json:health`
+}
+
+func EtcdHealthCheck(data []byte) error {
+	obj := etcdHealth{}
+	if err := json.Unmarshal(data, &obj); err != nil {
+		return err
+	}
+	if obj.Health != "true" {
+		return fmt.Errorf("Unhealthy status: %s", obj.Health)
+	}
+	return nil
+}
diff --git a/pkg/tools/etcd_helper_test.go b/pkg/tools/etcd_helper_test.go
index 5c39002d4cc..95beb529ea0 100644
--- a/pkg/tools/etcd_helper_test.go
+++ b/pkg/tools/etcd_helper_test.go
@@ -856,3 +856,32 @@ func TestPrefixEtcdKey(t *testing.T) {
 
 	assert.Equal(t, keyBefore, keyAfter, "Prefix incorrectly added by EtcdHelper")
 }
+
+func TestEtcdHealthCheck(t *testing.T) {
+	tests := []struct {
+		data      string
+		expectErr bool
+	}{
+		{
+			data:      "{\"health\": \"true\"}",
+			expectErr: false,
+		},
+		{
+			data:      "{\"health\": \"false\"}",
+			expectErr: true,
+		},
+		{
+			data:      "invalid json",
+			expectErr: true,
+		},
+	}
+	for _, test := range tests {
+		err := EtcdHealthCheck([]byte(test.data))
+		if err != nil && !test.expectErr {
+			t.Errorf("unexpected error: %v", err)
+		}
+		if err == nil && test.expectErr {
+			t.Error("unexpected non-error")
+		}
+	}
+}