fix: make CometDiskBlockWriter spill registry per-task instead of executor-global

.github/workflows/pr_build_linux.yml

@@ -333,6 +333,7 @@ jobs:
               org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleReaderSuite
               org.apache.spark.sql.comet.execution.shuffle.CometCelebornShufflePlanningSuite
               org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite
+              org.apache.spark.sql.comet.execution.shuffle.CometDiskBlockWriterSuite
               org.apache.comet.exec.CometShuffleEncryptionSuite
               org.apache.comet.exec.CometShuffleManagerSuite
               org.apache.comet.exec.CometAsyncShuffleSuite

.github/workflows/pr_build_macos.yml

@@ -149,6 +149,7 @@ jobs:
               org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleReaderSuite
               org.apache.spark.sql.comet.execution.shuffle.CometCelebornShufflePlanningSuite
               org.apache.spark.sql.comet.execution.shuffle.CometNativeShuffleInputRDDSuite
+              org.apache.spark.sql.comet.execution.shuffle.CometDiskBlockWriterSuite
               org.apache.comet.exec.CometShuffleEncryptionSuite
               org.apache.comet.exec.CometShuffleManagerSuite
               org.apache.comet.exec.CometAsyncShuffleSuite

spark/src/main/java/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocator.java

@@ -21,8 +21,13 @@
 
 import java.io.IOException;
 import java.util.BitSet;
+import java.util.HashMap;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import org.apache.spark.SparkConf;
+import org.apache.spark.TaskContext;
 import org.apache.spark.memory.MemoryConsumer;
 import org.apache.spark.memory.MemoryMode;
 import org.apache.spark.memory.SparkOutOfMemoryError;
@@ -32,6 +37,7 @@
 import org.apache.spark.unsafe.memory.MemoryBlock;
 import org.apache.spark.unsafe.memory.UnsafeMemoryAllocator;
 
