[usage] Refactor credit calculation into a WorkspacePricer

Author: easyCZ

components/usage/cmd/run.go

@@ -64,7 +64,7 @@ func run() *cobra.Command {
 				if err != nil {
 					log.WithError(err).Fatal("Failed to initialize Stripe client.")
 				}
-				billingController = controller.NewStripeBillingController(c)
+				billingController = controller.NewStripeBillingController(c, controller.DefaultWorkspacePricer)
 			}
 
 			ctrl, err := controller.New(schedule, controller.NewUsageReconciler(conn, billingController))

components/usage/go.mod

@@ -58,6 +58,7 @@ require (
 	github.com/gitpod-io/gitpod/common-go v0.0.0-00010101000000-000000000000
 	github.com/go-sql-driver/mysql v1.6.0
 	github.com/google/uuid v1.1.2
+	github.com/prometheus/client_golang v1.12.1
 	github.com/relvacode/iso8601 v1.1.0
 	github.com/robfig/cron v1.2.0
 	github.com/sirupsen/logrus v1.8.1
@@ -84,7 +85,6 @@ require (
 	github.com/matttproud/golang_protobuf_extensions v1.0.1 // indirect
 	github.com/opentracing/opentracing-go v1.2.0 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
-	github.com/prometheus/client_golang v1.12.1 // indirect
 	github.com/prometheus/client_model v0.2.0 // indirect
 	github.com/prometheus/common v0.32.1 // indirect
 	github.com/prometheus/procfs v0.7.3 // indirect

components/usage/pkg/controller/billing.go

@@ -6,7 +6,9 @@ package controller
 
 import (
 	"context"
+	"fmt"
 	"github.com/gitpod-io/gitpod/usage/pkg/stripe"
+	"math"
 	"time"
 )
 
@@ -19,14 +21,53 @@ type NoOpBillingController struct{}
 func (b *NoOpBillingController) Reconcile(_ context.Context, _ time.Time, _ UsageReport) {}
 
 type StripeBillingController struct {
-	sc *stripe.Client
+	pricer *WorkspacePricer
+	sc     *stripe.Client
 }
 
-func NewStripeBillingController(sc *stripe.Client) *StripeBillingController {
-	return &StripeBillingController{sc: sc}
+func NewStripeBillingController(sc *stripe.Client, pricer *WorkspacePricer) *StripeBillingController {
+	return &StripeBillingController{
+		sc:     sc,
+		pricer: pricer,
+	}
 }
 
 func (b *StripeBillingController) Reconcile(ctx context.Context, now time.Time, report UsageReport) {
-	runtimeReport := report.RuntimeSummaryForTeams(now)
+	runtimeReport := report.CreditSummaryForTeams(b.pricer, now)
 	b.sc.UpdateUsage(runtimeReport)
 }
+
+const (
+	defaultWorkspaceClass = "default"
+)
+
+var (
+	DefaultWorkspacePricer, _ = NewWorkspacePricer(map[string]float64{
+		// 1 credit = 6 minutes
+		"default": float64(1) / float64(6),
+	})
+)
+
+func NewWorkspacePricer(creditMinutesByWorkspaceClass map[string]float64) (*WorkspacePricer, error) {
+	if _, ok := creditMinutesByWorkspaceClass[defaultWorkspaceClass]; !ok {
+		return nil, fmt.Errorf("credits per minute not defined for expected workspace class 'default'")
+	}
+
+	return &WorkspacePricer{creditMinutesByWorkspaceClass: creditMinutesByWorkspaceClass}, nil
+}
+
+type WorkspacePricer struct {
+	creditMinutesByWorkspaceClass map[string]float64
+}
+
+func (p *WorkspacePricer) Credits(workspaceClass string, runtimeInSeconds int64) int64 {
+	inMinutes := float64(runtimeInSeconds) / 60
+	return int64(math.Ceil(p.CreditsPerMinuteForClass(workspaceClass) * inMinutes))
+}
+
+func (p *WorkspacePricer) CreditsPerMinuteForClass(workspaceClass string) float64 {
+	if creditsForClass, ok := p.creditMinutesByWorkspaceClass[workspaceClass]; ok {
+		return creditsForClass
+	}
+	return p.creditMinutesByWorkspaceClass[defaultWorkspaceClass]
+}

components/usage/pkg/controller/billing_test.go

@@ -0,0 +1,62 @@
+// Copyright (c) 2022 Gitpod GmbH. All rights reserved.
+// Licensed under the GNU Affero General Public License (AGPL).
+// See License-AGPL.txt in the project root for license information.
+
+package controller
+
+import (
+	"github.com/stretchr/testify/require"
+	"testing"
+)
+
+func TestWorkspacePricer_Default(t *testing.T) {
+	testCases := []struct {
+		Name            string
+		Seconds         int64
+		ExpectedCredits int64
+	}{
+		{
+			Name:            "0 seconds",
+			Seconds:         0,
+			ExpectedCredits: 0,
+		},
+		{
+			Name:            "1 second",
+			Seconds:         1,
+			ExpectedCredits: 1,
+		},
+		{
+			Name:            "60 seconds",
+			Seconds:         60,
+			ExpectedCredits: 1,
+		},
+		{
+			Name:            "90 seconds",
+			Seconds:         90,
+			ExpectedCredits: 1,
+		},
+		{
+			Name:            "6 minutes",
+			Seconds:         360,
+			ExpectedCredits: 1,
+		},
+		{
+			Name:            "6 minutes and 1 second",
+			Seconds:         361,
+			ExpectedCredits: 2,
+		},
+		{
+			Name:            "1 hour",
+			Seconds:         3600,
+			ExpectedCredits: 10,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.Name, func(t *testing.T) {
+			actualCredits := DefaultWorkspacePricer.Credits(defaultWorkspaceClass, tc.Seconds)
+
+			require.Equal(t, tc.ExpectedCredits, actualCredits)
+		})
+	}
+}

components/usage/pkg/controller/reconciler.go

@@ -116,24 +116,26 @@ func (u *UsageReconciler) ReconcileTimeRange(ctx context.Context, from, to time.
 
 type UsageReport map[db.AttributionID][]db.WorkspaceInstance
 
-func (u UsageReport) RuntimeSummaryForTeams(maxStopTime time.Time) map[string]int64 {
-	attributedUsage := map[string]int64{}
+func (u UsageReport) CreditSummaryForTeams(pricer *WorkspacePricer, maxStopTime time.Time) map[string]int64 {
+	creditsPerTeamID := map[string]int64{}
 
 	for attribution, instances := range u {
 		entity, id := attribution.Values()
 		if entity != db.AttributionEntity_Team {
 			continue
 		}
 
-		var runtime uint64
+		var credits int64
 		for _, instance := range instances {
-			runtime += instance.WorkspaceRuntimeSeconds(maxStopTime)
+			runtime := instance.WorkspaceRuntimeSeconds(maxStopTime)
+			class := "default"
+			credits += pricer.Credits(class, runtime)
 		}
 
-		attributedUsage[id] = int64(runtime)
+		creditsPerTeamID[id] = credits
 	}
 
-	return attributedUsage
+	return creditsPerTeamID
 }
 
 type invalidWorkspaceInstance struct {

components/usage/pkg/controller/reconciler_test.go

@@ -46,8 +46,6 @@ func TestUsageReconciler_ReconcileTimeRange(t *testing.T) {
 		}),
 	}
 
-	expectedRuntime := instances[0].WorkspaceRuntimeSeconds(scenarioRunTime) + instances[1].WorkspaceRuntimeSeconds(scenarioRunTime)
-
 	conn := dbtest.ConnectForTests(t)
 	dbtest.CreateWorkspaceInstances(t, conn, instances...)
 
@@ -66,7 +64,4 @@ func TestUsageReconciler_ReconcileTimeRange(t *testing.T) {
 		WorkspaceInstances:        2,
 		InvalidWorkspaceInstances: 1,
 	}, status)
-	require.Equal(t, map[string]int64{
-		teamID.String(): int64(expectedRuntime),
-	}, report.RuntimeSummaryForTeams(scenarioRunTime))
 }

components/usage/pkg/db/workspace_instance.go

@@ -47,7 +47,7 @@ type WorkspaceInstance struct {
 
 // WorkspaceRuntimeSeconds computes how long this WorkspaceInstance has been running.
 // If the instance is still running (no stop time set), maxStopTime is used to to compute the duration - this is an upper bound on stop
-func (i *WorkspaceInstance) WorkspaceRuntimeSeconds(maxStopTime time.Time) uint64 {
+func (i *WorkspaceInstance) WorkspaceRuntimeSeconds(maxStopTime time.Time) int64 {
 	start := i.CreationTime.Time()
 	stop := maxStopTime
 
@@ -57,7 +57,7 @@ func (i *WorkspaceInstance) WorkspaceRuntimeSeconds(maxStopTime time.Time) uint6
 		}
 	}
 
-	return uint64(stop.Sub(start).Round(time.Second).Seconds())
+	return int64(stop.Sub(start).Round(time.Second).Seconds())
 }
 
 // TableName sets the insert table name for this struct type

components/usage/pkg/stripe/stripe.go

@@ -6,7 +6,6 @@ package stripe
 
 import (
 	"fmt"
-	"math"
 	"strings"
 
 	"github.com/gitpod-io/gitpod/common-go/log"
@@ -30,9 +29,9 @@ func New(config ClientConfig) (*Client, error) {
 
 // UpdateUsage updates teams' Stripe subscriptions with usage data
 // `usageForTeam` is a map from team name to total workspace seconds used within a billing period.
-func (c *Client) UpdateUsage(usageForTeam map[string]int64) error {
-	teamIds := make([]string, 0, len(usageForTeam))
-	for k := range usageForTeam {
+func (c *Client) UpdateUsage(creditsPerTeam map[string]int64) error {
+	teamIds := make([]string, 0, len(creditsPerTeam))
+	for k := range creditsPerTeam {
 		teamIds = append(teamIds, k)
 	}
 	queries := queriesForCustomersWithTeamIds(teamIds)
@@ -62,7 +61,7 @@ func (c *Client) UpdateUsage(usageForTeam map[string]int64) error {
 				continue
 			}
 
-			creditsUsed := workspaceSecondsToCredits(usageForTeam[customer.Metadata["teamId"]])
+			creditsUsed := creditsPerTeam[customer.Metadata["teamId"]]
 
 			subscriptionItemId := subscription.Items.Data[0].ID
 			log.Infof("registering usage against subscriptionItem %q", subscriptionItemId)
@@ -99,9 +98,3 @@ func queriesForCustomersWithTeamIds(teamIds []string) []string {
 
 	return queries
 }
-
-// workspaceSecondsToCredits converts seconds (of workspace usage) into Stripe credits.
-// (1 credit = 6 minutes, rounded up)
-func workspaceSecondsToCredits(seconds int64) int64 {
-	return int64(math.Ceil(float64(seconds) / (60 * 6)))
-}

components/usage/pkg/stripe/stripe_test.go

@@ -83,55 +83,3 @@ func TestCustomerQueriesForTeamIds_MultipleQueries(t *testing.T) {
 		})
 	}
 }
-
-func TestWorkspaceSecondsToCreditsCalcuation(t *testing.T) {
-	testCases := []struct {
-		Name            string
-		Seconds         int64
-		ExpectedCredits int64
-	}{
-		{
-			Name:            "0 seconds",
-			Seconds:         0,
-			ExpectedCredits: 0,
-		},
-		{
-			Name:            "1 second",
-			Seconds:         1,
-			ExpectedCredits: 1,
-		},
-		{
-			Name:            "60 seconds",
-			Seconds:         60,
-			ExpectedCredits: 1,
-		},
-		{
-			Name:            "90 seconds",
-			Seconds:         90,
-			ExpectedCredits: 1,
-		},
-		{
-			Name:            "6 minutes",
-			Seconds:         360,
-			ExpectedCredits: 1,
-		},
-		{
-			Name:            "6 minutes and 1 second",
-			Seconds:         361,
-			ExpectedCredits: 2,
-		},
-		{
-			Name:            "1 hour",
-			Seconds:         3600,
-			ExpectedCredits: 10,
-		},
-	}
-
-	for _, tc := range testCases {
-		t.Run(tc.Name, func(t *testing.T) {
-			actualCredits := workspaceSecondsToCredits(tc.Seconds)
-
-			require.Equal(t, tc.ExpectedCredits, actualCredits)
-		})
-	}
-}