8368174: Proactive initialization of @AOTSafeClassInitializer classes

src/hotspot/share/cds/aotClassInitializer.cpp

@@ -48,7 +48,14 @@ bool AOTClassInitializer::can_archive_initialized_mirror(InstanceKlass* ik) {
     ik = RegeneratedClasses::get_original_object(ik);
   }
 
+  check_aot_annotations(ik);
+
   if (!ik->is_initialized() && !ik->is_being_initialized()) {
+    if (ik->has_aot_safe_initializer()) {
+      ResourceMark rm;
+      log_info(aot, init)("Class %s is annotated with @AOTSafeClassInitializer but has not been initialized",
+                          ik->external_name());
+    }
     return false;
   }
 
@@ -245,6 +252,56 @@ void AOTClassInitializer::call_runtime_setup(JavaThread* current, InstanceKlass*
   }
 }
 
+template <typename FUNCTION>
+void require_annotation_for_super_types(InstanceKlass* ik, const char* annotation, FUNCTION func) {
+  if (log_is_enabled(Info, aot, init)) {
+    ResourceMark rm;
+    log_info(aot, init)("Found %s class %s", annotation, ik->external_name());
+  }
+
+  // Since ik has this annotation, we require that
+  //   - all super classes must have this annotation
+  //   - all super interfaces that are interface_needs_clinit_execution_as_super()
+  //     must have this annotation
+  // This avoid the situation where in the production run, we run the <clinit>
+  // of a supertype but not the <clinit> of ik
+
+  InstanceKlass* super = ik->java_super();
+  if (super != nullptr && !func(super)) {
+    ResourceMark rm;
+    log_error(aot, init)("Missing %s in superclass %s for class %s",
+                         annotation, super->external_name(), ik->external_name());
+    AOTMetaspace::unrecoverable_writing_error();
+  }
+
+  int len = ik->local_interfaces()->length();
+  for (int i = 0; i < len; i++) {
+    InstanceKlass* intf = ik->local_interfaces()->at(i);
+    if (intf->interface_needs_clinit_execution_as_super() && !func(intf)) {
+      ResourceMark rm;
+      log_error(aot, init)("Missing %s in superinterface %s for class %s",
+                           annotation, intf->external_name(), ik->external_name());
+      AOTMetaspace::unrecoverable_writing_error();
+    }
+  }
+}
+
+void AOTClassInitializer::check_aot_annotations(InstanceKlass* ik) {
+  if (ik->has_aot_safe_initializer()) {
+    require_annotation_for_super_types(ik, "@AOTSafeClassInitializer", [&] (const InstanceKlass* supertype) {
+      return supertype->has_aot_safe_initializer();
+    });
+  } else {
+    // @AOTRuntimeSetup only meaningful in @AOTSafeClassInitializer
+    if (ik->is_runtime_setup_required()) {
+      ResourceMark rm;
+      log_error(aot, init)("@AOTRuntimeSetup meaningless in non-@AOTSafeClassInitializer class %s",
+                           ik->external_name());
+    }
+  }
+}
+
+
 #ifdef ASSERT
 void AOTClassInitializer::init_test_class(TRAPS) {
   // -XX:AOTInitTestClass is used in regression tests for adding additional AOT-initialized classes

src/hotspot/share/cds/aotClassInitializer.hpp

@@ -31,6 +31,9 @@
 class InstanceKlass;
 
 class AOTClassInitializer : AllStatic {
+
+  static void check_aot_annotations(InstanceKlass* ik);
+
 public:
   // Called by heapShared.cpp to see if src_ik->java_mirror() can be archived in
   // the initialized state.

src/hotspot/share/cds/aotMetaspace.cpp

@@ -804,6 +804,23 @@ void AOTMetaspace::link_shared_classes(TRAPS) {
   AOTClassLinker::initialize();
   AOTClassInitializer::init_test_class(CHECK);
 
+  if (CDSConfig::is_dumping_final_static_archive()) {
+    // - Load and link all classes used in the training run.
+    // - Initialize @AOTSafeClassInitializer classes that were
+    //   initialized in the training run.
+    // - Perform per-class optimization such as AOT-resolution of
+    //   constant pool entries that were resolved during the training run.
+    FinalImageRecipes::apply_recipes(CHECK);
+
+    // Because the AOT assembly phase does not run the same exact code as in the
+    // training run (e.g., we use different lambda form invoker classes;
+    // generated lambda form classes are not recorded in FinalImageRecipes),
+    // the recipes do not cover all classes that have been loaded so far. As
+    // a result, we might have some unlinked classes at this point. Since we
+    // require cached classes to be linked, all such classes will be linked
+    // by the following step.
+  }
+
   link_all_loaded_classes(THREAD);
 
   // Eargerly resolve all string constants in constant pools
@@ -817,10 +834,6 @@ void AOTMetaspace::link_shared_classes(TRAPS) {
       AOTConstantPoolResolver::preresolve_string_cp_entries(ik, CHECK);
     }
   }
-
-  if (CDSConfig::is_dumping_final_static_archive()) {
-    FinalImageRecipes::apply_recipes(CHECK);
-  }
 }
 
 void AOTMetaspace::dump_static_archive(TRAPS) {

src/hotspot/share/cds/finalImageRecipes.cpp

@@ -57,7 +57,7 @@ void FinalImageRecipes::record_recipes_for_constantpool() {
   // ignored during the final image assembly.
 
   GrowableArray<Array<int>*> tmp_cp_recipes;
-  GrowableArray<int> tmp_cp_flags;
+  GrowableArray<int> tmp_flags;
 
   GrowableArray<Klass*>* klasses = ArchiveBuilder::current()->klasses();
   for (int i = 0; i < klasses->length(); i++) {
@@ -70,12 +70,16 @@ void FinalImageRecipes::record_recipes_for_constantpool() {
       ConstantPool* cp = ik->constants();
       ConstantPoolCache* cp_cache = cp->cache();
 
+      if (ik->is_initialized()) {
+        flags |= WAS_INITED;
+      }
+
       for (int cp_index = 1; cp_index < cp->length(); cp_index++) { // Index 0 is unused
         if (cp->tag_at(cp_index).value() == JVM_CONSTANT_Class) {
           Klass* k = cp->resolved_klass_at(cp_index);
           if (k->is_instance_klass()) {
             cp_indices.append(cp_index);
-            flags |= HAS_CLASS;
+            flags |= CP_RESOLVE_CLASS;
           }
         }
       }
@@ -88,7 +92,7 @@ void FinalImageRecipes::record_recipes_for_constantpool() {
             if (rfe->is_resolved(Bytecodes::_getfield) ||
                 rfe->is_resolved(Bytecodes::_putfield)) {
               cp_indices.append(rfe->constant_pool_index());
-              flags |= HAS_FIELD_AND_METHOD;
+              flags |= CP_RESOLVE_FIELD_AND_METHOD;
             }
           }
         }
@@ -103,7 +107,7 @@ void FinalImageRecipes::record_recipes_for_constantpool() {
                 rme->is_resolved(Bytecodes::_invokestatic) ||
                 rme->is_resolved(Bytecodes::_invokehandle)) {
               cp_indices.append(rme->constant_pool_index());
-              flags |= HAS_FIELD_AND_METHOD;
+              flags |= CP_RESOLVE_FIELD_AND_METHOD;
             }
           }
         }
@@ -115,7 +119,7 @@ void FinalImageRecipes::record_recipes_for_constantpool() {
             int cp_index = rie->constant_pool_index();
             if (rie->is_resolved()) {
               cp_indices.append(cp_index);
-              flags |= HAS_INDY;
+              flags |= CP_RESOLVE_INDY;
             }
           }
         }
@@ -127,22 +131,22 @@ void FinalImageRecipes::record_recipes_for_constantpool() {
     } else {
       tmp_cp_recipes.append(nullptr);
     }
-    tmp_cp_flags.append(flags);
+    tmp_flags.append(flags);
   }
 
   _cp_recipes = ArchiveUtils::archive_array(&tmp_cp_recipes);
   ArchivePtrMarker::mark_pointer(&_cp_recipes);
 
