From a68ecd0fe7cfccc7fc8f585daa6c1b4d6dd9ad6c Mon Sep 17 00:00:00 2001
From: Pierre Souchay <p.souchay@criteo.com>
Date: Wed, 26 Jun 2019 00:05:24 +0200
Subject: [PATCH] Properly truncate TTL checks

Check that OutputMaxSize is properly taken into account.

BUGFIX: properly update local state on update of TTL
---
 agent/agent.go               | 21 ++++++++++++++-------
 agent/agent_endpoint_test.go |  2 +-
 agent/agent_test.go          | 18 ++++++++++++++++--
 agent/checks/check.go        |  4 +++-
 4 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/agent/agent.go b/agent/agent.go
index 0ed64c82866d..e2f0985ebe8b 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -2259,10 +2259,11 @@ func (a *Agent) addCheck(check *structs.HealthCheck, chkType *structs.CheckType,
 			}
 
 			ttl := &checks.CheckTTL{
-				Notify:  a.State,
-				CheckID: check.CheckID,
-				TTL:     chkType.TTL,
-				Logger:  a.logger,
+				Notify:        a.State,
+				CheckID:       check.CheckID,
+				TTL:           chkType.TTL,
+				Logger:        a.logger,
+				OutputMaxSize: maxOutputSize,
 			}
 
 			// Restore persisted state, if any
@@ -2878,8 +2879,14 @@ func (a *Agent) updateTTLCheck(checkID types.CheckID, status, output string) err
 	}
 
 	// Set the status through CheckTTL to reset the TTL.
-	check.SetStatus(status, output)
-
+	outputTruncated := check.SetStatus(status, output)
+	stateCheck := a.State.Checks()[checkID]
+	if stateCheck != nil {
+		stateCheck.Output = outputTruncated
+		stateCheck.Status = status
+	} else {
+		a.logger.Printf("[WARN] agent: Unexpected missing state for check %v", checkID)
+	}
 	// We don't write any files in dev mode so bail here.
 	if a.config.DataDir == "" {
 		return nil
@@ -2887,7 +2894,7 @@ func (a *Agent) updateTTLCheck(checkID types.CheckID, status, output string) err
 
 	// Persist the state so the TTL check can come up in a good state after
 	// an agent restart, especially with long TTL values.
-	if err := a.persistCheckState(check, status, output); err != nil {
+	if err := a.persistCheckState(check, status, outputTruncated); err != nil {
 		return fmt.Errorf("failed persisting state for check %q: %s", checkID, err)
 	}
 
diff --git a/agent/agent_endpoint_test.go b/agent/agent_endpoint_test.go
index 44f619809397..18238dbe5f8a 100644
--- a/agent/agent_endpoint_test.go
+++ b/agent/agent_endpoint_test.go
@@ -2360,7 +2360,7 @@ func TestAgent_UpdateCheck(t *testing.T) {
 		// isn't super brittle.
 		state := a.State.Checks()["test"]
 		if state.Status != api.HealthPassing || len(state.Output) > 2*checks.DefaultBufSize {
-			t.Fatalf("bad: %v", state)
+			t.Fatalf("bad: %v, (len:=%d)", state, len(state.Output))
 		}
 	})
 
diff --git a/agent/agent_test.go b/agent/agent_test.go
index fecd1c208d12..226ee32355bc 100644
--- a/agent/agent_test.go
+++ b/agent/agent_test.go
@@ -1405,7 +1405,7 @@ func TestAgent_updateTTLCheck(t *testing.T) {
 	t.Parallel()
 	a := NewTestAgent(t, t.Name(), "")
 	defer a.Shutdown()
-
+	checkBufSize := 100
 	health := &structs.HealthCheck{
 		Node:    "foo",
 		CheckID: "mem",
@@ -1413,7 +1413,8 @@ func TestAgent_updateTTLCheck(t *testing.T) {
 		Status:  api.HealthCritical,
 	}
 	chk := &structs.CheckType{
-		TTL: 15 * time.Second,
+		TTL:           15 * time.Second,
+		OutputMaxSize: checkBufSize,
 	}
 
 	// Add check and update it.
@@ -1433,6 +1434,19 @@ func TestAgent_updateTTLCheck(t *testing.T) {
 	if status.Output != "foo" {
 		t.Fatalf("bad: %v", status)
 	}
+
+	if err := a.updateTTLCheck("mem", api.HealthCritical, strings.Repeat("--bad-- ", 5*checkBufSize)); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Ensure we have a check mapping.
+	status = a.State.Checks()["mem"]
+	if status.Status != api.HealthCritical {
+		t.Fatalf("bad: %v", status)
+	}
+	if len(status.Output) > checkBufSize*2 {
+		t.Fatalf("bad: %v", len(status.Output))
+	}
 }
 
 func TestAgent_PersistService(t *testing.T) {
diff --git a/agent/checks/check.go b/agent/checks/check.go
index 9997c0f807ca..8b6c4c971b3a 100644
--- a/agent/checks/check.go
+++ b/agent/checks/check.go
@@ -281,7 +281,8 @@ func (c *CheckTTL) getExpiredOutput() string {
 
 // SetStatus is used to update the status of the check,
 // and to renew the TTL. If expired, TTL is restarted.
-func (c *CheckTTL) SetStatus(status, output string) {
+// output is returned (might be truncated)
+func (c *CheckTTL) SetStatus(status, output string) string {
 	c.Logger.Printf("[DEBUG] agent: Check %q status is now %s", c.CheckID, status)
 	c.Notify.UpdateCheck(c.CheckID, status, output)
 	total := len(output)
@@ -295,6 +296,7 @@ func (c *CheckTTL) SetStatus(status, output string) {
 	c.lastOutputLock.Unlock()
 
 	c.timer.Reset(c.TTL)
+	return output
 }
 
 // CheckHTTP is used to periodically make an HTTP request to