diff --git a/components/usage/cmd/run.go b/components/usage/cmd/run.go index f3eab74c50444a..3558ed3f4d1752 100644 --- a/components/usage/cmd/run.go +++ b/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)) diff --git a/components/usage/go.mod b/components/usage/go.mod index 444d4d6458588f..1ac0cd88a164aa 100644 --- a/components/usage/go.mod +++ b/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 diff --git a/components/usage/pkg/controller/billing.go b/components/usage/pkg/controller/billing.go index 8fe1da66a1ce08..381e759f41f687 100644 --- a/components/usage/pkg/controller/billing.go +++ b/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] +} diff --git a/components/usage/pkg/controller/billing_test.go b/components/usage/pkg/controller/billing_test.go new file mode 100644 index 00000000000000..42d8ff8498ceca --- /dev/null +++ b/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) + }) + } +} diff --git a/components/usage/pkg/controller/reconciler.go b/components/usage/pkg/controller/reconciler.go index 32db89f829a010..299ed664a0facd 100644 --- a/components/usage/pkg/controller/reconciler.go +++ b/components/usage/pkg/controller/reconciler.go @@ -116,8 +116,8 @@ 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() @@ -125,15 +125,17 @@ func (u UsageReport) RuntimeSummaryForTeams(maxStopTime time.Time) map[string]in 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 { diff --git a/components/usage/pkg/controller/reconciler_test.go b/components/usage/pkg/controller/reconciler_test.go index dcff4bca4f39d2..aa7c1de7cdf29b 100644 --- a/components/usage/pkg/controller/reconciler_test.go +++ b/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)) } diff --git a/components/usage/pkg/db/workspace_instance.go b/components/usage/pkg/db/workspace_instance.go index f730bca2217cfe..eaa8dd3a6cfe2b 100644 --- a/components/usage/pkg/db/workspace_instance.go +++ b/components/usage/pkg/db/workspace_instance.go @@ -46,7 +46,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 @@ -56,7 +56,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 diff --git a/components/usage/pkg/stripe/stripe.go b/components/usage/pkg/stripe/stripe.go index 462cf45d402af5..57b58d2b1c2b9a 100644 --- a/components/usage/pkg/stripe/stripe.go +++ b/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))) -} diff --git a/components/usage/pkg/stripe/stripe_test.go b/components/usage/pkg/stripe/stripe_test.go index 3e88fb3084ffbd..c628e372577a59 100644 --- a/components/usage/pkg/stripe/stripe_test.go +++ b/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) - }) - } -}