diff --git a/docs/source/contributor-guide/memory_management.md b/docs/source/contributor-guide/memory_management.md
index af84e4a2786..5456c8e8c33 100644
--- a/docs/source/contributor-guide/memory_management.md
+++ b/docs/source/contributor-guide/memory_management.md
@@ -137,15 +137,10 @@ allocator before reading anything into the split.
 returns `CometUnifiedShuffleMemoryAllocator`, a Spark `MemoryConsumer` drawing from
 `spark.memory.offHeap.size`, so shuffle pages are arbitrated against Spark's other consumers in the
 same task like any other allocation. A page or pointer array allocation that waits in Spark for
-memory can wake up to find the task's entry gone from Spark's execution pool, because another
-consumer of the task released its last bytes in the meantime, and Spark then throws a
-`NoSuchElementException` ("key not found") instead of a grant
-([SPARK-59444](https://issues.apache.org/jira/browse/SPARK-59444)). The allocator retries that
-case, since the failed call was granted nothing, and after three attempts throws
-`SparkOutOfMemoryError`, which the shuffle writers handle like any other refused page. Only the
-shuffle allocator's callers are guarded this way. Spark's operators in the same task, such as its
-sorters and aggregates, can still hit the exception until Spark re-registers a waiting task in
-`ExecutionMemoryPool` ([apache/spark#58747](https://github.com/apache/spark/pull/58747)).
+memory can wake up to find the task's entry gone from Spark's execution pool; see
+[Constraints on a Comet memory consumer](#constraints-on-a-comet-memory-consumer). The allocator
+retries it, and after three attempts throws `SparkOutOfMemoryError`, which the shuffle writers
+handle like any other refused page.
 
 Which allocator each call site uses, and who ends up charged for the bytes:
 
@@ -216,12 +211,22 @@ that grant, and the only figure available for doing so is the task-wide one abov
 
 **A parked acquire can wake up to a missing task entry.** Spark removes the task's `memoryForTask`
 entry when its balance reaches zero, and an acquire that was waiting in `lock.wait()` and wakes
-afterwards throws a `NoSuchElementException` ("key not found" and the task id) instead of a grant.
-`CometTaskMemoryManager` retries that one case. A missing entry means the task held nothing from
-Spark at that moment, so the failed call has no partial grant to reconcile, and the retry registers
-the task again and waits for its share as the first call would have. After a few attempts it returns
-a zero grant, which the native side treats as a refusal and spills. JVM consumers that call
-`allocatePage` directly are not covered; that gap is tracked in #6304.
+afterwards throws a `NoSuchElementException` ("key not found" and the task id) instead of a grant
+([SPARK-59444](https://issues.apache.org/jira/browse/SPARK-59444)). Releases take no task monitor,
+so another native thread or plan, or a JVM consumer of the same task, can empty the balance while an
+acquire waits. Both of Comet's consumers that ask Spark directly, `CometTaskMemoryManager` and the
+JVM shuffle allocator, retry that one case through `MissingTaskEntryRetry`. A missing entry means
+the task held nothing from Spark at that moment, so the failed call has no partial grant to
+reconcile, and the retry registers the task again and waits for its share as the first call would
+have. After three attempts `CometTaskMemoryManager` returns a zero grant. `try_grow` treats that as
+a refusal and the operator spills, while `grow` records the request as overcommit. The shuffle
+allocator throws `SparkOutOfMemoryError` instead. Both refuse without retrying once the task has
+completed or been killed: releasing a native plan at the end of the task can empty the balance under
+a Tokio worker still waiting in Spark, and a retry would register the ended task again and could wait
+for other tasks while holding the `TaskMemoryManager` monitor that Spark's final cleanup needs.
+Spark's own operators in the same task, such as
+its sorters and aggregates, are not guarded and can still hit the exception until Spark re-registers
+a waiting task in `ExecutionMemoryPool` ([apache/spark#58747](https://github.com/apache/spark/pull/58747)).
 
 **A consumer whose `spill` returns zero takes budget it can never give back.**
 `NativeMemoryConsumer.spill` returns `0`, so Spark can select it as a spill victim and reclaim
diff --git a/spark/src/main/java/org/apache/comet/MissingTaskEntryRetry.java b/spark/src/main/java/org/apache/comet/MissingTaskEntryRetry.java
new file mode 100644
index 00000000000..ae35a4be923
--- /dev/null
+++ b/spark/src/main/java/org/apache/comet/MissingTaskEntryRetry.java
@@ -0,0 +1,162 @@
+/*
+ * 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.comet;
+
+import java.util.NoSuchElementException;
+import java.util.OptionalLong;
+import java.util.function.Function;
+import java.util.function.Supplier;
+
+import org.slf4j.Logger;
+
+import org.apache.spark.TaskContext;
+
+/**
+ * Retries a request for execution memory that Spark failed because the task's entry was gone from
+ * its execution pool.
+ *
+ * <p>Spark registers a task's entry in {@code ExecutionMemoryPool} once, before its wait loop, and
+ * removes it when the task's balance reaches zero. A request that waits for memory below the task's
+ * minimum share then fails when it wakes up with a NoSuchElementException ("key not found: " and
+ * the task id). Releases take no task monitor, so any other consumer of the same task, JVM or
+ * native, can empty the balance while a request waits. The failed call was granted nothing, because
+ * Spark only drops the entry at a zero balance, which cannot hold while the call has a partial
+ * grant. So trying again is safe: it registers the task again and waits for its share as Spark
+ * would have. Nothing here holds a lock across attempts, so a retry parks in Spark exactly as the
+ * first attempt did. The Spark issue is SPARK-59444.
+ *
+ * <p>A request of a task that has completed or been killed is refused instead of retried. Releasing
+ * a native plan when the task ends gives back the plan's memory while a Tokio worker may still wait
+ * in Spark for it, and when that empties the task's balance the worker loses the entry. Retried, it
+ * would register the ended task again and could wait for other tasks to release memory, holding the
+ * TaskMemoryManager monitor that Spark's final cleanup of the task needs.
+ */
+public final class MissingTaskEntryRetry {
+
+  /** Attempts at a request whose task entry Spark keeps losing. */
+  public static final int MAX_ATTEMPTS = 3;
+
+  private MissingTaskEntryRetry() {}
+
+  /**
+   * Runs {@code request}, a call that asks Spark for {@code required} bytes for task {@code
+   * taskAttemptId}, and retries it when it lost that task's entry. See {@link #retry(Logger,
+   * TaskContext, long, Supplier, Function)}.
+   */
+  public static <T> T retry(
+      Logger logger,
+      TaskContext context,
+      long taskAttemptId,
+      long required,
+      Supplier<T> request,
+      Function<NoSuchElementException, T> refuse) {
+    return retry(logger, context, OptionalLong.of(taskAttemptId), required, request, refuse);
+  }
+
+  /**
+   * Runs {@code request}, a call that asks Spark for {@code required} bytes, and retries it when it
+   * lost the task's entry. Without the task id, any task's missing entry counts. Other exceptions
+   * are rethrown. After {@link #MAX_ATTEMPTS} attempts that all lost the entry, returns what {@code
+   * refuse} makes of the last attempt's exception, or throws what it throws. It does the same
+   * without retrying once {@code context}, the task's context or null when there is none, says the
+   * task has completed or been killed. Both the retries and the refusal are logged to {@code
+   * logger}.
+   */
+  public static <T> T retry(
+      Logger logger,
+      TaskContext context,
+      long required,
+      Supplier<T> request,
+      Function<NoSuchElementException, T> refuse) {
+    return retry(logger, context, OptionalLong.empty(), required, request, refuse);
+  }
+
+  private static <T> T retry(
+      Logger logger,
+      TaskContext context,
+      OptionalLong taskAttemptId,
+      long required,
+      Supplier<T> request,
+      Function<NoSuchElementException, T> refuse) {
+    for (int attempt = 1; ; attempt++) {
+      try {
+        return request.get();
+      } catch (NoSuchElementException e) {
+        if (!isMissingTaskEntry(e, taskAttemptId)) {
+          throw e;
+        }
+        String task = taskAttemptId.isPresent() ? "Task " + taskAttemptId.getAsLong() : "The task";
+        if (hasEnded(context)) {
+          logger.info(
+              "{} lost its execution memory entry in Spark while waiting to acquire {} bytes after "
+                  + "it ended, refusing the request",
+              task,
+              required);
+          return refuse.apply(e);
+        }
+        if (attempt >= MAX_ATTEMPTS) {
+          logger.warn(
+              "{} lost its execution memory entry in Spark on {} attempts to acquire {} bytes, "
+                  + "refusing the request",
+              task,
+              attempt,
+              required,
+              e);
+          return refuse.apply(e);
+        }
+        logger.info(
+            "{} lost its execution memory entry in Spark while waiting to acquire {} bytes, "
+                + "trying again",
+            task,
+            required);
+      }
+    }
+  }
+
+  /**
+   * Whether the task has completed or been killed. TaskContextImpl.isCompleted takes the context's
+   * monitor, but Spark runs the task's completion listeners, which release native plans, without
+   * holding it.
+   */
+  private static boolean hasEnded(TaskContext context) {
+    return context != null && (context.isInterrupted() || context.isCompleted());
+  }
+
+  /**
+   * Whether {@code e} is Spark's execution pool failing to find the task's entry: the message names
+   * the task, or any task when its id is unknown, and the exception comes from the pool.
+   */
+  static boolean isMissingTaskEntry(NoSuchElementException e, OptionalLong taskAttemptId) {
+    String message = e.getMessage();
+    boolean matches =
+        taskAttemptId.isPresent()
+            ? ("key not found: " + taskAttemptId.getAsLong()).equals(message)
+            : message != null && message.startsWith("key not found: ");
+    if (!matches) {
+      return false;
+    }
+    for (StackTraceElement frame : e.getStackTrace()) {
+      if ("org.apache.spark.memory.ExecutionMemoryPool".equals(frame.getClassName())) {
+        return true;
+      }
+    }
+    return false;
+  }
+}
diff --git a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
index a468446bae7..1c76627c5b5 100644
--- a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
+++ b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
@@ -20,7 +20,6 @@
 package org.apache.spark;
 
 import java.io.IOException;
-import java.util.NoSuchElementException;
 import java.util.concurrent.atomic.AtomicLong;
 
 import org.slf4j.Logger;
@@ -30,6 +29,8 @@
 import org.apache.spark.memory.MemoryMode;
 import org.apache.spark.memory.TaskMemoryManager;
 
+import org.apache.comet.MissingTaskEntryRetry;
+
 /**
  * A adapter class that is used by Comet native to acquire & release memory through Spark's unified
  * memory manager. This assumes Spark's off-heap memory mode is enabled. All the native plans in a
@@ -39,14 +40,14 @@ public class CometTaskMemoryManager {
 
   private static final Logger logger = LoggerFactory.getLogger(CometTaskMemoryManager.class);
 
-  /** Attempts at an acquire whose task entry Spark keeps losing, see acquireFromSpark. */
-  private static final int MAX_ACQUIRE_ATTEMPTS = 3;
-
   /** The id of the native plan this memory manager was created for, the first in its task. */
   private final long id;
 
   private final long taskAttemptId;
 
+  /** The task's context, read from native threads, which have none of their own. */
+  private final TaskContext taskContext;
+
   public final TaskMemoryManager internal;
   private final NativeMemoryConsumer nativeMemoryConsumer;
   private final AtomicLong used = new AtomicLong();
@@ -54,7 +55,8 @@ public class CometTaskMemoryManager {
   public CometTaskMemoryManager(long id, long taskAttemptId) {
     this.id = id;
     this.taskAttemptId = taskAttemptId;
-    this.internal = TaskContext$.MODULE$.get().taskMemoryManager();
+    this.taskContext = TaskContext$.MODULE$.get();
+    this.internal = taskContext.taskMemoryManager();
     this.nativeMemoryConsumer = new NativeMemoryConsumer();
   }
 
@@ -99,54 +101,18 @@ public long acquireMemory(long size) {
   }
 
   /**
-   * Acquires from Spark, retrying when an acquire waiting in Spark lost the task's entry. Spark
-   * removes a task's entry from its execution pool when the task's balance reaches zero, and an
-   * acquire that was waiting for memory then fails when it wakes up with a NoSuchElementException
-   * ("key not found: " and the task id). Releases take no task monitor, so another native thread or
-   * a JVM consumer of the same task can empty the balance while an acquire waits. The failed call
-   * was granted nothing, because Spark only drops the entry at a zero balance, which cannot hold
-   * while this call has a partial grant. So trying again is safe: it registers the task again and
-   * waits for its share as Spark would have. After a few failed attempts the grant is refused. The
-   * helper holds no lock of its own across attempts, so a retry parks in Spark exactly as the first
-   * attempt did.
+   * Acquires from Spark, retrying when an acquire waiting in Spark lost the task's entry; see
+   * {@link MissingTaskEntryRetry}. When the retries run out, or the task has ended, the grant is
+   * refused, which the native pools treat like any other refusal.
    */
   private long acquireFromSpark(long size) {
-    for (int attempt = 1; ; attempt++) {
-      try {
-        return internal.acquireExecutionMemory(size, nativeMemoryConsumer);
-      } catch (NoSuchElementException e) {
-        if (!isMissingTaskEntry(e)) {
-          throw e;
-        }
-        if (attempt >= MAX_ACQUIRE_ATTEMPTS) {
-          logger.warn(
-              "Task {} lost its execution memory entry in Spark on {} attempts to acquire {} "
-                  + "bytes, refusing the request",
-              taskAttemptId,
-              attempt,
-              size);
-          return 0;
-        }
-        logger.info(
-            "Task {} lost its execution memory entry in Spark while waiting to acquire {} bytes, "
-                + "trying again",
-            taskAttemptId,
-            size);
-      }
-    }
-  }
-
-  /** Whether {@code e} is Spark's execution pool failing to find this task's entry. */
-  private boolean isMissingTaskEntry(NoSuchElementException e) {
-    if (!("key not found: " + taskAttemptId).equals(e.getMessage())) {
-      return false;
-    }
-    for (StackTraceElement frame : e.getStackTrace()) {
-      if ("org.apache.spark.memory.ExecutionMemoryPool".equals(frame.getClassName())) {
-        return true;
-      }
-    }
-    return false;
+    return MissingTaskEntryRetry.retry(
+        logger,
+        taskContext,
+        taskAttemptId,
+        size,
+        () -> internal.acquireExecutionMemory(size, nativeMemoryConsumer),
+        e -> 0L);
   }
 
   // Called by Comet native through JNI
diff --git a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java
index e3a2269cacc..f7bc13ea152 100644
--- a/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java
+++ b/spark/src/main/java/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocator.java
@@ -21,12 +21,12 @@
 
 import java.io.IOException;
 import java.util.Map;
-import java.util.NoSuchElementException;
 import java.util.function.Supplier;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import org.apache.spark.TaskContext;
 import org.apache.spark.memory.MemoryConsumer;
 import org.apache.spark.memory.MemoryMode;
 import org.apache.spark.memory.SparkOutOfMemoryError;
@@ -34,6 +34,8 @@
 import org.apache.spark.unsafe.array.LongArray;
 import org.apache.spark.unsafe.memory.MemoryBlock;
 
+import org.apache.comet.MissingTaskEntryRetry;
+
 /**
  * A simple memory allocator used by `CometShuffleExternalSorter` to allocate memory blocks which
  * store serialized rows. This class is simply an implementation of `MemoryConsumer` that delegates
@@ -49,11 +51,12 @@ public final class CometUnifiedShuffleMemoryAllocator extends CometShuffleMemory
   private static final Logger logger =
       LoggerFactory.getLogger(CometUnifiedShuffleMemoryAllocator.class);
 
-  /** Attempts at an allocation whose task entry Spark keeps losing, see retryMissingTaskEntry. */
-  private static final int MAX_ALLOCATE_ATTEMPTS = 3;
+  /** The task's context, or null outside a task; see retryMissingTaskEntry. */
+  private final TaskContext taskContext;
 
   CometUnifiedShuffleMemoryAllocator(TaskMemoryManager taskMemoryManager, long pageSize) {
     super(taskMemoryManager, pageSize, MemoryMode.OFF_HEAP);
+    this.taskContext = TaskContext.get();
     if (taskMemoryManager.getTungstenMemoryMode() != MemoryMode.OFF_HEAP) {
       throw new IllegalArgumentException(
           "CometUnifiedShuffleMemoryAllocator should be used with off-heap "
@@ -83,31 +86,17 @@ public LongArray allocateArray(long size) {
 
   /**
    * Runs an allocation through Spark, retrying when an allocation waiting in Spark lost the task's
-   * entry. Spark removes a task's entry from its execution pool when the task's balance reaches
-   * zero, and an allocation that was waiting for memory then fails when it wakes up with a
-   * NoSuchElementException ("key not found: " and the task id). Releases take no task monitor, so a
-   * native consumer or another JVM consumer of the same task can empty the balance while an
-   * allocation waits. The failed call was granted nothing and took no page, because Spark only
-   * drops the entry at a zero balance, so trying again is safe: it registers the task again and
-   * waits for its share as Spark would have. After a few failed attempts the allocation fails with
-   * a SparkOutOfMemoryError, which the shuffle writers handle like any other refused page, caused
-   * by the last attempt's exception. The Spark issue is SPARK-59444.
+   * entry; see {@link MissingTaskEntryRetry}. When the retries run out, or the task has ended, the
+   * allocation fails with a SparkOutOfMemoryError, which the shuffle writers handle like any other
+   * refused page, caused by the last attempt's exception.
    */
   private <T> T retryMissingTaskEntry(long required, Supplier<T> allocation) {
-    for (int attempt = 1; ; attempt++) {
-      try {
-        return allocation.get();
-      } catch (NoSuchElementException e) {
-        if (!isMissingTaskEntry(e)) {
-          throw e;
-        }
-        if (attempt >= MAX_ALLOCATE_ATTEMPTS) {
-          logger.warn(
-              "Lost the task's execution memory entry in Spark on {} attempts to allocate {} "
-                  + "bytes, refusing the allocation",
-              attempt,
-              required,
-              e);
+    return MissingTaskEntryRetry.retry(
+        logger,
+        taskContext,
+        required,
+        allocation,
+        e -> {
           SparkOutOfMemoryError error =
               new SparkOutOfMemoryError(
                   "UNABLE_TO_ACQUIRE_MEMORY",
@@ -116,30 +105,7 @@ private <T> T retryMissingTaskEntry(long required, Supplier<T> allocation) {
                       "receivedBytes", String.valueOf(0)));
           error.initCause(e);
           throw error;
-        }
-        logger.info(
-            "Lost the task's execution memory entry in Spark while waiting to allocate {} bytes, "
-                + "trying again",
-            required);
-      }
-    }
-  }
-
-  /**
-   * Whether {@code e} is Spark's execution pool failing to find the task's entry. The task id is
-   * not available here, so this matches the message prefix and the pool's stack frame.
-   */
-  private static boolean isMissingTaskEntry(NoSuchElementException e) {
-    String message = e.getMessage();
-    if (message == null || !message.startsWith("key not found: ")) {
-      return false;
-    }
-    for (StackTraceElement frame : e.getStackTrace()) {
-      if ("org.apache.spark.memory.ExecutionMemoryPool".equals(frame.getClassName())) {
-        return true;
-      }
-    }
-    return false;
+        });
   }
 
   public synchronized long free(MemoryBlock block) {
diff --git a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala
index efd899252ac..82c5731da7a 100644
--- a/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/CometTaskMemoryManagerSuite.scala
@@ -19,16 +19,14 @@
 
 package org.apache.spark
 
-import java.util.Properties
 import java.util.concurrent.TimeUnit
-import java.util.concurrent.atomic.{AtomicInteger, AtomicLong, AtomicReference}
+import java.util.concurrent.atomic.AtomicReference
 
-import org.apache.logging.log4j.{Level, LogManager}
-import org.apache.logging.log4j.core.{LogEvent, LoggerContext}
-import org.apache.spark.executor.TaskMetrics
-import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager, TestMemoryManager, UnifiedMemoryManager}
+import org.apache.logging.log4j.Level
+import org.apache.logging.log4j.core.LogEvent
+import org.apache.spark.memory.{MemoryConsumer, TaskMemoryManager, TestMemoryManager}
 
-class CometTaskMemoryManagerSuite extends SparkFunSuite {
+class CometTaskMemoryManagerSuite extends TaskMemoryTestUtils {
 
   test("native memory usage is visible to Spark's memory consumer") {
     withTaskMemoryManager { taskMemoryManager =>
@@ -136,62 +134,80 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite {
     }
   }
 
-  /** The events that Comet's and Spark's task memory managers log at `level` or above. */
-  private def logEvents(level: Level)(f: => Unit): Seq[LogEvent] = {
-    val appender = new LogAppender("task memory manager")
-    appender.setThreshold(level)
-    val loggers = Seq(classOf[CometTaskMemoryManager].getName, classOf[TaskMemoryManager].getName)
-    val context = LogManager.getContext(false).asInstanceOf[LoggerContext]
-    val unconfigured = loggers.filterNot(context.getConfiguration.getLoggers.containsKey)
-    try withLogAppender(appender, loggers, Some(level))(f)
-    finally {
-      // For a logger with no config of its own, withLogAppender adds one and never removes it.
-      // It copies the root config's additivity, which is off, so left in place it would keep the
-      // logger's events out of the test log for the rest of the run.
-      unconfigured.foreach(context.getConfiguration.removeLogger)
-      context.updateLoggers()
-    }
-    appender.loggingEvents.toSeq
-  }
-
   test("an acquire waiting in Spark survives a release that empties the task's balance") {
-    val memoryManager = offHeapMemoryManager()
-    val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L))
-
-    withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager =>
+    checkRequestWhileBalanceEmpties(loggers) { taskMemoryManager =>
       val manager = new CometTaskMemoryManager(1L, 0L)
-      assert(otherTask.acquireMemory(90L) == 90L)
-      assert(manager.acquireMemory(10L) == 10L)
-
       // Another native thread releases the task's last 10 bytes while the acquire waits.
-      acquireWhileBalanceEmpties(manager, taskMemoryManager, otherTask) {
-        manager.releaseMemory(10L)
-      }
+      WaitingRequest(
+        manager.acquireMemory(_),
+        manager.releaseMemory(_),
+        () => manager.acquireMemory(20L),
+        granted => {
+          assert(granted == 20L)
+          assert(manager.getUsed == 20L)
+          assert(taskMemoryManager.getMemoryConsumptionForThisTask == 20L)
+          manager.releaseMemory(20L)
+          assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L)
+        })
     }
   }
 
-  test("an acquire waiting in Spark survives a sibling consumer freeing the task's last bytes") {
-    val memoryManager = offHeapMemoryManager()
-    val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L))
+  for ((state, end) <- Seq[(String, TaskContextImpl => Unit)](
+      ("completed", _.markTaskCompleted(None)),
+      ("was killed", _.markInterrupted("killed")))) {
+    test(s"an acquire that loses its task entry after the task $state is not made again") {
+      val memoryManager = offHeapMemoryManager()
+      val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L))
 
-    withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager =>
-      val manager = new CometTaskMemoryManager(1L, 0L)
-      val sibling = new OffHeapConsumer(taskMemoryManager)
-      assert(otherTask.acquireMemory(90L) == 90L)
-      assert(sibling.acquireMemory(10L) == 10L)
-
-      // A JVM consumer of the same task frees with freeMemory, which takes no task monitor.
-      acquireWhileBalanceEmpties(manager, taskMemoryManager, otherTask) {
-        sibling.freeMemory(10L)
+      withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager =>
+        val manager = new CometTaskMemoryManager(1L, 0L)
+        assert(otherTask.acquireMemory(90L) == 90L)
+        assert(manager.acquireMemory(10L) == 10L)
+
+        val granted = new AtomicReference[java.lang.Long]()
+        val failure = new AtomicReference[Throwable]()
+        val acquire = new Thread(() =>
+          try granted.set(manager.acquireMemory(20L))
+          catch { case t: Throwable => failure.set(t) })
+        acquire.setDaemon(true)
+        // Spark's executor frees what the task still holds once the task has ended, under the
+        // task memory manager's monitor, which a parked acquire holds.
+        val cleanUp = new Thread(() => taskMemoryManager.cleanUpAllAllocatedMemory())
+        cleanUp.setDaemon(true)
+
+        try {
+          acquire.start()
+          awaitWaitingInSpark(acquire)
+          // The task ends and releasing its plan gives back the task's last bytes while a native
+          // thread still waits in Spark, as releasePlan does when it drops a plan's stream.
+          end(TaskContext.get().asInstanceOf[TaskContextImpl])
+          manager.releaseMemory(10L)
+          acquire.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds))
+          // Parked again, the acquire would wait for the other task while holding the monitor.
+          assert(!acquire.isAlive, s"the acquire is ${acquire.getState} after the task $state")
+          cleanUp.start()
+          cleanUp.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds))
+          assert(!cleanUp.isAlive, s"the clean up is ${cleanUp.getState}")
+        } finally {
+          // Free the other task's memory so that neither thread outlives a failed test.
+          otherTask.freeMemory(otherTask.getUsed)
+          acquire.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds))
+          cleanUp.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds))
+        }
+
+        assert(failure.get == null, s"the acquire failed: ${failure.get}")
+        assert(granted.get == 0L)
+        assert(manager.getUsed == 0L)
+        assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L)
       }
     }
   }
 
   test("a NoSuchElementException other than Spark's missing task entry is rethrown") {
-    val wrongTask = missingEntryError("key not found: 7", fromExecutionMemoryPool = true)
-    val notFromPool = missingEntryError("key not found: 0", fromExecutionMemoryPool = false)
+    val wrongTask = noSuchElement("key not found: 7", fromExecutionMemoryPool = true)
+    val notFromPool = noSuchElement("key not found: 0", fromExecutionMemoryPool = false)
     for (error <- Seq(wrongTask, notFromPool)) {
-      val taskMemoryManager = new FailingTaskMemoryManager(error)
+      val taskMemoryManager = new FailingTaskMemoryManager(() => error, failures = Int.MaxValue)
       withTaskContext(taskMemoryManager) { _ =>
         val manager = new CometTaskMemoryManager(1L, 0L)
         val thrown = intercept[NoSuchElementException](manager.acquireMemory(10L))
@@ -203,108 +219,30 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite {
   }
 
   test("an acquire that keeps losing its task entry gives up with a zero grant") {
-    val error = missingEntryError("key not found: 0", fromExecutionMemoryPool = true)
-    val taskMemoryManager = new FailingTaskMemoryManager(error)
+    // A new exception on every call, so the test can tell which one is logged.
+    val taskMemoryManager = new FailingTaskMemoryManager(
+      () => noSuchElement("key not found: 0", fromExecutionMemoryPool = true),
+      failures = Int.MaxValue)
     withTaskContext(taskMemoryManager) { _ =>
       val manager = new CometTaskMemoryManager(1L, 0L)
-      assert(manager.acquireMemory(10L) == 0L)
-      assert(taskMemoryManager.calls.get == maxAcquireAttempts)
+      val events = logEvents(Level.INFO) {
+        // A refusal rather than an exception, so native spills as for any refused reservation.
+        assert(manager.acquireMemory(10L) == 0L)
+      }
+      assert(events.map(_.getLevel) == Seq(Level.INFO, Level.INFO, Level.WARN), messages(events))
+      assert(events.last.getThrown eq taskMemoryManager.errors.last)
+      assert(taskMemoryManager.calls.get == MaxAcquireAttempts)
       assert(manager.getUsed == 0L)
       assert(nativeMemoryConsumer(manager).getUsed == 0L)
       assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L)
     }
   }
 
-  private val TimeoutSeconds = 10L
-
-  /**
-   * The task holds 10 bytes of a 100 byte pool and another task holds 90. An acquire of 20 bytes
-   * waits in Spark below its minimum share of 25 while `release` takes the task's balance to
-   * zero, which removes the task's entry from Spark's pool. Once the other task frees its memory
-   * the acquire must be granted all 20 bytes.
-   */
-  private def acquireWhileBalanceEmpties(
-      manager: CometTaskMemoryManager,
-      taskMemoryManager: TaskMemoryManager,
-      otherTask: OffHeapConsumer)(release: => Unit): Unit = {
-    val granted = new AtomicLong(-1L)
-    val failure = new AtomicReference[Throwable]()
-    val acquire = new Thread(() =>
-      try granted.set(manager.acquireMemory(20L))
-      catch { case t: Throwable => failure.set(t) })
-    acquire.setDaemon(true)
-
-    try {
-      acquire.start()
-      val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TimeoutSeconds)
-      while (!waitingInSpark(acquire) && System.nanoTime() < deadline) Thread.sleep(10)
-      assert(waitingInSpark(acquire), s"the acquire is ${acquire.getState}")
-      release
-    } finally {
-      // Free the other task's memory so that the acquire does not outlive a failed test.
-      otherTask.freeMemory(otherTask.getUsed)
-      acquire.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds))
-    }
-
-    assert(!acquire.isAlive, s"the acquire is ${acquire.getState}")
-    assert(failure.get == null, s"the acquire failed: ${failure.get}")
-    assert(granted.get == 20L)
-    assert(manager.getUsed == 20L)
-    assert(taskMemoryManager.getMemoryConsumptionForThisTask == 20L)
-    manager.releaseMemory(20L)
-    assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L)
-  }
-
-  private def waitingInSpark(thread: Thread): Boolean =
-    thread.getState == Thread.State.WAITING &&
-      thread.getStackTrace.exists(_.getClassName == "org.apache.spark.memory.ExecutionMemoryPool")
-
-  /** A 100 byte off-heap execution pool. */
-  private def offHeapMemoryManager(): UnifiedMemoryManager = {
-    val conf = new SparkConf()
-      .set("spark.memory.offHeap.enabled", "true")
-      .set("spark.memory.offHeap.size", "100")
-      .set("spark.memory.storageFraction", "0")
-    new UnifiedMemoryManager(conf, 1000L, 500L, 1)
-  }
-
-  /** An off-heap consumer that never spills. */
-  private class OffHeapConsumer(taskMemoryManager: TaskMemoryManager)
-      extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) {
-    override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
-  }
-
-  private val MaxFailingCalls = 10
-
-  /** A task memory manager whose every acquire throws `error`. */
-  private class FailingTaskMemoryManager(error: Throwable)
-      extends TaskMemoryManager(new TestMemoryManager(new SparkConf()), 0L) {
-    val calls = new AtomicInteger()
+  private val loggers =
+    Seq(classOf[CometTaskMemoryManager].getName, classOf[TaskMemoryManager].getName)
 
-    override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = {
-      // Past this bound a caller is retrying without a cap; fail the test instead of hanging it.
-      if (calls.incrementAndGet() > MaxFailingCalls) {
-        throw new IllegalStateException(
-          s"acquireExecutionMemory called more than $MaxFailingCalls times")
-      }
-      throw error
-    }
-  }
-
-  private def missingEntryError(
-      message: String,
-      fromExecutionMemoryPool: Boolean): NoSuchElementException = {
-    val error = new NoSuchElementException(message)
-    if (fromExecutionMemoryPool) {
-      val poolFrame = new StackTraceElement(
-        "org.apache.spark.memory.ExecutionMemoryPool",
-        "acquireMemory",
-        "ExecutionMemoryPool.scala",
-        115)
-      error.setStackTrace(poolFrame +: error.getStackTrace)
-    }
-    error
-  }
+  /** The events that Comet's and Spark's task memory managers log at `level` or above. */
+  private def logEvents(level: Level)(f: => Unit): Seq[LogEvent] = logEvents(level, loggers)(f)
 
   private def withTaskMemoryManager(f: TaskMemoryManager => Unit): Unit = {
     val memoryManager = new TestMemoryManager(new SparkConf())
@@ -312,40 +250,6 @@ class CometTaskMemoryManagerSuite extends SparkFunSuite {
     withTaskContext(new TaskMemoryManager(memoryManager, 0L))(f)
   }
 
-  private def withTaskContext(taskMemoryManager: TaskMemoryManager)(
-      f: TaskMemoryManager => Unit): Unit = {
-    val taskContext = new TaskContextImpl(
-      stageId = 0,
-      stageAttemptNumber = 0,
-      partitionId = 0,
-      numPartitions = 1,
-      taskAttemptId = 0L,
-      attemptNumber = 0,
-      taskMemoryManager = taskMemoryManager,
-      localProperties = new Properties,
-      metricsSystem = null,
-      taskMetrics = TaskMetrics.empty,
-      cpus = 1,
-      resources = Map.empty)
-
-    TaskContext.setTaskContext(taskContext)
-    try {
-      f(taskMemoryManager)
-    } finally {
-      try {
-        taskMemoryManager.cleanUpAllAllocatedMemory()
-      } finally {
-        TaskContext.unset()
-      }
-    }
-  }
-
-  private def maxAcquireAttempts: Int = {
-    val field = classOf[CometTaskMemoryManager].getDeclaredField("MAX_ACQUIRE_ATTEMPTS")
-    field.setAccessible(true)
-    field.getInt(null)
-  }
-
   private def nativeMemoryConsumer(manager: CometTaskMemoryManager): MemoryConsumer = {
     val field = classOf[CometTaskMemoryManager].getDeclaredField("nativeMemoryConsumer")
     field.setAccessible(true)
diff --git a/spark/src/test/scala/org/apache/spark/TaskMemoryTestUtils.scala b/spark/src/test/scala/org/apache/spark/TaskMemoryTestUtils.scala
new file mode 100644
index 00000000000..a82e2377e82
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/TaskMemoryTestUtils.scala
@@ -0,0 +1,219 @@
+/*
+ * 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
+
+import java.util.Properties
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.{AtomicInteger, AtomicReference}
+
+import scala.collection.mutable.ArrayBuffer
+
+import org.apache.logging.log4j.{Level, LogManager}
+import org.apache.logging.log4j.core.{LogEvent, LoggerContext}
+import org.apache.spark.executor.TaskMetrics
+import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager, UnifiedMemoryManager}
+
+import org.apache.comet.MissingTaskEntryRetry
+
+/**
+ * Helpers for suites that drive Spark's task memory manager from Comet's memory consumers, among
+ * them a reproduction of a request that waits in Spark's execution pool and loses the task's
+ * entry there (SPARK-59444).
+ */
+trait TaskMemoryTestUtils extends SparkFunSuite {
+
+  protected val TimeoutSeconds = 10L
+  protected val MaxAcquireAttempts: Int = MissingTaskEntryRetry.MAX_ATTEMPTS
+
+  /** A 100 byte off-heap execution pool. */
+  protected def offHeapMemoryManager(): UnifiedMemoryManager = {
+    val conf = new SparkConf()
+      .set("spark.memory.offHeap.enabled", "true")
+      .set("spark.memory.offHeap.size", "100")
+      .set("spark.memory.storageFraction", "0")
+    new UnifiedMemoryManager(conf, 1000L, 500L, 1)
+  }
+
+  /** An off-heap consumer that never spills. */
+  protected class OffHeapConsumer(taskMemoryManager: TaskMemoryManager)
+      extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) {
+    override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
+  }
+
+  private val MaxFailingCalls = 10
+
+  /**
+   * A task memory manager over a 100 byte off-heap pool whose first `failures` acquires throw the
+   * result of `newError`, recording each one in `errors`.
+   */
+  protected class FailingTaskMemoryManager(newError: () => Throwable, failures: Int)
+      extends TaskMemoryManager(offHeapMemoryManager(), 0L) {
+    val calls = new AtomicInteger()
+    val errors = new ArrayBuffer[Throwable]()
+
+    override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = {
+      val call = calls.incrementAndGet()
+      // Past this bound a caller is retrying without a cap; fail the test instead of hanging it.
+      if (call > MaxFailingCalls) {
+        throw new IllegalStateException(
+          s"acquireExecutionMemory called more than $MaxFailingCalls times")
+      }
+      if (call <= failures) {
+        val error = newError()
+        errors += error
+        throw error
+      }
+      super.acquireExecutionMemory(required, consumer)
+    }
+  }
+
+  protected def noSuchElement(
+      message: String,
+      fromExecutionMemoryPool: Boolean): NoSuchElementException = {
+    val error = new NoSuchElementException(message)
+    if (fromExecutionMemoryPool) {
+      val poolFrame = new StackTraceElement(
+        "org.apache.spark.memory.ExecutionMemoryPool",
+        "acquireMemory",
+        "ExecutionMemoryPool.scala",
+        115)
+      error.setStackTrace(poolFrame +: error.getStackTrace)
+    }
+    error
+  }
+
+  /** Whether `thread` is parked in Spark's execution pool waiting for memory. */
+  protected def waitingInSpark(thread: Thread): Boolean =
+    thread.getState == Thread.State.WAITING &&
+      thread.getStackTrace.exists(_.getClassName == "org.apache.spark.memory.ExecutionMemoryPool")
+
+  /** Waits until `thread` is parked in Spark's execution pool, failing after the timeout. */
+  protected def awaitWaitingInSpark(thread: Thread): Unit = {
+    val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TimeoutSeconds)
+    while (!waitingInSpark(thread) && System.nanoTime() < deadline) Thread.sleep(10)
+    assert(waitingInSpark(thread), s"the request is ${thread.getState}")
+  }
+
+  /**
+   * A request for memory made while the task's balance empties; see
+   * [[checkRequestWhileBalanceEmpties]]. `hold` and `release` take and give back the task's bytes
+   * through a consumer other than the one `request` asks through, and `request` returns what it
+   * was granted. `check` runs on the test thread once the request has returned.
+   */
+  protected case class WaitingRequest(
+      hold: Long => Long,
+      release: Long => Unit,
+      request: () => Long,
+      check: Long => Unit)
+
+  /**
+   * The task holds 10 bytes of a 100 byte pool through `hold` and another task holds 90. The
+   * request is made on another thread and waits in Spark below the task's minimum share of 25
+   * while `release` gives back the task's last 10 bytes, which removes the task's entry from
+   * Spark's pool. Once the other task frees its memory the request must not fail, and `check`
+   * sees what it was granted. The first attempt lost the task's entry, so `loggers` must have
+   * logged a single retry at INFO.
+   */
+  protected def checkRequestWhileBalanceEmpties(loggers: Seq[String])(
+      setup: TaskMemoryManager => WaitingRequest): Unit = {
+    val memoryManager = offHeapMemoryManager()
+    val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L))
+
+    withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager =>
+      val waiting = setup(taskMemoryManager)
+      // The other task goes first: acquiring its 90 bytes after the task's 10 would cap it at 50.
+      assert(otherTask.acquireMemory(90L) == 90L)
+      assert(waiting.hold(10L) == 10L)
+
+      val granted = new AtomicReference[java.lang.Long]()
+      val failure = new AtomicReference[Throwable]()
+      val thread = new Thread(() =>
+        try granted.set(waiting.request())
+        catch { case t: Throwable => failure.set(t) })
+      thread.setDaemon(true)
+
+      val events = logEvents(Level.INFO, loggers) {
+        try {
+          thread.start()
+          awaitWaitingInSpark(thread)
+          waiting.release(10L)
+        } finally {
+          // Free the other task's memory so that the request does not outlive a failed test.
+          otherTask.freeMemory(otherTask.getUsed)
+          thread.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds))
+        }
+      }
+
+      assert(!thread.isAlive, s"the request is ${thread.getState}")
+      assert(failure.get == null, s"the request failed: ${failure.get}")
+      waiting.check(granted.get)
+      assert(events.map(_.getLevel) == Seq(Level.INFO), messages(events))
+    }
+  }
+
+  /** The events that `loggers` log at `level` or above while `f` runs. */
+  protected def logEvents(level: Level, loggers: Seq[String])(f: => Unit): Seq[LogEvent] = {
+    val appender = new LogAppender(loggers.mkString(", "))
+    appender.setThreshold(level)
+    val context = LogManager.getContext(false).asInstanceOf[LoggerContext]
+    val unconfigured = loggers.filterNot(context.getConfiguration.getLoggers.containsKey)
+    try withLogAppender(appender, loggers, Some(level))(f)
+    finally {
+      // For a logger with no config of its own, withLogAppender adds one and never removes it.
+      // It copies the root config's additivity, which is off, so left in place it would keep the
+      // logger's events out of the test log for the rest of the run.
+      unconfigured.foreach(context.getConfiguration.removeLogger)
+      context.updateLoggers()
+    }
+    appender.loggingEvents.toSeq
+  }
+
+  protected def messages(events: Seq[LogEvent]): String =
+    events.map(_.getMessage.getFormattedMessage).mkString("\n")
+
+  /** Runs `f` as task 0 with `taskMemoryManager`, then frees whatever the task still holds. */
+  protected def withTaskContext(taskMemoryManager: TaskMemoryManager)(
+      f: TaskMemoryManager => Unit): Unit = {
+    val taskContext = new TaskContextImpl(
+      stageId = 0,
+      stageAttemptNumber = 0,
+      partitionId = 0,
+      numPartitions = 1,
+      taskAttemptId = 0L,
+      attemptNumber = 0,
+      taskMemoryManager = taskMemoryManager,
+      localProperties = new Properties,
+      metricsSystem = null,
+      taskMetrics = TaskMetrics.empty,
+      cpus = 1,
+      resources = Map.empty)
+
+    TaskContext.setTaskContext(taskContext)
+    try {
+      f(taskMemoryManager)
+    } finally {
+      try {
+        taskMemoryManager.cleanUpAllAllocatedMemory()
+      } finally {
+        TaskContext.unset()
+      }
+    }
+  }
+}
diff --git a/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala b/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala
index b978c272802..ecf2e47be21 100644
--- a/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/shuffle/comet/CometUnifiedShuffleMemoryAllocatorSuite.scala
@@ -19,23 +19,14 @@
 
 package org.apache.spark.shuffle.comet
 
