Usage Stats: Make the UsageStatsService extension point more flexible (#34778) (#34895)

* Usage Stats: Rename service to use a more idiomatic name

* Usage Stats: Update MetricsFunc definition and implementations

* Revert "Usage Stats: Rename service to use a more idiomatic name"

This reverts commit 910ecce3e8.

* Usage Stats: Update MetricsFunc definition and implementations

(cherry picked from commit f601921670)

pkg/infra/usagestats/service.go

@@ -6,32 +6,31 @@ import (
 	"time"
 
 	"github.com/grafana/grafana/pkg/bus"
+	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/login/social"
 	"github.com/grafana/grafana/pkg/models"
 	"github.com/grafana/grafana/pkg/plugins"
+	"github.com/grafana/grafana/pkg/registry"
 	"github.com/grafana/grafana/pkg/services/alerting"
 	"github.com/grafana/grafana/pkg/services/sqlstore"
-
-	"github.com/grafana/grafana/pkg/infra/log"
-	"github.com/grafana/grafana/pkg/registry"
 	"github.com/grafana/grafana/pkg/setting"
 )
 
-var metricsLogger log.Logger = log.New("metrics")
+var metricsLogger = log.New("metrics")
 
 func init() {
 	registry.RegisterService(&UsageStatsService{
 		log:             log.New("infra.usagestats"),
-		externalMetrics: make(map[string]MetricFunc),
+		externalMetrics: make([]MetricsFunc, 0),
 	})
 }
 
 type UsageStats interface {
-	GetUsageReport(ctx context.Context) (UsageReport, error)
-	RegisterMetric(name string, fn MetricFunc)
+	GetUsageReport(context.Context) (UsageReport, error)
+	RegisterMetricsFunc(MetricsFunc)
 }
 
-type MetricFunc func() (interface{}, error)
+type MetricsFunc func() (map[string]interface{}, error)
 
 type UsageStatsService struct {
 	Cfg                *setting.Cfg               `inject:""`
@@ -44,7 +43,7 @@ type UsageStatsService struct {
 	log log.Logger
 
 	oauthProviders           map[string]bool
-	externalMetrics          map[string]MetricFunc
+	externalMetrics          []MetricsFunc
 	concurrentUserStatsCache memoConcurrentUserStats
 }
 

pkg/infra/usagestats/usage_stats.go

@@ -240,18 +240,21 @@ func (uss *UsageStatsService) GetUsageReport(ctx context.Context) (UsageReport,
 }
 
 func (uss *UsageStatsService) registerExternalMetrics(metrics map[string]interface{}) {
-	for name, fn := range uss.externalMetrics {
-		result, err := fn()
+	for _, fn := range uss.externalMetrics {
+		fnMetrics, err := fn()
 		if err != nil {
-			metricsLogger.Error("Failed to fetch external metric", "name", name, "error", err)
+			metricsLogger.Error("Failed to fetch external metrics", "error", err)
 			continue
 		}
-		metrics[name] = result
+
+		for name, value := range fnMetrics {
+			metrics[name] = value
+		}
 	}
 }
 
-func (uss *UsageStatsService) RegisterMetric(name string, fn MetricFunc) {
-	uss.externalMetrics[name] = fn
+func (uss *UsageStatsService) RegisterMetricsFunc(fn MetricsFunc) {
+	uss.externalMetrics = append(uss.externalMetrics, fn)
 }
 
 func (uss *UsageStatsService) sendUsageStats(ctx context.Context) error {

pkg/infra/usagestats/usage_stats_test.go

@@ -5,25 +5,23 @@ import (
 	"context"
 	"errors"
 	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
 	"runtime"
 	"testing"
 	"time"
 
+	"github.com/grafana/grafana/pkg/bus"
+	"github.com/grafana/grafana/pkg/components/simplejson"
+	"github.com/grafana/grafana/pkg/models"
 	"github.com/grafana/grafana/pkg/plugins"
 	"github.com/grafana/grafana/pkg/plugins/manager"
 	"github.com/grafana/grafana/pkg/services/alerting"
 	"github.com/grafana/grafana/pkg/services/licensing"
-	"github.com/stretchr/testify/require"
-
-	"net/http"
-	"net/http/httptest"
-
-	"github.com/grafana/grafana/pkg/bus"
-	"github.com/grafana/grafana/pkg/components/simplejson"
-	"github.com/grafana/grafana/pkg/models"
 	"github.com/grafana/grafana/pkg/services/sqlstore"
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 // This is to ensure that the interface contract is held by the implementation
@@ -415,30 +413,13 @@ func TestMetrics(t *testing.T) {
 		metricName := "stats.test_metric.count"
 
 		t.Run("Adds a new metric to the external metrics", func(t *testing.T) {
-			uss.RegisterMetric(metricName, func() (interface{}, error) {
-				return 1, nil
+			uss.RegisterMetricsFunc(func() (map[string]interface{}, error) {
+				return map[string]interface{}{metricName: 1}, nil
 			})
 
-			metric, err := uss.externalMetrics[metricName]()
+			metrics, err := uss.externalMetrics[0]()
 			require.NoError(t, err)
-			assert.Equal(t, 1, metric)
-		})
-
-		t.Run("When metric already exists, the metric should be overridden", func(t *testing.T) {
-			uss.RegisterMetric(metricName, func() (interface{}, error) {
-				return 1, nil
-			})
-
-			metric, err := uss.externalMetrics[metricName]()
-			require.NoError(t, err)
-			assert.Equal(t, 1, metric)
-
-			uss.RegisterMetric(metricName, func() (interface{}, error) {
-				return 2, nil
-			})
-			newMetric, err := uss.externalMetrics[metricName]()
-			require.NoError(t, err)
-			assert.Equal(t, 2, newMetric)
+			assert.Equal(t, map[string]interface{}{metricName: 1}, metrics)
 		})
 	})
 
@@ -486,8 +467,8 @@ func TestMetrics(t *testing.T) {
 		})
 
 		t.Run("Should include external metrics", func(t *testing.T) {
-			uss.RegisterMetric(metricName, func() (interface{}, error) {
-				return 1, nil
+			uss.RegisterMetricsFunc(func() (map[string]interface{}, error) {
+				return map[string]interface{}{metricName: 1}, nil
 			})
 
 			report, err := uss.GetUsageReport(context.Background())
@@ -503,8 +484,8 @@ func TestMetrics(t *testing.T) {
 		metrics := map[string]interface{}{"stats.test_metric.count": 1, "stats.test_metric_second.count": 2}
 		extMetricName := "stats.test_external_metric.count"
 
-		uss.RegisterMetric(extMetricName, func() (interface{}, error) {
-			return 1, nil
+		uss.RegisterMetricsFunc(func() (map[string]interface{}, error) {
+			return map[string]interface{}{extMetricName: 1}, nil
 		})
 
 		uss.registerExternalMetrics(metrics)
@@ -512,14 +493,14 @@ func TestMetrics(t *testing.T) {
 		assert.Equal(t, 1, metrics[extMetricName])
 
 		t.Run("When loading a metric results to an error", func(t *testing.T) {
-			uss.RegisterMetric(extMetricName, func() (interface{}, error) {
-				return 1, nil
+			uss.RegisterMetricsFunc(func() (map[string]interface{}, error) {
+				return map[string]interface{}{extMetricName: 1}, nil
 			})
 			extErrorMetricName := "stats.test_external_metric_error.count"
 
 			t.Run("Should not add it to metrics", func(t *testing.T) {
-				uss.RegisterMetric(extErrorMetricName, func() (interface{}, error) {
-					return 1, errors.New("some error")
+				uss.RegisterMetricsFunc(func() (map[string]interface{}, error) {
+					return map[string]interface{}{extErrorMetricName: 1}, errors.New("some error")
 				})
 
 				uss.registerExternalMetrics(metrics)
@@ -619,7 +600,7 @@ func createService(t *testing.T, cfg setting.Cfg) *UsageStatsService {
 		SQLStore:           sqlstore.InitTestDB(t),
 		License:            &licensing.OSSLicensingService{},
 		AlertingUsageStats: &alertingUsageMock{},
-		externalMetrics:    make(map[string]MetricFunc),
+		externalMetrics:    make([]MetricsFunc, 0),
 		PluginManager:      &fakePluginManager{},
 	}
 }

pkg/services/accesscontrol/ossaccesscontrol/ossaccesscontrol.go

@@ -39,14 +39,19 @@ func (ac *OSSAccessControlService) IsDisabled() bool {
 }
 
 func (ac *OSSAccessControlService) registerUsageMetrics() {
-	ac.UsageStats.RegisterMetric("stats.oss.accesscontrol.enabled.count", ac.getUsageMetrics)
+	ac.UsageStats.RegisterMetricsFunc(func() (map[string]interface{}, error) {
+		return map[string]interface{}{
+			"stats.oss.accesscontrol.enabled.count": ac.getUsageMetrics(),
+		}, nil
+	})
 }
 
-func (ac *OSSAccessControlService) getUsageMetrics() (interface{}, error) {
+func (ac *OSSAccessControlService) getUsageMetrics() interface{} {
 	if ac.IsDisabled() {
-		return 0, nil
+		return 0
 	}
-	return 1, nil
+
+	return 1
 }
 
 // Evaluate evaluates access to the given resource

pkg/services/accesscontrol/ossaccesscontrol/ossaccesscontrol_test.go

@@ -4,15 +4,14 @@ import (
 	"context"
 	"testing"
 
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/infra/usagestats"
 	"github.com/grafana/grafana/pkg/models"
 	"github.com/grafana/grafana/pkg/registry"
 	"github.com/grafana/grafana/pkg/services/accesscontrol"
 	"github.com/grafana/grafana/pkg/setting"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func setupTestEnv(t testing.TB) *OSSAccessControlService {
@@ -23,7 +22,7 @@ func setupTestEnv(t testing.TB) *OSSAccessControlService {
 
 	ac := OSSAccessControlService{
 		Cfg:        cfg,
-		UsageStats: &usageStatsMock{metricFuncs: make(map[string]usagestats.MetricFunc)},
+		UsageStats: &usageStatsMock{metricsFuncs: make([]usagestats.MetricsFunc, 0)},
 		Log:        log.New("accesscontrol-test"),
 	}
 
@@ -33,22 +32,25 @@ func setupTestEnv(t testing.TB) *OSSAccessControlService {
 }
 
 type usageStatsMock struct {
-	t           *testing.T
-	metricFuncs map[string]usagestats.MetricFunc
+	t            *testing.T
+	metricsFuncs []usagestats.MetricsFunc
 }
 
-func (usm *usageStatsMock) RegisterMetric(name string, fn usagestats.MetricFunc) {
-	usm.metricFuncs[name] = fn
+func (usm *usageStatsMock) RegisterMetricsFunc(fn usagestats.MetricsFunc) {
+	usm.metricsFuncs = append(usm.metricsFuncs, fn)
 }
 
 func (usm *usageStatsMock) GetUsageReport(_ context.Context) (usagestats.UsageReport, error) {
-	metrics := make(map[string]interface{})
-	for name, fn := range usm.metricFuncs {
-		v, err := fn()
-		metrics[name] = v
+	all := make(map[string]interface{})
+	for _, fn := range usm.metricsFuncs {
+		fnMetrics, err := fn()
 		require.NoError(usm.t, err)
+
+		for name, value := range fnMetrics {
+			all[name] = value
+		}
 	}
-	return usagestats.UsageReport{Metrics: metrics}, nil
+	return usagestats.UsageReport{Metrics: all}, nil
 }
 
 type evaluatingPermissionsTestCase struct {
@@ -146,7 +148,7 @@ func TestUsageMetrics(t *testing.T) {
 
 			s := &OSSAccessControlService{
 				Cfg:        cfg,
-				UsageStats: &usageStatsMock{t: t, metricFuncs: make(map[string]usagestats.MetricFunc)},
+				UsageStats: &usageStatsMock{t: t, metricsFuncs: make([]usagestats.MetricsFunc, 0)},
 				Log:        log.New("accesscontrol-test"),
 			}