viirya commented on code in PR #59565:
URL: https://github.com/apache/spark/pull/59565#discussion_r4301621653
##########
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)
Review Comment:
#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.
##########
core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala:
##########
@@ -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)
Review Comment:
Done. The test now counts `write(int)` calls on the delegate, and it fails
without the override.
##########
core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala:
##########
@@ -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 = {}
Review Comment:
Done. The stub `open` now closes `dataOut`, so reverting L1076 fails while
the `ReaderInputStream` is being constructed.
##########
core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala:
##########
@@ -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 " +
Review Comment:
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.
##########
core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala:
##########
@@ -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)
Review Comment:
Dropped the clue there, and the `inWrite.await` assertion now prints `error`.
##########
core/src/test/scala/org/apache/spark/api/python/BasePythonRunnerSuite.scala:
##########
@@ -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 {
Review Comment:
Done. `resume.set(true)` is now the first statement of the `finally`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
