agent: Limit health check output to 4K. Fixes #83.

command/agent/check.go

@@ -1,8 +1,8 @@
 package agent
 
 import (
-	"bytes"
 	"fmt"
+	"github.com/armon/circbuf"
 	"github.com/hashicorp/consul/consul/structs"
 	"log"
 	"os/exec"
@@ -16,6 +16,11 @@ const (
 	// Do not allow for a interval below this value.
 	// Otherwise we risk fork bombing a system.
 	MinInterval = time.Second
+
+	// Limit the size of a check's output to the
+	// last CheckBufSize. Prevents an enormous buffer
+	// from being captured
+	CheckBufSize = 4 * 1024 // 4KB
 )
 
 // CheckType is used to create either the CheckMonitor
@@ -115,9 +120,9 @@ func (c *CheckMonitor) check() {
 	cmd := exec.Command(shell, flag, c.Script)
 
 	// Collect the output
-	var output bytes.Buffer
-	cmd.Stdout = &output
-	cmd.Stderr = &output
+	output, _ := circbuf.NewBuffer(CheckBufSize)
+	cmd.Stdout = output
+	cmd.Stderr = output
 
 	// Start the check
 	if err := cmd.Start(); err != nil {
@@ -137,7 +142,13 @@ func (c *CheckMonitor) check() {
 	}()
 	err := <-errCh
 
+	// Get the output, add a message about truncation
 	outputStr := string(output.Bytes())
+	if output.TotalWritten() > output.Size() {
+		outputStr = fmt.Sprintf("Captured %d of %d bytes\n...\n%s",
+			output.Size(), output.TotalWritten(), outputStr)
+	}
+
 	c.Logger.Printf("[DEBUG] agent: check '%s' script '%s' output: %s",
 		c.CheckID, c.Script, outputStr)
 

command/agent/check_test.go

@@ -11,18 +11,21 @@ import (
 type MockNotify struct {
 	state   map[string]string
 	updates map[string]int
+	output  map[string]string
 }
 
-func (m *MockNotify) UpdateCheck(id, status, note string) {
+func (m *MockNotify) UpdateCheck(id, status, output string) {
 	m.state[id] = status
 	old := m.updates[id]
 	m.updates[id] = old + 1
+	m.output[id] = output
 }
 
 func expectStatus(t *testing.T, script, status string) {
 	mock := &MockNotify{
 		state:   make(map[string]string),
 		updates: make(map[string]int),
+		output:  make(map[string]string),
 	}
 	check := &CheckMonitor{
 		Notify:   mock,
@@ -62,10 +65,35 @@ func TestCheckMonitor_BadCmd(t *testing.T) {
 	expectStatus(t, "foobarbaz", structs.HealthCritical)
 }
 
+func TestCheckMonitor_LimitOutput(t *testing.T) {
+	mock := &MockNotify{
+		state:   make(map[string]string),
+		updates: make(map[string]int),
+		output:  make(map[string]string),
+	}
+	check := &CheckMonitor{
+		Notify:   mock,
+		CheckID:  "foo",
+		Script:   "dd if=/dev/urandom bs=8192 count=10",
+		Interval: 25 * time.Millisecond,
+		Logger:   log.New(os.Stderr, "", log.LstdFlags),
+	}
+	check.Start()
+	defer check.Stop()
+
+	time.Sleep(50 * time.Millisecond)
+
+	// Allow for extra bytes for the truncation message
+	if len(mock.output["foo"]) > CheckBufSize+100 {
+		t.Fatalf("output size is too long")
+	}
+}
+
 func TestCheckTTL(t *testing.T) {
 	mock := &MockNotify{
 		state:   make(map[string]string),
 		updates: make(map[string]int),
+		output:  make(map[string]string),
 	}
 	check := &CheckTTL{
 		Notify:  mock,