From 4524f61e3a7aabc1c343198dfb08bbc8a2ee7824 Mon Sep 17 00:00:00 2001
From: Up Neck <163534172+up2neck@users.noreply.github.com>
Date: Wed, 22 Jan 2025 06:32:36 +0000
Subject: [PATCH] Refactor global labels removal in metrics aggregation process

Signed-off-by: Up Neck <163534172+up2neck@users.noreply.github.com>
---
 x-pack/apm-server/aggregation/lsm.go      | 23 +++++++++++--------
 x-pack/apm-server/aggregation/lsm_test.go | 28 +++++++++++++++++++++++
 2 files changed, 42 insertions(+), 9 deletions(-)
 create mode 100644 x-pack/apm-server/aggregation/lsm_test.go

diff --git a/x-pack/apm-server/aggregation/lsm.go b/x-pack/apm-server/aggregation/lsm.go
index 6e4d5c8cb64..077cd6c9c46 100644
--- a/x-pack/apm-server/aggregation/lsm.go
+++ b/x-pack/apm-server/aggregation/lsm.go
@@ -81,7 +81,12 @@ func (a *Aggregator) Stop(ctx context.Context) error {
 // so that aggregator can consume events from intake.
 func (a *Aggregator) ProcessBatch(ctx context.Context, b *modelpb.Batch) error {
 	for _, e := range *b {
-		removeRUMGlobalLabels(e)
+		// Remove global labels for RUM services to avoid explosion of metric groups
+		// to track for servicetxmetrics.
+		// For consistency, this will remove labels for other aggregated metrics as well.
+		if isRumAgent(e.Agent.GetName()) {
+			removeGlobalLabels(e)
+		}
 	}
 	return a.baseaggregator.AggregateBatch(ctx, [16]byte{}, b)
 }
@@ -108,17 +113,17 @@ func wrapNextProcessor(processor modelpb.BatchProcessor) aggregators.Processor {
 	}
 }
 
-func removeRUMGlobalLabels(event *modelpb.APMEvent) {
-	// Remove global labels for RUM services to avoid explosion of metric groups
-	// to track for servicetxmetrics.
-	// For consistency, this will remove labels for other aggregated metrics as well.
-	switch event.GetAgent().GetName() {
+// Check Elastic Agent belongs to rum-services
+func isRumAgent(name string) bool {
+	switch name {
 	case "rum-js", "js-base", "android/java", "iOS/swift":
-	default:
-		return
+		return true
 	}
+	return false
+}
 
-	// Setting the labels to non-global so that they are ignored by the aggregator.
+// Setting the labels to non-global so that they are ignored by the aggregator.
+func removeGlobalLabels(event *modelpb.APMEvent) {
 	for _, v := range event.Labels {
 		v.Global = false
 	}
diff --git a/x-pack/apm-server/aggregation/lsm_test.go b/x-pack/apm-server/aggregation/lsm_test.go
new file mode 100644
index 00000000000..30af52644f9
--- /dev/null
+++ b/x-pack/apm-server/aggregation/lsm_test.go
@@ -0,0 +1,28 @@
+package aggregation
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestIsRumAgent(t *testing.T) {
+	cases := []struct {
+		name   string
+		expect bool
+	}{
+		{"rum-js", true},
+		{"js-base", true},
+		{"android/java", true},
+		{"iOS/swift", true},
+		{"nodejs", false},
+		{"python", false},
+		{"java", false},
+		{"dotnet", false},
+	}
+
+	for _, tt := range cases {
+		got := isRumAgent(tt.name)
+		assert.Equal(t, tt.expect, got)
+	}
+}