JoshRosen opened a new pull request, #57710:
URL: https://github.com/apache/spark/pull/57710
### 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.
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 `ClosureCleaner` only modifies closures for Scala REPL and Ammonite
capturing classes; for compiled code the indylambda path is a validator. This
change removes neither validation nor cleaning -- only the redundancy.
### 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.
* 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
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
