From da7bcfcf6e45486a0427e0ceaba74d52acbd722f Mon Sep 17 00:00:00 2001
From: Jiangli Zhou <jiangli@openjdk.org>
Date: Wed, 29 Nov 2023 16:05:54 +0000
Subject: [PATCH] 8319935: Ensure only one JvmtiThreadState is created for one
 JavaThread associated with attached native thread

Reviewed-by: manc, dholmes, sspitsyn
---
 src/hotspot/share/prims/jvmtiExport.cpp       |  7 ++++++
 .../share/prims/jvmtiThreadState.inline.hpp   | 23 ++++++++++++++-----
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/src/hotspot/share/prims/jvmtiExport.cpp b/src/hotspot/share/prims/jvmtiExport.cpp
index 26b7abf6859..eee5b2a70fa 100644
--- a/src/hotspot/share/prims/jvmtiExport.cpp
+++ b/src/hotspot/share/prims/jvmtiExport.cpp
@@ -3139,6 +3139,13 @@ bool JvmtiSampledObjectAllocEventCollector::object_alloc_is_safe_to_sample() {
     return false;
   }
 
+  // If the current thread is attaching from native and its Java thread object
+  // is being allocated, things are not ready for allocation sampling.
+  JavaThread* jt = JavaThread::cast(thread);
+  if (jt->is_attaching_via_jni() && jt->threadObj() == nullptr) {
+    return false;
+  }
+
   if (MultiArray_lock->owner() == thread) {
     return false;
   }
diff --git a/src/hotspot/share/prims/jvmtiThreadState.inline.hpp b/src/hotspot/share/prims/jvmtiThreadState.inline.hpp
index 1737bfd6a9f..0df26a1b1c5 100644
--- a/src/hotspot/share/prims/jvmtiThreadState.inline.hpp
+++ b/src/hotspot/share/prims/jvmtiThreadState.inline.hpp
@@ -78,17 +78,27 @@ inline JvmtiThreadState* JvmtiThreadState::state_for_while_locked(JavaThread *th
 
   NoSafepointVerifier nsv;  // oop is safe to use.
 
-  if (thread_oop == nullptr) {  // Then thread should not be null (see assert above).
-    thread_oop = thread->jvmti_vthread() != nullptr ? thread->jvmti_vthread() : thread->threadObj();
-  }
-
   // In a case of unmounted virtual thread the thread can be null.
   JvmtiThreadState *state = thread == nullptr ? nullptr : thread->jvmti_thread_state();
 
-  if (state == nullptr && thread != nullptr && thread->is_exiting()) {
-    // Don't add a JvmtiThreadState to a thread that is exiting.
+  if (state == nullptr && thread != nullptr &&
+      (thread->is_exiting() || thread->is_attaching_via_jni())) {
+    // Don't add a JvmtiThreadState to a thread that is exiting or is attaching.
+    // When a thread is attaching, it may not have a Java level thread object
+    // created yet.
     return nullptr;
   }
+
+  // Make sure we don't see an incomplete state. An incomplete state can cause
+  // a duplicate JvmtiThreadState being created below and bound to the 'thread'
+  // incorrectly, which leads to stale JavaThread* from the JvmtiThreadState
+  // after the thread exits.
+  assert(state == nullptr || state->get_thread_oop() != nullptr,
+         "incomplete state");
+
+  if (thread_oop == nullptr) {  // Then thread should not be null (see assert above).
+    thread_oop = thread->jvmti_vthread() != nullptr ? thread->jvmti_vthread() : thread->threadObj();
+  }
   if (state == nullptr || state->get_thread_oop() != thread_oop) {
     // Check if java_lang_Thread already has a link to the JvmtiThreadState.
     if (thread_oop != nullptr) {  // thread_oop can be null during early VMStart.
@@ -98,6 +108,7 @@ inline JvmtiThreadState* JvmtiThreadState::state_for_while_locked(JavaThread *th
       state = new JvmtiThreadState(thread, thread_oop);
     }
   }
+  assert(state != nullptr, "sanity check");
   return state;
 }