-import java.util.Properties
-import java.util.concurrent.TimeUnit
-import java.util.concurrent.atomic.{AtomicInteger, AtomicReference}
+import org.apache.logging.log4j.Level
+import org.apache.logging.log4j.core.LogEvent
+import org.apache.spark.{CometTaskMemoryManager, TaskContext, TaskContextImpl, TaskMemoryTestUtils}
+import org.apache.spark.memory.SparkOutOfMemoryError
 
-import scala.collection.mutable.ArrayBuffer
-
-import org.apache.logging.log4j.{Level, LogManager}
-import org.apache.logging.log4j.core.{LogEvent, LoggerContext}
-import org.apache.spark.{CometTaskMemoryManager, SparkConf, SparkFunSuite, TaskContext, TaskContextImpl}
-import org.apache.spark.executor.TaskMetrics
-import org.apache.spark.memory.{MemoryConsumer, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager, UnifiedMemoryManager}
-
-class CometUnifiedShuffleMemoryAllocatorSuite extends SparkFunSuite {
+class CometUnifiedShuffleMemoryAllocatorSuite extends TaskMemoryTestUtils {
 
   private val PageSize = 20L
-  private val TimeoutSeconds = 10L
-  private val MaxAllocateAttempts = 3
 
   test("a page waiting in Spark survives a release that empties the task's balance") {
     checkWaitingAllocationSurvives(20L)(allocatePage)
@@ -79,25 +70,45 @@ class CometUnifiedShuffleMemoryAllocatorSuite extends SparkFunSuite {
       }
       assert(events.map(_.getLevel) == Seq(Level.INFO, Level.INFO, Level.WARN), messages(events))
       assert(events.last.getThrown eq taskMemoryManager.errors.last)
-      assert(taskMemoryManager.calls.get == MaxAllocateAttempts)
-      assert(taskMemoryManager.errors.distinct.size == MaxAllocateAttempts)
+      assert(taskMemoryManager.calls.get == MaxAcquireAttempts)
+      assert(taskMemoryManager.errors.distinct.size == MaxAcquireAttempts)
       assert(allocator.getUsed == 0L)
       assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L)
       assert(taskMemoryManager.cleanUpAllAllocatedMemory() == 0L)
     }
   }
 
