viirya commented on code in PR #59663:
URL: https://github.com/apache/spark/pull/59663#discussion_r4404165554
##########
core/src/test/scala/org/apache/spark/shuffle/streaming/StreamingShuffleSuite.scala:
##########
@@ -1372,68 +1372,81 @@ class StreamingShuffleSuite
}
}
- // Demonstrates why background threads must not call
context.markTaskFailed() directly:
- // doing so can cause the task thread's markTaskCompleted() to silently skip
completion
- // listeners. In production, cleanupResources() is a completion listener, so
skipping it
- // leaks threads. The reader avoids this by routing background errors
through a
- // CompletableFuture so markTaskFailed/markTaskCompleted are both called
from the task thread.
+ // Shows why background threads should not call context.markTaskFailed()
directly: the
+ // failure listeners then run on that thread, and the task thread's
markTaskCompleted() blocks
+ // until they finish before it runs the completion listeners. In production,
cleanupResources()
+ // is a completion listener, so it waits for whatever the background
thread's failure listeners
+ // wait for. The reader avoids this by routing background errors through an
ErrorNotifier so
+ // that markTaskFailed/markTaskCompleted are both called from the task
thread.
test("background markTaskFailed races with task thread markTaskCompleted") {
val readerContext = createTaskContext(sparkConf, partitionId = 0)
TaskContext.setTaskContext(readerContext)
Review Comment:
Removed both.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -217,30 +240,85 @@ private[spark] class TaskContextImpl(
// writer that adds listeners lazily while it pulls its input. And since
the completion
// listeners only run after `completed` is set, they run after every
pre-completion listener,
// even if another thread is invoking listeners at this point.
+ //
+ // The completion listeners then run on this thread too, and this method
returns only after
+ // they have run. If it returned while another thread is invoking
listeners (e.g. a kill thread
+ // running failure listeners), that thread would run them only after its
own listeners, after
+ // the task has released its memory and reported its result. So the
listener invocation is
+ // claimed together with setting `completed`, which also keeps a thread
that adds a completion
+ // listener at that moment from running them; if another thread holds it,
that thread hands it
+ // over to this one when done. Failure listeners pending at that point run
first (see
+ // `invokeListeners()`). A completion listener added after this method has
run the others runs
+ // on the thread adding it, or, if another thread is invoking listeners at
that moment, on that
+ // thread before it stops (see `invokeListeners()`).
+ //
+ // The wait for the other thread has no time bound: if its listeners
block, so does this task.
+ // It logs a warning periodically, the first one with that thread's stack
trace. The task waits
+ // here before it releases its block locks and memory (the executor does
that once `Task.run`
+ // returns), so a failure listener on another thread, e.g. a kill thread,
must not wait for
+ // either. If such a task is killed, `spark.task.reaper.enabled` with a
positive
+ // `spark.task.reaper.killTimeout` bounds the wait by killing the executor
JVM, except in
+ // local mode, where the reaper only logs an error, and except when the
blocked listeners run
+ // on the reaper thread itself, within its kill call.
+ val currentThread = Thread.currentThread()
+ var mode: InvocationMode = Claimed
+ var awaitHandover = false
+ // Whether this call registered this thread as the listener invocation
thread or as the next
+ // one, so that it must make sure not to leave it registered.
+ var registered = false
val preCompletionExceptions = new ArrayBuffer[Throwable](2)
def nextPreCompletionListener(): Option[TaskCompletionListener] =
synchronized {
if (onPreCompleteCallbacks.empty()) {
// Set together with finding no more pre-completion listeners, so that
one added later
- // is handed to `addTaskCompletionListener` instead of being left in
the stack.
+ // is handed to `addTaskCompletionListener` instead of being left in
the stack. The
+ // listener invocation is claimed in the same step (see above).
completed = true
+ if (listenerInvocationThread.isEmpty) {
+ listenerInvocationThread = Some(currentThread)
+ registered = true
+ } else if (listenerInvocationThread.contains(currentThread)) {
+ // Called from a listener on this thread: run the completion
listeners as part of that
+ // invocation.
+ mode = Nested
+ } else {
+ nextListenerInvocationThread = Some(currentThread)
+ awaitHandover = true
+ registered = true
+ }
None
} else {
Some(onPreCompleteCallbacks.pop())
}
}
var listener: Option[TaskCompletionListener] = None
- while ({listener = nextPreCompletionListener(); listener.nonEmpty}) {
- invokeTaskPreCompletionListener(listener.get, preCompletionExceptions)
- }
-
try {
- invokeTaskCompletionListeners(error)
- } catch {
- case e: TaskCompletionListenerException if
preCompletionExceptions.nonEmpty =>
- // Merge the completion listeners' own failures, not the exception
that wraps them. Any
- // other error comes from invokeListeners itself (e.g. an
OutOfMemoryError) and is left
- // to propagate, so that the executor still sees a fatal one.
- preCompletionExceptions ++= e.getSuppressed.toSeq
+ while ({listener = nextPreCompletionListener(); listener.nonEmpty}) {
+ invokeTaskPreCompletionListener(listener.get, preCompletionExceptions)
+ }
+ try {
+ if (awaitHandover) {
+ awaitListenerInvocationHandover(currentThread)
+ }
+ invokeTaskCompletionListeners(error, mode)
+ } catch {
+ case e: TaskCompletionListenerException if
preCompletionExceptions.nonEmpty =>
+ // Merge the completion listeners' own failures, not the exception
that wraps them. Any
+ // other error comes from invokeListeners itself (e.g. an
OutOfMemoryError) and is left
+ // to propagate, so that the executor still sees a fatal one.
+ preCompletionExceptions ++= e.getSuppressed.toSeq
+ }
+ } finally {
+ if (registered) synchronized {
+ // `invokeListeners()` releases the invocation itself once it has
started; this covers an
+ // error before that, or in the wait. A stale waiter would get the
invocation handed to
+ // it, and a stale holder would keep every other thread from running
listeners.
+ if (nextListenerInvocationThread.contains(currentThread)) {
+ nextListenerInvocationThread = None
+ }
+ if (listenerInvocationThread.contains(currentThread)) {
Review Comment:
The release at the end of the loop with a waiter is now covered by "an error
that escapes the holder's listener loop hands the invocation over": the
holder's failure listener throws, logging that fails, and the thread waiting in
`markTaskCompleted()` must get the invocation and run the pending listener.
Replacing `releaseListenerInvocation()` there with `listenerInvocationThread =
None` fails it.
For L315-320: with item 2, the wait no longer ends early, and the thread it
hands over to goes straight into `invokeListeners()`, which releases the
invocation itself. So an error can reach this `finally` with this thread still
registered, as waiter or as holder, only before `invokeListeners()` has
started, e.g. an `OutOfMemoryError`. I kept it as a defensive release, and I
don't see a way to test it without injecting an error there; the description
says it is not covered.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -217,30 +240,85 @@ private[spark] class TaskContextImpl(
// writer that adds listeners lazily while it pulls its input. And since
the completion
// listeners only run after `completed` is set, they run after every
pre-completion listener,
// even if another thread is invoking listeners at this point.
+ //
+ // The completion listeners then run on this thread too, and this method
returns only after
+ // they have run. If it returned while another thread is invoking
listeners (e.g. a kill thread
+ // running failure listeners), that thread would run them only after its
own listeners, after
+ // the task has released its memory and reported its result. So the
listener invocation is
Review Comment:
Thanks for checking. With items 1 and 2 the holder no longer runs the
completion listeners, so the sentence is now about what happens if
`markTaskCompleted()` returned early: the other thread would not run them, and
the task would release its memory with them pending. The sentence about the
executor exiting is gone with the early end of the wait. The description now
says that on master the task thread could release its block locks and memory,
but not report its result, before it blocked on the `TaskRunner` lock.
##########
core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala:
##########
@@ -216,23 +222,22 @@ class TaskContextSuite extends SparkFunSuite with
BeforeAndAfter with LocalSpark
})
// Interrupt the task from a separate thread and wait until the interrupt
listener starts.
- val interruptThread = new Thread(() => context.markInterrupted("test
interrupt"))
- interruptThread.start()
- interruptListenerStarted.acquire()
-
- // While the interrupt listener is running on the interrupt thread, mark
the task completed
- // on this thread. With the dedicated interrupt listener loop, this must
NOT be blocked.
- context.markTaskCompleted(None)
-
- // The completion listener should have been called even though the
interrupt listener is still
- // running. If `invokeListeners()` were shared between interrupt and
completion listeners,
- // the completion listener would be silently skipped because
`listenerInvocationThread` would
- // be held by the interrupt thread.
- verify(completionListener, times(1)).onTaskCompletion(any())
-
- // Release the interrupt listener and join the interrupt thread.
- interruptListenerRelease.release()
- interruptThread.join()
+ val interruptThread =
startDaemonThread("interrupt")(context.markInterrupted("test interrupt"))
+ try {
+ interruptListenerStarted.acquire()
+
+ // While the interrupt listener is running on the interrupt thread, mark
the task completed
+ // on another thread. With the dedicated interrupt listener loop, this
must NOT be blocked.
+ // If `invokeListeners()` were shared between interrupt and completion
listeners, the
+ // interrupt thread would hold `listenerInvocationThread`, and the
completion would wait
+ // for it, so the join below would time out.
+
joinThread(startDaemonThread("completion")(context.markTaskCompleted(None)))
Review Comment:
Done: `startDaemonThread` now returns a thread that records what its body
throws, and `joinThread` asserts that it threw nothing. The `taskThread` in
`StreamingShuffleSuite` records it too.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -331,10 +497,34 @@ private[spark] class TaskContextImpl(
}
}
+ /**
+ * Invokes and removes the listeners in `listeners`, in reverse order of
registration.
+ *
+ * Other threads cannot invoke listeners while this thread does, so they
queue completion and
+ * failure listeners and leave them to this thread. (Post-status-update
listeners that another
+ * thread invokes meanwhile are not left to it; they are dropped.) Failure
listeners left to it
+ * (the task has failed) run before the next listener of `listeners`, e.g.
so that a failure
+ * listener that aborts a state store update runs before the completion
listener that releases
+ * the state store. Unless `mode` is `Nested`, this thread releases the
invocation when it is
+ * done, but first runs the completion listeners left to it if the task has
completed, unless a
+ * thread waiting in `markTaskCompleted()` takes the invocation over and
runs them itself. Such
+ * leftover completion listeners were added after `markTaskCompleted()` had
run the others.
Review Comment:
The leftover branch is gone (item 1), and so is the sentence.
##########
core/src/main/scala/org/apache/spark/util/ErrorNotifier.scala:
##########
@@ -28,9 +28,10 @@ import org.apache.spark.internal.Logging
*
* A background thread records the first error via [[markError]] rather than
failing the task
* directly; the owning thread polls [[throwErrorIfExists]] (or [[getError]])
at a safe point and
- * re-throws on its own thread. This avoids a race where calling
`TaskContext.markTaskFailed` from
- * a background thread can run concurrently with the task thread's
`markTaskCompleted` and silently
- * skip completion listeners (including cleanup), leading to thread leaks.
+ * re-throws on its own thread. Calling `TaskContext.markTaskFailed` from a
background thread
Review Comment:
Done, with your wording.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -134,7 +149,7 @@ private[spark] class TaskContextImpl(
completed
}
if (needToCallListener) {
- invokeTaskCompletionListeners(None)
+ invokeTaskCompletionListeners(None, Claim)
Review Comment:
Updated: with another thread invoking listeners, the call is a no-op, and
the new listener runs only if that thread runs the completion listeners in
`markTaskCompleted()`, or hands the invocation over to a thread waiting to do
so; otherwise it is not run until a listener is added while no thread is
invoking listeners.
--
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]
