8354897: Support Soft/Weak Reference in AOT cache

src/hotspot/share/cds/aotArtifactFinder.cpp

@@ -25,6 +25,7 @@
 #include "cds/aotClassLinker.hpp"
 #include "cds/aotArtifactFinder.hpp"
 #include "cds/aotClassInitializer.hpp"
+#include "cds/aotReferenceObjSupport.hpp"
 #include "cds/dumpTimeClassInfo.inline.hpp"
 #include "cds/heapShared.hpp"
 #include "cds/lambdaProxyClassDictionary.hpp"
@@ -73,6 +74,7 @@ void AOTArtifactFinder::find_artifacts() {
   // Note, if a class is not excluded, it does NOT mean it will be automatically included
   // into the AOT cache -- that will be decided by the code below.
   SystemDictionaryShared::finish_exclusion_checks();
+  AOTReferenceObjSupport::init_keep_alive_objs_table();
 
   start_scanning_for_oops();
 

src/hotspot/share/cds/aotClassInitializer.cpp

@@ -338,7 +338,8 @@ bool AOTClassInitializer::can_archive_initialized_mirror(InstanceKlass* ik) {
 bool AOTClassInitializer::is_runtime_setup_required(InstanceKlass* ik) {
   return ik == vmClasses::Class_klass() ||
          ik == vmClasses::internal_Unsafe_klass() ||
-         ik == vmClasses::ConcurrentHashMap_klass();
+         ik == vmClasses::ConcurrentHashMap_klass() ||
+         ik == vmClasses::Reference_klass();
 }
 
 void AOTClassInitializer::call_runtime_setup(JavaThread* current, InstanceKlass* ik) {

src/hotspot/share/cds/aotReferenceObjSupport.cpp

@@ -0,0 +1,240 @@
+/*
+ * 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 "cds/aotReferenceObjSupport.hpp"
+#include "cds/heapShared.hpp"
+#include "classfile/javaClasses.hpp"
+#include "classfile/symbolTable.hpp"
+#include "classfile/systemDictionary.hpp"
+#include "classfile/vmSymbols.hpp"
+#include "logging/log.hpp"
+#include "memory/resourceArea.hpp"
+#include "memory/universe.hpp"
+#include "oops/oop.inline.hpp"
+#include "oops/oopHandle.inline.hpp"
+#include "runtime/fieldDescriptor.inline.hpp"
+#include "runtime/javaCalls.hpp"
+#include "utilities/resourceHash.hpp"
+
+// Handling of java.lang.ref.Reference objects in the AOT cache
+// ============================================================
+//
+// When AOTArtifactFinder finds an oop which is a instance of java.lang.ref.Reference:
+//
+// - We check if the oop is eligible to be stored in the AOT cache. If not, the AOT cache
+//   creation fails -- see AOTReferenceObjSupport::check_if_ref_obj()
+//
+// - Otherwise, we store the oop into the AOT cache, but we unconditionally reset its
+//   "next" and "discovered" fields to null. Otherwise, if AOTArtifactFinder follows these
+//   fields, it may found unrelated objects that we don't intend to cache.
+//
+// Eligibility
+// ===========
+//
+// [1] A reference that does not require special clean up (i.e., Reference::queue == ReferenceQueue.NULL_QUEUE)
+//     is eligible.
+//
+// [2] A reference that REQUIRE specials clean up (i.e., Reference::queue != ReferenceQueue.NULL_QUEUE)
+//     is eligible ONLY if its referent is not null.
+//
+// As of this version, the only oops in group [2] that can be found by AOTArtifactFinder are
+// the keys used by ReferencedKeyMap in the implementation of MethodType::internTable.
+// stabilize_cached_reference_objects() ensures that all keys found by AOTArtifactFinder are eligible.
+//
+// The purpose of the error check in check_if_ref_obj() is to guard against changes in the JDK core
+// libs that might introduce new types of oops in group [2] into the AOT cache.
+//
+// Reasons for the eligibility restrictions
+// ========================================
+//
+// Reference handling is complex. In this version, we implement only enough functionality to support
+// the use of Weak/Soft references used by java.lang.invoke.
+//
+// We intend to evolve the implementation in the future by
+// -- implementing more assemblySetup() operations for other use cases, and/or
+// -- relaxing the eligibility restrictions.
+//
+//
+// null referents for group [1]
+// ============================
+//
+// Any cached reference R1 of group [1] is allowed to have a null referent.
+// This can happen in the following situations:
+//    (a) R1.clear() was called by Java code during the assembly phase.
+//    (b) The referent has been collected, and R1 is in the "pending" state.
+// In case (b), the "next" and "discovered" fields of the cached copy of R1 will
+// be set to null. During the production run:
+//    - It would appear to the Java program as if immediately during VM start-up, the referent
+//      was collected and ReferenceThread completed processing of R1.
+//    - It would appear to the GC as if immediately during VM start-up, the Java program called
+//      R1.clear().
+
+#if INCLUDE_CDS_JAVA_HEAP
+
+class KeepAliveObjectsTable : public ResourceHashtable<oop, bool,
+    36137, // prime number
+    AnyObj::C_HEAP,
+    mtClassShared,
+    HeapShared::oop_hash> {};
+
+static KeepAliveObjectsTable* _keep_alive_objs_table;
+static OopHandle _keep_alive_objs_array;
+static OopHandle _null_queue;
+
+bool AOTReferenceObjSupport::is_enabled() {
+  // For simplicity, AOTReferenceObjSupport is enabled only when dumping method handles.
+  // Otherwise we won't see Reference objects in the AOT cache. Let's be conservative now.
+  return CDSConfig::is_dumping_method_handles();
+}
+
+void AOTReferenceObjSupport::initialize(TRAPS) {
+  if (!AOTReferenceObjSupport::is_enabled()) {
+    return;
+  }
+
+  TempNewSymbol class_name = SymbolTable::new_symbol("java/lang/ref/ReferenceQueue");
+  Klass* k = SystemDictionary::resolve_or_fail(class_name, true, CHECK);
+  InstanceKlass* ik = InstanceKlass::cast(k);
+  ik->initialize(CHECK);
+
+  TempNewSymbol field_name = SymbolTable::new_symbol("NULL_QUEUE");
+  fieldDescriptor fd;
+  bool found = ik->find_local_field(field_name, vmSymbols::referencequeue_signature(), &fd);
+  precond(found);
+  precond(fd.is_static());
+
+  _null_queue = OopHandle(Universe::vm_global(), ik->java_mirror()->obj_field(fd.offset()));
+}
+
+// Ensure that all group [2] references found by AOTArtifactFinder are eligible.
+void AOTReferenceObjSupport::stabilize_cached_reference_objects(TRAPS) {
+  if (AOTReferenceObjSupport::is_enabled()) {
+    // This assert means that the MethodType and MethodTypeForm tables won't be
+    // updated concurrently, so we can remove GC'ed entries ...
+    assert(CDSConfig::allow_only_single_java_thread(), "Required");
+
+    {
+      TempNewSymbol method_name = SymbolTable::new_symbol("assemblySetup");
+      JavaValue result(T_VOID);
+      JavaCalls::call_static(&result, vmClasses::MethodType_klass(),
+                           method_name,
+                           vmSymbols::void_method_signature(),
+                           CHECK);
+    }
+
+    {
+      Symbol* cds_name  = vmSymbols::jdk_internal_misc_CDS();
+      Klass* cds_klass = SystemDictionary::resolve_or_fail(cds_name, true /*throw error*/,  CHECK);
+      TempNewSymbol method_name = SymbolTable::new_symbol("getKeepAliveObjects");
+      TempNewSymbol method_sig = SymbolTable::new_symbol("()[Ljava/lang/Object;");
+      JavaValue result(T_OBJECT);
+      JavaCalls::call_static(&result, cds_klass, method_name, method_sig, CHECK);
+
+      _keep_alive_objs_array = OopHandle(Universe::vm_global(), result.get_oop());
+    }
+  }
+}
+
+void AOTReferenceObjSupport::init_keep_alive_objs_table() {
+  assert_at_safepoint(); // _keep_alive_objs_table uses raw oops
+  oop a = _keep_alive_objs_array.resolve();
+  if (a != nullptr) {
+    precond(a->is_objArray());
+    precond(AOTReferenceObjSupport::is_enabled());
+    objArrayOop array = objArrayOop(a);
+
+    _keep_alive_objs_table = new (mtClass)KeepAliveObjectsTable();
+    for (int i = 0; i < array->length(); i++) {
+      oop obj = array->obj_at(i);
+      _keep_alive_objs_table->put(obj, true); // The array may have duplicated entries but that's OK.
+    }
+  }
+}
+
+// Returns true IFF obj is an instance of java.lang.ref.Reference. If so, perform extra eligibility checks.
+bool AOTReferenceObjSupport::check_if_ref_obj(oop obj) {
+  // We have a single Java thread. This means java.lang.ref.Reference$ReferenceHandler thread
+  // is not running. Otherwise the checks for next/discovered may not work.
+  precond(CDSConfig::allow_only_single_java_thread());
+  assert_at_safepoint(); // _keep_alive_objs_table uses raw oops
+
+  if (obj->klass()->is_subclass_of(vmClasses::Reference_klass())) {
+    precond(AOTReferenceObjSupport::is_enabled());
+    precond(JavaClasses::is_supported_for_archiving(obj));
+    precond(_keep_alive_objs_table != nullptr);
+
+    // GC needs to know about this load, It will keep referent alive until the current safepoint ends.
+    oop referent = HeapAccess<ON_UNKNOWN_OOP_REF>::oop_load_at(obj, java_lang_ref_Reference::referent_offset());
+
+    oop queue = obj->obj_field(java_lang_ref_Reference::queue_offset());
+    oop next = java_lang_ref_Reference::next(obj);
+    oop discovered = java_lang_ref_Reference::discovered(obj);
+    bool needs_special_cleanup = (queue != _null_queue.resolve());
+
+    // If you see the errors below, you probably modified the implementation of java.lang.invoke.
+    // Please check the comments at the top of this file.
+    if (needs_special_cleanup && (referent == nullptr || !_keep_alive_objs_table->contains(referent))) {
+      ResourceMark rm;
+
+      log_error(cds, heap)("Cannot archive reference object " PTR_FORMAT " of class %s",
+                           p2i(obj), obj->klass()->external_name());
+      log_error(cds, heap)("referent = " PTR_FORMAT
+                           ", queue = " PTR_FORMAT
+                           ", next = " PTR_FORMAT
+                           ", discovered = " PTR_FORMAT,
+                           p2i(referent), p2i(queue), p2i(next), p2i(discovered));
+      log_error(cds, heap)("This object requires special clean up as its queue is not ReferenceQueue::N" "ULL ("
+                           PTR_FORMAT ")", p2i(_null_queue.resolve()));
+      log_error(cds, heap)("%s", (referent == nullptr) ?
+                           "referent cannot be null" : "referent is not registered with CDS.keepAlive()");
+      HeapShared::debug_trace();
+      MetaspaceShared::unrecoverable_writing_error();
+    }
+
+    if (log_is_enabled(Info, cds, ref)) {
+      ResourceMark rm;
+      log_info(cds, ref)("Reference obj:"
+                         " r=" PTR_FORMAT
+                         " q=" PTR_FORMAT
+                         " n=" PTR_FORMAT
+                         " d=" PTR_FORMAT
+                         " %s",
+                         p2i(referent),
+                         p2i(queue),
+                         p2i(next),
+                         p2i(discovered),
+                         obj->klass()->external_name());
+    }
+    return true;
+  } else {
+    return false;
+  }
+}
+
+bool AOTReferenceObjSupport::skip_field(int field_offset) {
+  return (field_offset == java_lang_ref_Reference::next_offset() ||
+          field_offset == java_lang_ref_Reference::discovered_offset());
+}
+
+#endif // INCLUDE_CDS_JAVA_HEAP

