From 3cc21f911715cc0112f9a2e80a199723b29262e1 Mon Sep 17 00:00:00 2001
Date: Fri, 29 Nov 2024 15:39:13 +0800
Subject: Optimize LazyAOT klasses sending strategy

---
 .../jbooster/client/clientMessageHandler.cpp  |  6 ++
 .../jbooster/client/clientMessageHandler.hpp  |  1 +
 .../share/jbooster/jClientArguments.cpp       | 15 +++
 .../share/jbooster/jClientArguments.hpp       |  1 +
 .../share/jbooster/jbooster_globals.hpp       |  4 +
 src/hotspot/share/jbooster/lazyAot.cpp        | 15 ++-
 src/hotspot/share/jbooster/lazyAot.hpp        |  2 +
 .../jbooster/net/communicationStream.cpp      | 15 ++-
 .../net/communicationStream.inline.hpp        |  1 +
 src/hotspot/share/jbooster/net/errorCode.hpp  |  1 +
 src/hotspot/share/jbooster/net/message.hpp    |  5 +
 .../share/jbooster/net/message.inline.hpp     |  2 +
 .../share/jbooster/net/messageType.hpp        |  1 +
 src/hotspot/share/jbooster/net/netCommon.hpp  |  8 +-
 .../jbooster/server/serverMessageHandler.cpp  | 10 +-
 .../jbooster/server/serverMessageHandler.hpp  |  3 +-
 .../share/jbooster/utilities/fileUtils.cpp    | 64 ++-----------
 .../share/jbooster/utilities/fileUtils.hpp    |  3 +-
 .../jbooster/JBoosterCompilationContext.java  |  5 +
 .../graalvm/compiler/java/BytecodeParser.java | 29 ++++++
 .../share/classes/jdk/jbooster/JBooster.java  |  8 +-
 .../JBoosterCompilationContextImpl.java       | 10 +-
 test/hotspot/gtest/jbooster/common.hpp        | 96 +++++++++++++++++++
 test/hotspot/gtest/jbooster/test_net.cpp      |  4 +-
 test/hotspot/gtest/jbooster/test_util.cpp     | 12 ++-
 test/jdk/tools/jbooster/JBoosterTestBase.java |  6 +-
 26 files changed, 244 insertions(+), 83 deletions(-)
 create mode 100644 test/hotspot/gtest/jbooster/common.hpp

diff --git a/src/hotspot/share/jbooster/client/clientMessageHandler.cpp b/src/hotspot/share/jbooster/client/clientMessageHandler.cpp
index 9119b897a..b8278d41c 100644
--- a/src/hotspot/share/jbooster/client/clientMessageHandler.cpp
+++ b/src/hotspot/share/jbooster/client/clientMessageHandler.cpp
@@ -289,6 +289,12 @@ int ClientMessageHandler::handle_ArrayKlasses() {
   return 0;
 }
 
+int ClientMessageHandler::handle_ResolveExtraKlasses() {
+  ResourceMark rm;
+  bool should_resolve_extra_Klasses = JBoosterResolveExtraKlasses;
+  JB_RETURN(cs().send_response(&should_resolve_extra_Klasses));
+  return 0;
+}
 // ---------------------------------- Some Tasks -----------------------------------
 
 int ClientMessageHandler::send_cache_file_sync_task() {
diff --git a/src/hotspot/share/jbooster/client/clientMessageHandler.hpp b/src/hotspot/share/jbooster/client/clientMessageHandler.hpp
index 66844247c..94ff7387d 100644
--- a/src/hotspot/share/jbooster/client/clientMessageHandler.hpp
+++ b/src/hotspot/share/jbooster/client/clientMessageHandler.hpp
@@ -38,6 +38,7 @@
   F(MethodLocators          )             \
   F(Profilinginfo           )             \
   F(ArrayKlasses            )             \
+  F(ResolveExtraKlasses   )               \
 
 class ArrayKlass;
 class ClassLoaderData;
diff --git a/src/hotspot/share/jbooster/jClientArguments.cpp b/src/hotspot/share/jbooster/jClientArguments.cpp
index b215913ea..42aba4181 100644
--- a/src/hotspot/share/jbooster/jClientArguments.cpp
+++ b/src/hotspot/share/jbooster/jClientArguments.cpp
@@ -171,6 +171,14 @@ const char* calc_java_commands_by_class(const char* full_cmd, int full_cmd_len)
   return StringUtils::copy_to_heap(start, mtJBooster);
 }
 
