dongjoon-hyun commented on code in PR #59282:
URL: https://github.com/apache/spark/pull/59282#discussion_r4235446543
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/HybridQueue.scala:
##########
@@ -171,10 +171,17 @@ abstract class HybridQueue[T, Q <: Queue[T]](
reading.close()
reading = null.asInstanceOf[Q]
}
- synchronized {
- while (!queues.isEmpty) {
- queues.remove().close()
- }
+ // Detach the queues under the lock but close them (which frees their
pages through the
+ // TaskMemoryManager) after releasing it. `spill` is called with the
TaskMemoryManager's
+ // monitor held and then takes this queue's lock, so freeing pages while
holding this lock
+ // could deadlock with a concurrent spill of this queue triggered from
another thread.
+ val detached = synchronized {
+ val current = queues
+ queues = new java.util.LinkedList[Q]()
+ current
+ }
+ while (!detached.isEmpty) {
+ detached.remove().close()
}
}
Review Comment:
This is not introduced by this PR, but since this PR reasons about `close()`
running concurrently with a pipelined writer: `close()` still frees the
in-memory pages, including the current `writing` page, without coordinating
with the queue's own writer thread. In lockFree mode, `add()` writes straight
into that page.
The pipelined writer listener in `PythonRunner.createPipelinedDataIn` is
meant to prevent this, because it runs before the queue-close listener under
LIFO. However, it calls `writerFuture.cancel(true)` and then
`writerFuture.get()`. After a successful `cancel(true)`, `FutureTask.get()`
throws `CancellationException` immediately and does not wait for the runnable
to exit. So the writer can still be inside `queue.add()` when this loop frees
the page. That is a use-after-free, and a JVM crash for off-heap pages.
Also, a `createNewQueue` that races with the detach adds its new buffer to
the fresh `queues` list, which nobody closes, so that page leaks until
`cleanUpAllAllocatedMemory`.
I think this deserves a separate JIRA. For example, the listener could wait
on a latch that the writer counts down in a `finally`, instead of on
`Future.get()`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/RowQueue.scala:
##########
@@ -181,8 +181,13 @@ private[python] case class DiskRowQueue(
*
* HybridRowQueue could be safely appended in one thread, and pulled in
another thread in the same
* time.
+ *
+ * This is intentionally not a case class. As a memory consumer it must use
identity equality:
+ * otherwise two queues created with the same arguments in one task are equal,
so only one of
+ * them is tracked (and offered for spilling) by the [[TaskMemoryManager]],
and each refuses to
Review Comment:
The root cause is that `TaskMemoryManager` keeps spillable consumers in a
value-equality `HashSet<MemoryConsumer>` (`TaskMemoryManager.java` L119, filled
by `consumers.add(requestingConsumer)` in `acquireExecutionMemory`).
Any other `MemoryConsumer` with a value-based `equals` (a Scala case class,
or a Java class with value semantics) is still silently deduplicated there. The
second instance is never offered for spilling, and its memory is missing from
the OOM memory breakdown. Using an identity set, e.g.
`Collections.newSetFromMap(new IdentityHashMap<>())`, would protect every
consumer in one place. This may be better as a separate JIRA.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/RowQueue.scala:
##########
@@ -181,8 +181,13 @@ private[python] case class DiskRowQueue(
*
* HybridRowQueue could be safely appended in one thread, and pulled in
another thread in the same
* time.
+ *
+ * This is intentionally not a case class. As a memory consumer it must use
identity equality:
Review Comment:
This fixes the leaf class, but the generic base still compares by value:
`HybridQueue.spill` checks `if (trigger == this)` (`HybridQueue.scala` L73).
`WindowSegmentTree.spill` uses `trigger eq this` for the same purpose.
Could we fix this in `HybridQueue` instead? For example:
```scala
final override def equals(other: Any): Boolean = other match {
case r: AnyRef => this eq r
case _ => false
}
final override def hashCode(): Int = System.identityHashCode(this)
```
plus `if (trigger eq this)` in `spill`. Scala does not synthesize
`equals`/`hashCode` for a case class when a base class other than `AnyRef`
already defines them concretely. So every current and future `HybridQueue`
subclass would get identity equality, and `HybridRowQueue` could even stay a
case class without the hand-written companion `apply`.
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/python/RowQueue.scala:
##########
@@ -181,8 +181,13 @@ private[python] case class DiskRowQueue(
*
* HybridRowQueue could be safely appended in one thread, and pulled in
another thread in the same
* time.
+ *
+ * This is intentionally not a case class. As a memory consumer it must use
identity equality:
+ * otherwise two queues created with the same arguments in one task are equal,
so only one of
+ * them is tracked (and offered for spilling) by the [[TaskMemoryManager]],
and each refuses to
+ * spill when the other one triggers it.
*/
-case class HybridRowQueue(
+class HybridRowQueue(
Review Comment:
Dropping `case class` also drops the generated `toString`, and
`MemoryConsumer` does not override it. The per-consumer memory breakdown (the
OOM message sent to the driver and UI) and the `TaskMemoryManager` debug logs
now show `HybridRowQueue@1a2b3c` instead of the width and the `lockFree` flag.
The PR description says the extra queues now show up in the memory usage
breakdown, but they are hard to tell apart there. Could we keep a short
`toString`? For example:
```scala
override def toString: String =
s"HybridRowQueue(numFields=$numFields, lockFree=$lockFree)@" +
Integer.toHexString(System.identityHashCode(this))
```
This is not needed if `HybridRowQueue` stays a case class with identity
`equals`/`hashCode` in the base class.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/RowQueueSuite.scala:
##########
@@ -170,4 +173,117 @@ class RowQueueSuite extends SparkFunSuite with
EncryptionFunSuite {
queue.close()
}
}
+
+ private def addRows(queue: HybridRowQueue, n: Int): Unit = {
+ val row = new UnsafeRow(1)
+ row.pointTo(new Array[Byte](16), 16)
+ (0 until n).foreach { i =>
+ row.setLong(0, i)
+ assert(queue.add(row) === QueueMode.IN_MEMORY)
+ }
+ }
+
+ private def drainAndClose(queue: HybridRowQueue, n: Int): Unit = {
+ (0 until n).foreach { i =>
+ assert(queue.remove().getLong(0) === i)
+ }
+ queue.close()
+ }
+
+ test("hybrid queues with the same fields are distinct memory consumers") {
+ val conf = new SparkConf(false).set(BUFFER_PAGESIZE, 1024L)
+ val serManager = createSerializerManager(conf)
+ val mem = new TestMemoryManager(conf)
+ val taskM = new TaskMemoryManager(mem, 0)
+ val tempDir = Utils.createTempDir().getCanonicalFile
+ val queue1 = HybridRowQueue(taskM, tempDir, 1, serManager)
+ val queue2 = HybridRowQueue(taskM, tempDir, 1, serManager)
+ assert(queue1 != queue2)
+
+ // queue1 is registered first and holds a single page, which it never
spills because the
+ // last queue is kept for writing. queue2 holds several pages and can
release all but one.
+ addRows(queue1, 1)
+ val n = 150
+ addRows(queue2, n)
+ assert(queue2.numQueues() > 2)
+ val queue2Used = queue2.getUsed
+
+ // Memory pressure from another consumer must be able to spill queue2.
+ mem.limit(0)
+ val consumer = new TestMemoryConsumer(taskM)
+ consumer.use(1024)
+ assert(consumer.getUsed === 1024)
+ assert(queue2.getUsed === 1024)
+ assert(queue2.getUsed < queue2Used)
+ consumer.free(1024)
+
+ drainAndClose(queue1, 1)
+ drainAndClose(queue2, n)
+ assert(taskM.cleanUpAllAllocatedMemory() === 0)
+ }
+
+ test("hybrid queue spills for another queue with the same fields") {
+ val conf = new SparkConf(false).set(BUFFER_PAGESIZE, 1024L)
+ val serManager = createSerializerManager(conf)
+ val taskM = new TaskMemoryManager(new TestMemoryManager(conf), 0)
+ val tempDir = Utils.createTempDir().getCanonicalFile
+ val queue1 = HybridRowQueue(taskM, tempDir, 1, serManager)
+ val queue2 = HybridRowQueue(taskM, tempDir, 1, serManager)
+ val n = 150
+ addRows(queue1, n)
+ assert(queue1.numQueues() > 2)
+ val used = queue1.getUsed
+ assert(queue1.spill(Long.MaxValue, queue2) === used - 1024)
+ assert(queue1.spill(Long.MaxValue, queue1) === 0)
+ drainAndClose(queue1, n)
+ queue2.close()
+ }
+
+ test("closing a hybrid queue does not deadlock with a concurrent spill of
it") {
+ val conf = new SparkConf(false).set(BUFFER_PAGESIZE, 1024L)
+ val serManager = createSerializerManager(conf)
+ val closeFreeingPage = new CountDownLatch(1)
+ val spillHoldingTaskMemoryManager = new CountDownLatch(1)
+ @volatile var closingThread: Thread = null
+ val taskM = new TaskMemoryManager(new TestMemoryManager(conf), 0) {
+ override def freePage(page: MemoryBlock, consumer: MemoryConsumer): Unit
= {
+ if (Thread.currentThread() eq closingThread) {
+ // Pause close() at its first page free until the spilling thread
holds this
+ // TaskMemoryManager's monitor, so the two lock orders are forced to
overlap.
+ closeFreeingPage.countDown()
+ spillHoldingTaskMemoryManager.await(10, TimeUnit.SECONDS)
+ }
+ super.freePage(page, consumer)
+ }
+ }
+ val tempDir = Utils.createTempDir().getCanonicalFile
+ val queue1 = HybridRowQueue(taskM, tempDir, 1, serManager)
+ val queue2 = HybridRowQueue(taskM, tempDir, 1, serManager)
+ addRows(queue2, 150)
+ assert(queue2.numQueues() > 2)
+
+ val pool = ThreadUtils.newDaemonFixedThreadPool(2, "row-queue-close-spill")
+ try {
+ val closing = pool.submit[Unit](() => {
+ closingThread = Thread.currentThread()
+ queue2.close()
+ })
+ // Mimic TaskMemoryManager.acquireExecutionMemory, which holds its
monitor while it asks
+ // another consumer to spill.
+ val spilling = pool.submit[Long](() => {
+ assert(closeFreeingPage.await(10, TimeUnit.SECONDS))
+ taskM.synchronized {
+ spillHoldingTaskMemoryManager.countDown()
+ queue2.spill(Long.MaxValue, queue1)
+ }
+ })
+ ThreadUtils.awaitResult(closing, 10.seconds)
+ ThreadUtils.awaitResult(spilling, 10.seconds)
Review Comment:
Two small gaps in this test:
- The `Long` returned by the spill is dropped. With these latches, the spill
always runs after `close()` has detached the list, so the expected result is
`0`. Asserting `ThreadUtils.awaitResult(spilling, 10.seconds) === 0L` would
document that.
- At L254, the result of `spillHoldingTaskMemoryManager.await(10,
TimeUnit.SECONDS)` is ignored. If it times out, `close()` just proceeds, and
the test can pass without having forced the two lock orders to overlap.
Asserting it would make a missed handshake fail the test.
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/python/RowQueueSuite.scala:
##########
@@ -170,4 +173,117 @@ class RowQueueSuite extends SparkFunSuite with
EncryptionFunSuite {
queue.close()
}
}
+
+ private def addRows(queue: HybridRowQueue, n: Int): Unit = {
+ val row = new UnsafeRow(1)
+ row.pointTo(new Array[Byte](16), 16)
+ (0 until n).foreach { i =>
+ row.setLong(0, i)
+ assert(queue.add(row) === QueueMode.IN_MEMORY)
+ }
+ }
+
+ private def drainAndClose(queue: HybridRowQueue, n: Int): Unit = {
+ (0 until n).foreach { i =>
+ assert(queue.remove().getLong(0) === i)
+ }
+ queue.close()
+ }
+
+ test("hybrid queues with the same fields are distinct memory consumers") {
+ val conf = new SparkConf(false).set(BUFFER_PAGESIZE, 1024L)
Review Comment:
nit: the three new tests repeat the same setup: the conf with
`BUFFER_PAGESIZE` 1024, the `SerializerManager`, the temp dir, and two queues
with the same arguments. Only the `TaskMemoryManager` differs. A small helper
that takes the `TaskMemoryManager` would remove the duplicated lines and keep
the page size in one place.
--
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]
