mirror of
https://github.com/labring/sealos.git
synced 2026-08-30 17:58:09 +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:
zijiren
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,114 @@
|
||||
package controllers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/go-logr/logr"
|
||||
"github.com/google/uuid"
|
||||
"github.com/labring/sealos/controllers/pkg/database"
|
||||
"github.com/labring/sealos/controllers/pkg/types"
|
||||
"github.com/labring/sealos/controllers/pkg/utils/maps"
|
||||
"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"
|
||||
)
|
||||
|
||||
type debtSnapshotAccountV2 struct {
|
||||
database.AccountV2
|
||||
db *gorm.DB
|
||||
}
|
||||
|
||||
func (a *debtSnapshotAccountV2) GetGlobalDB() *gorm.DB {
|
||||
return a.db
|
||||
}
|
||||
|
||||
func TestLoadDebtUsersAtUsesFirstLaterTransition(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.Debt{}, &types.DebtStatusRecord{}); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
endHourTime := time.Now().UTC().Add(-time.Hour).Truncate(time.Hour)
|
||||
historicalUserUID := uuid.New()
|
||||
currentDebtUserUID := uuid.New()
|
||||
createdLaterUserUID := uuid.New()
|
||||
debts := []types.Debt{
|
||||
{
|
||||
UserUID: historicalUserUID, CreatedAt: endHourTime.Add(-time.Hour),
|
||||
UpdatedAt: endHourTime.Add(20 * time.Minute), AccountDebtStatus: types.NormalPeriod,
|
||||
},
|
||||
{
|
||||
UserUID: currentDebtUserUID, CreatedAt: endHourTime.Add(-time.Hour),
|
||||
UpdatedAt: endHourTime.Add(-time.Hour), AccountDebtStatus: types.DebtPeriod,
|
||||
},
|
||||
{
|
||||
UserUID: createdLaterUserUID, CreatedAt: endHourTime.Add(time.Minute),
|
||||
UpdatedAt: endHourTime.Add(time.Minute), AccountDebtStatus: types.DebtPeriod,
|
||||
},
|
||||
}
|
||||
if err := db.Create(&debts).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
records := []types.DebtStatusRecord{
|
||||
{
|
||||
ID: uuid.New(), UserUID: historicalUserUID,
|
||||
LastStatus: types.NormalPeriod, CurrentStatus: types.DebtPeriod,
|
||||
CreateAt: endHourTime,
|
||||
},
|
||||
{
|
||||
ID: uuid.New(), UserUID: historicalUserUID,
|
||||
LastStatus: types.DebtPeriod, CurrentStatus: types.NormalPeriod,
|
||||
CreateAt: endHourTime.Add(20 * time.Minute),
|
||||
},
|
||||
}
|
||||
if err := db.Create(&records).Error; err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
DebtUserMap = maps.NewConcurrentNullValueMap()
|
||||
reconciler := &BillingReconciler{
|
||||
AccountV2: &debtSnapshotAccountV2{db: db},
|
||||
Logger: logr.Discard(),
|
||||
}
|
||||
if err := reconciler.loadDebtUsersAt(endHourTime); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, inDebt := DebtUserMap.Get(historicalUserUID.String()); inDebt {
|
||||
t.Fatal("historical user did not use the first later transition")
|
||||
}
|
||||
if _, inDebt := DebtUserMap.Get(currentDebtUserUID.String()); !inDebt {
|
||||
t.Fatal("current debt user without a later transition was omitted")
|
||||
}
|
||||
if _, inDebt := DebtUserMap.Get(createdLaterUserUID.String()); inDebt {
|
||||
t.Fatal("debt created after the billing hour was included")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,729 @@
|
||||
package controllers
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/go-logr/logr"
|
||||
"github.com/google/uuid"
|
||||
"github.com/labring/sealos/controllers/pkg/database"
|
||||
"github.com/labring/sealos/controllers/pkg/resources"
|
||||
"github.com/labring/sealos/controllers/pkg/types"
|
||||
"github.com/labring/sealos/controllers/pkg/utils/maps"
|
||||
userv1 "github.com/labring/sealos/controllers/user/api/v1"
|
||||
"github.com/prometheus/client_golang/prometheus/testutil"
|
||||
"k8s.io/apimachinery/pkg/runtime"
|
||||
"sigs.k8s.io/controller-runtime/pkg/client/fake"
|
||||
)
|
||||
|
||||
type billingTestAccount struct {
|
||||
database.Account
|
||||
mu sync.Mutex
|
||||
checkpoint time.Time
|
||||
hasCheckpoint bool
|
||||
checkpointWrites []time.Time
|
||||
monitorErr error
|
||||
existing map[string][]*resources.Billing
|
||||
unsettled map[string][]*resources.Billing
|
||||
generated map[string][]*resources.Billing
|
||||
generateInput map[string][]string
|
||||
saveErr error
|
||||
saved []*resources.Billing
|
||||
statuses map[string]resources.BillingStatus
|
||||
}
|
||||
|
||||
func (f *billingTestAccount) GetBillingCheckpoint() (time.Time, bool, error) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
return f.checkpoint, f.hasCheckpoint, nil
|
||||
}
|
||||
|
||||
func (f *billingTestAccount) SaveBillingCheckpoint(value time.Time) error {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.checkpoint, f.hasCheckpoint = value, true
|
||||
f.checkpointWrites = append(f.checkpointWrites, value)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *billingTestAccount) GetTimeUsedNamespaceList(time.Time, time.Time) ([]string, error) {
|
||||
return nil, f.monitorErr
|
||||
}
|
||||
|
||||
func (f *billingTestAccount) GetOwnerBillingsAt(
|
||||
[]string,
|
||||
time.Time,
|
||||
) (map[string][]*resources.Billing, error) {
|
||||
return f.existing, nil
|
||||
}
|
||||
|
||||
func (f *billingTestAccount) GetUnsettledBillingsAt(
|
||||
time.Time,
|
||||
) (map[string][]*resources.Billing, error) {
|
||||
return f.unsettled, nil
|
||||
}
|
||||
|
||||
func (f *billingTestAccount) GenerateBillingData(
|
||||
_ time.Time,
|
||||
_ time.Time,
|
||||
_ *resources.PropertyTypeLS,
|
||||
ownerListMap map[string][]string,
|
||||
) (map[string][]*resources.Billing, error) {
|
||||
f.mu.Lock()
|
||||
f.generateInput = make(map[string][]string, len(ownerListMap))
|
||||
for owner, namespaces := range ownerListMap {
|
||||
f.generateInput[owner] = append([]string(nil), namespaces...)
|
||||
}
|
||||
f.mu.Unlock()
|
||||
return f.generated, nil
|
||||
}
|
||||
|
||||
func (f *billingTestAccount) SaveBillings(billings ...*resources.Billing) error {
|
||||
if f.saveErr != nil {
|
||||
return f.saveErr
|
||||
}
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.saved = append(f.saved, billings...)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *billingTestAccount) UpdateBillingStatus(
|
||||
orderIDs []string,
|
||||
status resources.BillingStatus,
|
||||
) error {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
if f.statuses == nil {
|
||||
f.statuses = make(map[string]resources.BillingStatus)
|
||||
}
|
||||
for _, orderID := range orderIDs {
|
||||
f.statuses[orderID] = status
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type billingTestAccountV2 struct {
|
||||
database.AccountV2
|
||||
mu sync.Mutex
|
||||
err error
|
||||
deductions int
|
||||
amounts []int64
|
||||
}
|
||||
|
||||
func (f *billingTestAccountV2) AddDeductionBalance(
|
||||
_ *types.UserQueryOpts, amount int64,
|
||||
) error {
|
||||
if f.err != nil {
|
||||
return f.err
|
||||
}
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.deductions++
|
||||
f.amounts = append(f.amounts, amount)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (f *billingTestAccountV2) AddDeductionBalanceWithCreditsAt(
|
||||
_ *types.UserQueryOpts, amount int64, _ []string, _ time.Time,
|
||||
) error {
|
||||
if f.err != nil {
|
||||
return f.err
|
||||
}
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
f.deductions++
|
||||
f.amounts = append(f.amounts, amount)
|
||||
return nil
|
||||
}
|
||||
|
||||
func initBillingTestGlobals() {
|
||||
DebtUserMap = maps.NewConcurrentNullValueMap()
|
||||
SubscriptionWorkspaceMap = maps.NewConcurrentNullValueMap()
|
||||
}
|
||||
|
||||
func TestRunBillingReadsRunsIndependentTasksConcurrently(t *testing.T) {
|
||||
started := make(chan struct{}, 2)
|
||||
release := make(chan struct{})
|
||||
result := make(chan error, 1)
|
||||
wantErr := errors.New("read failed")
|
||||
|
||||
read := func(err error) func() error {
|
||||
return func() error {
|
||||
started <- struct{}{}
|
||||
<-release
|
||||
return err
|
||||
}
|
||||
}
|
||||
go func() {
|
||||
result <- runBillingReads(read(nil), read(wantErr))
|
||||
}()
|
||||
|
||||
for range 2 {
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(time.Second):
|
||||
close(release)
|
||||
t.Fatal("independent read did not start concurrently")
|
||||
}
|
||||
}
|
||||
close(release)
|
||||
if err := <-result; !errors.Is(err, wantErr) {
|
||||
t.Fatalf("error = %v, want %v", err, wantErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBillingMetricsTrackCheckpointLag(t *testing.T) {
|
||||
checkpoint := time.Date(2026, time.July, 7, 9, 0, 0, 0, time.UTC)
|
||||
target := checkpoint.Add(3 * time.Hour)
|
||||
setBillingTargetMetrics(target)
|
||||
setBillingCheckpointMetrics(checkpoint, target)
|
||||
|
||||
wantLag := 3 * float64(time.Hour/time.Second)
|
||||
if got := testutil.ToFloat64(billingCheckpointLagSeconds); got != wantLag {
|
||||
t.Fatalf("checkpoint lag seconds = %v, want %v", got, wantLag)
|
||||
}
|
||||
if got := testutil.ToFloat64(billingPendingCheckpoints); got != 3 {
|
||||
t.Fatalf("pending checkpoints = %v, want 3", got)
|
||||
}
|
||||
if got := testutil.ToFloat64(billingTargetTimestamp); got != float64(target.Unix()) {
|
||||
t.Fatalf("target timestamp = %v, want %v", got, target.Unix())
|
||||
}
|
||||
setBillingProcessingMetrics(target, true)
|
||||
if got := testutil.ToFloat64(billingProcessing); got != 1 {
|
||||
t.Fatalf("processing = %v, want 1", got)
|
||||
}
|
||||
if got := testutil.ToFloat64(billingProcessingStartedTimestamp); got <= 0 {
|
||||
t.Fatalf("processing started timestamp = %v, want positive timestamp", got)
|
||||
}
|
||||
setBillingProcessingMetrics(time.Time{}, false)
|
||||
if got := testutil.ToFloat64(billingProcessing); got != 0 {
|
||||
t.Fatalf("processing = %v, want 0", got)
|
||||
}
|
||||
if got := testutil.ToFloat64(billingProcessingStartedTimestamp); got != 0 {
|
||||
t.Fatalf("processing started timestamp = %v, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBillingHourScheduling(t *testing.T) {
|
||||
t.Run("normal hourly execution", func(t *testing.T) {
|
||||
checkpoint := time.Date(2026, time.July, 7, 9, 0, 0, 0, time.UTC)
|
||||
target := checkpoint.Add(time.Hour)
|
||||
hours := billingHoursAfter(checkpoint, target)
|
||||
if len(hours) != 1 || !hours[0].Equal(target) {
|
||||
t.Fatalf("hours = %v", hours)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("start after half hour", func(t *testing.T) {
|
||||
now := time.Date(2026, time.July, 7, 10, 37, 0, 0, time.UTC)
|
||||
want := time.Date(2026, time.July, 7, 10, 0, 0, 0, time.UTC)
|
||||
if got := latestReadyBillingHour(now); !got.Equal(want) {
|
||||
t.Fatalf("ready hour = %v, want %v", got, want)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("resume after several hours", func(t *testing.T) {
|
||||
start := time.Date(2026, time.July, 7, 7, 0, 0, 0, time.UTC)
|
||||
hours := billingHoursAfter(start, start.Add(4*time.Hour))
|
||||
if len(hours) != 4 {
|
||||
t.Fatalf("hours = %v", hours)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("task crosses next hour", func(t *testing.T) {
|
||||
started := time.Date(2026, time.July, 7, 10, 6, 0, 0, time.UTC)
|
||||
finished := started.Add(61 * time.Minute)
|
||||
if !latestReadyBillingHour(finished).After(latestReadyBillingHour(started)) {
|
||||
t.Fatalf("finished target did not advance")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestOwnerInDebtUsesOwnerMapping(t *testing.T) {
|
||||
initBillingTestGlobals()
|
||||
DebtUserMap.Set("owner-uid")
|
||||
reconciler := &BillingReconciler{
|
||||
debtOwnerMap: maps.NewConcurrentNullValueMap(),
|
||||
}
|
||||
if reconciler.ownerInDebt("owner-uid") {
|
||||
t.Fatal("user UID was treated as an owner")
|
||||
}
|
||||
reconciler.debtOwnerMap.Set("owner")
|
||||
if !reconciler.ownerInDebt("owner") {
|
||||
t.Fatal("debt owner was not detected")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteBillingTasksUntilPersistsEachSuccessfulHour(t *testing.T) {
|
||||
start := time.Date(2026, time.July, 7, 7, 0, 0, 0, time.UTC)
|
||||
db := &billingTestAccount{checkpoint: start, hasCheckpoint: true}
|
||||
var executed []time.Time
|
||||
reconciler := &BillingReconciler{
|
||||
DBClient: db,
|
||||
executeBillingHourFunc: func(hour time.Time) error {
|
||||
executed = append(executed, hour)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
if err := reconciler.ExecuteBillingTasksUntil(start.Add(3 * time.Hour)); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(executed) != 3 || len(db.checkpointWrites) != 3 {
|
||||
t.Fatalf("executed=%v checkpoints=%v", executed, db.checkpointWrites)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteBillingTasksUntilFirstStartProcessesLatestHour(t *testing.T) {
|
||||
t.Setenv(billingMaxCatchupDurationEnv, "72h")
|
||||
target := time.Date(2026, time.July, 7, 10, 0, 0, 0, time.UTC)
|
||||
db := &billingTestAccount{}
|
||||
var executed []time.Time
|
||||
reconciler := &BillingReconciler{
|
||||
DBClient: db,
|
||||
executeBillingHourFunc: func(hour time.Time) error {
|
||||
executed = append(executed, hour)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
if err := reconciler.ExecuteBillingTasksUntil(target); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(executed) != 1 || !executed[0].Equal(target) {
|
||||
t.Fatalf("executed = %v", executed)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteBillingTasksUntilLimitsHistoricalCatchup(t *testing.T) {
|
||||
t.Setenv(billingMaxCatchupDurationEnv, "48h")
|
||||
target := time.Date(2026, time.July, 7, 10, 0, 0, 0, time.UTC)
|
||||
db := &billingTestAccount{checkpoint: target.Add(-72 * time.Hour), hasCheckpoint: true}
|
||||
var executed []time.Time
|
||||
var pendingAtStart, checkpointAtStart float64
|
||||
reconciler := &BillingReconciler{
|
||||
DBClient: db,
|
||||
executeBillingHourFunc: func(hour time.Time) error {
|
||||
if len(executed) == 0 {
|
||||
pendingAtStart = testutil.ToFloat64(billingPendingCheckpoints)
|
||||
checkpointAtStart = testutil.ToFloat64(billingCheckpointTimestamp)
|
||||
}
|
||||
executed = append(executed, hour)
|
||||
return nil
|
||||
},
|
||||
}
|
||||
if err := reconciler.ExecuteBillingTasksUntil(target); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(executed) != 48 {
|
||||
t.Fatalf("executed %d hours, want 48", len(executed))
|
||||
}
|
||||
if pendingAtStart != 48 {
|
||||
t.Fatalf("pending checkpoints at start = %v, want 48", pendingAtStart)
|
||||
}
|
||||
if want := float64(target.Add(-72 * time.Hour).Unix()); checkpointAtStart != want {
|
||||
t.Fatalf("checkpoint metric at start = %v, want %v", checkpointAtStart, want)
|
||||
}
|
||||
if !executed[0].Equal(target.Add(-47 * time.Hour)) {
|
||||
t.Fatalf("first executed hour = %v, want %v", executed[0], target.Add(-47*time.Hour))
|
||||
}
|
||||
if !db.checkpoint.Equal(target) {
|
||||
t.Fatalf("checkpoint = %v, want %v", db.checkpoint, target)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLimitBillingCheckpointValidatesDuration(t *testing.T) {
|
||||
checkpoint := time.Date(2026, time.July, 4, 10, 0, 0, 0, time.UTC)
|
||||
target := time.Date(2026, time.July, 7, 10, 0, 0, 0, time.UTC)
|
||||
limited, err := limitBillingCheckpoint(checkpoint, target, 48*time.Hour)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if want := target.Add(-48 * time.Hour); !limited.Equal(want) {
|
||||
t.Fatalf("limited checkpoint = %v, want %v", limited, want)
|
||||
}
|
||||
if _, err := limitBillingCheckpoint(checkpoint, target, 90*time.Minute); err == nil {
|
||||
t.Fatal("expected non-whole-hour duration error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteBillingTasksUntilStopsAtFailedHour(t *testing.T) {
|
||||
start := time.Date(2026, time.July, 7, 7, 0, 0, 0, time.UTC)
|
||||
db := &billingTestAccount{checkpoint: start, hasCheckpoint: true}
|
||||
reconciler := &BillingReconciler{
|
||||
DBClient: db,
|
||||
executeBillingHourFunc: func(hour time.Time) error {
|
||||
if hour.Equal(start.Add(2 * time.Hour)) {
|
||||
return errors.New("monitor unavailable")
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
if err := reconciler.ExecuteBillingTasksUntil(start.Add(3 * time.Hour)); err == nil {
|
||||
t.Fatal("expected failed hour error")
|
||||
}
|
||||
if !db.checkpoint.Equal(start.Add(time.Hour)) {
|
||||
t.Fatalf("checkpoint = %v", db.checkpoint)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExecuteBillingTasksUntilAdvancesCheckpointWhenMonitorDataIsEmpty(t *testing.T) {
|
||||
start := time.Date(2026, time.July, 7, 7, 0, 0, 0, time.UTC)
|
||||
db := &billingTestAccount{checkpoint: start, hasCheckpoint: true}
|
||||
reconciler := &BillingReconciler{DBClient: db, Logger: logr.Discard()}
|
||||
reconciler.executeBillingHourFunc = func(hour time.Time) error {
|
||||
owners, _, err := reconciler.getRecentUsedOwnersAt(hour)
|
||||
if len(owners) != 0 {
|
||||
t.Fatalf("owners=%v", owners)
|
||||
}
|
||||
return err
|
||||
}
|
||||
target := start.Add(time.Hour)
|
||||
if err := reconciler.ExecuteBillingTasksUntil(target); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !db.checkpoint.Equal(target) || len(db.checkpointWrites) != 1 {
|
||||
t.Fatalf("checkpoint=%v writes=%v", db.checkpoint, db.checkpointWrites)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetRecentUsedOwnersReturnsMonitorError(t *testing.T) {
|
||||
db := &billingTestAccount{monitorErr: errors.New("query failed")}
|
||||
reconciler := &BillingReconciler{DBClient: db}
|
||||
_, _, err := reconciler.getRecentUsedOwnersAt(time.Now())
|
||||
if err == nil {
|
||||
t.Fatal("expected monitor query error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestActiveWorkspaceSubscriptionsAtBillingHour(t *testing.T) {
|
||||
hour := time.Date(2026, time.July, 7, 10, 0, 0, 0, time.UTC)
|
||||
activeStart := hour.Add(-24 * time.Hour)
|
||||
activeEnd := hour.Add(time.Hour)
|
||||
expiredEnd := hour.Add(-time.Minute)
|
||||
subscriptions := []types.WorkspaceSubscription{
|
||||
{
|
||||
Workspace: "active", CreateAt: activeStart,
|
||||
CurrentPeriodStartAt: activeStart, CurrentPeriodEndAt: activeEnd,
|
||||
},
|
||||
{
|
||||
Workspace: "purchased-later", CreateAt: hour.Add(time.Hour),
|
||||
CurrentPeriodStartAt: hour.Add(time.Hour), CurrentPeriodEndAt: hour.Add(48 * time.Hour),
|
||||
},
|
||||
{
|
||||
Workspace: "expired", CreateAt: activeStart,
|
||||
CurrentPeriodStartAt: activeStart, CurrentPeriodEndAt: expiredEnd,
|
||||
},
|
||||
{
|
||||
Workspace: "renewed-later", CreateAt: activeStart, UpdateAt: hour.Add(time.Hour),
|
||||
CurrentPeriodStartAt: activeStart, CurrentPeriodEndAt: hour.Add(48 * time.Hour),
|
||||
},
|
||||
}
|
||||
transactions := []types.WorkspaceSubscriptionTransaction{
|
||||
{
|
||||
ID: uuid.New(),
|
||||
Workspace: "previous-period",
|
||||
Operator: types.SubscriptionTransactionTypeCreated,
|
||||
Status: types.SubscriptionTransactionStatusCompleted,
|
||||
PayStatus: types.SubscriptionPayStatusPaid,
|
||||
StartAt: hour.Add(-12 * time.Hour),
|
||||
UpdatedAt: hour.Add(-11 * time.Hour),
|
||||
Period: types.DayPeriod(1),
|
||||
},
|
||||
{
|
||||
ID: uuid.New(),
|
||||
Workspace: "delayed-activation",
|
||||
Operator: types.SubscriptionTransactionTypeUpgraded,
|
||||
Status: types.SubscriptionTransactionStatusCompleted,
|
||||
PayStatus: types.SubscriptionPayStatusPaid,
|
||||
StartAt: hour.Add(-25 * time.Hour),
|
||||
UpdatedAt: hour.Add(-23 * time.Hour),
|
||||
Period: types.DayPeriod(1),
|
||||
},
|
||||
{
|
||||
ID: uuid.New(),
|
||||
Workspace: "completed-later",
|
||||
Operator: types.SubscriptionTransactionTypeCreated,
|
||||
Status: types.SubscriptionTransactionStatusCompleted,
|
||||
PayStatus: types.SubscriptionPayStatusPaid,
|
||||
StartAt: hour.Add(-12 * time.Hour),
|
||||
UpdatedAt: hour.Add(time.Hour),
|
||||
Period: types.DayPeriod(1),
|
||||
},
|
||||
{
|
||||
ID: uuid.New(),
|
||||
Workspace: "canceled",
|
||||
Operator: types.SubscriptionTransactionTypeCreated,
|
||||
Status: types.SubscriptionTransactionStatusCompleted,
|
||||
PayStatus: types.SubscriptionPayStatusPaid,
|
||||
StartAt: hour.Add(-12 * time.Hour),
|
||||
UpdatedAt: hour.Add(-11 * time.Hour),
|
||||
Period: types.DayPeriod(1),
|
||||
},
|
||||
{
|
||||
ID: uuid.New(),
|
||||
Workspace: "canceled",
|
||||
Operator: types.SubscriptionTransactionTypeCanceled,
|
||||
Status: types.SubscriptionTransactionStatusCompleted,
|
||||
PayStatus: types.SubscriptionPayStatusCanceled,
|
||||
StartAt: hour.Add(-time.Hour),
|
||||
UpdatedAt: hour.Add(-time.Hour),
|
||||
},
|
||||
}
|
||||
workspaces, err := activeWorkspaceSubscriptionsAt(hour, subscriptions, transactions)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
got := make(map[string]struct{}, len(workspaces))
|
||||
for _, workspace := range workspaces {
|
||||
got[workspace] = struct{}{}
|
||||
}
|
||||
if _, exists := got["active"]; !exists {
|
||||
t.Fatal("active current period is missing")
|
||||
}
|
||||
if _, exists := got["previous-period"]; !exists {
|
||||
t.Fatal("active historical transaction period is missing")
|
||||
}
|
||||
if _, exists := got["delayed-activation"]; !exists {
|
||||
t.Fatal("subscription completion time was not used as its period start")
|
||||
}
|
||||
if _, exists := got["purchased-later"]; exists {
|
||||
t.Fatal("future subscription was applied to historical billing")
|
||||
}
|
||||
if _, exists := got["expired"]; exists {
|
||||
t.Fatal("expired subscription was applied to billing")
|
||||
}
|
||||
if _, exists := got["renewed-later"]; exists {
|
||||
t.Fatal("later subscription snapshot was applied to historical billing")
|
||||
}
|
||||
if _, exists := got["completed-later"]; exists {
|
||||
t.Fatal("later transaction completion was applied to historical billing")
|
||||
}
|
||||
if _, exists := got["canceled"]; exists {
|
||||
t.Fatal("canceled subscription was applied to billing")
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetUsersForNamespacesUsesTargetedLookups(t *testing.T) {
|
||||
scheme := runtime.NewScheme()
|
||||
if err := userv1.AddToScheme(scheme); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
userA := &userv1.User{}
|
||||
userA.Name = "a"
|
||||
userA.Annotations = map[string]string{userv1.UserLabelOwnerKey: "owner-a"}
|
||||
userB := &userv1.User{}
|
||||
userB.Name = "b"
|
||||
userB.Annotations = map[string]string{userv1.UserLabelOwnerKey: "owner-b"}
|
||||
reconciler := &BillingReconciler{Client: fake.NewClientBuilder().
|
||||
WithScheme(scheme).
|
||||
WithObjects(userA, userB).
|
||||
Build()}
|
||||
|
||||
owners, err := reconciler.getUsersForNamespaces([]string{
|
||||
"ns-a", "ns-b", "ns-missing", "workspace-without-user-prefix",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(owners) != 2 || owners["ns-a"] != "owner-a" || owners["ns-b"] != "owner-b" {
|
||||
t.Fatalf("owners = %#v", owners)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReconcileOwnerListReturnsPartialOwnerFailure(t *testing.T) {
|
||||
initBillingTestGlobals()
|
||||
end := time.Date(2026, time.July, 7, 10, 0, 0, 0, time.UTC)
|
||||
db := &billingTestAccount{
|
||||
existing: map[string][]*resources.Billing{},
|
||||
generated: map[string][]*resources.Billing{
|
||||
"owner-a": {{OrderID: "a", Owner: "owner-a", Namespace: "ns-a", Time: end, Amount: 1}},
|
||||
"owner-b": {{OrderID: "b", Owner: "owner-b", Namespace: "ns-b", Time: end, Amount: 1}},
|
||||
},
|
||||
}
|
||||
reconciler := &BillingReconciler{
|
||||
DBClient: db,
|
||||
Logger: logr.Discard(),
|
||||
concurrentLimit: 2,
|
||||
reconcileBillingFunc: func(owner string, _ []*resources.Billing, _ time.Time) error {
|
||||
db.mu.Lock()
|
||||
if db.statuses == nil {
|
||||
db.statuses = make(map[string]resources.BillingStatus)
|
||||
}
|
||||
db.statuses[owner] = resources.Settled
|
||||
db.mu.Unlock()
|
||||
if owner == "owner-b" {
|
||||
return errors.New("save failed")
|
||||
}
|
||||
return nil
|
||||
},
|
||||
}
|
||||
err := reconciler.reconcileOwnerList(map[string][]string{
|
||||
"owner-a": {"ns-a"}, "owner-b": {"ns-b"},
|
||||
}, end)
|
||||
if err == nil {
|
||||
t.Fatal("expected partial owner failure")
|
||||
}
|
||||
if len(db.statuses) != 2 {
|
||||
t.Fatalf("reconciled owner count = %d", len(db.statuses))
|
||||
}
|
||||
}
|
||||
|
||||
func TestReconcileOwnerListRejectsNonPositiveConcurrency(t *testing.T) {
|
||||
reconciler := &BillingReconciler{concurrentLimit: 0}
|
||||
if err := reconciler.reconcileOwnerList(nil, time.Now()); err == nil {
|
||||
t.Fatal("expected invalid concurrency error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReconcileBillingSaveFailure(t *testing.T) {
|
||||
initBillingTestGlobals()
|
||||
db := &billingTestAccount{saveErr: errors.New("insert failed")}
|
||||
reconciler := &BillingReconciler{DBClient: db, AccountV2: &billingTestAccountV2{}}
|
||||
err := reconciler.reconcileBilling(
|
||||
"owner", []*resources.Billing{{OrderID: "id", Amount: 10}}, time.Now(),
|
||||
)
|
||||
if err == nil {
|
||||
t.Fatal("expected save failure")
|
||||
}
|
||||
}
|
||||
|
||||
func TestReconcileBillingDeductionFailureMayLeaveSettled(t *testing.T) {
|
||||
initBillingTestGlobals()
|
||||
db := &billingTestAccount{}
|
||||
account := &billingTestAccountV2{err: errors.New("deduction failed")}
|
||||
billing := &resources.Billing{OrderID: "id", Amount: 10, Status: resources.Settled}
|
||||
reconciler := &BillingReconciler{DBClient: db, AccountV2: account}
|
||||
if err := reconciler.reconcileBilling(
|
||||
"owner", []*resources.Billing{billing}, time.Now(),
|
||||
); err == nil {
|
||||
t.Fatal("expected deduction failure")
|
||||
}
|
||||
if len(db.saved) != 1 {
|
||||
t.Fatalf("saved billing count = %d", len(db.saved))
|
||||
}
|
||||
if db.saved[0].Status != resources.Settled {
|
||||
t.Fatalf("saved billing status = %v", db.saved[0].Status)
|
||||
}
|
||||
if status := db.statuses[billing.OrderID]; status != resources.Unsettled {
|
||||
t.Fatalf("failed deduction status = %v", status)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReconcileBillingReclassifiesSubscriptionBeforeDeduction(t *testing.T) {
|
||||
initBillingTestGlobals()
|
||||
SubscriptionWorkspaceMap.Set("ns-subscription")
|
||||
db := &billingTestAccount{}
|
||||
account := &billingTestAccountV2{}
|
||||
billing := &resources.Billing{
|
||||
OrderID: "subscription-order",
|
||||
Namespace: "ns-subscription",
|
||||
Amount: 10,
|
||||
Status: resources.Unsettled,
|
||||
}
|
||||
reconciler := &BillingReconciler{DBClient: db, AccountV2: account}
|
||||
if err := reconciler.reconcileBilling(
|
||||
"owner", []*resources.Billing{billing}, time.Now(),
|
||||
); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if account.deductions != 0 {
|
||||
t.Fatalf("subscription deductions = %d", account.deductions)
|
||||
}
|
||||
if billing.Status != resources.Subscription {
|
||||
t.Fatalf("billing status = %v", billing.Status)
|
||||
}
|
||||
}
|
||||
|
||||
func TestClassifyGeneratedBillingsDefaultsToSettled(t *testing.T) {
|
||||
initBillingTestGlobals()
|
||||
SubscriptionWorkspaceMap.Set("ns-subscription")
|
||||
generated := map[string][]*resources.Billing{"owner": {
|
||||
{OrderID: "new-subscription", Namespace: "ns-subscription", AppName: "new"},
|
||||
{OrderID: "new-usage", Namespace: "ns-usage", AppName: "usage"},
|
||||
}}
|
||||
existing := &resources.Billing{
|
||||
OrderID: "bh_existing", Namespace: "ns-subscription", AppName: "existing",
|
||||
Status: resources.Unsettled,
|
||||
}
|
||||
|
||||
classifyGeneratedBillings(generated)
|
||||
pending := pendingOwnerBillings(generated, map[string][]*resources.Billing{
|
||||
"owner": {existing},
|
||||
})["owner"]
|
||||
|
||||
statuses := make(map[string]resources.BillingStatus, len(pending))
|
||||
for _, billing := range pending {
|
||||
statuses[billing.OrderID] = billing.Status
|
||||
}
|
||||
if statuses["new-subscription"] != resources.Subscription {
|
||||
t.Fatalf("new subscription status = %v", statuses["new-subscription"])
|
||||
}
|
||||
if statuses["new-usage"] != resources.Settled {
|
||||
t.Fatalf("new usage status = %v", statuses["new-usage"])
|
||||
}
|
||||
if _, exists := statuses[existing.OrderID]; exists {
|
||||
t.Fatalf("existing billing was unexpectedly requeued: %v", existing.OrderID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReconcileBillingRetryCanDeductAgain(t *testing.T) {
|
||||
initBillingTestGlobals()
|
||||
db := &billingTestAccount{}
|
||||
account := &billingTestAccountV2{}
|
||||
billings := []*resources.Billing{
|
||||
{OrderID: "stable-a", Amount: 10, Status: resources.Settled},
|
||||
{OrderID: "stable-b", Amount: 15, Status: resources.Settled},
|
||||
}
|
||||
reconciler := &BillingReconciler{DBClient: db, AccountV2: account}
|
||||
for range 2 {
|
||||
if err := reconciler.reconcileBilling("owner", billings, time.Now()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
if account.deductions != 2 {
|
||||
t.Fatalf("deductions = %d", account.deductions)
|
||||
}
|
||||
if len(account.amounts) != 2 || account.amounts[0] != 25 || account.amounts[1] != 25 {
|
||||
t.Fatalf("deduction amounts = %#v", account.amounts)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPendingOwnerBillingsSkipsExistingBusinessKeys(t *testing.T) {
|
||||
generated := map[string][]*resources.Billing{"owner": {
|
||||
{OrderID: "new-a", Namespace: "ns", AppType: 1, AppName: "a"},
|
||||
{OrderID: "new-b", Namespace: "ns", AppType: 1, AppName: "b"},
|
||||
{OrderID: "new-c", Namespace: "ns", AppType: 1, AppName: "c"},
|
||||
}}
|
||||
existing := map[string][]*resources.Billing{"owner": {
|
||||
{OrderID: "legacy-a", Namespace: "ns", AppType: 1, AppName: "a", Status: resources.Settled},
|
||||
{
|
||||
OrderID: "bh_existing-b",
|
||||
Namespace: "ns",
|
||||
AppType: 1,
|
||||
AppName: "b",
|
||||
Status: resources.Unsettled,
|
||||
},
|
||||
{
|
||||
OrderID: "bh_orphan-d",
|
||||
Namespace: "ns",
|
||||
AppType: 1,
|
||||
AppName: "d",
|
||||
Status: resources.Unsettled,
|
||||
},
|
||||
}}
|
||||
pending := pendingOwnerBillings(generated, existing)["owner"]
|
||||
got := make(map[string]struct{}, len(pending))
|
||||
for _, billing := range pending {
|
||||
got[billing.OrderID] = struct{}{}
|
||||
}
|
||||
if len(got) != 1 {
|
||||
t.Fatalf("pending = %#v", pending)
|
||||
}
|
||||
for _, orderID := range []string{"new-c"} {
|
||||
if _, exists := got[orderID]; !exists {
|
||||
t.Fatalf("pending billing %s is missing: %#v", orderID, pending)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,116 @@
|
||||
/*
|
||||
Copyright 2023 sealos.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
*/
|
||||
|
||||
package controllers
|
||||
|
||||
import (
|
||||
"time"
|
||||
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
ctrlmetrics "sigs.k8s.io/controller-runtime/pkg/metrics"
|
||||
)
|
||||
|
||||
var (
|
||||
billingCheckpointTimestamp = prometheus.NewGauge(prometheus.GaugeOpts{
|
||||
Name: "sealos_account_billing_checkpoint_timestamp_seconds",
|
||||
Help: "Unix timestamp of the last successfully persisted account billing checkpoint.",
|
||||
})
|
||||
billingTargetTimestamp = prometheus.NewGauge(prometheus.GaugeOpts{
|
||||
Name: "sealos_account_billing_target_timestamp_seconds",
|
||||
Help: "Unix timestamp of the latest account billing hour ready for processing.",
|
||||
})
|
||||
billingCheckpointLagSeconds = prometheus.NewGauge(prometheus.GaugeOpts{
|
||||
Name: "sealos_account_billing_checkpoint_lag_seconds",
|
||||
Help: "Duration of ready account billing hours after the persisted checkpoint.",
|
||||
})
|
||||
billingPendingCheckpoints = prometheus.NewGauge(prometheus.GaugeOpts{
|
||||
Name: "sealos_account_billing_pending_checkpoints",
|
||||
Help: "Number of ready account billing hours after the persisted checkpoint.",
|
||||
})
|
||||
billingProcessing = prometheus.NewGauge(prometheus.GaugeOpts{
|
||||
Name: "sealos_account_billing_processing",
|
||||
Help: "Whether the account billing runner is processing a billing hour.",
|
||||
})
|
||||
billingProcessingTimestamp = prometheus.NewGauge(prometheus.GaugeOpts{
|
||||
Name: "sealos_account_billing_processing_timestamp_seconds",
|
||||
Help: "Unix timestamp of the account billing hour currently being processed.",
|
||||
})
|
||||
billingProcessingStartedTimestamp = prometheus.NewGauge(prometheus.GaugeOpts{
|
||||
Name: "sealos_account_billing_processing_started_timestamp_seconds",
|
||||
Help: "Unix timestamp when processing of the current account billing hour started.",
|
||||
})
|
||||
billingLastSuccessTimestamp = prometheus.NewGauge(prometheus.GaugeOpts{
|
||||
Name: "sealos_account_billing_last_success_timestamp_seconds",
|
||||
Help: "Unix timestamp when the account billing checkpoint was last advanced.",
|
||||
})
|
||||
billingReconcileFailures = prometheus.NewCounter(prometheus.CounterOpts{
|
||||
Name: "sealos_account_billing_reconcile_failures_total",
|
||||
Help: "Total account billing hour reconciliation failures.",
|
||||
})
|
||||
billingFailedOwners = prometheus.NewGauge(prometheus.GaugeOpts{
|
||||
Name: "sealos_account_billing_failed_owners",
|
||||
Help: "Number of owners that failed in the most recent account billing batch.",
|
||||
})
|
||||
)
|
||||
|
||||
func init() {
|
||||
ctrlmetrics.Registry.MustRegister(
|
||||
billingCheckpointTimestamp,
|
||||
billingTargetTimestamp,
|
||||
billingCheckpointLagSeconds,
|
||||
billingPendingCheckpoints,
|
||||
billingProcessing,
|
||||
billingProcessingTimestamp,
|
||||
billingProcessingStartedTimestamp,
|
||||
billingLastSuccessTimestamp,
|
||||
billingReconcileFailures,
|
||||
billingFailedOwners,
|
||||
)
|
||||
}
|
||||
|
||||
func setBillingTargetMetrics(target time.Time) {
|
||||
target = target.UTC().Truncate(time.Hour)
|
||||
billingTargetTimestamp.Set(float64(target.Unix()))
|
||||
}
|
||||
|
||||
func setBillingCheckpointMetrics(checkpoint, target time.Time) {
|
||||
checkpoint = checkpoint.UTC().Truncate(time.Hour)
|
||||
billingCheckpointTimestamp.Set(float64(checkpoint.Unix()))
|
||||
setBillingPendingCheckpointMetrics(checkpoint, target)
|
||||
}
|
||||
|
||||
func setBillingPendingCheckpointMetrics(checkpoint, target time.Time) {
|
||||
checkpoint = checkpoint.UTC().Truncate(time.Hour)
|
||||
target = target.UTC().Truncate(time.Hour)
|
||||
lag := target.Sub(checkpoint).Seconds()
|
||||
if lag < 0 {
|
||||
lag = 0
|
||||
}
|
||||
billingCheckpointLagSeconds.Set(lag)
|
||||
billingPendingCheckpoints.Set(lag / float64(time.Hour/time.Second))
|
||||
}
|
||||
|
||||
func setBillingProcessingMetrics(hour time.Time, processing bool) {
|
||||
if !processing {
|
||||
billingProcessing.Set(0)
|
||||
billingProcessingTimestamp.Set(0)
|
||||
billingProcessingStartedTimestamp.Set(0)
|
||||
return
|
||||
}
|
||||
billingProcessing.Set(1)
|
||||
billingProcessingTimestamp.Set(float64(hour.UTC().Truncate(time.Hour).Unix()))
|
||||
billingProcessingStartedTimestamp.Set(float64(time.Now().UTC().Unix()))
|
||||
}
|
||||
@@ -24,6 +24,7 @@ data:
|
||||
"DOMAIN" .Values.accountEnv.cloudDomain
|
||||
"PORT" .Values.accountEnv.cloudPort
|
||||
"ACCOUNT_API_JWT_SECRET" .Values.accountEnv.accountApiJwtSecret
|
||||
"BILLING_MAX_CATCHUP_DURATION" .Values.accountEnv.billingMaxCatchupDuration
|
||||
"BASE_BALANCE" .Values.accountEnv.baseBalance
|
||||
"QUOTA_LIMITS_CPU" .Values.accountEnv.quotaLimitsCpu
|
||||
"QUOTA_LIMITS_MEMORY" .Values.accountEnv.quotaLimitsMemory
|
||||
|
||||
@@ -60,6 +60,9 @@ accountEnv:
|
||||
# Authentication secrets (auto-configured from sealos-config)
|
||||
accountApiJwtSecret: "secret" # Auto-fetched from sealos-config.jwtInternal
|
||||
|
||||
# Maximum historical billing window replayed from an existing checkpoint.
|
||||
billingMaxCatchupDuration: "24h"
|
||||
|
||||
# Kubernetes API whitelist (auto-generated from cloudDomain)
|
||||
whitelistKubernetesHosts: "" # Auto-generated: https://${cloudDomain}:6443
|
||||
|
||||
|
||||
@@ -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