ivoson commented on code in PR #58518:
URL: https://github.com/apache/spark/pull/58518#discussion_r4002275801
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/SQLExecution.scala:
##########
@@ -252,67 +309,69 @@ object SQLExecution extends Logging {
ex = Some(e)
throw e
} finally {
- val endTime = System.nanoTime()
- val errorMessage = ex.map {
- case e: SparkThrowable =>
- SparkThrowableHelper.getMessage(e, ErrorMessageFormat.PRETTY)
- case e =>
- Utils.exceptionString(e)
- }
- if (queryExecution.shuffleCleanupMode != DoNotCleanup
- && isExecutedPlanAvailable) {
- val shuffleIds = queryExecution.executedPlan match {
- case command: V2CommandExec =>
- command.children.flatMap(extractShuffleIds)
- case dataWritingCommand: DataWritingCommandExec =>
- extractShuffleIds(dataWritingCommand.child)
- case plan =>
- extractShuffleIds(plan)
- }
- shuffleIds.foreach { shuffleId =>
- queryExecution.shuffleCleanupMode match {
- case RemoveShuffleFiles =>
- // Same as what we do in
ContextCleaner.doCleanupShuffle, but do not
- // unregister the shuffle on MapOutputTracker, so that
stage retries would be
- // triggered.
- // Set blocking to Utils.isTesting to deflake unit tests.
- sc.shuffleDriverComponents.removeShuffle(shuffleId,
Utils.isTesting)
- case SkipMigration =>
-
SparkEnv.get.blockManager.migratableResolver.addShuffleToSkip(shuffleId)
- case _ => // this should not happen
+ // `SparkContext.stop()` nulls `dagScheduler` before it stops
the listener bus, so the
+ // end event may still be posted after the scheduler is
unavailable. Keep observation
+ // completion in a `finally` so an error in this block never
leaves a waiter hung.
+ try {
+ val endTime = System.nanoTime()
+ val errorMessage = ex.map { e =>
+ try {
+ e match {
+ case st: SparkThrowable =>
+ SparkThrowableHelper.getMessage(st,
ErrorMessageFormat.PRETTY)
+ case _ =>
+ Utils.exceptionString(e)
+ }
+ } catch {
+ // Rendering a user throwable can itself throw (e.g. a
custom `getMessage`).
+ // Fall back to a safe value so the query's real failure
is still surfaced and
+ // the cleanup, event post and observation completion
below still run.
+ case NonFatal(t) =>
+ logWarning(log"Failed to render the error message for
execution " +
+ log"${MDC(EXECUTION_ID, executionId)}.", t)
+ e.getClass.getName
}
}
- }
- val event = SparkListenerSQLExecutionEnd(
- executionId,
- System.currentTimeMillis(),
- // Use empty string to indicate no error, as None may mean
events generated by old
- // versions of Spark.
- errorMessage.orElse(Some("")),
- Some(queryId))
- // Currently only `Dataset.withAction` and
`DataFrameWriter.runCommand` specify the
- // `name` parameter. The `ExecutionListenerManager` only watches
SQL executions with
- // name. We can specify the execution name in more places in the
future, so that
- // `QueryExecutionListener` can track more cases.
- event.executionName = name
- event.duration = endTime - startTime
- event.qe = queryExecution
- event.executionFailure = ex
- if (Utils.isTesting) {
- import scala.jdk.CollectionConverters._
- event.jobIds =
Option(sc.dagScheduler.activeQueryToJobs.get(executionId))
- .map(_.asScala.map(_.jobId).toSet)
- .getOrElse(Set.empty)
- }
-
- // Clean up jobs tracked by DAGScheduler for this query
execution.
- sc.dagScheduler.cleanupQueryJobs(executionId)
+ if (queryExecution.shuffleCleanupMode != DoNotCleanup &&
isExecutedPlanAvailable) {
+ cleanupShuffleDependencies(queryExecution, executionId)
+ }
+ val event = SparkListenerSQLExecutionEnd(
+ executionId,
+ System.currentTimeMillis(),
+ // Use empty string to indicate no error, as None may mean
events generated by old
+ // versions of Spark.
+ errorMessage.orElse(Some("")),
+ Some(queryId))
+ // Currently only `Dataset.withAction` and
`DataFrameWriter.runCommand` specify the
+ // `name` parameter. The `ExecutionListenerManager` only
watches SQL executions with
+ // name. We can specify the execution name in more places in
the future, so that
+ // `QueryExecutionListener` can track more cases.
+ event.executionName = name
+ event.duration = endTime - startTime
+ event.qe = queryExecution
+ event.executionFailure = ex
+ // Snapshot the `@volatile` `dagScheduler` once and share it
across both reads
+ // below; it is null once `SparkContext.stop()` has run.
+ val dagSchedulerOpt = Option(sc.dagScheduler)
+ if (Utils.isTesting) {
+ import scala.jdk.CollectionConverters._
+ // Only runs under `Utils.isTesting`; hits the same teardown
race as the job
+ // cleanup below.
+ event.jobIds = dagSchedulerOpt
+ .flatMap(ds =>
Option(ds.activeQueryToJobs.get(executionId)))
+ .map(_.asScala.map(_.jobId).toSet)
+ .getOrElse(Set.empty)
+ }
- sc.listenerBus.post(event)
+ // Clean up jobs tracked by DAGScheduler for this query
execution.
+ dagSchedulerOpt.foreach(_.cleanupQueryJobs(executionId))
- // Observation.tryComplete is called here to ensure the
observation is completed,
- // but it is not high priority, so it is fine to call it later.
- sparkSession.observationManager.tryComplete(queryExecution)
+ sc.listenerBus.post(event)
+ } finally {
+ // Complete the observation whatever the block above threw, so
an `Observation.get`
+ // waiter is never left hung. `promise.tryComplete` is
idempotent.
+ sparkSession.observationManager.tryComplete(queryExecution)
Review Comment:
Good point -- confirmed: `ObservationManager.tryComplete` only `Try`-wraps
`qe.observedMetrics`; the `foreachWithSubqueriesAndPruning` walk and
`setMetricsAndNotify` aren't, so a NonFatal throw there would escape the inner
`finally` and mask the body's `ex`. Latent as you note (the walk is over an
already-built plan and `promise.tryComplete` is idempotent), but I wrapped the
call in `catch { case NonFatal(e) => logWarning(...) }` carrying the execution
id, to line it up with the error-rendering fallback above. Kept it at the call
site since this is the only teardown-time caller.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/SQLExecutionSuite.scala:
##########
@@ -423,6 +427,261 @@ 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
+ }
+ }
+
+ private def withUnavailableSparkEnv[T](f: => T): T = {
+ val savedEnv = SparkEnv.get
+ SparkEnv.set(null)
Review Comment:
Thanks -- fixed by nulling `SparkEnv` around only the narrow call. The test
now calls `SQLExecution.cleanupShuffleDependencies(qe, ...)` directly instead
of going through `withNewExecutionId`, so no execution-end event is posted and
no async listener-bus thread can observe the process-wide null `SparkEnv`
(`cleanupShuffleDependencies` is now `private[execution]` for this). The full
`withNewExecutionId` cleanup path stays covered by the other tests; this one is
now a focused check of the SkipMigration branch's mode-specific warning.
--
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]
