tests: Make billing lifecycle retries idempotent

admin/billing.go

@@ -285,7 +285,7 @@ func (s *Service) StartCreditTrial(ctx context.Context, org *database.Organizati
 	if balance <= 0 {
 		// Only seed the initial credits if the customer does not already have a trial balance.
 		// This keeps retries idempotent when an earlier attempt granted credits before failing.
-		if err := s.Biller.GrantCustomerCredits(ctx, org.BillingCustomerID, CreditTrialAllocation, billing.CreditsCurrency, "Initial trial credits", nil); err != nil {
+		if err := s.Biller.GrantCustomerCredits(ctx, org.BillingCustomerID, CreditTrialAllocation, billing.CreditsCurrency, "Initial trial credits", nil, ""); err != nil {
 			return nil, nil, fmt.Errorf("failed to grant trial credits: %w", err)
 		}
 	}

admin/billing/biller.go

@@ -58,11 +58,11 @@ type Biller interface {
 	// CreateCustomerCreditAlerts registers credit-balance alerts for the customer in the given currency.
 	CreateCustomerCreditAlerts(ctx context.Context, customerID, currency string, lowThreshold float64) error
 
-	// GrantCustomerCredits increments credit ledger entry to the customer's balance in the given currency. Description is recorded on the ledger entry; expiryDate may be nil for credits that never expire.
-	GrantCustomerCredits(ctx context.Context, customerID string, amount float64, currency, description string, expiryDate *time.Time) error
+	// GrantCustomerCredits increments credit ledger entry to the customer's balance in the given currency. Description is recorded on the ledger entry; expiryDate may be nil for credits that never expire. When idempotencyKey is non-empty, retries must apply the logical grant at most once.
+	GrantCustomerCredits(ctx context.Context, customerID string, amount float64, currency, description string, expiryDate *time.Time, idempotencyKey string) error
 
-	// DebitCustomerCredits posts a `decrement` ledger entry against the customer's balance in the given currency.
-	DebitCustomerCredits(ctx context.Context, customerID string, amount float64, currency, description string) error
+	// DebitCustomerCredits posts a `decrement` ledger entry against the customer's balance in the given currency. When idempotencyKey is non-empty, retries must apply the logical debit at most once.
+	DebitCustomerCredits(ctx context.Context, customerID string, amount float64, currency, description, idempotencyKey string) error
 
 	// GetCustomerCreditBalance returns the customer's current credit balance in the given currency.
 	GetCustomerCreditBalance(ctx context.Context, customerID, currency string) (float64, error)
@@ -74,8 +74,8 @@ type Biller interface {
 	// CancelSubscriptionsForCustomer cancels all the subscriptions for the given organization and returns the end date of the subscription
 	CancelSubscriptionsForCustomer(ctx context.Context, customerID string, cancelOption SubscriptionCancellationOption) (time.Time, error)
 
-	// ChangeSubscriptionPlan changes the plan of the given subscription immediately and returns the updated subscription
-	ChangeSubscriptionPlan(ctx context.Context, subscriptionID string, plan *Plan) (*Subscription, error)
+	// ChangeSubscriptionPlan changes the plan of the given subscription immediately and returns the updated subscription. When idempotencyKey is non-empty, retries must apply the logical change at most once.
+	ChangeSubscriptionPlan(ctx context.Context, subscriptionID string, plan *Plan, idempotencyKey string) (*Subscription, error)
 	// UnscheduleCancellation cancels the scheduled cancellation for the given subscription and returns the updated subscription
 	UnscheduleCancellation(ctx context.Context, subscriptionID string) (*Subscription, error)
 

admin/billing/noop.go

@@ -77,11 +77,11 @@ func (n noop) CreateCustomerCreditAlerts(ctx context.Context, customerID, curren
 	return nil
 }
 
-func (n noop) GrantCustomerCredits(ctx context.Context, customerID string, amount float64, currency, description string, expiryDate *time.Time) error {
+func (n noop) GrantCustomerCredits(ctx context.Context, customerID string, amount float64, currency, description string, expiryDate *time.Time, idempotencyKey string) error {
 	return nil
 }
 
-func (n noop) DebitCustomerCredits(ctx context.Context, customerID string, amount float64, currency, description string) error {
+func (n noop) DebitCustomerCredits(ctx context.Context, customerID string, amount float64, currency, description, idempotencyKey string) error {
 	return nil
 }
 
@@ -97,7 +97,7 @@ func (n noop) GetActiveSubscription(ctx context.Context, customerID string) (*Su
 	return &Subscription{Customer: &Customer{}, Plan: &Plan{Quotas: Quotas{}}}, nil
 }
 
-func (n noop) ChangeSubscriptionPlan(ctx context.Context, subscriptionID string, plan *Plan) (*Subscription, error) {
+func (n noop) ChangeSubscriptionPlan(ctx context.Context, subscriptionID string, plan *Plan, idempotencyKey string) (*Subscription, error) {
 	return &Subscription{Customer: &Customer{}, Plan: &Plan{Quotas: Quotas{}}}, nil
 }
 

admin/billing/orb.go

@@ -269,7 +269,7 @@ func (o *Orb) CreateCustomerCreditAlerts(ctx context.Context, customerID, curren
 }
 
 // GrantCustomerCredits adds an `increment` ledger entry to the customer's credit balance in Orb in the given currency. per_unit_cost_basis is set to "0.00" so the grant itself does not produce an invoice line item.
-func (o *Orb) GrantCustomerCredits(ctx context.Context, customerID string, amount float64, currency, description string, expiryDate *time.Time) error {
+func (o *Orb) GrantCustomerCredits(ctx context.Context, customerID string, amount float64, currency, description string, expiryDate *time.Time, idempotencyKey string) error {
 	params := orb.CustomerCreditLedgerNewEntryByExternalIDParamsAddIncrementCreditLedgerEntryRequestParams{
 		Amount:           orb.F(amount),
 		Currency:         orb.String(currency),
@@ -280,21 +280,21 @@ func (o *Orb) GrantCustomerCredits(ctx context.Context, customerID string, amoun
 	if expiryDate != nil {
 		params.ExpiryDate = orb.F(*expiryDate)
 	}
-	_, err := o.client.Customers.Credits.Ledger.NewEntryByExternalID(ctx, customerID, params)
+	_, err := o.client.Customers.Credits.Ledger.NewEntryByExternalID(ctx, customerID, params, orbIdempotencyOption(idempotencyKey)...)
 	if err != nil {
 		return fmt.Errorf("creating credit ledger increment entry: %w", err)
 	}
 	return nil
 }
 
 // DebitCustomerCredits posts a `decrement` ledger entry against the customer's credit balance in Orb in the given currency. Used to expire/zero out a credit pool, e.g., to roll trial credits over to a different currency on upgrade.
-func (o *Orb) DebitCustomerCredits(ctx context.Context, customerID string, amount float64, currency, description string) error {
+func (o *Orb) DebitCustomerCredits(ctx context.Context, customerID string, amount float64, currency, description, idempotencyKey string) error {
 	_, err := o.client.Customers.Credits.Ledger.NewEntryByExternalID(ctx, customerID, orb.CustomerCreditLedgerNewEntryByExternalIDParamsAddDecrementCreditLedgerEntryRequestParams{
 		Amount:      orb.F(amount),
 		Currency:    orb.String(currency),
 		EntryType:   orb.F(orb.CustomerCreditLedgerNewEntryByExternalIDParamsAddDecrementCreditLedgerEntryRequestParamsEntryTypeDecrement),
 		Description: orb.String(description),
-	})
+	}, orbIdempotencyOption(idempotencyKey)...)
 	if err != nil {
 		return fmt.Errorf("creating credit ledger decrement entry: %w", err)
 	}
@@ -337,11 +337,11 @@ func (o *Orb) GetActiveSubscription(ctx context.Context, customerID string) (*Su
 	return subs[0], nil
 }
 
-func (o *Orb) ChangeSubscriptionPlan(ctx context.Context, subscriptionID string, plan *Plan) (*Subscription, error) {
+func (o *Orb) ChangeSubscriptionPlan(ctx context.Context, subscriptionID string, plan *Plan, idempotencyKey string) (*Subscription, error) {
 	s, err := o.client.Subscriptions.SchedulePlanChange(ctx, subscriptionID, orb.SubscriptionSchedulePlanChangeParams{
 		PlanID:       orb.String(plan.ID),
 		ChangeOption: orb.F(orb.SubscriptionSchedulePlanChangeParamsChangeOptionImmediate),
-	})
+	}, orbIdempotencyOption(idempotencyKey)...)
 	if err != nil {
 		return nil, err
 	}
@@ -359,6 +359,13 @@ func (o *Orb) ChangeSubscriptionPlan(ctx context.Context, subscriptionID string,
 	}, nil
 }
 
+func orbIdempotencyOption(key string) []option.RequestOption {
+	if key == "" {
+		return nil
+	}
+	return []option.RequestOption{option.WithHeader("Idempotency-Key", key)}
+}
+
 func (o *Orb) UnscheduleCancellation(ctx context.Context, subscriptionID string) (*Subscription, error) {
 	sub, err := o.client.Subscriptions.UnscheduleCancellation(ctx, subscriptionID)
 	if err != nil {

admin/database/postgres/postgres.go

@@ -3160,7 +3160,9 @@ func (c *connection) UpsertBillingIssue(ctx context.Context, opts *database.Upse
 }
 
 func (c *connection) UpdateBillingIssueOverdueAsProcessed(ctx context.Context, id string) error {
-	res, err := c.getDB(ctx).ExecContext(ctx, `UPDATE billing_issues SET overdue_processed = true WHERE id = $1`, id)
+	// Treat this update as an atomic claim. A concurrent lifecycle worker that
+	// already observed the stale row must not also proceed to an external side effect.
+	res, err := c.getDB(ctx).ExecContext(ctx, `UPDATE billing_issues SET overdue_processed = true WHERE id = $1 AND overdue_processed = false`, id)
 	return checkUpdateRow("billing issue", res, err)
 }
 

admin/jobs/river/biller_event_handlers.go

@@ -42,24 +42,43 @@ func (w *PaymentFailedWorker) Work(ctx context.Context, job *river.Job[PaymentFa
 		}
 		return fmt.Errorf("failed to find organization of billing customer id %q: %w", job.Args.BillingCustomerID, err)
 	}
+	invoice, err := w.admin.Biller.GetInvoice(ctx, job.Args.InvoiceID)
+	if err != nil {
+		return fmt.Errorf("failed to verify invoice %q after payment failure: %w", job.Args.InvoiceID, err)
+	}
+	if invoice != nil && (!w.admin.Biller.IsInvoiceValid(ctx, invoice) || w.admin.Biller.IsInvoicePaid(ctx, invoice)) {
+		// A delayed failure event must not resurrect an issue after the invoice has
+		// already been paid or voided.
+		w.logger.Info("ignoring stale invoice payment failure", zap.String("org_id", org.ID), zap.String("invoice_id", job.Args.InvoiceID), zap.String("invoice_status", invoice.Status))
+		return nil
+	}
 
 	be, err := w.admin.DB.FindBillingIssueByTypeForOrg(ctx, org.ID, database.BillingIssueTypePaymentFailed)
 	if err != nil {
 		if !errors.Is(err, database.ErrNotFound) {
 			return fmt.Errorf("failed to find billing errors: %w", err)
 		}
 	}
-	var metadata *database.BillingIssueMetadataPaymentFailed
+	invoices := make(map[string]*database.BillingIssueMetadataPaymentFailedMeta)
+	eventTime := job.Args.FailedAt
 	if be != nil {
-		metadata = be.Metadata.(*database.BillingIssueMetadataPaymentFailed)
-	} else {
-		metadata = &database.BillingIssueMetadataPaymentFailed{
-			Invoices: make(map[string]*database.BillingIssueMetadataPaymentFailedMeta),
+		metadata := be.Metadata.(*database.BillingIssueMetadataPaymentFailed)
+		for id, failedInvoice := range metadata.Invoices {
+			invoiceCopy := *failedInvoice
+			invoices[id] = &invoiceCopy
+		}
+		if existing := invoices[job.Args.InvoiceID]; existing != nil && !job.Args.FailedAt.After(existing.FailedOn) {
+			// The invoice entry itself is the durable marker for the at-most-once
+			// failure-email attempt. Duplicates and older deliveries are no-ops.
+			return nil
+		}
+		if be.EventTime.After(eventTime) {
+			eventTime = be.EventTime
 		}
 	}
 
 	gracePeriodEndDate := job.Args.DueDate.AddDate(0, 0, database.BillingGracePeriodDays)
-	metadata.Invoices[job.Args.InvoiceID] = &database.BillingIssueMetadataPaymentFailedMeta{
+	invoices[job.Args.InvoiceID] = &database.BillingIssueMetadataPaymentFailedMeta{
 		ID:                 job.Args.InvoiceID,
 		Number:             job.Args.InvoiceNumber,
 		URL:                job.Args.InvoiceURL,
@@ -74,8 +93,8 @@ func (w *PaymentFailedWorker) Work(ctx context.Context, job *river.Job[PaymentFa
 	_, err = w.admin.DB.UpsertBillingIssue(ctx, &database.UpsertBillingIssueOptions{
 		OrgID:     org.ID,
 		Type:      database.BillingIssueTypePaymentFailed,
-		Metadata:  &database.BillingIssueMetadataPaymentFailed{Invoices: metadata.Invoices},
-		EventTime: job.Args.FailedAt,
+		Metadata:  &database.BillingIssueMetadataPaymentFailed{Invoices: invoices},
+		EventTime: eventTime,
 	})
 	if err != nil {
 		return fmt.Errorf("failed to add billing error: %w", err)
@@ -110,6 +129,7 @@ type PaymentSuccessWorker struct {
 	river.WorkerDefaults[PaymentSuccessArgs]
 	admin  *admin.Service
 	logger *zap.Logger
+	now    func() time.Time
 }
 
 func (w *PaymentSuccessWorker) Work(ctx context.Context, job *river.Job[PaymentSuccessArgs]) error {
@@ -138,12 +158,29 @@ func (w *PaymentSuccessWorker) Work(ctx context.Context, job *river.Job[PaymentS
 		// invoice not found in the failed invoices, do nothing
 		return nil
 	}
+	invoice, err := w.admin.Biller.GetInvoice(ctx, job.Args.InvoiceID)
+	if err != nil {
+		return fmt.Errorf("failed to verify invoice %q after payment success: %w", job.Args.InvoiceID, err)
+	}
+	if invoice != nil && w.admin.Biller.IsInvoiceValid(ctx, invoice) && !w.admin.Biller.IsInvoicePaid(ctx, invoice) {
+		// A delayed success event from an earlier attempt must not clear a newer
+		// failure while the provider still considers the invoice unpaid.
+		w.logger.Info("ignoring stale invoice payment success", zap.String("org_id", org.ID), zap.String("invoice_id", job.Args.InvoiceID), zap.String("invoice_status", invoice.Status))
+		return nil
+	}
 
-	delete(failedInvoices, job.Args.InvoiceID)
+	remainingInvoices := make(map[string]*database.BillingIssueMetadataPaymentFailedMeta, len(failedInvoices)-1)
+	for id, failedInvoice := range failedInvoices {
+		if id == job.Args.InvoiceID {
+			continue
+		}
+		invoiceCopy := *failedInvoice
+		remainingInvoices[id] = &invoiceCopy
+	}
 	w.logger.Info("invoice payment success for a failed invoice", zap.String("org_id", org.ID), zap.String("org_name", org.Name), zap.String("invoice_id", job.Args.InvoiceID))
 
 	// if no more failed invoices, delete the billing error
-	if len(failedInvoices) == 0 {
+	if len(remainingInvoices) == 0 {
 		err = w.admin.DB.DeleteBillingIssue(ctx, be.ID)
 		if err != nil {
 			return fmt.Errorf("failed to delete billing error: %w", err)
@@ -153,7 +190,7 @@ func (w *PaymentSuccessWorker) Work(ctx context.Context, job *river.Job[PaymentS
 		_, err = w.admin.DB.UpsertBillingIssue(ctx, &database.UpsertBillingIssueOptions{
 			OrgID:     org.ID,
 			Type:      database.BillingIssueTypePaymentFailed,
-			Metadata:  &database.BillingIssueMetadataPaymentFailed{Invoices: failedInvoices},
+			Metadata:  &database.BillingIssueMetadataPaymentFailed{Invoices: remainingInvoices},
 			EventTime: be.EventTime,
 		})
 		if err != nil {
@@ -166,7 +203,7 @@ func (w *PaymentSuccessWorker) Work(ctx context.Context, job *river.Job[PaymentS
 		ToEmail:        org.BillingEmail,
 		ToName:         org.Name,
 		OrgName:        org.Name,
-		PaymentDate:    time.Now(),
+		PaymentDate:    billingLifecycleNow(w.now),
 		BillingPageURL: w.admin.URLs.Billing(org.Name, false),
 	})
 	if err != nil {
@@ -187,6 +224,7 @@ type PaymentFailedGracePeriodCheckWorker struct {
 	river.WorkerDefaults[PaymentFailedGracePeriodCheckArgs]
 	admin  *admin.Service
 	logger *zap.Logger
+	now    func() time.Time
 }
 
 func (w *PaymentFailedGracePeriodCheckWorker) Work(ctx context.Context, job *river.Job[PaymentFailedGracePeriodCheckArgs]) error {
@@ -199,71 +237,54 @@ func (w *PaymentFailedGracePeriodCheckWorker) Work(ctx context.Context, job *riv
 		return fmt.Errorf("failed to find organization with billing issue: %w", err)
 	}
 
+	now := billingLifecycleNow(w.now)
+	var workErr error
 	// failures are per org
 	for _, f := range failures {
-		overdue, err := w.checkFailedInvoicesForOrg(ctx, f)
-		if err != nil {
-			w.logger.Error("failed to check failed invoices for org", zap.String("org_id", f.OrgID), zap.Error(err))
-			continue // continue to next org
+		if err := w.processIssue(ctx, f, now); err != nil {
+			w.logger.Error("failed to process overdue invoices for org", zap.String("org_id", f.OrgID), zap.Error(err))
+			workErr = errors.Join(workErr, fmt.Errorf("process overdue invoices for org %q: %w", f.OrgID, err))
 		}
+	}
+	return workErr
+}
 
-		if !overdue {
-			continue // continue to next org
-		}
-
-		// hibernate projects
-		limit := 10
-		afterProjectName := ""
-		for {
-			projs, err := w.admin.DB.FindProjectsForOrganization(ctx, f.OrgID, afterProjectName, limit)
-			if err != nil {
-				return err
-			}
-
-			for _, proj := range projs {
-				_, err = w.admin.HibernateProject(ctx, proj)
-				if err != nil {
-					return fmt.Errorf("failed to hibernate project %q: %w", proj.Name, err)
-				}
-				afterProjectName = proj.Name
-			}
-
-			if len(projs) < limit {
-				break
-			}
-		}
-		org, err := w.admin.DB.FindOrganization(ctx, f.OrgID)
-		if err != nil {
-			return fmt.Errorf("failed to find organization: %w", err)
-		}
-
-		w.logger.Warn("projects hibernated due to unpaid invoice", zap.String("org_id", org.ID), zap.String("org_name", org.Name))
-
-		// send email
-		err = w.admin.Email.SendInvoiceUnpaid(&email.InvoiceUnpaid{
-			ToEmail:    org.BillingEmail,
-			ToName:     org.Name,
-			OrgName:    org.Name,
-			PaymentURL: w.admin.URLs.PaymentPortal(org.Name),
-		})
-		if err != nil {
-			return fmt.Errorf("failed to send project hibernated due to payment overdue email for org %q: %w", org.Name, err)
-		}
+func (w *PaymentFailedGracePeriodCheckWorker) processIssue(ctx context.Context, issue *database.BillingIssue, now time.Time) error {
+	overdue, err := w.checkFailedInvoicesForOrg(ctx, issue, now)
+	if err != nil || !overdue {
+		return err
+	}
+	if _, err := hibernateProjectsForOrganization(ctx, w.admin, issue.OrgID); err != nil {
+		return err
+	}
+	org, err := w.admin.DB.FindOrganization(ctx, issue.OrgID)
+	if err != nil {
+		return fmt.Errorf("failed to find organization: %w", err)
+	}
+	w.logger.Warn("projects hibernated due to unpaid invoice", zap.String("org_id", org.ID), zap.String("org_name", org.Name))
 
-		// mark the billing issue as processed
-		err = w.admin.DB.UpdateBillingIssueOverdueAsProcessed(ctx, f.ID)
-		if err != nil {
-			return fmt.Errorf("failed to mark billing issue as processed: %w", err)
-		}
+	// Mark first so a processed-marker failure cannot follow a successful email
+	// and cause a duplicate on retry.
+	if err := w.admin.DB.UpdateBillingIssueOverdueAsProcessed(ctx, issue.ID); err != nil {
+		return fmt.Errorf("failed to mark billing issue as processed: %w", err)
+	}
+	err = w.admin.Email.SendInvoiceUnpaid(&email.InvoiceUnpaid{
+		ToEmail:    org.BillingEmail,
+		ToName:     org.Name,
+		OrgName:    org.Name,
+		PaymentURL: w.admin.URLs.PaymentPortal(org.Name),
+	})
+	if err != nil {
+		w.logger.Error("failed to send invoice unpaid email after marking it attempted", zap.String("org_id", org.ID), zap.String("org_name", org.Name), zap.Error(err))
 	}
 	return nil
 }
 
 // reconciles failed payments for the org and returns true if any is overdue
-func (w *PaymentFailedGracePeriodCheckWorker) checkFailedInvoicesForOrg(ctx context.Context, orgPaymentFailures *database.BillingIssue) (bool, error) {
+func (w *PaymentFailedGracePeriodCheckWorker) checkFailedInvoicesForOrg(ctx context.Context, orgPaymentFailures *database.BillingIssue, now time.Time) (bool, error) {
 	hasOverdue := false
 	for invoiceID, failedInvoice := range orgPaymentFailures.Metadata.(*database.BillingIssueMetadataPaymentFailed).Invoices {
-		if time.Now().UTC().Before(failedInvoice.GracePeriodEndDate.AddDate(0, 0, 1)) {
+		if now.Before(failedInvoice.GracePeriodEndDate.AddDate(0, 0, 1)) {
 			continue
 		}