parthchandra commented on code in PR #1727:
URL: https://github.com/apache/datafusion-comet/pull/1727#discussion_r2082639273
##########
spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java:
##########
@@ -30,36 +34,41 @@
* 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 += acquired;
+ used.addAndGet(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.warn("Used memory is negative: " + newUsed);
Review Comment:
Should probably make this an error message (or maybe even throw an exception
since this is a pretty fatal condition?)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscr...@datafusion.apache.org
For additional commands, e-mail: github-h...@datafusion.apache.org
