diff --git a/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/DynamicOffHeapSizingMemoryTarget.java b/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/DynamicOffHeapSizingMemoryTarget.java index b1664223aaa0..3124f78fb79a 100644 --- a/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/DynamicOffHeapSizingMemoryTarget.java +++ b/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/DynamicOffHeapSizingMemoryTarget.java @@ -130,7 +130,11 @@ public class DynamicOffHeapSizingMemoryTarget implements MemoryTarget, KnownName private final String name = MemoryTargetUtil.toUniqueName("DynamicOffHeapSizing"); private final SimpleMemoryUsageRecorder recorder = new SimpleMemoryUsageRecorder(); - public DynamicOffHeapSizingMemoryTarget() {} + private final MemoryTarget target; + + public DynamicOffHeapSizingMemoryTarget(MemoryTarget target) { + this.target = target; + } @Override public long borrow(long size) { @@ -197,6 +201,7 @@ public void run() { USED_OFF_HEAP_BYTES.addAndGet(size); recorder.inc(size); + target.borrow(size); return size; } @@ -204,6 +209,7 @@ public void run() { public long repay(long size) { USED_OFF_HEAP_BYTES.addAndGet(-size); recorder.inc(-size); + target.repay(size); return size; } @@ -227,6 +233,10 @@ public MemoryUsageStats stats() { return recorder.toStats(); } + public MemoryTarget target() { + return target; + } + public static boolean isJava9OrLater() { String spec = System.getProperty("java.specification.version", "1.8"); // "1.8" → 8, "9" → 9, "11" → 11, etc. diff --git a/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/MemoryTargets.java b/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/MemoryTargets.java index b1e41bc9eba8..e6cab4b6a879 100644 --- a/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/MemoryTargets.java +++ b/gluten-core/src/main/java/org/apache/gluten/memory/memtarget/MemoryTargets.java @@ -52,7 +52,7 @@ public static MemoryTarget overAcquire( @Experimental public static MemoryTarget dynamicOffHeapSizingIfEnabled(MemoryTarget memoryTarget) { if (GlutenCoreConfig.get().dynamicOffHeapSizingEnabled()) { - return new DynamicOffHeapSizingMemoryTarget(); + return new DynamicOffHeapSizingMemoryTarget(memoryTarget); } return memoryTarget; diff --git a/gluten-core/src/main/scala/org/apache/spark/memory/SparkMemoryUtil.scala b/gluten-core/src/main/scala/org/apache/spark/memory/SparkMemoryUtil.scala index ae287cef3fe6..a902c448fe47 100644 --- a/gluten-core/src/main/scala/org/apache/spark/memory/SparkMemoryUtil.scala +++ b/gluten-core/src/main/scala/org/apache/spark/memory/SparkMemoryUtil.scala @@ -132,6 +132,7 @@ object SparkMemoryUtil { prettyPrintStats( "Dynamic off-heap sizing memory target stats: ", dynamicOffHeapSizingMemoryTarget) + dynamicOffHeapSizingMemoryTarget.target().accept(this) } override def visit(retryOnOomMemoryTarget: RetryOnOomMemoryTarget): String = {