Refactor tier switching code

lib/llvm-backend/src/backend.rs

@@ -287,11 +287,13 @@ impl LLVMBackend {
             )
         };
 
+        /*
+
         static SIGNAL_HANDLER_INSTALLED: Once = Once::new();
 
         SIGNAL_HANDLER_INSTALLED.call_once(|| unsafe {
             crate::platform::install_signal_handler();
-        });
+        });*/
 
         if res != LLVMResult::OK {
             panic!("failed to load object")

lib/runtime-core/Cargo.toml

@@ -57,3 +57,4 @@ trace = ["debug"]
 # backend flags used in conditional compilation of Backend::variants
 "backend-singlepass" = []
 "backend-llvm" = []
+managed = []

lib/runtime-core/src/fault.rs

@@ -9,6 +9,7 @@ mod raw {
 }
 
 use crate::codegen::{BreakpointInfo, BreakpointMap};
+use crate::state::CodeVersion;
 use crate::state::x64::{build_instance_image, read_stack, X64Register, GPR, XMM};
 use crate::vm;
 use libc::{mmap, mprotect, siginfo_t, MAP_ANON, MAP_PRIVATE, PROT_NONE, PROT_READ, PROT_WRITE};
@@ -17,7 +18,7 @@ use nix::sys::signal::{
     SIGSEGV, SIGTRAP,
 };
 use std::any::Any;
-use std::cell::{Cell, UnsafeCell};
+use std::cell::{Cell, UnsafeCell, RefCell};
 use std::ffi::c_void;
 use std::process;
 use std::sync::atomic::{AtomicBool, Ordering};
@@ -41,6 +42,7 @@ struct UnwindInfo {
 thread_local! {
     static UNWIND: UnsafeCell<Option<UnwindInfo>> = UnsafeCell::new(None);
     static CURRENT_CTX: UnsafeCell<*mut vm::Ctx> = UnsafeCell::new(::std::ptr::null_mut());
+    static CURRENT_CODE_VERSIONS: RefCell<Vec<CodeVersion>> = RefCell::new(vec![]);
     static WAS_SIGINT_TRIGGERED: Cell<bool> = Cell::new(false);
 }
 
@@ -83,6 +85,14 @@ pub unsafe fn with_ctx<R, F: FnOnce() -> R>(ctx: *mut vm::Ctx, cb: F) -> R {
     ret
 }
 
+pub fn push_code_version(version: CodeVersion) {
+    CURRENT_CODE_VERSIONS.with(|x| x.borrow_mut().push(version));
+}
+
+pub fn pop_code_version() -> Option<CodeVersion> {
+    CURRENT_CODE_VERSIONS.with(|x| x.borrow_mut().pop())
+}
+
 pub unsafe fn get_wasm_interrupt_signal_mem() -> *mut u8 {
     INTERRUPT_SIGNAL_MEM.0
 }
@@ -242,18 +252,15 @@ extern "C" fn signal_trap_handler(
             let ctx: &mut vm::Ctx = &mut **CURRENT_CTX.with(|x| x.get());
             let rsp = fault.known_registers[X64Register::GPR(GPR::RSP).to_index().0].unwrap();
 
-            let msm = (*ctx.module)
+            let es_image = CURRENT_CODE_VERSIONS.with(|versions| {
-                .runnable_module
+                let versions = versions.borrow();
-                .get_module_state_map()
+                read_stack(
-                .unwrap();
+                    || versions.iter(),
-            let code_base = (*ctx.module).runnable_module.get_code().unwrap().as_ptr() as usize;
-            let es_image = read_stack(
-                &msm,
-                code_base,
                     rsp as usize as *const u64,
                     fault.known_registers,
                     Some(fault.ip as usize as u64),
-            );
+                )
+            });
 
             if is_suspend_signal {
                 let image = build_instance_image(ctx, es_image);

lib/runtime-core/src/lib.rs

@@ -47,6 +47,8 @@ pub use trampoline_x64 as trampoline;
 #[cfg(all(unix, target_arch = "x86_64"))]
 pub mod fault;
 pub mod state;
+#[cfg(feature = "managed")]
+pub mod tiering;
 
 use self::error::CompileResult;
 #[doc(inline)]

lib/runtime-core/src/state.rs

@@ -102,6 +102,12 @@ pub struct InstanceImage {
     pub execution_state: ExecutionStateImage,
 }
 
+#[derive(Debug, Clone)]
+pub struct CodeVersion {
+    pub msm: ModuleStateMap,
+    pub base: usize,
+}
+
 impl ModuleStateMap {
     fn lookup_ip<F: FnOnce(&FunctionStateMap) -> &BTreeMap<usize, OffsetInfo>>(
         &self,
@@ -739,9 +745,8 @@ pub mod x64 {
     }
 
     #[warn(unused_variables)]
-    pub unsafe fn read_stack(
+    pub unsafe fn read_stack<'a, I: Iterator<Item = &'a CodeVersion>, F: Fn() -> I + 'a>(
-        msm: &ModuleStateMap,
+        versions: F,
-        code_base: usize,
         mut stack: *const u64,
         initially_known_registers: [Option<u64>; 24],
         mut initial_address: Option<u64>,
@@ -755,13 +760,27 @@ pub mod x64 {
                 stack = stack.offset(1);
                 x
             });
-            let (fsm, state) = match msm
+
-                .lookup_call_ip(ret_addr as usize, code_base)
+            let mut fsm_state: Option<(&FunctionStateMap, MachineState)> = None;
-                .or_else(|| msm.lookup_trappable_ip(ret_addr as usize, code_base))
+
-                .or_else(|| msm.lookup_loop_ip(ret_addr as usize, code_base))
+            for version in versions() {
+                match version.msm
+                    .lookup_call_ip(ret_addr as usize, version.base)
+                    .or_else(|| version.msm.lookup_trappable_ip(ret_addr as usize, version.base))
+                    .or_else(|| version.msm.lookup_loop_ip(ret_addr as usize, version.base))
                 {
-                Some(x) => x,
+                    Some(x) => {
-                _ => return ExecutionStateImage { frames: results },
+                        fsm_state = Some(x);
+                        break;
+                    },
+                    None => {},
+                };
+            }
+
+            let (fsm, state) = if let Some(x) = fsm_state {
+                x
+            } else {
+                return ExecutionStateImage { frames: results };
             };
 
             let mut wasm_stack: Vec<Option<u64>> = state

lib/runtime-core/src/tiering.rs

@@ -0,0 +1,236 @@
+use crate::backend::{Compiler, CompilerConfig};
+use crate::import::ImportObject;
+use crate::fault::{
+    catch_unsafe_unwind, ensure_sighandler, with_ctx, push_code_version, pop_code_version,
+};
+use crate::state::{
+    x64::invoke_call_return_on_stack, InstanceImage, CodeVersion
+};
+use crate::vm::Ctx;
+use crate::compile_with_config;
+use crate::instance::Instance;
+use crate::fault::{set_wasm_interrupt_on_ctx, was_sigint_triggered_fault};
+use crate::module::{Module, ModuleInfo};
+
+use std::sync::{Arc, Mutex};
+use std::cell::Cell;
+
+struct Defer<F: FnOnce()>(Option<F>);
+impl<F: FnOnce()> Drop for Defer<F> {
+    fn drop(&mut self) {
+        if let Some(f) = self.0.take() {
+            f();
+        }
+    }
+}
+
+pub enum ShellExitOperation {
+    ContinueWith(InstanceImage),
+}
+
+pub struct InteractiveShellContext {
+    pub image: Option<InstanceImage>,
+}
+
+struct OptimizationState {
+    outcome: Mutex<Option<OptimizationOutcome>>,
+}
+
+struct OptimizationOutcome {
+    module: Module,
+}
+
+#[repr(transparent)]
+struct CtxWrapper(*mut Ctx);
+unsafe impl Send for CtxWrapper {}
+unsafe impl Sync for CtxWrapper {}
+
+unsafe fn do_optimize(
+    binary: &[u8],
+    compiler: Box<dyn Compiler>,
+    ctx: &Mutex<CtxWrapper>,
+    state: &OptimizationState,
+) {
+    let module = match compile_with_config(
+        &binary[..],
+        &*compiler,
+        CompilerConfig {
+            symbol_map: None,
+            track_state: true,
+            ..Default::default()
+        },
+    ) {
+        Ok(x) => x,
+        Err(_) => return,
+    };
+
+    let ctx_inner = ctx.lock().unwrap();
+    if !ctx_inner.0.is_null() {
+        *state.outcome.lock().unwrap() = Some(OptimizationOutcome { module });
+        set_wasm_interrupt_on_ctx(ctx_inner.0);
+    }
+}
+
+pub fn run_tiering<F: Fn(InteractiveShellContext) -> ShellExitOperation>(
+    module_info: &ModuleInfo,
+    wasm_binary: &[u8],
+    mut resume_image: Option<InstanceImage>,
+    import_object: &ImportObject,
+    start_raw: extern "C" fn(&mut Ctx),
+    baseline: &mut Instance,
+    optimized_backends: Vec<Box<Fn() -> Box<dyn Compiler> + Send>>,
+    interactive_shell: F,
+) -> Result<(), String> {
+    unsafe {
+        ensure_sighandler();
+
+        let ctx_box =
+            Arc::new(Mutex::new(CtxWrapper(baseline.context_mut() as *mut _)));
+        // Ensure that the ctx pointer's lifetime is not longer than Instance's.
+        let _deferred_ctx_box_cleanup: Defer<_> = {
+            let ctx_box = ctx_box.clone();
+            Defer(Some(move || {
+                ctx_box.lock().unwrap().0 = ::std::ptr::null_mut();
+            }))
+        };
+        let opt_state = Arc::new(OptimizationState {
+            outcome: Mutex::new(None),
+        });
+
+        {
+            let wasm_binary = wasm_binary.to_vec();
+            let ctx_box = ctx_box.clone();
+            let opt_state = opt_state.clone();
+            ::std::thread::spawn(move || {
+                for backend in optimized_backends {
+                    if !ctx_box.lock().unwrap().0.is_null() {
+                        do_optimize(
+                            &wasm_binary,
+                            backend(),
+                            &ctx_box,
+                            &opt_state,
+                        );
+                    }
+                }
+            });
+        }
+
+        let mut optimized_instances: Vec<Instance> = vec![];
+
+        push_code_version(CodeVersion {
+            msm: baseline.module.runnable_module.get_module_state_map().unwrap(),
+            base: baseline.module.runnable_module.get_code().unwrap().as_ptr() as usize,
+        });
+        let n_versions: Cell<usize> = Cell::new(1);
+
+        let _deferred_pop_versions = Defer(Some(|| {
+            for _ in 0..n_versions.get() {
+                pop_code_version().unwrap();
+            }
+        }));
+
+        loop {
+            let new_optimized: Option<&mut Instance> = {
+                let mut outcome = opt_state.outcome.lock().unwrap();
+                if let Some(x) = outcome.take() {
+                    let instance =
+                        x.module.instantiate(&import_object).map_err(|e| {
+                            format!("Can't instantiate module: {:?}", e)
+                        })?;
+                    // Keep the optimized code alive.
+                    optimized_instances.push(instance);
+                    optimized_instances.last_mut()
+                } else {
+                    None
+                }
+            };
+            if let Some(optimized) = new_optimized {
+                let base = module_info.imported_functions.len();
+                let code_ptr = optimized
+                    .module
+                    .runnable_module
+                    .get_code()
+                    .unwrap()
+                    .as_ptr()
+                    as usize;
+                let target_addresses: Vec<usize> = optimized
+                    .module
+                    .runnable_module
+                    .get_local_function_offsets()
+                    .unwrap()
+                    .into_iter()
+                    .map(|x| code_ptr + x)
+                    .collect();
+                assert_eq!(
+                    target_addresses.len(),
+                    module_info.func_assoc.len() - base
+                );
+                for i in base..module_info.func_assoc.len() {
+                    baseline
+                        .module
+                        .runnable_module
+                        .patch_local_function(i - base, target_addresses[i - base]);
+                }
+
+                push_code_version(CodeVersion {
+                    msm: optimized.module.runnable_module.get_module_state_map().unwrap(),
+                    base: optimized.module.runnable_module.get_code().unwrap().as_ptr() as usize,
+                });
+                n_versions.set(n_versions.get() + 1);
+
+                eprintln!("Patched");
+            }
+            // TODO: Fix this for optimized version.
+            let breakpoints = baseline.module.runnable_module.get_breakpoints();
+            let ctx = baseline.context_mut() as *mut _;
+            let ret = with_ctx(ctx, || {
+                if let Some(image) = resume_image.take() {
+                    let msm = baseline
+                        .module
+                        .runnable_module
+                        .get_module_state_map()
+                        .unwrap();
+                    let code_base =
+                        baseline.module.runnable_module.get_code().unwrap().as_ptr()
+                            as usize;
+                    invoke_call_return_on_stack(
+                        &msm,
+                        code_base,
+                        image,
+                        baseline.context_mut(),
+                        breakpoints.clone(),
+                    )
+                    .map(|_| ())
+                } else {
+                    catch_unsafe_unwind(
+                        || start_raw(baseline.context_mut()),
+                        breakpoints.clone(),
+                    )
+                }
+            });
+            if let Err(e) = ret {
+                if let Some(new_image) = e.downcast_ref::<InstanceImage>() {
+                    // Tier switch event
+                    if !was_sigint_triggered_fault()
+                        && opt_state.outcome.lock().unwrap().is_some()
+                    {
+                        resume_image = Some(new_image.clone());
+                        continue;
+                    }
+                    let op = interactive_shell(InteractiveShellContext {
+                        image: Some(new_image.clone()),
+                    });
+                    match op {
+                        ShellExitOperation::ContinueWith(new_image) => {
+                            resume_image = Some(new_image);
+                        }
+                    }
+                } else {
+                    return Err("Error while executing WebAssembly".into());
+                }
+            } else {
+                return Ok(());
+            }
+        }
+    }
+}

src/bin/wasmer.rs

@@ -9,7 +9,6 @@ use std::io::Read;
 use std::path::PathBuf;
 use std::process::exit;
 use std::str::FromStr;
-use std::sync::{Arc, Mutex};
 
 use hashbrown::HashMap;
 use structopt::StructOpt;
@@ -26,14 +25,14 @@ use wasmer_runtime_core::{
     self,
     backend::{Backend, Compiler, CompilerConfig, MemoryBoundCheckMode},
     debug,
-    fault::{set_wasm_interrupt_on_ctx, was_sigint_triggered_fault},
     loader::{Instance as LoadedInstance, LocalLoader},
-    Instance, Module,
 };
 #[cfg(feature = "backend-singlepass")]
 use wasmer_singlepass_backend::SinglePassCompiler;
 #[cfg(feature = "wasi")]
 use wasmer_wasi;
+#[cfg(feature = "managed")]
+use wasmer_runtime_core::tiering::{InteractiveShellContext, ShellExitOperation, run_tiering};
 
 // stub module to make conditional compilation happy
 #[cfg(not(feature = "wasi"))]
@@ -190,55 +189,6 @@ struct Validate {
     path: PathBuf,
 }
 
-struct OptimizationState {
-    outcome: Mutex<Option<OptimizationOutcome>>,
-}
-
-struct OptimizationOutcome {
-    module: Module,
-}
-
-struct Defer<F: FnOnce()>(Option<F>);
-impl<F: FnOnce()> Drop for Defer<F> {
-    fn drop(&mut self) {
-        if let Some(f) = self.0.take() {
-            f();
-        }
-    }
-}
-
-#[repr(transparent)]
-struct CtxWrapper(*mut wasmer_runtime_core::vm::Ctx);
-unsafe impl Send for CtxWrapper {}
-unsafe impl Sync for CtxWrapper {}
-
-#[cfg(feature = "managed")]
-unsafe fn begin_optimize(
-    binary: Vec<u8>,
-    compiler: Box<dyn Compiler>,
-    ctx: Arc<Mutex<CtxWrapper>>,
-    state: Arc<OptimizationState>,
-) {
-    let module = match webassembly::compile_with_config_with(
-        &binary[..],
-        CompilerConfig {
-            symbol_map: None,
-            track_state: true,
-            ..Default::default()
-        },
-        &*compiler,
-    ) {
-        Ok(x) => x,
-        Err(_) => return,
-    };
-
-    let ctx_inner = ctx.lock().unwrap();
-    if !ctx_inner.0.is_null() {
-        *state.outcome.lock().unwrap() = Some(OptimizationOutcome { module });
-        set_wasm_interrupt_on_ctx(ctx_inner.0);
-    }
-}
-
 /// Read the contents of a file
 fn read_file_contents(path: &PathBuf) -> Result<Vec<u8>, io::Error> {
     let mut buffer: Vec<u8> = Vec::new();
@@ -556,194 +506,32 @@ fn execute_wasm(options: &Run) -> Result<(), String> {
                 .map_err(|e| format!("Can't instantiate module: {:?}", e))?;
 
             let start: Func<(), ()> = instance.func("_start").map_err(|e| format!("{:?}", e))?;
+            let start_raw: extern "C" fn(&mut wasmer_runtime_core::vm::Ctx) = unsafe {
+                ::std::mem::transmute(start.get_vm_func())
+            };
 
             #[cfg(feature = "managed")]
-            unsafe {
+            run_tiering(
-                if options.backend == Backend::Singlepass {
+                module.info(),
-                    use wasmer_runtime_core::fault::{
+                &wasm_binary,
-                        catch_unsafe_unwind, ensure_sighandler, with_ctx,
+                if let Some(ref path) = options.resume {
-                    };
-                    use wasmer_runtime_core::state::{
-                        x64::invoke_call_return_on_stack, InstanceImage,
-                    };
-                    use wasmer_runtime_core::vm::Ctx;
-
-                    ensure_sighandler();
-
-                    let start_raw: extern "C" fn(&mut Ctx) =
-                        ::std::mem::transmute(start.get_vm_func());
-
-                    let mut image: Option<InstanceImage> = if let Some(ref path) = options.resume {
                     let mut f = File::open(path).unwrap();
                     let mut out: Vec<u8> = vec![];
                     f.read_to_end(&mut out).unwrap();
-                        Some(InstanceImage::from_bytes(&out).expect("failed to decode image"))
+                    Some(wasmer_runtime_core::state::InstanceImage::from_bytes(&out).expect("failed to decode image"))
                 } else {
                     None
-                    };
-
-                    let ctx_box =
-                        Arc::new(Mutex::new(CtxWrapper(instance.context_mut() as *mut _)));
-                    // Ensure that the ctx pointer's lifetime is not longer than Instance's.
-                    let _deferred_ctx_box_cleanup: Defer<_> = {
-                        let ctx_box = ctx_box.clone();
-                        Defer(Some(move || {
-                            ctx_box.lock().unwrap().0 = ::std::ptr::null_mut();
-                        }))
-                    };
-                    let opt_state = Arc::new(OptimizationState {
-                        outcome: Mutex::new(None),
-                    });
-
-                    {
-                        let wasm_binary = wasm_binary.to_vec();
-                        let ctx_box = ctx_box.clone();
-                        let opt_state = opt_state.clone();
-                        ::std::thread::spawn(move || {
-                            // TODO: CLI option for optimized backend
-                            begin_optimize(
-                                wasm_binary,
-                                get_compiler_by_backend(Backend::LLVM).unwrap(),
-                                ctx_box,
-                                opt_state,
-                            );
-                        });
-                    }
-
-                    let mut patched = false;
-                    let mut optimized_instance: Option<Instance> = None;
-
-                    loop {
-                        let optimized: Option<&mut Instance> =
-                            if let Some(ref mut x) = optimized_instance {
-                                Some(x)
-                            } else {
-                                let mut outcome = opt_state.outcome.lock().unwrap();
-                                if let Some(x) = outcome.take() {
-                                    let instance =
-                                        x.module.instantiate(&import_object).map_err(|e| {
-                                            format!("Can't instantiate module: {:?}", e)
-                                        })?;
-                                    optimized_instance = Some(instance);
-                                    optimized_instance.as_mut()
-                                } else {
-                                    None
-                                }
-                            };
-                        if !patched && false {
-                            if let Some(optimized) = optimized {
-                                let base = module.info().imported_functions.len();
-                                let code_ptr = optimized
-                                    .module
-                                    .runnable_module
-                                    .get_code()
-                                    .unwrap()
-                                    .as_ptr()
-                                    as usize;
-                                let target_addresses: Vec<usize> = optimized
-                                    .module
-                                    .runnable_module
-                                    .get_local_function_offsets()
-                                    .unwrap()
-                                    .into_iter()
-                                    .map(|x| code_ptr + x)
-                                    .collect();
-                                assert_eq!(
-                                    target_addresses.len(),
-                                    module.info().func_assoc.len() - base
-                                );
-                                for i in base..module.info().func_assoc.len() {
-                                    instance
-                                        .module
-                                        .runnable_module
-                                        .patch_local_function(i - base, target_addresses[i - base]);
-                                }
-                                patched = true;
-                                eprintln!("Patched");
-                            }
-                        }
-                        let breakpoints = instance.module.runnable_module.get_breakpoints();
-                        let ctx = instance.context_mut() as *mut _;
-                        let ret = with_ctx(ctx, || {
-                            if let Some(image) = image.take() {
-                                let msm = instance
-                                    .module
-                                    .runnable_module
-                                    .get_module_state_map()
-                                    .unwrap();
-                                let code_base =
-                                    instance.module.runnable_module.get_code().unwrap().as_ptr()
-                                        as usize;
-                                invoke_call_return_on_stack(
-                                    &msm,
-                                    code_base,
-                                    image,
-                                    instance.context_mut(),
-                                    breakpoints.clone(),
-                                )
-                                .map(|_| ())
-                            } else {
-                                catch_unsafe_unwind(
-                                    || start_raw(instance.context_mut()),
-                                    breakpoints.clone(),
-                                )
-                            }
-                        });
-                        if let Err(e) = ret {
-                            if let Some(new_image) = e.downcast_ref::<InstanceImage>() {
-                                // Tier switch event
-                                if !was_sigint_triggered_fault()
-                                    && optimized_instance.is_none()
-                                    && opt_state.outcome.lock().unwrap().is_some()
-                                {
-                                    image = Some(new_image.clone());
-                                    continue;
-                                }
-                                let op = interactive_shell(InteractiveShellContext {
-                                    image: Some(new_image.clone()),
-                                });
-                                match op {
-                                    ShellExitOperation::ContinueWith(new_image, new_backend) => {
-                                        image = Some(new_image);
-                                        if let Some(new_backend) = new_backend {
-                                            let compiler =
-                                                match get_compiler_by_backend(new_backend) {
-                                                    Some(x) => x,
-                                                    None => {
-                                                        return Err(
-                                                            "the requested backend is not enabled"
-                                                                .into(),
-                                                        )
-                                                    }
-                                                };
-                                            let module = webassembly::compile_with_config_with(
-                                                &wasm_binary[..],
-                                                CompilerConfig {
-                                                    symbol_map: em_symbol_map.clone(),
-                                                    track_state,
-                                                    ..Default::default()
                 },
-                                                &*compiler,
+                &import_object,
-                                            )
+                start_raw,
-                                            .map_err(|e| {
+                &mut instance,
-                                                format!("Can't compile module: {:?}", e)
+                vec![
-                                            })?;
+                    Box::new(|| get_compiler_by_backend(Backend::LLVM).unwrap()),
-                                            instance = module.instantiate(&import_object).map_err(
+                ],
-                                                |e| format!("Can't instantiate module: {:?}", e),
+                interactive_shell,
             )?;
-                                        }
-                                    }
-                                }
-                            } else {
-                                return Err("Error while executing WebAssembly".into());
-                            }
-                        } else {
-                            return Ok(());
-                        }
-                    }
-                }
-            }
 
+            #[cfg(not(feature = "managed"))]
             {
                 use wasmer_runtime::error::RuntimeError;
                 let result = start.call();
@@ -783,17 +571,6 @@ fn execute_wasm(options: &Run) -> Result<(), String> {
     Ok(())
 }
 
-#[cfg(feature = "managed")]
-struct InteractiveShellContext {
-    image: Option<wasmer_runtime_core::state::InstanceImage>,
-}
-
-#[cfg(feature = "managed")]
-#[derive(Debug)]
-enum ShellExitOperation {
-    ContinueWith(wasmer_runtime_core::state::InstanceImage, Option<Backend>),
-}
-
 #[cfg(feature = "managed")]
 fn interactive_shell(mut ctx: InteractiveShellContext) -> ShellExitOperation {
     use std::io::Write;
@@ -844,7 +621,7 @@ fn interactive_shell(mut ctx: InteractiveShellContext) -> ShellExitOperation {
             }
             "continue" | "c" => {
                 if let Some(image) = ctx.image.take() {
-                    return ShellExitOperation::ContinueWith(image, None);
+                    return ShellExitOperation::ContinueWith(image);
                 } else {
                     println!("Program state not available, cannot continue execution");
                 }