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..33ccef7ab0b15 --- /dev/null +++ b/core/src/main/scala/org/apache/spark/api/python/CloseShieldOutputStream.scala @@ -0,0 +1,41 @@ +/* + * 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 `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(...)`, + * 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(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) + + 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..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 @@ -1059,16 +1059,22 @@ 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 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 * 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 - writer.open(new DataOutputStream(bufferStream)) + writer.open(new DataOutputStream(writerStream)) buffer = bufferStream.toByteBuffer override def read(): Int = { @@ -1228,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. `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 + } + } + /** * 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) { @@ -1241,7 +1263,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/CloseShieldOutputStreamSuite.scala b/core/src/test/scala/org/apache/spark/api/python/CloseShieldOutputStreamSuite.scala new file mode 100644 index 0000000000000..e34d34422557a --- /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, 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. + resume.set(true) + worker.stop() + accepted.close() + server.close() + } + } +}