8351748: Add class init barrier to AOT-cached Method/Var Handles

Reviewed-by: vlivanov, liach

Author: iklam

src/hotspot/share/cds/aotClassInitializer.cpp

@@ -26,11 +26,16 @@
 #include "cds/archiveBuilder.hpp"
 #include "cds/cdsConfig.hpp"
 #include "cds/heapShared.hpp"
+#include "classfile/symbolTable.hpp"
+#include "classfile/systemDictionaryShared.hpp"
 #include "classfile/vmSymbols.hpp"
 #include "oops/instanceKlass.inline.hpp"
 #include "oops/symbol.hpp"
+#include "runtime/java.hpp"
 #include "runtime/javaCalls.hpp"
 
+DEBUG_ONLY(InstanceKlass* _aot_init_class = nullptr;)
+
 // Detector for class names we wish to handle specially.
 // It is either an exact string match or a string prefix match.
 class AOTClassInitializer::AllowedSpec {
@@ -93,12 +98,12 @@ bool AOTClassInitializer::is_allowed(AllowedSpec* specs, InstanceKlass* ik) {
 
 
 bool AOTClassInitializer::can_archive_initialized_mirror(InstanceKlass* ik) {
-  assert(!ArchiveBuilder::current()->is_in_buffer_space(ik), "must be source klass");
+  assert(!ArchiveBuilder::is_active() || !ArchiveBuilder::current()->is_in_buffer_space(ik), "must be source klass");
   if (!CDSConfig::is_initing_classes_at_dump_time()) {
     return false;
   }
 
-  if (!ik->is_initialized()) {
+  if (!ik->is_initialized() && !ik->is_being_initialized()) {
     return false;
   }
 
@@ -293,9 +298,11 @@ bool AOTClassInitializer::can_archive_initialized_mirror(InstanceKlass* ik) {
       {"java/lang/invoke/LambdaForm"},
       {"java/lang/invoke/LambdaForm$Holder"},                 // UNSAFE.ensureClassInitialized()
       {"java/lang/invoke/LambdaForm$NamedFunction"},
+      {"java/lang/invoke/LambdaMetafactory"},
       {"java/lang/invoke/MethodHandle"},
       {"java/lang/invoke/MethodHandles"},
       {"java/lang/invoke/SimpleMethodHandle"},
+      {"java/lang/invoke/StringConcatFactory"},
       {"java/util/Collections"},
       {"java/util/stream/Collectors"},
       {"jdk/internal/constant/ConstantUtils"},
@@ -315,6 +322,12 @@ bool AOTClassInitializer::can_archive_initialized_mirror(InstanceKlass* ik) {
     }
   }
 
+#ifdef ASSERT
+  if (ik == _aot_init_class) {
+    return true;
+  }
+#endif
+
   return false;
 }
 
@@ -345,3 +358,33 @@ void AOTClassInitializer::call_runtime_setup(JavaThread* current, InstanceKlass*
     }
   }
 }
+
+#ifdef ASSERT
+void AOTClassInitializer::init_test_class(TRAPS) {
+  // -XX:AOTInitTestClass is used in regression tests for adding additional AOT-initialized classes
+  // and heap objects into the AOT cache. The tests must be carefully written to avoid including
+  // any classes that cannot be AOT-initialized.
+  //
+  // -XX:AOTInitTestClass is NOT a general mechanism for including user-defined objects into
+  // the AOT cache. Therefore, this option is NOT available in product JVM.
+  if (AOTInitTestClass != nullptr && CDSConfig::is_initing_classes_at_dump_time()) {
+    log_info(cds)("Debug build only: force initialization of AOTInitTestClass %s", AOTInitTestClass);
+    TempNewSymbol class_name = SymbolTable::new_symbol(AOTInitTestClass);
+    Handle app_loader(THREAD, SystemDictionary::java_system_loader());
+    Klass* k = SystemDictionary::resolve_or_null(class_name, app_loader, CHECK);
+    if (k == nullptr) {
+      vm_exit_during_initialization("AOTInitTestClass not found", AOTInitTestClass);
+    }
+    if (!k->is_instance_klass()) {
+      vm_exit_during_initialization("Invalid name for AOTInitTestClass", AOTInitTestClass);
+    }
+
+    _aot_init_class = InstanceKlass::cast(k);
+    _aot_init_class->initialize(CHECK);
+  }
+}
+
+bool AOTClassInitializer::has_test_class() {
+  return _aot_init_class != nullptr;
+}
+#endif

src/hotspot/share/cds/aotClassInitializer.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2024, 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
@@ -41,6 +41,10 @@ class AOTClassInitializer : AllStatic {
 
   static bool is_runtime_setup_required(InstanceKlass* ik);
   static void call_runtime_setup(JavaThread* current, InstanceKlass* ik);
+
+  // Support for regression testing. Available in debug builds only.
+  static void init_test_class(TRAPS) NOT_DEBUG_RETURN;
+  static bool has_test_class() NOT_DEBUG({ return false; });
 };
 
 #endif // SHARE_CDS_AOTCLASSINITIALIZER_HPP

src/hotspot/share/cds/cdsConfig.cpp

@@ -60,6 +60,7 @@ JavaThread* CDSConfig::_dumper_thread = nullptr;
 int CDSConfig::get_status() {
   assert(Universe::is_fully_initialized(), "status is finalized only after Universe is initialized");
   return (is_dumping_archive()              ? IS_DUMPING_ARCHIVE : 0) |
+         (is_dumping_method_handles()       ? IS_DUMPING_METHOD_HANDLES : 0) |
          (is_dumping_static_archive()       ? IS_DUMPING_STATIC_ARCHIVE : 0) |
          (is_logging_lambda_form_invokers() ? IS_LOGGING_LAMBDA_FORM_INVOKERS : 0) |
          (is_using_archive()                ? IS_USING_ARCHIVE : 0);

src/hotspot/share/cds/cdsConfig.hpp

@@ -68,9 +68,11 @@ class CDSConfig : public AllStatic {
 public:
   // Used by jdk.internal.misc.CDS.getCDSConfigStatus();
   static const int IS_DUMPING_ARCHIVE              = 1 << 0;
-  static const int IS_DUMPING_STATIC_ARCHIVE       = 1 << 1;
-  static const int IS_LOGGING_LAMBDA_FORM_INVOKERS = 1 << 2;
-  static const int IS_USING_ARCHIVE                = 1 << 3;
+  static const int IS_DUMPING_METHOD_HANDLES       = 1 << 1;
+  static const int IS_DUMPING_STATIC_ARCHIVE       = 1 << 2;
+  static const int IS_LOGGING_LAMBDA_FORM_INVOKERS = 1 << 3;
+  static const int IS_USING_ARCHIVE                = 1 << 4;
+
   static int get_status() NOT_CDS_RETURN_(0);
 
   // Initialization and command-line checking

src/hotspot/share/cds/cds_globals.hpp

@@ -71,6 +71,10 @@
           "\"archivedObjects\" of the specified class is stored in the "    \
           "CDS archive heap")                                               \
                                                                             \
+  develop(ccstr, AOTInitTestClass, nullptr,                                 \
+          "For JVM internal testing only. The specified class is stored "   \
+          "in the initialized state in the AOT cache ")                     \
+                                                                            \
   product(ccstr, DumpLoadedClassList, nullptr,                              \
           "Dump the names all loaded classes, that could be stored into "   \
           "the CDS archive, in the specified file")                         \

src/hotspot/share/cds/heapShared.cpp

@@ -138,7 +138,7 @@ static ArchivableStaticFieldInfo fmg_archive_subgraph_entry_fields[] = {
 KlassSubGraphInfo* HeapShared::_dump_time_special_subgraph;
 ArchivedKlassSubGraphInfoRecord* HeapShared::_run_time_special_subgraph;
 GrowableArrayCHeap<oop, mtClassShared>* HeapShared::_pending_roots = nullptr;
-GrowableArrayCHeap<OopHandle, mtClassShared>* HeapShared::_root_segments;
+GrowableArrayCHeap<OopHandle, mtClassShared>* HeapShared::_root_segments = nullptr;
 int HeapShared::_root_segment_max_size_elems;
 OopHandle HeapShared::_scratch_basic_type_mirrors[T_VOID+1];
 MetaspaceObjToOopHandleTable* HeapShared::_scratch_objects_table = nullptr;
@@ -225,10 +225,6 @@ int HeapShared::append_root(oop obj) {
   // No GC should happen since we aren't scanning _pending_roots.
   assert(Thread::current() == (Thread*)VMThread::vm_thread(), "should be in vm thread");
 
-  if (_pending_roots == nullptr) {
-    _pending_roots = new GrowableArrayCHeap<oop, mtClassShared>(500);
-  }
-
   return _pending_roots->append(obj);
 }
 
@@ -336,6 +332,12 @@ bool HeapShared::archive_object(oop obj, oop referrer, KlassSubGraphInfo* subgra
         if (mirror_k != nullptr) {
           AOTArtifactFinder::add_cached_class(mirror_k);
         }
+      } else if (java_lang_invoke_ResolvedMethodName::is_instance(obj)) {
+        Method* m = java_lang_invoke_ResolvedMethodName::vmtarget(obj);
+        if (m != nullptr) {
+          InstanceKlass* method_holder = m->method_holder();
+          AOTArtifactFinder::add_cached_class(method_holder);
+        }
       }
     }
 
@@ -402,6 +404,7 @@ objArrayOop HeapShared::scratch_resolved_references(ConstantPool* src) {
 
 void HeapShared::init_dumping() {
   _scratch_objects_table = new (mtClass)MetaspaceObjToOopHandleTable();
+  _pending_roots = new GrowableArrayCHeap<oop, mtClassShared>(500);
 }
 
 void HeapShared::init_scratch_objects_for_basic_type_mirrors(TRAPS) {
@@ -565,7 +568,7 @@ void HeapShared::copy_and_rescan_aot_inited_mirror(InstanceKlass* ik) {
     assert(success, "sanity");
   }
 
-  if (log_is_enabled(Info, cds, init)) {
+  if (log_is_enabled(Debug, cds, init)) {
     ResourceMark rm;
     log_debug(cds, init)("copied %3d field(s) in aot-initialized mirror %s%s%s", nfields, ik->external_name(),
                          ik->is_hidden() ? " (hidden)" : "",
@@ -781,8 +784,11 @@ void KlassSubGraphInfo::add_subgraph_object_klass(Klass* orig_k) {
 #ifdef ASSERT
     InstanceKlass* ik = InstanceKlass::cast(orig_k);
     if (CDSConfig::is_dumping_method_handles()) {
+      // -XX:AOTInitTestClass must be used carefully in regression tests to
+      // include only classes that are safe to aot-initialize.
       assert(ik->class_loader() == nullptr ||
-             HeapShared::is_lambda_proxy_klass(ik),
+             HeapShared::is_lambda_proxy_klass(ik) ||
+             AOTClassInitializer::has_test_class(),
             "we can archive only instances of boot classes or lambda proxy classes");
     } else {
       assert(ik->class_loader() == nullptr, "must be boot class");
@@ -827,6 +833,13 @@ void KlassSubGraphInfo::add_subgraph_object_klass(Klass* orig_k) {
 }
 
 void KlassSubGraphInfo::check_allowed_klass(InstanceKlass* ik) {
+#ifndef PRODUCT
+  if (AOTClassInitializer::has_test_class()) {
+    // The tests can cache arbitrary types of objects.
+    return;
+  }
+#endif
+
   if (ik->module()->name() == vmSymbols::java_base()) {
     assert(ik->package() != nullptr, "classes in java.base cannot be in unnamed package");
     return;

src/hotspot/share/cds/metaspaceShared.cpp

@@ -23,6 +23,7 @@
  */
 
 #include "cds/aotArtifactFinder.hpp"
+#include "cds/aotClassInitializer.hpp"
 #include "cds/aotClassLinker.hpp"
 #include "cds/aotClassLocation.hpp"
 #include "cds/aotConstantPoolResolver.hpp"
@@ -736,6 +737,7 @@ bool MetaspaceShared::link_class_for_cds(InstanceKlass* ik, TRAPS) {
 
 void MetaspaceShared::link_shared_classes(bool jcmd_request, TRAPS) {
   AOTClassLinker::initialize();
+  AOTClassInitializer::init_test_class(CHECK);
 
   if (!jcmd_request && !CDSConfig::is_dumping_final_static_archive()) {
     LambdaFormInvokers::regenerate_holder_classes(CHECK);

src/hotspot/share/include/jvm.h

@@ -218,6 +218,9 @@ JVM_DumpClassListToFile(JNIEnv* env, jstring fileName);
 JNIEXPORT void JNICALL
 JVM_DumpDynamicArchive(JNIEnv* env, jstring archiveName);
 
+JNIEXPORT jboolean JNICALL
+JVM_NeedsClassInitBarrierForCDS(JNIEnv* env, jclass cls);
+
 /*
  * java.lang.Throwable
  */

src/hotspot/share/prims/jvm.cpp

@@ -22,6 +22,7 @@
  *
  */
 
+#include "cds/aotClassInitializer.hpp"
 #include "cds/cdsConfig.hpp"
 #include "cds/classListParser.hpp"
 #include "cds/classListWriter.hpp"
@@ -3349,7 +3350,6 @@ JVM_END
 
 JVM_ENTRY(void, JVM_InitializeFromArchive(JNIEnv* env, jclass cls))
   Klass* k = java_lang_Class::as_Klass(JNIHandles::resolve(cls));
-  assert(k->is_klass(), "just checking");
   HeapShared::initialize_from_archived_subgraph(THREAD, k);
 JVM_END
 
@@ -3514,6 +3514,29 @@ JVM_ENTRY(void, JVM_DumpDynamicArchive(JNIEnv *env, jstring archiveName))
 #endif // INCLUDE_CDS
 JVM_END
 
+JVM_ENTRY(jboolean, JVM_NeedsClassInitBarrierForCDS(JNIEnv* env, jclass cls))
+#if INCLUDE_CDS
+  Klass* k = java_lang_Class::as_Klass(JNIHandles::resolve(cls));
+  if (!k->is_instance_klass()) {
+    return false;
+  } else {
+    if (InstanceKlass::cast(k)->is_enum_subclass() ||
+        AOTClassInitializer::can_archive_initialized_mirror(InstanceKlass::cast(k))) {
+      // This class will be cached in AOT-initialized state. No need for init barriers.
+      return false;
+    } else {
+      // If we cannot cache the class in AOT-initialized state, java.lang.invoke handles
+      // must emit barriers to ensure class initialization during production run.
+      ResourceMark rm(THREAD);
+      log_debug(cds)("NeedsClassInitBarrierForCDS: %s", k->external_name());
+      return true;
+    }
+  }
+#else
+  return false;
+#endif // INCLUDE_CDS
+JVM_END
+
 // Returns an array of all live Thread objects (VM internal JavaThreads,
 // jvmti agent threads, and JNI attaching threads  are skipped)
 // See CR 6404306 regarding JNI attaching threads

src/java.base/share/classes/java/lang/invoke/DirectMethodHandle.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2008, 2024, Oracle and/or its affiliates. All rights reserved.
+ * Copyright (c) 2008, 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
@@ -25,6 +25,7 @@
 
 package java.lang.invoke;
 
+import jdk.internal.misc.CDS;
 import jdk.internal.misc.Unsafe;
 import jdk.internal.vm.annotation.ForceInline;
 import jdk.internal.vm.annotation.Stable;
@@ -367,9 +368,9 @@ static boolean shouldBeInitialized(MemberName member) {
             // It is a system class.  It is probably in the process of
             // being initialized, but we will help it along just to be safe.
             UNSAFE.ensureClassInitialized(cls);
-            return false;
+            return CDS.needsClassInitBarrier(cls);
         }
-        return UNSAFE.shouldBeInitialized(cls);
+        return UNSAFE.shouldBeInitialized(cls) || CDS.needsClassInitBarrier(cls);
     }
 
     private void ensureInitialized() {