Alerting: Decouple quota configuration logic from API interfaces and add tests (#78930)

* Separate usage reporter from API

* Extract quota registration

* Decouple from API store interface

* Move to ngalert package and add tests

* linter

pkg/services/ngalert/api/api.go

@@ -148,34 +148,3 @@ func (api *API) RegisterAPIEndpoints(m *metrics.API) {
 		hist:   api.Historian,
 	}), m)
 }
-
-func (api *API) Usage(ctx context.Context, scopeParams *quota.ScopeParameters) (*quota.Map, error) {
-	u := &quota.Map{}
-
-	var orgID int64 = 0
-	if scopeParams != nil {
-		orgID = scopeParams.OrgID
-	}
-
-	if orgUsage, err := api.RuleStore.Count(ctx, orgID); err != nil {
-		return u, err
-	} else {
-		tag, err := quota.NewTag(models.QuotaTargetSrv, models.QuotaTarget, quota.OrgScope)
-		if err != nil {
-			return u, err
-		}
-		u.Set(tag, orgUsage)
-	}
-
-	if globalUsage, err := api.RuleStore.Count(ctx, 0); err != nil {
-		return u, err
-	} else {
-		tag, err := quota.NewTag(models.QuotaTargetSrv, models.QuotaTarget, quota.GlobalScope)
-		if err != nil {
-			return u, err
-		}
-		u.Set(tag, globalUsage)
-	}
-
-	return u, nil
-}

pkg/services/ngalert/api/persist.go

@@ -24,6 +24,4 @@ type RuleStore interface {
 
 	// IncreaseVersionForAllRulesInNamespace Increases version for all rules that have specified namespace. Returns all rules that belong to the namespace
 	IncreaseVersionForAllRulesInNamespace(ctx context.Context, orgID int64, namespaceUID string) ([]ngmodels.AlertRuleKeyWithVersionAndPauseStatus, error)
-
-	Count(ctx context.Context, orgID int64) (int64, error)
 }

pkg/services/ngalert/limits.go

@@ -0,0 +1,88 @@
+package ngalert
+
+import (
+	"context"
+
+	"github.com/grafana/grafana/pkg/services/ngalert/models"
+	"github.com/grafana/grafana/pkg/services/quota"
+	"github.com/grafana/grafana/pkg/setting"
+)
+
+type RuleUsageReader interface {
+	Count(ctx context.Context, orgID int64) (int64, error)
+}
+
+func RegisterQuotas(cfg *setting.Cfg, qs quota.Service, rules RuleUsageReader) error {
+	defaultLimits, err := readQuotaConfig(cfg)
+	if err != nil {
+		return err
+	}
+
+	return qs.RegisterQuotaReporter(&quota.NewUsageReporter{
+		TargetSrv:     models.QuotaTargetSrv,
+		DefaultLimits: defaultLimits,
+		Reporter:      UsageReporter(rules),
+	})
+}
+
+func UsageReporter(rules RuleUsageReader) quota.UsageReporterFunc {
+	return func(ctx context.Context, scopeParams *quota.ScopeParameters) (*quota.Map, error) {
+		u := &quota.Map{}
+
+		var orgID int64 = 0
+		if scopeParams != nil {
+			orgID = scopeParams.OrgID
+		}
+
+		if orgUsage, err := rules.Count(ctx, orgID); err != nil {
+			return u, err
+		} else {
+			tag, err := quota.NewTag(models.QuotaTargetSrv, models.QuotaTarget, quota.OrgScope)
+			if err != nil {
+				return u, err
+			}
+			u.Set(tag, orgUsage)
+		}
+
+		if globalUsage, err := rules.Count(ctx, 0); err != nil {
+			return u, err
+		} else {
+			tag, err := quota.NewTag(models.QuotaTargetSrv, models.QuotaTarget, quota.GlobalScope)
+			if err != nil {
+				return u, err
+			}
+			u.Set(tag, globalUsage)
+		}
+
+		return u, nil
+	}
+}
+
+func readQuotaConfig(cfg *setting.Cfg) (*quota.Map, error) {
+	limits := &quota.Map{}
+
+	if cfg == nil {
+		return limits, nil
+	}
+
+	var alertOrgQuota int64
+	var alertGlobalQuota int64
+
+	if cfg.UnifiedAlerting.IsEnabled() {
+		alertOrgQuota = cfg.Quota.Org.AlertRule
+		alertGlobalQuota = cfg.Quota.Global.AlertRule
+	}
+
+	globalQuotaTag, err := quota.NewTag(models.QuotaTargetSrv, models.QuotaTarget, quota.GlobalScope)
+	if err != nil {
+		return limits, err
+	}
+	orgQuotaTag, err := quota.NewTag(models.QuotaTargetSrv, models.QuotaTarget, quota.OrgScope)
+	if err != nil {
+		return limits, err
+	}
+
+	limits.Set(globalQuotaTag, alertGlobalQuota)
+	limits.Set(orgQuotaTag, alertOrgQuota)
+	return limits, nil
+}

