[GR-51851] Add initial support for native memory tracking (NMT mallocs)

sdk/src/org.graalvm.nativeimage/src/org/graalvm/nativeimage/UnmanagedMemory.java

@@ -76,6 +76,14 @@ public static <T extends PointerBase> T malloc(UnsignedWord size) {
         return result;
     }
 
+    public static <T extends PointerBase> T malloc(UnsignedWord size, int flag) {
+        T result = ImageSingletons.lookup(UnmanagedMemorySupport.class).malloc(size, flag);
+        if (result.isNull()) {
+            throw new OutOfMemoryError("malloc of unmanaged memory");
+        }
+        return result;
+    }
+
     /**
      * Allocates {@code size} bytes of unmanaged memory. The content of the memory is undefined.
      * <p>
@@ -138,6 +146,14 @@ public static <T extends PointerBase> T realloc(T ptr, UnsignedWord size) {
         return result;
     }
 
+    public static <T extends PointerBase> T realloc(T ptr, UnsignedWord size, int flag) {
+        T result = ImageSingletons.lookup(UnmanagedMemorySupport.class).realloc(ptr, size, flag);
+        if (result.isNull()) {
+            throw new OutOfMemoryError("realloc of unmanaged memory");
+        }
+        return result;
+    }
+
     /**
      * Frees unmanaged memory that was previously allocated using methods of this class.
      *
@@ -146,4 +162,13 @@ public static <T extends PointerBase> T realloc(T ptr, UnsignedWord size) {
     public static void free(PointerBase ptr) {
         ImageSingletons.lookup(UnmanagedMemorySupport.class).free(ptr);
     }
+
+    /**
+     * Will not attempt to perform any NMT operations. This is crucial for releasing memory
+     * allocated by C libraries which will not have NMT "malloc headers". If
+     * {@link UnmanagedMemory#free(PointerBase)} is used instead, a segfault will occur.
+     */
+    public static void untrackedFree(PointerBase ptr) {
+        ImageSingletons.lookup(UnmanagedMemorySupport.class).untrackedFree(ptr);
+    }
 }

sdk/src/org.graalvm.nativeimage/src/org/graalvm/nativeimage/impl/UnmanagedMemorySupport.java

@@ -47,9 +47,17 @@
 public interface UnmanagedMemorySupport {
     <T extends PointerBase> T malloc(UnsignedWord size);
 
+    <T extends PointerBase> T malloc(UnsignedWord size, int flag);
+
     <T extends PointerBase> T calloc(UnsignedWord size);
 
+    <T extends PointerBase> T calloc(UnsignedWord size, int flag);
+
     <T extends PointerBase> T realloc(T ptr, UnsignedWord size);
 
+    <T extends PointerBase> T realloc(T ptr, UnsignedWord size, int flag);
+
     void free(PointerBase ptr);
+
+    void untrackedFree(PointerBase ptr);
 }

substratevm/src/com.oracle.svm.core.posix/src/com/oracle/svm/core/posix/PosixProcessPropertiesSupport.java

@@ -36,12 +36,12 @@
 import org.graalvm.nativeimage.c.type.CCharPointer;
 import org.graalvm.nativeimage.c.type.CTypeConversion;
 import org.graalvm.nativeimage.c.type.CTypeConversion.CCharPointerHolder;
+import org.graalvm.nativeimage.UnmanagedMemory;
 import org.graalvm.word.PointerBase;
 import org.graalvm.word.WordFactory;
 
 import com.oracle.svm.core.BaseProcessPropertiesSupport;
 import com.oracle.svm.core.graal.stackvalue.UnsafeStackValue;
-import com.oracle.svm.core.headers.LibC;
 import com.oracle.svm.core.posix.headers.Dlfcn;
 import com.oracle.svm.core.posix.headers.Signal;
 import com.oracle.svm.core.posix.headers.Stdlib;
