8369238: Allow virtual thread preemption on some common class initialization paths

src/hotspot/cpu/aarch64/continuationFreezeThaw_aarch64.inline.hpp

@@ -200,6 +200,41 @@ inline void FreezeBase::patch_pd_unused(intptr_t* sp) {
   *fp_addr = badAddressVal;
 }
 
+inline intptr_t* AnchorMark::anchor_mark_set_pd() {
+  intptr_t* sp = _top_frame.sp();
+  if (_top_frame.is_interpreted_frame()) {
+    // In case the top frame is interpreted we need to set up the anchor using
+    // the last_sp saved in the frame (remove possible alignment added while
+    // thawing, see ThawBase::finish_thaw()). We also clear last_sp to match
+    // the behavior when calling the VM from the interpreter (we check for this
+    // in FreezeBase::prepare_freeze_interpreted_top_frame, which can be reached
+    // if preempting again at redo_vmcall()).
+    _last_sp_from_frame = _top_frame.interpreter_frame_last_sp();
+    assert(_last_sp_from_frame != nullptr, "");
+    _top_frame.interpreter_frame_set_last_sp(nullptr);
+    if (sp != _last_sp_from_frame) {
+      // We need to move up return pc and fp. They will be read next in
+      // set_anchor() and set as _last_Java_pc and _last_Java_fp respectively.
+      _last_sp_from_frame[-1] = (intptr_t)_top_frame.pc();
+      _last_sp_from_frame[-2] = (intptr_t)_top_frame.fp();
+    }
+    _is_interpreted = true;
+    sp = _last_sp_from_frame;
+  }
+  return sp;
+}
+
+inline void AnchorMark::anchor_mark_clear_pd() {
+  if (_is_interpreted) {
+    // Restore last_sp_from_frame and possibly overwritten pc.
+    _top_frame.interpreter_frame_set_last_sp(_last_sp_from_frame);
+    intptr_t* sp = _top_frame.sp();
+    if (sp != _last_sp_from_frame) {
+      sp[-1] = (intptr_t)_top_frame.pc();
+    }
+  }
+}
+
 //////// Thaw
 
 // Fast path
@@ -304,10 +339,17 @@ inline intptr_t* ThawBase::push_cleanup_continuation() {
   frame enterSpecial = new_entry_frame();
   intptr_t* sp = enterSpecial.sp();
 
+  // We only need to set the return pc. rfp will be restored back in gen_continuation_enter().
   sp[-1] = (intptr_t)ContinuationEntry::cleanup_pc();
-  sp[-2] = (intptr_t)enterSpecial.fp();
+  return sp;
+}
+
+inline intptr_t* ThawBase::push_preempt_adapter() {
+  frame enterSpecial = new_entry_frame();
+  intptr_t* sp = enterSpecial.sp();
 
-  log_develop_trace(continuations, preempt)("push_cleanup_continuation initial sp: " INTPTR_FORMAT " final sp: " INTPTR_FORMAT, p2i(sp + 2 * frame::metadata_words), p2i(sp));
+  // We only need to set the return pc. rfp will be restored back in generate_cont_preempt_stub().
+  sp[-1] = (intptr_t)StubRoutines::cont_preempt_stub();
   return sp;
 }
 

src/hotspot/cpu/aarch64/continuationHelper_aarch64.inline.hpp

@@ -52,6 +52,9 @@ static inline void patch_return_pc_with_preempt_stub(frame& f) {
     // The target will check for preemption once it returns to the interpreter
     // or the native wrapper code and will manually jump to the preempt stub.
     JavaThread *thread = JavaThread::current();
+    DEBUG_ONLY(Method* m = f.is_interpreted_frame() ? f.interpreter_frame_method() : f.cb()->as_nmethod()->method();)
+    assert(m->is_object_wait0() || thread->interp_at_preemptable_vmcall_cnt() > 0,
+           "preemptable VM call not using call_VM_preemptable");
     thread->set_preempt_alternate_return(StubRoutines::cont_preempt_stub());
   }
 }

