[usage] reduce possibility for races in ResetUsage

Author: svenefftinge

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,36 @@ 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("attributionId", id)
+	now := time.Now().UTC()
+	cc, err := getCostCenter(ctx, c.conn, id)
+	if err != nil {
+		return cc, err
+	}
 	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 +300,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
 }

components/gitpod-db/go/cost_center_test.go

@@ -408,12 +408,16 @@ func TestCostCenterManager_ResetUsage(t *testing.T) {
 			ForTeams: 0,
 			ForUsers: 500,
 		})
-		_, err := mnr.ResetUsage(context.Background(), db.CostCenter{
+		cc := &db.CostCenter{
 			ID:              db.NewUserAttributionID(uuid.New().String()),
 			CreationTime:    db.NewVarCharTime(time.Now()),
 			SpendingLimit:   500,
 			BillingStrategy: db.CostCenter_Stripe,
-		})
+		}
+
+		err := conn.Save(cc).Error
+		require.NoError(t, err)
+		_, err = mnr.ResetUsage(context.Background(), cc.ID)
 		require.Error(t, err)
 	})
 
@@ -423,23 +427,29 @@ 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),
+		oldCC := &db.CostCenter{
+			ID:                db.NewTeamAttributionID(uuid.New().String()),
+			CreationTime:      db.NewVarCharTime(time.Now()),
+			SpendingLimit:     0,
+			BillingStrategy:   db.CostCenter_Other,
+			NextBillingTime:   db.NewVarCharTime(ts),
+			BillingCycleStart: db.NewVarCharTime(ts.AddDate(0, -1, 0)),
 		}
-		newCC, err := mnr.ResetUsage(context.Background(), oldCC)
+		err := conn.Save(oldCC).Error
+		require.NoError(t, err)
+		t.Logf("time %s", ts)
+		newCC, err := mnr.ResetUsage(context.Background(), oldCC.ID)
+		t.Logf("time %s", time.Now())
 		require.NoError(t, err)
 		t.Cleanup(func() {
-			conn.Model(&db.CostCenter{}).Delete(newCC)
+			conn.Exec("DELETE from d_b_cost_center where id='%s'", oldCC.ID)
+			conn.Exec("DELETE from d_b_usage where attributionid='%s'", oldCC.ID)
 		})
 
 		require.Equal(t, oldCC.ID, newCC.ID)
 		require.EqualValues(t, 0, 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())
 
 	})
 

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