Option for host managed memory

crates/api/src/instance.rs

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

crates/api/src/runtime.rs

@@ -9,6 +9,7 @@ use wasmtime_environ::settings::{self, Configurable};
 use wasmtime_environ::CacheConfig;
 use wasmtime_jit::{native, CompilationStrategy, Compiler};
 use wasmtime_profiling::{JitDumpAgent, NullProfilerAgent, ProfilingAgent};
+use wasmtime_runtime::Allocator;
 
 // Runtime Environment
 
@@ -27,6 +28,7 @@ pub struct Config {
     pub(crate) strategy: CompilationStrategy,
     pub(crate) cache_config: CacheConfig,
     pub(crate) profiler: Arc<dyn ProfilingAgent>,
+    pub(crate) allocator: Option<Arc<dyn Allocator>>,
 }
 
 impl Config {
@@ -66,6 +68,7 @@ impl Config {
             strategy: CompilationStrategy::Auto,
             cache_config: CacheConfig::new_cache_disabled(),
             profiler: Arc::new(NullProfilerAgent),
+            allocator: None,
         }
     }
 
@@ -309,6 +312,12 @@ impl Config {
         self.cache_config = wasmtime_environ::CacheConfig::from_file(None)?;
         Ok(self)
     }
+
+    /// Sets a custom memory allocator
+    pub fn with_host_memory(&mut self, allocator: Arc<dyn Allocator>) -> &mut Self {
+        self.allocator = Some(allocator);
+        self
+    }
 }
 
 impl Default for Config {
@@ -486,6 +495,11 @@ impl Store {
         &self.inner.engine
     }
 
+    /// Returns an optional reference to a memory ['Allocator']
+    pub(crate) fn allocator(&self) -> Option<&dyn Allocator> {
+        self.engine().config.allocator.as_ref().map(|a| &**a)
+    }
+
     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.allocator(),
             &data_initializers,
             signatures.into_boxed_slice(),
             None,

crates/jit/src/instantiate.rs

@@ -20,7 +20,7 @@ use wasmtime_environ::{
 };
 use wasmtime_profiling::ProfilingAgent;
 use wasmtime_runtime::{
-    GdbJitImageRegistration, InstanceHandle, InstantiationError, SignatureRegistry,
+    Allocator, GdbJitImageRegistration, InstanceHandle, InstantiationError, SignatureRegistry,
     TrapRegistration, VMFunctionBody, VMSharedSignatureIndex, VMTrampoline,
 };
 
@@ -203,6 +203,7 @@ impl CompiledModule {
         is_bulk_memory: bool,
         resolver: &mut dyn Resolver,
         sig_registry: &SignatureRegistry,
+        allocator: Option<&dyn Allocator>,
     ) -> Result<InstanceHandle, InstantiationError> {
         let data_initializers = self
             .data_initializers
@@ -219,6 +220,7 @@ impl CompiledModule {
             self.finished_functions.clone(),
             self.trampolines.clone(),
             imports,
+            allocator,
             &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,
+    allocator: Option<&dyn Allocator>,
 ) -> Result<InstanceHandle, SetupError> {
     let instance = CompiledModule::new(compiler, data, debug_info, profiler)?.instantiate(
         is_bulk_memory,
         resolver,
         compiler.signatures(),
+        allocator,
     )?;
     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::{Allocator, DefaultAllocator, LinearMemory};
 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 LinearMemory>>,
 
     /// WebAssembly table data.
     tables: BoxedSlice<DefinedTableIndex, Table>,
@@ -861,14 +861,16 @@ impl InstanceHandle {
         finished_functions: BoxedSlice<DefinedFuncIndex, *mut [VMFunctionBody]>,
         trampolines: HashMap<VMSharedSignatureIndex, VMTrampoline>,
         imports: Imports,
+        allocator: Option<&dyn Allocator>,
         data_initializers: &[DataInitializer<'_>],
         vmshared_signatures: BoxedSlice<SignatureIndex, VMSharedSignatureIndex>,
         dbg_jit_registration: Option<Rc<GdbJitImageRegistration>>,
         is_bulk_memory: bool,
         host_state: Box<dyn Any>,
     ) -> Result<Self, InstantiationError> {
         let tables = create_tables(&module);
-        let memories = create_memories(&module)?;
+
+        let memories = create_memories(&module, allocator.unwrap_or(&DefaultAllocator {}))?;
 
         let vmctx_tables = tables
             .values()
@@ -878,7 +880,7 @@ impl InstanceHandle {
 
         let vmctx_memories = memories
             .values()
-            .map(LinearMemory::vmmemory)
+            .map(|a| a.vmmemory())
             .collect::<PrimaryMap<DefinedMemoryIndex, _>>()
             .into_boxed_slice();
 
@@ -1300,12 +1302,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> {
+    allocator: &dyn Allocator,
+) -> Result<BoxedSlice<DefinedMemoryIndex, Box<dyn LinearMemory>>, 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(
+            allocator
+                .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::{Allocator, LinearMemory};
 pub use crate::mmap::Mmap;
 pub use crate::sig_registry::SignatureRegistry;
 pub use crate::table::Table;

crates/runtime/src/memory.rs

@@ -9,9 +9,38 @@ use std::cell::RefCell;
 use std::convert::TryFrom;
 use wasmtime_environ::{MemoryPlan, MemoryStyle, WASM_MAX_PAGES, WASM_PAGE_SIZE};
 
+/// A memory allocator
+pub trait Allocator: Send + Sync {
+    /// Create new LinearMemory
+    fn new_memory(&self, plan: &MemoryPlan) -> Result<Box<dyn LinearMemory>, String>;
+}
+
+/// A default memory allocator used by Wasmtime
+pub struct DefaultAllocator;
+
+impl Allocator for DefaultAllocator {
+    /// Create new MmapMemory
+    fn new_memory(&self, plan: &MemoryPlan) -> Result<Box<dyn LinearMemory>, String> {
+        Ok(Box::new(MmapMemory::new(plan)?) as Box<dyn LinearMemory>)
+    }
+}
+
+/// A linear memory
+pub 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 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 +64,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,17 +106,19 @@ impl LinearMemory {
             needs_signal_handlers,
         })
     }
+}
 
+impl LinearMemory for MmapMemory {
     /// Returns the number of allocated wasm pages.
-    pub fn size(&self) -> u32 {
+    fn size(&self) -> u32 {
         self.mmap.borrow().size
     }
 
     /// Grow memory by the specified amount of wasm pages.
     ///
     /// 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 +174,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(),