dongjoon-hyun commented on code in PR #59332:
URL: https://github.com/apache/spark/pull/59332#discussion_r4256073416
##########
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:
**[Low] An interrupt that arrives during this wait is passed on to the
completion listeners, while `awaitExit` drops one for that reason.**
This restores the flag before the completion listeners run.
`stopAndAwaitExit` then sees it as pending on entry and restores it again, so
the ordinary completion listeners (the `HybridRowQueue` close, `worker.stop()`,
the upstream reader's close) run with the flag set.
`PipelinedWriterTask.awaitExit` (PythonRunner L333-336) drops an interrupt that
arrives during its own wait, because the flag "would make the blocking channel
and lock operations of the completion listeners that run next fail and skip
cleanup". So an interrupt that lands while the task thread waits here is passed
on, and the same interrupt landing a moment later, during the writer wait, is
dropped.
This wait can be interrupted, e.g. by a second kill with `interruptThread =
true`, or by the writer thread resolving a pending interrupt in
`createResourceUninterruptibly`, which interrupts the task thread. Could both
waits follow one policy, or could the comments say 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) {
+ // 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:
**[Medium] Failure listeners left to this thread run after every completion
listener.**
Suppose another thread fails the task while this thread runs the completion
listeners, e.g. a kill whose interrupt listener throws. Its `markTaskFailed`
finds the invocation taken and returns. This thread then first runs the
remaining completion listeners, which close streams, release state stores and
free the `HybridRowQueue` pages, and only after that the failure listeners,
here. `runTaskWithListeners` always runs the failure listeners before the
completion listeners, and listeners may rely on that. For example,
`RocksDBStateStoreProvider` registers a failure listener `if (!hasCommitted)
abort()` next to a completion listener that aborts the store and releases its
stamp, so the failure listener may now call `abort()` after the stamp has been
released.
The same happens when the release below hands the invocation to a thread
waiting in `markTaskCompleted` while failure listeners are still pending: that
thread runs all completion listeners first.
Exceptions from these failure listeners are also added to
`listenerExceptions` (L466) and thrown from `markTaskCompleted` on the task
thread. Before, they stayed on the thread that called `markTaskFailed`.
Running them late is still better than dropping them, as before. Could
pending failure listeners run first instead, i.e. before the handover below and
before the next completion listener? Otherwise, could the order be documented
here and in `addTaskFailureListener`?
##########
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:
**[Low] The release is not in a `finally`, and a stale holder now blocks
`markTaskCompleted` forever.**
If a `Throwable` escapes the loop above, this thread never gets here, and
`listenerInvocationThread` keeps pointing at it. The catch block can throw
itself: if a listener fails with `OutOfMemoryError`, `listenerExceptions += e`
or `logError` (L437-438) can fail again. Before this PR, a stale holder only
made `markTaskCompleted` return without running the completion listeners. Now
`markTaskCompleted` waits in `awaitListenerInvocationHandover` for a handover
that never comes, and interrupts do not end that wait.
The reverse case also exists: if `awaitListenerInvocationHandover` throws,
e.g. from the `logWarning` at L269, `nextListenerInvocationThread` stays set.
The holder later hands the invocation to a thread that has already left, and
every later listener is skipped.
Could the main loop be wrapped in a `try`/`finally`, so that the release, or
at least the handover, always runs?
##########
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())
+ if (nextListenerInvocationThread.isEmpty &&
pendingFailure.isDefined) {
+ Some(() => invokeTaskFailureListeners(pendingFailure.get, Nested))
+ } else if (nextListenerInvocationThread.isEmpty && completed &&
+ (!onPreCompleteCallbacks.empty() ||
!onCompleteCallbacks.empty())) {
+ Some(() => invokeTaskCompletionListeners(None, Nested))
Review Comment:
**[Low] Leftover completion listeners can run on a kill thread under the
global `TaskContext` lock.**
This runs the pending completion listeners on whichever thread holds the
invocation, not only on the completing thread. For example, after the task
thread has finished `markTaskCompleted`, a late kill whose interrupt listener
throws runs the failure listeners on the kill thread, inside `Task.kill`'s
JVM-wide `TaskContext.synchronized` block (`Task.scala` L235). A completion
listener that another thread adds at that point is left to the kill thread,
which runs it here under that lock. A slow close, e.g. of an S3 or JDBC
resource, then also blocks `Task.kill`, `interruptible()` and
`createResourceUninterruptibly` for every task on the executor. The new test
"completion listeners queued meanwhile run before failure invoker releases"
asserts this behavior on purpose.
This is rare, but the comment at L215, "The completion listeners always run
on this thread", now holds only for the listeners that are pending when
`markTaskCompleted` is called. Could it say so, and mention that the leftover
ones run on the thread that holds the invocation?
##########
core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala:
##########
@@ -757,6 +760,254 @@ class TaskContextSuite extends SparkFunSuite with
BeforeAndAfter with LocalSpark
assert(invocationOrder === Seq("C", "B", "A", "D"))
}
+ test("SPARK-60113: pre-completion listeners run before all completion
listeners") {
+ val context = TaskContext.empty()
+ val invocationOrder = ArrayBuffer.empty[String]
+
+ def makeLoggingListener(id: String): TaskCompletionListener = new
TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = {
+ invocationOrder += id
+ }
+ }
+ context.addTaskCompletionListener(makeLoggingListener("A"))
+ context.addTaskPreCompletionListener(makeLoggingListener("P1"))
+ context.addTaskPreCompletionListener(makeLoggingListener("P2"))
+ // Added after the pre-completion listeners, e.g. lazily by the thread
they stop. Had P1 and
+ // P2 been added with addTaskCompletionListener, B would run before them.
+ context.addTaskCompletionListener(makeLoggingListener("B"))
+
+ context.markTaskCompleted(None)
+ assert(invocationOrder === Seq("P2", "P1", "B", "A"))
+ // Called immediately once the task has completed, like completion
listeners.
+ context.addTaskPreCompletionListener(makeLoggingListener("P3"))
+ assert(invocationOrder === Seq("P2", "P1", "B", "A", "P3"))
+ }
+
+ test("SPARK-60113: completion listeners run even if a pre-completion
listener fails") {
+ val context = TaskContext.empty()
+ val listener = mock(classOf[TaskCompletionListener])
+ context.addTaskCompletionListener(listener)
+ context.addTaskPreCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = throw new
Exception("pre")
+ })
+
+ val e = intercept[TaskCompletionListenerException] {
+ context.markTaskCompleted(None)
+ }
+ assert(e.getMessage.contains("pre"))
+ verify(listener, times(1)).onTaskCompletion(any())
+ assert(context.isFailed())
+ }
+
+ test("SPARK-60113: failures of pre-completion and completion listeners are
both reported") {
+ val context = TaskContext.empty()
+ context.addTaskCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = throw new
Exception("post")
+ })
+ context.addTaskPreCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = throw new
Exception("pre")
+ })
+
+ val e = intercept[TaskCompletionListenerException] {
+ context.markTaskCompleted(None)
+ }
+ assert(e.getMessage.contains("pre"))
+ assert(e.getMessage.contains("post"))
+ }
+
+ test("SPARK-60113: a pre-completion listener added during completion runs
next") {
+ val context = TaskContext.empty()
+ val invocationOrder = ArrayBuffer.empty[String]
+
+ def makeLoggingListener(id: String): TaskCompletionListener = new
TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = {
+ invocationOrder += id
+ }
+ }
+ context.addTaskCompletionListener(makeLoggingListener("A"))
+ context.addTaskCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit = {
+ invocationOrder += "B"
+ // Added from a second thread while this thread runs the completion
listeners, so the
+ // call returns without running it; this thread must still run it.
+ val thread = new Thread(() => {
+ context.addTaskPreCompletionListener(makeLoggingListener("P"))
+ })
+ thread.start()
+ thread.join()
Review Comment:
**[Nit] These tests join helper threads without a timeout, some from inside
a listener that holds the invocation.**
Here and in "failure listeners queued during completion run before it
releases" (`failer.join()`, L970), the listener blocks in `join()` while its
thread holds the listener invocation. "completion waits for another thread
invoking listeners and runs them" has the same `adder.join()` (L945), and none
of these helper threads is a daemon thread. If a later change makes the helper
wait for the invocation, as `markTaskCompleted` now does, the test hangs until
the core job hits the 2h30m limit, as the `StreamingShuffleSuite` test does
now, instead of failing. Could they use `join(timeout)` with
`assert(!thread.isAlive)` and `setDaemon(true)`, like the other new tests?
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -426,3 +559,21 @@ private[spark] class TaskContextImpl(
}
}
}
+
+private object TaskContextImpl {
+
+ /** How an `invokeListeners()` call takes part in the listener invocation. */
+ sealed trait InvocationMode
+
+ /** Registers this thread as `listenerInvocationThread` unless another
thread is registered. */
+ case object Claim extends InvocationMode
+
+ /** The caller has already registered this thread as
`listenerInvocationThread`. */
+ case object Claimed extends InvocationMode
+
+ /** An outer invocation on this thread has registered it, and keeps it when
this one is done. */
+ case object Nested extends InvocationMode
+
+ /** How often to log while `markTaskCompleted()` waits for another thread
invoking listeners. */
+ val listenerHandoverWarnIntervalMs = 10000L
Review Comment:
**[Nit] The handover warning has no test, and this interval cannot be
shortened for one.**
The warning is the only sign that `markTaskCompleted` is stuck behind
another thread, but it fires only after 10 s, and none of the new tests wait
that long. A regression there, e.g. the wrong holder named or no stack
attached, would go unnoticed. `PipelinedWriterTask` takes `warnIntervalMs` as a
parameter, and `BasePythonRunnerSuite` sets it to 100 ms to exercise that path.
Could this interval be configurable in the same way, with a test that checks
that the warning names the holder and carries its stack?
##########
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)
Review Comment:
**[Nit] The warning walks the holder's stack and logs it every 10 s while
holding the `TaskContextImpl` lock.**
The holder needs this lock to hand over (L455-458), and `isCompleted()`,
`isFailed()` and every `add*Listener` call on other threads need it too, so
each warning stalls them behind a stack walk and a synchronous log write. A
task that stays stuck for an hour also logs 360 full stack traces as
`java.lang.Exception: Stack trace of thread ...`, which log scanners tend to
count as errors. Could the stack be attached only to the first warning, or
could the warnings come at growing intervals?
##########
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:
**[Discussion] The `TaskContextImpl` part could be its own `[CORE]` change.**
The rework here (the handover in `markTaskCompleted`, the release loop, and
the three invocation modes) applies to every task, not only to the opt-in
pipelined mode:
- `markTaskCompleted` can now block indefinitely on another thread.
- Failure-listener exceptions from other threads can now fail the task on
the task thread.
- Failure and completion listeners that used to be dropped now run, possibly
late and on other threads.
- An existing core test deadlocks (see the comment on L243).
The PR description answers "No" to the user-facing change question because
the pipelined mode is unreleased, but this part is not behind that flag. A
separate JIRA and PR for it would get its own review and tests, and a
regression in it would not force reverting the use-after-free fix together with
it. What do you think?
##########
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:
**[High] This wait deadlocks an existing `StreamingShuffleSuite` test, and
the core CI job times out on it.**
`StreamingShuffleSuite`'s "background markTaskFailed races with task thread
markTaskCompleted" (L1380-1433) has a background thread call `markTaskFailed`
and block in a failure listener on `failureListenerProceed`. The test thread
then calls `markTaskCompleted(Some(error))` (L1419). Before this PR, that call
returned right away. Now it waits here for the background thread to hand over
the invocation, but the test releases the failure listener only at L1422, on
the same thread, after `markTaskCompleted` returns. Neither thread can make
progress, and since this wait ignores interrupts, the test framework cannot
break it either.
The core module job shows it:
- 7ab7ae2e6af
([job](https://github.com/viirya/spark-1/actions/runs/38314820956/job/115033969680)):
logs "Test still running after 2 hours, 14 minutes" for this test and is
cancelled at the 2h30m limit.
- 1053b61d411
([job](https://github.com/viirya/spark-1/actions/runs/38285171547/job/114945456531)):
hangs in the same test and is cancelled the same way.
- 038c4b3420e, the last commit before the handover wait
([job](https://github.com/viirya/spark-1/actions/runs/38194049541/job/114635793761)):
passes in about 66 minutes.
So the core-module suites scheduled after it, `BasePythonRunnerSuite`
included, have not run in CI since 1053b61d411. `TaskContextSuite` runs earlier
and passes.
The test and its comments (L1375-1379, L1405-1406, L1416-1418, L1425-1430)
document the skipped-listener race that this PR removes. Could the test be
updated to the new behavior? For example, call `markTaskCompleted` on another
thread, wait until that thread blocks in this wait, release the failure
listener, and then assert that the completion listener ran, on that thread.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,180 @@ private[spark] object BasePythonRunner extends Logging {
ThreadUtils.newDaemonCachedThreadPool("python-udf-pipelined-writer")
}
+ /** How often to log while task completion waits for a pipelined writer to
exit. */
+ private val pipelinedWriterExitWarnIntervalMs = 10000L
+
+ /**
+ * Runs `writer` on [[pipelinedWriterThreadPool]] and registers a task
pre-completion listener
+ * that interrupts it and blocks until it has exited.
+ *
+ * Task completion listeners free memory that the writer may still read or
write, e.g. the
+ * pages of the HybridRowQueue that buffers the input rows, or the off-heap
column vectors of a
+ * vectorized reader. If the writer were still running at that point, it
would be a
+ * use-after-free that can crash the executor (SPARK-33277). The listener is
added with
+ * `addTaskPreCompletionListener` so that it runs before all of them,
including the ones that
+ * the writer thread adds lazily while it pulls the upstream iterator, which
would run first if
+ * this were an ordinary completion listener. `Future.get()` cannot be used
for this wait:
+ * once `cancel(true)` succeeds it throws `CancellationException`
immediately, without waiting
+ * for the interrupted runnable to return.
+ *
+ * The wait has no time bound, so a writer blocked in upstream I/O that
ignores interrupts
+ * (e.g. a socket read without a timeout) holds task completion until that
I/O returns. This
+ * includes I/O that only a close listener of the upstream reader would
unblock, since that
+ * listener now runs after the wait. 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.
+ */
+ private[python] def startPipelinedWriter(writer: Runnable, context:
TaskContext): Unit = {
+ val writerTask = new PipelinedWriterTask(writer, taskIdentifier(context))
+ val writerFuture = pipelinedWriterThreadPool.submit(writerTask)
+ context.addTaskPreCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit =
+ writerTask.stopAndAwaitExit(writerFuture)
+ })
+ }
+
+ /** The state of a [[PipelinedWriterTask]]. */
+ private sealed trait PipelinedWriterState
+ private case object PipelinedWriterNotStarted extends PipelinedWriterState
+ private case class PipelinedWriterRunning(thread: Thread) extends
PipelinedWriterState
+ /** The writer has exited, or was stopped before it started and will never
run. */
+ private case object PipelinedWriterFinished extends PipelinedWriterState
+
+ /**
+ * Wraps a pipelined writer so that task completion can wait until the
writer has either
+ * exited or is guaranteed never to run.
+ */
+ private[python] class PipelinedWriterTask(
+ writer: Runnable,
+ taskName: String,
+ warnIntervalMs: Long = pipelinedWriterExitWarnIntervalMs)
+ extends Runnable {
+ // Protected by `this`. NotStarted moves to Running when run() starts the
writer, or to
+ // Finished when stopAndAwaitExit() gets there first; Running moves to
Finished when the
+ // writer exits. `exited` is counted down whenever the state becomes
Finished.
+ private var state: PipelinedWriterState = PipelinedWriterNotStarted
+ private val exited = new CountDownLatch(1)
+
+ private def finish(): Unit = synchronized {
+ state = PipelinedWriterFinished
+ exited.countDown()
+ }
+
+ override def run(): Unit = {
+ val start = synchronized {
+ if (state == PipelinedWriterNotStarted) {
+ state = PipelinedWriterRunning(Thread.currentThread())
+ true
+ } else {
+ false
+ }
+ }
+ if (start) {
+ try {
+ writer.run()
+ } catch {
+ case t: Throwable =>
+ // The thread pool's FutureTask would keep this to itself, and the
result of a
+ // cancelled future cannot be retrieved, so log it here.
+ logError(log"Pipelined Python writer of ${MDC(TASK_NAME,
taskName)} failed", t)
+ throw t
+ } finally {
+ finish()
+ }
+ }
+ }
+
+ /** Interrupts the writer thread, but only while it still runs the writer.
*/
+ private def interruptWriter(): Unit = synchronized {
+ state match {
+ case PipelinedWriterRunning(thread) => thread.interrupt()
+ case _ =>
+ }
+ }
+
+ /**
+ * Stops the writer and waits until it has exited, or makes sure it never
runs. Leaves the
+ * interrupt status of the calling thread as it was on entry.
+ */
+ def stopAndAwaitExit(future: Future[_]): Unit = {
+ val wasInterrupted = Thread.interrupted()
+ try {
+ // Interrupts the writer thread if the writer is running. This
unblocks channel.write
+ // (the JDK closes the channel and throws ClosedByInterruptException).
Otherwise the
+ // writer sees the interrupt only between calls to
writeNextInputToStream, each of which
+ // can pull and serialize a whole batch of input rows (e.g. for the
Arrow runners).
+ future.cancel(true)
Review Comment:
**[Nit] `future.cancel(true)` is no longer needed with the new state.**
The state already covers both cases. `NotStarted -> Finished` keeps `run()`
from starting the writer. For `PipelinedWriterRunning(thread)`, interrupting
`thread` under `this`, as `interruptWriter()` does, reaches the same thread
that `cancel(true)` interrupts. `cancel(true)` does it outside the lock, so it
can still interrupt the pool thread after the writer has exited. On the writer
thread itself, it interrupts the calling thread, which is one reason for the
unconditional `Thread.interrupted()` in the `finally`.
Interrupting in the state match instead, skipped when `thread eq
Thread.currentThread()`, would remove the `future` parameter of
`stopAndAwaitExit`, the `writerFuture` in `startPipelinedWriter`, the
self-interrupt, and most of the `FutureTask` plumbing in the tests.
--
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]
