diff --git a/components/usage/pkg/controller/billing.go b/components/usage/pkg/controller/billing.go index cb6c7c1c5f142f..b2b3782704af81 100644 --- a/components/usage/pkg/controller/billing.go +++ b/components/usage/pkg/controller/billing.go @@ -7,9 +7,11 @@ package controller import ( "context" "fmt" - "github.com/gitpod-io/gitpod/usage/pkg/stripe" "math" "time" + + "github.com/gitpod-io/gitpod/usage/pkg/db" + "github.com/gitpod-io/gitpod/usage/pkg/stripe" ) type BillingController interface { @@ -67,6 +69,15 @@ type WorkspacePricer struct { creditMinutesByWorkspaceClass map[string]float64 } +func (p *WorkspacePricer) CreditsUsedByInstance(instance *db.WorkspaceInstanceForUsage, maxStopTime time.Time) int64 { + runtime := instance.WorkspaceRuntimeSeconds(maxStopTime) + class := defaultWorkspaceClass + if instance.WorkspaceClass != "" { + class = instance.WorkspaceClass + } + return p.Credits(class, runtime) +} + func (p *WorkspacePricer) Credits(workspaceClass string, runtimeInSeconds int64) int64 { inMinutes := float64(runtimeInSeconds) / 60 return int64(math.Ceil(p.CreditsPerMinuteForClass(workspaceClass) * inMinutes)) diff --git a/components/usage/pkg/controller/reconciler.go b/components/usage/pkg/controller/reconciler.go index d38987e1d3e31e..a651a4098ff681 100644 --- a/components/usage/pkg/controller/reconciler.go +++ b/components/usage/pkg/controller/reconciler.go @@ -33,11 +33,17 @@ func (f ReconcilerFunc) Reconcile() error { type UsageReconciler struct { nowFunc func() time.Time conn *gorm.DB + pricer *WorkspacePricer billingController BillingController } -func NewUsageReconciler(conn *gorm.DB, billingController BillingController) *UsageReconciler { - return &UsageReconciler{conn: conn, billingController: billingController, nowFunc: time.Now} +func NewUsageReconciler(conn *gorm.DB, pricer *WorkspacePricer, billingController BillingController) *UsageReconciler { + return &UsageReconciler{ + conn: conn, + pricer: pricer, + billingController: billingController, + nowFunc: time.Now, + } } type UsageReconcileStatus struct { @@ -86,7 +92,7 @@ func (u *UsageReconciler) Reconcile() (err error) { } log.Infof("Wrote usage report into %s", filepath.Join(dir, stat.Name())) - err = db.CreateUsageRecords(ctx, u.conn, usageReportToUsageRecords(report)) + err = db.CreateUsageRecords(ctx, u.conn, usageReportToUsageRecords(report, u.pricer, u.nowFunc().UTC())) if err != nil { return fmt.Errorf("failed to write usage records to database: %s", err) } @@ -136,12 +142,7 @@ func (u UsageReport) CreditSummaryForTeams(pricer *WorkspacePricer, maxStopTime var credits int64 for _, instance := range instances { - runtime := instance.WorkspaceRuntimeSeconds(maxStopTime) - class := defaultWorkspaceClass - if instance.WorkspaceClass != "" { - class = instance.WorkspaceClass - } - credits += pricer.Credits(class, runtime) + credits += pricer.CreditsUsedByInstance(&instance, maxStopTime) } creditsPerTeamID[id] = credits @@ -232,7 +233,7 @@ func groupInstancesByAttributionID(instances []db.WorkspaceInstanceForUsage) Usa return result } -func usageReportToUsageRecords(report UsageReport) []db.WorkspaceInstanceUsage { +func usageReportToUsageRecords(report UsageReport, pricer *WorkspacePricer, now time.Time) []db.WorkspaceInstanceUsage { var usageRecords []db.WorkspaceInstanceUsage for attributionId, instances := range report { @@ -247,7 +248,7 @@ func usageReportToUsageRecords(report UsageReport) []db.WorkspaceInstanceUsage { AttributionID: attributionId, StartedAt: instance.CreationTime.Time(), StoppedAt: stoppedAt, - CreditsUsed: 0, + CreditsUsed: float64(pricer.CreditsUsedByInstance(&instance, now)), GenerationId: 0, Deleted: false, }) diff --git a/components/usage/pkg/controller/reconciler_test.go b/components/usage/pkg/controller/reconciler_test.go index 97a77975c193ab..bb5512e1aaa4bd 100644 --- a/components/usage/pkg/controller/reconciler_test.go +++ b/components/usage/pkg/controller/reconciler_test.go @@ -149,6 +149,7 @@ func TestUsageReport_CreditSummaryForTeams(t *testing.T) { } func TestUsageReportConversionToDBUsageRecords(t *testing.T) { + maxStopTime := time.Date(2022, 05, 31, 23, 00, 00, 00, time.UTC) teamID := uuid.New().String() teamAttributionID := db.NewTeamAttributionID(teamID) instanceId := uuid.New() @@ -178,7 +179,7 @@ func TestUsageReportConversionToDBUsageRecords(t *testing.T) { AttributionID: teamAttributionID, StartedAt: creationTime.Time(), StoppedAt: sql.NullTime{Time: stoppedTime.Time(), Valid: true}, - CreditsUsed: 0, + CreditsUsed: 470, GenerationId: 0, }}, }, @@ -200,7 +201,7 @@ func TestUsageReportConversionToDBUsageRecords(t *testing.T) { AttributionID: teamAttributionID, StartedAt: creationTime.Time(), StoppedAt: sql.NullTime{}, - CreditsUsed: 0, + CreditsUsed: 470, GenerationId: 0, }}, }, @@ -208,7 +209,7 @@ func TestUsageReportConversionToDBUsageRecords(t *testing.T) { for _, s := range scenarios { t.Run(s.Name, func(t *testing.T) { - actual := usageReportToUsageRecords(s.Report) + actual := usageReportToUsageRecords(s.Report, DefaultWorkspacePricer, maxStopTime) require.Equal(t, s.Expected, actual) }) } diff --git a/components/usage/pkg/server/server.go b/components/usage/pkg/server/server.go index 88d5164dd5712c..22a9ee95eae952 100644 --- a/components/usage/pkg/server/server.go +++ b/components/usage/pkg/server/server.go @@ -43,6 +43,11 @@ func Start(cfg Config) error { return fmt.Errorf("failed to establish database connection: %w", err) } + pricer, err := controller.NewWorkspacePricer(cfg.CreditsPerMinuteByWorkspaceClass) + if err != nil { + return fmt.Errorf("failed to create workspace pricer: %w", err) + } + var billingController controller.BillingController = &controller.NoOpBillingController{} if cfg.StripeCredentialsFile != "" { @@ -56,11 +61,6 @@ func Start(cfg Config) error { return fmt.Errorf("failed to initialize stripe client: %w", err) } - pricer, err := controller.NewWorkspacePricer(cfg.CreditsPerMinuteByWorkspaceClass) - if err != nil { - return fmt.Errorf("failed to create workspace pricer: %w", err) - } - billingController = controller.NewStripeBillingController(c, pricer) } @@ -69,7 +69,7 @@ func Start(cfg Config) error { return fmt.Errorf("failed to parse schedule duration: %w", err) } - ctrl, err := controller.New(schedule, controller.NewUsageReconciler(conn, billingController)) + ctrl, err := controller.New(schedule, controller.NewUsageReconciler(conn, pricer, billingController)) if err != nil { return fmt.Errorf("failed to initialize usage controller: %w", err) }