diff --git a/components/usage/pkg/controller/billing.go b/components/usage/pkg/controller/billing.go index e8ea9347ab1161..8fe1da66a1ce08 100644 --- a/components/usage/pkg/controller/billing.go +++ b/components/usage/pkg/controller/billing.go @@ -4,15 +4,19 @@ package controller -import "github.com/gitpod-io/gitpod/usage/pkg/stripe" +import ( + "context" + "github.com/gitpod-io/gitpod/usage/pkg/stripe" + "time" +) type BillingController interface { - Reconcile(report []TeamUsage) + Reconcile(ctx context.Context, now time.Time, report UsageReport) } type NoOpBillingController struct{} -func (b *NoOpBillingController) Reconcile(report []TeamUsage) {} +func (b *NoOpBillingController) Reconcile(_ context.Context, _ time.Time, _ UsageReport) {} type StripeBillingController struct { sc *stripe.Client @@ -22,12 +26,7 @@ func NewStripeBillingController(sc *stripe.Client) *StripeBillingController { return &StripeBillingController{sc: sc} } -func (b *StripeBillingController) Reconcile(report []TeamUsage) { - // Convert the usage report to sum all entries for the same team. - var summedReport = make(map[string]int64) - for _, usageEntry := range report { - summedReport[usageEntry.TeamID] += usageEntry.WorkspaceSeconds - } - - b.sc.UpdateUsage(summedReport) +func (b *StripeBillingController) Reconcile(ctx context.Context, now time.Time, report UsageReport) { + runtimeReport := report.RuntimeSummaryForTeams(now) + b.sc.UpdateUsage(runtimeReport) } diff --git a/components/usage/pkg/controller/reconciler.go b/components/usage/pkg/controller/reconciler.go index 672b1beeaf0e65..724311fde6d37d 100644 --- a/components/usage/pkg/controller/reconciler.go +++ b/components/usage/pkg/controller/reconciler.go @@ -45,12 +45,6 @@ type UsageReconcileStatus struct { WorkspaceInstances int InvalidWorkspaceInstances int - - Workspaces int - - Teams int - - Report []TeamUsage } func (u *UsageReconciler) Reconcile() error { @@ -60,7 +54,7 @@ func (u *UsageReconciler) Reconcile() error { startOfCurrentMonth := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC) startOfNextMonth := startOfCurrentMonth.AddDate(0, 1, 0) - status, err := u.ReconcileTimeRange(ctx, startOfCurrentMonth, startOfNextMonth) + status, report, err := u.ReconcileTimeRange(ctx, startOfCurrentMonth, startOfNextMonth) if err != nil { return err } @@ -75,7 +69,7 @@ func (u *UsageReconciler) Reconcile() error { defer f.Close() enc := json.NewEncoder(f) - err = enc.Encode(status.Report) + err = enc.Encode(report) if err != nil { return fmt.Errorf("failed to marshal report to JSON: %w", err) } @@ -89,7 +83,7 @@ func (u *UsageReconciler) Reconcile() error { return nil } -func (u *UsageReconciler) ReconcileTimeRange(ctx context.Context, from, to time.Time) (*UsageReconcileStatus, error) { +func (u *UsageReconciler) ReconcileTimeRange(ctx context.Context, from, to time.Time) (*UsageReconcileStatus, UsageReport, error) { now := u.nowFunc().UTC() log.Infof("Gathering usage data from %s to %s", from, to) status := &UsageReconcileStatus{ @@ -98,7 +92,7 @@ func (u *UsageReconciler) ReconcileTimeRange(ctx context.Context, from, to time. } instances, invalidInstances, err := u.loadWorkspaceInstances(ctx, from, to) if err != nil { - return nil, fmt.Errorf("failed to load workspace instances: %w", err) + return nil, nil, fmt.Errorf("failed to load workspace instances: %w", err) } status.WorkspaceInstances = len(instances) status.InvalidWorkspaceInstances = len(invalidInstances) @@ -108,129 +102,38 @@ func (u *UsageReconciler) ReconcileTimeRange(ctx context.Context, from, to time. } log.WithField("workspace_instances", instances).Debug("Successfully loaded workspace instances.") - workspaces, err := u.loadWorkspaces(ctx, instances) - if err != nil { - return nil, fmt.Errorf("failed to load workspaces for workspace instances in time range: %w", err) - } - status.Workspaces = len(workspaces) - - // match workspaces to teams - teams, err := u.loadTeamsForWorkspaces(ctx, workspaces) - if err != nil { - return nil, fmt.Errorf("failed to load teams for workspaces: %w", err) - } - status.Teams = len(teams) - - report, err := generateUsageReport(teams, now) - if err != nil { - return nil, fmt.Errorf("failed to generate usage report: %w", err) - } - status.Report = report - - u.billingController.Reconcile(status.Report) - - return status, nil -} - -func generateUsageReport(teams []teamWithWorkspaces, maxStopTime time.Time) ([]TeamUsage, error) { - var report []TeamUsage - for _, team := range teams { - var teamTotalRuntime int64 - for _, workspace := range team.Workspaces { - for _, instance := range workspace.Instances { - teamTotalRuntime += instance.WorkspaceRuntimeSeconds(maxStopTime) - } - } + instancesByAttributionID := groupInstancesByAttributionID(instances) - report = append(report, TeamUsage{ - TeamID: team.TeamID.String(), - WorkspaceSeconds: teamTotalRuntime, - }) - } - return report, nil -} + u.billingController.Reconcile(ctx, now, instancesByAttributionID) -type teamWithWorkspaces struct { - TeamID uuid.UUID - Workspaces []workspaceWithInstances + return status, instancesByAttributionID, nil } -func (u *UsageReconciler) loadTeamsForWorkspaces(ctx context.Context, workspaces []workspaceWithInstances) ([]teamWithWorkspaces, error) { - // find owner IDs of these workspaces - var ownerIDs []uuid.UUID - for _, workspace := range workspaces { - ownerIDs = append(ownerIDs, workspace.Workspace.OwnerID) - } +type UsageReport map[db.AttributionID][]db.WorkspaceInstance - // Retrieve memberships. This gives a link between an Owner and a Team they belong to. - memberships, err := db.ListTeamMembershipsForUserIDs(ctx, u.conn, ownerIDs) - if err != nil { - return nil, fmt.Errorf("failed to list team memberships: %w", err) - } +func (u UsageReport) RuntimeSummaryForTeams(maxStopTime time.Time) map[string]int64 { + attributedUsage := map[string]int64{} - membershipsByUserID := map[uuid.UUID]db.TeamMembership{} - for _, membership := range memberships { - // User can belong to multiple teams. For now, we're choosing the membership at random. - membershipsByUserID[membership.UserID] = membership - } + for attribution, instances := range u { + entity, id := attribution.Values() + if entity != db.AttributionEntity_Team { + continue + } - // Convert workspaces into a lookup so that we can index into them by Owner ID, needed for joining Teams with Workspaces - workspacesByOwnerID := map[uuid.UUID][]workspaceWithInstances{} - for _, workspace := range workspaces { - workspacesByOwnerID[workspace.Workspace.OwnerID] = append(workspacesByOwnerID[workspace.Workspace.OwnerID], workspace) - } + var runtime uint64 + for _, instance := range instances { + runtime += instance.WorkspaceRuntimeSeconds(maxStopTime) + } - // Finally, join the datasets - // Because we iterate over memberships, and not workspaces, we're in effect ignoring Workspaces which are not in a team. - // This is intended as we focus on Team usage for now. - var teamsWithWorkspaces []teamWithWorkspaces - for userID, membership := range membershipsByUserID { - teamsWithWorkspaces = append(teamsWithWorkspaces, teamWithWorkspaces{ - TeamID: membership.TeamID, - Workspaces: workspacesByOwnerID[userID], - }) + attributedUsage[id] = int64(runtime) } - return teamsWithWorkspaces, nil -} - -type workspaceWithInstances struct { - Workspace db.Workspace - Instances []db.WorkspaceInstance + return attributedUsage } -func (u *UsageReconciler) loadWorkspaces(ctx context.Context, instances []db.WorkspaceInstance) ([]workspaceWithInstances, error) { - var workspaceIDs []string - for _, instance := range instances { - workspaceIDs = append(workspaceIDs, instance.WorkspaceID) - } - - workspaces, err := db.ListWorkspacesByID(ctx, u.conn, toSet(workspaceIDs)) - if err != nil { - return nil, fmt.Errorf("failed to find workspaces for provided workspace instances: %w", err) - } - - workspacesByID := map[string]db.Workspace{} - for _, workspace := range workspaces { - workspacesByID[workspace.ID] = workspace - } - - // We need to also add the instances to corresponding records, a single workspace can have multiple instances - instancesByWorkspaceID := map[string][]db.WorkspaceInstance{} - for _, instance := range instances { - instancesByWorkspaceID[instance.WorkspaceID] = append(instancesByWorkspaceID[instance.WorkspaceID], instance) - } - - // Flatten results into a list - var workspacesWithInstances []workspaceWithInstances - for workspaceID, workspace := range workspacesByID { - workspacesWithInstances = append(workspacesWithInstances, workspaceWithInstances{ - Workspace: workspace, - Instances: instancesByWorkspaceID[workspaceID], - }) - } - - return workspacesWithInstances, nil +type invalidWorkspaceInstance struct { + reason string + workspaceInstanceID uuid.UUID } func (u *UsageReconciler) loadWorkspaceInstances(ctx context.Context, from, to time.Time) ([]db.WorkspaceInstance, []invalidWorkspaceInstance, error) { @@ -246,11 +149,6 @@ func (u *UsageReconciler) loadWorkspaceInstances(ctx context.Context, from, to t return trimmed, invalid, nil } -type invalidWorkspaceInstance struct { - reason string - workspaceInstanceID uuid.UUID -} - func validateInstances(instances []db.WorkspaceInstance) (valid []db.WorkspaceInstance, invalid []invalidWorkspaceInstance) { for _, i := range instances { // i is a pointer to the current element, we need to assign it to ensure we're copying the value, not the current pointer. @@ -302,20 +200,15 @@ func trimStartStopTime(instances []db.WorkspaceInstance, maximumStart, minimumSt return updated } -func toSet(items []string) []string { - m := map[string]struct{}{} - for _, i := range items { - m[i] = struct{}{} - } +func groupInstancesByAttributionID(instances []db.WorkspaceInstance) map[db.AttributionID][]db.WorkspaceInstance { + result := map[db.AttributionID][]db.WorkspaceInstance{} + for _, instance := range instances { + if _, ok := result[instance.UsageAttributionID]; !ok { + result[instance.UsageAttributionID] = []db.WorkspaceInstance{} + } - var result []string - for s := range m { - result = append(result, s) + result[instance.UsageAttributionID] = append(result[instance.UsageAttributionID], instance) } - return result -} -type TeamUsage struct { - TeamID string `json:"team_id"` - WorkspaceSeconds int64 `json:"workspace_seconds"` + return result } diff --git a/components/usage/pkg/controller/reconciler_test.go b/components/usage/pkg/controller/reconciler_test.go index ad30cd7fb493be..dcff4bca4f39d2 100644 --- a/components/usage/pkg/controller/reconciler_test.go +++ b/components/usage/pkg/controller/reconciler_test.go @@ -15,153 +15,58 @@ import ( ) func TestUsageReconciler_ReconcileTimeRange(t *testing.T) { - type Scenario struct { - Name string - NowFunc func() time.Time - - Memberships []db.TeamMembership - Workspaces []db.Workspace - Instances []db.WorkspaceInstance - - Expected *UsageReconcileStatus - } - startOfMay := time.Date(2022, 05, 1, 0, 00, 00, 00, time.UTC) startOfJune := time.Date(2022, 06, 1, 0, 00, 00, 00, time.UTC) - scenarios := []Scenario{ - (func() Scenario { - teamID := uuid.New() - userID := uuid.New() - scenarioRunTime := time.Date(2022, 05, 31, 23, 00, 00, 00, time.UTC) - - membership := db.TeamMembership{ - ID: uuid.New(), - TeamID: teamID, - UserID: userID, - Role: db.TeamMembershipRole_Member, - } - workspace := dbtest.NewWorkspace(t, db.Workspace{ - ID: "gitpodio-gitpod-gyjr82jkfna", - OwnerID: userID, - }) - instances := []db.WorkspaceInstance{ - // Ran throughout the reconcile period - dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ - ID: uuid.New(), - WorkspaceID: workspace.ID, - CreationTime: db.NewVarcharTime(time.Date(2022, 05, 1, 00, 00, 00, 00, time.UTC)), - StartedTime: db.NewVarcharTime(time.Date(2022, 05, 1, 00, 00, 00, 00, time.UTC)), - StoppedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)), - }), - // Still running - dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ - ID: uuid.New(), - WorkspaceID: workspace.ID, - CreationTime: db.NewVarcharTime(time.Date(2022, 05, 30, 00, 00, 00, 00, time.UTC)), - StartedTime: db.NewVarcharTime(time.Date(2022, 05, 30, 00, 00, 00, 00, time.UTC)), - }), - // No creation time, invalid record - dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ - ID: uuid.New(), - WorkspaceID: workspace.ID, - StartedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)), - StoppedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)), - }), - } + teamID := uuid.New() + scenarioRunTime := time.Date(2022, 05, 31, 23, 00, 00, 00, time.UTC) - expectedRuntime := instances[0].WorkspaceRuntimeSeconds(scenarioRunTime) + instances[1].WorkspaceRuntimeSeconds(scenarioRunTime) - - return Scenario{ - Name: "one team with one workspace", - Memberships: []db.TeamMembership{membership}, - Workspaces: []db.Workspace{workspace}, - Instances: instances, - NowFunc: func() time.Time { return scenarioRunTime }, - Expected: &UsageReconcileStatus{ - StartTime: startOfMay, - EndTime: startOfJune, - WorkspaceInstances: 2, - InvalidWorkspaceInstances: 1, - Workspaces: 1, - Teams: 1, - Report: []TeamUsage{ - { - TeamID: teamID.String(), - WorkspaceSeconds: expectedRuntime, - }, - }, - }, - } - })(), - (func() Scenario { - runTime := time.Date(2022, 05, 31, 23, 59, 59, 999999, time.UTC) - teamID, userID := uuid.New(), uuid.New() - workspaceID := "gitpodio-gitpod-gyjr82jkfnd" - var instances []db.WorkspaceInstance - for i := 0; i < 100; i++ { - instances = append(instances, dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ - ID: uuid.New(), - WorkspaceID: workspaceID, - CreationTime: db.NewVarcharTime(time.Date(2022, 05, 01, 00, 00, 00, 00, time.UTC)), - StartedTime: db.NewVarcharTime(time.Date(2022, 05, 01, 00, 00, 00, 00, time.UTC)), - StoppedTime: db.NewVarcharTime(time.Date(2022, 05, 31, 23, 59, 59, 999999, time.UTC)), - })) - } - - return Scenario{ - Name: "many long running instances do not overflow number of seconds in usage", - NowFunc: func() time.Time { return runTime }, - Memberships: []db.TeamMembership{ - { - ID: uuid.New(), - TeamID: teamID, - UserID: userID, - Role: db.TeamMembershipRole_Member, - }, - }, - Workspaces: []db.Workspace{ - dbtest.NewWorkspace(t, db.Workspace{ - ID: workspaceID, - OwnerID: userID, - }), - }, - Instances: instances, - Expected: &UsageReconcileStatus{ - StartTime: startOfMay, - EndTime: startOfJune, - WorkspaceInstances: 100, - InvalidWorkspaceInstances: 0, - Workspaces: 1, - Teams: 1, - Report: []TeamUsage{ - { - TeamID: teamID.String(), - WorkspaceSeconds: instances[0].WorkspaceRuntimeSeconds(runTime) * 100, - }, - }, - }, - } - })(), + instances := []db.WorkspaceInstance{ + // Ran throughout the reconcile period + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + ID: uuid.New(), + UsageAttributionID: db.NewTeamAttributionID(teamID.String()), + CreationTime: db.NewVarcharTime(time.Date(2022, 05, 1, 00, 00, 00, 00, time.UTC)), + StartedTime: db.NewVarcharTime(time.Date(2022, 05, 1, 00, 00, 00, 00, time.UTC)), + StoppedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)), + }), + // Still running + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + ID: uuid.New(), + UsageAttributionID: db.NewTeamAttributionID(teamID.String()), + CreationTime: db.NewVarcharTime(time.Date(2022, 05, 30, 00, 00, 00, 00, time.UTC)), + StartedTime: db.NewVarcharTime(time.Date(2022, 05, 30, 00, 00, 00, 00, time.UTC)), + }), + // No creation time, invalid record + dbtest.NewWorkspaceInstance(t, db.WorkspaceInstance{ + ID: uuid.New(), + UsageAttributionID: db.NewTeamAttributionID(teamID.String()), + StartedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)), + StoppedTime: db.NewVarcharTime(time.Date(2022, 06, 1, 1, 0, 0, 0, time.UTC)), + }), } - for _, scenario := range scenarios { - t.Run(scenario.Name, func(t *testing.T) { - conn := dbtest.ConnectForTests(t) - require.NoError(t, conn.Create(scenario.Memberships).Error) - dbtest.CreateWorkspaces(t, conn, scenario.Workspaces...) - dbtest.CreateWorkspaceInstances(t, conn, scenario.Instances...) + expectedRuntime := instances[0].WorkspaceRuntimeSeconds(scenarioRunTime) + instances[1].WorkspaceRuntimeSeconds(scenarioRunTime) - reconciler := &UsageReconciler{ - billingController: &NoOpBillingController{}, - nowFunc: scenario.NowFunc, - conn: conn, - } - status, err := reconciler.ReconcileTimeRange(context.Background(), startOfMay, startOfJune) - require.NoError(t, err) - - require.Equal(t, scenario.Expected, status) - }) + conn := dbtest.ConnectForTests(t) + dbtest.CreateWorkspaceInstances(t, conn, instances...) + reconciler := &UsageReconciler{ + billingController: &NoOpBillingController{}, + nowFunc: func() time.Time { return scenarioRunTime }, + conn: conn, } + status, report, err := reconciler.ReconcileTimeRange(context.Background(), startOfMay, startOfJune) + require.NoError(t, err) + + require.Len(t, report, 1) + require.Equal(t, &UsageReconcileStatus{ + StartTime: startOfMay, + EndTime: startOfJune, + 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 396ac48a759218..f730bca2217cfe 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) int64 { +func (i *WorkspaceInstance) WorkspaceRuntimeSeconds(maxStopTime time.Time) uint64 { start := i.CreationTime.Time() stop := maxStopTime @@ -56,7 +56,7 @@ func (i *WorkspaceInstance) WorkspaceRuntimeSeconds(maxStopTime time.Time) int64 } } - return int64(stop.Sub(start).Round(time.Second).Seconds()) + return uint64(stop.Sub(start).Round(time.Second).Seconds()) } // TableName sets the insert table name for this struct type @@ -79,6 +79,7 @@ func ListWorkspaceInstancesInRange(ctx context.Context, conn *gorm.DB, from, to ). Where("creationTime < ?", TimeToISO8601(to)). Where("startedTime != ?", ""). + Where("usageAttributionId != ?", ""). FindInBatches(&instancesInBatch, 1000, func(_ *gorm.DB, _ int) error { instances = append(instances, instancesInBatch...) return nil