[SPARK-60305][PYTHON] Keep the Python runner's input buffer alive when an Arrow write is interrupted

[viirya]

What changes were proposed in this pull request?

BasePythonRunner.ReaderInputStream (the default, non-pipelined Python runner path) serializes the task input into a DirectByteBufferOutputStream. This PR hands the writer that buffer only through a new CloseShieldOutputStream, a FilterOutputStream that forwards write(b, off, len) and write(b) and whose close() only flushes. The buffer is now freed only on the thread that writes into it: after the last input is sent, as before, or, new in this PR, when an exception escapes writeAdditionalInputToPythonWorker() (a kill during a write, a writer failure, a socket write failure). Before, that path left the buffer to GC unless the interrupt close happened to free it, and with the shield that close no longer frees it at all. After freeing it, the stream also replaces buffer, which may be null or still point into the freed memory, with an empty one, so a later call fails fast in bufferStream instead of touching the freed memory.

#59332 (SPARK-60113) adds the same CloseShieldOutputStream file for the pipelined writer, so the two PRs stay independent: whichever merges second has no diff for it. If #59332 merges first, this PR's commit no longer adds the file, so its backports to branch-4.2, 4.1 and 4.0, which #59332 doesn't reach, need to bring it along. This PR does not touch PipelinedWriterRunnable.

Why are the changes needed?

The Arrow inputs (PythonArrowInput, ColumnarArrowPythonInput) write dataOut through Channels.newChannel(dataOut), via ArrowStreamWriter and WriteChannel. That channel is an AbstractInterruptibleChannel: each chunk is written between begin() and end(), and closing the channel closes dataOut. When another thread interrupts the task thread inside such a write, e.g. Task.kill(interruptThread = true) for a losing speculative attempt or a job cancelled with interrupt-on-cancel, Thread.interrupt() runs the channel's interruptor on the interrupting thread. That closes DataOutputStream -> DirectByteBufferOutputStream, which frees the direct buffer while the task thread may still be copying into it (or copying the old buffer in grow()). The result can be a JVM crash or silent memory corruption. A Java-level error from the write does not surface: the channel's end() throws ClosedByInterruptException, which replaces it, so the task still looks like a clean kill.

With the shield, the interrupt still closes the channel and the task thread stops with ClosedByInterruptException right after the current chunk, but the buffer is not freed under it.

This affects Arrow-based Python execution (pandas/Arrow UDFs, Arrow-optimized Python UDFs, mapInPandas/mapInArrow, applyInPandas, ...) since SPARK-44705 made the runner single-threaded (4.0.0). The pickled BatchEvalPython path writes dataOut directly and is not affected. Other streams in core and sql/core that are written through Channels.newChannel (ByteArrayOutputStream, ChunkedByteBufferOutputStream, ByteBufferOutputStream) do not free native memory on close, and PythonPlannerRunner's DirectByteBufferOutputStream is not written through a channel.

Thanks to @dongjoon-hyun, who found this in the review of #59332 (#59332 (comment)).

Does this PR introduce any user-facing change?

Yes, as a bug fix: killing a task that runs an Arrow-based Python UDF with thread interruption no longer risks freeing the input buffer while it is being written. There is no API or configuration change.

How was this patch tested?

New CloseShieldOutputStreamSuite:

• CloseShieldOutputStream forwards bulk writes as bulk writes and does not close the underlying stream.
• A channel write with the interrupt flag already set closes the channel on the writing thread, and the shielded buffer stays open.
• A test through the real ReaderInputStream, over a local socket, with a writer that writes through Channels.newChannel(dataOut) like the Arrow inputs. It interrupts the writing thread inside the channel write, then checks that the blocked write into the buffer still succeeds, that the thread stops with ClosedByInterruptException, and that the buffer is freed afterwards. A later read must then fail with Cannot call methods on a closed DirectByteBufferOutputStream, which checks both that the buffer was freed and that the freed buffer is not reused. Without the shield in writeBufferedInputToPythonWorker(), the blocked write fails with the same error, because the interrupting thread already freed the buffer; without the one on the open path, constructing the ReaderInputStream already fails with it, because the test writer's open closes its stream.

