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
This commit is contained in:
Alexander Weaver 2023-12-01 10:47:19 -06:00 committed by GitHub
parent ea7a179f2a
commit ab0ef5276f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 202 additions and 76 deletions

View File

@ -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
}

View File

@ -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)
}

View File

@ -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
}

View File

@ -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
}

View File

@ -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(&quota.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 := &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

View File

@ -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
}