Răsfoiți Sursa

add usage tracker that reports project usage

Alexander Belanger 4 ani în urmă
părinte
comite
9fa804c14f

+ 5 - 1
api/server/handlers/project/get_usage.go

@@ -30,7 +30,11 @@ func (p *ProjectGetUsageHandler) ServeHTTP(w http.ResponseWriter, r *http.Reques
 
 	res := &types.GetProjectUsageResponse{}
 
-	currUsage, limit, usageCache, err := usage.GetUsage(p.Config(), proj)
+	currUsage, limit, usageCache, err := usage.GetUsage(&usage.GetUsageOpts{
+		Project: proj,
+		DOConf:  p.Config().DOConf,
+		Repo:    p.Repo(),
+	})
 
 	if err != nil {
 		p.HandleAPIError(w, r, apierrors.NewErrInternal(err))

+ 5 - 1
api/server/router/middleware/usage.go

@@ -27,7 +27,11 @@ func (b *UsageMiddleware) Middleware(next http.Handler) http.Handler {
 		proj, _ := r.Context().Value(types.ProjectScope).(*models.Project)
 
 		// get the project usage limits
-		currentUsage, limit, _, err := usage.GetUsage(b.config, proj)
+		currentUsage, limit, _, err := usage.GetUsage(&usage.GetUsageOpts{
+			Project: proj,
+			DOConf:  b.config.DOConf,
+			Repo:    b.config.Repo,
+		})
 
 		if err != nil {
 			apierrors.HandleAPIError(

+ 3 - 3
ee/usage/limit.go

@@ -5,15 +5,15 @@ package usage
 import (
 	"errors"
 
-	"github.com/porter-dev/porter/api/server/shared/config"
 	"github.com/porter-dev/porter/api/types"
 	"github.com/porter-dev/porter/internal/models"
+	"github.com/porter-dev/porter/internal/repository"
 	"gorm.io/gorm"
 )
 
-func GetLimit(config *config.Config, proj *models.Project) (limit *types.ProjectUsage, err error) {
+func GetLimit(repo repository.Repository, proj *models.Project) (limit *types.ProjectUsage, err error) {
 	// query for the project limit; if not found, default to basic
-	limitModel, err := config.Repo.ProjectUsage().ReadProjectUsage(proj.ID)
+	limitModel, err := repo.ProjectUsage().ReadProjectUsage(proj.ID)
 
 	if err != nil && errors.Is(err, gorm.ErrRecordNotFound) {
 		copyBasic := types.BasicPlan

+ 2 - 2
internal/usage/limit_ce.go

@@ -3,12 +3,12 @@
 package usage
 
 import (
-	"github.com/porter-dev/porter/api/server/shared/config"
 	"github.com/porter-dev/porter/api/types"
 	"github.com/porter-dev/porter/internal/models"
+	"github.com/porter-dev/porter/internal/repository"
 )
 
-func GetLimit(config *config.Config, proj *models.Project) (limit *types.ProjectUsage, err error) {
+func GetLimit(repo repository.Repository, proj *models.Project) (limit *types.ProjectUsage, err error) {
 	copyLimit := types.BasicPlan
 
 	return &copyLimit, nil

+ 2 - 2
internal/usage/limit_ee.go

@@ -3,13 +3,13 @@
 package usage
 
 import (
-	"github.com/porter-dev/porter/api/server/shared/config"
 	"github.com/porter-dev/porter/api/types"
 	"github.com/porter-dev/porter/ee/usage"
 	"github.com/porter-dev/porter/internal/models"
+	"github.com/porter-dev/porter/internal/repository"
 )
 
-var GetLimit func(config *config.Config, proj *models.Project) (limit *types.ProjectUsage, err error)
+var GetLimit func(repo repository.Repository, proj *models.Project) (limit *types.ProjectUsage, err error)
 
 func init() {
 	GetLimit = usage.GetLimit

+ 23 - 15
internal/usage/usage.go

@@ -5,42 +5,48 @@ import (
 	"fmt"
 	"time"
 
-	"github.com/porter-dev/porter/api/server/authz"
-	"github.com/porter-dev/porter/api/server/shared/config"
 	"github.com/porter-dev/porter/api/types"
 	"github.com/porter-dev/porter/internal/kubernetes"
 	"github.com/porter-dev/porter/internal/kubernetes/nodes"
 	"github.com/porter-dev/porter/internal/models"
+	"github.com/porter-dev/porter/internal/repository"
+	"golang.org/x/oauth2"
 	"gorm.io/gorm"
 )
 
+type GetUsageOpts struct {
+	Repo    repository.Repository
+	DOConf  *oauth2.Config
+	Project *models.Project
+}
+
 // GetUsage gets a project's current usage and usage limit
-func GetUsage(config *config.Config, proj *models.Project) (
+func GetUsage(opts *GetUsageOpts) (
 	current, limit *types.ProjectUsage,
 	resourceUse *models.ProjectUsageCache,
 	err error,
 ) {
-	limit, err = GetLimit(config, proj)
+	limit, err = GetLimit(opts.Repo, opts.Project)
 
 	if err != nil {
 		return nil, nil, nil, err
 	}
 
 	// query for the linked cluster counts
-	clusters, err := config.Repo.Cluster().ListClustersByProjectID(proj.ID)
+	clusters, err := opts.Repo.Cluster().ListClustersByProjectID(opts.Project.ID)
 
 	if err != nil {
 		return nil, nil, nil, err
 	}
 
 	// query for the linked user counts
-	roles, err := config.Repo.Project().ListProjectRoles(proj.ID)
+	roles, err := opts.Repo.Project().ListProjectRoles(opts.Project.ID)
 
 	if err != nil {
 		return nil, nil, nil, err
 	}
 
-	usageCache, err := config.Repo.ProjectUsage().ReadProjectUsageCache(proj.ID)
+	usageCache, err := opts.Repo.ProjectUsage().ReadProjectUsageCache(opts.Project.ID)
 	isCacheFound := true
 
 	if isCacheFound = !errors.Is(err, gorm.ErrRecordNotFound); err != nil && isCacheFound {
@@ -50,7 +56,7 @@ func GetUsage(config *config.Config, proj *models.Project) (
 	// if the usage cache is 24 hours old, was not found, or usage is over limit,
 	// re-query for the usage
 	if !isCacheFound || usageCache.Is24HrOld() || usageCache.ResourceMemory > limit.ResourceMemory || usageCache.ResourceCPU > limit.ResourceCPU {
-		cpu, memory, err := getResourceUsage(config, clusters)
+		cpu, memory, err := getResourceUsage(opts, clusters)
 
 		if err != nil {
 			return nil, nil, nil, err
@@ -58,7 +64,7 @@ func GetUsage(config *config.Config, proj *models.Project) (
 
 		if !isCacheFound {
 			usageCache = &models.ProjectUsageCache{
-				ProjectID:      proj.ID,
+				ProjectID:      opts.Project.ID,
 				ResourceCPU:    cpu,
 				ResourceMemory: memory,
 			}
@@ -78,9 +84,9 @@ func GetUsage(config *config.Config, proj *models.Project) (
 		usageCache.Exceeded = isExceeded
 
 		if !isCacheFound {
-			usageCache, err = config.Repo.ProjectUsage().CreateProjectUsageCache(usageCache)
+			usageCache, err = opts.Repo.ProjectUsage().CreateProjectUsageCache(usageCache)
 		} else {
-			usageCache, err = config.Repo.ProjectUsage().UpdateProjectUsageCache(usageCache)
+			usageCache, err = opts.Repo.ProjectUsage().UpdateProjectUsageCache(usageCache)
 		}
 	}
 
@@ -93,13 +99,15 @@ func GetUsage(config *config.Config, proj *models.Project) (
 }
 
 // gets the total resource usage across all nodes in all clusters
-func getResourceUsage(config *config.Config, clusters []*models.Cluster) (uint, uint, error) {
-	// TODO; pass this in?
+func getResourceUsage(opts *GetUsageOpts, clusters []*models.Cluster) (uint, uint, error) {
 	var totCPU, totMem uint = 0, 0
-	getter := authz.NewOutOfClusterAgentGetter(config)
 
 	for _, cluster := range clusters {
-		ooc := getter.GetOutOfClusterConfig(cluster)
+		ooc := &kubernetes.OutOfClusterConfig{
+			Cluster:           cluster,
+			Repo:              opts.Repo,
+			DigitalOceanOAuth: opts.DOConf,
+		}
 
 		agent, err := kubernetes.GetAgentOutOfClusterConfig(ooc)
 

+ 97 - 0
services/usage/usage.go

@@ -0,0 +1,97 @@
+package usage
+
+import (
+	"time"
+
+	"github.com/porter-dev/porter/api/server/shared/config/env"
+	"github.com/porter-dev/porter/internal/adapter"
+	"github.com/porter-dev/porter/internal/models"
+	"github.com/porter-dev/porter/internal/repository"
+	"github.com/porter-dev/porter/internal/usage"
+	"golang.org/x/oauth2"
+	"gorm.io/gorm"
+
+	rgorm "github.com/porter-dev/porter/internal/repository/gorm"
+)
+
+type UsageTracker struct {
+	db     *gorm.DB
+	repo   repository.Repository
+	doConf *oauth2.Config
+}
+
+type UsageTrackerOpts struct {
+	DBConf *env.DBConf
+	DOConf *oauth2.Config
+}
+
+const stepSize = 100
+
+func NewUsageTracker(opts *UsageTrackerOpts) (*UsageTracker, error) {
+	db, err := adapter.New(opts.DBConf)
+
+	if err != nil {
+		return nil, err
+	}
+
+	var key [32]byte
+
+	for i, b := range []byte(opts.DBConf.EncryptionKey) {
+		key[i] = b
+	}
+
+	repo := rgorm.NewRepository(db, &key)
+
+	return &UsageTracker{db, repo, opts.DOConf}, nil
+}
+
+type UsageTrackerResponse struct {
+	ResourceCPU    uint
+	ResourceMemory uint
+	Exceeded       bool
+	ExceededSince  *time.Time
+}
+
+func (u *UsageTracker) GetProjectUsage() (map[uint]*UsageTrackerResponse, error) {
+	res := make(map[uint]*UsageTrackerResponse)
+
+	// get the count of the projects
+	var count int64
+
+	if err := u.db.Model(&models.Project{}).Count(&count).Error; err != nil {
+		return nil, err
+	}
+
+	// iterate (count / stepSize) + 1 times using Limit and Offset
+	for i := 0; i < (int(count)/stepSize)+1; i++ {
+		projects := []*models.Project{}
+
+		if err := u.db.Order("id asc").Offset(i * stepSize).Limit(stepSize).Find(&projects).Error; err != nil {
+			return nil, err
+		}
+
+		// go through each project
+		for _, project := range projects {
+			_, _, cache, err := usage.GetUsage(&usage.GetUsageOpts{
+				Repo:    u.repo,
+				DOConf:  u.doConf,
+				Project: project,
+			})
+
+			if err != nil {
+				continue
+			}
+
+			if cache.Exceeded {
+				res[project.ID] = &UsageTrackerResponse{
+					ResourceCPU:    cache.ResourceCPU,
+					ResourceMemory: cache.ResourceMemory,
+					Exceeded:       cache.Exceeded,
+					ExceededSince:  cache.ExceededSince,
+				}
+			}
+		}
+	}
+
+	return res, nil
+}