[GR-35484] Added support for ThreadMXBean.getThreadAllocatedBytes.

PullRequest: graal/10643

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

@@ -28,6 +28,7 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import com.oracle.svm.core.graal.snippets.SubstrateAllocationSnippets;
 import org.graalvm.compiler.api.replacements.Fold;
 import org.graalvm.compiler.core.common.NumUtil;
 import org.graalvm.compiler.core.common.SuppressFBWarnings;
@@ -181,7 +182,7 @@ public void suspendAllocation() {
 
     @Override
     public void resumeAllocation() {
-        ThreadLocalAllocation.resumeInCurrentThread();
+        // Nothing to do - the next allocation will refill the TLAB.
     }
 
     @Override
@@ -651,6 +652,32 @@ public void optionValueChanged(RuntimeOptionKey<?> key) {
         }
     }
 
+    @Override
+    @Uninterruptible(reason = "Called from uninterruptible code.", mayBeInlined = true)
+    public long getThreadAllocatedMemory(IsolateThread thread) {
+        UnsignedWord allocatedBytes = ThreadLocalAllocation.allocatedBytes.getVolatile(thread);
+
+        /*
+         * The current aligned chunk in the TLAB is only partially filled and therefore not yet
+         * accounted for in ThreadLocalAllocation.allocatedBytes. The reads below are unsynchronized
+         * and unordered with the thread updating its TLAB, so races may occur. We only use the read
+         * values if they are plausible and not obviously racy. We also accept that certain races
+         * can cause that the memory in the current aligned TLAB chunk is counted twice.
+         */
+        ThreadLocalAllocation.Descriptor tlab = ThreadLocalAllocation.getTlab(thread);
+        AlignedHeader alignedTlab = tlab.getAlignedChunk();
+        Pointer top = tlab.getAllocationTop(SubstrateAllocationSnippets.TLAB_TOP_IDENTITY);
+        Pointer start = AlignedHeapChunk.getObjectsStart(alignedTlab);
+
+        if (top.aboveThan(start)) {
+            UnsignedWord usedTlabSize = top.subtract(start);
+            if (usedTlabSize.belowOrEqual(HeapParameters.getAlignedHeapChunkSize())) {
+                return allocatedBytes.add(usedTlabSize).rawValue();
+            }
+        }
+        return allocatedBytes.rawValue();
+    }
+
     static Pointer getImageHeapStart() {
         int imageHeapOffsetInAddressSpace = Heap.getHeap().getImageHeapOffsetInAddressSpace();
         if (imageHeapOffsetInAddressSpace > 0) {

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

@@ -28,9 +28,11 @@
 import static com.oracle.svm.core.graal.snippets.SubstrateAllocationSnippets.TLAB_TOP_IDENTITY;
 
 import com.oracle.svm.core.SubstrateGCOptions;
+import com.oracle.svm.core.threadlocal.FastThreadLocalWord;
 import org.graalvm.compiler.api.replacements.Fold;
 import org.graalvm.compiler.replacements.AllocationSnippets.FillContent;
 import org.graalvm.compiler.word.Word;
+import org.graalvm.nativeimage.CurrentIsolate;
 import org.graalvm.nativeimage.IsolateThread;
 import org.graalvm.nativeimage.Platform;
 import org.graalvm.nativeimage.Platforms;
@@ -66,9 +68,12 @@
 import com.oracle.svm.core.util.VMError;
 
 /**
- * Bump-pointer allocation from thread-local top and end Pointers.
- *
- * Many of these methods are called from allocation snippets, so they can not do anything fancy.
+ * Bump-pointer allocation from thread-local top and end Pointers. Many of these methods are called
+ * from allocation snippets, so they can not do anything fancy.
+ * 
+ * It happens that prefetch instructions access memory outside the TLAB. At the moment, this is not
+ * an issue as we only support architectures where the prefetch instructions never cause a segfault,
+ * even if they try to access memory that is not accessible.
  */
 public final class ThreadLocalAllocation {
     @RawStructure
@@ -110,6 +115,13 @@ public interface Descriptor extends PointerBase {
         void setAllocationEnd(Pointer end, LocationIdentity endIdentity);
     }
 
+    /*
+     * Stores the number of bytes that this thread allocated in the past on the Java heap. This
+     * excludes allocations that were done in the latest not yet retired {@link AlignedHeapChunk} of
+     * the TLAB.
+     */
+    public static final FastThreadLocalWord<UnsignedWord> allocatedBytes = FastThreadLocalFactory.createWord("ThreadLocalAllocation.allocatedBytes");
+
     /**
      * Don't read this value directly, use the {@link Uninterruptible} accessor methods instead.
      * This is necessary to avoid races between the GC and code that accesses or modifies the TLAB.
@@ -276,6 +288,8 @@ private static Object allocateLargeArrayInNewTlab(DynamicHub hub, int length, Un
 
         HeapChunk.setNext(newTlabChunk, tlab.getUnalignedChunk());
         tlab.setUnalignedChunk(newTlabChunk);
+
+        allocatedBytes.set(allocatedBytes.get().add(size));
         HeapImpl.getHeapImpl().getAccounting().increaseEdenUsedBytes(size);
 
         Pointer memory = UnalignedHeapChunk.allocateMemory(newTlabChunk, size);
@@ -296,10 +310,9 @@ private static Object allocateLargeArrayInNewTlab(DynamicHub hub, int length, Un
 
     @Uninterruptible(reason = "Returns uninitialized memory, modifies TLAB", callerMustBe = true)
     private static Pointer allocateRawMemoryInNewTlab(UnsignedWord size, AlignedHeader newTlabChunk) {
-        ThreadLocalAllocation.Descriptor tlab = getTlab();
-        assert DeoptTester.enabled() || availableTlabMemory(tlab).belowThan(size) : "Slowpath allocation was used even though TLAB had sufficient space";
+        assert DeoptTester.enabled() || availableTlabMemory(getTlab()).belowThan(size) : "Slowpath allocation was used even though TLAB had sufficient space";
 
-        retireCurrentAllocationChunk(tlab);
+        Descriptor tlab = retireCurrentAllocationChunk(CurrentIsolate.getCurrentThread());
         registerNewAllocationChunk(tlab, newTlabChunk);
 
         return allocateRawMemoryInTlab(size, tlab);
@@ -351,7 +364,7 @@ static void disableAndFlushForAllThreads() {
 
     @Uninterruptible(reason = "Accesses TLAB")
     static void disableAndFlushForThread(IsolateThread vmThread) {
-        retireToSpace(getTlab(vmThread), HeapImpl.getHeapImpl().getYoungGeneration().getEden());
+        retireTlabToEden(vmThread);
     }
 
     @Uninterruptible(reason = "Called from uninterruptible code.", mayBeInlined = true)
@@ -373,52 +386,52 @@ private static void freeHeapChunks(Descriptor tlab) {
 
     @Uninterruptible(reason = "Accesses TLAB")
     static void suspendInCurrentThread() {
-        retireCurrentAllocationChunk(getTlab());
+        retireCurrentAllocationChunk(CurrentIsolate.getCurrentThread());
     }
 
     @Uninterruptible(reason = "Accesses TLAB")
-    static void resumeInCurrentThread() {
-        resumeAllocationInCurrentChunk(getTlab());
-    }
-
-    @Uninterruptible(reason = "Accesses TLAB")
-    static void retireToSpace(Descriptor tlab, Space space) {
+    private static void retireTlabToEden(IsolateThread thread) {
         VMThreads.guaranteeOwnsThreadMutex("Otherwise, we wouldn't be allowed to access the space.");
-        assert !space.isOldSpace() : "must not be moved to the old gen - otherwise a remembered set would have to be constructed";
-
-        retireCurrentAllocationChunk(tlab);
 
+        Descriptor tlab = retireCurrentAllocationChunk(thread);
         AlignedHeader alignedChunk = tlab.getAlignedChunk();
         UnalignedHeader unalignedChunk = tlab.getUnalignedChunk();
         tlab.setAlignedChunk(WordFactory.nullPointer());
         tlab.setUnalignedChunk(WordFactory.nullPointer());
 
+        Space eden = HeapImpl.getHeapImpl().getYoungGeneration().getEden();
         while (alignedChunk.isNonNull()) {
             AlignedHeader next = HeapChunk.getNext(alignedChunk);
             HeapChunk.setNext(alignedChunk, WordFactory.nullPointer());
-            space.appendAlignedHeapChunk(alignedChunk);
+            eden.appendAlignedHeapChunk(alignedChunk);
             alignedChunk = next;
         }
 
         while (unalignedChunk.isNonNull()) {
             UnalignedHeader next = HeapChunk.getNext(unalignedChunk);
             HeapChunk.setNext(unalignedChunk, WordFactory.nullPointer());
-            space.appendUnalignedHeapChunk(unalignedChunk);
+            eden.appendUnalignedHeapChunk(unalignedChunk);
             unalignedChunk = next;
         }
     }
 
     @Uninterruptible(reason = "Modifies TLAB")
     private static void registerNewAllocationChunk(Descriptor tlab, AlignedHeader newChunk) {
+        assert tlab.getAllocationTop(TLAB_TOP_IDENTITY).isNull();
+        assert tlab.getAllocationEnd(TLAB_END_IDENTITY).isNull();
+
         HeapChunk.setNext(newChunk, tlab.getAlignedChunk());
         tlab.setAlignedChunk(newChunk);
         HeapImpl.getHeapImpl().getAccounting().increaseEdenUsedBytes(HeapParameters.getAlignedHeapChunkSize());
 
-        resumeAllocationInCurrentChunk(tlab);
+        tlab.setAllocationTop(HeapChunk.getTopPointer(newChunk), TLAB_TOP_IDENTITY);
+        tlab.setAllocationEnd(HeapChunk.getEndPointer(newChunk), TLAB_END_IDENTITY);
+        HeapChunk.setTopPointer(newChunk, WordFactory.nullPointer());
     }
 
-    @Uninterruptible(reason = "Modifies TLAB")
-    private static void retireCurrentAllocationChunk(Descriptor tlab) {
+    @Uninterruptible(reason = "Modifies and returns TLAB", callerMustBe = true)
+    private static Descriptor retireCurrentAllocationChunk(IsolateThread thread) {
+        Descriptor tlab = getTlab(thread);
         Pointer allocationTop = tlab.getAllocationTop(TLAB_TOP_IDENTITY);
         if (allocationTop.isNonNull()) {
             AlignedHeader alignedChunk = tlab.getAlignedChunk();
@@ -434,24 +447,10 @@ private static void retireCurrentAllocationChunk(Descriptor tlab) {
             HeapChunk.setTopPointer(alignedChunk, allocationTop);
             tlab.setAllocationTop(WordFactory.nullPointer(), TLAB_TOP_IDENTITY);
             tlab.setAllocationEnd(WordFactory.nullPointer(), TLAB_END_IDENTITY);
-        }
-    }
-
-    @Uninterruptible(reason = "Modifies TLAB.")
-    static void resumeAllocationInCurrentChunk(Descriptor tlab) {
-        assert tlab.getAllocationTop(TLAB_TOP_IDENTITY).isNull();
-        assert tlab.getAllocationEnd(TLAB_END_IDENTITY).isNull();
 
-        AlignedHeader alignedChunk = tlab.getAlignedChunk();
-        if (alignedChunk.isNonNull()) {
-            tlab.setAllocationTop(HeapChunk.getTopPointer(alignedChunk), TLAB_TOP_IDENTITY);
-            /*
-             * It happens that prefetch instructions access memory outside the TLAB. At the moment,
-             * this is not an issue as we only support architectures where the prefetch instructions
-             * never cause a segfault, even if they try to access memory that is not accessible.
-             */
-            tlab.setAllocationEnd(HeapChunk.getEndPointer(alignedChunk), TLAB_END_IDENTITY);
-            HeapChunk.setTopPointer(alignedChunk, WordFactory.nullPointer());
+            UnsignedWord usedTlabSize = HeapChunk.getTopPointer(alignedChunk).subtract(AlignedHeapChunk.getObjectsStart(alignedChunk));
+            allocatedBytes.set(thread, allocatedBytes.get(thread).add(usedTlabSize));
         }
+        return tlab;
     }
 }

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/heap/Heap.java

