allow periodic flushes to Stop being emitted. (#146)

* allow periodic flushes to Stop being emitted.

This introduces a new Stop() method and changes the behavior of the
Start() method on a store just slightly.

store.Start(), if called while a periodic flush goroutine is currently
running, causes the existing goroutine to stop and be replaced with a
new goroutine. That is, the caller can now change the rate of flushing
with calls such as:

store.Start(time.NewTicker(1*time.Second))
store.Start(time.NewTicker(100*time.Second))

store.Stop() will stop any existing running periodic flush
goroutine. If called multiple times, it is a no-op.

* close instead of sending a bool

* use the values from the closure

stats.go

@@ -40,9 +40,19 @@ type Store interface {
 	// and flush all the Counters and Gauges registered with it.
 	Flush()
 
-	// Start a timer for periodic stat Flushes.
+	// Start a timer for periodic stat flushes.
+	//
+	// If Start is called multiple times, the previous ticker is
+	// stopped and a new replacement ticker is started. It is
+	// equivalent to calling Stop() and then Start() with the new
+	// ticker.
 	Start(*time.Ticker)
 
+	// Stop stops the running periodic flush timer.
+	//
+	// If no periodic flush is currently active, this is a no-op.
+	Stop()
+
 	// Add a StatGenerator to the Store that programatically generates stats.
 	AddStatGenerator(StatGenerator)
 	Scope
@@ -331,18 +341,78 @@ type statStore struct {
 	gauges   sync.Map
 	timers   sync.Map
 
-	genMtx         sync.RWMutex
+	mu             sync.RWMutex
 	statGenerators []StatGenerator
+	stop           chan bool
+	wg             *sync.WaitGroup
 
 	sink Sink
 }
 
+func (s *statStore) Start(ticker *time.Ticker) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	// if there is already a stop channel allocated, that means there
+	// is a ticker running - we will replace the ticker now.
+	if s.stop != nil {
+		s.stopLocked()
+	}
+
+	stopChan := make(chan bool, 1)
+	wg := &sync.WaitGroup{}
+	wg.Add(1)
+
+	s.stop = stopChan
+	s.wg = wg
+
+	go func() {
+		defer wg.Done()
+		for {
+			select {
+			case <-ticker.C:
+				s.Flush()
+			case <-stopChan:
+				return
+			}
+		}
+	}()
+}
+
+// stopLocked is the core of the stop implementation, but without a
+// lock.
+func (s *statStore) stopLocked() {
+	// if the stop channel is nil, there is no ticker running, so this
+	// is a no-op.
+	if s.stop == nil {
+		return
+	}
+
+	// close to make the flush goroutine stop
+	close(s.stop)
+
+	// wait for the flush goroutine to fully stop
+	s.wg.Wait()
+
+	// nil out the stop channel
+	s.stop = nil
+
+	// nil out the wait group for tidyness
+	s.wg = nil
+}
+
+func (s *statStore) Stop() {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	s.stopLocked()
+}
+
 func (s *statStore) Flush() {
-	s.genMtx.RLock()
+	s.mu.RLock()
 	for _, g := range s.statGenerators {
 		g.GenerateStats()
 	}
-	s.genMtx.RUnlock()
+	s.mu.RUnlock()
 
 	s.counters.Range(func(key, v interface{}) bool {
 		// do not flush counters that are set to zero
@@ -363,20 +433,10 @@ func (s *statStore) Flush() {
 	}
 }
 
-func (s *statStore) Start(ticker *time.Ticker) {
-	s.run(ticker)
-}
-
 func (s *statStore) AddStatGenerator(statGenerator StatGenerator) {
-	s.genMtx.Lock()
+	s.mu.Lock()
+	defer s.mu.Unlock()
 	s.statGenerators = append(s.statGenerators, statGenerator)
-	s.genMtx.Unlock()
-}
-
-func (s *statStore) run(ticker *time.Ticker) {
-	for range ticker.C {
-		s.Flush()
-	}
 }
 
 func (s *statStore) Store() Store {

stats_test.go

@@ -43,6 +43,72 @@ func TestStats(t *testing.T) {
 	wg.Wait()
 }
 
+// TestStatsStartMultipleTimes ensures starting a periodic flush twice
+// works as expected.
+func TestStatsStartMultipleTimes(t *testing.T) {
+	sink := &testStatSink{}
+	store := NewStore(sink, true)
+
+	// we check to make sure the two stop channels are different,
+	// that's the best we can do.
+	store.Start(time.NewTicker(1 * time.Minute))
+
+	// grab the stop channel
+	realStore := store.(*statStore)
+	stopChan1 := realStore.stop
+
+	store.Start(time.NewTicker(10 * time.Hour))
+
+	// grab the new stop channel
+	stopChan2 := realStore.stop
+
+	if stopChan1 == stopChan2 {
+		t.Error("two stop channels are the same")
+	}
+}
+
+// TestStatsStartStop ensures starting a periodic flush can be
+// stopped.
+func TestStatsStartStop(t *testing.T) {
+	sink := &testStatSink{}
+	store := NewStore(sink, true)
+
+	store.Start(time.NewTicker(1 * time.Minute))
+	store.Stop()
+
+	realStore := store.(*statStore)
+
+	// we check to make sure the two stop channel is nil'ed out, that
+	// is the best we can do to avoid flakey time based tests.
+	if realStore.stop != nil {
+		t.Errorf("expected stop channel to be nil")
+	}
+}
+
+// TestStatsStartStopMultipleTimes ensures starting a periodic flush
+// can be stopped, even if called multiple times.
+func TestStatsStartStopMultipleTimes(t *testing.T) {
+	sink := &testStatSink{}
+	store := NewStore(sink, true)
+
+	// ensure we can call it if no ticker was ever started
+	store.Stop()
+
+	// start one, and stop many times.
+	store.Start(time.NewTicker(1 * time.Minute))
+	store.Stop()
+	store.Stop()
+	store.Stop()
+
+	realStore := store.(*statStore)
+
+	// we check to make sure the two stop channel is nil'ed out, that
+	// is the best we can do to avoid flakey time based tests.
+	if realStore.stop != nil {
+		t.Errorf("expected stop channel to be nil")
+	}
+}
+
 // Ensure timers and timespans are working
 func TestTimer(t *testing.T) {
 	testDuration := time.Duration(9800000)