diff --git a/crates/tinywasm/src/engine.rs b/crates/tinywasm/src/engine.rs index 76acc19..7cb1d9a 100644 --- a/crates/tinywasm/src/engine.rs +++ b/crates/tinywasm/src/engine.rs @@ -1,7 +1,7 @@ use alloc::boxed::Box; +use tinywasm_types::Shared; use crate::ResourceLimiter; -use crate::shared::StoreShared; /// Global configuration for the WebAssembly interpreter /// @@ -122,7 +122,7 @@ pub struct Config { /// Fuel accounting policy used by budgeted execution. Defaults to [`FuelPolicy::PerInstruction`]. pub fuel_policy: FuelPolicy, /// Resource limiter shared across all stores created from this engine. Defaults to `None`. - pub resource_limiter: Option>, + pub resource_limiter: Option>, /// Initial number of GC heap bytes that triggers collection. /// Defaults to 1 MiB. pub gc_collection_threshold: usize, @@ -177,7 +177,7 @@ impl Config { /// The limiter is converted to TinyWasm's internal shared pointer. Pass the /// limiter value directly rather than wrapping it in `Rc` or `Arc`. pub fn with_resource_limiter(mut self, limiter: impl ResourceLimiter + 'static) -> Self { - self.resource_limiter = Some(StoreShared::from(Box::new(limiter) as Box)); + self.resource_limiter = Some(Shared::from(Box::new(limiter) as Box)); self } diff --git a/crates/tinywasm/src/instance.rs b/crates/tinywasm/src/instance.rs index f92765d..aedcef6 100644 --- a/crates/tinywasm/src/instance.rs +++ b/crates/tinywasm/src/instance.rs @@ -215,8 +215,7 @@ impl ModuleInstance { addrs.tags.extend(store.init_tags(&module.tags, &type_addrs)); let limiter = store.engine.config().resource_limiter.clone(); if !module.skip_local_memory_allocation { - let memories = - store.init_memories(&module.memory_types, |ty| MemoryInstance::new(ty, limiter.as_deref()))?; + let memories = store.init_memories(&module.memory_types, |ty| MemoryInstance::new(ty, limiter.clone()))?; addrs.memories.extend(memories); } diff --git a/crates/tinywasm/src/interpreter/executor/mod.rs b/crates/tinywasm/src/interpreter/executor/mod.rs index 4f2c27f..e6241a5 100644 --- a/crates/tinywasm/src/interpreter/executor/mod.rs +++ b/crates/tinywasm/src/interpreter/executor/mod.rs @@ -1268,7 +1268,7 @@ impl<'store> Executor<'store> { fn exec_memory_grow(&mut self, addr: u32) -> Result<(), Trap> { let mem_addr = self.mem_addr(addr); - let limiter = self.store.engine.config().resource_limiter.as_deref(); + let limiter = self.store.engine.config().resource_limiter.as_ref(); let is_64bit = self.store.state.memory_type(mem_addr).arch() == MemoryArch::I64; let pages_delta = match is_64bit { true => i64::stack_pop(&mut self.store.value_stack), diff --git a/crates/tinywasm/src/reference/store.rs b/crates/tinywasm/src/reference/store.rs index b0771ef..f35a64f 100644 --- a/crates/tinywasm/src/reference/store.rs +++ b/crates/tinywasm/src/reference/store.rs @@ -386,7 +386,7 @@ impl Memory { return Err(Error::UnsupportedFeature("too many memories")); } let limiter = store.engine.config().resource_limiter.clone(); - store.state.memories.push(MemoryInstance::new(ty, limiter.as_deref())?); + store.state.memories.push(MemoryInstance::new(ty, limiter)?); Ok(Self(StoreItem::new(store.id(), addr))) } @@ -487,9 +487,8 @@ impl Memory { /// Returns the previous size, or `None` if growth fails or is rejected by the resource limiter. /// A limiter-provided trap is returned as an error. pub fn grow(&self, store: &mut Store, delta_pages: i64) -> Result> { - let limiter = store.engine.config().resource_limiter.clone(); let mem = self.instance_mut(store)?; - mem.grow(delta_pages, limiter.as_deref()).map_err(Into::into) + mem.grow(delta_pages).map_err(Into::into) } /// Get the current size of the memory in pages. diff --git a/crates/tinywasm/src/shared.rs b/crates/tinywasm/src/shared.rs index 5fe65f2..81f6019 100644 --- a/crates/tinywasm/src/shared.rs +++ b/crates/tinywasm/src/shared.rs @@ -4,6 +4,6 @@ pub(crate) use alloc::rc::Rc as StoreShared; pub(crate) use tinywasm_types::Shared as StoreShared; #[cfg(not(feature = "portable-atomic"))] -pub(crate) use core::sync::atomic::{AtomicU32, Ordering, fence}; +pub(crate) use core::sync::atomic::{AtomicU32, AtomicUsize, Ordering, fence}; #[cfg(feature = "portable-atomic")] -pub(crate) use portable_atomic::{AtomicU32, Ordering, fence}; +pub(crate) use portable_atomic::{AtomicU32, AtomicUsize, Ordering, fence}; diff --git a/crates/tinywasm/src/store/gc/object.rs b/crates/tinywasm/src/store/gc/object.rs index c76d7e9..a292fe6 100644 --- a/crates/tinywasm/src/store/gc/object.rs +++ b/crates/tinywasm/src/store/gc/object.rs @@ -1,11 +1,10 @@ use alloc::{boxed::Box, vec::Vec}; use core::mem::size_of; -use tinywasm_types::{TagAddr, TypeAddr}; +use tinywasm_types::{Shared, TagAddr, TypeAddr}; use crate::engine::Config; use crate::interpreter::{RuntimeValue, ValueRef}; -use crate::shared::StoreShared; use crate::{ResourceLimiter, Trap}; use super::{AllocError, Arena, Handle, Trace}; @@ -34,7 +33,7 @@ pub(crate) struct GcHeap { objects: Arena, directory: Vec<(u32, Handle)>, next_key: u32, - resource_limiter: Option>, + resource_limiter: Option>, } impl Default for GcHeap { diff --git a/crates/tinywasm/src/store/memory/instance.rs b/crates/tinywasm/src/store/memory/instance.rs index 0222cc0..8c990cc 100644 --- a/crates/tinywasm/src/store/memory/instance.rs +++ b/crates/tinywasm/src/store/memory/instance.rs @@ -1,5 +1,7 @@ +use tinywasm_types::Shared; use tinywasm_types::{MemoryArch, MemoryType}; +use crate::shared::{AtomicUsize, Ordering}; use crate::{Error, ResourceLimiter, Result, Trap}; use super::{MemoryStorage, memory_oob}; @@ -13,6 +15,37 @@ pub(crate) struct MemoryInstance { pub(crate) kind: MemoryType, pub(crate) inner: MemoryStorage, pub(crate) page_count: usize, + // Fields drop in declaration order: release the charge after the backing storage. + pub(super) charge: Option, +} + +/// The part of a memory's logical size approved by one limiter. +pub(super) struct MemoryCharge { + limiter: Shared, + bytes: AtomicUsize, +} + +impl MemoryCharge { + pub(super) fn new(limiter: Shared, bytes: usize) -> Self { + Self { limiter, bytes: AtomicUsize::new(bytes) } + } + + pub(super) fn limiter(&self) -> &dyn ResourceLimiter { + self.limiter.as_ref() + } + + pub(super) fn add(&self, bytes: usize) { + self.bytes.fetch_add(bytes, Ordering::Relaxed); + } +} + +impl Drop for MemoryCharge { + fn drop(&mut self) { + let bytes = self.bytes.load(Ordering::Relaxed); + if bytes != 0 { + self.limiter.memory_dropped(bytes); + } + } } #[cfg(feature = "debug")] @@ -44,7 +77,15 @@ impl MemoryInstance { } } - pub(crate) fn new(kind: MemoryType, limiter: Option<&dyn ResourceLimiter>) -> Result { + pub(crate) fn new(kind: MemoryType, limiter: Option>) -> Result { + Self::new_with_storage(kind, limiter, MemoryStorage::try_new) + } + + fn new_with_storage( + kind: MemoryType, + limiter: Option>, + allocate: impl FnOnce(MemoryArch, usize, usize) -> core::result::Result, + ) -> Result { if kind.shared() && kind.page_count_max_declared().is_none() { return Err(Error::UnsupportedFeature("shared memory requires a maximum")); } @@ -71,15 +112,30 @@ impl MemoryInstance { ); if initial_len != 0 - && let Some(limiter) = limiter + && let Some(limiter) = limiter.as_deref() && !limiter.memory_growing(0, initial_len, Self::maximum_size(kind))? { return cold!(Err(Trap::OutOfMemory.into())); } let max_len = Self::host_size(kind, max_pages).unwrap_or(usize::MAX); - let storage = MemoryStorage::try_new(kind.arch(), initial_len, max_len)?; - Ok(Self { kind, inner: storage, page_count: kind.page_count_initial() as usize }) + let storage = match allocate(kind.arch(), initial_len, max_len) { + Ok(storage) => storage, + Err(error) => { + if initial_len != 0 + && let Some(limiter) = limiter.as_deref() + { + limiter.memory_grow_failed(0, initial_len); + } + return Err(error.into()); + } + }; + Ok(Self { + kind, + inner: storage, + page_count: kind.page_count_initial() as usize, + charge: limiter.map(|limiter| MemoryCharge::new(limiter, initial_len)), + }) } pub(crate) fn copy_from_memory( @@ -96,21 +152,41 @@ impl MemoryInstance { cold_err!(self.inner.copy_within(dst, src, len).ok_or_else(|| memory_oob(dst, len, self.inner.len()))) } - pub(crate) fn grow( - &mut self, + pub(crate) fn grow(&mut self, pages_delta: i64) -> Result, Trap> { + let before = self.inner.len(); + let result = Self::grow_storage( + self.kind, + &mut self.inner, + &mut self.page_count, + pages_delta, + self.charge.as_ref().map(MemoryCharge::limiter), + )?; + if result.is_some() + && let Some(charge) = &self.charge + { + charge.add(self.inner.len() - before); + } + Ok(result) + } + + /// Grows exclusively borrowed storage after checking limits and the host limiter. + pub(super) fn grow_storage( + kind: MemoryType, + inner: &mut MemoryStorage, + page_count: &mut usize, pages_delta: i64, limiter: Option<&dyn ResourceLimiter>, ) -> Result, Trap> { - Self::grow_storage(self.kind, &mut self.inner, &mut self.page_count, pages_delta, limiter) + Self::grow_storage_with(kind, inner, page_count, pages_delta, limiter, MemoryStorage::grow_to) } - /// Grows exclusively borrowed storage after checking limits and the host limiter. - pub(super) fn grow_storage( + fn grow_storage_with( kind: MemoryType, inner: &mut MemoryStorage, page_count: &mut usize, pages_delta: i64, limiter: Option<&dyn ResourceLimiter>, + grow: impl FnOnce(&mut MemoryStorage, usize) -> core::result::Result<(), Trap>, ) -> Result, Trap> { let current_pages = *page_count; let Some(new_pages) = usize::try_from(pages_delta).ok().and_then(|delta| current_pages.checked_add(delta)) @@ -129,20 +205,86 @@ impl MemoryInstance { let Some(new_size) = Self::host_size(kind, new_pages as u64) else { return cold!(Ok(None)); }; - if new_size == inner.len() { + let current_size = inner.len(); + if new_size == current_size { return Ok(i64::try_from(current_pages).ok()); } if let Some(limiter) = limiter - && !limiter.memory_growing(inner.len(), new_size, Self::maximum_size(kind))? + && !limiter.memory_growing(current_size, new_size, Self::maximum_size(kind))? { return cold!(Ok(None)); } - if inner.grow_to(new_size).is_err() { + if grow(inner, new_size).is_err() { + if let Some(limiter) = limiter { + limiter.memory_grow_failed(current_size, new_size); + } return cold!(Ok(None)); } *page_count = new_pages; Ok(i64::try_from(current_pages).ok()) } } + +#[cfg(test)] +mod tests { + use alloc::boxed::Box; + use tinywasm_types::Shared; + + use super::*; + use crate::shared::Ordering; + + struct ReservingLimiter { + used: Shared, + } + + impl ResourceLimiter for ReservingLimiter { + fn memory_growing(&self, current: usize, desired: usize, _maximum: Option) -> Result { + self.used.fetch_add(desired - current, Ordering::SeqCst); + Ok(true) + } + + fn memory_grow_failed(&self, current: usize, desired: usize) { + self.used.fetch_sub(desired - current, Ordering::SeqCst); + } + + fn memory_dropped(&self, charged_bytes: usize) { + self.used.fetch_sub(charged_bytes, Ordering::SeqCst); + } + } + + fn limiter(used: &Shared) -> Shared { + Shared::from(Box::new(ReservingLimiter { used: used.clone() }) as Box) + } + + #[test] + fn failed_initial_allocation_refunds_approved_reservation() { + let used = Shared::new(AtomicUsize::new(0)); + let ty = MemoryType::new(MemoryArch::I32, 1, None, None); + let result = MemoryInstance::new_with_storage(ty, Some(limiter(&used)), |_, _, _| Err(Trap::OutOfMemory)); + assert!(matches!(result, Err(Error::Trap(Trap::OutOfMemory)))); + assert_eq!(used.load(Ordering::SeqCst), 0); + } + + #[test] + fn failed_growth_refunds_only_the_attempted_delta() { + let used = Shared::new(AtomicUsize::new(0)); + let ty = MemoryType::new(MemoryArch::I32, 1, None, None); + let mut memory = MemoryInstance::new(ty, Some(limiter(&used))).unwrap(); + assert_eq!(used.load(Ordering::SeqCst), 65_536); + let result = MemoryInstance::grow_storage_with( + memory.kind, + &mut memory.inner, + &mut memory.page_count, + 1, + memory.charge.as_ref().map(MemoryCharge::limiter), + |_, _| Err(Trap::OutOfMemory), + ); + assert_eq!(result.unwrap(), None); + assert_eq!(memory.page_count, 1); + assert_eq!(used.load(Ordering::SeqCst), 65_536); + drop(memory); + assert_eq!(used.load(Ordering::SeqCst), 0); + } +} diff --git a/crates/tinywasm/src/store/memory/shared.rs b/crates/tinywasm/src/store/memory/shared.rs index fa68a47..be99f07 100644 --- a/crates/tinywasm/src/store/memory/shared.rs +++ b/crates/tinywasm/src/store/memory/shared.rs @@ -1,4 +1,4 @@ -use crate::std::sync::{Condvar, Mutex, MutexGuard}; +use crate::std::sync::{Condvar, Mutex, MutexGuard, OnceLock}; use crate::std::time::Duration; use alloc::{ collections::{BTreeMap, VecDeque}, @@ -7,7 +7,9 @@ use alloc::{ }; use core::sync::atomic::{AtomicUsize, Ordering}; use tinywasm_types::MemoryType; +use tinywasm_types::Shared; +use super::instance::MemoryCharge; use super::{MemoryInstance, MemoryStorage, memory_oob}; use crate::{ResourceLimiter, Result, Trap}; @@ -35,6 +37,8 @@ struct MemorySharedInstance { bytes: Mutex, // Lock order: bytes, waiters, then an individual waiter's notified flag. waiters: Mutex>>>, + // Dropped after `bytes`, and retained until the last shared handle is gone. + charge: OnceLock, } struct Waiter { @@ -61,10 +65,20 @@ pub struct MemorySharedGuard<'a> { impl MemoryShared { /// Moves a newly allocated memory into a shared backing. pub(crate) fn from_instance(instance: MemoryInstance) -> Self { - let kind = instance.kind; - let pages = AtomicUsize::new(instance.page_count); - let bytes = Mutex::new(instance.inner); - Self(Arc::new(MemorySharedInstance { kind, pages, bytes, waiters: Mutex::new(BTreeMap::new()) })) + let MemoryInstance { kind, inner, page_count, charge } = instance; + let pages = AtomicUsize::new(page_count); + let bytes = Mutex::new(inner); + let accounting = OnceLock::new(); + if let Some(charge) = charge { + let _ = accounting.set(charge); + } + Self(Arc::new(MemorySharedInstance { + kind, + pages, + bytes, + waiters: Mutex::new(BTreeMap::new()), + charge: accounting, + })) } /// Whether two handles reference the same instance, even across store slots. @@ -167,9 +181,13 @@ impl MemoryShared { pub(crate) fn grow_with_limiter( &self, pages: i64, - limiter: Option<&dyn ResourceLimiter>, + store_limiter: Option<&Shared>, ) -> Result, Trap> { - let Some(limiter) = limiter else { + // A shared backing is charged to its owner, not separately to every importing store. + // Host-created backings acquire an owner on their first limiter-approved guest growth; + // their initial allocation remains host-owned. + let Some(limiter) = self.0.charge.get().map(MemoryCharge::limiter).or_else(|| store_limiter.map(AsRef::as_ref)) + else { return self.lock().grow_inner(pages); }; let kind = self.ty(); @@ -193,12 +211,46 @@ impl MemoryShared { let mut guard = self.lock(); if self.page_count() != current_pages { // The limiter approved a different size transition. Ask again with the new size. + drop(guard); + if allowed && desired != current { + limiter.memory_grow_failed(current, desired); + } continue; } if !allowed { return Ok(None); } - return guard.grow_inner(pages); + let result = guard.grow_inner(pages); + match result { + Ok(Some(previous)) => { + if desired != current { + if let Some(charge) = self.0.charge.get() { + charge.add(desired - current); + } else if let Some(owner) = store_limiter { + // Only a successful growth can bind an uncharged host-created memory. + assert!( + self.0.charge.set(MemoryCharge::new(owner.clone(), desired - current)).is_ok(), + "shared memory accounting owner changed during growth" + ); + } + } + return Ok(Some(previous)); + } + Ok(None) => { + drop(guard); + if desired != current { + limiter.memory_grow_failed(current, desired); + } + return Ok(None); + } + Err(error) => { + drop(guard); + if desired != current { + limiter.memory_grow_failed(current, desired); + } + return Err(error); + } + } } } } diff --git a/crates/tinywasm/src/store/mod.rs b/crates/tinywasm/src/store/mod.rs index e89650a..79d0be0 100644 --- a/crates/tinywasm/src/store/mod.rs +++ b/crates/tinywasm/src/store/mod.rs @@ -48,6 +48,11 @@ pub(crate) use {data::*, element::*, function::*, global::*, table::*, tag::*}; /// trap. Allowing a request does not guarantee that the backing allocation will succeed. Rejected /// growth uses the operation's normal failed-growth result. Rejected initial allocation and GC /// allocation produce [`Trap::OutOfMemory`]. +/// A limiter that reserves the `desired - current` memory bytes on approval can refund failed +/// attempts through `memory_grow_failed` and live charges through `memory_dropped`. These are +/// logical linear-memory bytes, not allocator capacity or total process memory. A shared memory +/// retains its first charging limiter across store imports and releases its charge after the last +/// handle drops; host-created shared memory starts uncharged. /// /// # Example /// ```rust @@ -90,6 +95,20 @@ pub trait ResourceLimiter: Send + Sync { Ok(true) } + /// Reverses a successful `memory_growing` approval when backing allocation fails, or when + /// concurrent growth of a shared memory makes the approved transition stale. + /// + /// `current` and `desired` are the same byte sizes passed to `memory_growing`. This is never + /// called for a rejection or a zero-size transition. A limiter that reserves `desired - current` + /// in `memory_growing` can refund that amount here. + fn memory_grow_failed(&self, _current: usize, _desired: usize) {} + + /// Releases a memory's approved logical byte charges after its backing storage is freed. + /// + /// This can occur after the creating store is dropped when a shared-memory handle survives + /// it. Host-created shared memories are not initially charged to a store limiter. + fn memory_dropped(&self, _charged_bytes: usize) {} + /// Checks a nonzero table allocation or growth request. /// /// Sizes are in elements. `current` is zero for initial allocation. `maximum` is the declared @@ -633,31 +652,47 @@ impl Store { init: impl Fn(MemoryType) -> Result, ) -> Result> { let mut addresses = Vec::with_capacity(memories.len()); + let mut ordinary_count = self.state.memories.len(); + #[cfg(feature = "std")] + let mut shared_count = self.state.shared_memories.len(); for &ty in memories { - let instance = cold_err!(init(ty))?; if ty.shared() { #[cfg(feature = "std")] { - let index = MemAddr::try_from(self.state.shared_memories.len()) + let index = MemAddr::try_from(shared_count) .map_err(|_| Error::UnsupportedFeature("too many shared memories"))?; if index >= SHARED_MEM_BIT - 1 { return Err(Error::UnsupportedFeature("too many shared memories")); } - self.state.shared_memories.push(MemoryShared::from_instance(instance)); + shared_count += 1; addresses.push(index | SHARED_MEM_BIT); } #[cfg(not(feature = "std"))] unreachable!("shared memory instantiation requires std"); } else { - let index = MemAddr::try_from(self.state.memories.len()) - .map_err(|_| Error::UnsupportedFeature("too many memories"))?; + let index = + MemAddr::try_from(ordinary_count).map_err(|_| Error::UnsupportedFeature("too many memories"))?; if index >= SHARED_MEM_BIT { return Err(Error::UnsupportedFeature("too many memories")); } - self.state.memories.push(instance); + ordinary_count += 1; addresses.push(index); } } + + // Keep new memories local until every allocation succeeds. A failed later memory must + // not leave an unreachable, charged earlier memory in a reusable store. + let pending = memories.iter().map(|&ty| cold_err!(init(ty))).collect::>>()?; + for instance in pending { + if instance.kind.shared() { + #[cfg(feature = "std")] + self.state.shared_memories.push(MemoryShared::from_instance(instance)); + #[cfg(not(feature = "std"))] + unreachable!("shared memory instantiation requires std"); + } else { + self.state.memories.push(instance); + } + } Ok(addresses.into_iter()) } diff --git a/crates/tinywasm/src/store/state.rs b/crates/tinywasm/src/store/state.rs index 5026201..9dd9e93 100644 --- a/crates/tinywasm/src/store/state.rs +++ b/crates/tinywasm/src/store/state.rs @@ -4,6 +4,7 @@ use super::*; use crate::engine::Config; use crate::interpreter::{InternalValue, Value32, Value64, Value128}; use crate::store::memory::memory_oob; +use tinywasm_types::Shared; /// Global state that can be manipulated by WebAssembly programs /// @@ -81,13 +82,15 @@ impl State { &mut self, addr: MemAddr, pages: i64, - limiter: Option<&dyn ResourceLimiter>, + limiter: Option<&Shared>, ) -> Result, Trap> { + #[cfg(not(feature = "std"))] + let _ = limiter; #[cfg(feature = "std")] if addr & SHARED_MEM_BIT != 0 { return self.shared_memories[(addr & !SHARED_MEM_BIT) as usize].grow_with_limiter(pages, limiter); } - self.get_mem_mut(addr).grow(pages, limiter) + self.get_mem_mut(addr).grow(pages) } /// Copies between memories, locking at most one shared backing at a time. diff --git a/crates/tinywasm/tests/memory.rs b/crates/tinywasm/tests/memory.rs index cf90c13..898448e 100644 --- a/crates/tinywasm/tests/memory.rs +++ b/crates/tinywasm/tests/memory.rs @@ -1,6 +1,8 @@ +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use tinywasm::engine::Config; use tinywasm::types::{MemoryArch, MemoryType, RefType, RefValue, TableType}; -use tinywasm::{Engine, Memory, ModuleInstance, ResourceLimiter, Store, Table, Trap}; +use tinywasm::{Engine, Imports, Memory, ModuleInstance, ResourceLimiter, Store, Table, Trap}; type TestResult = Result>; @@ -9,6 +11,156 @@ fn store_with_limiter(limiter: impl ResourceLimiter + 'static) -> Store { Store::new(engine) } +#[derive(Clone)] +struct SharedQuota { + used: Arc, + limit: usize, +} + +impl SharedQuota { + fn new(limit: usize) -> Self { + Self { used: Arc::new(AtomicUsize::new(0)), limit } + } + + fn used(&self) -> usize { + self.used.load(Ordering::SeqCst) + } +} + +impl ResourceLimiter for SharedQuota { + fn memory_growing(&self, current: usize, desired: usize, _maximum: Option) -> Result { + let delta = desired - current; + self.used + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |used| { + (delta <= self.limit.saturating_sub(used)).then_some(used + delta) + }) + .map(|_| true) + .or_else(|_| Ok(false)) + } + + fn memory_grow_failed(&self, current: usize, desired: usize) { + self.used.fetch_sub(desired - current, Ordering::SeqCst); + } + + fn memory_dropped(&self, charged_bytes: usize) { + self.used.fetch_sub(charged_bytes, Ordering::SeqCst); + } +} + +#[test] +fn shared_quota_releases_store_memories_for_worker_restart() -> TestResult { + const PAGE: usize = 65_536; + let quota = SharedQuota::new(2 * PAGE); + let engine = Engine::new(Config::new().with_resource_limiter(quota.clone())); + let ty = MemoryType::new(MemoryArch::I32, 1, None, None); + + let mut main = Store::new(engine.clone()); + let _main_memory = Memory::try_new(&mut main, ty)?; + let mut worker = Store::new(engine.clone()); + let _worker_memory = Memory::try_new(&mut worker, ty)?; + assert_eq!(quota.used(), 2 * PAGE); + + let mut denied = Store::new(engine.clone()); + assert!(matches!(Memory::try_new(&mut denied, ty), Err(tinywasm::Error::Trap(Trap::OutOfMemory)))); + assert_eq!(quota.used(), 2 * PAGE); + + drop(worker); + assert_eq!(quota.used(), PAGE); + let _replacement_worker = Memory::try_new(&mut denied, ty)?; + assert_eq!(quota.used(), 2 * PAGE); + drop(denied); + drop(main); + assert_eq!(quota.used(), 0); + Ok(()) +} + +#[test] +fn shared_quota_counts_logical_bytes_for_custom_pages() -> TestResult { + let quota = SharedQuota::new(8); + let mut store = store_with_limiter(quota.clone()); + let ty = MemoryType::new(MemoryArch::I64, 1, Some(3), Some(4)); + let memory = Memory::try_new(&mut store, ty)?; + assert_eq!(quota.used(), 4); + assert_eq!(memory.grow(&mut store, 1)?, Some(1)); + assert_eq!(quota.used(), 8); + assert_eq!(memory.grow(&mut store, 1)?, None); + assert_eq!(quota.used(), 8); + drop(store); + assert_eq!(quota.used(), 0); + Ok(()) +} + +#[test] +fn imported_memory_is_only_charged_once() -> TestResult { + const PAGE: usize = 65_536; + let module = tinywasm::parse_bytes(&wat::parse_str("(module (import \"host\" \"memory\" (memory 1)))")?)?; + let quota = SharedQuota::new(PAGE); + let mut store = store_with_limiter(quota.clone()); + let memory = Memory::try_new(&mut store, MemoryType::new(MemoryArch::I32, 1, None, None))?; + let mut imports = Imports::new(); + imports.define("host", "memory", memory); + let _first = ModuleInstance::instantiate(&mut store, &module, Some(&imports))?; + let _second = ModuleInstance::instantiate(&mut store, &module, Some(&imports))?; + assert_eq!(quota.used(), PAGE); + drop(store); + assert_eq!(quota.used(), 0); + Ok(()) +} + +#[test] +fn failed_second_memory_does_not_strand_first_reservation() -> TestResult { + const PAGE: usize = 65_536; + let module = tinywasm::parse_bytes(&wat::parse_str( + "(module (memory (export \"first\") 1) (memory (export \"second\") 1))", + )?)?; + assert_eq!(module.memory_types.len(), 2); + let quota = SharedQuota::new(PAGE); + let mut store = store_with_limiter(quota.clone()); + let result = ModuleInstance::instantiate(&mut store, &module, None); + assert!(matches!(&result, Err(tinywasm::Error::Trap(Trap::OutOfMemory))), "{result:?}"); + assert_eq!(quota.used(), 0); + let _memory = Memory::try_new(&mut store, MemoryType::new(MemoryArch::I32, 1, None, None))?; + assert_eq!(quota.used(), PAGE); + drop(store); + assert_eq!(quota.used(), 0); + Ok(()) +} + +#[cfg(feature = "send")] +#[test] +fn concurrent_stores_cannot_both_reserve_the_last_page() { + const PAGE: usize = 65_536; + let quota = SharedQuota::new(PAGE); + let engine = Engine::new(Config::new().with_resource_limiter(quota.clone())); + let start = Arc::new(std::sync::Barrier::new(3)); + let release = Arc::new(std::sync::Barrier::new(3)); + let (sender, receiver) = std::sync::mpsc::channel(); + let workers: Vec<_> = (0..2) + .map(|_| { + let engine = engine.clone(); + let start = start.clone(); + let release = release.clone(); + let sender = sender.clone(); + std::thread::spawn(move || { + let mut store = Store::new(engine); + start.wait(); + let accepted = Memory::try_new(&mut store, MemoryType::new(MemoryArch::I32, 1, None, None)).is_ok(); + sender.send(accepted).unwrap(); + release.wait(); + }) + }) + .collect(); + start.wait(); + let accepted = usize::from(receiver.recv().unwrap()) + usize::from(receiver.recv().unwrap()); + assert_eq!(quota.used(), PAGE); + assert_eq!(accepted, 1); + release.wait(); + for worker in workers { + worker.join().unwrap(); + } + assert_eq!(quota.used(), 0); +} + #[test] fn memory_read_write_roundtrip() -> TestResult { let mut store = Store::default(); diff --git a/crates/tinywasm/tests/shared_memory.rs b/crates/tinywasm/tests/shared_memory.rs index 80641e9..f84b0a9 100644 --- a/crates/tinywasm/tests/shared_memory.rs +++ b/crates/tinywasm/tests/shared_memory.rs @@ -1,11 +1,171 @@ #![cfg(feature = "std")] +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; use tinywasm::engine::Config; use tinywasm::types::{MemoryArch, MemoryType, ModuleInner}; use tinywasm::{Engine, Imports, MemoryShared, ModuleInstance, ResourceLimiter, Store, Trap}; type TestResult = Result<(), Box>; +#[derive(Clone)] +struct MemoryLedger { + used: Arc, + failed: Arc, + approval_barrier: Option>, +} + +impl MemoryLedger { + fn new() -> Self { + Self { used: Arc::new(AtomicUsize::new(0)), failed: Arc::new(AtomicUsize::new(0)), approval_barrier: None } + } + + fn used(&self) -> usize { + self.used.load(Ordering::SeqCst) + } +} + +impl ResourceLimiter for MemoryLedger { + fn memory_growing(&self, current: usize, desired: usize, _maximum: Option) -> Result { + self.used.fetch_add(desired - current, Ordering::SeqCst); + if let Some(barrier) = &self.approval_barrier { + barrier.wait(); + } + Ok(true) + } + + fn memory_grow_failed(&self, current: usize, desired: usize) { + self.failed.fetch_add(1, Ordering::SeqCst); + self.used.fetch_sub(desired - current, Ordering::SeqCst); + } + + fn memory_dropped(&self, charged_bytes: usize) { + self.used.fetch_sub(charged_bytes, Ordering::SeqCst); + } +} + +#[test] +fn defined_shared_memory_releases_quota_after_last_handle() -> TestResult { + const PAGE: usize = 65_536; + let module = tinywasm::parse_bytes(&wat::parse_str( + r#"(module (memory (export "memory") 1 2 shared) + (func (export "grow") (result i32) i32.const 1 memory.grow))"#, + )?)?; + let quota = MemoryLedger::new(); + let mut store = Store::new(Engine::new(Config::new().with_resource_limiter(quota.clone()))); + let instance = ModuleInstance::instantiate(&mut store, &module, None)?; + assert_eq!(quota.used(), PAGE); + assert_eq!(instance.func::<(), i32>(&store, "grow")?.call(&mut store, ())?, 1); + assert_eq!(quota.used(), 2 * PAGE); + + let memory = instance.memory_shared("memory")?; + let other_handle = memory.clone(); + drop(instance); + drop(store); + assert_eq!(quota.used(), 2 * PAGE); + drop(memory); + assert_eq!(quota.used(), 2 * PAGE); + drop(other_handle); + assert_eq!(quota.used(), 0); + Ok(()) +} + +#[test] +fn imported_shared_memory_keeps_its_creators_quota() -> TestResult { + const PAGE: usize = 65_536; + let creator_module = tinywasm::parse_bytes(&wat::parse_str(r#"(module (memory (export "memory") 1 2 shared))"#)?)?; + let importer_module = tinywasm::parse_bytes(&wat::parse_str( + r#"(module (import "host" "memory" (memory 1 2 shared)) + (func (export "grow") (result i32) i32.const 1 memory.grow))"#, + )?)?; + let quota = MemoryLedger::new(); + let mut creator_store = Store::new(Engine::new(Config::new().with_resource_limiter(quota.clone()))); + let creator = ModuleInstance::instantiate(&mut creator_store, &creator_module, None)?; + let memory = creator.memory_shared("memory")?; + drop(creator); + drop(creator_store); + assert_eq!(quota.used(), PAGE); + + let mut imports = Imports::new(); + imports.define("host", "memory", memory.clone()); + let mut importer_store = Store::default(); + let importer = ModuleInstance::instantiate(&mut importer_store, &importer_module, Some(&imports))?; + assert_eq!(importer.func::<(), i32>(&importer_store, "grow")?.call(&mut importer_store, ())?, 1); + assert_eq!(quota.used(), 2 * PAGE); + drop(importer); + drop(importer_store); + drop(imports); + drop(memory); + assert_eq!(quota.used(), 0); + Ok(()) +} + +#[test] +fn host_created_shared_memory_charges_only_guest_growth() -> TestResult { + const PAGE: usize = 65_536; + let module = tinywasm::parse_bytes(&wat::parse_str( + r#"(module (import "host" "memory" (memory 1 2 shared)) + (func (export "grow") (result i32) i32.const 1 memory.grow))"#, + )?)?; + let memory = MemoryShared::try_new(MemoryType::new(MemoryArch::I32, 1, Some(2), None))?; + let quota = MemoryLedger::new(); + let mut imports = Imports::new(); + imports.define("host", "memory", memory.clone()); + let mut store = Store::new(Engine::new(Config::new().with_resource_limiter(quota.clone()))); + let instance = ModuleInstance::instantiate(&mut store, &module, Some(&imports))?; + assert_eq!(quota.used(), 0); + assert_eq!(instance.func::<(), i32>(&store, "grow")?.call(&mut store, ())?, 1); + assert_eq!(quota.used(), PAGE); + drop(instance); + drop(store); + drop(imports); + assert_eq!(quota.used(), PAGE); + drop(memory); + assert_eq!(quota.used(), 0); + Ok(()) +} + +#[cfg(feature = "send")] +#[test] +fn stale_shared_growth_approval_is_refunded() -> TestResult { + const PAGE: usize = 65_536; + let module = tinywasm::parse_bytes(&wat::parse_str( + r#"(module (import "host" "memory" (memory 1 2 shared)) + (func (export "grow") (result i32) i32.const 1 memory.grow))"#, + )?)?; + let memory = MemoryShared::try_new(MemoryType::new(MemoryArch::I32, 1, Some(2), None))?; + let mut imports = Imports::new(); + imports.define("host", "memory", memory.clone()); + let mut quota = MemoryLedger::new(); + quota.approval_barrier = Some(Arc::new(std::sync::Barrier::new(2))); + let engine = Engine::new(Config::new().with_resource_limiter(quota.clone())); + + let outcomes: Vec<_> = std::thread::scope(|scope| { + let workers: Vec<_> = (0..2) + .map(|_| { + let module = module.clone(); + let engine = engine.clone(); + let imports = &imports; + scope.spawn(move || -> Result { + let mut store = Store::new(engine); + let instance = ModuleInstance::instantiate(&mut store, &module, Some(imports))?; + instance.func::<(), i32>(&store, "grow")?.call(&mut store, ()) + }) + }) + .collect(); + workers.into_iter().map(|worker| worker.join().unwrap()).collect() + }); + let mut outcomes: Vec<_> = outcomes.into_iter().collect::>()?; + outcomes.sort_unstable(); + assert_eq!(outcomes, [-1, 1]); + assert_eq!(quota.failed.load(Ordering::SeqCst), 1); + assert_eq!(quota.used(), PAGE); + drop(imports); + drop(memory); + assert_eq!(quota.used(), 0); + Ok(()) +} + #[test] fn shared_memory_is_visible_across_stores() -> TestResult { let wasm = wat::parse_str(