8328306: AArch64: MacOS lazy JIT "write xor execute" switching

Co-authored-by: Dean Long <dlong@openjdk.org>
Reviewed-by: dlong, adinn

Author: Andrew Haley

src/hotspot/cpu/aarch64/assembler_aarch64.hpp

@@ -4329,6 +4329,7 @@ template<typename R, typename... Rx>
 #undef INSN
 
   Assembler(CodeBuffer* code) : AbstractAssembler(code) {
+    MACOS_AARCH64_ONLY(os::thread_wx_enable_write());
   }
 
   // Stack overflow checking

src/hotspot/cpu/aarch64/gc/shared/barrierSetNMethod_aarch64.cpp

@@ -209,6 +209,10 @@ void BarrierSetNMethod::set_guard_value(nmethod* nm, int value, int bit_mask) {
     bs_asm->increment_patching_epoch();
   }
 
+  // Enable WXWrite: the function is called directly from nmethod_entry_barrier
+  // stub.
+  MACOS_AARCH64_ONLY(ThreadWXEnable wx(WXWrite, Thread::current()));
+
   NativeNMethodBarrier barrier(nm);
   barrier.set_value(value, bit_mask);
 }

src/hotspot/cpu/aarch64/macroAssembler_aarch64.cpp

@@ -473,6 +473,7 @@ address MacroAssembler::target_addr_for_insn(address insn_addr) {
 // Patch any kind of instruction; there may be several instructions.
 // Return the total length (in bytes) of the instructions.
 int MacroAssembler::pd_patch_instruction_size(address insn_addr, address target) {
+  MACOS_AARCH64_ONLY(os::thread_wx_enable_write());
   return RelocActions<Patcher>::run(insn_addr, target);
 }
 
@@ -481,6 +482,8 @@ int MacroAssembler::patch_oop(address insn_addr, address o) {
   unsigned insn = *(unsigned*)insn_addr;
   assert(nativeInstruction_at(insn_addr+4)->is_movk(), "wrong insns in patch");
 
+  MACOS_AARCH64_ONLY(os::thread_wx_enable_write());
+
   // OOPs are either narrow (32 bits) or wide (48 bits).  We encode
   // narrow OOPs by setting the upper 16 bits in the first
   // instruction.
@@ -510,6 +513,8 @@ int MacroAssembler::patch_narrow_klass(address insn_addr, narrowKlass n) {
   assert(Instruction_aarch64::extract(insn->encoding(), 31, 21) == 0b11010010101 &&
          nativeInstruction_at(insn_addr+4)->is_movk(), "wrong insns in patch");
 
+  MACOS_AARCH64_ONLY(os::thread_wx_enable_write());
+
   Instruction_aarch64::patch(insn_addr, 20, 5, n >> 16);
   Instruction_aarch64::patch(insn_addr+4, 20, 5, n & 0xffff);
   return 2 * NativeInstruction::instruction_size;

src/hotspot/cpu/aarch64/nativeInst_aarch64.cpp

@@ -133,7 +133,6 @@ void NativeMovConstReg::verify() {
 
 
 intptr_t NativeMovConstReg::data() const {
-  // das(uint64_t(instruction_address()),2);
   address addr = MacroAssembler::target_addr_for_insn(instruction_address());
   if (maybe_cpool_ref(instruction_address())) {
     return *(intptr_t*)addr;
@@ -144,6 +143,7 @@ intptr_t NativeMovConstReg::data() const {
 
 void NativeMovConstReg::set_data(intptr_t x) {
   if (maybe_cpool_ref(instruction_address())) {
+    MACOS_AARCH64_ONLY(os::thread_wx_enable_write());
     address addr = MacroAssembler::target_addr_for_insn(instruction_address());
     *(intptr_t*)addr = x;
   } else {
@@ -350,8 +350,6 @@ bool NativeInstruction::is_stop() {
 
 //-------------------------------------------------------------------
 
-void NativeGeneralJump::verify() {  }
-
 // MT-safe patching of a long jump instruction.
 void NativeGeneralJump::replace_mt_safe(address instr_addr, address code_buffer) {
   ShouldNotCallThis();

src/hotspot/cpu/aarch64/nativeInst_aarch64.hpp

@@ -90,16 +90,18 @@ class NativeInstruction {
 
   s_char sbyte_at(int offset) const { return *(s_char*)addr_at(offset); }
   u_char ubyte_at(int offset) const { return *(u_char*)addr_at(offset); }
-  jint int_at(int offset) const { return *(jint*)addr_at(offset); }
-  juint uint_at(int offset) const { return *(juint*)addr_at(offset); }
-  address ptr_at(int offset) const { return *(address*)addr_at(offset); }
-  oop oop_at(int offset) const { return *(oop*)addr_at(offset); }
-
-  void set_char_at(int offset, char c) { *addr_at(offset) = (u_char)c; }
-  void set_int_at(int offset, jint i) { *(jint*)addr_at(offset) = i; }
-  void set_uint_at(int offset, jint i) { *(juint*)addr_at(offset) = i; }
-  void set_ptr_at(int offset, address ptr) { *(address*)addr_at(offset) = ptr; }
-  void set_oop_at(int offset, oop o) { *(oop*)addr_at(offset) = o; }
+  jint int_at(int offset) const     { return *(jint*)addr_at(offset); }
+  juint uint_at(int offset) const   { return *(juint*)addr_at(offset); }
+  address ptr_at(int offset) const  { return *(address*)addr_at(offset); }
+  oop oop_at(int offset) const      { return *(oop*)addr_at(offset); }
+
+#define MACOS_WX_WRITE MACOS_AARCH64_ONLY(os::thread_wx_enable_write())
+  void set_char_at(int offset, char c)     { MACOS_WX_WRITE;  *addr_at(offset) = (u_char)c; }
+  void set_int_at(int offset, jint i)      { MACOS_WX_WRITE;  *(jint*)addr_at(offset) = i; }
+  void set_uint_at(int offset, jint i)     { MACOS_WX_WRITE;  *(juint*)addr_at(offset) = i; }
+  void set_ptr_at(int offset, address ptr) { MACOS_WX_WRITE;  *(address*)addr_at(offset) = ptr; }
+  void set_oop_at(int offset, oop o)       { MACOS_WX_WRITE;  *(oop*)addr_at(offset) = o; }
+#undef MACOS_WX_WRITE
 
   void wrote(int offset);
 
@@ -380,7 +382,6 @@ class NativeGeneralJump: public NativeJump {
   void set_jump_destination(address dest);
 
   static void replace_mt_safe(address instr_addr, address code_buffer);
-  static void verify();
 };
 
 inline NativeGeneralJump* nativeGeneralJump_at(address address) {

src/hotspot/cpu/aarch64/stubGenerator_aarch64.cpp

@@ -11742,7 +11742,9 @@ class StubGenerator: public StubCodeGenerator {
     }
 #endif
 
-    StubRoutines::_unsafe_setmemory = generate_unsafe_setmemory();
+    if (vmIntrinsics::is_intrinsic_available(vmIntrinsics::_setMemory)) {
+      StubRoutines::_unsafe_setmemory = generate_unsafe_setmemory();
+    }
 
     StubRoutines::aarch64::set_completed(); // Inidicate that arraycopy and zero_blocks stubs are generated
   }

src/hotspot/cpu/aarch64/vm_version_aarch64.cpp

@@ -622,6 +622,22 @@ void VM_Version::initialize() {
 
   check_virtualizations();
 
+#ifdef __APPLE__
+  DefaultWXWriteMode = UseOldWX ? WXWrite : WXArmedForWrite;
+
+  if (TraceWXHealing) {
+    if (pthread_jit_write_protect_supported_np()) {
+      tty->print_cr("### TraceWXHealing is in use");
+      if (StressWXHealing) {
+        tty->print_cr("### StressWXHealing is in use");
+      }
+    } else {
+      tty->print_cr("WX Healing is not in use because MAP_JIT write protection "
+                    "does not work on this system.");
+    }
+  }
+#endif
+
   // Sync SVE related CPU features with flags
   if (UseSVE < 2) {
     clear_feature(CPU_SVE2);

src/hotspot/os/bsd/globals_bsd.hpp

@@ -28,16 +28,29 @@
 //
 // Declare Bsd specific flags. They are not available on other platforms.
 //
+#ifdef AARCH64
 #define RUNTIME_OS_FLAGS(develop,                                       \
                          develop_pd,                                    \
                          product,                                       \
                          product_pd,                                    \
                          range,                                         \
                          constraint)                                    \
                                                                         \
-  AARCH64_ONLY(develop(bool, AssertWXAtThreadSync, true,                \
-          "Conservatively check W^X thread state at possible safepoint" \
-          "or handshake"))
+  develop(bool, TraceWXHealing, false,                                  \
+          "track occurrences of W^X mode healing")                      \
+  develop(bool, UseOldWX, false,                                        \
+          "Choose old W^X implementation.")                             \
+  product(bool, StressWXHealing, false, DIAGNOSTIC,                     \
+          "Stress W xor X healing on MacOS")
+
+#else
+#define RUNTIME_OS_FLAGS(develop,                                       \
+                         develop_pd,                                    \
+                         product,                                       \
+                         product_pd,                                    \
+                         range,                                         \
+                         constraint)
+#endif
 
 // end of RUNTIME_OS_FLAGS
 

src/hotspot/os/bsd/os_bsd.cpp

@@ -841,6 +841,7 @@ jlong os::javaTimeNanos() {
   // We might also condition (c) on the magnitude of the delta between obsv and now.
   // Avoiding excessive CAS operations to hot RW locations is critical.
   // See https://blogs.oracle.com/dave/entry/cas_and_cache_trivia_invalidate
+  // https://web.archive.org/web/20131214182431/https://blogs.oracle.com/dave/entry/cas_and_cache_trivia_invalidate
   return (prev == obsv) ? now : obsv;
 }
 

src/hotspot/os_cpu/bsd_aarch64/os_bsd_aarch64.cpp

@@ -54,8 +54,11 @@
 #include "signals_posix.hpp"
 #include "utilities/align.hpp"
 #include "utilities/debug.hpp"
+#include "utilities/decoder.hpp"
 #include "utilities/events.hpp"
+#include "utilities/nativeStackPrinter.hpp"
 #include "utilities/vmError.hpp"
+#include "compiler/disassembler.hpp"
 
 // put OS-includes here
 # include <sys/types.h>
@@ -85,6 +88,8 @@
 #define SPELL_REG_SP "sp"
 
 #ifdef __APPLE__
+WXMode DefaultWXWriteMode;
+
 // see darwin-xnu/osfmk/mach/arm/_structs.h
 
 // 10.5 UNIX03 member name prefixes
@@ -233,19 +238,56 @@ NOINLINE frame os::current_frame() {
 
 bool PosixSignals::pd_hotspot_signal_handler(int sig, siginfo_t* info,
                                              ucontext_t* uc, JavaThread* thread) {
-  // Enable WXWrite: this function is called by the signal handler at arbitrary
-  // point of execution.
-  ThreadWXEnable wx(WXWrite, thread);
-
   // decide if this trap can be handled by a stub
   address stub = nullptr;
-
-  address pc          = nullptr;
+  address pc   = nullptr;
 
   //%note os_trap_1
   if (info != nullptr && uc != nullptr && thread != nullptr) {
     pc = (address) os::Posix::ucontext_get_pc(uc);
 
+#ifdef MACOS_AARCH64
+    // If we got a SIGBUS because we tried to write into the code
+    // cache, try enabling WXWrite mode.
+    if (sig == SIGBUS
+        && pc != info->si_addr
+        && CodeCache::contains(info->si_addr)
+        && os::address_is_in_vm(pc)) {
+      WXMode *entry_mode = thread->_cur_wx_mode;
+      if (entry_mode != nullptr && *entry_mode == WXArmedForWrite) {
+        if (TraceWXHealing) {
+          static const char *mode_names[3] = {"WXWrite", "WXExec", "WXArmedForWrite"};
+          tty->print("Healing WXMode %s at %p to WXWrite",
+                        mode_names[*entry_mode], entry_mode);
+          char name[128];
+          int offset = 0;
+          if (os::dll_address_to_function_name(pc, name, sizeof name, &offset)) {
+            tty->print_cr("  (%s+0x%x)", name, offset);
+          } else {
+            tty->cr();
+          }
+          if (Verbose) {
+            char buf[O_BUFLEN];
+            NativeStackPrinter nsp(thread);
+            nsp.print_stack(tty, buf, sizeof(buf), pc,
+                            true /* print_source_info */, -1 /* max stack */);
+          }
+        }
+#ifndef PRODUCT
+        guarantee(StressWXHealing,
+                  "We should not reach here unless StressWXHealing");
+#endif
+        *(thread->_cur_wx_mode) = WXWrite;
+        return thread->wx_enable_write();
+      }
+    }
+
+    // There may be cases where code after this point that we call
+    // from the signal handler changes WX state, so we protect against
+    // that by saving and restoring the state.
+    ThreadWXEnable wx(thread->get_wx_state(), thread);
+#endif
+
     // Handle ALL stack overflow variations here
     if (sig == SIGSEGV || sig == SIGBUS) {
       address addr = (address) info->si_addr;
@@ -515,11 +557,42 @@ int os::extra_bang_size_in_bytes() {
   return 0;
 }
 
-#ifdef __APPLE__
+#ifdef MACOS_AARCH64
+THREAD_LOCAL bool os::_jit_exec_enabled;
+
+// This is a wrapper around the standard library function
+// pthread_jit_write_protect_np(3). We keep track of the state of
+// per-thread write protection on the MAP_JIT region in the
+// thread-local variable os::_jit_exec_enabled
 void os::current_thread_enable_wx(WXMode mode) {
-  pthread_jit_write_protect_np(mode == WXExec);
+  bool exec_enabled = mode != WXWrite;
+  if (exec_enabled != _jit_exec_enabled NOT_PRODUCT( || DefaultWXWriteMode == WXWrite)) {
+    permit_forbidden_function::pthread_jit_write_protect_np(exec_enabled);
+    _jit_exec_enabled = exec_enabled;
+  }
 }
-#endif
+
+// If the current thread is in the WX state WXArmedForWrite, change
+// the state to WXWrite.
+bool Thread::wx_enable_write() {
+  if (_wx_state == WXArmedForWrite) {
+    _wx_state = WXWrite;
+    os::current_thread_enable_wx(WXWrite);
+    return true;
+  } else {
+    return false;
+  }
+}
+
+// A wrapper around wx_enable_write() for when the current thread is
+// not known.
+void os::thread_wx_enable_write_impl() {
+  if (!StressWXHealing) {
+    Thread::current()->wx_enable_write();
+  }
+}
+
+#endif // MACOS_AARCH64
 
 static inline void atomic_copy64(const volatile void *src, volatile void *dst) {
   *(jlong *) dst = *(const jlong *) src;