+  test("an allocation that loses its task entry after the task ended is not made again") {
+    for ((allocation, requested) <- allocations.zip(Seq("20", "24"))) {
+      val taskMemoryManager = new FailingTaskMemoryManager(
+        () => noSuchElement("key not found: 0", fromExecutionMemoryPool = true),
+        failures = Int.MaxValue)
+      withTaskContext(taskMemoryManager) { _ =>
+        val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize)
+        TaskContext.get().asInstanceOf[TaskContextImpl].markTaskCompleted(None)
+        val events = logEvents(Level.INFO) {
+          val thrown = intercept[SparkOutOfMemoryError](allocation(allocator))
+          assert(thrown.getMessageParameters.get("requestedBytes") == requested)
+          assert(thrown.getCause eq taskMemoryManager.errors.last)
+        }
+        assert(events.map(_.getLevel) == Seq(Level.INFO), messages(events))
+        assert(taskMemoryManager.calls.get == 1)
+        assert(allocator.getUsed == 0L)
+      }
+    }
+  }
+
   test("an allocation that loses its task entry twice is granted on the last attempt") {
     val error = noSuchElement("key not found: 0", fromExecutionMemoryPool = true)
     for ((allocation, requested) <- allocations.zip(Seq(20L, 24L))) {
       val taskMemoryManager =
-        new FailingTaskMemoryManager(() => error, failures = MaxAllocateAttempts - 1)
+        new FailingTaskMemoryManager(() => error, failures = MaxAcquireAttempts - 1)
       val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize)
       val events = logEvents(Level.INFO) {
         assert(allocation(allocator) == requested)
       }
       assert(events.map(_.getLevel) == Seq(Level.INFO, Level.INFO), messages(events))
-      assert(taskMemoryManager.calls.get == MaxAllocateAttempts)
+      assert(taskMemoryManager.calls.get == MaxAcquireAttempts)
       assert(allocator.getUsed == 0L)
       assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L)
       assert(taskMemoryManager.cleanUpAllAllocatedMemory() == 0L)
