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