+static int64_t calc_cds_file_size() {
+  char* default_archive_path = Arguments::get_default_shared_archive_path();
+  if (default_archive_path == nullptr) { return -1; }
+  int64_t cds_file_size = FileUtils::file_size(default_archive_path);
+  StringUtils::free_from_heap(default_archive_path);
+  return cds_file_size;
+}
+
 JClientArguments::JClientArguments(bool is_client) {
   if (is_client) {
     init_for_client();
@@ -224,6 +232,7 @@ void JClientArguments::init_for_client() {
   _classpath_name_hash = calc_classpath_name_hash(app_cp, app_cp_len);
   _classpath_timestamp_hash = calc_classpath_timestamp_hash(app_cp, app_cp_len);
   _agent_name_hash = calc_agent_name_hash();
+  _cds_file_size = calc_cds_file_size();
   if (JBoosterClientStrictMatch) {
     _java_commands = is_jar ? calc_java_commands_by_jar(full_cmd, app_cp_len)
                             : calc_java_commands_by_class(full_cmd, full_cmd_len);
@@ -248,6 +257,7 @@ uint32_t JClientArguments::calc_hash() {
   result = calc_new_hash(result, _classpath_name_hash);
   result = calc_new_hash(result, _classpath_timestamp_hash);
   result = calc_new_hash(result, _agent_name_hash);
+  result = calc_new_hash(result, _cds_file_size);
   result = calc_new_hash(result, StringUtils::hash_code(_java_commands));
   result = calc_new_hash(result, _related_flags->hash());
 
@@ -267,6 +277,7 @@ bool JClientArguments::equals(const JClientArguments* that) const {
   if (this->_classpath_name_hash      != that->_classpath_name_hash) return false;
   if (this->_classpath_timestamp_hash != that->_classpath_timestamp_hash) return false;
   if (this->_agent_name_hash          != that->_agent_name_hash) return false;
+  if (this->_cds_file_size            != that->_cds_file_size) return false;
   if (StringUtils::compare(this->_java_commands, that->_java_commands) != 0) return false;
   if (!this->_related_flags->equals(that->_related_flags)) return false;
   return true;
@@ -284,6 +295,7 @@ void JClientArguments::print_args(outputStream* st) const {
   st->print_cr("  classpath_name_hash:      %x",      _classpath_name_hash);
   st->print_cr("  classpath_timestamp_hash: %x",      _classpath_timestamp_hash);
   st->print_cr("  agent_name_hash:          %x",      _agent_name_hash);
+  st->print_cr("  cds_file_size:            %lu",     _cds_file_size);
   st->print_cr("  java_commands:            \"%s\"",  _java_commands);
   st->print_cr("  vm_flags:");
   st->print_cr("    hash: %u", _related_flags->hash());
@@ -307,6 +319,7 @@ int JClientArguments::serialize(MessageBuffer& buf) const {
   JB_RETURN(buf.serialize_no_meta(_classpath_name_hash));
   JB_RETURN(buf.serialize_no_meta(_classpath_timestamp_hash));
   JB_RETURN(buf.serialize_no_meta(_agent_name_hash));
+  JB_RETURN(buf.serialize_no_meta(_cds_file_size));
   JB_RETURN(buf.serialize_with_meta(&_java_commands));
   JB_RETURN(buf.serialize_with_meta(_related_flags));
 
@@ -343,6 +356,8 @@ int JClientArguments::deserialize(MessageBuffer& buf) {
 
   JB_RETURN(buf.deserialize_ref_no_meta(_agent_name_hash));
 
+  JB_RETURN(buf.deserialize_ref_no_meta(_cds_file_size));
+
   StringWrapper sw_java_commands;
   JB_RETURN(buf.deserialize_with_meta(&sw_java_commands));
   _java_commands = sw_java_commands.export_string();
diff --git a/src/hotspot/share/jbooster/jClientArguments.hpp b/src/hotspot/share/jbooster/jClientArguments.hpp
index 09c5dfdc0..491ac3159 100644
--- a/src/hotspot/share/jbooster/jClientArguments.hpp
+++ b/src/hotspot/share/jbooster/jClientArguments.hpp
@@ -92,6 +92,7 @@ private:
   uint32_t    _classpath_name_hash;
   uint32_t    _classpath_timestamp_hash;
   uint32_t    _agent_name_hash;
+  int64_t     _cds_file_size;
   const char* _java_commands;
   JClientVMFlags* _related_flags;
   // ========================= end =========================
diff --git a/src/hotspot/share/jbooster/jbooster_globals.hpp b/src/hotspot/share/jbooster/jbooster_globals.hpp
index 74968af75..c3cdd71f8 100644
--- a/src/hotspot/share/jbooster/jbooster_globals.hpp
+++ b/src/hotspot/share/jbooster/jbooster_globals.hpp
@@ -119,6 +119,10 @@
                                                                             \
   product(ccstr, JBoosterServerSSLRootCerts, NULL, EXPERIMENTAL,            \
           "The file path to save server SSL root certificate.")             \
+                                                                            \
+  product(bool, JBoosterResolveExtraKlasses, true, EXPERIMENTAL,            \
+          "Whether resolve additional klasses "                             \
+          "while collecting klasses for AOT.")                              \
 
 
 // end of JBOOSTER_FLAGS
diff --git a/src/hotspot/share/jbooster/lazyAot.cpp b/src/hotspot/share/jbooster/lazyAot.cpp
index 410d575e9..462b0c72b 100644
--- a/src/hotspot/share/jbooster/lazyAot.cpp
+++ b/src/hotspot/share/jbooster/lazyAot.cpp
@@ -199,7 +199,11 @@ void LazyAOT::collect_klasses_in_constant_pool(GrowableArray<InstanceKlass*>* ds
     int len = dst->length();
     for (int i = last_len; i < len; ++i) {
       ThreadInVMfromNative tiv(THREAD);
-      collect_klasses_in_constant_pool(dst, vis, dst->at(i), ALL_KLASSES, CHECK);
+      if (JBoosterResolveExtraKlasses) {
+        collect_klasses_in_constant_pool(dst, vis, dst->at(i), ALL_KLASSES, CHECK);
+      } else {
+        collect_klasses_in_constant_pool(dst, vis, dst->at(i), RESOLVED_KLASSES, CHECK);
+      }
     }
     last_len = len;
   }
@@ -527,6 +531,7 @@ bool LazyAOT::compile_classes_by_graal(int session_id,
                                        const char* file_path,
                                        GrowableArray<InstanceKlass*>* klasses,
                                        bool use_pgo,
+                                       bool resolve_no_extra_klasses,
                                        TRAPS) {
   DebugUtils::assert_thread_in_vm();
 
@@ -541,9 +546,11 @@ bool LazyAOT::compile_classes_by_graal(int session_id,
   java_args.push_int(session_id);
   java_args.push_oop(file_path_h);
   java_args.push_oop(hash_set_h);
+  java_args.push_int(use_pgo);
+  java_args.push_int(resolve_no_extra_klasses);
 
   TempNewSymbol compile_classes_name = SymbolTable::new_symbol("compileClasses");
-  TempNewSymbol compile_classes_signature = SymbolTable::new_symbol("(ILjava/lang/String;Ljava/util/Set;)Z");
+  TempNewSymbol compile_classes_signature = SymbolTable::new_symbol("(ILjava/lang/String;Ljava/util/Set;ZZ)Z");
   JavaCalls::call_static(&result, ServerDataManager::get().main_klass(),
                          compile_classes_name,
                          compile_classes_signature,
@@ -557,6 +564,7 @@ bool LazyAOT::compile_methods_by_graal(int session_id,
                                        GrowableArray<Method*>* methods_to_compile,
                                        GrowableArray<Method*>* methods_not_compile,
                                        bool use_pgo,
+                                       bool resolve_no_extra_klasses,
                                        TRAPS) {
   DebugUtils::assert_thread_in_vm();
 
@@ -580,9 +588,10 @@ bool LazyAOT::compile_methods_by_graal(int session_id,
   java_args.push_oop(method_name_set_h);
   java_args.push_oop(not_method_name_set_h);
   java_args.push_int(use_pgo);
+  java_args.push_int(resolve_no_extra_klasses);
 
   TempNewSymbol compile_methods_name = SymbolTable::new_symbol("compileMethods");
-  TempNewSymbol compile_methods_signature = SymbolTable::new_symbol("(ILjava/lang/String;Ljava/util/Set;Ljava/util/Set;Ljava/util/Set;Z)Z");
+  TempNewSymbol compile_methods_signature = SymbolTable::new_symbol("(ILjava/lang/String;Ljava/util/Set;Ljava/util/Set;Ljava/util/Set;ZZ)Z");
   JavaCalls::call_static(&result, ServerDataManager::get().main_klass(),
                          compile_methods_name,
                          compile_methods_signature,
diff --git a/src/hotspot/share/jbooster/lazyAot.hpp b/src/hotspot/share/jbooster/lazyAot.hpp
index 1d2a49351..5b4183118 100644
--- a/src/hotspot/share/jbooster/lazyAot.hpp
+++ b/src/hotspot/share/jbooster/lazyAot.hpp
@@ -112,6 +112,7 @@ public:
                                        const char* file_path,
                                        GrowableArray<InstanceKlass*>* klasses,
                                        bool use_pgo,
+                                       bool resolve_extra_klasses,
                                        TRAPS);
   static bool compile_methods_by_graal(int session_id,
                                        const char* file_path,
@@ -119,6 +120,7 @@ public:
                                        GrowableArray<Method*>* methods_to_compile,
                                        GrowableArray<Method*>* methods_not_compile,
                                        bool use_pgo,
+                                       bool resolve_extra_klasses,
                                        TRAPS);
 };
 
diff --git a/src/hotspot/share/jbooster/net/communicationStream.cpp b/src/hotspot/share/jbooster/net/communicationStream.cpp
index cdb4b8fa7..37bb45032 100644
--- a/src/hotspot/share/jbooster/net/communicationStream.cpp
+++ b/src/hotspot/share/jbooster/net/communicationStream.cpp
@@ -145,16 +145,18 @@ int CommunicationStream::recv_message() {
   Message& msg = _msg_recv;
   // read once (or memmove from the overflowed buffer) to get message size
   uint32_t read_size, msg_size;
+  uint16_t magic_num;
+  const uint32_t min_read_size = sizeof(msg_size) + sizeof(magic_num);
   if (msg.has_overflow()) {
     read_size = msg.move_overflow();
-    if (read_size < sizeof(msg_size)) {
+    if (read_size < min_read_size) {
       read_size += read_once_from_stream(msg.buf_beginning() + read_size, msg.buf_size() - read_size);
     }
   } else {
     read_size = read_once_from_stream(msg.buf_beginning(), msg.buf_size());
   }
 
-  if (read_size < sizeof(msg_size)) {
+  if (read_size < min_read_size) {
     if (!is_stream_closed()) {
       log_warning(jbooster, rpc)("Failed to read the size of the message (read_size=%d). stream_id=%u.",
                                  read_size, stream_id());
@@ -162,7 +164,16 @@ int CommunicationStream::recv_message() {
     return return_errno_or_flag(JBErr::BAD_MSG_SIZE);
   }
 
+  magic_num = msg.deserialize_magic_num_only();
+  if (magic_num != MessageConst::RPC_MAGIC) {
+    log_warning(jbooster, rpc)("Message not from JBooster client. stream_id=%u.", stream_id());
+    return return_errno_or_flag(JBErr::BAD_MAGIC_NUM);
+  }
   msg_size = msg.deserialize_size_only();
+  if (msg_size > MessageConst::MAX_MSG_SIZE) {
+    log_warning(jbooster, rpc)("Message size should be no more than 2GB. stream_id=%u.", stream_id());
+    return return_errno_or_flag(JBErr::BAD_MSG_SIZE);
+  }
   if (read_size > msg_size) {         // read too much
     msg.set_overflow(msg_size, read_size - msg_size);
   } else if (read_size < msg_size) {  // read the rest then
diff --git a/src/hotspot/share/jbooster/net/communicationStream.inline.hpp b/src/hotspot/share/jbooster/net/communicationStream.inline.hpp
index 45311c2be..bc5c0565f 100644
--- a/src/hotspot/share/jbooster/net/communicationStream.inline.hpp
+++ b/src/hotspot/share/jbooster/net/communicationStream.inline.hpp
@@ -58,6 +58,7 @@ inline bool CommunicationStream::check_received_message_size() {
 template <typename... Args>
 inline int CommunicationStream::send_request(MessageType type, const Args* const... args) {
   _msg_send.set_msg_type(type);
+  _msg_send.set_magic_num(MessageConst::RPC_MAGIC);
   _msg_send.set_cur_buf_offset_after_meta();
   JB_RETURN(_msg_send.serialize(args...));
   _msg_send.set_msg_size_based_on_cur_buf_offset();
diff --git a/src/hotspot/share/jbooster/net/errorCode.hpp b/src/hotspot/share/jbooster/net/errorCode.hpp
index 625ef6951..cbb4cc54c 100644
--- a/src/hotspot/share/jbooster/net/errorCode.hpp
+++ b/src/hotspot/share/jbooster/net/errorCode.hpp
@@ -30,6 +30,7 @@
   f(CONN_CLOSED,              "Connection has been closed"                        ) \
   f(CONN_CLOSED_BY_PEER,      "Connection is closed by the other end"             ) \
   f(BAD_SSL,                  "Unexpected SSL error during initialization"        ) \
+  f(BAD_MAGIC_NUM,            "Unexpected magic number of the received message"   ) \
   f(BAD_MSG_SIZE,             "Unexpected size of the received message"           ) \
   f(BAD_MSG_TYPE,             "Unexpected message type of the received message"   ) \
   f(BAD_MSG_DATA,             "Unexpected payload data of the received message"   ) \
diff --git a/src/hotspot/share/jbooster/net/message.hpp b/src/hotspot/share/jbooster/net/message.hpp
index 47d2634e2..413ff84a5 100644
--- a/src/hotspot/share/jbooster/net/message.hpp
+++ b/src/hotspot/share/jbooster/net/message.hpp
@@ -33,6 +33,7 @@ class Message: public MessageConst {
 private:
   struct {
     uint32_t msg_size;
+    uint16_t magic_num;
     MessageType msg_type;
   } _meta;
 
@@ -58,6 +59,7 @@ public:
   uint32_t msg_size() const { return _meta.msg_size; }
   void set_msg_size(uint32_t size) { _meta.msg_size = size; }
   void set_msg_size_based_on_cur_buf_offset() { set_msg_size(cur_buf_offset()); }
+  void set_magic_num(uint16_t magic_num) { _meta.magic_num = magic_num; }
   MessageType msg_type() const { return _meta.msg_type; }
   void set_msg_type(MessageType type) { _meta.msg_type = type; }
 
@@ -76,6 +78,9 @@ public:
   }
 
   uint32_t deserialize_size_only() { return *((uint32_t*)_buf.buf()); }
+  uint16_t deserialize_magic_num_only() { return *((uint16_t*)(_buf.buf() + sizeof(_meta.msg_size))); }
+
+  bool check_magic_num(uint16_t magic_num) { return magic_num == MessageConst::RPC_MAGIC; }
 
   template <typename... Args>
   int serialize(const Args* const... args);
diff --git a/src/hotspot/share/jbooster/net/message.inline.hpp b/src/hotspot/share/jbooster/net/message.inline.hpp
index 5b5add47f..b85362efe 100644
--- a/src/hotspot/share/jbooster/net/message.inline.hpp
+++ b/src/hotspot/share/jbooster/net/message.inline.hpp
@@ -78,6 +78,7 @@ inline int Message::deserialize(Args* const&... args) {
 inline void Message::serialize_meta() {
   _buf.set_cur_offset(0);
   _buf.serialize_no_meta(_meta.msg_size);
+  _buf.serialize_no_meta(_meta.magic_num);
   _buf.serialize_no_meta(_meta.msg_type);
   assert(cur_buf_offset() == meta_size, "sanity");
 }
@@ -85,6 +86,7 @@ inline void Message::serialize_meta() {
 inline void Message::deserialize_meta() {
   _buf.set_cur_offset(0);
   _buf.deserialize_ref_no_meta(_meta.msg_size);
+  _buf.deserialize_ref_no_meta(_meta.magic_num);
   _buf.deserialize_ref_no_meta(_meta.msg_type);
   assert(cur_buf_offset() == meta_size, "sanity");
 }
diff --git a/src/hotspot/share/jbooster/net/messageType.hpp b/src/hotspot/share/jbooster/net/messageType.hpp
index f8cb8f3e6..dbd89c07c 100644
--- a/src/hotspot/share/jbooster/net/messageType.hpp
+++ b/src/hotspot/share/jbooster/net/messageType.hpp
@@ -54,6 +54,7 @@
   f(ArrayKlasses,                       "from server"   ) \
   f(DataOfKlasses,                      "from server"   ) \
   f(MethodLocators,                     "from server"   ) \
+  f(ResolveExtraKlasses,                "from server"   ) \
                                                           \
   /* others */                                            \
   f(FileSegment,                        "from both"     ) \
diff --git a/src/hotspot/share/jbooster/net/netCommon.hpp b/src/hotspot/share/jbooster/net/netCommon.hpp
index 8706b7f22..2e30708cc 100644
--- a/src/hotspot/share/jbooster/net/netCommon.hpp
+++ b/src/hotspot/share/jbooster/net/netCommon.hpp
@@ -133,13 +133,15 @@
 
 class MessageConst {
 public:
+  static constexpr uint32_t MAX_MSG_SIZE = 0x80000000;
+  static constexpr uint16_t RPC_MAGIC = 0xB05E;
   enum: uint32_t {
     /**
      * The layout of the message in the buffer:
-     * | msg_size | msg_type | ... (all of the arguments) ... |
-     * | 4 bytes  | 2 bytes  |     msg_size - 4 - 2 bytes     |
+     * | msg_size | magic_num | msg_type | ... (all of the arguments) ... |
+     * | 4 bytes  |  2 bytes  | 2 bytes  |     msg_size - 4 - 2 bytes     |
      */
-    meta_size = sizeof(uint32_t) + sizeof(MessageType),
+    meta_size = sizeof(uint32_t) + sizeof(uint16_t) + sizeof(MessageType),
     /**
      * The layout of each argument in the buffer:
      * | arg_size | ... (payload of the argument) ... |
diff --git a/src/hotspot/share/jbooster/server/serverMessageHandler.cpp b/src/hotspot/share/jbooster/server/serverMessageHandler.cpp
index f46b839bb..713ea1374 100644
--- a/src/hotspot/share/jbooster/server/serverMessageHandler.cpp
+++ b/src/hotspot/share/jbooster/server/serverMessageHandler.cpp
@@ -319,11 +319,15 @@ int ServerMessageHandler::handle_lazy_aot_compilation_task(TRAPS) {
   GrowableArray<Method*>        methods_to_compile;
   GrowableArray<Method*>        methods_not_compile;
   bool compile_in_current_thread = false;
+  bool resolve_extra_klasses = true;
 
   JB_TRY {
     compile_in_current_thread = !aot_cache_state.is_cached() && aot_cache_state.set_being_generated();
     JB_THROW(ss().send_response(&compile_in_current_thread));
     if (compile_in_current_thread) {
+      JB_RETURN(ss().send_request(MessageType::ResolveExtraKlasses));
+      JB_RETURN(ss().recv_response(&resolve_extra_klasses));
+
       JB_THROW(request_missing_class_loaders(THREAD));
       JB_THROW(request_missing_klasses(THREAD));
       JB_THROW(request_methods_to_compile(&klasses_to_compile, &methods_to_compile, THREAD));
@@ -345,6 +349,7 @@ int ServerMessageHandler::handle_lazy_aot_compilation_task(TRAPS) {
                                       &methods_to_compile,
                                       &methods_not_compile,
                                       enabling_aot_pgo,
+                                      resolve_extra_klasses,
                                       THREAD));
   } else {  // not compile in current thread
     if (aot_cache_state.is_being_generated()) {
@@ -365,6 +370,7 @@ int ServerMessageHandler::try_to_compile_lazy_aot(GrowableArray<InstanceKlass*>*
                                                   GrowableArray<Method*>* methods_to_compile,
                                                   GrowableArray<Method*>* methods_not_compile,
                                                   bool enabling_aot_pgo,
+                                                  bool resolve_extra_klasses,
                                                   TRAPS) {
   JClientProgramData* pd = ss().session_data()->program_data();
   JClientCacheState& aot_cache_state =  enabling_aot_pgo ? pd->aot_pgo_cache_state() : pd->aot_static_cache_state();
@@ -382,10 +388,10 @@ int ServerMessageHandler::try_to_compile_lazy_aot(GrowableArray<InstanceKlass*>*
 
   ThreadInVMfromNative tiv(THREAD);
   if (methods_to_compile->is_empty()) {
-    successful = LazyAOT::compile_classes_by_graal(session_id, file_path, klasses_to_compile, enabling_aot_pgo, THREAD);
+    successful = LazyAOT::compile_classes_by_graal(session_id, file_path, klasses_to_compile, enabling_aot_pgo, resolve_extra_klasses, THREAD);
   } else {
     successful = LazyAOT::compile_methods_by_graal(session_id, file_path, klasses_to_compile,
-                                                   methods_to_compile, methods_not_compile, enabling_aot_pgo, THREAD);
+                                                   methods_to_compile, methods_not_compile, enabling_aot_pgo, resolve_extra_klasses, THREAD);
   }
 
   if (successful) {
diff --git a/src/hotspot/share/jbooster/server/serverMessageHandler.hpp b/src/hotspot/share/jbooster/server/serverMessageHandler.hpp
index 26926dd41..c761b9c44 100644
--- a/src/hotspot/share/jbooster/server/serverMessageHandler.hpp
+++ b/src/hotspot/share/jbooster/server/serverMessageHandler.hpp
@@ -59,7 +59,8 @@ private:
   int try_to_compile_lazy_aot(GrowableArray<InstanceKlass*>* klasses_to_compile,
                               GrowableArray<Method*>* methods_to_compile,
                               GrowableArray<Method*>* methods_not_compile,
-                              bool use_pgo,
+                              bool enabling_aot_pgo,
+                              bool resolve_extra_klasses,
                               TRAPS);
 public:
   ServerMessageHandler(ServerStream* server_stream);
diff --git a/src/hotspot/share/jbooster/utilities/fileUtils.cpp b/src/hotspot/share/jbooster/utilities/fileUtils.cpp
index f19bf8fb3..598013523 100644
--- a/src/hotspot/share/jbooster/utilities/fileUtils.cpp
+++ b/src/hotspot/share/jbooster/utilities/fileUtils.cpp
@@ -75,6 +75,13 @@ bool FileUtils::is_dir(const char* path) {
   return S_ISDIR(st.st_mode);
 }
 
+int64_t FileUtils::file_size(const char* path) {
+  struct stat st = {0};
+  if (os::stat(path, &st) != 0) return -1;
+  // We don't care if it is a regular file.
+  return st.st_size;
+}
+
 uint64_t FileUtils::modify_time(const char* path) {
   struct stat st = {0};
   if (os::stat(path, &st) != 0) return 0;
@@ -119,60 +126,3 @@ bool FileUtils::move(const char* path_from, const char* path_to) {
 bool FileUtils::remove(const char* path) {
   return ::remove(path) == 0;
 }
-
-bool FileUtils::is_same(const char* path1, const char* path2) {
-  bool res = false;
-  char* buf1 = nullptr;
-  char* buf2 = nullptr;
-  int fd1 = os::open(path1, O_BINARY | O_RDONLY, 0);
-  int fd2 = os::open(path2, O_BINARY | O_RDONLY, 0);
-  do {
-    if (fd1 < 0 || fd2 < 0) break;
-    int64_t size1 = os::lseek(fd1, 0, SEEK_END);
-    int64_t size2 = os::lseek(fd2, 0, SEEK_END);
-    if (size1 != size2) break;
-    int64_t size = size1;
-    os::lseek(fd1, 0, SEEK_SET);
-    os::lseek(fd2, 0, SEEK_SET);
-    // We don't use NEW_RESOURCE_ARRAY here as Thread::current() may
-    // not be initialized yet.
-    buf1 = NEW_C_HEAP_ARRAY(char, (size_t) size, mtJBooster);
-    buf2 = NEW_C_HEAP_ARRAY(char, (size_t) size, mtJBooster);
-    size1 = (int64_t) os::read(fd1, buf1, size);
-    size2 = (int64_t) os::read(fd2, buf2, size);
-    guarantee(size1 == size && size2 == size, "sanity");
-    res = memcmp(buf1, buf2, size) == 0;
-  } while (false);
-  if (fd1 >= 0) os::close(fd1);
-  if (fd2 >= 0) os::close(fd2);
-  if (buf1 != nullptr) {
-    FREE_C_HEAP_ARRAY(char, buf1);
-  }
-  if (buf2 != nullptr) {
-    FREE_C_HEAP_ARRAY(char, buf2);
-  }
-  return res;
-}
-
-bool FileUtils::is_same(const char* path, const char* content, int64_t size) {
-  bool res = false;
-  char* buf = nullptr;
-  int fd = os::open(path, O_BINARY | O_RDONLY, 0);
-  do {
-    if (fd < 0) break;
-    int64_t fsize = os::lseek(fd, 0, SEEK_END);
-    if (fsize != size) break;
-    os::lseek(fd, 0, SEEK_SET);
-    // We don't use NEW_RESOURCE_ARRAY here as Thread::current() may
-    // not be initialized yet.
-    buf = NEW_C_HEAP_ARRAY(char, (size_t) size, mtJBooster);
-    fsize = (int64_t) os::read(fd, buf, size);
-    guarantee(fsize == size, "sanity");
-    res = memcmp(content, buf, size) == 0;
-  } while (false);
-  if (fd >= 0) os::close(fd);
-  if (buf != nullptr) {
-    FREE_C_HEAP_ARRAY(char, buf);
-  }
-  return res;
-}
diff --git a/src/hotspot/share/jbooster/utilities/fileUtils.hpp b/src/hotspot/share/jbooster/utilities/fileUtils.hpp
index 2b5734754..ccc677af7 100644
--- a/src/hotspot/share/jbooster/utilities/fileUtils.hpp
+++ b/src/hotspot/share/jbooster/utilities/fileUtils.hpp
@@ -44,14 +44,13 @@ public:
   static bool exists(const char* path);
   static bool is_file(const char* path);
   static bool is_dir(const char* path);
+  static int64_t file_size(const char* path);
   static uint64_t modify_time(const char* path);
   static bool mkdir(const char* path);
   static bool mkdirs(const char* path);
   static bool rename(const char* path_from, const char* path_to);
   static bool move(const char* path_from, const char* path_to);
   static bool remove(const char* path);
-  static bool is_same(const char* path1, const char* path2);
-  static bool is_same(const char* path, const char* content, int64_t size);
 
   class ListDir: public StackObj {
     const char* _path;
diff --git a/src/jdk.internal.vm.ci/share/classes/jdk/vm/ci/jbooster/JBoosterCompilationContext.java b/src/jdk.internal.vm.ci/share/classes/jdk/vm/ci/jbooster/JBoosterCompilationContext.java
index b26786a52..85c9bf464 100644
--- a/src/jdk.internal.vm.ci/share/classes/jdk/vm/ci/jbooster/JBoosterCompilationContext.java
+++ b/src/jdk.internal.vm.ci/share/classes/jdk/vm/ci/jbooster/JBoosterCompilationContext.java
@@ -94,6 +94,11 @@ public interface JBoosterCompilationContext {
      */
     boolean usePGO();
 
+    /**
+     * return true if compilation uses extra-resolved klasses
+     */
+    boolean resolveExtraKlasses();
+
     /**
      * Get methodCount of CompiledMethodInfo.
      * (To support multi-task concurrent compilation of AOT)
diff --git a/src/jdk.internal.vm.compiler/share/classes/org.graalvm.compiler.java/src/org/graalvm/compiler/java/BytecodeParser.java b/src/jdk.internal.vm.compiler/share/classes/org.graalvm.compiler.java/src/org/graalvm/compiler/java/BytecodeParser.java
index e429e43e8..01ee855d3 100644
--- a/src/jdk.internal.vm.compiler/share/classes/org.graalvm.compiler.java/src/org/graalvm/compiler/java/BytecodeParser.java
+++ b/src/jdk.internal.vm.compiler/share/classes/org.graalvm.compiler.java/src/org/graalvm/compiler/java/BytecodeParser.java
@@ -263,6 +263,7 @@ import static org.graalvm.compiler.nodes.extended.BranchProbabilityNode.LUDICROU
 import static org.graalvm.compiler.nodes.graphbuilderconf.IntrinsicContext.CompilationContext.INLINE_DURING_PARSING;
 import static org.graalvm.compiler.nodes.type.StampTool.isPointerNonNull;
 
+import java.lang.annotation.Target;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Comparator;
@@ -4386,6 +4387,20 @@ public class BytecodeParser implements GraphBuilderContext {
     }
 
     protected void maybeEagerlyResolve(int cpi, int bytecode) {
+        JBoosterCompilationContext ctx = JBoosterCompilationContext.get();
+        if (ctx != null && !ctx.resolveExtraKlasses()) {
+            try{
+                maybeEagerlyResolveBase(cpi, bytecode);
+            } catch (NoClassDefFoundError e) {
+                debug.log("Cannot resolve all elements in constant pool");
+                return;
+            }
+        } else {
+            maybeEagerlyResolveBase(cpi, bytecode);
+        }
+    }
+
+    protected void maybeEagerlyResolveBase(int cpi, int bytecode) {
         if (intrinsicContext != null) {
             constantPool.loadReferencedType(cpi, bytecode);
         } else if (graphBuilderConfig.eagerResolving()) {
@@ -4420,6 +4435,20 @@ public class BytecodeParser implements GraphBuilderContext {
     }
 
     protected JavaType maybeEagerlyResolve(JavaType type, ResolvedJavaType accessingClass) {
+        JBoosterCompilationContext ctx = JBoosterCompilationContext.get();
+        if (ctx != null && !ctx.resolveExtraKlasses()) {
+            try{
+                return maybeEagerlyResolveBase(type, accessingClass);
+            } catch (NoClassDefFoundError e) {
+                debug.log("Cannot resolve all elements in constant pool");
+                return type;
+            }
+        } else {
+            return maybeEagerlyResolveBase(type, accessingClass);
+        }
+    }
+
+    protected JavaType maybeEagerlyResolveBase(JavaType type, ResolvedJavaType accessingClass) {
         if (graphBuilderConfig.eagerResolving() || parsingIntrinsic()) {
             return type.resolve(accessingClass);
         }
diff --git a/src/jdk.jbooster/share/classes/jdk/jbooster/JBooster.java b/src/jdk.jbooster/share/classes/jdk/jbooster/JBooster.java
index db5d916ac..76d62e9b2 100644
--- a/src/jdk.jbooster/share/classes/jdk/jbooster/JBooster.java
+++ b/src/jdk.jbooster/share/classes/jdk/jbooster/JBooster.java
@@ -175,15 +175,15 @@ public final class JBooster {
     /**
      * This method is invoked only in C++.
      */
-    private static boolean compileClasses(int sessionId, String filePath, Set<Class<?>> classes, boolean usePGO) {
-        return compileMethods(sessionId, filePath, classes, null, null, usePGO);
+    private static boolean compileClasses(int sessionId, String filePath, Set<Class<?>> classes, boolean usePGO, boolean resolveExtraKlasses) {
+        return compileMethods(sessionId, filePath, classes, null, null, usePGO, resolveExtraKlasses);
     }
 
     /**
      * This method is invoked only in C++.
      */
     private static boolean compileMethods(int sessionId, String filePath, Set<Class<?>> classes,
-                                          Set<String> methodsToCompile, Set<String> methodsNotToCompile, boolean usePGO) {
+                                          Set<String> methodsToCompile, Set<String> methodsNotToCompile, boolean usePGO, boolean resolveExtraKlasses) {
         LOGGER.log(INFO, "Compilation task received: classes_to_compile={0}, methods_to_compile={1}, methods_not_compile={2}, session_id={3}.",
                 classes.size(),
                 (methodsToCompile == null ? "all" : String.valueOf(methodsToCompile.size())),
@@ -191,7 +191,7 @@ public final class JBooster {
                 sessionId);
         try {
             JBoosterCompilationContextImpl ctx = new JBoosterCompilationContextImpl(
-                    sessionId, filePath, classes, methodsToCompile, methodsNotToCompile, usePGO);
+                    sessionId, filePath, classes, methodsToCompile, methodsNotToCompile, usePGO, resolveExtraKlasses);
             return new Main(ctx).compileForJBooster();
         } catch (Exception e) {
             e.printStackTrace();
diff --git a/src/jdk.jbooster/share/classes/jdk/jbooster/JBoosterCompilationContextImpl.java b/src/jdk.jbooster/share/classes/jdk/jbooster/JBoosterCompilationContextImpl.java
index 4918f8552..efdeee9e5 100644
--- a/src/jdk.jbooster/share/classes/jdk/jbooster/JBoosterCompilationContextImpl.java
+++ b/src/jdk.jbooster/share/classes/jdk/jbooster/JBoosterCompilationContextImpl.java
@@ -43,6 +43,7 @@ public class JBoosterCompilationContextImpl implements JBoosterCompilationContex
     private final Set<String> methodsToCompile;
     private final Set<String> methodsNotToCompile;
     private final boolean usePGO;
+    private final boolean resolveExtraKlasses;
 
     // These values are used to replace the static values in AOT classes.
     private final AtomicInteger compiledMethodInfoMethodsCount = new AtomicInteger(0);
@@ -63,13 +64,15 @@ public class JBoosterCompilationContextImpl implements JBoosterCompilationContex
             Set<Class<?>> classesToCompile,
             Set<String> methodsToCompile,
             Set<String> methodsNotCompile,
-            boolean usePGO) {
+            boolean usePGO,
+            boolean resolveExtraKlasses) {
         this.sessionId = sessionId;
         this.filePath = filePath;
         this.classesToCompile = classesToCompile;
         this.methodsToCompile = methodsToCompile;
         this.methodsNotToCompile = methodsNotCompile;
         this.usePGO = usePGO;
+        this.resolveExtraKlasses = resolveExtraKlasses;
     }
 
     @Override
@@ -102,6 +105,11 @@ public class JBoosterCompilationContextImpl implements JBoosterCompilationContex
         return usePGO;
     }
 
+    @Override
+    public boolean resolveExtraKlasses() {
+        return resolveExtraKlasses;
+    }
+
     @Override
     public AtomicInteger getCompiledMethodInfoMethodsCount() {
         return compiledMethodInfoMethodsCount;
diff --git a/test/hotspot/gtest/jbooster/common.hpp b/test/hotspot/gtest/jbooster/common.hpp
new file mode 100644
index 000000000..15b773786
--- /dev/null
+++ b/test/hotspot/gtest/jbooster/common.hpp
@@ -0,0 +1,96 @@
+/*
+ * Copyright (c) 2020, 2023, Huawei Technologies Co., Ltd. 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 GTEST_JBOOSTER_COMMON
+#define GTEST_JBOOSTER_COMMON
+
+#if INCLUDE_JBOOSTER
+
+// @see src/hotspot/share/memory/filemap.cpp
+#ifndef O_BINARY       // if defined (Win32) use binary files.
+#define O_BINARY 0     // otherwise do nothing.
+#endif
+
+class TestUtils {
+public:
+  static bool is_same(const char* path1, const char* path2) {
+    bool res = false;
+    char* buf1 = nullptr;
+    char* buf2 = nullptr;
+    int fd1 = os::open(path1, O_BINARY | O_RDONLY, 0);
+    int fd2 = os::open(path2, O_BINARY | O_RDONLY, 0);
+    do {
+      if (fd1 < 0 || fd2 < 0) break;
+      int64_t size1 = os::lseek(fd1, 0, SEEK_END);
+      int64_t size2 = os::lseek(fd2, 0, SEEK_END);
+      if (size1 != size2) break;
+      int64_t size = size1;
+      os::lseek(fd1, 0, SEEK_SET);
+      os::lseek(fd2, 0, SEEK_SET);
+      // We don't use NEW_RESOURCE_ARRAY here as Thread::current() may
+      // not be initialized yet.
+      buf1 = NEW_C_HEAP_ARRAY(char, (size_t) size, mtJBooster);
+      buf2 = NEW_C_HEAP_ARRAY(char, (size_t) size, mtJBooster);
+      size1 = (int64_t) os::read(fd1, buf1, size);
+      size2 = (int64_t) os::read(fd2, buf2, size);
+      guarantee(size1 == size && size2 == size, "sanity");
+      res = memcmp(buf1, buf2, size) == 0;
+    } while (false);
+    if (fd1 >= 0) os::close(fd1);
+    if (fd2 >= 0) os::close(fd2);
+    if (buf1 != nullptr) {
+      FREE_C_HEAP_ARRAY(char, buf1);
+    }
+    if (buf2 != nullptr) {
+      FREE_C_HEAP_ARRAY(char, buf2);
+    }
+    return res;
+  }
+
+  static bool is_same(const char* path, const char* content, int64_t size) {
+    bool res = false;
+    char* buf = nullptr;
+    int fd = os::open(path, O_BINARY | O_RDONLY, 0);
+    do {
+      if (fd < 0) break;
+      int64_t fsize = os::lseek(fd, 0, SEEK_END);
+      if (fsize != size) break;
+      os::lseek(fd, 0, SEEK_SET);
+      // We don't use NEW_RESOURCE_ARRAY here as Thread::current() may
+      // not be initialized yet.
+      buf = NEW_C_HEAP_ARRAY(char, (size_t) size, mtJBooster);
+      fsize = (int64_t) os::read(fd, buf, size);
+      guarantee(fsize == size, "sanity");
+      res = memcmp(content, buf, size) == 0;
+    } while (false);
+    if (fd >= 0) os::close(fd);
+    if (buf != nullptr) {
+      FREE_C_HEAP_ARRAY(char, buf);
+    }
+    return res;
+  }
+};
+
+#endif // INCLUDE_JBOOSTER
+
+#endif // GTEST_JBOOSTER_COMMON
\ No newline at end of file
diff --git a/test/hotspot/gtest/jbooster/test_net.cpp b/test/hotspot/gtest/jbooster/test_net.cpp
index bf20b0aa3..dd45dd65a 100644
--- a/test/hotspot/gtest/jbooster/test_net.cpp
+++ b/test/hotspot/gtest/jbooster/test_net.cpp
@@ -26,6 +26,8 @@
 
 #if INCLUDE_JBOOSTER
 
+#include "common.hpp"
+
 #include "jbooster/jClientArguments.hpp"
 #include "jbooster/net/message.inline.hpp"
 #include "jbooster/net/messageBuffer.inline.hpp"
@@ -431,7 +433,7 @@ TEST(JBoosterNet, serializationn_file_wrapper) {
     }
     EXPECT_EQ(times, 3);
   }
-  EXPECT_TRUE(FileUtils::is_same(file_name1, file_name2));
+  EXPECT_TRUE(TestUtils::is_same(file_name1, file_name2));
   FileUtils::remove(file_name1);
   FileUtils::remove(file_name2);
 }
diff --git a/test/hotspot/gtest/jbooster/test_util.cpp b/test/hotspot/gtest/jbooster/test_util.cpp
index cd65804be..461e3faa7 100644
--- a/test/hotspot/gtest/jbooster/test_util.cpp
+++ b/test/hotspot/gtest/jbooster/test_util.cpp
@@ -26,6 +26,8 @@
 
 #if INCLUDE_JBOOSTER
 
+#include "common.hpp"
+
 #include "classfile/symbolTable.hpp"
 #include "jbooster/utilities/concurrentHashMap.inline.hpp"
 #include "jbooster/utilities/debugUtils.inline.hpp"
@@ -95,13 +97,13 @@ TEST(JBoosterUtil, file) {
 
   write_file("gtest-jbooster-tmp5", "12345");
   write_file("gtest-jbooster-tmp6", "12345");
-  EXPECT_TRUE(FileUtils::is_same("gtest-jbooster-tmp5", "gtest-jbooster-tmp6"));
+  EXPECT_TRUE(TestUtils::is_same("gtest-jbooster-tmp5", "gtest-jbooster-tmp6"));
   write_file("gtest-jbooster-tmp6", "123456");
-  EXPECT_FALSE(FileUtils::is_same("gtest-jbooster-tmp5", "gtest-jbooster-tmp6"));
+  EXPECT_FALSE(TestUtils::is_same("gtest-jbooster-tmp5", "gtest-jbooster-tmp6"));
 
-  EXPECT_TRUE(FileUtils::is_same("gtest-jbooster-tmp5", "12345", 6));
-  EXPECT_FALSE(FileUtils::is_same("gtest-jbooster-tmp5", "12346", 6));
-  EXPECT_FALSE(FileUtils::is_same("gtest-jbooster-tmp5", "123456", 7));
+  EXPECT_TRUE(TestUtils::is_same("gtest-jbooster-tmp5", "12345", 6));
+  EXPECT_FALSE(TestUtils::is_same("gtest-jbooster-tmp5", "12346", 6));
+  EXPECT_FALSE(TestUtils::is_same("gtest-jbooster-tmp5", "123456", 7));
 
   EXPECT_FALSE(FileUtils::is_file("gtest-jbooster-tmp4"));
   EXPECT_TRUE(FileUtils::is_dir("gtest-jbooster-tmp4"));
diff --git a/test/jdk/tools/jbooster/JBoosterTestBase.java b/test/jdk/tools/jbooster/JBoosterTestBase.java
index 08792a77b..4c1dcbd4a 100644
--- a/test/jdk/tools/jbooster/JBoosterTestBase.java
+++ b/test/jdk/tools/jbooster/JBoosterTestBase.java
@@ -44,11 +44,12 @@ import jdk.test.lib.Utils;
  * @see JcmdBase
  */
 public class JBoosterTestBase {
-    public static final int WAIT_START_TIME = 2;
+    public static final int WAIT_START_TIME = 128;
     public static final int WAIT_SHORT_TIME = 8;
     public static final int WAIT_EXIT_TIME = 64;
 
     public static final int SERVER_PORT = 41567;
+    public static final int SERVER_CONNECTION_TIMEOUT = 256 * 1000;
     public static final String SERVER_PORT_STR = "41567";
 
     public static final String CLIENT_CACHE_PATH = "jbooster-cache-client";
@@ -80,7 +81,8 @@ public class JBoosterTestBase {
 
     public static final List<String> SERVER_STANDARD_ARGS = List.of(
             "--server-port=" + SERVER_PORT_STR,
-            "--cache-path=" + SERVER_CACHE_PATH
+            "--cache-path=" + SERVER_CACHE_PATH,
+            "--connection-timeout=" + SERVER_CONNECTION_TIMEOUT
     );
 
     private static final ProcessBuilder processBuilder = new ProcessBuilder();
-- 
2.22.0