Option for host managed memory (#1400)

* Option for host managed memory

* Rename Allocator to MemoryCreator

* Create LinearMemory and MemoryCreator traits in api

* Leave only one as_ptr function in LinearMemory trait

* Memory creator test

* Update comments/docs for LinearMemory and MemoryCreator traits

* Add guard page to the custom memory example

* Remove mut from LinearMemory trait as_ptr

* Host_memory_grow test

crates/api/src/externals.rs

@@ -853,3 +853,46 @@ impl Memory {
         }
     }
 }
+
+/// A linear memory. This trait provides an interface for raw memory buffers which are used
+/// by wasmtime, e.g. inside ['Memory']. Such buffers are in principle not thread safe.
+/// By implementing this trait together with MemoryCreator,
+/// one can supply wasmtime with custom allocated host managed memory.
+///
+/// # Safety
+/// The memory should be page aligned and a multiple of page size.
+/// To prevent possible silent overflows, the memory should be protected by a guard page.
+/// Additionally the safety concerns explained in ['Memory'], for accessing the memory
+/// apply here as well.
+///
+/// Note that this is a relatively new and experimental feature and it is recommended
+/// to be familiar with wasmtime runtime code to use it.
+pub unsafe trait LinearMemory {
+    /// Returns the number of allocated wasm pages.
+    fn size(&self) -> u32;
+
+    /// Grow memory by the specified amount of wasm pages.
+    ///
+    /// Returns `None` if memory can't be grown by the specified amount
+    /// of wasm pages.
+    fn grow(&self, delta: u32) -> Option<u32>;
+
+    /// Return the allocated memory as a mutable pointer to u8.
+    fn as_ptr(&self) -> *mut u8;
+}
+
+/// A memory creator. Can be used to provide a memory creator
+/// to wasmtime which supplies host managed memory.
+///
+/// # Safety
+/// This trait is unsafe, as the memory safety depends on proper implementation of
+/// memory management. Memories created by the MemoryCreator should always be treated
+/// as owned by wasmtime instance, and any modification of them outside of wasmtime
+/// invoked routines is unsafe and may lead to corruption.
+///
+/// Note that this is a relatively new and experimental feature and it is recommended
+/// to be familiar with wasmtime runtime code to use it.
+pub unsafe trait MemoryCreator: Send + Sync {
+    /// Create new LinearMemory
+    fn new_memory(&self, ty: MemoryType) -> Result<Box<dyn LinearMemory>, String>;
+}

crates/api/src/instance.rs