@@ -124,162 +135,27 @@ class CometUnifiedShuffleMemoryAllocatorSuite extends SparkFunSuite {
     Seq(allocatePage _, allocatePointerArray _)
 
   /**
-   * The task holds 10 bytes of a 100 byte pool through a native consumer and another task holds
-   * 90. `allocate` asks for `requested` bytes on another thread and waits in Spark below the
-   * task's minimum share of 25 while the native consumer releases the task's last 10 bytes, which
-   * removes the task's entry from Spark's pool. Once the other task frees its memory the
-   * allocation must be granted in full.
+   * An allocation of `requested` bytes waits in Spark while a native consumer of the task
+   * releases the task's last bytes, and must then be granted in full.
    */
   private def checkWaitingAllocationSurvives(requested: Long)(
-      allocate: CometUnifiedShuffleMemoryAllocator => Long): Unit = {
-    val memoryManager = offHeapMemoryManager()
-    val otherTask = new OffHeapConsumer(new TaskMemoryManager(memoryManager, 1L))
-
-    withTaskContext(new TaskMemoryManager(memoryManager, 0L)) { taskMemoryManager =>
-      val native = new CometTaskMemoryManager(1L, 0L)
-      val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize)
-      // The other task goes first: acquiring its 90 bytes after the task's 10 would cap it at 50.
-      assert(otherTask.acquireMemory(90L) == 90L)
-      assert(native.acquireMemory(10L) == 10L)
-
-      val granted = new AtomicReference[java.lang.Long]()
-      val failure = new AtomicReference[Throwable]()
-      val thread = new Thread(() =>
-        try granted.set(allocate(allocator))
-        catch { case t: Throwable => failure.set(t) })
-      thread.setDaemon(true)
-
-      val events = logEvents(Level.INFO) {
-        try {
-          thread.start()
-          val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(TimeoutSeconds)
-          while (!waitingInSpark(thread) && System.nanoTime() < deadline) Thread.sleep(10)
-          assert(waitingInSpark(thread), s"the allocation is ${thread.getState}")
-          native.releaseMemory(10L)
-        } finally {
-          // Free the other task's memory so that the allocation does not outlive a failed test.
-          otherTask.freeMemory(otherTask.getUsed)
-          thread.join(TimeUnit.SECONDS.toMillis(TimeoutSeconds))
-        }
-      }
-
-      assert(!thread.isAlive, s"the allocation is ${thread.getState}")
-      assert(failure.get == null, s"the allocation failed: ${failure.get}")
-      assert(granted.get == requested)
-      assert(allocator.getUsed == 0L)
-      assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L)
-      // The first attempt lost the task's entry and the allocator tried again.
-      assert(events.map(_.getLevel) == Seq(Level.INFO), messages(events))
+      allocate: CometUnifiedShuffleMemoryAllocator => Long): Unit =
+    checkRequestWhileBalanceEmpties(Seq(classOf[CometUnifiedShuffleMemoryAllocator].getName)) {
+      taskMemoryManager =>
+        val native = new CometTaskMemoryManager(1L, 0L)
+        val allocator = new CometUnifiedShuffleMemoryAllocator(taskMemoryManager, PageSize)
+        WaitingRequest(
+          native.acquireMemory(_),
+          native.releaseMemory(_),
+          () => allocate(allocator),
+          granted => {
+            assert(granted == requested)
+            assert(allocator.getUsed == 0L)
+            assert(taskMemoryManager.getMemoryConsumptionForThisTask == 0L)
+          })
     }
