viirya commented on code in PR #59332:
URL: https://github.com/apache/spark/pull/59332#discussion_r4262579640
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -169,37 +196,113 @@ private[spark] class TaskContextImpl(
}
private[spark] override def markTaskFailed(error: Throwable): Unit = {
+ markTaskFailed(error, Claim)
+ }
+
+ /**
+ * @param mode `Nested` if this thread is already invoking listeners, i.e. a
listener failed. The
+ * failure listeners then run as part of that invocation, which
this thread keeps.
+ */
+ private def markTaskFailed(error: Throwable, mode: InvocationMode): Unit = {
synchronized {
if (failureCauseOpt.isDefined) return
failureCauseOpt = Some(error)
}
- invokeTaskFailureListeners(error)
+ invokeTaskFailureListeners(error, mode)
}
private[spark] override def markTaskCompleted(error: Option[Throwable]):
Unit = {
- synchronized {
+ // The completion listeners always run on this thread, and this method
returns only after they
Review Comment:
Agreed. I moved it out in 00da9fb69d4: the handover wait in
`markTaskCompleted`, running failure listeners inside the current invocation,
running the listeners left to the invoking thread, and `InvocationMode` now go
to SPARK-60222 as a separate `[CORE]` PR. This PR keeps only what the
use-after-free fix needs from `TaskContextImpl`: `addTaskPreCompletionListener`
and the claim in `markTaskCompleted` that is set together with `completed` (as
in 038c4b3420e), plus the shared add helper.
For the cases that the follow-up closes (another thread holding the
invocation at completion, or the brief release while failure listeners run),
the stop listener keeps the same-thread guard here. If it runs on the writer
thread, it does not wait and logs an error, and the use-after-free from before
this PR remains possible on that path. The comment and the PR description now
say so and point to SPARK-60222. I also updated the PR description, which no
longer lists the moved changes.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -169,37 +196,113 @@ private[spark] class TaskContextImpl(
}
private[spark] override def markTaskFailed(error: Throwable): Unit = {
+ markTaskFailed(error, Claim)
+ }
+
+ /**
+ * @param mode `Nested` if this thread is already invoking listeners, i.e. a
listener failed. The
+ * failure listeners then run as part of that invocation, which
this thread keeps.
+ */
+ private def markTaskFailed(error: Throwable, mode: InvocationMode): Unit = {
synchronized {
if (failureCauseOpt.isDefined) return
failureCauseOpt = Some(error)
}
- invokeTaskFailureListeners(error)
+ invokeTaskFailureListeners(error, mode)
}
private[spark] override def markTaskCompleted(error: Option[Throwable]):
Unit = {
- synchronized {
+ // The completion listeners always run on this thread, and this method
returns only after they
+ // have run. If any other thread could run them, it could be one the
listeners are meant to
+ // stop and wait for, e.g. a pipelined Python UDF writer that adds
listeners lazily while it
+ // pulls its input, which would then wait for itself. And if this method
returned before they
+ // ran, the task could release its memory while that writer still uses it.
So the listener
+ // invocation is claimed together with setting `completed`; if another
thread holds it (e.g. a
+ // kill thread running failure listeners), that thread hands it over to
this one when done.
+ //
+ // That wait has no time bound: if the other thread's listeners block, so
does this task. It
+ // logs a warning with that thread's stack trace periodically. If such a
task is killed,
+ // `spark.task.reaper.enabled` with a positive
`spark.task.reaper.killTimeout` bounds how long
+ // it may stay in this state, by killing the executor JVM, but only if the
blocked listeners
+ // do not run within the reaper's own kill call: the reaper calls
`Task.kill` before it starts
+ // checking the timeout, and that call runs the failure listeners on the
reaper thread if an
+ // interrupt listener throws. Then only the periodic warning shows the
hang.
+ val mode = synchronized {
if (completed) return
completed = true
+ val currentThread = Thread.currentThread()
+ if (listenerInvocationThread.isEmpty) {
+ listenerInvocationThread = Some(currentThread)
+ Claimed
+ } else if (listenerInvocationThread.contains(currentThread)) {
+ // Called from a listener on this thread: run the completion listeners
as part of that
+ // invocation.
+ Nested
+ } else {
+ nextListenerInvocationThread = Some(currentThread)
+ awaitListenerInvocationHandover(currentThread)
Review Comment:
Thanks for tracking this down in the CI logs, and sorry about it. I had run
only the suites around the change. With the split (see L215), the handover wait
is no longer in this PR, so `TaskContextImpl` is back to what it was in
038c4b3420e, whose core job passed. On 00da9fb69d4 I ran
`StreamingShuffleSuite` and the other core suites that use `markTaskFailed`,
`markTaskCompleted` or `markInterrupted` locally, and they pass. The follow-up
will update that test to the new behavior along the lines you describe.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -320,21 +430,44 @@ private[spark] class TaskContextImpl(
// [[markTaskFailed]] here will have no effect and we simply
resume running the
// remaining completion handlers.
try {
- listenerInvocationThread = None
- markTaskFailed(e)
+ markTaskFailed(e, Nested)
} catch {
case t: Throwable => e.addSuppressed(t)
- } finally {
- synchronized {
- if (listenerInvocationThread.isEmpty) {
- listenerInvocationThread = Some(Thread.currentThread())
- }
- }
}
listenerExceptions += e
logError(log"Error in ${MDC(LISTENER, name)}", e)
}
}
+ if (mode != Nested) {
+ // We have executed all listeners of `listenerStacks` that have been
added so far. Run the
+ // listeners left to this thread (see above), then deregister this
thread as the callback
+ // invocation thread, or hand the invocation over to a thread waiting
for it.
+ var released = false
+ while (!released) {
+ val pendingListeners: Option[() => Unit] = synchronized {
+ val pendingFailure = failureCauseOpt.filter(_ =>
!onFailureCallbacks.empty())
Review Comment:
This code moved to SPARK-60222. I agree that the failure listeners left to
the completing thread should run before its remaining completion listeners,
including on the handover path, and I'll do that there rather than only
document the order.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -169,37 +196,113 @@ private[spark] class TaskContextImpl(
}
private[spark] override def markTaskFailed(error: Throwable): Unit = {
+ markTaskFailed(error, Claim)
+ }
+
+ /**
+ * @param mode `Nested` if this thread is already invoking listeners, i.e. a
listener failed. The
+ * failure listeners then run as part of that invocation, which
this thread keeps.
+ */
+ private def markTaskFailed(error: Throwable, mode: InvocationMode): Unit = {
synchronized {
if (failureCauseOpt.isDefined) return
failureCauseOpt = Some(error)
}
- invokeTaskFailureListeners(error)
+ invokeTaskFailureListeners(error, mode)
}
private[spark] override def markTaskCompleted(error: Option[Throwable]):
Unit = {
- synchronized {
+ // The completion listeners always run on this thread, and this method
returns only after they
+ // have run. If any other thread could run them, it could be one the
listeners are meant to
+ // stop and wait for, e.g. a pipelined Python UDF writer that adds
listeners lazily while it
+ // pulls its input, which would then wait for itself. And if this method
returned before they
+ // ran, the task could release its memory while that writer still uses it.
So the listener
+ // invocation is claimed together with setting `completed`; if another
thread holds it (e.g. a
+ // kill thread running failure listeners), that thread hands it over to
this one when done.
+ //
+ // That wait has no time bound: if the other thread's listeners block, so
does this task. It
+ // logs a warning with that thread's stack trace periodically. If such a
task is killed,
+ // `spark.task.reaper.enabled` with a positive
`spark.task.reaper.killTimeout` bounds how long
+ // it may stay in this state, by killing the executor JVM, but only if the
blocked listeners
+ // do not run within the reaper's own kill call: the reaper calls
`Task.kill` before it starts
+ // checking the timeout, and that call runs the failure listeners on the
reaper thread if an
+ // interrupt listener throws. Then only the periodic warning shows the
hang.
+ val mode = synchronized {
if (completed) return
completed = true
+ val currentThread = Thread.currentThread()
+ if (listenerInvocationThread.isEmpty) {
+ listenerInvocationThread = Some(currentThread)
+ Claimed
+ } else if (listenerInvocationThread.contains(currentThread)) {
+ // Called from a listener on this thread: run the completion listeners
as part of that
+ // invocation.
+ Nested
+ } else {
+ nextListenerInvocationThread = Some(currentThread)
+ awaitListenerInvocationHandover(currentThread)
+ Claimed
+ }
+ }
+ invokeTaskCompletionListeners(error, mode)
+ }
+
+ /**
+ * Waits until another thread invoking listeners hands the invocation over
to `currentThread`.
+ * Interrupts do not stop the wait and are restored afterwards. Must hold
the TaskContext lock.
+ */
+ private def awaitListenerInvocationHandover(currentThread: Thread): Unit = {
+ val startNs = System.nanoTime()
+ var nextWarningNs = startNs +
TimeUnit.MILLISECONDS.toNanos(listenerHandoverWarnIntervalMs)
+ var interrupted = false
+ while (!listenerInvocationThread.contains(currentThread)) {
+ try {
+ wait(math.max(1L, TimeUnit.NANOSECONDS.toMillis(nextWarningNs -
System.nanoTime())))
+ } catch {
+ case _: InterruptedException => interrupted = true
+ }
+ val nowNs = System.nanoTime()
+ if (nowNs >= nextWarningNs &&
!listenerInvocationThread.contains(currentThread)) {
+ listenerInvocationThread.foreach { holder =>
+ val holderStack = new Exception(s"Stack trace of thread
${holder.getName}")
+ holderStack.setStackTrace(holder.getStackTrace)
+ logWarning(log"Task ${MDC(TASK_ATTEMPT_ID, taskAttemptId)} has been
waiting for " +
+ log"${MDC(TOTAL_TIME, TimeUnit.NANOSECONDS.toMillis(nowNs -
startNs))} ms to run " +
+ log"its completion listeners until thread ${MDC(THREAD_NAME,
holder.getName)} " +
+ log"finishes running its task listeners.", holderStack)
+ }
+ nextWarningNs = nowNs +
TimeUnit.MILLISECONDS.toNanos(listenerHandoverWarnIntervalMs)
+ }
+ }
+ if (interrupted) {
Review Comment:
Moved to SPARK-60222 with the handover wait. I'll make the two waits follow
one policy there, or explain why they differ.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -320,21 +430,44 @@ private[spark] class TaskContextImpl(
// [[markTaskFailed]] here will have no effect and we simply
resume running the
// remaining completion handlers.
try {
- listenerInvocationThread = None
- markTaskFailed(e)
+ markTaskFailed(e, Nested)
} catch {
case t: Throwable => e.addSuppressed(t)
- } finally {
- synchronized {
- if (listenerInvocationThread.isEmpty) {
- listenerInvocationThread = Some(Thread.currentThread())
- }
- }
}
listenerExceptions += e
logError(log"Error in ${MDC(LISTENER, name)}", e)
}
}
+ if (mode != Nested) {
Review Comment:
Moved to SPARK-60222. I'll put the release and the handover in a `finally`
there, so that a stale holder or a stale waiter cannot be left behind.
--
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]
