Backtrace WebAssembly function JIT frames (#759)

* Create backtrace

* Extend unwind information with FDE data.

* Expose backtrace via API/Trap

* wasmtime_call returns not-str

* Return Arc<JITFrameTag>

* rename frame -> function

* Fix windows crashes and unwrap UNWIND_HISTORY_TABLE

* mmaps -> entries

* pass a backtrace in ActionOutcome

* add test_trap_stack_overflow

* Update cranelift version.

crates/runtime/src/backtrace.rs

@@ -0,0 +1,148 @@
+//! Backtrace object and utilities.
+
+use crate::jit_function_registry;
+use std::sync::Arc;
+
+/// Information about backtrace frame.
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct BacktraceFrame {
+    pc: usize,
+}
+
+impl Default for BacktraceFrame {
+    fn default() -> Self {
+        Self { pc: 0 }
+    }
+}
+
+impl BacktraceFrame {
+    /// Current PC or IP pointer for the frame.
+    pub fn pc(&self) -> usize {
+        self.pc
+    }
+    /// Additinal frame information.
+    pub fn tag(&self) -> Option<Arc<jit_function_registry::JITFunctionTag>> {
+        jit_function_registry::find(self.pc)
+    }
+}
+
+const BACKTRACE_LIMIT: usize = 32;
+
+/// Backtrace during WebAssembly trap.
+#[derive(Copy, Clone, PartialEq, Eq)]
+pub struct Backtrace {
+    len: usize,
+    frames: [BacktraceFrame; BACKTRACE_LIMIT],
+}
+
+impl Backtrace {
+    fn new() -> Self {
+        Self {
+            len: 0,
+            frames: [Default::default(); BACKTRACE_LIMIT],
+        }
+    }
+    fn full(&self) -> bool {
+        self.len >= BACKTRACE_LIMIT
+    }
+    fn push(&mut self, frame: BacktraceFrame) {
+        assert!(self.len < BACKTRACE_LIMIT);
+        self.frames[self.len] = frame;
+        self.len += 1;
+    }
+    /// Amount of the backtrace frames.
+    pub fn len(&self) -> usize {
+        self.len
+    }
+}
+
+impl std::ops::Index<usize> for Backtrace {
+    type Output = BacktraceFrame;
+    fn index(&self, index: usize) -> &Self::Output {
+        assert!(index < self.len);
+        &self.frames[index]
+    }
+}
+
+impl std::fmt::Debug for Backtrace {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        writeln!(f, "Backtrace![")?;
+        for i in 0..self.len() {
+            let frame = &self.frames[i];
+            writeln!(f, "  {:x}: {:?}", frame.pc(), frame.tag())?;
+        }
+        write!(f, "]")?;
+        Ok(())
+    }
+}
+
+#[cfg(not(all(target_os = "windows", target_arch = "x86_64")))]
+fn capture_stack<F>(mut f: F)
+where
+    F: FnMut(usize) -> bool,
+{
+    use backtrace::trace;
+    trace(|frame| {
+        let pc = frame.ip() as usize;
+        f(pc)
+    });
+}
+
+#[cfg(all(target_os = "windows", target_arch = "x86_64"))]
+fn capture_stack<F>(mut f: F)
+where
+    F: FnMut(usize) -> bool,
+{
+    use std::mem::MaybeUninit;
+    use std::ptr;
+    use winapi::um::winnt::{
+        RtlCaptureContext, RtlLookupFunctionEntry, RtlVirtualUnwind, CONTEXT, UNW_FLAG_NHANDLER,
+    };
+
+    #[repr(C, align(16))]
+    struct WrappedContext(CONTEXT);
+
+    unsafe {
+        let mut ctx = WrappedContext(MaybeUninit::uninit().assume_init());
+        RtlCaptureContext(&mut ctx.0);
+        let mut unwind_history_table = MaybeUninit::zeroed().assume_init();
+        while ctx.0.Rip != 0 {
+            let cont = f(ctx.0.Rip as usize);
+            if !cont {
+                break;
+            }
+
+            let mut image_base: u64 = 0;
+            let mut handler_data: *mut core::ffi::c_void = ptr::null_mut();
+            let mut establisher_frame: u64 = 0;
+            let rf = RtlLookupFunctionEntry(ctx.0.Rip, &mut image_base, &mut unwind_history_table);
+            if rf.is_null() {
+                ctx.0.Rip = ptr::read(ctx.0.Rsp as *const u64);
+                ctx.0.Rsp += 8;
+            } else {
+                RtlVirtualUnwind(
+                    UNW_FLAG_NHANDLER,
+                    image_base,
+                    ctx.0.Rip,
+                    rf,
+                    &mut ctx.0,
+                    &mut handler_data,
+                    &mut establisher_frame,
+                    ptr::null_mut(),
+                );
+            }
+        }
+    }
+}
+
+/// Returns current backtrace. Only registered wasmtime functions will be listed.
+pub fn get_backtrace() -> Backtrace {
+    let mut frames = Backtrace::new();
+    capture_stack(|pc| {
+        if let Some(_) = jit_function_registry::find(pc) {
+            frames.push(BacktraceFrame { pc });
+        }
+        !frames.full()
+    });
+    frames
+}

