andygrove commented on code in PR #1727:
URL: https://github.com/apache/datafusion-comet/pull/1727#discussion_r2084735443
##########
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:
I changed this to use error logging rather than a warning. I'd prefer not to
throw an exception because many jobs will run to completion without issue even
if this variable does go negative.
--
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