Each test fails when the part of the change it covers is reverted. The first two tests port unchanged to all branches; the third only needs the ReaderInputStream constructor call adjusted. Also ran BasePythonRunnerSuite, DirectByteBufferOutputStreamSuite, PythonRDDSuite, and the PySpark tests pyspark.sql.tests.arrow.test_arrow_python_udf, pyspark.sql.tests.pandas.test_pandas_udf_scalar, pyspark.sql.tests.arrow.test_arrow_map and pyspark.sql.tests.pandas.test_pandas_grouped_map, on JDK 17.

Backport plan: branch-4.3, branch-4.2, branch-4.1 and branch-4.0 are affected and will get backports after this is merged.

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

Generated-by: Claude Code (Claude Opus 5.5)

This pull request and its description were written by Isaac.

[dongjoon-hyun]

Thank you for the fix, @viirya. I reviewed b3c7fce and left 9 inline comments. The shield itself looks right to me: with it, the thread that runs the channel's interruptor only reaches no-op flush() calls and never touches the buffer, on JDK 17, 21 and 25. Here is a summary, ordered by priority.

Coverage

1. PythonRunner L1066: PipelinedWriterRunnable (L1318) still hands the Arrow writers an unshielded stream, so how is branch-4.3 covered if #59332 doesn't go there?

Buffer lifetime

2. PythonRunner L1065: on the kill path the buffer is now freed only by GC, because L1269-L1273 never runs, and this comment and the description describe only normal completion.

Tests

3. BasePythonRunnerSuite L227: deleting the bulk write(b, off, len) override keeps both tests green.
4. BasePythonRunnerSuite L254: open is stubbed, so reverting PythonRunner.scala L1076 also keeps both tests green.
5. BasePythonRunnerSuite L237: the runner test doesn't port to branch-4.2/4.1/4.0 and conflicts on branch-4.3.
6. BasePythonRunnerSuite L313: the null clue here and the missing clue at L303 hide the failure cause.
7. BasePythonRunnerSuite L314: the finally doesn't release the spin at L260.

Minor

8. BasePythonRunnerSuite L249: TestRunner could be the existing PythonRunner.
9. CloseShieldOutputStream L34: the out constructor parameter shadows FilterOutputStream.out.

PR description

10. "The result can be a JVM crash, silent memory corruption, or an NPE / internal error instead of a clean task kill": the NPE / internal error part can't surface. The failing write runs inside the try of WritableByteChannelImpl.write, and the finally's end() waits for the interruptor's close and then throws ClosedByInterruptException, which replaces it, so the task still ends with TaskKilledException. That is also why the new test needs writeError. The real outcomes are the native ones: a JVM crash or silent memory corruption.

Items 1 and 2 matter most.

[dongjoon-hyun]

This covers the default path, but PipelinedWriterRunnable still hands the Arrow writers an unshielded new DataOutputStream(bufferStream) (L1318, outside this diff). In pipelined mode (spark.python.udf.pipelined.enabled=true), any task that completes while the pool writer is still serializing (LIMIT/show(), a UDF error, a kill) runs writerFuture.cancel(true) at L554. If the writer is inside an Arrow Channels.newChannel(dataOut) chunk at that moment, Thread.interrupt() runs the channel's interruptor on the task thread, which closes dataOut and frees the direct buffer while the writer is still in buffer.put or grow(). writerFuture.get() returns right after a successful cancel, so nothing orders the two.

#59332 shields this stream on master, but its description doesn't say whether it goes to branch-4.3, where v4.3.0-rc2 ships the same code. If #59332 is not going to branch-4.3, could this PR also use new DataOutputStream(new CloseShieldOutputStream(bufferStream)) at L1318, so that the 4.3 backport covers both paths? That line only closes the buffer use-after-free; waiting for the writer before the later listeners run would stay with #59332.

[viirya]

