Retain Continuation.vthread until the J9VMContinuation is freed

This PR is needed to enable the Skynet benchmark.

The changes in #18180 pass threadObject to
walkContinuationStackFrames.

The GC can walk the continuations until the associated native
J9VMContinuation structure is freed.

Since we rely on Continuation.vthread in walkContinuationStackFrames,
Continuation.vthread can only be unset once the native continuation
structure is freed.

Now, we set Continuation.vthread to NULL after the continuation is
freed. This will prevent a segfault which can occur due to
#18180.

Continuation.vthread is also used in JVMTI to not suspend virtual
threads once they enter the last unmount phase. This won't work anymore
because Continuation.vthread is no longer set to NULL at the start of
the last unmount phase. To address this issue, another continuation
state has been added to indicate that a virtual thread's continuation
has entered the last unmount phase.

Signed-off-by: Babneet Singh <sbabneet@ca.ibm.com>

runtime/j9vm/javanextvmi.cpp

@@ -296,14 +296,15 @@ exitVThreadTransitionCritical(J9VMThread *currentThread, j9object_t vthread)
 }
 
 static void
-unsetParentVthread(J9VMThread *currentThread, jobject thread)
+setContinuationStateToLastUnmount(J9VMThread *currentThread, jobject thread)
 {
 	enterVThreadTransitionCritical(currentThread, thread);
 	/* Re-fetch reference as enterVThreadTransitionCritical may release VMAccess. */
 	j9object_t threadObj = J9_JNI_UNWRAP_REFERENCE(thread);
 	j9object_t continuationObj = J9VMJAVALANGVIRTUALTHREAD_CONT(currentThread, threadObj);
-	/* Add reverse link from Continuation object to VirtualThread object, this let JVMTI code. */
-	J9VMJDKINTERNALVMCONTINUATION_SET_VTHREAD(currentThread, continuationObj, NULL);
+	ContinuationState volatile *continuationStatePtr = VM_ContinuationHelpers::getContinuationStateAddress(currentThread, continuationObj);
+	/* Used in JVMTI to not suspend the virtual thread once it enters the last unmount phase. */
+	VM_ContinuationHelpers::setLastUnmount(continuationStatePtr);
 	exitVThreadTransitionCritical(currentThread, threadObj);
 }
 
@@ -501,7 +502,7 @@ JVM_VirtualThreadUnmountBegin(JNIEnv *env, jobject thread, jboolean lastUnmount)
 
 	if (lastUnmount) {
 		TRIGGER_J9HOOK_VM_VIRTUAL_THREAD_END(vm->hookInterface, currentThread);
-		unsetParentVthread((J9VMThread *)env, thread);
+		setContinuationStateToLastUnmount((J9VMThread *)env, thread);
 	}
 	virtualThreadUnmountBegin(env, thread);
 
@@ -661,7 +662,7 @@ JVM_VirtualThreadEnd(JNIEnv *env, jobject vthread)
 	vmFuncs->internalEnterVMFromJNI(currentThread);
 
 	TRIGGER_J9HOOK_VM_VIRTUAL_THREAD_END(vm->hookInterface, currentThread);
-	unsetParentVthread((J9VMThread *)env, vthread);
+	setContinuationStateToLastUnmount((J9VMThread *)env, vthread);
 	virtualThreadUnmountBegin(env, vthread);
 
 	vmFuncs->internalExitVMToJNI(currentThread);

runtime/jvmti/jvmtiThread.c

@@ -1324,8 +1324,11 @@ jvmtiGetCurrentThread(jvmtiEnv *env,
 static jvmtiIterationControl
 jvmtiSuspendResumeCallBack(J9VMThread *vmThread, J9MM_IterateObjectDescriptor *object, void *userData)
 {
-	j9object_t vthread = J9VMJDKINTERNALVMCONTINUATION_VTHREAD(vmThread, object->object);
-	if (NULL != vthread) {
+	j9object_t continuationObj = object->object;
+	j9object_t vthread = J9VMJDKINTERNALVMCONTINUATION_VTHREAD(vmThread, continuationObj);
+	ContinuationState continuationState = J9VMJDKINTERNALVMCONTINUATION_STATE(vmThread, continuationObj);;
+
+	if ((NULL != vthread) && J9_ARE_NO_BITS_SET(continuationState, J9_GC_CONTINUATION_STATE_LAST_UNMOUNT)) {
 		jvmtiVThreadCallBackData *data = (jvmtiVThreadCallBackData*)userData;
 		BOOLEAN is_excepted = FALSE;
 		for (jint i = 0; i < data->except_count; ++i) {

runtime/oti/ContinuationHelpers.hpp

@@ -108,6 +108,18 @@ class VM_ContinuationHelpers {
 		*continuationStatePtr |= J9_GC_CONTINUATION_STATE_STARTED;
 	}
 
+	static VMINLINE bool
+	isLastUnmount(ContinuationState continuationState)
+	{
+		return J9_ARE_ALL_BITS_SET(continuationState, J9_GC_CONTINUATION_STATE_LAST_UNMOUNT);
+	}
+
+	static VMINLINE void
+	setLastUnmount(ContinuationState volatile *continuationStatePtr)
+	{
+		*continuationStatePtr |= J9_GC_CONTINUATION_STATE_LAST_UNMOUNT;
+	}
+
 	static VMINLINE bool
 	isFinished(ContinuationState continuationState)
 	{

runtime/oti/j9consts.h

@@ -501,6 +501,7 @@ extern "C" {
 #define J9_GC_CONTINUATION_STATE_PENDING_TO_BE_MOUNTED 0x4
 #define J9_GC_CONTINUATION_STATE_CONCURRENT_SCAN_LOCAL 0x8
 #define J9_GC_CONTINUATION_STATE_CONCURRENT_SCAN_GLOBAL 0x10
+#define J9_GC_CONTINUATION_STATE_LAST_UNMOUNT 0x20
 #define J9_GC_CONTINUATION_STATE_CONCURRENT_SCAN_ANY (J9_GC_CONTINUATION_STATE_CONCURRENT_SCAN_LOCAL | J9_GC_CONTINUATION_STATE_CONCURRENT_SCAN_GLOBAL)
 #define J9_GC_CONTINUATION_STATE_CARRIERID_MASK (~(uintptr_t)0xff)
 

runtime/vm/ContinuationHelpers.cpp

@@ -343,6 +343,7 @@ freeContinuation(J9VMThread *currentThread, j9object_t continuationObject, BOOLE
 
 		/* Update Continuation object's vmRef field. */
 		J9VMJDKINTERNALVMCONTINUATION_SET_VMREF(currentThread, continuationObject, NULL);
+		J9VMJDKINTERNALVMCONTINUATION_SET_VTHREAD(currentThread, continuationObject, NULL);
 
 		recycleContinuation(currentThread->javaVM, currentThread, continuation, skipLocalCache);
 	}