Usage stats: Fix concurrency issue (#84993)

* Usage stats: Fix concurrency issue

pkg/infra/usagestats/service/usage_stats.go

@@ -29,16 +29,16 @@ var usageStatsURL = "https://stats.grafana.org/grafana-usage-report"
 
 func (uss *UsageStats) GetUsageReport(ctx context.Context) (usagestats.Report, error) {
 	version := strings.ReplaceAll(uss.Cfg.BuildVersion, ".", "_")
-	metrics := map[string]any{}
+	metrics := sync.Map{}
 	start := time.Now()
 
 	edition := "oss"
 	if uss.Cfg.IsEnterprise {
 		edition = "enterprise"
 	}
+
 	report := usagestats.Report{
 		Version:      version,
-		Metrics:      metrics,
 		Os:           runtime.GOOS,
 		Arch:         runtime.GOARCH,
 		Edition:      edition,
@@ -46,20 +46,28 @@ func (uss *UsageStats) GetUsageReport(ctx context.Context) (usagestats.Report, e
 		UsageStatsId: uss.GetUsageStatsId(ctx),
 	}
 
-	uss.gatherMetrics(ctx, metrics)
+	uss.gatherMetrics(ctx, &metrics)
 
 	// must run after registration of external metrics
-	if v, ok := metrics["stats.valid_license.count"]; ok {
+	if v, ok := metrics.Load("stats.valid_license.count"); ok {
 		report.HasValidLicense = v == 1
 	} else {
-		metrics["stats.valid_license.count"] = 0
+		metrics.Store("stats.valid_license.count", 0)
 	}
 
-	uss.log.FromContext(ctx).Debug("Collected usage stats", "metricCount", len(metrics), "version", report.Version, "os", report.Os, "arch", report.Arch, "edition", report.Edition, "duration", time.Since(start))
+	report.Metrics = make(map[string]any)
+	metricCount := 0
+	metrics.Range(func(key, value any) bool {
+		report.Metrics[key.(string)] = value
+		metricCount++
+		return true
+	})
+
+	uss.log.FromContext(ctx).Debug("Collected usage stats", "metricCount", metricCount, "version", report.Version, "os", report.Os, "arch", report.Arch, "edition", report.Edition, "duration", time.Since(start))
 	return report, nil
 }
 
-func (uss *UsageStats) gatherMetrics(ctx context.Context, metrics map[string]any) {
+func (uss *UsageStats) gatherMetrics(ctx context.Context, metrics *sync.Map) {
 	ctxTracer, span := uss.tracer.Start(ctx, "UsageStats.GatherLoop")
 	defer span.End()
 	totC, errC := 0, 0
@@ -86,14 +94,14 @@ func (uss *UsageStats) gatherMetrics(ctx context.Context, metrics map[string]any
 			}
 
 			for name, value := range fnMetrics {
-				metrics[name] = value
+				metrics.Store(name, value)
 			}
 		}(fn)
 	}
 
 	wg.Wait()
-	metrics["stats.usagestats.debug.collect.total.count"] = totC
-	metrics["stats.usagestats.debug.collect.error.count"] = errC
+	metrics.Store("stats.usagestats.debug.collect.total.count", totC)
+	metrics.Store("stats.usagestats.debug.collect.error.count", errC)
 }
 
 func (uss *UsageStats) runMetricsFunc(ctx context.Context, fn usagestats.MetricsFunc) (map[string]any, error) {

pkg/infra/usagestats/service/usage_stats_test.go

@@ -9,6 +9,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"runtime"
+	"sync"
 	"testing"
 	"time"
 
@@ -175,7 +176,9 @@ func TestRegisterMetrics(t *testing.T) {
 
 	sqlStore := dbtest.NewFakeDB()
 	uss := createService(t, sqlStore, false)
-	metrics := map[string]any{"stats.test_metric.count": 1, "stats.test_metric_second.count": 2}
+	metrics := sync.Map{}
+	metrics.Store("stats.test_metric.count", 1)
+	metrics.Store("stats.test_metric_second.count", 2)
 
 	uss.RegisterMetricsFunc(func(context.Context) (map[string]any, error) {
 		return map[string]any{goodMetricName: 1}, nil
@@ -187,9 +190,15 @@ func TestRegisterMetrics(t *testing.T) {
 		assert.Equal(t, map[string]any{goodMetricName: 1}, extMetrics)
 	}
 
-	uss.gatherMetrics(context.Background(), metrics)
-	assert.Equal(t, 1, metrics[goodMetricName])
-	metricsCount := len(metrics)
+	uss.gatherMetrics(context.Background(), &metrics)
+	v, ok := metrics.Load(goodMetricName)
+	assert.True(t, ok)
+	assert.Equal(t, 1, v)
+	metricsCountBefore := 0
+	metrics.Range(func(_, _ any) bool {
+		metricsCountBefore++
+		return true
+	})
 
 	t.Run("do not add metrics that return an error when fetched", func(t *testing.T) {
 		const badMetricName = "stats.test_external_metric_error.count"
@@ -197,15 +206,26 @@ func TestRegisterMetrics(t *testing.T) {
 		uss.RegisterMetricsFunc(func(context.Context) (map[string]any, error) {
 			return map[string]any{badMetricName: 1}, errors.New("some error")
 		})
-		uss.gatherMetrics(context.Background(), metrics)
+		uss.gatherMetrics(context.Background(), &metrics)
 
-		extErrorMetric := metrics[badMetricName]
-		extMetric := metrics[goodMetricName]
+		extErrorMetric, ok := metrics.Load(badMetricName)
+		assert.False(t, ok)
+		extMetric, ok := metrics.Load(goodMetricName)
+		assert.True(t, ok)
 
 		require.Nil(t, extErrorMetric, "Invalid metric should not be added")
 		assert.Equal(t, 1, extMetric)
-		assert.Len(t, metrics, metricsCount, "Expected same number of metrics before and after collecting bad metric")
-		assert.EqualValues(t, 1, metrics["stats.usagestats.debug.collect.error.count"])
+
+		metricsCountAfter := 0
+		metrics.Range(func(_, _ any) bool {
+			metricsCountAfter++
+			return true
+		})
+
+		assert.Equal(t, metricsCountAfter, metricsCountBefore, "Expected same number of metrics before and after collecting bad metric")
+		errCount, ok := metrics.Load("stats.usagestats.debug.collect.error.count")
+		assert.True(t, ok)
+		assert.EqualValues(t, 1, errCount)
 	})
 }