From dd4fadfc9ed52d26b04e3cb109be079b120c5034 Mon Sep 17 00:00:00 2001
From: Ryan Uber <ru@ryanuber.com>
Date: Sat, 29 Nov 2014 12:25:01 -0800
Subject: [PATCH] agent: persist CheckType with health checks

---
 command/agent/agent.go      | 57 ++++++++++++++++++++-----------------
 command/agent/agent_test.go | 18 +++++++++---
 command/agent/check.go      |  7 +++++
 3 files changed, 52 insertions(+), 30 deletions(-)

diff --git a/command/agent/agent.go b/command/agent/agent.go
index bb2c9d63cb..facb260b7d 100644
--- a/command/agent/agent.go
+++ b/command/agent/agent.go
@@ -564,24 +564,29 @@ func (a *Agent) restoreServices() error {
 }
 
 // persistCheck saves a check definition to the local agent's state directory
-func (a *Agent) persistCheck(check *structs.HealthCheck) error {
+func (a *Agent) persistCheck(check *structs.HealthCheck, chkType *CheckType) error {
 	checkPath := filepath.Join(a.config.DataDir, checksDir, check.CheckID)
-	if _, err := os.Stat(checkPath); os.IsNotExist(err) {
-		encoded, err := json.Marshal(check)
-		if err != nil {
-			return nil
-		}
-		if err := os.MkdirAll(filepath.Dir(checkPath), 0700); err != nil {
-			return err
-		}
-		fh, err := os.OpenFile(checkPath, os.O_CREATE|os.O_WRONLY, 0600)
-		if err != nil {
-			return err
-		}
-		defer fh.Close()
-		if _, err := fh.Write(encoded); err != nil {
-			return err
-		}
+	if _, err := os.Stat(checkPath); !os.IsNotExist(err) {
+		return err
+	}
+
+	// Create the persisted check
+	p := persistedCheck{check, chkType}
+
+	encoded, err := json.Marshal(p)
+	if err != nil {
+		return nil
+	}
+	if err := os.MkdirAll(filepath.Dir(checkPath), 0700); err != nil {
+		return err
+	}
+	fh, err := os.OpenFile(checkPath, os.O_CREATE|os.O_WRONLY, 0600)
+	if err != nil {
+		return err
+	}
+	defer fh.Close()
+	if _, err := fh.Write(encoded); err != nil {
+		return err
 	}
 	return nil
 }
@@ -619,23 +624,23 @@ func (a *Agent) restoreChecks() error {
 			return err
 		}
 
-		var check *structs.HealthCheck
-		if err := json.Unmarshal(content, &check); err != nil {
+		var p persistedCheck
+		if err := json.Unmarshal(content, &p); err != nil {
 			return err
 		}
 
-		if _, ok := a.state.checks[check.CheckID]; ok {
+		if _, ok := a.state.checks[p.Check.CheckID]; ok {
 			// Purge previously persisted check. This allows config to be
 			// preferred over persisted checks from the API.
-			a.logger.Printf("[DEBUG] Check %s exists, not restoring", check.CheckID)
-			return a.purgeCheck(check.CheckID)
+			a.logger.Printf("[DEBUG] Check %s exists, not restoring", p.Check.CheckID)
+			return a.purgeCheck(p.Check.CheckID)
 		} else {
 			// Default check to critical to avoid placing potentially unhealthy
 			// services into the active pool
-			check.Status = structs.HealthCritical
+			p.Check.Status = structs.HealthCritical
 
-			a.logger.Printf("[DEBUG] Restored health check: %s", check.CheckID)
-			return a.AddCheck(check, nil, false)
+			a.logger.Printf("[DEBUG] Restored health check: %s", p.Check.CheckID)
+			return a.AddCheck(p.Check, p.ChkType, false)
 		}
 	})
 	return err
@@ -766,7 +771,7 @@ func (a *Agent) AddCheck(check *structs.HealthCheck, chkType *CheckType, persist
 
 	// Persist the check
 	if persist {
-		return a.persistCheck(check)
+		return a.persistCheck(check, chkType)
 	}
 
 	return nil
diff --git a/command/agent/agent_test.go b/command/agent/agent_test.go
index b7571a0855..8e78f24a44 100644
--- a/command/agent/agent_test.go
+++ b/command/agent/agent_test.go
@@ -534,11 +534,15 @@ func TestAgent_PersistCheck(t *testing.T) {
 		ServiceID:   "redis",
 		ServiceName: "redis",
 	}
+	chkType := &CheckType{
+		Script:   "/bin/true",
+		Interval: 10 * time.Second,
+	}
 
 	file := filepath.Join(agent.config.DataDir, checksDir, check.CheckID)
 
 	// Not persisted if not requested
-	if err := agent.AddCheck(check, nil, false); err != nil {
+	if err := agent.AddCheck(check, chkType, false); err != nil {
 		t.Fatalf("err: %v", err)
 	}
 	if _, err := os.Stat(file); err == nil {
@@ -546,7 +550,7 @@ func TestAgent_PersistCheck(t *testing.T) {
 	}
 
 	// Should persist if requested
-	if err := agent.AddCheck(check, nil, true); err != nil {
+	if err := agent.AddCheck(check, chkType, true); err != nil {
 		t.Fatalf("err: %v", err)
 	}
 
@@ -554,7 +558,8 @@ func TestAgent_PersistCheck(t *testing.T) {
 		t.Fatalf("err: %s", err)
 	}
 
-	expected, err := json.Marshal(check)
+	p := persistedCheck{check, chkType}
+	expected, err := json.Marshal(p)
 	if err != nil {
 		t.Fatalf("err: %s", err)
 	}
@@ -574,13 +579,18 @@ func TestAgent_PersistCheck(t *testing.T) {
 	}
 	defer agent2.Shutdown()
 
-	result, ok := agent2.state.checks[check.CheckID]
+	result, ok := agent2.state.checks[p.Check.CheckID]
 	if !ok {
 		t.Fatalf("bad: %#v", agent2.state.checks)
 	}
 	if result.Status != structs.HealthCritical {
 		t.Fatalf("bad: %#v", result)
 	}
+
+	// Should have restored the monitor
+	if _, ok := agent2.checkMonitors[p.Check.CheckID]; !ok {
+		t.Fatalf("bad: %#v", agent2.checkMonitors)
+	}
 }
 
 func TestAgent_PurgeCheck(t *testing.T) {
diff --git a/command/agent/check.go b/command/agent/check.go
index a2cb588dd6..180f6db13a 100644
--- a/command/agent/check.go
+++ b/command/agent/check.go
@@ -234,3 +234,10 @@ func (c *CheckTTL) SetStatus(status, output string) {
 	c.Notify.UpdateCheck(c.CheckID, status, output)
 	c.timer.Reset(c.TTL)
 }
+
+// persistedCheck is used to serialize a check and write it to disk
+// so that it may be restored later on.
+type persistedCheck struct {
+	Check   *structs.HealthCheck
+	ChkType *CheckType
+}