[SPARK-60113][CORE][SQL] Wait for the pipelined Python UDF writer to exit before task completion cleanup

[viirya]

What changes were proposed in this pull request?

In pipelined Python UDF mode (spark.python.udf.pipelined.enabled), the listener that stops the writer thread now waits until the writer has actually exited, and it runs before every task completion listener.

1. Add TaskContext.addTaskPreCompletionListener (private[spark], abstract; BarrierTaskContext delegates). TaskContextImpl.markTaskCompleted runs these listeners on its own thread before it marks the task completed, outside the listener invocation used for the other listeners. Only then does it set completed (in the same synchronized block that finds no more pre-completion listeners) and run the completion listeners as before. While the pre-completion listeners run, a listener that another thread adds sees the task as not completed and is only queued, so that thread can never start running the listeners. A pre-completion listener added after completion is handed to addTaskCompletionListener, so it runs like a completion listener added at that point. A failing one marks the task failed, and its failure is reported together with those of the completion listeners. Tasks without pre-completion listeners behave as before.
2. The writer is wrapped in BasePythonRunner.PipelinedWriterTask. Its state (not started, running on a thread, or finished) is kept in a single field guarded by the task, and a latch is counted down on every transition to finished. The stop listener is added with addTaskPreCompletionListener, and the writer is submitted only if the task has not completed by then. The listener interrupts the writer thread directly and waits on that latch. If it runs before the pool has started the writer, it marks the writer finished, so the writer never runs.
3. The wait has no time bound, the same as the WriterThread.join() that the writer thread used before SPARK-44705 (added for SPARK-33277). While it waits, it logs a warning every 10 seconds that names the writer thread (with its stack trace the first time), and it interrupts the writer again, in case code it called consumed the interrupt. Like Task.kill, it does not interrupt while context.interruptible() is false, which is while some thread of the task, not necessarily the writer, creates a resource through TaskContext.createResourceUninterruptibly; the writer is interrupted at a later warning then. A failure to interrupt the writer or to log the wait does not end the wait; it is thrown once the writer has exited. An interrupt of the listener thread that arrives during the wait does not stop it and is not passed on to the completion listeners that run next (the kill is already recorded in the TaskContext); an interrupt that was pending before the wait is kept. If the listener ever ran on the writer thread itself, it would not wait, since the writer can only exit after it returns, and would log an error.
4. The writer writes through a CloseShieldOutputStream (added by SPARK-60305) around its DirectByteBufferOutputStream. The Arrow writers write through Channels.newChannel(dataOut), an interruptible channel. When the stop listener interrupts the writer during such a write, the JDK closes that channel and the stream under it on the interrupting thread, and the buffer was freed while the writer still copied into it. Now the buffer is freed only in the writer's own finally.
5. A fatal error of the writer (e.g. OutOfMemoryError), including one while it sets up its buffer, is passed to the reader through writer.setException with a socket shutdownOutput(), and it is logged. Previously the thread pool's FutureTask kept it, and the reader waited for the Python worker.

Why are the changes needed?

The existing listener does:

writerFuture.cancel(true)
try { writerFuture.get() } catch { case _: CancellationException | ... => }

Its comment says this waits for the writer to exit before the listeners that run after it (registered earlier, LIFO) free the memory backing the input rows. But once cancel(true) succeeds, FutureTask.get() throws CancellationException right away and does not wait for the interrupted runnable to return. So the writer can still be serializing a row, or adding it to a HybridRowQueue page, when the queue's close listener frees that page. That is a use-after-free, and with off-heap memory it can crash the executor. A HybridQueue.createNewQueue racing with close() can also add a page that nobody closes until cleanUpAllAllocatedMemory.

A task completes while the writer is still running when the output is not fully consumed, for example with a LIMIT or when the task fails.

Ordering the writer listener in the completion listener stack is not enough. Listeners run in reverse order of registration, and the writer thread itself adds completion listeners lazily while it pulls the upstream iterator, e.g. ParquetReaderCallback.initIfNotAlready of the DSv2 Parquet reader (which closes the vectorized reader and its off-heap column vectors) or the UnsafeExternalSorter that ExternalAppendOnlyUnsafeRowArray creates once it spills. Those would run before the writer listener. Running the stop listener inside the usual listener invocation is not enough either: TaskContextImpl can let another thread run the pending listeners, e.g. when a kill thread is running failure listeners as the task completes (SPARK-60222). If that thread were the writer, it would run its own stop listener. Running the pre-completion listeners before the task is marked completed avoids both.

The wait cannot be bounded safely: if the listener gave up, the listeners that run next would still free the memory, and throwing from the listener does not stop them either. Even a writer blocked in upstream I/O that ignores interrupts (e.g. a slow JDBC fetch) adds the row it gets next to the HybridRowQueue, whose InMemoryRowQueue writes through the page's base object captured at construction, so a freed page would be corrupted even on-heap (HeapMemoryAllocator pools pages for reuse). The price is that task completion is held for as long as such I/O blocks, including I/O that only a close listener of the upstream reader would unblock, since that listener now runs after the wait; spark.task.reaper.enabled with a positive spark.task.reaper.killTimeout bounds this for killed tasks, by killing the executor JVM, except in local mode, where the reaper only logs an error. This is documented in the scaladoc and the warning. Usually the writer exits soon after the interrupt: channel.write fails with ClosedByInterruptException, and otherwise the writer loop checks the interrupt flag between writeNextInputToStream calls, each of which can pull and serialize a whole batch (e.g. for the Arrow runners).

The pipelined mode was added in 4.3.0 by SPARK-56642, so branch-4.3 and branch-4.x have the same issue. The same interrupt-closes-the-buffer problem also affects the default (non-pipelined) Arrow UDF path, where Task.kill interrupts the task thread while it serializes into ReaderInputStream's buffer; SPARK-60305 fixed that separately, since it needs backports to 4.0 through 4.3, and added the CloseShieldOutputStream that this change also uses.

Does this PR introduce any user-facing change?

No. This fixes a bug in an opt-in mode that is not released yet (4.3.0). For tasks with no pre-completion listener, listener behavior is unchanged.

How was this patch tested?

New tests in BasePythonRunnerSuite:

• Task completion stays blocked until a running writer that ignores interrupts exits. Both a cleanup listener added on the task thread before the writer starts and one that the writer thread adds after the stop listener see that the writer has exited. The tests first wait until task completion has reached the writer listener.
• A listener that the writer thread adds after task completion has started is queued and runs after the writer exits.
• Stopping interrupts a running writer right away; an interrupt of the stopping thread, whether it arrives during the wait or while stopping, does not stop it and is not passed on; an interrupt pending before stopping is kept.
• The writer is not interrupted while it creates a resource uninterruptibly, and is interrupted afterwards.
• A writer stopped before it starts never runs, also when it is stopped twice, when it is started after the task completed, or when it is started while another thread runs the completion listeners. Stopping on the writer thread itself does not wait.
• A failure to interrupt the writer does not end the wait, and is thrown once the writer has exited.
• A writer that consumed the first interrupt is interrupted again.
• An interrupted Arrow-style channel write in a later buffer leaves the buffer usable, through a real PipelinedWriterRunnable (with a socket pair standing in for the Python worker).
• A warning that fails on every tick does not skip the repeated interrupt.
• A fatal error of a real PipelinedWriterRunnable is rethrown, logged, and passed to the reader, and the Python side sees EOF.

New tests in TaskContextSuite:

• Pre-completion listener ordering and failure handling, and a pre-completion listener added after completion running like a completion listener, or failing the task and throwing if it fails. Failures of pre-completion and completion listeners are merged into one exception.
• While a pre-completion listener runs, listeners added from another thread are only queued and run on the completing thread afterwards, and failure listeners triggered from another thread still run.

Before the fix, a test with the old listener code showed markTaskCompleted returning while the writer was still running. Each of the following changes, reverted on its own, makes its test fail: the cancel-then-get() wait, adding the stop listener as a pre-completion listener, running the pre-completion listeners before setting completed, the close shield, the immediate interrupt, skipping the interrupt during uninterruptible resource creation, the repeated interrupt, clearing and keeping the interrupt status, and passing fatal errors to the reader.

Also ran the whole core test suite, the sql suites that use markTaskFailed, markTaskCompleted or markInterrupted (JdbcTaskInterruptSuite, SupportsArchiveFormatSuite, JsonSuite, ParquetStorageFilterSuite, ValueStateSuite, RocksDBStateStoreSuite, RocksDBStateStoreLockHardeningSuite, RocksDBCheckpointFailureInjectionSuite), and pyspark.sql.tests.pandas.test_pipelined_udf.

Was this patch authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Claude Opus 5.5)

[dongjoon-hyun]

Thank you for working on this, @viirya. I left 8 inline comments. Here is a summary.

