Merge pull request #24265 from tomal-majumder/jvmti-retain-jni-ids

Extend JNI ID retention to JVMTI GetStackTrace stack walks

Author: gacholio

runtime/j9vm/asgct.cpp

@@ -123,7 +123,7 @@ asyncFrameIterator(J9VMThread * currentThread, J9StackWalkState * walkState)
 		frame->lineno = (jint)walkState->bytecodePCOffset;
 	}
 	walkState->userData1 = (void*)(frame + 1);
-	declaringClass->classLoader->asyncGetCallTraceUsed = 1;
+	declaringClass->classLoader->keepJNIIDs = 1;
 	return J9_STACKWALK_KEEP_ITERATING;
 }
 

runtime/jvmti/jvmtiExtensionMechanism.c

@@ -2076,6 +2076,7 @@ jvmtiInternalGetStackTraceIteratorExtended(J9VMThread *currentThread, J9StackWal
 	UDATA frameCount = 0;
 	J9JVMTIStackTraceType type = (J9JVMTIStackTraceType)(UDATA)walkState->userData2;
 	J9Method *method = walkState->method;
+	J9Class *declaringClass = J9_CLASS_FROM_METHOD(method);
 
 #if JAVA_SPEC_VERSION >= 20
 	J9ROMMethod *romMethod = NULL;
@@ -2118,7 +2119,9 @@ jvmtiInternalGetStackTraceIteratorExtended(J9VMThread *currentThread, J9StackWal
 		} 
 
 		frame_buffer->method = methodID;
-		
+
+		declaringClass->classLoader->keepJNIIDs = 1;
+
 		if (J9_ARE_ANY_BITS_SET(type, J9JVMTI_STACK_TRACE_EXTRA_FRAME_INFO)) {
 			/* Fill in the extended data. */
 #if defined(J9VM_INTERP_NATIVE_SUPPORT)

runtime/jvmti/jvmtiStackFrame.cpp

@@ -709,6 +709,7 @@ jvmtiInternalGetStackTraceIterator(J9VMThread *currentThread, J9StackWalkState *
 	jmethodID methodID = NULL;
 	UDATA rc = J9_STACKWALK_KEEP_ITERATING;
 	J9Method *method = walkState->method;
+	J9Class *declaringClass = J9_CLASS_FROM_METHOD(method);
 
 #if JAVA_SPEC_VERSION >= 20
 	J9ROMMethod *romMethod = NULL;
@@ -747,6 +748,7 @@ jvmtiInternalGetStackTraceIterator(J9VMThread *currentThread, J9StackWalkState *
 			}
 		}
 
+		declaringClass->classLoader->keepJNIIDs = 1;
 		walkState->userData1 = frame_buffer + 1;
 	}
 

runtime/oti/j9.h

@@ -434,8 +434,8 @@ static_assert((LITERAL_STRLEN(J9_UNMODIFIABLE_CLASS_ANNOTATION) < (size_t)'/'),
 
 #define J9VM_NUM_OF_ENTRIES_IN_CLASS_JNIID_TABLE(romclass) ((romclass)->romMethodCount + (romclass)->romFieldCount)
 
-#define J9VM_SHOULD_CLEAR_JNIIDS_FOR_ASGCT(vm, classLoader) (J9_ARE_NO_BITS_SET((vm)->extendedRuntimeFlags2, J9_EXTENDED_RUNTIME2_NEVER_KEEP_JNI_IDS) \
-		&& ((classLoader)->asyncGetCallTraceUsed || J9_ARE_ANY_BITS_SET((vm)->extendedRuntimeFlags2, J9_EXTENDED_RUNTIME2_ALWAYS_KEEP_JNI_IDS)))
+#define J9VM_SHOULD_KEEP_JNIIDS(vm, classLoader) (J9_ARE_NO_BITS_SET((vm)->extendedRuntimeFlags2, J9_EXTENDED_RUNTIME2_NEVER_KEEP_JNI_IDS) \
+		&& ((classLoader)->keepJNIIDs || J9_ARE_ANY_BITS_SET((vm)->extendedRuntimeFlags2, J9_EXTENDED_RUNTIME2_ALWAYS_KEEP_JNI_IDS)))
 
 #define J9_IS_GCCONTAINERHEURISTICS(vm) \
 		J9_ARE_ANY_BITS_SET((vm)->extendedRuntimeFlags3, J9_EXTENDED_RUNTIME3_GCCONTAINERHEURISTICS)

runtime/oti/j9nonbuilder.h

@@ -3893,7 +3893,7 @@ typedef struct J9ClassLoader {
 	omrthread_monitor_t hotFieldPoolMutex;
 	omrthread_rwmutex_t cpEntriesMutex;
 	UDATA initClassPathEntryCount;
-	UDATA asyncGetCallTraceUsed;
+	UDATA keepJNIIDs;
 	omrthread_monitor_t mapCacheMutex;
 	struct J9HashTable* localmapCache;
 	struct J9HashTable* argsbitsCache;

runtime/vm/classallocation.c

@@ -430,7 +430,7 @@ freeClassLoader(J9ClassLoader *classLoader, J9JavaVM *javaVM, J9VMThread *vmThre
 	}
 
 	if (NULL != classLoader->jniIDs) {
-		if (J9VM_SHOULD_CLEAR_JNIIDS_FOR_ASGCT(javaVM, classLoader)) {
+		if (J9VM_SHOULD_KEEP_JNIIDS(javaVM, classLoader)) {
 			pool_state idWalkState;
 			J9GenericJNIID *jniID = pool_startDo(classLoader->jniIDs, &idWalkState);
 

runtime/vm/jvminit.c

@@ -8777,6 +8777,19 @@ freeClassNativeMemory(J9HookInterface** hook, UDATA eventNum, void* eventData, v
 	J9Class * clazz = data->clazz;
 	PORT_ACCESS_FROM_JAVAVM(vm);
 
+	if (NULL != clazz->jniIDs) {
+		J9ClassLoader *classLoader = clazz->classLoader;
+		if (J9VM_SHOULD_KEEP_JNIIDS(vm, classLoader)) {
+			UDATA size = J9VM_NUM_OF_ENTRIES_IN_CLASS_JNIID_TABLE(clazz->romClass);
+			for (UDATA i = 0; i < size; i++) {
+				J9GenericJNIID *id = (J9GenericJNIID *)(clazz->jniIDs[i]);
+				if (NULL != id) {
+					memset(id, -1, sizeof(J9GenericJNIID));
+				}
+			}
+		}
+	}
+
 	/* Free the ID table for this class, but do not free any of the IDs.  They will be freed by killing their
 		pools when the class loader is unloaded.
 	*/
@@ -8870,7 +8883,7 @@ vmHookAnonClassesUnload(J9HookInterface** hook, UDATA eventNum, void* eventData,
 		J9ClassLoader *classLoader = j9clazz->classLoader;
 		/* Anon classes are unloaded piecemeal, so clear the map cache where the anon maps may be cached */
 		freeMapCaches(classLoader);
-		if (J9VM_SHOULD_CLEAR_JNIIDS_FOR_ASGCT(vm, classLoader)) {
+		if (J9VM_SHOULD_KEEP_JNIIDS(vm, classLoader)) {
 			void **jniIDs = j9clazz->jniIDs;
 			if (NULL != jniIDs) {
 				UDATA size = J9VM_NUM_OF_ENTRIES_IN_CLASS_JNIID_TABLE(j9clazz->romClass);