ivoson commented on code in PR #58518:
URL: https://github.com/apache/spark/pull/58518#discussion_r3965948857
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala:
##########
@@ -423,6 +425,120 @@ class SQLExecutionSuite extends SparkFunSuite with
SQLConfHelper {
spark.stop()
}
}
+
+ /**
+ * Runs `f` with `spark`'s `dagScheduler` nulled out, standing in for a
`SparkContext` that has
+ * already been stopped. `SparkContext.stop()` nulls `_dagScheduler` before
it stops the listener
+ * bus, so a query really can unwind through `withNewExecutionId`'s
`finally` in this state.
+ */
+ private def withStoppedDagScheduler[T](spark: SparkSession)(f: => T): T = {
+ val sc = spark.sparkContext
+ val savedDagScheduler = sc.dagScheduler
+ sc.dagScheduler = null
+ try {
+ f
+ } finally {
+ sc.dagScheduler = savedDagScheduler
+ }
+ }
+
+ /**
+ * Runs `f` with `spark`'s `BlockManagerMaster.driverEndpoint` nulled out,
standing in for a
+ * `SparkContext` whose `SparkEnv` has been stopped. `SparkContext.stop()`
stops `SparkEnv` (which
+ * nulls that endpoint) after nulling `dagScheduler`, so the shuffle cleanup
in
+ * `withNewExecutionId`'s `finally` -- which runs before the `dagScheduler`
cleanup -- really can
+ * hit a stopped `BlockManagerMaster` and NPE while a query unwinds during
teardown.
+ */
+ private def withStoppedBlockManagerMaster[T](spark: SparkSession)(f: => T):
T = {
+ val master = spark.sparkContext.env.blockManager.master
+ val savedEndpoint = master.driverEndpoint
+ master.driverEndpoint = null
+ try {
+ f
+ } finally {
+ master.driverEndpoint = savedEndpoint
+ }
+ }
+
+ test("SPARK-59242: withNewExecutionId surfaces the body's failure when the
SparkContext " +
+ "has been stopped") {
+ val spark =
SparkSession.builder().master("local[*]").appName("test").getOrCreate()
+ try {
+ val qe = spark.range(1, 10).queryExecution
+ val bodyFailure = new IllegalStateException("body failed")
+ // Without the null guards, the `finally` dereferences `dagScheduler`
while it is null:
+ // under `Utils.isTesting` the `activeQueryToJobs` read throws first,
and the
+ // `cleanupQueryJobs` cleanup would do the same. Because the NPE is
thrown from a `finally`,
+ // it replaces `bodyFailure` entirely -- destroying the only record of
why the query failed.
+ val thrown = intercept[IllegalStateException] {
+ withStoppedDagScheduler(spark) {
+ SQLExecution.withNewExecutionId(qe) {
+ throw bodyFailure
+ }
+ }
+ }
+ assert(thrown eq bodyFailure)
+ } finally {
+ spark.stop()
+ }
+ }
+
+ test("SPARK-59242: withNewExecutionId completes normally when the
SparkContext has been " +
+ "stopped") {
+ val spark =
SparkSession.builder().master("local[*]").appName("test").getOrCreate()
+ try {
+ // Attach an observation to pin down that it is completed. `tryComplete`
runs at the end of
+ // the `finally`, after the guarded cleanup, so a cleanup failure would
skip it and leave the
+ // observation uncompleted. Assert on `future.isCompleted`
(non-blocking) so a regression
+ // fails fast, rather than calling `get`, which would block until the
suite timeout.
+ val observation = new Observation("obs")
+ val df = spark.range(1, 10).observe(observation, count(lit(1)).as("cnt"))
+ val qe = df.queryExecution
+ withStoppedDagScheduler(spark) {
+ assert(SQLExecution.withNewExecutionId(qe)("result") === "result")
+ }
+ assert(observation.future.isCompleted)
+ } finally {
+ spark.stop()
+ }
+ }
+
+ test("SPARK-59242: withNewExecutionId tolerates shuffle cleanup failing when
the " +
+ "SparkContext is stopping") {
+ val spark =
SparkSession.builder().master("local[*]").appName("test").getOrCreate()
+ try {
+ // Disable AQE so the shuffle id is materialized from the plan below
without running a job;
+ // the `finally`'s shuffle cleanup then actually calls `removeShuffle`
for it.
+ spark.conf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, false)
Review Comment:
Done. Switched to `withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key ->
"false", ...)` spanning from `val observation` through the last assert (with
the `fileCleanup` conf in the same block), so AQE is off at planning time and
nothing leaks to a shared session.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala:
##########
@@ -298,15 +305,23 @@ object SQLExecution extends Logging {
event.duration = endTime - startTime
event.qe = queryExecution
event.executionFailure = ex
+ // Snapshot the `@volatile` `dagScheduler` once and share it
across both reads below.
+ // `SparkContext.stop()` nulls `dagScheduler` before it stops
the listener bus, so a
+ // query unwinding here while the context tears down would
otherwise NPE. As this runs
+ // in a `finally`, that NPE would replace the query's real
failure and skip the event
+ // post and observation completion below, so tolerate an
already-stopped context.
+ val dagSchedulerOpt = Option(sc.dagScheduler)
Review Comment:
Done. Trimmed the repeated teardown explanation -- the observation-`finally`
comment now carries only the non-repeated fact (`SparkContext.stop()` nulls
`dagScheduler` before stopping the listener bus), and the helper scaladocs/test
comments no longer restate the full mechanism. The two cleanups are named
distinctly (`cleanupShuffleDependencies` vs `cleanupQueryJobs`), so "the
cleanup" is unambiguous.
--
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]
