diff --git a/components/usage/cmd/run.go b/components/usage/cmd/run.go index e23c044fe7fc78..1e1d719ed4fd8e 100644 --- a/components/usage/cmd/run.go +++ b/components/usage/cmd/run.go @@ -8,11 +8,13 @@ import ( "bytes" "encoding/json" "fmt" + "os" + "path" + "time" + "github.com/gitpod-io/gitpod/common-go/log" "github.com/gitpod-io/gitpod/usage/pkg/server" "github.com/spf13/cobra" - "os" - "path" ) func init() { @@ -66,5 +68,9 @@ func parseConfig(path string) (server.Config, error) { return server.Config{}, fmt.Errorf("failed to parse config from %s: %w", path, err) } + if cfg.BillInstancesAfter == nil { + cfg.BillInstancesAfter = &time.Time{} + } + return cfg, nil } diff --git a/components/usage/pkg/apiv1/billing.go b/components/usage/pkg/apiv1/billing.go index 0cdb031cf43500..c1bbadd23fb0a2 100644 --- a/components/usage/pkg/apiv1/billing.go +++ b/components/usage/pkg/apiv1/billing.go @@ -7,29 +7,33 @@ package apiv1 import ( "context" "fmt" + "math" + "time" + "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 { +func NewBillingService(stripeClient *stripe.Client, billInstancesAfter time.Time) *BillingService { return &BillingService{ - stripeClient: stripeClient, + stripeClient: stripeClient, + billInstancesAfter: billInstancesAfter, } } type BillingService struct { - stripeClient *stripe.Client + stripeClient *stripe.Client + billInstancesAfter time.Time v1.UnimplementedBillingServiceServer } func (s *BillingService) UpdateInvoices(ctx context.Context, in *v1.UpdateInvoicesRequest) (*v1.UpdateInvoicesResponse, error) { - credits, err := creditSummaryForTeams(in.GetSessions()) + credits, err := s.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") @@ -44,10 +48,14 @@ func (s *BillingService) UpdateInvoices(ctx context.Context, in *v1.UpdateInvoic return &v1.UpdateInvoicesResponse{}, nil } -func creditSummaryForTeams(sessions []*v1.BilledSession) (map[string]int64, error) { +func (s *BillingService) creditSummaryForTeams(sessions []*v1.BilledSession) (map[string]int64, error) { creditsPerTeamID := map[string]float64{} for _, session := range sessions { + if session.StartTime.AsTime().Before(s.billInstancesAfter) { + continue + } + attributionID, err := db.ParseAttributionID(session.AttributionId) if err != nil { return nil, fmt.Errorf("failed to parse attribution ID: %w", err) diff --git a/components/usage/pkg/apiv1/billing_test.go b/components/usage/pkg/apiv1/billing_test.go index 8ce7ef7a5e8c4e..b8b7d45190bd96 100644 --- a/components/usage/pkg/apiv1/billing_test.go +++ b/components/usage/pkg/apiv1/billing_test.go @@ -5,11 +5,15 @@ package apiv1 import ( + "testing" + "time" + 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" "github.com/google/uuid" "github.com/stretchr/testify/require" - "testing" + timestamppb "google.golang.org/protobuf/types/known/timestamppb" ) func TestCreditSummaryForTeams(t *testing.T) { @@ -17,17 +21,20 @@ func TestCreditSummaryForTeams(t *testing.T) { teamAttributionID_A, teamAttributionID_B := db.NewTeamAttributionID(teamID_A), db.NewTeamAttributionID(teamID_B) scenarios := []struct { - Name string - Sessions []*v1.BilledSession - Expected map[string]int64 + Name string + Sessions []*v1.BilledSession + BillSessionsAfter time.Time + Expected map[string]int64 }{ { - Name: "no instances in report, no summary", - Sessions: []*v1.BilledSession{}, - Expected: map[string]int64{}, + Name: "no instances in report, no summary", + BillSessionsAfter: time.Time{}, + Sessions: []*v1.BilledSession{}, + Expected: map[string]int64{}, }, { - Name: "skips user attributions", + Name: "skips user attributions", + BillSessionsAfter: time.Time{}, Sessions: []*v1.BilledSession{ { AttributionId: string(db.NewUserAttributionID(uuid.New().String())), @@ -36,7 +43,8 @@ func TestCreditSummaryForTeams(t *testing.T) { Expected: map[string]int64{}, }, { - Name: "two workspace instances", + Name: "two workspace instances", + BillSessionsAfter: time.Time{}, Sessions: []*v1.BilledSession{ { // has 1 day and 23 hours of usage @@ -55,7 +63,8 @@ func TestCreditSummaryForTeams(t *testing.T) { }, }, { - Name: "multiple teams", + Name: "multiple teams", + BillSessionsAfter: time.Time{}, Sessions: []*v1.BilledSession{ { // has 12 hours of usage @@ -74,11 +83,33 @@ func TestCreditSummaryForTeams(t *testing.T) { teamID_B: 240, }, }, + { + Name: "two instances, same team, one of which started too early to be considered", + BillSessionsAfter: time.Now().AddDate(0, 0, -2), + Sessions: []*v1.BilledSession{ + { + // has 12 hours of usage, started yesterday + AttributionId: string(teamAttributionID_A), + Credits: (12) * 10, + StartTime: timestamppb.New(time.Now().AddDate(0, 0, -1)), + }, + { + // has 1 day of usage, but started three days ago + AttributionId: string(teamAttributionID_A), + Credits: (24) * 10, + StartTime: timestamppb.New(time.Now().AddDate(0, 0, -3)), + }, + }, + Expected: map[string]int64{ + teamID_A: 120, + }, + }, } for _, s := range scenarios { t.Run(s.Name, func(t *testing.T) { - actual, err := creditSummaryForTeams(s.Sessions) + svc := NewBillingService(&stripe.Client{}, s.BillSessionsAfter) + actual, err := svc.creditSummaryForTeams(s.Sessions) require.NoError(t, err) require.Equal(t, s.Expected, actual) }) diff --git a/components/usage/pkg/server/server.go b/components/usage/pkg/server/server.go index 5835dd85e906ac..78c0a034d9aab8 100644 --- a/components/usage/pkg/server/server.go +++ b/components/usage/pkg/server/server.go @@ -34,6 +34,10 @@ type Config struct { ContentServiceAddress string `json:"contentServiceAddress,omitempty"` + // billInstancesAfter sets the date after which instances should be considered for billing - + // instances started before `billInstancesAfter` will not be considered by the billing controller. + BillInstancesAfter *time.Time `json:"billInstancesAfter,omitempty"` + Server *baseserver.Configuration `json:"server,omitempty"` } @@ -113,7 +117,7 @@ func Start(cfg Config) error { } defer ctrl.Stop() - err = registerGRPCServices(srv, conn, stripeClient, reportGenerator, contentService) + err = registerGRPCServices(srv, conn, stripeClient, reportGenerator, contentService, *cfg.BillInstancesAfter) if err != nil { return fmt.Errorf("failed to register gRPC services: %w", err) } @@ -131,12 +135,12 @@ func Start(cfg Config) error { return nil } -func registerGRPCServices(srv *baseserver.Server, conn *gorm.DB, stripeClient *stripe.Client, reportGenerator *apiv1.ReportGenerator, contentSvc contentservice.Interface) error { +func registerGRPCServices(srv *baseserver.Server, conn *gorm.DB, stripeClient *stripe.Client, reportGenerator *apiv1.ReportGenerator, contentSvc contentservice.Interface, billInstancesAfter time.Time) error { v1.RegisterUsageServiceServer(srv.GRPC(), apiv1.NewUsageService(conn, reportGenerator, contentSvc)) if stripeClient == nil { v1.RegisterBillingServiceServer(srv.GRPC(), &apiv1.BillingServiceNoop{}) } else { - v1.RegisterBillingServiceServer(srv.GRPC(), apiv1.NewBillingService(stripeClient)) + v1.RegisterBillingServiceServer(srv.GRPC(), apiv1.NewBillingService(stripeClient, billInstancesAfter)) } return nil }