@@ -224,4 +224,12 @@ public List<Class<?>> getLoadedClasses() {
      * Notify the GC that the value of a GC-relevant option changed.
      */
     public abstract void optionValueChanged(RuntimeOptionKey<?> key);
+
+    /**
+     * Returns the number of bytes that were allocated by the given thread. The caller of this
+     * method must ensure that the given {@link IsolateThread} remains alive during the execution of
+     * this method.
+     */
+    @Uninterruptible(reason = "Called from uninterruptible code.", mayBeInlined = true)
+    public abstract long getThreadAllocatedMemory(IsolateThread thread);
 }

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jdk/management/SubstrateThreadMXBean.java

@@ -26,6 +26,8 @@
 
 import java.lang.management.ManagementFactory;
 import java.lang.management.ThreadInfo;
+import java.util.Arrays;
+import java.util.Objects;
 
 import javax.management.ObjectName;
 
@@ -35,11 +37,12 @@
 import com.oracle.svm.core.annotate.Uninterruptible;
 import com.oracle.svm.core.jdk.UninterruptibleUtils.AtomicInteger;
 import com.oracle.svm.core.jdk.UninterruptibleUtils.AtomicLong;
+import com.oracle.svm.core.thread.JavaThreads;
 import com.oracle.svm.core.util.VMError;
 
 import sun.management.Util;
 
