fix: share the missing task entry retry and stop retrying it after the task has ended

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

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;
+  }
+}

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,22 +40,23 @@ 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();
 
   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