diff --git a/components/gitpod-db/go/cost_center.go b/components/gitpod-db/go/cost_center.go index fb9654157f4aba..f37ed3f0967e78 100644 --- a/components/gitpod-db/go/cost_center.go +++ b/components/gitpod-db/go/cost_center.go @@ -104,7 +104,7 @@ func (c *CostCenterManager) GetOrCreateCostCenter(ctx context.Context, attributi // This can happen in the following scenario: // * User accesses gitpod just after their CostCenter expired, but just before our periodic CostCenter reset kicks in. if result.BillingStrategy == CostCenter_Other && result.IsExpired() { - cc, err := c.ResetUsage(ctx, result) + cc, err := c.ResetUsage(ctx, result.ID) if err != nil { logger.WithError(err).Error("Failed to reset expired usage.") return CostCenter{}, fmt.Errorf("failed to reset usage for expired cost center ID: %s: %w", result.ID, err) @@ -262,31 +262,37 @@ func (c *CostCenterManager) ListLatestCostCentersWithBillingTimeBefore(ctx conte return results, nil } -func (c *CostCenterManager) ResetUsage(ctx context.Context, cc CostCenter) (CostCenter, error) { +func (c *CostCenterManager) ResetUsage(ctx context.Context, id AttributionID) (CostCenter, error) { + logger := log.WithField("attribution_id", id) + now := time.Now().UTC() + cc, err := getCostCenter(ctx, c.conn, id) + if err != nil { + return cc, err + } + logger = logger.WithField("cost_center", cc) if cc.BillingStrategy != CostCenter_Other { return CostCenter{}, fmt.Errorf("cannot reset usage for Billing Strategy %s for Cost Center ID: %s", cc.BillingStrategy, cc.ID) } + if !cc.IsExpired() { + logger.Info("Skipping ResetUsage because next billing cycle is in the future.") + return cc, nil + } - now := time.Now().UTC() - + logger.Info("Running `ResetUsage`.") // Default to 1 month from now, if there's no nextBillingTime set on the record. + billingCycleStart := now nextBillingTime := now.AddDate(0, 1, 0) if cc.NextBillingTime.IsSet() { + billingCycleStart = cc.NextBillingTime.Time() nextBillingTime = cc.NextBillingTime.Time().AddDate(0, 1, 0) } - // Create a synthetic Invoice Usage record, to reset usage - err := c.BalanceOutUsage(ctx, cc.ID) - if err != nil { - return CostCenter{}, fmt.Errorf("failed to compute invocie usage record for AttributonID: %s: %w", cc.ID, err) - } - // All fields on the new cost center remain the same, except for BillingCycleStart, NextBillingTime, and CreationTime newCostCenter := CostCenter{ ID: cc.ID, SpendingLimit: cc.SpendingLimit, BillingStrategy: cc.BillingStrategy, - BillingCycleStart: NewVarCharTime(now), + BillingCycleStart: NewVarCharTime(billingCycleStart), NextBillingTime: NewVarCharTime(nextBillingTime), CreationTime: NewVarCharTime(now), } @@ -295,5 +301,11 @@ func (c *CostCenterManager) ResetUsage(ctx context.Context, cc CostCenter) (Cost return CostCenter{}, fmt.Errorf("failed to store new cost center for AttribtuonID: %s: %w", cc.ID, err) } + // Create a synthetic Invoice Usage record, to reset usage + err = c.BalanceOutUsage(ctx, cc.ID) + if err != nil { + return CostCenter{}, fmt.Errorf("failed to compute invocie usage record for AttributonID: %s: %w", cc.ID, err) + } + return newCostCenter, nil } diff --git a/components/gitpod-db/go/cost_center_test.go b/components/gitpod-db/go/cost_center_test.go index 043dcc3aa9629f..a874baef32c516 100644 --- a/components/gitpod-db/go/cost_center_test.go +++ b/components/gitpod-db/go/cost_center_test.go @@ -408,12 +408,14 @@ func TestCostCenterManager_ResetUsage(t *testing.T) { ForTeams: 0, ForUsers: 500, }) - _, err := mnr.ResetUsage(context.Background(), db.CostCenter{ + cc := dbtest.CreateCostCenters(t, conn, db.CostCenter{ ID: db.NewUserAttributionID(uuid.New().String()), CreationTime: db.NewVarCharTime(time.Now()), SpendingLimit: 500, BillingStrategy: db.CostCenter_Stripe, - }) + })[0] + + _, err := mnr.ResetUsage(context.Background(), cc.ID) require.Error(t, err) }) @@ -423,23 +425,20 @@ func TestCostCenterManager_ResetUsage(t *testing.T) { ForTeams: 0, ForUsers: 500, }) - oldCC := db.CostCenter{ - ID: db.NewTeamAttributionID(uuid.New().String()), - CreationTime: db.NewVarCharTime(time.Now()), - SpendingLimit: 0, - BillingStrategy: db.CostCenter_Other, - NextBillingTime: db.NewVarCharTime(ts), - } - newCC, err := mnr.ResetUsage(context.Background(), oldCC) + oldCC := dbtest.CreateCostCenters(t, conn, db.CostCenter{ + ID: db.NewTeamAttributionID(uuid.New().String()), + CreationTime: db.NewVarCharTime(time.Now()), + SpendingLimit: 10, + BillingStrategy: db.CostCenter_Other, + NextBillingTime: db.NewVarCharTime(ts), + BillingCycleStart: db.NewVarCharTime(ts.AddDate(0, -1, 0)), + })[0] + newCC, err := mnr.ResetUsage(context.Background(), oldCC.ID) require.NoError(t, err) - t.Cleanup(func() { - conn.Model(&db.CostCenter{}).Delete(newCC) - }) - require.Equal(t, oldCC.ID, newCC.ID) - require.EqualValues(t, 0, newCC.SpendingLimit) + require.EqualValues(t, 10, newCC.SpendingLimit) require.Equal(t, db.CostCenter_Other, newCC.BillingStrategy) - require.Equal(t, ts.AddDate(0, 1, 0), newCC.NextBillingTime.Time()) + require.Equal(t, db.NewVarCharTime(ts.AddDate(0, 1, 0)).Time(), newCC.NextBillingTime.Time()) }) diff --git a/components/usage/pkg/apiv1/usage.go b/components/usage/pkg/apiv1/usage.go index f4835c213eae33..57fff29793c348 100644 --- a/components/usage/pkg/apiv1/usage.go +++ b/components/usage/pkg/apiv1/usage.go @@ -250,7 +250,7 @@ func (s *UsageService) ResetUsage(ctx context.Context, req *v1.ResetUsageRequest var errors []error for _, cc := range costCentersToUpdate { - _, err = s.costCenterManager.ResetUsage(ctx, cc) + _, err = s.costCenterManager.ResetUsage(ctx, cc.ID) if err != nil { errors = append(errors, err) }