agent: persist CheckType with health checks

command/agent/agent.go

@@ -564,10 +564,16 @@ 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 := 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
 	}
@@ -582,7 +588,6 @@ func (a *Agent) persistCheck(check *structs.HealthCheck) error {
 	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

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) {

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
+}