dongjoon-hyun commented on code in PR #59565:
URL: https://github.com/apache/spark/pull/59565#discussion_r4300834721
##########
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:
**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:
```scala
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)
```
##########
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`.
Review Comment:
**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.
##########
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:
**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.
##########
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:
**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.
##########
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) {
Review Comment:
**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.
##########
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))),
Review Comment:
**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.
##########
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:
**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:
```scala
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.
##########
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:
**`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.
##########
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:
**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.
--
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]