+import org.apache.comet.CometConf$;
 import org.apache.comet.CometSparkSessionExtensions$;
 
 /**
@@ -51,12 +57,21 @@
  * Columnar Shuffle and execution apart from Spark's on-heap memory configuration.
  */
 public final class CometBoundedShuffleMemoryAllocator extends CometShuffleMemoryAllocatorTrait {
+  private static final Logger logger =
+      LoggerFactory.getLogger(CometBoundedShuffleMemoryAllocator.class);
+
   private final UnsafeMemoryAllocator allocator = new UnsafeMemoryAllocator();
 
   private final long pageSize;
   private final long totalMemory;
   private long allocatedMemory = 0L;
 
+  /** How often a thread blocked in {@link #allocateBlocking(long)} logs that it is waiting. */
+  private static final long WAIT_LOG_INTERVAL_MS = 30_000L;
+
+  /** How often a blocked thread checks for cooperative task cancellation. */
+  private static final long TASK_KILL_POLL_INTERVAL_MS = 1_000L;
+
   /** The number of bits used to address the page table. */
   private static final int PAGE_NUMBER_BITS = 13;
 
@@ -66,6 +81,15 @@ public final class CometBoundedShuffleMemoryAllocator extends CometShuffleMemory
   private final MemoryBlock[] pageTable = new MemoryBlock[PAGE_TABLE_SIZE];
   private final BitSet allocatedPages = new BitSet(PAGE_TABLE_SIZE);
 
+  /** The thread that allocated each page, used to decide whether a blocked wait can succeed. */
+  private final Thread[] pageOwners = new Thread[PAGE_TABLE_SIZE];
+
+  /** Pool memory currently retained by each thread. */
+  private final HashMap<Thread, Long> retainedMemory = new HashMap<>();
+
+  /** Threads currently blocked in {@link #allocateBlocking(long)} and their request sizes. */
+  private final HashMap<Thread, Long> waitingThreads = new HashMap<>();
+
   private static final int OFFSET_BITS = 51;
   private static final long MASK_LONG_LOWER_51_BITS = 0x7FFFFFFFFFFFFL;
 
@@ -123,6 +147,122 @@ public synchronized MemoryBlock allocate(long required) {
     return allocateMemoryBlock(size);
   }
 
+  /**
+   * Like {@link #allocate(long)}, but waits for other tasks of this shared pool to free memory,
+   * mirroring how Spark's unified memory manager blocks a task until memory becomes available.
+   * Callers must first spill buffered data they can cheaply release; memory this thread still
+   * retains (e.g. the sorter's pointer array or sibling writers' pages) is included in the liveness
+   * checks below. The wait fails fast when it can never succeed: when the request does not fit next
+   * to the requester's retained memory, or when all allocated memory is retained by blocked threads
+   * and none of their requests fits in the free pool. Because the holders it depends on may in turn
+   * be blocked on resources outside this pool that only a task waiting here can release, the wait
+   * is also bounded by `spark.comet.shuffle.jvm.memoryWaitTimeout`, after which the managed
+   * allocation error is thrown and Spark's task retry can recover. Task cancellation or Java
+   * interruption aborts the wait.
+   */
+  @Override
+  public synchronized MemoryBlock allocateBlocking(long required) {
+    long memoryWaitTimeoutMs =
+        (long) CometConf$.MODULE$.COMET_SHUFFLE_JVM_MEMORY_WAIT_TIMEOUT().get();
+    long size = Math.max(pageSize, required);
+    Thread self = Thread.currentThread();
+    TaskContext taskContext = TaskContext.get();
+    long waitStart = 0;
+    long lastLog = 0;
+    try {
+      while (true) {
+        if (taskContext != null) {
+          taskContext.killTaskIfInterrupted();
+        }
+        try {
+          return allocateMemoryBlock(size);
+        } catch (SparkOutOfMemoryError e) {
+          if (waitingThreads.put(self, size) == null) {
+            // Wake existing waiters so they re-evaluate the deadlock check against the enlarged
+            // waiting set.
+            notifyAll();
+          }
+          // This thread cannot free what it retains while it waits, so a request that does not
+          // fit next to its own retained memory can never be satisfied.
+          if (size > totalMemory - retainedMemory.getOrDefault(self, 0L)) {
+            throw e;
+          }
+          // The allocation just failed, so the request does not fit in the unallocated pool.
+          // Waiting can only succeed while some thread can still free memory: either a thread
+          // outside the waiting set retains pool memory, or another waiter's request fits in the
+          // free pool, in which case that waiter can proceed and eventually free what it retains.
+          if (allocatedMemory <= retainedByWaitingThreads() && !anyWaiterCanProceed()) {
+            throw e;
+          }
+          // The holders this wait depends on may themselves be blocked on resources outside
+          // this pool (Spark execution memory, locks, I/O) that only a task waiting here can
+          // release - a cycle this allocator cannot observe. Bound the wait so such cycles
+          // unwind with the managed allocation error instead of hanging the executor; Spark's
+          // task retry can then recover.
+          long now = System.currentTimeMillis();
+          if (waitStart == 0) {
+            waitStart = now;
+            lastLog = now;
+            logger.warn(
+                "Waiting for other tasks to free up {} bytes of Comet shuffle pool memory", size);
+          } else if (now - waitStart >= memoryWaitTimeoutMs) {
+            logger.warn(
+                "Giving up after waiting {} ms for {} bytes of Comet shuffle pool memory "
+                    + "(see {})",
+                now - waitStart,
+                size,
+                CometConf$.MODULE$.COMET_SHUFFLE_JVM_MEMORY_WAIT_TIMEOUT().key());
+            throw e;
+          } else if (now - lastLog >= WAIT_LOG_INTERVAL_MS) {
+            lastLog = now;
+            logger.warn(
+                "Still waiting ({} ms so far) for {} bytes of Comet shuffle pool memory; "
+                    + "{} bytes free, {} thread(s) waiting",
+                now - waitStart,
+                size,
+                totalMemory - allocatedMemory,
+                waitingThreads.size());
+          }
+          try {
+            wait(
+                Math.max(
+                    1L,
+                    Math.min(TASK_KILL_POLL_INTERVAL_MS, memoryWaitTimeoutMs - (now - waitStart))));
+          } catch (InterruptedException ie) {
+            Thread.currentThread().interrupt();
+            // Not an allocation failure: stay non-fatal so that an intentional task kill is
+            // classified as TaskKilled rather than ExceptionFailure (Spark's killed-task handler
+            // only matches `InterruptedException | NonFatal(_)`).
+            throw new RuntimeException(
+                "Interrupted while waiting for Comet shuffle pool memory", ie);
+          }
+        }
+      }
+    } finally {
+      if (waitingThreads.remove(self) != null) {
+        notifyAll();
+      }
+    }
+  }
+
+  private long retainedByWaitingThreads() {
+    long retained = 0;
+    for (Thread thread : waitingThreads.keySet()) {
+      retained += retainedMemory.getOrDefault(thread, 0L);
+    }
+    return retained;
+  }
+
+  private boolean anyWaiterCanProceed() {
+    long free = totalMemory - allocatedMemory;
+    for (long requested : waitingThreads.values()) {
+      if (requested <= free) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   private synchronized MemoryBlock allocateMemoryBlock(long required) {
     if (required > TaskMemoryManager.MAXIMUM_PAGE_SIZE_BYTES) {
       throw new TooLargePageException(required);
@@ -153,6 +293,8 @@ private synchronized MemoryBlock allocateMemoryBlock(long required) {
     block.pageNumber = pageNumber;
     pageTable[pageNumber] = block;
     allocatedPages.set(pageNumber);
+    pageOwners[pageNumber] = Thread.currentThread();
+    retainedMemory.merge(Thread.currentThread(), got, Long::sum);
 
     return block;
   }
@@ -166,11 +308,19 @@ public synchronized long free(MemoryBlock block) {
     long blockSize = block.size();
     allocatedMemory -= blockSize;
 
+    Thread owner = pageOwners[block.pageNumber];
+    pageOwners[block.pageNumber] = null;
+    if (owner != null) {
+      retainedMemory.computeIfPresent(owner, (t, v) -> v - blockSize <= 0 ? null : v - blockSize);
+    }
+
     pageTable[block.pageNumber] = null;
     allocatedPages.clear(block.pageNumber);
     block.pageNumber = MemoryBlock.FREED_IN_TMM_PAGE_NUMBER;
 
     allocator.free(block);
+    // Wake up tasks waiting in `allocateBlocking`.
+    notifyAll();
     return blockSize;
   }
 

spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocatorTrait.java

@@ -33,6 +33,17 @@ protected CometShuffleMemoryAllocatorTrait(
 
   public abstract MemoryBlock allocate(long required);
 
+  /**
+   * Like {@link #allocate(long)}, but may wait for memory freed by other tasks instead of failing
+   * immediately. Callers must first spill buffered data they can cheaply release; an implementation
+   * that waits must account for any memory the caller retains while blocked. The default
+   * implementation does not wait: `CometUnifiedShuffleMemoryAllocator` delegates to Spark's memory
+   * manager, which already arbitrates memory between tasks.
+   */
+  public MemoryBlock allocateBlocking(long required) {
+    return allocate(required);
+  }
+
   public abstract long free(MemoryBlock block);
 
   public abstract long getOffsetInPage(long pagePlusOffsetAddress);

spark/src/main/java/org/apache/spark/shuffle/sort/SpillSorter.java

@@ -131,13 +131,28 @@ public SpillSorter(
               + "https://github.com/apache/arrow-datafusion-comet?tab=readme-ov-file#enable-comet-shuffle",
           e);
     }
-    sorterArray = allocator.allocateArray(initialSize);
-    this.inMemSorter.expandPointerArray(sorterArray);
-
-    this.allocatedPages = new LinkedList<>();
-
-    this.nativeLib = new Native();
-    this.dataTypes = serializeSchema(schema);
+    boolean adopted = false;
+    try {
+      sorterArray = allocator.allocateArray(initialSize);
+      this.inMemSorter.expandPointerArray(sorterArray);
+      adopted = true;
+
+      this.allocatedPages = new LinkedList<>();
+
+      this.nativeLib = new Native();
+      this.dataTypes = serializeSchema(schema);
+    } catch (Throwable t) {
+      // This writer is never handed to Spark when its constructor fails, so nothing else could
+      // reclaim what was allocated so far; free it here to keep the shared pool leak-free.
+      // Before adoption the sorter still owns only its initial one-entry array; after adoption
+      // it owns `sorterArray`, and `free()` releases whichever it holds.
+      if (!adopted && sorterArray != null) {
+        allocator.freeArray(sorterArray);
+        sorterArray = null;
+      }
+      this.inMemSorter.free();
+      throw t;
+    }
   }
 
   /** Frees allocated memory pages of this writer */

spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/CometBypassMergeSortShuffleWriter.java

@@ -24,6 +24,7 @@
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.channels.FileChannel;
+import java.util.LinkedList;
 import java.util.Optional;
 import javax.annotation.Nullable;
 
@@ -176,6 +177,10 @@ public void write(Iterator<Product2<K, V>> records) throws IOException {
                   CometShuffleExternalSorter.MAXIMUM_PAGE_SIZE_BYTES,
                   memoryManager.pageSizeBytes()));
 
+      // This task's disk writers. Under memory pressure a writer spills its sibling writers in
+      // this list, never those of other tasks.
+      final LinkedList<CometDiskBlockWriter> taskWriters = new LinkedList<>();
+
       // Allocate the disk writers, and open the files that we'll be writing to
       for (int i = 0; i < numPartitions; i++) {
         final Tuple2<TempShuffleBlockId, File> tempShuffleBlockIdPlusFile =
@@ -190,7 +195,8 @@ public void write(Iterator<Product2<K, V>> records) throws IOException {
                 schema,
                 writeMetrics,
                 conf,
-                tracingEnabled);
+                tracingEnabled,
+                taskWriters);
         if (partitionChecksums.length > 0) {
           writer.setChecksum(partitionChecksums[i]);
           writer.setChecksumAlgo(checksumAlgorithm);
@@ -253,7 +259,13 @@ public void write(Iterator<Product2<K, V>> records) throws IOException {
       // TODO: We probably can move checksum generation here when concatenating partition files
       partitionLengths = writePartitionedData(mapOutputWriter);
       mapStatus = MapStatusHelper.apply(blockManager.shuffleServerId(), partitionLengths, mapId);
-    } catch (Exception e) {
+    } catch (Throwable e) {
+      // Spark only calls stop(false) when write() throws an Exception; a fatal error such as
+      // SparkOutOfMemoryError skips it, and the buffered pages are invisible to Spark's
+      // task-memory cleanup (in on-heap mode they live in an executor-shared pool). Free them
+      // and delete their temp files here so they cannot starve other tasks' allocations or leak
+      // disk space.
+      cleanupPartitionWriters(e);
       try {
         mapOutputWriter.abort(e);
       } catch (Exception e2) {
@@ -264,6 +276,40 @@ public void write(Iterator<Product2<K, V>> records) throws IOException {
     }
   }
 
+  private void cleanupPartitionWriters(@Nullable Throwable failure) {
+    if (partitionWriters == null) {
+      return;
+    }
+    try {
+      for (CometDiskBlockWriter writer : partitionWriters) {
+        if (writer == null) {
+          continue;
+        }
+        try {
+          writer.freeMemory();
+        } catch (Exception e) {
+          logger.error("Failed to free memory of partition writer", e);
+          if (failure != null) {
+            failure.addSuppressed(e);
+          }
+        }
+        try {
+          File file = writer.getFile();
+          if (file.exists() && !file.delete()) {
+            logger.error("Error while deleting file {}", file.getAbsolutePath());
+          }
+        } catch (Exception e) {
+          logger.error("Failed to delete file of partition writer", e);
+          if (failure != null) {
+            failure.addSuppressed(e);
+          }
+        }
+      }
+    } finally {
+      partitionWriters = null;
+    }
+  }
+
   @Override
   public long[] getPartitionLengths() {
     return partitionLengths;
@@ -366,20 +412,7 @@ public Option<MapStatus> stop(boolean success) {
         return Option.apply(mapStatus);
       } else {
         // The map task failed, so delete our output data.
-        if (partitionWriters != null) {
-          try {
-            for (CometDiskBlockWriter writer : partitionWriters) {
-              writer.freeMemory();
-
-              File file = writer.getFile();
-              if (!file.delete()) {
-                logger.error("Error while deleting file {}", file.getAbsolutePath());
-              }
-            }
-          } finally {
-            partitionWriters = null;
-          }
-        }
+        cleanupPartitionWriters(null);
         return None$.empty();
       }
     }