bug: prevent race condition in Stop method on gauge

Author: mrdg

metriks/gauge.go

@@ -2,7 +2,6 @@ package metriks
 
 import (
 	"context"
-	"sync"
 	"sync/atomic"
 	"time"
 
@@ -24,7 +23,7 @@ type PersistentGauge struct {
 	ticker *time.Ticker
 	cancel context.CancelFunc
 	dur    time.Duration
-	wg     sync.WaitGroup
+	donec  chan struct{}
 }
 
 // Set will replace the value with a new one, it returns the old value
@@ -48,11 +47,10 @@ func (g *PersistentGauge) report(v int32) {
 }
 
 func (g *PersistentGauge) start(ctx context.Context) {
-	g.wg.Add(1)
 	for {
 		select {
 		case <-ctx.Done():
-			g.wg.Done()
+			close(g.donec)
 			return
 		case <-g.ticker.C:
 			g.report(atomic.LoadInt32(&g.value))
@@ -64,7 +62,7 @@ func (g *PersistentGauge) start(ctx context.Context) {
 // to the metrics collector
 func (g *PersistentGauge) Stop() {
 	g.cancel()
-	g.wg.Wait()
+	<-g.donec
 }
 
 // NewPersistentGauge will create and start a PersistentGauge that reports the current value every 10s
@@ -81,6 +79,7 @@ func NewPersistentGaugeWithDuration(name string, dur time.Duration, tags ...metr
 		ticker: time.NewTicker(dur),
 		cancel: cancel,
 		dur:    dur,
+		donec:  make(chan struct{}),
 	}
 	go g.start(ctx)
 	return &g

metriks/gauge_test.go

@@ -37,6 +37,13 @@ func TestPersistentGauge(t *testing.T) {
 	}
 }
 
+func TestPersistentGaugeRace(t *testing.T) {
+	for n := 0; n < 10; n++ {
+		g := NewPersistentGaugeWithDuration("some_gauge", time.Second, L("a", "b"))
+		g.Stop()
+	}
+}
+
 func TestScheduledGauge(t *testing.T) {
 	var callCount int32
 	cb := func() int32 {