viirya commented on code in PR #59332:
URL: https://github.com/apache/spark/pull/59332#discussion_r4262581044
##########
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:
Moved to SPARK-60222. I'll document there which listeners run on which
thread, and look at keeping leftover completion listeners off a thread that
holds the global `TaskContext` lock.
##########
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:
Done in 00da9fb69d4. `stopAndAwaitExit` no longer takes the future: it
interrupts the `PipelinedWriterRunning` thread under the lock, and skips that
when it runs on the writer thread itself. `startPipelinedWriter` no longer
keeps the future, and the tests no longer need `FutureTask`.
##########
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:
Moved to SPARK-60222 with the handover warning. I'll take the stack outside
the lock there, and attach it only to the first warning.
##########
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:
These tests moved to SPARK-60222 with the code they cover. I'll use daemon
threads and `join(timeout)` with `assert(!thread.isAlive)` there.
##########
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:
Moved to SPARK-60222. I'll make the interval configurable there and add a
test that checks the holder's name and 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
Review Comment:
Thanks. The handover wait moved to SPARK-60222, and I'll add that test
there: interrupt the completing thread while it waits, then check that the wait
continues, that the interrupt is kept, and that all completion listeners run on
the completing thread.
##########
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))
+ } else {
+ listenerInvocationThread = nextListenerInvocationThread
+ nextListenerInvocationThread = None
+ notifyAll()
+ released = true
+ None
+ }
+ }
+ pendingListeners.foreach { invoke =>
+ try {
+ invoke()
+ } catch {
+ case e: Throwable => listenerExceptions += e
Review Comment:
Thanks. The release loop moved to SPARK-60222, and I'll add a test there
with a deferred listener that throws. It will check that the invocation is
still released (or handed over) and that the failure reaches the caller.
--
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]