@@ -103,7 +103,7 @@ public String getObjectFile(PointerBase symbolAddress) {
         try {
             return CTypeConversion.toJavaString(realpath);
         } finally {
-            LibC.free(realpath);
+            UnmanagedMemory.untrackedFree(realpath);
         }
     }
 
@@ -164,7 +164,7 @@ protected static String realpath(String path) {
             } else {
                 /* Success */
                 final String result = CTypeConversion.toJavaString(realpathPointer);
-                LibC.free(realpathPointer);
+                UnmanagedMemory.untrackedFree(realpathPointer);
                 return result;
             }
         }

substratevm/src/com.oracle.svm.core.posix/src/com/oracle/svm/core/posix/UnmanagedMemorySupportImpl.java

@@ -24,38 +24,112 @@
  */
 package com.oracle.svm.core.posix;
 
+import jdk.graal.compiler.api.replacements.Fold;
+import org.graalvm.nativeimage.ImageSingletons;
 import org.graalvm.nativeimage.impl.UnmanagedMemorySupport;
+
 import org.graalvm.word.PointerBase;
+import org.graalvm.word.Pointer;
 import org.graalvm.word.UnsignedWord;
 import org.graalvm.word.WordFactory;
+import org.graalvm.nativeimage.StackValue;
 
 import com.oracle.svm.core.Uninterruptible;
 import com.oracle.svm.core.feature.AutomaticallyRegisteredImageSingleton;
