From 5d63160bb66df4699bce194b4c4a820a5d7735db Mon Sep 17 00:00:00 2001
From: Ryan Uber <ru@ryanuber.com>
Date: Sat, 29 Aug 2015 11:10:32 -0700
Subject: [PATCH] consul/state: node checks are removed with their associated
 nodes

---
 consul/state/state_store.go      | 22 ++++++++++++++++++++--
 consul/state/state_store_test.go | 25 ++++++++++++++++++++++++-
 2 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/consul/state/state_store.go b/consul/state/state_store.go
index 5a25dcb40..33e8d2d8e 100644
--- a/consul/state/state_store.go
+++ b/consul/state/state_store.go
@@ -173,7 +173,7 @@ func (s *StateStore) deleteNodeTxn(idx uint64, nodeID string, tx *memdb.Txn) err
 		return fmt.Errorf("node lookup failed: %s", err)
 	}
 
-	// Delete all services associated with the node
+	// Delete all services associated with the node and update the service index
 	services, err := tx.Get("services", "node", nodeID)
 	if err != nil {
 		return fmt.Errorf("failed service lookup: %s", err)
@@ -181,9 +181,27 @@ func (s *StateStore) deleteNodeTxn(idx uint64, nodeID string, tx *memdb.Txn) err
 	for service := services.Next(); service != nil; service = services.Next() {
 		svc := service.(*structs.ServiceNode)
 		if err := s.deleteNodeServiceTxn(idx, nodeID, svc.ServiceID, tx); err != nil {
-			return fmt.Errorf("failed removing node service: %s", err)
+			return err
 		}
 	}
+	if err := tx.Insert("index", &IndexEntry{"services", idx}); err != nil {
+		return fmt.Errorf("failed updating index: %s", err)
+	}
+
+	// Delete all checks associated with the node and update the check index
+	checks, err := tx.Get("checks", "node", nodeID)
+	if err != nil {
+		return fmt.Errorf("failed check lookup: %s", err)
+	}
+	for check := checks.Next(); check != nil; check = checks.Next() {
+		chk := check.(*structs.HealthCheck)
+		if err := s.deleteCheckTxn(idx, nodeID, chk.CheckID, tx); err != nil {
+			return err
+		}
+	}
+	if err := tx.Insert("index", &IndexEntry{"checks", idx}); err != nil {
+		return fmt.Errorf("failed updating index: %s", err)
+	}
 
 	// Delete the node and update the index
 	if err := tx.Delete("nodes", node); err != nil {
diff --git a/consul/state/state_store_test.go b/consul/state/state_store_test.go
index 8fc09f88d..c5533e1b6 100644
--- a/consul/state/state_store_test.go
+++ b/consul/state/state_store_test.go
@@ -154,12 +154,26 @@ func TestStateStore_DeleteNode(t *testing.T) {
 		t.Fatalf("bad: %#v (err: %s)", services.Services, err)
 	}
 
+	// Register a check with the node service
+	chk := &structs.HealthCheck{
+		Node:    "node1",
+		CheckID: "check1",
+	}
+	if err := s.EnsureCheck(3, chk); err != nil {
+		t.Fatalf("err: %s", err)
+	}
+
+	// Check exists
+	if checks, err := s.NodeChecks("node1"); err != nil || len(checks) != 1 {
+		t.Fatalf("bad: %#v (err: %s)", checks, err)
+	}
+
 	// Delete the node
 	if err := s.DeleteNode(3, "node1"); err != nil {
 		t.Fatalf("err: %s", err)
 	}
 
-	// The node and service are gone and the index was updated
+	// The node was removed
 	if n, err := s.GetNode("node1"); err != nil || n != nil {
 		t.Fatalf("bad: %#v (err: %#v)", node, err)
 	}
@@ -176,6 +190,15 @@ func TestStateStore_DeleteNode(t *testing.T) {
 		t.Fatalf("bad: %#v", s)
 	}
 
+	// Associated health check was removed.
+	checks, err := tx.Get("checks", "id", "node1", "check1")
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if c := checks.Next(); c != nil {
+		t.Fatalf("bad: %#v", c)
+	}
+
 	// Indexes were updated.
 	for _, tbl := range []string{"nodes", "services"} {
 		if idx := s.maxIndex(tbl); idx != 3 {