viirya commented on code in PR #59332:
URL: https://github.com/apache/spark/pull/59332#discussion_r4244149268
##########
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:
Good catch, confirmed. Fixed in 038c4b3420e: `markTaskCompleted` now claims
`listenerInvocationThread` in the same `synchronized` block that sets
`completed`, unless another thread already holds it, in which case the existing
return-early behavior is kept. Only the `markTaskCompleted` entry point
pre-claims; failure and post-status-update listeners are unchanged. A listener
the writer thread adds once completion has started is therefore only queued,
and the task thread runs it after the writer's stop listener.
Tests:
- `TaskContextSuite`: a thread adds a completion listener as soon as
`isCompleted()` turns true, repeated 500 times, and a listener registered
before completion must run on the task thread. Without the pre-claim, this
fails on the first run.
- `BasePythonRunnerSuite`: the pipelined writer adds a completion listener
after task completion has reached its stop listener. Task completion keeps
waiting, and the listener runs on the completing thread after the writer has
exited.
One hand-off point remains: when a listener throws, `invokeListeners` clears
`listenerInvocationThread` while `markTaskFailed` runs the failure listeners.
The writer could take over there only if a pre-completion listener that runs
before its stop listener throws, and none of the current pre-completion
listeners (the writer stop listeners) throw. I left that path as is, since
changing it would affect how failure listeners nest inside listener invocation.
--
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]
