Wasmtime: Introduce {Module,Component}::resources_required

crates/wasmtime/src/component/component.rs

@@ -1,6 +1,6 @@
 use crate::code::CodeObject;
 use crate::signatures::SignatureCollection;
-use crate::{Engine, Module};
+use crate::{Engine, Module, ResourcesRequired};
 use anyhow::{bail, Context, Result};
 use serde::{Deserialize, Serialize};
 use std::fs;
@@ -9,7 +9,8 @@ use std::path::Path;
 use std::ptr::NonNull;
 use std::sync::Arc;
 use wasmtime_environ::component::{
-    AllCallFunc, ComponentTypes, StaticModuleIndex, TrampolineIndex, Translator,
+    AllCallFunc, ComponentTypes, GlobalInitializer, InstantiateModule, StaticModuleIndex,
+    TrampolineIndex, Translator,
 };
 use wasmtime_environ::{FunctionLoc, ObjectKind, PrimaryMap, ScopeVec};
 use wasmtime_jit::{CodeMemory, CompiledModuleInfo};
@@ -361,6 +362,92 @@ impl Component {
             ..*dtor.func_ref()
         }
     }
+
+    /// Returns a summary of the resources required to instantiate this
+    /// [`Component`][crate::component::Component].
+    ///
+    /// Note that when a component imports and instantiates another component or
+    /// core module, we cannot determine ahead of time how many resources
+    /// instantiating this component will require, and therefore this method
+    /// will return `None` in these scenarios.
+    ///
+    /// Potential uses of the returned information:
+    ///
+    /// * Determining whether your pooling allocator configuration supports
+    ///   instantiating this component.
+    ///
+    /// * Deciding how many of which `Component` you want to instantiate within
+    ///   a fixed amount of resources, e.g. determining whether to create 5
+    ///   instances of component X or 10 instances of component Y.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # fn main() -> wasmtime::Result<()> {
+    /// use wasmtime::{Config, Engine, component::Component};
+    ///
+    /// let mut config = Config::new();
+    /// config.wasm_multi_memory(true);
+    /// config.wasm_component_model(true);
+    /// let engine = Engine::new(&config)?;
+    ///
+    /// let component = Component::new(&engine, &r#"
+    ///     (component
+    ///         ;; Define a core module that uses two memories.
+    ///         (core module $m
+    ///             (memory 1)
+    ///             (memory 6)
+    ///         )
+    ///
+    ///         ;; Instantiate that core module three times.
+    ///         (core instance $i1 (instantiate (module $m)))
+    ///         (core instance $i2 (instantiate (module $m)))
+    ///         (core instance $i3 (instantiate (module $m)))
+    ///     )
+    /// "#)?;
+    ///
+    /// let resources = component.resources_required()
+    ///     .expect("this component does not import any core modules or instances");
+    ///
+    /// // Instantiating the component will require allocating two memories per
+    /// // core instance, and there are three instances, so six total memories.
+    /// assert_eq!(resources.num_memories, 6);
+    /// assert_eq!(resources.max_initial_memory_size, Some(6));
+    ///
+    /// // The component doesn't need any tables.
+    /// assert_eq!(resources.num_tables, 0);
+    /// assert_eq!(resources.max_initial_table_size, None);
+    /// # Ok(()) }
+    /// ```
+    pub fn resources_required(&self) -> Option<ResourcesRequired> {
+        let mut resources = ResourcesRequired {
+            num_memories: 0,
+            max_initial_memory_size: None,
+            num_tables: 0,
+            max_initial_table_size: None,
+        };
+        for init in &self.env_component().initializers {
+            match init {
+                GlobalInitializer::InstantiateModule(inst) => match inst {
+                    InstantiateModule::Static(index, _) => {
+                        let module = self.static_module(*index);
+                        resources.add(&module.resources_required());
+                    }
+                    InstantiateModule::Import(_, _) => {
+                        // We can't statically determine the resources required
+                        // to instantiate this component.
+                        return None;
+                    }
+                },
+                GlobalInitializer::LowerImport { .. }
+                | GlobalInitializer::ExtractMemory(_)
+                | GlobalInitializer::ExtractRealloc(_)
+                | GlobalInitializer::ExtractPostReturn(_)
+                | GlobalInitializer::Resource(_) => {}
+            }
+        }
+        Some(resources)
+    }
 }
 
 impl ComponentRuntimeInfo for ComponentInner {

crates/wasmtime/src/lib.rs

@@ -404,6 +404,7 @@ mod memory;
 mod module;
 mod profiling;
 mod r#ref;
+mod resources;
 mod signatures;
 mod store;
 mod trampoline;
@@ -422,6 +423,7 @@ pub use crate::memory::*;
 pub use crate::module::Module;
 pub use crate::profiling::GuestProfiler;
 pub use crate::r#ref::ExternRef;
+pub use crate::resources::*;
 #[cfg(feature = "async")]
 pub use crate::store::CallHookHandler;
 pub use crate::store::{

crates/wasmtime/src/module.rs

@@ -1,5 +1,6 @@
-use crate::code::CodeObject;
 use crate::{
+    code::CodeObject,
+    resources::ResourcesRequired,
     signatures::SignatureCollection,
     types::{ExportType, ExternType, ImportType},
     Engine,
@@ -904,6 +905,75 @@ impl Module {
         &self.inner.engine
     }
 
+    /// Returns a summary of the resources required to instantiate this
+    /// [`Module`].
+    ///
+    /// Potential uses of the returned information:
+    ///
+    /// * Determining whether your pooling allocator configuration supports
+    ///   instantiating this module.
+    ///
+    /// * Deciding how many of which `Module` you want to instantiate within a
+    ///   fixed amount of resources, e.g. determining whether to create 5
+    ///   instances of module X or 10 instances of module Y.
+    ///
+    /// # Example
+    ///
+    /// ```
+    /// # fn main() -> wasmtime::Result<()> {
+    /// use wasmtime::{Config, Engine, Module};
+    ///
+    /// let mut config = Config::new();
+    /// config.wasm_multi_memory(true);
+    /// let engine = Engine::new(&config)?;
+    ///
+    /// let module = Module::new(&engine, r#"
+    ///     (module
+    ///         ;; Import a memory. Doesn't count towards required resources.
+    ///         (import "a" "b" (memory 10))
+    ///         ;; Define two local memories. These count towards the required
+    ///         ;; resources.
+    ///         (memory 1)
+    ///         (memory 6)
+    ///     )
+    /// "#)?;
+    ///
+    /// let resources = module.resources_required();
+    ///
+    /// // Instantiating the module will require allocating two memories, and
+    /// // the maximum initial memory size is six Wasm pages.
+    /// assert_eq!(resources.num_memories, 2);
+    /// assert_eq!(resources.max_initial_memory_size, Some(6));
+    ///
+    /// // The module doesn't need any tables.
+    /// assert_eq!(resources.num_tables, 0);
+    /// assert_eq!(resources.max_initial_table_size, None);
+    /// # Ok(()) }
+    /// ```
+    pub fn resources_required(&self) -> ResourcesRequired {
+        let em = self.env_module();
+        let num_memories = u32::try_from(em.memory_plans.len() - em.num_imported_memories).unwrap();
+        let max_initial_memory_size = em
+            .memory_plans
+            .values()
+            .skip(em.num_imported_memories)
+            .map(|plan| plan.memory.minimum)
+            .max();
+        let num_tables = u32::try_from(em.table_plans.len() - em.num_imported_tables).unwrap();
+        let max_initial_table_size = em
+            .table_plans
+            .values()
+            .skip(em.num_imported_tables)
+            .map(|plan| plan.table.minimum)
+            .max();
+        ResourcesRequired {
+            num_memories,
+            max_initial_memory_size,
+            num_tables,
+            max_initial_table_size,
+        }
+    }
+
     /// Returns the `ModuleInner` cast as `ModuleRuntimeInfo` for use
     /// by the runtime.
     pub(crate) fn runtime_info(&self) -> Arc<dyn wasmtime_runtime::ModuleRuntimeInfo> {

crates/wasmtime/src/resources.rs

@@ -0,0 +1,33 @@
+/// A summary of the amount of resources required to instantiate a particular
+/// [`Module`][crate::Module] or [`Component`][crate::component::Component].
+///
+/// Example uses of this information:
+///
+/// * Determining whether your pooling allocator configuration supports
+///   instantiating this module.
+///
+/// * Deciding how many of which `Module` you want to instantiate within a
+///   fixed amount of resources, e.g. determining whether to create 5
+///   instances of module `X` or 10 instances of module `Y`.
+pub struct ResourcesRequired {
+    /// The number of memories that are required.
+    pub num_memories: u32,
+    /// The maximum initial size required by any memory, in units of Wasm pages.
+    pub max_initial_memory_size: Option<u64>,
+    /// The number of tables that are required.
+    pub num_tables: u32,
+    /// The maximum initial size required by any table.
+    pub max_initial_table_size: Option<u32>,
+}
+
+impl ResourcesRequired {
+    #[cfg(feature = "component-model")]
+    pub(crate) fn add(&mut self, other: &ResourcesRequired) {
+        self.num_memories += other.num_memories;
+        self.max_initial_memory_size =
+            std::cmp::max(self.max_initial_memory_size, other.max_initial_memory_size);
+        self.num_tables += other.num_tables;
+        self.max_initial_table_size =
+            std::cmp::max(self.max_initial_table_size, other.max_initial_table_size);
+    }
+}