Allow instance allocators control over module compilation.

This commit introduces two new methods on `InstanceAllocator`:

* `validate_module` - this method is used to validate a module after
  translation but before compilation. It will be used for the upcoming pooling
  allocator to ensure a module being compiled adheres to the limits of the
  allocator.

* `adjust_tunables` - this method is used to adjust the `Tunables` given the
  JIT compiler.  The pooling allocator will use this to force all memories to
  be static during compilation.

crates/cache/src/lib.rs

@@ -43,7 +43,7 @@ impl<'config> ModuleCacheEntry<'config> {
     }
 
     /// Gets cached data if state matches, otherwise calls the `compute`.
-    pub fn get_data<T, U, E>(&self, state: T, compute: fn(T) -> Result<U, E>) -> Result<U, E>
+    pub fn get_data<T, U, E>(&self, state: T, compute: impl Fn(T) -> Result<U, E>) -> Result<U, E>
     where
         T: Hash,
         U: Serialize + for<'a> Deserialize<'a>,

crates/cache/src/tests.rs

@@ -65,28 +65,68 @@ fn test_write_read_cache() {
     let entry1 = ModuleCacheEntry::from_inner(ModuleCacheEntryInner::new(compiler1, &cache_config));
     let entry2 = ModuleCacheEntry::from_inner(ModuleCacheEntryInner::new(compiler2, &cache_config));
 
-    entry1.get_data::<_, i32, i32>(1, |_| Ok(100)).unwrap();
-    entry1.get_data::<_, i32, i32>(1, |_| panic!()).unwrap();
+    entry1
+        .get_data(1, |_| -> Result<i32, ()> { Ok(100) })
+        .unwrap();
+    entry1
+        .get_data(1, |_| -> Result<i32, ()> { panic!() })
+        .unwrap();
 
-    entry1.get_data::<_, i32, i32>(2, |_| Ok(100)).unwrap();
-    entry1.get_data::<_, i32, i32>(1, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32, i32>(2, |_| panic!()).unwrap();
+    entry1
+        .get_data(2, |_| -> Result<i32, ()> { Ok(100) })
+        .unwrap();
+    entry1
+        .get_data(1, |_| -> Result<i32, ()> { panic!() })
+        .unwrap();
+    entry1
+        .get_data(2, |_| -> Result<i32, ()> { panic!() })
+        .unwrap();
 
-    entry1.get_data::<_, i32, i32>(3, |_| Ok(100)).unwrap();
-    entry1.get_data::<_, i32, i32>(1, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32, i32>(2, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32, i32>(3, |_| panic!()).unwrap();
+    entry1
+        .get_data(3, |_| -> Result<i32, ()> { Ok(100) })
+        .unwrap();
+    entry1
+        .get_data(1, |_| -> Result<i32, ()> { panic!() })
+        .unwrap();
+    entry1
+        .get_data(2, |_| -> Result<i32, ()> { panic!() })
+        .unwrap();
+    entry1
+        .get_data(3, |_| -> Result<i32, ()> { panic!() })
+        .unwrap();
 
-    entry1.get_data::<_, i32, i32>(4, |_| Ok(100)).unwrap();
-    entry1.get_data::<_, i32, i32>(1, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32, i32>(2, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32, i32>(3, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32, i32>(4, |_| panic!()).unwrap();
+    entry1
+        .get_data(4, |_| -> Result<i32, ()> { Ok(100) })
+        .unwrap();
+    entry1
+        .get_data(1, |_| -> Result<i32, ()> { panic!() })
+        .unwrap();
+    entry1
+        .get_data(2, |_| -> Result<i32, ()> { panic!() })
+        .unwrap();
+    entry1
+        .get_data(3, |_| -> Result<i32, ()> { panic!() })
+        .unwrap();
+    entry1
+        .get_data(4, |_| -> Result<i32, ()> { panic!() })
+        .unwrap();
 
-    entry2.get_data::<_, i32, i32>(1, |_| Ok(100)).unwrap();
-    entry1.get_data::<_, i32, i32>(1, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32, i32>(2, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32, i32>(3, |_| panic!()).unwrap();
-    entry1.get_data::<_, i32, i32>(4, |_| panic!()).unwrap();
-    entry2.get_data::<_, i32, i32>(1, |_| panic!()).unwrap();
+    entry2
+        .get_data(1, |_| -> Result<i32, ()> { Ok(100) })
+        .unwrap();
+    entry1
+        .get_data(1, |_| -> Result<i32, ()> { panic!() })
+        .unwrap();
+    entry1
+        .get_data(2, |_| -> Result<i32, ()> { panic!() })
+        .unwrap();
+    entry1
+        .get_data(3, |_| -> Result<i32, ()> { panic!() })
+        .unwrap();
+    entry1
+        .get_data(4, |_| -> Result<i32, ()> { panic!() })
+        .unwrap();
+    entry2
+        .get_data(1, |_| -> Result<i32, ()> { panic!() })
+        .unwrap();
 }

crates/environ/src/module.rs

@@ -42,8 +42,16 @@ impl MemoryStyle {
         // A heap with a maximum that doesn't exceed the static memory bound specified by the
         // tunables make it static.
         //
-        // If the module doesn't declare an explicit maximum treat it as 4GiB.
-        let maximum = memory.maximum.unwrap_or(WASM_MAX_PAGES);
+        // If the module doesn't declare an explicit maximum treat it as 4GiB when not
+        // requested to use the static memory bound itself as the maximum.
+        let maximum = memory
+            .maximum
+            .unwrap_or(if tunables.static_memory_bound_is_maximum {
+                tunables.static_memory_bound
+            } else {
+                WASM_MAX_PAGES
+            });
+
         if maximum <= tunables.static_memory_bound {
             assert_ge!(tunables.static_memory_bound, memory.minimum);
             return (

crates/environ/src/tunables.rs

@@ -27,6 +27,9 @@ pub struct Tunables {
     /// Whether or not fuel is enabled for generated code, meaning that fuel
     /// will be consumed every time a wasm instruction is executed.
     pub consume_fuel: bool,
+
+    /// Whether or not to treat the static memory bound as the maximum for unbounded heaps.
+    pub static_memory_bound_is_maximum: bool,
 }
 
 impl Default for Tunables {
@@ -62,6 +65,7 @@ impl Default for Tunables {
             parse_wasm_debuginfo: true,
             interruptable: false,
             consume_fuel: false,
+            static_memory_bound_is_maximum: false,
         }
     }
 }

crates/jit/src/instantiate.rs

@@ -101,6 +101,7 @@ impl CompilationArtifacts {
     pub fn build(
         compiler: &Compiler,
         data: &[u8],
+        validate: impl Fn(&ModuleTranslation) -> Result<(), String> + Sync,
     ) -> Result<(usize, Vec<CompilationArtifacts>, TypeTables), SetupError> {
         let (main_module, translations, types) = ModuleEnvironment::new(
             compiler.frontend_config(),
@@ -112,6 +113,8 @@ impl CompilationArtifacts {
 
         let list = maybe_parallel!(translations.(into_iter | into_par_iter))
             .map(|mut translation| {
+                validate(&translation).map_err(|e| SetupError::Validate(e))?;
+
                 let Compilation {
                     obj,
                     unwind_info,

crates/runtime/src/instance/allocator.rs

@@ -24,7 +24,9 @@ use wasmtime_environ::wasm::{
     DefinedFuncIndex, DefinedMemoryIndex, DefinedTableIndex, FuncIndex, GlobalInit, SignatureIndex,
     TableElementType, WasmType,
 };
-use wasmtime_environ::{ir, Module, ModuleType, OwnedDataInitializer, TableElements, VMOffsets};
+use wasmtime_environ::{
+    ir, Module, ModuleTranslation, ModuleType, OwnedDataInitializer, TableElements, VMOffsets,
+};
 
 /// Represents a request for a new runtime instance.
 pub struct InstanceAllocationRequest<'a> {
@@ -80,6 +82,21 @@ pub enum InstantiationError {
 ///
 /// This trait is unsafe as it requires knowledge of Wasmtime's runtime internals to implement correctly.
 pub unsafe trait InstanceAllocator: Send + Sync {
+    /// Validates a module translation.
+    ///
+    /// This is used to ensure a module being compiled is supported by the instance allocator.
+    fn validate_module(&self, translation: &ModuleTranslation) -> Result<(), String> {
+        drop(translation);
+        Ok(())
+    }
+
+    /// Adjusts the tunables prior to creation of any JIT compiler.
+    ///
+    /// This method allows the instance allocator control over tunables passed to a `wasmtime_jit::Compiler`.
+    fn adjust_tunables(&self, tunables: &mut wasmtime_environ::Tunables) {
+        drop(tunables);
+    }
+
     /// Allocates an instance for the given allocation request.
     ///
     /// # Safety

crates/runtime/src/memory.rs

@@ -52,10 +52,6 @@ pub struct MmapMemory {
     // Size in bytes of extra guard pages after the end to optimize loads and stores with
     // constant offsets.
     offset_guard_size: usize,
-
-    // Records whether we're using a bounds-checking strategy which requires
-    // handlers to catch trapping accesses.
-    pub(crate) needs_signal_handlers: bool,
 }
 
 #[derive(Debug)]
@@ -75,15 +71,6 @@ impl MmapMemory {
 
         let offset_guard_bytes = plan.offset_guard_size as usize;
 
-        // If we have an offset guard, or if we're doing the static memory
-        // allocation strategy, we need signal handlers to catch out of bounds
-        // acceses.
-        let needs_signal_handlers = offset_guard_bytes > 0
-            || match plan.style {
-                MemoryStyle::Dynamic => false,
-                MemoryStyle::Static { .. } => true,
-            };
-
         let minimum_pages = match plan.style {
             MemoryStyle::Dynamic => plan.memory.minimum,
             MemoryStyle::Static { bound } => {
@@ -105,7 +92,6 @@ impl MmapMemory {
             mmap: mmap.into(),
             maximum: plan.memory.maximum,
             offset_guard_size: offset_guard_bytes,
-            needs_signal_handlers,
         })
     }
 }

crates/wasmtime/src/config.rs

@@ -764,7 +764,9 @@ impl Config {
 
     pub(crate) fn build_compiler(&self) -> Compiler {
         let isa = self.target_isa();
-        Compiler::new(isa, self.strategy, self.tunables.clone(), self.features)
+        let mut tunables = self.tunables.clone();
+        self.instance_allocator().adjust_tunables(&mut tunables);
+        Compiler::new(isa, self.strategy, tunables, self.features)
     }
 
     pub(crate) fn instance_allocator(&self) -> &dyn InstanceAllocator {

crates/wasmtime/src/module.rs

@@ -307,15 +307,22 @@ impl Module {
     /// # }
     /// ```
     pub fn from_binary(engine: &Engine, binary: &[u8]) -> Result<Module> {
+        // Check with the instance allocator to see if the given module is supported
+        let allocator = engine.config().instance_allocator();
+
         #[cfg(feature = "cache")]
-        let (main_module, artifacts, types) =
-            ModuleCacheEntry::new("wasmtime", engine.cache_config())
-                .get_data((engine.compiler(), binary), |(compiler, binary)| {
-                    CompilationArtifacts::build(compiler, binary)
-                })?;
+        let (main_module, artifacts, types) = ModuleCacheEntry::new(
+            "wasmtime",
+            engine.cache_config(),
+        )
+        .get_data((engine.compiler(), binary), |(compiler, binary)| {
+            CompilationArtifacts::build(compiler, binary, |m| allocator.validate_module(m))
+        })?;
         #[cfg(not(feature = "cache"))]
         let (main_module, artifacts, types) =
-            CompilationArtifacts::build(engine.compiler(), binary)?;
+            CompilationArtifacts::build(engine.compiler(), binary, |m| {
+                allocator.validate_module(m)
+            })?;
 
         let mut modules = CompiledModule::from_artifacts_list(
             artifacts,