pkg/services/ngalert/limits_test.go

@@ -0,0 +1,113 @@
+package ngalert
+
+import (
+	"context"
+	"testing"
+
+	"github.com/grafana/grafana/pkg/services/ngalert/models"
+	"github.com/grafana/grafana/pkg/services/quota"
+	"github.com/grafana/grafana/pkg/setting"
+	"github.com/stretchr/testify/require"
+)
+
+func TestUsageReporter(t *testing.T) {
+	t.Run("reports org usage", func(t *testing.T) {
+		rules := newFakeUsageReader(map[int64]int64{1: 10, 2: 20})
+		params := quota.ScopeParameters{
+			OrgID: 1,
+		}
+
+		res, err := UsageReporter(rules)(context.Background(), &params)
+
+		require.NoError(t, err)
+		rulesOrg, _ := quota.NewTag(models.QuotaTargetSrv, models.QuotaTarget, quota.OrgScope)
+		val, ok := res.Get(rulesOrg)
+		require.True(t, ok, "reporter did not report on org 1 rules usage")
+		require.Equal(t, int64(10), val)
+	})
+
+	t.Run("reports global usage", func(t *testing.T) {
+		rules := newFakeUsageReader(map[int64]int64{1: 10, 2: 20})
+		params := quota.ScopeParameters{
+			OrgID: 1,
+		}
+
+		res, err := UsageReporter(rules)(context.Background(), &params)
+
+		require.NoError(t, err)
+		rulesGlobal, _ := quota.NewTag(models.QuotaTargetSrv, models.QuotaTarget, quota.GlobalScope)
+		val, ok := res.Get(rulesGlobal)
+		require.True(t, ok, "reporter did not report on global rules usage")
+		require.Equal(t, int64(30), val)
+	})
+
+	t.Run("reports global usage if scope params are nil", func(t *testing.T) {
+		rules := newFakeUsageReader(map[int64]int64{1: 10, 2: 20})
+
+		res, err := UsageReporter(rules)(context.Background(), nil)
+
+		require.NoError(t, err)
+		rulesGlobal, _ := quota.NewTag(models.QuotaTargetSrv, models.QuotaTarget, quota.GlobalScope)
+		val, ok := res.Get(rulesGlobal)
+		require.True(t, ok, "reporter did not report on global rules usage")
+		require.Equal(t, int64(30), val)
+	})
+}
+
+func TestReadQuotaConfig(t *testing.T) {
+	cfg := &setting.Cfg{
+		Quota: setting.QuotaSettings{
+			Org: setting.OrgQuota{
+				AlertRule: 30,
+			},
+			Global: setting.GlobalQuota{
+				AlertRule: 50,
+			},
+		},
+	}
+
+	t.Run("registers per-org rule quota from config", func(t *testing.T) {
+		res, err := readQuotaConfig(cfg)
+
+		require.NoError(t, err)
+		rulesOrg, _ := quota.NewTag(models.QuotaTargetSrv, models.QuotaTarget, quota.OrgScope)
+		val, ok := res.Get(rulesOrg)
+		require.True(t, ok, "did not configure per-org rules quota")
+		require.Equal(t, int64(30), val)
+	})
+
+	t.Run("registers global rule quota from config", func(t *testing.T) {
+		res, err := readQuotaConfig(cfg)
+
+		require.NoError(t, err)
+		rulesGlobal, _ := quota.NewTag(models.QuotaTargetSrv, models.QuotaTarget, quota.GlobalScope)
+		val, ok := res.Get(rulesGlobal)
+		require.True(t, ok, "did not configure global rules quota")
+		require.Equal(t, int64(50), val)
+	})
+}
+
+type fakeUsageReader struct {
+	usage map[int64]int64 // orgID -> count
+}
+
+func newFakeUsageReader(usage map[int64]int64) fakeUsageReader {
+	return fakeUsageReader{
+		usage: usage,
+	}
+}
+
+func (f fakeUsageReader) Count(_ context.Context, orgID int64) (int64, error) {
+	if orgID == 0 {
+		total := int64(0)
+		for _, count := range f.usage {
+			total += count
+		}
+		return total, nil
+	}
+
+	if c, ok := f.usage[orgID]; ok {
+		return c, nil
+	}
+	return 0, nil
+}

