Add PoolingAllocatorMetrics

This exposes some basic runtime metrics derived from the internal state
of a `PoolingInstanceAllocator`.

Two new atomics were added to PoolingInstanceAllocator: `live_memories`
and `live_tables`. While these counts could be derived from existing
state it would require acquiring mutexes on some inner state.

Author: lann

crates/wasmtime/src/engine.rs

@@ -720,6 +720,13 @@ impl Engine {
         crate::runtime::vm::tls_eager_initialize();
     }
 
+    /// Returns a [`PoolingAllocatorMetrics`] if this engine was configured with
+    /// [`InstanceAllocationStrategy::Pooling`].
+    #[cfg(feature = "pooling-allocator")]
+    pub fn pooling_allocator_metrics(&self) -> Option<crate::vm::PoolingAllocatorMetrics> {
+        crate::runtime::vm::PoolingAllocatorMetrics::new(self)
+    }
+
     pub(crate) fn allocator(&self) -> &dyn crate::runtime::vm::InstanceAllocator {
         self.inner.allocator.as_ref()
     }

crates/wasmtime/src/runtime.rs

@@ -95,7 +95,7 @@ pub use values::*;
 pub(crate) use uninhabited::*;
 
 #[cfg(feature = "pooling-allocator")]
-pub use vm::PoolConcurrencyLimitError;
+pub use vm::{PoolConcurrencyLimitError, PoolingAllocatorMetrics};
 
 #[cfg(feature = "profiling")]
 mod profiling;

crates/wasmtime/src/runtime/vm.rs

@@ -102,7 +102,7 @@ pub use crate::runtime::vm::instance::{
 };
 #[cfg(feature = "pooling-allocator")]
 pub use crate::runtime::vm::instance::{
-    InstanceLimits, PoolConcurrencyLimitError, PoolingInstanceAllocator,
+    InstanceLimits, PoolConcurrencyLimitError, PoolingAllocatorMetrics, PoolingInstanceAllocator,
     PoolingInstanceAllocatorConfig,
 };
 pub use crate::runtime::vm::interpreter::*;

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

@@ -32,7 +32,7 @@ pub use self::on_demand::OnDemandInstanceAllocator;
 mod pooling;
 #[cfg(feature = "pooling-allocator")]
 pub use self::pooling::{
-    InstanceLimits, PoolConcurrencyLimitError, PoolingInstanceAllocator,
+    InstanceLimits, PoolConcurrencyLimitError, PoolingAllocatorMetrics, PoolingInstanceAllocator,
     PoolingInstanceAllocatorConfig,
 };
 
@@ -357,6 +357,12 @@ pub unsafe trait InstanceAllocatorImpl {
 
     /// Allow access to memory regions protected by any protection key.
     fn allow_all_pkeys(&self);
+
+    /// Returns `Some(&PoolingInstanceAllocator)` if this is one.
+    #[cfg(feature = "pooling-allocator")]
+    fn as_pooling(&self) -> Option<&PoolingInstanceAllocator> {
+        None
+    }
 }
 
 /// A thing that can allocate instances.

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

@@ -21,6 +21,7 @@
 mod decommit_queue;
 mod index_allocator;
 mod memory_pool;
+mod metrics;
 mod table_pool;
 
 #[cfg(feature = "gc")]
@@ -54,6 +55,7 @@ use crate::runtime::vm::{
     mpk::{self, ProtectionKey, ProtectionMask},
     sys::vm::PageMap,
 };
+use core::sync::atomic::AtomicUsize;
 use std::borrow::Cow;
 use std::fmt::Display;
 use std::sync::{Mutex, MutexGuard};
@@ -65,6 +67,8 @@ use wasmtime_environ::{
     DefinedMemoryIndex, DefinedTableIndex, HostPtr, Module, Tunables, VMOffsets,
 };
 
+pub use self::metrics::PoolingAllocatorMetrics;
+
 #[cfg(feature = "gc")]
 use super::GcHeapAllocationIndex;
 #[cfg(feature = "gc")]
@@ -305,8 +309,12 @@ pub struct PoolingInstanceAllocator {
     live_component_instances: AtomicU64,
 
     decommit_queue: Mutex<DecommitQueue>,
+
     memories: MemoryPool,
+    live_memories: AtomicUsize,
+
     tables: TablePool,
+    live_tables: AtomicUsize,
 
     #[cfg(feature = "gc")]
     gc_heaps: GcHeapPool,
@@ -335,6 +343,8 @@ impl Drop for PoolingInstanceAllocator {
 
         debug_assert_eq!(self.live_component_instances.load(Ordering::Acquire), 0);
         debug_assert_eq!(self.live_core_instances.load(Ordering::Acquire), 0);
+        debug_assert_eq!(self.live_memories.load(Ordering::Acquire), 0);
+        debug_assert_eq!(self.live_tables.load(Ordering::Acquire), 0);
 
         debug_assert!(self.memories.is_empty());
         debug_assert!(self.tables.is_empty());
@@ -357,7 +367,9 @@ impl PoolingInstanceAllocator {
             live_core_instances: AtomicU64::new(0),
             decommit_queue: Mutex::new(DecommitQueue::default()),
             memories: MemoryPool::new(config, tunables)?,
+            live_memories: AtomicUsize::new(0),
             tables: TablePool::new(config)?,
+            live_tables: AtomicUsize::new(0),
             #[cfg(feature = "gc")]
             gc_heaps: GcHeapPool::new(config)?,
             #[cfg(feature = "async")]
@@ -651,7 +663,12 @@ unsafe impl InstanceAllocatorImpl for PoolingInstanceAllocator {
         tunables: &Tunables,
         memory_index: Option<DefinedMemoryIndex>,
     ) -> Result<(MemoryAllocationIndex, Memory)> {
+        self.live_memories.fetch_add(1, Ordering::AcqRel);
         self.with_flush_and_retry(|| self.memories.allocate(request, ty, tunables, memory_index))
+            .map_err(|err| {
+                self.live_memories.fetch_sub(1, Ordering::AcqRel);
+                err
+            })
     }
 
     unsafe fn deallocate_memory(
@@ -660,6 +677,7 @@ unsafe impl InstanceAllocatorImpl for PoolingInstanceAllocator {
         allocation_index: MemoryAllocationIndex,
         memory: Memory,
     ) {
+        self.live_memories.fetch_sub(1, Ordering::AcqRel);
         // Reset the image slot. If there is any error clearing the
         // image, just drop it here, and let the drop handler for the
         // slot unmap in a way that retains the address space
@@ -696,7 +714,12 @@ unsafe impl InstanceAllocatorImpl for PoolingInstanceAllocator {
         tunables: &Tunables,
         _table_index: DefinedTableIndex,
     ) -> Result<(super::TableAllocationIndex, Table)> {
+        self.live_tables.fetch_add(1, Ordering::AcqRel);
         self.with_flush_and_retry(|| self.tables.allocate(request, ty, tunables))
+            .map_err(|err| {
+                self.live_tables.fetch_sub(1, Ordering::AcqRel);
+                err
+            })
     }
 
     unsafe fn deallocate_table(
@@ -705,6 +728,7 @@ unsafe impl InstanceAllocatorImpl for PoolingInstanceAllocator {
         allocation_index: TableAllocationIndex,
         mut table: Table,
     ) {
+        self.live_tables.fetch_sub(1, Ordering::AcqRel);
         let mut queue = DecommitQueue::default();
         // SAFETY: This table is no longer in use by the allocator when this
         // method is called and additionally all image ranges are pushed with
@@ -786,6 +810,10 @@ unsafe impl InstanceAllocatorImpl for PoolingInstanceAllocator {
     ) -> (MemoryAllocationIndex, Memory) {
         self.gc_heaps.deallocate(allocation_index, gc_heap)
     }
+
+    fn as_pooling(&self) -> Option<&PoolingInstanceAllocator> {
+        Some(self)
+    }
 }
 
 #[cfg(test)]

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

@@ -0,0 +1,114 @@
+use core::sync::atomic::Ordering;
+
+use crate::{Engine, vm::PoolingInstanceAllocator};
+
+/// `PoolingAllocatorMetrics` provides access to runtime metrics of a pooling
+/// allocator configured with [`crate::InstanceAllocationStrategy::Pooling`].
+///
+/// This is a cheap cloneable handle which can be obtained with
+/// [`Engine::pooling_allocator_metrics`].
+#[derive(Clone)]
+pub struct PoolingAllocatorMetrics {
+    engine: Engine,
+}
+
+impl PoolingAllocatorMetrics {
+    pub(crate) fn new(engine: &Engine) -> Option<Self> {
+        engine.allocator().as_pooling().map(|_| Self {
+            engine: engine.clone(),
+        })
+    }
+
+    /// Returns the number of core (module) instances currently allocated.
+    pub fn core_instances(&self) -> u64 {
+        self.allocator().live_core_instances.load(Ordering::Acquire)
+    }
+
+    /// Returns the number of component instances currently allocated.
+    pub fn component_instances(&self) -> u64 {
+        self.allocator()
+            .live_component_instances
+            .load(Ordering::Acquire)
+    }
+
+    /// Returns the number of WebAssembly memories currently allocated.
+    pub fn memories(&self) -> usize {
+        self.allocator().live_memories.load(Ordering::Acquire)
+    }
+
+    /// Returns the number of WebAssembly tables currently allocated.
+    pub fn tables(&self) -> usize {
+        self.allocator().live_tables.load(Ordering::Acquire)
+    }
+
+    fn allocator(&self) -> &PoolingInstanceAllocator {
+        self.engine
+            .allocator()
+            .as_pooling()
+            .expect("engine should have pooling allocator")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::{
+        Config, InstanceAllocationStrategy, Store,
+        component::{Component, Linker},
+    };
+
+    use super::*;
+
+    // A component with 1 core instance, 1 memory, 1 table
+    const TEST_COMPONENT: &[u8] = b"
+        (component
+            (core module $m
+                (memory 1)
+                (table 1 funcref)
+            )
+            (core instance (instantiate (module $m)))
+        )
+    ";
+
+    #[test]
+    fn smoke_test() {
+        // Start with nothing
+        let engine =
+            Engine::new(&Config::new().allocation_strategy(InstanceAllocationStrategy::pooling()))
+                .unwrap();
+        let metrics = engine.pooling_allocator_metrics().unwrap();
+
+        assert_eq!(metrics.core_instances(), 0);
+        assert_eq!(metrics.component_instances(), 0);
+        assert_eq!(metrics.memories(), 0);
+        assert_eq!(metrics.tables(), 0);
+
+        // Instantiate one of each
+        let mut store = Store::new(&engine, ());
+        let component = Component::new(&engine, TEST_COMPONENT).unwrap();
+        let linker = Linker::new(&engine);
+        let instance = linker.instantiate(&mut store, &component).unwrap();
+
+        assert_eq!(metrics.core_instances(), 1);
+        assert_eq!(metrics.component_instances(), 1);
+        assert_eq!(metrics.memories(), 1);
+        assert_eq!(metrics.tables(), 1);
+
+        // Back to nothing
+        let _ = (instance, store);
+
+        assert_eq!(metrics.core_instances(), 0);
+        assert_eq!(metrics.component_instances(), 0);
+        assert_eq!(metrics.memories(), 0);
+        assert_eq!(metrics.tables(), 0);
+    }
+
+    #[test]
+    fn test_non_pooling_allocator() {
+        let engine =
+            Engine::new(&Config::new().allocation_strategy(InstanceAllocationStrategy::OnDemand))
+                .unwrap();
+
+        let maybe_metrics = engine.pooling_allocator_metrics();
+        assert!(maybe_metrics.is_none());
+    }
+}