src/hotspot/share/prims/jvmtiExport.cpp

@@ -1364,14 +1364,13 @@
  
    JvmtiThreadState *state = get_jvmti_thread_state(thread);
    if (state == nullptr) {
      return;
    }
-   if (thread->is_in_tmp_VTMS_transition()) {
-     return; // skip ClassLoad events in tmp VTMS transition
+   if (thread->is_in_any_VTMS_transition()) {
+     return; // skip ClassLoad events in VTMS transition
    }
-   assert(!thread->is_in_any_VTMS_transition(), "class load events are not allowed in any VTMS transition");
  
    EVT_TRIG_TRACE(JVMTI_EVENT_CLASS_LOAD, ("[%s] Trg Class Load triggered",
                        JvmtiTrace::safe_get_thread_name(thread)));
    JvmtiEnvThreadStateIterator it(state);
    for (JvmtiEnvThreadState* ets = it.first(); ets != nullptr; ets = it.next(ets)) {

@@ -1402,14 +1401,13 @@
  
    JvmtiThreadState *state = get_jvmti_thread_state(thread);
    if (state == nullptr) {
      return;
    }
-   if (thread->is_in_tmp_VTMS_transition()) {
-     return; // skip ClassPrepare events in tmp VTMS transition
+   if (thread->is_in_any_VTMS_transition()) {
+     return; // skip ClassPrepare events in VTMS transition
    }
-   assert(!thread->is_in_any_VTMS_transition(), "class prepare events are not allowed in any VTMS transition");
  
    EVT_TRIG_TRACE(JVMTI_EVENT_CLASS_PREPARE, ("[%s] Trg Class Prepare triggered",
                        JvmtiTrace::safe_get_thread_name(thread)));
    JvmtiEnvThreadStateIterator it(state);
    for (JvmtiEnvThreadState* ets = it.first(); ets != nullptr; ets = it.next(ets)) {

@@ -1676,12 +1674,13 @@
    }
    JavaThread *thread = JavaThread::current();
    HandleMark hm(thread);
    EVT_TRIG_TRACE(EXT_EVENT_VIRTUAL_THREAD_UNMOUNT, ("[%p] Trg Virtual Thread Unmount event triggered", vthread));
  
-   JvmtiThreadState *state = get_jvmti_thread_state(thread);
-   if (state == nullptr) {
+   // On preemption JVMTI state rebinding has already happened so get it always direclty from the oop.
+   JvmtiThreadState *state = java_lang_Thread::jvmti_thread_state(JNIHandles::resolve(vthread));
+   if (state == NULL) {
      return;
    }
  
    if (state->is_enabled((jvmtiEvent)EXT_EVENT_VIRTUAL_THREAD_UNMOUNT)) {
      JvmtiEnvThreadStateIterator it(state);

@@ -1696,11 +1695,11 @@
  
          JvmtiVirtualThreadEventMark jem(thread);
          JvmtiJavaThreadEventTransition jet(thread);
          jvmtiExtensionEvent callback = env->ext_callbacks()->VirtualThreadUnmount;
          if (callback != nullptr) {
-           (*callback)(env->jvmti_external(), jem.jni_env(), jem.jni_thread());
+           (*callback)(env->jvmti_external(), jem.jni_env(), vthread);
          }
        }
      }
    }
  }