From 4b8958b35b1b1fc33037f3b48fa4a50be27967b7 Mon Sep 17 00:00:00 2001
From: Preetha Appan <preetha@hashicorp.com>
Date: Tue, 18 Jul 2017 16:06:37 -0500
Subject: [PATCH] Removed unit test, added clarifying comment and returned a
 friendlier error message similar to the one in agent's AddService method

Fixes #3297
---
 agent/agent.go      |  5 +++-
 agent/agent_test.go | 67 ---------------------------------------------
 agent/local.go      | 24 ++++++++--------
 agent/local_test.go |  1 +
 4 files changed, 18 insertions(+), 79 deletions(-)

diff --git a/agent/agent.go b/agent/agent.go
index 0ae753b1c..c2529bb99 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -1787,7 +1787,10 @@ func (a *Agent) AddCheck(check *structs.HealthCheck, chkType *structs.CheckType,
 	}
 
 	// Add to the local state for anti-entropy
-	a.state.AddCheck(check, token)
+	err := a.state.AddCheck(check, token)
+	if err != nil {
+		return err
+	}
 
 	// Persist the check
 	if persist && !a.config.DevMode {
diff --git a/agent/agent_test.go b/agent/agent_test.go
index a47768243..c19ab9832 100644
--- a/agent/agent_test.go
+++ b/agent/agent_test.go
@@ -14,8 +14,6 @@ import (
 	"testing"
 	"time"
 
-	"sync"
-
 	"github.com/hashicorp/consul/agent/consul"
 	"github.com/hashicorp/consul/agent/consul/structs"
 	"github.com/hashicorp/consul/api"
@@ -628,71 +626,6 @@ func TestAgent_RemoveServiceRemovesAllChecks(t *testing.T) {
 	}
 }
 
-func TestAgent_AddCheck_RemoveServiceRace(t *testing.T) {
-	t.Parallel()
-	cfg := TestConfig()
-	cfg.NodeName = "node1"
-	a := NewTestAgent(t.Name(), cfg)
-	defer a.Shutdown()
-
-	// NOTE - trying a few times reproduces the race condition where a check can be added while a service is being deleted.
-	// This test may return a false positive if Addcheck happens to execute before RemoveService
-	for i := 0; i < 50; i++ {
-		addServiceAndChecks(a, t)
-		//try removing the service and adding a check for that service in two different go routines
-		var wg sync.WaitGroup
-		wg.Add(2)
-
-		chk1 := &structs.CheckType{CheckID: "chk1", Name: "chk1", TTL: time.Minute}
-		hchk1 := &structs.HealthCheck{Node: "node1", CheckID: "chk1", Name: "chk1", Status: "critical", ServiceID: "redis", ServiceName: "redis"}
-
-		go func() {
-			a.AddCheck(hchk1, chk1, false, "")
-			wg.Done()
-		}()
-
-		go func() {
-			a.RemoveService("redis", false)
-			wg.Done()
-		}()
-
-		wg.Wait()
-		checks := a.state.Checks()
-		if len(checks) > 0 {
-			t.Fatalf("Expected checks map to be empty")
-		}
-	}
-
-}
-
-func addServiceAndChecks(a *TestAgent, t *testing.T) {
-	svc := &structs.NodeService{ID: "redis", Service: "redis", Port: 8000}
-	chk1 := &structs.CheckType{CheckID: "chk1", Name: "chk1", TTL: time.Minute}
-	chk2 := &structs.CheckType{CheckID: "chk2", Name: "chk2", TTL: 2 * time.Minute}
-	hchk1 := &structs.HealthCheck{Node: "node1", CheckID: "chk1", Name: "chk1", Status: "critical", ServiceID: "redis", ServiceName: "redis"}
-	hchk2 := &structs.HealthCheck{Node: "node1", CheckID: "chk2", Name: "chk2", Status: "critical", ServiceID: "redis", ServiceName: "redis"}
-	// register service with chk1
-	if err := a.AddService(svc, []*structs.CheckType{chk1}, false, ""); err != nil {
-		t.Fatal("Failed to register service", err)
-	}
-	// verify chk1 exists
-	if a.state.Checks()["chk1"] == nil {
-		t.Fatal("Could not find health check chk1")
-	}
-	// update the service with chk2
-	if err := a.AddService(svc, []*structs.CheckType{chk2}, false, ""); err != nil {
-		t.Fatal("Failed to update service", err)
-	}
-	// check that both checks are there
-	if got, want := a.state.Checks()["chk1"], hchk1; !verify.Values(t, "", got, want) {
-		t.FailNow()
-	}
-	if got, want := a.state.Checks()["chk2"], hchk2; !verify.Values(t, "", got, want) {
-		t.FailNow()
-	}
-
-}
-
 func TestAgent_AddCheck(t *testing.T) {
 	t.Parallel()
 	cfg := TestConfig()
diff --git a/agent/local.go b/agent/local.go
index e2cb8e37a..decf7acdd 100644
--- a/agent/local.go
+++ b/agent/local.go
@@ -247,23 +247,25 @@ func (l *localState) checkToken(checkID types.CheckID) string {
 // AddCheck is used to add a health check to the local state.
 // This entry is persistent and the agent will make a best effort to
 // ensure it is registered
-func (l *localState) AddCheck(check *structs.HealthCheck, token string) {
+func (l *localState) AddCheck(check *structs.HealthCheck, token string) error {
 	// Set the node name
 	check.Node = l.config.NodeName
 
 	l.Lock()
 	defer l.Unlock()
-	allow := true
-	if check.ServiceID != "" { //if there is a serviceID associated with the check, make sure it exists
-		_, allow = l.services[check.ServiceID]
-	}
-	if allow {
-		l.checks[check.CheckID] = check
-		l.checkStatus[check.CheckID] = syncStatus{}
-		l.checkTokens[check.CheckID] = token
-		delete(l.checkCriticalTime, check.CheckID)
-		l.changeMade()
+
+	// if there is a serviceID associated with the check, make sure it exists before adding it
+	// NOTE - This logic may be moved to be handled within the Agent's Addcheck method after a refactor
+	if check.ServiceID != "" && l.services[check.ServiceID] == nil {
+		return fmt.Errorf("ServiceID %q does not exist", check.ServiceID)
 	}
+
+	l.checks[check.CheckID] = check
+	l.checkStatus[check.CheckID] = syncStatus{}
+	l.checkTokens[check.CheckID] = token
+	delete(l.checkCriticalTime, check.CheckID)
+	l.changeMade()
+	return nil
 }
 
 // RemoveCheck is used to remove a health check from the local state.
diff --git a/agent/local_test.go b/agent/local_test.go
index 2a6217f3a..e1ae506c0 100644
--- a/agent/local_test.go
+++ b/agent/local_test.go
@@ -1477,6 +1477,7 @@ func TestAgent_checkCriticalTime(t *testing.T) {
 
 	svc := &structs.NodeService{ID: "redis", Service: "redis", Port: 8000}
 	l.AddService(svc, "")
+
 	// Add a passing check and make sure it's not critical.
 	checkID := types.CheckID("redis:1")
 	chk := &structs.HealthCheck{