-  }
-
-  private def waitingInSpark(thread: Thread): Boolean =
-    thread.getState == Thread.State.WAITING &&
-      thread.getStackTrace.exists(_.getClassName == "org.apache.spark.memory.ExecutionMemoryPool")
-
-  /** A 100 byte off-heap execution pool. */
-  private def offHeapMemoryManager(): UnifiedMemoryManager = {
-    val conf = new SparkConf()
-      .set("spark.memory.offHeap.enabled", "true")
-      .set("spark.memory.offHeap.size", "100")
-      .set("spark.memory.storageFraction", "0")
-    new UnifiedMemoryManager(conf, 1000L, 500L, 1)
-  }
-
-  /** An off-heap consumer that never spills. */
-  private class OffHeapConsumer(taskMemoryManager: TaskMemoryManager)
-      extends MemoryConsumer(taskMemoryManager, 0L, MemoryMode.OFF_HEAP) {
-    override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
-  }
-
-  private val MaxFailingCalls = 10
-
-  /**
-   * A task memory manager over a 100 byte off-heap pool whose first `failures` acquires throw the
-   * result of `newError`, recording each one in `errors`.
-   */
-  private class FailingTaskMemoryManager(newError: () => Throwable, failures: Int)
-      extends TaskMemoryManager(offHeapMemoryManager(), 0L) {
-    val calls = new AtomicInteger()
-    val errors = new ArrayBuffer[Throwable]()
-
-    override def acquireExecutionMemory(required: Long, consumer: MemoryConsumer): Long = {
-      val call = calls.incrementAndGet()
-      // Past this bound a caller is retrying without a cap; fail the test instead of hanging it.
-      if (call > MaxFailingCalls) {
-        throw new IllegalStateException(
-          s"acquireExecutionMemory called more than $MaxFailingCalls times")
-      }
-      if (call <= failures) {
-        val error = newError()
-        errors += error
-        throw error
-      }
-      super.acquireExecutionMemory(required, consumer)
-    }
-  }
-
-  private def noSuchElement(
-      message: String,
-      fromExecutionMemoryPool: Boolean): NoSuchElementException = {
-    val error = new NoSuchElementException(message)
-    if (fromExecutionMemoryPool) {
-      val poolFrame = new StackTraceElement(
-        "org.apache.spark.memory.ExecutionMemoryPool",
-        "acquireMemory",
-        "ExecutionMemoryPool.scala",
-        115)
-      error.setStackTrace(poolFrame +: error.getStackTrace)
-    }
-    error
-  }
 
   /** The events that the allocator logs at `level` or above. */