src/hotspot/share/cds/aotReferenceObjSupport.hpp

@@ -0,0 +1,45 @@
+/*
+ * 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.
+ *
+ */
+
+#ifndef SHARE_CDS_AOTREFERENCEOBJSUPPORT_HPP
+#define SHARE_CDS_AOTREFERENCEOBJSUPPORT_HPP
+
+#include "memory/allStatic.hpp"
+#include "oops/oopsHierarchy.hpp"
+#include "utilities/exceptions.hpp"
+
+// Support for ahead-of-time allocated instances of java.lang.ref.Reference
+
+class AOTReferenceObjSupport : AllStatic {
+
+public:
+  static void initialize(TRAPS);
+  static void stabilize_cached_reference_objects(TRAPS);
+  static void init_keep_alive_objs_table() NOT_CDS_JAVA_HEAP_RETURN;
+  static bool check_if_ref_obj(oop obj);
+  static bool skip_field(int field_offset);
+  static bool is_enabled();
+};
+
+#endif // SHARE_CDS_AOTREFERENCEOBJSUPPORT_HPP

src/hotspot/share/cds/archiveHeapWriter.cpp

@@ -22,6 +22,7 @@
  *
  */
 
+#include "cds/aotReferenceObjSupport.hpp"
 #include "cds/archiveHeapWriter.hpp"
 #include "cds/cdsConfig.hpp"
 #include "cds/filemap.hpp"
