Custom loader support

src/hotspot/share/cds/aotArtifactFinder.cpp

@@ -253,7 +253,7 @@ void AOTArtifactFinder::add_cached_instance_class(InstanceKlass* ik) {
       add_cached_instance_class(nest_host);
     }
 
-    if (CDSConfig::is_dumping_final_static_archive() && ik->defined_by_other_loaders()) {
+    if (CDSConfig::is_dumping_final_static_archive() && ik->defined_by_other_loaders() && !ik->is_defined_by_aot_safe_custom_loader() ) {
       // The following are not appliable to unregistered classes
       return;
     }

src/hotspot/share/cds/aotClassLinker.cpp

@@ -45,6 +45,12 @@ AOTClassLinker::ClassesTable* AOTClassLinker::_vm_classes = nullptr;
 AOTClassLinker::ClassesTable* AOTClassLinker::_candidates = nullptr;
 GrowableArrayCHeap<InstanceKlass*, mtClassShared>* AOTClassLinker::_sorted_candidates = nullptr;
 
+static const unsigned INITIAL_TABLE_SIZE = 997; // prime number
+static const unsigned MAX_TABLE_SIZE     = 10000;
+typedef GrowableArrayCHeap<InstanceKlass*, mtClassShared> ClassList;
+typedef ResizeableHashTable<Symbol*, ClassList*, AnyObj::C_HEAP, mtClass> ClassLoaderIdToPrelinkedTable;
+ClassLoaderIdToPrelinkedTable* _custom_loader_prelinked_table;
+
 #ifdef ASSERT
 bool AOTClassLinker::is_initialized() {
   assert(CDSConfig::is_dumping_archive(), "AOTClassLinker is for CDS dumping only");
@@ -59,6 +65,8 @@ void AOTClassLinker::initialize() {
   _candidates = new (mtClass)ClassesTable();
   _sorted_candidates = new GrowableArrayCHeap<InstanceKlass*, mtClassShared>(1000);
 
+  _custom_loader_prelinked_table = new (mtClass) ClassLoaderIdToPrelinkedTable(INITIAL_TABLE_SIZE, MAX_TABLE_SIZE);
+
   for (auto id : EnumRange<vmClassID>{}) {
     add_vm_class(vmClasses::klass_at(id));
   }
@@ -113,6 +121,19 @@ void AOTClassLinker::add_new_candidate(InstanceKlass* ik) {
   _candidates->put_when_absent(ik, true);
   _sorted_candidates->append(ik);
 
+  Symbol* loader_id;
+  loader_id = ik->cl_aot_identity();
+  if (loader_id != nullptr) {
+    ClassList** class_list_ptr = _custom_loader_prelinked_table->get(loader_id);
+    ClassList* class_list = nullptr;
+    if (class_list_ptr != nullptr) {
+      class_list = *class_list_ptr;
+    } else {
+      class_list = new ClassList(1000);
+      _custom_loader_prelinked_table->put(loader_id, class_list);
+    }
+    class_list->append(ik);
+  }
   if (log_is_enabled(Info, aot, link)) {
     ResourceMark rm;
     log_info(aot, link)("%s %s %p", class_category_name(ik), ik->external_name(), ik);
@@ -127,7 +148,7 @@ bool AOTClassLinker::try_add_candidate(InstanceKlass* ik) {
   assert(is_initialized(), "sanity");
   assert(CDSConfig::is_dumping_aot_linked_classes(), "sanity");
 
-  if (!SystemDictionaryShared::is_builtin(ik)) {
+  if (!SystemDictionaryShared::is_builtin(ik) && !ik->is_defined_by_aot_safe_custom_loader()) {
     // not loaded by a class loader which we know about
     return false;
   }
@@ -190,6 +211,66 @@ void AOTClassLinker::add_candidates() {
   }
 }
 
+static inline bool prelinked_table_equals(AOTLinkedClassTableForCustomLoader* table, Symbol* loader_id, int len_unused) {
+  return table->loader_id()->equals(loader_id);
+}
+
+class ArchivedCustomLoaderPrelinkedTable : public OffsetCompactHashtable<Symbol*, AOTLinkedClassTableForCustomLoader*,
+                                                                         prelinked_table_equals> {};
+ArchivedCustomLoaderPrelinkedTable _archived_custom_loader_prelinked_table;
+
+AOTLinkedClassTableForCustomLoader* AOTClassLinker::get_prelinked_table(Symbol* aot_id) {
+  unsigned int hash = Symbol::symbol_hash(aot_id);
+  return _archived_custom_loader_prelinked_table.lookup(aot_id, hash, 0 /* ignored */);
+}
+
+void AOTClassLinker::all_symbols_do(MetaspaceClosure* it) {
+  _custom_loader_prelinked_table->iterate_all([&](Symbol* loader_id, ClassList* class_list) {
+    it->push(&loader_id);
+  });
+}
+
+void AOTClassLinker::serialize_prelinked_table_header(SerializeClosure* soc) {
+  _archived_custom_loader_prelinked_table.serialize_header(soc);
+}
+
+void AOTClassLinker::print_archived_custom_loader_prelinked_table() {
+  if (log_is_enabled(Info, aot, link)) {
+    ResourceMark rm;
+    _archived_custom_loader_prelinked_table.iterate([&](AOTLinkedClassTableForCustomLoader* table) {
+      Array<InstanceKlass*>* class_list = table->class_list();
+      log_info(aot, link)("Class loader \"%s\" has %d classes in prelinked table", table->loader_id()->as_C_string(), class_list->length());
+      for (int i = 0; i < class_list->length(); i++) {
+        InstanceKlass* ik = class_list->at(i);
+        log_info(aot, link)("  %s", ik->external_name());
+      }
+    });
+  }
+}
+
+class CopyPrelinkTableToArchive : StackObj {
+private:
+  CompactHashtableWriter* _writer;
+  ArchiveBuilder* _builder;
+public:
+  CopyPrelinkTableToArchive(CompactHashtableWriter* writer) : _writer(writer),
+                                                              _builder(ArchiveBuilder::current())
+  {}
+
+  bool do_entry(Symbol* loader_id, ClassList* class_list) {
+    AOTLinkedClassTableForCustomLoader* tableForLoader = (AOTLinkedClassTableForCustomLoader*)ArchiveBuilder::ro_region_alloc(sizeof(AOTLinkedClassTableForCustomLoader));
+    assert(_builder->has_been_archived(loader_id), "must be");
+    Symbol* buffered_sym = _builder->get_buffered_addr(loader_id);
+    tableForLoader->init(buffered_sym, ArchiveUtils::archive_array(class_list));
+    ArchivePtrMarker::mark_pointer(tableForLoader->loader_id_addr());
+    ArchivePtrMarker::mark_pointer(tableForLoader->class_list_addr());
+    unsigned int hash = Symbol::symbol_hash(loader_id);
+    u4 delta = _builder->buffer_to_offset_u4((address)tableForLoader);
+    _writer->add(hash, delta);
+    return true;
+  }
+};
+
 void AOTClassLinker::write_to_archive() {
   assert(is_initialized(), "sanity");
   assert_at_safepoint();
@@ -200,6 +281,24 @@ void AOTClassLinker::write_to_archive() {
     table->set_boot2(write_classes(nullptr, false));
     table->set_platform(write_classes(SystemDictionary::java_platform_loader(), false));
     table->set_app(write_classes(SystemDictionary::java_system_loader(), false));
+
+    CompactHashtableStats stats;
+    CompactHashtableWriter writer(_custom_loader_prelinked_table->number_of_entries(), &stats);
+    CopyPrelinkTableToArchive archiver(&writer);
+    _custom_loader_prelinked_table->iterate(&archiver);
+    writer.dump(&_archived_custom_loader_prelinked_table, "archived prelinked table");
+
+    if (log_is_enabled(Info, aot, link)) {
+      ResourceMark rm;
+      _custom_loader_prelinked_table->iterate_all([&](Symbol* loader_id, ClassList* class_list) {
+        log_info(aot, link)("Class loader \"%s\" has %d classes in prelinked table", loader_id->as_C_string(), class_list->length());
+        for (int i = 0; i < class_list->length(); i++) {
+          InstanceKlass* ik = class_list->at(i);
+          log_info(aot, link)("  %s", ik->external_name());
+        }
+        return true;
+      });
+    }
   }
 }
 
@@ -228,6 +327,8 @@ Array<InstanceKlass*>* AOTClassLinker::write_classes(oop class_loader, bool is_j
   }
 }
 
+
+
 int AOTClassLinker::num_platform_initiated_classes() {
   if (CDSConfig::is_dumping_aot_linked_classes()) {
     // AOTLinkedClassBulkLoader will initiate loading of all public boot classes in the platform loader.
@@ -281,6 +382,8 @@ const char* AOTClassLinker::class_category_name(Klass* k) {
         return "plat";
       } else if (loader == SystemDictionary::java_system_loader()) {
         return "app";
+      } else if (k->cl_aot_identity() != nullptr) {
+        return "aotsafe_custom_lodaer";
       } else {
         return "unreg";
       }

src/hotspot/share/cds/aotClassLinker.hpp

@@ -40,6 +40,25 @@ class SerializeClosure;
 template <typename T> class Array;
 enum class AOTLinkedClassCategory : int;
 
+class AOTLinkedClassTableForCustomLoader {
+private:
+  Symbol* _loader_id;
+  Array<InstanceKlass*>* _prelinked_classes;
+  bool _loaded;
+public:
+  void init(Symbol* aot_id, Array<InstanceKlass*>* class_list) {
+    _loader_id = aot_id;
+    _prelinked_classes = class_list;
+    _loaded = false;
+  }
+  Symbol* loader_id() const { return _loader_id; }
+  address* loader_id_addr() const { return (address*)&_loader_id; }
+  Array<InstanceKlass*>* class_list() const { return _prelinked_classes; }
+  address* class_list_addr() const { return (address*)&_prelinked_classes; }
+  void set_loaded(bool value) { _loaded = value; }
+  bool is_loaded() const { return _loaded; }
+};
+
 // AOTClassLinker is used during the AOTCache Assembly Phase.
 // It links eligible classes before they are written into the AOTCache
 //
@@ -111,6 +130,11 @@ class AOTClassLinker :  AllStatic {
   static int num_app_initiated_classes();
   static int num_platform_initiated_classes();
 
+  static void all_symbols_do(MetaspaceClosure* it);
+  static void serialize_prelinked_table_header(SerializeClosure* soc);
+  static void print_archived_custom_loader_prelinked_table();
+  static AOTLinkedClassTableForCustomLoader* get_prelinked_table(Symbol* aot_id);
+
   // Used in logging: "boot1", "boot2", "plat", "app" and "unreg";
   static const char* class_category_name(AOTLinkedClassCategory category);
   static const char* class_category_name(Klass* k);

src/hotspot/share/cds/aotClassLocation.cpp

@@ -26,13 +26,16 @@
 #include "cds/aotLogging.hpp"
 #include "cds/aotMetaspace.hpp"
 #include "cds/archiveBuilder.hpp"
+#include "cds/archiveUtils.inline.hpp"
 #include "cds/cdsConfig.hpp"
 #include "cds/dynamicArchive.hpp"
 #include "cds/filemap.hpp"
 #include "cds/serializeClosure.hpp"
 #include "classfile/classLoader.hpp"
 #include "classfile/classLoaderData.hpp"
 #include "classfile/javaClasses.hpp"
+#include "classfile/compactHashtable.hpp"
+#include "classfile/symbolTable.hpp"
 #include "logging/log.hpp"
 #include "logging/logStream.hpp"
 #include "memory/metadataFactory.hpp"
@@ -43,6 +46,7 @@
 #include "runtime/arguments.hpp"
 #include "utilities/classpathStream.hpp"
 #include "utilities/formatBuffer.hpp"
+#include "utilities/resizableHashTable.hpp"
 #include "utilities/stringUtils.hpp"
 
 #include <errno.h>
@@ -156,6 +160,13 @@ class AllClassLocationStreams {
   }
 };
 
+class URLClassLoaderClassLocationStream : public ClassLocationStream {
+public:
+  URLClassLoaderClassLocationStream(const char* classpath) : ClassLocationStream() {
+    add_paths_in_classpath(classpath);
+  }
+};
+
 static bool has_jar_suffix(const char* filename) {
   // In jdk.internal.module.ModulePath.readModule(), it checks for the ".jar" suffix.
   // Performing the same check here.
@@ -1087,3 +1098,119 @@ void AOTClassLocationConfig::print_on(outputStream* st) const {
     st->print_cr("(%-6s) [%d] = %s", type, i, path);
   }
 }
+
+typedef ResizeableHashTable<Symbol*, GrowableClassLocationArray*, AnyObj::C_HEAP, mtClass> AOTIdToURLLoaderClasspath;
+static AOTIdToURLLoaderClasspath* _aot_id_to_classpath = nullptr;
+
+static inline bool ucc_symbol_equals(URLClassLoaderClasspath* ucc, Symbol* loader_id, int len_unused) {
+  return ucc->loader_id()->equals(loader_id);
+}
+
+class ArchivedAOTIdToClasspathMap : public OffsetCompactHashtable<Symbol*, URLClassLoaderClasspath*,
+                                                                         ucc_symbol_equals> {};
+static ArchivedAOTIdToClasspathMap _archived_aot_id_to_classpath;
+
+void URLClassLoaderClasspathSupport::init() {
+  _aot_id_to_classpath = new (mtClass) AOTIdToURLLoaderClasspath(11, 1000);
+  if (CDSConfig::is_dumping_final_static_archive()) {
+    reload_runtime_map();
+  }
+}
+
+void URLClassLoaderClasspathSupport::reload_runtime_map() {
+  _archived_aot_id_to_classpath.iterate([&](URLClassLoaderClasspath* ucc) {
+    GrowableClassLocationArray* locations = new GrowableClassLocationArray(ucc->num_entries());
+    for (int i = 0; i < ucc->num_entries(); i++) {
+      locations->append(ucc->class_location_at(i));
+      log_info(class, path)("path [%d] = %s", i, ucc->class_location_at(i)->path());
+    }
+    _aot_id_to_classpath->put(ucc->loader_id(), locations);
+  });
+}
+
+void URLClassLoaderClasspathSupport::add_urlclassloader_classpath(ClassLoaderData* loader_data, const char* classpath) {
+  assert(_aot_id_to_classpath != nullptr, "sanity check");
+  Symbol* aot_id = loader_data->aot_identity();
+  if (_aot_id_to_classpath->contains(aot_id)) {
+    // cannot allow aot_id clash; return without doing anything
+    return;
+  }
+  GrowableClassLocationArray* locations = new GrowableClassLocationArray(10);
+  URLClassLoaderClassLocationStream css(classpath);
+  for (css.start(); css.has_next(); ) {
+    const char* path = css.get_next();
+    AOTClassLocation* cs = AOTClassLocation::allocate(JavaThread::current(), path, locations->length(), AOTClassLocation::Group::URLCLASSLOADER_CLASSPATH, false);
+    log_info(class, path)("path [%d] = %s", locations->length(), path);
+    locations->append(cs);
+  }
+  _aot_id_to_classpath->put(aot_id, locations);
+}
+
+class URLClassLoaderClasspathArchiver : StackObj {
+private:
+  CompactHashtableWriter* _writer;
+  ArchiveBuilder* _builder;
+public:
+  URLClassLoaderClasspathArchiver(CompactHashtableWriter* writer) : _writer(writer),
+                                                              _builder(ArchiveBuilder::current())
+  {}
+
+  bool do_entry(Symbol* loader_id, GrowableClassLocationArray* class_locations) {
+    Array<AOTClassLocation*>* archived_copy = ArchiveBuilder::new_ro_array<AOTClassLocation*>(class_locations->length());
+    for (int i = 0; i < class_locations->length(); i++) {
+      archived_copy->at_put(i, class_locations->at(i)->write_to_archive());
+      ArchivePtrMarker::mark_pointer((address*)archived_copy->adr_at(i));
+    }
+    URLClassLoaderClasspath* ucc = (URLClassLoaderClasspath*)ArchiveBuilder::ro_region_alloc(sizeof(URLClassLoaderClasspath));
+    assert(_builder->has_been_archived(loader_id), "must be");
+    Symbol* buffered_sym = _builder->get_buffered_addr(loader_id);
+    ucc->init(buffered_sym, archived_copy);
+    ArchivePtrMarker::mark_pointer(ucc->loader_id_addr());
+    ArchivePtrMarker::mark_pointer(ucc->class_locations_addr());
+    unsigned int hash = Symbol::symbol_hash(loader_id);
+    u4 delta = _builder->buffer_to_offset_u4((address)ucc);
+    _writer->add(hash, delta);
+    return true;
+  }
+};
+
+void URLClassLoaderClasspathSupport::archive_classpath_map() {
+  CompactHashtableStats stats;
+  CompactHashtableWriter writer(_aot_id_to_classpath->number_of_entries(), &stats);
+  URLClassLoaderClasspathArchiver archiver(&writer);
+  _aot_id_to_classpath->iterate(&archiver);
+  writer.dump(&_archived_aot_id_to_classpath, "archived prelinked table");
+}
+
+void URLClassLoaderClasspathSupport::serialize_classpath_map_table_header(SerializeClosure* soc) {
+  _archived_aot_id_to_classpath.serialize_header(soc);
+}
+
+bool URLClassLoaderClasspathSupport::verify_archived_classpath(ClassLoaderData* loader_data, const char* classpath) {
+  ResourceMark rm;
+  Symbol* aot_id = loader_data->aot_identity();
+  assert(aot_id != nullptr, "sanity check");
+  const char* aot_id_str = aot_id->as_C_string();
+  unsigned int hash = Symbol::symbol_hash(aot_id);
+  URLClassLoaderClasspath* archived_classpath = _archived_aot_id_to_classpath.lookup(aot_id, hash, /*len*/0); // len is ignored
+  if (archived_classpath == nullptr) {
+    aot_log_warning(aot)("URLClassLoader (id=%s) classpath validation failed (reason: no archived entry found)", aot_id_str);
+    return false;
+  }
+  URLClassLoaderClassLocationStream uccs(classpath);
+  uccs.start();
+  for (int i = 0; i < archived_classpath->num_entries(); i++) {
+    AOTClassLocation* location = archived_classpath->class_location_at(i);
+    const char* archived_path = location->path();
+    if (!uccs.has_next()) {
+      aot_log_warning(aot)("URLClassLoader (id=%s) classpath validation failed (reason: classpath has fewer elements than expected)", aot_id_str);
+      return false;
+    }
+    const char* runtime_path = uccs.get_next();
+    if (!location->check(runtime_path, true)) {
+      aot_log_warning(aot)("URLClassLoader (id=%s) classpath validation failed", aot_id_str);
+      return false;
+    }
+  }
+  return true;
+}