Add copy-on-write based instance reuse mechanism

crates/runtime/src/instance.rs

@@ -88,6 +88,17 @@ pub(crate) struct Instance {
     /// allocation, but some host-defined objects will store their state here.
     host_state: Box<dyn Any + Send + Sync>,
 
+    /// A flag which stores whether this instance is reusable.
+    ///
+    /// If this is true then [`InstanceHandle::create_snapshot`] and
+    /// [`InstanceHandle::restore_snapshot`] can be called on this instance.
+    #[cfg(target_os = "linux")]
+    is_reusable: bool,
+
+    /// The values of globals saved in [`InstanceHandle::create_snapshot`].
+    #[cfg(target_os = "linux")]
+    saved_globals: Vec<(GlobalIndex, VMGlobalDefinition)>,
+
     /// Additional context used by compiled wasm code. This field is last, and
     /// represents a dynamically-sized array that extends beyond the nominal
     /// end of the struct (similar to a flexible array member).
@@ -843,4 +854,71 @@ impl InstanceHandle {
             instance: self.instance,
         }
     }
+
+    /// Creates an internal snapshot of this instance's memory and tables
+    /// to be restored at a later time.
+    ///
+    /// Should only be called once.
+    #[cfg(target_os = "linux")]
+    pub fn create_snapshot(&mut self) -> Result<(), anyhow::Error> {
+        let instance = self.instance_mut();
+        if !instance.is_reusable {
+            anyhow::bail!("instance is not reusable");
+        }
+
+        for memory in instance.memories.values_mut() {
+            memory.create_snapshot()?;
+        }
+
+        for table in instance.tables.values_mut() {
+            table.create_snapshot()?;
+        }
+
+        for index in instance.module.globals.keys() {
+            let value = unsafe { std::ptr::read(instance.defined_or_imported_global_ptr(index)) };
+            instance.saved_globals.push((index, value));
+        }
+
+        Ok(())
+    }
+
+    /// Resets this instance's memory, tables and globals to their initial
+    /// state from when the [`InstanceHandle::create_snapshot`] was called.
+    ///
+    /// Can be called any number of times.
+    #[cfg(target_os = "linux")]
+    pub fn restore_snapshot(&mut self) -> Result<(), anyhow::Error> {
+        let instance = self.instance_mut();
+        if !instance.is_reusable {
+            anyhow::bail!("instance is not reusable");
+        }
+
+        for index in (0..instance.memories.len() as u32).map(DefinedMemoryIndex::from_u32) {
+            let memory = &mut instance.memories[index];
+            memory.restore_snapshot()?;
+
+            let vmmemory = memory.vmmemory();
+            instance.set_memory(index, vmmemory);
+        }
+
+        for index in (0..instance.tables.len() as u32).map(DefinedTableIndex::from_u32) {
+            let table = &mut instance.tables[index];
+            table.restore_snapshot()?;
+
+            let vmtable = table.vmtable();
+            instance.set_table(index, vmtable);
+        }
+
+        for (index, value) in &instance.saved_globals {
+            unsafe {
+                std::ptr::copy_nonoverlapping(
+                    value,
+                    instance.defined_or_imported_global_ptr(*index),
+                    1,
+                );
+            }
+        }
+
+        Ok(())
+    }
 }

crates/runtime/src/instance/allocator.rs

@@ -31,6 +31,18 @@ pub use self::pooling::{
     InstanceLimits, ModuleLimits, PoolingAllocationStrategy, PoolingInstanceAllocator,
 };
 
+/// The source of memory which the instance allocator is supposed to use.
+#[derive(Copy, Clone, Debug)]
+pub enum MemorySource {
+    /// Use the `MemoryCreator` provided when configuring the engine.
+    FromCreator,
+
+    /// Ignore the `MemoryCreator` provided when configuring the engine;
+    /// instead directly use `mmap` to allocate memory appropriate for a reusable instance.
+    #[cfg(target_os = "linux")]
+    CopyOnWriteInitialize,
+}
+
 /// Represents a request for a new runtime instance.
 pub struct InstanceAllocationRequest<'a> {
     /// The module being instantiated.
@@ -77,6 +89,9 @@ pub struct InstanceAllocationRequest<'a> {
     /// responsibility of the callee when allocating to ensure that this data
     /// outlives the instance.
     pub wasm_data: *const [u8],
+
+    /// The memory source to use when allocating a new instance.
+    pub memory_source: MemorySource,
 }
 
 /// A pointer to a Store. This Option<*mut dyn Store> is wrapped in a struct
@@ -131,6 +146,10 @@ pub enum InstantiationError {
     /// A limit on how many instances are supported has been reached.
     #[error("Limit of {0} concurrent instances has been reached")]
     Limit(u32),
+
+    /// The configured allocation strategy is incompatible with this kind of an instance.
+    #[error("Incompatible allocation strategy")]
+    IncompatibleAllocationStrategy,
 }
 
 /// An error while creating a fiber stack.
