8371046: Segfault in compiler/whitebox/StressNMethodRelocation.java with -XX:+UseZGC

Reviewed-by: kvn, eastigeevich

src/hotspot/cpu/aarch64/relocInfo_aarch64.cpp

@@ -85,7 +85,7 @@ void Relocation::pd_set_call_destination(address x) {
   } else {
     MacroAssembler::pd_patch_instruction(addr(), x);
   }
-  assert(pd_call_destination(addr()) == x, "fail in reloc");
+  guarantee(pd_call_destination(addr()) == x, "fail in reloc");
 }
 
 void trampoline_stub_Relocation::pd_fix_owner_after_move() {

src/hotspot/share/asm/codeBuffer.cpp

@@ -90,7 +90,7 @@ typedef CodeBuffer::csize_t csize_t;  // file-local definition
 
 // External buffer, in a predefined CodeBlob.
 // Important: The code_start must be taken exactly, and not realigned.
-CodeBuffer::CodeBuffer(CodeBlob* blob) DEBUG_ONLY(: Scrubber(this, sizeof(*this))) {
+CodeBuffer::CodeBuffer(const CodeBlob* blob) DEBUG_ONLY(: Scrubber(this, sizeof(*this))) {
   // Provide code buffer with meaningful name
   initialize_misc(blob->name());
   initialize(blob->content_begin(), blob->content_size());

src/hotspot/share/asm/codeBuffer.hpp

@@ -672,7 +672,7 @@ class CodeBuffer: public StackObj DEBUG_ONLY(COMMA private Scrubber) {
   }
 
   // (2) CodeBuffer referring to pre-allocated CodeBlob.
-  CodeBuffer(CodeBlob* blob);
+  CodeBuffer(const CodeBlob* blob);
 
   // (3) code buffer allocating codeBlob memory for code & relocation
   // info but with lazy initialization.  The name must be something

src/hotspot/share/code/nmethod.cpp

@@ -1498,6 +1498,40 @@ nmethod::nmethod(const nmethod &nm) : CodeBlob(nm._name, nm._kind, nm._size, nm.
   //   - OOP table
   memcpy(consts_begin(), nm.consts_begin(), nm.data_end() - nm.consts_begin());
 
+  // Fix relocation
+  RelocIterator iter(this);
+  CodeBuffer src(&nm);
+  CodeBuffer dst(this);
+  while (iter.next()) {
+#ifdef USE_TRAMPOLINE_STUB_FIX_OWNER
+    // After an nmethod is moved, some direct call sites may end up out of range.
+    // CallRelocation::fix_relocation_after_move() assumes the target is always
+    // reachable and does not check branch range. Calling it without range checks
+    // could cause us to write an offset too large for the instruction.
+    //
+    // If a call site has a trampoline, we skip the normal call relocation. The
+    // associated trampoline_stub_Relocation will handle the call and the
+    // trampoline, including range checks and updating the branch as needed.
+    //
+    // If no trampoline exists, we can assume the call target is always
+    // reachable and therefore within direct branch range, so calling
+    // CallRelocation::fix_relocation_after_move() is safe.
+    if (iter.reloc()->is_call()) {
+      address trampoline = trampoline_stub_Relocation::get_trampoline_for(iter.reloc()->addr(), this);
+      if (trampoline != nullptr) {
+        continue;
+      }
+    }
+#endif
+
+    iter.reloc()->fix_relocation_after_move(&src, &dst);
+  }
+
+  {
+    MutexLocker ml(NMethodState_lock, Mutex::_no_safepoint_check_flag);
+    clear_inline_caches();
+  }
+
   post_init();
 }
 
@@ -1521,25 +1555,6 @@ nmethod* nmethod::relocate(CodeBlobType code_blob_type) {
     return nullptr;
   }
 
-  // Fix relocation
-  RelocIterator iter(nm_copy);
-  CodeBuffer src(this);
-  CodeBuffer dst(nm_copy);
-  while (iter.next()) {
-#ifdef USE_TRAMPOLINE_STUB_FIX_OWNER
-    // Direct calls may no longer be in range and the use of a trampoline may now be required.
-    // Instead, allow trampoline relocations to update their owners and perform the necessary checks.
-    if (iter.reloc()->is_call()) {
-      address trampoline = trampoline_stub_Relocation::get_trampoline_for(iter.reloc()->addr(), nm_copy);
-      if (trampoline != nullptr) {
-        continue;
-      }
-    }
-#endif
-
-    iter.reloc()->fix_relocation_after_move(&src, &dst);
-  }
-
   // To make dependency checking during class loading fast, record
   // the nmethod dependencies in the classes it is dependent on.
   // This allows the dependency checking code to simply walk the
@@ -1569,8 +1584,6 @@ nmethod* nmethod::relocate(CodeBlobType code_blob_type) {
   if (!is_marked_for_deoptimization() && is_in_use()) {
     assert(method() != nullptr && method()->code() == this, "should be if is in use");
 
-    nm_copy->clear_inline_caches();
-
     // Attempt to start using the copy
     if (nm_copy->make_in_use()) {
       ICache::invalidate_range(nm_copy->code_begin(), nm_copy->code_size());
@@ -1578,7 +1591,7 @@ nmethod* nmethod::relocate(CodeBlobType code_blob_type) {
       methodHandle mh(Thread::current(), nm_copy->method());
       nm_copy->method()->set_code(mh, nm_copy);
 
-      make_not_used();
+      make_not_entrant(InvalidationReason::RELOCATED);
 
       nm_copy->post_compiled_method_load_event();
 

src/hotspot/share/code/nmethod.hpp

@@ -499,6 +499,7 @@ public:
     UNCOMMON_TRAP,
     WHITEBOX_DEOPTIMIZATION,
     ZOMBIE,
+    RELOCATED,
     INVALIDATION_REASONS_COUNT
   };
 
@@ -543,6 +544,8 @@ public:
         return "whitebox deoptimization";
       case InvalidationReason::ZOMBIE:
         return "zombie";
+      case InvalidationReason::RELOCATED:
+        return "relocated";
       default: {
         assert(false, "Unhandled reason");
         return "Unknown";

src/hotspot/share/jvmci/vmStructs_jvmci.cpp

@@ -586,6 +586,7 @@
   declare_constant(nmethod::InvalidationReason::UNCOMMON_TRAP)                            \
   declare_constant(nmethod::InvalidationReason::WHITEBOX_DEOPTIMIZATION)                  \
   declare_constant(nmethod::InvalidationReason::ZOMBIE)                                   \
+  declare_constant(nmethod::InvalidationReason::RELOCATED)                                \
                                                                                           \
   declare_constant(CodeInstaller::VERIFIED_ENTRY)                         \
   declare_constant(CodeInstaller::UNVERIFIED_ENTRY)                       \

src/hotspot/share/prims/whitebox.cpp

@@ -1678,7 +1678,7 @@ WB_ENTRY(void, WB_RelocateNMethodFromAddr(JNIEnv* env, jobject o, jlong addr, ji
   CodeBlob* blob = CodeCache::find_blob(address);
   if (blob != nullptr && blob->is_nmethod()) {
     nmethod* code = blob->as_nmethod();
-    if (code->is_in_use()) {
+    if (code->is_in_use() && !code->is_unloading()) {
       CompiledICLocker ic_locker(code);
       code->relocate(static_cast<CodeBlobType>(blob_type));
     }