#59332 covers the pipelined path on branch-4.3 and branch-4.x too: its description says the pipelined mode was added in 4.3.0 by SPARK-56642, so both branches have the same issue there. So I'd keep PipelinedWriterRunnable in #59332 and leave this PR to the default path, which is the one that also needs branch-4.2 through branch-4.0.

[dongjoon-hyun]

On the kill path this PR targets, bufferStream.close() below is never reached, so the direct buffer now stays allocated until GC runs its Cleaner. L1272 only runs when !hasInput && !buffer.hasRemaining, i.e. after all input has been sent. When a task is killed with interruptThread = true during an Arrow write, the ClosedByInterruptException leaves writer.writeNextInputToStream (L1253) before L1269 and becomes a TaskKilledException. Before this PR, the interrupt close freed the buffer right away when the Arrow write channel was the first to see the interrupt: racily when the killing thread closed the channel mid-chunk (the bug fixed here), but safely when the interrupt landed outside a chunk and the next chunk's begin() closed the channel on the task thread itself. The shield now makes both closes no-ops. (Kills noticed elsewhere first already relied on GC.)

With the default 64MB spark.sql.execution.arrow.maxBytesPerBatch, this buffer can grow to about 128 MiB for wide rows (grow() doubles and reset() never shrinks it), and Platform.allocateDirectBuffer allocates outside MaxDirectMemorySize, so nothing forces that GC. Since spark.sql.execution.interruptOnCancel is on by default, each cancelled Arrow UDF task can hold one until then.

Could we free it on the thread that drives this stream when an exception escapes writeAdditionalInputToPythonWorker (clear buffer, call bufferStream.close() and rethrow), like PipelinedWriterRunnable does in its finally? A task completion listener wouldn't be safe here, because PythonForeachWriter drives this stream from its own WriterThread, which can outlive the task. Overriding ReaderInputStream.close() wouldn't be either, because an interrupted ArrowStreamReader can close its channel, and this stream with it, on the killing thread. Either way, "It is freed only by bufferStream.close() below" here, and "The buffer is still freed exactly once, by ReaderInputStream's own bufferStream.close() after the last input is sent" in the description, describe only normal completion.

[viirya]

Good catch. writeAdditionalInputToPythonWorker() now wraps the original body, renamed to writeBufferedInputToPythonWorker(). When an exception escapes it, it frees bufferStream on the thread that writes into it and sets hasInput = false and an empty buffer, so a later call does not touch the freed buffer. That covers a kill during a write as well as writer and socket write failures. The interrupting thread only reaches the no-op flush() calls, so freeing here does not race with it. I updated this comment and the description, and the ReaderInputStream test now also checks that the buffer is freed after the interrupted write.

[dongjoon-hyun]

Deleting the write(b, off, len) override (CloseShieldOutputStream.scala L37) keeps both new tests green. This test only compares the bytes, and FilterOutputStream's default loop produces the same bytes one write(int) at a time. Without the override, every bulk write into writerStream (each Arrow chunk of up to 8 KiB, pickled rows, even the 4 bytes of writeInt) would go through DirectByteBufferOutputStream.write(int), each call running checkNotClosed(), ensureCapacity() and a single-byte put. Could we count the per-byte writes on the delegate? For example:

var byteWrites = 0
val bufferStream = new DirectByteBufferOutputStream() {
  override def write(b: Int): Unit = { byteWrites += 1; super.write(b) }
}
val shielded = new CloseShieldOutputStream(bufferStream)
shielded.write(Array[Byte](1, 2, 3), 1, 2)
assert(byteWrites === 0)

[viirya]

Done. The test now counts write(int) calls on the delegate, and it fails without the override.

[dongjoon-hyun]

Since open is stubbed here, reverting PythonRunner.scala L1076 to writer.open(new DataOutputStream(bufferStream)) still passes both tests. How about override def open(dataOut: DataOutputStream): Unit = dataOut.close()? With the fix, the close only flushes. With L1076 reverted, it frees bufferStream, and buffer = bufferStream.toByteBuffer at L1077 then fails with "Cannot call methods on a closed DirectByteBufferOutputStream" while the ReaderInputStream is being constructed, with no threads involved. Related: in this test only assert(writeError === null) tells fixed from unfixed, because end(true) in the channel's finally throws ClosedByInterruptException either way, so that assertion is the one to keep if the test is reworked.

