Expose linear memory tracking APIs

crates/wasmtime/src/runtime/limits.rs

@@ -90,6 +90,14 @@ pub trait ResourceLimiter: Send {
         Ok(())
     }
 
+    /// Notifies the resource limiter that a linear-memory growth permitted by
+    /// `memory_growing` has successfully committed.
+    ///
+    /// This is not called for a memory's initial allocation or for shared
+    /// memories. `current` and `desired` are the memory's old and new sizes in
+    /// bytes.
+    fn memory_grown(&mut self, _current: usize, _desired: usize) {}
+
     /// Notifies the resource limiter that an instance's table has been
     /// requested to grow.
     ///
@@ -185,6 +193,9 @@ pub trait ResourceLimiterAsync: Send {
         Ok(())
     }
 
+    /// Identical to [`ResourceLimiter::memory_grown`].
+    fn memory_grown(&mut self, _current: usize, _desired: usize) {}
+
     /// Asynchronous version of [`ResourceLimiter::table_growing`]
     async fn table_growing(
         &mut self,

crates/wasmtime/src/runtime/memory.rs

@@ -14,6 +14,19 @@ use wasmtime_environ::DefinedMemoryIndex;
 
 pub use crate::runtime::vm::WaitResult;
 
+/// A unique linear-memory backing allocated in a [`Store`](crate::Store).
+///
+/// Values returned by [`Store::linear_memories`](crate::Store::linear_memories)
+/// include memories that are not exported. Imports and multiple exports that
+/// refer to the same backing do not produce duplicate values.
+#[derive(Clone, Debug)]
+pub enum StoreMemory {
+    /// An unshared linear memory owned by the store.
+    Unshared(Memory),
+    /// A shared linear memory visible to the store.
+    Shared(SharedMemory),
+}
+
 /// Error for out of bounds [`Memory`] access.
 #[derive(Debug)]
 #[non_exhaustive]
@@ -862,6 +875,10 @@ pub struct SharedMemory {
 }
 
 impl SharedMemory {
+    pub(crate) fn same_backing(&self, other: &Self) -> bool {
+        self.vm.same_backing(&other.vm)
+    }
+
     /// Construct a [`SharedMemory`] by providing both the `minimum` and
     /// `maximum` number of 64K-sized pages. This call allocates the necessary
     /// pages on the system.
@@ -1098,6 +1115,123 @@ impl fmt::Debug for SharedMemory {
 #[cfg(test)]
 mod tests {
     use crate::*;
+    use alloc::vec::Vec;
+
+    #[derive(Default)]
+    struct SuccessfulGrowths(Vec<(usize, usize)>);
+
+    impl ResourceLimiter for SuccessfulGrowths {
+        fn memory_growing(
+            &mut self,
+            _current: usize,
+            _desired: usize,
+            _maximum: Option<usize>,
+        ) -> Result<bool> {
+            Ok(true)
+        }
+
+        fn memory_grown(&mut self, current: usize, desired: usize) {
+            self.0.push((current, desired));
+        }
+
+        fn table_growing(
+            &mut self,
+            _current: usize,
+            _desired: usize,
+            _maximum: Option<usize>,
+        ) -> Result<bool> {
+            Ok(true)
+        }
+    }
+
+    #[test]
+    fn limiter_observes_only_successful_post_instantiation_growth() -> Result<()> {
+        let engine = Engine::default();
+        let module = Module::new(&engine, r#"(module (memory (export "m") 1 2))"#)?;
+        let mut store = Store::new(&engine, SuccessfulGrowths::default());
+        store.limiter(|state| state);
+        let instance = Instance::new(&mut store, &module, &[])?;
+        assert!(store.data().0.is_empty());
+
+        let memory = instance.get_memory(&mut store, "m").unwrap();
+        assert_eq!(memory.grow(&mut store, 1)?, 1);
+        assert_eq!(store.data().0, [(65536, 2 * 65536)]);
+
+        assert!(memory.grow(&mut store, 1).is_err());
+        assert_eq!(store.data().0, [(65536, 2 * 65536)]);
+        Ok(())
+    }
+
+    #[test]
+    fn linear_memories_include_unique_non_exported_backings() -> Result<()> {
+        let engine = Engine::default();
+        let module = Module::new(
+            &engine,
+            r#"(module
+                (memory $aliased 2 3)
+                (export "a" (memory $aliased))
+                (export "b" (memory $aliased))
+                (memory 4 5)
+            )"#,
+        )?;
+        let mut store = Store::new(&engine, ());
+        Instance::new(&mut store, &module, &[])?;
+
+        let mut sizes = store
+            .linear_memories()
+            .into_iter()
+            .map(|memory| match memory {
+                StoreMemory::Unshared(memory) => memory.data_size(&store),
+                StoreMemory::Shared(_) => unreachable!(),
+            })
+            .collect::<Vec<_>>();
+        sizes.sort_unstable();
+
+        assert_eq!(sizes, [2 * 65536, 4 * 65536]);
+        Ok(())
+    }
+
+    #[cfg(feature = "threads")]
+    #[test]
+    fn linear_memories_include_shared_and_imported_backings_once() -> Result<()> {
+        let mut config = Config::new();
+        config.wasm_threads(true).shared_memory(true);
+        let engine = Engine::new(&config)?;
+        let module = Module::new(
+            &engine,
+            r#"(module
+                (import "env" "imported" (memory $imported 5 10 shared))
+                (memory $owned 2 3)
+                (export "a" (memory $owned))
+                (export "b" (memory $owned))
+                (memory 1 2 shared)
+            )"#,
+        )?;
+        let mut store = Store::new(&engine, ());
+        let imported = SharedMemory::new(&engine, MemoryType::shared(5, 10))?;
+        Instance::new(&mut store, &module, &[imported.clone().into()])?;
+        Instance::new(&mut store, &module, &[imported.into()])?;
+
+        let memories = store.linear_memories();
+        assert_eq!(memories.len(), 5);
+        assert_eq!(
+            memories
+                .iter()
+                .filter(|memory| matches!(memory, StoreMemory::Unshared(_)))
+                .count(),
+            2
+        );
+        let mut shared_sizes = memories
+            .iter()
+            .filter_map(|memory| match memory {
+                StoreMemory::Shared(memory) => Some(memory.data_size()),
+                StoreMemory::Unshared(_) => None,
+            })
+            .collect::<Vec<_>>();
+        shared_sizes.sort_unstable();
+        assert_eq!(shared_sizes, [65536, 65536, 5 * 65536]);
+        Ok(())
+    }
 
     // Assert that creating a memory via `Memory::new` respects the limits/tunables
     // in `Config`.

crates/wasmtime/src/runtime/store.rs

@@ -97,7 +97,7 @@ use crate::trampoline::VMHostGlobalContext;
 #[cfg(feature = "debug")]
 use crate::{BreakpointState, DebugHandler, FrameDataCache};
 use crate::{Engine, Module, Val, ValRaw, module::ModuleRegistry};
-use crate::{Global, Instance, Table};
+use crate::{Global, Instance, SharedMemory, StoreMemory, Table};
 use core::convert::Infallible;
 use core::fmt;
 #[cfg(any(feature = "async", feature = "gc"))]
@@ -346,6 +346,14 @@ impl StoreResourceLimiter<'_> {
         }
     }
 
+    pub(crate) fn memory_grown(&mut self, current: usize, desired: usize) {
+        match self {
+            Self::Sync(s) => s.memory_grown(current, desired),
+            #[cfg(feature = "async")]
+            Self::Async(s) => s.memory_grown(current, desired),
+        }
+    }
+
     pub(crate) async fn table_growing(
         &mut self,
         current: usize,
@@ -1000,6 +1008,29 @@ impl<T> Store<T> {
         self.inner.engine()
     }
 
+    /// Returns every unique linear-memory backing allocated in this store.
+    ///
+    /// This includes non-exported memories and host-created memories. Imported
+    /// aliases and multiple exports of one backing are returned only once.
+    pub fn linear_memories(&self) -> Vec<StoreMemory> {
+        let mut memories = Vec::new();
+        for memory in self.inner.all_memories() {
+            match memory {
+                ExportMemory::Unshared(memory) => memories.push(StoreMemory::Unshared(memory)),
+                ExportMemory::Shared(memory, _) => {
+                    let memory = SharedMemory::from_raw(memory, self.engine().clone());
+                    if !memories.iter().any(|existing| match existing {
+                        StoreMemory::Shared(existing) => existing.same_backing(&memory),
+                        StoreMemory::Unshared(_) => false,
+                    }) {
+                        memories.push(StoreMemory::Shared(memory));
+                    }
+                }
+            }
+        }
+        memories
+    }
+
     /// Returns the amount fuel in this [`Store`]. When fuel is enabled, it must
     /// be configured via [`Store::set_fuel`].
     ///

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

@@ -688,23 +688,30 @@ unsafe impl InstanceAllocator for PoolingInstanceAllocator {
         request: &'a mut InstanceAllocationRequest<'b, 'c>,
         ty: &'a wasmtime_environ::Memory,
         memory_index: Option<DefinedMemoryIndex>,
-        _memory_kind: MemoryKind,
+        memory_kind: MemoryKind,
     ) -> Pin<Box<dyn Future<Output = Result<(MemoryAllocationIndex, Memory)>> + Send + 'a>> {
         crate::runtime::box_future(async move {
             async {
                 // FIXME(rust-lang/rust#145127) this should ideally use a version of
                 // `with_flush_and_retry` but adapted for async closures instead of only
                 // sync closures. Right now that won't compile though so this is the
                 // manually expanded version of the method.
-                let e = match self.memories.allocate(request, ty, memory_index).await {
+                let e = match self
+                    .memories
+                    .allocate(request, ty, memory_index, memory_kind)
+                    .await
+                {
                     Ok(result) => return Ok(result),
                     Err(e) => e,
                 };
 
                 if e.is::<PoolConcurrencyLimitError>() {
                     let queue = self.decommit_queue.lock().unwrap();
                     if self.flush_decommit_queue(queue) {
-                        return self.memories.allocate(request, ty, memory_index).await;
+                        return self
+                            .memories
+                            .allocate(request, ty, memory_index, memory_kind)
+                            .await;
                     }
                 }
 

crates/wasmtime/src/runtime/vm/instance/allocator/pooling/memory_pool.rs

@@ -352,9 +352,10 @@ impl MemoryPool {
         request: &mut InstanceAllocationRequest<'_, '_>,
         ty: &wasmtime_environ::Memory,
         memory_index: Option<DefinedMemoryIndex>,
+        memory_kind: MemoryKind,
     ) -> Result<(MemoryAllocationIndex, Memory)> {
         let tunables = request.store.engine().tunables();
-        let memory_tunables = MemoryTunables::new(tunables, MemoryKind::LinearMemory);
+        let memory_tunables = MemoryTunables::new(tunables, memory_kind);
         let stripe_index = if let Some(pkey) = request.store.get_pkey() {
             pkey.as_stripe()
         } else {
@@ -427,6 +428,7 @@ impl MemoryPool {
         let memory = Memory::new_static(
             ty,
             &memory_tunables,
+            memory_kind,
             MemoryBase::Mmap(base),
             base_capacity.byte_count(),
             slot,

crates/wasmtime/src/runtime/vm/memory.rs

@@ -251,7 +251,7 @@ impl Memory {
         let memory_tunables = MemoryTunables::new(tunables, kind);
         let allocation = creator.new_memory(ty, &memory_tunables, minimum, maximum)?;
 
-        let memory = LocalMemory::new(ty, &memory_tunables, allocation, memory_image)?;
+        let memory = LocalMemory::new(ty, &memory_tunables, allocation, memory_image, kind)?;
         Ok(if ty.shared {
             Memory::Shared(SharedMemory::wrap(engine, ty, memory)?)
         } else {
@@ -264,6 +264,7 @@ impl Memory {
     pub async fn new_static(
         ty: &wasmtime_environ::Memory,
         memory_tunables: &MemoryTunables<'_>,
+        kind: MemoryKind,
         base: MemoryBase,
         base_capacity: usize,
         memory_image: MemoryImageSlot,
@@ -277,7 +278,7 @@ impl Memory {
         // `LocalMemory` structure created, notably we already have
         // `memory_image` and regardless of configuration settings this memory
         // can't move its base pointer since it's a fixed allocation.
-        let mut memory = LocalMemory::new(ty, memory_tunables, allocation, None)?;
+        let mut memory = LocalMemory::new(ty, memory_tunables, allocation, None, kind)?;
         assert!(memory.memory_image.is_none());
         memory.memory_image = Some(memory_image);
         memory.memory_may_move = false;
@@ -536,6 +537,7 @@ impl Memory {
 pub struct LocalMemory {
     alloc: Box<dyn RuntimeLinearMemory>,
     ty: wasmtime_environ::Memory,
+    kind: MemoryKind,
     memory_may_move: bool,
     memory_guard_size: usize,
     memory_reservation: usize,
@@ -551,6 +553,7 @@ impl LocalMemory {
         memory_tunables: &MemoryTunables<'_>,
         alloc: Box<dyn RuntimeLinearMemory>,
         memory_image: Option<&Arc<MemoryImage>>,
+        kind: MemoryKind,
     ) -> Result<LocalMemory> {
         // If a memory image was specified, try to create the MemoryImageSlot on
         // top of our mmap.
@@ -584,6 +587,7 @@ impl LocalMemory {
         };
         Ok(LocalMemory {
             ty: *ty,
+            kind,
             alloc,
             memory_may_move: ty.memory_may_move(memory_tunables),
             memory_image,
@@ -710,6 +714,12 @@ impl LocalMemory {
                     assert_eq!(base_ptr_before, self.alloc.base().as_mut_ptr());
                 }
 
+                if matches!(self.kind, MemoryKind::LinearMemory)
+                    && let Some(limiter) = limiter
+                {
+                    limiter.memory_grown(old_byte_size, new_byte_size);
+                }
+
                 Ok(Some((old_byte_size, new_byte_size)))
             }
             Err(e) => {

crates/wasmtime/src/runtime/vm/memory/shared_memory.rs

@@ -30,6 +30,10 @@ struct SharedMemoryInner {
 }
 
 impl SharedMemory {
+    pub(crate) fn same_backing(&self, other: &Self) -> bool {
+        Arc::ptr_eq(&self.0, &other.0)
+    }
+
     /// Construct a new [`SharedMemory`].
     pub fn new(engine: &Engine, ty: &wasmtime_environ::Memory) -> Result<Self> {
         let tunables = engine.tunables();
@@ -46,7 +50,13 @@ impl SharedMemory {
         Self::wrap(
             engine,
             ty,
-            LocalMemory::new(ty, &memory_tunables, boxed, None)?,
+            LocalMemory::new(
+                ty,
+                &memory_tunables,
+                boxed,
+                None,
+                wasmtime_environ::MemoryKind::LinearMemory,
+            )?,
         )
     }
 

crates/wasmtime/src/runtime/vm/memory/shared_memory_disabled.rs

@@ -11,6 +11,10 @@ use wasmtime_environ::Trap;
 pub enum SharedMemory {}
 
 impl SharedMemory {
+    pub(crate) fn same_backing(&self, _other: &Self) -> bool {
+        unreachable!()
+    }
+
     pub fn wrap(_: &Engine, _ty: &wasmtime_environ::Memory, _memory: LocalMemory) -> Result<Self> {
         bail!("support for shared memories was disabled at compile time");
     }