From 3047ab73f5ac8b711b55884f55ab809e96e68fbb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maciej=20Skocze=C5=84?= <mskoczen@google.com>
Date: Thu, 5 Sep 2024 13:39:32 +0000
Subject: [PATCH] Reset only metrics configured in collector before the
 createPodsOp

---
 .../scheduler_perf/scheduler_perf.go            | 14 +++++---------
 test/integration/scheduler_perf/util.go         | 17 ++++++++++++++---
 2 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/test/integration/scheduler_perf/scheduler_perf.go b/test/integration/scheduler_perf/scheduler_perf.go
index d08c643b172..183f56bc3bb 100644
--- a/test/integration/scheduler_perf/scheduler_perf.go
+++ b/test/integration/scheduler_perf/scheduler_perf.go
@@ -1139,7 +1139,10 @@ func runWorkload(tCtx ktesting.TContext, tc *testCase, w *workload, informerFact
 				for _, collector := range collectors {
 					// Need loop-local variable for function below.
 					collector := collector
-					collector.init()
+					err = collector.init()
+					if err != nil {
+						tCtx.Fatalf("op %d: Failed to initialize data collector: %v", opIndex, err)
+					}
 					collectorWG.Add(1)
 					go func() {
 						defer collectorWG.Done()
@@ -1205,13 +1208,6 @@ func runWorkload(tCtx ktesting.TContext, tc *testCase, w *workload, informerFact
 				}()
 			}
 
-			if !concreteOp.SkipWaitToCompletion {
-				// SkipWaitToCompletion=false indicates this step has waited for the Pods to be scheduled.
-				// So we reset the metrics in global registry; otherwise metrics gathered in this step
-				// will be carried over to next step.
-				legacyregistry.Reset()
-			}
-
 		case *churnOp:
 			var namespace string
 			if concreteOp.Namespace != nil {
@@ -1376,7 +1372,7 @@ func createNamespaceIfNotPresent(tCtx ktesting.TContext, namespace string, podsP
 }
 
 type testDataCollector interface {
-	init()
+	init() error
 	run(tCtx ktesting.TContext)
 	collect() []DataItem
 }
diff --git a/test/integration/scheduler_perf/util.go b/test/integration/scheduler_perf/util.go
index 1994b299e4a..2e93d42a62c 100644
--- a/test/integration/scheduler_perf/util.go
+++ b/test/integration/scheduler_perf/util.go
@@ -263,9 +263,19 @@ func newMetricsCollector(config *metricsCollectorConfig, labels map[string]strin
 	}
 }
 
-func (mc *metricsCollector) init() {
+func (mc *metricsCollector) init() error {
 	// Reset the metrics so that the measurements do not interfere with those collected during the previous steps.
-	legacyregistry.Reset()
+	m, err := legacyregistry.DefaultGatherer.Gather()
+	if err != nil {
+		return fmt.Errorf("failed to gather metrics to reset: %w", err)
+	}
+	for _, mFamily := range m {
+		// Reset only metrics defined in the collector.
+		if _, ok := mc.Metrics[mFamily.GetName()]; ok {
+			mFamily.Reset()
+		}
+	}
+	return nil
 }
 
 func (*metricsCollector) run(tCtx ktesting.TContext) {
@@ -381,7 +391,8 @@ func newThroughputCollector(podInformer coreinformers.PodInformer, labels map[st
 	}
 }
 
-func (tc *throughputCollector) init() {
+func (tc *throughputCollector) init() error {
+	return nil
 }
 
 func (tc *throughputCollector) run(tCtx ktesting.TContext) {