From 0022639a2f9157d9d60c8f54b1913a19ae866c92 Mon Sep 17 00:00:00 2001
From: Tim Hockin <thockin@google.com>
Date: Mon, 10 Apr 2017 21:38:13 -0700
Subject: [PATCH] Zero-out healthchecks when no endpoints

---
 pkg/proxy/healthcheck/healthcheck.go      |  6 +++++
 pkg/proxy/healthcheck/healthcheck_test.go | 31 +++++++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/pkg/proxy/healthcheck/healthcheck.go b/pkg/proxy/healthcheck/healthcheck.go
index faaf001fad..999bc1b8ee 100644
--- a/pkg/proxy/healthcheck/healthcheck.go
+++ b/pkg/proxy/healthcheck/healthcheck.go
@@ -196,6 +196,7 @@ func (h hcHandler) ServeHTTP(resp http.ResponseWriter, req *http.Request) {
 	h.hcs.lock.Lock()
 	count := h.hcs.services[h.name].endpoints
 	h.hcs.lock.Unlock()
+
 	resp.Header().Set("Content-Type", "application/json")
 	if count == 0 {
 		resp.WriteHeader(http.StatusServiceUnavailable)
@@ -225,5 +226,10 @@ func (hcs *server) SyncEndpoints(newEndpoints map[types.NamespacedName]int) erro
 		glog.V(3).Infof("Reporting %d endpoints for healthcheck %q", count, nsn.String())
 		hcs.services[nsn].endpoints = count
 	}
+	for nsn, hci := range hcs.services {
+		if _, found := newEndpoints[nsn]; !found {
+			hci.endpoints = 0
+		}
+	}
 	return nil
 }
diff --git a/pkg/proxy/healthcheck/healthcheck_test.go b/pkg/proxy/healthcheck/healthcheck_test.go
index 208ffb18ad..a90a8a386e 100644
--- a/pkg/proxy/healthcheck/healthcheck_test.go
+++ b/pkg/proxy/healthcheck/healthcheck_test.go
@@ -174,6 +174,14 @@ func TestServer(t *testing.T) {
 	// test the handler
 	testHandler(hcs, nsn, http.StatusServiceUnavailable, 0, t)
 
+	// put the endpoint back
+	hcs.SyncEndpoints(map[types.NamespacedName]int{nsn: 11})
+	if len(hcs.services) != 1 {
+		t.Errorf("expected 1 service, got %d", len(hcs.services))
+	}
+	if hcs.services[nsn].endpoints != 11 {
+		t.Errorf("expected 18 endpoints, got %d", hcs.services[nsn].endpoints)
+	}
 	// sync nil endpoints
 	hcs.SyncEndpoints(nil)
 	if len(hcs.services) != 1 {
@@ -227,6 +235,7 @@ func TestServer(t *testing.T) {
 	// test the handlers
 	testHandler(hcs, nsn1, http.StatusServiceUnavailable, 0, t)
 	testHandler(hcs, nsn2, http.StatusServiceUnavailable, 0, t)
+	testHandler(hcs, nsn3, http.StatusServiceUnavailable, 0, t)
 
 	// sync endpoints
 	hcs.SyncEndpoints(map[types.NamespacedName]int{
@@ -298,6 +307,28 @@ func TestServer(t *testing.T) {
 	testHandler(hcs, nsn2, http.StatusOK, 3, t)
 	testHandler(hcs, nsn3, http.StatusOK, 7, t)
 	testHandler(hcs, nsn4, http.StatusOK, 6, t)
+
+	// sync endpoints, missing nsn2
+	hcs.SyncEndpoints(map[types.NamespacedName]int{
+		nsn3: 7,
+		nsn4: 6,
+	})
+	if len(hcs.services) != 3 {
+		t.Errorf("expected 3 services, got %d", len(hcs.services))
+	}
+	if hcs.services[nsn2].endpoints != 0 {
+		t.Errorf("expected 0 endpoints, got %d", hcs.services[nsn2].endpoints)
+	}
+	if hcs.services[nsn3].endpoints != 7 {
+		t.Errorf("expected 7 endpoints, got %d", hcs.services[nsn3].endpoints)
+	}
+	if hcs.services[nsn4].endpoints != 6 {
+		t.Errorf("expected 6 endpoints, got %d", hcs.services[nsn4].endpoints)
+	}
+	// test the handlers
+	testHandler(hcs, nsn2, http.StatusServiceUnavailable, 0, t)
+	testHandler(hcs, nsn3, http.StatusOK, 7, t)
+	testHandler(hcs, nsn4, http.StatusOK, 6, t)
 }
 
 func testHandler(hcs *server, nsn types.NamespacedName, status int, endpoints int, t *testing.T) {