(datafusion-comet) branch main updated: fix: Fix data race in memory profiling (#1727)

[ag...@apache.org]

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 108f8377b fix: Fix data race in memory profiling (#1727)
108f8377b is described below

commit 108f8377b9976ce4e3319f74eae0e08f6bf28fe8
Author: Andy Grove <ag...@apache.org>
AuthorDate: Mon May 12 19:04:11 2025 -0600

    fix: Fix data race in memory profiling (#1727)
---
 .../org/apache/spark/CometTaskMemoryManager.java     | 20 +++++++++++++++-----
 1 file changed, 15 insertions(+), 5 deletions(-)

diff --git a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
index 7646fcdfc..219bc67c2 100644
--- a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
+++ b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
@@ -20,6 +20,10 @@
 package org.apache.spark;
 
 import java.io.IOException;
+import java.util.concurrent.atomic.AtomicLong;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import org.apache.spark.memory.MemoryConsumer;
 import org.apache.spark.memory.MemoryMode;
@@ -30,40 +34,46 @@ import org.apache.spark.memory.TaskMemoryManager;
  * memory manager. This assumes Spark's off-heap memory mode is enabled.
  */
 public class CometTaskMemoryManager {
+
+  private static final Logger logger = LoggerFactory.getLogger(CometTaskMemoryManager.class);
+
   /** The id uniquely identifies the native plan this memory manager is associated to */
   private final long id;
 
   private final TaskMemoryManager internal;
   private final NativeMemoryConsumer nativeMemoryConsumer;
-  private long used;
+  private final AtomicLong used = new AtomicLong();
 
   public CometTaskMemoryManager(long id) {
     this.id = id;
     this.internal = TaskContext$.MODULE$.get().taskMemoryManager();
     this.nativeMemoryConsumer = new NativeMemoryConsumer();
-    this.used = 0;
   }
 
   // Called by Comet native through JNI.
   // Returns the actual amount of memory (in bytes) granted.
   public long acquireMemory(long size) {
     long acquired = internal.acquireExecutionMemory(size, nativeMemoryConsumer);
+    used.addAndGet(acquired);
     if (acquired < size) {
       // If memory manager is not able to acquire the requested size, log memory usage
       internal.showMemoryUsage();
     }
-    used += acquired;
     return acquired;
   }
 
   // Called by Comet native through JNI
   public void releaseMemory(long size) {
-    used -= size;
+    long newUsed = used.addAndGet(-size);
+    if (newUsed < 0) {
+      logger.error(
+          "Used memory is negative: " + newUsed + " after releasing memory chunk of: " + size);
+    }
     internal.releaseExecutionMemory(size, nativeMemoryConsumer);
   }
 
   public long getUsed() {
-    return used;
+    return used.get();
   }
 
   /**


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org