wasm-scheduler: catch and log scheduler panics nicely

example-configs/wasm-minimal/src/main.rs

@@ -2,13 +2,39 @@
 
 use std::mem;
 use std::slice;
+use std::panic::set_hook;
 use std::fmt::Write;
 use std::os::raw::{c_void};
 
 use serde_json::{Value, from_slice, to_vec};
 
+extern {
+    fn log_panic(payload_ptr: *const u8, payload_len: usize,
+                        file_ptr: *const u8, file_len: usize, line: u32);
+}
+
 
 fn main() {
+    set_hook(Box::new(|panic_info| {
+        let payload = panic_info.payload();
+        let (ptr, len) = if let Some(s) = payload.downcast_ref::<&str>() {
+            (s.as_bytes().as_ptr(), s.len())
+        } else if let Some(s) = payload.downcast_ref::<String>() {
+            (s.as_bytes().as_ptr(), s.len())
+        } else {
+            (0 as *const u8, 0)
+        };
+        let (file_ptr, file_len, line) = match panic_info.location() {
+            Some(loc) => {
+                let file = loc.file().as_bytes();
+                (file.as_ptr(), file.len(), loc.line())
+            }
+            None => (0 as *const u8, 0, 0),
+        };
+        unsafe {
+            log_panic(ptr, len, file_ptr, file_len, line);
+        }
+    }));
 }
 
 #[no_mangle]

src/daemon/scheduler/wasm.rs

@@ -5,23 +5,35 @@ use std::path::{Path};
 use failure::{Error, err_msg};
 use serde_json::{Value as Json, to_vec, de};
 use serde::Serialize;
-use wasmi::{load_from_buffer, ModuleInstance, ImportsBuilder, ModuleRef};
+use wasmi::{self, load_from_buffer, ModuleInstance, ImportsBuilder, ModuleRef};
 use wasmi::RuntimeValue::I32;
-use wasmi::{MemoryRef, NopExternals};
+use wasmi::{MemoryRef};
+use wasmi::{Externals, RuntimeValue, RuntimeArgs, ModuleImportResolver};
+use wasmi::{FuncRef, ValueType, Signature, FuncInstance};
+
 
 const PAGE_SIZE: usize = 65536; // for some reason it's not in interpreter
+const PANIC_INDEX: usize = 0;
 
 
 pub(in scheduler) struct Scheduler {
     module: ModuleRef,
     memory: MemoryRef,
+    scheduler_util: Util,
 }
 
 struct ReadMemory {
     memory: MemoryRef,
     offset: usize,
 }
 
+struct Resolver;
+
+struct Util {
+    memory: MemoryRef,
+}
+
+
 pub(in scheduler) fn read(dir: &Path)
     -> Result<Scheduler, Error>
 {
@@ -36,14 +48,16 @@ pub(in scheduler) fn read(dir: &Path)
         .map_err(|e| err_msg(format!("error decoding wasm: {:?}", e)))?;
     let module = ModuleInstance::new(
             &module,
-            &ImportsBuilder::default(),
-        ).map_err(|e| err_msg(format!("error adding wasm module: {}", e)))?
-        .run_start(&mut NopExternals)
-        .map_err(|e| err_msg(format!("error starting wasm module: {}", e)))?;
-    let memory = module.export_by_name("memory")
+            &ImportsBuilder::new()
+            .with_resolver("env", &Resolver),
+        ).map_err(|e| err_msg(format!("error adding wasm module: {}", e)))?;
+    let memory = module.not_started_instance().export_by_name("memory")
         .and_then(|x| x.as_memory().map(Clone::clone))
         .ok_or_else(|| err_msg("no memory exported"))?;
