peter-toth commented on code in PR #58747:
URL: https://github.com/apache/spark/pull/58747#discussion_r3999304919
##########
core/src/main/scala/org/apache/spark/memory/ExecutionMemoryPool.scala:
##########
@@ -164,15 +202,16 @@ private[memory] class ExecutionMemoryPool(
}
if (memoryForTask.contains(taskAttemptId)) {
memoryForTask(taskAttemptId) -= memoryToFree
- if (memoryForTask(taskAttemptId) <= 0) {
+ if (memoryForTask(taskAttemptId) <= 0 &&
!hasWaitingAcquisition(taskAttemptId)) {
Review Comment:
**Finding 2.** This line is the whole lifecycle change - a task at zero
bytes stays registered while an acquisition waits. `releaseMemory`'s scaladoc
at `:190` is unchanged, while the two callers that only delegate to it both got
the note. The rule should be stated where it is enforced:
```scala
/**
* Release `numBytes` of memory acquired by the given task. The task is
deregistered once it
* has neither reserved bytes nor a waiting acquisition.
*/
```
Two more places still describe the old lifecycle:
- `memoryForTask` at
`core/src/main/scala/org/apache/spark/memory/ExecutionMemoryPool.scala:54` -
"Map from taskAttemptId -> memory consumption in bytes" does not say an entry
can be a zero-byte placeholder for a waiter.
- the class scaladoc at `:34`, "we keep track of the set of active tasks",
which is the set `numActiveTasks` divides by.
Same area as
[r3997983841](https://github.com/apache/spark/pull/58747#discussion_r3997983841),
so it is probably one edit.
##########
core/src/test/scala/org/apache/spark/memory/ExecutionMemoryPoolSuite.scala:
##########
@@ -0,0 +1,242 @@
+/*
+ * 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.memory
+
+import java.util.concurrent.{CompletableFuture, ExecutionException, TimeUnit}
+import java.util.concurrent.atomic.AtomicBoolean
+
+import scala.collection.mutable.ArrayBuffer
+import scala.util.control.NonFatal
+
+import org.scalatest.concurrent.Eventually
+import org.scalatest.time.SpanSugar._
+
+import org.apache.spark.{SparkConf, SparkFunSuite}
+
+class ExecutionMemoryPoolSuite extends SparkFunSuite with Eventually {
+ private class WaitingAcquire(acquire: () => Long) {
+ val result = new CompletableFuture[Long]()
+ val thread = new Thread("execution-memory-waiter") {
+ override def run(): Unit = {
+ try {
+ result.complete(acquire())
+ } catch {
+ case e: InterruptedException => result.completeExceptionally(e)
+ case NonFatal(e) => result.completeExceptionally(e)
+ }
+ }
+ }
+ thread.setDaemon(true)
+
+ def awaitWaiting(): Unit = eventually(timeout(10.seconds)) {
+ assert(!result.isDone, "acquisition completed instead of waiting")
+ assert(thread.getState == Thread.State.WAITING)
+ assert(thread.getStackTrace.exists { frame =>
+ frame.getClassName == classOf[ExecutionMemoryPool].getName &&
+ frame.getMethodName == "acquireMemory"
+ })
+ }
+
+ def acquired(): Long = result.get(10, TimeUnit.SECONDS)
+
+ def failure(): Throwable = {
+ val error = intercept[ExecutionException] {
+ result.get(10, TimeUnit.SECONDS)
+ }
+ error.getCause
+ }
+
+ def interrupt(): Unit = {
+ thread.interrupt()
+ assert(failure().isInstanceOf[InterruptedException])
+ }
+ }
+
+ private val waiters = new ArrayBuffer[WaitingAcquire]()
+
+ override protected def afterEach(): Unit = {
+ try {
+ waiters.foreach(_.thread.interrupt())
+ waiters.foreach(_.thread.join(10000))
+ assert(waiters.forall(!_.thread.isAlive), "memory acquisition thread did
not terminate")
+ } finally {
+ waiters.clear()
+ super.afterEach()
+ }
+ }
+
+ private def acquireAsync(
+ pool: ExecutionMemoryPool,
+ bytes: Long,
+ maybeGrowPool: Long => Unit = _ => ()): WaitingAcquire = {
+ acquireAsync(pool.acquireMemory(bytes, 1L, maybeGrowPool))
+ }
+
+ private def acquireAsync(acquire: => Long): WaitingAcquire = {
+ val waiter = new WaitingAcquire(() => acquire)
+ waiters += waiter
+ waiter.thread.start()
+ waiter.awaitWaiting()
+ waiter
+ }
+
+ private def newPool(mode: MemoryMode): ExecutionMemoryPool = {
+ val pool = new ExecutionMemoryPool(new Object, mode)
+ pool.incrementPoolSize(1000L)
+ assert(pool.acquireMemory(900L, 2L) == 900L)
+ pool
+ }
+
+ for (mode <- Seq(MemoryMode.ON_HEAP, MemoryMode.OFF_HEAP)) {
+ test(s"retain task registration across two consumers of one
TaskMemoryManager ($mode)") {
+ val conf = new SparkConf(false)
+ .set("spark.memory.offHeap.enabled", "true")
+ .set("spark.memory.offHeap.size", "1000")
+ val memory = new UnifiedMemoryManager(conf, 1000L, 500L, 1)
+ val task = new TaskMemoryManager(memory, 1L)
+ val peerTask = new TaskMemoryManager(memory, 2L)
+ val owner = new TestMemoryConsumer(task, mode)
+ val requester = new TestMemoryConsumer(task, mode)
+ val peer = new TestMemoryConsumer(peerTask, mode)
+ assert(peer.acquireMemory(900L) == 900L)
+ assert(owner.acquireMemory(100L) == 100L)
+ val waiter = acquireAsync(requester.acquireMemory(300L))
+
+ // Release through another consumer of the waiting task, not through the
pool directly.
+ owner.freeMemory(100L)
+ peer.freeMemory(300L)
+ assert(waiter.acquired() == 300L)
+ assert(owner.getUsed() == 0L)
+ assert(requester.getUsed() == 300L)
+ assert(task.getMemoryConsumptionForThisTask() == 300L)
+ assert(peerTask.getMemoryConsumptionForThisTask() == 600L)
+ assert(memory.executionMemoryUsed == 900L)
+
+ requester.freeMemory(300L)
+ peer.freeMemory(600L)
+ assert(task.cleanUpAllAllocatedMemory() == 0L)
+ assert(peerTask.cleanUpAllAllocatedMemory() == 0L)
+ assert(memory.executionMemoryUsed == 0L)
+ }
+
+ for (releaseAll <- Seq(false, true)) {
+ test(s"retain a waiting task after its last release ($mode,
releaseAll=$releaseAll)") {
+ val pool = newPool(mode)
+ assert(pool.acquireMemory(100L, 1L) == 100L)
+ val waiter = acquireAsync(pool, 300L)
+
+ if (releaseAll) {
+ assert(pool.releaseAllMemoryForTask(1L) == 100L)
+ } else {
+ pool.releaseMemory(100L, 1L)
+ }
+ pool.releaseMemory(300L, 2L)
+
+ assert(waiter.acquired() == 300L)
+ assert(pool.getMemoryUsageForTask(1L) == 300L)
+ assert(pool.memoryUsed == 900L)
Review Comment:
**Finding 1.** The suite pins the *retain* half of the waiter bookkeeping
thoroughly and never checks the *release* half. Nothing asserts that a task
which waited is fully deregistered once it has no waiter and no reservation
left.
I measured the gap rather than guessing. Leaking the waiter count on the
success path only:
```scala
var registeredWaiter = false
var succeeded = false
...
memoryForTask(taskAttemptId) += toGrant
succeeded = true
return toGrant
...
} finally {
if (registeredWaiter && !succeeded) {
```
All 18 new cases still pass. The leak is permanent: `waitingAcquisitions`
keeps the entry forever, so `releaseMemory` never removes the task, and a
phantom zero-byte participant stays in `memoryForTask` for the life of the JVM.
Every other task's `maxMemoryPerTask` is halved from then on.
Three lines at the end of this test close it:
```suggestion
assert(pool.memoryUsed == 900L)
// The waiter is gone, so releasing the reservation must deregister
the task.
pool.releaseMemory(300L, 1L)
// numActiveTasks == 1 caps task 2 at the whole pool. A phantom
entry would cap it
// at 500 < 600 and this would return 0.
assert(pool.acquireMemory(100L, 2L) == 100L)
```
I ran that check as a standalone case: it passes on the PR as-is and fails
on the leak above.
The same hole exists across tasks. All 18 cases only ever have task 1
waiting, so nothing covers task 1's last waiter leaving while task 2 still has
one - the `waitingAcquisitions.isEmpty` / `null` reset path. Worth a case if
the waiter map survives
[r3998051507](https://github.com/apache/spark/pull/58747#discussion_r3998051507).
--
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]