-import com.oracle.svm.core.headers.LibC;
+import com.oracle.svm.core.headers.LibCSupport;
+import com.oracle.svm.core.nmt.NmtFlag;
+import com.oracle.svm.core.nmt.NativeMemoryTracking;
+import com.oracle.svm.core.nmt.ReturnAddress;
 
 @AutomaticallyRegisteredImageSingleton(UnmanagedMemorySupport.class)
 class UnmanagedMemorySupportImpl implements UnmanagedMemorySupport {
     @Override
     @Uninterruptible(reason = "Called from uninterruptible code.", mayBeInlined = true)
     public <T extends PointerBase> T malloc(UnsignedWord size) {
-        return LibC.malloc(size);
+        return malloc(size, NmtFlag.mtNone.ordinal());
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    @Uninterruptible(reason = "Called from uninterruptible code.", mayBeInlined = true)
+    public <T extends PointerBase> T malloc(UnsignedWord size, int flag) {
+        ReturnAddress ra = StackValue.get(ReturnAddress.class);
+        if (NativeMemoryTracking.handlePreInitMallocs(size, flag, ra)) {
+            return (T) ra.get();
+        }
+        Pointer outerPointer = libc().malloc(size.add(NativeMemoryTracking.getHeaderSize()));
+        return (T) NativeMemoryTracking.recordMalloc(outerPointer, size, flag);
     }
 
     @Override
     @Uninterruptible(reason = "Called from uninterruptible code.", mayBeInlined = true)
     public <T extends PointerBase> T calloc(UnsignedWord size) {
-        return LibC.calloc(WordFactory.unsigned(1), size);
+        return calloc(size, NmtFlag.mtNone.ordinal());
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    @Uninterruptible(reason = "Called from uninterruptible code.", mayBeInlined = true)
+    public <T extends PointerBase> T calloc(UnsignedWord size, int flag) {
+        ReturnAddress ra = StackValue.get(ReturnAddress.class);
+        if (NativeMemoryTracking.handlePreInitCallocs(size, flag, ra)) {
+            return (T) ra.get();
+        }
+        Pointer outerPointer = libc().calloc(WordFactory.unsigned(1), size.add(NativeMemoryTracking.getHeaderSize()));
+        return (T) NativeMemoryTracking.recordMalloc(outerPointer, size, flag);
     }
 
     @Override
     @Uninterruptible(reason = "Called from uninterruptible code.", mayBeInlined = true)
     public <T extends PointerBase> T realloc(T ptr, UnsignedWord size) {
-        return LibC.realloc(ptr, size);
+        return realloc(ptr, size, NmtFlag.mtNone.ordinal());
+    }
+
+    @Override
+    @SuppressWarnings("unchecked")
+    @Uninterruptible(reason = "Called from uninterruptible code.", mayBeInlined = true)
+    public <T extends PointerBase> T realloc(T ptr, UnsignedWord size, int flag) {
+
+        ReturnAddress ra = StackValue.get(ReturnAddress.class);
+        if (NativeMemoryTracking.handlePreInitReallocs((Pointer) ptr, size, flag, ra)) {
+            return (T) ra.get();
+        }
+
+        // Retrieve necessary data from the old block
+        Pointer oldOuterPointer = ((Pointer) ptr).subtract(NativeMemoryTracking.getHeaderSize());
+        long oldSize = NativeMemoryTracking.getAllocationSize(oldOuterPointer);
+        int oldCategory = NativeMemoryTracking.getAllocationCategory(oldOuterPointer);
+
+        // Perform the realloc
+        Pointer newOuterPointer = libc().realloc(oldOuterPointer, size.add(NativeMemoryTracking.getHeaderSize()));
+
+        // Only deaccount the old block, if we were successful.
+        if (newOuterPointer.isNonNull()) {
+            NativeMemoryTracking.deaccountMalloc(oldSize, oldCategory);
+        }
+
+        // Account the new block and overwrite the header with the new tracking data
+        return (T) NativeMemoryTracking.recordMalloc(newOuterPointer, size, flag);
     }
 
     @Override
     @Uninterruptible(reason = "Called from uninterruptible code.", mayBeInlined = true)
     public void free(PointerBase ptr) {
-        LibC.free(ptr);
+        if (NativeMemoryTracking.handlePreInitFrees((Pointer) ptr)) {
+            return;
+        }
+        NativeMemoryTracking.deaccountMalloc(ptr);
+        libc().free(((Pointer) ptr).subtract(NativeMemoryTracking.getHeaderSize()));
+    }
+
+    @Override
+    @Uninterruptible(reason = "Called from uninterruptible code.", mayBeInlined = true)
+    public void untrackedFree(PointerBase ptr) {
+        libc().free(ptr);
+    }
+
+    @Fold
+    static LibCSupport libc() {
+        return ImageSingletons.lookup(LibCSupport.class);
     }
 }

substratevm/src/com.oracle.svm.core.posix/src/com/oracle/svm/core/posix/darwin/DarwinSystemPropertiesSupport.java

@@ -31,13 +31,13 @@
 import org.graalvm.nativeimage.c.type.CTypeConversion;
 import org.graalvm.nativeimage.c.type.CTypeConversion.CCharPointerHolder;
 import org.graalvm.nativeimage.impl.RuntimeSystemPropertiesSupport;
+import org.graalvm.nativeimage.UnmanagedMemory;
 import org.graalvm.word.UnsignedWord;
 import org.graalvm.word.WordFactory;
 
 import com.oracle.svm.core.feature.AutomaticallyRegisteredFeature;
 import com.oracle.svm.core.feature.InternalFeature;
 import com.oracle.svm.core.graal.stackvalue.UnsafeStackValue;
-import com.oracle.svm.core.headers.LibC;
 import com.oracle.svm.core.jdk.SystemPropertiesSupport;
 import com.oracle.svm.core.posix.PosixSystemPropertiesSupport;
 import com.oracle.svm.core.posix.headers.Limits;
@@ -105,7 +105,7 @@ protected String osVersionValue() {
                 CCharPointer osVersionStr = Foundation.systemVersionPlatform();
                 if (osVersionStr.isNonNull()) {
                     osVersionValue = CTypeConversion.toJavaString(osVersionStr);
-                    LibC.free(osVersionStr);
+                    UnmanagedMemory.untrackedFree(osVersionStr);
                     return osVersionValue;
                 }
             } else {
@@ -120,7 +120,7 @@ protected String osVersionValue() {
         CCharPointer osVersionStr = Foundation.systemVersionPlatformFallback();
         if (osVersionStr.isNonNull()) {
             osVersionValue = CTypeConversion.toJavaString(osVersionStr);
-            LibC.free(osVersionStr);
+            UnmanagedMemory.untrackedFree(osVersionStr);
             return osVersionValue;
         }
         return osVersionValue = "Unknown";

substratevm/src/com.oracle.svm.core.windows/src/com/oracle/svm/core/windows/WindowsSystemPropertiesSupport.java

@@ -38,14 +38,14 @@
 import org.graalvm.nativeimage.c.type.VoidPointer;
 import org.graalvm.nativeimage.c.type.WordPointer;
 import org.graalvm.nativeimage.impl.RuntimeSystemPropertiesSupport;
+import org.graalvm.nativeimage.impl.UnmanagedMemorySupport;
 import org.graalvm.word.UnsignedWord;
 import org.graalvm.word.WordFactory;
 
 import com.oracle.svm.core.c.NonmovableArrays;
 import com.oracle.svm.core.feature.AutomaticallyRegisteredFeature;
 import com.oracle.svm.core.feature.InternalFeature;
 import com.oracle.svm.core.graal.stackvalue.UnsafeStackValue;
-import com.oracle.svm.core.headers.LibC;
 import com.oracle.svm.core.jdk.SystemPropertiesSupport;
 import com.oracle.svm.core.util.VMError;
 import com.oracle.svm.core.windows.headers.FileAPI;
@@ -245,29 +245,30 @@ public Pair<String, String> getOsNameAndVersion() {
                 break;
             }
 
-            VoidPointer versionInfo = LibC.malloc(WordFactory.unsigned(versionSize));
+            VoidPointer versionInfo = ImageSingletons.lookup(UnmanagedMemorySupport.class).malloc(WordFactory.unsigned(versionSize));
             if (versionInfo.isNull()) {
                 break;
             }
+            try {
 
-            if (WinVer.GetFileVersionInfoW(kernel32Path, 0, versionSize, versionInfo) == 0) {
-                LibC.free(versionInfo);
-                break;
-            }
+                if (WinVer.GetFileVersionInfoW(kernel32Path, 0, versionSize, versionInfo) == 0) {
+                    break;
+                }
 
-            WindowsLibC.WCharPointer rootPath = NonmovableArrays.addressOf(NonmovableArrays.fromImageHeap(ROOT_PATH), 0);
-            WordPointer fileInfoPointer = UnsafeStackValue.get(WordPointer.class);
-            CIntPointer lengthPointer = UnsafeStackValue.get(CIntPointer.class);
-            if (WinVer.VerQueryValueW(versionInfo, rootPath, fileInfoPointer, lengthPointer) == 0) {
-                LibC.free(versionInfo);
-                break;
-            }
+                WindowsLibC.WCharPointer rootPath = NonmovableArrays.addressOf(NonmovableArrays.fromImageHeap(ROOT_PATH), 0);
+                WordPointer fileInfoPointer = UnsafeStackValue.get(WordPointer.class);
+                CIntPointer lengthPointer = UnsafeStackValue.get(CIntPointer.class);
+                if (WinVer.VerQueryValueW(versionInfo, rootPath, fileInfoPointer, lengthPointer) == 0) {
+                    break;
+                }
 
-            VerRsrc.VS_FIXEDFILEINFO fileInfo = fileInfoPointer.read();
-            majorVersion = (short) (fileInfo.dwProductVersionMS() >> 16); // HIWORD
-            minorVersion = (short) fileInfo.dwProductVersionMS(); // LOWORD
-            buildNumber = (short) (fileInfo.dwProductVersionLS() >> 16); // HIWORD
-            LibC.free(versionInfo);
+                VerRsrc.VS_FIXEDFILEINFO fileInfo = fileInfoPointer.read();
+                majorVersion = (short) (fileInfo.dwProductVersionMS() >> 16); // HIWORD
+                minorVersion = (short) fileInfo.dwProductVersionMS(); // LOWORD
+                buildNumber = (short) (fileInfo.dwProductVersionLS() >> 16); // HIWORD
+            } finally {
+                ImageSingletons.lookup(UnmanagedMemorySupport.class).free(versionInfo);
+            }
         } while (false);
 
         String osVersion = majorVersion + "." + minorVersion;

substratevm/src/com.oracle.svm.core.windows/src/com/oracle/svm/core/windows/WindowsUnmanagedMemorySupportImpl.java

@@ -24,38 +24,66 @@
  */
 package com.oracle.svm.core.windows;
 
+import jdk.graal.compiler.api.replacements.Fold;
+
+import org.graalvm.nativeimage.ImageSingletons;
 import org.graalvm.nativeimage.impl.UnmanagedMemorySupport;
 import org.graalvm.word.PointerBase;
 import org.graalvm.word.UnsignedWord;
 import org.graalvm.word.WordFactory;
 
 import com.oracle.svm.core.Uninterruptible;
 import com.oracle.svm.core.feature.AutomaticallyRegisteredImageSingleton;
-import com.oracle.svm.core.headers.LibC;
+import com.oracle.svm.core.headers.LibCSupport;
 
 @AutomaticallyRegisteredImageSingleton(UnmanagedMemorySupport.class)
 class WindowsUnmanagedMemorySupportImpl implements UnmanagedMemorySupport {
     @Override
     @Uninterruptible(reason = "Called from uninterruptible code.", mayBeInlined = true)
     public <T extends PointerBase> T malloc(UnsignedWord size) {
-        return LibC.malloc(size);
+        return libc().malloc(size);
+    }
+
+    @Override
+    public <T extends PointerBase> T malloc(UnsignedWord size, int flag) {
+        return libc().malloc(size);
     }
 
     @Override
     @Uninterruptible(reason = "Called from uninterruptible code.", mayBeInlined = true)
     public <T extends PointerBase> T calloc(UnsignedWord size) {
-        return LibC.calloc(WordFactory.unsigned(1), size);
+        return libc().calloc(WordFactory.unsigned(1), size);
+    }
+
+    @Override
+    public <T extends PointerBase> T calloc(UnsignedWord size, int flag) {
+        return libc().calloc(WordFactory.unsigned(1), size);
     }
 
     @Override
     @Uninterruptible(reason = "Called from uninterruptible code.", mayBeInlined = true)
     public <T extends PointerBase> T realloc(T ptr, UnsignedWord size) {
-        return LibC.realloc(ptr, size);
+        return libc().realloc(ptr, size);
+    }
+
+    @Override
+    public <T extends PointerBase> T realloc(T ptr, UnsignedWord size, int flag) {
+        return libc().realloc(ptr, size);
     }
 
     @Override
     @Uninterruptible(reason = "Called from uninterruptible code.", mayBeInlined = true)
     public void free(PointerBase ptr) {
-        LibC.free(ptr);
+        libc().free(ptr);
+    }
+
+    @Override
+    public void untrackedFree(PointerBase ptr) {
+        libc().free(ptr);
+    }
+
+    @Fold
+    static LibCSupport libc() {
+        return ImageSingletons.lookup(LibCSupport.class);
     }
 }

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/SubstrateOptions.java

@@ -912,6 +912,9 @@ public Boolean getValue(OptionValues values) {
     @Option(help = "file:doc-files/FlightRecorderLoggingHelp.txt")//
     public static final RuntimeOptionKey<String> FlightRecorderLogging = new RuntimeOptionKey<>("all=warning", Immutable);
 
+    @Option(help = "Enable native memory tracking")//
+    public static final RuntimeOptionKey<Boolean> NativeMemoryTracking = new RuntimeOptionKey<>(false);
+
     public static String reportsPath() {
         Path reportsPath = ImageSingletons.lookup(ReportingSupport.class).reportsPath;
         if (reportsPath.isAbsolute()) {