fix: Fix data race issue on metrics Counters and Histograms (#139)

openfga · Oct 22, 2024 · d70965e · d70965e
1 parent 9c2614b
commit d70965e
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 4 deletions.
diff --git a/telemetry/metrics.go b/telemetry/metrics.go
@@ -3,16 +3,19 @@ package telemetry
 import (
 	"context"
 	"net/http"
+	"sync"
 	"time"
 
 	"go.opentelemetry.io/otel/metric"
 )
 
 type Metrics struct {
-	Meter         metric.Meter
-	Counters      map[string]metric.Int64Counter
-	Histograms    map[string]metric.Float64Histogram
-	Configuration *MetricsConfiguration
+	Meter          metric.Meter
+	countersLock   sync.Mutex
+	Counters       map[string]metric.Int64Counter
+	histogramsLock sync.Mutex
+	Histograms     map[string]metric.Float64Histogram
+	Configuration  *MetricsConfiguration
 }
 
 type MetricsInterface interface {
@@ -25,6 +28,9 @@ type MetricsInterface interface {
 }
 
 func (m *Metrics) GetCounter(name string, description string) (metric.Int64Counter, error) {
+	m.countersLock.Lock()
+	defer m.countersLock.Unlock()
+
 	if counter, exists := m.Counters[name]; exists {
 		return counter, nil
 	}
@@ -34,6 +40,9 @@ func (m *Metrics) GetCounter(name string, description string) (metric.Int64Count
 }
 
 func (m *Metrics) GetHistogram(name string, description string, unit string) (metric.Float64Histogram, error) {
+	m.histogramsLock.Lock()
+	defer m.histogramsLock.Unlock()
+
 	if histogram, exists := m.Histograms[name]; exists {
 		return histogram, nil
 	}

diff --git a/telemetry/metrics_test.go b/telemetry/metrics_test.go
@@ -2,6 +2,7 @@ package telemetry
 
 import (
 	"context"
+	"sync"
 	"testing"
 
 	"go.opentelemetry.io/otel/metric"
@@ -81,6 +82,41 @@ func TestGetCounter(t *testing.T) {
 	}
 }
 
+// Run this test with the "-race" flag.
+//
+//	go test -race -run ^TestGetCounterRace github.com/openfga/go-sdk/telemetry
+func TestGetCounterRace(t *testing.T) {
+
+	mockMeter := &MockMeter{
+		counters:   make(map[string]metric.Int64Counter),
+		histograms: make(map[string]metric.Float64Histogram),
+	}
+	metrics := &Metrics{
+		Meter:    mockMeter,
+		Counters: make(map[string]metric.Int64Counter),
+	}
+
+	t.Parallel()
+
+	var wg sync.WaitGroup
+
+	for i := 0; i < 10; i++ {
+		wg.Add(1)
+
+		go func() {
+			defer wg.Done()
+			_, err := metrics.GetCounter("test_counter", "A test counter")
+			if err != nil {
+				t.Errorf("Expected no error, but got %v", err)
+				return
+			}
+			Get(TelemetryFactoryParameters{})
+		}()
+	}
+
+	wg.Wait()
+}
+
 func TestGetHistogram(t *testing.T) {
 	mockMeter := &MockMeter{
 		counters:   make(map[string]metric.Int64Counter),
@@ -110,6 +146,37 @@ func TestGetHistogram(t *testing.T) {
 	}
 }
 
+// go test -race -run ^TestGetHistogramRace github.com/openfga/go-sdk/telemetry
+func TestGetHistogramRace(t *testing.T) {
+	mockMeter := &MockMeter{
+		counters:   make(map[string]metric.Int64Counter),
+		histograms: make(map[string]metric.Float64Histogram),
+	}
+	metrics := &Metrics{
+		Meter:      mockMeter,
+		Histograms: make(map[string]metric.Float64Histogram),
+	}
+
+	t.Parallel()
+
+	var wg sync.WaitGroup
+
+	for i := 0; i < 10; i++ {
+		wg.Add(1)
+
+		go func() {
+			defer wg.Done()
+			_, err := metrics.GetHistogram("test_histogram", "A test histogram", "ms")
+			if err != nil {
+				t.Errorf("Expected no error, but got %v", err)
+				return
+			}
+		}()
+	}
+
+	wg.Wait()
+}
+
 func TestCredentialsRequest(t *testing.T) {
 	mockMeter := &MockMeter{
 		counters:   make(map[string]metric.Int64Counter),