[usage] Implement BillingService

Author: easyCZ

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
+}

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
+}

components/usage/pkg/apiv1/billing_test.go

@@ -0,0 +1,66 @@
+// 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 := uuid.New().String()
+	teamAttributionID := db.NewTeamAttributionID(teamID)
+
+	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),
+					Credits:       (24 + 23) * 10,
+				},
+				{
+					// has 1 hour of usage
+					AttributionId: string(teamAttributionID),
+					Credits:       10,
+				},
+			},
+			Expected: map[string]int64{
+				// total of 2 days runtime, at 10 credits per hour, that's 480 credits
+				teamID: 480,
+			},
+		},
+	}
+
+	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)
+		})
+	}
+}

components/usage/pkg/controller/billing.go renamed to 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"
 )

components/usage/pkg/controller/billing_test.go renamed to components/usage/pkg/controller/pricer_test.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:        int64(instance.CreditsUsed),
+		})
 	}
 
-	return creditsPerTeamID
+	return sessions
 }
 
+type UsageReport []db.WorkspaceInstanceUsage
+
 type invalidWorkspaceInstance struct {
 	reason              string
 	workspaceInstanceID uuid.UUID