From 8ef01236e1ce9730a42e4539e84a5d1bd33ff4e0 Mon Sep 17 00:00:00 2001
From: Ryan Uber <ru@ryanuber.com>
Date: Tue, 5 May 2015 22:08:03 -0700
Subject: [PATCH] agent: allow persisted services to be updated on disk

---
 command/agent/agent.go      | 43 +++++++++++++----------------
 command/agent/agent_test.go | 54 +++++++++++++++++++++++++++++++++----
 2 files changed, 68 insertions(+), 29 deletions(-)

diff --git a/command/agent/agent.go b/command/agent/agent.go
index 0163e44ad4..17c7ca5ca5 100644
--- a/command/agent/agent.go
+++ b/command/agent/agent.go
@@ -535,26 +535,24 @@ func (a *Agent) ResumeSync() {
 // persistService saves a service definition to a JSON file in the data dir
 func (a *Agent) persistService(service *structs.NodeService) error {
 	svcPath := filepath.Join(a.config.DataDir, servicesDir, stringHash(service.ID))
-	if _, err := os.Stat(svcPath); os.IsNotExist(err) {
-		wrapped := persistedService{
-			Token:   a.state.ServiceToken(service.ID),
-			Service: service,
-		}
-		encoded, err := json.Marshal(wrapped)
-		if err != nil {
-			return nil
-		}
-		if err := os.MkdirAll(filepath.Dir(svcPath), 0700); err != nil {
-			return err
-		}
-		fh, err := os.OpenFile(svcPath, os.O_CREATE|os.O_WRONLY, 0600)
-		if err != nil {
-			return err
-		}
-		defer fh.Close()
-		if _, err := fh.Write(encoded); err != nil {
-			return err
-		}
+	wrapped := persistedService{
+		Token:   a.state.ServiceToken(service.ID),
+		Service: service,
+	}
+	encoded, err := json.Marshal(wrapped)
+	if err != nil {
+		return nil
+	}
+	if err := os.MkdirAll(filepath.Dir(svcPath), 0700); err != nil {
+		return err
+	}
+	fh, err := os.OpenFile(svcPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
+	if err != nil {
+		return err
+	}
+	defer fh.Close()
+	if _, err := fh.Write(encoded); err != nil {
+		return err
 	}
 	return nil
 }
@@ -571,9 +569,6 @@ func (a *Agent) purgeService(serviceID string) error {
 // persistCheck saves a check definition to the local agent's state directory
 func (a *Agent) persistCheck(check *structs.HealthCheck, chkType *CheckType) error {
 	checkPath := filepath.Join(a.config.DataDir, checksDir, stringHash(check.CheckID))
-	if _, err := os.Stat(checkPath); !os.IsNotExist(err) {
-		return err
-	}
 
 	// Create the persisted check
 	wrapped := persistedCheck{
@@ -589,7 +584,7 @@ func (a *Agent) persistCheck(check *structs.HealthCheck, chkType *CheckType) err
 	if err := os.MkdirAll(filepath.Dir(checkPath), 0700); err != nil {
 		return err
 	}
-	fh, err := os.OpenFile(checkPath, os.O_CREATE|os.O_WRONLY, 0600)
+	fh, err := os.OpenFile(checkPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0600)
 	if err != nil {
 		return err
 	}
diff --git a/command/agent/agent_test.go b/command/agent/agent_test.go
index 1a2aed3a1b..a2d3fdc092 100644
--- a/command/agent/agent_test.go
+++ b/command/agent/agent_test.go
@@ -533,11 +533,9 @@ func TestAgent_PersistService(t *testing.T) {
 	if err := agent.AddService(svc, nil, true, "mytoken"); err != nil {
 		t.Fatalf("err: %v", err)
 	}
-
 	if _, err := os.Stat(file); err != nil {
 		t.Fatalf("err: %s", err)
 	}
-
 	expected, err := json.Marshal(persistedService{
 		Token:   "mytoken",
 		Service: svc,
@@ -552,6 +550,26 @@ func TestAgent_PersistService(t *testing.T) {
 	if !bytes.Equal(expected, content) {
 		t.Fatalf("bad: %s", string(content))
 	}
+
+	// Updates service definition on disk
+	svc.Port = 8001
+	if err := agent.AddService(svc, nil, true, "mytoken"); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	expected, err = json.Marshal(persistedService{
+		Token:   "mytoken",
+		Service: svc,
+	})
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	content, err = ioutil.ReadFile(file)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if !bytes.Equal(expected, content) {
+		t.Fatalf("bad: %s", string(content))
+	}
 	agent.Shutdown()
 
 	// Should load it back during later start
@@ -561,12 +579,16 @@ func TestAgent_PersistService(t *testing.T) {
 	}
 	defer agent2.Shutdown()
 
-	if _, ok := agent2.state.services[svc.ID]; !ok {
+	restored, ok := agent2.state.services[svc.ID]
+	if !ok {
 		t.Fatalf("bad: %#v", agent2.state.services)
 	}
 	if agent2.state.serviceTokens[svc.ID] != "mytoken" {
 		t.Fatalf("bad: %#v", agent2.state.services[svc.ID])
 	}
+	if restored.Port != 8001 {
+		t.Fatalf("bad: %#v", restored)
+	}
 }
 
 func TestAgent_persistedService_compat(t *testing.T) {
@@ -731,11 +753,9 @@ func TestAgent_PersistCheck(t *testing.T) {
 	if err := agent.AddCheck(check, chkType, true, "mytoken"); err != nil {
 		t.Fatalf("err: %v", err)
 	}
-
 	if _, err := os.Stat(file); err != nil {
 		t.Fatalf("err: %s", err)
 	}
-
 	expected, err := json.Marshal(persistedCheck{
 		Check:   check,
 		ChkType: chkType,
@@ -751,6 +771,27 @@ func TestAgent_PersistCheck(t *testing.T) {
 	if !bytes.Equal(expected, content) {
 		t.Fatalf("bad: %s", string(content))
 	}
+
+	// Updates the check definition on disk
+	check.Name = "mem1"
+	if err := agent.AddCheck(check, chkType, true, "mytoken"); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+	expected, err = json.Marshal(persistedCheck{
+		Check:   check,
+		ChkType: chkType,
+		Token:   "mytoken",
+	})
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	content, err = ioutil.ReadFile(file)
+	if err != nil {
+		t.Fatalf("err: %s", err)
+	}
+	if !bytes.Equal(expected, content) {
+		t.Fatalf("bad: %s", string(content))
+	}
 	agent.Shutdown()
 
 	// Should load it back during later start
@@ -767,6 +808,9 @@ func TestAgent_PersistCheck(t *testing.T) {
 	if result.Status != structs.HealthCritical {
 		t.Fatalf("bad: %#v", result)
 	}
+	if result.Name != "mem1" {
+		t.Fatalf("bad: %#v", result)
+	}
 
 	// Should have restored the monitor
 	if _, ok := agent2.checkMonitors[check.CheckID]; !ok {