From 76f7ebec4caed6f3f063ba17f9c261d2599b1cfd Mon Sep 17 00:00:00 2001
From: Hyukjin Kwon <gurwls223@apache.org>
Date: Mon, 6 Jul 2026 12:46:27 +0900
Subject: [PATCH] [SPARK][CORE][TESTS] Retry retryOnOOM a few times to reduce
 SorterSuite flaky OOM

SorterSuite's 'SPARK-11072 ...' test allocates a single ~1GB byte array and wraps
it in retryOnOOM (SPARK-57037), which does one GC + one retry. In memory-constrained
Maven CI this still OOMs intermittently (~1 in 5 core-module runs): a single
System.gc() after the first failure does not always reclaim the prior test's large
array in time, so the lone retry allocates against a still-pressured heap and fails.

Retry a few times (default 3), running a full runGC() (which waits for a weak
reference to clear) between attempts, to give the JVM more chances to reclaim memory.
This is a flakiness reduction, not a correctness change: an unbounded/too-large
allocation still fails once attempts are exhausted, so genuine OOMs are not masked.

Co-authored-by: Isaac
---
 .../org/apache/spark/SparkFunSuite.scala      | 29 ++++++++++++++-----
 .../spark/util/collection/SorterSuite.scala   |  4 +--
 2 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/core/src/test/scala/org/apache/spark/SparkFunSuite.scala b/core/src/test/scala/org/apache/spark/SparkFunSuite.scala
index a0f17f8af3f33..a18b1ec5dbf94 100644
--- a/core/src/test/scala/org/apache/spark/SparkFunSuite.scala
+++ b/core/src/test/scala/org/apache/spark/SparkFunSuite.scala
@@ -113,13 +113,28 @@ abstract class SparkFunSuite
     }
   }
 
-  /** Run `body`; if it throws OutOfMemoryError, force a GC and retry once. */
-  protected def retryOnOOM[T](body: => T): T = {
-    try body
-    catch {
-      case _: OutOfMemoryError =>
-        runGC()
-        body
+  /**
+   * Run `body`; if it throws OutOfMemoryError, force a GC and retry, up to `maxAttempts`
+   * total attempts (default 3). A single GC+retry is not always enough in memory-constrained
+   * CI: a large array retained by a previous test may not be reclaimed by the first
+   * System.gc(), so the retried allocation OOMs again. Retrying a few times (with a full
+   * runGC() that waits for a weak reference to clear between attempts) gives the JVM
+   * additional chances to reclaim memory. This does not mask genuine leaks: an unbounded
+   * allocation still fails once the attempts are exhausted.
+   */
+  protected def retryOnOOM[T](maxAttempts: Int = 3)(body: => T): T = {
+    var attempt = 1
+    while (true) {
+      try {
+        return body
+      } catch {
+        case oom: OutOfMemoryError =>
+          if (attempt >= maxAttempts) throw oom
+          runGC()
+          attempt += 1
+      }
     }
+    // Unreachable: the loop either returns a value or rethrows the OOM.
+    throw new IllegalStateException("retryOnOOM exited its retry loop unexpectedly")
   }
 }
diff --git a/core/src/test/scala/org/apache/spark/util/collection/SorterSuite.scala b/core/src/test/scala/org/apache/spark/util/collection/SorterSuite.scala
index 2767769924bc8..f654501ab436d 100644
--- a/core/src/test/scala/org/apache/spark/util/collection/SorterSuite.scala
+++ b/core/src/test/scala/org/apache/spark/util/collection/SorterSuite.scala
@@ -141,8 +141,8 @@ class SorterSuite extends SparkFunSuite {
     val arrayToSortSize = 1091482190
     // Memory held by the previous test (e.g. the ~256 MB int array in "SPARK-5984
     // TimSort bug") may not be reclaimed before this >1 GB allocation, causing flaky
-    // OOM in CI. Force a GC and retry once on OOM.
-    val arrayToSort = retryOnOOM(new Array[Byte](arrayToSortSize))
+    // OOM in CI. Force a GC and retry a few times on OOM.
+    val arrayToSort = retryOnOOM()(new Array[Byte](arrayToSortSize))
     var sum: Int = -1
     for (i <- runLengths) {
       sum += i