@@ -647,6 +666,7 @@ impl OnDemandInstanceAllocator {
         &self,
         module: &Module,
         store: &mut StorePtr,
+        memory_source: &MemorySource,
     ) -> Result<PrimaryMap<DefinedMemoryIndex, Memory>, InstantiationError> {
         let creator = self
             .mem_creator
@@ -655,16 +675,38 @@ impl OnDemandInstanceAllocator {
         let num_imports = module.num_imported_memories;
         let mut memories: PrimaryMap<DefinedMemoryIndex, _> =
             PrimaryMap::with_capacity(module.memory_plans.len() - num_imports);
-        for plan in &module.memory_plans.values().as_slice()[num_imports..] {
-            memories.push(
-                Memory::new_dynamic(plan, creator, unsafe {
-                    store
-                        .get()
-                        .expect("if module has memory plans, store is not empty")
-                })
-                .map_err(InstantiationError::Resource)?,
-            );
+
+        let plans = &module.memory_plans.values().as_slice()[num_imports..];
+        if plans.is_empty() {
+            return Ok(memories);
+        }
+
+        let store = unsafe {
+            store
+                .get()
+                .expect("if module has memory plans, store is not empty")
+        };
+
+        match memory_source {
+            MemorySource::FromCreator => {
+                for plan in plans {
+                    memories.push(
+                        Memory::new_dynamic(plan, creator, store)
+                            .map_err(InstantiationError::Resource)?,
+                    );
+                }
+            }
+            #[cfg(target_os = "linux")]
+            MemorySource::CopyOnWriteInitialize => {
+                for plan in plans {
+                    memories.push(
+                        Memory::new_snapshottable(plan, store)
+                            .map_err(InstantiationError::Resource)?,
+                    );
+                }
+            }
         }
+
         Ok(memories)
     }
 }
@@ -684,7 +726,7 @@ unsafe impl InstanceAllocator for OnDemandInstanceAllocator {
         &self,
         mut req: InstanceAllocationRequest,
     ) -> Result<InstanceHandle, InstantiationError> {
-        let memories = self.create_memories(&req.module, &mut req.store)?;
+        let memories = self.create_memories(&req.module, &mut req.store, &req.memory_source)?;
         let tables = Self::create_tables(&req.module, &mut req.store)?;
 
         let host_state = std::mem::replace(&mut req.host_state, Box::new(()));
@@ -702,6 +744,15 @@ unsafe impl InstanceAllocator for OnDemandInstanceAllocator {
                 vmctx: VMContext {
                     _marker: marker::PhantomPinned,
                 },
+
+                #[cfg(target_os = "linux")]
+                is_reusable: match req.memory_source {
+                    MemorySource::FromCreator => false,
+                    MemorySource::CopyOnWriteInitialize => true,
+                },
+
+                #[cfg(target_os = "linux")]
+                saved_globals: Default::default(),
             };
             let layout = instance.alloc_layout();
             let instance_ptr = alloc::alloc(layout) as *mut Instance;

crates/runtime/src/instance/allocator/pooling.rs

@@ -9,7 +9,7 @@
 
 use super::{
     initialize_instance, initialize_vmcontext, InstanceAllocationRequest, InstanceAllocator,
-    InstanceHandle, InstantiationError,
+    InstanceHandle, InstantiationError, MemorySource,
 };
 use crate::{instance::Instance, Memory, Mmap, Table, VMContext};
 use anyhow::{anyhow, bail, Context, Result};
@@ -367,6 +367,12 @@ impl InstancePool {
                     vmctx: VMContext {
                         _marker: marker::PhantomPinned,
                     },
+
+                    #[cfg(target_os = "linux")]
+                    is_reusable: false,
+
+                    #[cfg(target_os = "linux")]
+                    saved_globals: Default::default(),
                 },
             );
         }
@@ -414,6 +420,14 @@ impl InstancePool {
         strategy: PoolingAllocationStrategy,
         req: InstanceAllocationRequest,
     ) -> Result<InstanceHandle, InstantiationError> {
+        #[cfg(target_os = "linux")]
+        match req.memory_source {
+            MemorySource::FromCreator => {}
+            MemorySource::CopyOnWriteInitialize => {
+                return Err(InstantiationError::IncompatibleAllocationStrategy);
+            }
+        }
+
         let index = {
             let mut free_list = self.free_list.lock().unwrap();
             if free_list.is_empty() {
@@ -1049,7 +1063,7 @@ unsafe impl InstanceAllocator for PoolingInstanceAllocator {
 #[cfg(test)]
 mod test {
     use super::*;
-    use crate::{Imports, StorePtr, VMSharedSignatureIndex};
+    use crate::{Imports, MemorySource, StorePtr, VMSharedSignatureIndex};
     use wasmtime_environ::{
         EntityRef, Global, GlobalInit, Memory, MemoryPlan, ModuleType, SignatureIndex, Table,
         TablePlan, TableStyle, WasmType,
@@ -1413,6 +1427,7 @@ mod test {
                             host_state: Box::new(()),
                             store: StorePtr::empty(),
                             wasm_data: &[],
+                            memory_source: MemorySource::FromCreator,
                         },
                     )
                     .expect("allocation should succeed"),
@@ -1437,6 +1452,7 @@ mod test {
                 host_state: Box::new(()),
                 store: StorePtr::empty(),
                 wasm_data: &[],
+                memory_source: MemorySource::FromCreator,
             },
         ) {
             Err(InstantiationError::Limit(3)) => {}

crates/runtime/src/lib.rs

@@ -41,7 +41,7 @@ pub use crate::externref::*;
 pub use crate::imports::Imports;
 pub use crate::instance::{
     InstanceAllocationRequest, InstanceAllocator, InstanceHandle, InstantiationError, LinkError,
-    OnDemandInstanceAllocator, StorePtr,
+    MemorySource, OnDemandInstanceAllocator, StorePtr,
 };
 #[cfg(feature = "pooling-allocator")]
 pub use crate::instance::{