viirya commented on code in PR #59332:
URL: https://github.com/apache/spark/pull/59332#discussion_r4252748737
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -169,37 +193,86 @@ private[spark] class TaskContextImpl(
}
private[spark] override def markTaskFailed(error: Throwable): Unit = {
+ markTaskFailed(error, nested = false)
+ }
+
+ /**
+ * @param nested whether 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, nested: Boolean): Unit = {
synchronized {
if (failureCauseOpt.isDefined) return
failureCauseOpt = Some(error)
}
- invokeTaskFailureListeners(error)
+ invokeTaskFailureListeners(error, nested)
}
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.
+ val nested = synchronized {
if (completed) return
completed = true
+ val currentThread = Thread.currentThread()
+ if (listenerInvocationThread.isEmpty) {
+ listenerInvocationThread = Some(currentThread)
+ false
+ } else if (listenerInvocationThread.contains(currentThread)) {
+ // Called from a listener on this thread: run the completion listeners
as part of that
+ // invocation.
+ true
+ } else {
+ nextListenerInvocationThread = Some(currentThread)
+ var interrupted = false
+ while (!listenerInvocationThread.contains(currentThread)) {
Review Comment:
Done in 7ab7ae2e6af. The handover wait now uses `wait(timeout)` and logs a
warning every 10 seconds while it lasts. The warning names the thread holding
the invocation and includes that thread's current stack trace, so a blocked
failure listener shows up in the log. As before, it has no time bound, and an
interrupt that arrives during it is restored afterwards. The comment and the
scaladoc now point at `spark.task.reaper.enabled` with a positive
`spark.task.reaper.killTimeout` as what bounds it for a killed task.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,148 @@ 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`
+ * bounds how long it may stay in this state.
+ */
+ 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)
+ })
+ }
+
+ /**
+ * 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 {
+ // Set by whichever of run() and stopAndAwaitExit() gets here first. If
run() wins, the
+ // writer runs and stopAndAwaitExit() waits for it; otherwise run() skips
the writer.
+ private val claimed = new AtomicBoolean(false)
+ private val exited = new CountDownLatch(1)
+ // The thread running the writer, while it does. Protected by `this`, so
that it is only
+ // interrupted while it still runs the writer, not after it has moved on
to other work.
+ private var writerThread: Thread = _
+
+ override def run(): Unit = {
+ if (claimed.compareAndSet(false, true)) {
+ synchronized {
+ writerThread = Thread.currentThread()
+ }
+ 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 {
+ synchronized {
+ writerThread = null
+ }
+ exited.countDown()
+ }
+ }
+ }
+
+ private def isWriterThread: Boolean = synchronized {
+ writerThread eq Thread.currentThread()
+ }
+
+ private def interruptWriter(): Unit = synchronized {
+ if (writerThread != null) {
+ writerThread.interrupt()
+ }
+ }
+
+ def stopAndAwaitExit(future: Future[_]): Unit = {
+ // 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)
+ if (claimed.compareAndSet(false, true)) {
+ // The writer never started and now never will.
+ return
+ }
+ if (isWriterThread) {
Review Comment:
You're right: on that thread `cancel(true)` only interrupts itself, and the
writer keeps going once the remaining listeners return. Changed as you
suggested. That path now logs at ERROR, saying that it may lead to a
use-after-free. `stopAndAwaitExit` clears the interrupt that `cancel(true)` set
on the thread, so the remaining listeners can clean up. The comment states that
reaching this path means the `TaskContextImpl` invariant is broken, and that
the writer is not stopped before the remaining listeners run. The test now runs
the writer through a `FutureTask`, so that `cancel(true)` really interrupts the
writer thread, and it checks that the flag is cleared when the call returns.
##########
core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala:
##########
@@ -50,6 +62,355 @@ class BasePythonRunnerSuite extends SparkFunSuite {
}
}
+ /**
+ * A writer that, once started, runs `onStart` and then keeps running until
`release` is
+ * counted down, ignoring interrupts the way a writer in the middle of
copying a row into a
+ * HybridRowQueue page does.
+ */
+ private class UninterruptibleWriter(
+ onStart: UninterruptibleWriter => Unit = _ => ()) extends Runnable {
+ val started = new CountDownLatch(1)
+ val release = new CountDownLatch(1)
+ @volatile var exited = false
+
+ override def run(): Unit = {
+ try {
+ onStart(this)
+ started.countDown()
+ awaitIgnoringInterrupts(release)
+ } finally {
+ exited = true
+ }
+ }
+ }
+
+ /** Waits for `latch`, ignoring interrupts like a writer in the middle of a
row copy. */
+ private def awaitIgnoringInterrupts(latch: CountDownLatch): Unit = {
+ var done = false
+ while (!done) {
+ try {
+ done = latch.await(1, TimeUnit.MINUTES)
+ } catch {
+ case _: InterruptedException => // keep going
+ }
+ }
+ }
+
+ /** Adds a completion listener that records whether `writer` had exited when
it ran. */
+ private def addCleanupListener(
+ context: TaskContext,
+ writer: UninterruptibleWriter): AtomicReference[Option[Boolean]] = {
+ val exitedBeforeCleanup = new AtomicReference[Option[Boolean]](None)
+ context.addTaskCompletionListener[Unit] { _ =>
exitedBeforeCleanup.set(Some(writer.exited)) }
+ exitedBeforeCleanup
+ }
+
+ /**
+ * Adds a pre-completion listener that runs right before the one added by
+ * `startPipelinedWriter`, so that the test knows task completion has
reached the writer.
+ */
+ private def addReachedListener(context: TaskContext): CountDownLatch = {
+ val reached = new CountDownLatch(1)
+ context.addTaskPreCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit =
reached.countDown()
+ })
+ reached
+ }
+
+ private def startThread(body: => Unit): Thread = {
+ val thread = new Thread(() => body)
+ thread.setDaemon(true)
+ thread.start()
+ thread
+ }
+
+ test("SPARK-60113: task completion waits for the pipelined writer to exit") {
+ val context = TaskContext.empty()
+ @volatile var cleanupAddedByWriter: AtomicReference[Option[Boolean]] = null
+ val writer = new UninterruptibleWriter(self => {
+ // Like an upstream reader that adds its close listener when the writer
first pulls it.
+ cleanupAddedByWriter = addCleanupListener(context, self)
+ })
+ // Like the listener that closes the HybridRowQueue, added on the task
thread before the
+ // writer starts.
+ val cleanupAddedBefore = addCleanupListener(context, writer)
+ BasePythonRunner.startPipelinedWriter(writer, context)
+ assert(writer.started.await(1, TimeUnit.MINUTES))
+ val reached = addReachedListener(context)
+
+ val completion = startThread(context.markTaskCompleted(None))
+ try {
+ assert(reached.await(1, TimeUnit.MINUTES))
+ // Task completion must stay blocked while the writer is still running.
+ completion.join(500)
+ assert(completion.isAlive)
+ assert(cleanupAddedBefore.get.isEmpty)
+ assert(cleanupAddedByWriter.get.isEmpty)
+ } finally {
+ writer.release.countDown()
+ }
+ completion.join(TimeUnit.MINUTES.toMillis(1))
+ assert(!completion.isAlive)
+ assert(cleanupAddedBefore.get === Some(true))
+ assert(cleanupAddedByWriter.get === Some(true))
+ }
+
+ test("SPARK-60113: a listener the writer adds once completion has started
waits for it") {
+ val context = TaskContext.empty()
+ val addNow = new CountDownLatch(1)
+ val added = new CountDownLatch(1)
+ // Whether the writer had exited, and the thread, when the listener added
by the writer ran.
+ val lateCleanup = new AtomicReference[Option[(Boolean, Thread)]](None)
+ val writerStarted = new CountDownLatch(1)
+ val release = new CountDownLatch(1)
+ // Like a writer that adds a listener when it pulls more input after
completion has started.
+ val lateWriter: Runnable = new Runnable {
+ @volatile var exited = false
+ override def run(): Unit = {
+ try {
+ writerStarted.countDown()
+ awaitIgnoringInterrupts(addNow)
+ context.addTaskCompletionListener[Unit] { _ =>
+ lateCleanup.set(Some((exited, Thread.currentThread())))
+ }
+ added.countDown()
+ awaitIgnoringInterrupts(release)
+ } finally {
+ exited = true
+ }
+ }
+ }
+ BasePythonRunner.startPipelinedWriter(lateWriter, context)
+ assert(writerStarted.await(1, TimeUnit.MINUTES))
+ val reached = addReachedListener(context)
+
+ val completion = startThread(context.markTaskCompleted(None))
+ try {
+ assert(reached.await(1, TimeUnit.MINUTES))
+ addNow.countDown()
+ assert(added.await(1, TimeUnit.MINUTES))
+ // The writer thread only queued the listener: task completion still
waits for the writer.
+ completion.join(500)
+ assert(completion.isAlive)
+ assert(lateCleanup.get.isEmpty)
+ } finally {
+ release.countDown()
+ }
+ completion.join(TimeUnit.MINUTES.toMillis(1))
+ assert(!completion.isAlive)
+ assert(lateCleanup.get === Some((true, completion)))
+ }
+
+ test("SPARK-60113: an interrupt does not stop the wait and is not passed to
later listeners") {
+ val context = TaskContext.empty()
+ val writer = new UninterruptibleWriter
+ val cleanup = addCleanupListener(context, writer)
+ val interruptedInCleanup = new AtomicReference[Option[Boolean]](None)
+ context.addTaskCompletionListener[Unit] { _ =>
+ interruptedInCleanup.set(Some(Thread.currentThread().isInterrupted))
+ }
+ BasePythonRunner.startPipelinedWriter(writer, context)
+ assert(writer.started.await(1, TimeUnit.MINUTES))
+ val reached = addReachedListener(context)
+
+ val completion = startThread(context.markTaskCompleted(None))
+ try {
+ assert(reached.await(1, TimeUnit.MINUTES))
+ // Give the completion thread time to block in the wait, then interrupt
it, e.g. a task
+ // kill that lands while the listeners run.
+ completion.join(200)
Review Comment:
Thanks, changed to wait with `eventually` until the completion thread is in
`TIMED_WAITING` before interrupting it.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -267,13 +352,19 @@ private[spark] class TaskContextImpl(
}
def getNextListenerOrDeregisterThread(): Option[T] = synchronized {
- if (listeners.empty()) {
- // We have executed all listeners that have been added so far.
Deregister this thread as the
- // callback invocation thread.
- listenerInvocationThread = None
- None
- } else {
- Some(listeners.pop())
+ listenerStacks.find(!_.empty()) match {
+ case Some(listeners) =>
+ Some(listeners.pop())
+ case None =>
+ // We have executed all listeners that have been added so far.
Deregister this thread as
+ // the callback invocation thread, unless an outer invocation on
this thread continues. A
+ // thread waiting in `markTaskCompleted()` takes over directly.
+ if (!nested) {
Review Comment:
Fixed. Before a thread releases the listener invocation, it now also runs
the listeners that other threads queued in the meantime and left to it. Those
are the failure listeners if the task has failed, and the completion listeners
if the task has completed. They run in the nested mode, so the invocation stays
with this thread until nothing is pending. If a thread is waiting in
`markTaskCompleted`, the invocation is handed over to it instead, and it runs
whatever is pending when it releases. This also covers the reverse case:
completion listeners queued while a failure-listener thread holds the
invocation after the task has completed. New tests cover both directions.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,148 @@ 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`
Review Comment:
Thanks, corrected in both places (`startPipelinedWriter` and
`markTaskCompleted`). They now say that `spark.task.reaper.enabled` with a
positive `spark.task.reaper.killTimeout` bounds how long a killed task may stay
in this state, by killing the executor JVM.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -169,37 +193,86 @@ private[spark] class TaskContextImpl(
}
private[spark] override def markTaskFailed(error: Throwable): Unit = {
+ markTaskFailed(error, nested = false)
+ }
+
+ /**
+ * @param nested whether 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, nested: Boolean): Unit = {
synchronized {
if (failureCauseOpt.isDefined) return
failureCauseOpt = Some(error)
}
- invokeTaskFailureListeners(error)
+ invokeTaskFailureListeners(error, nested)
}
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.
Review Comment:
Updated both. The `ErrorNotifier` doc now says that calling `markTaskFailed`
from a background thread runs the failure listeners on that thread, and that a
concurrent `markTaskCompleted` on the task thread then blocks until they finish
before it runs the completion listeners. The
`StreamingShuffleReader.checkTaskFailure` comment now says that surfacing the
error on the task thread lets the failure and completion listeners run there
without waiting for another thread. I left out the dropped-failure-listener
hazard, since the change for L362 removes it.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -239,10 +312,22 @@ private[spark] class TaskContextImpl(
}
}
+ /**
+ * Invokes and removes the listeners in `listenerStacks`, each in reverse
order of registration.
+ * A listener is always taken from the first non-empty stack, so all
listeners of a stack run
+ * before those of the stacks after it that are still pending.
+ *
+ * @param claimed whether the caller has already registered this thread as
+ * `listenerInvocationThread`.
+ * @param nested whether that registration belongs to an outer invocation on
this thread, which
+ * keeps it when this one is done. Implies `claimed`.
+ */
private def invokeListeners[T](
- listeners: Stack[T],
+ listenerStacks: Seq[Stack[T]],
name: String,
- error: Option[Throwable])(
+ error: Option[Throwable],
+ claimed: Boolean = false,
Review Comment:
Done: replaced them with a single `InvocationMode` (`Claim`, `Claimed`,
`Nested`) in a private `TaskContextImpl` companion object.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,148 @@ 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`
+ * bounds how long it may stay in this state.
+ */
+ 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)
+ })
+ }
+
+ /**
+ * 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 {
+ // Set by whichever of run() and stopAndAwaitExit() gets here first. If
run() wins, the
+ // writer runs and stopAndAwaitExit() waits for it; otherwise run() skips
the writer.
+ private val claimed = new AtomicBoolean(false)
Review Comment:
Done. The `AtomicBoolean`, the latch-only completion and the separate
`writerThread` are replaced by a single state guarded by `this`: `NotStarted`,
`Running(thread)`, or `Finished`. The latch is counted down on every transition
to `Finished`, including when the stop listener gets there before the writer
starts. A second `stopAndAwaitExit` therefore returns at once, and the test now
calls it twice.
--
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]