-  _cp_flags = ArchiveUtils::archive_array(&tmp_cp_flags);
-  ArchivePtrMarker::mark_pointer(&_cp_flags);
+  _flags = ArchiveUtils::archive_array(&tmp_flags);
+  ArchivePtrMarker::mark_pointer(&_flags);
 }
 
 void FinalImageRecipes::apply_recipes_for_constantpool(JavaThread* current) {
   assert(CDSConfig::is_dumping_final_static_archive(), "must be");
 
   for (int i = 0; i < _all_klasses->length(); i++) {
     Array<int>* cp_indices = _cp_recipes->at(i);
-    int flags = _cp_flags->at(i);
+    int flags = _flags->at(i);
     if (cp_indices != nullptr) {
       InstanceKlass* ik = InstanceKlass::cast(_all_klasses->at(i));
       if (ik->is_loaded()) {
@@ -152,13 +156,13 @@ void FinalImageRecipes::apply_recipes_for_constantpool(JavaThread* current) {
         for (int j = 0; j < cp_indices->length(); j++) {
           preresolve_list.at_put(cp_indices->at(j), true);
         }
-        if ((flags & HAS_CLASS) != 0) {
+        if ((flags & CP_RESOLVE_CLASS) != 0) {
           AOTConstantPoolResolver::preresolve_class_cp_entries(current, ik, &preresolve_list);
         }
-        if ((flags & HAS_FIELD_AND_METHOD) != 0) {
+        if ((flags & CP_RESOLVE_FIELD_AND_METHOD) != 0) {
           AOTConstantPoolResolver::preresolve_field_and_method_cp_entries(current, ik, &preresolve_list);
         }
-        if ((flags & HAS_INDY) != 0) {
+        if ((flags & CP_RESOLVE_INDY) != 0) {
           AOTConstantPoolResolver::preresolve_indy_cp_entries(current, ik, &preresolve_list);
         }
       }
@@ -171,6 +175,7 @@ void FinalImageRecipes::load_all_classes(TRAPS) {
   Handle class_loader(THREAD, SystemDictionary::java_system_loader());
   for (int i = 0; i < _all_klasses->length(); i++) {
     Klass* k = _all_klasses->at(i);
+    int flags = _flags->at(i);
     if (k->is_instance_klass()) {
       InstanceKlass* ik = InstanceKlass::cast(k);
       if (ik->defined_by_other_loaders()) {
@@ -188,6 +193,11 @@ void FinalImageRecipes::load_all_classes(TRAPS) {
         }
         assert(ik->is_loaded(), "must be");
         ik->link_class(CHECK);
+
+        if (ik->has_aot_safe_initializer() && (flags & WAS_INITED) != 0) {
+          assert(ik->class_loader() == nullptr, "supported only for boot classes for now");
+          ik->initialize(CHECK);
+        }
       }
     }
   }

src/hotspot/share/cds/finalImageRecipes.hpp

@@ -42,20 +42,21 @@ template <typename T> class Array;
 //   - The list of all classes that are stored in the AOTConfiguration file.
 //   - The list of all classes that require AOT resolution of invokedynamic call sites.
 class FinalImageRecipes {
-  static constexpr int HAS_CLASS            = 0x1;
-  static constexpr int HAS_FIELD_AND_METHOD = 0x2;
-  static constexpr int HAS_INDY             = 0x4;
+  static constexpr int CP_RESOLVE_CLASS            = 0x1 << 0; // CP has preresolved class entries
+  static constexpr int CP_RESOLVE_FIELD_AND_METHOD = 0x1 << 1; // CP has preresolved field/method entries
+  static constexpr int CP_RESOLVE_INDY             = 0x1 << 2; // CP has preresolved indy entries
+  static constexpr int WAS_INITED                  = 0x1 << 3; // Class was initialized during training run
 
   // A list of all the archived classes from the preimage. We want to transfer all of these
   // into the final image.
   Array<Klass*>* _all_klasses;
 
-  // For each klass k _all_klasses->at(i), _cp_recipes->at(i) lists all the {klass,field,method,indy}
-  // cp indices that were resolved for k during the training run.
+  // For each klass k _all_klasses->at(i): _cp_recipes->at(i) lists all the {klass,field,method,indy}
+  // cp indices that were resolved for k during the training run; _flags->at(i) has extra info about k.
   Array<Array<int>*>* _cp_recipes;
-  Array<int>* _cp_flags;
+  Array<int>* _flags;
 
-  FinalImageRecipes() : _all_klasses(nullptr), _cp_recipes(nullptr), _cp_flags(nullptr) {}
+  FinalImageRecipes() : _all_klasses(nullptr), _cp_recipes(nullptr), _flags(nullptr) {}
 
   void* operator new(size_t size) throw();
 

src/hotspot/share/classfile/classFileParser.cpp

@@ -5163,46 +5163,6 @@ void ClassFileParser::fill_instance_klass(InstanceKlass* ik,
   if (_parsed_annotations->has_any_annotations())
     _parsed_annotations->apply_to(ik);
 
-  // AOT-related checks.
-  // Note we cannot check this in general due to instrumentation or module patching
-  if (CDSConfig::is_initing_classes_at_dump_time()) {
-    // Check the aot initialization safe status.
-    // @AOTSafeClassInitializer is used only to support ahead-of-time initialization of classes
-    // in the AOT assembly phase.
-    if (ik->has_aot_safe_initializer()) {
-      // If a type is included in the tables inside can_archive_initialized_mirror(), we require that
-      //   - all super classes must be included
-      //   - all super interfaces that have <clinit> must be included.
-      // This ensures that in the production run, we don't run the <clinit> of a supertype but skips
-      // ik's <clinit>.
-      if (_super_klass != nullptr) {
-        guarantee_property(_super_klass->has_aot_safe_initializer(),
-                           "Missing @AOTSafeClassInitializer in superclass %s for class %s",
-                           _super_klass->external_name(),
-                           CHECK);
-      }
-
-      int len = _local_interfaces->length();
-      for (int i = 0; i < len; i++) {
-        InstanceKlass* intf = _local_interfaces->at(i);
-        guarantee_property(intf->class_initializer() == nullptr || intf->has_aot_safe_initializer(),
-                           "Missing @AOTSafeClassInitializer in superinterface %s for class %s",
-                           intf->external_name(),
-                           CHECK);
-      }
-
-      if (log_is_enabled(Info, aot, init)) {
-        ResourceMark rm;
-        log_info(aot, init)("Found @AOTSafeClassInitializer class %s", ik->external_name());
-      }
-    } else {
-      // @AOTRuntimeSetup only meaningful in @AOTClassInitializer
-      guarantee_property(!ik->is_runtime_setup_required(),
-                         "@AOTRuntimeSetup meaningless in non-@AOTSafeClassInitializer class %s",
-                         CHECK);
-    }
-  }
-
   apply_parsed_class_attributes(ik);
 
   // Miranda methods

src/java.base/share/classes/jdk/internal/math/MathUtils.java

@@ -26,12 +26,14 @@
 package jdk.internal.math;
 
 import jdk.internal.vm.annotation.Stable;
+import jdk.internal.vm.annotation.AOTSafeClassInitializer;
 
 /**
  * This class exposes package private utilities for other classes.
  * Thus, all methods are assumed to be invoked with correct arguments,
  * so these are not checked at all.
  */
+@AOTSafeClassInitializer
 final class MathUtils {
     /*
      * For full details about this code see the following reference: