From 76155a56b1bac16e24ee2564767fadf9cc8c477c Mon Sep 17 00:00:00 2001
From: Hyukjin Kwon <gurwls223@apache.org>
Date: Mon, 6 Jul 2026 18:42:54 +0900
Subject: [PATCH] [SPARK-57959][SQL][TESTS] Deflake
 MetricsFailureInjectionSuite by retrying until fetch-failure injection fires

The 'Non-deterministic stage block failure injection - injectFailure=true' test
is flaky (~1/6 in the full suite; passes 10/10 in isolation). macOS diagnostics
show the failing run has stage1=300 stage2=5 (a clean run, no recompute) and
finishes in ~370ms vs ~860ms on a passing run, and the test's own shuffle
receives no injected FetchFailed.

INJECT_SHUFFLE_FETCH_FAILURES corrupts mapper-0 of the shuffle map stage, but
whether the downstream reducer observes the resulting FetchFailed - and thus
forces the stage-1 recompute that inflates the raw metric - depends on task
scheduling in the shared SparkContext, and it occasionally does not fire, leaving
stage1Metric at exactly 300 and failing 'value > 300'.

Re-run the query (resetting metrics each attempt) until the injection forces a
recompute, up to 10 attempts, then run the original assertions. Test-only; does
not touch the shared DAGScheduler injection machinery used by other suites.

Generated-by: Claude Code
---
 .../metric/MetricsFailureInjectionSuite.scala | 70 +++++++++++++------
 1 file changed, 48 insertions(+), 22 deletions(-)

diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala
index 005734d14713d..a7ad1da466fbd 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/metric/MetricsFailureInjectionSuite.scala
@@ -317,28 +317,54 @@ class MetricsFailureInjectionSuite
         setUpTestTable("test_table")
         withSparkContextConf(
             config.Tests.INJECT_SHUFFLE_FETCH_FAILURES.key -> injectFailure.toString) {
-          val stage1MetricsExpr = incrementMetrics(Seq(stage1Metric, stage1SLAMetric))
-          val udfRand =
-            udf {
-              () => {
-                new Random().nextDouble()
-              }
-            }.asNondeterministic().apply().expr
-          val stage1 = spark.read.table("test_table")
-            .withColumn("non_deterministic_col", Column(udfRand))
-            .filter(Column(stage1MetricsExpr))
-          val stage2MetricsExpr = incrementMetrics(Seq(stage2Metric, stage2SLAMetric))
-          val stage2 = stage1
-            .groupBy("low_cardinality_col")
-            .avg("non_deterministic_col")
-            .filter(Column(stage2MetricsExpr))
-          // Add an extra stage with a single task to avoid flaky failures. If a ResultTask
-          // returns non-deterministic results to the client, it forces the query to abort
-          // instead of retrying the input stages.
-          val finalDf = stage2.repartition(1).as[(Int, Double)]
-          val result = finalDf.collect()
-          // Don't compare the second value, since it's random.
-          assert(result.map(_._1).toSet === (0 until 5).toSet)
+          def runOnce(): Dataset[_] = {
+            val stage1MetricsExpr = incrementMetrics(Seq(stage1Metric, stage1SLAMetric))
+            val udfRand =
+              udf {
+                () => {
+                  new Random().nextDouble()
+                }
+              }.asNondeterministic().apply().expr
+            val stage1 = spark.read.table("test_table")
+              .withColumn("non_deterministic_col", Column(udfRand))
+              .filter(Column(stage1MetricsExpr))
+            val stage2MetricsExpr = incrementMetrics(Seq(stage2Metric, stage2SLAMetric))
+            val stage2 = stage1
+              .groupBy("low_cardinality_col")
+              .avg("non_deterministic_col")
+              .filter(Column(stage2MetricsExpr))
+            // Add an extra stage with a single task to avoid flaky failures. If a ResultTask
+            // returns non-deterministic results to the client, it forces the query to abort
+            // instead of retrying the input stages.
+            val finalDf = stage2.repartition(1).as[(Int, Double)]
+            val result = finalDf.collect()
+            // Don't compare the second value, since it's random.
+            assert(result.map(_._1).toSet === (0 until 5).toSet)
+            finalDf
+          }
+
+          // The INJECT_SHUFFLE_FETCH_FAILURES machinery corrupts mapper-0 of the first successful
+          // attempt of the shuffle map stage. Whether the downstream reducer observes the resulting
+          // FetchFailed (and thus forces the stage-1 recompute that inflates the raw metric) depends
+          // on task scheduling within the shared SparkContext; across the suite it occasionally does
+          // not fire, leaving stage1Metric at exactly 300 and failing "value > 300" (a ~1/6 flake,
+          // more frequent on slower runners such as macOS arm64). When we require a recompute
+          // (injectFailure = true), re-run the query until the injection actually fires. Each attempt
+          // resets the metrics, so a successful attempt is indistinguishable from a first-try success.
+          var finalDf = runOnce()
+          if (injectFailure) {
+            var attempts = 1
+            while (stage1Metric.value <= 300 && attempts < 10) {
+              stage1Metric.reset()
+              stage2Metric.reset()
+              stage1SLAMetric.reset()
+              stage2SLAMetric.reset()
+              finalDf = runOnce()
+              attempts += 1
+            }
+            assert(stage1Metric.value > 300,
+              s"fetch-failure injection did not force a recompute after $attempts attempts")
+          }
           postRunChecks(finalDf)
           stage1Metric.reset()
           stage2Metric.reset()