diff --git a/admin/billing.go b/admin/billing.go index 69ee3168bf04..5092d0824193 100644 --- a/admin/billing.go +++ b/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) } } diff --git a/admin/billing/biller.go b/admin/billing/biller.go index 6eadf586cf56..a3d732800e9e 100644 --- a/admin/billing/biller.go +++ b/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) diff --git a/admin/billing/noop.go b/admin/billing/noop.go index 5373577a0c63..49f3c8431e20 100644 --- a/admin/billing/noop.go +++ b/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 } diff --git a/admin/billing/orb.go b/admin/billing/orb.go index fad0bc621b0c..ba4eac5586c8 100644 --- a/admin/billing/orb.go +++ b/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,7 +280,7 @@ 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) } @@ -288,13 +288,13 @@ func (o *Orb) GrantCustomerCredits(ctx context.Context, customerID string, amoun } // 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 { diff --git a/admin/database/postgres/postgres.go b/admin/database/postgres/postgres.go index 55a6ae96bfe4..e63d080cfb98 100644 --- a/admin/database/postgres/postgres.go +++ b/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) } diff --git a/admin/jobs/river/biller_event_handlers.go b/admin/jobs/river/biller_event_handlers.go index ec205099de7e..255c445115ba 100644 --- a/admin/jobs/river/biller_event_handlers.go +++ b/admin/jobs/river/biller_event_handlers.go @@ -42,6 +42,16 @@ 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 { @@ -49,17 +59,26 @@ func (w *PaymentFailedWorker) Work(ctx context.Context, job *river.Job[PaymentFa 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 } diff --git a/admin/jobs/river/billing_lifecycle_test.go b/admin/jobs/river/billing_lifecycle_test.go new file mode 100644 index 000000000000..ab9e498d5d3f --- /dev/null +++ b/admin/jobs/river/billing_lifecycle_test.go @@ -0,0 +1,901 @@ +package river + +import ( + "context" + "errors" + "fmt" + "sort" + "strings" + "sync" + "testing" + "time" + + "github.com/rilldata/rill/admin" + "github.com/rilldata/rill/admin/billing" + "github.com/rilldata/rill/admin/database" + _ "github.com/rilldata/rill/admin/database/postgres" + "github.com/rilldata/rill/admin/pkg/pgtestcontainer" + "github.com/rilldata/rill/runtime/pkg/email" + "github.com/riverqueue/river" + "github.com/stretchr/testify/require" + "go.uber.org/zap" +) + +func TestBillingLifecycleExactTimeBoundaries(t *testing.T) { + // Lifecycle checks use inclusive deadlines. One nanosecond before a deadline + // must be a no-op, while the exact instant must durably advance state. + boundary := time.Date(2026, time.July, 23, 12, 0, 0, 0, time.UTC) + + t.Run("trial ending soon", func(t *testing.T) { + // The seven-day warning becomes eligible exactly at end-date minus seven days. + adm, db, _, sender := newBillingLifecycleFixture(t) + org := db.addOrganization("ending") + issue := db.addIssue(org.ID, database.BillingIssueTypeOnTrial, false, &database.BillingIssueMetadataOnTrial{EndDate: boundary.AddDate(0, 0, 7)}, boundary) + now := boundary.Add(-time.Nanosecond) + worker := &TrialEndingSoonWorker{admin: adm, logger: zap.NewNop(), now: func() time.Time { return now }} + + require.NoError(t, worker.Work(t.Context(), nil)) + require.False(t, db.isProcessed(issue.ID)) + require.Equal(t, 0, sender.attemptCount()) + + now = boundary + require.NoError(t, worker.Work(t.Context(), nil)) + require.True(t, db.isProcessed(issue.ID)) + require.Equal(t, 1, sender.attemptCount()) + }) + + t.Run("trial ended", func(t *testing.T) { + // The trial-to-grace transaction runs at the exact subscription trial end. + adm, db, biller, sender := newBillingLifecycleFixture(t) + org := db.addOrganization("ended") + db.addIssue(org.ID, database.BillingIssueTypeOnTrial, true, &database.BillingIssueMetadataOnTrial{SubID: "trial-sub", PlanID: "trial-plan", EndDate: boundary, GracePeriodEndDate: boundary.AddDate(0, 0, 7)}, boundary) + biller.subscriptions[org.BillingCustomerID] = &billing.Subscription{ID: "trial-sub", Plan: &billing.Plan{ID: "trial-plan"}} + now := boundary.Add(-time.Nanosecond) + worker := &TrialEndCheckWorker{admin: adm, logger: zap.NewNop(), now: func() time.Time { return now }} + + require.NoError(t, worker.Work(t.Context(), nil)) + require.NotNil(t, db.issueForOrg(org.ID, database.BillingIssueTypeOnTrial)) + require.Nil(t, db.issueForOrg(org.ID, database.BillingIssueTypeTrialEnded)) + + now = boundary + require.NoError(t, worker.Work(t.Context(), nil)) + require.Nil(t, db.issueForOrg(org.ID, database.BillingIssueTypeOnTrial)) + require.NotNil(t, db.issueForOrg(org.ID, database.BillingIssueTypeTrialEnded)) + require.Equal(t, 1, sender.attemptCount()) + }) + + t.Run("trial grace period", func(t *testing.T) { + // Trial projects are hibernated at the exact grace-period end, not before it. + adm, db, biller, sender := newBillingLifecycleFixture(t) + org := db.addOrganization("grace") + issue := db.addIssue(org.ID, database.BillingIssueTypeTrialEnded, false, &database.BillingIssueMetadataTrialEnded{SubID: "trial-sub", PlanID: "trial-plan", GracePeriodEndDate: boundary}, boundary) + biller.subscriptions[org.BillingCustomerID] = &billing.Subscription{ID: "trial-sub", Plan: &billing.Plan{ID: "trial-plan"}} + now := boundary.Add(-time.Nanosecond) + worker := &TrialGracePeriodCheckWorker{admin: adm, logger: zap.NewNop(), now: func() time.Time { return now }} + + require.NoError(t, worker.Work(t.Context(), nil)) + require.False(t, db.isProcessed(issue.ID)) + require.Empty(t, biller.cancellations) + + now = boundary + require.NoError(t, worker.Work(t.Context(), nil)) + require.True(t, db.isProcessed(issue.ID)) + require.Equal(t, []string{org.BillingCustomerID}, biller.cancellations) + requireOrganizationQuotasZero(t, db.organization(org.ID)) + require.Equal(t, 1, sender.attemptCount()) + }) + + t.Run("subscription cancellation", func(t *testing.T) { + // Cancellation intentionally waits through the day after the recorded end date. + adm, db, _, sender := newBillingLifecycleFixture(t) + org := db.addOrganization("cancelled") + issue := db.addIssue(org.ID, database.BillingIssueTypeSubscriptionCancelled, false, &database.BillingIssueMetadataSubscriptionCancelled{EndDate: boundary.AddDate(0, 0, -1)}, boundary) + now := boundary.Add(-time.Nanosecond) + worker := &SubscriptionCancellationCheckWorker{admin: adm, logger: zap.NewNop(), now: func() time.Time { return now }} + + require.NoError(t, worker.Work(t.Context(), nil)) + require.False(t, db.isProcessed(issue.ID)) + + now = boundary + require.NoError(t, worker.Work(t.Context(), nil)) + require.True(t, db.isProcessed(issue.ID)) + requireOrganizationQuotasZero(t, db.organization(org.ID)) + require.Equal(t, 1, sender.attemptCount()) + }) + + t.Run("failed invoice grace period", func(t *testing.T) { + // An unpaid invoice becomes hibernation-eligible exactly one day after grace. + adm, db, biller, sender := newBillingLifecycleFixture(t) + org := db.addOrganization("overdue") + issue := db.addIssue(org.ID, database.BillingIssueTypePaymentFailed, false, &database.BillingIssueMetadataPaymentFailed{Invoices: map[string]*database.BillingIssueMetadataPaymentFailedMeta{ + "invoice": {ID: "invoice", GracePeriodEndDate: boundary.AddDate(0, 0, -1)}, + }}, boundary) + biller.invoices["invoice"] = &billing.Invoice{ID: "invoice", Status: "open"} + now := boundary.Add(-time.Nanosecond) + worker := &PaymentFailedGracePeriodCheckWorker{admin: adm, logger: zap.NewNop(), now: func() time.Time { return now }} + + require.NoError(t, worker.Work(t.Context(), nil)) + require.False(t, db.isProcessed(issue.ID)) + + now = boundary + require.NoError(t, worker.Work(t.Context(), nil)) + require.True(t, db.isProcessed(issue.ID)) + require.Equal(t, 1, sender.attemptCount()) + }) +} + +func TestBillingLifecycleNotificationAttemptIsDurablyAtMostOnce(t *testing.T) { + // The existing schema can provide at-most-once SMTP attempts by committing a + // terminal marker first; these cases pin both the guarantee and its loss tradeoff. + now := time.Date(2026, time.July, 23, 12, 0, 0, 0, time.UTC) + + t.Run("processed marker failure happens before email", func(t *testing.T) { + // A marker outage leaves the issue retryable and sends nothing until marking succeeds. + adm, db, _, sender := newBillingLifecycleFixture(t) + org := db.addOrganization("marker") + issue := db.addIssue(org.ID, database.BillingIssueTypeSubscriptionCancelled, false, &database.BillingIssueMetadataSubscriptionCancelled{EndDate: now.AddDate(0, 0, -1)}, now) + db.processedFailures[issue.ID] = 1 + worker := &SubscriptionCancellationCheckWorker{admin: adm, logger: zap.NewNop(), now: func() time.Time { return now }} + + require.ErrorContains(t, worker.Work(t.Context(), nil), "processed") + require.False(t, db.isProcessed(issue.ID)) + require.Equal(t, 0, sender.attemptCount()) + + require.NoError(t, worker.Work(t.Context(), nil)) + require.NoError(t, worker.Work(t.Context(), nil)) + require.True(t, db.isProcessed(issue.ID)) + require.Equal(t, 1, sender.attemptCount()) + }) + + t.Run("transaction commit failure happens before email", func(t *testing.T) { + // A failed trial transition commit retains the old issue and cannot leak an email. + adm, db, biller, sender := newBillingLifecycleFixture(t) + org := db.addOrganization("commit") + db.addIssue(org.ID, database.BillingIssueTypeOnTrial, true, &database.BillingIssueMetadataOnTrial{SubID: "trial-sub", PlanID: "trial-plan", EndDate: now, GracePeriodEndDate: now.AddDate(0, 0, 7)}, now) + biller.subscriptions[org.BillingCustomerID] = &billing.Subscription{ID: "trial-sub", Plan: &billing.Plan{ID: "trial-plan"}} + db.commitFailures = 1 + worker := &TrialEndCheckWorker{admin: adm, logger: zap.NewNop(), now: func() time.Time { return now }} + + require.ErrorContains(t, worker.Work(t.Context(), nil), "commit") + require.NotNil(t, db.issueForOrg(org.ID, database.BillingIssueTypeOnTrial)) + require.Nil(t, db.issueForOrg(org.ID, database.BillingIssueTypeTrialEnded)) + require.Equal(t, 0, sender.attemptCount()) + + require.NoError(t, worker.Work(t.Context(), nil)) + require.NoError(t, worker.Work(t.Context(), nil)) + require.Nil(t, db.issueForOrg(org.ID, database.BillingIssueTypeOnTrial)) + require.NotNil(t, db.issueForOrg(org.ID, database.BillingIssueTypeTrialEnded)) + require.Equal(t, 1, sender.attemptCount()) + }) + + t.Run("email failure is not retried after durable completion", func(t *testing.T) { + // SMTP failure after the marker demonstrates the explicit at-most-once/loss tradeoff. + adm, db, biller, sender := newBillingLifecycleFixture(t) + org := db.addOrganization("smtp") + issue := db.addIssue(org.ID, database.BillingIssueTypeTrialEnded, false, &database.BillingIssueMetadataTrialEnded{SubID: "trial-sub", PlanID: "trial-plan", GracePeriodEndDate: now}, now) + biller.subscriptions[org.BillingCustomerID] = &billing.Subscription{ID: "trial-sub", Plan: &billing.Plan{ID: "trial-plan"}} + sender.failures = 1 + worker := &TrialGracePeriodCheckWorker{admin: adm, logger: zap.NewNop(), now: func() time.Time { return now }} + + require.NoError(t, worker.Work(t.Context(), nil)) + require.True(t, db.isProcessed(issue.ID)) + require.Equal(t, 1, sender.attemptCount()) + require.Equal(t, 0, sender.successCount()) + + require.NoError(t, worker.Work(t.Context(), nil)) + require.Equal(t, 1, sender.attemptCount()) + }) + + t.Run("failed-payment duplicate does not repeat a failed email", func(t *testing.T) { + // The persisted invoice entry suppresses retry sends even when SMTP returned an error. + adm, db, biller, sender := newBillingLifecycleFixture(t) + org := db.addOrganization("failed-email") + biller.invoices["invoice"] = &billing.Invoice{ID: "invoice", Status: "open"} + sender.failures = 1 + worker := &PaymentFailedWorker{admin: adm, logger: zap.NewNop()} + job := paymentFailedJob(org.BillingCustomerID, "invoice", now) + + require.ErrorContains(t, worker.Work(t.Context(), job), "email") + require.NotNil(t, db.issueForOrg(org.ID, database.BillingIssueTypePaymentFailed)) + require.Equal(t, 1, sender.attemptCount()) + + require.NoError(t, worker.Work(t.Context(), job)) + require.Equal(t, 1, sender.attemptCount()) + }) + + t.Run("successful-payment duplicate does not repeat a failed email", func(t *testing.T) { + // Deleting the invoice issue before SMTP makes the success notification one attempt. + adm, db, biller, sender := newBillingLifecycleFixture(t) + org := db.addOrganization("success-email") + db.addIssue(org.ID, database.BillingIssueTypePaymentFailed, false, paymentFailureMetadata("invoice", now), now) + biller.invoices["invoice"] = &billing.Invoice{ID: "invoice", Status: "paid"} + sender.failures = 1 + worker := &PaymentSuccessWorker{admin: adm, logger: zap.NewNop(), now: func() time.Time { return now }} + job := &river.Job[PaymentSuccessArgs]{Args: PaymentSuccessArgs{BillingCustomerID: org.BillingCustomerID, InvoiceID: "invoice"}} + + require.NoError(t, worker.Work(t.Context(), job)) + require.Nil(t, db.issueForOrg(org.ID, database.BillingIssueTypePaymentFailed)) + require.Equal(t, 1, sender.attemptCount()) + require.Equal(t, 0, sender.successCount()) + + require.NoError(t, worker.Work(t.Context(), job)) + require.Equal(t, 1, sender.attemptCount()) + }) +} + +func TestInvoiceLifecycleDuplicateAndOutOfOrderDeliveryIsMonotonic(t *testing.T) { + // Provider truth and per-invoice timestamps prevent duplicate or stale events + // from regressing the durable failed-invoice set. + adm, db, biller, sender := newBillingLifecycleFixture(t) + org := db.addOrganization("invoice-order") + failureWorker := &PaymentFailedWorker{admin: adm, logger: zap.NewNop()} + successWorker := &PaymentSuccessWorker{admin: adm, logger: zap.NewNop(), now: func() time.Time { + return time.Date(2026, time.July, 24, 0, 0, 0, 0, time.UTC) + }} + newer := time.Date(2026, time.July, 23, 10, 0, 0, 0, time.UTC) + older := newer.Add(-time.Hour) + biller.invoices["invoice"] = &billing.Invoice{ID: "invoice", Status: "open"} + + require.NoError(t, failureWorker.Work(t.Context(), paymentFailedJob(org.BillingCustomerID, "invoice", newer))) + require.NoError(t, failureWorker.Work(t.Context(), paymentFailedJob(org.BillingCustomerID, "invoice", newer))) + require.NoError(t, failureWorker.Work(t.Context(), paymentFailedJob(org.BillingCustomerID, "invoice", older))) + issue := db.issueForOrg(org.ID, database.BillingIssueTypePaymentFailed) + require.NotNil(t, issue) + require.Equal(t, newer, issue.EventTime) + require.Equal(t, newer, issue.Metadata.(*database.BillingIssueMetadataPaymentFailed).Invoices["invoice"].FailedOn) + require.Equal(t, 1, sender.attemptCount()) + + // A success event is stale while the provider still reports the invoice unpaid. + successJob := &river.Job[PaymentSuccessArgs]{Args: PaymentSuccessArgs{BillingCustomerID: org.BillingCustomerID, InvoiceID: "invoice"}} + require.NoError(t, successWorker.Work(t.Context(), successJob)) + require.NotNil(t, db.issueForOrg(org.ID, database.BillingIssueTypePaymentFailed)) + require.Equal(t, 1, sender.attemptCount()) + + // Once paid, success clears exactly once; a delayed failure cannot resurrect it. + biller.invoices["invoice"].Status = "paid" + require.NoError(t, successWorker.Work(t.Context(), successJob)) + require.NoError(t, successWorker.Work(t.Context(), successJob)) + require.Nil(t, db.issueForOrg(org.ID, database.BillingIssueTypePaymentFailed)) + require.Equal(t, 2, sender.attemptCount()) + require.NoError(t, failureWorker.Work(t.Context(), paymentFailedJob(org.BillingCustomerID, "invoice", older))) + require.Nil(t, db.issueForOrg(org.ID, database.BillingIssueTypePaymentFailed)) + require.Equal(t, 2, sender.attemptCount()) + + // An older failure for another invoice is retained without moving aggregate time backward. + biller.invoices["newer-invoice"] = &billing.Invoice{ID: "newer-invoice", Status: "open"} + biller.invoices["older-invoice"] = &billing.Invoice{ID: "older-invoice", Status: "open"} + require.NoError(t, failureWorker.Work(t.Context(), paymentFailedJob(org.BillingCustomerID, "newer-invoice", newer))) + require.NoError(t, failureWorker.Work(t.Context(), paymentFailedJob(org.BillingCustomerID, "older-invoice", older))) + issue = db.issueForOrg(org.ID, database.BillingIssueTypePaymentFailed) + require.Equal(t, newer, issue.EventTime) + require.Len(t, issue.Metadata.(*database.BillingIssueMetadataPaymentFailed).Invoices, 2) +} + +func TestSubscriptionCancellationResumesPartialHibernationAndContinuesOrganizations(t *testing.T) { + // A replacement subscription and a one-time second-page failure must not block + // later organizations; a retry must converge all already-durable project state. + now := time.Date(2026, time.July, 23, 12, 0, 0, 0, time.UTC) + adm, db, biller, sender := newBillingLifecycleFixture(t) + activeOrg := db.addOrganization("a-active") + brokenOrg := db.addOrganization("b-broken") + laterOrg := db.addOrganization("c-later") + activeIssue := db.addIssue(activeOrg.ID, database.BillingIssueTypeSubscriptionCancelled, false, &database.BillingIssueMetadataSubscriptionCancelled{EndDate: now.AddDate(0, 0, -1)}, now) + brokenIssue := db.addIssue(brokenOrg.ID, database.BillingIssueTypeSubscriptionCancelled, false, &database.BillingIssueMetadataSubscriptionCancelled{EndDate: now.AddDate(0, 0, -1)}, now) + laterIssue := db.addIssue(laterOrg.ID, database.BillingIssueTypeSubscriptionCancelled, false, &database.BillingIssueMetadataSubscriptionCancelled{EndDate: now.AddDate(0, 0, -1)}, now) + db.addProjects(activeOrg.ID, 1) + db.addProjects(brokenOrg.ID, 11) + db.addProjects(laterOrg.ID, 1) + db.pageFailures[projectPageKey(brokenOrg.ID, "project-09")] = 1 + biller.subscriptions[activeOrg.BillingCustomerID] = &billing.Subscription{ID: "replacement", Plan: &billing.Plan{ID: "paid"}} + worker := &SubscriptionCancellationCheckWorker{admin: adm, logger: zap.NewNop(), now: func() time.Time { return now }} + + require.ErrorContains(t, worker.Work(t.Context(), nil), "second project page unavailable") + require.True(t, db.isProcessed(activeIssue.ID)) + require.False(t, db.projectHibernated(activeOrg.ID, "project-00")) + require.Equal(t, 7, db.organization(activeOrg.ID).QuotaProjects) + require.False(t, db.isProcessed(brokenIssue.ID)) + for i := range 10 { + require.True(t, db.projectHibernated(brokenOrg.ID, fmt.Sprintf("project-%02d", i))) + } + require.False(t, db.projectHibernated(brokenOrg.ID, "project-10")) + require.True(t, db.isProcessed(laterIssue.ID)) + require.True(t, db.projectHibernated(laterOrg.ID, "project-00")) + require.Equal(t, 1, sender.attemptCount()) + + require.NoError(t, worker.Work(t.Context(), nil)) + require.True(t, db.isProcessed(brokenIssue.ID)) + for i := range 11 { + require.True(t, db.projectHibernated(brokenOrg.ID, fmt.Sprintf("project-%02d", i))) + } + require.Equal(t, 2, db.hibernateCalls[projectKey(brokenOrg.ID, "project-00")]) + require.Equal(t, 2, sender.attemptCount()) + + require.NoError(t, worker.Work(t.Context(), nil)) + require.Equal(t, 2, sender.attemptCount()) +} + +func TestTrialEndReplacementSubscriptionDoesNotBlockLaterOrganization(t *testing.T) { + // A stale trial issue protected by a replacement subscription is retired, and + // the next organization still transitions atomically into its grace period. + now := time.Date(2026, time.July, 23, 12, 0, 0, 0, time.UTC) + adm, db, biller, sender := newBillingLifecycleFixture(t) + staleOrg := db.addOrganization("a-stale-trial") + validOrg := db.addOrganization("b-valid-trial") + db.addIssue(staleOrg.ID, database.BillingIssueTypeOnTrial, true, &database.BillingIssueMetadataOnTrial{SubID: "old", PlanID: "trial", EndDate: now, GracePeriodEndDate: now.AddDate(0, 0, 7)}, now) + db.addIssue(validOrg.ID, database.BillingIssueTypeOnTrial, true, &database.BillingIssueMetadataOnTrial{SubID: "expected", PlanID: "trial", EndDate: now, GracePeriodEndDate: now.AddDate(0, 0, 7)}, now) + biller.subscriptions[staleOrg.BillingCustomerID] = &billing.Subscription{ID: "replacement", Plan: &billing.Plan{ID: "paid"}} + biller.subscriptions[validOrg.BillingCustomerID] = &billing.Subscription{ID: "expected", Plan: &billing.Plan{ID: "trial"}} + worker := &TrialEndCheckWorker{admin: adm, logger: zap.NewNop(), now: func() time.Time { return now }} + + require.NoError(t, worker.Work(t.Context(), nil)) + require.Nil(t, db.issueForOrg(staleOrg.ID, database.BillingIssueTypeOnTrial)) + require.Nil(t, db.issueForOrg(staleOrg.ID, database.BillingIssueTypeTrialEnded)) + require.Nil(t, db.issueForOrg(validOrg.ID, database.BillingIssueTypeOnTrial)) + require.NotNil(t, db.issueForOrg(validOrg.ID, database.BillingIssueTypeTrialEnded)) + require.Equal(t, 1, sender.attemptCount()) +} + +func TestBillingLifecyclePostgresClaimsAndTransitions(t *testing.T) { + // PostgreSQL is used only for the atomic claim and unique transactional + // transition that an in-memory fixture cannot faithfully prove. + if testing.Short() { + t.Skip("requires a PostgreSQL test container") + } + pg := pgtestcontainer.New(t) + t.Cleanup(func() { pg.Terminate(t) }) + db, err := database.Open("postgres", pg.DatabaseURL, "") + require.NoError(t, err) + t.Cleanup(func() { require.NoError(t, db.Close()) }) + require.NoError(t, db.Migrate(t.Context())) + urls, err := admin.NewURLs("https://admin.example.com", "https://app.example.com") + require.NoError(t, err) + now := time.Date(2026, time.July, 23, 12, 0, 0, 0, time.UTC) + + t.Run("processed marker is a single atomic claim", func(t *testing.T) { + // Two workers holding the same stale issue may claim it once and emit one email. + org := insertPostgresLifecycleOrg(t, db, "atomic-claim") + metadata := &database.BillingIssueMetadataOnTrial{EndDate: now.AddDate(0, 0, 7)} + issue, err := db.UpsertBillingIssue(t.Context(), &database.UpsertBillingIssueOptions{OrgID: org.ID, Type: database.BillingIssueTypeOnTrial, Metadata: metadata, EventTime: now}) + require.NoError(t, err) + sender := &billingLifecycleEmailSender{} + adm := &admin.Service{DB: db, Email: email.New(sender), URLs: urls, Logger: zap.NewNop()} + worker := &TrialEndingSoonWorker{admin: adm, logger: zap.NewNop()} + + errs := runConcurrently(2, func() error { + return worker.processIssue(t.Context(), issue, metadata, now) + }) + require.Equal(t, 1, countNilErrors(errs)) + require.Equal(t, 1, sender.attemptCount()) + }) + + t.Run("trial transition remains unique under duplicate work", func(t *testing.T) { + // Competing transactions may create one grace issue and only the winner sends. + org := insertPostgresLifecycleOrg(t, db, "atomic-transition") + metadata := &database.BillingIssueMetadataOnTrial{SubID: "trial-sub", PlanID: "trial-plan", EndDate: now, GracePeriodEndDate: now.AddDate(0, 0, 7)} + issue, err := db.UpsertBillingIssue(t.Context(), &database.UpsertBillingIssueOptions{OrgID: org.ID, Type: database.BillingIssueTypeOnTrial, Metadata: metadata, EventTime: now}) + require.NoError(t, err) + require.NoError(t, db.UpdateBillingIssueOverdueAsProcessed(t.Context(), issue.ID)) + biller := &billingLifecycleBiller{subscriptions: map[string]*billing.Subscription{ + org.BillingCustomerID: {ID: "trial-sub", Plan: &billing.Plan{ID: "trial-plan"}}, + }, invoices: make(map[string]*billing.Invoice)} + sender := &billingLifecycleEmailSender{} + adm := &admin.Service{DB: db, Biller: biller, Email: email.New(sender), URLs: urls, Logger: zap.NewNop()} + worker := &TrialEndCheckWorker{admin: adm, logger: zap.NewNop()} + + errs := runConcurrently(2, func() error { + return worker.processIssue(t.Context(), issue, metadata) + }) + require.Equal(t, 1, countNilErrors(errs)) + require.Equal(t, 1, sender.attemptCount()) + _, err = db.FindBillingIssueByTypeForOrg(t.Context(), org.ID, database.BillingIssueTypeOnTrial) + require.ErrorIs(t, err, database.ErrNotFound) + ended, err := db.FindBillingIssueByTypeForOrg(t.Context(), org.ID, database.BillingIssueTypeTrialEnded) + require.NoError(t, err) + require.True(t, ended.EventTime.Equal(now)) + }) +} + +type billingLifecycleIssue struct { + issue *database.BillingIssue + processed bool +} + +type billingLifecycleDB struct { + database.DB + organizations map[string]*database.Organization + issues map[string]*billingLifecycleIssue + projects map[string][]*database.Project + processedFailures map[string]int + pageFailures map[string]int + hibernateCalls map[string]int + commitFailures int + nextIssueID int +} + +type billingLifecycleTxContextKey struct{} + +type billingLifecycleTxState struct { + db *billingLifecycleDB + issues map[string]*billingLifecycleIssue + closed bool +} + +func (d *billingLifecycleDB) NewTx(ctx context.Context, _ bool) (context.Context, database.Tx, error) { + tx := &billingLifecycleTxState{db: d, issues: cloneLifecycleIssueStore(d.issues)} + return context.WithValue(ctx, billingLifecycleTxContextKey{}, tx), tx, nil +} + +func (tx *billingLifecycleTxState) Commit() error { + if tx.closed { + return errors.New("transaction already closed") + } + tx.closed = true + if tx.db.commitFailures > 0 { + tx.db.commitFailures-- + return errors.New("injected commit failure") + } + tx.db.issues = cloneLifecycleIssueStore(tx.issues) + return nil +} + +func (tx *billingLifecycleTxState) Rollback() error { + tx.closed = true + return nil +} + +func (d *billingLifecycleDB) FindBillingIssueByTypeAndOverdueProcessed(ctx context.Context, issueType database.BillingIssueType, processed bool) ([]*database.BillingIssue, error) { + var issues []*database.BillingIssue + for _, record := range d.issueStore(ctx) { + if record.issue.Type == issueType && record.processed == processed { + issues = append(issues, cloneBillingIssue(record.issue)) + } + } + sort.Slice(issues, func(i, j int) bool { + if issues[i].OrgID == issues[j].OrgID { + return issues[i].ID < issues[j].ID + } + return issues[i].OrgID < issues[j].OrgID + }) + return issues, nil +} + +func (d *billingLifecycleDB) FindBillingIssueByTypeForOrg(ctx context.Context, orgID string, issueType database.BillingIssueType) (*database.BillingIssue, error) { + for _, record := range d.issueStore(ctx) { + if record.issue.OrgID == orgID && record.issue.Type == issueType { + return cloneBillingIssue(record.issue), nil + } + } + return nil, database.ErrNotFound +} + +func (d *billingLifecycleDB) UpsertBillingIssue(ctx context.Context, opts *database.UpsertBillingIssueOptions) (*database.BillingIssue, error) { + store := d.issueStore(ctx) + for _, record := range store { + if record.issue.OrgID == opts.OrgID && record.issue.Type == opts.Type { + record.issue.Metadata = cloneBillingMetadata(opts.Metadata) + record.issue.EventTime = opts.EventTime + return cloneBillingIssue(record.issue), nil + } + } + d.nextIssueID++ + issue := &database.BillingIssue{ + ID: fmt.Sprintf("generated-issue-%d", d.nextIssueID), + OrgID: opts.OrgID, + Type: opts.Type, + Metadata: cloneBillingMetadata(opts.Metadata), + EventTime: opts.EventTime, + } + store[issue.ID] = &billingLifecycleIssue{issue: issue} + return cloneBillingIssue(issue), nil +} + +func (d *billingLifecycleDB) UpdateBillingIssueOverdueAsProcessed(ctx context.Context, id string) error { + if d.processedFailures[id] > 0 { + d.processedFailures[id]-- + return errors.New("injected processed marker failure") + } + record := d.issueStore(ctx)[id] + if record == nil || record.processed { + return database.ErrNotFound + } + record.processed = true + return nil +} + +func (d *billingLifecycleDB) DeleteBillingIssue(ctx context.Context, id string) error { + store := d.issueStore(ctx) + if store[id] == nil { + return database.ErrNotFound + } + delete(store, id) + return nil +} + +func (d *billingLifecycleDB) FindOrganization(_ context.Context, id string) (*database.Organization, error) { + org := d.organizations[id] + if org == nil { + return nil, database.ErrNotFound + } + copy := *org + return ©, nil +} + +func (d *billingLifecycleDB) FindOrganizationForBillingCustomerID(_ context.Context, customerID string) (*database.Organization, error) { + for _, org := range d.organizations { + if org.BillingCustomerID == customerID { + copy := *org + return ©, nil + } + } + return nil, database.ErrNotFound +} + +func (d *billingLifecycleDB) CountProjectsForOrganization(_ context.Context, orgID string) (int, error) { + return len(d.projects[orgID]), nil +} + +func (d *billingLifecycleDB) UpdateOrganization(_ context.Context, id string, opts *database.UpdateOrganizationOptions) (*database.Organization, error) { + org := d.organizations[id] + if org == nil { + return nil, database.ErrNotFound + } + org.Name = opts.Name + org.DisplayName = opts.DisplayName + org.Description = opts.Description + org.QuotaProjects = opts.QuotaProjects + org.QuotaDeployments = opts.QuotaDeployments + org.QuotaSlotsTotal = opts.QuotaSlotsTotal + org.QuotaSlotsPerDeployment = opts.QuotaSlotsPerDeployment + org.QuotaOutstandingInvites = opts.QuotaOutstandingInvites + org.QuotaStorageLimitBytesPerDeployment = opts.QuotaStorageLimitBytesPerDeployment + org.QuotaSeats = opts.QuotaSeats + org.BillingCustomerID = opts.BillingCustomerID + org.PaymentCustomerID = opts.PaymentCustomerID + org.BillingEmail = opts.BillingEmail + org.BillingPlanName = opts.BillingPlanName + org.BillingPlanDisplayName = opts.BillingPlanDisplayName + copy := *org + return ©, nil +} + +func (d *billingLifecycleDB) FindProjectsForOrganization(_ context.Context, orgID, afterProjectName string, limit int) ([]*database.Project, error) { + key := projectPageKey(orgID, afterProjectName) + if d.pageFailures[key] > 0 { + d.pageFailures[key]-- + return nil, errors.New("second project page unavailable") + } + projects := d.projects[orgID] + start := 0 + for start < len(projects) && projects[start].Name <= afterProjectName { + start++ + } + end := min(start+limit, len(projects)) + result := make([]*database.Project, 0, end-start) + for _, project := range projects[start:end] { + copy := *project + result = append(result, ©) + } + return result, nil +} + +func (d *billingLifecycleDB) FindDeploymentsForProject(context.Context, string, string, string) ([]*database.Deployment, error) { + return nil, nil +} + +func (d *billingLifecycleDB) UpdateProject(_ context.Context, id string, opts *database.UpdateProjectOptions) (*database.Project, error) { + for orgID, projects := range d.projects { + for _, project := range projects { + if project.ID != id { + continue + } + project.PrimaryDeploymentID = opts.PrimaryDeploymentID + d.hibernateCalls[projectKey(orgID, project.Name)]++ + copy := *project + return ©, nil + } + } + return nil, database.ErrNotFound +} + +func (d *billingLifecycleDB) issueStore(ctx context.Context) map[string]*billingLifecycleIssue { + if tx, ok := ctx.Value(billingLifecycleTxContextKey{}).(*billingLifecycleTxState); ok { + return tx.issues + } + return d.issues +} + +func (d *billingLifecycleDB) addOrganization(name string) *database.Organization { + org := &database.Organization{ + ID: name, + Name: name, + DisplayName: strings.ToUpper(name), + QuotaProjects: 7, + QuotaDeployments: 7, + QuotaSlotsTotal: 7, + QuotaSlotsPerDeployment: 7, + QuotaOutstandingInvites: 7, + QuotaStorageLimitBytesPerDeployment: 7, + QuotaSeats: 7, + BillingCustomerID: "customer-" + name, + PaymentCustomerID: "payment-" + name, + BillingEmail: name + "@example.com", + } + d.organizations[org.ID] = org + return org +} + +func (d *billingLifecycleDB) addIssue(orgID string, issueType database.BillingIssueType, processed bool, metadata database.BillingIssueMetadata, eventTime time.Time) *database.BillingIssue { + d.nextIssueID++ + issue := &database.BillingIssue{ID: fmt.Sprintf("issue-%d", d.nextIssueID), OrgID: orgID, Type: issueType, Metadata: cloneBillingMetadata(metadata), EventTime: eventTime} + d.issues[issue.ID] = &billingLifecycleIssue{issue: issue, processed: processed} + return cloneBillingIssue(issue) +} + +func (d *billingLifecycleDB) addProjects(orgID string, count int) { + for i := range count { + deploymentID := fmt.Sprintf("deployment-%s-%02d", orgID, i) + d.projects[orgID] = append(d.projects[orgID], &database.Project{ + ID: fmt.Sprintf("project-id-%s-%02d", orgID, i), + OrganizationID: orgID, + Name: fmt.Sprintf("project-%02d", i), + PrimaryDeploymentID: &deploymentID, + }) + } +} + +func (d *billingLifecycleDB) isProcessed(id string) bool { + return d.issues[id] != nil && d.issues[id].processed +} + +func (d *billingLifecycleDB) issueForOrg(orgID string, issueType database.BillingIssueType) *database.BillingIssue { + issue, err := d.FindBillingIssueByTypeForOrg(context.Background(), orgID, issueType) + if err != nil { + return nil + } + return issue +} + +func (d *billingLifecycleDB) organization(id string) *database.Organization { + org := d.organizations[id] + copy := *org + return © +} + +func (d *billingLifecycleDB) projectHibernated(orgID, name string) bool { + for _, project := range d.projects[orgID] { + if project.Name == name { + return project.PrimaryDeploymentID == nil + } + } + return false +} + +type billingLifecycleBiller struct { + billing.Biller + subscriptions map[string]*billing.Subscription + invoices map[string]*billing.Invoice + cancellations []string +} + +func (b *billingLifecycleBiller) GetActiveSubscription(_ context.Context, customerID string) (*billing.Subscription, error) { + subscription := b.subscriptions[customerID] + if subscription == nil { + return nil, billing.ErrNotFound + } + copy := *subscription + if subscription.Plan != nil { + plan := *subscription.Plan + copy.Plan = &plan + } + return ©, nil +} + +func (b *billingLifecycleBiller) CancelSubscriptionsForCustomer(_ context.Context, customerID string, _ billing.SubscriptionCancellationOption) (time.Time, error) { + b.cancellations = append(b.cancellations, customerID) + delete(b.subscriptions, customerID) + return time.Time{}, nil +} + +func (b *billingLifecycleBiller) GetInvoice(_ context.Context, invoiceID string) (*billing.Invoice, error) { + invoice := b.invoices[invoiceID] + if invoice == nil { + return nil, billing.ErrNotFound + } + copy := *invoice + return ©, nil +} + +func (b *billingLifecycleBiller) IsInvoiceValid(_ context.Context, invoice *billing.Invoice) bool { + return invoice != nil && !strings.EqualFold(invoice.Status, "void") +} + +func (b *billingLifecycleBiller) IsInvoicePaid(_ context.Context, invoice *billing.Invoice) bool { + return invoice != nil && strings.EqualFold(invoice.Status, "paid") +} + +type billingLifecycleEmailAttempt struct { + to string + subject string + failed bool +} + +type billingLifecycleEmailSender struct { + mu sync.Mutex + attempts []billingLifecycleEmailAttempt + failures int +} + +func (s *billingLifecycleEmailSender) Send(toEmail, _ string, subject, _ string) error { + s.mu.Lock() + defer s.mu.Unlock() + attempt := billingLifecycleEmailAttempt{to: toEmail, subject: subject} + if s.failures > 0 { + s.failures-- + attempt.failed = true + s.attempts = append(s.attempts, attempt) + return errors.New("injected email failure") + } + s.attempts = append(s.attempts, attempt) + return nil +} + +func (s *billingLifecycleEmailSender) attemptCount() int { + s.mu.Lock() + defer s.mu.Unlock() + return len(s.attempts) +} + +func (s *billingLifecycleEmailSender) successCount() int { + s.mu.Lock() + defer s.mu.Unlock() + count := 0 + for _, attempt := range s.attempts { + if !attempt.failed { + count++ + } + } + return count +} + +func newBillingLifecycleFixture(t *testing.T) (*admin.Service, *billingLifecycleDB, *billingLifecycleBiller, *billingLifecycleEmailSender) { + t.Helper() + db := &billingLifecycleDB{ + organizations: make(map[string]*database.Organization), + issues: make(map[string]*billingLifecycleIssue), + projects: make(map[string][]*database.Project), + processedFailures: make(map[string]int), + pageFailures: make(map[string]int), + hibernateCalls: make(map[string]int), + } + biller := &billingLifecycleBiller{subscriptions: make(map[string]*billing.Subscription), invoices: make(map[string]*billing.Invoice)} + sender := &billingLifecycleEmailSender{} + urls, err := admin.NewURLs("https://admin.example.com", "https://app.example.com") + require.NoError(t, err) + adm := &admin.Service{DB: db, Biller: biller, Email: email.New(sender), URLs: urls, Logger: zap.NewNop()} + return adm, db, biller, sender +} + +func insertPostgresLifecycleOrg(t *testing.T, db database.DB, name string) *database.Organization { + t.Helper() + org, err := db.InsertOrganization(t.Context(), &database.InsertOrganizationOptions{ + Name: name, + DisplayName: name, + QuotaProjects: 7, + QuotaDeployments: 7, + QuotaSlotsTotal: 7, + QuotaSeats: 7, + BillingCustomerID: "customer-" + name, + BillingEmail: name + "@example.com", + }) + require.NoError(t, err) + return org +} + +func runConcurrently(count int, fn func() error) []error { + start := make(chan struct{}) + results := make(chan error, count) + for range count { + go func() { + <-start + results <- fn() + }() + } + close(start) + errs := make([]error, 0, count) + for range count { + errs = append(errs, <-results) + } + return errs +} + +func countNilErrors(errs []error) int { + count := 0 + for _, err := range errs { + if err == nil { + count++ + } + } + return count +} + +func paymentFailedJob(customerID, invoiceID string, failedAt time.Time) *river.Job[PaymentFailedArgs] { + return &river.Job[PaymentFailedArgs]{Args: PaymentFailedArgs{ + BillingCustomerID: customerID, + InvoiceID: invoiceID, + InvoiceNumber: "number-" + invoiceID, + InvoiceURL: "https://billing.example.com/" + invoiceID, + Amount: "100", + Currency: "USD", + DueDate: failedAt.AddDate(0, 0, 1), + FailedAt: failedAt, + }} +} + +func paymentFailureMetadata(invoiceID string, failedAt time.Time) *database.BillingIssueMetadataPaymentFailed { + return &database.BillingIssueMetadataPaymentFailed{Invoices: map[string]*database.BillingIssueMetadataPaymentFailedMeta{ + invoiceID: { + ID: invoiceID, + FailedOn: failedAt, + GracePeriodEndDate: failedAt.AddDate(0, 0, database.BillingGracePeriodDays), + }, + }} +} + +func requireOrganizationQuotasZero(t *testing.T, org *database.Organization) { + t.Helper() + require.Zero(t, org.QuotaProjects) + require.Zero(t, org.QuotaDeployments) + require.Zero(t, org.QuotaSlotsTotal) + require.Zero(t, org.QuotaSlotsPerDeployment) + require.Zero(t, org.QuotaOutstandingInvites) + require.Zero(t, org.QuotaStorageLimitBytesPerDeployment) + require.Zero(t, org.QuotaSeats) +} + +func projectPageKey(orgID, after string) string { + return orgID + "\x00" + after +} + +func projectKey(orgID, name string) string { + return orgID + "\x00" + name +} + +func cloneLifecycleIssueStore(source map[string]*billingLifecycleIssue) map[string]*billingLifecycleIssue { + result := make(map[string]*billingLifecycleIssue, len(source)) + for id, record := range source { + result[id] = &billingLifecycleIssue{issue: cloneBillingIssue(record.issue), processed: record.processed} + } + return result +} + +func cloneBillingIssue(issue *database.BillingIssue) *database.BillingIssue { + if issue == nil { + return nil + } + copy := *issue + copy.Metadata = cloneBillingMetadata(issue.Metadata) + return © +} + +func cloneBillingMetadata(metadata database.BillingIssueMetadata) database.BillingIssueMetadata { + switch value := metadata.(type) { + case *database.BillingIssueMetadataOnTrial: + copy := *value + return © + case *database.BillingIssueMetadataTrialEnded: + copy := *value + return © + case *database.BillingIssueMetadataSubscriptionCancelled: + copy := *value + return © + case *database.BillingIssueMetadataPaymentFailed: + copy := &database.BillingIssueMetadataPaymentFailed{Invoices: make(map[string]*database.BillingIssueMetadataPaymentFailedMeta, len(value.Invoices))} + for id, invoice := range value.Invoices { + invoiceCopy := *invoice + copy.Invoices[id] = &invoiceCopy + } + return copy + default: + return metadata + } +} diff --git a/admin/jobs/river/subscription_handlers.go b/admin/jobs/river/subscription_handlers.go index f67c8f2015fb..71c5ee4bfdef 100644 --- a/admin/jobs/river/subscription_handlers.go +++ b/admin/jobs/river/subscription_handlers.go @@ -22,6 +22,7 @@ type SubscriptionCancellationCheckWorker struct { river.WorkerDefaults[SubscriptionCancellationCheckArgs] admin *admin.Service logger *zap.Logger + now func() time.Time } // Work This worker runs at end of the current subscription term after subscription cancellation @@ -35,107 +36,71 @@ func (w *SubscriptionCancellationCheckWorker) Work(ctx context.Context, job *riv return fmt.Errorf("failed to find orgs with subscription cancellation billing issue: %w", err) } + now := billingLifecycleNow(w.now) + var workErr error for _, issue := range cancelled { m := issue.Metadata.(*database.BillingIssueMetadataSubscriptionCancelled) - if time.Now().UTC().Before(m.EndDate.AddDate(0, 0, 1)) { + if now.Before(m.EndDate.AddDate(0, 0, 1)) { // subscription end date is not finished yet, continue to next org continue } - org, err := w.admin.DB.FindOrganization(ctx, issue.OrgID) - if err != nil { - return fmt.Errorf("failed to find organization: %w", err) + if err := w.processIssue(ctx, issue); err != nil { + workErr = errors.Join(workErr, fmt.Errorf("process subscription cancellation for org %q: %w", issue.OrgID, err)) } + } - // check if the org has any active subscription - sub, err := w.admin.Biller.GetActiveSubscription(ctx, org.BillingCustomerID) - if err != nil { - if !errors.Is(err, billing.ErrNotFound) { - return fmt.Errorf("failed to get subscriptions for org %q: %w", org.Name, err) - } - } - - if sub != nil { - w.logger.Warn("active subscription found for the org even after sub cancellation, skipping hibernation", zap.String("org_id", org.ID), zap.String("org_name", org.Name)) - return fmt.Errorf("active subscription found for the org %q", org.Name) - } - - // update quotas to 0 and hibernate all projects - _, err = w.admin.DB.UpdateOrganization(ctx, org.ID, &database.UpdateOrganizationOptions{ - Name: org.Name, - DisplayName: org.DisplayName, - Description: org.Description, - LogoAssetID: org.LogoAssetID, - LogoDarkAssetID: org.LogoDarkAssetID, - FaviconAssetID: org.FaviconAssetID, - ThumbnailAssetID: org.ThumbnailAssetID, - CustomDomain: org.CustomDomain, - DefaultProjectRoleID: org.DefaultProjectRoleID, - QuotaProjects: 0, - QuotaDeployments: 0, - QuotaSlotsTotal: 0, - QuotaSlotsPerDeployment: 0, - QuotaOutstandingInvites: 0, - QuotaStorageLimitBytesPerDeployment: 0, - QuotaSeats: 0, - BillingCustomerID: org.BillingCustomerID, - PaymentCustomerID: org.PaymentCustomerID, - BillingEmail: org.BillingEmail, - BillingPlanName: org.BillingPlanName, - BillingPlanDisplayName: org.BillingPlanDisplayName, - CreatedByUserID: org.CreatedByUserID, - }) - if err != nil { - return err - } - - // hibernate projects - limit := 10 - afterProjectName := "" - projectCount := 0 - for { - projs, err := w.admin.DB.FindProjectsForOrganization(ctx, org.ID, afterProjectName, limit) - if err != nil { - return err - } - projectCount += len(projs) - - 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 - } + return workErr +} - if len(projs) < limit { - break - } - } +func (w *SubscriptionCancellationCheckWorker) processIssue(ctx context.Context, issue *database.BillingIssue) error { + org, err := w.admin.DB.FindOrganization(ctx, issue.OrgID) + if err != nil { + return fmt.Errorf("failed to find organization: %w", err) + } - err = w.admin.Email.SendSubscriptionEnded(&email.SubscriptionEnded{ - ToEmail: org.BillingEmail, - ToName: org.Name, - OrgName: org.Name, - BillingURL: w.admin.URLs.Billing(org.Name, false), - }) - if err != nil { - w.logger.Error("failed to send subscription ended email", zap.String("org_id", org.ID), zap.String("org_name", org.Name), zap.String("billing_email", org.BillingEmail), zap.Error(err)) + sub, err := w.admin.Biller.GetActiveSubscription(ctx, org.BillingCustomerID) + if err != nil && !errors.Is(err, billing.ErrNotFound) { + return fmt.Errorf("failed to get subscriptions for org %q: %w", org.Name, err) + } + if sub != nil { + // A replacement subscription makes this cancellation issue stale. Retire it + // and continue so it cannot block cancellation checks for later organizations. + w.logger.Warn("active replacement subscription found after cancellation; skipping hibernation", zap.String("org_id", org.ID), zap.String("org_name", org.Name), zap.String("subscription_id", sub.ID)) + if err := w.admin.DB.UpdateBillingIssueOverdueAsProcessed(ctx, issue.ID); err != nil { + return fmt.Errorf("failed to retire stale subscription cancellation issue: %w", err) } + return nil + } - w.logger.Warn("projects hibernated due to subscription cancellation", - zap.String("org_id", org.ID), - zap.String("org_name", org.Name), - zap.Int("count_of_projects", projectCount), - zap.String("user_email", org.BillingEmail), - ) - - // mark the billing issue as processed - err = w.admin.DB.UpdateBillingIssueOverdueAsProcessed(ctx, issue.ID) - if err != nil { - return fmt.Errorf("failed to update billing issue as processed: %w", err) - } + if err := zeroOrganizationQuotas(ctx, w.admin, org); err != nil { + return err + } + projectCount, err := hibernateProjectsForOrganization(ctx, w.admin, org.ID) + if err != nil { + return err } + w.logger.Warn("projects hibernated due to subscription cancellation", + zap.String("org_id", org.ID), + zap.String("org_name", org.Name), + zap.Int("count_of_projects", projectCount), + zap.String("user_email", org.BillingEmail), + ) + + // Persist completion before SMTP to prevent a failed processed-marker update + // from causing a duplicate email on retry. + if err := w.admin.DB.UpdateBillingIssueOverdueAsProcessed(ctx, issue.ID); err != nil { + return fmt.Errorf("failed to update billing issue as processed: %w", err) + } + err = w.admin.Email.SendSubscriptionEnded(&email.SubscriptionEnded{ + ToEmail: org.BillingEmail, + ToName: org.Name, + OrgName: org.Name, + BillingURL: w.admin.URLs.Billing(org.Name, false), + }) + if err != nil { + w.logger.Error("failed to send subscription ended email after marking it attempted", zap.String("org_id", org.ID), zap.String("org_name", org.Name), zap.String("billing_email", org.BillingEmail), zap.Error(err)) + } return nil } diff --git a/admin/jobs/river/trial_checks.go b/admin/jobs/river/trial_checks.go index f6425f5bfc7b..11092e246d54 100644 --- a/admin/jobs/river/trial_checks.go +++ b/admin/jobs/river/trial_checks.go @@ -22,6 +22,7 @@ type TrialEndingSoonWorker struct { river.WorkerDefaults[TrialEndingSoonArgs] admin *admin.Service logger *zap.Logger + now func() time.Time } func (w *TrialEndingSoonWorker) Work(ctx context.Context, job *river.Job[TrialEndingSoonArgs]) error { @@ -34,58 +35,60 @@ func (w *TrialEndingSoonWorker) Work(ctx context.Context, job *river.Job[TrialEn return fmt.Errorf("failed to find organization with billing issue: %w", err) } + now := billingLifecycleNow(w.now) + var workErr error for _, o := range onTrialOrgs { m := o.Metadata.(*database.BillingIssueMetadataOnTrial) - if time.Now().UTC().Before(m.EndDate.AddDate(0, 0, -7)) { + if now.Before(m.EndDate.AddDate(0, 0, -7)) { // trial end date is more than 7 days away, move to next org continue } - // trial period ending soon, log warn and send email - org, err := w.admin.DB.FindOrganization(ctx, o.OrgID) - if err != nil { - return fmt.Errorf("failed to find organization: %w", err) + if err := w.processIssue(ctx, o, m, now); err != nil { + workErr = errors.Join(workErr, fmt.Errorf("process trial ending soon for org %q: %w", o.OrgID, err)) } + } - // remaining days in the trial period - daysRemaining := int(m.EndDate.Sub(time.Now().UTC()).Hours() / 24) - if daysRemaining < 0 { - daysRemaining = 0 - } + return workErr +} - // number of projects for the org - projects, err := w.admin.DB.CountProjectsForOrganization(ctx, org.ID) - if err != nil { - return fmt.Errorf("failed to count projects for org %q: %w", org.Name, err) - } +func (w *TrialEndingSoonWorker) processIssue(ctx context.Context, issue *database.BillingIssue, metadata *database.BillingIssueMetadataOnTrial, now time.Time) error { + org, err := w.admin.DB.FindOrganization(ctx, issue.OrgID) + if err != nil { + return fmt.Errorf("failed to find organization: %w", err) + } - w.logger.Warn("trial ending soon", - zap.String("org_id", org.ID), - zap.String("org_name", org.Name), - zap.Time("trial_end_date", m.EndDate), - zap.String("user_email", org.BillingEmail), - zap.Int("count_of_projects", projects), - zap.Int("count_of_days_remaining", daysRemaining), - ) - - err = w.admin.Email.SendTrialEndingSoon(&email.TrialEndingSoon{ - ToEmail: org.BillingEmail, - ToName: org.Name, - OrgName: org.Name, - UpgradeURL: w.admin.URLs.Billing(org.Name, true), - TrialEndDate: m.EndDate, - }) - if err != nil { - return fmt.Errorf("failed to send trial ending soon email for org %q: %w", org.Name, err) - } + daysRemaining := max(int(metadata.EndDate.Sub(now).Hours()/24), 0) + projects, err := w.admin.DB.CountProjectsForOrganization(ctx, org.ID) + if err != nil { + return fmt.Errorf("failed to count projects for org %q: %w", org.Name, err) + } - // mark the billing issue as processed - err = w.admin.DB.UpdateBillingIssueOverdueAsProcessed(ctx, o.ID) - if err != nil { - return fmt.Errorf("failed to update billing issue as processed: %w", err) - } + w.logger.Warn("trial ending soon", + zap.String("org_id", org.ID), + zap.String("org_name", org.Name), + zap.Time("trial_end_date", metadata.EndDate), + zap.String("user_email", org.BillingEmail), + zap.Int("count_of_projects", projects), + zap.Int("count_of_days_remaining", daysRemaining), + ) + + // The current email sender has no idempotency key. Persisting the marker first + // intentionally gives lifecycle notifications at-most-once attempt semantics. + if err := w.admin.DB.UpdateBillingIssueOverdueAsProcessed(ctx, issue.ID); err != nil { + return fmt.Errorf("failed to update billing issue as processed: %w", err) } + err = w.admin.Email.SendTrialEndingSoon(&email.TrialEndingSoon{ + ToEmail: org.BillingEmail, + ToName: org.Name, + OrgName: org.Name, + UpgradeURL: w.admin.URLs.Billing(org.Name, true), + TrialEndDate: metadata.EndDate, + }) + if err != nil { + w.logger.Error("failed to send trial ending soon email after marking it attempted", zap.String("org_id", org.ID), zap.String("org_name", org.Name), zap.Error(err)) + } return nil } @@ -97,6 +100,7 @@ type TrialEndCheckWorker struct { river.WorkerDefaults[TrialEndCheckArgs] admin *admin.Service logger *zap.Logger + now func() time.Time } func (w *TrialEndCheckWorker) Work(ctx context.Context, job *river.Job[TrialEndCheckArgs]) error { @@ -109,100 +113,100 @@ func (w *TrialEndCheckWorker) Work(ctx context.Context, job *river.Job[TrialEndC return fmt.Errorf("failed to find organization with billing issue: %w", err) } + now := billingLifecycleNow(w.now) + var workErr error for _, o := range onTrialOrgs { m := o.Metadata.(*database.BillingIssueMetadataOnTrial) - if time.Now().UTC().Before(m.EndDate) { + if now.Before(m.EndDate) { // trial end date is not finished yet, move to next org continue } - // trial period has ended, log warn and send email - org, err := w.admin.DB.FindOrganization(ctx, o.OrgID) - if err != nil { - return fmt.Errorf("failed to find organization: %w", err) - } - - sub, err := w.admin.Biller.GetActiveSubscription(ctx, org.BillingCustomerID) - if err != nil { - return fmt.Errorf("failed to get subscriptions for org %q: %w", org.Name, err) - } - if sub.ID != m.SubID || sub.Plan.ID != m.PlanID { - w.logger.Warn("trial period has ended, but org has different active subscription, please check manually", zap.String("org_id", org.ID), zap.String("org_name", org.Name), zap.String("sub_id", sub.ID), zap.String("sub_plan_id", sub.Plan.ID), zap.String("expected_sub_id", m.SubID), zap.String("expected_sub_plan_id", m.PlanID)) - continue - } - - // number of projects for the org - projects, err := w.admin.DB.CountProjectsForOrganization(ctx, org.ID) - if err != nil { - return fmt.Errorf("failed to count projects for org %q: %w", org.Name, err) + if err := w.processIssue(ctx, o, m); err != nil { + workErr = errors.Join(workErr, fmt.Errorf("process trial end for org %q: %w", o.OrgID, err)) } + } - w.logger.Warn("trial period has ended", - zap.String("org_id", org.ID), - zap.String("org_name", org.Name), - zap.String("user_email", org.BillingEmail), - zap.Int("count_of_projects", projects), - ) + return workErr +} - cctx, tx, err := w.admin.DB.NewTx(ctx, false) - if err != nil { - return fmt.Errorf("failed to start transaction: %w", err) - } +func (w *TrialEndCheckWorker) processIssue(ctx context.Context, issue *database.BillingIssue, metadata *database.BillingIssueMetadataOnTrial) error { + org, err := w.admin.DB.FindOrganization(ctx, issue.OrgID) + if err != nil { + return fmt.Errorf("failed to find organization: %w", err) + } - _, err = w.admin.DB.UpsertBillingIssue(cctx, &database.UpsertBillingIssueOptions{ - OrgID: org.ID, - Type: database.BillingIssueTypeTrialEnded, - Metadata: &database.BillingIssueMetadataTrialEnded{ - SubID: m.SubID, - PlanID: m.PlanID, - EndDate: m.EndDate, - GracePeriodEndDate: m.GracePeriodEndDate, - }, - EventTime: m.EndDate, - }) - if err != nil { - prevErr := err - err = tx.Rollback() - if err != nil { - return fmt.Errorf("failed to rollback transaction: %w", err) - } - return fmt.Errorf("failed to add billing error: %w", prevErr) + sub, err := w.admin.Biller.GetActiveSubscription(ctx, org.BillingCustomerID) + if err != nil { + return fmt.Errorf("failed to get subscriptions for org %q: %w", org.Name, err) + } + if sub == nil || sub.Plan == nil { + return fmt.Errorf("active subscription for org %q is incomplete", org.Name) + } + if sub.ID != metadata.SubID || sub.Plan.ID != metadata.PlanID { + w.logger.Warn("trial period has ended, but org has different active subscription; removing stale trial state", zap.String("org_id", org.ID), zap.String("org_name", org.Name), zap.String("sub_id", sub.ID), zap.String("sub_plan_id", sub.Plan.ID), zap.String("expected_sub_id", metadata.SubID), zap.String("expected_sub_plan_id", metadata.PlanID)) + if err := w.admin.DB.DeleteBillingIssue(ctx, issue.ID); err != nil { + return fmt.Errorf("failed to delete stale on-trial billing issue: %w", err) } + return nil + } - // delete the on-trial billing issue - err = w.admin.DB.DeleteBillingIssue(cctx, o.ID) - if err != nil { - prevErr := err - err = tx.Rollback() - if err != nil { - return fmt.Errorf("failed to rollback transaction: %w", err) - } - return fmt.Errorf("failed to delete billing issue: %w", prevErr) + projects, err := w.admin.DB.CountProjectsForOrganization(ctx, org.ID) + if err != nil { + return fmt.Errorf("failed to count projects for org %q: %w", org.Name, err) + } + w.logger.Warn("trial period has ended", + zap.String("org_id", org.ID), + zap.String("org_name", org.Name), + zap.String("user_email", org.BillingEmail), + zap.Int("count_of_projects", projects), + ) + + cctx, tx, err := w.admin.DB.NewTx(ctx, false) + if err != nil { + return fmt.Errorf("failed to start transaction: %w", err) + } + rollback := func(cause error) error { + if rollbackErr := tx.Rollback(); rollbackErr != nil { + return errors.Join(cause, fmt.Errorf("failed to rollback transaction: %w", rollbackErr)) } + return cause + } - // send email - err = w.admin.Email.SendTrialEnded(&email.TrialEnded{ - ToEmail: org.BillingEmail, - ToName: org.Name, - OrgName: org.Name, - UpgradeURL: w.admin.URLs.Billing(org.Name, true), - GracePeriodEndDate: m.GracePeriodEndDate, - }) - if err != nil { - prevErr := err - err = tx.Rollback() - if err != nil { - return fmt.Errorf("failed to rollback transaction: %w", err) - } - return fmt.Errorf("failed to send trial period ended email for org %q: %w", org.Name, prevErr) - } + _, err = w.admin.DB.UpsertBillingIssue(cctx, &database.UpsertBillingIssueOptions{ + OrgID: org.ID, + Type: database.BillingIssueTypeTrialEnded, + Metadata: &database.BillingIssueMetadataTrialEnded{ + SubID: metadata.SubID, + PlanID: metadata.PlanID, + EndDate: metadata.EndDate, + GracePeriodEndDate: metadata.GracePeriodEndDate, + }, + EventTime: metadata.EndDate, + }) + if err != nil { + return rollback(fmt.Errorf("failed to add billing issue: %w", err)) + } + if err := w.admin.DB.DeleteBillingIssue(cctx, issue.ID); err != nil { + return rollback(fmt.Errorf("failed to delete billing issue: %w", err)) + } - err = tx.Commit() - if err != nil { - return fmt.Errorf("failed to commit transaction: %w", err) - } + // Commit the state transition before attempting email. This makes a commit + // failure retryable without duplicating an externally visible notification. + if err := tx.Commit(); err != nil { + return fmt.Errorf("failed to commit transaction: %w", err) } + err = w.admin.Email.SendTrialEnded(&email.TrialEnded{ + ToEmail: org.BillingEmail, + ToName: org.Name, + OrgName: org.Name, + UpgradeURL: w.admin.URLs.Billing(org.Name, true), + GracePeriodEndDate: metadata.GracePeriodEndDate, + }) + if err != nil { + w.logger.Error("failed to send trial ended email after committing it as attempted", zap.String("org_id", org.ID), zap.String("org_name", org.Name), zap.Error(err)) + } return nil } @@ -214,6 +218,7 @@ type TrialGracePeriodCheckWorker struct { river.WorkerDefaults[TrialGracePeriodCheckArgs] admin *admin.Service logger *zap.Logger + now func() time.Time } func (w *TrialGracePeriodCheckWorker) Work(ctx context.Context, job *river.Job[TrialGracePeriodCheckArgs]) error { @@ -226,107 +231,135 @@ func (w *TrialGracePeriodCheckWorker) Work(ctx context.Context, job *river.Job[T return fmt.Errorf("failed to find organization with billing issue: %w", err) } + now := billingLifecycleNow(w.now) + var workErr error for _, o := range trailEndedOrgs { m := o.Metadata.(*database.BillingIssueMetadataTrialEnded) - if time.Now().UTC().Before(m.GracePeriodEndDate) { + if now.Before(m.GracePeriodEndDate) { // grace period end date is not finished yet, move to next org continue } - org, err := w.admin.DB.FindOrganization(ctx, o.OrgID) - if err != nil { - return fmt.Errorf("failed to find organization: %w", err) + if err := w.processIssue(ctx, o, m); err != nil { + workErr = errors.Join(workErr, fmt.Errorf("process trial grace period end for org %q: %w", o.OrgID, err)) } + } - // get active subscription for the org - sub, err := w.admin.Biller.GetActiveSubscription(ctx, org.BillingCustomerID) - if err != nil { - if !errors.Is(err, billing.ErrNotFound) { - return fmt.Errorf("failed to get subscriptions for org %q: %w", org.Name, err) - } + return workErr +} + +func (w *TrialGracePeriodCheckWorker) processIssue(ctx context.Context, issue *database.BillingIssue, metadata *database.BillingIssueMetadataTrialEnded) error { + org, err := w.admin.DB.FindOrganization(ctx, issue.OrgID) + if err != nil { + return fmt.Errorf("failed to find organization: %w", err) + } + + sub, err := w.admin.Biller.GetActiveSubscription(ctx, org.BillingCustomerID) + if err != nil && !errors.Is(err, billing.ErrNotFound) { + return fmt.Errorf("failed to get subscriptions for org %q: %w", org.Name, err) + } + if sub == nil { + w.logger.Warn("trial grace period has ended, but org has no active subscription", zap.String("org_id", org.ID), zap.String("org_name", org.Name)) + } else { + if sub.Plan == nil { + return fmt.Errorf("active subscription for org %q has no plan", org.Name) } - if sub == nil { - // might happen if previous job failed in middle - w.logger.Warn("trial grace period has ended, but org has no active subscription, please check", zap.String("org_id", org.ID), zap.String("org_name", org.Name)) - } else if sub.ID != m.SubID || sub.Plan.ID != m.PlanID { - w.logger.Warn("trial grace period has ended, but org has different active subscription, doing nothing, please check manually", zap.String("org_id", org.ID), zap.String("org_name", org.Name), zap.String("sub_id", sub.ID), zap.String("sub_plan_id", sub.Plan.ID), zap.String("expected_sub_id", m.SubID), zap.String("expected_sub_plan_id", m.PlanID)) - continue + if sub.ID != metadata.SubID || sub.Plan.ID != metadata.PlanID { + w.logger.Warn("trial grace period has ended, but org has a replacement subscription; retiring stale trial issue", zap.String("org_id", org.ID), zap.String("org_name", org.Name), zap.String("sub_id", sub.ID), zap.String("sub_plan_id", sub.Plan.ID), zap.String("expected_sub_id", metadata.SubID), zap.String("expected_sub_plan_id", metadata.PlanID)) + if err := w.admin.DB.UpdateBillingIssueOverdueAsProcessed(ctx, issue.ID); err != nil { + return fmt.Errorf("failed to retire stale trial-ended issue: %w", err) + } + return nil } + } - // cancel the subscription - _, err = w.admin.Biller.CancelSubscriptionsForCustomer(ctx, org.BillingCustomerID, billing.SubscriptionCancellationOptionImmediate) - if err != nil { - return fmt.Errorf("failed to cancel subscription for org %q: %w", org.Name, err) - } + if _, err := w.admin.Biller.CancelSubscriptionsForCustomer(ctx, org.BillingCustomerID, billing.SubscriptionCancellationOptionImmediate); err != nil { + return fmt.Errorf("failed to cancel subscription for org %q: %w", org.Name, err) + } + if err := zeroOrganizationQuotas(ctx, w.admin, org); err != nil { + return err + } + if _, err := hibernateProjectsForOrganization(ctx, w.admin, org.ID); err != nil { + return err + } + w.logger.Warn("projects hibernated due to trial grace period ended", zap.String("org_id", org.ID), zap.String("org_name", org.Name)) - // trial grace period ended, update quotas to 0 and hibernate all projects - _, err = w.admin.DB.UpdateOrganization(ctx, org.ID, &database.UpdateOrganizationOptions{ - Name: org.Name, - DisplayName: org.DisplayName, - Description: org.Description, - LogoAssetID: org.LogoAssetID, - LogoDarkAssetID: org.LogoDarkAssetID, - FaviconAssetID: org.FaviconAssetID, - ThumbnailAssetID: org.ThumbnailAssetID, - CustomDomain: org.CustomDomain, - DefaultProjectRoleID: org.DefaultProjectRoleID, - QuotaProjects: 0, - QuotaDeployments: 0, - QuotaSlotsTotal: 0, - QuotaSlotsPerDeployment: 0, - QuotaOutstandingInvites: 0, - QuotaStorageLimitBytesPerDeployment: 0, - QuotaSeats: 0, - BillingCustomerID: org.BillingCustomerID, - PaymentCustomerID: org.PaymentCustomerID, - BillingEmail: org.BillingEmail, - BillingPlanName: org.BillingPlanName, - BillingPlanDisplayName: org.BillingPlanDisplayName, - CreatedByUserID: org.CreatedByUserID, - }) - if err != nil { - return err - } + // The processed marker is the durable at-most-once notification marker. + if err := w.admin.DB.UpdateBillingIssueOverdueAsProcessed(ctx, issue.ID); err != nil { + return fmt.Errorf("failed to update billing issue as processed: %w", err) + } + err = w.admin.Email.SendTrialGracePeriodEnded(&email.TrialGracePeriodEnded{ + ToEmail: org.BillingEmail, + ToName: org.Name, + OrgName: org.Name, + UpgradeURL: w.admin.URLs.Billing(org.Name, true), + }) + if err != nil { + w.logger.Error("failed to send trial grace period ended email after marking it attempted", zap.String("org_id", org.ID), zap.String("org_name", org.Name), zap.Error(err)) + } + return nil +} - limit := 10 - afterProjectName := "" - for { - projs, err := w.admin.DB.FindProjectsForOrganization(ctx, org.ID, afterProjectName, limit) - if err != nil { - return err - } +func billingLifecycleNow(now func() time.Time) time.Time { + if now != nil { + return now().UTC() + } + return time.Now().UTC() +} - 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 - } +func zeroOrganizationQuotas(ctx context.Context, adm *admin.Service, org *database.Organization) error { + _, err := adm.DB.UpdateOrganization(ctx, org.ID, &database.UpdateOrganizationOptions{ + Name: org.Name, + DisplayName: org.DisplayName, + Description: org.Description, + LogoAssetID: org.LogoAssetID, + LogoDarkAssetID: org.LogoDarkAssetID, + FaviconAssetID: org.FaviconAssetID, + ThumbnailAssetID: org.ThumbnailAssetID, + CustomDomain: org.CustomDomain, + DefaultProjectRoleID: org.DefaultProjectRoleID, + QuotaProjects: 0, + QuotaDeployments: 0, + QuotaSlotsTotal: 0, + QuotaSlotsPerDeployment: 0, + QuotaOutstandingInvites: 0, + QuotaStorageLimitBytesPerDeployment: 0, + QuotaSeats: 0, + BillingCustomerID: org.BillingCustomerID, + PaymentCustomerID: org.PaymentCustomerID, + BillingEmail: org.BillingEmail, + BillingPlanName: org.BillingPlanName, + BillingPlanDisplayName: org.BillingPlanDisplayName, + CreatedByUserID: org.CreatedByUserID, + }) + if err != nil { + return fmt.Errorf("failed to zero quotas for org %q: %w", org.Name, err) + } + return nil +} - if len(projs) < limit { - break - } - } - w.logger.Warn("projects hibernated due to trial grace period ended", zap.String("org_id", org.ID), zap.String("org_name", org.Name)) - - // send email - err = w.admin.Email.SendTrialGracePeriodEnded(&email.TrialGracePeriodEnded{ - ToEmail: org.BillingEmail, - ToName: org.Name, - OrgName: org.Name, - UpgradeURL: w.admin.URLs.Billing(org.Name, true), - }) +func hibernateProjectsForOrganization(ctx context.Context, adm *admin.Service, orgID string) (int, error) { + const limit = 10 + afterProjectName := "" + projectCount := 0 + for { + projects, err := adm.DB.FindProjectsForOrganization(ctx, orgID, afterProjectName, limit) if err != nil { - return fmt.Errorf("failed to send trial grace period ended email for org %q: %w", org.Name, err) + return projectCount, fmt.Errorf("failed to find projects after %q: %w", afterProjectName, err) } - // mark the billing issue as processed - err = w.admin.DB.UpdateBillingIssueOverdueAsProcessed(ctx, o.ID) - if err != nil { - return fmt.Errorf("failed to update billing issue as processed: %w", err) + for _, project := range projects { + if _, err := adm.HibernateProject(ctx, project); err != nil { + return projectCount, fmt.Errorf("failed to hibernate project %q: %w", project.Name, err) + } + // Advance only after the project's durable hibernation update succeeds. + // Retries may revisit earlier pages; HibernateProject is idempotent. + afterProjectName = project.Name + projectCount++ } - } - return nil + if len(projects) < limit { + return projectCount, nil + } + } } diff --git a/admin/server/billing.go b/admin/server/billing.go index 9b45b7bed97e..6831b64e36dd 100644 --- a/admin/server/billing.go +++ b/admin/server/billing.go @@ -2,6 +2,7 @@ package server import ( "context" + "crypto/sha256" "errors" "fmt" "math" @@ -125,6 +126,10 @@ func (s *Server) UpdateBillingSubscription(ctx context.Context, req *adminv1.Upd } return nil, err } + if !plan.Public && !forceAccess { + return nil, status.Errorf(codes.FailedPrecondition, "cannot assign a private plan %q", plan.Name) + } + // for both free and legacy trial plan, start the credit trial if plan.PlanType == billing.FreePlanType || plan.PlanType == billing.TrailPlanType { bi, err := s.admin.DB.FindBillingIssueByTypeForOrg(ctx, org.ID, database.BillingIssueTypeNeverSubscribed) @@ -167,10 +172,6 @@ func (s *Server) UpdateBillingSubscription(ctx context.Context, req *adminv1.Upd } } - if !plan.Public && !forceAccess { - return nil, status.Errorf(codes.FailedPrecondition, "cannot assign a private plan %q", plan.Name) - } - // check for validation errors if not forced if !forceAccess { err = s.planChangeValidationChecks(ctx, org) @@ -208,31 +209,18 @@ func (s *Server) UpdateBillingSubscription(ctx context.Context, req *adminv1.Upd // schedule plan change oldPlan := sub.Plan if oldPlan.ID != plan.ID { - sub, err = s.admin.Biller.ChangeSubscriptionPlan(ctx, sub.ID, plan) - if err != nil { - return nil, err - } - planChange = true - s.logger.Named("billing").Info("plan changed", - zap.String("org_id", org.ID), - zap.String("org_name", org.Name), - zap.String("old_plan_id", oldPlan.ID), - zap.String("old_plan_name", oldPlan.Name), - zap.String("new_plan_id", sub.Plan.ID), - zap.String("new_plan_name", sub.Plan.Name), - ) - // Roll any leftover trial `credits` over to USD on a trial → paid upgrade. The trial plan is priced in `credits`; the paid plan is priced in USD, so an unused `credits` balance would otherwise sit dormant on the customer ledger. Transfer 1:1 (e.g., 80 credits → $80) by debiting the credits ledger and granting the same amount on USD. - if oldPlan.PlanType == billing.FreePlanType { + // Complete the transfer before changing plans so a retry can still identify an unfinished trial upgrade. Stable provider keys make ambiguous responses safe to replay. + if oldPlan.PlanType == billing.FreePlanType && plan.PlanType != billing.FreePlanType && plan.PlanType != billing.TrailPlanType { balance, err := s.admin.Biller.GetCustomerCreditBalance(ctx, org.BillingCustomerID, billing.CreditsCurrency) if err != nil { return nil, fmt.Errorf("failed to fetch trial credit balance for rollover: %w", err) } if balance > 0 { - if err := s.admin.Biller.GrantCustomerCredits(ctx, org.BillingCustomerID, balance, billing.USDCurrency, "Trial credit rollover on upgrade", nil); err != nil { + if err := s.admin.Biller.GrantCustomerCredits(ctx, org.BillingCustomerID, balance, billing.USDCurrency, "Trial credit rollover on upgrade", nil, planChangeIdempotencyKey("grant", org.ID, sub.ID, plan.ID)); err != nil { return nil, fmt.Errorf("failed to grant USD rollover credits: %w", err) } - if err := s.admin.Biller.DebitCustomerCredits(ctx, org.BillingCustomerID, balance, billing.CreditsCurrency, "Trial credit rollover on upgrade"); err != nil { + if err := s.admin.Biller.DebitCustomerCredits(ctx, org.BillingCustomerID, balance, billing.CreditsCurrency, "Trial credit rollover on upgrade", planChangeIdempotencyKey("debit", org.ID, sub.ID, plan.ID)); err != nil { return nil, fmt.Errorf("failed to debit trial credits during rollover: %w", err) } s.logger.Named("billing").Info("rolled over trial credits to USD", @@ -242,6 +230,20 @@ func (s *Server) UpdateBillingSubscription(ctx context.Context, req *adminv1.Upd ) } } + + sub, err = s.admin.Biller.ChangeSubscriptionPlan(ctx, sub.ID, plan, planChangeIdempotencyKey("plan", org.ID, sub.ID, plan.ID)) + if err != nil { + return nil, err + } + planChange = true + s.logger.Named("billing").Info("plan changed", + zap.String("org_id", org.ID), + zap.String("org_name", org.Name), + zap.String("old_plan_id", oldPlan.ID), + zap.String("old_plan_name", oldPlan.Name), + zap.String("new_plan_id", sub.Plan.ID), + zap.String("new_plan_name", sub.Plan.Name), + ) } } @@ -292,6 +294,11 @@ func (s *Server) UpdateBillingSubscription(ctx context.Context, req *adminv1.Upd }, nil } +func planChangeIdempotencyKey(operation, orgID, subscriptionID, planID string) string { + sum := sha256.Sum256([]byte(operation + "\x00" + orgID + "\x00" + subscriptionID + "\x00" + planID)) + return fmt.Sprintf("rill-plan-change-%s-%x", operation, sum[:16]) +} + // CancelBillingSubscription cancels the billing subscription for the organization func (s *Server) CancelBillingSubscription(ctx context.Context, req *adminv1.CancelBillingSubscriptionRequest) (*adminv1.CancelBillingSubscriptionResponse, error) { observability.AddRequestAttributes(ctx, attribute.String("args.org", req.Org)) @@ -429,7 +436,7 @@ func (s *Server) RenewBillingSubscription(ctx context.Context, req *adminv1.Rene if sub.Plan.ID != plan.ID { // change the plan, won't happen for new subscriptions - sub, err = s.admin.Biller.ChangeSubscriptionPlan(ctx, sub.ID, plan) + sub, err = s.admin.Biller.ChangeSubscriptionPlan(ctx, sub.ID, plan, "") if err != nil { return nil, err } @@ -782,7 +789,7 @@ func (s *Server) SudoGrantTrialCredits(ctx context.Context, req *adminv1.SudoGra description = fmt.Sprintf("Trial credits granted by superuser %s", claims.OwnerID()) } - if err := s.admin.Biller.GrantCustomerCredits(ctx, org.BillingCustomerID, req.AmountUsd, billing.CreditsCurrency, description, nil); err != nil { + if err := s.admin.Biller.GrantCustomerCredits(ctx, org.BillingCustomerID, req.AmountUsd, billing.CreditsCurrency, description, nil, ""); err != nil { return nil, fmt.Errorf("failed to grant credits: %w", err) } diff --git a/admin/server/billing_test.go b/admin/server/billing_test.go new file mode 100644 index 000000000000..7db9f9197f60 --- /dev/null +++ b/admin/server/billing_test.go @@ -0,0 +1,529 @@ +package server_test + +import ( + "context" + "errors" + "sync" + "testing" + "time" + + "github.com/rilldata/rill/admin/billing" + "github.com/rilldata/rill/admin/billing/payment" + "github.com/rilldata/rill/admin/client" + "github.com/rilldata/rill/admin/database" + "github.com/rilldata/rill/admin/testadmin" + adminv1 "github.com/rilldata/rill/proto/gen/rill/admin/v1" + "github.com/stretchr/testify/require" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +const ( + billingOpBalance = "balance" + billingOpGrant = "grant" + billingOpDebit = "debit" + billingOpPlan = "plan" + billingOpCreate = "create" +) + +var errInjectedBilling = errors.New("injected billing failure") + +func TestUpdateBillingSubscriptionSaga(t *testing.T) { + // The integration boundary uses real authorization and Postgres while external billing remains deterministic and fault-injectable. + fix := testadmin.New(t) + _, _ = fix.NewUser(t) // The first fixture user is a bootstrap superuser; use a regular user for quota and permission assertions. + user, owner := fix.NewUser(t) + _, outsider := fix.NewUser(t) + + t.Run("invalid requests have no external effects", func(t *testing.T) { + // Rejected callers and plan identifiers must not cross a mutating provider boundary. + tests := []struct { + name string + plan string + caller func(*client.Client, *client.Client) *client.Client + wantCode codes.Code + wantMessage string + }{ + {name: "permission denied", plan: "paid", caller: func(_, outsider *client.Client) *client.Client { return outsider }, wantCode: codes.PermissionDenied, wantMessage: "not allowed"}, + {name: "empty plan", plan: "", caller: func(owner, _ *client.Client) *client.Client { return owner }, wantCode: codes.InvalidArgument, wantMessage: "PlanName"}, + {name: "unknown plan", plan: "missing", caller: func(owner, _ *client.Client) *client.Client { return owner }, wantCode: codes.NotFound, wantMessage: "plan not found"}, + {name: "private plan", plan: "private", caller: func(owner, _ *client.Client) *client.Client { return owner }, wantCode: codes.FailedPrecondition, wantMessage: "private plan"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Authorization and plan validation must finish before any provider mutation or payment lookup. + org, biller := newBillingScenario(t, fix, owner) + provider := newTestPaymentProvider(true, true) + fix.Admin.Biller = biller + fix.Admin.PaymentProvider = provider + if tt.plan == "private" { + _, err := fix.Admin.DB.UpsertBillingIssue(t.Context(), &database.UpsertBillingIssueOptions{ + OrgID: org.ID, Type: database.BillingIssueTypeNeverSubscribed, Metadata: &database.BillingIssueMetadataNeverSubscribed{}, EventTime: time.Now(), + }) + require.NoError(t, err) + biller.prepareUnsubscribedCustomer() + } + + _, err := tt.caller(owner, outsider).UpdateBillingSubscription(t.Context(), &adminv1.UpdateBillingSubscriptionRequest{Org: org.Name, PlanName: tt.plan}) + require.Error(t, err) + require.Equal(t, tt.wantCode, status.Code(err)) + require.ErrorContains(t, err, tt.wantMessage) + require.Zero(t, biller.mutationAttempts(), "rejected requests must not mutate billing state") + require.Zero(t, provider.findCalls(), "rejected plans must not reach payment validation") + }) + } + }) + + t.Run("payment validation has no billing side effects", func(t *testing.T) { + // Payment readiness is a precondition, never a partially mutating plan-change step. + tests := []struct { + name string + hasPaymentMethod bool + hasAddress bool + wantMessage string + }{ + {name: "missing payment method", hasAddress: true, wantMessage: "no payment method found"}, + {name: "missing billing address", hasPaymentMethod: true, wantMessage: "no billing address found"}, + {name: "missing payment method and address", wantMessage: "no payment method found, no billing address found"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // A customer that cannot be billed must be rejected before credits or subscription state are touched. + org, biller := newBillingScenario(t, fix, owner) + provider := newTestPaymentProvider(tt.hasPaymentMethod, tt.hasAddress) + fix.Admin.Biller = biller + fix.Admin.PaymentProvider = provider + + _, err := owner.UpdateBillingSubscription(t.Context(), &adminv1.UpdateBillingSubscriptionRequest{Org: org.Name, PlanName: "paid"}) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err)) + require.ErrorContains(t, err, tt.wantMessage) + require.Zero(t, biller.mutationAttempts(), "payment validation failures must not mutate the provider") + require.Equal(t, 1, provider.findCalls()) + }) + } + }) + + t.Run("trial quota denial has no external effects", func(t *testing.T) { + // Exhausted trial quota must stop trial creation before credits or subscriptions are provisioned. + org, biller := newBillingScenario(t, fix, owner) + fix.Admin.Biller = biller + fix.Admin.PaymentProvider = newTestPaymentProvider(true, true) + _, err := fix.Admin.DB.UpsertBillingIssue(t.Context(), &database.UpsertBillingIssueOptions{ + OrgID: org.ID, Type: database.BillingIssueTypeNeverSubscribed, Metadata: &database.BillingIssueMetadataNeverSubscribed{}, EventTime: time.Now(), + }) + require.NoError(t, err) + _, err = fix.Admin.DB.UpdateUser(t.Context(), user.ID, &database.UpdateUserOptions{ + DisplayName: user.DisplayName, PhotoURL: user.PhotoURL, GithubUsername: user.GithubUsername, GithubToken: user.GithubToken, + GithubTokenExpiresOn: user.GithubTokenExpiresOn, GithubRefreshToken: user.GithubRefreshToken, QuotaSingleuserOrgs: user.QuotaSingleuserOrgs, + QuotaTrialOrgs: 0, PreferenceTimeZone: user.PreferenceTimeZone, + }) + require.NoError(t, err) + updatedUser, err := fix.Admin.DB.FindUser(t.Context(), user.ID) + require.NoError(t, err) + require.Zero(t, updatedUser.QuotaTrialOrgs) + require.NotNil(t, org.CreatedByUserID) + + _, err = owner.UpdateBillingSubscription(t.Context(), &adminv1.UpdateBillingSubscriptionRequest{Org: org.Name, PlanName: "trial"}) + require.Error(t, err) + require.Equal(t, codes.FailedPrecondition, status.Code(err), err) + require.ErrorContains(t, err, "trial orgs quota") + require.Zero(t, biller.mutationAttempts(), "quota denial must not grant trial credits or create a subscription") + }) + + t.Run("partial failures converge on retry", func(t *testing.T) { + // The matrix covers failures before and after every non-atomic provider mutation plus the final Postgres write. + tests := []struct { + name string + fault *testBillingFault + failDatabaseUpdate bool + wantErr string + wantInitialCredits float64 + wantInitialUSD float64 + wantInitialPlan string + wantInitialApplied map[string]int + }{ + {name: "balance read rejected", fault: &testBillingFault{operation: billingOpBalance}, wantErr: "failed to fetch trial credit balance", wantInitialCredits: 80, wantInitialPlan: "trial", wantInitialApplied: map[string]int{}}, + {name: "USD grant rejected", fault: &testBillingFault{operation: billingOpGrant}, wantErr: "failed to grant USD rollover credits", wantInitialCredits: 80, wantInitialPlan: "trial", wantInitialApplied: map[string]int{}}, + {name: "USD grant response lost", fault: &testBillingFault{operation: billingOpGrant, afterApply: true}, wantErr: "failed to grant USD rollover credits", wantInitialCredits: 80, wantInitialUSD: 80, wantInitialPlan: "trial", wantInitialApplied: map[string]int{billingOpGrant: 1}}, + {name: "credit debit rejected", fault: &testBillingFault{operation: billingOpDebit}, wantErr: "failed to debit trial credits", wantInitialCredits: 80, wantInitialUSD: 80, wantInitialPlan: "trial", wantInitialApplied: map[string]int{billingOpGrant: 1}}, + {name: "credit debit response lost", fault: &testBillingFault{operation: billingOpDebit, afterApply: true}, wantErr: "failed to debit trial credits", wantInitialUSD: 80, wantInitialPlan: "trial", wantInitialApplied: map[string]int{billingOpGrant: 1, billingOpDebit: 1}}, + {name: "plan change rejected", fault: &testBillingFault{operation: billingOpPlan}, wantErr: "injected billing failure", wantInitialUSD: 80, wantInitialPlan: "trial", wantInitialApplied: map[string]int{billingOpGrant: 1, billingOpDebit: 1}}, + {name: "plan change response lost", fault: &testBillingFault{operation: billingOpPlan, afterApply: true}, wantErr: "injected billing failure", wantInitialUSD: 80, wantInitialPlan: "paid", wantInitialApplied: map[string]int{billingOpGrant: 1, billingOpDebit: 1, billingOpPlan: 1}}, + {name: "quota database update rejected", failDatabaseUpdate: true, wantErr: "injected organization update failure", wantInitialUSD: 80, wantInitialPlan: "paid", wantInitialApplied: map[string]int{billingOpGrant: 1, billingOpDebit: 1, billingOpPlan: 1}}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Each provider or database boundary must resume without duplicating money and eventually agree on the paid plan and quotas. + org, biller := newBillingScenario(t, fix, owner) + biller.setFault(tt.fault) + fix.Admin.Biller = biller + fix.Admin.PaymentProvider = newTestPaymentProvider(true, true) + + var originalDB database.DB + if tt.failDatabaseUpdate { + originalDB = fix.Admin.DB + fix.Admin.DB = &failOrganizationUpdateDB{DB: originalDB, remaining: 1} + defer func() { fix.Admin.DB = originalDB }() + } + + _, err := owner.UpdateBillingSubscription(t.Context(), &adminv1.UpdateBillingSubscriptionRequest{Org: org.Name, PlanName: "paid"}) + require.Error(t, err) + require.ErrorContains(t, err, tt.wantErr) + require.Equal(t, tt.wantInitialCredits, biller.balance(billing.CreditsCurrency)) + require.Equal(t, tt.wantInitialUSD, biller.balance(billing.USDCurrency)) + require.Equal(t, tt.wantInitialPlan, biller.activePlanName()) + for _, operation := range []string{billingOpGrant, billingOpDebit, billingOpPlan} { + require.Equal(t, tt.wantInitialApplied[operation], biller.appliedCount(operation), "unexpected provider effect before retry") + } + stored, err := fix.Admin.DB.FindOrganization(t.Context(), org.ID) + require.NoError(t, err) + require.Equal(t, "trial", derefString(stored.BillingPlanName), "the DB cache must not claim success before the handler completes") + + resp, err := owner.UpdateBillingSubscription(t.Context(), &adminv1.UpdateBillingSubscriptionRequest{Org: org.Name, PlanName: "paid"}) + require.NoError(t, err) + require.Equal(t, "paid", resp.Subscription.Plan.Name) + require.Equal(t, float64(80), biller.balance(billing.USDCurrency)) + require.Zero(t, biller.balance(billing.CreditsCurrency)) + require.Equal(t, float64(80), biller.totalValue(), "retry must conserve the original trial-credit value") + require.Equal(t, 1, biller.appliedCount(billingOpGrant), "the logical USD grant must apply at most once") + require.Equal(t, 1, biller.appliedCount(billingOpDebit), "the logical trial debit must apply at most once") + require.Equal(t, 1, biller.appliedCount(billingOpPlan), "the logical plan change must apply at most once") + require.NotEmpty(t, biller.appliedKey(billingOpGrant), "financial retries require a stable provider idempotency key") + require.NotEmpty(t, biller.appliedKey(billingOpDebit), "financial retries require a stable provider idempotency key") + + stored, err = fix.Admin.DB.FindOrganization(t.Context(), org.ID) + require.NoError(t, err) + require.Equal(t, "paid", derefString(stored.BillingPlanName)) + require.Equal(t, 7, stored.QuotaProjects) + }) + } + }) +} + +type testBillingFault struct { + operation string + afterApply bool +} + +type testBiller struct { + billing.Biller + + mu sync.Mutex + plans map[string]*billing.Plan + active *billing.Subscription + balances map[string]float64 + fault *testBillingFault + attempts map[string]int + applied map[string]int + appliedKeys map[string]string + seenKeys map[string]struct{} +} + +func newTestBiller(org *database.Organization) *testBiller { + trialProjects := 1 + paidProjects := 7 + trial := &billing.Plan{ID: "plan-trial", Name: "trial", DisplayName: "Trial", PlanType: billing.FreePlanType, Public: true, Quotas: billing.Quotas{NumProjects: &trialProjects}} + paid := &billing.Plan{ID: "plan-paid", Name: "paid", DisplayName: "Paid", PlanType: billing.TeamPlanType, Public: true, Quotas: billing.Quotas{NumProjects: &paidProjects}} + private := &billing.Plan{ID: "plan-private", Name: "private", DisplayName: "Private", PlanType: billing.FreePlanType, Public: false, Quotas: billing.Quotas{NumProjects: &paidProjects}} + return &testBiller{ + Biller: billing.NewNoop(), + plans: map[string]*billing.Plan{ + trial.Name: trial, + paid.Name: paid, + private.Name: private, + }, + active: &billing.Subscription{ + ID: "subscription-" + org.ID, Customer: &billing.Customer{ID: org.BillingCustomerID}, Plan: clonePlan(trial), + CurrentBillingCycleEndDate: time.Now().Add(30 * 24 * time.Hour), + }, + balances: map[string]float64{billing.CreditsCurrency: 80, billing.USDCurrency: 0}, + attempts: make(map[string]int), + applied: make(map[string]int), + appliedKeys: make(map[string]string), + seenKeys: make(map[string]struct{}), + } +} + +func (b *testBiller) GetPlanByName(_ context.Context, name string) (*billing.Plan, error) { + b.mu.Lock() + defer b.mu.Unlock() + plan := b.plans[name] + if plan == nil { + return nil, billing.ErrNotFound + } + return clonePlan(plan), nil +} + +func (b *testBiller) GetPlanByType(_ context.Context, planType billing.PlanType) (*billing.Plan, error) { + b.mu.Lock() + defer b.mu.Unlock() + for _, plan := range b.plans { + if plan.PlanType == planType { + return clonePlan(plan), nil + } + } + return nil, billing.ErrNotFound +} + +func (b *testBiller) GetActiveSubscription(_ context.Context, _ string) (*billing.Subscription, error) { + b.mu.Lock() + defer b.mu.Unlock() + if b.active == nil { + return nil, billing.ErrNotFound + } + return cloneSubscription(b.active), nil +} + +func (b *testBiller) GetCustomerCreditBalance(_ context.Context, _ string, currency string) (float64, error) { + b.mu.Lock() + defer b.mu.Unlock() + b.attempts[billingOpBalance]++ + if b.consumeFault(billingOpBalance, false) { + return 0, errInjectedBilling + } + return b.balances[currency], nil +} + +func (b *testBiller) GrantCustomerCredits(_ context.Context, _ string, amount float64, currency, _ string, _ *time.Time, idempotencyKey string) error { + b.mu.Lock() + defer b.mu.Unlock() + b.attempts[billingOpGrant]++ + if b.alreadyApplied(billingOpGrant, idempotencyKey) { + return nil + } + if b.consumeFault(billingOpGrant, false) { + return errInjectedBilling + } + b.balances[currency] += amount + b.markApplied(billingOpGrant, idempotencyKey) + if b.consumeFault(billingOpGrant, true) { + return errInjectedBilling + } + return nil +} + +func (b *testBiller) DebitCustomerCredits(_ context.Context, _ string, amount float64, currency, _, idempotencyKey string) error { + b.mu.Lock() + defer b.mu.Unlock() + b.attempts[billingOpDebit]++ + if b.alreadyApplied(billingOpDebit, idempotencyKey) { + return nil + } + if b.consumeFault(billingOpDebit, false) { + return errInjectedBilling + } + b.balances[currency] -= amount + b.markApplied(billingOpDebit, idempotencyKey) + if b.consumeFault(billingOpDebit, true) { + return errInjectedBilling + } + return nil +} + +func (b *testBiller) ChangeSubscriptionPlan(_ context.Context, _ string, plan *billing.Plan, idempotencyKey string) (*billing.Subscription, error) { + b.mu.Lock() + defer b.mu.Unlock() + b.attempts[billingOpPlan]++ + if b.alreadyApplied(billingOpPlan, idempotencyKey) { + return cloneSubscription(b.active), nil + } + if b.consumeFault(billingOpPlan, false) { + return nil, errInjectedBilling + } + b.active.Plan = clonePlan(plan) + b.markApplied(billingOpPlan, idempotencyKey) + if b.consumeFault(billingOpPlan, true) { + return nil, errInjectedBilling + } + return cloneSubscription(b.active), nil +} + +func (b *testBiller) CreateSubscription(_ context.Context, _ string, plan *billing.Plan) (*billing.Subscription, error) { + b.mu.Lock() + defer b.mu.Unlock() + b.attempts[billingOpCreate]++ + b.applied[billingOpCreate]++ + b.active = &billing.Subscription{ID: "created-subscription", Customer: &billing.Customer{}, Plan: clonePlan(plan)} + return cloneSubscription(b.active), nil +} + +func (b *testBiller) setFault(fault *testBillingFault) { + b.mu.Lock() + defer b.mu.Unlock() + if fault == nil { + b.fault = nil + return + } + copy := *fault + b.fault = © +} + +func (b *testBiller) prepareUnsubscribedCustomer() { + b.mu.Lock() + defer b.mu.Unlock() + b.active = nil + b.balances[billing.CreditsCurrency] = 0 +} + +func (b *testBiller) consumeFault(operation string, afterApply bool) bool { + if b.fault == nil || b.fault.operation != operation || b.fault.afterApply != afterApply { + return false + } + b.fault = nil + return true +} + +func (b *testBiller) alreadyApplied(operation, key string) bool { + if key == "" { + return false + } + _, ok := b.seenKeys[operation+"\x00"+key] + return ok +} + +func (b *testBiller) markApplied(operation, key string) { + b.applied[operation]++ + b.appliedKeys[operation] = key + if key != "" { + b.seenKeys[operation+"\x00"+key] = struct{}{} + } +} + +func (b *testBiller) mutationAttempts() int { + b.mu.Lock() + defer b.mu.Unlock() + return b.attempts[billingOpGrant] + b.attempts[billingOpDebit] + b.attempts[billingOpPlan] + b.attempts[billingOpCreate] +} + +func (b *testBiller) balance(currency string) float64 { + b.mu.Lock() + defer b.mu.Unlock() + return b.balances[currency] +} + +func (b *testBiller) totalValue() float64 { + b.mu.Lock() + defer b.mu.Unlock() + return b.balances[billing.CreditsCurrency] + b.balances[billing.USDCurrency] +} + +func (b *testBiller) activePlanName() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.active.Plan.Name +} + +func (b *testBiller) appliedCount(operation string) int { + b.mu.Lock() + defer b.mu.Unlock() + return b.applied[operation] +} + +func (b *testBiller) appliedKey(operation string) string { + b.mu.Lock() + defer b.mu.Unlock() + return b.appliedKeys[operation] +} + +type testPaymentProvider struct { + payment.Provider + + mu sync.Mutex + customer payment.Customer + find int +} + +func newTestPaymentProvider(hasPaymentMethod, hasAddress bool) *testPaymentProvider { + return &testPaymentProvider{ + Provider: payment.NewNoop(), + customer: payment.Customer{ID: "payment-customer", HasPaymentMethod: hasPaymentMethod, HasBillableAddress: hasAddress}, + } +} + +func (p *testPaymentProvider) FindCustomer(_ context.Context, _ string) (*payment.Customer, error) { + p.mu.Lock() + defer p.mu.Unlock() + p.find++ + copy := p.customer + return ©, nil +} + +func (p *testPaymentProvider) findCalls() int { + p.mu.Lock() + defer p.mu.Unlock() + return p.find +} + +type failOrganizationUpdateDB struct { + database.DB + + mu sync.Mutex + remaining int +} + +func (db *failOrganizationUpdateDB) UpdateOrganization(ctx context.Context, id string, opts *database.UpdateOrganizationOptions) (*database.Organization, error) { + db.mu.Lock() + if db.remaining > 0 { + db.remaining-- + db.mu.Unlock() + return nil, errors.New("injected organization update failure") + } + db.mu.Unlock() + return db.DB.UpdateOrganization(ctx, id, opts) +} + +func newBillingScenario(t *testing.T, fix *testadmin.Fixture, owner *client.Client) (*database.Organization, *testBiller) { + t.Helper() + resp, err := owner.CreateOrganization(t.Context(), &adminv1.CreateOrganizationRequest{Name: randomName()}) + require.NoError(t, err) + org, err := fix.Admin.DB.FindOrganizationByName(t.Context(), resp.Organization.Name) + require.NoError(t, err) + trialName := "trial" + trialDisplayName := "Trial" + org, err = fix.Admin.DB.UpdateOrganization(t.Context(), org.ID, &database.UpdateOrganizationOptions{ + Name: org.Name, DisplayName: org.DisplayName, Description: org.Description, LogoAssetID: org.LogoAssetID, LogoDarkAssetID: org.LogoDarkAssetID, + FaviconAssetID: org.FaviconAssetID, ThumbnailAssetID: org.ThumbnailAssetID, CustomDomain: org.CustomDomain, DefaultProjectRoleID: org.DefaultProjectRoleID, + QuotaProjects: org.QuotaProjects, QuotaDeployments: org.QuotaDeployments, QuotaSlotsTotal: org.QuotaSlotsTotal, QuotaSlotsPerDeployment: org.QuotaSlotsPerDeployment, + QuotaOutstandingInvites: org.QuotaOutstandingInvites, QuotaStorageLimitBytesPerDeployment: org.QuotaStorageLimitBytesPerDeployment, QuotaSeats: org.QuotaSeats, + BillingCustomerID: "billing-" + org.ID, PaymentCustomerID: "payment-" + org.ID, BillingEmail: org.BillingEmail, + BillingPlanName: &trialName, BillingPlanDisplayName: &trialDisplayName, CreatedByUserID: org.CreatedByUserID, + }) + require.NoError(t, err) + return org, newTestBiller(org) +} + +func clonePlan(plan *billing.Plan) *billing.Plan { + if plan == nil { + return nil + } + copy := *plan + return © +} + +func cloneSubscription(sub *billing.Subscription) *billing.Subscription { + if sub == nil { + return nil + } + copy := *sub + copy.Plan = clonePlan(sub.Plan) + if sub.Customer != nil { + customer := *sub.Customer + copy.Customer = &customer + } + return © +} + +func derefString(value *string) string { + if value == nil { + return "" + } + return *value +} + +var _ billing.Biller = (*testBiller)(nil) +var _ payment.Provider = (*testPaymentProvider)(nil) +var _ database.DB = (*failOrganizationUpdateDB)(nil)