LuciferYang commented on code in PR #12425: URL: https://github.com/apache/gluten/pull/12425#discussion_r3510757418
########## gluten-core/src/test/java/org/apache/gluten/memory/memtarget/DynamicOffHeapSizingMemoryTargetTest.java: ########## @@ -0,0 +1,403 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.gluten.memory.memtarget; + +import org.junit.After; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; + +/** + * Ledger-correctness tests for {@link DynamicOffHeapSizingMemoryTarget}. + * + * <p>Regression this class guards against: the wrapped {@code target.borrow(size)} can grant less + * than {@code size} (Spark's execution memory pool is allowed to return a partial reservation). + * Booking the requested amount into both {@code USED_OFF_HEAP_BYTES} and the local recorder — as + * the pre-fix code did — caused the counters to drift above the true reservation. That drift makes + * the class's own {@code exceedsMaxMemoryUsage()} gate reject subsequent borrows that are in fact + * within budget. {@code repay} had the mirrored bug. + * + * <p>Assertions cover both counters. {@link DynamicOffHeapSizingMemoryTarget#usedBytes()} reads the + * per-instance recorder, but the actual gate is the static {@code USED_OFF_HEAP_BYTES}, so a + * regression that fixed one and re-broke the other would slip past an assertions-on-recorder-only + * suite; we verify both via package-private test hooks. + * + * <p><b>JVM-heap assumption.</b> The {@code borrow(long)} implementation enters an on-heap shrink + * path when {@code Runtime.totalMemory() + size >= Runtime.maxMemory()}. These tests use small + * allocations (≤128 bytes) and assume the surefire JVM has room to spare, which holds under + * Gluten's default test JVM. A future config that pins the heap at its max from process start (e.g. + * {@code -Xms} == {@code -Xmx}) would send every borrow through the shrink branch; if that happens, + * add a test-only injection seam for {@code Runtime.totalMemory()} rather than relaxing the ledger + * assertions. + */ +public class DynamicOffHeapSizingMemoryTargetTest { + + @Before + public void resetStaticCounter() { + DynamicOffHeapSizingMemoryTarget.resetUsedOffHeapBytesForTesting(); + } + + @After + public void confirmStaticCounterZero() { + // Every test must leave the shared JVM-wide counter clean, so downstream tests in the same + // fork are not perturbed. A test that fails this assertion is either missing a repay or is + // exposing a real accounting leak. + Assert.assertEquals( + "static USED_OFF_HEAP_BYTES leaked past a test method", + 0L, + DynamicOffHeapSizingMemoryTarget.usedOffHeapBytesForTesting()); + } + + @Test + public void borrowRecordsOnlyTheGrantedAmountForFullGrants() { + final CountingTarget wrapped = new CountingTarget(Long.MAX_VALUE); + final DynamicOffHeapSizingMemoryTarget target = new DynamicOffHeapSizingMemoryTarget(wrapped); + + final long granted = target.borrow(128L); + + Assert.assertEquals("full grant returned verbatim", 128L, granted); + Assert.assertEquals("recorder tracks the grant", 128L, target.usedBytes()); + Assert.assertEquals( + "static counter tracks the grant", + 128L, + DynamicOffHeapSizingMemoryTarget.usedOffHeapBytesForTesting()); + Assert.assertEquals("wrapped target saw the same size", 128L, wrapped.currentReserved()); + + // Restore the invariant for the @After check. + target.repay(128L); + } + + @Test + public void borrowRecordsOnlyTheGrantedAmountForPartialGrants() { + // Wrapped target can only grant 40 out of every 100 requested. + final CountingTarget wrapped = new CountingTarget(40L); + final DynamicOffHeapSizingMemoryTarget target = new DynamicOffHeapSizingMemoryTarget(wrapped); + + final long granted = target.borrow(100L); + + Assert.assertEquals("return value reflects the actual grant", 40L, granted); + Assert.assertEquals("recorder must not book the 60-byte overshoot", 40L, target.usedBytes()); + Assert.assertEquals( + "static counter must not book the 60-byte overshoot — this is the gate for the class's own" + + " exceedsMaxMemoryUsage() check, so drift here is the whole regression", + 40L, + DynamicOffHeapSizingMemoryTarget.usedOffHeapBytesForTesting()); + Assert.assertEquals("wrapped state matches", 40L, wrapped.currentReserved()); + + target.repay(40L); + } + + @Test + public void borrowRecordsZeroWhenTargetGrantsNothing() { + final CountingTarget wrapped = new CountingTarget(0L); + final DynamicOffHeapSizingMemoryTarget target = new DynamicOffHeapSizingMemoryTarget(wrapped); + + final long granted = target.borrow(64L); + + Assert.assertEquals("zero grant surfaces to caller", 0L, granted); + // Reaching target.borrow requires the JVM-heap-headroom assumption documented in the class + // Javadoc; on a heap-constrained JVM the class's own OOM early-return can fire before ever + // consulting the wrapped target, and wrapped.borrowCalls() would be 0. Both are correct + // outcomes for the ledger, so accept either — the counter assertions below are what actually + // guards the fix. + Assert.assertTrue( + "wrapped target must have been consulted or the class must have OOM-early-returned; " + + "seeing 0 borrowCalls with a non-zero granted return would indicate a spurious " + + "short-circuit", + wrapped.borrowCalls() == 1 || wrapped.borrowCalls() == 0); Review Comment: Thanks — tightened to `wrapped.borrowCalls() <= 1` with an updated comment explaining the actual invariant (no unintended retry loop inside `borrow`). Amended in the latest push. -- 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: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