pkg/services/ngalert/ngalert.go

@@ -304,16 +304,7 @@ func (ng *AlertNG) init() error {
 	}
 	ng.api.RegisterAPIEndpoints(ng.Metrics.GetAPIMetrics())
 
-	defaultLimits, err := readQuotaConfig(ng.Cfg)
+	if err := RegisterQuotas(ng.Cfg, ng.QuotaService, ng.store); err != nil {
-	if err != nil {
-		return err
-	}
-
-	if err := ng.QuotaService.RegisterQuotaReporter(&quota.NewUsageReporter{
-		TargetSrv:     models.QuotaTargetSrv,
-		DefaultLimits: defaultLimits,
-		Reporter:      ng.api.Usage,
-	}); err != nil {
 		return err
 	}
 
@@ -393,35 +384,6 @@ func (ng *AlertNG) GetHooks() *api.Hooks {
 	return ng.api.Hooks
 }
 
-func readQuotaConfig(cfg *setting.Cfg) (*quota.Map, error) {
-	limits := &quota.Map{}
-
-	if cfg == nil {
-		return limits, nil
-	}
-
-	var alertOrgQuota int64
-	var alertGlobalQuota int64
-
-	if cfg.UnifiedAlerting.IsEnabled() {
-		alertOrgQuota = cfg.Quota.Org.AlertRule
-		alertGlobalQuota = cfg.Quota.Global.AlertRule
-	}
-
-	globalQuotaTag, err := quota.NewTag(models.QuotaTargetSrv, models.QuotaTarget, quota.GlobalScope)
-	if err != nil {
-		return limits, err
-	}
-	orgQuotaTag, err := quota.NewTag(models.QuotaTargetSrv, models.QuotaTarget, quota.OrgScope)
-	if err != nil {
-		return limits, err
-	}
-
-	limits.Set(globalQuotaTag, alertGlobalQuota)
-	limits.Set(orgQuotaTag, alertOrgQuota)
-	return limits, nil
-}
-
 type Historian interface {
 	api.Historian
 	state.Historian

pkg/services/ngalert/tests/fakes/rules.go

@@ -344,10 +344,6 @@ func (f *RuleStore) IncreaseVersionForAllRulesInNamespace(_ context.Context, org
 	return result, nil
 }
 
-func (f *RuleStore) Count(ctx context.Context, orgID int64) (int64, error) {
-	return 0, nil
-}
-
 func (f *RuleStore) CountInFolder(ctx context.Context, orgID int64, folderUID string, u identity.Requester) (int64, error) {
 	return 0, nil
 }