Merge pull request #35 from InVisionApp/dselans/leak-fix

fixed goroutine leak after .Stop() is called

health.go

@@ -79,7 +79,7 @@ type Health struct {
 	configs    []*Config
 	states     map[string]State
 	statesLock sync.Mutex
-	tickers    map[string]*time.Ticker // contains map of actively running tickers
+	runners    map[string]chan struct{} // contains map of active runners w/ a stop channel
 }
 
 // New returns a new instance of the Health struct.
@@ -88,7 +88,7 @@ func New() *Health {
 		Logger:     loggers.NewBasic(),
 		configs:    make([]*Config, 0),
 		states:     make(map[string]State, 0),
-		tickers:    make(map[string]*time.Ticker, 0),
+		runners:    make(map[string]chan struct{}, 0),
 		active:     newBool(),
 		failed:     newBool(), // init as false
 		statesLock: sync.Mutex{},
@@ -136,12 +136,13 @@ func (h *Health) Start() error {
 	for _, c := range h.configs {
 		h.Logger.Debug("Starting checker", map[string]interface{}{"name": c.Name})
 		ticker := time.NewTicker(c.Interval)
+		stop := make(chan struct{})
 
-		if err := h.startRunner(c, ticker); err != nil {
+		if err := h.startRunner(c, ticker, stop); err != nil {
 			return fmt.Errorf("Unable to create healthcheck runner '%v': %v", c.Name, err)
 		}
 
-		h.tickers[c.Name] = ticker
+		h.runners[c.Name] = stop
 	}
 
 	// Checkers are now actively running
@@ -157,13 +158,16 @@ func (h *Health) Stop() error {
 		return ErrAlreadyStopped
 	}
 
-	for name, ticker := range h.tickers {
+	for name, stop := range h.runners {
 		h.Logger.Debug("Stopping checker", map[string]interface{}{"name": name})
-		ticker.Stop()
+		close(stop)
 	}
 
-	// Reset ticker map
-	h.tickers = make(map[string]*time.Ticker, 0)
+	// Reset runner map
+	h.runners = make(map[string]chan struct{}, 0)
+
+	// Reset states
+	h.safeResetStates()
 
 	return nil
 }
@@ -186,7 +190,7 @@ func (h *Health) Failed() bool {
 	return h.failed.val()
 }
 
-func (h *Health) startRunner(cfg *Config, ticker *time.Ticker) error {
+func (h *Health) startRunner(cfg *Config, ticker *time.Ticker, stop <-chan struct{}) error {
 
 	// function to execute and collect check data
 	checkFunc := func() {
@@ -220,12 +224,20 @@ func (h *Health) startRunner(cfg *Config, ticker *time.Ticker) error {
 	}
 
 	go func() {
+		defer ticker.Stop()
+
 		// execute once so that it is immediate
 		checkFunc()
 
 		// all following executions
-		for range ticker.C {
-			checkFunc()
+	RunLoop:
+		for {
+			select {
+			case <-ticker.C:
+				checkFunc()
+			case <-stop:
+				break RunLoop
+			}
 		}
 
 		h.Logger.Debug("Checker exiting", map[string]interface{}{"name": cfg.Name})
@@ -234,6 +246,12 @@ func (h *Health) startRunner(cfg *Config, ticker *time.Ticker) error {
 	return nil
 }
 
+func (h *Health) safeResetStates() {
+	h.statesLock.Lock()
+	defer h.statesLock.Unlock()
+	h.states = make(map[string]State, 0)
+}
+
 func (h *Health) safeUpdateState(stateEntry *State) {
 	// update states here
 	h.statesLock.Lock()

health_shared_test.go

@@ -45,8 +45,8 @@ func setupRunners(cfgs []*Config, logger loggers.ILogger) (*Health, []*Config, e
 	}
 
 	// Correct number of runners/tickers were created
-	if len(h.tickers) != len(cfgs) {
-		return nil, nil, fmt.Errorf("Start() did not create the expected number of tickers")
+	if len(h.runners) != len(cfgs) {
+		return nil, nil, fmt.Errorf("Start() did not create the expected number of runners")
 	}
 
 	return h, cfgs, nil

health_test.go

@@ -24,7 +24,7 @@ func TestNew(t *testing.T) {
 
 		Expect(h.configs).ToNot(BeNil())
 		Expect(h.states).ToNot(BeNil())
-		Expect(h.tickers).ToNot(BeNil())
+		Expect(h.runners).ToNot(BeNil())
 	})
 }
 
@@ -274,12 +274,12 @@ func TestStart(t *testing.T) {
 
 		err = h.Start()
 		Expect(err).ToNot(HaveOccurred())
-		// Correct number of runners/tickers were created
-		Expect(len(h.tickers)).To(Equal(2))
+		// Correct number of runners were created
+		Expect(len(h.runners)).To(Equal(2))
 
-		// Tickers are created (and saved) based on their name
+		// Runners are created (and saved) based on their name
 		for _, v := range cfgs {
-			Expect(h.tickers).To(HaveKey(v.Name))
+			Expect(h.runners).To(HaveKey(v.Name))
 		}
 
 		// This is pretty brittle - will update if this is causing random test failures
@@ -325,20 +325,33 @@ func TestStop(t *testing.T) {
 		Expect(err).ToNot(HaveOccurred())
 		Expect(h).ToNot(BeNil())
 
+		// A bit brittle, but it'll do
+		time.Sleep(time.Duration(15) * time.Millisecond)
+		Expect(len(h.states)).To(Equal(2))
+
 		err = h.Stop()
 		Expect(err).ToNot(HaveOccurred())
 
-		// Tickers map should be reset
-		Expect(h.tickers).To(BeEmpty())
+		// Wait a bit to ensure goroutines have exited
+		time.Sleep(15 * time.Millisecond)
+
+		// Runners map should be reset
+		Expect(h.runners).To(BeEmpty())
 
 		// Ensure that logger captured the start and stop messages
-		Expect(fakeLogger.DebugCallCount()).To(Equal(4))
+		Expect(fakeLogger.DebugCallCount()).To(Equal(6))
 
 		for i := range cfgs {
 			// 3rd and 4th message should indicate goroutine exit
 			msg, _ := fakeLogger.DebugArgsForCall(i + 2)
 			Expect(msg).To(Equal("Stopping checker"))
+
+			exitMsg, _ := fakeLogger.DebugArgsForCall(i + 4)
+			Expect(exitMsg).To(Equal("Checker exiting"))
 		}
+
+		// Expect state map to be reset
+		Expect(len(h.states)).To(Equal(0))
 	})
 
 	t.Run("Should error if healthcheck is not running", func(t *testing.T) {