-final class SubstrateThreadMXBean implements com.sun.management.ThreadMXBean {
+public final class SubstrateThreadMXBean implements com.sun.management.ThreadMXBean {
 
     private static final String MSG = "ThreadMXBean methods";
 
@@ -48,8 +51,15 @@ final class SubstrateThreadMXBean implements com.sun.management.ThreadMXBean {
     private final AtomicInteger threadCount = new AtomicInteger(0);
     private final AtomicInteger daemonThreadCount = new AtomicInteger(0);
 
+    private boolean allocatedMemoryEnabled;
+
     @Platforms(Platform.HOSTED_ONLY.class)
     SubstrateThreadMXBean() {
+        /*
+         * We always track the amount of memory that is allocated by each thread, so this MX bean
+         * feature can be on by default.
+         */
+        this.allocatedMemoryEnabled = true;
     }
 
     @Uninterruptible(reason = "Called from uninterruptible code.", mayBeInlined = true)
@@ -86,12 +96,12 @@ public ObjectName getObjectName() {
 
     @Override
     public boolean isThreadAllocatedMemoryEnabled() {
-        return false;
+        return allocatedMemoryEnabled;
     }
 
     @Override
     public boolean isThreadAllocatedMemorySupported() {
-        return false;
+        return true;
     }
 
     @Override
@@ -231,13 +241,48 @@ public ThreadInfo[] dumpAllThreads(boolean lockedMonitors, boolean lockedSynchro
     }
 
     @Override
-    public long getThreadAllocatedBytes(long arg0) {
-        throw VMError.unsupportedFeature(MSG);
+    public long getThreadAllocatedBytes(long id) {
+        boolean valid = verifyThreadAllocatedMemory(id);
+        if (!valid) {
+            return -1;
+        }
+
+        return JavaThreads.getThreadAllocatedBytes(id);
     }
 
     @Override
-    public long[] getThreadAllocatedBytes(long[] arg0) {
-        throw VMError.unsupportedFeature(MSG);
+    public long[] getThreadAllocatedBytes(long[] ids) {
+        Objects.requireNonNull(ids);
+        boolean valid = verifyThreadAllocatedMemory(ids);
+
+        long[] sizes = new long[ids.length];
+        Arrays.fill(sizes, -1);
+        if (valid) {
+            JavaThreads.getThreadAllocatedBytes(ids, sizes);
+        }
+        return sizes;
+    }
+
+    private boolean verifyThreadAllocatedMemory(long id) {
+        verifyThreadId(id);
+        return isThreadAllocatedMemoryEnabled();
+    }
+
+    private boolean verifyThreadAllocatedMemory(long[] ids) {
+        verifyThreadIds(ids);
+        return isThreadAllocatedMemoryEnabled();
+    }
+
+    private static void verifyThreadId(long id) {
+        if (id <= 0) {
+            throw new IllegalArgumentException("Invalid thread ID parameter: " + id);
+        }
+    }
+
+    private static void verifyThreadIds(long[] ids) {
+        for (int i = 0; i < ids.length; i++) {
+            verifyThreadId(ids[i]);
+        }
     }
 
     @Override
@@ -251,7 +296,7 @@ public long[] getThreadUserTime(long[] arg0) {
     }
 
     @Override
-    public void setThreadAllocatedMemoryEnabled(boolean arg0) {
-        throw VMError.unsupportedFeature(MSG);
+    public void setThreadAllocatedMemoryEnabled(boolean value) {
+        allocatedMemoryEnabled = value;
     }
 }