diff --git a/components/usage/pkg/apiv1/billing.go b/components/usage/pkg/apiv1/billing.go new file mode 100644 index 00000000000000..0cdb031cf43500 --- /dev/null +++ b/components/usage/pkg/apiv1/billing.go @@ -0,0 +1,74 @@ +// 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 apiv1 + +import ( + "context" + "fmt" + "github.com/gitpod-io/gitpod/common-go/log" + v1 "github.com/gitpod-io/gitpod/usage-api/v1" + "github.com/gitpod-io/gitpod/usage/pkg/db" + "github.com/gitpod-io/gitpod/usage/pkg/stripe" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" + "math" +) + +func NewBillingService(stripeClient *stripe.Client) *BillingService { + return &BillingService{ + stripeClient: stripeClient, + } +} + +type BillingService struct { + stripeClient *stripe.Client + + v1.UnimplementedBillingServiceServer +} + +func (s *BillingService) UpdateInvoices(ctx context.Context, in *v1.UpdateInvoicesRequest) (*v1.UpdateInvoicesResponse, error) { + credits, err := creditSummaryForTeams(in.GetSessions()) + if err != nil { + log.Log.WithError(err).Errorf("Failed to compute credit summary.") + return nil, status.Errorf(codes.InvalidArgument, "failed to compute credit summary") + } + + err = s.stripeClient.UpdateUsage(ctx, credits) + if err != nil { + log.Log.WithError(err).Errorf("Failed to update stripe invoices.") + return nil, status.Errorf(codes.Internal, "failed to update stripe invoices") + } + + return &v1.UpdateInvoicesResponse{}, nil +} + +func creditSummaryForTeams(sessions []*v1.BilledSession) (map[string]int64, error) { + creditsPerTeamID := map[string]float64{} + + for _, session := range sessions { + attributionID, err := db.ParseAttributionID(session.AttributionId) + if err != nil { + return nil, fmt.Errorf("failed to parse attribution ID: %w", err) + } + + entity, id := attributionID.Values() + if entity != db.AttributionEntity_Team { + continue + } + + if _, ok := creditsPerTeamID[id]; !ok { + creditsPerTeamID[id] = 0 + } + + creditsPerTeamID[id] += session.GetCredits() + } + + rounded := map[string]int64{} + for teamID, credits := range creditsPerTeamID { + rounded[teamID] = int64(math.Ceil(credits)) + } + + return rounded, nil +} diff --git a/components/usage/pkg/apiv1/billing_noop.go b/components/usage/pkg/apiv1/billing_noop.go new file mode 100644 index 00000000000000..6172e9d78d3896 --- /dev/null +++ b/components/usage/pkg/apiv1/billing_noop.go @@ -0,0 +1,21 @@ +// 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 apiv1 + +import ( + "context" + "github.com/gitpod-io/gitpod/common-go/log" + v1 "github.com/gitpod-io/gitpod/usage-api/v1" +) + +// BillingServiceNoop is used for Self-Hosted installations +type BillingServiceNoop struct { + v1.UnimplementedBillingServiceServer +} + +func (s *BillingServiceNoop) UpdateInvoices(_ context.Context, _ *v1.UpdateInvoicesRequest) (*v1.UpdateInvoicesResponse, error) { + log.Log.Infof("UpdateInvoices RPC invoked in no-op mode, no invoices will be updated.") + return &v1.UpdateInvoicesResponse{}, nil +} diff --git a/components/usage/pkg/apiv1/billing_test.go b/components/usage/pkg/apiv1/billing_test.go new file mode 100644 index 00000000000000..8ce7ef7a5e8c4e --- /dev/null +++ b/components/usage/pkg/apiv1/billing_test.go @@ -0,0 +1,86 @@ +// 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 apiv1 + +import ( + v1 "github.com/gitpod-io/gitpod/usage-api/v1" + "github.com/gitpod-io/gitpod/usage/pkg/db" + "github.com/google/uuid" + "github.com/stretchr/testify/require" + "testing" +) + +func TestCreditSummaryForTeams(t *testing.T) { + teamID_A, teamID_B := uuid.New().String(), uuid.New().String() + teamAttributionID_A, teamAttributionID_B := db.NewTeamAttributionID(teamID_A), db.NewTeamAttributionID(teamID_B) + + scenarios := []struct { + Name string + Sessions []*v1.BilledSession + Expected map[string]int64 + }{ + { + Name: "no instances in report, no summary", + Sessions: []*v1.BilledSession{}, + Expected: map[string]int64{}, + }, + { + Name: "skips user attributions", + Sessions: []*v1.BilledSession{ + { + AttributionId: string(db.NewUserAttributionID(uuid.New().String())), + }, + }, + Expected: map[string]int64{}, + }, + { + Name: "two workspace instances", + Sessions: []*v1.BilledSession{ + { + // has 1 day and 23 hours of usage + AttributionId: string(teamAttributionID_A), + Credits: (24 + 23) * 10, + }, + { + // has 1 hour of usage + AttributionId: string(teamAttributionID_A), + Credits: 10, + }, + }, + Expected: map[string]int64{ + // total of 2 days runtime, at 10 credits per hour, that's 480 credits + teamID_A: 480, + }, + }, + { + Name: "multiple teams", + Sessions: []*v1.BilledSession{ + { + // has 12 hours of usage + AttributionId: string(teamAttributionID_A), + Credits: (12) * 10, + }, + { + // has 1 day of usage + AttributionId: string(teamAttributionID_B), + Credits: (24) * 10, + }, + }, + Expected: map[string]int64{ + // total of 2 days runtime, at 10 credits per hour, that's 480 credits + teamID_A: 120, + teamID_B: 240, + }, + }, + } + + for _, s := range scenarios { + t.Run(s.Name, func(t *testing.T) { + actual, err := creditSummaryForTeams(s.Sessions) + require.NoError(t, err) + require.Equal(t, s.Expected, actual) + }) + } +} diff --git a/components/usage/pkg/controller/billing.go b/components/usage/pkg/controller/pricer.go similarity index 69% rename from components/usage/pkg/controller/billing.go rename to components/usage/pkg/controller/pricer.go index 22fd5bc8ed3843..cbd2291a59e699 100644 --- a/components/usage/pkg/controller/billing.go +++ b/components/usage/pkg/controller/pricer.go @@ -5,44 +5,12 @@ package controller import ( - "context" "fmt" "time" "github.com/gitpod-io/gitpod/usage/pkg/db" - "github.com/gitpod-io/gitpod/usage/pkg/stripe" ) -type BillingController interface { - Reconcile(ctx context.Context, report UsageReport) error -} - -type NoOpBillingController struct{} - -func (b *NoOpBillingController) Reconcile(_ context.Context, _ UsageReport) error { - return nil -} - -type StripeBillingController struct { - sc *stripe.Client -} - -func NewStripeBillingController(sc *stripe.Client) *StripeBillingController { - return &StripeBillingController{ - sc: sc, - } -} - -func (b *StripeBillingController) Reconcile(ctx context.Context, report UsageReport) error { - runtimeReport := report.CreditSummaryForTeams() - - err := b.sc.UpdateUsage(ctx, runtimeReport) - if err != nil { - return fmt.Errorf("failed to update usage: %w", err) - } - return nil -} - const ( defaultWorkspaceClass = "default" ) diff --git a/components/usage/pkg/controller/billing_test.go b/components/usage/pkg/controller/pricer_test.go similarity index 100% rename from components/usage/pkg/controller/billing_test.go rename to components/usage/pkg/controller/pricer_test.go diff --git a/components/usage/pkg/controller/reconciler.go b/components/usage/pkg/controller/reconciler.go index 6d52f6e4ce44d1..799bbe5312fa95 100644 --- a/components/usage/pkg/controller/reconciler.go +++ b/components/usage/pkg/controller/reconciler.go @@ -8,7 +8,8 @@ import ( "context" "database/sql" "fmt" - "math" + v1 "github.com/gitpod-io/gitpod/usage-api/v1" + "google.golang.org/protobuf/types/known/timestamppb" "time" "github.com/gitpod-io/gitpod/common-go/log" @@ -29,20 +30,20 @@ func (f ReconcilerFunc) Reconcile() error { } type UsageReconciler struct { - nowFunc func() time.Time - conn *gorm.DB - pricer *WorkspacePricer - billingController BillingController - contentService contentservice.Interface + nowFunc func() time.Time + conn *gorm.DB + pricer *WorkspacePricer + billingService v1.BillingServiceClient + contentService contentservice.Interface } -func NewUsageReconciler(conn *gorm.DB, pricer *WorkspacePricer, billingController BillingController, contentService contentservice.Interface) *UsageReconciler { +func NewUsageReconciler(conn *gorm.DB, pricer *WorkspacePricer, billingClient v1.BillingServiceClient, contentService contentservice.Interface) *UsageReconciler { return &UsageReconciler{ - conn: conn, - pricer: pricer, - billingController: billingController, - contentService: contentService, - nowFunc: time.Now, + conn: conn, + pricer: pricer, + billingService: billingClient, + contentService: contentService, + nowFunc: time.Now, } } @@ -106,11 +107,14 @@ func (u *UsageReconciler) ReconcileTimeRange(ctx context.Context, from, to time. log.WithField("workspace_instances", instances).Debug("Successfully loaded workspace instances.") usageRecords := instancesToUsageRecords(instances, u.pricer, now) - //instancesByAttributionID := groupInstancesByAttributionID(instances) - err = u.billingController.Reconcile(ctx, usageRecords) + _, err = u.billingService.UpdateInvoices(ctx, &v1.UpdateInvoicesRequest{ + StartTime: timestamppb.New(from), + EndTime: timestamppb.New(to), + Sessions: instancesToBilledSessions(usageRecords), + }) if err != nil { - return nil, nil, fmt.Errorf("failed to reconcile billing: %w", err) + return nil, nil, fmt.Errorf("failed to update invoices: %w", err) } return status, usageRecords, nil @@ -148,31 +152,36 @@ func instancesToUsageRecords(instances []db.WorkspaceInstanceForUsage, pricer *W return usageRecords } -type UsageReport []db.WorkspaceInstanceUsage - -func (u UsageReport) CreditSummaryForTeams() map[string]int64 { - creditsPerTeamID := map[string]int64{} +func instancesToBilledSessions(instances []db.WorkspaceInstanceUsage) []*v1.BilledSession { + var sessions []*v1.BilledSession - for _, instance := range u { - entity, id := instance.AttributionID.Values() - if entity != db.AttributionEntity_Team { - continue - } + for _, instance := range instances { + var endTime *timestamppb.Timestamp - if _, ok := creditsPerTeamID[id]; !ok { - creditsPerTeamID[id] = 0 + if instance.StoppedAt.Valid { + endTime = timestamppb.New(instance.StoppedAt.Time) } - creditsPerTeamID[id] += int64(instance.CreditsUsed) - } - - for teamID, credits := range creditsPerTeamID { - creditsPerTeamID[teamID] = int64(math.Ceil(float64(credits))) + sessions = append(sessions, &v1.BilledSession{ + AttributionId: string(instance.AttributionID), + UserId: instance.UserID.String(), + TeamId: "", + WorkspaceId: instance.WorkspaceID, + WorkspaceType: string(instance.WorkspaceType), + ProjectId: instance.ProjectID, + InstanceId: instance.InstanceID.String(), + WorkspaceClass: instance.WorkspaceClass, + StartTime: timestamppb.New(instance.StartedAt), + EndTime: endTime, + Credits: instance.CreditsUsed, + }) } - return creditsPerTeamID + return sessions } +type UsageReport []db.WorkspaceInstanceUsage + type invalidWorkspaceInstance struct { reason string workspaceInstanceID uuid.UUID diff --git a/components/usage/pkg/controller/reconciler_test.go b/components/usage/pkg/controller/reconciler_test.go index d713e5db4bef09..2c85694964d765 100644 --- a/components/usage/pkg/controller/reconciler_test.go +++ b/components/usage/pkg/controller/reconciler_test.go @@ -7,6 +7,8 @@ package controller import ( "context" "database/sql" + v1 "github.com/gitpod-io/gitpod/usage-api/v1" + "google.golang.org/grpc" "testing" "time" @@ -52,10 +54,10 @@ func TestUsageReconciler_ReconcileTimeRange(t *testing.T) { dbtest.CreateWorkspaceInstances(t, conn, instances...) reconciler := &UsageReconciler{ - billingController: &NoOpBillingController{}, - nowFunc: func() time.Time { return scenarioRunTime }, - conn: conn, - pricer: DefaultWorkspacePricer, + billingService: &NoOpBillingServiceClient{}, + nowFunc: func() time.Time { return scenarioRunTime }, + conn: conn, + pricer: DefaultWorkspacePricer, } status, report, err := reconciler.ReconcileTimeRange(context.Background(), startOfMay, startOfJune) require.NoError(t, err) @@ -69,88 +71,10 @@ func TestUsageReconciler_ReconcileTimeRange(t *testing.T) { }, status) } -func TestUsageReport_CreditSummaryForTeams(t *testing.T) { - teamID := uuid.New().String() - teamAttributionID := db.NewTeamAttributionID(teamID) - - scenarios := []struct { - Name string - Report UsageReport - Expected map[string]int64 - }{ - { - Name: "no instances in report, no summary", - Report: []db.WorkspaceInstanceUsage{}, - Expected: map[string]int64{}, - }, - { - Name: "skips user attributions", - Report: []db.WorkspaceInstanceUsage{ - { - AttributionID: db.NewUserAttributionID(uuid.New().String()), - }, - }, - Expected: map[string]int64{}, - }, - { - Name: "two workspace instances", - Report: []db.WorkspaceInstanceUsage{ - { - // has 1 day and 23 hours of usage - AttributionID: teamAttributionID, - WorkspaceClass: defaultWorkspaceClass, - StartedAt: time.Date(2022, 05, 30, 00, 00, 00, 00, time.UTC), - StoppedAt: sql.NullTime{ - Time: time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC), - Valid: true, - }, - CreditsUsed: (24 + 23) * 10, - }, - { - // has 1 hour of usage - AttributionID: teamAttributionID, - WorkspaceClass: defaultWorkspaceClass, - StartedAt: time.Date(2022, 05, 30, 00, 00, 00, 00, time.UTC), - StoppedAt: sql.NullTime{ - Time: time.Date(2022, 05, 30, 1, 0, 0, 0, time.UTC), - Valid: true, - }, - CreditsUsed: 10, - }, - }, - Expected: map[string]int64{ - // total of 2 days runtime, at 10 credits per hour, that's 480 credits - teamID: 480, - }, - }, - { - Name: "unknown workspace class uses default", - Report: []db.WorkspaceInstanceUsage{ - // has 1 hour of usage - { - WorkspaceClass: "yolo-workspace-class", - AttributionID: teamAttributionID, - StartedAt: time.Date(2022, 05, 30, 00, 00, 00, 00, time.UTC), - StoppedAt: sql.NullTime{ - Time: time.Date(2022, 05, 30, 1, 0, 0, 0, time.UTC), - Valid: true, - }, - CreditsUsed: 10, - }, - }, - Expected: map[string]int64{ - // total of 1 hour usage, at default cost of 10 credits per hour - teamID: 10, - }, - }, - } +type NoOpBillingServiceClient struct{} - for _, s := range scenarios { - t.Run(s.Name, func(t *testing.T) { - actual := s.Report.CreditSummaryForTeams() - require.Equal(t, s.Expected, actual) - }) - } +func (c *NoOpBillingServiceClient) UpdateInvoices(ctx context.Context, in *v1.UpdateInvoicesRequest, opts ...grpc.CallOption) (*v1.UpdateInvoicesResponse, error) { + return &v1.UpdateInvoicesResponse{}, nil } func TestInstanceToUsageRecords(t *testing.T) { diff --git a/components/usage/pkg/db/workspace_instance.go b/components/usage/pkg/db/workspace_instance.go index fb6a54232c0ec6..4346ce211ea872 100644 --- a/components/usage/pkg/db/workspace_instance.go +++ b/components/usage/pkg/db/workspace_instance.go @@ -120,6 +120,22 @@ func (a AttributionID) Values() (entity string, identifier string) { return tokens[0], tokens[1] } +func ParseAttributionID(s string) (AttributionID, error) { + tokens := strings.Split(s, ":") + if len(tokens) != 2 { + return "", fmt.Errorf("attribution ID (%s) does not have two parts", s) + } + + switch tokens[0] { + case AttributionEntity_Team: + return NewTeamAttributionID(tokens[1]), nil + case AttributionEntity_User: + return NewUserAttributionID(tokens[1]), nil + default: + return "", fmt.Errorf("unknown attribution ID type: %s", s) + } +} + const ( WorkspaceClass_Default = "default" ) diff --git a/components/usage/pkg/server/server.go b/components/usage/pkg/server/server.go index 641308a09d68e2..414e2024827778 100644 --- a/components/usage/pkg/server/server.go +++ b/components/usage/pkg/server/server.go @@ -6,6 +6,8 @@ package server import ( "fmt" + "google.golang.org/grpc" + "google.golang.org/grpc/credentials/insecure" "net" "os" "time" @@ -47,13 +49,26 @@ func Start(cfg Config) error { return fmt.Errorf("failed to establish database connection: %w", err) } + var serverOpts []baseserver.Option + if cfg.Server != nil { + serverOpts = append(serverOpts, baseserver.WithConfig(cfg.Server)) + } + srv, err := baseserver.New("usage", serverOpts...) + if err != nil { + return fmt.Errorf("failed to initialize usage server: %w", err) + } + + selfConnection, err := grpc.Dial(srv.GRPCAddress(), grpc.WithTransportCredentials(insecure.NewCredentials())) + if err != nil { + return fmt.Errorf("failed to create self-connection to grpc server: %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{} - + var stripeClient *stripe.Client if cfg.StripeCredentialsFile != "" { config, err := stripe.ReadConfigFromFile(cfg.StripeCredentialsFile) if err != nil { @@ -65,7 +80,7 @@ func Start(cfg Config) error { return fmt.Errorf("failed to initialize stripe client: %w", err) } - billingController = controller.NewStripeBillingController(c) + stripeClient = c } schedule, err := time.ParseDuration(cfg.ControllerSchedule) @@ -78,7 +93,7 @@ func Start(cfg Config) error { contentService = contentservice.New(cfg.ContentServiceAddress) } - ctrl, err := controller.New(schedule, controller.NewUsageReconciler(conn, pricer, billingController, contentService)) + ctrl, err := controller.New(schedule, controller.NewUsageReconciler(conn, pricer, v1.NewBillingServiceClient(selfConnection), contentService)) if err != nil { return fmt.Errorf("failed to initialize usage controller: %w", err) } @@ -89,15 +104,7 @@ func Start(cfg Config) error { } defer ctrl.Stop() - var serverOpts []baseserver.Option - if cfg.Server != nil { - serverOpts = append(serverOpts, baseserver.WithConfig(cfg.Server)) - } - srv, err := baseserver.New("usage", serverOpts...) - if err != nil { - return fmt.Errorf("failed to initialize usage server: %w", err) - } - err = registerGRPCServices(srv, conn) + err = registerGRPCServices(srv, conn, stripeClient) if err != nil { return fmt.Errorf("failed to register gRPC services: %w", err) } @@ -115,7 +122,12 @@ func Start(cfg Config) error { return nil } -func registerGRPCServices(srv *baseserver.Server, conn *gorm.DB) error { +func registerGRPCServices(srv *baseserver.Server, conn *gorm.DB, stripeClient *stripe.Client) error { v1.RegisterUsageServiceServer(srv.GRPC(), apiv1.NewUsageService(conn)) + if stripeClient == nil { + v1.RegisterBillingServiceServer(srv.GRPC(), &apiv1.BillingServiceNoop{}) + } else { + v1.RegisterBillingServiceServer(srv.GRPC(), apiv1.NewBillingService(stripeClient)) + } return nil }