From fe4107019e0cafb37950e37fd26c2f7f17ba3b2e Mon Sep 17 00:00:00 2001
From: Matt Shean <mshean@yahoo-inc.com>
Date: Thu, 25 Feb 2016 19:18:20 -0800
Subject: [PATCH] add Timeout field to CheckMonitor

---
 command/agent/agent.go      |  1 +
 command/agent/check.go      |  7 ++++++-
 command/agent/check_test.go | 30 ++++++++++++++++++++++++++++++
 3 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/command/agent/agent.go b/command/agent/agent.go
index fe495fcf91..621d0ad0a5 100644
--- a/command/agent/agent.go
+++ b/command/agent/agent.go
@@ -968,6 +968,7 @@ func (a *Agent) AddCheck(check *structs.HealthCheck, chkType *CheckType, persist
 				CheckID:  check.CheckID,
 				Script:   chkType.Script,
 				Interval: chkType.Interval,
+				Timeout:  chkType.Timeout,
 				Logger:   a.logger,
 				ReapLock: &a.reapLock,
 			}
diff --git a/command/agent/check.go b/command/agent/check.go
index cc024b7fc1..515fade406 100644
--- a/command/agent/check.go
+++ b/command/agent/check.go
@@ -101,6 +101,7 @@ type CheckMonitor struct {
 	CheckID  string
 	Script   string
 	Interval time.Duration
+	Timeout  time.Duration
 	Logger   *log.Logger
 	ReapLock *sync.RWMutex
 
@@ -180,7 +181,11 @@ func (c *CheckMonitor) check() {
 		errCh <- cmd.Wait()
 	}()
 	go func() {
-		time.Sleep(30 * time.Second)
+		if c.Timeout > 0 {
+			time.Sleep(c.Timeout)
+		} else {
+			time.Sleep(30 * time.Second)
+		}
 		errCh <- fmt.Errorf("Timed out running check '%s'", c.Script)
 	}()
 	err = <-errCh
diff --git a/command/agent/check_test.go b/command/agent/check_test.go
index 053eb5148a..d2d5816e82 100644
--- a/command/agent/check_test.go
+++ b/command/agent/check_test.go
@@ -82,6 +82,36 @@ func TestCheckMonitor_BadCmd(t *testing.T) {
 	expectStatus(t, "foobarbaz", structs.HealthCritical)
 }
 
+func TestCheckMonitor_Timeout(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:   "sleep 1 && exit 0",
+		Interval: 10 * time.Millisecond,
+		Timeout:  5 * time.Millisecond,
+		Logger:   log.New(os.Stderr, "", log.LstdFlags),
+		ReapLock: &sync.RWMutex{},
+	}
+	check.Start()
+	defer check.Stop()
+
+	time.Sleep(50 * time.Millisecond)
+
+	// Should have at least 2 updates
+	if mock.updates["foo"] < 2 {
+		t.Fatalf("should have at least 2 updates %v", mock.updates)
+	}
+
+	if mock.state["foo"] != "critical" {
+		t.Fatalf("should be critical %v", mock.state)
+	}
+}
+
 func TestCheckMonitor_RandomStagger(t *testing.T) {
 	mock := &MockNotify{
 		state:   make(map[string]string),