-    Ok(Scheduler { module, memory })
+    let mut scheduler_util = Util { memory: memory.clone() };
+    let module = module.run_start(&mut scheduler_util)
+        .map_err(|e| err_msg(format!("error starting wasm module: {}", e)))?;
+    Ok(Scheduler { module, memory, scheduler_util })
 }
 
 impl Scheduler {
@@ -58,7 +72,7 @@ impl Scheduler {
         let off = match
             self.module.invoke_export("alloc",
                 &[I32(bytes.len() as i32)],
-                &mut NopExternals)
+                &mut self.scheduler_util)
         {
             Ok(Some(I32(off))) => off as u32,
             Ok(x) => return (
@@ -76,10 +90,10 @@ impl Scheduler {
         };
         let roff = self.module.invoke_export("scheduler",
             &[I32(off as i32), I32(bytes.len() as i32)],
-            &mut NopExternals);
+            &mut self.scheduler_util);
         match self.module.invoke_export("dealloc",
             &[I32(off as i32)],
-            &mut NopExternals)
+            &mut self.scheduler_util)
         {
             Ok(_) => {}
             Err(e) => return (
@@ -121,7 +135,7 @@ impl Scheduler {
                 String::from("deserialize error")),
         };
         match self.module.invoke_export("dealloc",
-            &[I32(roff as i32)], &mut NopExternals)
+            &[I32(roff as i32)], &mut self.scheduler_util)
         {
             Ok(_) => {}
             Err(e) => return (
@@ -151,3 +165,73 @@ impl io::Read for ReadMemory {
         }
     }
 }
+
+impl ModuleImportResolver for Resolver {
+    fn resolve_func(&self, field_name: &str, signature: &Signature)
+        -> Result<FuncRef, wasmi::Error>
+    {
+        match field_name {
+            "log_panic"
+            if signature == &Signature::new(&[ValueType::I32; 5][..], None)
+            => Ok(FuncInstance::alloc_host(signature.clone(), PANIC_INDEX)),
+            "log_panic" => Err(wasmi::Error::Instantiation(
+                format!("Export {} expects invalid signature {:?}",
+                    field_name, signature)
+            )),
+            _ => Err(wasmi::Error::Instantiation(
+                format!("Export {} not found", field_name),
+            ))
+        }
+    }
+}
+
+impl Externals for Util {
+    fn invoke_index(&mut self, index: usize, args: RuntimeArgs)
+        -> Result<Option<RuntimeValue>, wasmi::Error>
+    {
+        match index {
+            PANIC_INDEX => {
+                // panic(payload_str, payload_len, file_ptr, file_len, line);
+                let payload_ptr: Result<u32, _> = args.nth(0);
+                let payload_len: Result<u32, _> = args.nth(1);
+                let file_ptr: Result<u32, _> = args.nth(2);
+                let file_len: Result<u32, _> = args.nth(3);
+                let line: Result<u32, _> = args.nth(4);
+                let payload = match (payload_ptr, payload_len) {
+                    (_, Ok(0)) | (_, Err(_)) | (Err(_), _) => {
+                        "<non-str payload>".into()
+                    }
+                    (Ok(off), Ok(len)) => {
+                        let mut buf = vec![0u8; len as usize];
+                        match self.memory.get_into(off, &mut buf) {
+                            Ok(()) => String::from_utf8_lossy(&buf).to_string(),
+                            Err(e) => {
+                                debug!("Error reading panic payload: {}", e);
+                                "<can't read payload>".into()
+                            }
+                        }
+                    }
+                };
+                let filename = match (file_ptr, file_len) {
+                    (_, Ok(0)) | (_, Err(_)) | (Err(_), _) => {
+                        "<unknown file>".into()
+                    }
+                    (Ok(off), Ok(len)) => {
+                        let mut buf = vec![0u8; len as usize];
+                        match self.memory.get_into(off, &mut buf) {
+                            Ok(()) => String::from_utf8_lossy(&buf).to_string(),
+                            Err(e) => {
+                                debug!("Error reading panic filename: {}", e);
+                                "<can't read filename>".into()
+                            }
+                        }
+                    }
+                };
+                error!("Scheduler panicked at {:?}:{}: {:?}",
+                    filename, line.unwrap_or(0), payload);
+                Ok(None)
+            }
+            _ => panic!("Unimplemented function at {}", index),
+        }
+    }
+}