NMT extended to Java

Author: jdksjolen

src/hotspot/share/prims/jvm.cpp

@@ -3837,6 +3837,11 @@ JVM_LEAF(jboolean, JVM_PrintWarningAtDynamicAgentLoad(void))
   return (EnableDynamicAgentLoading && !FLAG_IS_CMDLINE(EnableDynamicAgentLoading)) ? JNI_TRUE : JNI_FALSE;
 JVM_END
 
+
+/*
+ * Native Memory Tracking
+ */
+
 JNIEXPORT arena_t JNICALL
 JVM_MakeArena(const char *name) {
   return static_cast<arena_t>(MemTagFactory::tag(name));

src/hotspot/share/prims/nativeLookup.cpp

@@ -199,6 +199,7 @@ char* NativeLookup::long_jni_name(const methodHandle& method) {
 }
 
 extern "C" {
+  void JNICALL JVM_RegisterNativeMemoryTrackingMethods(JNIEnv *env, jclass unsafecls);
   void JNICALL JVM_RegisterMethodHandleMethods(JNIEnv *env, jclass unsafecls);
   void JNICALL JVM_RegisterReferencesMethods(JNIEnv *env, jclass unsafecls);
   void JNICALL JVM_RegisterUpcallHandlerMethods(JNIEnv *env, jclass unsafecls);
@@ -227,6 +228,7 @@ static JNINativeMethod lookup_special_native_methods[] = {
   { CC"Java_sun_hotspot_WhiteBox_registerNatives",                 nullptr, FN_PTR(JVM_RegisterWhiteBoxMethods)     },
   { CC"Java_jdk_test_whitebox_WhiteBox_registerNatives",           nullptr, FN_PTR(JVM_RegisterWhiteBoxMethods)     },
   { CC"Java_jdk_internal_vm_vector_VectorSupport_registerNatives", nullptr, FN_PTR(JVM_RegisterVectorSupportMethods)},
+  { CC"Java_lang_foreign_NativeMemoryTracking_registerNatives", nullptr, FN_PTR(JVM_RegisterNativeMemoryTrackingMethods)},
 #if INCLUDE_JVMCI
   { CC"Java_jdk_vm_ci_runtime_JVMCI_initializeRuntime",            nullptr, FN_PTR(JVM_GetJVMCIRuntime)             },
   { CC"Java_jdk_vm_ci_services_Services_readSystemPropertiesInfo", nullptr, FN_PTR(JVM_ReadSystemPropertiesInfo)    },

src/hotspot/share/prims/nativememorytracking.cpp

@@ -0,0 +1,50 @@
+#include "jni.h"
+#include "runtime/interfaceSupport.inline.hpp"
+#include "runtime/jniHandles.hpp"
+#include "nmt/memTagFactory.hpp"
+#include "nmt/memTracker.hpp"
+
+#define CC (char*)  /*cast a literal from (const char*)*/
+#define FN_PTR(f) CAST_FROM_FN_PTR(void*, &f)
+
+JVM_ENTRY(jlong, NMT_makeTag(JNIEnv *env, jobject ignored_this, jobject tag_name_string)) {
+  oop tn_oop = JNIHandles::resolve(tag_name_string);
+  Handle tnh(THREAD, tn_oop);
+  if (tnh.is_null()) {
+    // TODO: Throw exception
+  }
+
+  if (!java_lang_String::is_instance(tn_oop)) {
+    // TODO: Throw exception
+  }
+
+  if (!MemTracker::enabled()) {
+    return (jlong)mtNone;
+  }
+
+  ResourceMark rm;
+  const char* tag_name = java_lang_String::as_utf8_string(tn_oop);
+  return (jlong)MemTagFactory::tag(tag_name);
+}
+JVM_END
+
+JVM_ENTRY(jlong, NMT_allocate(JNIEnv *env, jobject ignored_this, jlong size, jlong mem_tag)) {
+  return (jlong)os::malloc((size_t)size, (MemTag)mem_tag);
+}
+JVM_ENTRY(void, NMT_free(JNIEnv *env, jobject ignored_this, jlong ptr)) {
+  os::free((void*)ptr);
+}
+
+static JNINativeMethod NMT_methods[] = {
+  {CC "makeTag", CC "(Ljava/lang/String)J", FN_PTR(NMT_makeTag)},
+  {CC "allocate", CC "(JJ)J", FN_PTR(NMT_allocate)},
+  {CC "freeMemory" CC "(J)V", FN_PTR(NMT_free)}
+};
+
+JVM_ENTRY(void, JVM_RegisterNativeMemoryTrackingMethods(JNIEnv *env, jclass NMT_class)) {
+  ThreadToNativeFromVM ttndv(thread);
+
+  int status = env->RegisterNatives(NMT_class, NMT_methods, sizeof(NMT_methods) / sizeof(JNINativeMethod));
+  guarantee(status == JNI_OK && !env->ExceptionCheck(),
+            "register java.lang.invoke.MethodHandleNative natives");
+}

src/java.base/share/classes/java/lang/foreign/NativeMemoryTracking.java

@@ -0,0 +1,45 @@
+package java.lang.foreign;
+
+public class NativeMemoryTracking {
+    public static native long makeTag(String name);
+    public static native long allocate(long size, long mem_tag);
+    public static native void freeMemory(long addr);
+
+    private static native void registerNatives();
+    static {
+        runtimeSetup();
+    }
+
+    // Unsafe.java has this comment: Called from JVM when loading an AOT cache
+    // so we imitate the same pattern, maybe it's necessary for us as well???
+    private static void runtimeSetup() {
+        registerNatives();
+    }
+
+    public static class NMTArena implements Arena {
+        private long mem_tag;
+
+        public NMTArena(String name) {
+            this.mem_tag = NativeMemoryTracking.makeTag(name);
+        }
+
+        final Arena arena = Arena.ofConfined();
+
+        @Override
+        public Scope scope() {
+            return arena.scope();
+        }
+
+        @Override
+        public void close() {
+            arena.close();
+        }
+
+        @Override
+        public MemorySegment allocate(long byteSize, long byteAlignment) {
+            long address = NativeMemoryTracking.allocate(byteSize, mem_tag);
+            return MemorySegment.ofAddress(address)
+                    .reinterpret(byteSize, arena, ms -> NativeMemoryTracking.freeMemory(ms.address()));
+        }
+    }
+}