Chore: Refactor quota service (#58643)

Chore: Refactor quota service (#57586)

* Chore: refactore quota service

* Apply suggestions from code review

pkg/services/auth/auth_token.go

@@ -12,6 +12,7 @@ import (
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/infra/serverlock"
 	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/services/quota"
 	"github.com/grafana/grafana/pkg/services/user"
 	"github.com/grafana/grafana/pkg/setting"
 	"github.com/grafana/grafana/pkg/util"
@@ -41,19 +42,38 @@ type UserAuthTokenService struct {
 	log               log.Logger
 }
 
+type ActiveTokenService interface {
+	ActiveTokenCount(ctx context.Context, _ *quota.ScopeParameters) (*quota.Map, error)
+}
+
 type ActiveAuthTokenService struct {
 	cfg      *setting.Cfg
 	sqlStore db.DB
 }
 
-func ProvideActiveAuthTokenService(cfg *setting.Cfg, sqlStore db.DB) *ActiveAuthTokenService {
-	return &ActiveAuthTokenService{
+func ProvideActiveAuthTokenService(cfg *setting.Cfg, sqlStore db.DB, quotaService quota.Service) (*ActiveAuthTokenService, error) {
+	s := &ActiveAuthTokenService{
 		cfg:      cfg,
 		sqlStore: sqlStore,
 	}
+
+	defaultLimits, err := readQuotaConfig(cfg)
+	if err != nil {
+		return s, err
+	}
+
+	if err := quotaService.RegisterQuotaReporter(&quota.NewUsageReporter{
+		TargetSrv:     QuotaTargetSrv,
+		DefaultLimits: defaultLimits,
+		Reporter:      s.ActiveTokenCount,
+	}); err != nil {
+		return s, err
+	}
+
+	return s, nil
 }
 
-func (a *ActiveAuthTokenService) ActiveTokenCount(ctx context.Context) (int64, error) {
+func (a *ActiveAuthTokenService) ActiveTokenCount(ctx context.Context, _ *quota.ScopeParameters) (*quota.Map, error) {
 	var count int64
 	var err error
 	err = a.sqlStore.WithDbSession(ctx, func(dbSession *db.Session) error {
@@ -66,7 +86,14 @@ func (a *ActiveAuthTokenService) ActiveTokenCount(ctx context.Context) (int64, e
 		return err
 	})
 
-	return count, err
+	tag, err := quota.NewTag(QuotaTargetSrv, QuotaTarget, quota.GlobalScope)
+	if err != nil {
+		return nil, err
+	}
+	u := &quota.Map{}
+	u.Set(tag, count)
+
+	return u, err
 }
 
 func (s *UserAuthTokenService) CreateToken(ctx context.Context, user *user.User, clientIP net.IP, userAgent string) (*models.UserToken, error) {
@@ -472,3 +499,19 @@ func hashToken(token string) string {
 	hashBytes := sha256.Sum256([]byte(token + setting.SecretKey))
 	return hex.EncodeToString(hashBytes[:])
 }
+
+func readQuotaConfig(cfg *setting.Cfg) (*quota.Map, error) {
+	limits := &quota.Map{}
+
+	if cfg == nil {
+		return limits, nil
+	}
+
+	globalQuotaTag, err := quota.NewTag(QuotaTargetSrv, QuotaTarget, quota.GlobalScope)
+	if err != nil {
+		return limits, err
+	}
+
+	limits.Set(globalQuotaTag, cfg.Quota.Global.Session)
+	return limits, nil
+}

pkg/services/auth/auth_token_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/grafana/grafana/pkg/infra/db"
 	"github.com/grafana/grafana/pkg/infra/log"
 	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/services/quota"
 	"github.com/grafana/grafana/pkg/services/user"
 	"github.com/grafana/grafana/pkg/setting"
 )
@@ -40,8 +41,12 @@ func TestUserAuthToken(t *testing.T) {
 		userToken := createToken()
 
 		t.Run("Can count active tokens", func(t *testing.T) {
-			count, err := ctx.activeTokenService.ActiveTokenCount(context.Background())
+			m, err := ctx.activeTokenService.ActiveTokenCount(context.Background(), &quota.ScopeParameters{})
 			require.Nil(t, err)
+			tag, err := quota.NewTag(QuotaTargetSrv, QuotaTarget, quota.GlobalScope)
+			require.NoError(t, err)
+			count, ok := m.Get(tag)
+			require.True(t, ok)
 			require.Equal(t, int64(1), count)
 		})
 
@@ -208,8 +213,12 @@ func TestUserAuthToken(t *testing.T) {
 			require.Nil(t, notGood)
 
 			t.Run("should not find active token when expired", func(t *testing.T) {
-				count, err := ctx.activeTokenService.ActiveTokenCount(context.Background())
+				m, err := ctx.activeTokenService.ActiveTokenCount(context.Background(), &quota.ScopeParameters{})
 				require.Nil(t, err)
+				tag, err := quota.NewTag(QuotaTargetSrv, QuotaTarget, quota.GlobalScope)
+				require.NoError(t, err)
+				count, ok := m.Get(tag)
+				require.True(t, ok)
 				require.Equal(t, int64(0), count)
 			})
 		})

pkg/services/auth/model.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 
 	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/services/quota"
 )
 
 type userAuthToken struct {
@@ -71,3 +72,8 @@ func (uat *userAuthToken) toUserToken(ut *models.UserToken) error {
 
 	return nil
 }
+
+const (
+	QuotaTargetSrv quota.TargetSrv = "auth"
+	QuotaTarget    quota.Target    = "session"
+)