From e812891355f63c71e9550e7f0ea480ed1bd31efa Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Thu, 3 Sep 2026 18:47:38 +0800 Subject: [PATCH] perf(account): optimize consumption amount aggregation (#7297) perf(account): optimize consumption amount aggregation (#7294) * perf(account): optimize consumption amount aggregation * fix(account): correct consumption app type filtering Co-authored-by: Yun Pan --- .github/workflows/service-build.yml | 2 +- service/account/dao/consumption_test.go | 77 ++++++ service/account/dao/interface.go | 173 ++++++------- .../dao/workspace_consumption_runtime_test.go | 231 +++++++++++++++++- 4 files changed, 395 insertions(+), 88 deletions(-) create mode 100644 service/account/dao/consumption_test.go diff --git a/.github/workflows/service-build.yml b/.github/workflows/service-build.yml index 90209a7ef..f75eba2c2 100644 --- a/.github/workflows/service-build.yml +++ b/.github/workflows/service-build.yml @@ -87,7 +87,7 @@ jobs: working-directory: service/account env: TESTCONTAINERS_REQUIRED: "true" - run: go test ./dao -run '^TestGetWorkspaceConsumptionAmountWithMongoRuntime$' -count=1 -v + run: go test ./dao -run '^TestGet(Workspace)?ConsumptionAmountWithMongoRuntime$' -count=1 -v image-build: strategy: diff --git a/service/account/dao/consumption_test.go b/service/account/dao/consumption_test.go new file mode 100644 index 000000000..059c4a9c4 --- /dev/null +++ b/service/account/dao/consumption_test.go @@ -0,0 +1,77 @@ +package dao + +import ( + "testing" + "time" + + "github.com/labring/sealos/controllers/pkg/resources" + "github.com/labring/sealos/service/account/helper" + "go.mongodb.org/mongo-driver/bson" +) + +func consumptionRequest() helper.ConsumptionRecordReq { + return helper.ConsumptionRecordReq{ + TimeRange: helper.TimeRange{ + StartTime: time.Date(2026, time.January, 1, 0, 0, 0, 0, time.UTC), + EndTime: time.Date(2026, time.January, 2, 0, 0, 0, 0, time.UTC), + }, + AuthBase: helper.AuthBase{ + Auth: &helper.Auth{Owner: "owner-test"}, + }, + } +} + +func TestBuildConsumptionAmountPipeline(t *testing.T) { + tests := []struct { + name string + request helper.ConsumptionRecordReq + }{ + { + name: "all consumption", + request: consumptionRequest(), + }, + { + name: "namespace and app filters", + request: func() helper.ConsumptionRecordReq { + req := consumptionRequest() + req.Namespace = "ns-test" + req.AppType = resources.APP + req.AppName = "app-test" + return req + }(), + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + pipeline := buildConsumptionAmountPipeline(test.request) + if len(pipeline) != 3 { + t.Fatalf("pipeline stage count = %d, want 3", len(pipeline)) + } + if !hasConsumptionStage(pipeline[0], "$match") { + t.Fatal("pipeline does not start with $match") + } + if !hasConsumptionStage(pipeline[1], "$project") { + t.Fatal("pipeline does not project one amount per billing record") + } + if hasConsumptionStage(pipeline[1], "$facet") { + t.Fatal("pipeline should not use $facet") + } + if hasConsumptionStage(pipeline[1], "$unwind") { + t.Fatal("pipeline should not use $unwind") + } + if !hasConsumptionStage(pipeline[2], "$group") { + t.Fatal("pipeline does not end with $group") + } + }) + } +} + +func hasConsumptionStage(stage bson.D, key string) bool { + for _, element := range stage { + if element.Key == key { + return true + } + } + return false +} diff --git a/service/account/dao/interface.go b/service/account/dao/interface.go index 5c4383845..f7477f1b6 100644 --- a/service/account/dao/interface.go +++ b/service/account/dao/interface.go @@ -2218,126 +2218,129 @@ func (m *Account) Disconnect(ctx context.Context) error { return nil } -func (m *MongoDB) GetConsumptionAmount(req helper.ConsumptionRecordReq) (int64, error) { - owner, namespace, appType, appName, startTime, endTime := req.Owner, req.Namespace, req.AppType, req.AppName, req.StartTime, req.EndTime +func buildConsumptionAmountPipeline(req helper.ConsumptionRecordReq) mongo.Pipeline { + appType := strings.ToUpper(strings.TrimSpace(req.AppType)) timeMatchValue := bson.D{ - primitive.E{Key: "$gte", Value: startTime}, - primitive.E{Key: "$lte", Value: endTime}, + primitive.E{Key: "$gte", Value: req.StartTime}, + primitive.E{Key: "$lte", Value: req.EndTime}, } - - // Build base match conditions for app_costs (sub-consumption type) matchValue := bson.D{ - primitive.E{Key: "time", Value: timeMatchValue}, + primitive.E{Key: "owner", Value: req.Owner}, primitive.E{Key: "status", Value: resources.Settled}, - primitive.E{Key: "owner", Value: owner}, + primitive.E{Key: "time", Value: timeMatchValue}, + } + if req.Namespace != "" { + matchValue = append(matchValue, primitive.E{Key: "namespace", Value: req.Namespace}) } if appType != "" { matchValue = append( matchValue, - primitive.E{Key: "app_type", Value: resources.AppType[strings.ToUpper(appType)]}, + primitive.E{Key: "app_type", Value: resources.AppType[appType]}, ) } - if namespace != "" { - matchValue = append(matchValue, primitive.E{Key: "namespace", Value: namespace}) + + appCostsInput := bson.M{"$ifNull": bson.A{"$app_costs", bson.A{}}} + appCostsAmount := bson.M{ + "$reduce": bson.M{ + "input": appCostsInput, + "initialValue": int64(0), + "in": bson.M{"$add": bson.A{ + "$$value", + bson.M{"$ifNull": bson.A{"$$this.amount", int64(0)}}, + }}, + }, } - unwindMatchValue := bson.D{ - primitive.E{Key: "time", Value: timeMatchValue}, - } - if appType != "" && appName != "" { + + // Preserve the legacy app_costs matching semantics while avoiding $unwind. + nestedAmount := any(appCostsAmount) + if appType != "" && req.AppName != "" { if appType != resources.AppStore { - unwindMatchValue = append( - unwindMatchValue, - primitive.E{Key: "app_costs.name", Value: appName}, - ) + filteredAppCosts := bson.M{ + "$filter": bson.M{ + "input": appCostsInput, + "as": "appCost", + "cond": bson.M{"$eq": bson.A{"$$appCost.name", req.AppName}}, + }, + } + nestedAmount = bson.M{ + "$reduce": bson.M{ + "input": filteredAppCosts, + "initialValue": int64(0), + "in": bson.M{"$add": bson.A{ + "$$value", + bson.M{"$ifNull": bson.A{"$$this.amount", int64(0)}}, + }}, + }, + } } else { - unwindMatchValue = append( - unwindMatchValue, - primitive.E{Key: "app_name", Value: appName}, - ) + nestedAmount = bson.M{ + "$cond": bson.A{ + bson.M{"$eq": bson.A{"$app_name", req.AppName}}, + appCostsAmount, + int64(0), + }, + } } } - // Build match conditions for direct consumption (AppStore and LLMToken) - directMatchValue := bson.D{ - primitive.E{Key: "time", Value: timeMatchValue}, - primitive.E{Key: "status", Value: resources.Settled}, - primitive.E{Key: "owner", Value: owner}, - } - if namespace != "" { - directMatchValue = append(directMatchValue, primitive.E{Key: "namespace", Value: namespace}) - } - // For direct consumption, match app_type to AppStore or LLMToken if not specified + directCondition := any(bson.M{ + "$in": bson.A{"$app_type", bson.A{ + resources.AppType[resources.AppStore], + resources.AppType[resources.LLMToken], + }}, + }) if appType != "" { - directMatchValue = append( - directMatchValue, - primitive.E{Key: "app_type", Value: resources.AppType[strings.ToUpper(appType)]}, - ) - } else { - // If no appType specified, match both AppStore and LLMToken - directMatchValue = append( - directMatchValue, - primitive.E{Key: "app_type", Value: bson.D{{Key: "$in", Value: bson.A{ - resources.AppType[resources.AppStore], - resources.AppType[resources.LLMToken], - }}}}, - ) + directCondition = appType == resources.AppStore || appType == resources.LLMToken } - if appName != "" { - directMatchValue = append(directMatchValue, primitive.E{Key: "app_name", Value: appName}) + if req.AppName != "" { + directCondition = bson.M{"$and": bson.A{ + directCondition, + bson.M{"$eq": bson.A{"$app_name", req.AppName}}, + }} + } + directAmount := bson.M{ + "$cond": bson.A{ + directCondition, + bson.M{"$ifNull": bson.A{"$amount", int64(0)}}, + int64(0), + }, } - // Use $facet to query both types in parallel - pipeline := bson.A{ - bson.D{{Key: "$facet", Value: bson.M{ - "appCosts": bson.A{ - bson.D{{Key: "$match", Value: matchValue}}, - bson.D{{Key: "$unwind", Value: "$app_costs"}}, - bson.D{{Key: "$match", Value: unwindMatchValue}}, - bson.D{{Key: "$group", Value: bson.M{ - "_id": nil, - "total": bson.M{"$sum": "$app_costs.amount"}, - }}}, - }, - "directAmount": bson.A{ - bson.D{{Key: "$match", Value: directMatchValue}}, - bson.D{{Key: "$group", Value: bson.M{ - "_id": nil, - "total": bson.M{"$sum": "$amount"}, - }}}, - }, + return mongo.Pipeline{ + {{Key: "$match", Value: matchValue}}, + {{Key: "$project", Value: bson.D{ + {Key: "amount", Value: bson.M{"$add": bson.A{nestedAmount, directAmount}}}, + }}}, + {{Key: "$group", Value: bson.D{ + {Key: "_id", Value: nil}, + {Key: "total", Value: bson.D{{Key: "$sum", Value: "$amount"}}}, }}}, } +} - cursor, err := m.getBillingCollection().Aggregate(context.Background(), pipeline) +func (m *MongoDB) GetConsumptionAmount(req helper.ConsumptionRecordReq) (int64, error) { + pipeline := buildConsumptionAmountPipeline(req) + + ctx := context.Background() + cursor, err := m.getBillingCollection().Aggregate(ctx, pipeline) if err != nil { return 0, fmt.Errorf("failed to aggregate billing collection: %w", err) } - defer cursor.Close(context.Background()) + defer cursor.Close(ctx) var result struct { - AppCosts []struct { - Total int64 `bson:"total"` - } `bson:"appCosts"` - DirectAmount []struct { - Total int64 `bson:"total"` - } `bson:"directAmount"` + Total int64 `bson:"total"` } - if cursor.Next(context.Background()) { + if cursor.Next(ctx) { if err := cursor.Decode(&result); err != nil { return 0, fmt.Errorf("failed to decode result: %w", err) } + } else if err := cursor.Err(); err != nil { + return 0, fmt.Errorf("failed to iterate aggregate result: %w", err) } - totalAmount := int64(0) - if len(result.AppCosts) > 0 { - totalAmount += result.AppCosts[0].Total - } - if len(result.DirectAmount) > 0 { - totalAmount += result.DirectAmount[0].Total - } - - return totalAmount, nil + return result.Total, nil } func normalizeWorkspaceConsumptionAppType(appType string) (string, uint8, error) { diff --git a/service/account/dao/workspace_consumption_runtime_test.go b/service/account/dao/workspace_consumption_runtime_test.go index 0112b873e..a85db7398 100644 --- a/service/account/dao/workspace_consumption_runtime_test.go +++ b/service/account/dao/workspace_consumption_runtime_test.go @@ -12,6 +12,7 @@ import ( "github.com/labring/sealos/service/account/helper" "github.com/testcontainers/testcontainers-go" "github.com/testcontainers/testcontainers-go/wait" + "go.mongodb.org/mongo-driver/bson" "go.mongodb.org/mongo-driver/mongo" "go.mongodb.org/mongo-driver/mongo/options" ) @@ -72,11 +73,22 @@ func newWorkspaceConsumptionMongo(tb testing.TB) (*MongoDB, context.Context) { } }) - return &MongoDB{ + mongoDB := &MongoDB{ Client: client, AccountDBName: workspaceConsumptionTestDB, BillingConn: workspaceConsumptionTestColl, - }, ctx + } + if _, err := mongoDB.getBillingCollection().Indexes().CreateOne(ctx, mongo.IndexModel{ + Keys: bson.D{ + {Key: "owner", Value: 1}, + {Key: "status", Value: 1}, + {Key: "time", Value: 1}, + }, + }); err != nil { + tb.Fatalf("create billing query index: %v", err) + } + + return mongoDB, ctx } func skipIfWorkspaceConsumptionDockerIsNotHealthy(tb testing.TB) { @@ -230,6 +242,131 @@ func TestGetWorkspaceConsumptionAmountWithMongoRuntime(t *testing.T) { } } +func TestGetConsumptionAmountWithMongoRuntime(t *testing.T) { + mongoDB, ctx := newWorkspaceConsumptionMongo(t) + startTime := time.Date(2026, time.January, 1, 0, 0, 0, 0, time.UTC) + endTime := time.Date(2026, time.January, 2, 0, 0, 0, 0, time.UTC) + documents := []any{ + resources.Billing{ + Time: startTime, OrderID: "nested-consumption", Type: resources.Consumption, + Namespace: "ns-a", AppCosts: []resources.AppCost{ + {Name: "app-a", Amount: 30}, + {Name: "app-b", Amount: 5}, + }, + AppType: resources.AppType[resources.APP], Amount: 100, + Owner: workspaceConsumptionTestOwner, Status: resources.Settled, + }, + resources.Billing{ + Time: endTime, OrderID: "llm-consumption", Type: resources.SubConsumption, + Namespace: "ns-b", AppName: "llm-a", AppType: resources.AppType[resources.LLMToken], + Amount: 20, Owner: workspaceConsumptionTestOwner, Status: resources.Settled, + }, + resources.Billing{ + Time: endTime, OrderID: "app-store-consumption", Type: resources.Consumption, + Namespace: "ns-c", AppName: "store-a", AppType: resources.AppType[resources.AppStore], + Amount: 40, Owner: workspaceConsumptionTestOwner, Status: resources.Settled, + }, + resources.Billing{ + Time: endTime, OrderID: "unsettled-consumption", Type: resources.Consumption, + Namespace: "ns-ignored", AppCosts: []resources.AppCost{{Name: "app-a", Amount: 1000}}, + AppType: resources.AppType[resources.APP], Amount: 1000, + Owner: workspaceConsumptionTestOwner, Status: resources.Unsettled, + }, + resources.Billing{ + Time: endTime, OrderID: "other-owner-consumption", Type: resources.Consumption, + Namespace: "ns-ignored", AppCosts: []resources.AppCost{{Name: "app-a", Amount: 2000}}, + AppType: resources.AppType[resources.APP], Amount: 2000, + Owner: "other-owner", Status: resources.Settled, + }, + resources.Billing{ + Time: endTime.Add(time.Hour), + OrderID: "outside-consumption", Type: resources.Consumption, + Namespace: "ns-ignored", AppCosts: []resources.AppCost{{Name: "app-a", Amount: 3000}}, + AppType: resources.AppType[resources.APP], Amount: 3000, + Owner: workspaceConsumptionTestOwner, Status: resources.Settled, + }, + } + if _, err := mongoDB.getBillingCollection().InsertMany(ctx, documents); err != nil { + t.Fatalf("insert billing fixtures: %v", err) + } + + baseRequest := helper.ConsumptionRecordReq{ + TimeRange: helper.TimeRange{StartTime: startTime, EndTime: endTime}, + AuthBase: helper.AuthBase{Auth: &helper.Auth{Owner: workspaceConsumptionTestOwner}}, + } + tests := []struct { + name string + req helper.ConsumptionRecordReq + want int64 + }{ + {name: "all settled consumption", req: baseRequest, want: 95}, + { + name: "namespace filter", + req: withWorkspaceConsumptionRequest( + baseRequest, + func(req *helper.ConsumptionRecordReq) { + req.Namespace = "ns-a" + }, + ), + want: 35, + }, + { + name: "nested app type filter", + req: withWorkspaceConsumptionRequest( + baseRequest, + func(req *helper.ConsumptionRecordReq) { + req.AppType = " app " + }, + ), + want: 35, + }, + { + name: "nested app name filter", + req: withWorkspaceConsumptionRequest( + baseRequest, + func(req *helper.ConsumptionRecordReq) { + req.AppType = resources.APP + req.AppName = "app-a" + }, + ), + want: 30, + }, + { + name: "direct app name filter", + req: withWorkspaceConsumptionRequest( + baseRequest, + func(req *helper.ConsumptionRecordReq) { + req.AppName = "store-a" + }, + ), + want: 75, + }, + { + name: "direct app type and name filter", + req: withWorkspaceConsumptionRequest( + baseRequest, + func(req *helper.ConsumptionRecordReq) { + req.AppType = resources.AppStore + req.AppName = "store-a" + }, + ), + want: 40, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + got, err := mongoDB.GetConsumptionAmount(test.req) + if err != nil { + t.Fatalf("get consumption amount: %v", err) + } + if got != test.want { + t.Fatalf("consumption amount = %d, want %d", got, test.want) + } + }) + } +} + func BenchmarkGetWorkspaceConsumptionAmountWithMongoRuntime(b *testing.B) { mongoDB, ctx := newWorkspaceConsumptionMongo(b) startTime := time.Date(2026, time.January, 1, 0, 0, 0, 0, time.UTC) @@ -306,6 +443,96 @@ func BenchmarkGetWorkspaceConsumptionAmountWithMongoRuntime(b *testing.B) { } } +func BenchmarkGetConsumptionAmountWithMongoRuntime(b *testing.B) { + mongoDB, ctx := newWorkspaceConsumptionMongo(b) + startTime := time.Date(2026, time.January, 1, 0, 0, 0, 0, time.UTC) + endTime := startTime.Add(24 * time.Hour) + documents := make([]any, 0, workspaceConsumptionBenchmarkRecords) + for i := range workspaceConsumptionBenchmarkRecords { + document := resources.Billing{ + Time: startTime.Add(time.Duration(i%24) * time.Hour), + OrderID: fmt.Sprintf("consumption-benchmark-%d", i), + Namespace: fmt.Sprintf("ns-%02d", i%32), + Owner: workspaceConsumptionTestOwner, + Status: resources.Settled, + } + switch i % 4 { + case 0: + document.Type = resources.Consumption + document.AppType = resources.AppType[resources.AppStore] + document.AppName = "store-a" + document.Amount = int64(i%100 + 1) + case 1: + document.Type = resources.SubConsumption + document.AppType = resources.AppType[resources.LLMToken] + document.AppName = "llm-a" + document.Amount = int64(i%100 + 1) + default: + document.Type = resources.Consumption + document.AppType = resources.AppType[resources.APP] + document.AppCosts = []resources.AppCost{ + {Name: "app-a", Amount: int64(i%100 + 1)}, + {Name: "app-b", Amount: 5}, + } + } + documents = append(documents, document) + } + if _, err := mongoDB.getBillingCollection().InsertMany(ctx, documents); err != nil { + b.Fatalf("insert billing benchmark fixtures: %v", err) + } + + baseRequest := helper.ConsumptionRecordReq{ + TimeRange: helper.TimeRange{StartTime: startTime, EndTime: endTime}, + AuthBase: helper.AuthBase{Auth: &helper.Auth{Owner: workspaceConsumptionTestOwner}}, + } + benchmarks := []struct { + name string + req helper.ConsumptionRecordReq + }{ + {name: "all", req: baseRequest}, + { + name: "namespace", + req: withWorkspaceConsumptionRequest( + baseRequest, + func(req *helper.ConsumptionRecordReq) { + req.Namespace = "ns-07" + }, + ), + }, + { + name: "app_type", + req: withWorkspaceConsumptionRequest( + baseRequest, + func(req *helper.ConsumptionRecordReq) { + req.AppType = resources.APP + }, + ), + }, + { + name: "app_name", + req: withWorkspaceConsumptionRequest( + baseRequest, + func(req *helper.ConsumptionRecordReq) { + req.AppType = resources.APP + req.AppName = "app-a" + }, + ), + }, + } + + for _, benchmark := range benchmarks { + b.Run(benchmark.name, func(b *testing.B) { + b.ReportAllocs() + b.ResetTimer() + for range b.N { + if _, err := mongoDB.GetConsumptionAmount(benchmark.req); err != nil { + b.Fatalf("get consumption amount: %v", err) + } + } + }) + } +} + func withWorkspaceConsumptionRequest( base helper.ConsumptionRecordReq, update func(*helper.ConsumptionRecordReq),