1. (PythonRunner.scala L203) The wait only covers listeners registered before startPipelinedWriter. Listeners that the writer thread itself registers lazily while pulling upstream input (e.g. ParquetReaderCallback.initIfNotAlready of the DSv2 Parquet reader, or UnsafeExternalSorter's cleanupResources() when ExternalAppendOnlyUnsafeRowArray switches to the spillable sorter) are pushed later, so they run first under LIFO and can free memory while the writer is still running and not even interrupted yet.
2. (L239) The unbounded wait can block task completion indefinitely when the writer is stuck in upstream I/O that ignores interrupts (e.g. JDBC ResultSet.next(), HTTP/S3 stream reads). The task reaper is disabled by default and only handles killed tasks, not a normal LIMIT completion.
3. (L259) Restoring the interrupt flag means all remaining (earlier-registered) cleanup listeners run with the flag set, so NIO/lock-based cleanup in them can fail.
4. (BasePythonRunnerSuite.scala L94) The blocking assertions can pass even when the completion thread has not reached the listener yet, so the test may not catch a regression on a slow CI machine.
5. (BasePythonRunnerSuite.scala L141) The race where the pool thread enters FutureTask.run and wins the claim after cancel(true) is not tested.
6. (L240) The hand-written interrupt-preserving loop can be replaced with Guava Uninterruptibles.awaitUninterruptibly, which is already used in core.
7. (L204) A fatal error thrown by the writer is captured by the FutureTask and never logged or propagated.
8. (L231) The comment says the writer checks the interrupt flag between rows, but it checks only between writeNextInputToStream calls (a whole batch for Arrow runners).

[dongjoon-hyun]

This orders the writer listener only against the listeners registered before this line. In pipelined mode, the upstream iterator is first pulled on the writer pool thread, so any completion listener that upstream operators register lazily is pushed after this one and runs before it under LIFO, while the writer is still running and not even interrupted yet. For example:

• DSv2 Parquet: ParquetReaderCallback.initIfNotAlready registers closeCurrent() when the first file is opened (ParquetPartitionReaderFactory.scala L391), which closes the vectorized reader (off-heap column vectors when spark.sql.columnVector.offheap.enabled=true).
• ExternalAppendOnlyUnsafeRowArray (SMJ / Window upstream) creates an UnsafeExternalSorter lazily once the threshold is exceeded, and its constructor registers cleanupResources().

In both cases, a LIMIT that completes the task frees the memory under the running writer, which is the same use-after-free that this PR is trying to prevent. The new tests register the cleanup listener on the task thread before the writer starts, so they don't cover this. Could we stop and wait for the writer before any completion listener runs (e.g. a pre-completion hook in TaskContextImpl), instead of relying on the position in the LIFO stack?

[viirya]

Good catch, confirmed. DataSourceRDD.compute creates the reader inside flatMap, so ParquetReaderCallback.initIfNotAlready runs on the writer thread and its closeCurrent() lands above the writer listener in the stack.

Added TaskContext.addTaskPreCompletionListener (private[spark], default falls back to addTaskCompletionListener). TaskContextImpl.invokeListeners now always takes the next listener from the pre-completion stack first, so these run before every completion listener, including ones added later or from other threads, and also when a pre-completion listener is added while another thread is already running the completion listeners. The writer listener uses it. The test now adds a cleanup listener from the writer thread itself, and it fails if the writer listener is switched back to an ordinary completion listener.

As far as I can tell, the WriterThread listener before SPARK-44705 had the same gap, since it was also registered before the writer thread started pulling upstream.

[dongjoon-hyun]

Since the listener no longer looks at the Future, a fatal Throwable escaping PipelinedWriterRunnable.run (NonFatal excludes VirtualMachineError, LinkageError, etc.) is stored in the FutureTask and silently dropped. In that case the writer also doesn't call shutdownOutput() or writer.setException(...), so the Python worker keeps waiting for input and the reader waits for output with no root cause recorded. Since the listener now knows the writer has exited, could we log the cause here (e.g. future.get() after the latch, ignoring CancellationException)?

[viirya]

Fixed in two places. PipelinedWriterRunnable now also calls writer.setException(t) and shutdownOutput() for fatal errors before rethrowing them, so the reader fails with the cause instead of waiting for Python. PipelinedWriterTask.run logs anything that escapes the writer. A future.get() after the latch would not be enough here: once cancel(true) has succeeded, the future's outcome is no longer available.

[dongjoon-hyun]

nit. checks the interrupt flag between rows is not accurate. The loop checks it only between writeNextInputToStream calls. For Arrow runners, one call serializes a whole batch (up to arrowMaxRecordsPerBatch rows) and pulls from upstream, and ColumnarArrowEvalPythonEvaluatorFactory Path 2 copies all rows of an input ColumnarBatch into the HybridRowQueue in one step without checking the interrupt. The same wording is in the PR description. Since the wait below is unbounded, it would be good to describe it precisely.

[viirya]

Fixed the comment: it now says that the interrupt is seen only between writeNextInputToStream calls, each of which can pull and serialize a whole batch. Will fix the PR description as well.

[dongjoon-hyun]

This unbounded wait is a behavior change for writers blocked in upstream I/O that does not respond to interrupts, e.g. JDBC ResultSet.next() on a slow fetch, or an HTTP/S3 stream read with a long (or no) socket timeout. When a LIMIT is satisfied, the task completes normally, cancel(true) cannot unblock the read, and markTaskCompleted now waits until that read returns, holding the task slot and only logging a warning every 10 seconds. Previously the listener returned immediately.

Also, the task reaper only applies to killed tasks and spark.task.reaper.enabled is false by default, so nothing recovers the normal completion case. In these examples the rows are on-heap, so the wait doesn't buy any safety. Could we mention this trade-off explicitly, or consider limiting the wait to the cases where the input memory can actually be freed under the writer?

[viirya]

The trade-off is real, and I documented it in the comment. I still kept the wait unbounded, because in these examples the wait does buy safety. In BatchEvalPython / ArrowEvalPython, the row the writer gets next from the blocked upstream read goes to queue.add before the loop checks the interrupt again. InMemoryRowQueue writes through the page's base object that it captured at construction. So even for on-heap memory, a freed page (at least the page size, which HeapMemoryAllocator pools and hands to other consumers) would be silently corrupted, and off-heap it would crash. I couldn't find a reliable way to tell when giving up would be safe. The WriterThread.join() before SPARK-44705 was also unbounded, for every Python UDF.

[dongjoon-hyun]

nit. The interrupted flag, the outer try/finally, and the inner catch InterruptedException re-implement Guava's Uninterruptibles.awaitUninterruptibly(CountDownLatch, long, TimeUnit), which also restores the caller's interrupt status before returning. core already uses Uninterruptibles (e.g. AsyncEventQueue.scala). This could be simplified to something like:

while (!Uninterruptibles.awaitUninterruptibly(
    exited, pipelinedWriterExitWarnIntervalMs, TimeUnit.MILLISECONDS)) {
  logWarning(...)
}

[viirya]

Following your next comment, the wait no longer restores the interrupt status, while Uninterruptibles.awaitUninterruptibly always restores it. So I kept the loop, which is now just a catch-and-continue without the outer try/finally.

[dongjoon-hyun]

Restoring the interrupt flag here means that every listener registered earlier (which runs after this one) runs with the flag set, e.g. worker.stop() (L560), queue.close() of the HybridRowQueue (which deletes spill files), and the reader/allocator closes. Since this listener can now take much longer, a task-kill interrupt is more likely to land during this wait and be carried over. Any blocking NIO channel operation or lock wait in those listeners would then fail with ClosedByInterruptException / InterruptedException, and the cleanup would be skipped or logged as a listener failure. Previously, the old listener didn't wait at all. Is this intended? If so, could we add a comment about the impact on the following listeners? The new test asserts the propagation, so it would lock in this behavior.

[viirya]

Not intended. Changed it to not restore the flag, with a comment: the kill is already recorded in the TaskContext, and the old listener also swallowed the interrupt (get() threw InterruptedException). The test now asserts that a later completion listener does not see the flag.

[dongjoon-hyun]

completion.isAlive and exitedBeforeCleanup.isEmpty are also true when the completion thread has not reached the writer listener yet. On a slow CI machine where the thread isn't scheduled within 500 ms, this test would pass even with the old cancel(true) + get() listener. How about registering a listener after startPipelinedWriter (which runs first under LIFO) that counts down a latch, and waiting on it before checking isAlive? The same applies to completion.join(200) + interrupt() at L121, which doesn't verify that the interrupt arrives during the wait.

[viirya]

Fixed. The tests add a pre-completion listener after startPipelinedWriter, which runs right before the writer listener and counts down a latch. The tests wait for that latch before checking isAlive, and the interrupt test also interrupts only after it.

[dongjoon-hyun]

This FutureTask is never run by a pool thread, so the test covers only the case where FutureTask.run sees the cancelled state and never calls PipelinedWriterTask.run. The real race isn't covered: the pool thread has already entered FutureTask.run, cancel(true) interrupts it, and run() still wins the claim and starts the writer with the interrupt flag already set. In that case the listener must wait for the latch, so a test for it would protect against a future change that breaks the countdown on that path.

[viirya]

Added "stopping waits for a writer that started after the interrupt": the writer is claimed by run() on a thread whose interrupt flag is already set, and stopAndAwaitExit must block until it exits.

[HyukjinKwon]

Review summary

All eight points from the earlier review are addressed in the current head. The writer stop now runs from a pre-completion stack ahead of every completion listener. The unbounded wait and its trade-off are documented. The interrupt is no longer restored. The blocking tests now wait until completion reaches the writer listener, and the claim-after-interrupt path has its own test. The hand-written wait loop is justified. Fatal writer errors are propagated and logged. The interrupt-granularity comment is corrected.

One new issue comes from the new wait itself (see the inline comment on stopAndAwaitExit). If the writer thread registers a completion listener in the short gap between markTaskCompleted setting completed and claiming the listener-invocation slot, the writer becomes the invoker, runs its own stop listener, and waits forever for itself. The task thread then skips every remaining cleanup listener. The race is narrow, but when it hits, it leaks a thread and a Python worker permanently instead of causing a delay.

Findings

1 total: 0 P0, 0 P1, 1 P2, 0 P3.

Non-blocking (P2)

• The pipelined writer can wait forever on itself when it becomes the listener invocation thread — core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:262 — see inline.

Existing discussions

• existing discussion — The current head addresses the eight items in this summary. (1) TaskContextImpl adds onPreCompleteCallbacks, which invokeListeners always drains first, and the writer stop is registered with addTaskPreCompletionListener. (2) The unbounded wait is kept, but the trade-off is documented in the stopAndAwaitExit comment and the PR description. (3) The interrupt is no longer restored after the wait, and a test asserts that later listeners do not see it. (4) The tests wait on a reached latch from addReachedListener before checking isAlive. (5) 'stopping waits for a writer that started after the interrupt' covers the claim-after-interrupt path. (6) The author explained why the hand-written loop stays (Uninterruptibles would restore the flag). (7) PipelinedWriterTask.run logs fatal errors, and PipelinedWriterRunnable now calls setException and shutdownOutput before rethrowing. (8) The comment and PR description now say that the interrupt is seen between writeNextInputToStream calls. The fix for item 1, however, exposes a related new defect. Item 1 identified listeners that the writer thread registers lazily. When such a registration happens after completed is set and before the task thread claims the invocation slot, the writer thread becomes the invocation thread and runs the new pre-completion stop listener on itself, waiting forever for its own exit.

Generated by Omnigent on Databricks.

[HyukjinKwon]

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.

[viirya]

Good catch, confirmed. Fixed in 038c4b3: 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.

[HyukjinKwon]

Thanks, I checked 038c4b3. markTaskCompleted now takes listenerInvocationThread in the same synchronized block that sets completed, only that entry point pre-claims, and the new TaskContextSuite and BasePythonRunnerSuite tests cover a listener added just as completion starts and one added by the writer after completion has reached its stop listener. I agree the listener-failure handoff you describe is fine to leave: it needs a pre-completion listener that throws, and the stop listener does not throw in practice.

[HyukjinKwon]

Review summary

The earlier P2 is fixed for the window it described. markTaskCompleted now claims the listener-invocation slot together with completed, and the new tests cover a listener added just as completion starts and one added by the writer after completion has reached its stop listener. The same self-wait is still reachable on the remaining branch: when another thread already holds the slot at markTaskCompleted, for example the kill thread running failure listeners after an interrupt listener throws. That residual is a non-blocking P2.

Two P3 items concern tests. The fatal-error branch of PipelinedWriterRunnable has no regression test, and one comment in TaskContextSuite contradicts the order its test asserts. The eight items from the first review remain addressed. There is also one question below about what happens if the writer's interrupt is lost.

Cross-cutting observations

• Both remaining hangs come from the unbounded stop-listener wait. The only things that can end it are the writer's own exit and, when the writer is stuck on the Python socket, worker.stop(), and both can end up on the far side of the wait. The self-wait does this through the listener handoff, the lost-interrupt cycle through a writer blocked on a socket whose closer is queued behind the wait.
  Separation: Different triggers and remedies. The self-wait needs a same-thread guard and is a defect, while the cycle depends on whether a supported upstream loses the interrupt, which only the author can settle. Neither repair closes the other, and the accepted unbounded-wait decision stays as it is.

Findings

3 total: 0 P0, 0 P1, 1 P2, 2 P3.

Non-blocking (P2)

• The writer can still wait forever on itself when another thread holds the listener slot at completion — core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:239 — see inline.

Nit (P3)

• The fatal-error branch of PipelinedWriterRunnable has no test — core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:1412 — see inline.
• Test comment says the opposite of the asserted order — core/src/test/scala/org/apache/spark/scheduler/TaskContextSuite.scala:772 — see inline.

Re-review status

Prior AI findings: 0 addressed, 1 still present; additional unresolved findings in this review: 2.

New attribution: 0 newly introduced, 2 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

• The writer can still wait forever on itself when another thread holds the listener slot at completion — core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala:239

Existing discussions

• existing discussion — The trade-off is now documented in the stopAndAwaitExit comment and the PR description, and the author explained why the wait stays unbounded. A separate case is raised as a new question: a lost interrupt leaves the writer blocked on the Python socket, and worker.stop() is queued behind the wait.
• existing discussion — All eight items are addressed in the current head. (1) The stop listener is registered with addTaskPreCompletionListener and runs before every completion listener. (2) The unbounded-wait trade-off is documented in stopAndAwaitExit and in the PR body. (3) The interrupt flag is no longer restored. (4) The tests wait on addReachedListener before checking isAlive. (5) A test covers a writer claimed after the interrupt. (6) The hand-written wait is kept, with a reason. (7) Fatal writer errors are now propagated and logged. (8) The comment describes the writeNextInputToStream granularity correctly. Item 1's mechanism, the writer thread lazily registering listeners, still leads to a self-wait on the unclaimed markTaskCompleted path, so it is bound to the new finding as a related new defect.

Decision challenges

Can a lost writer interrupt turn the unbounded wait into a permanent hang?

stopAndAwaitExit interrupts the writer only once, through future.cancel(true), and then waits with no bound while ignoring interrupts. Everything that could otherwise unstick the writer now comes after this wait: worker.stop() (which closes the worker channel) and the upstream close listeners are queued behind it, and MonitorThread backs off once isCompleted() is true. The writer loop stops only on the interrupt flag; nothing checks task completion (ContextAwareIterator has no users). So suppose upstream code on the writer thread catches the InterruptedException from cancel(true) and carries on without restoring the flag. The writer then keeps pulling rows and writing batches. The task thread has stopped reading Python's output, so Python blocks writing it and stops reading input, and the writer blocks in channel.write with no interrupt pending. That hold is permanent, not bounded by upstream I/O, and the task never reports its result. Before this PR, worker.stop() ran right after cancel(true) and broke this cycle.

I verified each step from the code except one: whether a supported upstream source actually swallows the interrupt. Do you know of one, or is that safe to rule out? If it can happen, re-interrupting the writer thread (or shutting down the worker channel) on each 10-second warning tick would keep the wait unbounded but break the cycle, without freeing any task memory.

Generated by Omnigent on Databricks.

[HyukjinKwon]

Non-blocking (P2): The new claim in markTaskCompleted closes the window from the last round when the slot is empty. When another thread already holds listenerInvocationThread, though, the task thread still returns without running any completion listener, and whichever thread adds a listener next inherits them. That can happen when Task.kill runs markInterrupted on the kill thread and a TaskInterruptListener throws (addTaskInterruptListener is public), so markTaskFailed runs the failure listeners on that thread. It can also happen when a failure listener is added from another thread after the task failed. If the pipelined writer, which nobody has stopped yet, then lazily adds a completion listener (e.g. ParquetReaderCallback.initIfNotAlready), it claims the slot and pops this pre-completion listener first. cancel(true) interrupts the writer itself, the claimed CAS fails, and the loop waits on exited, which only the finally of run() further up the same stack can count down. The pool thread and the Python worker leak permanently, and worker.stop(), the HybridRowQueue close, and the reader closes never run. The old cancel(true) + get() returned at once on this path. The listener-failure handoff you mentioned ends in the same wait if it is ever hit.

Recommended change: In PipelinedWriterTask, remember the thread that claimed and is running the writer. In stopAndAwaitExit, after cancel(true), if the current thread is that writer thread, log that task completion cleanup is running on the writer thread and return without waiting, because the writer can only exit after this listener returns. Add Python runner unit-test coverage that invokes the stop listener on the writer's own thread while the writer is still running.

Why this works: On the writer thread, the latch can only be counted down by run()'s finally block further up the same stack, so waiting there can never succeed. Skipping the wait on that one thread restores the merge-target behavior for this handoff: the listener returns, the remaining cleanup runs, and the writer exits on its self-interrupt. Every other thread keeps the unbounded wait.

Scope: Add a same-thread guard to the pipelined writer stop listener and a regression test for it.

Compatibility: The unbounded wait from the task thread and the run()/stopAndAwaitExit claim protocol stay as they are.

Risks: On the self-invocation path, the remaining completion listeners still run on the writer thread while it is in the middle of pulling input, as in the merge target. The guard prevents the hang; it does not add a use-after-free safety net for that narrow path.

Constraints: Keep the unbounded, interrupt-ignoring wait unchanged for every thread other than the writer. Keep the claim protocol between run() and stopAndAwaitExit, including the never-started case.

Success: stopAndAwaitExit never blocks when it is invoked on the thread that is running this writer. Invoked from any other thread, it still interrupts the writer and returns only after the writer has exited or is guaranteed never to run.

Verification:

• Regression: Invoking the stop listener on the writer's own thread while that writer is still running returns instead of hanging, and the writer then finishes and counts down its exit latch.
• Behavior: Task completion from another thread still waits for a running writer, skips a writer that never started, and waits for a writer claimed after the interrupt.

[viirya]

Thanks, confirmed. Fixed in 1053b61 on both sides. stopAndAwaitExit now has the same-thread guard you suggested: PipelinedWriterTask records the thread running the writer, and if the stop listener runs on that thread, it logs a warning and returns after cancel(true) instead of waiting. The new test "stopping on the writer thread itself does not wait for the writer" covers it.

The handoff that led there is also closed in TaskContextImpl (see dongjoon's comments on L322 and L208). When another thread holds the listener invocation, markTaskCompleted now waits for it and gets the invocation handed over directly, instead of returning. A failing listener no longer releases the invocation while the failure listeners run. So the guard is a safety net rather than the main mechanism.

[HyukjinKwon]

Thanks, I checked 7ab7ae2. stopAndAwaitExit no longer waits when it runs on the thread recorded in PipelinedWriterRunning, and that branch now logs at ERROR and clears the self-interrupt; "stopping on the writer thread itself does not wait for the writer" exercises it through a real FutureTask. The handoff that led there is closed at the source too: markTaskCompleted now waits for the holder and takes the invocation over directly, and a failing listener keeps the invocation while the failure listeners run, so the guard is only a safety net now.

[HyukjinKwon]

Nit (P3): With this change, a fatal error goes through setException, shutdownOutput(), and the rethrow below, but no test runs this branch. a fatal error of the pipelined writer is logged wraps a bare Runnable in PipelinedWriterTask, so narrowing this back to NonFatal(t), or dropping the rethrow, would keep every test green. The reader would then again wait on Python with no cause recorded, which is the hang raised earlier in review.

Verification:

• Regression: With the fatal-error branch narrowed back to NonFatal, or with the rethrow removed, the new coverage fails.

[viirya]

Added "a fatal error of the pipelined writer is passed to the reader". It runs a real PipelinedWriterRunnable against a socket pair, with a Writer whose writeNextInputToStream throws a LinkageError, and checks that the error is rethrown, that writer.exception holds it, and that the Python side reads EOF. It fails if the branch is narrowed back to NonFatal or if the rethrow is removed.

[HyukjinKwon]

Thanks, verified. The new test drives a real PipelinedWriterRunnable against a socket pair and checks both writer.exception and the EOF on the Python side. Narrowing the branch back to NonFatal would leave writer.exception empty, and dropping the rethrow would fail the intercept[LinkageError], so both regressions are caught.

[HyukjinKwon]

Nit (P3): B is an ordinary completion listener added at this exact point, and the assertion below expects it to run after P2 and P1. Read literally, this sentence says the opposite of what the test checks. Something like: "Had P1 and P2 been added with addTaskCompletionListener, B would run before them."

Verification:

• Inspection: The comment names the condition (P1 and P2 being ordinary completion listeners) under which B would run before them.

[viirya]

Thanks, reworded as suggested.

[HyukjinKwon]

Thanks, the comment now matches the asserted order.

[dongjoon-hyun]

Thank you for the update, @viirya. addTaskPreCompletionListener and claiming the invocation in markTaskCompleted close the main gap from the previous round. I left 8 inline comments on the new head (038c4b3420e). Summary:

1. Failure path can still hand the invocation to the writer

   invokeListeners still clears listenerInvocationThread before markTaskFailed. A still-running writer that adds a listener in that window takes over and runs its own stop listener, so it waits for itself.

2. An unclaimed markTaskCompleted runs no completion listeners

   If another thread holds the invocation (e.g. the kill thread running failure listeners), the stop listener and all completion listeners are skipped, and task memory can be freed while the writer is still running.

3. The wait now runs before listeners that could unblock the writer

   Close listeners that the writer registers lazily used to run first and could unblock a writer stuck in non-interruptible upstream I/O. Now they run only after the unbounded wait.

4. The TaskContext doc overstates the ordering guarantee

   "Regardless of when" is contradicted by the new test that expects Seq("B", "P", "A").

5. Interrupt clearing depends on whether the writer was claimed

   CountDownLatch.await throws on a pending interrupt even when the latch is already zero, so later listeners see different interrupt states depending on a race.

6. The takeover regression test does not force the race

   The 500-iteration busy-spin probably passes without the fix too, and the failure-path window in item 1 is not covered.

7. Fatal writer errors are now reported twice

   They go through setException and shutdownOutput, are rethrown, and are logged by PipelinedWriterTask while the reader rethrows them on the task thread.

8. addTaskPreCompletionListener duplicates addTaskCompletionListener

   A private helper that takes the target stack would keep the two in sync.

[dongjoon-hyun]

The claim in markTaskCompleted closes the takeover window at completion start, but the failure path still opens one. When a listener throws, the catch block (L370) sets listenerInvocationThread = None before markTaskFailed(e) and only reclaims it in the finally.

Suppose a task has two pipelined writers, W1 and W2. W2's stop listener (or any pre-completion listener registered after W1's) throws while W1 is still running. If W1 then lazily adds a listener in that window, it sees completed and an empty slot, becomes the invocation thread, and pops its own stop listener. stopAndAwaitExit then waits on W1's own exited latch forever. The task thread also keeps popping listeners at the same time, which breaks the sequential guarantee.

Could the catch block keep this thread registered while it calls markTaskFailed (e.g. let the failure invocation reuse the already-claimed slot) instead of clearing it?

[viirya]

Done as you suggested. When a listener throws, the catch block no longer clears listenerInvocationThread. It calls markTaskFailed in a nested mode, which runs the failure listeners as part of the current invocation and keeps the registration when they are done. Since nothing else cleared it, the old reclaim in finally is gone, and so is the window in which two threads could pop listeners at once. In the old code, markTaskFailed reclaimed the invocation for the failure listeners right away, so the windows were the gaps around that reclaim rather than the whole time failure listeners run. The new test "a failing listener does not let another thread take over the listeners" therefore uses contention: two pre-completion listeners throw while another thread keeps adding listeners, repeated 200 times, and the pending listener must run on the task thread. With the old catch block restored, it fails on the first run.

[dongjoon-hyun]

When another thread already holds the invocation, claimedInvocation is false, so invokeListeners returns immediately. The holder only drains onFailureCallbacks, which means neither the pre-completion stop listener nor any completion listener runs.

For example, a kill arrives and an interrupt listener fails. markInterrupted then calls markTaskFailed on the kill thread, which runs the failure listeners. If the task thread reaches markTaskCompleted(None) meanwhile, it returns without running anything. Task.run finishes, and the executor's cleanUpAllAllocatedMemory frees the pages while the pipelined writer may still be copying rows (SPARK-33277). If the writer adds a listener later, it takes over and waits for itself.

This predates the PR, but the PR now depends on the stop listener always running on the completing thread. Should markTaskCompleted wait for the current holder and then run the completion stacks, or should the holder also drain the completion stacks once completed is set?

[viirya]

Went with the first option. If another thread holds the invocation, markTaskCompleted registers itself as the next invoker and waits. The holder hands the invocation to it directly when it finishes, so no third thread (such as the writer) can take it in between. The completing thread then runs the pre-completion and completion listeners, and markTaskCompleted returns only after they have run, so Task.run no longer reaches cleanUpAllAllocatedMemory first. If markTaskCompleted is called from a listener on the thread that already holds the invocation, it runs the completion listeners as part of that invocation instead of waiting. An interrupt during the wait is restored afterwards.

The one risk is a holder that itself waits for the task thread. Only threads running failure listeners can hold the invocation at that point (e.g. the kill thread after an interrupt listener failed), and a failure listener that blocks on the task thread seems unlikely. The new test "completion waits for another thread invoking listeners and runs them" reproduces your example: a failing interrupt listener on a kill thread, with a failure listener that blocks. markTaskCompleted stays blocked, a listener added from a third thread does not take over, and once the failure listener returns, all completion listeners run on the completing thread.

[dongjoon-hyun]

Putting this listener first also puts it ahead of the close listeners that the writer thread registers lazily (a DSv2/JDBC/S3 reader opened on its first pull). Before this PR those ran first (LIFO), and closing the stream was what unblocked a writer stuck in a socket read that ignores Thread.interrupt.

Now stopAndAwaitExit waits without bound first, logging every 10s, and the close listener never gets to run. Since the task reaper is off by default, the task never completes and the executor slot is lost.

The UAF fix needs this ordering, so I'm not suggesting a revert. Could we at least document this hang mode (and the reaper as the escape hatch) in the scaladoc and the warning message, or add a configurable bound after which the listener fails loudly?

[viirya]

Agreed, and documented. The scaladoc of startPipelinedWriter now says that a writer blocked in upstream I/O that ignores interrupts holds task completion until that I/O returns, including I/O that only the upstream reader's close listener would unblock, and that spark.task.reaper.enabled bounds how long a killed task can stay there. The 10-second warning says the same. I didn't add a configurable bound: when it expires, failing the listener would not stop the listeners after it from freeing memory under the writer, so the bound would only trade the hang for the use-after-free.

For writers that are blocked because an interrupt was consumed (HyukjinKwon's question in the review summary), the wait now interrupts the writer again on every warning tick, as long as the writer is still running.

[dongjoon-hyun]

"regardless of when or from which thread that listener was added" claims more than the implementation does. A pre-completion listener added after completion has started runs only before the remaining completion listeners. The new test a pre-completion listener added during completion runs next asserts Seq("B", "P", "A"), so B runs before P.

Maybe: "before every pending listener added with addTaskCompletionListener. Completion listeners that have already run are not affected." That makes it clear a pre-completion listener must be registered before completion starts to protect everything.

[viirya]

Thanks, reworded along those lines: "before every pending listener added with addTaskCompletionListener, from any thread and even if that listener was added later. Completion listeners that have already run are not affected, so to run before all of them, the listener must be added before the task completes."

[dongjoon-hyun]

CountDownLatch.await checks Thread.interrupted() before it checks the count, so when the writer was claimed, a kill interrupt that is already set is consumed here even if the writer has already exited. If the writer was never claimed, the early return at L247 leaves the flag set.

As a result, the listeners that run next see the interrupt flag set or not depending on a race. The comment above says the flag would make their blocking channel and lock operations fail and skip cleanup, which is the case this is trying to avoid.

If the intent is to never pass the interrupt on, Thread.interrupted() could be called on the early-return path as well. Otherwise, could the comment describe the actual behavior?

[viirya]

Good catch. The wait now leaves the caller's interrupt status as it found it. It reads and clears the flag before waiting and restores it afterwards. Interrupts that arrive during the wait are still not passed on. Both paths (writer never claimed, writer claimed) now leave a pending interrupt in place, and the comment says so. The new test "an interrupt pending before stopping the writer is kept" covers both paths.

[dongjoon-hyun]

This depends on the busy-spin landing in the few-instruction window between completed = true and the claim, without forcing that interleaving. If the markTaskCompleted claim is reverted, the adder almost always arrives after the task thread has registered itself, so I expect all 500 iterations to still pass.

Could the test force the interleaving deterministically, for example with a TaskContextImpl test subclass or hook that blocks between setting completed and invoking the listeners? A case for the failure-path window (inline comment on L322) would also help.

[viirya]

When I reverted the markTaskCompleted claim while keeping the test, it failed on the first run, so in practice the spin lands in the window well within 500 iterations. With the claim in place there is no window left between setting completed and claiming the invocation, so there is nothing to block on there. For the new handoffs, the test with another thread holding the invocation is deterministic (the holder blocks in a failure listener). The failure-path window is a few instructions wide in the old code as well, so that test also relies on contention; it too fails on the first run with the old code.

[dongjoon-hyun]

With NonFatal(t) widened to Throwable, a fatal error such as an OOM now goes through writer.setException and shutdownOutput (both may allocate near OOM) before it is rethrown. PipelinedWriterTask then logs it at ERROR, while ReaderIterator rethrows the same error on the task thread, so one fatal error is reported from two threads.

Is the setException path needed for fatal errors, or would rethrowing them directly (as before) and letting PipelinedWriterTask log them be enough? If the goal is to unblock the reader, a short note on the expected double report would help.

[viirya]

The setException path is what keeps the task from waiting on the Python worker with no cause recorded, so I kept it, and added a comment that the error is then reported twice (the writer thread's log and the reader's rethrow on the task thread). setException is a field store, and shutdownOutput() goes through Utils.tryLog, which catches and logs any failure, so the original error is still rethrown and logged afterwards.

[dongjoon-hyun]

nit: this copies addTaskCompletionListener with only the stack changed (the comment even says "Same as addTaskCompletionListener"). A private helper such as addCompletionListenerTo(stack, listener) used by both would keep the add/invoke protocol in one place.

[viirya]

Done: both now go through a private addCompletionListenerTo(stack, listener).

[viirya]

@HyukjinKwon On the decision challenge in your review summary (a lost writer interrupt): I can't rule out an upstream iterator that catches the InterruptedException and carries on, so the wait now interrupts the writer again on every 10-second warning while the writer is still running. The interrupt is guarded so that a pool thread that has already left the writer is never interrupted. A writer blocked in channel.write then fails with ClosedByInterruptException, which also closes the worker channel. The new test "the writer is interrupted again if it consumed the interrupt" covers it, and it fails without the repeated interrupt.

[dongjoon-hyun]

Thank you for the update, @viirya. 1053b61d411 addresses everything from the previous round: the failure path keeps the invocation, markTaskCompleted waits for the holder and takes over from it directly, a pending interrupt is kept, and the hang mode is documented. I left 8 inline comments on the new head, mostly on the new wait and the safety net. Summary:

1. The handoff wait in markTaskCompleted is unbounded and silent

   If the thread holding the invocation blocks in a failure listener, the task thread waits forever, ignores kill interrupts, and logs nothing. A periodic warning that names the holder, like the one in stopAndAwaitExit, would make this diagnosable.

2. The writer-thread safety net does not stop the writer

   cancel(true) only interrupts the current thread. The remaining listeners then run with the interrupt flag set and free memory that the writer uses again as soon as they return, so this path trades the hang for the use-after-free.

3. The interrupt test depends on a 200 ms sleep

   Now that a pending interrupt is kept, an interrupt that arrives before the wait starts leaves interruptedInCleanup at Some(true), so the test can fail under CI load.

4. Failure listeners queued during completion are dropped

   The completing thread only drains the completion stacks, and it now holds the invocation for the whole writer wait.

5. spark.task.reaper.enabled alone does not bound the hang

   With the default spark.task.reaper.killTimeout=-1, the reaper never ends the task.

6. ErrorNotifier's doc describes the old behavior

   markTaskCompleted no longer skips the completion listeners when another thread holds the invocation.

7. nit: claimed and nested allow an invalid combination

   A single three-state parameter would rule out nested without claimed.

8. nit: PipelinedWriterTask spreads its state over three primitives

   A second stopAndAwaitExit on a writer that never started would wait forever.

[dongjoon-hyun]

This wait has no time bound, ignores interrupts, and logs nothing. Before this PR, markTaskCompleted returned right away in this case, so it could not hang here.

The holder here is typically the kill thread. Task.kill runs markInterrupted while holding the global TaskContext lock, and with the task reaper off (the default) Executor.killTask runs it on the thread that received the kill request. If an interrupt listener throws, that thread runs the failure listeners. If one of them blocks, the task thread stays in wait() forever: kill interrupts are caught at L237 and the loop keeps waiting. Nothing in the log points at the holder, and with the default spark.task.reaper.killTimeout=-1 nothing ever ends the task.

You mentioned this risk in your reply, and I agree it is unlikely. But stopAndAwaitExit already handles the same kind of wait with a 10s warning. Could this loop do the same, e.g. wait(timeout) plus a periodic warning that names the holder thread (and maybe logs its stack)? Then this case would at least be diagnosable.

[viirya]

Done in 7ab7ae2. 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.

[dongjoon-hyun]

I don't think "The interrupt above stops the writer afterwards" holds. On this thread, cancel(true) interrupts the FutureTask runner, which is this thread itself, so all it does is set our own interrupt flag. After this listener returns:

• The remaining completion listeners run on the writer thread with that flag set. The comment at L296-299 says this makes their blocking channel and lock operations fail and skip cleanup.
• Those listeners close the HybridRowQueue pages and the upstream reader's off-heap vectors. When they return, control goes back to the upstream code that triggered them (e.g. ParquetReaderCallback.initIfNotAlready). That code keeps using the closed reader, and the writer adds the next row to the freed queue page before the loop at L1445 checks isInterrupted.

So when this path is taken, it trades the hang for the SPARK-33277 use-after-free, and only logs a WARN. I agree it should not be reachable with TaskContextImpl now. Since reaching it means that invariant is broken, could it log at ERROR, clear the self-interrupt (Thread.interrupted()) before returning, and have the comment say that the writer is not stopped before the remaining listeners run?

[viirya]

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.

[dongjoon-hyun]

Since 1053b61d411 keeps an interrupt that is pending when the wait starts, this test now depends on timing. The completion thread may not reach Thread.interrupted() in stopAndAwaitExit within these 200 ms. It can still be between reached.countDown() and the stop listener, e.g. after a GC pause on a busy CI machine. The interrupt then counts as pending, the finally restores it, and interruptedInCleanup ends up Some(true), so the assertion at L232 fails.

Could this wait until the thread blocks in the wait instead, as the other new tests do?

eventually(timeout(1.minute)) {
  assert(completion.getState == Thread.State.TIMED_WAITING)
}
completion.interrupt()

[viirya]

Thanks, changed to wait with eventually until the completion thread is in TIMED_WAITING before interrupting it.

[dongjoon-hyun]

Failure listeners that another thread queues while this thread holds the invocation never run. markTaskFailed and addTaskFailureListener from another thread find listenerInvocationThread set and return (L342-345). This loop only pops the completion stacks, so it deregisters here and leaves onFailureCallbacks untouched.

This predates the PR, but the window used to be short, because the old writerFuture.get() returned right away. Now the task thread holds the invocation for the whole wait in stopAndAwaitExit. For example, if a kill arrives during that wait and an interrupt listener throws, markInterrupted records the failure on the kill thread, but no failure listener ever runs.

Would it make sense to run pending failure listeners before deregistering when failureCauseOpt is set? Or at least to note the gap in the scaladoc?

[viirya]

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.

[viirya]

An update, since this thread still ends at "Fixed.": the release loop that ran these listeners moved to SPARK-60222 with the split. Since 43dff78, the pipelined writer wait no longer holds the listener invocation (it runs as a pre-completion listener, outside it), so failure listeners triggered from another thread during that wait run on that thread. Failure listeners queued while the completion listeners run are still dropped, as on master, and SPARK-60222 covers that.

[dongjoon-hyun]

spark.task.reaper.enabled alone does not bound this. With the default spark.task.reaper.killTimeout=-1, TaskReaper only monitors the killed task (and optionally takes thread dumps). timeoutExceeded() is false whenever the timeout is not positive, so the reaper never ends the task.

Could this say something like "spark.task.reaper.enabled together with a positive spark.task.reaper.killTimeout bounds how long it may stay in this state, by killing the executor JVM"? The same applies to the new wait in markTaskCompleted.

[viirya]

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.

[dongjoon-hyun]

With this change, the class doc of ErrorNotifier (L31-33) is out of date. It says that calling markTaskFailed from a background thread "can run concurrently with the task thread's markTaskCompleted and silently skip completion listeners". Now markTaskCompleted waits for that thread and then runs them.

The hazards are different now: the task thread blocks until the background thread's failure listeners finish, and failure listeners that another thread queues during completion are dropped (see the comment on L362). The comment in StreamingShuffleReader.checkTaskFailure ("run without contention") is in the same situation. Could both be updated here, since this PR changes the behavior they describe?

[viirya]

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.

[dongjoon-hyun]

nit: nested implies claimed, but the two booleans still allow nested = true, claimed = false. That combination would register the thread at L350 and then never deregister it, because the nested branch at L362 skips that. listenerInvocationThread would stay set, and a later markTaskCompleted from another thread would wait forever.

Every caller passes a valid combination today ((nested, nested) and (true, nested)), so this is not a bug. A single three-state parameter (e.g. claim / already claimed / nested) would make the invalid case unrepresentable, and invokeTaskCompletionListeners would not need its own defaults.

[viirya]

Done: replaced them with a single InvocationMode (Claim, Claimed, Nested) in a private TaskContextImpl companion object.

[dongjoon-hyun]

nit: claimed is set by whichever side gets there first, so afterwards it cannot tell "the writer is running" from "the stopper cancelled it". A second stopAndAwaitExit on a writer that the first call stopped before it started would fail the CAS, skip the early return, and wait on exited forever, because run() never counts it down when it loses the CAS.

Only one listener calls this today, so it cannot happen yet. One state field guarded by this (not started / Running(thread) / done or cancelled), with the latch counted down on every terminal transition, would remove this case and also the split between the CAS and the synchronized writerThread.

[viirya]

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.

[HyukjinKwon]

Review summary

Everything from the previous round is addressed in 7ab7ae2.

• Self-wait (previous P2). The residual is now closed at its source. markTaskCompleted waits for the holder and takes the invocation over directly, and a failing listener keeps the invocation while the failure listeners run. The same-thread guard in stopAndAwaitExit stays as a logged safety net.
• Fatal-error branch (previous P3). It now has a regression test that runs a real PipelinedWriterRunnable against a socket pair.
• TaskContextSuite comment (previous P3). It now matches the asserted order.
• Lost writer interrupt (previous decision challenge). The wait now re-interrupts a still-running writer on every warning tick, guarded by the writer state.
• dongjoon-hyun's three rounds. All items are addressed in the current head.

One new non-blocking item: the corrected reaper sentence in the markTaskCompleted comment still does not hold for the comment's own example, a kill thread running blocking failure listeners, because with the reaper enabled that thread is the reaper itself.

Findings

1 total: 0 P0, 0 P1, 0 P2, 1 P3.

Nit (P3)

• The reaper bound in the handover-wait comment does not hold when the reaper itself runs the failure listeners — core/src/main/scala/org/apache/spark/TaskContextImpl.scala:225 — see inline.

Re-review status

Prior AI findings: 3 addressed, 0 still present; additional unresolved findings in this review: 1.

New attribution: 1 newly introduced, 0 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

No prior AI findings remain.

Existing discussions

• existing discussion — Both comments now name spark.task.reaper.enabled together with a positive spark.task.reaper.killTimeout, as requested. For the writer wait in startPipelinedWriter, that statement holds. For the markTaskCompleted handover wait it still does not hold in the comment's own example: with the reaper enabled, the reaper thread runs the kill, and a failing interrupt listener makes it run the failure listeners before it ever reaches its killTimeout loop. That residual is filed as a related new defect.

PR description suggestions

• Item 3 still says that an AtomicBoolean claim makes sure a writer stopped before it starts is skipped. The current head replaces that with a single PipelinedWriterState (NotStarted, Running, Finished) guarded by the task, and the exit latch is counted down on every transition to Finished, including a stop before the writer starts.

Generated by Omnigent on Databricks.

[HyukjinKwon]

Nit (P3): This bound does not hold for the example holder named above, a kill thread running failure listeners. With spark.task.reaper.enabled, Executor.killTask runs the kill on a TaskReaper, and TaskReaper.run calls taskRunner.kill() synchronously before it enters its killTimeout loop. kill() goes through Task.kill to markInterrupted, and if an interrupt listener throws, markTaskFailed runs the failure listeners on that same reaper thread. If one of them blocks, kill() never returns, timeoutExceeded() is never checked, and this thread stays in awaitListenerInvocationHandover with no JVM kill. The bound does hold when the holder is some other thread, and for the writer-exit wait in startPipelinedWriter, where the task thread holds the invocation. Could the comment say that the reaper bounds this wait only when the blocked listeners do not run inside the reaper's own kill call, and that otherwise only the periodic warning shows it?

Verification:

• Inspection: The comment's recovery statement agrees with the TaskReaper sequence, where kill() runs synchronously before the killTimeout loop, and with markInterrupted running failure listeners on the kill thread when an interrupt listener throws.

[viirya]

Good catch, confirmed: TaskReaper.run calls taskRunner.kill() before it starts checking killTimeout, so when the blocked failure listeners run within that call, the reaper never gets to the timeout. Reworded the comment in d05f954: the reaper bounds the wait only if the blocked listeners do not run within its own kill call, and otherwise only the periodic warning shows the hang. Also updated item 3 of the PR description to describe the single PipelinedWriterState instead of the old AtomicBoolean claim.

[HyukjinKwon]

Thanks, verified in d05f954. The markTaskCompleted comment now says that the reaper bounds the handover wait only when the blocked listeners do not run within its own kill call, and that otherwise only the periodic warning shows the hang. Item 3 of the PR description now describes the single PipelinedWriterState as well.

[HyukjinKwon]

Thanks, I checked 7ab7ae2. awaitExit now calls interruptWriter() on every warning tick, and interruptWriter interrupts only while the state is PipelinedWriterRunning, under the same lock that finish() takes, so a pool thread that has already left the writer is never interrupted. A writer that swallowed the first interrupt and then blocks in channel.write gets ClosedByInterruptException on the next tick, which closes the worker channel and breaks the cycle. The new test "the writer is interrupted again if it consumed the interrupt" only stops after a second interrupt, so it would hang without the repeated one. That resolves the question from my earlier review.

[dongjoon-hyun]

Thank you for the update, @viirya. d05f9547fea addresses everything from the previous round: the handover wait now warns with the holder's stack, the listeners left to the invoking thread run instead of being dropped, the reaper bound is documented correctly, InvocationMode and PipelinedWriterState replace the boolean pairs, the writer-thread branch clears its own interrupt, and the interrupt test waits for TIMED_WAITING.

However, the core module CI job has been timing out since 1053b61, because the new handover wait deadlocks an existing StreamingShuffleSuite test (1). I left 10 inline comments on this head. Summary:

1. The handover wait deadlocks an existing StreamingShuffleSuite test, so the core CI job times out (TaskContextImpl.scala L243)

   "background markTaskFailed races with task thread markTaskCompleted" releases the background failure listener only after markTaskCompleted returns, and markTaskCompleted now waits for it. The core job was cancelled at the 2h30m limit on 1053b61 and 7ab7ae2, so BasePythonRunnerSuite has not run in CI since then.

2. Failure listeners left to the completing thread run after every completion listener (TaskContextImpl.scala L448)

   This reverses the failure-before-completion order, and their exceptions now fail the task from markTaskCompleted.

3. The handover wait passes on an interrupt that awaitExit would drop (TaskContextImpl.scala L277)

   The completion listeners then run with the interrupt flag set.

4. A stale holder now blocks markTaskCompleted forever (TaskContextImpl.scala L441)

   The release is not in a finally, and the new wait never ends if it is skipped.

5. Leftover completion listeners can run on a kill thread under the global TaskContext lock (TaskContextImpl.scala L453)

   "The completion listeners always run on this thread" holds only for the listeners pending at completion.

6. nit: future.cancel(true) is redundant with the new state (PythonRunner.scala L289)

   Interrupting the PipelinedWriterRunning thread under the lock does the same, without the self-interrupt.

7. nit: the handover warning walks the holder's stack and logs it every 10 s under the lock (TaskContextImpl.scala L268)

8. nit: the new TaskContextSuite tests join helper threads without a timeout (TaskContextSuite.scala L837)

9. nit: the handover warning has no test, and its interval is hard-coded (TaskContextImpl.scala L578)

10. Discussion: the TaskContextImpl rework could be its own [CORE] PR (TaskContextImpl.scala L215)

    It changes listener behavior for every task, not only for the pipelined mode.

[dongjoon-hyun]

[High] This wait deadlocks an existing StreamingShuffleSuite test, and the core CI job times out on it.

StreamingShuffleSuite's "background markTaskFailed races with task thread markTaskCompleted" (L1380-1433) has a background thread call markTaskFailed and block in a failure listener on failureListenerProceed. The test thread then calls markTaskCompleted(Some(error)) (L1419). Before this PR, that call returned right away. Now it waits here for the background thread to hand over the invocation, but the test releases the failure listener only at L1422, on the same thread, after markTaskCompleted returns. Neither thread can make progress, and since this wait ignores interrupts, the test framework cannot break it either.

The core module job shows it:

• 7ab7ae2 (job): logs "Test still running after 2 hours, 14 minutes" for this test and is cancelled at the 2h30m limit.
• 1053b61 (job): hangs in the same test and is cancelled the same way.
• 038c4b3, the last commit before the handover wait (job): passes in about 66 minutes.

So the core-module suites scheduled after it, BasePythonRunnerSuite included, have not run in CI since 1053b61. TaskContextSuite runs earlier and passes.

The test and its comments (L1375-1379, L1405-1406, L1416-1418, L1425-1430) document the skipped-listener race that this PR removes. Could the test be updated to the new behavior? For example, call markTaskCompleted on another thread, wait until that thread blocks in this wait, release the failure listener, and then assert that the completion listener ran, on that thread.

[viirya]

Thanks for tracking this down in the CI logs, and sorry about it. I had run only the suites around the change. With the split (see L215), the handover wait is no longer in this PR, so TaskContextImpl is back to what it was in 038c4b3, whose core job passed. On 00da9fb I ran StreamingShuffleSuite and the other core suites that use markTaskFailed, markTaskCompleted or markInterrupted locally, and they pass. The follow-up will update that test to the new behavior along the lines you describe.

[dongjoon-hyun]

[Medium] Failure listeners left to this thread run after every completion listener.

Suppose another thread fails the task while this thread runs the completion listeners, e.g. a kill whose interrupt listener throws. Its markTaskFailed finds the invocation taken and returns. This thread then first runs the remaining completion listeners, which close streams, release state stores and free the HybridRowQueue pages, and only after that the failure listeners, here. runTaskWithListeners always runs the failure listeners before the completion listeners, and listeners may rely on that. For example, RocksDBStateStoreProvider registers a failure listener if (!hasCommitted) abort() next to a completion listener that aborts the store and releases its stamp, so the failure listener may now call abort() after the stamp has been released.

The same happens when the release below hands the invocation to a thread waiting in markTaskCompleted while failure listeners are still pending: that thread runs all completion listeners first.

Exceptions from these failure listeners are also added to listenerExceptions (L466) and thrown from markTaskCompleted on the task thread. Before, they stayed on the thread that called markTaskFailed.

Running them late is still better than dropping them, as before. Could pending failure listeners run first instead, i.e. before the handover below and before the next completion listener? Otherwise, could the order be documented here and in addTaskFailureListener?

[viirya]

This code moved to SPARK-60222. I agree that the failure listeners left to the completing thread should run before its remaining completion listeners, including on the handover path, and I'll do that there rather than only document the order.

[dongjoon-hyun]

[Low] An interrupt that arrives during this wait is passed on to the completion listeners, while awaitExit drops one for that reason.

This restores the flag before the completion listeners run. stopAndAwaitExit then sees it as pending on entry and restores it again, so the ordinary completion listeners (the HybridRowQueue close, worker.stop(), the upstream reader's close) run with the flag set. PipelinedWriterTask.awaitExit (PythonRunner L333-336) drops an interrupt that arrives during its own wait, because the flag "would make the blocking channel and lock operations of the completion listeners that run next fail and skip cleanup". So an interrupt that lands while the task thread waits here is passed on, and the same interrupt landing a moment later, during the writer wait, is dropped.

This wait can be interrupted, e.g. by a second kill with interruptThread = true, or by the writer thread resolving a pending interrupt in createResourceUninterruptibly, which interrupts the task thread. Could both waits follow one policy, or could the comments say why they differ?

[viirya]

Moved to SPARK-60222 with the handover wait. I'll make the two waits follow one policy there, or explain why they differ.

[dongjoon-hyun]

[Low] The release is not in a finally, and a stale holder now blocks markTaskCompleted forever.

If a Throwable escapes the loop above, this thread never gets here, and listenerInvocationThread keeps pointing at it. The catch block can throw itself: if a listener fails with OutOfMemoryError, listenerExceptions += e or logError (L437-438) can fail again. Before this PR, a stale holder only made markTaskCompleted return without running the completion listeners. Now markTaskCompleted waits in awaitListenerInvocationHandover for a handover that never comes, and interrupts do not end that wait.

The reverse case also exists: if awaitListenerInvocationHandover throws, e.g. from the logWarning at L269, nextListenerInvocationThread stays set. The holder later hands the invocation to a thread that has already left, and every later listener is skipped.

Could the main loop be wrapped in a try/finally, so that the release, or at least the handover, always runs?

[viirya]

Moved to SPARK-60222. I'll put the release and the handover in a finally there, so that a stale holder or a stale waiter cannot be left behind.

[dongjoon-hyun]

[Low] Leftover completion listeners can run on a kill thread under the global TaskContext lock.

This runs the pending completion listeners on whichever thread holds the invocation, not only on the completing thread. For example, after the task thread has finished markTaskCompleted, a late kill whose interrupt listener throws runs the failure listeners on the kill thread, inside Task.kill's JVM-wide TaskContext.synchronized block (Task.scala L235). A completion listener that another thread adds at that point is left to the kill thread, which runs it here under that lock. A slow close, e.g. of an S3 or JDBC resource, then also blocks Task.kill, interruptible() and createResourceUninterruptibly for every task on the executor. The new test "completion listeners queued meanwhile run before failure invoker releases" asserts this behavior on purpose.

This is rare, but the comment at L215, "The completion listeners always run on this thread", now holds only for the listeners that are pending when markTaskCompleted is called. Could it say so, and mention that the leftover ones run on the thread that holds the invocation?

[viirya]

Moved to SPARK-60222. I'll document there which listeners run on which thread, and look at keeping leftover completion listeners off a thread that holds the global TaskContext lock.

[dongjoon-hyun]

[Nit] future.cancel(true) is no longer needed with the new state.

The state already covers both cases. NotStarted -> Finished keeps run() from starting the writer. For PipelinedWriterRunning(thread), interrupting thread under this, as interruptWriter() does, reaches the same thread that cancel(true) interrupts. cancel(true) does it outside the lock, so it can still interrupt the pool thread after the writer has exited. On the writer thread itself, it interrupts the calling thread, which is one reason for the unconditional Thread.interrupted() in the finally.

Interrupting in the state match instead, skipped when thread eq Thread.currentThread(), would remove the future parameter of stopAndAwaitExit, the writerFuture in startPipelinedWriter, the self-interrupt, and most of the FutureTask plumbing in the tests.

[viirya]

Done in 00da9fb. stopAndAwaitExit no longer takes the future: it interrupts the PipelinedWriterRunning thread under the lock, and skips that when it runs on the writer thread itself. startPipelinedWriter no longer keeps the future, and the tests no longer need FutureTask.

[dongjoon-hyun]

[Nit] The warning walks the holder's stack and logs it every 10 s while holding the TaskContextImpl lock.

The holder needs this lock to hand over (L455-458), and isCompleted(), isFailed() and every add*Listener call on other threads need it too, so each warning stalls them behind a stack walk and a synchronous log write. A task that stays stuck for an hour also logs 360 full stack traces as java.lang.Exception: Stack trace of thread ..., which log scanners tend to count as errors. Could the stack be attached only to the first warning, or could the warnings come at growing intervals?

[viirya]

Moved to SPARK-60222 with the handover warning. I'll take the stack outside the lock there, and attach it only to the first warning.

[dongjoon-hyun]

[Nit] These tests join helper threads without a timeout, some from inside a listener that holds the invocation.

Here and in "failure listeners queued during completion run before it releases" (failer.join(), L970), the listener blocks in join() while its thread holds the listener invocation. "completion waits for another thread invoking listeners and runs them" has the same adder.join() (L945), and none of these helper threads is a daemon thread. If a later change makes the helper wait for the invocation, as markTaskCompleted now does, the test hangs until the core job hits the 2h30m limit, as the StreamingShuffleSuite test does now, instead of failing. Could they use join(timeout) with assert(!thread.isAlive) and setDaemon(true), like the other new tests?

[viirya]

These tests moved to SPARK-60222 with the code they cover. I'll use daemon threads and join(timeout) with assert(!thread.isAlive) there.

[dongjoon-hyun]

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.

[viirya]

Thanks. That test changed with the L205 change: a pre-completion listener added after completion is now called immediately on the adding thread, so the test ("a pre-completion listener added after completion runs immediately") checks that, and it now uses a daemon thread with join(timeout) and assert(!thread.isAlive). I'll leave the older "listener registers another listener using a second thread" test to SPARK-60222.

[dongjoon-hyun]

[Nit] The handover warning has no test, and this interval cannot be shortened for one.

The warning is the only sign that markTaskCompleted is stuck behind another thread, but it fires only after 10 s, and none of the new tests wait that long. A regression there, e.g. the wrong holder named or no stack attached, would go unnoticed. PipelinedWriterTask takes warnIntervalMs as a parameter, and BasePythonRunnerSuite sets it to 100 ms to exercise that path. Could this interval be configurable in the same way, with a test that checks that the warning names the holder and carries its stack?

[viirya]

Moved to SPARK-60222. I'll make the interval configurable there and add a test that checks the holder's name and stack.

[dongjoon-hyun]

[Discussion] The TaskContextImpl part could be its own [CORE] change.

The rework here (the handover in markTaskCompleted, the release loop, and the three invocation modes) applies to every task, not only to the opt-in pipelined mode:

• markTaskCompleted can now block indefinitely on another thread.
• Failure-listener exceptions from other threads can now fail the task on the task thread.
• Failure and completion listeners that used to be dropped now run, possibly late and on other threads.
• An existing core test deadlocks (see the comment on L243).

The PR description answers "No" to the user-facing change question because the pipelined mode is unreleased, but this part is not behind that flag. A separate JIRA and PR for it would get its own review and tests, and a regression in it would not force reverting the use-after-free fix together with it. What do you think?

[viirya]

Agreed. I moved it out in 00da9fb: the handover wait in markTaskCompleted, running failure listeners inside the current invocation, running the listeners left to the invoking thread, and InvocationMode now go to SPARK-60222 as a separate [CORE] PR. This PR keeps only what the use-after-free fix needs from TaskContextImpl: addTaskPreCompletionListener and the claim in markTaskCompleted that is set together with completed (as in 038c4b3), plus the shared add helper.

For the cases that the follow-up closes (another thread holding the invocation at completion, or the brief release while failure listeners run), the stop listener keeps the same-thread guard here. If it runs on the writer thread, it does not wait and logs an error, and the use-after-free from before this PR remains possible on that path. The comment and the PR description now say so and point to SPARK-60222. I also updated the PR description, which no longer lists the moved changes.

[HyukjinKwon]

Review summary

The previous round's P3 is addressed in d05f954. The markTaskCompleted comment now says the task reaper bounds the handover wait only when the blocked failure listeners do not run inside the reaper's own kill call, and that otherwise only the periodic warning shows the hang. The PR description's item 3 now describes the single PipelinedWriterState. The earlier P2 and P3 items, the lost-interrupt question, and dongjoon-hyun's three rounds all remain addressed in this head.

Two new non-blocking P3 items, both late catches in code unchanged since 7ab7ae2, concern failure paths of the new TaskContextImpl handover that no test exercises. Nothing interrupts the thread waiting in markTaskCompleted for a handover, and no deferred listener throws inside the release loop. The current code handles both correctly, but a later change to either path could reintroduce skipped cleanup or a permanently held listener invocation without any test failing.

Findings

2 total: 0 P0, 0 P1, 0 P2, 2 P3.

Nit (P3)

• The interrupt handling of the markTaskCompleted handover wait has no test — core/src/main/scala/org/apache/spark/TaskContextImpl.scala:262 — see inline.
• The release loop's handling of a throwing deferred listener has no test — core/src/main/scala/org/apache/spark/TaskContextImpl.scala:466 — see inline.

Re-review status

Prior AI findings: 1 addressed, 0 still present; additional unresolved findings in this review: 2.

New attribution: 0 newly introduced, 2 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

No prior AI findings remain.

Generated by Omnigent on Databricks.

[HyukjinKwon]

Nit (P3): This catch keeps the handover wait going on an interrupt, and the restore below hands the interrupt back once the invocation arrives. Both matter on the kill path: Task.kill interrupts the task thread, which can land while it waits here for a kill thread that is running failure listeners. If a later change let the InterruptedException escape, markTaskCompleted would throw with completed already set and nextListenerInvocationThread still naming this thread. The holder would then hand the invocation to a thread that has left, listenerInvocationThread would stay set, and none of the completion listeners (including the pipelined writer stop and the memory cleanup) would run. If the restore were dropped, the kill's interrupt would be silently lost. "completion waits for another thread invoking listeners and runs them" never interrupts the completion thread, so no current assertion fails in either case.

Verification:

• Regression: The test fails when the handover wait propagates InterruptedException or does not restore the interrupt flag, and passes with the current implementation.
• Behavior: After the interrupted wait, every pending completion listener runs on the completing thread.

[viirya]

Thanks. The handover wait moved to SPARK-60222, and I'll add that test there: interrupt the completing thread while it waits, then check that the wait continues, that the interrupt is kept, and that all completion listeners run on the completing thread.

[HyukjinKwon]

Nit (P3): This catch keeps the release loop reaching released = true and notifyAll() when a deferred listener throws, and it puts that failure into the reported exception. Both new release-loop tests ("failure listeners queued during completion run before it releases" and "completion listeners queued meanwhile run before failure invoker releases") use listeners that do not throw, so nothing exercises this branch. If a later change dropped the catch, the nested invocation's TaskCompletionListenerException would leave the loop early with this thread still registered as listenerInvocationThread. Later invocations, including invokePostStatusUpdateListeners, would return without running anything, and a markTaskCompleted waiting for the handover would never be woken. If it swallowed the exception instead, the failure would never reach the caller. Neither change would fail a current test.

Verification:

• Regression: The test fails when the release loop lets a deferred listener's exception escape before releasing the invocation, or swallows it, and passes with the current implementation.

[viirya]

Thanks. The release loop moved to SPARK-60222, and I'll add a test there with a deferred listener that throws. It will check that the invocation is still released (or handed over) and that the failure reaches the caller.

[dongjoon-hyun]

Thank you for the update, @viirya. 00da9fb69d4 moves the TaskContextImpl rework to SPARK-60222 as discussed, so the StreamingShuffleSuite deadlock is gone, and stopAndAwaitExit now interrupts the running writer directly instead of through Future.cancel(true). The description now also lists the paths that SPARK-60222 still has to close.

I left 14 inline comments on this head, plus a reply in the round-4 thread on TaskContextSuite.scala:834. Item 1 matters most: the interrupt that stops the writer can itself free the Arrow writer's output buffer while the writer is still copying into it. It predates this PR, but it breaks the invariant that this PR establishes. Summary:

1. [High] The stop interrupt can free the Arrow writer's output buffer while the writer is copying into it (PythonRunner.scala:295)

   The JDK runs the interruptor of Channels.newChannel(dataOut) on the stopping thread, which closes the DirectByteBufferOutputStream under the writer. A close-shielded stream for the writer would avoid it.

2. [Medium] On the writer-thread path nothing stops the writer any more (PythonRunner.scala:309)

   Master's cancel(true) self-interrupt stopped it after the current call. Now it keeps adding rows to the freed queue page until its next flush. A !context.isCompleted() check in the writer loop would restore the old bound.

3. [Low] A fatal error before the try still leaves the reader waiting for Python (PythonRunner.scala:1497)

4. [Low] Failure listeners that another thread triggers during the writer wait are dropped again, and nothing says so (TaskContextImpl.scala:319)

5. [Low] The writer interrupts ignore context.interruptible(), which Task.kill respects (PythonRunner.scala:271)

6. [Low] No test fails if the immediate interrupt, or the Thread.interrupted() in the finally, is removed (PythonRunner.scala:294)

7. [Low] The only test that catches reverting addTaskPreCompletionListener in startPipelinedWriter depends on a race (BasePythonRunnerSuite.scala:132)

8. [Discussion] Running the pre-completion listeners before completed is set would close the deferred writer paths in this PR (TaskContextImpl.scala:205)

9. [Low] The BarrierTaskContext delegation is untested (BarrierTaskContext.scala:211)

10. [Low] The fatal-error test's read blocks forever on a regression instead of failing (BasePythonRunnerSuite.scala:404)

11. [Nit] "which would then wait for itself" no longer matches the same-thread guard (TaskContextImpl.scala:204)

12. [Nit] The scaladoc still justifies the latch with Future.cancel(true) (PythonRunner.scala:199)

13. [Nit] The wait warning doesn't identify the writer thread (PythonRunner.scala:348)

14. [Nit] addNow is not counted down in the finally (BasePythonRunnerSuite.scala:197)

15. [Nit] The test that the round-4 join() comment is anchored on stayed in this PR (TaskContextSuite.scala:834, reply in that thread)

[dongjoon-hyun]

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.

[viirya]

Great find, confirmed against the JDK 17 sources: Thread.interrupt() runs blocker.interrupt(this) on the interrupting thread, the WritableByteChannelImpl interruptor calls implCloseChannel(), i.e. out.close(), and DirectByteBufferOutputStream.close() disposes the buffer. Writing the test also showed that begin() only closes the channel and then lets out.write copy the bytes, and only end() throws. Fixed in 43dff78 as you suggested: the pipelined writer writes through a new CloseShieldOutputStream around bufferStream, which forwards writes and only flushes on close, and bufferStream is freed only in the writer's own finally. Tests: an interrupted channel write through the shield leaves the buffer usable, both directly and through a real PipelinedWriterRunnable with an Arrow-style write. Each fails without the shield.

The sync path has the same exposure through Task.kill, e.g. when a speculative attempt is killed with interruptThread = true, and it has been there since SPARK-44705 (4.0.0). I'm fixing that separately in SPARK-60305 with the same CloseShieldOutputStream, so that it can be backported to 4.0 through 4.3.

[dongjoon-hyun]

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()?

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 says the comment points to SPARK-60222, but no code comment mentions it yet.

[viirya]

With the change for your L205 comment (see there), this path is no longer reachable through TaskContextImpl. The pre-completion listeners run before the task is marked completed, so the writer's lazy add*Listener calls only queue, and the writer has exited before any completion listener runs. context.isCompleted() is therefore never true while the writer runs, and a check in the writer loop would never fire, so I didn't add it. Instead, the comment on that branch now says that it should not happen, and that if it does, nothing stops the writer, which goes back to pulling and writing input once the remaining listeners have freed its memory. Its logError now carries the current stack. The wrong description of the failure-listener window and the missing SPARK-60222 reference in that comment are moot now, since the comment no longer refers to those paths.

[dongjoon-hyun]

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?

[viirya]

Fixed in 43dff78: the thread-locals, the DirectByteBufferOutputStream and the DataOutputStream are now set up inside the try, and the finally closes bufferStream only if it was created.

[dongjoon-hyun]

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 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?

[viirya]

With the L205 change this gap is gone for the writer wait. The pre-completion listeners, including the writer wait, run outside the listener invocation, so a markTaskFailed from a kill thread during the wait claims the invocation and runs the failure listeners itself. The new test "failure listeners run while pre-completion listeners run" covers this. The remaining case, failure listeners queued while the completion listeners run, is as on master and stays with SPARK-60222. I updated the round-3 thread accordingly.

[dongjoon-hyun]

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?

[viirya]

Done in 43dff78, for both the first interrupt and the repeated one. The writer is interrupted only under TaskContext.synchronized while context.interruptible() is true; otherwise it is skipped, and the next warning tick tries again. The new test "the writer is not interrupted while it creates a resource uninterruptibly" runs the writer inside context.createResourceUninterruptibly with a 100 ms tick and checks that no interrupt reaches it there, and that the writer is still stopped afterwards. Whether PipedRDD and other interrupters should do the same is a separate question that I didn't change here.

[dongjoon-hyun]

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:

pythonSide.socket().setSoTimeout(TimeUnit.MINUTES.toMillis(1).toInt)
assert(pythonSide.socket().getInputStream.read() === -1)

[viirya]

Fixed: it now reads through the socket adaptor with a 1-minute SO_TIMEOUT.

[dongjoon-hyun]

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.

[viirya]

That comment went away with the L205 change. The new comment in markTaskCompleted says why the pre-completion listeners run before completed is set, and the TaskContext.addTaskPreCompletionListener scaladoc now says which thread runs them: the one that marks the task completed, before isCompleted() becomes true, with completion listeners added meanwhile from other threads only queued; one added after completion is called immediately.

[dongjoon-hyun]

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."

[viirya]

Thanks, reworded along the lines you suggested.

[dongjoon-hyun]

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.

[viirya]

Done: every warning names the writer thread, and the first one also carries its stack trace. The stack is taken outside the task's monitor and dropped if the writer has exited in the meantime. The same-thread logError carries the current stack as well.

[dongjoon-hyun]

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.

[viirya]

Fixed: the finally blocks now release every latch, and the started assertions moved inside the try blocks, so a failing assertion no longer leaves a writer behind.

[dongjoon-hyun]

24. [Nit] One more started assertion is still before its try

Thanks. I missed one place when I listed L138 and L184: in "an interrupt does not stop the wait and is not passed to later listeners", assert(writer.started.await(1, TimeUnit.MINUTES)) (now L223) is still before the try at L227, whose finally is the only place that counts down writer.release. If it fails, the writer stays in awaitIgnoringInterrupts(release) for the rest of the core test JVM and keeps a pipelined-writer pool thread busy, as in the other two tests. Could L222-L226 move into the try, with var completion: Thread = null before it, as L142-L149 and L191-L197 now do?

[viirya]

Thanks, moved into the try, with var completion: Thread = null before it, like the other two tests.

[HyukjinKwon]

Review summary

Both P3 items from the previous review are moot: they covered the interrupt handling of the markTaskCompleted handover wait and a throwing deferred listener in the release loop. That handover wait and release loop are no longer in this PR. They moved to SPARK-60222, and invokeListeners is back to the merge-base code. The new design closes the deferred writer paths structurally. Pre-completion listeners run on the completing thread before completed is set, so the writer thread can only queue listeners and can never run its own stop listener. A kill thread can still claim the invocation and run failure listeners while the writer wait is in progress. dongjoon-hyun's latest round is addressed in this head, including the CloseShieldOutputStream fix for the interrupt that closed the Arrow writer's buffer, the setup inside the try, the interruptible() check, the abstract addTaskPreCompletionListener, and the deterministic test ordering.

One new non-blocking P3 item concerns the new pre-completion API itself. The branch that runs a pre-completion listener added after the task has completed has failure handling (marking the task failed and throwing a TaskCompletionListenerException) that no test exercises.

Findings

1 total: 0 P0, 0 P1, 0 P2, 1 P3.

Nit (P3)

• The failure path of a pre-completion listener added after completion has no test — General
  The new after-completion branch of TaskContextImpl.addTaskPreCompletionListener (callListenerNow) runs the listener at once through invokeTaskPreCompletionListener. If the listener throws, the branch calls markTaskFailed and throws a TaskCompletionListenerException with the error suppressed. No test reaches that failure handling. The TaskContextSuite cases that add a pre-completion listener after completion (P3 in "pre-completion listeners run before all completion listeners", and P in "a pre-completion listener added after completion runs immediately") only use listeners that succeed. The failing-listener tests add their listener before markTaskCompleted, so they only cover the in-loop path. If the branch stopped throwing, skipped markTaskFailed, or dropped the cause, every test would still pass. A small test that completes a TaskContext.empty(), adds a throwing pre-completion listener, and checks the thrown exception, its suppressed cause, and context.isFailed() would pin it. This only affects late registration on the internal API: startPipelinedWriter registers before completion.

  Verification:

  • Regression: The new test fails if the after-completion branch of addTaskPreCompletionListener stops throwing, stops calling markTaskFailed, or drops the listener's exception, and it passes with the current implementation.

Re-review status

Prior AI findings: 2 addressed, 0 still present; additional unresolved findings in this review: 1.

New attribution: 1 newly introduced, 0 late catch, 0 previously raised, 0 unattributed.

Remaining prior AI findings

No prior AI findings remain.

[viirya]

@HyukjinKwon Thanks, added in 820c7ac: "a failing pre-completion listener added after completion fails the task" completes a TaskContext.empty(), adds a throwing pre-completion listener, and checks the thrown TaskCompletionListenerException, its suppressed cause and context.isFailed(). It fails if the after-completion branch stops throwing, skips markTaskFailed, or drops the listener's exception.

[dongjoon-hyun]

Thank you for the update, @viirya. 43dff78152f runs the pre-completion listeners before completed is set, as discussed in item 8, which closes the deferred writer paths, and CloseShieldOutputStream fixes item 1. 820c7acead2 adds the test for a failing pre-completion listener added after completion. Items 1-15 are addressed; item 24 below adds one place that I missed in item 14.

I left 9 inline comments on this head, plus a reply in the item-14 thread, continuing the numbering from items 1-15. Item 16 matters most: 43dff78152f moved the first interrupt out of the state check, so a writer that starts in between is not interrupted until the first 10-second warning. Summary:

16. [Low] A writer that starts between interruptWriter() and the state check is not interrupted until the first warning (PythonRunner.scala:313)

    On master and at 00da9fb69d4, the writer was either kept from starting or interrupted. Moving interruptWriter() into the PipelinedWriterRunning case restores that.

17. [Low] The writer is submitted before its stop listener is registered (PythonRunner.scala:216)

    A compute() on a helper thread that straddles task completion, e.g. PipedRDD's stdin writer, can let the writer start while the task thread frees its input. This predates the PR; registering the listener first closes it.

18. [Low] awaitExit can still end early, or put off its re-interrupt (PythonRunner.scala:376)

19. [Low] Correction to my item 5: context.interruptible() is task-wide, not per writer (PythonRunner.scala:290)

20. [Low] The merged TaskCompletionListenerException nests the completion listeners' exception (TaskContextImpl.scala:246)

21. [Nit] The round-2 qualification is gone from the addTaskPreCompletionListener scaladoc again (TaskContext.scala:152)

22. [Nit] Correction to my round-3 suggestion: the task reaper does not bound this wait in local mode (PythonRunner.scala:211)

23. [Nit] The ClosedByInterruptException comment is stale (PythonRunner.scala:1546)

24. [Nit] One more started assertion is still before its try (BasePythonRunnerSuite.scala:223, reply in the item-14 thread)

25. [Nit] interruptWriter()'s result is unused, and writerThread re-reads a known thread (PythonRunner.scala:286)

[dongjoon-hyun]

16. [Low] A writer that starts between interruptWriter() and the state check is not interrupted until the first warning

43dff78152f moved the first interrupt out of the state match into its own synchronized block (this line), ahead of the one that reads the state (L314-L323). If the pool thread enters run() between the two, interruptWriter() still sees PipelinedWriterNotStarted and does nothing (L297), and L320 then sees PipelinedWriterRunning and goes to awaitExit(). The writer runs uninterrupted: it sends the init message and pulls and serializes input until it blocks in channel.write, since nobody reads the worker's output anymore. Its first interrupt comes from the warning tick after 10 s (L369-L373), and that warning blames I/O that ignores interrupts, which is misleading here.

The window is narrow (the task has to complete just as the pool starts the writer, e.g. a task that fails or is killed right after compute()), but it did not exist before: 00da9fb69d4 interrupted inside case PipelinedWriterRunning(thread) of the same block, and on master cancel(true) either kept the runnable from starting or interrupted its runner. No test covers it, since "stopping interrupts a running writer right away" waits for WAITING first.

Could interruptWriter() move into that case, with the comment at L308-L312? The monitor is reentrant, the lock order (this, then TaskContext) stays the same, and the interruptible() check and the same-thread guard stay in interruptWriter().

            case PipelinedWriterRunning(thread) =>
              interruptWriter()
              Some(thread)

[viirya]

Good catch, my regression in 43dff78. Fixed in 9b0a1f8 as you suggested: the first interrupt now happens inside case PipelinedWriterRunning(thread) of the same synchronized block. interruptWriter takes that thread, and keeps the same-thread guard and the interruptible() check. I don't have a deterministic test for this window, since it needs the pool thread to enter run() between two statements of the stopping thread, and the window no longer exists.

[dongjoon-hyun]

17. [Low] The writer is submitted before its stop listener is registered

This submits the writer before L217 registers the stop listener. On the task thread the order does not matter, but compute() can also run on a helper thread that task completion does not stop, e.g. PipedRDD's stdin writer, which computes the parent partition on its own thread (PipedRDD.scala:149) and is only interrupted at completion. If such a call registers its worker.stop listener (L703) before the task thread sets completed (TaskContextImpl.scala:227), and reaches L217 after it, the stop listener runs right away on the helper thread, while the task thread is already running worker.stop and the HybridRowQueue close. In between, the pool can start the writer, which flushes the init message and adds rows to the queue that is being freed. This is the one path left where the writer can still run after completed is set.

It predates this PR (master also submitted first), so this is only a hardening, but it is a one-line move. Could the listener be registered before the submit? A listener that runs before the writer starts then finds PipelinedWriterNotStarted and calls finish(), so the writer never runs, and the queued worker.stop still unblocks the helper's reader.

[viirya]

Done in 9b0a1f8: the stop listener is registered before the writer is submitted, with a comment on why. New test "a writer started after the task completed never runs": it completes the task first and then calls startPipelinedWriter, so the listener runs at once, finds the writer not started, and the writer never runs. It fails when the submit comes first and the pool starts the writer before the listener is added.

[dongjoon-hyun]

18. [Low] awaitExit can still end early, or put off its re-interrupt

Two things in this loop (L367-L378):

• It absorbs only InterruptedException. Any other Throwable from a tick, or from the first interruptWriter() at L313, ends the wait while the writer may still run, e.g. an OutOfMemoryError while warnStillWaiting takes the writer's stack or builds the message (L388-L400), or an exception out of thread.interrupt() (L291), which runs the interruptor of the channel the writer is blocked in on this thread and swallows only IOException there. invokeTaskPreCompletionListener records it (TaskContextImpl.scala:258-269), and markTaskCompleted goes on to the completion listeners, which is the use-after-free that L349-L352 rules out. The error only ends up suppressed in a TaskCompletionListenerException, so the executor does not treat it as fatal either. (YARN kills the JVM on OOM by default, so this matters mostly elsewhere.)
• Each InterruptedException restarts a full warnIntervalMs. When the first interrupt was skipped (L290) or consumed, the re-interrupt at L373 is what stops the writer, and every kill that interrupts the task thread pushes it, and the warning, back by up to another 10 s.

Could the tick keep its own deadline, and could the warning and the interrupts be best effort, e.g. catch everything but InterruptedException around L313 and L371-L373, keep the first failure, and rethrow it once exited is zero? NonFatal alone would miss the OOM, and the catch should not cover exited.await itself, since the wait can allocate and a lasting OOM would make the loop spin.

[viirya]

Both done in 9b0a1f8:

• The warning and the interrupts are now best effort. interruptWriter returns a failure to interrupt instead of throwing it. The per-tick warning and interrupt are wrapped in a catch of everything but InterruptedException, and exited.await is not covered by it, as you noted. The first failure is kept and thrown once the writer has exited. The new test "a failure to interrupt the writer does not end the wait" runs the writer on a thread whose interrupt() throws, and checks that stopping keeps waiting and throws that failure only after the writer has exited. That test also found a bug in my first version: failure.orElse(awaitExit(thread)) skipped the wait entirely once the first interrupt had failed, since orElse takes its argument by name.
• The warnings keep their own deadline, so an interrupt of the waiting thread no longer pushes the next warning, and its re-interrupt, back.

[dongjoon-hyun]

27. [Low] Correction to my item 18: a failing warning skips that tick's re-interrupt

I suggested catching everything but InterruptedException "around L313 and L371-L373", and L371-L373 then were the warning, warned = true and the re-interrupt, so 9b0a1f86199 put all three into one try (now L388-L401). If warnStillWaiting throws (L389), the re-interrupt at L391-L397 is skipped for that tick, and since warned stays false, the next tick takes the writer's stack again (L417) and can fail the same way.

So if the warning fails on every tick, e.g. through an appender configured with ignoreExceptions="false" whose target is down, a logging binding that throws, or a heap too full for the stack trace, a writer that needs a second interrupt never gets one: one that consumed the first interrupt, the case that L365-L368 describe, or one whose first interrupt was skipped because context.interruptible() was false (L290). Such a writer blocks in channel.write (L1626), and unless the task is killed, the wait does not end. Sorry for the imprecise suggestion.

Could the warning get its own try, with the tick's interrupt in a separate one after it?

            try {
              warnStillWaiting(thread, startNs, withStackTrace = !warned)
              warned = true
            } catch {
              case t: Throwable if !t.isInstanceOf[InterruptedException] =>
                failure = failure.orElse(Some(t))
            }
            try {
              val interruptFailure = synchronized {
                state match {
                  case PipelinedWriterRunning(_) => interruptWriter(thread)
                  case _ => None
                }
              }
              failure = failure.orElse(interruptFailure)
            } catch {
              case t: Throwable if !t.isInstanceOf[InterruptedException] =>
                failure = failure.orElse(Some(t))
            }
            nextWarningNs = System.nanoTime() + intervalNs

exited.await stays outside both, as before. Running the writer of "the writer is interrupted again if it consumed the interrupt" on a thread whose getStackTrace throws would pin this; that test fails at this head.

[viirya]

Thanks, fixed in 27d5537: the warning and the tick's interrupt now have separate best-effort try blocks, and exited.await stays outside both. The new test "a failing warning does not skip the repeated interrupt" runs the writer that consumes the first interrupt on a thread whose getStackTrace throws, and checks that the writer still gets its second interrupt and that stopping returns, throwing the warning's failure. With the two blocks merged back, it times out.

[dongjoon-hyun]

19. [Low] Correction to my item 5: context.interruptible() is task-wide, not per writer

I suggested skipping "a writer that is inside createResourceUninterruptibly" by checking context.interruptible(). But _interruptible (TaskContextImpl.scala:107-108) is a single flag for the whole task: any thread that has the TaskContext set clears it when it enters createResourceUninterruptibly (TaskContextImpl.scala:508) and sets it back to true when it leaves (TaskContextImpl.scala:516), with no count and no thread identity. So this check does not tell whether this writer is in a builder:

• It also skips this writer while another thread of the task creates a resource, although this writer creates nothing, e.g. while the inner writer of two stacked Python UDFs with different eval types opens its next CSV file, or a TRANSFORM stdin writer does. This writer then blocks in channel.write and is interrupted only at the first warning at which no thread of the task is in a builder.
• When two creations overlap, the first one to finish sets the flag back to true while the other is still building, so a writer inside its own builder can still be interrupted there.

Task.kill has had the same gap since SPARK-50768, and both cases are narrower than master's unconditional cancel(true), so I don't think this PR needs a per-thread or counted flag; that would be a separate [CORE] JIRA. But the scaladoc at L282-L284, the comment at L311-L312 and item 3 of the PR description say it is about the writer thread. Could they describe the check instead, e.g. "not while context.interruptible() is false, which is only while some thread of the task, not necessarily the writer, creates a resource uninterruptibly"?

[viirya]

Thanks for the correction. Done in 9b0a1f8: the interruptWriter scaladoc, the comment at the first interrupt and item 3 of the PR description now describe the check itself. The writer is not interrupted while context.interruptible() is false, which is while some thread of the task, not necessarily the writer, creates a resource uninterruptibly. I agree that a per-thread or counted flag would be a separate [CORE] change.

[dongjoon-hyun]

20. [Low] The merged TaskCompletionListenerException nests the completion listeners' exception

When a pre-completion listener and a completion listener both fail, completionException is the TaskCompletionListenerException that invokeListeners already built from the completion listeners' failures and error (L431-L435), and this adds it as a single entry of a second one. The outer message then embeds the inner one: with two failing completion listeners it reads Exception 0: pre / Exception 1: Exception 0: c1 / Exception 1: c0, and if the task body failed, the Previous exception in task: ... block with the task error's stack appears twice. runTaskWithListeners attaches the outer exception to that error as suppressed, so the reported ExceptionFailure carries the error's stack four times instead of twice, and the completion listeners' own exceptions are only reachable through the inner exception's getSuppressed. The classification of the failure is the same, and today this needs the writer's stop listener to throw.

Could the completion listeners' exceptions themselves be merged? For example, this would also drop the Option round trip:

    try {
      invokeTaskCompletionListeners(error)
    } catch {
      case e: Throwable if preCompletionExceptions.nonEmpty =>
        preCompletionExceptions ++= (e match {
          case tcle: TaskCompletionListenerException => tcle.getSuppressed.toSeq
          case _ => Seq(e)
        })
    }
    if (preCompletionExceptions.nonEmpty) {
      throw newTaskCompletionListenerException(preCompletionExceptions.toSeq, error)
    }

The test at TaskContextSuite.scala:800-814 passes either way; keeping references to the two exceptions there and asserting e.getSuppressed.toSeq === Seq(pre, post) would pin it.

[viirya]

Done in 9b0a1f8, along the lines of your snippet: the completion listeners' own failures (the suppressed exceptions of the TaskCompletionListenerException that invokeListeners built) are merged, and the Option round trip is gone. The test now keeps references to both exceptions and asserts e.getSuppressed.toSeq === Seq(pre, post).

[dongjoon-hyun]

21. [Nit] The round-2 qualification is gone from this scaladoc again

This again says the listener runs "before every listener added with addTaskCompletionListener ... even if that listener was added later". After round 2 it read "before every pending listener ... Completion listeners that have already run are not affected, so to run before all of them, the listener must be added before the task completes", and 43dff78152f dropped that. It still applies: markTaskCompleted sets completed (TaskContextImpl.scala:227) before it runs the completion listeners (TaskContextImpl.scala:239), and a pre-completion listener added from then on runs right away on the adding thread (TaskContextImpl.scala:152-157), outside invokeListeners. So it runs after the completion listeners that have already run, and it can overlap one that is still running on the completing thread, whereas at 00da9fb69d4 it was queued and ran next on that thread. The new test pins exactly this (TaskContextSuite.scala:838 and 843), and L156 mentions neither.

Could it say so, e.g. "..., as long as this listener is added before the task completes. ... A listener added after the task completes, including while the completion listeners run, is called immediately on the calling thread: it runs after the completion listeners that have already run, and may overlap one that is still running."? Alternatively, handing a late listener to addTaskCompletionListener would keep it sequential, as at 00da9fb69d4.

[viirya]

I went with your alternative in 9b0a1f8: a pre-completion listener added after the task completed is now handed to addTaskCompletionListener. So it runs sequentially like a completion listener added at that point, queued if the completion listeners are still running, and the separate run-now branch with its own failure handling is gone. The scaladoc now says the ordering holds only if the listener is added before the task completes, and that a later one is handled like one added with addTaskCompletionListener and runs after the completion listeners that have already run. The test "a pre-completion listener added after completion runs as a completion one" checks that one added from another thread while the completion listeners run is queued and runs next on the completing thread, and that one added after all listeners have run runs right away. HyukjinKwon's test for a failing late listener still passes, through invokeListeners.

[dongjoon-hyun]

22. [Nit] Correction to my round-3 suggestion: the task reaper does not bound this wait in local mode

L210-L212 follow my round-3 wording, which is wrong in local mode: once the kill timeout is exceeded, TaskReaper only logs "not killing JVM because we are running in local mode" (Executor.scala:1324-1327), since LocalSchedulerBackend creates its executor with isLocal = true. Pipelined mode is not limited to cluster deployments, and the PR's PySpark tests run it on local[4], so in local mode, e.g. in a PySpark shell, nothing bounds this hang. Sorry for the imprecise suggestion. Could it add "except in local mode, where the reaper only logs an error"?

[viirya]

Thanks, added "except in local mode, where the reaper only logs an error" to the scaladoc and the PR description.

[dongjoon-hyun]

23. [Nit] The ClosedByInterruptException comment is stale

L1543 was updated, but this still says "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." The interrupt now comes only from the stop listener (interruptWriter() at L313 and L373), not from a cancel. The JDK closes the socket only when it lands in flushToSocket's worker.channel.write. The Arrow writers can also be interrupted in the in-memory Channels.newChannel(dataOut) they write through, and closing that only closes dataOut, which CloseShieldOutputStream keeps away from the buffer, so worker.channel stays open and Python gets no EOF from it (the new test at BasePythonRunnerSuite.scala:541-570 still reads later bytes from the socket). Either way no reader is waiting: the task thread that reads the worker's output is the one running the stop listener. Could it say something like "Interrupted the same way, during a write to an interruptible channel, which the JDK closes: worker.channel in flushToSocket, or the Channels.newChannel(dataOut) that the Arrow writers write through, in which case the socket stays open (see CloseShieldOutputStream)"?

[viirya]

Thanks, reworded along the lines you suggested.

[dongjoon-hyun]

25. [Nit] interruptWriter()'s result is unused, and writerThread re-reads a known thread

This documents and returns whether it interrupted the writer (L284), but both callers drop the result (L313, L373), and as a private method nothing else can use it. writerThread (L273-L279) is used only at L387, to re-read the thread that stopAndAwaitExit already has at L320. Once the state is PipelinedWriterRunning(thread), it only becomes PipelinedWriterFinished, together with the count-down of exited (L245-L248), so during the wait writerThread is None only when exited is zero, which L389 checks anyway. Could interruptWriter() return Unit, and the thread be passed down (awaitExit(thread), warnStillWaiting(thread, ...))? The stack would still be taken outside the lock and dropped if the writer exited meanwhile.

[viirya]

Done: interruptWriter now takes the thread and returns only a failure to interrupt (for #18), writerThread is gone, and awaitExit and warnStillWaiting take the thread. The stack is still taken outside the lock and dropped if the writer exited meanwhile.

[dongjoon-hyun]

Could you resolve the conflicts once more, @viirya ?

[viirya]

@dongjoon-hyun I resolved the conflicts. Thanks!

[dongjoon-hyun]

Thank you for the update, @viirya. 9b0a1f86199 addresses items 16-25, and fc2437733a0 takes CloseShieldOutputStream verbatim from #59565, which is merged now.

I left 9 inline comments on this head (27ac39752da), plus a reply in the item-18 thread, continuing the numbering from items 16-25. Item 26 matters most: the alternative I offered in item 21 queues a late stop listener while another thread runs the listeners, so the register-before-submit from item 17 no longer keeps a late writer from starting. Summary:

26. [Low] Correction to my items 17 and 21: a writer started after task completion can still run (PythonRunner.scala:222)

    While another thread invokes the listeners, a late stop listener is only queued and the writer is submitted anyway, so the close listeners that the writer adds lazily run before the stop listener. Submitting only if !context.isCompleted() closes it.

27. [Low] Correction to my item 18: a failing warning skips that tick's re-interrupt (PythonRunner.scala:389, reply in the item-18 thread)

    With the warning and the re-interrupt in one try, a warning that fails on every tick leaves a writer that needs a second interrupt running, and the wait does not end.

28. [Nit] Correction to my item 20: the case _ => Seq(e) arm hides a fatal error (TaskContextImpl.scala:243)

29. [Nit] This comment still says a late listener is called immediately (TaskContextImpl.scala:223)

30. [Nit] This test duplicates CloseShieldOutputStreamSuite from #59565 (BasePythonRunnerSuite.scala:501)

31. [Nit] This builds the function differently from the SPARK-60343 test below (BasePythonRunnerSuite.scala:537)

32. [Nit] This is Guava's Uninterruptibles.awaitUninterruptibly (BasePythonRunnerSuite.scala:89)

33. [Nit] ConcurrentLinkedQueue is spelled out although the import changed (TaskContextSuite.scala:822)

34. [Nit] These two helpers now split the mark* overrides (TaskContextImpl.scala:255)

35. [Nit] awaitExit could take the first failure (PythonRunner.scala:340)

[dongjoon-hyun]

26. [Low] Correction to my items 17 and 21: a writer started after task completion can still run

The register-before-submit that I asked for in item 17 relies on a late stop listener running at once: "A listener that runs before the writer starts then finds PipelinedWriterNotStarted and calls finish(), so the writer never runs". The alternative I offered in item 21, which 9b0a1f86199 adopted, removed that. A pre-completion listener added after completed is set now goes to addTaskCompletionListener (TaskContextImpl.scala:152-156), which only pushes it while another thread is invoking listeners (TaskContextImpl.scala:351-353). This line then submits the writer anyway, and run() starts it from PipelinedWriterNotStarted (L253-L261). So the comment at L216-L218 ("the listener runs at once and keeps the writer from ever starting") holds only when no thread is invoking listeners.

In the interleaving that item 17 describes, the queued listener still lands above worker.stop and the HybridRowQueue close, so it runs before them. It can still come too late:

• The task thread completes while a helper thread that computes the partition (e.g. PipedRDD's stdin writer, PipedRDD.scala:149) is still in env.createPythonWorker (L721). If the task thread has already run the queue and scan close listeners and is in an older listener when the helper reaches L219, worker.stop (L731) and the stop listener are both only queued, configureBlocking(true) (L790) succeeds, and the writer starts, sends the init message and pulls input. The readers that it opens register their close listeners on the writer thread, e.g. ParquetReaderCallback.initIfNotAlready of the DSv2 Parquet reader (ParquetPartitionReaderFactory.scala:389-391) or createResourceUninterruptibly (TaskContextImpl.scala:512). Those land above the stop listener, so the task thread runs them first, while the writer still reads through them; with off-heap column vectors, that is the use-after-free that this PR fixes.
• A kill thread holds the invocation, running failure listeners after completion. It pops only the failure listeners, so nothing runs the queued stop listener, and the writer keeps running after markTaskCompleted returns and cleanUpAllAllocatedMemory frees the task's memory. The handover itself is SPARK-60222's, but here it lets a writer run after completion.

"a writer started after the task completed never runs" (BasePythonRunnerSuite.scala:249) completes the context first, so no thread is invoking listeners. compute() hardly reaches that case, since worker.stop would already have closed the channel, and even there, swapping L219-L221 and L222 back fails the test only when the pool wins the race.

Could the writer be submitted only if the task has not completed yet?

    context.addTaskPreCompletionListener(new TaskCompletionListener {
      override def onTaskCompletion(context: TaskContext): Unit = writerTask.stopAndAwaitExit()
    })
    if (!context.isCompleted()) {
      pipelinedWriterThreadPool.submit(writerTask)
    }

completed only goes from false to true, and only once the pre-completion stack is empty (TaskContextImpl.scala:221-226). If it is still false after the registration, the listener is a pre-completion one and stops the writer before completed is set. If it is true, the writer is never submitted, and the queued worker.stop still unblocks the helper's reader. L216-L218 would then describe the check. A test could call startPipelinedWriter from another thread inside a completion listener, which holds the invocation, and check that the writer never runs; it fails at this head.

[viirya]

Thanks, confirmed: with item 21, a late stop listener is only queued while another thread runs the listeners, so register-before-submit no longer kept the writer from starting. Fixed in 27d5537 as you suggested: the writer is submitted only if !context.isCompleted() once the listener is added, and the comment explains why that is enough. The new test "a writer started while another thread runs the listeners never runs" calls startPipelinedWriter from a helper thread inside a completion listener, and checks that the writer never runs. It fails without the check.

[dongjoon-hyun]

28. [Nit] Correction to my item 20: the case _ => Seq(e) arm hides a fatal error

invokeListeners catches every Throwable that a listener throws (L380) and throws only the TaskCompletionListenerException it builds (L431-L436), so this arm only sees errors from its own bookkeeping, in practice an OutOfMemoryError in addSuppressed, listenerExceptions += e, logError or the exception construction (L419-L435). Without a failed pre-completion listener, such an error propagates from markTaskCompleted as it is, and after a successful task body the executor handles it as fatal (Executor.scala:1144). With one, the arm from my item-20 snippet adds it to the suppressed exceptions of the merged TaskCompletionListenerException, and Executor.isFatalError (Executor.scala:1654-1665) follows only the cause chain, so the executor keeps running after the OOM. This needs two failures at once, and merging any Throwable here predates item 20, so it is minor. Could the case match only the wrapper, so that fatal errors propagate as on master?

      case e: TaskCompletionListenerException if preCompletionExceptions.nonEmpty =>
        // Merge the completion listeners' own failures, not the exception that wraps them.
        preCompletionExceptions ++= e.getSuppressed.toSeq

[viirya]

Thanks, done: the case matches only TaskCompletionListenerException now, with a comment that other errors come from invokeListeners itself and propagate as on master.

[dongjoon-hyun]

29. [Nit] This comment still says a late listener is called immediately

"so that one added later is called immediately instead of being left in the stack" describes the callListenerNow branch of 43dff78152f. Since 9b0a1f86199, such a listener goes to addTaskCompletionListener (L152-L156) and is queued while this thread runs the completion listeners right after this (L237), as L153-L154, TaskContext.scala:156-159 and the comment at TaskContextSuite.scala:833-834 now say. Could it read, e.g., "..., so that one added later is handed to addTaskCompletionListener instead of being left in the stack"?

[viirya]

Thanks, reworded: "so that one added later is handed to addTaskCompletionListener instead of being left in the stack".

[dongjoon-hyun]

30. [Nit] This test duplicates CloseShieldOutputStreamSuite from #59565

With #59565 merged (058c570f99e on master, also on branch-4.0 through branch-4.3 and branch-4.x, with the same CloseShieldOutputStream.scala blob as this PR), its "SPARK-60305: an interrupted channel write does not close the shielded buffer" does what this test does: a channel over DataOutputStream(CloseShieldOutputStream(bufferStream)), a pending interrupt, ClosedByInterruptException, and a check that the buffer still works. This PR's own case is covered by "an interrupted Arrow-style write leaves the writer's buffer usable" (L589-L618), which also fails without the shield, since writeInt(42) then hits the closed buffer and sets writer.exception. Could this test go, together with the DirectByteBufferOutputStream import that only it uses (L37)? Merging master would also drop CloseShieldOutputStream.scala from this diff. Separately, newDataOut() (PythonRunner.scala:1552-1553) wraps a new CloseShieldOutputStream around bufferStream for every buffer, while ReaderInputStream on master keeps one in writerStream; the shield has no state, so one per writer would do.

[viirya]

Done. I merged master, so CloseShieldOutputStream.scala is no longer in the diff, and I dropped the duplicate test with its import. The Arrow-style test now does the interrupted write in a later buffer, so it also covers the stream that the writer creates after each flush (SPARK-60343); it fails if that stream is not shielded. The writer also creates one CloseShieldOutputStream now, as ReaderInputStream does, and wraps a new DataOutputStream around it for every buffer.

[dongjoon-hyun]

31. [Nit] This builds the function differently from the SPARK-60343 test below

new JHashMap[String, String](), new JArrayList[String]() and new JArrayList[Broadcast[PythonBroadcast]]() are the only uses of the JArrayList/JHashMap aliases (L25) and of the Broadcast import (L36). The SPARK-60343 test that the master merge brought in passes Collections.emptyMap() and Collections.emptyList() for the same arguments (L621-L622), and nothing on these paths modifies them (only compute() puts into envVars). Could this do the same, so that L25 matches master again and L36 can go? Also, readFromPythonSide (L567-L570) has one caller, while the Arrow-style test sets the same timeout inline (L608-L611). Setting SO_TIMEOUT once in this helper, right after server.accept() (L544), would keep the reads through the socket adaptor from item 10 for both tests.

[viirya]

Done: the helper passes Collections.emptyMap() / Collections.emptyList(), so the aliases and the Broadcast import are gone and the import block matches master apart from what this PR uses. SO_TIMEOUT is set once in the helper right after server.accept(), and readFromPythonSide and the inline timeout are gone.

[dongjoon-hyun]

32. [Nit] This is Guava's Uninterruptibles.awaitUninterruptibly

Apart from restoring the interrupt status at the end, this loop is Uninterruptibles.awaitUninterruptibly(latch), which core tests already use for a callback that keeps waiting through an interrupt (EventLoopSuite.scala:230, DAGSchedulerSuite.scala:8025). None of the four callers needs the interrupt dropped: after these waits the writers only queue a listener, count down a latch or set exited, and the pool clears a leftover interrupt before its next task. Unlike the wait in awaitExit, these test writers may keep it. Could they use it? Also, try { new CountDownLatch(1).await() } catch { case _: InterruptedException => } appears three times (L388-L392, L416-L420, L478-L482); a small helper next to this one would cover those.

[viirya]

Done: awaitIgnoringInterrupts is Uninterruptibles.awaitUninterruptibly(latch) now, and the three "wait until interrupted" blocks use a small awaitInterrupt() helper.

[dongjoon-hyun]

33. [Nit] ConcurrentLinkedQueue is spelled out although the import changed

L822 and L871 write new java.util.concurrent.ConcurrentLinkedQueue[(String, Thread)](), while this PR already edits the java.util.concurrent import at L21. Could it be imported there? The tests at L820 and L869 also define the same makeLoggingListener (L824-L828, L875-L879), and the tests at L869 and L911 share about 20 lines: the two latches, the blocking pre-completion listener, the daemon completion thread, the release in finally and the timed join. A helper that runs a body while the pre-completion listener blocks would keep both tests as they are.

[viirya]

Done: ConcurrentLinkedQueue is imported, the new tests share a loggingListener(invocations, id) helper, and the two tests that block a pre-completion listener use completeWhilePreCompletionListenerBlocks(context) { ... }, which returns the completing thread.

[dongjoon-hyun]

34. [Nit] These two helpers now split the mark* overrides

Since 9b0a1f86199 removed the late-listener branch, invokeTaskPreCompletionListener and newTaskCompletionListenerException each have one caller (L233, L247), and they sit between markTaskCompleted and getTaskFailure, which were adjacent at the merge base. AGENTS.md asks to "follow the sectioning the file already uses and put the new member with the code it belongs with". Could they move below getTaskFailure, next to invokeTaskCompletionListeners and the other private invoke* helpers? newTaskCompletionListenerException would then also sit next to the same construction in invokeTaskInterruptListeners (L327-L331) and invokeListeners (L431-L435).

[viirya]

Done: invokeTaskPreCompletionListener and newTaskCompletionListenerException moved below getTaskFailure, next to invokeTaskCompletionListeners and the other private helpers.

[dongjoon-hyun]

35. [Nit] awaitExit could take the first failure

L340-L342 keep the first interrupt failure out of awaitExit, so they need this comment and a temporary to avoid the by-name orElse that, as you wrote in the item-18 thread, skipped the wait in your first version. Passing the failure in would rule that out, e.g. failure = awaitExit(thread, failure) here and var failure = firstFailure in awaitExit (L379). The first failure still wins, and the wait still runs either way.

[viirya]

Done: awaitExit(thread, failure) starts from the first failure, so the temporary and its comment are gone.

[dongjoon-hyun]

+1, LGTM (Pending CIs) with the above comments.

[dongjoon-hyun]

Merge Summary:

• merged into master ca1f5c5
• merged into branch-4.x 4407a0b

Posted by merge_spark_pr.py

[dongjoon-hyun]

Could you make a backporting PR to branch-4.3? There was a conflict on branch-4.3, @viirya .

[viirya]

Thank you @dongjoon-hyun! I opened the backport PR for branch-4.3.