[viirya]

Done. The stub open now closes dataOut, so reverting L1076 fails while the ReaderInputStream is being constructed.

[dongjoon-hyun]

The fix itself applies cleanly to branch-4.3, 4.2, 4.1 and 4.0, but this test won't. BasePythonRunnerSuite doesn't exist on branch-4.2/4.1/4.0. PythonWorkerHandle (SPARK-58119) is 4.3+ only, so the newReaderIterator override takes pid: Option[Int] there, and ReaderInputStream takes 7, 6 (differently typed) and 4 parameters on 4.2, 4.1 and 4.0. Even on branch-4.3, the cherry-pick conflicts in the imports and at the end of this file. Since the description plans backports down to 4.0, could the new tests go into their own suite file, plus a runner-free check that ports unchanged? With the interrupt flag already set, begin() closes the channel and its stream on the writing thread before the write, so this fails at size() without the shield and passes with it on JDK 17, 21 and 25:

test("SPARK-60305: an interrupted channel write does not close the shielded buffer") {
  val buffer = new DirectByteBufferOutputStream()
  val channel = Channels.newChannel(new DataOutputStream(new CloseShieldOutputStream(buffer)))
  Thread.currentThread().interrupt()
  try {
    intercept[ClosedByInterruptException](channel.write(ByteBuffer.wrap(new Array[Byte](16))))
  } finally {
    Thread.interrupted()
  }
  assert(buffer.size() === 16)
  buffer.close()
}

It doesn't replace this test, since it covers neither the cross-thread close nor the ReaderInputStream wiring.

[viirya]

Moved the tests to a new CloseShieldOutputStreamSuite and added your runner-free test. The first two tests port unchanged; the ReaderInputStream test only needs its constructor call adjusted on each branch.

[dongjoon-hyun]

ScalaTest's assert(cond, clue) rejects a null clue before it evaluates cond, since macroAssert calls requireNonNull(clue) first. So if input.read ever returns normally here, this fails with NullArgumentException: clue was null instead of saying that no exception was thrown. The clue also adds nothing when error is non-null, because the isInstanceOf macro already prints the left operand. Could we drop it, or use a null-safe string such as s"expected ClosedByInterruptException but got: $error"? Similarly, if the task thread fails before it reaches the write, assert(inWrite.await(1, TimeUnit.MINUTES)) at L303 waits a minute and reports only that the await returned false; s"task thread failed before the write: $error" as its clue would show the cause.

[viirya]

Dropped the clue there, and the inWrite.await assertion now prints error.

[dongjoon-hyun]

