diff --git a/crates/runtime/src/instance.rs b/crates/runtime/src/instance.rs index 64afecd76048..05a96e6042df 100644 --- a/crates/runtime/src/instance.rs +++ b/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, + /// 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(()) + } } diff --git a/crates/runtime/src/instance/allocator.rs b/crates/runtime/src/instance/allocator.rs index 55ffd3181822..fb1a76a938ef 100644 --- a/crates/runtime/src/instance/allocator.rs +++ b/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, InstantiationError> { let creator = self .mem_creator @@ -655,16 +675,38 @@ impl OnDemandInstanceAllocator { let num_imports = module.num_imported_memories; let mut memories: PrimaryMap = 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 { - 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; diff --git a/crates/runtime/src/instance/allocator/pooling.rs b/crates/runtime/src/instance/allocator/pooling.rs index 76614137d5e9..6e9d17e8605e 100644 --- a/crates/runtime/src/instance/allocator/pooling.rs +++ b/crates/runtime/src/instance/allocator/pooling.rs @@ -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 { + #[cfg(target_os = "linux")] + match req.memory_source { + super::MemorySource::FromCreator => {} + super::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)) => {} diff --git a/crates/runtime/src/instance/allocator/pooling/uffd.rs b/crates/runtime/src/instance/allocator/pooling/uffd.rs index 843a03f49129..2bdd2177c2f9 100644 --- a/crates/runtime/src/instance/allocator/pooling/uffd.rs +++ b/crates/runtime/src/instance/allocator/pooling/uffd.rs @@ -435,7 +435,7 @@ impl Drop for PageFaultHandler { mod test { use super::*; use crate::{ - Imports, InstanceAllocationRequest, InstanceLimits, ModuleLimits, + Imports, InstanceAllocationRequest, InstanceLimits, MemorySource, ModuleLimits, PoolingAllocationStrategy, Store, StorePtr, VMSharedSignatureIndex, }; use std::sync::Arc; @@ -582,6 +582,7 @@ mod test { host_state: Box::new(()), store: StorePtr::new(&mut mock_store), wasm_data: &[], + memory_source: MemorySource::FromCreator, }, ) .expect("instance should allocate"), diff --git a/crates/runtime/src/lib.rs b/crates/runtime/src/lib.rs index ba8b544b3a8f..5d5dfe118222 100644 --- a/crates/runtime/src/lib.rs +++ b/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::{ diff --git a/crates/runtime/src/memory.rs b/crates/runtime/src/memory.rs index 07c0c619cc53..471d0ab06467 100644 --- a/crates/runtime/src/memory.rs +++ b/crates/runtime/src/memory.rs @@ -86,6 +86,10 @@ pub struct MmapMemory { // optimize loads and stores with constant offsets. pre_guard_size: usize, offset_guard_size: usize, + + // The initial value of the `accessible` from when the memory was fossilized. + #[cfg(target_os = "linux")] + original_accessible: usize, } impl MmapMemory { @@ -122,6 +126,7 @@ impl MmapMemory { .ok_or_else(|| format_err!("cannot allocate {} with guard regions", minimum))?; let mut mmap = Mmap::accessible_reserved(0, request_bytes)?; + if minimum > 0 { mmap.make_accessible(pre_guard_bytes, minimum)?; } @@ -133,8 +138,35 @@ impl MmapMemory { pre_guard_size: pre_guard_bytes, offset_guard_size: offset_guard_bytes, extra_to_reserve_on_growth, + + #[cfg(target_os = "linux")] + original_accessible: 0, }) } + + #[cfg(target_os = "linux")] + fn create_snapshot(&mut self) -> Result<()> { + self.mmap + .create_snapshot(self.pre_guard_size, self.accessible)?; + self.original_accessible = self.accessible; + + Ok(()) + } + + #[cfg(target_os = "linux")] + fn restore_snapshot(&mut self) -> Result<()> { + unsafe { + self.mmap.reset(self.pre_guard_size, self.accessible)?; + if self.accessible > self.original_accessible { + self.mmap.make_inaccessible( + self.pre_guard_size + self.original_accessible, + self.accessible - self.original_accessible, + )?; + self.accessible = self.original_accessible; + } + Ok(()) + } + } } impl RuntimeLinearMemory for MmapMemory { @@ -159,13 +191,12 @@ impl RuntimeLinearMemory for MmapMemory { .and_then(|s| s.checked_add(self.offset_guard_size)) .ok_or_else(|| format_err!("overflow calculating size of memory allocation"))?; - let mut new_mmap = Mmap::accessible_reserved(0, request_bytes)?; - new_mmap.make_accessible(self.pre_guard_size, new_size)?; - - new_mmap.as_mut_slice()[self.pre_guard_size..][..self.accessible] - .copy_from_slice(&self.mmap.as_slice()[self.pre_guard_size..][..self.accessible]); - - self.mmap = new_mmap; + self.mmap.reallocate( + self.pre_guard_size, + self.accessible, + request_bytes, + new_size, + )?; } else { // If the new size of this heap fits within the existing allocation // then all we need to do is to make the new pages accessible. This @@ -219,6 +250,10 @@ pub enum Memory { /// A "dynamic" memory whose data is managed at runtime and lifetime is tied /// to this instance. Dynamic(Box), + + /// A memory for which [`Memory::create_snapshot`] can be called. + #[cfg(target_os = "linux")] + Snapshottable(MmapMemory), } impl Memory { @@ -232,6 +267,40 @@ impl Memory { Ok(Memory::Dynamic(creator.new_memory(plan, minimum, maximum)?)) } + /// Creates a new memory instance for the specified plan, with the intent to + /// fossilize a snapshot of it later with [`Memory::create_snapshot`]. + #[cfg(target_os = "linux")] + pub fn new_snapshottable(plan: &MemoryPlan, store: &mut dyn Store) -> Result { + let (minimum, maximum) = Self::limit_new(plan, store)?; + let memory = MmapMemory::new(plan, minimum, maximum)?; + Ok(Memory::Snapshottable(memory)) + } + + /// Creates an internal snapshot of this memory to be restored at a later time. + /// + /// Can only be called for memories created with [`Memory::new_snapshottable`]. + /// Should only be called once. + #[cfg(target_os = "linux")] + pub fn create_snapshot(&mut self) -> Result<()> { + match self { + Memory::Static { .. } => unreachable!(), + Memory::Dynamic(..) => unreachable!(), + Memory::Snapshottable(ref mut memory) => memory.create_snapshot(), + } + } + + /// Resets this memory to its initial state from when the [`Memory::create_snapshot`] was called. + /// + /// Can be called multiple times. + #[cfg(target_os = "linux")] + pub fn restore_snapshot(&mut self) -> Result<()> { + match self { + Memory::Static { .. } => unreachable!(), + Memory::Dynamic(..) => unreachable!(), + Memory::Snapshottable(ref mut memory) => memory.restore_snapshot(), + } + } + /// Create a new static (immovable) memory instance for the specified plan. pub fn new_static( plan: &MemoryPlan, @@ -347,6 +416,8 @@ impl Memory { match self { Memory::Static { size, .. } => *size, Memory::Dynamic(mem) => mem.byte_size(), + #[cfg(target_os = "linux")] + Memory::Snapshottable(mem) => mem.byte_size(), } } @@ -360,6 +431,8 @@ impl Memory { match self { Memory::Static { base, .. } => Some(base.len()), Memory::Dynamic(mem) => mem.maximum_byte_size(), + #[cfg(target_os = "linux")] + Memory::Snapshottable(mem) => mem.maximum_byte_size(), } } @@ -471,6 +544,13 @@ impl Memory { return Ok(None); } } + #[cfg(target_os = "linux")] + Memory::Snapshottable(mem) => { + if let Err(e) = mem.grow_to(new_byte_size) { + store.memory_grow_failed(&e); + return Ok(None); + } + } } Ok(Some(old_byte_size)) } @@ -483,6 +563,8 @@ impl Memory { current_length: *size, }, Memory::Dynamic(mem) => mem.vmmemory(), + #[cfg(target_os = "linux")] + Memory::Snapshottable(mem) => mem.vmmemory(), } } @@ -507,6 +589,10 @@ impl Memory { Memory::Dynamic(_) => { unreachable!("dynamic memories should not have guard page faults") } + #[cfg(target_os = "linux")] + Memory::Snapshottable(_) => { + unreachable!("snapshottable memories should not have guard page faults") + } } } @@ -528,6 +614,10 @@ impl Memory { Memory::Dynamic(_) => { unreachable!("dynamic memories should not have guard page faults") } + #[cfg(target_os = "linux")] + Memory::Snapshottable(_) => { + unreachable!("snapshottable memories should not have guard page faults") + } } Ok(()) diff --git a/crates/runtime/src/mmap.rs b/crates/runtime/src/mmap.rs index 09b9c8843fc1..a11badd9386a 100644 --- a/crates/runtime/src/mmap.rs +++ b/crates/runtime/src/mmap.rs @@ -22,6 +22,9 @@ pub struct Mmap { ptr: usize, len: usize, file: Option, + + #[cfg(target_os = "linux")] + memfd: Option<(rustix::io::OwnedFd, usize, usize)>, } impl Mmap { @@ -35,6 +38,9 @@ impl Mmap { ptr: empty.as_ptr() as usize, len: 0, file: None, + + #[cfg(target_os = "linux")] + memfd: None, } } @@ -78,6 +84,9 @@ impl Mmap { ptr: ptr as usize, len, file: Some(file), + + #[cfg(target_os = "linux")] + memfd: None, }) } @@ -183,6 +192,9 @@ impl Mmap { ptr: ptr as usize, len: mapping_size, file: None, + + #[cfg(target_os = "linux")] + memfd: None, } } else { // Reserve the mapping size. @@ -200,6 +212,9 @@ impl Mmap { ptr: ptr as usize, len: mapping_size, file: None, + + #[cfg(target_os = "linux")] + memfd: None, }; if accessible_size != 0 { @@ -326,6 +341,224 @@ impl Mmap { Ok(()) } + /// Returns a page-aligned offset + length pair delimiting the memory pages which + /// are currently populated without generating any extraneous page faults. + #[cfg(target_os = "linux")] + fn populated_range( + &self, + accessible_offset: usize, + accessible: usize, + ) -> Result<(usize, usize)> { + // Docs: https://www.kernel.org/doc/Documentation/vm/pagemap.txt + use std::io::{Read, Seek}; + const PAGE_SIZE: usize = 4096; + + assert_eq!(rustix::process::page_size(), PAGE_SIZE); + + assert_eq!(accessible_offset % PAGE_SIZE, 0); + assert_eq!(accessible % PAGE_SIZE, 0); + + let mut page_last_index = 0; + let mut page_first_index = None; + unsafe { + let mut fp = std::fs::File::open("/proc/self/pagemap") + .context("failed to open /proc/self/pagemap")?; + + let offset = (self.as_ptr() as usize + accessible_offset) / PAGE_SIZE + * std::mem::size_of::(); + fp.seek(std::io::SeekFrom::Start(offset as u64)) + .context("failed to seek inside of /proc/self/pagemap")?; + + union Buffer { + as_u8: [u8; PAGE_SIZE], + as_u64: [u64; PAGE_SIZE / std::mem::size_of::()], + } + + let mut buffer: Buffer = std::mem::zeroed(); + + let total_page_count = accessible / PAGE_SIZE; + let mut current_page_offset = 0; + while current_page_offset < total_page_count { + let page_count = + std::cmp::min(buffer.as_u64.len(), total_page_count - current_page_offset); + fp.read(&mut buffer.as_u8[..page_count * std::mem::size_of::()]) + .context("failed to read from /proc/self/pagemap")?; + + for relative_page_index in 0..page_count { + let is_populated = (buffer.as_u64[relative_page_index] & (0b11 << 62)) != 0; + if is_populated { + page_last_index = current_page_offset + relative_page_index; + if page_first_index.is_none() { + page_first_index = Some(page_last_index); + } + } + } + + current_page_offset += page_count; + } + } + + if let Some(page_first_index) = page_first_index { + let (data_offset, data_length) = ( + accessible_offset + page_first_index * PAGE_SIZE, + (page_last_index - page_first_index + 1) * PAGE_SIZE, + ); + + assert!(data_offset + data_length <= self.len); + Ok((data_offset, data_length)) + } else { + Ok((accessible_offset, 0)) + } + } + + /// Saves a snapshot of the current contents of the mapping in an memfd, + /// and replaces that part of the mapping with a copy-on-write copy. + /// + /// Can only be used once. + #[cfg(target_os = "linux")] + pub fn create_snapshot(&mut self, accessible_offset: usize, accessible: usize) -> Result<()> { + assert!(self.memfd.is_none()); + assert!(accessible_offset + accessible_offset <= self.len); + + // Here we narrow down the exact range of memory which is populated. + // + // We could in theory not do this, however resetting copy-on-write + // pages is visibly slower (since that actually copies memory) than + // resetting those which are not copy-on-write, so we want to only + // snapshot as narrow of a memory region as possible. + let (data_offset, data_length) = self.populated_range(accessible_offset, accessible)?; + + debug_assert!(self.as_slice()[accessible_offset..data_offset] + .iter() + .all(|&byte| byte == 0)); + + debug_assert!( + self.as_slice()[data_offset + data_length..accessible_offset + accessible] + .iter() + .all(|&byte| byte == 0) + ); + + if data_length == 0 { + // Memory is completely empty, so no point in doing anything. + return Ok(()); + } + + unsafe { + let memfd = rustix::fs::memfd_create("wasmtime", rustix::fs::MemfdFlags::CLOEXEC) + .context("memfd_create failed")?; + + rustix::fs::ftruncate(&memfd, data_length as u64) + .context("failed to enlarge memfd: ftruncate failed")?; + + // In theory we could just map the memfd in memory and do a direct copy, + // but simply using `write` is going to have a lower overhead. + use rustix::fd::AsRawFd; + let bytes_written = libc::write( + memfd.as_raw_fd(), + (self.ptr as *const u8) + .add(data_offset) + .cast::(), + data_length, + ); + if bytes_written < 0 { + anyhow::bail!( + "failed to copy memory contents into memfd: write failed: {}", + std::io::Error::last_os_error() + ); + } + if bytes_written as usize != data_length { + anyhow::bail!("failed to copy memory contents into memfd: managed to write only {} bytes; expected {} bytes", bytes_written, data_length); + } + + rustix::io::mmap( + (self.ptr as *mut u8) + .add(data_offset) + .cast::(), + data_length, + rustix::io::ProtFlags::READ | rustix::io::ProtFlags::WRITE, + rustix::io::MapFlags::PRIVATE | rustix::io::MapFlags::FIXED, + &memfd, + 0, + ) + .context("failed to attach the memfd: mmap failed")?; + + self.memfd = Some((memfd, data_offset, data_length)); + }; + + Ok(()) + } + + /// Resets the memory contents within the given range. + /// + /// The memory will be filled with its original contents from + /// when [`Mmap::create_snapshot`] was called, or cleared + /// with zeros for non-memfd mappings. + #[cfg(target_os = "linux")] + pub unsafe fn reset(&mut self, offset: usize, length: usize) -> Result<()> { + rustix::io::madvise( + (self.ptr as *mut u8) + .add(offset) + .cast::(), + length, + rustix::io::Advice::LinuxDontNeed, + )?; + + Ok(()) + } + + /// Makes the memory within the given range completely inaccessible. + #[cfg(target_os = "linux")] + pub unsafe fn make_inaccessible(&mut self, offset: usize, length: usize) -> Result<()> { + rustix::io::mprotect( + (self.ptr as *mut u8) + .add(offset) + .cast::(), + length, + rustix::io::MprotectFlags::empty(), + )?; + Ok(()) + } + + /// Reallocates this mapping preserving its contents. + pub fn reallocate( + &mut self, + accessible_offset: usize, + old_accessible_size: usize, + new_mapping_size: usize, + new_accessible_size: usize, + ) -> Result<()> { + let mut new_mmap = Self::accessible_reserved(0, new_mapping_size)?; + + #[cfg(target_os = "linux")] + if let Some((ref memfd, data_offset, data_length)) = self.memfd { + unsafe { + rustix::io::mmap( + (new_mmap.ptr as *mut u8) + .add(data_offset) + .cast::(), + data_length, + rustix::io::ProtFlags::empty(), + rustix::io::MapFlags::PRIVATE | rustix::io::MapFlags::FIXED, + memfd, + 0, + ) + .context("failed to map the memfd: mmap failed")?; + } + } + + new_mmap.make_accessible(accessible_offset, new_accessible_size)?; + new_mmap.as_mut_slice()[accessible_offset..][..old_accessible_size] + .copy_from_slice(&self.as_slice()[accessible_offset..][..old_accessible_size]); + + #[cfg(target_os = "linux")] + { + new_mmap.memfd = self.memfd.take(); + } + + std::mem::swap(self, &mut new_mmap); + Ok(()) + } + /// Return the allocated memory as a slice of u8. pub fn as_slice(&self) -> &[u8] { unsafe { slice::from_raw_parts(self.ptr as *const u8, self.len) } @@ -450,3 +683,64 @@ fn _assert() { fn _assert_send_sync() {} _assert_send_sync::(); } + +#[cfg(target_os = "linux")] +#[cfg(test)] +mod tests { + use super::Mmap; + use anyhow::Result; + + // A few helper functions to make the tests easier to read. + fn pages(count: usize) -> usize { + count * 4096 + } + fn first_byte_of_page(count: usize) -> usize { + count * 4096 + } + fn last_byte_of_page(count: usize) -> usize { + count * 4096 + 4095 + } + + #[test] + fn create_snapshot_and_reset() -> Result<()> { + let mut mmap = Mmap::accessible_reserved(0, pages(4))?; + mmap.make_accessible(pages(1), pages(3))?; + mmap.as_mut_slice()[first_byte_of_page(1)] = 1; + mmap.create_snapshot(pages(1), pages(2))?; + + mmap.as_mut_slice()[first_byte_of_page(1)] = 10; + mmap.as_mut_slice()[first_byte_of_page(2)] = 100; + unsafe { + mmap.reset(pages(1), pages(3))?; + } + + assert_eq!(mmap.as_slice()[first_byte_of_page(1)], 1); + assert_eq!(mmap.as_slice()[first_byte_of_page(2)], 0); + Ok(()) + } + + #[test] + fn populated_range() -> Result<()> { + let mut mmap = Mmap::with_at_least(pages(16))?; + assert_eq!(mmap.populated_range(0, pages(16))?, (0, 0)); + assert_eq!(mmap.populated_range(pages(1), pages(15))?, (pages(1), 0)); + + mmap.as_mut_slice()[last_byte_of_page(1)] = 1; + assert_eq!(mmap.populated_range(0, pages(16))?, (pages(1), pages(1))); + + mmap.as_mut_slice()[first_byte_of_page(1)] = 1; + assert_eq!(mmap.populated_range(0, pages(16))?, (pages(1), pages(1))); + + mmap.as_mut_slice()[last_byte_of_page(3)] = 1; + assert_eq!(mmap.populated_range(0, pages(16))?, (pages(1), pages(3))); + + mmap.as_mut_slice()[first_byte_of_page(1)] = 1; + assert_eq!(mmap.populated_range(0, pages(16))?, (pages(1), pages(3))); + + assert_eq!( + mmap.populated_range(pages(2), pages(14))?, + (pages(3), pages(1)) + ); + Ok(()) + } +} diff --git a/crates/runtime/src/table.rs b/crates/runtime/src/table.rs index 5aed2333a6a4..95dba1444e8b 100644 --- a/crates/runtime/src/table.rs +++ b/crates/runtime/src/table.rs @@ -124,6 +124,9 @@ pub enum Table { ty: TableElementType, /// Maximum size that `elements` can grow to. maximum: Option, + + /// Fossilized storage for this table. + fossilized_elements: Vec, }, } @@ -147,6 +150,8 @@ impl Table { elements, ty, maximum, + + fossilized_elements: Vec::new(), }) } @@ -497,6 +502,38 @@ impl Table { } } } + + /// Saves a snapshot of the current table (size and contents) to be restored + /// in the future using `restore_snapshot`. + pub fn create_snapshot(&mut self) -> Result<()> { + match self { + Table::Static { .. } => bail!("static tables cannot be snapshotted"), + Table::Dynamic { + ref mut elements, + ref mut fossilized_elements, + .. + } => { + *fossilized_elements = elements.clone(); + Ok(()) + } + } + } + + /// Restores the table to a previously saved state. + pub fn restore_snapshot(&mut self) -> Result<()> { + match self { + Table::Static { .. } => bail!("static tables cannot be restored from snapshot"), + Table::Dynamic { + ref mut elements, + ref mut fossilized_elements, + .. + } => { + elements.clear(); + elements.extend_from_slice(fossilized_elements); + Ok(()) + } + } + } } impl Drop for Table { diff --git a/crates/wasmtime/src/instance.rs b/crates/wasmtime/src/instance.rs index aec6c1ba06f6..596e15805d53 100644 --- a/crates/wasmtime/src/instance.rs +++ b/crates/wasmtime/src/instance.rs @@ -15,8 +15,8 @@ use wasmtime_environ::{ }; use wasmtime_jit::TypeTables; use wasmtime_runtime::{ - Imports, InstanceAllocationRequest, InstantiationError, StorePtr, VMContext, VMFunctionBody, - VMFunctionImport, VMGlobalImport, VMMemoryImport, VMTableImport, + Imports, InstanceAllocationRequest, InstantiationError, MemorySource, StorePtr, VMContext, + VMFunctionBody, VMFunctionImport, VMGlobalImport, VMMemoryImport, VMTableImport, }; /// An instantiated WebAssembly module. @@ -124,7 +124,12 @@ impl Instance { let mut store = store.as_context_mut(); let mut i = unsafe { typecheck_externs(store.0, module, imports)?; - Instantiator::new(store.0, module, ImportSource::Externs(imports))? + Instantiator::new( + store.0, + module, + ImportSource::Externs(imports), + MemorySource::FromCreator, + )? }; assert!( !store.0.async_support(), @@ -165,7 +170,12 @@ impl Instance { let mut store = store.as_context_mut(); let mut i = unsafe { typecheck_externs(store.0, module, imports)?; - Instantiator::new(store.0, module, ImportSource::Externs(imports))? + Instantiator::new( + store.0, + module, + ImportSource::Externs(imports), + MemorySource::FromCreator, + )? }; let mut store = store.as_context_mut(); assert!( @@ -422,11 +432,27 @@ impl Instance { pub fn get_global(&self, store: impl AsContextMut, name: &str) -> Option { self.get_export(store, name)?.into_global() } + + /// Resets this instance's memory and tables to their initial state + /// from after the instantiation. + /// + /// Can only be used on an instance created through [`InstancePre::instantiate_reusable`]. + #[cfg(target_os = "linux")] + pub fn reset(&self, mut store: impl AsContextMut) -> Result<()> { + let store = store.as_context_mut().0; + let id = match &store[self.0] { + InstanceData::Synthetic(..) => unreachable!(), + InstanceData::Instantiated { id, .. } => *id, + }; + + store.instance_mut(id).restore_snapshot() + } } struct Instantiator<'a> { in_progress: Vec>, cur: ImportsBuilder<'a>, + memory_source: MemorySource, } struct ImportsBuilder<'a> { @@ -472,6 +498,7 @@ impl<'a> Instantiator<'a> { store: &StoreOpaque, module: &Module, imports: ImportSource<'a>, + memory_source: MemorySource, ) -> Result> { if !Engine::same(store.engine(), module.engine()) { bail!("cross-`Engine` instantiation is not currently supported"); @@ -480,6 +507,7 @@ impl<'a> Instantiator<'a> { Ok(Instantiator { in_progress: Vec::new(), cur: ImportsBuilder::new(module, imports), + memory_source, }) } @@ -490,6 +518,19 @@ impl<'a> Instantiator<'a> { Instantiator::start_raw(store, instance, start)?; } if toplevel { + match self.memory_source { + MemorySource::FromCreator => {} + #[cfg(target_os = "linux")] + MemorySource::CopyOnWriteInitialize => { + let id = match &store.0.store_data()[instance.0] { + InstanceData::Instantiated { id, .. } => *id, + InstanceData::Synthetic(_) => unreachable!(), + }; + + store.0.instance_mut(id).create_snapshot()?; + } + } + break Ok(instance); } } @@ -714,6 +755,7 @@ impl<'a> Instantiator<'a> { host_state: Box::new(Instance(instance_to_be)), store: StorePtr::new(store.traitobj()), wasm_data: compiled_module.wasm_data(), + memory_source: self.memory_source, })?; // The instance still has lots of setup, for example @@ -945,7 +987,42 @@ impl InstancePre { /// `store`, or if `store` has async support enabled. Additionally this /// function will panic if the `store` provided comes from a different /// [`Engine`] than the [`InstancePre`] originally came from. - pub fn instantiate(&self, mut store: impl AsContextMut) -> Result { + pub fn instantiate(&self, store: impl AsContextMut) -> Result { + self.instantiate_impl(store, MemorySource::FromCreator) + } + + /// Instantiates an instance which can be reset and reused with [`Instance::reset`]. + /// + /// Unlike normal instantiation it will not use the [`MemoryCreator`](crate::memory::MemoryCreator) registered + /// through [`Config::with_host_memory`](crate::config::Config::with_host_memory) + /// + /// Incompatible with any strategy other than [`InstanceAllocationStrategy::OnDemand`](crate::config::InstanceAllocationStrategy::OnDemand) + /// when set through [`Config::strategy`](crate::config::Config::strategy). + /// + /// Incompatible with imported memories and tables. + /// + /// For more information about instantiation see [`InstancePre::instantiate`]. + #[cfg(target_os = "linux")] + pub fn instantiate_reusable(&self, store: impl AsContextMut) -> Result { + for import in self.module.imports() { + match import.ty() { + ExternType::Table(..) => { + bail!("`instantiate_reusable` cannot be used for modules which import a table") + } + ExternType::Memory(..) => { + bail!("`instantiate_reusable` cannot be used for modules which import a memory") + } + _ => {} + } + } + self.instantiate_impl(store, MemorySource::CopyOnWriteInitialize) + } + + fn instantiate_impl( + &self, + mut store: impl AsContextMut, + memory_source: MemorySource, + ) -> Result { // For the unsafety here the typecheck happened at creation time of this // structure and then othrewise the `T` of `InstancePre` connects any // host functions we have in our definition list to the `store` that was @@ -957,6 +1034,7 @@ impl InstancePre { store.0, &self.module, ImportSource::Definitions(&self.items), + memory_source, )? }; instantiator.run(&mut store) @@ -989,6 +1067,7 @@ impl InstancePre { store.0, &self.module, ImportSource::Definitions(&self.items), + MemorySource::FromCreator, )? }; diff --git a/crates/wasmtime/src/store.rs b/crates/wasmtime/src/store.rs index 38bda4a734ef..96b5394634e7 100644 --- a/crates/wasmtime/src/store.rs +++ b/crates/wasmtime/src/store.rs @@ -91,7 +91,7 @@ use std::ptr; use std::sync::Arc; use std::task::{Context, Poll}; use wasmtime_runtime::{ - InstanceAllocationRequest, InstanceAllocator, InstanceHandle, ModuleInfo, + InstanceAllocationRequest, InstanceAllocator, InstanceHandle, MemorySource, ModuleInfo, OnDemandInstanceAllocator, SignalHandler, StorePtr, VMCallerCheckedAnyfunc, VMContext, VMExternRef, VMExternRefActivationsTable, VMInterrupts, VMSharedSignatureIndex, VMTrampoline, }; @@ -409,6 +409,7 @@ impl Store { module: Arc::new(wasmtime_environ::Module::default()), store: StorePtr::empty(), wasm_data: &[], + memory_source: MemorySource::FromCreator, }) .expect("failed to allocate default callee") }; diff --git a/crates/wasmtime/src/trampoline.rs b/crates/wasmtime/src/trampoline.rs index c1f8038a5ace..b0dbd66f82c9 100644 --- a/crates/wasmtime/src/trampoline.rs +++ b/crates/wasmtime/src/trampoline.rs @@ -18,8 +18,8 @@ use std::any::Any; use std::sync::Arc; use wasmtime_environ::{EntityIndex, GlobalIndex, MemoryIndex, Module, TableIndex}; use wasmtime_runtime::{ - Imports, InstanceAllocationRequest, InstanceAllocator, OnDemandInstanceAllocator, StorePtr, - VMFunctionImport, VMSharedSignatureIndex, + Imports, InstanceAllocationRequest, InstanceAllocator, MemorySource, OnDemandInstanceAllocator, + StorePtr, VMFunctionImport, VMSharedSignatureIndex, }; fn create_handle( @@ -48,6 +48,7 @@ fn create_handle( host_state, store: StorePtr::new(store.traitobj()), wasm_data: &[], + memory_source: MemorySource::FromCreator, }, )?; diff --git a/crates/wasmtime/src/trampoline/func.rs b/crates/wasmtime/src/trampoline/func.rs index fc217377beab..40930c1ea63b 100644 --- a/crates/wasmtime/src/trampoline/func.rs +++ b/crates/wasmtime/src/trampoline/func.rs @@ -8,7 +8,7 @@ use std::sync::Arc; use wasmtime_environ::{EntityIndex, Module, ModuleType, PrimaryMap, SignatureIndex}; use wasmtime_jit::{CodeMemory, MmapVec}; use wasmtime_runtime::{ - Imports, InstanceAllocationRequest, InstanceAllocator, InstanceHandle, + Imports, InstanceAllocationRequest, InstanceAllocator, InstanceHandle, MemorySource, OnDemandInstanceAllocator, StorePtr, VMContext, VMFunctionBody, VMSharedSignatureIndex, VMTrampoline, }; @@ -134,6 +134,7 @@ pub unsafe fn create_raw_function( host_state, store: StorePtr::empty(), wasm_data: &[], + memory_source: MemorySource::FromCreator, })?, ) } diff --git a/tests/all/main.rs b/tests/all/main.rs index b1ec333981a5..06bf3b9565d2 100644 --- a/tests/all/main.rs +++ b/tests/all/main.rs @@ -25,6 +25,8 @@ mod module_serialize; mod name; mod pooling_allocator; mod relocs; +#[cfg(target_os = "linux")] +mod reusable; mod stack_overflow; mod store; mod table; diff --git a/tests/all/reusable.rs b/tests/all/reusable.rs new file mode 100644 index 000000000000..d0e1bc24efce --- /dev/null +++ b/tests/all/reusable.rs @@ -0,0 +1,577 @@ +use anyhow::Result; +use std::sync::Arc; +use wasmtime::*; + +fn dynamic_memory_config() -> Config { + let mut config = Config::new(); + config.static_memory_maximum_size(0); + config.dynamic_memory_reserved_for_growth(0); + config +} + +fn clears_memory_on_reset(config: &Config) -> Result<()> { + let engine = Engine::new(&config)?; + let linker = Linker::new(&engine); + + let module = Module::new( + &engine, + r#"(module + (memory $0 1) + (export "write_memory" (func $write_memory)) + (export "read_memory" (func $read_memory)) + (func $write_memory + (i32.store offset=1024 + (i32.const 0) + (i32.const 10) + ) + ) + (func $read_memory (result i32) + (i32.load offset=1024 (i32.const 0)) + ) + )"#, + )?; + + let mut store = Store::new(&engine, ()); + let instance_pre = linker.instantiate_pre(&mut store, &module)?; + let instance = instance_pre.instantiate_reusable(&mut store)?; + + let read_memory = instance + .get_func(&mut store, "read_memory") + .unwrap() + .typed::<(), i32, _>(&store)?; + let write_memory = instance + .get_func(&mut store, "write_memory") + .unwrap() + .typed::<(), (), _>(&store)?; + + // Modify the memory and make sure it's actually modified. + assert_eq!(read_memory.call(&mut store, ())?, 0); + write_memory.call(&mut store, ())?; + assert_eq!(read_memory.call(&mut store, ())?, 10); + + // Reset the instance and make sure the memory was actually cleared. + instance.reset(&mut store)?; + assert_eq!(read_memory.call(&mut store, ())?, 0); + + // Do it all over again to make sure the instance can be reset multiple times. + write_memory.call(&mut store, ())?; + assert_eq!(read_memory.call(&mut store, ())?, 10); + instance.reset(&mut store)?; + assert_eq!(read_memory.call(&mut store, ())?, 0); + + Ok(()) +} + +#[test] +fn clears_memory_on_reset_static_memory() -> Result<()> { + clears_memory_on_reset(&Config::new()) +} + +#[test] +fn clears_memory_on_reset_dynamic_memory() -> Result<()> { + clears_memory_on_reset(&dynamic_memory_config()) +} + +fn restores_original_data_segment_on_reset(config: &Config) -> Result<()> { + let engine = Engine::new(&config)?; + let linker = Linker::new(&engine); + + let module = Module::new( + &engine, + r#"(module + (memory $0 1) + (data $.data (i32.const 1024) "A\00\00\00") + (export "write_memory" (func $write_memory)) + (export "grow_memory" (func $grow_memory)) + (export "read_memory" (func $read_memory)) + (func $write_memory + (i32.store offset=1024 + (i32.const 0) + (i32.const 66) + ) + ) + (func $grow_memory (param $0 i32) (result i32) + (memory.grow + (local.get $0) + ) + ) + (func $read_memory (result i32) + (i32.load offset=1024 (i32.const 0)) + ) + )"#, + )?; + + let mut store = Store::new(&engine, ()); + let instance_pre = linker.instantiate_pre(&mut store, &module)?; + let instance = instance_pre.instantiate_reusable(&mut store)?; + + let grow_memory = instance + .get_func(&mut store, "grow_memory") + .unwrap() + .typed::(&store)?; + let read_memory = instance + .get_func(&mut store, "read_memory") + .unwrap() + .typed::<(), i32, _>(&store)?; + let write_memory = instance + .get_func(&mut store, "write_memory") + .unwrap() + .typed::<(), (), _>(&store)?; + + // Modify the memory and make sure it's actually modified. + assert_eq!(read_memory.call(&mut store, ())?, 65); + write_memory.call(&mut store, ())?; + assert_eq!(read_memory.call(&mut store, ())?, 66); + + // Reset the memory and make sure it was actually reset to its initial value. + instance.reset(&mut store)?; + assert_eq!(read_memory.call(&mut store, ())?, 65); + + // Do it all over again to make sure the instance can be reset multiple times. + write_memory.call(&mut store, ())?; + assert_eq!(read_memory.call(&mut store, ())?, 66); + instance.reset(&mut store)?; + assert_eq!(read_memory.call(&mut store, ())?, 65); + + // Do it once again, but this time grow the memory after writing to it. + // + // If the memory is dynamic this will reallocate it from scratch somewhere + // else in memory, so we want to make sure that it still can be reset. + write_memory.call(&mut store, ())?; + assert_eq!(grow_memory.call(&mut store, 1)?, 1); + assert_eq!(read_memory.call(&mut store, ())?, 66); + instance.reset(&mut store)?; + assert_eq!(read_memory.call(&mut store, ())?, 65); + + Ok(()) +} + +#[test] +fn restores_original_data_segment_on_reset_static_memory() -> Result<()> { + restores_original_data_segment_on_reset(&Config::new()) +} + +#[test] +fn restores_original_data_segment_on_reset_dynamic_memory() -> Result<()> { + restores_original_data_segment_on_reset(&dynamic_memory_config()) +} + +fn restores_memory_size_on_reset_after_grow(config: &Config) -> Result<()> { + let engine = Engine::new(&config)?; + let linker = Linker::new(&engine); + + let module = Module::new( + &engine, + r#"(module + (memory $0 1) + (export "write_memory" (func $write_memory)) + (export "grow_memory" (func $grow_memory)) + (export "read_memory" (func $read_memory)) + (func $write_memory + (i32.store offset=65536 + (i32.const 0) + (i32.const 10) + ) + ) + (func $grow_memory (param $0 i32) (result i32) + (memory.grow + (local.get $0) + ) + ) + (func $read_memory (result i32) + (i32.load offset=65536 (i32.const 0)) + ) + )"#, + )?; + + let mut store = Store::new(&engine, ()); + let instance_pre = linker.instantiate_pre(&mut store, &module)?; + let instance = instance_pre.instantiate_reusable(&mut store)?; + + let grow_memory = instance + .get_func(&mut store, "grow_memory") + .unwrap() + .typed::(&store)?; + let write_memory = instance + .get_func(&mut store, "write_memory") + .unwrap() + .typed::<(), (), _>(&store)?; + let read_memory = instance + .get_func(&mut store, "read_memory") + .unwrap() + .typed::<(), i32, _>(&store)?; + + // The memory should initially be one WASM page big. + assert_eq!(grow_memory.call(&mut store, 0)?, 1); + assert_eq!( + read_memory + .call(&mut store, ()) + .unwrap_err() + .trap_code() + .unwrap(), + TrapCode::MemoryOutOfBounds + ); + assert_eq!( + write_memory + .call(&mut store, ()) + .unwrap_err() + .trap_code() + .unwrap(), + TrapCode::MemoryOutOfBounds + ); + + // ...then we grow it, which should make it accessible. + assert_eq!(grow_memory.call(&mut store, 1)?, 1); + + // Make sure that we can access it and that it's actually zero'd. + assert_eq!(read_memory.call(&mut store, ())?, 0); + + // Now we can dirty it. + write_memory.call(&mut store, ())?; + assert_eq!(read_memory.call(&mut store, ())?, 10); + + // Then we reset it to its initial state (both its size and its contents). + instance.reset(&mut store)?; + + // Make sure that the memory size was reset, and that it isn't actually accessible. + assert_eq!(grow_memory.call(&mut store, 0)?, 1); + assert_eq!( + read_memory + .call(&mut store, ()) + .unwrap_err() + .trap_code() + .unwrap(), + TrapCode::MemoryOutOfBounds + ); + assert_eq!( + write_memory + .call(&mut store, ()) + .unwrap_err() + .trap_code() + .unwrap(), + TrapCode::MemoryOutOfBounds + ); + + // Now we can grow it again. + assert_eq!(grow_memory.call(&mut store, 1)?, 1); + + // Let's make sure the old value doesn't linger (that is - that the memory wasn't simply only made unaccessible). + assert_eq!(read_memory.call(&mut store, ())?, 0); + + // Now make sure we can write to it again and read the value back. + write_memory.call(&mut store, ())?; + assert_eq!(read_memory.call(&mut store, ())?, 10); + + Ok(()) +} + +#[test] +fn restores_memory_size_on_reset_after_grow_static_memory() -> Result<()> { + restores_memory_size_on_reset_after_grow(&Config::new()) +} + +#[test] +fn restores_memory_size_on_reset_after_grow_dynamic_memory() -> Result<()> { + restores_memory_size_on_reset_after_grow(&dynamic_memory_config()) +} + +#[test] +fn restores_table_on_reset() -> Result<()> { + let engine = Engine::default(); + let linker = Linker::new(&engine); + + let module = Module::new( + &engine, + r#"(module + (type $none_=>_i32 (func (result i32))) + (memory $0 1) + (table $table 1 2 funcref) + (export "return_100" (func $return_100)) + (export "return_200" (func $return_200)) + (export "call_table" (func $call_table)) + (export "table" (table $table)) + (elem (i32.const 0) $return_100) + (func $call_table (result i32) + (call_indirect (type $none_=>_i32) + (i32.const 0) + ) + ) + (func $return_100 (result i32) + (i32.const 100) + ) + (func $return_200 (result i32) + (i32.const 200) + ) + )"#, + )?; + + let mut store = Store::new(&engine, ()); + let instance_pre = linker.instantiate_pre(&mut store, &module)?; + let instance = instance_pre.instantiate_reusable(&mut store)?; + + let call_table = instance + .get_func(&mut store, "call_table") + .unwrap() + .typed::<(), i32, _>(&store)?; + let return_200 = instance.get_func(&mut store, "return_200").unwrap(); + let table = instance.get_table(&mut store, "table").unwrap(); + + // Sanity checks on the initial state. + assert_eq!(call_table.call(&mut store, ())?, 100); + assert_eq!(table.size(&mut store), 1); + + // Replace the function in the table, and grow the table. + table.set(&mut store, 0, return_200.into())?; + assert_eq!(call_table.call(&mut store, ())?, 200); + assert_eq!(table.grow(&mut store, 1, return_200.into())?, 1); + + // Reset the instance, and make sure the table was restored (both its size and its contents). + instance.reset(&mut store)?; + assert_eq!(call_table.call(&mut store, ())?, 100); + assert_eq!(table.size(&mut store), 1); + + Ok(()) +} + +#[test] +fn snapshot_is_taken_after_start() -> Result<()> { + let engine = Engine::default(); + let linker = Linker::new(&engine); + + let module = Module::new( + &engine, + r#"(module + (memory $0 1) + (export "write_memory" (func $write_memory)) + (export "read_memory" (func $read_memory)) + (start $start) + (func $write_memory + (i32.store offset=1024 + (i32.const 0) + (i32.const 10) + ) + ) + (func $read_memory (result i32) + (i32.load offset=1024 (i32.const 0)) + ) + (func $start + (i32.store offset=1024 + (i32.const 0) + (i32.const 5) + ) + ) + )"#, + )?; + + let mut store = Store::new(&engine, ()); + let instance_pre = linker.instantiate_pre(&mut store, &module)?; + let instance = instance_pre.instantiate_reusable(&mut store)?; + + let read_memory = instance + .get_func(&mut store, "read_memory") + .unwrap() + .typed::<(), i32, _>(&store)?; + let write_memory = instance + .get_func(&mut store, "write_memory") + .unwrap() + .typed::<(), (), _>(&store)?; + + // Make sure the start function was actually run. + assert_eq!(read_memory.call(&mut store, ())?, 5); + + // Overwrite the memory. + write_memory.call(&mut store, ())?; + assert_eq!(read_memory.call(&mut store, ())?, 10); + + // Reset the memory and make sure it was restored to what the start function wrote. + instance.reset(&mut store)?; + assert_eq!(read_memory.call(&mut store, ())?, 5); + + Ok(()) +} + +#[test] +fn globals_are_restored_on_reset() -> Result<()> { + let engine = Engine::default(); + let linker = Linker::new(&engine); + + let module = Module::new( + &engine, + r#"(module + (memory $0 1) + (global $1 (mut i32) (i32.const 123)) + (export "write_global" (func $write_global)) + (export "read_global" (func $read_global)) + (func $write_global + (global.set $1 (i32.const 124)) + ) + (func $read_global (result i32) + (global.get $1) + ) + )"#, + )?; + + let mut store = Store::new(&engine, ()); + let instance_pre = linker.instantiate_pre(&mut store, &module)?; + let instance = instance_pre.instantiate_reusable(&mut store)?; + + let read_global = instance + .get_func(&mut store, "read_global") + .unwrap() + .typed::<(), i32, _>(&store)?; + let write_global = instance + .get_func(&mut store, "write_global") + .unwrap() + .typed::<(), (), _>(&store)?; + + // Make sure the global was properly initialized. + assert_eq!(read_global.call(&mut store, ())?, 123); + + // Overwrite the global. + write_global.call(&mut store, ())?; + assert_eq!(read_global.call(&mut store, ())?, 124); + + // Reset the instance and make sure the global was also restored. + instance.reset(&mut store)?; + assert_eq!(read_global.call(&mut store, ())?, 123); + + Ok(()) +} + +#[test] +fn non_reusable_instance_cannot_be_reset() -> Result<()> { + let engine = Engine::default(); + let linker = Linker::new(&engine); + + let module = Module::new(&engine, r#"(module)"#)?; + + let mut store = Store::new(&engine, ()); + let instance_pre = linker.instantiate_pre(&mut store, &module)?; + let instance = instance_pre.instantiate(&mut store)?; + + assert!(instance.reset(&mut store).is_err()); + + Ok(()) +} + +#[test] +fn reusable_instances_are_not_compatible_with_instance_pooling_strategy() -> Result<()> { + let mut config = Config::new(); + config.allocation_strategy(InstanceAllocationStrategy::Pooling { + strategy: Default::default(), + module_limits: Default::default(), + instance_limits: Default::default(), + }); + + let engine = Engine::new(&config)?; + let linker = Linker::new(&engine); + + let module = Module::new(&engine, r#"(module)"#)?; + + let mut store = Store::new(&engine, ()); + let instance_pre = linker.instantiate_pre(&mut store, &module)?; + assert!(instance_pre.instantiate_reusable(&mut store).is_err()); + + Ok(()) +} + +struct DummyMemoryCreator; + +unsafe impl MemoryCreator for DummyMemoryCreator { + fn new_memory( + &self, + _ty: MemoryType, + _minimum: usize, + _maximum: Option, + _reserved_size_in_bytes: Option, + _guard_size_in_bytes: usize, + ) -> Result, String> { + unimplemented!(); + } +} + +#[test] +fn reusable_instances_do_not_use_a_custom_memory_creator() -> Result<()> { + let mut config = Config::new(); + config.with_host_memory(Arc::new(DummyMemoryCreator)); + + let engine = Engine::new(&config)?; + let linker = Linker::new(&engine); + + let module = Module::new( + &engine, + r#"(module + (memory $0 1) + (export "grow_memory" (func $grow_memory)) + (func $grow_memory (param $0 i32) (result i32) + (memory.grow + (local.get $0) + ) + ) + )"#, + )?; + + let mut store = Store::new(&engine, ()); + let instance_pre = linker.instantiate_pre(&mut store, &module)?; + let instance = instance_pre.instantiate_reusable(&mut store)?; + + let grow_memory = instance + .get_func(&mut store, "grow_memory") + .unwrap() + .typed::(&store)?; + + assert_eq!(grow_memory.call(&mut store, 1)?, 1); + assert_eq!(grow_memory.call(&mut store, 1)?, 2); + + instance.reset(&mut store)?; + + Ok(()) +} + +#[test] +fn incompatible_with_imported_memories() -> Result<()> { + let engine = Engine::default(); + let mut linker = Linker::new(&engine); + + let module = Module::new( + &engine, + r#"(module + (import "env" "memory" (memory $memoryimport 1)) + )"#, + )?; + + let mut store = Store::new(&engine, ()); + let memory_ty = MemoryType::new(1, None); + let memory = Memory::new(&mut store, memory_ty)?; + linker.define("env", "memory", memory)?; + + let instance_pre = linker.instantiate_pre(&mut store, &module)?; + assert!(instance_pre.instantiate(&mut store).is_ok()); + assert!(instance_pre.instantiate_reusable(&mut store).is_err()); + + Ok(()) +} + +#[test] +fn incompatible_with_imported_tables() -> Result<()> { + let engine = Engine::default(); + let mut linker = Linker::new(&engine); + + let module = Module::new( + &engine, + r#"(module + (import "env" "table" (table $tableimport 1 2 funcref)) + )"#, + )?; + + let mut store = Store::new(&engine, ()); + let table_ty = TableType::new(ValType::FuncRef, 1, Some(2)); + let table = Table::new(&mut store, table_ty, Val::FuncRef(None))?; + linker.define("env", "table", table)?; + + let instance_pre = linker.instantiate_pre(&mut store, &module)?; + assert!(instance_pre.instantiate(&mut store).is_ok()); + assert!(instance_pre.instantiate_reusable(&mut store).is_err()); + + Ok(()) +}