fix: [branch-0.9] Avoid double free in CometUnifiedShuffleMemoryAllocator (#2201)

Author: andygrove

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

@@ -146,18 +146,21 @@ private synchronized MemoryBlock allocateMemoryBlock(long required) {
     return block;
   }
 
-  public synchronized void free(MemoryBlock block) {
-    if (block.pageNumber == MemoryBlock.FREED_IN_ALLOCATOR_PAGE_NUMBER) {
+  public synchronized long free(MemoryBlock block) {
+    if (block.pageNumber == MemoryBlock.FREED_IN_ALLOCATOR_PAGE_NUMBER
+        || block.pageNumber == MemoryBlock.FREED_IN_TMM_PAGE_NUMBER) {
       // Already freed block
-      return;
+      return 0;
     }
-    allocatedMemory -= block.size();
+    long blockSize = block.size();
+    allocatedMemory -= blockSize;
 
     pageTable[block.pageNumber] = null;
     allocatedPages.clear(block.pageNumber);
     block.pageNumber = MemoryBlock.FREED_IN_TMM_PAGE_NUMBER;
 
     allocator.free(block);
+    return blockSize;
   }
 
   /**

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

@@ -33,7 +33,7 @@ protected CometShuffleMemoryAllocatorTrait(
 
   public abstract MemoryBlock allocate(long required);
 
-  public abstract void free(MemoryBlock block);
+  public abstract long free(MemoryBlock block);
 
   public abstract long getOffsetInPage(long pagePlusOffsetAddress);
 

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

@@ -57,8 +57,15 @@ public synchronized MemoryBlock allocate(long required) {
     return this.allocatePage(required);
   }
 
-  public synchronized void free(MemoryBlock block) {
+  public synchronized long free(MemoryBlock block) {
+    if (block.pageNumber == MemoryBlock.FREED_IN_ALLOCATOR_PAGE_NUMBER
+        || block.pageNumber == MemoryBlock.FREED_IN_TMM_PAGE_NUMBER) {
+      // Already freed block
+      return 0;
+    }
+    long blockSize = block.size();
     this.freePage(block);
+    return blockSize;
   }
 
   /**

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

@@ -218,8 +218,7 @@ protected long doSpilling(
   public long freeMemory() {
     long freed = 0L;
     for (MemoryBlock block : allocatedPages) {
-      freed += block.size();
-      allocator.free(block);
+      freed += allocator.free(block);
     }
     allocatedPages.clear();
     currentPage = null;