monitoring: fix race condition on New* variable functions (#339)

Author: AndersonQ

monitoring/metrics.go

@@ -43,6 +43,13 @@ type Int struct{ i atomic.Int64 }
 // variable will be available via expvars package as well, but can not be removed
 // anymore.
 func NewInt(r *Registry, name string, opts ...Option) *Int {
+	rr := r
+	if rr == nil {
+		rr = Default
+	}
+	rr.txMu.Lock()
+	defer rr.txMu.Unlock()
+
 	existingVar, r := setupMetric(r, name, opts)
 	if existingVar != nil {
 		cast, ok := existingVar.(*Int)
@@ -77,6 +84,13 @@ type Uint struct{ u atomic.Uint64 }
 // variable will be available via expvars package as well, but can not be removed
 // anymore.
 func NewUint(r *Registry, name string, opts ...Option) *Uint {
+	rr := r
+	if rr == nil {
+		rr = Default
+	}
+	rr.txMu.Lock()
+	defer rr.txMu.Unlock()
+
 	existingVar, r := setupMetric(r, name, opts)
 	if existingVar != nil {
 		cast, ok := existingVar.(*Uint)
@@ -114,6 +128,13 @@ type Float struct{ f atomic.Uint64 }
 // variable will be available via expvars package as well, but can not be removed
 // anymore.
 func NewFloat(r *Registry, name string, opts ...Option) *Float {
+	rr := r
+	if rr == nil {
+		rr = Default
+	}
+	rr.txMu.Lock()
+	defer rr.txMu.Unlock()
+
 	existingVar, r := setupMetric(r, name, opts)
 	if existingVar != nil {
 		cast, ok := existingVar.(*Float)
@@ -155,6 +176,13 @@ type Bool struct{ f atomic.Bool }
 // variable will be available via expvars package as well, but can not be removed
 // anymore.
 func NewBool(r *Registry, name string, opts ...Option) *Bool {
+	rr := r
+	if rr == nil {
+		rr = Default
+	}
+	rr.txMu.Lock()
+	defer rr.txMu.Unlock()
+
 	existingVar, r := setupMetric(r, name, opts)
 	if existingVar != nil {
 		cast, ok := existingVar.(*Bool)
@@ -188,6 +216,13 @@ type String struct {
 // variable will be available via expvars package as well, but can not be removed
 // anymore.
 func NewString(r *Registry, name string, opts ...Option) *String {
+	rr := r
+	if rr == nil {
+		rr = Default
+	}
+	rr.txMu.Lock()
+	defer rr.txMu.Unlock()
+
 	existingVar, r := setupMetric(r, name, opts)
 	if existingVar != nil {
 		cast, ok := existingVar.(*String)
@@ -239,6 +274,13 @@ type Func struct {
 }
 
 func NewFunc(r *Registry, name string, f func(Mode, Visitor), opts ...Option) *Func {
+	rr := r
+	if rr == nil {
+		rr = Default
+	}
+	rr.txMu.Lock()
+	defer rr.txMu.Unlock()
+
 	existingVar, r := setupMetric(r, name, opts)
 	if existingVar != nil {
 		cast, ok := existingVar.(*Func)
@@ -282,6 +324,13 @@ type Timestamp struct {
 
 // NewTimestamp creates and registers a new timestamp variable.
 func NewTimestamp(r *Registry, name string, opts ...Option) *Timestamp {
+	rr := r
+	if rr == nil {
+		rr = Default
+	}
+	rr.txMu.Lock()
+	defer rr.txMu.Unlock()
+
 	existingVar, r := setupMetric(r, name, opts)
 	if existingVar != nil {
 		cast, ok := existingVar.(*Timestamp)

monitoring/metrics_test.go

@@ -18,24 +18,150 @@
 package monitoring
 
 import (
+	"sync"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
 func TestSafeVars(t *testing.T) {
-	uintValName := "testUint"
-	testReg := Default.NewRegistry("safe_registry")
-	testUint := NewUint(testReg, uintValName)
-	testUint.Set(5)
-	// Add the first time
-	require.NotNil(t, testUint)
+	t.Run("no concurrency", func(t *testing.T) {
+		uintValName := "testUint"
+		testReg := NewRegistry().NewRegistry("safe_registry")
+		testUint := NewUint(testReg, uintValName)
+		testUint.Set(5)
+		// Add the first time
+		require.NotNil(t, testUint)
+
+		// Add the metric a second time
+		testSecondUint := NewUint(testReg, uintValName)
+		require.NotNil(t, testSecondUint)
+		// make sure we fetch the same unit
+		require.Equal(t, uint64(5), testSecondUint.Get())
+	})
+
+	t.Run("with concurrency", func(t *testing.T) {
+		t.Run("NewInt", func(t *testing.T) {
+			reg := NewRegistry()
+			name := "foo"
+
+			wg := sync.WaitGroup{}
+			assert.NotPanics(t, func() {
+				for i := 0; i < 1000; i++ {
+					wg.Add(1)
+					go func() {
+						defer wg.Done()
+						NewInt(reg, name)
+					}()
+				}
+			})
+			wg.Wait()
+		})
+
+		t.Run("NewUint", func(t *testing.T) {
+			reg := NewRegistry()
+			name := "foo"
+
+			wg := sync.WaitGroup{}
+			assert.NotPanics(t, func() {
+				for i := 0; i < 1000; i++ {
+					wg.Add(1)
+					go func() {
+						defer wg.Done()
+						NewUint(reg, name)
+					}()
+				}
+			})
+			wg.Wait()
+		})
+
+		t.Run("NewFloat", func(t *testing.T) {
+			reg := NewRegistry()
+			name := "foo"
 
-	// Add the metric a second time
-	testSecondUint := NewUint(testReg, uintValName)
-	require.NotNil(t, testSecondUint)
-	// make sure we fetch the same unit
-	require.Equal(t, uint64(5), testSecondUint.Get())
+			wg := sync.WaitGroup{}
+			assert.NotPanics(t, func() {
+				for i := 0; i < 1000; i++ {
+					wg.Add(1)
+					go func() {
+						defer wg.Done()
+						NewFloat(reg, name)
+					}()
+				}
+			})
+			wg.Wait()
+		})
+
+		t.Run("NewBool", func(t *testing.T) {
+			reg := NewRegistry()
+			name := "foo"
+
+			wg := sync.WaitGroup{}
+			assert.NotPanics(t, func() {
+				for i := 0; i < 1000; i++ {
+					wg.Add(1)
+					go func() {
+						defer wg.Done()
+						NewBool(reg, name)
+					}()
+				}
+			})
+			wg.Wait()
+		})
+
+		t.Run("NewString", func(t *testing.T) {
+			reg := NewRegistry()
+			name := "foo"
+
+			wg := sync.WaitGroup{}
+			assert.NotPanics(t, func() {
+				for i := 0; i < 1000; i++ {
+					wg.Add(1)
+					go func() {
+						defer wg.Done()
+						NewString(reg, name)
+					}()
+				}
+			})
+			wg.Wait()
+		})
+
+		t.Run("NewFunc", func(t *testing.T) {
+			reg := NewRegistry()
+			name := "foo"
+			dummyFunc := func(m Mode, v Visitor) {}
+
+			wg := sync.WaitGroup{}
+			assert.NotPanics(t, func() {
+				for i := 0; i < 1000; i++ {
+					wg.Add(1)
+					go func() {
+						defer wg.Done()
+						NewFunc(reg, name, dummyFunc)
+					}()
+				}
+			})
+			wg.Wait()
+		})
+		
+		t.Run("NewTimestamp", func(t *testing.T) {
+			reg := NewRegistry()
+			name := "foo"
+
+			wg := sync.WaitGroup{}
+			assert.NotPanics(t, func() {
+				for i := 0; i < 1000; i++ {
+					wg.Add(1)
+					go func() {
+						defer wg.Done()
+						NewTimestamp(reg, name)
+					}()
+				}
+			})
+			wg.Wait()
+		})
+	})
 }
 
 func TestVarsTypes(t *testing.T) {

monitoring/registry.go

@@ -28,7 +28,10 @@ import (
 // When adding or retrieving variables, all names are split on the `.`-symbol and
 // intermediate registries will be generated.
 type Registry struct {
-	mu sync.RWMutex
+	// txMu is a transaction mutex for the New* functions which create new
+	// variables on a registry as they are not goroutine safe.
+	txMu sync.Mutex
+	mu   sync.RWMutex
 
 	name    string
 	entries map[string]entry