From b7aca45de7e40b7e99949480bfdfa3e188a0344b Mon Sep 17 00:00:00 2001
From: Ben Ye <benye@amazon.com>
Date: Thu, 31 Oct 2024 00:30:22 -0700
Subject: [PATCH] fix round function ignoring enableDelayedNameRemoval feature
 flag

Signed-off-by: Ben Ye <benye@amazon.com>
---
 promql/engine_test.go                         | 62 +++++++++++++++++++
 promql/functions.go                           |  3 +
 .../testdata/name_label_dropping.test         |  4 ++
 3 files changed, 69 insertions(+)

diff --git a/promql/engine_test.go b/promql/engine_test.go
index db399d8656..8da27a73c5 100644
--- a/promql/engine_test.go
+++ b/promql/engine_test.go
@@ -3501,3 +3501,65 @@ histogram {{sum:4 count:4 buckets:[2 2]}} {{sum:6 count:6 buckets:[3 3]}} {{sum:
 		},
 	})
 }
+
+func TestEvaluationWithDelayedNameRemovalDisabled(t *testing.T) {
+	opts := promql.EngineOpts{
+		Logger:                   nil,
+		Reg:                      nil,
+		EnableAtModifier:         true,
+		MaxSamples:               10000,
+		Timeout:                  10 * time.Second,
+		EnableDelayedNameRemoval: false,
+	}
+	engine := promqltest.NewTestEngineWithOpts(t, opts)
+
+	promqltest.RunTest(t, `
+load 5m
+	metric{env="1"}	0 60 120
+	another_metric{env="1"}	60 120 180
+
+# Does not drop __name__ for vector selector
+eval instant at 15m metric{env="1"}
+	metric{env="1"} 120
+
+# Drops __name__ for unary operators
+eval instant at 15m -metric
+	{env="1"} -120
+
+# Drops __name__ for binary operators
+eval instant at 15m metric + another_metric
+	{env="1"} 300
+
+# Does not drop __name__ for binary comparison operators
+eval instant at 15m metric <= another_metric
+	metric{env="1"} 120
+
+# Drops __name__ for binary comparison operators with "bool" modifier
+eval instant at 15m metric <= bool another_metric
+	{env="1"} 1
+
+# Drops __name__ for vector-scalar operations
+eval instant at 15m metric * 2
+	{env="1"} 240
+
+# Drops __name__ for instant-vector functions
+eval instant at 15m clamp(metric, 0, 100)
+	{env="1"} 100
+
+# Drops __name__ for round function
+eval instant at 15m round(metric)
+	{env="1"} 120
+
+# Drops __name__ for range-vector functions
+eval instant at 15m rate(metric{env="1"}[10m])
+	{env="1"} 0.2
+
+# Does not drop __name__ for last_over_time function
+eval instant at 15m last_over_time(metric{env="1"}[10m])
+	metric{env="1"} 120
+
+# Drops name for other _over_time functions
+eval instant at 15m max_over_time(metric{env="1"}[10m])
+	{env="1"} 120
+`, engine)
+}
diff --git a/promql/functions.go b/promql/functions.go
index cc5f19dae7..a509f783fa 100644
--- a/promql/functions.go
+++ b/promql/functions.go
@@ -534,6 +534,9 @@ func funcRound(vals []parser.Value, args parser.Expressions, enh *EvalNodeHelper
 
 	for _, el := range vec {
 		f := math.Floor(el.F*toNearestInverse+0.5) / toNearestInverse
+		if !enh.enableDelayedNameRemoval {
+			el.Metric = el.Metric.DropMetricName()
+		}
 		enh.Out = append(enh.Out, Sample{
 			Metric:   el.Metric,
 			F:        f,
diff --git a/promql/promqltest/testdata/name_label_dropping.test b/promql/promqltest/testdata/name_label_dropping.test
index 1f1dac3602..c8c0eb285f 100644
--- a/promql/promqltest/testdata/name_label_dropping.test
+++ b/promql/promqltest/testdata/name_label_dropping.test
@@ -31,6 +31,10 @@ eval instant at 15m metric * 2
 eval instant at 15m clamp(metric, 0, 100)
 	{env="1"} 100
 
+# Drops __name__ for round function
+eval instant at 15m round(metric)
+	{env="1"} 120
+
 # Drops __name__ for range-vector functions
 eval instant at 15m rate(metric{env="1"}[10m])
 	{env="1"} 0.2