fix

viirya · Nov 13, 2024 · 0f3989d · 0f3989d
1 parent 5425868
commit 0f3989d
Showing 1 changed file with 16 additions and 13 deletions.
diff --git a/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocator.java b/spark/src/main/java/org/apache/spark/shuffle/comet/CometShuffleMemoryAllocator.java
@@ -52,22 +52,25 @@ public static synchronized CometShuffleMemoryAllocatorTrait getInstance(
         (boolean)
             CometConf$.MODULE$.COMET_COLUMNAR_SHUFFLE_UNIFIED_MEMORY_ALLOCATOR_IN_TEST().get();
 
-    if (INSTANCE == null) {
-      if (isSparkTesting && !useUnifiedMemAllocator) {
+    if (isSparkTesting && !useUnifiedMemAllocator) {
+      if (INSTANCE == null) {
+        // CometTestShuffleMemoryAllocator handles pages by itself so it can be a singleton.
         INSTANCE = new CometTestShuffleMemoryAllocator(conf, taskMemoryManager, pageSize);
-      } else {
-        if (taskMemoryManager.getTungstenMemoryMode() != MemoryMode.OFF_HEAP) {
-          throw new IllegalArgumentException(
-              "CometShuffleMemoryAllocator should be used with off-heap "
-                  + "memory mode, but got "
-                  + taskMemoryManager.getTungstenMemoryMode());
-        }
-
-        INSTANCE = new CometShuffleMemoryAllocator(taskMemoryManager, pageSize);
       }
-    }
 
-    return INSTANCE;
+      return INSTANCE;
+    } else {
+      if (taskMemoryManager.getTungstenMemoryMode() != MemoryMode.OFF_HEAP) {
+        throw new IllegalArgumentException(
+            "CometShuffleMemoryAllocator should be used with off-heap "
+                + "memory mode, but got "
+                + taskMemoryManager.getTungstenMemoryMode());
+      }
+
+      // CometShuffleMemoryAllocator stores pages in TaskMemoryManager which is not singleton,
+      // but one instance per task. So we need to create a new instance for each task.
+      return new CometShuffleMemoryAllocator(taskMemoryManager, pageSize);
+    }
   }
 
   CometShuffleMemoryAllocator(TaskMemoryManager taskMemoryManager, long pageSize) {