[SPARK-58507][CORE] Avoid redundant closure-class parsing in ClosureCleaner's indylambda path

### What changes were proposed in this pull request?

Two changes to `ClosureCleaner`'s indylambda path:

1. Hoist the `getCapturedArgCount == 0` check above `Class.forName` and the ASM parse. It is an O(1) read of a `SerializedLambda` the method already holds, and when it is zero there is nothing to clean, so the expensive work was being done only to be discarded. Skipping the return-statement fail-fast for such closures is safe: a non-local return compiles to `throw new NonLocalReturnControl(key, value)` where `key` is allocated in the enclosing method, so a closure containing one necessarily captures that key. The null-`getCapturedArg(0)` bail-out, by contrast, deliberately stays _below_ the return-statement check: a closure with a non-local return can capture a null value as its first captured argument, so hoisting that bail-out too would silently skip the fail-fast (see the new regression test, which fails in that configuration).

2. Memoize the return-statement check per class in a `java.lang.ClassValue`. `ReturnStatementFinder` is replaced by a collecting variant (`ReturnStatementCollector`) so one parse answers for every method on the class, preserving the existing `$adapted` matching rule exactly.

    Scanning granularity is unchanged: a classfile cannot be parsed per-method (ASM's `accept` walks the whole file), and both the old and new visitors decode instructions only for `apply`/`$anonfun$` methods. The previous code paid that full-class walk on every `clean()` call; the memoized version pays it once per class. The only per-parse work lost is the old throw-on-first-match early abort, which fired only when the closure was about to be rejected anyway. The stored result is a set of offending method names, or the shared empty-set singleton for any class with no non-local returns.

    Why `ClassValue` rather than a map keyed on `Class`? A `ConcurrentHashMap[Class[_], _]` would strongly reference its keys and prevent classloader unloading, which matters here because Spark's REPL and `ChildFirstURLClassLoader` discard loaders routinely. `ClassValue` associates the state with the `Class` itself, so an entry becomes unreachable precisely when its class does; no unloading is blocked. Verified empirically: a populated entry does not prevent class or classloader collection, and 20000 sequentially created classes each given an entry left zero alive after GC with metaspace flat. `computeValue` may be invoked concurrently and more than once per class, which is safe because return-statement detection is a pure function of the class bytes.

As a side effect this removes a latent NPE: `getClassReader` has always been allowed to return `null` (bytecode is not resource-accessible for e.g. LambdaMetafactory-generated classes, which is why SPARK-14540 added the same guard in `getInnerClosureClasses`), and the two return-statement call sites dereferenced it unconditionally. At these sites the argument is a capturing class, which in practice has readable bytecode, so the NPE is not known to fire; the memoized check now honors the contract and skips such classes with a debug log. The fail-fast is best-effort, so skipping cannot cause incorrect execution, only a later, less friendly error if the closure really contains a non-local return.

### Why are the changes needed?

`RDD.collect()` passes `(iter: Iterator[T]) => iter.toArray` to `runJob`, which captures nothing, yet `SparkContext.runJob` cleans it unconditionally -- loading a class, reading a class file out of a JAR and running a full ASM parse, on every collect.

Profiling a Spark test JVM (`SQLQueryTestSuite`) showed:

* `ClosureCleaner` on the stack for 13.2% of CPU samples and 25.2% of all allocation samples;
* 3251 `getClassReader` invocations over 21 distinct classes -- a 155x repeat ratio, because the lambdas passed to `runJob` are declared by a handful of Spark's own classes (`WholeStageCodegenExec`, `SparkContext`, `RDD`, `Dataset`), so every job re-parses the same bytecode;
* 14.7% of indylambda cleans are non-capturing.

The cost is a per-job overhead, measured at 9.8-13.2% of test-JVM CPU across six suites including `core/RDDSuite`, which involves no SQL at all.

Note that the ClosureCleaner's indylambda path only modifies closures whose first captured argument is a Scala REPL line object or an instance of a class defined in an Ammonite session; for all other code it is purely a validator. Since modification is gated on a captured argument existing, a closure with `getCapturedArgCount == 0` can never be modified, so the hoist cannot remove cleaning, and the legacy non-indylambda path's cleaning is untouched by this PR (only its validator is memoized).

### Does this PR introduce _any_ user-facing change?

No. Behaviour is unchanged, including the fail-fast `ReturnStatementInClosureException` for every capturing closure.

### How was this patch tested?

* New regression test "return statements in closures capturing a null value are identified at cleaning time": its closure's first captured argument is a null local that precedes the `NonLocalReturnControl` key in the capture order (verified via javap: `$anonfun$run$16(String, Object, int)`), pinning the requirement that the capture-count hoist must not skip the return-statement check for capturing closures. It fails if the null-capture bail-out is hoisted above the check.
* New unit test "hasReturnStatement identifies non-local returns per method": exercises the any-method query, the exact impl-method match, the `$adapted`-suffix resolution rule (scala/scala-dev#109), a non-matching method name, and a class with no non-local returns.
* `ClosureCleanerSuite` + `ClosureCleanerSuite2`: 18/18 (`ClosureCleanerSuite`, 12/12, rerun locally against the final patch).
* `ReplSuite` + `SingletonReplSuite`: 36/36, and `AmmoniteReplE2ESuite`: 1/1. These are the two paths where cleaning actually modifies closures: the indylambda path only rewrites Scala REPL and Ammonite closures, so every other suite exercises the branch where a regression would be invisible.
* `DataFrameSuite`: 176/176.
* Effect confirmed by profiling the patched build on the same suite: `xbean.asm9` falls from 11.02% of samples to 0.17%, `getClassReader` to zero.
* The above tests / profiling was on a Linux VPS box. Independently reproduced on a second machine (macOS, JDK 21, JFR `profile` settings), comparing baseline and patched `ClosureCleaner` swapped in front of an otherwise identical classpath:
  - `core/RDDSuite` end-to-end (79/79 green in both): cleaning-related frames fell from 8.1% of execution samples to 0.0%, `getClassReader` calls from 1306 over 11 distinct classes (a 119x repeat ratio) to 13, and suite wall clock dropped ~8.5%.
  - A loop of minimal `collect()`/`count()` jobs showed where the time goes: the driver re-parses `SparkContext.class` (204 KB) and `RDD.class` (189 KB) about five times per job; per-iteration wall time halved (11.1 ms -> 5.5 ms).
  - Microbenchmark of `SparkClosureCleaner.clean()` alone (baseline cost scales with the capturing class's file size; the memoized path is size-independent): for a 9 KB capturing class, 46.4 us -> 2.0 us per call; for a 113 KB one, 340 us -> 4.8 us. Non-capturing: 44.9 us -> 0.28 us. The memo's worst case (a never-repeated closure class) pays the same single parse as before plus ~360 bytes; measured repeat ratios on real suites were 119x (`RDDSuite`) and 155x (`DataFrameSuite`).

### Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Opus 5; refined with Claude Fable 5

Closes #57710 from JoshRosen/SPARK-58507-closurecleaner-redundant-parsing.

Authored-by: Josh Rosen <rosenville@gmail.com>
Signed-off-by: yangjie01 <yangjie01@baidu.com>

Author: JoshRosen

common/utils/src/main/scala/org/apache/spark/util/ClosureCleaner.scala

@@ -21,6 +21,7 @@ import java.io.{ByteArrayInputStream, ByteArrayOutputStream}
 import java.lang.invoke.{LambdaMetafactory, MethodHandle, MethodHandleInfo, MethodHandles, MethodType, SerializedLambda}
 import java.lang.reflect.{Field, Modifier}
 
+import scala.collection.immutable
 import scala.collection.mutable.{ArrayBuffer, Map, Queue, Set, Stack}
 import scala.jdk.CollectionConverters._
 
@@ -34,6 +35,40 @@ import org.apache.spark.internal.Logging
  * A cleaner that renders closures serializable if they can be done so safely.
  */
 private[spark] object ClosureCleaner extends Logging {
+  /**
+   * Per-class memo of which closure methods contain a non-local return, i.e. allocate a
+   * `scala/runtime/NonLocalReturnControl`. The verdict is a pure function of the class's
+   * immutable bytecode, so one ASM parse per class answers for every `clean()` call.
+   */
+  private val methodsWithNonLocalReturn = new ClassValue[immutable.Set[String]] {
+    override def computeValue(cls: Class[_]): immutable.Set[String] = {
+      val collector = new ReturnStatementCollector
+      val reader = getClassReader(cls)
+      if (reader != null) {
+        reader.accept(collector, 0)
+      } else {
+        logDebug(s"Cannot get class bytes for ${cls.getName}; skipping return-statement check")
+      }
+      collector.found.toSet
+    }
+  }
+
+  /** Whether `implMethodName` (any closure method, if `None`) of `cls` has a non-local return. */
+  private[util] def hasReturnStatement(cls: Class[_], implMethodName: Option[String]): Boolean = {
+    val found = methodsWithNonLocalReturn.get(cls)
+    implMethodName match {
+      case None => found.nonEmpty
+      case Some(target) =>
+        // Some lambdas get an "$adapted" boxing bridge (e.g. { _: Int => return; Seq() }) while
+        // others do not ({ _: Int => return; true }, fully specialized). When the bridge
+        // exists, the SerializedLambda's impl method is the bridge, which only delegates: the
+        // `new NonLocalReturnControl` instruction is emitted in the underlying unadapted
+        // method, so also match with the suffix stripped.
+        // See https://github.com/scala/scala-dev/issues/109.
+        found.contains(target) || found.contains(target.stripSuffix("$adapted"))
+    }
+  }
+
   // Get an ASM class reader for a given class from the JAR that loaded it
   private[util] def getClassReader(cls: Class[_]): ClassReader = {
     // Copy data over, before delegating to ClassReader - else we can run out of open file handles.
@@ -226,6 +261,13 @@ private[spark] object ClosureCleaner extends Logging {
 
       logDebug(s"Cleaning indylambda closure: $implMethodName")
 
+      // A closure with no captured arguments needs no cleaning, and cannot contain a non-local
+      // return (the `NonLocalReturnControl` key is allocated in the enclosing method and would
+      // be captured), so return before loading and parsing the capturing class.
+      if (lambdaProxy.getCapturedArgCount == 0) {
+        return None
+      }
+
       // capturing class is the class that declared this lambda
       val capturingClassName = lambdaProxy.getCapturingClass.replace('/', '.')
       val classLoader = func.getClass.getClassLoader // this is the safest option
@@ -234,16 +276,13 @@ private[spark] object ClosureCleaner extends Logging {
       // scalastyle:on classforname
 
       // Fail fast if we detect return statements in closures
-      val capturingClassReader = getClassReader(capturingClass)
-      capturingClassReader.accept(new ReturnStatementFinder(Option(implMethodName)), 0)
-
-      val outerThis = if (lambdaProxy.getCapturedArgCount > 0) {
-        // only need to clean when there is an enclosing non-null "this" captured by the closure
-        Option(lambdaProxy.getCapturedArg(0)).getOrElse(return None)
-      } else {
-        return None
+      if (hasReturnStatement(capturingClass, Option(implMethodName))) {
+        throw new ReturnStatementInClosureException
       }
 
+      // only need to clean when there is an enclosing non-null "this" captured by the closure
+      val outerThis = Option(lambdaProxy.getCapturedArg(0)).getOrElse(return None)
+
       // clean only if enclosing "this" is something cleanable, i.e. a Scala REPL line object or
       // Ammonite command helper object.
       // For Ammonite closures, we do not care about actual capturing class name,
@@ -312,7 +351,9 @@ private[spark] object ClosureCleaner extends Logging {
     }
 
     // Fail fast if we detect return statements in closures
-    getClassReader(func.getClass).accept(new ReturnStatementFinder(), 0)
+    if (hasReturnStatement(func.getClass, None)) {
+      throw new ReturnStatementInClosureException
+    }
 
     // If accessed fields is not populated yet, we assume that
     // the closure we are trying to clean is the starting one
@@ -1075,26 +1116,19 @@ private[spark] object IndylambdaScalaClosures extends Logging {
 private[spark] class ReturnStatementInClosureException
   extends SparkException("Return statements aren't allowed in Spark closures")
 
-private class ReturnStatementFinder(targetMethodName: Option[String] = None)
-  extends ClassVisitor(Opcodes.ASM9) {
+/** Collects the names of all closure methods that contain a non-local return. */
+private class ReturnStatementCollector extends ClassVisitor(Opcodes.ASM9) {
+  val found = Set.empty[String]
+
   override def visitMethod(access: Int, name: String, desc: String,
       sig: String, exceptions: Array[String]): MethodVisitor = {
 
     // $anonfun$ covers indylambda closures
     if (name.contains("apply") || name.contains("$anonfun$")) {
-      // A method with suffix "$adapted" will be generated in cases like
-      // { _:Int => return; Seq()} but not { _:Int => return; true}
-      // closure passed is $anonfun$t$1$adapted while actual code resides in $anonfun$s$1
-      // visitor will see only $anonfun$s$1$adapted, so we remove the suffix, see
-      // https://github.com/scala/scala-dev/issues/109
-      val isTargetMethod = targetMethodName.isEmpty ||
-        name == targetMethodName.get || name == targetMethodName.get.stripSuffix("$adapted")
-
       new MethodVisitor(Opcodes.ASM9) {
         override def visitTypeInsn(op: Int, tp: String): Unit = {
-          if (op == Opcodes.NEW && tp.contains("scala/runtime/NonLocalReturnControl") &&
-              isTargetMethod) {
-            throw new ReturnStatementInClosureException
+          if (op == Opcodes.NEW && tp.contains("scala/runtime/NonLocalReturnControl")) {
+            found += name
           }
         }
       }

core/src/test/scala/org/apache/spark/util/ClosureCleanerSuite.scala

@@ -60,11 +60,40 @@ class ClosureCleanerSuite extends SparkFunSuite {
     }
   }
 
+  test("return statements in closures capturing a null value are identified at cleaning time") {
+    intercept[ReturnStatementInClosureException] {
+      TestObjectWithBogusReturnsAndNullCapture.run()
+    }
+  }
+
   test("return statements from named functions nested in closures don't raise exceptions") {
     val result = TestObjectWithNestedReturns.run()
     assert(result === 1)
   }
 
+  test("hasReturnStatement identifies non-local returns per method") {
+    TestObjectWithReturnInClosure.run()
+    val cls = TestObjectWithReturnInClosure.getClass
+    val implMethodName = {
+      val proxy =
+        IndylambdaScalaClosures.getSerializationProxy(TestObjectWithReturnInClosure.lastClosure)
+      assert(proxy.isDefined)
+      proxy.get.getImplMethodName
+    }
+    // Any-method query.
+    assert(ClosureCleaner.hasReturnStatement(cls, None))
+    // Targeted query with the exact impl method name.
+    assert(ClosureCleaner.hasReturnStatement(cls, Some(implMethodName)))
+    // An "$adapted" wrapper name resolves to the underlying method that holds the closure body
+    // (see https://github.com/scala/scala-dev/issues/109).
+    assert(ClosureCleaner.hasReturnStatement(cls, Some(implMethodName + "$adapted")))
+    // A method that does not exist on the class must not match.
+    assert(!ClosureCleaner.hasReturnStatement(cls, Some("$anonfun$doesNotExist$1")))
+    // A class whose closures contain no non-local returns.
+    TestObjectWithoutReturnInClosure.run()
+    assert(!ClosureCleaner.hasReturnStatement(TestObjectWithoutReturnInClosure.getClass, None))
+  }
+
   test("user provided closures are actually cleaned") {
 
     // We use return statements as an indication that a closure is actually being cleaned
@@ -199,6 +228,40 @@ object TestObjectWithBogusReturns {
   }
 }
 
+object TestObjectWithBogusReturnsAndNullCapture {
+  def run(): Int = {
+    withSpark(new SparkContext("local", "test")) { sc =>
+      val nums = sc.parallelize(Array(1, 2, 3, 4).toImmutableArraySeq)
+      val s: String = null
+      // The closure's first captured argument may be the (null) `s` rather than the non-local
+      // return's key: the cleaner must still detect the invalid return rather than bail out on
+      // the null capture.
+      nums.map { x => if (s != null) return 1; x * 2 }
+      1
+    }
+  }
+}
+
+object TestObjectWithReturnInClosure {
+  // The non-local `return` forces the enclosing method to have type Int, so it cannot return the
+  // closure to the caller directly; stash it for the test to inspect instead.
+  var lastClosure: Int => Int = null
+  def run(): Int = {
+    val f = (x: Int) => { if (x < 0) return -1; x }
+    lastClosure = f
+    f(1)
+  }
+}
+
+object TestObjectWithoutReturnInClosure {
+  var lastClosure: Int => Int = null
+  def run(): Int = {
+    val f = (x: Int) => x * 2
+    lastClosure = f
+    f(1)
+  }
+}
+
 object TestObjectWithNestedReturns {
   def run(): Int = {
     withSpark(new SparkContext("local", "test")) { sc =>