viirya commented on code in PR #59282:
URL: https://github.com/apache/spark/pull/59282#discussion_r4235494447
##########
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:
Done in 7ee7678. `HybridQueue` now defines final identity
`equals`/`hashCode`, and `spill` checks `trigger eq this`. I kept
`HybridRowQueue` as a regular class. Copying or destructuring a memory consumer
by its constructor arguments is not meaningful, and the generated case-class
`toString` would print the same string for two queues with the same arguments,
so it would need an override anyway.
##########
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:
Agreed. I checked the other `MemoryConsumer` subclasses, and none of them
defines value equality today, so this protects future consumers. Since it
changes core, I filed SPARK-60112 for it.
##########
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:
Good catch, confirmed. After a successful `cancel(true)`, `FutureTask.get()`
throws `CancellationException` right away without waiting for the runnable, so
the listener's assumption does not hold. This came with the pipelined mode in
SPARK-56642. I filed SPARK-60113 to make the listener wait for the writer to
actually exit. That also covers the page leaked by a `createNewQueue` that
races with the detach.
##########
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:
Done in 7ee7678 with the `toString` you suggested.
##########
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:
Done in 7ee7678. The handshake `await` is now asserted inside `freePage`,
and the spill result is asserted to be 0. I also checked that the test still
times out with the old `close()`.
##########
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:
Done in 7ee7678. The new tests now share `smallPageConf` and a
`queuesWithSameArguments(taskM)` helper.
--
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]
