Update TLAB logic to JDK 25+25.

Fix gcWaste computation.

Author: christianhaeubl

substratevm/src/com.oracle.svm.core.genscavenge/src/com/oracle/svm/core/genscavenge/ThreadLocalAllocation.java

@@ -25,16 +25,11 @@
 package com.oracle.svm.core.genscavenge;
 
 import static com.oracle.svm.core.Uninterruptible.CALLED_FROM_UNINTERRUPTIBLE_CODE;
-import static com.oracle.svm.core.genscavenge.TlabSupport.computeMinSizeOfNewTlab;
-import static com.oracle.svm.core.genscavenge.TlabSupport.computeSizeOfNewTlab;
-import static com.oracle.svm.core.genscavenge.TlabSupport.fillTlab;
-import static com.oracle.svm.core.genscavenge.TlabSupport.recordSlowAllocation;
-import static com.oracle.svm.core.genscavenge.TlabSupport.retireTlabBeforeAllocation;
-import static com.oracle.svm.core.genscavenge.TlabSupport.shouldRetainTlab;
 import static com.oracle.svm.core.graal.snippets.SubstrateAllocationSnippets.TLAB_END_IDENTITY;
 import static com.oracle.svm.core.graal.snippets.SubstrateAllocationSnippets.TLAB_START_IDENTITY;
 import static com.oracle.svm.core.graal.snippets.SubstrateAllocationSnippets.TLAB_TOP_IDENTITY;
 
+import org.graalvm.nativeimage.CurrentIsolate;
 import org.graalvm.nativeimage.IsolateThread;
 import org.graalvm.nativeimage.Platform;
 import org.graalvm.nativeimage.Platforms;
@@ -371,6 +366,7 @@ private static Pointer allocateRawMemory(UnsignedWord size, BooleanPointer alloc
         return allocateRawMemoryOutsideTlab(size, allocatedOutsideTlab);
     }
 
