Usage Stats: Decouple from alerting.UsageStatsQuerier (#39515)

* Usage Stats: Decouple from alerting.UsageStatsQuerier

Co-authored-by: Tania B <yalyna.ts@gmail.com>

pkg/services/alerting/engine.go

@@ -9,6 +9,7 @@ import (
 	"github.com/benbjohnson/clock"
 	"github.com/grafana/grafana/pkg/bus"
 	"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/plugins"
 	"github.com/grafana/grafana/pkg/services/rendering"
@@ -29,13 +30,14 @@ type AlertEngine struct {
 	DataService      plugins.DataRequestHandler
 	Cfg              *setting.Cfg
 
-	execQueue     chan *Job
-	ticker        *Ticker
-	scheduler     scheduler
-	evalHandler   evalHandler
-	ruleReader    ruleReader
-	log           log.Logger
-	resultHandler resultHandler
+	execQueue         chan *Job
+	ticker            *Ticker
+	scheduler         scheduler
+	evalHandler       evalHandler
+	ruleReader        ruleReader
+	log               log.Logger
+	resultHandler     resultHandler
+	usageStatsService usagestats.Service
 }
 
 // IsDisabled returns true if the alerting service is disable for this instance.
@@ -45,13 +47,14 @@ func (e *AlertEngine) IsDisabled() bool {
 
 // ProvideAlertEngine returns a new AlertEngine.
 func ProvideAlertEngine(renderer rendering.Service, bus bus.Bus, requestValidator models.PluginRequestValidator,
-	dataService plugins.DataRequestHandler, cfg *setting.Cfg) *AlertEngine {
+	dataService plugins.DataRequestHandler, usageStatsService usagestats.Service, cfg *setting.Cfg) *AlertEngine {
 	e := &AlertEngine{
-		Cfg:              cfg,
-		RenderService:    renderer,
-		Bus:              bus,
-		RequestValidator: requestValidator,
-		DataService:      dataService,
+		Cfg:               cfg,
+		RenderService:     renderer,
+		Bus:               bus,
+		RequestValidator:  requestValidator,
+		DataService:       dataService,
+		usageStatsService: usageStatsService,
 	}
 	e.ticker = NewTicker(time.Now(), time.Second*0, clock.New(), 1)
 	e.execQueue = make(chan *Job, 1000)
@@ -61,6 +64,8 @@ func ProvideAlertEngine(renderer rendering.Service, bus bus.Bus, requestValidato
 	e.log = log.New("alerting.engine")
 	e.resultHandler = newResultHandler(e.RenderService)
 
+	e.registerUsageMetrics()
+
 	return e
 }
 
@@ -237,3 +242,27 @@ func (e *AlertEngine) processJob(attemptID int, attemptChan chan int, cancelChan
 		close(attemptChan)
 	}()
 }
+
+func (e *AlertEngine) registerUsageMetrics() {
+	e.usageStatsService.RegisterMetricsFunc(func() (map[string]interface{}, error) {
+		alertingUsageStats, err := e.QueryUsageStats()
+		if err != nil {
+			return nil, err
+		}
+
+		alertingOtherCount := 0
+		metrics := map[string]interface{}{}
+
+		for dsType, usageCount := range alertingUsageStats.DatasourceUsage {
+			if e.usageStatsService.ShouldBeReported(dsType) {
+				metrics[fmt.Sprintf("stats.alerting.ds.%s.count", dsType)] = usageCount
+			} else {
+				alertingOtherCount += usageCount
+			}
+		}
+
+		metrics["stats.alerting.ds.other.count"] = alertingOtherCount
+
+		return metrics, nil
+	})
+}

pkg/services/alerting/engine_integration_test.go

@@ -18,7 +18,8 @@ import (
 
 func TestEngineTimeouts(t *testing.T) {
 	Convey("Alerting engine timeout tests", t, func() {
-		engine := ProvideAlertEngine(nil, nil, nil, nil, setting.NewCfg())
+		usMock := &usageStatsMock{t: t}
+		engine := ProvideAlertEngine(nil, nil, nil, nil, usMock, setting.NewCfg())
 		setting.AlertingNotificationTimeout = 30 * time.Second
 		setting.AlertingMaxAttempts = 3
 		engine.resultHandler = &FakeResultHandler{}

pkg/services/alerting/engine_test.go

@@ -5,11 +5,15 @@ import (
 	"errors"
 	"math"
 	"testing"
-
 	"time"
 
+	"github.com/grafana/grafana/pkg/bus"
+	"github.com/grafana/grafana/pkg/components/simplejson"
+	"github.com/grafana/grafana/pkg/infra/usagestats"
+	"github.com/grafana/grafana/pkg/models"
 	"github.com/grafana/grafana/pkg/setting"
 	. "github.com/smartystreets/goconvey/convey"
+	"github.com/stretchr/testify/require"
 )
 
 type FakeEvalHandler struct {
@@ -37,15 +41,65 @@ func (handler *FakeResultHandler) handle(evalContext *EvalContext) error {
 	return nil
 }
 
+type usageStatsMock struct {
+	t            *testing.T
+	metricsFuncs []usagestats.MetricsFunc
+}
+
+func (usm *usageStatsMock) RegisterMetricsFunc(fn usagestats.MetricsFunc) {
+	usm.metricsFuncs = append(usm.metricsFuncs, fn)
+}
+
+func (usm *usageStatsMock) GetUsageReport(_ context.Context) (usagestats.Report, error) {
+	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.Report{Metrics: all}, nil
+}
+
+func (usm *usageStatsMock) ShouldBeReported(_ string) bool {
+	return true
+}
+
 func TestEngineProcessJob(t *testing.T) {
 	Convey("Alerting engine job processing", t, func() {
-		engine := ProvideAlertEngine(nil, nil, nil, nil, setting.NewCfg())
+		bus := bus.New()
+		usMock := &usageStatsMock{t: t}
+		engine := ProvideAlertEngine(nil, bus, nil, nil, usMock, setting.NewCfg())
 		setting.AlertingEvaluationTimeout = 30 * time.Second
 		setting.AlertingNotificationTimeout = 30 * time.Second
 		setting.AlertingMaxAttempts = 3
 		engine.resultHandler = &FakeResultHandler{}
 		job := &Job{running: true, Rule: &Rule{}}
 
+		Convey("Should register usage metrics func", func() {
+			bus.AddHandler(func(q *models.GetAllAlertsQuery) error {
+				settings, err := simplejson.NewJson([]byte(`{"conditions": [{"query": { "datasourceId": 1}}]}`))
+				if err != nil {
+					return err
+				}
+				q.Result = []*models.Alert{{Settings: settings}}
+				return nil
+			})
+
+			bus.AddHandler(func(q *models.GetDataSourceQuery) error {
+				q.Result = &models.DataSource{Id: 1, Type: models.DS_PROMETHEUS}
+				return nil
+			})
+
+			report, err := usMock.GetUsageReport(context.Background())
+			So(err, ShouldBeNil)
+
+			So(report.Metrics["stats.alerting.ds.prometheus.count"], ShouldEqual, 1)
+			So(report.Metrics["stats.alerting.ds.other.count"], ShouldEqual, 0)
+		})
+
 		Convey("Should trigger retry if needed", func() {
 			Convey("error + not last attempt -> retry", func() {
 				engine.evalHandler = NewFakeEvalHandler(0)