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 <dinoallo@netc.it>

.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:

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
+}

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) {

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),