mirror of
https://github.com/labring/sealos.git
synced 2026-09-21 05:44:43 +08:00
fix(account): make hourly billing reconciliation resumable and idempotent (#7126)
* fix(account): make hourly billing reconciliation resumable and idempotent - schedule every ready billing hour from a persisted Mongo checkpoint - retry failed owner reconciliation and preserve checkpoint progress - use stable billing IDs with Mongo upsert semantics - recover stable unsettled billings independently from monitor data - make Cockroach balance and credits deductions idempotent - reconstruct historical debt and subscription state using transaction UpdatedAt - scope subscription history queries to active workspaces in the target hour - add focused unit and Testcontainers runtime coverage * perf * fix(account): restore hourly billing deduction semantics * feat(account): bound billing checkpoint catch-up window - add BILLING_MAX_CATCHUP_DURATION with a default of 24h - limit historical billing replay to the configured duration - preserve first-start behavior to process only the latest ready hour - log skipped hours when the persisted checkpoint exceeds the replay window - expose durable checkpoint lag and windowed pending checkpoint metrics - wire the setting into the account controller Helm chart - add bounded catch-up, first-start, validation, and metric tests * fix: ci * perf(account): optimize indexed historical billing queries - replace full DebtStatusRecord loading with per-user LATERAL lookup - use half-open billing time boundaries for historical debt state - add composite indexes for debt, subscription, and credits queries - add Testcontainers runtime and execution-plan coverage for billing paths - verify Mongo billing indexes and repeated initialization behavior
This commit is contained in:
@@ -1065,44 +1065,52 @@ func (c *Cockroach) AddDeductionBalanceWithCredits(
|
||||
deductionAmount int64,
|
||||
orderIDs []string,
|
||||
) error {
|
||||
return c.AddDeductionBalanceWithCreditsAt(
|
||||
ops, deductionAmount, orderIDs, time.Now().UTC(),
|
||||
)
|
||||
}
|
||||
|
||||
func (c *Cockroach) AddDeductionBalanceWithCreditsAt(
|
||||
ops *types.UserQueryOpts,
|
||||
deductionAmount int64,
|
||||
_ []string,
|
||||
at time.Time,
|
||||
) error {
|
||||
if at.IsZero() {
|
||||
at = time.Now().UTC()
|
||||
}
|
||||
err := RetryTransaction(3, 2*time.Second, c.DB, func(tx *gorm.DB) error {
|
||||
remainingAmount := deductionAmount
|
||||
userUID, dErr := c.GetUserUID(ops)
|
||||
if dErr != nil {
|
||||
return fmt.Errorf("failed to get user uid: %w", dErr)
|
||||
}
|
||||
var credits []types.Credits
|
||||
if dErr = c.DB.Where("user_uid = ? AND expire_at > ? AND status = ?", userUID, time.Now().UTC(), types.CreditsStatusActive).Order("expire_at ASC").Find(&credits).Error; dErr != nil {
|
||||
if dErr = tx.Where(
|
||||
"user_uid = ? AND start_at <= ? AND expire_at > ? AND status = ?",
|
||||
userUID,
|
||||
at,
|
||||
at,
|
||||
types.CreditsStatusActive,
|
||||
).Order("expire_at ASC").Find(&credits).Error; dErr != nil {
|
||||
return fmt.Errorf("failed to get credits: %w", dErr)
|
||||
}
|
||||
now := time.Now().UTC()
|
||||
accountTransactionID := uuid.New()
|
||||
accountTransaction := types.AccountTransaction{
|
||||
ID: accountTransactionID,
|
||||
RegionUID: c.LocalRegion.UID,
|
||||
Type: "RESOURCE_BILLING",
|
||||
UserUID: userUID,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
BillingIDList: orderIDs,
|
||||
}
|
||||
var updateCredits []types.Credits
|
||||
var updateCreditsIDs []string
|
||||
// var creditTransactions []types.CreditsTransaction
|
||||
var creditUsedAmountAll int64
|
||||
for i := range credits {
|
||||
creditAmt := credits[i].Amount - credits[i].UsedAmount
|
||||
if creditAmt > 0 && deductionAmount > 0 {
|
||||
if creditAmt > 0 && remainingAmount > 0 {
|
||||
var usedAmount int64
|
||||
if creditAmt > deductionAmount {
|
||||
credits[i].UsedAmount += deductionAmount
|
||||
usedAmount = deductionAmount
|
||||
if creditAmt > remainingAmount {
|
||||
credits[i].UsedAmount += remainingAmount
|
||||
usedAmount = remainingAmount
|
||||
} else {
|
||||
credits[i].UsedAmount = credits[i].Amount
|
||||
credits[i].Status = types.CreditsStatusUsedUp
|
||||
usedAmount = creditAmt
|
||||
}
|
||||
creditUsedAmountAll += usedAmount
|
||||
deductionAmount -= usedAmount
|
||||
remainingAmount -= usedAmount
|
||||
// creditTransactions = append(creditTransactions, types.CreditsTransaction{
|
||||
// ID: uuid.New(),
|
||||
// UserUID: userUID,
|
||||
@@ -1115,7 +1123,6 @@ func (c *Cockroach) AddDeductionBalanceWithCredits(
|
||||
// })
|
||||
credits[i].UpdatedAt = now
|
||||
updateCredits = append(updateCredits, credits[i])
|
||||
updateCreditsIDs = append(updateCreditsIDs, credits[i].ID.String())
|
||||
}
|
||||
}
|
||||
if len(updateCredits) > 0 {
|
||||
@@ -1124,20 +1131,12 @@ func (c *Cockroach) AddDeductionBalanceWithCredits(
|
||||
return fmt.Errorf("failed to update credits: %w", dErr)
|
||||
}
|
||||
}
|
||||
accountTransaction.DeductionCredit = creditUsedAmountAll
|
||||
accountTransaction.CreditIDList = updateCreditsIDs
|
||||
}
|
||||
if deductionAmount > 0 {
|
||||
if dErr = c.updateBalance(tx, ops, deductionAmount, true, true); dErr != nil {
|
||||
if remainingAmount > 0 {
|
||||
if dErr = c.updateBalance(tx, ops, remainingAmount, true, true); dErr != nil {
|
||||
return fmt.Errorf("failed to update balance: %w", dErr)
|
||||
}
|
||||
accountTransaction.DeductionBalance = deductionAmount
|
||||
} else {
|
||||
accountTransaction.DeductionBalance = 0
|
||||
}
|
||||
// if dErr = tx.Create(&accountTransaction).Error; dErr != nil {
|
||||
// return fmt.Errorf("failed to create account transaction: %v", dErr)
|
||||
//}
|
||||
// if len(creditTransactions) > 0 {
|
||||
// if dErr = tx.Create(&creditTransactions).Error; dErr != nil {
|
||||
// return fmt.Errorf("failed to create credit transactions: %v", dErr)
|
||||
@@ -2257,6 +2256,9 @@ func (c *Cockroach) InitTables() error {
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create index on WorkspaceSubscriptionTransaction: %w", err)
|
||||
}
|
||||
if err := ensureBillingQueryIndexes(c.DB); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// TODO: remove this after migration
|
||||
if err := c.migrateColumns(); err != nil {
|
||||
@@ -2265,6 +2267,38 @@ func (c *Cockroach) InitTables() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// ensureBillingQueryIndexes covers the billing queries that operate on
|
||||
// growing historical tables. Equality columns lead each index so a billing
|
||||
// request can avoid scanning rows belonging to other users or workspaces.
|
||||
func ensureBillingQueryIndexes(db *gorm.DB) error {
|
||||
statements := []struct {
|
||||
name string
|
||||
sql string
|
||||
}{
|
||||
{
|
||||
name: "DebtStatusRecord first status",
|
||||
sql: `CREATE INDEX IF NOT EXISTS idx_debt_record_user_time
|
||||
ON "DebtStatusRecord" (user_uid, create_at, id, last_status);`,
|
||||
},
|
||||
{
|
||||
name: "WorkspaceSubscriptionTransaction history",
|
||||
sql: `CREATE INDEX IF NOT EXISTS idx_workspace_subscription_billing_history
|
||||
ON "WorkspaceSubscriptionTransaction" (region_domain, workspace, status, updated_at);`,
|
||||
},
|
||||
{
|
||||
name: "Credits active period",
|
||||
sql: `CREATE INDEX IF NOT EXISTS idx_credits_active_period
|
||||
ON "Credits" (user_uid, status, expire_at, start_at);`,
|
||||
},
|
||||
}
|
||||
for _, statement := range statements {
|
||||
if err := db.Exec(statement.sql).Error; err != nil {
|
||||
return fmt.Errorf("failed to create billing history index on %s: %w", statement.name, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Cockroach) migratorPaymentRefundTable() error {
|
||||
// If the id field does not exist or exists but is not the primary key and is not not null, skip the migration
|
||||
if !c.DB.Migrator().HasConstraint(&types.PaymentRefund{}, "PaymentRefund_pkey") {
|
||||
|
||||
@@ -0,0 +1,337 @@
|
||||
package cockroach
|
||||
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/labring/sealos/controllers/pkg/types"
|
||||
"github.com/testcontainers/testcontainers-go"
|
||||
postgrescontainer "github.com/testcontainers/testcontainers-go/modules/postgres"
|
||||
"github.com/testcontainers/testcontainers-go/wait"
|
||||
"gorm.io/driver/postgres"
|
||||
"gorm.io/gorm"
|
||||
)
|
||||
|
||||
func TestBillingDeductionWithPostgresRuntimeDoesNotCreateTransaction(t *testing.T) {
|
||||
testcontainers.SkipIfProviderIsNotHealthy(t)
|
||||
ctx := context.Background()
|
||||
container, err := postgrescontainer.Run(ctx, "postgres:16-alpine",
|
||||
postgrescontainer.WithDatabase("account"),
|
||||
postgrescontainer.WithUsername("account"),
|
||||
postgrescontainer.WithPassword("account"),
|
||||
testcontainers.WithWaitStrategy(
|
||||
wait.ForLog("database system is ready to accept connections").WithOccurrence(2),
|
||||
),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
if err := container.Terminate(ctx); err != nil {
|
||||
t.Errorf("terminate PostgreSQL: %v", err)
|
||||
}
|
||||
})
|
||||
dsn, err := container.ConnectionString(ctx, "sslmode=disable")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.AutoMigrate(&types.Account{}, &types.AccountTransaction{}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
userUID := uuid.New()
|
||||
if err := db.Create(&types.Account{
|
||||
UserUID: userUID, CreateRegionID: "test", Balance: 1000,
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
regionUID := uuid.New()
|
||||
account := &Cockroach{
|
||||
DB: db, Localdb: db, LocalRegion: &types.Region{UID: regionUID, Domain: "test"},
|
||||
ownerUsrUIDMap: &sync.Map{}, ownerUsrIDMap: &sync.Map{},
|
||||
}
|
||||
for range 2 {
|
||||
if err := account.AddDeductionBalance(
|
||||
&types.UserQueryOpts{UID: userUID}, 125,
|
||||
); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
var stored types.Account
|
||||
if err := db.First(&stored, `"userUid" = ?`, userUID).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if stored.DeductionBalance != 250 {
|
||||
t.Fatalf("deduction balance = %d", stored.DeductionBalance)
|
||||
}
|
||||
var transactionCount int64
|
||||
if err := db.Model(&types.AccountTransaction{}).Count(&transactionCount).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if transactionCount != 0 {
|
||||
t.Fatalf("transaction count = %d", transactionCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHistoricalBillingDeductionUsesCreditsActiveAtBillingTime(t *testing.T) {
|
||||
testcontainers.SkipIfProviderIsNotHealthy(t)
|
||||
ctx := context.Background()
|
||||
container, err := postgrescontainer.Run(ctx, "postgres:16-alpine",
|
||||
postgrescontainer.WithDatabase("account"),
|
||||
postgrescontainer.WithUsername("account"),
|
||||
postgrescontainer.WithPassword("account"),
|
||||
testcontainers.WithWaitStrategy(
|
||||
wait.ForLog("database system is ready to accept connections").WithOccurrence(2),
|
||||
),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
if err := container.Terminate(ctx); err != nil {
|
||||
t.Errorf("terminate PostgreSQL: %v", err)
|
||||
}
|
||||
})
|
||||
dsn, err := container.ConnectionString(ctx, "sslmode=disable")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.AutoMigrate(
|
||||
&types.Account{}, &types.AccountTransaction{}, &types.Credits{},
|
||||
); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
userUID := uuid.New()
|
||||
if err := db.Create(&types.Account{
|
||||
UserUID: userUID, CreateRegionID: "test", Balance: 1000,
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
creditID := uuid.New()
|
||||
billingTime := time.Now().UTC().Add(-2 * time.Hour)
|
||||
if err := db.Create(&types.Credits{
|
||||
ID: creditID, UserUID: userUID, Amount: 100, Status: types.CreditsStatusActive,
|
||||
StartAt: billingTime.Add(-time.Hour), ExpireAt: billingTime.Add(time.Hour),
|
||||
}).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
regionUID := uuid.New()
|
||||
account := &Cockroach{
|
||||
DB: db, Localdb: db, LocalRegion: &types.Region{UID: regionUID, Domain: "test"},
|
||||
ownerUsrUIDMap: &sync.Map{}, ownerUsrIDMap: &sync.Map{},
|
||||
}
|
||||
if err := account.AddDeductionBalanceWithCreditsAt(
|
||||
&types.UserQueryOpts{UID: userUID},
|
||||
125,
|
||||
[]string{"historical-credit-order"},
|
||||
billingTime,
|
||||
); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
var storedCredit types.Credits
|
||||
if err := db.First(&storedCredit, "id = ?", creditID).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if storedCredit.UsedAmount != 100 {
|
||||
t.Fatalf("used credits = %d", storedCredit.UsedAmount)
|
||||
}
|
||||
var storedAccount types.Account
|
||||
if err := db.First(&storedAccount, `"userUid" = ?`, userUID).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if storedAccount.DeductionBalance != 25 {
|
||||
t.Fatalf("deduction balance = %d", storedAccount.DeductionBalance)
|
||||
}
|
||||
var transactionCount int64
|
||||
if err := db.Model(&types.AccountTransaction{}).Count(&transactionCount).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if transactionCount != 0 {
|
||||
t.Fatalf("transaction count = %d", transactionCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBillingHistoryQueriesUseCompositeIndexes(t *testing.T) {
|
||||
testcontainers.SkipIfProviderIsNotHealthy(t)
|
||||
ctx := context.Background()
|
||||
container, err := postgrescontainer.Run(ctx, "postgres:16-alpine",
|
||||
postgrescontainer.WithDatabase("account"),
|
||||
postgrescontainer.WithUsername("account"),
|
||||
postgrescontainer.WithPassword("account"),
|
||||
testcontainers.WithWaitStrategy(
|
||||
wait.ForLog("database system is ready to accept connections").WithOccurrence(2),
|
||||
),
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
if err := container.Terminate(ctx); err != nil {
|
||||
t.Errorf("terminate PostgreSQL: %v", err)
|
||||
}
|
||||
})
|
||||
dsn, err := container.ConnectionString(ctx, "sslmode=disable")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
db, err := gorm.Open(postgres.Open(dsn), &gorm.Config{})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Exec(`
|
||||
CREATE TABLE "Debt" (
|
||||
user_uid TEXT PRIMARY KEY,
|
||||
account_debt_status TEXT NOT NULL,
|
||||
created_at TIMESTAMPTZ NOT NULL
|
||||
);
|
||||
CREATE TABLE "DebtStatusRecord" (
|
||||
id BIGINT PRIMARY KEY,
|
||||
user_uid TEXT NOT NULL,
|
||||
last_status TEXT NOT NULL,
|
||||
create_at TIMESTAMPTZ NOT NULL
|
||||
);
|
||||
CREATE TABLE "WorkspaceSubscriptionTransaction" (
|
||||
id BIGINT PRIMARY KEY,
|
||||
region_domain TEXT NOT NULL,
|
||||
workspace TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
pay_status TEXT NOT NULL,
|
||||
operator TEXT NOT NULL,
|
||||
updated_at TIMESTAMPTZ NOT NULL
|
||||
);
|
||||
CREATE TABLE "Credits" (
|
||||
id BIGINT PRIMARY KEY,
|
||||
user_uid TEXT NOT NULL,
|
||||
amount BIGINT NOT NULL,
|
||||
used_amount BIGINT NOT NULL,
|
||||
expire_at TIMESTAMPTZ NOT NULL,
|
||||
start_at TIMESTAMPTZ NOT NULL,
|
||||
status TEXT NOT NULL
|
||||
);`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := ensureBillingQueryIndexes(db); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Exec(`
|
||||
INSERT INTO "Debt" (user_uid, account_debt_status, created_at)
|
||||
SELECT CASE WHEN i = 1 THEN 'target-user' ELSE 'user-' || i::TEXT END,
|
||||
'NormalPeriod',
|
||||
now() - INTERVAL '1 day'
|
||||
FROM generate_series(1, 1000) AS series(i);`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Exec(`
|
||||
INSERT INTO "DebtStatusRecord" (id, user_uid, last_status, create_at)
|
||||
SELECT i,
|
||||
CASE WHEN i % 1000 = 0 THEN 'target-user' ELSE 'user-' || i::TEXT END,
|
||||
'NormalPeriod',
|
||||
now() - (i || ' seconds')::INTERVAL
|
||||
FROM generate_series(1, 100000) AS series(i);`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Exec(`
|
||||
INSERT INTO "WorkspaceSubscriptionTransaction" (
|
||||
id, region_domain, workspace, status, pay_status, operator, updated_at
|
||||
)
|
||||
SELECT i,
|
||||
'test-region',
|
||||
CASE WHEN i % 1000 = 0 THEN 'target-workspace' ELSE 'workspace-' || i::TEXT END,
|
||||
'completed',
|
||||
CASE WHEN i % 2 = 0 THEN 'paid' ELSE 'no_need' END,
|
||||
CASE WHEN i % 1000 = 0 THEN 'canceled' WHEN i % 2 = 0 THEN 'created' ELSE 'canceled' END,
|
||||
now() - (i || ' seconds')::INTERVAL
|
||||
FROM generate_series(1, 100000) AS series(i);`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := db.Exec(`
|
||||
INSERT INTO "Credits" (id, user_uid, amount, used_amount, expire_at, start_at, status)
|
||||
SELECT i,
|
||||
CASE WHEN i % 1000 = 0 THEN 'target-user' ELSE 'user-' || i::TEXT END,
|
||||
100,
|
||||
0,
|
||||
now() + INTERVAL '1 day',
|
||||
now() - INTERVAL '1 day',
|
||||
'active'
|
||||
FROM generate_series(1, 100000) AS series(i);`).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
plans := map[string]string{
|
||||
"debt first status": `EXPLAIN (ANALYZE, COSTS OFF)
|
||||
SELECT d.user_uid,
|
||||
COALESCE(first_record.last_status, d.account_debt_status) AS account_debt_status
|
||||
FROM "Debt" AS d
|
||||
LEFT JOIN LATERAL (
|
||||
SELECT r.last_status
|
||||
FROM "DebtStatusRecord" AS r
|
||||
WHERE r.user_uid = d.user_uid
|
||||
AND r.create_at > now() - INTERVAL '2 hours'
|
||||
ORDER BY r.create_at ASC, r.id ASC
|
||||
LIMIT 1
|
||||
) AS first_record ON TRUE
|
||||
WHERE d.created_at <= now()`,
|
||||
"subscription period": `EXPLAIN (ANALYZE, COSTS OFF)
|
||||
SELECT * FROM "WorkspaceSubscriptionTransaction"
|
||||
WHERE region_domain = 'test-region'
|
||||
AND workspace IN ('target-workspace')
|
||||
AND status = 'completed'
|
||||
AND pay_status IN ('paid', 'no_need')
|
||||
AND updated_at <= now()`,
|
||||
"subscription terminal": `EXPLAIN (ANALYZE, COSTS OFF)
|
||||
SELECT * FROM "WorkspaceSubscriptionTransaction"
|
||||
WHERE region_domain = 'test-region'
|
||||
AND workspace IN ('target-workspace')
|
||||
AND status = 'completed'
|
||||
AND operator IN ('canceled', 'deleted')
|
||||
AND updated_at <= now()`,
|
||||
"credits active period": `EXPLAIN (ANALYZE, COSTS OFF)
|
||||
SELECT * FROM "Credits"
|
||||
WHERE user_uid = 'target-user'
|
||||
AND start_at <= now()
|
||||
AND expire_at > now()
|
||||
AND status = 'active'
|
||||
ORDER BY expire_at ASC`,
|
||||
}
|
||||
wantIndexes := map[string]string{
|
||||
"debt first status": "idx_debt_record_user_time",
|
||||
"subscription period": "idx_workspace_subscription_billing_history",
|
||||
"subscription terminal": "idx_workspace_subscription_billing_history",
|
||||
"credits active period": "idx_credits_active_period",
|
||||
}
|
||||
for name, query := range plans {
|
||||
rows, err := db.Raw(query).Rows()
|
||||
if err != nil {
|
||||
t.Fatalf("%s explain: %v", name, err)
|
||||
}
|
||||
var plan strings.Builder
|
||||
for rows.Next() {
|
||||
var line string
|
||||
if err := rows.Scan(&line); err != nil {
|
||||
rows.Close()
|
||||
t.Fatalf("%s scan explain: %v", name, err)
|
||||
}
|
||||
plan.WriteString(line)
|
||||
}
|
||||
if err := rows.Close(); err != nil {
|
||||
t.Fatalf("%s close explain: %v", name, err)
|
||||
}
|
||||
planText := plan.String()
|
||||
t.Logf("%s: %s", name, planText)
|
||||
if !strings.Contains(planText, wantIndexes[name]) {
|
||||
t.Fatalf("%s did not use %s: %s", name, wantIndexes[name], planText)
|
||||
}
|
||||
}
|
||||
t.Logf("indexed historical billing queries completed against 100000-row tables")
|
||||
}
|
||||
@@ -47,6 +47,13 @@ type CVM interface {
|
||||
type Account interface {
|
||||
GetBillingLastUpdateTime(owner string, _type common.Type) (bool, time.Time, error)
|
||||
GetOwnersRecentUpdates(ownerList []string, checkTime time.Time) ([]string, error)
|
||||
GetOwnerBillingsAt(
|
||||
ownerList []string,
|
||||
billingTime time.Time,
|
||||
) (map[string][]*resources.Billing, error)
|
||||
GetUnsettledBillingsAt(billingTime time.Time) (map[string][]*resources.Billing, error)
|
||||
GetBillingCheckpoint() (time.Time, bool, error)
|
||||
SaveBillingCheckpoint(billingTime time.Time) error
|
||||
GetTimeUsedNamespaceList(startTime, endTime time.Time) ([]string, error)
|
||||
SaveBillings(billing ...*resources.Billing) error
|
||||
SaveObjTraffic(obs ...*types.ObjectStorageTraffic) error
|
||||
@@ -141,6 +148,12 @@ type AccountV2 interface {
|
||||
) error
|
||||
AddBalance(user *types.UserQueryOpts, balance int64) error
|
||||
AddDeductionBalanceWithCredits(ops *types.UserQueryOpts, amount int64, orderIDs []string) error
|
||||
AddDeductionBalanceWithCreditsAt(
|
||||
ops *types.UserQueryOpts,
|
||||
amount int64,
|
||||
orderIDs []string,
|
||||
at time.Time,
|
||||
) error
|
||||
ReduceBalance(ops *types.UserQueryOpts, amount int64) error
|
||||
ReduceDeductionBalance(ops *types.UserQueryOpts, amount int64) error
|
||||
NewAccount(user *types.UserQueryOpts) (*types.Account, error)
|
||||
|
||||
@@ -16,11 +16,11 @@ package mongo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"errors"
|
||||
"fmt"
|
||||
"math"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -30,7 +30,6 @@ import (
|
||||
"github.com/labring/sealos/controllers/pkg/types"
|
||||
"github.com/labring/sealos/controllers/pkg/utils/env"
|
||||
"github.com/labring/sealos/controllers/pkg/utils/logger"
|
||||
gonanoid "github.com/matoous/go-nanoid/v2"
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
"go.mongodb.org/mongo-driver/mongo"
|
||||
@@ -47,18 +46,19 @@ const (
|
||||
)
|
||||
|
||||
const (
|
||||
DefaultAccountDBName = "sealos-resources"
|
||||
DefaultTrafficDBName = "sealos-networkmanager"
|
||||
DefaultAuthDBName = "sealos-auth"
|
||||
DefaultCVMDBName = "sealos-cvm"
|
||||
DefaultCVMConn = "cvm"
|
||||
DefaultMeteringConn = "metering"
|
||||
DefaultMonitorConn = "monitor"
|
||||
DefaultBillingConn = "billing"
|
||||
DefaultObjTrafficConn = "objectstorage-traffic"
|
||||
DefaultUserConn = "user"
|
||||
DefaultPricesConn = "prices"
|
||||
DefaultPropertiesConn = "properties"
|
||||
DefaultAccountDBName = "sealos-resources"
|
||||
DefaultTrafficDBName = "sealos-networkmanager"
|
||||
DefaultAuthDBName = "sealos-auth"
|
||||
DefaultCVMDBName = "sealos-cvm"
|
||||
DefaultCVMConn = "cvm"
|
||||
DefaultMeteringConn = "metering"
|
||||
DefaultMonitorConn = "monitor"
|
||||
DefaultBillingConn = "billing"
|
||||
DefaultObjTrafficConn = "objectstorage-traffic"
|
||||
DefaultUserConn = "user"
|
||||
DefaultPricesConn = "prices"
|
||||
DefaultPropertiesConn = "properties"
|
||||
DefaultBillingCheckpointConn = "billing-checkpoint"
|
||||
// TODO fix
|
||||
DefaultTrafficConn = "traffic"
|
||||
)
|
||||
@@ -216,6 +216,103 @@ func (m *mongoDB) GetOwnersRecentUpdates(
|
||||
return updatedOwners, nil
|
||||
}
|
||||
|
||||
func (m *mongoDB) GetOwnerBillingsAt(
|
||||
ownerList []string,
|
||||
billingTime time.Time,
|
||||
) (map[string][]*resources.Billing, error) {
|
||||
result := make(map[string][]*resources.Billing)
|
||||
if len(ownerList) == 0 {
|
||||
return result, nil
|
||||
}
|
||||
filter := bson.M{
|
||||
"owner": ownerListFilter(ownerList),
|
||||
"time": billingTime,
|
||||
"type": common.Consumption,
|
||||
"app_type": bson.M{"$nin": []int{
|
||||
int(resources.AppType[resources.CVM]),
|
||||
int(resources.AppType[resources.LLMToken]),
|
||||
}},
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
|
||||
defer cancel()
|
||||
cursor, err := m.getBillingCollection().Find(ctx, filter)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("find owner billings: %w", err)
|
||||
}
|
||||
defer cursor.Close(ctx)
|
||||
var billings []*resources.Billing
|
||||
if err := cursor.All(ctx, &billings); err != nil {
|
||||
return nil, fmt.Errorf("decode owner billings: %w", err)
|
||||
}
|
||||
for _, billing := range billings {
|
||||
result[billing.Owner] = append(result[billing.Owner], billing)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func (m *mongoDB) GetUnsettledBillingsAt(
|
||||
billingTime time.Time,
|
||||
) (map[string][]*resources.Billing, error) {
|
||||
filter := bson.M{
|
||||
"time": billingTime,
|
||||
"type": common.Consumption,
|
||||
"status": resources.Unsettled,
|
||||
"order_id": primitive.Regex{Pattern: "^bh_"},
|
||||
"app_type": bson.M{"$nin": []int{
|
||||
int(resources.AppType[resources.CVM]),
|
||||
int(resources.AppType[resources.LLMToken]),
|
||||
}},
|
||||
}
|
||||
cursor, err := m.getBillingCollection().Find(context.Background(), filter)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("find unsettled billings: %w", err)
|
||||
}
|
||||
defer cursor.Close(context.Background())
|
||||
var billings []*resources.Billing
|
||||
if err := cursor.All(context.Background(), &billings); err != nil {
|
||||
return nil, fmt.Errorf("decode unsettled billings: %w", err)
|
||||
}
|
||||
result := make(map[string][]*resources.Billing)
|
||||
for _, billing := range billings {
|
||||
result[billing.Owner] = append(result[billing.Owner], billing)
|
||||
}
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func ownerListFilter(ownerList []string) bson.M {
|
||||
return bson.M{"$in": ownerList}
|
||||
}
|
||||
|
||||
const billingCheckpointID = "account-hourly-billing"
|
||||
|
||||
func (m *mongoDB) GetBillingCheckpoint() (time.Time, bool, error) {
|
||||
var checkpoint resources.BillingCheckpoint
|
||||
err := m.getBillingCheckpointCollection().FindOne(
|
||||
context.Background(),
|
||||
bson.M{"_id": billingCheckpointID},
|
||||
).Decode(&checkpoint)
|
||||
if errors.Is(err, mongo.ErrNoDocuments) {
|
||||
return time.Time{}, false, nil
|
||||
}
|
||||
if err != nil {
|
||||
return time.Time{}, false, fmt.Errorf("get billing checkpoint: %w", err)
|
||||
}
|
||||
return checkpoint.Time.UTC(), true, nil
|
||||
}
|
||||
|
||||
func (m *mongoDB) SaveBillingCheckpoint(billingTime time.Time) error {
|
||||
_, err := m.getBillingCheckpointCollection().UpdateOne(
|
||||
context.Background(),
|
||||
bson.M{"_id": billingCheckpointID},
|
||||
bson.M{"$set": bson.M{"time": billingTime.UTC(), "updated_at": time.Now().UTC()}},
|
||||
options.Update().SetUpsert(true),
|
||||
)
|
||||
if err != nil {
|
||||
return fmt.Errorf("save billing checkpoint: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (m *mongoDB) GetTimeUsedNamespaceList(startTime, endTime time.Time) ([]string, error) {
|
||||
pipeline := mongo.Pipeline{
|
||||
{
|
||||
@@ -308,11 +405,19 @@ func (m *mongoDB) UpdateBillingStatus(orderIDs []string, status resources.Billin
|
||||
}
|
||||
|
||||
func (m *mongoDB) SaveBillings(billing ...*resources.Billing) error {
|
||||
billings := make([]any, len(billing))
|
||||
for i, b := range billing {
|
||||
billings[i] = b
|
||||
if len(billing) == 0 {
|
||||
return nil
|
||||
}
|
||||
_, err := m.getBillingCollection().InsertMany(context.Background(), billings)
|
||||
models := make([]mongo.WriteModel, 0, len(billing))
|
||||
for _, b := range billing {
|
||||
models = append(models, mongo.NewUpdateOneModel().
|
||||
SetFilter(bson.M{"owner": b.Owner, "order_id": b.OrderID}).
|
||||
SetUpdate(bson.M{"$setOnInsert": b}).
|
||||
SetUpsert(true))
|
||||
}
|
||||
_, err := m.getBillingCollection().BulkWrite(
|
||||
context.Background(), models, options.BulkWrite().SetOrdered(false),
|
||||
)
|
||||
return err
|
||||
}
|
||||
|
||||
@@ -709,15 +814,15 @@ func (m *mongoDB) FetchOwnerMonitorRecords(
|
||||
startTime, endTime time.Time,
|
||||
ownerToNS map[string][]string,
|
||||
) (map[string][]resources.Monitor, error) {
|
||||
// collect all namespaces to avoid repetition
|
||||
nsSet := make(map[string]struct{})
|
||||
for _, nsList := range ownerToNS {
|
||||
// Build the reverse index once so each monitor record has an O(1) owner lookup.
|
||||
nsToOwner := make(map[string]string)
|
||||
for owner, nsList := range ownerToNS {
|
||||
for _, ns := range nsList {
|
||||
nsSet[ns] = struct{}{}
|
||||
nsToOwner[ns] = owner
|
||||
}
|
||||
}
|
||||
namespaces := make([]string, 0, len(nsSet))
|
||||
for ns := range nsSet {
|
||||
namespaces := make([]string, 0, len(nsToOwner))
|
||||
for ns := range nsToOwner {
|
||||
namespaces = append(namespaces, ns)
|
||||
}
|
||||
|
||||
@@ -740,21 +845,20 @@ func (m *mongoDB) FetchOwnerMonitorRecords(
|
||||
return nil, fmt.Errorf("failed to decode monitor records: %w", err)
|
||||
}
|
||||
|
||||
// build the mapping of owner monitor data
|
||||
return groupMonitorRecordsByOwner(allRecords, nsToOwner), nil
|
||||
}
|
||||
|
||||
func groupMonitorRecordsByOwner(
|
||||
records []resources.Monitor,
|
||||
nsToOwner map[string]string,
|
||||
) map[string][]resources.Monitor {
|
||||
ownerMonitorRecords := make(map[string][]resources.Monitor)
|
||||
for _, record := range allRecords {
|
||||
for owner, nsList := range ownerToNS {
|
||||
// Only the records of the namespace that belong to the owner are saved
|
||||
for _, ns := range nsList {
|
||||
if record.Category == ns {
|
||||
ownerMonitorRecords[owner] = append(ownerMonitorRecords[owner], record)
|
||||
break // avoid duplicate additions
|
||||
}
|
||||
}
|
||||
for _, record := range records {
|
||||
if owner, ok := nsToOwner[record.Category]; ok {
|
||||
ownerMonitorRecords[owner] = append(ownerMonitorRecords[owner], record)
|
||||
}
|
||||
}
|
||||
|
||||
return ownerMonitorRecords, nil
|
||||
return ownerMonitorRecords
|
||||
}
|
||||
|
||||
func GenerateBillingDataFromRecords(
|
||||
@@ -803,9 +907,12 @@ func GenerateBillingDataFromRecords(
|
||||
}
|
||||
|
||||
// 存储最终计费数据
|
||||
// map[namespace]map[app_type | parent_type/parent_name][]resources.AppCost
|
||||
appCostsMap := make(map[string]map[string][]resources.AppCost)
|
||||
nsTypeAmount := make(map[string]map[string]int64)
|
||||
type billingGroupKey struct {
|
||||
appType uint8
|
||||
appName string
|
||||
}
|
||||
appCostsMap := make(map[string]map[billingGroupKey][]resources.AppCost)
|
||||
nsTypeAmount := make(map[string]map[billingGroupKey]int64)
|
||||
|
||||
calculateFinalUsed := func(values map[uint8][]int64, prols *resources.PropertyTypeLS, minutes float64) map[uint8]int64 {
|
||||
finalUsed := make(map[uint8]int64)
|
||||
@@ -844,18 +951,18 @@ func GenerateBillingDataFromRecords(
|
||||
continue
|
||||
}
|
||||
appCost.Amount = totalAmount
|
||||
groupKey := strconv.Itoa(int(agg.Type))
|
||||
groupKey := billingGroupKey{appType: agg.Type}
|
||||
if agg.ParentType != 0 && agg.ParentName != "" {
|
||||
groupKey = strconv.Itoa(int(agg.ParentType)) + "/" + agg.ParentName
|
||||
groupKey = billingGroupKey{appType: agg.ParentType, appName: agg.ParentName}
|
||||
}
|
||||
ns := agg.Category
|
||||
if _, ok := nsTypeAmount[ns]; !ok {
|
||||
nsTypeAmount[ns] = make(map[string]int64)
|
||||
nsTypeAmount[ns] = make(map[billingGroupKey]int64)
|
||||
}
|
||||
nsTypeAmount[ns][groupKey] += totalAmount
|
||||
|
||||
if _, ok := appCostsMap[ns]; !ok {
|
||||
appCostsMap[ns] = make(map[string][]resources.AppCost)
|
||||
appCostsMap[ns] = make(map[billingGroupKey][]resources.AppCost)
|
||||
}
|
||||
appCostsMap[ns][groupKey] = append(appCostsMap[ns][groupKey], appCost)
|
||||
}
|
||||
@@ -863,27 +970,19 @@ func GenerateBillingDataFromRecords(
|
||||
|
||||
// 生成 Billing 数据
|
||||
for ns, appCostMap := range appCostsMap {
|
||||
for tp, appCostList := range appCostMap {
|
||||
amount := nsTypeAmount[ns][tp]
|
||||
for groupKey, appCostList := range appCostMap {
|
||||
amount := nsTypeAmount[ns][groupKey]
|
||||
if amount <= 0 {
|
||||
continue
|
||||
}
|
||||
id, err := gonanoid.New(12)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("generate billing id error: %w", err)
|
||||
}
|
||||
parts := strings.Split(tp, "/")
|
||||
appType, _ := strconv.Atoi(parts[0])
|
||||
appName := ""
|
||||
if len(parts) > 1 {
|
||||
appName = parts[1]
|
||||
}
|
||||
billings = append(billings, &resources.Billing{
|
||||
OrderID: id,
|
||||
OrderID: stableBillingOrderID(
|
||||
owner, endTime, ns, groupKey.appType, groupKey.appName,
|
||||
),
|
||||
Type: Consumption,
|
||||
Namespace: ns,
|
||||
AppType: uint8(appType), // #nosec G115
|
||||
AppName: appName,
|
||||
AppType: groupKey.appType,
|
||||
AppName: groupKey.appName,
|
||||
AppCosts: appCostList,
|
||||
Amount: amount,
|
||||
Owner: owner,
|
||||
@@ -895,6 +994,25 @@ func GenerateBillingDataFromRecords(
|
||||
return billings, nil
|
||||
}
|
||||
|
||||
func stableBillingOrderID(
|
||||
owner string,
|
||||
endTime time.Time,
|
||||
namespace string,
|
||||
appType uint8,
|
||||
appName string,
|
||||
) string {
|
||||
key := fmt.Sprintf(
|
||||
"%s\x00%s\x00%s\x00%d\x00%s",
|
||||
owner,
|
||||
endTime.UTC().Format(time.RFC3339),
|
||||
namespace,
|
||||
appType,
|
||||
appName,
|
||||
)
|
||||
sum := sha256.Sum256([]byte(key))
|
||||
return fmt.Sprintf("bh_%x", sum[:12])
|
||||
}
|
||||
|
||||
func computeUsedValue(usedValues []int64, prop resources.PropertyType, minutes float64) int64 {
|
||||
switch prop.PriceType {
|
||||
case resources.DIF:
|
||||
@@ -1009,6 +1127,10 @@ func (m *mongoDB) getBillingCollection() *mongo.Collection {
|
||||
return m.Client.Database(m.AccountDB).Collection(m.BillingConn)
|
||||
}
|
||||
|
||||
func (m *mongoDB) getBillingCheckpointCollection() *mongo.Collection {
|
||||
return m.Client.Database(m.AccountDB).Collection(DefaultBillingCheckpointConn)
|
||||
}
|
||||
|
||||
func (m *mongoDB) getObjTrafficCollection() *mongo.Collection {
|
||||
return m.Client.Database(m.AccountDB).Collection(m.ObjTrafficConn)
|
||||
}
|
||||
@@ -1018,13 +1140,15 @@ func (m *mongoDB) getPropertiesCollection() *mongo.Collection {
|
||||
}
|
||||
|
||||
func (m *mongoDB) CreateBillingIfNotExist() error {
|
||||
if exist, err := m.collectionExist(m.AccountDB, m.BillingConn); exist || err != nil {
|
||||
ctx := context.Background()
|
||||
exist, err := m.collectionExist(m.AccountDB, m.BillingConn)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
ctx := context.Background()
|
||||
err := m.Client.Database(m.AccountDB).CreateCollection(ctx, m.BillingConn)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create collection for billing: %w", err)
|
||||
if !exist {
|
||||
if err := m.Client.Database(m.AccountDB).CreateCollection(ctx, m.BillingConn); err != nil {
|
||||
return fmt.Errorf("failed to create collection for billing: %w", err)
|
||||
}
|
||||
}
|
||||
|
||||
// create index
|
||||
@@ -1045,6 +1169,15 @@ func (m *mongoDB) CreateBillingIfNotExist() error {
|
||||
primitive.E{Key: "type", Value: 1},
|
||||
},
|
||||
},
|
||||
{
|
||||
// recover stable unsettled billings for one billing hour
|
||||
Keys: bson.D{
|
||||
primitive.E{Key: "time", Value: 1},
|
||||
primitive.E{Key: "status", Value: 1},
|
||||
primitive.E{Key: "type", Value: 1},
|
||||
primitive.E{Key: "order_id", Value: 1},
|
||||
},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to create index for billing: %w", err)
|
||||
|
||||
@@ -30,6 +30,41 @@ import (
|
||||
|
||||
var testTime = time.Date(2023, time.May, 9, 5, 0, 0, 0, time.UTC)
|
||||
|
||||
func TestGenerateBillingDataPreservesTypedGroupKey(t *testing.T) {
|
||||
start := time.Date(2026, time.July, 29, 1, 0, 0, 0, time.UTC)
|
||||
end := start.Add(time.Hour)
|
||||
properties := resources.NewPropertyTypeLS([]resources.PropertyType{
|
||||
{
|
||||
Name: "cpu", Enum: 0, PriceType: resources.AVG, UnitPrice: 1,
|
||||
},
|
||||
})
|
||||
records := []resources.Monitor{
|
||||
{
|
||||
Time: start, Category: "ns-owner", Type: 1,
|
||||
ParentType: 255, ParentName: "parent/name", Name: "child",
|
||||
Used: resources.EnumUsedMap{0: 60},
|
||||
},
|
||||
}
|
||||
|
||||
billings, err := GenerateBillingDataFromRecords(
|
||||
records, properties, start, end, "owner",
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(billings) != 1 {
|
||||
t.Fatalf("billing count = %d, want 1", len(billings))
|
||||
}
|
||||
if billings[0].AppType != 255 || billings[0].AppName != "parent/name" {
|
||||
t.Fatalf(
|
||||
"billing group = (%d, %q), want (255, %q)",
|
||||
billings[0].AppType,
|
||||
billings[0].AppName,
|
||||
"parent/name",
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMongoDB_SaveBillingsWithAccountBalance(t *testing.T) {
|
||||
type fields struct {
|
||||
URL string
|
||||
|
||||
@@ -0,0 +1,296 @@
|
||||
package mongo
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/labring/sealos/controllers/pkg/resources"
|
||||
"github.com/testcontainers/testcontainers-go"
|
||||
"github.com/testcontainers/testcontainers-go/wait"
|
||||
"go.mongodb.org/mongo-driver/bson"
|
||||
"go.mongodb.org/mongo-driver/bson/primitive"
|
||||
)
|
||||
|
||||
func explainBillingFind(ctx context.Context, account *mongoDB, filter bson.D) (bson.M, error) {
|
||||
var explain bson.M
|
||||
err := account.Client.Database(account.AccountDB).RunCommand(ctx, bson.D{
|
||||
{Key: "explain", Value: bson.D{
|
||||
{Key: "find", Value: account.BillingConn},
|
||||
{Key: "filter", Value: filter},
|
||||
}},
|
||||
{Key: "verbosity", Value: "executionStats"},
|
||||
}).Decode(&explain)
|
||||
return explain, err
|
||||
}
|
||||
|
||||
func TestBillingPersistenceWithMongoRuntime(t *testing.T) {
|
||||
testcontainers.SkipIfProviderIsNotHealthy(t)
|
||||
ctx := context.Background()
|
||||
container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
|
||||
ContainerRequest: testcontainers.ContainerRequest{
|
||||
Image: "mongo:7.0",
|
||||
ExposedPorts: []string{"27017/tcp"},
|
||||
WaitingFor: wait.ForListeningPort("27017/tcp").
|
||||
WithStartupTimeout(2 * time.Minute),
|
||||
},
|
||||
Started: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
if err := container.Terminate(ctx); err != nil {
|
||||
t.Errorf("terminate MongoDB: %v", err)
|
||||
}
|
||||
})
|
||||
host, err := container.Host(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
port, err := container.MappedPort(ctx, "27017/tcp")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
account, err := NewMongoInterface(ctx, "mongodb://"+net.JoinHostPort(host, port.Port()))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
if err := account.Disconnect(ctx); err != nil {
|
||||
t.Errorf("disconnect MongoDB: %v", err)
|
||||
}
|
||||
})
|
||||
mongoAccount, ok := account.(*mongoDB)
|
||||
if !ok {
|
||||
t.Fatalf("account type = %T", account)
|
||||
}
|
||||
// Simulate an existing production collection so initialization must add
|
||||
// indexes during an upgrade.
|
||||
if err := mongoAccount.Client.Database(mongoAccount.AccountDB).
|
||||
CreateCollection(ctx, mongoAccount.BillingConn); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := account.CreateBillingIfNotExist(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := account.CreateBillingIfNotExist(); err != nil {
|
||||
t.Fatalf("repeat index initialization: %v", err)
|
||||
}
|
||||
|
||||
end := time.Date(2026, time.July, 7, 10, 0, 0, 0, time.UTC)
|
||||
billing := &resources.Billing{
|
||||
Time: end, OrderID: "stable-order", Owner: "owner", Namespace: "ns-owner",
|
||||
Type: Consumption, AppType: 1, Amount: 100, Status: resources.Unsettled,
|
||||
}
|
||||
if err := account.SaveBillings(billing); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
billing.Status = resources.Settled
|
||||
if err := account.SaveBillings(billing); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
indexSpecs, err := mongoAccount.getBillingCollection().Indexes().ListSpecifications(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(indexSpecs) != 4 {
|
||||
t.Fatalf("billing index count = %d, want 4", len(indexSpecs))
|
||||
}
|
||||
monitorTime := end.Add(-time.Hour)
|
||||
namespaces, err := account.GetTimeUsedNamespaceList(monitorTime, end)
|
||||
if err != nil || len(namespaces) != 0 {
|
||||
t.Fatalf("missing monitor collection namespaces=%v err=%v", namespaces, err)
|
||||
}
|
||||
count, err := mongoAccount.getBillingCollection().CountDocuments(ctx, bson.M{
|
||||
"owner": "owner", "order_id": "stable-order",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if count != 1 {
|
||||
t.Fatalf("billing count = %d", count)
|
||||
}
|
||||
var stored resources.Billing
|
||||
if err := mongoAccount.getBillingCollection().FindOne(ctx, bson.M{
|
||||
"owner": "owner", "order_id": "stable-order",
|
||||
}).Decode(&stored); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if stored.Status != resources.Unsettled {
|
||||
t.Fatalf("stored status = %v", stored.Status)
|
||||
}
|
||||
recoveryBillings := []*resources.Billing{
|
||||
{
|
||||
Time: end, OrderID: "bh_recover", Owner: "recover-owner", Namespace: "ns-recover",
|
||||
Type: Consumption, AppType: 1, Amount: 50, Status: resources.Unsettled,
|
||||
},
|
||||
{
|
||||
Time: end, OrderID: "bh_settled", Owner: "settled-owner", Namespace: "ns-settled",
|
||||
Type: Consumption, AppType: 1, Amount: 50, Status: resources.Settled,
|
||||
},
|
||||
{
|
||||
Time: end, OrderID: "legacy-random", Owner: "legacy-owner", Namespace: "ns-legacy",
|
||||
Type: Consumption, AppType: 1, Amount: 50, Status: resources.Unsettled,
|
||||
},
|
||||
{
|
||||
Time: end.Add(time.Hour),
|
||||
OrderID: "bh_other-hour",
|
||||
Owner: "other-owner",
|
||||
Namespace: "ns-other",
|
||||
Type: Consumption, AppType: 1, Amount: 50, Status: resources.Unsettled,
|
||||
},
|
||||
}
|
||||
if err := account.SaveBillings(recoveryBillings...); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
unsettled, err := account.GetUnsettledBillingsAt(end)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(unsettled) != 1 || len(unsettled["recover-owner"]) != 1 ||
|
||||
unsettled["recover-owner"][0].OrderID != "bh_recover" {
|
||||
t.Fatalf("unsettled billings = %#v", unsettled)
|
||||
}
|
||||
|
||||
if err := account.SaveBillingCheckpoint(end); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
checkpoint, exists, err := account.GetBillingCheckpoint()
|
||||
if err != nil || !exists || !checkpoint.Equal(end) {
|
||||
t.Fatalf("checkpoint=%v exists=%v err=%v", checkpoint, exists, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBillingQueriesUseIndexesWithMongoRuntime(t *testing.T) {
|
||||
testcontainers.SkipIfProviderIsNotHealthy(t)
|
||||
ctx := context.Background()
|
||||
container, err := testcontainers.GenericContainer(ctx, testcontainers.GenericContainerRequest{
|
||||
ContainerRequest: testcontainers.ContainerRequest{
|
||||
Image: "mongo:7.0",
|
||||
ExposedPorts: []string{"27017/tcp"},
|
||||
WaitingFor: wait.ForListeningPort("27017/tcp").
|
||||
WithStartupTimeout(2 * time.Minute),
|
||||
},
|
||||
Started: true,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
if err := container.Terminate(ctx); err != nil {
|
||||
t.Errorf("terminate MongoDB: %v", err)
|
||||
}
|
||||
})
|
||||
host, err := container.Host(ctx)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
port, err := container.MappedPort(ctx, "27017/tcp")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
account, err := NewMongoInterface(ctx, "mongodb://"+net.JoinHostPort(host, port.Port()))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() {
|
||||
if err := account.Disconnect(ctx); err != nil {
|
||||
t.Errorf("disconnect MongoDB: %v", err)
|
||||
}
|
||||
})
|
||||
if err := account.CreateBillingIfNotExist(); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
mongoAccount, ok := account.(*mongoDB)
|
||||
if !ok {
|
||||
t.Fatalf("account type = %T", account)
|
||||
}
|
||||
end := time.Date(2026, time.July, 7, 10, 0, 0, 0, time.UTC)
|
||||
documents := make([]any, 0, 50000)
|
||||
for i := 0; i < 50000; i++ {
|
||||
owner := "owner-other"
|
||||
billingTime := end.Add(time.Duration(i%24) * time.Hour)
|
||||
status := resources.Settled
|
||||
if i%1000 == 0 {
|
||||
owner = "owner-target"
|
||||
billingTime = end
|
||||
status = resources.Unsettled
|
||||
}
|
||||
documents = append(documents, &resources.Billing{
|
||||
Time: billingTime,
|
||||
OrderID: fmt.Sprintf("bh_%05d", i),
|
||||
Type: Consumption,
|
||||
Namespace: "ns-owner",
|
||||
AppType: 1,
|
||||
Amount: 1,
|
||||
Owner: owner,
|
||||
Status: status,
|
||||
})
|
||||
}
|
||||
if _, err := mongoAccount.getBillingCollection().InsertMany(ctx, documents); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
ownerFilter := bson.D{
|
||||
{Key: "owner", Value: bson.M{"$in": []string{"owner-target"}}},
|
||||
{Key: "time", Value: end},
|
||||
{Key: "type", Value: Consumption},
|
||||
{Key: "app_type", Value: bson.M{"$nin": []int{int(resources.AppType[resources.CVM]), int(resources.AppType[resources.LLMToken])}}},
|
||||
}
|
||||
unsettledFilter := bson.D{
|
||||
{Key: "time", Value: end},
|
||||
{Key: "status", Value: resources.Unsettled},
|
||||
{Key: "type", Value: Consumption},
|
||||
{Key: "order_id", Value: primitive.Regex{Pattern: "^bh_"}},
|
||||
{Key: "app_type", Value: bson.M{"$nin": []int{int(resources.AppType[resources.CVM]), int(resources.AppType[resources.LLMToken])}}},
|
||||
}
|
||||
for name, filter := range map[string]bson.D{
|
||||
"owner billing lookup": ownerFilter,
|
||||
"unsettled recovery lookup": unsettledFilter,
|
||||
} {
|
||||
explain, err := explainBillingFind(ctx, mongoAccount, filter)
|
||||
if err != nil {
|
||||
t.Fatalf("%s explain: %v", name, err)
|
||||
}
|
||||
plan := fmt.Sprint(explain["queryPlanner"])
|
||||
if !strings.Contains(plan, "IXSCAN") {
|
||||
t.Fatalf("%s did not use an index: %s", name, plan)
|
||||
}
|
||||
stats := fmt.Sprint(explain["executionStats"])
|
||||
t.Logf("%s: %s", name, stats)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStableBillingOrderID(t *testing.T) {
|
||||
end := time.Date(2026, time.July, 7, 10, 0, 0, 0, time.UTC)
|
||||
first := stableBillingOrderID("owner", end, "ns-owner", 1, "app")
|
||||
second := stableBillingOrderID("owner", end, "ns-owner", 1, "app")
|
||||
if first != second {
|
||||
t.Fatalf("IDs differ: %q %q", first, second)
|
||||
}
|
||||
if first == stableBillingOrderID("owner", end.Add(time.Hour), "ns-owner", 1, "app") {
|
||||
t.Fatal("different billing windows share an ID")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupMonitorRecordsByOwner(t *testing.T) {
|
||||
records := []resources.Monitor{
|
||||
{Category: "ns-a", Name: "a-1"},
|
||||
{Category: "ns-b", Name: "b-1"},
|
||||
{Category: "ns-a", Name: "a-2"},
|
||||
{Category: "unmapped", Name: "ignored"},
|
||||
}
|
||||
grouped := groupMonitorRecordsByOwner(records, map[string]string{
|
||||
"ns-a": "owner-a",
|
||||
"ns-b": "owner-b",
|
||||
})
|
||||
if len(grouped) != 2 || len(grouped["owner-a"]) != 2 || len(grouped["owner-b"]) != 1 {
|
||||
t.Fatalf("grouped records = %#v", grouped)
|
||||
}
|
||||
}
|
||||
@@ -137,6 +137,12 @@ type Billing struct {
|
||||
// UserUID uuid.UUID `json:"user_uid" bson:"user_uid,omitempty"`
|
||||
}
|
||||
|
||||
type BillingCheckpoint struct {
|
||||
ID string `json:"id" bson:"_id"`
|
||||
Time time.Time `json:"time" bson:"time"`
|
||||
UpdatedAt time.Time `json:"updated_at" bson:"updated_at"`
|
||||
}
|
||||
|
||||
type Payment struct {
|
||||
Method string `json:"method" bson:"method"`
|
||||
UserID string `json:"user_id" bson:"user_id"`
|
||||
|
||||
Reference in New Issue
Block a user