src/hotspot/cpu/aarch64/interp_masm_aarch64.cpp

@@ -479,6 +479,14 @@ void InterpreterMacroAssembler::remove_activation(TosState state,
   // result check if synchronized method
   Label unlocked, unlock, no_unlock;
 
+#ifdef ASSERT
+  Label not_preempted;
+  ldr(rscratch1, Address(rthread, JavaThread::preempt_alternate_return_offset()));
+  cbz(rscratch1, not_preempted);
+  stop("remove_activation: should not have alternate return address set");
+  bind(not_preempted);
+#endif /* ASSERT */
+
   // get the value of _do_not_unlock_if_synchronized into r3
   const Address do_not_unlock_if_synchronized(rthread,
     in_bytes(JavaThread::do_not_unlock_if_synchronized_offset()));
@@ -1371,6 +1379,7 @@ void InterpreterMacroAssembler::call_VM_leaf_base(address entry_point,
 void InterpreterMacroAssembler::call_VM_base(Register oop_result,
                                              Register java_thread,
                                              Register last_java_sp,
+                                             Label*   return_pc,
                                              address  entry_point,
                                              int      number_of_arguments,
                                              bool     check_exceptions) {
@@ -1394,39 +1403,58 @@ void InterpreterMacroAssembler::call_VM_base(Register oop_result,
 #endif /* ASSERT */
   // super call
   MacroAssembler::call_VM_base(oop_result, noreg, last_java_sp,
-                               entry_point, number_of_arguments,
-                     check_exceptions);
+                               return_pc, entry_point,
+                               number_of_arguments, check_exceptions);
 // interpreter specific
   restore_bcp();
   restore_locals();
 }
 
-void InterpreterMacroAssembler::call_VM_preemptable(Register oop_result,
-                                                    address entry_point,
-                                                    Register arg_1) {
-  assert(arg_1 == c_rarg1, "");
+void InterpreterMacroAssembler::call_VM_preemptable_helper(Register oop_result,
+                                                           address entry_point,
+                                                           int number_of_arguments,
+                                                           bool check_exceptions) {
+  assert(InterpreterRuntime::is_preemptable_call(entry_point), "VM call not preemptable, should use call_VM()");
   Label resume_pc, not_preempted;
 
 #ifdef ASSERT
   {
-    Label L;
+    Label L1, L2;
     ldr(rscratch1, Address(rthread, JavaThread::preempt_alternate_return_offset()));
-    cbz(rscratch1, L);
-    stop("Should not have alternate return address set");
-    bind(L);
+    cbz(rscratch1, L1);
+    stop("call_VM_preemptable_helper: Should not have alternate return address set");
+    bind(L1);
+    // We check this counter in patch_return_pc_with_preempt_stub() during freeze.
+    incrementw(Address(rthread, JavaThread::interp_at_preemptable_vmcall_cnt_offset()));
+    ldrw(rscratch1, Address(rthread, JavaThread::interp_at_preemptable_vmcall_cnt_offset()));
+    cmpw(rscratch1, 0);
+    br(Assembler::GT, L2);
+    stop("call_VM_preemptable_helper: should be > 0");
+    bind(L2);
   }
 #endif /* ASSERT */
 
   // Force freeze slow path.
   push_cont_fastpath();
 
   // Make VM call. In case of preemption set last_pc to the one we want to resume to.
-  adr(rscratch1, resume_pc);
-  str(rscratch1, Address(rthread, JavaThread::last_Java_pc_offset()));
-  call_VM_base(oop_result, noreg, noreg, entry_point, 1, false /*check_exceptions*/);
+  // Note: call_VM_base will use resume_pc label to set last_Java_pc.
+  call_VM_base(noreg, noreg, noreg, &resume_pc, entry_point, number_of_arguments, false /*check_exceptions*/);
 
   pop_cont_fastpath();
 
+#ifdef ASSERT
+  {
+    Label L;
+    decrementw(Address(rthread, JavaThread::interp_at_preemptable_vmcall_cnt_offset()));
+    ldrw(rscratch1, Address(rthread, JavaThread::interp_at_preemptable_vmcall_cnt_offset()));
+    cmpw(rscratch1, 0);
+    br(Assembler::GE, L);
+    stop("call_VM_preemptable_helper: should be >= 0");
+    bind(L);
+  }
+#endif /* ASSERT */
+
   // Check if preempted.
   ldr(rscratch1, Address(rthread, JavaThread::preempt_alternate_return_offset()));
   cbz(rscratch1, not_preempted);
@@ -1438,6 +1466,51 @@ void InterpreterMacroAssembler::call_VM_preemptable(Register oop_result,
   restore_after_resume(false /* is_native */);
 
   bind(not_preempted);
+  if (check_exceptions) {
+    // check for pending exceptions
+    ldr(rscratch1, Address(rthread, in_bytes(Thread::pending_exception_offset())));
+    Label ok;
+    cbz(rscratch1, ok);
+    lea(rscratch1, RuntimeAddress(StubRoutines::forward_exception_entry()));
+    br(rscratch1);
+    bind(ok);
+  }
+
+  // get oop result if there is one and reset the value in the thread
+  if (oop_result->is_valid()) {
+    get_vm_result_oop(oop_result, rthread);
+  }
+}
+
+static void pass_arg1(MacroAssembler* masm, Register arg) {
+  if (c_rarg1 != arg ) {
+    masm->mov(c_rarg1, arg);
+  }
+}
+
+static void pass_arg2(MacroAssembler* masm, Register arg) {
+  if (c_rarg2 != arg ) {
+    masm->mov(c_rarg2, arg);
+  }
+}
+
+void InterpreterMacroAssembler::call_VM_preemptable(Register oop_result,
+                                         address entry_point,
+                                         Register arg_1,
+                                         bool check_exceptions) {
+  pass_arg1(this, arg_1);
+  call_VM_preemptable_helper(oop_result, entry_point, 1, check_exceptions);
+}
+
+void InterpreterMacroAssembler::call_VM_preemptable(Register oop_result,
+                                         address entry_point,
+                                         Register arg_1,
+                                         Register arg_2,
+                                         bool check_exceptions) {
+  LP64_ONLY(assert_different_registers(arg_1, c_rarg2));
+  pass_arg2(this, arg_2);
+  pass_arg1(this, arg_1);
+  call_VM_preemptable_helper(oop_result, entry_point, 2, check_exceptions);
 }
 
 void InterpreterMacroAssembler::restore_after_resume(bool is_native) {

src/hotspot/cpu/aarch64/interp_masm_aarch64.hpp

@@ -45,6 +45,7 @@ class InterpreterMacroAssembler: public MacroAssembler {
   virtual void call_VM_base(Register oop_result,
                             Register java_thread,
                             Register last_java_sp,
+                            Label*   return_pc,
                             address  entry_point,
                             int number_of_arguments,
                             bool check_exceptions);
@@ -58,11 +59,24 @@ class InterpreterMacroAssembler: public MacroAssembler {
 
   void load_earlyret_value(TosState state);
 
+  // Use for vthread preemption
   void call_VM_preemptable(Register oop_result,
                            address entry_point,
-                           Register arg_1);
+                           Register arg_1,
+                           bool check_exceptions = true);
+  void call_VM_preemptable(Register oop_result,
+                           address entry_point,
+                           Register arg_1,
+                           Register arg_2,
+                           bool check_exceptions = true);
   void restore_after_resume(bool is_native);
+ private:
+  void call_VM_preemptable_helper(Register oop_result,
+                                  address entry_point,
+                                  int number_of_arguments,
+                                  bool check_exceptions);
 
+ public:
   void jump_to_entry(address entry);
 
   virtual void check_and_handle_popframe(Register java_thread);

src/hotspot/cpu/aarch64/macroAssembler_aarch64.cpp

@@ -743,13 +743,10 @@ static void pass_arg3(MacroAssembler* masm, Register arg) {
   }
 }
 
-static bool is_preemptable(address entry_point) {
-  return entry_point == CAST_FROM_FN_PTR(address, InterpreterRuntime::monitorenter);
-}
-
 void MacroAssembler::call_VM_base(Register oop_result,
                                   Register java_thread,
                                   Register last_java_sp,
+                                  Label*   return_pc,
                                   address  entry_point,
                                   int      number_of_arguments,
                                   bool     check_exceptions) {
@@ -782,12 +779,7 @@ void MacroAssembler::call_VM_base(Register oop_result,
   assert(last_java_sp != rfp, "can't use rfp");
 
   Label l;
-  if (is_preemptable(entry_point)) {
-    // skip setting last_pc since we already set it to desired value.
-    set_last_Java_frame(last_java_sp, rfp, noreg, rscratch1);
-  } else {
-    set_last_Java_frame(last_java_sp, rfp, l, rscratch1);
-  }
+  set_last_Java_frame(last_java_sp, rfp, return_pc != nullptr ? *return_pc : l, rscratch1);
 
   // do the call, remove parameters
   MacroAssembler::call_VM_leaf_base(entry_point, number_of_arguments, &l);
@@ -822,7 +814,7 @@ void MacroAssembler::call_VM_base(Register oop_result,
 }
 
 void MacroAssembler::call_VM_helper(Register oop_result, address entry_point, int number_of_arguments, bool check_exceptions) {
-  call_VM_base(oop_result, noreg, noreg, entry_point, number_of_arguments, check_exceptions);
+  call_VM_base(oop_result, noreg, noreg, nullptr, entry_point, number_of_arguments, check_exceptions);
 }
 
 // Check the entry target is always reachable from any branch.
@@ -1080,7 +1072,7 @@ void MacroAssembler::call_VM(Register oop_result,
                              address entry_point,
                              int number_of_arguments,
                              bool check_exceptions) {
-  call_VM_base(oop_result, rthread, last_java_sp, entry_point, number_of_arguments, check_exceptions);
+  call_VM_base(oop_result, rthread, last_java_sp, nullptr, entry_point, number_of_arguments, check_exceptions);
 }
 
 void MacroAssembler::call_VM(Register oop_result,

src/hotspot/cpu/aarch64/macroAssembler_aarch64.hpp

@@ -81,6 +81,7 @@ class MacroAssembler: public Assembler {
     Register oop_result,               // where an oop-result ends up if any; use noreg otherwise
     Register java_thread,              // the thread if computed before     ; use noreg otherwise
     Register last_java_sp,             // to set up last_Java_frame in stubs; use noreg otherwise
+    Label*   return_pc,                // to set up last_Java_frame; use nullptr otherwise
     address  entry_point,              // the entry point
     int      number_of_arguments,      // the number of arguments (w/o thread) to pop after the call
     bool     check_exceptions          // whether to check for pending exceptions after return

src/hotspot/cpu/aarch64/smallRegisterMap_aarch64.inline.hpp

@@ -28,18 +28,17 @@
 #include "runtime/frame.inline.hpp"
 #include "runtime/registerMap.hpp"
 
+class SmallRegisterMap;
+
 // Java frames don't have callee saved registers (except for rfp), so we can use a smaller RegisterMap
-class SmallRegisterMap {
-  constexpr SmallRegisterMap() = default;
-  ~SmallRegisterMap() = default;
-  NONCOPYABLE(SmallRegisterMap);
+template <bool IncludeArgs>
+class SmallRegisterMapType {
+  friend SmallRegisterMap;
+
+  constexpr SmallRegisterMapType() = default;
+  ~SmallRegisterMapType() = default;
+  NONCOPYABLE(SmallRegisterMapType);
 
-public:
-  static const SmallRegisterMap* instance() {
-    static constexpr SmallRegisterMap the_instance{};
-    return &the_instance;
-  }
-private:
   static void assert_is_rfp(VMReg r) NOT_DEBUG_RETURN
                                      DEBUG_ONLY({ assert (r == rfp->as_VMReg() || r == rfp->as_VMReg()->next(), "Reg: %s", r->name()); })
 public:
@@ -71,7 +70,7 @@ class SmallRegisterMap {
 
   bool update_map()    const { return false; }
   bool walk_cont()     const { return false; }
-  bool include_argument_oops() const { return false; }
+  bool include_argument_oops() const { return IncludeArgs; }
   void set_include_argument_oops(bool f)  {}
   bool in_cont()       const { return false; }
   stackChunkHandle stack_chunk() const { return stackChunkHandle(); }

src/hotspot/cpu/aarch64/stackChunkFrameStream_aarch64.inline.hpp

@@ -108,17 +108,14 @@ inline int StackChunkFrameStream<frame_kind>::interpreter_frame_stack_argsize()
 }
 
 template <ChunkFrames frame_kind>
-inline int StackChunkFrameStream<frame_kind>::interpreter_frame_num_oops() const {
+template <typename RegisterMapT>
+inline int StackChunkFrameStream<frame_kind>::interpreter_frame_num_oops(RegisterMapT* map) const {
   assert_is_interpreted_and_frame_type_mixed();
   ResourceMark rm;
-  InterpreterOopMap mask;
   frame f = to_frame();
-  f.interpreted_frame_oop_map(&mask);
-  return  mask.num_oops()
-        + 1 // for the mirror oop
-        + (f.interpreter_frame_method()->is_native() ? 1 : 0) // temp oop slot
-        + pointer_delta_as_int((intptr_t*)f.interpreter_frame_monitor_begin(),
-              (intptr_t*)f.interpreter_frame_monitor_end())/BasicObjectLock::size();
+  InterpreterOopCount closure;
+  f.oops_interpreted_do(&closure, map);
+  return closure.count();
 }
 
 template<>

src/hotspot/cpu/aarch64/templateTable_aarch64.cpp

@@ -2304,7 +2304,7 @@ void TemplateTable::resolve_cache_and_index_for_method(int byte_no,
   // Class initialization barrier slow path lands here as well.
   address entry = CAST_FROM_FN_PTR(address, InterpreterRuntime::resolve_from_cache);
   __ mov(temp, (int) code);
-  __ call_VM(noreg, entry, temp);
+  __ call_VM_preemptable(noreg, entry, temp);
 
   // Update registers with resolved info
   __ load_method_entry(Rcache, index);
@@ -2356,7 +2356,7 @@ void TemplateTable::resolve_cache_and_index_for_field(int byte_no,
   // Class initialization barrier slow path lands here as well.
   address entry = CAST_FROM_FN_PTR(address, InterpreterRuntime::resolve_from_cache);
   __ mov(temp, (int) code);
-  __ call_VM(noreg, entry, temp);
+  __ call_VM_preemptable(noreg, entry, temp);
 
   // Update registers with resolved info
   __ load_field_entry(Rcache, index);
@@ -3700,7 +3700,7 @@ void TemplateTable::_new() {
   __ bind(slow_case);
   __ get_constant_pool(c_rarg1);
   __ get_unsigned_2_byte_index_at_bcp(c_rarg2, 1);
-  call_VM(r0, CAST_FROM_FN_PTR(address, InterpreterRuntime::_new), c_rarg1, c_rarg2);
+  __ call_VM_preemptable(r0, CAST_FROM_FN_PTR(address, InterpreterRuntime::_new), c_rarg1, c_rarg2);
   __ verify_oop(r0);
 
   // continue