-  private def logEvents(level: Level)(f: => Unit): Seq[LogEvent] = {
-    val appender = new LogAppender("unified shuffle memory allocator")
-    appender.setThreshold(level)
-    val loggers = Seq(classOf[CometUnifiedShuffleMemoryAllocator].getName)
-    val context = LogManager.getContext(false).asInstanceOf[LoggerContext]
-    val unconfigured = loggers.filterNot(context.getConfiguration.getLoggers.containsKey)
-    try withLogAppender(appender, loggers, Some(level))(f)
-    finally {
-      // For a logger with no config of its own, withLogAppender adds one and never removes it.
-      // It copies the root config's additivity, which is off, so left in place it would keep the
-      // logger's events out of the test log for the rest of the run.
-      unconfigured.foreach(context.getConfiguration.removeLogger)
-      context.updateLoggers()
-    }
-    appender.loggingEvents.toSeq
-  }
-
-  private def messages(events: Seq[LogEvent]): String =
-    events.map(_.getMessage.getFormattedMessage).mkString("\n")
-
-  private def withTaskContext(taskMemoryManager: TaskMemoryManager)(
-      f: TaskMemoryManager => Unit): Unit = {
-    val taskContext = new TaskContextImpl(
-      stageId = 0,
-      stageAttemptNumber = 0,
-      partitionId = 0,
-      numPartitions = 1,
-      taskAttemptId = 0L,
-      attemptNumber = 0,
-      taskMemoryManager = taskMemoryManager,
-      localProperties = new Properties,
-      metricsSystem = null,
-      taskMetrics = TaskMetrics.empty,
-      cpus = 1,
-      resources = Map.empty)
-
-    TaskContext.setTaskContext(taskContext)
-    try {
-      f(taskMemoryManager)
-    } finally {
-      try {
-        taskMemoryManager.cleanUpAllAllocatedMemory()
-      } finally {
-        TaskContext.unset()
-      }
-    }
-  }
+  private def logEvents(level: Level)(f: => Unit): Seq[LogEvent] =
+    logEvents(level, Seq(classOf[CometUnifiedShuffleMemoryAllocator].getName))(f)
 }