fix(profiler): pre-register vtable receiver classes via JVMTI (PROF-14618)

ddprof-lib/src/main/cpp/profiler.cpp

@@ -1084,6 +1084,9 @@ Error Profiler::start(Arguments &args, bool reset) {
     // it is being cleaned up
     _class_map_lock.lock();
     _class_map.clear();
+    if (args._features.vtable_target && VMStructs::hasClassNames()) {
+      preregisterLoadedClasses(VM::jvmti());
+    }
     _class_map_lock.unlock();
 
     // Reset call trace storage
@@ -1399,6 +1402,9 @@ Error Profiler::dump(const char *path, const int length) {
     // Reset classmap
     _class_map_lock.lock();
     _class_map.clear();
+    if (_features.vtable_target && VMStructs::hasClassNames()) {
+      preregisterLoadedClasses(VM::jvmti());
+    }
     _class_map_lock.unlock();
 
     _thread_info.clearAll(thread_ids);
@@ -1539,6 +1545,55 @@ int Profiler::lookupClass(const char *key, size_t length) {
   return -1;
 }
 
+void Profiler::preregisterLoadedClasses(jvmtiEnv* jvmti) {
+  if (jvmti == nullptr) {
+    return;
+  }
+  JNIEnv* jni = VM::jni();
+  jint class_count = 0;
+  jclass* classes = nullptr;
+  jvmtiError err = jvmti->GetLoadedClasses(&class_count, &classes);
+  if (err != JVMTI_ERROR_NONE) {
+    Log::warn("preregisterLoadedClasses: GetLoadedClasses failed (%d) — vtable-target frames will miss until ClassPrepare fires", err);
+    return;
+  }
+  if (classes == nullptr) {
+    return;
+  }
+  // Each classes[i] is a JNI local reference allocated by JVMTI; delete it at
+  // every loop exit so the local-reference table does not grow across repeated
+  // start/dump calls on large applications.
+  for (jint i = 0; i < class_count; i++) {
+    char* sig = nullptr;
+    if (jvmti->GetClassSignature(classes[i], &sig, nullptr) != JVMTI_ERROR_NONE) {
+      if (jni != nullptr) jni->DeleteLocalRef(classes[i]);
+      continue;
+    }
+    if (sig == nullptr) {
+      if (jni != nullptr) jni->DeleteLocalRef(classes[i]);
+      continue;
+    }
+    // Only 'L'-type (reference) signatures can ever match vtable_target lookup
+    // keys; skip primitives, arrays, and other non-reference signatures.
+    if (sig[0] != 'L') {
+      jvmti->Deallocate(reinterpret_cast<unsigned char*>(sig));
+      if (jni != nullptr) jni->DeleteLocalRef(classes[i]);
+      continue;
+    }
+    const char* slice = nullptr;
+    size_t slice_len = 0;
+    if (ObjectSampler::normalizeClassSignature(sig, &slice, &slice_len)) {
+      // Caller holds _class_map_lock exclusively — call _class_map.lookup
+      // directly. Do NOT call lookupClass(), which would re-attempt
+      // tryLockShared() and deadlock.
+      (void)_class_map.lookup(slice, slice_len);
+    }
+    jvmti->Deallocate(reinterpret_cast<unsigned char*>(sig));
+    if (jni != nullptr) jni->DeleteLocalRef(classes[i]);
+  }
+  jvmti->Deallocate(reinterpret_cast<unsigned char*>(classes));
+}
+
 int Profiler::status(char* status, int max_len) {
   return snprintf(status, max_len,
     "== Java-Profiler Status ===\n"

ddprof-lib/src/main/cpp/profiler.h

@@ -213,6 +213,13 @@ class alignas(alignof(SpinLock)) Profiler {
 
   const char* cstack() const;
   int lookupClass(const char *key, size_t length);
+  // Pre-populate _class_map with all currently-loaded 'L'-type (reference)
+  // class signatures so that signal-safe lookups in walkVM (vtable_target) can
+  // resolve them without ever needing to malloc. Primitives and arrays are
+  // skipped — they never match vtable lookup keys. Caller MUST hold
+  // _class_map_lock EXCLUSIVELY. Runs on a JVM thread (never in a signal
+  // handler).
+  void preregisterLoadedClasses(jvmtiEnv* jvmti);
   void processCallTraces(std::function<void(const std::unordered_set<CallTrace*>&)> processor) {
     if (!_omit_stacktraces) {
       _call_trace_storage.processTraces(processor);

ddprof-lib/src/main/cpp/vmEntry.cpp

@@ -575,6 +575,44 @@ void VM::loadAllMethodIDs(jvmtiEnv *jvmti, JNIEnv *jni) {
     }
 }
 
+void JNICALL VM::ClassPrepare(jvmtiEnv* jvmti, JNIEnv* jni, jthread thread,
+                               jclass klass) {
+  loadMethodIDs(jvmti, jni, klass);
+
+  // Pre-populate _class_map for vtable_target signal-safe lookup. ClassPrepare
+  // runs on a JVM thread (never a signal handler), so we take a blocking
+  // shared lock via classMapSharedGuard() rather than the best-effort
+  // tryLockShared() in lookupClass(). This ensures newly prepared classes are
+  // never silently skipped while an exclusive dump/reset is in flight. Gate on
+  // vtable_target to avoid overhead when the feature is disabled.
+  Profiler* profiler = Profiler::instance();
+  if (profiler == nullptr || !profiler->stackWalkFeatures().vtable_target) {
+    return;
+  }
+  char* sig = nullptr;
+  jvmtiError err = jvmti->GetClassSignature(klass, &sig, nullptr);
+  if (err != JVMTI_ERROR_NONE) {
+    Log::warn("ClassPrepare: GetClassSignature failed (%d) — class skipped for vtable-target pre-registration", err);
+    return;
+  }
+  if (sig == nullptr) {
+    return;
+  }
+  // Only 'L'-type (reference) signatures can ever match vtable_target lookup
+  // keys; skip primitives and arrays.
+  if (sig[0] != 'L') {
+    jvmti->Deallocate(reinterpret_cast<unsigned char*>(sig));
+    return;
+  }
+  const char* slice = nullptr;
+  size_t slice_len = 0;
+  if (ObjectSampler::normalizeClassSignature(sig, &slice, &slice_len)) {
+    SharedLockGuard guard = profiler->classMapSharedGuard();
+    (void)profiler->classMap()->lookup(slice, slice_len);
+  }
+  jvmti->Deallocate(reinterpret_cast<unsigned char*>(sig));
+}
+
 void JNICALL VM::VMInit(jvmtiEnv* jvmti, JNIEnv* jni, jthread thread) {
     ready(jvmti, jni);
     loadAllMethodIDs(jvmti, jni);

ddprof-lib/src/main/cpp/vmEntry.h

@@ -198,10 +198,8 @@ class VM {
     // Needed only for AsyncGetCallTrace support
   }
 
-  static void JNICALL ClassPrepare(jvmtiEnv *jvmti, JNIEnv *jni, jthread thread,
-                                   jclass klass) {
-    loadMethodIDs(jvmti, jni, klass);
-  }
+  static void JNICALL ClassPrepare(jvmtiEnv* jvmti, JNIEnv* jni, jthread thread,
+                                   jclass klass);
 
   static jvmtiError JNICALL
   RedefineClassesHook(jvmtiEnv *jvmti, jint class_count,

ddprof-lib/src/test/cpp/dictionary_concurrent_ut.cpp

@@ -250,6 +250,41 @@ TEST(DictionaryConcurrent, SignalHandlerBoundedLookupVsDumpClear) {
     EXPECT_GT(total_clears.load(), 0L);
 }
 
+// (4a) Single-threaded bulk insert (mimicking the inserts preregisterLoadedClasses
+// performs while holding _class_map_lock exclusively in production) followed by
+// read-only bounded_lookup(0) on each inserted key. Verifies the pre-registration
+// contract: every key inserted via Dictionary::lookup() is subsequently visible
+// to bounded_lookup(0). This test takes no external lock — the production
+// exclusive-lock protocol is exercised by the other tests in this file.
+TEST(DictionaryConcurrent, BulkInsertThenBoundedLookupHitsMirrorInsertedIds) {
+    Dictionary dict(/*id=*/0);
+
+    constexpr int kBulk = 50;
+    char keys[kBulk][64];
+    unsigned int inserted_ids[kBulk];
+
+    // Bulk insert — mirrors the per-class lookup() calls preregisterLoadedClasses
+    // issues from a JVM thread.
+    for (int i = 0; i < kBulk; i++) {
+        snprintf(keys[i], sizeof(keys[i]), "java/util/BulkClass%d", i);
+        inserted_ids[i] = dict.lookup(keys[i], strlen(keys[i]));
+        ASSERT_NE(0u, inserted_ids[i]);
+        ASSERT_NE(static_cast<unsigned int>(INT_MAX), inserted_ids[i]);
+    }
+
+    // Read back with bounded_lookup(0) — must return the same id, no malloc.
+    for (int i = 0; i < kBulk; i++) {
+        unsigned int found = dict.bounded_lookup(keys[i], strlen(keys[i]), 0);
+        EXPECT_EQ(inserted_ids[i], found)
+            << "bounded_lookup returned wrong id for key " << keys[i];
+    }
+
+    // A never-inserted key must return INT_MAX.
+    const char* absent = "java/util/NeverInserted";
+    EXPECT_EQ(static_cast<unsigned int>(INT_MAX),
+              dict.bounded_lookup(absent, strlen(absent), 0));
+}
+
 // (4) Same race as (3) but using BoundedOptionalSharedLockGuard, which is the
 // guard classMapTrySharedGuard() now returns in hotspotSupport.cpp. The bounded
 // variant may fail spuriously under reader pressure (≤5 CAS attempts); this