fix: balance linear-memory limiter reservations

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<StoreShared<dyn ResourceLimiter>>,
+    pub resource_limiter: Option<Shared<dyn ResourceLimiter>>,
     /// 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<dyn ResourceLimiter>));
+        self.resource_limiter = Some(Shared::from(Box::new(limiter) as Box<dyn ResourceLimiter>));
         self
     }
 

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);
         }
 

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),

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<Option<i64>> {
-        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.

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};

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<GcObject>,
     directory: Vec<(u32, Handle)>,
     next_key: u32,
-    resource_limiter: Option<StoreShared<dyn ResourceLimiter>>,
+    resource_limiter: Option<Shared<dyn ResourceLimiter>>,
 }
 
 impl Default for GcHeap {

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<MemoryCharge>,
+}
+
+/// The part of a memory's logical size approved by one limiter.
+pub(super) struct MemoryCharge {
+    limiter: Shared<dyn ResourceLimiter>,
+    bytes: AtomicUsize,
+}
+
+impl MemoryCharge {
+    pub(super) fn new(limiter: Shared<dyn ResourceLimiter>, 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<Self> {
+    pub(crate) fn new(kind: MemoryType, limiter: Option<Shared<dyn ResourceLimiter>>) -> Result<Self> {
+        Self::new_with_storage(kind, limiter, MemoryStorage::try_new)
+    }
+
+    fn new_with_storage(
+        kind: MemoryType,
+        limiter: Option<Shared<dyn ResourceLimiter>>,
+        allocate: impl FnOnce(MemoryArch, usize, usize) -> core::result::Result<MemoryStorage, Trap>,
+    ) -> Result<Self> {
         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<Option<i64>, 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<Option<i64>, 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<Option<i64>, 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<AtomicUsize>,
+    }
+
+    impl ResourceLimiter for ReservingLimiter {
+        fn memory_growing(&self, current: usize, desired: usize, _maximum: Option<usize>) -> Result<bool, Trap> {
+            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<AtomicUsize>) -> Shared<dyn ResourceLimiter> {
+        Shared::from(Box::new(ReservingLimiter { used: used.clone() }) as Box<dyn ResourceLimiter>)
+    }
+
+    #[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);
+    }
+}