resume.set(true) is reached only on the happy path, and nothing in this finally releases the spin at L260. That loop has no deadline and ignores interrupts. If inWrite.await at L303 ever times out after the writer has already passed the selector checks, the writer spins a core until the forked test JVM exits (worker.stop() only stops a writer that hasn't reached them yet). It is unlikely and only happens on a run that has already failed, but resume.set(true) as the first statement of this finally (it is idempotent) would make the writer always terminate.

[viirya]

Done. resume.set(true) is now the first statement of the finally.

[dongjoon-hyun]

nit: PythonRunner is already a concrete BasePythonRunner with exactly these constructor arguments (PythonRunner.scala L1403-L1406). new PythonRunner(Seq(ChainedPythonFunctions(Seq(func))), None) plus new runner.Writer(SparkEnv.get, worker, Iterator.empty, 0, null) { ... } and new runner.ReaderInputStream(worker, writer, None, 0, false, null) could replace TestRunner, its never-called newReaderIterator override, the newReaderInputStream helper and the TaskContext import, the same way #59332 builds its runner test in this suite. It would also ease the backports: the dead override's handle: Option[PythonWorkerHandle] is pid: Option[Int] on branch-4.2/4.1/4.0, while the PythonRunner and Writer constructors are the same on all four branches.

[viirya]

Done. The test now uses PythonRunner with new runner.Writer(...) and new runner.ReaderInputStream(...).

[dongjoon-hyun]

nit: since write(b, off, len) refers to out, Scala resolves it to this constructor parameter and keeps it as a second field next to the inherited FilterOutputStream.out. Both point to the same object today, but L37 then doesn't use the field that write(int) and flush() use. Naming the parameter differently, e.g. class CloseShieldOutputStream(underlying: OutputStream) extends FilterOutputStream(underlying), makes out at L37 the inherited field and drops the extra one.

[viirya]

Renamed the parameter to underlying.

[viirya]

Thank you for the detailed review, @dongjoon-hyun. I addressed the inline comments in a13f656. On the description: you're right that the NPE / internal error can't surface, since end() replaces it with ClosedByInterruptException. I changed it to list only the native outcomes, a JVM crash or silent memory corruption, and updated how the buffer is freed.

[dongjoon-hyun]

Thank you for addressing items 1-10, @viirya. I reviewed the latest commit, a13f656, and left 4 inline comments, continuing the numbering from items 1-10. The new catch frees the buffer on the thread that writes into it, as I asked in item 2, and the new suite now fails when the bulk-write override, either shield or the free in the catch is reverted. Here is a summary, ordered by priority.

Buffer lifetime

11. PythonRunner L1246: the empty buffer, not hasInput = false, keeps a later call off the freed memory; hasInput = false would turn a later read() into an unbounded wait, and neither line is tested.

Coordination with #59332

12. CloseShieldOutputStream L34: this copy and #59332's are no longer the same class, so if #59332 lands first, the underlying rename is lost and this PR's backports lose the file.

Minor

13. PythonRunner L1074: the scaladoc above (L1072) still points to writeAdditionalInputToPythonWorker() for the sizing details, which moved to writeBufferedInputToPythonWorker().
14. CloseShieldOutputStreamSuite L134: the reflective check can be a write through the writer's own dataOut.

PR description

15. "Without the fix, the write fails with Cannot call methods on a closed DirectByteBufferOutputStream, because the interrupting thread already freed the buffer": since the stub open now closes dataOut, reverting the whole fix fails earlier, at buffer = bufferStream.toByteBuffer (PythonRunner L1078), while the test thread is still constructing the ReaderInputStream and before any write or interrupt. The sentence holds only when the shield at L1266 alone is reverted. How about "Without the shield in writeBufferedInputToPythonWorker(), the write fails with ..., because the interrupting thread already freed the buffer; without the one on the open path, constructing the ReaderInputStream already fails with the same error, because open closes its stream."?

Items 11 and 12 matter most.

[dongjoon-hyun]

buffer = ByteBuffer.allocate(0) on the next line, not hasInput = false, is what keeps a later call off the freed memory, and if read() were ever called again after this rethrow, hasInput = false would turn a fail-fast into an unbounded wait. Without the empty buffer, a later call would hit a null buffer (set at L1264) after a writer failure, or pass the freed memory to the socket after a failed worker.channel.write(buffer) (L1281). Every bufferStream method that call could reach (reset(), size(), toByteBuffer, write) already fails in checkNotClosed(). So without hasInput = false, a later read() fails fast at bufferStream.reset() (L1262) with "Cannot call methods on a closed DirectByteBufferOutputStream". With it, the call takes the end-of-input branch (L1286-L1289): it switches the key to OP_READ as if the end of input had been sent, which it never was, and waits in select for output from a worker that is itself waiting for its next input. With the default spark.python.worker.idleTimeoutSeconds=0, that lasts until the task is killed. Nothing reads again after this exception today (every handleException case throws, and PythonForeachWriter's WriterThread calls hasNext once), so this is latent, but the comment above and the PR description ("the stream stops sending input, so a later call does not touch the freed buffer") credit the safety to the wrong line.

Neither line is covered either: test 3 never calls read() again, so deleting L1246 or L1247 keeps all three tests green. Could we drop hasInput = false and end test 3 with a read that must fail? Closing the peer first makes a regression fail instead of hang:

      // A later read fails instead of reusing the freed buffer.
      accepted.close()
      val e2 = intercept[SparkException](input.read(new Array[Byte](1)))
      assert(e2.getMessage.contains("closed DirectByteBufferOutputStream"))

Without buffer = ByteBuffer.allocate(0) this read throws a NullPointerException, and with hasInput = false back it returns -1, so either change fails the test. (If you'd rather keep hasInput = false, assert(input.read(new Array[Byte](1)) === -1) after accepted.close() covers both lines instead.)

[viirya]

You're right, the empty buffer is the line that matters. I removed hasInput = false, so a later call now fails fast at bufferStream.reset(), and reworded the comment. The ReaderInputStream test now ends with your later read after closing the peer, which fails with a NullPointerException without the empty buffer, returns -1 with hasInput = false back, and returns -1 without the free.

[dongjoon-hyun]

This copy and #59332's are no longer the same class. Since a13f656, this one takes underlying: OutputStream (my item 9), while #59332's copy (L33 at 820c7ac) still declares class CloseShieldOutputStream(out: OutputStream) extends FilterOutputStream(out), with a different scaladoc and without the one-byte comment above write. If #59332 merges first and this PR "drops its copy on rebase", as the description plans, the rename is lost. This PR's commit would then also no longer add the file, so its cherry-picks to branch-4.2, 4.1 and 4.0, which #59332 doesn't reach, wouldn't compile unless the file is brought along. Could #59332 take this file verbatim? Then either PR can land second cleanly, and "the same class" here and "Both changes add the same CloseShieldOutputStream" in #59332's description are true again.

[viirya]

Agreed, #59332 will take this file verbatim, so the two copies are the same class again. One caveat for the backports: if #59332 lands first, this PR's squashed commit no longer adds the file, so the backports to branch-4.2, 4.1 and 4.0 have to bring it along. I noted that in the description.

[dongjoon-hyun]

nit: after this round's rename, the scaladoc of this field still says "See writeAdditionalInputToPythonWorker() for details." (L1072), but that name is now the try/catch wrapper. The sizing loop it refers to ("Try not to grow the buffer much beyond bufferSize", L1267) moved to writeBufferedInputToPythonWorker(). Could it point there instead? The other uses of the old name (the writerStream comment above and PythonForeachWriter) are still accurate, since they refer to the entry point that read() calls.

[viirya]

Fixed, it now points to writeBufferedInputToPythonWorker().

[dongjoon-hyun]

nit: instead of reflection, this can check the free through the stream the runner hands the writer. DataOutputStream.write(int) goes through the shield to DirectByteBufferOutputStream.write(int), which runs checkNotClosed(), and neither DataOutputStream nor FilterOutputStream checks its own closed flag, so it doesn't matter that the interrupt closed that DataOutputStream. The write throws only if the catch freed the buffer, and it also shows that the writer's own stream now fails instead of writing into freed memory:

    @volatile var writeError: Throwable = null
    @volatile var inputOut: DataOutputStream = null
...
        override def writeNextInputToStream(dataOut: DataOutputStream): Boolean = {
          inputOut = dataOut
...
      // The failed write frees the buffer on the task thread instead of leaving it to GC.
      val e = intercept[SparkException](inputOut.write(0))
      assert(e.getMessage.contains("closed DirectByteBufferOutputStream"))

[viirya]

Removed the reflection. The later read from item 11 already covers it: it fails with the closed error only if the catch freed the buffer, and returns -1 without the free, so I didn't add a separate write through dataOut.

[viirya]

Thanks, @dongjoon-hyun. I addressed items 11-14 in bc2638a, and changed the testing paragraph of the description as you suggested in item 15, plus the sentence on how a later call is kept off the freed buffer.

[dongjoon-hyun]

+1, LGTM (Pending CIs). Thank you for addressing all the comments, @viirya.