porter/services/usage/usage.go

// +build ee

package usage

import (
	"fmt"
	"sync"
	"time"

	"github.com/porter-dev/porter/api/server/shared/config/env"
	"github.com/porter-dev/porter/api/types"
	"github.com/porter-dev/porter/ee/integrations/vault"
	"github.com/porter-dev/porter/internal/adapter"
	"github.com/porter-dev/porter/internal/models"
	"github.com/porter-dev/porter/internal/oauth"
	"github.com/porter-dev/porter/internal/repository"
	"github.com/porter-dev/porter/internal/usage"
	"golang.org/x/oauth2"
	"gorm.io/gorm"

	"github.com/porter-dev/porter/internal/repository/credentials"
	rgorm "github.com/porter-dev/porter/internal/repository/gorm"
)

type UsageTracker struct {
	db               *gorm.DB
	repo             repository.Repository
	doConf           *oauth2.Config
	whitelistedUsers map[uint]uint
}

type UsageTrackerOpts struct {
	DBConf           *env.DBConf
	DOClientID       string
	DOClientSecret   string
	DOScopes         []string
	ServerURL        string
	WhitelistedUsers map[uint]uint
}

const stepSize = 100

func NewUsageTracker(opts *UsageTrackerOpts) (*UsageTracker, error) {
	db, err := adapter.New(opts.DBConf)

	if err != nil {
		return nil, err
	}

	var credBackend credentials.CredentialStorage

	if opts.DBConf.VaultAPIKey != "" && opts.DBConf.VaultServerURL != "" && opts.DBConf.VaultPrefix != "" {
		credBackend = vault.NewClient(
			opts.DBConf.VaultServerURL,
			opts.DBConf.VaultAPIKey,
			opts.DBConf.VaultPrefix,
		)
	}

	var key [32]byte

	for i, b := range []byte(opts.DBConf.EncryptionKey) {
		key[i] = b
	}

	repo := rgorm.NewRepository(db, &key, credBackend)

	doConf := oauth.NewDigitalOceanClient(&oauth.Config{
		ClientID:     opts.DOClientID,
		ClientSecret: opts.DOClientSecret,
		Scopes:       opts.DOScopes,
		BaseURL:      opts.ServerURL,
	})

	return &UsageTracker{db, repo, doConf, opts.WhitelistedUsers}, nil
}

type UsageTrackerResponse struct {
	CPULimit      uint
	CPUUsage      uint
	MemoryLimit   uint
	MemoryUsage   uint
	UserLimit     uint
	UserUsage     uint
	ClusterLimit  uint
	ClusterUsage  uint
	Exceeded      bool
	ExceededSince time.Time
	Project       models.Project
	AdminEmails   []string
}

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
	}

	var mu sync.Mutex
	var wg sync.WaitGroup

	worker := func(project *models.Project) {
		defer wg.Done()

		current, limit, cache, err := usage.GetUsage(&usage.GetUsageOpts{
			Repo:             u.repo,
			DOConf:           u.doConf,
			Project:          project,
			WhitelistedUsers: u.whitelistedUsers,
		})

		if err != nil {
			fmt.Printf("Project %d: error getting usage: %v\n", project.ID, err)
			return
		}

		// get the admin emails for the project
		roles, err := u.repo.Project().ListProjectRoles(project.ID)

		if err != nil {
			fmt.Printf("Project %d: error getting admin emails: %v\n", project.ID, err)
			return
		}

		adminEmails := make([]string, 0)

		for _, role := range roles {
			if role.Kind == types.RoleAdmin {
				user, err := u.repo.User().ReadUser(role.UserID)

				if err != nil {
					continue
				}

				adminEmails = append(adminEmails, user.Email)
			}
		}

		exceededSince := cache.ExceededSince

		if exceededSince == nil {
			now := time.Now()
			exceededSince = &now
		}

		mu.Lock()
		res[project.ID] = &UsageTrackerResponse{
			CPUUsage:      cache.ResourceCPU,
			CPULimit:      limit.ResourceCPU,
			MemoryUsage:   cache.ResourceMemory,
			MemoryLimit:   limit.ResourceMemory,
			UserUsage:     current.Users,
			UserLimit:     limit.Users,
			ClusterUsage:  current.Clusters,
			ClusterLimit:  limit.Clusters,
			Exceeded:      cache.Exceeded,
			ExceededSince: *exceededSince,
			Project:       *project,
			AdminEmails:   adminEmails,
		}
		mu.Unlock()

		return
	}

	// 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 {
			wg.Add(1)
			go worker(project)
		}

		wg.Wait()
	}

	return res, nil
}