+    // TEMP (chaeubl): move this to TlabSupport and make most TlabSupport methods private
     @BasedOnJDKFile("https://github.com/openjdk/jdk/blob/jdk-25+25/src/hotspot/share/gc/shared/memAllocator.cpp#L257-L329")
     @Uninterruptible(reason = "Holds uninitialized memory.")
     private static Pointer allocateRawMemoryInTlabSlow(UnsignedWord size) {
@@ -380,19 +376,18 @@ private static Pointer allocateRawMemoryInTlabSlow(UnsignedWord size) {
          * Retain tlab and allocate object as an heap allocation if the amount free in the tlab is
          * too large to discard.
          */
-        if (shouldRetainTlab(tlab)) {
-            recordSlowAllocation();
+        if (TlabSupport.shouldRetainTlab(tlab)) {
+            TlabSupport.recordSlowAllocation();
             return Word.nullPointer();
         }
 
-        /*
-         * Discard tlab and allocate a new one. To minimize fragmentation, the last tlab may be
-         * smaller than the rest.
-         */
-        UnsignedWord newTlabSize = computeSizeOfNewTlab(size);
+        /* Discard tlab and allocate a new one. */
 
-        retireTlabBeforeAllocation();
+        TlabSupport.recordRefillWaste();
+        TlabSupport.retireTlab(CurrentIsolate.getCurrentThread(), false);
 
+        /* To minimize fragmentation, the last tlab may be smaller than the rest. */
+        UnsignedWord newTlabSize = TlabSupport.computeSizeOfNewTlab(size);
         if (newTlabSize.equal(0)) {
             return Word.nullPointer();
         }
@@ -401,8 +396,7 @@ private static Pointer allocateRawMemoryInTlabSlow(UnsignedWord size) {
          * Allocate a new TLAB requesting newTlabSize. Any size between minimal and newTlabSize is
          * accepted.
          */
-
-        UnsignedWord computedMinSize = computeMinSizeOfNewTlab(size);
+        UnsignedWord computedMinSize = TlabSupport.computeMinSizeOfNewTlab(size);
 
         WordPointer allocatedTlabSize = StackValue.get(WordPointer.class);
         Pointer memory = YoungGeneration.getHeapAllocation().allocateNewTlab(computedMinSize, newTlabSize, allocatedTlabSize);
@@ -412,7 +406,7 @@ private static Pointer allocateRawMemoryInTlabSlow(UnsignedWord size) {
         }
         assert Word.unsigned(0).notEqual(allocatedTlabSize.read()) : "Allocation succeeded but actual size not updated.";
 
-        fillTlab(memory, memory.add(size), allocatedTlabSize);
+        TlabSupport.fillTlab(memory, memory.add(size), allocatedTlabSize);
         return memory;
     }
 

substratevm/src/com.oracle.svm.core.genscavenge/src/com/oracle/svm/core/genscavenge/TlabSupport.java

@@ -70,7 +70,7 @@
  * and another thread allocates nothing. Between GC 23 and GC 24 the allocation behaviour of these
  * two threads switches. The allocation average and the TLAB size adapt to the new allocation
  * behavior.
- * 
+ *
  * <pre>
  * +-----+---------------------------------------++---------------------------------------+
  * | #GC | Thread 1                              || Thread 2                              |
@@ -87,7 +87,7 @@
  * |  29 |           0B |   270,44kB |    5,41kB ||       3,55MB |     3,28MB |   67,14kB |
  * +-----+--------------+------------+-----------++--------------+------------+-----------+
  * </pre>
- *
+ * <p>
  * A thread allocating a very large amount of memory will also have a high
  * {@link #allocatedBytesAvg}. If such a thread later changes its allocation behaviour and only
  * allocates a small amount of memory the {@link #allocatedBytesAvg} starts decreasing with the next
@@ -109,25 +109,23 @@ public class TlabSupport {
     @BasedOnJDKFile("https://github.com/openjdk/jdk/blob/jdk-23-ga/src/hotspot/share/gc/shared/tlab_globals.hpp#L82-L85")//
     private static final long TLAB_WASTE_INCREMENT = 4;
 
-    // The desired size of the TLAB, including the reserve for filling the unused memory.
+    /* The desired size of the TLAB, including the reserve for filling the unused memory. */
     private static final FastThreadLocalWord<UnsignedWord> desiredSize = FastThreadLocalFactory.createWord("TlabSupport.desiredSize");
-
     private static final FastThreadLocalWord<UnsignedWord> tlabAllocatedAlignedBytesBeforeLastGC = FastThreadLocalFactory.createWord("TlabSupport.tlabAllocatedAlignedBytesBeforeLastGC");
-
     private static final FastThreadLocalInt numberOfRefills = FastThreadLocalFactory.createInt("TlabSupport.numberOfRefills");
     private static final FastThreadLocalInt refillWaste = FastThreadLocalFactory.createInt("TlabSupport.refillWaste");
     private static final FastThreadLocalInt gcWaste = FastThreadLocalFactory.createInt("TlabSupport.gcWaste");
 
-    // Average of allocated bytes in TLABs of this thread.
+    /* Average of allocated bytes in TLABs of this thread. */
     private static final FastThreadLocalBytes<AdaptiveWeightedAverageStruct.Data> allocatedBytesAvg = FastThreadLocalFactory
                     .createBytes(() -> SizeOf.get(AdaptiveWeightedAverageStruct.Data.class), "TlabSupport.allocatedBytesAvg");
 
-    // Hold onto the TLAB if availableTlabMemory() is larger than this.
+    /* Hold onto the TLAB if availableTlabMemory() is larger than this. */
     private static final FastThreadLocalWord<UnsignedWord> refillWasteLimit = FastThreadLocalFactory.createWord("TlabSupport.refillWasteLimit");
 
     private static final FastThreadLocalInt slowAllocations = FastThreadLocalFactory.createInt("TlabSupport.slowAllocations");
 
-    // Expected number of refills between GCs.
+    /* Expected number of refills between GCs. */
     private static UnsignedWord targetRefills = Word.unsigned(1);
 
     private static boolean initialized;
@@ -160,9 +158,11 @@ public static void initialize(IsolateThread thread) {
         resetStatistics(thread);
     }
 
+    @BasedOnJDKFile("https://github.com/openjdk/jdk/blob/jdk-25+25/src/hotspot/share/runtime/thread.cpp#L168-L174")
     @BasedOnJDKFile("https://github.com/openjdk/jdk/blob/jdk-23-ga/src/hotspot/share/gc/shared/threadLocalAllocBuffer.cpp#L183-L195")
     @Uninterruptible(reason = "Accesses TLAB")
     static void fillTlab(Pointer start, Pointer top, WordPointer newSize) {
+        /* Fill the TLAB. */
         numberOfRefills.set(numberOfRefills.get() + 1);
 
         Pointer hardEnd = start.add(newSize.read());
@@ -172,39 +172,38 @@ static void fillTlab(Pointer start, Pointer top, WordPointer newSize) {
 
         initialize(getTlab(), start, top, end);
 
-        // Reset amount of internal fragmentation
+        /* Reset amount of internal fragmentation. */
         refillWasteLimit.set(initialRefillWasteLimit());
     }
 
     @BasedOnJDKFile("https://github.com/openjdk/jdk/blob/jdk-25+25/src/hotspot/share/gc/shared/threadLocalAllocBuffer.cpp#L143-L145")
     @Uninterruptible(reason = "Accesses TLAB")
-    static void retireTlabBeforeAllocation() {
+    static void recordRefillWaste() {
         long availableTlabMemory = availableTlabMemory(getTlab()).rawValue();
         refillWaste.set(refillWaste.get() + UninterruptibleUtils.NumUtil.safeToInt(availableTlabMemory));
-        retireCurrentTlab(CurrentIsolate.getCurrentThread(), false);
     }
 
+    @BasedOnJDKFile("https://github.com/openjdk/jdk/blob/jdk-25+25/src/hotspot/share/runtime/thread.cpp#L157-L166")
     @BasedOnJDKFile("https://github.com/openjdk/jdk/blob/jdk-25+25/src/hotspot/share/gc/shared/threadLocalAllocBuffer.cpp#L131-L141")
     @Uninterruptible(reason = "Accesses TLAB")
-    private static void retireCurrentTlab(IsolateThread thread, boolean calculateStats) {
+    static void retireTlab(IsolateThread thread, boolean calculateStats) {
+        /* Sampling and serviceability support. */
         ThreadLocalAllocation.Descriptor tlab = getTlab(thread);
-
         if (tlab.getAllocationEnd(TLAB_END_IDENTITY).isNonNull()) {
-            assert checkInvariants(tlab);
-
-            UnsignedWord usedTlabSize = getUsedTlabSize(tlab);
-            allocatedAlignedBytes.set(thread, allocatedAlignedBytes.get(thread).add(usedTlabSize));
-            insertFiller(tlab);
-            initialize(tlab, Word.nullPointer(), Word.nullPointer(), Word.nullPointer());
+            UnsignedWord usedBytes = getUsedTlabSize(tlab);
+            allocatedAlignedBytes.set(thread, allocatedAlignedBytes.get(thread).add(usedBytes));
         }
 
-        /*
-         * Collect statistics after the TLAB has been retired. Otherwise, the current TLAB is
-         * excluded from the statistics.
-         */
+        /* Retire the TLAB. */
         if (calculateStats) {
             accumulateAndResetStatistics(thread);
         }
+
+        if (tlab.getAllocationEnd(TLAB_END_IDENTITY).isNonNull()) {
+            assert checkInvariants(tlab);
+            insertFiller(tlab);
+            initialize(tlab, Word.nullPointer(), Word.nullPointer(), Word.nullPointer());
+        }
     }
 
     @Uninterruptible(reason = "Accesses TLAB")
@@ -238,7 +237,7 @@ private static boolean checkInvariants(Descriptor tlab) {
     @Uninterruptible(reason = "Accesses TLAB")
     static void suspendAllocationInCurrentThread() {
         /* The statistics for this thread will be updated later. */
-        retireCurrentTlab(CurrentIsolate.getCurrentThread(), false);
+        retireTlab(CurrentIsolate.getCurrentThread(), false);
     }
 
     @Uninterruptible(reason = CALLED_FROM_UNINTERRUPTIBLE_CODE, mayBeInlined = true)
@@ -268,7 +267,7 @@ static void disableAndFlushForThread(IsolateThread vmThread) {
     private static void retireTlabToEden(IsolateThread thread) {
         VMThreads.guaranteeOwnsThreadMutex("Otherwise, we wouldn't be allowed to access the space.", true);
 
-        retireCurrentTlab(thread, true);
+        retireTlab(thread, true);
 
         Descriptor tlab = getTlab(thread);
         UnalignedHeapChunk.UnalignedHeader unalignedChunk = tlab.getUnalignedChunk();
@@ -436,15 +435,16 @@ static UnsignedWord maxSize() {
     @BasedOnJDKFile("https://github.com/openjdk/jdk/blob/jdk-23-ga/src/hotspot/share/gc/shared/threadLocalAllocBuffer.cpp#L76-L117")
     @Uninterruptible(reason = CALLED_FROM_UNINTERRUPTIBLE_CODE, mayBeInlined = true)
     private static void accumulateAndResetStatistics(IsolateThread thread) {
-        gcWaste.set(thread, gcWaste.get() + UninterruptibleUtils.NumUtil.safeToInt(availableTlabMemory(getTlab(thread)).rawValue()));
+        UnsignedWord remaining = availableTlabMemory(getTlab());
+        gcWaste.set(thread, gcWaste.get() + UnsignedUtils.safeToInt(remaining));
+
         UnsignedWord totalAlignedAllocated = ThreadLocalAllocation.getAlignedAllocatedBytes(thread);
         UnsignedWord allocatedAlignedSinceLastGC = totalAlignedAllocated.subtract(tlabAllocatedAlignedBytesBeforeLastGC.get(thread));
         tlabAllocatedAlignedBytesBeforeLastGC.set(thread, totalAlignedAllocated);
 
         AdaptiveWeightedAverageStruct.sample(allocatedBytesAvg.getAddress(thread), allocatedAlignedSinceLastGC.rawValue());
 
         printStats(thread, allocatedAlignedSinceLastGC);
-
         resetStatistics(thread);
     }