crates/runtime/src/instance.rs

@@ -9,7 +9,7 @@ use crate::memory::LinearMemory;
 use crate::mmap::Mmap;
 use crate::signalhandlers::{wasmtime_init_eager, wasmtime_init_finish};
 use crate::table::Table;
-use crate::traphandlers::wasmtime_call;
+use crate::traphandlers::{wasmtime_call, TrapMessageAndStack};
 use crate::vmcontext::{
     VMBuiltinFunctionsArray, VMCallerCheckedAnyfunc, VMContext, VMFunctionBody, VMFunctionImport,
     VMGlobalDefinition, VMGlobalImport, VMMemoryDefinition, VMMemoryImport, VMSharedSignatureIndex,
@@ -568,7 +568,7 @@ impl Instance {
 
         // Make the call.
         unsafe { wasmtime_call(callee_vmctx, callee_address) }
-            .map_err(InstantiationError::StartTrap)
+            .map_err(|TrapMessageAndStack(msg, _)| InstantiationError::StartTrap(msg))
     }
 
     /// Invoke the WebAssembly start function of the instance, if one is present.

crates/runtime/src/jit_function_registry.rs

@@ -0,0 +1,83 @@
+#![allow(missing_docs)]
+
+use lazy_static::lazy_static;
+use std::collections::BTreeMap;
+use std::sync::{Arc, RwLock};
+
+lazy_static! {
+    static ref REGISTRY: RwLock<JITFunctionRegistry> = RwLock::new(JITFunctionRegistry::default());
+}
+
+#[derive(Clone)]
+pub struct JITFunctionTag {
+    pub module_id: Option<String>,
+    pub func_index: usize,
+}
+
+impl std::fmt::Debug for JITFunctionTag {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if let Some(ref module_id) = self.module_id {
+            write!(f, "{}", module_id)?;
+        } else {
+            write!(f, "(module)")?;
+        }
+        write!(f, ":{}", self.func_index)
+    }
+}
+
+struct JITFunctionRegistry {
+    ranges: BTreeMap<usize, (usize, Arc<JITFunctionTag>)>,
+}
+
+impl Default for JITFunctionRegistry {
+    fn default() -> Self {
+        Self {
+            ranges: Default::default(),
+        }
+    }
+}
+
+impl JITFunctionRegistry {
+    fn register(&mut self, fn_start: usize, fn_end: usize, tag: JITFunctionTag) {
+        self.ranges.insert(fn_end, (fn_start, Arc::new(tag)));
+    }
+
+    fn unregister(&mut self, fn_end: usize) {
+        self.ranges.remove(&fn_end);
+    }
+
+    fn find(&self, pc: usize) -> Option<&Arc<JITFunctionTag>> {
+        self.ranges
+            .range(pc..)
+            .next()
+            .and_then(|(end, (start, s))| {
+                if *start <= pc && pc < *end {
+                    Some(s)
+                } else {
+                    None
+                }
+            })
+    }
+}
+
+pub fn register(fn_start: usize, fn_end: usize, tag: JITFunctionTag) {
+    REGISTRY
+        .write()
+        .expect("jit function registry lock got poisoned")
+        .register(fn_start, fn_end, tag);
+}
+
+pub fn unregister(_fn_start: usize, fn_end: usize) {
+    REGISTRY
+        .write()
+        .expect("jit function registry lock got poisoned")
+        .unregister(fn_end);
+}
+
+pub fn find(pc: usize) -> Option<Arc<JITFunctionTag>> {
+    REGISTRY
+        .read()
+        .expect("jit function registry lock got poisoned")
+        .find(pc)
+        .cloned()
+}

crates/runtime/src/lib.rs

@@ -21,6 +21,7 @@
     )
 )]
 
+mod backtrace;
 mod export;
 mod imports;
 mod instance;
@@ -34,8 +35,10 @@ mod trap_registry;
 mod traphandlers;
 mod vmcontext;
 
+pub mod jit_function_registry;
 pub mod libcalls;
 
+pub use crate::backtrace::{get_backtrace, Backtrace, BacktraceFrame};
 pub use crate::export::Export;
 pub use crate::imports::Imports;
 pub use crate::instance::{InstanceHandle, InstantiationError, LinkError};
@@ -44,7 +47,7 @@ pub use crate::mmap::Mmap;
 pub use crate::sig_registry::SignatureRegistry;
 pub use crate::signalhandlers::{wasmtime_init_eager, wasmtime_init_finish};
 pub use crate::trap_registry::{get_mut_trap_registry, get_trap_registry, TrapRegistrationGuard};
-pub use crate::traphandlers::{wasmtime_call, wasmtime_call_trampoline};
+pub use crate::traphandlers::{wasmtime_call, wasmtime_call_trampoline, TrapMessageAndStack};
 pub use crate::vmcontext::{
     VMCallerCheckedAnyfunc, VMContext, VMFunctionBody, VMFunctionImport, VMGlobalDefinition,
     VMGlobalImport, VMInvokeArgument, VMMemoryDefinition, VMMemoryImport, VMSharedSignatureIndex,

crates/runtime/src/traphandlers.rs

@@ -1,6 +1,7 @@
 //! WebAssembly trap handling, which is built on top of the lower-level
 //! signalhandling mechanisms.
 
+use crate::backtrace::{get_backtrace, Backtrace};
 use crate::trap_registry::get_trap_registry;
 use crate::trap_registry::TrapDescription;
 use crate::vmcontext::{VMContext, VMFunctionBody};
@@ -18,7 +19,7 @@ extern "C" {
 }
 
 thread_local! {
-    static RECORDED_TRAP: Cell<Option<TrapDescription>> = Cell::new(None);
+    static RECORDED_TRAP: Cell<Option<(TrapDescription, Backtrace)>> = Cell::new(None);
     static JMP_BUF: Cell<*const u8> = Cell::new(ptr::null());
     static RESET_GUARD_PAGE: Cell<bool> = Cell::new(false);
 }
@@ -50,6 +51,8 @@ pub extern "C" fn RecordTrap(pc: *const u8, reset_guard_page: bool) {
             trap_code: ir::TrapCode::StackOverflow,
         });
 
+    let trap_backtrace = get_backtrace();
+
     if reset_guard_page {
         RESET_GUARD_PAGE.with(|v| v.set(true));
     }
@@ -60,7 +63,7 @@ pub extern "C" fn RecordTrap(pc: *const u8, reset_guard_page: bool) {
             None,
             "Only one trap per thread can be recorded at a moment!"
         );
-        data.set(Some(trap_desc))
+        data.set(Some((trap_desc, trap_backtrace)))
     });
 }
 
@@ -108,15 +111,22 @@ fn reset_guard_page() {
 #[cfg(not(target_os = "windows"))]
 fn reset_guard_page() {}
 
-fn trap_message() -> String {
+/// Stores trace message with backtrace.
+#[derive(Debug)]
+pub struct TrapMessageAndStack(pub String, pub Backtrace);
+
+fn trap_message_and_stack() -> TrapMessageAndStack {
     let trap_desc = RECORDED_TRAP
         .with(|data| data.replace(None))
         .expect("trap_message must be called after trap occurred");
 
-    format!(
-        "wasm trap: {}, source location: {}",
-        trap_code_to_expected_string(trap_desc.trap_code),
-        trap_desc.source_loc,
+    TrapMessageAndStack(
+        format!(
+            "wasm trap: {}, source location: {}",
+            trap_code_to_expected_string(trap_desc.0.trap_code),
+            trap_desc.0.source_loc,
+        ),
+        trap_desc.1,
     )
 }
 
@@ -146,9 +156,9 @@ pub unsafe extern "C" fn wasmtime_call_trampoline(
     vmctx: *mut VMContext,
     callee: *const VMFunctionBody,
     values_vec: *mut u8,
-) -> Result<(), String> {
+) -> Result<(), TrapMessageAndStack> {
     if WasmtimeCallTrampoline(vmctx as *mut u8, callee, values_vec) == 0 {
-        Err(trap_message())
+        Err(trap_message_and_stack())
     } else {
         Ok(())
     }
@@ -160,9 +170,9 @@ pub unsafe extern "C" fn wasmtime_call_trampoline(
 pub unsafe extern "C" fn wasmtime_call(
     vmctx: *mut VMContext,
     callee: *const VMFunctionBody,
-) -> Result<(), String> {
+) -> Result<(), TrapMessageAndStack> {
     if WasmtimeCall(vmctx as *mut u8, callee) == 0 {
-        Err(trap_message())
+        Err(trap_message_and_stack())
     } else {
         Ok(())
     }