From b3c7fceae9ab41da5c44b62f39b2b39961cfdefe Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Fri, 16 Oct 2026 14:54:16 -0700 Subject: [PATCH 1/3] [SPARK-60305][PYTHON] Keep the Python runner's input buffer alive when an Arrow write is interrupted ReaderInputStream serializes the task input into a DirectByteBufferOutputStream. The Arrow inputs write it through Channels.newChannel(dataOut), an interruptible channel, so interrupting the task thread inside a channel write (e.g. Task.kill(interruptThread = true)) closes dataOut, and so frees the direct buffer, on the interrupting thread while the task thread may still be writing into it. Hand the writer a CloseShieldOutputStream over the buffer, whose close() only flushes. The buffer is still freed only by ReaderInputStream itself, and the writer still stops with ClosedByInterruptException. Co-authored-by: Claude Code --- .../api/python/CloseShieldOutputStream.scala | 40 +++++++ .../spark/api/python/PythonRunner.scala | 9 +- .../api/python/BasePythonRunnerSuite.scala | 112 +++++++++++++++++- 3 files changed, 156 insertions(+), 5 deletions(-) create mode 100644 core/src/main/scala/org/apache/spark/api/python/CloseShieldOutputStream.scala diff --git a/core/src/main/scala/org/apache/spark/api/python/CloseShieldOutputStream.scala b/core/src/main/scala/org/apache/spark/api/python/CloseShieldOutputStream.scala new file mode 100644 index 0000000000000..d9ab4b4de2a84 --- /dev/null +++ b/core/src/main/scala/org/apache/spark/api/python/CloseShieldOutputStream.scala @@ -0,0 +1,40 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.api.python + +import java.io.{FilterOutputStream, OutputStream} + +/** + * An output stream that forwards writes to `out` but does not close it on [[close()]]. + * + * Python runners serialize their input into a `DirectByteBufferOutputStream`, whose `close()` + * frees its native buffer. Arrow writers write that stream through `Channels.newChannel(...)`, + * an interruptible channel: when another thread interrupts the writing thread inside a channel + * write (e.g. a task kill with `interruptThread = true`), `Thread.interrupt()` closes the + * channel, and so the wrapped stream, on the interrupting thread while the writing thread may + * still be copying into the buffer. Handing the writer this stream keeps the buffer alive; the + * writing thread still stops with a `ClosedByInterruptException`, and the buffer is freed only + * by its owner. + */ +private[python] class CloseShieldOutputStream(out: OutputStream) extends FilterOutputStream(out) { + + // FilterOutputStream writes an array one byte at a time. + override def write(b: Array[Byte], off: Int, len: Int): Unit = out.write(b, off, len) + + override def close(): Unit = flush() +} diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala index 0f58098d79b3e..4144720817437 100644 --- a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala +++ b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala @@ -1059,6 +1059,11 @@ private[spark] abstract class BasePythonRunner[IN, OUT]( private[this] var writerIfbhThreadLocalValue: Object = null private[this] val temp = new Array[Byte](1) private[this] val bufferStream = new DirectByteBufferOutputStream() + // The writer gets `bufferStream` only through this stream. Arrow writers write through an + // interruptible channel, so a task kill that interrupts this thread closes the stream they + // write to on the killing thread; that must not free `bufferStream` under the writer. It is + // freed only by `bufferStream.close()` below. See `CloseShieldOutputStream`. + private[this] val writerStream = new CloseShieldOutputStream(bufferStream) /** * Buffers data to be written to the Python worker until the socket is * available for write. @@ -1068,7 +1073,7 @@ private[spark] abstract class BasePythonRunner[IN, OUT]( private[this] var buffer: ByteBuffer = _ private[this] var hasInput = true - writer.open(new DataOutputStream(bufferStream)) + writer.open(new DataOutputStream(writerStream)) buffer = bufferStream.toByteBuffer override def read(): Int = { @@ -1241,7 +1246,7 @@ private[spark] abstract class BasePythonRunner[IN, OUT]( // Set the `buffer` to null to make it eligible for GC buffer = null - val dataOut = new DataOutputStream(bufferStream) + val dataOut = new DataOutputStream(writerStream) // Try not to grow the buffer much beyond `bufferSize`. This is inevitable for large // input rows. while (bufferStream.size() < bufferSize && hasInput && !shouldFlush()) { diff --git a/core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala b/core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala index f7675880e01f0..b838786f0b887 100644 --- a/core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala +++ b/core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala @@ -17,13 +17,21 @@ package org.apache.spark.api.python -import java.io.{ByteArrayInputStream, ByteArrayOutputStream, DataInputStream, DataOutputStream, EOFException} +import java.io.{ByteArrayInputStream, ByteArrayOutputStream, DataInputStream, DataOutputStream, + EOFException, FilterOutputStream} +import java.net.{InetAddress, InetSocketAddress} +import java.nio.ByteBuffer +import java.nio.channels.{Channels, ClosedByInterruptException, ServerSocketChannel, SocketChannel} import java.nio.charset.StandardCharsets +import java.util.{ArrayList => JArrayList, HashMap => JHashMap} import java.util.concurrent.{CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.AtomicBoolean -import org.apache.spark.{SparkException, SparkFunSuite} +import org.apache.spark.{LocalSparkContext, SparkContext, SparkEnv, SparkException, + SparkFunSuite, TaskContext} +import org.apache.spark.util.DirectByteBufferOutputStream -class BasePythonRunnerSuite extends SparkFunSuite { +class BasePythonRunnerSuite extends SparkFunSuite with LocalSparkContext { test("SPARK-58192: pipelined writer pool admits more concurrent writers than processors") { // A fractional spark.task.cpus admits more concurrent tasks -- each holding one writer @@ -211,4 +219,102 @@ class BasePythonRunnerSuite extends SparkFunSuite { assert(BasePythonRunner.readWorkerMetrics(stream) == expectedMetrics) checkEndMarkers(stream) } + + test("SPARK-60305: CloseShieldOutputStream forwards writes and does not close the " + + "underlying stream") { + val bufferStream = new DirectByteBufferOutputStream() + val shielded = new CloseShieldOutputStream(bufferStream) + shielded.write(Array[Byte](1, 2, 3), 1, 2) + shielded.write(4) + shielded.close() + assert(bufferStream.size() === 3) + val bytes = new Array[Byte](3) + bufferStream.toByteBuffer.get(bytes) + assert(bytes.toSeq === Seq[Byte](2, 3, 4)) + bufferStream.close() + } + + test("SPARK-60305: interrupting the task thread inside a channel write keeps the input " + + "buffer alive") { + sc = new SparkContext("local", "test") + val func = SimplePythonFunction(Seq.empty[Byte], new JHashMap[String, String](), + new JArrayList[String](), "python3", "3", new JArrayList(), null) + val inWrite = new CountDownLatch(1) + val resume = new AtomicBoolean(false) + @volatile var writeError: Throwable = null + + // Writes the way the Arrow inputs do: through `Channels.newChannel(dataOut)`, whose writes + // the JVM closes, together with `dataOut`, when the writing thread is interrupted. The + // write blocks until the test has interrupted the thread, then forwards to `dataOut`. + class TestRunner extends BasePythonRunner[Int, Int](Seq(ChainedPythonFunctions(Seq(func))), + PythonEvalType.NON_UDF, Array(Array(0)), None, Map.empty) { + override protected def newWriter(env: SparkEnv, worker: PythonWorker, + inputIterator: Iterator[Int], partitionIndex: Int, context: TaskContext): Writer = { + new Writer(env, worker, inputIterator, partitionIndex, context) { + override def open(dataOut: DataOutputStream): Unit = {} + override protected def writeCommand(dataOut: DataOutputStream): Unit = {} + override def writeNextInputToStream(dataOut: DataOutputStream): Boolean = { + val channel = Channels.newChannel(new FilterOutputStream(dataOut) { + override def write(b: Array[Byte], off: Int, len: Int): Unit = { + inWrite.countDown() + while (!resume.get()) Thread.onSpinWait() + try { + out.write(b, off, len) + } catch { + case t: Throwable => + writeError = t + throw t + } + } + }) + channel.write(ByteBuffer.wrap(new Array[Byte](16))) + false + } + } + } + + override protected def newReaderIterator(stream: DataInputStream, writer: Writer, + startTime: Long, env: SparkEnv, worker: PythonWorker, + handle: Option[PythonWorkerHandle], releasedOrClosed: AtomicBoolean, + context: TaskContext): Iterator[Int] = throw new UnsupportedOperationException() + + def newReaderInputStream(worker: PythonWorker): ReaderInputStream = new ReaderInputStream( + worker, newWriter(SparkEnv.get, worker, Iterator.empty, 0, null), None, 0, false, null) + } + val runner = new TestRunner + + val server = ServerSocketChannel.open() + server.bind(new InetSocketAddress(InetAddress.getLoopbackAddress, 0)) + val client = SocketChannel.open(server.getLocalAddress) + val accepted = server.accept() + client.configureBlocking(false) + val worker = PythonWorker(client).refresh() + try { + val input = runner.newReaderInputStream(worker) + @volatile var error: Throwable = null + val taskThread = new Thread(() => { + try { + input.read(new Array[Byte](1)) + } catch { + case t: Throwable => error = t + } + }) + taskThread.start() + assert(inWrite.await(1, TimeUnit.MINUTES)) + // Like Task.kill(interruptThread = true). This closes the channel and `dataOut` on this + // thread while the task thread is still inside the write. + taskThread.interrupt() + resume.set(true) + taskThread.join(TimeUnit.MINUTES.toMillis(1)) + assert(!taskThread.isAlive) + // The write into the buffer still succeeds: the buffer was not freed under the writer. + // The interrupt then surfaces as the channel's ClosedByInterruptException. + assert(writeError === null) + assert(error.isInstanceOf[ClosedByInterruptException], error) + } finally { + worker.stop() + accepted.close() + server.close() + } + } } From a13f6569487afa73513c7e8da12a90cd3b210e11 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Fri, 16 Oct 2026 22:48:42 -0700 Subject: [PATCH 2/3] Free the input buffer when writing fails and move the tests to their own suite On a kill during a write, the shield also turns the interrupt close into a no-op, so the buffer was left to GC. writeAdditionalInputToPythonWorker() now frees it on the writing thread when an exception escapes, and stops sending input so that a later call does not touch it. Move the tests to CloseShieldOutputStreamSuite so that they port to older branches, add a runner-free interrupt test, cover the open path and bulk writes, and use PythonRunner instead of a test runner. Rename the CloseShieldOutputStream parameter so it does not shadow FilterOutputStream.out. Co-authored-by: Claude Code --- .../api/python/CloseShieldOutputStream.scala | 5 +- .../spark/api/python/PythonRunner.scala | 21 ++- .../api/python/BasePythonRunnerSuite.scala | 112 +------------ .../python/CloseShieldOutputStreamSuite.scala | 147 ++++++++++++++++++ 4 files changed, 172 insertions(+), 113 deletions(-) create mode 100644 core/src/test/scala/org/apache/spark/api/python/CloseShieldOutputStreamSuite.scala diff --git a/core/src/main/scala/org/apache/spark/api/python/CloseShieldOutputStream.scala b/core/src/main/scala/org/apache/spark/api/python/CloseShieldOutputStream.scala index d9ab4b4de2a84..33ccef7ab0b15 100644 --- a/core/src/main/scala/org/apache/spark/api/python/CloseShieldOutputStream.scala +++ b/core/src/main/scala/org/apache/spark/api/python/CloseShieldOutputStream.scala @@ -20,7 +20,7 @@ package org.apache.spark.api.python import java.io.{FilterOutputStream, OutputStream} /** - * An output stream that forwards writes to `out` but does not close it on [[close()]]. + * An output stream that forwards writes to `underlying` but does not close it on [[close()]]. * * Python runners serialize their input into a `DirectByteBufferOutputStream`, whose `close()` * frees its native buffer. Arrow writers write that stream through `Channels.newChannel(...)`, @@ -31,7 +31,8 @@ import java.io.{FilterOutputStream, OutputStream} * writing thread still stops with a `ClosedByInterruptException`, and the buffer is freed only * by its owner. */ -private[python] class CloseShieldOutputStream(out: OutputStream) extends FilterOutputStream(out) { +private[python] class CloseShieldOutputStream(underlying: OutputStream) + extends FilterOutputStream(underlying) { // FilterOutputStream writes an array one byte at a time. override def write(b: Array[Byte], off: Int, len: Int): Unit = out.write(b, off, len) diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala index 4144720817437..754f6a658c368 100644 --- a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala +++ b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala @@ -1062,7 +1062,8 @@ private[spark] abstract class BasePythonRunner[IN, OUT]( // The writer gets `bufferStream` only through this stream. Arrow writers write through an // interruptible channel, so a task kill that interrupts this thread closes the stream they // write to on the killing thread; that must not free `bufferStream` under the writer. It is - // freed only by `bufferStream.close()` below. See `CloseShieldOutputStream`. + // freed only on this thread, by `writeAdditionalInputToPythonWorker()`, once all input is + // sent or when writing fails. See `CloseShieldOutputStream`. private[this] val writerStream = new CloseShieldOutputStream(bufferStream) /** * Buffers data to be written to the Python worker until the socket is @@ -1233,11 +1234,27 @@ private[spark] abstract class BasePythonRunner[IN, OUT]( } } + private def writeAdditionalInputToPythonWorker(): Unit = { + try { + writeBufferedInputToPythonWorker() + } catch { + case t: Throwable => + // Writing failed or the task was killed mid-write, so the close at the end of + // `writeBufferedInputToPythonWorker()` is never reached. Free the buffer here, on the + // thread that writes into it, instead of leaving it to GC, and stop sending input so + // that a later call does not touch it. + hasInput = false + buffer = ByteBuffer.allocate(0) + bufferStream.close() + throw t + } + } + /** * Reads input data from `writer.inputIterator` into `buffer` and writes the buffer to the * Python worker if the socket is available for writing. */ - private def writeAdditionalInputToPythonWorker(): Unit = { + private def writeBufferedInputToPythonWorker(): Unit = { var acceptsInput = true while (acceptsInput && (hasInput || buffer.hasRemaining)) { if (!buffer.hasRemaining && hasInput) { diff --git a/core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala b/core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala index b838786f0b887..f7675880e01f0 100644 --- a/core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala +++ b/core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala @@ -17,21 +17,13 @@ package org.apache.spark.api.python -import java.io.{ByteArrayInputStream, ByteArrayOutputStream, DataInputStream, DataOutputStream, - EOFException, FilterOutputStream} -import java.net.{InetAddress, InetSocketAddress} -import java.nio.ByteBuffer -import java.nio.channels.{Channels, ClosedByInterruptException, ServerSocketChannel, SocketChannel} +import java.io.{ByteArrayInputStream, ByteArrayOutputStream, DataInputStream, DataOutputStream, EOFException} import java.nio.charset.StandardCharsets -import java.util.{ArrayList => JArrayList, HashMap => JHashMap} import java.util.concurrent.{CountDownLatch, TimeUnit} -import java.util.concurrent.atomic.AtomicBoolean -import org.apache.spark.{LocalSparkContext, SparkContext, SparkEnv, SparkException, - SparkFunSuite, TaskContext} -import org.apache.spark.util.DirectByteBufferOutputStream +import org.apache.spark.{SparkException, SparkFunSuite} -class BasePythonRunnerSuite extends SparkFunSuite with LocalSparkContext { +class BasePythonRunnerSuite extends SparkFunSuite { test("SPARK-58192: pipelined writer pool admits more concurrent writers than processors") { // A fractional spark.task.cpus admits more concurrent tasks -- each holding one writer @@ -219,102 +211,4 @@ class BasePythonRunnerSuite extends SparkFunSuite with LocalSparkContext { assert(BasePythonRunner.readWorkerMetrics(stream) == expectedMetrics) checkEndMarkers(stream) } - - test("SPARK-60305: CloseShieldOutputStream forwards writes and does not close the " + - "underlying stream") { - val bufferStream = new DirectByteBufferOutputStream() - val shielded = new CloseShieldOutputStream(bufferStream) - shielded.write(Array[Byte](1, 2, 3), 1, 2) - shielded.write(4) - shielded.close() - assert(bufferStream.size() === 3) - val bytes = new Array[Byte](3) - bufferStream.toByteBuffer.get(bytes) - assert(bytes.toSeq === Seq[Byte](2, 3, 4)) - bufferStream.close() - } - - test("SPARK-60305: interrupting the task thread inside a channel write keeps the input " + - "buffer alive") { - sc = new SparkContext("local", "test") - val func = SimplePythonFunction(Seq.empty[Byte], new JHashMap[String, String](), - new JArrayList[String](), "python3", "3", new JArrayList(), null) - val inWrite = new CountDownLatch(1) - val resume = new AtomicBoolean(false) - @volatile var writeError: Throwable = null - - // Writes the way the Arrow inputs do: through `Channels.newChannel(dataOut)`, whose writes - // the JVM closes, together with `dataOut`, when the writing thread is interrupted. The - // write blocks until the test has interrupted the thread, then forwards to `dataOut`. - class TestRunner extends BasePythonRunner[Int, Int](Seq(ChainedPythonFunctions(Seq(func))), - PythonEvalType.NON_UDF, Array(Array(0)), None, Map.empty) { - override protected def newWriter(env: SparkEnv, worker: PythonWorker, - inputIterator: Iterator[Int], partitionIndex: Int, context: TaskContext): Writer = { - new Writer(env, worker, inputIterator, partitionIndex, context) { - override def open(dataOut: DataOutputStream): Unit = {} - override protected def writeCommand(dataOut: DataOutputStream): Unit = {} - override def writeNextInputToStream(dataOut: DataOutputStream): Boolean = { - val channel = Channels.newChannel(new FilterOutputStream(dataOut) { - override def write(b: Array[Byte], off: Int, len: Int): Unit = { - inWrite.countDown() - while (!resume.get()) Thread.onSpinWait() - try { - out.write(b, off, len) - } catch { - case t: Throwable => - writeError = t - throw t - } - } - }) - channel.write(ByteBuffer.wrap(new Array[Byte](16))) - false - } - } - } - - override protected def newReaderIterator(stream: DataInputStream, writer: Writer, - startTime: Long, env: SparkEnv, worker: PythonWorker, - handle: Option[PythonWorkerHandle], releasedOrClosed: AtomicBoolean, - context: TaskContext): Iterator[Int] = throw new UnsupportedOperationException() - - def newReaderInputStream(worker: PythonWorker): ReaderInputStream = new ReaderInputStream( - worker, newWriter(SparkEnv.get, worker, Iterator.empty, 0, null), None, 0, false, null) - } - val runner = new TestRunner - - val server = ServerSocketChannel.open() - server.bind(new InetSocketAddress(InetAddress.getLoopbackAddress, 0)) - val client = SocketChannel.open(server.getLocalAddress) - val accepted = server.accept() - client.configureBlocking(false) - val worker = PythonWorker(client).refresh() - try { - val input = runner.newReaderInputStream(worker) - @volatile var error: Throwable = null - val taskThread = new Thread(() => { - try { - input.read(new Array[Byte](1)) - } catch { - case t: Throwable => error = t - } - }) - taskThread.start() - assert(inWrite.await(1, TimeUnit.MINUTES)) - // Like Task.kill(interruptThread = true). This closes the channel and `dataOut` on this - // thread while the task thread is still inside the write. - taskThread.interrupt() - resume.set(true) - taskThread.join(TimeUnit.MINUTES.toMillis(1)) - assert(!taskThread.isAlive) - // The write into the buffer still succeeds: the buffer was not freed under the writer. - // The interrupt then surfaces as the channel's ClosedByInterruptException. - assert(writeError === null) - assert(error.isInstanceOf[ClosedByInterruptException], error) - } finally { - worker.stop() - accepted.close() - server.close() - } - } } diff --git a/core/src/test/scala/org/apache/spark/api/python/CloseShieldOutputStreamSuite.scala b/core/src/test/scala/org/apache/spark/api/python/CloseShieldOutputStreamSuite.scala new file mode 100644 index 0000000000000..2200fbe8ec8cd --- /dev/null +++ b/core/src/test/scala/org/apache/spark/api/python/CloseShieldOutputStreamSuite.scala @@ -0,0 +1,147 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.apache.spark.api.python + +import java.io.{DataOutputStream, FilterOutputStream} +import java.net.{InetAddress, InetSocketAddress} +import java.nio.ByteBuffer +import java.nio.channels.{Channels, ClosedByInterruptException, ServerSocketChannel, SocketChannel} +import java.util.{ArrayList => JArrayList, HashMap => JHashMap} +import java.util.concurrent.{CountDownLatch, TimeUnit} +import java.util.concurrent.atomic.AtomicBoolean + +import org.apache.spark.{LocalSparkContext, SparkContext, SparkEnv, SparkException, SparkFunSuite} +import org.apache.spark.util.DirectByteBufferOutputStream + +class CloseShieldOutputStreamSuite extends SparkFunSuite with LocalSparkContext { + + test("SPARK-60305: forwards bulk writes and does not close the underlying stream") { + 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) + shielded.write(4) + assert(byteWrites === 1) + shielded.close() + assert(bufferStream.size() === 3) + val bytes = new Array[Byte](3) + bufferStream.toByteBuffer.get(bytes) + assert(bytes.toSeq === Seq[Byte](2, 3, 4)) + bufferStream.close() + } + + test("SPARK-60305: an interrupted channel write does not close the shielded buffer") { + val bufferStream = new DirectByteBufferOutputStream() + val channel = + Channels.newChannel(new DataOutputStream(new CloseShieldOutputStream(bufferStream))) + // With the interrupt flag already set, the channel closes itself, and the stream it writes + // to, on this thread before it writes. + Thread.currentThread().interrupt() + try { + intercept[ClosedByInterruptException](channel.write(ByteBuffer.wrap(new Array[Byte](16)))) + } finally { + Thread.interrupted() + } + assert(bufferStream.size() === 16) + bufferStream.close() + } + + test("SPARK-60305: interrupting the task thread inside a channel write keeps the " + + "ReaderInputStream buffer alive") { + sc = new SparkContext("local", "test") + val func = SimplePythonFunction(Seq.empty[Byte], new JHashMap[String, String](), + new JArrayList[String](), "python3", "3", new JArrayList(), null) + val runner = new PythonRunner(Seq(ChainedPythonFunctions(Seq(func))), None) + val inWrite = new CountDownLatch(1) + val resume = new AtomicBoolean(false) + @volatile var writeError: Throwable = null + + val server = ServerSocketChannel.open() + server.bind(new InetSocketAddress(InetAddress.getLoopbackAddress, 0)) + val client = SocketChannel.open(server.getLocalAddress) + val accepted = server.accept() + client.configureBlocking(false) + val worker = PythonWorker(client).refresh() + try { + val writer = new runner.Writer(SparkEnv.get, worker, Iterator.empty, 0, null) { + // An interrupt closes the stream the writer writes to; with the shield, this only + // flushes. + override def open(dataOut: DataOutputStream): Unit = dataOut.close() + override protected def writeCommand(dataOut: DataOutputStream): Unit = {} + // Writes the way the Arrow inputs do: through `Channels.newChannel(dataOut)`, which the + // JVM closes, together with `dataOut`, when the writing thread is interrupted. The + // write blocks until the test has interrupted the thread, then forwards to `dataOut`. + override def writeNextInputToStream(dataOut: DataOutputStream): Boolean = { + val channel = Channels.newChannel(new FilterOutputStream(dataOut) { + override def write(b: Array[Byte], off: Int, len: Int): Unit = { + inWrite.countDown() + while (!resume.get()) Thread.onSpinWait() + try { + out.write(b, off, len) + } catch { + case t: Throwable => + writeError = t + throw t + } + } + }) + channel.write(ByteBuffer.wrap(new Array[Byte](16))) + false + } + } + val input = new runner.ReaderInputStream(worker, writer, None, 0, false, null) + @volatile var error: Throwable = null + val taskThread = new Thread(() => { + try { + input.read(new Array[Byte](1)) + } catch { + case t: Throwable => error = t + } + }) + taskThread.start() + assert(inWrite.await(1, TimeUnit.MINUTES), s"task thread failed before the write: $error") + // Like Task.kill(interruptThread = true). This closes the channel and `dataOut` on this + // thread while the task thread is still inside the write. + taskThread.interrupt() + resume.set(true) + taskThread.join(TimeUnit.MINUTES.toMillis(1)) + assert(!taskThread.isAlive) + // The write into the buffer still succeeds: the buffer was not freed under the writer. + assert(writeError === null) + assert(error.isInstanceOf[ClosedByInterruptException]) + // The failed write frees the buffer on the task thread instead of leaving it to GC. + val bufferField = input.getClass.getDeclaredField("bufferStream") + bufferField.setAccessible(true) + val bufferStream = bufferField.get(input).asInstanceOf[DirectByteBufferOutputStream] + val e = intercept[SparkException](bufferStream.size()) + assert(e.getMessage.contains("closed DirectByteBufferOutputStream")) + } finally { + // Never leave the writer spinning, even if the test failed before releasing it. + resume.set(true) + worker.stop() + accepted.close() + server.close() + } + } +} From bc2638aa6077c51688c8708e45662a50ef1124b0 Mon Sep 17 00:00:00 2001 From: Liang-Chi Hsieh Date: Mon, 19 Oct 2026 13:48:35 -0700 Subject: [PATCH 3/3] Keep a later read off the freed input buffer by failing fast Replacing buffer with an empty one is what keeps a later call off the freed memory; drop hasInput = false, which would instead turn a later read into an unbounded wait. Check the free and the fail-fast with a later read instead of reflection, and point the buffer scaladoc to writeBufferedInputToPythonWorker(). Co-authored-by: Claude Code --- .../org/apache/spark/api/python/PythonRunner.scala | 8 ++++---- .../api/python/CloseShieldOutputStreamSuite.scala | 10 +++++----- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala index 754f6a658c368..1031a7a00d1cf 100644 --- a/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala +++ b/core/src/main/scala/org/apache/spark/api/python/PythonRunner.scala @@ -1069,7 +1069,7 @@ private[spark] abstract class BasePythonRunner[IN, OUT]( * Buffers data to be written to the Python worker until the socket is * available for write. * A best-effort attempt is made to not grow the buffer beyond "spark.buffer.size". See - * `writeAdditionalInputToPythonWorker()` for details. + * `writeBufferedInputToPythonWorker()` for details. */ private[this] var buffer: ByteBuffer = _ private[this] var hasInput = true @@ -1241,9 +1241,9 @@ private[spark] abstract class BasePythonRunner[IN, OUT]( case t: Throwable => // Writing failed or the task was killed mid-write, so the close at the end of // `writeBufferedInputToPythonWorker()` is never reached. Free the buffer here, on the - // thread that writes into it, instead of leaving it to GC, and stop sending input so - // that a later call does not touch it. - hasInput = false + // thread that writes into it, instead of leaving it to GC. `buffer` may be null or + // still point into that memory, so replace it with an empty one: a later call then + // fails fast in `bufferStream` instead of touching the freed memory. buffer = ByteBuffer.allocate(0) bufferStream.close() throw t diff --git a/core/src/test/scala/org/apache/spark/api/python/CloseShieldOutputStreamSuite.scala b/core/src/test/scala/org/apache/spark/api/python/CloseShieldOutputStreamSuite.scala index 2200fbe8ec8cd..e34d34422557a 100644 --- a/core/src/test/scala/org/apache/spark/api/python/CloseShieldOutputStreamSuite.scala +++ b/core/src/test/scala/org/apache/spark/api/python/CloseShieldOutputStreamSuite.scala @@ -130,11 +130,11 @@ class CloseShieldOutputStreamSuite extends SparkFunSuite with LocalSparkContext // The write into the buffer still succeeds: the buffer was not freed under the writer. assert(writeError === null) assert(error.isInstanceOf[ClosedByInterruptException]) - // The failed write frees the buffer on the task thread instead of leaving it to GC. - val bufferField = input.getClass.getDeclaredField("bufferStream") - bufferField.setAccessible(true) - val bufferStream = bufferField.get(input).asInstanceOf[DirectByteBufferOutputStream] - val e = intercept[SparkException](bufferStream.size()) + // The failed write frees the buffer on the task thread instead of leaving it to GC, and + // a later read fails instead of reusing the freed buffer. Closing the peer first makes a + // regression return -1 rather than wait for output. + accepted.close() + val e = intercept[SparkException](input.read(new Array[Byte](1))) assert(e.getMessage.contains("closed DirectByteBufferOutputStream")) } finally { // Never leave the writer spinning, even if the test failed before releasing it.