@@ -607,18 +608,27 @@ class ArchiveHeapWriter::EmbeddedOopRelocator: public BasicOopIterateClosure {
   oop _src_obj;
   address _buffered_obj;
   CHeapBitMap* _oopmap;
-
+  bool _is_java_lang_ref;
 public:
   EmbeddedOopRelocator(oop src_obj, address buffered_obj, CHeapBitMap* oopmap) :
-    _src_obj(src_obj), _buffered_obj(buffered_obj), _oopmap(oopmap) {}
+    _src_obj(src_obj), _buffered_obj(buffered_obj), _oopmap(oopmap)
+  {
+    _is_java_lang_ref = AOTReferenceObjSupport::check_if_ref_obj(src_obj);
+  }
 
   void do_oop(narrowOop *p) { EmbeddedOopRelocator::do_oop_work(p); }
   void do_oop(      oop *p) { EmbeddedOopRelocator::do_oop_work(p); }
 
 private:
   template <class T> void do_oop_work(T *p) {
-    size_t field_offset = pointer_delta(p, _src_obj, sizeof(char));
-    ArchiveHeapWriter::relocate_field_in_buffer<T>((T*)(_buffered_obj + field_offset), _oopmap);
+    int field_offset = pointer_delta_as_int((char*)p, cast_from_oop<char*>(_src_obj));
+    T* field_addr = (T*)(_buffered_obj + field_offset);
+    if (_is_java_lang_ref && AOTReferenceObjSupport::skip_field(field_offset)) {
+      // Do not copy these fields. Set them to null
+      *field_addr = (T)0x0;
+    } else {
+      ArchiveHeapWriter::relocate_field_in_buffer<T>(field_addr, _oopmap);
+    }
   }
 };
 

src/hotspot/share/cds/cdsConfig.cpp

@@ -536,9 +536,6 @@ bool CDSConfig::check_vm_args_consistency(bool patch_mod_javabase, bool mode_fla
     // run to another which resulting in non-determinstic CDS archives.
     // Disable UseStringDeduplication while dumping CDS archive.
     UseStringDeduplication = false;
-
-    // Don't use SoftReferences so that objects used by java.lang.invoke tables can be archived.
-    Arguments::PropertyList_add(new SystemProperty("java.lang.invoke.MethodHandleNatives.USE_SOFT_CACHE", "false", false));
   }
 
   // RecordDynamicDumpInfo is not compatible with ArchiveClassesAtExit