Fix race in Flush and AddStatGenerator (#8)

stats.go

@@ -299,13 +299,13 @@ type statStore struct {
 }
 
 func (s *statStore) Flush() {
+	s.Lock()
+	defer s.Unlock()
+
 	for _, g := range s.statGenerators {
 		g.GenerateStats()
 	}
 
-	s.Lock()
-	defer s.Unlock()
-
 	for name, cv := range s.counters {
 		value := cv.latch()
 
@@ -328,6 +328,9 @@ func (s *statStore) Start(ticker *time.Ticker) {
 }
 
 func (s *statStore) AddStatGenerator(statGenerator StatGenerator) {
+	s.Lock()
+	defer s.Unlock()
+
 	s.statGenerators = append(s.statGenerators, statGenerator)
 }
 

stats_test.go

@@ -0,0 +1,31 @@
+package stats
+
+import (
+	"sync"
+	"testing"
+)
+
+// Ensure flushing and adding generators does not race
+func TestStats(t *testing.T) {
+	sink := &testStatSink{}
+	store := NewStore(sink, false)
+
+	scope := store.Scope("runtime")
+	g := NewRuntimeStats(scope)
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		store.AddStatGenerator(g)
+		store.Flush()
+		wg.Done()
+	}()
+
+	go func() {
+		store.AddStatGenerator(g)
+		store.Flush()
+		wg.Done()
+	}()
+
+	wg.Wait()
+}