8372039: post_sampled_object_alloc is called while lock is handled

Reviewed-by: sspitsyn, eosterlund, amenkov

src/hotspot/share/cds/aotStreamedHeapLoader.cpp

@@ -184,6 +184,7 @@ static size_t archive_object_size(oopDesc* archive_object) {
 oop AOTStreamedHeapLoader::allocate_object(oopDesc* archive_object, markWord mark, size_t size, TRAPS) {
   assert(!archive_object->is_stackChunk(), "no such objects are archived");
 
+  NoJvmtiEventsMark njem;
   oop heap_object;
 
   Klass* klass = archive_object->klass();

src/hotspot/share/cds/aotThread.cpp

@@ -63,7 +63,7 @@ void AOTThread::initialize() {
   // This is important because this thread runs before JVMTI monitors are set up appropriately.
   // Therefore, callbacks would not work as intended. JVMTI has no business peeking at how we
   // materialize primordial objects from the AOT cache.
-  thread->toggle_is_disable_suspend();
+  thread->disable_jvmti_events();
 #endif
 
   JavaThread::vm_exit_on_osthread_failure(thread);

src/hotspot/share/prims/jvmtiExport.cpp

@@ -3159,31 +3159,19 @@ void JvmtiObjectAllocEventCollector::record_allocation(oop obj) {
   _allocated->push(OopHandle(JvmtiExport::jvmti_oop_storage(), obj));
 }
 
-// Disable collection of VMObjectAlloc events
-NoJvmtiVMObjectAllocMark::NoJvmtiVMObjectAllocMark() : _collector(nullptr) {
-  // a no-op if VMObjectAlloc event is not enabled
-  if (!JvmtiExport::should_post_vm_object_alloc()) {
-    return;
-  }
+NoJvmtiEventsMark::NoJvmtiEventsMark() {
   Thread* thread = Thread::current_or_null();
   if (thread != nullptr && thread->is_Java_thread())  {
     JavaThread* current_thread = JavaThread::cast(thread);
-    JvmtiThreadState *state = current_thread->jvmti_thread_state();
-    if (state != nullptr) {
-      JvmtiVMObjectAllocEventCollector *collector;
-      collector = state->get_vm_object_alloc_event_collector();
-      if (collector != nullptr && collector->is_enabled()) {
-        _collector = collector;
-        _collector->set_enabled(false);
-      }
-    }
+    current_thread->disable_jvmti_events();
   }
 }
 
-// Re-Enable collection of VMObjectAlloc events (if previously enabled)
-NoJvmtiVMObjectAllocMark::~NoJvmtiVMObjectAllocMark() {
-  if (was_enabled()) {
-    _collector->set_enabled(true);
+NoJvmtiEventsMark::~NoJvmtiEventsMark() {
+  Thread* thread = Thread::current_or_null();
+  if (thread != nullptr && thread->is_Java_thread())  {
+    JavaThread* current_thread = JavaThread::cast(thread);
+    current_thread->enable_jvmti_events();
   }
 };
 

src/hotspot/share/prims/jvmtiExport.hpp

@@ -585,29 +585,12 @@ class JvmtiSampledObjectAllocEventCollector : public JvmtiObjectAllocEventCollec
   static bool object_alloc_is_safe_to_sample() NOT_JVMTI_RETURN_(false);
 };
 
-// Marker class to disable the posting of VMObjectAlloc events
-// within its scope.
-//
-// Usage :-
-//
-// {
-//   NoJvmtiVMObjectAllocMark njm;
-//   :
-//   // VMObjAlloc event will not be posted
-//   JvmtiExport::vm_object_alloc_event_collector(obj);
-//   :
-// }
-
-class NoJvmtiVMObjectAllocMark : public StackObj {
- private:
-  // enclosing collector if enabled, null otherwise
-  JvmtiVMObjectAllocEventCollector *_collector;
-
-  bool was_enabled()    { return _collector != nullptr; }
+// Marker class to temporary disable posting of jvmti events.
+class NoJvmtiEventsMark : public StackObj {
 
  public:
-  NoJvmtiVMObjectAllocMark() NOT_JVMTI_RETURN;
-  ~NoJvmtiVMObjectAllocMark() NOT_JVMTI_RETURN;
+  NoJvmtiEventsMark() NOT_JVMTI_RETURN;
+  ~NoJvmtiEventsMark() NOT_JVMTI_RETURN;
 };
 
 

src/hotspot/share/runtime/javaThread.cpp

@@ -451,6 +451,7 @@ JavaThread::JavaThread(MemTag mem_tag) :
   _is_in_VTMS_transition(false),
   _is_disable_suspend(false),
   _is_in_java_upcall(false),
+  _jvmti_events_disabled(0),
   _VTMS_transition_mark(false),
   _on_monitor_waited_event(false),
   _contended_entered_monitor(nullptr),

src/hotspot/share/runtime/javaThread.hpp

@@ -325,6 +325,7 @@ class JavaThread: public Thread {
   bool                  _is_in_VTMS_transition;          // thread is in virtual thread mount state transition
   bool                  _is_disable_suspend;             // JVMTI suspend is temporarily disabled; used on current thread only
   bool                  _is_in_java_upcall;              // JVMTI is doing a Java upcall, so JVMTI events must be hidden
+  int                   _jvmti_events_disabled;          // JVMTI events disabled manually
   bool                  _VTMS_transition_mark;           // used for sync between VTMS transitions and disablers
   bool                  _on_monitor_waited_event;        // Avoid callee arg processing for enterSpecial when posting waited event
   ObjectMonitor*        _contended_entered_monitor;      // Monitor for pending monitor_contended_entered callback
@@ -748,10 +749,13 @@ public:
   void set_is_in_VTMS_transition(bool val);
 
   bool is_disable_suspend() const                { return _is_disable_suspend; }
-  void toggle_is_disable_suspend()               { _is_disable_suspend = !_is_disable_suspend; };
+  void toggle_is_disable_suspend()               { _is_disable_suspend = !_is_disable_suspend; }
 
   bool is_in_java_upcall() const                 { return _is_in_java_upcall; }
-  void toggle_is_in_java_upcall()                { _is_in_java_upcall = !_is_in_java_upcall; };
+  void toggle_is_in_java_upcall()                { _is_in_java_upcall = !_is_in_java_upcall; }
+
+  void disable_jvmti_events()                    { _jvmti_events_disabled++; }
+  void enable_jvmti_events()                     { _jvmti_events_disabled--; }
 
   bool VTMS_transition_mark() const              { return AtomicAccess::load(&_VTMS_transition_mark); }
   void set_VTMS_transition_mark(bool val)        { AtomicAccess::store(&_VTMS_transition_mark, val); }
@@ -760,7 +764,9 @@ public:
   // - is in a VTMS transition (_is_in_VTMS_transition)
   // - is in an interruptLock or similar critical section (_is_disable_suspend)
   // - JVMTI is making a Java upcall (_is_in_java_upcall)
-  bool should_hide_jvmti_events() const          { return _is_in_VTMS_transition || _is_disable_suspend || _is_in_java_upcall; }
+  bool should_hide_jvmti_events() const {
+    return _is_in_VTMS_transition || _is_disable_suspend || _is_in_java_upcall || _jvmti_events_disabled != 0;
+  }
 
   bool on_monitor_waited_event()             { return _on_monitor_waited_event; }
   void set_on_monitor_waited_event(bool val) { _on_monitor_waited_event = val; }

test/hotspot/jtreg/serviceability/jvmti/events/SampledObjectAlloc/SamplingDuringInit/SamplingDuringInit.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright (c) 2025, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+
+/**
+ * @test
+ * @bug 8372039
+ * @summary The test verifies that object allocation sampling is disabled during AOT.
+ *
+ * Don't remove 'modules' line, it triggers the crash.
+ * @modules java.management
+ *
+ * @run main/othervm/native -agentlib:SamplingDuringInit SamplingDuringInit
+ * @run main/othervm/native -agentlib:SamplingDuringInit -XX:-UseCompressedOops SamplingDuringInit
+ */
+
+public class SamplingDuringInit {
+
+    public static Object[] tmp = new Object[1000];
+    public static void main(String[] args) throws Exception {
+        // Allocate some objects to trigger Sampling even if
+        // all JDK classes are preloaded.
+        for (int i = 0; i < tmp.length; i++) {
+            tmp[i] = new String("tmp" + i);
+        }
+    }
+}

test/hotspot/jtreg/serviceability/jvmti/events/SampledObjectAlloc/SamplingDuringInit/libSamplingDuringInit.cpp

@@ -0,0 +1,95 @@
+/*
+ * Copyright (c) 2025, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+
+#include "jvmti.h"
+#include "jvmti_common.hpp"
+
+#include <atomic>
+
+extern "C" {
+
+// SampledObjectAlloc event might be triggered on any thread
+static std::atomic<int> events_counter(0);
+
+JNIEXPORT void JNICALL
+SampledObjectAlloc(jvmtiEnv *jvmti, JNIEnv* jni, jthread thread, jobject object, jclass object_klass, jlong size) {
+  events_counter++;
+  LOG("Sampled object, events_counter = %d\n", events_counter.load());
+}
+
+void JNICALL
+VMDeath(jvmtiEnv *jvmti, JNIEnv* jni) {
+  if (events_counter == 0) {
+    fatal(jni, "SampledObjectAlloc events counter shouldn't be zero");
+  }
+}
+
+jint Agent_Initialize(JavaVM *jvm, char *options, void *reserved) {
+  jvmtiEnv* jvmti = nullptr;
+  jvmtiCapabilities caps;
+  jvmtiEventCallbacks callbacks;
+  jvmtiError err;
+  jint res;
+
+  LOG("AGENT INIT");
+  res = jvm->GetEnv((void **) &jvmti, JVMTI_VERSION_9);
+  if (res != JNI_OK || jvmti == nullptr) {
+    LOG("Wrong result of a valid call to GetEnv!\n");
+    return JNI_ERR;
+  }
+
+  memset(&caps, 0, sizeof(caps));
+  caps.can_generate_sampled_object_alloc_events = 1;
+  if (jvmti->AddCapabilities(&caps) != JVMTI_ERROR_NONE) {
+    return JNI_ERR;
+  }
+
+  memset(&callbacks, 0, sizeof(callbacks));
+  callbacks.SampledObjectAlloc = &SampledObjectAlloc;
+  callbacks.VMDeath = &VMDeath;
+
+  err = jvmti->SetEventCallbacks(&callbacks, sizeof(callbacks));
+  check_jvmti_error(err, "SetEventCallbacks");
+  /*
+   * Interval should be small enough to triggger sampling event while objects are init by VM.
+   */
+  err = jvmti->SetHeapSamplingInterval(10);
+  check_jvmti_error(err, "SetHeapSamplingInterval");
+
+  err = jvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_VM_DEATH, nullptr);
+  check_jvmti_error(err, "SetEventNotificationMode");
+
+  err = jvmti->SetEventNotificationMode(JVMTI_ENABLE, JVMTI_EVENT_SAMPLED_OBJECT_ALLOC, nullptr);
+  check_jvmti_error(err, "SetEventNotificationMode");
+
+  return JNI_OK;
+}
+
+JNIEXPORT jint JNICALL Agent_OnLoad(JavaVM *jvm, char *options, void *reserved) {
+  return Agent_Initialize(jvm, options, reserved);
+}
+
+JNIEXPORT jint JNICALL Agent_OnAttach(JavaVM *jvm, char *options, void *reserved) {
+  return Agent_Initialize(jvm, options, reserved);
+}
+}