diff --git a/pkg/services/ngalert/api/api.go b/pkg/services/ngalert/api/api.go index 400d667674c..7b755d9a6dc 100644 --- a/pkg/services/ngalert/api/api.go +++ b/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 := "a.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 -} diff --git a/pkg/services/ngalert/api/persist.go b/pkg/services/ngalert/api/persist.go index 264d9faff54..82217a5f8cf 100644 --- a/pkg/services/ngalert/api/persist.go +++ b/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) } diff --git a/pkg/services/ngalert/limits.go b/pkg/services/ngalert/limits.go new file mode 100644 index 00000000000..ec459d08673 --- /dev/null +++ b/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("a.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 := "a.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 := "a.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 +} diff --git a/pkg/services/ngalert/limits_test.go b/pkg/services/ngalert/limits_test.go new file mode 100644 index 00000000000..28895574648 --- /dev/null +++ b/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(), ¶ms) + + 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(), ¶ms) + + 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 +} diff --git a/pkg/services/ngalert/ngalert.go b/pkg/services/ngalert/ngalert.go index d0eb359484f..48d6e5d4249 100644 --- a/pkg/services/ngalert/ngalert.go +++ b/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 != nil { - return err - } - - if err := ng.QuotaService.RegisterQuotaReporter("a.NewUsageReporter{ - TargetSrv: models.QuotaTargetSrv, - DefaultLimits: defaultLimits, - Reporter: ng.api.Usage, - }); err != nil { + if err := RegisterQuotas(ng.Cfg, ng.QuotaService, ng.store); 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 := "a.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 diff --git a/pkg/services/ngalert/tests/fakes/rules.go b/pkg/services/ngalert/tests/fakes/rules.go index 0f8f063227b..05654557d92 100644 --- a/pkg/services/ngalert/tests/fakes/rules.go +++ b/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 }