This is an automated email from the ASF dual-hosted git repository.
agrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 517fde2a6 fix: Avoid double free in CometUnifiedShuffleMemoryAllocator
(#2122)
517fde2a6 is described below
commit 517fde2a62bd917dc41b5a172735810ed8f520e5
Author: Andy Grove <agr...@apache.org>
AuthorDate: Mon Aug 11 18:15:49 2025 -0600
fix: Avoid double free in CometUnifiedShuffleMemoryAllocator (#2122)
---
.../shuffle/comet/CometBoundedShuffleMemoryAllocator.java | 11 +++++++----
.../spark/shuffle/comet/CometShuffleMemoryAllocatorTrait.java | 2 +-
.../shuffle/comet/CometUnifiedShuffleMemoryAllocator.java | 9 ++++++++-
.../apache/spark/sql/comet/execution/shuffle/SpillWriter.java | 3 +--
4 files changed, 17 insertions(+), 8 deletions(-)
diff --git
a/spark/src/main/java/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocator.java
b/spark/src/main/java/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocator.java
index 54e9dc684..dae55c04b 100644
---
a/spark/src/main/java/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocator.java
+++
b/spark/src/main/java/org/apache/spark/shuffle/comet/CometBoundedShuffleMemoryAllocator.java
@@ -146,18 +146,21 @@ public final class CometBoundedShuffleMemoryAllocator
extends CometShuffleMemory
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;
}
/**
diff --git
a/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocatorTrait.java
b/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocatorTrait.java
index 6831396b3..36fa9d2ff 100644
---
a/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocatorTrait.java
+++
b/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocatorTrait.java
@@ -33,7 +33,7 @@ public abstract class CometShuffleMemoryAllocatorTrait
extends MemoryConsumer {
public abstract MemoryBlock allocate(long required);
- public abstract void free(MemoryBlock block);
+ public abstract long free(MemoryBlock block);
public abstract long getOffsetInPage(long pagePlusOffsetAddress);
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 917d96f0f..aa8de6f17 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
@@ -57,8 +57,15 @@ public final class CometUnifiedShuffleMemoryAllocator
extends CometShuffleMemory
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;
}
/**
diff --git
a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java
b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java
index c8f845c5d..044c7842f 100644
---
a/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java
+++
b/spark/src/main/java/org/apache/spark/sql/comet/execution/shuffle/SpillWriter.java
@@ -218,8 +218,7 @@ public abstract class SpillWriter {
public long freeMemory() {
long freed = 0L;
for (MemoryBlock block : allocatedPages) {
- freed += block.size();
- allocator.free(block);
+ freed += allocator.free(block);
}
allocatedPages.clear();
currentPage = null;
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@datafusion.apache.org
For additional commands, e-mail: commits-h...@datafusion.apache.org
