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;