HyukjinKwon commented on code in PR #59332:
URL: https://github.com/apache/spark/pull/59332#discussion_r4242065501
##########
core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:
##########
@@ -183,6 +183,97 @@ 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.
+ */
+ 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)
+ 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)
+
+ override def run(): Unit = {
+ if (claimed.compareAndSet(false, true)) {
+ 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 {
+ exited.countDown()
+ }
+ }
+ }
+
+ 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
+ }
+ // 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.
+ //
+ // An interrupt of this thread (e.g. a task kill) is not restored after
the wait: 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(pipelinedWriterExitWarnIntervalMs,
TimeUnit.MILLISECONDS)
Review Comment:
**Non-blocking (P2):** This wait can block forever when `stopAndAwaitExit`
runs on the writer thread itself, which can happen through the existing
`TaskContextImpl` handoff. `markTaskCompleted` sets `completed = true` in one
`synchronized` block, and `invokeListeners` claims `listenerInvocationThread`
in a later one. The writer thread has the `TaskContext` set
(`PipelinedWriterRunnable.run`) and registers completion listeners lazily while
it pulls input, e.g. `ParquetPartitionReaderFactory` L391 or
`UnsafeExternalSorter` L206. If one of those registrations lands between the
two blocks, `addTaskCompletionListener` sees `completed = true` and the writer
thread becomes the invoker. Pre-completion listeners are popped first, so the
writer then runs its own stop listener:
- `future.cancel(true)` interrupts the writer thread itself, because it is
the `FutureTask` runner.
- `claimed.compareAndSet(false, true)` fails, because `run()` already
claimed the writer.
- This loop swallows the `InterruptedException` and waits on `exited`, which
only the `finally` of `run()`, further up the same stack, can count down.
Meanwhile the task thread's `invokeListeners` sees another invoker and
returns. The task reports success, but none of the remaining completion
listeners run (`worker.stop()`, the `HybridRowQueue` close, the reader closes),
and `MonitorThread` does not step in because `isCompleted()` is already true.
Each occurrence permanently leaks a `python-udf-pipelined-writer` thread and a
Python worker, and this warning is logged every 10 seconds forever. The old
`cancel(true)` + `get()` listener did not hang on the same interleaving,
because `get()` throws `CancellationException` right away. The window is
narrow, but it sits on exactly the path this PR targets (lazy listener
registration on the writer thread), and none of the new tests registers a
listener from the writer after `completed` is set.
**Recommended change:** In TaskContextImpl, have markTaskCompleted claim the
completion-listener invocation slot for the calling thread in the same
synchronized block that sets completed, when no thread already holds it, and
run the listeners under that pre-claimed slot. Add TaskContextSuite coverage
for a listener added from another thread at the moment completion starts, and a
BasePythonRunnerSuite case in which the pipelined writer thread registers a
completion listener after markTaskCompleted has begun.
**Why this works:** Once completed is visible to other threads, the task
thread already owns listenerInvocationThread. A listener added by the writer
thread at that point is only pushed, so the task thread pops and runs the
writer's stop listener, waits for the real writer to exit, and then runs the
newly queued listener together with the rest.
**Scope:** Make the task thread the guaranteed owner of completion-listener
invocation from the moment the task is marked completed, and cover the late
writer-thread registration.
**Compatibility:** A listener added after the task's listeners have all run
is still called immediately, and pre-completion listeners still run before
every pending completion listener.
**Risks:** invokeListeners also serves failure and post-status-update
listeners. Pre-claiming must apply only to the markTaskCompleted entry point,
or those invocations could stop running. If another thread already holds the
slot when markTaskCompleted runs, the existing return-early behavior must be
kept.
**Constraints:** Listeners must still run one at a time and never while
holding the TaskContext lock. Pre-completion listeners must still run before
every pending completion listener.
**Success:** When the task thread calls markTaskCompleted and no other
thread is invoking listeners, the task thread holds the listener-invocation
slot as soon as completed becomes visible, so a listener added from another
thread at that point is only queued. The pipelined writer's stop listener never
runs on the writer thread during markTaskCompleted. markTaskCompleted returns
only after the writer has exited and every queued pre-completion and completion
listener has run. A listener added after all listeners have run still runs
immediately on the calling thread.
**Verification:**
- **Regression:** A completion listener registered from the running
pipelined writer thread after task completion has started is queued, not
executed on the writer thread. Task completion then finishes after the writer
exits, and the listener observes that the writer has exited.
- **Behavior:** A listener added from another thread while markTaskCompleted
is starting runs on the task thread, after the pending pre-completion listeners.
- **Compatibility:** Existing listener ordering, add-during-completion,
add-after-completion, and failure-aggregation behavior of task completion,
failure, and post-status-update listeners is unchanged.
--
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]