@@ -31,6 +31,7 @@ fn instantiate(
                 config.validating_config.operator_config.enable_bulk_memory,
                 &mut resolver,
                 sig_registry,
+                config.memory_creator.as_ref().map(|a| a as _),
             )
             .map_err(|e| -> Error {
                 match e {

crates/api/src/runtime.rs

@@ -1,3 +1,5 @@
+use crate::externals::MemoryCreator;
+use crate::trampoline::MemoryCreatorProxy;
 use anyhow::Result;
 use std::cell::RefCell;
 use std::fmt;
@@ -9,6 +11,7 @@ use wasmtime_environ::settings::{self, Configurable};
 use wasmtime_environ::CacheConfig;
 use wasmtime_jit::{native, CompilationStrategy, Compiler};
 use wasmtime_profiling::{JitDumpAgent, NullProfilerAgent, ProfilingAgent, VTuneAgent};
+use wasmtime_runtime::RuntimeMemoryCreator;
 
 // Runtime Environment
 
@@ -27,6 +30,7 @@ pub struct Config {
     pub(crate) strategy: CompilationStrategy,
     pub(crate) cache_config: CacheConfig,
     pub(crate) profiler: Arc<dyn ProfilingAgent>,
+    pub(crate) memory_creator: Option<MemoryCreatorProxy>,
 }
 
 impl Config {
@@ -66,6 +70,7 @@ impl Config {
             strategy: CompilationStrategy::Auto,
             cache_config: CacheConfig::new_cache_disabled(),
             profiler: Arc::new(NullProfilerAgent),
+            memory_creator: None,
         }
     }
 
@@ -326,6 +331,12 @@ impl Config {
         self.cache_config = wasmtime_environ::CacheConfig::from_file(None)?;
         Ok(self)
     }
+
+    /// Sets a custom memory creator
+    pub fn with_host_memory(&mut self, mem_creator: Arc<dyn MemoryCreator>) -> &mut Self {
+        self.memory_creator = Some(MemoryCreatorProxy { mem_creator });
+        self
+    }
 }
 
 impl Default for Config {
@@ -506,6 +517,11 @@ impl Store {
         &self.inner.engine
     }
 
+    /// Returns an optional reference to a ['RuntimeMemoryCreator']
+    pub(crate) fn memory_creator(&self) -> Option<&dyn RuntimeMemoryCreator> {
+        self.engine().config.memory_creator.as_ref().map(|x| x as _)
+    }
+
     pub(crate) fn compiler(&self) -> std::cell::Ref<'_, Compiler> {
         self.inner.compiler.borrow()
     }

crates/api/src/trampoline/create_handle.rs

@@ -43,6 +43,7 @@ pub(crate) fn create_handle(
             finished_functions.into_boxed_slice(),
             trampolines,
             imports,
+            store.memory_creator(),
             &data_initializers,
             signatures.into_boxed_slice(),
             None,

crates/api/src/trampoline/memory.rs

@@ -1,10 +1,15 @@
 use super::create_handle::create_handle;
-use crate::MemoryType;
+use crate::externals::{LinearMemory, MemoryCreator};
 use crate::Store;
+use crate::{Limits, MemoryType};
 use anyhow::Result;
 use wasmtime_environ::entity::PrimaryMap;
-use wasmtime_environ::{wasm, Module};
-use wasmtime_runtime::InstanceHandle;
+use wasmtime_environ::{wasm, MemoryPlan, Module, WASM_PAGE_SIZE};
+use wasmtime_runtime::{
+    InstanceHandle, RuntimeLinearMemory, RuntimeMemoryCreator, VMMemoryDefinition,
+};
+
+use std::sync::Arc;
 
 pub fn create_handle_with_memory(store: &Store, memory: &MemoryType) -> Result<InstanceHandle> {
     let mut module = Module::new();
@@ -31,3 +36,38 @@ pub fn create_handle_with_memory(store: &Store, memory: &MemoryType) -> Result<I
         Box::new(()),
     )
 }
+
+struct LinearMemoryProxy {
+    mem: Box<dyn LinearMemory>,
+}
+
+impl RuntimeLinearMemory for LinearMemoryProxy {
+    fn size(&self) -> u32 {
+        self.mem.size()
+    }
+
+    fn grow(&self, delta: u32) -> Option<u32> {
+        self.mem.grow(delta)
+    }
+
+    fn vmmemory(&self) -> VMMemoryDefinition {
+        VMMemoryDefinition {
+            base: self.mem.as_ptr(),
+            current_length: self.mem.size() as usize * WASM_PAGE_SIZE as usize,
+        }
+    }
+}
+
+#[derive(Clone)]
+pub(crate) struct MemoryCreatorProxy {
+    pub(crate) mem_creator: Arc<dyn MemoryCreator>,
+}
+
+impl RuntimeMemoryCreator for MemoryCreatorProxy {
+    fn new_memory(&self, plan: &MemoryPlan) -> Result<Box<dyn RuntimeLinearMemory>, String> {
+        let ty = MemoryType::new(Limits::new(plan.memory.minimum, plan.memory.maximum));
+        self.mem_creator
+            .new_memory(ty)
+            .map(|mem| Box::new(LinearMemoryProxy { mem }) as Box<dyn RuntimeLinearMemory>)
+    }
+}

crates/api/src/trampoline/mod.rs

@@ -6,6 +6,8 @@ mod global;
 mod memory;
 mod table;
 
+pub(crate) use memory::MemoryCreatorProxy;
+
 use self::func::create_handle_with_function;
 use self::global::create_global;
 use self::memory::create_handle_with_memory;

crates/api/tests/memory_creator.rs

@@ -0,0 +1,192 @@
+#[cfg(not(target_os = "windows"))]
+mod not_for_windows {
+    use wasmtime::*;
+    use wasmtime_environ::{WASM_MAX_PAGES, WASM_PAGE_SIZE};
+
+    use libc::c_void;
+    use libc::MAP_FAILED;
+    use libc::{mmap, mprotect, munmap};
+    use libc::{sysconf, _SC_PAGESIZE};
+    use libc::{MAP_ANON, MAP_PRIVATE, PROT_NONE, PROT_READ, PROT_WRITE};
+
+    use std::cell::RefCell;
+    use std::io::Error;
+    use std::ptr::null_mut;
+    use std::sync::{Arc, Mutex};
+
+    struct CustomMemory {
+        mem: *mut c_void,
+        size: usize,
+        used_wasm_pages: RefCell<u32>,
+        glob_page_counter: Arc<Mutex<u64>>,
+    }
+
+    impl CustomMemory {
+        unsafe fn new(
+            num_wasm_pages: u32,
+            max_wasm_pages: u32,
+            glob_counter: Arc<Mutex<u64>>,
+        ) -> Self {
+            let page_size = sysconf(_SC_PAGESIZE) as usize;
+            let guard_size = page_size;
+            let size = max_wasm_pages as usize * WASM_PAGE_SIZE as usize + guard_size;
+            let used_size = num_wasm_pages as usize * WASM_PAGE_SIZE as usize;
+            assert_eq!(size % page_size, 0); // we rely on WASM_PAGE_SIZE being multiple of host page size
+
+            let mem = mmap(null_mut(), size, PROT_NONE, MAP_PRIVATE | MAP_ANON, -1, 0);
+            assert_ne!(mem, MAP_FAILED, "mmap failed: {}", Error::last_os_error());
+
+            let r = mprotect(mem, used_size, PROT_READ | PROT_WRITE);
+            assert_eq!(r, 0, "mprotect failed: {}", Error::last_os_error());
+            *glob_counter.lock().unwrap() += num_wasm_pages as u64;
+
+            Self {
+                mem,
+                size,
+                used_wasm_pages: RefCell::new(num_wasm_pages),
+                glob_page_counter: glob_counter,
+            }
+        }
+    }
+
+    impl Drop for CustomMemory {
+        fn drop(&mut self) {
+            let n = *self.used_wasm_pages.borrow() as u64;
+            *self.glob_page_counter.lock().unwrap() -= n;
+            let r = unsafe { munmap(self.mem, self.size) };
+            assert_eq!(r, 0, "munmap failed: {}", Error::last_os_error());
+        }
+    }
+
+    unsafe impl LinearMemory for CustomMemory {
+        fn size(&self) -> u32 {
+            *self.used_wasm_pages.borrow()
+        }
+
+        fn grow(&self, delta: u32) -> Option<u32> {
+            let delta_size = (delta as usize).checked_mul(WASM_PAGE_SIZE as usize)?;
+
+            let prev_pages = *self.used_wasm_pages.borrow();
+            let prev_size = (prev_pages as usize).checked_mul(WASM_PAGE_SIZE as usize)?;
+
+            let new_pages = prev_pages.checked_add(delta)?;
+            let new_size = (new_pages as usize).checked_mul(WASM_PAGE_SIZE as usize)?;
+
+            let guard_size = unsafe { sysconf(_SC_PAGESIZE) as usize };
+
+            if new_size > self.size - guard_size {
+                return None;
+            }
+            unsafe {
+                let start = (self.mem as *mut u8).add(prev_size) as _;
+                let r = mprotect(start, delta_size, PROT_READ | PROT_WRITE);
+                assert_eq!(r, 0, "mprotect failed: {}", Error::last_os_error());
+            }
+
+            *self.glob_page_counter.lock().unwrap() += delta as u64;
+            *self.used_wasm_pages.borrow_mut() = new_pages;
+            Some(prev_pages)
+        }
+
+        fn as_ptr(&self) -> *mut u8 {
+            self.mem as *mut u8
+        }
+    }
+
+    struct CustomMemoryCreator {
+        pub num_created_memories: Mutex<usize>,
+        pub num_total_pages: Arc<Mutex<u64>>,
+    }
+
+    impl CustomMemoryCreator {
+        pub fn new() -> Self {
+            Self {
+                num_created_memories: Mutex::new(0),
+                num_total_pages: Arc::new(Mutex::new(0)),
+            }
+        }
+    }
+
+    unsafe impl MemoryCreator for CustomMemoryCreator {
+        fn new_memory(&self, ty: MemoryType) -> Result<Box<dyn LinearMemory>, String> {
+            let max = ty.limits().max().unwrap_or(WASM_MAX_PAGES);
+            unsafe {
+                let mem = Box::new(CustomMemory::new(
+                    ty.limits().min(),
+                    max,
+                    self.num_total_pages.clone(),
+                ));
+                *self.num_created_memories.lock().unwrap() += 1;
+                Ok(mem)
+            }
+        }
+    }
+
+    #[test]
+    fn host_memory() -> anyhow::Result<()> {
+        let mem_creator = Arc::new(CustomMemoryCreator::new());
+        let mut config = Config::default();
+        config.with_host_memory(mem_creator.clone());
+        let engine = Engine::new(&config);
+        let store = Store::new(&engine);
+
+        let module = Module::new(
+            &store,
+            r#"
+            (module
+                (memory (export "memory") 1)
+            )
+        "#,
+        )?;
+        Instance::new(&module, &[])?;
+
+        assert_eq!(*mem_creator.num_created_memories.lock().unwrap(), 1);
+
+        Ok(())
+    }
+
+    #[test]
+    fn host_memory_grow() -> anyhow::Result<()> {
+        let mem_creator = Arc::new(CustomMemoryCreator::new());
+        let mut config = Config::default();
+        config.with_host_memory(mem_creator.clone());
+        let engine = Engine::new(&config);
+        let store = Store::new(&engine);
+
+        let module = Module::new(
+            &store,
+            r#"
+            (module
+                (func $f (drop (memory.grow (i32.const 1))))
+                (memory (export "memory") 1 2)
+                (start $f)
+            )
+        "#,
+        )?;
+
+        let instance1 = Instance::new(&module, &[])?;
+        let instance2 = Instance::new(&module, &[])?;
+
+        assert_eq!(*mem_creator.num_created_memories.lock().unwrap(), 2);
+
+        assert_eq!(
+            instance2
+                .get_export("memory")
+                .unwrap()
+                .memory()
+                .unwrap()
+                .size(),
+            2
+        );
+
+        // we take the lock outside the assert, so it won't get poisoned on assert failure
+        let tot_pages = *mem_creator.num_total_pages.lock().unwrap();
+        assert_eq!(tot_pages, 4);
+
+        drop(instance1);
+        let tot_pages = *mem_creator.num_total_pages.lock().unwrap();
+        assert_eq!(tot_pages, 2);
+
+        Ok(())
+    }
+}

crates/jit/src/instantiate.rs

@@ -20,8 +20,8 @@ use wasmtime_environ::{
 };
 use wasmtime_profiling::ProfilingAgent;
 use wasmtime_runtime::{
-    GdbJitImageRegistration, InstanceHandle, InstantiationError, SignatureRegistry,
-    TrapRegistration, VMFunctionBody, VMSharedSignatureIndex, VMTrampoline,
+    GdbJitImageRegistration, InstanceHandle, InstantiationError, RuntimeMemoryCreator,
+    SignatureRegistry, TrapRegistration, VMFunctionBody, VMSharedSignatureIndex, VMTrampoline,
 };
 
 /// An error condition while setting up a wasm instance, be it validation,
@@ -203,6 +203,7 @@ impl CompiledModule {
         is_bulk_memory: bool,
         resolver: &mut dyn Resolver,
         sig_registry: &SignatureRegistry,
+        mem_creator: Option<&dyn RuntimeMemoryCreator>,
     ) -> Result<InstanceHandle, InstantiationError> {
         let data_initializers = self
             .data_initializers
@@ -219,6 +220,7 @@ impl CompiledModule {
             self.finished_functions.clone(),
             self.trampolines.clone(),
             imports,
+            mem_creator,
             &data_initializers,
             self.signatures.clone(),
             self.dbg_jit_registration.as_ref().map(|r| Rc::clone(&r)),
@@ -283,11 +285,13 @@ pub unsafe fn instantiate(
     debug_info: bool,
     is_bulk_memory: bool,
     profiler: &dyn ProfilingAgent,
+    mem_creator: Option<&dyn RuntimeMemoryCreator>,
 ) -> Result<InstanceHandle, SetupError> {
     let instance = CompiledModule::new(compiler, data, debug_info, profiler)?.instantiate(
         is_bulk_memory,
         resolver,
         compiler.signatures(),
+        mem_creator,
     )?;
     Ok(instance)
 }

crates/runtime/src/instance.rs

@@ -5,7 +5,7 @@
 use crate::export::Export;
 use crate::imports::Imports;
 use crate::jit_int::GdbJitImageRegistration;
-use crate::memory::LinearMemory;
+use crate::memory::{DefaultMemoryCreator, RuntimeLinearMemory, RuntimeMemoryCreator};
 use crate::table::Table;
 use crate::traphandlers;
 use crate::traphandlers::{catch_traps, Trap};
@@ -82,7 +82,7 @@ pub(crate) struct Instance {
     offsets: VMOffsets,
 
     /// WebAssembly linear memory data.
-    memories: BoxedSlice<DefinedMemoryIndex, LinearMemory>,
+    memories: BoxedSlice<DefinedMemoryIndex, Box<dyn RuntimeLinearMemory>>,
 
     /// WebAssembly table data.
     tables: BoxedSlice<DefinedTableIndex, Table>,
@@ -861,6 +861,7 @@ impl InstanceHandle {
         finished_functions: BoxedSlice<DefinedFuncIndex, *mut [VMFunctionBody]>,
         trampolines: HashMap<VMSharedSignatureIndex, VMTrampoline>,
         imports: Imports,
+        mem_creator: Option<&dyn RuntimeMemoryCreator>,
         data_initializers: &[DataInitializer<'_>],
         vmshared_signatures: BoxedSlice<SignatureIndex, VMSharedSignatureIndex>,
         dbg_jit_registration: Option<Rc<GdbJitImageRegistration>>,
@@ -868,7 +869,7 @@ impl InstanceHandle {
         host_state: Box<dyn Any>,
     ) -> Result<Self, InstantiationError> {
         let tables = create_tables(&module);
-        let memories = create_memories(&module)?;
+        let memories = create_memories(&module, mem_creator.unwrap_or(&DefaultMemoryCreator {}))?;
 
         let vmctx_tables = tables
             .values()
@@ -878,7 +879,7 @@ impl InstanceHandle {
 
         let vmctx_memories = memories
             .values()
-            .map(LinearMemory::vmmemory)
+            .map(|a| a.vmmemory())
             .collect::<PrimaryMap<DefinedMemoryIndex, _>>()
             .into_boxed_slice();
 
@@ -1300,12 +1301,17 @@ fn initialize_passive_elements(instance: &Instance) {
 /// Allocate memory for just the memories of the current module.
 fn create_memories(
     module: &Module,
-) -> Result<BoxedSlice<DefinedMemoryIndex, LinearMemory>, InstantiationError> {
+    mem_creator: &dyn RuntimeMemoryCreator,
+) -> Result<BoxedSlice<DefinedMemoryIndex, Box<dyn RuntimeLinearMemory>>, InstantiationError> {
     let num_imports = module.imported_memories.len();
     let mut memories: PrimaryMap<DefinedMemoryIndex, _> =
         PrimaryMap::with_capacity(module.local.memory_plans.len() - num_imports);
     for plan in &module.local.memory_plans.values().as_slice()[num_imports..] {
-        memories.push(LinearMemory::new(plan).map_err(InstantiationError::Resource)?);
+        memories.push(
+            mem_creator
+                .new_memory(plan)
+                .map_err(InstantiationError::Resource)?,
+        );
     }
     Ok(memories.into_boxed_slice())
 }

crates/runtime/src/lib.rs

@@ -39,6 +39,7 @@ pub use crate::export::*;
 pub use crate::imports::Imports;
 pub use crate::instance::{InstanceHandle, InstantiationError, LinkError};
 pub use crate::jit_int::GdbJitImageRegistration;
+pub use crate::memory::{RuntimeLinearMemory, RuntimeMemoryCreator};
 pub use crate::mmap::Mmap;
 pub use crate::sig_registry::SignatureRegistry;
 pub use crate::table::Table;

crates/runtime/src/memory.rs

@@ -1,6 +1,6 @@
 //! Memory management for linear memories.
 //!
-//! `LinearMemory` is to WebAssembly linear memories what `Table` is to WebAssembly tables.
+//! `RuntimeLinearMemory` is to WebAssembly linear memories what `Table` is to WebAssembly tables.
 
 use crate::mmap::Mmap;
 use crate::vmcontext::VMMemoryDefinition;
@@ -9,9 +9,40 @@ use std::cell::RefCell;
 use std::convert::TryFrom;
 use wasmtime_environ::{MemoryPlan, MemoryStyle, WASM_MAX_PAGES, WASM_PAGE_SIZE};
 
+/// A memory allocator
+pub trait RuntimeMemoryCreator: Send + Sync {
+    /// Create new RuntimeLinearMemory
+    fn new_memory(&self, plan: &MemoryPlan) -> Result<Box<dyn RuntimeLinearMemory>, String>;
+}
+
+/// A default memory allocator used by Wasmtime
+pub struct DefaultMemoryCreator;
+
+impl RuntimeMemoryCreator for DefaultMemoryCreator {
+    /// Create new MmapMemory
+    fn new_memory(&self, plan: &MemoryPlan) -> Result<Box<dyn RuntimeLinearMemory>, String> {
+        Ok(Box::new(MmapMemory::new(plan)?) as Box<dyn RuntimeLinearMemory>)
+    }
+}
+
+/// A linear memory
+pub trait RuntimeLinearMemory {
+    /// Returns the number of allocated wasm pages.
+    fn size(&self) -> u32;
+
+    /// Grow memory by the specified amount of wasm pages.
+    ///
+    /// Returns `None` if memory can't be grown by the specified amount
+    /// of wasm pages.
+    fn grow(&self, delta: u32) -> Option<u32>;
+
+    /// Return a `VMMemoryDefinition` for exposing the memory to compiled wasm code.
+    fn vmmemory(&self) -> VMMemoryDefinition;
+}
+
 /// A linear memory instance.
 #[derive(Debug)]
-pub struct LinearMemory {
+pub struct MmapMemory {
     // The underlying allocation.
     mmap: RefCell<WasmMmap>,
 
@@ -35,7 +66,7 @@ struct WasmMmap {
     size: u32,
 }
 
-impl LinearMemory {
+impl MmapMemory {
     /// Create a new linear memory instance with specified minimum and maximum number of wasm pages.
     pub fn new(plan: &MemoryPlan) -> Result<Self, String> {
         // `maximum` cannot be set to more than `65536` pages.
@@ -77,9 +108,11 @@ impl LinearMemory {
             needs_signal_handlers,
         })
     }
+}
 
+impl RuntimeLinearMemory for MmapMemory {
     /// Returns the number of allocated wasm pages.
-    pub fn size(&self) -> u32 {
+    fn size(&self) -> u32 {
         self.mmap.borrow().size
     }
 
@@ -87,7 +120,7 @@ impl LinearMemory {
     ///
     /// Returns `None` if memory can't be grown by the specified amount
     /// of wasm pages.
-    pub fn grow(&self, delta: u32) -> Option<u32> {
+    fn grow(&self, delta: u32) -> Option<u32> {
         // Optimization of memory.grow 0 calls.
         let mut mmap = self.mmap.borrow_mut();
         if delta == 0 {
@@ -143,7 +176,7 @@ impl LinearMemory {
     }
 
     /// Return a `VMMemoryDefinition` for exposing the memory to compiled wasm code.
-    pub fn vmmemory(&self) -> VMMemoryDefinition {
+    fn vmmemory(&self) -> VMMemoryDefinition {
         let mut mmap = self.mmap.borrow_mut();
         VMMemoryDefinition {
             base: mmap.alloc.as_mut_ptr(),