Implement signal-handling for MMU epochs and an epoch-ending method on Store.

epoch_mmu_trap_via_signal_handler() now passes. It shows the newly added nop version of the asm trampoline works.

Author: erikrose

crates/wasmtime/src/runtime/code_memory.rs

@@ -33,6 +33,7 @@ pub struct CodeMemory {
     text: Range<usize>,
     unwind: Range<usize>,
     trap_data: Range<usize>,
+    epoch_check_data: Range<usize>,
     wasm_data: Range<usize>,
     address_map_data: Range<usize>,
     stack_map_data: Range<usize>,
@@ -127,6 +128,7 @@ impl CodeMemory {
         #[cfg(feature = "debug-builtins")]
         let mut has_native_debug_info = false;
         let mut trap_data = 0..0;
+        let mut epoch_check_data = 0..0;
         let mut exception_data = 0..0;
         let mut frame_tables_data = 0..0;
         let mut wasm_data = 0..0;
@@ -178,6 +180,7 @@ impl CodeMemory {
                 obj::ELF_WASMTIME_ADDRMAP => address_map_data = range,
                 obj::ELF_WASMTIME_STACK_MAP => stack_map_data = range,
                 obj::ELF_WASMTIME_TRAPS => trap_data = range,
+                obj::ELF_WASMTIME_EPOCH_CHECKS => epoch_check_data = range,
                 obj::ELF_WASMTIME_EXCEPTIONS => exception_data = range,
                 obj::ELF_WASMTIME_FRAMES => frame_tables_data = range,
                 obj::ELF_NAME_DATA => func_name_data = range,
@@ -222,6 +225,7 @@ impl CodeMemory {
             text,
             unwind,
             trap_data,
+            epoch_check_data,
             address_map_data,
             stack_map_data,
             exception_data,

crates/wasmtime/src/runtime/store.rs

@@ -1178,6 +1178,12 @@ impl<T> Store<T> {
         self.inner.epoch_deadline_callback(Box::new(callback));
     }
 
+    /// Iff epoch_interruption_via_mmu() is on, end the current epoch, causing
+    /// Wasm code belonging to this store to yield to a different task.
+    pub fn end_mmu_epoch(&self) {
+        self.inner.vm_store_context().protect_interrupt_page();
+    }
+
     /// Set an exception as the currently pending exception, and
     /// return an error that propagates the throw.
     ///

crates/wasmtime/src/runtime/vm/sys/unix/signals.rs

@@ -1,7 +1,9 @@
 //! Trap handling on Unix based on POSIX signals.
 
 use crate::prelude::*;
+use crate::runtime::module::lookup_code;
 use crate::runtime::vm::traphandlers::{TrapRegisters, TrapTest, tls};
+use core::arch::naked_asm;
 use std::cell::RefCell;
 use std::io;
 use std::mem;
@@ -18,6 +20,10 @@ static mut PREV_SIGBUS: libc::sigaction = UNINIT_SIGACTION;
 static mut PREV_SIGILL: libc::sigaction = UNINIT_SIGACTION;
 static mut PREV_SIGFPE: libc::sigaction = UNINIT_SIGACTION;
 
+// From signal.h. Not yet exposed in libc. Value is valid for Linux and Mac; not
+// sure about elsewhere.
+#[cfg(all(target_os = "linux"))]
+const SEGV_ACCERR: libc::c_int = 2;
 pub struct TrapHandler;
 
 impl TrapHandler {
@@ -128,6 +134,26 @@ now.
     }
 }
 
+/// Switches tasks in response to a signal thrown under MMU-based epoch
+/// interruption.
+///
+/// Saves register state, makes a host call to switch tasks, restores state, and
+/// returns.
+#[unsafe(naked)]
+unsafe extern "C" fn task_switch_trampoline(_vmctx: usize) {
+    naked_asm!(
+        "
+        // Save regs.
+        // Call hostcall to do task switch, passing in vmctx.
+        // Restore regs (including r10).
+
+        // Resume right after the load instruction that triggered the signal
+        // handler.
+        jmp r10
+        "
+    );
+}
+
 unsafe extern "C" fn trap_handler(
     signum: libc::c_int,
     siginfo: *mut libc::siginfo_t,
@@ -148,6 +174,37 @@ unsafe extern "C" fn trap_handler(
             None => return false,
         };
 
+        // Check for segfaults meant as cues to end an epoch.
+        #[cfg(all(target_os = "linux", target_arch = "x86_64"))]
+        if signum == libc::SIGSEGV && unsafe { (*siginfo).si_code } == SEGV_ACCERR {
+            // Compare it with the offsets of epoch-check instructions as stored
+            // in the object file.
+            let ucontext = unsafe { &mut *(context as *mut libc::ucontext_t) };
+            let pc = ucontext.uc_mcontext.gregs[libc::REG_RIP as usize] as usize;
+            // Now things get expensive: we call lookup_code(), which takes a global lock.
+            if let Some((code_memory, offset_within_code)) = lookup_code(pc) {
+                // We're within a 'target_arch = "x86_64"', so we can just treat
+                // the stored little-endians as native u32s.
+                if let Some(return_address) = code_memory.return_address_for_epoch_check(
+                    offset_within_code
+                        .try_into()
+                        .expect("epoch-check location should fit in 32 bits"),
+                ) {
+                    // It is an epoch check. Arrange to resume at asm trampoline
+                    // after signal handler exits:
+                    ucontext.uc_mcontext.gregs[libc::REG_RIP as usize] =
+                        task_switch_trampoline as *const () as i64;
+                    // Put original resumption address in R10:
+                    ucontext.uc_mcontext.gregs[libc::REG_R10 as usize] = return_address as i64;
+                    // Trampoline can expect the vmctx in RDI.
+
+                    return true;
+                }
+            }
+            // Else it is an ordinary trap or the epochchecks section is somehow
+            // missing from the binary; continue on.
+        }
+
         // If we hit an exception while handling a previous trap, that's
         // quite bad, so bail out and let the system handle this
         // recursive segfault.

tests/all/epoch_mmu.rs

@@ -1,7 +1,7 @@
 #![cfg(not(miri))]
 
 use object::{Object, ObjectSection};
-use wasmtime::{Config, Engine};
+use wasmtime::{Config, Engine, Instance, Module, Store};
 use wasmtime_environ::obj::ELF_WASMTIME_EPOCH_CHECKS;
 
 /// Asserts that each epoch-check offset encoded into the binary points to the
@@ -49,3 +49,40 @@ fn epoch_check_offsets() {
         "There should be 2 epoch checks (function prologue & loop backedge). The offset after the prologue's dead load should be 15, and the one after the loop's backedge should be 18."
     );
 }
+
+/// Runs a wasm function that loops forever with MMU-based epoch interruption
+/// enabled. Incrementing the epoch past the deadline triggers the signal
+/// handler (`trap_handler`) which converts the SIGSEGV into a trap.
+#[test]
+fn epoch_mmu_trap_via_signal_handler() {
+    let mut config = Config::new();
+    config.epoch_interruption_via_mmu(true);
+    let engine = Engine::new(&config).unwrap();
+    let module = Module::new(
+        &engine,
+        r#"(module
+             (memory 0)
+             (func (export "answer") (result i32)
+                i32.const 42
+             )
+           )"#,
+    )
+    .unwrap();
+
+    // Trap as soon as the first epoch check is encountered, in the function
+    // prologue. Recall that MMU-based epochs don't operate based on a numeric
+    // deadline but on an external entity protecting the memory page, typically
+    // on a timer.
+    let mut store = Store::new(&engine, ());
+    store.epoch_deadline_trap(); // Allegedly the default.
+    // Protect that page:
+    store.end_mmu_epoch();
+
+    let instance = Instance::new(&mut store, &module, &[]).unwrap();
+    let func = instance
+        .get_typed_func::<(), i32>(&mut store, "answer")
+        .unwrap();
+
+    let result = func.call(&mut store, ()).unwrap();
+    assert_eq!(result, 42);
+}