[SPARK-35486][CORE] TaskMemoryManager: retry if other task takes memory freed by partial self-spill

### What changes were proposed in this pull request?

When a memory reservation triggers a self-spill, `ExecutionMemoryPool#releaseMemory()` will immediately notify waiting tasks that memory has been freed. If there are any waiting tasks with less than 1/2N of the memory pool, they may acquire the newly-freed memory before the current task has a chance to do so. This will cause the original memory reservation to fail. If the initial spill did not release all available memory, the reservation could have been satisfied by asking it to spill again.

This PR adds logic to TaskMemoryManager to detect this case and retry.

### Why are the changes needed?

This bug affects queries with a MemoryConsumer that can spill part of its memory, such as BytesToBytesMap. If the MemoryConsumer is using all available memory and there is a waiting task, then attempting to acquire more memory on the MemoryConsumer will trigger a partial self-spill. However, because the waiting task gets priority, the attempt to acquire memory will fail even if it could have been satisfied by another spill.

### Does this PR introduce _any_ user-facing change?

No.

### How was this patch tested?

Added a test to MemoryManagerSuite that previously failed and now passes.

Closes #32625 from ankurdave/SPARK-35486.

Authored-by: Ankur Dave <ankurdave@gmail.com>
Signed-off-by: yi.wu <yi.wu@databricks.com>

Author: ankurdave

core/src/main/java/org/apache/spark/memory/TaskMemoryManager.java

@@ -202,14 +202,22 @@ public long acquireExecutionMemory(long required, MemoryConsumer consumer) {
         }
       }
 
-      // call spill() on itself
-      if (got < required) {
+      // Attempt to free up memory by self-spilling.
+      //
+      // When our spill handler releases memory, `ExecutionMemoryPool#releaseMemory()` will
+      // immediately notify other tasks that memory has been freed, and they may acquire the
+      // newly-freed memory before we have a chance to do so (SPARK-35486). In that case, we will
+      // try again in the next loop iteration.
+      while (got < required) {
         try {
           long released = consumer.spill(required - got, consumer);
           if (released > 0) {
             logger.debug("Task {} released {} from itself ({})", taskAttemptId,
               Utils.bytesToString(released), consumer);
             got += memoryManager.acquireExecutionMemory(required - got, taskAttemptId, mode);
+          } else {
+            // Self-spilling could not free up any more memory.
+            break;
           }
         } catch (ClosedByInterruptException e) {
           // This called by user to kill a task (e.g: speculative task).

core/src/test/java/org/apache/spark/memory/TestPartialSpillingMemoryConsumer.java

@@ -0,0 +1,48 @@
+/*
+ * 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.io.IOException;
+
+/**
+ * A TestMemoryConsumer which, when asked to spill, releases only enough memory to satisfy the
+ * request rather than releasing all its memory.
+ */
+public class TestPartialSpillingMemoryConsumer extends TestMemoryConsumer {
+  private long spilledBytes = 0L;
+
+  public TestPartialSpillingMemoryConsumer(TaskMemoryManager memoryManager, MemoryMode mode) {
+    super(memoryManager, mode);
+  }
+  public TestPartialSpillingMemoryConsumer(TaskMemoryManager memoryManager) {
+    super(memoryManager);
+  }
+
+  @Override
+  public long spill(long size, MemoryConsumer trigger) throws IOException {
+    long used = getUsed();
+    long released = Math.min(used, size);
+    free(released);
+    spilledBytes += released;
+    return released;
+  }
+
+  public long getSpilledBytes() {
+    return spilledBytes;
+  }
+}

core/src/test/scala/org/apache/spark/memory/MemoryManagerSuite.scala

@@ -240,6 +240,40 @@ private[memory] trait MemoryManagerSuite extends SparkFunSuite with BeforeAndAft
     assert(ThreadUtils.awaitResult(t2Result2, 200.millis) === 0L)
   }
 
+  test("SPARK-35486: memory freed by self-spilling is taken by another task") {
+    val memoryManager = createMemoryManager(1000L)
+    val t1MemManager = new TaskMemoryManager(memoryManager, 1)
+    val t2MemManager = new TaskMemoryManager(memoryManager, 2)
+    val c1 = new TestPartialSpillingMemoryConsumer(t1MemManager)
+    val c2 = new TestMemoryConsumer(t2MemManager)
+    val futureTimeout: Duration = 20.seconds
+
+    // t1 acquires 1000 bytes. This should succeed immediately.
+    val t1Result1 = Future { c1.acquireMemory(1000L) }
+    assert(ThreadUtils.awaitResult(t1Result1, futureTimeout) === 1000L)
+    assert(c1.getUsed() === 1000L)
+    assert(c1.getSpilledBytes() === 0L)
+
+    // t2 attempts to acquire 500 bytes. This should block since there is no memory available.
+    val t2Result1 = Future { c2.acquireMemory(500L) }
+    Thread.sleep(300)
+    assert(!t2Result1.isCompleted)
+    assert(c2.getUsed() === 0L)
+
+    // t1 attempts to acquire 500 bytes, causing its existing reservation to spill partially. After
+    // the spill, t1 is still at its fair share of 500 bytes, so it cannot acquire memory and t2
+    // gets the freed memory instead. t1 must try again, causing the rest of the reservation to
+    // spill.
+    val t1Result2 = Future { c1.acquireMemory(500L) }
+
+    // The spill should release enough memory for both t1's and t2's reservations to be satisfied.
+    assert(ThreadUtils.awaitResult(t2Result1, futureTimeout) === 500L)
+    assert(ThreadUtils.awaitResult(t1Result2, futureTimeout) === 500L)
+    assert(c1.getSpilledBytes() === 1000L)
+    assert(c1.getUsed() === 500L)
+    assert(c2.getUsed() === 500L)
+  }
+
   test("TaskMemoryManager.cleanUpAllAllocatedMemory") {
     val memoryManager = createMemoryManager(1000L)
     val t1MemManager = new TaskMemoryManager(memoryManager, 1)