dongjoon-hyun commented on code in PR #59332:
URL: https://github.com/apache/spark/pull/59332#discussion_r4264575990
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,182 @@ 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 `Future.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))
+ pipelinedWriterThreadPool.submit(writerTask)
+ context.addTaskPreCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit =
writerTask.stopAndAwaitExit()
+ })
+ }
+
+ /** 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, 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(): Unit = {
+ val wasInterrupted = Thread.interrupted()
+ try {
+ val runningThread = synchronized {
+ state match {
+ case PipelinedWriterNotStarted =>
+ // The writer never started and now never will.
+ finish()
+ None
+ case PipelinedWriterRunning(thread) =>
+ // Interrupts the writer thread. 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).
+ if (thread ne Thread.currentThread()) {
+ thread.interrupt()
Review Comment:
**1. [High] The stop interrupt can free the Arrow writer's output buffer
while the writer is copying into it**
When this interrupt lands while an Arrow writer is serializing a batch, the
JDK closes the writer's `DataOutputStream`, and with it the
`DirectByteBufferOutputStream`, on this thread, while the writer thread is
still copying into that buffer:
- The Arrow inputs write `dataOut` through `Channels.newChannel(dataOut)`:
`new ArrowStreamWriter(root, null, dataOut)` (`PythonArrowInput.scala` L103)
and `new WriteChannel(Channels.newChannel(dataOut))` (L144, L238, and
`ColumnarArrowPythonInput.scala` L125, L148). In pipelined mode, `dataOut` is
`new DataOutputStream(new DirectByteBufferOutputStream(bufferSize))`
(L1470-L1471 in this file).
- `Channels.newChannel(OutputStream)` returns `WritableByteChannelImpl`, an
`AbstractInterruptibleChannel`. Its `write` wraps each chunk of up to 8 KiB in
`begin(); out.write(...); end()`, and its `implCloseChannel()` is `out.close()`.
- `Thread.interrupt()` on a thread that is between `begin()` and `end()`
runs the channel's interruptor on the calling thread. Here that is
`dataOut.close()` -> `DirectByteBufferOutputStream.close()` ->
`StorageUtils.dispose(buffer)`, which frees the native memory right away.
`FilterOutputStream.close()` does not take the `DataOutputStream` monitor, so
the synchronized `write` does not keep it out.
- Meanwhile the writer is in `buffer.put(b, off, len)`, or in `grow()`
copying the old buffer, and writes to freed memory. Only after that does
`end()` throw `ClosedByInterruptException`.
This happens before `awaitExit()` starts, so the invariant this PR
establishes does not hold for the writer's own buffer: the executor can crash,
or memory that another allocation reused can be corrupted silently. A single
stop has a small chance, since the interrupt has to land during a chunk copy,
but any early completion of an Arrow UDF while its writer is busy serializing
can hit it, and the repeated interrupt at L271 adds more chances. I checked
each step against the JDK 17/21 sources and the Arrow 19.0.0 bytecode.
It predates this PR: the old `cancel(true)` interrupted the same thread. The
sync path has the same exposure when `Task.kill` interrupts a task thread that
is serializing in `ReaderInputStream`. Could the writer get a stream whose
`close()` does not free the buffer? For example, wrap `bufferStream` in a
`FilterOutputStream` that forwards `write(b, off, len)` and ignores `close()`,
and free the buffer only in the writer's own `finally`, which already calls
`bufferStream.close()`. The writer still stops right after the current chunk,
since `end()` throws `ClosedByInterruptException` either way.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -1329,16 +1487,16 @@ private[spark] abstract class BasePythonRunner[IN, OUT](
}
} catch {
case _: InterruptedException =>
- // Task cancelled via Future.cancel(true)
+ // Interrupted by the task completion listener that stops the writer.
Thread.currentThread().interrupt()
case _: java.nio.channels.ClosedByInterruptException =>
// Task cancelled while blocked in channel.write(). The channel is
// automatically closed by the JVM, which will cause Python to
receive
// EOF and the reader thread to get IOException.
Thread.currentThread().interrupt()
- case NonFatal(t) =>
+ case t: Throwable =>
Review Comment:
**3. [Low] A fatal error before this `try` still leaves the reader waiting
for Python**
This catch does not cover the four statements before the `try`
(L1468-L1471): `TaskContext.setTaskContext`,
`InputFileBlockHolder.setThreadLocalValue`, `new
DirectByteBufferOutputStream(bufferSize)` and `new DataOutputStream(...)`. If
the direct buffer allocation fails (`Platform.allocateDirectBuffer` ->
`Unsafe.allocateMemory`), the `OutOfMemoryError` skips `writer.setException`,
`shutdownOutput()` and the `finally`. `PipelinedWriterTask` logs it, and the
discarded `FutureTask` swallows it. The Python worker then waits for its init
message, and the reader finds `writer.exception` empty and blocks in
`readInt()` (the idle timeout is 0 by default) until the task is killed. The
thread-locals also stay set on the pool thread.
It is rare (a 64 KiB allocation), and the order predates this PR, but item 6
of the description says that fatal writer errors are now passed to the reader.
Could the four statements move into the `try`, with a nullable `bufferStream`
that the `finally` closes only if it was created?
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,182 @@ 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 `Future.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))
+ pipelinedWriterThreadPool.submit(writerTask)
+ context.addTaskPreCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit =
writerTask.stopAndAwaitExit()
+ })
+ }
+
+ /** 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, 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(): Unit = {
+ val wasInterrupted = Thread.interrupted()
+ try {
+ val runningThread = synchronized {
+ state match {
+ case PipelinedWriterNotStarted =>
+ // The writer never started and now never will.
+ finish()
+ None
+ case PipelinedWriterRunning(thread) =>
+ // Interrupts the writer thread. 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).
+ if (thread ne Thread.currentThread()) {
Review Comment:
**6. [Low] No test fails if this interrupt, or the `Thread.interrupted()` in
the `finally`, is removed**
Deleting L294-L296, or inverting the guard, keeps every test green:
- "the writer is interrupted again if it consumed the interrupt" uses
`warnIntervalMs = 100`, so two warning-tick interrupts at about 100 ms and 200
ms satisfy `interrupts.get >= 2`.
- The `running` writer in "an interrupt pending before stopping the writer
is kept" exits at the first 10 s tick, within its 1-minute join.
In production, every task that completes early while its writer is blocked
in `channel.write` would then hold completion for 10 s and log a warning. The
same goes for the unconditional `Thread.interrupted()` at L318: round 4's
same-thread test pinned it through the `FutureTask` self-interrupt, and nothing
replaced that test.
"stopping waits for a writer that started after the interrupt" (L252) sets
the flag before `run()` claims the writer, which production no longer does: the
pool clears the flag before each task, and nothing cancels the future. It could
pin this interrupt instead, with a writer that exits on interrupt, a long
`warnIntervalMs` (e.g. 10 minutes), and an assertion that stopping returns
within the 1-minute join. For L318, a deterministic test could hold the monitor
of a `PipelinedWriterTask` that never started, start a stopper, wait until it
is `BLOCKED`, interrupt it, and release the monitor: the stopper must return
with the flag cleared.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -177,29 +194,48 @@ private[spark] class TaskContextImpl(
}
private[spark] override def markTaskCompleted(error: Option[Throwable]):
Unit = {
- synchronized {
+ val claimedInvocation = synchronized {
if (completed) return
completed = true
+ // Claim the listener invocation together with setting `completed`,
unless another thread is
+ // already invoking listeners. Otherwise a thread that adds a completion
listener in between
+ // would see `completed` and run the listeners itself. That thread can
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.
+ if (listenerInvocationThread.isEmpty) {
Review Comment:
**8. [Discussion] Running the pre-completion listeners before `completed` is
set would close the deferred writer paths here**
Instead of claiming the invocation here and popping the pre-completion stack
first in `invokeListeners`, `markTaskCompleted` could run the pre-completion
listeners on this thread *before* it sets `completed`, outside
`listenerInvocationThread`, as `invokeTaskInterruptListeners` already does.
Then it would set `completed` and invoke the completion listeners as it does
now.
While the stop listener waits, the writer's lazy `add*Listener` calls see
`completed == false` and only push, so the writer can never become the invoker.
And the writer has exited before any completion listener runs, and before
`Task.run` reaches `cleanUpAllAllocatedMemory`, even when another thread holds
the invocation. That closes all three writer paths that the description defers
to SPARK-60222: another holder at completion, the catch-block window, and the
same-thread branch. A task without pre-completion listeners behaves exactly as
on master, so this touches fewer tasks than the claim, which applies to every
task.
It is not less code. Since `completed` is set later, the method needs its
own "completion started" flag, and failures of pre-completion listeners have to
be merged into the `TaskCompletionListenerException`. `isCompleted()` is false
during the wait; I checked `Writer.open`, `MonitorThread`,
`BlockManager.releaseLock` and `StreamingShuffleWriter`, and none of them is
harmed. The 500-iteration claim test would also go. What do you think?
##########
core/src/main/scala/org/apache/spark/BarrierTaskContext.scala:
##########
@@ -208,6 +208,12 @@ class BarrierTaskContext private[spark] (
this
}
+ override private[spark] def addTaskPreCompletionListener(
Review Comment:
**9. [Low] No test covers this delegation, and dropping it would silently
restore the LIFO order in barrier stages**
If this override were removed, `TaskContext`'s default (`TaskContext.scala`
L164-L165) would forward to `addTaskCompletionListener`, so everything would
compile and all tests would pass, while barrier stages get the use-after-free
back. Pipelined mode does run with a `BarrierTaskContext`:
`mapInPandas`/`mapInArrow(..., barrier=True)` goes through `MapInBatchExec`
(`child.execute().barrier()`) with
`SQL_MAP_PANDAS_ITER_UDF`/`SQL_MAP_ARROW_ITER_UDF`, which `usePipelined` does
not exclude.
Could `BarrierTaskContextSuite` get a `local[2]` barrier job that adds a
pre-completion listener and then a completion listener and checks their order?
Alternatively, declaring `addTaskPreCompletionListener` abstract in
`TaskContext`, like `addPostStatusUpdateListener` (SPARK-57491), would make it
impossible for a wrapper to miss it.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,182 @@ 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:
Review Comment:
**12. [Nit] This sentence still justifies the latch with
`Future.cancel(true)`, which the code no longer calls**
Since `00da9fb69d4`, the future from `submit` is discarded (L212) and
nothing cancels it, and an uncancelled `FutureTask.get()` does wait for `run()`
to return. So "`Future.get()` cannot be used for this wait" no longer holds as
written, and it could even read as a reason to go back to the future. The
actual reasons are that a `Future` exposes neither the thread to interrupt
(directly, and again on every tick) nor the thread to compare with the current
one, and that a stop before the pool starts the task should not wait for a pool
thread.
Maybe: "The listener does not use the `Future` returned by the pool:
interrupting the writer through it takes `Future.cancel(true)`, after which
`Future.get()` throws `CancellationException` immediately instead of waiting
for the runnable to return. `PipelinedWriterTask` records the thread running
the writer to interrupt it directly, and counts down a latch when the writer
exits or is stopped before it starts."
##########
core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala:
##########
@@ -50,6 +62,356 @@ 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))
+ // Interrupt the completion thread once it blocks in the wait, e.g. a
task kill that lands
+ // while the listeners run.
+ eventually(timeout(1.minute)) {
+ assert(completion.getState == Thread.State.TIMED_WAITING)
+ }
+ completion.interrupt()
+ completion.join(500)
+ assert(completion.isAlive)
+ assert(cleanup.get.isEmpty)
+ } finally {
+ writer.release.countDown()
+ }
+ completion.join(TimeUnit.MINUTES.toMillis(1))
+ assert(!completion.isAlive)
+ assert(cleanup.get === Some(true))
+ assert(interruptedInCleanup.get === Some(false))
+ }
+
+ test("SPARK-60113: a pipelined writer stopped before it starts never runs") {
+ val writer = new UninterruptibleWriter
+ val writerTask = new BasePythonRunner.PipelinedWriterTask(writer, "test
task")
+ // The pool has not started the writer yet: stopping must not wait for it.
+ writerTask.stopAndAwaitExit()
+ // Stopping again does not wait either.
+ writerTask.stopAndAwaitExit()
+ // Even if the pool thread reaches the task afterwards, the writer must
not run. Release it
+ // up front so that a writer that wrongly runs fails the test instead of
hanging it.
+ writer.release.countDown()
+ writerTask.run()
+ assert(writer.started.getCount === 1)
+ assert(!writer.exited)
+ }
+
+ test("SPARK-60113: stopping waits for a writer that started after the
interrupt") {
+ // The writer starts with the interrupt flag of its thread already set,
and ignores it.
+ val writer = new UninterruptibleWriter
+ val writerTask = new BasePythonRunner.PipelinedWriterTask(writer, "test
task")
+ val writerThread = startThread {
+ Thread.currentThread().interrupt()
+ writerTask.run()
+ }
+ assert(writer.started.await(1, TimeUnit.MINUTES))
+
+ val stopper = startThread(writerTask.stopAndAwaitExit())
+ try {
+ stopper.join(500)
+ assert(stopper.isAlive)
+ assert(!writer.exited)
+ } finally {
+ writer.release.countDown()
+ }
+ stopper.join(TimeUnit.MINUTES.toMillis(1))
+ assert(!stopper.isAlive)
+ assert(writer.exited)
+ writerThread.join(TimeUnit.MINUTES.toMillis(1))
+ }
+
+ test("SPARK-60113: a fatal error of the pipelined writer is logged") {
+ val writerTask = new BasePythonRunner.PipelinedWriterTask(
+ // scalastyle:off throwerror
+ () => throw new LinkageError("fatal writer error"), "test task")
+ // scalastyle:on throwerror
+ val logAppender = new LogAppender("pipelined writer failure")
+ withLogAppender(logAppender) {
+ intercept[LinkageError] {
+ writerTask.run()
+ }
+ }
+ assert(logAppender.loggingEvents.exists { event =>
+ event.getLevel == Level.ERROR &&
+ event.getMessage.getFormattedMessage.contains("Pipelined Python writer
of test task") &&
+ event.getThrown.getMessage == "fatal writer error"
+ })
+ // The writer has exited, so stopping does not wait.
+ writerTask.stopAndAwaitExit()
+ }
+
+ test("SPARK-60113: stopping on the writer thread itself does not wait for
the writer") {
+ val taskRef = new AtomicReference[BasePythonRunner.PipelinedWriterTask]()
+ val stopReturned = new CountDownLatch(1)
+ val interruptedAfterStop = new AtomicReference[Option[Boolean]](None)
+ val writerTask = new BasePythonRunner.PipelinedWriterTask(() => {
+ // Like task completion listeners handed to the writer thread while it
pulls its input.
+ taskRef.get.stopAndAwaitExit()
+ interruptedAfterStop.set(Some(Thread.currentThread().isInterrupted))
+ stopReturned.countDown()
+ }, "test task")
+ taskRef.set(writerTask)
+ val writerThread = startThread(writerTask.run())
+ assert(stopReturned.await(1, TimeUnit.MINUTES))
+ // Stopping leaves the interrupt status of the thread it runs on as it
was, so that the
+ // listeners that run next on it can clean up.
+ assert(interruptedAfterStop.get === Some(false))
+ writerThread.join(TimeUnit.MINUTES.toMillis(1))
+ assert(!writerThread.isAlive)
+ // The writer has exited, so stopping from another thread does not wait
either.
+ writerTask.stopAndAwaitExit()
+ }
+
+ test("SPARK-60113: the writer is interrupted again if it consumed the
interrupt") {
+ val interrupts = new AtomicInteger(0)
+ val writerTask = new BasePythonRunner.PipelinedWriterTask(() => {
+ // Like upstream code that catches the first InterruptedException and
carries on, and then
+ // blocks again (e.g. in channel.write), where the next interrupt stops
it.
+ val blocker = new CountDownLatch(1)
+ var stopped = false
+ while (!stopped) {
+ try {
+ blocker.await()
+ } catch {
+ case _: InterruptedException =>
+ stopped = interrupts.incrementAndGet() >= 2
+ }
+ }
+ }, "test task", warnIntervalMs = 100)
+ val writerThread = startThread(writerTask.run())
+ eventually(timeout(1.minute)) {
+ assert(writerThread.getState == Thread.State.WAITING)
+ }
+
+ val stopper = startThread(writerTask.stopAndAwaitExit())
+ stopper.join(TimeUnit.MINUTES.toMillis(1))
+ assert(!stopper.isAlive)
+ assert(interrupts.get >= 2)
+ writerThread.join(TimeUnit.MINUTES.toMillis(1))
+ }
+
+ test("SPARK-60113: an interrupt pending before stopping the writer is kept")
{
+ // The writer never started.
+ val notStarted = new BasePythonRunner.PipelinedWriterTask(() => (), "test
task")
+ // The writer runs and exits once interrupted.
+ val running = new BasePythonRunner.PipelinedWriterTask(() => {
+ try {
+ new CountDownLatch(1).await()
+ } catch {
+ case _: InterruptedException =>
+ }
+ }, "test task")
+ val writerThread = startThread(running.run())
+ eventually(timeout(1.minute)) {
+ assert(writerThread.getState == Thread.State.WAITING)
+ }
+
+ Seq(notStarted, running).foreach { writerTask =>
+ val interruptedAfterStop = new AtomicReference[Option[Boolean]](None)
+ val stopper = startThread {
+ Thread.currentThread().interrupt()
+ writerTask.stopAndAwaitExit()
+ interruptedAfterStop.set(Some(Thread.currentThread().isInterrupted))
+ }
+ stopper.join(TimeUnit.MINUTES.toMillis(1))
+ assert(interruptedAfterStop.get === Some(true))
+ }
+ writerThread.join(TimeUnit.MINUTES.toMillis(1))
+ }
+
+ test("SPARK-60113: a fatal error of the pipelined writer is passed to the
reader") {
+ LocalSparkContext.withSpark(new SparkContext("local", "test")) { _ =>
+ val func = SimplePythonFunction(
+ Seq.empty, new JHashMap[String, String](), new JArrayList[String](),
"python3", "3",
+ new JArrayList[Broadcast[PythonBroadcast]](), null)
+ val runner = new PythonRunner(Seq(ChainedPythonFunctions(Seq(func))),
None)
+ val server = ServerSocketChannel.open()
+ try {
+ server.bind(new InetSocketAddress(InetAddress.getLoopbackAddress, 0))
+ val channel = SocketChannel.open(server.getLocalAddress)
+ val pythonSide = server.accept()
+ try {
+ val context = TaskContext.empty()
+ val worker = PythonWorker(channel)
+ val writer = new runner.Writer(SparkEnv.get, worker, Iterator.empty,
0, context) {
+ override def open(dataOut: DataOutputStream): Unit = {}
+ override protected def writeCommand(dataOut: DataOutputStream):
Unit = {}
+ override def writeNextInputToStream(dataOut: DataOutputStream):
Boolean = {
+ // scalastyle:off throwerror
+ throw new LinkageError("fatal writer error")
+ // scalastyle:on throwerror
+ }
+ }
+ intercept[LinkageError] {
+ new runner.PipelinedWriterRunnable(worker, writer, 1024,
context).run()
+ }
+ // The reader rethrows the error instead of waiting for the Python
worker, which sees
+ // the end of its input.
+ assert(writer.exception.map(_.getMessage) === Some("fatal writer
error"))
+ assert(pythonSide.read(ByteBuffer.allocate(1)) === -1)
Review Comment:
**10. [Low] This read blocks forever if `shutdownOutput()` regresses,
instead of failing**
`pythonSide` is a blocking `SocketChannel`, and `SocketChannel.read` honors
no timeout. If a change keeps `writer.setException` but skips
`shutdownOutput()`, e.g. by moving the rethrow above it, the assertions above
still pass and this read never returns, since `channel` is closed only in the
`finally` after it. `SparkFunSuite`'s 20-minute `failAfter` uses the default
`DoNotSignal`, so the core job would run into its 2h30m limit, as it did with
`StreamingShuffleSuite`.
Reading through the socket adaptor, which honors `SO_TIMEOUT`, would make it
fail instead:
```scala
pythonSide.socket().setSoTimeout(TimeUnit.MINUTES.toMillis(1).toInt)
assert(pythonSide.socket().getInputStream.read() === -1)
```
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -267,13 +310,14 @@ 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.
+ listenerInvocationThread = None
Review Comment:
**4. [Low] Failure listeners that another thread triggers during the writer
wait are dropped again, and nothing says so**
With the release loop moved to SPARK-60222, the gap from [round
3](https://github.com/apache/spark/pull/59332#discussion_r4252547639) is back.
While this thread holds the invocation, which now includes the wait for the
writer, a `markTaskFailed` or `addTaskFailureListener` from another thread
returns at L300-L303. This loop drains only the pre-completion and completion
stacks before it deregisters here, so those failure listeners never run,
although `isFailed()` is true.
For example, a task is killed while its writer is blocked in upstream I/O,
and a `TaskInterruptListener` throws: `markInterrupted` calls `markTaskFailed`
on the kill thread, which only records the failure. Master has the same gap,
but only while the completion listeners run; this PR extends it to the whole
writer wait. The in-tree interrupt listeners catch `NonFatal`, so in practice
it takes a third-party listener.
I'm not asking to bring the release loop back here. But the description's
SPARK-60222 paragraph lists only the completion-listener cases, and the round-3
thread still ends at "Fixed.". Could the paragraph, or the
`startPipelinedWriter` scaladoc, mention this case, so that SPARK-60222 covers
it too?
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,182 @@ 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 `Future.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))
+ pipelinedWriterThreadPool.submit(writerTask)
+ context.addTaskPreCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit =
writerTask.stopAndAwaitExit()
+ })
+ }
+
+ /** 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, 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(): Unit = {
+ val wasInterrupted = Thread.interrupted()
+ try {
+ val runningThread = synchronized {
+ state match {
+ case PipelinedWriterNotStarted =>
+ // The writer never started and now never will.
+ finish()
+ None
+ case PipelinedWriterRunning(thread) =>
+ // Interrupts the writer thread. 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).
+ if (thread ne Thread.currentThread()) {
+ thread.interrupt()
+ }
+ Some(thread)
+ case PipelinedWriterFinished => None
+ }
+ }
+ runningThread.foreach { thread =>
+ if (thread eq Thread.currentThread()) {
+ // TaskContextImpl can still let a thread that adds a listener run
the pending
+ // completion listeners: when another thread is invoking listeners
as the task
+ // completes, or while a failing listener's failure listeners run.
If that thread is
+ // this writer, in the middle of pulling its input, the writer can
only exit after
+ // this returns, so waiting would never end. It is not stopped
either: the remaining
+ // listeners run on this thread right after this one and free
memory that the writer
+ // uses again when they return, as before this listener waited at
all.
Review Comment:
**2. [Medium] On the writer-thread path nothing stops the writer any more,
so it keeps writing after the cleanup**
Since `00da9fb69d4` this branch is reachable again (the paths the
description defers to SPARK-60222), and "as before this listener waited at all"
is not quite what happens.
Before this PR, `cancel(true)` on this path interrupted this thread itself,
since it is the `FutureTask` runner. Once the remaining listeners returned, the
writer stopped at the `isInterrupted` check at L1479 after the current
`writeNextInputToStream` call, or at the next `channel.write`. Now nothing sets
that flag: L294 skips this thread, and the `finally` restores the state on
entry. Clearing the self-interrupt was my round-3 suggestion, back when this
path was unreachable.
So after `queue.close()` and `worker.stop()` run here, the writer goes back
to the upstream code and keeps adding rows to the freed `HybridRowQueue` page:
`HybridQueue.close()` does not reset `writing`, so `add` keeps writing through
it. It stops only when `bufferStream.size() >= bufferSize` (64 KiB by default)
and the flush hits the closed channel, or when upstream throws. On the row path
that is up to 64 KiB of rows, against at most the current pickled batch (100
rows) before.
Could the loop also check `!context.isCompleted()`?
```scala
while (hasInput && !Thread.currentThread().isInterrupted &&
!context.isCompleted()) {
```
It is always true on this path, since the takeover requires `completed`. On
the normal path the stop listener interrupts the writer right after anyway. And
it leaves the interrupt status of the listeners alone, which was the reason to
clear the flag. Otherwise, could the comment say that the writer keeps going?
Also on this comment: the invocation is released around a failing listener's
failure listeners, not while they run. The windows are between
`listenerInvocationThread = None` (`TaskContextImpl.scala` L367) and the claim
in `markTaskFailed`, and between their deregistration and the reclaim in the
`finally`. And [your reply on the
split](https://github.com/apache/spark/pull/59332#discussion_r4262579640) says
the comment points to SPARK-60222, but no code comment mentions it yet.
##########
core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala:
##########
@@ -50,6 +62,356 @@ 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)
Review Comment:
**7. [Low] The only test that catches reverting
`addTaskPreCompletionListener` in `startPipelinedWriter` depends on a race**
If `startPipelinedWriter` went back to `addTaskCompletionListener` (the LIFO
bug this PR fixes), this is the only test that would fail, and only when this
push comes after the stop listener is registered. `startPipelinedWriter`
submits the writer first (L212) and registers the stop listener next (L213).
When the pool thread pushes `cleanupAddedByWriter` first, the reverted stop
listener lands on top, runs first, and L146-L155 all pass. That order is likely
enough: an idle thread left by the SPARK-58192 test picks the task up quickly,
while the test thread may still be loading the listener class.
Could `onStart` wait for a latch that the test counts down after
`startPipelinedWriter` returns? Then the writer's listener is always added
after the stop listener, and the test fails deterministically on that
regression.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,182 @@ 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 `Future.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))
+ pipelinedWriterThreadPool.submit(writerTask)
+ context.addTaskPreCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit =
writerTask.stopAndAwaitExit()
+ })
+ }
+
+ /** 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, 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()
Review Comment:
**5. [Low] The writer interrupts ignore the uninterruptible resource
creation that `Task.kill` respects**
The writer thread has the `TaskContext` set (L1468), so
`Utils.createResourceUninterruptiblyIfInTaskThread` applies on it when an
upstream scan opens a file: CSV, JSON, Text, XML, `NewHadoopRDD`, and the DSv2
Parquet and Text readers. `Task.kill` defers its interrupt while
`context.interruptible()` is false, which SPARK-50768 added to avoid leaking
half-built streams. This interrupt and the one at L295 do not check it.
If the interrupt lands in the builder, or the flag is already set when the
writer enters it, the builder can fail after it opened a stream that nothing
closes. For example, an S3A open through `FutureIO.awaitFuture` turns the
interrupt into `InterruptedIOException` and leaves the stream that the pool
thread opened unclosed. And `HadoopCodecStreams.createInputStream` does not
close `fs.open(file)` when `codec.createInputStream` fails.
The first interrupt behaves like master's `cancel(true)`, and `PipedRDD`
interrupts its stdin writer the same way, so this may deserve its own JIRA.
What this PR adds is the repeated interrupt here. Could it at least skip a
writer that is inside `createResourceUninterruptibly`, e.g. by checking
`context.interruptible()` under `TaskContext.synchronized` and retrying at the
next tick?
##########
core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala:
##########
@@ -50,6 +62,356 @@ 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()
Review Comment:
**14. [Nit] `addNow` is not counted down in this `finally`**
If `reached.await` (L189) fails, the writer stays in
`awaitIgnoringInterrupts(addNow)` for the rest of the core test JVM and keeps a
pool thread busy. If completion already reached the stop listener, that thread
also logs a warning and re-interrupts the writer every 10 s. `finally {
addNow.countDown(); release.countDown() }` would release both. The `started`
assertions before the `try` blocks (L138, L184) leave a writer behind in the
same way.
##########
core/src/main/scala/org/apache/spark/TaskContextImpl.scala:
##########
@@ -177,29 +194,48 @@ private[spark] class TaskContextImpl(
}
private[spark] override def markTaskCompleted(error: Option[Throwable]):
Unit = {
- synchronized {
+ val claimedInvocation = synchronized {
if (completed) return
completed = true
+ // Claim the listener invocation together with setting `completed`,
unless another thread is
+ // already invoking listeners. Otherwise a thread that adds a completion
listener in between
+ // would see `completed` and run the listeners itself. That thread can
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.
Review Comment:
**11. [Nit] "which would then wait for itself" no longer matches what
happens**
This was true in `038c4b3420e`, before the same-thread guard. With the guard
in `stopAndAwaitExit`, such a writer now skips the wait, and the listeners
after it free memory it still uses, as the comment in `PythonRunner.scala`
(L303-L312) and the description say. So the two comments describe different
outcomes for the same path. This comment also doesn't say that the `else`
branch below still allows the takeover when another thread holds the invocation.
Maybe: "..., e.g. a pipelined Python UDF writer that adds listeners lazily
while it pulls its input. That thread cannot stop itself, so the listeners
after the stop listener would free memory it still uses. If another thread
already holds the invocation, this does not help (SPARK-60222)."
The scaladoc of `TaskContext.addTaskPreCompletionListener` could also say
which thread runs these listeners: normally the one that marks the task
completed, otherwise the next thread that adds a listener, which can be the
helper thread itself.
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,182 @@ 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 `Future.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))
+ pipelinedWriterThreadPool.submit(writerTask)
+ context.addTaskPreCompletionListener(new TaskCompletionListener {
+ override def onTaskCompletion(context: TaskContext): Unit =
writerTask.stopAndAwaitExit()
+ })
+ }
+
+ /** 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, 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(): Unit = {
+ val wasInterrupted = Thread.interrupted()
+ try {
+ val runningThread = synchronized {
+ state match {
+ case PipelinedWriterNotStarted =>
+ // The writer never started and now never will.
+ finish()
+ None
+ case PipelinedWriterRunning(thread) =>
+ // Interrupts the writer thread. 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).
+ if (thread ne Thread.currentThread()) {
+ thread.interrupt()
+ }
+ Some(thread)
+ case PipelinedWriterFinished => None
+ }
+ }
+ runningThread.foreach { thread =>
+ if (thread eq Thread.currentThread()) {
+ // TaskContextImpl can still let a thread that adds a listener run
the pending
+ // completion listeners: when another thread is invoking listeners
as the task
+ // completes, or while a failing listener's failure listeners run.
If that thread is
+ // this writer, in the middle of pulling its input, the writer can
only exit after
+ // this returns, so waiting would never end. It is not stopped
either: the remaining
+ // listeners run on this thread right after this one and free
memory that the writer
+ // uses again when they return, as before this listener waited at
all.
+ logError(log"Task completion listeners of ${MDC(TASK_NAME,
taskName)} run on its " +
+ log"pipelined Python writer thread while the writer is running;
not waiting " +
+ log"for the writer to exit, which may lead to a use-after-free.")
+ } else {
+ awaitExit()
+ }
+ }
+ } finally {
+ Thread.interrupted()
+ if (wasInterrupted) {
+ Thread.currentThread().interrupt()
+ }
+ }
+ }
+
+ private def awaitExit(): Unit = {
+ // Wait without a time bound. The completion listeners that run next
free memory that
+ // the writer may still access: even a writer blocked in upstream I/O
that ignores
+ // interrupts adds the row it gets next to the HybridRowQueue. Giving up
here would allow
+ // that use-after-free, and failing this listener would not stop those
listeners either.
+ // The price is that task completion is held for as long as such
upstream I/O blocks.
+ //
+ // The writer is interrupted again on every warning, in case code it
called consumed the
+ // interrupt without stopping the writer, e.g. an upstream iterator that
caught an
+ // InterruptedException. The writer would then go on and block in
channel.write, since
+ // nothing reads the Python worker's output anymore.
+ //
+ // An interrupt of this thread that arrives during the wait (e.g. a task
kill) does not stop
+ // it and is not passed on (see stopAndAwaitExit): the kill is already
recorded in the
+ // TaskContext, and the flag would make the blocking channel and lock
operations of the
+ // completion listeners that run next fail and skip cleanup.
+ val startNs = System.nanoTime()
+ var done = false
+ while (!done) {
+ try {
+ done = exited.await(warnIntervalMs, TimeUnit.MILLISECONDS)
+ if (!done) {
+ val waitedMs = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() -
startNs)
+ logWarning(log"Still waiting for the pipelined Python writer of " +
Review Comment:
**13. [Nit] The warning doesn't identify the writer thread or where it is
blocked**
The writer threads are named `python-udf-pipelined-writer-N` and are reused
across tasks, the default log layout has no thread name, and the task reaper
dumps only the task thread, which shows `CountDownLatch.await`. So when this
warning repeats, the logs say neither which thread to look at nor where it is
stuck.
The thread is at hand in `PipelinedWriterRunning(thread)`. Could each
warning include its name, and the first one also its stack? The stack can be
taken outside the monitor and dropped if `exited` reached zero meanwhile, since
the pool thread may already run another task. The `logError` in the same-thread
branch (L310) could likewise carry a throwable with the current stack, to show
which lazy `add*Listener` call handed the listeners to the writer.
##########
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:
**15. [Nit] This test stayed in the PR**
This comment is anchored on "a pre-completion listener added during
completion runs next", which is still in this PR: its listener joins a
non-daemon helper without a timeout (`thread.join()` at L834) while it holds
the invocation. It cannot hang today, since the helper's add returns at once,
but it would if SPARK-60222 makes an add wait for the holder. Could it use
`setDaemon(true)`, `join(timeout)` and `assert(!thread.isAlive)`, like the test
below (L858, L862)? The older "listener registers another listener using a
second thread" test (L685) has the same pattern and could be changed together
with SPARK-60222.
--
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]
